feat(integrations): add support for embed_content methods in GoogleGenAI integration

sentry_sdk/integrations/google_genai/__init__.py

@@ -26,6 +26,9 @@
     set_span_data_for_response,
     _capture_exception,
     prepare_generate_content_args,
+    prepare_embed_content_args,
+    set_span_data_for_embed_request,
+    set_span_data_for_embed_response,
 )
 from .streaming import (
     set_span_data_for_streaming_response,
@@ -49,6 +52,7 @@ def setup_once():
         Models.generate_content_stream = _wrap_generate_content_stream(
             Models.generate_content_stream
         )
+        Models.embed_content = _wrap_embed_content(Models.embed_content)
 
         # Patch async methods
         AsyncModels.generate_content = _wrap_async_generate_content(
@@ -57,6 +61,7 @@ def setup_once():
         AsyncModels.generate_content_stream = _wrap_async_generate_content_stream(
             AsyncModels.generate_content_stream
         )
+        AsyncModels.embed_content = _wrap_async_embed_content(AsyncModels.embed_content)
 
 
 def _wrap_generate_content_stream(f):
@@ -299,3 +304,73 @@ async def new_async_generate_content(self, *args, **kwargs):
                 return response
 
     return new_async_generate_content
+
+
+def _wrap_embed_content(f):
+    # type: (Callable[..., Any]) -> Callable[..., Any]
+    @wraps(f)
+    def new_embed_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_name, contents = prepare_embed_content_args(args, kwargs)
+
+        with sentry_sdk.start_span(
+            op=OP.GEN_AI_EMBEDDINGS,
+            name=f"embeddings {model_name}",
+            origin=ORIGIN,
+        ) as span:
+            span.set_data(SPANDATA.GEN_AI_OPERATION_NAME, "embeddings")
+            span.set_data(SPANDATA.GEN_AI_SYSTEM, GEN_AI_SYSTEM)
+            span.set_data(SPANDATA.GEN_AI_REQUEST_MODEL, model_name)
+            set_span_data_for_embed_request(span, integration, contents, kwargs)
+
+            try:
+                response = f(self, *args, **kwargs)
+            except Exception as exc:
+                _capture_exception(exc)
+                span.set_status(SPANSTATUS.INTERNAL_ERROR)
+                raise
+
+            set_span_data_for_embed_response(span, integration, response)
+
+            return response
+
+    return new_embed_content
+
+
+def _wrap_async_embed_content(f):
+    # type: (Callable[..., Any]) -> Callable[..., Any]
+    @wraps(f)
+    async def new_async_embed_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_name, contents = prepare_embed_content_args(args, kwargs)
+
+        with sentry_sdk.start_span(
+            op=OP.GEN_AI_EMBEDDINGS,
+            name=f"embeddings {model_name}",
+            origin=ORIGIN,
+        ) as span:
+            span.set_data(SPANDATA.GEN_AI_OPERATION_NAME, "embeddings")
+            span.set_data(SPANDATA.GEN_AI_SYSTEM, GEN_AI_SYSTEM)
+            span.set_data(SPANDATA.GEN_AI_REQUEST_MODEL, model_name)
+            set_span_data_for_embed_request(span, integration, contents, kwargs)
+
+            try:
+                response = await f(self, *args, **kwargs)
+            except Exception as exc:
+                _capture_exception(exc)
+                span.set_status(SPANSTATUS.INTERNAL_ERROR)
+                raise
+
+            set_span_data_for_embed_response(span, integration, response)
+
+            return response
+
+    return new_async_embed_content

sentry_sdk/integrations/google_genai/utils.py

@@ -36,6 +36,7 @@
         ContentListUnion,
         Tool,
         Model,
+        EmbedContentResponse,
     )
 
 
@@ -574,3 +575,70 @@ def prepare_generate_content_args(args, kwargs):
         kwargs["config"] = wrapped_config
 
     return model, contents, model_name
+
+
+def prepare_embed_content_args(args, kwargs):
+    # type: (tuple[Any, ...], dict[str, Any]) -> tuple[str, Any]
+    """Extract and prepare common arguments for embed_content methods.
+
+    Returns:
+        tuple: (model_name, contents)
+    """
+    model = kwargs.get("model", "unknown")
+    contents = kwargs.get("contents")
+    model_name = get_model_name(model)
+
+    return model_name, contents
+
+
+def set_span_data_for_embed_request(span, integration, contents, kwargs):
+    # type: (Span, Any, Any, dict[str, Any]) -> None
+    """Set span data for embedding request."""
+    # Include input contents if PII is allowed
+    if should_send_default_pii() and integration.include_prompts:
+        if contents:
+            # For embeddings, contents is typically a list of strings/texts
+            input_texts = []
+
+            # Handle various content formats
+            if isinstance(contents, str):
+                input_texts = [contents]
+            elif isinstance(contents, list):
+                for item in contents:
+                    text = extract_contents_text(item)
+                    if text:
+                        input_texts.append(text)
+            else:
+                text = extract_contents_text(contents)
+                if text:
+                    input_texts = [text]
+
+            if input_texts:
+                set_data_normalized(
+                    span,
+                    SPANDATA.GEN_AI_EMBEDDINGS_INPUT,
+                    input_texts,
+                    unpack=False,
+                )
+
+
+def set_span_data_for_embed_response(span, integration, response):
+    # type: (Span, Any, EmbedContentResponse) -> None
+    """Set span data for embedding response."""
+    if not response:
+        return
+
+    # Extract token counts from embeddings statistics (Vertex AI only)
+    # Each embedding has its own statistics with token_count
+    if hasattr(response, "embeddings") and response.embeddings:
+        total_tokens = 0
+
+        for embedding in response.embeddings:
+            if hasattr(embedding, "statistics") and embedding.statistics:
+                token_count = getattr(embedding.statistics, "token_count", None)
+                if token_count is not None:
+                    total_tokens += int(token_count)
+
+        # Set token count if we found any
+        if total_tokens > 0:
+            span.set_data(SPANDATA.GEN_AI_USAGE_INPUT_TOKENS, total_tokens)