Rename main2.py to main.py

main2.py → main.py

@@ -3,9 +3,11 @@ from fastapi import FastAPI, HTTPException
 from pydantic import BaseModel
 import google.generativeai as genai
 from fastapi.middleware.cors import CORSMiddleware
+from embed.py import ingest_documents,search_knowledge_base
 
 # --- 0. Config ---
 GEMINI_API_KEY = os.getenv("GEMINI_API_KEY")
+print("GEMINI API KEY",GEMINI_API_KEY)
 if not GEMINI_API_KEY:
     raise RuntimeError("GEMINI_API_KEY is not set in environment.")
 
@@ -37,6 +39,13 @@ app.add_middleware(
 class ChatInput(BaseModel):
     user_message: str
 
+@app.on_event("startup")
+def startup_ingest():
+    try:
+        ingest_documents()
+    except Exception as e:
+        print(f"KB ingestion failed: {e}")
+
 # --- 4. Health check ---
 @app.get("/")
 async def health_check():
@@ -45,19 +54,56 @@ async def health_check():
 # --- 5. Chat endpoint using direct Gemini SDK ---
 @app.post("/chat")
 async def chat_with_ai(input_data: ChatInput):
-    """
-    Handle chat interactions using the official Google Generative AI SDK.
-    """
     try:
 
         # Call Gemini directly via SDK
-        resp = model.generate_content(
-            input_data.user_message,
-        )
-        print("resp",resp)
-        bot_response = getattr(resp, "text", None) or "No response text."
-        print("bot_response",bot_response)
-        return {"bot_response": bot_response}
+        #resp = model.generate_content(
+        #    input_data.user_message,
+        #)
+        #print("resp",resp)
+        #bot_response = getattr(resp, "text", None) or "No response text."
+        #print("bot_response",bot_response)
+        #return {"bot_response": bot_response}
+                
+        # Retrieve relevant documents from knowledge base
+
+        kb_results = search_knowledge_base(input_data.user_message, top_k=10)
+        print(f"kb_results are: {kb_results}")
+
+        # Extract relevant context from search results
+        context = ""
+        relevant_docs = []
+        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:"
+
+        response = model.generate_content(enhanced_prompt)
+        print("response",response)
+        # Extract Gemini's response
+        bot_response = getattr(response, "text", None) or "No response 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}
 
     except Exception as e:
         raise HTTPException(status_code=500, detail=str(e))