Update data_planner.py

data_planner.py

@@ -1,4 +1,3 @@
-# 🔧 data_planner.py
 import json
 from openai_client import ask_gpt
 
@@ -10,10 +9,8 @@ def determine_data_requirements(user_prompt: str, company_code: str):
             "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"
-            "Always include the filter 'companyCode': '<provided>' in your output.\n"
-            "Strictly return only a valid JSON object with keys: collections, filters, instruction.\n"
-            "Never return markdown, explanations, or preamble.\n"
-            "Return ONLY a raw JSON object."
+            f"Always include the filter 'companyCode': '{company_code}' in your output.\n"
+            "Strictly return only a valid JSON object with keys: collections, filters, instruction."
         )
     }
 
@@ -24,13 +21,11 @@ def determine_data_requirements(user_prompt: str, company_code: str):
 
     try:
         result = ask_gpt([system_msg, user_msg])
-        print("🧠 Raw planning response from Gemini:")
-        print(result)
+        print("🧠 Raw planning response from Gemini:\n", result)
 
-        parsed = json.loads(result)
-        if "collections" not in parsed:
-            raise ValueError("Response JSON missing 'collections'")
-        return parsed
+        # Strip markdown formatting (``` or ```json)
+        cleaned = result.strip().removeprefix("```json").removeprefix("```").removesuffix("```").strip()
+
+        return json.loads(cleaned)
     except Exception as e:
-        print("❌ Planning error:", str(e))
-        return {"error": f"Planning failed: {str(e)}"}
+        return {"error": f"Planning failed: {str(e)}"}