File size: 7,742 Bytes
116524e
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
"""ACE — full adaptive pipeline runner."""

from __future__ import annotations

from collections.abc import Iterable, Sequence
from pathlib import Path
from types import MappingProxyType
from typing import Any

from pipeline import Pipeline
from pipeline.protocol import SampleResult, StepProtocol

from ..core.context import ACEStepContext, SkillbookView
from ..core.environments import Sample, TaskEnvironment
from ..core.insight_source import TRACE_IDENTITY_METADATA_KEY, infer_trace_identity
from ..protocols import (
    AgentLike,
    DeduplicationManagerLike,
    ReflectorLike,
    SkillManagerLike,
)
from ..core.skillbook import Skillbook
from ..steps import AgentStep, EvaluateStep, learning_tail
from .base import ACERunner


class ACE(ACERunner):
    """Live adaptive pipeline: Agent -> Evaluate -> Reflect -> Tag -> Update -> Apply.



    The full ACE loop.  An agent executes, the environment evaluates, the

    reflector analyses, and the skill manager updates the skillbook.



    A single class handles both single-pass (``epochs=1``) and multi-epoch

    batch training (``epochs > 1``).



    Use when you are building a new agent from scratch and want

    closed-loop learning where the agent improves in real time.

    """

    @classmethod
    def build_steps(

        cls,

        *,

        agent: AgentLike,

        reflector: ReflectorLike,

        skill_manager: SkillManagerLike,

        environment: TaskEnvironment | None = None,

        skillbook: Skillbook | None = None,

        dedup_manager: DeduplicationManagerLike | None = None,

        dedup_interval: int = 10,

        checkpoint_dir: str | Path | None = None,

        checkpoint_interval: int = 10,

        extra_steps: list[StepProtocol] | None = None,

    ) -> list[StepProtocol]:
        """Return the steps that ``from_roles()`` would compose.



        Use this to inspect, modify, or extend the pipeline before

        constructing it yourself::



            steps = ACE.build_steps(agent=agent, reflector=reflector, ...)

            steps.insert(2, MyCustomStep())

            pipe = Pipeline(steps)

            runner = ACERunner(pipeline=pipe, skillbook=skillbook)



        Args:

            agent: Agent role for producing answers.

            reflector: Reflector role for analysing execution.

            skill_manager: SkillManager role for update operations.

            environment: Optional task environment for evaluation feedback.

            skillbook: Starting skillbook.  Creates an empty one if ``None``.

            dedup_manager: Optional deduplication manager.

            dedup_interval: Samples between deduplication runs.

            checkpoint_dir: Directory for checkpoint files.

            checkpoint_interval: Samples between checkpoint saves.

            extra_steps: Additional steps appended after the learning

                tail (e.g. ``OpikStep``).

        """
        skillbook = skillbook or Skillbook()
        steps: list[StepProtocol[ACEStepContext]] = [
            AgentStep(agent, skillbook),
            EvaluateStep(environment),
            *learning_tail(
                reflector,
                skill_manager,
                skillbook,
                dedup_manager=dedup_manager,
                dedup_interval=dedup_interval,
                checkpoint_dir=checkpoint_dir,
                checkpoint_interval=checkpoint_interval,
            ),
        ]
        if extra_steps:
            steps.extend(extra_steps)
        return steps

    @classmethod
    def from_roles(

        cls,

        *,

        agent: AgentLike,

        reflector: ReflectorLike,

        skill_manager: SkillManagerLike,

        environment: TaskEnvironment | None = None,

        skillbook: Skillbook | None = None,

        dedup_manager: DeduplicationManagerLike | None = None,

        dedup_interval: int = 10,

        checkpoint_dir: str | Path | None = None,

        checkpoint_interval: int = 10,

        extra_steps: list[StepProtocol] | None = None,

    ) -> ACE:
        """Construct from pre-built role instances.



        Args:

            agent: Agent role for producing answers.

            reflector: Reflector role for analysing execution.

            skill_manager: SkillManager role for update operations.

            environment: Optional task environment for evaluation feedback.

                When provided, ``EvaluateStep`` generates feedback that

                enriches the trace.  When omitted, the trace still contains

                the agent's output, question, context, and ground truth.

            skillbook: Starting skillbook.  Creates an empty one if ``None``.

            dedup_manager: Optional deduplication manager.

            dedup_interval: Samples between deduplication runs.

            checkpoint_dir: Directory for checkpoint files.

            checkpoint_interval: Samples between checkpoint saves.

            extra_steps: Additional steps appended after the learning

                tail (e.g. ``OpikStep``).

        """
        skillbook = skillbook or Skillbook()
        steps = cls.build_steps(
            agent=agent,
            reflector=reflector,
            skill_manager=skill_manager,
            environment=environment,
            skillbook=skillbook,
            dedup_manager=dedup_manager,
            dedup_interval=dedup_interval,
            checkpoint_dir=checkpoint_dir,
            checkpoint_interval=checkpoint_interval,
            extra_steps=extra_steps,
        )
        return cls(pipeline=Pipeline(steps), skillbook=skillbook)

    def run(

        self,

        samples: Sequence[Sample] | Iterable[Sample],

        epochs: int = 1,

        *,

        wait: bool = True,

    ) -> list[SampleResult]:
        """Run the adaptive pipeline over samples.



        Args:

            samples: Input samples.  Must be a ``Sequence`` for

                ``epochs > 1``.  ``Iterable`` is accepted when

                ``epochs=1`` (consumed once).

            epochs: Number of passes over the samples.

            wait: If ``True``, block until background learning completes.



        Returns:

            List of ``SampleResult``, one per sample per epoch.



        Raises:

            ValueError: If ``epochs > 1`` and *samples* is not a

                ``Sequence``.

        """
        return self._run(samples, epochs=epochs, wait=wait)

    def _build_context(  # type: ignore[override]

        self,

        sample: Sample,

        *,

        epoch: int,

        total_epochs: int,

        index: int,

        total: int | None,

        global_sample_index: int,

        **_: Any,

    ) -> ACEStepContext:
        """Map a ``Sample`` to an ``ACEStepContext`` for the full pipeline.



        Sets ``sample`` and ``skillbook`` on the context.  The environment

        (if any) is injected into ``EvaluateStep`` at construction time —

        it does not appear on the context.

        """
        return ACEStepContext(
            sample=sample,
            metadata=MappingProxyType(
                {
                    TRACE_IDENTITY_METADATA_KEY: infer_trace_identity(
                        sample=sample,
                        metadata=sample.metadata,
                        default_source_system="sample",
                    ).to_dict()
                }
            ),
            skillbook=SkillbookView(self.skillbook),
            epoch=epoch,
            total_epochs=total_epochs,
            step_index=index,
            total_steps=total,
            global_sample_index=global_sample_index,
        )