fix: preserve replay generator typing and cover switchyard boundaries

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.
 

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)

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")

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,
+        },
+    )

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"),
+    )

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()