Update main.py

main.py

@@ -1,12 +1,13 @@
 from flask import Flask, request, jsonify
 from firestore_client import get_firestore_client
 from openai_client import ask_gpt
+from prompt_instructions import build_system_message
 from role_access import get_allowed_collections
+from data_fetcher import fetch_data_from_collections  # <-- new function we'll build
 from flask_cors import CORS
 
-
 app = Flask(__name__)
-CORS(app)  # Allow access from your React app
+CORS(app)
 
 db = get_firestore_client()
 
@@ -15,22 +16,44 @@ def chat():
     data = request.json
     role = data.get('role')
     user_input = data.get('message')
-    company_code = data.get('companyCode')  # <-- New addition
+    company_code = data.get('companyCode')
 
     if not role or not user_input or not company_code:
         return jsonify({"error": "Missing role, message, or companyCode"}), 400
-        
-  
-    system_msg = build_system_message(company_code)
 
+    # 🔹 First step: instruct Gemini to tell us what to fetch
+    planner_prompt = {
+        "role": "system",
+        "content": (
+            "Your job is to act like a planner. Based on the user query, output a list of Firestore collections and filters needed to answer it."
+            "Return your result strictly in this JSON format:\n\n"
+            "{ collections: [ { name: 'collectionName', filters: { field: value } } ] }\n\n"
+            "Only include collections that are actually relevant. The user's companyCode is "
+            f"'{company_code}' and their role is '{role}'. Here is the question:\n\n{user_input}"
+        )
+    }
+
+    user_msg = { "role": "user", "content": user_input }
+    planning_response = ask_gpt([planner_prompt, user_msg])
+
+    # 🔹 Second step: fetch data based on the plan
+    try:
+        import json
+        plan = json.loads(planning_response)
+        collections = plan.get("collections", [])
+        fetched_data = fetch_data_from_collections(db, collections)
+    except Exception as e:
+        return jsonify({ "reply": f"⚠️ Failed to plan or fetch data: {str(e)}" })
 
-    user_msg = {
-        "role": "user",
-        "content": user_input
+    # 🔹 Third step: answer the original question using real data
+    system_msg = build_system_message(company_code)
+    data_msg = {
+        "role": "system",
+        "content": f"The following data was retrieved from Firestore:\n\n{json.dumps(fetched_data)}"
     }
 
-    answer = ask_gpt([system_msg, user_msg])
-    return jsonify({"reply": answer})
+    final_response = ask_gpt([system_msg, data_msg, user_msg])
+    return jsonify({ "reply": final_response })
 
 if __name__ == "__main__":
     app.run(host="0.0.0.0", port=7860)