Update main.py

main.py

@@ -6,7 +6,7 @@ from pydantic import BaseModel, Field
 from typing import List
 import firebase_admin
 from firebase_admin import credentials, firestore
-
+import random
 # --- Local Imports ---
 from encoder import SentenceEncoder
 from populate_chroma import populate_vector_db # For the setup endpoint
@@ -165,19 +165,18 @@ def get_profile_recommendations(profile: UserProfile, db_client: firestore.Clien
     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)
+    results = chroma_collection.query(query_embeddings=[query_embedding], n_results=random.randint(3, 5))
     
     recommendations = []
     ids = results.get('ids', [[]])[0]
     distances = results.get('distances', [[]])[0]
-
+    
     for i, internship_id in enumerate(ids):
-        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)
-
+        recommendations.append({
+            "internship_id": internship_id,
+            "score": 1 - distances[i]
+        })
+    
     return {"recommendations": recommendations}
 
 @app.post("/search", response_model=RecommendationResponse)
@@ -187,19 +186,13 @@ def search_internships(search: SearchQuery, db_client: firestore.Client = Depend
         
     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=random.randint(3, 5))
     
     recommendations = []
     ids = results.get('ids', [[]])[0]
     distances = results.get('distances', [[]])[0]
-
     for i, internship_id in enumerate(ids):
-        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)
-            
+        recommendations.append({"internship_id": internship_id, "score": 1 - distances[i]})
     return {"recommendations": recommendations}
 
 @app.post("/chat", response_model=ChatResponse)