diff --git a/guardrails/telemetry/guard_tracing.py b/guardrails/telemetry/guard_tracing.py index cd850acd5..b6dbb604b 100644 --- a/guardrails/telemetry/guard_tracing.py +++ b/guardrails/telemetry/guard_tracing.py @@ -179,10 +179,11 @@ def trace_guard_execution( if isinstance(result, Iterator) and not isinstance( result, ValidationOutcome ): - return trace_stream_guard(guard_span, result, history) - add_guard_attributes(guard_span, history, result) - add_user_attributes(guard_span) - return result + for res in trace_stream_guard(guard_span, result, history): + yield res + else: + add_guard_attributes(guard_span, history, result) + add_user_attributes(guard_span) except Exception as e: guard_span.set_status(status=StatusCode.ERROR, description=str(e)) raise e diff --git a/guardrails/telemetry/runner_tracing.py b/guardrails/telemetry/runner_tracing.py index d45c6ee4f..27cf27a97 100644 --- a/guardrails/telemetry/runner_tracing.py +++ b/guardrails/telemetry/runner_tracing.py @@ -265,6 +265,11 @@ def trace_call_wrapper(*args, **kwargs): ) as call_span: try: response = fn(*args, **kwargs) + if isinstance(response, LLMResponse) and ( + response.async_stream_output or response.stream_output + ): + # TODO: Iterate, add a call attr each time + return response add_call_attributes(call_span, response, *args, **kwargs) return response except Exception as e: