Update app.py

app.py

@@ -1,180 +1,277 @@
+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
-import pandas as pd
-from typing import List, Dict, Any
+from dotenv import load_dotenv
 
-# --- LlamaIndex & LangChain Imports ---
-from llama_index.core import VectorStoreIndex, Document, Settings
-from llama_index.llms.openai import OpenAI as LlamaOpenAI
-from llama_index.embeddings.openai import OpenAIEmbedding
-from langchain_experimental.agents.agent_toolkits import create_pandas_dataframe_agent
-from langchain_openai import ChatOpenAI
-from langchain.agents.agent_types import AgentType
+# Ignore warnings for a cleaner interface
+warnings.filterwarnings('ignore')
+# Load environment variables from .env file
+load_dotenv()
 
-class HybridExcelQuerySystem:
+class ExcelAIQuerySystem:
+    """
+    A system to query Excel files using natural language, powered by OpenAI and LangChain.
+    """
     def __init__(self, openai_api_key: str):
         os.environ["OPENAI_API_KEY"] = openai_api_key
-        Settings.llm = LlamaOpenAI(model="gpt-4o")
-        Settings.embed_model = OpenAIEmbedding(model="text-embedding-3-small")
-        self.agent_llm = ChatOpenAI(temperature=0, model="gpt-4o")
-        
-        self.dataframes: Dict[str, pd.DataFrame] = {}
-        self.vector_stores: Dict[str, VectorStoreIndex] = {}
+        self.llm = OpenAI(temperature=0)
+        self.embeddings = OpenAIEmbeddings()
+        self.excel_data = {}
+        self.sheet_descriptions = {}
+        self.vectorstore = None
         self.logs = []
-        self.sheet_names = []
-
-    def _pivot_numerical_data(self, df: pd.DataFrame) -> pd.DataFrame:
-        """
-        --- NEW, MORE ROBUST VERSION ---
-        Intelligently reconstructs the 'Numerical Data' sheet into a clean, tidy format.
-        """
-        # Define expected months to correctly identify the data block
-        months = ['Apr', 'May', 'Jun', 'Jul', 'Aug', 'Sep', 'Oct', 'Nov', 'Dec', 'Jan', 'Feb', 'Mar']
-        
-        # Find the first row that contains at least one month name
-        header_row_index = df[df.isin(months)].dropna(how='all').index[0]
-        header_series = df.iloc[header_row_index]
-        
-        # Filter this row to get only the actual month names
-        actual_month_headers = [h for h in header_series if h in months]
-        
-        # Find the start and end column positions of the month data
-        start_col_pos = header_series.tolist().index(actual_month_headers[0])
-        end_col_pos = header_series.tolist().index(actual_month_headers[-1])
-        
-        # Find the start row of the financial metrics
-        metric_col = df.iloc[:, 0].dropna()
-        start_row_pos = metric_col.index[1] # The metrics start after "Profit & Loss Account"
-
-        # Slice the core data block
-        data = df.iloc[start_row_pos:, start_col_pos:end_col_pos+1]
-        metrics = df.iloc[start_row_pos:, 0]
-        
-        # Create a new, clean DataFrame
-        clean_df = pd.DataFrame(data.values, index=metrics, columns=actual_month_headers)
-        
-        # Transpose so months are rows
-        clean_df = clean_df.T
-        clean_df = clean_df.reset_index().rename(columns={'index': 'Month'})
-        
-        # Clean column names
-        clean_df.columns = clean_df.columns.str.strip().str.replace(r'[^a-zA-Z0-9_%]', '_', regex=True).str.replace('__', '_')
-        
-        # Convert all columns except 'Month' to numeric
-        for col in clean_df.columns:
-            if col != 'Month':
-                if clean_df[col].dtype == 'object':
-                    clean_df[col] = clean_df[col].astype(str).str.replace('%', '', regex=False)
-                clean_df[col] = pd.to_numeric(clean_df[col], errors='coerce')
-        
-        return clean_df
-
 
     def load_excel_file(self, file_path: str) -> str:
