sentry-sdk 0.7.5__py2.py3-none-any.whl → 2.46.0__py2.py3-none-any.whl

sentry_sdk/integrations/google_genai/__init__.py

@@ -0,0 +1,301 @@
+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)
+        chat_span.set_data(SPANDATA.GEN_AI_AGENT_NAME, model_name)
+
+        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)
+        chat_span.set_data(SPANDATA.GEN_AI_AGENT_NAME, model_name)
+
+        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)
+                chat_span.set_data(SPANDATA.GEN_AI_AGENT_NAME, 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}"

sentry_sdk/integrations/google_genai/streaming.py

@@ -0,0 +1,155 @@
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    List,
+    TypedDict,
+    Optional,
+)
+
+from sentry_sdk.ai.utils import set_data_normalized
+from sentry_sdk.consts import SPANDATA
+from sentry_sdk.scope import should_send_default_pii
+from sentry_sdk.utils import (
+    safe_serialize,
+)
+from .utils import (
+    extract_tool_calls,
+    extract_finish_reasons,
+    extract_contents_text,
+    extract_usage_data,
+    UsageData,
+)
+
+if TYPE_CHECKING:
+    from sentry_sdk.tracing import Span
+    from google.genai.types import GenerateContentResponse
+
+
+class AccumulatedResponse(TypedDict):
+    id: Optional[str]
+    model: Optional[str]
+    text: str
+    finish_reasons: List[str]
+    tool_calls: List[dict[str, Any]]
+    usage_metadata: UsageData
+
+
+def accumulate_streaming_response(chunks):
+    # type: (List[GenerateContentResponse]) -> AccumulatedResponse
+    """Accumulate streaming chunks into a single response-like object."""
+    accumulated_text = []
+    finish_reasons = []
+    tool_calls = []
+    total_input_tokens = 0
+    total_output_tokens = 0
+    total_tokens = 0
+    total_cached_tokens = 0
+    total_reasoning_tokens = 0
+    response_id = None
+    model = None
+
+    for chunk in chunks:
+        # Extract text and tool calls
+        if getattr(chunk, "candidates", None):
+            for candidate in getattr(chunk, "candidates", []):
+                if hasattr(candidate, "content") and getattr(
+                    candidate.content, "parts", []
+                ):
+                    extracted_text = extract_contents_text(candidate.content)
+                    if extracted_text:
+                        accumulated_text.append(extracted_text)
+
+        extracted_finish_reasons = extract_finish_reasons(chunk)
+        if extracted_finish_reasons:
+            finish_reasons.extend(extracted_finish_reasons)
+
+        extracted_tool_calls = extract_tool_calls(chunk)
+        if extracted_tool_calls:
+            tool_calls.extend(extracted_tool_calls)
+
+        # Accumulate token usage
+        extracted_usage_data = extract_usage_data(chunk)
+        total_input_tokens += extracted_usage_data["input_tokens"]
+        total_output_tokens += extracted_usage_data["output_tokens"]
+        total_cached_tokens += extracted_usage_data["input_tokens_cached"]
+        total_reasoning_tokens += extracted_usage_data["output_tokens_reasoning"]
+        total_tokens += extracted_usage_data["total_tokens"]
+
+    accumulated_response = AccumulatedResponse(
+        text="".join(accumulated_text),
+        finish_reasons=finish_reasons,
+        tool_calls=tool_calls,
+        usage_metadata=UsageData(
+            input_tokens=total_input_tokens,
+            output_tokens=total_output_tokens,
+            input_tokens_cached=total_cached_tokens,
+            output_tokens_reasoning=total_reasoning_tokens,
+            total_tokens=total_tokens,
+        ),
+        id=response_id,
+        model=model,
+    )
+
+    return accumulated_response
+
+
+def set_span_data_for_streaming_response(span, integration, accumulated_response):
+    # type: (Span, Any, AccumulatedResponse) -> None
+    """Set span data for accumulated streaming response."""
+    if (
+        should_send_default_pii()
+        and integration.include_prompts
+        and accumulated_response.get("text")
+    ):
+        span.set_data(
+            SPANDATA.GEN_AI_RESPONSE_TEXT,
+            safe_serialize([accumulated_response["text"]]),
+        )
+
+    if accumulated_response.get("finish_reasons"):
+        set_data_normalized(
+            span,
+            SPANDATA.GEN_AI_RESPONSE_FINISH_REASONS,
+            accumulated_response["finish_reasons"],
+        )
+
+    if accumulated_response.get("tool_calls"):
+        span.set_data(
+            SPANDATA.GEN_AI_RESPONSE_TOOL_CALLS,
+            safe_serialize(accumulated_response["tool_calls"]),
+        )
+
+    if accumulated_response.get("id"):
+        span.set_data(SPANDATA.GEN_AI_RESPONSE_ID, accumulated_response["id"])
+    if accumulated_response.get("model"):
+        span.set_data(SPANDATA.GEN_AI_RESPONSE_MODEL, accumulated_response["model"])
+
+    if accumulated_response["usage_metadata"]["input_tokens"]:
+        span.set_data(
+            SPANDATA.GEN_AI_USAGE_INPUT_TOKENS,
+            accumulated_response["usage_metadata"]["input_tokens"],
+        )
+
+    if accumulated_response["usage_metadata"]["input_tokens_cached"]:
+        span.set_data(
+            SPANDATA.GEN_AI_USAGE_INPUT_TOKENS_CACHED,
+            accumulated_response["usage_metadata"]["input_tokens_cached"],
+        )
+
+    if accumulated_response["usage_metadata"]["output_tokens"]:
+        span.set_data(
+            SPANDATA.GEN_AI_USAGE_OUTPUT_TOKENS,
+            accumulated_response["usage_metadata"]["output_tokens"],
+        )
+
+    if accumulated_response["usage_metadata"]["output_tokens_reasoning"]:
+        span.set_data(
+            SPANDATA.GEN_AI_USAGE_OUTPUT_TOKENS_REASONING,
+            accumulated_response["usage_metadata"]["output_tokens_reasoning"],
+        )
+
+    if accumulated_response["usage_metadata"]["total_tokens"]:
+        span.set_data(
+            SPANDATA.GEN_AI_USAGE_TOTAL_TOKENS,
+            accumulated_response["usage_metadata"]["total_tokens"],
+        )