ace/steps/reflect.py

"""ReflectStep — analyses a trace to produce a ReflectorOutput."""

from __future__ import annotations

import logging

from ..core.context import ACEStepContext
from ..core.outputs import AgentOutput
from ..protocols import ReflectorLike

logger = logging.getLogger(__name__)


class ReflectStep:
    """Run the Reflector role against the trace and current skillbook.

    Receives ``ctx.trace`` — a dict from EvaluateStep (standard ACE pipeline),
    a raw object from TraceAnalyser, or any integration-produced trace.  When
    the trace is a dict with known keys, the step extracts them and calls the
    Reflector's existing API.  For raw/opaque traces, it passes them as
    keyword arguments for the Reflector to handle.

    Declares ``async_boundary = True`` — everything from this step onward
    runs in a background thread pool when the pipeline has background
    execution enabled.

    Pure — produces a reflection object, no side effects.
    """

    requires = frozenset({"trace", "skillbook"})
    provides = frozenset({"reflections"})

    async_boundary = True
    max_workers = 3

    def __init__(self, reflector: ReflectorLike) -> None:
        self.reflector = reflector

    @staticmethod
    def _is_batch_container(trace: dict) -> bool:
        for key in ("items", "tasks"):
            if isinstance(trace.get(key), list):
                return True

        steps = trace.get("steps")
        return (
            isinstance(steps, list)
            and bool(steps)
            and all(
                isinstance(step, dict)
                and step.get("role") == "conversation"
                and isinstance(step.get("content"), dict)
                for step in steps
            )
        )

    def __call__(self, ctx: ACEStepContext) -> ACEStepContext:
        trace = ctx.trace

        if isinstance(trace, dict) and not self._is_batch_container(trace):
            # Structured trace from EvaluateStep — extract known fields
            agent_output = AgentOutput(
                reasoning=trace.get("reasoning", ""),
                final_answer=trace.get("answer", ""),
            )
            reflection = self.reflector.reflect(
                question=trace.get("question", ""),
                agent_output=agent_output,
                skillbook=ctx.skillbook,
                ground_truth=trace.get("ground_truth"),
                feedback=trace.get("feedback"),
                injected_skill_ids=ctx.injected_skill_ids,
                mode=ctx.mode,
            )
        else:
            # Raw trace from TraceAnalyser or integration — pass as-is
            # The Reflector must handle the trace type via **kwargs
            reflection = self.reflector.reflect(
                question="",
                agent_output=AgentOutput(reasoning="", final_answer=""),
                skillbook=ctx.skillbook,
                trace=trace,
                injected_skill_ids=ctx.injected_skill_ids,
                mode=ctx.mode,
            )

        return ctx.replace(reflections=(reflection,))