feat(sdk): add step_type as input parameter to the control decorator and evaluate_controls function

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

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"

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 = [

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]:

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):

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.")