Update app.py

app.py

@@ -1,289 +1,139 @@
-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
+import os
+import pandas as pd
+import tempfile
+import nest_asyncio
+from llama_index.llms.openai import OpenAI
+from llama_index.core import VectorStoreIndex
+from llama_parse import LlamaParse
+from llama_index.core.node_parser import MarkdownElementNodeParser
 
-# Ignore warnings for a cleaner interface
-warnings.filterwarnings('ignore')
-# Load environment variables from .env file
-load_dotenv()
+# Apply nest_asyncio to handle async operations in Gradio
+nest_asyncio.apply()
 
-class ExcelAIQuerySystem:
+def get_sheet_names(file):
     """
-    A system to query Excel files using natural language, powered by OpenAI and LangChain.
+    Reads an uploaded Excel file and returns its sheet names.
+    This function is triggered when a file is uploaded.
     """
-    def __init__(self, openai_api_key: str):
-        os.environ["OPENAI_API_KEY"] = openai_api_key
-        self.llm = OpenAI(temperature=0)
-        self.embeddings = OpenAIEmbeddings()
-        self.excel_data = {}
-        self.sheet_descriptions = {}
-        self.vectorstore = None
-        self.logs = []
-
-    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:
-            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 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._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)}")
+    if file is None:
+        # No file, so return an empty, non-interactive dropdown
+        return gr.Dropdown(choices=[], interactive=False)
+    try:
+        # Use pandas to quickly get sheet names without loading all data
+        xls = pd.ExcelFile(file.name)
+        sheet_names = xls.sheet_names
+        # Return an updated, interactive dropdown with the sheet names
+        return gr.Dropdown(choices=sheet_names, value=sheet_names[0], interactive=True)
+    except Exception as e:
+        gr.Warning(f"Could not read Excel file: {e}")
+        return gr.Dropdown(choices=[], interactive=False)
 
-    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:
-            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 _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))}"
-
-    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 process_and_query(openai_api_key, llama_parse_api_key, uploaded_file, selected_sheet, query_text):
+    """
+    Processes a selected sheet from an uploaded Excel file and answers a user's query.
 
-    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())
+    Args:
+        openai_api_key (str): The user's OpenAI API key.
+        llama_parse_api_key (str): The user's LlamaParse API key.
+        uploaded_file (gradio.File): The uploaded Excel file object.
+        selected_sheet (str): The name of the sheet to process.
+        query_text (str): The question to ask about the document.
 
-    def query_data(self, query: str, selected_sheet: str = None) -> Dict[str, Any]:
-        """
-        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.
-        """
-        results = {'query': query, 'relevant_sheets': [], 'sheet_results': {}, 'summary': '', 'insights': []}
-        try:
-            # 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)
-            
-            results['relevant_sheets'] = relevant_sheets
+    Returns:
+        str: The answer to the query or an error message.
+    """
+    if not all([openai_api_key, llama_parse_api_key, uploaded_file, selected_sheet, query_text]):
+        return "Error: Please provide all inputs - both API keys, a file, a selected sheet, and a query."
 
-            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:
-            results['summary'] = f"Error processing query: {str(e)}"
-            return results
+    # Create a temporary file to store the selected sheet's data
+    temp_file = None
+    try:
+        os.environ["OPENAI_API_KEY"] = openai_api_key
+        llm = OpenAI(model="gpt-4o-mini", api_key=openai_api_key)
 
-    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."
+        # Read the selected sheet using pandas
+        df = pd.read_excel(uploaded_file.name, sheet_name=selected_sheet)
         
-        combined_responses = "\n\n".join(
-            f"--- Analysis from Sheet '{name}' ---\n{res['response']}"
-            for name, res in sheet_results.items()
+        # Save the sheet's data to a temporary CSV file for LlamaParse
+        with tempfile.NamedTemporaryFile(mode='w+', delete=False, suffix=".csv", encoding='utf-8') as temp_file:
+            df.to_csv(temp_file.name, index=False)
+            temp_file_path = temp_file.name
+
+        # Initialize LlamaParse
+        parser = LlamaParse(
+            api_key=llama_parse_api_key,
+            result_type="markdown",
+            verbose=True
         )
-        prompt = f"""
-        Based on the following analyses, provide a final, consolidated answer to the query.
-        Original Query: {query}
         
-        {combined_responses}
+        # Load data from the temporary file containing only the selected sheet
+        documents = parser.load_data(temp_file_path)
         
-        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)
+        # Initialize the parser and process the documents
+        node_parser = MarkdownElementNodeParser(llm=llm, num_workers=4)
+        nodes = node_parser.get_nodes_from_documents(documents)
+        base_nodes, objects = node_parser.get_nodes_and_objects(nodes)
         
-        # Get sheet names for the dropdown
-        sheet_names = list(excel_system.excel_data.keys())
-        dropdown_choices = ["Auto-Select based on Query"] + sheet_names
+        # Create the index and query engine
+        recursive_index = VectorStoreIndex(nodes=base_nodes + objects, llm=llm)
+        query_engine = recursive_index.as_query_engine(similarity_top_k=5, llm=llm)
         
-        return (
-            loading_logs, 
-            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_selection, 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.")
-    
-    try:
-        # 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', []))
+        # Execute the query
+        response = query_engine.query(query_text)
         
-        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 ---
+        return str(response)
 
-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.")
+    except Exception as e:
+        return f"An error occurred: {e}"
+    finally:
+        # Clean up the temporary file
+        if temp_file and os.path.exists(temp_file.name):
+            os.unlink(temp_file.name)
 
+# --- Gradio Interface ---
+with gr.Blocks(
+    title="Excel Sheet Q&A",
+    theme=gr.themes.Soft(),
+    css=".gradio-container {background: linear-gradient(to right, #C9D6FF, #E2E2E2)}"
+) as iface:
+    gr.Markdown(
+        """
+        # 📊 Excel Sheet Q&A
+        
+        1.  Enter your API keys.
+        2.  Upload an Excel file.
+        3.  Choose a specific sheet from the dropdown.
+        4.  Ask a question about the data in that sheet.
+        """
+    )
+    
     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=5)
+            openai_key_input = gr.Textbox(label="OpenAI API Key", type="password")
+            llamaparse_key_input = gr.Textbox(label="LlamaParse API Key", type="password")
+            file_uploader = gr.File(label="Upload Excel Databook", file_types=[".xlsx", ".xls"])
+            sheet_selector = gr.Dropdown(label="Choose a Sheet", interactive=False)
         
         with gr.Column(scale=2):
-            gr.Markdown("### 2. Ask a Question")
-            sheet_selector = gr.Dropdown(
-                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 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 ---
+            query_input = gr.Textbox(label="Your Question", placeholder="e.g., What were the total revenues in 2022?", lines=8)
+            submit_button = gr.Button("Ask Question", variant="primary")
+            output_display = gr.Markdown(label="Answer")
     
-    load_button.click(
-        fn=process_file,
-        inputs=[api_key_input, file_input],
-        outputs=[status_output, system_state, sheet_selector, query_input, ask_button, results_accordion]
+    # Define the interactive workflow
+    file_uploader.upload(
+        fn=get_sheet_names,
+        inputs=[file_uploader],
+        outputs=[sheet_selector]
     )
     
-    ask_button.click(
-        fn=generate_response,
-        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
+    submit_button.click(
+        fn=process_and_query,
+        inputs=[openai_key_input, llamaparse_key_input, file_uploader, sheet_selector, query_input],
+        outputs=output_display
     )
 
 if __name__ == "__main__":
-    demo.launch(share=True)
+    iface.launch(share=True)
+