Create main.py

main.py

@@ -0,0 +1,103 @@
+import os
+os.environ["POSTHOG_DISABLED"] = "true"  # Disable PostHog telemetry
+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
+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"],
+    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}"
+
+# --- 5. Endpoints ---
+@app.get("/")
+async def health_check():
+    return {"status": "ok"}
+
+@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:"
+
+        headers = {"Content-Type": "application/json"}
+        payload = {
+            "contents": [
+                {
+                    "parts": [{"text": enhanced_prompt}]
+                }
+            ]
+        }
+
+        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}
+
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=str(e))