feat(gemini): trace google gemini Integration

.github/CODEOWNERS

@@ -127,13 +127,16 @@ 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                               @DataDog/ml-observability
+ddtrace/contrib/internal/google_generativeai            @DataDog/ml-observability
+ddtrace/contrib/google_generativeai                     @DataDog/ml-observability
 tests/llmobs                                            @DataDog/ml-observability
 tests/contrib/openai                                    @DataDog/ml-observability
 tests/contrib/langchain                                 @DataDog/ml-observability
 tests/contrib/botocore/test_bedrock.py                  @DataDog/ml-observability
 tests/contrib/botocore/test_bedrock_llmobs.py           @DataDog/ml-observability
 tests/contrib/botocore/bedrock_cassettes                @DataDog/ml-observability
 tests/contrib/anthropic                                 @DataDog/ml-observability
+tests/contrib/google_generativeai                       @DataDog/ml-observability
 .gitlab/tests/llmobs.yml                                @DataDog/ml-observability
 
 # Remote Config

.gitlab/tests/llmobs.yml

@@ -22,3 +22,7 @@ anthropic:
   variables:
     SUITE_NAME: "anthropic"
 
+google_generativeai:
+  extends: .test_base_riot_snapshot
+  variables:
+    SUITE_NAME: "google_generativeai"

.riot/requirements/1e15a25.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/1e15a25.in
+#
+annotated-types==0.7.0
+attrs==24.2.0
+cachetools==5.5.0
+certifi==2024.8.30
+charset-normalizer==3.3.2
+coverage[toml]==7.6.1
+google-ai-generativelanguage==0.6.9
+google-api-core[grpc]==2.19.2
+google-api-python-client==2.145.0
+google-auth==2.34.0
+google-auth-httplib2==0.2.0
+google-generativeai==0.8.0
+googleapis-common-protos==1.65.0
+grpcio==1.66.1
+grpcio-status==1.66.1
+httplib2==0.22.0
+hypothesis==6.45.0
+idna==3.8
+iniconfig==2.0.0
+mock==5.1.0
+opentracing==2.4.0
+packaging==24.1
+pillow==10.4.0
+pluggy==1.5.0
+proto-plus==1.24.0
+protobuf==5.28.0
+pyasn1==0.6.0
+pyasn1-modules==0.4.0
+pydantic==2.9.1
+pydantic-core==2.23.3
+pyparsing==3.1.4
+pytest==8.3.3
+pytest-asyncio==0.24.0
+pytest-cov==5.0.0
+pytest-mock==3.14.0
+requests==2.32.3
+rsa==4.9
+sortedcontainers==2.4.0
+tqdm==4.66.5
+typing-extensions==4.12.2
+uritemplate==4.1.1
+urllib3==2.2.2

.riot/requirements/1f54e6b.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/1f54e6b.in
+#
+annotated-types==0.7.0
+attrs==24.2.0
+cachetools==5.5.0
+certifi==2024.8.30
+charset-normalizer==3.3.2
+coverage[toml]==7.6.1
+exceptiongroup==1.2.2
+google-ai-generativelanguage==0.6.9
+google-api-core[grpc]==2.19.2
+google-api-python-client==2.145.0
+google-auth==2.34.0
+google-auth-httplib2==0.2.0
+google-generativeai==0.8.0
+googleapis-common-protos==1.65.0
+grpcio==1.66.1
+grpcio-status==1.66.1
+httplib2==0.22.0
+hypothesis==6.45.0
+idna==3.8
+iniconfig==2.0.0
+mock==5.1.0
+opentracing==2.4.0
+packaging==24.1
+pillow==10.4.0
+pluggy==1.5.0
+proto-plus==1.24.0
+protobuf==5.28.0
+pyasn1==0.6.0
+pyasn1-modules==0.4.0
+pydantic==2.9.1
+pydantic-core==2.23.3
+pyparsing==3.1.4
+pytest==8.3.3
+pytest-asyncio==0.24.0
+pytest-cov==5.0.0
+pytest-mock==3.14.0
+requests==2.32.3
+rsa==4.9
+sortedcontainers==2.4.0
+tomli==2.0.1
+tqdm==4.66.5
+typing-extensions==4.12.2
+uritemplate==4.1.1
+urllib3==2.2.2

