feat(litellm): emit streaming latency span attrs

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

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import inspect
+import time
 from types import SimpleNamespace
 
 from opentelemetry import context as context_api
@@ -11,7 +12,9 @@
     CompletionTextStreamGroup,
 )
 from opentelemetry.instrumentation.litellm.safety import PROVIDER, extract_text_content
-from opentelemetry.semconv_ai import SpanAttributes
+
+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"
 
 
 def is_sync_streaming_response(kwargs, response) -> bool:
@@ -50,12 +53,27 @@ def wrap_sync_streaming_response(span, response, request_type, span_name, set_re
         request_type=request_type,
     )
     complete_response = {"choices": [], "model": None, "usage": None}
+    stream_started_at = time.time()
+    first_token_at = None
 
     try:
         for chunk in response:
             _mask_streaming_chunk(streams, chunk)
             _accumulate_streaming_chunk(complete_response, chunk)
+            if first_token_at is None and _chunk_has_output_text(chunk):
+                first_token_at = time.time()
+                _set_streaming_latency_attr(
+                    span,
+                    FR_STREAMING_TIME_TO_FIRST_TOKEN_MS,
+                    first_token_at - stream_started_at,
+                )
             yield chunk
+        if first_token_at is not None:
+            _set_streaming_latency_attr(
+                span,
+                FR_STREAMING_TIME_TO_GENERATE_MS,
+                time.time() - first_token_at,
+            )
         _finalize_streaming_span(span, complete_response, set_response_attributes)
     except Exception as exc:
         _close_streaming_response(response)
@@ -88,12 +106,27 @@ async def wrap_async_streaming_response(
         request_type=request_type,
     )
     complete_response = {"choices": [], "model": None, "usage": None}
+    stream_started_at = time.time()
+    first_token_at = None
 
     try:
         async for chunk in response:
             _mask_streaming_chunk(streams, chunk)
             _accumulate_streaming_chunk(complete_response, chunk)
+            if first_token_at is None and _chunk_has_output_text(chunk):
+                first_token_at = time.time()
+                _set_streaming_latency_attr(
+                    span,
+                    FR_STREAMING_TIME_TO_FIRST_TOKEN_MS,
+                    first_token_at - stream_started_at,
+                )
             yield chunk
+        if first_token_at is not None:
+            _set_streaming_latency_attr(
+                span,
+                FR_STREAMING_TIME_TO_GENERATE_MS,
+                time.time() - first_token_at,
+            )
         _finalize_streaming_span(span, complete_response, set_response_attributes)
     except Exception as exc:
         await _aclose_streaming_response(response)
@@ -213,3 +246,28 @@ async def _aclose_streaming_response(response) -> None:
     aclose = getattr(response, "aclose", None)
     if callable(aclose):
         await aclose()
+
+
+def _set_streaming_latency_attr(span, key: str, seconds: float) -> None:
+    if span.is_recording():
+        span.set_attribute(key, int(round(seconds * 1000)))
+
+
+def _chunk_has_output_text(chunk) -> bool:
+    for choice in get_object_value(chunk, "choices") or []:
+        if _value_has_text(get_object_value(choice, "text")):
+            return True
+        if _value_has_text(get_object_value(choice, "content")):
+            return True
+        message = get_object_value(choice, "message")
+        if message is not None and _value_has_text(get_object_value(message, "content")):
+            return True
+        delta = get_object_value(choice, "delta")
+        if delta is not None and _value_has_text(get_object_value(delta, "content")):
+            return True
+    return False
+
+
+def _value_has_text(value) -> bool:
+    text = extract_text_content(value)
+    return isinstance(text, str) and text != ""

packages/opentelemetry-instrumentation-litellm/tests/test_safety_hooks.py

@@ -32,7 +32,10 @@
     _resolve_masked_text,
 )
 from opentelemetry.instrumentation.litellm.streaming_safety import (
+    FR_STREAMING_TIME_TO_FIRST_TOKEN_MS,
+    FR_STREAMING_TIME_TO_GENERATE_MS,
     _accumulate_streaming_chunk,
+    _chunk_has_output_text,
     _mask_streaming_chunk,
     is_async_streaming_response,
     is_sync_streaming_response,
@@ -647,7 +650,10 @@ def test_safety_helpers_cover_args_blocks_and_text_extraction():
         None, args, {}, "chat", "litellm.completion"
     )
     assert unchanged_kwargs == {}
