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:
@@ -889,10 +889,6 @@ def traced_llm(func: Callable | None = None, *, name: str | None = None) -> Call
|
|||||||
fn_called = True
|
fn_called = True
|
||||||
result = await f(self, context, *args, **kwargs)
|
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
|
return result
|
||||||
|
|
||||||
finally:
|
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
|
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
|
# Update TTFB metric
|
||||||
ttfb: float | None = getattr(getattr(self, "_metrics", None), "ttfb", None)
|
ttfb: float | None = getattr(getattr(self, "_metrics", None), "ttfb", None)
|
||||||
if ttfb is not None:
|
if ttfb is not None:
|
||||||
|
|||||||
Reference in New Issue
Block a user