Update main.py

main.py

@@ -4,14 +4,18 @@ import chromadb
 import math # ✅ Add the math library for ceiling division
 from fastapi import FastAPI, HTTPException, Depends, Query
 from pydantic import BaseModel, Field
-from typing import List
+from typing import List, Optional
 import firebase_admin
 from firebase_admin import credentials, firestore
 
 # --- Local Imports ---
 from encoder import SentenceEncoder
 from populate_chroma import populate_vector_db
-from llm_handler import initialize_llm, get_rag_response
+from llm_handler import (
+    initialize_llm, get_rag_response, create_chat_session, 
+    clear_chat_session, delete_chat_session, get_chat_history,
+    get_chat_session_count
+)
 import llm_handler
 
 # --------------------------------------------------------------------
@@ -44,7 +48,6 @@ class SimpleRecommendation(BaseModel):
     internship_id: str
     score: float
 
-# --- ✅ UPDATED RESPONSE MODEL ---
 class RecommendationResponse(BaseModel):
     recommendations: List[SimpleRecommendation]
     page: int
@@ -54,19 +57,34 @@ class StatusResponse(BaseModel):
     status: str
     internship_id: str
 
+# --- ✅ UPDATED CHAT MODELS ---
 class ChatMessage(BaseModel):
     query: str
+    session_id: Optional[str] = Field(None, description="Chat session ID (optional)")
 
 class ChatResponse(BaseModel):
     response: str
+    session_id: str
+
+class NewChatSessionResponse(BaseModel):
+    session_id: str
+    message: str
+
+class ChatHistoryResponse(BaseModel):
+    session_id: str
+    history: List[dict]
+
+class ClearChatResponse(BaseModel):
+    session_id: str
+    message: str
 
 # --------------------------------------------------------------------
 # FastAPI App
 # --------------------------------------------------------------------
 app = FastAPI(
     title="Internship Recommendation & Chatbot API",
-    description="An API using Firestore for metadata, ChromaDB for vector search, and an LLM chatbot.",
-    version="3.1.0",
+    description="An API using Firestore for metadata, ChromaDB for vector search, and an LLM chatbot with memory.",
+    version="3.2.0",
     root_path=root_path
 )
 
@@ -118,13 +136,11 @@ def load_model_and_data():
         raise
 
 # --------------------------------------------------------------------
-# Endpoints
+# Existing Endpoints
 # --------------------------------------------------------------------
 @app.get("/")
 def read_root():
-    return {"message": "Welcome to the Internship Recommendation API!"}
-
-# ... (setup, add-internship, and healthz endpoints are unchanged)
+    return {"message": "Welcome to the Internship Recommendation API with Chat Memory!"}
 
 @app.post("/setup")
 def run_initial_setup(secret_key: str = Query(..., example="your_secret_password")):
@@ -155,8 +171,6 @@ def add_internship(internship: InternshipData, db_client: firestore.Client = Dep
     print(f"✅ Added internship to Firestore and ChromaDB: {internship.id}")
     return {"status": "success", "internship_id": internship.id}
 
-
-# --- ✅ ENDPOINT UPDATED FOR PAGINATION ---
 @app.post("/profile-recommendations", response_model=RecommendationResponse)
 def get_profile_recommendations(profile: UserProfile, page: int = 1, page_size: int = 4):
     if chroma_collection is None or encoder is None:
@@ -165,16 +179,12 @@ def get_profile_recommendations(profile: UserProfile, page: int = 1, page_size:
     query_text = f"Skills: {', '.join(profile.skills)}. Sectors: {', '.join(profile.sectors)}"
     query_embedding = encoder.encode([query_text])[0].tolist()
     
-    # Query for all results to sort them, then paginate
-    # This is less efficient at scale, but simple and effective for this project
-    # A more advanced approach would use ChromaDB's offset/limit if available or other methods.
     total_items = chroma_collection.count()
     results = chroma_collection.query(
         query_embeddings=[query_embedding], 
-        n_results=total_items # Get all results to sort
+        n_results=total_items
     )
     
-    # Process and sort all hits by score
     ids = results.get('ids', [[]])[0]
     distances = results.get('distances', [[]])[0]
     all_recommendations = [
@@ -182,7 +192,6 @@ def get_profile_recommendations(profile: UserProfile, page: int = 1, page_size:
         for i in range(len(ids))
     ]
 
-    # --- PAGINATION LOGIC ---
     start_index = (page - 1) * page_size
     end_index = start_index + page_size
     paginated_results = all_recommendations[start_index:end_index]
@@ -195,9 +204,8 @@ def get_profile_recommendations(profile: UserProfile, page: int = 1, page_size:
         "total_pages": total_pages
     }
 
-# --- ✅ ENDPOINT UPDATED FOR PAGINATION ---
 @app.post("/search", response_model=RecommendationResponse)
-def search_internships(search: SearchQuery, page: int = 1, page_size: int =4 ):
+def search_internships(search: SearchQuery, page: int = 1, page_size: int = 4):
     if chroma_collection is None or encoder is None:
         raise HTTPException(status_code=503, detail="Server is not ready.")
     
@@ -227,7 +235,88 @@ def search_internships(search: SearchQuery, page: int = 1, page_size: int =4 ):
         "total_pages": total_pages
     }
 
+# --------------------------------------------------------------------
+# ✅ NEW CHAT ENDPOINTS WITH MEMORY
+# --------------------------------------------------------------------
+
+@app.post("/chat/new-session", response_model=NewChatSessionResponse)
+def create_new_chat_session():
+    """Create a new chat session."""
+    session_id = create_chat_session()
+    return {
+        "session_id": session_id,
+        "message": "New chat session created successfully"
+    }
+
 @app.post("/chat", response_model=ChatResponse)
 def chat_with_bot(message: ChatMessage):
-    response = get_rag_response(message.query)
-    return {"response": response}
+    """Chat with the bot. Maintains memory within the session."""
+    try:
+        response, session_id = get_rag_response(message.query, message.session_id)
+        return {
+            "response": response,
+            "session_id": session_id
+        }
+    except Exception as e:
+        # Handle the case where get_rag_response returns only response (backward compatibility)
+        if isinstance(e, ValueError):
+            response = get_rag_response(message.query, message.session_id)
+            return {
+                "response": response,
+                "session_id": message.session_id or "unknown"
+            }
+        raise HTTPException(status_code=500, detail=str(e))
+
+@app.get("/chat/{session_id}/history", response_model=ChatHistoryResponse)
+def get_session_history(session_id: str):
+    """Get the chat history for a specific session."""
+    history = get_chat_history(session_id)
+    if history is None:
+        raise HTTPException(status_code=404, detail="Chat session not found")
+    return {
+        "session_id": session_id,
+        "history": history
+    }
+
+@app.delete("/chat/{session_id}/clear", response_model=ClearChatResponse)
+def clear_session_history(session_id: str):
+    """Clear the chat history for a specific session."""
+    success = clear_chat_session(session_id)
+    if not success:
+        raise HTTPException(status_code=404, detail="Chat session not found")
+    return {
+        "session_id": session_id,
+        "message": "Chat history cleared successfully"
+    }
+
+@app.delete("/chat/{session_id}/delete", response_model=ClearChatResponse)
+def delete_session(session_id: str):
+    """Delete a chat session completely."""
+    success = delete_chat_session(session_id)
+    if not success:
+        raise HTTPException(status_code=404, detail="Chat session not found")
+    return {
+        "session_id": session_id,
+        "message": "Chat session deleted successfully"
+    }
+
+@app.get("/chat/sessions/count")
+def get_active_sessions():
+    """Get the number of active chat sessions."""
+    count = get_chat_session_count()
+    return {
+        "active_sessions": count,
+        "message": f"There are {count} active chat sessions"
+    }
+
+# Health check endpoint
+@app.get("/healthz")
+def health_check():
+    status = {
+        "status": "healthy",
+        "encoder_ready": encoder is not None,
+        "chroma_ready": chroma_collection is not None,
+        "firebase_ready": db is not None,
+        "active_chat_sessions": get_chat_session_count()
+    }
+    return status