Update chatbot.py

chatbot.py

@@ -2,23 +2,32 @@ import streamlit as st
 import pandas as pd
 import os
 from datetime import datetime
-import google.generativeai as genai
 
-# Configure Gemini
-GEMINI_API_KEY = os.environ.get('GEMINI_API_KEY')
-genai.configure(api_key=GEMINI_API_KEY)
-model = genai.GenerativeModel('gemini-pro')
+try:
+    import google.generativeai as genai
+    GEMINI_AVAILABLE = True
+except ImportError:
+    GEMINI_AVAILABLE = False
 
-class GeminiDataChatbot:
+class ChatbotManager:
     def __init__(self):
+        if GEMINI_AVAILABLE and 'GEMINI_API_KEY' in os.environ:
+            genai.configure(api_key=os.environ['GEMINI_API_KEY'])
+            self.model = genai.GenerativeModel('gemini-pro')
+        else:
+            self.model = None
+        
         if 'uploaded_df' not in st.session_state:
             st.session_state.uploaded_df = None
         if 'chat_history' not in st.session_state:
             st.session_state.chat_history = []
     
-    def render_interface(self):
-        st.title("📊 Data Analysis Chatbot")
-        st.write("Upload your CSV file and ask questions about your data")
+    def render_chat_interface(self):
+        """Render the main chat interface"""
+        st.header("📊 Data Analysis Chatbot")
+        
+        if not GEMINI_AVAILABLE:
+            st.warning("Gemini API not available - running in limited mode")
         
         # File upload section
         uploaded_file = st.file_uploader("Choose a CSV file", type="csv")
@@ -31,6 +40,7 @@ class GeminiDataChatbot:
             self._render_chat_window()
     
     def _process_uploaded_file(self, uploaded_file):
+        """Process the uploaded CSV file"""
         try:
             df = pd.read_csv(uploaded_file)
             st.session_state.uploaded_df = df
@@ -39,16 +49,10 @@ class GeminiDataChatbot:
             with st.expander("View Data Preview"):
                 st.dataframe(df.head())
             
-            # Initial analysis prompt
-            initial_prompt = f"""
-            I have uploaded a dataset with {len(df)} rows and {len(df.columns)} columns.
-            Columns: {', '.join(df.columns)}.
-            First give a very brief (2-3 sentence) overview of what this data might contain.
-            Then suggest 3 specific questions I could ask about this data.
-            """
-            
-            with st.spinner("Analyzing your data..."):
-                response = self._generate_gemini_response(initial_prompt, df)
+            # Initial analysis
+            if self.model:
+                initial_prompt = f"Briefly describe this dataset with {len(df)} rows and {len(df.columns)} columns."
+                response = self._generate_response(initial_prompt)
                 st.session_state.chat_history.append({
                     "role": "assistant", 
                     "content": response
@@ -58,6 +62,7 @@ class GeminiDataChatbot:
             st.error(f"Error processing file: {str(e)}")
     
     def _render_chat_window(self):
+        """Render the chat conversation window"""
         st.subheader("Chat About Your Data")
         
         # Display chat history
@@ -67,55 +72,44 @@ class GeminiDataChatbot:
         
         # User input
         if prompt := st.chat_input("Ask about your data..."):
-            # Add user message to chat history
-            st.session_state.chat_history.append({"role": "user", "content": prompt})
-            
-            # Display user message
-            with st.chat_message("user"):
-                st.markdown(prompt)
-            
-            # Generate and display assistant response
-            with st.chat_message("assistant"):
-                with st.spinner("Thinking..."):
-                    response = self._generate_gemini_response(prompt, st.session_state.uploaded_df)
-                    st.markdown(response)
-            
-            # Add assistant response to chat history
-            st.session_state.chat_history.append({"role": "assistant", "content": response})
+            self._handle_user_input(prompt)
     
-    def _generate_gemini_response(self, prompt: str, df: pd.DataFrame) -> str:
-        """Generate response using Gemini API with data context"""
-        try:
-            # Create data summary for context
-            data_summary = f"""
-            Data Summary:
-            - Shape: {df.shape}
-            - Columns: {', '.join(df.columns)}
-            - First 5 rows:
-            {df.head().to_markdown()}
-            """
-            
-            # Create prompt with context
-            full_prompt = f"""
-            You are a data analysis assistant. The user has uploaded a dataset with the following characteristics:
-            {data_summary}
-            
-            User Question: {prompt}
-            
-            Provide a detailed response answering their question about the data. If appropriate, include:
-            - Relevant statistics
-            - Potential visualizations that would help
-            - Any data quality issues to consider
-            - Business insights if applicable
-            """
-            
-            response = model.generate_content(full_prompt)
-            return response.text
+    def _handle_user_input(self, prompt):
+        """Handle user input and generate response"""
+        # Add user message to chat history
+        st.session_state.chat_history.append({"role": "user", "content": prompt})
         
-        except Exception as e:
-            return f"Sorry, I encountered an error processing your request: {str(e)}"
-
-# Initialize and run the chatbot
-if __name__ == "__main__":
-    chatbot = GeminiDataChatbot()
-    chatbot.render_interface()
+        # Display user message
+        with st.chat_message("user"):
+            st.markdown(prompt)
+        
+        # Generate and display assistant response
+        with st.chat_message("assistant"):
+            with st.spinner("Thinking..."):
+                response = self._generate_response(prompt)
+                st.markdown(response)
+        
+        # Add assistant response to chat history
+        st.session_state.chat_history.append({"role": "assistant", "content": response})
+    
+    def _generate_response(self, prompt: str) -> str:
+        """Generate response using available backend"""
+        df = st.session_state.uploaded_df
+        
+        if self.model:
+            # Use Gemini if available
+            try:
+                data_summary = f"Data: {len(df)} rows, columns: {', '.join(df.columns)}"
+                full_prompt = f"{data_summary}\n\nUser question: {prompt}"
+                response = self.model.generate_content(full_prompt)
+                return response.text
+            except Exception as e:
+                return f"Gemini error: {str(e)}"
+        else:
+            # Fallback basic analysis
+            if "summary" in prompt.lower():
+                return f"Basic summary:\n{df.describe().to_markdown()}"
+            elif "columns" in prompt.lower():
+                return f"Columns: {', '.join(df.columns)}"
+            else:
+                return "I can provide basic info about columns and summary statistics."