Update app1.py

app1.py

@@ -1,21 +1,29 @@
 import streamlit as st
 import pandas as pd
 from pandasai import SmartDataframe
-from pandasai.llm import OpenAI
+from pandasai.llm import OpenAI, ChatGroq
 from dotenv import load_dotenv
 from datasets import load_dataset
 import os
 
 # Load environment variables
 load_dotenv()
-OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")
 
-if not OPENAI_API_KEY:
-    st.error("OpenAI API Key is missing. Make sure you set it in a .env file.")
-    st.stop()
+def initialize_llm(model_choice):
+    """Initialize the chosen LLM based on the user's selection."""
+    groq_api_key = os.getenv("GROQ_API_KEY")
+    openai_api_key = os.getenv("OPENAI_API_KEY")
 
-# Initialize OpenAI LLM
-llm = OpenAI(api_token=OPENAI_API_KEY)
+    if model_choice == "llama-3.3-70b":
+        if not groq_api_key:
+            st.error("Groq API key is missing. Please set the GROQ_API_KEY environment variable.")
+            return None
+        return ChatGroq(groq_api_key=groq_api_key, model="groq/llama-3.3-70b-versatile")
+    elif model_choice == "GPT-4o":
+        if not openai_api_key:
+            st.error("OpenAI API key is missing. Please set the OPENAI_API_KEY environment variable.")
+            return None
+        return OpenAI(api_token=openai_api_key)
 
 # App title and description
 st.title("Patent Analytics: Chat With Your Dataset")
@@ -31,31 +39,36 @@ st.markdown(
 if "df" not in st.session_state:
     st.session_state.df = None
 
-# Dataset input options
-input_option = st.sidebar.radio(
-    "Choose Dataset Input Method",
-    options=["Use Hugging Face Dataset", "Upload CSV File"],
-    index=0
-)
+# Select the model
+model_choice = st.radio("Select LLM", ["GPT-4o", "llama-3.3-70b"], index=0, horizontal=True)
+llm = initialize_llm(model_choice)
+if not llm:
+    st.stop()
 
-# Dataset loading logic
-if input_option == "Use Hugging Face Dataset":
-    dataset_name = st.sidebar.text_input("Enter Hugging Face Dataset Name:", value="HUPD/hupd")
-    if st.sidebar.button("Load Dataset"):
-        try:
-            dataset = load_dataset(dataset_name, name="sample", split="train", trust_remote_code=True)
-            st.session_state.df = pd.DataFrame(dataset)
-            st.sidebar.success(f"Dataset '{dataset_name}' loaded successfully!")
-        except Exception as e:
-            st.sidebar.error(f"Error loading dataset: {e}")
-elif input_option == "Upload CSV File":
-    uploaded_file = st.sidebar.file_uploader("Upload CSV File:", type=["csv"])
-    if uploaded_file:
-        try:
-            st.session_state.df = pd.read_csv(uploaded_file)
-            st.sidebar.success("File uploaded successfully!")
-        except Exception as e:
-            st.sidebar.error(f"Error loading file: {e}")
+def load_dataset_into_session():
+    """Load dataset based on user input."""
+    input_option = st.radio("Choose Dataset Input Method", ["Use Hugging Face Dataset", "Upload CSV File"], index=0)
+
+    if input_option == "Use Hugging Face Dataset":
+        dataset_name = st.text_input("Enter Hugging Face Dataset Name:", value="HUPD/hupd")
+        if st.button("Load Dataset"):
+            try:
+                dataset = load_dataset(dataset_name, split="train", trust_remote_code=True)
+                st.session_state.df = pd.DataFrame(dataset)
+                st.success(f"Dataset '{dataset_name}' loaded successfully!")
+            except Exception as e:
+                st.error(f"Error loading dataset: {e}")
+    elif input_option == "Upload CSV File":
+        uploaded_file = st.file_uploader("Upload CSV File:", type=["csv"])
+        if uploaded_file:
+            try:
+                st.session_state.df = pd.read_csv(uploaded_file)
+                st.success("File uploaded successfully!")
+            except Exception as e:
+                st.error(f"Error loading file: {e}")
+
+# Load dataset
+load_dataset_into_session()
 
 # Show the loaded dataframe preview
 if st.session_state.df is not None:
@@ -77,14 +90,14 @@ if st.session_state.df is not None:
                 # Chat with the dataframe
                 response = chat_df.chat(question)
 
-                # Detect visualizations in the query
-                if "plot" in question.lower() or "graph" in question.lower():
-                    st.write("### Visualization")
+                # Display response
+                if isinstance(response, pd.DataFrame):
+                    st.write("### Response")
+                    st.dataframe(response)
                 else:
                     st.write("### Response")
-                
-                # Display response or plot
-                st.write(response)
+                    st.write(response)
+
                 st.success("Request processed successfully!")
             except Exception as e:
                 st.error(f"An error occurred: {e}")