.riot/requirements/e8247d6.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/e8247d6.in
+#
+annotated-types==0.7.0
+attrs==24.2.0
+cachetools==5.5.0
+certifi==2024.8.30
+charset-normalizer==3.3.2
+coverage[toml]==7.6.1
+google-ai-generativelanguage==0.6.9
+google-api-core[grpc]==2.19.2
+google-api-python-client==2.145.0
+google-auth==2.34.0
+google-auth-httplib2==0.2.0
+google-generativeai==0.8.0
+googleapis-common-protos==1.65.0
+grpcio==1.66.1
+grpcio-status==1.66.1
+httplib2==0.22.0
+hypothesis==6.45.0
+idna==3.8
+iniconfig==2.0.0
+mock==5.1.0
+opentracing==2.4.0
+packaging==24.1
+pillow==10.4.0
+pluggy==1.5.0
+proto-plus==1.24.0
+protobuf==5.28.0
+pyasn1==0.6.0
+pyasn1-modules==0.4.0
+pydantic==2.9.1
+pydantic-core==2.23.3
+pyparsing==3.1.4
+pytest==8.3.3
+pytest-asyncio==0.24.0
+pytest-cov==5.0.0
+pytest-mock==3.14.0
+requests==2.32.3
+rsa==4.9
+sortedcontainers==2.4.0
+tqdm==4.66.5
+typing-extensions==4.12.2
+uritemplate==4.1.1
+urllib3==2.2.2

.riot/requirements/ebe4ea5.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/ebe4ea5.in
+#
+annotated-types==0.7.0
+attrs==24.2.0
+cachetools==5.5.0
+certifi==2024.8.30
+charset-normalizer==3.3.2
+coverage[toml]==7.6.1
+exceptiongroup==1.2.2
+google-ai-generativelanguage==0.6.9
+google-api-core[grpc]==2.19.2
+google-api-python-client==2.145.0
+google-auth==2.34.0
+google-auth-httplib2==0.2.0
+google-generativeai==0.8.0
+googleapis-common-protos==1.65.0
+grpcio==1.66.1
+grpcio-status==1.66.1
+httplib2==0.22.0
+hypothesis==6.45.0
+idna==3.8
+iniconfig==2.0.0
+mock==5.1.0
+opentracing==2.4.0
+packaging==24.1
+pillow==10.4.0
+pluggy==1.5.0
+proto-plus==1.24.0
+protobuf==5.28.0
+pyasn1==0.6.0
+pyasn1-modules==0.4.0
+pydantic==2.9.1
+pydantic-core==2.23.3
+pyparsing==3.1.4
+pytest==8.3.3
+pytest-asyncio==0.24.0
+pytest-cov==5.0.0
+pytest-mock==3.14.0
+requests==2.32.3
+rsa==4.9
+sortedcontainers==2.4.0
+tomli==2.0.1
+tqdm==4.66.5
+typing-extensions==4.12.2
+uritemplate==4.1.1
+urllib3==2.2.2

ddtrace/_monkey.py

@@ -39,6 +39,7 @@
     "elasticsearch": True,
     "algoliasearch": True,
     "futures": True,
+    "google_generativeai": True,
     "gevent": True,
     "graphql": True,
     "grpc": True,
@@ -139,6 +140,7 @@
     "aws_lambda": ("datadog_lambda",),
     "httplib": ("http.client",),
     "kafka": ("confluent_kafka",),
+    "google_generativeai": ("google.generativeai",),
 }
 
 

ddtrace/contrib/google_generativeai/__init__.py

