Update main.py

main.py

@@ -3,8 +3,9 @@ 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 data_fetcher import fetch_data_from_firestore  # 🔄 using your function
 from flask_cors import CORS
+import json
 
 app = Flask(__name__)
 CORS(app)
@@ -21,35 +22,33 @@ def chat():
     if not role or not user_input or not company_code:
         return jsonify({"error": "Missing role, message, or companyCode"}), 400
 
-    # 🔹 First step: instruct Gemini to tell us what to fetch
+    # ✨ Ask Gemini to generate the Firestore fetch plan
     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}"
+            "You are a planning agent. Given a user query and context, return a list of Firestore collections to query.\n"
+            "Respond ONLY in this JSON format:\n"
+            "{ \"collections\": [ { \"name\": \"collectionName\", \"filters\": [ {\"field\": \"\", \"op\": \"==\", \"value\": \"\"} ], \"limit\": 50 } ] }\n"
+            f"The companyCode is '{company_code}' and role is '{role}'.\n"
+            "Do not generate explanations. Only the raw JSON.\n\n"
+            f"User query: {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)
+        firestore_data = fetch_data_from_firestore(plan)
     except Exception as e:
-        return jsonify({ "reply": f"⚠️ Failed to plan or fetch data: {str(e)}" })
+        return jsonify({ "reply": f"⚠️ Planning or Firestore error: {str(e)}" })
 
-    # 🔹 Third step: answer the original question using real data
+    # ✨ Ask Gemini to analyze data + answer user
     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)}"
+        "content": f"Here is the data from Firestore:\n{json.dumps(firestore_data)}"
     }
 
     final_response = ask_gpt([system_msg, data_msg, user_msg])