Update main.py

main.py

@@ -5,122 +5,107 @@ import numpy as np
 import pickle
 from flask import Flask, request, jsonify
 from flask_cors import CORS
-import google.generativeai as genai
 import firebase_admin
 from firebase_admin import credentials, firestore
 from dotenv import load_dotenv
 
+from google import genai
+from google.genai import types
+
 load_dotenv()
 
-# --------- Flask Setup ---------
+# --------- Flask & Firebase Setup ---------
 app = Flask(__name__)
 CORS(app)
 
-# --------- Firebase Initialization ---------
 cred_json = os.environ.get("FIREBASE")
 if cred_json:
     cred = credentials.Certificate(json.loads(cred_json))
     firebase_admin.initialize_app(cred)
 fs = firestore.client()
 
-# --------- Gemini Configuration ---------
-genai.configure(api_key=os.getenv("Gemini"))
-chat_model = genai.GenerativeModel("gemini-2.0-flash-thinking-exp")
-# --------- Paths for Cached Index ---------
-INDEX_PATH = "vector.index"
-DOCS_PATH = "documents.pkl"
-
-# --------- Load Documents from Firestore ---------
-def fetch_documents():
-    documents = []
-
-    for doc in fs.collection("participants").stream():
-        d = doc.to_dict()
-        documents.append(f"{d.get('name')} ({d.get('enterpriseName')}), sector: {d.get('sector')}, stage: {d.get('stage')}, type: {d.get('developmentType')}.")
-
-    for doc in fs.collection("interventions").stream():
-        d = doc.to_dict()
-        for item in d.get("interventions", []):
-            documents.append(f"Intervention: {item.get('title')} under {d.get('area')}.")
+# --------- Google GenAI Client ---------
+client = genai.Client(api_key=os.getenv("Gemini"))
 
-    for doc in fs.collection("feedbacks").stream():
-        d = doc.to_dict()
-        documents.append(f"Feedback on {d.get('interventionTitle')} by {d.get('smeName')}: {d.get('comment')}")
-
-    for doc in fs.collection("complianceDocuments").stream():
-        d = doc.to_dict()
-        documents.append(f"Compliance document '{d.get('documentType')}' for {d.get('participantName')} is {d.get('status')} and expires on {d.get('expiryDate')}.")
-
-    for doc in fs.collection("assignedInterventions").stream():
-        d = doc.to_dict()
-        documents.append(f"Assigned intervention '{d.get('interventionTitle')}' for {d.get('smeName')} by consultant {d.get('consultantId')} with status {d.get('status')}.")
-
-    for doc in fs.collection("consultants").stream():
-        d = doc.to_dict()
-        documents.append(f"Consultant {d.get('name')} with expertise in {', '.join(d.get('expertise', []))} and rating {d.get('rating')}.")
-
-    return documents
+# --------- FAISS Cache Paths ---------
+INDEX_PATH = "vector.index"
+DOCS_PATH  = "documents.pkl"
+
+# --------- Fetch & Summarize Firestore Docs ---------
+def fetch_documents() -> list[str]:
+    docs = []
+    for col in [
+        ("participants", lambda d: f"{d['name']} ({d['enterpriseName']}), sector: {d['sector']}, stage: {d['stage']}, type: {d['developmentType']}."),
+        ("interventions", lambda d: [f"Intervention: {i['title']} under {d['area']}." for i in d.get("interventions", [])]),
+        ("feedbacks", lambda d: f"Feedback on {d['interventionTitle']} by {d['smeName']}: {d['comment']}"),
+        ("complianceDocuments", lambda d: f"Compliance document '{d['documentType']}' for {d['participantName']} is {d['status']} (expires {d['expiryDate']})."),
+        ("assignedInterventions", lambda d: f"Assigned '{d['interventionTitle']}' for {d['smeName']} by consultant {d['consultantId']} ({d['status']})."),
+        ("consultants", lambda d: f"Consultant {d['name']} – expertise: {', '.join(d.get('expertise', []))}, rating: {d['rating']}.")
+    ]:
+        for snap in fs.collection(col[0]).stream():
+            entry = snap.to_dict()
+            out = col[1](entry)
+            # flatten lists
+            if isinstance(out, list):
+                docs.extend(out)
+            else:
+                docs.append(out)
+    return docs
+
+# --------- Embedding Helper ---------
+def get_embeddings(texts: list[str]) -> list[list[float]]:
+    resp = client.models.embed_content(
+        model="text-embedding-004",
+        contents=texts
+        # , config=types.EmbedContentConfig(output_dimensionality=512)
+    )
+    return [emb.values for emb in resp.embeddings]
 
