From 3d1b6b996b0c06ec593e281c8d0b6892ce7ec298 Mon Sep 17 00:00:00 2001 From: sayonfortify Date: Fri, 19 Jun 2026 23:19:06 +0530 Subject: [PATCH] feat(llamaindex): emit streaming latency span attrs --- .../llamaindex/streaming_safety.py | 43 +++++++++++ .../tests/test_streaming_safety.py | 73 +++++++++++++++++++ 2 files changed, 116 insertions(+) diff --git a/packages/opentelemetry-instrumentation-llamaindex/opentelemetry/instrumentation/llamaindex/streaming_safety.py b/packages/opentelemetry-instrumentation-llamaindex/opentelemetry/instrumentation/llamaindex/streaming_safety.py index 90fe09fd38..7e345e8cb0 100644 --- a/packages/opentelemetry-instrumentation-llamaindex/opentelemetry/instrumentation/llamaindex/streaming_safety.py +++ b/packages/opentelemetry-instrumentation-llamaindex/opentelemetry/instrumentation/llamaindex/streaming_safety.py @@ -9,6 +9,7 @@ """ from __future__ import annotations +import time from typing import Any, Generator from opentelemetry.instrumentation.fortifyroot.text_streaming import ( @@ -16,6 +17,8 @@ ) PROVIDER = "LlamaIndex" +FR_STREAMING_TIME_TO_FIRST_TOKEN_MS = "fortifyroot.llm.streaming.time_to_first_token_ms" +FR_STREAMING_TIME_TO_GENERATE_MS = "fortifyroot.llm.streaming.time_to_generate_ms" class LlamaIndexStreamingSafety: @@ -26,6 +29,9 @@ class LlamaIndexStreamingSafety: """ def __init__(self, span: Any, span_name: str, request_type: str) -> None: + self._span = span + self._stream_started_at = time.perf_counter() + self._first_token_at: float | None = None self._streams = CompletionTextStreamGroup( span=span, provider=PROVIDER, @@ -62,6 +68,27 @@ def flush(self, *, segment_index: int = 0) -> str: """ return self._streams.flush(key=segment_index) or "" + def record_first_token(self) -> None: + """Record TTFT once, when the first output-bearing chunk arrives.""" + if self._first_token_at is not None: + return + self._first_token_at = time.perf_counter() + _set_streaming_latency_attribute( + self._span, + FR_STREAMING_TIME_TO_FIRST_TOKEN_MS, + self._first_token_at - self._stream_started_at, + ) + + def record_completion(self) -> None: + """Record STTG if at least one output-bearing chunk arrived.""" + if self._first_token_at is None: + return + _set_streaming_latency_attribute( + self._span, + FR_STREAMING_TIME_TO_GENERATE_MS, + time.perf_counter() - self._first_token_at, + ) + # --------------------------------------------------------------------------- # Internal helpers (best-effort, never raise) @@ -91,6 +118,16 @@ def _flush_into_last(response: Any, safety: LlamaIndexStreamingSafety) -> None: pass +def _set_streaming_latency_attribute(span: Any, name: str, seconds: float) -> None: + if span is not None and span.is_recording() and seconds is not None: + span.set_attribute(name, int(round(max(0, seconds) * 1000))) + + +def _response_has_output_delta(response: Any) -> bool: + delta = getattr(response, "delta", None) + return isinstance(delta, str) and delta != "" + + # --------------------------------------------------------------------------- # Sync streaming wrapper # --------------------------------------------------------------------------- @@ -106,10 +143,13 @@ def wrap_stream(gen: Generator, safety: LlamaIndexStreamingSafety) -> Generator: for response in gen: if pending is not None: yield pending + if _response_has_output_delta(response): + safety.record_first_token() _patch_delta(response, safety) pending = response if pending is not None: _flush_into_last(pending, safety) + safety.record_completion() yield pending @@ -132,10 +172,13 @@ async def _gen(): async for response in agen: if pending is not None: yield pending + if _response_has_output_delta(response): + safety.record_first_token() _patch_delta(response, safety) pending = response if pending is not None: _flush_into_last(pending, safety) + safety.record_completion() yield pending return _gen() diff --git a/packages/opentelemetry-instrumentation-llamaindex/tests/test_streaming_safety.py b/packages/opentelemetry-instrumentation-llamaindex/tests/test_streaming_safety.py index 79a435c22c..452d51b03c 100644 --- a/packages/opentelemetry-instrumentation-llamaindex/tests/test_streaming_safety.py +++ b/packages/opentelemetry-instrumentation-llamaindex/tests/test_streaming_safety.py @@ -41,6 +41,8 @@ uninstrument_llm_safety_wrappers, ) from opentelemetry.instrumentation.llamaindex.streaming_safety import ( + FR_STREAMING_TIME_TO_FIRST_TOKEN_MS, + FR_STREAMING_TIME_TO_GENERATE_MS, LlamaIndexStreamingSafety, make_async_stream, wrap_stream, @@ -203,6 +205,27 @@ def test_pending_item_pattern_yields_every_item(self): result = list(wrap_stream(iter(responses), safety)) assert len(result) == 5 + def test_streaming_latency_hooks_fire_for_output_deltas(self): + """TTFT is first output chunk; STTG is stream exhaustion after first output.""" + responses = [_chat_response("a"), _chat_response("b")] + safety = _mock_safety(flush_return="") + + result = list(wrap_stream(iter(responses), safety)) + + assert [r.delta for r in result] == ["a", "b"] + assert safety.record_first_token.call_count == 2 + safety.record_completion.assert_called_once() + + def test_streaming_latency_hooks_skip_empty_delta_stream(self): + responses = [_chat_response(None), _chat_response("")] + safety = _mock_safety(flush_return="") + + result = list(wrap_stream(iter(responses), safety)) + + assert len(result) == 2 + safety.record_first_token.assert_not_called() + safety.record_completion.assert_called_once() + # =========================================================================== # streaming_safety.py: make_async_stream @@ -265,6 +288,20 @@ async def _source(): assert result == [] safety.flush.assert_not_called() + @pytest.mark.asyncio + async def test_streaming_latency_hooks_fire_for_output_deltas(self): + async def _source(): + yield _chat_response("a") + yield _chat_response("b") + + safety = _mock_safety(flush_return="") + + result = [r async for r in make_async_stream(_source(), safety)] + + assert [r.delta for r in result] == ["a", "b"] + assert safety.record_first_token.call_count == 2 + safety.record_completion.assert_called_once() + # =========================================================================== # safety.py: llm_stream_chat_wrapper @@ -664,3 +701,39 @@ def test_flush_returns_empty_string_when_none(self): safety = LlamaIndexStreamingSafety(span, "test.stream", "CHAT") assert safety.flush() == "" + + def test_streaming_latency_span_attrs_are_recorded_as_ms(self, monkeypatch): + span = MagicMock() + span.is_recording.return_value = True + ticks = iter([10.0, 10.125, 11.25]) + monkeypatch.setattr( + "opentelemetry.instrumentation.llamaindex.streaming_safety.time.perf_counter", + lambda: next(ticks), + ) + with patch( + "opentelemetry.instrumentation.llamaindex.streaming_safety.CompletionTextStreamGroup", + autospec=True, + ): + safety = LlamaIndexStreamingSafety(span, "test.stream", "COMPLETION") + safety.record_first_token() + safety.record_first_token() + safety.record_completion() + + span.set_attribute.assert_any_call(FR_STREAMING_TIME_TO_FIRST_TOKEN_MS, 125) + span.set_attribute.assert_any_call(FR_STREAMING_TIME_TO_GENERATE_MS, 1125) + + def test_streaming_latency_span_attrs_skip_empty_stream(self, monkeypatch): + span = MagicMock() + span.is_recording.return_value = True + monkeypatch.setattr( + "opentelemetry.instrumentation.llamaindex.streaming_safety.time.perf_counter", + lambda: 10.0, + ) + with patch( + "opentelemetry.instrumentation.llamaindex.streaming_safety.CompletionTextStreamGroup", + autospec=True, + ): + safety = LlamaIndexStreamingSafety(span, "test.stream", "COMPLETION") + safety.record_completion() + + span.set_attribute.assert_not_called()