Update app.py

app.py

@@ -1,40 +1,38 @@
 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 ExcelPandasAgent:
+class ExcelAIQuerySystem:
     """
-    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.
+    A system to query Excel files using natural language, powered by OpenAI and LangChain.
     """
     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.excel_data: dict[str, pd.DataFrame] = {}
+        self.embeddings = OpenAIEmbeddings()
+        self.excel_data = {}
+        self.sheet_descriptions = {}
+        self.vectorstore = None
         self.logs = []
 
-    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.
-        """
+    def load_excel_file(self, file_path: str) -> str:
+        """Loads and processes an Excel file, generating descriptions and a vector store."""
         self.logs.clear()
-        self.excel_data.clear()
         try:
             excel_file = pd.ExcelFile(file_path)
             sheet_names = excel_file.sheet_names
@@ -43,65 +41,143 @@ class ExcelPandasAgent:
             for sheet_name in sheet_names:
                 try:
                     df = pd.read_excel(file_path, sheet_name=sheet_name)
-                    # The cleaning function is called here for each sheet
                     df = self._clean_dataframe(df)
                     self.excel_data[sheet_name] = df
-                    self.logs.append(f"  - Indexed and cleaned sheet '{sheet_name}' ({df.shape[0]} rows × {df.shape[1]} columns)")
+
+                    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)")
                 except Exception as e:
                     self.logs.append(f"⚠️ Error loading sheet '{sheet_name}': {str(e)}")
                     continue
             
-            self.logs.append("✅ All sheets processed and indexed.")
-            return "\n".join(self.logs), sheet_names
+            self._create_vectorstore()
+            self.logs.append("✅ Vector store created successfully.")
+            return "\n".join(self.logs)
         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 robustly converting types.
-        """
+        """Cleans a DataFrame by removing empty rows/columns and converting data types."""
         df = df.dropna(how='all').dropna(axis=1, how='all').reset_index(drop=True)
         for col in df.columns:
-            # Apply to object columns that might contain mixed numeric/text data
             if df[col].dtype == 'object':