-    assert updated_args[1][0]["content"][:2] == ["[MASKED:prompt-a]", {"type": "input_text", "text": "[MASKED:prompt-b]"}]
+    assert updated_args[1][0]["content"][:2] == [
+        "[MASKED:prompt-a]",
+        {"type": "input_text", "text": "[MASKED:prompt-b]"},
+    ]
     assert _get_messages(updated_args, {})[1] == "args"
     assert apply_prompt_safety(None, (), {"messages": "invalid"}, "chat", "litellm.completion") == (
         (),
@@ -733,6 +739,122 @@ async def _agen():
     assert is_async_streaming_response({}, _agen()) is False
 
 
+def test_streaming_chunk_output_detection_handles_delta_message_and_text():
+    assert _chunk_has_output_text(
+        SimpleNamespace(
+            choices=[SimpleNamespace(delta=SimpleNamespace(content="hello"))]
+        )
+    )
+    assert _chunk_has_output_text(
+        SimpleNamespace(
+            choices=[
+                SimpleNamespace(message=SimpleNamespace(content="hello"))
+            ]
+        )
+    )
+    assert _chunk_has_output_text(
+        SimpleNamespace(choices=[SimpleNamespace(text="hello")])
+    )
+    assert not _chunk_has_output_text(
+        SimpleNamespace(
+            choices=[
+                SimpleNamespace(
+                    delta=SimpleNamespace(role="assistant", content="")
+                )
+            ]
+        )
+    )
+
+
+def test_sync_streaming_wrapper_sets_streaming_latency_attrs():
+    exporter, tracer = _test_tracer()
+    span = tracer.start_span("litellm.completion")
+    chunks = [
+        SimpleNamespace(
+            choices=[
+                SimpleNamespace(
+                    delta=SimpleNamespace(role="assistant", content="")
+                )
+            ]
+        ),
+        SimpleNamespace(
+            choices=[SimpleNamespace(delta=SimpleNamespace(content="hello"))]
+        ),
+    ]
+
+    with patch(
+        "opentelemetry.instrumentation.litellm.streaming_safety.time.time",
+        side_effect=[100.0, 100.25, 100.9],
+    ):
+        list(
+            wrap_sync_streaming_response(
+                span,
+                iter(chunks),
+                "chat",
+                "litellm.completion",
+                lambda *_: None,
+            )
+        )
+
+    attrs = exporter.get_finished_spans()[0].attributes
+    assert attrs[FR_STREAMING_TIME_TO_FIRST_TOKEN_MS] == 250
+    assert attrs[FR_STREAMING_TIME_TO_GENERATE_MS] == 650
+
+
+@pytest.mark.asyncio
+async def test_async_streaming_wrapper_sets_streaming_latency_attrs():
+    exporter, tracer = _test_tracer()
+    span = tracer.start_span("litellm.completion")
+
+    async def _response():
+        yield SimpleNamespace(
+            choices=[SimpleNamespace(delta=SimpleNamespace(content="hello"))]
+        )
+
+    with patch(
+        "opentelemetry.instrumentation.litellm.streaming_safety.time.time",
+        side_effect=[200.0, 200.125, 200.5],
+    ):
+        yielded = [
+            chunk
+            async for chunk in wrap_async_streaming_response(
+                span,
+                _response(),
+                "chat",
+                "litellm.completion",
+                lambda *_: None,
+            )
+        ]
+
+    assert len(yielded) == 1
+    attrs = exporter.get_finished_spans()[0].attributes
+    assert attrs[FR_STREAMING_TIME_TO_FIRST_TOKEN_MS] == 125
+    assert attrs[FR_STREAMING_TIME_TO_GENERATE_MS] == 375
+
+
+def test_streaming_wrapper_leaves_latency_attrs_unset_for_empty_stream():
+    exporter, tracer = _test_tracer()
+    span = tracer.start_span("litellm.completion")
+
+    with patch(
+        "opentelemetry.instrumentation.litellm.streaming_safety.time.time",
+        side_effect=[300.0],
+    ):
+        list(
+            wrap_sync_streaming_response(
+                span,
+                iter(()),
+                "chat",
+                "litellm.completion",
+                lambda *_: None,
+            )
+        )
+
+    attrs = exporter.get_finished_spans()[0].attributes
+    assert FR_STREAMING_TIME_TO_FIRST_TOKEN_MS not in attrs
+    assert FR_STREAMING_TIME_TO_GENERATE_MS not in attrs
+
+
 def test_set_request_attributes_does_not_treat_text_prompt_as_model():
     exporter, tracer = _test_tracer()
     span = tracer.start_span("litellm.text_completion")