Chat

main.py

@@ -1,47 +1,126 @@
 import os
-import io
 import json
-import hashlib
-import uuid
-from datetime import datetime, time, timedelta
-from PIL import Image
-import pytz
-from flask import Flask, request, jsonify, send_file
+import faiss
+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, db, storage, auth, firestore
-import pandas as pd
-import numpy as np
-import requests
-from urllib.parse import urlparse, unquote
+from firebase_admin import credentials, firestore
+from dotenv import load_dotenv
+
+load_dotenv()
 
+# --------- Flask Setup ---------
 app = Flask(__name__)
 CORS(app)
 
-# ---------- Firebase initialization ----------
-Firebase_Storage = os.getenv("Firebase_Storage")
+# --------- 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")
+embed_model = genai.EmbeddingModel("models/embedding-001")
+
+# --------- 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')}.")
+
+    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')}.")
 
-try:
-    cred_json = os.environ.get("FIREBASE")
-    if cred_json:
-        cred = credentials.Certificate(json.loads(cred_json))
-        firebase_admin.initialize_app(cred, {
-            'storageBucket': Firebase_Storage
-        })
+    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 Caching ---------
+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("FIREBASE secret not set.")
-except Exception as e:
-    print(f"Firebase init error: {e}")
+        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)
 
-bucket = storage.bucket()
-fs = firestore.client()
+        # Save cache
+        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 = embed_model.get_embeddings(texts=texts)
+    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 ---------
+@app.route("/chat", methods=["POST"])
+def chat():
+    data = request.get_json(force=True)
+    user_query = data.get("user_query")
 
-# ---------- Helper functions ----------
-def configure_gemini():
-    genai.configure(api_key=os.getenv("Gemini"))
-    return genai.GenerativeModel('gemini-2.0-flash-thinking-exp')
+    if not user_query:
+        return jsonify({"error": "Missing user_query"}), 400
 
+    try:
+        reply = retrieve_and_respond(user_query)
+        return jsonify({"reply": reply})
+    except Exception as e:
+        return jsonify({"error": str(e)}), 500
 
-if __name__ == '__main__':
-    app.run(debug=True, host="0.0.0.0", port=7860)
+# --------- Run Flask Server ---------
+if __name__ == "__main__":
+    app.run(host="0.0.0.0", port=7860, debug=True)