feat(google_genai): [MLOB-2932] add apm tracing for google-genai

.github/CODEOWNERS

@@ -147,6 +147,8 @@ ddtrace/contrib/internal/botocore/services/bedrock.py   @DataDog/ml-observabilit
 ddtrace/contrib/botocore/services/bedrock.py            @DataDog/ml-observability
 ddtrace/contrib/internal/anthropic                      @DataDog/ml-observability
 ddtrace/contrib/_anthropic.py                           @DataDog/ml-observability
+ddtrace/contrib/internal/google_genai                   @DataDog/ml-observability
+ddtrace/contrib/_google_genai.py                        @DataDog/ml-observability
 ddtrace/contrib/internal/google_generativeai            @DataDog/ml-observability
 ddtrace/contrib/_google_generativeai.py                 @DataDog/ml-observability
 ddtrace/contrib/internal/vertexai                       @DataDog/ml-observability
@@ -167,6 +169,7 @@ tests/contrib/botocore/test_bedrock_llmobs.py           @DataDog/ml-observabilit
 tests/contrib/botocore/bedrock_utils.py                 @DataDog/ml-observability
 tests/contrib/botocore/bedrock_cassettes                @DataDog/ml-observability
 tests/contrib/anthropic                                 @DataDog/ml-observability
+tests/contrib/google_genai                              @DataDog/ml-observability
 tests/contrib/google_generativeai                       @DataDog/ml-observability
 tests/contrib/vertexai                                  @DataDog/ml-observability
 tests/contrib/langgraph                                 @DataDog/ml-observability

.riot/requirements/1de4a65.txt

@@ -0,0 +1,48 @@
+#
+# This file is autogenerated by pip-compile with Python 3.11
+# by the following command:
+#
+#    pip-compile --allow-unsafe --no-annotate .riot/requirements/1de4a65.in
+#
+annotated-types==0.7.0
+anyio==4.9.0
+attrs==25.3.0
+cachetools==5.5.2
+certifi==2025.4.26
+charset-normalizer==3.4.2
+coverage[toml]==7.9.0
+google-auth==2.40.3
+google-genai==1.20.0
+h11==0.16.0
+httpcore==1.0.9
+httpx==0.28.1
+hypothesis==6.45.0
+idna==3.10
+iniconfig==2.1.0
+mock==5.2.0
+multidict==6.4.4
+opentracing==2.4.0
+packaging==25.0
+pluggy==1.6.0
+propcache==0.3.2
+pyasn1==0.6.1
+pyasn1-modules==0.4.2
+pydantic==2.11.5
+pydantic-core==2.33.2
+pygments==2.19.1
+pytest==8.4.0
+pytest-asyncio==1.0.0
+pytest-cov==6.2.1
+pytest-mock==3.14.1
+pyyaml==6.0.2
+requests==2.32.4
+rsa==4.9.1
+sniffio==1.3.1
+sortedcontainers==2.4.0
+typing-extensions==4.14.0
+typing-inspection==0.4.1
+urllib3==2.4.0
+vcrpy==7.0.0
+websockets==15.0.1
+wrapt==1.17.2
+yarl==1.20.1

.riot/requirements/7d83e7d.txt

@@ -0,0 +1,50 @@
+#
+# This file is autogenerated by pip-compile with Python 3.10
+# by the following command:
+#
+#    pip-compile --allow-unsafe --no-annotate .riot/requirements/7d83e7d.in
+#
+annotated-types==0.7.0
+anyio==4.9.0
+attrs==25.3.0
+cachetools==5.5.2
+certifi==2025.4.26
+charset-normalizer==3.4.2
+coverage[toml]==7.9.0
+exceptiongroup==1.3.0
+google-auth==2.40.3
+google-genai==1.20.0
+h11==0.16.0
+httpcore==1.0.9
+httpx==0.28.1
+hypothesis==6.45.0
+idna==3.10
+iniconfig==2.1.0
+mock==5.2.0
+multidict==6.4.4
+opentracing==2.4.0
+packaging==25.0
+pluggy==1.6.0
+propcache==0.3.2
+pyasn1==0.6.1
+pyasn1-modules==0.4.2
+pydantic==2.11.5
+pydantic-core==2.33.2
+pygments==2.19.1
+pytest==8.4.0
+pytest-asyncio==1.0.0
+pytest-cov==6.2.1
+pytest-mock==3.14.1
+pyyaml==6.0.2
+requests==2.32.4
+rsa==4.9.1
+sniffio==1.3.1
+sortedcontainers==2.4.0
+tomli==2.2.1
+typing-extensions==4.14.0
+typing-inspection==0.4.1
+urllib3==2.4.0
+vcrpy==7.0.0
+websockets==15.0.1
+wrapt==1.17.2
+yarl==1.20.1

