src/agents/slow_path/schemas.py

"""Slow-path execution + output contracts.

The seams between the three slow-path stages:
- TaskRunner writes `RunState` (a blackboard of `TaskResult`s) — §8.2.
- Assembler reads `RunState` + `BusinessContext` and produces `AssembledOutput`
  (`chat_answer` + `AnalysisRecord`) — §8.3.

`ToolOutput` (the tool -> agent envelope) is reused from the planner contracts so
there is exactly one definition across the layer.

Note on authorship (§8.3): the Assembler LLM authors only the *narrative* fields
(`AssemblerNarrative`). The `AnalysisRecord`'s structured pass-through fields
(`results_snapshot`, `tasks_run`) and metadata are copied from `RunState` by code,
never re-authored by the model — that is the source of truth the report generator
renders from (INV-4).

See AGENT_ARCHITECTURE_CONTEXT_new.md §8.2 / §8.3.
"""

from __future__ import annotations

from datetime import datetime
from typing import Literal
from uuid import uuid4

from pydantic import BaseModel, Field

from ..planner.contracts import ToolOutput
from ..planner.schemas import CrispStage

TaskStatus = Literal["success", "partial", "failure"]


# --------------------------------------------------------------------------- #
# Execution state (TaskRunner -> Assembler) — §8.2
# --------------------------------------------------------------------------- #


class TaskResult(BaseModel):
    task_id: str
    stage: CrispStage  # copied from the plan Task; carries CRISP-DM grouping to the report
    status: TaskStatus
    objective: str
    outputs: list[ToolOutput] = Field(default_factory=list)  # one per tool_call
    note: str | None = None
    error: str | None = None


class RunState(BaseModel):
    plan_id: str
    business_context_id: str
    results: dict[str, TaskResult] = Field(default_factory=dict)  # task_id -> result
    open_questions: list[str] = Field(default_factory=list)


# --------------------------------------------------------------------------- #
# Assembled output (Assembler -> Orchestrator / memory) — §8.3
# --------------------------------------------------------------------------- #


class TaskSummary(BaseModel):
    task_id: str
    stage: CrispStage  # lets the report group the method appendix by CRISP-DM phase
    objective: str
    status: TaskStatus
    tools_used: list[str] = Field(default_factory=list)


class AnalysisRecord(BaseModel):
    # Identity. `record_id` is the unit the report cites and snapshots
    # (`record_ids`); `analysis_id`/`user_id` scope the record to one analysis
    # session + owner and are stamped by the composition root / ReportInputStore at
    # persist time (they depend on the Analysis State that lives outside the slow
    # path), so they default to None when the Assembler first builds the record.
    record_id: str = Field(default_factory=lambda: uuid4().hex)
    analysis_id: str | None = None
    user_id: str | None = None
    # Narrative fields — authored by the Assembler LLM.
    goal_restated: str
    findings: list[str] = Field(default_factory=list)
    caveats: list[str] = Field(default_factory=list)
    data_used: list[str] = Field(default_factory=list)
    open_questions: list[str] = Field(default_factory=list)
    # Structured pass-through — NOT re-authored; copied from RunState.
    tasks_run: list[TaskSummary] = Field(default_factory=list)
    results_snapshot: dict[str, TaskResult] = Field(default_factory=dict)
    # Metadata.
    plan_id: str
    business_context_id: str
    created_at: datetime


class AssembledOutput(BaseModel):
    chat_answer: str  # FIRST field — streams via SSE; markdown prose + tables
    analysis_record: AnalysisRecord


class AssemblerNarrative(BaseModel):
    """The subset of `AnalysisRecord` the Assembler LLM actually authors.

    Kept separate from `AssembledOutput` so the model never emits the structured
    pass-through fields (which would invite hallucinated numbers); `Assembler`
    code merges this with the real `RunState` to build the final record.
    """

    chat_answer: str
    goal_restated: str
    findings: list[str] = Field(default_factory=list)
    caveats: list[str] = Field(default_factory=list)
    data_used: list[str] = Field(default_factory=list)
    open_questions: list[str] = Field(default_factory=list)