feat(google-genai): emit streaming latency span attrs

packages/opentelemetry-instrumentation-google-generativeai/opentelemetry/instrumentation/google_generativeai/__init__.py

@@ -257,11 +257,11 @@ async def _awrap(
     if response:
         if is_streaming_response(response):
             return _build_from_streaming_response(
-                span, response, llm_model, event_logger, token_histogram
+                span, response, llm_model, event_logger, token_histogram, start_time
             )
         elif is_async_streaming_response(response):
             return _abuild_from_streaming_response(
-                span, response, llm_model, event_logger, token_histogram
+                span, response, llm_model, event_logger, token_histogram, start_time
             )
         else:
             await asyncio.to_thread(
@@ -343,11 +343,11 @@ def _wrap(
     if response:
         if is_streaming_response(response):
             return _build_from_streaming_response(
-                span, response, llm_model, event_logger, token_histogram
+                span, response, llm_model, event_logger, token_histogram, start_time
             )
         elif is_async_streaming_response(response):
             return _abuild_from_streaming_response(
-                span, response, llm_model, event_logger, token_histogram
+                span, response, llm_model, event_logger, token_histogram, start_time
             )
         else:
             _apply_completion_safety(span, response, name, response_model=llm_model)

packages/opentelemetry-instrumentation-google-generativeai/opentelemetry/instrumentation/google_generativeai/streaming_runtime.py

@@ -6,7 +6,9 @@
 )
 
 
-def build_streaming_response_delegate(span, response, llm_model, event_logger, token_histogram):
+def build_streaming_response_delegate(
+    span, response, llm_model, event_logger, token_histogram, start_time=None
+):
     from opentelemetry.instrumentation.google_generativeai import (
         emit_choice_events,
         set_model_response_attributes,
@@ -31,11 +33,12 @@ def build_streaming_response_delegate(span, response, llm_model, event_logger, t
         span=span,
         llm_model=llm_model,
         finalize_response=finalize_response,
+        start_time=start_time,
     )
 
 
 async def build_async_streaming_response_delegate(
-    span, response, llm_model, event_logger, token_histogram
+    span, response, llm_model, event_logger, token_histogram, start_time=None
 ):
     from opentelemetry.instrumentation.google_generativeai import (
         emit_choice_events,
@@ -61,5 +64,6 @@ async def build_async_streaming_response_delegate(
         span=span,
         llm_model=llm_model,
         finalize_response=finalize_response,
+        start_time=start_time,
     ):
         yield item

packages/opentelemetry-instrumentation-google-generativeai/opentelemetry/instrumentation/google_generativeai/streaming_safety.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import time
 from collections.abc import AsyncIterator, Iterator
 
 from opentelemetry.instrumentation.fortifyroot import get_object_value, set_object_value
@@ -9,6 +10,9 @@
 from opentelemetry.instrumentation.google_generativeai.safety import PROVIDER
 from opentelemetry.semconv_ai import LLMRequestTypeValues
 
+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 GoogleGenerativeAIStreamingSafety:
     def __init__(self, span, span_name: str):
@@ -65,20 +69,35 @@ def _sync_item_text(self, item, parts):
         set_object_value(item, "text", combined)
 
 
+def _set_streaming_latency_attribute(span, name, seconds):
+    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 build_streaming_response(
     response,
     *,
     span,
     llm_model,
     finalize_response,
+    start_time=None,
 ) -> Iterator:
     complete_response = ""
     last_chunk = None
     pending_item = None
+    first_token_time = None
     streaming_safety = GoogleGenerativeAIStreamingSafety(
         span, "gemini.generate_content"
     )
     for item in response:
+        if first_token_time is None:
+            first_token_time = time.perf_counter()
+            if start_time is not None:
+                _set_streaming_latency_attribute(
+                    span,
+                    FR_STREAMING_TIME_TO_FIRST_TOKEN_MS,
+                    first_token_time - start_time,
+                )
         item = streaming_safety.process_item(item)
         if pending_item is not None:
             yield pending_item
@@ -92,6 +111,12 @@ def build_streaming_response(
         complete_response += str(pending_item.text)
         last_chunk = pending_item
 
+    if first_token_time is not None:
+        _set_streaming_latency_attribute(
+            span,
+            FR_STREAMING_TIME_TO_GENERATE_MS,
+            time.perf_counter() - first_token_time,
+        )
     finalize_response(complete_response, last_chunk or response, llm_model)
 
 
@@ -101,14 +126,24 @@ async def build_async_streaming_response(
     span,
     llm_model,
     finalize_response,
+    start_time=None,
 ) -> AsyncIterator:
     complete_response = ""
     last_chunk = None
     pending_item = None
+    first_token_time = None
     streaming_safety = GoogleGenerativeAIStreamingSafety(
         span, "gemini.generate_content"
     )
     async for item in response:
+        if first_token_time is None:
+            first_token_time = time.perf_counter()
+            if start_time is not None:
+                _set_streaming_latency_attribute(
+                    span,
+                    FR_STREAMING_TIME_TO_FIRST_TOKEN_MS,
+                    first_token_time - start_time,
+                )
         item = streaming_safety.process_item(item)
         if pending_item is not None:
             yield pending_item
@@ -122,4 +157,10 @@ async def build_async_streaming_response(
         complete_response += str(pending_item.text)
         last_chunk = pending_item
 
+    if first_token_time is not None:
+        _set_streaming_latency_attribute(
+            span,
+            FR_STREAMING_TIME_TO_GENERATE_MS,
+            time.perf_counter() - first_token_time,
+        )
     finalize_response(complete_response, last_chunk or response, llm_model)

packages/opentelemetry-instrumentation-google-generativeai/tests/test_safety_unit.py

@@ -8,6 +8,10 @@
     build_async_streaming_response_delegate,
     build_streaming_response_delegate,
 )
+from opentelemetry.instrumentation.google_generativeai.streaming_safety import (
+    FR_STREAMING_TIME_TO_FIRST_TOKEN_MS,
+    FR_STREAMING_TIME_TO_GENERATE_MS,
+)
 from opentelemetry.sdk.trace import TracerProvider
 
 
@@ -168,6 +172,77 @@ def test_streaming_runtime_delegate_masks_chunk_and_finalizes(monkeypatch):
     clear_completion_safety_stream_factory()
 
 
+def test_streaming_runtime_delegate_sets_latency_span_attrs(monkeypatch):
+    from opentelemetry.instrumentation import google_generativeai as google_module
+    from opentelemetry.instrumentation.fortifyroot import (
+        clear_completion_safety_stream_factory,
+        register_completion_safety_stream_factory,
+    )
+    from opentelemetry.instrumentation.google_generativeai import streaming_safety
+
+    clear_completion_safety_stream_factory()
+    register_completion_safety_stream_factory(lambda _: _FakeStreamSession(["token"]))
+    ticks = iter([11.25, 13.0])
+    monkeypatch.setattr(streaming_safety.time, "perf_counter", lambda: next(ticks))
+    monkeypatch.setattr(google_module, "should_emit_events", lambda: False)
+    monkeypatch.setattr(google_module, "set_response_attributes", lambda *args: None)
+    monkeypatch.setattr(google_module, "set_model_response_attributes", lambda *args: None)
+
+    tracer = TracerProvider().get_tracer(__name__)
+    span = tracer.start_span("gemini.generate_content")
+    response = [
+        SimpleNamespace(
+            text="raw",
+            candidates=[
+                SimpleNamespace(
+                    content=SimpleNamespace(parts=[SimpleNamespace(text="secret")])
+                )
+            ],
+        )
+    ]
+
+    list(
+        build_streaming_response_delegate(
+            span,
+            response,
+            "gemini-1.5",
+            event_logger=None,
+            token_histogram=None,
+            start_time=10.0,
+        )
+    )
+
+    assert span.attributes[FR_STREAMING_TIME_TO_FIRST_TOKEN_MS] == 1250
+    assert span.attributes[FR_STREAMING_TIME_TO_GENERATE_MS] == 1750
+
+    clear_completion_safety_stream_factory()
+
+
+def test_streaming_runtime_delegate_leaves_empty_stream_latency_unset(monkeypatch):
+    from opentelemetry.instrumentation import google_generativeai as google_module
+
+    monkeypatch.setattr(google_module, "should_emit_events", lambda: False)
+    monkeypatch.setattr(google_module, "set_response_attributes", lambda *args: None)
+    monkeypatch.setattr(google_module, "set_model_response_attributes", lambda *args: None)
+
+    tracer = TracerProvider().get_tracer(__name__)
+    span = tracer.start_span("gemini.generate_content")
+
+    list(
+        build_streaming_response_delegate(
+            span,
+            [],
+            "gemini-1.5",
+            event_logger=None,
+            token_histogram=None,
+            start_time=10.0,
+        )
+    )
+
+    assert FR_STREAMING_TIME_TO_FIRST_TOKEN_MS not in span.attributes
+    assert FR_STREAMING_TIME_TO_GENERATE_MS not in span.attributes
+
+
 @pytest.mark.asyncio
 async def test_async_streaming_runtime_delegate_masks_chunk_and_finalizes(monkeypatch):
     from opentelemetry.instrumentation import google_generativeai as google_module
@@ -229,3 +304,51 @@ async def _response():
     assert not span.is_recording()
 
     clear_completion_safety_stream_factory()
+
+
+@pytest.mark.asyncio
+async def test_async_streaming_runtime_delegate_sets_latency_span_attrs(monkeypatch):
+    from opentelemetry.instrumentation import google_generativeai as google_module
+    from opentelemetry.instrumentation.fortifyroot import (
+        clear_completion_safety_stream_factory,
+        register_completion_safety_stream_factory,
+    )
+    from opentelemetry.instrumentation.google_generativeai import streaming_safety
+
+    clear_completion_safety_stream_factory()
+    register_completion_safety_stream_factory(lambda _: _FakeStreamSession(["token"]))
+    ticks = iter([21.1, 22.4])
+    monkeypatch.setattr(streaming_safety.time, "perf_counter", lambda: next(ticks))
+    monkeypatch.setattr(google_module, "should_emit_events", lambda: False)
+    monkeypatch.setattr(google_module, "set_response_attributes", lambda *args: None)
+    monkeypatch.setattr(google_module, "set_model_response_attributes", lambda *args: None)
+
+    async def _response():
+        yield SimpleNamespace(
+            text="raw",
+            candidates=[
+                SimpleNamespace(
+                    content=SimpleNamespace(parts=[SimpleNamespace(text="secret")])
+                )
+            ],
+        )
+
+    tracer = TracerProvider().get_tracer(__name__)
+    span = tracer.start_span("gemini.generate_content")
+    yielded = [
+        item
+        async for item in build_async_streaming_response_delegate(
+            span,
+            _response(),
+            "gemini-1.5",
+            event_logger=None,
+            token_histogram=None,
+            start_time=20.0,
+        )
+    ]
+
+    assert yielded[0].text == "token"
+    assert span.attributes[FR_STREAMING_TIME_TO_FIRST_TOKEN_MS] == 1100
+    assert span.attributes[FR_STREAMING_TIME_TO_GENERATE_MS] == 1300
+
+    clear_completion_safety_stream_factory()