Skip to content

Commit c669667

Browse files
feat(sdk): add StepRecorder for incremental trace/session Step trees
Manually populating Step.children for a trace/session control today means hand-building a side-channel record per leaf call and converting it into Step objects afterward (see ace-demo's build_agent_control_step). Add agent_control.record_step()/StepRecorder: an incremental builder that constructs the Step tree directly as children execute, reusing @control()'s capture logic (_create_evaluation_payload) for .call()/ .acall(). Extract evaluate_step() out of evaluate_controls so both the manual API and the new recorder share the same Step -> evaluation path. Co-Authored-By: Claude Code <noreply@anthropic.com>
1 parent 1f07d7c commit c669667

7 files changed

Lines changed: 656 additions & 30 deletions

File tree

‎sdks/python/README.md‎

Lines changed: 48 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -84,6 +84,54 @@ async with agent_control.AgentControlClient() as client:
8484
The existing `evaluate_controls` helper remains available for callers that
8585
prefer its field-based convenience arguments.
8686

87+
## Building trace/session steps
88+
89+
Trace- and session-level controls evaluate the aggregate of several
90+
already-executed child steps, so `Step.children` has to be populated by the
91+
caller. Building that tree by hand means hand-rolling a side-channel record
92+
for every leaf call and converting it into `Step` objects afterward:
93+
94+
```python
95+
# Before: a hand-built dict record next to the real return value
96+
def execute_policy_search(query):
97+
docs = search_policy_documents(query)
98+
record = {"type": "retriever", "query": query, "docs": docs}
99+
return StepExecution(value=docs, record=record)
100+
101+
spans = []
102+
result = execute_policy_search(query)
103+
spans.append(result.record)
104+
...
105+
trace_step = build_agent_control_step({"type": "trace", "spans": spans, ...})
106+
```
107+
108+
`agent_control.record_step()` builds the same tree incrementally, in `Step`
109+
vocabulary, with no intermediate dict and no converter:
110+
111+
```python
112+
with agent_control.record_step("trace", "banking_trace", input={"request": req}) as trace:
113+
with trace.child("retriever", "policy_lookup", input=query) as span:
114+
span.output = search_policy_documents(query)
115+
116+
account = trace.call(lookup_account, account_id="acct-1001", step_type="tool")
117+
plan = trace.call(run_banking_model, req, step_type="llm", tools=TOOL_DEFINITIONS)
118+
trace.output = {"status": "planned", "message": plan["content"]}
119+
120+
result = await trace.evaluate(stage="post")
121+
```
122+
123+
`trace.call(...)`/`await trace.acall(...)` run the function and record it as
124+
a child using the same capture logic as `@control()` (input from bound
125+
arguments, output from the return value), then return the real result -
126+
removing the need for a separate `StepExecution`-style wrapper. A failed
127+
call is still recorded (with the error in `context`) before the exception is
128+
re-raised. `trace.child(...)` nests another recorder the same way, so
129+
sessions nest traces with `session.child("trace", ...)`.
130+
131+
`trace.build()` produces the frozen `Step`; `trace.evaluate(...)` builds it
132+
and evaluates it in one call via `evaluate_step()` - the same function
133+
`evaluate_controls()` uses internally once its `Step` is built.
134+
87135
## Sharing an OpenTelemetry provider with Google ADK
88136

89137
When Google ADK and Agent Control should export through the same OpenTelemetry

‎sdks/python/src/agent_control/__init__.py‎

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -92,7 +92,7 @@ async def handle_input(user_message: str) -> str:
9292
)
9393
from .client import AgentControlClient
9494
from .control_decorators import ControlSteerError, ControlViolationError, control
95-
from .evaluation import check_evaluation_with_local, evaluate_controls
95+
from .evaluation import check_evaluation_with_local, evaluate_controls, evaluate_step
9696
from .observability import (
9797
LogConfig,
9898
add_event,
@@ -116,6 +116,7 @@ async def handle_input(user_message: str) -> str:
116116
)
117117
from .otel_sink import control_event_to_otel_span
118118
from .runtime_auth import validate_http_field_name
119+
from .step_recorder import StepRecorder, record_step
119120
from .tracing import (
120121
get_current_span_id,
121122
get_current_trace_id,
@@ -1619,6 +1620,10 @@ async def main():
16191620
# Local evaluation
16201621
"check_evaluation_with_local",
16211622
"evaluate_controls",
1623+
"evaluate_step",
1624+
# Step recorder (incremental trace/session Step tree builder)
1625+
"record_step",
1626+
"StepRecorder",
16221627
# Tracing
16231628
"get_trace_and_span_ids",
16241629
"get_current_trace_id",

‎sdks/python/src/agent_control/evaluation.py‎

Lines changed: 63 additions & 27 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,6 @@
11
"""Evaluation check operations for Agent Control SDK."""
22

3-
from collections.abc import Awaitable, Callable
3+
from collections.abc import Awaitable, Callable, Mapping, Sequence
44
from dataclasses import dataclass
55
from inspect import iscoroutinefunction
66
from typing import Any, Literal, cast
@@ -515,6 +515,58 @@ def _with_parse_errors(result: EvaluationResult) -> EvaluationResult:
515515
return _with_parse_errors(EvaluationResult(is_safe=True, confidence=1.0))
516516

517517

518+
async def evaluate_step(
519+
step: Step,
520+
*,
521+
agent_name: str,
522+
stage: Literal["pre", "post"] = "pre",
523+
target_type: str | None = None,
524+
target_id: str | None = None,
525+
trace_id: str | None = None,
526+
span_id: str | None = None,
527+
) -> EvaluationResult:
528+
"""Evaluate controls for an already-built ``Step``.
529+
530+
This is the shared tail of :func:`evaluate_controls`: resolve the
531+
session target, open a client, and run local/server evaluation. Use
532+
it directly when the ``Step`` - including any ``children`` - was
533+
already assembled, e.g. via :class:`~agent_control.step_recorder.StepRecorder`.
534+
535+
When ``target_type`` and ``target_id`` are both supplied, the request
536+
is target-bearing: the server merges target bindings into the
537+
effective control set. If they are omitted, the SDK falls back to the
538+
target context fixed at ``init()`` time when present. A per-call
539+
override that disagrees with the session target is rejected because
540+
the cached controls were fetched for the session target and would
541+
otherwise drive stale local-first evaluation.
542+
"""
543+
if state.server_url is None:
544+
raise RuntimeError("Server URL not configured. Call agent_control.init() first.")
545+
546+
target_type, target_id = _resolve_session_target(target_type, target_id)
547+
resolved_controls = state.server_controls or []
548+
549+
async with AgentControlClient(
550+
base_url=state.server_url,
551+
api_key=state.api_key,
552+
api_key_header=state.api_key_header,
553+
runtime_token_header=state.runtime_token_header,
554+
runtime_token_cache=state.runtime_token_cache,
555+
) as client:
556+
return await check_evaluation_with_local(
557+
client=client,
558+
agent_name=agent_name,
559+
step=step,
560+
stage=stage,
561+
controls=resolved_controls,
562+
target_type=target_type,
563+
target_id=target_id,
564+
trace_id=trace_id,
565+
span_id=span_id,
566+
event_agent_name=agent_name,
567+
)
568+
569+
518570
async def evaluate_controls(
519571
step_name: str,
520572
*,
@@ -523,7 +575,7 @@ async def evaluate_controls(
523575
context: dict[str, Any] | None = None,
524576
tools: list[dict[str, JSONValue]] | None = None,
525577
ground_truth: JSONValue | None = None,
526-
children: list[Step] | None = None,
578+
children: Sequence[Step | Mapping[str, Any]] | None = None,
527579
step_type: str = "llm",
528580
stage: Literal["pre", "post"] = "pre",
529581
agent_name: str,
@@ -544,11 +596,6 @@ async def evaluate_controls(
544596
"""
545597
step_type = ensure_step_type(step_type)
546598

547-
if state.server_url is None:
548-
raise RuntimeError("Server URL not configured. Call agent_control.init() first.")
549-
550-
target_type, target_id = _resolve_session_target(target_type, target_id)
551-
552599
default_value = {} if step_type == "tool" else ""
553600
step_dict: dict[str, Any] = {
554601
"type": step_type,
@@ -566,24 +613,13 @@ async def evaluate_controls(
566613
step_dict["children"] = children
567614

568615
step_obj = Step(**step_dict) # type: ignore[arg-type]
569-
resolved_controls = state.server_controls or []
570616

571-
async with AgentControlClient(
572-
base_url=state.server_url,
573-
api_key=state.api_key,
574-
api_key_header=state.api_key_header,
575-
runtime_token_header=state.runtime_token_header,
576-
runtime_token_cache=state.runtime_token_cache,
577-
) as client:
578-
return await check_evaluation_with_local(
579-
client=client,
580-
agent_name=agent_name,
581-
step=step_obj,
582-
stage=stage,
583-
controls=resolved_controls,
584-
target_type=target_type,
585-
target_id=target_id,
586-
trace_id=trace_id,
587-
span_id=span_id,
588-
event_agent_name=agent_name,
589-
)
617+
return await evaluate_step(
618+
step_obj,
619+
agent_name=agent_name,
620+
stage=stage,
621+
target_type=target_type,
622+
target_id=target_id,
623+
trace_id=trace_id,
624+
span_id=span_id,
625+
)

0 commit comments

Comments
 (0)