Update app.py

app.py

@@ -1,38 +1,39 @@
 import pandas as pd
-import numpy as np
-from langchain_openai import OpenAI
-from langchain_core.documents import Document
-from langchain_community.vectorstores import FAISS
-from langchain_openai import OpenAIEmbeddings
-from langchain_text_splitters import RecursiveCharacterTextSplitter
-import re
 import os
-from typing import Dict, List, Any
 import warnings
 import gradio as gr
 from dotenv import load_dotenv
 
+# New imports for the Pandas Agent
+from langchain_openai import OpenAI
+from langchain_experimental.agents.agent_toolkits import create_pandas_dataframe_agent
+
 # Ignore warnings for a cleaner interface
 warnings.filterwarnings('ignore')
 # Load environment variables from .env file
 load_dotenv()
 
-class ExcelAIQuerySystem:
+class ExcelPandasAgent:
     """
-    A system to query Excel files using natural language, powered by OpenAI and LangChain.
+    An agent-based system to query Excel files using natural language,
+    powered by an OpenAI LLM and a Pandas DataFrame Agent.
+    This version can perform mathematical calculations, comparisons, and data analysis.
     """
     def __init__(self, openai_api_key: str):
+        """Initializes the system with the OpenAI API key."""
         os.environ["OPENAI_API_KEY"] = openai_api_key
+        # Using a temperature of 0 for deterministic, factual answers.
         self.llm = OpenAI(temperature=0)
-        self.embeddings = OpenAIEmbeddings()
-        self.excel_data = {}
-        self.sheet_descriptions = {}
-        self.vectorstore = None
+        self.excel_data: dict[str, pd.DataFrame] = {}
         self.logs = []
 
-    def load_excel_file(self, file_path: str) -> str:
-        """Loads and processes an Excel file, generating descriptions and a vector store."""
+    def load_excel_file(self, file_path: str) -> tuple[str, list]:
+        """
+        Loads and processes an Excel file into multiple pandas DataFrames,
+        one for each sheet.
+        """
         self.logs.clear()
+        self.excel_data.clear()
         try:
             excel_file = pd.ExcelFile(file_path)
             sheet_names = excel_file.sheet_names
@@ -43,132 +44,57 @@ class ExcelAIQuerySystem:
                     df = pd.read_excel(file_path, sheet_name=sheet_name)
                     df = self._clean_dataframe(df)
                     self.excel_data[sheet_name] = df
-
-                    description = self._generate_sheet_description(sheet_name, df)
-                    self.sheet_descriptions[sheet_name] = description
-                    self.logs.append(f"  - Loaded and described sheet '{sheet_name}' ({df.shape[0]} rows × {df.shape[1]} columns)")
+                    self.logs.append(f"  - Indexed sheet '{sheet_name}' ({df.shape[0]} rows × {df.shape[1]} columns)")
                 except Exception as e:
                     self.logs.append(f"⚠️ Error loading sheet '{sheet_name}': {str(e)}")
                     continue
             
-            self._create_vectorstore()
-            self.logs.append("✅ Vector store created successfully.")
-            return "\n".join(self.logs)
+            self.logs.append("✅ All sheets processed and indexed.")
+            return "\n".join(self.logs), sheet_names
         except Exception as e:
             raise Exception(f"Error loading Excel file: {str(e)}")
 
     def _clean_dataframe(self, df: pd.DataFrame) -> pd.DataFrame:
-        """Cleans a DataFrame by removing empty rows/columns and converting data types."""
+        """Cleans a DataFrame by removing empty rows/columns."""
         df = df.dropna(how='all').dropna(axis=1, how='all').reset_index(drop=True)
+        # Attempt to convert object columns to numeric or datetime where possible
         for col in df.columns:
             if df[col].dtype == 'object':
                 try:
-                    df[col] = pd.to_datetime(df[col], errors='ignore')
+                    df[col] = pd.to_numeric(df[col], errors='ignore')
                 except:
                     pass
                 try:
-                    df[col] = pd.to_numeric(df[col], errors='ignore')
+                    df[col] = pd.to_datetime(df[col], errors='ignore')
                 except:
                     pass
         return df
 
-    def _generate_sheet_description(self, sheet_name: str, df: pd.DataFrame) -> str:
-        """Generates a text description of a DataFrame using an LLM."""
-        sample_data = df.head(3).to_string()
-        prompt = f"""
-        Analyze this Excel sheet and provide a concise one-paragraph summary.
-        Sheet Name: {sheet_name}
-        Columns: {list(df.columns)}
-        Sample Data:
-        {sample_data}
-        
-        Focus on the main purpose of the data, key metrics, and the time period covered.
+    def query_sheet(self, query: str, sheet_name: str) -> str:
         """
