Update app.py

app.py

@@ -720,13 +720,14 @@ async def chat_general(
 
 
 # Custom AI Chat
+# Custom AI Chat - FIXED VERSION with Fallback
 @app.post("/chat/custom/{ai_id}", response_model=QuestionResponse)
 async def chat_custom_ai(
         ai_id: str,
         request: QuestionRequest,
         current_user: dict = Depends(get_current_user)
 ):
-    """Chat with custom AI"""
+    """Chat with custom AI using fallback modes"""
     if not lightrag_manager:
         raise HTTPException(status_code=503, detail="LightRAG system not initialized")
 
@@ -734,20 +735,198 @@ async def chat_custom_ai(
         # Generate conversation ID if not provided
         conversation_id = request.conversation_id or str(uuid.uuid4())
 
-        # Query with persistent LightRAG
-        response = await lightrag_manager.query_with_memory(
+        # Query with fallback - SAME AS FIRE SAFETY
+        result = await query_rag_with_fallback(
+            lightrag_manager=lightrag_manager,
             ai_type="custom",
             question=request.question,
             conversation_id=conversation_id,
             user_id=current_user["id"],
+            ai_id=ai_id,  # This is the key difference - pass ai_id
+            preferred_mode=request.mode or "hybrid"
+        )
+
+        await update_message_stats(request.question, result["answer"])
+        return QuestionResponse(
+            answer=result["answer"],
+            mode=result["mode"],
+            status=result["status"],
+            conversation_id=conversation_id
+        )
+
+    except Exception as e:
+        logger.error(f"Custom AI chat error: {e}")
+        raise HTTPException(status_code=500, detail=f"Chat error: {str(e)}")
+
+
+# Alternative: Enhanced fallback system specifically for custom AIs
+async def query_custom_ai_with_fallback(
+        lightrag_manager: PersistentLightRAGManager,
+        ai_id: str,
+        question: str,
+        conversation_id: str,
+        user_id: str,
+        preferred_mode: str = "hybrid"
+) -> Dict[str, Any]:
+    """Query custom AI with automatic fallback to working modes and Cloudflare backup"""
+
+    # Try modes in order of preference
+    fallback_modes = ["hybrid", "local", "global", "naive"]
+
+    # Start with user's preferred mode
+    if preferred_mode in fallback_modes:
+        fallback_modes.remove(preferred_mode)
+        fallback_modes.insert(0, preferred_mode)
+
+    last_error = None
+
+    for mode in fallback_modes:
+        try:
+            logger.info(f"🔍 Trying {mode} mode for custom AI {ai_id}")
+
+            response = await lightrag_manager.query_with_memory(
+                ai_type="custom",
+                question=question,
+                conversation_id=conversation_id,
+                user_id=user_id,
+                ai_id=ai_id,
+                mode=mode
+            )
+
+            # Check if response is valid (not a "Sorry" message)
+            if response and len(response.strip()) > 20 and not response.startswith("Sorry"):
+                logger.info(f"✅ {mode} mode worked for custom AI {ai_id}")
+                return {
+                    "answer": response,
+                    "mode": mode,
+                    "status": "success",
+                    "fallback_used": mode != preferred_mode
+                }
+            else:
+                logger.warning(f"⚠️ {mode} mode returned empty/error response for custom AI {ai_id}")
+                last_error = f"{mode} mode returned: {response[:100]}..."
+
+        except Exception as e:
+            logger.warning(f"⚠️ {mode} mode failed for custom AI {ai_id}: {e}")
+            last_error = str(e)
+            continue
+
+    # If all LightRAG modes fail, fallback to Cloudflare AI
+    try:
+        logger.info(f"🔄 All LightRAG modes failed for custom AI {ai_id}, falling back to Cloudflare AI")
+
+        # Get custom AI details for context
+        async with lightrag_manager.db.pool.acquire() as conn:
+            ai_details = await conn.fetchrow("""
+                SELECT name, description FROM rag_instances 
+                WHERE ai_id = $1 AND ai_type = 'custom' AND status = 'active'
+            """, ai_id)
+
+        # Create context-aware system prompt
+        ai_name = ai_details['name'] if ai_details else "Custom AI"
+        ai_description = ai_details['description'] if ai_details else "a custom AI assistant"
+
+        system_prompt = f"""You are {ai_name}, {ai_description}. 
+
+        Although you don't have access to your specific knowledge base right now, 
+        provide the best general assistance you can. Be helpful, accurate, and 
+        acknowledge that you're operating in a general mode without your specialized knowledge."""
+
+        fallback_response = await lightrag_manager.cloudflare_worker.query(
+            question,
+            system_prompt
+        )
+
+        # Save fallback response to database
+        await lightrag_manager.db.save_conversation_message(
+            conversation_id, "user", question, {
+                "user_id": user_id,
+                "ai_type": "custom",
+                "ai_id": ai_id
+            }
+        )
+        await lightrag_manager.db.save_conversation_message(
+            conversation_id, "assistant", fallback_response, {
+                "mode": "cloudflare_fallback",
+                "ai_type": "custom",
+                "ai_id": ai_id,
+                "user_id": user_id,
+                "error": last_error
+            }
+        )
+
+        return {
+            "answer": fallback_response,
+            "mode": "cloudflare_fallback",
+            "status": "success",
+            "fallback_used": True
+        }
+
+    except Exception as fallback_error:
+        logger.error(f"❌ Even Cloudflare fallback failed for custom AI {ai_id}: {fallback_error}")
+
+        # Last resort - return informative error
+        error_response = f"I'm having trouble accessing my knowledge base right now. Please try again in a moment, or contact support if the issue persists. (Error: {last_error})"
+
+        return {
+            "answer": error_response,
+            "mode": "error",
+            "status": "error",
+            "fallback_used": True
+        }
+
+
+@app.post("/admin/rebuild-custom-ai/{ai_id}")
+async def rebuild_custom_ai(
+        ai_id: str,
+        current_user: dict = Depends(get_current_user)
+):
+    """Force rebuild a corrupted custom AI"""
+
+    if not lightrag_manager:
+        raise HTTPException(status_code=503, detail="LightRAG system not initialized")
+
+    try:
+        success = await lightrag_manager.force_rebuild_custom_ai(ai_id, current_user["id"])
+
+        if success:
+            return {"message": f"Successfully rebuilt custom AI {ai_id}"}
+        else:
+            raise HTTPException(status_code=500, detail="Failed to rebuild custom AI")
+
+    except Exception as e:
+        logger.error(f"Rebuild custom AI error: {e}")
+        raise HTTPException(status_code=500, detail=f"Rebuild error: {str(e)}")
+
+@app.post("/chat/custom/{ai_id}", response_model=QuestionResponse)
+async def chat_custom_ai_enhanced(
+        ai_id: str,
+        request: QuestionRequest,
+        current_user: dict = Depends(get_current_user)
+):
+    """Chat with custom AI using enhanced fallback system"""
+    if not lightrag_manager:
+        raise HTTPException(status_code=503, detail="LightRAG system not initialized")
+
+    try:
+        # Generate conversation ID if not provided
+        conversation_id = request.conversation_id or str(uuid.uuid4())
+
+        # Query with enhanced custom AI fallback
+        result = await query_custom_ai_with_fallback(
+            lightrag_manager=lightrag_manager,
             ai_id=ai_id,
-            mode=request.mode or "hybrid"
+            question=request.question,
+            conversation_id=conversation_id,
+            user_id=current_user["id"],
+            preferred_mode=request.mode or "hybrid"
         )
-        await update_message_stats(request.question, response)
+
+        await update_message_stats(request.question, result["answer"])
         return QuestionResponse(
-            answer=response,
-            mode=request.mode or "hybrid",
-            status="success",
+            answer=result["answer"],
+            mode=result["mode"],
+            status=result["status"],
             conversation_id=conversation_id
         )