Update main.py

main.py

@@ -7,20 +7,17 @@ from typing import List
 import firebase_admin
 from firebase_admin import credentials, firestore
 
+# --- Local Imports ---
 from encoder import SentenceEncoder
-from populate_chroma import populate_vector_db
+from populate_chroma import populate_vector_db # For the setup endpoint
+from llm_handler import initialize_llm, get_rag_response
+import llm_handler # To pass global objects
 
 # --------------------------------------------------------------------
-# Cache setup (store HF models in /data for persistence on Hugging Face)
+# Cache & Root Path Setup
 # --------------------------------------------------------------------
 os.environ["HF_HOME"] = "/data/cache"
 os.environ["SENTENCE_TRANSFORMERS_HOME"] = "/data/cache"
-
-# --------------------------------------------------------------------
-# Determine root_path dynamically
-#   Locally: root_path = ""
-#   On Hugging Face Spaces: root_path = "/username/space-name"
-# --------------------------------------------------------------------
 root_path = os.getenv("HF_SPACE_ROOT_PATH", "")
 
 # --------------------------------------------------------------------
@@ -42,24 +39,31 @@ class InternshipData(BaseModel):
     createdAt: str
     stipend: int = None
 
+# --- NEW: More descriptive response models ---
+class InternshipRecommendation(InternshipData):
+    score: float
+
 class RecommendationResponse(BaseModel):
-    recommendations: List[dict]
+    recommendations: List[InternshipRecommendation]
 
 class StatusResponse(BaseModel):
     status: str
     internship_id: str
 
+class ChatMessage(BaseModel):
+    query: str
+
+class ChatResponse(BaseModel):
+    response: str
+
 # --------------------------------------------------------------------
 # FastAPI App
 # --------------------------------------------------------------------
 app = FastAPI(
-    title="Internship Recommendation API",
-    description="An API using Firestore for metadata, and ChromaDB for vector search.",
-    version="2.2.0",
-    docs_url="/docs",            # Swagger UI
-    redoc_url="/redoc",          # ReDoc
-    openapi_url="/openapi.json", # OpenAPI schema
-    root_path=root_path          # ✅ Fix for Hugging Face Spaces subpath issue
+    title="Internship Recommendation & Chatbot API",
+    description="An API using Firestore for metadata, ChromaDB for vector search, and an LLM chatbot.",
+    version="3.0.0",
+    root_path=root_path
 )
 
 # --------------------------------------------------------------------
@@ -97,16 +101,20 @@ def load_model_and_data():
     print("🚀 Loading sentence encoder model...")
     encoder = SentenceEncoder()
     
-    # Point ChromaDB to the persistent /data storage path
     chroma_db_path = "/data/chroma_db"
-    
     try:
         client = chromadb.PersistentClient(path=chroma_db_path)
         chroma_collection = client.get_or_create_collection(name="internships")
         print("✅ ChromaDB client initialized and collection is ready.")
         print(f"   - Internships in DB: {chroma_collection.count()}")
+
+        # Pass the loaded models to the llm_handler module
+        llm_handler.encoder = encoder
+        llm_handler.chroma_collection = chroma_collection
+        initialize_llm() # Initialize the LLM for the chatbot
+
     except Exception as e:
-        print(f"❌ Error initializing ChromaDB: {e}")
+        print(f"❌ Error initializing ChromaDB or LLM: {e}")
         raise
 
 # --------------------------------------------------------------------
@@ -116,20 +124,11 @@ def load_model_and_data():
 def read_root():
     return {"message": "Welcome to the Internship Recommendation API!"}
 
-@app.get("/healthz")
-def health_check():
-    return {"status": "ok"}
-
 @app.post("/setup")
 def run_initial_setup(secret_key: str = Query(..., example="your_secret_password")):
-    """
-    A secret endpoint to run the initial database setup.
-    This should only be run once after deployment.
-    """
     correct_key = os.getenv("SETUP_SECRET_KEY")
     if not correct_key or secret_key != correct_key:
         raise HTTPException(status_code=403, detail="Invalid secret key.")
-    
     try:
         print("--- RUNNING DATABASE POPULATION SCRIPT ---")
         populate_vector_db()
@@ -155,35 +154,51 @@ def add_internship(internship: InternshipData, db_client: firestore.Client = Dep
     return {"status": "success", "internship_id": internship.id}
 
 @app.post("/profile-recommendations", response_model=RecommendationResponse)
-def get_profile_recommendations(profile: UserProfile):
+def get_profile_recommendations(profile: UserProfile, db_client: firestore.Client = Depends(get_db)):
     if chroma_collection is None or encoder is None:
         raise HTTPException(status_code=503, detail="Server is not ready.")
     
     query_text = f"Skills: {', '.join(profile.skills)}. Sectors: {', '.join(profile.sectors)}"
     query_embedding = encoder.encode([query_text])[0].tolist()
+    
     results = chroma_collection.query(query_embeddings=[query_embedding], n_results=3)
     
     recommendations = []
     ids = results.get('ids', [[]])[0]
     distances = results.get('distances', [[]])[0]
-    
+
     for i, internship_id in enumerate(ids):
-        recommendations.append({
-            "internship_id": internship_id,
-            "score": 1 - distances[i]
-        })
-    
+        doc_ref = db_client.collection('internships').document(internship_id).get()
+        if doc_ref.exists:
+            internship_data = doc_ref.to_dict()
+            internship_data['score'] = 1 - distances[i]
+            recommendations.append(internship_data)
+
     return {"recommendations": recommendations}
 
 @app.post("/search", response_model=RecommendationResponse)
-def search_internships(search: SearchQuery):
+def search_internships(search: SearchQuery, db_client: firestore.Client = Depends(get_db)):
     if chroma_collection is None or encoder is None:
         raise HTTPException(status_code=503, detail="Server is not ready.")
+        
     query_embedding = encoder.encode([search.query])[0].tolist()
+    
     results = chroma_collection.query(query_embeddings=[query_embedding], n_results=3)
+    
     recommendations = []
     ids = results.get('ids', [[]])[0]
     distances = results.get('distances', [[]])[0]
+
     for i, internship_id in enumerate(ids):
-        recommendations.append({"internship_id": internship_id, "score": 1 - distances[i]})
-    return {"recommendations": recommendations}
+        doc_ref = db_client.collection('internships').document(internship_id).get()
+        if doc_ref.exists:
+            internship_data = doc_ref.to_dict()
+            internship_data['score'] = 1 - distances[i]
+            recommendations.append(internship_data)
+            
+    return {"recommendations": recommendations}
+
+@app.post("/chat", response_model=ChatResponse)
+def chat_with_bot(message: ChatMessage):
+    response = get_rag_response(message.query)
+    return {"response": response}