Skip to content

Commit 9d12df0

Browse files
feat(sdk): add step_type as input parameter to the control decorator and evaluate_controls function (#262)
## Summary - Added support for arbitrary custom step types such as retriever, trace, and session across the Python SDK, telemetry, observability models, TypeScript SDK, and UI API types. Existing llm/tool inference remains unchanged. This is to accomodate controls at different levels especially Luna but not restricted to Galileo levels only. ## Scope User-facing/API changes:- Added step_type: str | None to @agent_control.control(). - Widened evaluate_controls(step_type) to accept custom strings. - Updated generated TypeScript/UI API artifacts. Internal changes:- Preserved custom decorator arguments and telemetry classifications. - Added step-type normalization and validation. - Added endpoint and regression tests. Out of scope:- Google ADK and Strands adapters remain intentionally limited to LLM/tool lifecycle events. ## Risk and Rollout Risk level: low Rollback plan: Revert this commit. ## Testing - [ x] Added or updated automated tests - [ x] Ran `make check` (or explained why not) - [x ] Manually verified behavior ## Checklist - [ ] Linked issue/spec (if applicable) - [ x] Updated docs/examples for user-facing changes - [ ] Included any required follow-up tasks
1 parent bd90753 commit 9d12df0

24 files changed

Lines changed: 409 additions & 89 deletions

File tree

‎engine/tests/test_core.py‎

Lines changed: 35 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1085,6 +1085,41 @@ def register_payload_evaluator(self):
10851085
except ValueError:
10861086
pass
10871087

1088+
@pytest.mark.asyncio
1089+
async def test_custom_step_type_filters_controls(self):
1090+
"""Controls match custom step types without falling back to LLM."""
1091+
controls = [
1092+
make_control(
1093+
1,
1094+
"retriever-control",
1095+
"test-deny",
1096+
action="deny",
1097+
config_value="retriever",
1098+
step_types=["retriever"],
1099+
),
1100+
make_control(
1101+
2,
1102+
"llm-control",
1103+
"test-deny",
1104+
action="deny",
1105+
config_value="llm",
1106+
step_types=["llm"],
1107+
),
1108+
]
1109+
engine = ControlEngine(controls)
1110+
1111+
result = await engine.process(
1112+
EvaluationRequest(
1113+
agent_name="00000000-0000-0000-0000-000000000001",
1114+
step=Step(type="retriever", name="retrieve", input="query"),
1115+
stage="pre",
1116+
)
1117+
)
1118+
1119+
assert [match.control_name for match in result.matches or []] == [
1120+
"retriever-control"
1121+
]
1122+
10881123
@pytest.mark.asyncio
10891124
async def test_step_names_filters_tasks(self):
10901125
# Given: two controls scoped to different steps

‎models/src/agent_control_models/observability.py‎

Lines changed: 8 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -46,7 +46,7 @@ class ControlExecutionEvent(BaseModel):
4646
control_id: Database ID of the control
4747
control_name: Name of the control (denormalized for queries)
4848
check_stage: "pre" (before execution) or "post" (after execution)
49-
applies_to: "llm_call" or "tool_call"
49+
applies_to: "llm_call", "tool_call", or a custom step type with a "_call" suffix
5050
action: The action taken (deny, steer, observe)
5151
matched: Whether the control evaluator matched
5252
confidence: Confidence score from the evaluator (0.0-1.0)
@@ -90,8 +90,10 @@ class ControlExecutionEvent(BaseModel):
9090
check_stage: Literal["pre", "post"] = Field(
9191
..., description="Check stage: 'pre' or 'post'"
9292
)
93-
applies_to: Literal["llm_call", "tool_call"] = Field(
94-
..., description="Type of call: 'llm_call' or 'tool_call'"
93+
applies_to: str = Field(
94+
...,
95+
min_length=1,
96+
description="Type of call or custom step type with a '_call' suffix",
9597
)
9698

9799
# Result
@@ -278,7 +280,7 @@ class EventQueryRequest(BaseModel):
278280
actions: Filter by actions (deny, steer, observe)
279281
matched: Filter by matched status
280282
check_stages: Filter by check stages (pre, post)
281-
applies_to: Filter by call type (llm_call, tool_call)
283+
applies_to: Filter by call type or custom call type
282284
start_time: Filter events after this time
283285
end_time: Filter events before this time
284286
limit: Maximum number of events to return
@@ -310,8 +312,8 @@ class EventQueryRequest(BaseModel):
310312
check_stages: list[Literal["pre", "post"]] | None = Field(
311313
default=None, description="Filter by check stages"
312314
)
313-
applies_to: list[Literal["llm_call", "tool_call"]] | None = Field(
314-
default=None, description="Filter by call types"
315+
applies_to: list[str] | None = Field(
316+
default=None, description="Filter by call types or custom call types"
315317
)
316318
start_time: datetime | None = Field(
317319
default=None, description="Filter events after this time"

‎pyproject.toml‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -11,7 +11,7 @@
1111
name = "agent-control"
1212
version = "8.7.0"
1313
description = "Agent Control - protect your AI agents with controls"
14-
requires-python = ">=3.12"
14+
requires-python = ">=3.12,<3.15"
1515

1616
[tool.uv.workspace]
1717
members = [

‎sdks/python/src/agent_control/_control_registry.py‎

Lines changed: 17 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,7 @@
1717
from typing import Any, NotRequired, TypedDict
1818

1919
from ._schema_derivation import derive_schemas
20+
from .validation import ensure_step_type
2021

2122
logger = logging.getLogger(__name__)
2223

@@ -70,7 +71,11 @@ class _RegisteredControl:
7071
# ---------------------------------------------------------------------------
7172

7273

73-
def register(func: Callable[..., Any], policy: str | None = None) -> None:
74+
def register(
75+
func: Callable[..., Any],
76+
policy: str | None = None,
77+
step_type: str | None = None,
78+
) -> None:
7479
"""Register a decorated function's step schema in the registry.
7580
7681
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:
8388
Args:
8489
func: The original (unwrapped) function being decorated.
8590
policy: Optional policy name (stored as metadata).
91+
step_type: Optional explicit step type. When omitted, tool-like
92+
functions are registered as ``tool`` and other functions as ``llm``.
8693
"""
8794
# Determine step name -- tools typically have .name or .tool_name
8895
tool_name = getattr(func, "name", None) or getattr(func, "tool_name", None)
8996
step_name: str = tool_name if isinstance(tool_name, str) else func.__name__
90-
step_type: str = "tool" if isinstance(tool_name, str) else "llm"
97+
resolved_step_type = (
98+
ensure_step_type(step_type)
99+
if step_type is not None
100+
else ("tool" if isinstance(tool_name, str) else "llm")
101+
)
91102

92103
# Extract description from docstring (first line only)
93104
description: str | None = None
@@ -100,10 +111,10 @@ def register(func: Callable[..., Any], policy: str | None = None) -> None:
100111
if policy is not None:
101112
metadata["policy"] = policy
102113

103-
key = _step_key(step_type, step_name)
114+
key = _step_key(resolved_step_type, step_name)
104115
registered = _RegisteredControl(
105116
func=func,
106-
step_type=step_type,
117+
step_type=resolved_step_type,
107118
step_name=step_name,
108119
description=description,
109120
metadata=metadata,
@@ -114,10 +125,10 @@ def register(func: Callable[..., Any], policy: str | None = None) -> None:
114125
logger.debug(
115126
"Overwriting previously registered step '%s' (type=%s)",
116127
step_name,
117-
step_type,
128+
resolved_step_type,
118129
)
119130
_registered_steps[key] = registered
120-
logger.debug("Registered step schema: %s (type=%s)", step_name, step_type)
131+
logger.debug("Registered step schema: %s (type=%s)", step_name, resolved_step_type)
121132

122133

123134
def get_registered_steps() -> list[StepSchemaDict]:

‎sdks/python/src/agent_control/control_decorators.py‎

Lines changed: 70 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -53,6 +53,7 @@ async def chat(message: str) -> str:
5353
)
5454
from agent_control.settings import get_settings
5555
from agent_control.tracing import _generate_span_id, get_current_trace_id, get_trace_and_span_ids
56+
from agent_control.validation import ensure_step_type
5657

5758
logger = get_logger(__name__)
5859

@@ -96,6 +97,7 @@ class ControlContext:
9697
span_id: str
9798
start_time: float
9899
step_name: str | None = None
100+
step_type: str | None = None
99101

100102
# Stats (mutually exclusive: errors vs matches vs non_matches)
101103
total_executions: int = 0
@@ -126,13 +128,23 @@ def log_end(self) -> None:
126128
def pre_payload(self) -> dict[str, Any]:
127129
"""Build payload for pre-execution check (supports tool call detection)."""
128130
return _create_evaluation_payload(
129-
self.func, self.args, self.kwargs, output=None, step_name=self.step_name
131+
self.func,
132+
self.args,
133+
self.kwargs,
134+
output=None,
135+
step_name=self.step_name,
136+
explicit_step_type=self.step_type,
130137
)
131138

132139
def post_payload(self, output: Any) -> dict[str, Any]:
133140
"""Build payload for post-execution check (supports tool call detection)."""
134141
return _create_evaluation_payload(
135-
self.func, self.args, self.kwargs, output=output, step_name=self.step_name
142+
self.func,
143+
self.args,
144+
self.kwargs,
145+
output=output,
146+
step_name=self.step_name,
147+
explicit_step_type=self.step_type,
136148
)
137149

138150
def process_result(self, result: dict[str, Any], check_stage: str) -> None:
@@ -478,7 +490,8 @@ def _create_evaluation_payload(
478490
args: tuple,
479491
kwargs: dict,
480492
output: Any = None,
481-
step_name: str | None = None
493+
step_name: str | None = None,
494+
explicit_step_type: str | None = None,
482495
) -> dict[str, Any]:
483496
"""
484497
Create evaluation payload for server, detecting if it's a tool step or LLM step.
@@ -491,6 +504,7 @@ def _create_evaluation_payload(
491504
kwargs: Function keyword arguments
492505
output: Function output (None for pre-execution)
493506
step_name: Optional explicit step name to override auto-detection
507+
explicit_step_type: Optional explicit step type to override inference
494508
"""
495509
sig = inspect.signature(func)
496510
bound = sig.bind(*args, **kwargs)
@@ -505,21 +519,40 @@ def _create_evaluation_payload(
505519
getattr(func, "name", None) is not None
506520
or getattr(func, "tool_name", None) is not None
507521
)
508-
step_type = "tool" if is_tool else "llm"
522+
inferred_type = "tool" if is_tool else "llm"
509523
else:
510524
# Auto-detect: Check if function has tool_name from @tool decorator
511525
tool_name = getattr(func, "name", None) or getattr(func, "tool_name", None)
512526
if tool_name:
513527
determined_name = tool_name
514-
step_type = "tool"
528+
inferred_type = "tool"
515529
else:
516530
determined_name = func.__name__
517-
step_type = "llm"
531+
inferred_type = "llm"
532+
533+
step_type = (
534+
ensure_step_type(explicit_step_type)
535+
if explicit_step_type is not None
536+
else inferred_type
537+
)
538+
539+
if step_type not in ("llm", "tool"):
540+
# Custom step types use the complete bound argument mapping. Unlike an
541+
# LLM step, there is no generic convention for selecting one argument
542+
# as the prompt/input for a retriever, trace, session, etc.
543+
return {
544+
"type": step_type,
545+
"name": determined_name,
546+
"input": dict(bound.arguments),
547+
"output": output if isinstance(output, (str, int, float, bool, dict, list)) else (
548+
None if output is None else str(output)
549+
),
550+
}
518551

519552
if step_type == "tool":
520553
# This is a tool step
521554
return {
522-
"type": "tool",
555+
"type": step_type,
523556
"name": determined_name,
524557
"input": dict(bound.arguments),
525558
"output": output if isinstance(output, (str, int, float, bool, dict, list)) else (
@@ -530,7 +563,7 @@ def _create_evaluation_payload(
530563
# This is an LLM step
531564
input_data = _extract_input_from_args(func, args, kwargs)
532565
return {
533-
"type": "llm",
566+
"type": step_type,
534567
"name": determined_name,
535568
"input": input_data,
536569
"output": output if isinstance(output, (str, int, float, bool, dict, list)) else (
@@ -699,6 +732,7 @@ async def _execute_with_control(
699732
kwargs: dict,
700733
is_async: bool,
701734
step_name: str | None = None,
735+
step_type: str | None = None,
702736
) -> Any:
703737
"""
704738
Core control execution logic for both async and sync functions.
@@ -717,6 +751,7 @@ async def _execute_with_control(
717751
kwargs: Keyword arguments for the function
718752
is_async: Whether the wrapped function is async
719753
step_name: Optional explicit step name for control matching
754+
step_type: Optional explicit step type for control matching
720755
721756
Returns:
722757
The result of the wrapped function
@@ -751,6 +786,7 @@ async def _execute_with_control(
751786
span_id=span_id,
752787
start_time=time.perf_counter(),
753788
step_name=step_name,
789+
step_type=step_type,
754790
)
755791
ctx.log_start()
756792

@@ -772,7 +808,11 @@ async def _execute_with_control(
772808
ctx.log_end()
773809

774810

775-
def control(policy: str | None = None, step_name: str | None = None) -> Callable[[F], F]:
811+
def control(
812+
policy: str | None = None,
813+
step_name: str | None = None,
814+
step_type: str | None = None,
815+
) -> Callable[[F], F]:
776816
"""
777817
Decorator to apply server-defined controls at this code location.
778818
@@ -784,6 +824,9 @@ def control(policy: str | None = None, step_name: str | None = None) -> Callable
784824
clarity in code when multiple policies exist.
785825
step_name: Optional custom name for this step. If not provided, uses
786826
the function name.
827+
step_type: Optional custom type for this step. If not provided, tool-like
828+
functions are classified as ``tool`` and other functions as
829+
``llm``.
787830
788831
Returns:
789832
Decorated function
@@ -835,6 +878,9 @@ async def handle_user_input(user_message: str) -> str:
835878
POST /api/v1/agents/{agent_name}/policies/{policy_id}
836879
POST /api/v1/agents/{agent_name}/controls/{control_id}
837880
"""
881+
if step_type is not None:
882+
step_type = ensure_step_type(step_type)
883+
838884
# The policy parameter is for documentation only - the server evaluates
839885
# controls associated with the agent via policy and direct links.
840886
_ = policy
@@ -843,12 +889,17 @@ def decorator(func: F) -> F:
843889
# Register this function's step schema for auto-discovery by init()
844890
from agent_control._control_registry import register
845891

846-
register(func, policy)
892+
register(func, policy=policy, step_type=step_type)
847893

848894
@functools.wraps(func)
849895
async def async_wrapper(*args: Any, **kwargs: Any) -> Any:
850896
return await _execute_with_control(
851-
func, args, kwargs, is_async=True, step_name=step_name
897+
func,
898+
args,
899+
kwargs,
900+
is_async=True,
901+
step_name=step_name,
902+
step_type=step_type,
852903
)
853904

854905
# 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:
862913
@functools.wraps(func)
863914
def sync_wrapper(*args: Any, **kwargs: Any) -> Any:
864915
return asyncio.run(
865-
_execute_with_control(func, args, kwargs, is_async=False, step_name=step_name)
916+
_execute_with_control(
917+
func,
918+
args,
919+
kwargs,
920+
is_async=False,
921+
step_name=step_name,
922+
step_type=step_type,
923+
)
866924
)
867925

868926
if inspect.iscoroutinefunction(func):

‎sdks/python/src/agent_control/evaluation.py‎

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -24,7 +24,7 @@
2424
from .evaluation_events import build_control_execution_events, enqueue_observability_events
2525
from .observability import is_observability_enabled
2626
from .tracing import get_trace_and_span_ids
27-
from .validation import ensure_agent_name
27+
from .validation import ensure_agent_name, ensure_step_type
2828

2929
_RuntimePostEvaluation = Callable[..., Awaitable[httpx.Response]]
3030

@@ -523,7 +523,7 @@ async def evaluate_controls(
523523
context: dict[str, Any] | None = None,
524524
tools: list[dict[str, JSONValue]] | None = None,
525525
ground_truth: JSONValue | None = None,
526-
step_type: Literal["tool", "llm"] = "llm",
526+
step_type: str = "llm",
527527
stage: Literal["pre", "post"] = "pre",
528528
agent_name: str,
529529
target_type: str | None = None,
@@ -541,6 +541,8 @@ async def evaluate_controls(
541541
the cached controls were fetched for the session target and would
542542
otherwise drive stale local-first evaluation.
543543
"""
544+
step_type = ensure_step_type(step_type)
545+
544546
if state.server_url is None:
545547
raise RuntimeError("Server URL not configured. Call agent_control.init() first.")
546548

0 commit comments

Comments
 (0)