mypy errors

Author: alexander-alderman-webb

sentry_sdk/integrations/google_genai/__init__.py

@@ -303,7 +303,7 @@ async def new_async_generate_content(
                     response = await f(self, *args, **kwargs)
                 except Exception as exc:
                     _capture_exception(exc)
-                    chat_span.set_status(SPANSTATUS.INTERNAL_ERROR)
+                    chat_span.status = SpanStatus.ERROR
                     raise
 
                 set_span_data_for_response(chat_span, integration, response)

sentry_sdk/integrations/google_genai/streaming.py

@@ -130,10 +130,13 @@ def set_span_data_for_streaming_response(
             safe_serialize(accumulated_response["tool_calls"]),
         )
 
-    if accumulated_response.get("id"):
-        set_on_span(SPANDATA.GEN_AI_RESPONSE_ID, accumulated_response["id"])
-    if accumulated_response.get("model"):
-        set_on_span(SPANDATA.GEN_AI_RESPONSE_MODEL, accumulated_response["model"])
+    response_id = accumulated_response.get("id")
+    if response_id is not None:
+        set_on_span(SPANDATA.GEN_AI_RESPONSE_ID, response_id)
+
+    response_model = accumulated_response.get("model")
+    if response_model is not None:
+        set_on_span(SPANDATA.GEN_AI_RESPONSE_MODEL, response_model)
 
     if accumulated_response["usage_metadata"] is None:
         return

sentry_sdk/integrations/google_genai/utils.py

@@ -1053,7 +1053,10 @@ def prepare_embed_content_args(
 
 
 def set_span_data_for_embed_request(
-    span: "Span", integration: "Any", contents: "Any", kwargs: "dict[str, Any]"
+    span: "Union[Span, StreamedSpan]",
+    integration: "Any",
+    contents: "Any",
+    kwargs: "dict[str, Any]",
 ) -> None:
     """Set span data for embedding request."""
     # Include input contents if PII is allowed