From 2d918fccdcc9a884c3cc1ee1303a788deafd62f8 Mon Sep 17 00:00:00 2001 From: Yixin Date: Wed, 7 Oct 2026 13:46:43 -0700 Subject: [PATCH 1/6] 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 --- sdks/python/README.md | 48 ++++ sdks/python/src/agent_control/__init__.py | 7 +- sdks/python/src/agent_control/evaluation.py | 90 +++++--- .../python/src/agent_control/step_recorder.py | 215 ++++++++++++++++++ sdks/python/tests/test_evaluation.py | 54 ++++- sdks/python/tests/test_step_recorder.py | 185 +++++++++++++++ .../test_step_recorder_galileo_integration.py | 87 +++++++ 7 files changed, 656 insertions(+), 30 deletions(-) create mode 100644 sdks/python/src/agent_control/step_recorder.py create mode 100644 sdks/python/tests/test_step_recorder.py create mode 100644 sdks/python/tests/test_step_recorder_galileo_integration.py diff --git a/sdks/python/README.md b/sdks/python/README.md index 9522cf2e..94785665 100644 --- a/sdks/python/README.md +++ b/sdks/python/README.md @@ -84,6 +84,54 @@ async with agent_control.AgentControlClient() as client: The existing `evaluate_controls` helper remains available for callers that prefer its field-based convenience arguments. +## Building trace/session steps + +Trace- and session-level controls evaluate the aggregate of several +already-executed child steps, so `Step.children` has to be populated by the +caller. Building that tree by hand means hand-rolling a side-channel record +for every leaf call and converting it into `Step` objects afterward: + +```python +# Before: a hand-built dict record next to the real return value +def execute_policy_search(query): + docs = search_policy_documents(query) + record = {"type": "retriever", "query": query, "docs": docs} + return StepExecution(value=docs, record=record) + +spans = [] +result = execute_policy_search(query) +spans.append(result.record) +... +trace_step = build_agent_control_step({"type": "trace", "spans": spans, ...}) +``` + +`agent_control.record_step()` builds the same tree incrementally, in `Step` +vocabulary, with no intermediate dict and no converter: + +```python +with agent_control.record_step("trace", "banking_trace", input={"request": req}) as trace: + with trace.child("retriever", "policy_lookup", input=query) as span: + span.output = search_policy_documents(query) + + account = trace.call(lookup_account, account_id="acct-1001", step_type="tool") + plan = trace.call(run_banking_model, req, step_type="llm", tools=TOOL_DEFINITIONS) + trace.output = {"status": "planned", "message": plan["content"]} + +result = await trace.evaluate(stage="post") +``` + +`trace.call(...)`/`await trace.acall(...)` run the function and record it as +a child using the same capture logic as `@control()` (input from bound +arguments, output from the return value), then return the real result - +removing the need for a separate `StepExecution`-style wrapper. A failed +call is still recorded (with the error in `context`) before the exception is +re-raised. `trace.child(...)` nests another recorder the same way, so +sessions nest traces with `session.child("trace", ...)`. + +`trace.build()` produces the frozen `Step`; `trace.evaluate(...)` builds it +and evaluates it in one call via `evaluate_step()` - the same function +`evaluate_controls()` uses internally once its `Step` is built. + ## Sharing an OpenTelemetry provider with Google ADK When Google ADK and Agent Control should export through the same OpenTelemetry diff --git a/sdks/python/src/agent_control/__init__.py b/sdks/python/src/agent_control/__init__.py index 04d040f6..9cd90511 100644 --- a/sdks/python/src/agent_control/__init__.py +++ b/sdks/python/src/agent_control/__init__.py @@ -92,7 +92,7 @@ async def handle_input(user_message: str) -> str: ) from .client import AgentControlClient from .control_decorators import ControlSteerError, ControlViolationError, control -from .evaluation import check_evaluation_with_local, evaluate_controls +from .evaluation import check_evaluation_with_local, evaluate_controls, evaluate_step from .observability import ( LogConfig, add_event, @@ -116,6 +116,7 @@ async def handle_input(user_message: str) -> str: ) from .otel_sink import control_event_to_otel_span from .runtime_auth import validate_http_field_name +from .step_recorder import StepRecorder, record_step from .tracing import ( get_current_span_id, get_current_trace_id, @@ -1619,6 +1620,10 @@ async def main(): # Local evaluation "check_evaluation_with_local", "evaluate_controls", + "evaluate_step", + # Step recorder (incremental trace/session Step tree builder) + "record_step", + "StepRecorder", # Tracing "get_trace_and_span_ids", "get_current_trace_id", diff --git a/sdks/python/src/agent_control/evaluation.py b/sdks/python/src/agent_control/evaluation.py index 800b4bf4..749c4f00 100644 --- a/sdks/python/src/agent_control/evaluation.py +++ b/sdks/python/src/agent_control/evaluation.py @@ -1,6 +1,6 @@ """Evaluation check operations for Agent Control SDK.""" -from collections.abc import Awaitable, Callable +from collections.abc import Awaitable, Callable, Mapping, Sequence from dataclasses import dataclass from inspect import iscoroutinefunction from typing import Any, Literal, cast @@ -515,6 +515,58 @@ def _with_parse_errors(result: EvaluationResult) -> EvaluationResult: return _with_parse_errors(EvaluationResult(is_safe=True, confidence=1.0)) +async def evaluate_step( + step: Step, + *, + agent_name: str, + stage: Literal["pre", "post"] = "pre", + target_type: str | None = None, + target_id: str | None = None, + trace_id: str | None = None, + span_id: str | None = None, +) -> EvaluationResult: + """Evaluate controls for an already-built ``Step``. + + This is the shared tail of :func:`evaluate_controls`: resolve the + session target, open a client, and run local/server evaluation. Use + it directly when the ``Step`` - including any ``children`` - was + already assembled, e.g. via :class:`~agent_control.step_recorder.StepRecorder`. + + When ``target_type`` and ``target_id`` are both supplied, the request + is target-bearing: the server merges target bindings into the + effective control set. If they are omitted, the SDK falls back to the + target context fixed at ``init()`` time when present. A per-call + override that disagrees with the session target is rejected because + the cached controls were fetched for the session target and would + otherwise drive stale local-first evaluation. + """ + if state.server_url is None: + raise RuntimeError("Server URL not configured. Call agent_control.init() first.") + + target_type, target_id = _resolve_session_target(target_type, target_id) + resolved_controls = state.server_controls or [] + + async with AgentControlClient( + base_url=state.server_url, + api_key=state.api_key, + api_key_header=state.api_key_header, + runtime_token_header=state.runtime_token_header, + runtime_token_cache=state.runtime_token_cache, + ) as client: + return await check_evaluation_with_local( + client=client, + agent_name=agent_name, + step=step, + stage=stage, + controls=resolved_controls, + target_type=target_type, + target_id=target_id, + trace_id=trace_id, + span_id=span_id, + event_agent_name=agent_name, + ) + + async def evaluate_controls( step_name: str, *, @@ -523,7 +575,7 @@ async def evaluate_controls( context: dict[str, Any] | None = None, tools: list[dict[str, JSONValue]] | None = None, ground_truth: JSONValue | None = None, - children: list[Step] | None = None, + children: Sequence[Step | Mapping[str, Any]] | None = None, step_type: str = "llm", stage: Literal["pre", "post"] = "pre", agent_name: str, @@ -544,11 +596,6 @@ async def evaluate_controls( """ step_type = ensure_step_type(step_type) - if state.server_url is None: - raise RuntimeError("Server URL not configured. Call agent_control.init() first.") - - target_type, target_id = _resolve_session_target(target_type, target_id) - default_value = {} if step_type == "tool" else "" step_dict: dict[str, Any] = { "type": step_type, @@ -566,24 +613,13 @@ async def evaluate_controls( step_dict["children"] = children step_obj = Step(**step_dict) # type: ignore[arg-type] - resolved_controls = state.server_controls or [] - async with AgentControlClient( - base_url=state.server_url, - api_key=state.api_key, - api_key_header=state.api_key_header, - runtime_token_header=state.runtime_token_header, - runtime_token_cache=state.runtime_token_cache, - ) as client: - return await check_evaluation_with_local( - client=client, - agent_name=agent_name, - step=step_obj, - stage=stage, - controls=resolved_controls, - target_type=target_type, - target_id=target_id, - trace_id=trace_id, - span_id=span_id, - event_agent_name=agent_name, - ) + return await evaluate_step( + step_obj, + agent_name=agent_name, + stage=stage, + target_type=target_type, + target_id=target_id, + trace_id=trace_id, + span_id=span_id, + ) diff --git a/sdks/python/src/agent_control/step_recorder.py b/sdks/python/src/agent_control/step_recorder.py new file mode 100644 index 00000000..160d5ed0 --- /dev/null +++ b/sdks/python/src/agent_control/step_recorder.py @@ -0,0 +1,215 @@ +"""Incremental builder for ``Step`` trees (trace/session-level evaluation). + +Manually populating ``Step.children`` for a trace or session means building +the whole tree by hand before calling ``evaluate_controls``/``evaluate_step``. +``StepRecorder`` lets callers build that tree as they execute instead: each +nested call is appended to its parent as it happens, with no intermediate +dict representation and no separate converter. +""" + +from __future__ import annotations + +from collections.abc import Awaitable, Callable, Mapping +from typing import Any, Literal, TypeVar + +from agent_control_models import EvaluationResult, JSONObject, JSONValue, Step + +from ._state import state +from .control_decorators import ToolsConfiguration, _create_evaluation_payload +from .evaluation import evaluate_step +from .validation import ensure_step_type + +T = TypeVar("T") + +# Step types whose children are reported as a span/trace tree; an absent +# recording means "nothing happened", which the record factory treats +# differently from "not applicable" (None). See models/agent.py and +# records/factory.py for the trace/session child-shape rules this mirrors. +_EMPTY_CHILDREN_STEP_TYPES = frozenset({"trace", "session"}) + + +def record_step( + step_type: str, + name: str, + *, + input: Any | None = None, + context: dict[str, Any] | None = None, + tools: list[JSONObject] | None = None, + ground_truth: JSONValue | None = None, +) -> StepRecorder: + """Start building a ``Step`` tree. Use as a context manager.""" + return StepRecorder( + step_type, + name, + input=input, + context=context, + tools=tools, + ground_truth=ground_truth, + ) + + +class StepRecorder: + """Mutable builder for one ``Step`` node and its children. + + ``Step`` itself is frozen, so this accumulates state and only + constructs the real (frozen) ``Step`` in :meth:`build`. + """ + + def __init__( + self, + step_type: str, + name: str, + *, + input: Any | None = None, + context: dict[str, Any] | None = None, + tools: list[JSONObject] | None = None, + ground_truth: JSONValue | None = None, + _parent: StepRecorder | None = None, + ) -> None: + self.type = ensure_step_type(step_type) + self.name = name + self.input = input + self.output: Any | None = None + self.context: dict[str, Any] | None = dict(context) if context is not None else None + self.tools = tools + self.ground_truth = ground_truth + self._children: list[Step] = [] + self._parent = _parent + + def __enter__(self) -> StepRecorder: + return self + + def __exit__(self, *exc_info: object) -> None: + if self._parent is not None: + self._parent._children.append(self.build()) + + def child( + self, + step_type: str, + name: str, + *, + input: Any | None = None, + context: dict[str, Any] | None = None, + tools: list[JSONObject] | None = None, + ground_truth: JSONValue | None = None, + ) -> StepRecorder: + """Start a nested recorder; it attaches to this node on ``__exit__``.""" + return StepRecorder( + step_type, + name, + input=input, + context=context, + tools=tools, + ground_truth=ground_truth, + _parent=self, + ) + + def add(self, step: Step | Mapping[str, Any]) -> None: + """Attach an already-built child step.""" + self._children.append(step if isinstance(step, Step) else Step.model_validate(step)) + + def _record_call_result( + self, + func: Callable[..., Any], + args: tuple[Any, ...], + kwargs: dict[str, Any], + output: Any, + step_name: str | None, + step_type: str | None, + tools: ToolsConfiguration, + error: BaseException | None, + ) -> None: + payload = _create_evaluation_payload( + func, args, kwargs, output, step_name, step_type, tools + ) + if error is not None: + payload["context"] = {**(payload.get("context") or {}), "error": repr(error)} + self._children.append(Step.model_validate(payload)) + + def call( + self, + func: Callable[..., T], + *args: Any, + step_type: str | None = None, + step_name: str | None = None, + tools: ToolsConfiguration = None, + **kwargs: Any, + ) -> T: + """Run ``func``, record it via ``@control()``'s capture logic, return its output.""" + try: + output = func(*args, **kwargs) + except Exception as exc: + self._record_call_result(func, args, kwargs, None, step_name, step_type, tools, exc) + raise + self._record_call_result(func, args, kwargs, output, step_name, step_type, tools, None) + return output + + async def acall( + self, + func: Callable[..., Awaitable[T]], + *args: Any, + step_type: str | None = None, + step_name: str | None = None, + tools: ToolsConfiguration = None, + **kwargs: Any, + ) -> T: + """Async counterpart to :meth:`call`.""" + try: + output = await func(*args, **kwargs) + except Exception as exc: + self._record_call_result(func, args, kwargs, None, step_name, step_type, tools, exc) + raise + self._record_call_result(func, args, kwargs, output, step_name, step_type, tools, None) + return output + + def build(self) -> Step: + """Recursively construct the frozen ``Step`` for this node and its children.""" + step_dict: dict[str, Any] = { + "type": self.type, + "name": self.name, + "input": self.input, + "output": self.output, + } + if self.context is not None: + step_dict["context"] = self.context + if self.tools is not None: + step_dict["tools"] = self.tools + if self.ground_truth is not None: + step_dict["ground_truth"] = self.ground_truth + if self._children: + step_dict["children"] = list(self._children) + elif self.type in _EMPTY_CHILDREN_STEP_TYPES: + step_dict["children"] = [] + return Step(**step_dict) # type: ignore[arg-type] + + async def evaluate( + self, + *, + agent_name: str | None = None, + stage: Literal["pre", "post"] = "pre", + target_type: str | None = None, + target_id: str | None = None, + trace_id: str | None = None, + span_id: str | None = None, + ) -> EvaluationResult: + """Build this node into a ``Step`` and evaluate it via :func:`evaluate_step`. + + ``agent_name`` defaults to the agent registered via ``init()``. + """ + resolved_agent_name = agent_name or ( + state.current_agent.agent_name if state.current_agent is not None else None + ) + if resolved_agent_name is None: + raise RuntimeError( + "agent_name not supplied and no agent registered. " + "Call agent_control.init() first or pass agent_name explicitly." + ) + return await evaluate_step( + self.build(), + agent_name=resolved_agent_name, + stage=stage, + target_type=target_type, + target_id=target_id, + trace_id=trace_id, + span_id=span_id, + ) diff --git a/sdks/python/tests/test_evaluation.py b/sdks/python/tests/test_evaluation.py index 94f596fc..9657545a 100644 --- a/sdks/python/tests/test_evaluation.py +++ b/sdks/python/tests/test_evaluation.py @@ -4,11 +4,12 @@ from uuid import UUID import pytest -from agent_control import evaluation -from agent_control.evaluation import EvaluationResult from agent_control_models import Step from pydantic import ValidationError +from agent_control import evaluation +from agent_control.evaluation import EvaluationResult + @pytest.mark.asyncio async def test_check_evaluation_requires_step_name_before_server_call(): @@ -187,6 +188,55 @@ async def test_evaluate_controls_with_explicit_agent_name(monkeypatch): assert result.is_safe is True assert result.confidence == 1.0 + + +@pytest.mark.asyncio +async def test_evaluate_controls_delegates_to_evaluate_step(monkeypatch): + """evaluate_controls builds the Step then routes through evaluate_step.""" + mock_result = EvaluationResult(is_safe=True, confidence=1.0) + mock_evaluate_step = AsyncMock(return_value=mock_result) + monkeypatch.setattr(evaluation, "evaluate_step", mock_evaluate_step) + + result = await evaluation.evaluate_controls( + step_name="chat", + input="hello", + stage="pre", + agent_name="test-bot", + ) + + assert result is mock_result + mock_evaluate_step.assert_awaited_once() + args, kwargs = mock_evaluate_step.call_args + step_arg = args[0] + assert isinstance(step_arg, Step) + assert step_arg.name == "chat" + assert step_arg.input == "hello" + assert kwargs["agent_name"] == "test-bot" + assert kwargs["stage"] == "pre" + + +@pytest.mark.asyncio +async def test_evaluate_controls_coerces_dict_children_to_step(monkeypatch): + """children= accepts plain dicts; pydantic coerces them to Step instances.""" + mock_check = AsyncMock(return_value=EvaluationResult(is_safe=True, confidence=1.0)) + monkeypatch.setattr(evaluation, "check_evaluation_with_local", mock_check) + + with patch("agent_control.state.server_url", "http://localhost:8000"): + await evaluation.evaluate_controls( + step_name="trace1", + step_type="trace", + input={"request": "r1"}, + children=[{"type": "llm", "name": "x", "input": "hi", "output": "bye"}], + stage="post", + agent_name="test-bot", + ) + + step = mock_check.call_args.kwargs["step"] + assert step.children is not None + assert len(step.children) == 1 + assert isinstance(step.children[0], Step) + assert step.children[0].type == "llm" + assert step.children[0].name == "x" mock_check.assert_called_once() diff --git a/sdks/python/tests/test_step_recorder.py b/sdks/python/tests/test_step_recorder.py new file mode 100644 index 00000000..44ec33c2 --- /dev/null +++ b/sdks/python/tests/test_step_recorder.py @@ -0,0 +1,185 @@ +"""Tests for StepRecorder, the incremental Step tree builder.""" + +from unittest.mock import AsyncMock, patch + +import pytest +from agent_control_models import EvaluationResult, Step + +from agent_control import record_step, step_recorder +from agent_control.control_decorators import _create_evaluation_payload + + +def test_nested_trace_builds_expected_step(): + """A trace with span children and a leaf retriever builds the matching tree.""" + with record_step("trace", "banking_trace", input={"request": "r1"}) as trace: + with trace.child("retriever", "policy_lookup", input="refund policy?") as retriever: + retriever.output = {"hits": 1} + trace.output = {"status": "planned"} + + step = trace.build() + + assert step.type == "trace" + assert step.name == "banking_trace" + assert step.input == {"request": "r1"} + assert step.output == {"status": "planned"} + assert step.children is not None + assert len(step.children) == 1 + child = step.children[0] + assert child.type == "retriever" + assert child.name == "policy_lookup" + assert child.input == "refund policy?" + assert child.output == {"hits": 1} + + +def test_session_nests_traces(): + """Sessions nest traces the same way traces nest spans.""" + with record_step("session", "banking_session") as session: + with session.child("trace", "turn_1") as trace: + with trace.child("llm", "respond") as llm: + llm.output = "hi" + + step = session.build() + + assert step.type == "session" + assert len(step.children) == 1 + assert step.children[0].type == "trace" + assert step.children[0].name == "turn_1" + assert len(step.children[0].children) == 1 + assert step.children[0].children[0].type == "llm" + + +def test_call_matches_decorator_payload_for_tool(): + """StepRecorder.call() must capture the same payload _create_evaluation_payload would.""" + + def lookup_account(account_id: str) -> dict: + return {"balance": 100} + + with record_step("trace", "t") as trace: + result = trace.call( + lookup_account, account_id="acct-1", step_type="tool", step_name="lookup_account" + ) + + assert result == {"balance": 100} + expected_payload = _create_evaluation_payload( + lookup_account, (), {"account_id": "acct-1"}, result, "lookup_account", "tool", None + ) + [child] = trace.build().children + assert child.model_dump(exclude_none=True) == { + k: v for k, v in expected_payload.items() if v is not None + } + + +@pytest.mark.asyncio +async def test_acall_runs_async_function_and_records_child(): + """StepRecorder.acall() awaits the function and records the result.""" + + async def run_banking_model(prompt: str) -> str: + return f"answer: {prompt}" + + with record_step("trace", "t") as trace: + result = await trace.acall( + run_banking_model, "hello", step_type="llm", step_name="respond" + ) + + assert result == "answer: hello" + [child] = trace.build().children + assert child.type == "llm" + assert child.name == "respond" + assert child.input == "hello" + assert child.output == "answer: hello" + + +def test_call_records_exception_and_reraises(): + """A failing call is still recorded as a child (with the error in context), then re-raised.""" + + def flaky(x: int) -> int: + raise ValueError("boom") + + with record_step("trace", "t") as trace: + with pytest.raises(ValueError, match="boom"): + trace.call(flaky, 1, step_type="tool", step_name="flaky") + + [child] = trace.build().children + assert child.name == "flaky" + assert child.output is None + assert child.context == {"error": "ValueError('boom')"} + + +def test_empty_trace_and_session_build_empty_children_list(): + """A trace/session with no recorded children gets children=[], not None.""" + with record_step("trace", "empty_trace") as trace: + pass + with record_step("session", "empty_session") as session: + pass + + assert trace.build().children == [] + assert session.build().children == [] + + +def test_llm_step_without_children_stays_none(): + """Non trace/session step types are untouched when no children were recorded.""" + with record_step("llm", "respond", input="hi") as leaf: + leaf.output = "hello" + + assert leaf.build().children is None + + +def test_add_attaches_prebuilt_step_or_dict(): + """add() accepts both a real Step and a plain dict (coerced via Step.model_validate).""" + with record_step("trace", "t") as trace: + trace.add(Step(type="llm", name="a", input="x", output="y")) + trace.add({"type": "tool", "name": "b", "input": {}, "output": {}}) + + children = trace.build().children + assert [c.name for c in children] == ["a", "b"] + assert all(isinstance(c, Step) for c in children) + + +@pytest.mark.asyncio +async def test_evaluate_forwards_to_evaluate_step(): + """StepRecorder.evaluate() builds the Step and delegates to evaluate_step().""" + mock_result = EvaluationResult(is_safe=True, confidence=1.0) + mock_evaluate_step = AsyncMock(return_value=mock_result) + + with record_step("trace", "t", input="hi") as trace: + pass + + with patch.object(step_recorder, "evaluate_step", mock_evaluate_step): + result = await trace.evaluate(agent_name="test-bot", stage="post") + + assert result is mock_result + mock_evaluate_step.assert_awaited_once() + _, kwargs = mock_evaluate_step.call_args + assert kwargs["agent_name"] == "test-bot" + assert kwargs["stage"] == "post" + + +@pytest.mark.asyncio +async def test_evaluate_defaults_agent_name_from_current_agent(): + """When agent_name is omitted, evaluate() falls back to the agent registered via init().""" + mock_result = EvaluationResult(is_safe=True, confidence=1.0) + mock_evaluate_step = AsyncMock(return_value=mock_result) + + with record_step("trace", "t") as trace: + pass + + class _FakeAgent: + agent_name = "test-bot-0123456789" + + with patch.object(step_recorder, "evaluate_step", mock_evaluate_step): + with patch.object(step_recorder.state, "current_agent", _FakeAgent()): + await trace.evaluate(stage="post") + + _, kwargs = mock_evaluate_step.call_args + assert kwargs["agent_name"] == "test-bot-0123456789" + + +@pytest.mark.asyncio +async def test_evaluate_without_agent_name_or_init_raises(): + """Without an explicit agent_name or a prior init(), evaluate() fails loudly.""" + with record_step("trace", "t") as trace: + pass + + with patch.object(step_recorder.state, "current_agent", None): + with pytest.raises(RuntimeError, match="agent_name not supplied"): + await trace.evaluate(stage="post") diff --git a/sdks/python/tests/test_step_recorder_galileo_integration.py b/sdks/python/tests/test_step_recorder_galileo_integration.py new file mode 100644 index 00000000..6b6d3bc2 --- /dev/null +++ b/sdks/python/tests/test_step_recorder_galileo_integration.py @@ -0,0 +1,87 @@ +"""Cross-package check: StepRecorder trees feed the Galileo record factory. + +The galileo extras are normally installed in the dev environment (see +``test_evaluators_optional_imports.py``), so this skips cleanly when they +are not available instead of failing the suite. +""" + +from __future__ import annotations + +import importlib.util + +import pytest + +from agent_control import record_step + + +def _module_available(name: str) -> bool: + try: + return importlib.util.find_spec(name) is not None + except (ImportError, ValueError): + return False + + +_GALILEO_INSTALLED = _module_available("agent_control_evaluator_galileo.records") + +pytestmark = pytest.mark.skipif( + not _GALILEO_INSTALLED, + reason="agent-control-evaluator-galileo extras not installed in this environment", +) + + +def _record_from_step(step): + from agent_control_evaluator_galileo.records.factory import record_from_step + + return record_from_step(step) + + +def test_recorder_built_trace_matches_demo_shape(): + """A trace with one llm/tool/retriever child each builds a 3-span Trace.""" + + def policy_search(query: str) -> dict: + return {"docs": ["policy-1"]} + + def account_lookup(account_id: str) -> dict: + return {"balance": 500} + + def banking_llm(prompt: str) -> str: + return f"response: {prompt}" + + with record_step("trace", "banking_trace", input={"request": "r1"}) as trace: + trace.call( + policy_search, query="refund policy", step_type="retriever", step_name="policy_search" + ) + trace.call( + account_lookup, account_id="acct-1", step_type="tool", step_name="account_lookup" + ) + trace.call(banking_llm, "draft a reply", step_type="llm", step_name="banking_llm") + trace.output = {"status": "done"} + + record = _record_from_step(trace.build()) + + assert type(record).__name__ == "Trace" + assert len(record.spans) == 3 + assert {type(span).__name__ for span in record.spans} == { + "LlmSpan", + "ToolSpan", + "RetrieverSpan", + } + + +def test_recorder_built_session_matches_demo_shape(): + """A session with 2 traces, each with >=2 spans, builds the matching Session.""" + with record_step("session", "banking_session") as session: + for i in range(2): + with session.child("trace", f"turn_{i}", input={"turn": i}) as trace: + with trace.child("llm", "respond", input="hi") as llm: + llm.output = "hello" + with trace.child("tool", "lookup", input={}) as tool: + tool.output = {} + trace.output = {"turn": i, "status": "ok"} + + record = _record_from_step(session.build()) + + assert type(record).__name__ == "Session" + assert len(record.traces) == 2 + for sub_trace in record.traces: + assert len(sub_trace.spans) >= 2 From c6ba3927b70950782405429e8d2cfd176a513ee3 Mon Sep 17 00:00:00 2001 From: Yixin Date: Wed, 7 Oct 2026 15:14:31 -0700 Subject: [PATCH 2/6] test(sdk): cover StepRecorder.acall() exceptions and build() passthrough Codecov flagged 6 uncovered lines in step_recorder.py: acall()'s exception-recording path (mirrored from call() but never exercised) and build()'s context/tools/ground_truth passthrough branches. Add the two missing tests; step_recorder.py is now at 100% line coverage. Co-Authored-By: Claude Code --- sdks/python/tests/test_step_recorder.py | 35 +++++++++++++++++++++++++ 1 file changed, 35 insertions(+) diff --git a/sdks/python/tests/test_step_recorder.py b/sdks/python/tests/test_step_recorder.py index 44ec33c2..ca0e7e6e 100644 --- a/sdks/python/tests/test_step_recorder.py +++ b/sdks/python/tests/test_step_recorder.py @@ -105,6 +105,41 @@ def flaky(x: int) -> int: assert child.context == {"error": "ValueError('boom')"} +@pytest.mark.asyncio +async def test_acall_records_exception_and_reraises(): + """acall()'s failure path mirrors call(): record (error in context), then re-raise.""" + + async def flaky(x: int) -> int: + raise ValueError("boom") + + with record_step("trace", "t") as trace: + with pytest.raises(ValueError, match="boom"): + await trace.acall(flaky, 1, step_type="tool", step_name="flaky") + + [child] = trace.build().children + assert child.name == "flaky" + assert child.output is None + assert child.context == {"error": "ValueError('boom')"} + + +def test_build_includes_context_tools_ground_truth(): + """build() carries context/tools/ground_truth through when supplied.""" + with record_step( + "llm", + "respond", + input="hi", + context={"locale": "en-US"}, + tools=[{"type": "function", "function": {"name": "search"}}], + ground_truth="hello", + ) as leaf: + leaf.output = "hello" + + step = leaf.build() + assert step.context == {"locale": "en-US"} + assert step.tools == [{"type": "function", "function": {"name": "search"}}] + assert step.ground_truth == "hello" + + def test_empty_trace_and_session_build_empty_children_list(): """A trace/session with no recorded children gets children=[], not None.""" with record_step("trace", "empty_trace") as trace: From fb631825538e8ad6c01d84a7caa22aeeb1f45cea Mon Sep 17 00:00:00 2001 From: Yixin Date: Thu, 8 Oct 2026 17:57:55 -0700 Subject: [PATCH 3/6] fix(sdk): address PR review feedback on StepRecorder OSS/Galileo boundary - record_step()/.child() take an opt-in container_types param instead of hardcoding "trace"/"session" as the step types that serialize empty children as [] - step types are caller-defined strings, not a core vocabulary. - Move the galileo-dependent StepRecorder integration test out of the SDK into evaluators/contrib/galileo, where its dependency belongs. - Replace the README's trace/session example with a generic workflow/lookup one so the OSS SDK docs don't lean on Galileo's hierarchy vocabulary. Co-Authored-By: Claude Code --- evaluators/contrib/galileo/pyproject.toml | 6 +++ .../test_step_recorder_galileo_integration.py | 33 +----------- sdks/python/README.md | 52 +++++++++++-------- .../python/src/agent_control/step_recorder.py | 35 +++++++++---- sdks/python/tests/test_step_recorder.py | 35 ++++++++++--- 5 files changed, 91 insertions(+), 70 deletions(-) rename {sdks/python => evaluators/contrib/galileo}/tests/test_step_recorder_galileo_integration.py (71%) diff --git a/evaluators/contrib/galileo/pyproject.toml b/evaluators/contrib/galileo/pyproject.toml index fbea1f3a..4d23e030 100644 --- a/evaluators/contrib/galileo/pyproject.toml +++ b/evaluators/contrib/galileo/pyproject.toml @@ -16,6 +16,9 @@ dependencies = [ [project.optional-dependencies] dev = [ + "agent-control-sdk>=8.11.0", + "agent-control-telemetry>=8.6.0", + "agent-control-engine>=8.6.0", "pytest>=8.0.0", "pytest-asyncio>=0.23.0", "pytest-cov>=4.0.0", @@ -37,3 +40,6 @@ packages = ["src/agent_control_evaluator_galileo"] [tool.uv.sources] agent-control-evaluators = { path = "../../builtin", editable = true } agent-control-models = { path = "../../../models", editable = true } +agent-control-sdk = { path = "../../../sdks/python", editable = true } +agent-control-telemetry = { path = "../../../telemetry", editable = true } +agent-control-engine = { path = "../../../engine", editable = true } diff --git a/sdks/python/tests/test_step_recorder_galileo_integration.py b/evaluators/contrib/galileo/tests/test_step_recorder_galileo_integration.py similarity index 71% rename from sdks/python/tests/test_step_recorder_galileo_integration.py rename to evaluators/contrib/galileo/tests/test_step_recorder_galileo_integration.py index 6b6d3bc2..f0d6b8b6 100644 --- a/sdks/python/tests/test_step_recorder_galileo_integration.py +++ b/evaluators/contrib/galileo/tests/test_step_recorder_galileo_integration.py @@ -1,38 +1,9 @@ -"""Cross-package check: StepRecorder trees feed the Galileo record factory. - -The galileo extras are normally installed in the dev environment (see -``test_evaluators_optional_imports.py``), so this skips cleanly when they -are not available instead of failing the suite. -""" +"""Cross-package check: StepRecorder trees feed the Galileo record factory.""" from __future__ import annotations -import importlib.util - -import pytest - from agent_control import record_step - - -def _module_available(name: str) -> bool: - try: - return importlib.util.find_spec(name) is not None - except (ImportError, ValueError): - return False - - -_GALILEO_INSTALLED = _module_available("agent_control_evaluator_galileo.records") - -pytestmark = pytest.mark.skipif( - not _GALILEO_INSTALLED, - reason="agent-control-evaluator-galileo extras not installed in this environment", -) - - -def _record_from_step(step): - from agent_control_evaluator_galileo.records.factory import record_from_step - - return record_from_step(step) +from agent_control_evaluator_galileo.records.factory import record_from_step as _record_from_step def test_recorder_built_trace_matches_demo_shape(): diff --git a/sdks/python/README.md b/sdks/python/README.md index 94785665..980cce6d 100644 --- a/sdks/python/README.md +++ b/sdks/python/README.md @@ -84,53 +84,59 @@ async with agent_control.AgentControlClient() as client: The existing `evaluate_controls` helper remains available for callers that prefer its field-based convenience arguments. -## Building trace/session steps +## Building a step with children -Trace- and session-level controls evaluate the aggregate of several -already-executed child steps, so `Step.children` has to be populated by the -caller. Building that tree by hand means hand-rolling a side-channel record -for every leaf call and converting it into `Step` objects afterward: +A control's scope can target any step type, including one that aggregates +several already-executed child steps - `Step.children` just has to be +populated by the caller. Building that tree by hand means hand-rolling a +side-channel record for every leaf call and converting it into `Step` +objects afterward: ```python # Before: a hand-built dict record next to the real return value def execute_policy_search(query): docs = search_policy_documents(query) - record = {"type": "retriever", "query": query, "docs": docs} + record = {"type": "lookup", "query": query, "docs": docs} return StepExecution(value=docs, record=record) -spans = [] +children = [] result = execute_policy_search(query) -spans.append(result.record) +children.append(result.record) ... -trace_step = build_agent_control_step({"type": "trace", "spans": spans, ...}) +workflow_step = build_agent_control_step({"type": "workflow", "children": children, ...}) ``` `agent_control.record_step()` builds the same tree incrementally, in `Step` vocabulary, with no intermediate dict and no converter: ```python -with agent_control.record_step("trace", "banking_trace", input={"request": req}) as trace: - with trace.child("retriever", "policy_lookup", input=query) as span: +with agent_control.record_step("workflow", "banking_workflow", input={"request": req}) as workflow: + with workflow.child("lookup", "policy_lookup", input=query) as span: span.output = search_policy_documents(query) - account = trace.call(lookup_account, account_id="acct-1001", step_type="tool") - plan = trace.call(run_banking_model, req, step_type="llm", tools=TOOL_DEFINITIONS) - trace.output = {"status": "planned", "message": plan["content"]} + account = workflow.call(lookup_account, account_id="acct-1001", step_type="tool") + plan = workflow.call(run_banking_model, req, step_type="llm", tools=TOOL_DEFINITIONS) + workflow.output = {"status": "planned", "message": plan["content"]} -result = await trace.evaluate(stage="post") +result = await workflow.evaluate(stage="post") ``` -`trace.call(...)`/`await trace.acall(...)` run the function and record it as -a child using the same capture logic as `@control()` (input from bound +`workflow.call(...)`/`await workflow.acall(...)` run the function and record +it as a child using the same capture logic as `@control()` (input from bound arguments, output from the return value), then return the real result - removing the need for a separate `StepExecution`-style wrapper. A failed call is still recorded (with the error in `context`) before the exception is -re-raised. `trace.child(...)` nests another recorder the same way, so -sessions nest traces with `session.child("trace", ...)`. - -`trace.build()` produces the frozen `Step`; `trace.evaluate(...)` builds it -and evaluates it in one call via `evaluate_step()` - the same function -`evaluate_controls()` uses internally once its `Step` is built. +re-raised. `workflow.child(...)` nests another recorder the same way, so a +parent step can nest children of any depth, e.g. `session.child("trace", +...)` for callers that model multi-turn sessions of traces. + +`workflow.build()` produces the frozen `Step`, with `children=[]` instead of +omitted (`None`) for a step type passed via `container_types` even when no +children were recorded - e.g. `record_step("workflow", ..., container_types= +{"workflow"})` for a workflow that legitimately ran with zero children. +`workflow.evaluate(...)` builds the step and evaluates it in one call via +`evaluate_step()` - the same function `evaluate_controls()` uses internally +once its `Step` is built. ## Sharing an OpenTelemetry provider with Google ADK diff --git a/sdks/python/src/agent_control/step_recorder.py b/sdks/python/src/agent_control/step_recorder.py index 160d5ed0..a4376cb8 100644 --- a/sdks/python/src/agent_control/step_recorder.py +++ b/sdks/python/src/agent_control/step_recorder.py @@ -9,7 +9,7 @@ from __future__ import annotations -from collections.abc import Awaitable, Callable, Mapping +from collections.abc import Awaitable, Callable, Collection, Mapping from typing import Any, Literal, TypeVar from agent_control_models import EvaluationResult, JSONObject, JSONValue, Step @@ -21,12 +21,6 @@ T = TypeVar("T") -# Step types whose children are reported as a span/trace tree; an absent -# recording means "nothing happened", which the record factory treats -# differently from "not applicable" (None). See models/agent.py and -# records/factory.py for the trace/session child-shape rules this mirrors. -_EMPTY_CHILDREN_STEP_TYPES = frozenset({"trace", "session"}) - def record_step( step_type: str, @@ -36,8 +30,18 @@ def record_step( context: dict[str, Any] | None = None, tools: list[JSONObject] | None = None, ground_truth: JSONValue | None = None, + container_types: Collection[str] = (), ) -> StepRecorder: - """Start building a ``Step`` tree. Use as a context manager.""" + """Start building a ``Step`` tree. Use as a context manager. + + ``container_types`` names step types whose ``children`` should be + serialized as ``[]`` rather than omitted when none were recorded - e.g. + a trace/session that legitimately ran with zero children, as opposed to + a leaf type for which children isn't a meaningful concept at all. Step + types are caller-defined strings, so this isn't assumed; opt in with + e.g. ``container_types={"trace", "session"}``. Inherited by nested + ``.child()`` recorders unless overridden there. + """ return StepRecorder( step_type, name, @@ -45,6 +49,7 @@ def record_step( context=context, tools=tools, ground_truth=ground_truth, + container_types=frozenset(container_types), ) @@ -64,6 +69,7 @@ def __init__( context: dict[str, Any] | None = None, tools: list[JSONObject] | None = None, ground_truth: JSONValue | None = None, + container_types: frozenset[str] = frozenset(), _parent: StepRecorder | None = None, ) -> None: self.type = ensure_step_type(step_type) @@ -75,6 +81,7 @@ def __init__( self.ground_truth = ground_truth self._children: list[Step] = [] self._parent = _parent + self._container_types = container_types def __enter__(self) -> StepRecorder: return self @@ -92,8 +99,13 @@ def child( context: dict[str, Any] | None = None, tools: list[JSONObject] | None = None, ground_truth: JSONValue | None = None, + container_types: Collection[str] | None = None, ) -> StepRecorder: - """Start a nested recorder; it attaches to this node on ``__exit__``.""" + """Start a nested recorder; it attaches to this node on ``__exit__``. + + Inherits this node's ``container_types`` unless ``container_types`` + is given explicitly. + """ return StepRecorder( step_type, name, @@ -101,6 +113,9 @@ def child( context=context, tools=tools, ground_truth=ground_truth, + container_types=( + self._container_types if container_types is None else frozenset(container_types) + ), _parent=self, ) @@ -178,7 +193,7 @@ def build(self) -> Step: step_dict["ground_truth"] = self.ground_truth if self._children: step_dict["children"] = list(self._children) - elif self.type in _EMPTY_CHILDREN_STEP_TYPES: + elif self.type in self._container_types: step_dict["children"] = [] return Step(**step_dict) # type: ignore[arg-type] diff --git a/sdks/python/tests/test_step_recorder.py b/sdks/python/tests/test_step_recorder.py index ca0e7e6e..03ffaa6f 100644 --- a/sdks/python/tests/test_step_recorder.py +++ b/sdks/python/tests/test_step_recorder.py @@ -140,25 +140,48 @@ def test_build_includes_context_tools_ground_truth(): assert step.ground_truth == "hello" -def test_empty_trace_and_session_build_empty_children_list(): - """A trace/session with no recorded children gets children=[], not None.""" - with record_step("trace", "empty_trace") as trace: +def test_container_types_opts_in_to_empty_children_list(): + """A declared container type with no recorded children gets children=[], not None.""" + with record_step("trace", "empty_trace", container_types={"trace", "session"}) as trace: pass - with record_step("session", "empty_session") as session: + with record_step("session", "empty_session", container_types={"trace", "session"}) as session: pass assert trace.build().children == [] assert session.build().children == [] -def test_llm_step_without_children_stays_none(): - """Non trace/session step types are untouched when no children were recorded.""" +def test_step_type_without_container_types_stays_none(): + """Step types are caller-defined strings; children stays None unless the + caller opts a type into container_types - no type is a container by default.""" + with record_step("trace", "empty_trace") as trace: + pass with record_step("llm", "respond", input="hi") as leaf: leaf.output = "hello" + assert trace.build().children is None assert leaf.build().children is None +def test_child_inherits_parent_container_types(): + """A nested .child() recorder inherits the parent's container_types unless overridden.""" + with record_step("session", "s", container_types={"trace", "session"}) as session: + with session.child("trace", "empty_trace") as trace: + pass + + assert trace.build().children == [] + assert session.build().children == [trace.build()] + + +def test_child_can_override_container_types(): + """A nested .child() recorder can override the inherited container_types.""" + with record_step("session", "s", container_types={"trace", "session"}) as session: + with session.child("trace", "empty_trace", container_types=()) as trace: + pass + + assert trace.build().children is None + + def test_add_attaches_prebuilt_step_or_dict(): """add() accepts both a real Step and a plain dict (coerced via Step.model_validate).""" with record_step("trace", "t") as trace: From 7f2991484f8ca072743d8d5d5a01f76f3543d194 Mon Sep 17 00:00:00 2001 From: Yixin Date: Thu, 8 Oct 2026 18:13:55 -0700 Subject: [PATCH 4/6] fix(ci): install galileo evaluator's dev extra when running its tests The moved StepRecorder integration test needs agent-control-sdk, which only lives in the galileo evaluator's dev extra. The test target never requested that extra, so CI installed a venv without it and failed to import agent_control. Co-Authored-By: Claude Code --- evaluators/contrib/galileo/Makefile | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/evaluators/contrib/galileo/Makefile b/evaluators/contrib/galileo/Makefile index 89d61ac0..bfc15c91 100644 --- a/evaluators/contrib/galileo/Makefile +++ b/evaluators/contrib/galileo/Makefile @@ -16,7 +16,7 @@ sync: uv sync test: - uv run --with pytest --with pytest-asyncio --with pytest-cov --package $(PACKAGE) pytest tests --cov=src --cov-report=xml:../../../coverage-evaluators-galileo.xml -q + uv run --extra dev --with pytest --with pytest-asyncio --with pytest-cov --package $(PACKAGE) pytest tests --cov=src --cov-report=xml:../../../coverage-evaluators-galileo.xml -q lint: uv run --with ruff --package $(PACKAGE) ruff check --config ../../../pyproject.toml src/ From 31da9f2b8aaa1fa2ccb5fca5a9ee62d901d0048e Mon Sep 17 00:00:00 2001 From: Yixin Date: Fri, 9 Oct 2026 08:25:42 -0700 Subject: [PATCH 5/6] feat(sdk): add step_context to StepRecorder.call()/acall() call()/acall() run func(*args, **kwargs) and forward every unrecognized keyword straight to func, so a plain `context=` parameter on call() would collide with (and silently steal) a same-named kwarg on func - exactly the case for functions like run_banking_model(..., context=...). step_context is merged into the recorded child's Step.context instead, without touching what gets forwarded to func. On failure, the error key added to context takes precedence over the same key in step_context. Co-Authored-By: Claude Code --- .../python/src/agent_control/step_recorder.py | 33 +++++++-- sdks/python/tests/test_step_recorder.py | 67 +++++++++++++++++++ 2 files changed, 94 insertions(+), 6 deletions(-) diff --git a/sdks/python/src/agent_control/step_recorder.py b/sdks/python/src/agent_control/step_recorder.py index a4376cb8..60676454 100644 --- a/sdks/python/src/agent_control/step_recorder.py +++ b/sdks/python/src/agent_control/step_recorder.py @@ -132,13 +132,17 @@ def _record_call_result( step_name: str | None, step_type: str | None, tools: ToolsConfiguration, + step_context: dict[str, Any] | None, error: BaseException | None, ) -> None: payload = _create_evaluation_payload( func, args, kwargs, output, step_name, step_type, tools ) + context = {**(payload.get("context") or {}), **(step_context or {})} if error is not None: - payload["context"] = {**(payload.get("context") or {}), "error": repr(error)} + context["error"] = repr(error) + if context: + payload["context"] = context self._children.append(Step.model_validate(payload)) def call( @@ -148,15 +152,27 @@ def call( step_type: str | None = None, step_name: str | None = None, tools: ToolsConfiguration = None, + step_context: dict[str, Any] | None = None, **kwargs: Any, ) -> T: - """Run ``func``, record it via ``@control()``'s capture logic, return its output.""" + """Run ``func``, record it via ``@control()``'s capture logic, return its output. + + ``step_context`` is merged into the recorded child's ``context`` (an + ``error`` key added on failure takes precedence over the same key in + ``step_context``). It is named ``step_context`` rather than ``context`` + so it can't collide with a ``context`` keyword argument on ``func`` + itself, which ``**kwargs`` would otherwise forward unchanged. + """ try: output = func(*args, **kwargs) except Exception as exc: - self._record_call_result(func, args, kwargs, None, step_name, step_type, tools, exc) + self._record_call_result( + func, args, kwargs, None, step_name, step_type, tools, step_context, exc + ) raise - self._record_call_result(func, args, kwargs, output, step_name, step_type, tools, None) + self._record_call_result( + func, args, kwargs, output, step_name, step_type, tools, step_context, None + ) return output async def acall( @@ -166,15 +182,20 @@ async def acall( step_type: str | None = None, step_name: str | None = None, tools: ToolsConfiguration = None, + step_context: dict[str, Any] | None = None, **kwargs: Any, ) -> T: """Async counterpart to :meth:`call`.""" try: output = await func(*args, **kwargs) except Exception as exc: - self._record_call_result(func, args, kwargs, None, step_name, step_type, tools, exc) + self._record_call_result( + func, args, kwargs, None, step_name, step_type, tools, step_context, exc + ) raise - self._record_call_result(func, args, kwargs, output, step_name, step_type, tools, None) + self._record_call_result( + func, args, kwargs, output, step_name, step_type, tools, step_context, None + ) return output def build(self) -> Step: diff --git a/sdks/python/tests/test_step_recorder.py b/sdks/python/tests/test_step_recorder.py index 03ffaa6f..ba083838 100644 --- a/sdks/python/tests/test_step_recorder.py +++ b/sdks/python/tests/test_step_recorder.py @@ -1,5 +1,6 @@ """Tests for StepRecorder, the incremental Step tree builder.""" +from typing import Any from unittest.mock import AsyncMock, patch import pytest @@ -89,6 +90,72 @@ async def run_banking_model(prompt: str) -> str: assert child.output == "answer: hello" +def test_call_step_context_is_recorded(): + """step_context is merged into the recorded child's context.""" + + def lookup_account(account_id: str) -> dict: + return {"balance": 100} + + with record_step("trace", "t") as trace: + trace.call( + lookup_account, + account_id="acct-1", + step_type="tool", + step_name="lookup_account", + step_context={"executed_function": "lookup_account"}, + ) + + [child] = trace.build().children + assert child.context == {"executed_function": "lookup_account"} + + +def test_call_step_context_does_not_collide_with_func_context_kwarg(): + """step_context is merged into Step.context, not forwarded to func - a plain + `context` kwarg on call() would instead be swallowed by **kwargs and sent + to func, silently dropping the caller's intended Step metadata.""" + + received: dict[str, Any] = {} + + def run_banking_model(prompt: str, *, context: dict[str, Any] | None = None) -> str: + received["context"] = context + return f"answer: {prompt}" + + with record_step("trace", "t") as trace: + result = trace.call( + run_banking_model, + "hello", + context={"account": "acct-1"}, + step_type="llm", + step_name="respond", + step_context={"executed_function": "run_banking_model"}, + ) + + assert result == "answer: hello" + assert received["context"] == {"account": "acct-1"} + [child] = trace.build().children + assert child.context == {"executed_function": "run_banking_model"} + + +def test_call_step_context_and_error_merge_on_failure(): + """On failure, the error key is added to step_context's dict and wins on collision.""" + + def flaky(x: int) -> int: + raise ValueError("boom") + + with record_step("trace", "t") as trace: + with pytest.raises(ValueError, match="boom"): + trace.call( + flaky, + 1, + step_type="tool", + step_name="flaky", + step_context={"attempt": 1, "error": "overwritten"}, + ) + + [child] = trace.build().children + assert child.context == {"attempt": 1, "error": "ValueError('boom')"} + + def test_call_records_exception_and_reraises(): """A failing call is still recorded as a child (with the error in context), then re-raised.""" From cc83d63c595c8872846dde611c0e2a4cc7618615 Mon Sep 17 00:00:00 2001 From: Yixin Date: Mon, 12 Oct 2026 17:20:40 -0700 Subject: [PATCH 6/6] fix(step_recorder): default tool-step input and record errors on exit - build() defaults a tool step's input to {} when omitted, matching evaluate_controls()'s existing behavior and avoiding Step's "tool steps require object input" validation error. - __exit__ now receives the real exception value and merges {"error": repr(exc)} into context, consistent with call()/acall()'s failure path, so a child() block that raises is no longer indistinguishable from one that simply has no output. Addresses PR #287 review comments. Co-Authored-By: Claude Code --- .../python/src/agent_control/step_recorder.py | 15 ++++- sdks/python/tests/test_step_recorder.py | 64 +++++++++++++++++++ 2 files changed, 76 insertions(+), 3 deletions(-) diff --git a/sdks/python/src/agent_control/step_recorder.py b/sdks/python/src/agent_control/step_recorder.py index 60676454..1fe8a66a 100644 --- a/sdks/python/src/agent_control/step_recorder.py +++ b/sdks/python/src/agent_control/step_recorder.py @@ -12,7 +12,7 @@ from collections.abc import Awaitable, Callable, Collection, Mapping from typing import Any, Literal, TypeVar -from agent_control_models import EvaluationResult, JSONObject, JSONValue, Step +from agent_control_models import STEP_TYPE_TOOL, EvaluationResult, JSONObject, JSONValue, Step from ._state import state from .control_decorators import ToolsConfiguration, _create_evaluation_payload @@ -86,7 +86,11 @@ def __init__( def __enter__(self) -> StepRecorder: return self - def __exit__(self, *exc_info: object) -> None: + def __exit__( + self, exc_type: object, exc_value: BaseException | None, traceback: object + ) -> None: + if exc_value is not None: + self.context = {**(self.context or {}), "error": repr(exc_value)} if self._parent is not None: self._parent._children.append(self.build()) @@ -200,10 +204,15 @@ async def acall( def build(self) -> Step: """Recursively construct the frozen ``Step`` for this node and its children.""" + input_value = self.input + if input_value is None and self.type == STEP_TYPE_TOOL: + # Tool steps require object input (Step.validate_builtin_types); default to + # {} so an omitted input doesn't raise, matching evaluate_controls(). + input_value = {} step_dict: dict[str, Any] = { "type": self.type, "name": self.name, - "input": self.input, + "input": input_value, "output": self.output, } if self.context is not None: diff --git a/sdks/python/tests/test_step_recorder.py b/sdks/python/tests/test_step_recorder.py index ba083838..f8a67f70 100644 --- a/sdks/python/tests/test_step_recorder.py +++ b/sdks/python/tests/test_step_recorder.py @@ -189,6 +189,43 @@ async def flaky(x: int) -> int: assert child.context == {"error": "ValueError('boom')"} +def test_child_block_exception_is_recorded_in_context_and_reraised(): + """A raising child() block still attaches to its parent, with the error + recorded in context - consistent with call()/acall()'s failure path.""" + with record_step("trace", "t") as trace: + with pytest.raises(ValueError, match="boom"): + with trace.child("tool", "flaky", input={}): + raise ValueError("boom") + + [child] = trace.build().children + assert child.name == "flaky" + assert child.context == {"error": "ValueError('boom')"} + + +def test_child_block_exception_merges_with_existing_context(): + """An error added by a raising child() block merges into, rather than + replacing, any context already supplied, and wins on key collision.""" + with record_step("trace", "t") as trace: + with pytest.raises(ValueError, match="boom"): + with trace.child( + "tool", "flaky", input={}, context={"attempt": 1, "error": "overwritten"} + ): + raise ValueError("boom") + + [child] = trace.build().children + assert child.context == {"attempt": 1, "error": "ValueError('boom')"} + + +def test_top_level_block_exception_is_recorded_in_context(): + """A top-level record_step() block (no parent) gets the same error-in-context + treatment, so a caller who catches and later calls .build() can see it.""" + with pytest.raises(ValueError, match="boom"): + with record_step("trace", "t") as trace: + raise ValueError("boom") + + assert trace.build().context == {"error": "ValueError('boom')"} + + def test_build_includes_context_tools_ground_truth(): """build() carries context/tools/ground_truth through when supplied.""" with record_step( @@ -207,6 +244,33 @@ def test_build_includes_context_tools_ground_truth(): assert step.ground_truth == "hello" +def test_build_defaults_missing_tool_input_to_empty_dict(): + """A tool step with no input omitted would otherwise fail Step's + `tool steps require object input` validation; build() defaults it to {}, + matching evaluate_controls()'s existing behavior for tool-type input.""" + with record_step("tool", "lookup") as tool: + tool.output = {"balance": 100} + + step = tool.build() + assert step.input == {} + + +def test_build_preserves_explicit_tool_input(): + """An explicitly supplied tool input is never overridden by the default.""" + with record_step("tool", "lookup", input={"account_id": "acct-1"}) as tool: + tool.output = {"balance": 100} + + assert tool.build().input == {"account_id": "acct-1"} + + +def test_build_does_not_default_input_for_non_tool_types(): + """The {} default is scoped to tool steps; other step types keep input=None.""" + with record_step("llm", "respond") as leaf: + leaf.output = "hello" + + assert leaf.build().input is None + + def test_container_types_opts_in_to_empty_children_list(): """A declared container type with no recorded children gets children=[], not None.""" with record_step("trace", "empty_trace", container_types={"trace", "session"}) as trace: