Update app.py

app.py

@@ -8,88 +8,147 @@ import pickle
 import os
 
 # Step 2: Configure API Key from Hugging Face Secrets
+GOOGLE_API_KEY = None
+api_configured = False
+
 try:
     GOOGLE_API_KEY = os.environ.get('GOOGLE_API_KEY')
     if GOOGLE_API_KEY is None:
-        raise ValueError("GOOGLE_API_KEY not found in environment variables.")
-    genai.configure(api_key=GOOGLE_API_KEY)
-    print("API Key configured successfully.")
+        print("WARNING: GOOGLE_API_KEY not found in environment variables.")
+        print("Please add it to your Hugging Face Space secrets.")
+    else:
+        genai.configure(api_key=GOOGLE_API_KEY)
+        api_configured = True
+        print("✅ API Key configured successfully.")
 except Exception as e:
-    print(f"ERROR: Could not configure API key. Please ensure 'GOOGLE_API_KEY' is set in your Hugging Face Space secrets. Details: {e}")
+    print(f"ERROR: Could not configure API key. Details: {e}")
+    api_configured = False
 
 # Step 3: Define Data Path
 DATA_PATH = "data"
 vector_store_file = os.path.join(DATA_PATH, "vector_store.index")
 data_file = os.path.join(DATA_PATH, "data.pkl")
 
-# Step 4: Load Models and Pre-processed Data
+# Step 4: Initialize variables
 vector_store_data = None
+embedding_model = None
+llm = None
+texts = []
+sources = []
+index = None
+
+# Step 4: Load Models and Pre-processed Data
 if os.path.exists(vector_store_file) and os.path.exists(data_file):
     try:
         print(f"Loading pre-processed data from the '{DATA_PATH}' directory...")
+        
+        # Load FAISS index
         index = faiss.read_index(vector_store_file)
+        
+        # Load texts and sources
         with open(data_file, "rb") as f:
             stored_data = pickle.load(f)
-            texts = stored_data["texts"]
-            sources = stored_data["sources"]
-        print("✅ Data loaded successfully.")
-
-        vector_store_data = (index, texts, sources)
-
-        print("Loading AI and embedding models...")
+            texts = stored_data.get("texts", [])
+            sources = stored_data.get("sources", [])
+        
+        print(f"✅ Data loaded successfully. Found {len(texts)} text chunks.")
+        
+        # Load embedding model
+        print("Loading embedding model...")
         embedding_model = SentenceTransformer('BAAI/bge-large-en-v1.5')
-        llm = genai.GenerativeModel('gemini-1.5-flash-latest')
-        print("✅ Models loaded successfully.")
-
+        print("✅ Embedding model loaded successfully.")
+        
+        # Load LLM if API is configured
+        if api_configured:
+            print("Loading Gemini model...")
+            llm = genai.GenerativeModel('gemini-1.5-flash-latest')
+            print("✅ Gemini model loaded successfully.")
+            vector_store_data = (index, texts, sources)
+        else:
+            print("⚠️ Gemini model not loaded due to missing API key.")
+            
     except Exception as e:
         print(f"❌ ERROR: An error occurred during data or model loading: {e}")
+        import traceback
+        traceback.print_exc()
         vector_store_data = None
 else:
     print(f"❌ ERROR: Pre-processed data not found in the '{DATA_PATH}' directory.")
     print(f"Please make sure '{vector_store_file}' and '{data_file}' exist.")
-
+    if not os.path.exists(DATA_PATH):
+        print(f"The '{DATA_PATH}' directory does not exist.")
+    else:
+        print(f"Contents of '{DATA_PATH}' directory: {os.listdir(DATA_PATH) if os.path.exists(DATA_PATH) else 'N/A'}")
 
 # Step 5: RAG and Chat Functions (with Updated System Prompt)
-def get_relevant_context(query, index, top_k=5):
-    query_embedding = embedding_model.encode([query])
-    distances, indices = index.search(query_embedding, top_k)
-    context = []
-    for i in indices[0]:
-        if i < len(texts):
-            context.append({"text": texts[i], "source": sources[i]})
-    return context
+def get_relevant_context(query, index, texts, sources, top_k=5):
+    """Retrieve relevant context from the vector store."""
+    if embedding_model is None:
+        return []
+    
+    try:
+        query_embedding = embedding_model.encode([query])
+        distances, indices = index.search(query_embedding, top_k)
+        context = []
+        for i in indices[0]:
+            if 0 <= i < len(texts):  # Ensure index is valid
+                context.append({"text": texts[i], "source": sources[i]})
+        return context
+    except Exception as e:
+        print(f"Error in get_relevant_context: {e}")
+        return []
 
 def chat_with_rag(message, history, vector_store_data):
