Update main.py

main.py

@@ -1,82 +1,62 @@
 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
+import google.generativeai as genai
 
-logging.basicConfig(level=logging.INFO)
+# --- 0. Config ---
+GEMINI_API_KEY = os.getenv("GEMINI_API_KEY")
+if not GEMINI_API_KEY:
+    raise RuntimeError("GEMINI_API_KEY is not set in environment.")
+
+# Configure the SDK
+genai.configure(api_key=GEMINI_API_KEY)
 
-# Load environment variables from the .env file
-load_dotenv()
+# Choose the model
+MODEL_NAME = "gemini-2.5-flash-lite"
+model = genai.GenerativeModel(MODEL_NAME)
 
-# --- 1. Initialize FastAPI ---
+# --- 1. FastAPI ---
 app = FastAPI()
 
 # --- 2. Configure CORS ---
-#origins = [
-#    "http://localhost:5173",
-#    "http://localhost:3000",
-#]
+origins = [
+    "https://jaita-chatbot-react-frontend-v1.hf.space"
+    "https://jaita-chatbot-fastapi-backend.hf.space/chat",
+]
 
 app.add_middleware(
     CORSMiddleware,
-    allow_origins=["https://jaita-chatbot-fastapi-backend.hf.space/chat","https://jaita-chatbot-react-frontend-v1.hf.space"],
+    allow_origins=origins,
     allow_credentials=True,
     allow_methods=["*"],
     allow_headers=["*"],
 )
 
-# --- 3. Define the Request Data Structure ---
+# --- 3. Schemas ---
 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 ---
+# --- 4. Health check ---
 @app.get("/")
 async def health_check():
     return {"status": "ok"}
 
-
+# --- 5. Chat endpoint using direct Gemini SDK ---
 @app.post("/chat")
 async def chat_with_ai(input_data: ChatInput):
-    """Handle chat interactions using Google Generative AI via requests."""
+    """
+    Handle chat interactions using the official Google Generative AI SDK.
+    """
     try:
-        # Gemini expects contents array with role and parts
-        payload = {
-            "contents": [
-                {
-                    "role": "user",
-                    "parts": [{"text": input_data.user_message}],
-                }
-            ],
-            #"generationConfig": {
-            #    "temperature": 0.7,
-            #    "maxOutputTokens": 512
-            #}
-        }
-
-        # Make POST request to Gemini API
-        response = requests.post(GEMINI_URL, json=payload,verify=False)
-        response.raise_for_status()
-        data = response.json()
-
-        # 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."}
 
+        # 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}
 
     except Exception as e:
-        raise HTTPException(status_code=500, detail=str(e))
+        raise HTTPException(status_code=500, detail=str(e))