Skip to content

Commit 9abe566

Browse files
gustavocidornelasgustavocidornelas
committed
chore: update tracer implementation
1 parent d91fd55 commit 9abe566

File tree

1 file changed

+143
-9
lines changed

1 file changed

+143
-9
lines changed

‎src/openlayer/lib/tracing/tracer.py‎

Lines changed: 143 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -540,14 +540,84 @@ async def async_function_wrapper(*func_args, **func_kwargs):
540540
**step_kwargs,
541541
) as step:
542542
output = exception = None
543+
guardrail_metadata = {}
543544

544545
try:
545-
output = await func(*func_args, **func_kwargs)
546+
# Apply input guardrails if provided
547+
if guardrails:
548+
try:
549+
inputs = _extract_function_inputs(
550+
func_signature=func_signature,
551+
func_args=func_args,
552+
func_kwargs=func_kwargs,
553+
context_kwarg=context_kwarg,
554+
)
555+
556+
# Process inputs through guardrails
557+
modified_inputs, input_metadata = (
558+
_apply_guardrails_to_inputs(
559+
guardrails,
560+
inputs,
561+
)
562+
)
563+
guardrail_metadata.update(input_metadata)
564+
565+
# Execute function with potentially modified inputs
566+
if modified_inputs != inputs:
567+
# Reconstruct function arguments from modified inputs
568+
bound = func_signature.bind(
569+
*func_args, **func_kwargs
570+
)
571+
bound.apply_defaults()
572+
573+
# Update bound arguments with modified values
574+
for (
575+
param_name,
576+
modified_value,
577+
) in modified_inputs.items():
578+
if param_name in bound.arguments:
579+
bound.arguments[param_name] = (
580+
modified_value
581+
)
582+
583+
output = await func(*bound.args, **bound.kwargs)
584+
else:
585+
output = await func(*func_args, **func_kwargs)
586+
except Exception as e:
587+
# Log guardrail errors but don't fail function execution
588+
logger.error("Guardrail error: %s", e)
589+
output = await func(*func_args, **func_kwargs)
590+
else:
591+
output = await func(*func_args, **func_kwargs)
592+
546593
except Exception as exc:
547594
_log_step_exception(step, exc)
548595
exception = exc
549596
raise
550597

598+
# Apply output guardrails if provided
599+
if guardrails and output is not None:
600+
try:
601+
final_output, output_metadata = (
602+
_apply_guardrails_to_outputs(
603+
guardrails,
604+
output,
605+
_extract_function_inputs(
606+
func_signature=func_signature,
607+
func_args=func_args,
608+
func_kwargs=func_kwargs,
609+
context_kwarg=context_kwarg,
610+
),
611+
)
612+
)
613+
guardrail_metadata.update(output_metadata)
614+
615+
if final_output != output:
616+
output = final_output
617+
except Exception as e:
618+
# Log guardrail errors but don't fail function execution
619+
logger.error("Output guardrail error: %s", e)
620+
551621
# Extract inputs and finalize logging
552622
_process_wrapper_inputs_and_outputs(
553623
step=step,
@@ -556,7 +626,7 @@ async def async_function_wrapper(*func_args, **func_kwargs):
556626
func_kwargs=func_kwargs,
557627
context_kwarg=context_kwarg,
558628
output=output,
559-
guardrail_metadata={}, # TODO: Add guardrail support
629+
guardrail_metadata=guardrail_metadata,
560630
)
561631

562632
return output
@@ -572,12 +642,80 @@ def sync_wrapper(*func_args, **func_kwargs):
572642
**step_kwargs,
573643
) as step:
574644
output = exception = None
645+
guardrail_metadata = {}
575646
try:
576-
output = func(*func_args, **func_kwargs)
647+
# Apply input guardrails if provided
648+
if guardrails:
649+
try:
650+
inputs = _extract_function_inputs(
651+
func_signature=func_signature,
652+
func_args=func_args,
653+
func_kwargs=func_kwargs,
654+
context_kwarg=context_kwarg,
655+
)
656+
657+
# Process inputs through guardrails
658+
modified_inputs, input_metadata = (
659+
_apply_guardrails_to_inputs(
660+
guardrails,
661+
inputs,
662+
)
663+
)
664+
guardrail_metadata.update(input_metadata)
665+
666+
# Execute function with potentially modified inputs
667+
if modified_inputs != inputs:
668+
# Reconstruct function arguments from modified inputs
669+
bound = func_signature.bind(
670+
*func_args, **func_kwargs
671+
)
672+
bound.apply_defaults()
673+
674+
# Update bound arguments with modified values
675+
for (
676+
param_name,
677+
modified_value,
678+
) in modified_inputs.items():
679+
if param_name in bound.arguments:
680+
bound.arguments[param_name] = modified_value
681+
682+
output = func(*bound.args, **bound.kwargs)
683+
else:
684+
output = func(*func_args, **func_kwargs)
685+
except Exception as e:
686+
# Log guardrail errors but don't fail function execution
687+
logger.error("Guardrail error: %s", e)
688+
output = func(*func_args, **func_kwargs)
689+
else:
690+
output = func(*func_args, **func_kwargs)
691+
577692
except Exception as exc:
578693
_log_step_exception(step, exc)
579694
exception = exc
580695

696+
# Apply output guardrails if provided
697+
if guardrails and output is not None:
698+
try:
699+
final_output, output_metadata = (
700+
_apply_guardrails_to_outputs(
701+
guardrails,
702+
output,
703+
_extract_function_inputs(
704+
func_signature=func_signature,
705+
func_args=func_args,
706+
func_kwargs=func_kwargs,
707+
context_kwarg=context_kwarg,
708+
),
709+
)
710+
)
711+
guardrail_metadata.update(output_metadata)
712+
713+
if final_output != output:
714+
output = final_output
715+
except Exception as e:
716+
# Log guardrail errors but don't fail function execution
717+
logger.error("Output guardrail error: %s", e)
718+
581719
# Extract inputs and finalize logging
582720
_process_wrapper_inputs_and_outputs(
583721
step=step,
@@ -586,7 +724,7 @@ def sync_wrapper(*func_args, **func_kwargs):
586724
func_kwargs=func_kwargs,
587725
context_kwarg=context_kwarg,
588726
output=output,
589-
guardrail_metadata={}, # TODO: Add guardrail support
727+
guardrail_metadata=guardrail_metadata,
590728
)
591729

592730
if exception is not None:
@@ -987,11 +1125,7 @@ def _finalize_async_generator_step(
9871125
"""Finalize async generator step - called when generator is consumed."""
9881126
_current_step.reset(token)
9891127
_finalize_step_logging(
990-
step=step,
991-
inputs=inputs,
992-
output=output,
993-
start_time=step.start_time,
994-
guardrail_metadata={}, # TODO: Add guardrail support for async generators
1128+
step=step, inputs=inputs, output=output, start_time=step.start_time
9951129
)
9961130
_handle_trace_completion(
9971131
is_root_step=is_root_step,

0 commit comments

Comments
 (0)