Capture partial LLM output on interruption

traced_llm only attached the aggregated ``output`` attribute to the
span after the wrapped function returned successfully. When the LLM
call was cancelled mid-stream (e.g. interruption during generation),
the accumulated text was discarded — the span had no ``output``.

Moved the attribute assignment into the ``finally`` block alongside
the existing TTFB write so the partial text we already captured via
the patched ``push_frame`` lands on the span regardless of whether
``f`` returned normally, raised, or was cancelled.
This commit is contained in:
Aleix Conchillo Flaqué
2026-05-12 12:04:08 -07:00
parent e70ee603b2
commit a3ce963b54

View File

@@ -889,10 +889,6 @@ def traced_llm(func: Callable | None = None, *, name: str | None = None) -> Call
fn_called = True
result = await f(self, context, *args, **kwargs)
# Add aggregated output after function completes, if available
if output_text:
current_span.set_attribute("output", output_text)
return result
finally:
@@ -905,6 +901,15 @@ def traced_llm(func: Callable | None = None, *, name: str | None = None) -> Call
):
self.start_llm_usage_metrics = original_start_llm_usage_metrics
# Attach whatever output text we accumulated so
# far. Doing this in finally captures partial
# output when ``f`` is cancelled or raises mid-
# stream (e.g. interruption during LLM
# generation), rather than only on clean
# completion.
if output_text:
current_span.set_attribute("output", output_text)
# Update TTFB metric
ttfb: float | None = getattr(getattr(self, "_metrics", None), "ttfb", None)
if ttfb is not None: