diff --git a/src/openlayer/lib/__init__.py b/src/openlayer/lib/__init__.py index 35ea566a..55ede7e3 100644 --- a/src/openlayer/lib/__init__.py +++ b/src/openlayer/lib/__init__.py @@ -25,6 +25,7 @@ "trace_claude_agent_sdk", "traced_claude_agent_sdk_query", "trace_gemini", + "trace_google_genai", "update_current_trace", "update_current_step", # Offline buffer management functions @@ -162,9 +163,7 @@ def trace_azure_content_understanding(client): from .integrations import azure_content_understanding_tracer if not isinstance(client, ContentUnderstandingClient): - raise ValueError( - "Invalid client. Please provide a ContentUnderstandingClient." - ) + raise ValueError("Invalid client. Please provide a ContentUnderstandingClient.") return azure_content_understanding_tracer.trace_azure_content_understanding(client) @@ -236,14 +235,14 @@ def trace_portkey(): >>> from openlayer.lib import trace_portkey >>> # Enable openlayer tracing for all Portkey completions >>> trace_portkey() - >>> # Basic portkey client initialization + >>> # Basic portkey client initialization >>> portkey = Portkey( >>> api_key = os.environ['PORTKEY_API_KEY'], >>> config = "YOUR_PORTKEY_CONFIG_ID", # optional your portkey config id >>> ) >>> # use portkey normally - tracing happens automatically >>> response = portkey.chat.completions.create( - >>> #model = "@YOUR_PORTKEY_SLUG/YOUR_MODEL_NAME", # optional if giving config + >>> # model = "@YOUR_PORTKEY_SLUG/YOUR_MODEL_NAME", # optional if giving config >>> messages = [ >>> {"role": "system", "content": "You are a helpful assistant."}, >>> {"role": "user", "content": "Write a poem on Argentina, least 100 words."} @@ -395,18 +394,98 @@ def traced_claude_agent_sdk_query(*, prompt, options=None, inference_pipeline_id # -------------------------------- Google Gemini --------------------------------- # -def trace_gemini(client): - """Trace Google Gemini chat completions.""" +def _legacy_gemini_model_class(): + """``google.generativeai.GenerativeModel``, or None if not installed.""" + # pylint: disable=import-outside-toplevel + try: + import google.generativeai as legacy_genai + + return legacy_genai.GenerativeModel + except ImportError: + return None + + +def _google_genai_client_class(): + """``google.genai.Client``, or None if not installed.""" # pylint: disable=import-outside-toplevel try: - import google.generativeai as genai + from google import genai as google_genai + + return google_genai.Client except ImportError: + return None + + +def trace_gemini(client): + """Trace Google Gemini chat completions for either Google SDK. + + Accepts both Gemini clients and dispatches to the matching tracer: + + - ``google.genai.Client`` (package ``google-genai``, the current unified SDK, + including ``vertexai=True`` mode) -> ``trace_google_genai`` + - ``google.generativeai.GenerativeModel`` (package ``google-generativeai``, + legacy / maintenance mode) -> the legacy Gemini tracer + + Only the SDK matching the client passed in is imported, so having just one of + the two installed is fine. + + Example: + >>> from google import genai + >>> from openlayer.lib import trace_gemini + >>> client = trace_gemini(genai.Client(vertexai=True, project="p", location="us-central1")) + >>> client.models.generate_content(model="gemini-2.5-flash", contents="hi") + """ + # pylint: disable=import-outside-toplevel + legacy_class = _legacy_gemini_model_class() + unified_class = _google_genai_client_class() + + if legacy_class is not None and isinstance(client, legacy_class): + from .integrations import gemini_tracer + + return gemini_tracer.trace_gemini(client) + + if unified_class is not None and isinstance(client, unified_class): + from .integrations import google_genai_tracer + + return google_genai_tracer.trace_google_genai(client) + + if legacy_class is None and unified_class is None: raise ImportError( - "google-generativeai is required for Gemini tracing. Install with: pip install google-generativeai" + "Gemini tracing requires either google-genai (recommended) or " + "google-generativeai. Install with: pip install google-genai" + ) + + raise ValueError( + "Invalid client. Please provide a google.genai.Client (package " + "google-genai) or a google.generativeai.GenerativeModel (package " + f"google-generativeai). Got: {type(client).__module__}.{type(client).__qualname__}" + ) + + +def trace_google_genai(client): + """Trace Google Gen AI (``google-genai``) content generation. + + Covers ``client.models.generate_content`` and ``generate_content_stream``, + their async equivalents under ``client.aio.models``, and chat sessions via + ``client.chats`` — in both Vertex AI and AI Studio mode. + + Example: + >>> from google import genai + >>> from openlayer.lib import trace_google_genai + >>> client = trace_google_genai(genai.Client(vertexai=True, project="p", location="us-central1")) + >>> client.models.generate_content(model="gemini-2.5-flash", contents="hi") + """ + # pylint: disable=import-outside-toplevel + unified_class = _google_genai_client_class() + if unified_class is None: + raise ImportError("google-genai is required for Google Gen AI tracing. Install with: pip install google-genai") + + if not isinstance(client, unified_class): + raise ValueError( + "Invalid client. Please provide a google.genai.Client instance. Got: " + f"{type(client).__module__}.{type(client).__qualname__}" ) - from .integrations import gemini_tracer + from .integrations import google_genai_tracer - if not isinstance(client, genai.GenerativeModel): - raise ValueError("Invalid client. Please provide a google.generativeai.GenerativeModel instance.") - return gemini_tracer.trace_gemini(client) + return google_genai_tracer.trace_google_genai(client) diff --git a/src/openlayer/lib/integrations/__init__.py b/src/openlayer/lib/integrations/__init__.py index b91776ef..e6931fd8 100644 --- a/src/openlayer/lib/integrations/__init__.py +++ b/src/openlayer/lib/integrations/__init__.py @@ -38,3 +38,10 @@ __all__.extend(["trace_gemini"]) except ImportError: pass + +try: + from .google_genai_tracer import trace_google_genai + + __all__.extend(["trace_google_genai"]) +except ImportError: + pass diff --git a/src/openlayer/lib/integrations/_auto.py b/src/openlayer/lib/integrations/_auto.py index 418c8a24..81ef2f02 100644 --- a/src/openlayer/lib/integrations/_auto.py +++ b/src/openlayer/lib/integrations/_auto.py @@ -138,18 +138,25 @@ def _do_patch() -> None: _patch_via("groq_tracer", "_patch_groq"), _patch_via("groq_tracer", "_unpatch_groq"), ), - # TODO: This targets the LEGACY Google Generative AI SDK (package - # `google-generativeai`, module `google.generativeai`), which is in - # maintenance mode. The new Google Gen AI SDK (package `google-genai`, - # module `google.genai`, client `genai.Client()`) is NOT yet covered — - # users on it get no auto-instrumentation. Follow-up: add a - # google_genai_tracer and register a second entry probing `google.genai`. + # Two Google Gemini entries, because the two SDKs are separate packages that + # can be installed side by side: + # - `gemini` -> LEGACY google-generativeai / genai.GenerativeModel + # (maintenance mode) + # - `google_genai` -> CURRENT google-genai / genai.Client(), incl. Vertex + # `_is_installed` probes the full dotted path, which matters here: `google` is + # a namespace package, so one can be importable without the other. IntegrationSpec( "gemini", "google.generativeai", _patch_via("gemini_tracer", "_patch_gemini"), _patch_via("gemini_tracer", "_unpatch_gemini"), ), + IntegrationSpec( + "google_genai", + "google.genai", + _patch_via("google_genai_tracer", "_patch_google_genai"), + _patch_via("google_genai_tracer", "_unpatch_google_genai"), + ), IntegrationSpec( "oci", "oci", diff --git a/src/openlayer/lib/integrations/gemini_tracer.py b/src/openlayer/lib/integrations/gemini_tracer.py index ceaa81f5..d8e8d459 100644 --- a/src/openlayer/lib/integrations/gemini_tracer.py +++ b/src/openlayer/lib/integrations/gemini_tracer.py @@ -2,11 +2,10 @@ NOTE: This targets the LEGACY Google Generative AI SDK only — package ``google-generativeai``, module ``google.generativeai``, client -``genai.GenerativeModel``. That SDK is in maintenance mode. The new Google +``genai.GenerativeModel``. That SDK is in maintenance mode. The current Google Gen AI SDK (package ``google-genai``, module ``google.genai``, client -``genai.Client()`` with ``client.models.generate_content``) is NOT handled -here and is not auto-instrumented; supporting it needs a separate tracer + -registry entry (see the TODO in ``_auto.py``). +``genai.Client()`` with ``client.models.generate_content``) is handled by +``google_genai_tracer.py``, which is registered separately in ``_auto.py``. """ import json @@ -29,23 +28,66 @@ logger = logging.getLogger(__name__) +# Backend cost lookup is an exact ``(provider.lower(), model)`` match with no +# aliasing, so ``provider`` must BE a real cost-table slug. The former "Google" +# label is OpenRouter-sourced: it misses newer models entirely (e.g. +# gemini-3-pro-preview prices at $0) and carries stale prices for the ``-latest`` +# aliases. Kept in sync with ``google_genai_tracer.PROVIDER``. +PROVIDER = "gemini" + + +def _extract_token_counts(usage: Any) -> tuple: + """Return ``(prompt_tokens, completion_tokens, total_tokens)`` from usage. + + ``candidates_token_count`` EXCLUDES ``thoughts_token_count``, and thinking is + on by default for 2.5-series models, so ``prompt + candidates`` badly + undercounts — measured 14 tokens for a call that consumed 757 (prompt 8, + candidates 6, thoughts 743). + + ``total_token_count`` is authoritative when present. Thinking tokens are + folded into ``completion_tokens`` because they are billed as output and the + backend prices cost from ``completionTokens``. + """ + if usage is None: + return 0, 0, 0 + + def _count(name: str) -> Optional[int]: + value = getattr(usage, name, None) + return value if isinstance(value, int) else None + + prompt_tokens = _count("prompt_token_count") or 0 + candidates_tokens = _count("candidates_token_count") or 0 + thoughts_tokens = _count("thoughts_token_count") or 0 + tool_use_tokens = _count("tool_use_prompt_token_count") or 0 + total_tokens = _count("total_token_count") + + completion_tokens = candidates_tokens + thoughts_tokens + if total_tokens is None: + total_tokens = prompt_tokens + completion_tokens + tool_use_tokens + + return prompt_tokens, completion_tokens, total_tokens + def _clean_model_name(model_name: str) -> str: - """Remove 'models/' prefix from Gemini model names. - + """Reduce a Gemini model name to the bare slug the cost table is keyed on. + Parameters ---------- model_name : str The raw model name from the client (e.g., "models/gemini-pro"). - + Returns ------- str The cleaned model name (e.g., "gemini-pro"). """ - if model_name and model_name.startswith("models/"): - return model_name[7:] # Remove "models/" prefix (7 characters) - return model_name + if not model_name or not isinstance(model_name, str): + return model_name + # Last path segment, so full Vertex resource paths + # ("projects/p/locations/l/publishers/google/models/gemini-2.5-flash") reduce + # correctly too, not just the "models/" prefix. Cost rows are keyed on bare + # slugs, so a prefixed name silently prices at $0. + return model_name.rsplit("/", 1)[-1] or model_name def trace_gemini( @@ -209,6 +251,7 @@ def stream_chunks( first_token_time = None num_of_completion_tokens = 0 num_of_prompt_tokens = 0 + num_of_total_tokens = 0 latency = None try: @@ -227,13 +270,14 @@ def stream_chunks( if hasattr(chunk, "text"): collected_output_data.append(chunk.text) - # Extract token counts if available - if hasattr(chunk, "usage_metadata"): - usage = chunk.usage_metadata - if hasattr(usage, "prompt_token_count"): - num_of_prompt_tokens = usage.prompt_token_count - if hasattr(usage, "candidates_token_count"): - num_of_completion_tokens = usage.candidates_token_count + # Extract token counts if available. usage_metadata is cumulative + # across chunks, so the last chunk's values win rather than summing. + if getattr(chunk, "usage_metadata", None) is not None: + ( + num_of_prompt_tokens, + num_of_completion_tokens, + num_of_total_tokens, + ) = _extract_token_counts(chunk.usage_metadata) yield chunk @@ -253,7 +297,7 @@ def stream_chunks( inputs={"prompt": _format_input_messages(contents)}, output=output_data, latency=latency, - tokens=num_of_prompt_tokens + num_of_completion_tokens, + tokens=num_of_total_tokens, prompt_tokens=num_of_prompt_tokens, completion_tokens=num_of_completion_tokens, model=model_name, @@ -320,11 +364,15 @@ async def stream_chunks_async( first_token_time = None num_of_completion_tokens = 0 num_of_prompt_tokens = 0 + num_of_total_tokens = 0 latency = None try: i = 0 - async for i, chunk in enumerate(chunks): + # NOTE: a manual counter, not ``enumerate`` — ``enumerate`` returns a plain + # iterator, so ``async for i, chunk in enumerate(chunks)`` raises + # TypeError. This path could never have worked before. + async for chunk in chunks: # Store raw output try: raw_outputs.append(_serialize_chunk(chunk)) @@ -333,18 +381,20 @@ async def stream_chunks_async( if i == 0: first_token_time = time.time() + i += 1 # Extract text content from chunk if hasattr(chunk, "text"): collected_output_data.append(chunk.text) - # Extract token counts if available - if hasattr(chunk, "usage_metadata"): - usage = chunk.usage_metadata - if hasattr(usage, "prompt_token_count"): - num_of_prompt_tokens = usage.prompt_token_count - if hasattr(usage, "candidates_token_count"): - num_of_completion_tokens = usage.candidates_token_count + # Extract token counts if available. usage_metadata is cumulative + # across chunks, so the last chunk's values win rather than summing. + if getattr(chunk, "usage_metadata", None) is not None: + ( + num_of_prompt_tokens, + num_of_completion_tokens, + num_of_total_tokens, + ) = _extract_token_counts(chunk.usage_metadata) yield chunk @@ -364,7 +414,7 @@ async def stream_chunks_async( inputs={"prompt": _format_input_messages(contents)}, output=output_data, latency=latency, - tokens=num_of_prompt_tokens + num_of_completion_tokens, + tokens=num_of_total_tokens, prompt_tokens=num_of_prompt_tokens, completion_tokens=num_of_completion_tokens, model=model_name, @@ -414,22 +464,20 @@ def handle_non_streaming_generate( try: output_data = parse_non_streaming_output_data(response) - # Extract token counts - num_of_prompt_tokens = 0 - num_of_completion_tokens = 0 - if hasattr(response, "usage_metadata"): - usage = response.usage_metadata - if hasattr(usage, "prompt_token_count"): - num_of_prompt_tokens = usage.prompt_token_count - if hasattr(usage, "candidates_token_count"): - num_of_completion_tokens = usage.candidates_token_count + # Extract token counts. Thinking tokens are folded into the completion + # count and the reported total comes from total_token_count. + ( + num_of_prompt_tokens, + num_of_completion_tokens, + num_of_total_tokens, + ) = _extract_token_counts(getattr(response, "usage_metadata", None)) trace_args = create_trace_args( end_time=end_time, inputs={"prompt": _format_input_messages(args[0] if args else kwargs.get("contents"))}, output=output_data, latency=(end_time - start_time) * 1000, - tokens=num_of_prompt_tokens + num_of_completion_tokens, + tokens=num_of_total_tokens, prompt_tokens=num_of_prompt_tokens, completion_tokens=num_of_completion_tokens, model=model_name, @@ -477,22 +525,20 @@ async def handle_non_streaming_generate_async( try: output_data = parse_non_streaming_output_data(response) - # Extract token counts - num_of_prompt_tokens = 0 - num_of_completion_tokens = 0 - if hasattr(response, "usage_metadata"): - usage = response.usage_metadata - if hasattr(usage, "prompt_token_count"): - num_of_prompt_tokens = usage.prompt_token_count - if hasattr(usage, "candidates_token_count"): - num_of_completion_tokens = usage.candidates_token_count + # Extract token counts. Thinking tokens are folded into the completion + # count and the reported total comes from total_token_count. + ( + num_of_prompt_tokens, + num_of_completion_tokens, + num_of_total_tokens, + ) = _extract_token_counts(getattr(response, "usage_metadata", None)) trace_args = create_trace_args( end_time=end_time, inputs={"prompt": _format_input_messages(args[0] if args else kwargs.get("contents"))}, output=output_data, latency=(end_time - start_time) * 1000, - tokens=num_of_prompt_tokens + num_of_completion_tokens, + tokens=num_of_total_tokens, prompt_tokens=num_of_prompt_tokens, completion_tokens=num_of_completion_tokens, model=model_name, @@ -614,7 +660,7 @@ def create_trace_args( def add_to_trace(**kwargs) -> None: """Add a chat completion step to the trace.""" - tracer.add_chat_completion_step_to_trace(**kwargs, name="Gemini Generation", provider="Google") + tracer.add_chat_completion_step_to_trace(**kwargs, name="Gemini Generation", provider=PROVIDER) def _format_input_messages(contents: Any) -> list: diff --git a/src/openlayer/lib/integrations/google_adk_tracer.py b/src/openlayer/lib/integrations/google_adk_tracer.py index fc724d25..5c92ec27 100644 --- a/src/openlayer/lib/integrations/google_adk_tracer.py +++ b/src/openlayer/lib/integrations/google_adk_tracer.py @@ -898,6 +898,13 @@ async def new_function(): ) as step: # Set ChatCompletionStep attributes step.model = model_name + # NOTE: deliberately NOT the "gemini" cost slug used by the Gemini + # tracers. This wrapper is model-agnostic — BaseLlmFlow._call_llm_async + # resolves the LLM via __get_llm() and calls generate_content_async on + # whatever it gets, so LiteLlm and Claude-on-Vertex agents reach here + # too. Hardcoding "gemini" would pair it with e.g. a claude-* model and + # miss the cost lookup entirely. Fixing this properly means deriving the + # slug per-model; tracked separately. step.provider = "Google" step.model_parameters = model_parameters diff --git a/src/openlayer/lib/integrations/google_genai_tracer.py b/src/openlayer/lib/integrations/google_genai_tracer.py new file mode 100644 index 00000000..dc855619 --- /dev/null +++ b/src/openlayer/lib/integrations/google_genai_tracer.py @@ -0,0 +1,818 @@ +"""Module with methods used to trace the unified Google Gen AI SDK. + +Targets the CURRENT Google Gen AI SDK — package ``google-genai``, module +``google.genai``, client ``genai.Client()`` — in both AI Studio mode and Vertex +AI mode (``vertexai=True``). Generation lives on a sub-object here +(``client.models.generate_content``), not on the client itself. + +The LEGACY ``google-generativeai`` SDK (``genai.GenerativeModel``) is handled by +``gemini_tracer.py`` instead. Both packages can be installed side by side, so +both tracers are registered independently in ``_auto.py``. + +Coverage: ``generate_content`` and ``generate_content_stream``, sync and async +(``client.aio.models.*``). ``client.chats`` is covered transitively — +``Chats(modules=client.models)`` shares the patched ``Models`` instance and +``Chat.send_message`` delegates to ``generate_content``. +""" + +import json +import logging +import sys +import time +from functools import wraps +from typing import Any, Dict, List, Optional, Tuple, Union, TYPE_CHECKING + +try: + from google import genai + + HAVE_GOOGLE_GENAI = True +except ImportError: + HAVE_GOOGLE_GENAI = False + +if TYPE_CHECKING: + from google import genai + +from ..tracing import tracer + +logger = logging.getLogger(__name__) + +# Cost lookup on the backend is an exact ``(provider.lower(), model)`` match with +# no aliasing, so ``provider`` has to BE a real cost-table slug rather than a +# display label. "gemini" is the LiteLLM-sourced slug; the friendlier "Google" +# label misses newer models (e.g. gemini-3-pro-preview) and carries stale +# OpenRouter-sourced prices for the ``-latest`` aliases. +PROVIDER = "gemini" + +# Kept identical to the legacy tracer's step name so spans from the two Google +# SDKs don't fragment dashboards. +STEP_NAME = "Gemini Generation" + +# Checked via sys.modules rather than imported: google_adk_tracer is ~66KB, and a +# guarded lazy import would construct an ImportError on every LLM call in the +# common case where ADK isn't installed. +_ADK_TRACER_MODULE = "openlayer.lib.integrations.google_adk_tracer" + +_MODEL_PARAM_KEYS = ( + "temperature", + "top_p", + "top_k", + "max_output_tokens", + "candidate_count", + "stop_sequences", + "seed", + "presence_penalty", + "frequency_penalty", + "response_mime_type", +) + + +# ----------------------------- Public API ----------------------------- # +def trace_google_genai( + client: "genai.Client", +) -> "genai.Client": + """Patch a ``google.genai.Client`` to trace content generation. + + The following information is collected for each generation: + - start_time / end_time / latency + - tokens, prompt_tokens, completion_tokens (thinking tokens included, see + ``_extract_usage``) + - model, model_parameters + - inputs, output, raw_output + - metadata: ``llm_system`` (``google_vertex`` or ``google_ai_studio``), GCP + project/location in Vertex mode, the raw token split, and + ``timeToFirstToken`` when streaming. + + Parameters + ---------- + client : genai.Client + The Google Gen AI client to patch. Works for both + ``genai.Client(api_key=...)`` and + ``genai.Client(vertexai=True, project=..., location=...)``. + + Returns + ------- + genai.Client + The same client, patched in place. + """ + if not HAVE_GOOGLE_GENAI: + raise ImportError("google-genai is required for Google Gen AI tracing. Install with: pip install google-genai") + + if getattr(client, "_openlayer_patched", False) is True: + return client + + # `client.models` / `client.aio` are read-only properties backed by instances + # built eagerly in Client.__init__, so we patch the attribute ON the + # sub-object. Assigning `client.models = ...` would raise. + _patch_models(getattr(client, "models", None), client, is_async=False) + try: + aio = getattr(client, "aio", None) + _patch_models(getattr(aio, "models", None), client, is_async=True) + except Exception as e: # pylint: disable=broad-except + logger.debug("Openlayer: could not patch the async Google Gen AI surface: %s", e) + + try: + client._openlayer_patched = True # pylint: disable=protected-access + except Exception: # pylint: disable=broad-except + # Marking the sub-objects (done in _patch_models) already makes this + # idempotent, so a read-only client is not fatal. + logger.debug("Openlayer: could not mark the Google Gen AI client as patched") + + return client + + +def _patch_google_genai() -> None: + """Patch ``google.genai.Client.__init__`` so every newly-constructed client + is auto-traced. Idempotent.""" + if not HAVE_GOOGLE_GENAI: + return + # pylint: disable=import-outside-toplevel + from ._auto import _patch_class_init + + _patch_class_init(genai.Client, trace_google_genai) + + +def _unpatch_google_genai() -> None: + if not HAVE_GOOGLE_GENAI: + return + # pylint: disable=import-outside-toplevel + from ._auto import _unpatch_class_init + + _unpatch_class_init(genai.Client) + + +# ----------------------------- Patching ----------------------------- # +def _patch_models(models_obj: Any, client: Any, is_async: bool) -> None: + """Wrap ``generate_content`` / ``generate_content_stream`` on a ``Models`` or + ``AsyncModels`` instance. Idempotent per instance.""" + if models_obj is None or getattr(models_obj, "_openlayer_patched", False) is True: + return + + original_generate = getattr(models_obj, "generate_content", None) + original_stream = getattr(models_obj, "generate_content_stream", None) + + if original_generate is not None: + models_obj.generate_content = ( + _wrap_generate_content_async(original_generate, client) + if is_async + else _wrap_generate_content(original_generate, client) + ) + + if original_stream is not None: + models_obj.generate_content_stream = ( + _wrap_stream_async(original_stream, client) if is_async else _wrap_stream(original_stream, client) + ) + + models_obj._openlayer_patched = True # pylint: disable=protected-access + + +def _adk_span_active() -> bool: + """True when the Google ADK tracer already has an LLM step open. + + ADK's ``Gemini.api_client`` is a ``google.genai.Client``, and its + ``generate_content_async`` calls ``client.aio.models.generate_content``. Since + ``google-adk`` depends on ``google-genai``, both integrations activate for ADK + users, and without this check every ADK LLM call would emit two spans and be + billed twice. + + The ADK tracer sets ``_current_llm_step`` inside the ``create_step`` block + that wraps the call, so our patched method runs in a context where it's + visible (async generators have no independent context, so each ``__anext__`` + runs in the driver's context). + """ + module = sys.modules.get(_ADK_TRACER_MODULE) + if module is None: + return False + context_var = getattr(module, "_current_llm_step", None) + if context_var is None: + return False + try: + return context_var.get() is not None + except LookupError: + return False + + +def _wrap_generate_content(original: Any, client: Any) -> Any: + @wraps(original) + def traced_generate_content(*args, **kwargs): + # Popped before any call-through so the SDK never sees the extra kwarg. + inference_id = kwargs.pop("inference_id", None) + if _adk_span_active(): + return original(*args, **kwargs) + + start_time = time.time() + response = original(*args, **kwargs) + end_time = time.time() + + try: + _add_span( + client=client, + kwargs=kwargs, + output=parse_non_streaming_output_data(response), + usage=getattr(response, "usage_metadata", None), + raw_output=_serialize(response), + model_version=getattr(response, "model_version", None), + start_time=start_time, + end_time=end_time, + inference_id=inference_id, + ) + except Exception as e: # pylint: disable=broad-except + logger.error("Failed to trace the generate content request with Openlayer. %s", e) + + return response + + return traced_generate_content + + +def _wrap_generate_content_async(original: Any, client: Any) -> Any: + @wraps(original) + async def traced_generate_content_async(*args, **kwargs): + inference_id = kwargs.pop("inference_id", None) + if _adk_span_active(): + return await original(*args, **kwargs) + + start_time = time.time() + response = await original(*args, **kwargs) + end_time = time.time() + + try: + _add_span( + client=client, + kwargs=kwargs, + output=parse_non_streaming_output_data(response), + usage=getattr(response, "usage_metadata", None), + raw_output=_serialize(response), + model_version=getattr(response, "model_version", None), + start_time=start_time, + end_time=end_time, + inference_id=inference_id, + ) + except Exception as e: # pylint: disable=broad-except + logger.error("Failed to trace the async generate content request with Openlayer. %s", e) + + return response + + return traced_generate_content_async + + +def _wrap_stream(original: Any, client: Any) -> Any: + @wraps(original) + def traced_generate_content_stream(*args, **kwargs): + inference_id = kwargs.pop("inference_id", None) + if _adk_span_active(): + return original(*args, **kwargs) + stream = original(*args, **kwargs) + return stream_chunks(stream, client=client, kwargs=kwargs, inference_id=inference_id) + + return traced_generate_content_stream + + +def _wrap_stream_async(original: Any, client: Any) -> Any: + @wraps(original) + async def traced_generate_content_stream_async(*args, **kwargs): + # ``AsyncModels.generate_content_stream`` is a coroutine function that + # returns an AsyncIterator — callers ``await`` it and *then* ``async for``. + # So this wrapper must be ``async def`` returning an async generator; an + # async generator here would break the ``await``. + inference_id = kwargs.pop("inference_id", None) + if _adk_span_active(): + return await original(*args, **kwargs) + stream = await original(*args, **kwargs) + return stream_chunks_async(stream, client=client, kwargs=kwargs, inference_id=inference_id) + + return traced_generate_content_stream_async + + +# ----------------------------- Streaming ----------------------------- # +def stream_chunks( + stream: Any, + client: Any, + kwargs: Dict[str, Any], + inference_id: Optional[str] = None, +): + """Yield chunks through unchanged while accumulating a span.""" + collected_output_data: List[str] = [] + raw_outputs: List[Any] = [] + usage = None + model_version = None + start_time = time.time() + first_token_time = None + end_time = None + + try: + for i, chunk in enumerate(stream): + if i == 0: + first_token_time = time.time() + + # usage_metadata is present on EVERY chunk and is cumulative, so the + # last one wins. Summing would multiply-count. + chunk_usage = getattr(chunk, "usage_metadata", None) + if chunk_usage is not None: + usage = chunk_usage + + chunk_model_version = getattr(chunk, "model_version", None) + if chunk_model_version: + model_version = chunk_model_version + + text = _safe_text(chunk) + if text: + collected_output_data.append(text) + + try: + raw_outputs.append(_serialize(chunk)) + except Exception as e: # pylint: disable=broad-except + logger.debug("Failed to serialize chunk: %s", e) + + yield chunk + + end_time = time.time() + finally: + # ``finally`` so the span is still emitted when the consumer abandons the + # iterator part-way (GeneratorExit) or the stream raises. + if end_time is None: + end_time = time.time() + try: + _add_span( + client=client, + kwargs=kwargs, + output="".join(collected_output_data), + usage=usage, + raw_output=raw_outputs, + model_version=model_version, + start_time=start_time, + end_time=end_time, + inference_id=inference_id, + metadata_extra={ + "timeToFirstToken": ((first_token_time - start_time) * 1000 if first_token_time else None) + }, + ) + except Exception as e: # pylint: disable=broad-except + logger.error("Failed to trace the streaming generate content request with Openlayer. %s", e) + + +async def stream_chunks_async( + stream: Any, + client: Any, + kwargs: Dict[str, Any], + inference_id: Optional[str] = None, +): + """Async counterpart of :func:`stream_chunks`.""" + collected_output_data: List[str] = [] + raw_outputs: List[Any] = [] + usage = None + model_version = None + start_time = time.time() + first_token_time = None + end_time = None + + try: + i = 0 + # NOTE: a manual counter, not ``enumerate`` — ``enumerate`` returns a + # plain iterator and cannot be driven by ``async for``. + async for chunk in stream: + if i == 0: + first_token_time = time.time() + i += 1 + + chunk_usage = getattr(chunk, "usage_metadata", None) + if chunk_usage is not None: + usage = chunk_usage + + chunk_model_version = getattr(chunk, "model_version", None) + if chunk_model_version: + model_version = chunk_model_version + + text = _safe_text(chunk) + if text: + collected_output_data.append(text) + + try: + raw_outputs.append(_serialize(chunk)) + except Exception as e: # pylint: disable=broad-except + logger.debug("Failed to serialize chunk: %s", e) + + yield chunk + + end_time = time.time() + finally: + if end_time is None: + end_time = time.time() + try: + _add_span( + client=client, + kwargs=kwargs, + output="".join(collected_output_data), + usage=usage, + raw_output=raw_outputs, + model_version=model_version, + start_time=start_time, + end_time=end_time, + inference_id=inference_id, + metadata_extra={ + "timeToFirstToken": ((first_token_time - start_time) * 1000 if first_token_time else None) + }, + ) + except Exception as e: # pylint: disable=broad-except + logger.error("Failed to trace the async streaming generate content request with Openlayer. %s", e) + + +# ----------------------------- Span assembly ----------------------------- # +def _add_span( + client: Any, + kwargs: Dict[str, Any], + output: Union[str, Dict[str, Any], None], + usage: Any, + raw_output: Any, + model_version: Optional[str], + start_time: float, + end_time: float, + inference_id: Optional[str] = None, + metadata_extra: Optional[Dict[str, Any]] = None, +) -> None: + """Build the trace arguments and add the chat completion step.""" + requested_model = kwargs.get("model") + model = _normalize_model_name(requested_model) + prompt_tokens, completion_tokens, total_tokens, token_extras = _extract_usage(usage) + + metadata = _mode_metadata(client) + metadata.update(token_extras) + if model_version: + metadata["modelVersion"] = model_version + if requested_model and requested_model != model: + # Keeps the caller's original string recoverable after normalization. + metadata["requestedModel"] = requested_model + if metadata_extra: + metadata.update(metadata_extra) + + trace_args = create_trace_args( + end_time=end_time, + inputs={"prompt": _format_input_messages(kwargs.get("contents"), kwargs.get("config"))}, + output=output, + latency=(end_time - start_time) * 1000, + tokens=total_tokens, + prompt_tokens=prompt_tokens, + completion_tokens=completion_tokens, + model=model, + model_parameters=get_model_parameters(kwargs.get("config")), + raw_output=raw_output, + id=inference_id, + metadata=metadata, + ) + add_to_trace(**trace_args) + + +def create_trace_args( + end_time: float, + inputs: Dict[str, Any], + output: Union[str, Dict[str, Any], None], + latency: float, + tokens: int, + prompt_tokens: int, + completion_tokens: int, + model: str, + model_parameters: Optional[Dict[str, Any]] = None, + metadata: Optional[Dict[str, Any]] = None, + raw_output: Optional[Any] = None, + id: Optional[str] = None, # pylint: disable=redefined-builtin +) -> Dict[str, Any]: + """Returns a dictionary with the trace arguments.""" + trace_args = { + "end_time": end_time, + "inputs": inputs, + "output": output, + "latency": latency, + "tokens": tokens, + "prompt_tokens": prompt_tokens, + "completion_tokens": completion_tokens, + "model": model, + "model_parameters": model_parameters, + "raw_output": raw_output, + "metadata": metadata if metadata else {}, + } + if id: + trace_args["id"] = id + return trace_args + + +def add_to_trace(**kwargs) -> None: + """Add a chat completion step to the trace.""" + tracer.add_chat_completion_step_to_trace(**kwargs, name=STEP_NAME, provider=PROVIDER) + + +def _mode_metadata(client: Any) -> Dict[str, Any]: + """Distinguish Vertex mode from AI Studio mode. + + The same ``Client`` class serves both, so the mode is only knowable from the + instance. ``llm_system`` matches what the Google ADK tracer already emits. + """ + metadata: Dict[str, Any] = {} + try: + is_vertex = bool(getattr(client, "vertexai", False)) + except Exception: # pylint: disable=broad-except + is_vertex = False + + metadata["llm_system"] = "google_vertex" if is_vertex else "google_ai_studio" + + if is_vertex: + api_client = getattr(client, "_api_client", None) + if api_client is not None: + project = getattr(api_client, "project", None) + location = getattr(api_client, "location", None) + if project: + metadata["gcp_project"] = project + if location: + metadata["gcp_location"] = location + + return metadata + + +# ----------------------------- Extraction helpers ----------------------------- # +def _normalize_model_name(model: Optional[str]) -> str: + """Reduce a Gemini model identifier to the bare slug the cost table uses. + + Cost rows are keyed on bare slugs, so anything else silently prices at $0 + (verified: ``models/gemini-2.5-flash`` returned $0.00 where bare + ``gemini-2.5-flash`` returned $2.80 on identical tokens). + + Handles every shape the API accepts:: + + gemini-2.5-flash -> gemini-2.5-flash + models/gemini-2.5-flash -> gemini-2.5-flash + publishers/google/models/gemini-2.5-flash -> gemini-2.5-flash + projects/p/locations/l/publishers/google/models/gemini-2.5-flash -> gemini-2.5-flash + + Version suffixes are deliberately preserved — upstream coverage for them is + inconsistent, and rewriting them would misreport which model actually ran. + """ + if not model or not isinstance(model, str): + return "unknown" + return model.rsplit("/", 1)[-1] or model + + +def _extract_usage(usage: Any) -> Tuple[int, int, int, Dict[str, Any]]: + """Return ``(prompt_tokens, completion_tokens, total_tokens, extras)``. + + Thinking tokens are the subtlety here. ``candidates_token_count`` EXCLUDES + ``thoughts_token_count``, and thinking is on by default for 2.5-series + models, so the naive ``prompt + candidates`` badly undercounts — measured 14 + for a call that really consumed 757 (prompt 8, candidates 6, thoughts 743). + + So ``tokens`` uses ``total_token_count`` (documented as prompt + candidates + + tool_use_prompt + thoughts), and thinking tokens are folded into + ``completion_tokens`` because Vertex bills them as output and the backend + prices from ``completionTokens``. The raw split is preserved in metadata. + """ + if usage is None: + return 0, 0, 0, {} + + def _count(name: str) -> Optional[int]: + value = getattr(usage, name, None) + return value if isinstance(value, int) else None + + prompt_tokens = _count("prompt_token_count") or 0 + candidates_tokens = _count("candidates_token_count") or 0 + thoughts_tokens = _count("thoughts_token_count") or 0 + tool_use_tokens = _count("tool_use_prompt_token_count") or 0 + cached_tokens = _count("cached_content_token_count") or 0 + total_tokens = _count("total_token_count") + + completion_tokens = candidates_tokens + thoughts_tokens + if total_tokens is None: + total_tokens = prompt_tokens + completion_tokens + tool_use_tokens + + extras: Dict[str, Any] = {} + if candidates_tokens: + extras["candidatesTokens"] = candidates_tokens + if thoughts_tokens: + extras["thoughtsTokens"] = thoughts_tokens + if tool_use_tokens: + extras["toolUsePromptTokens"] = tool_use_tokens + if cached_tokens: + extras["cachedContentTokens"] = cached_tokens + + return prompt_tokens, completion_tokens, total_tokens, extras + + +def _config_to_dict(config: Any) -> Dict[str, Any]: + """Normalize a ``GenerateContentConfig`` (pydantic) or its dict form.""" + if config is None: + return {} + if isinstance(config, dict): + return dict(config) + model_dump = getattr(config, "model_dump", None) + if callable(model_dump): + try: + return model_dump(exclude_none=True) + except Exception: # pylint: disable=broad-except + pass + return dict(getattr(config, "__dict__", None) or {}) + + +def get_model_parameters(config: Any) -> Dict[str, Any]: + """Gets the model parameters from the request config. + + Unlike the legacy SDK's loose ``generation_config`` kwarg, the unified SDK + carries these on a pydantic ``GenerateContentConfig``. + """ + config_dict = _config_to_dict(config) + parameters: Dict[str, Any] = {key: config_dict.get(key) for key in _MODEL_PARAM_KEYS} + + thinking_config = config_dict.get("thinking_config") + if isinstance(thinking_config, dict): + if thinking_config.get("thinking_budget") is not None: + parameters["thinking_budget"] = thinking_config["thinking_budget"] + if thinking_config.get("include_thoughts") is not None: + parameters["include_thoughts"] = thinking_config["include_thoughts"] + + return parameters + + +def _safe_text(obj: Any) -> Optional[str]: + """Read ``.text`` without letting the SDK's accessor raise or warn through. + + ``GenerateContentResponse.text`` raises when the response holds no text part + (function calls, safety blocks), so it can never be read bare. + """ + try: + text = getattr(obj, "text", None) + except Exception: # pylint: disable=broad-except + return None + return text if isinstance(text, str) else None + + +def parse_non_streaming_output_data( + response: Any, +) -> Union[str, Dict[str, Any], None]: + """Parses the output data from a non-streaming generation.""" + text = _safe_text(response) + if text and text.strip(): + return text.strip() + + try: + candidates = getattr(response, "candidates", None) or [] + for candidate in candidates: + content = getattr(candidate, "content", None) + parts = getattr(content, "parts", None) or [] + text_parts: List[str] = [] + for part in parts: + part_text = getattr(part, "text", None) + if isinstance(part_text, str): + text_parts.append(part_text) + continue + function_call = getattr(part, "function_call", None) + if function_call is not None: + return { + "name": getattr(function_call, "name", None) or "", + "arguments": dict(getattr(function_call, "args", None) or {}), + } + if text_parts: + return " ".join(text_parts).strip() + except Exception as e: # pylint: disable=broad-except + logger.debug("Could not parse Google Gen AI output data: %s", e) + + return None + + +def _format_input_messages(contents: Any, config: Any = None) -> List[Dict[str, Any]]: + """Format request contents into a messages array. + + ``config.system_instruction`` is surfaced as a leading system message, since + the unified SDK carries it on the config rather than in ``contents``. + """ + messages: List[Dict[str, Any]] = [] + + system_instruction = _config_to_dict(config).get("system_instruction") + if system_instruction: + system_message = _to_message(system_instruction) + system_message["role"] = "system" + messages.append(system_message) + + if contents is None: + return messages + + if isinstance(contents, (str, bytes)): + messages.append({"role": "user", "content": _as_text(contents)}) + return messages + + if isinstance(contents, list): + for item in contents: + messages.append(_to_message(item)) + return messages + + messages.append(_to_message(contents)) + return messages + + +def _as_text(value: Any) -> str: + if isinstance(value, bytes): + try: + return value.decode("utf-8", errors="replace") + except Exception: # pylint: disable=broad-except + return str(value) + return value if isinstance(value, str) else str(value) + + +def _to_message(item: Any) -> Dict[str, Any]: + """Convert a ``Content``, ``Part``, dict or scalar into a chat message.""" + if isinstance(item, (str, bytes)): + return {"role": "user", "content": _as_text(item)} + + role = None + parts = None + if isinstance(item, dict): + role = item.get("role") + parts = item.get("parts") + if parts is None and "content" in item: + return {"role": role or "user", "content": item["content"]} + else: + role = getattr(item, "role", None) + parts = getattr(item, "parts", None) + + if parts: + text_parts: List[str] = [] + for part in parts: + part_text = _part_to_text(part) + if part_text: + text_parts.append(part_text) + if text_parts: + return {"role": role or "user", "content": " ".join(text_parts)} + + # A bare Part (no role/parts of its own) is still worth capturing. + single = _part_to_text(item) + if single: + return {"role": role or "user", "content": single} + + return {"role": role or "user", "content": str(item)} + + +def _part_to_text(part: Any) -> Optional[str]: + """Best-effort text for a single ``Part``-like object.""" + if isinstance(part, (str, bytes)): + return _as_text(part) + + if isinstance(part, dict): + if isinstance(part.get("text"), str): + return part["text"] + if part.get("function_call") is not None: + return json.dumps({"function_call": part["function_call"]}, default=str) + if part.get("function_response") is not None: + return json.dumps({"function_response": part["function_response"]}, default=str) + if part.get("inline_data") is not None: + return "" + return None + + text = getattr(part, "text", None) + if isinstance(text, str): + return text + + function_call = getattr(part, "function_call", None) + if function_call is not None: + return json.dumps( + { + "function_call": { + "name": getattr(function_call, "name", None) or "", + "args": dict(getattr(function_call, "args", None) or {}), + } + }, + default=str, + ) + + function_response = getattr(part, "function_response", None) + if function_response is not None: + return json.dumps( + { + "function_response": { + "name": getattr(function_response, "name", None) or "", + "response": getattr(function_response, "response", None), + } + }, + default=str, + ) + + if getattr(part, "inline_data", None) is not None: + return "" + + return None + + +def _serialize(obj: Any) -> Any: + """Serialize a response or chunk to something JSON-safe.""" + model_dump = getattr(obj, "model_dump", None) + if callable(model_dump): + try: + # mode="json" is what makes the result safe to publish — it coerces + # datetimes, enums and bytes that the default mode leaves as objects. + return model_dump(mode="json", exclude_none=True) + except Exception: # pylint: disable=broad-except + pass + + model_dump_json = getattr(obj, "model_dump_json", None) + if callable(model_dump_json): + try: + return json.loads(model_dump_json(exclude_none=True)) + except Exception: # pylint: disable=broad-except + pass + + if isinstance(obj, dict): + return obj + + return {"response": str(obj)} diff --git a/tests/test_google_genai_adk_dedup.py b/tests/test_google_genai_adk_dedup.py new file mode 100644 index 00000000..037dc476 --- /dev/null +++ b/tests/test_google_genai_adk_dedup.py @@ -0,0 +1,189 @@ +"""Real-integration test for ADK / google-genai double-instrumentation. + +``google-adk`` depends on ``google-genai``, and ``Gemini.api_client`` IS a +``google.genai.Client`` whose ``aio.models.generate_content`` the ADK flow calls +directly. So registering a ``google.genai`` auto-instrument entry activates it for +every ADK user, and without contextual suppression each ADK LLM call would emit +two chat_completion spans and be billed twice. + +Unit tests for ``_adk_span_active`` inject a fake module into ``sys.modules``, +which proves the *check* works but not that the real ``_current_llm_step`` +contextvar is visible at the moment our patched method runs. This file drives a +real ``Runner`` + ``LlmAgent`` with both tracers active and asserts the span +count, which is the only thing that discriminates. + +No network: ``AsyncModels.generate_content`` is stubbed on the CLASS before the +class-init patch, so ADK's real call path is exercised end to end. +""" + +# Neither google-adk nor google-genai is installed in the lint env, so imports +# from the `google` namespace package don't resolve there. +# pyright: reportMissingImports=false, reportUnknownMemberType=false, reportUnknownVariableType=false, reportUnknownArgumentType=false, reportUnknownParameterType=false, reportMissingParameterType=false, reportUnusedFunction=false +# pyright: reportMissingTypeStubs=false, reportAttributeAccessIssue=false, reportCallIssue=false + +import asyncio +from typing import Any, Dict, List +from unittest.mock import patch + +import pytest + +pytest.importorskip("google.adk") +pytest.importorskip("google.genai") +pytest.importorskip("wrapt") + +from openlayer.lib.integrations import google_genai_tracer as gg + +APP_NAME = "openlayer_open_11891" +USER_ID = "u" +SESSION_ID = "s" + + +@pytest.fixture(autouse=True) +def _disable_publish(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv("OPENLAYER_DISABLE_PUBLISH", "true") + monkeypatch.setenv("OPENLAYER_API_KEY", "fake") + monkeypatch.setenv("GOOGLE_API_KEY", "fake") + + from openlayer.lib.tracing import tracer as _tracer + + monkeypatch.setattr(_tracer, "_publish", False, raising=False) + + +@pytest.fixture(autouse=True) +def _reset_patches(): + """Both tracers keep module/class-level state; clear it between tests.""" + yield + from openlayer.lib.integrations.google_adk_tracer import _unpatch_google_adk + + _unpatch_google_adk() + gg._unpatch_google_genai() + + +@pytest.fixture +def stub_model(monkeypatch: pytest.MonkeyPatch): + """Stub ``AsyncModels.generate_content`` at the class level. + + Installed BEFORE the class-init patch so the per-instance wrapper wraps the + stub — i.e. ADK -> genai.Client -> our wrapper -> stub, the real chain. + """ + from google.genai import types, models as genai_models + + response = types.GenerateContentResponse( + candidates=[ + types.Candidate( + content=types.Content(role="model", parts=[types.Part(text="hello from stub")]), + finish_reason=types.FinishReason.STOP, + ) + ], + usage_metadata=types.GenerateContentResponseUsageMetadata( + prompt_token_count=5, + candidates_token_count=3, + total_token_count=8, + ), + model_version="gemini-2.5-flash", + ) + + calls: List[Dict[str, Any]] = [] + + async def _impl(_self: Any, **kwargs: Any) -> Any: + calls.append(kwargs) + return response + + monkeypatch.setattr(genai_models.AsyncModels, "generate_content", _impl) + return calls + + +def _run_one_agent_turn() -> None: + from google.genai import types + from google.adk.agents import LlmAgent + from google.adk.runners import Runner + from google.adk.sessions import InMemorySessionService + + agent = LlmAgent(model="gemini-2.5-flash", name="OpenlayerDedupAgent", instruction="test") + session_service = InMemorySessionService() + runner = Runner(agent=agent, app_name=APP_NAME, session_service=session_service) + + async def _drive() -> None: + await session_service.create_session(app_name=APP_NAME, user_id=USER_ID, session_id=SESSION_ID) + message = types.Content(role="user", parts=[types.Part(text="hi")]) + async for _event in runner.run_async(user_id=USER_ID, session_id=SESSION_ID, new_message=message): + pass + + asyncio.run(_drive()) + + +def _count_chat_completion_steps(step: Any) -> int: + from openlayer.lib.tracing import enums + + count = 1 if getattr(step, "step_type", None) == enums.StepType.CHAT_COMPLETION else 0 + for nested in getattr(step, "steps", None) or []: + count += _count_chat_completion_steps(nested) + return count + + +@pytest.mark.filterwarnings("ignore::DeprecationWarning") +class TestAdkDoubleInstrumentation: + def test_genai_span_suppressed_when_adk_tracer_is_active(self, stub_model: Any) -> None: + """Both tracers on: the ADK tracer owns the span, ours must stand down.""" + from openlayer.lib.integrations import trace_google_adk + + trace_google_adk() + gg._patch_google_genai() + + with patch.object(gg, "add_to_trace") as mock_add: + _run_one_agent_turn() + + assert stub_model, "the ADK flow must actually have reached genai.Client" + assert mock_add.call_count == 0, ( + "google_genai_tracer emitted a span inside an ADK LLM step -- " + "every ADK call would be double-counted and double-billed" + ) + + def test_exactly_one_chat_completion_step_in_the_trace(self, stub_model: Any) -> None: + """The user-visible property: one LLM call produces one span.""" + from openlayer.lib.tracing import tracer as _tracer + from openlayer.lib.integrations import trace_google_adk + + trace_google_adk() + gg._patch_google_genai() + + captured: Dict[str, Any] = {} + original_handle = _tracer._handle_trace_completion + + def _capture(*args: Any, **kwargs: Any) -> Any: + current = _tracer.get_current_trace() + if current is not None and "trace" not in captured: + captured["trace"] = current + return original_handle(*args, **kwargs) + + with patch.object(_tracer, "_handle_trace_completion", _capture): + _run_one_agent_turn() + + assert len(stub_model) == 1, "the agent turn should make exactly one LLM call" + + trace = captured.get("trace") + assert trace is not None, "no trace was captured" + total = sum(_count_chat_completion_steps(step) for step in trace.steps) + assert total == 1, f"expected exactly 1 chat_completion step, found {total}" + + def test_span_emitted_when_only_the_genai_tracer_is_active(self, stub_model: Any) -> None: + """Control: suppression must be CONTEXTUAL, not 'is ADK importable'. + + ``google_adk_tracer`` stays in ``sys.modules`` once any earlier test has + imported it, so a module-presence check would wrongly suppress here. + """ + import sys + + gg._patch_google_genai() + + with patch.object(gg, "add_to_trace") as mock_add: + _run_one_agent_turn() + + assert stub_model, "the ADK flow must actually have reached genai.Client" + assert mock_add.call_count >= 1, ( + "with no ADK tracer active, the genai tracer is the only thing that " + "can trace this call and must not stay silent" + ) + # Guard the premise of this test: if the module were absent, the control + # would pass trivially rather than proving contextual behaviour. + assert gg._ADK_TRACER_MODULE in sys.modules diff --git a/tests/test_google_genai_integration.py b/tests/test_google_genai_integration.py new file mode 100644 index 00000000..aa3cc178 --- /dev/null +++ b/tests/test_google_genai_integration.py @@ -0,0 +1,556 @@ +"""Tests for the unified Google Gen AI (``google-genai``) tracer. + +Covers OPEN-11891: ``trace_gemini`` previously supported only the legacy +``google-generativeai`` SDK, so users of ``genai.Client()`` (including Vertex +mode) got no auto-instrumentation and a hard error on explicit tracing. + +No network calls — a real ``genai.Client`` is constructed with a fake key and its +``generate_content`` methods are replaced with stubs *before* tracing, so the +tracer wraps the stub. Spans are asserted by patching the tracer module's +``add_to_trace``, matching tests/test_bedrock_integration.py. +""" + +# google-genai isn't installed in the lint env, and pytest fixtures hide autouse +# functions from static analysis. The last two matter in BOTH directions: without +# the package, `from google import genai` is an unresolved attribute on a +# namespace package; with it, `inference_id` is an extra kwarg the SDK's own +# signature doesn't declare (that's the whole point -- the tracer pops it). +# pyright: reportMissingImports=false, reportUnknownMemberType=false, reportUnknownVariableType=false, reportUnknownArgumentType=false, reportUnknownParameterType=false, reportMissingParameterType=false, reportUnusedFunction=false +# pyright: reportMissingTypeStubs=false, reportAttributeAccessIssue=false, reportCallIssue=false + +import asyncio +import contextvars +from types import SimpleNamespace +from typing import Any, Dict, List, Optional +from unittest.mock import patch + +import pytest + +pytest.importorskip("google.genai") + +from openlayer.lib.integrations import google_genai_tracer as gg + + +# ------------------------------- fixtures ------------------------------- # +@pytest.fixture(autouse=True) +def _disable_publish(monkeypatch: pytest.MonkeyPatch) -> None: + """Keep every tracer publish path off.""" + monkeypatch.setenv("OPENLAYER_DISABLE_PUBLISH", "true") + monkeypatch.setenv("OPENLAYER_API_KEY", "fake") + + from openlayer.lib.tracing import tracer as _tracer + + monkeypatch.setattr(_tracer, "_publish", False, raising=False) + + +@pytest.fixture(autouse=True) +def _reset_client_class_patch(): + """Undo any class-level ``Client.__init__`` patch so the module-level + idempotency marker doesn't leak between tests.""" + yield + gg._unpatch_google_genai() + + +# ------------------------------- helpers ------------------------------- # +def _usage( + prompt: Optional[int] = None, + candidates: Optional[int] = None, + thoughts: Optional[int] = None, + total: Optional[int] = None, + tool_use: Optional[int] = None, + cached: Optional[int] = None, +) -> SimpleNamespace: + return SimpleNamespace( + prompt_token_count=prompt, + candidates_token_count=candidates, + thoughts_token_count=thoughts, + total_token_count=total, + tool_use_prompt_token_count=tool_use, + cached_content_token_count=cached, + ) + + +class FakeResponse: + """Minimal stand-in for ``types.GenerateContentResponse``.""" + + def __init__(self, text: str = "hello", usage: Any = None, model_version: Optional[str] = None) -> None: + self.text = text + self.usage_metadata = usage + self.model_version = model_version + self.candidates = [] + # Read by the real Chat.send_message when recording history. + self.automatic_function_calling_history = None + + def model_dump(self, **_kwargs: Any) -> Dict[str, Any]: + return {"text": self.text} + + +def _make_client(vertexai: bool = False): + """A real ``genai.Client`` that will never reach the network.""" + from google import genai + + if vertexai: + return genai.Client(vertexai=True, project="test-project", location="us-central1") + return genai.Client(api_key="fake-key") + + +def _stub_sync(client, response: Any) -> List[Dict[str, Any]]: + """Replace sync generate_content with a stub; return the recorded calls.""" + calls: List[Dict[str, Any]] = [] + + def _impl(**kwargs: Any) -> Any: + calls.append(kwargs) + return response + + client.models.generate_content = _impl + return calls + + +def _stub_sync_stream(client, chunks: List[Any]) -> None: + def _impl(**_kwargs: Any): + return iter(chunks) + + client.models.generate_content_stream = _impl + + +def _stub_async(client, response: Any) -> None: + async def _impl(**_kwargs: Any) -> Any: + return response + + client.aio.models.generate_content = _impl + + +def _stub_async_stream(client, chunks: List[Any]) -> None: + async def _impl(**_kwargs: Any): + async def _gen(): + for chunk in chunks: + yield chunk + + return _gen() + + client.aio.models.generate_content_stream = _impl + + +# ------------------------------- model names ------------------------------- # +class TestNormalizeModelName: + """Prefixed model strings price at $0, so this is load-bearing.""" + + @pytest.mark.parametrize( + "raw", + [ + "gemini-2.5-flash", + "models/gemini-2.5-flash", + "publishers/google/models/gemini-2.5-flash", + "projects/p/locations/us-central1/publishers/google/models/gemini-2.5-flash", + ], + ) + def test_all_four_shapes_reduce_to_bare_slug(self, raw: str) -> None: + assert gg._normalize_model_name(raw) == "gemini-2.5-flash" + + def test_version_suffix_is_preserved(self) -> None: + """Rewriting suffixes would misreport which model actually ran.""" + assert gg._normalize_model_name("gemini-2.0-flash-001") == "gemini-2.0-flash-001" + + @pytest.mark.parametrize("raw", [None, "", 123]) + def test_missing_or_non_string_is_unknown(self, raw: Any) -> None: + assert gg._normalize_model_name(raw) == "unknown" + + +# ------------------------------- token accounting ------------------------------- # +class TestExtractUsage: + """OPEN-11891: prompt+candidates undercounts 54x on thinking models.""" + + def test_thinking_tokens_measured_case(self) -> None: + """The real 2.5-flash measurement: 8 / 6 / 743 -> 757, not 14.""" + prompt, completion, total, extras = gg._extract_usage(_usage(prompt=8, candidates=6, thoughts=743, total=757)) + assert prompt == 8 + assert completion == 749, "thinking tokens must be billed as output" + assert total == 757, "must use total_token_count, not prompt + candidates" + assert prompt + completion == total + assert extras["thoughtsTokens"] == 743 + assert extras["candidatesTokens"] == 6 + + def test_total_is_derived_when_absent(self) -> None: + prompt, completion, total, _ = gg._extract_usage(_usage(prompt=10, candidates=5, thoughts=2, tool_use=3)) + assert (prompt, completion, total) == (10, 7, 20) + + def test_no_thinking_tokens_behaves_conventionally(self) -> None: + prompt, completion, total, extras = gg._extract_usage(_usage(prompt=8, candidates=4, total=12)) + assert (prompt, completion, total) == (8, 4, 12) + assert "thoughtsTokens" not in extras + + def test_none_usage_is_zeroed(self) -> None: + assert gg._extract_usage(None) == (0, 0, 0, {}) + + def test_cached_and_tool_use_land_in_extras(self) -> None: + _, _, _, extras = gg._extract_usage(_usage(prompt=5, candidates=5, total=20, tool_use=4, cached=6)) + assert extras["toolUsePromptTokens"] == 4 + assert extras["cachedContentTokens"] == 6 + + +# ------------------------------- sync paths ------------------------------- # +class TestSyncGeneration: + def test_non_streaming_emits_span_with_slug_and_bare_model(self) -> None: + client = _make_client() + calls = _stub_sync(client, FakeResponse("hi there", _usage(prompt=8, candidates=6, thoughts=743, total=757))) + gg.trace_google_genai(client) + + with patch.object(gg, "add_to_trace") as mock_add: + response = client.models.generate_content(model="models/gemini-2.5-flash", contents="hi") + + assert response.text == "hi there" + mock_add.assert_called_once() + kwargs = mock_add.call_args.kwargs + assert kwargs["model"] == "gemini-2.5-flash", "must be normalized for cost lookup" + assert kwargs["tokens"] == 757 + assert kwargs["completion_tokens"] == 749 + assert kwargs["output"] == "hi there" + assert kwargs["inputs"] == {"prompt": [{"role": "user", "content": "hi"}]} + assert kwargs["metadata"]["requestedModel"] == "models/gemini-2.5-flash" + # inference_id must never reach the SDK + assert "inference_id" not in calls[0] + + def test_provider_is_the_cost_slug(self) -> None: + """provider must be a real slug; 'Google' silently prices at $0.""" + assert gg.PROVIDER == "gemini" + + client = _make_client() + _stub_sync(client, FakeResponse()) + gg.trace_google_genai(client) + + from openlayer.lib.tracing import tracer as _tracer + + with patch.object(_tracer, "add_chat_completion_step_to_trace") as mock_step: + client.models.generate_content(model="gemini-2.5-flash", contents="hi") + + assert mock_step.call_args.kwargs["provider"] == "gemini" + + def test_inference_id_is_forwarded_as_step_id(self) -> None: + client = _make_client() + _stub_sync(client, FakeResponse()) + gg.trace_google_genai(client) + + with patch.object(gg, "add_to_trace") as mock_add: + client.models.generate_content(model="gemini-2.5-flash", contents="hi", inference_id="abc-123") + + assert mock_add.call_args.kwargs["id"] == "abc-123" + + def test_streaming_passes_chunks_through_and_emits_once(self) -> None: + chunks = [ + FakeResponse("Hel", _usage(prompt=8, candidates=2, total=10)), + FakeResponse("lo!", _usage(prompt=8, candidates=5, total=13)), + ] + client = _make_client() + _stub_sync_stream(client, chunks) + gg.trace_google_genai(client) + + with patch.object(gg, "add_to_trace") as mock_add: + received = list(client.models.generate_content_stream(model="gemini-2.5-flash", contents="hi")) + + assert [c.text for c in received] == ["Hel", "lo!"] + mock_add.assert_called_once() + kwargs = mock_add.call_args.kwargs + assert kwargs["output"] == "Hello!" + # usage_metadata is cumulative per chunk: last wins, never summed + assert kwargs["tokens"] == 13, "summing chunks would give 23" + assert kwargs["metadata"]["timeToFirstToken"] is not None + + def test_streaming_span_emitted_when_consumer_abandons_iterator(self) -> None: + chunks = [FakeResponse("a", _usage(prompt=1, candidates=1, total=2)) for _ in range(5)] + client = _make_client() + _stub_sync_stream(client, chunks) + gg.trace_google_genai(client) + + with patch.object(gg, "add_to_trace") as mock_add: + stream = client.models.generate_content_stream(model="gemini-2.5-flash", contents="hi") + next(iter(stream)) + del stream + + mock_add.assert_called_once() + + +# ------------------------------- async paths ------------------------------- # +class TestAsyncGeneration: + def test_non_streaming_async(self) -> None: + client = _make_client() + _stub_async(client, FakeResponse("async hi", _usage(prompt=3, candidates=4, total=7))) + gg.trace_google_genai(client) + + with patch.object(gg, "add_to_trace") as mock_add: + response = asyncio.run(client.aio.models.generate_content(model="gemini-2.5-flash", contents="hi")) + + assert response.text == "async hi" + assert mock_add.call_args.kwargs["tokens"] == 7 + + def test_streaming_async_is_await_then_async_for(self) -> None: + """AsyncModels.generate_content_stream is a coroutine returning an + AsyncIterator — the wrapper must preserve that two-step shape.""" + chunks = [ + FakeResponse("x", _usage(prompt=2, candidates=1, total=3)), + FakeResponse("y", _usage(prompt=2, candidates=3, total=5)), + ] + client = _make_client() + _stub_async_stream(client, chunks) + gg.trace_google_genai(client) + + async def _drive() -> List[str]: + stream = await client.aio.models.generate_content_stream(model="gemini-2.5-flash", contents="hi") + return [chunk.text async for chunk in stream if chunk.text] + + with patch.object(gg, "add_to_trace") as mock_add: + texts = asyncio.run(_drive()) + + assert texts == ["x", "y"] + mock_add.assert_called_once() + assert mock_add.call_args.kwargs["output"] == "xy" + assert mock_add.call_args.kwargs["tokens"] == 5 + + +# ------------------------------- chats ------------------------------- # +class TestChatsCoverage: + def test_chat_send_message_is_traced_via_patched_models(self) -> None: + """client.chats wires to the same Models instance, so it traces for free. + Locked in so a refactor can't silently drop chat coverage.""" + client = _make_client() + + def _impl(**_kwargs: Any) -> Any: + return FakeResponse("chat reply", _usage(prompt=4, candidates=6, total=10)) + + client.models.generate_content = _impl + gg.trace_google_genai(client) + + with patch.object(gg, "add_to_trace") as mock_add: + chat = client.chats.create(model="gemini-2.5-flash") + chat.send_message("hello") + + mock_add.assert_called_once() + assert mock_add.call_args.kwargs["model"] == "gemini-2.5-flash" + + +# ------------------------------- Vertex vs AI Studio ------------------------------- # +class TestModeMetadata: + def test_ai_studio_mode(self) -> None: + client = _make_client(vertexai=False) + _stub_sync(client, FakeResponse()) + gg.trace_google_genai(client) + + with patch.object(gg, "add_to_trace") as mock_add: + client.models.generate_content(model="gemini-2.5-flash", contents="hi") + + assert mock_add.call_args.kwargs["metadata"]["llm_system"] == "google_ai_studio" + + def test_vertex_mode_records_system_and_project(self) -> None: + client = _make_client(vertexai=True) + _stub_sync(client, FakeResponse()) + gg.trace_google_genai(client) + + with patch.object(gg, "add_to_trace") as mock_add: + client.models.generate_content( + model="projects/test-project/locations/us-central1/publishers/google/models/gemini-2.5-flash", + contents="hi", + ) + + metadata = mock_add.call_args.kwargs["metadata"] + assert metadata["llm_system"] == "google_vertex" + assert metadata["gcp_project"] == "test-project" + assert metadata["gcp_location"] == "us-central1" + # The full Vertex resource path must still reduce to the bare slug + assert mock_add.call_args.kwargs["model"] == "gemini-2.5-flash" + + +# ------------------------------- ADK dedup ------------------------------- # +class TestAdkSuppression: + def test_no_span_when_adk_llm_step_is_active(self) -> None: + """ADK's api_client IS a genai.Client, so both integrations activate for + ADK users. Without suppression every ADK call double-counts.""" + client = _make_client() + _stub_sync(client, FakeResponse("adk")) + gg.trace_google_genai(client) + + fake_adk = SimpleNamespace(_current_llm_step=contextvars.ContextVar("test_llm_step", default=None)) + fake_adk._current_llm_step.set(object()) + + with patch.dict("sys.modules", {gg._ADK_TRACER_MODULE: fake_adk}): + with patch.object(gg, "add_to_trace") as mock_add: + response = client.models.generate_content(model="gemini-2.5-flash", contents="hi") + + assert response.text == "adk", "the underlying call must still run" + mock_add.assert_not_called() + + def test_span_emitted_when_adk_step_is_not_active(self) -> None: + client = _make_client() + _stub_sync(client, FakeResponse("no adk")) + gg.trace_google_genai(client) + + fake_adk = SimpleNamespace(_current_llm_step=contextvars.ContextVar("test_llm_step_idle", default=None)) + + with patch.dict("sys.modules", {gg._ADK_TRACER_MODULE: fake_adk}): + with patch.object(gg, "add_to_trace") as mock_add: + client.models.generate_content(model="gemini-2.5-flash", contents="hi") + + mock_add.assert_called_once() + + def test_absent_adk_module_is_not_treated_as_active(self) -> None: + import sys + + assert gg._ADK_TRACER_MODULE not in sys.modules or gg._adk_span_active() is False + + def test_inference_id_stripped_even_on_suppressed_path(self) -> None: + """The SDK rejects unknown kwargs, so the pop must precede the ADK check.""" + client = _make_client() + calls = _stub_sync(client, FakeResponse()) + gg.trace_google_genai(client) + + fake_adk = SimpleNamespace(_current_llm_step=contextvars.ContextVar("test_llm_step_pop", default=None)) + fake_adk._current_llm_step.set(object()) + + with patch.dict("sys.modules", {gg._ADK_TRACER_MODULE: fake_adk}): + client.models.generate_content(model="gemini-2.5-flash", contents="hi", inference_id="x") + + assert "inference_id" not in calls[0] + + +# ------------------------------- idempotency & dispatch ------------------------------- # +class TestIdempotency: + def test_double_patch_wraps_once(self) -> None: + client = _make_client() + _stub_sync(client, FakeResponse()) + + gg.trace_google_genai(client) + first = client.models.generate_content + gg.trace_google_genai(client) + + assert client.models.generate_content is first + + with patch.object(gg, "add_to_trace") as mock_add: + client.models.generate_content(model="gemini-2.5-flash", contents="hi") + + assert mock_add.call_count == 1, "a double wrap would emit two spans" + + def test_auto_instrument_traces_newly_constructed_clients(self, monkeypatch: pytest.MonkeyPatch) -> None: + """The stub is installed on the CLASS before the class-init patch, so the + per-instance wrapper wraps the stub rather than being clobbered by it.""" + from google.genai import models as genai_models + + def _impl(_self: Any, **_kwargs: Any) -> Any: + return FakeResponse("auto-traced", _usage(prompt=2, candidates=3, total=5)) + + monkeypatch.setattr(genai_models.Models, "generate_content", _impl) + + gg._patch_google_genai() + client = _make_client() + + with patch.object(gg, "add_to_trace") as mock_add: + response = client.models.generate_content(model="gemini-2.5-flash", contents="hi") + + assert response.text == "auto-traced" + mock_add.assert_called_once() + assert mock_add.call_args.kwargs["tokens"] == 5 + + def test_auto_instrument_then_explicit_call_does_not_rewrap(self) -> None: + gg._patch_google_genai() + client = _make_client() + + # Construction already traced this instance via the patched __init__. + assert getattr(client.models, "_openlayer_patched", False) is True + wrapped = client.models.generate_content + + assert gg.trace_google_genai(client) is client + assert client.models.generate_content is wrapped, "explicit call must not re-wrap" + + +class TestTraceGeminiDispatch: + def test_client_routes_to_the_unified_tracer(self) -> None: + from openlayer.lib import trace_gemini + + client = _make_client() + _stub_sync(client, FakeResponse()) + assert trace_gemini(client) is client + assert getattr(client, "_openlayer_patched", False) is True + + def test_trace_google_genai_public_alias(self) -> None: + from openlayer.lib import trace_google_genai + + client = _make_client() + _stub_sync(client, FakeResponse()) + assert trace_google_genai(client) is client + + def test_bad_type_names_both_accepted_clients(self) -> None: + from openlayer.lib import trace_gemini + + with pytest.raises(ValueError) as excinfo: + trace_gemini(object()) + + message = str(excinfo.value) + assert "google.genai.Client" in message + assert "GenerativeModel" in message + + +class TestAutoRegistry: + def test_google_genai_entry_registered_and_probe_resolves(self) -> None: + from openlayer.lib.integrations._auto import REGISTRY, _is_installed + + spec = next((s for s in REGISTRY if s.name == "google_genai"), None) + assert spec is not None, "auto_instrument() must cover the unified SDK" + assert spec.probe == "google.genai" + assert spec.unpatch is not None, "unpatch_all() must be symmetric" + assert _is_installed(spec.probe) is True + + def test_legacy_gemini_entry_is_kept(self) -> None: + """Both packages can be installed side by side.""" + from openlayer.lib.integrations._auto import REGISTRY + + assert any(s.name == "gemini" and s.probe == "google.generativeai" for s in REGISTRY) + + +# ------------------------------- input formatting ------------------------------- # +class TestInputFormatting: + def test_plain_string(self) -> None: + assert gg._format_input_messages("hi") == [{"role": "user", "content": "hi"}] + + def test_none_contents(self) -> None: + assert gg._format_input_messages(None) == [] + + def test_content_objects_with_parts(self) -> None: + contents = [ + SimpleNamespace(role="user", parts=[SimpleNamespace(text="first")]), + SimpleNamespace(role="model", parts=[SimpleNamespace(text="second")]), + ] + assert gg._format_input_messages(contents) == [ + {"role": "user", "content": "first"}, + {"role": "model", "content": "second"}, + ] + + def test_dict_message_form(self) -> None: + contents = [{"role": "user", "parts": [{"text": "from dict"}]}] + assert gg._format_input_messages(contents) == [{"role": "user", "content": "from dict"}] + + def test_system_instruction_becomes_leading_system_message(self) -> None: + messages = gg._format_input_messages("hi", {"system_instruction": "be terse"}) + assert messages[0] == {"role": "system", "content": "be terse"} + assert messages[1] == {"role": "user", "content": "hi"} + + +class TestModelParameters: + def test_reads_pydantic_config(self) -> None: + from google.genai import types + + config = types.GenerateContentConfig(temperature=0.25, max_output_tokens=64, top_p=0.9) + params = gg.get_model_parameters(config) + assert params["temperature"] == 0.25 + assert params["max_output_tokens"] == 64 + assert params["top_p"] == 0.9 + + def test_reads_dict_config(self) -> None: + params = gg.get_model_parameters({"temperature": 0.5, "top_k": 3}) + assert params["temperature"] == 0.5 + assert params["top_k"] == 3 + + def test_thinking_budget_is_surfaced(self) -> None: + from google.genai import types + + config = types.GenerateContentConfig(thinking_config=types.ThinkingConfig(thinking_budget=0)) + assert gg.get_model_parameters(config)["thinking_budget"] == 0 + + def test_none_config(self) -> None: + assert gg.get_model_parameters(None)["temperature"] is None