Update main.py

main.py

@@ -49,34 +49,35 @@ async def health_check():
 
 @app.post("/chat")
 async def chat_with_gemini(input_data: ChatInput):
-    try:
-        # Retrieve relevant documents from knowledge base
-
-        kb_results = search_knowledge_base(input_data.user_message, top_k=10)
-        logging.info(f"kb_results are: {kb_results}")
-
-        # Extract relevant context from search results
-        context = ""
-        if kb_results and kb_results.get('documents'):
-            # Limit context to avoid token limits - take top 2 most relevant
-            relevant_docs = kb_results['documents'][0][:2]
-            context = "\n\n".join(relevant_docs)
-          
-        # Construct enhanced prompt with context
-        if context:
-            enhanced_prompt = f"""Use the following knowledge base context to answer the user's question accurately.
-If the context contains relevant information, base your answer on it.
-If the context doesn't help, say you're raising a ticket and provide a ticket number.
-
-Knowledge Base Context:
-{context}
-
-User Question: {input_data.user_message}
-
-Answer:"""
-        else:
-            enhanced_prompt = f"User Question: {input_data.user_message}\n\nAnswer:"
-
+ 
+    # 1. Search Knowledge Base
+    kb_results = search_knowledge_base(input_data.user_message, top_k=10)
+    logging.info(kb_results)
+ 
+    context = ""
+    relevant_docs = []
+ 
+    # 2. Extract relevant KB docs
+    if kb_results and kb_results.get("documents"):
+        relevant_docs = kb_results["documents"][0][:2]
+        context = "\n\n".join(relevant_docs)
+ 
+    # 3. If KB contains direct answer → RETURN it (No LLM call)
+    if context:
+        kb_answer = f"From knowledge base:\n\n{context}"
+        return {
+            "bot_response": kb_answer,
+            "debug_info": f"Context found: YES, docs used: {len(relevant_docs)}"
+        }
+ 
+    # 4. If KB empty → fallback to Gemini
+    enhanced_prompt = (
+        f"User question: {input_data.user_message}\n\n"
+        "No relevant KB found. You must raise a ticket.\n"
+        "Say: 'I'm raising a ticket. Ticket# 12345'."
+    )
+ 
+    # 5. Call Gemini
         headers = {"Content-Type": "application/json"}
         payload = {
             "contents": [
@@ -85,19 +86,16 @@ Answer:"""
                 }
             ]
         }
-
-        response = requests.post(GEMINI_URL, headers=headers, json=payload, verify=False)  # SSL disabled for testing
-        result = response.json()
-
-        # Extract Gemini's response
-        bot_response = result["candidates"][0]["content"]["parts"][0]["text"]
-
-        # Include debug info in response
-        debug_info = f"Context found: {'Yes' if context else 'No'}"
-        if context:
-            debug_info += f" (Top {len(relevant_docs)} documents used)"
-
-        return {"bot_response": bot_response, "debug": debug_info}
+ 
+    response = requests.post(GEMINI_URL, headers=headers, json=payload)
+    result = response.json()
+ 
+    bot_response = result["candidates"][0]["content"]["parts"][0]["text"]
+ 
+    return {
+        "bot_response": bot_response,
+        "debug_info": "Context found: NO"
+    }
 
     except Exception as e:
         raise HTTPException(status_code=500, detail=str(e))