feat(ai): Add python-genai integration

.github/workflows/test-integrations-ai.yml

@@ -82,6 +82,10 @@ jobs:
         run: |
           set -x # print commands that are executed
           ./scripts/runtox.sh "py${{ matrix.python-version }}-langgraph"
+      - name: Test google_genai
+        run: |
+          set -x # print commands that are executed
+          ./scripts/runtox.sh "py${{ matrix.python-version }}-google_genai"
       - name: Test openai_agents
         run: |
           set -x # print commands that are executed

pyproject.toml

@@ -118,6 +118,10 @@ ignore_missing_imports = true
 module = "langgraph.*"
 ignore_missing_imports = true
 
+[[tool.mypy.overrides]]
+module = "google.genai.*"
+ignore_missing_imports = true
+
 [[tool.mypy.overrides]]
 module = "executing.*"
 ignore_missing_imports = true

scripts/populate_tox/config.py

@@ -142,6 +142,13 @@
         "package": "gql[all]",
         "num_versions": 2,
     },
+    "google_genai": {
+        "package": "google-genai",
+        "deps": {
+            "*": ["pytest-asyncio"],
+        },
+        "python": ">=3.9",
+    },
     "graphene": {
         "package": "graphene",
         "deps": {

scripts/split_tox_gh_actions/split_tox_gh_actions.py

@@ -78,6 +78,7 @@
         "openai-base",
         "openai-notiktoken",
         "langgraph",
+        "google_genai",
         "openai_agents",
         "huggingface_hub",
     ],

sentry_sdk/integrations/__init__.py

@@ -140,6 +140,7 @@ def iter_default_integrations(with_auto_enabling_integrations):
     "flask": (1, 1, 4),
     "gql": (3, 4, 1),
     "graphene": (3, 3),
+    "google_genai": (1, 29, 0),  # google-genai
     "grpc": (1, 32, 0),  # grpcio
     "httpx": (0, 16, 0),
     "huggingface_hub": (0, 24, 7),

sentry_sdk/integrations/google_genai/__init__.py

@@ -0,0 +1,298 @@
+from functools import wraps
+from typing import (
+    Any,
+    AsyncIterator,
+    Callable,
+    Iterator,
+    List,
+)
+
+import sentry_sdk
+from sentry_sdk.ai.utils import get_start_span_function
+from sentry_sdk.integrations import DidNotEnable, Integration
+from sentry_sdk.consts import OP, SPANDATA
+from sentry_sdk.tracing import SPANSTATUS
+
+
+try:
+    from google.genai.models import Models, AsyncModels
+except ImportError:
+    raise DidNotEnable("google-genai not installed")
+
+
+from .consts import IDENTIFIER, ORIGIN, GEN_AI_SYSTEM
+from .utils import (
+    set_span_data_for_request,
+    set_span_data_for_response,
+    _capture_exception,
+    prepare_generate_content_args,
+)
+from .streaming import (
+    set_span_data_for_streaming_response,
+    accumulate_streaming_response,
+)
+
+
+class GoogleGenAIIntegration(Integration):
+    identifier = IDENTIFIER
+    origin = ORIGIN
+
+    def __init__(self, include_prompts=True):
+        # type: (GoogleGenAIIntegration, bool) -> None
+        self.include_prompts = include_prompts
+
+    @staticmethod
+    def setup_once():
+        # type: () -> None
+        # Patch sync methods
+        Models.generate_content = _wrap_generate_content(Models.generate_content)
+        Models.generate_content_stream = _wrap_generate_content_stream(
+            Models.generate_content_stream
+        )
+
+        # Patch async methods
+        AsyncModels.generate_content = _wrap_async_generate_content(
+            AsyncModels.generate_content
+        )
+        AsyncModels.generate_content_stream = _wrap_async_generate_content_stream(
+            AsyncModels.generate_content_stream
+        )
+
+
+def _wrap_generate_content_stream(f):
+    # type: (Callable[..., Any]) -> Callable[..., Any]
+    @wraps(f)
+    def new_generate_content_stream(self, *args, **kwargs):
+        # type: (Any, Any, Any) -> Any
+        integration = sentry_sdk.get_client().get_integration(GoogleGenAIIntegration)
+        if integration is None:
+            return f(self, *args, **kwargs)
+
+        _model, contents, model_name = prepare_generate_content_args(args, kwargs)
+
+        span = get_start_span_function()(
+            op=OP.GEN_AI_INVOKE_AGENT,
+            name="invoke_agent",
+            origin=ORIGIN,
+        )
+        span.__enter__()
+        span.set_data(SPANDATA.GEN_AI_AGENT_NAME, model_name)
+        span.set_data(SPANDATA.GEN_AI_OPERATION_NAME, "invoke_agent")
+        set_span_data_for_request(span, integration, model_name, contents, kwargs)
+        span.set_data(SPANDATA.GEN_AI_RESPONSE_STREAMING, True)
+
+        chat_span = sentry_sdk.start_span(
+            op=OP.GEN_AI_CHAT,
+            name=f"chat {model_name}",
+            origin=ORIGIN,
+        )
+        chat_span.__enter__()
+        chat_span.set_data(SPANDATA.GEN_AI_OPERATION_NAME, "chat")
+        chat_span.set_data(SPANDATA.GEN_AI_SYSTEM, GEN_AI_SYSTEM)
+        chat_span.set_data(SPANDATA.GEN_AI_REQUEST_MODEL, model_name)
+        set_span_data_for_request(chat_span, integration, model_name, contents, kwargs)
+        chat_span.set_data(SPANDATA.GEN_AI_RESPONSE_STREAMING, True)
+
+        try:
+            stream = f(self, *args, **kwargs)
+
+            # Create wrapper iterator to accumulate responses
+            def new_iterator():
+                # type: () -> Iterator[Any]
+                chunks = []  # type: List[Any]
+                try:
+                    for chunk in stream:
+                        chunks.append(chunk)
+                        yield chunk
+                except Exception as exc:
+                    _capture_exception(exc)
+                    chat_span.set_status(SPANSTATUS.ERROR)
+                    raise
+                finally:
+                    # Accumulate all chunks and set final response data on spans
+                    if chunks:
+                        accumulated_response = accumulate_streaming_response(chunks)
+                        set_span_data_for_streaming_response(
+                            chat_span, integration, accumulated_response
+                        )
+                        set_span_data_for_streaming_response(
+                            span, integration, accumulated_response
+                        )
+                    chat_span.__exit__(None, None, None)
+                    span.__exit__(None, None, None)
+
+            return new_iterator()
+
+        except Exception as exc:
+            _capture_exception(exc)
+            chat_span.__exit__(None, None, None)
+            span.__exit__(None, None, None)
+            raise
+
+    return new_generate_content_stream
+
+
+def _wrap_async_generate_content_stream(f):
+    # type: (Callable[..., Any]) -> Callable[..., Any]
+    @wraps(f)
+    async def new_async_generate_content_stream(self, *args, **kwargs):
+        # type: (Any, Any, Any) -> Any
+        integration = sentry_sdk.get_client().get_integration(GoogleGenAIIntegration)
+        if integration is None:
+            return await f(self, *args, **kwargs)
+
+        _model, contents, model_name = prepare_generate_content_args(args, kwargs)
+
+        span = get_start_span_function()(
+            op=OP.GEN_AI_INVOKE_AGENT,
+            name="invoke_agent",
+            origin=ORIGIN,
+        )
+        span.__enter__()
+        span.set_data(SPANDATA.GEN_AI_AGENT_NAME, model_name)
+        span.set_data(SPANDATA.GEN_AI_OPERATION_NAME, "invoke_agent")
+        set_span_data_for_request(span, integration, model_name, contents, kwargs)
+        span.set_data(SPANDATA.GEN_AI_RESPONSE_STREAMING, True)
+
+        chat_span = sentry_sdk.start_span(
+            op=OP.GEN_AI_CHAT,
+            name=f"chat {model_name}",
+            origin=ORIGIN,
+        )
+        chat_span.__enter__()
+        chat_span.set_data(SPANDATA.GEN_AI_OPERATION_NAME, "chat")
+        chat_span.set_data(SPANDATA.GEN_AI_SYSTEM, GEN_AI_SYSTEM)
+        chat_span.set_data(SPANDATA.GEN_AI_REQUEST_MODEL, model_name)
+        set_span_data_for_request(chat_span, integration, model_name, contents, kwargs)
+        chat_span.set_data(SPANDATA.GEN_AI_RESPONSE_STREAMING, True)
+
+        try:
+            stream = await f(self, *args, **kwargs)
+
+            # Create wrapper async iterator to accumulate responses
+            async def new_async_iterator():
+                # type: () -> AsyncIterator[Any]
+                chunks = []  # type: List[Any]
+                try:
+                    async for chunk in stream:
+                        chunks.append(chunk)
+                        yield chunk
+                except Exception as exc:
+                    _capture_exception(exc)
+                    chat_span.set_status(SPANSTATUS.ERROR)
+                    raise
+                finally:
+                    # Accumulate all chunks and set final response data on spans
+                    if chunks:
+                        accumulated_response = accumulate_streaming_response(chunks)
+                        set_span_data_for_streaming_response(
+                            chat_span, integration, accumulated_response
+                        )
+                        set_span_data_for_streaming_response(
+                            span, integration, accumulated_response
+                        )
+                    chat_span.__exit__(None, None, None)
+                    span.__exit__(None, None, None)
+
+            return new_async_iterator()
+
+        except Exception as exc:
+            _capture_exception(exc)
+            chat_span.__exit__(None, None, None)
+            span.__exit__(None, None, None)
+            raise
+
+    return new_async_generate_content_stream
+
+
+def _wrap_generate_content(f):
+    # type: (Callable[..., Any]) -> Callable[..., Any]
+    @wraps(f)
+    def new_generate_content(self, *args, **kwargs):
+        # type: (Any, Any, Any) -> Any
+        integration = sentry_sdk.get_client().get_integration(GoogleGenAIIntegration)
+        if integration is None:
+            return f(self, *args, **kwargs)
+
+        model, contents, model_name = prepare_generate_content_args(args, kwargs)
+
+        with get_start_span_function()(
+            op=OP.GEN_AI_INVOKE_AGENT,
+            name="invoke_agent",
+            origin=ORIGIN,
+        ) as span:
+            span.set_data(SPANDATA.GEN_AI_AGENT_NAME, model_name)
+            span.set_data(SPANDATA.GEN_AI_OPERATION_NAME, "invoke_agent")
+            set_span_data_for_request(span, integration, model_name, contents, kwargs)
+
+            with sentry_sdk.start_span(
+                op=OP.GEN_AI_CHAT,
+                name=f"chat {model_name}",
+                origin=ORIGIN,
+            ) as chat_span:
+                chat_span.set_data(SPANDATA.GEN_AI_OPERATION_NAME, "chat")
+                chat_span.set_data(SPANDATA.GEN_AI_SYSTEM, GEN_AI_SYSTEM)
+                chat_span.set_data(SPANDATA.GEN_AI_REQUEST_MODEL, model_name)
+                set_span_data_for_request(
+                    chat_span, integration, model_name, contents, kwargs
+                )
+
+                try:
+                    response = f(self, *args, **kwargs)
+                except Exception as exc:
+                    _capture_exception(exc)
+                    chat_span.set_status(SPANSTATUS.ERROR)
+                    raise
+
+                set_span_data_for_response(chat_span, integration, response)
+                set_span_data_for_response(span, integration, response)
+
+                return response
+
+    return new_generate_content
+
+
+def _wrap_async_generate_content(f):
+    # type: (Callable[..., Any]) -> Callable[..., Any]
+    @wraps(f)
+    async def new_async_generate_content(self, *args, **kwargs):
+        # type: (Any, Any, Any) -> Any
+        integration = sentry_sdk.get_client().get_integration(GoogleGenAIIntegration)
+        if integration is None:
+            return await f(self, *args, **kwargs)
+
+        model, contents, model_name = prepare_generate_content_args(args, kwargs)
+
+        with get_start_span_function()(
+            op=OP.GEN_AI_INVOKE_AGENT,
+            name="invoke_agent",
+            origin=ORIGIN,
+        ) as span:
+            span.set_data(SPANDATA.GEN_AI_AGENT_NAME, model_name)
+            span.set_data(SPANDATA.GEN_AI_OPERATION_NAME, "invoke_agent")
+            set_span_data_for_request(span, integration, model_name, contents, kwargs)
+
+            with sentry_sdk.start_span(
+                op=OP.GEN_AI_CHAT,
+                name=f"chat {model_name}",
+                origin=ORIGIN,
+            ) as chat_span:
+                chat_span.set_data(SPANDATA.GEN_AI_OPERATION_NAME, "chat")
+                chat_span.set_data(SPANDATA.GEN_AI_SYSTEM, GEN_AI_SYSTEM)
+                chat_span.set_data(SPANDATA.GEN_AI_REQUEST_MODEL, model_name)
+                set_span_data_for_request(
+                    chat_span, integration, model_name, contents, kwargs
+                )
+                try:
+                    response = await f(self, *args, **kwargs)
+                except Exception as exc:
+                    _capture_exception(exc)
+                    chat_span.set_status(SPANSTATUS.ERROR)
+                    raise
+
+                set_span_data_for_response(chat_span, integration, response)
+                set_span_data_for_response(span, integration, response)
+
+                return response
+
+    return new_async_generate_content

sentry_sdk/integrations/google_genai/consts.py

@@ -0,0 +1,16 @@
+GEN_AI_SYSTEM = "gcp.gemini"
+
+# Mapping of tool attributes to their descriptions
+# These are all tools that are available in the Google GenAI API
+TOOL_ATTRIBUTES_MAP = {
+    "google_search_retrieval": "Google Search retrieval tool",
+    "google_search": "Google Search tool",
+    "retrieval": "Retrieval tool",
+    "enterprise_web_search": "Enterprise web search tool",
+    "google_maps": "Google Maps tool",
+    "code_execution": "Code execution tool",
+    "computer_use": "Computer use tool",
+}
+
+IDENTIFIER = "google_genai"
+ORIGIN = f"auto.ai.{IDENTIFIER}"