"""PlannerValidator — checks a TaskList before it reaches the TaskRunner. Runs the 8 checks from AGENT_ARCHITECTURE_CONTEXT_new.md §7.3. On failure it raises `PlannerValidationError` with a message specific enough that the planner can be re-prompted to self-correct (the retry loop lives in service.py). Check #1 (Pydantic parse) is enforced at the structured-output boundary — by the time a `TaskList` reaches here it has already parsed; this validator additionally rejects structurally-invalid plans (duplicate ids, dangling edges, cycles). """ from __future__ import annotations from pydantic import ValidationError from ...catalog.models import Catalog from ...query.ir.models import QueryIR from ...query.ir.validator import IRValidationError, IRValidator from .contracts import ToolRegistry from .errors import PlannerValidationError from .inputs import Constraints from .schemas import PLACEHOLDER_RE, TaskList # Heuristic: a checkable success_criteria mentions a measurable signal. _CHECKABLE_TOKENS = ("rate", "count", "match", "produced", "above", "below", "equal") # DFS colors for cycle detection. _WHITE, _GREY, _BLACK = 0, 1, 2 class PlannerValidator: def __init__(self, ir_validator: IRValidator | None = None) -> None: self._ir_validator = ir_validator or IRValidator() def validate( self, task_list: TaskList, registry: ToolRegistry, catalog: Catalog, constraints: Constraints, ) -> None: tasks = task_list.tasks # Check 6 — plan non-empty and within the task cap. if not tasks: raise PlannerValidationError("plan is empty: at least one task is required") if len(tasks) > constraints.max_tasks: raise PlannerValidationError( f"plan has {len(tasks)} tasks, exceeds max_tasks={constraints.max_tasks}" ) ids = [t.id for t in tasks] if len(set(ids)) != len(ids): dupes = sorted({i for i in ids if ids.count(i) > 1}) raise PlannerValidationError(f"duplicate task id(s): {dupes}") id_set = set(ids) tasks_by_id = {t.id: t for t in tasks} known_tools = registry.names() known_sources = {s.source_id for s in catalog.sources} for task in tasks: for call in task.tool_calls: # Check 2 — every tool exists in the registry. if call.tool not in known_tools: raise PlannerValidationError( f"task {task.id}: tool {call.tool!r} not in registry " f"(known: {sorted(known_tools)})" ) spec = registry.get(call.tool) assert spec is not None # guaranteed by the membership check above # Check 8a — args carry the required keys and no unknown keys. required = set(spec.input_schema.get("required", [])) allowed = set(spec.input_schema.get("properties", {}).keys()) | required missing = required - set(call.args.keys()) if missing: raise PlannerValidationError( f"task {task.id}: tool {call.tool!r} missing required arg(s): " f"{sorted(missing)}" ) unknown = set(call.args.keys()) - allowed if unknown: raise PlannerValidationError( f"task {task.id}: tool {call.tool!r} has unknown arg(s): " f"{sorted(unknown)} (allowed: {sorted(allowed)})" ) # Check 3 — concrete source_id args must exist in the catalog. src = call.args.get("source_id") if isinstance(src, str) and not _is_placeholder(src): if src not in known_sources: raise PlannerValidationError( f"task {task.id}: tool {call.tool!r} references unknown " f"source_id {src!r} (known: {sorted(known_sources)})" ) # Check 8b — inline retrieve_data IR validates against the catalog. if call.tool == "retrieve_data": self._validate_inline_ir(task.id, call.args, catalog) # Check 7 — success_criteria is checkable. if not _is_checkable(task.success_criteria): raise PlannerValidationError( f"task {task.id}: success_criteria is not checkable — include a " f"measurable signal (one of {list(_CHECKABLE_TOKENS)}); " f"got {task.success_criteria!r}" ) # Check 4 — DAG: edges resolve, placeholders resolve, no cycles. self._validate_dag(tasks_by_id, id_set) def _validate_inline_ir(self, task_id: str, args: dict, catalog: Catalog) -> None: raw_ir = args.get("ir") if not isinstance(raw_ir, dict): raise PlannerValidationError( f"task {task_id}: retrieve_data.args.ir must be an inline QueryIR " f"object, got {type(raw_ir).__name__}" ) try: ir = QueryIR.model_validate(raw_ir) except ValidationError as e: raise PlannerValidationError( f"task {task_id}: retrieve_data.args.ir is not a valid QueryIR: {e}" ) from e try: self._ir_validator.validate(ir, catalog) except IRValidationError as e: raise PlannerValidationError( f"task {task_id}: retrieve_data IR failed catalog validation: {e}" ) from e @staticmethod def _validate_dag(tasks_by_id: dict, id_set: set[str]) -> None: for task in tasks_by_id.values(): for dep in task.depends_on: if dep not in id_set: raise PlannerValidationError( f"task {task.id}: depends_on references unknown task {dep!r}" ) if dep == task.id: raise PlannerValidationError( f"task {task.id}: depends_on includes itself" ) cycle = _find_cycle(tasks_by_id) if cycle: raise PlannerValidationError(f"cycle detected in depends_on: {' -> '.join(cycle)}") # On an acyclic graph, a placeholder is safe iff its target is a # transitive ancestor — i.e. guaranteed to have completed before this # task runs. Requiring a *direct* depends_on would wrongly reject valid # plans that depend on the target through an intermediate task. ancestors = _all_ancestors(tasks_by_id) for task in tasks_by_id.values(): for ref in _placeholder_refs(task): if ref not in id_set: raise PlannerValidationError( f"task {task.id}: placeholder '${{{ref}}}' references unknown task" ) if ref not in ancestors[task.id]: raise PlannerValidationError( f"task {task.id}: placeholder '${{{ref}}}' used but {ref!r} is " f"not a (transitive) dependency — add it to depends_on" ) def _is_placeholder(value: str) -> bool: return bool(PLACEHOLDER_RE.fullmatch(value.strip())) def _placeholder_refs(task) -> set[str]: refs: set[str] = set() for call in task.tool_calls: for value in call.args.values(): if isinstance(value, str): refs.update(PLACEHOLDER_RE.findall(value)) return refs def _is_checkable(text: str) -> bool: low = text.lower() return any(tok in low for tok in _CHECKABLE_TOKENS) def _find_cycle(tasks_by_id: dict) -> list[str] | None: color = {tid: _WHITE for tid in tasks_by_id} stack: list[str] = [] def dfs(node: str) -> list[str] | None: color[node] = _GREY stack.append(node) for dep in tasks_by_id[node].depends_on: if color.get(dep) == _GREY: idx = stack.index(dep) return stack[idx:] + [dep] if color.get(dep) == _WHITE: found = dfs(dep) if found: return found stack.pop() color[node] = _BLACK return None for tid in tasks_by_id: if color[tid] == _WHITE: found = dfs(tid) if found: return found return None def _all_ancestors(tasks_by_id: dict) -> dict[str, set[str]]: """ancestors[id] = all tasks reachable by following depends_on edges.""" cache: dict[str, set[str]] = {} def visit(node: str, seen: set[str]) -> set[str]: if node in cache: return cache[node] acc: set[str] = set() for dep in tasks_by_id[node].depends_on: if dep in seen or dep not in tasks_by_id: continue acc.add(dep) acc |= visit(dep, seen | {dep}) cache[node] = acc return acc return {tid: visit(tid, {tid}) for tid in tasks_by_id}