Update main.py

main.py

@@ -16,145 +16,154 @@ from dotenv import load_dotenv
 from firebase_admin import credentials, firestore, storage
 from google import genai
 
+import os
+import json
+import pickle
+import numpy as np
+from flask import Flask, request, jsonify
+from flask_cors import CORS
+from dotenv import load_dotenv
+from firebase_admin import credentials, firestore, storage, initialize_app
+from google import genai
+import faiss
+
 load_dotenv()
 
-# --------- Flask & Firebase Setup ---------
+# --- Flask Setup ---
 app = Flask(__name__)
 CORS(app)
 
-# Initialize Firebase with Firestore + Storage
+# --- Firebase Initialization ---
 cred_json = os.environ.get("FIREBASE")
 if not cred_json:
     raise RuntimeError("Missing FIREBASE env var")
 cred = credentials.Certificate(json.loads(cred_json))
-firebase_admin.initialize_app(cred, {
-    "storageBucket": os.environ.get("Firebase_Storage")
-})
-fs      = firestore.client()
-bucket  = storage.bucket()
+initialize_app(cred, {"storageBucket": os.environ.get("Firebase_Storage")})
+
+fs = firestore.client()
+bucket = storage.bucket()
 
-# --------- Google GenAI Client ---------
+# --- Gemini Client ---
 client = genai.Client(api_key=os.getenv("Gemini"))
 
-# --------- FAISS Cache Paths (unchanged) ---------
+# --- FAISS Setup ---
 INDEX_PATH = "vector.index"
-DOCS_PATH  = "documents.pkl"
-
-
+DOCS_PATH = "documents.pkl"
 
-# --------- Fetch & Summarize Firestore Docs ---------
-def fetch_documents() -> list[str]:
-    docs: list[str] = []
+# --- Role-Aware Firestore Fetch ---
+def fetch_documents(role: str, user_id: str) -> list[str]:
+    docs = []
 
-    # 1) Participants
+    # 1) participants
     for snap in fs.collection("participants").stream():
         d = snap.to_dict()
-        name    = d.get('name', 'Unknown Participant')
-        ent     = d.get('enterpriseName', 'Unknown Enterprise')
-        sector  = d.get('sector', 'Unknown Sector')
-        stage   = d.get('stage', 'Unknown Stage')
-        devtype = d.get('developmentType', 'Unknown Type')
-        docs.append(
-            f"{name} ({ent}), sector: {sector}, stage: {stage}, type: {devtype}."
-        )
-
-    # 2) Interventions
-    for snap in fs.collection("interventions").stream():
+        owner_id = snap.id
+        if role == "incubatee" and owner_id != user_id:
+            continue
+        if role == "consultant" and user_id not in d.get("assignedConsultants", []):
+            continue
+        name = d.get('beneficiaryName', 'Unknown')
+        ent = d.get('enterpriseName', 'Unknown')
+        sector = d.get('sector', 'Unknown')
+        stage = d.get('stage', 'Unknown')
+        devtype = d.get('developmentType', 'Unknown')
+        docs.append(f"{name} ({ent}), sector: {sector}, stage: {stage}, type: {devtype}.")
+
+    # 2) consultants
+    for snap in fs.collection("consultants").stream():
+        d = snap.to_dict()
+        if role == "consultant" and snap.id != user_id:
+            continue
+        name = d.get("name", "Unknown")
+        expertise = ", ".join(d.get("expertise", [])) or "no listed expertise"
+        rating = d.get("rating", "no rating")
+        docs.append(f"Consultant {name} with expertise in {expertise} and rating {rating}.")
+
+    # 3) programs
+    if role in ["admin", "operations", "funder", "incubatee"]:
+        for snap in fs.collection("programs").stream():
+            d = snap.to_dict()
+            docs.append(f"Program {d.get('name')} ({d.get('status')}): {d.get('type')} - Budget {d.get('budget')}")
+
+    # 4) interventions
+    if role in ["admin", "operations", "incubatee"]:
+        for snap in fs.collection("interventions").stream():
+            d = snap.to_dict()
+            for item in d.get('interventions', []):
+                title = item.get("title")
+                area = d.get("areaOfSupport", "General")
+                if title:
+                    docs.append(f"Intervention: {title} under {area}.")
+
+    # 5) assignedInterventions
+    for snap in fs.collection("assignedInterventions").stream():
         d = snap.to_dict()
-        area = d.get('area', 'General')
-        for item in d.get('interventions', []):
-            title = item.get('title')
-            if title:
-                docs.append(f"Intervention: {title} under {area}.")
+        if role == "consultant" and user_id not in d.get("consultantId", []):
+            continue
+        if role == "incubatee" and d.get("participantId") != user_id:
+            continue
+        title = d.get("interventionTitle", "Unknown")
+        sme = d.get("smeName", "Unknown")
+        status = d.get("status", "Unknown")
+        docs.append(f"Assigned intervention '{title}' for {sme} ({status})")
 
-    # 3) Feedbacks
+    # 6) feedbacks
     for snap in fs.collection("feedbacks").stream():
         d = snap.to_dict()
