Update app.py

app.py

@@ -313,7 +313,7 @@ class ProductionRAGSystem:
         return content[:200] + "..."
     
     def generate_answer(self, query, search_results):
-        """Generate both AI and extracted answers"""
+        """Generate both AI and extracted answers with better error handling"""
         if not search_results:
             return {
                 'ai_answer': "No information found in documents.",
@@ -330,7 +330,7 @@ class ProductionRAGSystem:
         # Always generate extracted answer
         extracted_answer = self.extract_direct_answer(query, best_result['content'])
         
-        # Try AI answer if API key available
+        # Try AI answer with better error handling
         ai_answer = None
         openrouter_key = os.environ.get("OPENROUTER_API_KEY")
         
@@ -339,11 +339,14 @@ class ProductionRAGSystem:
             prompt = f"Answer briefly: {query}\n\nContext: {context}\n\nAnswer (1 sentence):"
             
             try:
+                st.write("DEBUG: Calling OpenRouter API...")
                 response = requests.post(
                     "https://openrouter.ai/api/v1/chat/completions",
                     headers={
                         "Authorization": f"Bearer {openrouter_key}",
-                        "Content-Type": "application/json"
+                        "Content-Type": "application/json",
+                        "HTTP-Referer": "https://huggingface.co/spaces",
+                        "X-Title": "RAG Chatbot"
                     },
                     json={
                         "model": "openai/gpt-3.5-turbo",
@@ -354,10 +357,19 @@ class ProductionRAGSystem:
                     timeout=10
                 )
                 
+                st.write(f"DEBUG: API Status Code: {response.status_code}")
+                
                 if response.status_code == 200:
                     ai_answer = response.json()['choices'][0]['message']['content'].strip()
+                    st.write("DEBUG: AI answer received successfully")
+                else:
+                    error_text = response.text
+                    st.error(f"API Error {response.status_code}: {error_text}")
+                    
             except Exception as e:
-                st.warning(f"AI API error: {e}")
+                st.error(f"API Exception: {str(e)}")
+        else:
+            st.warning("No OpenRouter API key found in environment variables")
         
         return {
             'ai_answer': ai_answer,
@@ -367,7 +379,41 @@ class ProductionRAGSystem:
             'has_both': ai_answer is not None
         }
 
-# ================= UTILITY FUNCTIONS =================
+def get_general_ai_response(query):
+    """Get AI response for general questions not related to documents"""
+    openrouter_key = os.environ.get("OPENROUTER_API_KEY")
+    
+    if not openrouter_key:
+        return "I can only answer questions about your uploaded documents. Please add an OpenRouter API key for general conversations."
+    
+    try:
+        response = requests.post(
+            "https://openrouter.ai/api/v1/chat/completions",
+            headers={
+                "Authorization": f"Bearer {openrouter_key}",
+                "Content-Type": "application/json",
+                "HTTP-Referer": "https://huggingface.co/spaces",
+                "X-Title": "RAG Chatbot"
+            },
+            json={
+                "model": "openai/gpt-3.5-turbo",
+                "messages": [
+                    {"role": "system", "content": "You are a helpful AI assistant. Be concise and friendly."},
+                    {"role": "user", "content": query}
+                ],
+                "max_tokens": 150,
+                "temperature": 0.7
+            },
+            timeout=15
+        )
+        
+        if response.status_code == 200:
+            return response.json()['choices'][0]['message']['content'].strip()
+        else:
+            return f"Sorry, I encountered an error (Status: {response.status_code}). Please try again."
+            
+    except Exception as e:
+        return f"Sorry, I encountered an error: {str(e)}"
 
 def get_user_id():
     """Get unique ID for this user session"""
@@ -527,12 +573,35 @@ with st.sidebar:
     # Settings
     st.header("⚙️ Settings")
     
-    # API Status
+    # API Status with better checking
     openrouter_key = os.environ.get("OPENROUTER_API_KEY")
     if openrouter_key:
-        st.success("🟢 AI API Connected")
+        st.success("🟢 OpenRouter API Connected")
+        # Quick API test
+        if st.button("Test API Connection", use_container_width=True):
+            try:
+                test_response = requests.post(
+                    "https://openrouter.ai/api/v1/chat/completions",
+                    headers={
+                        "Authorization": f"Bearer {openrouter_key}",
+                        "Content-Type": "application/json"
+                    },
+                    json={
+                        "model": "openai/gpt-3.5-turbo",
+                        "messages": [{"role": "user", "content": "test"}],
+                        "max_tokens": 5
+                    },
+                    timeout=5
+                )
+                if test_response.status_code == 200:
+                    st.success("✅ API working correctly!")
+                else:
+                    st.error(f"❌ API Error: {test_response.status_code}")
+            except Exception as e:
+                st.error(f"❌ API Test Failed: {str(e)}")
     else:
-        st.warning("⚠️ No AI API Key (using extracted answers only)")
+        st.error("❌ No OpenRouter API Key")
+        st.info("Add OPENROUTER_API_KEY in Hugging Face Space settings → Variables and secrets")
     
     # RAG Settings
     use_ai_enhancement = st.checkbox("Use AI Enhancement", value=bool(openrouter_key))
@@ -602,11 +671,12 @@ if prompt := st.chat_input("Ask questions about your documents..."):
     # Get RAG response
     with st.chat_message("assistant"):
         if rag_system and rag_system.model and rag_system.get_collection_count() > 0:
-            # Search documents
+            # Search documents first
             search_results = rag_system.search(prompt, n_results=3)
             
-            if search_results:
-                # Generate answer
+            # Check if we found relevant documents (confidence > 0.05)
+            if search_results and search_results[0]['similarity'] > 0.05:
+                # Generate document-based answer
                 result = rag_system.generate_answer(prompt, search_results)
                 
                 # Display AI answer or extracted answer
@@ -615,7 +685,7 @@ if prompt := st.chat_input("Ask questions about your documents..."):
                     st.markdown(f"🤖 **AI Answer:** {answer_text}")
                 else:
                     answer_text = result['extracted_answer']
-                    st.markdown(f"📄 **Answer:** {answer_text}")
+                    st.markdown(f"📄 **Document Answer:** {answer_text}")
                 
                 # Show RAG info
                 if show_sources and result['sources']:
@@ -640,25 +710,31 @@ if prompt := st.chat_input("Ask questions about your documents..."):
                 }
                 
             else:
-                # No relevant documents found
-                no_info_msg = "I couldn't find relevant information in your documents. Try rephrasing your question or check if the information exists in your uploaded documents."
-                st.markdown(no_info_msg)
+                # No relevant documents found - use general AI
+                st.info("No relevant documents found. Switching to general AI mode...")
+                general_response = get_general_ai_response(prompt)
+                st.markdown(f"💬 **General AI:** {general_response}")
                 
                 assistant_message = {
                     "role": "assistant", 
-                    "content": no_info_msg,
-                    "rag_info": {"sources": [], "confidence": 0}
+                    "content": general_response,
+                    "rag_info": {"sources": [], "confidence": 0, "mode": "general"}
                 }
         
         else:
-            # RAG system not ready
-            error_msg = "Document system not ready. Please upload and index documents first."
-            st.error(error_msg)
+            # RAG system not ready - use general AI
+            if rag_system and rag_system.get_collection_count() == 0:
+                st.warning("No documents indexed. Using general AI mode...")
+            else:
+                st.error("RAG system not ready. Using general AI mode...")
+            
+            general_response = get_general_ai_response(prompt)
+            st.markdown(f"💬 **General AI:** {general_response}")
             
             assistant_message = {
                 "role": "assistant", 
-                "content": error_msg,
-                "rag_info": {"sources": [], "confidence": 0}
+                "content": general_response,
+                "rag_info": {"sources": [], "confidence": 0, "mode": "general"}
             }
     
     # Add assistant message to history