+        Processes a user query against a specific sheet using the Pandas Agent.
+        """
+        if sheet_name not in self.excel_data:
+            return f"Error: Sheet '{sheet_name}' not found. Please select a valid sheet."
+        
+        df = self.excel_data[sheet_name]
+        
         try:
-            return self.llm.invoke(prompt)
-        except Exception:
-            return f"Sheet: {sheet_name}, Columns: {', '.join(list(df.columns))}"
-
-    def _create_vectorstore(self):
-        """Creates a FAISS vector store from sheet descriptions for similarity search."""
-        documents = [
-            Document(page_content=desc, metadata={"sheet_name": name})
-            for name, desc in self.sheet_descriptions.items()
-        ]
-        text_splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200)
-        splits = text_splitter.split_documents(documents)
-        self.vectorstore = FAISS.from_documents(splits, self.embeddings)
-
-    def identify_relevant_sheets(self, query: str) -> List[str]:
-        """Identifies the most relevant sheets for a given query using the vector store."""
-        if not self.vectorstore:
-            return list(self.excel_data.keys())
-        try:
-            docs = self.vectorstore.similarity_search(query, k=3)
-            sheet_names = [doc.metadata['sheet_name'] for doc in docs if 'sheet_name' in doc.metadata]
-            return list(dict.fromkeys(sheet_names))[:5]
-        except Exception:
-            return list(self.excel_data.keys())
-
-    def query_data(self, query: str) -> Dict[str, Any]:
-        """Processes a user query against the loaded Excel data."""
-        results = {'query': query, 'relevant_sheets': [], 'sheet_results': {}, 'summary': '', 'insights': []}
-        try:
-            relevant_sheets = self.identify_relevant_sheets(query)
-            results['relevant_sheets'] = relevant_sheets
-
-            for sheet_name in relevant_sheets:
-                if sheet_name not in self.excel_data:
-                    continue
-                df = self.excel_data[sheet_name]
-                analysis_prompt = f"""
-                Analyze the data from sheet '{sheet_name}' to answer the query: "{query}"
-                Columns: {list(df.columns)}
-                Sample Data:
-                {df.head(5).to_string()}
-                
-                Provide a direct answer, including key numbers, trends, or patterns.
-                """
-                response = self.llm.invoke(analysis_prompt)
-                results['sheet_results'][sheet_name] = {'response': response}
+            # Create a new pandas agent for each query.
+            # verbose=True will print the agent's thought process to the console.
+            pandas_agent = create_pandas_dataframe_agent(
+                self.llm,
+                df,
+                verbose=True,
+                agent_executor_kwargs={"handle_parsing_errors": True} # Helps with robustness
+            )
+            # Invoke the agent with the user's query.
+            response = pandas_agent.invoke(query)
             
-            results['summary'] = self._generate_summary(query, results['sheet_results'])
-            results['insights'] = self._extract_insights(results['sheet_results'])
-            return results
+            # The final answer is in the 'output' key of the response dictionary.
+            return response.get('output', 'Sorry, I could not generate an answer.')
         except Exception as e:
-            results['summary'] = f"Error processing query: {str(e)}"
-            return results
-
-    def _generate_summary(self, query: str, sheet_results: Dict) -> str:
-        """Generates a final, consolidated summary from individual sheet analyses."""
-        if not sheet_results:
-            return "No relevant data found to answer the query."
-        
-        combined_responses = "\n\n".join(
-            f"--- Analysis from Sheet '{name}' ---\n{res['response']}"
-            for name, res in sheet_results.items()
-        )
-        prompt = f"""
-        Based on the following analyses, provide a final, consolidated answer to the query.
-        Original Query: {query}
-        
-        {combined_responses}
-        
-        Synthesize these findings into a clear and direct summary.
-        """
-        return self.llm.invoke(prompt)
-
-    def _extract_insights(self, sheet_results: Dict) -> List[str]:
-        """Extracts simple, actionable insights from the analysis results."""
-        insights = set()
-        for sheet_name, result in sheet_results.items():
-            response = result.get('response', '').lower()
-            if re.search(r'\b\d+\.?\d*\b', response):
-                insights.add(f"Numerical data found in '{sheet_name}'")
-            trend_keywords = ['increase', 'decrease', 'growth', 'decline', 'trend', 'pattern']
-            if any(keyword in response for keyword in trend_keywords):
-                insights.add(f"Trend analysis available in '{sheet_name}'")
-        return list(insights)
+            return f"An error occurred while querying the agent: {str(e)}"
 
 # --- Gradio Interface ---
 
@@ -179,47 +105,50 @@ def process_file(api_key, file_obj):
     if file_obj is None:
         raise gr.Error("Please upload an Excel file.")
     try:
-        excel_system = ExcelAIQuerySystem(api_key)
-        loading_logs = excel_system.load_excel_file(file_obj.name)
+        # Instantiate the agent system
+        agent_system = ExcelPandasAgent(api_key)
+        loading_logs, sheet_names = agent_system.load_excel_file(file_obj.name)
         
+        # Return updates to the UI components
         return (
             loading_logs, 
-            excel_system, 
+            agent_system, 
+            # Populate and show the sheet selector dropdown
+            gr.update(choices=sheet_names, value=sheet_names[0] if sheet_names else None, visible=True),
+            # Show the query box and button
             gr.update(visible=True), 
             gr.update(visible=True),
-            gr.update(visible=True)
+            # Hide the results from any previous run
+            gr.update(visible=False, open=False)
         )
     except Exception as e:
         raise gr.Error(f"Failed to process file: {e}")
 
-def generate_response(query, system_state):
+def generate_response(query, sheet_name, system_state):
     """Gradio function to handle user queries and display results."""
     if not query:
-        raise gr.Error("Please enter a query.")
+        raise gr.Error("Please enter a question.")
+    if not sheet_name:
+        raise gr.Error("Please select a sheet to query from the dropdown.")
     if system_state is None:
         raise gr.Error("File not loaded. Please upload and load a file first.")
     
     try:
-        result = system_state.query_data(query)
-        summary = result.get('summary', 'No summary available.')
-        sheets = ", ".join(result.get('relevant_sheets', []))
-        insights = ", ".join(result.get('insights', []))
+        # Call the agent's query method
+        answer = system_state.query_sheet(query, sheet_name)
         
-        details = f"**🔍 Relevant Sheets Identified:**\n{sheets}\n\n"
-        if insights:
-            details += f"**💡 Key Insights:**\n{insights}"
-            
-        return summary, details
+        # Return the answer and make the results accordion visible and open
+        return answer, gr.update(visible=True, open=True)
     except Exception as e:
         raise gr.Error(f"Error during query: {e}")
 
 # --- UI Layout ---
 
-with gr.Blocks(theme=gr.themes.Soft(), title="Excel AI Query System") as demo:
+with gr.Blocks(theme=gr.themes.Soft(primary_hue="blue"), title="Excel AI Agent") as demo:
     system_state = gr.State(None)
 
-    gr.Markdown("# 📊 Excel AI Query System")
-    gr.Markdown("Upload an Excel file, and ask questions about your data in plain English.")
+    gr.Markdown("# 🤖 Excel AI Agent (Pandas Edition)")
+    gr.Markdown("This version uses a **Pandas Agent** to answer questions by executing code, allowing for mathematical calculations and data analysis.")
 
     with gr.Row():
         with gr.Column(scale=1):
@@ -232,46 +161,43 @@ with gr.Blocks(theme=gr.themes.Soft(), title="Excel AI Query System") as demo:
             )
             file_input = gr.File(label="Upload Excel File", file_types=[".xlsx", ".xls"])
             load_button = gr.Button("Load File", variant="primary")
-            status_output = gr.Textbox(label="Loading Status", interactive=False, lines=5)
+            status_output = gr.Textbox(label="Indexing Status", interactive=False, lines=8)
         
         with gr.Column(scale=2):
             gr.Markdown("### 2. Ask a Question")
+            # This dropdown is now connected to the backend
+            sheet_selector = gr.Dropdown(
+                label="Select a sheet to query",
+                interactive=True,
+                visible=False
+            )
             query_input = gr.Textbox(
                 label="Your Question", 
-                placeholder="e.g., 'What were the total sales in Q3?' or 'Show me the performance trend for Product X.'",
-                visible=False
+                placeholder="e.g., 'What is the sum of the sales column?' or 'Which product had the highest profit in March?'",
+                visible=False,
+                lines=3
             )
             ask_button = gr.Button("Get Answer", variant="primary", visible=False)
             
+            # Simplified results area
             results_accordion = gr.Accordion("Results", open=False, visible=False)
             with results_accordion:
-                summary_output = gr.Markdown(label="Summary")
-                details_output = gr.Markdown(label="Details")
+                answer_output = gr.Markdown(label="Answer")
 
     # --- Event Handlers ---
     
     load_button.click(
         fn=process_file,
         inputs=[api_key_input, file_input],
-        outputs=[status_output, system_state, query_input, ask_button, results_accordion]
+        outputs=[status_output, system_state, sheet_selector, query_input, ask_button, results_accordion]
     )
     
     ask_button.click(
         fn=generate_response,
-        inputs=[query_input, system_state],
-        outputs=[summary_output, details_output]
-    ).then(
-        lambda: gr.update(open=True),
-        outputs=results_accordion
+        inputs=[query_input, sheet_selector, system_state],
+        outputs=[answer_output, results_accordion]
     )
 
 
 if __name__ == "__main__":
-    demo.launch(share=True)
-
-# # --- To this ---
-# if __name__ == "__main__":
-#     # Render provides the PORT environment variable
-#     port = int(os.environ.get('PORT', 10000))
-#     # Launch on 0.0.0.0 to make it accessible outside the container
-#     demo.launch(server_name="0.0.0.0", server_port=port)
+    demo.launch(share=True)