src/agents/planner/service.py

"""PlannerService — single LLM call: context + catalog + tools + question -> TaskList.

Mirrors `query/planner/service.py` (chain construction) and `query/service.py`
(validate-and-retry loop). The planner LLM emits a `TaskList` via structured
output; the `PlannerValidator` runs the 8 checks; on failure the planner is
re-prompted with the error context, up to `max_retries` (default 3). No
replanning happens at execution time — this loop only hardens the *initial*
static plan.

The service takes the full `Catalog` (not just a `CatalogSummary`): it derives
the PII-safe `CatalogSummary` for the prompt, but validation needs the full
catalog so the existing `IRValidator` can check inline `retrieve_data` IRs.

See AGENT_ARCHITECTURE_CONTEXT_new.md §7.3.
"""

from __future__ import annotations

from pathlib import Path

from langchain_core.messages import SystemMessage
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.runnables import Runnable
from langchain_openai import AzureChatOpenAI

from src.middlewares.logging import get_logger

from ...catalog.models import Catalog
from .contracts import BusinessContext, ToolRegistry
from .errors import PlannerError, PlannerValidationError
from .inputs import CatalogSummary, Constraints
from .prompt import build_planner_prompt
from .schemas import TaskList
from .validator import PlannerValidator

logger = get_logger("planner_agent")

_PROMPT_PATH = (
    Path(__file__).resolve().parent.parent.parent / "config" / "prompts" / "planner.md"
)


def _load_prompt_text() -> str:
    return _PROMPT_PATH.read_text(encoding="utf-8")


def _build_default_chain() -> Runnable:
    from src.config.settings import settings

    llm = AzureChatOpenAI(
        azure_deployment=settings.azureai_deployment_name_4o,
        openai_api_version=settings.azureai_api_version_4o,
        azure_endpoint=settings.azureai_endpoint_url_4o,
        api_key=settings.azureai_api_key_4o,
        temperature=0,
    )
    prompt = ChatPromptTemplate.from_messages(
        [
            SystemMessage(content=_load_prompt_text()),
            ("human", "{human_content}"),
        ]
    )
    return prompt | llm.with_structured_output(TaskList)


_default_chain: Runnable | None = None


def _get_default_chain() -> Runnable:
    global _default_chain
    if _default_chain is None:
        _default_chain = _build_default_chain()
    return _default_chain


class PlannerService:
    """Wraps the planner LLM call + the validate-and-retry loop.

    Inject `structured_chain` and/or `validator` for tests.
    """

    def __init__(
        self,
        structured_chain: Runnable | None = None,
        validator: PlannerValidator | None = None,
        max_retries: int = 3,
    ) -> None:
        self._chain = structured_chain
        self._validator = validator or PlannerValidator()
        self._max_retries = max(1, max_retries)

    def _ensure_chain(self) -> Runnable:
        if self._chain is None:
            self._chain = _get_default_chain()
        return self._chain

    async def plan(
        self,
        context: BusinessContext,
        catalog: Catalog,
        tools: ToolRegistry,
        query: str,
        constraints: Constraints,
        callbacks: list | None = None,
    ) -> TaskList:
        summary = CatalogSummary.from_catalog(catalog)
        chain = self._ensure_chain()
        previous_errors: list[str] = []

        for attempt in range(1, self._max_retries + 1):
            human_content = build_planner_prompt(
                context, summary, tools, query, constraints, previous_errors
            )
            # All retry attempts share `callbacks`, so each shows up under the same
            # trace — that is how retry token cost becomes visible.
            if callbacks:
                task_list: TaskList = await chain.ainvoke(
                    {"human_content": human_content}, config={"callbacks": callbacks}
                )
            else:
                task_list = await chain.ainvoke({"human_content": human_content})
            try:
                self._validator.validate(task_list, tools, catalog, constraints)
            except PlannerValidationError as e:
                # Accumulate the full error history (oldest first) so the next
                # attempt sees every prior failure and can't fix one by
                # reintroducing another (the observed value_type -> arg -> value_type
                # whack-a-mole).
                previous_errors.append(f"attempt {attempt}: {e}")
                logger.warning(
                    "planner validation failed",
                    project_id=context.project_id,
                    plan_id=task_list.plan_id,
                    attempt=attempt,
                    error=str(e),
                )
                continue

            logger.info(
                "analysis planned",
                project_id=context.project_id,
                plan_id=task_list.plan_id,
                n_tasks=len(task_list.tasks),
                retry=attempt > 1,
                # Compact plan dump: task id, its tools, deps, and each tool's arg keys
                # (+ any analyze_* column refs) — enough to see how retrieve_data feeds
                # the analyze_* tools without dumping full inline IRs.
                plan=[
                    {
                        "id": t.id,
                        "depends_on": t.depends_on,
                        "tools": [
                            {
                                "tool": c.tool,
                                "args": sorted(c.args.keys()),
                                "cols": c.args.get("column_ids") or c.args.get("column"),
                                "data": c.args.get("data"),
                            }
                            for c in t.tool_calls
                        ],
                    }
                    for t in task_list.tasks
                ],
            )
            return task_list

        raise PlannerError(
            f"planner failed validation after {self._max_retries} attempts; "
            f"last error: {previous_errors[-1] if previous_errors else 'unknown'}"
        )


async def plan_analysis(
    context: BusinessContext,
    catalog: Catalog,
    tools: ToolRegistry,
    query: str,
    constraints: Constraints,
) -> TaskList:
    """Convenience entry point using the default chain + validator."""
    return await PlannerService().plan(context, catalog, tools, query, constraints)