@@ -573,6 +573,12 @@ def decorator(func):
573573 step_kwargs ["name" ] = func .__name__
574574 step_name = step_kwargs ["name" ]
575575
576+ _raw_step_type = step_kwargs .get ("step_type" , enums .StepType .USER_CALL )
577+ if isinstance (_raw_step_type , str ):
578+ _raw_step_type = enums .StepType (_raw_step_type )
579+ step_type = _raw_step_type
580+ step_kwargs ["step_type" ] = step_type
581+
576582 # Check if it's a generator function
577583 if inspect .isgeneratorfunction (func ):
578584 # For sync generators, use class-based approach to delay trace creation
@@ -601,7 +607,7 @@ def __next__(self):
601607 self ._step , self ._is_root_step , self ._token = (
602608 _create_and_initialize_step (
603609 step_name = step_name ,
604- step_type = enums . StepType . USER_CALL ,
610+ step_type = step_type ,
605611 inputs = None ,
606612 output = None ,
607613 metadata = None ,
@@ -847,6 +853,12 @@ def decorator(func):
847853 step_kwargs ["name" ] = func .__name__
848854 step_name = step_kwargs ["name" ]
849855
856+ _raw_step_type = step_kwargs .get ("step_type" , enums .StepType .USER_CALL )
857+ if isinstance (_raw_step_type , str ):
858+ _raw_step_type = enums .StepType (_raw_step_type )
859+ step_type = _raw_step_type
860+ step_kwargs ["step_type" ] = step_type
861+
850862 if asyncio .iscoroutinefunction (func ) or inspect .isasyncgenfunction (func ):
851863 # Check if it's specifically an async generator function
852864 if inspect .isasyncgenfunction (func ):
@@ -873,7 +885,7 @@ async def __anext__(self):
873885 self ._step , self ._is_root_step , self ._token = (
874886 _create_and_initialize_step (
875887 step_name = step_name ,
876- step_type = enums . StepType . USER_CALL ,
888+ step_type = step_type ,
877889 inputs = None ,
878890 output = None ,
879891 metadata = None ,
@@ -1772,6 +1784,9 @@ def _process_wrapper_inputs_and_outputs(
17721784 context_kwarg = context_kwarg ,
17731785 question_kwarg = question_kwarg ,
17741786 )
1787+ if isinstance (step , steps .ToolStep ):
1788+ step .function_name = step .name
1789+ step .arguments = inputs
17751790 _finalize_step_logging (
17761791 step = step ,
17771792 inputs = inputs ,
0 commit comments