fix stream generation

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

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: