diff --git a/engine/tests/test_core.py b/engine/tests/test_core.py index aae3f8fc..e16d437f 100644 --- a/engine/tests/test_core.py +++ b/engine/tests/test_core.py @@ -1085,6 +1085,41 @@ def register_payload_evaluator(self): except ValueError: pass + @pytest.mark.asyncio + async def test_custom_step_type_filters_controls(self): + """Controls match custom step types without falling back to LLM.""" + controls = [ + make_control( + 1, + "retriever-control", + "test-deny", + action="deny", + config_value="retriever", + step_types=["retriever"], + ), + make_control( + 2, + "llm-control", + "test-deny", + action="deny", + config_value="llm", + step_types=["llm"], + ), + ] + engine = ControlEngine(controls) + + result = await engine.process( + EvaluationRequest( + agent_name="00000000-0000-0000-0000-000000000001", + step=Step(type="retriever", name="retrieve", input="query"), + stage="pre", + ) + ) + + assert [match.control_name for match in result.matches or []] == [ + "retriever-control" + ] + @pytest.mark.asyncio async def test_step_names_filters_tasks(self): # Given: two controls scoped to different steps diff --git a/models/src/agent_control_models/observability.py b/models/src/agent_control_models/observability.py index dbd11fac..41cadd27 100644 --- a/models/src/agent_control_models/observability.py +++ b/models/src/agent_control_models/observability.py @@ -46,7 +46,7 @@ class ControlExecutionEvent(BaseModel): control_id: Database ID of the control control_name: Name of the control (denormalized for queries) check_stage: "pre" (before execution) or "post" (after execution) - applies_to: "llm_call" or "tool_call" + applies_to: "llm_call", "tool_call", or a custom step type with a "_call" suffix action: The action taken (deny, steer, observe) matched: Whether the control evaluator matched confidence: Confidence score from the evaluator (0.0-1.0) @@ -90,8 +90,10 @@ class ControlExecutionEvent(BaseModel): check_stage: Literal["pre", "post"] = Field( ..., description="Check stage: 'pre' or 'post'" ) - applies_to: Literal["llm_call", "tool_call"] = Field( - ..., description="Type of call: 'llm_call' or 'tool_call'" + applies_to: str = Field( + ..., + min_length=1, + description="Type of call or custom step type with a '_call' suffix", ) # Result @@ -278,7 +280,7 @@ class EventQueryRequest(BaseModel): actions: Filter by actions (deny, steer, observe) matched: Filter by matched status check_stages: Filter by check stages (pre, post) - applies_to: Filter by call type (llm_call, tool_call) + applies_to: Filter by call type or custom call type start_time: Filter events after this time end_time: Filter events before this time limit: Maximum number of events to return @@ -310,8 +312,8 @@ class EventQueryRequest(BaseModel): check_stages: list[Literal["pre", "post"]] | None = Field( default=None, description="Filter by check stages" ) - applies_to: list[Literal["llm_call", "tool_call"]] | None = Field( - default=None, description="Filter by call types" + applies_to: list[str] | None = Field( + default=None, description="Filter by call types or custom call types" ) start_time: datetime | None = Field( default=None, description="Filter events after this time" diff --git a/pyproject.toml b/pyproject.toml index d252755f..af242982 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -11,7 +11,7 @@ name = "agent-control" version = "8.7.0" description = "Agent Control - protect your AI agents with controls" -requires-python = ">=3.12" +requires-python = ">=3.12,<3.15" [tool.uv.workspace] members = [ diff --git a/sdks/python/src/agent_control/_control_registry.py b/sdks/python/src/agent_control/_control_registry.py index 5b37c47c..1f7ee2a3 100644 --- a/sdks/python/src/agent_control/_control_registry.py +++ b/sdks/python/src/agent_control/_control_registry.py @@ -17,6 +17,7 @@ from typing import Any, NotRequired, TypedDict from ._schema_derivation import derive_schemas +from .validation import ensure_step_type logger = logging.getLogger(__name__) @@ -70,7 +71,11 @@ class _RegisteredControl: # --------------------------------------------------------------------------- -def register(func: Callable[..., Any], policy: str | None = None) -> None: +def register( + func: Callable[..., Any], + policy: str | None = None, + step_type: str | None = None, +) -> None: """Register a decorated function's step schema in the registry. Extracts step metadata from the function and stores it for later retrieval @@ -83,11 +88,17 @@ def register(func: Callable[..., Any], policy: str | None = None) -> None: Args: func: The original (unwrapped) function being decorated. policy: Optional policy name (stored as metadata). + step_type: Optional explicit step type. When omitted, tool-like + functions are registered as ``tool`` and other functions as ``llm``. """ # Determine step name -- tools typically have .name or .tool_name tool_name = getattr(func, "name", None) or getattr(func, "tool_name", None) step_name: str = tool_name if isinstance(tool_name, str) else func.__name__ - step_type: str = "tool" if isinstance(tool_name, str) else "llm" + resolved_step_type = ( + ensure_step_type(step_type) + if step_type is not None + else ("tool" if isinstance(tool_name, str) else "llm") + ) # Extract description from docstring (first line only) description: str | None = None @@ -100,10 +111,10 @@ def register(func: Callable[..., Any], policy: str | None = None) -> None: if policy is not None: metadata["policy"] = policy - key = _step_key(step_type, step_name) + key = _step_key(resolved_step_type, step_name) registered = _RegisteredControl( func=func, - step_type=step_type, + step_type=resolved_step_type, step_name=step_name, description=description, metadata=metadata, @@ -114,10 +125,10 @@ def register(func: Callable[..., Any], policy: str | None = None) -> None: logger.debug( "Overwriting previously registered step '%s' (type=%s)", step_name, - step_type, + resolved_step_type, ) _registered_steps[key] = registered - logger.debug("Registered step schema: %s (type=%s)", step_name, step_type) + logger.debug("Registered step schema: %s (type=%s)", step_name, resolved_step_type) def get_registered_steps() -> list[StepSchemaDict]: diff --git a/sdks/python/src/agent_control/control_decorators.py b/sdks/python/src/agent_control/control_decorators.py index 66e70917..fee78333 100644 --- a/sdks/python/src/agent_control/control_decorators.py +++ b/sdks/python/src/agent_control/control_decorators.py @@ -53,6 +53,7 @@ async def chat(message: str) -> str: ) from agent_control.settings import get_settings from agent_control.tracing import _generate_span_id, get_current_trace_id, get_trace_and_span_ids +from agent_control.validation import ensure_step_type logger = get_logger(__name__) @@ -96,6 +97,7 @@ class ControlContext: span_id: str start_time: float step_name: str | None = None + step_type: str | None = None # Stats (mutually exclusive: errors vs matches vs non_matches) total_executions: int = 0 @@ -126,13 +128,23 @@ def log_end(self) -> None: def pre_payload(self) -> dict[str, Any]: """Build payload for pre-execution check (supports tool call detection).""" return _create_evaluation_payload( - self.func, self.args, self.kwargs, output=None, step_name=self.step_name + self.func, + self.args, + self.kwargs, + output=None, + step_name=self.step_name, + explicit_step_type=self.step_type, ) def post_payload(self, output: Any) -> dict[str, Any]: """Build payload for post-execution check (supports tool call detection).""" return _create_evaluation_payload( - self.func, self.args, self.kwargs, output=output, step_name=self.step_name + self.func, + self.args, + self.kwargs, + output=output, + step_name=self.step_name, + explicit_step_type=self.step_type, ) def process_result(self, result: dict[str, Any], check_stage: str) -> None: @@ -478,7 +490,8 @@ def _create_evaluation_payload( args: tuple, kwargs: dict, output: Any = None, - step_name: str | None = None + step_name: str | None = None, + explicit_step_type: str | None = None, ) -> dict[str, Any]: """ Create evaluation payload for server, detecting if it's a tool step or LLM step. @@ -491,6 +504,7 @@ def _create_evaluation_payload( kwargs: Function keyword arguments output: Function output (None for pre-execution) step_name: Optional explicit step name to override auto-detection + explicit_step_type: Optional explicit step type to override inference """ sig = inspect.signature(func) bound = sig.bind(*args, **kwargs) @@ -505,21 +519,40 @@ def _create_evaluation_payload( getattr(func, "name", None) is not None or getattr(func, "tool_name", None) is not None ) - step_type = "tool" if is_tool else "llm" + inferred_type = "tool" if is_tool else "llm" else: # Auto-detect: Check if function has tool_name from @tool decorator tool_name = getattr(func, "name", None) or getattr(func, "tool_name", None) if tool_name: determined_name = tool_name - step_type = "tool" + inferred_type = "tool" else: determined_name = func.__name__ - step_type = "llm" + inferred_type = "llm" + + step_type = ( + ensure_step_type(explicit_step_type) + if explicit_step_type is not None + else inferred_type + ) + + if step_type not in ("llm", "tool"): + # Custom step types use the complete bound argument mapping. Unlike an + # LLM step, there is no generic convention for selecting one argument + # as the prompt/input for a retriever, trace, session, etc. + return { + "type": step_type, + "name": determined_name, + "input": dict(bound.arguments), + "output": output if isinstance(output, (str, int, float, bool, dict, list)) else ( + None if output is None else str(output) + ), + } if step_type == "tool": # This is a tool step return { - "type": "tool", + "type": step_type, "name": determined_name, "input": dict(bound.arguments), "output": output if isinstance(output, (str, int, float, bool, dict, list)) else ( @@ -530,7 +563,7 @@ def _create_evaluation_payload( # This is an LLM step input_data = _extract_input_from_args(func, args, kwargs) return { - "type": "llm", + "type": step_type, "name": determined_name, "input": input_data, "output": output if isinstance(output, (str, int, float, bool, dict, list)) else ( @@ -699,6 +732,7 @@ async def _execute_with_control( kwargs: dict, is_async: bool, step_name: str | None = None, + step_type: str | None = None, ) -> Any: """ Core control execution logic for both async and sync functions. @@ -717,6 +751,7 @@ async def _execute_with_control( kwargs: Keyword arguments for the function is_async: Whether the wrapped function is async step_name: Optional explicit step name for control matching + step_type: Optional explicit step type for control matching Returns: The result of the wrapped function @@ -751,6 +786,7 @@ async def _execute_with_control( span_id=span_id, start_time=time.perf_counter(), step_name=step_name, + step_type=step_type, ) ctx.log_start() @@ -772,7 +808,11 @@ async def _execute_with_control( ctx.log_end() -def control(policy: str | None = None, step_name: str | None = None) -> Callable[[F], F]: +def control( + policy: str | None = None, + step_name: str | None = None, + step_type: str | None = None, +) -> Callable[[F], F]: """ Decorator to apply server-defined controls at this code location. @@ -784,6 +824,9 @@ def control(policy: str | None = None, step_name: str | None = None) -> Callable clarity in code when multiple policies exist. step_name: Optional custom name for this step. If not provided, uses the function name. + step_type: Optional custom type for this step. If not provided, tool-like + functions are classified as ``tool`` and other functions as + ``llm``. Returns: Decorated function @@ -835,6 +878,9 @@ async def handle_user_input(user_message: str) -> str: POST /api/v1/agents/{agent_name}/policies/{policy_id} POST /api/v1/agents/{agent_name}/controls/{control_id} """ + if step_type is not None: + step_type = ensure_step_type(step_type) + # The policy parameter is for documentation only - the server evaluates # controls associated with the agent via policy and direct links. _ = policy @@ -843,12 +889,17 @@ def decorator(func: F) -> F: # Register this function's step schema for auto-discovery by init() from agent_control._control_registry import register - register(func, policy) + register(func, policy=policy, step_type=step_type) @functools.wraps(func) async def async_wrapper(*args: Any, **kwargs: Any) -> Any: return await _execute_with_control( - func, args, kwargs, is_async=True, step_name=step_name + func, + args, + kwargs, + is_async=True, + step_name=step_name, + step_type=step_type, ) # Copy over ALL attributes from the original function (important for LangChain tools) @@ -862,7 +913,14 @@ async def async_wrapper(*args: Any, **kwargs: Any) -> Any: @functools.wraps(func) def sync_wrapper(*args: Any, **kwargs: Any) -> Any: return asyncio.run( - _execute_with_control(func, args, kwargs, is_async=False, step_name=step_name) + _execute_with_control( + func, + args, + kwargs, + is_async=False, + step_name=step_name, + step_type=step_type, + ) ) if inspect.iscoroutinefunction(func): diff --git a/sdks/python/src/agent_control/evaluation.py b/sdks/python/src/agent_control/evaluation.py index 2a045bb5..3a8ec91b 100644 --- a/sdks/python/src/agent_control/evaluation.py +++ b/sdks/python/src/agent_control/evaluation.py @@ -24,7 +24,7 @@ from .evaluation_events import build_control_execution_events, enqueue_observability_events from .observability import is_observability_enabled from .tracing import get_trace_and_span_ids -from .validation import ensure_agent_name +from .validation import ensure_agent_name, ensure_step_type _RuntimePostEvaluation = Callable[..., Awaitable[httpx.Response]] @@ -523,7 +523,7 @@ async def evaluate_controls( context: dict[str, Any] | None = None, tools: list[dict[str, JSONValue]] | None = None, ground_truth: JSONValue | None = None, - step_type: Literal["tool", "llm"] = "llm", + step_type: str = "llm", stage: Literal["pre", "post"] = "pre", agent_name: str, target_type: str | None = None, @@ -541,6 +541,8 @@ async def evaluate_controls( the cached controls were fetched for the session target and would otherwise drive stale local-first evaluation. """ + step_type = ensure_step_type(step_type) + if state.server_url is None: raise RuntimeError("Server URL not configured. Call agent_control.init() first.") diff --git a/sdks/python/src/agent_control/evaluation_events.py b/sdks/python/src/agent_control/evaluation_events.py index 8db75f63..3a2c1a60 100644 --- a/sdks/python/src/agent_control/evaluation_events.py +++ b/sdks/python/src/agent_control/evaluation_events.py @@ -2,7 +2,6 @@ from collections.abc import Mapping from datetime import UTC, datetime -from typing import Literal from agent_control_models import ( ControlDefinition, @@ -62,9 +61,13 @@ def observability_metadata( ) -def map_applies_to(step_type: str) -> Literal["llm_call", "tool_call"]: +def map_applies_to(step_type: str) -> str: """Map Agent Control step types to observability applies_to values.""" - return "tool_call" if step_type == "tool" else "llm_call" + if step_type == "tool": + return "tool_call" + if step_type == "llm": + return "llm_call" + return f"{step_type}_call" def _resolve_event_trace_context( diff --git a/sdks/python/src/agent_control/integrations/_core.py b/sdks/python/src/agent_control/integrations/_core.py index c9587dcd..570fb395 100644 --- a/sdks/python/src/agent_control/integrations/_core.py +++ b/sdks/python/src/agent_control/integrations/_core.py @@ -53,7 +53,7 @@ async def _evaluate_and_enforce( context: dict[str, Any] | None = None, tools: list[dict[str, JSONValue]] | None = None, ground_truth: JSONValue | None = None, - step_type: Literal["tool", "llm"] = "llm", + step_type: str = "llm", stage: Literal["pre", "post"] = "pre", ) -> EvaluationResult: """Evaluate controls and enforce fail-closed blocking semantics.""" 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 4f977440..28345bcd 100644 --- a/sdks/python/src/agent_control/integrations/google_adk/plugin.py +++ b/sdks/python/src/agent_control/integrations/google_adk/plugin.py @@ -58,7 +58,12 @@ class AgentControlPlugin(BasePlugin): - """Google ADK plugin that enforces Agent Control across model and tool hooks.""" + """Google ADK plugin that enforces Agent Control across model and tool hooks. + + Google ADK lifecycle callbacks expose model and tool operations only, so + this adapter intentionally uses the built-in ``llm`` and ``tool`` step + types. Custom step types remain available through the core SDK APIs. + """ name = "agent-control-google-adk" diff --git a/sdks/python/src/agent_control/integrations/strands/plugin.py b/sdks/python/src/agent_control/integrations/strands/plugin.py index 367a673c..12dd66da 100644 --- a/sdks/python/src/agent_control/integrations/strands/plugin.py +++ b/sdks/python/src/agent_control/integrations/strands/plugin.py @@ -71,6 +71,9 @@ class AgentControlPlugin(Plugin): The Agent Control server is required for control distribution and policy assignment. Controls may specify execution="sdk" or execution="server". + Strands lifecycle events handled by this adapter are model and tool events, + so its internal step-type annotations intentionally remain limited to + ``llm`` and ``tool``. Custom types remain available through the core SDK APIs. """ name = "agent-control-plugin" diff --git a/sdks/python/src/agent_control/validation.py b/sdks/python/src/agent_control/validation.py index 874b0072..dc3d2a6f 100644 --- a/sdks/python/src/agent_control/validation.py +++ b/sdks/python/src/agent_control/validation.py @@ -20,3 +20,10 @@ def ensure_agent_name(value: str, field_name: str = "agent_name") -> str: f"{field_name} may only contain lowercase letters, digits, ':', '_' or '-'" ) return normalized + + +def ensure_step_type(value: str, field_name: str = "step_type") -> str: + """Return a normalized non-empty step type or raise ``ValueError``.""" + if not isinstance(value, str) or not value.strip(): + raise ValueError(f"{field_name} must be a non-empty string") + return value.strip() diff --git a/sdks/python/tests/test_control_decorators.py b/sdks/python/tests/test_control_decorators.py index ca059c55..808c30c4 100644 --- a/sdks/python/tests/test_control_decorators.py +++ b/sdks/python/tests/test_control_decorators.py @@ -190,6 +190,57 @@ async def chat(message: str) -> str: result = await chat("Hello!") assert result == "Response to: Hello!" + @pytest.mark.asyncio + async def test_explicit_step_type_is_used_for_pre_and_post_payloads( + self, mock_agent, mock_safe_response + ): + """Explicit custom types reach both decorator evaluation stages.""" + payloads = [] + + async def mock_evaluate(*args, **kwargs): + payloads.append(args[1]) + return mock_safe_response + + with patch("agent_control.control_decorators._get_current_agent", return_value=mock_agent), \ + patch("agent_control.control_decorators._evaluate", side_effect=mock_evaluate): + + @control(step_name="retrieve_documents", step_type=" retriever ") + async def retrieve(query: str, top_k: int, filters: dict[str, str]) -> str: + return query + + await retrieve("hello", 5, {"language": "en"}) + + assert [payload["type"] for payload in payloads] == ["retriever", "retriever"] + assert payloads[0]["name"] == "retrieve_documents" + assert payloads[0]["input"] == { + "query": "hello", + "top_k": 5, + "filters": {"language": "en"}, + } + + def test_explicit_step_type_overrides_tool_inference( + self, mock_agent, mock_safe_response + ): + """An explicit LLM type wins even when function metadata looks like a tool.""" + payloads = [] + + async def mock_evaluate(*args, **kwargs): + payloads.append(args[1]) + return mock_safe_response + + def search(query: str) -> str: + return query + + search.name = "search" # type: ignore[attr-defined] + + with patch("agent_control.control_decorators._get_current_agent", return_value=mock_agent), \ + patch("agent_control.control_decorators._evaluate", side_effect=mock_evaluate): + guarded_search = control(step_type="llm")(search) + guarded_search("hello") + + assert [payload["type"] for payload in payloads] == ["llm", "llm"] + assert payloads[0]["input"] == "hello" + # ============================================================================= # CONTROL NAME TESTS diff --git a/sdks/python/tests/test_control_registry.py b/sdks/python/tests/test_control_registry.py index 3ba2d320..a7872400 100644 --- a/sdks/python/tests/test_control_registry.py +++ b/sdks/python/tests/test_control_registry.py @@ -67,6 +67,16 @@ def search_db(query: str, limit: int = 10) -> str: assert steps[0]["type"] == "tool" assert steps[0]["name"] == "search_db" + def test_register_explicit_custom_type(self) -> None: + """An explicit type overrides the default LLM inference.""" + + def retrieve(query: str) -> str: + ... + + register(retrieve, step_type=" retriever ") + + assert get_registered_steps()[0]["type"] == "retriever" + def test_register_with_policy(self) -> None: # Given a typed function and an explicit policy value at registration time. def my_func(x: str) -> str: @@ -286,6 +296,28 @@ def _lookup(query: str) -> str: assert steps[0]["type"] == "tool" assert steps[0]["name"] == "lookup_tool" + def test_decorator_registers_explicit_custom_type(self) -> None: + """A decorator-supplied type is retained by auto-discovery.""" + from agent_control.control_decorators import control + + @control(step_type="retriever") + async def retrieve(query: str) -> str: + return query + + steps = get_registered_steps() + + assert steps[0]["type"] == "retriever" + + @pytest.mark.parametrize("step_type", ["", " ", 123]) + def test_explicit_type_must_be_non_empty_string(self, step_type: object) -> None: + """Explicit registry types reject invalid values.""" + + def step(query: str) -> str: + ... + + with pytest.raises(ValueError, match="non-empty string"): + register(step, step_type=step_type) # type: ignore[arg-type] + def test_stacked_decorators_deduplicate(self) -> None: """Stacking @control() twice on the same function deduplicates by name.""" diff --git a/sdks/python/tests/test_evaluation.py b/sdks/python/tests/test_evaluation.py index 885a69a7..d598c118 100644 --- a/sdks/python/tests/test_evaluation.py +++ b/sdks/python/tests/test_evaluation.py @@ -73,6 +73,40 @@ def json(self) -> dict[str, object]: ) +@pytest.mark.asyncio +async def test_check_evaluation_sends_custom_step_type_to_server(): + """The server path sends one generic Step inside EvaluationRequest.""" + + class DummyResponse: + def raise_for_status(self) -> None: + return None + + def json(self) -> dict[str, object]: + return {"is_safe": True, "confidence": 1.0} + + client = MagicMock() + client.http_client = MagicMock() + client.http_client.post = AsyncMock(return_value=DummyResponse()) + + await evaluation.check_evaluation( + client=client, + agent_name="Agent-Example_01", + step={"type": "trace", "name": "trace-check", "input": "hello"}, + stage="post", + ) + + request = client.http_client.post.await_args.kwargs["json"] + assert request["step"] == { + "type": "trace", + "name": "trace-check", + "input": "hello", + "output": None, + "context": None, + "tools": None, + "ground_truth": None, + } + + @pytest.mark.asyncio async def test_evaluate_controls_requires_server_url(): """evaluate_controls should require server_url to be configured.""" @@ -127,6 +161,45 @@ async def test_evaluate_controls_with_context(monkeypatch): assert mock_check.call_args is not None +@pytest.mark.parametrize("step_type", ["trace", "session"]) +@pytest.mark.asyncio +async def test_evaluate_controls_accepts_custom_step_type(monkeypatch, step_type): + """Direct evaluation preserves custom types in the runtime Step.""" + 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="trace_check", + step_type=step_type, + input="trace input", + output="trace output", + stage="post", + agent_name="test-bot", + ) + + step = mock_check.call_args.kwargs["step"] + assert step.type == step_type + assert step.name == "trace_check" + + +@pytest.mark.asyncio +async def test_evaluate_controls_rejects_empty_step_type(monkeypatch): + """Direct evaluation rejects empty custom types before evaluation.""" + 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"): + with pytest.raises(ValueError, match="non-empty string"): + await evaluation.evaluate_controls( + step_name="trace_check", + step_type="", + agent_name="test-bot", + ) + + mock_check.assert_not_called() + + @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.""" diff --git a/sdks/python/tests/test_observability_updates.py b/sdks/python/tests/test_observability_updates.py index f9bdf6a1..475705c2 100644 --- a/sdks/python/tests/test_observability_updates.py +++ b/sdks/python/tests/test_observability_updates.py @@ -51,6 +51,21 @@ def test_maps_tool_to_tool_call(self): def test_maps_llm_to_llm_call(self): assert map_applies_to("llm") == "llm_call" + def test_maps_custom_step_type_to_call_type(self): + assert map_applies_to("retriever") == "retriever_call" + assert map_applies_to("trace") == "trace_call" + assert map_applies_to("session") == "session_call" + + def test_custom_step_type_round_trips_in_observability_query(self): + from agent_control_models import EventQueryRequest + + request = EventQueryRequest(applies_to=["retriever_call", "trace_call"]) + + assert request.model_dump(mode="json")["applies_to"] == [ + "retriever_call", + "trace_call", + ] + class TestMergeResults: def _make_response(self, **kwargs): @@ -238,6 +253,32 @@ def test_builds_events_with_trace_context(self): assert event.evaluator_name == "regex" assert event.selector_path == "input" + def test_builds_custom_type_event_without_llm_classification(self): + response = self._make_response(matches=[self._make_match(1, "ctrl-1")]) + request = self._make_request(step_type="retriever") + control_lookup = { + 1: self._make_control( + 1, + "ctrl-1", + { + "evaluator": {"name": "regex", "config": {"pattern": "test"}}, + "selector": {"path": "input"}, + }, + ).control + } + + events = build_control_execution_events( + response, + request, + control_lookup, + "trace123", + "span456", + "test-agent", + ) + + assert events[0].applies_to == "retriever_call" + assert events[0].model_dump(mode="json")["applies_to"] == "retriever_call" + def test_uses_safe_selected_data_preview_as_event_input(self): response = self._make_response( matches=[ diff --git a/sdks/typescript/src/generated/funcs/observability-query-events.ts b/sdks/typescript/src/generated/funcs/observability-query-events.ts index 92b93bae..81b7dc24 100644 --- a/sdks/typescript/src/generated/funcs/observability-query-events.ts +++ b/sdks/typescript/src/generated/funcs/observability-query-events.ts @@ -41,7 +41,7 @@ import { Result } from "../types/fp.js"; * - actions: Filter by actions (deny, steer, observe) * - matched: Filter by matched status * - check_stages: Filter by check stage (pre, post) - * - applies_to: Filter by call type (llm_call, tool_call) + * - applies_to: Filter by call type or custom call type * - start_time/end_time: Filter by time range * * Results are paginated with limit/offset. diff --git a/sdks/typescript/src/generated/models/control-execution-event.ts b/sdks/typescript/src/generated/models/control-execution-event.ts index b9f4dd68..a452e60c 100644 --- a/sdks/typescript/src/generated/models/control-execution-event.ts +++ b/sdks/typescript/src/generated/models/control-execution-event.ts @@ -16,20 +16,6 @@ import { } from "./action-decision.js"; import { SDKValidationError } from "./errors/sdk-validation-error.js"; -/** - * Type of call: 'llm_call' or 'tool_call' - */ -export const ControlExecutionEventAppliesTo = { - LlmCall: "llm_call", - ToolCall: "tool_call", -} as const; -/** - * Type of call: 'llm_call' or 'tool_call' - */ -export type ControlExecutionEventAppliesTo = OpenEnum< - typeof ControlExecutionEventAppliesTo ->; - /** * Check stage: 'pre' or 'post' */ @@ -62,7 +48,7 @@ export type CheckStage = OpenEnum; * control_id: Database ID of the control * control_name: Name of the control (denormalized for queries) * check_stage: "pre" (before execution) or "post" (after execution) - * applies_to: "llm_call" or "tool_call" + * applies_to: "llm_call", "tool_call", or a custom step type with a "_call" suffix * action: The action taken (deny, steer, observe) * matched: Whether the control evaluator matched * confidence: Confidence score from the evaluator (0.0-1.0) @@ -80,9 +66,9 @@ export type ControlExecutionEvent = { */ agentName: string; /** - * Type of call: 'llm_call' or 'tool_call' + * Type of call or custom step type with a '_call' suffix */ - appliesTo: ControlExecutionEventAppliesTo; + appliesTo: string; /** * Check stage: 'pre' or 'post' */ @@ -141,17 +127,6 @@ export type ControlExecutionEvent = { traceId: string; }; -/** @internal */ -export const ControlExecutionEventAppliesTo$inboundSchema: z.ZodMiniType< - ControlExecutionEventAppliesTo, - unknown -> = openEnums.inboundSchema(ControlExecutionEventAppliesTo); -/** @internal */ -export const ControlExecutionEventAppliesTo$outboundSchema: z.ZodMiniType< - string, - ControlExecutionEventAppliesTo -> = openEnums.outboundSchema(ControlExecutionEventAppliesTo); - /** @internal */ export const CheckStage$inboundSchema: z.ZodMiniType = openEnums.inboundSchema(CheckStage); @@ -167,7 +142,7 @@ export const ControlExecutionEvent$inboundSchema: z.ZodMiniType< z.object({ action: ActionDecision$inboundSchema, agent_name: types.string(), - applies_to: ControlExecutionEventAppliesTo$inboundSchema, + applies_to: types.string(), check_stage: CheckStage$inboundSchema, confidence: types.number(), control_execution_id: types.optional(types.string()), @@ -229,7 +204,7 @@ export const ControlExecutionEvent$outboundSchema: z.ZodMiniType< z.object({ action: ActionDecision$outboundSchema, agentName: z.string(), - appliesTo: ControlExecutionEventAppliesTo$outboundSchema, + appliesTo: z.string(), checkStage: CheckStage$outboundSchema, confidence: z.number(), controlExecutionId: z.optional(z.string()), diff --git a/sdks/typescript/src/generated/models/event-query-request.ts b/sdks/typescript/src/generated/models/event-query-request.ts index 3484a7db..a40a9de7 100644 --- a/sdks/typescript/src/generated/models/event-query-request.ts +++ b/sdks/typescript/src/generated/models/event-query-request.ts @@ -10,12 +10,6 @@ import { ActionDecision$outboundSchema, } from "./action-decision.js"; -export const AppliesTo = { - LlmCall: "llm_call", - ToolCall: "tool_call", -} as const; -export type AppliesTo = ClosedEnum; - export const CheckStages = { Pre: "pre", Post: "post", @@ -38,7 +32,7 @@ export type CheckStages = ClosedEnum; * actions: Filter by actions (deny, steer, observe) * matched: Filter by matched status * check_stages: Filter by check stages (pre, post) - * applies_to: Filter by call type (llm_call, tool_call) + * applies_to: Filter by call type or custom call type * start_time: Filter events after this time * end_time: Filter events before this time * limit: Maximum number of events to return @@ -54,9 +48,9 @@ export type EventQueryRequest = { */ agentName?: string | null | undefined; /** - * Filter by call types + * Filter by call types or custom call types */ - appliesTo?: Array | null | undefined; + appliesTo?: Array | null | undefined; /** * Filter by check stages */ @@ -99,11 +93,6 @@ export type EventQueryRequest = { traceId?: string | null | undefined; }; -/** @internal */ -export const AppliesTo$outboundSchema: z.ZodMiniEnum = z.enum( - AppliesTo, -); - /** @internal */ export const CheckStages$outboundSchema: z.ZodMiniEnum = z .enum(CheckStages); @@ -133,7 +122,7 @@ export const EventQueryRequest$outboundSchema: z.ZodMiniType< z.object({ actions: z.optional(z.nullable(z.array(ActionDecision$outboundSchema))), agentName: z.optional(z.nullable(z.string())), - appliesTo: z.optional(z.nullable(z.array(AppliesTo$outboundSchema))), + appliesTo: z.optional(z.nullable(z.array(z.string()))), checkStages: z.optional(z.nullable(z.array(CheckStages$outboundSchema))), controlExecutionId: z.optional(z.nullable(z.string())), controlIds: z.optional(z.nullable(z.array(z.int()))), diff --git a/sdks/typescript/src/generated/sdk/observability.ts b/sdks/typescript/src/generated/sdk/observability.ts index f11f88ce..6c13482f 100644 --- a/sdks/typescript/src/generated/sdk/observability.ts +++ b/sdks/typescript/src/generated/sdk/observability.ts @@ -54,7 +54,7 @@ export class Observability extends ClientSDK { * - actions: Filter by actions (deny, steer, observe) * - matched: Filter by matched status * - check_stages: Filter by check stage (pre, post) - * - applies_to: Filter by call type (llm_call, tool_call) + * - applies_to: Filter by call type or custom call type * - start_time/end_time: Filter by time range * * Results are paginated with limit/offset. diff --git a/sdks/typescript/tests/generated-smoke.test.ts b/sdks/typescript/tests/generated-smoke.test.ts index dda98e3a..2196eda7 100644 --- a/sdks/typescript/tests/generated-smoke.test.ts +++ b/sdks/typescript/tests/generated-smoke.test.ts @@ -3,6 +3,7 @@ import path from "node:path"; import { describe, expect, it } from "vitest"; +import { eventQueryRequestToJSON } from "../src/generated/models/event-query-request"; import { stepToJSON } from "../src/generated/models/step"; describe("generated client layout", () => { @@ -58,4 +59,12 @@ describe("generated client layout", () => { ], }); }); + + it("serializes custom observability step types", () => { + const serialized = eventQueryRequestToJSON({ + appliesTo: ["retriever_call", "trace_call"], + }); + + expect(JSON.parse(serialized).applies_to).toEqual(["retriever_call", "trace_call"]); + }); }); diff --git a/server/pyproject.toml b/server/pyproject.toml index 90dfa736..4f938ef4 100644 --- a/server/pyproject.toml +++ b/server/pyproject.toml @@ -2,7 +2,7 @@ name = "agent-control-server" version = "8.7.0" description = "Server for Agent Control - manage and evaluate controls for AI agents" -requires-python = ">=3.12" +requires-python = ">=3.12,<3.15" # Note: agent-control-models, agent-control-engine, and agent-control-telemetry # are bundled at build time # Note: agent-control-evaluators is a runtime dependency (NOT vendored) to avoid diff --git a/server/src/agent_control_server/endpoints/observability.py b/server/src/agent_control_server/endpoints/observability.py index 4e52377b..e9437526 100644 --- a/server/src/agent_control_server/endpoints/observability.py +++ b/server/src/agent_control_server/endpoints/observability.py @@ -149,7 +149,7 @@ async def query_events( - actions: Filter by actions (deny, steer, observe) - matched: Filter by matched status - check_stages: Filter by check stage (pre, post) - - applies_to: Filter by call type (llm_call, tool_call) + - applies_to: Filter by call type or custom call type - start_time/end_time: Filter by time range Results are paginated with limit/offset. diff --git a/server/tests/test_evaluation_e2e.py b/server/tests/test_evaluation_e2e.py index fbe9a483..e4422591 100644 --- a/server/tests/test_evaluation_e2e.py +++ b/server/tests/test_evaluation_e2e.py @@ -63,6 +63,31 @@ def test_evaluation_no_policy(client: TestClient): assert not resp.json()["matches"] +def test_evaluation_accepts_custom_step_type(client: TestClient): + """The evaluation endpoint accepts custom Step.type values.""" + agent_name = f"agent-{uuid.uuid4().hex[:12]}" + client.post( + "/api/v1/agents/initAgent", + json={"agent": {"agent_name": agent_name}, "steps": []}, + ) + + request = EvaluationRequest( + agent_name=agent_name, + step=Step( + type="retriever", + name="retrieve-documents", + input={"query": "anything"}, + ), + stage="pre", + ) + + response = client.post("/api/v1/evaluation", json=request.model_dump(mode="json")) + + assert response.status_code == 200 + assert response.json()["is_safe"] is True + assert not response.json()["matches"] + + def test_evaluation_empty_policy(client: TestClient): """Test that an agent with an empty policy is safe.""" # Given: an empty policy diff --git a/ui/src/core/api/generated/api-types.ts b/ui/src/core/api/generated/api-types.ts index fe859d89..c7ca0ef1 100644 --- a/ui/src/core/api/generated/api-types.ts +++ b/ui/src/core/api/generated/api-types.ts @@ -748,7 +748,7 @@ export interface paths { * - actions: Filter by actions (deny, steer, observe) * - matched: Filter by matched status * - check_stages: Filter by check stage (pre, post) - * - applies_to: Filter by call type (llm_call, tool_call) + * - applies_to: Filter by call type or custom call type * - start_time/end_time: Filter by time range * * Results are paginated with limit/offset. @@ -1526,7 +1526,7 @@ export interface components { * control_id: Database ID of the control * control_name: Name of the control (denormalized for queries) * check_stage: "pre" (before execution) or "post" (after execution) - * applies_to: "llm_call" or "tool_call" + * applies_to: "llm_call", "tool_call", or a custom call type * action: The action taken (deny, steer, observe) * matched: Whether the control evaluator matched * confidence: Confidence score from the evaluator (0.0-1.0) @@ -1564,10 +1564,9 @@ export interface components { agent_name: string; /** * Applies To - * @description Type of call: 'llm_call' or 'tool_call' - * @enum {string} + * @description Type of call or custom call type */ - applies_to: 'llm_call' | 'tool_call'; + applies_to: string; /** * Check Stage * @description Check stage: 'pre' or 'post' @@ -2295,7 +2294,7 @@ export interface components { * actions: Filter by actions (deny, steer, observe) * matched: Filter by matched status * check_stages: Filter by check stages (pre, post) - * applies_to: Filter by call type (llm_call, tool_call) + * applies_to: Filter by call type or custom call type * start_time: Filter events after this time * end_time: Filter events before this time * limit: Maximum number of events to return @@ -2326,9 +2325,9 @@ export interface components { agent_name?: string | null; /** * Applies To - * @description Filter by call types + * @description Filter by call types or custom call types */ - applies_to?: ('llm_call' | 'tool_call')[] | null; + applies_to?: string[] | null; /** * Check Stages * @description Filter by check stages