Update app.py

app.py

@@ -9,6 +9,10 @@ from dotenv import load_dotenv
 import anthropic
 import ast
 import re
+from langchain.agents import AgentType, initialize_agent
+from langchain.tools import Tool
+from langchain.chat_models import ChatOpenAI
+from langchain.memory import ConversationBufferMemory
 
 # Load environment variables
 load_dotenv()
@@ -58,37 +62,27 @@ st.markdown(
 
 st.title("Excel Q&A Chatbot 📊")
 
-# Model Selection
-model_choice = st.selectbox("Select LLM Model", ["OpenAI GPT-3.5", "Claude 3 Haiku", "Mistral-7B"])
-
-# Load appropriate model based on selection
-if model_choice == "Mistral-7B":
-    model_name = "mistralai/Mistral-7B-Instruct"
-    tokenizer = AutoTokenizer.from_pretrained(model_name)
-    model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=torch.float16)
-    def ask_mistral(query):
-        inputs = tokenizer(query, return_tensors="pt").to("cuda")
-        output = model.generate(**inputs)
-        return tokenizer.decode(output[0])
+# Initialize LangChain Agent with Multi-step Reasoning and Memory
+def execute_query(query):
+    memory = ConversationBufferMemory(memory_key="chat_history", return_messages=True)
+    
+    tool = Tool(
+        name="Pandas Query Executor",
+        func=lambda q: eval(q, {"df": df, "pd": pd}),
+        description="Executes Pandas-based queries on uploaded data"
+    )
 
-elif model_choice == "Claude 3 Haiku":
-    client = anthropic.Anthropic(api_key=os.environ["ANTHROPIC_API_KEY"])
-    def ask_claude(query):
-        response = client.messages.create(
-            model="claude-3-haiku",
-            max_tokens=512,
-            messages=[{"role": "user", "content": query}]
-        )
-        return response.content[0]["text"]
+    agent = initialize_agent(
+        tools=[tool],
+        llm=ChatOpenAI(model_name="gpt-3.5-turbo", temperature=0),
+        agent=AgentType.CONVERSATIONAL_REACT_DESCRIPTION,
+        memory=memory,
+        verbose=True
+    )
+    return agent.run(query)
 
-else:
-    client = openai.OpenAI()
-    def ask_gpt(query):
-        response = client.chat.completions.create(
-            model="gpt-3.5-turbo",
-            messages=[{"role": "user", "content": query}]
-        )
-        return response.choices[0].message.content
+# Model Selection
+model_choice = st.selectbox("Select LLM Model", ["OpenAI GPT-3.5", "Claude 3 Haiku", "Mistral-7B"])
 
 # File Upload with validation
 uploaded_file = st.file_uploader("Upload a file", type=["csv", "xlsx", "xls", "json", "tsv"])
@@ -130,47 +124,10 @@ if uploaded_file is not None:
     
     if st.button("Submit Query"):
         if query:
-            # Interpret the query using selected LLM
-            if model_choice == "Mistral-7B":
-                parsed_query = ask_mistral(f"Convert this question into a structured data operation: {query}")
-            elif model_choice == "Claude 3 Haiku":
-                parsed_query = ask_claude(f"Convert this question into a structured data operation: {query}")
-            else:
-                parsed_query = ask_gpt(f"Convert this question into a structured data operation: {query}")
-            
-            # Validate and clean query
-            #parsed_query = re.sub(r"[^a-zA-Z0-9_()\[\]"'., ]", "", parsed_query.strip())
-            parsed_query = re.sub(r"[^a-zA-Z0-9_()\[\]\'., ]", "", parsed_query.strip())
-
-
-
-            st.write(f"Parsed Query: `{parsed_query}`")
-            
-            # Predefined Safe Execution Methods
-            SAFE_OPERATIONS = {
-                "sum": lambda col: df[col].sum(),
-                "mean": lambda col: df[col].mean(),
-                "max": lambda col: df[col].max(),
-                "groupby_sum": lambda col, group_by: df.groupby(group_by)[col].sum()
-            }
-            
-            # Safe Execution
             try:
-                exec_result = eval(parsed_query, {"df": df, "pd": pd, "SAFE_OPERATIONS": SAFE_OPERATIONS})
+                exec_result = execute_query(query)
                 st.write("### Result:")
-                st.write(exec_result if isinstance(exec_result, pd.DataFrame) else str(exec_result))
-                
-                # If numerical data, show a visualization dynamically
-                if isinstance(exec_result, pd.Series):
-                    fig, ax = plt.subplots()
-                    if exec_result.dtype in ["int64", "float64"]:
-                        exec_result.plot(kind="bar", ax=ax)
-                    elif exec_result.dtype == "object":
-                        exec_result.value_counts().plot(kind="bar", ax=ax)
-                    st.pyplot(fig)
-                
-            except SyntaxError as e:
-                st.error(f"Syntax Error in parsed query: {str(e)}")
+                st.write(exec_result)
             except Exception as e:
                 st.error(f"Error executing query: {str(e)}")