+        """Loads and processes an Excel file, generating descriptions and a vector store."""
         self.logs.clear()
         try:
-            xls = pd.ExcelFile(file_path)
-            self.sheet_names = xls.sheet_names
-            self.logs.append(f"✅ Found {len(self.sheet_names)} sheets: {', '.join(self.sheet_names)}")
+            excel_file = pd.ExcelFile(file_path)
+            sheet_names = excel_file.sheet_names
+            self.logs.append(f"✅ Found {len(sheet_names)} sheets: {', '.join(sheet_names)}")
 
-            for sheet_name in self.sheet_names:
-                df = pd.read_excel(file_path, sheet_name=sheet_name, header=None)
-                
-                if sheet_name == "Numerical Data":
-                    agent_df = self._pivot_numerical_data(df.copy())
-                else:
-                    agent_df = self._clean_dataframe_for_agent(df.copy())
-                self.dataframes[sheet_name] = agent_df
-
-                rag_df = self._clean_dataframe_for_rag(df.copy())
-                markdown_text = rag_df.to_markdown(index=False)
-                doc = Document(text=markdown_text, metadata={"source": sheet_name})
-                self.vector_stores[sheet_name] = VectorStoreIndex.from_documents([doc])
-                
-                self.logs.append(f"  - Prepared sheet '{sheet_name}' for both Lookup and Calculation.")
+            for sheet_name in 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)")
+                except Exception as e:
+                    self.logs.append(f"⚠️ Error loading sheet '{sheet_name}': {str(e)}")
+                    continue
             
-            self.logs.append("✅ All sheets are ready.")
+            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: {e}")
+            raise Exception(f"Error loading Excel file: {str(e)}")
 
-    def _clean_dataframe_for_agent(self, df: pd.DataFrame) -> pd.DataFrame:
-        df.columns = [f"Col_{i}" for i in range(len(df.columns))]
-        return df
-
-    def _clean_dataframe_for_rag(self, df: pd.DataFrame) -> pd.DataFrame:
+    def _clean_dataframe(self, df: pd.DataFrame) -> pd.DataFrame:
+        """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:
-            df[col] = df[col].astype(str)
+            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
         return df
 
-    def _classify_query(self, query: str) -> str:
-        prompt = f"""Classify the user's query about an Excel sheet as either "lookup" or "calculation". "lookup": for questions asking for specific data or summaries. "calculation": for questions requiring math, sorting, or filtering. Query: "{query}" Classification:"""
-        response = self.agent_llm.invoke(prompt)
-        classification = response.content.strip().lower()
-        return "calculation" if "calculation" in classification else "lookup"
-
-    def query(self, query: str, selected_sheet: str) -> Dict[str, Any]:
-        if not selected_sheet:
-            return {"answer": "Error: Please select a sheet first.", "tool_used": "None"}
-        classification = self._classify_query(query)
-        if classification == "calculation":
-            return self._execute_agent_query(query, selected_sheet)
-        else:
-            return self._execute_rag_query(query, selected_sheet)
-
-    def _execute_rag_query(self, query: str, sheet_name: str) -> Dict[str, Any]:
+    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:
-            query_engine = self.vector_stores[sheet_name].as_query_engine()
-            response = query_engine.query(query)
-            return {"answer": str(response), "tool_used": "Lookup (RAG Search)"}
-        except Exception as e:
-            return {"answer": f"Error during lookup: {e}", "tool_used": "Lookup (RAG Search)"}
+            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 _execute_agent_query(self, query: str, sheet_name: str) -> Dict[str, Any]:
+    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:
-            df = self.dataframes[sheet_name]
-            agent = create_pandas_dataframe_agent(self.agent_llm, df, agent_type=AgentType.ZERO_SHOT_REACT_DESCRIPTION, verbose=True, allow_dangerous_code=True, max_iterations=15, handle_parsing_errors=True)
-            response = agent.invoke(query)
-            return {"answer": response['output'], "tool_used": "Calculation (Pandas Agent)"}
+            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}
+            
+            results['summary'] = self._generate_summary(query, results['sheet_results'])
+            results['insights'] = self._extract_insights(results['sheet_results'])
+            return results
         except Exception as e:
