Update main.py

main.py

@@ -1,9 +1,17 @@
 import os
+os.environ["POSTHOG_DISABLED"] = "true"  # Disable PostHog telemetry
+import requests
 from fastapi import FastAPI, HTTPException
-from pydantic import BaseModel
-import google.generativeai as genai
 from fastapi.middleware.cors import CORSMiddleware
-from embed import ingest_documents,search_knowledge_base,collection
+from pydantic import BaseModel
+from dotenv import load_dotenv
+#from kb_embed import search_knowledge_base
+from services.kb_creation import collection, ingest_documents, search_knowledge_base
+import logging
+from contextlib import asynccontextmanager
+
+
+logging.basicConfig(level=logging.INFO)
 
 # --- 0. Config ---
 GEMINI_API_KEY = os.getenv("GEMINI_API_KEY")
@@ -18,10 +26,29 @@ genai.configure(api_key=GEMINI_API_KEY)
 MODEL_NAME = "gemini-2.5-flash-lite"
 model = genai.GenerativeModel(MODEL_NAME)
 
-# --- 1. FastAPI ---
-app = FastAPI()
 
-# --- 2. Configure CORS ---
+
+# Load environment variables from the .env file
+load_dotenv()
+
+# --- 1. Initialize FastAPI ---
+#app = FastAPI()
+@asynccontextmanager
+async def lifespan(app: FastAPI):
+    try:
+        folder_path = os.path.join(os.getcwd(), "documents")
+        if collection.count() == 0:
+                print("🔍 KB empty. Running ingestion...")
+                ingest_documents(folder_path)
+        else:
+                print(f"✅ KB already populated with {collection.count()} entries. Skipping ingestion.")
+    except Exception as e:
+        print(f"⚠️ KB ingestion failed: {e}")
+    yield
+
+app = FastAPI(lifespan=lifespan)
+
+# --- Configure CORS ---
 origins = [
     "https://jaita-chatbot-react-frontend-v1.hf.space"
     "https://jaita-chatbot-fastapi-backend.hf.space/chat",
@@ -35,55 +62,51 @@ app.add_middleware(
     allow_headers=["*"],
 )
 
-# --- 3. Schemas ---
+# --- 3. Define the Request Data Structure ---
 class ChatInput(BaseModel):
     user_message: str
 
-@app.on_event("startup")
-def startup_ingest():
-    try:
-        folder_path = os.path.join(os.getcwd(), "documents")
-        if collection.count() == 0:
-            print("🔍 KB empty. Running ingestion...")
-            ingest_documents(folder_path)
-        else:
-            print(f"✅ KB already populated with {collection.count()} entries. Skipping ingestion.")
-
-    except Exception as e:
-        print(f"KB ingestion failed: {e}")
+# --- 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}"
 
-# --- 4. Health check ---
+# --- 5. Endpoints ---
 @app.get("/")
 async def health_check():
     return {"status": "ok"}
 
-# --- 5. Chat endpoint using direct Gemini SDK ---
+#@app.on_event("startup")
+#async def startup():
+#    try:
+#        folder_path = os.path.join(os.getcwd(), "documents")
+#        if collection.count() == 0:
+#                print("🔍 KB empty. Running ingestion...")
+#                ingest_all_documents(folder_path)
+#        else:
+#                print(f"✅ KB already populated with {collection.count()} entries. Skipping ingestion.")
+#    except Exception as e:
+#        print(f"⚠️ KB ingestion failed: {e}")
+
+         
 @app.post("/chat")
 async def chat_with_ai(input_data: ChatInput):
+    """Handle chat interactions using Google Generative AI via requests."""
     try:
+        #folder_path = os.path.join(os.getcwd(), "documents")
+        #print("folder_path",folder_path)
 
-        # 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}
-                
         # 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}")
+        #print(f"kb_results are: {kb_results}")
 
-        # Extract relevant context from search results
+# Extract relevant context from search results
         context = ""
-        relevant_docs = []
+        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.
@@ -98,11 +121,20 @@ 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)
+        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()
+        #print("result",result)
         # Extract Gemini's response
-        bot_response = getattr(response, "text", None) or "No response text."
+        bot_response = result["candidates"][0]["content"]["parts"][0]["text"]
 
         # Include debug info in response
         debug_info = f"Context found: {'Yes' if context else 'No'}"
@@ -111,5 +143,23 @@ Answer:"""
 
         return {"bot_response": bot_response, "debug": debug_info}
 
+        # Make POST request to Gemini API
+        #response = requests.post(GEMINI_URL, json=payload,verify=False)
+        #if(response.status_code==200):
+        #    print("response",response.status_code)
+        #    data = response.json()
+        #    #print("data",data)
+
+        # Extract text from response
+        #bot_response = ""
+        #if "candidates" in data and data["candidates"]:
+        #    parts = data["candidates"][0].get("content", {}).get("parts", [])
+        #    for part in parts:
+        #        if "text" in part:
+        #            bot_response += part["text"]
+#
+        #return {"bot_response": bot_response or "No response text."}
+
+
     except Exception as e:
         raise HTTPException(status_code=500, detail=str(e))