From 81d0474d90090fa6a7d848e0752a6029c0753e09 Mon Sep 17 00:00:00 2001 From: Nabin Mulepati Date: Mon, 9 Mar 2026 11:29:11 -0600 Subject: [PATCH 1/5] monkey patch litellm to make index optional --- .../engine/models/litellm_overrides.py | 16 +++++++++++++++- .../engine/models/test_litellm_overrides.py | 14 ++++++++++++++ 2 files changed, 29 insertions(+), 1 deletion(-) 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..7874a664d 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,25 @@ 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] + + 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..31b646cbb 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 @@ -14,6 +14,7 @@ CustomRouter, ThreadSafeCache, apply_litellm_patches, + patch_image_url_list_item, ) @@ -139,3 +140,16 @@ 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: + patch_image_url_list_item() + + 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" From ea45df53bfbe44af754d9429a3b80a7eadc54bce Mon Sep 17 00:00:00 2001 From: Nabin Mulepati Date: Mon, 9 Mar 2026 11:58:27 -0600 Subject: [PATCH 2/5] address greptile feedbaack --- .../engine/models/litellm_overrides.py | 2 + .../engine/models/test_litellm_overrides.py | 37 ++++++++++++++----- 2 files changed, 30 insertions(+), 9 deletions(-) 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 7874a664d..897536e6c 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 @@ -177,6 +177,8 @@ def patch_image_url_list_item(): 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"} def apply_litellm_patches(): 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 31b646cbb..03d124252 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 @@ -12,6 +12,7 @@ from data_designer.engine.models.litellm_overrides import ( DEFAULT_MAX_CALLBACKS, CustomRouter, + ImageURLListItem, ThreadSafeCache, apply_litellm_patches, patch_image_url_list_item, @@ -143,13 +144,31 @@ def test_custom_router_calculate_exponential_backoff_with_jitter(mock_uniform): def test_patch_image_url_list_item_makes_index_optional() -> None: - patch_image_url_list_item() + 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"} - 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" + assert "index" in ImageURLListItem.__required_keys__ + + 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 From 6132fabf1e628d80c736c4bfe359b372208be2c5 Mon Sep 17 00:00:00 2001 From: Nabin Mulepati Date: Mon, 9 Mar 2026 12:04:45 -0600 Subject: [PATCH 3/5] fix tests --- .../src/data_designer/engine/models/litellm_overrides.py | 4 ++++ 1 file changed, 4 insertions(+) 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 897536e6c..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 @@ -180,6 +180,10 @@ def patch_image_url_list_item(): 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() From f327ea2c03b3b874a11dc96020ef823c3c45465c Mon Sep 17 00:00:00 2001 From: Nabin Mulepati Date: Mon, 9 Mar 2026 12:26:45 -0600 Subject: [PATCH 4/5] fix: add model_rebuild calls for proper test isolation Rebuild the Pydantic Message model after restoring the TypedDict to unpatched state and in the finally block so cached schemas don't leak across tests. Also assert that Message construction fails without the patch to prove it is necessary. Made-with: Cursor --- .../tests/engine/models/test_litellm_overrides.py | 9 +++++++++ 1 file changed, 9 insertions(+) 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 03d124252..5581453aa 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 @@ -152,9 +152,17 @@ def test_patch_image_url_list_item_makes_index_optional() -> None: 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(Exception): + 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__ @@ -172,3 +180,4 @@ def test_patch_image_url_list_item_makes_index_optional() -> None: ImageURLListItem.__annotations__["index"] = original_annotation ImageURLListItem.__required_keys__ = original_required ImageURLListItem.__optional_keys__ = original_optional + litellm.Message.model_rebuild(force=True) From 0cef98df94710e1f0539a08f1d2c4e464e57b9ee Mon Sep 17 00:00:00 2001 From: Nabin Mulepati Date: Mon, 9 Mar 2026 12:36:19 -0600 Subject: [PATCH 5/5] fix greptile's weak suggestion --- .../tests/engine/models/test_litellm_overrides.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) 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 5581453aa..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,6 +7,7 @@ import litellm import pytest +from pydantic import ValidationError from data_designer.engine.models import litellm_overrides from data_designer.engine.models.litellm_overrides import ( @@ -156,7 +157,7 @@ def test_patch_image_url_list_item_makes_index_optional() -> None: assert "index" in ImageURLListItem.__required_keys__ - with pytest.raises(Exception): + with pytest.raises(ValidationError): litellm.Message( content=None, role="assistant",