diff --git a/packages/data-designer-engine/src/data_designer/engine/models/litellm_overrides.py b/packages/data-designer-engine/src/data_designer/engine/models/litellm_overrides.py index 92070def8..b60fea081 100644 --- a/packages/data-designer-engine/src/data_designer/engine/models/litellm_overrides.py +++ b/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") diff --git a/packages/data-designer-engine/tests/engine/models/test_litellm_overrides.py b/packages/data-designer-engine/tests/engine/models/test_litellm_overrides.py index 6ccbe63a6..aa547e187 100644 --- a/packages/data-designer-engine/tests/engine/models/test_litellm_overrides.py +++ b/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)