From ee42fc8f0662a0bed4dc6add0f76acc2b891c8db Mon Sep 17 00:00:00 2001 From: Joshua Valluru <326636767+joshua-berri@users.noreply.github.com> Date: Sat, 31 Oct 2026 16:46:04 -0700 Subject: [PATCH] fix: preserve replay generator typing and cover switchyard boundaries --- .../convert_dict_to_response.py | 8 +- .../switchyard_router/test_config.py | 29 +++++ .../test_switchyard_router.py | 35 ++++- .../switchyard_router/test_translation.py | 120 ++++++++++++++++++ .../test_auto_router_model_naming.py | 41 ++++++ tests/unit/test_router/test_router.py | 54 ++++++++ 6 files changed, 282 insertions(+), 5 deletions(-) diff --git a/litellm/litellm_core_utils/llm_response_utils/convert_dict_to_response.py b/litellm/litellm_core_utils/llm_response_utils/convert_dict_to_response.py index 2b97c30c5c7..6b3fd9511c2 100644 --- a/litellm/litellm_core_utils/llm_response_utils/convert_dict_to_response.py +++ b/litellm/litellm_core_utils/llm_response_utils/convert_dict_to_response.py @@ -3,7 +3,7 @@ import json import re import time import traceback -from collections.abc import AsyncIterator, Generator, Iterator, Mapping, Sequence +from collections.abc import AsyncGenerator, Generator, Iterator, Mapping, Sequence from typing import Final, Literal, cast, overload from pydantic import ConfigDict, TypeAdapter, ValidationError, with_config @@ -269,16 +269,16 @@ def _invalid_choices_message(response_object: Mapping[str, object]) -> str: @overload def convert_to_streaming_response_async( response_object: Mapping[str, object], -) -> AsyncIterator[ModelResponseStream]: ... +) -> AsyncGenerator[ModelResponseStream, None]: ... @overload -def convert_to_streaming_response_async(response_object: None = None) -> AsyncIterator[ModelResponseStream]: ... +def convert_to_streaming_response_async(response_object: None = None) -> AsyncGenerator[ModelResponseStream, None]: ... async def convert_to_streaming_response_async( response_object: Mapping | None = None, -) -> AsyncIterator[ModelResponseStream]: +) -> AsyncGenerator[ModelResponseStream, None]: """ Asynchronously converts a response object to a streaming response. diff --git a/tests/unit/router_strategy/switchyard_router/test_config.py b/tests/unit/router_strategy/switchyard_router/test_config.py index 69dea2bc61d..0f23d079857 100644 --- a/tests/unit/router_strategy/switchyard_router/test_config.py +++ b/tests/unit/router_strategy/switchyard_router/test_config.py @@ -87,3 +87,32 @@ def test_random_registry_keeps_weighted_candidates_even_with_extra_groups() -> N ) assert config.runtime_models() == groups assert config.weights == (0, 1, 0) + + +@pytest.mark.parametrize( + "config", + [ + {"algorithm": "capability", "models": {"efficient": ["fast"], "capable": ["strong"]}}, + {"algorithm": "custom", "models": {"judge": ["judge"], "general": ["fast"]}}, + { + "algorithm": "custom", + "models": {"judge": ["judge"], "general": ["fast"]}, + "custom": {"prompt": "Route", "response_schema": {}, "selector": "/target", "default_target": "missing"}, + }, + { + "algorithm": "custom", + "models": {"judge": ["judge"], "general": ["fast"]}, + "custom": { + "prompt": "Route", + "response_schema": {}, + "selector": "/target", + "default_target": "general", + "message_hash_fallback": True, + }, + }, + {"algorithm": "stage", "models": {"efficient": ["fast"], "capable": ["strong"]}, "weights": [1]}, + ], +) +def test_incomplete_or_incompatible_algorithm_configuration_is_rejected(config: dict[str, object]) -> None: + with pytest.raises(ValidationError): + SwitchyardRouterConfig.model_validate(config) diff --git a/tests/unit/router_strategy/switchyard_router/test_switchyard_router.py b/tests/unit/router_strategy/switchyard_router/test_switchyard_router.py index c2da26ab59c..ccc00684592 100644 --- a/tests/unit/router_strategy/switchyard_router/test_switchyard_router.py +++ b/tests/unit/router_strategy/switchyard_router/test_switchyard_router.py @@ -4,7 +4,7 @@ from typing import Final import pytest -from litellm.exceptions import APIError, AuthenticationError, ServiceUnavailableError +from litellm.exceptions import APIError, AuthenticationError, ContextWindowExceededError, ServiceUnavailableError from litellm.router_strategy.switchyard_router import SwitchyardAutoRouter, SwitchyardRouterConfig from litellm.types.utils import ModelResponse @@ -333,3 +333,36 @@ async def test_cancelling_judge_cancels_provider_work() -> None: with pytest.raises(asyncio.CancelledError): await task assert cancelled.is_set() + + +@pytest.mark.asyncio +async def test_disabling_fallback_preserves_selected_provider_failure() -> None: + error: Final = ServiceUnavailableError("unavailable", "fake", "fast") + provider: Final = ScriptedProvider({"fast": error}) + router: Final = SwitchyardAutoRouter(classifier_config("stage", fallback_enabled=False)) + with pytest.raises(ServiceUnavailableError) as raised: + await router.execute({"messages": [{"role": "user", "content": "hello"}]}, provider) + assert raised.value is error + assert tuple(model for model, _ in provider.calls) == ("fast",) + + +@pytest.mark.asyncio +async def test_escalation_context_overflow_executes_capable_without_judge() -> None: + answer: Final = response("strong", "fits") + provider: Final = ScriptedProvider( + { + "fast": ContextWindowExceededError("too large", "fast", "fake"), + "strong": answer, + } + ) + router: Final = SwitchyardAutoRouter(classifier_config("escalation", escalation={"confirmations": 1})) + actual: Final = await router.execute({"messages": [{"role": "user", "content": "hello"}]}, provider) + assert actual is answer + assert tuple(model for model, _ in provider.calls) == ("fast", "strong") + + +@pytest.mark.asyncio +async def test_model_selection_hook_rejects_unsupported_endpoint_dispatch() -> None: + router: Final = SwitchyardAutoRouter(classifier_config("capability")) + with pytest.raises(ValueError, match="async chat completions only"): + await router.async_pre_routing_hook("route", {}, input="embedding input") diff --git a/tests/unit/router_strategy/switchyard_router/test_translation.py b/tests/unit/router_strategy/switchyard_router/test_translation.py index 93f17ff8804..51dc6357d21 100644 --- a/tests/unit/router_strategy/switchyard_router/test_translation.py +++ b/tests/unit/router_strategy/switchyard_router/test_translation.py @@ -93,3 +93,123 @@ def test_judge_rewrite_removes_user_tools_limits_and_stop_sequences() -> None: def test_unsupported_content_is_rejected_before_execution(content: object) -> None: with pytest.raises(ValueError, match="text and function-tool"): normalize_request({"messages": [{"role": "user", "content": content}]}) + + +@pytest.mark.parametrize("choice", ["auto", "none", "required"]) +def test_standard_tool_choice_and_text_blocks_survive_roundtrip(choice: str) -> None: + request: Final = { + "messages": [{"role": "user", "content": [{"type": "text", "text": "hello"}]}], + "tool_choice": choice, + } + assert chat_request(normalize_request(request), request, model="fast", stream=False, judge=False) == { + **request, + "model": "fast", + "stream": False, + } + + +@pytest.mark.parametrize( + ("payload", "message"), + [ + ({"messages": []}, "at least one message"), + ({"messages": [{"role": "user", "content": "hello"}], "n": 2}, "n=1"), + ({"messages": [{"role": "function", "content": "hello"}]}, "unsupported chat role"), + ( + {"messages": [{"role": "user", "content": "hello"}], "tools": [{"type": "web_search"}]}, + "function tools only", + ), + ( + {"messages": [{"role": "user", "content": "hello"}], "tool_choice": {"type": "web_search"}}, + "function tool_choice", + ), + ( + { + "messages": [ + { + "role": "assistant", + "tool_calls": [{"id": "t", "type": "function", "function": {"name": "read", "arguments": {}}}], + } + ] + }, + "JSON arguments", + ), + ], +) +def test_unsupported_chat_inputs_fail_explicitly(payload: dict[str, object], message: str) -> None: + with pytest.raises(ValueError, match=message): + normalize_request(payload) + + +def test_rewritten_tool_calls_results_and_refusals_are_preserved() -> None: + normalized: Final = { + "messages": [ + { + "role": "assistant", + "content": [ + {"type": "tool_call", "id": "t", "name": "read", "arguments": {"path": "file"}}, + {"type": "refusal", "text": "Cannot modify"}, + ], + }, + { + "role": "tool", + "content": [ + { + "type": "tool_result", + "tool_call_id": "t", + "content": [{"type": "text", "text": "result"}], + "is_error": None, + } + ], + }, + ] + } + assert chat_request(normalized, {}, model="fast", stream=False, judge=False) == { + "model": "fast", + "stream": False, + "messages": [ + { + "role": "assistant", + "content": None, + "refusal": "Cannot modify", + "tool_calls": ( + {"id": "t", "type": "function", "function": {"name": "read", "arguments": '{"path": "file"}'}}, + ), + }, + {"role": "tool", "tool_call_id": "t", "content": ({"type": "text", "text": "result"},)}, + ], + } + + +@pytest.mark.parametrize( + ("normalized", "message"), + [ + ({"messages": [{"role": "tool", "content": []}]}, "unsupported tool result"), + ({"messages": [{"role": "assistant", "content": [{"type": "image", "source": {}}]}]}, "unsupported normalized"), + ({"messages": [], "tool_choice": {"type": "raw", "data": {}}}, "unsupported tool_choice"), + ], +) +def test_unsupported_native_rewrites_are_not_silently_dropped(normalized: dict[str, object], message: str) -> None: + with pytest.raises(ValueError, match=message): + chat_request(normalized, {}, model="fast", stream=False, judge=False) + + +@pytest.mark.parametrize( + ("finish", "expected"), [("length", "max_tokens"), ("tool_calls", "tool_use"), ("content_filter", "content_filter")] +) +def test_response_finish_reason_and_refusal_are_normalized(finish: str, expected: str) -> None: + response: Final = ModelResponse( + model="fast", + choices=[ + { + "message": {"role": "assistant", "content": None, "refusal": "Cannot comply"}, + "finish_reason": finish, + } + ], + ) + assert normalize_response(response, "internal-id")["outputs"] == ( + { + "role": "assistant", + "content": ({"type": "refusal", "text": "Cannot comply"},), + "stop_reason": expected, + }, + ) diff --git a/tests/unit/router_strategy/test_auto_router_model_naming.py b/tests/unit/router_strategy/test_auto_router_model_naming.py index 3df3ac377e2..5512f627656 100644 --- a/tests/unit/router_strategy/test_auto_router_model_naming.py +++ b/tests/unit/router_strategy/test_auto_router_model_naming.py @@ -15,6 +15,7 @@ from litellm.router_strategy.auto_router.auto_router_model_naming import ( validate_complexity_router_config_placement, validate_complexity_router_config_write, validate_strategy_router_model_write, + validate_switchyard_router_config_write, ) from litellm.router_strategy.complexity_router.fuse_presets import get_fuse_presets from litellm.router_strategy.complexity_router.jev_classifier import DEFAULT_JEV_INSTRUCTIONS @@ -807,3 +808,43 @@ def test_switchyard_dependencies_include_targets_and_judges(): validate_strategy_router_model_write("auto_router/switchyard", frozenset({"switchyard_router_config"})) is None ) assert validate_strategy_router_model_write("auto_router/switchyard", frozenset()) is not None + + +@pytest.mark.parametrize( + "config", + [None, {"algorithm": "random", "models": {"any": ["fast", "strong"]}}], +) +def test_switchyard_write_accepts_valid_config_and_unrelated_patch(config: Mapping[str, object] | None): + assert validate_switchyard_router_config_write(config) is None + + +@pytest.mark.parametrize( + "config, expected_location", + [ + ({"algorithm": "random", "models": {"any": ["fast"]}, "weights": [-1]}, "switchyard_router_config"), + ({"algorithm": "capability", "models": {"efficient": ["fast"]}}, "switchyard_router_config"), + ({"algorithm": "random", "models": {"any": ["fast"]}, "seed": -1}, "seed"), + ], +) +def test_switchyard_write_rejects_unloadable_configuration(config: Mapping[str, object], expected_location: str): + violation: Final = validate_switchyard_router_config_write(config) + assert violation is not None + assert f"invalid at {expected_location}:" in violation + + +@pytest.mark.parametrize("config", [None, "invalid", ["fast"], {"models": "invalid"}, {"models": None}]) +def test_switchyard_dependency_inspection_tolerates_malformed_saved_config(config: object): + assert strategy_router_dependencies({"model": "auto_router/switchyard", "switchyard_router_config": config}) == () + + +def test_switchyard_dependency_inspection_ignores_invalid_pool_entries(): + dependencies: Final = strategy_router_dependencies( + { + "model": "auto_router/switchyard", + "switchyard_router_config": {"models": {"any": ["fast", None, 4, ""], "judge": "judge", "invalid": 5}}, + } + ) + assert tuple((dependency.model_name, dependency.role) for dependency in dependencies) == ( + ("fast", "tier"), + ("judge", "classifier"), + ) diff --git a/tests/unit/test_router/test_router.py b/tests/unit/test_router/test_router.py index 9996ffcb098..6d605bf3cea 100644 --- a/tests/unit/test_router/test_router.py +++ b/tests/unit/test_router/test_router.py @@ -31882,3 +31882,57 @@ def test_image_routing_inspects_native_request_content_without_translation(reque original: Final = copy.deepcopy(request_kwargs) assert _candidate_ids(router, request_kwargs=request_kwargs) == ["vision"] assert request_kwargs == original + + +@pytest.mark.asyncio +@pytest.mark.parametrize("stream", [False, True]) +async def test_switchyard_escalation_returns_accepted_answer_through_router(stream: bool): + router: Final = Router( + model_list=[ + { + "model_name": "routed", + "litellm_params": { + "model": "auto_router/switchyard_router", + "switchyard_router_config": { + "algorithm": "escalation", + "models": {"efficient": ["fast"], "capable": ["strong"], "judge": ["judge"]}, + "escalation": {"confirmations": 1}, + }, + }, + }, + *( + { + "model_name": name, + "litellm_params": { + "model": "openai/gpt-4o-mini", + "api_key": "test-key", + "mock_response": text, + }, + } + for name, text in ( + ("fast", "accepted answer"), + ("strong", "strong answer"), + ("judge", '{"escalate":false,"category":"none","new_evidence":false,"reason":"progress"}'), + ) + ), + ], + num_retries=0, + ) + try: + response: Final = await router.acompletion( + model="routed", + messages=[{"role": "user", "content": "hello"}], + stream=stream, + stream_options={"include_usage": True} if stream else None, + ) + if stream: + chunks: Final = [chunk async for chunk in response] + assert ( + "".join(chunk.choices[0].delta.content or "" for chunk in chunks if chunk.choices) == "accepted answer" + ) + assert chunks[-1].usage.total_tokens > 0 + else: + assert response.choices[0].message.content == "accepted answer" + assert response.usage.total_tokens > 0 + finally: + router.reset()