Update app.py

app.py

@@ -1,15 +1,11 @@
 import pandas as pd
-import numpy as np
-from langchain_openai import OpenAIEmbeddings, ChatOpenAI # MODIFIED
+from langchain_openai import OpenAIEmbeddings, ChatOpenAI
 from langchain_core.documents import Document
 from langchain_community.vectorstores import FAISS
 from langchain_text_splitters import RecursiveCharacterTextSplitter
-from langchain.agents.agent_types import AgentType
-from langchain_experimental.agents.agent_toolkits import create_pandas_dataframe_agent
-import re
+from langchain.chains.qa_with_sources import load_qa_with_sources_chain
 import os
-import io
-from typing import Dict, List, Any
+from typing import Dict, Any
 import warnings
 import gradio as gr
 from dotenv import load_dotenv
@@ -18,37 +14,48 @@ warnings.filterwarnings('ignore')
 load_dotenv()
 
 class ExcelAIQuerySystem:
+    """
+    A system to query Excel files using a reliable "Chunk and Search" (RAG) method.
+    This method is good for lookups but not for mathematical aggregations.
+    """
     def __init__(self, openai_api_key: str):
         os.environ["OPENAI_API_KEY"] = openai_api_key
-        # --- USE A MORE CAPABLE CHAT MODEL ---
         self.llm = ChatOpenAI(temperature=0, model_name="gpt-3.5-turbo") 
         self.embeddings = OpenAIEmbeddings()
-        self.excel_data = {}
-        self.sheet_descriptions = {}
-        self.vectorstore = None
+        self.sheet_data_stores: Dict[str, FAISS] = {} # Store a vector store for each sheet
         self.logs = []
+        self.sheet_names = []
 
     def load_excel_file(self, file_path: str) -> str:
         self.logs.clear()
         try:
             excel_file = pd.ExcelFile(file_path)
-            sheet_names = excel_file.sheet_names
-            self.logs.append(f"✅ Found {len(sheet_names)} sheets: {', '.join(sheet_names)}")
+            self.sheet_names = excel_file.sheet_names
+            self.logs.append(f"✅ Found {len(self.sheet_names)} sheets: {', '.join(self.sheet_names)}")
 
-            for sheet_name in sheet_names:
+            text_splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200)
+
+            for sheet_name in self.sheet_names:
                 try:
                     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)")
+                    
+                    # Convert dataframe to a single text document
+                    # Using markdown format for better structure
+                    markdown_text = df.to_markdown(index=False)
+                    
+                    # Create documents and split them into chunks
+                    doc = Document(page_content=markdown_text, metadata={"sheet_name": sheet_name})
+                    chunks = text_splitter.split_documents([doc])
+                    
+                    # Create a FAISS vector store for the chunks
+                    self.sheet_data_stores[sheet_name] = FAISS.from_documents(chunks, self.embeddings)
+                    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)}")
+                    self.logs.append(f"⚠️ Error processing sheet '{sheet_name}': {str(e)}")
                     continue
             
-            self._create_vectorstore()
-            self.logs.append("✅ Vector store created successfully.")
+            self.logs.append("✅ All sheets processed and indexed.")
             return "\n".join(self.logs)
         except Exception as e:
             raise Exception(f"Error loading Excel file: {str(e)}")
@@ -56,102 +63,60 @@ class ExcelAIQuerySystem:
     def _clean_dataframe(self, df: pd.DataFrame) -> pd.DataFrame:
         df = df.dropna(how='all').dropna(axis=1, how='all').reset_index(drop=True)
         df.columns = df.columns.str.strip().str.replace(r'[^a-zA-Z0-9_]', '_', regex=True)
+        # Convert all data to string to ensure consistency for text processing
         for col in df.columns:
-            if df[col].dtype == 'object':
-                try: df[col] = pd.to_datetime(df[col], errors='ignore')
-                except: pass
-                try: df[col] = pd.to_numeric(df[col], errors='ignore')
-                except: pass
+            df[col] = df[col].astype(str)
         return df
 
-    def _generate_sheet_description(self, sheet_name: str, df: pd.DataFrame) -> str:
-        buffer = io.StringIO()
-        df.info(buf=buffer)
-        prompt = f"""
-        Analyze the metadata of this Excel sheet to provide a concise, one-paragraph summary.
-        Sheet Name: {sheet_name}
-        Dataframe Info: {buffer.getvalue()}
-        First 3 Rows: {df.head(3).to_string()}
-        Summary Stats: {df.describe().to_string()}
-        Based on all the metadata, summarize the sheet's main purpose and the types of data it contains.
+    def query_data(self, query: str, target_sheet: str) -> Dict[str, Any]:
         """
