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 from services.kb_creation import collection, ingest_documents, search_knowledge_base from contextlib import asynccontextmanager import google.generativeai as genai # --- 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.") # Configure the SDK genai.configure(api_key=GEMINI_API_KEY) # Choose the model MODEL_NAME = "gemini-2.5-flash-lite" model = genai.GenerativeModel(MODEL_NAME) # 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", ] app.add_middleware( CORSMiddleware, allow_origins=origins, 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.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) # 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:" 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 = 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} # 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))