fix: tests and lint

Author: carlos-marchal-ph

posthog/ai/gemini/gemini_converter.py

@@ -175,7 +175,8 @@ def _format_object_message(item: Any) -> FormattedMessage:
         content = item.content
 
         if isinstance(content, list):
-            content = _extract_text_from_parts(content)
+            content_blocks = _format_parts_as_content_blocks(content)
+            return {"role": role, "content": content_blocks}
 
         elif not isinstance(content, str):
             content = str(content)

posthog/ai/sanitization.py

@@ -8,7 +8,11 @@
 
 def _is_multimodal_enabled() -> bool:
     """Check if multimodal capture is enabled via environment variable."""
-    return os.environ.get("_INTERNAL_LLMA_MULTIMODAL", "").lower() in ("true", "1", "yes")
+    return os.environ.get("_INTERNAL_LLMA_MULTIMODAL", "").lower() in (
+        "true",
+        "1",
+        "yes",
+    )
 
 
 def is_base64_data_url(text: str) -> bool:

posthog/test/ai/gemini/test_gemini.py

@@ -407,7 +407,9 @@ def test_new_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()
@@ -418,7 +420,15 @@ def test_new_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()

posthog/test/ai/test_sanitization.py

@@ -409,7 +409,11 @@ def test_gemini_multimodal_enabled(self):
         os.environ["_INTERNAL_LLMA_MULTIMODAL"] = "true"
 
         input_data = [
-            {"parts": [{"inline_data": {"mime_type": "image/jpeg", "data": "base64data"}}]}
+            {
+                "parts": [
+                    {"inline_data": {"mime_type": "image/jpeg", "data": "base64data"}}
+                ]
+            }
         ]
 
         result = sanitize_gemini(input_data)
@@ -509,7 +513,9 @@ def test_gemini_audio_preserved_with_flag(self):
         ]
 
         result = sanitize_gemini(input_data)
-        self.assertEqual(result[0]["parts"][0]["inline_data"]["data"], "base64audiodata")
+        self.assertEqual(
+            result[0]["parts"][0]["inline_data"]["data"], "base64audiodata"
+        )
 
 
 if __name__ == "__main__":