Skip to content

Commit 41de182

Browse files
chore(closes OPEN-10613): assign step type in trace decorator
1 parent 60c4841 commit 41de182

1 file changed

Lines changed: 17 additions & 2 deletions

File tree

src/openlayer/lib/tracing/tracer.py

Lines changed: 17 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -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

Comments
 (0)