-                # This is the key change. It attempts to convert the column
-                # to numbers. Any value that fails (like 'Apr') becomes NaN (Not a Number).
-                df[col] = pd.to_numeric(df[col], errors='coerce')
+                try:
+                    df[col] = pd.to_datetime(df[col], errors='ignore')
+                except:
+                    pass
+                try:
+                    df[col] = pd.to_numeric(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.
+        """
+        try:
+            return self.llm.invoke(prompt)
+        except Exception:
+            return f"Sheet: {sheet_name}, Columns: {', '.join(list(df.columns))}"
 
-        # Now, find all numeric columns (including those just converted)
-        # and fill any resulting NaN values with 0. This prepares them for calculations.
-        for col in df.select_dtypes(include=np.number).columns:
-            df[col] = df[col].fillna(0)
+    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)
 
-        return df
+    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_sheet(self, query: str, sheet_name: str) -> str:
+    def query_data(self, query: str, selected_sheet: str = None) -> Dict[str, Any]:
         """
-        Processes a user query against a specific sheet using the Pandas Agent.
+        Processes a user query against the loaded Excel data.
+        If a sheet is selected, it queries that sheet directly.
+        Otherwise, it identifies the most relevant sheets.
         """
-        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]
-        
+        results = {'query': query, 'relevant_sheets': [], 'sheet_results': {}, 'summary': '', 'insights': []}
         try:
-            # Create a new pandas agent for each query.
-            pandas_agent = create_pandas_dataframe_agent(
-                self.llm,
-                df,
-                verbose=True,
-                max_iterations=50,
-                max_execution_time=300,
-                agent_executor_kwargs={"handle_parsing_errors": True},
-                allow_dangerous_code=True
-            )
-            # Invoke the agent with the user's query.
-            response = pandas_agent.invoke(query)
+            # If a specific sheet is selected (and it's not the default auto-select), use it.
+            if selected_sheet and selected_sheet != "Auto-Select based on Query":
+                relevant_sheets = [selected_sheet]
+            else:
+                relevant_sheets = self.identify_relevant_sheets(query)
             
-            # The final answer is in the 'output' key of the response dictionary.
-            return response.get('output', 'Sorry, I could not generate an answer.')
+            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}
+            
+            results['summary'] = self._generate_summary(query, results['sheet_results'])
+            results['insights'] = self._extract_insights(results['sheet_results'])
+            return results
         except Exception as e:
-            return f"An error occurred while querying the agent: {str(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)
 
 # --- Gradio Interface ---
 
@@ -112,47 +188,53 @@ def process_file(api_key, file_obj):
     if file_obj is None:
         raise gr.Error("Please upload an Excel file.")
     try:
-        # Instantiate the agent system
-        agent_system = ExcelPandasAgent(api_key)
-        loading_logs, sheet_names = agent_system.load_excel_file(file_obj.name)
+        excel_system = ExcelAIQuerySystem(api_key)
+        loading_logs = excel_system.load_excel_file(file_obj.name)
+        
+        # Get sheet names for the dropdown
+        sheet_names = list(excel_system.excel_data.keys())
+        dropdown_choices = ["Auto-Select based on Query"] + sheet_names
         
-        # Return updates to the UI components
         return (
             loading_logs, 
-            agent_system, 
-            gr.update(choices=sheet_names, value=sheet_names[0] if sheet_names else None, visible=True),
-            gr.update(visible=True), 
-            gr.update(visible=True),
-            gr.update(visible=False, open=False)
+            excel_system, 
+            gr.update(choices=dropdown_choices, value=dropdown_choices[0], visible=True), # Update dropdown
+            gr.update(visible=True), # Query input
+            gr.update(visible=True), # Ask button
+            gr.update(visible=True)  # Results accordion
         )
     except Exception as e:
         raise gr.Error(f"Failed to process file: {e}")
 
-def generate_response(query, sheet_name, system_state):
+def generate_response(query, sheet_selection, system_state):
     """Gradio function to handle user queries and display results."""
     if not query:
-        raise gr.Error("Please enter a question.")
-    if not sheet_name:
-        raise gr.Error("Please select a sheet to query from the dropdown.")
+        raise gr.Error("Please enter a query.")
     if system_state is None:
         raise gr.Error("File not loaded. Please upload and load a file first.")
     
     try:
-        # Call the agent's query method
-        answer = system_state.query_sheet(query, sheet_name)
+        # Pass the selected sheet to the query function
+        result = system_state.query_data(query, selected_sheet=sheet_selection)
+        summary = result.get('summary', 'No summary available.')
+        sheets = ", ".join(result.get('relevant_sheets', []))
+        insights = ", ".join(result.get('insights', []))
         
-        # Return the answer and make the results accordion visible and open
-        return answer, gr.update(visible=True, open=True)
+        details = f"**🔍 Relevant Sheets Identified:**\n{sheets}\n\n"
+        if insights:
+            details += f"**💡 Key Insights:**\n{insights}"
+            
+        return summary, details
     except Exception as e:
         raise gr.Error(f"Error during query: {e}")
 
 # --- UI Layout ---
 
-with gr.Blocks(theme=gr.themes.Soft(primary_hue="blue"), title="Excel AI Agent") as demo:
+with gr.Blocks(theme=gr.themes.Soft(), title="Excel AI Query System") as demo:
     system_state = gr.State(None)
 
-    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.")
+    gr.Markdown("# 📊 Excel AI Query System")
+    gr.Markdown("Upload an Excel file, and ask questions about your data in plain English.")
 
     with gr.Row():
         with gr.Column(scale=1):
@@ -165,26 +247,26 @@ with gr.Blocks(theme=gr.themes.Soft(primary_hue="blue"), title="Excel AI Agent")
             )
             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="Indexing Status", interactive=False, lines=8)
+            status_output = gr.Textbox(label="Loading Status", interactive=False, lines=5)
         
         with gr.Column(scale=2):
             gr.Markdown("### 2. Ask a Question")
             sheet_selector = gr.Dropdown(
-                label="Select a sheet to query",
-                interactive=True,
+                label="📄 Select a Sheet to Query",
+                info="Choose a specific sheet, or let the AI decide automatically.",
                 visible=False
             )
             query_input = gr.Textbox(
                 label="Your Question", 
-                placeholder="e.g., 'What is the sum of the sales column?' or 'Which product had the highest profit in March?'",
-                visible=False,
-                lines=3
+                placeholder="e.g., 'What were the total sales in Q3?' or 'Show me the performance trend for Product X.'",
+                visible=False
             )
             ask_button = gr.Button("Get Answer", variant="primary", visible=False)
             
             results_accordion = gr.Accordion("Results", open=False, visible=False)
             with results_accordion:
-                answer_output = gr.Markdown(label="Answer")
+                summary_output = gr.Markdown(label="Summary")
+                details_output = gr.Markdown(label="Details")
 
     # --- Event Handlers ---
     
@@ -196,10 +278,12 @@ with gr.Blocks(theme=gr.themes.Soft(primary_hue="blue"), title="Excel AI Agent")
     
     ask_button.click(
         fn=generate_response,
-        inputs=[query_input, sheet_selector, system_state],
-        outputs=[answer_output, results_accordion]
+        inputs=[query_input, sheet_selector, system_state], # Add sheet_selector as an input
+        outputs=[summary_output, details_output]
+    ).then(
+        lambda: gr.update(open=True),
+        outputs=results_accordion
     )
 
-
 if __name__ == "__main__":
     demo.launch(share=True)