-    index, texts, sources = vector_store_data
-    relevant_context = get_relevant_context(message, index)
-    context_str = "\n\n".join([f"Source: {c['source']}\nContent: {c['text']}" for c in relevant_context])
-
-    # --- NEW, ROBUST SYSTEM PROMPT ---
-    prompt = f"""
-    You are a friendly and engaging science communicator for the Halassa Lab at MIT.
-    Your goal is to explain the lab's complex computational neuroscience research to a general audience that has little to no scientific background.
-
-    Follow these rules strictly:
-    1.  **Simplify, Don't Dumb Down:** Break down complex topics into simple, easy-to-understand concepts. Use analogies and real-world examples (e.g., "think of the thalamus as a busy switchboard operator for the brain").
-    2.  **Be Engaging and Accessible:** Write in a clear, conversational, and friendly tone. Avoid jargon at all costs. If you must use a technical term, explain it immediately in simple terms.
-    3.  **Synthesize and Explain:** Do not just copy-paste from the provided research papers. Read the relevant context and formulate a comprehensive, well-written answer in your own words.
-    4.  **Cite Your Sources:** At the end of your response, if you used information from the provided papers, you MUST include a "Sources:" list. Use the format [filename - Page X]. This adds credibility.
-    5.  **Stay Focused:** Only answer questions related to the Halassa Lab's work based on the provided context. If the context doesn't contain the answer, state that the information isn't available in the provided documents.
-
-    Context from the Halassa Lab's papers:
-    ---
-    {context_str}
-    ---
-
-    User Question: {message}
-
-    Your Friendly Explanation:
-    """
+    """Generate response using RAG."""
+    if vector_store_data is None:
+        return "❌ The AI system is not properly initialized. Please check the configuration."
+    
+    if not api_configured or llm is None:
+        return "❌ The Google API key is not configured. Please add GOOGLE_API_KEY to your Hugging Face Space secrets."
+    
     try:
+        index, texts, sources = vector_store_data
+        
+        # Get relevant context
+        relevant_context = get_relevant_context(message, index, texts, sources)
+        
+        if not relevant_context:
+            return "I couldn't find relevant information in the Halassa Lab's papers to answer your question. Could you try rephrasing or asking about a different aspect of the lab's research?"
+        
+        context_str = "\n\n".join([f"Source: {c['source']}\nContent: {c['text']}" for c in relevant_context])
+
+        # System prompt
+        prompt = f"""
+You are a friendly and engaging science communicator for the Halassa Lab at MIT.
+Your goal is to explain the lab's complex computational neuroscience research to a general audience that has little to no scientific background.
+
+Follow these rules strictly:
+1. **Simplify, Don't Dumb Down:** Break down complex topics into simple, easy-to-understand concepts. Use analogies and real-world examples (e.g., "think of the thalamus as a busy switchboard operator for the brain").
+2. **Be Engaging and Accessible:** Write in a clear, conversational, and friendly tone. Avoid jargon at all costs. If you must use a technical term, explain it immediately in simple terms.
+3. **Synthesize and Explain:** Do not just copy-paste from the provided research papers. Read the relevant context and formulate a comprehensive, well-written answer in your own words.
+4. **Cite Your Sources:** At the end of your response, if you used information from the provided papers, you MUST include a "Sources:" list. Use the format [filename - Page X]. This adds credibility.
+5. **Stay Focused:** Only answer questions related to the Halassa Lab's work based on the provided context. If the context doesn't contain the answer, state that the information isn't available in the provided documents.
+
+Context from the Halassa Lab's papers:
+---
+{context_str}
+---
+
+User Question: {message}
+
+Your Friendly Explanation:
+"""
+        
+        # Generate response
         response = llm.generate_content(prompt)
         return response.text
+        
     except Exception as e:
-        return f"An error occurred with the AI model: {str(e)}"
+        error_msg = f"An error occurred while generating the response: {str(e)}"
+        print(error_msg)
+        import traceback
+        traceback.print_exc()
+        return f"❌ {error_msg}"
 
 # Step 6: Polished Gradio User Interface
 # Define a professional, clean theme for a science lab
@@ -117,44 +176,119 @@ example_questions = [
     "Explain the role of the thalamic reticular nucleus like I'm a high school student.",
 ]
 
-with gr.Blocks(theme=theme, title="Halassa Lab AI Explainer") as demo:
-    with gr.Column():
-        # --- Header Section ---
-        gr.Image("https://d39w22sdwnt1s2.cloudfront.net/wp-content/uploads/2023/10/Halassa-M-2023-Option-2-1200x675.jpg", show_label=False, show_download_button=False, container=False)
-        gr.Markdown("# Hello! ")
-        gr.Markdown("Ask me anything about the work we do in the Halassa Lab!")
+# Create the Gradio interface
+def create_demo():
+    """Create the Gradio demo interface."""
+    with gr.Blocks(theme=theme, title="Halassa Lab AI Explainer") as demo:
+        with gr.Column():
+            # --- Header Section ---
+            try:
+                gr.Image(
+                    "https://d39w22sdwnt1s2.cloudfront.net/wp-content/uploads/2023/10/Halassa-M-2023-Option-2-1200x675.jpg", 
+                    show_label=False, 
+                    show_download_button=False, 
+                    container=False
+                )
+            except:
+                # If image fails to load, continue without it
+                pass
+                
+            gr.Markdown("# Hello! 👋")
+            gr.Markdown("Ask me anything about the work we do in the Halassa Lab!")
+            
+            # Show status
+            if vector_store_data is not None and api_configured:
+                gr.Markdown("✅ **System Status:** Ready to answer your questions!")
+            elif vector_store_data is not None and not api_configured:
+                gr.Markdown("⚠️ **System Status:** Data loaded but Google API key is missing. Please configure it in Space secrets.")
+            else:
+                gr.Markdown("❌ **System Status:** System not properly initialized. Check logs for details.")
 
