diff --git a/packages/opentelemetry-instrumentation-litellm/opentelemetry/instrumentation/litellm/streaming_safety.py b/packages/opentelemetry-instrumentation-litellm/opentelemetry/instrumentation/litellm/streaming_safety.py index fbe1d30390..fea4a688dc 100644 --- a/packages/opentelemetry-instrumentation-litellm/opentelemetry/instrumentation/litellm/streaming_safety.py +++ b/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 != "" diff --git a/packages/opentelemetry-instrumentation-litellm/tests/test_safety_hooks.py b/packages/opentelemetry-instrumentation-litellm/tests/test_safety_hooks.py index e5607fb2c4..cdbfaae1b3 100644 --- a/packages/opentelemetry-instrumentation-litellm/tests/test_safety_hooks.py +++ b/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")