Update main.py

main.py

@@ -8,15 +8,27 @@ import firebase_admin
 from firebase_admin import credentials, firestore
 
 from encoder import SentenceEncoder
-from llm_handler import initialize_llm, get_rag_response
-import llm_handler
-from populate_chroma import populate_vector_db # For the setup endpoint
-
-# --- Pydantic Models ---
+from populate_chroma import populate_vector_db
+
+# --------------------------------------------------------------------
+# Cache setup (store HF models in /data for persistence on Hugging Face)
+# --------------------------------------------------------------------
+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", "")
+
+# --------------------------------------------------------------------
+# Pydantic Models
+# --------------------------------------------------------------------
 class UserProfile(BaseModel):
-    user_id: str
     skills: List[str] = Field(..., example=["python", "data analysis"])
-    interests: List[str] = Field(..., example=["machine learning", "web development"])
+    sectors: List[str] = Field(..., example=["machine learning", "web development"])
 
 class SearchQuery(BaseModel):
     query: str = Field(..., example="marketing internship in mumbai")
@@ -37,20 +49,22 @@ class StatusResponse(BaseModel):
     status: str
     internship_id: str
 
-class ChatMessage(BaseModel):
-    query: str
-
-class ChatResponse(BaseModel):
-    response: str
-
-# --- FastAPI App Initialization ---
+# --------------------------------------------------------------------
+# FastAPI App
+# --------------------------------------------------------------------
 app = FastAPI(
-    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"
+    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
 )
 
-# --- Firebase Initialization ---
+# --------------------------------------------------------------------
+# Firebase Initialization
+# --------------------------------------------------------------------
 db = None
 try:
     firebase_creds = os.getenv("FIREBASE_CREDS_JSON")
@@ -62,66 +76,67 @@ try:
         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')
-        if not firebase_admin._apps:
-            firebase_admin.initialize_app(cred)
-        db = firestore.client()
-        print("✅ Firebase initialized with local key file.")
+        raise Exception("FIREBASE_CREDS_JSON not found")
 except Exception as e:
-    print(f"❌ Could not initialize Firebase. Error: {e}")
+    print(f"❌ Could not initialize Firebase: {e}")
 
 def get_db():
     if db is None:
         raise HTTPException(status_code=503, detail="Firestore connection not available.")
     return db
 
-# --- Global Variables ---
+# --------------------------------------------------------------------
+# Global Variables (encoder + chroma)
+# --------------------------------------------------------------------
 encoder = None
 chroma_collection = None
 
 @app.on_event("startup")
 def load_model_and_data():
     global encoder, chroma_collection
-    
     print("🚀 Loading sentence encoder model...")
     encoder = SentenceEncoder()
     
-    # --- THIS IS THE FIX ---
-    # Point ChromaDB to the correct writable persistent storage path on Hugging Face
+    # Point ChromaDB to the persistent /data storage path
     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()
+    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()}")
+    except Exception as e:
+        print(f"❌ Error initializing ChromaDB: {e}")
+        raise
 
-# --- API Endpoints ---
+# --------------------------------------------------------------------
+# Endpoints
+# --------------------------------------------------------------------
 @app.get("/")
 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(...)):
+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()
         print("--- SETUP COMPLETE ---")
         return {"status": "Setup completed successfully."}
     except Exception as e:
-        raise HTTPException(status_code=500, detail=f"An error occurred: {str(e)}")
-
+        raise HTTPException(status_code=500, detail=f"An error occurred during setup: {str(e)}")
 
 @app.post("/add-internship", response_model=StatusResponse)
 def add_internship(internship: InternshipData, db_client: firestore.Client = Depends(get_db)):
@@ -143,14 +158,21 @@ def add_internship(internship: InternshipData, db_client: firestore.Client = Dep
 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_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]})
+        recommendations.append({
+            "internship_id": internship_id,
+            "score": 1 - distances[i]
+        })
+    
     return {"recommendations": recommendations}
 
 @app.post("/search", response_model=RecommendationResponse)
@@ -164,9 +186,4 @@ def search_internships(search: SearchQuery):
     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}
-
-@app.post("/chat", response_model=ChatResponse)
-def chat_with_bot(message: ChatMessage):
-    response = get_rag_response(message.query)
-    return {"response": response}
+    return {"recommendations": recommendations}