@@ -53,6 +53,7 @@ async def chat(message: str) -> str:
5353)
5454from agent_control .settings import get_settings
5555from 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
5758logger = 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 ):
0 commit comments