.riot/requirements/ce785c0.txt

@@ -0,0 +1,48 @@
+#
+# This file is autogenerated by pip-compile with Python 3.12
+# by the following command:
+#
+#    pip-compile --allow-unsafe --no-annotate .riot/requirements/ce785c0.in
+#
+annotated-types==0.7.0
+anyio==4.9.0
+attrs==25.3.0
+cachetools==5.5.2
+certifi==2025.4.26
+charset-normalizer==3.4.2
+coverage[toml]==7.9.0
+google-auth==2.40.3
+google-genai==1.20.0
+h11==0.16.0
+httpcore==1.0.9
+httpx==0.28.1
+hypothesis==6.45.0
+idna==3.10
+iniconfig==2.1.0
+mock==5.2.0
+multidict==6.4.4
+opentracing==2.4.0
+packaging==25.0
+pluggy==1.6.0
+propcache==0.3.2
+pyasn1==0.6.1
+pyasn1-modules==0.4.2
+pydantic==2.11.5
+pydantic-core==2.33.2
+pygments==2.19.1
+pytest==8.4.0
+pytest-asyncio==1.0.0
+pytest-cov==6.2.1
+pytest-mock==3.14.1
+pyyaml==6.0.2
+requests==2.32.4
+rsa==4.9.1
+sniffio==1.3.1
+sortedcontainers==2.4.0
+typing-extensions==4.14.0
+typing-inspection==0.4.1
+urllib3==2.4.0
+vcrpy==7.0.0
+websockets==15.0.1
+wrapt==1.17.2
+yarl==1.20.1

.riot/requirements/f5e518d.txt

@@ -0,0 +1,50 @@
+#
+# This file is autogenerated by pip-compile with Python 3.9
+# by the following command:
+#
+#    pip-compile --allow-unsafe --no-annotate .riot/requirements/f5e518d.in
+#
+annotated-types==0.7.0
+anyio==4.9.0
+attrs==25.3.0
+cachetools==5.5.2
+certifi==2025.4.26
+charset-normalizer==3.4.2
+coverage[toml]==7.9.0
+exceptiongroup==1.3.0
+google-auth==2.40.3
+google-genai==1.20.0
+h11==0.16.0
+httpcore==1.0.9
+httpx==0.28.1
+hypothesis==6.45.0
+idna==3.10
+iniconfig==2.1.0
+mock==5.2.0
+multidict==6.4.4
+opentracing==2.4.0
+packaging==25.0
+pluggy==1.6.0
+propcache==0.3.2
+pyasn1==0.6.1
+pyasn1-modules==0.4.2
+pydantic==2.11.5
+pydantic-core==2.33.2
+pygments==2.19.1
+pytest==8.4.0
+pytest-asyncio==1.0.0
+pytest-cov==6.2.1
+pytest-mock==3.14.1
+pyyaml==6.0.2
+requests==2.32.4
+rsa==4.9.1
+sniffio==1.3.1
+sortedcontainers==2.4.0
+tomli==2.2.1
+typing-extensions==4.14.0
+typing-inspection==0.4.1
+urllib3==1.26.20
+vcrpy==7.0.0
+websockets==15.0.1
+wrapt==1.17.2
+yarl==1.20.1

ddtrace/_monkey.py

@@ -50,6 +50,7 @@
     "futures": True,
     "freezegun": True,
     "google_generativeai": True,
+    "google_genai": True,
     "gevent": True,
     "graphql": True,
     "grpc": True,
@@ -160,6 +161,7 @@
     "httplib": ("http.client",),
     "kafka": ("confluent_kafka",),
     "google_generativeai": ("google.generativeai",),
