Update data_planner.py

data_planner.py

@@ -1,16 +1,20 @@
 import json
 from openai_client import ask_gpt
+from schema_map import schema  # 👈 Import the full schema
 
 def determine_data_requirements(user_prompt: str, company_code: str):
     system_msg = {
         "role": "system",
         "content": (
-            "You are a planning assistant for a Firestore chatbot. Your job is to analyze the user's question and return ONLY a JSON object with:\n"
+            "You are a planning assistant for a Firestore chatbot.\n"
+            "Your job is to analyze the user's question and return ONLY a JSON object with:\n"
             "1. 'collections': a list of Firestore collections needed.\n"
-            "2. 'filters': a dictionary of filters to apply, such as companyCode.\n"
-            "3. 'instruction': a clear summary of what to do with the data (e.g., summarize interventions, count tasks).\n"
-            f"Always include the filter 'companyCode': '{company_code}' in your output.\n"
-            "Strictly return only a valid JSON object with keys: collections, filters, instruction."
+            "2. 'filters': a dictionary of filters to apply (e.g., companyCode).\n"
+            "3. 'instruction': a clear summary of what to do with the data (e.g., summarize interventions).\n\n"
+            f"Always include the filter 'companyCode': '{company_code}'.\n"
+            "Use this schema for all reasoning. It includes field names, types, and aliases:\n"
+            f"{json.dumps(schema)}\n\n"
+            "Strictly return only a valid JSON object with keys: collections, filters, instruction. No explanation."
         )
     }
 
@@ -23,9 +27,7 @@ def determine_data_requirements(user_prompt: str, company_code: str):
         result = ask_gpt([system_msg, user_msg])
         print("🧠 Raw planning response from Gemini:\n", result)
 
-        # Strip markdown formatting (``` or ```json)
         cleaned = result.strip().removeprefix("```json").removeprefix("```").removesuffix("```").strip()
-
         return json.loads(cleaned)
     except Exception as e:
-        return {"error": f"Planning failed: {str(e)}"}
+        return {"error": f"Planning failed: {str(e)}"}