Update main.py

main.py

@@ -1,101 +1,104 @@
+
+# main.py
 import os
-os.environ["POSTHOG_DISABLED"] = "true"  # Disable PostHog telemetry
+os.environ["POSTHOG_DISABLED"] = "true"
 import requests
 from fastapi import FastAPI, HTTPException
 from fastapi.middleware.cors import CORSMiddleware
 from pydantic import BaseModel
 from dotenv import load_dotenv
-from kb_embed import search_knowledge_base
+from kb_embed import search_knowledge_base, ingest_documents, collection, DOCS_DIR
 import logging
 
-
-
 logging.basicConfig(level=logging.INFO)
-
-
-
-# Load environment variables from the .env file
 load_dotenv()
 
-# --- 1. Initialize FastAPI ---
 app = FastAPI()
 
-# --- 2. Configure CORS ---
-#origins = [
-#    "http://localhost:5173",
-#    "http://localhost:3000",
-#]
-
 app.add_middleware(
     CORSMiddleware,
-    allow_origins=["https://jaita-chatbot-react-frontend-v1.hf.space","https://jaita-chatbot-fastapi-backend.hf.space/chat"],
+    allow_origins=[
+        "https://jaita-chatbot-react-frontend-v1.hf.space",  # frontend space origin
+        "https://jaita-chatbot-fastapi-backend.hf.space",     # backend space origin
+    ],
     allow_credentials=True,
     allow_methods=["*"],
     allow_headers=["*"],
 )
 
-# --- 3. Define the Request Data Structure ---
 class ChatInput(BaseModel):
     user_message: str
 
-# --- 4. Gemini API Setup ---
 GEMINI_API_KEY = os.getenv("GEMINI_API_KEY")
-GEMINI_URL = f"https://generativelanguage.googleapis.com/v1beta/models/gemini-2.5-flash-lite:generateContent?key={GEMINI_API_KEY}"
+GEMINI_URL = (
+    f"https://generativelanguage.googleapis.com/v1beta/models/"
+    f"gemini-2.5-flash-lite:generateContent?key={GEMINI_API_KEY}"
+)
+
+@app.on_event("startup")
+def startup_ingest():
+    try:
+        # Make sure DOCS_DIR exists and has .docx files, then ingest
+        if DOCS_DIR.exists():
+            logging.info(f"Starting KB ingestion from: {DOCS_DIR}")
+            ingest_documents(str(DOCS_DIR))
+        else:
+            logging.warning(f"Docs directory not found: {DOCS_DIR}")
+        logging.info(f"Chroma collection count after startup: {collection.count()}")
+    except Exception as e:
+        logging.exception(f"KB ingestion failed: {e}")
 
-# --- 5. Endpoints ---
 @app.get("/")
 async def health_check():
-    return {"status": "ok"}
+    return {
+        "status": "ok",
+        "kb_count": collection.count(),
+        "docs_dir_exists": DOCS_DIR.exists()
+    }
 
 @app.post("/chat")
 async def chat_with_gemini(input_data: ChatInput):
- 
     # 1. Search Knowledge Base
     kb_results = search_knowledge_base(input_data.user_message, top_k=10)
-    logging.info(kb_results)
- 
+    logging.info(f"KB query results keys: {list(kb_results.keys()) if kb_results else 'None'}")
+
     context = ""
     relevant_docs = []
- 
-    # 2. Extract relevant KB docs
-    if kb_results and kb_results.get("documents"):
-        relevant_docs = kb_results["documents"][0][:2]
+
+    # 2. Extract relevant KB docs if present
+    if kb_results and kb_results.get("documents") and len(kb_results["documents"]) > 0:
+        # kb_results["documents"] is a list of lists (one per query)
+        first_query_docs = kb_results["documents"][0]
+        relevant_docs = first_query_docs[:2]
         context = "\n\n".join(relevant_docs)
- 
+
     # 3. If KB contains direct answer → RETURN it (No LLM call)
-    if context:
+    if context.strip():
         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)}"
+            "debug_info": f"Context found: YES, docs used: {len(relevant_docs)}, kb_count: {collection.count()}"
         }
- 
+
     # 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": [
-            {
-                "parts": [{"text": enhanced_prompt}]
-            }
-        ]
-    }
- 
-    response = requests.post(GEMINI_URL, headers=headers, json=payload)
-    result = response.json()
- 
-    bot_response = result["candidates"][0]["content"]["parts"][0]["text"]
- 
+    payload = {"contents": [{"parts": [{"text": enhanced_prompt}]}]}
+
+    try:
+        response = requests.post(GEMINI_URL, headers=headers, json=payload)
+        result = response.json()
+        bot_response = result["candidates"][0]["content"]["parts"][0]["text"]
+    except Exception as e:
+        logging.exception(f"Gemini call failed: {e}")
+        raise HTTPException(status_code=500, detail="LLM call failed")
+
     return {
         "bot_response": bot_response,
-        "debug_info": "Context found: NO"
+        "debug_info": f"Context found: NO, kb_count: {collection.count()}"
     }
-
-    #except Exception as e:
-    #    raise HTTPException(status_code=500, detail=str(e))