ace/core/outputs.py

"""Output types produced by ACE roles."""

from __future__ import annotations

from typing import Any, Dict, List, Optional

from pydantic import BaseModel, ConfigDict, Field, model_validator

from .skillbook import UpdateBatch


class AgentOutput(BaseModel):
    """Output from the Agent role containing reasoning and answer."""

    model_config = ConfigDict(arbitrary_types_allowed=True)

    reasoning: str = Field(..., description="Step-by-step reasoning process")
    final_answer: str = Field(..., description="The final answer to the question")
    skill_ids: List[str] = Field(
        default_factory=list, description="IDs of strategies cited in reasoning"
    )
    raw: Dict[str, Any] = Field(
        default_factory=dict, description="Raw LLM response data"
    )
    trace_context: Optional[Any] = Field(
        default=None,
        exclude=True,
        description="Pre-built TraceContext from integration (bypasses auto-detection)",
    )


class ExtractedLearning(BaseModel):
    """A single learning extracted by the Reflector from task execution."""

    learning: str = Field(..., description="The extracted learning or insight")
    evidence: str = Field(
        default="",
        description=(
            "Specific traces/items where this pattern was observed. "
            "Cite task IDs or item indices, e.g. 'task_2, task_16, task_29'."
        ),
    )
    justification: str = Field(
        default="",
        description=(
            "Why this is worth remembering: how many traces exhibited this pattern, "
            "whether it's recurring or a one-off, and why it generalizes beyond these examples."
        ),
    )


class ReflectorOutput(BaseModel):
    """Output from the Reflector role containing pure analysis.

    Reflector reports what it found; downstream (SkillManager) decides how
    to act. No tagging fields, no prescriptive output.
    """

    model_config = ConfigDict(arbitrary_types_allowed=True)

    reasoning: str = Field(..., description="Overall reasoning about the outcome")
    error_identification: str = Field(
        default="", description="Description of what went wrong (if applicable)"
    )
    root_cause_analysis: str = Field(
        default="", description="Analysis of why errors occurred"
    )
    correct_approach: str = Field(
        ..., description="What the correct approach should be"
    )
    key_insight: str = Field(
        ..., description="The main lesson learned from this iteration"
    )
    raw: Dict[str, Any] = Field(
        default_factory=dict, description="Raw LLM response data"
    )


class SkillManagerOutput(BaseModel):
    """Output from the SkillManager role containing skillbook update operations.

    Accepts both nested ``{"update": {"reasoning": ..., "operations": [...]}}``
    and the flat shape the LLM actually returns:
    ``{"reasoning": ..., "operations": [...]}``.
    """

    model_config = ConfigDict(arbitrary_types_allowed=True)

    update: UpdateBatch = Field(
        ..., description="Batch of update operations to apply to skillbook"
    )
    raw: Dict[str, Any] = Field(
        default_factory=dict, description="Raw LLM response data"
    )

    @model_validator(mode="before")
    @classmethod
    def _accept_flat_shape(cls, data: Any) -> Any:
        """If the LLM returns {reasoning, operations, ...} without an 'update'
        wrapper, nest it automatically so Pydantic can validate."""
        if isinstance(data, dict) and "update" not in data and "operations" in data:
            reasoning = data.pop("reasoning", "")
            operations = data.pop("operations", [])
            data["update"] = UpdateBatch.from_json(
                {"reasoning": reasoning, "operations": operations}
            )
        return data