diff --git a/packages/opentelemetry-instrumentation-google-generativeai/opentelemetry/instrumentation/google_generativeai/__init__.py b/packages/opentelemetry-instrumentation-google-generativeai/opentelemetry/instrumentation/google_generativeai/__init__.py index a1d7149710..015b451a7a 100644 --- a/packages/opentelemetry-instrumentation-google-generativeai/opentelemetry/instrumentation/google_generativeai/__init__.py +++ b/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) diff --git a/packages/opentelemetry-instrumentation-google-generativeai/opentelemetry/instrumentation/google_generativeai/streaming_runtime.py b/packages/opentelemetry-instrumentation-google-generativeai/opentelemetry/instrumentation/google_generativeai/streaming_runtime.py index 25ead716c1..817fe9447d 100644 --- a/packages/opentelemetry-instrumentation-google-generativeai/opentelemetry/instrumentation/google_generativeai/streaming_runtime.py +++ b/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 diff --git a/packages/opentelemetry-instrumentation-google-generativeai/opentelemetry/instrumentation/google_generativeai/streaming_safety.py b/packages/opentelemetry-instrumentation-google-generativeai/opentelemetry/instrumentation/google_generativeai/streaming_safety.py index 4068aa0ee5..f8fb35e2e2 100644 --- a/packages/opentelemetry-instrumentation-google-generativeai/opentelemetry/instrumentation/google_generativeai/streaming_safety.py +++ b/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) diff --git a/packages/opentelemetry-instrumentation-google-generativeai/tests/test_safety_unit.py b/packages/opentelemetry-instrumentation-google-generativeai/tests/test_safety_unit.py index 99a3c8b975..f54ae5ece8 100644 --- a/packages/opentelemetry-instrumentation-google-generativeai/tests/test_safety_unit.py +++ b/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()