feat(llamaindex): emit streaming latency span attrs

packages/opentelemetry-instrumentation-llamaindex/opentelemetry/instrumentation/llamaindex/streaming_safety.py

@@ -9,13 +9,16 @@
 """
 from __future__ import annotations
 
+import time
 from typing import Any, Generator
 
 from opentelemetry.instrumentation.fortifyroot.text_streaming import (
     CompletionTextStreamGroup,
 )
 
 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()

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()