fix: non passing tests

Author: carlos-marchal-ph

posthog/ai/gemini/gemini_converter.py

@@ -53,12 +53,15 @@ def _format_parts_as_content_blocks(parts: List[Any]) -> List[FormattedContentIt
         elif isinstance(part, str):
             content_blocks.append({"type": "text", "text": part})
 
-        # Handle dict with inline_data (images)
+        # Handle dict with inline_data (images, documents, etc.)
         elif isinstance(part, dict) and "inline_data" in part:
             inline_data = part["inline_data"]
+            mime_type = inline_data.get("mime_type", "")
+            content_type = "image" if mime_type.startswith("image/") else "document"
+
             content_blocks.append(
                 {
-                    "type": "image",
+                    "type": content_type,
                     "inline_data": inline_data,
                 }
             )
@@ -74,11 +77,15 @@ def _format_parts_as_content_blocks(parts: List[Any]) -> List[FormattedContentIt
             inline_data = part.inline_data
             # Convert to dict if needed
             if hasattr(inline_data, "mime_type") and hasattr(inline_data, "data"):
+                # Determine type based on mime_type
+                mime_type = inline_data.mime_type
+                content_type = "image" if mime_type.startswith("image/") else "document"
+
                 content_blocks.append(
                     {
-                        "type": "image",
+                        "type": content_type,
                         "inline_data": {
-                            "mime_type": inline_data.mime_type,
+                            "mime_type": mime_type,
                             "data": inline_data.data,
                         },
                     }

posthog/test/ai/gemini/test_gemini_async.py

@@ -392,7 +392,9 @@ async def test_async_client_different_input_formats(
     )
     call_args = mock_client.capture.call_args[1]
     props = call_args["properties"]
-    assert props["$ai_input"] == [{"role": "user", "content": "hey"}]
+    assert props["$ai_input"] == [
+        {"role": "user", "content": [{"type": "text", "text": "hey"}]}
+    ]
 
     # Test multiple parts in the parts array
     mock_client.reset_mock()
@@ -403,7 +405,15 @@ async def test_async_client_different_input_formats(
     )
     call_args = mock_client.capture.call_args[1]
     props = call_args["properties"]
-    assert props["$ai_input"] == [{"role": "user", "content": "Hello world"}]
+    assert props["$ai_input"] == [
+        {
+            "role": "user",
+            "content": [
+                {"type": "text", "text": "Hello "},
+                {"type": "text", "text": "world"},
+            ],
+        }
+    ]
 
     # Test list input with string
     mock_client.capture.reset_mock()