Update main.py

main.py

@@ -1,14 +1,15 @@
 import os
 import json
 import chromadb
-from fastapi import FastAPI, HTTPException, Depends
+from fastapi import FastAPI, HTTPException, Depends, Query
 from pydantic import BaseModel, Field
 from typing import List
 import firebase_admin
 from firebase_admin import credentials, firestore
-from llm_handler import initialize_llm, get_rag_response
-import llm_handler
+
 from encoder import SentenceEncoder
+from llm_handler import initialize_llm, get_rag_response, llm_handler # Make sure llm_handler is imported
+from populate_chroma import populate_vector_db # For the setup endpoint
 
 # --- Pydantic Models ---
 class UserProfile(BaseModel):
@@ -19,8 +20,6 @@ class UserProfile(BaseModel):
 class SearchQuery(BaseModel):
     query: str = Field(..., example="marketing internship in mumbai")
 
-# --- SCHEMA CHANGED HERE ---
-# Reverted to use 'id' and 'skills'
 class InternshipData(BaseModel):
     id: str = Field(..., example="int_021")
     title: str
@@ -43,35 +42,40 @@ class ChatMessage(BaseModel):
 class ChatResponse(BaseModel):
     response: str
 
-# --- FastAPI App & Firebase Initialization ---
+# --- FastAPI App Initialization ---
 app = FastAPI(
-    title="Internship Recommendation API",
-    description="An API using Firestore for metadata, and ChromaDB for vector search.",
-    version="2.1.0"
+    title="Internship Recommendation & Chatbot API",
+    description="An API using Firestore for metadata, ChromaDB for vector search, and an LLM for chat.",
+    version="3.0.0"
 )
 
-# Initialize Firebase ONCE at startup
+# --- Firebase Initialization ---
+db = None
 try:
-    if 'FIREBASE_CREDS_JSON' in os.environ:
-        creds_dict = json.loads(os.environ.get('FIREBASE_CREDS_JSON'))
+    firebase_creds = os.getenv("FIREBASE_CREDS_JSON")
+    if firebase_creds:
+        creds_dict = json.loads(firebase_creds)
         cred = credentials.Certificate(creds_dict)
+        if not firebase_admin._apps:
+            firebase_admin.initialize_app(cred)
+        db = firestore.client()
+        print("✅ Firebase initialized with Hugging Face secret.")
     else:
+        # Fallback for local development if the secret isn't set
         cred = credentials.Certificate('serviceAccountKey.json')
-    
-    firebase_admin.initialize_app(cred)
-    db = firestore.client()
-    print("✅ Firebase connection initialized.")
+        if not firebase_admin._apps:
+            firebase_admin.initialize_app(cred)
+        db = firestore.client()
+        print("✅ Firebase initialized with local key file.")
 except Exception as e:
     print(f"❌ Could not initialize Firebase. Error: {e}")
-    db = None
 
-# Dependency to provide the db client
 def get_db():
     if db is None:
         raise HTTPException(status_code=503, detail="Firestore connection not available.")
     return db
 
-# --- Global Variables for Model and ChromaDB ---
+# --- Global Variables ---
 encoder = None
 chroma_collection = None
 
@@ -82,13 +86,21 @@ def load_model_and_data():
     print("🚀 Loading sentence encoder model...")
     encoder = SentenceEncoder()
     
-    client = chromadb.PersistentClient(path="/content/chroma_db")
+    # --- THIS IS THE FIX ---
+    # Point ChromaDB to the correct writable persistent storage path on Hugging Face
+    chroma_db_path = "/data/chroma_db"
+    
+    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 the LLM
     initialize_llm()
 
 # --- API Endpoints ---
@@ -96,90 +108,63 @@ def load_model_and_data():
 def read_root():
     return {"message": "Welcome to the Internship Recommendation API!"}
 
+@app.post("/setup")
+def run_initial_setup(secret_key: str = Query(...)):
+    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()
+        print("--- SETUP COMPLETE ---")
+        return {"status": "Setup completed successfully."}
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=f"An error occurred: {str(e)}")
+
+
 @app.post("/add-internship", response_model=StatusResponse)
 def add_internship(internship: InternshipData, 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.")
-
-    # --- SCHEMA CHANGED HERE ---
-    # Using internship.id
     doc_ref = db_client.collection('internships').document(internship.id)
     if doc_ref.get().exists:
         raise HTTPException(status_code=400, detail="Internship ID already exists.")
-
-    # Save to Firestore
     doc_ref.set(internship.dict())
-
-    # --- SCHEMA CHANGED HERE ---
-    # Using internship.skills
     text_to_encode = f"{internship.title}. {internship.description}. Skills: {', '.join(internship.skills)}"
     embedding = encoder.encode([text_to_encode])[0].tolist()
-
-    # --- CRITICAL FIX RE-APPLIED HERE ---
-    # Prepare metadata for ChromaDB, converting skills list to a JSON string
     metadata_for_chroma = internship.dict()
     metadata_for_chroma['skills'] = json.dumps(metadata_for_chroma['skills'])
-
-    chroma_collection.add(
-        # --- SCHEMA CHANGED HERE ---
-        # Using internship.id
-        ids=[internship.id],
-        embeddings=[embedding],
-        metadatas=[metadata_for_chroma]
-    )
-
+    chroma_collection.add(ids=[internship.id], embeddings=[embedding], metadatas=[metadata_for_chroma])
     print(f"✅ Added internship to Firestore and ChromaDB: {internship.id}")
-    # --- SCHEMA CHANGED HERE ---
     return {"status": "success", "internship_id": internship.id}
 
 @app.post("/profile-recommendations", response_model=RecommendationResponse)
 def get_profile_recommendations(profile: UserProfile):
     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)}. Interests: {', '.join(profile.interests)}"
     query_embedding = encoder.encode([query_text])[0].tolist()
-    
-    results = chroma_collection.query(
-        query_embeddings=[query_embedding],
-        n_results=3
-    )
-    
+    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]
-        })
-
+        recommendations.append({"internship_id": internship_id, "score": 1 - distances[i]})
     return {"recommendations": recommendations}
 
 @app.post("/search", response_model=RecommendationResponse)
 def search_internships(search: SearchQuery):
     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
-    )
-    
+    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]
-        })
-        
+        recommendations.append({"internship_id": internship_id, "score": 1 - distances[i]})
     return {"recommendations": recommendations}
+
 @app.post("/chat", response_model=ChatResponse)
 def chat_with_bot(message: ChatMessage):
     response = get_rag_response(message.query)