-        try:
-            return self.llm.invoke(prompt).content
-        except Exception:
-            return f"Sheet: {sheet_name}, Columns: {', '.join(list(df.columns))}"
-
-    def _create_vectorstore(self):
-        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]:
-        if not self.vectorstore: return list(self.excel_data.keys())
-        try:
-            docs = self.vectorstore.similarity_search(query, k=5)
-            sheet_names = [doc.metadata['sheet_name'] for doc in docs if 'sheet_name' in doc.metadata]
-            return list(dict.fromkeys(sheet_names))
-        except Exception:
-            return list(self.excel_data.keys())
-
-    def query_data(self, query: str, target_sheet: str = "Auto-Select") -> Dict[str, Any]:
-        results = {'query': query, 'relevant_sheets': [], 'sheet_results': {}, 'summary': ''}
+        --- NEW LOGIC ---
+        Searches for relevant data chunks and uses an LLM to answer based on them.
+        """
+        results = {'query': query, 'summary': ''}
         
-        try:
-            if target_sheet and target_sheet != "Auto-Select":
-                relevant_sheets = [target_sheet]
-                if target_sheet not in self.excel_data:
-                    results['summary'] = f"Error: The selected sheet '{target_sheet}' was not found or could not be loaded."
-                    return results
-            else:
-                relevant_sheets = self.identify_relevant_sheets(query)
-
-            results['relevant_sheets'] = relevant_sheets
+        if not target_sheet or target_sheet not in self.sheet_data_stores:
+            results['summary'] = "Error: Please select a valid sheet to query."
+            return results
 
-            for sheet_name in relevant_sheets:
-                if sheet_name not in self.excel_data: continue
-                
-                df = self.excel_data[sheet_name]
-                # --- INCREASE THE ITERATION LIMIT ---
-                pandas_agent = create_pandas_dataframe_agent(
-                    self.llm,
-                    df,
-                    agent_type=AgentType.ZERO_SHOT_REACT_DESCRIPTION,
-                    verbose=True,
-                    allow_dangerous_code=True,
-                    max_iterations=50
-                )
-                response = pandas_agent.invoke(query)
-                results['sheet_results'][sheet_name] = {'response': response['output']}
+        try:
+            vector_store = self.sheet_data_stores[target_sheet]
+            
+            # Find the most relevant data chunks for the query
+            relevant_docs = vector_store.similarity_search(query, k=5)
+            
+            # Create a Question-Answering chain
+            qa_chain = load_qa_with_sources_chain(self.llm, chain_type="stuff")
             
-            results['summary'] = self._generate_summary(query, results['sheet_results'])
+            # Run the chain with the relevant docs
+            response = qa_chain.invoke(
+                {"input_documents": relevant_docs, "question": query},
+                return_only_outputs=True
+            )
+            
+            results['summary'] = response.get('output_text', "Could not find an answer in the data.")
             return results
         except Exception as e:
             results['summary'] = f"An error occurred while querying the data: {str(e)}"
             return results
 
-    def _generate_summary(self, query: str, sheet_results: Dict) -> str:
-        if not sheet_results: return "No relevant data found to answer the query."
-        if len(sheet_results) == 1: return list(sheet_results.values())[0]['response']
+# --- Gradio Interface ---
+# Simplified to work with the new RAG logic
 
-        combined_responses = "\n\n".join([f"--- Analysis from Sheet '{name}' ---\n{res['response']}" for name, res in sheet_results.items()])
-        prompt = f"The following are answers to the query '{query}' from different data sheets. Synthesize them into a single, cohesive final answer.\n\n{combined_responses}\n\nProvide a final, consolidated answer."
-        return self.llm.invoke(prompt).content
-
-# --- Gradio Interface (No changes needed) ---
 def process_file(api_key, file_obj):
     if not api_key: raise gr.Error("OpenAI API Key is required.")
     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)
-        sheet_names = ["Auto-Select"] + list(excel_system.excel_data.keys())
+        
+        # Now a sheet must be selected, so we don't include "Auto-Select"
+        sheet_names = excel_system.sheet_names
         
         return (
             loading_logs, 
             excel_system, 
-            gr.update(choices=sheet_names, value="Auto-Select", visible=True),
+            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=True)
@@ -162,39 +127,39 @@ def process_file(api_key, file_obj):
 def generate_response(query, selected_sheet, system_state):
     if not query: 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.")
+    if not selected_sheet: raise gr.Error("Please select a sheet to query.")
     
     try:
         result = system_state.query_data(query, target_sheet=selected_sheet)
         summary = result.get('summary', 'No summary available.')
-        sheets = ", ".join(result.get('relevant_sheets', []))
-        details = f"**🔍 Sheets Queried:**\n{sheets}"
+        details = f"**🔍 Searched in Sheet:**\n{selected_sheet}"
         return summary, details
     except Exception as e:
         raise gr.Error(f"Error during query: {e}")
 
 with gr.Blocks(theme=gr.themes.Soft(), title="Excel AI Query System") as demo:
     system_state = gr.State(None)
-    gr.Markdown("# 📊 Excel AI Query System")
-    gr.Markdown("Upload an Excel file, choose a specific sheet or let the AI decide, and ask questions about your data.")
+    gr.Markdown("# 📊 Excel AI Query System (Chunk & Search Edition)")
+    gr.Markdown("This version finds specific information in your Excel file. It is not designed for math or whole-dataset calculations.")
     with gr.Row():
         with gr.Column(scale=1):
             gr.Markdown("### 1. Setup")
             api_key_input = gr.Textbox(label="OpenAI API Key", type="password", placeholder="Enter your OpenAI API key...", value=os.getenv("OPENAI_API_KEY", ""))
             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=10)
+            status_output = gr.Textbox(label="Indexing Status", interactive=False, lines=10)
         with gr.Column(scale=2):
             gr.Markdown("### 2. Ask a Question")
             sheet_selector = gr.Dropdown(
                 label="Select a sheet to query",
-                info="Choose 'Auto-Select' to let the AI find the best sheet.",
+                info="You must select a sheet.",
                 visible=False,
                 interactive=True
             )
-            query_input = gr.Textbox(label="Your Question", placeholder="e.g., 'What is the average revenue?'", visible=False)
+            query_input = gr.Textbox(label="Your Question", placeholder="e.g., 'What are the details for order #12345?'", visible=False)
             ask_button = gr.Button("Get Answer", variant="primary", visible=False)
             with gr.Accordion("Results", open=False, visible=False) as results_accordion:
-                summary_output = gr.Markdown(label="Summary")
+                summary_output = gr.Markdown(label="Answer")
                 details_output = gr.Markdown(label="Details")
     load_button.click(
         fn=process_file,