-# --------- FAISS Caching ---------
+# --------- Build or Load FAISS Index ---------
 def build_or_load_index():
     if os.path.exists(INDEX_PATH) and os.path.exists(DOCS_PATH):
-        print("Loading FAISS index and documents from cache...")
         with open(DOCS_PATH, "rb") as f:
             documents = pickle.load(f)
         index = faiss.read_index(INDEX_PATH)
     else:
-        print("Building FAISS index...")
         documents = fetch_documents()
-        embeddings = np.array(get_embeddings(documents), dtype="float32")
-        dimension = len(embeddings[0])
-        index = faiss.IndexFlatIP(dimension)
-        index.add(embeddings)
-
-        # Save cache
+        embs = np.array(get_embeddings(documents), dtype="float32")
+        dim = embs.shape[1]
+        index = faiss.IndexFlatIP(dim)
+        index.add(embs)
+        # cache to disk
         with open(DOCS_PATH, "wb") as f:
             pickle.dump(documents, f)
         faiss.write_index(index, INDEX_PATH)
     return documents, index
 
-def get_embeddings(texts):
-    response= genai.embed_content(
-            model="models/text-embedding-004", content=texts, output_dimensionality=10
-        )
-    return [e.values for e in response.embeddings]
-
 documents, index = build_or_load_index()
 
-# --------- Helper Function: RAG Chat ---------
-def retrieve_and_respond(user_query, top_k=3):
-    query_embedding = np.array(get_embeddings([user_query]), dtype="float32")
-    distances, indices = index.search(query_embedding, top_k)
-    retrieved_docs = [documents[i] for i in indices[0]]
-
-    prompt = (
-        "Use the following context to answer the question.\n\n"
-        + "\n\n".join(retrieved_docs)
-        + f"\n\nQuestion: {user_query}\nAnswer:"
-    )
-
-    chat = chat_model.start_chat()
-    response = chat.send_message(prompt)
-    return getattr(response, "text", response.last.text)
-
-# --------- Flask Chat Endpoint ---------
+# --------- RAG Chat Helper ---------
+def retrieve_and_respond(user_query: str, top_k: int = 3) -> str:
+    # 1) Embed query
+    q_emb = np.array(get_embeddings([user_query]), dtype="float32")
+    # 2) Search index
+    _, idxs = index.search(q_emb, top_k)
+    ctx = "\n\n".join(documents[i] for i in idxs[0])
+    # 3) Build prompt
+    prompt = f"Use the context below to answer:\n\n{ctx}\n\nQuestion: {user_query}\nAnswer:"
+    # 4) Chat
+    chat = client.chats.create(model="gemini-2.0-flash-thinking-exp")
+    resp = chat.send_message(prompt)
+    return resp.text
+
+# --------- Flask Endpoint ---------
 @app.route("/chat", methods=["POST"])
-def chat():
+def chat_endpoint():
     data = request.get_json(force=True)
-    user_query = data.get("user_query")
-
-    if not user_query:
+    q = data.get("user_query")
+    if not q:
         return jsonify({"error": "Missing user_query"}), 400
-
     try:
-        reply = retrieve_and_respond(user_query)
-        return jsonify({"reply": reply})
+        return jsonify({"reply": retrieve_and_respond(q)})
     except Exception as e:
         return jsonify({"error": str(e)}), 500
 
-# --------- Run Flask Server ---------
 if __name__ == "__main__":
     app.run(host="0.0.0.0", port=7860, debug=True)