+    "google_genai": ("google.genai",),
     "langgraph": (
         "langgraph",
         "langgraph.graph",

ddtrace/contrib/_google_genai.py

@@ -0,0 +1,56 @@
+"""
+The Google GenAI integration instruments the Google GenAI Python SDK to trace for requests made to Google models.
+
+All traces submitted from the Google GenAI integration are tagged by:
+
+- ``service``, ``env``, ``version``: see the `Unified Service Tagging docs <https://docs.datadoghq.com/getting_started/tagging/unified_service_tagging>`_.
+- model used in the request.
+- provider used in the request.
+
+(beta) Prompt and Completion Sampling
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+Prompt texts and completion content are collected in span tags with a default sampling rate of ``1.0``
+for the following methods:
+
+- ``generate_content/generate_content_stream`` of the ``google.genai.models.Models`` class
+
+These tags will have truncation applied if the text exceeds the configured character limit.
+
+Enabling
+~~~~~~~~
+
+The Google GenAI integration is enabled automatically when you use
+:ref:`ddtrace-run<ddtracerun>` or :ref:`import ddtrace.auto<ddtraceauto>`.
+
+Alternatively, use :func:`patch() <ddtrace.patch>` to manually enable the Google GenAI integration::
+
+    from ddtrace import config, patch
+
+    patch(google_genai=True)
+
+Global Configuration
+~~~~~~~~~~~~~~~~~~~~
+
+.. py:data:: ddtrace.config.google_genai["service"]
+
+   The service name reported by default for Google GenAI requests.
+
+   Alternatively, you can set this option with the ``DD_SERVICE`` or ``DD_VERTEXAI_SERVICE`` environment
+   variables.
+
+   Default: ``DD_SERVICE``
+
+
+Instance Configuration
+~~~~~~~~~~~~~~~~~~~~~~
+
+To configure the Google GenAI integration on a per-instance basis use the
+``Pin`` API::
+
+    from google import genai
+    from ddtrace import config
+    from ddtrace.trace import Pin
+
+    Pin.override(genai, service="my-google-genai-service")
+"""

ddtrace/contrib/integration_registry/registry.yaml

@@ -383,6 +383,16 @@ integrations:
       min: 20.12.1
       max: 24.11.1
 
+- integration_name: google_genai
+  is_external_package: true
+  is_tested: true
+  dependency_names:
+  - google-genai
+  tested_versions_by_dependency:
+    google-genai:
+      min: 1.20.0
+      max: 1.20.0
+
 - integration_name: google_generativeai
   is_external_package: true
   is_tested: true

ddtrace/contrib/internal/google_genai/_utils.py

@@ -0,0 +1,50 @@
+import sys
+
+
+# https://cloud.google.com/vertex-ai/generative-ai/docs/model-garden/quickstart
+# for vertex, it seems like the best way to associate provider name with each call is based on the model name prefix
+MODEL_PREFIX_TO_PROVIDER = {
+    "gemini": "google",
+    "imagen": "google",
+    "veo": "google",
+    "jamba": "ai21labs",
+    "claude": "anthropic",
+    "llama": "meta",
+    "mistral": "mistral",
+}
+
+
+def extract_provider_and_model_name_genai(kwargs):
+    model_name = kwargs.get("model", "").split("/")[-1]
+    provider_name = "other"
+    for prefix in MODEL_PREFIX_TO_PROVIDER.keys():
+        if model_name.startswith(prefix):
+            provider_name = MODEL_PREFIX_TO_PROVIDER[prefix]
+    return provider_name, model_name if len(model_name) > 0 else "unknown"
+
+
+class BaseTracedGoogleGenAIStreamResponse:
+    def __init__(self, generation_response, integration, span, args, kwargs):
+        self._generation_response = generation_response
+        self._integration = integration
+        self._dd_span = span
+        self._args = args
+        self._kwargs = kwargs
+        self._chunks = []
+
+
+class TracedGoogleGenAIStreamResponse(BaseTracedGoogleGenAIStreamResponse):
+    # generate_content_stream returns Iterator[GenerateContentResponse]
+    def __iter__(self):
+        return self
+
+    def __next__(self):
+        try:
+            chunk = self._generation_response.__next__()
+            self._chunks.append(chunk)
+            return chunk
+        except StopIteration:
+            raise
+        except Exception:
+            self._dd_span.set_exc_info(*sys.exc_info())
+            raise