-        # --- Chatbot Interface ---
-        chatbot_ui = gr.Chatbot(label="Conversation", height=500, layout="panel", bubble_full_width=False)
-        message_box = gr.Textbox(label="Ask your question here...", lines=3, placeholder="e.g., How does the brain filter out distractions?")
-        
-        with gr.Row():
-            submit_button = gr.Button("Submit", variant="primary", scale=2)
-            clear_button = gr.ClearButton(components=[chatbot_ui, message_box], value="Clear Conversation", scale=1)
-
-        # --- Example Questions Section ---
-        gr.Examples(
-            examples=example_questions,
-            inputs=message_box,
-            label="Click an example to get started:"
-        )
+            # --- Chatbot Interface ---
+            chatbot_ui = gr.Chatbot(
+                label="Conversation", 
+                height=500, 
+                layout="panel", 
+                bubble_full_width=False
+            )
+            message_box = gr.Textbox(
+                label="Ask your question here...", 
+                lines=3, 
+                placeholder="e.g., How does the brain filter out distractions?"
+            )
+            
+            with gr.Row():
+                submit_button = gr.Button("Submit", variant="primary", scale=2)
+                clear_button = gr.ClearButton(
+                    components=[chatbot_ui, message_box], 
+                    value="Clear Conversation", 
+                    scale=1
+                )
 
-    # --- Logic to make the chatbot respond ---
-    def respond(message, history):
-        response_text = chat_with_rag(message, history, vector_store_data)
-        history.append((message, response_text))
-        return "", history
+            # --- Example Questions Section ---
+            gr.Examples(
+                examples=example_questions,
+                inputs=message_box,
+                label="Click an example to get started:"
+            )
 
-    submit_button.click(respond, inputs=[message_box, chatbot_ui], outputs=[message_box, chatbot_ui])
-    message_box.submit(respond, inputs=[message_box, chatbot_ui], outputs=[message_box, chatbot_ui])
+        # --- Logic to make the chatbot respond ---
+        def respond(message, history):
+            """Handle user messages and generate responses."""
+            if not message.strip():
+                return "", history
+            
+            response_text = chat_with_rag(message, history, vector_store_data)
+            history.append((message, response_text))
+            return "", history
 
+        # Connect event handlers
+        submit_button.click(
+            respond, 
+            inputs=[message_box, chatbot_ui], 
+            outputs=[message_box, chatbot_ui]
+        )
+        message_box.submit(
+            respond, 
+            inputs=[message_box, chatbot_ui], 
+            outputs=[message_box, chatbot_ui]
+        )
+    
+    return demo
 
 # Step 7: Launch the app
-if vector_store_data:
-    demo.launch()
-else:
-    # Display a clear error message in the UI if data loading failed
-    with gr.Blocks() as demo:
-        gr.Markdown("## ⚠️ Application Error")
-        gr.Markdown("Could not load the necessary data and models. Please check the Hugging Face Space logs for details. This may be due to missing data files or an incorrect API key.")
+if __name__ == "__main__":
+    if vector_store_data is not None:
+        print("Launching application with full functionality...")
+        demo = create_demo()
+    else:
+        print("Launching application in error mode...")
+        # Display a clear error message in the UI if data loading failed
+        with gr.Blocks(theme=theme) as demo:
+            gr.Markdown("## ⚠️ Application Error")
+            gr.Markdown("""
+            Could not load the necessary data and models. 
+            
+            **Possible issues:**
+            1. Missing data files (`vector_store.index` and `data.pkl` in the `data/` directory)
+            2. Missing or incorrect Google API key in Hugging Face Space secrets
+            3. Model loading errors
+            
+            Please check the Hugging Face Space logs for detailed error messages.
+            """)
+            
+            # Show current status
+            gr.Markdown("### Debug Information:")
+            status_text = []
+            status_text.append(f"- Data directory exists: {os.path.exists(DATA_PATH)}")
+            status_text.append(f"- Vector store file exists: {os.path.exists(vector_store_file)}")
+            status_text.append(f"- Data pickle file exists: {os.path.exists(data_file)}")
+            status_text.append(f"- Google API key configured: {api_configured}")
+            
+            if os.path.exists(DATA_PATH):
+                files = os.listdir(DATA_PATH)
+                status_text.append(f"- Files in data directory: {files if files else 'Empty directory'}")
+            
+            gr.Markdown("\n".join(status_text))
+    
     demo.launch()