diff --git a/engine/src/agent_control_engine/core.py b/engine/src/agent_control_engine/core.py index b2cd81b3..7dbba3c7 100644 --- a/engine/src/agent_control_engine/core.py +++ b/engine/src/agent_control_engine/core.py @@ -320,7 +320,7 @@ async def _evaluate_leaf( timeout = DEFAULT_EVALUATOR_TIMEOUT result = await asyncio.wait_for( - evaluator.evaluate(data), + evaluator.evaluate_with_context(data, request.step), timeout=timeout, ) except TimeoutError: diff --git a/engine/tests/test_core.py b/engine/tests/test_core.py index baa46bab..aae3f8fc 100644 --- a/engine/tests/test_core.py +++ b/engine/tests/test_core.py @@ -39,13 +39,15 @@ class SimpleConfig(BaseModel): # Shared state for coordination between test evaluators _execution_log: list[str] = [] _blocker_event: asyncio.Event | None = None +_context_calls: list[tuple[Any, Step]] = [] def reset_test_state() -> None: """Reset shared test state.""" - global _execution_log, _blocker_event + global _execution_log, _blocker_event, _context_calls _execution_log = [] _blocker_event = asyncio.Event() + _context_calls = [] class AllowEvaluator(Evaluator[SimpleConfig]): @@ -163,6 +165,24 @@ async def evaluate(self, data: Any) -> EvaluatorResult: return result +class ContextEvaluator(Evaluator[SimpleConfig]): + """Evaluator that records selector data and full request context.""" + + metadata = EvaluatorMetadata( + name="test-context", + version="1.0.0", + description="Records contextual calls", + ) + config_model = SimpleConfig + + async def evaluate(self, data: Any) -> EvaluatorResult: + raise AssertionError("engine should call evaluate_with_context") + + async def evaluate_with_context(self, data: Any, step: Step) -> EvaluatorResult: + _context_calls.append((data, step)) + return EvaluatorResult(matched=False, confidence=1.0, message="context received") + + @dataclass class MockControlWithIdentity: """Mock control for testing.""" @@ -185,6 +205,7 @@ def setup_test_evaluators(): BlockerEvaluator, SlowEvaluator, MetadataEvaluator, + ContextEvaluator, ]: try: register_evaluator(evaluator_cls) @@ -253,6 +274,69 @@ def make_control( ) +@pytest.mark.asyncio +async def test_context_evaluator_receives_selected_data_and_complete_step() -> None: + # Given: a selector targeting output and a step with structured context + engine = ControlEngine( + [make_control(1, "context", "test-context", action="observe", path="output")] + ) + step = Step( + type="llm", + name="test-step", + input="question", + output="answer", + context={"conversation_id": "c-1"}, + tools=[{"name": "search", "description": "Search", "input_schema": {}}], + ground_truth="expected answer", + ) + + # When: evaluating the request through the engine + await engine.process( + EvaluationRequest( + agent_name="00000000-0000-0000-0000-000000000001", + step=step, + stage="pre", + ) + ) + + # Then: selector behavior is unchanged and the complete Step is separate + assert _context_calls == [("answer", step)] + + +@pytest.mark.asyncio +async def test_cached_context_evaluator_handles_concurrent_steps_without_retaining_state() -> None: + # Given: one cached evaluator configuration and two independent requests + engine = ControlEngine( + [make_control(1, "context", "test-context", action="observe", path="input")] + ) + first = Step(type="llm", name="test-step", input="first", ground_truth="one") + second = Step(type="llm", name="test-step", input="second", ground_truth="two") + + # When: the requests are evaluated concurrently + await asyncio.gather( + engine.process( + EvaluationRequest( + agent_name="00000000-0000-0000-0000-000000000001", + step=first, + stage="pre", + ) + ), + engine.process( + EvaluationRequest( + agent_name="00000000-0000-0000-0000-000000000001", + step=second, + stage="pre", + ) + ), + ) + + # Then: each selected value remains paired with its own full Step + assert {(data, step.ground_truth) for data, step in _context_calls} == { + ("first", "one"), + ("second", "two"), + } + + # ============================================================================= # Test: Parallel Execution # ============================================================================= diff --git a/evaluators/builtin/src/agent_control_evaluators/_base.py b/evaluators/builtin/src/agent_control_evaluators/_base.py index c32b92a5..d83c0c3a 100644 --- a/evaluators/builtin/src/agent_control_evaluators/_base.py +++ b/evaluators/builtin/src/agent_control_evaluators/_base.py @@ -7,7 +7,7 @@ from dataclasses import dataclass from typing import TYPE_CHECKING, Any, ClassVar, Generic, TypeVar -from agent_control_models import EvaluatorResult +from agent_control_models import EvaluatorResult, Step from agent_control_models.base import BaseModel if TYPE_CHECKING: @@ -161,6 +161,23 @@ async def evaluate(self, data: Any) -> EvaluatorResult: """ pass + async def evaluate_with_context(self, data: Any, step: Step) -> EvaluatorResult: + """Evaluate selected data with access to the complete runtime step. + + The default implementation preserves compatibility with evaluators that + implement only :meth:`evaluate`. Evaluators must treat ``step`` as + immutable request-scoped context because evaluator instances are cached + and may be invoked concurrently. + + Args: + data: Data extracted by the configured selector. + step: Complete runtime step for the current request. + + Returns: + EvaluatorResult produced by this evaluator. + """ + return await self.evaluate(data) + def get_timeout_seconds(self) -> float: """Get timeout in seconds from config or metadata default.""" timeout_ms: int = getattr(self.config, "timeout_ms", self.metadata.timeout_ms) diff --git a/evaluators/builtin/tests/test_base.py b/evaluators/builtin/tests/test_base.py index 776a8d01..d81cc3f5 100644 --- a/evaluators/builtin/tests/test_base.py +++ b/evaluators/builtin/tests/test_base.py @@ -3,11 +3,12 @@ Architecture: Evaluators take config at __init__, evaluate() only takes data. """ -import pytest from typing import Any +import pytest + from agent_control_evaluators import Evaluator, EvaluatorConfig, EvaluatorMetadata -from agent_control_models import EvaluatorResult +from agent_control_models import EvaluatorResult, Step class MockConfig(EvaluatorConfig): @@ -106,6 +107,20 @@ async def test_mock_evaluator_evaluate_no_match(self): assert result.matched is False + @pytest.mark.asyncio + async def test_contextual_evaluation_delegates_to_existing_evaluate(self): + """Existing evaluators receive selected data through the default hook.""" + # Given: an evaluator that implements only evaluate(data) + evaluator = MockEvaluator.from_dict({"should_match": True}) + step = Step(type="llm", name="answer", input="full input") + + # When: the engine-facing contextual hook is called + result = await evaluator.evaluate_with_context("selected data", step) + + # Then: the legacy evaluate implementation handles the selected data + assert result.matched is True + assert result.metadata == {"data": "selected data"} + def test_evaluator_config_stored(self): """Test that evaluator stores config.""" evaluator = MockEvaluator.from_dict({"should_match": True}) diff --git a/evaluators/contrib/galileo/src/agent_control_evaluator_galileo/__init__.py b/evaluators/contrib/galileo/src/agent_control_evaluator_galileo/__init__.py index 5606bf5d..93eb8567 100644 --- a/evaluators/contrib/galileo/src/agent_control_evaluator_galileo/__init__.py +++ b/evaluators/contrib/galileo/src/agent_control_evaluator_galileo/__init__.py @@ -25,6 +25,7 @@ LunaEvaluator, LunaEvaluatorConfig, LunaOperator, + ScorerInvokeRecord, ScorerInvokeRequest, ScorerInvokeResponse, ) @@ -32,6 +33,7 @@ __all__ = [ "GalileoLunaClient", "ScorerInvokeRequest", + "ScorerInvokeRecord", "ScorerInvokeResponse", "LunaEvaluator", "LunaEvaluatorConfig", diff --git a/evaluators/contrib/galileo/src/agent_control_evaluator_galileo/luna/__init__.py b/evaluators/contrib/galileo/src/agent_control_evaluator_galileo/luna/__init__.py index b26feaac..39f95016 100644 --- a/evaluators/contrib/galileo/src/agent_control_evaluator_galileo/luna/__init__.py +++ b/evaluators/contrib/galileo/src/agent_control_evaluator_galileo/luna/__init__.py @@ -3,6 +3,7 @@ from agent_control_evaluator_galileo.luna.client import ( GalileoLunaClient, ScorerInvokeInputs, + ScorerInvokeRecord, ScorerInvokeRequest, ScorerInvokeResponse, ) @@ -12,6 +13,7 @@ __all__ = [ "GalileoLunaClient", "ScorerInvokeInputs", + "ScorerInvokeRecord", "ScorerInvokeRequest", "ScorerInvokeResponse", "LunaEvaluatorConfig", diff --git a/evaluators/contrib/galileo/src/agent_control_evaluator_galileo/luna/client.py b/evaluators/contrib/galileo/src/agent_control_evaluator_galileo/luna/client.py index a1fc4d71..e0b759b7 100644 --- a/evaluators/contrib/galileo/src/agent_control_evaluator_galileo/luna/client.py +++ b/evaluators/contrib/galileo/src/agent_control_evaluator_galileo/luna/client.py @@ -14,7 +14,7 @@ from urllib.parse import urlsplit import httpx -from agent_control_models import JSONObject, JSONValue +from agent_control_models import JSONObject, JSONValue, Step from pydantic import BaseModel, Field, PrivateAttr, model_validator logger = logging.getLogger(__name__) @@ -154,7 +154,18 @@ class ScorerInvokeInputs(BaseModel): query: JSONValue = "" response: JSONValue = "" ground_truth: JSONValue = None - tools: JSONValue = None + tools: list[JSONObject] | None = None + + +class ScorerInvokeRecord(BaseModel): + """Structured runtime record sent alongside legacy scorer inputs.""" + + type: str = Field(min_length=1) + input: JSONValue = None + output: JSONValue = None + context: JSONObject | None = None + tools: list[JSONObject] | None = None + dataset_output: JSONValue = None class ScorerInvokeRequest(BaseModel): @@ -172,6 +183,7 @@ class ScorerInvokeRequest(BaseModel): scorer_version_id: str | None = Field(default=None, min_length=1) scorer_label: str | None = Field(default=None, min_length=1) inputs: ScorerInvokeInputs + record: ScorerInvokeRecord | None = None config: JSONObject = Field(default_factory=dict) @model_validator(mode="after") @@ -357,6 +369,7 @@ async def invoke( scorer_label: str | None = None, input: JSONValue = None, output: JSONValue = None, + step: Step | None = None, config: JSONObject | None = None, timeout: float = DEFAULT_TIMEOUT_SECS, headers: dict[str, str] | None = None, @@ -369,6 +382,7 @@ async def invoke( scorer_label: Optional display/metadata label. input: Optional user/system prompt text. output: Optional model response text. + step: Optional complete runtime step used for structured dual-write. config: Optional scorer-specific configuration. timeout: Request timeout in seconds. headers: Additional request headers. @@ -390,7 +404,22 @@ async def invoke( scorer_version_id=scorer_version_id, scorer_label=scorer_label, inputs=ScorerInvokeInputs( - query="" if input is None else input, response="" if output is None else output + query="" if input is None else input, + response="" if output is None else output, + ground_truth=step.ground_truth if step is not None else None, + tools=step.tools if step is not None else None, + ), + record=( + ScorerInvokeRecord( + type=step.type, + input=step.input, + output=step.output, + context=step.context, + tools=step.tools, + dataset_output=step.ground_truth, + ) + if step is not None + else None ), config=config if config is not None else {}, ).to_dict() diff --git a/evaluators/contrib/galileo/src/agent_control_evaluator_galileo/luna/evaluator.py b/evaluators/contrib/galileo/src/agent_control_evaluator_galileo/luna/evaluator.py index 777c22cf..7a837aaa 100644 --- a/evaluators/contrib/galileo/src/agent_control_evaluator_galileo/luna/evaluator.py +++ b/evaluators/contrib/galileo/src/agent_control_evaluator_galileo/luna/evaluator.py @@ -10,7 +10,7 @@ import httpx from agent_control_evaluators import Evaluator, EvaluatorMetadata, register_evaluator -from agent_control_models import EvaluatorResult, JSONValue +from agent_control_models import EvaluatorResult, JSONValue, Step from .client import GalileoLunaClient, ScorerInvokeResponse from .config import LunaEvaluatorConfig, coerce_number @@ -198,6 +198,22 @@ async def evaluate(self, data: Any) -> EvaluatorResult: Returns: EvaluatorResult with local threshold decision and scorer metadata. """ + return await self._evaluate(data, step=None) + + async def evaluate_with_context(self, data: Any, step: Step) -> EvaluatorResult: + """Evaluate selected data while dual-writing the complete runtime step. + + Args: + data: Data selected by the configured control selector. + step: Complete runtime step for structured scorer context. + + Returns: + EvaluatorResult with local threshold decision and scorer metadata. + """ + return await self._evaluate(data, step=step) + + async def _evaluate(self, data: Any, *, step: Step | None) -> EvaluatorResult: + """Run a Luna evaluation with optional structured runtime context.""" input_text, output_text = self._prepare_payload(data) if not (_has_text(input_text) or _has_text(output_text)): return EvaluatorResult( @@ -209,6 +225,8 @@ async def evaluate(self, data: Any) -> EvaluatorResult: try: scorer_kwargs = self._scorer_kwargs() + if step is not None: + scorer_kwargs["step"] = step response = await self._get_client().invoke( **scorer_kwargs, input=input_text if _has_text(input_text) else None, diff --git a/evaluators/contrib/galileo/tests/test_luna_evaluator.py b/evaluators/contrib/galileo/tests/test_luna_evaluator.py index 6c605c9c..cb2a5f11 100644 --- a/evaluators/contrib/galileo/tests/test_luna_evaluator.py +++ b/evaluators/contrib/galileo/tests/test_luna_evaluator.py @@ -10,7 +10,7 @@ import httpx import pytest -from agent_control_models import EvaluatorResult +from agent_control_models import EvaluatorResult, Step from pydantic import ValidationError LUNA_ENV = { @@ -515,6 +515,65 @@ def handler(request: httpx.Request) -> httpx.Response: assert payload["internal"] is True assert payload["scope"] == "scorers.invoke" + @pytest.mark.asyncio + async def test_client_dual_writes_legacy_inputs_and_structured_record(self) -> None: + from agent_control_evaluator_galileo.luna import GalileoLunaClient + + captured: dict[str, object] = {} + + def handler(request: httpx.Request) -> httpx.Response: + captured["body"] = json.loads(request.content.decode()) + return httpx.Response( + 200, + json={"score": 0.9, "status": "success", "additive_field": "ignored"}, + ) + + # Given: selected legacy values and a complete structured runtime Step + step = Step( + type="llm", + name="answer", + input={"messages": [{"role": "user", "content": "question"}]}, + output={"text": "answer"}, + context={"session": "s-1"}, + tools=[{"name": "search", "description": "Search", "input_schema": {}}], + ground_truth={"text": "expected"}, + ) + with patch.dict(os.environ, LUNA_ENV, clear=True): + client = GalileoLunaClient() + client._client = httpx.AsyncClient(transport=httpx.MockTransport(handler)) + + # When: invoking the rollout-compatible Runners endpoint + try: + response = await client.invoke( + scorer_id="scorer-123", + input="selected question", + output="selected answer", + step=step, + ) + finally: + await client.close() + + # Then: old inputs and the expanded record are sent together + assert response.score == 0.9 + assert captured["body"] == { + "scorer_id": "scorer-123", + "inputs": { + "query": "selected question", + "response": "selected answer", + "ground_truth": {"text": "expected"}, + "tools": [{"name": "search", "description": "Search", "input_schema": {}}], + }, + "record": { + "type": "llm", + "input": {"messages": [{"role": "user", "content": "question"}]}, + "output": {"text": "answer"}, + "context": {"session": "s-1"}, + "tools": [{"name": "search", "description": "Search", "input_schema": {}}], + "dataset_output": {"text": "expected"}, + }, + "config": {}, + } + @pytest.mark.asyncio async def test_client_forwards_scorer_version_id_when_configured(self) -> None: from agent_control_evaluator_galileo.luna import GalileoLunaClient @@ -656,6 +715,41 @@ async def test_evaluator_applies_threshold_locally_to_raw_score(self) -> None: timeout=5.0, ) + @patch.dict(os.environ, LUNA_ENV) + @pytest.mark.asyncio + async def test_evaluator_contextual_hook_forwards_complete_step(self) -> None: + from agent_control_evaluator_galileo.luna import LunaEvaluator, ScorerInvokeResponse + from agent_control_evaluator_galileo.luna.client import GalileoLunaClient + + # Given: selected scorer data and complete structured runtime context + evaluator = LunaEvaluator.from_dict( + {"scorer_id": "scorer-123", "threshold": 0.5, "operator": "gte"} + ) + step = Step( + type="llm", + name="answer", + input="full input", + output="full output", + tools=[{"name": "search", "description": "Search", "input_schema": {}}], + ground_truth="expected", + ) + + # When: evaluating through the contextual hook + with patch.object(GalileoLunaClient, "invoke", new_callable=AsyncMock) as mock_invoke: + mock_invoke.return_value = ScorerInvokeResponse(score=0.8, status="success") + result = await evaluator.evaluate_with_context("selected input", step) + + # Then: selector-selected data and the complete Step are both forwarded + assert result.matched is True + mock_invoke.assert_awaited_once_with( + scorer_id="scorer-123", + step=step, + input="selected input", + output=None, + config=None, + timeout=10.0, + ) + @patch.dict(os.environ, LUNA_ENV) @pytest.mark.asyncio async def test_evaluator_forwards_configured_scorer_version_id(self) -> None: diff --git a/models/src/agent_control_models/agent.py b/models/src/agent_control_models/agent.py index 6a0eedba..e60573c9 100644 --- a/models/src/agent_control_models/agent.py +++ b/models/src/agent_control_models/agent.py @@ -142,6 +142,8 @@ def validate_type(cls, v: str) -> str: class Step(BaseModel): """Runtime payload for an agent step invocation.""" + model_config = {"frozen": True} + type: str = Field( ..., min_length=1, @@ -159,6 +161,13 @@ class Step(BaseModel): context: JSONObject | None = Field( None, description="Optional context (conversation history, metadata, etc.)" ) + tools: list[JSONObject] | None = Field( + None, + description="Complete structured definitions of tools available to the LLM", + ) + ground_truth: JSONValue | None = Field( + None, description="Optional expected or reference output for this step" + ) @field_validator("type") @classmethod diff --git a/models/tests/test_step_runtime_context.py b/models/tests/test_step_runtime_context.py new file mode 100644 index 00000000..2b9f80f6 --- /dev/null +++ b/models/tests/test_step_runtime_context.py @@ -0,0 +1,51 @@ +"""Tests for optional structured runtime context on Step.""" + +from agent_control_models import Step +from pydantic import ValidationError + + +def test_step_accepts_and_serializes_structured_tools_and_ground_truth() -> None: + # Given: a runtime LLM step with structured scorer context + payload = { + "type": "llm", + "name": "answer", + "input": {"question": "Capital of France?"}, + "output": "Paris", + "tools": [ + { + "name": "search", + "description": "Search documents", + "input_schema": {"type": "object"}, + } + ], + "ground_truth": {"answer": "Paris"}, + } + + # When: validating and serializing the public Step model + serialized = Step.model_validate(payload).model_dump(mode="json") + + # Then: structured values survive unchanged + assert serialized["tools"] == payload["tools"] + assert serialized["ground_truth"] == payload["ground_truth"] + + +def test_existing_step_payload_remains_valid() -> None: + # Given/When: an existing payload without structured context + step = Step(type="llm", name="answer", input="hello") + + # Then: new fields remain optional + assert step.tools is None + assert step.ground_truth is None + + +def test_step_is_immutable_runtime_context() -> None: + # Given: a validated runtime step + step = Step(type="llm", name="answer", input="hello") + + # When/Then: evaluators cannot replace request-scoped fields + try: + step.input = "changed" + except ValidationError: + pass + else: + raise AssertionError("Step must be frozen") diff --git a/sdks/python/src/agent_control/evaluation.py b/sdks/python/src/agent_control/evaluation.py index e79b736b..2a045bb5 100644 --- a/sdks/python/src/agent_control/evaluation.py +++ b/sdks/python/src/agent_control/evaluation.py @@ -15,6 +15,7 @@ EvaluationResponse, EvaluationResult, EvaluatorResult, + JSONValue, Step, ) @@ -520,6 +521,8 @@ async def evaluate_controls( input: Any | None = None, output: Any | None = None, context: dict[str, Any] | None = None, + tools: list[dict[str, JSONValue]] | None = None, + ground_truth: JSONValue | None = None, step_type: Literal["tool", "llm"] = "llm", stage: Literal["pre", "post"] = "pre", agent_name: str, @@ -552,6 +555,10 @@ async def evaluate_controls( } if context is not None: step_dict["context"] = context + if tools is not None: + step_dict["tools"] = tools + if ground_truth is not None: + step_dict["ground_truth"] = ground_truth step_obj = Step(**step_dict) # type: ignore[arg-type] resolved_controls = state.server_controls or [] diff --git a/sdks/python/src/agent_control/integrations/_core.py b/sdks/python/src/agent_control/integrations/_core.py index 27693dbf..c9587dcd 100644 --- a/sdks/python/src/agent_control/integrations/_core.py +++ b/sdks/python/src/agent_control/integrations/_core.py @@ -4,7 +4,7 @@ from typing import Any, Literal -from agent_control_models import EvaluationResult +from agent_control_models import EvaluationResult, JSONValue import agent_control from agent_control import ControlSteerError, ControlViolationError @@ -51,6 +51,8 @@ async def _evaluate_and_enforce( input: Any | None = None, output: Any | None = None, context: dict[str, Any] | None = None, + tools: list[dict[str, JSONValue]] | None = None, + ground_truth: JSONValue | None = None, step_type: Literal["tool", "llm"] = "llm", stage: Literal["pre", "post"] = "pre", ) -> EvaluationResult: @@ -61,6 +63,8 @@ async def _evaluate_and_enforce( input=input, output=output, context=context, + tools=tools, + ground_truth=ground_truth, step_type=step_type, stage=stage, agent_name=agent_name, diff --git a/sdks/python/src/agent_control/integrations/_tools.py b/sdks/python/src/agent_control/integrations/_tools.py new file mode 100644 index 00000000..84f55f90 --- /dev/null +++ b/sdks/python/src/agent_control/integrations/_tools.py @@ -0,0 +1,52 @@ +"""Normalization helpers for complete framework tool registries.""" + +from __future__ import annotations + +from typing import Any + +from agent_control_models import JSONObject, JSONValue + + +def normalized_tool_definition( + *, + name: str, + description: str | None, + input_schema: dict[str, Any] | None, +) -> JSONObject: + """Return the framework-neutral available-tool representation.""" + definition: dict[str, JSONValue] = { + "name": name, + "description": description or "", + "input_schema": input_schema or {}, + } + return definition + + +def normalize_strands_tool_specs(specs: object) -> list[JSONObject] | None: + """Normalize a complete Strands tool-spec collection, if available.""" + if not isinstance(specs, list): + return None + + normalized: list[JSONObject] = [] + for spec in specs: + if not isinstance(spec, dict): + return None + name = spec.get("name") + if not isinstance(name, str) or not name: + return None + description = spec.get("description") + raw_input_schema = spec.get("inputSchema") + if isinstance(raw_input_schema, dict) and isinstance(raw_input_schema.get("json"), dict): + input_schema = raw_input_schema["json"] + elif isinstance(raw_input_schema, dict): + input_schema = raw_input_schema + else: + input_schema = None + normalized.append( + normalized_tool_definition( + name=name, + description=description if isinstance(description, str) else None, + input_schema=input_schema, + ) + ) + return normalized diff --git a/sdks/python/src/agent_control/integrations/google_adk/plugin.py b/sdks/python/src/agent_control/integrations/google_adk/plugin.py index 28e59698..75266623 100644 --- a/sdks/python/src/agent_control/integrations/google_adk/plugin.py +++ b/sdks/python/src/agent_control/integrations/google_adk/plugin.py @@ -19,6 +19,7 @@ from agent_control._schema_derivation import derive_schemas from agent_control._state import state from agent_control.integrations._core import _evaluate_and_enforce +from agent_control.integrations._tools import normalized_tool_definition from agent_control.validation import ensure_agent_name try: @@ -109,11 +110,13 @@ def __init__( self._known_steps: dict[tuple[str, str], StepSchemaDict] = {} self._synced_step_keys: set[tuple[str, str]] = set() self._step_sync_tasks: dict[tuple[str, str], asyncio.Task[None]] = {} + self._available_tools_by_step: dict[str, list[dict[str, Any]]] = {} def bind(self, agent: Any) -> None: """Pre-register known ADK steps before the runner starts.""" steps = self._discover_steps(agent) + self._remember_available_tools(agent) self._remember_steps(steps) self._sync_steps_blocking(steps, raise_on_error=True) @@ -176,6 +179,7 @@ async def before_model_callback( step_name, input=request_text, context=context, + tools=self._available_tools_by_step.get(step_name), step_type="llm", stage="pre", ) @@ -227,6 +231,7 @@ async def after_model_callback( input=input_text, output=output_text, context=context, + tools=self._available_tools_by_step.get(step_name), step_type="llm", stage="post", ) @@ -662,6 +667,33 @@ def _iter_tools(self, agent: Any) -> Iterable[Any]: return tools return [] + def _remember_available_tools(self, root_agent: Any) -> None: + """Capture each ADK agent's complete bound tool set as structured JSON.""" + available_tools: dict[str, list[dict[str, Any]]] = {} + for agent in self._iter_agents(root_agent): + agent_name = getattr(agent, "name", None) + if not isinstance(agent_name, str) or not agent_name: + continue + step_name = self._resolve_step_name( + agent_name, + step_type="llm", + callback_context=None, + agent=agent, + ) + definitions: list[dict[str, Any]] = [] + for tool in self._iter_tools(agent): + tool_name = self._resolve_tool_step_name(tool, agent_step_name=step_name) + schema = self._build_tool_step_schema(tool, tool_name) + definitions.append( + normalized_tool_definition( + name=resolve_tool_name(tool), + description=schema.get("description"), + input_schema=schema.get("input_schema"), + ) + ) + available_tools[step_name] = definitions + self._available_tools_by_step = available_tools + def _remember_steps(self, steps: Iterable[StepSchemaDict]) -> None: for step in steps: key = (step["type"], step["name"]) diff --git a/sdks/python/src/agent_control/integrations/strands/plugin.py b/sdks/python/src/agent_control/integrations/strands/plugin.py index 1aa503cd..367a673c 100644 --- a/sdks/python/src/agent_control/integrations/strands/plugin.py +++ b/sdks/python/src/agent_control/integrations/strands/plugin.py @@ -10,6 +10,7 @@ import agent_control from agent_control import ControlSteerError, ControlViolationError +from agent_control.integrations._tools import normalize_strands_tool_specs try: from strands.hooks import ( # type: ignore[import-not-found] @@ -86,6 +87,7 @@ def __init__( self.event_control_list = event_control_list self.on_violation_callback = on_violation_callback self.enable_logging = enable_logging + self._tool_registry: Any | None = None def _invoke_callback(self, control_name: str, stage: str, result: EvaluationResult) -> None: if self.on_violation_callback: @@ -109,6 +111,7 @@ async def _evaluate_and_enforce( input: Any | None = None, output: Any | None = None, context: dict[str, Any] | None = None, + tools: list[dict[str, Any]] | None = None, step_type: Literal["tool", "llm"] = "llm", stage: Literal["pre", "post"] = "pre", use_runtime_error: bool = False, @@ -118,6 +121,7 @@ async def _evaluate_and_enforce( input=input, output=output, context=context, + tools=tools, step_type=step_type, stage=stage, agent_name=self.agent_name, @@ -181,6 +185,7 @@ async def _evaluate_and_enforce( ) def init_agent(self, agent: Any) -> None: + self._tool_registry = getattr(agent, "tool_registry", None) event_map = { BeforeInvocationEvent: self.check_before_invocation, BeforeModelCallEvent: self.check_before_model, @@ -211,6 +216,7 @@ async def check_before_invocation(self, event: BeforeInvocationEvent) -> None: await self._evaluate_and_enforce( step_name="check_before_invocation", input=input_text, + tools=self._available_tools(), step_type="llm", stage="pre", ) @@ -220,6 +226,7 @@ async def check_before_model(self, event: BeforeModelCallEvent) -> None: await self._evaluate_and_enforce( step_name="check_before_model", input=input_text, + tools=self._available_tools(), step_type="llm", stage="pre", ) @@ -232,6 +239,7 @@ async def check_after_model(self, event: AfterModelCallEvent) -> None: input=input_text, output=output_text, context=context, + tools=self._available_tools(), step_type="llm", stage="post", ) @@ -270,6 +278,7 @@ async def check_before_node(self, event: BeforeNodeCallEvent) -> None: step_name=node_id, input=input_text, context=context, + tools=self._available_tools(), step_type="llm", stage="pre", ) @@ -283,10 +292,22 @@ async def check_after_node(self, event: AfterNodeCallEvent) -> None: input=input_text, output=output_text, context=context, + tools=self._available_tools(), step_type="llm", stage="post", ) + def _available_tools(self) -> list[dict[str, Any]] | None: + """Return the complete current Strands registry in normalized form.""" + get_specs = getattr(self._tool_registry, "get_all_tool_specs", None) + if not callable(get_specs): + return None + try: + return normalize_strands_tool_specs(get_specs()) + except Exception: + logger.warning("Unable to capture complete Strands tool definitions", exc_info=True) + return None + def _extract_user_message_from_list(self, messages: list | None, reverse: bool = False) -> str: if not messages: return "" diff --git a/sdks/python/tests/test_evaluation.py b/sdks/python/tests/test_evaluation.py index 2fb92555..885a69a7 100644 --- a/sdks/python/tests/test_evaluation.py +++ b/sdks/python/tests/test_evaluation.py @@ -60,8 +60,10 @@ def json(self) -> dict[str, object]: "type": "llm", "name": "chat", "input": "hello", - "output": None, - "context": None, + "output": None, + "context": None, + "tools": None, + "ground_truth": None, }, "stage": "pre", "target_type": None, @@ -125,6 +127,34 @@ async def test_evaluate_controls_with_context(monkeypatch): assert mock_check.call_args is not None +@pytest.mark.asyncio +async def test_evaluate_controls_preserves_explicit_tools_and_ground_truth(monkeypatch): + """Explicit structured scorer context is preserved on the SDK Step.""" + # Given: a configured SDK and local evaluation boundary + mock_check = AsyncMock(return_value=EvaluationResult(is_safe=True, confidence=1.0)) + monkeypatch.setattr(evaluation, "check_evaluation_with_local", mock_check) + tools = [ + {"name": "search", "description": "Search", "input_schema": {"type": "object"}} + ] + + # When: a direct SDK caller supplies tools and ground truth + with patch("agent_control.state.server_url", "http://localhost:8000"): + await evaluation.evaluate_controls( + step_name="chat", + input="question", + output="answer", + tools=tools, + ground_truth={"answer": "expected"}, + stage="post", + agent_name="test-bot", + ) + + # Then: both fields are retained as structured JSON + step = mock_check.call_args.kwargs["step"] + assert step.tools == tools + assert step.ground_truth == {"answer": "expected"} + + @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_google_adk_plugin.py b/sdks/python/tests/test_google_adk_plugin.py index f68bd341..b50499a1 100644 --- a/sdks/python/tests/test_google_adk_plugin.py +++ b/sdks/python/tests/test_google_adk_plugin.py @@ -540,6 +540,34 @@ def test_bind_discovers_root_sub_agents_and_tools(plugin_module): mock_sync.assert_called_once() +@pytest.mark.asyncio +async def test_bound_agent_passes_complete_normalized_tools_to_llm(plugin_module): + # Given: an ADK agent whose complete tool list is available during binding + plugin = plugin_module.AgentControlPlugin(agent_name="test-agent01") + root = SimpleNamespace( + name="planner", + tools=[MockTool("search_docs", "Search documentation")], + ) + with patch.object(plugin, "_sync_steps_blocking"): + plugin.bind(root) + + # When: an LLM callback is evaluated + with patch.object( + plugin_module, "_evaluate_and_enforce", AsyncMock(return_value=MagicMock()) + ) as mock_eval: + await plugin.before_model_callback( + callback_context=MockCallbackContext("planner"), + llm_request=MockLlmRequest("hello"), + ) + + # Then: available definitions are normalized, not inferred from a call + definitions = mock_eval.await_args.kwargs["tools"] + assert len(definitions) == 1 + assert definitions[0]["name"] == "search_docs" + assert definitions[0]["description"] == "Search documentation" + assert definitions[0]["input_schema"]["properties"]["city"]["type"] == "string" + + def test_bind_keeps_duplicate_tool_names_distinct_across_sub_agents(plugin_module): plugin = plugin_module.AgentControlPlugin(agent_name="test-agent01") root = SimpleNamespace( @@ -557,6 +585,18 @@ def test_bind_keeps_duplicate_tool_names_distinct_across_sub_agents(plugin_modul assert ("tool", "writer.search_docs") in plugin._known_steps +def test_available_tool_capture_skips_unnamed_agents(plugin_module): + # Given: a framework object that does not identify an ADK agent + plugin = plugin_module.AgentControlPlugin(agent_name="test-agent01") + unnamed_agent = SimpleNamespace(tools=[MockTool("search_docs")]) + + # When: available tools are captured from the bound hierarchy + plugin._remember_available_tools(unnamed_agent) + + # Then: no incomplete tool set is guessed + assert plugin._available_tools_by_step == {} + + @pytest.mark.asyncio async def test_lazy_step_sync_when_bind_skipped(plugin_module): plugin = plugin_module.AgentControlPlugin(agent_name="test-agent01") diff --git a/sdks/python/tests/test_integration_tools.py b/sdks/python/tests/test_integration_tools.py new file mode 100644 index 00000000..8f23eb4f --- /dev/null +++ b/sdks/python/tests/test_integration_tools.py @@ -0,0 +1,47 @@ +"""Tests for framework-neutral available-tool normalization.""" + +import pytest + +from agent_control.integrations._tools import normalize_strands_tool_specs + + +@pytest.mark.parametrize( + "specs", + [ + {"search": {}}, + ["not-a-tool-spec"], + [{"description": "missing name"}], + [{"name": ""}], + ], +) +def test_normalizer_rejects_incomplete_strands_registries(specs: object) -> None: + # Given/When: the purported complete registry has an invalid shape + normalized = normalize_strands_tool_specs(specs) + + # Then: the integration leaves tools absent rather than guessing + assert normalized is None + + +@pytest.mark.parametrize( + ("input_schema", "expected"), + [ + ({"type": "object", "properties": {}}, {"type": "object", "properties": {}}), + (None, {}), + ], +) +def test_normalizer_supports_plain_or_missing_strands_input_schema( + input_schema: object, + expected: dict[str, object], +) -> None: + # Given: valid Strands specs from supported schema variants + spec: dict[str, object] = {"name": "search", "description": 123} + if input_schema is not None: + spec["inputSchema"] = input_schema + + # When: normalizing the complete registry + normalized = normalize_strands_tool_specs([spec]) + + # Then: a stable JSON-only tool definition is produced + assert normalized == [ + {"name": "search", "description": "", "input_schema": expected} + ] diff --git a/sdks/python/tests/test_strands_plugin.py b/sdks/python/tests/test_strands_plugin.py index 3f52848e..2210fda4 100644 --- a/sdks/python/tests/test_strands_plugin.py +++ b/sdks/python/tests/test_strands_plugin.py @@ -424,6 +424,47 @@ def test_hook_initialization(): assert hook.enable_logging is False +def test_init_agent_captures_complete_strands_tool_registry(agent_control_hook): + # Given: a Strands agent exposing its complete normalized registry + registry = MagicMock() + registry.get_all_tool_specs.return_value = [ + { + "name": "search_docs", + "description": "Search documentation", + "inputSchema": {"json": {"type": "object", "properties": {}}}, + } + ] + agent = MagicMock(tool_registry=registry) + + # When: the plugin is initialized against that agent + agent_control_hook.init_agent(agent) + + # Then: the full registry is normalized without executable objects + assert agent_control_hook._available_tools() == [ + { + "name": "search_docs", + "description": "Search documentation", + "input_schema": {"type": "object", "properties": {}}, + } + ] + + +def test_available_tools_fails_closed_when_strands_registry_raises( + agent_control_hook, caplog +): + # Given: a Strands registry that cannot provide a complete tool set + registry = MagicMock() + registry.get_all_tool_specs.side_effect = RuntimeError("registry unavailable") + agent_control_hook._tool_registry = registry + + # When: tool definitions are requested + tools = agent_control_hook._available_tools() + + # Then: tools remain absent and the integration records the failure + assert tools is None + assert "Unable to capture complete Strands tool definitions" in caplog.text + + def test_hook_with_callback(): """Test AgentControlPlugin with violation callback.""" from agent_control.integrations.strands.plugin import AgentControlPlugin diff --git a/sdks/typescript/src/generated/models/step.ts b/sdks/typescript/src/generated/models/step.ts index 132cf9c9..e5dd9775 100644 --- a/sdks/typescript/src/generated/models/step.ts +++ b/sdks/typescript/src/generated/models/step.ts @@ -3,6 +3,7 @@ */ import * as z from "zod/v4-mini"; +import { remap as remap$ } from "../lib/primitives.js"; /** * Runtime payload for an agent step invocation. @@ -12,6 +13,10 @@ export type Step = { * Optional context (conversation history, metadata, etc.) */ context?: { [k: string]: any } | null | undefined; + /** + * Optional expected or reference output for this step + */ + groundTruth?: any | null | undefined; /** * Any JSON value */ @@ -24,6 +29,10 @@ export type Step = { * Output content for this step (None for pre-checks) */ output?: any | null | undefined; + /** + * Complete structured definitions of tools available to the LLM + */ + tools?: Array<{ [k: string]: any }> | null | undefined; /** * Step type (e.g., 'tool', 'llm') */ @@ -33,21 +42,30 @@ export type Step = { /** @internal */ export type Step$Outbound = { context?: { [k: string]: any } | null | undefined; + ground_truth?: any | null | undefined; input: any; name: string; output?: any | null | undefined; + tools?: Array<{ [k: string]: any }> | null | undefined; type: string; }; /** @internal */ -export const Step$outboundSchema: z.ZodMiniType = z.object( - { +export const Step$outboundSchema: z.ZodMiniType = z.pipe( + z.object({ context: z.optional(z.nullable(z.record(z.string(), z.any()))), + groundTruth: z.optional(z.nullable(z.any())), input: z.any(), name: z.string(), output: z.optional(z.nullable(z.any())), + tools: z.optional(z.nullable(z.array(z.record(z.string(), z.any())))), type: z.string(), - }, + }), + z.transform((v) => { + return remap$(v, { + groundTruth: "ground_truth", + }); + }), ); export function stepToJSON(step: Step): string { diff --git a/sdks/typescript/tests/generated-smoke.test.ts b/sdks/typescript/tests/generated-smoke.test.ts index 453267aa..dda98e3a 100644 --- a/sdks/typescript/tests/generated-smoke.test.ts +++ b/sdks/typescript/tests/generated-smoke.test.ts @@ -3,6 +3,8 @@ import path from "node:path"; import { describe, expect, it } from "vitest"; +import { stepToJSON } from "../src/generated/models/step"; + describe("generated client layout", () => { it("has a generated index entrypoint", () => { const generatedIndex = path.resolve(process.cwd(), "src/generated/index.ts"); @@ -24,4 +26,36 @@ describe("generated client layout", () => { } } }); + + it("serializes structured Step scorer context", () => { + const serialized = stepToJSON({ + type: "llm", + name: "answer", + input: "question", + output: "answer", + groundTruth: "expected", + tools: [ + { + name: "search", + description: "Search documents", + input_schema: { type: "object" }, + }, + ], + }); + + expect(JSON.parse(serialized)).toEqual({ + type: "llm", + name: "answer", + input: "question", + output: "answer", + ground_truth: "expected", + tools: [ + { + name: "search", + description: "Search documents", + input_schema: { type: "object" }, + }, + ], + }); + }); });