Update app.py

app.py

@@ -1,217 +1,222 @@
-import os
-import gradio as gr
-import tempfile
-from pathlib import Path
-import base64
-from PIL import Image
-import io
-import time
-
-# Import our components
-from models.llm_setup import setup_llm
-from indexes.csv_index_builder import EnhancedCSVReader
-from indexes.index_manager import CSVIndexManager
-from indexes.query_engine import CSVQueryEngine
-from tools.data_tools import PandasDataTools
-from tools.visualization import VisualizationTools
-from tools.export import ExportTools
-
-# Setup temporary directory for uploaded files
-UPLOAD_DIR = Path(tempfile.mkdtemp())
-EXPORT_DIR = Path(tempfile.mkdtemp())
-
-class CSVChatApp:
-    """Main application class for CSV chatbot."""
-    
-    def __init__(self):
-        """Initialize the application components."""
-        # Initialize the language model
-        self.llm = setup_llm()
-        
-        # Initialize the index manager
-        self.index_manager = CSVIndexManager()
-        
-        # Initialize tools
-        self.data_tools = PandasDataTools(str(UPLOAD_DIR))
-        self.viz_tools = VisualizationTools(str(UPLOAD_DIR))
-        self.export_tools = ExportTools(str(EXPORT_DIR))
-        
-        # Initialize query engine with tools
-        self.query_engine = self._setup_query_engine()
-        
-        # Track conversation history
-        self.chat_history = []
-        self.uploaded_files = []
-    
-    def _setup_query_engine(self):
-        """Set up the query engine with tools."""
-        # Get all tools
-        tools = (
-            self.data_tools.get_tools() + 
-            self.viz_tools.get_tools() + 
-            self.export_tools.get_tools()
-        )
-        
-        # Create query engine with tools
-        query_engine = CSVQueryEngine(self.index_manager, self.llm)
-        
-        return query_engine
-    
-    def handle_file_upload(self, files):
-        """Process uploaded CSV files."""
-        file_info = []
-        
-        for file in files:
-            if file is None:
-                continue
-                
-            # Get file path
-            file_path = Path(file.name)
-            
-            # Only process CSV files
-            if not file_path.suffix.lower() == '.csv':
-                continue
-                
-            # Copy to upload directory
-            dest_path = UPLOAD_DIR / file_path.name
-            with open(dest_path, 'wb') as f:
-                f.write(file_path.read_bytes())
-            
-            # Create index for this file
-            try:
-                self.index_manager.create_index(str(dest_path))
-                file_info.append(f"✅ Indexed: {file_path.name}")
-                self.uploaded_files.append(str(dest_path))
-            except Exception as e:
-                file_info.append(f"❌ Failed to index {file_path.name}: {str(e)}")
-        
-        # Return information about processed files
-        if file_info:
-            return "\n".join(file_info)
-        else:
-            return "No CSV files were uploaded."
-    
-    def process_query(self, query, history):
-        """Process a user query and generate a response."""
-        if not self.uploaded_files:
-            return "Please upload CSV files before asking questions."
-        
-        # Add user message to history
-        self.chat_history.append({"role": "user", "content": query})
-        
-        # Process the query
-        try:
-            response = self.query_engine.query(query)
-            answer = response["answer"]
-            
-            # Check if response contains an image
-            if isinstance(answer, dict) and "image" in answer:
-                # Handle image in response
-                img_data = answer["image"]
-                img = Image.open(io.BytesIO(base64.b64decode(img_data)))
-                img_path = EXPORT_DIR / f"viz_{int(time.time())}.png"
-                img.save(img_path)
-                
-                # Update answer to include image path
-                text_response = answer.get("text", "Generated visualization")
-                answer = (text_response, str(img_path))
-            
-            # Add assistant message to history
-            self.chat_history.append({"role": "assistant", "content": answer})
-            
-            return answer
-            
-        except Exception as e:
-            error_msg = f"Error processing query: {str(e)}"
-            self.chat_history.append({"role": "assistant", "content": error_msg})
-            return error_msg
-    
-    def export_conversation(self):
-        """Export the conversation as a report."""
-        if not self.chat_history:
-            return "No conversation to export."
-        
-        # Extract content for report
-        title = "CSV Chat Conversation Report"
-        content = ""
-        images = []
-        
-        for msg in self.chat_history:
-            role = msg["role"]
-            content_text = msg["content"]
-            
-            # Handle content that might contain images
-            if isinstance(content_text, tuple) and len(content_text) == 2:
-                text, img_path = content_text
-                content += f"\n\n{'User' if role == 'user' else 'Assistant'}: {text}"
-                
-                # Add image to report
-                try:
-                    with open(img_path, "rb") as img_file:
-                        img_data = base64.b64encode(img_file.read()).decode('utf-8')
-                        images.append(img_data)
-                except Exception:
-                    pass
-            else:
-                content += f"\n\n{'User' if role == 'user' else 'Assistant'}: {content_text}"
-        
-        # Generate report
-        result = self.export_tools.generate_report(title, content, images)
-        
-        if result["success"]:
-            return f"Report exported to: {result['report_path']}"
-        else:
-            return "Failed to export report."
-
-# Create the Gradio interface
-def create_interface():
-    """Create the Gradio web interface."""
-    app = CSVChatApp()
-    
-    with gr.Blocks(title="CSV Chat Assistant") as interface:
-        gr.Markdown("# CSV Chat Assistant")
-        gr.Markdown("Upload CSV files and ask questions in natural language.")
-        
-        with gr.Row():
-            with gr.Column(scale=1):
-                file_upload = gr.File(
-                    label="Upload CSV Files",
-                    file_count="multiple",
-                    type="file"
-                )
-                upload_button = gr.Button("Process Files")
-                file_status = gr.Textbox(label="File Status")
-                
-                export_button = gr.Button("Export Conversation")
-                export_status = gr.Textbox(label="Export Status")
-            
-            with gr.Column(scale=2):
-                chatbot = gr.Chatbot(label="Conversation")
-                msg = gr.Textbox(label="Your Question")
-                submit_button = gr.Button("Submit")
-        
-        # Set up event handlers
-        upload_button.click(
-            fn=app.handle_file_upload,
-            inputs=[file_upload],
-            outputs=[file_status]
-        )
-        
-        submit_button.click(
-            fn=app.process_query,
-            inputs=[msg, chatbot],
-            outputs=[chatbot]
-        )
-        
-        export_button.click(
-            fn=app.export_conversation,
-            inputs=[],
-            outputs=[export_status]
-        )
-    
-    return interface
-
-# Launch the app
-if __name__ == "__main__":
-    interface = create_interface()
-    interface.launch()
+import os
+import gradio as gr
+import tempfile
+from pathlib import Path
+import base64
+from PIL import Image
+import io
+import time
+
+
+current_dir = os.path.dirname(os.path.abspath(__file__))
+sys.path.append(current_dir)
+
+
+# Import our components
+from models.llm_setup import setup_llm
+from indexes.csv_index_builder import EnhancedCSVReader
+from indexes.index_manager import CSVIndexManager
+from indexes.query_engine import CSVQueryEngine
+from tools.data_tools import PandasDataTools
+from tools.visualization import VisualizationTools
+from tools.export import ExportTools
+
+# Setup temporary directory for uploaded files
+UPLOAD_DIR = Path(tempfile.mkdtemp())
+EXPORT_DIR = Path(tempfile.mkdtemp())
+
+class CSVChatApp:
+    """Main application class for CSV chatbot."""
+    
+    def __init__(self):
+        """Initialize the application components."""
+        # Initialize the language model
+        self.llm = setup_llm()
+        
+        # Initialize the index manager
+        self.index_manager = CSVIndexManager()
+        
+        # Initialize tools
+        self.data_tools = PandasDataTools(str(UPLOAD_DIR))
+        self.viz_tools = VisualizationTools(str(UPLOAD_DIR))
+        self.export_tools = ExportTools(str(EXPORT_DIR))
+        
+        # Initialize query engine with tools
+        self.query_engine = self._setup_query_engine()
+        
+        # Track conversation history
+        self.chat_history = []
+        self.uploaded_files = []
+    
+    def _setup_query_engine(self):
+        """Set up the query engine with tools."""
+        # Get all tools
+        tools = (
+            self.data_tools.get_tools() + 
+            self.viz_tools.get_tools() + 
+            self.export_tools.get_tools()
+        )
+        
+        # Create query engine with tools
+        query_engine = CSVQueryEngine(self.index_manager, self.llm)
+        
+        return query_engine
+    
+    def handle_file_upload(self, files):
+        """Process uploaded CSV files."""
+        file_info = []
+        
+        for file in files:
+            if file is None:
+                continue
+                
+            # Get file path
+            file_path = Path(file.name)
+            
+            # Only process CSV files
+            if not file_path.suffix.lower() == '.csv':
+                continue
+                
+            # Copy to upload directory
+            dest_path = UPLOAD_DIR / file_path.name
+            with open(dest_path, 'wb') as f:
+                f.write(file_path.read_bytes())
+            
+            # Create index for this file
+            try:
+                self.index_manager.create_index(str(dest_path))
+                file_info.append(f"✅ Indexed: {file_path.name}")
+                self.uploaded_files.append(str(dest_path))
+            except Exception as e:
+                file_info.append(f"❌ Failed to index {file_path.name}: {str(e)}")
+        
+        # Return information about processed files
+        if file_info:
+            return "\n".join(file_info)
+        else:
+            return "No CSV files were uploaded."
+    
+    def process_query(self, query, history):
+        """Process a user query and generate a response."""
+        if not self.uploaded_files:
+            return "Please upload CSV files before asking questions."
+        
+        # Add user message to history
+        self.chat_history.append({"role": "user", "content": query})
+        
+        # Process the query
+        try:
+            response = self.query_engine.query(query)
+            answer = response["answer"]
+            
+            # Check if response contains an image
+            if isinstance(answer, dict) and "image" in answer:
+                # Handle image in response
+                img_data = answer["image"]
+                img = Image.open(io.BytesIO(base64.b64decode(img_data)))
+                img_path = EXPORT_DIR / f"viz_{int(time.time())}.png"
+                img.save(img_path)
+                
+                # Update answer to include image path
+                text_response = answer.get("text", "Generated visualization")
+                answer = (text_response, str(img_path))
+            
+            # Add assistant message to history
+            self.chat_history.append({"role": "assistant", "content": answer})
+            
+            return answer
+            
+        except Exception as e:
+            error_msg = f"Error processing query: {str(e)}"
+            self.chat_history.append({"role": "assistant", "content": error_msg})
+            return error_msg
+    
+    def export_conversation(self):
+        """Export the conversation as a report."""
+        if not self.chat_history:
+            return "No conversation to export."
+        
+        # Extract content for report
+        title = "CSV Chat Conversation Report"
+        content = ""
+        images = []
+        
+        for msg in self.chat_history:
+            role = msg["role"]
+            content_text = msg["content"]
+            
+            # Handle content that might contain images
+            if isinstance(content_text, tuple) and len(content_text) == 2:
+                text, img_path = content_text
+                content += f"\n\n{'User' if role == 'user' else 'Assistant'}: {text}"
+                
+                # Add image to report
+                try:
+                    with open(img_path, "rb") as img_file:
+                        img_data = base64.b64encode(img_file.read()).decode('utf-8')
+                        images.append(img_data)
+                except Exception:
+                    pass
+            else:
+                content += f"\n\n{'User' if role == 'user' else 'Assistant'}: {content_text}"
+        
+        # Generate report
+        result = self.export_tools.generate_report(title, content, images)
+        
+        if result["success"]:
+            return f"Report exported to: {result['report_path']}"
+        else:
+            return "Failed to export report."
+
+# Create the Gradio interface
+def create_interface():
+    """Create the Gradio web interface."""
+    app = CSVChatApp()
+    
+    with gr.Blocks(title="CSV Chat Assistant") as interface:
+        gr.Markdown("# CSV Chat Assistant")
+        gr.Markdown("Upload CSV files and ask questions in natural language.")
+        
+        with gr.Row():
+            with gr.Column(scale=1):
+                file_upload = gr.File(
+                    label="Upload CSV Files",
+                    file_count="multiple",
+                    type="file"
+                )
+                upload_button = gr.Button("Process Files")
+                file_status = gr.Textbox(label="File Status")
+                
+                export_button = gr.Button("Export Conversation")
+                export_status = gr.Textbox(label="Export Status")
+            
+            with gr.Column(scale=2):
+                chatbot = gr.Chatbot(label="Conversation")
+                msg = gr.Textbox(label="Your Question")
+                submit_button = gr.Button("Submit")
+        
+        # Set up event handlers
+        upload_button.click(
+            fn=app.handle_file_upload,
+            inputs=[file_upload],
+            outputs=[file_status]
+        )
+        
+        submit_button.click(
+            fn=app.process_query,
+            inputs=[msg, chatbot],
+            outputs=[chatbot]
+        )
+        
+        export_button.click(
+            fn=app.export_conversation,
+            inputs=[],
+            outputs=[export_status]
+        )
+    
+    return interface
+
+# Launch the app
+if __name__ == "__main__":
+    interface = create_interface()
+    interface.launch()