@@ -0,0 +1,91 @@
+"""
+The Gemini integration instruments the Google Gemini Python API to traces for requests made to Google models.
+
+All traces submitted from the Gemini integration are tagged by:
+
+- ``service``, ``env``, ``version``: see the `Unified Service Tagging docs <https://docs.datadoghq.com/getting_started/tagging/unified_service_tagging>`_.
+- ``google_generativeai.request.model``: Google model used in the request.
+- ``google_generativeai.request.api_key``: Google Gemini API key used to make the request (obfuscated to match the Google AI Studio UI representation ``...XXXX`` where ``XXXX`` is the last 4 digits of the key).
+
+
+(beta) Prompt and Completion Sampling
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+Prompt texts and completion content for the ``generateContent`` endpoint are collected in span tags with a default sampling rate of ``1.0``.
+These tags will have truncation applied if the text exceeds the configured character limit.
+
+
+Enabling
+~~~~~~~~
+
+The Gemini 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 Gemini integration::
+
+    from ddtrace import config, patch
+
+    patch(google_generativeai=True)
+
+
+Global Configuration
+~~~~~~~~~~~~~~~~~~~~
+
+.. py:data:: ddtrace.config.google_generativeai["service"]
+
+   The service name reported by default for Gemini requests.
+
+   Alternatively, you can set this option with the ``DD_SERVICE`` or ``DD_GOOGLE_GENERATIVEAI_SERVICE`` environment
+   variables.
+
+   Default: ``DD_SERVICE``
+
+
+.. py:data:: (beta) ddtrace.config.google_generativeai["span_char_limit"]
+
+   Configure the maximum number of characters for the following data within span tags:
+
+   - Text inputs and completions
+
+   Text exceeding the maximum number of characters is truncated to the character limit
+   and has ``...`` appended to the end.
+
+   Alternatively, you can set this option with the ``DD_GOOGLE_GENERATIVEAI_SPAN_CHAR_LIMIT`` environment
+   variable.
+
+   Default: ``128``
+
+
+.. py:data:: (beta) ddtrace.config.google_generativeai["span_prompt_completion_sample_rate"]
+
+   Configure the sample rate for the collection of prompts and completions as span tags.
+
+   Alternatively, you can set this option with the ``DD_GOOGLE_GENERATIVEAI_SPAN_PROMPT_COMPLETION_SAMPLE_RATE`` environment
+   variable.
+
+   Default: ``1.0``
+
+
+Instance Configuration
+~~~~~~~~~~~~~~~~~~~~~~
+
+To configure the Gemini integration on a per-instance basis use the
+``Pin`` API::
+
+    import google.generativeai as genai
+    from ddtrace import Pin, config
+
+    Pin.override(genai, service="my-gemini-service")
+"""  # noqa: E501
+from ...internal.utils.importlib import require_modules
+
+
+required_modules = ["google.generativeai"]
+
+with require_modules(required_modules) as missing_modules:
+    if not missing_modules:
+        from ..internal.google_generativeai.patch import get_version
+        from ..internal.google_generativeai.patch import patch
+        from ..internal.google_generativeai.patch import unpatch
+
+        __all__ = ["patch", "unpatch", "get_version"]

ddtrace/contrib/internal/anthropic/_streaming.py

@@ -9,7 +9,7 @@
 
 from ddtrace.contrib.internal.anthropic.utils import tag_tool_use_output_on_span
 from ddtrace.internal.logger import get_logger
-from ddtrace.llmobs._integrations.anthropic import _get_attr
+from ddtrace.llmobs._utils import _get_attr
 
 
 log = get_logger(__name__)

ddtrace/contrib/internal/anthropic/patch.py

@@ -17,7 +17,7 @@
 from ddtrace.internal.logger import get_logger
 from ddtrace.internal.utils import get_argument_value
 from ddtrace.llmobs._integrations import AnthropicIntegration
-from ddtrace.llmobs._integrations.anthropic import _get_attr
+from ddtrace.llmobs._utils import _get_attr
 from ddtrace.pin import Pin
 
 

ddtrace/contrib/internal/anthropic/utils.py

@@ -3,7 +3,7 @@
 from typing import Optional
 
 from ddtrace.internal.logger import get_logger
-from ddtrace.llmobs._integrations.anthropic import _get_attr
+from ddtrace.llmobs._utils import _get_attr
 
 
 log = get_logger(__name__)