diff --git a/models/src/agent_control_models/agent.py b/models/src/agent_control_models/agent.py index e60573c9..aa18e4f3 100644 --- a/models/src/agent_control_models/agent.py +++ b/models/src/agent_control_models/agent.py @@ -168,6 +168,11 @@ class Step(BaseModel): ground_truth: JSONValue | None = Field( None, description="Optional expected or reference output for this step" ) + children: list[Step] | None = Field( + None, + exclude_if=lambda value: value is None, + description="Optional child steps associated with this invocation", + ) @field_validator("type") @classmethod @@ -182,3 +187,6 @@ def validate_builtin_types(self) -> Step: if not isinstance(self.input, dict): raise ValueError("tool steps require object input") return self + + +Step.model_rebuild() diff --git a/sdks/python/src/agent_control/evaluation.py b/sdks/python/src/agent_control/evaluation.py index 3a8ec91b..800b4bf4 100644 --- a/sdks/python/src/agent_control/evaluation.py +++ b/sdks/python/src/agent_control/evaluation.py @@ -523,6 +523,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, step_type: str = "llm", stage: Literal["pre", "post"] = "pre", agent_name: str, @@ -561,6 +562,8 @@ async def evaluate_controls( step_dict["tools"] = tools if ground_truth is not None: step_dict["ground_truth"] = ground_truth + if children is not None: + step_dict["children"] = children step_obj = Step(**step_dict) # type: ignore[arg-type] resolved_controls = state.server_controls or [] diff --git a/sdks/python/tests/test_evaluation.py b/sdks/python/tests/test_evaluation.py index d598c118..f4d5915d 100644 --- a/sdks/python/tests/test_evaluation.py +++ b/sdks/python/tests/test_evaluation.py @@ -6,6 +6,7 @@ import pytest from agent_control import evaluation from agent_control.evaluation import EvaluationResult +from agent_control_models import Step from pydantic import ValidationError @@ -228,6 +229,85 @@ async def test_evaluate_controls_preserves_explicit_tools_and_ground_truth(monke assert step.ground_truth == {"answer": "expected"} +def test_step_legacy_serialization_is_unchanged_without_children(): + """A Step without children keeps the existing serialized payload.""" + step = Step(type="llm", name="chat", input="hello") + + assert step.model_dump(mode="json") == { + "type": "llm", + "name": "chat", + "input": "hello", + "output": None, + "context": None, + "tools": None, + "ground_truth": None, + } + assert Step(type="llm", name="chat", input="hello", children=None).children is None + + +def test_step_accepts_and_preserves_recursive_children(): + """Nested child steps retain all provider-neutral runtime fields.""" + child_payload = { + "type": "llm", + "name": "child", + "input": {"prompt": "question"}, + "output": "answer", + "context": {"conversation": ["earlier"]}, + "tools": [{"name": "lookup", "input_schema": {"type": "object"}}], + "ground_truth": {"answer": "expected"}, + "children": [ + {"type": "retriever", "name": "search", "input": {"query": "q"}} + ], + } + parent = Step(type="llm", name="parent", input="hello", children=[child_payload]) + + assert parent.children is not None + child = parent.children[0] + assert child.type == child_payload["type"] + assert child.name == child_payload["name"] + assert child.input == child_payload["input"] + assert child.output == child_payload["output"] + assert child.context == child_payload["context"] + assert child.tools == child_payload["tools"] + assert child.ground_truth == child_payload["ground_truth"] + assert child.children is not None + assert child.children[0].input == {"query": "q"} + + +def test_step_rejects_invalid_child_values_clearly(): + with pytest.raises(ValidationError, match="children"): + Step(type="llm", name="parent", input="hello", children=["not a step"]) + + +@pytest.mark.asyncio +async def test_evaluate_controls_passes_children_without_recursively_evaluating(monkeypatch): + """Children are data on the current Step and are not evaluated separately.""" + mock_check = AsyncMock(return_value=EvaluationResult(is_safe=True, confidence=1.0)) + monkeypatch.setattr(evaluation, "check_evaluation_with_local", mock_check) + children = [ + Step( + type="tool", + name="search", + input={"query": "nested"}, + children=[Step(type="llm", name="nested", input="prompt")], + ) + ] + + with patch("agent_control.state.server_url", "http://localhost:8000"): + await evaluation.evaluate_controls( + step_name="parent", + input="hello", + children=children, + stage="pre", + agent_name="test-bot", + ) + + mock_check.assert_awaited_once() + passed_step = mock_check.call_args.kwargs["step"] + assert passed_step.children == children + assert passed_step.children[0].children == children[0].children + + @pytest.mark.asyncio async def test_evaluate_controls_uses_session_api_key_header(monkeypatch): """evaluate_controls should pass init's API-key header into the client.""" diff --git a/sdks/python/tests/test_local_evaluation.py b/sdks/python/tests/test_local_evaluation.py index 1763a136..0c033c3e 100644 --- a/sdks/python/tests/test_local_evaluation.py +++ b/sdks/python/tests/test_local_evaluation.py @@ -13,6 +13,11 @@ import httpx import pytest +from agent_control.client import AgentControlClient +from agent_control.evaluation import ( + _merge_results, + check_evaluation_with_local, +) from agent_control_models import ( ControlMatch, EvaluationResponse, @@ -20,12 +25,6 @@ Step, ) -from agent_control.client import AgentControlClient -from agent_control.evaluation import ( - _merge_results, - check_evaluation_with_local, -) - # ============================================================================= # Test Fixtures # ============================================================================= @@ -1316,3 +1315,29 @@ async def test_local_evaluation_includes_steering_context(self, agent_name, llm_ from agent_control_models.controls import SteeringContext as SteeringContextModel assert isinstance(match.steering_context, SteeringContextModel) assert match.steering_context.message == "Please rephrase your input" + + +@pytest.mark.asyncio +async def test_child_content_does_not_change_current_step_control_matching( + agent_name: str, +) -> None: + """Local controls still select against the current step, not its children.""" + parent = Step( + type="llm", + name="parent", + input="safe parent input", + children=[Step(type="llm", name="child", input="test secret")], + ) + client = _HttpOnlyDuckClient() + + result = await check_evaluation_with_local( + client=client, # Given: local execution with a matching pattern only in a child + agent_name=agent_name, + step=parent, + stage="pre", + controls=[make_control_dict(1, "parent-match", execution="sdk", pattern="secret")], + ) + + assert result.is_safe is True + assert result.matches is None + client.http_client.post.assert_not_awaited() diff --git a/sdks/typescript/src/generated/models/step.ts b/sdks/typescript/src/generated/models/step.ts index e5dd9775..cee84a13 100644 --- a/sdks/typescript/src/generated/models/step.ts +++ b/sdks/typescript/src/generated/models/step.ts @@ -9,6 +9,10 @@ import { remap as remap$ } from "../lib/primitives.js"; * Runtime payload for an agent step invocation. */ export type Step = { + /** + * Optional child steps associated with this invocation + */ + children?: Array | null | undefined; /** * Optional context (conversation history, metadata, etc.) */ @@ -41,6 +45,7 @@ export type Step = { /** @internal */ export type Step$Outbound = { + children?: Array | null | undefined; context?: { [k: string]: any } | null | undefined; ground_truth?: any | null | undefined; input: any; @@ -53,6 +58,9 @@ export type Step$Outbound = { /** @internal */ export const Step$outboundSchema: z.ZodMiniType = z.pipe( z.object({ + children: z.optional( + z.nullable(z.array(z.lazy(() => Step$outboundSchema))), + ), context: z.optional(z.nullable(z.record(z.string(), z.any()))), groundTruth: z.optional(z.nullable(z.any())), input: z.any(),