data_planner.py

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,
    user_id: str,
    participant_email: str | None,
    participant_id: str | None,
):
    token_participant = "{{{{participantId}}}}"   # renders as {{participantId}} inside f-string

    system_msg = {
        "role": "system",
        "content": (
            "You are a planning assistant for a Firestore chatbot.\n"
            "Return ONLY JSON with keys: collections, filters, instruction.\n"
            "Rules:\n"
            "- Always include filters.companyCode set to the current company.\n"
            "- The 'userId' is an Auth UID (users/{uid}).\n"
            "- 'applications.participantId' is the doc ID of /participants/{participantId}.\n"
            "- To reference the current user's participant-linked data, use this participantId.\n"
            f"- If you need the current participant, use token {token_participant} in filters.\n"
            "- If participantId is unavailable, prefer filtering by email where applicable.\n\n"
            f"Current companyCode: '{company_code}'\n"
            f"Current uid: '{user_id}'\n"
            f"Current user email: '{participant_email or ''}'\n"
            f"Current participantId: '{participant_id or ''}'\n"
            "Use this schema for all reasoning. It includes field names, types, and aliases:\n"
            f"{json.dumps(schema)}\n\n"
            "PLANNING RULES:\n"
            "- Choose collections ONLY if they contain the fields needed to answer.\n"
            "- If the user asks about headcount/revenue history, choose 'participants' (fields: headcountHistory, revenueHistory).\n"
            "- Apply companyCode ONLY on collections that have companyCode in the schema (do NOT add companyCode to 'participants').\n"
            "- If an alias is used in the question, map it to the canonical field using $aliases.\n"
            "- Return 'collections' as a list of objects [{'name': <collection>, 'fields': [<field1>, ...]}].\n"
            "- 'filters' must only include fields that exist in the chosen collections.\n"
            "Strictly return only a valid JSON object with keys: collections, filters, instruction. No explanation."
        ),
    }

    user_msg = {
        "role": "user",
        "content": (
            f"User question: {user_prompt}\n"
            f"CompanyCode: {company_code}\n"
            f"UserID: {user_id}"
        ),
    }

    try:
        result = ask_gpt([system_msg, user_msg])
        print("🧠 Raw planning response from Gemini:\n", result)

        cleaned = (
            result.strip()
            .removeprefix("```json")
            .removeprefix("```")
            .removesuffix("```")
            .strip()
        )
        return json.loads(cleaned)
    except Exception as e:
        return {"error": f"Planning failed: {str(e)}"}