-            return {"answer": f"Error during calculation: {e}", "tool_used": "Calculation (Pandas Agent)"}
+            results['summary'] = f"Error processing query: {str(e)}"
+            return results
 
-# --- Gradio UI ---
-def process_excel(api_key: str, file_obj: gr.File):
-    if not api_key: raise gr.Error("Please provide your OpenAI API key.")
-    if not file_obj: raise gr.Error("Please upload an Excel file.")
-    
-    system = HybridExcelQuerySystem(openai_api_key=api_key)
-    logs = system.load_excel_file(file_obj.name)
-    sheet_names = system.sheet_names
+    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 ---
+
+def process_file(api_key, file_obj):
+    """Gradio function to load the file and prepare the system."""
+    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)
+        
+        return (
+            loading_logs, 
+            excel_system, 
+            gr.update(visible=True), 
+            gr.update(visible=True),
+            gr.update(visible=True)
+        )
+    except Exception as e:
+        raise gr.Error(f"Failed to process file: {e}")
+
+def generate_response(query, system_state):
+    """Gradio function to handle user queries and display results."""
+    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.")
     
-    return (logs, system, gr.update(choices=sheet_names, value=sheet_names[0] if sheet_names else None, visible=True), gr.update(visible=True))
-
-def user_interaction(question: str, history: List, system_state: HybridExcelQuerySystem, selected_sheet: str):
-    if not system_state: raise gr.Error("Please upload and process a file first.")
-    result = system_state.query(question, selected_sheet)
-    answer = result.get("answer", "No response.")
-    tool_used = result.get("tool_used", "Unknown")
-    full_response = f"{answer}\n\n*Tool Used: {tool_used}*"
-    return full_response
-
-with gr.Blocks(theme=gr.themes.Soft(), title="Hybrid Excel Analyzer") as demo:
-    system_state = gr.State()
-    gr.Markdown("# 🤖 Hybrid Excel Analyzer")
-    gr.Markdown("This app automatically chooses the best AI tool to answer your questions about an Excel file.")
+    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', []))
+        
+        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(), title="Excel AI Query System") 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.")
+
     with gr.Row():
         with gr.Column(scale=1):
             gr.Markdown("### 1. Setup")
-            openai_api_key = gr.Textbox(label="OpenAI API Key", type="password", value=os.getenv("OPENAI_API_KEY", ""))
-            excel_upload = gr.File(label="Upload Excel File", file_types=[".xlsx", ".xls"])
-            process_button = gr.Button("Process File", variant="primary")
-            status_text = gr.Textbox(label="Processing Status", interactive=False, lines=8)
+            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=5)
+        
         with gr.Column(scale=2):
             gr.Markdown("### 2. Ask a Question")
-            with gr.Group(visible=False) as query_ui:
-                sheet_selector = gr.Dropdown(label="Select a Sheet")
-                chat_interface = gr.ChatInterface(fn=user_interaction, additional_inputs=[system_state, sheet_selector], title="Chat with your Excel Data")
-    process_button.click(fn=process_excel, inputs=[openai_api_key, excel_upload], outputs=[status_text, system_state, sheet_selector, query_ui])
+            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
+            )
+            ask_button = gr.Button("Get Answer", variant="primary", visible=False)
+            
+            results_accordion = gr.Accordion("Results", open=False, visible=False)
+            with results_accordion:
+                summary_output = gr.Markdown(label="Summary")
+                details_output = gr.Markdown(label="Details")
+
+    # --- 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]
+    )
+    
+    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
+    )
+
 
 if __name__ == "__main__":
-    demo.launch()
+    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)