Update main.py

main.py

@@ -1,11 +1,13 @@
 from flask import Flask, request, jsonify
+from flask_cors import CORS
+import json
+
 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_firestore  # 🔄 using your function
-from flask_cors import CORS
-import json
+from data_fetcher import fetch_data_from_firestore
+from data_planner import determine_data_requirements  # 🧠 Replaces inline planner
 
 app = Flask(__name__)
 CORS(app)
@@ -22,34 +24,26 @@ def chat():
     if not role or not user_input or not company_code:
         return jsonify({"error": "Missing role, message, or companyCode"}), 400
 
-    # ✨ Ask Gemini to generate the Firestore fetch plan
-    planner_prompt = {
-        "role": "system",
-        "content": (
-            "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}"
-        )
-    }
+    # 🧠 Step 1: Use Gemini to plan the data to fetch
+    planning_result = determine_data_requirements(user_input, company_code)
+    print("📋 Data Plan:", planning_result)
 
-    user_msg = { "role": "user", "content": user_input }
-    planning_response = ask_gpt([planner_prompt, user_msg])
+    if "error" in planning_result:
+        return jsonify({"reply": f"⚠️ Planning error: {planning_result['error']}"})
 
     try:
-        plan = json.loads(planning_response)
-        firestore_data = fetch_data_from_firestore(plan)
+        firestore_data = fetch_data_from_firestore(planning_result)
+        print("📦 Firestore Fetched Data:", firestore_data)
     except Exception as e:
-        return jsonify({ "reply": f"⚠️ Planning or Firestore error: {str(e)}" })
+        return jsonify({"reply": f"⚠️ Firestore fetch error: {str(e)}"})
 
-    # ✨ Ask Gemini to analyze data + answer user
+    # 💬 Step 2: Ask Gemini to answer using that data
     system_msg = build_system_message(company_code)
     data_msg = {
         "role": "system",
         "content": f"Here is the data from Firestore:\n{json.dumps(firestore_data)}"
     }
+    user_msg = { "role": "user", "content": user_input }
 
     final_response = ask_gpt([system_msg, data_msg, user_msg])
     return jsonify({ "reply": final_response })