fix: patch litellm ImageURLListItem to make index field optional (#384)

packages/data-designer-engine/src/data_designer/engine/models/litellm_overrides.py

@@ -29,8 +29,9 @@
 from litellm.caching.in_memory_cache import InMemoryCache
 from litellm.litellm_core_utils.logging_callback_manager import LoggingCallbackManager
 from litellm.router import Router
+from litellm.types.llms.openai import ImageURLListItem
 from pydantic import BaseModel, Field
-from typing_extensions import override
+from typing_extensions import NotRequired, override
 
 from data_designer.logging import quiet_noisy_logger
 
@@ -168,12 +169,31 @@ def calculate_exponential_backoff(initial_retry_after_s: float, current_retry: i
         return sleep_s * jitter
 
 
+def patch_image_url_list_item():
+    """Make ImageURLListItem.index optional.
+
+    Some providers (e.g. OpenRouter) return image objects without the
+    ``index`` field. LiteLLM's TypedDict marks it as required, causing
+    a Pydantic validation error when constructing ``Message``.
+    """
+    ImageURLListItem.__annotations__["index"] = NotRequired[int]
+    ImageURLListItem.__required_keys__ = ImageURLListItem.__required_keys__ - {"index"}
+    ImageURLListItem.__optional_keys__ = ImageURLListItem.__optional_keys__ | {"index"}
+
+    # Pydantic v2 compiles TypedDict schemas at class definition time,
+    # so we must rebuild the Message model to pick up the annotation change.
+    litellm.Message.model_rebuild(force=True)
+
+
 def apply_litellm_patches():
     litellm.in_memory_llm_clients_cache = ThreadSafeCache()
 
     # Workaround for the litellm issue described in https://github.com/BerriAI/litellm/issues/9792
     LoggingCallbackManager.MAX_CALLBACKS = DEFAULT_MAX_CALLBACKS
 
+    # Workaround for missing 'index' field in image responses from some providers
+    patch_image_url_list_item()
+
     quiet_noisy_logger("httpx")
     quiet_noisy_logger("LiteLLM")
     quiet_noisy_logger("LiteLLM Router")

packages/data-designer-engine/tests/engine/models/test_litellm_overrides.py

@@ -7,13 +7,16 @@
 
 import litellm
 import pytest
+from pydantic import ValidationError
 
 from data_designer.engine.models import litellm_overrides
 from data_designer.engine.models.litellm_overrides import (
     DEFAULT_MAX_CALLBACKS,
     CustomRouter,
+    ImageURLListItem,
     ThreadSafeCache,
     apply_litellm_patches,
+    patch_image_url_list_item,
 )
 
 
@@ -139,3 +142,43 @@ def test_custom_router_calculate_exponential_backoff_with_jitter(mock_uniform):
     assert result >= 4.0
     assert result <= 4.4
     mock_uniform.assert_called_once_with(-0.2, 0.2)
+
+
+def test_patch_image_url_list_item_makes_index_optional() -> None:
+    original_annotation = ImageURLListItem.__annotations__["index"]
+    original_required = ImageURLListItem.__required_keys__
+    original_optional = ImageURLListItem.__optional_keys__
+    try:
+        # Restore to unpatched state in case prior tests already applied the patch
+        ImageURLListItem.__annotations__["index"] = int
+        ImageURLListItem.__required_keys__ = original_required | {"index"}
+        ImageURLListItem.__optional_keys__ = original_optional - {"index"}
+        litellm.Message.model_rebuild(force=True)
+
+        assert "index" in ImageURLListItem.__required_keys__
+
+        with pytest.raises(ValidationError):
+            litellm.Message(
+                content=None,
+                role="assistant",
+                images=[{"type": "image_url", "image_url": {"url": "data:image/png;base64,abc"}}],
+            )
+
+        patch_image_url_list_item()
+
+        assert "index" not in ImageURLListItem.__required_keys__
+        assert "index" in ImageURLListItem.__optional_keys__
+
+        message = litellm.Message(
+            content=None,
+            role="assistant",
+            images=[{"type": "image_url", "image_url": {"url": "data:image/png;base64,abc"}}],
+        )
+        assert message.images is not None
+        assert len(message.images) == 1
+        assert message.images[0]["type"] == "image_url"
+    finally:
+        ImageURLListItem.__annotations__["index"] = original_annotation
+        ImageURLListItem.__required_keys__ = original_required
+        ImageURLListItem.__optional_keys__ = original_optional
+        litellm.Message.model_rebuild(force=True)