Handle bytes in google genai messages (#1231)

Author: alexmojaki

logfire/_internal/integrations/google_genai.py

@@ -1,15 +1,23 @@
 from __future__ import annotations
 
+import base64
 import json
 from typing import Any
 
 from opentelemetry._events import Event, EventLogger, EventLoggerProvider
 from opentelemetry.instrumentation.google_genai import GoogleGenAiSdkInstrumentor
 from opentelemetry.trace import get_current_span
+from typing_extensions import TypeAlias
 
 import logfire
 from logfire._internal.utils import handle_internal_errors
 
+Part: TypeAlias = 'dict[str, Any] | str'
+
+
+def default_json(x: Any) -> str:
+    return base64.b64encode(x).decode('utf-8') if isinstance(x, bytes) else x
+
 
 class SpanEventLogger(EventLogger):
     @handle_internal_errors
@@ -18,18 +26,25 @@ def emit(self, event: Event) -> None:
         assert isinstance(event.body, dict)
         body: dict[str, Any] = {**event.body}
         if event.name == 'gen_ai.choice':
-            parts = body.pop('content')['parts']
-            new_parts: list[dict[str, Any] | str] = []
-            for part in parts:
-                new_part: str | dict[str, Any] = {k: v for k, v in part.items() if v is not None}
-                if list(new_part.keys()) == ['text']:  # pragma: no branch
-                    new_part = new_part['text']
-                new_parts.append(new_part)
-            body['message'] = {'role': 'assistant', 'content': new_parts}
+            if 'content' in body:  # pragma: no branch
+                parts = body.pop('content')['parts']
+                new_parts = [transform_part(part) for part in parts]
+                body['message'] = {'role': 'assistant', 'content': new_parts}
         else:
+            if 'content' in body:  # pragma: no branch
+                body['content'] = transform_part(body['content'])
             body['role'] = body.get('role', event.name.split('.')[1])
 
-        span.add_event(event.name, attributes={'event_body': json.dumps(body)})
+        span.add_event(event.name, attributes={'event_body': json.dumps(body, default=default_json)})
+
+
+def transform_part(part: Part) -> Part:
+    if isinstance(part, str):
+        return part
+    new_part = {k: v for k, v in part.items() if v is not None}
+    if list(new_part.keys()) == ['text']:
+        return new_part['text']
+    return new_part
 
 
 class SpanEventLoggerProvider(EventLoggerProvider):

tests/otel_integrations/test_google_genai.py

@@ -43,7 +43,10 @@ def get_current_weather(location: str) -> str:
 
     response = client.models.generate_content(  # type: ignore
         model='gemini-2.0-flash-001',
-        contents='What is the weather like in Boston?',
+        contents=[
+            'What is the weather like in Boston?',
+            types.Part.from_bytes(data=b'123', mime_type='text/plain'),
+        ],
         config=types.GenerateContentConfig(
             tools=[get_current_weather],
         ),
@@ -57,7 +60,7 @@ def get_current_weather(location: str) -> str:
                 'context': {'trace_id': 1, 'span_id': 1, 'is_remote': False},
                 'parent': None,
                 'start_time': IsInt(),
-                'end_time': 3000000000,
+                'end_time': 4000000000,
                 'attributes': {
                     'code.function.name': 'google.genai.Models.generate_content',
                     'gen_ai.system': 'gemini',
@@ -71,6 +74,12 @@ def get_current_weather(location: str) -> str:
                     'logfire.metrics': IsPartialDict(),
                     'events': [
                         {'content': 'What is the weather like in Boston?', 'role': 'user'},
+                        {
+                            'content': {
+                                'inline_data': {'display_name': None, 'data': 'MTIz', 'mime_type': 'text/plain'}
+                            },
+                            'role': 'user',
+                        },
                         {
                             'index': 0,
                             'finish_reason': 'STOP',