chore: Update _flatten_ollama_content return type and add tests

The return type of _flatten_ollama_content is now strictly str | None

Co-authored-by: George Weale <gweale@google.com>
PiperOrigin-RevId: 845963722

src/google/adk/models/lite_llm.py

@@ -640,7 +640,7 @@ def _is_ollama_chat_provider(
 
 def _flatten_ollama_content(
     content: OpenAIMessageContent | str | None,
-) -> str | OpenAIMessageContent | None:
+) -> str | None:
   """Flattens multipart content to text for ollama_chat compatibility.
 
   Ollama's chat endpoint rejects arrays for `content`. We keep textual parts,

tests/unittests/models/test_litellm.py

@@ -1560,6 +1560,58 @@ def test_flatten_ollama_content_accepts_tuple_blocks():
   assert flattened == "first\nsecond"
 
 
+@pytest.mark.parametrize(
+    "content, expected",
+    [
+        (None, None),
+        ("hello", "hello"),
+        (
+            [
+                {"type": "text", "text": "first"},
+                {"type": "text", "text": "second"},
+            ],
+            "first\nsecond",
+        ),
+        (
+            [
+                {"type": "text", "text": "Describe this image."},
+                {
+                    "type": "image_url",
+                    "image_url": {"url": "http://example.com"},
+                },
+            ],
+            "Describe this image.",
+        ),
+    ],
+)
+def test_flatten_ollama_content_returns_str_or_none(content, expected):
+  from google.adk.models.lite_llm import _flatten_ollama_content
+
+  flattened = _flatten_ollama_content(content)
+  assert flattened == expected
+  assert flattened is None or isinstance(flattened, str)
+
+
+def test_flatten_ollama_content_serializes_non_text_blocks_to_json():
+  from google.adk.models.lite_llm import _flatten_ollama_content
+
+  blocks = [
+      {"type": "image_url", "image_url": {"url": "http://example.com"}},
+  ]
+  flattened = _flatten_ollama_content(blocks)
+  assert isinstance(flattened, str)
+  assert json.loads(flattened) == blocks
+
+
+def test_flatten_ollama_content_serializes_dict_to_json():
+  from google.adk.models.lite_llm import _flatten_ollama_content
+
+  content = {"type": "image_url", "image_url": {"url": "http://example.com"}}
+  flattened = _flatten_ollama_content(content)
+  assert isinstance(flattened, str)
+  assert json.loads(flattened) == content
+
+
 @pytest.mark.asyncio
 async def test_content_to_message_param_user_message():
   content = types.Content(