From 10170bba00b2a7cbd82dcd857e3de4cc37bcb119 Mon Sep 17 00:00:00 2001 From: sayonfortify Date: Fri, 19 Jun 2026 01:19:17 +0530 Subject: [PATCH] feat(openai): emit streaming latency span attrs --- .../openai/shared/chat_wrappers.py | 86 +++++++++-- .../test_streaming_latency_span_attrs.py | 145 ++++++++++++++++++ 2 files changed, 220 insertions(+), 11 deletions(-) create mode 100644 packages/opentelemetry-instrumentation-openai/tests/test_streaming_latency_span_attrs.py diff --git a/packages/opentelemetry-instrumentation-openai/opentelemetry/instrumentation/openai/shared/chat_wrappers.py b/packages/opentelemetry-instrumentation-openai/opentelemetry/instrumentation/openai/shared/chat_wrappers.py index 6cd8cc88d6..23db70fd97 100644 --- a/packages/opentelemetry-instrumentation-openai/opentelemetry/instrumentation/openai/shared/chat_wrappers.py +++ b/packages/opentelemetry-instrumentation-openai/opentelemetry/instrumentation/openai/shared/chat_wrappers.py @@ -76,6 +76,13 @@ LLM_REQUEST_TYPE = LLMRequestTypeValues.CHAT +# FR: stable internal span attributes carrying streaming latency as integer +# milliseconds, so the backend can extract TTFT/STTG into RDS llm_usage_events. +# Additive to (not a replacement for) the Mimir streaming histograms. Contract: +# fr-backend/docs/development/STREAMING_LATENCY_TTFT_STTG_PLAN.md +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" + logger = logging.getLogger(__name__) @@ -773,13 +780,22 @@ def _process_item(self, item): self._span.add_event( name=f"{SpanAttributes.LLM_CONTENT_COMPLETION_CHUNK}") - if self._first_token and self._streaming_time_to_first_token: + if self._first_token: self._time_of_first_token = time.time() - self._streaming_time_to_first_token.record( - self._time_of_first_token - self._start_time, - attributes=self._shared_attributes(), - ) + if self._streaming_time_to_first_token: + self._streaming_time_to_first_token.record( + self._time_of_first_token - self._start_time, + attributes=self._shared_attributes(), + ) self._first_token = False + # FR: also persist TTFT (ms, int) on the span so the backend can + # extract it into RDS (the Mimir histogram above is unchanged). + if self._span and self._span.is_recording(): + _set_span_attribute( + self._span, + FR_STREAMING_TIME_TO_FIRST_TOKEN_MS, + round((self._time_of_first_token - self._start_time) * 1000), + ) item = self._streaming_safety.process_chunk(item) _accumulate_stream_items(item, self._complete_response) @@ -847,6 +863,20 @@ def _process_complete_response(self): _set_completions( self._span, self._complete_response.get("choices")) + # FR: persist STTG (ms, int) on the span for RDS extraction, only + # when a first token actually arrived (mirrors the TTFT capture). + if ( + self._span + and self._span.is_recording() + and not self._first_token + and self._time_of_first_token + ): + _set_span_attribute( + self._span, + FR_STREAMING_TIME_TO_GENERATE_MS, + round((time.time() - self._time_of_first_token) * 1000), + ) + self._span.set_status(Status(StatusCode.OK)) self._span.end() self._cleanup_completed = True @@ -950,11 +980,19 @@ def _build_from_streaming_response( item_to_yield = streaming_safety.process_chunk(item) - if first_token and streaming_time_to_first_token: + if first_token: time_of_first_token = time.time() - streaming_time_to_first_token.record( - time_of_first_token - start_time) + if streaming_time_to_first_token: + streaming_time_to_first_token.record( + time_of_first_token - start_time) first_token = False + # FR: also persist TTFT (ms, int) on the span for RDS extraction. + if span and span.is_recording(): + _set_span_attribute( + span, + FR_STREAMING_TIME_TO_FIRST_TOKEN_MS, + round((time_of_first_token - start_time) * 1000), + ) _accumulate_stream_items(item, complete_response) @@ -990,6 +1028,15 @@ def _build_from_streaming_response( if streaming_time_to_generate and time_of_first_token: streaming_time_to_generate.record(time.time() - time_of_first_token) + # FR: persist STTG (ms, int) on the span for RDS extraction, only when a + # first token actually arrived (not first_token). + if span and span.is_recording() and not first_token and time_of_first_token: + _set_span_attribute( + span, + FR_STREAMING_TIME_TO_GENERATE_MS, + round((time.time() - time_of_first_token) * 1000), + ) + _set_response_attributes(span, complete_response) if should_emit_events(): for choice in complete_response.get("choices", []): @@ -1027,11 +1074,19 @@ async def _abuild_from_streaming_response( item_to_yield = streaming_safety.process_chunk(item) - if first_token and streaming_time_to_first_token: + if first_token: time_of_first_token = time.time() - streaming_time_to_first_token.record( - time_of_first_token - start_time) + if streaming_time_to_first_token: + streaming_time_to_first_token.record( + time_of_first_token - start_time) first_token = False + # FR: also persist TTFT (ms, int) on the span for RDS extraction. + if span and span.is_recording(): + _set_span_attribute( + span, + FR_STREAMING_TIME_TO_FIRST_TOKEN_MS, + round((time_of_first_token - start_time) * 1000), + ) _accumulate_stream_items(item, complete_response) @@ -1067,6 +1122,15 @@ async def _abuild_from_streaming_response( if streaming_time_to_generate and time_of_first_token: streaming_time_to_generate.record(time.time() - time_of_first_token) + # FR: persist STTG (ms, int) on the span for RDS extraction, only when a + # first token actually arrived (not first_token). + if span and span.is_recording() and not first_token and time_of_first_token: + _set_span_attribute( + span, + FR_STREAMING_TIME_TO_GENERATE_MS, + round((time.time() - time_of_first_token) * 1000), + ) + _set_response_attributes(span, complete_response) if should_emit_events(): for choice in complete_response.get("choices", []): diff --git a/packages/opentelemetry-instrumentation-openai/tests/test_streaming_latency_span_attrs.py b/packages/opentelemetry-instrumentation-openai/tests/test_streaming_latency_span_attrs.py new file mode 100644 index 0000000000..fcb099c892 --- /dev/null +++ b/packages/opentelemetry-instrumentation-openai/tests/test_streaming_latency_span_attrs.py @@ -0,0 +1,145 @@ +# NOTE: +# This file has been added by FortifyRoot. +# +# Unit test for the FortifyRoot streaming-latency span attributes added to the +# OpenAI streaming instrumentation: +# - fortifyroot.llm.streaming.time_to_first_token_ms (TTFT) +# - fortifyroot.llm.streaming.time_to_generate_ms (STTG) +# +# These are the RDS-extractable span counterparts of the existing Mimir +# streaming histograms (which are left unchanged). It exercises the v0 streaming +# generator directly with fake chunks + stub histograms, so it needs NO OpenAI +# API key and NO VCR cassette and is fully deterministic. +# +# Contract: fr-backend/docs/development/STREAMING_LATENCY_TTFT_STTG_PLAN.md +import time + +from opentelemetry.sdk.trace import TracerProvider +from opentelemetry.sdk.trace.export import SimpleSpanProcessor +from opentelemetry.sdk.trace.export.in_memory_span_exporter import ( + InMemorySpanExporter, +) + +from opentelemetry.instrumentation.openai.shared import chat_wrappers as cw + +TTFT = "fortifyroot.llm.streaming.time_to_first_token_ms" +STTG = "fortifyroot.llm.streaming.time_to_generate_ms" + + +class _StubHistogram: + """Truthy histogram stub so the wrapper's record()/attr branches run.""" + + def __init__(self): + self.records = [] + + def record(self, value, attributes=None): + self.records.append(value) + + +class _StubStreamingSafety: + def __init__(self, span, span_name): + pass + + def process_chunk(self, item): + return item + + +def _neutralize_helpers(monkeypatch): + """Isolate the streaming loop + new span-attr logic from the OpenAI-typed + accumulation/response helpers so the generator runs on fake chunks.""" + monkeypatch.setattr(cw, "_accumulate_stream_items", lambda *a, **k: None) + monkeypatch.setattr(cw, "_set_streaming_token_metrics", lambda *a, **k: None) + monkeypatch.setattr(cw, "_set_response_attributes", lambda *a, **k: None) + monkeypatch.setattr(cw, "_set_completions", lambda *a, **k: None) + monkeypatch.setattr(cw, "_get_openai_base_url", lambda *a, **k: "") + monkeypatch.setattr(cw, "metric_shared_attributes", lambda *a, **k: {}) + monkeypatch.setattr(cw, "should_emit_events", lambda: False) + monkeypatch.setattr(cw, "should_send_prompts", lambda: False) + monkeypatch.setattr(cw, "OpenAIChatStreamingSafety", _StubStreamingSafety) + + +def _recording_span(): + exporter = InMemorySpanExporter() + provider = TracerProvider() + provider.add_span_processor(SimpleSpanProcessor(exporter)) + span = provider.get_tracer("fr-streaming-latency-test").start_span("openai.chat") + return span, exporter + + +def test_streaming_sets_positive_integer_latency_attrs(monkeypatch): + """A streaming response that yields tokens sets both attributes as + non-negative integer milliseconds on the single LLM span.""" + _neutralize_helpers(monkeypatch) + span, exporter = _recording_span() + + start = time.time() - 0.05 # 50ms ago => TTFT/STTG are clearly > 0 + ttft_hist, sttg_hist = _StubHistogram(), _StubHistogram() + chunks = [object(), object(), object()] + + consumed = list( + cw._build_from_streaming_response( + span, + iter(chunks), + streaming_time_to_first_token=ttft_hist, + streaming_time_to_generate=sttg_hist, + start_time=start, + ) + ) + + assert len(consumed) == len(chunks) + finished = exporter.get_finished_spans() + assert len(finished) == 1 + attrs = finished[0].attributes + + assert isinstance(attrs.get(TTFT), int) and attrs[TTFT] >= 0 + assert isinstance(attrs.get(STTG), int) and attrs[STTG] >= 0 + # The existing Mimir histograms still record exactly once each (unchanged). + assert len(ttft_hist.records) == 1 + assert len(sttg_hist.records) == 1 + + +def test_streaming_latency_attrs_do_not_require_metrics_histograms(monkeypatch): + """Tracing/RDS span attrs must not depend on the Mimir histogram objects.""" + _neutralize_helpers(monkeypatch) + span, exporter = _recording_span() + + consumed = list( + cw._build_from_streaming_response( + span, + iter([object(), object()]), + streaming_time_to_first_token=None, + streaming_time_to_generate=None, + start_time=time.time() - 0.05, + ) + ) + + assert len(consumed) == 2 + finished = exporter.get_finished_spans() + assert len(finished) == 1 + attrs = finished[0].attributes + assert isinstance(attrs.get(TTFT), int) and attrs[TTFT] >= 0 + assert isinstance(attrs.get(STTG), int) and attrs[STTG] >= 0 + + +def test_empty_stream_sets_no_latency_attrs(monkeypatch): + """No chunk => no first token => neither attribute is set, so the backend + leaves the RDS columns NULL (they are nullable by design).""" + _neutralize_helpers(monkeypatch) + span, exporter = _recording_span() + + consumed = list( + cw._build_from_streaming_response( + span, + iter([]), + streaming_time_to_first_token=_StubHistogram(), + streaming_time_to_generate=_StubHistogram(), + start_time=time.time(), + ) + ) + + assert consumed == [] + finished = exporter.get_finished_spans() + assert len(finished) == 1 + attrs = finished[0].attributes + assert TTFT not in attrs + assert STTG not in attrs