-        intervention = d.get('interventionTitle', 'Unknown Intervention')
-        smeName      = d.get('smeName', 'Unknown SME')
-        comment      = d.get('comment')
+        if role == "consultant" and d.get("consultantId") != user_id:
+            continue
+        intervention = d.get("interventionTitle", "Unknown")
+        comment = d.get("comment")
         if comment:
-            docs.append(f"Feedback on {intervention} by {smeName}: {comment}")
+            docs.append(f"Feedback on {intervention}: {comment}")
 
-    # 4) Compliance Documents
+    # 7) complianceDocuments
     for snap in fs.collection("complianceDocuments").stream():
         d = snap.to_dict()
-        pName   = d.get('participantName', 'Unknown Participant')
-        docType = d.get('documentType', 'Unknown Type')
-        status  = d.get('status', 'Unknown Status')
-        expiry  = d.get('expiryDate', 'Unknown Expiry')
-        docs.append(
-            f"Compliance document '{docType}' for {pName} is {status} (expires {expiry})."
-        )
-
-    # 5) Assigned Interventions
-    for snap in fs.collection("assignedInterventions").stream():
-        d = snap.to_dict()
-        title   = d.get('interventionTitle', 'Unknown Intervention')
-        smeName = d.get('smeName', 'Unknown SME')
-        cons    = d.get('consultantId', 'Unknown Consultant')
-        status  = d.get('status', 'Unknown Status')
-        docs.append(
-            f"Assigned intervention '{title}' for {smeName} by consultant {cons} ({status})."
-        )
-
-    # 6) Consultants
-    for snap in fs.collection("consultants").stream():
-        d = snap.to_dict()
-        name     = d.get('name', 'Unknown Consultant')
-        expertise= d.get('expertise', [])
-        rating   = d.get('rating')
-        exp_txt  = ", ".join(expertise) if expertise else "no listed expertise"
-        rating_txt = f"rating {rating}" if rating is not None else "no rating"
-        docs.append(f"Consultant {name} with expertise in {exp_txt} and {rating_txt}.")
+        if role == "incubatee" and d.get("participantId") != user_id:
+            continue
+        docs.append(f"Compliance document '{d.get('documentType')}' for {d.get('participantName')} is {d.get('status')} (expires {d.get('expiryDate')})")
+
+    # 8) interventionDatabase
+    if role in ["admin", "operations", "director", "funder"]:
+        for snap in fs.collection("interventionDatabase").stream():
+            d = snap.to_dict()
+            title = d.get("interventionTitle", "Unknown")
+            status = d.get("status", "Unknown")
+            feedback = d.get("feedback", "")
+            docs.append(f"Finalized intervention '{title}' ({status}): {feedback}")
 
     return docs
 
-# --------- Embedding Helper ---------
+# --- Embedding ---
 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)
-    )
+    resp = client.models.embed_content(model="text-embedding-004", contents=texts)
     return [emb.values for emb in resp.embeddings]
 
-# --------- Build or Load FAISS Index ---------
-def build_or_load_index():
-    if os.path.exists(INDEX_PATH) and os.path.exists(DOCS_PATH):
-        with open(DOCS_PATH, "rb") as f:
-            documents = pickle.load(f)
-        index = faiss.read_index(INDEX_PATH)
-    else:
-        documents = fetch_documents()
-        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
-
-documents, index = build_or_load_index()
-
-# --------- RAG Chat Helper ---------
-def retrieve_and_respond(user_query: str, top_k: int = 3) -> str:
-    # 1) Embed query
+# --- Dynamic Index ---
+def build_faiss_index(docs: list[str]):
+    embs = np.array(get_embeddings(docs), dtype="float32")
+    dim = embs.shape[1]
+    index = faiss.IndexFlatIP(dim)
+    index.add(embs)
+    return index
+
+# --- Retrieval Helper ---
+def retrieve_and_respond(user_query: str, role: str, user_id: str) -> str:
+    docs = fetch_documents(role, user_id)
+    if not docs:
+        return "No relevant data found for your role or access level."
+
+    index = build_faiss_index(docs)
     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
+    _, idxs = index.search(q_emb, 3)
+    ctx = "\n\n".join(docs[i] for i in idxs[0])
     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")
+    chat = client.chats.create(model="gemini-2.0-flash-thinking-exp")  
     resp = chat.send_message(prompt)
     return resp.text
 
 
-
 # --------- Helpers for Bank-Statement Processing ---------
 
 def read_pdf_pages(file_obj):
@@ -250,13 +259,18 @@ def process_pdf_pages(pdf_file):
 def chat_endpoint():
     data = request.get_json(force=True)
     q = data.get("user_query")
-    if not q:
-        return jsonify({"error": "Missing user_query"}), 400
+    role = data.get("role")
+    user_id = data.get("user_id")
+
+    if not q or not role or not user_id:
+        return jsonify({"error": "Missing user_query, role, or user_id"}), 400
+
     try:
-        return jsonify({"reply": retrieve_and_respond(q)})
+        reply = retrieve_and_respond(q, role.lower(), user_id)
+        return jsonify({"reply": reply})
     except Exception as e:
         return jsonify({"error": str(e)}), 500
-
+        
 # --------- Endpoint: Upload & Store Bank Statements ---------
 
 @app.route("/upload_statements", methods=["POST"])