From 9752869d8415150ee754165503a23129aeb29f97 Mon Sep 17 00:00:00 2001 From: Nate Armstrong Date: Wed, 4 Nov 2026 21:47:58 -0800 Subject: [PATCH] feat(router): add the Rust router core with the host as its only boundary (#49513) Adds the litellm-router-core crate, its step-driven PyO3 bridge (RouterCoreHandle / RouterCall) and RustCore (litellm/router_rust/), served behind LITELLM_RUST through the ROUTER route rule. The host call-out is the only boundary: every attempt goes through Host::invoke, there is no native executor path, and RustCore never builds, wraps or forwards to PythonCore, so router_python_core.py stays deletable. Attempt preparation lives in litellm/router_attempt.py so it outlives router_python_core.py. RustCore matches PythonCore's behavior, including its known bugs (fixed in follow-up PRs). Routing decisions (strategy picks, routing groups, PTU, pre-routing, pre-call checks, fallback walk, health routing, access groups, weights, cooldowns, allowed fails, usage) live in the crate; work that needs live Python objects (request shaping, endpoints, streaming, model info facts, callbacks, alerting, prompt management) lives in the host-side litellm/router_rust/ package (steps/, invoke/, surface/) driven by Rust steps. The crate owns the registry. It runs the deployment build loop (secret references, ids, organization fan-out, install order) and asks the host for what needs Python, one deployment at a time so errors surface where Python raised them: secrets, validate, provider and added steps. It holds the model list, the index maps and the pattern routers with the Python core's history rules (host_index.rs), and answers the registry lookups (lookup.rs). The Python views of that state in litellm/router_rust/surface/ (model_list, index(), names(), patterns()) read the handle, and every write to them becomes a changeset. Model info, deployment model info, configured limits and model group info are merged in the crate (model_info.rs, group_info.rs) from catalog rows and facts the host answers on a model_info step asked at query time, so price reloads and register_model stay visible. The usage cache key format is defined once, in the crate. The crate records each attempt's outcome itself rather than through litellm callbacks. A streamed attempt is deferred: it counts its request when the stream opens, and the host's stream wrappers report the success with the final usage, or a failure after the first chunk, when the stream ends, as the Python core's callbacks do. tests/unit/test_router/core characterizes PythonCore's public behavior and runs every test on both cores, with nothing pending; only the names of the Python core's own outcome callbacks are asserted on Python alone. litellm.Router is now final: subclassing it raises TypeError at the class statement, and the machinery that kept subclass overrides working is gone. The Rust core never honoured subclass overrides, so they would have silently stopped running there. Custom routing goes through register_plugin(), settings through apply() and outcomes through observe(). This is a breaking change for anyone subclassing Router. Fuzz parity fixes on top of the rebase: the Rust core keeps num_retries, max_fallbacks, allowed_fails, retry_after, timeout, the tag settings, routing_strategy_args and the fallback lists as given (the crate reads them as numbers, update_settings casts as Python does); a deprecated `model_list = [...]` write reaches the Rust core; a failed set_model_list is all or nothing on both cores (the deployments served before the call are put back and the error raised); config_deployments() is a generator and `del` of a plain attribute works on the Rust core. Least-busy pick reservations from main now work on the Rust core: its selectors share one reservation store, routed calls reserve and free their unstarted picks, and select() and batch cancel do not reserve. On the Python core a dropped upsert takes its name back off deployment_names, and get_deployment_order reads every numeric spelling as one int so mixed orders no longer raise TypeError. Rulings on the remaining differences, the Python core changed where it was the odd one out: sync picks and picks by deployment id apply the team filter as async picks do, so another team's deployment is refused; the fallback lists and default_litellm_params in RouterSettings are copies, so an in-place edit is ignored and settings change only through apply(); one health record with an unreadable timestamp drops out alone; an empty fallback entry is skipped; a bool cooldown_time is refused at construction, a bool or string cooldown_time on a deployment or a call falls back to the router default, and a failure status that no HTTP status can be is recorded as 500; an infinite or NaN cost and a whole number wider than 64 bits (ptu_shares) are refused on both cores with the same error (litellm/router_utils/deployment_numbers.py); and a /v1/messages stream whose provider iterator raises cools the deployment and retries like a non-stream call, with the Rust core waiting for that cooldown before its retry picks. Differences the proxy parity fuzz and its regression tests found on the Rust core: a file upload or a batch on an uploaded file falls back only within its own model group, as on the Python core, since the file lives with one provider (the host sends same_group_fallbacks_only and the walk skips targets in other groups); Decisions calls drop deployments whose provider has no Decisions config, and deployments that refuse the audio or videos a call carries, unless that would leave none, as the Python core does, while the retry count still reads the whole group; a file upload that fails on one order level is created on the next level's deployments rather than the first level again; and a response that cannot be converted (a text_completion body without a message) reports only its own traceback, so a rate limit handled by the attempt before it no longer turns the error into a 429, and both cores return 500. --- .github/workflows/test-unit.yml | 1 + litellm-rust/Cargo.lock | 20 + litellm-rust/Cargo.toml | 3 + litellm-rust/crates/python-bridge/Cargo.toml | 2 + litellm-rust/crates/python-bridge/src/lib.rs | 5 + .../crates/python-bridge/src/router/config.rs | 408 ++ .../crates/python-bridge/src/router/host.rs | 1569 ++++++ .../crates/python-bridge/src/router/mod.rs | 2051 +++++++ .../python-bridge/src/router/retries.rs | 55 + .../crates/python-bridge/src/router/stream.rs | 94 + litellm-rust/crates/router-core/Cargo.toml | 18 + litellm-rust/crates/router-core/src/access.rs | 236 + litellm-rust/crates/router-core/src/build.rs | 386 ++ litellm-rust/crates/router-core/src/cache.rs | 159 + litellm-rust/crates/router-core/src/call.rs | 1555 ++++++ .../crates/router-core/src/call/stream.rs | 456 ++ litellm-rust/crates/router-core/src/config.rs | 355 ++ .../crates/router-core/src/eligibility.rs | 240 + litellm-rust/crates/router-core/src/error.rs | 473 ++ .../crates/router-core/src/fallback.rs | 330 ++ .../crates/router-core/src/group_info.rs | 646 +++ litellm-rust/crates/router-core/src/health.rs | 960 ++++ .../crates/router-core/src/health_routing.rs | 126 + litellm-rust/crates/router-core/src/host.rs | 726 +++ .../crates/router-core/src/host_index.rs | 727 +++ litellm-rust/crates/router-core/src/lib.rs | 709 ++- litellm-rust/crates/router-core/src/lookup.rs | 883 +++ litellm-rust/crates/router-core/src/mock.rs | 44 + .../crates/router-core/src/model_info.rs | 716 +++ .../crates/router-core/src/outcome.rs | 125 + .../crates/router-core/src/precall.rs | 300 + .../crates/router-core/src/prerouting.rs | 53 + litellm-rust/crates/router-core/src/ptu.rs | 96 + .../crates/router-core/src/registry.rs | 1548 ++++++ .../crates/router-core/src/resolve.rs | 585 ++ litellm-rust/crates/router-core/src/retry.rs | 219 + .../crates/router-core/src/routing_groups.rs | 48 + .../crates/router-core/src/routing_plugins.rs | 114 + litellm-rust/crates/router-core/src/select.rs | 795 +++ .../crates/router-core/src/settings.rs | 133 + .../crates/router-core/src/strategy.rs | 110 + .../crates/router-core/src/weights.rs | 81 + .../crates/router-core/tests/call_loop.rs | 1757 ++++++ litellm/caching/dual_cache.py | 2 +- .../add_retry_fallback_headers.py | 2 +- litellm/litellm_core_utils/core_helpers.py | 11 + .../convert_dict_to_response.py | 6 +- .../routing_common_utils.py | 28 +- .../messages/adapters/streaming_iterator.py | 7 + .../responses_adapters/streaming_iterator.py | 7 + litellm/router.py | 166 +- litellm/router_api.py | 118 +- litellm/router_attempt.py | 215 + litellm/router_config.py | 7 +- litellm/router_core.py | 34 +- litellm/router_deprecated.py | 487 +- litellm/router_python_core.py | 107 +- litellm/router_python_host.py | 91 +- litellm/router_rust/__init__.py | 13 + litellm/router_rust/core.py | 4943 +++++++++++++++++ litellm/router_rust/invoke/README.md | 13 + litellm/router_rust/invoke/__init__.py | 1 + litellm/router_rust/invoke/endpoints.py | 2454 ++++++++ .../router_rust/invoke/prompt_management.py | 129 + litellm/router_rust/invoke/shaping.py | 635 +++ litellm/router_rust/invoke/slots.py | 107 + litellm/router_rust/invoke/streaming.py | 1839 ++++++ litellm/router_rust/invoke/walk.py | 398 ++ litellm/router_rust/invoke/websocket.py | 127 + litellm/router_rust/steps/README.md | 12 + litellm/router_rust/steps/__init__.py | 1 + litellm/router_rust/steps/access.py | 85 + litellm/router_rust/steps/callbacks.py | 650 +++ litellm/router_rust/steps/plugins.py | 87 + litellm/router_rust/steps/precall.py | 494 ++ litellm/router_rust/steps/prerouting.py | 1107 ++++ litellm/router_rust/steps/selection.py | 445 ++ litellm/router_rust/surface/README.md | 14 + litellm/router_rust/surface/__init__.py | 1 + litellm/router_rust/surface/alerting.py | 33 + litellm/router_rust/surface/fallbacks.py | 129 + litellm/router_rust/surface/health.py | 558 ++ litellm/router_rust/surface/health_routing.py | 160 + litellm/router_rust/surface/indexes.py | 514 ++ litellm/router_rust/surface/legacy_members.py | 547 ++ litellm/router_rust/surface/members.py | 163 + litellm/router_rust/surface/model_info.py | 871 +++ litellm/router_rust/surface/model_list.py | 375 ++ litellm/router_rust/surface/registry.py | 2191 ++++++++ litellm/router_rust/surface/replay.py | 30 + litellm/router_rust/surface/retries.py | 295 + litellm/router_rust/surface/settings.py | 236 + litellm/router_utils/cooldown_handlers.py | 14 +- litellm/router_utils/deployment_numbers.py | 90 + litellm/router_utils/endpoint_eligibility.py | 12 +- .../router_utils/fallback_event_handlers.py | 2 + .../router_utils/global_callback_holders.py | 53 + litellm/router_utils/health_state_cache.py | 12 +- litellm/rust_bridge/_native.pyi | 494 +- litellm/rust_bridge/catalog.py | 2 + litellm/types/router.py | 8 +- litellm/utils.py | 19 +- .../code_coverage_tests/recursive_detector.py | 3 + .../router_litellm_globals_check.py | 1 + .../router_surface_check.py | 8 +- tests/unit/conftest.py | 6 + .../test_invalid_response_traceback.py | 54 + .../litellm_core_utils/test_core_helpers.py | 19 + .../proxy/auth/test_fallback_model_access.py | 13 +- .../auth/test_model_access_group_budgets.py | 38 +- .../test_auto_router_endpoints.py | 42 +- tests/unit/router_rust/__init__.py | 0 .../router_rust/test_surface_fallbacks.py | 11 + .../test_fallback_event_handlers.py | 37 +- tests/unit/test_private_usage_aliases.py | 1 - tests/unit/test_router/core/conftest.py | 49 +- .../core/construction_precall_helpers.py | 98 + .../test_router/core/endpoints_helpers.py | 212 + .../test_router/core/fallbacks_helpers.py | 105 + .../test_router/core/gaps_l11100_helpers.py | 112 + .../unit/test_router/core/registry_helpers.py | 47 + .../core/retries_health_helpers.py | 198 + .../test_router/core/rust_parity_pending.txt | 2 + tests/unit/test_router/core/scripted_host.py | 98 +- .../test_router/core/selection_helpers.py | 88 + .../test_router/core/streaming_helpers.py | 474 ++ ...t_construction_precall_characterization.py | 1302 +++++ ...uction_precall_characterization_precall.py | 707 +++ .../core/test_deprecated_members.py | 665 +++ tests/unit/test_router/core/test_endpoints.py | 20 +- .../test_endpoints_characterization_batch.py | 525 ++ ...test_endpoints_characterization_generic.py | 850 +++ ...st_endpoints_characterization_inference.py | 648 +++ tests/unit/test_router/core/test_fallbacks.py | 177 + .../core/test_fallbacks_characterization.py | 584 ++ ...est_fallbacks_characterization_outcomes.py | 690 +++ ...test_fallbacks_characterization_refusal.py | 101 + .../core/test_fallbacks_walk_parity.py | 353 ++ .../core/test_gaps_adapter_timeout.py | 34 + .../test_router/core/test_gaps_fuzz_access.py | 267 + .../unit/test_router/core/test_gaps_l0001.py | 445 ++ .../core/test_gaps_l11100_registry.py | 732 +++ .../core/test_gaps_l11100_strategy.py | 438 ++ .../unit/test_router/core/test_gaps_l14000.py | 299 + .../core/test_gaps_l14000_selection.py | 491 ++ .../unit/test_router/core/test_gaps_l2800.py | 431 ++ .../core/test_gaps_l2800_streaming.py | 177 + .../unit/test_router/core/test_gaps_l5600.py | 479 ++ .../core/test_gaps_l5600_anthropic.py | 103 + .../core/test_gaps_l8400_precall.py | 178 + .../core/test_gaps_l8400_pricing.py | 240 + .../core/test_gaps_l8400_retries.py | 138 + .../test_gaps_pick_input_and_compaction.py | 253 + .../core/test_gaps_request_overrides.py | 227 + .../core/test_gaps_router_core_crate.py | 283 + .../core/test_gaps_router_core_crate_picks.py | 372 ++ .../core/test_health_registry_views.py | 174 + ..._host_settings_changes_characterization.py | 428 ++ .../test_router/core/test_legacy_members.py | 711 +++ .../test_parity_cooldown_service_model.py | 901 +++ .../core/test_parity_decisions_routing.py | 235 + .../core/test_parity_failure_logging.py | 647 +++ .../test_parity_malformed_after_rate_limit.py | 86 + ...test_parity_pick_count_and_fallback_log.py | 316 ++ .../test_parity_provider_scoped_fallbacks.py | 182 + .../test_router/core/test_parity_regress5.py | 482 ++ .../core/test_parity_regress5_errors.py | 501 ++ .../core/test_parity_regress5_settings.py | 473 ++ .../core/test_parity_regress5_streams.py | 297 + .../core/test_parity_response_headers.py | 205 + .../core/test_parity_rulings_2026_11.py | 286 + .../core/test_registry_changesets.py | 7 +- ..._registry_characterization_construction.py | 379 ++ .../test_registry_characterization_lookups.py | 682 +++ ...st_registry_characterization_model_info.py | 580 ++ .../test_registry_characterization_writes.py | 516 ++ .../core/test_request_scoped_outcomes.py | 466 ++ ...tries_health_characterization_cooldowns.py | 845 +++ ...s_health_characterization_failure_exits.py | 365 ++ ...retries_health_characterization_retries.py | 683 +++ tests/unit/test_router/core/test_selection.py | 39 +- .../test_selection_characterization_groups.py | 747 +++ ...t_selection_characterization_prerouting.py | 505 ++ ...election_characterization_session_repin.py | 119 + ...t_selection_characterization_strategies.py | 394 ++ ...test_selection_characterization_traffic.py | 145 + ...st_streaming_characterization_anthropic.py | 517 ++ .../test_streaming_characterization_chat.py | 737 +++ ...est_streaming_characterization_outcomes.py | 382 ++ ...st_streaming_characterization_responses.py | 458 ++ .../test_usage_and_failure_counting_parity.py | 455 ++ tests/unit/test_router/test_router.py | 100 +- .../test_router/test_router_entry_points.py | 109 +- .../unit/test_router_model_cost_isolation.py | 64 +- tests/unit/test_router_order_fallback.py | 6 +- 195 files changed, 69945 insertions(+), 772 deletions(-) create mode 100644 litellm-rust/crates/python-bridge/src/router/config.rs create mode 100644 litellm-rust/crates/python-bridge/src/router/host.rs create mode 100644 litellm-rust/crates/python-bridge/src/router/mod.rs create mode 100644 litellm-rust/crates/python-bridge/src/router/retries.rs create mode 100644 litellm-rust/crates/python-bridge/src/router/stream.rs create mode 100644 litellm-rust/crates/router-core/src/access.rs create mode 100644 litellm-rust/crates/router-core/src/build.rs create mode 100644 litellm-rust/crates/router-core/src/cache.rs create mode 100644 litellm-rust/crates/router-core/src/call.rs create mode 100644 litellm-rust/crates/router-core/src/call/stream.rs create mode 100644 litellm-rust/crates/router-core/src/config.rs create mode 100644 litellm-rust/crates/router-core/src/eligibility.rs create mode 100644 litellm-rust/crates/router-core/src/error.rs create mode 100644 litellm-rust/crates/router-core/src/fallback.rs create mode 100644 litellm-rust/crates/router-core/src/group_info.rs create mode 100644 litellm-rust/crates/router-core/src/health.rs create mode 100644 litellm-rust/crates/router-core/src/health_routing.rs create mode 100644 litellm-rust/crates/router-core/src/host.rs create mode 100644 litellm-rust/crates/router-core/src/host_index.rs create mode 100644 litellm-rust/crates/router-core/src/lookup.rs create mode 100644 litellm-rust/crates/router-core/src/mock.rs create mode 100644 litellm-rust/crates/router-core/src/model_info.rs create mode 100644 litellm-rust/crates/router-core/src/outcome.rs create mode 100644 litellm-rust/crates/router-core/src/precall.rs create mode 100644 litellm-rust/crates/router-core/src/prerouting.rs create mode 100644 litellm-rust/crates/router-core/src/ptu.rs create mode 100644 litellm-rust/crates/router-core/src/registry.rs create mode 100644 litellm-rust/crates/router-core/src/resolve.rs create mode 100644 litellm-rust/crates/router-core/src/retry.rs create mode 100644 litellm-rust/crates/router-core/src/routing_groups.rs create mode 100644 litellm-rust/crates/router-core/src/routing_plugins.rs create mode 100644 litellm-rust/crates/router-core/src/select.rs create mode 100644 litellm-rust/crates/router-core/src/settings.rs create mode 100644 litellm-rust/crates/router-core/src/strategy.rs create mode 100644 litellm-rust/crates/router-core/src/weights.rs create mode 100644 litellm-rust/crates/router-core/tests/call_loop.rs create mode 100644 litellm/router_attempt.py create mode 100644 litellm/router_rust/__init__.py create mode 100644 litellm/router_rust/core.py create mode 100644 litellm/router_rust/invoke/README.md create mode 100644 litellm/router_rust/invoke/__init__.py create mode 100644 litellm/router_rust/invoke/endpoints.py create mode 100644 litellm/router_rust/invoke/prompt_management.py create mode 100644 litellm/router_rust/invoke/shaping.py create mode 100644 litellm/router_rust/invoke/slots.py create mode 100644 litellm/router_rust/invoke/streaming.py create mode 100644 litellm/router_rust/invoke/walk.py create mode 100644 litellm/router_rust/invoke/websocket.py create mode 100644 litellm/router_rust/steps/README.md create mode 100644 litellm/router_rust/steps/__init__.py create mode 100644 litellm/router_rust/steps/access.py create mode 100644 litellm/router_rust/steps/callbacks.py create mode 100644 litellm/router_rust/steps/plugins.py create mode 100644 litellm/router_rust/steps/precall.py create mode 100644 litellm/router_rust/steps/prerouting.py create mode 100644 litellm/router_rust/steps/selection.py create mode 100644 litellm/router_rust/surface/README.md create mode 100644 litellm/router_rust/surface/__init__.py create mode 100644 litellm/router_rust/surface/alerting.py create mode 100644 litellm/router_rust/surface/fallbacks.py create mode 100644 litellm/router_rust/surface/health.py create mode 100644 litellm/router_rust/surface/health_routing.py create mode 100644 litellm/router_rust/surface/indexes.py create mode 100644 litellm/router_rust/surface/legacy_members.py create mode 100644 litellm/router_rust/surface/members.py create mode 100644 litellm/router_rust/surface/model_info.py create mode 100644 litellm/router_rust/surface/model_list.py create mode 100644 litellm/router_rust/surface/registry.py create mode 100644 litellm/router_rust/surface/replay.py create mode 100644 litellm/router_rust/surface/retries.py create mode 100644 litellm/router_rust/surface/settings.py create mode 100644 litellm/router_utils/deployment_numbers.py create mode 100644 litellm/router_utils/global_callback_holders.py create mode 100644 tests/unit/litellm_core_utils/llm_response_utils/test_invalid_response_traceback.py create mode 100644 tests/unit/router_rust/__init__.py create mode 100644 tests/unit/router_rust/test_surface_fallbacks.py create mode 100644 tests/unit/test_router/core/construction_precall_helpers.py create mode 100644 tests/unit/test_router/core/endpoints_helpers.py create mode 100644 tests/unit/test_router/core/fallbacks_helpers.py create mode 100644 tests/unit/test_router/core/gaps_l11100_helpers.py create mode 100644 tests/unit/test_router/core/registry_helpers.py create mode 100644 tests/unit/test_router/core/retries_health_helpers.py create mode 100644 tests/unit/test_router/core/rust_parity_pending.txt create mode 100644 tests/unit/test_router/core/selection_helpers.py create mode 100644 tests/unit/test_router/core/streaming_helpers.py create mode 100644 tests/unit/test_router/core/test_construction_precall_characterization.py create mode 100644 tests/unit/test_router/core/test_construction_precall_characterization_precall.py create mode 100644 tests/unit/test_router/core/test_deprecated_members.py create mode 100644 tests/unit/test_router/core/test_endpoints_characterization_batch.py create mode 100644 tests/unit/test_router/core/test_endpoints_characterization_generic.py create mode 100644 tests/unit/test_router/core/test_endpoints_characterization_inference.py create mode 100644 tests/unit/test_router/core/test_fallbacks_characterization.py create mode 100644 tests/unit/test_router/core/test_fallbacks_characterization_outcomes.py create mode 100644 tests/unit/test_router/core/test_fallbacks_characterization_refusal.py create mode 100644 tests/unit/test_router/core/test_fallbacks_walk_parity.py create mode 100644 tests/unit/test_router/core/test_gaps_adapter_timeout.py create mode 100644 tests/unit/test_router/core/test_gaps_fuzz_access.py create mode 100644 tests/unit/test_router/core/test_gaps_l0001.py create mode 100644 tests/unit/test_router/core/test_gaps_l11100_registry.py create mode 100644 tests/unit/test_router/core/test_gaps_l11100_strategy.py create mode 100644 tests/unit/test_router/core/test_gaps_l14000.py create mode 100644 tests/unit/test_router/core/test_gaps_l14000_selection.py create mode 100644 tests/unit/test_router/core/test_gaps_l2800.py create mode 100644 tests/unit/test_router/core/test_gaps_l2800_streaming.py create mode 100644 tests/unit/test_router/core/test_gaps_l5600.py create mode 100644 tests/unit/test_router/core/test_gaps_l5600_anthropic.py create mode 100644 tests/unit/test_router/core/test_gaps_l8400_precall.py create mode 100644 tests/unit/test_router/core/test_gaps_l8400_pricing.py create mode 100644 tests/unit/test_router/core/test_gaps_l8400_retries.py create mode 100644 tests/unit/test_router/core/test_gaps_pick_input_and_compaction.py create mode 100644 tests/unit/test_router/core/test_gaps_request_overrides.py create mode 100644 tests/unit/test_router/core/test_gaps_router_core_crate.py create mode 100644 tests/unit/test_router/core/test_gaps_router_core_crate_picks.py create mode 100644 tests/unit/test_router/core/test_health_registry_views.py create mode 100644 tests/unit/test_router/core/test_host_settings_changes_characterization.py create mode 100644 tests/unit/test_router/core/test_legacy_members.py create mode 100644 tests/unit/test_router/core/test_parity_cooldown_service_model.py create mode 100644 tests/unit/test_router/core/test_parity_decisions_routing.py create mode 100644 tests/unit/test_router/core/test_parity_failure_logging.py create mode 100644 tests/unit/test_router/core/test_parity_malformed_after_rate_limit.py create mode 100644 tests/unit/test_router/core/test_parity_pick_count_and_fallback_log.py create mode 100644 tests/unit/test_router/core/test_parity_provider_scoped_fallbacks.py create mode 100644 tests/unit/test_router/core/test_parity_regress5.py create mode 100644 tests/unit/test_router/core/test_parity_regress5_errors.py create mode 100644 tests/unit/test_router/core/test_parity_regress5_settings.py create mode 100644 tests/unit/test_router/core/test_parity_regress5_streams.py create mode 100644 tests/unit/test_router/core/test_parity_response_headers.py create mode 100644 tests/unit/test_router/core/test_parity_rulings_2026_11.py create mode 100644 tests/unit/test_router/core/test_registry_characterization_construction.py create mode 100644 tests/unit/test_router/core/test_registry_characterization_lookups.py create mode 100644 tests/unit/test_router/core/test_registry_characterization_model_info.py create mode 100644 tests/unit/test_router/core/test_registry_characterization_writes.py create mode 100644 tests/unit/test_router/core/test_request_scoped_outcomes.py create mode 100644 tests/unit/test_router/core/test_retries_health_characterization_cooldowns.py create mode 100644 tests/unit/test_router/core/test_retries_health_characterization_failure_exits.py create mode 100644 tests/unit/test_router/core/test_retries_health_characterization_retries.py create mode 100644 tests/unit/test_router/core/test_selection_characterization_groups.py create mode 100644 tests/unit/test_router/core/test_selection_characterization_prerouting.py create mode 100644 tests/unit/test_router/core/test_selection_characterization_session_repin.py create mode 100644 tests/unit/test_router/core/test_selection_characterization_strategies.py create mode 100644 tests/unit/test_router/core/test_selection_characterization_traffic.py create mode 100644 tests/unit/test_router/core/test_streaming_characterization_anthropic.py create mode 100644 tests/unit/test_router/core/test_streaming_characterization_chat.py create mode 100644 tests/unit/test_router/core/test_streaming_characterization_outcomes.py create mode 100644 tests/unit/test_router/core/test_streaming_characterization_responses.py create mode 100644 tests/unit/test_router/core/test_usage_and_failure_counting_parity.py diff --git a/.github/workflows/test-unit.yml b/.github/workflows/test-unit.yml index 94152aeb7a0..f1456b1fde8 100644 --- a/.github/workflows/test-unit.yml +++ b/.github/workflows/test-unit.yml @@ -419,6 +419,7 @@ jobs: - shard: router test-path: >- + tests/unit/router_rust tests/unit/test_router python-version: "3.12" workers: 4 diff --git a/litellm-rust/Cargo.lock b/litellm-rust/Cargo.lock index ba1131d7a92..23ebfc7d3cf 100644 --- a/litellm-rust/Cargo.lock +++ b/litellm-rust/Cargo.lock @@ -4288,6 +4288,7 @@ dependencies = [ name = "litellm-python-bridge" version = "0.1.0" dependencies = [ + "async-trait", "aws-sdk-secretsmanager", "bytes", "criterion", @@ -4312,6 +4313,7 @@ dependencies = [ "litellm-llms", "litellm-llms-types", "litellm-model-catalog", + "litellm-router-core", "litellm-router-types", "litellm-secret-value", "litellm-secrets", @@ -4393,6 +4395,24 @@ dependencies = [ [[package]] name = "litellm-router-core" version = "0.1.0" +dependencies = [ + "arc-swap", + "async-trait", + "jiff", + "litellm-cache", + "litellm-router-types", + "parking_lot", + "rand 0.8.7", + "rstest", + "serde", + "serde_json", + "serde_with", + "sha2 0.10.9", + "strum", + "thiserror 2.0.19", + "tokio", + "tracing", +] [[package]] name = "litellm-router-types" diff --git a/litellm-rust/Cargo.toml b/litellm-rust/Cargo.toml index 943e8576ae6..71914a0bffe 100644 --- a/litellm-rust/Cargo.toml +++ b/litellm-rust/Cargo.toml @@ -13,6 +13,7 @@ litellm-config = { path = "crates/config" } litellm-error = { path = "crates/error" } litellm-router = { path = "crates/router" } litellm-router-types = { path = "crates/router-types" } +litellm-router-core = { path = "crates/router-core" } litellm-tracing = { path = "crates/tracing" } litellm-spend-clickhouse = { path = "crates/spend-clickhouse" } litellm-storage-clickhouse = { path = "crates/storage-clickhouse" } @@ -76,6 +77,8 @@ litellm-cost = { path = "crates/cost" } litellm-model-catalog = { path = "crates/model-catalog" } litellm-model-catalog-refresh = { path = "crates/model-catalog-refresh" } +arc-swap = "1.9.2" +async-trait = "0.1.91" askama = { version = "0.16.1", default-features = false, features = ["derive", "std"] } tracing = "0.1" axum = { version = "0.8.9", default-features = false, features = ["http1", "tokio", "multipart"] } diff --git a/litellm-rust/crates/python-bridge/Cargo.toml b/litellm-rust/crates/python-bridge/Cargo.toml index fd826f549e5..ba892199dcf 100644 --- a/litellm-rust/crates/python-bridge/Cargo.toml +++ b/litellm-rust/crates/python-bridge/Cargo.toml @@ -25,6 +25,7 @@ litellm-tracing.workspace = true litellm-spend-clickhouse.workspace = true litellm-storage-clickhouse.workspace = true litellm-host.workspace = true +async-trait.workspace = true bytes.workspace = true futures-util.workspace = true litellm-cache.workspace = true @@ -47,6 +48,7 @@ litellm-secrets = { workspace = true, features = ["aws", "azure", "google", "has litellm-secrets-types.workspace = true litellm-llms-types.workspace = true litellm-router-types.workspace = true +litellm-router-core.workspace = true litellm-host-python.workspace = true litellm-model-catalog.workspace = true litellm-token-counter = { path = "../token-counter", default-features = false } diff --git a/litellm-rust/crates/python-bridge/src/lib.rs b/litellm-rust/crates/python-bridge/src/lib.rs index de35419308a..f4a4f7c2b26 100644 --- a/litellm-rust/crates/python-bridge/src/lib.rs +++ b/litellm-rust/crates/python-bridge/src/lib.rs @@ -12,6 +12,7 @@ mod marshal; mod model_registry; mod preflight; mod python_settings; +mod router; mod routes; mod secrets; mod tokenizer; @@ -36,6 +37,8 @@ mod _native { #[pymodule_export] use crate::model_registry::bundled_model_cost_map; #[pymodule_export] + use crate::router::{RouterCall, RouterCoreHandle}; + #[pymodule_export] use crate::routes::audio_transcription::{atranscription, transcription}; #[pymodule_export] use crate::routes::chat_completions::{acompletion, completion}; @@ -115,6 +118,8 @@ mod tests { "bundled_model_cost_map", "TokenCounter", "Tokenizer", + "RouterCoreHandle", + "RouterCall", "gil_stats", "process_state_started", "reserve_process_for_forking", diff --git a/litellm-rust/crates/python-bridge/src/router/config.rs b/litellm-rust/crates/python-bridge/src/router/config.rs new file mode 100644 index 00000000000..0b46f2208c5 --- /dev/null +++ b/litellm-rust/crates/python-bridge/src/router/config.rs @@ -0,0 +1,408 @@ +//! The documents `RouterCoreHandle` reads: the router's static and dynamic config, a call's +//! request and an `apply` change set. A config key router-core does not read declines the +//! handle, so Python keeps the router on its own core. + +use std::collections::BTreeMap; + +use litellm_router_core::select::PassThrough; +use litellm_router_core::{ + ChangeSet, ConfigError, DeploymentSpec, OutcomeCallbacks, RouterConfig, RouterDynamicConfig, + RouterRequest, + config::{Fallbacks, RetryPolicy}, + host::UncountedRetries, + mock::MockFailure, + registry::IndexOverrides, +}; +use serde::{Deserialize, de::DeserializeOwned}; +use serde_json::{Map, Value}; + +/// The `RouterConfig` keys router-core reads. +const STATIC_KEYS: &[&str] = &[ + "model_list", + "routing_strategy", + "routing_groups", + "routing_group_members", + "plugins", + "cooldown_time", + "disable_cooldowns", + "allowed_fails_policy", + "retry_policy", + "model_group_retry_policy", + "retry_after", + "timeout", + "request_timeout", + "drop_params", + "enable_pre_call_checks", + "routing_plugins", + "callback_filters", + "enable_health_check_routing", + "health_check_staleness_threshold", + "background_health_check_model_groups", + "enable_weighted_failover", + "fallback_on_unsupported_params", + "provider_default_deployment_ids", +]; + +/// The `RouterDynamicConfig` keys router-core reads. +const DYNAMIC_KEYS: &[&str] = &[ + "version", + "model_alias_map", + "model_group_alias", + "fallbacks", + "default_fallbacks", + "context_window_fallbacks", + "content_policy_fallbacks", + "num_retries", + "max_fallbacks", + "allowed_fails", + "allowed_fails_is_default", + "fallback_keys", + "default_deployment", + "default_deployment_raises", +]; + +/// Router arguments that do not change routing, accepted in either config and dropped. +const IGNORED_KEYS: &[&str] = &[ + "set_verbose", + "debug_level", + "alerting_config", + "client_ttl", +]; + +/// Fallback lists Python treats as unset when empty (`context_window_fallbacks or None`). +const EMPTY_MEANS_UNSET: &[&str] = &["context_window_fallbacks", "content_policy_fallbacks"]; + +/// Which document a failure is about. +#[derive(Clone, Copy, Debug, PartialEq, Eq, strum::Display)] +pub(crate) enum Document { + #[strum(serialize = "router config")] + Static, + #[strum(serialize = "dynamic router config")] + Dynamic, + #[strum(serialize = "router request")] + Request, + #[strum(serialize = "change set")] + Changes, +} + +#[derive(Debug, thiserror::Error)] +pub(crate) enum Error { + #[error("the {document} must be a JSON object")] + NotObject { document: Document }, + #[error("router-core does not read `{key}` in the {document}")] + UnknownKey { document: Document, key: String }, + #[error("the {document} does not have the shape router-core reads")] + Shape { + document: Document, + #[source] + source: serde_json::Error, + }, + #[error("the request names no model")] + MissingModel, + #[error("router-core rejected the router config")] + Core(#[from] ConfigError), +} + +impl Error { + /// Whether router-core cannot serve this config, so Python keeps its own core. + pub(crate) fn declined(&self) -> bool { + match self { + Self::UnknownKey { .. } | Self::Core(_) => true, + Self::Shape { document, .. } => { + matches!(document, Document::Static | Document::Dynamic) + } + Self::NotObject { .. } | Self::MissingModel => false, + } + } +} + +fn object(document: Document, value: Value) -> Result, Error> { + match value { + Value::Object(object) => Ok(object), + _ => Err(Error::NotObject { document }), + } +} + +fn shaped(document: Document, object: Map) -> Result { + serde_json::from_value(Value::Object(object)) + .map_err(|source| Error::Shape { document, source }) +} + +/// Drop unset and ignored keys, then reject any key outside `known`. +fn routing_keys( + document: Document, + value: Value, + known: &[&str], +) -> Result, Error> { + let mut object = object(document, value)?; + object.retain(|key, value| { + let empty_unset = EMPTY_MEANS_UNSET.contains(&key.as_str()) + && value.as_array().is_some_and(Vec::is_empty); + !(value.is_null() || empty_unset || IGNORED_KEYS.contains(&key.as_str())) + }); + if let Some(key) = object.keys().find(|key| !known.contains(&key.as_str())) { + return Err(Error::UnknownKey { + document, + key: key.clone(), + }); + } + Ok(object) +} + +pub(crate) fn parse_router_config(value: Value) -> Result { + let object = routing_keys(Document::Static, value, STATIC_KEYS)?; + shaped(Document::Static, object) +} + +pub(crate) fn parse_dynamic_config(value: Value) -> Result { + let object = routing_keys(Document::Dynamic, value, DYNAMIC_KEYS)?; + shaped(Document::Dynamic, object) +} + +/// A call's request as Python spells it. +#[derive(Deserialize)] +#[serde(deny_unknown_fields)] +struct RequestSpec { + model: Option, + team_id: Option, + listing_team_id: Option, + proxy_admin: Option, + allowed_models: Option>, + admin_model: Option, + specific_deployment: Option, + num_retries: Option, + uncounted_retries: Option, + disable_fallbacks: Option, + router_weights: Option, + pinned_deployment_id: Option, + routing_strategy: Option, + filter_tags: Option, + /// `listed` or `healthy`: which pass-through pick the request makes. + pass_through: Option, + target_order: Option, + excluded_ids: Option>, + retry_skipped_ids: Option>, + pre_routes: Option, + has_image: Option, + endpoint: Option, + web_search: Option, + route_handler: Option, + decisions_media: Option, + wildcard_request: Option, + unrouted: Option, + passthrough_on_no_deployment: Option, + skip_weighted_failover: Option, + same_group_fallbacks_only: Option, + model_group_retry_policy: Option>, + retry_policy_raises: Option>, + #[serde(default, deserialize_with = "present")] + fallbacks: Option>, + #[serde(default, deserialize_with = "present")] + context_window_fallbacks: Option>, + #[serde(default, deserialize_with = "present")] + content_policy_fallbacks: Option>, + max_fallbacks: Option, + mock_failure: Option, + cooldown_time: Option, + health_check_probe: Option, + outcome_callbacks: Option, + body: Option, +} + +/// A key the request passed, even as `None`: `Some(None)`. +fn present<'de, D, T>(deserializer: D) -> Result>, D::Error> +where + D: serde::Deserializer<'de>, + T: Deserialize<'de>, +{ + Option::::deserialize(deserializer).map(Some) +} + +/// Read a request; `model` replaces the request's own when given. +pub(crate) fn parse_request(value: Value, model: Option<&str>) -> Result { + let spec: RequestSpec = shaped(Document::Request, object(Document::Request, value)?)?; + let model = model + .map(str::to_owned) + .or(spec.model) + .ok_or(Error::MissingModel)?; + Ok(RouterRequest { + model, + body: spec.body.unwrap_or_default(), + team_id: spec.team_id, + listing_team_id: spec.listing_team_id, + proxy_admin: spec.proxy_admin.unwrap_or(false), + allowed_models: spec.allowed_models.unwrap_or_default(), + admin_model: spec.admin_model, + specific_deployment: spec.specific_deployment.unwrap_or(false), + num_retries: spec.num_retries, + uncounted_retries: spec.uncounted_retries, + disable_fallbacks: spec.disable_fallbacks.unwrap_or(false), + router_weights: spec.router_weights, + pinned_deployment_id: spec.pinned_deployment_id, + routing_strategy: spec.routing_strategy, + filter_tags: spec.filter_tags.unwrap_or(false), + pass_through: match spec.pass_through.as_deref() { + None => None, + Some("healthy") => Some(PassThrough::Healthy), + Some(_) => Some(PassThrough::Listed), + }, + target_order: spec.target_order, + excluded_ids: spec.excluded_ids.unwrap_or_default(), + retry_skipped_ids: spec.retry_skipped_ids.unwrap_or_default(), + pre_routes: spec.pre_routes.unwrap_or(false), + has_image: spec.has_image.unwrap_or(false), + endpoint: spec.endpoint, + web_search: spec.web_search.unwrap_or(false), + route_handler: spec.route_handler, + decisions_media: spec.decisions_media.unwrap_or(false), + wildcard_request: spec.wildcard_request, + unrouted: spec.unrouted.unwrap_or(false), + passthrough_on_no_deployment: spec.passthrough_on_no_deployment.unwrap_or(false), + skip_weighted_failover: spec.skip_weighted_failover.unwrap_or(false), + same_group_fallbacks_only: spec.same_group_fallbacks_only.unwrap_or(false), + model_group_retry_policy: spec.model_group_retry_policy, + retry_policy_raises: spec.retry_policy_raises.unwrap_or_default(), + fallbacks: spec.fallbacks, + context_window_fallbacks: spec.context_window_fallbacks, + content_policy_fallbacks: spec.content_policy_fallbacks, + max_fallbacks: spec.max_fallbacks, + mock_failure: spec.mock_failure, + cooldown_time: spec.cooldown_time, + health_check_probe: spec.health_check_probe.unwrap_or(false), + outcome_callbacks: spec.outcome_callbacks.unwrap_or_default(), + payload: None, + }) +} + +#[derive(Deserialize)] +#[serde(deny_unknown_fields)] +struct ChangesSpec { + /// Every deployment, replacing the current list before the upserts and deletes. + #[serde(default)] + replace: Option>, + /// Deployments added at the end, beside any that share their id. + #[serde(default)] + append: Vec, + #[serde(default)] + upsert: Vec, + #[serde(default)] + delete: Vec, + /// A new dynamic document (the same keys `RouterCoreHandle(static, dynamic)` takes). + #[serde(default)] + settings: Option, + /// The registry indexes the host wrote (`IndexOverrides`); `{}` drops them. + #[serde(default)] + index: Option, + /// The deployment at the id map's position for this id, removed with every index shifted. + #[serde(default)] + remove: Option, + /// A write to the indexes the host hands out (`HostIndexes`), field by field. + #[serde(default)] + host_index: Option, +} + +pub(crate) fn parse_changes(value: Value) -> Result { + let spec: ChangesSpec = shaped(Document::Changes, object(Document::Changes, value)?)?; + Ok(ChangeSet { + replace: spec.replace, + appends: spec.append, + upserts: spec.upsert, + deletes: spec.delete, + settings: spec.settings.map(parse_dynamic_config).transpose()?, + index: spec.index, + remove: spec.remove, + host_index: spec.host_index, + }) +} + +#[cfg(test)] +mod tests { + use rstest::rstest; + use serde_json::json; + + use super::*; + + #[rstest] + #[case::unknown_key(json!({"routing_strategy_args": {"ttl": 1}}), "routing_strategy_args")] + #[case::dynamic_key_in_static(json!({"num_retries": 2}), "num_retries")] + fn an_unread_static_key_declines_naming_the_key(#[case] json: Value, #[case] key: &str) { + let error = parse_router_config(json).unwrap_err(); + assert!(error.declined()); + assert!(matches!(&error, Error::UnknownKey { key: named, .. } if named == key)); + } + + #[rstest] + #[case::null_unknown(json!({"routing_strategy_args": null}))] + #[case::ignored(json!({"set_verbose": true, "debug_level": "INFO"}))] + #[case::pre_call_checks_off(json!({"enable_pre_call_checks": false}))] + #[case::pre_call_checks_on(json!({"enable_pre_call_checks": true}))] + fn unset_and_ignored_keys_are_accepted(#[case] json: Value) { + parse_router_config(json).unwrap(); + } + + #[rstest] + #[case::not_object(json!([]), false)] + #[case::bad_shape(json!({"retry_after": "soon"}), true)] + fn a_config_that_is_not_an_object_is_a_value_error_and_a_wrong_shape_declines( + #[case] json: Value, + #[case] declined: bool, + ) { + assert_eq!(parse_router_config(json).unwrap_err().declined(), declined); + } + + #[rstest] + fn python_router_defaults_read_as_unset() { + let dynamic = parse_dynamic_config(json!({ + "num_retries": null, "fallbacks": [], "context_window_fallbacks": [], + "content_policy_fallbacks": [], "model_group_alias": {"gpt": "gpt-4"}, + })) + .unwrap(); + assert_eq!(dynamic.num_retries, None); + assert!(dynamic.context_window_fallbacks.is_none()); + assert!(dynamic.content_policy_fallbacks.is_none()); + assert_eq!(dynamic.model_alias_map["gpt"], "gpt-4"); + } + + #[rstest] + #[case::own_model(json!({"model": "gpt", "num_retries": 1}), None, "gpt")] + #[case::override_model(json!({"model": "gpt"}), Some("other"), "other")] + fn a_request_reads_its_model( + #[case] json: Value, + #[case] model: Option<&str>, + #[case] expected: &str, + ) { + assert_eq!(parse_request(json, model).unwrap().model, expected); + } + + #[rstest] + #[case::negative("negative", UncountedRetries::Negative)] + #[case::not_positive("not_positive", UncountedRetries::NotPositive)] + #[case::raises("raises", UncountedRetries::Raises)] + fn a_request_reads_a_num_retries_that_is_no_count( + #[case] spelled: &str, + #[case] expected: UncountedRetries, + ) { + let request = + parse_request(json!({"model": "gpt", "uncounted_retries": spelled}), None).unwrap(); + assert_eq!(request.uncounted_retries, Some(expected)); + assert_eq!(request.num_retries, None); + } + + #[rstest] + #[case::unknown_field(json!({"model": "gpt", "temperature": 1}))] + #[case::no_model(json!({"team_id": "t"}))] + fn a_bad_request_is_a_value_error(#[case] json: Value) { + assert!(!parse_request(json, None).unwrap_err().declined()); + } + + #[rstest] + fn a_change_set_reads_upserts_and_deletes() { + let changes = parse_changes(json!({ + "upsert": [{"model_name": "m", "litellm_params": {"model": "openai/x"}}], + "delete": ["old"], + })) + .unwrap(); + assert_eq!(changes.upserts.len(), 1); + assert_eq!(changes.deletes, ["old"]); + } +} diff --git a/litellm-rust/crates/python-bridge/src/router/host.rs b/litellm-rust/crates/python-bridge/src/router/host.rs new file mode 100644 index 00000000000..16b425d69c6 --- /dev/null +++ b/litellm-rust/crates/python-bridge/src/router/host.rs @@ -0,0 +1,1569 @@ +//! The host behind `RouterCoreHandle`. Rust never awaits Python: each host operation the call +//! loop needs becomes a [`HostCall`] parked in the call's [`Exchange`], and the call future stays +//! pending until Python runs it and hands the [`HostCallResponse`] back through [`Driver::resume`]. +//! +//! The driver polls the call future with a no-op waker while the call's exchange is bound to the +//! current thread, so a host future finds its exchange without the router passing it through. +//! +//! The routing cache is a host operation too: [`StepCache`] parks every [`RouterCache`] op as a +//! `cache` step, and Python runs it against the router's own `DualCache` (memory plus Redis), so +//! Python and Rust read and write the same cooldowns, fail counters and usage. + +use std::{ + cell::RefCell, + future::{Future, poll_fn}, + pin::Pin, + sync::{Arc, Mutex, MutexGuard, PoisonError}, + task::{Context, Poll, Waker}, + time::Duration, +}; + +use async_trait::async_trait; +use litellm_router_core::{ + Attempt, AttemptError, Blocked, BuildAsk, BuildPath, Built, Deployment, Error, Host, Output, + Payload, PreRoute, PreRouted, Provided, Randomness, RouterCache, RouterEvent, RouterRequest, + RouterResponse, RoutingContext, SecretRef, SelectStage, ThreadRandomness, + build::BuiltDeployment, host::ModelInfoLookup, model_info::MergedModelInfo, + precall::PreCallFacts, +}; +use pyo3::{PyResult, Python, types::PyAnyMethods}; +use serde_json::{Map, Value}; + +/// The `kind` of a step dict. +#[derive(Clone, Copy, Debug, PartialEq, Eq, strum::Display, strum::IntoStaticStr)] +pub(crate) enum StepKind { + #[strum(serialize = "attempt")] + Attempt, + #[strum(serialize = "sleep")] + Sleep, + #[strum(serialize = "allow_fallback")] + AllowFallback, + #[strum(serialize = "refuse_fallbacks")] + RefuseFallbacks, + #[strum(serialize = "select")] + Select, + #[strum(serialize = "pre_route")] + PreRoute, + #[strum(serialize = "cache")] + Cache, + #[strum(serialize = "secrets")] + Secrets, + #[strum(serialize = "validate")] + Validate, + #[strum(serialize = "provider")] + Provider, + #[strum(serialize = "added")] + Added, + #[strum(serialize = "model_info")] + ModelInfo, + #[strum(serialize = "done")] + Done, + #[strum(serialize = "failed")] + Failed, +} + +/// Which selection stage Python is asked to run. +#[derive(Clone, Debug, PartialEq)] +pub(crate) enum Stage { + /// Pre-call checks and sub-routers: return the deployments that remain. + Plugins, + /// A strategy router-core does not implement: return the candidates with the pick first. + Strategy(String), + /// Tag routing: return the deployments that remain. + Tags, + /// `enable_pre_call_checks`: return the facts the filter decides by. + PreCall, + /// `_deployment_declares_no_vision`: return the candidates that declare no vision support. + TextOnly, + /// `_deployment_serves_route`: return the candidates that cannot serve `handler`'s call. + RouteIncapable { handler: String }, + /// `_decisions_deployment_refuses_media`: return the candidates that refuse the call's media. + MediaRefusing, + /// The routing plugins: run them when `run`, then return the candidate models they left. + RoutingPlugins { run: bool }, + /// The host's `CustomLogger` deployment filters: return the ids that remain, or a failure. + /// They read the ids this request already failed on, as Python's request kwargs carry them. + CallbackFilters { + excluded_ids: Vec, + skipped_ids: Vec, + target_order: Option, + }, + /// `f"{provider}/{model}"` for a name no wildcard pattern matches as spelled: return it as + /// the one entry of `models`. + ProviderRequest, +} + +/// The `op` of a `cache` step: one per [`RouterCache`] method. +#[derive(Clone, Copy, Debug, PartialEq, Eq, strum::Display, strum::IntoStaticStr)] +pub(crate) enum CacheOp { + /// `async_get_cache(key)`: answered with the stored JSON value, or `None`. + #[strum(serialize = "get")] + Get, + /// `async_set_cache(key, value, ttl)`: answered with `None`. + #[strum(serialize = "set")] + Set, + /// `async_increment(key, value, ttl)`: answered with the new number. The TTL applies when + /// the key is created. + #[strum(serialize = "increment")] + Increment, + /// `async_increment_usage(keys, value, ttl)`: `value` is one amount per key, added as + /// `_increment_deployment_usage` adds them. Answered with `None`. + #[strum(serialize = "increment_usage")] + IncrementUsage, + /// `async_batch_get_cache(keys)`: answered with one JSON value or `None` per key, in order. + #[strum(serialize = "batch_get")] + BatchGet, + /// `async_delete_cache(key)`: answered with `None`. + #[strum(serialize = "delete")] + Delete, +} + +/// Python's outcome of a `cache` step. +#[derive(Clone, Debug, PartialEq)] +pub(crate) enum CacheOutcome { + /// `get`: the stored value (`None` when missing). + Value(Option), + /// `batch_get`: one entry per key. + Values(Vec>), + /// `increment`: the counter after the increment. + Number(f64), + /// `set`, `delete` and `increment_usage`. + Stored, +} + +impl CacheOutcome { + /// The op this outcome answers (`Stored` answers both `set` and `delete`). + fn answers(&self, op: CacheOp) -> bool { + matches!( + (self, op), + (Self::Value(_), CacheOp::Get) + | (Self::Values(_), CacheOp::BatchGet) + | (Self::Number(_), CacheOp::Increment) + | ( + Self::Stored, + CacheOp::Set | CacheOp::Delete | CacheOp::IncrementUsage + ) + ) + } +} + +/// A fallback chain's dict targets, which can carry an api key. +#[derive(Clone, PartialEq, veil::Redact)] +pub(crate) struct DictTargets(#[redact(with = "[REDACTED]")] pub(crate) Vec>); + +/// What a call is blocked on, for Python to run. +#[derive(Clone, Debug, PartialEq)] +pub(crate) enum HostCall { + Attempt { + deployment_id: String, + /// The picked deployment as configured: `model_name`, `litellm_params`, `model_info` + /// (with its `id`). The host builds the request from it; it never looks the deployment up + /// in a Python router. + deployment: Value, + /// The name the attempt resolved: the requested name or the fallback target. + model: String, + model_group: String, + /// The `order` level of a same-group order fallback. + target_order: Option, + /// Every deployment `model` resolved to for this caller, the picked one included. + candidate_ids: Vec, + retry: u32, + fallback_depth: u32, + fallback_level: u32, + timeout: Option, + /// Request params the fallback walk's dict targets merged in. + params: Option>, + /// The walk level whose request metadata the attempt carries. + scope: u32, + max_retries: u32, + model_group_size: usize, + }, + Sleep { + seconds: f64, + }, + AllowFallback { + target: String, + }, + RefuseFallbacks { + /// The walk level whose chain the targets are. + level: u32, + targets: DictTargets, + }, + Select { + stage: Stage, + model_group: String, + /// The name the request resolved from, which strategies key their state by. + name: String, + /// Whose selector a strategy reads: `default`, `group` (with `group`) or `override`. + scope: &'static str, + group: Option, + /// The deployments the name resolved to, before health filters (the tag stage reads them). + pool: Vec, + candidates: Vec, + }, + /// Pre-routing before a pick resolves `model`: answered with the model group to resolve + /// instead (a string), a custom strategy's deployment document (a mapping), or `None`. + PreRoute { + model: String, + fallback_depth: u32, + retry: u32, + }, + /// One [`RouterCache`] op against the host's cache. `key` is set for every op but + /// `batch_get`, which sets `keys`; `value` is set for `set` (the JSON to store) and + /// `increment` (the amount); `ttl` is in seconds. + Cache { + op: CacheOp, + key: Option, + keys: Option>, + value: Option, + ttl: Option, + /// The op reads or writes this worker's in-memory tier only (`local_only=True`). + local: bool, + }, + /// The deployment build's `os.environ/` references: answered with one value per ref. + Secrets { + index: usize, + refs: Vec, + }, + /// The deployment build's `validate`, `provider` and `added` steps for one deployment. + Build { + kind: StepKind, + path: BuildPath, + index: usize, + entry: Value, + declared_id: Option, + duplicate_id: bool, + /// The provider the `provider` step resolved (`added` only). + provider: Option, + }, + /// One lookup of a model-info query ([`ModelInfoLookup`], as its `lookup` field and + /// arguments): answered as the lookup documents. + ModelInfo { + lookup: Value, + }, +} + +impl HostCall { + pub(crate) fn kind(&self) -> StepKind { + match self { + Self::Attempt { .. } => StepKind::Attempt, + Self::Sleep { .. } => StepKind::Sleep, + Self::AllowFallback { .. } => StepKind::AllowFallback, + Self::RefuseFallbacks { .. } => StepKind::RefuseFallbacks, + Self::Select { .. } => StepKind::Select, + Self::PreRoute { .. } => StepKind::PreRoute, + Self::Cache { .. } => StepKind::Cache, + Self::Secrets { .. } => StepKind::Secrets, + Self::Build { kind, .. } => *kind, + Self::ModelInfo { .. } => StepKind::ModelInfo, + } + } +} + +/// Python's response to a [`HostCall`]. +#[derive(Debug)] +pub(crate) enum HostCallResponse { + /// The attempt's total tokens, or its failure. + Attempt(Result), + /// A successful attempt whose response the provider blocked, or a stream whose outcome the + /// host reports when it ends (`deferred`), with its tokens. + Served { + total_tokens: u64, + blocked: Option, + deferred: bool, + tracks_usage: Option, + }, + Slept, + AllowFallback(bool), + /// Whether the chain a `refuse_fallbacks` step named is refused. + RefuseFallbacks(bool), + /// The candidate ids that remain, pick first. + Select(Vec), + /// The model group to resolve instead, a custom strategy's deployment, or `None`. + PreRoute(Option), + /// The facts a `pre_call` select step asked for. + PreCall(PreCallFacts), + /// The candidate models a `routing_plugins` select step reported. + Models(Option>), + /// The ids a `callback_filters` select step left, or the error a filter raised. + Filtered(Result, AttemptError>), + /// What a pick's select or pre-route step, or a fallback check, raised on the host instead of + /// answering: the walk takes it as the try's failure. + Raised(AttemptError), + Cache(CacheOutcome), + /// One value per ref of a `secrets` step. + Secrets(Vec), + Validated(Built), + Provided(Provided), + Added(Built), + /// A `model_info` lookup's response, as the host spelled it. + ModelInfo(Value), +} + +impl HostCallResponse { + fn kind(&self) -> StepKind { + match self { + Self::Attempt(_) | Self::Served { .. } => StepKind::Attempt, + // Answers whichever step raised it (`check`). + Self::Raised(_) => StepKind::Select, + Self::Slept => StepKind::Sleep, + Self::AllowFallback(_) => StepKind::AllowFallback, + Self::RefuseFallbacks(_) => StepKind::RefuseFallbacks, + Self::Select(_) | Self::PreCall(_) | Self::Models(_) | Self::Filtered(_) => { + StepKind::Select + } + Self::PreRoute(_) => StepKind::PreRoute, + Self::Cache(_) => StepKind::Cache, + Self::Secrets(_) => StepKind::Secrets, + Self::Validated(_) => StepKind::Validate, + Self::Provided(_) => StepKind::Provider, + Self::Added(_) => StepKind::Added, + Self::ModelInfo(_) => StepKind::ModelInfo, + } + } + + fn check(&self, call: &HostCall) -> Result<(), DriveError> { + if matches!(self, Self::Raised(_)) + && matches!( + call.kind(), + StepKind::Select | StepKind::PreRoute | StepKind::AllowFallback + ) + { + return Ok(()); + } + if self.kind() != call.kind() { + return Err(DriveError::WrongAnswer { + expected: call.kind(), + got: self.kind(), + }); + } + if let (Self::Select(ids) | Self::Filtered(Ok(ids)), HostCall::Select { candidates, .. }) = + (self, call) + && let Some(id) = ids.iter().find(|id| !candidates.contains(id)) + { + return Err(DriveError::UnknownCandidate { id: id.clone() }); + } + if let (Self::Secrets(values), HostCall::Secrets { refs, .. }) = (self, call) + && values.len() != refs.len() + { + return Err(DriveError::SecretsLength { + expected: refs.len(), + got: values.len(), + }); + } + if let (Self::Cache(outcome), HostCall::Cache { op, keys, .. }) = (self, call) { + if !outcome.answers(*op) { + return Err(DriveError::WrongCacheOutcome { op: *op }); + } + if let (CacheOutcome::Values(values), Some(keys)) = (outcome, keys) + && values.len() != keys.len() + { + return Err(DriveError::BatchLength { + expected: keys.len(), + got: values.len(), + }); + } + } + Ok(()) + } +} + +/// How a driven future ended well. +#[derive(Debug)] +pub(crate) enum Finish { + /// `start`: the call's response. + Call(RouterResponse), + /// `select`: the deployment picked and the model group it serves. + Selected { + deployment_id: String, + model_group: String, + /// The deployment as served: a priority group's member carries its priority as `order`. + deployment: Value, + }, + /// `healthy`: the deployments a pick would choose among, as served. + Healthy { + /// A deployment named by its id, which the listing returns alone rather than as a list. + single: bool, + deployment_ids: Vec, + deployments: Vec, + }, + /// `report_failure`: the cooldown applied, if the deployment was cooled down. + Reported { cooldown: Option }, + /// `report_success` / `count_success`: the success counter's key, if one was counted. + Counted { key: Option }, + /// `build`: the deployments it installed, in order. + Built(Vec), + /// `router_model_info`: the merged model info (`None` when the lookup answered `None`). + RouterModelInfo(Option), + /// `deployment_model_info`: the merged model info, if any. + DeploymentModelInfo(Option>), + /// A model-info query that ends with one JSON answer (`model_group_info`, + /// `configured_token_limits`, `model_listing_info`). + Value(Value), +} + +/// One step of a call: what it is blocked on, or how it ended. +#[derive(Debug)] +pub(crate) enum Step { + HostCall(HostCall), + Done(Finish), + Failed(Error), +} + +/// A misuse of the step protocol by the caller. +#[derive(Debug, thiserror::Error)] +pub(crate) enum DriveError { + #[error("the router call already finished")] + Finished, + #[error("the router call has not made a host call yet, so the outcome must be None")] + UnexpectedResponse, + #[error("the router call is waiting for an outcome of its `{expected}` step")] + MissingResponse { expected: StepKind }, + #[error("the router call is waiting for an outcome of its `{expected}` step, not `{got}`")] + WrongAnswer { expected: StepKind, got: StepKind }, + #[error("deployment `{id}` was not one of the select step's candidates")] + UnknownCandidate { id: String }, + #[error("the outcome does not answer a `{op}` cache step")] + WrongCacheOutcome { op: CacheOp }, + #[error("a `batch_get` cache step for {expected} keys was answered with {got} values")] + BatchLength { expected: usize, got: usize }, + #[error("a `secrets` step for {expected} references was answered with {got} values")] + SecretsLength { expected: usize, got: usize }, + #[error("the router call is pending on something other than a step")] + Stalled, +} + +impl DriveError { + /// Whether the caller passed a bad outcome, rather than driving a call that cannot go on. + pub(crate) fn is_request(&self) -> bool { + match self { + Self::UnexpectedResponse + | Self::MissingResponse { .. } + | Self::WrongAnswer { .. } + | Self::UnknownCandidate { .. } + | Self::WrongCacheOutcome { .. } + | Self::BatchLength { .. } + | Self::SecretsLength { .. } => true, + Self::Finished | Self::Stalled => false, + } + } +} + +/// The host call a router call is blocked on and the response that releases it. +#[derive(Default)] +struct Exchange { + call: Option, + response: Option, + /// Events the call reported since the last step, handed to Python with the next one. + events: Vec, +} + +type Shared = Arc>; + +fn lock(exchange: &Mutex) -> MutexGuard<'_, Exchange> { + exchange.lock().unwrap_or_else(PoisonError::into_inner) +} + +thread_local! { + static CURRENT: RefCell> = const { RefCell::new(None) }; +} + +/// Binds an exchange to this thread for one poll and restores the previous one after. +struct Bound(Option); + +impl Bound { + fn new(exchange: &Shared) -> Self { + Self(CURRENT.replace(Some(exchange.clone()))) + } +} + +impl Drop for Bound { + fn drop(&mut self) { + CURRENT.set(self.0.take()); + } +} + +/// Park the host call `call` in the bound exchange and wait for its response. +/// `response`, or the failure a step answered with instead ([`HostCallResponse::Raised`]). +fn answered(response: HostCallResponse) -> Result { + match response { + HostCallResponse::Raised(failure) => Err(failure), + response => Ok(response), + } +} + +async fn call_host(call: HostCall) -> HostCallResponse { + let mut call = Some(call); + poll_fn(move |_| { + CURRENT.with_borrow(|current| { + let Some(exchange) = current else { + return Poll::Pending; + }; + let mut exchange = lock(exchange); + if let Some(call) = call.take() { + exchange.call = Some(call); + return Poll::Pending; + } + exchange.response.take().map_or(Poll::Pending, Poll::Ready) + }) + }) + .await +} + +const CHECKED: &str = "the driver checks each response against its host call"; + +fn ttl_seconds(ttl: Option) -> Option { + ttl.map(|ttl| ttl.as_secs_f64()) +} + +/// An increment amount as Python's counters store it: an int when it is whole. +fn amount_value(amount: f64) -> Value { + if amount.fract() == 0.0 && amount.abs() < i64::MAX as f64 { + Value::from(amount as i64) + } else { + Value::from(amount) + } +} + +/// The host's cache: every op is a `cache` step Python runs against the router's `DualCache`. +pub(crate) struct StepCache; + +async fn cache_step( + op: CacheOp, + key: Option<&str>, + keys: Option<&[String]>, + value: Option, + ttl: Option, +) -> CacheOutcome { + call_cache(op, key, keys, value, ttl, false).await +} + +async fn call_cache( + op: CacheOp, + key: Option<&str>, + keys: Option<&[String]>, + value: Option, + ttl: Option, + local: bool, +) -> CacheOutcome { + let call = HostCall::Cache { + op, + key: key.map(str::to_owned), + keys: keys.map(<[String]>::to_vec), + value, + ttl: ttl_seconds(ttl), + local, + }; + let HostCallResponse::Cache(outcome) = call_host(call).await else { + unreachable!("{CHECKED}"); + }; + outcome +} + +#[async_trait] +impl RouterCache for StepCache { + async fn async_get_cache(&self, key: &str) -> Result, litellm_cache::Error> { + let CacheOutcome::Value(value) = + cache_step(CacheOp::Get, Some(key), None, None, None).await + else { + unreachable!("{CHECKED}"); + }; + Ok(value) + } + + async fn async_set_cache( + &self, + key: &str, + value: Value, + ttl: Option, + ) -> Result<(), litellm_cache::Error> { + cache_step(CacheOp::Set, Some(key), None, Some(value), ttl).await; + Ok(()) + } + + async fn async_increment( + &self, + key: &str, + amount: f64, + ttl: Option, + ) -> Result { + let value = Some(amount_value(amount)); + let CacheOutcome::Number(next) = + cache_step(CacheOp::Increment, Some(key), None, value, ttl).await + else { + unreachable!("{CHECKED}"); + }; + Ok(next) + } + + async fn async_batch_get_cache( + &self, + keys: &[String], + ) -> Result>, litellm_cache::Error> { + let CacheOutcome::Values(values) = + cache_step(CacheOp::BatchGet, None, Some(keys), None, None).await + else { + unreachable!("{CHECKED}"); + }; + Ok(values) + } + + async fn async_delete_cache(&self, key: &str) -> Result<(), litellm_cache::Error> { + cache_step(CacheOp::Delete, Some(key), None, None, None).await; + Ok(()) + } + + async fn async_increment_usage( + &self, + increments: &[(String, f64)], + ttl: Option, + ) -> Result<(), litellm_cache::Error> { + let keys: Vec = increments.iter().map(|(key, _)| key.clone()).collect(); + let amounts = increments + .iter() + .map(|(_, amount)| amount_value(*amount)) + .collect(); + let value = Some(Value::Array(amounts)); + cache_step(CacheOp::IncrementUsage, None, Some(&keys), value, ttl).await; + Ok(()) + } + + async fn async_get_local(&self, key: &str) -> Result, litellm_cache::Error> { + let CacheOutcome::Value(value) = + call_cache(CacheOp::Get, Some(key), None, None, None, true).await + else { + unreachable!("{CHECKED}"); + }; + Ok(value) + } + + async fn async_increment_local( + &self, + key: &str, + amount: f64, + ttl: Option, + ) -> Result { + let value = Some(amount_value(amount)); + let CacheOutcome::Number(next) = + call_cache(CacheOp::Increment, Some(key), None, value, ttl, true).await + else { + unreachable!("{CHECKED}"); + }; + Ok(next) + } +} + +/// A deployment as Python spells a `model_list` entry, with its id in `model_info`. +pub(crate) fn deployment_record(deployment: &Deployment) -> Value { + deployment.record() +} + +/// The process's `random` module: a Python router draws its shuffle and retry jitter there, so +/// a seeded process picks the same deployments on either core. The bare module functions share +/// one generator, so these calls interleave with every other draw the process makes. +pub(crate) struct PythonRandomness; + +impl Randomness for PythonRandomness { + fn next_f64(&self) -> f64 { + Python::attach(|py| -> PyResult { + py.import("random")?.call_method0("random")?.extract() + }) + .unwrap_or_else(|_| ThreadRandomness.next_f64()) + } + + fn index_below(&self, len: usize) -> usize { + // `random.choice` is `seq[self._randbelow(len(seq))]`, and so is `randrange(len)`. + Python::attach(|py| -> PyResult { + py.import("random")? + .call_method1("randrange", (len,))? + .extract() + }) + .unwrap_or_else(|_| ThreadRandomness.index_below(len)) + } +} + +/// The [`Host`] every `RouterCoreHandle` router runs on. +pub(crate) struct StepHost; + +#[async_trait] +impl Host for StepHost { + async fn invoke(&self, attempt: &Attempt) -> Result { + let call = HostCall::Attempt { + deployment_id: attempt.deployment.id.clone(), + deployment: deployment_record(&attempt.deployment), + model: attempt.name.clone(), + model_group: attempt.model_group.clone(), + target_order: attempt.target_order, + candidate_ids: attempt.candidate_ids.clone(), + retry: attempt.retry, + fallback_depth: attempt.fallback_depth, + fallback_level: attempt.fallback_level, + timeout: attempt.timeout.map(|timeout| timeout.as_secs_f64()), + params: attempt + .params + .as_ref() + .map(|params| Box::new(Value::Object(Map::clone(params)))), + scope: attempt.scope, + max_retries: attempt.max_retries, + model_group_size: attempt.model_group_size, + }; + let (result, blocked, deferred, tracks_usage) = match call_host(call).await { + HostCallResponse::Attempt(result) => (result, None, false, None), + HostCallResponse::Served { + total_tokens, + blocked, + deferred, + tracks_usage, + } => (Ok(total_tokens), blocked, deferred, tracks_usage), + _ => unreachable!("{CHECKED}"), + }; + result.map(|total_tokens| Output { + value: Payload(Arc::new(())), + total_tokens, + blocked, + deferred, + tracks_usage, + }) + } + + async fn select( + &self, + stage: &SelectStage, + ctx: &RoutingContext<'_>, + candidates: Vec>, + ) -> Result>, Error> { + let (scope, group) = ctx.scope.wire(); + let call = HostCall::Select { + stage: match stage { + SelectStage::Plugins => Stage::Plugins, + SelectStage::Strategy(strategy) => Stage::Strategy(strategy.to_string()), + SelectStage::Tags => Stage::Tags, + SelectStage::TextOnly => Stage::TextOnly, + SelectStage::RouteIncapable { handler } => Stage::RouteIncapable { + handler: handler.clone(), + }, + SelectStage::MediaRefusing => Stage::MediaRefusing, + }, + pool: ctx.pool.iter().map(|d| d.id.clone()).collect(), + model_group: ctx.model_group.to_owned(), + name: ctx.name.to_owned(), + scope, + group: group.map(str::to_owned), + candidates: candidates.iter().map(|d| d.id.clone()).collect(), + }; + let HostCallResponse::Select(ids) = answered(call_host(call).await)? else { + unreachable!("{CHECKED}"); + }; + Ok(ids + .iter() + .filter_map(|id| candidates.iter().find(|d| d.id == *id).cloned()) + .collect()) + } + + async fn pre_route(&self, pre: &PreRoute<'_>) -> Result, Error> { + let call = HostCall::PreRoute { + model: pre.name.to_owned(), + fallback_depth: pre.fallback_depth, + retry: pre.retry, + }; + let HostCallResponse::PreRoute(routed) = answered(call_host(call).await)? else { + unreachable!("{CHECKED}"); + }; + Ok(routed) + } + + async fn routing_plugin_models( + &self, + ctx: &RoutingContext<'_>, + run: bool, + ) -> Result>, Error> { + let (scope, group) = ctx.scope.wire(); + let call = HostCall::Select { + stage: Stage::RoutingPlugins { run }, + pool: Vec::new(), + model_group: ctx.model_group.to_owned(), + name: ctx.name.to_owned(), + scope, + group: group.map(str::to_owned), + candidates: Vec::new(), + }; + let HostCallResponse::Models(models) = answered(call_host(call).await)? else { + unreachable!("{CHECKED}"); + }; + Ok(models) + } + + async fn provider_request(&self, ctx: &RoutingContext<'_>) -> Result, Error> { + let (scope, group) = ctx.scope.wire(); + let call = HostCall::Select { + stage: Stage::ProviderRequest, + pool: Vec::new(), + model_group: ctx.model_group.to_owned(), + name: ctx.name.to_owned(), + scope, + group: group.map(str::to_owned), + candidates: Vec::new(), + }; + let HostCallResponse::Models(models) = answered(call_host(call).await)? else { + unreachable!("{CHECKED}"); + }; + Ok(models.into_iter().flatten().next()) + } + + async fn pre_call_facts( + &self, + ctx: &RoutingContext<'_>, + candidates: &[Arc], + ) -> Result { + let (scope, group) = ctx.scope.wire(); + let call = HostCall::Select { + stage: Stage::PreCall, + pool: ctx.pool.iter().map(|d| d.id.clone()).collect(), + model_group: ctx.model_group.to_owned(), + name: ctx.name.to_owned(), + scope, + group: group.map(str::to_owned), + candidates: candidates.iter().map(|d| d.id.clone()).collect(), + }; + let HostCallResponse::PreCall(facts) = answered(call_host(call).await)? else { + unreachable!("{CHECKED}"); + }; + Ok(facts) + } + + async fn callback_filters( + &self, + ctx: &RoutingContext<'_>, + candidates: Vec>, + ) -> Result>, Error> { + let (scope, group) = ctx.scope.wire(); + let joined = |own: &[String], walked: &[String]| [own, walked].concat(); + let call = HostCall::Select { + stage: Stage::CallbackFilters { + excluded_ids: joined(&ctx.request.excluded_ids, ctx.excluded_ids), + skipped_ids: joined(&ctx.request.retry_skipped_ids, ctx.skipped_ids), + target_order: ctx.target_order, + }, + pool: Vec::new(), + model_group: ctx.model_group.to_owned(), + name: ctx.name.to_owned(), + scope, + group: group.map(str::to_owned), + candidates: candidates.iter().map(|d| d.id.clone()).collect(), + }; + let HostCallResponse::Filtered(result) = call_host(call).await else { + unreachable!("{CHECKED}"); + }; + let ids = result.map_err(Error::Attempt)?; + Ok(ids + .iter() + .filter_map(|id| candidates.iter().find(|d| d.id == *id).cloned()) + .collect()) + } + + async fn allow_fallback( + &self, + _request: &RouterRequest, + target: &str, + ) -> Result { + let call = HostCall::AllowFallback { + target: target.to_owned(), + }; + let HostCallResponse::AllowFallback(allowed) = answered(call_host(call).await)? else { + unreachable!("{CHECKED}"); + }; + Ok(allowed) + } + + async fn refuses_fallbacks(&self, level: u32, targets: &[Map]) -> bool { + let call = HostCall::RefuseFallbacks { + level, + targets: DictTargets(targets.to_vec()), + }; + let HostCallResponse::RefuseFallbacks(refused) = call_host(call).await else { + unreachable!("{CHECKED}"); + }; + refused + } + + /// A zero sleep is asked too: Python's retry loop awaits `asyncio.sleep(0)` between sibling + /// retries, which yields to the event loop. + async fn sleep(&self, duration: Duration) { + let call = HostCall::Sleep { + seconds: duration.as_secs_f64(), + }; + let HostCallResponse::Slept = call_host(call).await else { + unreachable!("{CHECKED}"); + }; + } + + /// Kept for Python only when it acts on it: an exhausted retry loop is stamped on the + /// error Python raises, a fallback hop tells it which target a failed walk ended on, and a + /// cooldown or a pick is reported to the host's Prometheus and service loggers (a pick also + /// joins the router cache's per-deployment count, a failed walk the router's error log). + fn observe(&self, event: &RouterEvent) { + if !matches!( + event, + RouterEvent::RetriesExhausted { .. } + | RouterEvent::Retry { .. } + | RouterEvent::Fallback { .. } + | RouterEvent::FallingBack { .. } + | RouterEvent::OrderChainFailed { .. } + | RouterEvent::FallbacksFailed { .. } + | RouterEvent::LevelFailed { .. } + | RouterEvent::PickFailed { .. } + | RouterEvent::RetriesOpened { .. } + | RouterEvent::RetryBudget { .. } + | RouterEvent::MockFailed { .. } + | RouterEvent::RetryLogged { .. } + | RouterEvent::RetryPolicyRaised { .. } + | RouterEvent::RetryCountRaised { .. } + | RouterEvent::FallbackDone { .. } + | RouterEvent::Blocked { .. } + | RouterEvent::CooldownSet { .. } + | RouterEvent::Selected { .. } + ) { + return; + } + CURRENT.with_borrow(|current| { + if let Some(exchange) = current { + lock(exchange).events.push(event.clone()); + } + }); + } + + fn cache(&self) -> Arc { + Arc::new(StepCache) + } + + async fn secrets(&self, index: usize, refs: &[SecretRef]) -> Result, Error> { + let HostCallResponse::Secrets(values) = call_host(HostCall::Secrets { + index, + refs: refs.to_vec(), + }) + .await + else { + unreachable!("{CHECKED}"); + }; + Ok(values) + } + + async fn validate(&self, build: &BuildAsk<'_>) -> Result { + let HostCallResponse::Validated(built) = + call_host(build_host_call(StepKind::Validate, build, None)).await + else { + unreachable!("{CHECKED}"); + }; + Ok(built) + } + + async fn provider(&self, build: &BuildAsk<'_>) -> Result { + let HostCallResponse::Provided(provided) = + call_host(build_host_call(StepKind::Provider, build, None)).await + else { + unreachable!("{CHECKED}"); + }; + Ok(provided) + } + + async fn added(&self, build: &BuildAsk<'_>, provider: Option<&str>) -> Result { + let HostCallResponse::Added(built) = + call_host(build_host_call(StepKind::Added, build, provider)).await + else { + unreachable!("{CHECKED}"); + }; + Ok(built) + } + + async fn model_info(&self, lookup: &ModelInfoLookup) -> Result { + let lookup = serde_json::to_value(lookup).unwrap_or(Value::Null); + let HostCallResponse::ModelInfo(response) = call_host(HostCall::ModelInfo { lookup }).await + else { + unreachable!("{CHECKED}"); + }; + Ok(response) + } +} + +fn build_host_call(kind: StepKind, build: &BuildAsk<'_>, provider: Option<&str>) -> HostCall { + HostCall::Build { + kind, + path: build.path, + index: build.index, + entry: Value::Object(build.entry.clone()), + declared_id: build.declared_id.map(str::to_owned), + duplicate_id: build.duplicate_id, + provider: provider.map(str::to_owned), + } +} + +type CallFuture = Pin> + Send>>; + +/// Poll `future` once with `exchange` bound. +fn poll_bound(future: Pin<&mut F>, exchange: &Shared) -> Poll { + let _bound = Bound::new(exchange); + future.poll(&mut Context::from_waker(Waker::noop())) +} + +/// Drives one router call a step at a time. +pub(crate) struct Driver { + future: Option, + exchange: Shared, + pending_call: Option, +} + +impl Driver { + pub(crate) fn new(call: impl Future> + Send + 'static) -> Self { + Self { + future: Some(Box::pin(call)), + exchange: Shared::default(), + pending_call: None, + } + } + + /// The host call the router call is blocked on, if it made one. + pub(crate) fn pending_call(&self) -> Option<&HostCall> { + self.pending_call.as_ref() + } + + /// The events reported since the last call. + pub(crate) fn take_events(&self) -> Vec { + std::mem::take(&mut lock(&self.exchange).events) + } + + /// Hand back the response to the pending host call (`None` on the first call) and run the call to + /// its next step or its end. + pub(crate) fn resume( + &mut self, + response: Option, + ) -> Result { + let Some(future) = self.future.as_mut() else { + return Err(DriveError::Finished); + }; + match (&self.pending_call, response) { + (None, None) => {} + (None, Some(_)) => return Err(DriveError::UnexpectedResponse), + (Some(call), None) => { + return Err(DriveError::MissingResponse { + expected: call.kind(), + }); + } + (Some(call), Some(response)) => { + response.check(call)?; + lock(&self.exchange).response = Some(response); + } + } + self.pending_call = None; + match poll_bound(future.as_mut(), &self.exchange) { + Poll::Ready(result) => { + self.future = None; + Ok(match result { + Ok(finish) => Step::Done(finish), + Err(error) => Step::Failed(error), + }) + } + Poll::Pending => { + let Some(call) = lock(&self.exchange).call.take() else { + self.future = None; + return Err(DriveError::Stalled); + }; + self.pending_call = Some(call.clone()); + Ok(Step::HostCall(call)) + } + } + } +} + +#[cfg(test)] +mod tests { + use std::{ + collections::HashMap, + time::{SystemTime, UNIX_EPOCH}, + }; + + use litellm_router_core::{ + FailureKind, Op, Router, RouterConfig, RouterDynamicConfig, SeededRandomness, + }; + use rstest::{fixture, rstest}; + use serde_json::{Value, json}; + + use super::*; + + fn router(config: Value, dynamic: Value) -> Router { + let config: RouterConfig = serde_json::from_value(config).unwrap(); + let dynamic: RouterDynamicConfig = serde_json::from_value(dynamic).unwrap(); + Router::builder(config, dynamic, Arc::new(StepHost)) + .randomness(Arc::new(SeededRandomness::new(7))) + .build() + .unwrap() + } + + fn deployment(group: &str, id: &str) -> Value { + json!({"model_name": group, "litellm_params": {"model": format!("openai/{id}")}, "model_info": {"id": id}}) + } + + fn start(router: &Router, model: &str) -> Driver { + let router = router.clone(); + let request = RouterRequest { + model: model.to_owned(), + ..RouterRequest::default() + }; + Driver::new(async move { + router + .call(Op::from("acompletion"), request) + .await + .map(Finish::Call) + }) + } + + fn rate_limited() -> HostCallResponse { + HostCallResponse::Attempt(Err(AttemptError::new(FailureKind::RateLimit, Some(429)))) + } + + /// A dict-backed stand-in for Python's `DualCache` that answers `cache` steps and keeps + /// every one it answered. + #[derive(Default)] + struct Store { + entries: HashMap, + calls: Vec, + } + + impl Store { + fn respond(&mut self, call: &HostCall) -> HostCallResponse { + let HostCall::Cache { + op, + key, + keys, + value, + .. + } = call + else { + panic!("not a cache step: {call:?}"); + }; + self.calls.push(call.clone()); + let key = key.clone().unwrap_or_default(); + HostCallResponse::Cache(match op { + CacheOp::Get => CacheOutcome::Value(self.entries.get(&key).cloned()), + CacheOp::Set => { + self.entries.insert(key, value.clone().unwrap()); + CacheOutcome::Stored + } + CacheOp::Increment => { + let current = self.entries.get(&key).and_then(Value::as_f64); + let next = + current.unwrap_or(0.0) + value.as_ref().and_then(Value::as_f64).unwrap(); + self.entries.insert(key, amount_value(next)); + CacheOutcome::Number(next) + } + CacheOp::IncrementUsage => { + let amounts = value.as_ref().and_then(Value::as_array).unwrap(); + for (key, amount) in keys.iter().flatten().zip(amounts) { + let current = self.entries.get(key).and_then(Value::as_f64); + let next = current.unwrap_or(0.0) + amount.as_f64().unwrap(); + self.entries.insert(key.clone(), amount_value(next)); + } + CacheOutcome::Stored + } + CacheOp::BatchGet => CacheOutcome::Values( + keys.iter() + .flatten() + .map(|key| self.entries.get(key).cloned()) + .collect(), + ), + CacheOp::Delete => { + self.entries.remove(&key); + CacheOutcome::Stored + } + }) + } + + /// Resume with `response`, answering every cache step on the way, up to the next other + /// step. + fn resume(&mut self, driver: &mut Driver, response: Option) -> Step { + let mut step = driver.resume(response).unwrap(); + while let Step::HostCall(call @ HostCall::Cache { .. }) = &step { + let response = self.respond(call); + step = driver.resume(Some(response)).unwrap(); + } + step + } + } + + #[fixture] + fn pair() -> Router { + router( + json!({"model_list": [deployment("primary", "a1"), deployment("primary", "a2")]}), + json!({}), + ) + } + + #[rstest] + fn a_first_attempt_success_is_done(pair: Router) { + let mut store = Store::default(); + let mut driver = start(&pair, "primary"); + let Step::HostCall(HostCall::Attempt { + deployment_id, + model_group, + retry: 0, + fallback_depth: 0, + .. + }) = store.resume(&mut driver, None) + else { + panic!("the call should start with an attempt"); + }; + assert_eq!(model_group, "primary"); + let Step::Done(Finish::Call(response)) = + store.resume(&mut driver, Some(HostCallResponse::Attempt(Ok(12)))) + else { + panic!("a successful attempt should finish the call"); + }; + assert_eq!(response.deployment_id, deployment_id); + assert_eq!(response.output.total_tokens, 12); + assert_eq!( + (response.attempted_retries, response.attempted_fallbacks), + (0, 0) + ); + assert!(matches!(driver.resume(None), Err(DriveError::Finished))); + } + + #[rstest] + fn cooldown_reads_are_cache_steps(pair: Router) { + let mut driver = start(&pair, "primary"); + let step = driver.resume(None).unwrap(); + let Step::HostCall(HostCall::Cache { + op: CacheOp::BatchGet, + key: None, + keys: Some(keys), + value: None, + ttl: None, + local: false, + }) = step + else { + panic!("the call should read cooldowns first, got {step:?}"); + }; + assert_eq!(keys, ["deployment:a1:cooldown", "deployment:a2:cooldown"]); + } + + #[rstest] + #[case::lone_deployment_backs_off(vec![deployment("primary", "a1")], true)] + #[case::sibling_retries_at_once( + vec![deployment("primary", "a1"), deployment("primary", "a2")], + false + )] + fn a_rate_limit_retries_after_the_backoff( + #[case] model_list: Vec, + #[case] backs_off: bool, + ) { + let router = router(json!({"model_list": model_list}), json!({"num_retries": 1})); + let mut store = Store::default(); + let mut driver = start(&router, "primary"); + assert_eq!( + store.resume(&mut driver, None).unwrap_kind(), + StepKind::Attempt + ); + let step = store.resume(&mut driver, Some(rate_limited())); + let Step::HostCall(HostCall::Sleep { seconds }) = step else { + panic!("a retry should sleep first, got {step:?}"); + }; + assert_eq!(seconds > 0.0, backs_off, "slept {seconds}s"); + let step = store.resume(&mut driver, Some(HostCallResponse::Slept)); + assert!( + matches!(step, Step::HostCall(HostCall::Attempt { retry: 1, .. })), + "expected the retry, got {step:?}" + ); + let Step::Done(Finish::Call(response)) = + store.resume(&mut driver, Some(HostCallResponse::Attempt(Ok(1)))) + else { + panic!("the retry should finish the call"); + }; + assert_eq!(response.attempted_retries, 1); + } + + #[rstest] + fn a_cooldown_python_wrote_is_read_back_by_the_next_call() { + let pair_without_retries = router( + json!({"model_list": [deployment("primary", "a1"), deployment("primary", "a2")]}), + json!({"num_retries": 0}), + ); + let mut store = Store::default(); + let now = SystemTime::now() + .duration_since(UNIX_EPOCH) + .unwrap() + .as_secs_f64(); + store.entries.insert( + "deployment:a1:cooldown".into(), + json!({"exception_received": "slow down", "status_code": "429", "timestamp": now, "cooldown_time": 60}), + ); + for _ in 0..8 { + let mut call = start(&pair_without_retries, "primary"); + let Step::HostCall(HostCall::Attempt { deployment_id, .. }) = + store.resume(&mut call, None) + else { + panic!("the call should attempt the healthy deployment"); + }; + assert_eq!(deployment_id, "a2", "the cooled-down deployment was picked"); + } + } + + /// The host only runs the attempt; the core records its outcome (cooldowns, fail counters, + /// usage) through the host's cache. + #[rstest] + #[case::failure(rate_limited())] + #[case::success(HostCallResponse::Attempt(Ok(7)))] + fn attempt_outcomes_are_recorded_by_the_core(#[case] outcome: HostCallResponse) { + let tracked = |id: &str| { + json!({ + "model_name": "primary", + "litellm_params": {"model": format!("openai/{id}"), "tpm": 100, "cooldown_time": 7}, + "model_info": {"id": id, "allowed_fails": 3}, + }) + }; + let router = router( + json!({"model_list": [tracked("a1"), tracked("a2")]}), + json!({"num_retries": 0}), + ); + let mut store = Store::default(); + let mut driver = start(&router, "primary"); + store.resume(&mut driver, None); + let step = store.resume(&mut driver, Some(outcome)); + assert!( + matches!(step, Step::Done(_) | Step::Failed(_)), + "got {step:?}" + ); + let writes: Vec<&HostCall> = store + .calls + .iter() + .filter(|call| { + matches!(call, HostCall::Cache { op, .. } if !matches!(op, CacheOp::Get | CacheOp::BatchGet)) + }) + .collect(); + assert!(!writes.is_empty(), "the core wrote no outcome"); + } + + #[rstest] + fn a_failed_group_falls_back_once_python_allows_it() { + let router = router( + json!({"model_list": [deployment("primary", "a1"), deployment("backup", "b1")]}), + json!({"fallbacks": [{"primary": ["backup"]}]}), + ); + let mut store = Store::default(); + let mut driver = start(&router, "primary"); + assert_eq!( + store.resume(&mut driver, None).unwrap_kind(), + StepKind::Attempt + ); + let bad_request = + HostCallResponse::Attempt(Err(AttemptError::new(FailureKind::BadRequest, Some(400)))); + let step = store.resume(&mut driver, Some(bad_request)); + assert!( + matches!(&step, Step::HostCall(HostCall::AllowFallback { target }) if target == "backup"), + "expected the fallback check, got {step:?}" + ); + let step = store.resume(&mut driver, Some(HostCallResponse::AllowFallback(true))); + assert!( + matches!(&step, Step::HostCall(HostCall::Attempt { deployment_id, fallback_depth: 1, .. }) if deployment_id == "b1"), + "expected the fallback attempt, got {step:?}" + ); + let Step::Done(Finish::Call(response)) = + store.resume(&mut driver, Some(HostCallResponse::Attempt(Ok(3)))) + else { + panic!("the fallback should finish the call"); + }; + assert_eq!(response.model_group, "backup"); + assert_eq!(response.attempted_fallbacks, 1); + } + + #[rstest] + fn a_refused_fallback_fails_with_the_attempt_error() { + let router = router( + json!({"model_list": [deployment("primary", "a1"), deployment("backup", "b1")]}), + json!({"fallbacks": [{"primary": ["backup"]}]}), + ); + let mut store = Store::default(); + let mut driver = start(&router, "primary"); + store.resume(&mut driver, None); + let bad_request = + HostCallResponse::Attempt(Err(AttemptError::new(FailureKind::BadRequest, Some(400)))); + store.resume(&mut driver, Some(bad_request)); + let Step::Failed(error) = + store.resume(&mut driver, Some(HostCallResponse::AllowFallback(false))) + else { + panic!("a refused fallback should fail the call"); + }; + assert_eq!( + error.attempt().map(|a| a.kind), + Some(FailureKind::BadRequest) + ); + } + + #[rstest] + #[case::none_while_waiting(None, StepKind::Attempt)] + #[case::wrong_kind(Some(HostCallResponse::Slept), StepKind::Attempt)] + #[case::cache_outcome(Some(HostCallResponse::Cache(CacheOutcome::Stored)), StepKind::Attempt)] + fn a_mismatched_outcome_is_rejected_and_the_step_kept( + pair: Router, + #[case] response: Option, + #[case] expected: StepKind, + ) { + let mut store = Store::default(); + let mut driver = start(&pair, "primary"); + store.resume(&mut driver, None); + let error = driver.resume(response).unwrap_err(); + assert!(error.is_request(), "{error}"); + assert_eq!(driver.pending_call().map(HostCall::kind), Some(expected)); + assert!(matches!( + store.resume(&mut driver, Some(HostCallResponse::Attempt(Ok(1)))), + Step::Done(_) + )); + } + + #[rstest] + #[case::wrong_op(HostCallResponse::Cache(CacheOutcome::Stored))] + #[case::number_for_batch(HostCallResponse::Cache(CacheOutcome::Number(1.0)))] + #[case::short_batch(HostCallResponse::Cache(CacheOutcome::Values(vec![None])))] + #[case::not_a_cache_outcome(HostCallResponse::Slept)] + fn a_mismatched_cache_outcome_is_rejected_and_the_step_kept( + pair: Router, + #[case] response: HostCallResponse, + ) { + let mut store = Store::default(); + let mut driver = start(&pair, "primary"); + let step = driver.resume(None).unwrap(); + let Step::HostCall( + call @ HostCall::Cache { + op: CacheOp::BatchGet, + .. + }, + ) = &step + else { + panic!("expected the cooldown read, got {step:?}"); + }; + let error = driver.resume(Some(response)).unwrap_err(); + assert!(error.is_request(), "{error}"); + assert_eq!(driver.pending_call(), Some(call)); + let response = store.respond(call); + assert_eq!( + store.resume(&mut driver, Some(response)).unwrap_kind(), + StepKind::Attempt + ); + } + + #[rstest] + fn a_strategy_python_runs_asks_for_a_selection_and_keeps_the_pick() { + let router = router( + json!({ + "model_list": [deployment("primary", "a1"), deployment("primary", "a2")], + "routing_strategy": "least-busy", + }), + json!({}), + ); + let mut store = Store::default(); + let mut driver = start(&router, "primary"); + let step = store.resume(&mut driver, None); + let Step::HostCall(HostCall::Select { + stage: Stage::Strategy(strategy), + candidates, + .. + }) = step + else { + panic!("a host strategy should ask for a selection, got {step:?}"); + }; + assert_eq!(strategy, "least-busy"); + assert_eq!(candidates, ["a1", "a2"]); + let step = store.resume( + &mut driver, + Some(HostCallResponse::Select(vec!["a2".to_owned()])), + ); + assert!( + matches!(&step, Step::HostCall(HostCall::Attempt { deployment_id, .. }) if deployment_id == "a2"), + "expected the picked deployment, got {step:?}" + ); + } + + #[rstest] + fn a_selection_is_driven_through_cache_steps(pair: Router) { + let request = RouterRequest { + model: "primary".to_owned(), + ..RouterRequest::default() + }; + let mut store = Store::default(); + let mut driver = Driver::new(async move { + pair.select(&Op::from("acompletion"), &request) + .await + .map(|deployment| Finish::Selected { + deployment_id: deployment.id.clone(), + model_group: deployment.model_name.clone(), + deployment: deployment_record(&deployment), + }) + }); + let Step::Done(Finish::Selected { + deployment_id, + model_group, + .. + }) = store.resume(&mut driver, None) + else { + panic!("a selection should finish with the pick"); + }; + assert!(["a1", "a2"].contains(&deployment_id.as_str())); + assert_eq!(model_group, "primary"); + assert!(matches!( + store.calls.as_slice(), + [HostCall::Cache { + op: CacheOp::BatchGet, + .. + }] + )); + } + + #[rstest] + fn dict_targets_keep_their_api_key_out_of_debug() { + let target = json!({"model": "b", "api_key": "credential-value"}); + let Value::Object(target) = target else { + unreachable!("a JSON object"); + }; + let shown = format!("{:?}", DictTargets(vec![target])); + assert!(!shown.contains("credential-value")); + assert!(shown.contains("[REDACTED]")); + } + + #[rstest] + fn a_chain_with_a_dict_target_asks_python_before_any_hop() { + let router = router( + json!({"model_list": [deployment("primary", "a1"), deployment("backup", "b1")]}), + json!({"fallbacks": [{"primary": [{"model": "backup", "api_key": "k"}]}]}), + ); + let mut store = Store::default(); + let mut driver = start(&router, "primary"); + store.resume(&mut driver, None); + let bad_request = + HostCallResponse::Attempt(Err(AttemptError::new(FailureKind::BadRequest, Some(400)))); + let step = store.resume(&mut driver, Some(bad_request)); + assert!( + matches!(&step, Step::HostCall(HostCall::RefuseFallbacks { level: 0, targets }) if targets.0.len() == 1), + "expected the chain's check, got {step:?}" + ); + let Step::Failed(error) = + store.resume(&mut driver, Some(HostCallResponse::RefuseFallbacks(true))) + else { + panic!("a refused chain should fail the call"); + }; + assert_eq!( + error.attempt().map(|a| a.kind), + Some(FailureKind::BadRequest) + ); + } + + trait StepKindOf { + fn unwrap_kind(self) -> StepKind; + } + + impl StepKindOf for Step { + fn unwrap_kind(self) -> StepKind { + match self { + Step::HostCall(call) => call.kind(), + Step::Done(_) => StepKind::Done, + Step::Failed(_) => StepKind::Failed, + } + } + } +} diff --git a/litellm-rust/crates/python-bridge/src/router/mod.rs b/litellm-rust/crates/python-bridge/src/router/mod.rs new file mode 100644 index 00000000000..c69259fcf1c --- /dev/null +++ b/litellm-rust/crates/python-bridge/src/router/mod.rs @@ -0,0 +1,2051 @@ +//! `RouterCoreHandle`: router-core behind a step-driven handle. Python drives each call through +//! [`RouterCall::next`], running every attempt, sleep, fallback check, selection and cache op the +//! call loop asks for and handing the outcome back; Rust never awaits Python. `select` is driven +//! the same way, since reading cooldowns is a cache step. + +mod config; +mod host; +mod retries; +mod stream; + +use std::{ + sync::{Arc, Mutex, PoisonError}, + time::{Duration, SystemTime}, +}; + +use litellm_error::Report; +use litellm_host_python::{from_py_argument, json_loads, to_py}; +use litellm_router_core::lookup::{ListScope, Lookup, Row}; +use litellm_router_core::{ + AttemptError, Blocked, BuildPath, Built, Deployment, DeploymentSpec, Error, FailureKind, + FailureReport, Op, OutcomeCallbacks, PreRouted, Provided, RetryCause, Router, RouterEvent, + SecretAt, SuccessReport, + build::Install, + health::{UsageMetric, current_minute, usage_key}, + registry::{DeploymentFilter, ResolveKind, ResolveScope}, + strategy::routing_context, +}; +use pyo3::{ + exceptions::{PyRuntimeError, PyValueError}, + prelude::*, + types::{PyBool, PyDict, PyList, PyString}, +}; +use serde_json::Value; + +use self::host::{ + CacheOp, CacheOutcome, DriveError, Driver, Finish, HostCall, HostCallResponse, + PythonRandomness, Stage, Step, StepHost, StepKind, +}; +use crate::errors::RustBridgeDeclined; + +/// The `error_kind` of a `failed` step. +#[derive(Clone, Copy, Debug, strum::IntoStaticStr)] +enum ErrorKind { + #[strum(serialize = "config")] + Config, + #[strum(serialize = "unknown_model")] + UnknownModel, + #[strum(serialize = "no_healthy_deployments")] + NoHealthyDeployments, + #[strum(serialize = "attempt")] + Attempt, + #[strum(serialize = "cache")] + Cache, + #[strum(serialize = "retired")] + Retired, + #[strum(serialize = "paused")] + Paused, + #[strum(serialize = "not_pass_through")] + NotPassThrough, + #[strum(serialize = "no_pass_through_deployments")] + NoPassThroughDeployments, + #[strum(serialize = "only_markers")] + OnlyMarkers, + #[strum(serialize = "endpoint_not_supported")] + EndpointNotSupported, + #[strum(serialize = "ptu_reserved")] + PtuReserved, + #[strum(serialize = "context_window_exceeded")] + ContextWindowExceeded, + #[strum(serialize = "unsupported_params")] + UnsupportedParams, + #[strum(serialize = "routing_plugins_left_nothing")] + RoutingPluginsLeftNothing, + #[strum(serialize = "routing_plugins_need_async")] + RoutingPluginsNeedAsync, + #[strum(serialize = "multiple_teams")] + MultipleTeams, + #[strum(serialize = "reserved")] + Reserved, + #[strum(serialize = "invalid_router_weights")] + InvalidRouterWeights, + #[strum(serialize = "non_numeric_weight")] + NonNumericWeight, + #[strum(serialize = "missing_routing_group")] + MissingRoutingGroup, + #[strum(serialize = "default_deployment_raises")] + DefaultDeploymentRaises, + #[strum(serialize = "index_out_of_range")] + IndexOutOfRange, + #[strum(serialize = "model_info")] + ModelInfo, +} + +impl From<&Error> for ErrorKind { + fn from(error: &Error) -> Self { + match error { + Error::Config(_) => Self::Config, + Error::UnknownModel { .. } => Self::UnknownModel, + Error::NoHealthyDeployments { .. } => Self::NoHealthyDeployments, + Error::Attempt(_) => Self::Attempt, + Error::Cache(_) => Self::Cache, + Error::Retired => Self::Retired, + Error::Paused { .. } => Self::Paused, + Error::NotPassThrough { .. } => Self::NotPassThrough, + Error::NoPassThroughDeployments { .. } => Self::NoPassThroughDeployments, + Error::OnlyMarkers { .. } => Self::OnlyMarkers, + Error::EndpointNotSupported { .. } => Self::EndpointNotSupported, + Error::PtuReserved { .. } => Self::PtuReserved, + Error::ContextWindowExceeded { .. } => Self::ContextWindowExceeded, + Error::UnsupportedParams { .. } => Self::UnsupportedParams, + Error::RoutingPluginsLeftNothing { .. } => Self::RoutingPluginsLeftNothing, + Error::RoutingPluginsNeedAsync => Self::RoutingPluginsNeedAsync, + Error::MultipleTeams { .. } => Self::MultipleTeams, + Error::Reserved { .. } => Self::Reserved, + Error::InvalidRouterWeights => Self::InvalidRouterWeights, + Error::NonNumericWeight { .. } => Self::NonNumericWeight, + Error::MissingRoutingGroup { .. } => Self::MissingRoutingGroup, + Error::DefaultDeploymentRaises { .. } => Self::DefaultDeploymentRaises, + Error::IndexOutOfRange { .. } => Self::IndexOutOfRange, + Error::ModelInfo { .. } => Self::ModelInfo, + } + } +} + +fn config_error_to_pyerr(error: config::Error) -> PyErr { + let message = Report::new(&error).to_string(); + if error.declined() { + RustBridgeDeclined::new_err(message) + } else { + PyValueError::new_err(message) + } +} + +fn drive_error_to_pyerr(error: DriveError) -> PyErr { + let message = Report::new(&error).to_string(); + if error.is_request() { + PyValueError::new_err(message) + } else { + PyRuntimeError::new_err(message) + } +} + +/// A document Python passes as a mapping or as its JSON text. +fn document(value: &Bound<'_, PyAny>) -> PyResult { + match value.cast::() { + Ok(text) => { + let parsed = json_loads(value.py(), text.to_str()?.as_bytes())?; + from_py_argument(parsed.bind(value.py())) + } + Err(_) => from_py_argument(value), + } +} + +/// The `OutcomeCallbacks` a reported outcome came through; none is a report no callback makes. +fn outcome_callbacks(value: Option<&Bound<'_, PyAny>>) -> PyResult { + let Some(value) = value else { + return Ok(OutcomeCallbacks::default()); + }; + serde_json::from_value(document(value)?) + .map_err(|error| PyValueError::new_err(error.to_string())) +} + +/// A router-core router whose host operations Python runs. +#[pyclass(frozen)] +pub(crate) struct RouterCoreHandle { + router: Router, +} + +#[pymethods] +impl RouterCoreHandle { + #[new] + fn new(config: &Bound<'_, PyAny>, dynamic: &Bound<'_, PyAny>) -> PyResult { + let (config, dynamic) = (document(config)?, document(dynamic)?); + let build = || { + let config = config::parse_router_config(config)?; + let dynamic = config::parse_dynamic_config(dynamic)?; + Ok(Router::builder(config, dynamic, Arc::new(StepHost)) + .randomness(Arc::new(PythonRandomness)) + .build()?) + }; + build() + .map(|router| Self { router }) + .map_err(config_error_to_pyerr) + } + + /// Start a call; nothing runs until the first `next(None)`. + fn start(&self, op: &str, request: &Bound<'_, PyAny>) -> PyResult { + let op = Op::new(op); + let request = + config::parse_request(document(request)?, None).map_err(config_error_to_pyerr)?; + let router = self.router.clone(); + Ok(RouterCall::new(async move { + router.call(op, request).await.map(Finish::Call) + })) + } + + /// Pick the deployment `model` routes to now, without calling it. Driven like `start`: its + /// steps are the cache reads and selections the pick needs, and it ends with `done` + /// (`deployment_id`, `model_group`) or `failed`. + fn select(&self, model: &str, request: &Bound<'_, PyAny>) -> PyResult { + let request = config::parse_request(document(request)?, Some(model)) + .map_err(config_error_to_pyerr)?; + let router = self.router.clone(); + Ok(RouterCall::new(async move { + let deployment = router.select(&Op::from("acompletion"), &request).await?; + Ok(Finish::Selected { + deployment_id: deployment.id.clone(), + model_group: deployment.model_name.clone(), + deployment: host::deployment_record(&deployment), + }) + })) + } + + /// List the deployments a pick for `model` would choose among, driven like `select`. + fn healthy(&self, model: &str, request: &Bound<'_, PyAny>) -> PyResult { + let request = config::parse_request(document(request)?, Some(model)) + .map_err(config_error_to_pyerr)?; + let router = self.router.clone(); + Ok(RouterCall::new(async move { + let candidates = router.healthy(&Op::from("acompletion"), &request).await?; + // A deployment named by its id is listed alone, as a dict (`_core_async_get_healthy_deployments`). + let single = candidates.kind.is_single() && candidates.healthy.len() == 1; + let deployments = candidates.healthy; + Ok(Finish::Healthy { + single, + deployment_ids: deployments.iter().map(|d| d.id.clone()).collect(), + deployments: deployments + .iter() + .map(|d| host::deployment_record(d)) + .collect(), + }) + })) + } + + /// Apply `{"upsert": [deployments], "delete": [ids], "settings": {dynamic document}}` and + /// return the new registry version. + fn apply(&self, changes: &Bound<'_, PyAny>) -> PyResult { + let changes = config::parse_changes(document(changes)?).map_err(config_error_to_pyerr)?; + self.router + .apply(changes) + .map_err(|error| config_error_to_pyerr(error.into())) + } + + /// Build `entries` (model_list entries) through the host, driven like `start`: its steps + /// are `secrets`, `validate`, `provider` and `added`, and it ends with `done` (`built`: one + /// `{"index", "id", "provider"}` per deployment installed). `path` is `create` + /// (`set_model_list`) or `add` (`add_deployment`); `install` is `replace`, `append` or + /// `upsert`. + #[pyo3(signature = (entries, path, resolve_secrets, install))] + fn build( + &self, + entries: &Bound<'_, PyAny>, + path: &str, + resolve_secrets: bool, + install: &str, + ) -> PyResult { + let entries: Vec> = + serde_json::from_value(document(entries)?) + .map_err(|error| PyValueError::new_err(format!("bad model_list: {error}")))?; + let path: BuildPath = path + .parse() + .map_err(|_| PyValueError::new_err(format!("`{path}` is not a build path")))?; + let install: Install = install + .parse() + .map_err(|_| PyValueError::new_err(format!("`{install}` is not a build install")))?; + let router = self.router.clone(); + Ok(RouterCall::new(async move { + router + .build(path, entries, resolve_secrets, install) + .await + .map(Finish::Built) + })) + } + + /// `get_router_model_info` for the deployment `query` describes (a + /// `RouterModelInfoQuery` document), driven like `start`: its steps are `model_info` + /// lookups, and it ends with `done` (`info`: the merged map or None, `user_keys`: the keys + /// whose values are the deployment's own). + fn router_model_info(&self, query: &Bound<'_, PyAny>) -> PyResult { + let query: litellm_router_core::model_info::RouterModelInfoQuery = + serde_json::from_value(document(query)?) + .map_err(|error| PyValueError::new_err(format!("bad model-info query: {error}")))?; + let router = self.router.clone(); + Ok(RouterCall::new(async move { + router + .router_model_info(&query) + .await + .map(Finish::RouterModelInfo) + })) + } + + /// `get_deployment_model_info` for deployment `model_id` serving `model_name`, driven like + /// `router_model_info`; it ends with `done` (`info`: the merged map or None). + fn deployment_model_info( + &self, + model_id: &str, + model_name: &Bound<'_, PyAny>, + ) -> PyResult { + let model_id = model_id.to_owned(); + let model_name = from_py_argument::(model_name)?; + let router = self.router.clone(); + Ok(RouterCall::new(async move { + router + .deployment_model_info(&model_id, &model_name) + .await + .map(Finish::DeploymentModelInfo) + })) + } + + /// `_set_model_group_info` for the deployments `query` names (a `GroupInfoQuery` document), + /// driven like `router_model_info`; it ends with `done` (`value`: None when no deployment + /// matched, else `{"assignments": {field: value}}`, what the merge set on the group the host + /// built from its first deployment). + fn model_group_info(&self, query: &Bound<'_, PyAny>) -> PyResult { + let query: litellm_router_core::group_info::GroupInfoQuery = + serde_json::from_value(document(query)?) + .map_err(|error| PyValueError::new_err(format!("bad group-info query: {error}")))?; + let router = self.router.clone(); + Ok(RouterCall::new(async move { + let info = router.model_group_info(&query).await?; + Ok(Finish::Value(info.map_or( + Value::Null, + |info| serde_json::json!({"assignments": info.assignments}), + ))) + })) + } + + /// `get_configured_token_limits` for the deployment with `id` and its `configured` model_info, + /// driven like `router_model_info`; it ends with `done` (`value`: `[max_input, max_output]`). + fn configured_token_limits( + &self, + id: &Bound<'_, PyAny>, + configured: &Bound<'_, PyAny>, + ) -> PyResult { + let id = from_py_argument::(id)?; + let Value::Object(configured) = document(configured)? else { + return Err(PyValueError::new_err( + "the configured model_info must be a dict", + )); + }; + let router = self.router.clone(); + Ok(RouterCall::new(async move { + let (input, output) = router.configured_token_limits(id, &configured).await?; + Ok(Finish::Value(serde_json::json!([input, output]))) + })) + } + + /// `get_model_listing_info` over the `deployments` a listed name names, driven like + /// `router_model_info`; it ends with `done` (`value`: `{"cost_map_keys", "max_input_tokens", + /// "max_output_tokens"}`). + fn model_listing_info(&self, deployments: &Bound<'_, PyAny>) -> PyResult { + let deployments: Vec = from_py_argument(deployments)?; + let router = self.router.clone(); + Ok(RouterCall::new(async move { + let listing = router.model_listing_info(&deployments).await?; + Ok(Finish::Value(serde_json::json!({ + "cost_map_keys": listing.cost_map_keys, + "max_input_tokens": listing.max_input_tokens, + "max_output_tokens": listing.max_output_tokens, + }))) + })) + } + + /// `get_configured_mode` / `get_configured_display_name`: `value` when it is a string that is + /// not blank. + #[staticmethod] + fn configured_text(value: &Bound<'_, PyAny>) -> PyResult> { + let value = from_py_argument::(value).ok(); + Ok(litellm_router_core::group_info::configured_text( + value.as_ref(), + )) + } + + /// `get_configured_service_tiers`: each of `model_infos`' `service_tiers`, unvalidated. + #[staticmethod] + fn configured_service_tiers( + py: Python<'_>, + model_infos: &Bound<'_, PyAny>, + ) -> PyResult>> { + let model_infos: Vec> = from_py_argument(model_infos)?; + litellm_router_core::group_info::configured_service_tiers(&model_infos) + .iter() + .map(|tiers| to_py(py, tiers)) + .collect() + } + + /// The routing-cache key deployment `id` serving litellm `model` counts `metric` (`tpm`, + /// `rpm`, `itpm` or `otpm`) under in `minute` (`HH-MM`). + fn usage_key(&self, id: &str, model: &str, metric: &str, minute: &str) -> PyResult { + let metric: UsageMetric = metric + .parse() + .map_err(|_| PyValueError::new_err(format!("`{metric}` is not a usage metric")))?; + Ok(usage_key(id, model, metric, minute)) + } + + /// Keep the providers `other`'s build resolved for the deployments this router serves too. + fn adopt_providers(&self, other: &Self) { + self.router.adopt_providers(&other.router); + } + + fn registry_version(&self) -> u64 { + self.router.registry().version() + } + + /// The version the deployments before the newest appends last changed at: a list read at + /// or after it is current once it reads `deployments_from(len(list))`. + fn registry_rewritten(&self) -> u64 { + self.router.registry().rewritten() + } + + /// Record a failure the router did not see itself (a health check, an interceptor, a + /// logging callback), driven like `start`: its steps are cache ops, and it ends with `done` + /// (`cooled`, `cooldown_time`). `counts_request` also counts the failed request against the deployment's + /// per-minute requests, as an async call's failure callback does. `callbacks` is the logged + /// request's `OutcomeCallbacks`. + #[pyo3(signature = (id, failure, deployment_cooldown=false, count_only=false, requested_group=None, cooldown_time=None, uncounted=false, counts_request=false, usage_model=None, callbacks=None))] + #[allow(clippy::too_many_arguments)] + fn report_failure( + &self, + id: &str, + failure: &Bound<'_, PyDict>, + deployment_cooldown: bool, + count_only: bool, + requested_group: Option, + cooldown_time: Option, + uncounted: bool, + counts_request: bool, + usage_model: Option, + callbacks: Option<&Bound<'_, PyAny>>, + ) -> PyResult { + let failure = match attempt_outcome(failure)? { + Err(failure) => failure, + Ok(_) => { + return Err(PyValueError::new_err( + "a reported failure needs `ok: False`", + )); + } + }; + let router = self.router.clone(); + let id = id.to_owned(); + let report = FailureReport { + deployment_cooldown, + cooldown_time, + count_only, + requested_group, + uncounted, + counts_request, + usage_model, + callbacks: outcome_callbacks(callbacks)?, + }; + Ok(RouterCall::new(async move { + router + .report_failure(&id, &failure, report) + .await + .map(|cooldown| Finish::Reported { cooldown }) + })) + } + + /// Record a success a logging callback saw, driven like `report_failure`; it ends with + /// `done` (`key`: the success counter's key, or None for an unknown deployment). + /// `usage_model` is the model its usage is keyed by, and `tracks_usage` the host's answer to + /// whether it counts any (see `Output::tracks_usage`). `callbacks`: see `report_failure`. + #[pyo3(signature = (id, total_tokens, requests, usage_model=None, tracks_usage=None, callbacks=None))] + fn report_success( + &self, + id: &str, + total_tokens: f64, + requests: f64, + usage_model: Option, + tracks_usage: Option, + callbacks: Option<&Bound<'_, PyAny>>, + ) -> PyResult { + let router = self.router.clone(); + let id = id.to_owned(); + let report = SuccessReport { + tracks_usage, + usage_model, + callbacks: outcome_callbacks(callbacks)?, + }; + Ok(RouterCall::new(async move { + router + .report_success(&id, total_tokens, requests, report) + .await + .map(|key| Finish::Counted { key }) + })) + } + + /// Count a success against `id`, configured or not; ends with `done` (`key`). `callbacks`: + /// see `report_failure`. + #[pyo3(signature = (id, callbacks=None))] + fn count_success( + &self, + id: &str, + callbacks: Option<&Bound<'_, PyAny>>, + ) -> PyResult { + let router = self.router.clone(); + let id = id.to_owned(); + let callbacks = outcome_callbacks(callbacks)?; + Ok(RouterCall::new(async move { + router + .count_success(&id, callbacks) + .await + .map(|key| Finish::Counted { key }) + })) + } + + /// The usage keys of this minute for every deployment of `model_group`: + /// `[(id, tpm_key, rpm_key)]`, or None when the group has no deployment. + fn usage_keys(&self, model_group: &str) -> Option> { + let minute = current_minute(SystemTime::now()); + let group = self.router.registry().group(model_group)?; + Some( + group + .iter() + .map(|deployment| { + let model = &deployment.params.model; + ( + deployment.id.clone(), + usage_key(&deployment.id, model, UsageMetric::Tpm, &minute), + usage_key(&deployment.id, model, UsageMetric::Rpm, &minute), + ) + }) + .collect(), + ) + } + + /// Every configured deployment, as `model_list` entries with their ids, in config order. + fn deployments<'py>(&self, py: Python<'py>) -> PyResult> { + let registry = self.router.registry(); + let records = registry + .deployments(&DeploymentFilter::default()) + .iter() + .map(|deployment| to_py(py, &host::deployment_record(deployment))) + .collect::>>()?; + PyList::new(py, records) + } + + /// The deployments from position `start` on, as `deployments()` lists them. + fn deployments_from<'py>(&self, py: Python<'py>, start: usize) -> PyResult> { + let registry = self.router.registry(); + let records = registry + .deployments_from(start) + .iter() + .map(|deployment| to_py(py, &host::deployment_record(deployment))) + .collect::>>()?; + PyList::new(py, records) + } + + /// The deployment with this id, or None. + fn deployment(&self, py: Python<'_>, id: &str) -> PyResult>> { + self.router + .registry() + .deployment(id) + .map(|deployment| to_py(py, &host::deployment_record(&deployment))) + .transpose() + } + + // -- registry lookups (router-core `lookup.rs`): rows the host reads its stored entries at -- + + /// `_core_get_all_deployments`, as rows. + #[pyo3(signature = (model_name, team_id=None, rename=None))] + fn lookup_all<'py>( + &self, + py: Python<'py>, + model_name: &str, + team_id: Option<&str>, + rename: Option<&str>, + ) -> PyResult> { + let registry = self.router.registry(); + rows_to_py(py, &Lookup::new(®istry).all(model_name, team_id, rename)) + } + + /// `_team_deployments_across_teams`: the rows, and how many teams own them. + fn lookup_across_teams<'py>( + &self, + py: Python<'py>, + model: &str, + ) -> PyResult<(Bound<'py, PyAny>, usize)> { + let registry = self.router.registry(); + let (rows, teams) = Lookup::new(®istry).across_teams(model); + Ok((rows_to_py(py, &rows)?, teams)) + } + + /// `_core_config_deployments`, as rows. + fn lookup_config<'py>(&self, py: Python<'py>) -> PyResult> { + let registry = self.router.registry(); + rows_to_py(py, &Lookup::new(®istry).config()) + } + + /// `_core_get_model_ids`: the ids, or the row the read raises at. + #[pyo3(signature = (model_name=None, exclude_team_models=false))] + fn lookup_model_ids<'py>( + &self, + py: Python<'py>, + model_name: Option<&str>, + exclude_team_models: bool, + ) -> PyResult<(Vec, Option>)> { + let registry = self.router.registry(); + match Lookup::new(®istry).model_ids(model_name, exclude_team_models) { + Ok(ids) => Ok((ids, None)), + Err(raw) => Ok((Vec::new(), Some(row_to_py(py, &raw)?))), + } + } + + /// `_core_get_model_list`, as rows. + #[pyo3(signature = (model_name, team_id, scope))] + fn lookup_model_list<'py>( + &self, + py: Python<'py>, + model_name: Option<&str>, + team_id: Option<&str>, + scope: &Bound<'_, PyAny>, + ) -> PyResult> { + let scope = list_scope(scope)?; + let registry = self.router.registry(); + rows_to_py( + py, + &Lookup::new(®istry).model_list(model_name, team_id, &scope), + ) + } + + /// `_core_get_model_names`, or the row the read raises at. + fn lookup_model_names<'py>( + &self, + py: Python<'py>, + team_id: Option<&str>, + scope: &Bound<'_, PyAny>, + ) -> PyResult<(Vec, Option>)> { + let scope = list_scope(scope)?; + let registry = self.router.registry(); + match Lookup::new(®istry).model_names(team_id, &scope) { + Ok(names) => Ok((names, None)), + Err(raw) => Ok((Vec::new(), Some(row_to_py(py, &raw)?))), + } + } + + /// `_core_get_model_access_groups`: `(group, names)` pairs in the order the groups appear, + /// or the row the read raises at. + fn lookup_access_groups<'py>( + &self, + py: Python<'py>, + model_name: Option<&str>, + access_group: Option<&str>, + team_id: Option<&str>, + scope: &Bound<'_, PyAny>, + ) -> PyResult<(AccessGroups, Option>)> { + let scope = list_scope(scope)?; + let registry = self.router.registry(); + match Lookup::new(®istry).access_groups(model_name, access_group, team_id, &scope) { + Ok(groups) => Ok((groups, None)), + Err(raw) => Ok((Vec::new(), Some(row_to_py(py, &raw)?))), + } + } + + /// Whether the whole list lists any row, and the name of its first row whose litellm model + /// is `model_id` (or ends in `/model_id` or `:model_id`); or the row the read raises at. + fn lookup_name_serving<'py>( + &self, + py: Python<'py>, + model_id: &str, + scope: &Bound<'_, PyAny>, + ) -> PyResult<(bool, Option, Option>)> { + let scope = list_scope(scope)?; + let registry = self.router.registry(); + match Lookup::new(®istry).name_serving(model_id, &scope) { + Ok((listed, name)) => Ok((listed, name, None)), + Err(raw) => Ok((false, None, Some(row_to_py(py, &raw)?))), + } + } + + /// Whether `model_names` (the group names) holds `name`. + fn has_model_name(&self, name: &str) -> bool { + self.router + .registry() + .host_indexes() + .model_names + .contains(&Value::String(name.to_owned())) + } + + /// `_core_team_model_has_alternatives`, or the row the read raises at. + fn lookup_team_model_has_alternatives<'py>( + &self, + py: Python<'py>, + id: &str, + ) -> PyResult<(bool, Option>)> { + let registry = self.router.registry(); + match Lookup::new(®istry).team_model_has_alternatives(id) { + Ok(found) => Ok((found, None)), + Err(raw) => Ok((false, Some(row_to_py(py, &raw)?))), + } + } + + /// The position of the group's first deployment the team may use, or the row the read + /// raises at. + fn lookup_group_usable_by_team<'py>( + &self, + py: Python<'py>, + model_group: &str, + team_id: Option<&str>, + ) -> PyResult<(Option, Option>)> { + let registry = self.router.registry(); + match Lookup::new(®istry).group_deployment_usable_by_team(model_group, team_id) { + Ok(position) => Ok((position, None)), + Err(raw) => Ok((None, Some(row_to_py(py, &raw)?))), + } + } + + /// The position of the team public name's first deployment, or the row the read raises at. + fn lookup_team_public_name<'py>( + &self, + py: Python<'py>, + model_id: &str, + team_id: &str, + ) -> PyResult<(Option, Option>)> { + let registry = self.router.registry(); + match Lookup::new(®istry).team_public_name_deployment(model_id, team_id) { + Ok(position) => Ok((position, None)), + Err(raw) => Ok((None, Some(row_to_py(py, &raw)?))), + } + } + + /// `_routable_deployments`, as rows. + fn lookup_routable<'py>( + &self, + py: Python<'py>, + model_group: &str, + team_id: Option<&str>, + ) -> PyResult> { + let registry = self.router.registry(); + rows_to_py(py, &Lookup::new(®istry).routable(model_group, team_id)) + } + + /// `_core_wildcard_route_deployments`: the team router's routes, else the global router's. + fn lookup_wildcard_route<'py>( + &self, + py: Python<'py>, + model: &str, + team_id: Option<&str>, + requests: Vec, + ) -> PyResult> { + let registry = self.router.registry(); + rows_to_py( + py, + &Lookup::new(®istry).wildcard_route(model, team_id, &requests), + ) + } + + /// `_core_get_candidate_model_ids_for_route` for the resolved name: the members' ids when it + /// names a callable routing group (`members`), else for a name that is no model group the + /// team's or the wildcard route's, else the group's; or the row the read raises at. + fn lookup_candidate_ids<'py>( + &self, + py: Python<'py>, + resolved: &str, + team_id: Option<&str>, + members: Option>, + requests: Vec, + ) -> PyResult<(Vec, Option>)> { + let registry = self.router.registry(); + let lookup = Lookup::new(®istry); + match lookup.candidate_ids(resolved, team_id, members.as_deref(), &requests) { + Ok(ids) => Ok((ids, None)), + Err(raw) => Ok((Vec::new(), Some(row_to_py(py, &raw)?))), + } + } + + /// Whether the id map holds `id` (`_has_model_id`). + fn has_model_id(&self, id: &str) -> bool { + self.router.registry().host_indexes().has_id(id) + } + + /// What the id map holds for `id`, as written (a position, or what a caller wrote), or None. + fn id_position<'py>(&self, py: Python<'py>, id: &str) -> PyResult>> { + let registry = self.router.registry(); + registry + .host_indexes() + .id_entry(id) + .map(|held| to_py(py, held).map(|value| value.into_bound(py))) + .transpose() + } + + /// The id map's entries, in order: `(id, what it holds)`. + fn id_positions<'py>(&self, py: Python<'py>) -> PyResult> { + let registry = self.router.registry(); + let pairs: Vec = registry + .host_indexes() + .id_map + .iter() + .map(|(id, held)| Value::Array(vec![Value::String(id.clone()), held.clone()])) + .collect(); + to_py(py, &Value::Array(pairs)).map(|value| value.into_bound(py)) + } + + /// What the name index holds for `name`, as written, or None when it does not index it. + fn named_positions<'py>( + &self, + py: Python<'py>, + name: &str, + ) -> PyResult>> { + let registry = self.router.registry(); + registry + .host_indexes() + .named(name) + .map(|held| to_py(py, held).map(|value| value.into_bound(py))) + .transpose() + } + + /// What the team index holds for `(team, public)`, as written, or None. + fn team_positions<'py>( + &self, + py: Python<'py>, + team: &str, + public: &str, + ) -> PyResult>> { + let registry = self.router.registry(); + registry + .host_indexes() + .team_named(team, public) + .map(|held| to_py(py, held).map(|value| value.into_bound(py))) + .transpose() + } + + /// `team_public_model_names_of`: the public names the team index holds for `team`. + fn team_public_names_of(&self, team: &str) -> Vec { + self.router + .registry() + .host_indexes() + .team_public_names_of(team) + } + + /// The positions of the deployments whose `litellm_params.model` is `model`. + fn litellm_model_positions(&self, model: &str) -> Vec { + self.router + .registry() + .deployments(&DeploymentFilter::default()) + .iter() + .enumerate() + .filter(|(_, deployment)| { + deployment.raw_params.get("model").and_then(Value::as_str) == Some(model) + }) + .map(|(position, _)| position) + .collect() + } + + /// `PatternMatchRouter.route` over the wildcard router (or `team`'s): the first pattern's + /// deployments, each with the model the wildcards matched, or None. + #[pyo3(signature = (team, request, filtered=None))] + fn pattern_route<'py>( + &self, + py: Python<'py>, + team: Option<&str>, + request: &str, + filtered: Option>, + ) -> PyResult>> { + let registry = self.router.registry(); + let indexes = registry.host_indexes(); + let router = match team { + None => Some(&indexes.patterns), + Some(team) => indexes + .team_patterns + .iter() + .find(|(held, _)| held == team) + .map(|(_, router)| router), + }; + router + .and_then(|router| router.route(request, filtered.as_deref())) + .map(|routed| to_py(py, &Value::Array(routed)).map(|value| value.into_bound(py))) + .transpose() + } + + /// `get_models_by_pattern`: the models the first matching (request, pattern) pair routes to. + fn models_by_pattern(&self, team: Option<&str>, requests: Vec) -> Option> { + let registry = self.router.registry(); + let indexes = registry.host_indexes(); + let router = match team { + None => Some(&indexes.patterns), + Some(team) => indexes + .team_patterns + .iter() + .find(|(held, _)| held == team) + .map(|(_, router)| router), + }; + router.and_then(|router| router.models_by_pattern(&requests)) + } + + /// The indexes the host hands out (`HostIndexes`): id map, name and team indexes, group and + /// deployment names, team public names, provider default ids and the pattern routers. + fn host_indexes<'py>(&self, py: Python<'py>) -> PyResult> { + let indexes = serde_json::to_value(self.router.registry().host_indexes()) + .map_err(|error| PyValueError::new_err(error.to_string()))?; + to_py(py, &indexes).map(|value| value.into_bound(py)) + } + + /// `provider_default_deployment_ids` as the indexes hold it now: the ids of the wildcard + /// deployments the cooldown check skips. + fn provider_default_ids(&self) -> Vec { + self.router + .registry() + .host_indexes() + .provider_default_ids + .clone() + } + + /// The provider the host resolved when it built the deployment with this id (its + /// `provider` build step), or None when there is no such deployment or it never went + /// through a host build. + fn provider(&self, id: &str) -> Option { + self.router + .registry() + .deployment(id) + .and_then(|deployment| deployment.resolved_provider.clone()) + } + + /// What `name` resolves to: `{"model_group", "kind", "ids"}`, or None for an unknown name. + #[pyo3(signature = (name, specific_deployment=false))] + fn resolve<'py>( + &self, + py: Python<'py>, + name: &str, + specific_deployment: bool, + ) -> PyResult>> { + let registry = self.router.registry(); + let Some(resolution) = registry.resolve( + name, + ResolveScope { + specific_deployment, + }, + ) else { + return Ok(None); + }; + let dict = PyDict::new(py); + dict.set_item("model_group", resolution.group)?; + dict.set_item("kind", format!("{:?}", resolution.kind))?; + dict.set_item( + "ids", + PyList::new(py, resolution.deployments.iter().map(|d| d.id.clone()))?, + )?; + Ok(Some(dict)) + } + + /// The deployments a pick for `model` would resolve to for `request`, before any team, + /// cooldown or strategy filter (`_core_get_routable_deployments`); empty when it refuses. + fn routable<'py>( + &self, + py: Python<'py>, + model: &str, + request: &Bound<'py, PyAny>, + ) -> PyResult> { + let request = config::parse_request(document(request)?, Some(model)) + .map_err(config_error_to_pyerr)?; + let records = litellm_router_core::resolve::routable( + &self.router.registry(), + &self.router.settings().dynamic, + &request, + ) + .iter() + .map(|d| to_py(py, &host::deployment_record(d))) + .collect::>>()?; + PyList::new(py, records) + } + + /// `_get_routing_context` for a pick of `model`: `(strategy, scope, group)`, spelled as a + /// strategy select step spells them. A name that resolves to nothing reads its routing group + /// membership as is. + #[pyo3(signature = (model, routing_strategy=None))] + fn routing_context( + &self, + model: &str, + routing_strategy: Option<&str>, + ) -> (String, &'static str, Option) { + let registry = self.router.registry(); + let settings = self.router.settings(); + let (kind, group) = registry + .resolve(model, ResolveScope::default()) + .map_or((ResolveKind::Group, model.to_owned()), |resolution| { + (resolution.kind, resolution.group) + }); + let (strategy, scope) = routing_context( + ®istry, + &settings.routing_strategy, + routing_strategy, + &kind, + &group, + ); + let (scope, group) = scope.wire(); + (strategy.to_string(), scope, group.map(str::to_owned)) + } + + fn retire(&self) { + self.router.retire(); + } + + /// What may still take over a stream from `origin` that fails before content, and the + /// retry budget for a failure of `kind` (see `stream::policy`). + #[pyo3(signature = (request, origin, kind=None))] + fn stream_policy<'py>( + &self, + py: Python<'py>, + request: &Bound<'py, PyAny>, + origin: &Bound<'py, PyAny>, + kind: Option<&str>, + ) -> PyResult> { + stream::policy(py, &self.router, request, origin, kind) + } + + /// Retry a stream from `origin` that failed before content within its group, driven like + /// `start`: `done` when a retry returned a response, `failed` with the newest failure. + fn retry_stream( + &self, + op: &str, + request: &Bound<'_, PyAny>, + origin: &Bound<'_, PyAny>, + failure: &Bound<'_, PyDict>, + ) -> PyResult { + stream::recover( + &self.router, + stream::Recovery::Retry, + op, + request, + origin, + failure, + ) + } + + /// The fallback walk for a stream from `origin` that failed before content, driven like + /// `start`. + fn fallback_stream( + &self, + op: &str, + request: &Bound<'_, PyAny>, + origin: &Bound<'_, PyAny>, + failure: &Bound<'_, PyDict>, + ) -> PyResult { + stream::recover( + &self.router, + stream::Recovery::Fallback, + op, + request, + origin, + failure, + ) + } + + /// `get_num_retries_from_retry_policy` for a failure shaped like an attempt outcome. + #[pyo3(signature = (failure, model_group=None))] + fn retry_policy_retries( + &self, + failure: &Bound<'_, PyDict>, + model_group: Option<&str>, + ) -> PyResult> { + retries::policy_retries(&self.router, failure, model_group) + } + + /// `_core_get_allowed_fails_from_policy` for a failure shaped like an attempt outcome. + fn allowed_fails_for(&self, failure: &Bound<'_, PyDict>) -> PyResult> { + retries::allowed_fails(&self.router, failure) + } + + /// `should_retry_this_error` as a predicate, for a failure shaped like an attempt outcome. + fn should_retry( + &self, + failure: &Bound<'_, PyDict>, + counts: (usize, usize), + fallbacks: (bool, bool, bool), + ) -> PyResult { + retries::should_retry(failure, counts, fallbacks) + } +} + +/// One router call, driven a step at a time by Python. +#[pyclass(frozen)] +pub(crate) struct RouterCall { + driver: Mutex, +} + +impl RouterCall { + fn new( + future: impl std::future::Future> + Send + 'static, + ) -> Self { + Self { + driver: Mutex::new(Driver::new(future)), + } + } +} + +#[pymethods] +impl RouterCall { + /// Hand back the outcome of the last step (`None` on the first call) and return the next + /// step, or the call's end. + #[pyo3(signature = (outcome))] + fn next<'py>( + &self, + py: Python<'py>, + outcome: Option>, + ) -> PyResult> { + let mut driver = self.driver.lock().unwrap_or_else(PoisonError::into_inner); + let response = match driver.pending_call() { + None => match outcome { + None => None, + Some(_) => return Err(drive_error_to_pyerr(DriveError::UnexpectedResponse)), + }, + Some(call) => Some(host_call_response(call, outcome)?), + }; + let step = driver.resume(response).map_err(drive_error_to_pyerr)?; + let dict = step_dict(py, step)?; + dict.set_item("events", events_list(py, driver.take_events())?)?; + Ok(dict) + } +} + +/// `(group, names)` pairs, in the order the groups appear. +type AccessGroups = Vec<(String, Vec)>; + +/// Lookup rows as the host reads them. +fn rows_to_py<'py>(py: Python<'py>, rows: &[Row]) -> PyResult> { + let rows = Value::Array(rows.iter().map(Row::wire).collect()); + to_py(py, &rows).map(|value| value.into_bound(py)) +} + +fn row_to_py<'py>(py: Python<'py>, row: &Row) -> PyResult> { + to_py(py, &row.wire()).map(|value| value.into_bound(py)) +} + +/// `{"aliases": [{alias, target}], "groups": [{name, members}], "wildcard": [requests]}`. +fn list_scope(scope: &Bound<'_, PyAny>) -> PyResult { + serde_json::from_value(document(scope)?) + .map_err(|error| PyValueError::new_err(format!("bad lookup scope: {error}"))) +} + +/// Read Python's outcome for `call`. +fn host_call_response( + call: &HostCall, + outcome: Option>, +) -> PyResult { + let missing = || { + drive_error_to_pyerr(DriveError::MissingResponse { + expected: call.kind(), + }) + }; + // `{"raised": True, ...}`: the step's work raised, reported as an attempt's failure. + if let HostCall::Select { .. } | HostCall::PreRoute { .. } | HostCall::AllowFallback { .. } = + call + && let Some(failure) = outcome.as_ref().and_then(|o| o.cast::().ok()) + && failure.contains("raised")? + { + return match attempt_outcome(failure)? { + Err(error) => Ok(HostCallResponse::Raised(error)), + Ok(_) => Err(PyValueError::new_err("a `raised` outcome is a failure")), + }; + } + match call { + HostCall::Sleep { .. } => Ok(HostCallResponse::Slept), + HostCall::AllowFallback { .. } => Ok(HostCallResponse::AllowFallback( + outcome.ok_or_else(missing)?.extract()?, + )), + HostCall::RefuseFallbacks { .. } => Ok(HostCallResponse::RefuseFallbacks( + outcome.ok_or_else(missing)?.extract()?, + )), + HostCall::Select { + stage: Stage::PreCall, + .. + } => { + let facts = from_py_argument(&outcome.ok_or_else(missing)?)?; + serde_json::from_value(facts) + .map(HostCallResponse::PreCall) + .map_err(|error| PyValueError::new_err(format!("bad pre-call facts: {error}"))) + } + HostCall::Select { + stage: Stage::RoutingPlugins { .. } | Stage::ProviderRequest, + .. + } => { + // `{"models": [...] | None}`: a bare None would read as no outcome. + let outcome = outcome.ok_or_else(missing)?; + let models = outcome.cast::()?.get_item("models")?; + Ok(HostCallResponse::Models(match models { + Some(models) if !models.is_none() => Some(models.extract()?), + _ => None, + })) + } + HostCall::Select { + stage: Stage::CallbackFilters { .. }, + .. + } => { + // The remaining ids, or a failure mapping as an attempt reports one. + let outcome = outcome.ok_or_else(missing)?; + if let Ok(failure) = outcome.cast::() { + return match attempt_outcome(failure)? { + Err(error) => Ok(HostCallResponse::Filtered(Err(error))), + Ok(_) => Err(PyValueError::new_err( + "a `callback_filters` outcome is the remaining ids or a failure", + )), + }; + } + Ok(HostCallResponse::Filtered(Ok(outcome.extract()?))) + } + HostCall::Select { .. } => Ok(HostCallResponse::Select( + outcome.ok_or_else(missing)?.extract()?, + )), + HostCall::PreRoute { .. } => pre_route_outcome(outcome).map(HostCallResponse::PreRoute), + HostCall::Attempt { .. } => { + let outcome = outcome.ok_or_else(missing)?; + let outcome = outcome.cast::()?; + let blocked = match outcome.get_item("blocked")? { + Some(value) if !value.is_none() => { + let name: String = value.extract()?; + Some(name.parse::().map_err(|_| { + PyValueError::new_err(format!( + "`{name}` is not a way a response is blocked" + )) + })?) + } + Some(_) | None => None, + }; + let deferred = outcome + .get_item("deferred")? + .is_some_and(|flag| flag.is_truthy().unwrap_or(false)); + let tracks_usage = outcome + .get_item("tracks_usage")? + .map(|flag| flag.extract::>()) + .transpose()? + .flatten(); + attempt_outcome(outcome).map(|result| match result { + Ok(total_tokens) if blocked.is_some() || deferred || tracks_usage.is_some() => { + HostCallResponse::Served { + total_tokens, + blocked, + deferred, + tracks_usage, + } + } + result => HostCallResponse::Attempt(result), + }) + } + HostCall::Cache { op, .. } => cache_outcome(*op, outcome).map(HostCallResponse::Cache), + HostCall::Secrets { .. } => Ok(HostCallResponse::Secrets(from_py_argument( + &outcome.ok_or_else(missing)?, + )?)), + // `None` is a response here (no row, no entry), so a missing outcome reads as one. + HostCall::ModelInfo { .. } => Ok(HostCallResponse::ModelInfo(match outcome { + Some(outcome) if !outcome.is_none() => from_py_argument(&outcome)?, + _ => Value::Null, + })), + HostCall::Build { kind, .. } => { + let outcome = outcome.ok_or_else(missing)?; + let outcome = outcome.cast::()?; + let invalid = outcome + .get_item("invalid")? + .is_some_and(|flag| flag.is_truthy().unwrap_or(false)); + match kind { + StepKind::Provider => Ok(HostCallResponse::Provided(if invalid { + Provided::Invalid + } else { + Provided::Provider(match outcome.get_item("provider")? { + Some(provider) if !provider.is_none() => Some(provider.extract()?), + _ => None, + }) + })), + _ => { + let built = if invalid { + Built::Invalid + } else { + match outcome.get_item("entry")? { + Some(entry) if !entry.is_none() => { + Built::Entry(from_py_argument(&entry)?) + } + _ => Built::Skip, + } + }; + Ok(if *kind == StepKind::Validate { + HostCallResponse::Validated(built) + } else { + HostCallResponse::Added(built) + }) + } + } + } + } +} + +/// A pre-route response: `None`, a model group name, or a custom strategy's deployment document. +fn pre_route_outcome(outcome: Option>) -> PyResult> { + let Some(outcome) = outcome.filter(|outcome| !outcome.is_none()) else { + return Ok(None); + }; + if let Ok(model) = outcome.cast::() { + return Ok(Some(PreRouted::Model(model.to_str()?.to_owned()))); + } + let spec: DeploymentSpec = serde_json::from_value(document(&outcome)?) + .map_err(|error| PyValueError::new_err(error.to_string()))?; + let deployment = Deployment::from_spec(spec) + .map_err(|error| PyValueError::new_err(Report::new(&error).to_string()))?; + Ok(Some(PreRouted::Deployment(Arc::new(deployment)))) +} + +/// `get`: the JSON value or None; `batch_get`: a list of JSON values or None, one per key; +/// `increment`: the new int or float; `set` and `delete`: None. +fn cache_outcome(op: CacheOp, outcome: Option>) -> PyResult { + let outcome = outcome.filter(|outcome| !outcome.is_none()); + let json = |value: Option| value.filter(|value| !value.is_null()); + match op { + CacheOp::Get => outcome + .map(|outcome| from_py_argument::(&outcome)) + .transpose() + .map(|value| CacheOutcome::Value(json(value))), + CacheOp::BatchGet => { + let outcome = outcome.ok_or_else(|| { + PyValueError::new_err("a `batch_get` cache outcome must be a list, not None") + })?; + if !outcome.is_instance_of::() { + return Err(PyValueError::new_err( + "a `batch_get` cache outcome must be a list with one entry per key", + )); + } + let values: Vec> = from_py_argument(&outcome)?; + Ok(CacheOutcome::Values(values.into_iter().map(json).collect())) + } + CacheOp::Increment => outcome + .filter(|outcome| !outcome.is_instance_of::()) + .and_then(|outcome| outcome.extract::().ok()) + .map(CacheOutcome::Number) + .ok_or_else(|| { + PyValueError::new_err("an `increment` cache outcome must be the new int or float") + }), + CacheOp::Set | CacheOp::Delete | CacheOp::IncrementUsage => match outcome { + None => Ok(CacheOutcome::Stored), + Some(_) => Err(PyValueError::new_err(format!( + "a `{op}` cache step has no outcome, so pass None" + ))), + }, + } +} + +/// `{"ok": True, "total_tokens": int}` (an attempt's may add `"blocked": str`, `"deferred": bool` +/// and `"tracks_usage": bool`) or `{"ok": False, "kind": str, "status": int | None, +/// "retry_after": float | None, "num_retries": int | None, "detail": str | None, +/// "outside_deployment": bool, "caller_scoped": bool, "status_absent": bool, +/// "uncounted_request": bool}`. +pub(crate) fn attempt_outcome(outcome: &Bound<'_, PyDict>) -> PyResult> { + fn item<'py, T: FromPyObjectOwned<'py>>( + outcome: &Bound<'py, PyDict>, + key: &str, + ) -> PyResult> { + match outcome.get_item(key)? { + Some(value) if !value.is_none() => Ok(Some(value.extract().map_err(Into::into)?)), + Some(_) | None => Ok(None), + } + } + let ok: bool = item(outcome, "ok")? + .ok_or_else(|| PyValueError::new_err("an attempt outcome needs `ok`"))?; + if ok { + return Ok(Ok(item(outcome, "total_tokens")?.unwrap_or(0))); + } + let kind: String = item(outcome, "kind")? + .ok_or_else(|| PyValueError::new_err("a failed attempt outcome needs `kind`"))?; + let kind: FailureKind = kind + .parse() + .map_err(|_| PyValueError::new_err(format!("`{kind}` is not a router failure kind")))?; + let retry_after = item::(outcome, "retry_after")? + .map(|seconds| { + Duration::try_from_secs_f64(seconds) + .map_err(|_| PyValueError::new_err("`retry_after` must be a non-negative number")) + }) + .transpose()?; + let mut failure = AttemptError::new(kind, item(outcome, "status")?); + failure.retry_after = retry_after; + failure.num_retries = item(outcome, "num_retries")?; + failure.failed_deployment_id = item(outcome, "failed_deployment_id")?; + failure.detail = item::(outcome, "detail")?.map(Into::into); + failure.recorded = item::(outcome, "recorded")?.unwrap_or(false); + failure.names_provider = item::(outcome, "llm_provider")?.unwrap_or(false); + failure.agentic_followup = item::(outcome, "agentic_followup")?.unwrap_or(false); + failure.outside_deployment = item::(outcome, "outside_deployment")?.unwrap_or(false); + failure.status_absent = item::(outcome, "status_absent")?.unwrap_or(false); + failure.uncounted_request = item::(outcome, "uncounted_request")?.unwrap_or(false); + failure.caller_scoped = item::(outcome, "caller_scoped")?.unwrap_or(false); + failure.classes = item::>(outcome, "classes")? + .unwrap_or_default() + .into_iter() + .map(Into::into) + .collect(); + Ok(Err(failure)) +} + +/// `[{"kind": "retries_exhausted", "model_group": str, "attempted": int, "max": int} +/// | {"kind": "fallback", "to": str, "order": int | None, "level": int}]`. +/// The `cause` of a `retry_logged` or `retry_policy_raised` event, with what names it. +fn cause_items(dict: &Bound<'_, PyDict>, cause: RetryCause) -> PyResult<()> { + match cause { + RetryCause::Attempt => dict.set_item("cause", "attempt"), + RetryCause::Level => dict.set_item("cause", "level"), + RetryCause::Other => dict.set_item("cause", "other"), + RetryCause::Pick(picked) => { + dict.set_item("cause", "pick")?; + pick_items(dict, &picked.0) + } + } +} + +/// A pick's error, as the `failed` step names it (`error_kind` and what names it), with the +/// model its message names (`failed_model`) when the error carries one. +fn pick_items(dict: &Bound<'_, PyDict>, error: &Error) -> PyResult<()> { + error_items(dict, error)?; + match error { + Error::NoHealthyDeployments { model } | Error::UnknownModel { model } => { + dict.set_item("failed_model", model) + } + _ => Ok(()), + } +} + +/// `error_kind` and the fields that name the error. +fn error_items(dict: &Bound<'_, PyDict>, error: &Error) -> PyResult<()> { + dict.set_item("error", Report::new(error).to_string())?; + dict.set_item("error_kind", <&'static str>::from(ErrorKind::from(error)))?; + match error { + Error::EndpointNotSupported { model, message } => { + dict.set_item("error_model", model)?; + dict.set_item("error_detail", message) + } + Error::OnlyMarkers { model } + | Error::UnknownModel { model } + | Error::NoPassThroughDeployments { model } + | Error::PtuReserved { model } + | Error::MultipleTeams { model } + | Error::DefaultDeploymentRaises { model } + | Error::IndexOutOfRange { model } => dict.set_item("error_model", model), + Error::ContextWindowExceeded { model, detail } => { + dict.set_item("error_model", model)?; + dict.set_item("error_detail", detail) + } + Error::UnsupportedParams { model, params } => { + dict.set_item("error_model", model)?; + dict.set_item("error_params", PyList::new(dict.py(), params)?) + } + Error::RoutingPluginsLeftNothing { models } => { + dict.set_item("error_params", PyList::new(dict.py(), models)?) + } + Error::Reserved { + model, + until, + timezone, + } => { + dict.set_item("error_model", model)?; + dict.set_item("error_detail", format!("{until} {timezone}")) + } + Error::NonNumericWeight { metric } => dict.set_item("error_detail", metric), + Error::MissingRoutingGroup { group } => dict.set_item("error_model", group), + Error::ModelInfo { exception, detail } => { + dict.set_item("error_exception", *exception)?; + dict.set_item("error_detail", detail) + } + _ => Ok(()), + } +} + +fn events_list(py: Python<'_>, events: Vec) -> PyResult> { + let list = PyList::empty(py); + for event in events { + let dict = PyDict::new(py); + match event { + RouterEvent::RetriesExhausted { + model_group, + attempted, + max, + level, + picked, + } => { + dict.set_item("kind", "retries_exhausted")?; + dict.set_item("model_group", model_group)?; + dict.set_item("attempted", attempted)?; + dict.set_item("max", max)?; + dict.set_item("level", level)?; + dict.set_item("picked", picked)?; + } + RouterEvent::Retry { retry, .. } => { + dict.set_item("kind", "retry")?; + dict.set_item("retry", retry)?; + } + RouterEvent::FallingBack { level, model } => { + dict.set_item("kind", "falling_back")?; + dict.set_item("level", level)?; + dict.set_item("model", model)?; + } + RouterEvent::OrderChainFailed { level, last_hop } => { + dict.set_item("kind", "order_chain_failed")?; + dict.set_item("level", level)?; + dict.set_item("last_hop", last_hop)?; + } + RouterEvent::FallbacksFailed { + level, + chain, + last_hop, + } => { + dict.set_item("kind", "fallbacks_failed")?; + dict.set_item("level", level)?; + dict.set_item("chain", chain)?; + dict.set_item("last_hop", last_hop)?; + } + RouterEvent::Fallback { + from, + to, + order, + depth, + level, + scope, + parent, + max_fallbacks, + unroutable, + } => { + dict.set_item("kind", "fallback")?; + dict.set_item("from", from)?; + dict.set_item("to", to)?; + dict.set_item("order", order)?; + dict.set_item("depth", depth)?; + dict.set_item("level", level)?; + dict.set_item("scope", scope)?; + dict.set_item("parent", parent)?; + dict.set_item("max_fallbacks", max_fallbacks)?; + if let Some(model) = unroutable { + dict.set_item("unroutable", to_py(py, &model)?)?; + } + } + RouterEvent::Blocked { + level, + model, + blocked, + } => { + dict.set_item("kind", "blocked")?; + dict.set_item("level", level)?; + dict.set_item("model", model)?; + dict.set_item("blocked", <&'static str>::from(blocked))?; + } + RouterEvent::FallbackDone { + from, + to, + level, + scope, + parent, + succeeded, + } => { + dict.set_item("kind", "fallback_done")?; + dict.set_item("from", from)?; + dict.set_item("to", to)?; + dict.set_item("level", level)?; + dict.set_item("scope", scope)?; + dict.set_item("parent", parent)?; + dict.set_item("succeeded", succeeded)?; + } + RouterEvent::RetryLogged { + scope, + level, + model, + attempted_retries, + cause, + } => { + dict.set_item("kind", "retry_logged")?; + dict.set_item("scope", scope)?; + dict.set_item("level", level)?; + dict.set_item("model", model)?; + dict.set_item("attempted_retries", attempted_retries)?; + cause_items(&dict, cause)?; + } + RouterEvent::RetryPolicyRaised { + level, + model, + cause, + } => { + dict.set_item("kind", "retry_policy_raised")?; + dict.set_item("level", level)?; + dict.set_item("model", model)?; + cause_items(&dict, cause)?; + } + RouterEvent::RetryCountRaised { + level, + model, + cause, + } => { + dict.set_item("kind", "retry_count_raised")?; + dict.set_item("level", level)?; + dict.set_item("model", model)?; + cause_items(&dict, cause)?; + } + RouterEvent::RetriesOpened { + scope, + max_retries, + model_group_size, + } => { + dict.set_item("kind", "retries_opened")?; + dict.set_item("scope", scope)?; + dict.set_item("max_retries", max_retries)?; + dict.set_item("model_group_size", model_group_size)?; + } + RouterEvent::RetryBudget { scope, max_retries } => { + dict.set_item("kind", "retry_budget")?; + dict.set_item("scope", scope)?; + dict.set_item("max_retries", max_retries)?; + } + RouterEvent::PickFailed { + level, + model, + error, + } => { + dict.set_item("kind", "pick_failed")?; + dict.set_item("level", level)?; + dict.set_item("model", model)?; + pick_items(&dict, &error.0)?; + } + RouterEvent::LevelFailed { + level, + model, + error, + handling, + } => { + dict.set_item("kind", "level_failed")?; + dict.set_item("level", level)?; + dict.set_item("model", model)?; + dict.set_item("handling", handling)?; + pick_items(&dict, &error.0)?; + } + RouterEvent::MockFailed { + mock, + level, + model, + num_retries, + } => { + dict.set_item("kind", "mock_failed")?; + dict.set_item("mock", <&'static str>::from(mock))?; + dict.set_item("level", level)?; + dict.set_item("model", model)?; + dict.set_item("num_retries", num_retries)?; + } + RouterEvent::CooldownSet { + deployment_id, + cooldown, + } => { + dict.set_item("kind", "cooldown_set")?; + dict.set_item("deployment_id", deployment_id)?; + dict.set_item("cooldown_time", cooldown.as_secs_f64())?; + } + RouterEvent::Selected { + deployment_id, + pick, + .. + } => { + dict.set_item("kind", "selected")?; + dict.set_item("deployment_id", deployment_id)?; + dict.set_item("pick_seconds", pick.as_secs_f64())?; + } + _ => continue, + } + list.append(dict)?; + } + Ok(list) +} + +fn step_dict(py: Python<'_>, step: Step) -> PyResult> { + let dict = PyDict::new(py); + let kind = match &step { + Step::HostCall(call) => call.kind(), + Step::Done(_) => StepKind::Done, + Step::Failed(_) => StepKind::Failed, + }; + dict.set_item("kind", <&'static str>::from(kind))?; + match step { + Step::HostCall(call) => host_call_items(py, &dict, call)?, + Step::Done(finish) => done_items(&dict, finish)?, + Step::Failed(error) => { + let attempt = error.attempt(); + error_items(&dict, &error)?; + dict.set_item("attempt_failed", attempt.is_some())?; + dict.set_item("status", attempt.and_then(|failure| failure.status))?; + dict.set_item( + "failure_kind", + attempt.map(|failure| <&'static str>::from(failure.kind)), + )?; + } + } + Ok(dict) +} + +fn host_call_items(py: Python<'_>, dict: &Bound<'_, PyDict>, call: HostCall) -> PyResult<()> { + match call { + HostCall::Attempt { + deployment_id, + deployment, + model, + model_group, + target_order, + candidate_ids, + retry, + fallback_depth, + fallback_level, + timeout, + params, + scope, + max_retries, + model_group_size, + } => { + dict.set_item("scope", scope)?; + dict.set_item("max_retries", max_retries)?; + dict.set_item("model_group_size", model_group_size)?; + dict.set_item("deployment_id", deployment_id)?; + dict.set_item("deployment", to_py(py, &deployment)?)?; + dict.set_item("model", model)?; + dict.set_item("model_group", model_group)?; + dict.set_item("target_order", target_order)?; + dict.set_item("candidate_ids", PyList::new(py, candidate_ids)?)?; + dict.set_item("retry", retry)?; + dict.set_item("fallback_depth", fallback_depth)?; + dict.set_item("fallback_level", fallback_level)?; + dict.set_item("timeout", timeout)?; + dict.set_item("params", params.map(|p| to_py(py, &p)).transpose()?) + } + HostCall::Sleep { seconds } => dict.set_item("seconds", seconds), + HostCall::AllowFallback { target } => dict.set_item("target", target), + HostCall::RefuseFallbacks { level, targets } => { + dict.set_item("level", level)?; + let targets = targets.0.iter().cloned().map(Value::Object).collect(); + dict.set_item("targets", to_py(py, &Value::Array(targets))?) + } + HostCall::Select { + stage, + model_group, + name, + scope, + group, + pool, + candidates, + } => { + let (stage, strategy) = match stage { + Stage::Plugins => ("plugins", None), + Stage::Strategy(strategy) => ("strategy", Some(strategy)), + Stage::Tags => ("tags", None), + Stage::PreCall => ("pre_call", None), + Stage::TextOnly => ("text_only", None), + Stage::RouteIncapable { handler } => { + dict.set_item("handler", handler)?; + ("route_incapable", None) + } + Stage::MediaRefusing => ("media_refusing", None), + Stage::ProviderRequest => ("provider_request", None), + Stage::CallbackFilters { + excluded_ids, + skipped_ids, + target_order, + } => { + dict.set_item("excluded_ids", excluded_ids)?; + dict.set_item("skipped_ids", skipped_ids)?; + dict.set_item("target_order", target_order)?; + ("callback_filters", None) + } + Stage::RoutingPlugins { run } => { + dict.set_item("run", run)?; + ("routing_plugins", None) + } + }; + dict.set_item("pool", PyList::new(py, pool)?)?; + dict.set_item("stage", stage)?; + dict.set_item("strategy", strategy)?; + dict.set_item("model_group", model_group)?; + dict.set_item("model", name)?; + dict.set_item("scope", scope)?; + dict.set_item("group", group)?; + dict.set_item("candidates", PyList::new(py, candidates)?) + } + HostCall::PreRoute { + model, + fallback_depth, + retry, + } => { + dict.set_item("model", model)?; + dict.set_item("fallback_depth", fallback_depth)?; + dict.set_item("retry", retry) + } + HostCall::Cache { + op, + key, + keys, + value, + ttl, + local, + } => { + dict.set_item("op", <&'static str>::from(op))?; + dict.set_item("key", key)?; + dict.set_item("keys", keys.map(|keys| PyList::new(py, keys)).transpose()?)?; + dict.set_item("value", value.map(|value| to_py(py, &value)).transpose()?)?; + dict.set_item("ttl", ttl)?; + dict.set_item("local", local) + } + HostCall::Secrets { index, refs } => { + dict.set_item("index", index)?; + let refs = refs + .into_iter() + .map(|secret| { + let item = PyDict::new(py); + item.set_item("name", secret.name)?; + match secret.at { + SecretAt::Param(key) => { + item.set_item("key", key)?; + item.set_item("env_only", false)?; + } + SecretAt::DataSource { source, key } => { + item.set_item("source", source)?; + item.set_item("key", key)?; + item.set_item("env_only", true)?; + } + } + Ok(item) + }) + .collect::>>()?; + dict.set_item("refs", PyList::new(py, refs)?) + } + HostCall::Build { + kind: _, + path, + index, + entry, + declared_id, + duplicate_id, + provider, + } => { + dict.set_item("path", <&'static str>::from(path))?; + dict.set_item("index", index)?; + dict.set_item("entry", to_py(py, &entry)?)?; + dict.set_item("declared_id", declared_id)?; + dict.set_item("duplicate_id", duplicate_id)?; + dict.set_item("provider", provider) + } + HostCall::ModelInfo { lookup } => { + if let Value::Object(lookup) = lookup { + for (key, value) in lookup { + dict.set_item(key, to_py(py, &value)?)?; + } + } + Ok(()) + } + } +} + +fn done_items(dict: &Bound<'_, PyDict>, finish: Finish) -> PyResult<()> { + match finish { + Finish::Call(response) => { + dict.set_item("deployment_id", response.deployment_id)?; + dict.set_item("model_group", response.model_group)?; + dict.set_item("attempted_retries", response.attempted_retries)?; + dict.set_item("max_retries", response.max_retries)?; + dict.set_item("attempted_fallbacks", response.attempted_fallbacks)?; + dict.set_item("total_tokens", response.output.total_tokens) + } + Finish::Selected { + deployment_id, + model_group, + deployment, + } => { + dict.set_item("deployment_id", deployment_id)?; + dict.set_item("model_group", model_group)?; + dict.set_item("deployment", to_py(dict.py(), &deployment)?) + } + Finish::Reported { cooldown } => { + dict.set_item("cooled", cooldown.is_some())?; + dict.set_item("cooldown_time", cooldown.map(|c| c.as_secs_f64())) + } + Finish::Built(built) => { + let py = dict.py(); + let records = built + .into_iter() + .map(|deployment| { + let item = PyDict::new(py); + item.set_item("index", deployment.index)?; + item.set_item( + "id", + deployment.spec.model_info.get("id").and_then(Value::as_str), + )?; + item.set_item("provider", deployment.spec.provider)?; + Ok(item) + }) + .collect::>>()?; + dict.set_item("built", PyList::new(py, records)?) + } + Finish::Counted { key } => dict.set_item("key", key), + Finish::RouterModelInfo(merged) => { + let py = dict.py(); + match merged { + Some(merged) => { + dict.set_item("info", to_py(py, &Value::Object(merged.info))?)?; + dict.set_item("user_keys", merged.user_keys) + } + None => dict.set_item("info", py.None()), + } + } + Finish::DeploymentModelInfo(info) => { + let py = dict.py(); + match info { + Some(info) => dict.set_item("info", to_py(py, &Value::Object(info))?), + None => dict.set_item("info", py.None()), + } + } + Finish::Value(value) => dict.set_item("value", to_py(dict.py(), &value)?), + Finish::Healthy { + single, + deployment_ids, + deployments, + } => { + dict.set_item("single", single)?; + dict.set_item("deployment_ids", deployment_ids)?; + let records = deployments + .iter() + .map(|d| to_py(dict.py(), d)) + .collect::>>()?; + dict.set_item("deployments", PyList::new(dict.py(), records)?) + } + } +} + +#[cfg(test)] +mod tests { + use pyo3::{ffi::c_str, prelude::*, types::PyDict}; + use rstest::rstest; + + const CONFIG: &str = r#"{"model_list": [ + {"model_name": "primary", "litellm_params": {"model": "openai/a"}, "model_info": {"id": "a1"}} + ]}"#; + + const PAIR: &str = r#"{"model_list": [ + {"model_name": "primary", "litellm_params": {"model": "openai/a"}, "model_info": {"id": "a1"}}, + {"model_name": "primary", "litellm_params": {"model": "openai/b"}, "model_info": {"id": "a2"}} + ]}"#; + + #[rstest] + fn python_drives_calls_and_selections_through_the_handle() { + Python::initialize(); + Python::attach(|py| { + let locals = PyDict::new(py); + locals.set_item("native", crate::native_module(py)).unwrap(); + locals.set_item("config", CONFIG).unwrap(); + locals.set_item("pair", PAIR).unwrap(); + py.run( + c_str!( + r#" +import time +store = {} +asked = [] + +def answer_cache(step): + asked.append(step) + op, key = step["op"], step["key"] + if op == "get": + return store.get(key) + if op == "batch_get": + return [store.get(k) for k in step["keys"]] + if op == "set": + store[key] = step["value"] + return None + if op == "increment": + store[key] = store.get(key, 0) + step["value"] + return store[key] + if op == "delete": + store.pop(key, None) + return None + raise AssertionError(op) + +def drive(call, outcome): + step = call.next(outcome) + while step["kind"] == "cache": + step = call.next(answer_cache(step)) + return step + +handle = native.RouterCoreHandle(config, {"num_retries": 1}) +picked = drive(handle.select("primary", {}), None) +assert (picked["kind"], picked["deployment_id"], picked["model_group"]) == ("done", "a1", "primary"), picked +assert [step["op"] for step in asked] == ["batch_get"] +assert list(asked[0]) == ["kind", "op", "key", "keys", "value", "ttl", "local", "events"], list(asked[0]) +call = handle.start("acompletion", '{"model": "primary"}') +version = handle.registry_version() +assert handle.apply({"delete": ["gone"]}) == version + 1 +steps = [drive(call, None)] +steps.append(drive(call, {"ok": False, "kind": "RateLimitError", "status": 429})) +steps.append(drive(call, None)) +steps.append(drive(call, {"ok": True, "total_tokens": 5})) +kinds = [step["kind"] for step in steps] + +cooled = native.RouterCoreHandle(pair, {"num_retries": 0}) +call = cooled.start("acompletion", {"model": "primary"}) +attempt = drive(call, None) +tried = attempt["deployment_id"] +assert attempt["deployment"]["model_info"]["id"] == tried, attempt +assert attempt["deployment"]["litellm_params"]["model"].startswith("openai/"), attempt +assert [d["model_info"]["id"] for d in cooled.deployments()] == ["a1", "a2"] +assert cooled.deployment("a2")["model_name"] == "primary" +assert cooled.resolve("primary")["ids"] == ["a1", "a2"] +assert cooled.resolve("nope") is None +assert drive(call, {"ok": False, "kind": "RateLimitError", "status": 429})["kind"] == "failed" +# The core records the failure itself: the 429 cools the deployment down in the host's store. +assert f"deployment:{tried}:cooldown" in store, store +other = drive(cooled.select("primary", {}), None) +assert other["kind"] == "done" and other["deployment_id"] != tried, other + +call = handle.select("primary", {}) +step = call.next(None) +assert step["op"] == "batch_get" +for bad in ([], [None, None], "nope", None): + try: + call.next(bad) + except ValueError: + pass + else: + raise AssertionError(f"{bad!r} should not answer a batch_get of one key") +assert drive(call, [None])["kind"] == "done" +"# + ), + Some(&locals), + Some(&locals), + ) + .unwrap(); + let kinds: Vec = locals + .get_item("kinds") + .unwrap() + .unwrap() + .extract() + .unwrap(); + assert_eq!(kinds, ["attempt", "sleep", "attempt", "done"]); + }); + } + + #[rstest] + #[case::unknown_key(r#"{"routing_strategy_args": {}}"#, true)] + #[case::bad_json("{", false)] + #[case::not_an_object("[]", false)] + fn construction_failures_raise_the_class_python_reads( + #[case] config: &str, + #[case] declined: bool, + ) { + Python::initialize(); + Python::attach(|py| { + let error = crate::native_module(py) + .getattr("RouterCoreHandle") + .and_then(|handle| handle.call1((config, "{}"))) + .unwrap_err(); + assert_eq!( + error.is_instance_of::(py), + declined, + "{error}" + ); + assert_eq!( + error.is_instance_of::(py), + !declined, + "{error}" + ); + }); + } +} diff --git a/litellm-rust/crates/python-bridge/src/router/retries.rs b/litellm-rust/crates/python-bridge/src/router/retries.rs new file mode 100644 index 00000000000..949aae96b34 --- /dev/null +++ b/litellm-rust/crates/python-bridge/src/router/retries.rs @@ -0,0 +1,55 @@ +//! The retry decisions `RouterCoreHandle` exposes as members: `get_num_retries_from_retry_policy` +//! and `should_retry_this_error`, decided by router-core for a failure Python describes the way it +//! describes a failed attempt. + +use litellm_router_core::{ + AttemptError, Router, + retry::{RetryContext, resolve_retry_policy, should_retry as core_should_retry}, +}; +use pyo3::{exceptions::PyValueError, prelude::*, types::PyDict}; + +fn failure(outcome: &Bound<'_, PyDict>) -> PyResult { + match super::attempt_outcome(outcome)? { + Err(failure) => Ok(failure), + Ok(_) => Err(PyValueError::new_err("a failure needs `ok: False`")), + } +} + +pub(crate) fn policy_retries( + router: &Router, + outcome: &Bound<'_, PyDict>, + model_group: Option<&str>, +) -> PyResult> { + let failure = failure(outcome)?; + let settings = router.settings(); + let policy = match model_group { + Some(group) => resolve_retry_policy( + settings.retry_policy.as_ref(), + group, + Some(&settings.model_group_retry_policy), + ), + None => settings.retry_policy.as_ref(), + }; + Ok(policy.and_then(|policy| policy.retries_for(&failure))) +} + +pub(crate) fn allowed_fails(router: &Router, outcome: &Bound<'_, PyDict>) -> PyResult> { + Ok(router.settings().allowed_fails_for(&failure(outcome)?)) +} + +pub(crate) fn should_retry( + outcome: &Bound<'_, PyDict>, + (healthy, all): (usize, usize), + (context_window, content_policy, regular): (bool, bool, bool), +) -> PyResult { + Ok(core_should_retry( + &failure(outcome)?, + RetryContext { + healthy, + all, + has_context_window_fallbacks: context_window, + has_content_policy_fallbacks: content_policy, + has_regular_fallbacks: regular, + }, + )) +} diff --git a/litellm-rust/crates/python-bridge/src/router/stream.rs b/litellm-rust/crates/python-bridge/src/router/stream.rs new file mode 100644 index 00000000000..f85d73e0778 --- /dev/null +++ b/litellm-rust/crates/python-bridge/src/router/stream.rs @@ -0,0 +1,94 @@ +//! The stream-recovery entry points of `RouterCoreHandle`: what may still take over a stream that +//! failed before content, the same-group stream retry, and the fallback walk re-entered from a +//! stream. The host's stream wrappers call them; the calls are driven like `start`. + +use litellm_router_core::{ + FailureKind, Op, Router, + call::{StreamOrigin, StreamPolicy}, +}; +use pyo3::{exceptions::PyValueError, prelude::*, types::PyDict}; + +use super::{RouterCall, attempt_outcome, config, config_error_to_pyerr, document, host::Finish}; + +/// A stream recovery a host wrapper asks for. +#[derive(Clone, Copy)] +pub(super) enum Recovery { + /// `_astream_retry_same_group`. + Retry, + /// The fallback walk re-entered from the stream wrapper. + Fallback, +} + +fn origin(value: &Bound<'_, PyAny>) -> PyResult { + serde_json::from_value(document(value)?) + .map_err(|error| PyValueError::new_err(format!("a stream origin: {error}"))) +} + +fn kind(kind: Option<&str>) -> PyResult> { + kind.map(|kind| { + kind.parse() + .map_err(|_| PyValueError::new_err(format!("`{kind}` is not a router failure kind"))) + }) + .transpose() +} + +/// `{"can_retry", "can_fall_back", "refusal_fallback", "retry_budget": int | None, +/// "policy_applies": bool}`. +pub(super) fn policy<'py>( + py: Python<'py>, + router: &Router, + request: &Bound<'py, PyAny>, + stream_origin: &Bound<'py, PyAny>, + failure_kind: Option<&str>, +) -> PyResult> { + let request = config::parse_request(document(request)?, None).map_err(config_error_to_pyerr)?; + let StreamPolicy { + can_retry, + can_fall_back, + refusal_fallback, + retry_budget, + } = router.stream_policy(&request, &origin(stream_origin)?, kind(failure_kind)?); + let dict = PyDict::new(py); + dict.set_item("can_retry", can_retry)?; + dict.set_item("can_fall_back", can_fall_back)?; + dict.set_item("refusal_fallback", refusal_fallback)?; + dict.set_item("retry_budget", retry_budget.map(|(budget, _)| budget))?; + dict.set_item( + "policy_applies", + retry_budget.is_some_and(|(_, applies)| applies), + )?; + Ok(dict) +} + +/// A recovery call for a stream from `stream_origin` that failed with `failure`. +pub(super) fn recover( + router: &Router, + recovery: Recovery, + op: &str, + request: &Bound<'_, PyAny>, + stream_origin: &Bound<'_, PyAny>, + failure: &Bound<'_, PyDict>, +) -> PyResult { + let op = Op::new(op); + let request = config::parse_request(document(request)?, None).map_err(config_error_to_pyerr)?; + let stream_origin = origin(stream_origin)?; + let Err(failure) = attempt_outcome(failure)? else { + return Err(PyValueError::new_err("a stream failure needs `ok: False`")); + }; + let router = router.clone(); + Ok(RouterCall::new(async move { + let response = match recovery { + Recovery::Retry => { + router + .retry_stream(op, request, stream_origin, failure) + .await + } + Recovery::Fallback => { + router + .fallback_stream(op, request, stream_origin, failure) + .await + } + }; + response.map(Finish::Call) + })) +} diff --git a/litellm-rust/crates/router-core/Cargo.toml b/litellm-rust/crates/router-core/Cargo.toml index af97bd52c57..ff5e2e8161f 100644 --- a/litellm-rust/crates/router-core/Cargo.toml +++ b/litellm-rust/crates/router-core/Cargo.toml @@ -7,3 +7,21 @@ repository.workspace = true description = "The router core behind litellm.Router: registry, settings, health, selection and the call loop" [dependencies] +arc-swap.workspace = true +async-trait.workspace = true +litellm-cache.workspace = true +litellm-router-types.workspace = true +jiff = { version = "0.2.37", default-features = false, features = ["std", "tzdb-zoneinfo", "tzdb-bundle-always"] } +parking_lot.workspace = true +rand.workspace = true +serde.workspace = true +serde_with.workspace = true +serde_json = { workspace = true, features = ["preserve_order"] } +sha2.workspace = true +strum.workspace = true +thiserror.workspace = true +tracing.workspace = true + +[dev-dependencies] +rstest.workspace = true +tokio.workspace = true diff --git a/litellm-rust/crates/router-core/src/access.rs b/litellm-rust/crates/router-core/src/access.rs new file mode 100644 index 00000000000..1b91fa417da --- /dev/null +++ b/litellm-rust/crates/router-core/src/access.rs @@ -0,0 +1,236 @@ +//! Who may reach a deployment: model access windows (`litellm_core_utils/access_windows.py`, +//! a deployment reserved for some teams during a daily wall-clock window) and model access +//! groups (`_filter_deployments_by_model_access_groups`, a caller allowed only through access +//! groups is kept to its groups' deployments). + +use std::{collections::BTreeSet, sync::Arc}; + +use jiff::{Timestamp, tz::TimeZone}; +use serde_json::Value; + +use crate::registry::Deployment; + +/// One `model_info.access_windows` entry (`ModelAccessWindow`), validated when the deployment +/// was added. +#[derive(Clone, Debug, PartialEq)] +pub struct AccessWindow { + /// Microseconds after midnight. + pub start: u64, + pub end: u64, + pub timezone: String, + pub team_ids: Vec, +} + +const MICROS_PER_SECOND: u64 = 1_000_000; + +/// A `datetime.time` as pydantic reads it: `HH:MM[:SS[.ffffff]]`, or seconds after midnight. +fn wall_clock(value: &Value) -> Option { + match value { + Value::Number(seconds) => { + let seconds = seconds.as_f64()?; + (0.0..86_400.0) + .contains(&seconds) + .then(|| (seconds * MICROS_PER_SECOND as f64).round() as u64) + } + Value::String(text) => { + let mut parts = text.splitn(3, ':'); + let hour: u64 = parts.next()?.parse().ok()?; + let minute: u64 = parts.next()?.parse().ok()?; + let (second, micros) = match parts.next() { + None => (0, 0), + Some(rest) => { + let (whole, fraction) = rest.split_once('.').unwrap_or((rest, "")); + let digits: String = fraction.chars().take(6).collect(); + let micros = if digits.is_empty() { + 0 + } else { + format!("{digits:0<6}").parse().ok()? + }; + (whole.parse().ok()?, micros) + } + }; + (hour < 24 && minute < 60 && second < 60) + .then(|| ((hour * 60 + minute) * 60 + second) * MICROS_PER_SECOND + micros) + } + _ => None, + } +} + +impl AccessWindow { + fn parse(value: &Value) -> Option { + let window = value.as_object()?; + Some(Self { + start: wall_clock(window.get("start")?)?, + end: wall_clock(window.get("end")?)?, + timezone: window.get("timezone")?.as_str()?.to_owned(), + team_ids: window + .get("team_ids")? + .as_array()? + .iter() + .filter_map(Value::as_str) + .map(str::to_owned) + .collect(), + }) + } + + /// `is_window_active`: the window's wall clock contains `now`, wrapping past midnight when + /// it ends before it starts. + pub fn is_active(&self, now: Timestamp) -> bool { + let Ok(zone) = TimeZone::get(&self.timezone) else { + return false; + }; + let local = now.to_zoned(zone).time(); + let micros = ((u64::from(local.hour().unsigned_abs()) * 60 + + u64::from(local.minute().unsigned_abs())) + * 60 + + u64::from(local.second().unsigned_abs())) + * MICROS_PER_SECOND + + u64::from(local.subsec_nanosecond().unsigned_abs()) / 1_000; + if self.start < self.end { + self.start <= micros && micros < self.end + } else { + micros >= self.start || micros < self.end + } + } + + /// The window's end as `{end:%H:%M}`. + pub fn end_hhmm(&self) -> String { + let minutes = self.end / MICROS_PER_SECOND / 60; + format!("{:02}:{:02}", minutes / 60, minutes % 60) + } +} + +/// A deployment's access windows. +pub fn access_windows(deployment: &Deployment) -> Vec { + deployment + .model_info + .raw + .get("access_windows") + .and_then(Value::as_array) + .map(|windows| windows.iter().filter_map(AccessWindow::parse).collect()) + .unwrap_or_default() +} + +/// `_reservation_blocking_window`: the first active window, unless the caller's team is named +/// by one of the active windows. +fn blocking_window( + deployment: &Deployment, + team_id: Option<&str>, + now: Timestamp, +) -> Option { + let active: Vec = access_windows(deployment) + .into_iter() + .filter(|window| window.is_active(now)) + .collect(); + if let Some(team) = team_id + && active + .iter() + .any(|window| window.team_ids.iter().any(|id| id == team)) + { + return None; + } + active.into_iter().next() +} + +/// `filter_reserved_deployments`: the deployments no active window withholds from the caller, +/// and the first window that withheld one. +pub fn filter_reserved( + deployments: Vec>, + team_id: Option<&str>, + now: Timestamp, +) -> (Vec>, Option) { + let mut blocking = None; + let kept = deployments + .into_iter() + .filter( + |deployment| match blocking_window(deployment, team_id, now) { + None => true, + Some(window) => { + blocking.get_or_insert(window); + false + } + }, + ) + .collect(); + (kept, blocking) +} + +/// `_filter_deployments_by_model_access_groups`: a caller whose allowed models (`models` plus +/// `team_models`) reach `model` only through access groups is kept to the deployments in those +/// groups. `model_groups` are the access groups the rows listed for `model` carry. +pub fn filter_by_access_groups( + model: &str, + deployments: Vec>, + allowed_models: &[String], + model_groups: impl FnOnce() -> BTreeSet, +) -> Vec> { + if deployments.is_empty() || allowed_models.is_empty() { + return deployments; + } + let allowed = |name: &str| allowed_models.iter().any(|m| m == name); + if allowed(model) || allowed("*") || allowed("all-proxy-models") { + return deployments; + } + let groups: BTreeSet = model_groups() + .into_iter() + .filter(|group| allowed(group)) + .collect(); + if groups.is_empty() { + return deployments; + } + deployments + .into_iter() + .filter(|d| { + d.model_info + .access_groups + .iter() + .any(|g| groups.contains(g)) + }) + .collect() +} + +#[cfg(test)] +mod tests { + use rstest::rstest; + use serde_json::json; + + use super::*; + + #[rstest] + #[case::minutes(json!("09:30"), Some(34_200_000_000))] + #[case::seconds(json!("09:30:15"), Some(34_215_000_000))] + #[case::fraction(json!("00:00:00.5"), Some(500_000))] + #[case::number(json!(60), Some(60_000_000))] + #[case::bad(json!("25:00"), None)] + fn wall_clock_reads_pydantic_times(#[case] raw: Value, #[case] expected: Option) { + assert_eq!(wall_clock(&raw), expected); + } + + #[rstest] + #[case::inside("2026-01-01T10:00:00Z", true)] + #[case::at_end("2026-01-01T17:00:00Z", false)] + #[case::before("2026-01-01T08:59:59Z", false)] + fn a_window_is_active_inside_its_wall_clock(#[case] now: &str, #[case] expected: bool) { + let window = AccessWindow { + start: wall_clock(&json!("09:00")).unwrap(), + end: wall_clock(&json!("17:00")).unwrap(), + timezone: "UTC".to_owned(), + team_ids: vec!["t".to_owned()], + }; + assert_eq!(window.is_active(now.parse().unwrap()), expected); + assert_eq!(window.end_hhmm(), "17:00"); + } + + #[test] + fn an_overnight_window_wraps_past_midnight() { + let window = AccessWindow { + start: wall_clock(&json!("22:00")).unwrap(), + end: wall_clock(&json!("02:00")).unwrap(), + timezone: "UTC".to_owned(), + team_ids: vec!["t".to_owned()], + }; + assert!(window.is_active("2026-01-01T23:00:00Z".parse().unwrap())); + assert!(window.is_active("2026-01-01T01:00:00Z".parse().unwrap())); + assert!(!window.is_active("2026-01-01T12:00:00Z".parse().unwrap())); + } +} diff --git a/litellm-rust/crates/router-core/src/build.rs b/litellm-rust/crates/router-core/src/build.rs new file mode 100644 index 00000000000..36bad012b3a --- /dev/null +++ b/litellm-rust/crates/router-core/src/build.rs @@ -0,0 +1,386 @@ +//! The deployment build: how configured entries become the registry's deployments. +//! +//! The core owns the loop (`set_model_list`, `_create_deployment` and `_add_deployment`): which +//! values are secret references, the id of a deployment that declares none, one deployment per +//! `organization`, and the order everything happens in. What needs the embedding process is a +//! host step, asked for one deployment at a time so an error surfaces exactly where the Python +//! router raises it: +//! - [`Host::secrets`]: the `os.environ/` references the entry holds; +//! - [`Host::validate`]: the host's checks, and the entry as it normalizes it; +//! - [`Host::provider`]: the provider that serves it; +//! - [`Host::added`]: the host's side effects, and the entry as the registry stores it. + +use serde_json::{Map, Value}; + +use crate::{ + error::{ConfigError, Error}, + host::{BuildAsk, BuildPath, Built, Host, Provided, SecretAt, SecretRef}, + registry::{DeploymentSpec, deployment_id}, +}; + +/// The `litellm_params` keys whose `os.environ/` value is a secret pointer the identity source +/// dereferences at use time, so the build leaves it as written. The router's copy of +/// `_WIF_POINTER_FIELDS` (litellm/types/router.py), which the proxy's config loader also reads. +pub const SECRET_POINTER_PARAMS: &[&str] = &[ + "anthropic_issuer_signing_key_ref", + "anthropic_keycloak_client_secret_ref", +]; + +const SECRET_PREFIX: &str = "os.environ/"; + +/// How a build's deployments join the registry. +#[derive(Clone, Copy, Debug, PartialEq, Eq, strum::IntoStaticStr, strum::EnumString)] +#[strum(serialize_all = "snake_case")] +pub enum Install { + /// They become every deployment (`set_model_list`). + Replace, + /// They are added at the end (a deployment retried once its provider resolves). + Append, + /// Each replaces the deployment with its id, if any, and is added at the end + /// (`add_deployment`, an upsert). + Upsert, +} + +/// One deployment the build accepted, as the registry stores it. +#[derive(Clone, Debug, PartialEq)] +pub struct BuiltDeployment { + /// The entry index in the list the build was given. + pub index: usize, + pub spec: DeploymentSpec, +} + +/// Build `entries` (model_list entries) through `host`, in order. `resolve_secrets` resolves +/// top-level `os.environ/` references first, as `set_model_list` does; `add_deployment` and +/// upserts take their params as given. `install` takes each deployment as soon as it is +/// accepted, so it joins the registry before the next one is built, as the Python router adds +/// each in turn. +pub async fn build( + host: &dyn Host, + path: BuildPath, + entries: Vec>, + resolve_secrets: bool, + mut install: impl FnMut(&BuiltDeployment) -> Result<(), Error> + Send, +) -> Result, Error> { + let declared: Vec> = entries.iter().map(declared_id).collect(); + let mut built = Vec::new(); + for (index, mut entry) in entries.into_iter().enumerate() { + let model_name = text(entry.get("model_name")); + let mut params = match entry.remove("litellm_params") { + Some(Value::Object(params)) => params, + _ => Map::new(), + }; + if resolve_secrets { + let refs = param_refs(¶ms); + if !refs.is_empty() { + let values = host.secrets(index, &refs).await?; + for (secret, value) in refs.into_iter().zip(values) { + if let SecretAt::Param(key) = secret.at { + params.insert(key, value); + } + } + } + } + let mut info = match entry.remove("model_info") { + Some(Value::Object(info)) => info, + _ => Map::new(), + }; + let declared_id = declared[index].clone(); + if !info.contains_key("id") { + info.insert( + "id".to_owned(), + Value::String(deployment_id(&model_name, ¶ms)), + ); + } + // `set_model_list` adds one deployment per organization of a list, each with the + // organization set on the same params (so every one keeps the list's id). + let organizations = match (path, params.get("organization")) { + (BuildPath::Create, Some(Value::Array(organizations))) => organizations.clone(), + _ => vec![Value::Null], + }; + let duplicate_id = path == BuildPath::Create + && declared_id.as_ref().is_some_and(|id| { + declared + .iter() + .filter(|other| other.as_ref() == Some(id)) + .count() + > 1 + }); + for organization in organizations { + if !organization.is_null() { + params.insert("organization".to_owned(), organization); + } + let mut deployment = entry.clone(); + deployment.insert("model_name".to_owned(), Value::String(model_name.clone())); + deployment.insert("litellm_params".to_owned(), Value::Object(params.clone())); + deployment.insert("model_info".to_owned(), Value::Object(info.clone())); + if let Some(spec) = build_one( + host, + path, + index, + &deployment, + declared_id.as_deref(), + duplicate_id, + ) + .await? + { + let deployment = BuiltDeployment { index, spec }; + install(&deployment)?; + built.push(deployment); + } + } + } + Ok(built) +} + +/// One deployment's steps: validate, provider, the data-source references, added. +async fn build_one( + host: &dyn Host, + path: BuildPath, + index: usize, + entry: &Map, + declared_id: Option<&str>, + duplicate_id: bool, +) -> Result, Error> { + let ask = |entry| BuildAsk { + path, + index, + entry, + declared_id, + duplicate_id, + }; + let validated = match host.validate(&ask(entry)).await? { + Built::Entry(validated) => validated, + Built::Skip | Built::Invalid => return Ok(None), + }; + let provider = match host.provider(&ask(&validated)).await? { + Provided::Provider(provider) => provider, + Provided::Invalid => return Ok(None), + }; + let refs = data_source_refs(&validated); + if !refs.is_empty() { + let _ = host.secrets(index, &refs).await?; + } + let stored = match host.added(&ask(&validated), provider.as_deref()).await? { + Built::Entry(stored) => stored, + Built::Skip | Built::Invalid => return Ok(None), + }; + let model_name = text(stored.get("model_name")); + let mut spec: DeploymentSpec = serde_json::from_value(Value::Object(stored)) + .map_err(|source| ConfigError::InvalidDeployment { model_name, source })?; + spec.provider = provider; + Ok(Some(spec)) +} + +/// The top-level `litellm_params` values `set_model_list` resolves. +fn param_refs(params: &Map) -> Vec { + params + .iter() + .filter(|(key, _)| !SECRET_POINTER_PARAMS.contains(&key.as_str())) + .filter_map(|(key, value)| { + let name = value + .as_str() + .filter(|name| name.starts_with(SECRET_PREFIX))?; + Some(SecretRef { + name: name.to_owned(), + at: SecretAt::Param(key.clone()), + }) + }) + .collect() +} + +/// The `dataSources[].parameters.endpoint` / `key` references `_add_deployment` reads from the +/// environment. +fn data_source_refs(entry: &Map) -> Vec { + let sources = entry + .get("litellm_params") + .and_then(|params| params.get("dataSources")) + .and_then(Value::as_array); + let mut refs = Vec::new(); + for (source, data_source) in sources.into_iter().flatten().enumerate() { + let Some(parameters) = data_source.get("parameters").and_then(Value::as_object) else { + continue; + }; + for key in ["endpoint", "key"] { + if let Some(name) = parameters + .get(key) + .and_then(Value::as_str) + .filter(|name| name.starts_with(SECRET_PREFIX)) + { + refs.push(SecretRef { + name: name.to_owned(), + at: SecretAt::DataSource { + source, + key: key.to_owned(), + }, + }); + } + } + } + refs +} + +/// The `model_info.id` an entry declares, as Python's `str()` spells it. +fn declared_id(entry: &Map) -> Option { + match entry.get("model_info")?.get("id")? { + Value::Null => None, + Value::String(id) => Some(id.clone()), + other => Some(other.to_string()), + } +} + +/// A name as Python's `str()` spells it. +fn text(value: Option<&Value>) -> String { + match value { + Some(Value::String(text)) => text.clone(), + Some(Value::Null) | None => "None".to_owned(), + Some(other) => other.to_string(), + } +} + +#[cfg(test)] +mod tests { + use std::sync::{Arc, Mutex}; + + use async_trait::async_trait; + use serde_json::json; + + use super::*; + use crate::{ + cache::{InMemoryRouterCache, RouterCache}, + error::AttemptError, + host::{Attempt, Output, RouterEvent, RouterRequest, RoutingContext, SelectStage}, + registry::Deployment, + }; + + /// Records every build step and resolves `os.environ/X` to `secret-X`. + #[derive(Default)] + struct Recording { + steps: Mutex>, + } + + #[async_trait] + impl Host for Recording { + async fn invoke(&self, _attempt: &Attempt) -> Result { + unreachable!("the build runs no attempt") + } + async fn select( + &self, + _stage: &SelectStage, + _ctx: &RoutingContext<'_>, + candidates: Vec>, + ) -> Result>, Error> { + Ok(candidates) + } + async fn allow_fallback( + &self, + _request: &RouterRequest, + _target: &str, + ) -> Result { + Ok(true) + } + async fn sleep(&self, _duration: std::time::Duration) {} + fn observe(&self, _event: &RouterEvent) {} + fn cache(&self) -> Arc { + Arc::new(InMemoryRouterCache::default()) + } + async fn secrets(&self, _index: usize, refs: &[SecretRef]) -> Result, Error> { + self.steps.lock().unwrap().push(format!("secrets {refs:?}")); + Ok(refs + .iter() + .map(|r| json!(format!("secret-{}", &r.name[SECRET_PREFIX.len()..]))) + .collect()) + } + async fn validate(&self, ask: &BuildAsk<'_>) -> Result { + self.steps.lock().unwrap().push(format!( + "validate {} {}", + ask.index, ask.entry["litellm_params"] + )); + Ok(Built::Entry(ask.entry.clone())) + } + async fn provider(&self, ask: &BuildAsk<'_>) -> Result { + self.steps + .lock() + .unwrap() + .push(format!("provider {}", ask.index)); + Ok(Provided::Provider(Some("azure".to_owned()))) + } + async fn added(&self, ask: &BuildAsk<'_>, provider: Option<&str>) -> Result { + self.steps + .lock() + .unwrap() + .push(format!("added {} {provider:?}", ask.index)); + Ok(Built::Entry(ask.entry.clone())) + } + } + + fn entries(value: Value) -> Vec> { + serde_json::from_value(value).unwrap() + } + + #[tokio::test] + async fn create_resolves_secrets_then_builds_one_deployment_per_organization() { + let host = Recording::default(); + let built = build( + &host, + BuildPath::Create, + entries(json!([{ + "model_name": "g", + "litellm_params": { + "model": "azure/x", + "api_key": "os.environ/KEY", + "anthropic_issuer_signing_key_ref": "os.environ/POINTER", + "organization": ["a", "b"], + }, + }])), + true, + |_| Ok(()), + ) + .await + .unwrap(); + let steps = host.steps.lock().unwrap().clone(); + assert_eq!(steps.len(), 7, "{steps:?}"); + assert!(steps[0].starts_with("secrets") && steps[0].contains("KEY")); + assert!(!steps[0].contains("POINTER")); + assert!(steps[1].contains(r#""organization":"a""#)); + assert!(steps[1].contains("secret-KEY")); + assert_eq!(steps[3], "added 0 Some(\"azure\")"); + assert!(steps[4].contains(r#""organization":"b""#)); + assert_eq!(built.len(), 2); + assert_eq!(built[0].spec.provider.as_deref(), Some("azure")); + assert_eq!( + built[0].spec.model_info["id"], + built[1].spec.model_info["id"] + ); + let deployment = Deployment::from_spec(built[0].spec.clone()).unwrap(); + assert_eq!(deployment.provider(), "azure"); + } + + #[tokio::test] + async fn add_leaves_params_as_given_and_asks_for_data_source_references() { + let host = Recording::default(); + let built = build( + &host, + BuildPath::Add, + entries(json!([{ + "model_name": "g", + "litellm_params": { + "model": "azure/x", + "api_key": "os.environ/KEY", + "organization": ["a", "b"], + "dataSources": [{"parameters": {"endpoint": "os.environ/END", "key": "plain"}}], + }, + "model_info": {"id": "given"}, + }])), + false, + |_| Ok(()), + ) + .await + .unwrap(); + let steps = host.steps.lock().unwrap().clone(); + assert_eq!(steps.len(), 4, "{steps:?}"); + assert!(steps[0].contains("os.environ/KEY")); + assert!(steps[2].contains("DataSource { source: 0, key: \"endpoint\" }")); + assert_eq!(built.len(), 1); + assert_eq!(built[0].spec.model_info["id"], "given"); + } +} diff --git a/litellm-rust/crates/router-core/src/cache.rs b/litellm-rust/crates/router-core/src/cache.rs new file mode 100644 index 00000000000..9811e34a3e3 --- /dev/null +++ b/litellm-rust/crates/router-core/src/cache.rs @@ -0,0 +1,159 @@ +//! [`RouterCache`]: the DualCache-style store all routing state lives in, under the keys, +//! encodings and TTLs the Python router uses, so Python and Rust workers share it. Method names +//! follow `litellm-cache`'s `BaseCache`, `BatchCache`, `CounterCache` and `DeleteCache`. + +use std::{ + collections::HashMap, + sync::Arc, + time::{Duration, Instant}, +}; + +use async_trait::async_trait; +use litellm_cache::Error; +use parking_lot::Mutex; +use serde_json::Value; + +#[async_trait] +pub trait RouterCache: Send + Sync { + async fn async_get_cache(&self, key: &str) -> Result, Error>; + + async fn async_set_cache( + &self, + key: &str, + value: Value, + ttl: Option, + ) -> Result<(), Error>; + + /// Add `amount` to the number at `key` (0 when missing) and return the new value. The TTL + /// applies when the key is created, as Redis `INCRBYFLOAT` + `EXPIRE NX` does. + async fn async_increment( + &self, + key: &str, + amount: f64, + ttl: Option, + ) -> Result; + + async fn async_batch_get_cache(&self, keys: &[String]) -> Result>, Error>; + + async fn async_delete_cache(&self, key: &str) -> Result<(), Error>; + + /// Read this worker's in-memory tier only (`local_only=True`). + async fn async_get_local(&self, key: &str) -> Result, Error> { + self.async_get_cache(key).await + } + + /// Add to a counter in this worker's in-memory tier only (`local_only=True`). + async fn async_increment_local( + &self, + key: &str, + amount: f64, + ttl: Option, + ) -> Result { + self.async_increment(key, amount, ttl).await + } + + /// `_increment_deployment_usage`'s write: add each `(key, amount)` in one pipeline. With a + /// shared tier, this worker's in-memory tier then holds the shared totals, so its reads see + /// what other workers counted up to now. + async fn async_increment_usage( + &self, + increments: &[(String, f64)], + ttl: Option, + ) -> Result<(), Error> { + for (key, amount) in increments { + self.async_increment(key, *amount, ttl).await?; + } + Ok(()) + } +} + +#[derive(Default)] +struct Entry { + value: Value, + expires: Option, +} + +/// An in-process [`RouterCache`], for tests and for state Python keeps `local_only`. +#[derive(Clone, Default)] +pub struct InMemoryRouterCache { + entries: Arc>>, +} + +impl InMemoryRouterCache { + fn live(entries: &mut HashMap, key: &str) -> Option { + let expired = entries + .get(key) + .is_some_and(|entry| entry.expires.is_some_and(|at| at <= Instant::now())); + if expired { + entries.remove(key); + } + entries.get(key).map(|entry| entry.value.clone()) + } + + pub fn get(&self, key: &str) -> Option { + Self::live(&mut self.entries.lock(), key) + } + + pub fn increment(&self, key: &str, amount: f64, ttl: Option) -> f64 { + let mut entries = self.entries.lock(); + let current = Self::live(&mut entries, key).and_then(|value| value.as_f64()); + let next = current.unwrap_or(0.0) + amount; + let value = serde_json::Number::from_f64(next) + .map(|number| { + if next.fract() == 0.0 && next.abs() < i64::MAX as f64 { + Value::from(next as i64) + } else { + Value::Number(number) + } + }) + .unwrap_or(Value::Null); + let expires = match (current, entries.get(key)) { + (Some(_), Some(entry)) => entry.expires, + _ => ttl.map(|ttl| Instant::now() + ttl), + }; + entries.insert(key.to_owned(), Entry { value, expires }); + next + } +} + +#[async_trait] +impl RouterCache for InMemoryRouterCache { + async fn async_get_cache(&self, key: &str) -> Result, Error> { + Ok(self.get(key)) + } + + async fn async_set_cache( + &self, + key: &str, + value: Value, + ttl: Option, + ) -> Result<(), Error> { + let expires = ttl.map(|ttl| Instant::now() + ttl); + self.entries + .lock() + .insert(key.to_owned(), Entry { value, expires }); + Ok(()) + } + + async fn async_increment( + &self, + key: &str, + amount: f64, + ttl: Option, + ) -> Result { + Ok(self.increment(key, amount, ttl)) + } + + async fn async_batch_get_cache(&self, keys: &[String]) -> Result>, Error> { + let mut entries = self.entries.lock(); + Ok(keys + .iter() + .map(|key| Self::live(&mut entries, key)) + .collect()) + } + + async fn async_delete_cache(&self, key: &str) -> Result<(), Error> { + self.entries.lock().remove(key); + Ok(()) + } +} diff --git a/litellm-rust/crates/router-core/src/call.rs b/litellm-rust/crates/router-core/src/call.rs new file mode 100644 index 00000000000..287bb98d8b7 --- /dev/null +++ b/litellm-rust/crates/router-core/src/call.rs @@ -0,0 +1,1555 @@ +//! The call loop: select, run the attempt, then retry in the same group or fall back to another +//! until it succeeds or the policy is exhausted. Mirrors `Router.async_function_with_fallbacks`, +//! `async_function_with_fallbacks_common_utils`, `async_function_with_retries`, +//! `should_retry_this_error` and `_time_to_sleep_before_retry` in `litellm/router.py`, and +//! `run_async_fallback` in `router_utils/fallback_event_handlers.py`. + +use std::{ + borrow::Cow, + future::Future, + pin::Pin, + sync::Arc, + time::{Duration, Instant}, +}; + +use serde_json::{Map, Value}; + +use tracing::warn; + +use crate::{ + Core, RouterStateSnapshot, + config::Op, + error::{AttemptError, ConfigError, Error, FailureKind, Unsupported}, + fallback::{FallbackKeys, FallbackTarget, Fallbacks}, + health::{GroupShape, should_retry_status, team_model_has_alternatives}, + host::{ + Attempt, Blocked, Output, PickError, RetryCause, RouterEvent, RouterRequest, + UncountedRetries, + }, + host_index::python_str, + lookup::Lookup, + mock::MockFailure, + outcome::OutcomeWrite, + prerouting::{PreRoute, PreRouted}, + registry::{ResolveKind, ResolveScope}, + retry::{RetryContext, resolve_retry_policy, should_retry}, + select::{Candidates, SelectInputs}, +}; + +mod stream; + +pub use self::stream::{StreamOrigin, StreamPolicy}; + +/// `INITIAL_RETRY_DELAY`, `MAX_RETRY_DELAY` and `JITTER` in `litellm/constants.py`. +const INITIAL_RETRY_DELAY: f64 = 0.5; +const MAX_RETRY_DELAY: f64 = 8.0; +const JITTER: f64 = 0.75; +/// `calculate_retry_after` obeys a provider's retry-after up to a minute. +const MAX_HONOURED_RETRY_AFTER: Duration = Duration::from_secs(60); + +#[derive(Clone, Debug)] +pub struct RouterResponse { + pub output: Output, + pub deployment_id: String, + pub model_group: String, + pub attempted_retries: u32, + /// The retry budget of the loop a retry succeeded in (`x-litellm-max-retries`); `None` + /// when the first attempt succeeded. + pub max_retries: Option, + pub attempted_fallbacks: u32, +} + +/// A failed attempt with the deployment it ran on, when one was picked. +struct Failed { + error: Box, + deployment_id: Option, +} + +impl From for Failed { + fn from(error: Error) -> Self { + Self { + error: Box::new(error), + deployment_id: None, + } + } +} + +/// Where a level of the fallback chain sits: `depth` is `run_async_fallback`'s counter, one more +/// for every target tried, `level` how many fallbacks deep the chain is, and `scope` the level's +/// own copy of the request metadata (every hop copies its parent's). +#[derive(Clone, Copy, Debug, Default)] +struct FallbackHop { + depth: u32, + level: u32, + scope: u32, +} + +/// What the whole chain shares: the attempts made (`attempted_targets`) and the scopes handed out. +#[derive(Debug, Default)] +struct FallbackChain { + attempted: Vec, + scopes: u32, +} + +/// What one level of the chain serves: a group, or the same group at a later `order`, with the +/// request params the chain's dict targets merged so far (`kwargs.update(target)`). +#[derive(Clone, Debug)] +struct Target { + group: String, + order: Option, + params: Option>>, + /// `_excluded_deployment_ids` a weighted failover hop put on the request: never picked. + excluded: Arc<[String]>, + /// `_failover_excluded_ids` in the level's metadata: the deployments weighted failover + /// already excluded on the way here. + failover_excluded: Arc<[String]>, +} + +/// A step of a level's fallback chain: a configured target, or an order level of the group. +#[derive(Clone, Debug)] +enum Next { + Fallback(FallbackTarget), + Order(String, i64), + /// A weighted failover hop: the same group without the deployments already excluded. + Failover(String, Arc<[String]>), +} + +/// An attempt the chain already made (`fallback_attempt_key`). +#[derive(Clone, Debug, PartialEq)] +enum Tried { + Target(FallbackTarget), + Order(String, i64), + Failover(String, Vec), +} + +/// The fallback lists a call reads: the request's own where it passed them, else the router's. +struct Lists<'a> { + fallbacks: Option<&'a Fallbacks>, + context_window: Option<&'a Fallbacks>, + content_policy: Option<&'a Fallbacks>, + keys: &'a FallbackKeys, +} + +impl<'a> Lists<'a> { + fn of(snapshot: &'a RouterStateSnapshot, request: &'a RouterRequest) -> Self { + let dynamic = &snapshot.settings.dynamic; + let pick = |own: &'a Option>, router: Option<&'a Fallbacks>| match own { + Some(own) => own.as_ref(), + None => router, + }; + Self { + fallbacks: pick(&request.fallbacks, Some(&dynamic.fallbacks)), + context_window: pick( + &request.context_window_fallbacks, + dynamic.context_window_fallbacks.as_ref(), + ), + content_policy: pick( + &request.content_policy_fallbacks, + dynamic.content_policy_fallbacks.as_ref(), + ), + keys: &dynamic.fallback_keys, + } + } + + fn has_regular(&self) -> bool { + self.fallbacks.is_some_and(|list| !list.is_empty()) + } +} + +/// `has_unattempted_fallback_target` over the chain `list` resolves for the level's lookup +/// groups: before the first hop nothing counts as tried, and a dict target with its own params +/// is never skipped. +fn unattempted_chain( + list: &Fallbacks, + lists: &Lists<'_>, + request: &RouterRequest, + target: &Target, + at: FallbackHop, + attempted: &[Tried], +) -> bool { + let mut lookup = vec![target.group.as_str()]; + if request.model != target.group { + lookup.push(request.model.as_str()); + } + list.lookup_groups(&lookup, lists.keys) + .is_some_and(|targets| { + at.level == 0 + || targets.iter().any(|next| match next.identity() { + FallbackTarget::Params(_) => true, + identity => !attempted.contains(&Tried::Target(identity)), + }) + }) +} + +/// The deployments `_core_get_model_list(model_name=group, team_id=team_id)` lists: the name's own +/// (a team's public name among them), an alias's target's, a routing group's members, and only +/// when none of those serve it, the wildcard routes (`_get_wildcard_deployments`), the proxy-wide +/// pattern's and then the team's. +fn group_members( + snapshot: &RouterStateSnapshot, + group: &str, + team_id: Option<&str>, +) -> Vec> { + let registry = &snapshot.registry; + let all = |name: &str, team_id: Option<&str>| { + crate::resolve::all_deployments(registry, name, team_id).unwrap_or_default() + }; + let members = |routing_group: &crate::registry::RoutingGroupSpec| { + crate::routing_groups::members(routing_group, |member| all(member, None)) + }; + let mut rows = all(group, team_id); + if let Some(alias_target) = registry.alias_of(group) { + rows.extend(match registry.routing_group_named(alias_target) { + Some(routing_group) => members(routing_group), + None => all(alias_target, None), + }); + } + if let Some(routing_group) = registry.routing_group_named(group) { + rows.extend(members(routing_group)); + } + if rows.is_empty() { + if let Some((_, routed)) = registry.wildcard_for(group) { + rows.extend(routed); + } + if let Some(team) = team_id { + rows.extend(Lookup::new(registry).team_wildcard(group, team)); + } + } + rows +} + +/// What `_maybe_run_weighted_failover` catches and defers to the regular fallbacks: +/// `openai.APIError` (every litellm provider error, a failure the host reported no classes for +/// among them), `RouterRateLimitError` and `RouterRateLimitErrorBasic`. Anything else propagates. +fn defers_to_regular_fallbacks(error: &Error) -> bool { + retried_failure(error).is_some_and(|failure| { + failure.classes.is_empty() + || [ + "APIError", + "BadRequestError", + "ServiceUnavailableError", + "RouterRateLimitError", + "RouterRateLimitErrorBasic", + ] + .iter() + .any(|class| failure.is_a(class)) + }) +} + +/// `is_bad_request_without_fallback` with `fallback_on_unsupported_params` off: an unsupported +/// params error, or a provider's plain 400 bad request. +fn is_bad_request_without_fallback(failure: &AttemptError) -> bool { + if !failure.is_a("BadRequestError") { + return false; + } + if failure.is_a("UnsupportedParamsError") { + return true; + } + if [ + "ContextWindowExceededError", + "ContentPolicyViolationError", + "UnsupportedEndpointError", + ] + .iter() + .any(|class| failure.is_a(class)) + { + return false; + } + failure.status == Some(400) && failure.names_provider +} + +type BoxFuture<'a, T> = Pin + Send + 'a>>; + +/// `KeyError`'s classes, most specific first. +/// The failure a retry loop decides by: an attempt's, or the exception a failed pick raised +/// (Python's loop retries a pick's error like any other: no deployment available, a stale +/// routing group's `KeyError`, ...). +fn retried_failure(error: &Error) -> Option> { + match error { + Error::Attempt(failure) => Some(Cow::Borrowed(failure)), + _ => error.pick_failure().map(Cow::Owned), + } +} + +impl Core { + pub(crate) async fn call( + &self, + op: Op, + request: Arc, + ) -> Result { + let snapshot = self.snapshot.load_full(); + let mut chain = FallbackChain::default(); + let target = Target { + group: request.model.clone(), + order: None, + params: None, + excluded: Arc::from([]), + failover_excluded: Arc::from([]), + }; + let at = FallbackHop::default(); + // A fallbacks mock stands in for the requested group's retry loop: no pick, no attempt. + let Some(mock) = request.mock_failure.filter(|mock| mock.before_retries()) else { + return self + .with_fallbacks(&snapshot, &op, &request, target, at, &mut chain) + .await; + }; + let failed = Failed::from(self.mock_failed(mock, &target, at, None)); + self.fall_back(&snapshot, &op, &request, target, at, &mut chain, failed) + .await + } + + /// One level of `async_function_with_fallbacks`: the target's retry loop, then its + /// fallbacks if every attempt failed. + #[allow(clippy::too_many_arguments)] + fn with_fallbacks<'a>( + &'a self, + snapshot: &'a RouterStateSnapshot, + op: &'a Op, + request: &'a Arc, + target: Target, + at: FallbackHop, + chain: &'a mut FallbackChain, + ) -> BoxFuture<'a, Result> { + Box::pin(async move { + let failed = match self + .with_retries(snapshot, op, request, &target, at, &chain.attempted) + .await + { + Ok(response) => return Ok(response), + Err(failed) => failed, + }; + self.fall_back(snapshot, op, request, target, at, chain, failed) + .await + }) + } + + /// The level after its own failure: the no-fallback checks, weighted failover, then the + /// fallback targets. A level that fails raises its own error, whatever its fallbacks raised, + /// except an order chain's failure on any error but a context window one, which raises what + /// the chain's last hop raised ([`RouterEvent::OrderChainFailed`]). + /// A stream that fails before content re-enters here (`fallback_stream`), so the group's + /// retry loop does not run again. + #[allow(clippy::too_many_arguments)] + async fn fall_back( + &self, + snapshot: &RouterStateSnapshot, + op: &Op, + request: &Arc, + target: Target, + at: FallbackHop, + chain: &mut FallbackChain, + failed: Failed, + ) -> Result { + let Failed { + error, + deployment_id: failed_id, + } = failed; + let error = *error; + self.level_failed(at, &target, &error, false); + if self.raises_without_fallback(snapshot, request, &error) { + return Err(error); + } + // `async_function_with_fallbacks_common_utils` lists the group's deployments + // (`_core_get_model_list`) for its order fallbacks: a written index position past the + // model list raises there, before any fallback. + if let Err(raised) = crate::resolve::all_deployments( + &snapshot.registry, + &target.group, + request.listing_team_id.as_deref(), + ) { + self.level_failed(at, &target, &raised, true); + return Err(raised); + } + let (targets, order_path) = self.fallback_targets(snapshot, request, &target, &error); + if !order_path { + self.host.observe(&RouterEvent::FallingBack { + level: at.level, + model: target.group.clone(), + }); + } + let mut target = target; + // `_maybe_run_weighted_failover`: its own `run_async_fallback`, before the regular one. + let failover = if order_path { + None + } else { + match self + .weighted_failover(snapshot, request, &target, &error, failed_id.as_deref()) + .await + { + Ok(failover) => failover, + Err(raised) => { + // `_get_routing_context` raised inside the level's fallback handling: the + // level raises that error, before any regular fallback. + self.level_failed(at, &target, &raised, true); + return Err(raised); + } + } + }; + if let Some(failover) = failover { + let excluded = match &failover { + Next::Failover(_, excluded) => excluded.clone(), + _ => Arc::from([]), + }; + target.failover_excluded = excluded; + match self + .run_fallbacks(snapshot, op, request, &target, at, chain, vec![failover]) + .await + { + Ok(response) => return Ok(response), + Err(raised) if defers_to_regular_fallbacks(raised.as_ref().unwrap_or(&error)) => {} + // Outside the `try` around the regular fallbacks: the level raises it as it is. + Err(raised) => { + self.host.observe(&RouterEvent::OrderChainFailed { + level: at.level, + last_hop: raised.is_some(), + }); + return Err(raised.unwrap_or(error)); + } + } + } + let Some(targets) = targets else { + self.host.observe(&RouterEvent::FallbacksFailed { + level: at.level, + chain: false, + last_hop: false, + }); + return Err(error); + }; + let last_hop_error = match self + .run_fallbacks(snapshot, op, request, &target, at, chain, targets) + .await + { + Ok(response) => return Ok(response), + Err(raised) => raised, + }; + if !order_path { + self.host.observe(&RouterEvent::FallbacksFailed { + level: at.level, + chain: true, + last_hop: last_hop_error.is_some(), + }); + return Err(error); + } + // The order path's `raise`, outside the `try` that re-raises the original error: only a + // context window error stays the one the caller sees. + let context_window = retried_failure(&error) + .is_some_and(|failure| failure.is_a("ContextWindowExceededError")); + let raised = last_hop_error.filter(|_| !context_window); + self.host.observe(&RouterEvent::OrderChainFailed { + level: at.level, + last_hop: raised.is_some(), + }); + Err(raised.unwrap_or(error)) + } + + /// `run_async_fallback` over one chain: the depth is checked once on entry, then every + /// target tried counts one more. When the chain served nothing, the error its last hop + /// raised (`error_from_fallbacks`), or `None` when no hop ran. + #[allow(clippy::too_many_arguments)] + fn run_fallbacks<'a>( + &'a self, + snapshot: &'a RouterStateSnapshot, + op: &'a Op, + request: &'a Arc, + target: &'a Target, + at: FallbackHop, + chain: &'a mut FallbackChain, + targets: Vec, + ) -> BoxFuture<'a, Result>> { + Box::pin(async move { + let max_fallbacks = request + .max_fallbacks + .unwrap_or(snapshot.settings.max_fallbacks); + if at.depth >= max_fallbacks { + return Err(None); + } + let failed = Tried::Target(FallbackTarget::Name(target.group.clone())); + if !chain.attempted.contains(&failed) { + chain.attempted.push(failed); + } + // Before the loop, or the refusal would be one target's failure: no target of the + // chain is tried, and the walk raises a `ValueError` the host holds. + let dicts: Vec> = targets + .iter() + .filter_map(|next| match next { + Next::Fallback(FallbackTarget::Params(dict)) => Some(dict.clone()), + _ => None, + }) + .collect(); + if !dicts.is_empty() && self.host.refuses_fallbacks(at.level, &dicts).await { + let refused = AttemptError::host_raised(&["ValueError"]); + return Err(Some(Error::Attempt(refused))); + } + let mut params = target.params.clone(); + let mut hop = at.depth; + let mut last_hop_error = None; + // `kwargs["model"]` as a dict target's non-string `model` left it: later targets + // without a `model` keep it. + let mut unroutable: Option = None; + for next in targets { + let (tried, group, order) = match &next { + Next::Order(group, order) => ( + Tried::Order(group.clone(), *order), + group.clone(), + Some(*order), + ), + Next::Failover(group, excluded) => { + let mut key = excluded.to_vec(); + key.sort(); + (Tried::Failover(group.clone(), key), group.clone(), None) + } + Next::Fallback(FallbackTarget::Name(name)) if *name == target.group => continue, + // A provider-scoped request never leaves its group; a dict target with no + // `model` names no group, so it is skipped too. + Next::Fallback(fallback) + if request.same_group_fallbacks_only + && fallback.model() != Some(target.group.as_str()) => + { + continue; + } + Next::Fallback(fallback) => { + if let Some(model) = fallback.model() + && model != target.group + { + // A check that raises ends the loop outside its `try`, as the + // chain's own failure. + match self.host.allow_fallback(request, model).await { + Ok(true) => {} + Ok(false) => continue, + Err(raised) => return Err(Some(Error::Attempt(raised))), + } + } + let group = fallback.model().unwrap_or(&target.group).to_owned(); + (Tried::Target(fallback.identity()), group, None) + } + }; + if chain.attempted.contains(&tried) { + continue; + } + chain.attempted.push(tried); + if let Next::Fallback(fallback) = &next + && let Some(merged) = fallback.params() + { + let mut combined = params.as_deref().cloned().unwrap_or_default(); + combined.extend(merged.iter().map(|(k, v)| (k.clone(), v.clone()))); + params = Some(Arc::new(combined)); + } + match &next { + Next::Fallback(fallback) if fallback.model().is_none() => { + if let Some(model) = fallback.unroutable_model() { + unroutable = Some(model.clone()); + } + } + _ => unroutable = None, + } + let group = unroutable.as_ref().map_or(group, python_str); + hop += 1; + self.host.observe(&RouterEvent::RetryLogged { + scope: at.scope, + level: at.level, + model: target.group.clone(), + attempted_retries: None, + cause: RetryCause::Level, + }); + chain.scopes += 1; + self.host.observe(&RouterEvent::Fallback { + from: target.group.clone(), + to: group.clone(), + order, + depth: hop, + level: at.level + 1, + scope: chain.scopes, + parent: at.scope, + max_fallbacks, + unroutable: unroutable.clone(), + }); + let excluded = match &next { + Next::Failover(_, excluded) => excluded.clone(), + _ => target.excluded.clone(), + }; + let next = Target { + group: group.clone(), + order, + params: params.clone(), + excluded, + failover_excluded: target.failover_excluded.clone(), + }; + let parent = at; + let at = FallbackHop { + depth: hop, + level: at.level + 1, + scope: chain.scopes, + }; + // A non-string `model` is a hop that fails before any attempt, on an error its + // own level's fallback lookup raises too, so it walks no chain of its own. + let result = match &unroutable { + Some(_) => Err(Error::UnknownModel { + model: group.clone(), + }), + None => { + self.with_fallbacks(snapshot, op, request, next, at, chain) + .await + } + }; + self.host.observe(&RouterEvent::FallbackDone { + from: target.group.clone(), + to: group, + level: parent.level, + scope: at.scope, + parent: parent.scope, + succeeded: result.is_ok(), + }); + match result { + Ok(mut response) => { + response.attempted_fallbacks = hop; + return Ok(response); + } + Err(raised) => last_hop_error = Some(raised), + } + } + Err(last_hop_error) + }) + } + + /// `_maybe_run_weighted_failover`'s hop for a failed deployment, when it applies: the + /// simple-shuffle group re-picked without every deployment already excluded, while one + /// that is not cooling down remains. + async fn weighted_failover( + &self, + snapshot: &RouterStateSnapshot, + request: &RouterRequest, + target: &Target, + error: &Error, + failed_id: Option<&str>, + ) -> Result, Error> { + if !snapshot.settings.enable_weighted_failover { + return Ok(None); + } + if retried_failure(error).is_some_and(|failure| { + failure.is_a("ContextWindowExceededError") + || failure.is_a("ContentPolicyViolationError") + }) { + return Ok(None); + } + // `_get_routing_context` runs first, whatever failed: a stale routing-group member map + // raises its `KeyError` here. + snapshot.missing_routing_group(request, &target.group)?; + if request.skip_weighted_failover || error.attempt().is_none() { + return Ok(None); + } + let kind = snapshot + .registry + .resolve(&target.group, ResolveScope::default()) + .map_or(ResolveKind::Group, |resolution| resolution.kind); + let (strategy, _) = crate::strategy::routing_context( + &snapshot.registry, + &snapshot.settings.routing_strategy, + request.routing_strategy.as_deref(), + &kind, + &target.group, + ); + if strategy != crate::config::RoutingStrategy::SimpleShuffle { + return Ok(None); + } + let Some(failed_id) = failed_id else { + return Ok(None); + }; + let mut excluded: Vec = target.failover_excluded.to_vec(); + if !excluded.iter().any(|id| id == failed_id) { + excluded.push(failed_id.to_owned()); + } + let listing_team = request.listing_team_id.as_deref(); + let ids: Vec = group_members(snapshot, &target.group, listing_team) + .iter() + .filter(|d| d.listed_for_team(listing_team)) + .map(|d| d.id.clone()) + .collect(); + let cooling = self + .health + .active_cooldowns(&snapshot.registry.ids()) + .await + .unwrap_or_default(); + let remaining = ids + .iter() + .any(|id| !excluded.contains(id) && !cooling.iter().any(|(cooled, _)| cooled == id)); + Ok(remaining.then(|| Next::Failover(target.group.clone(), excluded.into()))) + } + + /// `async_function_with_fallbacks_common_utils` raises before looking for any fallback: + /// fallbacks are off, the error is a guardrail's, a bad request without + /// `fallback_on_unsupported_params`, or no fallback serves the error at all. + fn raises_without_fallback( + &self, + snapshot: &RouterStateSnapshot, + request: &RouterRequest, + error: &Error, + ) -> bool { + if request.disable_fallbacks { + return true; + } + // Python decides by the exception alone, so a pick's error falls back like an attempt's: + // an `UnsupportedParamsError` raises, and the refusals name no provider. Only an error no + // pick raises (no Python exception stands behind it) has no fallback. + retried_failure(error).is_none_or(|failure| { + failure.kind == FailureKind::GuardrailIntervention + || (!snapshot.settings.fallback_on_unsupported_params + && is_bad_request_without_fallback(&failure)) + }) + } + + /// The fallback targets for a failure, or `None` when it must be raised as is + /// (`async_function_with_fallbacks_common_utils`). + fn fallback_targets( + &self, + snapshot: &RouterStateSnapshot, + request: &RouterRequest, + target: &Target, + error: &Error, + ) -> (Option>, bool) { + // A pick's error decides by its exception's classes as an attempt's does: a pre-call + // context window refusal takes the context window fallbacks, any other the generic ones. + let failure = retried_failure(error); + let is_a = |name: &str| failure.as_ref().is_some_and(|f| f.is_a(name)); + let lists = Lists::of(snapshot, request); + // `fallback_lookup_groups`: the group being served, then the group the request was for. + let mut lookup = vec![target.group.as_str()]; + if request.model != target.group { + lookup.push(request.model.as_str()); + } + let chain = |list: &Fallbacks| { + if list.is_client_side(lists.keys) { + Some(list.as_chain()) + } else { + list.lookup_groups(&lookup, lists.keys) + } + }; + let fallbacks_of = |targets: Vec| { + targets.into_iter().map(Next::Fallback).collect::>() + }; + // `context_window_fallback_model_group`: only a chain the group matched takes precedence. + let context_window_targets = is_a("ContextWindowExceededError") + .then(|| lists.context_window?.lookup_groups(&lookup, lists.keys)) + .flatten(); + let skip_order = is_a("ContentPolicyViolationError") || context_window_targets.is_some(); + if !skip_order { + // A routing group's members, with a priority group's priorities as their order. + let listing_team = request.listing_team_id.as_deref(); + let mut orders: Vec = group_members(snapshot, &target.group, listing_team) + .iter() + .filter(|d| d.listed_for_team(listing_team)) + .filter_map(|d| d.order()) + .collect(); + orders.sort_unstable(); + orders.dedup(); + if orders.len() > 1 { + let current = target.order.unwrap_or(orders[0]); + let mut targets: Vec = orders + .into_iter() + .filter(|order| *order > current) + .map(|order| Next::Order(target.group.clone(), order)) + .collect(); + if let Some(list) = lists.fallbacks { + targets.extend(fallbacks_of(chain(list).unwrap_or_default())); + } + if !targets.is_empty() { + return (Some(targets), true); + } + } + } + if let Some(list) = lists.fallbacks + && list.is_client_side(lists.keys) + && context_window_targets.is_none() + { + return (Some(fallbacks_of(list.as_chain())), false); + } + if let Some(targets) = context_window_targets { + return (Some(fallbacks_of(targets)), false); + } + // A context window list with no match falls through to the regular fallbacks; a content + // policy list with no match raises. + if is_a("ContentPolicyViolationError") + && let Some(list) = lists.content_policy + { + return ( + list.lookup_groups(&lookup, lists.keys).map(fallbacks_of), + false, + ); + } + ( + lists + .fallbacks + .and_then(|list| list.lookup_groups(&lookup, lists.keys)) + .map(fallbacks_of), + false, + ) + } + + /// A level's first attempt, or the rate-limit mock in its place + /// (`_handle_mock_testing_rate_limit_error`), which the retry loop handles like a failed + /// attempt. + #[allow(clippy::too_many_arguments)] + async fn first_attempt( + &self, + snapshot: &RouterStateSnapshot, + op: &Op, + request: &Arc, + target: &Target, + inputs: &SelectInputs, + max_retries: u32, + place: (FallbackHop, &[Tried]), + ) -> Result { + if request.mock_failure != Some(MockFailure::RateLimit) { + return self + .attempt( + snapshot, + op, + request, + target, + inputs, + (0, max_retries), + place, + ) + .await; + } + // A lone deployment's `num_retries` rides on the mock error, as on a real one. + let lone = snapshot + .registry + .group(&target.group) + .filter(|group| group.len() == 1) + .and_then(|group| group[0].raw_params.get("num_retries")?.as_u64()) + .and_then(|retries| u32::try_from(retries).ok()); + Err(self + .mock_failed(MockFailure::RateLimit, target, place.0, lone) + .into()) + } + + async fn with_retries( + &self, + snapshot: &RouterStateSnapshot, + op: &Op, + request: &Arc, + target: &Target, + at: FallbackHop, + attempted: &[Tried], + ) -> Result { + let place = (at, attempted); + let mut inputs = SelectInputs { + target_order: target.order, + excluded_ids: target.excluded.clone(), + requested: Some(target.group.clone()), + fallback_hop: at.level > 0, + ..SelectInputs::default() + }; + let counting = SelectInputs { + count_only: true, + ..SelectInputs::default() + }; + let settings = &snapshot.settings; + let first_budget = request.num_retries.unwrap_or(settings.num_retries); + self.host.observe(&RouterEvent::RetriesOpened { + scope: at.scope, + max_retries: first_budget, + model_group_size: snapshot.registry.model_list_len(&target.group), + }); + let first = self + .first_attempt(snapshot, op, request, target, &inputs, first_budget, place) + .await; + let mut failed = match first { + Ok(response) => return Ok(response), + Err(failed) => failed, + }; + self.pick_failed(at, target, &failed.error); + // A guardrail's block and a failed agentic follow-up are re-raised as they are, before + // anything else looks at them (`async_function_with_retries`). + if matches!( + &*failed.error, + Error::Attempt(failure) + if failure.kind == FailureKind::GuardrailIntervention || failure.agentic_followup + ) { + return Err(failed); + } + // The first failure looks the group's retry policy up; a lookup that raises replaces the + // failure with its own error, before any retry (`async_function_with_retries`). + if request.retry_policy_raises.contains(&target.group) { + self.host.observe(&RouterEvent::RetryPolicyRaised { + level: at.level, + model: target.group.clone(), + cause: RetryCause::of(&failed.error), + }); + return Err(Failed { + error: Box::new(Error::Attempt(AttemptError::host_raised(&["Exception"]))), + deployment_id: None, + }); + } + let Some(first_failure) = retried_failure(&failed.error) else { + return Err(failed); + }; + let failure: &AttemptError = &first_failure; + // A request's `num_retries`, a count or not, wins over the deployment's and the router's. + let mut num_retries = if request.uncounted_retries.is_some() { + 0 + } else { + request + .num_retries + .or(failure.num_retries) + .unwrap_or(settings.num_retries) + }; + let group_policies = request + .model_group_retry_policy + .as_ref() + .unwrap_or(&settings.model_group_retry_policy); + let policy_retries = (request.num_retries != Some(0)) + .then(|| { + resolve_retry_policy( + settings.retry_policy.as_ref(), + &target.group, + Some(group_policies), + )? + .retries_for(failure) + }) + .flatten(); + let policy_applies = policy_retries.is_some(); + num_retries = policy_retries.unwrap_or(num_retries); + drop(first_failure); + // `_async_get_healthy_deployments` finds no group for a deployment id, so a call by id + // counts no deployments when deciding whether to retry. + let counted = |c: &Candidates, n: usize| if c.kind.is_single() { 0 } else { n }; + let all = self + .candidates(snapshot, op, request, &target.group, &counting) + .await + .map(|c| counted(&c, c.all.len())) + .unwrap_or(0); + + for retry in 0..=num_retries { + if retry > 0 { + self.log_retry(target, at, retry, &failed.error); + } + let healthy_ids: Vec = match self + .candidates(snapshot, op, request, &target.group, &counting) + .await + { + Ok(c) => { + let ids = c.healthy.iter().map(|d| d.id.clone()); + ids.take(counted(&c, c.healthy.len())).collect() + } + // `_async_get_healthy_deployments` swallows the lookup's error and still reads the + // cooldowns: a Redis round trip the refused pick's failure log runs across. + Err(_) => { + let _ = self.health.active_cooldowns(&snapshot.registry.ids()).await; + Vec::new() + } + }; + let healthy = healthy_ids.len(); + // A pick's error is retried as Python's loop retries any exception: a retry policy + // keeps going, else `should_retry_this_error` decides (it raises when nothing is + // healthy). + let Some(failure) = retried_failure(&failed.error) else { + return Err(failed); + }; + if !policy_applies && !self.should_retry(snapshot, request, &failure, healthy, all) { + return Err(failed); + } + let uncounted = request.uncounted_retries.filter(|_| !policy_applies); + if retry == 0 && uncounted.is_none() { + self.host.observe(&RouterEvent::RetryBudget { + scope: at.scope, + max_retries: num_retries, + }); + } + if uncounted == Some(UncountedRetries::Raises) { + self.host.observe(&RouterEvent::RetryCountRaised { + level: at.level, + model: target.group.clone(), + cause: RetryCause::of(&failed.error), + }); + return Err(Failed { + error: Box::new(Error::Attempt(AttemptError::host_raised(&["Exception"]))), + deployment_id: None, + }); + } + if num_retries == 0 { + return Err(failed); + } + if retry == 0 { + self.log_retry(target, at, 0, &failed.error); + } + let named = failure.failed_deployment_id.as_ref(); + if let (Some(status), Some(id)) = + (failure.status, named.or(failed.deployment_id.as_ref())) + && !should_retry_status(status) + && !inputs.skipped_ids.contains(id) + { + inputs.skipped_ids.push(id.clone()); + } + // An endpoint a deployment does not serve is retried on the others, until every + // one refused (`_every_deployment_refused`). + if failure.is_a("UnsupportedEndpointError") + && healthy_ids.iter().all(|id| inputs.skipped_ids.contains(id)) + { + return Err(failed); + } + let delay = self.time_to_sleep( + snapshot, + failure.retry_after, + num_retries - retry, + num_retries, + healthy, + all, + ); + self.host.sleep(delay).await; + if retry == num_retries { + self.host.observe(&RouterEvent::RetriesExhausted { + model_group: target.group.clone(), + attempted: num_retries, + max: num_retries, + level: at.level, + picked: failed.error.attempt().is_none(), + }); + break; + } + self.host.observe(&RouterEvent::Retry { + model_group: target.group.clone(), + retry: retry + 1, + delay, + }); + match self + .attempt( + snapshot, + op, + request, + target, + &inputs, + (retry + 1, num_retries), + place, + ) + .await + { + Ok(mut response) => { + response.max_retries = Some(num_retries); + return Ok(response); + } + Err(next) if next.error.attempt().is_some_and(|f| f.agentic_followup) => { + return Err(next); + } + Err(next) => { + self.pick_failed(at, target, &next.error); + failed = next; + } + } + } + Err(failed) + } + + /// `error` is a pick's: reported as it happens, whatever the level then does with it. + fn pick_failed(&self, at: FallbackHop, target: &Target, error: &Error) { + if let Some(picked) = PickError::of(error) { + self.host.observe(&RouterEvent::PickFailed { + level: at.level, + model: target.group.clone(), + error: picked, + }); + } + } + + /// The level serving `target` raises a pick's `error` ([`RouterEvent::LevelFailed`]); + /// `handling` when its fallback handling raised it. + fn level_failed(&self, at: FallbackHop, target: &Target, error: &Error, handling: bool) { + if let Some(picked) = PickError::of(error) { + self.host.observe(&RouterEvent::LevelFailed { + level: at.level, + model: target.group.clone(), + error: picked, + handling, + }); + } + } + + /// `_core_log_retry` for a failure in the level's retry loop. + fn log_retry(&self, target: &Target, at: FallbackHop, retry: u32, error: &Error) { + let cause = match RetryCause::of(error) { + RetryCause::Other => RetryCause::Attempt, + cause => cause, + }; + self.host.observe(&RouterEvent::RetryLogged { + scope: at.scope, + level: at.level, + model: target.group.clone(), + attempted_retries: Some(retry), + cause, + }); + } + + /// A `mock_testing_*` failure for the level's group, reported to the host, which raises + /// the matching exception. + fn mock_failed( + &self, + mock: MockFailure, + target: &Target, + at: FallbackHop, + num_retries: Option, + ) -> Error { + self.host.observe(&RouterEvent::MockFailed { + mock, + level: at.level, + model: target.group.clone(), + num_retries, + }); + Error::Attempt(mock.failure(num_retries)) + } + + fn should_retry( + &self, + snapshot: &RouterStateSnapshot, + request: &RouterRequest, + failure: &AttemptError, + healthy: usize, + all: usize, + ) -> bool { + let lists = Lists::of(snapshot, request); + should_retry( + failure, + RetryContext { + healthy, + all, + has_context_window_fallbacks: lists.context_window.is_some(), + has_content_policy_fallbacks: lists.content_policy.is_some(), + has_regular_fallbacks: lists.has_regular(), + }, + ) + } + + /// `_time_to_sleep_before_retry` and `calculate_retry_after`. + fn time_to_sleep( + &self, + snapshot: &RouterStateSnapshot, + retry_after: Option, + remaining_retries: u32, + num_retries: u32, + healthy: usize, + all: usize, + ) -> Duration { + if all != 1 && healthy > 0 { + return Duration::ZERO; + } + let jitter = JITTER * self.rng.next_f64(); + if let Some(retry_after) = retry_after + && !retry_after.is_zero() + && retry_after <= MAX_HONOURED_RETRY_AFTER + { + return retry_after + Duration::from_secs_f64(jitter); + } + let exponent = num_retries.saturating_sub(remaining_retries); + let backoff = (INITIAL_RETRY_DELAY * 2f64.powi(exponent as i32)) + .max(snapshot.settings.retry_after.as_secs_f64()) + .min(MAX_RETRY_DELAY); + Duration::from_secs_f64(backoff + jitter) + } + + /// One attempt: select a deployment, run it through the host, record the outcome. + #[allow(clippy::too_many_arguments)] + async fn attempt( + &self, + snapshot: &RouterStateSnapshot, + op: &Op, + request: &Arc, + target: &Target, + inputs: &SelectInputs, + (retry, max_retries): (u32, u32), + (at, attempted): (FallbackHop, &[Tried]), + ) -> Result { + if request.unrouted { + return self + .unrouted_attempt(snapshot, op, request, target, (retry, max_retries), at) + .await; + } + let pick_started = Instant::now(); + let (deployment, candidates) = match self + .picked(snapshot, op, request, target, inputs, retry, at) + .await + { + Ok(picked) => picked, + Err(failed) if request.passthrough_on_no_deployment => { + self.pick_failed(at, target, &failed.error); + return self + .unrouted_attempt(snapshot, op, request, target, (retry, max_retries), at) + .await; + } + Err(failed) => return Err(failed), + }; + self.host.observe(&RouterEvent::Selected { + model_group: candidates.model_group.clone(), + deployment_id: deployment.id.clone(), + pick: pick_started.elapsed(), + }); + if !self.host.supports(op, deployment.provider()) { + return Err(Error::Config(ConfigError::Unsupported(Unsupported::Op { + op: op.clone(), + provider: deployment.provider().to_owned(), + })) + .into()); + } + // A name no deployment serves is routed to the first default fallback + // (`_common_checks_available_deployment`), still stamped as the name. + let defaulted = candidates.model_group != target.group + && matches!( + crate::resolve::resolve_request( + &snapshot.registry, + &target.group, + request.specific_deployment, + crate::resolve::Caller { + team_id: request.team_id.as_deref(), + proxy_admin: request.proxy_admin, + allowed_models: &request.allowed_models, + wildcard_request: candidates.wildcard_request.as_deref(), + ..crate::resolve::Caller::default() + }, + None, + snapshot.settings.dynamic.default_ref(), + jiff::Timestamp::now(), + ), + Ok(None) + ); + let attempt = Attempt { + op: op.clone(), + request: request.clone(), + deployment: deployment.clone(), + name: target.group.clone(), + model_group: if defaulted { + target.group.clone() + } else { + candidates.model_group.clone() + }, + target_order: target.order, + candidate_ids: candidates.all.iter().map(|d| d.id.clone()).collect(), + retry, + fallback_depth: at.depth, + fallback_level: at.level, + timeout: snapshot.settings.timeout, + params: target.params.clone(), + scope: at.scope, + max_retries, + model_group_size: snapshot.registry.model_list_len(&target.group), + }; + // Every attempt runs through the host; the core records its outcome itself, as far as the + // request's `OutcomeCallbacks` reach, except a deferred stream's: its request counts now + // (`increment_deployment_usage_for_response`), its success or failure when the host reports + // the stream's end. + match self.host.invoke(&attempt).await { + Ok(output) => { + let tracks = output + .tracks_usage + .unwrap_or_else(|| deployment.tracks_usage()); + let recorded = match (output.deferred, tracks) { + (true, true) => { + let tokens = output.total_tokens as f64; + self.health.add_usage(&deployment, None, tokens, 1.0).await + } + (true, false) => Ok(()), + (false, _) => { + self.health + .record_success( + request.outcome_callbacks, + &deployment, + output.total_tokens, + tracks, + ) + .await + } + }; + if let Err(error) = recorded { + warn!(deployment_id = %deployment.id, error = %error, "usage was not recorded"); + } + self.host.observe(&RouterEvent::Succeeded { + deployment_id: deployment.id.clone(), + }); + // `_should_raise_content_policy_error`: a blocked response is raised only when + // a content-policy fallback would serve it, and counts as no failure. + let raised = + match output.blocked { + Some(Blocked::ContentFilter) => self + .has_content_policy_fallback(snapshot, request, target, at, attempted), + Some(Blocked::Refusal) => self + .refusal_fallback_available(snapshot, request, target, at, attempted), + None => false, + }; + if let Some(blocked) = output.blocked.filter(|_| raised) { + self.host.observe(&RouterEvent::Blocked { + level: at.level, + model: target.group.clone(), + blocked, + }); + let mut failure = + AttemptError::new(FailureKind::ContentPolicyViolation, Some(400)); + failure.classes = crate::retry::content_policy_classes(); + return Err(Failed { + error: Box::new(Error::Attempt(failure)), + deployment_id: Some(deployment.id.clone()), + }); + } + Ok(RouterResponse { + output, + deployment_id: deployment.id.clone(), + model_group: candidates.model_group, + attempted_retries: retry, + max_retries: None, + attempted_fallbacks: at.depth, + }) + } + Err(failure) => { + if !failure.recorded { + self.record_failure(snapshot, request, &candidates, &deployment, &failure) + .await; + } + Err(Failed { + error: Box::new(Error::Attempt(failure)), + deployment_id: Some(deployment.id.clone()), + }) + } + } + } + + /// The attempt's pick: the deployment and the candidates it was picked from. + #[allow(clippy::too_many_arguments)] + async fn picked( + &self, + snapshot: &RouterStateSnapshot, + op: &Op, + request: &Arc, + target: &Target, + inputs: &SelectInputs, + retry: u32, + at: FallbackHop, + ) -> Result<(Arc, Candidates), Failed> { + crate::routing_plugins::refuse_sync_pick( + snapshot.settings.routing_plugins, + &snapshot.settings.routing_strategy, + request, + )?; + let routed = self + .pre_routed(PreRoute { + op, + request, + name: &target.group, + fallback_depth: at.depth, + retry, + }) + .await?; + match routed { + PreRouted::Model(name) => { + snapshot.missing_routing_group(request, &name)?; + let candidates = self + .candidates(snapshot, op, request, &name, inputs) + .await?; + if candidates.healthy.is_empty() { + return Err(Error::NoHealthyDeployments { + model: candidates.failed_model(&snapshot.registry), + } + .into()); + } + let deployment = self.pick(snapshot, op, request, &candidates).await?; + Ok((deployment, candidates)) + } + // A custom strategy's pick is served as is, for the requested name. + PreRouted::Deployment(deployment) => { + let candidates = Candidates { + name: target.group.clone(), + model_group: target.group.clone(), + kind: ResolveKind::Group, + all: vec![deployment.clone()], + healthy: vec![deployment.clone()], + wildcard_request: None, + }; + Ok((deployment, candidates)) + } + } + } + + /// An attempt on a name that is no model group: the host serves it under that name, with no + /// deployment picked and no outcome recorded. + async fn unrouted_attempt( + &self, + snapshot: &RouterStateSnapshot, + op: &Op, + request: &Arc, + target: &Target, + (retry, max_retries): (u32, u32), + at: FallbackHop, + ) -> Result { + let deployment = + Arc::new(crate::registry::Deployment::unrouted(&target.group).map_err(Error::from)?); + let attempt = Attempt { + op: op.clone(), + request: request.clone(), + deployment, + name: target.group.clone(), + model_group: target.group.clone(), + target_order: target.order, + candidate_ids: Vec::new(), + retry, + fallback_depth: at.depth, + fallback_level: at.level, + timeout: snapshot.settings.timeout, + params: target.params.clone(), + scope: at.scope, + max_retries, + model_group_size: snapshot.registry.model_list_len(&target.group), + }; + match self.host.invoke(&attempt).await { + Ok(output) => Ok(RouterResponse { + output, + deployment_id: String::new(), + model_group: target.group.clone(), + attempted_retries: retry, + max_retries: None, + attempted_fallbacks: at.depth, + }), + Err(failure) => Err(Error::Attempt(failure).into()), + } + } + + /// `_has_content_policy_fallback`: the content-policy chain resolves for the level's + /// lookup groups and holds a target the chain has not tried; with no content-policy list, + /// any router `"*"` entry. + fn has_content_policy_fallback( + &self, + snapshot: &RouterStateSnapshot, + request: &RouterRequest, + target: &Target, + at: FallbackHop, + attempted: &[Tried], + ) -> bool { + let lists = Lists::of(snapshot, request); + match lists.content_policy { + Some(list) => unattempted_chain(list, &lists, request, target, at, attempted), + None => snapshot.settings.dynamic.fallbacks.has_star(), + } + } + + /// `_refusal_fallback_available`: the content-policy list decides when there is one; + /// else a router `"*"` entry, else a generic chain with a target the chain has not tried. + fn refusal_fallback_available( + &self, + snapshot: &RouterStateSnapshot, + request: &RouterRequest, + target: &Target, + at: FallbackHop, + attempted: &[Tried], + ) -> bool { + if request.disable_fallbacks { + return false; + } + let lists = Lists::of(snapshot, request); + if lists.content_policy.is_some() { + return self.has_content_policy_fallback(snapshot, request, target, at, attempted); + } + if snapshot.settings.dynamic.fallbacks.has_star() { + return true; + } + lists + .fallbacks + .is_some_and(|list| unattempted_chain(list, &lists, request, target, at, attempted)) + } + + /// `async_deployment_callback_on_failure`, when it runs for this failure. + async fn count_failed_request( + &self, + deployment: &crate::registry::Deployment, + failure: &AttemptError, + ) { + if failure.uncounted_request { + return; + } + if let Err(error) = self.health.count_failed_request(deployment, None).await { + warn!(deployment_id = %deployment.id, error = %error, "usage was not recorded"); + } + } + + async fn record_failure( + &self, + snapshot: &RouterStateSnapshot, + request: &RouterRequest, + candidates: &Candidates, + deployment: &crate::registry::Deployment, + failure: &AttemptError, + ) { + self.host.observe(&RouterEvent::AttemptFailed { + deployment_id: deployment.id.clone(), + kind: failure.kind, + status: failure.status, + }); + let recorded = request.outcome_callbacks.record( + OutcomeWrite::Failure, + self.attempt_failure(snapshot, request, candidates, deployment, failure), + ); + if let Err(error) = recorded.await { + warn!(deployment_id = %deployment.id, error = %error, "cooldown was not recorded"); + } + } + + /// What the failure callbacks record for a failed attempt. + async fn attempt_failure( + &self, + snapshot: &RouterStateSnapshot, + request: &RouterRequest, + candidates: &Candidates, + deployment: &crate::registry::Deployment, + failure: &AttemptError, + ) -> Result<(), Error> { + // The default deployment's copy is no configured deployment: Python's cooldown check finds + // no group for its id and records nothing, unless the id names a configured deployment, + // which then takes the failure. + let configured; + let deployment = if candidates.kind == ResolveKind::Default { + match snapshot.registry.deployment(&deployment.id) { + Some(found) => { + configured = found; + &*configured + } + None => { + self.count_failed_request(deployment, failure).await; + return Ok(()); + } + } + } else { + deployment + }; + let shape = GroupShape { + // The deployment's own group decides, also when the call named the deployment by id. + single_deployment: snapshot + .registry + .group(&deployment.model_name) + .map_or(candidates.all.len() == 1, |group| group.len() == 1) + // A routing group call with alternatives cools its member down so selection + // moves on (`routing_group_alternatives`). + && !(candidates.kind == ResolveKind::RoutingGroup && candidates.all.len() > 1) + // As does a team deployment with a sibling under its public name + // (`team_model_alternatives`). + && !team_model_has_alternatives(&snapshot.registry, deployment), + has_fallback_group: snapshot + .settings + .dynamic + .fallback_targets(&[candidates.model_group.as_str()]) + .is_some(), + }; + self.count_failed_request(deployment, failure).await; + match self + .health + .record_failure( + &snapshot.settings, + deployment, + shape, + request.cooldown_time, + failure, + ) + .await + { + Ok(Some(cooldown)) => self.host.observe(&RouterEvent::CooldownSet { + deployment_id: deployment.id.clone(), + cooldown, + }), + Ok(None) => {} + Err(error) => return Err(error), + } + Ok(()) + } +} diff --git a/litellm-rust/crates/router-core/src/call/stream.rs b/litellm-rust/crates/router-core/src/call/stream.rs new file mode 100644 index 00000000000..a0fce4372dd --- /dev/null +++ b/litellm-rust/crates/router-core/src/call/stream.rs @@ -0,0 +1,456 @@ +//! Recovering a stream that failed after its attempt returned it. The host wraps every streamed +//! response; when the stream fails before the caller received content, the wrapper asks the core +//! what may still take over and hands the failure back here. Mirrors the decisions of the Python +//! router's stream wrappers: `_stream_can_fall_back`, `_refusal_fallback_available`, +//! `_anthropic_messages_stream_can_retry`, `_stream_retry_budget`, `_astream_retry_same_group`, +//! and the `async_function_with_fallbacks_common_utils` re-entry of the mid-stream fallback. + +use std::sync::Arc; + +use serde::Deserialize; + +use super::{Failed, FallbackChain, FallbackHop, Lists, Target, Tried, group_members}; +use crate::{ + Router, RouterStateSnapshot, + config::Op, + error::{AttemptError, Error, FailureKind}, + fallback::FallbackTarget, + health::should_retry_status, + host::{RouterEvent, RouterRequest, UncountedRetries}, + select::{Candidates, SelectInputs}, +}; + +/// Where a streamed response came from: the attempt that returned it. +#[derive(Clone, Debug, Default, Deserialize)] +#[serde(default, deny_unknown_fields)] +pub struct StreamOrigin { + /// The name the attempt served: the requested group, or the fallback target. + pub group: String, + /// The `order` level an order fallback targeted. + pub order: Option, + /// `fallback_depth` of the attempt. + pub depth: u32, + pub level: u32, + /// The deployment the attempt ran on. + pub deployment_id: Option, + /// Same-group stream retries this request already spent (`attempted_retries_for_request`). + pub attempted_retries: u32, + /// The budget the first stream retry committed to (`committed_retry_budget_for_request`). + pub committed_budget: Option, +} + +/// What may still take over a stream that fails before content. +#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)] +pub struct StreamPolicy { + /// `_anthropic_messages_stream_can_retry`. + pub can_retry: bool, + /// `_stream_can_fall_back`. + pub can_fall_back: bool, + /// `_refusal_fallback_available`. + pub refusal_fallback: bool, + /// `_stream_retry_budget` for the failure kind asked about: the budget, and whether a retry + /// policy governs the retry. + pub retry_budget: Option<(u32, bool)>, +} + +impl StreamOrigin { + fn target(&self) -> Target { + Target { + group: self.group.clone(), + order: self.order, + params: None, + excluded: Arc::from([]), + failover_excluded: Arc::from([]), + } + } + + /// `fallback_lookup_groups`: the served group, then the group the request was for. + fn lookup<'a>(&'a self, request: &'a RouterRequest) -> Vec<&'a str> { + let mut lookup = vec![self.group.as_str()]; + if request.model != self.group { + lookup.push(request.model.as_str()); + } + lookup + } + + /// `has_unattempted_fallback_target`. The first attempt carries no attempted targets, so any + /// resolved chain counts; a fallback hop has already tried its own group and the original one. + fn has_unattempted(&self, request: &RouterRequest, chain: Option>) -> bool { + let Some(chain) = chain else { + return false; + }; + if self.depth == 0 { + return true; + } + chain + .iter() + .any(|target| *target != self.group && *target != request.model) + } +} + +/// The group names a chain targets (a dict target counts by its `model`). +fn names(chain: Option>) -> Option> { + chain.map(|targets| { + targets + .iter() + .filter_map(|target| target.model().map(str::to_owned)) + .collect() + }) +} + +/// The `content_policy_fallbacks` chain for the lookup groups: the request's list where it +/// passed one, else the router's (`_has_content_policy_fallback`). +fn content_policy_targets(lists: &Lists<'_>, lookup: &[&str]) -> Option>> { + lists + .content_policy + .map(|list| names(list.lookup_groups(lookup, lists.keys))) +} + +/// The router's `default_fallbacks`, which Python folds into its `fallbacks` as `{"*": [...]}`: a +/// request's own list replaces them. +fn folded_defaults<'a>(snapshot: &'a RouterStateSnapshot, request: &RouterRequest) -> &'a [String] { + match request.fallbacks { + Some(_) => &[], + None => &snapshot.settings.dynamic.default_fallbacks, + } +} + +/// The regular fallback targets for the lookup groups, then the folded `default_fallbacks`. +fn regular_targets( + snapshot: &RouterStateSnapshot, + request: &RouterRequest, + lists: &Lists<'_>, + lookup: &[&str], +) -> Option> { + let defaults = folded_defaults(snapshot, request); + names( + lists + .fallbacks + .and_then(|list| list.lookup_groups(lookup, lists.keys)), + ) + .or_else(|| (!defaults.is_empty()).then(|| defaults.to_vec())) +} + +/// `_has_default_fallbacks`: the router's list, whatever the request passed. +fn has_default_fallbacks(snapshot: &RouterStateSnapshot) -> bool { + let dynamic = &snapshot.settings.dynamic; + dynamic.fallbacks.has_star() || !dynamic.default_fallbacks.is_empty() +} + +impl Router { + /// What may still take over a stream from `origin` that fails before content, and the retry + /// budget for a failure of `kind`. + pub fn stream_policy( + &self, + request: &RouterRequest, + origin: &StreamOrigin, + kind: Option, + ) -> StreamPolicy { + let snapshot = self.core.snapshot.load_full(); + StreamPolicy { + can_retry: stream_can_retry(&snapshot, request, origin), + can_fall_back: stream_can_fall_back(&snapshot, request, origin), + refusal_fallback: refusal_fallback_available(&snapshot, request, origin), + retry_budget: kind.map(|kind| stream_retry_budget(&snapshot, request, origin, kind)), + } + } + + /// The fallback chain for a stream from `origin` that failed with `failure` before content: + /// `async_function_with_fallbacks_common_utils` re-entered from the stream wrapper, so the + /// group's own retry loop does not run again. + pub async fn fallback_stream( + &self, + op: Op, + request: RouterRequest, + origin: StreamOrigin, + failure: AttemptError, + ) -> Result { + self.live()?; + let snapshot = self.core.snapshot.load_full(); + let request = Arc::new(request); + let mut chain = FallbackChain::default(); + if request.model != origin.group { + chain + .attempted + .push(Tried::Target(FallbackTarget::Name(request.model.clone()))); + } + self.core + .fall_back( + &snapshot, + &op, + &request, + origin.target(), + FallbackHop { + depth: origin.depth, + level: origin.level, + scope: 0, + }, + &mut chain, + Failed::from(Error::Attempt(failure)), + ) + .await + } + + /// `_astream_retry_same_group`: re-run the attempt in the stream's own group until a retry + /// returns a response or the budget runs out. The failed result carries the newest failure. + pub async fn retry_stream( + &self, + op: Op, + request: RouterRequest, + origin: StreamOrigin, + failure: AttemptError, + ) -> Result { + self.live()?; + let snapshot = self.core.snapshot.load_full(); + let request = Arc::new(request); + let target = origin.target(); + let at = FallbackHop { + depth: origin.depth, + level: origin.level, + scope: 0, + }; + let counting = SelectInputs { + count_only: true, + ..SelectInputs::default() + }; + let counted = |c: &Candidates, n: usize| if c.kind.is_single() { 0 } else { n }; + let (healthy, all) = self + .core + .candidates(&snapshot, &op, &request, &target.group, &counting) + .await + .map(|c| (counted(&c, c.healthy.len()), counted(&c, c.all.len()))) + .unwrap_or((0, 0)); + let (budget, policy_applies) = + stream_retry_budget(&snapshot, &request, &origin, failure.kind); + let inputs = SelectInputs { + target_order: target.order, + requested: Some(target.group.clone()), + fallback_hop: at.level > 0, + ..SelectInputs::default() + }; + let mut last = failure; + for attempt in origin.attempted_retries..budget { + if last.agentic_followup { + break; + } + if !policy_applies + && !self + .core + .should_retry(&snapshot, &request, &last, healthy, all) + { + break; + } + let delay = self.core.time_to_sleep( + &snapshot, + last.retry_after, + budget - attempt, + budget, + healthy, + all, + ); + self.core.host.sleep(delay).await; + self.core.host.observe(&RouterEvent::Retry { + model_group: target.group.clone(), + retry: attempt + 1, + delay, + }); + match self + .core + .attempt( + &snapshot, + &op, + &request, + &target, + &inputs, + (attempt + 1, budget), + (at, &[]), + ) + .await + { + Ok(mut response) => { + response.attempted_retries = attempt + 1; + response.max_retries = Some(budget); + return Ok(response); + } + Err(failed) => match *failed.error { + // `_stream_fallback_error_for_raised`: a status the pre-stream retry rule + // would not retry ends the retries with that failure. + Error::Attempt(failure) => { + let stop = failure.status.is_some_and(|s| !should_retry_status(s)); + last = failure; + if stop { + break; + } + } + other => return Err(other), + }, + } + } + Err(Error::Attempt(last)) + } +} + +/// The request's `num_retries`, then the routed deployment's, then the router's +/// (`_anthropic_messages_plain_retry_budget`). +fn plain_retry_budget( + snapshot: &RouterStateSnapshot, + request: &RouterRequest, + origin: &StreamOrigin, +) -> u32 { + if let Some(retries) = request.num_retries { + return retries; + } + if request.uncounted_retries == Some(UncountedRetries::Negative) { + return 0; + } + let deployment = origin + .deployment_id + .as_deref() + .and_then(|id| snapshot.registry.deployment(id)); + let configured = deployment.and_then(|d| match d.raw_params.get("num_retries") { + Some(serde_json::Value::Number(n)) => n.as_u64().and_then(|n| u32::try_from(n).ok()), + Some(serde_json::Value::String(s)) if s.chars().all(|c| c.is_ascii_digit()) => { + s.parse().ok() + } + _ => None, + }); + configured.unwrap_or(snapshot.settings.num_retries) +} + +/// The retry policy for the request's model group (`_anthropic_messages_resolved_retry_policy`), +/// unless the request opted out with `num_retries=0`. A policy that does not read as one governs +/// nothing here, with no fall through to the router's `retry_policy`. +fn retry_policy<'a>( + snapshot: &'a RouterStateSnapshot, + request: &'a RouterRequest, + origin: &StreamOrigin, +) -> Option<&'a crate::config::RetryPolicy> { + if request.num_retries == Some(0) || request.retry_policy_raises.contains(&origin.group) { + return None; + } + let settings = &snapshot.settings; + crate::retry::resolve_retry_policy( + settings.retry_policy.as_ref(), + &origin.group, + Some( + request + .model_group_retry_policy + .as_ref() + .unwrap_or(&settings.model_group_retry_policy), + ), + ) +} + +fn stream_retry_budget( + snapshot: &RouterStateSnapshot, + request: &RouterRequest, + origin: &StreamOrigin, + kind: FailureKind, +) -> (u32, bool) { + let failure = AttemptError::new(kind, None); + let policy_retries = + retry_policy(snapshot, request, origin).and_then(|p| p.retries_for(&failure)); + if origin.attempted_retries > 0 + && let Some(committed) = origin.committed_budget + { + return (committed, policy_retries.is_some()); + } + match policy_retries { + Some(retries) => (retries, true), + None => (plain_retry_budget(snapshot, request, origin), false), + } +} + +fn stream_can_retry( + snapshot: &RouterStateSnapshot, + request: &RouterRequest, + origin: &StreamOrigin, +) -> bool { + let attempted = origin.attempted_retries; + if attempted > 0 + && let Some(committed) = origin.committed_budget + { + return committed > attempted; + } + let plain = plain_retry_budget(snapshot, request, origin); + // `_retry_policy_ceiling`: every field the policy dumps. + let ceiling = retry_policy(snapshot, request, origin).map_or(plain, |policy| { + let granted = [ + policy.BadRequestErrorRetries, + policy.AuthenticationErrorRetries, + policy.TimeoutErrorRetries, + policy.RateLimitErrorRetries, + policy.ContentPolicyViolationErrorRetries, + policy.InternalServerErrorRetries, + policy.ServiceUnavailableErrorRetries, + policy.NotFoundErrorRetries, + policy.DefaultRetries, + ]; + plain.max(granted.into_iter().flatten().max().unwrap_or(0)) + }); + ceiling > attempted +} + +fn stream_can_fall_back( + snapshot: &RouterStateSnapshot, + request: &RouterRequest, + origin: &StreamOrigin, +) -> bool { + if request.disable_fallbacks { + return false; + } + if snapshot.settings.enable_weighted_failover { + return true; + } + // `_stream_order_levels`, the levels the fallback walk sees for the group. + let listing_team = request.listing_team_id.as_deref(); + let mut orders: Vec = group_members(snapshot, &origin.group, listing_team) + .iter() + .filter(|d| d.listed_for_team(listing_team)) + .filter_map(|d| d.order()) + .collect(); + orders.sort_unstable(); + orders.dedup(); + if orders.len() > 1 { + let current = origin.order.unwrap_or(orders[0]); + if orders.iter().any(|order| *order > current) { + return true; + } + } + let lists = Lists::of(snapshot, request); + let lookup = origin.lookup(request); + if let Some(chain) = content_policy_targets(&lists, &lookup) + && origin.has_unattempted(request, chain) + { + return true; + } + if !lists.has_regular() && folded_defaults(snapshot, request).is_empty() { + return false; + } + if lists + .fallbacks + .is_some_and(|list| list.is_client_side(lists.keys)) + { + return true; + } + origin.has_unattempted(request, regular_targets(snapshot, request, &lists, &lookup)) +} + +fn refusal_fallback_available( + snapshot: &RouterStateSnapshot, + request: &RouterRequest, + origin: &StreamOrigin, +) -> bool { + if request.disable_fallbacks { + return false; + } + let lists = Lists::of(snapshot, request); + let lookup = origin.lookup(request); + if let Some(chain) = content_policy_targets(&lists, &lookup) { + return origin.has_unattempted(request, chain); + } + if has_default_fallbacks(snapshot) { + return true; + } + origin.has_unattempted(request, regular_targets(snapshot, request, &lists, &lookup)) +} diff --git a/litellm-rust/crates/router-core/src/config.rs b/litellm-rust/crates/router-core/src/config.rs new file mode 100644 index 00000000000..fb9b7af215f --- /dev/null +++ b/litellm-rust/crates/router-core/src/config.rs @@ -0,0 +1,355 @@ +//! What the router is told: [`RouterConfig`], fixed for the router's life, and +//! [`RouterDynamicConfig`], a versioned snapshot read once per call. The router never reads +//! process globals. + +use std::{collections::BTreeMap, time::Duration}; + +use serde::Deserialize; + +use crate::{ + error::AttemptError, + registry::{DeploymentSpec, RoutingGroupSpec}, +}; + +/// `DEFAULT_COOLDOWN_TIME_SECONDS` in `litellm/constants.py`. +pub const DEFAULT_COOLDOWN_TIME: Duration = Duration::from_secs(5); +/// `ROUTER_MAX_FALLBACKS` in `litellm/constants.py`. +pub const DEFAULT_MAX_FALLBACKS: u32 = 5; + +/// The router operation a call runs: the Python router's endpoint method name (`acompletion`, +/// `aembedding`, ...). The core runs every op alike; what differs per endpoint is the host's. +#[derive(Clone, Debug, PartialEq, Eq, Hash, Deserialize)] +#[serde(transparent)] +pub struct Op(String); + +impl Op { + pub fn new(name: impl Into) -> Self { + Self(name.into()) + } + + pub fn as_str(&self) -> &str { + &self.0 + } +} + +impl From<&str> for Op { + fn from(name: &str) -> Self { + Self::new(name) + } +} + +impl std::fmt::Display for Op { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.write_str(&self.0) + } +} + +/// `routing_strategy`. Only simple-shuffle is native; the rest go through `Host::select`. +#[derive( + Clone, + Debug, + Default, + PartialEq, + Eq, + strum::EnumString, + strum::Display, + serde_with::DeserializeFromStr, +)] +pub enum RoutingStrategy { + #[default] + #[strum(serialize = "simple-shuffle")] + SimpleShuffle, + #[strum(serialize = "least-busy")] + LeastBusy, + #[strum(serialize = "usage-based-routing")] + UsageBased, + #[strum(serialize = "usage-based-routing-v2")] + UsageBasedV2, + #[strum(serialize = "latency-based-routing")] + LatencyBased, + #[strum(serialize = "cost-based-routing")] + CostBased, + /// `routing_strategy=None` passed explicitly (sent as `"None"`): stored as None, so no + /// strategy arm matches and every pick finds no deployment. + #[strum(serialize = "None")] + Unset, + #[strum(default, transparent)] + Other(String), +} + +impl RoutingStrategy { + pub fn is_native(&self) -> bool { + matches!(self, Self::SimpleShuffle | Self::Unset) + } +} + +/// `PluginKind` in `litellm/router_api.py`. +#[derive(Clone, Copy, Debug, PartialEq, Eq, Deserialize, strum::Display)] +pub enum PluginKind { + #[serde(rename = "auto")] + #[strum(serialize = "auto")] + Auto, + #[serde(rename = "complexity")] + #[strum(serialize = "complexity")] + Complexity, + #[serde(rename = "quality")] + #[strum(serialize = "quality")] + Quality, + #[serde(rename = "adaptive")] + #[strum(serialize = "adaptive")] + Adaptive, + #[serde(rename = "routing")] + #[strum(serialize = "routing")] + Routing, + #[serde(rename = "budget_limiter")] + #[strum(serialize = "budget_limiter")] + BudgetLimiter, + #[serde(rename = "prompt_caching_pins")] + #[strum(serialize = "prompt_caching_pins")] + PromptCachingPins, +} + +/// One registered plugin. Its behaviour lives in the host until it is ported. +#[derive(Clone, Debug, PartialEq, Deserialize)] +pub struct PluginSpec { + pub kind: PluginKind, + pub name: String, + #[serde(default)] + pub config: serde_json::Value, +} + +/// `RetryPolicy` in `litellm/types/router.py`: retries per exception type. +#[derive(Clone, Debug, Default, PartialEq, Deserialize)] +#[allow(non_snake_case)] +pub struct RetryPolicy { + pub BadRequestErrorRetries: Option, + pub AuthenticationErrorRetries: Option, + pub TimeoutErrorRetries: Option, + pub RateLimitErrorRetries: Option, + pub ContentPolicyViolationErrorRetries: Option, + pub InternalServerErrorRetries: Option, + pub ServiceUnavailableErrorRetries: Option, + pub NotFoundErrorRetries: Option, + pub DefaultRetries: Option, +} + +/// `AllowedFailsPolicy` in `litellm/types/router.py`: allowed fails per minute per exception type. +#[derive(Clone, Debug, Default, PartialEq, Deserialize)] +#[allow(non_snake_case)] +pub struct AllowedFailsPolicy { + pub BadRequestErrorAllowedFails: Option, + pub AuthenticationErrorAllowedFails: Option, + pub TimeoutErrorAllowedFails: Option, + pub RateLimitErrorAllowedFails: Option, + pub ContentPolicyViolationErrorAllowedFails: Option, + pub InternalServerErrorAllowedFails: Option, + pub ServiceUnavailableErrorAllowedFails: Option, + pub BadGatewayErrorAllowedFails: Option, + pub NotFoundErrorAllowedFails: Option, +} + +impl AllowedFailsPolicy { + /// `_core_get_allowed_fails_from_policy`: the first set field whose exception class the + /// failure is an instance of, in Python's order. + pub fn allowed_fails_for(&self, failure: &AttemptError) -> Option { + [ + ("AuthenticationError", self.AuthenticationErrorAllowedFails), + ("Timeout", self.TimeoutErrorAllowedFails), + ("RateLimitError", self.RateLimitErrorAllowedFails), + ( + "ContentPolicyViolationError", + self.ContentPolicyViolationErrorAllowedFails, + ), + ("BadRequestError", self.BadRequestErrorAllowedFails), + ("InternalServerError", self.InternalServerErrorAllowedFails), + ( + "ServiceUnavailableError", + self.ServiceUnavailableErrorAllowedFails, + ), + ("BadGatewayError", self.BadGatewayErrorAllowedFails), + ("NotFoundError", self.NotFoundErrorAllowedFails), + ] + .into_iter() + .find_map(|(class, allowed)| allowed.filter(|_| failure.is_a(class))) + } +} + +pub use crate::fallback::{FallbackKeys, Fallbacks}; + +/// Static config, fixed for the router's life. +#[derive(Clone, Debug, Default, Deserialize)] +#[serde(default)] +pub struct RouterConfig { + pub model_list: Vec, + pub routing_strategy: RoutingStrategy, + pub routing_groups: Vec, + /// `_model_to_group` as Python last built it (member model group to routing group), set + /// only when a raw `routing_groups` write replaced the groups since: member calls keep + /// routing by it, as the Python router's stale map does. + pub routing_group_members: Option>, + pub plugins: Vec, + /// Router default `cooldown_time`; `None` is `DEFAULT_COOLDOWN_TIME_SECONDS`. + pub cooldown_time: Option, + pub disable_cooldowns: bool, + pub allowed_fails_policy: Option, + pub retry_policy: Option, + /// `model_group_retry_policy`: a group's own policy, which wins over `retry_policy`. + pub model_group_retry_policy: BTreeMap, + /// `retry_after`: the minimum backoff between retries, in seconds. + pub retry_after: f64, + pub timeout: Option, + pub request_timeout: Option, + pub drop_params: Option, + pub enable_pre_call_checks: bool, + /// `Router(plugins=[...])` is set: the host runs the routing-plugin pipeline. + pub routing_plugins: bool, + /// The host's `CustomLogger` deployment filters run on every async candidate list + /// (`async_callback_filter_deployments`). + pub callback_filters: bool, + pub enable_health_check_routing: bool, + /// How old (seconds) a background health check result may be and still count + /// (`health_check_staleness_threshold`, else the default interval times its multiplier). + pub health_check_staleness_threshold: Option, + /// `background_health_check_model_groups`: when set, only these groups are health-filtered. + pub background_health_check_model_groups: Option>, + /// `enable_weighted_failover`: a failed deployment's group is re-picked without it before + /// any cross-group fallback (simple-shuffle only). + pub enable_weighted_failover: bool, + /// `fallback_on_unsupported_params`; `None` is Python's default, `True`. + pub fallback_on_unsupported_params: Option, + /// Deployment ids that are a provider's default and never cool down + /// (`provider_default_deployment_ids`). + pub provider_default_deployment_ids: Vec, + /// The operations this router must serve: `Host::supports` must accept each one for every + /// configured provider, or the router fails to build. + pub served_ops: Vec, +} + +/// Dynamic config: a versioned snapshot read once per call. +#[derive(Clone, Debug, Deserialize)] +#[serde(default)] +pub struct RouterDynamicConfig { + pub version: u64, + /// `model_group_alias` as Python spells it: each alias maps to a group name, or to + /// `{"model": , "hidden": }`. + #[serde(alias = "model_group_alias", deserialize_with = "alias_targets")] + pub model_alias_map: BTreeMap, + pub fallbacks: Fallbacks, + pub default_fallbacks: Vec, + pub context_window_fallbacks: Option, + pub content_policy_fallbacks: Option, + pub num_retries: Option, + pub max_fallbacks: u32, + /// `allowed_fails`, as the router copied it when it was built. + pub allowed_fails: Option, + /// `allowed_fails` is only the process default: it was not passed to the router and still + /// equals `litellm.allowed_fails`. Anything else switches cooldowns from the failure-ratio + /// rule to the allowed-fails counter (`_is_allowed_fails_set_on_router`). + pub allowed_fails_is_default: bool, + /// Other spellings a fallback key matches. + pub fallback_keys: FallbackKeys, + /// `default_deployment`: what a name no deployment, alias, routing group, team route or + /// wildcard serves is routed to, as a copy whose `litellm_params.model` is the name + /// (`_try_early_resolve_deployments_for_model_not_in_names`). + pub default_deployment: Option, + /// The `default_deployment` is a value Python cannot copy for a name: resolving to it + /// raises the copy's error ([`crate::Error::DefaultDeploymentRaises`]). + pub default_deployment_raises: bool, +} + +/// `model_group_alias` as `resolve_model_group_alias` reads it: an alias maps to a non-empty +/// group name, or to `{"model": , ...}`; any other value rewrites nothing. +fn alias_targets<'de, D>(deserializer: D) -> Result, D::Error> +where + D: serde::Deserializer<'de>, +{ + let aliases = BTreeMap::::deserialize(deserializer)?; + Ok(aliases + .into_iter() + .filter_map(|(alias, target)| { + let group = match &target { + serde_json::Value::Object(entry) => entry.get("model")?, + other => other, + }; + group + .as_str() + .filter(|group| !group.is_empty()) + .map(|group| (alias, group.to_owned())) + }) + .collect()) +} + +impl Default for RouterDynamicConfig { + fn default() -> Self { + Self { + version: 0, + model_alias_map: BTreeMap::new(), + fallbacks: Fallbacks::default(), + default_fallbacks: Vec::new(), + context_window_fallbacks: None, + content_policy_fallbacks: None, + num_retries: None, + max_fallbacks: DEFAULT_MAX_FALLBACKS, + allowed_fails: None, + allowed_fails_is_default: false, + fallback_keys: FallbackKeys::default(), + default_deployment: None, + default_deployment_raises: false, + } + } +} + +impl RouterDynamicConfig { + /// The `default_deployment` as resolution serves it. + pub fn default_ref(&self) -> Option> { + if self.default_deployment_raises { + return Some(crate::resolve::DefaultDeployment::Raises); + } + self.default_deployment + .as_ref() + .map(crate::resolve::DefaultDeployment::Spec) + } + + /// The regular fallback targets for the lookup groups, then `default_fallbacks`, which + /// Python folds into `fallbacks` as `{"*": [...]}`: what `_has_fallback_group` reads. + pub fn fallback_targets(&self, groups: &[&str]) -> Option> { + if groups + .iter() + .any(|group| self.fallbacks.has_keyed_chain(group, &self.fallback_keys)) + { + return Some(Vec::new()); + } + (!self.default_fallbacks.is_empty()).then(|| self.default_fallbacks.clone()) + } +} + +#[cfg(test)] +mod tests { + use rstest::rstest; + use serde_json::json; + + use super::*; + + #[rstest] + #[case::group_name(json!({"model_group_alias": {"gpt": "gpt-4"}}))] + #[case::hidden_entry(json!({"model_group_alias": {"gpt": {"model": "gpt-4", "hidden": true}}}))] + #[case::native_name(json!({"model_alias_map": {"gpt": "gpt-4"}}))] + fn model_group_alias_reads_both_python_spellings(#[case] wire: serde_json::Value) { + let dynamic: RouterDynamicConfig = serde_json::from_value(wire).unwrap(); + assert_eq!( + dynamic.model_alias_map, + BTreeMap::from([("gpt".to_owned(), "gpt-4".to_owned())]) + ); + } + + #[rstest] + #[case::native("simple-shuffle", RoutingStrategy::SimpleShuffle)] + #[case::unknown("my-strategy", RoutingStrategy::Other("my-strategy".to_owned()))] + fn routing_strategy_parses_its_wire_name( + #[case] wire: &str, + #[case] expected: RoutingStrategy, + ) { + assert_eq!(wire.parse::().unwrap(), expected); + assert_eq!(expected.to_string(), wire); + } +} diff --git a/litellm-rust/crates/router-core/src/eligibility.rs b/litellm-rust/crates/router-core/src/eligibility.rs new file mode 100644 index 00000000000..5cedf218b9a --- /dev/null +++ b/litellm-rust/crates/router-core/src/eligibility.rs @@ -0,0 +1,240 @@ +//! Which deployments may serve a request's endpoint: a port of +//! `litellm/router_utils/endpoint_eligibility.py`. A deployment that sets +//! `enforce_supported_endpoints` and lists none of the chat-shaped endpoints serves none of them; +//! the host names the endpoint for the calls `endpoint_selection` wraps. + +use std::{collections::BTreeSet, sync::Arc}; + +use serde_json::{Map, Value}; + +use crate::{error::Error, registry::Deployment, resolve::Caller}; + +/// `CHAT_ENDPOINTS`: `_endpoint_path` of each value of `CHAT_ENDPOINT_BY_CALL`. +const CHAT_ENDPOINTS: [&str; 8] = [ + "/chat/completions", + "/completions", + "/messages", + "/responses", + "/generatecontent", + "/streamgeneratecontent", + "/interactions", + "/responses/compact", +]; + +/// `_endpoint_path`. +fn endpoint_path(endpoint: &str) -> String { + let rooted = format!("/{}", endpoint.trim().trim_matches('/').to_lowercase()); + let path = rooted.strip_prefix("/openai").unwrap_or(&rooted); + let path = path.strip_prefix("/v1beta").unwrap_or(path); + let path = path.strip_prefix("/v1").unwrap_or(path); + match path.rsplit_once(':') { + Some((_, action)) => format!("/{action}"), + None => path.to_owned(), + } +} + +/// The flag as pydantic's lax `bool | None` reads it: `None` for a value it refuses, which fails +/// the whole declaration. +fn enforce_flag(value: Option<&Value>) -> Option { + match value { + None | Some(Value::Null) => Some(false), + Some(Value::Bool(flag)) => Some(*flag), + Some(Value::Number(number)) => match number.as_f64() { + Some(0.0) => Some(false), + Some(1.0) => Some(true), + _ => None, + }, + Some(Value::String(text)) => match text.to_lowercase().as_str() { + "0" | "off" | "f" | "false" | "n" | "no" => Some(false), + "1" | "on" | "t" | "true" | "y" | "yes" => Some(true), + _ => None, + }, + Some(_) => None, + } +} + +/// `_enforced_endpoints`: the declared endpoints of a deployment that enforces them. One endpoint +/// may be written as a string; a declaration that does not validate enforces nothing. +fn enforced_endpoints(model_info: &Map) -> Option> { + if !enforce_flag(model_info.get("enforce_supported_endpoints"))? { + return None; + } + match model_info.get("supported_endpoints")? { + Value::String(endpoint) => Some(vec![endpoint.clone()]), + Value::Array(endpoints) => endpoints + .iter() + .map(|endpoint| endpoint.as_str().map(str::to_owned)) + .collect(), + _ => None, + } +} + +/// `non_chat_endpoints`: what a deployment serves when it serves no chat-shaped endpoint, else +/// `None`. +fn non_chat_endpoints(deployment: &Deployment) -> Option> { + let declared = enforced_endpoints(&deployment.model_info.raw)?; + let chat = declared + .iter() + .any(|endpoint| CHAT_ENDPOINTS.contains(&endpoint_path(endpoint).as_str())); + (!chat).then(|| declared.into_iter().collect()) +} + +/// `eligible_deployments`: the deployments that may serve the caller's endpoint. When none may, +/// the refusal names what they do serve. +pub(crate) fn eligible_deployments( + model: &str, + caller: Caller<'_>, + deployments: Vec>, +) -> Result>, Error> { + let Some(endpoint) = caller.endpoint.filter(|_| !deployments.is_empty()) else { + return Ok(deployments); + }; + let declared: Vec>> = deployments + .iter() + .map(|deployment| non_chat_endpoints(deployment)) + .collect(); + let eligible: Vec> = deployments + .into_iter() + .zip(&declared) + .filter(|(_, endpoints)| endpoints.is_none()) + .map(|(deployment, _)| deployment) + .collect(); + if !eligible.is_empty() { + return Ok(eligible); + } + let supported: BTreeSet = declared.into_iter().flatten().flatten().collect(); + let supported = if supported.is_empty() { + "none".to_owned() + } else { + Vec::from_iter(supported).join(", ") + }; + let requested = if caller.requested.is_empty() { + model + } else { + caller.requested + }; + Err(Error::EndpointNotSupported { + model: model.to_owned(), + message: format!( + "{endpoint} is not supported by model={requested}. Supported endpoints: {supported}" + ), + }) +} + +#[cfg(test)] +mod tests { + use rstest::rstest; + use serde_json::json; + + use super::*; + use crate::registry::DeploymentSpec; + + fn deployment(id: &str, info: Value) -> Arc { + let mut model_info = info; + model_info["id"] = json!(id); + let spec: DeploymentSpec = serde_json::from_value(json!({ + "model_name": "g", "litellm_params": {"model": "openai/a"}, "model_info": model_info, + })) + .unwrap(); + Arc::new(Deployment::from_spec(spec).unwrap()) + } + + fn chat<'a>() -> Caller<'a> { + Caller { + endpoint: Some("/v1/chat/completions"), + requested: "g", + ..Caller::default() + } + } + + #[rstest] + #[case::v1_prefix("/v1/chat/completions", "/chat/completions")] + #[case::openai_prefix("/openai/v1/responses/", "/responses")] + #[case::bare_gemini_action("generateContent", "/generatecontent")] + #[case::beta_prefix("/v1beta/interactions", "/interactions")] + #[case::gemini_model_path( + "/v1beta/models/gemini:streamGenerateContent", + "/streamgeneratecontent" + )] + fn endpoint_paths_normalize(#[case] endpoint: &str, #[case] expected: &str) { + assert_eq!(endpoint_path(endpoint), expected); + } + + #[rstest] + #[case::embeddings_only( + json!({"enforce_supported_endpoints": true, "supported_endpoints": ["/v1/embeddings"]}), + false + )] + #[case::one_endpoint_as_a_string( + json!({"enforce_supported_endpoints": true, "supported_endpoints": "/v1/embeddings"}), + false + )] + #[case::string_flag( + json!({"enforce_supported_endpoints": "yes", "supported_endpoints": ["/v1/embeddings"]}), + false + )] + #[case::an_empty_list_serves_nothing( + json!({"enforce_supported_endpoints": true, "supported_endpoints": []}), + false + )] + #[case::any_chat_shaped_endpoint_serves_every_one( + json!({"enforce_supported_endpoints": true, "supported_endpoints": ["/v1/responses"]}), + true + )] + #[case::not_enforced( + json!({"enforce_supported_endpoints": false, "supported_endpoints": ["/v1/embeddings"]}), + true + )] + #[case::unset_flag(json!({"supported_endpoints": ["/v1/embeddings"]}), true)] + #[case::enforced_without_a_list(json!({"enforce_supported_endpoints": true}), true)] + #[case::an_invalid_flag_enforces_nothing( + json!({"enforce_supported_endpoints": "maybe", "supported_endpoints": ["/v1/embeddings"]}), + true + )] + #[case::an_invalid_list_enforces_nothing( + json!({"enforce_supported_endpoints": true, "supported_endpoints": [1]}), + true + )] + fn a_declaration_keeps_or_drops_a_deployment(#[case] info: Value, #[case] kept: bool) { + let open = deployment("open", json!({})); + let declared = deployment("declared", info); + let eligible = eligible_deployments("g", chat(), vec![declared, open]).unwrap(); + assert_eq!(eligible.len(), if kept { 2 } else { 1 }); + } + + #[test] + fn no_endpoint_keeps_everything() { + let declared = deployment( + "d", + json!({"enforce_supported_endpoints": true, "supported_endpoints": []}), + ); + let eligible = eligible_deployments("g", Caller::default(), vec![declared]).unwrap(); + assert_eq!(eligible.len(), 1); + } + + #[rstest] + #[case::lists_what_they_serve( + json!(["/v1/rerank", "/v1/embeddings"]), + "/v1/chat/completions is not supported by model=g. \ + Supported endpoints: /v1/embeddings, /v1/rerank" + )] + #[case::nothing_declared( + json!([]), + "/v1/chat/completions is not supported by model=g. Supported endpoints: none" + )] + fn every_deployment_refusing_names_the_endpoint( + #[case] endpoints: Value, + #[case] expected: &str, + ) { + let declared = deployment( + "d", + json!({"enforce_supported_endpoints": true, "supported_endpoints": endpoints}), + ); + let refused = eligible_deployments("resolved", chat(), vec![declared]); + assert!(matches!( + refused, + Err(Error::EndpointNotSupported { model, message }) + if model == "resolved" && message == expected + )); + } +} diff --git a/litellm-rust/crates/router-core/src/error.rs b/litellm-rust/crates/router-core/src/error.rs new file mode 100644 index 00000000000..b4112501c36 --- /dev/null +++ b/litellm-rust/crates/router-core/src/error.rs @@ -0,0 +1,473 @@ +//! The router's errors: [`ConfigError`] when a router is built or changed, [`Error`] when a call +//! or a selection fails, and [`AttemptError`], the typed facts of one failed attempt. + +use std::{any::Any, fmt, sync::Arc, time::Duration}; + +use crate::config::{Op, RoutingStrategy}; + +/// What a host cannot serve, found when the router is built rather than partway through a call. +#[derive(Clone, Debug, PartialEq, thiserror::Error)] +pub enum Unsupported { + #[error("routing strategy `{0}` needs a host that can select")] + Strategy(RoutingStrategy), + #[error("routing plugin `{0}` needs a host that can select")] + Plugin(String), + #[error("the host cannot run operation `{op}` on provider `{provider}`")] + Op { op: Op, provider: String }, +} + +#[derive(Debug, thiserror::Error)] +pub enum ConfigError { + #[error("the router config asks for something this host cannot serve")] + Unsupported(#[from] Unsupported), + #[error("deployment `{model_name}` has invalid litellm_params")] + InvalidDeployment { + model_name: String, + #[source] + source: serde_json::Error, + }, + #[error("deployment id `{id}` is configured more than once")] + DuplicateId { id: String }, + /// A removal the host's indexes could not finish ([`crate::host_index::RemovalError`]). + #[error("{0}")] + Removal(String), +} + +#[derive(Debug, thiserror::Error)] +pub enum Error { + #[error("the router config is invalid")] + Config(#[from] ConfigError), + #[error("no deployment serves model `{model}`")] + UnknownModel { model: String }, + #[error("no deployment of `{model}` is available")] + NoHealthyDeployments { model: String }, + #[error("the attempt failed")] + Attempt(#[from] AttemptError), + #[error("the routing cache failed")] + Cache(#[from] litellm_cache::Error), + #[error("the router was retired")] + Retired, + /// A blocked deployment named by its id (`model` is its litellm model). + #[error("deployment model `{model}` is paused")] + Paused { model: String }, + /// A pass-through pick named a deployment id without `use_in_pass_through`. + #[error("deployment `{id}` does not support pass-through")] + NotPassThrough { id: String }, + /// A pass-through pick of a group with healthy deployments, none of them pass-through. + #[error("model `{model}` has no pass-through deployment")] + NoPassThroughDeployments { model: String }, + /// Every deployment `model` resolved to is a pre-routing strategy marker (`auto_router/...`). + #[error("every deployment of `{model}` is a strategy router marker")] + OnlyMarkers { model: String }, + /// Every deployment `model` resolved to refuses the request's endpoint + /// (`eligible_deployments`); `message` names what they serve. + #[error("{message}")] + EndpointNotSupported { model: String, message: String }, + /// Every deployment `model` resolved to is split into PTU shares the caller holds none of. + #[error("deployment `{model}` is reserved for the teams holding a PTU share of it")] + PtuReserved { model: String }, + /// A proxy admin without a team named a public model name several teams use. + #[error("model name `{model}` matches deployments from multiple teams")] + MultipleTeams { model: String }, + /// Every deployment `model` resolved to is reserved for other teams by an active access + /// window, the first of which ends at `until` (`HH:MM`) in `timezone`. + #[error("deployment `{model}` is reserved for another team until {until} {timezone}")] + Reserved { + model: String, + until: String, + timezone: String, + }, + /// The pre-call checks found no deployment whose context window fits the prompt; `detail` + /// names each one dropped for it. + #[error("no deployment of `{model}` has a context window large enough for this call")] + ContextWindowExceeded { model: String, detail: String }, + /// The pre-call checks found no deployment of `model` that takes the requested `params`. + #[error("no deployment of `{model}` supports {params:?}")] + UnsupportedParams { model: String, params: Vec }, + /// The routing plugins left none of the healthy deployments' models. + #[error("no deployments left after routing-plugin filtering: {models:?}")] + RoutingPluginsLeftNothing { models: Vec }, + /// The request's `_router_weights` failed `validate_router_weights` when simple-shuffle picked. + #[error("the request's router weights are invalid")] + InvalidRouterWeights, + /// A deployment's `litellm_params` `metric` (weight, rpm or tpm) is not a number. + #[error("deployment {metric} must be numeric")] + NonNumericWeight { metric: &'static str }, + /// A router with routing plugins was asked for a pick on the sync path, which never runs them. + #[error("routing plugins only run on the async selection path")] + RoutingPluginsNeedAsync, + /// The stale routing-group member map (`_model_to_group`, kept as it was when a raw + /// `routing_groups` write replaced the groups) names a group that is no longer configured: + /// Python's `_get_routing_context` raises `KeyError(group)` for every pick of the name. + #[error("routing group `{group}` is no longer configured")] + MissingRoutingGroup { group: String }, + /// The router's `default_deployment` is not a value Python can copy for `model` + /// (`_try_early_resolve_deployments_for_model_not_in_names` raises the copy's error). + #[error("the default deployment cannot be served for `{model}`")] + DefaultDeploymentRaises { model: String }, + /// A written model-name or team index names a position past the end of the model list: + /// Python's lookup of `model` raises `IndexError`. + #[error("a written index of `{model}` names a position past the model list")] + IndexOutOfRange { model: String }, + /// A model-info query raised where the Python router raises: `exception` names the Python + /// exception class (`KeyError`, `TypeError`, `AttributeError`, `ValueError`) and `detail` its + /// argument. + #[error("the model-info query raised {exception}({detail:?})")] + ModelInfo { + exception: &'static str, + detail: String, + }, +} + +impl Error { + /// The failed attempt behind this error, when a deployment was called. + pub fn attempt(&self) -> Option<&AttemptError> { + match self { + Self::Attempt(attempt) => Some(attempt), + Self::Config(_) + | Self::UnknownModel { .. } + | Self::NoHealthyDeployments { .. } + | Self::Cache(_) + | Self::Retired + | Self::Paused { .. } + | Self::NotPassThrough { .. } + | Self::NoPassThroughDeployments { .. } + | Self::OnlyMarkers { .. } + | Self::EndpointNotSupported { .. } + | Self::PtuReserved { .. } + | Self::MultipleTeams { .. } + | Self::Reserved { .. } + | Self::ContextWindowExceeded { .. } + | Self::UnsupportedParams { .. } + | Self::RoutingPluginsLeftNothing { .. } + | Self::InvalidRouterWeights + | Self::NonNumericWeight { .. } + | Self::RoutingPluginsNeedAsync + | Self::MissingRoutingGroup { .. } + | Self::DefaultDeploymentRaises { .. } + | Self::IndexOutOfRange { .. } + | Self::ModelInfo { .. } => None, + } + } + + /// The Python exception a pick that failed with this error raises, as the retry loop + /// decides by it (its class names, most specific first, and its status): `None` for an + /// attempt's failure and for errors no pick raises. + pub fn pick_failure(&self) -> Option { + const BAD_REQUEST: &[&str] = &["BadRequestError"]; + let (classes, status): (&[&str], Option) = match self { + Self::UnknownModel { .. } + | Self::OnlyMarkers { .. } + | Self::PtuReserved { .. } + | Self::NotPassThrough { .. } + | Self::NoPassThroughDeployments { .. } + | Self::MultipleTeams { .. } + | Self::Reserved { .. } => (BAD_REQUEST, Some(400)), + Self::EndpointNotSupported { .. } => ( + &["EndpointNotSupportedByModelError", "BadRequestError"], + Some(400), + ), + Self::ContextWindowExceeded { .. } => ( + &["ContextWindowExceededError", "BadRequestError"], + Some(400), + ), + Self::UnsupportedParams { .. } => { + (&["UnsupportedParamsError", "BadRequestError"], Some(400)) + } + Self::Paused { .. } => (&["ServiceUnavailableError"], Some(503)), + Self::NoHealthyDeployments { .. } => (&["RouterRateLimitError", "ValueError"], None), + Self::RoutingPluginsLeftNothing { .. } | Self::RoutingPluginsNeedAsync => { + (&["ValueError"], None) + } + Self::InvalidRouterWeights => (&["ValidationError", "ValueError"], None), + Self::NonNumericWeight { .. } => (&["TypeError"], None), + Self::MissingRoutingGroup { .. } => (&["KeyError", "LookupError"], None), + Self::IndexOutOfRange { .. } => (&["IndexError", "LookupError"], None), + Self::DefaultDeploymentRaises { .. } => (&["Exception"], None), + Self::Attempt(_) + | Self::Config(_) + | Self::Cache(_) + | Self::Retired + | Self::ModelInfo { .. } => return None, + }; + let mut failure = AttemptError::host_raised(classes); + failure.status = status; + Some(failure) + } + + /// A copy of a pick's error, for the events that name it; `None` for any other error. + pub fn pick_copy(&self) -> Option { + Some(match self { + Self::UnknownModel { model } => Self::UnknownModel { + model: model.clone(), + }, + Self::NoHealthyDeployments { model } => Self::NoHealthyDeployments { + model: model.clone(), + }, + Self::Paused { model } => Self::Paused { + model: model.clone(), + }, + Self::NotPassThrough { id } => Self::NotPassThrough { id: id.clone() }, + Self::NoPassThroughDeployments { model } => Self::NoPassThroughDeployments { + model: model.clone(), + }, + Self::OnlyMarkers { model } => Self::OnlyMarkers { + model: model.clone(), + }, + Self::EndpointNotSupported { model, message } => Self::EndpointNotSupported { + model: model.clone(), + message: message.clone(), + }, + Self::PtuReserved { model } => Self::PtuReserved { + model: model.clone(), + }, + Self::MultipleTeams { model } => Self::MultipleTeams { + model: model.clone(), + }, + Self::Reserved { + model, + until, + timezone, + } => Self::Reserved { + model: model.clone(), + until: until.clone(), + timezone: timezone.clone(), + }, + Self::ContextWindowExceeded { model, detail } => Self::ContextWindowExceeded { + model: model.clone(), + detail: detail.clone(), + }, + Self::UnsupportedParams { model, params } => Self::UnsupportedParams { + model: model.clone(), + params: params.clone(), + }, + Self::RoutingPluginsLeftNothing { models } => Self::RoutingPluginsLeftNothing { + models: models.clone(), + }, + Self::InvalidRouterWeights => Self::InvalidRouterWeights, + Self::NonNumericWeight { metric } => Self::NonNumericWeight { metric }, + Self::RoutingPluginsNeedAsync => Self::RoutingPluginsNeedAsync, + Self::MissingRoutingGroup { group } => Self::MissingRoutingGroup { + group: group.clone(), + }, + Self::DefaultDeploymentRaises { model } => Self::DefaultDeploymentRaises { + model: model.clone(), + }, + Self::IndexOutOfRange { model } => Self::IndexOutOfRange { + model: model.clone(), + }, + Self::Attempt(_) + | Self::Config(_) + | Self::Cache(_) + | Self::Retired + | Self::ModelInfo { .. } => return None, + }) + } +} + +/// The class of a failed attempt, as litellm's exception mapping names it. +#[derive( + Clone, + Copy, + Debug, + PartialEq, + Eq, + Hash, + strum::Display, + strum::EnumString, + strum::IntoStaticStr, + strum::VariantArray, +)] +pub enum FailureKind { + #[strum(serialize = "RateLimitError")] + RateLimit, + #[strum(serialize = "ContextWindowExceededError")] + ContextWindowExceeded, + #[strum(serialize = "ContentPolicyViolationError")] + ContentPolicyViolation, + #[strum(serialize = "AuthenticationError")] + Authentication, + #[strum(serialize = "PermissionDeniedError")] + PermissionDenied, + #[strum(serialize = "NotFoundError")] + NotFound, + #[strum(serialize = "BadRequestError")] + BadRequest, + #[strum(serialize = "Timeout")] + Timeout, + #[strum(serialize = "APIConnectionError")] + ApiConnection, + #[strum(serialize = "InternalServerError")] + InternalServer, + #[strum(serialize = "ServiceUnavailableError")] + ServiceUnavailable, + #[strum(serialize = "BadGatewayError")] + BadGateway, + /// A guardrail blocked the request: never retried, never a fallback. + #[strum(serialize = "GuardrailIntervention")] + GuardrailIntervention, + /// litellm rejected the request before calling the deployment (`MalformedRequestFieldError`, + /// `UnsupportedEndpointError`): never cools the deployment down. + #[strum(serialize = "Declined")] + Declined, + /// The caller's own stored credential was rejected, not the deployment's + /// (`CallerCredentialAuthenticationError`): never cools the deployment down. + #[strum(serialize = "CallerCredentialAuthenticationError")] + CallerCredentialAuthentication, + /// A per-user credential exchange was rate limited (`CallerCredentialRateLimitError`): never + /// cools the deployment down. + #[strum(serialize = "CallerCredentialRateLimitError")] + CallerCredentialRateLimit, + #[strum(serialize = "APIError")] + Other, +} + +/// An opaque value the host hands through the router unchanged, such as the original Python +/// exception or response object. +#[derive(Clone)] +pub struct Payload(pub Arc); + +impl fmt::Debug for Payload { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + formatter.write_str("Payload(..)") + } +} + +/// One failed attempt, as typed facts the retry, fallback and cooldown rules decide by. +#[derive(Clone, Debug, thiserror::Error)] +#[error("the deployment call failed with {kind}")] +pub struct AttemptError { + pub kind: FailureKind, + pub status: Option, + /// The provider's `retry-after` header (`_get_retry_after_from_exception_header`). + pub retry_after: Option, + /// A deployment-level `num_retries` the failure carries (Python's `exception.num_retries`). + pub num_retries: Option, + /// The deployment the exception itself names (`_deployment_ids_to_skip_on_retry` reads its + /// `retry_skip_deployment_id`, then `failed_deployment_id`): an exception raised again keeps + /// the id its first failure stamped, whichever deployment raised it now. + pub failed_deployment_id: Option, + /// The provider's error text, kept only for the cooldown record Python writes + /// (`exception_received`, masked past 50 characters). + pub detail: Option>, + /// The exception's class names, most specific first (its `__mro__`), for the rules Python + /// decides by `isinstance`. Empty when the host reports none; the kind then stands in. + pub classes: Vec>, + /// The host already recorded what this failure does to the deployment's health (a routing + /// pre-call check's failure, which Python's attempt handles itself: a rate limit cools the + /// deployment down, anything else records nothing), or no failure callback runs for it (it + /// was raised outside the litellm call, or a shared logging object already logged a failure), + /// so the core records nothing. + pub recorded: bool, + /// The exception names its provider (`llm_provider` is set), which marks a provider's own + /// 400 for `fallback_on_unsupported_params`. + pub names_provider: bool, + /// An agentic follow-up call failed after the first call was served + /// (`is_agentic_followup_failure`): never retried, since a retry reruns every paid round. The + /// failure keeps its own class, which decides its fallbacks and cooldown as for any failure. + pub agentic_followup: bool, + /// One of `deployment_callback_on_failure`'s early exits holds: the failure says nothing about + /// the deployment, so it is neither counted nor cooled down (a failed request still counts + /// against the minute's requests, which `async_deployment_callback_on_failure` does apart). + pub outside_deployment: bool, + /// One of `_should_run_cooldown_logic`'s per-caller exits holds (`_is_per_user_provider_request` + /// or a per-user Copilot deployment): the failure is counted but never cools the deployment. + pub caller_scoped: bool, + /// The exception has no `status_code` (Python's status `""`): `_is_cooldown_required` holds + /// no cooldown for it, unlike a status the host could not read, which cools down as a 500. + pub status_absent: bool, + /// The failure counts against no per-minute requests: `async_deployment_callback_on_failure` + /// does not run for a sync call, and returns at once for a batch retrieve. + pub uncounted_request: bool, + pub payload: Option, +} + +impl AttemptError { + pub fn new(kind: FailureKind, status: Option) -> Self { + Self { + kind, + status, + retry_after: None, + num_retries: None, + failed_deployment_id: None, + detail: None, + classes: Vec::new(), + recorded: false, + names_provider: false, + agentic_followup: false, + outside_deployment: false, + caller_scoped: false, + status_absent: false, + uncounted_request: false, + payload: None, + } + } + + /// The exception's class names, most specific first: the host's, else the kind's. + pub fn class_names(&self) -> Vec<&str> { + if self.classes.is_empty() { + return crate::retry::kind_classes(self.kind).to_vec(); + } + self.classes.iter().map(AsRef::as_ref).collect() + } + + /// `isinstance(error, )`. + pub fn is_a(&self, name: &str) -> bool { + self.class_names().contains(&name) + } + + #[must_use] + pub fn with_retry_after(self, retry_after: Duration) -> Self { + Self { + retry_after: Some(retry_after), + ..self + } + } + + /// A plain Python exception the host raised outside any attempt (a `KeyError` from a stale + /// routing-group lookup, an error a retry policy lookup raised): no status, no provider, + /// nothing recorded against a deployment. `classes` are its class names, most specific first. + pub fn host_raised(classes: &[&str]) -> Self { + Self { + classes: classes.iter().map(|class| (*class).into()).collect(), + recorded: true, + ..Self::new(FailureKind::Other, None) + } + } + + /// The status Python's cooldown rules read: the provider's, else 500. + pub fn cooldown_status(&self) -> u16 { + self.status.unwrap_or(500) + } +} + +#[cfg(test)] +mod tests { + use rstest::rstest; + + use super::*; + + #[rstest] + #[case::rate_limit("RateLimitError", FailureKind::RateLimit)] + #[case::context_window("ContextWindowExceededError", FailureKind::ContextWindowExceeded)] + #[case::caller_credential_auth( + "CallerCredentialAuthenticationError", + FailureKind::CallerCredentialAuthentication + )] + #[case::caller_credential_rate_limit( + "CallerCredentialRateLimitError", + FailureKind::CallerCredentialRateLimit + )] + #[case::other("APIError", FailureKind::Other)] + fn failure_kind_round_trips_its_exception_name(#[case] wire: &str, #[case] kind: FailureKind) { + assert_eq!(wire.parse::().unwrap(), kind); + assert_eq!(<&'static str>::from(kind), wire); + } + + #[rstest] + #[case::empty("")] + #[case::lowercase("ratelimiterror")] + fn failure_kind_rejects_unknown_names(#[case] wire: &str) { + assert!(wire.parse::().is_err()); + } +} diff --git a/litellm-rust/crates/router-core/src/fallback.rs b/litellm-rust/crates/router-core/src/fallback.rs new file mode 100644 index 00000000000..a1755bd6ad5 --- /dev/null +++ b/litellm-rust/crates/router-core/src/fallback.rs @@ -0,0 +1,330 @@ +//! Fallback lists as Python spells them, and how a failed group's chain is found in one. +//! Mirrors `get_fallback_model_group`, `get_fallback_model_group_for_lookup_groups`, +//! `check_non_standard_fallback_format` and `fallback_attempt_key` in +//! `router_utils/fallback_event_handlers.py`. +//! +//! A list holds keyed entries (`{"gpt-4": ["gpt-3.5", {"model": "x", "temperature": 0}]}`), +//! bare names (`"claude"`, a chain for every group) and, in a client-side list, dict targets +//! (`{"model": "x", "messages": [...]}`), which override request params for their attempt. + +use serde::Deserialize; +use serde_json::{Map, Value}; + +use std::collections::BTreeMap; + +/// One item of a fallback list, kept as Python holds it. +#[derive(Clone, Debug, PartialEq)] +enum Entry { + Name(String), + Dict(Map), + /// Anything else: Python skips it when it looks a chain up. + Other, +} + +/// A fallback list, in order. `"*"` keys the default chain. +#[derive(Clone, Debug, Default, PartialEq, Deserialize)] +#[serde(from = "Vec")] +pub struct Fallbacks(Vec); + +impl From> for Fallbacks { + fn from(entries: Vec) -> Self { + Self( + entries + .into_iter() + .map(|entry| match entry { + Value::String(name) => Entry::Name(name), + Value::Object(dict) => Entry::Dict(dict), + _ => Entry::Other, + }) + .collect(), + ) + } +} + +/// One target of a chain: a group name, or a dict target that also overrides request params. +#[derive(Clone, Debug, PartialEq)] +pub enum FallbackTarget { + Name(String), + Params(Map), +} + +impl FallbackTarget { + fn of(value: &Value) -> Option { + match value { + Value::String(name) => Some(Self::Name(name.clone())), + Value::Object(dict) => Some(Self::Params(dict.clone())), + _ => None, + } + } + + /// `_get_fallback_target_model_group`: the group the target names, if any. + pub fn model(&self) -> Option<&str> { + match self { + Self::Name(name) => Some(name), + Self::Params(dict) => dict.get("model").and_then(Value::as_str), + } + } + + /// The `model` of a dict target that sets a non-string one: the hop routes that value, which + /// no group answers (`_get_fallback_target_model_group` finds no group to check, but + /// `kwargs.update(target)` still rebinds the request's model). + pub fn unroutable_model(&self) -> Option<&Value> { + match self { + Self::Params(dict) => dict.get("model").filter(|model| !model.is_string()), + Self::Name(_) => None, + } + } + + /// `fallback_attempt_key`: a bare name and `{"model": name}` are the same attempt; a dict + /// carrying anything else is its own. + pub fn identity(&self) -> FallbackTarget { + match self { + Self::Params(dict) if dict.len() == 1 => match dict.get("model") { + Some(Value::String(name)) => Self::Name(name.clone()), + _ => self.clone(), + }, + _ => self.clone(), + } + } + + /// The params a dict target merges into the request (`kwargs.update(target)`). + pub fn params(&self) -> Option<&Map> { + match self { + Self::Params(dict) => Some(dict), + Self::Name(_) => None, + } + } +} + +/// What fallback keys match besides the exact group name: Python's `provider_list`, for +/// `"openai/gpt-4o"` to use a `"gpt-4o"` entry, the provider-prefixed spelling of each bare +/// group name (`_provider_prefixed_model_group`), for `"gpt-4o"` to use an `"openai/gpt-4o"` +/// entry, and the `LiteLLMParamsTypedDict` keys that mark a client-side dict target. +#[derive(Clone, Debug, Default, PartialEq, Deserialize)] +#[serde(default)] +pub struct FallbackKeys { + pub providers: Vec, + pub provider_prefixed: BTreeMap, + pub param_keys: Vec, +} + +impl FallbackKeys { + /// `_check_stripped_model_group`. + fn is_stripped(&self, group: &str, key: &str) -> bool { + self.providers.iter().any(|provider| { + let prefix = format!("{provider}/"); + group.starts_with(&prefix) && group.replace(&prefix, "") == key + }) + } +} + +/// A chain found for a group, and whether it is the `"*"` chain itself. +struct Found { + targets: Vec, + generic: bool, +} + +/// The targets of a keyed entry's value: a list of names and dicts. A bare string is walked +/// character by character, as Python iterates it. +fn chain_of(value: &Value) -> Vec { + match value { + Value::Array(items) => items.iter().filter_map(FallbackTarget::of).collect(), + Value::String(text) => text + .chars() + .map(|c| FallbackTarget::Name(c.to_string())) + .collect(), + _ => Vec::new(), + } +} + +impl Fallbacks { + pub fn is_empty(&self) -> bool { + self.0.is_empty() + } + + /// `check_non_standard_fallback_format`: a client-side list of names, or of dicts that set + /// a litellm param, is one chain for every group. + pub fn is_client_side(&self, keys: &FallbackKeys) -> bool { + if self.0.is_empty() { + return false; + } + if self.0.iter().all(|entry| matches!(entry, Entry::Name(_))) { + return true; + } + let dicts: Vec<_> = self + .0 + .iter() + .filter_map(|entry| match entry { + Entry::Dict(dict) => Some(dict), + _ => None, + }) + .collect(); + dicts.len() == self.0.len() + && dicts.iter().any(|dict| { + keys.param_keys + .iter() + .any(|key| dict.get(key).is_some_and(|value| !value.is_array())) + }) + } + + /// The whole list as one chain (a client-side list). + pub fn as_chain(&self) -> Vec { + self.0 + .iter() + .filter_map(|entry| match entry { + Entry::Name(name) => Some(FallbackTarget::Name(name.clone())), + Entry::Dict(dict) => Some(FallbackTarget::Params(dict.clone())), + Entry::Other => None, + }) + .collect() + } + + /// `_get_first_default_fallback`: the first target of the first `"*"` entry. + pub fn default_target(&self) -> Option<&str> { + self.0.iter().find_map(|entry| match entry { + Entry::Dict(dict) => match dict.get("*") { + Some(Value::Array(targets)) => targets.first().and_then(Value::as_str), + _ => None, + }, + _ => None, + }) + } + + /// `_has_default_fallbacks`: some entry keys `"*"`. + pub fn has_star(&self) -> bool { + self.0 + .iter() + .any(|entry| matches!(entry, Entry::Dict(dict) if dict.contains_key("*"))) + } + + /// `get_fallback_model_group`. + fn lookup(&self, group: &str, keys: &FallbackKeys) -> Option { + let prefixed = keys.provider_prefixed.get(group).map(String::as_str); + let mut generic = None; + let mut stripped = None; + let mut found = None; + for entry in &self.0 { + match entry { + Entry::Dict(dict) => { + let Some((key, value)) = dict.iter().next() else { + continue; + }; + if key == group { + found = Some(chain_of(value)); + break; + } else if Some(key.as_str()) == prefixed || keys.is_stripped(group, key) { + stripped = Some(chain_of(value)); + } else if key == "*" { + generic = Some(value); + } + } + Entry::Name(name) => found = Some(vec![FallbackTarget::Name(name.clone())]), + Entry::Other => {} + } + } + found + .or(stripped) + .map(|targets| Found { + targets, + generic: false, + }) + .or_else(|| { + generic.map(|value| Found { + targets: chain_of(value), + generic: true, + }) + }) + } + + /// `get_fallback_model_group_for_lookup_groups`: the first lookup group with its own chain + /// wins; `"*"` applies only after every group missed. + pub fn lookup_groups( + &self, + groups: &[&str], + keys: &FallbackKeys, + ) -> Option> { + let results: Vec = groups + .iter() + .filter_map(|group| self.lookup(group, keys)) + .collect(); + let index = results + .iter() + .position(|found| !found.generic) + .or_else(|| (!results.is_empty()).then_some(0))?; + results.into_iter().nth(index).map(|found| found.targets) + } + + /// `_has_fallback_group`'s view: only keyed entries whose targets are names or dicts count. + pub fn has_keyed_chain(&self, group: &str, keys: &FallbackKeys) -> bool { + let keyed = Fallbacks( + self.0 + .iter() + .filter(|entry| matches!(entry, Entry::Dict(_))) + .cloned() + .collect(), + ); + keyed + .lookup(group, keys) + .is_some_and(|found| !found.targets.is_empty()) + } +} + +#[cfg(test)] +mod tests { + use rstest::rstest; + use serde_json::json; + + use super::*; + + #[rstest] + #[case::an_int(json!({"model": 5}), true)] + #[case::none(json!({"model": null}), true)] + #[case::a_string(json!({"model": "b"}), false)] + #[case::no_model(json!({"temperature": 0}), false)] + fn only_a_set_non_string_model_is_unroutable(#[case] target: Value, #[case] unroutable: bool) { + let target = FallbackTarget::of(&target).unwrap(); + assert_eq!(target.unroutable_model().is_some(), unroutable); + } + + fn names(targets: Option>) -> Option> { + targets.map(|targets| { + targets + .iter() + .map(|t| t.model().unwrap_or_default().to_owned()) + .collect() + }) + } + + #[rstest] + #[case::own_entry(json!([{"a": ["b"]}, {"*": ["d"]}]), "a", Some(vec!["b"]))] + #[case::star_entry(json!([{"a": ["b"]}, {"*": ["d"]}]), "x", Some(vec!["d"]))] + #[case::string_is_every_groups_chain(json!(["b", {"x": ["c"]}]), "a", Some(vec!["b"]))] + #[case::keyed_entry_wins(json!(["c", {"a": ["b"]}]), "a", Some(vec!["b"]))] + #[case::string_value_is_iterated(json!([{"a": "bc"}]), "a", Some(vec!["b", "c"]))] + #[case::none(json!([{"x": ["c"]}]), "a", None)] + fn chains_follow_python( + #[case] list: Value, + #[case] group: &str, + #[case] expected: Option>, + ) { + let fallbacks: Fallbacks = serde_json::from_value(list).unwrap(); + assert_eq!( + names(fallbacks.lookup_groups(&[group], &FallbackKeys::default())), + expected.map(|e| e.into_iter().map(str::to_owned).collect()) + ); + } + + #[rstest] + #[case::names(json!(["a", "b"]), true)] + #[case::param_dicts(json!([{"model": "a", "temperature": 1}]), true)] + #[case::keyed(json!([{"a": ["b"]}]), false)] + #[case::mixed(json!(["a", {"a": ["b"]}]), false)] + fn client_side_lists_follow_python(#[case] list: Value, #[case] expected: bool) { + let fallbacks: Fallbacks = serde_json::from_value(list).unwrap(); + let keys = FallbackKeys { + param_keys: vec!["model".to_owned(), "temperature".to_owned()], + ..FallbackKeys::default() + }; + assert_eq!(fallbacks.is_client_side(&keys), expected); + } +} diff --git a/litellm-rust/crates/router-core/src/group_info.rs b/litellm-rust/crates/router-core/src/group_info.rs new file mode 100644 index 00000000000..9a49858ca18 --- /dev/null +++ b/litellm-rust/crates/router-core/src/group_info.rs @@ -0,0 +1,646 @@ +//! Model group info and the configured-value reads: what the router reports a model group, or a +//! concrete deployment behind a listed name, can do. +//! +//! The rules are the Python router core's (`_set_model_group_info`, `get_model_listing_info`, +//! `get_configured_token_limits`, `get_configured_mode`, `get_configured_display_name`, +//! `get_configured_service_tiers`). The host picks the deployments (the registry lookups) and +//! answers what needs its objects as [`Host::model_info`] lookups asked while the query runs: a +//! deployment's validated params, its merged model info and provider, its capabilities, and the +//! group's first entry as its model validates it. The core decides which deployments count and +//! how their values combine. + +use serde_json::{Map, Number, Value}; + +use crate::{ + error::Error, + host::{Host, ModelInfoLookup}, + model_info::{ + discovered, is_opaque, lookup, malformed, python_float, python_int, raised, truthy, + without_digit_underscores, + }, +}; + +/// One `_set_model_group_info` query. +#[derive(Clone, Debug, Default, serde::Deserialize)] +#[serde(default)] +pub struct GroupInfoQuery { + /// The model group the deployments serve (an alias already resolved). + pub model_group: String, + /// The name the group is reported as. + pub user_facing_model_group_name: String, + /// The group is a team's public model name the team's deployments serve. + pub serves_team_deployments: bool, + /// A wildcard pattern routes `model_group` (asked only when an entry needs it). + pub pattern_routes: bool, + /// The first string `model_info.description` among the deployments. + pub description: Value, + /// The deployments the group routes to, as stored. + pub deployments: Vec, +} + +/// The group info: `None` when no deployment matched, else the fields the merge assigned after +/// the host built the group from its first deployment ([`ModelInfoLookup::GroupSeed`]). +#[derive(Clone, Debug, Default, PartialEq)] +pub struct GroupInfo { + pub assignments: Map, +} + +/// The type name Python reports for a value, in its error messages. +pub(crate) fn type_name(value: &Value) -> &'static str { + match value { + Value::Null => "NoneType", + Value::Bool(_) => "bool", + Value::Number(number) if number.is_f64() => "float", + Value::Number(_) => "int", + Value::String(_) => "str", + Value::Array(_) => "list", + Value::Object(_) if is_opaque(value) => "object", + Value::Object(_) => "dict", + } +} + +/// A number as Python's arithmetic reads it (a bool is 0 or 1), or `None`. +fn arithmetic(value: &Value) -> Option { + match value { + Value::Bool(flag) => Some(Number::from(u8::from(*flag))), + Value::Number(number) => Some(number.clone()), + _ => None, + } +} + +/// `total += value`, with `total` starting at `0`; adding a non-number raises `TypeError`. +fn add(total: Option, value: &Value) -> Result { + let total = total.unwrap_or_else(|| Value::Number(Number::from(0))); + let (Some(left), Some(right)) = (arithmetic(&total), arithmetic(value)) else { + return Err(raised( + "TypeError", + format!( + "unsupported operand type(s) for +=: '{}' and '{}'", + type_name(&total), + type_name(value) + ), + )); + }; + if let (Some(left), Some(right)) = (left.as_i64(), right.as_i64()) + && let Some(sum) = left.checked_add(right) + { + return Ok(Value::Number(Number::from(sum))); + } + let sum = left.as_f64().unwrap_or(0.0) + right.as_f64().unwrap_or(0.0); + Ok(Number::from_f64(sum).map_or(Value::Null, Value::Number)) +} + +/// `left > right`, for numbers or strings; any other pair raises `TypeError`. +fn greater(left: &Value, right: &Value) -> Result { + if let (Some(left), Some(right)) = (arithmetic(left), arithmetic(right)) { + return Ok(left.as_f64().unwrap_or(0.0) > right.as_f64().unwrap_or(0.0)); + } + if let (Value::String(left), Value::String(right)) = (left, right) { + return Ok(left > right); + } + Err(raised( + "TypeError", + format!( + "'>' not supported between instances of '{}' and '{}'", + type_name(left), + type_name(right) + ), + )) +} + +/// `_cost_value_as_float`: `float(value)`, or `None` when it is unset or not a number. +fn cost_as_float(value: Option<&Value>) -> Option { + match value? { + Value::Bool(flag) => Some(f64::from(u8::from(*flag))), + Value::Number(number) => number.as_f64(), + Value::String(text) => { + without_digit_underscores(text.trim()).and_then(|digits| python_float(&digits)) + } + _ => None, + } +} + +fn float(value: Option) -> Value { + value + .and_then(Number::from_f64) + .map_or(Value::Null, Value::Number) +} + +/// `value.get(key)` on a dict; a value that is no dict raises `AttributeError`, as `.get` does. +fn get<'a>(value: &'a Value, key: &str) -> Result, Error> { + match value { + Value::Object(map) if !is_opaque(value) => Ok(map.get(key)), + other => Err(raised( + "AttributeError", + format!("'{}' object has no attribute 'get'", type_name(other)), + )), + } +} + +/// `model.get(key)`, then `litellm_params.get(key)`, then `model_info.get(key)`: the first that is +/// not `None`. +fn first_set( + model: &Value, + params: &Value, + info: &Value, + key: &str, +) -> Result, Error> { + for layer in [model, params, info] { + if let Some(value) = get(layer, key)?.filter(|value| !value.is_null()) { + return Ok(Some(value.clone())); + } + } + Ok(None) +} + +fn object(answer: Value, kind: &str) -> Result, Error> { + match answer { + Value::Object(map) => Ok(map), + _ => Err(malformed(kind)), + } +} + +/// The group being merged: the host's validated first entry, and what the merge assigned since. +struct Merged { + current: Map, + assignments: Map, +} + +impl Merged { + fn get(&self, field: &str) -> &Value { + self.current.get(field).unwrap_or(&Value::Null) + } + + fn set(&mut self, field: &str, value: Value) { + self.current.insert(field.to_owned(), value.clone()); + self.assignments.insert(field.to_owned(), value); + } + + /// `if info.get(field) is not None and info[field] > group.field (or group.field is None)`. + fn widen(&mut self, field: &str, info: &Map) -> Result<(), Error> { + let Some(value) = info.get(field).filter(|value| !value.is_null()) else { + return Ok(()); + }; + let current = self.get(field); + if current.is_null() || greater(value, current)? { + self.set(field, value.clone()); + } + Ok(()) + } + + /// The cost `field` widened to `float(info[field])` when that exceeds `group.field or 0.0`. + fn widen_cost(&mut self, field: &str, info: &Map) { + let Some(cost) = cost_as_float(info.get(field)) else { + return; + }; + let current = self.get(field); + let floor = if truthy(current) { + current.as_f64().unwrap_or(0.0) + } else { + 0.0 + }; + if current.is_null() || cost > floor { + self.set(field, float(Some(cost))); + } + } + + /// `if info.get(field) is True: group.field = True`. + fn any(&mut self, field: &str, info: &Map) { + if info.get(field) == Some(&Value::Bool(true)) { + self.set(field, Value::Bool(true)); + } + } +} + +/// `tuple(dict.fromkeys((*current, *added))) or None`. +fn union(current: &Value, added: &[Value]) -> Value { + let mut kept: Vec = Vec::new(); + let existing = current.as_array().map(Vec::as_slice).unwrap_or_default(); + for value in existing.iter().chain(added) { + if !kept.contains(value) { + kept.push(value.clone()); + } + } + if kept.is_empty() { + Value::Null + } else { + Value::Array(kept) + } +} + +/// `_set_model_group_info`: the group's info merged across the deployments that serve it. +pub async fn model_group_info( + host: &dyn Host, + query: &GroupInfoQuery, +) -> Result, Error> { + let mut merged: Option = None; + let mut totals: [(&str, Option); 4] = + [("tpm", None), ("rpm", None), ("itpm", None), ("otpm", None)]; + let mut clientside_auth_params = Value::Null; + let mut efforts_initialized = false; + let mut efforts_unknown = false; + let empty = Value::Object(Map::new()); + + for (index, model) in query.deployments.iter().enumerate() { + let named = get(model, "model_name")?; + let is_match = query.serves_team_deployments + || named.is_some_and(|name| name.as_str() == Some(query.model_group.as_str())) + || (named.is_some() && query.pattern_routes); + if !is_match { + continue; + } + let params = object( + lookup(host, ModelInfoLookup::GroupParams { index }).await?, + "group_params", + )?; + clientside_auth_params = params + .get("configurable_clientside_auth_params") + .cloned() + .unwrap_or(Value::Null); + let model_params = get(model, "litellm_params")?.unwrap_or(&empty); + let model_info_dict = get(model, "model_info")?.unwrap_or(&empty); + let mut deployment = [None, None, None, None]; + for (slot, (metric, _)) in deployment.iter_mut().zip(totals.iter()) { + *slot = first_set(model, model_params, model_info_dict, metric)?; + } + + let facts = object( + lookup( + host, + ModelInfoLookup::GroupDeployment { + index, + model: params.get("model").cloned().unwrap_or(Value::Null), + }, + ) + .await?, + "group_deployment", + )?; + let mapped = facts.get("mapped") == Some(&Value::Bool(true)); + let litellm_model = facts.get("litellm_model").cloned().unwrap_or(Value::Null); + let llm_provider = facts.get("llm_provider").cloned().unwrap_or(Value::Null); + let model_info = match facts.get("model_info") { + Some(Value::Object(info)) if !is_opaque(&Value::Object(info.clone())) => info.clone(), + Some(Value::Null) | None => { + // No model info: what the deployment's own (database) model_info says. + let db_info = get(model, "model_info")?.unwrap_or(&empty); + let mut info = Map::new(); + info.insert("key".to_owned(), Value::String(query.model_group.clone())); + for key in ["max_tokens", "max_input_tokens", "max_output_tokens"] { + info.insert(key.to_owned(), Value::Null); + } + for key in ["input_cost_per_token", "output_cost_per_token"] { + info.insert(key.to_owned(), float(cost_as_float(get(db_info, key)?))); + } + info.insert("litellm_provider".to_owned(), llm_provider.clone()); + info.insert( + "mode".to_owned(), + match get(db_info, "mode")? { + Some(mode) => mode.clone(), + None => Value::String("chat".to_owned()), + }, + ); + info.insert( + "supported_openai_params".to_owned(), + match facts.get("supported_openai_params") { + Some(Value::Null) | None => Value::Array(Vec::new()), + Some(params) => params.clone(), + }, + ); + info.insert("supports_system_messages".to_owned(), Value::Null); + info + } + Some(_) => return Err(malformed("group_deployment")), + }; + + match merged.as_mut() { + None => { + let mut seed = Map::new(); + seed.insert( + "model_group".to_owned(), + Value::String(query.user_facing_model_group_name.clone()), + ); + seed.insert( + "providers".to_owned(), + Value::Array(vec![llm_provider.clone()]), + ); + seed.insert("description".to_owned(), query.description.clone()); + seed.extend(model_info.clone()); + seed.insert( + "supports_continuation".to_owned(), + Value::Bool( + model_info.get("supports_continuation") == Some(&Value::Bool(true)), + ), + ); + seed.insert("supports_fast_mode".to_owned(), Value::Bool(true)); + seed.insert("supported_reasoning_efforts".to_owned(), Value::Null); + seed.insert("supported_endpoints".to_owned(), Value::Null); + let current = object( + lookup(host, ModelInfoLookup::GroupSeed { info: seed }).await?, + "group_seed", + )?; + merged = Some(Merged { + current, + assignments: Map::new(), + }); + } + Some(group) => { + let providers = group + .get("providers") + .as_array() + .cloned() + .unwrap_or_default(); + if !providers.contains(&llm_provider) { + let mut providers = providers; + providers.push(llm_provider.clone()); + group.set("providers", Value::Array(providers)); + } + group.widen("max_input_tokens", &model_info)?; + group.widen("max_output_tokens", &model_info)?; + group.widen_cost("input_cost_per_token", &model_info); + group.widen_cost("output_cost_per_token", &model_info); + for field in [ + "supports_parallel_function_calling", + "supports_vision", + "supports_function_calling", + "supports_web_search", + "supports_url_context", + "supports_continuation", + "supports_reasoning", + ] { + group.any(field, &model_info); + } + if let Some(params) = model_info + .get("supported_openai_params") + .filter(|params| !params.is_null()) + { + group.set("supported_openai_params", params.clone()); + } + for (slot, metric) in deployment.iter_mut().zip(["tpm", "rpm"]) { + if slot.is_none() + && let Some(value) = model_info.get(metric).filter(|v| !v.is_null()) + { + *slot = Some(value.clone()); + } + } + } + } + let Some(group) = merged.as_mut() else { + unreachable!("the group was just built"); + }; + + let capabilities = object( + lookup( + host, + ModelInfoLookup::GroupCapabilities { + index, + model_info: Value::Object(model_info), + litellm_model, + llm_provider, + mapped, + }, + ) + .await?, + "group_capabilities", + )?; + let endpoints = capabilities + .get("supported_endpoints") + .and_then(Value::as_array) + .cloned() + .unwrap_or_default(); + let endpoints = union(group.get("supported_endpoints"), &endpoints); + group.set("supported_endpoints", endpoints); + let fast_mode = truthy(group.get("supports_fast_mode")) + && capabilities.get("supports_fast_mode") == Some(&Value::Bool(true)); + group.set("supports_fast_mode", Value::Bool(fast_mode)); + match capabilities + .get("reasoning_efforts") + .cloned() + .unwrap_or(Value::Null) + { + Value::Null => { + efforts_unknown = true; + group.set("supported_reasoning_efforts", Value::Null); + } + resolved if !efforts_initialized => { + efforts_initialized = true; + if !efforts_unknown { + group.set("supported_reasoning_efforts", resolved); + } + } + resolved if !efforts_unknown => { + let intersected = lookup( + host, + ModelInfoLookup::ReasoningIntersect { + current: group.get("supported_reasoning_efforts").clone(), + resolved, + }, + ) + .await?; + group.set("supported_reasoning_efforts", intersected); + } + _ => {} + } + + for ((_, total), value) in totals.iter_mut().zip(deployment) { + if let Some(value) = value { + *total = Some(add(total.take(), &value)?); + } + } + } + + let Some(mut group) = merged else { + return Ok(None); + }; + for (metric, total) in totals { + if let Some(total) = total { + group.set(metric, total); + } + } + if !clientside_auth_params.is_null() { + group.set( + "configurable_clientside_auth_params", + clientside_auth_params, + ); + } + Ok(Some(GroupInfo { + assignments: group.assignments, + })) +} + +/// `coerce_token_limit`: a token limit as an int, `None` when it is unset or not a usable number +/// (a bool never is). +pub fn coerce_token_limit(value: Option<&Value>) -> Option { + match value? { + Value::Number(number) => number.as_i64().or_else(|| { + number + .as_f64() + .filter(|n| n.is_finite()) + .map(|n| n.trunc() as i64) + }), + Value::String(text) => without_digit_underscores(text.trim()) + .and_then(|digits| python_int(&digits)) + .and_then(|number| number.as_i64()), + _ => None, + } +} + +/// The deployment's discovered limits under its configured `model_info` values that are set. +async fn effective_model_info( + host: &dyn Host, + id: Value, + configured: &Map, +) -> Result, Error> { + let mut info = discovered(host, id).await?; + info.extend( + configured + .iter() + .filter(|(_, value)| !value.is_null()) + .map(|(key, value)| (key.clone(), value.clone())), + ); + Ok(info) +} + +/// `get_configured_token_limits`: the (max input, max output) tokens a concrete deployment +/// configures, or its refresh discovered, from its `model_info` (`exclude_none` dump) and id. +pub async fn configured_token_limits( + host: &dyn Host, + id: Value, + configured: &Map, +) -> Result<(Option, Option), Error> { + let info = effective_model_info(host, id, configured).await?; + Ok(( + coerce_token_limit(info.get("max_input_tokens")), + coerce_token_limit(info.get("max_output_tokens")), + )) +} + +/// `get_configured_mode` / `get_configured_display_name`: a configured string that is not blank. +pub fn configured_text(value: Option<&Value>) -> Option { + value + .and_then(Value::as_str) + .filter(|text| !text.trim().is_empty()) + .map(str::to_owned) +} + +/// `get_configured_service_tiers`: each deployment's `model_info.service_tiers`, unvalidated. +pub fn configured_service_tiers(model_infos: &[Map]) -> Vec { + model_infos + .iter() + .map(|info| info.get("service_tiers").cloned().unwrap_or(Value::Null)) + .collect() +} + +/// What a listed name's deployments contribute to its `/v1/models` entry. +#[derive(Clone, Debug, Default, PartialEq)] +pub struct ListingInfo { + pub cost_map_keys: Vec, + pub max_input_tokens: Option, + pub max_output_tokens: Option, +} + +fn named_text(value: Option<&Value>) -> Option<&str> { + value + .and_then(Value::as_str) + .filter(|text| !text.is_empty()) +} + +/// `deployment_pricing_models(deployment)[0]`: `model_info.base_model`, then +/// `litellm_params.base_model`, then `litellm_params.model`; malformed blocks count as unset. +fn pricing_key(deployment: &Value) -> Option { + let block = |key: &str| { + deployment + .get(key) + .and_then(Value::as_object) + .filter(|_| !is_opaque(deployment.get(key).unwrap_or(&Value::Null))) + }; + let params = block("litellm_params"); + let info = block("model_info"); + let base_model = info + .and_then(|info| info.get("base_model")) + .filter(|base| truthy(base)) + .or_else(|| params.and_then(|params| params.get("base_model"))); + named_text(base_model) + .or_else(|| named_text(params.and_then(|params| params.get("model")))) + .map(str::to_owned) +} + +/// `get_model_listing_info` over the deployments a listed name names: the cost-map keys of their +/// models, deduplicated, and the widest token limits they configure or discovered. +pub async fn model_listing_info( + host: &dyn Host, + deployments: &[Value], +) -> Result { + let mut listing = ListingInfo::default(); + let empty = Value::Object(Map::new()); + for deployment in deployments { + let info = match get(deployment, "model_info")? { + Some(info) if truthy(info) => info, + _ => &empty, + }; + let id = get(info, "id")?.cloned().unwrap_or(Value::Null); + let configured = match info { + Value::Object(map) => map.clone(), + _ => Map::new(), + }; + let effective = effective_model_info(host, id, &configured).await?; + for (field, widest) in [ + ("max_input_tokens", &mut listing.max_input_tokens), + ("max_output_tokens", &mut listing.max_output_tokens), + ] { + if let Some(limit) = coerce_token_limit(effective.get(field)) { + *widest = Some(widest.map_or(limit, |current| current.max(limit))); + } + } + } + for key in deployments.iter().filter_map(pricing_key) { + if !listing.cost_map_keys.contains(&key) { + listing.cost_map_keys.push(key); + } + } + Ok(listing) +} + +#[cfg(test)] +mod tests { + use serde_json::json; + + use super::*; + + #[test] + fn coerce_token_limit_reads_python_ints() { + assert_eq!(coerce_token_limit(Some(&json!(128))), Some(128)); + assert_eq!(coerce_token_limit(Some(&json!(12.9))), Some(12)); + assert_eq!(coerce_token_limit(Some(&json!(" 1_000 "))), Some(1000)); + assert_eq!(coerce_token_limit(Some(&json!("128,000"))), None); + assert_eq!(coerce_token_limit(Some(&json!(true))), None); + assert_eq!(coerce_token_limit(None), None); + } + + #[test] + fn totals_add_like_python() { + assert_eq!(add(None, &json!(5)).unwrap(), json!(5)); + assert_eq!(add(Some(json!(5)), &json!(true)).unwrap(), json!(6)); + assert_eq!(add(Some(json!(5)), &json!(0.5)).unwrap(), json!(5.5)); + assert!(add(Some(json!(5)), &json!("1")).is_err()); + } + + #[test] + fn union_keeps_first_seen_order_and_empty_reads_as_none() { + assert_eq!( + union(&json!(["a", "b"]), &[json!("b"), json!("c")]), + json!(["a", "b", "c"]) + ); + assert_eq!(union(&Value::Null, &[]), Value::Null); + } + + #[test] + fn pricing_key_prefers_the_base_model() { + let deployment = json!({ + "litellm_params": {"model": "azure/x", "base_model": "params-base"}, + "model_info": {"base_model": "info-base"}, + }); + assert_eq!(pricing_key(&deployment).as_deref(), Some("info-base")); + let deployment = json!({"litellm_params": {"model": "azure/x"}, "model_info": "junk"}); + assert_eq!(pricing_key(&deployment).as_deref(), Some("azure/x")); + } +} diff --git a/litellm-rust/crates/router-core/src/health.rs b/litellm-rust/crates/router-core/src/health.rs new file mode 100644 index 00000000000..020a1a0f47d --- /dev/null +++ b/litellm-rust/crates/router-core/src/health.rs @@ -0,0 +1,960 @@ +//! [`Health`]: cooldowns, allowed-fails counters and per-minute usage, kept in the shared +//! [`RouterCache`] under the Python router's keys, encodings and TTLs: +//! +//! | state | key | value | TTL | +//! |---|---|---|---| +//! | cooldown | `deployment:{id}:cooldown` | `{exception_received, status_code, timestamp, cooldown_time}` | cooldown time | +//! | allowed fails | `deployment:{id}:allowed_fails[:{suffix}]` | int counter | cooldown time | +//! | usage | `global_router:{id}:{model}:tpm:{HH-MM}` and `:rpm:` | int counter | 60s | +//! | outcomes (local only) | `{id}:successes`, `{id}:fails` | int counter | 60s | +//! +//! From `router_utils/cooldown_cache.py`, `router_utils/cooldown_handlers.py`, +//! `router_utils/router_callbacks/track_deployment_metrics.py`, `RouterCacheEnum` in +//! `types/router.py` and `Router._increment_usage` / `deployment_callback_on_failure` in +//! `router.py`. + +use std::{ + sync::Arc, + time::{Duration, SystemTime, UNIX_EPOCH}, +}; + +use serde::{Deserialize, Serialize}; +use tracing::debug; + +use crate::{ + cache::RouterCache, + error::{AttemptError, Error, FailureKind}, + outcome::{OutcomeCallbacks, OutcomeWrite}, + registry::{Deployment, Registry}, + settings::Settings, +}; + +/// `RoutingArgs.ttl`: per-minute usage keys expire after a minute. +pub const USAGE_TTL: Duration = Duration::from_secs(60); +/// `DEFAULT_FAILURE_THRESHOLD_PERCENT`. +pub const FAILURE_THRESHOLD_PERCENT: f64 = 0.5; +/// `DEFAULT_FAILURE_THRESHOLD_MINIMUM_REQUESTS`. +pub const FAILURE_THRESHOLD_MINIMUM_REQUESTS: u64 = 5; +/// `SINGLE_DEPLOYMENT_TRAFFIC_FAILURE_THRESHOLD`. +pub const SINGLE_DEPLOYMENT_TRAFFIC_FAILURE_THRESHOLD: u64 = 1000; +/// `exception_masker` in `CooldownCache`: the visible prefix of `exception_received`. +const EXCEPTION_VISIBLE_PREFIX: usize = 50; + +pub trait Clock: Send + Sync { + fn now(&self) -> SystemTime; +} + +pub struct SystemClock; + +impl Clock for SystemClock { + fn now(&self) -> SystemTime { + SystemTime::now() + } +} + +fn unix_seconds(time: SystemTime) -> f64 { + time.duration_since(UNIX_EPOCH) + .map(|elapsed| elapsed.as_secs_f64()) + .unwrap_or_default() +} + +/// `get_utc_datetime().strftime("%H-%M")`. +pub fn current_minute(time: SystemTime) -> String { + let seconds = time + .duration_since(UNIX_EPOCH) + .map(|elapsed| elapsed.as_secs()) + .unwrap_or_default() + % 86_400; + format!("{:02}-{:02}", seconds / 3600, (seconds % 3600) / 60) +} + +pub fn cooldown_key(id: &str) -> String { + format!("deployment:{id}:cooldown") +} + +pub fn allowed_fails_key(id: &str, suffix: Option<&str>) -> String { + match suffix { + Some(suffix) => format!("deployment:{id}:allowed_fails:{suffix}"), + None => format!("deployment:{id}:allowed_fails"), + } +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq, strum::Display, strum::EnumString)] +pub enum UsageMetric { + #[strum(serialize = "tpm")] + Tpm, + #[strum(serialize = "rpm")] + Rpm, + /// Input tokens per minute, which the io-token rate limit check counts. + #[strum(serialize = "itpm")] + Itpm, + /// Output tokens per minute. + #[strum(serialize = "otpm")] + Otpm, +} + +/// `RouterCacheEnum.TPM` / `RPM` / `ITPM` / `OTPM`: `model` is the deployment's +/// `litellm_params.model`. +pub fn usage_key(id: &str, model: &str, metric: UsageMetric, minute: &str) -> String { + format!("global_router:{id}:{model}:{metric}:{minute}") +} + +/// `CooldownCacheValue`. +#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)] +pub struct CooldownValue { + pub exception_received: String, + pub status_code: String, + pub timestamp: f64, + /// Whole seconds are written as an int, as Python's default `cooldown_time` is, so the + /// "Try again in N seconds" text reads the same. + #[serde(serialize_with = "whole_as_int")] + pub cooldown_time: f64, +} + +fn whole_as_int(value: &f64, serializer: S) -> Result { + if value.fract() == 0.0 && value.abs() < i64::MAX as f64 { + serializer.serialize_i64(*value as i64) + } else { + serializer.serialize_f64(*value) + } +} + +/// `litellm.should_retry`. +pub fn should_retry_status(status: u16) -> bool { + matches!(status, 408 | 409 | 429) || status >= 500 +} + +/// `_is_cooldown_required`: a 4xx cools down only for 429, 401, 402, 404 and 408, every other +/// status does, and a connection error or an empty status never does. +pub fn is_cooldown_required(failure: &AttemptError) -> bool { + if failure.kind == FailureKind::ApiConnection || failure.status_absent { + return false; + } + match failure.cooldown_status() { + status @ 400..500 => matches!(status, 429 | 401 | 402 | 408 | 404), + _ => true, + } +} + +/// `SensitiveDataMasker(visible_prefix=50, visible_suffix=0, mask_short_values=False)`. +fn mask_exception(text: &str) -> String { + let length = text.chars().count(); + if length <= EXCEPTION_VISIBLE_PREFIX { + return text.to_owned(); + } + let mut masked: String = text.chars().take(EXCEPTION_VISIBLE_PREFIX).collect(); + masked.extend(std::iter::repeat_n('*', length - EXCEPTION_VISIBLE_PREFIX)); + masked +} + +/// `_EXCEPTION_POLICY_FIELDS`: a deployment policy's fields in the order Python matches them +/// (a content-policy violation is a bad request, so it comes first). +const DEPLOYMENT_POLICY_FIELDS: [(&str, &str); 9] = [ + ( + "ContentPolicyViolationError", + "ContentPolicyViolationErrorAllowedFails", + ), + ("BadRequestError", "BadRequestErrorAllowedFails"), + ("AuthenticationError", "AuthenticationErrorAllowedFails"), + ("Timeout", "TimeoutErrorAllowedFails"), + ("RateLimitError", "RateLimitErrorAllowedFails"), + ("InternalServerError", "InternalServerErrorAllowedFails"), + ( + "ServiceUnavailableError", + "ServiceUnavailableErrorAllowedFails", + ), + ("BadGatewayError", "BadGatewayErrorAllowedFails"), + ("NotFoundError", "NotFoundErrorAllowedFails"), +]; + +/// `_resolve_allowed_fails_from_policy` over a deployment's `model_info.allowed_fails_policy`: +/// the first field set for a class the failure is an instance of. +pub fn deployment_allowed_fails( + policy: Option<&serde_json::Map>, + failure: &AttemptError, +) -> Option { + let policy = policy?; + DEPLOYMENT_POLICY_FIELDS + .iter() + .filter(|(class, _)| failure.is_a(class)) + .find_map(|(_, field)| policy.get(*field).and_then(serde_json::Value::as_f64)) +} + +/// What the cooldown rules need to know about the deployment's group. +#[derive(Clone, Copy, Debug, Default)] +pub struct GroupShape { + /// The deployment is the only one in its model group. + pub single_deployment: bool, + /// The requested group has a fallback configured. + pub has_fallback_group: bool, +} + +/// `_core_team_model_has_alternatives`: another routable deployment serves the deployment's team +/// public name. +pub(crate) fn team_model_has_alternatives(registry: &Registry, deployment: &Deployment) -> bool { + let info = &deployment.model_info; + let (Some(team_id), Some(public_name)) = (&info.team_id, &info.team_public_model_name) else { + return false; + }; + registry + .team_named(team_id, public_name) + .is_some_and(|siblings| { + siblings + .iter() + .filter(|sibling| !sibling.model_info.blocked) + .count() + > 1 + }) +} + +/// `_core_routing_group_has_alternatives`: `model_group` names a callable routing group whose +/// members span more than one deployment. +pub(crate) fn routing_group_has_alternatives(registry: &Registry, model_group: &str) -> bool { + let resolved = registry.alias_of(model_group).unwrap_or(model_group); + registry.routing_group_named(resolved).is_some_and(|group| { + group + .models + .iter() + .map(|member| registry.named(member).map_or(0, |members| members.len())) + .sum::() + > 1 + }) +} + +/// An outcome the router did not see itself (`Health.report`). +#[derive(Clone, Debug)] +pub enum Outcome { + Success { total_tokens: u64 }, + Failure(AttemptError), +} + +#[derive(Clone)] +pub struct Health { + cache: Arc, + clock: Arc, +} + +impl Health { + pub fn new(cache: Arc, clock: Arc) -> Self { + Self { cache, clock } + } + + pub fn cache(&self) -> &Arc { + &self.cache + } + + pub(crate) fn now_time(&self) -> SystemTime { + self.clock.now() + } + + fn now(&self) -> f64 { + unix_seconds(self.clock.now()) + } + + /// The ids among `ids` with an unexpired cooldown (`async_get_active_cooldowns`). + pub async fn active_cooldowns( + &self, + ids: &[String], + ) -> Result, Error> { + let keys: Vec = ids.iter().map(|id| cooldown_key(id)).collect(); + let values = self.cache.async_batch_get_cache(&keys).await?; + let now = self.now(); + Ok(ids + .iter() + .zip(values) + .filter_map(|(id, value)| { + let value: CooldownValue = serde_json::from_value(value?).ok()?; + (value.timestamp + value.cooldown_time - now > 0.0).then(|| (id.clone(), value)) + }) + .collect()) + } + + /// `add_deployment_to_cooldown`. + pub async fn add_cooldown( + &self, + id: &str, + failure: &AttemptError, + cooldown: Duration, + ) -> Result<(), Error> { + let value = CooldownValue { + exception_received: mask_exception(failure.detail.as_deref().unwrap_or_default()), + status_code: failure.cooldown_status().to_string(), + timestamp: self.now(), + cooldown_time: cooldown.as_secs_f64(), + }; + let value = serde_json::to_value(value).unwrap_or_default(); + self.cache + .async_set_cache(&cooldown_key(id), value, Some(cooldown)) + .await?; + Ok(()) + } + + async fn local_count(&self, key: &str) -> Result { + Ok(self + .cache + .async_get_local(key) + .await? + .and_then(|value| value.as_f64()) + .map_or(0, |count| count as u64)) + } + + /// The per-minute (successes, fails) this worker saw (`{id}:successes`, `{id}:fails`, kept + /// `local_only`). + pub async fn outcomes_this_minute(&self, id: &str) -> Result<(u64, u64), Error> { + Ok(( + self.local_count(&format!("{id}:successes")).await?, + self.local_count(&format!("{id}:fails")).await?, + )) + } + + /// `increment_deployment_successes_for_current_minute`: returns the key it counted under. + pub async fn count_success(&self, id: &str) -> Result { + let key = format!("{id}:successes"); + self.cache + .async_increment_local(&key, 1.0, Some(USAGE_TTL)) + .await?; + Ok(key) + } + + /// `increment_deployment_failures_for_current_minute`. + pub async fn count_failure(&self, id: &str) -> Result<(), Error> { + self.cache + .async_increment_local(&format!("{id}:fails"), 1.0, Some(USAGE_TTL)) + .await?; + Ok(()) + } + + /// Cooldown duration precedence from `deployment_callback_on_failure`: the deployment's + /// `cooldown_time`, then the provider's retry-after, then the router default. + pub fn cooldown_duration( + settings: &Settings, + deployment: &Deployment, + failure: &AttemptError, + ) -> Duration { + Self::cooldown_for_call(settings, deployment, None, failure) + } + + /// `cooldown_duration` for a call whose own params carry `cooldown_time` (a request's + /// value wins over the deployment's `litellm_params`, not over its `model_info`). + pub fn cooldown_for_call( + settings: &Settings, + deployment: &Deployment, + call_cooldown: Option, + failure: &AttemptError, + ) -> Duration { + let configured = deployment + .model_info + .cooldown_time + .or(call_cooldown) + .or_else(|| { + deployment + .raw_params + .get("cooldown_time") + .and_then(|v| v.as_f64()) + }) + .filter(|seconds| *seconds >= 0.0) + .and_then(|seconds| Duration::try_from_secs_f64(seconds).ok()); + configured + .or(failure.retry_after) + .unwrap_or(settings.cooldown_time) + } + + /// The deployment's own `cooldown_time`: `model_info` first, then `litellm_params` + /// (`_first_present(model_info, litellm_params, key="cooldown_time")`). It is both the + /// cooldown and the TTL of the deployment's `allowed_fails:generic` counter. + fn configured_cooldown(deployment: &Deployment) -> Option { + deployment + .model_info + .cooldown_time + .or_else(|| { + deployment + .raw_params + .get("cooldown_time") + .and_then(|v| v.as_f64()) + }) + .filter(|seconds| *seconds >= 0.0) + .and_then(|seconds| Duration::try_from_secs_f64(seconds).ok()) + } + + /// Record a failed attempt and cool the deployment down when the rules say so. Returns the + /// cooldown applied, if any. + pub async fn record_failure( + &self, + settings: &Settings, + deployment: &Deployment, + shape: GroupShape, + call_cooldown: Option, + failure: &AttemptError, + ) -> Result, Error> { + if failure.outside_deployment { + return Ok(None); + } + self.count_failure(&deployment.id).await?; + let cooldown = Self::cooldown_for_call(settings, deployment, call_cooldown, failure); + self.cool_down_after(settings, deployment, shape, failure, cooldown) + .await + } + + /// `set_cooldown_deployments` for a failure already counted: cool the deployment down for + /// `cooldown` when the rules say so. Returns the cooldown applied, if any. + pub async fn cool_down_after( + &self, + settings: &Settings, + deployment: &Deployment, + shape: GroupShape, + failure: &AttemptError, + cooldown: Duration, + ) -> Result, Error> { + if !Self::should_run_cooldown_logic(settings, deployment, failure, cooldown) { + return Ok(None); + } + if !self + .should_cooldown(settings, deployment, shape, failure) + .await? + { + return Ok(None); + } + debug!(deployment_id = %deployment.id, cooldown_seconds = cooldown.as_secs_f64(), "cooling deployment down"); + self.add_cooldown(&deployment.id, failure, cooldown).await?; + Ok(Some(cooldown)) + } + + /// `_should_run_cooldown_logic`. + fn should_run_cooldown_logic( + settings: &Settings, + deployment: &Deployment, + failure: &AttemptError, + cooldown: Duration, + ) -> bool { + !matches!( + failure.kind, + FailureKind::Declined + | FailureKind::CallerCredentialAuthentication + | FailureKind::CallerCredentialRateLimit + ) && !failure.caller_scoped + && !cooldown.is_zero() + && !settings.disable_cooldowns + && (is_cooldown_required(failure) + || deployment_allowed_fails( + deployment.model_info.allowed_fails_policy.as_ref(), + failure, + ) + .is_some()) + && !settings + .provider_default_deployment_ids + .contains(&deployment.id) + } + + /// `_should_cooldown_deployment`. + async fn should_cooldown( + &self, + settings: &Settings, + deployment: &Deployment, + shape: GroupShape, + failure: &AttemptError, + ) -> Result { + let status = failure.cooldown_status(); + let policy_fails = settings.allowed_fails_for(failure); + let lone = shape.single_deployment && !shape.has_fallback_group; + let deployment_policy = deployment.model_info.allowed_fails_policy.as_ref(); + let deployment_fails = deployment_allowed_fails(deployment_policy, failure); + if shape.single_deployment + && status == 402 + && deployment_fails.is_none() + && policy_fails.is_none() + { + return Ok(false); + } + let deployment_allowed = deployment.model_info.allowed_fails; + if deployment_policy.is_some() || deployment_allowed.is_some() { + // `_should_cooldown_based_on_deployment_policy`: a deployment-wide `allowed_fails` + // keeps the lone-deployment safety net; a policy entry naming the error does not. + if deployment_fails.is_none() + && lone + && (deployment_allowed.is_some() || policy_fails.is_none()) + { + return Ok(false); + } + let (allowed, suffix) = match deployment_fails { + Some(allowed) => ( + Some(allowed), + failure.class_names().first().map(|name| (*name).to_owned()), + ), + None => ( + deployment_allowed.map(f64::from), + deployment_allowed.map(|_| "generic".to_owned()), + ), + }; + let allowed = allowed + .or_else(|| policy_fails.map(f64::from)) + .unwrap_or_else(|| f64::from(settings.effective_allowed_fails())); + let ttl = + Self::configured_cooldown(deployment).unwrap_or_else(|| settings.fails_window()); + return self + .exceeds_allowed_fails(&deployment.id, suffix.as_deref(), allowed, ttl) + .await; + } + if settings.allowed_fails_policy.is_none() && !settings.allowed_fails_set() { + let (successes, fails) = self.outcomes_this_minute(&deployment.id).await?; + let total = successes + fails; + let percent_fails = if total > 0 { + fails as f64 / total as f64 + } else { + 0.0 + }; + return Ok((status == 429 && !shape.single_deployment) + || (percent_fails == 1.0 && total >= SINGLE_DEPLOYMENT_TRAFFIC_FAILURE_THRESHOLD) + || (percent_fails > FAILURE_THRESHOLD_PERCENT + && total >= FAILURE_THRESHOLD_MINIMUM_REQUESTS + && !shape.single_deployment) + || !should_retry_status(status)); + } + if lone && policy_fails.is_none() { + return Ok(false); + } + let allowed = policy_fails.unwrap_or_else(|| settings.effective_allowed_fails()); + self.exceeds_allowed_fails( + &deployment.id, + None, + f64::from(allowed), + settings.fails_window(), + ) + .await + } + + /// `should_cooldown_based_on_allowed_fails_policy`: count this failure, then compare. + async fn exceeds_allowed_fails( + &self, + id: &str, + suffix: Option<&str>, + allowed: f64, + ttl: Duration, + ) -> Result { + let fails = self + .cache + .async_increment(&allowed_fails_key(id, suffix), 1.0, Some(ttl)) + .await?; + Ok(fails > allowed) + } + + /// Count a success, and the tokens and request against the deployment's per-minute usage + /// when `tracks` says it has limits to track. The count is the success callback's; the usage + /// is `make_call`'s, which runs whatever the callbacks do. + pub async fn record_success( + &self, + callbacks: OutcomeCallbacks, + deployment: &Deployment, + total_tokens: u64, + tracks: bool, + ) -> Result<(), Error> { + callbacks + .record(OutcomeWrite::Success, self.count_success(&deployment.id)) + .await?; + if !tracks { + return Ok(()); + } + callbacks + .record( + OutcomeWrite::AttemptUsage, + self.add_usage(deployment, None, total_tokens as f64, 1.0), + ) + .await + } + + /// `_increment_deployment_usage`, for a deployment that tracks usage: tokens and requests + /// under `model` (metadata's `deployment`, which a wildcard route fills; else the + /// deployment's own), unless there is nothing to add. + pub async fn add_usage( + &self, + deployment: &Deployment, + model: Option<&str>, + total_tokens: f64, + requests: f64, + ) -> Result<(), Error> { + if total_tokens <= 0.0 && requests <= 0.0 { + return Ok(()); + } + let minute = current_minute(self.clock.now()); + let model = model.unwrap_or(&deployment.params.model); + self.cache + .async_increment_usage( + &[ + ( + usage_key(&deployment.id, model, UsageMetric::Tpm, &minute), + total_tokens, + ), + ( + usage_key(&deployment.id, model, UsageMetric::Rpm, &minute), + requests, + ), + ], + Some(USAGE_TTL), + ) + .await?; + Ok(()) + } + + /// `async_deployment_callback_on_failure`: a failed request counts against the deployment's + /// per-minute requests under `model` (see [`Self::add_usage`]), whether or not it tracks usage. + pub async fn count_failed_request( + &self, + deployment: &Deployment, + model: Option<&str>, + ) -> Result<(), Error> { + let minute = current_minute(self.clock.now()); + self.cache + .async_increment( + &usage_key( + &deployment.id, + model.unwrap_or(&deployment.params.model), + UsageMetric::Rpm, + &minute, + ), + 1.0, + Some(USAGE_TTL), + ) + .await?; + Ok(()) + } + + /// This minute's (tpm, rpm) for a deployment. + pub async fn usage(&self, deployment: &Deployment) -> Result<(u64, u64), Error> { + let minute = current_minute(self.clock.now()); + let model = &deployment.params.model; + let keys = [UsageMetric::Tpm, UsageMetric::Rpm] + .map(|metric| usage_key(&deployment.id, model, metric, &minute)); + let values = self.cache.async_batch_get_cache(&keys).await?; + let count = |index: usize| { + values + .get(index) + .cloned() + .flatten() + .and_then(|value| value.as_f64()) + .map_or(0, |count| count as u64) + }; + Ok((count(0), count(1))) + } +} + +#[cfg(test)] +mod tests { + use std::collections::BTreeMap; + + use rstest::rstest; + + use super::*; + use crate::{ + cache::InMemoryRouterCache, + config::{RouterConfig, RouterDynamicConfig}, + registry::{DeploymentSpec, RoutingGroupSpec}, + }; + + #[rstest] + #[case::rate_limit(FailureKind::RateLimit, Some(429), true)] + #[case::auth(FailureKind::Authentication, Some(401), true)] + #[case::payment(FailureKind::Other, Some(402), true)] + #[case::not_found(FailureKind::NotFound, Some(404), true)] + #[case::timeout(FailureKind::Timeout, Some(408), true)] + #[case::bad_request(FailureKind::BadRequest, Some(400), false)] + #[case::forbidden(FailureKind::PermissionDenied, Some(403), false)] + #[case::server(FailureKind::InternalServer, Some(500), true)] + #[case::unreadable_status(FailureKind::Other, None, true)] + #[case::connection(FailureKind::ApiConnection, Some(500), false)] + fn cooldown_required_follows_python( + #[case] kind: FailureKind, + #[case] status: Option, + #[case] expected: bool, + ) { + assert_eq!( + is_cooldown_required(&AttemptError::new(kind, status)), + expected + ); + } + + /// An exception with no `status_code` is Python's status `""`, which no rule cools down. + #[rstest] + fn an_empty_status_requires_no_cooldown() { + let mut failure = AttemptError::new(FailureKind::Other, None); + failure.status_absent = true; + assert!(!is_cooldown_required(&failure)); + } + + #[rstest] + #[case::midnight(0, "00-00")] + #[case::afternoon(13 * 3600 + 7 * 60 + 59, "13-07")] + fn current_minute_is_utc_hours_and_minutes(#[case] seconds: u64, #[case] expected: &str) { + let time = UNIX_EPOCH + Duration::from_secs(86_400 * 3 + seconds); + assert_eq!(current_minute(time), expected); + } + + #[rstest] + #[case::content_policy_first(&["ContentPolicyViolationError", "BadRequestError"], Some(1.0))] + #[case::bad_request(&["BadRequestError"], Some(2.0))] + #[case::unnamed(&["Timeout"], None)] + fn a_deployment_policy_matches_in_python_order( + #[case] classes: &[&str], + #[case] expected: Option, + ) { + let policy = serde_json::json!({ + "BadRequestErrorAllowedFails": 2, + "ContentPolicyViolationErrorAllowedFails": 1, + }); + let mut failure = AttemptError::new(FailureKind::BadRequest, Some(400)); + failure.classes = classes.iter().map(|class| (*class).into()).collect(); + assert_eq!( + deployment_allowed_fails(policy.as_object(), &failure), + expected + ); + } + + fn deployment(id: &str, info: serde_json::Value) -> DeploymentSpec { + let mut model_info = serde_json::json!({"id": id}); + model_info + .as_object_mut() + .unwrap() + .extend(info.as_object().unwrap().clone()); + serde_json::from_value(serde_json::json!({ + "model_name": id, + "litellm_params": {"model": format!("openai/{id}")}, + "model_info": model_info, + })) + .unwrap() + } + + fn health() -> Health { + Health::new( + Arc::new(InMemoryRouterCache::default()), + Arc::new(SystemClock), + ) + } + + fn settings(dynamic: RouterDynamicConfig) -> Settings { + Settings::new(&RouterConfig::default(), Arc::new(dynamic)) + } + + #[rstest] + #[case::held(false, false, 1, true)] + #[case::caller_scoped(false, true, 1, false)] + #[case::outside_the_deployment(true, false, 0, false)] + #[tokio::test] + async fn the_failure_exits_decide_what_is_counted_and_cooled( + #[case] outside_deployment: bool, + #[case] caller_scoped: bool, + #[case] counted: u64, + #[case] cooled: bool, + ) { + let health = health(); + let deployment = Deployment::from_spec(deployment("a", serde_json::json!({}))).unwrap(); + let mut failure = AttemptError::new(FailureKind::RateLimit, Some(429)); + failure.outside_deployment = outside_deployment; + failure.caller_scoped = caller_scoped; + + let cooldown = health + .record_failure( + &settings(RouterDynamicConfig::default()), + &deployment, + GroupShape::default(), + None, + &failure, + ) + .await + .unwrap(); + + assert_eq!(cooldown.is_some(), cooled); + assert_eq!(health.outcomes_this_minute("a").await.unwrap().1, counted); + } + + #[rstest] + #[case::authentication( + FailureKind::CallerCredentialAuthentication, + 401, + "AuthenticationError" + )] + #[case::rate_limit(FailureKind::CallerCredentialRateLimit, 429, "RateLimitError")] + #[tokio::test] + async fn a_caller_credential_failure_counts_without_cooling( + #[case] kind: FailureKind, + #[case] status: u16, + #[case] base_class: &str, + ) { + let health = health(); + let deployment = Deployment::from_spec(deployment("a", serde_json::json!({}))).unwrap(); + let failure = AttemptError::new(kind, Some(status)); + + let cooldown = health + .record_failure( + &settings(RouterDynamicConfig::default()), + &deployment, + GroupShape::default(), + None, + &failure, + ) + .await + .unwrap(); + + assert_eq!(cooldown, None); + assert!(failure.is_a(base_class)); + assert_eq!(health.outcomes_this_minute("a").await.unwrap().1, 1); + } + + /// `litellm.allowed_fails` written after the build: the copied 3 starts counting instead of + /// the failure-ratio rule, which cools a 429 down at once. + #[rstest] + #[case::still_the_process_default(true, true)] + #[case::the_process_default_moved(false, false)] + #[tokio::test] + async fn the_copied_allowed_fails_counts_once_it_differs_from_the_process_default( + #[case] allowed_fails_is_default: bool, + #[case] cooled: bool, + ) { + let settings = settings(RouterDynamicConfig { + allowed_fails: Some(3), + allowed_fails_is_default, + ..RouterDynamicConfig::default() + }); + let deployment = Deployment::from_spec(deployment("a", serde_json::json!({}))).unwrap(); + let failure = AttemptError::new(FailureKind::RateLimit, Some(429)); + + let cooldown = health() + .record_failure( + &settings, + &deployment, + GroupShape::default(), + None, + &failure, + ) + .await + .unwrap(); + + assert_eq!(cooldown.is_some(), cooled); + } + + fn team(id: &str, team_id: &str, blocked: bool) -> DeploymentSpec { + deployment( + id, + serde_json::json!({"team_id": team_id, "team_public_model_name": "tp", "blocked": blocked}), + ) + } + + #[rstest] + #[case::a_sibling(team("b", "t1", false), true)] + #[case::a_blocked_sibling(team("b", "t1", true), false)] + #[case::another_teams_deployment(team("b", "t2", false), false)] + #[case::no_team(deployment("b", serde_json::json!({})), false)] + fn a_team_deployment_has_alternatives_under_its_public_name( + #[case] other: DeploymentSpec, + #[case] expected: bool, + ) { + let registry = Registry::build( + vec![team("a", "t1", false), other], + BTreeMap::new(), + Vec::new(), + ) + .unwrap(); + let first = registry.deployment("a").unwrap(); + assert_eq!(team_model_has_alternatives(®istry, &first), expected); + } + + #[rstest] + #[case::spans_two_deployments(&["a", "b"], "rg", true)] + #[case::one_deployment(&["a"], "rg", false)] + #[case::through_its_alias(&["a", "b"], "alias", true)] + #[case::a_member_called_directly(&["a", "b"], "a", false)] + fn a_routing_group_has_alternatives_past_one_deployment( + #[case] members: &[&str], + #[case] requested: &str, + #[case] expected: bool, + ) { + let group: RoutingGroupSpec = serde_json::from_value(serde_json::json!({ + "group_name": "rg", "models": members, "routing_strategy": "simple-shuffle", + })) + .unwrap(); + let registry = Registry::build( + vec![ + deployment("a", serde_json::json!({})), + deployment("b", serde_json::json!({})), + ], + BTreeMap::from([("alias".to_owned(), "rg".to_owned())]), + vec![group], + ) + .unwrap(); + assert_eq!( + routing_group_has_alternatives(®istry, requested), + expected + ); + } + + /// An empty status is counted and cools nothing down, unless the deployment's own policy + /// names the error. + #[rstest] + #[case::no_policy(serde_json::json!({}), false)] + #[case::policy_names_the_error( + serde_json::json!({"allowed_fails_policy": {"BadRequestErrorAllowedFails": 0}}), + true + )] + #[tokio::test] + async fn an_empty_status_on_the_attempt_path_needs_the_deployments_policy_to_cool( + #[case] info: serde_json::Value, + #[case] cooled: bool, + ) { + let health = health(); + let deployment = Deployment::from_spec(deployment("a", info)).unwrap(); + let mut failure = AttemptError::new(FailureKind::BadRequest, None); + failure.status_absent = true; + + let cooldown = health + .record_failure( + &settings(RouterDynamicConfig::default()), + &deployment, + GroupShape::default(), + None, + &failure, + ) + .await + .unwrap(); + + assert_eq!(cooldown.is_some(), cooled); + assert_eq!(health.outcomes_this_minute("a").await.unwrap().1, 1); + } + + #[rstest] + #[case::the_deployments_own_model(None, "openai/a")] + #[case::a_filled_wildcard(Some("openai/gpt-4o"), "openai/gpt-4o")] + #[tokio::test] + async fn usage_is_keyed_by_the_model_served(#[case] model: Option<&str>, #[case] keyed: &str) { + let cache = InMemoryRouterCache::default(); + let clock = Arc::new(SystemClock); + let health = Health::new(Arc::new(cache.clone()), clock.clone()); + let deployment = Deployment::from_spec(deployment("a", serde_json::json!({}))).unwrap(); + + health + .add_usage(&deployment, model, 15.0, 1.0) + .await + .unwrap(); + + let minute = current_minute(clock.now()); + assert_eq!( + ( + cache.get(&usage_key("a", keyed, UsageMetric::Tpm, &minute)), + cache.get(&usage_key("a", keyed, UsageMetric::Rpm, &minute)), + ), + (Some(serde_json::json!(15)), Some(serde_json::json!(1))), + ); + } + + #[rstest] + fn keys_match_python() { + assert_eq!( + ( + cooldown_key("abc"), + allowed_fails_key("abc", Some("generic")), + usage_key("abc", "openai/gpt-4o", UsageMetric::Tpm, "09-05"), + ), + ( + "deployment:abc:cooldown".to_owned(), + "deployment:abc:allowed_fails:generic".to_owned(), + "global_router:abc:openai/gpt-4o:tpm:09-05".to_owned(), + ) + ); + } +} diff --git a/litellm-rust/crates/router-core/src/health_routing.rs b/litellm-rust/crates/router-core/src/health_routing.rs new file mode 100644 index 00000000000..bc3e43a6e92 --- /dev/null +++ b/litellm-rust/crates/router-core/src/health_routing.rs @@ -0,0 +1,126 @@ +//! Health-check-driven routing: the background health checks' results, kept in the shared +//! [`RouterCache`](crate::cache::RouterCache) under one key as Python's `DeploymentHealthCache` +//! writes them, and the filter `_async_filter_health_check_unhealthy_deployments` / +//! `_filter_health_check_unhealthy_deployments` run before the cooldown filter. + +use std::{collections::BTreeSet, sync::Arc, time::UNIX_EPOCH}; + +use serde_json::Value; + +use crate::{health::Health, registry::Deployment, settings::Settings}; + +/// `DeploymentHealthCache.CACHE_KEY`. +pub const HEALTH_STATE_KEY: &str = "litellm:health_check:deployment_health_state"; + +/// `_extract_unhealthy_ids`: the ids marked unhealthy whose result is younger than `staleness`. +pub fn unhealthy_ids(raw: Option<&Value>, now: f64, staleness: f64) -> BTreeSet { + let Some(Value::Object(states)) = raw else { + return BTreeSet::new(); + }; + states + .iter() + .filter(|(_, state)| { + let Value::Object(state) = state else { + return false; + }; + // `not state.get("is_healthy", True)`: any falsy value is unhealthy. + let unhealthy = match state.get("is_healthy") { + None => false, + Some(Value::Null) => true, + Some(Value::Bool(value)) => !value, + Some(Value::Number(value)) => value.as_f64() == Some(0.0), + Some(Value::String(value)) => value.is_empty(), + Some(Value::Array(value)) => value.is_empty(), + Some(Value::Object(value)) => value.is_empty(), + }; + let timestamp = state + .get("timestamp") + .and_then(Value::as_f64) + .unwrap_or(0.0); + unhealthy && now - timestamp < staleness + }) + .map(|(id, _)| id.clone()) + .collect() +} + +/// What the filter does when it would leave nothing. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum WhenAllUnhealthy { + /// The safety net: every candidate is kept. + KeepAll, + /// A health check probe sees no deployment. + Nothing, +} + +impl Health { + /// The deployments the background health checks have not marked unhealthy, as the + /// router's health-check routing settings say. + pub async fn health_check_filter( + &self, + settings: &Settings, + deployments: Vec>, + when_all_unhealthy: WhenAllUnhealthy, + ) -> Vec> { + if !settings.enable_health_check_routing { + return deployments; + } + let scoped = settings.background_health_check_model_groups.as_deref(); + // With an allowed_fails_policy cooldowns are the only exclusion, unless the filter is + // scoped to listed groups. + if settings.allowed_fails_policy.is_some() && scoped.is_none() { + return deployments; + } + // A read that fails counts as no health state. + let raw = self + .cache() + .async_get_cache(HEALTH_STATE_KEY) + .await + .ok() + .flatten(); + let now = self + .now_time() + .duration_since(UNIX_EPOCH) + .map(|elapsed| elapsed.as_secs_f64()) + .unwrap_or_default(); + let unhealthy = unhealthy_ids( + raw.as_ref(), + now, + settings.health_check_staleness.as_secs_f64(), + ); + if unhealthy.is_empty() { + return deployments; + } + let filtered: Vec> = deployments + .iter() + .filter(|d| { + !unhealthy.contains(&d.id) + || scoped.is_some_and(|groups| !groups.contains(&d.model_name)) + }) + .cloned() + .collect(); + match (filtered.is_empty(), when_all_unhealthy) { + (false, _) => filtered, + (true, WhenAllUnhealthy::KeepAll) => deployments, + (true, WhenAllUnhealthy::Nothing) => Vec::new(), + } + } +} + +#[cfg(test)] +mod tests { + use rstest::rstest; + use serde_json::json; + + use super::*; + + #[rstest] + #[case::fresh_unhealthy(json!({"a": {"is_healthy": false, "timestamp": 95.0}}), vec!["a"])] + #[case::stale(json!({"a": {"is_healthy": false, "timestamp": 50.0}}), vec![])] + #[case::healthy(json!({"a": {"is_healthy": true, "timestamp": 95.0}}), vec![])] + #[case::no_flag(json!({"a": {"timestamp": 95.0}}), vec![])] + #[case::not_a_dict(json!({"a": 1}), vec![])] + fn unhealthy_ids_follow_python(#[case] raw: Value, #[case] expected: Vec<&str>) { + let ids = unhealthy_ids(Some(&raw), 100.0, 10.0); + assert_eq!(ids.iter().map(String::as_str).collect::>(), expected); + } +} diff --git a/litellm-rust/crates/router-core/src/host.rs b/litellm-rust/crates/router-core/src/host.rs new file mode 100644 index 00000000000..0a72be8f980 --- /dev/null +++ b/litellm-rust/crates/router-core/src/host.rs @@ -0,0 +1,726 @@ +//! [`Host`]: the router core's only boundary. Everything the core needs from the embedding +//! process goes through it, and every attempt runs through [`Host::invoke`]: the core never runs +//! an operation itself. `PythonHost` serves it in a Python process (attempts call `litellm.`, +//! and the deployment build asks it for secrets, validation, providers and side effects). + +use std::{collections::BTreeMap, sync::Arc, time::Duration}; + +use async_trait::async_trait; +use serde::Deserialize; +use serde_json::{Map, Value}; + +use crate::{ + cache::RouterCache, + config::{Op, RetryPolicy, RoutingStrategy}, + error::{AttemptError, Error, FailureKind, Payload}, + fallback::Fallbacks, + health::Health, + mock::MockFailure, + outcome::OutcomeCallbacks, + precall::PreCallFacts, + prerouting::{PreRoute, PreRouted}, + registry::{Deployment, Registry}, + settings::Settings, + strategy::StrategyScope, +}; + +/// One call into the router. +#[derive(Clone, Debug, Default)] +pub struct RouterRequest { + /// The requested name: a model group, alias, routing group, deployment id or wildcard match. + pub model: String, + /// The request params, for the plugins and host steps that read them. + pub body: Value, + /// `user_api_key_team_id`, from the first metadata bucket that holds one + /// (`get_request_team_id`). + pub team_id: Option, + /// The `metadata` bucket's `user_api_key_team_id` alone: the team a failed level lists the + /// group's deployments for (`async_function_with_fallbacks_common_utils`, + /// `_stream_order_levels`). `None` lists every team's. + pub listing_team_id: Option, + /// The caller's `user_role` is `proxy_admin`. + pub proxy_admin: bool, + /// The caller's `models` plus `team_models` (`user_api_key_auth`), which narrow a caller + /// allowed only through access groups to those groups' deployments. + pub allowed_models: Vec, + /// For a proxy admin, the model the request names (`model`, else the metadata + /// `model_group`), which keeps one team's deployments listed under it (`filter_team_based_models`). + pub admin_model: Option, + pub specific_deployment: bool, + /// A per-request `num_retries`, which wins over the router's. + pub num_retries: Option, + /// A per-request `num_retries` that is no count, which replaces the router's as well. + pub uncounted_retries: Option, + pub disable_fallbacks: bool, + /// `_router_weights` as the request passed it: model group to deployment id to weight, + /// validated when simple-shuffle picks ([`crate::weights::validate`]). + pub router_weights: Option, + /// The deployment a call on an existing resource (a video id) names: picked whenever the + /// filters keep it (`_narrow_to_video_deployment`). + pub pinned_deployment_id: Option, + /// A per-request `routing_strategy` override, as Python normalizes it. + pub routing_strategy: Option, + /// Whether tag routing may apply to this request (router, request or model setting), so the + /// host's tag stage runs. + pub filter_tags: bool, + /// A pass-through pick: only `use_in_pass_through` deployments may serve it. + pub pass_through: Option, + /// `_target_order`: the order level to pick from instead of the lowest. + pub target_order: Option, + /// `_excluded_deployment_ids`: never picked, even when that leaves nothing. + pub excluded_ids: Vec, + /// `_retry_skipped_deployment_ids`: skipped unless that leaves nothing. + pub retry_skipped_ids: Vec, + /// The endpoint `endpoint_selection` names around the call, which deployments that declare + /// `enforce_supported_endpoints` may refuse ([`crate::eligibility`]). + pub endpoint: Option, + /// A web search request on an async pick: `filter_web_search_deployments` drops the + /// deployments that declare no support. The sync pick never runs that filter. + pub web_search: bool, + /// The handler name of a call only some providers serve (`_ROUTE_CAPABILITY_RESOLVERS`): + /// the deployments whose provider cannot serve it are dropped unless that leaves nothing. + pub route_handler: Option, + /// A Decisions call carrying audio or videos: the deployments that refuse them are skipped + /// unless that leaves nothing (`_without_decisions_deployments_refusing_media`). + pub decisions_media: bool, + /// `f"{provider}/{model}"` for a lookup that cannot ask the host + /// (`_core_get_routable_deployments`); a pick asks [`Host::provider_request`]. + pub wildcard_request: Option, + /// The request carries an image (`_request_carries_image`), so deployments that declare no + /// vision are skipped while any other remains. + pub has_image: bool, + /// An async pick: the host's pre-routing step runs before each pick resolves its name. + pub pre_routes: bool, + /// Which of the router's outcome callbacks run for this request. + pub outcome_callbacks: OutcomeCallbacks, + /// The requested name is no model group (a search tool name): nothing is selected, every + /// attempt is handed to the host as is, and no outcome is recorded for a deployment + /// (`_asearch_with_fallbacks`, whose helper picks among the router's search tools). + pub unrouted: bool, + /// `passthrough_on_no_deployment`: a pick that fails hands the attempt to the host under the + /// requested name, inside the level's retries and fallbacks, with no outcome recorded for a + /// deployment (`_ageneric_api_call_with_fallbacks_helper`). + pub passthrough_on_no_deployment: bool, + /// The endpoint's attempt stamps no `failed_deployment_id` on its error (file and batch + /// creation, whose result belongs to the deployment that served it), so weighted failover + /// never runs for it. + pub skip_weighted_failover: bool, + /// The request creates or names a resource only one provider holds (a file upload, a batch + /// on an uploaded file): it falls back within its own model group only + /// (`run_async_fallback`'s `same_model_group_only`). + pub same_group_fallbacks_only: bool, + /// Fallback lists the request passed (`Some(None)` when it passed `None`), which replace + /// the router's. + pub fallbacks: Option>, + pub context_window_fallbacks: Option>, + pub content_policy_fallbacks: Option>, + /// A per-request `max_fallbacks`, which replaces the router's. + pub max_fallbacks: Option, + /// A per-request `cooldown_time`: the call's own params carry it, so it cools a failed + /// deployment down for that long unless its `model_info` sets one. + pub cooldown_time: Option, + /// A `mock_testing_*` flag the request set. + pub mock_failure: Option, + /// A per-request `model_group_retry_policy`, which replaces the router's. + pub model_group_retry_policy: Option>, + /// The model groups whose `model_group_retry_policy` lookup raises (the policy is no mapping + /// of policies, or the group's entry is not one): the first failure of a hop serving one ends + /// the hop with the host's error instead ([`RouterEvent::RetryPolicyRaised`]), no retry. + pub retry_policy_raises: Vec, + /// `health_check_probe`: a health check's own pick, which sees no deployment when every + /// one is marked unhealthy (the safety net keeps them all for any other request). + pub health_check_probe: bool, + /// What the host needs to run the call itself, such as the Python kwargs. + pub payload: Option, +} + +/// A per-request `num_retries` that is no count of retries. `async_function_with_retries` only +/// reads it once a hop's first failure was not taken over by a retry policy: a call that succeeds +/// first time never does. +#[derive(Clone, Copy, Debug, PartialEq, Eq, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum UncountedRetries { + /// A negative int: `num_retries > 0` is false, so the failure stands, and a stream's plain + /// budget is spent (`isinstance(num_retries, int)`). + Negative, + /// Any other value `num_retries > 0` is false for (a negative float, a NaN): the failure + /// stands, and a stream falls through to the deployment's budget, then the router's. + NotPositive, + /// `num_retries > 0` or `range(num_retries)` raises (a string, a positive float): the hop ends + /// with that error ([`RouterEvent::RetryCountRaised`]). + Raises, +} + +/// One `os.environ/` reference the deployment build found ([`crate::build`]). +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct SecretRef { + /// The value as written, `os.environ/`. + pub name: String, + /// Where the value sits in the deployment. + pub at: SecretAt, +} + +/// Where a [`SecretRef`] sits in a deployment. +#[derive(Clone, Debug, PartialEq, Eq)] +pub enum SecretAt { + /// A top-level `litellm_params` value: resolved through the host's secret manager. + Param(String), + /// `litellm_params.dataSources[source].parameters[key]` (`endpoint` or `key`): read from the + /// environment alone, empty when unset. + DataSource { source: usize, key: String }, +} + +/// Which of the router's adds a build runs. +#[derive(Clone, Copy, Debug, PartialEq, Eq, strum::IntoStaticStr, strum::EnumString)] +#[strum(serialize_all = "snake_case")] +pub enum BuildPath { + /// `set_model_list` (and the constructor's list): config entries, one deployment per + /// organization, ignored when invalid if the router ignores invalid deployments. + Create, + /// `add_deployment` (and upserts): one validated deployment. + Add, +} + +/// What the build asks the host about one deployment. +#[derive(Clone, Debug)] +pub struct BuildAsk<'a> { + pub path: BuildPath, + /// The entry's position in the list the build was given. + pub index: usize, + /// The entry as the build holds it: `model_name`, `litellm_params` (secrets resolved, the + /// organization of this deployment set), `model_info` (with its `id`) and any other keys. + pub entry: &'a Map, + /// The `model_info.id` the entry declared, before the build assigned one. + pub declared_id: Option<&'a str>, + /// Another entry of the same list declares the same id. + pub duplicate_id: bool, +} + +/// The host's answer about one deployment. +#[derive(Clone, Debug, PartialEq)] +pub enum Built { + Entry(Map), + /// Left out without an error (`model_info.supported_environments` excludes this one, or an + /// `add_deployment` id that already exists). + Skip, + /// The host failed the deployment and the router ignores invalid deployments. + Invalid, +} + +/// The host's provider answer. +#[derive(Clone, Debug, PartialEq, Eq)] +pub enum Provided { + Provider(Option), + Invalid, +} + +/// What the model-info merge ([`crate::model_info`]) asks the host while it answers one query. +/// Asked at query time, so a price reload or a registration the host made since the router was +/// built is what the next query reads. +#[derive(Clone, Debug, PartialEq, serde::Serialize)] +#[serde(tag = "lookup", rename_all = "snake_case")] +pub enum ModelInfoLookup { + /// The process cost map's row for `model` (`get_model_info(model=...)`), the host choosing + /// the keys it tries. Answered with the row, or `None`. A lookup that raises propagates, + /// unless `tolerate` is set: then the host answers `{"failed": True}`. `exact` also asks + /// whether the row came from a fallback-generalization rule (`{"row": ..., "generalized": + /// bool}`). + ModelInfo { + model: Value, + tolerate: bool, + exact: bool, + }, + /// The raw cost-map entry under `key` (`model_cost.get(key)`): answered with + /// `{"entry": value}`, or `None` when the key is absent. `process` reads the process global + /// (`litellm.model_cost`) rather than the host's live settings, where the Python router does. + CostEntry { key: String, process: bool }, + /// `get_llm_provider` for the queried deployment's `model` and validated params: answered + /// with `[model, provider]`. + Provider { model: String }, + /// The deployments a wildcard pattern routes `name` to: answered with `[{id, model}]`, or + /// `None` when no pattern matches. + PatternModels { name: String }, + /// The base model a bedrock model id names (`get_bedrock_base_model`). + BedrockBase { model_id: String }, + /// The limits the background model-info refresh discovered for deployment `id`. + Discovered { id: Value }, + /// The queried deployment's own `model_info`, read when the merge reaches it (a malformed one + /// raises there). + UserModelInfo, + /// A log record, emitted where the Python router logs it. + Log { level: LogLevel, message: String }, + /// Group info: deployment `index`'s validated params (`LiteLLM_Params(**litellm_params)`, which + /// raises on bad params): answered with `{"model", "configurable_clientside_auth_params"}`. + GroupParams { index: usize }, + /// Group info: what the host knows of deployment `index` serving `model`: answered with + /// `{"model_info"}` (its merged model info, `None` when that fails), `"mapped"` + /// (`deployment_is_catalog_mapped`), `"litellm_model"` and `"llm_provider"`, and, when it has + /// no model info, `"supported_openai_params"`. + GroupDeployment { index: usize, model: Value }, + /// Group info: deployment `index`'s capabilities under `model_info`: answered with + /// `{"supported_endpoints"}` (empty when the team may not use it), `"supports_fast_mode"` and + /// `"reasoning_efforts"` (`resolve_supported_reasoning_efforts`, or `None`). + GroupCapabilities { + index: usize, + model_info: Value, + litellm_model: Value, + llm_provider: Value, + mapped: bool, + }, + /// Group info: the group's first deployment as the host builds it (`ModelGroupInfo(**info)`, + /// whose validation may raise): answered with the fields the merge compares, as validated. + GroupSeed { info: Map }, + /// `intersect_supported_reasoning_efforts(current, resolved)`, in the host's advertisement order. + ReasoningIntersect { current: Value, resolved: Value }, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq, serde::Serialize)] +#[serde(rename_all = "snake_case")] +pub enum LogLevel { + Debug, + Error, +} + +/// One attempt against one deployment. +#[derive(Clone, Debug)] +pub struct Attempt { + pub op: Op, + pub request: Arc, + pub deployment: Arc, + /// The name this attempt resolved: the requested name, or the fallback target. + pub name: String, + /// The model group being served, after fallbacks. + pub model_group: String, + /// The `order` level an order fallback targets in the same group. + pub target_order: Option, + /// Every deployment `name` resolved to for this caller, the picked one included. + pub candidate_ids: Vec, + /// 0 for the first attempt in this group, then the retry number. + pub retry: u32, + pub fallback_depth: u32, + /// How many fallbacks deep the chain is: 0 for the requested name. + pub fallback_level: u32, + pub timeout: Option, + /// Request params the fallback chain's dict targets merged in (`kwargs.update(target)`). + pub params: Option>>, + /// The chain level whose request metadata this attempt carries (see [`RouterEvent::Fallback`]). + pub scope: u32, + /// The retry budget the level's loop stamps on the request before this attempt. + pub max_retries: u32, + /// How many deployments `get_model_list` lists for the level's group (`model_group_size`). + pub model_group_size: usize, +} + +/// What a successful attempt returns. +#[derive(Clone, Debug)] +pub struct Output { + pub value: Payload, + /// Prompt plus completion tokens, counted against the deployment's per-minute usage. + pub total_tokens: u64, + /// A response the provider blocked, which the router may raise to fall back. + pub blocked: Option, + /// A stream the host reports the outcome of when it ends ([`crate::Router::report_success`], + /// [`crate::Router::report_failure`]), or a realtime session, whose success the host counts + /// only when it ended cleanly: the attempt counts its request and `total_tokens` now, and no + /// success. + pub deferred: bool, + /// Whether the attempt counts against the deployment's per-minute usage, when the host knows + /// better than the deployment's own limits ([`Deployment::tracks_usage`]): a limit only the + /// model catalog sets (`_deployment_tracks_usage`), or a realtime session, whose `None` + /// response counts nothing (`increment_deployment_usage_for_response`). + pub tracks_usage: Option, +} + +/// How a provider blocked a response the call otherwise served. +#[derive(Clone, Copy, Debug, PartialEq, Eq, strum::IntoStaticStr, strum::EnumString)] +#[strum(serialize_all = "snake_case")] +pub enum Blocked { + /// A chat response whose first choice was content-filtered (or that has no choices). + ContentFilter, + /// An Anthropic /v1/messages safeguard refusal carrying `stop_details`. + Refusal, +} + +/// What a plugin or strategy sees: never the router itself. +pub struct RoutingContext<'a> { + pub registry: &'a Registry, + pub settings: &'a Settings, + pub health: &'a Health, + pub request: &'a RouterRequest, + pub op: &'a Op, + pub model_group: &'a str, + /// The name the request resolved from (Python's `model`), which strategies key state by. + pub name: &'a str, + /// Whose selector a strategy stage reads. + pub scope: StrategyScope, + /// The deployments the name resolved to for this caller, before health filters + /// (`request_deployments`), for the tag stage. + pub pool: &'a [Arc], + /// `_excluded_deployment_ids` a weighted failover hop put on the request, which the callback + /// filters read (`deployment_ids_failed_in_request`). + pub excluded_ids: &'a [String], + /// `_retry_skipped_deployment_ids` the retry loop put on the request, as `excluded_ids`. + pub skipped_ids: &'a [String], + /// `_target_order` as the filters read it: the level an order fallback targets, else the one + /// the request named. + pub target_order: Option, +} + +/// Which selection step the host is asked to run. +#[derive(Clone, Debug, PartialEq)] +pub enum SelectStage { + /// Pre-call checks and sub-routers registered as plugins: return the deployments that remain. + Plugins, + /// A strategy not yet ported: return the candidates with the pick first. + Strategy(RoutingStrategy), + /// Tag routing (`get_deployments_for_tag`), not yet ported: return the deployments that remain. + Tags, + /// `_deployment_declares_no_vision`: return the candidates that declare no vision support. + TextOnly, + /// `_deployment_serves_route`: return the candidates whose provider cannot serve the call + /// `handler` names (a Decisions call on a provider with no Decisions config). + RouteIncapable { handler: String }, + /// `_decisions_deployment_refuses_media`: return the candidates that refuse the audio or + /// videos a Decisions call carries. + MediaRefusing, +} + +/// Routing events, delivered through [`Host::observe`]. +#[derive(Clone, Debug, PartialEq)] +pub enum RouterEvent { + Selected { + model_group: String, + deployment_id: String, + /// How long the pick took (`_acompletion`'s `async_get_available_deployment` service event). + pick: Duration, + }, + Succeeded { + deployment_id: String, + }, + AttemptFailed { + deployment_id: String, + kind: FailureKind, + status: Option, + }, + CooldownSet { + deployment_id: String, + cooldown: Duration, + }, + Retry { + model_group: String, + retry: u32, + delay: Duration, + }, + Fallback { + from: String, + to: String, + /// The order level of `to` a same-group hop walks to. + order: Option, + depth: u32, + /// The fallback level `to` is tried at: 1 for the request's own fallbacks. + level: u32, + /// The new level's request metadata scope, copied from `parent`'s at the hop. + scope: u32, + parent: u32, + /// The `max_fallbacks` the chain reads, stamped on the hop's request. + max_fallbacks: u32, + /// The non-string `model` of the dict target this hop routes + /// ([`crate::fallback::FallbackTarget::unroutable_model`]): the hop fails on it before any + /// attempt, and the host raises what Python's level raises for that value. + unroutable: Option, + }, + /// A hop of the fallback chain returned: `log_success_fallback_event` or + /// `log_failure_fallback_event` for the level at `parent`. + FallbackDone { + from: String, + to: String, + level: u32, + scope: u32, + parent: u32, + succeeded: bool, + }, + /// The router raised a `ContentPolicyViolationError` for a blocked response, since a + /// fallback would serve it (`_should_raise_content_policy_error`, + /// `_should_raise_anthropic_refusal_error`). + Blocked { + level: u32, + model: String, + blocked: Blocked, + }, + /// `log_retry`: the router recorded a failure in `scope`'s request metadata + /// (`previous_models`, `request_retry_count`) before retrying or falling back. + RetryLogged { + scope: u32, + level: u32, + /// The group the failed level served (`kwargs["model"]`). + model: String, + /// The retry number stamped before the failed try; `None` keeps the scope's own. + attempted_retries: Option, + cause: RetryCause, + }, + /// `async_function_with_retries` opened a level's retry loop, before any pick: `scope`'s + /// request metadata tracks no retry yet, `max_retries` of them, and the group's size. + RetriesOpened { + scope: u32, + max_retries: u32, + model_group_size: usize, + }, + /// The level's first failure settled its overrides: the budget `scope` then tracks. + RetryBudget { + scope: u32, + max_retries: u32, + }, + /// A pick of the async selection or the sync one failed on `error`: the async one's `except` + /// logs every refusal to the request's logging object (`_async_get_available_deployment`). + PickFailed { + level: u32, + model: String, + error: PickError, + }, + /// A level of the fallback chain serving `model` ended on an error a pick raised (no + /// deployment for the caller, none available, a stale routing group, ...): Python raises + /// that error for the level, whatever its earlier attempts raised. + LevelFailed { + level: u32, + model: String, + error: PickError, + /// Raised by the level's fallback handling, while the level's own error was handled + /// (so it chains to that error), rather than by the level's retry loop. + handling: bool, + }, + /// The router raised a `mock_testing_*` failure for `model` instead of calling it. + MockFailed { + mock: MockFailure, + level: u32, + model: String, + num_retries: Option, + }, + /// The first failure of a hop serving `model` looked up a `model_group_retry_policy` that + /// raises for it: the hop ends with the error the lookup raised (the host raises it, from the + /// failure `cause` names), without retrying, and falls back as from any plain exception. + RetryPolicyRaised { + level: u32, + model: String, + cause: RetryCause, + }, + /// The first failure of a hop serving `model` read a request `num_retries` that raises + /// ([`UncountedRetries::Raises`]): the hop ends with that error (the host raises it, from the + /// failure `cause` names), and falls back as from any plain exception. + RetryCountRaised { + level: u32, + model: String, + cause: RetryCause, + }, + /// The level serving `model` reached the fallback handling that appends the context window + /// and content policy hints to its error, before any hop runs + /// (`async_function_with_fallbacks_common_utils`). The order path never gets there. + FallingBack { + level: u32, + model: String, + }, + /// The order fallback chain of the level at `level` served nothing, so the level raises from + /// `async_function_with_fallbacks_common_utils`' order path, with no outcome message: what the + /// chain's last hop raised when `last_hop`, else its own error (a context window error, or a + /// chain no hop of which ran). Weighted failover raising what it does not catch + /// (`_maybe_run_weighted_failover`) ends the level the same way. + OrderChainFailed { + level: u32, + last_hop: bool, + }, + /// The regular fallback handling of the level at `level` served nothing, where + /// `async_function_with_fallbacks_common_utils` reaches its `except` and logs the failure. + /// `chain` when the level had targets to walk: the error it logs is what the last hop raised + /// when `last_hop`, else its own. With no `chain`, the host decides: only the host holds + /// whether a level with no match had a fallback list at all, which Python does not log. + /// The order path and weighted failover are outside that `try` and never report this. + FallbacksFailed { + level: u32, + chain: bool, + last_hop: bool, + }, + /// A group's retry loop ran every retry and the last one failed: Python stamps + /// `num_retries` and `max_retries` on that error. + RetriesExhausted { + model_group: String, + attempted: u32, + max: u32, + /// The level the retry loop served. + level: u32, + /// The last failure was a pick's error, which the level then raises + /// ([`RouterEvent::LevelFailed`]), rather than an attempt's. + picked: bool, + }, +} + +/// A pick's error, as the events that name it carry it. +#[derive(Clone, Debug)] +pub struct PickError(pub Arc); + +impl PickError { + /// The copy of `error` events carry, when a pick raised it. + pub fn of(error: &Error) -> Option { + error.pick_copy().map(|copy| Self(Arc::new(copy))) + } +} + +impl PartialEq for PickError { + fn eq(&self, other: &Self) -> bool { + format!("{:?}", self.0) == format!("{:?}", other.0) + } +} + +/// Which error a [`RouterEvent::RetryLogged`] records. +#[derive(Clone, Debug, PartialEq)] +pub enum RetryCause { + /// The last attempt's failure. + Attempt, + /// The error the level ended on (`run_async_fallback`'s `original_exception`). + Level, + /// A pick raised this error: the group was unknown or had none available, the stale + /// routing-group member map named a group no longer configured, ... + Pick(PickError), + /// Any other error no attempt raised. + Other, +} + +impl RetryCause { + /// The cause that names `error`. + pub fn of(error: &Error) -> Self { + match error { + Error::Attempt(_) => Self::Attempt, + _ => PickError::of(error).map_or(Self::Other, Self::Pick), + } + } +} + +#[async_trait] +pub trait Host: Send + Sync { + /// Run one attempt against the picked deployment. The core records its outcome (cooldowns, + /// allowed-fails counters, usage) itself; the host only runs it. + async fn invoke(&self, attempt: &Attempt) -> Result; + + /// Run a selection stage the core does not implement. + async fn select( + &self, + stage: &SelectStage, + ctx: &RoutingContext<'_>, + candidates: Vec>, + ) -> Result>, Error>; + + /// Pre-routing before a pick resolves `ask.name`: the model group to resolve instead, a + /// custom strategy's deployment to serve, or `None` to resolve the name itself. Asked only + /// for picks that pre-route (`pre_routes`). + async fn pre_route(&self, ask: &PreRoute<'_>) -> Result, Error> { + let _ = ask; + Ok(None) + } + + /// Access and budget checks before a cross-group fallback target. + /// `Err` when a check raised: the chain ends there, outside any target's `try`. + async fn allow_fallback( + &self, + request: &RouterRequest, + target: &str, + ) -> Result; + + /// Whether `run_async_fallback` refuses a chain for one of its dict `targets`, a server-owned + /// federation field or an `oidc/` api key. The chain is walked by `level`. + async fn refuses_fallbacks(&self, _level: u32, _targets: &[Map]) -> bool { + false + } + + async fn sleep(&self, duration: Duration); + + fn observe(&self, event: &RouterEvent); + + /// The shared store routing state lives in. + fn cache(&self) -> Arc; + + /// Whether [`Self::invoke`] can run `op` on `provider`. Checked for every served op and + /// deployment when the router is built or changed, so an unsupported pair fails the config, + /// never a request. `PythonHost` runs every op. + fn supports(&self, _op: &Op, _provider: &str) -> bool { + true + } + + /// Whether [`Self::select`] can run strategies and plugins. + fn can_select(&self) -> bool { + true + } + + /// Run the routing plugins for `ctx.name` when `run` (the async pick's pre-routing hook), and + /// report the candidate models they left on the request, if any. + async fn routing_plugin_models( + &self, + _ctx: &RoutingContext<'_>, + _run: bool, + ) -> Result>, Error> { + Ok(None) + } + + /// The facts the `enable_pre_call_checks` filter decides by (token count, context windows, + /// param support), for `candidates` and the deployments in `ctx.pool`. A host without them + /// reports no input, so nothing is checked. + async fn pre_call_facts( + &self, + _ctx: &RoutingContext<'_>, + _candidates: &[Arc], + ) -> Result { + Ok(PreCallFacts::default()) + } + + /// `f"{provider}/{model}"` for `ctx.name`, which `PatternMatchRouter.get_pattern` tries when + /// no pattern matches the name itself: only the host can infer a provider. `None` from a + /// host that infers none. + async fn provider_request(&self, _ctx: &RoutingContext<'_>) -> Result, Error> { + Ok(None) + } + + /// `async_callback_filter_deployments`: the host's routing callbacks narrow `candidates` in + /// turn. A filter that raises fails the try with [`Error::Attempt`], which the call retries and + /// falls back from like any attempt failure. A host without filters keeps every candidate. + async fn callback_filters( + &self, + _ctx: &RoutingContext<'_>, + candidates: Vec>, + ) -> Result>, Error> { + Ok(candidates) + } + /// The values of `refs`, one per ref, in order: what the deployment build substitutes for + /// each `os.environ/` reference it found in entry `index` ([`crate::build`]). Asked once per + /// deployment, as the build reaches it. A host without secrets leaves every reference as + /// written. + async fn secrets(&self, index: usize, refs: &[SecretRef]) -> Result, Error> { + let _ = index; + Ok(refs.iter().map(|r| Value::String(r.name.clone())).collect()) + } + + /// Check one deployment before it is added, and answer the entry as the host normalized it, + /// a skip (the deployment is left out), or [`Built::Invalid`] (it failed and the router + /// ignores invalid deployments). A host without checks accepts the entry as given. + async fn validate(&self, ask: &BuildAsk<'_>) -> Result { + Ok(Built::Entry(ask.entry.clone())) + } + + /// The provider that serves one validated deployment (`get_llm_provider`), `None` for one + /// served without a provider (a prompt-management model), or [`Provided::Invalid`]. A host + /// without provider resolution answers `None`, so [`Deployment::provider`] reads the model. + async fn provider(&self, ask: &BuildAsk<'_>) -> Result { + let _ = ask; + Ok(Provided::Provider(None)) + } + + /// Run the host's side effects for an accepted deployment and answer the entry as the + /// registry stores it. A host without side effects stores the entry it validated. + async fn added(&self, ask: &BuildAsk<'_>, provider: Option<&str>) -> Result { + let _ = provider; + Ok(Built::Entry(ask.entry.clone())) + } + + /// Answer one model-info lookup ([`ModelInfoLookup`]). A host without a cost map answers + /// `None` to everything, so no deployment has model info. + async fn model_info(&self, lookup: &ModelInfoLookup) -> Result { + let _ = lookup; + Ok(Value::Null) + } +} diff --git a/litellm-rust/crates/router-core/src/host_index.rs b/litellm-rust/crates/router-core/src/host_index.rs new file mode 100644 index 00000000000..3748f7b4d19 --- /dev/null +++ b/litellm-rust/crates/router-core/src/host_index.rs @@ -0,0 +1,727 @@ +//! [`HostIndexes`]: the registry indexes the Python router hands out, kept the way it keeps them. +//! +//! The Python router keeps its indexes by history rather than deriving them from the deployment +//! list: `model_names` is the set of names `set_model_list` was given (invalid entries included), +//! `deployment_names` only ever grows, assigning `model_list` reindexes nothing, a removal shifts +//! the positions every index holds (a caller-written one included), and the wildcard pattern +//! routers keep the copy of a deployment they were given. The registry keeps one of these per +//! version and applies those rules on its own writes; the host reads them back as the live +//! objects `registry().index()`, `names()` and `patterns()` return, and hands a caller's write to +//! any of them back as a changeset ([`IndexWrite`]). +//! +//! Values are JSON so whatever a caller wrote reads back as written, junk included. + +use serde::{Deserialize, Serialize}; +use serde_json::{Map, Value}; + +/// One wildcard pattern router (`PatternMatchRouter`): regex to the deployments it routes to, +/// most specific pattern first. +#[derive(Clone, Debug, Default, PartialEq, Serialize, Deserialize)] +#[serde(transparent)] +pub struct PatternRouter { + pub patterns: Vec<(String, Vec)>, +} + +/// The indexes, as the host reads and writes them. +#[derive(Clone, Debug, Default, PartialEq, Serialize, Deserialize)] +#[serde(default)] +pub struct HostIndexes { + /// `model_id_to_deployment_index_map`: id to position, in insertion order. + pub id_map: Map, + /// `model_name_to_deployment_indices`: name to its positions. + pub model_name: Map, + /// `team_model_to_deployment_indices`: `[team_id, public name, positions]`. + pub team_model: Vec<(String, String, Value)>, + /// `model_names`: a set, its members in the order they joined. + pub model_names: Vec, + pub deployment_names: Vec, + pub team_public_model_names: Vec, + pub provider_default_ids: Vec, + /// `pattern_router`. + pub patterns: PatternRouter, + /// `team_pattern_routers`: team id to its router. + pub team_patterns: Vec<(String, PatternRouter)>, +} + +/// A host write to the indexes: each field it names replaces the one held. +#[derive(Clone, Debug, Default, PartialEq, Deserialize)] +#[serde(default)] +pub struct IndexWrite { + pub id_map: Option>, + pub model_name: Option>, + pub team_model: Option>, + pub model_names: Option>, + pub deployment_names: Option>, + pub team_public_model_names: Option>, + pub provider_default_ids: Option>, + pub patterns: Option, + pub team_patterns: Option>, +} + +/// A removal the Python router could not finish: an index held a position it cannot compare +/// (`_without_indices` raised `TypeError`). The indexes keep what it shifted before that. +#[derive(Clone, Debug, PartialEq, Eq, thiserror::Error)] +#[error("'>' not supported between instances of '{found}' and 'int'")] +pub struct RemovalError { + pub found: &'static str, +} + +impl HostIndexes { + /// Apply a host write. + pub fn write(&mut self, write: IndexWrite) { + let IndexWrite { + id_map, + model_name, + team_model, + model_names, + deployment_names, + team_public_model_names, + provider_default_ids, + patterns, + team_patterns, + } = write; + if let Some(value) = id_map { + self.id_map = value; + } + if let Some(value) = model_name { + self.model_name = value; + } + if let Some(value) = team_model { + self.team_model = value; + } + if let Some(value) = model_names { + self.model_names = unique(value); + } + if let Some(value) = deployment_names { + self.deployment_names = value; + } + if let Some(value) = team_public_model_names { + self.team_public_model_names = value; + } + if let Some(value) = provider_default_ids { + self.provider_default_ids = value; + } + if let Some(value) = patterns { + self.patterns = value; + } + if let Some(value) = team_patterns { + self.team_patterns = value; + } + } + + /// `set_model_list` starting over: the id, name and team indexes and the team routers go; + /// the deployment names, the wildcard router and the provider default ids stay. + pub fn reset(&mut self) { + self.id_map.clear(); + self.model_name.clear(); + self.team_model.clear(); + self.team_patterns.clear(); + self.team_public_model_names.clear(); + } + + /// `set_model_list`'s `model_names`: the names it was given. + pub fn set_model_names(&mut self, names: Vec) { + self.model_names = unique(names.into_iter().map(Value::String).collect()); + } + + /// A deployment joined the list at `position` (`_add_deployment` then + /// `_add_model_to_list_and_index_map`); `add` is `add_deployment`, which also names its group. + pub fn added(&mut self, entry: &Map, position: usize, add: bool) { + let params = entry.get("litellm_params").and_then(Value::as_object); + let info = entry.get("model_info").and_then(Value::as_object); + let model_name = entry.get("model_name").and_then(Value::as_str); + let id = info + .and_then(|info| info.get("id")) + .filter(|id| !id.is_null()); + self.deployment_names.push( + params + .and_then(|params| params.get("model")) + .cloned() + .unwrap_or(Value::Null), + ); + if let Some(name) = model_name.filter(|name| name.contains('*')) { + self.patterns.add(name, Value::Object(entry.clone())); + if let Some(id) = id.and_then(Value::as_str).filter(|id| !id.is_empty()) { + self.provider_default_ids.push(id.to_owned()); + } + } + let team_id = info + .and_then(|info| info.get("team_id")) + .filter(|v| !v.is_null()); + let public = info + .and_then(|info| info.get("team_public_model_name")) + .filter(|v| !v.is_null()); + if let (Some(team_id), Some(Value::String(public))) = (team_id, public) + && public.contains('*') + { + let team = python_str(team_id); + let router = match self.team_patterns.iter_mut().find(|(t, _)| *t == team) { + Some((_, router)) => router, + None => { + self.team_patterns.push((team, PatternRouter::default())); + &mut self.team_patterns.last_mut().expect("just pushed").1 + } + }; + router.add(public, Value::Object(entry.clone())); + } + let position = Value::from(position); + if let Some(id) = id { + self.id_map.insert(python_str(id), position.clone()); + } + if let Some(name) = model_name.filter(|name| !name.is_empty()) { + push_position( + self.model_name + .entry(name.to_owned()) + .or_insert_with(|| Value::Array(Vec::new())), + &position, + ); + } + if add { + let name = Value::String(model_name.unwrap_or_default().to_owned()); + if !self.model_names.contains(&name) { + self.model_names.push(name); + } + } + // `_update_team_model_index`: a truthy team id and public name. + if let (Some(team_id), Some(public)) = (team_id, public) + && truthy(team_id) + && truthy(public) + { + let (team, public) = (python_str(team_id), python_str(public)); + if !self.team_public_model_names.contains(&public) { + self.team_public_model_names.push(public.clone()); + } + let held = match self + .team_model + .iter_mut() + .find(|(t, p, _)| *t == team && *p == public) + { + Some((_, _, held)) => held, + None => { + self.team_model + .push((team, public, Value::Array(Vec::new()))); + &mut self.team_model.last_mut().expect("just pushed").2 + } + }; + if let Value::Array(positions) = held + && !positions.contains(&position) + { + positions.push(position); + } + } + } + + /// The deployment `id` left the list from `position` (`_update_deployment_indices_after_removal`). + pub fn removed(&mut self, id: &str, position: usize) -> Result<(), RemovalError> { + let mut id_map = Map::new(); + for (key, held) in std::mem::take(&mut self.id_map) { + if key == id { + continue; + } + let shifted = shift(&held, position)?.unwrap_or(held); + id_map.insert(key, shifted); + } + self.id_map = id_map; + let names: Vec = self.model_name.keys().cloned().collect(); + for name in names { + let kept = without(&self.model_name[&name], position)?; + if kept.is_empty() { + self.model_name.remove(&name); + self.model_names + .retain(|held| held.as_str() != Some(name.as_str())); + } else { + self.model_name.insert(name, Value::Array(kept)); + } + } + let mut team_model = Vec::new(); + let held = std::mem::take(&mut self.team_model); + let mut rest = held.into_iter(); + for (team, public, positions) in rest.by_ref() { + match without(&positions, position) { + Ok(kept) if kept.is_empty() => {} + Ok(kept) => team_model.push((team, public, Value::Array(kept))), + Err(error) => { + team_model.push((team, public, positions)); + team_model.extend(rest); + self.team_model = team_model; + return Err(error); + } + } + } + self.team_model = team_model; + let mut publics: Vec = Vec::new(); + for (_, public, _) in &self.team_model { + if !publics.contains(public) { + publics.push(public.clone()); + } + } + self.team_public_model_names = publics; + self.patterns.remove(id); + for (_, router) in &mut self.team_patterns { + router.remove(id); + } + self.team_patterns + .retain(|(_, router)| !router.patterns.is_empty()); + Ok(()) + } + + /// Whether the id map holds `id` (`_has_model_id`). + pub fn has_id(&self, id: &str) -> bool { + self.id_map.contains_key(id) + } + + /// What the id map holds for `id`, as written (`model_id_to_deployment_index_map.get(id)`). + pub fn id_entry(&self, id: &str) -> Option<&Value> { + self.id_map.get(id) + } + + /// What the name index holds for `name`, as written. + pub fn named(&self, name: &str) -> Option<&Value> { + self.model_name.get(name) + } + + /// What the team index holds for `(team, public)`, as written. + pub fn team_named(&self, team: &str, public: &str) -> Option<&Value> { + self.team_model + .iter() + .find(|(t, p, _)| t == team && p == public) + .map(|(_, _, positions)| positions) + } + + /// `team_public_model_names_of`: the public names the team index holds for `team`. + pub fn team_public_names_of(&self, team: &str) -> Vec { + let mut names: Vec = Vec::new(); + for (t, public, _) in &self.team_model { + if t == team && !names.contains(public) { + names.push(public.clone()); + } + } + names + } + + /// The position the id map holds for `id`, when it is a position. + pub fn position_of(&self, id: &str) -> Option { + self.id_map + .get(id) + .and_then(Value::as_u64) + .and_then(|position| usize::try_from(position).ok()) + } +} + +impl PatternRouter { + /// `add_pattern`: a new regex joins in specificity order (a stable sort, most specific + /// first); a known one gains the deployment. + pub fn add(&mut self, pattern: &str, deployment: Value) { + let regex = pattern_to_regex(pattern); + if let Some((_, deployments)) = self.patterns.iter_mut().find(|(held, _)| *held == regex) { + deployments.push(deployment); + return; + } + self.patterns.push((regex, vec![deployment])); + self.patterns + .sort_by_key(|(regex, _)| std::cmp::Reverse(regex_specificity(regex))); + } + + /// `remove_deployment`: every deployment with this id goes, and so does a pattern left empty. + pub fn remove(&mut self, id: &str) { + for (_, deployments) in &mut self.patterns { + deployments.retain(|deployment| { + deployment + .get("model_info") + .and_then(|info| info.get("id")) + .and_then(Value::as_str) + != Some(id) + }); + } + self.patterns + .retain(|(_, deployments)| !deployments.is_empty()); + } + + /// `route`: the deployments of the first pattern `request` matches, each a copy whose + /// `litellm_params.model` takes the parts the wildcards matched; `filtered` keeps to the + /// patterns of those names. + pub fn route(&self, request: &str, filtered: Option<&[String]>) -> Option> { + let filtered: Option> = + filtered.map(|names| names.iter().map(|name| pattern_to_regex(name)).collect()); + for (regex, deployments) in &self.patterns { + if filtered + .as_ref() + .is_some_and(|names| !names.contains(regex)) + { + continue; + } + let Some(groups) = regex_match(regex, request) else { + continue; + }; + let mut routed = Vec::with_capacity(deployments.len()); + for deployment in deployments { + let mut copy = deployment.clone(); + let model = deployment + .get("litellm_params") + .and_then(|params| params.get("model")); + let Some(Value::String(model)) = model else { + // `deployment["litellm_params"]["model"]` raising ends the route with None. + return None; + }; + let named = deployment_model_name(&groups, request, model); + if model.contains('*') + && let Some(params) = + deployment.get("litellm_params").and_then(Value::as_object) + && let Some(info) = copy.get_mut("model_info").and_then(Value::as_object_mut) + { + for flag in cache_flags_left_to_the_catalog(params) { + info.remove(*flag); + } + } + if let Some(params) = copy + .get_mut("litellm_params") + .and_then(Value::as_object_mut) + { + params.insert("model".to_owned(), Value::String(named)); + } + routed.push(copy); + } + return Some(routed); + } + None + } + + /// `get_models_by_pattern`'s pick: the first (request, pattern) pair that matches, trying + /// `model` against every pattern, then `provider/model`. + pub fn models_by_pattern(&self, requests: &[String]) -> Option> { + for request in requests { + for (regex, deployments) in &self.patterns { + if let Some(groups) = regex_match(regex, request) { + return Some( + deployments + .iter() + .filter_map(|deployment| { + deployment + .get("litellm_params") + .and_then(|params| params.get("model")) + .and_then(Value::as_str) + .map(|model| deployment_model_name(&groups, request, model)) + }) + .collect(), + ); + } + } + } + None + } +} + +/// `_cache_flags_left_to_the_catalog`: the `model_info` flags a wildcard deployment does not speak +/// for on the models it matches. An azure_ai wildcard's `supports_prompt_cache_breakpoint` does. +/// `inferred_provider` names `azure_ai` only for a model carrying that prefix. +pub(crate) fn cache_flags_left_to_the_catalog( + params: &Map, +) -> &'static [&'static str] { + let provider = params + .get("custom_llm_provider") + .filter(|provider| truthy(provider)) + .and_then(Value::as_str) + .or_else(|| { + params + .get("model") + .and_then(Value::as_str) + .and_then(|model| model.split_once('/')) + .map(|(prefix, _)| prefix) + }); + if provider == Some("azure_ai") { + &["supports_prompt_caching"] + } else { + &[ + "supports_prompt_cache_breakpoint", + "supports_prompt_caching", + ] + } +} + +/// `re.escape(pattern).replace(r"\*", "(.*)")`. +pub fn pattern_to_regex(pattern: &str) -> String { + let mut regex = String::new(); + for c in pattern.chars() { + match c { + '*' => regex.push_str("(.*)"), + '(' | ')' | '[' | ']' | '{' | '}' | '?' | '+' | '-' | '|' | '^' | '$' | '\\' | '.' + | '&' | '~' | '#' | ' ' | '\t' | '\n' | '\r' | '\x0b' | '\x0c' => { + regex.push('\\'); + regex.push(c); + } + _ => regex.push(c), + } + } + regex +} + +/// `PatternUtils.calculate_pattern_specificity`: the regex's length, then how much regex syntax +/// it has. +pub fn regex_specificity(regex: &str) -> (usize, usize) { + let complexity = regex + .chars() + .filter(|c| matches!(c, '*' | '+' | '?' | '\\' | '^' | '$' | '|' | '(' | ')')) + .count(); + (regex.chars().count(), complexity) +} + +/// A pattern regex read back as its literal runs between wildcards, or None when it holds regex +/// syntax `pattern_to_regex` never writes (a caller wrote it). +fn literal_runs(regex: &str) -> Option> { + let mut runs = vec![String::new()]; + let mut chars = regex.chars().peekable(); + while let Some(c) = chars.next() { + match c { + '\\' => runs.last_mut()?.push(chars.next()?), + '(' => { + let rest: String = chars.by_ref().take(3).collect(); + if rest != ".*)" { + return None; + } + runs.push(String::new()); + } + '.' | '*' | '+' | '?' | '[' | ']' | '{' | '}' | '|' | '^' | '$' | ')' => return None, + other => runs.last_mut()?.push(other), + } + } + Some(runs) +} + +/// `re.match(regex, request)` for a pattern regex: anchored at the start only, each `(.*)` +/// greedy. The groups it captured, or None when it does not match. +pub fn regex_match(regex: &str, request: &str) -> Option> { + let runs = literal_runs(regex)?; + let rest = request.strip_prefix(runs[0].as_str())?; + let mut groups = Vec::new(); + greedy(&runs[1..], rest, &mut groups).then_some(groups) +} + +/// Match `runs` (each a literal that follows a group) against `rest`, longest group first. +fn greedy(runs: &[String], rest: &str, groups: &mut Vec) -> bool { + let Some((literal, later)) = runs.split_first() else { + return true; + }; + let mut ends: Vec = rest.char_indices().map(|(at, _)| at).collect(); + ends.push(rest.len()); + for &end in ends.iter().rev() { + let (group, after) = rest.split_at(end); + let Some(after) = after.strip_prefix(literal.as_str()) else { + continue; + }; + groups.push(group.to_owned()); + if greedy(later, after, groups) { + return true; + } + groups.pop(); + } + false +} + +/// `set_deployment_model_name`: the deployment's model with each wildcard replaced by the group +/// matched for it, the request itself when more groups matched than it has wildcards. +pub fn deployment_model_name(groups: &[String], request: &str, model: &str) -> String { + if !model.contains('*') { + return model.to_owned(); + } + if groups.len() > model.matches('*').count() { + return request.to_owned(); + } + let mut named = model.to_owned(); + for group in groups { + named = named.replacen('*', group, 1); + } + named +} + +fn unique(values: Vec) -> Vec { + let mut kept: Vec = Vec::with_capacity(values.len()); + for value in values { + if !kept.contains(&value) { + kept.push(value); + } + } + kept +} + +fn push_position(held: &mut Value, position: &Value) { + if let Value::Array(positions) = held { + positions.push(position.clone()); + } +} + +/// Python's truthiness for a JSON value. +pub(crate) fn truthy(value: &Value) -> bool { + match value { + Value::Null => false, + Value::Bool(flag) => *flag, + Value::Number(number) => number.as_f64() != Some(0.0), + Value::String(text) => !text.is_empty(), + Value::Array(items) => !items.is_empty(), + Value::Object(items) => !items.is_empty(), + } +} + +/// Python's `str()` for an id or team id. +pub(crate) fn python_str(value: &Value) -> String { + match value { + Value::String(text) => text.clone(), + Value::Bool(true) => "True".to_owned(), + Value::Bool(false) => "False".to_owned(), + Value::Null => "None".to_owned(), + other => other.to_string(), + } +} + +fn type_name(value: &Value) -> &'static str { + match value { + Value::Null => "NoneType", + Value::Bool(_) => "bool", + Value::Number(number) if number.is_f64() => "float", + Value::Number(_) => "int", + Value::String(_) => "str", + Value::Array(_) => "list", + Value::Object(_) => "dict", + } +} + +/// `idx - 1 if idx > removal_idx else idx`: None keeps `held` as is. +fn shift(held: &Value, position: usize) -> Result, RemovalError> { + let removal = position as f64; + let number = match held { + Value::Bool(flag) => f64::from(u8::from(*flag)), + Value::Number(number) => number.as_f64().unwrap_or_default(), + other => { + return Err(RemovalError { + found: type_name(other), + }); + } + }; + if number <= removal { + return Ok(None); + } + Ok(Some(match held { + Value::Number(n) if n.is_i64() || n.is_u64() => Value::from(n.as_i64().unwrap_or(0) - 1), + _ => Value::from(number - 1.0), + })) +} + +/// `_without_indices`: the positions without `position`, the later ones shifted down. +fn without(held: &Value, position: usize) -> Result, RemovalError> { + let Value::Array(positions) = held else { + return Ok(Vec::new()); + }; + let removal = Value::from(position); + let mut kept = Vec::with_capacity(positions.len()); + for index in positions { + if *index == removal { + continue; + } + kept.push(shift(index, position)?.unwrap_or_else(|| index.clone())); + } + Ok(kept) +} + +#[cfg(test)] +mod tests { + use rstest::rstest; + use serde_json::json; + + use super::*; + + fn entry(name: &str, id: &str, model: &str) -> Map { + serde_json::from_value(json!({ + "model_name": name, + "litellm_params": {"model": model}, + "model_info": {"id": id}, + })) + .unwrap() + } + + #[test] + fn removal_shifts_every_index_and_drops_emptied_names() { + let mut indexes = HostIndexes::default(); + indexes.added(&entry("gpt", "a", "openai/a"), 0, false); + indexes.added(&entry("gpt", "b", "openai/b"), 1, false); + indexes.added(&entry("claude", "c", "anthropic/c"), 2, false); + indexes.set_model_names(vec!["gpt".into(), "claude".into(), "broken".into()]); + indexes.removed("a", 0).unwrap(); + assert_eq!( + indexes.id_map, + json!({"b": 0, "c": 1}).as_object().unwrap().clone() + ); + assert_eq!(indexes.model_name["gpt"], json!([0])); + indexes.removed("c", 1).unwrap(); + assert!(!indexes.model_name.contains_key("claude")); + assert_eq!(indexes.model_names, vec![json!("gpt"), json!("broken")]); + assert_eq!(indexes.deployment_names.len(), 3); + } + + #[test] + fn a_pattern_routes_with_the_matched_parts() { + let mut router = PatternRouter::default(); + router.add( + "llmengine/*", + json!({"litellm_params": {"model": "openai/*"}, "model_info": {"id": "w"}}), + ); + router.add( + "*meta.llama3*", + json!({"litellm_params": {"model": "bedrock/meta.llama3*"}, "model_info": {"id": "m"}}), + ); + let routed = router.route("llmengine/foo", None).unwrap(); + assert_eq!(routed[0]["litellm_params"]["model"], "openai/foo"); + let routed = router.route("hello-world-meta.llama3-70b", None).unwrap(); + assert_eq!( + routed[0]["litellm_params"]["model"], + "hello-world-meta.llama3-70b" + ); + assert_eq!( + regex_match(r"a(.*)b(.*)", "axbyb").unwrap(), + vec!["xby", ""] + ); + router.remove("w"); + assert!(router.route("llmengine/foo", None).is_none()); + } + + #[rstest] + #[case::any_wildcard( + json!({"model": "openai/*"}), + json!({"id": "w", "keep": 1}), + )] + #[case::azure_ai_keeps_the_breakpoint_flag( + json!({"model": "azure_ai/*"}), + json!({"id": "w", "keep": 1, "supports_prompt_cache_breakpoint": true}), + )] + #[case::azure_ai_named_by_custom_llm_provider( + json!({"model": "*", "custom_llm_provider": "azure_ai"}), + json!({"id": "w", "keep": 1, "supports_prompt_cache_breakpoint": true}), + )] + #[case::a_fixed_model_speaks_for_itself( + json!({"model": "openai/gpt-4o"}), + json!({ + "id": "w", + "keep": 1, + "supports_prompt_cache_breakpoint": true, + "supports_prompt_caching": true, + }), + )] + fn a_wildcard_leaves_the_cache_flags_to_the_catalog( + #[case] params: Value, + #[case] expected: Value, + ) { + let mut router = PatternRouter::default(); + router.add( + "any/*", + json!({ + "litellm_params": params, + "model_info": { + "id": "w", + "keep": 1, + "supports_prompt_cache_breakpoint": true, + "supports_prompt_caching": true, + }, + }), + ); + let routed = router.route("any/gpt-4o", None).unwrap(); + assert_eq!(routed[0]["model_info"], expected); + } +} diff --git a/litellm-rust/crates/router-core/src/lib.rs b/litellm-rust/crates/router-core/src/lib.rs index 05d95d02c84..5f15b6b8df0 100644 --- a/litellm-rust/crates/router-core/src/lib.rs +++ b/litellm-rust/crates/router-core/src/lib.rs @@ -43,21 +43,706 @@ //! - `call(op, request)`: select, run the attempt, and retry or fall back until it succeeds or the //! policy is exhausted. //! -//! # Execution -//! Each attempt resolves an `Executor` after a deployment is picked. A Rust-native operation runs -//! its `inference-*` route machine in-process; anything else goes through `Host::invoke`. A -//! started attempt is never replayed on a different executor. +//! # Hosts: the only boundary +//! The core reaches outside itself only through [`Host`], and every attempt runs through +//! `Host::invoke`: the core picks the deployment, the host runs the call, the core records the +//! outcome. There is no second, "native" attempt path in the core. Which code runs an attempt is +//! the host's business: +//! - `PythonHost` (inside a Python process) builds the request for the picked deployment and calls +//! `litellm.`. It never calls back into a Python router. //! -//! # Hosts -//! `Host` is what the embedding process provides: `invoke` for non-native operations, `select` -//! for strategies and sub-routers not yet ported, `allow_fallback` for access and budget checks, -//! `sleep`, `observe` for routing events, and `cache`, a DualCache-style store that holds all -//! routing state under the same keys, encodings and TTLs the Python router uses. `PythonHost` -//! serves them from the existing Python code when the router runs inside Python; `NativeHost` -//! serves them in the Rust server, where an operation without a Rust implementation makes the -//! config fail at build time, never partway through a request. +//! An op or provider a host cannot run fails `Host::supports` when the router is built or +//! changed, never partway through a request. +//! +//! The core also builds its deployments through the host ([`build`]): for each configured entry +//! it finds the secret references and asks `secrets`, then `validate`, `provider` and `added`. +//! +//! Besides `invoke` and `supports`, a host provides `select` for strategies and plugins not yet +//! ported, `allow_fallback` for access and budget checks, `refuses_fallbacks` for a chain's dict +//! targets, `sleep`, `observe` for routing events, and `cache`, a DualCache-style store that +//! holds all routing state under the same keys, encodings and TTLs the Python router uses. The +//! core writes cooldowns, allowed-fails counters and usage itself, through that cache. //! //! # Plugins //! Routing strategies, pre-call checks and sub-routers (auto, complexity, quality, adaptive) //! register through `register_plugin` and see a `RoutingContext` (registry, settings, health and //! the current request), never the router itself. +//! +//! # Modules +//! - [`registry`]: deployments and their indices, `deployment_id`, `ChangeSet`. +//! - [`config`] / [`settings`]: static and dynamic config, and the frozen `Settings` view. +//! - [`cache`] / [`health`]: the shared routing cache, cooldowns and usage. +//! - [`select`]: resolve, filter chain, simple-shuffle. +//! - [`call`]: retries and fallbacks. +//! - [`eligibility`]: the endpoints a deployment refuses to serve. +//! - [`precall`]: the `enable_pre_call_checks` filter. +//! - [`model_info`]: model info merged from the cost map, asked of the host at query time. +//! - [`host`]: the `Host` boundary. + +pub mod access; +pub mod build; +pub mod cache; +pub mod call; +pub mod config; +pub mod eligibility; +pub mod error; +pub mod fallback; +pub mod group_info; +pub mod health; +pub mod health_routing; +pub mod host; +pub mod host_index; +pub mod lookup; +pub mod mock; +pub mod model_info; +pub mod outcome; +pub mod precall; +pub mod prerouting; +pub mod ptu; +pub mod registry; +pub mod resolve; +pub mod retry; +pub mod routing_groups; +pub mod routing_plugins; +pub mod select; +pub mod settings; +pub mod strategy; +pub mod weights; + +use std::sync::{ + Arc, + atomic::{AtomicBool, Ordering}, +}; + +use arc_swap::ArcSwap; +use parking_lot::Mutex; + +pub use crate::{ + cache::{InMemoryRouterCache, RouterCache}, + call::RouterResponse, + config::{Op, RouterConfig, RouterDynamicConfig, RoutingStrategy}, + error::{AttemptError, ConfigError, Error, FailureKind, Payload, Unsupported}, + health::{Clock, Health, Outcome, SystemClock}, + host::{ + Attempt, Blocked, BuildAsk, BuildPath, Built, Host, Output, Provided, RetryCause, + RouterEvent, RouterRequest, RoutingContext, SecretAt, SecretRef, SelectStage, + }, + outcome::{OutcomeCallbacks, OutcomeWrite}, + prerouting::{PreRoute, PreRouted}, + registry::{ChangeSet, Deployment, DeploymentSpec, Registry, RegistryVersion}, + select::{Randomness, SeededRandomness, ThreadRandomness}, + settings::Settings, +}; +use crate::{ + health::{GroupShape, routing_group_has_alternatives, team_model_has_alternatives}, + registry::DeploymentFilter, + select::SelectInputs, +}; + +/// The registry and settings one call reads, swapped atomically by `apply`. +pub(crate) struct RouterStateSnapshot { + pub registry: Arc, + pub settings: Arc, +} + +impl RouterStateSnapshot { + /// `_get_routing_context` for a pick of `name`: unless the request overrides the strategy, + /// a stale member map naming a group no longer configured raises `KeyError(group)`. + pub(crate) fn missing_routing_group( + &self, + request: &RouterRequest, + name: &str, + ) -> Result<(), Error> { + if strategy::request_override(request.routing_strategy.as_deref()).is_some() { + return Ok(()); + } + match self.registry.missing_member_group(name) { + Some(group) => Err(Error::MissingRoutingGroup { + group: group.to_owned(), + }), + None => Ok(()), + } + } +} + +pub(crate) struct Core { + host: Arc, + rng: Arc, + health: Health, + config: RouterConfig, + snapshot: ArcSwap, + writer: Mutex<()>, + retired: AtomicBool, +} + +impl Core { + /// Reject what this host cannot serve, before any request. + fn validate(&self, registry: &Registry, settings: &Settings) -> Result<(), ConfigError> { + let can_select = self.host.can_select(); + let strategies = std::iter::once(&settings.routing_strategy).chain( + registry + .routing_groups() + .iter() + .map(|group| &group.routing_strategy), + ); + for strategy in strategies { + if !strategy.is_native() && !can_select { + return Err(Unsupported::Strategy(strategy.clone()).into()); + } + } + if let Some(plugin) = settings.plugins.first() + && !can_select + { + return Err(Unsupported::Plugin(plugin.name.clone()).into()); + } + for deployment in registry.deployments(&DeploymentFilter::default()) { + self.validate_deployment(&deployment)?; + } + Ok(()) + } + + /// Reject a deployment whose provider this host cannot run a served op on. + fn validate_deployment(&self, deployment: &Deployment) -> Result<(), ConfigError> { + for op in &self.config.served_ops { + if !self.host.supports(op, deployment.provider()) { + return Err(Unsupported::Op { + op: op.clone(), + provider: deployment.provider().to_owned(), + } + .into()); + } + } + Ok(()) + } +} + +/// How a reported failure is recorded. +#[derive(Clone, Debug, Default)] +pub struct FailureReport { + /// Cool down for the deployment's own `cooldown_time` (then the provider's retry-after), + /// as the failure callback does, rather than the router default. + pub deployment_cooldown: bool, + /// The `cooldown_time` the failed call's own params carried (`model_info`, then + /// `litellm_params`), which wins over the provider's retry-after for a callback report. + pub cooldown_time: Option, + /// Count the failure only: its status is one no cooldown rule reads (an empty status). + pub count_only: bool, + /// The model group the request was for, whose fallbacks keep a lone deployment cooling. + pub requested_group: Option, + /// Cool down without counting the failure (`set_cooldown_deployments` alone). + pub uncounted: bool, + /// Also count the failed request against the deployment's per-minute requests + /// (`async_deployment_callback_on_failure`). + pub counts_request: bool, + /// The model the request's usage is keyed by (metadata's `deployment`, which a wildcard route + /// fills), when it is not the deployment's own. + pub usage_model: Option, + /// The failure callbacks of a logged failure; a report no callback makes keeps the default. + pub callbacks: OutcomeCallbacks, +} + +/// How a reported success is recorded. +#[derive(Clone, Debug, Default)] +pub struct SuccessReport { + /// See [`host::Output::tracks_usage`]. + pub tracks_usage: Option, + /// See [`FailureReport::usage_model`]. + pub usage_model: Option, + /// See [`FailureReport::callbacks`]. + pub callbacks: OutcomeCallbacks, +} + +/// Builds a [`Router`], with test seams. +pub struct RouterBuilder { + config: RouterConfig, + dynamic: RouterDynamicConfig, + host: Arc, + rng: Arc, + clock: Arc, +} + +impl RouterBuilder { + #[must_use] + pub fn randomness(self, rng: Arc) -> Self { + Self { rng, ..self } + } + + #[must_use] + pub fn clock(self, clock: Arc) -> Self { + Self { clock, ..self } + } + + pub fn build(self) -> Result { + let registry = Registry::build( + self.config.model_list.clone(), + self.dynamic.model_alias_map.clone(), + self.config.routing_groups.clone(), + )? + .with_member_groups(self.config.routing_group_members.clone()); + let settings = Settings::new(&self.config, Arc::new(self.dynamic)); + let core = Core { + health: Health::new(self.host.cache(), self.clock), + host: self.host, + rng: self.rng, + config: self.config, + snapshot: ArcSwap::from_pointee(RouterStateSnapshot { + registry: Arc::new(Registry::default()), + settings: Arc::new(settings.clone()), + }), + writer: Mutex::new(()), + retired: AtomicBool::new(false), + }; + core.validate(®istry, &settings)?; + core.snapshot.store(Arc::new(RouterStateSnapshot { + registry: Arc::new(registry), + settings: Arc::new(settings), + })); + Ok(Router { + core: Arc::new(core), + }) + } +} + +/// The router: the entry points `litellm.Router` delegates to. +#[derive(Clone)] +pub struct Router { + core: Arc, +} + +impl Router { + pub fn new( + config: RouterConfig, + dynamic: RouterDynamicConfig, + host: Arc, + ) -> Result { + Self::builder(config, dynamic, host).build() + } + + pub fn builder( + config: RouterConfig, + dynamic: RouterDynamicConfig, + host: Arc, + ) -> RouterBuilder { + RouterBuilder { + config, + dynamic, + host, + rng: Arc::new(ThreadRandomness), + clock: Arc::new(SystemClock), + } + } + + fn live(&self) -> Result<(), Error> { + if self.core.retired.load(Ordering::Acquire) { + return Err(Error::Retired); + } + Ok(()) + } + + /// Select, run the attempt, and retry or fall back until it succeeds or the policy is + /// exhausted. + pub async fn call(&self, op: Op, request: RouterRequest) -> Result { + self.live()?; + self.core.call(op, Arc::new(request)).await + } + + /// Resolve, filter and pick a deployment without calling it. + pub async fn select(&self, op: &Op, request: &RouterRequest) -> Result, Error> { + self.live()?; + let snapshot = self.core.snapshot.load_full(); + routing_plugins::refuse_sync_pick( + snapshot.settings.routing_plugins, + &snapshot.settings.routing_strategy, + request, + )?; + let routed = self + .core + .pre_routed(PreRoute { + op, + request, + name: &request.model, + fallback_depth: 0, + retry: 0, + }) + .await?; + let name = match routed { + PreRouted::Model(name) => name, + PreRouted::Deployment(deployment) => return Ok(deployment), + }; + snapshot.missing_routing_group(request, &name)?; + let candidates = self + .core + .candidates(&snapshot, op, request, &name, &SelectInputs::default()) + .await?; + self.core.pick(&snapshot, op, request, &candidates).await + } + + /// The deployments a pick for `request` would choose among: resolved, team- and + /// cooldown-filtered, through the plugins and the lowest `order`. + pub async fn healthy( + &self, + op: &Op, + request: &RouterRequest, + ) -> Result { + self.live()?; + let snapshot = self.core.snapshot.load_full(); + let candidates = self + .core + .candidates( + &snapshot, + op, + request, + &request.model, + &SelectInputs::default(), + ) + .await?; + // `_core_async_get_healthy_deployments` raises when nothing is left. + if candidates.healthy.is_empty() { + return Err(Error::NoHealthyDeployments { + model: candidates.failed_model(&snapshot.registry), + }); + } + Ok(candidates) + } + + pub fn registry(&self) -> Arc { + self.core.snapshot.load().registry.clone() + } + + pub fn settings(&self) -> Arc { + self.core.snapshot.load().settings.clone() + } + + pub fn health(&self) -> &Health { + &self.core.health + } + + /// Record an outcome the router did not see itself, such as a health check. Returns whether + /// the deployment was cooled down. + pub async fn report(&self, id: &str, outcome: Outcome) -> Result { + match outcome { + Outcome::Success { total_tokens } => self + .report_success(id, total_tokens as f64, 1.0, SuccessReport::default()) + .await + .map(|_| false), + Outcome::Failure(failure) => self + .report_failure(id, &failure, FailureReport::default()) + .await + .map(|cooldown| cooldown.is_some()), + } + } + + /// A success a logging callback saw (`deployment_callback_on_success`): counted for a + /// configured deployment, with its tokens and requests when it tracks usage. Returns the + /// success counter's key, or `None` for an unknown deployment. + pub async fn report_success( + &self, + id: &str, + total_tokens: f64, + requests: f64, + report: SuccessReport, + ) -> Result, Error> { + report + .callbacks + .record( + OutcomeWrite::Success, + self.reported_success(id, total_tokens, requests, &report), + ) + .await + } + + async fn reported_success( + &self, + id: &str, + total_tokens: f64, + requests: f64, + report: &SuccessReport, + ) -> Result, Error> { + let snapshot = self.core.snapshot.load_full(); + let Some(deployment) = snapshot.registry.deployment(id) else { + return Ok(None); + }; + let key = self.core.health.count_success(id).await?; + if report + .tracks_usage + .unwrap_or_else(|| deployment.tracks_usage()) + { + self.core + .health + .add_usage( + &deployment, + report.usage_model.as_deref(), + total_tokens, + requests, + ) + .await?; + } + Ok(Some(key)) + } + + /// `sync_deployment_callback_on_success`: count a success against `id`, configured or not, + /// and return the counter's key. + pub async fn count_success( + &self, + id: &str, + callbacks: OutcomeCallbacks, + ) -> Result, Error> { + callbacks + .record(OutcomeWrite::Success, async { + self.core.health.count_success(id).await.map(Some) + }) + .await + } + + /// A failure the router did not see itself: counted against `id` (known or not), then + /// cooled down by the rules (`record_failure`, `deployment_callback_on_failure`). Returns the + /// cooldown applied, if any. + pub async fn report_failure( + &self, + id: &str, + failure: &AttemptError, + report: FailureReport, + ) -> Result, Error> { + report + .callbacks + .record( + OutcomeWrite::Failure, + self.reported_failure(id, failure, &report), + ) + .await + } + + async fn reported_failure( + &self, + id: &str, + failure: &AttemptError, + report: &FailureReport, + ) -> Result, Error> { + if !report.uncounted && !failure.outside_deployment { + self.core.health.count_failure(id).await?; + } + let snapshot = self.core.snapshot.load_full(); + let Some(deployment) = snapshot.registry.deployment(id) else { + return Ok(None); + }; + if report.counts_request { + self.core + .health + .count_failed_request(&deployment, report.usage_model.as_deref()) + .await?; + } + if failure.outside_deployment { + return Ok(None); + } + // An empty status is no cooldown (`_is_cooldown_required`), unless the deployment's + // own policy names the error. + if report.count_only + && crate::health::deployment_allowed_fails( + deployment.model_info.allowed_fails_policy.as_ref(), + failure, + ) + .is_none() + { + return Ok(None); + } + let cooldown = if report.deployment_cooldown { + report + .cooldown_time + .filter(|seconds| *seconds >= 0.0) + .and_then(|seconds| std::time::Duration::try_from_secs_f64(seconds).ok()) + .or(failure.retry_after) + .unwrap_or(snapshot.settings.cooldown_time) + } else { + snapshot.settings.cooldown_time + }; + let shape = GroupShape { + single_deployment: snapshot + .registry + .group(&deployment.model_name) + .is_some_and(|group| group.len() == 1) + && !report + .requested_group + .as_deref() + .is_some_and(|group| routing_group_has_alternatives(&snapshot.registry, group)) + && !team_model_has_alternatives(&snapshot.registry, &deployment), + has_fallback_group: snapshot + .settings + .dynamic + .fallback_targets(&[report + .requested_group + .as_deref() + .unwrap_or(&deployment.model_name)]) + .is_some(), + }; + self.core + .health + .cool_down_after(&snapshot.settings, &deployment, shape, failure, cooldown) + .await + } + + /// Apply one change set atomically and return the new registry version. + pub fn apply(&self, changes: ChangeSet) -> Result { + let _writer = self.core.writer.lock(); + let current = self.core.snapshot.load_full(); + let aliases = changes + .settings + .as_ref() + .map(|dynamic| dynamic.model_alias_map.clone()); + let registry = current.registry.apply(&changes, aliases)?; + let settings = match changes.settings { + Some(dynamic) => Settings::new(&self.core.config, Arc::new(dynamic)), + None => Settings::clone(¤t.settings), + }; + self.core.validate(®istry, &settings)?; + let version = registry.version(); + self.core.snapshot.store(Arc::new(RouterStateSnapshot { + registry: Arc::new(registry), + settings: Arc::new(settings), + })); + Ok(version) + } + + /// `get_router_model_info` for one deployment, its cost-map reads asked of the host now + /// ([`model_info::router_model_info`]). + pub async fn router_model_info( + &self, + query: &model_info::RouterModelInfoQuery, + ) -> Result, Error> { + model_info::router_model_info(self.core.host.as_ref(), query).await + } + + /// `_set_model_group_info` for the deployments `query` names ([`group_info::model_group_info`]). + pub async fn model_group_info( + &self, + query: &group_info::GroupInfoQuery, + ) -> Result, Error> { + group_info::model_group_info(self.core.host.as_ref(), query).await + } + + /// `get_configured_token_limits` for deployment `id` ([`group_info::configured_token_limits`]). + pub async fn configured_token_limits( + &self, + id: serde_json::Value, + configured: &serde_json::Map, + ) -> Result<(Option, Option), Error> { + group_info::configured_token_limits(self.core.host.as_ref(), id, configured).await + } + + /// `get_model_listing_info` over a listed name's deployments ([`group_info::model_listing_info`]). + pub async fn model_listing_info( + &self, + deployments: &[serde_json::Value], + ) -> Result { + group_info::model_listing_info(self.core.host.as_ref(), deployments).await + } + + /// `get_deployment_model_info` for deployment `model_id` serving `model_name` + /// ([`model_info::deployment_model_info`]). + pub async fn deployment_model_info( + &self, + model_id: &str, + model_name: &serde_json::Value, + ) -> Result>, Error> { + model_info::deployment_model_info(self.core.host.as_ref(), model_id, model_name).await + } + + /// Build `entries` through the host ([`build::build`]), installing each deployment it + /// accepts as soon as it does, as `install` says: a replace empties the registry first (and + /// takes `set_model_list`'s `model_names` once every entry is built), an append or upsert + /// adds at the end. When a step fails, what was installed before it stays, as the Python + /// router keeps the deployments it added before the one that raised. + /// + /// The build installs into a registry of its own, indexed in place, and serves it once it + /// ends (or fails), so n deployments build in O(n). A write that lands while it runs (a step + /// answer that applies a changeset) is kept: what the build installed is added over it. + pub async fn build( + &self, + path: BuildPath, + entries: Vec>, + resolve_secrets: bool, + install: build::Install, + ) -> Result, Error> { + let model_names: Vec = entries + .iter() + .map(|entry| match entry.get("model_name") { + Some(serde_json::Value::String(name)) => name.clone(), + None | Some(serde_json::Value::Null) => "None".to_owned(), + Some(other) => other.to_string(), + }) + .collect(); + let replace = install == build::Install::Replace; + let base = self.core.snapshot.load_full().registry.clone(); + let mut working = if replace { + base.emptied() + } else { + Registry::clone(&base) + }; + let mut installed: Vec = Vec::new(); + let add = path == BuildPath::Add; + let core = &self.core; + let built = build::build( + core.host.as_ref(), + path, + entries, + resolve_secrets, + |deployment| { + working.push_built(deployment.spec.clone(), add)?; + installed.push(deployment.spec.clone()); + core.validate_deployment(working.last().expect("just installed")) + .map_err(Error::from) + }, + ) + .await; + if built.is_ok() && replace && path == BuildPath::Create { + working.set_model_names(model_names); + } + self.serve_built(&base, working, installed, add)?; + built + } + + /// Serve the registry a build made from `base`. When another write replaced `base` while + /// the build ran, the build's deployments are added over that write instead. + fn serve_built( + &self, + base: &Arc, + built: Registry, + installed: Vec, + add: bool, + ) -> Result<(), ConfigError> { + let _writer = self.core.writer.lock(); + let current = self.core.snapshot.load_full(); + let registry = if Arc::ptr_eq(¤t.registry, base) { + built + } else { + let mut next = Registry::clone(¤t.registry); + for spec in installed { + next.push_built(spec, add)?; + } + next + }; + self.core.snapshot.store(Arc::new(RouterStateSnapshot { + registry: Arc::new(registry), + settings: current.settings.clone(), + })); + Ok(()) + } + + /// Keep the providers `from`'s host resolved for the deployments this router serves too (a + /// router rebuilt over the same deployments for new static settings). + pub fn adopt_providers(&self, from: &Self) { + let _writer = self.core.writer.lock(); + let current = self.core.snapshot.load_full(); + let registry = current + .registry + .with_providers_of(&from.core.snapshot.load().registry); + self.core.snapshot.store(Arc::new(RouterStateSnapshot { + registry: Arc::new(registry), + settings: current.settings.clone(), + })); + } + + /// Stop serving: every later call and selection fails with [`Error::Retired`]. + pub fn retire(&self) { + self.core.retired.store(true, Ordering::Release); + } +} diff --git a/litellm-rust/crates/router-core/src/lookup.rs b/litellm-rust/crates/router-core/src/lookup.rs new file mode 100644 index 00000000000..39c235d6f7b --- /dev/null +++ b/litellm-rust/crates/router-core/src/lookup.rs @@ -0,0 +1,883 @@ +//! The registry lookups the host serves as `registry()` queries, answered over the deployments and +//! the host indexes ([`crate::host_index::HostIndexes`]) the way the Python router answers them. +//! +//! A lookup answers with [`Row`]s: positions in the deployment list (each with how the row renames +//! the deployment), wildcard routes, or a position the Python list could not be read at. The host +//! turns a position into the entry object it stores, so identity and in-place reads stay its own; +//! a [`Row::Raw`] is handed back for the host to index its list by, which raises exactly where the +//! Python router's own read raised (an index a caller wrote junk into). + +use std::sync::Arc; + +use serde_json::{Map, Value, json}; + +use crate::{ + host_index::{HostIndexes, python_str, truthy}, + registry::{Deployment, Registry}, +}; + +/// One row of a lookup's answer. +#[derive(Clone, Debug, PartialEq)] +pub enum Row { + /// The deployment at `position`, under `name` when the row renames it (an alias or a routing + /// group), as a routing group's member `member` of `group` when `group` is set (its row drops + /// the member's `access_groups` and takes the group's priority as its `order`). + At { + position: usize, + name: Option, + group: Option<(String, String)>, + }, + /// A wildcard route: a copy of a pattern router's deployment, under `name`. + Pattern { entry: Value, name: String }, + /// A position the list cannot be read at: the host indexes its list by it, and raises. + Raw(Value), + /// A position container that is not a list: the host iterates it as its list's positions. + RawPositions(Value), +} + +impl Row { + /// The row as the host reads it. + pub fn wire(&self) -> Value { + match self { + Self::At { + position, + name, + group, + } => { + let mut row = Map::new(); + row.insert("at".to_owned(), json!(position)); + if let Some(name) = name { + row.insert("name".to_owned(), json!(name)); + } + if let Some((group, member)) = group { + row.insert("group".to_owned(), json!(group)); + row.insert("member".to_owned(), json!(member)); + } + Value::Object(row) + } + Self::Pattern { entry, name } => json!({"pattern": entry, "name": name}), + Self::Raw(value) => json!({"raw": value}), + Self::RawPositions(value) => json!({"raw_positions": value}), + } + } +} + +/// An alias `model_list` lists: its name and the group it targets (`hidden` aliases already left +/// out by the host when the whole list is asked for). +#[derive(Clone, Debug, PartialEq, serde::Deserialize)] +pub struct AliasItem { + pub alias: String, + pub target: String, +} + +/// A routing group callable by its name (the host leaves out one a deployment or an alias +/// shadows): its name and member model groups. +#[derive(Clone, Debug, PartialEq, serde::Deserialize)] +pub struct CallableGroup { + pub name: String, + pub members: Vec, +} + +/// What `get_model_list` composes beyond the deployments: the aliases and callable routing groups +/// it lists, and the wildcard requests a name no deployment serves is routed by. +#[derive(Clone, Debug, Default, PartialEq, serde::Deserialize)] +#[serde(default)] +pub struct ListScope { + pub aliases: Vec, + pub groups: Vec, + /// `get_deployments_by_pattern`'s requests: the name, then the name under its provider. + pub wildcard: Vec, + /// Only one part of the list: `"aliases"` (`_core_get_model_list_from_model_alias`) or + /// `"groups"` (`_core_get_model_list_from_routing_groups`). + pub only: Option, +} + +/// A row read for a lookup's filters: the deployment's fields as the row shows them. +struct View<'a> { + name: &'a str, + info: &'a Map, + params: &'a Map, + strips_access_groups: bool, +} + +/// The lookups over one registry version. +pub struct Lookup<'a> { + deployments: &'a [Arc], + host: &'a HostIndexes, +} + +impl<'a> Lookup<'a> { + pub fn new(registry: &'a Registry) -> Self { + Self { + deployments: registry.deployments_from(0), + host: registry.host_indexes(), + } + } + + /// `list[value]` as Python reads it: an int (or bool) in range, negative from the end. + fn position(&self, value: &Value) -> Option { + let len = i64::try_from(self.deployments.len()).ok()?; + let index = match value { + Value::Bool(flag) => i64::from(*flag), + Value::Number(number) => number.as_i64()?, + _ => return None, + }; + let index = if index < 0 { len + index } else { index }; + (0..len) + .contains(&index) + .then(|| usize::try_from(index).ok()) + .flatten() + } + + /// Every position `held` (what an index holds) names, read before any is filtered, as the + /// Python router reads them into a tuple first: `Err` with the row that raises. + fn positions(&self, held: &Value) -> Result, Row> { + let items = match held { + Value::Array(items) => items, + // Iterating an empty string or dict reads no position; anything else raises. + Value::String(text) if text.is_empty() => return Ok(Vec::new()), + Value::Object(map) if map.is_empty() => return Ok(Vec::new()), + _ => return Err(Row::RawPositions(held.clone())), + }; + items + .iter() + .map(|item| self.position(item).ok_or_else(|| Row::Raw(item.clone()))) + .collect() + } + + fn view(&self, position: usize) -> View<'a> { + let deployment = &self.deployments[position]; + View { + name: &deployment.model_name, + info: &deployment.model_info.raw, + params: &deployment.raw_params, + strips_access_groups: false, + } + } + + fn row_view<'r>(&self, row: &'r Row) -> Option> + where + 'a: 'r, + { + match row { + Row::At { + position, + name, + group, + } => { + let base = self.view(*position); + Some(View { + name: name.as_deref().unwrap_or(base.name), + strips_access_groups: group.is_some(), + ..base + }) + } + Row::Pattern { entry, name } => Some(View { + name, + info: entry + .get("model_info") + .and_then(Value::as_object) + .unwrap_or(empty()), + params: entry + .get("litellm_params") + .and_then(Value::as_object) + .unwrap_or(empty()), + strips_access_groups: false, + }), + Row::Raw(_) | Row::RawPositions(_) => None, + } + } + + // -- deployments by name ------------------------------------------------------------------- + + /// `_core_get_all_deployments`: the team's deployments under public name `model_name`, else + /// the group's, else (for a team) every deployment the name includes; each under `rename`. + pub fn all(&self, model_name: &str, team_id: Option<&str>, rename: Option<&str>) -> Vec { + let rows = |positions: Vec| -> Vec { + positions + .into_iter() + .filter(|position| should_include(model_name, &self.view(*position), team_id)) + .map(|position| Row::At { + position, + name: rename.map(str::to_owned), + group: None, + }) + .collect() + }; + if let Some(team) = team_id + && let Some(held) = self.host.team_named(team, model_name) + { + match self.positions(held) { + Ok(positions) => { + let found = rows(positions); + if !found.is_empty() { + return found; + } + } + Err(raw) => return vec![raw], + } + } + if let Some(held) = self.host.named(model_name) { + return match self.positions(held) { + Ok(positions) => rows(positions), + Err(raw) => vec![raw], + }; + } + if team_id.is_some() { + return rows((0..self.deployments.len()).collect()); + } + Vec::new() + } + + /// `_team_deployments_across_teams`: every team's deployments under public name `model`, and + /// how many distinct teams own them. + pub fn across_teams(&self, model: &str) -> (Vec, usize) { + let mut rows = Vec::new(); + for (_, public, held) in &self.host.team_model { + if public != model { + continue; + } + match self.positions(held) { + Ok(positions) => rows.extend(positions.into_iter().map(|position| Row::At { + position, + name: None, + group: None, + })), + Err(raw) => return (vec![raw], 0), + } + } + let mut teams: Vec<&Value> = Vec::new(); + for row in &rows { + if let Some(view) = self.row_view(row) + && let Some(team) = view.info.get("team_id").filter(|team| !team.is_null()) + && !teams.contains(&team) + { + teams.push(team); + } + } + let count = teams.len(); + (rows, count) + } + + /// `_core_config_deployments`: the deployments a config declared (no `db_model: true`). + pub fn config(&self) -> Vec { + (0..self.deployments.len()) + .filter(|position| { + self.view(*position).info.get("db_model") != Some(&Value::Bool(true)) + }) + .map(|position| Row::At { + position, + name: None, + group: None, + }) + .collect() + } + + /// `_core_get_model_ids`: the ids of the group's deployments, or (no group) of the id map's + /// entries, without team deployments when `exclude_team_models`. `Err` with the row that raises. + pub fn model_ids( + &self, + model_name: Option<&str>, + exclude_team_models: bool, + ) -> Result, Row> { + let kept = |view: &View<'_>| { + view.info.contains_key("id") + && !(exclude_team_models && view.info.get("team_id").is_some_and(truthy)) + }; + if let Some(name) = model_name { + let Some(held) = self.host.named(name) else { + return Ok(Vec::new()); + }; + return Ok(self + .positions(held)? + .into_iter() + .map(|position| self.view(position)) + .filter(kept) + .map(|view| python_str(&view.info["id"])) + .collect()); + } + let mut ids = Vec::new(); + for (id, held) in &self.host.id_map { + let position = self.position(held).ok_or_else(|| Row::Raw(held.clone()))?; + if kept(&self.view(position)) { + ids.push(id.clone()); + } + } + Ok(ids) + } + + // -- the model list -------------------------------------------------------------------------- + + /// `_core_get_model_list`: the name's deployments, its alias rows and routing group rows, its + /// wildcard routes when nothing else serves it, or (no name) every alias and group row and + /// then every deployment. + pub fn model_list( + &self, + model_name: Option<&str>, + team_id: Option<&str>, + scope: &ListScope, + ) -> Vec { + let only = scope.only.as_deref(); + let mut rows = Vec::new(); + if let Some(name) = model_name + && only.is_none() + { + rows.extend(self.all(name, team_id, None)); + } + let aliases: &[AliasItem] = if only == Some("groups") { + &[] + } else { + &scope.aliases + }; + for alias in aliases { + match scope.groups.iter().find(|group| group.name == alias.target) { + Some(group) => { + rows.extend(self.group_rows(group).into_iter().map(|row| match row { + Row::At { + position, group, .. + } => Row::At { + position, + name: Some(alias.alias.clone()), + group, + }, + other => other, + })) + } + None => rows.extend(self.all(&alias.target, None, Some(&alias.alias))), + } + } + match (model_name, only) { + (_, Some("aliases")) => {} + (Some(name), _) => { + if let Some(group) = scope.groups.iter().find(|group| group.name == name) { + rows.extend(self.group_rows(group)); + } + } + (None, _) => { + for group in &scope.groups { + rows.extend(self.group_rows(group)); + } + } + } + if only.is_some() { + return rows; + } + if let Some(name) = model_name + && rows.is_empty() + { + rows.extend(self.wildcard(name, team_id, &scope.wildcard)); + } + if model_name.is_none() { + rows.extend((0..self.deployments.len()).map(|position| Row::At { + position, + name: None, + group: None, + })); + } + rows + } + + /// The whole list's rows, or the first row the host's read of the list raises at. + fn readable(&self, rows: Vec) -> Result, Row> { + match rows + .iter() + .find(|row| matches!(row, Row::Raw(_) | Row::RawPositions(_))) + { + Some(raw) => Err(raw.clone()), + None => Ok(rows), + } + } + + /// A callable routing group's rows: each member's deployments under the group's name. + fn group_rows(&self, group: &CallableGroup) -> Vec { + let mut rows = Vec::new(); + for member in &group.members { + for row in self.all(member, None, Some(&group.name)) { + rows.push(match row { + Row::At { position, name, .. } => Row::At { + position, + name, + group: Some((group.name.clone(), member.clone())), + }, + other => other, + }); + } + } + rows + } + + /// `_get_wildcard_deployments`: the global router's routes for the requests, then the team's. + fn wildcard(&self, name: &str, team_id: Option<&str>, requests: &[String]) -> Vec { + let route = |router: &crate::host_index::PatternRouter| -> Vec { + requests + .iter() + .find_map(|request| { + router + .route(request, None) + .filter(|routed| !routed.is_empty()) + }) + .unwrap_or_default() + }; + let mut routed = route(&self.host.patterns); + if let Some(team) = team_id + && let Some((_, router)) = self.host.team_patterns.iter().find(|(t, _)| t == team) + { + routed.extend(route(router)); + } + routed + .into_iter() + .map(|entry| Row::Pattern { + entry, + name: name.to_owned(), + }) + .collect() + } + + /// The deployments `team_id`'s own wildcard pattern routes `name` to + /// (`_core_team_pattern_routers`, keyed by `team_public_model_name`). + pub fn team_wildcard(&self, name: &str, team_id: &str) -> Vec> { + let Some((_, router)) = self.host.team_patterns.iter().find(|(t, _)| t == team_id) else { + return Vec::new(); + }; + router + .route(name, None) + .unwrap_or_default() + .iter() + .filter_map(|entry| entry.get("model_info")?.get("id").map(python_str)) + .filter_map(|id| self.deployments.iter().find(|d| d.id == id).cloned()) + .collect() + } + + /// `_core_get_model_names`: per row of the whole list, a team deployment's public name for + /// its own team, else the row's name. + pub fn model_names( + &self, + team_id: Option<&str>, + scope: &ListScope, + ) -> Result, Row> { + let rows = self.readable(self.model_list(None, None, scope))?; + let mut names = Vec::new(); + for row in &rows { + let Some(view) = self.row_view(row) else { + continue; + }; + match view.info.get("team_id").filter(|team| truthy(team)) { + Some(team) => { + if team_id.is_some_and(|id| team == &Value::String(id.to_owned())) + && let Some(public) = view + .info + .get("team_public_model_name") + .filter(|public| truthy(public)) + { + names.push(python_str(public)); + } + } + None => names.push(view.name.to_owned()), + } + } + Ok(names) + } + + /// `_core_get_model_access_groups`: access group to the names of the rows that declare it, + /// in the order they appear. + pub fn access_groups( + &self, + model_name: Option<&str>, + access_group: Option<&str>, + team_id: Option<&str>, + scope: &ListScope, + ) -> Result)>, Row> { + let mut groups: Vec<(String, Vec)> = Vec::new(); + for row in self.readable(self.model_list(model_name, team_id, scope))? { + let Some(view) = self.row_view(&row) else { + continue; + }; + if view.info.is_empty() || view.strips_access_groups { + continue; + } + // `access_groups or []`, iterated: a list's items, a string's characters. + let declared: Vec = match view.info.get("access_groups") { + Some(Value::Array(items)) => items.iter().map(python_str).collect(), + Some(Value::String(text)) => text.chars().map(String::from).collect(), + _ => continue, + }; + for group in declared { + if access_group.is_some_and(|wanted| wanted != group) { + continue; + } + match groups.iter_mut().find(|(held, _)| *held == group) { + Some((_, names)) => names.push(view.name.to_owned()), + None => groups.push((group, vec![view.name.to_owned()])), + } + } + } + Ok(groups) + } + + /// `_core_resolve_model_name_from_model_id`'s scan of the whole list: whether it lists any + /// row, and the name of the first row whose `litellm_params.model` is `model_id`, or ends in + /// `/model_id` or `:model_id`. + pub fn name_serving( + &self, + model_id: &str, + scope: &ListScope, + ) -> Result<(bool, Option), Row> { + let rows = self.readable(self.model_list(None, None, scope))?; + let (slash, colon) = (format!("/{model_id}"), format!(":{model_id}")); + let name = rows.iter().find_map(|row| { + let view = self.row_view(row)?; + let model = view.params.get("model").and_then(Value::as_str)?; + (model == model_id || model.ends_with(&slash) || model.ends_with(&colon)) + .then_some(view.name) + .filter(|name| !name.is_empty()) + .map(str::to_owned) + }); + Ok((!rows.is_empty(), name)) + } + + // -- teams and credentials ------------------------------------------------------------------- + + /// `_core_team_model_has_alternatives`: whether the team's public name of deployment `id` has + /// another unpaused deployment of that team. + pub fn team_model_has_alternatives(&self, id: &str) -> Result { + let Some(held) = self.host.id_entry(id) else { + return Ok(false); + }; + let position = self.position(held).ok_or_else(|| Row::Raw(held.clone()))?; + let info = self.view(position).info; + let (Some(Value::String(team)), Some(Value::String(public))) = + (info.get("team_id"), info.get("team_public_model_name")) + else { + return Ok(false); + }; + let Some(siblings) = self.host.team_named(team, public) else { + return Ok(false); + }; + let unpaused = self + .positions(siblings)? + .into_iter() + .filter(|position| self.view(*position).info.get("blocked") != Some(&Value::Bool(true))) + .count(); + Ok(unpaused > 1) + } + + /// `_get_model_group_deployment_usable_by_team`: the first deployment of the group the team + /// may use. + pub fn group_deployment_usable_by_team( + &self, + model_group: &str, + team_id: Option<&str>, + ) -> Result, Row> { + // `.get(name) or ()`: nothing held, or a falsy value, reads no position. + let Some(held) = self.host.named(model_group).filter(|held| truthy(held)) else { + return Ok(None); + }; + let Value::Array(items) = held else { + return Err(Row::RawPositions(held.clone())); + }; + for item in items { + let position = self.position(item).ok_or_else(|| Row::Raw(item.clone()))?; + if usable_by_team(&self.view(position), team_id) { + return Ok(Some(position)); + } + } + Ok(None) + } + + /// `_get_team_public_name_deployment`: the first deployment of the team's public name. + pub fn team_public_name_deployment( + &self, + model_id: &str, + team_id: &str, + ) -> Result, Row> { + match self.host.team_named(team_id, model_id) { + Some(Value::Array(items)) => match items.first() { + Some(item) => self + .position(item) + .map(Some) + .ok_or_else(|| Row::Raw(item.clone())), + None => Ok(None), + }, + Some(held) if truthy(held) => Err(Row::RawPositions(held.clone())), + Some(_) | None => Ok(None), + } + } + + // -- routable deployments -------------------------------------------------------------------- + + /// `_routable_deployments`: the routable group's deployments for the team (those it may use + /// when any), without paused ones. + pub fn routable(&self, model_group: &str, team_id: Option<&str>) -> Vec { + let named = self.all(model_group, team_id, None); + if named + .iter() + .any(|row| matches!(row, Row::Raw(_) | Row::RawPositions(_))) + { + return named; + } + let usable: Vec<&Row> = named + .iter() + .filter(|row| { + self.row_view(row) + .is_some_and(|view| usable_by_team(&view, team_id)) + }) + .collect(); + let pool: Vec<&Row> = if usable.is_empty() { + named.iter().collect() + } else { + usable + }; + pool.into_iter() + .filter(|row| { + self.row_view(row) + .is_some_and(|view| view.info.get("blocked") != Some(&Value::Bool(true))) + }) + .cloned() + .collect() + } + + /// `_core_get_candidate_model_ids_for_route`'s ids of `rows`, in first-seen order. + pub fn ids_of(&self, rows: &[Row]) -> Result, Row> { + let mut ids: Vec = Vec::new(); + for row in rows { + let view = match row { + Row::Raw(_) | Row::RawPositions(_) => return Err(row.clone()), + _ => self.row_view(row).expect("a readable row"), + }; + if let Some(id) = view.info.get("id").filter(|id| !id.is_null()) { + let id = python_str(id); + if !ids.contains(&id) { + ids.push(id); + } + } + } + Ok(ids) + } + + /// `_core_get_candidate_model_ids_for_route` once the alias resolved `resolved`: the + /// members' ids when it names a callable routing group (`members`), else, for a name that is + /// no model group, the team's deployments under it or its wildcard routes, else the group's. + pub fn candidate_ids( + &self, + resolved: &str, + team_id: Option<&str>, + members: Option<&[String]>, + requests: &[String], + ) -> Result, Row> { + if let Some(members) = members { + let mut ids: Vec = Vec::new(); + for member in members { + for id in self.ids_of(&self.all(member, team_id, None))? { + if !ids.contains(&id) { + ids.push(id); + } + } + } + return Ok(ids); + } + let named = Value::String(resolved.to_owned()); + if !self.host.model_names.contains(&named) { + if team_id.is_some() { + let team_rows = self.all(resolved, team_id, None); + if !team_rows.is_empty() { + return self.ids_of(&team_rows); + } + } + let wildcard = self.wildcard_route(resolved, team_id, requests); + if !wildcard.is_empty() { + return self.ids_of(&wildcard); + } + } + self.ids_of(&self.all(resolved, team_id, None)) + } + + /// `_core_wildcard_route_deployments`: the team router's routes for `model`, else the global + /// router's. + pub fn wildcard_route( + &self, + model: &str, + team_id: Option<&str>, + requests: &[String], + ) -> Vec { + let route = |router: &crate::host_index::PatternRouter| -> Vec { + requests + .iter() + .find_map(|request| { + router + .route(request, None) + .filter(|routed| !routed.is_empty()) + }) + .unwrap_or_default() + }; + let team_owned = team_id + .and_then(|team| self.host.team_patterns.iter().find(|(t, _)| t == team)) + .map(|(_, router)| route(router)) + .unwrap_or_default(); + let routed = if team_owned.is_empty() { + route(&self.host.patterns) + } else { + team_owned + }; + routed + .into_iter() + .map(|entry| Row::Pattern { + entry, + name: model.to_owned(), + }) + .collect() + } +} + +/// `_core_should_include_deployment`: the team's deployment under its public name, or a +/// deployment of the group the team may use. +fn should_include(model_name: &str, view: &View<'_>, team_id: Option<&str>) -> bool { + let team_value = team_id.map(|team| Value::String(team.to_owned())); + let model_team = view.info.get("team_id"); + if let Some(team) = &team_value + && model_team == Some(team) + && view.info.get("team_public_model_name") == Some(&Value::String(model_name.to_owned())) + { + return true; + } + if view.name == model_name { + return team_value.is_none() + || model_team.is_none_or(Value::is_null) + || model_team == team_value.as_ref(); + } + false +} + +/// `deployment_usable_by_team`: a deployment no team owns, or one the team owns. +fn usable_by_team(view: &View<'_>, team_id: Option<&str>) -> bool { + match view.info.get("team_id").filter(|team| !team.is_null()) { + None => true, + Some(owner) => team_id.is_some_and(|team| owner == &Value::String(team.to_owned())), + } +} + +fn empty() -> &'static Map { + static EMPTY: std::sync::OnceLock> = std::sync::OnceLock::new(); + EMPTY.get_or_init(Map::new) +} + +#[cfg(test)] +mod tests { + use serde_json::json; + + use super::*; + use crate::registry::DeploymentSpec; + + fn registry(entries: Value) -> Registry { + let mut registry = Registry::default(); + for entry in entries.as_array().unwrap() { + let spec: DeploymentSpec = serde_json::from_value(entry.clone()).unwrap(); + registry.push_built(spec, false).unwrap(); + } + registry + } + + #[test] + fn a_teams_wildcard_routes_only_the_names_its_public_pattern_matches() { + let registry = registry(json!([ + {"model_name": "t1-w", "litellm_params": {"model": "openai/*"}, + "model_info": {"id": "tw1", "team_id": "T1", "team_public_model_name": "tw-*"}}, + {"model_name": "g", "litellm_params": {"model": "openai/g"}, "model_info": {"id": "gp"}}, + ])); + let lookup = Lookup::new(®istry); + let ids = |name: &str, team: &str| -> Vec { + let routed = lookup.team_wildcard(name, team); + routed.iter().map(|d| d.id.clone()).collect() + }; + assert_eq!(ids("tw-gpt-4o", "T1"), ["tw1"]); + assert!(ids("gpt-4o", "T1").is_empty()); + assert!(ids("tw-gpt-4o", "T2").is_empty()); + } + + fn entry(name: &str, id: &str, info: Value) -> Value { + let mut info = info.as_object().cloned().unwrap_or_default(); + info.insert("id".to_owned(), json!(id)); + json!({"model_name": name, "litellm_params": {"model": format!("openai/{id}")}, "model_info": info}) + } + + fn at(position: usize) -> Row { + Row::At { + position, + name: None, + group: None, + } + } + + #[test] + fn all_prefers_the_teams_public_name_then_the_group() { + let registry = registry(json!([ + entry("gpt", "a", json!({})), + entry( + "team-gpt", + "t", + json!({"team_id": "T", "team_public_model_name": "gpt"}) + ), + entry("gpt", "b", json!({"team_id": "U"})), + ])); + let lookup = Lookup::new(®istry); + assert_eq!(lookup.all("gpt", Some("T"), None), vec![at(1)]); + assert_eq!(lookup.all("gpt", Some("V"), None), vec![at(0)]); + assert_eq!(lookup.all("gpt", None, None), vec![at(0), at(2)]); + assert_eq!(lookup.all("missing", None, None), vec![]); + } + + #[test] + fn a_junk_position_is_handed_back_to_raise() { + let mut registry = registry(json!([entry("gpt", "a", json!({}))])); + let mut model_name = registry.host_indexes().model_name.clone(); + model_name.insert("gpt".to_owned(), json!([0, "x"])); + let changes = crate::registry::ChangeSet { + host_index: Some(crate::host_index::IndexWrite { + model_name: Some(model_name), + ..Default::default() + }), + ..Default::default() + }; + registry = registry.apply(&changes, None).unwrap(); + let lookup = Lookup::new(®istry); + assert_eq!(lookup.all("gpt", None, None), vec![Row::Raw(json!("x"))]); + } + + #[test] + fn the_model_list_lists_aliases_and_groups_then_every_deployment() { + let registry = registry(json!([ + entry("gpt", "a", json!({"access_groups": ["g1"]})), + entry("claude", "b", json!({})), + ])); + let lookup = Lookup::new(®istry); + let scope = ListScope { + aliases: vec![AliasItem { + alias: "fast".to_owned(), + target: "gpt".to_owned(), + }], + groups: vec![CallableGroup { + name: "pool".to_owned(), + members: vec!["gpt".to_owned(), "claude".to_owned()], + }], + wildcard: Vec::new(), + only: None, + }; + let rows = lookup.model_list(None, None, &scope); + assert_eq!(rows.len(), 5); + assert_eq!( + rows[0], + Row::At { + position: 0, + name: Some("fast".to_owned()), + group: None + } + ); + assert_eq!( + lookup.model_names(None, &scope).unwrap(), + ["fast", "pool", "pool", "gpt", "claude"] + ); + assert_eq!( + lookup.access_groups(None, None, None, &scope).unwrap(), + vec![("g1".to_owned(), vec!["fast".to_owned(), "gpt".to_owned()])] + ); + } +} diff --git a/litellm-rust/crates/router-core/src/mock.rs b/litellm-rust/crates/router-core/src/mock.rs new file mode 100644 index 00000000000..898782e6203 --- /dev/null +++ b/litellm-rust/crates/router-core/src/mock.rs @@ -0,0 +1,44 @@ +//! The request's `mock_testing_*` flags: a failure the router raises itself, without calling a +//! deployment (`_handle_mock_testing_fallbacks` and `_handle_mock_testing_rate_limit_error`). + +use serde::Deserialize; + +use crate::error::{AttemptError, FailureKind}; + +/// Which mock failure a request asks for. +#[derive( + Clone, Copy, Debug, PartialEq, Eq, Deserialize, strum::IntoStaticStr, strum::EnumString, +)] +#[serde(rename_all = "snake_case")] +#[strum(serialize_all = "snake_case")] +pub enum MockFailure { + /// `mock_testing_fallbacks`: an `InternalServerError` before the first level's retries. + Fallbacks, + /// `mock_testing_context_fallbacks`: a `ContextWindowExceededError`, likewise. + ContextWindow, + /// `mock_testing_content_policy_fallbacks`: a `ContentPolicyViolationError`, likewise. + ContentPolicy, + /// `mock_testing_rate_limit_error`: a `RateLimitError` in place of every level's first + /// attempt, which the retry loop handles like a failed attempt. + RateLimit, +} + +impl MockFailure { + /// Whether the failure stands in for the first level's retry loop (raised before its retries). + pub fn before_retries(self) -> bool { + !matches!(self, Self::RateLimit) + } + + /// The failure as the rules see it. + pub fn failure(self, num_retries: Option) -> AttemptError { + let (kind, status) = match self { + Self::Fallbacks => (FailureKind::InternalServer, 500), + Self::ContextWindow => (FailureKind::ContextWindowExceeded, 400), + Self::ContentPolicy => (FailureKind::ContentPolicyViolation, 400), + Self::RateLimit => (FailureKind::RateLimit, 429), + }; + let mut failure = AttemptError::new(kind, Some(status)); + failure.num_retries = num_retries; + failure + } +} diff --git a/litellm-rust/crates/router-core/src/model_info.rs b/litellm-rust/crates/router-core/src/model_info.rs new file mode 100644 index 00000000000..861388f36ba --- /dev/null +++ b/litellm-rust/crates/router-core/src/model_info.rs @@ -0,0 +1,716 @@ +//! Model info: what the router reports a deployment can do and costs, merged from the process +//! cost map, what the background refresh discovered, and the deployment's own `model_info`. +//! +//! The merge rules are the Python router core's (`get_router_model_info`, +//! `get_deployment_model_info`); every cost-map read is a [`Host::model_info`] lookup asked while +//! the query runs, so a price reload or a registration the host made since the router was built is +//! what the next query reads. The host keeps choosing the keys a lookup tries (`get_model_info`); +//! the core chooses which lookups a query makes. +//! +//! Values the host cannot spell as JSON travel as opaque tokens ([`is_opaque`]): the merge moves +//! them between layers without reading them, and the host restores them. + +use serde_json::{Map, Number, Value}; + +use crate::{ + error::Error, + host::{Host, LogLevel, ModelInfoLookup}, +}; + +/// The key the host wraps a value it cannot spell as JSON in. +pub const OPAQUE_KEY: &str = "__litellm_python_value__"; + +/// A value the host passed as an opaque token: the merge moves it, never reads it. +pub fn is_opaque(value: &Value) -> bool { + matches!(value, Value::Object(map) if map.len() == 1 && map.contains_key(OPAQUE_KEY)) +} + +/// One `get_router_model_info` query, as the host reads it off the deployment. +#[derive(Clone, Debug, Default, serde::Deserialize)] +#[serde(default)] +pub struct RouterModelInfoQuery { + /// `model_info.id`. + pub deployment_id: Value, + /// `model_info.base_model`, then `litellm_params.base_model`: the first that is not `None` + /// names the model whose limits and costs the deployment reports. + pub model_info_base_model: Value, + pub params_base_model: Value, + /// The validated `litellm_params.model`. + pub params_model: String, + /// `base_model` among the validated params' extra fields. + pub extra_base_model: Value, + /// The name the caller asked for, which a wildcard deployment matches. + pub received_model_name: String, +} + +/// A merged model info, and the keys whose values are the deployment's own (the last layer). +#[derive(Clone, Debug, PartialEq)] +pub struct MergedModelInfo { + pub info: Map, + pub user_keys: Vec, +} + +pub(crate) fn raised(exception: &'static str, detail: impl Into) -> Error { + Error::ModelInfo { + exception, + detail: detail.into(), + } +} + +pub(crate) fn malformed(lookup: &str) -> Error { + raised( + "ValueError", + format!("the host answered a `{lookup}` lookup malformed"), + ) +} + +/// Python truthiness of a JSON value. +pub(crate) fn truthy(value: &Value) -> bool { + match value { + Value::Null => false, + Value::Bool(flag) => *flag, + Value::Number(number) => number.as_f64().is_some_and(|n| n != 0.0), + Value::String(text) => !text.is_empty(), + Value::Array(items) => !items.is_empty(), + Value::Object(map) => is_opaque(value) || !map.is_empty(), + } +} + +/// `(value or 0) > 0`, for a value that is a number (anything else reads as not positive). +fn positive_or_zero(value: Option<&Value>) -> bool { + value + .filter(|value| truthy(value)) + .and_then(Value::as_f64) + .is_some_and(|n| n > 0.0) +} + +pub(crate) async fn lookup(host: &dyn Host, lookup: ModelInfoLookup) -> Result { + host.model_info(&lookup).await +} + +/// `get_model_info(model=...)`: the row, or `None` when the lookup tolerates a failure and failed. +async fn catalog_row( + host: &dyn Host, + model: Value, + tolerate: bool, +) -> Result>, Error> { + match lookup( + host, + ModelInfoLookup::ModelInfo { + model, + tolerate, + exact: false, + }, + ) + .await? + { + Value::Null => Ok(None), + Value::Object(row) if tolerate && row.get("failed") == Some(&Value::Bool(true)) => Ok(None), + Value::Object(row) => Ok(Some(row)), + _ => Err(malformed("model_info")), + } +} + +/// `model_cost.get(key)`: `None` when the key is absent, else the entry as stored. +async fn cost_entry(host: &dyn Host, key: &str, process: bool) -> Result, Error> { + match lookup( + host, + ModelInfoLookup::CostEntry { + key: key.to_owned(), + process, + }, + ) + .await? + { + Value::Null => Ok(None), + Value::Object(mut answer) => Ok(Some(answer.remove("entry").unwrap_or(Value::Null))), + _ => Err(malformed("cost_entry")), + } +} + +pub(crate) async fn discovered(host: &dyn Host, id: Value) -> Result, Error> { + match lookup(host, ModelInfoLookup::Discovered { id }).await? { + Value::Null => Ok(Map::new()), + Value::Object(limits) => Ok(limits), + _ => Err(malformed("discovered")), + } +} + +async fn log(host: &dyn Host, level: LogLevel, message: String) -> Result<(), Error> { + lookup(host, ModelInfoLookup::Log { level, message }) + .await + .map(|_| ()) +} + +/// `_provider_model_id`: the model id a cost-map key or base model names, bedrock base model ids +/// normalized. +async fn provider_model_id(host: &dyn Host, name: &str, provider: &str) -> Result { + let model_id = name.rsplit('/').next().unwrap_or(name).to_owned(); + if !provider.starts_with("bedrock") { + return Ok(model_id); + } + match lookup(host, ModelInfoLookup::BedrockBase { model_id }).await? { + Value::String(base) => Ok(base), + _ => Err(malformed("bedrock_base")), + } +} + +/// `_exact_base_model_max_input_tokens`: the base model's own window, unless its row came from a +/// fallback-generalization rule or it is unmapped. +async fn exact_base_model_max_input_tokens( + host: &dyn Host, + base_model: &str, +) -> Result, Error> { + let answer = lookup( + host, + ModelInfoLookup::ModelInfo { + model: Value::String(base_model.to_owned()), + tolerate: true, + exact: true, + }, + ) + .await?; + let Value::Object(mut answer) = answer else { + return Ok(None); + }; + if answer.get("failed") == Some(&Value::Bool(true)) + || answer.get("generalized") == Some(&Value::Bool(true)) + { + return Ok(None); + } + let Some(Value::Object(row)) = answer.remove("row") else { + return Ok(None); + }; + Ok(row + .get("max_input_tokens") + .filter(|tokens| !tokens.is_null()) + .cloned()) +} + +/// `_base_model_widened_window`: a base model naming the same model as the route can only widen +/// its context window. +async fn base_model_widened_window( + host: &dyn Host, + route: &Map, + base_model: Option<&str>, + provider: &str, +) -> Result, Error> { + let Some(route_max) = route + .get("max_input_tokens") + .filter(|tokens| !tokens.is_null()) + else { + return Ok(None); + }; + let Some(base_model) = base_model else { + return Ok(None); + }; + let route_key = match route.get("key") { + Some(Value::String(key)) => key.clone(), + Some(_) => return Err(raised("AttributeError", "key")), + None => return Err(raised("KeyError", "key")), + }; + let route_id = provider_model_id(host, &route_key, provider).await?; + let base_id = provider_model_id(host, base_model, provider).await?; + let unversioned = base_id + .split_once('@') + .map_or(base_id.as_str(), |(id, _)| id); + if route_id != base_id && route_id != unversioned { + return Ok(None); + } + let Some(base_max) = exact_base_model_max_input_tokens(host, base_model).await? else { + return Ok(None); + }; + match (base_max.as_f64(), route_max.as_f64()) { + (Some(base), Some(route)) if base > route => Ok(Some(base_max)), + (Some(_), Some(_)) => Ok(None), + _ => Err(raised( + "TypeError", + "'<=' not supported between the base model's and the route's max_input_tokens", + )), + } +} + +/// `_declared_base_model`: a `base_model` that is a string. +fn declared_base_model(value: &Value) -> Option<&str> { + value.as_str() +} + +/// `get_router_model_info` after the host validated the deployment's params: the cost-map row of +/// the model the deployment serves, widened by its base model, overlaid with what the background +/// refresh discovered and then with every value the deployment sets itself. `None` when the +/// lookup answers `None`; an unmapped model raises as the lookup raises. +pub async fn router_model_info( + host: &dyn Host, + query: &RouterModelInfoQuery, +) -> Result, Error> { + let base_model = if query.model_info_base_model.is_null() { + query.params_base_model.clone() + } else { + query.model_info_base_model.clone() + }; + let mut model = base_model.clone(); + + let (resolved_model, provider) = match lookup( + host, + ModelInfoLookup::Provider { + model: query.params_model.clone(), + }, + ) + .await? + { + Value::Array(pair) => match pair.as_slice() { + [Value::String(model), Value::String(provider), ..] => { + (model.clone(), provider.clone()) + } + _ => return Err(malformed("provider")), + }, + _ => return Err(malformed("provider")), + }; + + if provider == "azure" && base_model.is_null() { + // Router init registers every deployment name into the cost map as a zeroed stub, so + // membership alone cannot tell a resolvable name apart; it needs usable limits or costs. + let fallback_key = if resolved_model.starts_with("azure/") { + resolved_model.clone() + } else { + format!("azure/{resolved_model}") + }; + let resolves = match cost_entry(host, &fallback_key, false).await? { + Some(Value::Object(entry)) if !is_opaque(&Value::Object(entry.clone())) => { + positive_or_zero(entry.get("max_input_tokens")) + || positive_or_zero(entry.get("input_cost_per_token")) + } + _ => false, + }; + if resolves { + log( + host, + LogLevel::Debug, + format!( + "Azure deployment '{resolved_model}' has no base_model set; using '{fallback_key}' from the model cost map for max tokens, cost tracking, etc." + ), + ) + .await?; + } else { + log( + host, + LogLevel::Error, + format!( + "Could not identify azure model '{resolved_model}'. Set azure 'base_model' for accurate max tokens, cost tracking, etc.- https://docs.litellm.ai/docs/proxy/cost_tracking#spend-tracking-for-azure-openai-models" + ), + ) + .await?; + } + } else if provider != "azure" { + model = Value::String(resolved_model.clone()); + if resolved_model.contains('*') { + // Only a wildcard deployment asks which model its pattern routed the name to. + if let Value::Array(routed) = lookup( + host, + ModelInfoLookup::PatternModels { + name: query.received_model_name.clone(), + }, + ) + .await? + && let Some(routed) = routed + .iter() + .find(|routed| routed.get("id").unwrap_or(&Value::Null) == &query.deployment_id) + { + model = routed.get("model").cloned().unwrap_or(Value::Null); + } + } + } + + if model.is_null() { + // An azure deployment without a base_model is looked up by its own model. + model = Value::String(resolved_model); + } + let Value::String(model) = model else { + return Err(raised( + "AttributeError", + format!( + "'{}' object has no attribute 'startswith'", + crate::group_info::type_name(&model) + ), + )); + }; + let model_info_name = if model.starts_with(&format!("{provider}/")) { + model + } else { + format!("{provider}/{model}") + }; + + let Some(row) = catalog_row(host, Value::String(model_info_name), false).await? else { + return Ok(None); + }; + + let user_model_info = match lookup(host, ModelInfoLookup::UserModelInfo).await? { + Value::Object(info) => info, + Value::Null => Map::new(), + _ => return Err(malformed("user_model_info")), + }; + let declared = user_model_info + .get("base_model") + .and_then(declared_base_model) + .or_else(|| declared_base_model(&query.extra_base_model)); + let widened = base_model_widened_window(host, &row, declared, &provider).await?; + let discovered = discovered(host, query.deployment_id.clone()).await?; + + let mut info = row; + if let Some(max_input_tokens) = widened { + info.insert("max_input_tokens".to_owned(), max_input_tokens); + } + info.extend(discovered); + let mut user_keys = Vec::new(); + for (key, value) in &user_model_info { + if !value.is_null() { + info.insert(key.clone(), value.clone()); + user_keys.push(key.clone()); + } + } + Ok(Some(MergedModelInfo { info, user_keys })) +} + +/// `_convert_stringified_numbers`: a string that reads as a number becomes one (a float when it +/// has an exponent or a point, else an int); any other string stays as written. +fn convert_stringified_number(text: &str) -> Value { + let trimmed = text.trim(); + let digits = without_digit_underscores(trimmed); + let lower = text.to_lowercase(); + if lower.contains('e') || text.contains('.') { + return digits + .and_then(|digits| python_float(&digits)) + .and_then(Number::from_f64) + .map_or_else(|| Value::String(text.to_owned()), Value::Number); + } + digits + .and_then(|digits| python_int(&digits)) + .map_or_else(|| Value::String(text.to_owned()), Value::Number) +} + +/// The text with the underscores Python's `int()`/`float()` accept between digits removed, or +/// `None` when an underscore sits anywhere else. +pub(crate) fn without_digit_underscores(text: &str) -> Option { + let chars: Vec = text.chars().collect(); + let mut out = String::with_capacity(text.len()); + for (index, ch) in chars.iter().enumerate() { + if *ch == '_' { + let before = index.checked_sub(1).and_then(|i| chars.get(i)); + let after = chars.get(index + 1); + if !(before.is_some_and(char::is_ascii_digit) + && after.is_some_and(char::is_ascii_digit)) + { + return None; + } + continue; + } + out.push(*ch); + } + Some(out) +} + +pub(crate) fn python_float(text: &str) -> Option { + let lower = text.to_ascii_lowercase(); + let unsigned = lower.trim_start_matches(['+', '-']); + if matches!(unsigned, "inf" | "infinity" | "nan") { + // Not finite: JSON cannot carry it, so the value stays as written. + return None; + } + let well_formed = text + .chars() + .all(|ch| ch.is_ascii_digit() || matches!(ch, '+' | '-' | '.' | 'e' | 'E')); + if !well_formed { + return None; + } + text.parse::().ok().filter(|n| n.is_finite()) +} + +pub(crate) fn python_int(text: &str) -> Option { + let unsigned = text.strip_prefix(['+', '-']).unwrap_or(text); + if unsigned.is_empty() || !unsigned.chars().all(|ch| ch.is_ascii_digit()) { + return None; + } + let text = text.strip_prefix('+').unwrap_or(text); + text.parse::() + .map(Number::from) + .ok() + .or_else(|| text.parse::().map(Number::from).ok()) +} + +/// `update_dictionary`: `new`'s values over `existing`, skipping `None`s; a string that reads as a +/// number becomes one, and a dict merges one level into the dict it replaces. +pub fn update_dictionary(existing: &mut Map, new: &Map) { + for (key, value) in new { + match value { + Value::Null => {} + Value::String(text) => { + existing.insert(key.clone(), convert_stringified_number(text)); + } + Value::Object(nested) if !is_opaque(value) => { + let merged = match existing.get(key) { + Some(current @ Value::Object(current_map)) if !is_opaque(current) => { + let mut merged = current_map.clone(); + merged.extend(nested.clone()); + merged + } + _ => nested.clone(), + }; + existing.insert(key.clone(), Value::Object(merged)); + } + _ => { + existing.insert(key.clone(), value.clone()); + } + } + } +} + +/// `cost_map_omits_token_price`: the cost-map entries under `keys` exist, and none carries a +/// per-token price. +async fn cost_map_omits_token_price(host: &dyn Host, keys: &[Option<&str>]) -> Result { + let mut entries = 0; + for key in keys.iter().flatten() { + if let Some(entry) = cost_entry(host, key, true).await? + && !is_opaque(&entry) + && let Value::Object(map) = entry + { + entries += 1; + if map.contains_key("input_cost_per_token") || map.contains_key("output_cost_per_token") + { + return Ok(false); + } + } + } + Ok(entries > 0) +} + +/// `get_deployment_model_info`: deployment `model_id`'s own cost-map entry (with what the refresh +/// discovered, over its base model's row), merged over the row of its litellm model `model_name`. +/// A lookup that fails leaves its layer out, as the Python router ignores it. +pub async fn deployment_model_info( + host: &dyn Host, + model_id: &str, + model_name: &Value, +) -> Result>, Error> { + let mut custom: Option> = None; + if let Some(entry) = cost_entry(host, model_id, false).await? { + let entry = if truthy(&entry) { + match entry { + Value::Object(map) if !is_opaque(&Value::Object(map.clone())) => Some(map), + // Not a mapping: the Python router's TypeError, which it ignores. + _ => None, + } + } else { + Some(Map::new()) + }; + if let Some(mut entry) = entry { + entry.extend(discovered(host, Value::String(model_id.to_owned())).await?); + custom = Some(entry); + } + } + + let builtin = catalog_row(host, model_name.clone(), true).await?; + + let mut base_model_key: Option = None; + if let Some(current) = &custom + && let Some(base_model) = current.get("base_model").filter(|base| !base.is_null()) + && let Some(mut base_row) = catalog_row(host, base_model.clone(), true).await? + { + base_model_key = base_row + .get("key") + .and_then(Value::as_str) + .map(str::to_owned); + // The base model provides defaults; the deployment's own entry overrides them. + update_dictionary(&mut base_row, current); + custom = Some(base_row); + } + + let builtin_key = builtin + .as_ref() + .and_then(|row| row.get("key")) + .and_then(Value::as_str) + .map(str::to_owned); + let info = match (custom, builtin) { + (Some(custom), Some(mut builtin)) => { + update_dictionary(&mut builtin, &custom); + builtin + } + (None, Some(builtin)) => builtin, + (Some(custom), None) => custom, + (None, None) => return Ok(None), + }; + if cost_map_omits_token_price( + host, + &[ + Some(model_id), + builtin_key.as_deref(), + base_model_key.as_deref(), + ], + ) + .await? + { + let mut info = info; + info.insert("input_cost_per_token".to_owned(), Value::Null); + info.insert("output_cost_per_token".to_owned(), Value::Null); + return Ok(Some(info)); + } + Ok(Some(info)) +} + +#[cfg(test)] +mod tests { + use std::{ + collections::BTreeMap, + sync::{Arc, Mutex}, + }; + + use async_trait::async_trait; + use serde_json::json; + + use super::*; + use crate::{ + cache::{InMemoryRouterCache, RouterCache}, + error::AttemptError, + host::{Attempt, Output, RouterEvent, RouterRequest, RoutingContext, SelectStage}, + registry::Deployment, + }; + + /// A cost map and a provider, answering lookups the way `PythonHost` does. + #[derive(Default)] + struct Catalog { + rows: BTreeMap, + provider: (String, String), + user_model_info: Value, + lookups: Mutex>, + } + + #[async_trait] + impl Host for Catalog { + async fn invoke(&self, _attempt: &Attempt) -> Result { + unreachable!("a model-info query runs no attempt") + } + async fn select( + &self, + _stage: &SelectStage, + _ctx: &RoutingContext<'_>, + candidates: Vec>, + ) -> Result>, Error> { + Ok(candidates) + } + async fn allow_fallback( + &self, + _request: &RouterRequest, + _target: &str, + ) -> Result { + Ok(true) + } + async fn sleep(&self, _duration: std::time::Duration) {} + fn observe(&self, _event: &RouterEvent) {} + fn cache(&self) -> Arc { + Arc::new(InMemoryRouterCache::default()) + } + async fn model_info(&self, lookup: &ModelInfoLookup) -> Result { + self.lookups + .lock() + .unwrap() + .push(serde_json::to_value(lookup).unwrap()); + Ok(match lookup { + ModelInfoLookup::ModelInfo { + model, tolerate, .. + } => match model.as_str().and_then(|model| self.rows.get(model)) { + Some(row) => row.clone(), + None if *tolerate => json!({"failed": true}), + None => return Err(raised("ValueError", "unmapped")), + }, + ModelInfoLookup::CostEntry { key, .. } => self + .rows + .get(key) + .map_or(Value::Null, |entry| json!({"entry": entry})), + ModelInfoLookup::Provider { .. } => { + json!([self.provider.0, self.provider.1]) + } + ModelInfoLookup::UserModelInfo => self.user_model_info.clone(), + _ => Value::Null, + }) + } + } + + #[test] + fn update_dictionary_converts_numbers_and_merges_nested_dicts_one_level() { + let mut existing = json!({"a": 1, "n": {"x": 1, "y": 2}, "k": "keep"}) + .as_object() + .unwrap() + .clone(); + let new = json!({"a": None::, "n": {"y": 3}, "i": "8_192", "f": "3e-07", "s": "abc"}) + .as_object() + .unwrap() + .clone(); + update_dictionary(&mut existing, &new); + assert_eq!( + Value::Object(existing), + json!({"a": 1, "n": {"x": 1, "y": 3}, "k": "keep", "i": 8192, "f": 3e-07, "s": "abc"}) + ); + } + + #[tokio::test] + async fn router_model_info_overlays_the_deployments_own_values_on_the_row() { + let host = Catalog { + rows: BTreeMap::from([( + "openai/gpt-4o".to_owned(), + json!({"key": "gpt-4o", "max_input_tokens": 100, "input_cost_per_token": 1.0}), + )]), + provider: ("gpt-4o".to_owned(), "openai".to_owned()), + user_model_info: json!({"id": "d", "input_cost_per_token": 2.0, "mode": None::}), + ..Catalog::default() + }; + let query = RouterModelInfoQuery { + deployment_id: json!("d"), + params_model: "openai/gpt-4o".to_owned(), + received_model_name: "g".to_owned(), + ..RouterModelInfoQuery::default() + }; + let merged = router_model_info(&host, &query).await.unwrap().unwrap(); + assert_eq!( + Value::Object(merged.info), + json!({"key": "gpt-4o", "max_input_tokens": 100, "input_cost_per_token": 2.0, "id": "d"}) + ); + assert_eq!(merged.user_keys, vec!["id", "input_cost_per_token"]); + } + + #[tokio::test] + async fn deployment_model_info_merges_the_entry_over_the_builtin_row() { + let host = Catalog { + rows: BTreeMap::from([ + ( + "gpt-4o".to_owned(), + json!({"key": "gpt-4o", "max_tokens": 10, "input_cost_per_token": 1.0}), + ), + ("d".to_owned(), json!({"max_tokens": "20"})), + ]), + ..Catalog::default() + }; + let info = deployment_model_info(&host, "d", &json!("gpt-4o")) + .await + .unwrap() + .unwrap(); + assert_eq!( + Value::Object(info), + json!({"key": "gpt-4o", "max_tokens": 20, "input_cost_per_token": 1.0}) + ); + } + + #[tokio::test] + async fn an_unpriced_entry_reports_no_token_prices() { + let host = Catalog { + rows: BTreeMap::from([("d".to_owned(), json!({"max_tokens": 5}))]), + ..Catalog::default() + }; + let info = deployment_model_info(&host, "d", &json!("unmapped")) + .await + .unwrap() + .unwrap(); + assert_eq!( + Value::Object(info), + json!({"max_tokens": 5, "input_cost_per_token": null, "output_cost_per_token": null}) + ); + } +} diff --git a/litellm-rust/crates/router-core/src/outcome.rs b/litellm-rust/crates/router-core/src/outcome.rs new file mode 100644 index 00000000000..59e0d14b2c2 --- /dev/null +++ b/litellm-rust/crates/router-core/src/outcome.rs @@ -0,0 +1,125 @@ +//! Which of a request's outcomes reach the deployment's health. The Python router learns them from +//! logging callbacks it registers on the process lists; `_core_serve_through_request_logging` and +//! `_core_discard` take those off, and `_core_bind_to_request_logging` puts them back on one +//! request's logging object. This core registers none, so the host answers, for each request, what +//! would have run ([`OutcomeCallbacks`]), and every write a callback owns passes through +//! [`OutcomeCallbacks::record`]. +//! +//! What the Python router writes outside its callbacks is not gated, and reaches [`crate::Health`] +//! with the default value: `make_call`'s usage, `async_routing_strategy_pre_call_checks`' +//! cooldown, `_trigger_cooldown_for_failed_deployment`, `RouterHealth.record_failure` and a +//! realtime session's success. + +use std::future::Future; + +use serde::Deserialize; + +use crate::error::Error; + +/// Which of the router's outcome callbacks litellm's logging runs for a request. The default is a +/// router whose callbacks are on the process lists. +#[derive(Clone, Copy, Debug, Deserialize, PartialEq, Eq)] +#[serde(default, deny_unknown_fields)] +pub struct OutcomeCallbacks { + /// `deployment_callback_on_success`, `sync_deployment_callback_on_success` for a sync call. + /// `_core_bind_to_request_logging` binds only the former. + pub success: bool, + /// `deployment_callback_on_failure` and `async_deployment_callback_on_failure`. + pub failure: bool, + /// The attempt runs through `make_call`, which counts its usage whatever the callbacks do. + pub make_call: bool, +} + +impl Default for OutcomeCallbacks { + fn default() -> Self { + Self { + success: true, + failure: true, + make_call: true, + } + } +} + +/// What an outcome writes. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum OutcomeWrite { + /// The success counter, or the usage the success callback adds at the end of a stream. + Success, + /// The tokens and request of an attempt that returned: `make_call`'s write, or the success + /// callback's for an endpoint that skips `make_call`. + AttemptUsage, + /// The failure counter, the allowed-fails counter, the cooldown and the failed request. + Failure, +} + +impl OutcomeCallbacks { + /// Run `record` when `write` happens for this request, else record nothing. + pub(crate) async fn record( + self, + write: OutcomeWrite, + record: impl Future>, + ) -> Result { + let runs = match write { + OutcomeWrite::Success => self.success, + OutcomeWrite::AttemptUsage => self.success || self.make_call, + OutcomeWrite::Failure => self.failure, + }; + if runs { record.await } else { Ok(T::default()) } + } +} + +#[cfg(test)] +mod tests { + use rstest::rstest; + + use super::*; + + const REGISTERED: OutcomeCallbacks = OutcomeCallbacks { + success: true, + failure: true, + make_call: true, + }; + const UNBOUND: OutcomeCallbacks = OutcomeCallbacks { + success: false, + failure: false, + make_call: true, + }; + const UNBOUND_WITHOUT_MAKE_CALL: OutcomeCallbacks = OutcomeCallbacks { + success: false, + failure: false, + make_call: false, + }; + const BOUND_SYNC: OutcomeCallbacks = OutcomeCallbacks { + success: false, + failure: true, + make_call: true, + }; + + #[rstest] + #[case::registered_success(REGISTERED, OutcomeWrite::Success, true)] + #[case::registered_usage(REGISTERED, OutcomeWrite::AttemptUsage, true)] + #[case::registered_failure(REGISTERED, OutcomeWrite::Failure, true)] + #[case::unbound_success(UNBOUND, OutcomeWrite::Success, false)] + #[case::unbound_usage(UNBOUND, OutcomeWrite::AttemptUsage, true)] + #[case::unbound_failure(UNBOUND, OutcomeWrite::Failure, false)] + #[case::no_make_call_usage(UNBOUND_WITHOUT_MAKE_CALL, OutcomeWrite::AttemptUsage, false)] + #[case::bound_sync_success(BOUND_SYNC, OutcomeWrite::Success, false)] + #[case::bound_sync_failure(BOUND_SYNC, OutcomeWrite::Failure, true)] + #[tokio::test] + async fn a_write_runs_when_its_callback_does( + #[case] callbacks: OutcomeCallbacks, + #[case] write: OutcomeWrite, + #[case] runs: bool, + ) { + let recorded = callbacks.record(write, async { Ok(true) }).await; + + assert_eq!(recorded.ok(), Some(runs)); + } + + #[test] + fn a_router_on_the_process_lists_is_the_default() { + let parsed: OutcomeCallbacks = serde_json::from_str("{}").expect("an empty document"); + + assert_eq!(parsed, REGISTERED); + } +} diff --git a/litellm-rust/crates/router-core/src/precall.rs b/litellm-rust/crates/router-core/src/precall.rs new file mode 100644 index 00000000000..04dd22f1c90 --- /dev/null +++ b/litellm-rust/crates/router-core/src/precall.rs @@ -0,0 +1,300 @@ +//! `enable_pre_call_checks` (`_pre_call_checks`): before a pick, drop the deployments whose +//! context window the prompt exceeds, those outside the request's `allowed_model_region`, and +//! those that cannot take a requested param, then decide which error a group with nothing left +//! raises. The facts that need live model data (the prompt's token count, each deployment's +//! context window, which params it takes) come from [`Host::pre_call_facts`]; every decision is +//! made here. +//! +//! [`Host::pre_call_facts`]: crate::host::Host::pre_call_facts + +use std::{collections::BTreeMap, sync::Arc}; + +use serde::Deserialize; +use serde_json::Value; + +use crate::{error::Error, registry::Deployment}; + +/// What the host knows about one deployment for the pre-call checks. +#[derive(Clone, Debug, Default, Deserialize, PartialEq)] +#[serde(default)] +pub struct DeploymentFacts { + /// The deployment's context window (`get_router_model_info`'s `max_input_tokens`), when it + /// declares one and resolves. + pub max_input_tokens: Option, + /// `_deployment_lacks_params`: it does not support one of the requested params. + pub lacks_params: bool, + /// `_deployment_rejects_params`: it lacks one the request and deployment neither drop nor allow. + pub rejects_params: bool, +} + +/// What the host knows about a request for the pre-call checks. +#[derive(Clone, Debug, Default, Deserialize, PartialEq)] +#[serde(default)] +pub struct PreCallFacts { + /// The request carries `messages` or `input`: without either no check runs. + pub has_input: bool, + /// The prompt may be counted against a context window (no compaction is pending). + pub countable: bool, + /// The prompt's token count, when it was counted. + pub input_tokens: Option, + /// The request's `allowed_model_region`. + pub allowed_model_region: Option, + /// The requested params the check filters on, sorted (`_pre_call_check_requested_chat_params`). + pub requested_params: Vec, + /// Facts by deployment id, for the candidates and the group's deployments. + pub deployments: BTreeMap, +} + +impl PreCallFacts { + fn of(&self, deployment: &Deployment) -> DeploymentFacts { + self.deployments + .get(&deployment.id) + .cloned() + .unwrap_or_default() + } +} + +/// The name a context-window error gives a deployment: its base model, else its model. +fn deployment_model(deployment: &Deployment) -> String { + let truthy = |value: Option<&Value>| match value { + Some(Value::String(text)) if !text.is_empty() => Some(text.clone()), + Some(Value::Null | Value::Bool(false)) | None => None, + Some(Value::String(_)) => None, + Some(other) => Some(python_str(other)), + }; + truthy(deployment.model_info.raw.get("base_model")) + .or_else(|| truthy(deployment.raw_params.get("base_model"))) + .unwrap_or_else(|| deployment.params.model.clone()) +} + +/// `str(value)` for the scalar values a base model may hold. +fn python_str(value: &Value) -> String { + match value { + Value::Bool(true) => "True".to_owned(), + Value::Number(number) => number.to_string(), + other => other.to_string(), + } +} + +/// `is_region_allowed`: the deployment's `region_name` equals the allowed region. +fn in_region(deployment: &Deployment, allowed: &Value) -> bool { + deployment.raw_params.get("region_name") == Some(allowed) +} + +/// `_pre_call_checks` over `candidates`; `pool` is the group's deployments before health filters +/// (`request_deployments`), which decide whether nothing left is an unsupported-params error. +pub fn filter( + model: &str, + candidates: Vec>, + pool: &[Arc], + facts: &PreCallFacts, +) -> Result>, Error> { + if !facts.has_input { + return Ok(candidates); + } + let mut invalid = vec![false; candidates.len()]; + let mut context_window_error = false; + let mut detail = String::new(); + for (index, deployment) in candidates.iter().enumerate() { + let found = facts.of(deployment); + if let (Some(limit), true, Some(tokens)) = + (found.max_input_tokens, facts.countable, facts.input_tokens) + && tokens > limit + { + invalid[index] = true; + context_window_error = true; + detail.push_str(&format!( + "Model={}, Max Input Tokens={limit}, Got={tokens}", + deployment_model(deployment) + )); + continue; + } + if let Some(allowed) = &facts.allowed_model_region + && !allowed.is_null() + && !in_region(deployment, allowed) + { + invalid[index] = true; + } + } + let params_checked = !facts.requested_params.is_empty(); + let param_candidates: Vec = (0..candidates.len()).filter(|i| !invalid[*i]).collect(); + let lacking: Vec = param_candidates + .iter() + .map(|i| params_checked && facts.of(&candidates[*i]).lacks_params) + .collect(); + let none_takes = !lacking.is_empty() && lacking.iter().all(|lacks| *lacks); + let failing: Vec = if none_takes { + param_candidates + .iter() + .map(|i| params_checked && facts.of(&candidates[*i]).rejects_params) + .collect() + } else { + lacking + }; + for (index, fails) in param_candidates.iter().zip(failing) { + if fails { + invalid[*index] = true; + } + } + let excluded = invalid.iter().filter(|x| **x).count(); + if excluded == candidates.len() { + if !context_window_error + && !pool.is_empty() + && pool + .iter() + .all(|d| params_checked && facts.of(d).rejects_params) + { + return Err(Error::UnsupportedParams { + model: model.to_owned(), + params: facts.requested_params.clone(), + }); + } + if context_window_error { + return Err(Error::ContextWindowExceeded { + model: model.to_owned(), + detail, + }); + } + } + Ok(candidates + .into_iter() + .zip(invalid) + .filter_map(|(deployment, invalid)| (!invalid).then_some(deployment)) + .collect()) +} + +#[cfg(test)] +mod tests { + use rstest::rstest; + use serde_json::json; + + use super::*; + use crate::registry::DeploymentSpec; + + fn deployment(id: &str, params: Value, model_info: Value) -> Arc { + let mut info = model_info; + info["id"] = json!(id); + let spec: DeploymentSpec = serde_json::from_value(json!({ + "model_name": "g", "litellm_params": params, "model_info": info, + })) + .unwrap(); + Arc::new(Deployment::from_spec(spec).unwrap()) + } + + fn facts(entries: &[(&str, DeploymentFacts)]) -> PreCallFacts { + PreCallFacts { + has_input: true, + countable: true, + deployments: entries + .iter() + .map(|(id, f)| ((*id).to_owned(), f.clone())) + .collect(), + ..PreCallFacts::default() + } + } + + fn window(limit: i64) -> DeploymentFacts { + DeploymentFacts { + max_input_tokens: Some(limit), + ..DeploymentFacts::default() + } + } + + #[rstest] + fn a_prompt_over_a_context_window_drops_that_deployment() { + let small = deployment("small", json!({"model": "openai/a"}), json!({})); + let big = deployment("big", json!({"model": "openai/b"}), json!({})); + let mut found = facts(&[("small", window(5)), ("big", window(10_000))]); + found.input_tokens = Some(100); + let kept = filter("g", vec![small, big], &[], &found).unwrap(); + assert_eq!( + kept.iter().map(|d| d.id.as_str()).collect::>(), + ["big"] + ); + } + + #[rstest] + fn no_window_left_names_each_deployment_back_to_back() { + let one = deployment("one", json!({"model": "openai/one"}), json!({})); + let two = deployment( + "two", + json!({"model": "openai/two"}), + json!({"base_model": "azure/base"}), + ); + let mut found = facts(&[("one", window(5)), ("two", window(6))]); + found.input_tokens = Some(9); + let error = filter("g", vec![one, two], &[], &found).unwrap_err(); + let Error::ContextWindowExceeded { model, detail } = error else { + panic!("expected a context window error, got {error:?}"); + }; + assert_eq!(model, "g"); + assert_eq!( + detail, + "Model=openai/one, Max Input Tokens=5, Got=9Model=azure/base, Max Input Tokens=6, Got=9" + ); + } + + #[rstest] + fn without_input_nothing_is_checked() { + let small = deployment("small", json!({"model": "openai/a"}), json!({})); + let mut found = facts(&[("small", window(5))]); + found.input_tokens = Some(100); + found.has_input = false; + assert_eq!(filter("g", vec![small], &[], &found).unwrap().len(), 1); + } + + #[rstest] + #[case::in_region("eu", 1)] + #[case::out_of_region("ap", 0)] + fn the_allowed_region_keeps_only_its_deployments(#[case] region: &str, #[case] kept: usize) { + let eu = deployment( + "eu", + json!({"model": "openai/a", "region_name": "eu"}), + json!({}), + ); + let none = deployment("none", json!({"model": "openai/a"}), json!({})); + let mut found = facts(&[]); + found.allowed_model_region = Some(json!(region)); + assert_eq!( + filter("g", vec![eu, none], &[], &found).unwrap().len(), + kept + ); + } + + #[rstest] + fn a_group_where_every_deployment_rejects_the_param_is_unsupported() { + let old = deployment("old", json!({"model": "openai/old"}), json!({})); + let mut found = facts(&[( + "old", + DeploymentFacts { + lacks_params: true, + rejects_params: true, + ..DeploymentFacts::default() + }, + )]); + found.requested_params = vec!["response_format".to_owned()]; + let error = filter("g", vec![old.clone()], &[old], &found).unwrap_err(); + assert!( + matches!(error, Error::UnsupportedParams { params, .. } if params == ["response_format"]) + ); + } + + #[rstest] + fn a_deployment_that_takes_the_param_wins_over_one_that_would_drop_it() { + let old = deployment("old", json!({"model": "openai/old"}), json!({})); + let new = deployment("new", json!({"model": "openai/new"}), json!({})); + let mut found = facts(&[( + "old", + DeploymentFacts { + lacks_params: true, + ..DeploymentFacts::default() + }, + )]); + found.requested_params = vec!["response_format".to_owned()]; + let kept = filter("g", vec![old, new], &[], &found).unwrap(); + assert_eq!( + kept.iter().map(|d| d.id.as_str()).collect::>(), + ["new"] + ); + } +} diff --git a/litellm-rust/crates/router-core/src/prerouting.rs b/litellm-rust/crates/router-core/src/prerouting.rs new file mode 100644 index 00000000000..d6e3df05e7a --- /dev/null +++ b/litellm-rust/crates/router-core/src/prerouting.rs @@ -0,0 +1,53 @@ +//! Pre-routing: before an async pick resolves its name, the host may route the request to another +//! model group (Python's `_core_async_pre_routing_hook`). The complexity, auto, adaptive and +//! quality routers registered from `auto_router/...` marker deployments, and the Claude Code +//! session binding, live host-side; the core asks once per pick, before resolution, and resolves +//! the name the host answers. The attempt still serves the requested name (its `model_group`): +//! only the deployments come from the routed group. +//! +//! A custom routing strategy (`register_plugin`, `set_custom_routing_strategy`, lar1) replaces the +//! whole pick (`_core_async_get_available_deployment` and `get_available_deployment`): the host +//! answers the same step with the deployment document the strategy returned, which may name a +//! deployment the registry does not hold, and the core serves that deployment as is. + +use std::sync::Arc; + +use crate::{Core, config::Op, error::Error, host::RouterRequest, registry::Deployment}; + +/// What the host's pre-routing step sees. +#[derive(Clone, Copy, Debug)] +pub struct PreRoute<'a> { + pub op: &'a Op, + pub request: &'a RouterRequest, + /// The name this pick resolves: the requested name or the fallback target. + pub name: &'a str, + /// How many fallbacks deep the chain is (Python's `fallback_depth`): only a first attempt + /// binds a Claude Code session. + pub fallback_depth: u32, + /// 0 for the first attempt in this group, then the retry number. + pub retry: u32, +} + +/// The host's answer to a pre-routing step. +#[derive(Clone, Debug)] +pub enum PreRouted { + /// Resolve this model group instead of the requested name. + Model(String), + /// A custom strategy's pick: serve this deployment, without resolving or filtering. + Deployment(Arc), +} + +impl Core { + /// What a pick of `ask.name` serves: the host's pre-routed group or custom pick, else the + /// name itself. Only picks that ask for it pre-route (`request.pre_routes`). + pub(crate) async fn pre_routed(&self, ask: PreRoute<'_>) -> Result { + if !ask.request.pre_routes { + return Ok(PreRouted::Model(ask.name.to_owned())); + } + Ok(self + .host + .pre_route(&ask) + .await? + .unwrap_or_else(|| PreRouted::Model(ask.name.to_owned()))) + } +} diff --git a/litellm-rust/crates/router-core/src/ptu.rs b/litellm-rust/crates/router-core/src/ptu.rs new file mode 100644 index 00000000000..9110ac6a67f --- /dev/null +++ b/litellm-rust/crates/router-core/src/ptu.rs @@ -0,0 +1,96 @@ +//! PTU shares (`litellm_core_utils/ptu_shares.py`): a deployment whose `model_info.ptu_shares` +//! splits it across teams serves only those teams, and the teams holding a share are kept to +//! their share. + +use std::{collections::BTreeSet, sync::Arc}; + +use serde_json::{Map, Value}; + +use crate::registry::Deployment; + +/// `parsed_ptu_shares`: the teams of a non-empty `ptu_shares` map whose every entry is a team id +/// and a positive whole number of PTUs (`2.0` reads as 2, a bool is no number), else None. +pub fn share_teams(model_info: &Map) -> Option> { + let shares = model_info.get("ptu_shares")?.as_object()?; + if shares.is_empty() { + return None; + } + shares + .iter() + .map(|(team, share)| (!team.is_empty() && whole_ptus(share)).then(|| team.clone())) + .collect() +} + +/// `_whole_ptus`: a positive whole number. +fn whole_ptus(share: &Value) -> bool { + match share { + Value::Number(number) => number.as_i64().map_or_else( + || { + number + .as_u64() + .map(|whole| whole > 0) + .or_else(|| number.as_f64().map(|f| f.fract() == 0.0 && f > 0.0)) + .unwrap_or(false) + }, + |whole| whole > 0, + ), + _ => false, + } +} + +/// `filter_ptu_shared_deployments`: the deployments `team_id` holds a share on when it holds one, +/// else the ones split into no shares; a caller without a team holds none. Also whether any +/// deployment was withheld. +pub fn filter_shared( + deployments: Vec>, + team_id: Option<&str>, +) -> (Vec>, bool) { + let team = team_id.unwrap_or(""); + let checks: Vec<(Arc, Option>)> = deployments + .into_iter() + .map(|deployment| { + let shares = share_teams(&deployment.model_info.raw); + (deployment, shares) + }) + .collect(); + let holds_share = checks + .iter() + .any(|(_, shares)| shares.as_ref().is_some_and(|teams| teams.contains(team))); + let total = checks.len(); + let kept: Vec> = checks + .into_iter() + .filter(|(_, shares)| match shares { + Some(teams) => teams.contains(team), + None => !holds_share, + }) + .map(|(deployment, _)| deployment) + .collect(); + let withheld = kept.len() < total; + (kept, withheld) +} + +#[cfg(test)] +mod tests { + use rstest::rstest; + use serde_json::json; + + use super::*; + + #[rstest] + #[case(json!({"TA": 50}), Some(vec!["TA"]))] + #[case(json!({"TA": 2.0, "TB": 1}), Some(vec!["TA", "TB"]))] + #[case(json!({"TA": 2.5}), None)] + #[case(json!({"TA": true}), None)] + #[case(json!({"TA": 0}), None)] + #[case(json!({}), None)] + #[case(json!({"": 1}), None)] + #[case(json!("TA"), None)] + fn reads_the_share_teams(#[case] shares: Value, #[case] expected: Option>) { + let mut model_info = Map::new(); + model_info.insert("ptu_shares".to_owned(), shares); + assert_eq!( + share_teams(&model_info), + expected.map(|teams| teams.into_iter().map(str::to_owned).collect()) + ); + } +} diff --git a/litellm-rust/crates/router-core/src/registry.rs b/litellm-rust/crates/router-core/src/registry.rs new file mode 100644 index 00000000000..78c7d3b67e3 --- /dev/null +++ b/litellm-rust/crates/router-core/src/registry.rs @@ -0,0 +1,1548 @@ +//! The configured deployments, indexed by id, model group, alias, wildcard pattern, routing +//! group and team. A [`Registry`] is immutable; [`Registry::apply`] builds the next version. + +use std::{ + collections::{BTreeMap, BTreeSet, HashMap}, + fmt::Write as _, + sync::Arc, +}; + +use litellm_router_types::{LitellmParams, Spelled}; +use serde::Deserialize; +use serde_json::{Map, Value}; +use sha2::{Digest, Sha256}; + +use crate::{ + error::ConfigError, + host_index::{HostIndexes, IndexWrite, cache_flags_left_to_the_catalog, deployment_model_name}, +}; + +/// A registry's version; every change produces a higher one. +pub type RegistryVersion = u64; + +/// A deployment as a config or an `apply` spells it. +#[derive(Clone, Debug, Default, PartialEq, Deserialize)] +pub struct DeploymentSpec { + pub model_name: String, + pub litellm_params: Map, + #[serde(default)] + pub model_info: Map, + /// The entry's other top-level keys (a config's `tpm`, ...), kept so the host's view of the + /// deployment list reads back the entries it stored. + #[serde(flatten)] + pub extra: Map, + /// The provider the host resolved when it built the deployment ([`crate::build`]); never + /// read from a config or an `apply`. + #[serde(skip)] + pub provider: Option, +} + +/// `routing_groups`: a named subset of model groups sharing a routing strategy. +#[derive(Clone, Debug, PartialEq, Deserialize)] +pub struct RoutingGroupSpec { + pub group_name: String, + pub models: Vec, + pub routing_strategy: crate::config::RoutingStrategy, + /// A priority group's member priorities, served as each member's `order`. + #[serde(default)] + pub model_priorities: Option>, +} + +/// The `model_info` fields the router reads; every other key stays in [`Self::raw`]. +#[derive(Clone, Debug, Default, PartialEq)] +pub struct ModelInfo { + pub team_id: Option, + pub team_public_model_name: Option, + pub access_groups: Vec, + pub blocked: bool, + pub cooldown_time: Option, + pub allowed_fails: Option, + /// `model_info.allowed_fails_policy` when it is a dict, read as Python reads it: field name + /// to allowed fails, values as given. + pub allowed_fails_policy: Option>, + pub raw: Map, +} + +impl ModelInfo { + fn from_raw(raw: Map) -> Self { + let text = |key: &str| raw.get(key).and_then(Value::as_str).map(str::to_owned); + Self { + team_id: text("team_id"), + team_public_model_name: text("team_public_model_name"), + access_groups: raw + .get("access_groups") + .and_then(Value::as_array) + .map(|groups| { + groups + .iter() + .filter_map(Value::as_str) + .map(str::to_owned) + .collect() + }) + .unwrap_or_default(), + blocked: raw.get("blocked").and_then(Value::as_bool) == Some(true), + cooldown_time: raw.get("cooldown_time").and_then(Value::as_f64), + allowed_fails: raw + .get("allowed_fails") + .and_then(Value::as_u64) + .and_then(|fails| u32::try_from(fails).ok()), + allowed_fails_policy: raw + .get("allowed_fails_policy") + .and_then(Value::as_object) + .cloned(), + raw, + } + } +} + +/// One configured deployment. +#[derive(Clone, Debug, PartialEq)] +pub struct Deployment { + pub id: String, + pub model_name: String, + pub params: LitellmParams, + /// `litellm_params` as written, in order: what the id hashes and the host receives. + pub raw_params: Map, + pub model_info: ModelInfo, + /// The entry's other top-level keys, as written. + pub extra: Map, + /// The provider the host resolved for it, when it was built through the host. + pub resolved_provider: Option, +} + +impl Deployment { + pub fn from_spec(spec: DeploymentSpec) -> Result { + let params: LitellmParams = serde_json::from_value(Value::Object( + spec.litellm_params.clone(), + )) + .map_err(|source| ConfigError::InvalidDeployment { + model_name: spec.model_name.clone(), + source, + })?; + let id = match spec.model_info.get("id") { + Some(Value::String(id)) => id.clone(), + Some(Value::Number(id)) => id.to_string(), + _ => deployment_id(&spec.model_name, &spec.litellm_params), + }; + Ok(Self { + id, + model_name: spec.model_name, + params, + raw_params: spec.litellm_params, + model_info: ModelInfo::from_raw(spec.model_info), + extra: spec.extra, + resolved_provider: spec.provider, + }) + } + + /// The deployment as a `model_list` entry: its other top-level keys, then `model_name`, + /// `litellm_params` and `model_info` (with its id). + pub fn record(&self) -> Value { + let mut model_info = self.model_info.raw.clone(); + model_info.insert("id".to_owned(), Value::String(self.id.clone())); + let mut record = self.extra.clone(); + record.insert( + "model_name".to_owned(), + Value::String(self.model_name.clone()), + ); + record.insert( + "litellm_params".to_owned(), + Value::Object(self.raw_params.clone()), + ); + record.insert("model_info".to_owned(), Value::Object(model_info)); + Value::Object(record) + } + + /// The stand-in an unrouted attempt carries: the requested name as its group and model, no id. + pub fn unrouted(name: &str) -> Result { + let mut params = Map::new(); + params.insert("model".to_owned(), Value::String(name.to_owned())); + let mut info = Map::new(); + info.insert("id".to_owned(), Value::String(String::new())); + Self::from_spec(DeploymentSpec { + model_name: name.to_owned(), + litellm_params: params, + model_info: info, + extra: Map::new(), + provider: None, + }) + } + + /// The provider: the one the host resolved when it built the deployment, else + /// `custom_llm_provider`, else the `model` prefix before `/`. + pub fn provider(&self) -> &str { + self.resolved_provider + .as_deref() + .or(self.params.custom_llm_provider.as_deref()) + .or_else(|| { + self.params + .model + .split_once('/') + .map(|(provider, _)| provider) + }) + .unwrap_or("openai") + } + + /// A numeric `litellm_params` value, as `simple_shuffle._metric_weight` reads it. + pub fn metric(&self, key: &str) -> Option { + let spelled = |value: &Option>| match value { + Some(Spelled::Value(value)) => Some(*value), + Some(Spelled::Text(_)) | None => None, + }; + match key { + "rpm" => spelled(&self.params.rpm), + "tpm" => spelled(&self.params.tpm), + _ => self.raw_params.get(key).and_then(Value::as_f64), + } + } + + /// `simple_shuffle._metric_weight`: 0 when unset, a number (a bool counts as 0 or 1), or + /// `None` when the value is not numeric. A `rpm`/`tpm` spelled as text reads as unset. + pub fn metric_weight(&self, key: &str) -> Option { + if key != "weight" { + return Some(self.metric(key).unwrap_or(0.0)); + } + match self.raw_params.get(key) { + None | Some(Value::Null) => Some(0.0), + Some(Value::Bool(flag)) => Some(f64::from(u8::from(*flag))), + Some(Value::Number(number)) => number.as_f64(), + Some(_) => None, + } + } + + /// `get_deployment_order`: `litellm_params.order`, else `model_info.order`, which is where a + /// team model added through the API keeps it. Neither is validated as an int, so a bool, a + /// whole float and a numeric string are read as the level they compare as. + pub fn order(&self) -> Option { + let order = self + .raw_params + .get("order") + .filter(|order| !order.is_null()) + .or_else(|| self.model_info.raw.get("order"))?; + match order { + Value::Bool(flag) => Some(i64::from(*flag)), + Value::Number(number) => number.as_i64().or_else(|| { + number + .as_f64() + .filter(|float| float.fract() == 0.0) + .map(|float| float as i64) + }), + Value::String(text) => text.parse().ok(), + Value::Null | Value::Array(_) | Value::Object(_) => None, + } + } + + /// `_deployment_supports_web_search`: `model_info.supports_web_search` when it is set, read + /// for truthiness, else true. + pub(crate) fn supports_web_search(&self) -> bool { + self.model_info + .raw + .get("supports_web_search") + .is_none_or(|flag| flag.is_null() || crate::host_index::truthy(flag)) + } + + /// Whether `team_id` may use this deployment (`deployment_usable_by_team`). + pub fn usable_by_team(&self, team_id: Option<&str>) -> bool { + self.model_info + .team_id + .as_deref() + .is_none_or(|owner| Some(owner) == team_id) + } + + /// Whether `_core_get_model_list(team_id=team_id)` lists this deployment: no team lists + /// every team's (`_core_should_include_deployment`). + pub fn listed_for_team(&self, team_id: Option<&str>) -> bool { + team_id.is_none() || self.usable_by_team(team_id) + } + + /// Whether the deployment declares a limit the router counts per-minute usage for + /// (`_deployment_tracks_usage` and `deployment_has_io_token_limits`, without the limits only + /// the model catalog knows, which the host decides by): `tpm` or `rpm`, `itpm` or `otpm`, on + /// the entry itself, in `litellm_params` or in `model_info`. + pub fn tracks_usage(&self) -> bool { + let declared = |section: &Map| { + ["tpm", "rpm", "itpm", "otpm"] + .iter() + .any(|key| section.get(*key).is_some_and(|limit| !limit.is_null())) + }; + self.params.tpm.is_some() + || self.params.rpm.is_some() + || self.params.itpm.is_some() + || self.params.otpm.is_some() + || declared(&self.extra) + || declared(&self.model_info.raw) + } +} + +/// `_return_pattern_matched_deployments` for one deployment: a copy serving `request`, its +/// `litellm_params.model` filled with the `groups` the pattern matched and the prompt-cache flags +/// left to the catalog. A deployment whose model holds no wildcard is served as it is. +pub(crate) fn routed_deployment( + deployment: &Arc, + groups: &[String], + request: &str, +) -> Arc { + let model = &deployment.params.model; + if !model.contains('*') { + return deployment.clone(); + } + let named = deployment_model_name(groups, request, model); + let mut routed = Deployment::clone(deployment); + for flag in cache_flags_left_to_the_catalog(&deployment.raw_params) { + routed.model_info.raw.remove(*flag); + } + routed + .raw_params + .insert("model".to_owned(), Value::String(named.clone())); + routed.params.model = named; + Arc::new(routed) +} + +/// The connection params of one deployment (`credentials(id)`). +#[derive(Clone, Debug, PartialEq)] +pub struct Credentials { + pub provider: String, + pub api_base: Option, + pub api_version: Option, + pub credential_name: Option, + pub params: LitellmParams, +} + +/// What a name resolved through. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum ResolveKind { + SpecificDeployment, + Id, + /// A name nothing else serves, routed to the router's `default_deployment` (a copy whose + /// `litellm_params.model` is the name). Served like a deployment named by its id. + Default, + Alias, + RoutingGroup, + Group, + Wildcard, +} + +impl ResolveKind { + /// The name resolved to one deployment that is served as is: a deployment id, or the + /// default deployment (`_common_checks_available_deployment` returns a dict for both). + pub fn is_single(&self) -> bool { + matches!(self, Self::Id | Self::Default) + } +} + +#[derive(Clone, Debug, PartialEq)] +pub struct Resolution { + pub kind: ResolveKind, + /// The model group the deployments are routed as. + pub group: String, + pub deployments: Vec>, +} + +#[derive(Clone, Copy, Debug, Default)] +pub struct ResolveScope { + /// The caller named a deployment's `litellm_params.model` (`specific_deployment=True`). + pub specific_deployment: bool, +} + +#[derive(Clone, Debug, Default)] +pub struct DeploymentFilter<'a> { + pub model_name: Option<&'a str>, + pub team_id: Option<&'a str>, + pub provider: Option<&'a str>, +} + +#[derive(Clone, Debug, Default)] +pub struct NameFilter { + pub include_aliases: bool, +} + +/// The registry indexes as the host wrote them over the ones the deployments build (the +/// deprecated `model_names`, `model_name_to_deployment_indices` and +/// `team_model_to_deployment_indices` writes), which resolution reads instead from then on. +/// Indices point into the deployment list; one past its end is skipped. +#[derive(Clone, Debug, Default, PartialEq, Deserialize)] +#[serde(default, deny_unknown_fields)] +pub struct IndexOverrides { + /// `_core_model_names`: the names `name in model_names` finds. + pub model_names: Option>, + /// `_core_model_name_to_deployment_indices`: a model name's deployments. + pub model_name: Option>>, + /// `_core_team_model_to_deployment_indices`, as `(team_id, public name, indices)` entries. + pub team_model: Option)>>, +} + +impl IndexOverrides { + fn is_empty(&self) -> bool { + self.model_names.is_none() && self.model_name.is_none() && self.team_model.is_none() + } +} + +/// One atomic change: optionally a whole new deployment list, then upserts, deletes and, +/// optionally, a new dynamic config. +#[derive(Clone, Debug, Default)] +pub struct ChangeSet { + /// Every deployment, in order, replacing the current list (the host registry's `model_list` + /// after a write; it may hold one id more than once, as the Python router's list can). + pub replace: Option>, + /// Deployments added at the end, beside any that share their id (a deployment the host + /// retries after its provider resolved). + pub appends: Vec, + pub upserts: Vec, + pub deletes: Vec, + pub settings: Option, + /// The indexes the host wrote ([`IndexOverrides`]), replacing the ones held; an empty one + /// goes back to the indexes the deployments build. + pub index: Option, + /// The deployment at the position the id map holds for this id, removed with every index + /// shifted (`_delete_deployment`, and an upsert's removal of the deployment it edits). + pub remove: Option, + /// A host write to the indexes it hands out ([`HostIndexes`]). + pub host_index: Option, +} + +#[derive(Clone, Debug, Default)] +pub struct Registry { + version: RegistryVersion, + /// The version the deployments before the newest appends last changed at: a host that read + /// the list at or after it reads only what was appended since ([`Self::deployments_from`]). + rewritten: RegistryVersion, + deployments: Vec>, + by_id: HashMap, + by_group: BTreeMap>, + by_team: HashMap>, + wildcards: Vec<(String, Vec)>, + aliases: BTreeMap, + routing_groups: Vec, + overrides: Option>, + /// The indexes the host hands out, kept by the Python router's history rules. + host: Arc, + /// `_model_to_group` as Python last built it, when a raw `routing_groups` write replaced the + /// groups since: member calls keep routing by it (see [`Self::member_routing_group`]). + member_groups: Option>>, +} + +impl Registry { + pub fn build( + specs: Vec, + aliases: BTreeMap, + routing_groups: Vec, + ) -> Result { + // One id may name several deployments (a per-organization fan-out, or two entries + // declaring one id): each is kept, and the id finds the last, as the Python router keeps them. + let deployments = specs + .into_iter() + .map(|spec| Deployment::from_spec(spec).map(Arc::new)) + .collect::, _>>()?; + // The indexes `set_model_list` leaves, as a build installing each deployment in turn does. + let mut host = HostIndexes::default(); + for (position, deployment) in deployments.iter().enumerate() { + let record = deployment.record(); + host.added( + record.as_object().expect("a record is an object"), + position, + false, + ); + } + host.set_model_names(deployments.iter().map(|d| d.model_name.clone()).collect()); + Ok(Self { + host: Arc::new(host), + ..Self::index(0, deployments, aliases, routing_groups) + }) + } + + fn index( + version: RegistryVersion, + deployments: Vec>, + aliases: BTreeMap, + routing_groups: Vec, + ) -> Self { + let mut by_id = HashMap::new(); + let mut by_group: BTreeMap> = BTreeMap::new(); + let mut by_team: HashMap> = HashMap::new(); + for (index, deployment) in deployments.iter().enumerate() { + by_id.insert(deployment.id.clone(), index); + by_group + .entry(deployment.model_name.clone()) + .or_default() + .push(index); + if let Some(team) = &deployment.model_info.team_id { + by_team.entry(team.clone()).or_default().push(index); + } + } + // In the order each pattern first appears, then most specific first, as + // `PatternUtils.sorted_patterns` keeps them: a stable sort, longest regex first, + // then the most regex syntax. + let mut wildcards: Vec<(String, Vec)> = Vec::new(); + for deployment in &deployments { + let name = &deployment.model_name; + if name.contains('*') && !wildcards.iter().any(|(pattern, _)| pattern == name) { + wildcards.push(( + name.clone(), + by_group.get(name).cloned().unwrap_or_default(), + )); + } + } + wildcards.sort_by_key(|(pattern, _)| std::cmp::Reverse(pattern_specificity(pattern))); + Self { + version, + rewritten: version, + deployments, + by_id, + by_group, + by_team, + wildcards, + aliases, + routing_groups, + overrides: None, + host: Arc::default(), + member_groups: None, + } + } + + /// This registry with the stale routing-group member map (`None` derives it from the groups). + #[must_use] + pub fn with_member_groups(self, member_groups: Option>) -> Self { + Self { + member_groups: member_groups.map(Arc::new), + ..self + } + } + + pub fn version(&self) -> RegistryVersion { + self.version + } + + /// The version the deployments before the newest appends last changed at (see `rewritten`). + pub fn rewritten(&self) -> RegistryVersion { + self.rewritten + } + + /// Every deployment id, in list order (`names(kind="ids")`). + pub fn ids(&self) -> Vec { + self.deployments.iter().map(|d| d.id.clone()).collect() + } + + /// The deployments from position `start` on, in list order. + pub fn deployments_from(&self, start: usize) -> &[Arc] { + self.deployments.get(start..).unwrap_or_default() + } + + /// The next registry: upserts, then deletes (`apply_changes`). An upserted deployment whose + /// id exists is removed and appended at the end, as the Python router does on an edit. + pub fn apply( + &self, + changes: &ChangeSet, + aliases: Option>, + ) -> Result { + let mut deployments = match &changes.replace { + Some(specs) => specs + .iter() + .map(|spec| self.carried(spec.clone()).map(Arc::new)) + .collect::, _>>()?, + None => self.deployments.clone(), + }; + for spec in &changes.appends { + deployments.push(Arc::new(self.carried(spec.clone())?)); + } + for spec in &changes.upserts { + let deployment = self.carried(spec.clone())?; + deployments.retain(|existing| existing.id != deployment.id); + deployments.push(Arc::new(deployment)); + } + deployments.retain(|deployment| !changes.deletes.contains(&deployment.id)); + let mut host = HostIndexes::clone(&self.host); + if let Some(id) = &changes.remove + && let Some(position) = host.position_of(id) + && position < deployments.len() + { + deployments.remove(position); + host.removed(id, position) + .map_err(|error| ConfigError::Removal(error.to_string()))?; + } + if let Some(write) = &changes.host_index { + host.write(write.clone()); + } + let overrides = match &changes.index { + Some(index) => (!index.is_empty()).then(|| Arc::new(index.clone())), + None => self.overrides.clone(), + }; + // Appends alone leave every deployment before them as it was. + let appended_only = changes.replace.is_none() + && changes.upserts.is_empty() + && changes.deletes.is_empty() + && changes.remove.is_none(); + let next = Self::index( + self.version + 1, + deployments, + aliases.unwrap_or_else(|| self.aliases.clone()), + self.routing_groups.clone(), + ); + Ok(Self { + overrides, + host: Arc::new(host), + member_groups: self.member_groups.clone(), + rewritten: if appended_only { + self.rewritten + } else { + next.rewritten + }, + ..next + }) + } + + /// The next registry for a build starting over (`set_model_list`): no deployments, and the + /// indexes [`HostIndexes::reset`] clears cleared. + #[must_use] + pub fn emptied(&self) -> Self { + let mut host = HostIndexes::clone(&self.host); + host.reset(); + Self { + overrides: self.overrides.clone(), + host: Arc::new(host), + member_groups: self.member_groups.clone(), + ..Self::index( + self.version + 1, + Vec::new(), + self.aliases.clone(), + self.routing_groups.clone(), + ) + } + } + + /// Add one deployment a build accepted at the end and index it in place (`add` is + /// `add_deployment`, which also names its group). A build installs into a registry of its + /// own this way, so a list of n deployments builds in O(n) rather than reindexing per + /// deployment. + pub fn push_built(&mut self, spec: DeploymentSpec, add: bool) -> Result<(), ConfigError> { + let deployment = Arc::new(self.carried(spec)?); + let index = self.deployments.len(); + let record = deployment.record(); + Arc::make_mut(&mut self.host).added( + record.as_object().expect("a record is an object"), + index, + add, + ); + self.by_id.insert(deployment.id.clone(), index); + let name = &deployment.model_name; + self.by_group.entry(name.clone()).or_default().push(index); + if let Some(team) = &deployment.model_info.team_id { + self.by_team.entry(team.clone()).or_default().push(index); + } + if name.contains('*') { + match self + .wildcards + .iter_mut() + .find(|(pattern, _)| pattern == name) + { + Some((_, indices)) => indices.push(index), + None => { + // A new pattern joins after every one it ties with, as `index` orders them: + // first appearance, then a stable sort by specificity. + self.wildcards.push((name.clone(), vec![index])); + self.wildcards.sort_by_key(|(pattern, _)| { + std::cmp::Reverse(pattern_specificity(pattern)) + }); + } + } + } + self.deployments.push(deployment); + self.version += 1; + Ok(()) + } + + /// The last deployment added. + pub fn last(&self) -> Option<&Arc> { + self.deployments.last() + } + + /// Take `set_model_list`'s `model_names`: the names it was given. + pub fn set_model_names(&mut self, names: Vec) { + Arc::make_mut(&mut self.host).set_model_names(names); + self.version += 1; + } + + /// The indexes the host hands out. + pub fn host_indexes(&self) -> &HostIndexes { + &self.host + } + + /// The deployment `spec` describes, keeping the provider the host resolved for the + /// deployment it replaces (same id, same params) when the spec carries none: a host that + /// hands back its whole list after a write does not resolve every provider again. + fn carried(&self, mut spec: DeploymentSpec) -> Result { + let deployment = Deployment::from_spec(spec.clone())?; + if spec.provider.is_none() + && let Some(previous) = self + .by_id + .get(&deployment.id) + .and_then(|index| self.deployments.get(*index)) + && previous.raw_params == deployment.raw_params + && previous.model_name == deployment.model_name + { + spec.provider.clone_from(&previous.resolved_provider); + return Deployment::from_spec(spec); + } + Ok(deployment) + } + + /// This registry with the providers `from` resolved for the same deployments (same id, + /// same params) and `from`'s host indexes: a router rebuilt for new static settings over the + /// same list keeps what its host resolved and the indexes their history built. + #[must_use] + pub fn with_providers_of(&self, from: &Self) -> Self { + let deployments = self + .deployments + .iter() + .map(|deployment| { + let previous = from + .by_id + .get(&deployment.id) + .and_then(|index| from.deployments.get(*index)) + .filter(|previous| { + previous.raw_params == deployment.raw_params + && previous.model_name == deployment.model_name + }); + match previous { + Some(previous) if deployment.resolved_provider.is_none() => { + Arc::new(Deployment { + resolved_provider: previous.resolved_provider.clone(), + ..Deployment::clone(deployment) + }) + } + _ => deployment.clone(), + } + }) + .collect(); + Self { + deployments, + host: from.host.clone(), + ..self.clone() + } + } + + fn at(&self, indices: &[usize]) -> Vec> { + indices + .iter() + .filter_map(|index| self.deployments.get(*index).cloned()) + .collect() + } + + /// `_core_deployments_named(name)`: the deployments the model-name index holds for `name` + /// (whatever their own name), or `None` when it does not index the name. + pub(crate) fn named(&self, name: &str) -> Option>> { + match self.overrides.as_ref().and_then(|o| o.model_name.as_ref()) { + Some(index) => index.get(name).map(|indices| self.at(indices)), + None => self.by_group.get(name).map(|indices| self.at(indices)), + } + } + + /// Whether the model-name index holds `name` (`name in model_name_to_deployment_indices`). + fn indexes_name(&self, name: &str) -> bool { + match self.overrides.as_ref().and_then(|o| o.model_name.as_ref()) { + Some(index) => index.contains_key(name), + None => self.by_group.contains_key(name), + } + } + + /// `_team_deployments_named(team_id, name)`: the deployments the team index holds for the + /// pair, or `None` when it does not index it. + pub(crate) fn team_named(&self, team_id: &str, name: &str) -> Option>> { + match self.overrides.as_ref().and_then(|o| o.team_model.as_ref()) { + Some(entries) => entries + .iter() + .find(|(team, public, _)| team == team_id && public == name) + .map(|(_, _, indices)| self.at(indices)), + None => { + let own: Vec> = self + .deployments + .iter() + .filter(|d| { + d.model_info.team_id.as_deref() == Some(team_id) + && d.model_info.team_public_model_name.as_deref() == Some(name) + }) + .cloned() + .collect(); + (!own.is_empty()).then_some(own) + } + } + } + + /// Every deployment the team index holds under public name `name`, any team's + /// (`_team_deployments_across_teams`). + pub(crate) fn team_named_across_teams(&self, name: &str) -> Vec> { + match self.overrides.as_ref().and_then(|o| o.team_model.as_ref()) { + Some(entries) => entries + .iter() + .filter(|(_, public, _)| public == name) + .flat_map(|(_, _, indices)| self.at(indices)) + .collect(), + None => self + .deployments + .iter() + .filter(|d| { + d.model_info.team_id.is_some() + && d.model_info.team_public_model_name.as_deref() == Some(name) + }) + .cloned() + .collect(), + } + } + + /// Whether a written index entry names a position past the model list. + fn overflows(&self, indices: &[usize]) -> bool { + indices.iter().any(|index| *index >= self.deployments.len()) + } + + /// The written model-name index entry for `name` names a position past the model list. + pub(crate) fn name_index_overflows(&self, name: &str) -> bool { + self.overrides + .as_ref() + .and_then(|o| o.model_name.as_ref()) + .and_then(|index| index.get(name)) + .is_some_and(|indices| self.overflows(indices)) + } + + /// The written team index entry for `(team_id, name)` names a position past the model list. + pub(crate) fn team_index_overflows(&self, team_id: &str, name: &str) -> bool { + self.overrides + .as_ref() + .and_then(|o| o.team_model.as_ref()) + .and_then(|entries| { + entries + .iter() + .find(|(team, public, _)| team == team_id && public == name) + }) + .is_some_and(|(_, _, indices)| self.overflows(indices)) + } + + /// A written team index entry under public name `name` names a position past the list. + pub(crate) fn across_teams_index_overflows(&self, name: &str) -> bool { + self.overrides + .as_ref() + .and_then(|o| o.team_model.as_ref()) + .is_some_and(|entries| { + entries + .iter() + .any(|(_, public, indices)| public == name && self.overflows(indices)) + }) + } + + /// The group a pick of `name` reads its strategy from, when the stale member map names a + /// group that is no longer configured (`_get_routing_context` raises `KeyError`): `name` + /// after its alias, unless it is a callable routing group. + pub fn missing_member_group(&self, name: &str) -> Option<&str> { + let members = self.member_groups.as_ref()?; + let resolved = self.alias_of(name).unwrap_or(name); + if self.routing_group_named(resolved).is_some() { + return None; + } + let group = members.get(resolved)?; + (!self.routing_groups.iter().any(|g| g.group_name == *group)).then_some(group.as_str()) + } + + /// The routing group a member call of `model_group` routes by: the stale member map's + /// group when there is one (`_model_to_group`), else the non-priority group listing it. + pub(crate) fn member_routing_group(&self, model_group: &str) -> Option<&RoutingGroupSpec> { + match &self.member_groups { + Some(members) => members + .get(model_group) + .and_then(|name| self.routing_groups.iter().find(|g| g.group_name == *name)), + None => self + .routing_groups + .iter() + .filter(|group| !crate::routing_groups::is_priority(group)) + .find(|group| group.models.iter().any(|model| model == model_group)), + } + } + + pub fn deployment(&self, id: &str) -> Option> { + self.by_id + .get(id) + .map(|index| self.deployments[*index].clone()) + } + + pub fn deployments(&self, filter: &DeploymentFilter<'_>) -> Vec> { + let candidates = match (filter.model_name, filter.team_id) { + (Some(name), _) => self + .by_group + .get(name) + .map(|i| self.at(i)) + .unwrap_or_default(), + (None, Some(team)) => self + .by_team + .get(team) + .map(|i| self.at(i)) + .unwrap_or_default(), + (None, None) => self.deployments.clone(), + }; + candidates + .into_iter() + .filter(|d| { + filter + .team_id + .is_none_or(|team| d.model_info.team_id.as_deref() == Some(team)) + }) + .filter(|d| { + filter + .provider + .is_none_or(|provider| d.provider() == provider) + }) + .collect() + } + + /// The model group names, in name order, with aliases when asked. + pub fn names(&self, filter: &NameFilter) -> Vec { + let mut names: Vec = self.by_group.keys().cloned().collect(); + if filter.include_aliases { + names.extend(self.aliases.keys().cloned()); + } + names + } + + pub fn group(&self, name: &str) -> Option>> { + if self + .overrides + .as_ref() + .is_some_and(|o| o.model_name.is_some()) + { + return self.named(name).map(|deployments| { + deployments + .into_iter() + .filter(|d| d.model_name == name) + .collect() + }); + } + self.by_group.get(name).map(|indices| self.at(indices)) + } + + /// How many rows `_core_get_model_list(model_name=name)` lists: the group's deployments, + /// an alias's target group, a routing group's members, and only when none of those, the + /// deployments of the wildcard the name matches. + pub fn model_list_len(&self, name: &str) -> usize { + let count = |group: &str| self.group(group).map_or(0, |deployments| deployments.len()); + let listed = count(name) + + self.aliases.get(name).map_or(0, |target| count(target)) + + self + .routing_groups + .iter() + .find(|group| group.group_name == name) + .map_or(0, |group| { + group.models.iter().map(|model| count(model)).sum() + }); + if listed > 0 { + return listed; + } + self.wildcards + .iter() + .find(|(pattern, _)| wildcard_matches(pattern, name)) + .map_or(0, |(_, indices)| indices.len()) + } + + /// The routing group a model group belongs to, if any. + pub fn routing_group_of(&self, model_name: &str) -> Option<&RoutingGroupSpec> { + self.routing_groups + .iter() + .find(|group| group.models.iter().any(|model| model == model_name)) + } + + /// The strategy a routing group named `group_name` declares. + pub fn resolve_routing_group_strategy( + &self, + group_name: &str, + ) -> Option { + self.routing_groups + .iter() + .find(|group| group.group_name == group_name) + .map(|group| group.routing_strategy.clone()) + } + + pub(crate) fn routing_groups(&self) -> &[RoutingGroupSpec] { + &self.routing_groups + } + + /// Every deployment, in configured order. + pub(crate) fn all(&self) -> &[Arc] { + &self.deployments + } + + /// Whether `name` is a deployment `model_name` (`name in _core_model_names`). + pub(crate) fn has_name(&self, name: &str) -> bool { + match self.overrides.as_ref().and_then(|o| o.model_names.as_ref()) { + Some(names) => names.contains(name), + None => self.by_group.contains_key(name), + } + } + + /// The group `model_group_alias` rewrites `name` to. + pub(crate) fn alias_of(&self, name: &str) -> Option<&str> { + self.aliases.get(name).map(String::as_str) + } + + /// The callable routing group named `name`: a model group of that name shadows it. + pub(crate) fn routing_group_named(&self, name: &str) -> Option<&RoutingGroupSpec> { + (!self.indexes_name(name)) + .then(|| self.routing_groups.iter().find(|g| g.group_name == name)) + .flatten() + } + + /// Whether any deployment is listed under a wildcard pattern. + pub(crate) fn has_wildcards(&self) -> bool { + !self.wildcards.is_empty() + } + + /// The most specific wildcard pattern matching `name`, with its deployments. + pub(crate) fn wildcard_for(&self, name: &str) -> Option<(String, Vec>)> { + self.wildcards + .iter() + .find(|(pattern, _)| wildcard_matches(pattern, name)) + .map(|(pattern, indices)| (pattern.clone(), self.at(indices))) + } + + /// Resolve a requested name, in the order of `_common_checks_available_deployment`: + /// a specific deployment, an id, an alias, a routing group, a model group, a wildcard. + pub fn resolve(&self, name: &str, scope: ResolveScope) -> Option { + let resolution = |kind, group: &str, deployments: Vec>| { + (!deployments.is_empty()).then(|| Resolution { + kind, + group: group.to_owned(), + deployments, + }) + }; + if scope.specific_deployment { + let matching = self + .deployments + .iter() + .filter(|d| d.params.model == name) + .cloned() + .collect(); + // An empty match still resolves: the pick then finds no deployment rather than + // falling back to a default (`_common_checks_available_deployment` returns `[]`). + return Some(Resolution { + kind: ResolveKind::SpecificDeployment, + group: name.to_owned(), + deployments: matching, + }); + } + if !self.by_group.contains_key(name) + && let Some(deployment) = self.deployment(name) + { + let group = deployment.model_name.clone(); + return resolution(ResolveKind::Id, &group, vec![deployment]); + } + let routing_group_named = |name: &str| { + // A model group of the same name shadows a routing group (`get_routing_group`). + (!self.by_group.contains_key(name)) + .then(|| self.routing_groups.iter().find(|g| g.group_name == name)) + .flatten() + }; + if let Some(target) = self.aliases.get(name) { + if let Some(routing_group) = routing_group_named(target) { + let deployments = crate::routing_groups::members(routing_group, |model| { + self.by_group + .get(model) + .map(|indices| self.at(indices)) + .unwrap_or_default() + }); + return resolution(ResolveKind::RoutingGroup, target, deployments); + } + return self + .group(target) + .and_then(|deployments| resolution(ResolveKind::Alias, target, deployments)); + } + if let Some(routing_group) = routing_group_named(name) { + let deployments = crate::routing_groups::members(routing_group, |model| { + self.by_group + .get(model) + .map(|indices| self.at(indices)) + .unwrap_or_default() + }); + return resolution(ResolveKind::RoutingGroup, name, deployments); + } + if let Some(indices) = self.by_group.get(name) { + return resolution(ResolveKind::Group, name, self.at(indices)); + } + self.wildcards + .iter() + .find(|(pattern, _)| wildcard_matches(pattern, name)) + .and_then(|(pattern, indices)| { + resolution(ResolveKind::Wildcard, pattern, self.at(indices)) + }) + } + + /// Access group name to the model groups in it. + pub fn access_groups(&self) -> BTreeMap> { + let mut groups: BTreeMap> = BTreeMap::new(); + for deployment in &self.deployments { + for group in &deployment.model_info.access_groups { + groups + .entry(group.clone()) + .or_default() + .insert(deployment.model_name.clone()); + } + } + groups + } + + pub fn credentials(&self, id: &str) -> Option { + self.deployment(id).map(|deployment| Credentials { + provider: deployment.provider().to_owned(), + api_base: deployment.params.api_base.clone(), + api_version: deployment.params.api_version.clone(), + credential_name: deployment.params.litellm_credential_name.clone(), + params: deployment.params.clone(), + }) + } +} + +/// `PatternUtils.calculate_pattern_specificity` over the pattern's regex +/// (`re.escape(pattern).replace(r"\*", "(.*)")`): its length, then how much regex syntax it has. +fn pattern_specificity(pattern: &str) -> (usize, usize) { + let mut regex = String::new(); + for c in pattern.chars() { + match c { + '*' => regex.push_str("(.*)"), + '(' | ')' | '[' | ']' | '{' | '}' | '?' | '+' | '-' | '|' | '^' | '$' | '\\' | '.' + | '&' | '~' | '#' | ' ' | '\t' | '\n' | '\r' | '\x0b' | '\x0c' => { + regex.push('\\'); + regex.push(c); + } + _ => regex.push(c), + } + } + let complexity = regex + .chars() + .filter(|c| matches!(c, '*' | '+' | '?' | '\\' | '^' | '$' | '|' | '(' | ')')) + .count(); + (regex.chars().count(), complexity) +} + +/// `*` matches any run of characters, as the pattern router's `(.*)` does under `re.match`: +/// anchored at the start only, so the name may run on past the pattern's last literal. +fn wildcard_matches(pattern: &str, name: &str) -> bool { + let mut parts = pattern.split('*'); + let first = parts.next().unwrap_or_default(); + let Some(mut rest) = name.strip_prefix(first) else { + return false; + }; + let parts: Vec<&str> = parts.collect(); + for part in parts { + match rest.find(part) { + Some(at) => rest = &rest[at + part.len()..], + None => return false, + } + } + true +} + +/// `deployment_id` in `litellm/router_api.py`: sha256 over the model group and every +/// `litellm_params` key and value in order. A string is taken as is, a dict as `json.dumps`, and +/// anything else as Python's `str()`. +pub fn deployment_id(model_group: &str, litellm_params: &Map) -> String { + let mut text = String::from(model_group); + for (key, value) in litellm_params { + text.push_str(key); + match value { + Value::String(value) => text.push_str(value), + Value::Object(_) => json_dumps(value, &mut text), + Value::Null | Value::Bool(_) | Value::Number(_) | Value::Array(_) => { + python_str(value, &mut text); + } + } + } + let digest = Sha256::digest(text.as_bytes()); + let mut hex = String::with_capacity(64); + for byte in digest { + let _ = write!(hex, "{byte:02x}"); + } + hex +} + +/// Python's `repr(float)`: the shortest round-trip digits, positional for exponents in +/// `-4..16`, scientific otherwise. +fn python_float(value: f64, out: &mut String) { + if value.is_nan() { + out.push_str("nan"); + return; + } + if value.is_infinite() { + out.push_str(if value > 0.0 { "inf" } else { "-inf" }); + return; + } + let scientific = format!("{value:e}"); + let (mantissa, exponent) = scientific.split_once('e').unwrap_or((&scientific, "0")); + let exponent: i32 = exponent.parse().unwrap_or(0); + let (sign, mantissa) = mantissa + .strip_prefix('-') + .map_or(("", mantissa), |m| ("-", m)); + let digits: String = mantissa.chars().filter(char::is_ascii_digit).collect(); + out.push_str(sign); + if (-4..16).contains(&exponent) { + if exponent < 0 { + out.push_str("0."); + out.extend(std::iter::repeat_n('0', (-exponent - 1) as usize)); + out.push_str(&digits); + } else { + let point = exponent as usize + 1; + if digits.len() <= point { + out.push_str(&digits); + out.extend(std::iter::repeat_n('0', point - digits.len())); + out.push_str(".0"); + } else { + out.push_str(&digits[..point]); + out.push('.'); + out.push_str(&digits[point..]); + } + } + } else { + out.push_str(&digits[..1]); + if digits.len() > 1 { + out.push('.'); + out.push_str(&digits[1..]); + } + let _ = write!( + out, + "e{}{:02}", + if exponent < 0 { '-' } else { '+' }, + exponent.abs() + ); + } +} + +fn number(value: &serde_json::Number, out: &mut String) { + match value.as_f64() { + Some(float) if !value.is_i64() && !value.is_u64() => python_float(float, out), + Some(_) | None => out.push_str(&value.to_string()), + } +} + +/// `json.dumps` with its defaults: `", "` and `": "` separators and `ensure_ascii`. +fn json_dumps(value: &Value, out: &mut String) { + match value { + Value::Null => out.push_str("null"), + Value::Bool(flag) => out.push_str(if *flag { "true" } else { "false" }), + Value::Number(value) => number(value, out), + Value::String(text) => json_string(text, out), + Value::Array(items) => { + out.push('['); + for (index, item) in items.iter().enumerate() { + if index > 0 { + out.push_str(", "); + } + json_dumps(item, out); + } + out.push(']'); + } + Value::Object(entries) => { + out.push('{'); + for (index, (key, item)) in entries.iter().enumerate() { + if index > 0 { + out.push_str(", "); + } + json_string(key, out); + out.push_str(": "); + json_dumps(item, out); + } + out.push('}'); + } + } +} + +fn json_string(text: &str, out: &mut String) { + out.push('"'); + for ch in text.chars() { + match ch { + '"' => out.push_str("\\\""), + '\\' => out.push_str("\\\\"), + '\n' => out.push_str("\\n"), + '\r' => out.push_str("\\r"), + '\t' => out.push_str("\\t"), + '\u{08}' => out.push_str("\\b"), + '\u{0c}' => out.push_str("\\f"), + ch if ch.is_ascii() && !ch.is_ascii_control() => out.push(ch), + ch => { + let mut units = [0u16; 2]; + for unit in ch.encode_utf16(&mut units) { + let _ = write!(out, "\\u{unit:04x}"); + } + } + } + } + out.push('"'); +} + +/// Python's `str()` of a JSON value: `repr` for the items of a list or dict. +fn python_str(value: &Value, out: &mut String) { + match value { + Value::Null => out.push_str("None"), + Value::Bool(flag) => out.push_str(if *flag { "True" } else { "False" }), + Value::Number(value) => number(value, out), + Value::String(text) => python_repr_str(text, out), + Value::Array(items) => { + out.push('['); + for (index, item) in items.iter().enumerate() { + if index > 0 { + out.push_str(", "); + } + python_str(item, out); + } + out.push(']'); + } + Value::Object(entries) => { + out.push('{'); + for (index, (key, item)) in entries.iter().enumerate() { + if index > 0 { + out.push_str(", "); + } + python_repr_str(key, out); + out.push_str(": "); + python_str(item, out); + } + out.push('}'); + } + } +} + +fn python_repr_str(text: &str, out: &mut String) { + let quote = if text.contains('\'') && !text.contains('"') { + '"' + } else { + '\'' + }; + out.push(quote); + for ch in text.chars() { + match ch { + '\\' => out.push_str("\\\\"), + '\n' => out.push_str("\\n"), + '\r' => out.push_str("\\r"), + '\t' => out.push_str("\\t"), + ch if ch == quote => { + out.push('\\'); + out.push(ch); + } + ch if (ch as u32) < 0x20 || ch as u32 == 0x7f => { + let _ = write!(out, "\\x{:02x}", ch as u32); + } + ch => out.push(ch), + } + } + out.push(quote); +} + +#[cfg(test)] +mod tests { + use rstest::rstest; + use serde_json::json; + + use super::*; + + fn spec(model_name: &str, params: Value, info: Value) -> DeploymentSpec { + serde_json::from_value(json!({ + "model_name": model_name, "litellm_params": params, "model_info": info, + })) + .unwrap() + } + + #[rstest] + #[case::none(json!({}), json!({}), json!({}), false)] + #[case::params_rpm(json!({"rpm": 1}), json!({}), json!({}), true)] + #[case::params_itpm(json!({"itpm": 1}), json!({}), json!({}), true)] + #[case::model_info_rpm(json!({}), json!({"rpm": 1}), json!({}), true)] + #[case::model_info_tpm(json!({}), json!({"tpm": 1}), json!({}), true)] + #[case::model_info_otpm(json!({}), json!({"otpm": 1}), json!({}), true)] + #[case::model_info_null(json!({}), json!({"rpm": null}), json!({}), false)] + #[case::top_level_tpm(json!({}), json!({}), json!({"tpm": 1}), true)] + fn a_limit_anywhere_on_the_deployment_tracks_usage( + #[case] params: Value, + #[case] info: Value, + #[case] entry: Value, + #[case] tracks: bool, + ) { + let mut spec = json!({ + "model_name": "g", + "litellm_params": {"model": "openai/a"}, + "model_info": {"id": "a"}, + }); + let extend = |section: &mut Value, more: Value| { + section + .as_object_mut() + .unwrap() + .extend(more.as_object().unwrap().clone()); + }; + extend(&mut spec["litellm_params"], params); + extend(&mut spec["model_info"], info); + extend(&mut spec, entry); + let spec: DeploymentSpec = serde_json::from_value(spec).unwrap(); + assert_eq!(Deployment::from_spec(spec).unwrap().tracks_usage(), tracks); + } + + #[rstest] + #[case::params_int(json!({"model": "m", "order": 2}), json!({}), Some(2))] + #[case::params_win(json!({"model": "m", "order": 2}), json!({"order": 1}), Some(2))] + #[case::model_info(json!({"model": "m"}), json!({"order": 3}), Some(3))] + #[case::params_null(json!({"model": "m", "order": null}), json!({"order": 3}), Some(3))] + #[case::numeric_string(json!({"model": "m", "order": "4"}), json!({}), Some(4))] + #[case::model_info_string(json!({"model": "m"}), json!({"order": "5"}), Some(5))] + #[case::whole_float(json!({"model": "m", "order": 6.0}), json!({}), Some(6))] + #[case::none(json!({"model": "m"}), json!({}), None)] + fn order_reads_as_get_deployment_order( + #[case] params: Value, + #[case] info: Value, + #[case] expected: Option, + ) { + let deployment = Deployment::from_spec(spec("g", params, info)).unwrap(); + assert_eq!(deployment.order(), expected); + } + + #[rstest] + #[case::unset(json!({}), true)] + #[case::declared(json!({"supports_web_search": true}), true)] + #[case::refused(json!({"supports_web_search": false}), false)] + #[case::null_is_unset(json!({"supports_web_search": null}), true)] + fn web_search_support_is_assumed_unless_declared(#[case] info: Value, #[case] expected: bool) { + let deployment = Deployment::from_spec(spec("g", json!({"model": "m"}), info)).unwrap(); + assert_eq!(deployment.supports_web_search(), expected); + } + + #[test] + fn a_routed_wildcard_fills_its_model_and_drops_the_cache_flags() { + let deployment = Arc::new( + Deployment::from_spec(spec( + "anthropic/*", + json!({"model": "anthropic/*"}), + json!({"id": "w", "supports_prompt_caching": true, "keep": 1}), + )) + .unwrap(), + ); + let routed = routed_deployment( + &deployment, + &["claude-3-haiku".to_owned()], + "anthropic/claude-3-haiku", + ); + assert_eq!(routed.params.model, "anthropic/claude-3-haiku"); + assert_eq!(routed.raw_params["model"], "anthropic/claude-3-haiku"); + assert_eq!(routed.record()["model_info"], json!({"id": "w", "keep": 1})); + } + + #[rstest] + #[case::no_team_lists_a_team_owned_one(Some("t1"), None, true)] + #[case::no_team_lists_an_unowned_one(None, None, true)] + #[case::the_owner_lists_it(Some("t1"), Some("t1"), true)] + #[case::another_team_does_not(Some("t1"), Some("t2"), false)] + #[case::a_team_lists_an_unowned_one(None, Some("t2"), true)] + fn listed_for_team_lets_no_team_see_every_deployment( + #[case] owner: Option<&str>, + #[case] team: Option<&str>, + #[case] listed: bool, + ) { + let deployment = Deployment::from_spec(spec( + "g", + json!({"model": "openai/a"}), + json!({"id": "a", "team_id": owner}), + )) + .unwrap(); + assert_eq!(deployment.listed_for_team(team), listed); + } + + /// Expected ids computed with `litellm.router_api.deployment_id`. + #[rstest] + #[case::strings( + "g", + json!({"model": "openai/gpt-4o", "api_key": "sk-1"}), + "30256ba697cc98ce2aff32c3b73c4b90514cb199876aef429ffcd91ae826d025" + )] + #[case::scalars( + "claude", + json!({"model": "anthropic/claude-3", "rpm": 10, "timeout": 1.5, "stream": true, "x": null}), + "645fd1504089c1084316573ccfc5dd1dd9622623b169026657484fb64a2c89a2" + )] + #[case::list_and_dict( + "m", + json!({"model": "azure/x", "tags": ["a", "b'c"], "extra_body": {"k": [1, 2.0, "é"], "z": {"q": false}}}), + "64f11e16dabed321268bd92132de2020290ef193c0cf554b2c74faee8badbb96" + )] + fn deployment_id_matches_python( + #[case] group: &str, + #[case] params: Value, + #[case] expected: &str, + ) { + let Value::Object(params) = params else { + unreachable!() + }; + assert_eq!(deployment_id(group, ¶ms), expected); + } + + #[rstest] + #[case::integral(2.0, "2.0")] + #[case::small(0.00001, "1e-05")] + #[case::large(1e16, "1e+16")] + #[case::fraction(0.1, "0.1")] + fn floats_print_as_python_repr(#[case] value: f64, #[case] expected: &str) { + let mut out = String::new(); + python_float(value, &mut out); + assert_eq!(out, expected); + } + + #[rstest] + #[case::group("gpt-4", ResolveKind::Group, "gpt-4")] + #[case::alias("smart", ResolveKind::Alias, "gpt-4")] + #[case::id("dep-b", ResolveKind::Id, "gpt-4")] + #[case::wildcard("anthropic/claude-x", ResolveKind::Wildcard, "anthropic/*")] + fn resolve_follows_python_order( + #[case] name: &str, + #[case] kind: ResolveKind, + #[case] group: &str, + ) { + let registry = Registry::build( + vec![ + spec( + "gpt-4", + json!({"model": "openai/gpt-4"}), + json!({"id": "dep-a"}), + ), + spec( + "gpt-4", + json!({"model": "azure/gpt-4"}), + json!({"id": "dep-b"}), + ), + spec("anthropic/*", json!({"model": "anthropic/*"}), json!({})), + ], + BTreeMap::from([("smart".to_owned(), "gpt-4".to_owned())]), + Vec::new(), + ) + .unwrap(); + let resolution = registry.resolve(name, ResolveScope::default()).unwrap(); + assert_eq!((resolution.kind, resolution.group.as_str()), (kind, group)); + } + + #[rstest] + fn apply_moves_an_edited_deployment_to_the_end() { + let registry = Registry::build( + vec![ + spec("a", json!({"model": "openai/a"}), json!({"id": "1"})), + spec("a", json!({"model": "openai/b"}), json!({"id": "2"})), + ], + BTreeMap::new(), + Vec::new(), + ) + .unwrap(); + let next = registry + .apply( + &ChangeSet { + upserts: vec![spec("a", json!({"model": "openai/c"}), json!({"id": "1"}))], + ..ChangeSet::default() + }, + None, + ) + .unwrap(); + let ids: Vec = next + .group("a") + .unwrap() + .iter() + .map(|d| d.id.clone()) + .collect(); + assert_eq!( + (next.version(), ids), + (1, vec!["2".to_owned(), "1".to_owned()]) + ); + } +} diff --git a/litellm-rust/crates/router-core/src/resolve.rs b/litellm-rust/crates/router-core/src/resolve.rs new file mode 100644 index 00000000000..9075fc36c53 --- /dev/null +++ b/litellm-rust/crates/router-core/src/resolve.rs @@ -0,0 +1,585 @@ +//! A requested name to the deployments that may serve this caller: a port of +//! `_common_checks_available_deployment`. In order: a specific deployment, a deployment id, the +//! alias, a routing group, then for a name no deployment has, a team's public name, every team's +//! public name for a proxy admin, a wildcard; else the group (team-first), kept to the caller's +//! access groups, then deployments whose litellm model is the name, then the first default +//! fallback. Every arm drops strategy markers and deployments that refuse the request's endpoint, +//! and withholds deployments reserved for other teams. + +use std::{collections::BTreeSet, sync::Arc}; + +use jiff::Timestamp; + +use serde_json::Value; + +use crate::{ + access::{filter_by_access_groups, filter_reserved}, + eligibility::eligible_deployments, + error::Error, + host_index::{pattern_to_regex, regex_match}, + registry::{Deployment, DeploymentSpec, Registry, Resolution, ResolveKind, routed_deployment}, + select::is_strategy_marker, +}; + +/// Who is asking, as the request's auth metadata says. +#[derive(Clone, Copy, Debug, Default)] +pub struct Caller<'a> { + /// `user_api_key_team_id`. + pub team_id: Option<&'a str>, + /// The caller's `user_role` is `proxy_admin`. + pub proxy_admin: bool, + /// The caller's `models` plus `team_models`. + pub allowed_models: &'a [String], + /// The endpoint the pick runs under (`endpoint_selection`), which deployments may refuse. + pub endpoint: Option<&'a str>, + /// The model `endpoint_selection` was given: the name an endpoint refusal reads. + pub requested: &'a str, + /// `f"{provider}/{model}"`, the spelling `PatternMatchRouter.get_pattern` tries when no + /// pattern matches the name itself. + pub wildcard_request: Option<&'a str>, +} + +/// `_core_should_include_deployment`: the team's own deployment under its public name, or a +/// deployment named `model_name` the team may use. +fn included(deployment: &Deployment, model_name: &str, team_id: Option<&str>) -> bool { + let info = &deployment.model_info; + if team_id.is_some() + && info.team_id.as_deref() == team_id + && info.team_public_model_name.as_deref() == Some(model_name) + { + return true; + } + deployment.model_name == model_name + && (team_id.is_none() || info.team_id.is_none() || info.team_id.as_deref() == team_id) +} + +/// `_core_get_all_deployments(model_name, team_id)`: the team's deployments under that public +/// name when it has any, else the deployments named `model_name` the team may use; a name the +/// model-name index does not hold is looked for in every deployment, for a team. +pub fn all_deployments( + registry: &Registry, + model_name: &str, + team_id: Option<&str>, +) -> Result>, Error> { + let keep = |deployments: Vec>| -> Vec> { + deployments + .into_iter() + .filter(|d| included(d, model_name, team_id)) + .collect() + }; + // A written index position past the model list raises `IndexError` where it is read. + let out_of_range = || Error::IndexOutOfRange { + model: model_name.to_owned(), + }; + if let Some(team) = team_id { + if registry.team_index_overflows(team, model_name) { + return Err(out_of_range()); + } + if let Some(own) = registry.team_named(team, model_name) { + let own = keep(own); + if !own.is_empty() { + return Ok(own); + } + } + } + if registry.name_index_overflows(model_name) { + return Err(out_of_range()); + } + Ok(match registry.named(model_name) { + Some(named) => keep(named), + None if team_id.is_some() => keep(registry.all().to_vec()), + None => Vec::new(), + }) +} + +/// `_team_deployments_across_teams`: every team's deployments under public name `model`, for a +/// proxy admin calling without a team; a name several teams share is refused. +fn across_teams(registry: &Registry, model: &str) -> Result>, Error> { + if registry.across_teams_index_overflows(model) { + return Err(Error::IndexOutOfRange { + model: model.to_owned(), + }); + } + let deployments = registry.team_named_across_teams(model); + let teams: BTreeSet<&str> = deployments + .iter() + .filter_map(|d| d.model_info.team_id.as_deref()) + .collect(); + if teams.len() > 1 { + return Err(Error::MultipleTeams { + model: model.to_owned(), + }); + } + Ok(deployments) +} + +/// The access groups of the rows `_core_get_model_list(model_name, team_id)` lists: the +/// group's deployments, an alias's target deployments, a routing group's rows (which carry no +/// access groups), and only when there are none of those, the wildcard's deployments. +fn model_access_groups( + registry: &Registry, + model: &str, + team_id: Option<&str>, +) -> BTreeSet { + let mut rows = all_deployments(registry, model, team_id).unwrap_or_default(); + let mut stripped_rows = 0; + if let Some(target) = registry.alias_of(model) { + match registry.routing_group_named(target) { + Some(group) => stripped_rows += routing_group_rows(registry, group, None).len(), + None => rows.extend(all_deployments(registry, target, None).unwrap_or_default()), + } + } + if let Some(group) = registry.routing_group_named(model) { + stripped_rows += routing_group_rows(registry, group, None).len(); + } + if rows.is_empty() && stripped_rows == 0 { + rows = registry + .wildcard_for(model) + .map(|(_, deployments)| deployments) + .unwrap_or_default(); + } + rows.iter() + .flat_map(|d| d.model_info.access_groups.iter().cloned()) + .collect() +} + +fn routing_group_rows( + registry: &Registry, + group: &crate::registry::RoutingGroupSpec, + team_id: Option<&str>, +) -> Vec> { + crate::routing_groups::members(group, |member| { + all_deployments(registry, member, team_id).unwrap_or_default() + }) +} + +/// `_filter_reserved_deployments` for deployments resolved as `model`: markers dropped (only +/// markers is refused), those that refuse the caller's endpoint dropped (`_drop_strategy_markers`), +/// then the deployments an active access window withholds; when that leaves nothing, the window +/// is named in the refusal. +fn reserved( + model: &str, + deployments: Vec>, + caller: Caller<'_>, + now: Timestamp, +) -> Result>, Error> { + let any = !deployments.is_empty(); + let selectable: Vec> = deployments + .into_iter() + .filter(|d| !is_strategy_marker(d)) + .collect(); + if any && selectable.is_empty() { + return Err(Error::OnlyMarkers { + model: model.to_owned(), + }); + } + let selectable = eligible_deployments(model, caller, selectable)?; + let (kept, blocking) = filter_reserved(selectable, caller.team_id, now); + match blocking { + Some(window) if kept.is_empty() => Err(Error::Reserved { + model: model.to_owned(), + until: window.end_hhmm(), + timezone: window.timezone, + }), + _ => Ok(kept), + } +} + +/// `PatternMatchRouter.route`: copies of a pattern's deployments serving `request`. +fn routed_wildcard( + pattern: &str, + request: &str, + deployments: &[Arc], +) -> Vec> { + let groups = regex_match(&pattern_to_regex(pattern), request).unwrap_or_default(); + deployments + .iter() + .map(|deployment| routed_deployment(deployment, &groups, request)) + .collect() +} + +/// `get_pattern` for `model`: the most specific pattern matching it as spelled, else its +/// provider-prefixed spelling, with the pattern's deployments serving that request. +fn wildcard_route( + registry: &Registry, + model: &str, + wildcard_request: Option<&str>, +) -> Option<(String, Vec>)> { + let route = |request: &str| { + registry + .wildcard_for(request) + .map(|(pattern, deployments)| { + let routed = routed_wildcard(&pattern, request, &deployments); + (pattern, routed) + }) + }; + route(model).or_else(|| route(wildcard_request?)) +} + +/// The name resolution would look for under its provider (`get_pattern`'s second spelling), which +/// only the host can infer: one that reaches the pattern routers and that no pattern matches as +/// spelled. +pub(crate) fn wildcard_name<'a>( + registry: &'a Registry, + name: &'a str, + specific_deployment: bool, +) -> Option<&'a str> { + let model = registry.alias_of(name).unwrap_or(name); + (registry.has_wildcards() + && !specific_deployment + && !registry.has_name(model) + && (registry.has_name(name) || registry.deployment(name).is_none()) + && registry.routing_group_named(model).is_none() + && registry.wildcard_for(model).is_none()) + .then_some(model) +} + +/// The first default fallback a name no deployment serves is routed to. +pub fn default_target(dynamic: &crate::config::RouterDynamicConfig) -> Option<&str> { + dynamic + .fallbacks + .default_target() + .or_else(|| dynamic.default_fallbacks.first().map(String::as_str)) +} + +/// The router's `default_deployment`, as resolution serves it. +#[derive(Clone, Copy, Debug)] +pub enum DefaultDeployment<'a> { + Spec(&'a DeploymentSpec), + /// A value Python cannot copy for a name: resolving to it raises. + Raises, +} + +/// The default deployment as served for `model`: a copy whose `litellm_params.model` is the name. +fn default_for(spec: &DeploymentSpec, model: &str) -> Option> { + let mut spec = spec.clone(); + spec.litellm_params + .insert("model".to_owned(), Value::String(model.to_owned())); + Deployment::from_spec(spec).ok().map(Arc::new) +} + +/// `_core_get_routable_deployments`: the deployments resolution would consider for `request`, +/// before any team, cooldown or strategy filter; nothing when it resolves to nothing or refuses. +pub fn routable( + registry: &Registry, + dynamic: &crate::config::RouterDynamicConfig, + request: &crate::host::RouterRequest, +) -> Vec> { + let caller = Caller { + team_id: request.team_id.as_deref(), + proxy_admin: request.proxy_admin, + allowed_models: &request.allowed_models, + endpoint: request.endpoint.as_deref(), + requested: &request.model, + wildcard_request: request.wildcard_request.as_deref(), + }; + match resolve_request( + registry, + &request.model, + false, + caller, + default_target(dynamic), + dynamic.default_ref(), + Timestamp::now(), + ) { + Ok(Some(resolution)) => resolution.deployments, + Ok(None) | Err(_) => Vec::new(), + } +} + +/// Resolve `name` for `caller`. `None` when nothing serves it (`no_healthy_deployments`). +pub fn resolve_request( + registry: &Registry, + name: &str, + specific_deployment: bool, + caller: Caller<'_>, + default_target: Option<&str>, + default_deployment: Option>, + now: Timestamp, +) -> Result, Error> { + let team_id = caller.team_id; + let resolved = |kind, group: &str, deployments| { + Ok(Some(Resolution { + kind, + group: group.to_owned(), + deployments, + })) + }; + if specific_deployment { + let matching = registry + .all() + .iter() + .filter(|d| d.params.model == name) + .cloned() + .collect(); + let kept = reserved(name, matching, caller, now)?; + return resolved(ResolveKind::SpecificDeployment, name, kept); + } + if !registry.has_name(name) + && let Some(deployment) = registry.deployment(name) + { + let group = deployment.model_name.clone(); + let litellm_model = deployment.params.model.clone(); + let kept = reserved(&litellm_model, vec![deployment], caller, now)?; + return resolved(ResolveKind::Id, &group, kept); + } + let aliased = registry.alias_of(name); + let model = aliased.unwrap_or(name); + let routing_group = registry + .routing_group_named(model) + .map(|group| routing_group_rows(registry, group, team_id)); + if routing_group.is_none() && !registry.has_name(model) { + let early = match team_id { + Some(team) => Some(all_deployments(registry, model, Some(team))?), + None if caller.proxy_admin => Some(across_teams(registry, model)?), + None => None, + } + .filter(|deployments| !deployments.is_empty()); + if let Some(deployments) = early { + let kept = reserved(model, deployments, caller, now)?; + return resolved(ResolveKind::Group, model, kept); + } + // `_core_wildcard_route_deployments`: the team's own pattern, else the proxy-wide one. + let team_owned = team_id + .map(|team| crate::lookup::Lookup::new(registry).team_wildcard(model, team)) + .filter(|deployments| !deployments.is_empty()); + if let Some(deployments) = team_owned { + let routed = deployments.iter().flat_map(|deployment| { + let public = deployment.model_info.team_public_model_name.as_deref(); + let pattern = public.unwrap_or(model); + routed_wildcard(pattern, model, std::slice::from_ref(deployment)) + }); + let kept = reserved(model, routed.collect(), caller, now)?; + return resolved(ResolveKind::Wildcard, model, kept); + } + if let Some((pattern, deployments)) = + wildcard_route(registry, model, caller.wildcard_request) + { + let kept = reserved(model, deployments, caller, now)?; + return resolved(ResolveKind::Wildcard, &pattern, kept); + } + match default_deployment { + Some(DefaultDeployment::Raises) => { + return Err(Error::DefaultDeploymentRaises { + model: model.to_owned(), + }); + } + Some(DefaultDeployment::Spec(spec)) => { + if let Some(deployment) = default_for(spec, model) { + let kept = reserved(model, vec![deployment], caller, now)?; + return resolved(ResolveKind::Default, model, kept); + } + } + None => {} + } + } + let (kind, listed) = match routing_group { + Some(deployments) => (ResolveKind::RoutingGroup, deployments), + None if aliased.is_some() => ( + ResolveKind::Alias, + all_deployments(registry, model, team_id)?, + ), + None => ( + ResolveKind::Group, + all_deployments(registry, model, team_id)?, + ), + }; + let narrow = |model: &str, deployments: Vec>| { + let allowed = filter_by_access_groups(model, deployments, caller.allowed_models, || { + model_access_groups(registry, model, team_id) + }); + reserved(model, allowed, caller, now) + }; + let listed_any = !listed.is_empty(); + let deployments = narrow(model, listed)?; + if !deployments.is_empty() { + return resolved(kind, model, deployments); + } + let mut access_emptied = listed_any; + if !listed_any { + // The caller may have named a deployment's litellm model instead. + let by_model: Vec> = registry + .all() + .iter() + .filter(|d| d.params.model == model) + .cloned() + .collect(); + let found = !by_model.is_empty(); + let deployments = narrow(model, by_model)?; + if !deployments.is_empty() { + return resolved(ResolveKind::Group, model, deployments); + } + access_emptied = found; + } + if let Some(target) = default_target.filter(|target| !access_emptied && !target.is_empty()) { + let deployments = narrow(target, all_deployments(registry, target, team_id)?)?; + if !deployments.is_empty() { + return resolved(ResolveKind::Group, target, deployments); + } + // `_common_checks_available_deployment` re-assigned the name: its error names the target. + return Err(Error::UnknownModel { + model: target.to_owned(), + }); + } + Ok(None) +} + +#[cfg(test)] +mod tests { + use std::collections::BTreeMap; + + use serde_json::{Value, json}; + + use super::*; + use crate::registry::DeploymentSpec; + + fn registry(entries: Value) -> Registry { + let specs: Vec = serde_json::from_value(entries).unwrap(); + Registry::build(specs, BTreeMap::new(), Vec::new()).unwrap() + } + + fn ids(resolution: Option) -> Vec { + resolution + .map(|r| r.deployments.iter().map(|d| d.id.clone()).collect()) + .unwrap_or_default() + } + + #[test] + fn an_access_group_only_caller_is_kept_to_its_groups() { + let registry = registry(json!([ + {"model_name": "g", "litellm_params": {"model": "openai/a"}, "model_info": {"id": "a", "access_groups": ["A"]}}, + {"model_name": "g", "litellm_params": {"model": "openai/b"}, "model_info": {"id": "b", "access_groups": ["B"]}}, + ])); + let allowed = vec!["A".to_owned()]; + let caller = Caller { + allowed_models: &allowed, + ..Caller::default() + }; + let resolved = resolve_request(®istry, "g", false, caller, None, None, Timestamp::now()); + assert_eq!(ids(resolved.unwrap()), ["a"]); + // The litellm model string is not narrowed (python-fixes-to-make.md). + let by_model = resolve_request( + ®istry, + "openai/b", + false, + caller, + None, + None, + Timestamp::now(), + ); + assert_eq!(ids(by_model.unwrap()), ["b"]); + } + + #[test] + fn a_proxy_admin_is_refused_a_public_name_several_teams_share() { + let registry = registry(json!([ + {"model_name": "x", "litellm_params": {"model": "openai/x"}, "model_info": {"id": "x", "team_id": "tx", "team_public_model_name": "pub"}}, + {"model_name": "y", "litellm_params": {"model": "openai/y"}, "model_info": {"id": "y", "team_id": "ty", "team_public_model_name": "pub"}}, + ])); + let caller = Caller { + proxy_admin: true, + ..Caller::default() + }; + let refused = resolve_request( + ®istry, + "pub", + false, + caller, + None, + None, + Timestamp::now(), + ); + assert!(matches!(refused, Err(Error::MultipleTeams { model }) if model == "pub")); + } + + fn wildcard_group() -> Value { + json!([{ + "model_name": "openai/*", + "litellm_params": {"model": "openai/*"}, + "model_info": {"id": "w", "supports_prompt_cache_breakpoint": true}, + }]) + } + + #[test] + fn a_wildcard_serves_the_matched_model_without_its_cache_flags() { + let registry = registry(wildcard_group()); + let resolved = resolve_request( + ®istry, + "openai/gpt-4o", + false, + Caller::default(), + None, + None, + Timestamp::now(), + ); + let resolution = resolved.unwrap().unwrap(); + assert_eq!(resolution.kind, ResolveKind::Wildcard); + assert_eq!(resolution.deployments[0].params.model, "openai/gpt-4o"); + assert_eq!( + resolution.deployments[0].record()["model_info"], + json!({"id": "w"}) + ); + } + + #[test] + fn an_alias_to_a_wildcard_fills_the_model_from_its_target() { + let specs: Vec = serde_json::from_value(wildcard_group()).unwrap(); + let aliases = BTreeMap::from([("al".to_owned(), "openai/gpt-4o".to_owned())]); + let registry = Registry::build(specs, aliases, Vec::new()).unwrap(); + let resolved = resolve_request( + ®istry, + "al", + false, + Caller::default(), + None, + None, + Timestamp::now(), + ); + assert_eq!( + resolved.unwrap().unwrap().deployments[0].params.model, + "openai/gpt-4o" + ); + } + + #[test] + fn a_name_no_pattern_matches_is_tried_under_its_provider() { + let registry = registry(wildcard_group()); + assert_eq!(wildcard_name(®istry, "gpt-4o", false), Some("gpt-4o")); + assert_eq!(wildcard_name(®istry, "openai/gpt-4o", false), None); + let caller = Caller { + wildcard_request: Some("openai/gpt-4o"), + ..Caller::default() + }; + let resolved = resolve_request( + ®istry, + "gpt-4o", + false, + caller, + None, + None, + Timestamp::now(), + ); + let resolution = resolved.unwrap().unwrap(); + assert_eq!(resolution.group, "openai/*"); + assert_eq!(resolution.deployments[0].params.model, "openai/gpt-4o"); + } + + #[test] + fn a_deployment_named_by_id_refuses_the_endpoint_like_any_other() { + let registry = registry(json!([ + {"model_name": "g", "litellm_params": {"model": "openai/a"}, "model_info": { + "id": "a", "enforce_supported_endpoints": true, "supported_endpoints": ["/v1/embeddings"], + }}, + ])); + let caller = Caller { + endpoint: Some("/v1/chat/completions"), + requested: "a", + ..Caller::default() + }; + let refused = resolve_request(®istry, "a", false, caller, None, None, Timestamp::now()); + assert!(matches!( + refused, + Err(Error::EndpointNotSupported { model, .. }) if model == "openai/a" + )); + } +} diff --git a/litellm-rust/crates/router-core/src/retry.rs b/litellm-rust/crates/router-core/src/retry.rs new file mode 100644 index 00000000000..3de0ad601e3 --- /dev/null +++ b/litellm-rust/crates/router-core/src/retry.rs @@ -0,0 +1,219 @@ +//! Retry decisions: how many retries a `RetryPolicy` grants for a failure, and whether a failure +//! is retried in the same group at all. Mirrors `get_num_retries_from_retry_policy` and +//! `resolve_retry_policy` in `litellm_core_utils/get_retry_from_policy.py`, and +//! `Router.should_retry_this_error`. + +use std::collections::BTreeMap; + +use crate::{ + config::RetryPolicy, + error::{AttemptError, FailureKind}, + health::should_retry_status, +}; + +/// `resolve_retry_policy`: the group's own policy when `model_group_retry_policy` has one, else +/// the router's `retry_policy`. +pub fn resolve_retry_policy<'a>( + retry_policy: Option<&'a RetryPolicy>, + model_group: &str, + model_group_retry_policy: Option<&'a BTreeMap>, +) -> Option<&'a RetryPolicy> { + model_group_retry_policy + .and_then(|policies| policies.get(model_group)) + .or(retry_policy) +} + +impl RetryPolicy { + /// `get_num_retries_from_retry_policy`: `DefaultRetries` for an `UnsupportedEndpointError`; + /// else `NotFoundErrorRetries` for any 404 answer, then the most specific exception class + /// the policy names, then `DefaultRetries`. + pub fn retries_for(&self, failure: &AttemptError) -> Option { + if failure.is_a("UnsupportedEndpointError") { + return self.DefaultRetries; + } + let not_found = (failure.status == Some(404)) + .then_some(self.NotFoundErrorRetries) + .flatten(); + not_found + .or_else(|| { + failure + .class_names() + .iter() + .find_map(|class| self.field_for(class)) + }) + .or(self.DefaultRetries) + } + + /// `_RETRIES_BY_EXCEPTION_TYPE`: the policy field an exception class reads. + fn field_for(&self, class: &str) -> Option { + match class { + "AuthenticationError" => self.AuthenticationErrorRetries, + "Timeout" => self.TimeoutErrorRetries, + "RateLimitError" => self.RateLimitErrorRetries, + "ContentPolicyViolationError" => self.ContentPolicyViolationErrorRetries, + "BadRequestError" => self.BadRequestErrorRetries, + "ServiceUnavailableError" => self.ServiceUnavailableErrorRetries, + "InternalServerError" => self.InternalServerErrorRetries, + _ => None, + } + } +} + +/// What `should_retry_this_error` knows besides the failure. +#[derive(Clone, Copy, Debug, Default)] +pub struct RetryContext { + pub healthy: usize, + pub all: usize, + /// `context_window_fallbacks is not None`. + pub has_context_window_fallbacks: bool, + /// `content_policy_fallbacks is not None`. + pub has_content_policy_fallbacks: bool, + /// `regular_fallbacks is not None and len(regular_fallbacks) > 0`. + pub has_regular_fallbacks: bool, +} + +/// `should_retry_this_error`, as a predicate: `false` where Python raises the error. +pub fn should_retry(failure: &AttemptError, ctx: RetryContext) -> bool { + if failure.is_a("ContextWindowExceededError") && ctx.has_context_window_fallbacks { + return false; + } + if failure.is_a("ContentPolicyViolationError") && ctx.has_content_policy_fallbacks { + return false; + } + if let Some(status) = failure.status + && !should_retry_status(status) + && !matches!(status, 401 | 403) + && !failure.is_a("UnsupportedEndpointError") + { + return false; + } + if failure.is_a("NotFoundError") { + return false; + } + if failure.is_a("RateLimitError") && ctx.healthy == 0 && ctx.has_regular_fallbacks { + return false; + } + if (failure.is_a("AuthenticationError") || failure.is_a("PermissionDeniedError")) + && ctx.all <= 1 + { + return false; + } + ctx.healthy > 0 +} + +/// `ContentPolicyViolationError`'s classes, most specific first. +pub(crate) fn content_policy_classes() -> Vec> { + [ + "ContentPolicyViolationError", + "BadRequestError", + "APIStatusError", + "APIError", + "OpenAIError", + "Exception", + "BaseException", + "object", + ] + .into_iter() + .map(Into::into) + .collect() +} + +/// The exception classes a failure kind stands for, most specific first, for a host that +/// reports no class names (`litellm/exceptions.py`). +pub(crate) fn kind_classes(kind: FailureKind) -> &'static [&'static str] { + match kind { + FailureKind::RateLimit => &["RateLimitError"], + FailureKind::ContextWindowExceeded => &["ContextWindowExceededError", "BadRequestError"], + FailureKind::ContentPolicyViolation => &["ContentPolicyViolationError", "BadRequestError"], + FailureKind::Authentication => &["AuthenticationError"], + FailureKind::PermissionDenied => &["PermissionDeniedError"], + FailureKind::NotFound => &["NotFoundError"], + FailureKind::BadRequest => &["BadRequestError"], + FailureKind::Timeout => &["Timeout", "APITimeoutError", "APIConnectionError"], + FailureKind::ApiConnection => &["APIConnectionError"], + FailureKind::InternalServer => &["InternalServerError"], + FailureKind::ServiceUnavailable => &["ServiceUnavailableError"], + FailureKind::BadGateway => &["BadGatewayError"], + FailureKind::CallerCredentialAuthentication => { + &["CallerCredentialAuthenticationError", "AuthenticationError"] + } + FailureKind::CallerCredentialRateLimit => { + &["CallerCredentialRateLimitError", "RateLimitError"] + } + FailureKind::GuardrailIntervention | FailureKind::Declined | FailureKind::Other => &[], + } +} + +#[cfg(test)] +mod tests { + use rstest::rstest; + + use super::*; + + fn failure(kind: FailureKind, status: Option, classes: &[&str]) -> AttemptError { + let mut failure = AttemptError::new(kind, status); + failure.classes = classes.iter().map(|class| (*class).into()).collect(); + failure + } + + #[rstest] + #[case::own_field(failure(FailureKind::InternalServer, Some(500), &[]), Some(4))] + #[case::default(failure(FailureKind::BadGateway, Some(502), &[]), Some(7))] + #[case::not_found_status(failure(FailureKind::Other, Some(404), &["APIError"]), Some(1))] + #[case::unsupported_endpoint( + failure(FailureKind::Declined, Some(400), &["UnsupportedEndpointError", "BadRequestError"]), + Some(7) + )] + #[case::malformed_reads_bad_request( + failure(FailureKind::Declined, Some(400), &["MalformedRequestFieldError", "BadRequestError"]), + Some(3) + )] + fn retries_follow_the_most_specific_class( + #[case] failure: AttemptError, + #[case] expected: Option, + ) { + let policy = RetryPolicy { + InternalServerErrorRetries: Some(4), + BadRequestErrorRetries: Some(3), + NotFoundErrorRetries: Some(1), + DefaultRetries: Some(7), + ..RetryPolicy::default() + }; + assert_eq!(policy.retries_for(&failure), expected); + } + + #[rstest] + #[case::rate_limit_with_healthy(FailureKind::RateLimit, Some(429), 1, 2, false, true)] + #[case::rate_limit_none_healthy_with_fallbacks( + FailureKind::RateLimit, + Some(429), + 0, + 2, + true, + false + )] + #[case::bad_request(FailureKind::BadRequest, Some(400), 2, 2, false, false)] + #[case::auth_lone_deployment(FailureKind::Authentication, Some(401), 1, 1, false, false)] + #[case::auth_with_siblings(FailureKind::Authentication, Some(401), 1, 2, false, true)] + #[case::not_found(FailureKind::NotFound, Some(404), 2, 2, false, false)] + #[case::server_error_no_healthy(FailureKind::InternalServer, Some(500), 0, 2, false, false)] + fn should_retry_follows_python( + #[case] kind: FailureKind, + #[case] status: Option, + #[case] healthy: usize, + #[case] all: usize, + #[case] has_regular_fallbacks: bool, + #[case] expected: bool, + ) { + let ctx = RetryContext { + healthy, + all, + has_regular_fallbacks, + ..RetryContext::default() + }; + assert_eq!( + should_retry(&AttemptError::new(kind, status), ctx), + expected + ); + } +} diff --git a/litellm-rust/crates/router-core/src/routing_groups.rs b/litellm-rust/crates/router-core/src/routing_groups.rs new file mode 100644 index 00000000000..fcf6c2febf3 --- /dev/null +++ b/litellm-rust/crates/router-core/src/routing_groups.rs @@ -0,0 +1,48 @@ +//! Routing groups: a callable name over the deployments of its member model groups +//! (`_get_routing_group_deployments`). A priority group serves each member's priority as the +//! member deployments' `order` (`apply_routing_group_priority`), so the order filter and order +//! fallbacks walk the priorities. + +use std::sync::Arc; + +use serde_json::Value; + +use crate::registry::{Deployment, RoutingGroupSpec}; + +/// The member deployments of `group`, in member order, with a priority group's priorities as +/// their `order`. +pub fn members( + group: &RoutingGroupSpec, + deployments_of: impl Fn(&str) -> Vec>, +) -> Vec> { + let priorities = group + .model_priorities + .as_ref() + .filter(|_| is_priority(group)); + group + .models + .iter() + .flat_map(|member| { + let priority = priorities.and_then(|p| p.get(member)).copied(); + deployments_of(member) + .into_iter() + .map(move |deployment| match priority { + Some(order) => with_order(&deployment, order), + None => deployment, + }) + }) + .collect() +} + +/// Whether `group` is a priority group. +pub fn is_priority(group: &RoutingGroupSpec) -> bool { + matches!(&group.routing_strategy, crate::config::RoutingStrategy::Other(name) if name == "priority") +} + +fn with_order(deployment: &Arc, order: i64) -> Arc { + let mut prioritized = Deployment::clone(deployment); + prioritized + .raw_params + .insert("order".to_owned(), Value::from(order)); + Arc::new(prioritized) +} diff --git a/litellm-rust/crates/router-core/src/routing_plugins.rs b/litellm-rust/crates/router-core/src/routing_plugins.rs new file mode 100644 index 00000000000..236de6fc521 --- /dev/null +++ b/litellm-rust/crates/router-core/src/routing_plugins.rs @@ -0,0 +1,114 @@ +//! `Router(plugins=[...])`: the routing-plugin pipeline. The host runs the plugins (live Python +//! objects) in the async pick's pre-routing hook and reports the candidate models they left on +//! the request (`_routing_plugin_candidate_models`); router-core decides which deployments that +//! keeps (`_filter_by_routing_plugin_candidates`) and refuses the paths that never run the +//! pipeline (`_get_available_deployment_unresolved`). + +use std::sync::Arc; + +use crate::{config::RoutingStrategy, error::Error, host::RouterRequest, registry::Deployment}; + +/// Whether the router's own strategy has an async selector: the others pick on the sync path +/// even for async callers (`_async_get_available_deployment`). +pub(crate) fn picks_async(strategy: &RoutingStrategy) -> bool { + matches!( + strategy, + RoutingStrategy::UsageBasedV2 + | RoutingStrategy::SimpleShuffle + | RoutingStrategy::CostBased + | RoutingStrategy::LatencyBased + | RoutingStrategy::LeastBusy + ) +} + +/// The sync pick never runs the pipeline, so a router with plugins refuses it rather than +/// letting a policy plugin be bypassed. `request` is a pick or a call, never a listing. +pub(crate) fn refuse_sync_pick( + has_plugins: bool, + strategy: &RoutingStrategy, + request: &RouterRequest, +) -> Result<(), Error> { + if has_plugins + && request.pass_through.is_none() + && (!request.pre_routes || !picks_async(strategy)) + { + return Err(Error::RoutingPluginsNeedAsync); + } + Ok(()) +} + +/// `_filter_by_routing_plugin_candidates`: keep the deployments whose litellm model the plugins +/// left; nothing left is an error, never a fallback to the unfiltered pool. +pub(crate) fn narrow( + healthy: Vec>, + models: Option<&[String]>, +) -> Result>, Error> { + let Some(models) = models else { + return Ok(healthy); + }; + let kept: Vec> = healthy + .into_iter() + .filter(|d| models.contains(&d.params.model)) + .collect(); + if kept.is_empty() { + return Err(Error::RoutingPluginsLeftNothing { + models: models.to_vec(), + }); + } + Ok(kept) +} + +#[cfg(test)] +mod tests { + use rstest::rstest; + use serde_json::json; + + use super::*; + use crate::registry::DeploymentSpec; + + fn deployment(id: &str, model: &str) -> Arc { + let spec: DeploymentSpec = serde_json::from_value(json!({ + "model_name": "g", "litellm_params": {"model": model}, "model_info": {"id": id}, + })) + .unwrap(); + Arc::new(Deployment::from_spec(spec).unwrap()) + } + + #[rstest] + fn the_plugins_models_narrow_the_pool() { + let pool = vec![deployment("a", "openai/a"), deployment("b", "openai/b")]; + let kept = narrow(pool, Some(&["openai/b".to_owned()])).unwrap(); + assert_eq!( + kept.iter().map(|d| d.id.as_str()).collect::>(), + ["b"] + ); + } + + #[rstest] + fn narrowing_to_nothing_is_an_error() { + let pool = vec![deployment("a", "openai/a")]; + assert!(matches!( + narrow(pool, Some(&[])), + Err(Error::RoutingPluginsLeftNothing { models }) if models.is_empty() + )); + } + + #[rstest] + #[case::async_pick(true, RoutingStrategy::SimpleShuffle, false)] + #[case::sync_pick(false, RoutingStrategy::SimpleShuffle, true)] + #[case::v1_strategy(true, RoutingStrategy::UsageBased, true)] + fn only_the_async_pick_runs_plugins( + #[case] pre_routes: bool, + #[case] strategy: RoutingStrategy, + #[case] refused: bool, + ) { + let request = RouterRequest { + pre_routes, + ..RouterRequest::default() + }; + assert_eq!( + refuse_sync_pick(true, &strategy, &request).is_err(), + refused + ); + } +} diff --git a/litellm-rust/crates/router-core/src/select.rs b/litellm-rust/crates/router-core/src/select.rs new file mode 100644 index 00000000000..fefa7f25ad1 --- /dev/null +++ b/litellm-rust/crates/router-core/src/select.rs @@ -0,0 +1,795 @@ +//! Selection: resolve a name, filter its deployments in the order of +//! `Router._core_async_get_healthy_deployments` (team, web search, cooldown, blocked, plugins, +//! pre-call checks, tags, order, retry skips), then pick one with the strategy. Simple-shuffle is +//! native (`router_strategy/simple_shuffle.py`); other strategies go through `Host::select`. + +use std::{collections::BTreeMap, sync::Arc}; + +use parking_lot::Mutex; +use rand::{Rng, SeedableRng, rngs::StdRng}; + +use crate::{ + Core, RouterStateSnapshot, + config::{Op, RoutingStrategy}, + error::Error, + health_routing::WhenAllUnhealthy, + host::{RouterRequest, RoutingContext, SelectStage}, + registry::{Deployment, Registry, ResolveKind}, + resolve::{Caller, resolve_request}, + strategy::{StrategyScope, routing_context}, +}; + +/// Randomness for the shuffle and retry jitter, injectable so tests are deterministic. +pub trait Randomness: Send + Sync { + /// A float in `[0, 1)`, as Python's `random.random()`. + fn next_f64(&self) -> f64; + + /// An index below `len`, as Python's `random.choice` draws one. + fn index_below(&self, len: usize) -> usize { + ((self.next_f64() * len as f64) as usize).min(len - 1) + } +} + +pub struct ThreadRandomness; + +impl Randomness for ThreadRandomness { + fn next_f64(&self) -> f64 { + rand::thread_rng().r#gen() + } +} + +pub struct SeededRandomness(Mutex); + +impl SeededRandomness { + pub fn new(seed: u64) -> Self { + Self(Mutex::new(StdRng::seed_from_u64(seed))) + } +} + +impl Randomness for SeededRandomness { + fn next_f64(&self) -> f64 { + self.0.lock().r#gen() + } +} + +/// Which pass-through pick a request makes. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum PassThrough { + /// `get_available_deployment_for_pass_through` (and `select(kind="pass_through")`): the + /// pass-through deployments are chosen first, then cooldown and blocked filters, and only a + /// priority group filters by order. + Listed, + /// `async_get_available_deployment_for_pass_through`: the full healthy chain, then the + /// pass-through filter. + Healthy, +} + +/// A group's deployments before and after filtering. +#[derive(Clone, Debug)] +pub struct Candidates { + /// The name that was resolved (Python's `model`). + pub name: String, + pub model_group: String, + pub kind: ResolveKind, + /// Every deployment the caller may use (after the team filter). + pub all: Vec>, + pub healthy: Vec>, + /// The provider-prefixed spelling the name resolved under (`get_pattern`), when the host + /// was asked for one. + pub wildcard_request: Option, +} + +impl Candidates { + /// The `model` a no-deployments error names: what `_common_checks_available_deployment` + /// returned, so the alias-resolved name a wildcard serves and not its pattern. + pub(crate) fn failed_model(&self, registry: &Registry) -> String { + let model = match self.kind { + ResolveKind::Wildcard => registry.alias_of(&self.name).unwrap_or(&self.name), + _ => self.model_group.as_str(), + }; + model.to_owned() + } +} + +/// Per-attempt selection inputs the call loop carries. +#[derive(Clone, Debug, Default)] +pub(crate) struct SelectInputs { + pub target_order: Option, + pub skipped_ids: Vec, + /// A weighted failover hop's `_excluded_deployment_ids`: never picked. + pub excluded_ids: Arc<[String]>, + /// Stop before the plugins, order and skip filters: the group less cooldowns and blocked + /// deployments, which `_async_get_healthy_deployments` counts when deciding whether to retry. + pub count_only: bool, + /// The level's name, which an endpoint refusal reads (`endpoint_selection(model=...)`): the + /// name before pre-routing replaced it. + pub requested: Option, + /// A fallback hop: `run_async_fallback` pops the request's `_target_order`, so only the + /// hop's own order applies. + pub fallback_hop: bool, +} + +/// `simple_shuffle`: weighted by `_router_weights[group]`, else `weight`, else `rpm`, else +/// `tpm`, each normalized by its maximum, else uniform. A metric is read only when every set +/// before it weighed nothing, and a non-numeric one fails the pick then (`_metric_weight`). +pub fn simple_shuffle( + deployments: &[Arc], + model_group: &str, + router_weights: Option<&BTreeMap>>, + rng: &dyn Randomness, +) -> Result>, Error> { + if deployments.is_empty() { + return Ok(None); + } + let scoped = router_weights + .and_then(|weights| weights.get(model_group)) + .filter(|weights| !weights.is_empty()) + .map(|weights| { + deployments + .iter() + .map(|d| weights.get(&d.id).copied().unwrap_or(0.0)) + .collect::>() + }); + let metric_sets = ["weight", "rpm", "tpm"].into_iter().map(|metric| { + deployments + .iter() + .map(|d| d.metric_weight(metric)) + .collect::>>() + .ok_or(Error::NonNumericWeight { metric }) + }); + for weights in scoped.into_iter().map(Ok).chain(metric_sets) { + let weights = weights?; + let largest = weights.iter().copied().fold(0.0, f64::max); + if largest <= 0.0 { + continue; + } + let normalized: Vec = weights.iter().map(|weight| weight / largest).collect(); + let total: f64 = normalized.iter().sum(); + if total <= 0.0 { + continue; + } + // `random.choices`: bisect_right over the cumulative weights. + let target = rng.next_f64() * total; + let mut cumulative = 0.0; + let index = normalized + .iter() + .position(|weight| { + cumulative += weight; + target < cumulative + }) + .unwrap_or(deployments.len() - 1); + return Ok(Some(deployments[index].clone())); + } + Ok(Some( + deployments[rng.index_below(deployments.len())].clone(), + )) +} + +/// A pre-routing strategy marker (`classify_strategy_router_model`): its model is `auto_router/...`. +pub fn is_strategy_marker(deployment: &Deployment) -> bool { + deployment.params.model.starts_with("auto_router/") +} + +/// `litellm_params.use_in_pass_through`, read for truthiness as Python does. +fn uses_pass_through(deployment: &Deployment) -> bool { + match deployment.raw_params.get("use_in_pass_through") { + None | Some(serde_json::Value::Null) => false, + Some(serde_json::Value::Bool(value)) => *value, + Some(serde_json::Value::Number(value)) => value.as_f64().is_some_and(|n| n != 0.0), + Some(serde_json::Value::String(value)) => !value.is_empty(), + Some(serde_json::Value::Array(value)) => !value.is_empty(), + Some(serde_json::Value::Object(value)) => !value.is_empty(), + } +} + +/// `filter_team_based_models`: a team-owned deployment serves only its team, except that a proxy +/// admin without a team keeps deployments that are all team-owned and listed under the model the +/// request names, unless they belong to several teams. +fn filter_team_based( + deployments: Vec>, + request: &RouterRequest, +) -> Result>, Error> { + let team_id = request.team_id.as_deref(); + if team_id.is_none() + && request.proxy_admin + && let Some(requested) = request.admin_model.as_deref() + && !deployments.is_empty() + && deployments.iter().all(|d| { + d.model_info.team_id.is_some() + && (d.model_name == requested + || d.model_info.team_public_model_name.as_deref() == Some(requested)) + }) + { + let teams: std::collections::BTreeSet<&str> = deployments + .iter() + .filter_map(|d| d.model_info.team_id.as_deref()) + .collect(); + if teams.len() > 1 { + return Err(Error::MultipleTeams { + model: requested.to_owned(), + }); + } + return Ok(deployments); + } + Ok(deployments + .into_iter() + .filter(|d| d.usable_by_team(team_id)) + .collect()) +} + +impl Core { + fn context<'a>( + &'a self, + snapshot: &'a RouterStateSnapshot, + op: &'a Op, + request: &'a RouterRequest, + model_group: &'a str, + name: &'a str, + scope: StrategyScope, + ) -> RoutingContext<'a> { + RoutingContext { + registry: &snapshot.registry, + settings: &snapshot.settings, + health: &self.health, + request, + op, + model_group, + name, + scope, + pool: &[], + excluded_ids: &[], + skipped_ids: &[], + target_order: None, + } + } + + /// The eligible deployments for `name`, in filter-chain order. + pub(crate) async fn candidates( + &self, + snapshot: &RouterStateSnapshot, + op: &Op, + request: &RouterRequest, + name: &str, + inputs: &SelectInputs, + ) -> Result { + let registry = &snapshot.registry; + // The async pick's pre-routing hook runs the routing plugins before anything resolves; + // every pick but a listed pass-through one reads what they left on the request. + let plugin_models = if snapshot.settings.routing_plugins + && !inputs.count_only + && request.pass_through != Some(PassThrough::Listed) + { + let ctx = self.context(snapshot, op, request, name, name, StrategyScope::Default); + self.host + .routing_plugin_models(&ctx, request.pre_routes) + .await? + } else { + None + }; + // No deployment for the name: route to the first default fallback. + let default_target = crate::resolve::default_target(&snapshot.settings.dynamic); + // `get_pattern`: a name no pattern matches is tried again under its inferred provider. + let wildcard_name = + crate::resolve::wildcard_name(registry, name, request.specific_deployment); + let wildcard_request = match wildcard_name { + Some(model) => { + let scope = StrategyScope::Default; + let ctx = self.context(snapshot, op, request, model, model, scope); + self.host.provider_request(&ctx).await? + } + None => None, + }; + // `endpoint_selection` wraps the attempt, not the count a retry decision reads. + let endpoint = request.endpoint.as_deref().filter(|_| !inputs.count_only); + let caller = Caller { + team_id: request.team_id.as_deref(), + proxy_admin: request.proxy_admin, + allowed_models: &request.allowed_models, + endpoint, + requested: inputs.requested.as_deref().unwrap_or(name), + wildcard_request: wildcard_request.as_deref(), + }; + let resolution = resolve_request( + registry, + name, + request.specific_deployment, + caller, + default_target, + snapshot.settings.dynamic.default_ref(), + jiff::Timestamp::now(), + )? + .ok_or_else(|| Error::UnknownModel { + // The name `_common_checks_available_deployment` looked for, past the alias. + model: registry.alias_of(name).unwrap_or(name).to_owned(), + })?; + let first_model = resolution + .deployments + .first() + .map(|d| d.params.model.clone()); + let team_id = request.team_id.as_deref(); + let usable = filter_team_based(resolution.deployments.clone(), request)?; + // `_filter_reserved_deployments`: PTU shares keep holders to their share and everyone + // else to the deployments split into none. + let (mut all, withheld) = crate::ptu::filter_shared(usable, team_id); + if withheld && all.is_empty() { + return Err(Error::PtuReserved { + model: match resolution.kind { + ResolveKind::Id => first_model.unwrap_or_else(|| name.to_owned()), + ResolveKind::Wildcard | ResolveKind::SpecificDeployment => name.to_owned(), + _ => resolution.group.clone(), + }, + }); + } + // `filter_web_search_deployments` runs before `request_deployments` is read, so every later + // stage sees the group without the deployments it dropped. + if request.web_search && !inputs.count_only && !resolution.kind.is_single() { + all.retain(|d| d.supports_web_search()); + } + // `_without_deployments_that_cannot_serve_route`, right after the web search filter: the + // deployments whose provider cannot serve the call, unless that leaves nothing. Only a pick + // runs it (`_route_capability_pin`): the retry count reads the whole group. + if let Some(handler) = request.route_handler.as_deref() + && !inputs.count_only + && !resolution.kind.is_single() + && !all.is_empty() + { + let ctx = self.context( + snapshot, + op, + request, + &resolution.group, + name, + StrategyScope::Default, + ); + let stage = SelectStage::RouteIncapable { + handler: handler.to_owned(), + }; + let incapable = self.host.select(&stage, &ctx, all.clone()).await?; + let capable: Vec> = all + .iter() + .filter(|d| !incapable.iter().any(|i| i.id == d.id)) + .cloned() + .collect(); + if !capable.is_empty() { + all = capable; + } + } + let listed = + request.pass_through == Some(PassThrough::Listed) && !resolution.kind.is_single(); + if listed { + all.retain(|d| uses_pass_through(d)); + if all.is_empty() { + // The name `_common_checks_available_deployment` resolved: an alias's target, a + // routing group's name, the alias-resolved name a wildcard serves. + let model = match resolution.kind { + ResolveKind::Wildcard => registry.alias_of(name).unwrap_or(name), + ResolveKind::SpecificDeployment => name, + _ => resolution.group.as_str(), + }; + return Err(Error::NoPassThroughDeployments { + model: model.to_owned(), + }); + } + } + + // A deployment named by its id is served even while it cools down + // (`get_deployment_by_model_id` returns before the cooldown filter). + let settings = &snapshot.settings; + let cooling = if resolution.kind.is_single() { + Vec::new() + } else { + // `async_get_cooldown_deployments` reads every deployment's, as one throttled batch. + self.health + .active_cooldowns(&snapshot.registry.ids()) + .await? + }; + // Health-check filtering runs before the cooldown filter, except for a deployment named + // by its id and the healthy count a retry decision reads. + let checked = if resolution.kind.is_single() || inputs.count_only { + all.clone() + } else { + let when_all_unhealthy = if request.health_check_probe { + WhenAllUnhealthy::Nothing + } else { + WhenAllUnhealthy::KeepAll + }; + self.health + .health_check_filter(settings, all.clone(), when_all_unhealthy) + .await + }; + let mut healthy: Vec> = checked + .iter() + .filter(|d| !cooling.iter().any(|(id, _)| *id == d.id)) + .cloned() + .collect(); + // The safety net when health-check routing drives cooldowns (an allowed_fails_policy): + // never on a listed pass-through pick or the healthy count. + if healthy.is_empty() + && !listed + && !inputs.count_only + && settings.enable_health_check_routing + && settings.allowed_fails_policy.is_some() + { + healthy = checked; + } + if resolution.kind.is_single() + && let Some(deployment) = all.first() + { + // `_core_async_get_healthy_deployments`: a deployment named by id is checked alone. + if deployment.model_info.blocked { + return Err(Error::Paused { + model: deployment.params.model.clone(), + }); + } + if request.pass_through.is_some() && !uses_pass_through(deployment) { + return Err(Error::NotPassThrough { + id: deployment.id.clone(), + }); + } + } + healthy.retain(|d| !d.model_info.blocked); + // `continuation_deployments`: a mid-stream continuation never lands on a deployment the host + // named as unable to take its prefill. + if let Some(excluded) = request + .body + .get("continuation_excluded_ids") + .and_then(serde_json::Value::as_array) + { + healthy.retain(|d| !excluded.iter().any(|id| id.as_str() == Some(d.id.as_str()))); + } + let request_order = request.target_order.filter(|_| !inputs.fallback_hop); + // `async_callback_filter_deployments` runs after the cooldown and blocked filters on every + // async candidate list but a deployment named by id and a listed pass-through pick. + if settings.callback_filters + && !inputs.count_only + && !resolution.kind.is_single() + && request.pass_through != Some(PassThrough::Listed) + { + let mut ctx = self.context( + snapshot, + op, + request, + &resolution.group, + name, + StrategyScope::Default, + ); + ctx.excluded_ids = &inputs.excluded_ids; + ctx.skipped_ids = &inputs.skipped_ids; + ctx.target_order = inputs.target_order.or(request_order); + healthy = self.host.callback_filters(&ctx, healthy).await?; + } + if inputs.count_only { + return Ok(Candidates { + name: name.to_owned(), + model_group: resolution.group, + kind: resolution.kind, + all, + healthy, + wildcard_request, + }); + } + + if listed { + if settings.enable_pre_call_checks && !resolution.kind.is_single() { + healthy = self + .pre_call_checks( + snapshot, + op, + request, + &resolution.group, + name, + &all, + healthy, + ) + .await?; + } + // Only a priority group narrows a listed pass-through pick by order. + let priority = resolution.kind == ResolveKind::RoutingGroup + && registry.routing_groups().iter().any(|g| { + g.group_name == resolution.group && crate::routing_groups::is_priority(g) + }); + if priority { + let order = request + .target_order + .or_else(|| healthy.iter().filter_map(|d| d.order()).min()); + if let Some(order) = order { + healthy.retain(|d| d.order() == Some(order)); + } + } + return Ok(Candidates { + name: name.to_owned(), + model_group: resolution.group, + kind: resolution.kind, + all, + healthy, + wildcard_request, + }); + } + + if !settings.plugins.is_empty() && !healthy.is_empty() { + let ctx = self.context( + snapshot, + op, + request, + &resolution.group, + name, + StrategyScope::Default, + ); + healthy = self + .host + .select(&SelectStage::Plugins, &ctx, healthy) + .await?; + } + + if settings.enable_pre_call_checks && !resolution.kind.is_single() { + healthy = self + .pre_call_checks( + snapshot, + op, + request, + &resolution.group, + name, + &all, + healthy, + ) + .await?; + } + + if request.filter_tags && !healthy.is_empty() { + let mut ctx = self.context( + snapshot, + op, + request, + &resolution.group, + name, + StrategyScope::Default, + ); + ctx.pool = &all; + healthy = self.host.select(&SelectStage::Tags, &ctx, healthy).await?; + } + + if !resolution.kind.is_single() { + healthy = crate::routing_plugins::narrow(healthy, plugin_models.as_deref())?; + } + + // `_without_declared_text_only_deployments`: an image request skips the deployments that + // declare no vision, unless that leaves nothing. + if request.has_image && !resolution.kind.is_single() && !healthy.is_empty() { + let ctx = self.context( + snapshot, + op, + request, + &resolution.group, + name, + StrategyScope::Default, + ); + let text_only = self + .host + .select(&SelectStage::TextOnly, &ctx, healthy.clone()) + .await?; + let vision: Vec> = healthy + .iter() + .filter(|d| !text_only.iter().any(|t| t.id == d.id)) + .cloned() + .collect(); + if !vision.is_empty() { + healthy = vision; + } + } + + // `_without_decisions_deployments_refusing_media`: a Decisions call carrying audio or + // videos skips the deployments that refuse them, unless that leaves nothing. + if request.decisions_media && !resolution.kind.is_single() && !healthy.is_empty() { + let ctx = self.context( + snapshot, + op, + request, + &resolution.group, + name, + StrategyScope::Default, + ); + let refusing = self + .host + .select(&SelectStage::MediaRefusing, &ctx, healthy.clone()) + .await?; + let taking: Vec> = healthy + .iter() + .filter(|d| !refusing.iter().any(|r| r.id == d.id)) + .cloned() + .collect(); + if !taking.is_empty() { + healthy = taking; + } + } + + if let Some(pinned) = request.pinned_deployment_id.as_deref() + && healthy.iter().any(|d| d.id == pinned) + { + healthy.retain(|d| d.id == pinned); + } + + let order = inputs + .target_order + .or(request_order) + .or_else(|| healthy.iter().filter_map(|d| d.order()).min()); + if let Some(order) = order { + healthy.retain(|d| d.order() == Some(order)); + } + + healthy.retain(|d| { + !request.excluded_ids.contains(&d.id) && !inputs.excluded_ids.contains(&d.id) + }); + + if !inputs.skipped_ids.is_empty() || !request.retry_skipped_ids.is_empty() { + let kept: Vec> = healthy + .iter() + .filter(|d| { + !inputs.skipped_ids.contains(&d.id) + && !request.retry_skipped_ids.contains(&d.id) + }) + .cloned() + .collect(); + if !kept.is_empty() { + healthy = kept; + } + } + + if request.pass_through == Some(PassThrough::Healthy) + && !resolution.kind.is_single() + && !healthy.is_empty() + { + healthy.retain(|d| uses_pass_through(d)); + if healthy.is_empty() { + return Err(Error::NoPassThroughDeployments { + model: name.to_owned(), + }); + } + } + + Ok(Candidates { + name: name.to_owned(), + model_group: resolution.group, + kind: resolution.kind, + all, + healthy, + wildcard_request, + }) + } + + /// `_pre_call_checks` over `healthy`, with the host's facts. + #[allow(clippy::too_many_arguments)] + async fn pre_call_checks( + &self, + snapshot: &RouterStateSnapshot, + op: &Op, + request: &RouterRequest, + model_group: &str, + name: &str, + pool: &[Arc], + healthy: Vec>, + ) -> Result>, Error> { + let mut ctx = self.context( + snapshot, + op, + request, + model_group, + name, + StrategyScope::Default, + ); + ctx.pool = pool; + let facts = self.host.pre_call_facts(&ctx, &healthy).await?; + crate::precall::filter(model_group, healthy, pool, &facts) + } + + /// Pick one of `candidates.healthy` with the group's strategy. + pub(crate) async fn pick( + &self, + snapshot: &RouterStateSnapshot, + op: &Op, + request: &RouterRequest, + candidates: &Candidates, + ) -> Result, Error> { + let no_deployment = || Error::NoHealthyDeployments { + model: candidates.failed_model(&snapshot.registry), + }; + // A deployment named by its id is returned as is, before any strategy runs. + if candidates.kind.is_single() { + return candidates + .healthy + .first() + .cloned() + .ok_or_else(no_deployment); + } + let (strategy, scope) = routing_context( + &snapshot.registry, + &snapshot.settings.routing_strategy, + request.routing_strategy.as_deref(), + &candidates.kind, + &candidates.model_group, + ); + if strategy == RoutingStrategy::Unset { + return Err(no_deployment()); + } + if strategy == RoutingStrategy::SimpleShuffle { + let weights = crate::weights::validate(request.router_weights.as_ref()) + .map_err(|_| Error::InvalidRouterWeights)?; + return simple_shuffle( + &candidates.healthy, + &candidates.model_group, + weights.as_ref(), + self.rng.as_ref(), + )? + .ok_or_else(no_deployment); + } + let ctx = self.context( + snapshot, + op, + request, + &candidates.model_group, + &candidates.name, + scope, + ); + self.host + .select( + &SelectStage::Strategy(strategy), + &ctx, + candidates.healthy.clone(), + ) + .await? + .into_iter() + .next() + .ok_or_else(no_deployment) + } +} + +#[cfg(test)] +mod tests { + use rstest::rstest; + use serde_json::json; + + use super::*; + use crate::registry::DeploymentSpec; + + struct Fixed(f64); + + impl Randomness for Fixed { + fn next_f64(&self) -> f64 { + self.0 + } + } + + fn deployment(id: &str, params: serde_json::Value) -> Arc { + let spec: DeploymentSpec = serde_json::from_value(json!({ + "model_name": "g", "litellm_params": params, "model_info": {"id": id}, + })) + .unwrap(); + Arc::new(Deployment::from_spec(spec).unwrap()) + } + + #[rstest] + #[case::weight_low_draw(0.1, "a")] + #[case::weight_high_draw(0.9, "b")] + fn shuffle_uses_weight_before_rpm(#[case] draw: f64, #[case] expected: &str) { + let deployments = [ + deployment("a", json!({"model": "openai/a", "weight": 1, "rpm": 1000})), + deployment("b", json!({"model": "openai/b", "weight": 3})), + ]; + let picked = simple_shuffle(&deployments, "g", None, &Fixed(draw)) + .unwrap() + .unwrap(); + assert_eq!(picked.id, expected); + } + + #[rstest] + fn shuffle_prefers_request_weights() { + let deployments = [ + deployment("a", json!({"model": "openai/a", "weight": 100})), + deployment("b", json!({"model": "openai/b"})), + ]; + let weights = BTreeMap::from([("g".to_owned(), BTreeMap::from([("b".to_owned(), 1.0)]))]); + let picked = simple_shuffle(&deployments, "g", Some(&weights), &Fixed(0.99)) + .unwrap() + .unwrap(); + assert_eq!(picked.id, "b"); + } +} diff --git a/litellm-rust/crates/router-core/src/settings.rs b/litellm-rust/crates/router-core/src/settings.rs new file mode 100644 index 00000000000..f55b059d71c --- /dev/null +++ b/litellm-rust/crates/router-core/src/settings.rs @@ -0,0 +1,133 @@ +//! [`Settings`]: one typed, frozen view of the configuration a caller may read, built from the +//! static config and the current dynamic snapshot. + +use std::{collections::BTreeMap, sync::Arc, time::Duration}; + +use crate::{ + config::{ + AllowedFailsPolicy, DEFAULT_COOLDOWN_TIME, PluginSpec, RetryPolicy, RouterConfig, + RouterDynamicConfig, RoutingStrategy, + }, + error::AttemptError, +}; + +#[derive(Clone, Debug)] +pub struct Settings { + pub routing_strategy: RoutingStrategy, + pub plugins: Arc<[PluginSpec]>, + pub cooldown_time: Duration, + pub disable_cooldowns: bool, + pub allowed_fails: Option, + pub allowed_fails_policy: Option, + pub retry_policy: Option, + pub model_group_retry_policy: Arc>, + pub num_retries: u32, + pub max_fallbacks: u32, + pub retry_after: Duration, + pub timeout: Option, + pub drop_params: Option, + pub enable_pre_call_checks: bool, + pub routing_plugins: bool, + pub callback_filters: bool, + pub enable_health_check_routing: bool, + pub health_check_staleness: Duration, + pub background_health_check_model_groups: Option>, + pub enable_weighted_failover: bool, + pub fallback_on_unsupported_params: bool, + pub provider_default_deployment_ids: Arc<[String]>, + pub dynamic: Arc, +} + +/// `DEFAULT_HEALTH_CHECK_INTERVAL * DEFAULT_HEALTH_CHECK_STALENESS_MULTIPLIER`. +pub const DEFAULT_HEALTH_CHECK_STALENESS: Duration = Duration::from_secs(300 * 2); + +fn seconds(value: f64) -> Duration { + Duration::try_from_secs_f64(value).unwrap_or_default() +} + +impl Settings { + pub fn new(config: &RouterConfig, dynamic: Arc) -> Self { + Self { + routing_strategy: config.routing_strategy.clone(), + plugins: config.plugins.clone().into(), + cooldown_time: config.cooldown_time.map_or(DEFAULT_COOLDOWN_TIME, seconds), + disable_cooldowns: config.disable_cooldowns, + allowed_fails: dynamic.allowed_fails, + allowed_fails_policy: config.allowed_fails_policy.clone(), + retry_policy: config.retry_policy.clone(), + model_group_retry_policy: Arc::new(config.model_group_retry_policy.clone()), + num_retries: dynamic.num_retries.unwrap_or(0), + max_fallbacks: dynamic.max_fallbacks, + retry_after: seconds(config.retry_after), + timeout: config.request_timeout.or(config.timeout).map(seconds), + drop_params: config.drop_params, + enable_pre_call_checks: config.enable_pre_call_checks, + routing_plugins: config.routing_plugins, + callback_filters: config.callback_filters, + enable_health_check_routing: config.enable_health_check_routing, + health_check_staleness: config + .health_check_staleness_threshold + .map_or(DEFAULT_HEALTH_CHECK_STALENESS, seconds), + background_health_check_model_groups: config + .background_health_check_model_groups + .clone() + .map(Into::into), + enable_weighted_failover: config.enable_weighted_failover, + fallback_on_unsupported_params: config.fallback_on_unsupported_params.unwrap_or(true), + provider_default_deployment_ids: config.provider_default_deployment_ids.clone().into(), + dynamic, + } + } + + /// `allowed_fails_for`: the policy's count for this failure. + pub fn allowed_fails_for(&self, failure: &AttemptError) -> Option { + self.allowed_fails_policy + .as_ref() + .and_then(|policy| policy.allowed_fails_for(failure)) + } + + /// How long an allowed-fails counter lives: `settings().cooldown_time or + /// DEFAULT_COOLDOWN_TIME_SECONDS`, so a router cooldown of 0 still counts over the default. + pub fn fails_window(&self) -> Duration { + if self.cooldown_time.is_zero() { + DEFAULT_COOLDOWN_TIME + } else { + self.cooldown_time + } + } + + /// `_is_allowed_fails_set_on_router`. + pub fn allowed_fails_set(&self) -> bool { + self.allowed_fails.is_some() && !self.dynamic.allowed_fails_is_default + } + + /// Python's `allowed_fails` default when set without a value (`litellm.allowed_fails`). + pub fn effective_allowed_fails(&self) -> u32 { + self.allowed_fails.unwrap_or(3) + } +} + +#[cfg(test)] +mod tests { + use rstest::rstest; + + use super::*; + + #[rstest] + #[case::passed_to_the_router(Some(3), false, true)] + #[case::the_process_default(Some(3), true, false)] + #[case::none(None, false, false)] + fn allowed_fails_counts_as_set_unless_it_is_the_process_default( + #[case] allowed_fails: Option, + #[case] allowed_fails_is_default: bool, + #[case] expected: bool, + ) { + let dynamic = RouterDynamicConfig { + allowed_fails, + allowed_fails_is_default, + ..RouterDynamicConfig::default() + }; + let settings = Settings::new(&RouterConfig::default(), Arc::new(dynamic)); + assert_eq!(settings.allowed_fails_set(), expected); + } +} diff --git a/litellm-rust/crates/router-core/src/strategy.rs b/litellm-rust/crates/router-core/src/strategy.rs new file mode 100644 index 00000000000..0e1f3b49fa8 --- /dev/null +++ b/litellm-rust/crates/router-core/src/strategy.rs @@ -0,0 +1,110 @@ +//! Which routing strategy picks for a request, and whose state it reads: a port of +//! `Router._get_routing_context`. A per-request `routing_strategy` override wins, then the +//! routing group the resolved model belongs to (called as the group, or as a member of a +//! non-priority group), then the router's own strategy. + +use crate::{ + config::RoutingStrategy, + registry::{Registry, ResolveKind}, +}; + +/// The per-request overrides `_OVERRIDABLE_ROUTING_STRATEGIES` honours: simple-shuffle and every +/// strategy with a default selector. Anything else is ignored. +pub const OVERRIDABLE: [&str; 6] = [ + "simple-shuffle", + "least-busy", + "usage-based-routing", + "usage-based-routing-v2", + "latency-based-routing", + "cost-based-routing", +]; + +/// Whose selector state a strategy pick reads. +#[derive(Clone, Debug, PartialEq, Eq)] +pub enum StrategyScope { + /// The router's own strategy (the implicit `"default"` group). + Default, + /// A routing group's own selector. + Group(String), + /// A per-request override: the default selector when it names the router's strategy, else + /// one selector per strategy kept for the router's life. + Override, +} + +impl StrategyScope { + /// The scope as the host step spells it. + pub fn wire(&self) -> (&'static str, Option<&str>) { + match self { + Self::Default => ("default", None), + Self::Group(name) => ("group", Some(name)), + Self::Override => ("override", None), + } + } +} + +/// A request's `routing_strategy` override, when it is one Python honours. +pub fn request_override(raw: Option<&str>) -> Option { + let raw = raw?; + OVERRIDABLE + .contains(&raw) + .then(|| raw.parse().unwrap_or_default()) +} + +/// The strategy and selector scope for a pick of `model_group` (the resolved group, after +/// aliases), reached through a resolution of `kind`. +pub fn routing_context( + registry: &Registry, + default: &RoutingStrategy, + override_strategy: Option<&str>, + kind: &ResolveKind, + model_group: &str, +) -> (RoutingStrategy, StrategyScope) { + if let Some(strategy) = request_override(override_strategy) { + return (strategy, StrategyScope::Override); + } + let group = match kind { + ResolveKind::RoutingGroup => registry + .routing_groups() + .iter() + .find(|group| group.group_name == model_group), + ResolveKind::Alias | ResolveKind::Group => registry.member_routing_group(model_group), + ResolveKind::SpecificDeployment + | ResolveKind::Id + | ResolveKind::Default + | ResolveKind::Wildcard => None, + }; + match group { + None => (default.clone(), StrategyScope::Default), + Some(group) if is_priority(&group.routing_strategy) => { + (RoutingStrategy::SimpleShuffle, StrategyScope::Default) + } + Some(group) => ( + group.routing_strategy.clone(), + StrategyScope::Group(group.group_name.clone()), + ), + } +} + +fn is_priority(strategy: &RoutingStrategy) -> bool { + matches!(strategy, RoutingStrategy::Other(name) if name == "priority") +} + +#[cfg(test)] +mod tests { + use rstest::rstest; + + use super::*; + + #[rstest] + #[case::honoured(Some("least-busy"), Some(RoutingStrategy::LeastBusy))] + #[case::shuffle(Some("simple-shuffle"), Some(RoutingStrategy::SimpleShuffle))] + #[case::lar1_ignored(Some("lar1"), None)] + #[case::unknown_ignored(Some("bogus"), None)] + #[case::unset(None, None)] + fn only_overridable_strategies_override( + #[case] raw: Option<&str>, + #[case] expected: Option, + ) { + assert_eq!(request_override(raw), expected); + } +} diff --git a/litellm-rust/crates/router-core/src/weights.rs b/litellm-rust/crates/router-core/src/weights.rs new file mode 100644 index 00000000000..3e5d76844ab --- /dev/null +++ b/litellm-rust/crates/router-core/src/weights.rs @@ -0,0 +1,81 @@ +//! `_router_weights` as `validate_router_weights` reads it (`litellm/types/router_weights.py`): +//! model group to deployment id to a finite weight of at least 0 (an int or a float, never a +//! bool or a string), with non-blank ids and at least one positive weight in every non-empty group. + +use std::collections::BTreeMap; + +use serde_json::Value; + +/// Validated request weights. +pub type RouterWeights = BTreeMap>; + +/// The request's weights didn't validate; the host raises the validator's own error. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub struct InvalidRouterWeights; + +fn identifier(text: &str) -> bool { + text.chars().any(|c| !c.is_whitespace()) +} + +fn weight(value: &Value) -> Option { + match value { + Value::Number(number) => number.as_f64().filter(|w| w.is_finite() && *w >= 0.0), + _ => None, + } +} + +/// `validate_router_weights(raw)`: `None` when the request passed none. +pub fn validate(raw: Option<&Value>) -> Result, InvalidRouterWeights> { + let groups = match raw { + None | Some(Value::Null) => return Ok(None), + Some(Value::Object(groups)) => groups, + Some(_) => return Err(InvalidRouterWeights), + }; + let mut validated = RouterWeights::new(); + for (group, weights) in groups { + let Value::Object(weights) = weights else { + return Err(InvalidRouterWeights); + }; + if !identifier(group) { + return Err(InvalidRouterWeights); + } + let mut entries = BTreeMap::new(); + for (id, value) in weights { + if !identifier(id) { + return Err(InvalidRouterWeights); + } + entries.insert(id.clone(), weight(value).ok_or(InvalidRouterWeights)?); + } + if !entries.is_empty() && !entries.values().any(|w| *w > 0.0) { + return Err(InvalidRouterWeights); + } + validated.insert(group.clone(), entries); + } + Ok(Some(validated)) +} + +#[cfg(test)] +mod tests { + use rstest::rstest; + use serde_json::json; + + use super::*; + + #[rstest] + #[case::all_zero(json!({"g": {"a": 0.0, "b": 0.0}}))] + #[case::negative(json!({"g": {"a": -1.0}}))] + #[case::text(json!({"g": {"a": "1"}}))] + #[case::boolean(json!({"g": {"a": true}}))] + #[case::blank_id(json!({"g": {" ": 1.0}}))] + #[case::not_a_mapping(json!(["g"]))] + fn invalid_weights_are_refused(#[case] raw: Value) { + assert_eq!(validate(Some(&raw)), Err(InvalidRouterWeights)); + } + + #[rstest] + #[case::empty_group(json!({"g": {}}))] + #[case::int_weight(json!({"g": {"a": 1, "b": 0}}))] + fn valid_weights_are_kept(#[case] raw: Value) { + assert!(validate(Some(&raw)).unwrap().is_some()); + } +} diff --git a/litellm-rust/crates/router-core/tests/call_loop.rs b/litellm-rust/crates/router-core/tests/call_loop.rs new file mode 100644 index 00000000000..9f41371f255 --- /dev/null +++ b/litellm-rust/crates/router-core/tests/call_loop.rs @@ -0,0 +1,1757 @@ +//! A call through the router end to end: selection, retries, fallbacks and cooldowns against a +//! scripted host. + +use std::{collections::VecDeque, sync::Arc, time::Duration}; + +use async_trait::async_trait; +use litellm_router_core::{ + Attempt, AttemptError, Deployment, Error, FailureKind, FailureReport, Host, + InMemoryRouterCache, Op, Output, Payload, Router, RouterCache, RouterConfig, + RouterDynamicConfig, RouterEvent, RouterRequest, RoutingContext, SeededRandomness, SelectStage, + SuccessReport, + call::StreamOrigin, + fallback::Fallbacks, + health::{UsageMetric, current_minute, usage_key}, + host::UncountedRetries, + select::PassThrough, +}; +use parking_lot::Mutex; +use rstest::{fixture, rstest}; +use serde_json::json; + +#[derive(Default)] +struct ScriptedHost { + /// Failures handed out in order, keyed by model group; a group with none left succeeds. + script: Mutex)>>, + events: Mutex>, + cache: InMemoryRouterCache, + /// Answer every success as a stream whose outcome the host reports when it ends. + deferred: bool, + /// [`Output::tracks_usage`] on every success. + tracks_usage: Option, + /// Refuse every fallback chain with a dict target. + refuses_dict_targets: bool, + /// Every fallback check raises this exception class, as a host's sync check does. + fallback_check_raises: Option<&'static str>, + /// Failures the tag stage raises in turn, before it answers. + pick_failures: Mutex>, + /// The next attempt's failure, whatever group it serves. + first_failure: Mutex>, +} + +impl ScriptedHost { + fn failing(group: &str, failures: Vec) -> Self { + Self::scripted(vec![(group, failures)]) + } + + fn scripted(groups: Vec<(&str, Vec)>) -> Self { + Self { + script: Mutex::new( + groups + .into_iter() + .map(|(group, failures)| (group.to_owned(), failures.into())) + .collect(), + ), + ..Self::default() + } + } + + fn selected(&self) -> Vec { + self.events + .lock() + .iter() + .filter_map(|event| match event { + RouterEvent::Selected { deployment_id, .. } => Some(deployment_id.clone()), + _ => None, + }) + .collect() + } + + fn retries(&self) -> usize { + self.events + .lock() + .iter() + .filter(|event| matches!(event, RouterEvent::Retry { .. })) + .count() + } +} + +#[async_trait] +impl Host for ScriptedHost { + async fn invoke(&self, attempt: &Attempt) -> Result { + if let Some(failure) = self.first_failure.lock().take() { + return Err(failure); + } + let next = self + .script + .lock() + .iter_mut() + .find(|(group, _)| *group == attempt.model_group) + .and_then(|(_, failures)| failures.pop_front()); + match next { + Some(failure) => Err(failure), + None => Ok(Output { + value: Payload(Arc::new(attempt.deployment.id.clone())), + total_tokens: 10, + blocked: None, + deferred: self.deferred, + tracks_usage: self.tracks_usage, + }), + } + } + + async fn select( + &self, + _stage: &SelectStage, + _ctx: &RoutingContext<'_>, + candidates: Vec>, + ) -> Result>, Error> { + match self.pick_failures.lock().pop_front() { + Some(failure) => Err(failure.into()), + None => Ok(candidates), + } + } + + async fn allow_fallback( + &self, + _request: &RouterRequest, + _target: &str, + ) -> Result { + match self.fallback_check_raises { + Some(class) => Err(AttemptError::host_raised(&[class])), + None => Ok(true), + } + } + + async fn refuses_fallbacks( + &self, + _level: u32, + targets: &[serde_json::Map], + ) -> bool { + self.refuses_dict_targets && !targets.is_empty() + } + + async fn sleep(&self, _duration: Duration) {} + + fn observe(&self, event: &RouterEvent) { + self.events.lock().push(event.clone()); + } + + fn cache(&self) -> Arc { + Arc::new(self.cache.clone()) + } +} + +#[fixture] +fn config() -> RouterConfig { + serde_json::from_value(json!({ + "model_list": [ + {"model_name": "primary", "litellm_params": {"model": "openai/a"}, "model_info": {"id": "a1"}}, + {"model_name": "primary", "litellm_params": {"model": "openai/b"}, "model_info": {"id": "a2"}}, + {"model_name": "backup", "litellm_params": {"model": "anthropic/c"}, "model_info": {"id": "b1"}}, + ] + })) + .unwrap() +} + +fn router(config: RouterConfig, dynamic: serde_json::Value, host: Arc) -> Router { + let dynamic: RouterDynamicConfig = serde_json::from_value(dynamic).unwrap(); + Router::builder(config, dynamic, host) + .randomness(Arc::new(SeededRandomness::new(7))) + .build() + .unwrap() +} + +fn request() -> RouterRequest { + RouterRequest { + model: "primary".to_owned(), + ..RouterRequest::default() + } +} + +/// `primary` at order 1, its second deployment at order 2 when `tiered`, and two other groups. +fn tiers(tiered: bool) -> RouterConfig { + let second = if tiered { "primary" } else { "unused" }; + serde_json::from_value(json!({ + "model_list": [ + {"model_name": "primary", "litellm_params": {"model": "openai/p1", "order": 1}, "model_info": {"id": "p1"}}, + {"model_name": second, "litellm_params": {"model": "openai/p2", "order": 2}, "model_info": {"id": "p2"}}, + {"model_name": "backup", "litellm_params": {"model": "anthropic/b1"}, "model_info": {"id": "b1"}}, + {"model_name": "context", "litellm_params": {"model": "anthropic/c1"}, "model_info": {"id": "c1"}}, + ] + })) + .unwrap() +} + +fn failure(kind: FailureKind, status: u16) -> AttemptError { + AttemptError::new(kind, Some(status)) +} + +fn context_window() -> AttemptError { + failure(FailureKind::ContextWindowExceeded, 400) +} + +fn server_error() -> AttemptError { + failure(FailureKind::InternalServer, 500) +} + +fn agentic_followup() -> AttemptError { + let mut failed = server_error(); + failed.agentic_followup = true; + failed +} + +#[rstest] +#[tokio::test] +async fn a_rate_limit_retries_in_the_same_group(config: RouterConfig) { + let host = Arc::new(ScriptedHost::failing( + "primary", + vec![AttemptError::new(FailureKind::RateLimit, Some(429))], + )); + let router = router(config, json!({"num_retries": 2}), host.clone()); + let response = router + .call(Op::from("acompletion"), request()) + .await + .unwrap(); + assert_eq!( + ( + response.model_group.as_str(), + response.attempted_retries, + response.attempted_fallbacks + ), + ("primary", 1, 0) + ); +} + +#[rstest] +#[tokio::test] +async fn a_bad_request_falls_back_without_retrying(config: RouterConfig) { + let host = Arc::new(ScriptedHost::failing( + "primary", + vec![AttemptError::new(FailureKind::BadRequest, Some(400))], + )); + let router = router( + config, + json!({"num_retries": 3, "fallbacks": [{"primary": ["backup"]}]}), + host.clone(), + ); + let response = router + .call(Op::from("acompletion"), request()) + .await + .unwrap(); + let retried = host + .events + .lock() + .iter() + .any(|event| matches!(event, RouterEvent::Retry { .. })); + assert_eq!( + ( + response.deployment_id.as_str(), + response.attempted_fallbacks, + retried + ), + ("b1", 1, false) + ); +} + +#[rstest] +#[tokio::test] +async fn a_rate_limit_cools_the_deployment_down(config: RouterConfig) { + let host = Arc::new(ScriptedHost::failing( + "primary", + vec![AttemptError::new(FailureKind::RateLimit, Some(429))], + )); + let router = router(config, json!({}), host.clone()); + let error = router + .call(Op::from("acompletion"), request()) + .await + .unwrap_err(); + let cooled = router + .health() + .active_cooldowns(&["a1".to_owned(), "a2".to_owned()]) + .await + .unwrap(); + assert_eq!( + (error.attempt().map(|failure| failure.kind), cooled.len()), + (Some(FailureKind::RateLimit), 1) + ); +} + +#[rstest] +#[case::now(false, (Some(1.0), (10, 1)))] +#[case::deferred(true, (None, (10, 1)))] +#[tokio::test] +async fn a_deferred_attempt_counts_its_request_and_tokens_but_no_success( + #[case] deferred: bool, + #[case] expected: (Option, (u64, u64)), +) { + let config: RouterConfig = serde_json::from_value(json!({ + "model_list": [ + {"model_name": "primary", "litellm_params": {"model": "openai/a", "rpm": 100}, "model_info": {"id": "a1"}}, + ] + })) + .unwrap(); + let host = Arc::new(ScriptedHost { + deferred, + ..ScriptedHost::default() + }); + let router = router(config, json!({}), host.clone()); + router + .call(Op::from("acompletion"), request()) + .await + .unwrap(); + let deployment = router.registry().deployment("a1").unwrap(); + let successes = host + .cache + .get("a1:successes") + .and_then(|count| count.as_f64()); + let usage = router.health().usage(&deployment).await.unwrap(); + assert_eq!((successes, usage), expected); +} + +/// `litellm_params`, `model_info` and a limit only the host knows of all count; a realtime +/// session (`Some(false)` on a deferred answer) counts nothing, whatever limits it has. +#[rstest] +#[case::in_params(json!({"rpm": 100}), json!({}), false, None, (10, 1))] +#[case::in_model_info(json!({}), json!({"rpm": 100}), false, None, (10, 1))] +#[case::in_the_catalog(json!({}), json!({}), false, Some(true), (10, 1))] +#[case::no_limit(json!({}), json!({}), false, None, (0, 0))] +#[case::a_stream_in_model_info(json!({}), json!({"rpm": 100}), true, None, (10, 1))] +#[case::a_realtime_session(json!({"rpm": 100}), json!({}), true, Some(false), (0, 0))] +#[tokio::test] +async fn usage_counts_for_a_deployment_that_tracks_it( + #[case] params: serde_json::Value, + #[case] info: serde_json::Value, + #[case] deferred: bool, + #[case] tracks_usage: Option, + #[case] expected: (u64, u64), +) { + let mut litellm_params = json!({"model": "openai/a"}); + litellm_params + .as_object_mut() + .unwrap() + .extend(params.as_object().unwrap().clone()); + let mut model_info = json!({"id": "a1"}); + model_info + .as_object_mut() + .unwrap() + .extend(info.as_object().unwrap().clone()); + let config: RouterConfig = serde_json::from_value(json!({ + "model_list": [ + {"model_name": "primary", "litellm_params": litellm_params, "model_info": model_info}, + ] + })) + .unwrap(); + let host = Arc::new(ScriptedHost { + deferred, + tracks_usage, + ..ScriptedHost::default() + }); + let router = router(config, json!({}), host.clone()); + router + .call(Op::from("acompletion"), request()) + .await + .unwrap(); + let deployment = router.registry().deployment("a1").unwrap(); + assert_eq!(router.health().usage(&deployment).await.unwrap(), expected); +} + +/// A stream end's report of a wildcard deployment counts under the model the route filled, the +/// key its request was counted under, not the configured `*`. +#[rstest] +#[tokio::test] +async fn a_reported_success_keeps_the_filled_wildcard_model() { + let config: RouterConfig = serde_json::from_value(json!({ + "model_list": [ + {"model_name": "openai/*", "litellm_params": {"model": "openai/*", "rpm": 100}, "model_info": {"id": "w1"}}, + ] + })) + .unwrap(); + let host = Arc::new(ScriptedHost::default()); + let router = router(config, json!({}), host.clone()); + let report = SuccessReport { + usage_model: Some("openai/gpt-4o".to_owned()), + ..SuccessReport::default() + }; + router.report_success("w1", 7.0, 0.0, report).await.unwrap(); + let minute = current_minute(std::time::SystemTime::now()); + let key = |model: &str| usage_key("w1", model, UsageMetric::Tpm, &minute); + assert_eq!( + ( + host.cache.get(&key("openai/gpt-4o")), + host.cache.get(&key("openai/*")), + ), + (Some(json!(7)), None), + ); +} + +/// `async_deployment_callback_on_failure`'s request is written only when it runs: a failure the +/// host says ran no failure callback (`recorded`) writes none, nor does a sync call's. +#[rstest] +#[case::the_first_logged_failure(false, false, 1, Some(1.0))] +#[case::a_later_failure_on_the_shared_logging_object(true, false, 0, None)] +#[case::a_sync_call(false, true, 0, Some(1.0))] +#[tokio::test] +async fn a_failed_attempt_counts_its_request_only_when_its_callback_ran( + #[case] recorded: bool, + #[case] uncounted_request: bool, + #[case] requests: u64, + #[case] fails: Option, +) { + let config: RouterConfig = serde_json::from_value(json!({ + "model_list": [ + {"model_name": "primary", "litellm_params": {"model": "openai/a"}, "model_info": {"id": "a1"}}, + ] + })) + .unwrap(); + let mut failure = AttemptError::new(FailureKind::InternalServer, Some(500)); + failure.recorded = recorded; + failure.uncounted_request = uncounted_request; + let host = Arc::new(ScriptedHost::failing("primary", vec![failure])); + let router = router(config, json!({"num_retries": 0}), host.clone()); + let _ = router.call(Op::from("acompletion"), request()).await; + let deployment = router.registry().deployment("a1").unwrap(); + let usage = router.health().usage(&deployment).await.unwrap(); + let counted = host.cache.get("a1:fails").and_then(|count| count.as_f64()); + assert_eq!((usage.1, counted), (requests, fails)); +} + +#[rstest] +#[case::counted(true, 2)] +#[case::uncounted(false, 1)] +#[tokio::test] +async fn a_reported_failure_counts_its_request_when_asked( + #[case] counts_request: bool, + #[case] requests: u64, +) { + let config: RouterConfig = serde_json::from_value(json!({ + "model_list": [ + {"model_name": "primary", "litellm_params": {"model": "openai/a", "rpm": 100}, "model_info": {"id": "a1"}}, + ] + })) + .unwrap(); + let host = Arc::new(ScriptedHost { + deferred: true, + ..ScriptedHost::default() + }); + let router = router(config, json!({}), host.clone()); + router + .call(Op::from("acompletion"), request()) + .await + .unwrap(); + let report = FailureReport { + counts_request, + ..FailureReport::default() + }; + router + .report_failure( + "a1", + &AttemptError::new(FailureKind::RateLimit, Some(429)), + report, + ) + .await + .unwrap(); + let deployment = router.registry().deployment("a1").unwrap(); + let fails = host.cache.get("a1:fails").and_then(|count| count.as_f64()); + let usage = router.health().usage(&deployment).await.unwrap(); + assert_eq!((fails, usage.1), (Some(1.0), requests)); +} + +#[rstest] +#[tokio::test] +async fn a_raising_retry_policy_lookup_ends_the_hop_and_falls_back(config: RouterConfig) { + let host = Arc::new(ScriptedHost::failing( + "primary", + vec![AttemptError::new(FailureKind::RateLimit, Some(429))], + )); + let router = router( + config, + json!({"num_retries": 3, "fallbacks": [{"primary": ["backup"]}]}), + host.clone(), + ); + let request = RouterRequest { + retry_policy_raises: vec!["primary".to_owned()], + ..request() + }; + let response = router.call(Op::from("acompletion"), request).await.unwrap(); + let events = host.events.lock(); + assert!( + !events + .iter() + .any(|e| matches!(e, RouterEvent::Retry { .. })) + ); + assert!(events.iter().any(|e| matches!( + e, + RouterEvent::RetryPolicyRaised { level: 0, model, .. } if model == "primary" + ))); + assert_eq!(response.model_group, "backup"); +} + +#[rstest] +#[tokio::test] +async fn a_num_retries_that_raises_ends_the_hop_and_falls_back(config: RouterConfig) { + let host = Arc::new(ScriptedHost::failing("primary", vec![server_error()])); + let router = router( + config, + json!({"num_retries": 3, "fallbacks": [{"primary": ["backup"]}]}), + host.clone(), + ); + let request = RouterRequest { + uncounted_retries: Some(UncountedRetries::Raises), + ..request() + }; + let response = router.call(Op::from("acompletion"), request).await.unwrap(); + let events = host.events.lock(); + assert!( + !events + .iter() + .any(|e| matches!(e, RouterEvent::Retry { .. })) + ); + assert!(events.iter().any(|e| matches!( + e, + RouterEvent::RetryCountRaised { level: 0, model, .. } if model == "primary" + ))); + assert_eq!(response.model_group, "backup"); +} + +#[rstest] +#[case::negative(UncountedRetries::Negative)] +#[case::not_positive(UncountedRetries::NotPositive)] +#[tokio::test] +async fn a_num_retries_below_one_leaves_the_failure_standing( + config: RouterConfig, + #[case] uncounted: UncountedRetries, +) { + let host = Arc::new(ScriptedHost::failing("primary", vec![server_error()])); + let router = router(config, json!({"num_retries": 3}), host.clone()); + let request = RouterRequest { + uncounted_retries: Some(uncounted), + ..request() + }; + let refused = router.call(Op::from("acompletion"), request).await; + assert!(matches!(refused, Err(Error::Attempt(failed)) if failed.status == Some(500))); + assert_eq!(host.retries(), 0); + assert!( + !host + .events + .lock() + .iter() + .any(|event| matches!(event, RouterEvent::RetryCountRaised { .. })) + ); +} + +#[rstest] +#[tokio::test] +async fn a_retry_policy_that_grants_retries_outranks_a_num_retries_that_raises( + config: RouterConfig, +) { + let host = Arc::new(ScriptedHost::failing("primary", vec![server_error()])); + let router = router( + RouterConfig { + retry_policy: serde_json::from_value(json!({"DefaultRetries": 1})).unwrap(), + ..config + }, + json!({}), + host.clone(), + ); + let request = RouterRequest { + uncounted_retries: Some(UncountedRetries::Raises), + ..request() + }; + let response = router.call(Op::from("acompletion"), request).await.unwrap(); + assert_eq!(response.model_group, "primary"); + assert_eq!(host.retries(), 1); +} + +#[rstest] +#[tokio::test] +async fn a_name_nothing_serves_routes_to_the_default_deployment(config: RouterConfig) { + let host = Arc::new(ScriptedHost::default()); + let router = router( + config, + json!({"default_deployment": { + "model_name": "dflt", "litellm_params": {"model": "openai/x"}, "model_info": {"id": "dd"} + }}), + host.clone(), + ); + let request = RouterRequest { + model: "other".to_owned(), + ..RouterRequest::default() + }; + let response = router + .call(Op::from("acompletion"), request.clone()) + .await + .unwrap(); + assert_eq!( + ( + response.deployment_id.as_str(), + response.model_group.as_str() + ), + ("dd", "other") + ); + let picked = router + .select(&Op::from("acompletion"), &request) + .await + .unwrap(); + assert_eq!(picked.params.model, "other"); +} + +#[rstest] +#[tokio::test] +async fn a_stale_member_map_naming_a_dropped_group_raises_on_every_pick(mut config: RouterConfig) { + config.routing_group_members = Some([("primary".to_owned(), "gone".to_owned())].into()); + let router = router(config, json!({}), Arc::new(ScriptedHost::default())); + let refused = router.select(&Op::from("acompletion"), &request()).await; + assert!(matches!(refused, Err(Error::MissingRoutingGroup { group }) if group == "gone")); + let other = RouterRequest { + model: "backup".to_owned(), + ..RouterRequest::default() + }; + assert!( + router + .select(&Op::from("acompletion"), &other) + .await + .is_ok() + ); +} + +#[rstest] +#[tokio::test] +async fn written_indexes_replace_the_ones_the_deployments_build(config: RouterConfig) { + let router = router(config, json!({}), Arc::new(ScriptedHost::default())); + let index: litellm_router_core::registry::IndexOverrides = serde_json::from_value(json!({ + "model_names": ["backup"], + "model_name": {"primary": [1], "backup": [2]}, + })) + .unwrap(); + router + .apply(litellm_router_core::ChangeSet { + index: Some(index), + ..litellm_router_core::ChangeSet::default() + }) + .unwrap(); + for _ in 0..4 { + let picked = router + .select(&Op::from("acompletion"), &request()) + .await + .unwrap(); + assert_eq!(picked.id, "a2"); + } + router + .apply(litellm_router_core::ChangeSet { + index: Some(litellm_router_core::registry::IndexOverrides::default()), + ..litellm_router_core::ChangeSet::default() + }) + .unwrap(); + assert_eq!(router.registry().group("primary").unwrap().len(), 2); +} + +#[rstest] +#[tokio::test] +async fn a_retry_policy_retries_a_pick_that_found_no_deployment(config: RouterConfig) { + let host = Arc::new(ScriptedHost::default()); + let mut config = config; + config.retry_policy = serde_json::from_value(json!({"DefaultRetries": 2})).unwrap(); + let router = router(config, json!({"num_retries": 0}), host.clone()); + let request = RouterRequest { + model: "nothing".to_owned(), + ..RouterRequest::default() + }; + let refused = router.call(Op::from("acompletion"), request).await; + assert!(matches!(refused, Err(Error::UnknownModel { .. }))); + let events = host.events.lock(); + let retries = events + .iter() + .filter(|event| matches!(event, RouterEvent::Retry { .. })) + .count(); + assert_eq!(retries, 2); + assert!(events.iter().any(|event| matches!( + event, + RouterEvent::RetriesExhausted { + picked: true, + level: 0, + .. + } + ))); +} + +#[rstest] +#[tokio::test] +async fn a_default_deployment_python_cannot_copy_raises_for_a_name_nothing_serves( + config: RouterConfig, +) { + let router = router( + config, + json!({"default_deployment_raises": true}), + Arc::new(ScriptedHost::default()), + ); + let request = RouterRequest { + model: "other".to_owned(), + ..RouterRequest::default() + }; + let refused = router.select(&Op::from("acompletion"), &request).await; + assert!(matches!(refused, Err(Error::DefaultDeploymentRaises { model }) if model == "other")); + assert!( + router + .select(&Op::from("acompletion"), &self::request()) + .await + .is_ok() + ); +} + +#[rstest] +#[tokio::test] +async fn a_written_index_position_past_the_list_raises_for_that_name(config: RouterConfig) { + let router = router(config, json!({}), Arc::new(ScriptedHost::default())); + let index: litellm_router_core::registry::IndexOverrides = + serde_json::from_value(json!({"model_name": {"primary": [0, 9], "backup": [2]}})).unwrap(); + router + .apply(litellm_router_core::ChangeSet { + index: Some(index), + ..litellm_router_core::ChangeSet::default() + }) + .unwrap(); + let refused = router.select(&Op::from("acompletion"), &request()).await; + assert!(matches!(refused, Err(Error::IndexOutOfRange { model }) if model == "primary")); + let backup = RouterRequest { + model: "backup".to_owned(), + ..RouterRequest::default() + }; + assert!( + router + .select(&Op::from("acompletion"), &backup) + .await + .is_ok() + ); +} + +#[rstest] +#[case::no_list_takes_the_order_tier(json!(null), true, false, "p2")] +#[case::unmatched_list_takes_the_order_tier(json!([{"unrelated": ["context"]}]), true, false, "p2")] +#[case::unmatched_list_falls_through_to_regular(json!([{"unrelated": ["context"]}]), false, false, "b1")] +#[case::unmatched_list_leaves_the_client_side_list(json!([{"unrelated": ["context"]}]), false, true, "b1")] +#[case::matched_list_beats_the_order_tier(json!([{"primary": ["context"]}]), true, false, "c1")] +#[case::matched_star_beats_the_order_tier(json!([{"*": ["context"]}]), true, false, "c1")] +#[case::matched_list_beats_the_client_side_list(json!([{"primary": ["context"]}]), false, true, "c1")] +#[case::matched_list_beats_both(json!([{"primary": ["context"]}]), true, true, "c1")] +#[tokio::test] +async fn a_context_window_error_takes_its_own_chain_only_when_the_group_matches( + #[case] context_window_fallbacks: serde_json::Value, + #[case] tiered: bool, + #[case] client_side: bool, + #[case] expected: &str, +) { + let host = Arc::new(ScriptedHost::failing("primary", vec![context_window()])); + let router = router( + tiers(tiered), + json!({ + "fallbacks": [{"primary": ["backup"]}], + "context_window_fallbacks": context_window_fallbacks, + }), + host.clone(), + ); + let call = RouterRequest { + fallbacks: client_side.then(|| Some(serde_json::from_value(json!(["backup"])).unwrap())), + ..request() + }; + let response = router.call(Op::from("acompletion"), call).await.unwrap(); + assert_eq!(response.deployment_id, expected); +} + +#[rstest] +#[case::context_window_keeps_its_own_error( + context_window(), + FailureKind::ContextWindowExceeded, + false +)] +#[case::any_other_raises_the_last_hop(server_error(), FailureKind::ServiceUnavailable, true)] +#[tokio::test] +async fn a_failed_order_chain_raises_by_the_error_that_started_it( + #[case] first: AttemptError, + #[case] raised: FailureKind, + #[case] last_hop: bool, +) { + let host = Arc::new(ScriptedHost::scripted(vec![ + ("primary", vec![first, failure(FailureKind::RateLimit, 429)]), + ( + "backup", + vec![failure(FailureKind::ServiceUnavailable, 503)], + ), + ])); + let router = router( + tiers(true), + json!({ + "fallbacks": [{"primary": ["backup"]}], + "context_window_fallbacks": [{"unrelated": ["context"]}], + }), + host.clone(), + ); + let refused = router.call(Op::from("acompletion"), request()).await; + assert!(matches!(refused, Err(Error::Attempt(failed)) if failed.kind == raised)); + assert_eq!(host.selected(), ["p1", "p2", "b1"]); + assert!( + host.events + .lock() + .contains(&RouterEvent::OrderChainFailed { level: 0, last_hop }) + ); +} + +#[rstest] +#[case::context_window_skips_it(context_window(), "backup")] +#[case::a_rate_limit_takes_it(failure(FailureKind::RateLimit, 429), "primary")] +#[tokio::test] +async fn weighted_failover_serves_every_failure_but_a_context_window_error( + #[case] first: AttemptError, + #[case] served_by: &str, +) { + let host = Arc::new(ScriptedHost::failing("primary", vec![first])); + let config: RouterConfig = serde_json::from_value(json!({ + "model_list": [ + {"model_name": "primary", "litellm_params": {"model": "openai/a"}, "model_info": {"id": "a1"}}, + {"model_name": "primary", "litellm_params": {"model": "openai/b"}, "model_info": {"id": "a2"}}, + {"model_name": "backup", "litellm_params": {"model": "anthropic/c"}, "model_info": {"id": "b1"}}, + ], + "enable_weighted_failover": true, + })) + .unwrap(); + let router = router( + config, + json!({ + "fallbacks": [{"primary": ["backup"]}], + "context_window_fallbacks": [{"unrelated": ["backup"]}], + }), + host.clone(), + ); + let response = router + .call(Op::from("acompletion"), request()) + .await + .unwrap(); + assert_eq!(response.model_group, served_by); +} + +#[rstest] +#[tokio::test] +async fn a_failed_agentic_followup_is_not_retried_but_still_falls_back(config: RouterConfig) { + let host = Arc::new(ScriptedHost::failing("primary", vec![agentic_followup()])); + let router = router( + config, + json!({"num_retries": 3, "fallbacks": [{"primary": ["backup"]}]}), + host.clone(), + ); + let response = router + .call(Op::from("acompletion"), request()) + .await + .unwrap(); + assert_eq!(response.deployment_id, "b1"); + assert_eq!(host.retries(), 0); + assert_eq!(host.selected().len(), 2); +} + +#[rstest] +#[tokio::test] +async fn a_failed_agentic_followup_skips_the_retry_policy_lookup(config: RouterConfig) { + let host = Arc::new(ScriptedHost::failing("primary", vec![agentic_followup()])); + let router = router(config, json!({"num_retries": 3}), host.clone()); + let call = RouterRequest { + retry_policy_raises: vec!["primary".to_owned()], + ..request() + }; + let refused = router.call(Op::from("acompletion"), call).await; + assert!(matches!(refused, Err(Error::Attempt(failed)) if failed.agentic_followup)); + assert!( + !host + .events + .lock() + .iter() + .any(|event| matches!(event, RouterEvent::RetryPolicyRaised { .. })) + ); +} + +#[rstest] +#[tokio::test] +async fn a_retry_that_fails_its_agentic_followup_ends_the_retries(config: RouterConfig) { + let host = Arc::new(ScriptedHost::failing( + "primary", + vec![server_error(), agentic_followup()], + )); + let router = router(config, json!({"num_retries": 3}), host.clone()); + let refused = router.call(Op::from("acompletion"), request()).await; + assert!(matches!(refused, Err(Error::Attempt(failed)) if failed.agentic_followup)); + assert_eq!((host.retries(), host.selected().len()), (1, 2)); +} + +fn stream_origin() -> StreamOrigin { + StreamOrigin { + group: "primary".to_owned(), + ..StreamOrigin::default() + } +} + +fn stream_request() -> RouterRequest { + RouterRequest { + num_retries: Some(3), + ..request() + } +} + +#[rstest] +#[tokio::test] +async fn a_stream_that_failed_its_agentic_followup_is_not_retried(config: RouterConfig) { + let host = Arc::new(ScriptedHost::default()); + let router = router(config, json!({}), host.clone()); + let refused = router + .retry_stream( + Op::from("acompletion"), + stream_request(), + stream_origin(), + agentic_followup(), + ) + .await; + assert!(matches!(refused, Err(Error::Attempt(failed)) if failed.agentic_followup)); + assert_eq!((host.retries(), host.selected().len()), (0, 0)); +} + +#[rstest] +#[tokio::test] +async fn a_stream_retry_that_fails_its_agentic_followup_ends_the_retries(config: RouterConfig) { + let host = Arc::new(ScriptedHost::failing("primary", vec![agentic_followup()])); + let router = router(config, json!({}), host.clone()); + let refused = router + .retry_stream( + Op::from("acompletion"), + stream_request(), + stream_origin(), + server_error(), + ) + .await; + assert!(matches!(refused, Err(Error::Attempt(failed)) if failed.agentic_followup)); + assert_eq!((host.retries(), host.selected().len()), (1, 1)); +} + +#[rstest] +#[tokio::test] +async fn a_stream_retry_is_reported_before_its_attempt(config: RouterConfig) { + let host = Arc::new(ScriptedHost::default()); + let router = router(config, json!({}), host.clone()); + let response = router + .retry_stream( + Op::from("acompletion"), + stream_request(), + stream_origin(), + server_error(), + ) + .await + .unwrap(); + assert_eq!((response.attempted_retries, host.retries()), (1, 1)); +} + +fn fallback_list(value: serde_json::Value) -> Fallbacks { + serde_json::from_value(value).unwrap() +} + +#[rstest] +#[case::names(json!(["backup"]))] +#[case::keyed(json!([{"primary": ["backup"]}]))] +#[case::star(json!([{"*": ["backup"]}]))] +fn a_request_fallback_list_can_take_over_a_stream( + config: RouterConfig, + #[case] list: serde_json::Value, +) { + let router = router(config, json!({}), Arc::new(ScriptedHost::default())); + let request = RouterRequest { + fallbacks: Some(Some(fallback_list(list))), + ..request() + }; + let policy = router.stream_policy(&request, &stream_origin(), None); + assert!(policy.can_fall_back); + assert!(policy.refusal_fallback); +} + +#[rstest] +fn a_request_content_policy_list_can_take_over_a_stream(config: RouterConfig) { + let router = router(config, json!({}), Arc::new(ScriptedHost::default())); + let request = RouterRequest { + content_policy_fallbacks: Some(Some(fallback_list(json!([{"primary": ["backup"]}])))), + ..request() + }; + let policy = router.stream_policy(&request, &stream_origin(), None); + assert!(policy.can_fall_back); + assert!(policy.refusal_fallback); +} + +#[rstest] +#[case::cleared(None)] +#[case::another_group(Some(json!([{"other": ["backup"]}])))] +fn a_request_fallback_list_replaces_the_router_list_for_a_stream( + config: RouterConfig, + #[case] list: Option, +) { + let router = router( + config, + json!({"fallbacks": [{"primary": ["backup"]}], "default_fallbacks": ["backup"]}), + Arc::new(ScriptedHost::default()), + ); + let request = RouterRequest { + fallbacks: Some(list.map(fallback_list)), + ..request() + }; + let policy = router.stream_policy(&request, &stream_origin(), None); + assert!(!policy.can_fall_back); +} + +#[rstest] +#[case::readable(vec![], 5, true)] +#[case::malformed(vec!["primary".to_owned()], 1, false)] +fn a_group_policy_that_does_not_read_governs_no_stream_retry( + mut config: RouterConfig, + #[case] retry_policy_raises: Vec, + #[case] budget: u32, + #[case] applies: bool, +) { + config.retry_policy = serde_json::from_value(json!({"InternalServerErrorRetries": 5})).unwrap(); + let router = router( + config, + json!({"num_retries": 1}), + Arc::new(ScriptedHost::default()), + ); + let request = RouterRequest { + retry_policy_raises, + ..request() + }; + let policy = router.stream_policy( + &request, + &stream_origin(), + Some(FailureKind::InternalServer), + ); + assert_eq!(policy.retry_budget, Some((budget, applies))); +} + +#[rstest] +#[tokio::test] +async fn a_level_reports_its_fallback_handling_before_any_hop(config: RouterConfig) { + let host = Arc::new(ScriptedHost::failing("primary", vec![context_window()])); + let router = router( + config, + json!({"num_retries": 0, "fallbacks": [{"primary": ["backup"]}]}), + host.clone(), + ); + router + .call(Op::from("acompletion"), request()) + .await + .unwrap(); + let events = host.events.lock(); + let falling_back = RouterEvent::FallingBack { + level: 0, + model: "primary".to_owned(), + }; + let reported = events.iter().position(|e| *e == falling_back).unwrap(); + let hopped = events + .iter() + .position(|e| matches!(e, RouterEvent::Fallback { .. })) + .unwrap(); + assert!(reported < hopped); +} + +#[rstest] +#[tokio::test] +async fn the_order_path_reports_no_fallback_handling() { + let host = Arc::new(ScriptedHost::failing("primary", vec![server_error()])); + let router = router(tiers(true), json!({"num_retries": 0}), host.clone()); + router + .call(Op::from("acompletion"), request()) + .await + .unwrap(); + assert!( + !host + .events + .lock() + .iter() + .any(|e| matches!(e, RouterEvent::FallingBack { .. })) + ); +} + +fn fallbacks_failed(host: &ScriptedHost) -> Vec<(u32, bool, bool)> { + host.events + .lock() + .iter() + .filter_map(|event| match event { + RouterEvent::FallbacksFailed { + level, + chain, + last_hop, + } => Some((*level, *chain, *last_hop)), + _ => None, + }) + .collect() +} + +#[rstest] +#[tokio::test] +async fn a_chain_that_fails_reports_every_level_innermost_first(config: RouterConfig) { + let host = Arc::new(ScriptedHost::scripted(vec![ + ("primary", vec![server_error()]), + ("backup", vec![server_error()]), + ])); + let router = router( + config, + json!({"num_retries": 0, "fallbacks": [{"primary": ["backup"]}]}), + host.clone(), + ); + router + .call(Op::from("acompletion"), request()) + .await + .unwrap_err(); + assert_eq!( + fallbacks_failed(&host), + vec![(1, true, false), (0, true, true)] + ); +} + +#[rstest] +#[tokio::test] +async fn a_chain_past_max_fallbacks_reports_the_levels_own_failure(config: RouterConfig) { + let host = Arc::new(ScriptedHost::failing("primary", vec![server_error()])); + let router = router( + config, + json!({"num_retries": 0, "max_fallbacks": 0, "fallbacks": [{"primary": ["backup"]}]}), + host.clone(), + ); + router + .call(Op::from("acompletion"), request()) + .await + .unwrap_err(); + assert_eq!(fallbacks_failed(&host), vec![(0, true, false)]); +} + +#[rstest] +#[tokio::test] +async fn a_served_fallback_and_a_disabled_walk_report_no_failure(config: RouterConfig) { + let host = Arc::new(ScriptedHost::failing("primary", vec![server_error()])); + let dynamic = json!({"num_retries": 0, "fallbacks": [{"primary": ["backup"]}]}); + let served = router(config.clone(), dynamic.clone(), host.clone()); + served + .call(Op::from("acompletion"), request()) + .await + .unwrap(); + assert_eq!(fallbacks_failed(&host), vec![]); + + let host = Arc::new(ScriptedHost::failing("primary", vec![server_error()])); + let disabled = router(config, dynamic, host.clone()); + let request = RouterRequest { + disable_fallbacks: true, + ..request() + }; + disabled + .call(Op::from("acompletion"), request) + .await + .unwrap_err(); + assert_eq!(fallbacks_failed(&host), vec![]); +} + +#[tokio::test] +async fn the_order_path_reports_no_failed_fallbacks() { + let host = Arc::new(ScriptedHost::failing("primary", vec![server_error()])); + let router = router(tiers(true), json!({"num_retries": 0}), host.clone()); + router + .call(Op::from("acompletion"), request()) + .await + .unwrap(); + assert_eq!(fallbacks_failed(&host), vec![]); +} + +fn refusing_host(group: &str) -> Arc { + Arc::new(ScriptedHost { + refuses_dict_targets: true, + ..ScriptedHost::failing(group, vec![server_error()]) + }) +} + +#[rstest] +#[case::dict_target_first(json!([{"model": "backup", "api_key": "x"}, "backup"]))] +#[case::dict_target_last(json!(["backup", {"model": "backup", "api_key": "x"}]))] +#[tokio::test] +async fn a_refused_chain_tries_none_of_its_targets( + config: RouterConfig, + #[case] chain: serde_json::Value, +) { + let host = refusing_host("primary"); + let router = router( + config, + json!({"num_retries": 0, "fallbacks": [{"primary": chain}]}), + host.clone(), + ); + let raised = router.call(Op::from("acompletion"), request()).await; + assert!(matches!(raised, Err(Error::Attempt(failed)) if failed.is_a("InternalServerError"))); + assert_eq!(host.selected().len(), 1); +} + +#[rstest] +#[tokio::test] +async fn a_chain_of_names_is_never_refused(config: RouterConfig) { + let router = router( + config, + json!({"num_retries": 0, "fallbacks": [{"primary": ["backup"]}]}), + refusing_host("primary"), + ); + let response = router + .call(Op::from("acompletion"), request()) + .await + .unwrap(); + assert_eq!(response.deployment_id, "b1"); +} + +#[rstest] +#[tokio::test] +async fn a_refused_chain_on_the_order_path_raises_the_refusal() { + let host = refusing_host("primary"); + let router = router( + tiers(true), + json!({"num_retries": 0, "fallbacks": [{"primary": [{"model": "backup", "api_key": "x"}]}]}), + host.clone(), + ); + let raised = router.call(Op::from("acompletion"), request()).await; + assert!(matches!(raised, Err(Error::Attempt(failed)) if failed.is_a("ValueError"))); + assert_eq!(host.selected(), ["p1"]); +} + +#[rstest] +#[tokio::test] +async fn a_refused_chain_fails_only_its_own_level(mut config: RouterConfig) { + let host = Arc::new(ScriptedHost { + refuses_dict_targets: true, + ..ScriptedHost::scripted(vec![ + ("primary", vec![server_error()]), + ("backup", vec![server_error()]), + ]) + }); + let dynamic = json!({"num_retries": 0, "fallbacks": [ + {"primary": ["backup", "third"]}, + {"backup": [{"model": "third", "api_key": "x"}]}, + ]}); + config.model_list.push( + serde_json::from_value(json!({ + "model_name": "third", "litellm_params": {"model": "openai/t"}, "model_info": {"id": "t1"}, + })) + .unwrap(), + ); + let response = router(config, dynamic, host.clone()) + .call(Op::from("acompletion"), request()) + .await + .unwrap(); + assert_eq!(response.deployment_id, "t1"); +} + +fn shared_public_name() -> RouterConfig { + serde_json::from_value(json!({ + "model_list": [ + {"model_name": "backup", "litellm_params": {"model": "openai/g"}, "model_info": {"id": "b1"}}, + {"model_name": "t1-pub", "litellm_params": {"model": "openai/m"}, + "model_info": {"id": "p1", "team_id": "T1", "team_public_model_name": "pub"}}, + {"model_name": "t2-pub", "litellm_params": {"model": "openai/m"}, + "model_info": {"id": "p2", "team_id": "T2", "team_public_model_name": "pub"}}, + ] + })) + .unwrap() +} + +fn proxy_admin_of_pub() -> RouterRequest { + RouterRequest { + model: "pub".to_owned(), + proxy_admin: true, + admin_model: Some("pub".to_owned()), + ..RouterRequest::default() + } +} + +#[rstest] +#[tokio::test] +async fn a_name_several_teams_share_falls_back() { + let router = router( + shared_public_name(), + json!({"fallbacks": [{"pub": ["backup"]}]}), + Arc::new(ScriptedHost::default()), + ); + let response = router + .call(Op::from("acompletion"), proxy_admin_of_pub()) + .await + .unwrap(); + assert_eq!( + ( + response.deployment_id.as_str(), + response.attempted_fallbacks + ), + ("b1", 1) + ); +} + +#[rstest] +#[tokio::test] +async fn a_name_several_teams_share_raises_its_refusal_when_the_fallback_fails() { + let router = router( + shared_public_name(), + json!({"fallbacks": [{"pub": ["backup"]}]}), + Arc::new(ScriptedHost::failing("backup", vec![server_error()])), + ); + let raised = router + .call(Op::from("acompletion"), proxy_admin_of_pub()) + .await; + assert!(matches!(raised, Err(Error::MultipleTeams { model }) if model == "pub")); +} + +/// A group whose second order level belongs to another team. +fn team_tiers() -> RouterConfig { + serde_json::from_value(json!({ + "model_list": [ + {"model_name": "primary", "litellm_params": {"model": "openai/p1", "order": 1}, + "model_info": {"id": "p1", "team_id": "t1"}}, + {"model_name": "primary", "litellm_params": {"model": "openai/p2", "order": 2}, + "model_info": {"id": "p2", "team_id": "t2"}}, + ] + })) + .unwrap() +} + +#[rstest] +#[case::listed_for_the_team(Some("t1"), false)] +#[case::listed_for_no_team(None, true)] +#[tokio::test] +async fn a_failed_level_lists_its_order_levels_for_the_metadata_team( + #[case] listing_team_id: Option<&str>, + #[case] walks_the_order_chain: bool, +) { + let host = Arc::new(ScriptedHost::failing("primary", vec![server_error()])); + let router = router(team_tiers(), json!({"num_retries": 0}), host.clone()); + let request = RouterRequest { + team_id: Some("t1".to_owned()), + listing_team_id: listing_team_id.map(str::to_owned), + ..request() + }; + let raised = router.call(Op::from("acompletion"), request).await; + assert_eq!( + matches!(raised, Err(Error::NoHealthyDeployments { .. })), + walks_the_order_chain + ); + assert_eq!(host.selected(), ["p1"]); +} + +fn tagged_request() -> RouterRequest { + RouterRequest { + filter_tags: true, + ..request() + } +} + +#[rstest] +#[tokio::test] +async fn a_failed_pick_retries_like_a_failed_attempt(config: RouterConfig) { + let host = Arc::new(ScriptedHost { + pick_failures: Mutex::new(VecDeque::from([failure(FailureKind::RateLimit, 429)])), + ..ScriptedHost::default() + }); + let response = router(config, json!({"num_retries": 1}), host.clone()) + .call(Op::from("acompletion"), tagged_request()) + .await + .unwrap(); + assert_eq!( + (response.model_group.as_str(), response.attempted_retries), + ("primary", 1) + ); +} + +#[rstest] +#[tokio::test] +async fn a_failed_pick_falls_back_with_no_deployment_to_blame(config: RouterConfig) { + let host = Arc::new(ScriptedHost { + pick_failures: Mutex::new(VecDeque::from([failure(FailureKind::RateLimit, 429)])), + ..ScriptedHost::default() + }); + let router = router( + config, + json!({"num_retries": 0, "fallbacks": [{"primary": ["backup"]}]}), + host.clone(), + ); + let response = router + .call(Op::from("acompletion"), tagged_request()) + .await + .unwrap(); + assert_eq!( + ( + response.deployment_id.as_str(), + response.attempted_fallbacks + ), + ("b1", 1) + ); + assert_eq!(host.selected(), ["b1"]); +} + +#[rstest] +#[tokio::test] +async fn a_fallback_check_that_raises_ends_the_chain_with_the_levels_own_error( + config: RouterConfig, +) { + let host = Arc::new(ScriptedHost { + fallback_check_raises: Some("TypeError"), + ..ScriptedHost::failing("primary", vec![failure(FailureKind::RateLimit, 429)]) + }); + let router = router( + config, + json!({"num_retries": 0, "fallbacks": [{"primary": ["backup"]}]}), + host.clone(), + ); + let raised = router.call(Op::from("acompletion"), request()).await; + assert!(matches!(raised, Err(Error::Attempt(failed)) if failed.kind == FailureKind::RateLimit)); + assert_eq!(host.selected().len(), 1); +} + +#[rstest] +#[tokio::test] +async fn a_fallback_check_that_raises_on_the_order_path_raises_its_own_error() { + let host = Arc::new(ScriptedHost { + fallback_check_raises: Some("TypeError"), + ..ScriptedHost::scripted(vec![("primary", vec![server_error(), server_error()])]) + }); + let router = router( + tiers(true), + json!({"num_retries": 0, "fallbacks": [{"primary": ["backup"]}]}), + host.clone(), + ); + let raised = router.call(Op::from("acompletion"), request()).await; + assert!(matches!(raised, Err(Error::Attempt(failed)) if failed.is_a("TypeError"))); +} + +#[rstest] +#[case::an_int(json!(5))] +#[case::none(json!(null))] +#[tokio::test] +async fn a_dict_target_with_a_non_string_model_is_a_hop_that_attempts_nothing( + config: RouterConfig, + #[case] model: serde_json::Value, +) { + let host = Arc::new(ScriptedHost::failing("primary", vec![server_error()])); + let router = router( + config, + json!({"num_retries": 0, "fallbacks": [{"primary": [{"model": model}, "backup"]}]}), + host.clone(), + ); + let response = router + .call(Op::from("acompletion"), request()) + .await + .unwrap(); + assert_eq!( + ( + response.deployment_id.as_str(), + response.attempted_fallbacks + ), + ("b1", 2) + ); + assert_eq!(host.selected().len(), 2); +} + +#[rstest] +#[tokio::test] +async fn a_dict_target_without_a_model_serves_the_group_that_failed_again(config: RouterConfig) { + let host = Arc::new(ScriptedHost::failing("primary", vec![server_error()])); + let router = router( + config, + json!({"num_retries": 0, "fallbacks": [{"primary": [{"temperature": 0.5}]}]}), + host.clone(), + ); + let response = router + .call(Op::from("acompletion"), request()) + .await + .unwrap(); + assert_eq!( + (response.model_group.as_str(), response.attempted_fallbacks), + ("primary", 1) + ); +} + +fn weighted_failover() -> RouterConfig { + serde_json::from_value(json!({ + "model_list": [ + {"model_name": "primary", "litellm_params": {"model": "openai/a"}, "model_info": {"id": "a1"}}, + {"model_name": "primary", "litellm_params": {"model": "openai/b"}, "model_info": {"id": "a2"}}, + ], + "enable_weighted_failover": true, + })) + .unwrap() +} + +#[rstest] +#[case::a_bug_propagates(AttemptError::host_raised(&["TypeError", "Exception"]), "TypeError")] +#[case::a_provider_error_defers(failure(FailureKind::RateLimit, 429), "InternalServerError")] +#[case::no_deployment_defers( + AttemptError::host_raised(&["RouterRateLimitError", "ValueError"]), + "InternalServerError" +)] +#[tokio::test] +async fn weighted_failover_swallows_only_api_and_rate_limit_errors( + #[case] hop_failure: AttemptError, + #[case] raised_class: &str, +) { + let host = Arc::new(ScriptedHost::failing( + "primary", + vec![server_error(), hop_failure], + )); + let raised = router(weighted_failover(), json!({"num_retries": 0}), host.clone()) + .call(Op::from("acompletion"), request()) + .await; + assert!(matches!(raised, Err(Error::Attempt(failed)) if failed.is_a(raised_class))); + assert_eq!(host.selected().len(), 2); +} + +fn wildcard_tiers() -> RouterConfig { + serde_json::from_value(json!({ + "model_list": [ + {"model_name": "openai/*", "litellm_params": {"model": "openai/*", "order": 1}, "model_info": {"id": "w1"}}, + {"model_name": "openai/*", "litellm_params": {"model": "openai/*", "order": 2}, "model_info": {"id": "w2"}}, + ] + })) + .unwrap() +} + +fn alias_tiers() -> RouterConfig { + serde_json::from_value(json!({ + "model_list": [ + {"model_name": "gg", "litellm_params": {"model": "openai/a", "order": 1}, "model_info": {"id": "w1"}}, + {"model_name": "gg", "litellm_params": {"model": "openai/b", "order": 2}, "model_info": {"id": "w2"}}, + ] + })) + .unwrap() +} + +fn team_public_tiers() -> RouterConfig { + serde_json::from_value(json!({ + "model_list": [ + {"model_name": "t1a", "litellm_params": {"model": "openai/a", "order": 1}, + "model_info": {"id": "w1", "team_id": "T1", "team_public_model_name": "tp"}}, + {"model_name": "t1b", "litellm_params": {"model": "openai/b", "order": 2}, + "model_info": {"id": "w2", "team_id": "T1", "team_public_model_name": "tp"}}, + ] + })) + .unwrap() +} + +#[rstest] +#[case::a_wildcard(wildcard_tiers(), json!({"num_retries": 0}), "openai/gpt-4o", None)] +#[case::an_alias( + alias_tiers(), + json!({"num_retries": 0, "model_group_alias": {"al": "gg"}}), + "al", + None +)] +#[case::a_team_public_name(team_public_tiers(), json!({"num_retries": 0}), "tp", Some("T1"))] +#[tokio::test] +async fn order_levels_walk_through_any_name_the_model_list_lists( + #[case] config: RouterConfig, + #[case] dynamic: serde_json::Value, + #[case] model: &str, + #[case] team: Option<&str>, +) { + let host = Arc::new(ScriptedHost::default()); + *host.first_failure.lock() = Some(failure(FailureKind::RateLimit, 429)); + let request = RouterRequest { + model: model.to_owned(), + team_id: team.map(str::to_owned), + listing_team_id: team.map(str::to_owned), + ..RouterRequest::default() + }; + let response = router(config, dynamic, host.clone()) + .call(Op::from("acompletion"), request) + .await + .unwrap(); + assert_eq!(response.deployment_id, "w2"); +} + +#[rstest] +#[tokio::test] +async fn a_team_wildcard_public_name_routes_for_its_team() { + let config: RouterConfig = serde_json::from_value(json!({ + "model_list": [ + {"model_name": "t1-w", "litellm_params": {"model": "openai/*"}, + "model_info": {"id": "tw1", "team_id": "T1", "team_public_model_name": "tw-*"}}, + {"model_name": "g", "litellm_params": {"model": "openai/g"}, "model_info": {"id": "gp"}}, + ] + })) + .unwrap(); + let request = RouterRequest { + model: "tw-gpt-4o".to_owned(), + team_id: Some("T1".to_owned()), + ..RouterRequest::default() + }; + let response = router(config, json!({}), Arc::new(ScriptedHost::default())) + .call(Op::from("acompletion"), request) + .await + .unwrap(); + assert_eq!(response.deployment_id, "tw1"); +} + +#[rstest] +#[tokio::test] +async fn a_paused_deployment_named_by_id_falls_back() { + let config: RouterConfig = serde_json::from_value(json!({ + "model_list": [ + {"model_name": "primary", "litellm_params": {"model": "openai/a"}, + "model_info": {"id": "a1", "blocked": true}}, + {"model_name": "backup", "litellm_params": {"model": "openai/b"}, "model_info": {"id": "b1"}}, + ] + })) + .unwrap(); + let request = RouterRequest { + model: "a1".to_owned(), + ..RouterRequest::default() + }; + let response = router( + config, + json!({"fallbacks": [{"a1": ["backup"]}]}), + Arc::new(ScriptedHost::default()), + ) + .call(Op::from("acompletion"), request) + .await + .unwrap(); + assert_eq!(response.deployment_id, "b1"); +} + +/// `primary` with `a1`'s `model_info` as given, and an unremarkable `a2`. +fn with_first_info(info: serde_json::Value) -> RouterConfig { + serde_json::from_value(json!({ + "model_list": [ + {"model_name": "primary", "litellm_params": {"model": "openai/a"}, "model_info": info}, + {"model_name": "primary", "litellm_params": {"model": "openai/b"}, "model_info": {"id": "a2"}}, + ] + })) + .unwrap() +} + +#[rstest] +#[tokio::test] +async fn a_web_search_request_skips_the_deployments_that_declare_no_support() { + let config = with_first_info(json!({"id": "a1", "supports_web_search": false})); + let router = router(config, json!({}), Arc::new(ScriptedHost::default())); + let request = RouterRequest { + web_search: true, + ..request() + }; + for _ in 0..16 { + let picked = router + .select(&Op::from("acompletion"), &request) + .await + .unwrap(); + assert_eq!(picked.id, "a2"); + } +} + +#[rstest] +#[tokio::test] +async fn an_endpoint_a_deployment_refuses_routes_to_the_others() { + let config = with_first_info(json!({ + "id": "a1", + "enforce_supported_endpoints": true, + "supported_endpoints": ["/v1/embeddings"], + })); + let router = router(config, json!({}), Arc::new(ScriptedHost::default())); + let request = RouterRequest { + endpoint: Some("/v1/chat/completions".to_owned()), + ..request() + }; + for _ in 0..16 { + let picked = router + .select(&Op::from("acompletion"), &request) + .await + .unwrap(); + assert_eq!(picked.id, "a2"); + } +} + +#[rstest] +#[tokio::test] +async fn an_endpoint_every_deployment_refuses_fails_the_pick() { + let config: RouterConfig = serde_json::from_value(json!({ + "model_list": [{ + "model_name": "primary", + "litellm_params": {"model": "openai/a"}, + "model_info": { + "id": "a1", + "enforce_supported_endpoints": true, + "supported_endpoints": ["/v1/embeddings"], + }, + }] + })) + .unwrap(); + let router = router(config, json!({}), Arc::new(ScriptedHost::default())); + let request = RouterRequest { + endpoint: Some("/v1/chat/completions".to_owned()), + ..request() + }; + let refused = router + .call(Op::from("acompletion"), request) + .await + .unwrap_err(); + assert!(matches!( + refused, + Error::EndpointNotSupported { model, message } + if model == "primary" && message + == "/v1/chat/completions is not supported by model=primary. \ + Supported endpoints: /v1/embeddings" + )); +} + +#[rstest] +#[tokio::test] +async fn an_order_in_model_info_narrows_the_pick() { + let mut config = with_first_info(json!({"id": "a1", "order": 2})); + config.model_list[1] + .model_info + .insert("order".to_owned(), json!("1")); + let router = router(config, json!({}), Arc::new(ScriptedHost::default())); + for _ in 0..16 { + let picked = router + .select(&Op::from("acompletion"), &request()) + .await + .unwrap(); + assert_eq!(picked.id, "a2"); + } +} + +#[rstest] +#[tokio::test] +async fn a_fallback_hop_drops_the_order_the_request_named() { + let host = Arc::new(ScriptedHost::failing( + "primary", + vec![failure(FailureKind::InternalServer, 500)], + )); + let router = router( + tiers(false), + json!({"num_retries": 0, "fallbacks": [{"primary": ["backup"]}]}), + host, + ); + let request = RouterRequest { + target_order: Some(1), + ..request() + }; + let response = router.call(Op::from("acompletion"), request).await.unwrap(); + assert_eq!(response.deployment_id, "b1"); +} + +#[rstest] +#[tokio::test] +async fn a_listed_pass_through_pick_names_the_group_an_alias_resolves_to(config: RouterConfig) { + let router = router( + config, + json!({"model_alias_map": {"al": "primary"}}), + Arc::new(ScriptedHost::default()), + ); + let request = RouterRequest { + model: "al".to_owned(), + pass_through: Some(PassThrough::Listed), + ..RouterRequest::default() + }; + let refused = router + .select(&Op::from("acompletion"), &request) + .await + .unwrap_err(); + assert!(matches!( + refused, + Error::NoPassThroughDeployments { model } if model == "primary" + )); +} diff --git a/litellm/caching/dual_cache.py b/litellm/caching/dual_cache.py index c82d25c9425..af084a464ee 100644 --- a/litellm/caching/dual_cache.py +++ b/litellm/caching/dual_cache.py @@ -620,7 +620,7 @@ class DualCache(BaseCache): increment_list: list["RedisPipelineIncrementOperation"], local_only: bool = False, parent_otel_span: Span | None = None, - **kwargs, + **kwargs: object, # kwargs-ok: accepted and unused, as the cache interface's other calls take them ) -> list[float] | None: result: list[float] | None = None try: diff --git a/litellm/litellm_core_utils/add_retry_fallback_headers.py b/litellm/litellm_core_utils/add_retry_fallback_headers.py index 12b26954e4c..4cf2d1d58d8 100644 --- a/litellm/litellm_core_utils/add_retry_fallback_headers.py +++ b/litellm/litellm_core_utils/add_retry_fallback_headers.py @@ -202,7 +202,7 @@ def response_total_token_count(response: object) -> int: def apply_remaining_usage_headers( additional_headers: dict[str, object], - remaining_usage: dict[str, int], + remaining_usage: Mapping[str, object], ) -> None: for header, value in remaining_usage.items(): if value is not None and header not in additional_headers: diff --git a/litellm/litellm_core_utils/core_helpers.py b/litellm/litellm_core_utils/core_helpers.py index 43c586e7245..b28d5a4ea21 100644 --- a/litellm/litellm_core_utils/core_helpers.py +++ b/litellm/litellm_core_utils/core_helpers.py @@ -108,6 +108,17 @@ def as_object_list(value: object) -> Sequence[object] | None: return value if isinstance(value, list) else None +_Key = TypeVar("_Key") # rebind-ok: a TypeVar cannot be Final +_Value = TypeVar("_Value") # rebind-ok: a TypeVar cannot be Final + + +def without( + mapping: Mapping[_Key, _Value], *keys: object +) -> dict[_Key, _Value]: # mutable-ok: a fresh copy the caller owns + """``mapping`` without ``keys``, as a new dict in the same order; ``mapping`` is unchanged.""" + return {key: value for key, value in mapping.items() if key not in keys} + + def decoded_json_object(value: object) -> Mapping[str, object]: mapping: Final = as_str_mapping(value) if mapping is not None: 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 9a0a35979db..08ab2d1a2e5 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 @@ -1081,4 +1081,8 @@ def convert_to_model_response_object( end_time=end_time, convert_tool_call_to_json_mode=convert_tool_call_to_json_mode, ) - raise Exception(f"Invalid response object {traceback.format_exc()}\n\nreceived_args={received_args}") + # The conversion's own traceback, without the errors being handled when it ran: a router + # retry runs inside the handler of the attempt before it, and that attempt's error text + # (a rate limit) would otherwise decide how this one is classified. + own_traceback: Final = "".join(traceback.format_exception(type(e), e, e.__traceback__, chain=False)) + raise Exception(f"Invalid response object {own_traceback}\n\nreceived_args={received_args}") diff --git a/litellm/litellm_core_utils/routing_common_utils.py b/litellm/litellm_core_utils/routing_common_utils.py index a06860e7786..62bdc1c10a7 100644 --- a/litellm/litellm_core_utils/routing_common_utils.py +++ b/litellm/litellm_core_utils/routing_common_utils.py @@ -378,7 +378,8 @@ def filter_team_based_models( ids_to_remove: Final = set() if isinstance(healthy_deployments, dict): - return healthy_deployments + # A deployment asked for by id is filtered like any other: another team's is not served. + return [] if _owner_team(healthy_deployments) not in (None, request_team_id) else healthy_deployments # pyright: ignore[reportUnknownArgumentType, reportUnknownVariableType] # the legacy untyped deployment dict for deployment in healthy_deployments: _model_info = deployment.get("model_info") or {} model_team_id = _model_info.get("team_id") @@ -431,6 +432,16 @@ def _deployment_supports_web_search(deployment: dict) -> bool: return True +def is_web_search_request(request_kwargs: Mapping[str, object]) -> bool: + tools: Final[object] = request_kwargs.get("tools") + return isinstance(tools, (list, tuple)) and any( + isinstance(tool, dict) + and isinstance(tool_type := tool.get("type"), str) + and tool_type in _WEB_SEARCH_TOOL_TYPES + for tool in tools + ) + + def filter_web_search_deployments( healthy_deployments: list[dict] | dict, request_kwargs: dict | None = None, @@ -445,14 +456,7 @@ def filter_web_search_deployments( if isinstance(healthy_deployments, dict): return healthy_deployments - tools: Final[object] = request_kwargs.get("tools") - is_web_search_request: Final = isinstance(tools, (list, tuple)) and any( - isinstance(tool, dict) - and isinstance(tool_type := tool.get("type"), str) - and tool_type in _WEB_SEARCH_TOOL_TYPES - for tool in tools - ) - if not is_web_search_request: + if not is_web_search_request(request_kwargs): return healthy_deployments # Filter out deployments that don't support web search @@ -607,3 +611,9 @@ def warn_on_provider_credential_mismatch(model_name: str, litellm_params: Mappin ) verbose_router_logger.warning(warning) return warning + + +def _owner_team(deployment: Mapping[str, object]) -> object: + """The team a deployment belongs to (``model_info.team_id``), or None.""" + model_info: Final = deployment.get("model_info") + return model_info.get("team_id") if isinstance(model_info, Mapping) else None # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] # a deployment dict's model_info diff --git a/litellm/llms/anthropic/messages/adapters/streaming_iterator.py b/litellm/llms/anthropic/messages/adapters/streaming_iterator.py index c41deb00f1d..87be996ebbf 100644 --- a/litellm/llms/anthropic/messages/adapters/streaming_iterator.py +++ b/litellm/llms/anthropic/messages/adapters/streaming_iterator.py @@ -670,6 +670,8 @@ class AnthropicStreamWrapper(AdapterCompletionStreamWrapper): # Synthesized compaction block from compact_20260112 polyfill (streaming). self.compaction_block = compaction_block self.iterations_usage = iterations_usage + # The provider's error behind this stream's error frame, for a router that accounts for it. + self.stream_failure: Exception | None = None self._refusal_text: str = "" self._paused_turn: bool = False self._claude_refusal_stop_details: AnthropicStopDetails | None = None @@ -1396,6 +1398,7 @@ class AnthropicStreamWrapper(AdapterCompletionStreamWrapper): verbose_logger.exception("Anthropic Adapter - mid-stream error: %s", e) logging_obj: Final = self.litellm_logging_obj provider_error: Final = _provider_error(e) + self.stream_failure = provider_error if logging_obj is not None and logging_obj.on_detached_stream_failure is not None: if provider_error is e: raise @@ -1619,6 +1622,10 @@ class AnthropicSSEStream(AsyncIterator[bytes]): self._byte_stream: Final[AsyncIterator[bytes]] = anthropic_wrapper.async_anthropic_sse_wrapper() self._hidden_params: dict[str, object] = {} + @property + def stream_failure(self) -> Exception | None: + return self._anthropic_wrapper.stream_failure + @property def chunks(self) -> "list[ModelResponseStream] | None": return self._anthropic_wrapper.chunks diff --git a/litellm/llms/anthropic/messages/responses_adapters/streaming_iterator.py b/litellm/llms/anthropic/messages/responses_adapters/streaming_iterator.py index dae96ba4ce1..f0c3316c509 100644 --- a/litellm/llms/anthropic/messages/responses_adapters/streaming_iterator.py +++ b/litellm/llms/anthropic/messages/responses_adapters/streaming_iterator.py @@ -141,6 +141,8 @@ class AnthropicResponsesStreamWrapper: if litellm_logging_obj is not None: litellm_logging_obj.record_streamed_anthropic_message_id(self._message_id) self._unreported_failure: Exception | None = None + # The failure behind this stream's error frame, for a router that accounts for it. + self.stream_failure: Exception | None = None self._current_block_index: int = -1 # Map item_id -> content_block_index so we can stop the right block later self._item_id_to_block_index: dict[str, int] = {} @@ -467,6 +469,7 @@ class AnthropicResponsesStreamWrapper: llm_provider=getattr(self.responses_stream, "custom_llm_provider", None), model=self.model, ) + self.stream_failure = self._unreported_failure self._chunk_queue.append(_anthropic_error_chunk(status_code, message)) async def _report_unreported_failure(self) -> None: @@ -578,5 +581,9 @@ class AnthropicResponsesSSEByteStream: async def aclose(self) -> None: await self._frames.aclose() + @property + def stream_failure(self) -> Exception | None: + return self._wrapper.stream_failure + def partial_response_on_disconnect(self) -> "ResponseIncompleteEvent | None": return self._wrapper.partial_response_on_disconnect() diff --git a/litellm/router.py b/litellm/router.py index 4ee138fc9fb..6a60d5d958e 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -11,7 +11,7 @@ outside code used before the entry points live in ``DeprecatedRouterMembers`` un import inspect import weakref from collections.abc import Mapping -from types import CellType, FunctionType, MappingProxyType +from types import FunctionType, MappingProxyType from typing import TYPE_CHECKING, Final from litellm import router_python_core @@ -35,7 +35,7 @@ from litellm.types.router import ( RetryPolicy, ) -_CORE_SIGNATURE: Final = inspect.signature(PythonCore.__init__) +_CORE_SIGNATURE: Final = inspect.signature(PythonCore.__init__) # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType] # reflection reads only the parameters' names def router_arguments( @@ -65,142 +65,40 @@ else: _replaced_core_members: Final[ # mutable-ok: originals of patched core members, restored on delete dict[tuple[type, str], list[object]] ] = {} - _FACADE_ONLY: Final = frozenset( - { - "__module__", - "__qualname__", - "__doc__", - "__dict__", - "__weakref__", - "__init__", - "__init_subclass__", - "__annotations__", - "__firstlineno__", - "__static_attributes__", - "__classcell__", - "_core_class", - } - ) - - def _with_class_cell(value: object, cell: "CellType") -> object: - """``value`` with its ``__class__`` cell (what zero-argument ``super()`` reads) replaced by ``cell``.""" - if isinstance(value, staticmethod | classmethod): - return type(value)(_function_with_class_cell(value.__func__, cell)) - if isinstance(value, property): - return property( - *( - None if f is None else _function_with_class_cell(f, cell) - for f in (value.fget, value.fset, value.fdel) - ), - value.__doc__, - ) - return _function_with_class_cell(value, cell) - - def _function_with_class_cell(value: object, cell: "CellType") -> object: - """A plain function with its ``__class__`` cell replaced by ``cell``; anything else unchanged.""" - if not isinstance(value, FunctionType) or "__class__" not in value.__code__.co_freevars: - return value - closure: Final = tuple( - cell if name == "__class__" else original - for name, original in zip(value.__code__.co_freevars, value.__closure__ or (), strict=True) - ) - rebound: Final = FunctionType(value.__code__, value.__globals__, value.__name__, value.__defaults__, closure) - rebound.__kwdefaults__ = value.__kwdefaults__ - rebound.__qualname__ = value.__qualname__ - rebound.__dict__.update(value.__dict__) - return rebound - - class _Absent: - """A facade-only member, hidden from the core class: reading it raises the usual AttributeError.""" - - def __set_name__(self, owner: type, name: str) -> None: - self._name = name - - def __get__(self, instance: object, owner: type | None = None) -> object: - if instance is None: - raise AttributeError( - f"type object {getattr(owner, '__name__', owner)!r} has no attribute {self._name!r}" - ) - raise AttributeError(f"{type(instance).__name__!r} object has no attribute {self._name!r}") - - def _core_class_for(facade: type, parent: type) -> type: - """ - A core class carrying a Router subclass's own members, so the core runs its overrides. It - subclasses the facade too, so an override's ``super(Sub, self)`` finds the core it runs on. - """ - cell: Final = CellType() - members: Final = { - name: _with_class_cell(value, cell) for name, value in vars(facade).items() if name not in _FACADE_ONLY - } - namespace: Final = { - # The facade's own plumbing sits in the core class's MRO ahead of the parent core: put - # the parent core's versions back, and hide the facade-only members. - "__init__": parent.__init__, - "__getattribute__": object.__getattribute__, - "__setattr__": parent.__setattr__, - "__delattr__": parent.__delattr__, - "__dir__": parent.__dir__, - **{name: _Absent() for name in ("_attach", "_core", "from_config", "_core_class")}, - **members, - "__module__": facade.__module__, - } - core_class: Final = _CoreType(f"{facade.__name__}Core", (facade, parent), namespace) - cell.cell_contents = core_class - for name in members: - type.__delattr__(facade, name) - return core_class class _RouterType(type): """ - Class-level reads of core members (``Router._helper(...)``) go to the router's core class, and - so do class-level patches, which put the original back when the patch is removed. - ``isinstance`` and ``dir`` see the core class too. + Class-level reads of core members (``Router._helper(...)``) go to ``PythonCore``, and so do + class-level patches, which put the original back when the patch is removed. + ``isinstance`` and ``dir`` see ``PythonCore`` too. """ def __getattr__(cls, name: str) -> object: - if name == "_core_class": - raise AttributeError(name) - return getattr(cls._core_class, name) + return getattr(PythonCore, name) def __setattr__(cls, name: str, value: object) -> None: - core_class: Final = vars(cls).get("_core_class") or getattr(cls, "_core_class", None) - if core_class is None or any(name in vars(base) for base in cls.__mro__) or not hasattr(core_class, name): + if any(name in vars(base) for base in cls.__mro__) or not hasattr(PythonCore, name): super().__setattr__(name, value) return - _replaced_core_members.setdefault((core_class, name), []).append(vars(core_class).get(name, _MISSING)) - setattr(core_class, name, value) + _replaced_core_members.setdefault((PythonCore, name), []).append(vars(PythonCore).get(name, _MISSING)) + setattr(PythonCore, name, value) def __delattr__(cls, name: str) -> None: - core_class: Final = getattr(cls, "_core_class", None) - replaced: Final = _replaced_core_members.get((core_class, name)) if core_class is not None else None + replaced: Final = _replaced_core_members.get((PythonCore, name)) if not replaced or any(name in vars(base) for base in cls.__mro__): super().__delattr__(name) return original: Final = replaced.pop() if original is _MISSING: - delattr(core_class, name) + delattr(PythonCore, name) return - setattr(core_class, name, original) + setattr(PythonCore, name, original) def __instancecheck__(cls, instance: object) -> bool: - return type.__instancecheck__(cls, instance) or isinstance(instance, cls._core_class) + return type.__instancecheck__(cls, instance) or isinstance(instance, PythonCore) def __dir__(cls) -> list[str]: # mutable-ok: dir() returns a list - return sorted({*type.__dir__(cls), *dir(cls._core_class)}) - - class _CoreType(_RouterType): - """The metaclass of a Router subclass's core class: a plain class, none of the facade's redirects.""" - - def __getattr__(cls, name: str) -> object: - raise AttributeError(f"type object {cls.__name__!r} has no attribute {name!r}") - - __setattr__ = type.__setattr__ - __delattr__ = type.__delattr__ - __instancecheck__ = type.__instancecheck__ - __dir__ = type.__dir__ - - def _is_core(instance: object) -> bool: - return isinstance(type(instance), _CoreType) + return sorted({*type.__dir__(cls), *dir(PythonCore)}) def _patchable_method(core: object, name: str) -> bool: """A method of the core: an instance-level replacement of it is also kept on the facade, so @@ -212,27 +110,20 @@ else: return hasattr(type(core), "__getattr__") and inspect.ismethod(getattr(core, name, None)) class Router(DeprecatedRouterMembers, metaclass=_RouterType): - _core_class: type[PythonCore] = PythonCore - - def __init_subclass__(cls, **kwargs: object) -> None: # kwargs-ok: class keywords, passed on to super() - if isinstance(cls, _CoreType): - # A core class built below: the parent core's hooks wire it, as for any core class. - super().__init_subclass__(**kwargs) - return - # DeprecatedRouterMembers wires a subclass's overrides of old members; here that runs - # on the core class built from it, which is where the router's own calls land. - super(DeprecatedRouterMembers, cls).__init_subclass__(**kwargs) - type.__setattr__(cls, "_core_class", _core_class_for(cls, cls._core_class)) + def __init_subclass__(cls, **kwargs: object) -> None: # kwargs-ok: class keywords, refused with the class + raise TypeError( + "litellm.Router cannot be subclassed. Use the router's entry points instead: " + "register_plugin() for custom routing, apply() for settings, observe() for outcomes. " + "See https://docs.litellm.ai/docs/routing_plugins." + ) def __init__(self, *args: object, **kwargs: object) -> None: # kwargs-ok: Router's constructor arguments - self._attach( - make_core(build_router_config(router_arguments(*args, **kwargs)), python_host, type(self)._core_class) - ) + self._attach(make_core(build_router_config(router_arguments(*args, **kwargs)), python_host)) @classmethod def from_config(cls, config: RouterConfig, host: PythonHost = python_host) -> "Router": router: Final = cls.__new__(cls) - router._attach(make_core(config, host, cls._core_class)) + router._attach(make_core(config, host)) return router def _attach(self, core: RouterCore) -> None: @@ -253,8 +144,7 @@ else: try: return object.__getattribute__(self, name) except AttributeError: - if _is_core(self): - raise + pass facade: Final = object.__getattribute__(self, "__dict__") core: Final = facade.get("_core") if core is None: @@ -273,9 +163,6 @@ else: ) from None def __setattr__(self, name: str, value: object) -> None: - if _is_core(self): - super().__setattr__(name, value) - return if _class_descriptor(type(self), name): object.__setattr__(self, name, value) return @@ -289,9 +176,6 @@ else: self.__dict__[name] = value def __delattr__(self, name: str) -> None: - if _is_core(self): - super().__delattr__(name) - return if _class_descriptor(type(self), name): object.__delattr__(self, name) return @@ -299,9 +183,7 @@ else: delattr(self._core, name) def __dir__(self) -> list[str]: # mutable-ok: dir() returns a list - if _is_core(self): - return super().__dir__() - return sorted({*super().__dir__(), *dir(self.__dict__.get("_core", type(self)._core_class))}) + return sorted({*super().__dir__(), *dir(self.__dict__.get("_core", PythonCore))}) Router.__init__.__signature__ = _CORE_SIGNATURE.replace( # pyright: ignore[reportFunctionMemberAccess] # inspect reads it parameters=[parameter for parameter in _CORE_SIGNATURE.parameters.values() if parameter.name != "_host"] diff --git a/litellm/router_api.py b/litellm/router_api.py index f22e9b1f351..82660439ae9 100644 --- a/litellm/router_api.py +++ b/litellm/router_api.py @@ -32,6 +32,7 @@ as they were. # members are still untyped (bare ``dict`` / ``list``) or mutable (the live objects); the overloads # repeat them unchanged, and typing the members themselves is the core's job. +import copy import hashlib import inspect import json @@ -45,6 +46,7 @@ from typing import TYPE_CHECKING, Final, Literal, Protocol, TypeAlias, overload from litellm.router_config import detached if TYPE_CHECKING: + from functools import _lru_cache_wrapper from typing import Any # noqa: TID251 # only to spell the cores' untyped (Unknown) containers and one Any type from opentelemetry.trace import Span @@ -179,9 +181,9 @@ def deployment_usable_by_team(deployment: "Mapping[str, object] | Deployment", t Whether ``team_id`` may use ``deployment``: a team-scoped deployment (``model_info.team_id`` set) only by that team, any other deployment by everyone. Needs no router instance. """ - from litellm.router_python_core import PythonCore + from litellm.router_rust.surface.registry import HostRegistry - return PythonCore._core_deployment_usable_by_team(deployment, team_id) + return HostRegistry._core_deployment_usable_by_team(deployment, team_id) def deployment_id(model_group: str, litellm_params: Mapping[str, object]) -> str: @@ -201,11 +203,77 @@ def deployment_id(model_group: str, litellm_params: Mapping[str, object]) -> str return hashlib.sha256("".join(parts).encode()).hexdigest() +# The arguments ``Router.__init__`` accepts, in its signature's order. Every core reads this list +# (the Python core's ``get_valid_args`` reads its constructor's signature, and the core suite checks +# that the two agree). +ROUTER_INIT_ARGS: Final[tuple[str, ...]] = ( + "model_list", + "assistants_config", + "search_tools", + "guardrail_list", + "redis_url", + "redis_host", + "redis_port", + "redis_password", + "redis_db", + "cache_responses", + "cache_kwargs", + "caching_groups", + "client_ttl", + "polling_interval", + "default_priority", + "num_retries", + "max_fallbacks", + "timeout", + "stream_timeout", + "default_litellm_params", + "default_max_parallel_requests", + "set_verbose", + "debug_level", + "default_fallbacks", + "fallbacks", + "context_window_fallbacks", + "content_policy_fallbacks", + "model_group_alias", + "enable_pre_call_checks", + "enable_tag_filtering", + "tag_filtering_match_any", + "tag_routing_prefix", + "plugins", + "retry_after", + "retry_policy", + "model_group_retry_policy", + "allowed_fails", + "allowed_fails_policy", + "cooldown_time", + "disable_cooldowns", + "fallback_on_unsupported_params", + "routing_strategy", + "optional_pre_call_checks", + "routing_strategy_args", + "routing_groups", + "provider_budget_config", + "alerting_config", + "router_general_settings", + "deployment_affinity_ttl_seconds", + "model_group_affinity_config", + "ignore_invalid_deployments", + "enable_health_check_routing", + "health_check_staleness_threshold", + "health_check_ignore_transient_errors", + "background_health_check_model_groups", + "enable_weighted_failover", + "enable_mid_stream_fallback_continuation", + "fallback_access_check", + "fallback_budget_check", + "auto_router_capability_limit", + "session_affinity_repin_on_failover", +) + + def router_init_args() -> tuple[str, ...]: """The arguments ``Router.__init__`` accepts (was ``Router.get_valid_args``). Needs no router instance.""" - from litellm.router_python_core import PythonCore - - return tuple(PythonCore.get_valid_args()) + return ROUTER_INIT_ARGS def mark_advisor_orchestration_failure(exception: BaseException) -> None: @@ -269,9 +337,9 @@ def deployment_metadata(request_kwargs: Mapping[str, object]) -> Mapping[str, ob def _json_default_stable_id(obj: object) -> str: - from litellm.router_python_core import PythonCore + from litellm.router_rust.steps.prerouting import _json_default_stable_id as stable_id - return PythonCore._json_default_stable_id(obj) + return stable_id(obj) class RouterRegistry: @@ -316,6 +384,13 @@ class RouterRegistry: return self._router._core_cached_deployment_model_info(id, model_name) return self._router._core_get_deployment_model_info(model_id=id, model_name=model_name) + def deployment_model_info_cache(self) -> "_lru_cache_wrapper[ModelMapInfo | None]": + """ + The lru cache ``deployment_model_info(..., cached=True)`` reads, called as ``(id, model_name)`` + and cleared with ``cache_clear()`` (was the ``cached_deployment_model_info`` attribute). + """ + return self._router._core_cached_deployment_model_info + def router_model_info( self, deployment: "dict[str, object] | Deployment | None", # mutable-ok: the core takes a dict @@ -532,10 +607,8 @@ class RouterRegistry: ``should_include_deployment``), or, without ``model_name``, whether the team may use it at all (was ``deployment_usable_by_team``). """ - from litellm.router_python_core import PythonCore - if model_name is None: - return PythonCore._core_deployment_usable_by_team(deployment, team_id) + return deployment_usable_by_team(deployment, team_id) return self._router._core_should_include_deployment(model_name=model_name, model=deployment, team_id=team_id) @overload @@ -1074,13 +1147,22 @@ PLAIN_SETTING_ATTRIBUTES: Final[frozenset[str]] = frozenset({"num_retries", "max ASSIGNABLE_ALIASES: Final[Mapping[str, str]] = MappingProxyType({"model_group_alias": "_core_model_group_alias"}) _MISSING: Final = object() +# The settings fields ``RouterSettings`` copies, so an in-place edit of what a read returns never +# reaches the router (only ``apply()`` changes settings). +_COPIED_SETTING_FIELDS: Final = ( + "fallbacks", + "context_window_fallbacks", + "content_policy_fallbacks", + "default_litellm_params", +) @dataclass(frozen=True) class RouterSettings: """ The router's configuration as one frozen snapshot (``settings()``). Fields hold the router's - own objects, not copies: mutating a list in here mutates the router's. + own objects, except the fallback lists and ``default_litellm_params``, which are copies: settings + change only through ``apply()``, so editing one of those in place changes nothing on any core. A field the router never set reads as ``None``. Each field has the type of the Router attribute it reads (``_core_``), the type the old @@ -1119,6 +1201,12 @@ class RouterSettings: _router: "Router | None" = field(default=None, repr=False, compare=False) _unset: frozenset[str] = field(default=frozenset(), repr=False, compare=False) + def __post_init__(self) -> None: + for name in _COPIED_SETTING_FIELDS: + value: object = getattr(self, name) # pyright: ignore[reportAny] # a settings field + if isinstance(value, list | dict): + object.__setattr__(self, name, copy.deepcopy(value)) # pyright: ignore[reportUnknownArgumentType] # a settings list or dict + @classmethod def of(cls, router: "Router") -> "RouterSettings": read: Final = {name: getattr(router, "_core_" + name, _MISSING) for name in SETTING_FIELDS} @@ -1188,16 +1276,16 @@ class RouterHealth: @staticmethod def is_blocked(deployment: "Deployment") -> bool: """Whether an admin paused this deployment (was ``is_deployment_blocked``).""" - from litellm.router_python_core import PythonCore + from litellm.router_rust.surface.health_routing import is_blocked as deployment_is_blocked - return PythonCore._core_is_deployment_blocked(deployment) + return deployment_is_blocked(deployment) @staticmethod def all_blocked(deployments: "Sequence[DeploymentTypedDict]") -> bool: """Whether every deployment in the list is paused (was ``are_all_deployments_blocked``).""" - from litellm.router_python_core import PythonCore + from litellm.router_rust.surface.health_routing import all_blocked as deployments_all_blocked - return PythonCore._core_are_all_deployments_blocked(deployments) + return deployments_all_blocked(deployments) def fully_blocked_model_names(self) -> set[str]: # mutable-ok: the core's type, a fresh set """Groups whose every deployment is paused (was ``get_fully_blocked_model_names``).""" diff --git a/litellm/router_attempt.py b/litellm/router_attempt.py new file mode 100644 index 00000000000..33e0a19c530 --- /dev/null +++ b/litellm/router_attempt.py @@ -0,0 +1,215 @@ +""" +One router attempt, as the host runs it: the request for the deployment the router core picked. + +The Rust router core picks a deployment and hands the attempt to its host; ``PythonHost`` builds +the keyword arguments for that one deployment here and calls ``litellm.`` with them. Nothing +here selects, retries, falls back or records cooldowns: the core does all of that. Nothing here +reads a Python router either, so this module outlives ``router_python_core.py``. +""" + +import re +import uuid +from collections.abc import Mapping +from types import MappingProxyType +from typing import Final + +from litellm.litellm_core_utils.batch_utils import get_router_metadata_variable_name +from litellm.litellm_core_utils.get_llm_provider_logic import inferred_provider +from litellm.litellm_core_utils.pattern_match_deployments import PatternMatchRouter +from litellm.litellm_core_utils.routing_common_utils import without_router_only_kwargs + +# Request keys the router reads to route the call; the deployment call never receives them. +ROUTING_KEYS: Final = frozenset( + { + "model", + "num_retries", + "disable_fallbacks", + "specific_deployment", + "_router_weights", + "_retry_skipped_deployment_ids", + "_target_order", + "model_group_retry_policy", + "fallbacks", + "context_window_fallbacks", + "content_policy_fallbacks", + "mock_testing_fallbacks", + "mock_testing_context_fallbacks", + "mock_testing_content_policy_fallbacks", + "mock_testing_rate_limit_error", + } +) + +# The positional parameters of each endpoint, in litellm's order: a call that passes them +# positionally is read back into keywords before it is routed. +POSITIONAL: Final = MappingProxyType( + { + "acompletion": ("model", "messages"), + "completion": ("model", "messages"), + "atext_completion": ("model", "prompt"), + "text_completion": ("model", "prompt"), + "aembedding": ("model", "input"), + "embedding": ("model", "input"), + "aimage_generation": ("prompt", "model"), + "image_generation": ("prompt", "model"), + "atranscription": ("model", "file"), + "aspeech": ("model", "input", "voice"), + "arerank": ("model", "query", "documents"), + "amoderation": ("input", "model"), + "aresponses": ("input", "model"), + "aanthropic_messages": (), + } +) + + +def request_from_call(op: str, args: tuple[object, ...], kwargs: Mapping[str, object]) -> Mapping[str, object]: + """The call's arguments as one keyword request. Raises ``TypeError`` like the endpoint would.""" + names: Final = POSITIONAL.get(op, ()) + if len(args) > len(names): + raise TypeError(f"{op}() takes {len(names)} positional arguments but {len(args)} were given") + request: Final = dict(kwargs) + for name, value in zip(names, args, strict=False): + if name in request: + raise TypeError(f"{op}() got multiple values for argument {name!r}") + request[name] = value + return request + + +# Endpoints the Python router serves through ``_ageneric_api_call_with_fallbacks``, which writes +# the routing metadata under ``litellm_metadata``. +_GENERIC_API_OPS: Final = frozenset({"aanthropic_messages", "aresponses"}) + + +def metadata_field(op: str) -> str: + """``metadata`` or ``litellm_metadata``: where the router writes what it routed the call as.""" + return "litellm_metadata" if op in _GENERIC_API_OPS else get_router_metadata_variable_name(function_name=op) + + +def prepare_call(op: str, request: Mapping[str, object], model_group_alias: str | None) -> Mapping[str, object]: + """What the router adds to a request once, before its first attempt + (``_update_kwargs_before_fallbacks``): a trace id and the model group it was asked for.""" + field: Final = metadata_field(op) + metadata: Final = dict(_as_mapping(request.get(field))) + metadata.update({"model_group": request.get("model"), "model_group_alias": model_group_alias}) + return {**request, "litellm_trace_id": request.get("litellm_trace_id") or str(uuid.uuid4()), field: metadata} + + +def _wildcard_match(pattern: object, requested: str) -> "re.Match[str] | None": + if not isinstance(pattern, str) or "*" not in pattern: + return None + regex: Final = PatternMatchRouter().pattern_to_regex(pattern) + return re.match(regex, requested) or re.match(regex, f"{inferred_provider(requested)}/{requested}") + + +def for_requested_model(deployment: Mapping[str, object], requested: str | None) -> Mapping[str, object]: + """A wildcard deployment as the router serves ``requested`` through it: its ``model`` with + the wildcards filled in from the requested name (``PatternMatchRouter.route``), by the team's + public pattern (``_core_team_pattern_routers``) or its own. Any other deployment comes back + unchanged.""" + params: Final = _as_mapping(deployment.get("litellm_params")) + model: Final = params.get("model") + if requested is None or not isinstance(model, str): + return deployment + public_name: Final = _as_mapping(deployment.get("model_info")).get("team_public_model_name") + match: Final = _wildcard_match(public_name, requested) or _wildcard_match(deployment.get("model_name"), requested) + if match is None: + return deployment + return { + **deployment, + "litellm_params": {**params, "model": PatternMatchRouter.set_deployment_model_name(match, model)}, + } + + +def attempt_kwargs( + op: str, + request: Mapping[str, object], + deployment: Mapping[str, object], + *, + model_group: str, + default_litellm_params: Mapping[str, object], + timeout: float | None, +) -> Mapping[str, object]: + """The keyword arguments for ``litellm.`` against one deployment + (``_update_kwargs_with_deployment`` and the call sites that merge ``litellm_params``).""" + params: Final = _as_mapping(deployment.get("litellm_params")) + model_info: Final = _model_info_for(request, deployment, model_group) + field: Final = metadata_field(op) + metadata: Final = dict(_as_mapping(request.get(field))) + metadata.update( + { + "model_group": model_group, + "deployment": params.get("model"), + "model_info": model_info, + "api_base": params.get("api_base"), + "deployment_model_name": deployment.get("model_name"), + } + ) + tags: Final = _as_list(params.get("tags")) + if tags: + merged: Final = list(_as_list(metadata.get("tags"))) + merged.extend(tag for tag in tags if tag not in merged) + metadata["tags"] = merged + call: Final = {key: value for key, value in request.items() if key not in ROUTING_KEYS} + for key, value in default_litellm_params.items(): + call.setdefault(key, value) + call[field] = metadata + call["model_info"] = model_info + if "timeout" not in call: + deployment_timeout: Final = params.get("timeout") + chosen: Final = deployment_timeout if deployment_timeout is not None else timeout + if chosen is not None: + call["timeout"] = chosen + return without_router_only_kwargs({**params, **call, "model": params.get("model")}) + + +def _model_info_for( + request: Mapping[str, object], deployment: Mapping[str, object], model_group: str +) -> dict[str, object]: # mutable-ok: a fresh dict for the call + """The deployment's ``model_info``, or, for a request carrying its own credentials, the + request-scoped deployment's (``_handle_clientside_credential``): its own stable id, and the + configured deployment's id as ``original_model_id``.""" + from litellm.litellm_core_utils import clientside_credential_handler as clientside + from litellm.router_api import deployment_id + + model_info: Final = dict(_as_mapping(deployment.get("model_info"))) + if not clientside.is_clientside_credential(request_kwargs=dict(request)): # pyright: ignore[reportUnknownMemberType] # an untyped helper + return model_info + dynamic: Final = _as_mapping( + clientside.get_dynamic_litellm_params( # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType] # an untyped helper + litellm_params=dict(_as_mapping(deployment.get("litellm_params"))), request_kwargs=dict(request) + ) + ) + original: Final = model_info.get("id") + model_info["id"] = deployment_id(model_group, dynamic) + model_info["original_model_id"] = original + return model_info + + +def stamp_attempt_error(error: BaseException, deployment: Mapping[str, object]) -> None: + """What the router stamps on a failed attempt's error: the deployment's own ``num_retries`` + (``_set_deployment_num_retries_on_exception``) and its id (``failed_deployment_id``).""" + params: Final = _as_mapping(deployment.get("litellm_params")) + if getattr(error, "num_retries", None) is None and params.get("num_retries") is not None: + try: + error.num_retries = int(str(params["num_retries"])) # pyright: ignore[reportAttributeAccessIssue] # the router stamps the same attribute # rebind-ok: the router stamps the error in place + except ValueError: + pass + deployment_id: Final = _as_mapping(deployment.get("model_info")).get("id") + if deployment_id and not getattr(error, "failed_deployment_id", None): + try: + error.failed_deployment_id = deployment_id # pyright: ignore[reportAttributeAccessIssue] # the router stamps the same attribute # rebind-ok: see above + except AttributeError: + pass + + +def _as_mapping(value: object) -> Mapping[str, object]: + if isinstance(value, Mapping): + entries: Final[Mapping[object, object]] = value # pyright: ignore[reportUnknownVariableType] # a request's untyped mapping + return {str(key): item for key, item in entries.items()} + return MappingProxyType({}) + + +def _as_list(value: object) -> list[object]: # mutable-ok: a fresh list + if isinstance(value, list): + items: Final[list[object]] = value # pyright: ignore[reportUnknownVariableType] # a request's untyped list # mutable-ok: read, then copied + return list(items) + return [] diff --git a/litellm/router_config.py b/litellm/router_config.py index 75c79068191..36154af296c 100644 --- a/litellm/router_config.py +++ b/litellm/router_config.py @@ -79,7 +79,12 @@ def detached(value: object) -> object: def build_router_config(arguments: Mapping[str, object]) -> RouterConfig: - """The constructor arguments, detached from the caller's objects: settings are static.""" + """The constructor arguments, detached from the caller's objects: settings are static. + + A bool `cooldown_time` is refused: `True` is not a number of seconds, and + reading it as 1 s or as the default would both be a guess.""" + if isinstance(arguments.get("cooldown_time"), bool): + raise TypeError("Router cooldown_time must be a number of seconds, not a bool") return RouterConfig(arguments=MappingProxyType({name: detached(value) for name, value in arguments.items()})) diff --git a/litellm/router_core.py b/litellm/router_core.py index acbe82a0643..c077b3f0015 100644 --- a/litellm/router_core.py +++ b/litellm/router_core.py @@ -1,14 +1,16 @@ """ ``RouterCore``: the interface ``litellm.Router`` delegates to, and ``make_core``, which picks one. -``PythonCore`` (``litellm/router_python_core.py``) is today's routing logic. The Rust core -(litellm-rust/crates/router-core) implements the same operations. The core is picked once, when -the router is built, from its ``RouterConfig``. +``PythonCore`` (``litellm/router_python_core.py``) is today's routing logic. ``RustCore`` +(``litellm/router_rust/core.py`` over litellm-rust/crates/router-core) is a separate implementation +that reaches outside itself only through its host; it does not use ``PythonCore``. The core is +picked once, when the router is built, from its ``RouterConfig``. """ from collections.abc import Mapping -from typing import TYPE_CHECKING, Protocol, runtime_checkable +from typing import TYPE_CHECKING, Final, Protocol, runtime_checkable +from litellm._logging import verbose_router_logger from litellm.router_config import RouterConfig if TYPE_CHECKING: @@ -20,7 +22,6 @@ if TYPE_CHECKING: RouterSettings, RoutingObservers, ) - from litellm.router_python_core import PythonCore from litellm.router_python_host import PythonHost from litellm.types.router import CustomRoutingStrategyBase @@ -50,7 +51,26 @@ class RouterCore(Protocol): def retire(self) -> None: ... -def make_core(config: RouterConfig, host: "PythonHost", core_class: "type[PythonCore] | None" = None) -> RouterCore: +def make_core(config: RouterConfig, host: "PythonHost") -> RouterCore: + """ + The Rust core when the route catalog's router rule allows it and the Rust core serves this + config, else the Python core. Decided once, for the router's life, from the config alone: the + Rust core never builds, wraps or falls back to a Python core. + """ from litellm.router_python_core import PythonCore + from litellm.rust_bridge.catalog import Route, RouteContext, decision + from litellm.rust_bridge.configuration import Decision - return (core_class or PythonCore)(**config.arguments, _host=host) # pyright: ignore[reportArgumentType] # the arguments were bound against PythonCore's own signature + chosen: Final = decision(RouteContext(Route.ROUTER)) + if chosen is not Decision.PYTHON: + from litellm.router_rust.core import RustCore + from litellm.rust_bridge._native import RustBridgeDeclined + + try: + return RustCore.build(config, host) # pyright: ignore[reportReturnType] # RustCore serves the RouterCore entry points; other members raise RustCoreUnsupported + except RustBridgeDeclined as declined: + # A config the Rust core cannot take keeps the router on Python while both cores exist. + if chosen is Decision.RUST_REQUIRED: + raise + verbose_router_logger.info("Router served by the Python core: %s", declined) + return PythonCore(**config.arguments, _host=host) # pyright: ignore[reportArgumentType] # the arguments were bound against PythonCore's own signature diff --git a/litellm/router_deprecated.py b/litellm/router_deprecated.py index 0a02408a717..9d915fe0982 100644 --- a/litellm/router_deprecated.py +++ b/litellm/router_deprecated.py @@ -5,8 +5,10 @@ DeprecatedRouterMembers: the Router members that outside code used before the en entry points (``registry()``, ``settings()``, ``health()``, ``select()``, ``apply()``, ``observe()``, ``register_plugin()``, ``retire()``; see litellm/router_api.py). Every member here returns what it returned before, with the type it had before, and names its replacement in a -``DeprecationWarning``. Reads go to the same Router implementation (``Router._core_``) the -entry point uses; writes go through ``apply`` / ``aapply`` so the registry version moves. +``DeprecationWarning``. Reads go through that entry point, so they work on either router core +(Python or Rust) and, on the Python core, return the same ``_core_`` object as before. +Settings writes go through ``apply`` / ``aapply`` so the registry version moves; any other old +attribute write sets the core's ``_core_``, which each core keeps where its reads find it. Access from the router itself and from its own modules (``litellm.router*``) is not counted: those are the router's internals and move behind the core interface separately. Every other access @@ -17,11 +19,10 @@ deleted (planned about a month after the core swap). # pyright: reportPrivateUsage=false, reportMissingTypeArgument=false, reportUnknownParameterType=false # pyright: reportUnknownMemberType=false, reportUnknownVariableType=false, reportUnknownArgumentType=false # pyright: reportUnknownLambdaType=false -# This shim is part of the router (it reads Router's ``_core_*`` members, as Router's own methods do) +# This shim is part of the router (it writes Router's ``_core_*`` members, as Router's own methods do) # and repeats each old member's signature exactly, untyped parts included (``dict``, ``list``), so # callers keep the types they had. Typing those parts is the core's job, not the shim's. -import contextvars import functools import inspect import sys @@ -31,7 +32,7 @@ from collections.abc import Awaitable, Callable, Coroutine, Iterator, Mapping, S from types import MappingProxyType from typing import TYPE_CHECKING, Final, Generic, ParamSpec, TypeVar, overload -from litellm.router_api import ASSIGNABLE_ALIASES, SETTING_FIELDS, ChangeSet +from litellm.router_api import ASSIGNABLE_ALIASES, SETTING_FIELDS, ChangeSet, RoutingContext from litellm.router_strategy.least_busy import picks_reserve if TYPE_CHECKING: @@ -40,6 +41,7 @@ if TYPE_CHECKING: from opentelemetry.trace import Span from litellm.caching.caching import RedisCache + from litellm.integrations.SlackAlerting.slack_alerting import SlackAlerting from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging from litellm.litellm_core_utils.ptu_shares import PTUTeamCeiling from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler @@ -131,6 +133,9 @@ class _DeprecatedAttribute(Generic[_T]): _note(self.name, self.replacement) router: Final = _as_router(obj) try: + if obj.__dict__.get("_core") is None: + # No core yet (a Router made without __init__): the value written through this name. + return getattr(obj, "_core_" + self.name) # pyright: ignore[reportAny] # the pending write, as given return self.read(router) except AttributeError as e: raise _public_attribute_error(router, "_core_" + self.name, self.name, e) from None @@ -178,243 +183,174 @@ def _without_reserving(lookup: Callable[_P, Awaitable[_T]]) -> Callable[_P, Coro return looked_up -# (id(router), member name) pairs whose subclass override is running in this context. -_OVERRIDES_RUNNING: Final[contextvars.ContextVar[frozenset[tuple[int, str]]]] = contextvars.ContextVar( - "router_overrides_running", default=frozenset() -) - - -def _override_key(router: object, name: str) -> tuple[int, str]: - """The running-override key for ``router``: a Router and its core share one, so an old member - entered on the Router is seen by the override routing that runs on the core.""" - core: Final = getattr(router, "__dict__", {}).get("_core") - return (id(router if core is None else core), name) - - -def _enter_override(name: str, override: Callable[..., object]) -> Callable[..., object]: +def _async_picker(router: "Router") -> Callable[..., Awaitable[object]]: """ - A direct call of a subclass override (``router.()``) marks it running, so its - ``super().()`` reaches the parent once instead of re-entering the override through - ``_core_``. + The router's async pick (the old ``async_get_available_deployment`` attribute). A core that + keeps its pick as a member (the Python core's, or one assigned through this attribute) returns + that member, so a patched pick reads back as patched; any other core picks through ``select``. """ - if inspect.iscoroutinefunction(override): - async_override: Final[Callable[..., Awaitable[object]]] = override + member: Final[object] = getattr(router, "_core_async_get_available_deployment", None) + if callable(member): + return member # pyright: ignore[reportReturnType] # the core's own async pick - @functools.wraps(override) - async def entered_async( - router: object, - *args: object, - **kwargs: object, # kwargs-ok: passes the override's arguments through unchanged - ) -> object: - token: Final = _OVERRIDES_RUNNING.set(_OVERRIDES_RUNNING.get() | {_override_key(router, name)}) - try: - return await async_override(router, *args, **kwargs) - finally: - _OVERRIDES_RUNNING.reset(token) - - return entered_async - - @functools.wraps(override) - def entered( - router: object, - *args: object, - **kwargs: object, # kwargs-ok: passes the override's arguments through unchanged + async def async_get_available_deployment( + model: str, + request_kwargs: dict, # mutable-ok: the old parameter type, unchanged + messages: list[dict[str, str]] | None = None, # mutable-ok: the old parameter type, unchanged + input: str | list | None = None, # mutable-ok: the old parameter type, unchanged + specific_deployment: bool | None = False, ) -> object: - token: Final = _OVERRIDES_RUNNING.set(_OVERRIDES_RUNNING.get() | {_override_key(router, name)}) - try: - return override(router, *args, **kwargs) - finally: - _OVERRIDES_RUNNING.reset(token) + return await router.select( + model, request_kwargs, messages=messages, input=input, specific_deployment=specific_deployment + ) - return entered + return async_get_available_deployment -def _private_alias(name: str, member: Callable[..., object]) -> Callable[..., object]: +def _unreserved_picker(router: "Router") -> Callable[..., Awaitable[object]]: + """``_async_picker``, picking without holding a least-busy reservation (a lookup makes no + call). A pick assigned through the attribute reads back as the object assigned.""" + picker: Final = _async_picker(router) + if inspect.ismethod(picker) or picker.__name__ == "async_get_available_deployment": + return _without_reserving(picker) + return picker + + +def _slack_alerting_logger(router: "Router") -> "SlackAlerting": + """``observe().slack_alerting_logger``; a router without one has no such attribute, as before.""" + logger: Final = router.observe().slack_alerting_logger + if logger is None: + raise AttributeError( + f"'{type(router).__name__}' object has no attribute 'slack_alerting_logger'", + name="slack_alerting_logger", + obj=router, + ) + return logger + + +def _run_ready(step: Coroutine[object, object, _T]) -> _T: """ - A private alias of an old member (``_get_all_deployments``). It was bound to Router's own - method, so it never reached a subclass override of ``name``, and still doesn't. + The result of a coroutine that finishes without suspending: ``select(kind='pass_through')`` + picks synchronously on every core, so the old synchronous member can return its result. """ - return _enter_override(name, member) - - -def _route_override(cls: type, name: str, override: Callable[..., object]) -> Callable[..., object]: - """ - Make the router's own calls to ``_core_`` reach a subclass override of ````. - - When the override calls ``super().()``, the call comes back to ``_core_``; it is - then sent on to the parent's implementation instead of back into the override. - """ - core: Final = "_core_" + name - - def parent(router: object) -> Callable[..., object]: - return getattr(super(cls, router), core) - - if inspect.iscoroutinefunction(override): - async_override: Final[Callable[..., Awaitable[object]]] = override - - @functools.wraps(override) - async def routed_async( - router: object, - *args: object, - **kwargs: object, # kwargs-ok: passes the override's arguments through unchanged - ) -> object: - key: Final = _override_key(router, name) - running: Final = _OVERRIDES_RUNNING.get() - if key in running: - parent_async: Final[Callable[..., Awaitable[object]]] = getattr(super(cls, router), core) - return await parent_async(*args, **kwargs) - token: Final = _OVERRIDES_RUNNING.set(running | {key}) - try: - return await async_override(router, *args, **kwargs) - finally: - _OVERRIDES_RUNNING.reset(token) - - return routed_async - - @functools.wraps(override) - def routed( - router: object, - *args: object, - **kwargs: object, # kwargs-ok: passes the override's arguments through unchanged - ) -> object: - key: Final = _override_key(router, name) - running: Final = _OVERRIDES_RUNNING.get() - if key in running: - return parent(router)(*args, **kwargs) - token: Final = _OVERRIDES_RUNNING.set(running | {key}) - try: - return override(router, *args, **kwargs) - finally: - _OVERRIDES_RUNNING.reset(token) - - return routed + try: + step.send(None) + except StopIteration as done: + return done.value # pyright: ignore[reportAny] # the coroutine's own result + step.close() + raise RuntimeError("select(kind='pass_through') suspended; it must pick without awaiting") class DeprecatedRouterMembers: """Old Router members, kept for about a month. Each names its replacement.""" - def __init_subclass__(cls, **kwargs: object) -> None: # kwargs-ok: class keywords, passed on to super() - # A Router subclass that overrides an old member (say async_pre_routing_hook) keeps - # having its override called by the router, which now calls _core_ internally. - super().__init_subclass__(**kwargs) - for name, value in list(vars(cls).items()): - if ( - name.startswith("_core_") - or name not in vars(DeprecatedRouterMembers) - or not hasattr(cls, "_core_" + name) - ): - continue - if isinstance(value, staticmethod): - # Its super() call reaches the old member here, which calls Router's own implementation. - setattr(cls, "_core_" + name, value) - elif inspect.isfunction(value): - setattr(cls, "_core_" + name, _route_override(cls, name, value)) - setattr(cls, name, _enter_override(name, value)) - # -- registry: attributes -------------------------------------------------------------------- - model_list = _attr("registry().all()", lambda r: r._core_model_list) - model_names = _attr("registry().names(kind='group_set')", lambda r: r._core_model_names) - deployment_names = _attr("registry().names(kind='deployment_names')", lambda r: r._core_deployment_names) - team_public_model_names = _attr("registry().names(kind='team_public')", lambda r: r._core_team_public_model_names) + model_list = _attr("registry().all()", lambda r: r.registry().all()) + model_names = _attr("registry().names(kind='group_set')", lambda r: r.registry().names(kind="group_set")) + deployment_names = _attr( + "registry().names(kind='deployment_names')", lambda r: r.registry().names(kind="deployment_names") + ) + team_public_model_names = _attr( + "registry().names(kind='team_public')", lambda r: r.registry().names(kind="team_public") + ) model_name_to_deployment_indices = _attr( - "registry().index('model_name')", lambda r: r._core_model_name_to_deployment_indices + "registry().index('model_name')", lambda r: r.registry().index("model_name") ) team_model_to_deployment_indices = _attr( - "registry().index('team_model')", lambda r: r._core_team_model_to_deployment_indices + "registry().index('team_model')", lambda r: r.registry().index("team_model") ) cached_deployment_model_info = _attr( "registry().deployment_model_info(id, model_name, cached=True)", - lambda r: r._core_cached_deployment_model_info, + lambda r: r.registry().deployment_model_info_cache(), ) - model_group_alias = _attr("registry().aliases()", lambda r: r._core_model_group_alias) - pattern_router = _attr("registry().patterns()", lambda r: r._core_pattern_router) - team_pattern_routers = _attr("registry().patterns(team=True)", lambda r: r._core_team_pattern_routers) - routing_groups = _attr("registry().routing_groups()", lambda r: r._core_routing_groups) + model_group_alias = _attr("registry().aliases()", lambda r: r.registry().aliases()) + pattern_router = _attr("registry().patterns()", lambda r: r.registry().patterns()) + team_pattern_routers = _attr("registry().patterns(team=True)", lambda r: r.registry().patterns(team=True)) + routing_groups = _attr("registry().routing_groups()", lambda r: r.registry().routing_groups()) _routing_groups = routing_groups - auto_routers = _attr("registry().plugins('auto')", lambda r: r._core_auto_routers) - complexity_routers = _attr("registry().plugins('complexity')", lambda r: r._core_complexity_routers) - quality_routers = _attr("registry().plugins('quality')", lambda r: r._core_quality_routers) - adaptive_routers = _attr("registry().plugins('adaptive')", lambda r: r._core_adaptive_routers) - routing_plugins = _attr("registry().plugins('routing')", lambda r: r._core_routing_plugins) + auto_routers = _attr("registry().plugins('auto')", lambda r: r.registry().plugins("auto")) + complexity_routers = _attr("registry().plugins('complexity')", lambda r: r.registry().plugins("complexity")) + quality_routers = _attr("registry().plugins('quality')", lambda r: r.registry().plugins("quality")) + adaptive_routers = _attr("registry().plugins('adaptive')", lambda r: r.registry().plugins("adaptive")) + routing_plugins = _attr("registry().plugins('routing')", lambda r: r.registry().plugins("routing")) prompt_caching_pin_cache = _attr( - "registry().plugins('prompt_caching_pins')", lambda r: r._core_prompt_caching_pin_cache + "registry().plugins('prompt_caching_pins')", lambda r: r.registry().plugins("prompt_caching_pins") ) # -- health and selection: attributes -------------------------------------------------------- - cache = _attr("health().store", lambda r: r._core_cache) - cooldown_cache = _attr("health().cooldowns", lambda r: r._core_cooldown_cache) - health_state_cache = _attr("health().health_checks", lambda r: r._core_health_state_cache) - deployment_latency_map = _attr("health().latency", lambda r: r._core_deployment_latency_map) - async_get_available_deployment = _attr( - "select(model, request_kwargs)", - lambda r: _without_reserving(r._core_async_get_available_deployment), - ) + cache = _attr("health().store", lambda r: r.health().store) + cooldown_cache = _attr("health().cooldowns", lambda r: r.health().cooldowns) + health_state_cache = _attr("health().health_checks", lambda r: r.health().health_checks) + deployment_latency_map = _attr("health().latency", lambda r: r.health().latency) + async_get_available_deployment = _attr("select(model, request_kwargs)", lambda r: _unreserved_picker(r)) # -- observers ------------------------------------------------------------------------------ - slack_alerting_logger = _attr("observe().slack_alerting_logger", lambda r: r._core_slack_alerting_logger) + slack_alerting_logger = _attr("observe().slack_alerting_logger", _slack_alerting_logger) # -- settings ------------------------------------------------------------------------------- - routing_strategy = _setting("routing_strategy", lambda r: r._core_routing_strategy) - context_window_fallbacks = _setting("context_window_fallbacks", lambda r: r._core_context_window_fallbacks) - content_policy_fallbacks = _setting("content_policy_fallbacks", lambda r: r._core_content_policy_fallbacks) - default_deployment = _setting("default_deployment", lambda r: r._core_default_deployment) - default_litellm_params = _setting("default_litellm_params", lambda r: r._core_default_litellm_params) + routing_strategy = _setting("routing_strategy", lambda r: r.settings().routing_strategy) + context_window_fallbacks = _setting("context_window_fallbacks", lambda r: r.settings().context_window_fallbacks) + content_policy_fallbacks = _setting("content_policy_fallbacks", lambda r: r.settings().content_policy_fallbacks) + default_deployment = _setting("default_deployment", lambda r: r.settings().default_deployment) + default_litellm_params = _setting("default_litellm_params", lambda r: r.settings().default_litellm_params) default_max_parallel_requests = _setting( - "default_max_parallel_requests", lambda r: r._core_default_max_parallel_requests + "default_max_parallel_requests", lambda r: r.settings().default_max_parallel_requests + ) + cooldown_time = _setting("cooldown_time", lambda r: r.settings().cooldown_time) + allowed_fails = _setting("allowed_fails", lambda r: r.settings().allowed_fails) + allowed_fails_policy = _setting("allowed_fails_policy", lambda r: r.settings().allowed_fails_policy) + disable_cooldowns = _setting("disable_cooldowns", lambda r: r.settings().disable_cooldowns) + enable_pre_call_checks = _setting("enable_pre_call_checks", lambda r: r.settings().enable_pre_call_checks) + enable_health_check_routing = _setting( + "enable_health_check_routing", lambda r: r.settings().enable_health_check_routing ) - cooldown_time = _setting("cooldown_time", lambda r: r._core_cooldown_time) - allowed_fails = _setting("allowed_fails", lambda r: r._core_allowed_fails) - allowed_fails_policy = _setting("allowed_fails_policy", lambda r: r._core_allowed_fails_policy) - disable_cooldowns = _setting("disable_cooldowns", lambda r: r._core_disable_cooldowns) - enable_pre_call_checks = _setting("enable_pre_call_checks", lambda r: r._core_enable_pre_call_checks) - enable_health_check_routing = _setting("enable_health_check_routing", lambda r: r._core_enable_health_check_routing) health_check_ignore_transient_errors = _setting( - "health_check_ignore_transient_errors", lambda r: r._core_health_check_ignore_transient_errors + "health_check_ignore_transient_errors", lambda r: r.settings().health_check_ignore_transient_errors ) background_health_check_model_groups = _setting( - "background_health_check_model_groups", lambda r: r._core_background_health_check_model_groups + "background_health_check_model_groups", lambda r: r.settings().background_health_check_model_groups ) - enable_tag_filtering = _setting("enable_tag_filtering", lambda r: r._core_enable_tag_filtering) - tag_filtering_match_any = _setting("tag_filtering_match_any", lambda r: r._core_tag_filtering_match_any) - tag_routing_prefix = _setting("tag_routing_prefix", lambda r: r._core_tag_routing_prefix) - router_general_settings = _setting("router_general_settings", lambda r: r._core_router_general_settings) - provider_budget_config = _setting("provider_budget_config", lambda r: r._core_provider_budget_config) + enable_tag_filtering = _setting("enable_tag_filtering", lambda r: r.settings().enable_tag_filtering) + tag_filtering_match_any = _setting("tag_filtering_match_any", lambda r: r.settings().tag_filtering_match_any) + tag_routing_prefix = _setting("tag_routing_prefix", lambda r: r.settings().tag_routing_prefix) + router_general_settings = _setting("router_general_settings", lambda r: r.settings().router_general_settings) + provider_budget_config = _setting("provider_budget_config", lambda r: r.settings().provider_budget_config) provider_default_deployment_ids = _setting( - "provider_default_deployment_ids", lambda r: r._core_provider_default_deployment_ids + "provider_default_deployment_ids", lambda r: r.settings().provider_default_deployment_ids ) - fallback_access_check = _setting("fallback_access_check", lambda r: r._core_fallback_access_check) - fallback_budget_check = _setting("fallback_budget_check", lambda r: r._core_fallback_budget_check) + fallback_access_check = _setting("fallback_access_check", lambda r: r.settings().fallback_access_check) + fallback_budget_check = _setting("fallback_budget_check", lambda r: r.settings().fallback_budget_check) # These three are declared with the attribute type Router gave them rather than the descriptor's: # protocols in router_utils and the proxy match Router on them, and protocol matching does not # see through a descriptor. Reads and writes still go through _setting at runtime. fallbacks: list | list[dict[str, list[str] | list | None]] | None = _setting( # pyright: ignore[reportAssignmentType] # see above # mutable-ok: the old type - "fallbacks", lambda r: r._core_fallbacks + "fallbacks", lambda r: r.settings().fallbacks ) kubernetes_pod_discovery: "KubernetesPodDiscovery" = _setting( # pyright: ignore[reportAssignmentType] # see above - "kubernetes_pod_discovery", lambda r: r._core_kubernetes_pod_discovery + "kubernetes_pod_discovery", lambda r: r.settings().kubernetes_pod_discovery ) search_tools: "list[SearchToolTypedDict]" = _setting( # pyright: ignore[reportAssignmentType] # see above # mutable-ok: the old type - "search_tools", lambda r: r._core_search_tools + "search_tools", lambda r: r.settings().search_tools ) - guardrail_list = _setting("guardrail_list", lambda r: r._core_guardrail_list) - cache_responses = _setting("cache_responses", lambda r: r._core_cache_responses) + guardrail_list = _setting("guardrail_list", lambda r: r.settings().guardrail_list) + cache_responses = _setting("cache_responses", lambda r: r.settings().cache_responses) # -- registry: one deployment --------------------------------------------------------------- def get_deployment(self, model_id: str) -> "Deployment | None": _note("get_deployment", "registry().deployment(id)") - return _as_router(self)._core_get_deployment(model_id=model_id) + return _as_router(self).registry().deployment(model_id) def has_model_id(self, candidate_id: str) -> bool: _note("has_model_id", "registry().has_deployment(id)") - return _as_router(self)._core_has_model_id(candidate_id) + return _as_router(self).registry().has_deployment(candidate_id) def get_model_info(self, id: str) -> dict | None: # mutable-ok: the old return type, unchanged _note("get_model_info", "registry().deployment_dict(id)") - return _as_router(self)._core_get_model_info(id=id) + return _as_router(self).registry().deployment_dict(id) def get_deployment_model_info(self, model_id: str, model_name: str) -> "ModelMapInfo | None": _note("get_deployment_model_info", "registry().deployment_model_info(id, model_name)") - return _as_router(self)._core_get_deployment_model_info(model_id=model_id, model_name=model_name) + return _as_router(self).registry().deployment_model_info(model_id, model_name) @overload def get_router_model_info( @@ -434,62 +370,56 @@ class DeprecatedRouterMembers: id: str | None = None, ) -> "ModelMapInfo": _note("get_router_model_info", "registry().router_model_info(...)") - router: Final = _as_router(self) - # The core's overloads split (deployment, id) by which one is None; this forwards either form unchanged. - return router._core_get_router_model_info( # pyright: ignore[reportCallIssue] # see above - deployment=deployment, # pyright: ignore[reportArgumentType] # see above - received_model_name=received_model_name, - id=id, # pyright: ignore[reportArgumentType] # see above - ) + return _as_router(self).registry().router_model_info(deployment, received_model_name, id) def get_discovered_model_info(self, model_id: str | None) -> Mapping[str, int]: _note("get_discovered_model_info", "registry().discovered_model_info(id)") - return _as_router(self)._core_get_discovered_model_info(model_id) + return _as_router(self).registry().discovered_model_info(model_id) # -- registry: many deployments ------------------------------------------------------------- def get_model_list( self, model_name: str | None = None, team_id: str | None = None ) -> "list[DeploymentTypedDict] | None": # mutable-ok: the old return type, unchanged _note("get_model_list", "registry().deployments(model_name, team_id=...)") - return _as_router(self)._core_get_model_list(model_name=model_name, team_id=team_id) + return _as_router(self).registry().deployments(model_name, team_id=team_id) def get_all_deployments( self, model_name: str, model_alias: str | None = None, team_id: str | None = None ) -> "list[DeploymentTypedDict]": # mutable-ok: the old return type, unchanged _note("get_all_deployments", "registry().deployments(model_name, source='all')") - return _as_router(self)._core_get_all_deployments( - model_name=model_name, model_alias=model_alias, team_id=team_id + return ( + _as_router(self).registry().deployments(model_name, source="all", model_alias=model_alias, team_id=team_id) ) - _get_all_deployments = _private_alias("get_all_deployments", get_all_deployments) + _get_all_deployments = get_all_deployments def config_deployments(self) -> Iterator[Mapping[str, object]]: _note("config_deployments", "registry().deployments(source='config')") - return _as_router(self)._core_config_deployments() + return _as_router(self).registry().deployments(source="config") def deployments_named(self, model_name: str) -> "tuple[DeploymentTypedDict, ...] | None": _note("deployments_named", "registry().deployments(model_name, source='named')") - return _as_router(self)._core_deployments_named(model_name) + return _as_router(self).registry().deployments(model_name, source="named") def get_model_list_from_model_alias( self, model_name: str | None = None ) -> "list[DeploymentTypedDict]": # mutable-ok: the old return type, unchanged _note("get_model_list_from_model_alias", "registry().deployments(model_name, source='alias')") - return _as_router(self)._core_get_model_list_from_model_alias(model_name=model_name) + return _as_router(self).registry().deployments(model_name, source="alias") def get_model_list_from_routing_groups(self, model_name: str | None = None) -> "Sequence[DeploymentTypedDict]": _note("get_model_list_from_routing_groups", "registry().deployments(model_name, source='routing_groups')") - return _as_router(self)._core_get_model_list_from_routing_groups(model_name=model_name) + return _as_router(self).registry().deployments(model_name, source="routing_groups") def get_model_list_of_routed_group( self, model_group: str ) -> "list[DeploymentTypedDict]": # mutable-ok: the old return type, unchanged _note("get_model_list_of_routed_group", "registry().deployments(model_group, source='routed_group')") - return _as_router(self)._core_get_model_list_of_routed_group(model_group) + return _as_router(self).registry().deployments(model_group, source="routed_group") def wildcard_route_deployments(self, model: str, team_id: str | None) -> Sequence[Mapping[str, object]]: _note("wildcard_route_deployments", "registry().deployments(model, source='wildcard')") - return _as_router(self)._core_wildcard_route_deployments(model, team_id) + return _as_router(self).registry().deployments(model, source="wildcard", team_id=team_id) def get_routable_deployments( self, @@ -500,17 +430,21 @@ class DeprecatedRouterMembers: _note( "get_routable_deployments", "registry().deployments(model, source='routable', request=..., call_type=...)" ) - return _as_router(self)._core_get_routable_deployments(model, request_kwargs, call_type) + return ( + _as_router(self) + .registry() + .deployments(model, source="routable", request=request_kwargs, call_type=call_type) + ) def get_deployment_by_model_group_name(self, model_group_name: str) -> "Deployment | None": _note("get_deployment_by_model_group_name", "registry().deployments(name, source='first_in_group')") - return _as_router(self)._core_get_deployment_by_model_group_name(model_group_name=model_group_name) + return _as_router(self).registry().deployments(model_group_name, source="first_in_group") def deployments_for_request( self, model: str, request_kwargs: Mapping[str, object] ) -> "Sequence[DeploymentTypedDict]": _note("deployments_for_request", "registry().deployments(model, source='for_request', request=...)") - return _as_router(self)._core_deployments_for_request(model, request_kwargs) + return _as_router(self).registry().deployments(model, source="for_request", request=request_kwargs) def should_include_deployment( self, @@ -519,135 +453,137 @@ class DeprecatedRouterMembers: team_id: str | None = None, ) -> bool: _note("should_include_deployment", "registry().includes(deployment, model_name=..., team_id=...)") - return _as_router(self)._core_should_include_deployment(model_name=model_name, model=model, team_id=team_id) + return _as_router(self).registry().includes(model, model_name=model_name, team_id=team_id) @staticmethod def deployment_usable_by_team(model: "Mapping[str, object] | Deployment", team_id: str | None) -> bool: _note("deployment_usable_by_team", "registry().includes(deployment, team_id=...)") - from litellm.router_python_core import PythonCore + from litellm.router_api import deployment_usable_by_team as usable_by_team - return PythonCore._core_deployment_usable_by_team(model, team_id) + return usable_by_team(model, team_id) _deployment_usable_by_team = deployment_usable_by_team # -- registry: names ------------------------------------------------------------------------- def get_model_names(self, team_id: str | None = None) -> list[str]: # mutable-ok: the old return type, unchanged _note("get_model_names", "registry().names(team_id=...)") - return _as_router(self)._core_get_model_names(team_id=team_id) + return _as_router(self).registry().names(team_id=team_id) def get_model_ids( self, model_name: str | None = None, exclude_team_models: bool = False ) -> list[str]: # mutable-ok: the old return type, unchanged _note("get_model_ids", "registry().names(model_name, kind='ids')") - return _as_router(self)._core_get_model_ids(model_name=model_name, exclude_team_models=exclude_team_models) + return _as_router(self).registry().names(model_name, kind="ids", exclude_team_models=exclude_team_models) def resolved_litellm_models( self, model_name: str, team_id: str | None = None, *, use_base_model: bool = False ) -> tuple[str, ...]: _note("resolved_litellm_models", "registry().names(model_name, kind='litellm_models' or 'litellm_base_models')") - return _as_router(self)._core_resolved_litellm_models( - model_name, team_id=team_id, use_base_model=use_base_model + return ( + _as_router(self) + .registry() + .names(model_name, kind="litellm_base_models" if use_base_model else "litellm_models", team_id=team_id) ) def wildcard_routed_models(self, model: str, team_id: str | None) -> frozenset[str]: _note("wildcard_routed_models", "registry().names(model, kind='wildcard')") - return _as_router(self)._core_wildcard_routed_models(model, team_id) + return _as_router(self).registry().names(model, kind="wildcard", team_id=team_id) def get_candidate_model_ids_for_route(self, model: str, team_id: str | None = None) -> frozenset[str]: _note("get_candidate_model_ids_for_route", "registry().names(model, kind='route_candidates')") - return _as_router(self)._core_get_candidate_model_ids_for_route(model, team_id=team_id) + return _as_router(self).registry().names(model, kind="route_candidates", team_id=team_id) # -- registry: resolving names --------------------------------------------------------------- def routable_model_group(self, model_name: str) -> str: _note("routable_model_group", "registry().resolve(name)") - return _as_router(self)._core_routable_model_group(model_name) + return _as_router(self).registry().resolve(model_name) def get_model_from_alias(self, model: str) -> str | None: _note("get_model_from_alias", "registry().resolve(name, kind='alias')") - return _as_router(self)._core_get_model_from_alias(model=model) + return _as_router(self).registry().resolve(model, kind="alias") - _get_model_from_alias = _private_alias("get_model_from_alias", get_model_from_alias) + _get_model_from_alias = get_model_from_alias def resolve_served_model_name(self, model_name: str, served_model: str) -> str: _note("resolve_served_model_name", "registry().resolve(name, kind='served', served_model=...)") - return _as_router(self)._core_resolve_served_model_name(model_name, served_model) + return _as_router(self).registry().resolve(model_name, kind="served", served_model=served_model) def resolve_model_name_from_model_id( self, model_id: str | None, custom_llm_provider: str | None = None ) -> str | None: _note("resolve_model_name_from_model_id", "registry().resolve(id, kind='from_id')") - return _as_router(self)._core_resolve_model_name_from_model_id( - model_id, custom_llm_provider=custom_llm_provider - ) + return _as_router(self).registry().resolve(model_id, kind="from_id", custom_llm_provider=custom_llm_provider) def map_team_model( self, team_model_name: str | None, team_id: str, custom_llm_provider: str | None = None ) -> str | None: _note("map_team_model", "registry().resolve(name, kind='team_model', team_id=...)") - return _as_router(self)._core_map_team_model(team_model_name, team_id, custom_llm_provider=custom_llm_provider) + return ( + _as_router(self) + .registry() + .resolve(team_model_name, kind="team_model", team_id=team_id, custom_llm_provider=custom_llm_provider) + ) def is_recognized_model(self, model: str) -> bool: _note("is_recognized_model", "registry().resolve(name, kind='recognized')") - return _as_router(self)._core_is_recognized_model(model) + return _as_router(self).registry().resolve(model, kind="recognized") def get_routable_upstream_model(self, model_name: str, team_id: str | None = None) -> str | None: _note("get_routable_upstream_model", "registry().resolve(name, kind='upstream')") - return _as_router(self)._core_get_routable_upstream_model(model_name, team_id=team_id) + return _as_router(self).registry().resolve(model_name, kind="upstream", team_id=team_id) def get_search_tool_support(self, model_name: str, team_id: str | None = None) -> tuple[bool | None, ...]: _note("get_search_tool_support", "registry().resolve(name, kind='search_tool_support')") - return _as_router(self)._core_get_search_tool_support(model_name, team_id=team_id) + return _as_router(self).registry().resolve(model_name, kind="search_tool_support", team_id=team_id) def get_model_group(self, id: str) -> list | None: # mutable-ok: the old return type, unchanged _note("get_model_group", "registry().resolve(id, kind='group_of_id')") - return _as_router(self)._core_get_model_group(id=id) + return _as_router(self).registry().resolve(id, kind="group_of_id") def is_model_access_group_for_wildcard_route(self, model_access_group: str) -> bool: _note("is_model_access_group_for_wildcard_route", "registry().resolve(group, kind='wildcard_access_group')") - return _as_router(self)._core_is_model_access_group_for_wildcard_route(model_access_group=model_access_group) + return _as_router(self).registry().resolve(model_access_group, kind="wildcard_access_group") - _is_model_access_group_for_wildcard_route = _private_alias( - "is_model_access_group_for_wildcard_route", is_model_access_group_for_wildcard_route - ) + _is_model_access_group_for_wildcard_route = is_model_access_group_for_wildcard_route def team_model_has_alternatives(self, deployment_id: str) -> bool: _note("team_model_has_alternatives", "registry().resolve(id, kind='team_model_alternatives')") - return _as_router(self)._core_team_model_has_alternatives(deployment_id) + return _as_router(self).registry().resolve(deployment_id, kind="team_model_alternatives") def routing_group_has_alternatives(self, model_group: str | None) -> bool: _note("routing_group_has_alternatives", "registry().resolve(group, kind='routing_group_alternatives')") - return _as_router(self)._core_routing_group_has_alternatives(model_group) + return _as_router(self).registry().resolve(model_group, kind="routing_group_alternatives") # -- registry: groups, access groups, credentials, plugins ----------------------------------- def get_model_group_info( self, model_group: str, *, include_hidden: bool = False, team_id: str | None = None ) -> "ModelGroupInfo | None": _note("get_model_group_info", "registry().group(name)") - return _as_router(self)._core_get_model_group_info(model_group, include_hidden=include_hidden, team_id=team_id) + return _as_router(self).registry().group(model_group, include_hidden=include_hidden, team_id=team_id) def cached_model_group_info(self, model_group: str) -> "ModelGroupInfo | None": _note("cached_model_group_info", "registry().group(name, field='cached_info')") - return _as_router(self)._core_cached_model_group_info(model_group) + return _as_router(self).registry().group(model_group, field="cached_info") def get_configured_display_name(self, model_name: str) -> str | None: _note("get_configured_display_name", "registry().group(name, field='display_name')") - return _as_router(self)._core_get_configured_display_name(model_name) + return _as_router(self).registry().group(model_name, field="display_name") def get_configured_mode(self, model_name: str) -> str | None: _note("get_configured_mode", "registry().group(name, field='mode')") - return _as_router(self)._core_get_configured_mode(model_name) + return _as_router(self).registry().group(model_name, field="mode") def get_configured_service_tiers(self, model_name: str, team_id: str | None = None) -> tuple[object, ...]: _note("get_configured_service_tiers", "registry().group(name, field='service_tiers')") - return _as_router(self)._core_get_configured_service_tiers(model_name, team_id=team_id) + return _as_router(self).registry().group(model_name, field="service_tiers", team_id=team_id) def get_configured_token_limits(self, model_name: str) -> tuple[int | None, int | None]: _note("get_configured_token_limits", "registry().group(name, field='token_limits')") - return _as_router(self)._core_get_configured_token_limits(model_name) + return _as_router(self).registry().group(model_name, field="token_limits") def get_model_listing_info(self, model_name: str) -> "DeploymentModelListingInfo | None": _note("get_model_listing_info", "registry().group(name, field='listing_info')") - return _as_router(self)._core_get_model_listing_info(model_name) + return _as_router(self).registry().group(model_name, field="listing_info") def get_model_access_groups( self, @@ -656,31 +592,31 @@ class DeprecatedRouterMembers: team_id: str | None = None, ) -> dict[str, list[str]]: # mutable-ok: the old return type, unchanged _note("get_model_access_groups", "registry().access_groups(...)") - return _as_router(self)._core_get_model_access_groups( - model_name=model_name, model_access_group=model_access_group, team_id=team_id + return ( + _as_router(self) + .registry() + .access_groups(model_name=model_name, model_access_group=model_access_group, team_id=team_id) ) def get_deployment_credentials(self, model_id: str) -> dict | None: # mutable-ok: the old return type, unchanged _note("get_deployment_credentials", "registry().credentials(id)") - return _as_router(self)._core_get_deployment_credentials(model_id=model_id) + return _as_router(self).registry().credentials(model_id) def get_deployment_credentials_with_provider( self, model_id: str, team_id: str | None = None ) -> "dict[str, Any] | None": # pyright: ignore[reportExplicitAny] # the old return type, unchanged # mutable-ok: same _note("get_deployment_credentials_with_provider", "registry().credentials(id, form='with_provider')") - return _as_router(self)._core_get_deployment_credentials_with_provider(model_id=model_id, team_id=team_id) + return _as_router(self).registry().credentials(model_id, form="with_provider", team_id=team_id) def get_credential_deployment(self, model_id: str, team_id: str | None = None) -> "Deployment | None": _note("get_credential_deployment", "registry().credentials(id, form='deployment')") - return _as_router(self)._core_get_credential_deployment(model_id=model_id, team_id=team_id) + return _as_router(self).registry().credentials(model_id, form="deployment", team_id=team_id) def get_router_deployment_budget_limiter(self) -> "RouterBudgetLimiting | None": _note("get_router_deployment_budget_limiter", "registry().plugins('budget_limiter')") - return _as_router(self)._core_get_router_deployment_budget_limiter() + return _as_router(self).registry().plugins("budget_limiter") - _get_router_deployment_budget_limiter = _private_alias( - "get_router_deployment_budget_limiter", get_router_deployment_budget_limiter - ) + _get_router_deployment_budget_limiter = get_router_deployment_budget_limiter @staticmethod def generate_model_id(model_group: str, litellm_params: dict) -> str: # mutable-ok: the old parameter type @@ -692,15 +628,15 @@ class DeprecatedRouterMembers: # -- settings: methods ----------------------------------------------------------------------- def get_settings(self) -> dict[str, object]: # mutable-ok: Router._core_get_settings's type, a fresh dict _note("get_settings", "settings().as_dict()") - return _as_router(self)._core_get_settings() + return dict(_as_router(self).settings().as_dict()) def get_allowed_fails_from_policy(self, exception: Exception) -> int | None: _note("get_allowed_fails_from_policy", "settings().allowed_fails_for(exception)") - return _as_router(self)._core_get_allowed_fails_from_policy(exception=exception) + return _as_router(self).settings().allowed_fails_for(exception) def get_available_guardrail(self, guardrail_name: str) -> "GuardrailTypedDict": _note("get_available_guardrail", "settings().guardrail(name)") - return _as_router(self)._core_get_available_guardrail(guardrail_name) + return _as_router(self).settings().guardrail(guardrail_name) # -- health ---------------------------------------------------------------------------------- @staticmethod @@ -723,31 +659,31 @@ class DeprecatedRouterMembers: def get_fully_blocked_model_names(self) -> set[str]: # mutable-ok: the old return type, unchanged _note("get_fully_blocked_model_names", "health().fully_blocked_model_names()") - return _as_router(self)._core_get_fully_blocked_model_names() + return _as_router(self).health().fully_blocked_model_names() async def async_get_fully_unhealthy_model_names(self) -> set[str]: # mutable-ok: the old return type, unchanged _note("async_get_fully_unhealthy_model_names", "health().fully_unhealthy_model_names()") - return await _as_router(self)._core_async_get_fully_unhealthy_model_names() + return await _as_router(self).health().fully_unhealthy_model_names() async def get_model_group_usage(self, model_group: str) -> tuple[int | None, int | None]: _note("get_model_group_usage", "health().usage(group)") - return await _as_router(self)._core_get_model_group_usage(model_group=model_group) + return await _as_router(self).health().usage(model_group) async def get_remaining_model_group_usage( self, model_group: str ) -> dict[str, int]: # mutable-ok: the old return type, unchanged _note("get_remaining_model_group_usage", "health().usage(group, kind='remaining')") - return await _as_router(self)._core_get_remaining_model_group_usage(model_group=model_group) + return await _as_router(self).health().usage(model_group, kind="remaining") def team_ptu_ceiling(self, team_id: str, model_group: str) -> "PTUTeamCeiling | None": _note("team_ptu_ceiling", "health().team_ptu_ceiling(team_id, group)") - return _as_router(self)._core_team_ptu_ceiling(team_id, model_group) + return _as_router(self).health().team_ptu_ceiling(team_id, model_group) def split_request_priority( self, model: str, kwargs: Mapping[str, object] ) -> tuple[int | None, Mapping[str, object]]: _note("split_request_priority", "health().split_priority(model, kwargs)") - return _as_router(self)._core_split_request_priority(model, kwargs) + return _as_router(self).health().split_priority(model, kwargs) # -- selection ------------------------------------------------------------------------------- async def async_get_healthy_deployments( @@ -761,9 +697,10 @@ class DeprecatedRouterMembers: health_check_probe: bool = False, ) -> list[dict] | dict: # mutable-ok: the old return type, unchanged _note("async_get_healthy_deployments", "select(model, request_kwargs, kind='healthy')") - return await _as_router(self)._core_async_get_healthy_deployments( - model=model, - request_kwargs=request_kwargs, + return await _as_router(self).select( + model, + request_kwargs, + kind="healthy", messages=messages, input=input, specific_deployment=specific_deployment, @@ -771,7 +708,7 @@ class DeprecatedRouterMembers: health_check_probe=health_check_probe, ) - def get_available_deployment_for_pass_through( # noqa: ANN201 # the old member had no return annotation; its inferred type is kept by inferring it from the core + def get_available_deployment_for_pass_through( # noqa: ANN201 # the old member had no return annotation; its inferred type is kept by inferring it from select self, model: str, messages: list[dict[str, str]] | None = None, # mutable-ok: the old parameter type, unchanged @@ -780,12 +717,15 @@ class DeprecatedRouterMembers: request_kwargs: dict | None = None, # mutable-ok: the old parameter type, unchanged ): _note("get_available_deployment_for_pass_through", "select(model, kind='pass_through')") - return _as_router(self)._core_get_available_deployment_for_pass_through( - model=model, - messages=messages, - input=input, - specific_deployment=specific_deployment, - request_kwargs=request_kwargs, + return _run_ready( + _as_router(self).select( + model, + request_kwargs, + kind="pass_through", + messages=messages, + input=input, + specific_deployment=specific_deployment, + ) ) # -- writes ---------------------------------------------------------------------------------- @@ -805,7 +745,7 @@ class DeprecatedRouterMembers: _note("update_redis_cache", "apply(ChangeSet(redis_cache=cache))") _as_router(self).apply(ChangeSet(redis_cache=cache)) - _update_redis_cache = _private_alias("update_redis_cache", update_redis_cache) + _update_redis_cache = update_redis_cache async def arefresh_model_info(self, *, client: "AsyncHTTPHandler | None" = None) -> None: _note("arefresh_model_info", "aapply(ChangeSet(refresh_model_info=True))") @@ -822,8 +762,15 @@ class DeprecatedRouterMembers: def bind_to_request_logging(self, logging_obj: "LiteLLMLogging") -> None: _note("bind_to_request_logging", "observe(logging_obj)") - # Straight to the core, as before: observe() skips binding when logging_obj is None, the old member did not. - _as_router(self)._core_bind_to_request_logging(logging_obj) + router: Final = _as_router(self) + if logging_obj is None: # pyright: ignore[reportUnnecessaryComparison] # callers pass None despite the type + # observe() skips binding when logging_obj is None; the old member bound it, and failed on None. + raise AttributeError( + "'NoneType' object has no attribute 'add_dynamic_function_callbacks'", + name="add_dynamic_function_callbacks", + obj=None, + ) + router.observe(logging_obj) def discard(self) -> None: _note("discard", "retire()") @@ -836,7 +783,7 @@ class DeprecatedRouterMembers: request_kwargs: dict[str, object] | None = None, # mutable-ok: the old parameter type, unchanged ) -> None: _note("arm_routing_read_prefetch", "RoutingContext.arm_read_prefetch") - _as_router(self)._core_arm_routing_read_prefetch(model, request_kwargs) + RoutingContext(_as_router(self)).arm_read_prefetch(model, request_kwargs) async def async_pre_routing_hook( self, @@ -847,7 +794,7 @@ class DeprecatedRouterMembers: specific_deployment: bool | None = False, ) -> "PreRoutingHookResponse | None": _note("async_pre_routing_hook", "RoutingContext.pre_routing_hook") - return await _as_router(self)._core_async_pre_routing_hook( + return await RoutingContext(_as_router(self)).pre_routing_hook( model=model, request_kwargs=request_kwargs, messages=messages, @@ -857,7 +804,7 @@ class DeprecatedRouterMembers: def log_retry(self, kwargs: dict, e: Exception) -> dict: # mutable-ok: the old signature, unchanged _note("log_retry", "RoutingContext.log_retry") - return _as_router(self)._core_log_retry(kwargs=kwargs, e=e) + return RoutingContext(_as_router(self)).log_retry(kwargs=kwargs, e=e) # Router-family functions and classes that took the Router (``litellm_router_instance=``, diff --git a/litellm/router_python_core.py b/litellm/router_python_core.py index 4b5ecc82581..8ca2afb9a99 100644 --- a/litellm/router_python_core.py +++ b/litellm/router_python_core.py @@ -127,6 +127,7 @@ from litellm.litellm_core_utils.clientside_credential_handler import ( ) from litellm.litellm_core_utils.completion_timeout import CompletionTimeout from litellm.litellm_core_utils.core_helpers import ( + as_str_mapping, coerce_token_limit, get_litellm_metadata_from_kwargs, get_metadata_variable_name_from_kwargs, @@ -322,6 +323,7 @@ from litellm.router_utils.cooldown_handlers import ( mark_deployment_selection_failure, set_cooldown_deployments, ) +from litellm.router_utils.deployment_numbers import refuse_unservable_numbers from litellm.router_utils.embedding_batch_size_check import ( async_log_rejection_before_provider_call, clear_rejected_before_provider_call, @@ -366,6 +368,11 @@ from litellm.router_utils.fallback_event_handlers import ( routed_deployment_id, run_async_fallback, ) +from litellm.router_utils.global_callback_holders import ( + live_check_owners, + routers_holding_global_callbacks, + routers_holding_global_callbacks_lock, +) from litellm.router_utils.handle_error import ( async_raise_no_deployment_exception, send_llm_exception_alert, @@ -1233,8 +1240,9 @@ class RoutingArgs(enum.Enum): # as the per-request one built from a caller-supplied user_config, drops out on its # own rather than leaving entries behind that nothing can withdraw. _live_routers: Final["weakref.WeakSet[PythonCore]"] = weakref.WeakSet() -_routers_holding_global_callbacks: Final["weakref.WeakSet[PythonCore]"] = weakref.WeakSet() -_routers_holding_global_callbacks_lock: Final = threading.Lock() +# Shared with the other router cores of the process (litellm/router_utils/global_callback_holders.py). +_routers_holding_global_callbacks: Final["weakref.WeakSet[PythonCore]"] = routers_holding_global_callbacks # pyright: ignore[reportAssignmentType] # every member serves the holder members +_routers_holding_global_callbacks_lock: Final = routers_holding_global_callbacks_lock def _replay_live_router_model_cost() -> None: @@ -1878,6 +1886,7 @@ class PythonCore(DeprecatedRouterMembers): # cost-map entries from the list it is serving at that moment, whether # those deployments came from model_list or were added later from the DB. _live_routers.add(self) + live_check_owners.add(self) if allowed_fails is not None: self._core_allowed_fails = allowed_fails @@ -2103,6 +2112,7 @@ class PythonCore(DeprecatedRouterMembers): # Stop contributing to cost-map rebuilds straight away rather than waiting # for this router to be collected. _live_routers.discard(self) + live_check_owners.discard(self) self.retire_routing_state() self._release_global_callbacks() @@ -3414,7 +3424,7 @@ class PythonCore(DeprecatedRouterMembers): self.optional_callbacks = [cb for cb in self.optional_callbacks if type(cb) is not callback_cls] if any( router is not self and any(type(cb) is callback_cls for cb in (router.optional_callbacks or [])) - for router in tuple(_live_routers) + for router in (*_live_routers, *(owner for owner in live_check_owners if not isinstance(owner, PythonCore))) ): return for cb in tuple(self._host.callback_list("callbacks")): @@ -7150,6 +7160,8 @@ class PythonCore(DeprecatedRouterMembers): or self._stream_can_fall_back(model, initial_kwargs) ) buffered_lifecycle_chunks: tuple[bytes, ...] = () # rebind-ok: flushed once committed or on decline + # An error this loop raises off a frame, not the provider's. + raised_here: BaseException | None = None # rebind-ok: set when the loop raises off a frame try: async for chunk in source_iterator: if _anthropic_stream_forwards_ping_live(chunk, has_generated_content): @@ -7180,25 +7192,27 @@ class PythonCore(DeprecatedRouterMembers): and self._refusal_fallback_available(model, initial_kwargs) ): refusal_error = safeguard_refusal_error(model=model, stop_details=refusal_stop_details) - raise MidStreamFallbackError( + raised_here = MidStreamFallbackError( message=refusal_error.message, model=model, llm_provider="anthropic", original_exception=refusal_error, is_pre_first_chunk=True, ) + raise raised_here if not has_generated_content and error_event is None: buffered_lifecycle_chunks = (*buffered_lifecycle_chunks, chunk) continue if recoverable_frame_error is not None: assert error_event is not None - raise MidStreamFallbackError( + raised_here = MidStreamFallbackError( message=error_event[1], model=model, llm_provider="anthropic", original_exception=recoverable_frame_error, is_pre_first_chunk=True, ) + raise raised_here for buffered_chunk in buffered_lifecycle_chunks: yield buffered_chunk buffered_lifecycle_chunks = () @@ -7206,6 +7220,8 @@ class PythonCore(DeprecatedRouterMembers): for buffered_chunk in buffered_lifecycle_chunks: yield buffered_chunk except Exception as stream_error: # noqa: BLE001 # any raised provider error must reach the fallback gate + if stream_error is not raised_here: + await self._areport_raised_messages_stream_failure(stream_error, initial_kwargs, source_iterator) async for item in self._aanthropic_messages_recover_stream_error( stream_error, has_generated_content, @@ -7225,6 +7241,52 @@ class PythonCore(DeprecatedRouterMembers): wrapper: Final = FallbackAwareAnthropicMessagesStream(stream_with_fallbacks(), source_iterator) return wrapper + async def _areport_raised_messages_stream_failure( + self, stream_error: Exception, attempt_kwargs: Mapping[str, object], source: object + ) -> None: + """A /v1/messages stream whose provider iterator raised, counted against the deployment + that served it and cooled down by the rules, the same as a non-stream call's failure, so a + retry picks another deployment. A stream that logs itself (litellm's own stream logging + reaches the router's failure callback), or a failure the request's logging object already + logged, is not counted twice.""" + from litellm.exceptions import MidStreamFallbackError + from litellm.litellm_core_utils.litellm_logging import Logging + + if isinstance(getattr(source, "logging_obj", None), Logging): + return + logging_obj: Final = attempt_kwargs.get("litellm_logging_obj") + if isinstance(logging_obj, Logging) and logging_obj.model_call_details.get("has_logged_async_failure") is True: + return + bucket: Final = as_str_mapping(attempt_kwargs.get(get_metadata_variable_name_from_kwargs(attempt_kwargs))) or {} + model_info: Final = as_str_mapping(bucket.get("model_info")) or {} + deployment_id: Final = model_info.get("id") + if deployment_id is None: + return + deployment: Final = self.get_deployment(model_id=str(deployment_id)) + deployment_params: Final = deployment.litellm_params.model_dump(exclude_none=True) if deployment else {} + provider_error: Final = ( + stream_error.original_exception + if isinstance(stream_error, MidStreamFallbackError) and stream_error.original_exception is not None + else stream_error + ) + now: Final = datetime.now() # noqa: DTZ005 # naive, as the logging pipeline's times + logged: Final = { + "call_type": "anthropic_messages", + "exception": provider_error, + "litellm_params": { + **{ + name: deployment_params[name] for name in ("api_base", "cooldown_time") if name in deployment_params + }, + "metadata": dict(bucket), + "model_info": dict(model_info), + }, + } + try: + _ = self.deployment_callback_on_failure(logged, None, now, now) # pyright: ignore[reportUnknownMemberType] # an untyped callback + await self.async_deployment_callback_on_failure(logged, None, now, now) # pyright: ignore[reportUnknownMemberType] # see above + except Exception as error: # noqa: BLE001 # a recording never fails the stream + verbose_router_logger.warning("a stream's failure was not recorded: %s", error) + async def _aanthropic_messages_recover_stream_error( self, stream_error: Exception, @@ -11235,6 +11297,7 @@ class PythonCore(DeprecatedRouterMembers): model_info=_model_info, ) ) + refuse_unservable_numbers(deployment.to_json(exclude_none=True, warnings=False)) # pyright: ignore[reportUnknownMemberType] # pydantic's to_json PythonCore._register_deployment_pricing(deployment) ## Check if LLM Deployment is allowed for this deployment @@ -11738,6 +11801,16 @@ class PythonCore(DeprecatedRouterMembers): return model_info_is_active_for_environment(model_info=deployment.model_info) def set_model_list(self, model_list: list): + """Replace every deployment. All or nothing: when an entry raises, the deployments served + before the call are put back and the error is raised.""" + previous: Final = copy.deepcopy(list(self.__dict__.get("_core_model_list") or ())) + try: + self._replace_model_list(model_list) # pyright: ignore[reportUnknownMemberType] # the legacy untyped model list + except Exception: + self._replace_model_list(previous) # pyright: ignore[reportUnknownMemberType] # see above + raise + + def _replace_model_list(self, model_list: list) -> None: # pyright: ignore[reportUnknownParameterType, reportMissingTypeArgument] # the legacy untyped model list original_model_list: Final = copy.deepcopy(model_list) self._discovered_model_info_cache.flush_cache() with self._deployment_index_lock: @@ -12031,6 +12104,7 @@ class PythonCore(DeprecatedRouterMembers): # add to model list _deployment: Final = served.to_json(exclude_none=True) + refuse_unservable_numbers(_deployment) # initialize client self._add_deployment(deployment=served) @@ -12174,6 +12248,9 @@ class PythonCore(DeprecatedRouterMembers): _deployment_model_id: Final = deployment.model_info.id or "" _deployment_on_router: Final[Deployment | None] = self._core_get_deployment(model_id=_deployment_model_id) served: Final = without_model_capability_flags(deployment) + # `_create_deployment` lists the model name before it can still fail; a dropped deployment + # takes its name back off the list. + names_before: Final = len(self._core_deployment_names) try: if _deployment_on_router is not None: # deployment with this model_id exists on the router @@ -12223,6 +12300,7 @@ class PythonCore(DeprecatedRouterMembers): deployment.model_info.id, e, ) + del self._core_deployment_names[names_before:] self._restore_deployment_after_failed_upsert( previous_deployment=_deployment_on_router, model_id=_deployment_model_id ) @@ -14713,9 +14791,7 @@ class PythonCore(DeprecatedRouterMembers): """ if team_id is not None and self._team_deployments_named(team_id, model_group): return self._core_get_all_deployments(model_name=model_group, team_id=team_id) - own: Final = cast( # cast-ok: the private alias returns what get_all_deployments returns - list[DeploymentTypedDict], self._get_all_deployments(model_name=model_group) - ) + own: Final = self._get_all_deployments(model_name=model_group) named: Final = [*own, *self._core_get_model_list_from_routing_groups(model_name=model_group)] return named or self._get_wildcard_deployments(model_name=model_group) @@ -15923,10 +15999,15 @@ class PythonCore(DeprecatedRouterMembers): # IF TEAM ID SPECIFIED ON MODEL, AND REQUEST CONTAINS USER_API_KEY_TEAM_ID, FILTER OUT MODELS THAT ARE NOT IN THE TEAM ## THIS PREVENTS WRITING FILES OF OTHER TEAMS TO MODELS THAT ARE TEAM-ONLY MODELS + by_id: Final = healthy_deployments if isinstance(healthy_deployments, dict) else None # pyright: ignore[reportUnknownVariableType] # the legacy deployment dict healthy_deployments = filter_team_based_models( healthy_deployments=healthy_deployments, request_kwargs=request_kwargs, ) + if by_id is not None and not healthy_deployments: + # Another team's deployment asked for by id: refused under its group's name. + group_name: Final = str(by_id.get("model_name")) # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType] # see above + model = group_name # rebind-ok: the error names the group if verbose_router_logger.isEnabledFor(logging.DEBUG): verbose_router_logger.debug("healthy_deployments after team filter: %s", healthy_deployments) @@ -17236,6 +17317,16 @@ class PythonCore(DeprecatedRouterMembers): model, ) ) + # As the async pick: a team-owned deployment serves only its own team's callers, also when + # asked for by id (then refused under its group's name). + by_id: Final = healthy_deployments if isinstance(healthy_deployments, dict) else None # pyright: ignore[reportUnknownVariableType] # the legacy deployment dict + healthy_deployments = filter_team_based_models( # rebind-ok: the team filter narrows # pyright: ignore[reportUnknownVariableType] # legacy selector type + healthy_deployments=healthy_deployments, # pyright: ignore[reportArgumentType] # legacy selector type + request_kwargs=request_kwargs, + ) + if by_id is not None and not healthy_deployments: + group_name: Final = str(by_id.get("model_name")) # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType] # see above + model = group_name # rebind-ok: the error names the group strategy, strategy_selector = self._get_routing_context(model, request_kwargs) routing_kwargs: Final = _ROUTING_KWARGS_ADAPTER.validate_python(request_kwargs) pick_attributes: Final = _deployment_pick_attributes(model, routing_kwargs) diff --git a/litellm/router_python_host.py b/litellm/router_python_host.py index 8156740144e..8ab2c06b82f 100644 --- a/litellm/router_python_host.py +++ b/litellm/router_python_host.py @@ -7,7 +7,10 @@ come from ``dynamic()`` (``litellm/router_config.py``). A Rust core gets the sam """ import asyncio -from collections.abc import Awaitable, Callable, Sequence +import importlib +import time +import types +from collections.abc import Awaitable, Callable, Mapping, Sequence from typing import ( TYPE_CHECKING, Final, @@ -19,12 +22,13 @@ from typing import ( from typing_extensions import overload import litellm -from litellm.router_config import RouterDynamicConfig, live_dynamic_config +from litellm.router_config import RouterConfig, RouterDynamicConfig, live_dynamic_config if TYPE_CHECKING: from pydantic import BaseModel - from litellm.caching.caching import Cache + from litellm.caching.caching import Cache, RedisCache + from litellm.caching.dual_cache import DualCache from litellm.litellm_core_utils.logging_callback_manager import LoggingCallbackManager from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper from litellm.text_completions.streaming_iterator import TextCompletionStreamWrapper @@ -97,7 +101,7 @@ class PythonHost: def invoke(self, op: str) -> Endpoint: ... def invoke(self, op: str) -> Endpoint: """The litellm function named ``op`` (``acompletion``, ``get_model_info``, ...), looked up at call time.""" - endpoint: Final[object] = getattr(litellm, op, None) + endpoint: Final[object] = getattr(litellm, op, None) if "." not in op else _dotted(op) if endpoint is None: raise AttributeError(f"litellm has no endpoint {op!r}") if not callable(endpoint): @@ -107,6 +111,10 @@ class PythonHost: def sleep(self, seconds: float) -> Awaitable[None]: return asyncio.sleep(seconds) + def sleep_sync(self, seconds: float) -> None: + """The backoff of a sync endpoint's retry.""" + time.sleep(seconds) + def dynamic(self) -> RouterDynamicConfig: """The process settings, read live on each access (``current_dynamic_config()`` takes a snapshot).""" return cast(RouterDynamicConfig, live_dynamic_config) # cast-ok: the live view reads the same fields @@ -116,7 +124,7 @@ class PythonHost: return litellm.logging_callback_manager def callback_list(self, name: CallbackList) -> list[object]: # mutable-ok: the process callback list itself - value: Final = getattr(litellm, name) + value: Final[object] = getattr(litellm, name) # pyright: ignore[reportAny] # an untyped litellm global return cast(list[object], value) # cast-ok: the process callback list is an untyped litellm global def set_callback_list(self, name: CallbackList, value: list[object]) -> None: # mutable-ok: assigned as given @@ -131,5 +139,78 @@ class PythonHost: def suppress_debug_info(self) -> None: litellm.suppress_debug_info = True + def routing_cache(self, config: RouterConfig) -> "DualCache": + """The store a router keeps its routing state in (cooldowns, fail counts, usage): in memory, + plus Redis when the router's arguments name one. The host builds and owns it; a Rust core + reads and writes it through cache steps.""" + from litellm.caching.caching import DualCache, InMemoryCache + + return DualCache(redis_cache=redis_cache_for(config), in_memory_cache=InMemoryCache()) + + +def redis_cache_for(config: RouterConfig) -> "RedisCache | None": + """The Redis cache the router's ``redis_*`` and ``cache_kwargs`` arguments describe, or None.""" + import json + + from litellm.caching.caching import RedisCache, RedisClusterCache + from litellm.utils import get_secret + + url: Final = config.get("redis_url") + host: Final = config.get("redis_host") + port: Final = config.get("redis_port") + if url is None and (host is None or port is None): + return None + settings: Final[dict[str, object]] = {} # mutable-ok: built here, then passed as keywords + if url is not None: + settings["url"] = url + if host is not None: + settings["host"] = host + if port is not None: + settings["port"] = str(port) + if config.get("redis_password") is not None: + settings["password"] = config.get("redis_password") + if config.get("redis_db") is not None: + settings["db"] = str(config.get("redis_db")) + extra: Final = config.get("cache_kwargs") + if isinstance(extra, dict): + settings.update({str(key): value for key, value in extra.items()}) # pyright: ignore[reportUnknownVariableType, reportUnknownArgumentType] # the cache_kwargs argument is untyped + nodes: object = settings.get("startup_nodes") # rebind-ok: the env cluster nodes when the arguments name none + if not nodes: + env_nodes: Final[object] = get_secret("REDIS_CLUSTER_NODES") # pyright: ignore[reportUnknownVariableType] # get_secret is untyped + if isinstance(env_nodes, str): + parsed: Final[object] = json.loads(env_nodes) # pyright: ignore[reportAny] # plain JSON + nodes = parsed # rebind-ok: see above + if nodes: + return RedisClusterCache(**{**settings, "startup_nodes": nodes}) # pyright: ignore[reportArgumentType] # RedisCache's keyword arguments + return RedisCache(**settings) # pyright: ignore[reportArgumentType] # RedisCache's keyword arguments + python_host: Final = PythonHost() + + +def _dotted(op: str) -> object: + """``op`` named by its path under ``litellm`` (``vector_stores.main.asearch``), read at call time. + + Each segment is a module attribute, a submodule, or a key of a mapping such as a module's + table of generated endpoints. + """ + node: object = litellm # rebind-ok: walks the path one segment at a time + path: str = "litellm" # rebind-ok: the module path walked so far + for segment in op.split("."): + path = f"{path}.{segment}" + if isinstance(node, Mapping): + node = cast(Mapping[str, object], node).get(segment) # cast-ok: endpoint tables are keyed by name + elif isinstance(node, types.ModuleType) and (submodule := _submodule(path)) is not None: + node = submodule + else: + node = cast(object, getattr(node, segment, None)) # cast-ok: any attribute + if node is None: + return None + return node + + +def _submodule(path: str) -> types.ModuleType | None: + try: + return importlib.import_module(path) + except ModuleNotFoundError: + return None diff --git a/litellm/router_rust/__init__.py b/litellm/router_rust/__init__.py new file mode 100644 index 00000000000..42cfc56aca4 --- /dev/null +++ b/litellm/router_rust/__init__.py @@ -0,0 +1,13 @@ +""" +The Rust router core's Python side: ``RustCore`` (``core``) and the host-side modules router-core +(litellm-rust/crates/router-core) reaches through its steps. + +- ``steps``: answers to host steps that need live Python objects (strategies, pre-routing, + pre-call facts, plugins, callbacks, caller access) +- ``invoke``: the host side of ``Host::invoke``, one attempt against one deployment (request + shaping, endpoints, streams, slots, retry and fallback metadata, prompt management, websockets) +- ``surface``: the Python ``Router`` API kept over router-core's state (registry views, model + info, settings, health, retries, fallbacks, alerting); a Rust host needs none of it + +Nothing is imported here, so importing one module does not import the core. +""" diff --git a/litellm/router_rust/core.py b/litellm/router_rust/core.py new file mode 100644 index 00000000000..395b580d9b2 --- /dev/null +++ b/litellm/router_rust/core.py @@ -0,0 +1,4943 @@ +""" +``RustCore``: the router core served by litellm-rust/crates/router-core through ``_native``. + +The Rust core owns the registry, selection, retries, fallbacks and cooldowns. It reaches outside +itself only through its host, and the host never calls back into a Python router: + +- every attempt runs through the host: ``router_attempt`` builds the request for the one + deployment Rust picked and ``PythonHost.invoke(op)`` calls ``litellm.`` with it +- the core records every outcome (cooldowns, allowed-fails counters, usage) itself, in the routing + cache the host builds and owns (``PythonHost.routing_cache``), as far as the outcome callbacks the + Python router would have running for the request reach (``OptionalChecks.outcome_callbacks``) +- sleeps, fallback access and budget checks, and cache reads and writes are host steps + +``RustCore`` does not build, wrap or read a ``PythonCore``. Settings are resolved from the +``RouterConfig`` and the host's process settings; the views (``registry()``, ``settings()``, +``health()``) are read from the Rust handle. A member the Rust core does not serve raises +``RustCoreUnsupported`` (an ``AttributeError``) instead of reaching a Python core, so deleting +``router_python_core.py`` cannot change what this core does. +""" + +import contextlib +import dataclasses +import inspect +import itertools +import json +from collections.abc import Awaitable, Callable, Coroutine, Iterable, Mapping, Sequence +from datetime import date, datetime, time +from functools import partial +from time import perf_counter +from time import time as unix_time +from types import MappingProxyType +from typing import TYPE_CHECKING, Final, Literal, Protocol, TypeAlias, final + +from litellm._internal_context import service_target +from litellm.constants import DEFAULT_COOLDOWN_TIME_SECONDS, HEALTH_CHECKS_TARGET +from litellm.integrations.custom_guardrail import is_guardrail_intervention +from litellm.litellm_core_utils.add_retry_fallback_headers import ( + add_fallback_headers_to_response, + add_retry_headers_to_response, + apply_quality_router_decision_headers, + apply_remaining_usage_headers, + apply_response_model_id, + complexity_router_decision_headers, + ensure_response_additional_headers, + prepare_response_for_header_attachment, +) +from litellm.litellm_core_utils.agentic_followup_failure import is_agentic_followup_failure +from litellm.litellm_core_utils.batch_utils import is_batch_retrieve_call_type +from litellm.litellm_core_utils.core_helpers import as_object_list, as_str_mapping, without +from litellm.litellm_core_utils.routing_common_utils import ( + format_content_policy_fallback_hint, + format_fallback_outcome_message, + format_no_fallback_group_message, + get_request_team_id, + is_web_search_request, + truncate_fallback_error_detail, +) +from litellm.litellm_core_utils.routing_groups import validate_routing_strategy +from litellm.litellm_core_utils.secret_redaction import redact_string +from litellm.router_api import ( + ROUTER_INIT_ARGS, + SETTING_FIELDS, + ApplyResult, + ChangeSet, + RouterHealth, + RouterRegistry, + RouterSettings, + RoutingContext, + RoutingObservers, +) +from litellm.router_attempt import for_requested_model, request_from_call +from litellm.router_config import RouterConfig, RouterDynamicConfig, detached +from litellm.router_python_host import PythonHost +from litellm.router_rust.invoke.endpoints import ( + AttemptCall, + Plan, + RustEndpoints, + arm_compaction_state, + attempt_call, + clean_realtime_session_deployment, + compacted_attempt, + failover_fields, + generic_plan, + log_refused_call, + log_refused_call_sync, + pinned_video_deployment, + plan_for, + reports_own_success, + scoped_fallback_fields, + selection_endpoint, + stamp_error, + unrouted_fields, +) +from litellm.router_rust.invoke.shaping import Shaping +from litellm.router_rust.invoke.slots import enter_deployment_slot, max_parallel_requests_limit +from litellm.router_rust.invoke.streaming import ( + StreamAttempt, + continuation_flag, + continuation_routing, + open_attempt_stream, + open_attempt_stream_sync, + register_continuation_prefill_check, + shape_attempt_kwargs, +) +from litellm.router_rust.invoke.walk import FallbackOutcomes, WalkMetadata, await_all, log_retry_into +from litellm.router_rust.steps.access import caller_fields, selection_refusal +from litellm.router_rust.steps.callbacks import OptionalChecks, RustOptionalChecks +from litellm.router_rust.steps.plugins import SYNC_PATH_MESSAGE, candidate_models, run_plugins +from litellm.router_rust.steps.precall import ( + DECISIONS_MEDIA_CALL_TYPES, + PreCallFacts, + cannot_serve_route, + declares_no_vision, + refuses_decisions_media, + request_carries_image, + serves_route_capability, +) +from litellm.router_rust.steps.prerouting import PreRouter, PreRoutingStrategies, picks_async +from litellm.router_rust.steps.selection import ( + SELECTOR_ATTRIBUTES, + StrategySelectors, + group_documents, + group_settings, + member_map, + normalize_strategy, + parse_groups, + pick_async, + pick_sync, + picked_id, + routing_pre_call_checks, +) +from litellm.router_rust.surface.alerting import initialize_alerting +from litellm.router_rust.surface.fallbacks import ( + RustFallbacks, + mock_error, + mock_failure, + refuse_malformed_fallbacks, + rejected_fallback_target, +) +from litellm.router_rust.surface.health import ( + RustHealthCallbacks, + failure_already_logged, + failure_status, + first_cooldown, + logged_call, +) +from litellm.router_rust.surface.health_routing import ( + HEALTH_STATE_KEY, + HealthRouting, + failover_document, + fully_unhealthy_model_names, + resolved_allowed_fails_policy, +) +from litellm.router_rust.surface.health_routing import all_blocked as deployments_all_blocked +from litellm.router_rust.surface.health_routing import fully_blocked_model_names as blocked_model_names +from litellm.router_rust.surface.health_routing import is_blocked as deployment_is_blocked +from litellm.router_rust.surface.legacy_members import RustLegacyMembers +from litellm.router_rust.surface.members import RustRegistryMembers +from litellm.router_rust.surface.model_list import typed_deployment +from litellm.router_rust.surface.registry import ( + DEFAULT_DEPLOYMENT_COPIED_ID, + DEFAULT_DEPLOYMENT_FAILING_ID, + INDEX_ATTRIBUTES, + BuildInstall, + HostRegistry, + RustRegistry, + default_deployment_call_error, + default_deployment_copy, + default_deployment_copy_error, + hosted_registry_view, + index_document, + served_default_deployment, +) +from litellm.router_rust.surface.replay import track_live_registry, untrack_live_registry +from litellm.router_rust.surface.retries import ( + GroupPolicies, + RustRetries, + group_policies, + is_policy_per_group, + policy_lookup_error, + raising_groups, + readable_retry_policy, + reads_as_policy, + request_retry_fields, + resolved_retry_policy, + retry_count_error, + static_group_policies, +) +from litellm.router_rust.surface.settings import ( + MISSING, + PLAIN_DEFAULTS, + SettingsUpdate, + apply_construction_effects, + argument_attribute, + default_litellm_params, + initial_plain, + read_settings_update, + request_timeout, + settings_view, +) +from litellm.router_strategy.least_busy import CALL_IDENTITY, call_id_of, owning_picks, pick_reserves, picks_reserve +from litellm.router_utils.cooldown_cache import ROUTER_COOLDOWNS_TARGET +from litellm.router_utils.cooldown_handlers import get_cooldown_deployment_info, mark_deployment_selection_failure +from litellm.router_utils.endpoint_eligibility import endpoint_for_call, selected_endpoint +from litellm.router_utils.kubernetes_pod_discovery import ( + KubernetesPodDiscovery, + async_resolve_pods_after_bound, + resolve_pods_after_bound, +) +from litellm.router_utils.pre_call_checks.encrypted_content_affinity_check import EncryptedContentAffinityCheck +from litellm.router_utils.router_callbacks.track_deployment_metrics import find_deployment_metadata +from litellm.router_utils.routing_read_batch import ROUTER_USAGE_TARGET, RoutingReadBatch +from litellm.types.router import ( + Deployment, + RouterErrors, + RouterGeneralSettings, + RouterRateLimitError, + RoutingGroup, + RoutingStrategy, +) +from litellm.types.utils import LlmProviders + +if TYPE_CHECKING: + from opentelemetry.trace import Span + + from litellm.caching.dual_cache import DualCache + from litellm.caching.redis_cache import RedisCache + from litellm.integrations.SlackAlerting.slack_alerting import SlackAlerting + from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging + from litellm.litellm_core_utils.ptu_shares import PTUTeamCeiling + from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler + from litellm.router_api import UntypedDict + from litellm.router_utils.cooldown_cache import CooldownCache + from litellm.router_utils.health_state_cache import DeploymentHealthCache + from litellm.rust_bridge._native import ( + RouterCall, + RouterCoreHandle, + _CallSteps, # pyright: ignore[reportPrivateUsage] # the stub's step shapes + _FailureReportSteps, # pyright: ignore[reportPrivateUsage] # see above + _SuccessReportSteps, # pyright: ignore[reportPrivateUsage] # see above + ) + from litellm.types.router import ( + CustomRoutingStrategyBase, + DeploymentTypedDict, + GuardrailTypedDict, + PreRoutingHookResponse, + RoutingPlugin, + ) + +_FAILURE_KINDS: Final = ( + "RateLimitError", + "ContextWindowExceededError", + "ContentPolicyViolationError", + "AuthenticationError", + "PermissionDeniedError", + "NotFoundError", + "BadRequestError", + "Timeout", + "APIConnectionError", + "InternalServerError", + "ServiceUnavailableError", + "BadGatewayError", + "GuardrailIntervention", + "CallerCredentialAuthenticationError", + "CallerCredentialRateLimitError", +) +_DECLINED: Final = frozenset({"MalformedRequestFieldError", "UnsupportedEndpointError"}) +_PROVIDERS: Final = tuple(provider.value for provider in LlmProviders) +_DEFAULT_COOLDOWN_SECONDS: Final = 5.0 +# Past this a retry-after no longer fits the Rust core's durations; it cools down as good as forever. +_LONGEST_RETRY_AFTER_SECONDS: Final = 2**63 + +# Constructor arguments this core reads. Any other argument set to something other than its +# unset value keeps the router on the Python core while both exist (make_core decides that from +# the config alone, before any core is built). +_SERVED_ARGUMENTS: Final = frozenset( + { + "model_list", + "routing_strategy", + "routing_strategy_args", + "routing_groups", + "enable_tag_filtering", + "tag_filtering_match_any", + "tag_routing_prefix", + "model_group_alias", + "fallbacks", + "default_fallbacks", + "context_window_fallbacks", + "content_policy_fallbacks", + "num_retries", + "max_fallbacks", + "allowed_fails", + "cooldown_time", + "disable_cooldowns", + "allowed_fails_policy", + "retry_after", + "timeout", + "default_litellm_params", + "fallback_access_check", + "fallback_budget_check", + "enable_mid_stream_fallback_continuation", + "redis_url", + "redis_host", + "redis_port", + "redis_password", + "redis_db", + "cache_kwargs", + "set_verbose", + "debug_level", + # Read host-side by the endpoint plans (litellm/router_rust/invoke/endpoints.py). + "stream_timeout", + "cache_responses", + "assistants_config", + "search_tools", + "guardrail_list", + "default_priority", + "polling_interval", + "ignore_invalid_deployments", + # Read host-side by the pre-routing strategies (litellm/router_rust/steps/prerouting.py). + "auto_router_capability_limit", + # Decided in router-core (precall.rs); its facts come from litellm/router_rust/steps/precall.py. + "enable_pre_call_checks", + # Run host-side (litellm/router_rust/steps/plugins.py); router-core filters on what they leave. + "plugins", + "retry_policy", + "model_group_retry_policy", + # Health-check routing and the failover flags (litellm/router_rust/surface/health_routing.py). + "enable_health_check_routing", + "health_check_staleness_threshold", + "health_check_ignore_transient_errors", + "background_health_check_model_groups", + "enable_weighted_failover", + "fallback_on_unsupported_params", + # Built host-side as CustomLogger checks (litellm/router_rust/steps/callbacks.py); router-core runs + # their deployment filters through its callback_filters stage. + "optional_pre_call_checks", + "provider_budget_config", + "deployment_affinity_ttl_seconds", + "session_affinity_repin_on_failover", + "model_group_affinity_config", + # Kept as given (litellm/router_rust/surface/settings.py) and read host-side by the endpoint plans + # (default_litellm_params' caching groups) and the slot limits. + "client_ttl", + "caching_groups", + "default_max_parallel_requests", + # Built host-side into a Slack alerting logger (litellm/router_rust/surface/alerting.py). + "alerting_config", + # Kept as given and reported by settings(); nothing in the router reads its fields. + "router_general_settings", + } +) +# Arguments that may be passed at their default without asking for anything. +_UNSET_DEFAULTS: Final[Mapping[str, object]] = MappingProxyType( + { + "cache_responses": False, + "enable_pre_call_checks": False, + "enable_tag_filtering": False, + "retry_policy": None, + "model_group_retry_policy": {}, + } +) +# Settings a changeset may change. Static ones rebuild the Rust core (settings are fixed for a +# core's life); dynamic ones are handed to it as a new dynamic document. +_STATIC_SETTINGS: Final = frozenset( + { + "routing_strategy", + "routing_groups", + "cooldown_time", + "disable_cooldowns", + "allowed_fails_policy", + "retry_after", + "timeout", + "enable_pre_call_checks", + } +) +_DYNAMIC_SETTINGS: Final = frozenset( + { + "default_deployment", + "fallbacks", + "context_window_fallbacks", + "content_policy_fallbacks", + "num_retries", + "max_fallbacks", + "allowed_fails", + "model_group_alias", + "default_litellm_params", + "fallback_access_check", + "fallback_budget_check", + "routing_strategy_args", + "enable_tag_filtering", + "tag_filtering_match_any", + "tag_routing_prefix", + } +) +# Settings only the host reads (the endpoint plans, the slot limits, the budget limiter, the +# settings view). A changeset assigns them as given, the Rust core never sees them, and +# `update_settings` ignores them (none is runtime-updatable). +_HOST_SETTINGS: Final = frozenset( + { + "cache_responses", + "search_tools", + "guardrail_list", + "default_max_parallel_requests", + "provider_budget_config", + "health_check_ignore_transient_errors", + "router_general_settings", + "provider_default_deployment_ids", + "kubernetes_pod_discovery", + # Read by router-core's health-check filter: assigning one rebuilds the Rust core. + "enable_health_check_routing", + "background_health_check_model_groups", + } +) +# The constructor's default `router_general_settings`: one object every router built without the +# argument shares, as the Python core's default argument is. +_DEFAULT_GENERAL_SETTINGS: Final = RouterGeneralSettings() + + +# Plain settings `update_settings` sets that the Rust core's retry loop reads. +_RETRY_POLICY_SETTINGS: Final = ("retry_policy", "model_group_retry_policy") +# Attributes a `del` unsets beyond the setting fields, whose readers then raise AttributeError. +_DELETABLE_ATTRIBUTES: Final = frozenset({"model_group_alias", "num_retries", "scheduler"}) + + +# Calls whose pick least-busy does not hold (`picks_reserve(False)`): a batch cancel makes no +# call that logs a start. +_UNRESERVED_OPS: Final = frozenset({"acancel_batch"}) + + +# Settings converted by `Resolved._selection_setting`. +_SELECTION_SETTINGS: Final = frozenset( + { + "routing_strategy", + "routing_strategy_args", + "enable_tag_filtering", + "tag_filtering_match_any", + "tag_routing_prefix", + } +) +_NUMBER_SETTINGS: Final = frozenset( + {"num_retries", "max_fallbacks", "allowed_fails", "cooldown_time", "retry_after", "timeout"} +) +# The settings `_core_update_settings` casts with `int(value)`. +_UPDATE_INT_SETTINGS: Final = frozenset({"timeout", "num_retries", "retry_after", "allowed_fails", "cooldown_time"}) + + +class RustCoreUnsupported(AttributeError): + """A member or request the Rust core does not serve. It is never forwarded to a Python core.""" + + +def _cache_target(key: str) -> str: + """The service target a routing cache read or write is attributed to, as the Python core's.""" + if key == HEALTH_STATE_KEY: + return HEALTH_CHECKS_TARGET + return ROUTER_COOLDOWNS_TARGET if key.endswith(":cooldown") else ROUTER_USAGE_TARGET + + +# --- failures ----------------------------------------------------------------------------------- + + +def failure_kind(error: BaseException) -> str: + names: Final = tuple(cls.__name__ for cls in type(error).__mro__) + if not _DECLINED.isdisjoint(names): + return "Declined" + if is_guardrail_intervention(error): # pyright: ignore[reportArgumentType] # every attempt error is an Exception + return "GuardrailIntervention" + return next((name for name in names if name in _FAILURE_KINDS), "APIError") + + +def retry_after_seconds(error: BaseException) -> float | None: + """The header cooldown of ``deployment_callback_on_failure``: ``retry-after`` alone, as a whole + number or a date, with a negative or unreadable one read as none.""" + from litellm.litellm_core_utils.exception_mapping_utils import ( + _get_response_headers, # pyright: ignore[reportPrivateUsage] # the failure callback's own read + ) + from litellm.utils import _get_retry_after_from_exception_header # pyright: ignore[reportPrivateUsage] # see above + + headers: Final = _get_response_headers(original_exception=error) # pyright: ignore[reportArgumentType] # every attempt error is an Exception + if headers is None: + return None + seconds: Final = _get_retry_after_from_exception_header(response_headers=headers) # pyright: ignore[reportUnknownVariableType] # an untyped helper's int or None + return float(min(seconds, _LONGEST_RETRY_AFTER_SECONDS)) if seconds is not None and seconds >= 0 else None + + +def attempt_failure(error: BaseException) -> Mapping[str, object]: + status: Final = getattr(error, "status_code", None) + num_retries: Final = getattr(error, "num_retries", None) + # `_deployment_ids_to_skip_on_retry` reads the id off the exception. + named_deployment: Final = getattr(error, "retry_skip_deployment_id", None) or getattr( + error, "failed_deployment_id", None + ) + return { + "ok": False, + "kind": failure_kind(error), + "status": status if isinstance(status, int) and 100 <= status <= 599 else None, + "status_absent": getattr(error, "status_code", "") == "", + "retry_after": retry_after_seconds(error), + "num_retries": num_retries if isinstance(num_retries, int) and num_retries >= 0 else None, + "failed_deployment_id": named_deployment if isinstance(named_deployment, str) else None, + "detail": str(error), + "classes": [cls.__name__ for cls in type(error).__mro__], + "llm_provider": bool(getattr(error, "llm_provider", "")), + "agentic_followup": is_agentic_followup_failure(error), + } + + +def total_tokens(response: object) -> int: + usage: Final = getattr(response, "usage", None) + tokens: Final = getattr(usage, "total_tokens", None) + return tokens if isinstance(tokens, int) else 0 + + +# --- settings ----------------------------------------------------------------------------------- + + +def _json_default(value: object) -> object: + # ISO 8601, which ``typed_deployment`` reads back as the value it was. + if isinstance(value, date | time): + return value.isoformat() + dump: Final = getattr(value, "model_dump", None) + return dump(exclude_none=True) if callable(dump) else str(value) + + +def _plain(value: object) -> object: + """``value`` as plain JSON values, the shape the Rust core reads.""" + loaded: Final[object] = json.loads(json.dumps(value, default=_json_default)) # pyright: ignore[reportAny] # plain JSON + return loaded + + +def _document(value: object) -> Mapping[str, object]: + """A mapping as the plain JSON document the Rust core reads.""" + return as_str_mapping(_plain(value)) or {} + + +def _as_list(value: object) -> Sequence[object]: + return value if isinstance(value, list) else () # pyright: ignore[reportUnknownVariableType] # a list's items are read as objects + + +def _strategy_name(value: object) -> str | None: + """A strategy as a setting stores it: ``None`` stays no strategy (``_normalize_strategy``).""" + return None if value is None else normalize_strategy(value) or "simple-shuffle" + + +def _lenient_int(value: object, default: int) -> int: + """A count setting kept as given, as the crate reads it: an int, or a number spelled another way + (``1.5``, ``"3"``) truncated to one; ``default`` when it is no number. The Python core stores the + value as given and never checks it at construction.""" + if isinstance(value, bool | int): + return int(value) + try: + return int(float(str(value))) + except (ValueError, OverflowError): + return default + + +def _lenient_float(value: object, default: float | None) -> float | None: + """A seconds setting kept as given, as the crate reads it (see ``_lenient_int``).""" + if isinstance(value, bool | int | float): + return float(value) + try: + return float(str(value)) + except (ValueError, OverflowError): + return default + + +# Settings the Python core stores as given, so that a read returns the value passed (``"1"``, +# ``None``, ``[1]``); ``Resolved`` keeps them in ``given`` and the typed fields hold what the crate reads. +_KEPT_AS_GIVEN: Final = frozenset( + { + "num_retries", + "max_fallbacks", + "allowed_fails", + "retry_after", + "timeout", + "disable_cooldowns", + "enable_tag_filtering", + "tag_filtering_match_any", + "tag_routing_prefix", + "routing_strategy_args", + "fallbacks", + "context_window_fallbacks", + "content_policy_fallbacks", + } +) + + +@dataclasses.dataclass(frozen=True, slots=True) +class Resolved: + """The router's settings, resolved from its arguments and the process settings the way the + Python router's constructor resolves them. Replaced, never changed, by ``apply``.""" + + routing_strategy: str | None + routing_strategy_args: Mapping[str, object] + num_retries: int + max_fallbacks: int + allowed_fails: int + explicit_allowed_fails: int | None + cooldown_time: float + disable_cooldowns: bool | None + allowed_fails_policy: object + retry_after: float + timeout: float | None + fallbacks: list[object] | None # mutable-ok: the setting as Python spells it + context_window_fallbacks: list[object] | None # mutable-ok: same + content_policy_fallbacks: list[object] | None # mutable-ok: same + model_group_alias: Mapping[str, object] + default_litellm_params: Mapping[str, object] + fallback_access_check: object + fallback_budget_check: object + enable_mid_stream_fallback_continuation: bool = False + routing_groups: "tuple[RoutingGroup, ...]" = () + enable_tag_filtering: bool = False + tag_filtering_match_any: bool = True + tag_routing_prefix: str = "" + enable_pre_call_checks: bool = False + routing_plugins: "tuple[RoutingPlugin, ...]" = () + retry_policy: object = None + model_group_retry_policy: object = None + health_routing: HealthRouting = dataclasses.field(default_factory=HealthRouting) + # `_core_default_deployment`, as assigned (detached). + default_deployment: object = None + # `_model_to_group` as the Python core last built it, kept once a raw `routing_groups` write + # replaced the groups without rebuilding it. + routing_group_members: Mapping[str, str] | None = None + # `_routing_groups_input`: the groups last passed to the constructor or `update_settings`, which + # a strategy change through `update_settings` re-inits the routing groups from. + routing_groups_input: object = None + assigned: frozenset[str] = frozenset() + # The `_KEPT_AS_GIVEN` settings as the Python core stores them, for reads. + given: Mapping[str, object] = dataclasses.field(default_factory=lambda: MappingProxyType({})) + + def read(self, name: str) -> object: + """A setting as the Python core reads it back: as given when it keeps it so.""" + return self.given[name] if name in self.given else getattr(self, name) # pyright: ignore[reportAny] # a setting field + + @classmethod + def of(cls, config: RouterConfig, dynamic: RouterDynamicConfig) -> "Resolved": + get: Final = config.get + # The process list is the router's own, as in the Python router: default_fallbacks are + # appended to it in place. + configured: Final[list[object] | None] = ( # mutable-ok: the router's live list, as Python keeps it + list(_as_list(get("fallbacks"))) or (dynamic.fallbacks if dynamic.fallbacks else None) # pyright: ignore[reportAssignmentType] # the process list as given + ) + refuse_malformed_fallbacks(configured) + default_fallbacks: Final = get("default_fallbacks") or dynamic.default_fallbacks + if isinstance(default_fallbacks, str): + raise TypeError( + f"default_fallbacks must be a list of model names, got '{default_fallbacks}'. Use ['{default_fallbacks}']" + ) + if default_fallbacks is not None and configured is not None: + configured.append({"*": default_fallbacks}) + fallbacks: Final[list[object] | None] = ( # mutable-ok: see above + [{"*": default_fallbacks}] if default_fallbacks is not None and configured is None else configured + ) + content_policy: Final = list(_as_list(get("content_policy_fallbacks"))) or ( + list(dynamic.content_policy_fallbacks) if dynamic.content_policy_fallbacks else None + ) + refuse_malformed_fallbacks(content_policy) + num_retries: Final = get("num_retries") + max_fallbacks: Final = get("max_fallbacks") + allowed_fails: Final = get("allowed_fails") + timeout: Final = get("timeout") + strategy: Final = get("routing_strategy") + validate_routing_strategy(strategy) # pyright: ignore[reportArgumentType] # validated as given + # The values the Python core's constructor stores (`self.num_retries = num_retries`, ...). + given_num_retries: Final[object] = ( + num_retries if num_retries is not None else dynamic.num_retries if dynamic.num_retries is not None else 2 + ) + given_max_fallbacks: Final[object] = ( + max_fallbacks + if max_fallbacks is not None + else dynamic.max_fallbacks + if dynamic.max_fallbacks is not None + else dynamic.router_max_fallbacks + ) + given_allowed_fails: Final[object] = allowed_fails if allowed_fails is not None else dynamic.allowed_fails + given_timeout: Final[object] = timeout or dynamic.request_timeout + given_retry_after: Final[object] = get("retry_after") if "retry_after" in config.arguments else 0 + kept: Final = { + "num_retries": given_num_retries, + "max_fallbacks": given_max_fallbacks, + "allowed_fails": given_allowed_fails, + "timeout": given_timeout, + "retry_after": given_retry_after, + "disable_cooldowns": get("disable_cooldowns"), + **{ + name: get(name) + for name in ("enable_tag_filtering", "tag_filtering_match_any", "tag_routing_prefix") + if name in config.arguments + }, + **( + {"routing_strategy_args": get("routing_strategy_args")} + if "routing_strategy_args" in config.arguments + else {} + ), + } + return cls( + # An explicit `routing_strategy=None` is stored as None (no strategy picks). + routing_strategy=None + if strategy is None and "routing_strategy" in config.arguments + else normalize_strategy(strategy) or "simple-shuffle", + # A plain dict, as `self.routing_strategy_args`: `/router/settings` serializes it as read. + routing_strategy_args=dict(as_str_mapping(get("routing_strategy_args")) or {}), + num_retries=_lenient_int(given_num_retries, 2), + max_fallbacks=_lenient_int(given_max_fallbacks, dynamic.router_max_fallbacks), + allowed_fails=_lenient_int(given_allowed_fails, dynamic.allowed_fails), + explicit_allowed_fails=_lenient_int(allowed_fails, dynamic.allowed_fails) + if allowed_fails is not None + else None, + cooldown_time=_lenient_float(get("cooldown_time") or _DEFAULT_COOLDOWN_SECONDS, _DEFAULT_COOLDOWN_SECONDS) + or 0.0, + disable_cooldowns=_optional_bool(get("disable_cooldowns")), + allowed_fails_policy=resolved_allowed_fails_policy(get("allowed_fails_policy")), + retry_after=_lenient_float(given_retry_after or 0, 0.0) or 0.0, + timeout=_lenient_float(given_timeout, None) if given_timeout else None, + given=MappingProxyType(kept), + fallbacks=fallbacks, + context_window_fallbacks=list(_as_list(get("context_window_fallbacks"))) + or (list(dynamic.context_window_fallbacks) if dynamic.context_window_fallbacks else None), + content_policy_fallbacks=content_policy, + model_group_alias=MappingProxyType(dict(as_str_mapping(get("model_group_alias")) or {})), + default_litellm_params=MappingProxyType(dict(as_str_mapping(get("default_litellm_params")) or {})), + fallback_access_check=get("fallback_access_check"), + fallback_budget_check=get("fallback_budget_check"), + enable_mid_stream_fallback_continuation=continuation_flag(get("enable_mid_stream_fallback_continuation")), + enable_tag_filtering=bool(get("enable_tag_filtering")), + tag_filtering_match_any=_true_unless_given(get("tag_filtering_match_any")), + tag_routing_prefix=str(get("tag_routing_prefix") or ""), + enable_pre_call_checks=bool(get("enable_pre_call_checks")), + routing_plugins=tuple(_as_list(get("plugins"))), # pyright: ignore[reportArgumentType] # the plugins as passed + routing_groups_input=get("routing_groups"), + routing_groups=parse_groups( + get("routing_groups"), + frozenset( + str(name) + for entry in _as_list(get("model_list")) + if (name := (as_str_mapping(entry) or {}).get("model_name")) + ), + frozenset(as_str_mapping(get("model_group_alias")) or {}), + ), + retry_policy=resolved_retry_policy(get("retry_policy")), + model_group_retry_policy=get("model_group_retry_policy"), + health_routing=HealthRouting.of(config), + ) + + def changed(self, values: Mapping[str, object], updated: frozenset[str] = frozenset()) -> "Resolved": + """``values`` applied; ``updated`` names those passed as ``ChangeSet.settings`` (which go + through ``update_settings``: only they may name ``routing_strategy_args``).""" + from litellm.router_api import ASSIGNABLE_ALIASES, PLAIN_SETTING_ATTRIBUTES, SETTING_FIELDS + + unknown: Final = sorted( + set(values) + - {*SETTING_FIELDS, *PLAIN_SETTING_ATTRIBUTES, *ASSIGNABLE_ALIASES} + - (updated & {"routing_strategy_args", "routing_groups", "enable_tag_filtering", "tag_routing_prefix"}) + ) + if unknown: + raise ValueError(f"not assignable router settings: {unknown}") + unserved: Final = sorted(set(values) - _STATIC_SETTINGS - _DYNAMIC_SETTINGS) + if unserved: + raise RustCoreUnsupported(f"the Rust router core does not serve the setting(s) {', '.join(unserved)} yet") + updates: Final[dict[str, object]] = {} # mutable-ok: built here, then applied once + for name, value in values.items(): + if name in {"model_group_alias", "default_litellm_params"}: + updates[name] = MappingProxyType(dict(as_str_mapping(value) or {})) + elif name in _NUMBER_SETTINGS: + updates[name] = self._number_setting(name, value) + if name == "allowed_fails": + updates["explicit_allowed_fails"] = updates[name] + elif name in _SELECTION_SETTINGS: + updates[name] = self._selection_setting(name, value) + elif name == "disable_cooldowns": + updates[name] = _optional_bool(value) + elif name == "enable_pre_call_checks": + updates[name] = bool(value) + elif name in ("fallbacks", "context_window_fallbacks", "content_policy_fallbacks"): + updates[name] = list(_as_list(value)) or None + elif name == "default_deployment": + updates[name] = detached(value) + else: + updates[name] = value + if "routing_groups" in values: + # `_init_routing_groups` rebuilds the member map from the new groups. + updates["routing_group_members"] = None + # `_core_update_settings` casts its int settings (`int(value)`, which raises as Python's + # does); every other write is stored as given. + kept: Final = { + name: int(value) if name in updated and name in _UPDATE_INT_SETTINGS else value # pyright: ignore[reportArgumentType] # int() of the value as given, as Python casts it + for name, value in values.items() + if name in _KEPT_AS_GIVEN + } + return dataclasses.replace( + self, + **updates, + assigned=self.assigned | frozenset(values), + given=MappingProxyType({**self.given, **kept}), + ) # pyright: ignore[reportArgumentType] # each value converted to its field's type above + + def _number_setting(self, name: str, value: object) -> object: + """A numeric setting as the crate reads it; the value as given is kept in ``given``.""" + if name in ("num_retries", "max_fallbacks", "allowed_fails"): + return _lenient_int(value, getattr(self, name)) # pyright: ignore[reportAny] # the field's current int + if name == "cooldown_time": + # Stored as given (`_core_cooldown_time`), so a failure cools down for that long: + # 0 for no time at all; None leaves the cooldown cache's own default, set by the caller. + return _lenient_float(value, self.cooldown_time) if value is not None else self.cooldown_time + if name == "retry_after": + return _lenient_float(value or 0, self.retry_after) + return _lenient_float(value, self.timeout) if value is not None else None + + def _selection_setting(self, name: str, value: object) -> object: + """A selection setting converted to its field's type.""" + if name == "routing_strategy": + return _strategy_name(value) + if name == "routing_strategy_args": + return dict(as_str_mapping(value) or {}) + if name == "tag_routing_prefix": + return str(value or "") + return bool(value) + + def static_document(self, model_list: Sequence[object]) -> Mapping[str, object]: + values: Final = { + "model_list": _plain(list(model_list)), + "routing_strategy": self.routing_strategy, + "cooldown_time": self.cooldown_time, + "disable_cooldowns": self.disable_cooldowns, + "allowed_fails_policy": _plain(self.allowed_fails_policy), + "retry_after": self.retry_after, + "timeout": self.timeout, + "routing_groups": group_documents(self.routing_groups) if self.routing_groups else None, + "enable_pre_call_checks": True if self.enable_pre_call_checks else None, + "routing_plugins": True if self.routing_plugins else None, + # `async_callback_filter_deployments` runs on every async candidate list. + "callback_filters": True, + "retry_policy": _plain(self.retry_policy), + "model_group_retry_policy": _plain(static_group_policies(self.model_group_retry_policy)) or None, + "routing_group_members": dict(self.routing_group_members) + if self.routing_group_members is not None + else None, + **self.health_routing.static_document(), + } + kept: Final = {name: value for name, value in values.items() if value is not None} + # An explicit `routing_strategy=None` reaches the Rust core as "None": no strategy picks. + return kept if self.routing_strategy is not None else {**kept, "routing_strategy": "None"} + + def dynamic_document(self, version: int, model_names: Sequence[str]) -> Mapping[str, object]: + values: Final = { + "version": version, + "model_group_alias": _plain(dict(self.model_group_alias)), + "fallbacks": _plain(self.fallbacks), + "context_window_fallbacks": _plain(self.context_window_fallbacks), + "content_policy_fallbacks": _plain(self.content_policy_fallbacks), + "num_retries": self.num_retries, + "max_fallbacks": self.max_fallbacks, + "allowed_fails": self.allowed_fails, + "allowed_fails_is_default": True if self.allowed_fails_is_default() else None, + "fallback_keys": self._fallback_keys(model_names), + "default_deployment": _plain(served_default_deployment(self.default_deployment)[0]), + "default_deployment_raises": True if served_default_deployment(self.default_deployment)[1] else None, + } + return {name: value for name, value in values.items() if value is not None} + + def allowed_fails_is_default(self) -> bool: + """``allowed_fails`` does not count as set on the router (``_is_allowed_fails_set_on_router``): + not passed explicitly, and the value copied at init still equals the live process default.""" + import litellm + + return self.explicit_allowed_fails is None and self.allowed_fails == litellm.allowed_fails + + def _fallback_keys(self, model_names: Sequence[str]) -> Mapping[str, object]: + """The spellings a fallback key matches besides the group name (`get_fallback_model_group`).""" + from litellm.router_utils.fallback_event_handlers import ( + _provider_prefixed_model_group, # pyright: ignore[reportPrivateUsage] # the fallback walk's own spelling rule + ) + + keys: Final = tuple( # comprehension-ok: the first group of each fallback entry, across the three lists + str(next(iter(mapping))) + for entries in (self.fallbacks, self.context_window_fallbacks, self.content_policy_fallbacks) + for entry in _as_list(entries) + if (mapping := as_str_mapping(entry)) + ) + names: Final = {*model_names, *self.model_group_alias} + prefixed: Final = { + name: spelled for name in names if (spelled := _provider_prefixed_model_group(name, keys)) is not None + } + from litellm.types.router import LiteLLMParamsTypedDict + + return { + "providers": _PROVIDERS, + "provider_prefixed": prefixed, + "param_keys": list(LiteLLMParamsTypedDict.__annotations__), + } + + +def _alias_target(target: object) -> str | None: + mapping: Final = as_str_mapping(target) + if mapping is not None: + model: Final = mapping.get("model") + return model if isinstance(model, str) else None + return target if isinstance(target, str) else None + + +def _true_unless_given(value: object) -> bool: + return True if value is None else bool(value) + + +def _group_selectors(groups: "Sequence[RoutingGroup]") -> Sequence[tuple[str, str, Mapping[str, object]]]: + return [ + ( + group.group_name, + group.routing_strategy, + as_str_mapping(group.routing_strategy_args) or {}, # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType] # RoutingGroup's untyped dict + ) + for group in groups + ] + + +def _optional_bool(value: object) -> bool | None: + return None if value is None else bool(value) + + +def unsupported_arguments(config: RouterConfig) -> frozenset[str]: + """Arguments the Rust core does not serve.""" + unserved: Final = frozenset( + name + for name, value in config.arguments.items() + if name not in _SERVED_ARGUMENTS and value not in (None, _UNSET_DEFAULTS.get(name)) + ) + return unserved + + +def _listing_team_id(request: Mapping[str, object]) -> str | None: + """``(kwargs.get("metadata", {}) or {}).get("user_api_key_team_id")``: the team + ``async_function_with_fallbacks_common_utils`` lists a group for reads this bucket alone.""" + metadata: Final = as_str_mapping(request.get("metadata")) + team_id: Final = metadata.get("user_api_key_team_id") if metadata is not None else None + return team_id if isinstance(team_id, str) else None + + +# Request fallback settings, which replace the router's for the call. +_REQUEST_FALLBACK_KEYS: Final = ("fallbacks", "context_window_fallbacks", "content_policy_fallbacks", "max_fallbacks") + + +def _strings(value: object) -> "tuple[str, ...]": + """Every string in ``value``, a name or a fallback list of names and dict entries.""" + if isinstance(value, str): + return (value,) + if isinstance(value, Mapping): + entries: Final[Mapping[object, object]] = value # pyright: ignore[reportUnknownVariableType] # read as objects + return tuple(itertools.chain.from_iterable(map(_strings, itertools.chain.from_iterable(entries.items())))) + if isinstance(value, list | tuple): + items: Final[Sequence[object]] = value # pyright: ignore[reportUnknownVariableType] # see above + return tuple(itertools.chain.from_iterable(map(_strings, items))) # pyright: ignore[reportUnknownArgumentType] # read as objects + return () + + +def _routed_model(request: Mapping[str, object]) -> object: + """The request's ``model`` as a recovery of its stream names it (``RustCore.routed_name``).""" + model: Final = request.get("model") + return model if model is None or isinstance(model, str) else str(model) + + +def _start_fields(request: Mapping[str, object]) -> Mapping[str, object]: + """The request fields the Rust core routes on.""" + fields: Final = { + **request_retry_fields(request.get("num_retries")), + **({"cooldown_time": request["cooldown_time"]} if request.get("cooldown_time") is not None else {}), + } + # `_common_checks_available_deployment`: only True names a litellm model. + specific: Final = {"specific_deployment": True} if request.get("specific_deployment") is True else {} + # `record_disable_fallbacks(kwargs, disable_fallbacks is True)`: only True turns them off. + disabled: Final = {"disable_fallbacks": True} if request.get("disable_fallbacks") is True else {} + lists: Final = {key: _plain(request[key]) for key in _REQUEST_FALLBACK_KEYS if key in request} + weights: Final = request.get("_router_weights") + team_id: Final = get_request_team_id(request) + listing_team_id: Final = _listing_team_id(request) + strategy: Final = request.get("routing_strategy") + override: Final = normalize_strategy(strategy) if isinstance(strategy, str | RoutingStrategy) else None + return { + **fields, + **specific, + **({"routing_strategy": override} if override is not None else {}), + **disabled, + **({"router_weights": weights} if weights is not None else {}), + **lists, + **({"team_id": team_id} if team_id else {}), + **({"listing_team_id": listing_team_id} if listing_team_id is not None else {}), + **caller_fields(request), + } + + +# --- views -------------------------------------------------------------------------------------- + + +class RustSettings(RouterSettings): + """``settings()``: the resolved settings as a frozen snapshot.""" + + def guardrail(self, name: str) -> "GuardrailTypedDict": + router: Final = self._router + if not isinstance(router, _SettingsSource): + return super().guardrail(name) + return router.available_guardrail(name) # pyright: ignore[reportReturnType] # the configured guardrail entry + + def as_dict(self) -> Mapping[str, object]: + router: Final = self._router + if not isinstance(router, _SettingsSource): + return super().as_dict() + return router.settings_dict() + + +class _SettingsSource: + def settings_dict(self) -> Mapping[str, object]: + raise NotImplementedError + + def available_guardrail(self, name: str) -> object: + raise NotImplementedError + + +class RustHealth(RouterHealth): + """``health()``: cooldowns, fail counts and usage, read through the host's routing cache and + written through the Rust core.""" + + __slots__ = ("_core",) + + def __init__(self, core: "RustCore") -> None: # pyright: ignore[reportMissingSuperCall] # reads the Rust core, not a Router + self._core = core + + def __getattribute__(self, name: str) -> object: + if name in _HEALTH_SERVED or name.startswith("__") or name == "_core": + found: Final[object] = object.__getattribute__(self, name) # pyright: ignore[reportAny] # the view's own member + return found + raise RustCoreUnsupported(f"health().{name} is not served by the Rust router core", name=name, obj=self) + + @property + def store(self) -> "DualCache": + return self._core._cache # pyright: ignore[reportPrivateUsage] # the health view reads its own core + + @property + def cooldowns(self) -> "CooldownCache": + return self._core._cooldowns # pyright: ignore[reportPrivateUsage] # see above + + def failures_this_minute(self, deployment_id: str) -> int: + count: Final[object] = self.store.get_cache(f"{deployment_id}:fails", local_only=True) # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] # DualCache's reads are untyped + return int(count) if isinstance(count, int | float) else 0 + + def record_failure(self, deployment_id: str, exception: Exception, exception_status: str | int) -> bool: + self._core._sync_allowed_fails() # pyright: ignore[reportPrivateUsage] # see above + failure: Final = { + **attempt_failure(exception), + "status": failure_status(exception_status), + "caller_scoped": self._core.per_user_deployment(deployment_id), + } + handle: Final = self._core._handle # pyright: ignore[reportPrivateUsage] # see above + call: Final = handle.report_failure(deployment_id, failure, count_only=exception_status == "") + done: Final = self._core._drive_sync_steps(call) # pyright: ignore[reportPrivateUsage] # see above + return self._core.cooled_down(done, deployment_id, exception_status) + + def team_ptu_ceiling(self, team_id: str, model_group: str) -> "PTUTeamCeiling | None": + return self._core._hosted._core_team_ptu_ceiling(team_id, model_group) # pyright: ignore[reportPrivateUsage] # the host registry's ported ceiling + + def split_priority(self, model: str, kwargs: Mapping[str, object]) -> tuple[int | None, Mapping[str, object]]: + """``_core_split_request_priority``, as the endpoints resolve a request's priority.""" + return self._core._endpoints._split_priority(model, kwargs) # pyright: ignore[reportPrivateUsage] # see store + + def priority_header_fields( + self, model: str, kwargs: Mapping[str, object], header: str | None + ) -> Mapping[str, int | str]: + """``priority_header_update``, as the endpoints read a priority header.""" + return self._core._endpoints.priority_header_update(model=model, kwargs=kwargs, header=header) # pyright: ignore[reportPrivateUsage] # see store + + @property + def health_checks(self) -> "DeploymentHealthCache": + return self._core._health_checks # pyright: ignore[reportPrivateUsage] # see store + + @property + def latency(self) -> "UntypedDict": # the live latency map, as the Python core keeps it + return self._core._hosted._core_deployment_latency_map # pyright: ignore[reportPrivateUsage] # see team_ptu_ceiling + + @staticmethod + def is_blocked(deployment: "Deployment") -> bool: + return deployment_is_blocked(deployment) + + @staticmethod + def all_blocked(deployments: "Sequence[DeploymentTypedDict]") -> bool: + return deployments_all_blocked(deployments) + + def fully_blocked_model_names(self) -> set[str]: # mutable-ok: the core's type, a fresh set + return blocked_model_names(self._core._hosted._core_get_model_list() or []) # pyright: ignore[reportPrivateUsage, reportArgumentType] # see team_ptu_ceiling + + async def fully_unhealthy_model_names(self) -> set[str]: # mutable-ok: the core's type, a fresh set + core: Final = self._core + return await fully_unhealthy_model_names( + core._hosted._core_get_model_list() or [], # pyright: ignore[reportPrivateUsage, reportArgumentType] # see team_ptu_ceiling + core._health_checks, # pyright: ignore[reportPrivateUsage] # see store + core._resolved.allowed_fails_policy, # pyright: ignore[reportPrivateUsage] # see store + ) + + async def usage( # pyright: ignore[reportIncompatibleMethodOverride] # the same two kinds + self, model_group: str, *, kind: Literal["current", "remaining"] = "current" + ) -> tuple[int | None, int | None] | Mapping[str, int]: + """The group's usage this minute, read from the routing cache over the host registry's + live model list, as ``get_model_group_usage`` / ``get_remaining_model_group_usage`` read it.""" + hosted: Final = self._core._hosted # pyright: ignore[reportPrivateUsage] # see above + if kind == "current": + return await hosted._core_get_model_group_usage(model_group) # pyright: ignore[reportPrivateUsage] # the ported usage read + if kind == "remaining": + return await hosted._core_get_remaining_model_group_usage(model_group) # pyright: ignore[reportPrivateUsage] # see above + raise ValueError(f"unknown usage kind {kind!r}") + + +# Deprecated attribute writes (``router. = ...`` reaches the core as ``_core_``) of +# registry state this core keeps on its host registry, and of routing state it rebinds. +_HOSTED_STATE: Final = frozenset( + { + "_core_model_names", + "_core_deployment_names", + "_core_team_public_model_names", + "_core_model_name_to_deployment_indices", + "_core_team_model_to_deployment_indices", + "_core_cached_deployment_model_info", + "_core_pattern_router", + "_core_team_pattern_routers", + "_core_routing_groups", + "_core_auto_routers", + "_core_complexity_routers", + "_core_quality_routers", + "_core_adaptive_routers", + "_core_routing_plugins", + "_core_deployment_latency_map", + } +) +_REBOUND_STATE: Final = MappingProxyType( + { + "_core_cache": "_cache", + "_core_cooldown_cache": "_cooldowns", + "_core_health_state_cache": "_health_checks", + "_core_slack_alerting_logger": "_slack_alerting_logger", + } +) + + +_HEALTH_SERVED: Final = frozenset( + { + "store", + "cooldowns", + "failures_this_minute", + "record_failure", + "usage", + "team_ptu_ceiling", + "split_priority", + "priority_header_fields", + "health_checks", + "latency", + "is_blocked", + "all_blocked", + "fully_blocked_model_names", + "fully_unhealthy_model_names", + } +) + + +# --- the core ----------------------------------------------------------------------------------- + + +_NO_INPUT: Final = object() +_EndpointImpl: TypeAlias = Callable[..., Awaitable[object]] +# A strategy select step's strategy, selector, keyed name and candidate deployments. +_StrategyInputs: TypeAlias = tuple[str, object, str, list[dict[str, object]]] + + +@final +class _PickCounts(Protocol): + """The router cache as ``track_deployment_pick`` uses it: each deployment id holds its pick count.""" + + def get_cache( + self, key: object, parent_otel_span: "Span | None" = None, local_only: bool = False + ) -> int | None: ... + + def set_cache(self, key: object, value: int, local_only: bool = False, *, ttl: int = ...) -> None: ... + + +class _EndpointPort: + """``RustCore`` as its endpoints see it (``EndpointCore``): each member forwards to the core's + own state or call drivers, read live, since the core replaces its handle and settings. The + port is the core's own adapter, so it is the one place its private members are read from + outside the class; every ``reportPrivateUsage`` ignore below is for that reason.""" + + __slots__ = ("_core",) + + def __init__(self, core: "RustCore") -> None: + self._core: Final = core + + @property + def router(self) -> "RustCore": + return self._core + + @property + def host(self) -> PythonHost: + return self._core._host # pyright: ignore[reportPrivateUsage] # see the class docstring + + @property + def handle(self) -> "RouterCoreHandle": + return self._core._handle # pyright: ignore[reportPrivateUsage] # see the class docstring + + @property + def resolved(self) -> "Resolved": + return self._core._resolved # pyright: ignore[reportPrivateUsage] # see the class docstring + + @property + def shaping(self) -> Shaping: + return self._core._shaping # pyright: ignore[reportPrivateUsage] # see the class docstring + + @property + def routing_cache(self) -> "DualCache": + return self._core._cache # pyright: ignore[reportPrivateUsage] # see the class docstring + + @property + def selectors(self) -> StrategySelectors: + return self._core._selectors # pyright: ignore[reportPrivateUsage] # see the class docstring + + @property + def endpoints(self) -> RustEndpoints: + return self._core._endpoints # pyright: ignore[reportPrivateUsage] # see the class docstring + + async def drive( + self, + plan: Plan, + request: Mapping[str, object], + overrides: Mapping[str, object] | None = None, + *, + passthrough: Callable[[Mapping[str, object], str], Awaitable[object]] | None = None, + prepared_already: bool = False, + ) -> object: + return await self._core._drive( # pyright: ignore[reportPrivateUsage] # see the class docstring + plan, request, overrides, passthrough=passthrough, prepared_already=prepared_already + ) + + def call_sync( + self, plan: Plan, request: Mapping[str, object], overrides: Mapping[str, object] | None = None + ) -> object: + return self._core._call_sync(plan, request, overrides) # pyright: ignore[reportPrivateUsage] # see the class docstring + + def routing_callbacks(self) -> Iterable[object]: + return self._core._checks.routing_callbacks() # pyright: ignore[reportPrivateUsage] # see the class docstring + + def resolve_alias(self, model: str) -> str | None: + return self._core._hosted._core_get_model_from_alias(model) # pyright: ignore[reportPrivateUsage] # see the class docstring + + def group_deployments(self, model_name: str) -> Sequence[Mapping[str, object]] | None: + return self._core._hosted._core_get_model_list(model_name=model_name) # pyright: ignore[reportPrivateUsage] # see the class docstring + + def routed_name(self, model: object) -> str: + return self._core.routed_name(model) + + def service_logged(self, call_type: str, start_time: float, end_time: float, request: Mapping[str, object]) -> None: + self._core.service_logged(call_type, start_time, end_time, request) + + async def set_response_headers( + self, response: object, model_group: str | None, request_kwargs: Mapping[str, object] + ) -> object: + return await self._core.set_response_headers(response, model_group, request_kwargs) + + def configured_deployment(self, model_id: str) -> "Deployment | None": + return self._core._hosted._core_get_deployment(model_id) # pyright: ignore[reportPrivateUsage] # see the class docstring + + def registry(self) -> RouterRegistry: + return self._core.registry() + + async def select( + self, + model: str, + request_kwargs: dict[str, object] | None = None, # mutable-ok: the RouterCore signature + *, + kind: str = "pick", + specific_deployment: bool | None = None, + messages: object = None, + ) -> Mapping[str, object] | Sequence[Mapping[str, object]]: + return await self._core.select( + model, request_kwargs, kind=kind, specific_deployment=specific_deployment, messages=messages + ) + + async def acompletion(self, *args: object, **kwargs: object) -> object: # kwargs-ok: endpoint payload + return await self._core.acompletion(*args, **kwargs) + + +class RustCore( + _SettingsSource, + RustLegacyMembers, + RustRegistryMembers, + RustRetries, + RustFallbacks, + RustHealthCallbacks, + RustOptionalChecks, +): + """A ``RouterCore`` served by the Rust router core over a ``PythonHost``. It holds no Python core.""" + + def __init__(self, config: RouterConfig, host: PythonHost, resolved: Resolved) -> None: + from litellm._service_logger import ServiceLogging + from litellm.router_utils.cooldown_cache import CooldownCache + + self._config: Final = config + self._host: Final = host + self.service_logger_obj: ServiceLogging = ServiceLogging() + apply_construction_effects(config, host) + self._resolved: Resolved = resolved + # `cooldown_time` as the Python core stores it (`cooldown_time or DEFAULT_COOLDOWN_TIME_SECONDS`, + # then each value assigned or updated, as given): what `settings()` reads back. + self._cooldown_given: object = config.get("cooldown_time") or DEFAULT_COOLDOWN_TIME_SECONDS + # Settings the Python core keeps as plain attributes, and settings deleted with `del`. + self._plain: Final = initial_plain(config) + self._deleted: frozenset[str] = frozenset() + self._request_timeout: Final = request_timeout(config) + legacy: Final[dict[str, object]] = {"deployment_stats": {}} # mutable-ok: the core's own legacy attributes + self._legacy: Final = legacy + self._health_view: RustHealth | None = None + self._cache: DualCache = host.routing_cache(config) + self._cooldowns: CooldownCache = CooldownCache( + cache=self._cache, + default_cooldown_time=self._cooldown_given, # pyright: ignore[reportArgumentType] # the constructor's value, as the Python core passes it + ) + self._health_checks: DeploymentHealthCache = resolved.health_routing.health_checks(self._cache) + self._version: int = 1 + self._settings: RustSettings | None = None + self._tag_filtering_models: tuple[int, bool] | None = None + # The registry indexes written through the deprecated attributes (`_core_`), and the + # document router-core last received for them. + self._index_written: frozenset[str] = frozenset() + self._index_sent: Mapping[str, object] | None = None + self._selectors: Final = StrategySelectors(host, self._cache) + self._selectors.init_default(resolved.routing_strategy, resolved.routing_strategy_args) + self._selectors.init_groups(_group_selectors(resolved.routing_groups)) + self._port: Final = _EndpointPort(self) + self._endpoints: Final = RustEndpoints(self._port, config) + # The resolver every picked deployment with `kubernetes_pod_discovery` goes through. + self._pod_discovery: KubernetesPodDiscovery = KubernetesPodDiscovery() + general: Final = config.get("router_general_settings", _DEFAULT_GENERAL_SETTINGS) + self._general_settings: object = general or RouterGeneralSettings() + self._registry_view: RouterRegistry | None = None + self._hosted: Final = HostRegistry( + host, + model_group_alias=dict(resolved.model_group_alias), # pyright: ignore[reportArgumentType] # the configured alias map + ignore_invalid_deployments=bool(config.get("ignore_invalid_deployments")), + on_change=self._sync_registry, + ) + self._hosted._core_cache = self._cache # pyright: ignore[reportPrivateUsage] # usage reads come from the cache the Rust core writes + ctx: Final = RoutingContext(self) # pyright: ignore[reportArgumentType] # the strategies read registry(), settings(), health() and select() + strategies: Final = PreRoutingStrategies( + self._hosted, host, ctx, config.get("auto_router_capability_limit"), self._cache + ) + self._hosted.strategies = strategies + self._prerouter: Final = PreRouter(strategies, self._hosted, host, ctx, self._cache) + # `register_plugin`: a custom strategy that replaces every pick (the host answers the + # Rust core's pre_route step with its deployment). + self._custom_strategy: CustomRoutingStrategyBase | None = None + self._custom_discovery: KubernetesPodDiscovery | None = None + # router-core builds the deployments, asking the host registry about each one: the handle + # starts empty, and is rebuilt over the built list for the settings the list's names feed. + self._handle_default_ids: tuple[str, ...] = () # the provider default ids the handle was built with + self._allowed_fails_default_sent: bool = False # what the handle's dynamic document last said + self._walked_fallbacks: Mapping[str, object] = {} # the fallback lists that document held + self._handle: RouterCoreHandle = self._new_handle([]) + self._hosted.build_runner = self._run_build + self._hosted.missing_attribute = self._no_router_attribute + self._hosted.changes_runner = self._apply_registry_changes + self._hosted.handle_lookup = self._live_handle + self._hosted.build(_as_list(config.get("model_list")) if config.get("model_list") is not None else None) + self._hosted.set_routing_groups(resolved.routing_groups) + track_live_registry(self._hosted) + self._hosted.routable_lookup = self._routable + self._hosted.max_parallel_requests_lookup = self._max_parallel_requests_client + self._hosted._core_routing_plugins = list(resolved.routing_plugins) # pyright: ignore[reportPrivateUsage] # the core fills its own host registry + self._rebuild() + self._join_callback_holders() + self._checks: Final = OptionalChecks(self, host, self._cache, self._hosted, ctx, config.arguments) + self._hosted.checks = self._checks + strategies.ensure_affinity = self._checks.ensure_deployment_affinity_callback + self._checks.build(config.get("optional_pre_call_checks"), config.get("model_list")) + self._slack_alerting_logger: SlackAlerting | None = initialize_alerting(host, config.get("alerting_config")) # pyright: ignore[reportArgumentType] # the argument as passed + if self._slack_alerting_logger is not None: + self._checks.slack_alerting_loggers = (self._slack_alerting_logger,) + if resolved.enable_mid_stream_fallback_continuation: + register_continuation_prefill_check(host) + if resolved.routing_strategy == "lar1": + from litellm.router_strategy.lar1_routing import apply_lar1_routing_strategy + + apply_lar1_routing_strategy(ctx, resolved.routing_strategy_args) # pyright: ignore[reportArgumentType] # the core's RoutingContext + + @property + def _registry_state(self) -> HostRegistry: + return self._hosted + + @property + def _python_host(self) -> PythonHost: + return self._host + + @property + def _written_retry_policy(self) -> object: + return self._plain["retry_policy"] + + @classmethod + def build(cls, config: RouterConfig, host: PythonHost) -> "RustCore": + from litellm.rust_bridge._native import RustBridgeDeclined + + extra: Final = unsupported_arguments(config) + if extra: + raise RustBridgeDeclined(f"router arguments not served by the Rust core: {sorted(extra)}") + RustEndpoints.prepare_host(config, host) + return cls(config, host, Resolved.of(config, host.dynamic())) + + @property + def enable_mid_stream_fallback_continuation(self) -> bool: + return self._resolved.enable_mid_stream_fallback_continuation + + def _set_continuation(self, enabled: bool) -> None: + """``_set_mid_stream_fallback_continuation``: an enabled flag registers the process-wide prefill check.""" + self._resolved = dataclasses.replace(self._resolved, enable_mid_stream_fallback_continuation=enabled) + self._settings = None + if enabled: + register_continuation_prefill_check(self._host) + + def __getattr__(self, name: str) -> object: + if name in self.__dict__.get("_deleted", ()): + raise self._no_attribute(name) + if name in SELECTOR_ATTRIBUTES.values(): + return self._selectors.attribute(name) + if name == "_core_async_get_available_deployment" and self.__dict__.get("_custom_strategy") is not None: + return self._custom_async_pick() + # Only reached for members this class does not define: the endpoints served through the + # host, else refused. Never forwarded anywhere. + endpoints: Final[object] = self.__dict__.get("_endpoints") + endpoint: Final = endpoints.member(name) if isinstance(endpoints, RustEndpoints) else None + if endpoint is not None: + return endpoint + attribute: Final = self._settings_attribute(name) if "_resolved" in self.__dict__ else MISSING + if attribute is not MISSING: + return attribute + raise RustCoreUnsupported(f"'{name}' is not served by the Rust router core", name=name, obj=self) + + @property + def _unset_attributes(self) -> frozenset[str]: + """The instance attributes ``del`` removed, which ``_settings_attribute`` no longer serves.""" + unset: Final[object] = self.__dict__.get("_unset_attributes") + return unset if isinstance(unset, frozenset) else frozenset() # pyright: ignore[reportUnknownVariableType] # only names are stored + + def _settings_attribute(self, name: str) -> object: + """A router attribute that reads a setting, the constructor's arguments or a constant.""" + if name in self._unset_attributes: + return MISSING + r: Final = self._resolved + resolved: Final[Mapping[str, object]] = { + "num_retries": r.read("num_retries"), + "max_fallbacks": r.read("max_fallbacks"), + "timeout": r.read("timeout"), + "retry_after": r.read("retry_after"), + "explicit_allowed_fails": r.given["allowed_fails"] + if r.explicit_allowed_fails is not None and "allowed_fails" in r.given + else r.explicit_allowed_fails, + "routing_strategy_args": r.read("routing_strategy_args"), + "request_timeout": self._request_timeout, + "default_priority": self._endpoints._default_priority, # pyright: ignore[reportPrivateUsage] # parsed by the endpoint plans + } + if name in resolved: + return resolved[name] + if name in self._plain: + return self._plain[name] + if name in self._legacy: + return self._legacy[name] + if name == "healthy_deployments": + return self._hosted._core_model_list # pyright: ignore[reportPrivateUsage] # the live list, as the core keeps it + if name == "get_model_group_io_token_usage": + return self._hosted.get_model_group_io_token_usage + if name == "scheduler": + return self._endpoints._get_scheduler() # pyright: ignore[reportPrivateUsage] # the endpoint plans' scheduler + if name == "chat": + from litellm.main import Chat + + self._legacy["chat"] = Chat(params=default_litellm_params(self._config), router_obj=self) # pyright: ignore[reportArgumentType] # Chat calls the router's endpoints, which this core serves + return self._legacy["chat"] + return argument_attribute(self._config, name) + + @property + def _shaping(self) -> Shaping: + return self._endpoints.shaping(self._resolved) + + # -- building -------------------------------------------------------------------------------- + + def _new_handle(self, model_list: Sequence[object]) -> "RouterCoreHandle": + from litellm.rust_bridge._native import RouterCoreHandle + + names: Final = [ + str(name) + for entry in model_list + if isinstance(name := (as_str_mapping(entry) or {}).get("model_name"), str) + ] + default_ids: Final = self._provider_default_ids() + self._handle_default_ids = default_ids + static: Final = { + **self._resolved.static_document(model_list), + **{name: value for name, value in failover_document(self._plain).items() if value is not None}, + **({"provider_default_deployment_ids": list(default_ids)} if default_ids else {}), + } + return RouterCoreHandle(static, self._dynamic_document(names)) + + def _dynamic_document(self, model_names: Sequence[str]) -> Mapping[str, object]: + self._allowed_fails_default_sent = self._resolved.allowed_fails_is_default() + # Settings are static: an in-place edit of a live list reaches neither the walk nor its notes. + resolved: Final = self._resolved + self._walked_fallbacks = { + "fallbacks": detached(resolved.fallbacks), + "context_window_fallbacks": detached(resolved.context_window_fallbacks), + "content_policy_fallbacks": detached(resolved.content_policy_fallbacks), + } + return resolved.dynamic_document(self._version, model_names) + + def _provider_default_ids(self) -> tuple[str, ...]: + """``provider_default_deployment_ids`` as the cooldown check reads it (``deployment in + settings().provider_default_deployment_ids``): the ids it holds now. router-core holds them + with its static settings.""" + handle: Final = self._live_handle() + return () if handle is None else tuple(handle.provider_default_ids()) + + def _pick_discovery(self, *, sync: bool) -> KubernetesPodDiscovery | None: + """The pod discovery a pick resolves through: the one bound when a custom strategy was + registered (``resolve_pods_after_bound``), else the router's own (``resolve_pods_after``); + none for an async pick assigned through the deprecated attribute, which Python calls bare.""" + if not sync and callable(self.__dict__.get("_core_async_get_available_deployment")): + return None + if self._custom_strategy is not None and self._custom_discovery is not None: + return self._custom_discovery + return self._pod_discovery + + # -- the plugin context (``RoutingContext``) ----------------------------------------------- + + def _core_arm_routing_read_prefetch(self, model: str, request_kwargs: Mapping[str, object] | None = None) -> None: + """``RoutingContext.arm_read_prefetch``: declare the cooldown reads (and, for + usage-based-routing-v2, the usage reads) the pick for ``model`` makes on the request's + Redis batch, so admission's flush carries them. A miss costs nothing.""" + from litellm.router_utils.routing_read_batch import RoutingPrefetch + + try: + usage_selector: Final = self._prefetch_usage_selector(model, request_kwargs) + deployments: Final = self._hosted._core_get_model_list(model_name=model) # pyright: ignore[reportPrivateUsage] # see __init__ + if deployments: + RoutingPrefetch.arm(RoutingContext(self), usage_selector, deployments) # pyright: ignore[reportArgumentType, reportUnknownMemberType] # the prefetch reads registry() and health(); its deployments list is untyped + except Exception as e: # noqa: BLE001 # a prefetch is an optimisation, never a reason to fail the request + from litellm._logging import verbose_router_logger + + verbose_router_logger.debug( + "routing read prefetch not armed for %s: %s", + str(model).replace("\r", "").replace("\n", ""), + str(e).replace("\r", "").replace("\n", ""), + ) + + def _prefetch_usage_selector(self, model: str, request_kwargs: Mapping[str, object] | None) -> object: + """``_get_routing_context``'s selector when it is a usage-based-routing-v2 one: the request's + strategy override, else the routing group ``model`` resolves into, else the router's own.""" + from litellm.router_strategy.lowest_tpm_rpm_v2 import LowestTPMLoggingHandler_v2 + + strategy, selector = self._prefetch_strategy_and_selector(model, request_kwargs or {}) + if strategy == "usage-based-routing-v2" and isinstance(selector, LowestTPMLoggingHandler_v2): + return selector + return None + + def _prefetch_strategy_and_selector(self, model: str, request: Mapping[str, object]) -> tuple[str | None, object]: + raw: Final = request.get("routing_strategy") + override: Final = normalize_strategy(raw) if isinstance(raw, str | RoutingStrategy) else None + if override is not None and override in {"simple-shuffle", *SELECTOR_ATTRIBUTES}: + return override, self._selectors.selector_for(override, "override", None, request) + groups: Final = self._hosted._core_routing_groups # pyright: ignore[reportPrivateUsage] # see __init__ + resolved: Final = self._hosted._core_get_model_from_alias(model=model) or model # pyright: ignore[reportPrivateUsage] # see __init__ + group_name: Final = ( + resolved + if self._hosted.get_routing_group(resolved) is not None + else next( + ( + group.group_name + for group in groups.values() + if group.routing_strategy != "priority" and resolved in group.models + ), + None, + ) + ) + if group_name is None: + return normalize_strategy(self._resolved.routing_strategy), self._selectors.default + if groups[group_name].routing_strategy == "priority": + return None, None + group_strategy: Final = normalize_strategy(groups[group_name].routing_strategy) + return group_strategy, self._selectors.selector_for(group_strategy or "", "group", group_name, None) + + async def _core_async_pre_routing_hook( + self, + model: str, + request_kwargs: dict[str, object], # mutable-ok: the request kwargs the hook stamps + messages: object = None, + input: object = None, + specific_deployment: bool | None = False, + ) -> "PreRoutingHookResponse | None": + """``RoutingContext.pre_routing_hook``: the pre-routing strategies' answer for ``model`` + (the tier group and its params), with the routing plugins run first.""" + r: Final = self._resolved + owner: Final[object] = self.__dict__.get("_core_owner") + public: Final[object] = owner() if callable(owner) else None + return await self._prerouter._hook( # pyright: ignore[reportPrivateUsage] # the core's own pre-router + model, + request_kwargs, + messages=messages, + input=input, + fallback_depth=0, + enable_tag_filtering=r.enable_tag_filtering, + match_any=r.tag_filtering_match_any, + public_router=public if public is not None else self, + run_plugins=partial(self._run_plugins, request_kwargs, messages) + if self._hosted._core_routing_plugins # pyright: ignore[reportPrivateUsage] # the live plugin list + else None, + ) + + def _core_log_retry( + self, + kwargs: dict[str, object], # mutable-ok: the request kwargs the record joins + e: Exception, + ) -> dict[str, object]: # mutable-ok: the same kwargs, as Python's returns them + """``RoutingContext.log_retry``.""" + return log_retry_into(kwargs, e, self._deployment_params_base_model) + + def _deployment_params_base_model(self, deployment_id: str) -> str | None: + """``_deployment_params_base_model``: the ``base_model`` a deployment declares under ``litellm_params``.""" + deployment: Final = self._hosted._core_get_deployment(model_id=deployment_id) # pyright: ignore[reportPrivateUsage] # see __init__ + if deployment is None: + return None + base_model: Final = (deployment.litellm_params.model_extra or {}).get("base_model") + return base_model if isinstance(base_model, str) else None + + def _sync_routing_state(self) -> None: + """Before a pick: sync the routing state that callers can change in place, which the Python + core reads live on every pick. This exists only for the deprecated writable attributes and + goes away with them. + + Only these cases send anything; otherwise both checks return without touching router-core:: + + router.model_names = {...} # assigning syncs at once (__setattr__) + router.model_names.add("new-group") # an in-place change is picked up here, on the next pick + router.model_name_to_deployment_indices["g"].append(3) # likewise + + router.routing_plugins.append(plugin) # first plugin: rebuilds the handle with the plugin stage + router.routing_plugins.clear() # list emptied: rebuilds without it + + Until one of the index attributes has been assigned, ``_sync_index`` sends nothing. After + that it sends the index document only when it differs from the last one it sent. + """ + self._sync_plugin_stage() + self._sync_index() + self._sync_allowed_fails() + + def _sync_allowed_fails(self) -> None: + """``_is_allowed_fails_set_on_router`` compares with the live ``litellm.allowed_fails`` on + every failure, so a write to the global can change the rule; ``_settings`` stays as copied.""" + if self._resolved.allowed_fails_is_default() != self._allowed_fails_default_sent: + names: Final = [str(d.get("model_name")) for d in self._handle.deployments()] + self._handle.apply({"settings": self._dynamic_document(names)}) + + def _sync_index(self) -> None: + """Hand router-core the registry indexes written through the deprecated attributes + (``model_names``, ``model_name_to_deployment_indices``, ``team_model_to_deployment_indices``) + as the host registry holds them now (its own deployment changes update the written objects, + as the Python core's do), whenever they differ from what it holds.""" + if not self._index_written: + return + document: Final = index_document(self._hosted, self._index_written) + if document != self._index_sent: + _ = self._handle.apply({"index": document}) + self._index_sent = document + + def _sync_plugin_stage(self) -> None: + """Give router-core the routing-plugin stage when the live plugin list (``routing_plugins``, + which callers may assign or append to) gains its first plugin, and drop it when the list + empties: the Python core checks the list on every pick.""" + live: Final = _as_list(self._hosted._core_routing_plugins) # pyright: ignore[reportPrivateUsage] # the live plugin list + if bool(live) == bool(self._resolved.routing_plugins): + return + self._resolved = dataclasses.replace(self._resolved, routing_plugins=tuple(live)) # pyright: ignore[reportArgumentType] # the plugins as listed + self._rebuild() + + def _rebuild(self) -> None: + """A new Rust core for changed static settings, over the same deployments and the + providers the old one's builds resolved for them.""" + old: Final = self._handle + model_list: Final = list(self._hosted._core_model_list) # pyright: ignore[reportPrivateUsage] # see __init__ + self._handle = self._new_handle(model_list) + self._handle.adopt_providers(old) + old.retire() + # The live list and indexes read the new handle's (the same entries, kept as the same objects). + self._hosted.refresh() + # The new handle builds its indexes from the deployments: written ones are handed again. + self._index_sent = None + + def _run_build( + self, + entries: Sequence[object], + path: str, + resolve_secrets: bool, + install: BuildInstall, + answer: Callable[[Mapping[str, object]], object], + ) -> None: + """Drive one router-core deployment build (``RouterCoreHandle.build``), answering each of + its steps with ``answer``; router-core installs each deployment as it accepts it, and the + host registry reads the list and indexes back before answering the next step.""" + call: Final = self._handle.build(_plain(list(entries)), path, resolve_secrets, install) # pyright: ignore[reportArgumentType] # plain JSON entries + outcome: object = None # rebind-ok: each step's answer + while True: + step = call.next(outcome) + self._hosted.refresh() + if step["kind"] == "done": + break + if step["kind"] == "failed": + raise ValueError(str(step.get("error"))) + outcome = _plain(answer(step)) + + def _live_handle(self) -> "RouterCoreHandle | None": + """The handle the live model_list reads, once the constructor made one.""" + from litellm.rust_bridge._native import RouterCoreHandle + + handle: Final[object] = self.__dict__.get("_handle") + return handle if isinstance(handle, RouterCoreHandle) else None + + def _apply_registry_changes(self, changes: Mapping[str, object]) -> None: + """Apply a host registry write router-core's build did not make (a removal is a delete, a + swapped list a replace) as a changeset.""" + _ = self._handle.apply(_plain(dict(changes))) # pyright: ignore[reportArgumentType] # a plain JSON changeset + + def _sync_registry(self) -> None: + """After a host registry write: router-core already holds it (the build installed what it + added, a removal came as a changeset), so only the settings it feeds move.""" + from litellm.rust_bridge._native import RouterCoreHandle + + self._tag_filtering_models = None + handle: Final[object] = self.__dict__.get("_handle") + # A wildcard deployment added an id the cooldown check skips. + if isinstance(handle, RouterCoreHandle) and self._provider_default_ids() != self._handle_default_ids: + self._rebuild() + + # -- entry points ---------------------------------------------------------------------------- + + def registry(self) -> RouterRegistry: + """The registry view, built once (``registry()`` returns the same object each time).""" + if self._registry_view is None: + self._registry_view = hosted_registry_view(self._hosted) + return self._registry_view + + def settings(self) -> RustSettings: + """A cached snapshot, replaced when a changeset changes a setting.""" + if self._settings is None: + self._settings = self._settings_snapshot() + return self._settings + + def _settings_snapshot(self) -> RustSettings: + r: Final = self._resolved + snapshot: Final = RustSettings( + routing_strategy=r.routing_strategy, + fallbacks=r.read("fallbacks"), # pyright: ignore[reportArgumentType] # the value as given, as the Python core stores it + context_window_fallbacks=r.read("context_window_fallbacks"), # pyright: ignore[reportArgumentType] # the value as given, as the Python core stores it + content_policy_fallbacks=r.read("content_policy_fallbacks"), # pyright: ignore[reportArgumentType] # the value as given, as the Python core stores it + default_deployment=r.default_deployment, + default_litellm_params=dict(r.default_litellm_params) + if "default_litellm_params" in r.assigned + else default_litellm_params(self._config), + default_max_parallel_requests=self._endpoints.default_max_parallel_requests, # pyright: ignore[reportArgumentType] # the setting as given + cooldown_time=self._cooldown_given, # pyright: ignore[reportArgumentType] # the value as given, as the Python core stores it + allowed_fails=r.read("allowed_fails"), # pyright: ignore[reportArgumentType] # the value as given, as the Python core stores it + allowed_fails_policy=r.allowed_fails_policy, # pyright: ignore[reportArgumentType] # the policy as passed + disable_cooldowns=r.read("disable_cooldowns"), # pyright: ignore[reportArgumentType] # the value as given, as the Python core stores it + enable_pre_call_checks=r.enable_pre_call_checks, + enable_health_check_routing=r.health_routing.enable_health_check_routing, + enable_tag_filtering=r.read("enable_tag_filtering"), # pyright: ignore[reportArgumentType] # the value as given, as the Python core stores it + tag_filtering_match_any=r.read("tag_filtering_match_any"), # pyright: ignore[reportArgumentType] # the value as given, as the Python core stores it + tag_routing_prefix=r.read("tag_routing_prefix"), # pyright: ignore[reportArgumentType] # the value as given, as the Python core stores it + fallback_access_check=r.fallback_access_check, # pyright: ignore[reportArgumentType] # the check as passed + fallback_budget_check=r.fallback_budget_check, # pyright: ignore[reportArgumentType] # the check as passed + cache_responses=self._endpoints.cache_responses, + search_tools=self._endpoints.search_tools, # pyright: ignore[reportArgumentType] # the configured entries, as passed + guardrail_list=self._endpoints.guardrail_list, # pyright: ignore[reportArgumentType] # same + provider_default_deployment_ids=self._hosted._core_provider_default_deployment_ids, # pyright: ignore[reportPrivateUsage] # the host registry's live list + health_check_ignore_transient_errors=r.health_routing.health_check_ignore_transient_errors, + background_health_check_model_groups=r.health_routing.background_health_check_model_groups, + router_general_settings=self._general_settings, # pyright: ignore[reportArgumentType] # the setting as given + provider_budget_config=self._checks.provider_budget_config, # pyright: ignore[reportArgumentType] # the budget limiter's live config + _router=self, # pyright: ignore[reportArgumentType] # the snapshot reads as_dict from this core + kubernetes_pod_discovery=self._pod_discovery, + ) + if not self._deleted: + return snapshot + # A deleted setting reads as unset. + return dataclasses.replace( + snapshot, + **dict.fromkeys(self._deleted.intersection(SETTING_FIELDS), None), # pyright: ignore[reportArgumentType] # settings fields only: an alias map or retry count is deletable too + _unset=snapshot._unset | self._deleted, # pyright: ignore[reportPrivateUsage] # the snapshot's own unset set + ) + + def available_guardrail(self, name: str) -> object: + return self._endpoints.available_guardrail(name) + + def settings_dict(self) -> Mapping[str, object]: + r: Final = self._resolved + latency: Final = self._selectors.latency_args() if r.routing_strategy == "latency-based-routing" else None + # `get_settings` reports the latency selector's full routing args. + args: Final = ( + latency if latency is not None else r.given.get("routing_strategy_args", dict(r.routing_strategy_args)) + ) + values: Final = { + "routing_strategy_args": args, + "routing_strategy": r.routing_strategy, + "allowed_fails": r.read("allowed_fails"), + "cooldown_time": self._cooldown_given, + "num_retries": r.read("num_retries"), + "timeout": r.read("timeout"), + "retry_after": r.read("retry_after"), + "fallbacks": r.read("fallbacks"), + "context_window_fallbacks": r.read("context_window_fallbacks"), + "model_group_alias": dict(r.model_group_alias), + "enable_mid_stream_fallback_continuation": r.enable_mid_stream_fallback_continuation, + "enable_tag_filtering": r.read("enable_tag_filtering"), + "tag_routing_prefix": r.read("tag_routing_prefix"), + "routing_groups": group_settings(r.routing_groups), + } + return settings_view(values, self._plain, self._deleted) + + def health(self) -> RustHealth: + """The health view, built once.""" + if self._health_view is None: + self._health_view = RustHealth(self) + return self._health_view + + def observe(self, logging_obj: "LiteLLMLogging | None" = None, *, request_scoped: bool = False) -> RoutingObservers: + """The routing observers: the Slack alerting logger ``alerting_config`` built, if any. + ``request_scoped=True`` takes the router's callbacks off the host lists; passing a request's + ``logging_obj`` binds its selectors and policies to it. This core registers no outcome + callbacks, so it binds none either: router-core reads what they would do for the request + (``OptionalChecks.outcome_callbacks``).""" + if request_scoped: + self._serve_through_request_logging() + if logging_obj is not None: + self._checks.bind_to_request_logging(logging_obj) + return RoutingObservers(slack_alerting_logger=self._slack_alerting_logger) + + def register_plugin(self, plugin: "CustomRoutingStrategyBase") -> None: + """``_core_set_custom_routing_strategy``: ``plugin`` picks every deployment from now on, until + the routing strategy changes.""" + self._custom_strategy = plugin + # `resolve_pods_after_bound`: its picks resolve pods through the discovery bound now. + self._custom_discovery = self._pod_discovery + # Python assigns the strategy's picks over its own: a caller of the pick reaches the + # strategy with the arguments it passed and no others. + self.__dict__["get_available_deployment"] = self._custom_sync_pick(plugin) + # It replaces an async pick assigned through the deprecated attribute, as Python's does. + if "_core_async_get_available_deployment" in self.__dict__: + del self.__dict__["_core_async_get_available_deployment"] + + def _custom_sync_pick(self, plugin: "CustomRoutingStrategyBase") -> object: + discovery: Final = self._custom_discovery + pick: Final = plugin.get_available_deployment + return pick if discovery is None else resolve_pods_after_bound(pick, discovery) + + def _custom_async_pick(self) -> object: + """``_core_set_custom_routing_strategy``'s ``_core_async_get_available_deployment``.""" + custom: Final = self._custom_strategy + discovery: Final = self._custom_discovery + pick: Final = custom.async_get_available_deployment # pyright: ignore[reportOptionalMemberAccess] # __getattr__ checked it + return pick if discovery is None else async_resolve_pods_after_bound(pick, discovery) + + def set_custom_routing_strategy(self, CustomRoutingStrategy: "CustomRoutingStrategyBase") -> None: + from litellm.router_deprecated import ( + _note, # pyright: ignore[reportPrivateUsage] # the deprecated alias notes its use like Python's + ) + + _note("set_custom_routing_strategy", "register_plugin(strategy)") + self.register_plugin(CustomRoutingStrategy) + + def retire(self) -> None: + """``_core_discard``: only the router's global registrations go; it keeps serving calls.""" + untrack_live_registry(self._hosted) + self._retire_callbacks() + if self._hosted.strategies is not None: + self._hosted.strategies.retire() + + def apply(self, changes: ChangeSet) -> ApplyResult: + """``apply_changes``: assigned settings, then ``update_settings``, upserts, deletes and the + Redis cache, in that order. An unknown name refuses the changeset; a setting + ``update_settings`` raises for ends it there, with the settings before it applied.""" + from litellm.router_api import ASSIGNABLE_ALIASES, PLAIN_SETTING_ATTRIBUTES, SETTING_FIELDS + + if changes.refresh_model_info: + raise TypeError("ChangeSet(refresh_model_info=True) needs aapply()") + unknown: Final = sorted(set(changes.assign) - {*SETTING_FIELDS, *PLAIN_SETTING_ATTRIBUTES, *ASSIGNABLE_ALIASES}) + if unknown: + raise ValueError(f"not assignable router settings: {unknown}") + unserved: Final = sorted(set(changes.assign) - _STATIC_SETTINGS - _DYNAMIC_SETTINGS - _HOST_SETTINGS) + if unserved: + raise RustCoreUnsupported(f"the Rust router core does not serve the setting(s) {', '.join(unserved)} yet") + update: Final = read_settings_update(changes.settings, self._resolved.routing_strategy) + # `apply_changes` assigns before it runs `update_settings`: a budget config assigned beside + # `optional_pre_call_checks` builds the budget limiter with the new config. + self._assign_host_settings(changes.assign) + if update.optional_pre_call_checks is not MISSING: + self._checks.set(update.optional_pre_call_checks) + if update.continuation is not MISSING: + self._set_continuation(continuation_flag(update.continuation)) + self._apply_settings(changes.assign, update) + if update.failure is not None: + raise update.failure + # Upserts first, then deletes (`apply_changes`), through the host registry, which hands + # each change to the Rust core. + upserted: Final = [self._hosted._core_upsert_deployment(deployment=d) for d in changes.upsert] # pyright: ignore[reportPrivateUsage] # see __init__ + deleted: Final = [self._hosted._core_delete_deployment(id=i) for i in changes.delete] # pyright: ignore[reportPrivateUsage] # see __init__ + if changes.redis_cache is not None: + # `_core_update_redis_cache`: the routing state's cache gains the Redis cache. + self._cache.attach_redis_cache(changes.redis_cache) + self._checks.prompt_caching_pin_cache.attach_redis_cache(changes.redis_cache) + return ApplyResult( + version=self._hosted._registry_version, # pyright: ignore[reportPrivateUsage] # see __init__ + upserted=upserted, + deleted=deleted, # pyright: ignore[reportArgumentType] # the removed model_list entries, as the core returns them + ) + + async def aapply(self, changes: ChangeSet) -> ApplyResult: + """``apply``, then the model info refresh, which needs I/O.""" + result: Final = self.apply(dataclasses.replace(changes, refresh_model_info=False, refresh_client=None)) + if changes.refresh_model_info: + await self._hosted.arefresh_model_info(changes.refresh_client) + return result + + def _apply_settings(self, assign: Mapping[str, object], update: SettingsUpdate) -> None: + values: Final = { + **without(assign, *_HOST_SETTINGS), + **update.values, + } + given: Final = dict(update.values) # the settings passed through `update_settings` + groups_input: Final[object] = ( + update.routing_groups if update.routing_groups is not MISSING else self._resolved.routing_groups_input + ) + # What `_core_update_settings` runs after its loop, which a setting that raised never reaches. + settled: Final = update.failure is None + groups_failure: Exception | None = None # rebind-ok: set when the groups passed do not parse + if settled and (update.routing_groups is not MISSING or self._changes_strategy(update.values)): + # `update_settings` re-inits the routing groups from the groups passed, else (on a + # strategy change) from the groups last passed, whatever a raw write assigned since. + try: + values["routing_groups"] = parse_groups( + groups_input, + frozenset(str(d.get("model_name")) for d in self._handle.deployments()), + frozenset(self._resolved.model_group_alias), + ) + except Exception as raised: # noqa: BLE001 # raised below: `_init_routing_groups` runs after the loop + groups_failure = raised # rebind-ok: None until the groups fail + else: + given["routing_groups"] = groups_input + self._store_plain(update.plain) + if self._deleted & set(values): + self._deleted = self._deleted - set(values) + self._settings = None + if values: + previous: Final = self._resolved + resolved: Final = dataclasses.replace( + previous.changed(values, frozenset(given)), + routing_groups_input=groups_input if "routing_groups" in values else previous.routing_groups_input, + ) + static_changed: Final = bool(_STATIC_SETTINGS & set(values)) + self._resolved = resolved + if set(values) & set(SETTING_FIELDS): + self._settings = None # the snapshot holds settings fields, not the plain attributes + if "model_group_alias" in values: + self._hosted._core_model_group_alias = dict(resolved.model_group_alias) # pyright: ignore[reportPrivateUsage, reportAttributeAccessIssue] # see __init__ + self._apply_strategy(previous, resolved, given, settled=settled) + if "routing_groups" in values: + self._settings = None + self._hosted.set_routing_groups(resolved.routing_groups) + if "cooldown_time" in values: + self._cooldown_given = values["cooldown_time"] + if values["cooldown_time"] is None: + # `add_deployment_to_cooldown(cooldown_time=None)` cools down for the cooldown + # cache's default, which keeps the constructor's value. + self._resolved = dataclasses.replace( + resolved, cooldown_time=float(self._cooldowns.default_cooldown_time) + ) + if static_changed: + self._rebuild() + else: + names: Final = [str(d.get("model_name")) for d in self._handle.deployments()] + self._handle.apply({"settings": self._dynamic_document(names)}) + if groups_failure is not None: + raise groups_failure + + def _store_plain(self, plain: Mapping[str, object]) -> None: + """Settings the Python core keeps as plain attributes, read from the next call on.""" + flags_before: Final = failover_document(self._plain) + self._plain.update(plain) + if failover_document(self._plain) != flags_before: + self._rebuild() + policies: Final = {name: plain[name] for name in _RETRY_POLICY_SETTINGS if name in plain} + if policies: + self._set_retry_policies(policies) + + def _changes_strategy(self, updated: Mapping[str, object]) -> bool: + """``update_settings`` names a ``routing_strategy`` other than the router's.""" + if "routing_strategy" not in updated: + return False + return _strategy_name(updated["routing_strategy"]) != self._resolved.routing_strategy + + def _assign_host_settings(self, assign: Mapping[str, object]) -> None: + """The assigned settings only the host reads, each stored as given (detached, as + ``apply_changes`` stores it) where its readers find it from the next call on.""" + given: Final = {name: detached(value) for name, value in assign.items() if name in _HOST_SETTINGS} + if not given: + return + for name, value in given.items(): + if name in ("enable_health_check_routing", "background_health_check_model_groups"): + self._resolved = dataclasses.replace( + self._resolved, + health_routing=dataclasses.replace(self._resolved.health_routing, **{name: value}), # pyright: ignore[reportArgumentType] # stored as given, as the Python core stores it + ) + elif name == "provider_default_deployment_ids": + # The list wildcard deployments added later are appended to, as on the Python core. + self._hosted._core_provider_default_deployment_ids = value # pyright: ignore[reportPrivateUsage, reportAttributeAccessIssue] # stored as given, as the Python core stores it + elif name == "cache_responses": + self._endpoints.set_cache_responses(bool(value)) + elif name == "provider_budget_config": + # Read when the budget limiter is built; a limiter built already keeps its config. + self._checks.provider_budget_config = value + elif name == "health_check_ignore_transient_errors": + self._resolved = dataclasses.replace( + self._resolved, + health_routing=dataclasses.replace( + self._resolved.health_routing, + health_check_ignore_transient_errors=value, # pyright: ignore[reportArgumentType] # stored as given, as the Python core stores it + ), + ) + elif name == "router_general_settings": + self._general_settings = value + elif name == "kubernetes_pod_discovery": + self._pod_discovery = value # pyright: ignore[reportAttributeAccessIssue] # stored as given, as the Python core stores it + else: + self._endpoints.assign_setting(name, value) + self._deleted = self._deleted - set(given) + self._settings = None + if {"enable_health_check_routing", "background_health_check_model_groups"} & set(given) or ( + self._provider_default_ids() != self._handle_default_ids + ): + self._rebuild() + + def _set_retry_policies(self, policies: Mapping[str, object]) -> None: + """The ``retry_policy`` and ``model_group_retry_policy`` the retry loop reads from then on, + from ``update_settings`` or an attribute write: router-core holds them with its static + settings.""" + previous: Final = self._resolved + resolved: Final = dataclasses.replace( + previous, + retry_policy=readable_retry_policy(policies.get("retry_policy", previous.retry_policy)), + model_group_retry_policy=policies.get("model_group_retry_policy", previous.model_group_retry_policy), + ) + if resolved != previous: + self._resolved = resolved + self._rebuild() + + def _apply_strategy( + self, previous: Resolved, resolved: Resolved, values: Mapping[str, object], *, settled: bool + ) -> None: + """``_core_update_settings``: a changed strategy re-inits the selectors with the args passed + alongside it (or none); changed args alone rebuild the default selector. What runs after + the loop needs ``settled``: no setting raised, and the groups are in ``values``.""" + if "routing_strategy" in values and resolved.routing_strategy != previous.routing_strategy: + if resolved.routing_strategy == "lar1": + # `_switch_routing_strategy`: LAR-1 is a custom strategy, registered in place of + # `routing_strategy_init`, so the selectors the router held stay as they were. + from litellm.router_strategy.lar1_routing import apply_lar1_routing_strategy + + apply_lar1_routing_strategy( + RoutingContext(self), # pyright: ignore[reportArgumentType] # the core's RoutingContext + as_str_mapping(values.get("routing_strategy_args")), + ) + else: + # `routing_strategy_init` drops a custom strategy (`_reset_custom_routing_strategy`), + # and an async pick assigned through the deprecated attribute. + self._custom_strategy = None + self.__dict__.pop("get_available_deployment", None) + if "_core_async_get_available_deployment" in self.__dict__: + del self.__dict__["_core_async_get_available_deployment"] + self._selectors.init_default( + resolved.routing_strategy, as_str_mapping(values.get("routing_strategy_args")) or {} + ) + if "routing_groups" in values: + self._selectors.init_groups(_group_selectors(resolved.routing_groups)) + elif "routing_groups" in values: + self._selectors.init_groups(_group_selectors(resolved.routing_groups)) + elif ( + settled + and "routing_strategy_args" in values + and resolved.routing_strategy_args != previous.routing_strategy_args + ): + if resolved.routing_strategy == "lar1": + from litellm.router_strategy.lar1_routing import apply_lar1_routing_strategy + + apply_lar1_routing_strategy(RoutingContext(self), resolved.routing_strategy_args) # pyright: ignore[reportArgumentType] # the core's RoutingContext + return + self._selectors.rebuild_default(resolved.routing_strategy_args) + + def upsert_deployment(self, deployment: Deployment) -> Deployment | None: + from litellm.router_deprecated import ( + _note, # pyright: ignore[reportPrivateUsage] # the router's deprecation notice + ) + + _note("upsert_deployment", "apply(ChangeSet(upsert=[deployment]))") + return self.apply(ChangeSet(upsert=[deployment])).upserted[0] + + def delete_deployment(self, id: str) -> Deployment | None: + from litellm.router_deprecated import _note # pyright: ignore[reportPrivateUsage] # see above + + _note("delete_deployment", "apply(ChangeSet(delete=[id]))") + return self.apply(ChangeSet(delete=[id])).deleted[0] + + def update_settings(self, **kwargs: object) -> None: # kwargs-ok: the deprecated signature, unchanged + from litellm.router_deprecated import _note # pyright: ignore[reportPrivateUsage] # see upsert_deployment + + _note("update_settings", "apply(ChangeSet(settings={...}))") + _ = self.apply(ChangeSet(settings=kwargs)) + + def update_redis_cache(self, cache: "RedisCache") -> None: + from litellm.router_deprecated import _note # pyright: ignore[reportPrivateUsage] # see upsert_deployment + + _note("update_redis_cache", "apply(ChangeSet(redis_cache=cache))") + _ = self.apply(ChangeSet(redis_cache=cache)) + + def get_model_list( + self, model_name: str | None = None, team_id: str | None = None + ) -> "list[DeploymentTypedDict] | None": # mutable-ok: the old return type, unchanged + from litellm.router_deprecated import _note # pyright: ignore[reportPrivateUsage] # see upsert_deployment + + _note("get_model_list", "registry().deployments(model_name, team_id=...)") + return self._hosted._core_get_model_list(model_name=model_name, team_id=team_id) # pyright: ignore[reportPrivateUsage, reportReturnType] # the host registry's listing, as the deprecated member returns it + + @staticmethod + def get_valid_args() -> list[str]: # mutable-ok: the old return type, a fresh list + """The arguments ``Router.__init__`` accepts.""" + return list(ROUTER_INIT_ARGS) + + async def arefresh_model_info(self, *, client: "AsyncHTTPHandler | None" = None) -> None: + from litellm.router_deprecated import _note # pyright: ignore[reportPrivateUsage] # see upsert_deployment + + _note("arefresh_model_info", "aapply(ChangeSet(refresh_model_info=True))") + _ = await self.aapply(ChangeSet(refresh_model_info=True, refresh_client=client)) + + def flush_cache(self) -> None: + """Drop the host's response cache and the routing state.""" + self._host.set_global_cache(None) + self._cache.flush_cache() + self._checks.prompt_caching_pin_cache.in_memory_cache.flush_cache() + + def reset(self) -> None: + """Empty the host's success and failure callback lists (every entry, not only this core's), + drop the retry policy and flush the caches.""" + for name in ("success_callback", "_async_success_callback", "failure_callback", "_async_failure_callback"): + self._host.set_callback_list(name, []) + self._plain["retry_policy"] = None + self._set_retry_policies({"retry_policy": None}) + self.flush_cache() + + def __delattr__(self, name: str) -> None: + # `del router.` unsets it: the settings view reports it unset from then on. + from litellm.router_api import SETTING_FIELDS + + # The deprecated attribute deletes the core's own name (``_core_``). + setting: Final = name.removeprefix("_core_") + if setting not in SETTING_FIELDS and setting not in _DELETABLE_ATTRIBUTES: + if name in self.__dict__ or name in self._unset_attributes: + object.__delattr__(self, name) + return + if self._settings_attribute(name) is MISSING: + object.__delattr__(self, name) + return + # An attribute the Python core keeps on the instance (`alerting_config`, an argument + # as passed): reading it raises AttributeError from now on. + self.__dict__["_unset_attributes"] = self._unset_attributes | {name} + return + if setting in self._deleted: + raise self._no_attribute(setting) + self._deleted = self._deleted | {setting} + self._settings = None + if setting == "model_group_alias": + del self._hosted._core_model_group_alias # pyright: ignore[reportPrivateUsage] # every read of it raises from now on, as on the Python core + elif setting == "scheduler": + self._endpoints.delete_scheduler() + + def _no_attribute(self, name: str) -> AttributeError: + """What reading an attribute deleted from the core raises (the facade renames it).""" + return AttributeError(f"'{type(self).__name__}' object has no attribute '{name}'", name=name, obj=self) + + def _no_router_attribute(self, name: str) -> AttributeError: + """``_no_attribute`` for a read through ``registry()``, which the deprecated attribute + renames when it is the router's own.""" + owner: Final[object] = self.__dict__.get("_core_owner") + public: Final[object] = owner() if callable(owner) else None + return AttributeError( + f"'{type(self).__name__}' object has no attribute '{name}'", + name=name, + obj=self if public is None else public, + ) + + def _refuse_deleted_reads(self, request: Mapping[str, object]) -> None: + """What a call reads of the router before any attempt, once deleted: routing reads + ``_core_model_group_alias``, and ``async_function_with_retries`` reads ``self.num_retries`` + when the request names none. Every hop raises it, so no fallback serves the call.""" + if "model_group_alias" in self._deleted: + raise self._no_attribute("_core_model_group_alias") + if "num_retries" in self._deleted and request.get("num_retries") is None: + raise self._no_attribute("num_retries") + + def __setattr__(self, name: str, value: object) -> None: + # A setting write is a changeset (`router.num_retries = 2` is `apply(ChangeSet(assign=...))`), + # a `model_list` write swaps the live list, and any other attribute is kept as written. + from litellm.router_api import ASSIGNABLE_ALIASES, PLAIN_SETTING_ATTRIBUTES, SETTING_FIELDS + + if name in PLAIN_SETTING_ATTRIBUTES or name in SETTING_FIELDS or name in ASSIGNABLE_ALIASES: + _ = self.apply(ChangeSet(assign={name: value})) + elif name in PLAIN_DEFAULTS or name == "max_retries": + # Plain attributes on the Python core, outside the setting lists: stored as given. + self._store_plain({name: value}) + elif name == "scheduler": + self._endpoints.set_scheduler(value) + self._deleted = self._deleted - {name} + elif name in ("model_list", "_core_model_list"): + # `router.model_list = [...]` reaches the core as `_core_model_list` through the shim. + self._assign_model_list(_as_list(value)) + elif name in _HOSTED_STATE: + # A deprecated attribute write of the registry state: it lands where this core keeps it. + setattr(self._hosted, name, value) + if name in INDEX_ATTRIBUTES: + self._index_written = self._index_written | {name} + self._sync_index() + elif name == "_core_routing_groups": + # The Python core routes on the assigned groups from the next call on, while its + # member map and group selectors stay as they were; router-core holds the groups and + # the member map with its static settings. + assigned: Final = as_str_mapping(value) or {} + previous: Final = self._resolved + self._resolved = dataclasses.replace( + previous, + routing_groups=tuple(g for g in assigned.values() if isinstance(g, RoutingGroup)), + routing_group_members=previous.routing_group_members + if previous.routing_group_members is not None + else member_map(previous.routing_groups), + ) + self._settings = None + self._rebuild() + elif name == "_core_routing_plugins": + self._sync_plugin_stage() + elif name == "_core_prompt_caching_pin_cache": + # Read when a prompt caching check is built, by `update_redis_cache` and by `flush_cache`. + self._checks.prompt_caching_pin_cache = value # pyright: ignore[reportAttributeAccessIssue] # stored as given, as the Python core stores it + elif name in _REBOUND_STATE: + # A deprecated attribute write of the routing state (``router.cache = ...``): this core + # reads the new object from now on, as the Python core reads its rebound attribute. + object.__setattr__(self, _REBOUND_STATE[name], value) + if name == "_core_cache": + self._hosted._core_cache = value # pyright: ignore[reportPrivateUsage, reportAttributeAccessIssue] # usage reads come from the routing cache + elif name == "_core_slack_alerting_logger": + from litellm.integrations.custom_logger import CustomLogger + + # `_slack_alerting_loggers` reads the attribute each time: the new logger is the one bound. + self._checks.slack_alerting_loggers = (value,) if isinstance(value, CustomLogger) else () + else: + # Any other attribute is kept as written; `_core_async_get_available_deployment` (the + # deprecated async pick) is then asked by every async pick (see `_pre_route`). + object.__setattr__(self, name, value) + + async def select( # pyright: ignore[reportIncompatibleMethodOverride] # the RouterCore selection + self, + model: str, + request_kwargs: dict[str, object] | None = None, # mutable-ok: the RouterCore signature + *, + kind: str = "pick", + specific_deployment: bool | None = None, + messages: object = None, + input: object = None, + health_check_probe: bool = False, + **_: object, # kwargs-ok: the RouterCore select options this core ignores + ) -> Mapping[str, object] | Sequence[Mapping[str, object]]: + """Pick a deployment for ``model`` without calling it (``pick``), list the deployments a + pick would choose among (``healthy``), or pick one for a pass-through route + (``pass_through``, which picks with the sync strategies, as Python does).""" + if kind not in ("pick", "healthy", "pass_through"): + raise ValueError(f"unknown select kind {kind!r}") + self._refuse_deleted_reads({"num_retries": 0}) # a pick reads no retry budget + request: Final = request_kwargs if request_kwargs is not None else {} + selection: Final = _Selection( + self, + model, + request, + kind, + specific_deployment, + messages, + input, + probe=kind == "healthy" and health_check_probe, + ) + # A pick that makes no call holds no least-busy reservation (`select` on the Python core). + with picks_reserve(False): + if kind == "pass_through": + return selection.run_sync() + return await selection.run() + + async def async_get_available_deployment_for_pass_through( + self, + model: str, + request_kwargs: dict[str, object], # mutable-ok: the legacy signature + messages: object = None, + input: object = None, + specific_deployment: bool | None = False, + ) -> Mapping[str, object]: + """The full healthy chain, then only ``use_in_pass_through`` deployments, picked with the + async strategies.""" + selection: Final = _Selection( + self, model, request_kwargs, "async_pass_through", specific_deployment, messages, input + ) + picked: Final = await selection.run() + return picked if isinstance(picked, Mapping) else {} + + def get_available_deployment( + self, + model: str, + messages: object = None, + input: object = None, + specific_deployment: bool | None = False, + request_kwargs: dict[str, object] | None = None, # mutable-ok: the legacy signature + ) -> Mapping[str, object]: + """The sync pick: no pre-routing, the sync strategies.""" + selection: Final = _Selection( + self, + model, + request_kwargs if request_kwargs is not None else {}, + "sync_pick", + specific_deployment, + messages, + input, + ) + picked: Final = selection.run_sync() + return picked if isinstance(picked, Mapping) else {} + + def get_available_deployment_for_pass_through( + self, + model: str, + messages: object = None, + input: object = None, + specific_deployment: bool | None = False, + request_kwargs: dict[str, object] | None = None, # mutable-ok: the legacy signature + ) -> Mapping[str, object]: + """Deprecated: ``select(model, kind="pass_through")``.""" + from litellm.router_deprecated import ( + _note, # pyright: ignore[reportPrivateUsage] # the router's deprecation notice + ) + + _note("get_available_deployment_for_pass_through", 'select(model, kind="pass_through")') + selection: Final = _Selection( + self, + model, + request_kwargs if request_kwargs is not None else {}, + "pass_through", + specific_deployment, + messages, + input, + ) + picked: Final = selection.run_sync() + return picked if isinstance(picked, Mapping) else {} + + # -- endpoints ------------------------------------------------------------------------------- + + async def call(self, op: str, request: Mapping[str, object]) -> object: + """Run one endpoint with the request's keyword payload, as ``getattr(core, op)(**request)``.""" + member: Final[object] = getattr(self, op) # pyright: ignore[reportAny] # litellm's untyped helpers and request values + if not callable(member): + raise TypeError(f"{op!r} is not an endpoint") + result: Final[object] = member(**request) + return await result if inspect.isawaitable(result) else result + + async def acompletion(self, *args: object, **kwargs: object) -> object: # kwargs-ok: endpoint payload + return await self._endpoints.acompletion(request_from_call("acompletion", args, kwargs)) + + async def atext_completion(self, *args: object, **kwargs: object) -> object: # kwargs-ok: endpoint payload + return await self._endpoints.atext_completion(request_from_call("atext_completion", args, kwargs)) + + async def aembedding(self, *args: object, **kwargs: object) -> object: # kwargs-ok: endpoint payload + return await self._drive(plan_for("aembedding"), request_from_call("aembedding", args, kwargs)) + + async def aresponses(self, *args: object, **kwargs: object) -> object: # kwargs-ok: endpoint payload + return await self._endpoints.generic(plan_for("aresponses"), request_from_call("aresponses", args, kwargs)) + + async def aanthropic_messages(self, **kwargs: object) -> object: # kwargs-ok: endpoint payload + return await self._endpoints.generic(plan_for("aanthropic_messages"), kwargs) + + async def anthropic_messages(self, **kwargs: object) -> object: # kwargs-ok: endpoint payload + """The proxy's ``/v1/messages`` route name for ``aanthropic_messages``; async, as in PythonCore.""" + return await self._endpoints.generic(plan_for("anthropic_messages"), kwargs) + + async def aimage_generation(self, *args: object, **kwargs: object) -> object: # kwargs-ok: endpoint payload + return await self._drive(plan_for("aimage_generation"), request_from_call("aimage_generation", args, kwargs)) + + async def atranscription(self, *args: object, **kwargs: object) -> object: # kwargs-ok: endpoint payload + return await self._drive(plan_for("atranscription"), request_from_call("atranscription", args, kwargs)) + + async def aspeech(self, *args: object, **kwargs: object) -> object: # kwargs-ok: endpoint payload + request: Final = dict(request_from_call("aspeech", args, kwargs)) + request.setdefault("voice", None) + return await self._drive(plan_for("aspeech"), request) + + async def arerank(self, *args: object, **kwargs: object) -> object: # kwargs-ok: endpoint payload + # `arerank` sets kwargs["input"] = input: the Python builtin, which reaches the deployment call. + request: Final = {**request_from_call("arerank", args, kwargs), "input": input} + return await self._drive(plan_for("arerank"), request) + + async def amoderation(self, *args: object, **kwargs: object) -> object: # kwargs-ok: endpoint payload + return await self._endpoints.amoderation(request_from_call("amoderation", args, kwargs)) + + def completion(self, *args: object, **kwargs: object) -> object: # kwargs-ok: endpoint payload + return self._call_sync(plan_for("completion"), request_from_call("completion", args, kwargs)) + + def text_completion(self, *args: object, **kwargs: object) -> object: # kwargs-ok: endpoint payload + return self._endpoints.text_completion(request_from_call("text_completion", args, kwargs)) + + def embedding( + self, + model: str, + input: str | list[object], # mutable-ok: the request field, as the Python core takes it + is_async: bool | None = False, + **kwargs: object, # kwargs-ok: endpoint payload + ) -> object: + """``embedding``'s own signature; as there, ``is_async`` is named and goes no further.""" + return self._call_sync(plan_for("embedding"), {**kwargs, "model": model, "input": input}) + + def image_generation(self, *args: object, **kwargs: object) -> object: # kwargs-ok: endpoint payload + return self._call_sync(plan_for("image_generation"), request_from_call("image_generation", args, kwargs)) + + # -- driving a call -------------------------------------------------------------------------- + + def _start( + self, + plan: Plan, + request: Mapping[str, object], + overrides: Mapping[str, object] | None = None, + *, + prepared_already: bool = False, + ) -> "tuple[RouterCall[_CallSteps], Mapping[str, object]]": + self._refuse_deleted_reads({**request, **(overrides or {})}) + self._sync_routing_state() + model: Final = request.get("model") + if model is None and "model" in request: + raise self._no_model_error(plan, request) + if "model" not in request: + raise TypeError(f"{plan.endpoint}() needs a model name") + name: Final = model if isinstance(model, str) else self.routed_name(model) + field: Final = plan.prepare_field or "metadata" + if request.get(field) is not None and not isinstance(request.get(field), Mapping): + raise TypeError(f"{field} must be a mapping") + prepared: Final = dict(request) if prepared_already else self._endpoints.prepare(plan, request, model) # pyright: ignore[reportArgumentType] # the name as given, which the metadata keeps + return self._handle.start( + plan.endpoint, self._routing_inputs(plan, request, prepared, name, overrides) + ), prepared + + def service_logged(self, call_type: str, start_time: float, end_time: float, request: Mapping[str, object]) -> None: + """The ``ServiceTypes.ROUTER`` success event of a timed step, as a task of the running loop.""" + import asyncio + + from litellm.litellm_core_utils.core_helpers import get_parent_otel_span_from_kwargs + from litellm.types.services import ServiceTypes + + _ = asyncio.create_task( # fire and forget, as the Python core schedules it + self.service_logger_obj.async_service_success_hook( # pyright: ignore[reportUnknownMemberType] # its unused `event_metadata` is an untyped dict + service=ServiceTypes.ROUTER, + duration=end_time - start_time, + call_type=call_type, + start_time=start_time, + end_time=end_time, + parent_otel_span=get_parent_otel_span_from_kwargs(request), + ) + ) + + def track_deployment_pick(self, deployment_id: object, request: Mapping[str, object]) -> None: + """``_track_deployment_metrics`` for a pick: ``_update_usage``'s count of the deployment's + picks in the router cache, on this worker only. A default deployment served under a + placeholder id counts under the id Python's copy of it holds.""" + from litellm._logging import verbose_router_logger + from litellm.litellm_core_utils.core_helpers import get_parent_otel_span_from_kwargs + + parent_otel_span: Final = get_parent_otel_span_from_kwargs(request) + try: + model_id: Final = ( + self._default_deployment_id() + if deployment_id in (DEFAULT_DEPLOYMENT_FAILING_ID, DEFAULT_DEPLOYMENT_COPIED_ID) + else deployment_id + ) + if model_id is None: + return + counts: Final = self._pick_counts() + count: Final = counts.get_cache(key=model_id, parent_otel_span=parent_otel_span, local_only=True) + if count is None: + counts.set_cache(key=model_id, value=1, local_only=True, ttl=60) # only store for 60s + else: + counts.set_cache(key=model_id, value=count + 1, local_only=True) + except Exception as e: # noqa: BLE001 # as `_track_deployment_metrics` catches it + verbose_router_logger.error("Error in _track_deployment_metrics: %s", e) + + def _pick_counts(self) -> _PickCounts: + return self._cache + + def _default_deployment_id(self) -> object: + return self._resolved.default_deployment.get("model_info", {}).get("id", None) # pyright: ignore[reportOptionalMemberAccess, reportAttributeAccessIssue, reportUnknownMemberType, reportUnknownVariableType] # read as Python reads it + + def routed_name(self, model: object, *, read_as_member: bool = False) -> str: + """The name a ``model`` that is no string routes under, for the native calls, which take strings + alone. Python's first reads raise for most values: the alias lookup hashes it (a sync pick reads + ``model in _core_model_names`` first, ``read_as_member``), then ``_core_get_model_list`` reaches + ``declared_provider``, which searches it. A value both accept (falsy, a tuple) is a name no + deployment serves, and ``str(model)`` is how every error names it.""" + from litellm.litellm_core_utils.get_llm_provider_logic import inferred_provider + + if read_as_member: + _ = model in frozenset[object]() + else: + _ = self._hosted._core_get_model_from_alias(model) # pyright: ignore[reportPrivateUsage, reportArgumentType] # raising is the point + _ = inferred_provider(model) # pyright: ignore[reportArgumentType] # see above + return str(model) + + def _no_model_error(self, plan: Plan, request: Mapping[str, object]) -> BaseException: + """``model=None``: the pick fails in ``_common_checks_available_deployment``, naming the + first default fallback when there is one, and with no group the fallback wrapper raises + that error as it is (``async_function_with_fallbacks_common_utils``).""" + import litellm.exceptions as litellm_exceptions + + named: Final = next( + ( + targets[0] + for entry in _as_list(self._resolved.fallbacks) + if (targets := _as_list((as_str_mapping(entry) or {}).get("*"))) + ), + None, + ) + error: Final = litellm_exceptions.BadRequestError( + message=f"You passed in model={named}. {RouterErrors.no_healthy_deployments.value}", + model=named, # pyright: ignore[reportArgumentType] # None, as the Python router passes it + llm_provider="", + ) + if not plan.sync: + _Driver(self, plan, request)._log_pick_failure(error) # pyright: ignore[reportPrivateUsage] # the async pick's own logging + return error + + def _routing_inputs( + self, + plan: Plan, + request: Mapping[str, object], + prepared: Mapping[str, object], + model: str, + overrides: Mapping[str, object] | None, + ) -> Mapping[str, object]: + """The document router-core starts a call from: what it routes on, then what this call changes.""" + fields: dict[str, object] = {"model": model} # mutable-ok: built up below, later keys win # rebind-ok: same + fields.update(self._routing_fields(prepared, plan.pick_input(prepared)[0])) + fields.update(pinned_video_deployment(plan, request)) + fields.update(unrouted_fields(plan)) + fields.update(failover_fields(plan)) + fields.update(scoped_fallback_fields(plan, request)) + fields.update(self._selection_fields(plan, prepared)) + fields.update(self._outcome_fields(plan, request)) + # The sync path never runs pre-routing; a custom strategy picks on both. + if not plan.sync or self._custom_strategy is not None: + fields["pre_routes"] = True + if (mock := mock_failure(request)) is not None: + fields["mock_failure"] = mock + fields.update(self._retry_policy_raises(request)) + if overrides is not None and "num_retries" in overrides: + fields.pop("uncounted_retries", None) + fields.update(overrides or {}) + return _document(fields) + + def _outcome_fields(self, plan: Plan | None, request: Mapping[str, object]) -> Mapping[str, object]: + """Which outcome callbacks run for the request: an endpoint that makes one attempt skips + ``make_call``, and a recovered stream is async.""" + sync: Final = plan is not None and plan.sync + make_call: Final = plan is None or not plan.single_attempt + callbacks: Final = self.outcome_callbacks(request, sync=sync, make_call=make_call) + return {"outcome_callbacks": callbacks} if callbacks else {} + + def _selection_fields(self, plan: Plan, request: Mapping[str, object]) -> Mapping[str, object]: + """What an endpoint's pick reads beyond the routing fields: the narrowing keys the caller + passed, the endpoint ``endpoint_selection`` names, and whether it asks for web search, + which only an async pick filters by.""" + endpoint: Final = selection_endpoint(self._port, plan) + web_search: Final = not plan.sync and self._picks_async() and is_web_search_request(request) + # The generic handler `_route_capability_pin` names: a call only some providers serve + # drops the others, and a Decisions call carrying media skips the deployments refusing it. + handler: Final = plan.host_op.rsplit(".", 1)[-1] + decisions_media: Final = handler in DECISIONS_MEDIA_CALL_TYPES and bool( + request.get("audio") or request.get("videos") + ) + return { + **_request_narrowing(request), + **({"endpoint": endpoint} if endpoint is not None else {}), + **({"web_search": True} if web_search else {}), + **({"route_handler": handler} if serves_route_capability(handler) else {}), + **({"decisions_media": True} if decisions_media else {}), + } + + def _picks_async(self) -> bool: + return picks_async(normalize_strategy(self._resolved.routing_strategy)) + + def _process_renamed(self, model: str, group: str) -> str: + """``model`` as the error names it (``_common_checks_available_deployment``): the group a name + resolved to by name, alias or routing group, renamed through the process ``model_alias_map``.""" + resolution: Final = self._handle.resolve(model) + if resolution is None or resolution["kind"] not in ("Group", "Alias", "RoutingGroup"): + return model + renames: Final = self._host.dynamic().model_alias_map + return renames[group] if group in renames else group # noqa: SIM401 # `renames` is any process mapping, read as Python reads it + + def _max_parallel_requests_client(self, deployment: Mapping[str, object], kwargs: Mapping[str, object]) -> object: + """``_get_client(client_type="max_parallel_requests")``: the deployment's slot limit.""" + return max_parallel_requests_limit(self._port, deployment, kwargs) + + def _routable( + self, model: str, request: Mapping[str, object], call_type: str | None = None + ) -> Sequence[Mapping[str, object]]: + """``_core_get_routable_deployments``: ``call_type`` names the endpoint it asks about.""" + endpoint: Final = endpoint_for_call(call_type) + return self._handle.routable( + model, + _document( + { + **_start_fields(request), + **({"endpoint": endpoint} if endpoint is not None else {}), + "wildcard_request": wildcard_request(model), + } + ), + ) + + def _routing_fields(self, request: Mapping[str, object], messages: object = None) -> Mapping[str, object]: + """The request fields the Rust core routes on, and whether its tag stage may apply: the + router's ``enable_tag_filtering``, the request's own, or a deployment's + ``model_info.enable_tag_filtering`` (``get_deployments_for_tag``'s effective setting).""" + filter_tags: Final = ( + self._resolved.enable_tag_filtering + or request.get("enable_tag_filtering") is True + or self._model_tag_filtering() + ) + fields: dict[str, object] = dict(_start_fields(request)) # mutable-ok: built up below # rebind-ok: same + fields.update(self._group_policy_fields(request)) + if filter_tags: + fields["filter_tags"] = True + if request_carries_image(messages, request): + fields["has_image"] = True + return fields + + def _group_policies(self, request: Mapping[str, object]) -> tuple[object, GroupPolicies]: + """The ``model_group_retry_policy`` the retry loop reads for this request (the request's own, + even ``None``, replaces the router's: ``kwargs.pop(..., self.model_group_retry_policy)``), + resolved for every group name the call may look up.""" + policies: Final = request.get("model_group_retry_policy", self._resolved.model_group_retry_policy) + names: Final = self._policy_lookup_names(request) + resolved: Final = group_policies(policies, names) + if reads_as_policy(self._written_retry_policy): + return policies, resolved + # A group no entry names looks up the router's `retry_policy`, and that lookup raises. + unnamed: Final = names - frozenset(resolved.document) + return policies, dataclasses.replace(resolved, raising=resolved.raising | unnamed) + + def _policy_lookup_names(self, request: Mapping[str, object]) -> frozenset[str]: + """The model names a retry policy lookup may run on (each hop looks up the name it calls): + the requested one and every name in the fallback lists the walk may take.""" + resolved: Final = self._resolved + lists: Final = ( + resolved.fallbacks, + resolved.context_window_fallbacks, + resolved.content_policy_fallbacks, + *(request.get(key) for key in _REQUEST_FALLBACK_KEYS), + ) + return frozenset((*_strings(request.get("model")), *_strings(lists))) + + def _retry_policy_raises(self, request: Mapping[str, object]) -> Mapping[str, object]: + """The groups whose ``model_group_retry_policy`` lookup raises on a hop's first failure, for + router-core to end that hop with the error the lookup raises (see ``_Driver.take_events``).""" + raising: Final = raising_groups(request, self._group_policies(request)[1]) + return {"retry_policy_raises": list(raising)} if raising else {} + + def _group_policy_fields(self, request: Mapping[str, object]) -> Mapping[str, object]: + """The request's ``model_group_retry_policy`` for router-core; a router policy that is not a + mapping of policies rides on every request, resolved for it.""" + policies, resolved = self._group_policies(request) + if "model_group_retry_policy" not in request and is_policy_per_group(policies): + return {} + return {"model_group_retry_policy": _plain(resolved.document)} + + def _model_tag_filtering(self) -> bool: + """Whether any deployment sets ``model_info.enable_tag_filtering``, cached per registry version.""" + cached: Final = self._tag_filtering_models + if cached is not None and cached[0] == self._version: + return cached[1] + found: Final = any( + (as_str_mapping((as_str_mapping(record) or {}).get("model_info")) or {}).get("enable_tag_filtering") + is not None + for record in self._handle.deployments() + ) + self._tag_filtering_models = (self._version, found) + return found + + async def _drive( + self, + plan: Plan, + request: Mapping[str, object], + overrides: Mapping[str, object] | None = None, + *, + passthrough: "Callable[[Mapping[str, object], str], Awaitable[object]] | None" = None, + prepared_already: bool = False, + ) -> object: + routed: Final = ( + overrides if passthrough is None else {**(overrides or {}), "passthrough_on_no_deployment": True} + ) + # `owning_picks`: least-busy frees every pick this request made but never started. + with owning_picks(self._selectors.reservations, call_id_of(CALL_IDENTITY.validate_python(request))): + call, prepared = self._start(plan, request, routed, prepared_already=prepared_already) + arm_compaction_state(self._port, plan, prepared) # pyright: ignore[reportArgumentType] # the call's own dict + with picks_reserve(plan.endpoint not in _UNRESERVED_OPS): + driver, step = await self._run(plan, call, prepared, fresh=True, passthrough=passthrough) + if step["kind"] == "failed": + raise driver.failed(step) + return await driver.finish_async(step) + + async def _run( + self, + plan: Plan, + call: "RouterCall[_CallSteps]", + prepared: Mapping[str, object], + cut_deployment: object = None, + stream_failure: tuple[int, BaseException] | None = None, + *, + fresh: bool = False, + passthrough: "Callable[[Mapping[str, object], str], Awaitable[object]] | None" = None, + ) -> "tuple[_Driver, Mapping[str, object]]": + """Drive ``call`` to its ``done`` or ``failed`` step, answering every step on the host. + ``cut_deployment`` is the attempt whose stream failed, for a recovery call, and + ``stream_failure`` the level it served with the failure a fallback walk re-enters from. + ``fresh`` is a call from the top, through the retry and fallback wrappers, and + ``passthrough`` serves an attempt whose pick failed.""" + driver: Final = _Driver( + self, plan, prepared, cut_deployment, stream_failure, fresh=fresh, passthrough=passthrough + ) + outcome: object = None # rebind-ok: the answer to the last step + while True: + step = call.next(outcome) + driver.take_events(step.get("events")) + # `run_async_fallback` stamps a served response before it logs the hop's success: `finish` + # runs what a served call still has pending. + if step["kind"] != "done": + await await_all(driver.outcomes.take_pending()) + match step["kind"]: + case "attempt": + outcome = await driver.attempt(step) + case "sleep": + await self._host.sleep(_sleep_seconds(step)) + outcome = None + case "allow_fallback": + outcome = await driver.allow_fallback(str(step["target"])) + case "refuse_fallbacks": + outcome = driver.refuse_fallbacks(step) + case "cache": + outcome = await self._answer_cache(step) + case "select": + outcome = await driver.select(step) + case "pre_route": + outcome = await driver.pre_route(step) + case "done" | "failed": + return driver, step + case other: # pyright: ignore[reportUnnecessaryComparison] # a kind the bridge returns but the stub does not list + raise RuntimeError(f"router-core returned an unknown step {other!r}") + + def _call_sync( + self, plan: Plan, request: Mapping[str, object], overrides: Mapping[str, object] | None = None + ) -> object: + # The sync routed calls pick with `picks_reserve(True)` inside `owning_picks`. + with owning_picks(self._selectors.reservations, call_id_of(CALL_IDENTITY.validate_python(request))): + call, prepared = self._start(plan, request, overrides) + with picks_reserve(plan.endpoint not in _UNRESERVED_OPS): + driver, step = self._run_sync(plan, call, prepared, fresh=True) + if step["kind"] == "failed": + raise driver.failed(step) + return driver.finish(step) + + def _run_sync( + self, + plan: Plan, + call: "RouterCall[_CallSteps]", + prepared: Mapping[str, object], + cut_deployment: object = None, + stream_failure: tuple[int, BaseException] | None = None, + *, + fresh: bool = False, + ) -> "tuple[_Driver, Mapping[str, object]]": + driver: Final = _Driver(self, plan, prepared, cut_deployment, stream_failure, sync=True, fresh=fresh) + outcome: object = None # rebind-ok: the answer to the last step + while True: + step = call.next(outcome) + driver.take_events(step.get("events")) + if step["kind"] != "done": # see _run + driver.log_outcomes_sync() + match step["kind"]: + case "attempt": + outcome = driver.attempt_sync(step) + case "sleep": + self._host.sleep_sync(_sleep_seconds(step)) + outcome = None + case "allow_fallback": + from litellm.litellm_core_utils import asyncify + + outcome = asyncify.run_async_function(driver.allow_fallback, str(step["target"])) # pyright: ignore[reportUnknownMemberType] # the untyped sync bridge + case "refuse_fallbacks": + outcome = driver.refuse_fallbacks(step) + case "cache": + outcome = self._answer_cache_sync(step) + case "select": + outcome = driver.select_sync(step) + case "pre_route": + outcome = driver.pre_route_sync(step) + case "done" | "failed": + return driver, step + case other: # pyright: ignore[reportUnnecessaryComparison] # a kind the bridge returns but the stub does not list + raise RuntimeError(f"router-core returned an unknown step {other!r}") + + # -- stream recovery (the host's stream wrappers, router_rust.invoke.streaming) --------------------- + + def _recovery_call( + self, + recovery: str, + op: str, + request: Mapping[str, object], + origin: Mapping[str, object], + failure: BaseException, + ) -> "RouterCall[_CallSteps]": + start: Final = self._handle.retry_stream if recovery == "retry" else self._handle.fallback_stream + plan: Final = generic_plan(op) + routing: Final = _document( + { + "model": _routed_model(request), + **self._routing_fields(request), + **(self._selection_fields(plan, request) if plan is not None else {}), + **self._retry_policy_raises(request), + **continuation_routing(self, request), + **self._outcome_fields(plan, request), + } + ) + return start(op, routing, _document(origin), attempt_failure(failure)) + + def stream_policy( + self, request: Mapping[str, object], origin: Mapping[str, object], trigger: BaseException | None + ) -> Mapping[str, object]: + routing: Final = _document( + { + "model": _routed_model(request), + **_start_fields(request), + **self._group_policy_fields(request), + **self._retry_policy_raises(request), + } + ) + kind: Final = None if trigger is None else str(attempt_failure(trigger)["kind"]) + return self._handle.stream_policy(routing, _document(origin), kind) + + async def recover_stream( + self, + recovery: str, + op: str, + request: Mapping[str, object], + origin: Mapping[str, object], + failure: BaseException, + ) -> object: + """The fallback walk for a stream that failed before content: the fallback's response with the + router's headers, or the error the walk ended on.""" + call: Final = self._recovery_call(recovery, op, request, origin, failure) + driver, step = await self._run( + plan_for(op), call, request, origin.get("deployment_id"), (_origin_level(origin), failure) + ) + if step["kind"] == "failed": + raise driver.failed(step) + return await driver.finish_async(step) + + def recover_stream_sync( + self, + recovery: str, + op: str, + request: Mapping[str, object], + origin: Mapping[str, object], + failure: BaseException, + ) -> object: + call: Final = self._recovery_call(recovery, op, request, origin, failure) + driver, step = self._run_sync( + plan_for(op), call, request, origin.get("deployment_id"), (_origin_level(origin), failure) + ) + if step["kind"] == "failed": + raise driver.failed(step) + return driver.finish(step) + + async def retry_stream( + self, op: str, request: Mapping[str, object], origin: Mapping[str, object], failure: BaseException + ) -> "tuple[object, int, int] | BaseException": + """A same-group retry of a stream that failed before content: the retried response (without the + router's headers) with the retry count and budget, or the newest failure (``failure`` itself + when no retry ran).""" + call: Final = self._recovery_call("retry", op, request, origin, failure) + driver, step = await self._run(plan_for(op), call, request, origin.get("deployment_id")) + await await_all(driver.outcomes.take_pending()) + if step["kind"] == "failed": + return driver.last_error if driver.last_error is not None else failure + attempted: Final = step.get("attempted_retries") + budget: Final = step.get("max_retries") + return ( + driver.response, + attempted if isinstance(attempted, int) else 0, + budget if isinstance(budget, int) else 0, + ) + + def params_sent_to_deployment(self, request: Mapping[str, object], deployment_id: object) -> Mapping[str, object]: + """The request as a deployment received it: its params, the router defaults, its tool choice, then + the request (``_params_sent_to_deployment``).""" + record: Final = self._handle.deployment(deployment_id) if isinstance(deployment_id, str) else None + params: Final = as_str_mapping((as_str_mapping(record) or {}).get("litellm_params")) or {} + defaults: Final = {k: v for k, v in self._resolved.default_litellm_params.items() if v is not None} + tool_choice: Final = {"tool_choice": params["tool_choice"]} if params.get("tool_choice") is not None else {} + return MappingProxyType({**params, **defaults, **tool_choice, **request}) + + def _drive_sync_steps(self, call: "RouterCall[_FailureReportSteps | _SuccessReportSteps]") -> Mapping[str, object]: + """Drive a call whose only steps are cache ops (``report_failure``) to its end.""" + outcome: object = None # rebind-ok: the answer to the last step + while True: + step = call.next(outcome) + match step["kind"]: + case "cache": + outcome = self._answer_cache_sync(step) + case "done": + return step + case "failed": + raise RuntimeError(f"router-core ({step.get('error_kind')}): {step.get('error')}") + case other: # pyright: ignore[reportUnnecessaryComparison] # a kind the bridge returns but the stub does not list + raise RuntimeError(f"router-core asked for {other!r} while recording an outcome") + + async def _drive_cache_steps( + self, call: "RouterCall[_FailureReportSteps | _SuccessReportSteps]" + ) -> Mapping[str, object]: + """``_drive_sync_steps`` answered through the async cache (``report_success``).""" + outcome: object = None # rebind-ok: the answer to the last step + while True: + step = call.next(outcome) + match step["kind"]: + case "cache": + outcome = await self._answer_cache(step) + case "done": + return step + case "failed": + raise RuntimeError(f"router-core ({step.get('error_kind')}): {step.get('error')}") + case other: # pyright: ignore[reportUnnecessaryComparison] # a kind the bridge returns but the stub does not list + raise RuntimeError(f"router-core asked for {other!r} while recording an outcome") + + # -- pre-route steps ------------------------------------------------------------------------- + + async def _pre_route( + self, + step: Mapping[str, object], + request: Mapping[str, object], + messages: object = _NO_INPUT, + input: object = _NO_INPUT, + specific_deployment: object = _NO_INPUT, + *, + pass_through: bool = False, + ) -> object: + """The host's pre-routing of one pick (router_rust.steps.prerouting): the model group to resolve + instead of the step's ``model``, a custom strategy's deployment, or ``None``. ``pass_through`` + is ``_async_get_available_deployment_for_pass_through``, which runs the hook whatever the + router's strategy and which no replaced pick answers.""" + assigned: Final[object] = None if pass_through else self.__dict__.get("_core_async_get_available_deployment") + if callable(assigned): + # The async pick assigned through the deprecated attribute answers, as Python's + # replaced `_core_async_get_available_deployment` does. + answer: Final[object] = assigned( + model=str(step.get("model")), + request_kwargs=request, + messages=request.get("messages") if messages is _NO_INPUT else messages, + input=request.get("input") if input is _NO_INPUT else input, + specific_deployment=request.get("specific_deployment") + if specific_deployment is _NO_INPUT + else specific_deployment, + ) + chosen: Final[object] = await answer if inspect.isawaitable(answer) else answer + return _custom_pick(chosen) + custom: Final = None if pass_through else self._custom_strategy + if custom is not None: + picked: Final[object] = await custom.async_get_available_deployment( # pyright: ignore[reportUnknownMemberType] # the strategy's untyped pick + model=str(step.get("model")), + messages=request.get("messages") if messages is _NO_INPUT else messages, # pyright: ignore[reportArgumentType] # the request's messages, as given + input=request.get("input") if input is _NO_INPUT else input, # pyright: ignore[reportArgumentType] # the request's input, as given + specific_deployment=request.get("specific_deployment") + if specific_deployment is _NO_INPUT + else specific_deployment, # pyright: ignore[reportArgumentType] # same + request_kwargs=request, # pyright: ignore[reportArgumentType] # the request kwargs, as given + ) + return _custom_pick(picked) + if not isinstance(request, dict): + return None + depth: Final = step.get("fallback_depth") + r: Final = self._resolved + owner: Final[object] = self.__dict__.get("_core_owner") + public: Final[object] = owner() if callable(owner) else None + return await self._prerouter.pre_route( + str(step.get("model")), + request, # pyright: ignore[reportUnknownArgumentType] # the request kwargs the hook stamps + messages=request.get("messages") if messages is _NO_INPUT else messages, + input=request.get("input") if input is _NO_INPUT else input, + fallback_depth=depth if isinstance(depth, int) else 0, + router_strategy=normalize_strategy(r.routing_strategy), + pass_through=pass_through, + enable_tag_filtering=r.enable_tag_filtering, + match_any=r.tag_filtering_match_any, + public_router=public if public is not None else self, + run_plugins=partial( + self._run_plugins, request, request.get("messages") if messages is _NO_INPUT else messages + ) # pyright: ignore[reportUnknownArgumentType] # the request kwargs the plugins stamp + if self._hosted._core_routing_plugins # pyright: ignore[reportPrivateUsage] # the live plugin list + else None, + ) + + async def _run_plugins( + self, + request: dict[str, object], # mutable-ok: the plugins stamp their outcome on the caller's request + messages: object, + model: str, + ) -> None: + """``_run_routing_plugins`` over the pre-routing hook's session-resolved ``model``.""" + await run_plugins( + tuple(self._hosted._core_routing_plugins), # pyright: ignore[reportPrivateUsage] # the live plugin list + self._hosted, # pyright: ignore[reportArgumentType] # the host registry serves the listing + self._resolved.model_group_alias, + model, + request, + messages, + ) + + def _custom_pick_sync( + self, + step: Mapping[str, object], + request: Mapping[str, object], + messages: object = _NO_INPUT, + input: object = _NO_INPUT, + specific_deployment: object = _NO_INPUT, + ) -> object: + """A custom strategy's sync pick (``get_available_deployment``); without one, no pre-routing.""" + custom: Final = self._custom_strategy + if custom is None: + return None + picked: Final[object] = custom.get_available_deployment( # pyright: ignore[reportUnknownMemberType] # the strategy's untyped pick + model=str(step.get("model")), + messages=request.get("messages") if messages is _NO_INPUT else messages, # pyright: ignore[reportArgumentType] # see _pre_route + input=request.get("input") if input is _NO_INPUT else input, # pyright: ignore[reportArgumentType] # same + specific_deployment=request.get("specific_deployment") + if specific_deployment is _NO_INPUT + else specific_deployment, # pyright: ignore[reportArgumentType] # same + request_kwargs=request, # pyright: ignore[reportArgumentType] # same + ) + return _custom_pick(picked) + + # -- select steps ---------------------------------------------------------------------------- + + def _strategy_inputs(self, step: Mapping[str, object], request: Mapping[str, object]) -> "_StrategyInputs | None": + """The strategy, its selector, the name it keys state by and the candidate deployments of a + ``strategy`` select step; ``None`` for a stage the host passes through.""" + if step.get("stage") != "strategy": + return None + strategy: Final = str(step.get("strategy")) + group: Final = step.get("group") + selector: Final = self._selectors.selector_for( + strategy, str(step.get("scope")), str(group) if group is not None else None, request + ) + candidates: Final = [ + dict(as_str_mapping(self._handle.deployment(str(id))) or {}) for id in _as_list(step.get("candidates")) + ] + return strategy, selector, str(step.get("model")), candidates + + def _records_of(self, ids: object) -> list[dict[str, object]]: # mutable-ok: the deployment dicts filters take + return [dict(as_str_mapping(self._handle.deployment(str(id))) or {}) for id in _as_list(ids)] + + async def _answer_tags(self, step: Mapping[str, object], request: Mapping[str, object]) -> object: + """``get_deployments_for_tag`` over the step's candidates (bridged: not ported to Rust yet).""" + from litellm.litellm_core_utils.core_helpers import get_metadata_variable_name_from_kwargs + from litellm.router_strategy.tag_based_routing import get_deployments_for_tag + + kept: Final = await get_deployments_for_tag( + ctx=RoutingContext(self), # pyright: ignore[reportArgumentType] # the filter reads registry() and settings() + model=str(step.get("model")), + healthy_deployments=self._records_of(step.get("candidates")), # pyright: ignore[reportArgumentType] # deployment dicts + request_kwargs=dict(request), # pyright: ignore[reportArgumentType] # the request kwargs + metadata_variable_name=get_metadata_variable_name_from_kwargs(dict(request)), + request_deployments=self._records_of(step.get("pool")), # pyright: ignore[reportArgumentType] # deployment dicts + ) + return [found for record in kept if (found := picked_id(record)) is not None] + + def _pre_call_facts( + self, step: Mapping[str, object], request: Mapping[str, object], messages: object, input: object + ) -> PreCallFacts: + hosted: Final = self._hosted + + def record(id: str) -> Mapping[str, object]: + return hosted._core_get_model_info(id) or as_str_mapping(self._handle.deployment(id)) or {} # pyright: ignore[reportPrivateUsage] # see __init__ + + return PreCallFacts(self._host, hosted, step, record, request, messages, input) # pyright: ignore[reportArgumentType] # the host registry serves the reads + + def _text_only(self, step: Mapping[str, object]) -> list[str]: # mutable-ok: the step's answer + """The candidates that declare no vision support (bridged: needs the cost map).""" + group: Final = str(step.get("model_group")) + return [ + str(id) + for id in _as_list(step.get("candidates")) + if declares_no_vision(as_str_mapping(self._handle.deployment(str(id))) or {}, group) + ] + + def _route_incapable(self, step: Mapping[str, object]) -> list[str]: # mutable-ok: the step's answer + """The candidates whose provider cannot serve the call the step's handler names.""" + group: Final = str(step.get("model_group")) + handler: Final = str(step.get("handler")) + return [ + str(id) + for id in _as_list(step.get("candidates")) + if cannot_serve_route(as_str_mapping(self._handle.deployment(str(id))) or {}, group, handler) + ] + + def _media_refusing( + self, step: Mapping[str, object], request: Mapping[str, object] + ) -> list[str]: # mutable-ok: the step's answer + """The candidates that refuse the audio or videos a Decisions call carries.""" + group: Final = str(step.get("model_group")) + defaults: Final = {k: v for k, v in self._resolved.default_litellm_params.items() if v is not None} + sent: Final = {**defaults, **request} + audio: Final = bool(request.get("audio")) + videos: Final = bool(request.get("videos")) + return [ + str(id) + for id in _as_list(step.get("candidates")) + if refuses_decisions_media( + as_str_mapping(self._handle.deployment(str(id))) or {}, group, sent, audio=audio, videos=videos + ) + ] + + async def _answer_select( + self, + step: Mapping[str, object], + request: Mapping[str, object], + messages: object, + input: object, + *, + pass_through: bool = False, + ) -> object: + if step.get("stage") == "pre_call": + return await self._pre_call_facts(step, request, messages, input).answer_async() + if step.get("stage") == "routing_plugins": + # The plugins already ran in this pick's pre_route step (see `_pre_route`), on the + # session-resolved name and before a strategy was selected, as Python's hook runs them. + return candidate_models(request) + if step.get("stage") == "tags": + return await self._answer_tags(step, request) + if step.get("stage") == "text_only": + return self._text_only(step) + if step.get("stage") == "route_incapable": + return self._route_incapable(step) + if step.get("stage") == "media_refusing": + return self._media_refusing(step, request) + if step.get("stage") == "provider_request": + return _provider_request(step) + inputs: Final = self._strategy_inputs(step, request) + if inputs is None: + return step["candidates"] + strategy, selector, model, candidates = inputs + pinned: Final = bool(request.get("_encrypted_content_affinity_pinned")) and len(candidates) == 1 + if selector is None and not pinned: + return [] + # The strategy's service event: `_async_get_available_deployment` times it for the strategies + # with an async pick, `_async_get_available_deployment_for_pass_through` for any, on another clock. + timed: Final = not pinned and strategy != "simple-shuffle" and (pass_through or self._picks_async()) + clock: Final = perf_counter if pass_through else unix_time + start_time: Final = clock() + # `_async_get_available_deployment`: a pin by encrypted content takes no strategy's pick. + picked: Final = ( + candidates[0] + if pinned + else await pick_async( + strategy, + selector, # pyright: ignore[reportArgumentType] # a selector from selector_for + model, + candidates, + messages, + input, + dict(request), + # `_async_get_available_deployment` reserves a request's pick unless told not to; + # the pass-through pick only when told to. + reserve=pick_reserves(default=not pass_through) and bool(request), + ) + ) + if picked is not None and selector is not None and self._selectors.is_override(strategy, selector): + await selector.async_pre_call_check(dict(as_str_mapping(picked) or {}), None) # pyright: ignore[reportAttributeAccessIssue, reportUnknownMemberType] # every selector is a CustomLogger + if picked is not None and timed: + self.service_logged(".async_get_available_deployments", start_time, clock(), request) + found: Final = picked_id(picked) + return [found] if found is not None else [] + + def _answer_select_sync( + self, step: Mapping[str, object], request: Mapping[str, object], messages: object, input: object + ) -> object: + if step.get("stage") == "pre_call": + return self._pre_call_facts(step, request, messages, input).answer() + if step.get("stage") == "routing_plugins": + return candidate_models(request) + if step.get("stage") == "callback_filters": + # Only the async candidate lists run the callback filters. + return step["candidates"] + if step.get("stage") == "text_only": + return self._text_only(step) + if step.get("stage") == "route_incapable": + return self._route_incapable(step) + if step.get("stage") == "media_refusing": + return self._media_refusing(step, request) + if step.get("stage") == "provider_request": + return _provider_request(step) + inputs: Final = self._strategy_inputs(step, request) + if inputs is None: + return step["candidates"] + strategy, selector, model, candidates = inputs + if selector is None: + return [] + picked: Final = pick_sync( + strategy, + selector, # pyright: ignore[reportArgumentType] # a selector from selector_for + model, + candidates, + messages, + input, + dict(request), + # The sync picks reserve only when a routed call site says so (`picks_reserve(True)`). + reserve=pick_reserves(default=False) and bool(request), + ) + if picked is not None and self._selectors.is_override(strategy, selector): + selector.pre_call_check(dict(as_str_mapping(picked) or {})) # pyright: ignore[reportAttributeAccessIssue, reportUnknownMemberType] # see above + found: Final = picked_id(picked) + return [found] if found is not None else [] + + # -- cache steps ----------------------------------------------------------------------------- + + def _store_for(self, key: str) -> "DualCache": + """Cooldowns live in the cooldown store; every other routing key in the routing cache.""" + return self._cooldowns.cooldown_store if key.endswith(":cooldown") else self._cache + + async def _batch_read(self, keys: list[str]) -> Sequence[object]: # mutable-ok: the cache's key list + """One batch read, as the Python core makes it: the keys memory cannot serve, at most once per + read interval. Cooldowns are ``async_get_cooldown_deployments``' read, through the request's + armed prefetch when there is one.""" + if not keys or not keys[0].endswith(":cooldown"): + read: Final[object] = await self._cache.async_batch_get_cache(keys=keys) # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] # DualCache's reads are untyped + return read if isinstance(read, list) else [None] * len(keys) # pyright: ignore[reportUnknownVariableType] # see above + store: Final = self._cooldowns.cooldown_store + with service_target(ROUTER_COOLDOWNS_TARGET): + prefetched: Final = await RoutingReadBatch._read_prefetched(((store, keys),)) # pyright: ignore[reportPrivateUsage] # the batch's own prefetch read + results: Final[object] = ( + prefetched[0] if prefetched is not None else await store.async_batch_get_cache(keys=keys) # pyright: ignore[reportUnknownMemberType] # see above + ) + return self._corrected_cooldowns(keys, results) + + def _batch_read_sync(self, keys: list[str]) -> Sequence[object]: # mutable-ok: the cache's key list + """``_batch_read`` for a sync pick: ``get_active_cooldowns``' read.""" + if not keys or not keys[0].endswith(":cooldown"): + read: Final[object] = self._cache.batch_get_cache(keys=keys) # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] # DualCache's reads are untyped + return read if isinstance(read, list) else [None] * len(keys) # pyright: ignore[reportUnknownVariableType] # see above + with service_target(ROUTER_COOLDOWNS_TARGET): + results: Final = as_object_list(self._cooldowns.cooldown_store.batch_get_cache(keys=keys)) # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType] # see above + return self._corrected_cooldowns(keys, results) + + def _corrected_cooldowns(self, keys: Sequence[str], results: object) -> Sequence[object]: + """The cooldowns a batch read found, once ``active_cooldowns_from_results`` corrected the in-memory + TTL of each it promoted from Redis and dropped each that expired.""" + if not isinstance(results, list): + return [None] * len(keys) + ids: Final = [key.removeprefix("deployment:").removesuffix(":cooldown") for key in keys] + _ = self._cooldowns.active_cooldowns_from_results(ids, results) # pyright: ignore[reportUnknownArgumentType] # see above + return results # pyright: ignore[reportUnknownVariableType] # see above + + async def _increment_usage(self, keys: Sequence[object], amounts: Sequence[object], ttl: float | None) -> None: + """``_increment_deployment_usage``'s write: one pipeline, then this worker's memory holds what + Redis answered, so its usage reads see what the other workers counted.""" + from litellm.types.caching import RedisPipelineIncrementOperation + + operations: Final = [ + RedisPipelineIncrementOperation( + key=str(key), increment_value=_amount(amount), ttl=None if ttl is None else int(ttl) + ) + for key, amount in zip(keys, amounts) + ] + values: Final = await self._cache.async_increment_cache_pipeline(increment_list=operations) + if values is None or self._cache.redis_cache is None: + return + for operation, value in zip(operations, values): + await self._cache.async_set_cache(operation["key"], int(value), local_only=True, ttl=operation["ttl"]) # pyright: ignore[reportUnknownMemberType] # DualCache's writes are untyped + + async def _answer_cache(self, step: Mapping[str, object]) -> object: + with service_target(_cache_target(str(step.get("key")))): + return await self._answer_cache_in_target(step) + + async def _answer_cache_in_target(self, step: Mapping[str, object]) -> object: + key: Final = str(step.get("key")) + ttl: Final = _ttl(step) + local: Final = step.get("local") is True + match step.get("op"): + case "get": + return _plain(await self._store_for(key).async_get_cache(key, local_only=local)) # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType] # DualCache's reads are untyped + case "batch_get": + return _plain(await self._batch_read([str(k) for k in _as_list(step.get("keys"))])) + case "set": + store: Final = self._store_for(key) + if ttl is None: + await store.async_set_cache(key, step.get("value")) # pyright: ignore[reportUnknownMemberType] # DualCache's writes are untyped + else: + await store.async_set_cache(key, step.get("value"), ttl=ttl) # pyright: ignore[reportUnknownMemberType] # see above + return None + case "increment": + by: Final = _amount(step.get("value")) + counter: Final = self._store_for(key) + total: Final = ( + await counter.async_increment_cache(key, by, local_only=local) # pyright: ignore[reportUnknownMemberType] # see above + if ttl is None + else await counter.async_increment_cache(key, by, local_only=local, ttl=ttl) # pyright: ignore[reportUnknownMemberType] # see above + ) + return by if total is None else total + case "increment_usage": + await self._increment_usage(_as_list(step.get("keys")), _as_list(step.get("value")), ttl) + return None + case "delete": + await self._store_for(key).async_delete_cache(key) + return None + case other: + raise RuntimeError(f"router-core asked for an unknown cache op {other!r}") + + def run_routing_cache_sync(self, read_or_write: "Callable[[], Coroutine[object, object, object]]") -> object: + """A routing-cache coroutine on a sync call, whose `make_call` runs on the one event loop + `run_async_function` opened for the whole call. Without Redis the coroutine never suspends, so + it runs here: a loop of its own would hand the GIL to the logging thread, which then reads the + response before the router's headers are on it.""" + if self._cache.redis_cache is not None: + from litellm.litellm_core_utils import asyncify + + return asyncify.run_async_function(read_or_write) # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] # the untyped sync bridge + running: Final = read_or_write() + try: + while True: + _ = running.send(None) + except StopIteration as finished: + return finished.value # pyright: ignore[reportAny] # the coroutine's result + + def _answer_cache_sync(self, step: Mapping[str, object]) -> object: + with service_target(_cache_target(str(step.get("key")))): + return self._answer_cache_sync_in_target(step) + + def _answer_cache_sync_in_target(self, step: Mapping[str, object]) -> object: + key: Final = str(step.get("key")) + ttl: Final = _ttl(step) + local: Final = step.get("local") is True + match step.get("op"): + case "get": + return _plain(self._store_for(key).get_cache(key, local_only=local)) # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType] # DualCache's reads are untyped + case "batch_get": + return _plain(self._batch_read_sync([str(k) for k in _as_list(step.get("keys"))])) + case "set": + store: Final = self._store_for(key) + if ttl is None: + store.set_cache(key, step.get("value")) # pyright: ignore[reportUnknownMemberType] # DualCache's writes are untyped + else: + store.set_cache(key, step.get("value"), ttl=ttl) # pyright: ignore[reportUnknownMemberType] # see above + return None + case "increment": + by: Final = _amount(step.get("value")) + counter: Final = self._store_for(key) + # DualCache.increment_cache is annotated int but adds floats as given. + return ( + counter.increment_cache(key, by, local_only=local) # pyright: ignore[reportUnknownMemberType, reportArgumentType] # see above + if ttl is None + else counter.increment_cache(key, by, local_only=local, ttl=ttl) # pyright: ignore[reportUnknownMemberType, reportArgumentType] # see above + ) + case "increment_usage": + keys: Final = _as_list(step.get("keys")) + amounts: Final = _as_list(step.get("value")) + self.run_routing_cache_sync(lambda: self._increment_usage(keys, amounts, ttl)) + return None + case "delete": + self._store_for(key).delete_cache(key) # pyright: ignore[reportUnknownMemberType] # DualCache's delete is untyped + return None + case other: + raise RuntimeError(f"router-core asked for an unknown cache op {other!r}") + + # -- errors the core raises when no attempt failed ------------------------------------------ + + def _failure_error( # noqa: C901 # one branch per failure Python tells apart + self, + step: Mapping[str, object], + model: str, + last: BaseException | None, + *, + pass_through: str = "", + request: Mapping[str, object] | None = None, + ) -> BaseException: + """The error a call or pick that failed before any attempt raises; ``pass_through`` names + the selection kind (``pass_through``/``async_pass_through``) whose errors it words.""" + import litellm.exceptions as litellm_exceptions + + kind: Final = step.get("error_kind") + aliases: Final = self._resolved.model_group_alias + refused: Final = selection_refusal(step, model, request or {}) + if refused is not None: + return refused + if kind == "paused": + # The message names the deployment's litellm model, not the id that was asked for; the + # default deployment's copy is named by the requested name. + record: Final = as_str_mapping(self._handle.deployment(model)) + litellm_model: Final = ( + str((as_str_mapping(record.get("litellm_params")) or {}).get("model")) if record is not None else model + ) + return litellm_exceptions.ServiceUnavailableError( + message=f"Model '{litellm_model}' is currently paused and cannot accept requests.", + model=litellm_model, + llm_provider="", + ) + if kind == "only_markers": + named_model: Final = str(step.get("error_model") or model) + return litellm_exceptions.BadRequestError( + message=f"You passed in model={named_model}. {RouterErrors.only_strategy_marker_deployments.value}", + model=named_model, + llm_provider="", + ) + if kind == "endpoint_not_supported": + return litellm_exceptions.EndpointNotSupportedByModelError( + message=str(step.get("error_detail")), + model=str(step.get("error_model") or model), + llm_provider="", + ) + if kind == "ptu_reserved": + reserved: Final = str(step.get("error_model") or model) + return litellm_exceptions.BadRequestError( + message=f"Deployment {reserved} is reserved for the teams holding a PTU share of it", + model=reserved, + llm_provider="", + ) + if kind == "not_pass_through": + # The sync pick names what `_common_checks_available_deployment` resolved the id to. + named_record: Final = ( + as_str_mapping(self._handle.deployment(model)) if pass_through == "pass_through" else None + ) + error_model: Final = ( + str((as_str_mapping(named_record.get("litellm_params")) or {}).get("model")) + if named_record is not None + else model + ) + return litellm_exceptions.BadRequestError( + message=f"Deployment {model} does not support pass-through endpoint (use_in_pass_through=False)", + model=error_model, + llm_provider="", + ) + if kind == "no_pass_through_deployments": + # The sync pick names what `_common_checks_available_deployment` resolved the name to. + named_group: Final = str(step.get("error_model") or model) + listed_model: Final = ( + named_group if pass_through == "async_pass_through" else self._process_renamed(model, named_group) + ) + message: Final = ( + f"Model {listed_model} has no deployments configured with use_in_pass_through=True. " + "Please add use_in_pass_through: true to the deployment configuration" + if pass_through == "async_pass_through" + else f"Model {listed_model} has no deployment configured with use_in_pass_through=True. " + "Please add use_in_pass_through: true in the deployment configuration" + ) + return litellm_exceptions.BadRequestError(message=message, model=listed_model, llm_provider="") + if kind == "context_window_exceeded": + return litellm_exceptions.ContextWindowExceededError( + message="litellm._pre_call_checks: Context Window exceeded for given call. No models have context " + f"window large enough for this call.\n{step.get('error_detail')}", + model=str(step.get("error_model") or model), + llm_provider="", + ) + if kind == "unsupported_params": + params: Final = sorted(str(param) for param in _as_list(step.get("error_params"))) + unsupported_model: Final = str(step.get("error_model") or model) + return litellm_exceptions.UnsupportedParamsError( + message=( + f"no deployment of model={unsupported_model} supports {params}. " + "To drop these, set `litellm.drop_params=True` or for proxy `litellm_settings: drop_params: true`, " + f"or send allowed_openai_params={params} to pass them through" + ), + model=unsupported_model, + ) + if kind == "routing_plugins_left_nothing": + left: Final = [str(name) for name in _as_list(step.get("error_params"))] + return ValueError(f"No deployments left after routing-plugin filtering. candidate_models={left}") + if kind == "routing_plugins_need_async": + return ValueError(SYNC_PATH_MESSAGE) + looked_up: Final = _lookup_error(kind, str(step.get("error_model") or model), self._resolved.default_deployment) + if looked_up is not None: + return looked_up + if kind == "unknown_model": + # router-core names what `_common_checks_available_deployment` looked for last. + aliased: Final = _alias_target(aliases.get(model)) or model if model in aliases else model + named: Final = str(step.get("error_model") or aliased) + return litellm_exceptions.BadRequestError( + message=f"You passed in model={named}. {RouterErrors.no_healthy_deployments.value}", + model=named, + llm_provider="", + ) + if kind == "no_healthy_deployments": + group: Final = self.registry().resolve(model) + group_ids: Final = self.registry().names(str(group), kind="ids") + # A listed pass-through pick names only the pass-through deployments. + ids: Final = ( + [id for id in group_ids if _uses_pass_through(self._handle.deployment(id))] + if pass_through == "pass_through" + else list(group_ids) + ) + active: Final = self._cooldowns.get_active_cooldowns(list(ids), None) # pyright: ignore[reportArgumentType] # a list of ids + # `async_get_cooldown_deployments_with_debug_info`: every deployment cooling, in any group. + listed: Final = self._cooldowns.get_active_cooldowns(list(self.registry().names(kind="ids")), None) # pyright: ignore[reportArgumentType] # a list of ids + cooling: Final = [id for id, _ in listed] + min_cooldown: Final = min( + (_seconds(entry.get("cooldown_time"), self._cooldowns.default_cooldown_time) for _, entry in active), + default=_seconds(None, self._cooldowns.default_cooldown_time), + ) + return RouterRateLimitError( + model=self._process_renamed(model, str(group)), + cooldown_time=min_cooldown, + enable_pre_call_checks=self._resolved.enable_pre_call_checks, + cooldown_list=cooling, + model_ids=list(ids), # pyright: ignore[reportArgumentType] # a list of ids + cooldown_deployments=get_cooldown_deployment_info( + RoutingContext(self), # pyright: ignore[reportArgumentType] # reads registry() + ids, + cooling, + ), + ) + if last is not None: + return last + return RuntimeError(f"router-core ({kind}): {step.get('error')}") + + +def wildcard_request(model: str) -> str: + """``f"{provider}/{model}"``, the spelling ``PatternMatchRouter.get_pattern`` tries second.""" + from litellm.litellm_core_utils.get_llm_provider_logic import inferred_provider + + return f"{inferred_provider(model)}/{model}" + + +def _provider_request(step: Mapping[str, object]) -> Mapping[str, object]: + return {"models": [wildcard_request(str(step.get("model")))]} + + +def _lookup_error(kind: object, model: str, default_deployment: object) -> BaseException | None: + """The plain Python error a pick's lookup raised for ``model``, or None for any other failure.""" + if kind == "default_deployment_raises": + return default_deployment_copy_error(default_deployment, model) + if kind == "index_out_of_range": + # `self._core_model_list[index]` for a written index position past the list. + return IndexError("list index out of range") + if kind == "missing_routing_group": + # `_get_routing_context` reads the stale member map's group from the replaced groups. + return KeyError(model) + return None + + +def _origin_level(origin: Mapping[str, object]) -> int: + level: Final = origin.get("level") + return level if isinstance(level, int) else 0 + + +def _sleep_seconds(step: Mapping[str, object]) -> float: + """A retry's sleep: ``_time_to_sleep_before_retry`` returns a plain ``0`` when another + deployment is healthy, else a backoff in float seconds.""" + seconds: Final = float(str(step["seconds"])) + return 0 if seconds == 0 else seconds + + +def _uses_pass_through(record: object) -> bool: + params: Final = as_str_mapping((as_str_mapping(record) or {}).get("litellm_params")) or {} + return bool(params.get("use_in_pass_through")) + + +def _seconds(value: object, default: float) -> float | int: + """A cooldown time as stored: an int stays an int, so error texts read ``60`` like Python's.""" + if isinstance(value, int | float) and not isinstance(value, bool): + return value + return int(default) if float(default).is_integer() else default + + +def _ttl(step: Mapping[str, object]) -> float | None: + ttl: Final = step.get("ttl") + return ttl if isinstance(ttl, int | float) else None + + +def _amount(value: object) -> float | int: + return value if isinstance(value, int | float) and not isinstance(value, bool) else 1 + + +# The plans whose kwargs carry `original_generic_function`, which the fallback cooldown reads. +_GENERIC_API_KINDS: Final = frozenset(("generic", "search")) + + +class _NotSynchronous(TypeError): + """A sync attempt's host endpoint returned an awaitable: a host bug, never an attempt failure.""" + + +_MEMBER_READ_PICKS: Final = frozenset({"healthy", "pass_through", "sync_pick"}) + + +class _Selection: + """One ``select``: the Rust selection call, its select and cache steps answered on the host.""" + + def __init__( + self, + core: RustCore, + model: object, + request: dict[str, object], # mutable-ok: the caller's request, whose narrowing keys are popped + kind: str, + specific_deployment: bool | None, + messages: object, + input: object, + *, + probe: bool = False, + ) -> None: + self._core: Final = core + # `_async_get_available_deployment` reads the alias map first, every other pick its names. + self._model: Final = ( + model if isinstance(model, str) else core.routed_name(model, read_as_member=kind in _MEMBER_READ_PICKS) + ) + self._request: Final = request + self._kind: Final = kind + self._messages: Final = messages + self._input: Final = input + fields: Final = { + **core._routing_fields(request, messages), # pyright: ignore[reportPrivateUsage] # the selection serves its own core + **_narrowing_fields(request, kind), + } # pyright: ignore[reportPrivateUsage] # the selection serves its own core + # The pick's argument alone: a `specific_deployment` key in its request is no argument, and + # `_common_checks_available_deployment` reads nothing but True. + fields.pop("specific_deployment", None) + if specific_deployment is True: + fields["specific_deployment"] = True + # `filter_web_search_deployments` is in the async chain, which a strategy without an async + # selector never reaches on a pick. + if (kind in ("healthy", "async_pass_through") or (kind == "pick" and core._picks_async())) and ( # pyright: ignore[reportPrivateUsage] # the selection serves its own core + is_web_search_request(request) + ): + fields["web_search"] = True + # `_text_completion` wraps its sync pick in `endpoint_selection`. + endpoint: Final = selected_endpoint() + if endpoint is not None: + fields["endpoint"] = endpoint + if probe: + fields["health_check_probe"] = True + if kind in ("pick", "async_pass_through") or (kind == "sync_pick" and core._custom_strategy is not None): # pyright: ignore[reportPrivateUsage] # the selection serves its own core + fields["pre_routes"] = True + self._specific_deployment: Final = specific_deployment + # The error a callback filter raised, which a failed selection raises. + self._filter_error: BaseException | None = None + if kind == "pass_through": + fields["pass_through"] = "listed" + elif kind == "async_pass_through": + fields["pass_through"] = "healthy" + core._sync_routing_state() # pyright: ignore[reportPrivateUsage] # see above + handle: Final = core._handle # pyright: ignore[reportPrivateUsage] # see above + self._call: Final = (handle.healthy if kind == "healthy" else handle.select)(self._model, _document(fields)) + + def _done(self, step: Mapping[str, object]) -> Mapping[str, object] | Sequence[Mapping[str, object]]: + if self._kind == "healthy": + if step.get("single"): + # A deployment named by its id is listed alone, as a dict. + (record,) = _as_list(step.get("deployments")) + return for_requested_model(typed_deployment(as_str_mapping(record) or {}), self._model) + # The deployments as served: a priority group's members carry their priority as order. + return [ + for_requested_model(typed_deployment(as_str_mapping(record) or {}), self._model) + for record in _as_list(step.get("deployments")) + ] + if step.get("deployment_id") in (DEFAULT_DEPLOYMENT_FAILING_ID, DEFAULT_DEPLOYMENT_COPIED_ID): + # A default deployment every call to fails, or whose id is not a string, is picked as Python copies it. + picked: Final = default_deployment_copy(self._core._resolved.default_deployment, self._model) # pyright: ignore[reportPrivateUsage] # the selection serves its own core + return picked # pyright: ignore[reportReturnType] # the assigned value's copy, as the Python core returns it + return for_requested_model(typed_deployment(as_str_mapping(step.get("deployment")) or {}), self._model) + + def _failed(self, step: Mapping[str, object]) -> BaseException: + core: Final = self._core + error: Final = core._failure_error( # pyright: ignore[reportPrivateUsage] # see above + step, self._model, self._filter_error, pass_through=self._kind, request=self._request + ) + if self._kind not in ("pass_through", "sync_pick"): + _log_refused_pick(self._request, error) + return error + + async def run(self) -> Mapping[str, object] | Sequence[Mapping[str, object]]: + core: Final = self._core + outcome: object = None # rebind-ok: the answer to the last step + while True: + step = self._call.next(outcome) + match step["kind"]: + case "cache": + outcome = await core._answer_cache(step) # pyright: ignore[reportPrivateUsage] # see above + case "select" if step.get("stage") == "callback_filters": + outcome, self._filter_error = await core._answer_callback_filters( # pyright: ignore[reportPrivateUsage] # see above + step, self._request, self._messages + ) + case "select": + outcome = await core._answer_select( # pyright: ignore[reportPrivateUsage] # see above + step, + self._request, + self._messages, + self._input, + pass_through=self._kind == "async_pass_through", + ) + case "pre_route": + outcome = await core._pre_route( # pyright: ignore[reportPrivateUsage] # see above + step, + self._request, + self._messages, + self._input, + self._specific_deployment, + pass_through=self._kind == "async_pass_through", + ) + case "done": + done = self._done(step) + if self._kind != "pick" or not isinstance(done, Mapping): + return done + # `_core_async_get_available_deployment` resolves its pick to a pod. + return await self._resolved_pod(core, done) + case "failed": + raise self._failed(step) + case other: # pyright: ignore[reportUnnecessaryComparison] # a kind the bridge returns but the stub does not list + raise RuntimeError(f"router-core returned an unknown step {other!r}") + + async def _resolved_pod(self, core: "RustCore", done: Mapping[str, object]) -> Mapping[str, object]: + discovery: Final = core._pick_discovery(sync=False) # pyright: ignore[reportPrivateUsage] # the selection serves its own core + if discovery is None: + return done + resolved: Final[object] = await discovery.async_resolve_deployment(done, self._request) + return resolved # pyright: ignore[reportReturnType] # the deployment, resolved + + def _resolved_pod_sync(self, core: "RustCore", done: Mapping[str, object]) -> Mapping[str, object]: + discovery: Final = core._pick_discovery(sync=True) # pyright: ignore[reportPrivateUsage] # see _resolved_pod + return done if discovery is None else discovery.resolve_deployment(done, self._request) + + def run_sync(self) -> Mapping[str, object] | Sequence[Mapping[str, object]]: + core: Final = self._core + outcome: object = None # rebind-ok: the answer to the last step + while True: + step = self._call.next(outcome) + match step["kind"]: + case "cache": + outcome = core._answer_cache_sync(step) # pyright: ignore[reportPrivateUsage] # see above + case "select": + outcome = core._answer_select_sync(step, self._request, self._messages, self._input) # pyright: ignore[reportPrivateUsage] # see above + case "pre_route": + outcome = core._custom_pick_sync( # pyright: ignore[reportPrivateUsage] # see above + step, self._request, self._messages, self._input, self._specific_deployment + ) + case "done": + done = self._done(step) + if self._kind != "sync_pick" or not isinstance(done, Mapping): + return done + # `get_available_deployment` resolves its pick to a pod. + return self._resolved_pod_sync(core, done) + case "failed": + raise self._failed(step) + case other: # pyright: ignore[reportUnnecessaryComparison] # a kind the bridge returns but the stub does not list + raise RuntimeError(f"router-core returned an unknown step {other!r}") + + +def _log_refused_pick(request: Mapping[str, object], error: BaseException) -> None: + """A refused pick is logged to the request's logging object (`_async_get_available_deployment`).""" + dispatch: Final[object] = getattr(request.get("litellm_logging_obj"), "dispatch_failure_handlers", None) + if not callable(dispatch): + return + import asyncio + import traceback + + trace: Final = "Traceback (most recent call last):\n" + "".join( + traceback.format_exception(type(error), error, error.__traceback__) + ) + pending: Final[object] = dispatch(exception=error, traceback_exception=trace, prefer_async_handlers=True) + if inspect.isawaitable(pending): + _ = asyncio.ensure_future(pending) + + +def _log_refused_image_pick(request: Mapping[str, object], error: BaseException) -> None: + """``_aimage_generation``'s refused pick, for a request with no logging object of its own.""" + from litellm._uuid import uuid + from litellm.router_rust.invoke.endpoints import ROUTING_KEYS + from litellm.utils import ( + Rules, + _restore_correlation_context_if_supported, # pyright: ignore[reportPrivateUsage] # same task-local cleanup as the Python router + function_setup, + ) + + kwargs: Final = without(request, *(ROUTING_KEYS - {"model"})) + kwargs.setdefault("litellm_call_id", str(uuid.uuid4())) + logging_obj, _ = function_setup("aimage_generation", Rules(), datetime.now(), **kwargs) # noqa: DTZ005 # SDK logging uses naive local timestamps # pyright: ignore[reportArgumentType] # opaque SDK options, as the Python router passes them + try: + _log_refused_pick({"litellm_logging_obj": logging_obj}, error) + finally: + _restore_correlation_context_if_supported(logging_obj) + + +def _custom_pick(picked: object) -> Mapping[str, object]: + """A custom strategy's answer as the deployment document the Rust core serves.""" + document: Final = as_str_mapping(picked) + if document is None: + raise TypeError(f"the custom routing strategy returned {type(picked).__name__}, not a deployment") + return _document(document) + + +def _with_target_order(kwargs: Mapping[str, object], order: object) -> Mapping[str, object]: + """The attempt's kwargs with the order level its hop targets, as `run_async_fallback` sets + `_target_order` before `_acreate_file` picks every healthy deployment of that level.""" + if not isinstance(order, int) or isinstance(order, bool): + return kwargs + return {**kwargs, "_target_order": order} + + +def _narrowing_fields(request: dict[str, object], kind: str) -> Mapping[str, object]: # mutable-ok: popped + """``_target_order``, ``_excluded_deployment_ids`` and ``_retry_skipped_deployment_ids``, popped + from the request as the healthy chain pops them (a listed pass-through pick only reads the order).""" + fields: Final = _request_narrowing(request) + if kind == "pass_through": + return {key: value for key, value in fields.items() if key == "target_order"} + for key in ("_target_order", "_excluded_deployment_ids", "_retry_skipped_deployment_ids"): + request.pop(key, None) + return fields + + +def _request_narrowing(request: Mapping[str, object]) -> Mapping[str, object]: + """The narrowing keys a request carries, as the fields router-core reads them.""" + from litellm.litellm_core_utils.routing_common_utils import ( + as_excluded_deployment_ids, + as_retry_skipped_deployment_ids, + ) + + target_order: Final = request.get("_target_order") + excluded: Final = as_excluded_deployment_ids(request.get("_excluded_deployment_ids")) + skipped: Final = as_retry_skipped_deployment_ids(request.get("_retry_skipped_deployment_ids")) + return { + **( + {"target_order": target_order} + if isinstance(target_order, int) and not isinstance(target_order, bool) + else {} + ), + **({"excluded_ids": list(excluded)} if excluded else {}), + **({"retry_skipped_ids": list(skipped)} if skipped else {}), + } + + +_NO_ROUTING_PRE_CALL_CHECKS: Final = frozenset({"search", "rerank", "moderation"}) +# Attempts whose deployment is no pick (`_core_async_get_available_deployment` / +# `get_available_deployment`), so Kubernetes pod discovery never resolves it. +_UNPICKED_KINDS: Final = frozenset({"search", "create_file"}) + + +class _Driver: + """Answers the attempt and fallback steps of one call, and shapes its result. + + Each attempt is one host call against the deployment Rust picked. The driver keeps the error + each fallback level ended on: a level that fails raises its own error, whatever its fallbacks + raised, as the Python router does, unless an order chain handed it another + (``order_chain_failed``). + """ + + def __init__( + self, + core: RustCore, + plan: Plan, + request: Mapping[str, object], + cut_deployment: object = None, + stream_failure: tuple[int, BaseException] | None = None, + *, + sync: bool = False, + fresh: bool = False, + passthrough: "Callable[[Mapping[str, object], str], Awaitable[object]] | None" = None, + ) -> None: + self._core: Final = core + self._sync: Final = sync + self._passthrough: Final = passthrough + self._plan: Final = plan + self._op: Final = plan.endpoint + self._request: Final = request + self._cut_deployment: Final = cut_deployment if isinstance(cut_deployment, str) else None + self._last_response: object = None + self._last_error: BaseException | None = None + # The level the call raises from: a fallback walk re-entered from a failed stream starts at + # the level that stream served, on the failure ``async_function_with_fallbacks_common_utils`` + # was handed, which is what it raises when no fallback answers. + self._base_level: Final = stream_failure[0] if stream_failure is not None else 0 + self._level_errors: Final[dict[int, BaseException]] = dict( # mutable-ok: filled as the call runs + (stream_failure,) if stream_failure else () + ) + # The retry counts an exhausted retry loop stamps on the pick error its level then raises. + self._pending_stamps: Final[dict[int, tuple[object, object]]] = {} # mutable-ok: filled as the call runs + # Each level's current hop's first failure, once it retries: the error its retry loop raises + # last is raised while that one is handled, so it chains to it. + self._hop_first: Final[dict[int, BaseException]] = {} # mutable-ok: filled as the call runs + # Each level's latest pick error: the one object it logged, recorded and raises. + self._level_picks: Final[dict[int, BaseException]] = {} # mutable-ok: filled as the call runs + self._served: tuple[str | None, Mapping[str, object]] = (None, MappingProxyType({})) + # The latest attempt's deployment and call, which a blocked response is raised from. + self._latest_call: tuple[Mapping[str, object], AttemptCall] | None = None + # A sync attempt's stream is reading its first chunk, so a failure now is the stream's to log. + self._reading_first_chunk: bool = False + self._attempts: int = 0 + # When the latest attempt started: Python reads it off the logging object, when one is set up. + self._attempt_started: datetime | None = None + # A single-attempt endpoint runs neither of the wrappers that stamp the walk on the request. + self._walk: Final = WalkMetadata( + request, plan.metadata_field, self._params_base_model, live=fresh and not plan.single_attempt + ) + self.outcomes: Final = FallbackOutcomes(request) + # The fallback level the walk is at, from its fallback events. + self._level: int = self._base_level + # The primary level reached the end of `async_function_with_fallbacks_common_utils`, the + # only place the outcome message is appended: not its early raises, nor its order path. + self._explains_failure: bool = False + + def _blocked(self, response: object, endpoint_name: str) -> str | None: + """How the provider blocked a served response, the first check of + ``_should_raise_content_policy_error`` (a chat response with no choices, or whose first + choice was content-filtered) and of ``_should_raise_anthropic_refusal_error`` (a + /v1/messages safeguard refusal carrying ``stop_details``).""" + from litellm.llms.anthropic.messages.native.utils import get_safeguard_refusal_stop_details + from litellm.types.utils import ModelResponse + + if endpoint_name == "anthropic_messages" and get_safeguard_refusal_stop_details(response) is not None: + return "refusal" + if self._op not in ("acompletion", "completion") or not isinstance(response, ModelResponse): + return None + return ( + "content_filter" if not response.choices or response.choices[0].finish_reason == "content_filter" else None + ) + + def _blocked_error(self, event: Mapping[str, object]) -> BaseException: + """The error the router raises for a blocked response a fallback would serve.""" + import litellm.exceptions as litellm_exceptions + from litellm.llms.anthropic.messages.native.utils import ( + get_safeguard_refusal_stop_details, + safeguard_refusal_error, + ) + + model: Final = str(event.get("model")) + stop_details: Final = get_safeguard_refusal_stop_details(self._last_response) + if event.get("blocked") == "refusal" and stop_details is not None: + return safeguard_refusal_error(model=model, stop_details=stop_details) + return litellm_exceptions.ContentPolicyViolationError( + message="Response output was blocked.", model=model, llm_provider="" + ) + + def _stamp_blocked(self, error: BaseException) -> None: + """The attempt helper's ``except`` runs for the error it raises over a blocked response: the + response then counts as a failure, not a success.""" + if self._latest_call is None: + return + deployment, call = self._latest_call + endpoints: Final = self._core._endpoints # pyright: ignore[reportPrivateUsage] # the driver serves its own core + stamp_error(self._plan, error, deployment, call.served) + endpoints.counters.retract_success(self._plan, call) + endpoints.attempt_failed(self._plan, error, deployment, call, call.group) + + def _params_base_model(self, deployment_id: str) -> str | None: + """``_deployment_params_base_model``: the ``base_model`` under the deployment's ``litellm_params``.""" + record: Final = as_str_mapping(self._core._handle.deployment(deployment_id)) or {} # pyright: ignore[reportPrivateUsage] # the driver serves its own core + base_model: Final = (as_str_mapping(record.get("litellm_params")) or {}).get("base_model") + return base_model if isinstance(base_model, str) else None + + def _logged_error(self, event: Mapping[str, object]) -> BaseException | None: + """The error a ``retry_logged`` event records.""" + return self._cause_error(event) + + def _policy_raised(self, event: Mapping[str, object]) -> None: + """A hop's first failure looked its group's retry policy up and the lookup raised: the hop + ends with that error, raised from the failure the event names.""" + model: Final = str(event.get("model")) + original: Final = self._cause_error(event) or RuntimeError(f"router-core: the hop serving {model} failed") + policies: Final = self._request.get( + "model_group_retry_policy", + self._core._resolved.model_group_retry_policy, # pyright: ignore[reportPrivateUsage] # the driver serves its own core + ) + raised: Final = policy_lookup_error(original, model, self._core._written_retry_policy, policies) # pyright: ignore[reportPrivateUsage] # see above + level: Final = event.get("level") + self._last_error = raised + self._level_errors[level if isinstance(level, int) else 0] = raised + + def _stamped_budget(self, event: Mapping[str, object]) -> Mapping[str, object]: + """A ``retries_opened`` event as ``_metadata["max_retries"]`` reads it: a request + ``num_retries`` that is no count is stamped as given.""" + given: Final = self._request.get("num_retries") + return {**event, "max_retries": given} if "uncounted_retries" in request_retry_fields(given) else event + + def _count_raised(self, event: Mapping[str, object]) -> None: + """A hop's first failure reached the retry loop's read of a request ``num_retries`` that + raises: the hop ends with that error, raised from the failure the event names.""" + model: Final = str(event.get("model")) + original: Final = self._cause_error(event) or RuntimeError(f"router-core: the hop serving {model} failed") + raised: Final = retry_count_error(original, self._request.get("num_retries")) + level: Final = event.get("level") + self._last_error = raised + self._level_errors[level if isinstance(level, int) else 0] = raised + + def _cause_error(self, event: Mapping[str, object]) -> BaseException | None: + """The error the ``cause`` of a ``retry_logged`` or ``retry_policy_raised`` event names.""" + match event.get("cause"): + case "level": + level: Final = event.get("level") + return self._level_errors.get(level if isinstance(level, int) else 0, self._last_error) + case "pick": + at_level: Final = event.get("level") + held: Final = self._level_picks.get(at_level if isinstance(at_level, int) else 0) + return held if held is not None else self._picked_error(event) + case _: + return self._last_error + + def _picked_error(self, event: Mapping[str, object]) -> BaseException: + """The error a pick an event names raised, as a failed step with that error raises it.""" + model: Final = str(event.get("failed_model") or event.get("model")) + return self._core._failure_error(event, model, None, request=self._request) # pyright: ignore[reportPrivateUsage] # the driver serves its own core + + def _pick_raised(self, error: Exception, *, logged: bool) -> Mapping[str, object]: + """A pick raised on the host (a strategy, a tag filter, a pre-routing hook): the try fails with + that error, which the Rust walk retries and falls back from like any attempt's. ``logged`` + when the async pick's ``except`` reports it, as ``_async_get_available_deployment`` does.""" + if isinstance(error, RustCoreUnsupported): + raise error + if logged: + self._log_pick_failure(error) + else: + mark_deployment_selection_failure(error) + self._chain_to_hop_first(error, self._level) + self._last_error = error + self._level_errors[self._level] = error + return {**attempt_failure(error), "raised": True} + + def _log_pick_failure(self, error: BaseException) -> None: + """The ``except`` of ``_async_get_available_deployment``, which only the async selection runs: + the router's strategy may have no async selector, and a custom strategy replaces the method. + ``_aimage_generation`` logs a refused pick itself when the request set up no logging object.""" + core: Final = self._core + assigned: Final = "_core_async_get_available_deployment" in core.__dict__ + custom: Final = assigned or core._custom_strategy is not None # pyright: ignore[reportPrivateUsage] # the driver serves its own core + if not custom and core._picks_async(): # pyright: ignore[reportPrivateUsage] # see above + mark_deployment_selection_failure(error) + _log_refused_pick(self._request, error) + if self._plan.kind == "image" and self._request.get("litellm_logging_obj") is None: + _log_refused_image_pick(self._request, error) + + def _pick_failed(self, event: Mapping[str, object]) -> None: + """A pick failed in the Rust core: its error is raised once, however many events name it.""" + error: Final = self._picked_error(event) + level: Final = event.get("level") + self._level_picks[level if isinstance(level, int) else 0] = error + if not self._sync: + self._log_pick_failure(error) + + def _unroutable_error(self, model: object) -> BaseException: + """What a hop to a dict target's non-string ``model`` raises: no group answers ``None``, and + any other value fails the level's fallback lookup, which hashes it and then searches it.""" + if model is None: + return self._core._failure_error({"error_kind": "unknown_model"}, "None", None) # pyright: ignore[reportPrivateUsage] # the driver serves its own core + try: + _ = model in {} or "/" in model # pyright: ignore[reportOperatorIssue] # raising is the point + except TypeError as error: + return error + return TypeError(f"{model!r} is not a model group") + + def _level_failed(self, event: Mapping[str, object]) -> None: + """A level raises the error a pick raised (``level_failed``).""" + level: Final = event.get("level") + at_level: Final = level if isinstance(level, int) else 0 + picked: Final = ( + self._picked_error(event) + if event.get("handling") is True + else self._level_picks.pop(at_level, None) or self._picked_error(event) + ) + if event.get("handling") is True: + # Raised by the level's fallback handling, while the level's error was handled. + picked.__context__ = self._level_errors.get(at_level, self._last_error) + self._level_errors[at_level] = picked + if at_level == self._base_level: + self._explains_failure = False + return + stamp: Final = self._pending_stamps.pop(at_level, None) + if stamp is not None and type(picked) in self._core._host.dynamic().exception_types: # pyright: ignore[reportPrivateUsage] # the driver serves its own core + # The exhausted retry loop raised this pick's error: stamped as `retries_exhausted` stamps. + picked.max_retries, picked.num_retries = stamp # pyright: ignore[reportAttributeAccessIssue] # Python's router stamps the same attributes + self._chain_to_hop_first(picked, at_level) + self._level_errors[at_level] = picked + + def _carry_superseded_usage(self) -> None: + """``carry_superseded_attempt_usage``, before each same-group retry and each fallback hop.""" + from litellm.router_utils.fallback_event_handlers import carry_superseded_attempt_usage + + carry_superseded_attempt_usage(self._request) + + def _order_chain_failed(self, event: Mapping[str, object]) -> None: + """The order path of ``async_function_with_fallbacks_common_utils`` raised: the chain's last + hop's error, which the hop's level ended on, unless the level keeps its own.""" + level: Final = event.get("level") + at: Final = level if isinstance(level, int) else 0 + last_hop: Final = self._level_errors.get(at + 1) + if event.get("last_hop") is True and last_hop is not None: + self._level_errors[at] = last_hop + if at == self._base_level: + self._explains_failure = False + + def _logged_off_the_loop(self) -> bool: + """Whether litellm logged the attempt's failure on a thread with no running loop. An async + call's mocked failure raises in ``completion``, which ``acompletion`` runs in an executor + thread; a sync stream's first chunk logs in a ``threading.Thread``.""" + from litellm.utils import should_run_mock_completion + + if self._sync: + return self._reading_first_chunk + kwargs: Final[Mapping[str, object]] = self._latest_call[1].kwargs if self._latest_call is not None else {} + return self._plan.kind in ("chat", "text") and should_run_mock_completion( + kwargs.get("mock_response"), kwargs.get("mock_tool_calls"), kwargs.get("mock_timeout") + ) + + def _cooldown_set(self, event: Mapping[str, object]) -> None: + """An attempt's failure cooled its deployment down: ``set_cooldown_deployments`` as the failure + callback runs it. Outside a running loop it raises, which ``CustomLogger.log_event`` swallows; + a sync walk with a private loop has one (``cooldown_event_on_own_loop``).""" + deployment_id: Final = event.get("deployment_id") + seconds: Final = event.get("cooldown_time") + if self._logged_off_the_loop() or not isinstance(deployment_id, str): + return + fire: Final = self._core.cooldown_event_on_own_loop if self._plan.on_its_own_loop else self._core.cooldown_event + with contextlib.suppress(RuntimeError): + fire( + deployment_id, + getattr(self._last_error, "status_code", ""), + self._python_cooldown_time(seconds) if isinstance(seconds, float) else None, + ) + + def _python_cooldown_time(self, seconds: float) -> float: + """``_time_to_cooldown`` of ``deployment_callback_on_failure`` for the attempt that failed.""" + deployment: Final[Mapping[str, object]] = self._latest_call[0] if self._latest_call is not None else {} + configured: Final = first_cooldown( + as_str_mapping(deployment.get("model_info")), + self._request, + as_str_mapping(deployment.get("litellm_params")), + ) + return self._core.python_cooldown(seconds, configured, self._last_error) + + def _selected(self, event: Mapping[str, object]) -> None: + """``_acompletion``'s service event for the pick it awaited, and its ``_track_deployment_metrics``.""" + seconds: Final = event.get("pick_seconds") + if self._plan.kind == "chat" and not self._sync and isinstance(seconds, float): + end_time: Final = unix_time() + self._core.service_logged("async_get_available_deployment", end_time - seconds, end_time, self._request) + self._core.track_deployment_pick(event.get("deployment_id"), self._request) + + def _fallbacks_failed(self, event: Mapping[str, object]) -> None: + """The ``except`` around the regular fallbacks of ``async_function_with_fallbacks_common_utils``: + it logs what the ``try`` raised, the last hop's error or the level's own. A level with no + target raises its own error into it, except where no fallback list was set to search.""" + import litellm.exceptions as litellm_exceptions + from litellm.litellm_core_utils.core_helpers import get_parent_otel_span_from_kwargs + + level: Final = event.get("level") + at: Final = level if isinstance(level, int) else 0 + own: Final = self._level_errors.get(at, self._last_error) + request: Final = self._request + walked: Final = self._core._walked_fallbacks # pyright: ignore[reportPrivateUsage] # the driver serves its own core + if event.get("chain") is not True: + policy_list: Final = request.get("content_policy_fallbacks", walked["content_policy_fallbacks"]) + if not isinstance(own, litellm_exceptions.ContentPolicyViolationError) or policy_list is None: + if request.get("fallbacks", walked["fallbacks"]) is None: + return + self._explain_no_fallback_group(own, at) + raised: Final = self._level_errors.get(at + 1, own) if event.get("last_hop") is True else own + self.outcomes.failed_fallbacks( + RoutingContext(self._core), # pyright: ignore[reportArgumentType] # the read uses registry() and health() + truncate_fallback_error_detail(redact_string(str(raised))), + get_parent_otel_span_from_kwargs(dict(request)), + ) + + def _explain_no_fallback_group(self, primary: BaseException | None, at: int) -> None: + """The primary level names the missing fallback group on its error before raising it into + the ``except``, so the log reads it and no outcome message follows.""" + if ( + at == self._base_level == 0 + and self._core._host.dynamic().expose_router_debug_in_errors # pyright: ignore[reportPrivateUsage] # the driver serves its own core + and primary is not None + and hasattr(primary, "message") + ): + primary.message += self._explanation(primary, str(self._request.get("model"))) # pyright: ignore[reportAttributeAccessIssue, reportUnknownMemberType] # checked above # rebind-ok: the message is what this explains + self._explains_failure = False + + def take_events(self, events: object) -> None: # noqa: C901 # one branch per router-core event + for raw_event in _as_list(events): + event = as_str_mapping(raw_event) + if event is None: + continue + if event.get("kind") == "retry": + self._carry_superseded_usage() + continue + if event.get("kind") == "cooldown_set": + self._cooldown_set(event) + continue + if event.get("kind") == "selected": + self._selected(event) + continue + if event.get("kind") == "retries_opened": + self._walk.retries_opened(self._stamped_budget(event)) + continue + if event.get("kind") == "retry_budget": + self._walk.retry_budget(event) + continue + if event.get("kind") == "pick_failed": + self._pick_failed(event) + continue + if event.get("kind") == "order_chain_failed": + self._order_chain_failed(event) + continue + if event.get("kind") == "mock_failed": + resolved = self._core._resolved # pyright: ignore[reportPrivateUsage] # the driver serves its own core + mocked = mock_error( + str(event.get("mock")), + str(event.get("model")), + event.get("num_retries"), + self._request, + { + "fallbacks": resolved.fallbacks, + "context_window_fallbacks": resolved.context_window_fallbacks, + "content_policy_fallbacks": resolved.content_policy_fallbacks, + }, + ) + level = event.get("level") + self._last_error = mocked + self._level_errors[level if isinstance(level, int) else 0] = mocked + continue + if event.get("kind") == "retry_policy_raised": + self._policy_raised(event) + continue + if event.get("kind") == "retry_count_raised": + self._count_raised(event) + continue + if event.get("kind") == "level_failed": + self._level_failed(event) + continue + if event.get("kind") == "fallbacks_failed": + self._fallbacks_failed(event) + continue + if event.get("kind") == "fallback": + hop_to = event.get("level") + self._hop_first.pop(hop_to if isinstance(hop_to, int) else 0, None) + self._walk.hop(event) + self._carry_superseded_usage() + hop_level = event.get("level") + self._level = hop_level if isinstance(hop_level, int) else self._level + if "unroutable" in event: + self._level_errors[self._level] = self._unroutable_error(event["unroutable"]) + continue + if event.get("kind") == "blocked": + level = event.get("level") + blocked = self._blocked_error(event) + self._stamp_blocked(blocked) + self._last_error = blocked + self._level_errors[level if isinstance(level, int) else 0] = blocked + continue + if event.get("kind") == "falling_back": + level = event.get("level") + at = level if isinstance(level, int) else 0 + self._explains_failure = self._explains_failure or at == self._base_level + self._hint_level_failure(self._level_errors.get(at, self._last_error), str(event.get("model"))) + continue + if event.get("kind") == "fallback_done": + level = event.get("level") + at = level if isinstance(level, int) else 0 + self._level = at + self.outcomes.done( + event, + self._level_errors.get(at, self._last_error), + self._level_errors.get(at + 1), + self._walk.hop_kwargs(event), + ) + if event.get("succeeded") is not True: + self._core.trigger_fallback_cooldown( + self._request, + self._level_errors.get(at + 1), + self._plan.kind in _GENERIC_API_KINDS, + own_loop=self._plan.on_its_own_loop, + ) + continue + if event.get("kind") == "retry_logged": + logged = self._logged_error(event) + self._walk.log_retry(event, logged) + logged_level = event.get("level") + if event.get("attempted_retries") == 0 and logged is not None: + self._hop_first[logged_level if isinstance(logged_level, int) else 0] = logged + continue + if event.get("kind") != "retries_exhausted": + continue + if event.get("picked") is True: + # The level raises the last retry's pick error, which its `level_failed` names. + exhausted_level = event.get("level") + self._pending_stamps[exhausted_level if isinstance(exhausted_level, int) else 0] = ( + event.get("max"), + event.get("attempted"), + ) + continue + error: BaseException | None = self._last_error + # `async_function_with_retries` stamps the error its exhausted retry loop raises. + if error is not None and type(error) in self._core._host.dynamic().exception_types: # pyright: ignore[reportPrivateUsage] # the driver serves its own core + error.max_retries = event.get("max") # pyright: ignore[reportAttributeAccessIssue] # Python's router stamps the same attributes + error.num_retries = event.get("attempted") # pyright: ignore[reportAttributeAccessIssue] # same + + def _call(self, step: Mapping[str, object]) -> "tuple[Mapping[str, object], AttemptCall]": + name: Final = step.get("model") + requested: Final = str(name) if name is not None else str(self._request.get("model")) + deployment: Final = ( + # A default deployment whose id is not a string is called as Python copies it, id and all. + as_str_mapping(default_deployment_copy(self._core._resolved.default_deployment, requested)) or {} # pyright: ignore[reportPrivateUsage] # the driver serves its own core + if step.get("deployment_id") == DEFAULT_DEPLOYMENT_COPIED_ID + else for_requested_model(typed_deployment(as_str_mapping(step.get("deployment")) or {}), requested) + ) + self._attempts += 1 + self._attempt_started = datetime.now() # noqa: DTZ005 # naive to match the logging pipeline's api_call_start_time + level: Final = step.get("fallback_level") + # The request params the walk's dict targets merged in (`kwargs.update(target)`). + params: Final = as_str_mapping(step.get("params")) + call: Final = attempt_call( + self._core._port, # pyright: ignore[reportPrivateUsage] # the driver serves its own core + self._plan, + self._request if params is None else {**self._request, **params}, + deployment, + group=requested, + requested=requested, + fallback_level=level if isinstance(level, int) else 0, + ) + self._strip_undecryptable_reasoning(step, deployment, call.kwargs) + shaped: Final = shape_attempt_kwargs(call.kwargs, deployment) + # The walk's metadata for the attempt's level (retries, fallback hops) and its hop kwargs. + call.kwargs = self._walk.stamp(step, shaped) + field: Final = self._plan.metadata_field + if isinstance(call.served, dict) and field in call.kwargs: + call.served[field] = call.kwargs[field] + self._served = (requested, call.served) + self._latest_call = (deployment, call) + if call.refused is None: + self._core._endpoints.counters.started(self._plan, call) # pyright: ignore[reportPrivateUsage] # the driver serves its own core + return deployment, call + + def _strip_undecryptable_reasoning( + self, step: Mapping[str, object], deployment: Mapping[str, object], kwargs: Mapping[str, object] + ) -> None: + """``_update_kwargs_with_deployment`` on a fallback hop (``_is_fallback_hop``): the encrypted + reasoning the target cannot decrypt is dropped; reasoning with no origin marker came from + the deployment that just failed.""" + depth: Final = step.get("fallback_depth") + if not isinstance(depth, int) or depth <= 0: + return + EncryptedContentAffinityCheck.strip_reasoning_the_targets_cannot_decrypt( + RoutingContext(self._core), # pyright: ignore[reportArgumentType] # the check reads registry() + kwargs.get("input"), + kwargs.get("messages"), + (deployment,), + unmarked_origin=self._walk.deployment_that_just_failed(step) or self._cut_deployment, + ) + + @property + def response(self) -> object: + return self._last_response + + @property + def last_error(self) -> BaseException | None: + return self._last_error + + def _stream_attempt( + self, step: Mapping[str, object], kwargs: Mapping[str, object], slot: "contextlib.ExitStack | None" = None + ) -> StreamAttempt: + def number(key: str) -> int: + value: Final = step.get(key) + return value if isinstance(value, int) else 0 + + order: Final = step.get("target_order") + deployment_id: Final = step.get("deployment_id") + return StreamAttempt( + core=self._core, + op=self._op, + request=self._request, + kwargs=kwargs, + group=str(step.get("model") or step.get("model_group")), + order=order if isinstance(order, int) else None, + depth=number("fallback_depth"), + level=number("fallback_level"), + deployment_id=deployment_id if isinstance(deployment_id, str) else None, + attempted_retries=number("retry"), + slot=slot, + ) + + def _failed( + self, + step: Mapping[str, object], + deployment: Mapping[str, object], + call: "AttemptCall | None", + error: BaseException, + ) -> Mapping[str, object]: + # What litellm's failure logging read, before the attempt helper's `except` wrote on it. + detail: Final = str(error) + served: Final = call.served if call is not None else {} # pyright: ignore[reportUnknownVariableType] # litellm's untyped helpers and request values + stamp_error(self._plan, error, deployment, served) # pyright: ignore[reportUnknownArgumentType] # litellm's untyped helpers and request values + self._core._endpoints.attempt_failed(self._plan, error, deployment, call, str(step.get("model"))) # pyright: ignore[reportPrivateUsage] # see above + self._last_error = error + level: Final = step.get("fallback_level") + retry: Final = step.get("retry") + if isinstance(retry, int) and retry > 0: + self._chain_to_hop_first(error, level if isinstance(level, int) else 0) + self._level_errors[level if isinstance(level, int) else 0] = error + logged: Final = logged_call( + call.host_op if call is not None else self._plan.host_op, + served or self._request, # pyright: ignore[reportUnknownArgumentType] # see above + self._attempt_started, + ) + model_info: Final = as_str_mapping(deployment.get("model_info")) or {} + return { + **attempt_failure(error), + "detail": detail, + **self._core.failure_exits(error, logged, model_info.get("id")), + # `async_deployment_callback_on_failure` does not run for a sync call, and skips a batch retrieve. + "uncounted_request": self._plan.sync or is_batch_retrieve_call_type(logged.get("call_type")), + } + + def _chain_to_hop_first(self, error: BaseException, level: int) -> None: + """A retry's failure is raised while the hop's first failure is handled (the retry loop + runs in its ``except`` block), so it chains to it.""" + first: Final = self._hop_first.get(level) + if first is not None and first is not error: + error.__context__ = first # rebind-ok: chained as the retry loop's raise chains it + + def _succeeded( + self, call: "AttemptCall", response: object, host_op: str, stream: "StreamAttempt | None" = None + ) -> Mapping[str, object]: + self._last_response = response + tokens: Final = total_tokens(response) + # A wrapped stream reports its success or failure when it ends; the attempt counts the + # request and the tokens read so far, which the end's report leaves out. + # A realtime attempt counts its session's success itself (`attempt`), when it ended cleanly. + deferred: Final = (stream is not None and stream.reports_outcome) or reports_own_success(call) + if stream is not None and deferred: + stream.counted_tokens = tokens # rebind-ok: what the end's report leaves out + return { + "ok": True, + "total_tokens": tokens, + "blocked": self._blocked(response, host_op), + "deferred": deferred, + "tracks_usage": self._tracks_usage(call, response, deferred), + } + + def _tracks_usage(self, call: "AttemptCall", response: object, deferred: bool) -> bool: + """Whether the attempt counts against its deployment's usage (``_deployment_tracks_usage``). A + ``None`` response counts nothing (``increment_deployment_usage_for_response``), and an attempt + that reports no success then is a realtime session, which logs none to count it either.""" + if deferred and response is None: + return False + model: Final = call.kwargs.get("model") + return self._core.deployment_tracks_usage( + (as_str_mapping(call.kwargs.get("model_info")) or {}).get("id"), + model if isinstance(model, str) else call.group, + ) + + async def _refused( + self, + step: Mapping[str, object], + deployment: Mapping[str, object], + call: "AttemptCall", + error: BaseException, + ) -> Mapping[str, object]: + """A routing pre-call check refused the attempt (``async_routing_strategy_pre_call_checks``): + a rate limit cools the deployment down (``set_cooldown_deployments``); the attempt's failure + is recorded by nothing else.""" + import litellm.exceptions as litellm_exceptions + + core: Final = self._core + if isinstance(error, litellm_exceptions.RateLimitError): + deployment_id: Final = (as_str_mapping(deployment.get("model_info")) or {}).get("id") + if isinstance(deployment_id, str): + failure: Final = {**attempt_failure(error), "caller_scoped": core.per_user_deployment(deployment_id)} + report: Final = core._handle.report_failure(deployment_id, failure, uncounted=True) # pyright: ignore[reportPrivateUsage] # the driver serves its own core + done: Final = await core._drive_cache_steps(report) # pyright: ignore[reportPrivateUsage] # see above + _ = core.cooled_down(done, deployment_id, error.status_code) + return {**self._failed(step, deployment, call, error), "recorded": True} + + def _pick_kwargs(self, step: Mapping[str, object]) -> Mapping[str, object]: + """The request kwargs the pick sees: the request with its level's retry count.""" + count: Final = self._walk.request_retry_count(step) + if count is None: + return self._request + field: Final = self._plan.metadata_field + metadata: Final = as_str_mapping(self._request.get(field)) or {} + return {**self._request, field: {**metadata, "request_retry_count": count}} + + def _with_deployment(self, step: Mapping[str, object], resolved: object) -> Mapping[str, object]: + return step if resolved is step.get("deployment") else {**step, "deployment": resolved} + + async def _resolve_pods(self, step: Mapping[str, object]) -> Mapping[str, object]: + """``async_resolve_pods_after``: a picked deployment with Kubernetes pod discovery on is + sent to one of its service's pods.""" + if self._plan.kind in _UNPICKED_KINDS or step.get("deployment") is None: + return step + discovery: Final = self._core._pick_discovery(sync=False) # pyright: ignore[reportPrivateUsage] # the driver serves its own core + if discovery is None: + return step + resolved: Final[object] = await discovery.async_resolve_deployment( + step.get("deployment"), self._pick_kwargs(step) + ) + return self._with_deployment(step, resolved) + + def _resolve_pods_sync(self, step: Mapping[str, object]) -> Mapping[str, object]: + """``resolve_pods_after``, on the sync pick.""" + if self._plan.kind in _UNPICKED_KINDS or step.get("deployment") is None: + return step + discovery: Final = self._core._pick_discovery(sync=True) # pyright: ignore[reportPrivateUsage] # see _resolve_pods + if discovery is None: + return step + return self._with_deployment( + step, discovery.resolve_deployment(step.get("deployment"), self._pick_kwargs(step)) + ) + + def _default_copy(self, step: Mapping[str, object]) -> Mapping[str, object]: + """The default deployment as Python copied it for the call that fails on it: the error is + stamped from this copy, never from the id router-core serves it under.""" + copied: Final = default_deployment_copy(self._core._resolved.default_deployment, str(step.get("model"))) # pyright: ignore[reportPrivateUsage] # the driver serves its own core + return as_str_mapping(copied) or {} + + def _default_call_error(self, step: Mapping[str, object]) -> BaseException | None: + """The error a call to a default deployment every call fails to raises (its first + reads of the copy, before any endpoint runs): nothing records it against a deployment.""" + if step.get("deployment_id") != DEFAULT_DEPLOYMENT_FAILING_ID: + return None + return default_deployment_call_error(self._core._resolved.default_deployment) # pyright: ignore[reportPrivateUsage] # the driver serves its own core + + async def _passthrough_attempt(self, step: Mapping[str, object]) -> Mapping[str, object]: + """``_ageneric_api_call_with_fallbacks_helper`` with ``passthrough_on_no_deployment``: the pick + failed, so the handler is called under the attempt's name. Nothing was picked, so nothing is + stamped on a failure or recorded against a deployment.""" + passthrough: Final = self._passthrough + if passthrough is None: + raise RuntimeError("router-core asked for an unrouted attempt the host did not offer") + name: Final = str(step.get("model")) + params: Final = as_str_mapping(step.get("params")) + self._attempts += 1 + self._attempt_started = datetime.now() # noqa: DTZ005 # see _call + try: + answered: Final = await passthrough(self._request if params is None else {**self._request, **params}, name) + except Exception as error: # noqa: BLE001 # handed to the Rust call loop like any attempt failure + return {**self._failed(step, {}, None, error), "recorded": True} + self._last_response = answered + self._served = (name, self._request) + return {"ok": True, "total_tokens": total_tokens(answered)} + + async def attempt(self, step: Mapping[str, object]) -> Mapping[str, object]: + if self._passthrough is not None and not step.get("deployment_id"): + return await self._passthrough_attempt(step) + step = await self._resolve_pods(step) # rebind-ok: the picked deployment, resolved to a pod + deployment: Mapping[str, object] = as_str_mapping(step.get("deployment")) or {} # rebind-ok: the picked one + call: AttemptCall | None = None # rebind-ok: built inside the try + # Whether a failure raised now runs the failure callbacks: none raised before the litellm + # call or after it answered does, nor any but the first on a shared logging object. + logs_failure: bool = False # rebind-ok: set around the litellm call + reading_stream: StreamAttempt | None = None # rebind-ok: set while the first chunk is read + refused_call: Final = self._default_call_error(step) + if refused_call is not None: + return {**self._failed(step, self._default_copy(step), None, refused_call), "recorded": True} + try: + deployment, call = self._call(step) # rebind-ok: rebound as the Python router rebinds it + if call.refused is not None: + await log_refused_call(call, call.refused) + raise call.refused + host: Final = self._core._host # pyright: ignore[reportPrivateUsage] # the driver serves its own core + try: + compacted: Final = await compacted_attempt(self._core._port, self._plan, deployment, call) # pyright: ignore[reportPrivateUsage] # see above + except Exception as refused: # noqa: BLE001 # raised before the provider call, so no cooldown reads it + return {**self._failed(step, deployment, call, refused), "recorded": True} + # The deployment's max_parallel_requests slot; an open chat stream takes it along. + port = self._core._port # pyright: ignore[reportPrivateUsage] # the driver serves its own core # rebind-ok: the same port on every attempt + with enter_deployment_slot(port, self._plan, deployment, call.kwargs) as slot: + # A search tool is no deployment: nothing to check; Python's `_arerank` and + # `amoderation` run no routing pre-call checks either. + if self._plan.kind not in _NO_ROUTING_PRE_CALL_CHECKS: + try: + await routing_pre_call_checks( + self._core._checks.routing_callbacks(), # pyright: ignore[reportPrivateUsage] # see above + as_str_mapping(step.get("deployment")) or {}, + ) + except Exception as refused: # noqa: BLE001 # handed to the Rust call loop as the attempt's failure + # Only `_acompletion` hands the check its logging object. + if self._plan.kind == "chat": + _log_refused_pick(self._request, refused) + return await self._refused(step, deployment, call, refused) + first_to_log: Final = not failure_already_logged(call.kwargs, "async_failure") + logs_failure = first_to_log # rebind-ok: whether a failure here is this attempt's to log + pending: Final[object] = ( + self._core._endpoints.create_file_everywhere( # pyright: ignore[reportPrivateUsage] # see above + _with_target_order(call.kwargs, step.get("target_order")), + call.model_name, + self._request.get("specific_deployment"), + ) + if self._plan.kind == "create_file" + else host.invoke(call.host_op)(**compacted) + ) + # The generic helper awaits whatever the handler returns, as the Python router does. + awaited: Final = self._plan.kind == "generic" or inspect.isawaitable(pending) + answered: Final[object] = await pending if awaited else pending # pyright: ignore[reportGeneralTypeIssues, reportUnknownVariableType] # see above + logs_failure = False # rebind-ok: see above + self._core._endpoints.check_response(self._plan, answered, call) # pyright: ignore[reportPrivateUsage, reportUnknownArgumentType] # see above + stream: Final = self._stream_attempt(step, call.kwargs, slot) + logs_failure = first_to_log # rebind-ok: the stream's own failure logging, for its first chunk + reading_stream = stream # rebind-ok: None until the attempt opens a stream + response: Final = await open_attempt_stream(stream, answered) # pyright: ignore[reportUnknownArgumentType] # see above + logs_failure = False # rebind-ok: see above + # `self.success_calls[model_name] += 1`, inside the try: a counter replaced with a + # plain dict fails the attempt. + self._core._endpoints.counters.succeeded(self._plan, call) # pyright: ignore[reportPrivateUsage] # see above + except Exception as error: # noqa: BLE001 # every attempt failure is classified and handed to the Rust call loop + failed: Final = self._failed(step, deployment, call, error) + if reading_stream is None or not logs_failure: + return {**failed, "recorded": not logs_failure} + # `_dispatch_stream_failure`: the stream logs its failure in a task of its own. + reading_stream.report_first_chunk_failure(error) + return {**failed, "recorded": True} + served: Final = self._succeeded(call, response, call.host_op, stream) + # The session re-pin the Python core's success callback runs; a wrapped stream runs it when it ends. + if not served.get("deferred") and self._core.outcome_callbacks(call.kwargs, sync=False).get("success", True): + await self._core.repin_after_success(stream.logged(), response) + session_deployment: Final = clean_realtime_session_deployment(call) + if session_deployment is not None: + await self._core.count_realtime_session_success(session_deployment) + return served + + def attempt_sync(self, step: Mapping[str, object]) -> Mapping[str, object]: + step = self._resolve_pods_sync(step) # rebind-ok: see attempt + deployment: Mapping[str, object] = as_str_mapping(step.get("deployment")) or {} # rebind-ok: see attempt + call: AttemptCall | None = None # rebind-ok: see attempt + logs_failure: bool = False # rebind-ok: see attempt + self._reading_first_chunk = False + refused_call: Final = self._default_call_error(step) + if refused_call is not None: + return {**self._failed(step, self._default_copy(step), None, refused_call), "recorded": True} + try: + deployment, call = self._call(step) # rebind-ok: rebound as the Python router rebinds it + if call.refused is not None: + log_refused_call_sync(call, call.refused) + raise call.refused + try: + self._core._endpoints.sync_pre_call_checks(self._plan, deployment, self._request) # pyright: ignore[reportPrivateUsage] # see above + except Exception as refused: # noqa: BLE001 # handed to the Rust call loop like any attempt failure + # The sync path records no outcome for a refused pre-call check: no cooldown. + return {**self._failed(step, deployment, call, refused), "recorded": True} + first_to_log: Final = not failure_already_logged(call.kwargs, "sync_failure") + logs_failure = first_to_log # rebind-ok: see attempt + answered: Final = self._core._host.invoke(call.host_op)(**call.kwargs) # pyright: ignore[reportPrivateUsage] # a sync endpoint + logs_failure = False # rebind-ok: see attempt + if inspect.isawaitable(answered): + raise _NotSynchronous(f"the host's {call.host_op} endpoint is not synchronous") + self._core._endpoints.check_response(self._plan, answered, call) # pyright: ignore[reportPrivateUsage] # see above + stream: Final = self._stream_attempt(step, call.kwargs) + logs_failure = first_to_log # rebind-ok: see attempt + self._reading_first_chunk = True + response: Final = open_attempt_stream_sync(stream, answered) + self._reading_first_chunk = False + logs_failure = False # rebind-ok: see attempt + self._core._endpoints.counters.succeeded(self._plan, call) # pyright: ignore[reportPrivateUsage] # see attempt + except _NotSynchronous: + raise + except Exception as error: # noqa: BLE001 # see attempt + return {**self._failed(step, deployment, call, error), "recorded": not logs_failure} + return self._succeeded(call, response, "", stream) + + async def select(self, step: Mapping[str, object]) -> object: + request: Final = self._request + messages, input = self._plan.pick_input(request) + if step.get("stage") == "callback_filters": + answer, error = await self._core._answer_callback_filters(step, request, messages) # pyright: ignore[reportPrivateUsage] # the driver serves its own core + if isinstance(error, Exception): + # The try failed with the filter's error: the level raises it if nothing else serves. + _ = self._pick_raised(error, logged=True) + return answer + try: + return await self._core._answer_select(step, request, messages, input) # pyright: ignore[reportPrivateUsage] # the driver serves its own core + except Exception as error: # noqa: BLE001 # handed to the Rust call loop as the try's failure + return self._pick_raised(error, logged=True) + + def select_sync(self, step: Mapping[str, object]) -> object: + request: Final = self._request + try: + return self._core._answer_select_sync(step, request, *self._plan.pick_input(request)) # pyright: ignore[reportPrivateUsage] # see above + except Exception as error: # noqa: BLE001 # see select + return self._pick_raised(error, logged=False) + + async def pre_route(self, step: Mapping[str, object]) -> object: + request: Final = self._request + try: + return await self._core._pre_route(step, request, *self._plan.pick_input(request)) # pyright: ignore[reportPrivateUsage] # see above + except Exception as error: # noqa: BLE001 # see select + return self._pick_raised(error, logged=True) + + def pre_route_sync(self, step: Mapping[str, object]) -> object: + request: Final = self._request + try: + return self._core._custom_pick_sync(step, request, *self._plan.pick_input(request)) # pyright: ignore[reportPrivateUsage] # see above + except Exception as error: # noqa: BLE001 # see select + return self._pick_raised(error, logged=False) + + async def allow_fallback(self, target: str) -> object: + """`_is_fallback_target_authorized`, then `_is_fallback_target_within_budget`. A check that raises + ends the chain outside any hop's ``try``, as a chain's own failure: the level raises what it + would have (``async_function_with_fallbacks_common_utils`` catches it), and its note or an + order path reads the error as the next level's.""" + from litellm.router_utils.fallback_event_handlers import get_pre_routing_selection + from litellm.types.router import FallbackSkipReason + + ctx: Final = RoutingContext(self._core) # pyright: ignore[reportArgumentType] # the checks read registry() and settings() + request: Final = dict(self._request) + try: + if target != get_pre_routing_selection(request) and not await ctx.fallback_access_allowed( + model=target, request_kwargs=request + ): + self._walk.targets().record_skip(target, FallbackSkipReason.NOT_AUTHORIZED) + return False + if not await ctx.fallback_budget_allowed(model=target, request_kwargs=request): + self._walk.targets().record_skip(target, FallbackSkipReason.OVER_BUDGET) + return False + except Exception as error: # noqa: BLE001 # handed to the Rust call loop, which ends the chain + self._level_errors[self._level + 1] = error + return {**attempt_failure(error), "raised": True} + return True + + def refuse_fallbacks(self, step: Mapping[str, object]) -> bool: + """``run_async_fallback``'s check of a chain's dict targets: a refused chain ends the walk + with the check's error, which the first hop's level would have raised.""" + refused: Final = rejected_fallback_target(_as_list(step.get("targets"))) + if refused is None: + return False + level: Final = step.get("level") + self._level_errors[(level if isinstance(level, int) else 0) + 1] = refused + return True + + def _head_response(self) -> object | None: + """``set_response_headers`` up to its usage read: the deployment id, the model group the + response was served as and the routers' decisions.""" + response: Final = prepare_response_for_header_attachment(self._last_response) + if response is None: + return None + name, kwargs = self._served + additional: Final = ensure_response_additional_headers(response) + apply_response_model_id(response, find_deployment_metadata(kwargs)) + additional["x-litellm-model-group"] = name + apply_quality_router_decision_headers(additional, kwargs) + additional.update(complexity_router_decision_headers(kwargs)) + return response + + def _count_response(self, response: object, step: Mapping[str, object]) -> object: + """What the wrappers above `make_call` stamp: ``async_function_with_retries``'s retry count, + then the fallback count of ``async_function_with_fallbacks`` and ``run_async_fallback``.""" + retries: Final = step.get("attempted_retries") + if isinstance(retries, int): + max_retries: Final = step.get("max_retries") + response = add_retry_headers_to_response( # rebind-ok: gains the counts + response=response, + attempted_retries=retries, + max_retries=max_retries if isinstance(max_retries, int) and retries > 0 else None, + ) + fallbacks: Final = step.get("attempted_fallbacks") + attempted: Final = fallbacks if isinstance(fallbacks, int) else 0 + return add_fallback_headers_to_response( + response=self.outcomes.add_headers(response, attempted), attempted_fallbacks=attempted + ) + + def log_outcomes_sync(self) -> None: + if pending := self.outcomes.take_pending(): + from litellm.litellm_core_utils import asyncify + + asyncify.run_async_function(await_all, pending) # pyright: ignore[reportUnknownMemberType] # the untyped sync bridge + + def finish(self, step: Mapping[str, object]) -> object: + """The response with the router's headers, in the order ``make_call`` and the wrappers above it + stamp them: a usage read that raises leaves the counts off.""" + if not self._plan.sets_headers: + self.log_outcomes_sync() + return self._last_response + response: Final = self._head_response() + if response is None: + self.log_outcomes_sync() + return self._last_response + name: Final = self._served[0] + if name is not None: + core: Final = self._core + remaining: Final = as_str_mapping( + core.run_routing_cache_sync(lambda: core.health().usage(name, kind="remaining")) + ) + if remaining is not None: + apply_remaining_usage_headers(ensure_response_additional_headers(response), dict(remaining)) + counted: Final = self._count_response(response, step) + self.log_outcomes_sync() + return counted + + async def finish_async(self, step: Mapping[str, object]) -> object: + """``finish``, with the usage read on the call's own loop.""" + if not self._plan.sets_headers: + await await_all(self.outcomes.take_pending()) + return self._last_response + response: Final = self._head_response() + if response is None: + await await_all(self.outcomes.take_pending()) + return self._last_response + name: Final = self._served[0] + if name is not None: + remaining: Final = await self._core.health().usage(name, kind="remaining") + if isinstance(remaining, Mapping): + apply_remaining_usage_headers(ensure_response_additional_headers(response), dict(remaining)) + counted: Final = self._count_response(response, step) + await await_all(self.outcomes.take_pending()) + return counted + + def failed(self, step: Mapping[str, object]) -> BaseException: + model: Final = str(self._request.get("model")) + primary: Final = self._level_errors.get(self._base_level) + if primary is None: + return self._core._failure_error(step, model, self._last_error, request=self._request) # pyright: ignore[reportPrivateUsage] # see above + if ( + self._base_level == 0 + and self._explains_failure + and self._core._host.dynamic().expose_router_debug_in_errors # pyright: ignore[reportPrivateUsage] # see above + and hasattr(primary, "message") + ): + primary.message += self._explanation(primary, model) # pyright: ignore[reportAttributeAccessIssue, reportUnknownMemberType] # checked above + return primary + + def _hint_level_failure(self, error: BaseException | None, model: str) -> None: + """A level that reaches its fallback handling on a context-window or content-policy error + with no list of its own appends the hint to set one, before any hop runs, as every level of + the Python walk does: a stream's failure event, still queued then, reads the hint.""" + from litellm.router_utils.fallback_event_handlers import check_non_standard_fallback_format + + if ( + error is None + or not hasattr(error, "message") + or not self._core._host.dynamic().expose_router_debug_in_errors # pyright: ignore[reportPrivateUsage] # the driver serves its own core + ): + return + walked: Final = self._core._walked_fallbacks # pyright: ignore[reportPrivateUsage] # see above + fallbacks: Final = self._request.get("fallbacks", walked["fallbacks"]) + if check_non_standard_fallback_format(fallbacks=_as_list(fallbacks)): + return + hint: Final = self._kind_hint(error, model, fallbacks) + if hint: + error.message += hint # pyright: ignore[reportAttributeAccessIssue, reportUnknownMemberType] # checked above # rebind-ok: the router appends to the error it raises + + def _kind_hint(self, error: BaseException, model: str, fallbacks: object) -> str: + """The hint a context-window or content-policy error gets when its own list is unset.""" + import litellm.exceptions as litellm_exceptions + + walked: Final = self._core._walked_fallbacks # pyright: ignore[reportPrivateUsage] # the driver serves its own core + if isinstance(error, litellm_exceptions.ContextWindowExceededError): + if self._request.get("context_window_fallbacks", walked["context_window_fallbacks"]) is not None: + return "" + return ( + f"\nmodel={model}. context_window_fallbacks=None. fallbacks={fallbacks}.\n\n" + "Set 'context_window_fallback' - https://docs.litellm.ai/docs/routing#fallbacks" + ) + if isinstance(error, litellm_exceptions.ContentPolicyViolationError): + if self._request.get("content_policy_fallbacks", walked["content_policy_fallbacks"]) is not None: + return "" + return "\n" + format_content_policy_fallback_hint(model, None, fallbacks) # pyright: ignore[reportArgumentType] # the list as configured + return "" + + def _explanation(self, primary: BaseException, model: str) -> str: + """What ``async_function_with_fallbacks_common_utils`` appends to the error it re-raises + under ``expose_router_debug_in_errors``, from the resolved settings and what the call did.""" + import litellm.exceptions as litellm_exceptions + from litellm.router_utils.fallback_event_handlers import ( + check_non_standard_fallback_format, + fallback_lookup_groups, + get_fallback_model_group_for_lookup_groups, + ) + + walked: Final = self._core._walked_fallbacks # pyright: ignore[reportPrivateUsage] # the driver serves its own core + request: Final = self._request + # The request's own lists replace the router's (`kwargs.get("fallbacks", router)`). + fallbacks: Final = request.get("fallbacks", walked["fallbacks"]) + window_list: Final = request.get("context_window_fallbacks", walked["context_window_fallbacks"]) + policy_list: Final = request.get("content_policy_fallbacks", walked["content_policy_fallbacks"]) + lookup: Final = fallback_lookup_groups(request, model) + + def group_of(source: object) -> tuple[Sequence[object] | None, int | None]: + return get_fallback_model_group_for_lookup_groups(_as_list(source), lookup) # pyright: ignore[reportArgumentType] # the fallbacks config shape + + def outcome(group: Sequence[object] | None) -> str: + return format_fallback_outcome_message( + model, group, self._fallback_detail(primary), skipped=self._walk.targets().skipped + ) + + # `context_window_fallback_model_group`: only a chain the group matched takes precedence. + window_group: Final = ( + group_of(window_list)[0] + if isinstance(primary, litellm_exceptions.ContextWindowExceededError) and window_list is not None + else None + ) + if window_group is None and check_non_standard_fallback_format(fallbacks=_as_list(fallbacks)): + return outcome(_as_list(fallbacks)) + if window_group is not None: + return outcome(window_group) + # A content-policy list decides on its own: one with no match raises into the `except`. + if isinstance(primary, litellm_exceptions.ContentPolicyViolationError) and policy_list is not None: + return outcome(group_of(policy_list)[0]) + if fallbacks is None or not lookup: + return outcome(None) + found, generic = group_of(fallbacks) + group: Final[Sequence[object] | None] = ( + found if found is not None or generic is None else _as_list(_as_list(fallbacks)[generic]["*"]) # pyright: ignore[reportIndexIssue, reportCallIssue, reportArgumentType, reportUnknownArgumentType] # the entry the lookup matched + ) + return format_no_fallback_group_message(lookup, _as_list(fallbacks)) if group is None else outcome(group) # pyright: ignore[reportArgumentType] # the fallbacks config shape + + def _fallback_detail(self, primary: BaseException) -> str: + """``run_async_fallback`` raises its last hop's error, or the primary's when no hop ran.""" + return truncate_fallback_error_detail(redact_string(str(self._level_errors.get(1, primary)))) + + +__all__ = [ + "Resolved", + "RustCore", + "RustCoreUnsupported", + "RustHealth", + "RustRegistry", + "RustSettings", + "attempt_failure", + "failure_kind", + "total_tokens", + "unsupported_arguments", +] diff --git a/litellm/router_rust/invoke/README.md b/litellm/router_rust/invoke/README.md new file mode 100644 index 00000000000..d8562f99359 --- /dev/null +++ b/litellm/router_rust/invoke/README.md @@ -0,0 +1,13 @@ +# invoke + +The host side of `Host::invoke`: running one attempt against the one deployment router-core picked. Nothing here selects, retries, falls back or records cooldowns. + +- `shaping`: the kwargs an attempt sends to its deployment +- `endpoints`: what each public router method does before, around and after the routed call +- `streaming`: stream wrappers, so a stream that fails before its first token fails the attempt +- `slots`: a deployment's `max_parallel_requests` +- `walk`: the retry and fallback metadata stamped on each attempt +- `prompt_management`: prompt-management groups +- `websocket`: admission for each turn of a Responses websocket + +A Rust host replaces this with native endpoint calls. diff --git a/litellm/router_rust/invoke/__init__.py b/litellm/router_rust/invoke/__init__.py new file mode 100644 index 00000000000..faa2c13198d --- /dev/null +++ b/litellm/router_rust/invoke/__init__.py @@ -0,0 +1 @@ +"""The host side of ``Host::invoke``: one router attempt against one deployment.""" diff --git a/litellm/router_rust/invoke/endpoints.py b/litellm/router_rust/invoke/endpoints.py new file mode 100644 index 00000000000..c1dc5162ccf --- /dev/null +++ b/litellm/router_rust/invoke/endpoints.py @@ -0,0 +1,2454 @@ +""" +The Rust router core's endpoints: what each public router method does before, around and after +the routed call, as the Python router's endpoint methods do it. + +An endpoint is either routed (the Rust core picks the deployment, retries and falls back, and the +host runs each attempt) or served directly (the Python router calls litellm's function without +routing: no model group, an id that names no deployment, a fan-out over every deployment). Every +call out goes through the host: ``host.invoke(op)``, where an op is a litellm function name or a +dotted path to one (``vector_stores.main.asearch``), looked up when it is called. +""" + +import asyncio +import copy +import dataclasses +import threading +import time +import uuid +from collections import defaultdict +from collections.abc import Awaitable, Callable, Iterable, Mapping, Sequence +from types import MappingProxyType +from typing import TYPE_CHECKING, Final, Literal, Protocol, TypeAlias + +from pydantic import TypeAdapter, ValidationError + +from litellm._internal_context import service_target +from litellm._logging import verbose_router_logger +from litellm.litellm_core_utils.add_retry_fallback_headers import ( + add_fallback_headers_to_response, + add_retry_headers_to_response, +) +from litellm.litellm_core_utils.batch_utils import get_router_metadata_variable_name +from litellm.litellm_core_utils.core_helpers import without +from litellm.litellm_core_utils.routing_common_utils import without_router_only_kwargs +from litellm.responses.retrieve_stream_query import without_retrieve_stream_defaults +from litellm.router_api import RoutingContext +from litellm.router_rust.invoke.shaping import ( + ROUTER_CLIENTS_TARGET, + STREAM_TIMEOUT_AWARE_GENERIC_HANDLERS, + Shaping, + as_dict, + deployment_params_with_request_reasoning_override, + drop_unsupported_classifier_reasoning_effort, + metadata_bucket, + stamp_deployment_num_retries, + stamp_failed_deployment_id, + stamp_retry_skip_deployment_id, + update_kwargs_with_deployment, +) +from litellm.router_strategy.complexity_router.context_compaction import ( + Surface, + compact_to_fit, + initialize_compaction_state, + surface_for_call, +) +from litellm.router_utils.endpoint_eligibility import endpoint_for_call, endpoint_selection + +if TYPE_CHECKING: + from litellm.caching.dual_cache import DualCache + from litellm.router_api import RouterRegistry + from litellm.router_rust.core import RustCore + from litellm.router_rust.steps.selection import StrategySelectors + from litellm.rust_bridge._native import RouterCoreHandle + from litellm.types.router import Deployment + + +class EndpointCore(Protocol): + """What the endpoints use of the core they serve: its host, Rust handle, settings, routing + cache, registry lookups and call drivers. ``RustCore`` provides it through ``_EndpointPort``; + a Rust host's endpoint layer would have to provide the same list.""" + + @property + def router(self) -> "RustCore": + """The core itself, for the provider handlers and helpers that take a router.""" + ... + + @property + def host(self) -> "PythonHost": ... + + @property + def handle(self) -> "RouterCoreHandle": ... + + @property + def resolved(self) -> "Resolved": ... + + @property + def shaping(self) -> Shaping: ... + + @property + def routing_cache(self) -> "DualCache": ... + + @property + def selectors(self) -> "StrategySelectors": ... + + @property + def endpoints(self) -> "RustEndpoints": ... + + async def drive( + self, + plan: "Plan", + request: Mapping[str, object], + overrides: Mapping[str, object] | None = None, + *, + passthrough: Callable[[Mapping[str, object], str], Awaitable[object]] | None = None, + prepared_already: bool = False, + ) -> object: ... + + def call_sync( + self, plan: "Plan", request: Mapping[str, object], overrides: Mapping[str, object] | None = None + ) -> object: ... + + def routing_callbacks(self) -> Iterable[object]: ... + + def resolve_alias(self, model: str) -> str | None: ... + + def group_deployments(self, model_name: str) -> Sequence[Mapping[str, object]] | None: + """``_core_get_model_list(model_name=...)`` on the host registry.""" + ... + + def routed_name(self, model: object) -> str: + """``RustCore.routed_name``: what a ``model`` that is no string raises, or the name it routes under.""" + ... + + def service_logged( + self, call_type: str, start_time: float, end_time: float, request: Mapping[str, object] + ) -> None: ... + + async def set_response_headers( + self, response: object, model_group: str | None, request_kwargs: Mapping[str, object] + ) -> object: ... + + def configured_deployment(self, model_id: str) -> "Deployment | None": + """``_core_get_deployment``: the deployment as the host registry holds it now.""" + ... + + def registry(self) -> "RouterRegistry": ... + + async def select( + self, + model: str, + request_kwargs: dict[str, object] | None = None, # mutable-ok: the RouterCore signature + *, + kind: str = "pick", + specific_deployment: bool | None = None, + messages: object = None, + ) -> Mapping[str, object] | Sequence[Mapping[str, object]]: ... + + async def acompletion(self, *args: object, **kwargs: object) -> object: ... # kwargs-ok: endpoint payload + + +Kind = Literal[ # rebind-ok: rebound as the Python router rebinds it + "chat", + "text", + "embedding", + "image", + "transcription", + "speech", + "rerank", + "adapter", + "generic", + "generic_sync", + "batch", + "cancel_batch", + "vector_store_create", + "moderation", + "create_file", + "search", +] +Stamp = Literal["deployment", "retry_skip", "none"] # rebind-ok: rebound as the Python router rebinds it +CountKey = Literal["deployment", "group", "empty"] # rebind-ok: rebound as the Python router rebinds it + +# The text of the user message each attempt helper hands its pick in place of the request's own +# (`_aimage_generation`, `_aadapter_completion`, `_acreate_batch`, ...). +_PICK_TEXTS: Final[Mapping[str, str]] = MappingProxyType( + { + "image": "prompt", + "transcription": "prompt", + "speech": "prompt", + "adapter": "default text", + "batch": "files-api-fake-text", + "create_file": "files-api-fake-text", + "cancel_batch": "batch-api-fake-text", + "vector_store_create": "vector-store-api-fake-text", + } +) + + +def _user_message(content: object) -> list[dict[str, object]]: # mutable-ok: a fresh list, as each pick builds it + return [{"role": "user", "content": content}] + + +# Request keys the router reads to route the call; the deployment call never receives them. +ROUTING_KEYS: Final = frozenset( + { + "model", + "num_retries", + "disable_fallbacks", + "specific_deployment", + "_router_weights", + "_retry_skipped_deployment_ids", + "_excluded_deployment_ids", + "_target_order", + "original_function", + # The walk's own settings, which `async_function_with_fallbacks` pops before the attempt. + "model_group_retry_policy", + "fallbacks", + "context_window_fallbacks", + "content_policy_fallbacks", + "mock_testing_fallbacks", + "mock_testing_context_fallbacks", + "mock_testing_content_policy_fallbacks", + "mock_testing_rate_limit_error", + } +) + + +@dataclasses.dataclass(frozen=True, slots=True) +class Plan: + """How one endpoint runs its routed attempts.""" + + endpoint: str + host_op: str + kind: Kind + sync: bool = False + # The bucket `_update_kwargs_before_fallbacks` writes the requested group into; None for the + # endpoints that write `metadata.model_group` alone (no trace id, no alias). + prepare_field: str | None = "metadata" + function_name: str | None = None + splits_stream_timeout: bool = False + stamp: Stamp = "retry_skip" + total_key: CountKey | None = "deployment" + success_key: CountKey | None = "deployment" + fail_key: CountKey | None = "deployment" + # One attempt, no retries and no fallbacks. + single_attempt: bool = False + # The named parameter the endpoint takes beside **kwargs, sent between the deployment params and + # the request (`{**data, name: value, **kwargs}`). + named: tuple[str, ...] = () + + @property + def on_its_own_loop(self) -> bool: + """``function_with_fallbacks`` runs these sync endpoints' whole walk on a private event loop.""" + return self.sync and self.kind in ("chat", "embedding", "image") + + def pick_input(self, request: Mapping[str, object]) -> tuple[object, object]: + """The ``messages`` and ``input`` the endpoint's attempt helper hands its pick, which every + selection stage reads: the pre-call checks run when either is set.""" + match self.kind: + case "chat": + return request.get("messages"), None + case "text": + return _user_message(request.get("prompt")), None + case "embedding": + return None, request.get("input") + case "rerank" | "moderation" | "search": + return None, None + case "generic" | "generic_sync": + return request.get("messages"), request.get("input") + case _: + return _user_message(_PICK_TEXTS[self.kind]), None + + @property + def sets_headers(self) -> bool: + """Whether the response gets the router's headers (the sync generic wrapper adds none).""" + return self.kind != "generic_sync" + + @property + def stamps_failed_deployment(self) -> bool: + """Whether the attempt stamps ``failed_deployment_id`` on its error, which weighted failover + reads: every endpoint but file and batch creation (whose result belongs to the deployment + that served it), moderation's pass-through and search.""" + return self.kind not in ("batch", "cancel_batch", "create_file", "moderation", "search") + + @property + def metadata_field(self) -> str: + """Where `_update_kwargs_with_deployment` writes the deployment it routed to.""" + return get_router_metadata_variable_name(function_name=self.function_name) + + +def _generic(endpoint: str, host_op: str | None = None, *, sync: bool = False) -> Plan: + if sync: + return Plan( + endpoint=endpoint, + host_op=host_op or endpoint, + kind="generic_sync", + sync=True, + prepare_field="litellm_metadata", + function_name="generic_api_call", + splits_stream_timeout=_handler_name(host_op or endpoint) in STREAM_TIMEOUT_AWARE_GENERIC_HANDLERS, + stamp="none", + fail_key="group", + single_attempt=True, + ) + return Plan( + endpoint=endpoint, + host_op=host_op or endpoint, + kind="generic", + prepare_field="litellm_metadata", + function_name="_ageneric_api_call_with_fallbacks", + splits_stream_timeout=_handler_name(host_op or endpoint) in STREAM_TIMEOUT_AWARE_GENERIC_HANDLERS, + stamp="deployment", + fail_key="group", + ) + + +def _handler_name(host_op: str) -> str: + return host_op.rsplit(".", 1)[-1] + + +def invoked_handler_name(core: EndpointCore, plan: Plan) -> str | None: + """The generic handler's ``__name__`` (`original_generic_function.__name__`): the function the + host serves for the op, which decides the stream-timeout split and the responses socket.""" + function: Final[object] = core.host.invoke(plan.host_op) # pyright: ignore[reportUnknownMemberType] # the core's host + name: Final[object] = getattr(function, "__name__", None) + return name if isinstance(name, str) else None + + +# The endpoints whose attempt helper runs inside `make_call`'s `endpoint_selection`, under the +# name of its `original_function`. +_SELECTION_CALLS: Final[Mapping[str, str]] = MappingProxyType( + { + "acompletion": "_acompletion", + "completion": "_completion", + "atext_completion": "_atext_completion", + } +) + + +def selection_endpoint(core: EndpointCore, plan: Plan) -> str | None: + """The endpoint ``endpoint_selection`` names around the plan's pick. A sync generic call never + reaches ``make_call``, and a generic attempt reads the handler's name and module.""" + call: Final = _SELECTION_CALLS.get(plan.endpoint) + if call is not None: + return endpoint_for_call(call) + if plan.kind != "generic": + return None + function: Final[object] = core.host.invoke(plan.host_op) # pyright: ignore[reportUnknownMemberType] # the core's host + return endpoint_for_call(getattr(function, "__name__", None), getattr(function, "__module__", None)) + + +PLANS: Final[Mapping[str, Plan]] = MappingProxyType( + { + "acompletion": Plan( + "acompletion", "acompletion", "chat", splits_stream_timeout=True, stamp="deployment", named=("messages",) + ), + "completion": Plan( + "completion", + "completion", + "chat", + sync=True, + splits_stream_timeout=True, + stamp="deployment", + total_key=None, + success_key=None, + fail_key=None, + named=("messages",), + ), + "atext_completion": Plan( + "atext_completion", + "atext_completion", + "text", + splits_stream_timeout=True, + fail_key="group", + named=("prompt",), + ), + # `_asearch_with_fallbacks`: the search tool name routes as the model, with no deployment + # picked; its helper picks among the router's search tools. + "asearch": Plan( + "asearch", + "search.asearch", + "search", + prepare_field="litellm_metadata", + stamp="none", + total_key=None, + success_key=None, + fail_key=None, + ), + "aembedding": Plan("aembedding", "aembedding", "embedding", named=("input",)), + "embedding": Plan("embedding", "embedding", "embedding", sync=True, named=("input",)), + "aimage_generation": Plan("aimage_generation", "aimage_generation", "image", named=("prompt",)), + "image_generation": Plan( + "image_generation", + "image_generation", + "image", + sync=True, + prepare_field=None, + total_key="empty", + success_key="empty", + fail_key="empty", + named=("prompt",), + ), + "atranscription": Plan( + "atranscription", + "atranscription", + "transcription", + total_key="group", + success_key="group", + fail_key="group", + named=("file",), + ), + "aspeech": Plan( + "aspeech", + "aspeech", + "speech", + total_key="group", + success_key="group", + fail_key="group", + named=("input",), + ), + "arerank": Plan("arerank", "arerank", "rerank"), + "aadapter_completion": Plan( + "aadapter_completion", + "aadapter_completion", + "adapter", + prepare_field=None, + splits_stream_timeout=True, + fail_key="group", + ), + "acreate_batch": Plan( + "acreate_batch", + "acreate_batch", + "batch", + prepare_field="litellm_metadata", + function_name="_acreate_batch", + fail_key="group", + ), + "acancel_batch": Plan( + "acancel_batch", + "acancel_batch", + "cancel_batch", + prepare_field="litellm_metadata", + function_name="_acancel_batch", + fail_key="group", + ), + "acreate_file": Plan( + "acreate_file", + "acreate_file", + "create_file", + function_name="acreate_file", + stamp="none", + total_key=None, + success_key=None, + fail_key="group", + ), + "avector_store_create": Plan( + "avector_store_create", + "vector_stores.acreate", + "vector_store_create", + prepare_field=None, + function_name="avector_store_create", + stamp="none", + fail_key="group", + single_attempt=True, + ), + "amoderation": Plan( + "amoderation", + "amoderation", + "moderation", + stamp="none", + total_key=None, + success_key=None, + fail_key=None, + single_attempt=True, + ), + # The Python router's anthropic messages host op is `anthropic_messages`. + "aanthropic_messages": _generic("aanthropic_messages", "anthropic_messages"), + "anthropic_messages": _generic("anthropic_messages", "anthropic_messages"), + "aresponses": _generic("aresponses"), + "responses": _generic("responses", sync=True), + } +) + +# --- the generic endpoints ------------------------------------------------------------------ + +# endpoint -> host op, for the async generic endpoints routed through +# `_ageneric_api_call_with_fallbacks` when they carry a model. +_HOST_GENERIC: Final[Mapping[str, str]] = MappingProxyType( + { + "afile_delete": "afile_delete", + "afile_content": "afile_content", + "afile_list": "afile_list", + "aget_responses": "aget_responses", + "acancel_responses": "acancel_responses", + "acompact_responses": "acompact_responses", + "adelete_responses": "adelete_responses", + "alist_input_items": "alist_input_items", + "_arealtime": "_arealtime", + "acreate_realtime_client_secret": "acreate_realtime_client_secret", + "arealtime_calls": "arealtime_calls", + "acreate_realtime_transcription_session": "acreate_realtime_transcription_session", + "acreate_realtime_translation_client_secret": "acreate_realtime_translation_client_secret", + "arealtime_translation_calls": "arealtime_translation_calls", + "_aresponses_websocket": "_aresponses_websocket", + "acreate_fine_tuning_job": "acreate_fine_tuning_job", + "acancel_fine_tuning_job": "acancel_fine_tuning_job", + "alist_fine_tuning_jobs": "alist_fine_tuning_jobs", + "aretrieve_fine_tuning_job": "aretrieve_fine_tuning_job", + "aimage_edit": "aimage_edit", + "allm_passthrough_route": "allm_passthrough_route", + "acreate_skill": "acreate_skill", + "alist_skills": "alist_skills", + "aget_skill": "aget_skill", + "adelete_skill": "adelete_skill", + "agenerate_content": "google_genai.agenerate_content", + "agenerate_content_stream": "google_genai.agenerate_content_stream", + "aocr": "ocr.aocr", + "adecisions": "decisions.adecisions", + "asystemone": "systemone.asystemone", + "aextraction": "extraction.aextraction", + **{f"avideo_{name}": f"videos.avideo_{name}" for name in ("generation", "list", "status", "content", "remix")}, + **{ + name: f"containers.{name}" + for name in ("acreate_container", "alist_containers", "aretrieve_container", "adelete_container") + }, + **{ + name: f"containers.endpoint_factory._generated_endpoints.{name}" + for name in ( + "aupload_container_file", + "alist_container_files", + "aretrieve_container_file", + "adelete_container_file", + "aretrieve_container_file_content", + ) + }, + **{ + f"avector_store_{name}": f"vector_stores.main.a{name}" + for name in ("search", "retrieve", "list", "update", "delete") + }, + **{ + f"avector_store_file_{name}": f"vector_store_files.main.a{target}" + for name, target in ( + ("create", "create"), + ("list", "list"), + ("retrieve", "retrieve"), + ("content", "retrieve_content"), + ("update", "update"), + ("delete", "delete"), + ) + }, + "acreate_interaction": "interactions.acreate", + "aget_interaction": "interactions.aget", + "adelete_interaction": "interactions.adelete", + "acancel_interaction": "interactions.acancel", + } +) + +_SYNC_GENERIC: Final[Mapping[str, str]] = MappingProxyType( + { + "generate_content": "google_genai.generate_content", + "generate_content_stream": "google_genai.generate_content_stream", + "ocr": "ocr.ocr", + # The sync wrapper routes a search tool name as a model group, unlike `asearch`. + "search": "search.search", + "decisions": "decisions.decisions", + "systemone": "systemone.systemone", + "extraction": "extraction.extraction", + **{f"video_{name}": f"videos.video_{name}" for name in ("generation", "list", "status", "content", "remix")}, + **{ + name: f"containers.{name}" + for name in ("create_container", "list_containers", "retrieve_container", "delete_container") + }, + **{ + f"vector_store_{name}": f"vector_stores.main.{name}" + for name in ("search", "create", "retrieve", "list", "update", "delete") + }, + } +) + +_CONTAINER_FILE_ROUTED: Final = frozenset( + { + "aupload_container_file", + "alist_container_files", + "aretrieve_container_file", + "adelete_container_file", + "aretrieve_container_file_content", + } +) +_VIDEO_DIRECT: Final = frozenset({"avideo_create_character", "avideo_get_character", "avideo_edit", "avideo_extension"}) +_SYNC_VIDEO_UNBRANCHED: Final = frozenset( + {"video_create_character", "video_get_character", "video_edit", "video_extension"} +) +_SYNC_CONTAINER_FILES_UNBRANCHED: Final = frozenset( + { + "upload_container_file", + "list_container_files", + "retrieve_container_file", + "delete_container_file", + "retrieve_container_file_content", + } +) +_UNBRANCHED: Final = frozenset( + { + "aadapter_generate_content", + *_SYNC_VIDEO_UNBRANCHED, + *_SYNC_CONTAINER_FILES_UNBRANCHED, + "get_interaction", + "delete_interaction", + "cancel_interaction", + } +) +_ASSISTANTS: Final = frozenset( + { + "acreate_assistants", + "adelete_assistant", + "aget_assistants", + "acreate_thread", + "aget_thread", + "a_add_message", + "aget_messages", + "arun_thread", + } +) +_MANAGED_AGENTS: Final[Mapping[str, str]] = MappingProxyType( + { # comprehension-ok: a flat walk of nested rows + f"{prefix}{name}": f"interactions.agents.{prefix}{target}" + for prefix in ("a", "") + for name, target in ( + ("create_agent", "create"), + ("list_agents", "list"), + ("get_agent", "get"), + ("delete_agent", "delete"), + ("list_agent_versions", "list_versions"), + ) + } +) +_SYNC_VECTOR_STORE_FILES: Final[Mapping[str, str]] = MappingProxyType( + { + f"vector_store_file_{name}": f"vector_store_files.main.{target}" + for name, target in ( + ("create", "create"), + ("list", "list"), + ("retrieve", "retrieve"), + ("content", "retrieve_content"), + ("update", "update"), + ("delete", "delete"), + ) + } +) + + +def plan_for(endpoint: str) -> Plan: + plan: Final = generic_plan(endpoint) + if plan is None: + raise KeyError(endpoint) + return plan + + +def generic_plan(endpoint: str) -> Plan | None: + if endpoint in PLANS: + return PLANS[endpoint] + if endpoint in _HOST_GENERIC: + return _generic(endpoint, _HOST_GENERIC[endpoint]) + if endpoint in _VIDEO_DIRECT: + return _generic(endpoint, f"videos.{endpoint}") + if endpoint in _SYNC_GENERIC: + return _generic(endpoint, _SYNC_GENERIC[endpoint], sync=True) + return None + + +# --- preparing the request ------------------------------------------------------------------- + + +def prepare_request( + plan: Plan, request: Mapping[str, object], model: str, aliases: Mapping[str, object] +) -> dict[str, object]: # mutable-ok: the call's kwargs + """What the router adds once before the first attempt (`_update_kwargs_before_fallbacks`): a + trace id and the requested group, in the caller's own metadata dict.""" + prepared: Final = dict(request) + if plan.kind == "vector_store_create": + return prepared # routed straight to a deployment, with no group recorded + if plan.prepare_field is None: + metadata_bucket(prepared, "metadata").update({"model_group": model}) + return prepared + prepared.setdefault("litellm_trace_id", str(uuid.uuid4())) + metadata_bucket(prepared, plan.prepare_field).update( + {"model_group": model, "model_group_alias": model if model in aliases else None} + ) + return prepared + + +# --- one attempt ---------------------------------------------------------------------------- + + +@dataclasses.dataclass(slots=True) +class AttemptCall: + """One attempt's host call: the op, its kwargs, and the request kwargs the router kept (the + ones its response headers are read from).""" + + host_op: str + kwargs: Mapping[str, object] + served: Mapping[str, object] + model_name: str + group: str + # `_embedding`'s batch-size refusal, raised by the driver in place of the host call. + refused: Exception | None = None + + +def reports_own_success(call: AttemptCall) -> bool: + """A realtime attempt's return says nothing about the session (a dropped one returns too), so + the host counts its success itself (`clean_realtime_session_deployment`).""" + return call.host_op == "_arealtime" + + +def clean_realtime_session_deployment(call: AttemptCall) -> str | None: + """`_ageneric_api_call_with_fallbacks`: the deployment a finished `_arealtime` session counts a + success against, when its success handlers ran and it did not end abnormally.""" + from litellm.constants import REALTIME_SESSION_ENDED_ABNORMALLY_KEY, REALTIME_SESSION_SUCCESS_LOGGED_KEY + from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging + + session_logging: Final = call.kwargs.get("litellm_logging_obj") + if not reports_own_success(call) or not isinstance(session_logging, LiteLLMLogging): + return None + details: Final = session_logging.model_call_details + if details.get(REALTIME_SESSION_SUCCESS_LOGGED_KEY) is not True: + return None + if details.get(REALTIME_SESSION_ENDED_ABNORMALLY_KEY) is True: + return None + deployment_id: Final = as_dict(call.kwargs.get("model_info")).get("id") + return None if deployment_id is None else str(deployment_id) + + +# `_scheduler` once `del router.scheduler` ran (None is a scheduler not built yet). +_NO_SCHEDULER: Final = object() + + +class Counters: + """`total_calls`, `success_calls` and `fail_calls`, keyed as each endpoint keys them.""" + + __slots__ = ("fail_calls", "success_calls", "total_calls") + + def __init__(self) -> None: + self.total_calls: defaultdict[str, int] = defaultdict(int) # mutable-ok: the router's live counters + self.success_calls: defaultdict[str, int] = defaultdict(int) # mutable-ok: same + self.fail_calls: defaultdict[str, int] = defaultdict(int) # mutable-ok: same + + @staticmethod + def _key(which: CountKey, model_name: str, group: str) -> str: + return model_name if which == "deployment" else group if which == "group" else "" + + def started(self, plan: Plan, call: AttemptCall) -> None: + if plan.total_key is not None: + self.total_calls[self._key(plan.total_key, call.model_name, call.group)] += 1 + + def succeeded(self, plan: Plan, call: AttemptCall) -> None: + if plan.success_key is not None: + self.success_calls[self._key(plan.success_key, call.model_name, call.group)] += 1 + + def retract_success(self, plan: Plan, call: AttemptCall) -> None: + """A served response the router then raised over: the attempt helper never reached its count.""" + if plan.success_key is not None: + self.success_calls[self._key(plan.success_key, call.model_name, call.group)] -= 1 + + def failed(self, plan: Plan, model_name: str | None, group: str) -> None: + if plan.fail_key is None: + return + if plan.fail_key == "deployment" and model_name is None: + return + self.fail_calls[self._key(plan.fail_key, model_name or "", group)] += 1 + + +def _compaction_surface(core: EndpointCore, plan: Plan) -> Surface | None: + """``surface_for_call`` over the function the Python router's fallback loop is handed.""" + if plan.kind == "generic" and surface_for_call(_handler_name(plan.host_op)) is not None: + return surface_for_call(invoked_handler_name(core, plan) or "") + return "chat" if plan.kind == "chat" and not plan.sync else None + + +def arm_compaction_state( + core: EndpointCore, + plan: Plan, + request: dict[str, object], # mutable-ok: the request kwargs, written in place as the Python router writes them +) -> None: + """``initialize_compaction_state`` at the top of ``async_function_with_fallbacks``.""" + surface: Final = _compaction_surface(core, plan) + if surface is not None: + request["_context_compaction_state"] = initialize_compaction_state(request, surface) # rebind-ok: its own + + +async def compacted_attempt( + core: EndpointCore, plan: Plan, deployment: Mapping[str, object], call: "AttemptCall" +) -> Mapping[str, object]: + """The attempt's kwargs after ``compact_to_fit``, which ``_acompletion`` and the generic helper + run on every attempt: it is also what refuses a compactor's own oversized child call.""" + if plan.kind != "generic" and (plan.kind != "chat" or plan.sync): + return call.kwargs + ctx: Final = RoutingContext(core.router) # pyright: ignore[reportArgumentType] # the checks read registry(), settings() and the token count + return await compact_to_fit(ctx, deployment, call.kwargs, _compaction_surface(core, plan)) + + +def attempt_call( + core: EndpointCore, + plan: Plan, + request: Mapping[str, object], + deployment: Mapping[str, object], + *, + group: str, + requested: str, + fallback_level: int = 0, +) -> AttemptCall: + """The host call for one attempt of ``plan`` against ``deployment``: the request (as the attempt + helper receives it), shaped for the deployment and merged with its litellm params.""" + shaping: Final = core.shaping + kwargs: Final = without(request, *ROUTING_KEYS) + named: Final = {name: kwargs.pop(name) for name in plan.named if name in kwargs} + data: Final = as_dict(deployment.get("litellm_params")) + model_name: Final = str(data.get("model")) + cache: Final = core.routing_cache + caching: Final = shaping.cache_responses + # A fallback hop routes the target group in a copy of the metadata (`run_async_fallback`). + if fallback_level > 0: + field: Final = plan.prepare_field or "metadata" + kwargs[field] = {**as_dict(kwargs.get(field)), "model_group": requested} + match plan.kind: + case "chat": + host = core.host # rebind-ok: rebound as the Python router rebinds it + drop_unsupported_classifier_reasoning_effort(deployment, requested, kwargs, host) + params = ( # rebind-ok: rebound as the Python router rebinds it + deployment_params_with_request_reasoning_override(data, kwargs) + ) + silent_model = params.pop("silent_model", None) # rebind-ok: rebound as the Python router rebinds it + start_silent_experiments(core, plan, silent_model, named.get("messages"), kwargs) + if "messages" in named: + kwargs.setdefault("messages", named["messages"]) + update_kwargs_with_deployment(shaping, deployment, kwargs, cache=cache, splits_stream_timeout=True) + kwargs.pop("silent_model", None) + final = without_router_only_kwargs( # rebind-ok: rebound as the Python router rebinds it + { + **params, + **named, + "caching": caching, + "client": model_client(core, plan, deployment, kwargs), + **kwargs, + } + ) + return AttemptCall(plan.host_op, final, kwargs, str(params.get("model")), group) + case "text" | "embedding" | "image" | "transcription" | "rerank" | "adapter": + update_kwargs_with_deployment( + shaping, deployment, kwargs, cache=cache, splits_stream_timeout=plan.splits_stream_timeout + ) + refused: Final = embedding_batch_refusal(plan, deployment, named.get("input"), kwargs, requested) + if refused is not None: + return AttemptCall(plan.host_op, kwargs, kwargs, model_name, group, refused) + final = { # rebind-ok: rebound as the Python router rebinds it + **data, + **named, + "caching": caching, + "client": model_client(core, plan, deployment, kwargs), + **kwargs, + } + return AttemptCall(plan.host_op, without_router_only_kwargs(final), kwargs, model_name, group) + case "speech": + update_kwargs_with_deployment(shaping, deployment, kwargs, cache=cache) + voice = kwargs.pop("voice", None) # rebind-ok: rebound as the Python router rebinds it + final = { # rebind-ok: rebound as the Python router rebinds it + **data, + **named, + "voice": data.get("voice") if voice is None else voice, + "client": model_client(core, plan, deployment, kwargs), + **kwargs, + } + return AttemptCall(plan.host_op, final, kwargs, model_name, group) + case "batch" | "cancel_batch" | "vector_store_create": + if plan.kind == "cancel_batch": + data.update(_resolved_credentials(core, deployment, requested)) + data.pop("litellm_credential_name", None) + update_kwargs_with_deployment(shaping, deployment, kwargs, cache=cache, function_name=plan.function_name) + final = { # rebind-ok: rebound as the Python router rebinds it + **data, + "custom_llm_provider": _provider(core, data), + "caching": caching, + "client": model_client(core, plan, deployment, kwargs), + **kwargs, + } + return AttemptCall(plan.host_op, final, kwargs, str(data.get("model")), group) + case "generic": + return _generic_attempt(core, plan, kwargs, deployment, data, requested, group) + case "generic_sync": + update_kwargs_with_deployment( + shaping, + deployment, + kwargs, + cache=cache, + function_name="generic_api_call", + splits_stream_timeout=invoked_handler_name(core, plan) in STREAM_TIMEOUT_AWARE_GENERIC_HANDLERS, + ) + endpoint = kwargs.get("endpoint") # rebind-ok: rebound as the Python router rebinds it + if endpoint and isinstance(endpoint, str): + kwargs["endpoint"] = endpoint.replace(requested, model_name) + kwargs["model"] = model_name + final = { # rebind-ok: rebound as the Python router rebinds it + **data, + "custom_llm_provider": _generic_provider(data), + "caching": caching, + **kwargs, + } + return AttemptCall(plan.host_op, final, kwargs, model_name, group) + case "moderation": + update_kwargs_with_deployment(shaping, deployment, kwargs, cache=cache) + kwargs["model"] = model_name + kwargs.update(data) + return AttemptCall(plan.host_op, kwargs, kwargs, model_name, group) + case "search": + return _search_attempt(core, plan, kwargs, group) + case "create_file": + # `_acreate_file` creates the file on every healthy deployment of the hop's name, not + # only the picked one: the driver hands the request to `create_file_everywhere`. + return AttemptCall(plan.host_op, kwargs, kwargs, requested, group) + + +def _generic_attempt( + core: EndpointCore, + plan: Plan, + kwargs: dict[str, object], # mutable-ok: the attempt's own kwargs + deployment: Mapping[str, object], + data: dict[str, object], # mutable-ok: the attempt's copy of the deployment params + requested: str, + group: str, +) -> AttemptCall: + """`_ageneric_api_call_with_fallbacks_helper` for one deployment.""" + from litellm.llms.base_llm.passthrough.transformation import replace_path_segment + from litellm.router_rust.invoke.websocket import turn_admission + + handler: Final = invoked_handler_name(core, plan) + socket: Final = handler == "_aresponses_websocket" + update_kwargs_with_deployment( + core.shaping, + deployment, + kwargs, + cache=core.routing_cache, + function_name="_ageneric_api_call_with_fallbacks", + splits_stream_timeout=handler in STREAM_TIMEOUT_AWARE_GENERIC_HANDLERS, + opens_responses_socket=socket, + ) + model_name: Final = str(data.get("model")) + endpoint: Final = kwargs.get("endpoint") + if endpoint: + try: + stripped, _, _, _ = core.host.invoke("get_llm_provider")( + model=model_name, + custom_llm_provider=kwargs.get("custom_llm_provider"), + api_base=kwargs.get("api_base"), + ) + replacement = str(stripped) # rebind-ok: rebound as the Python router rebinds it + except Exception: # noqa: BLE001 # as the Python router: an unresolvable model is used as is + replacement = model_name # rebind-ok: rebound as the Python router rebinds it + kwargs["endpoint"] = ( # rebind-ok: the request kwargs are shaped in place, as the Python router shapes them + replace_path_segment(str(endpoint), requested, replacement) + ) + provider: Final = _generic_provider(data) + final: Final = { + **without_retrieve_stream_defaults(data, kwargs), + "caching": core.shaping.cache_responses, + **kwargs, + "model": model_name, + **_session_with_model(kwargs.get("session"), model_name), + **({"turn_admission": turn_admission(core, requested, deployment, kwargs)} if socket else {}), + } + if provider is not None: + final["custom_llm_provider"] = provider + return AttemptCall(plan.host_op, final, kwargs, model_name, group) + + +_SEARCH_PROTECTED_PARAMS: Final = frozenset({"search_provider", "api_key", "api_base"}) + + +def _search_attempt( + core: EndpointCore, + plan: Plan, + kwargs: dict[str, object], # mutable-ok: the attempt's own kwargs + group: str, +) -> AttemptCall: + """`SearchAPIRouter.async_search_with_fallbacks_helper`: one of the search tools named + ``group`` at random, called with its provider, credentials and params under the request's.""" + import random + + tools: Final = [tool for tool in core.endpoints.search_tools if as_dict(tool).get("search_tool_name") == group] + if not tools: + raise ValueError(f"Search tool '{group}' not found in router.search_tools") + selected: Final = random.choice(tools) # load balancing, as the Python router picks + params: Final = as_dict(as_dict(selected).get("litellm_params")) + search_provider: Final = params.get("search_provider") + if not search_provider: + raise ValueError(f"search_provider not found in litellm_params for search tool '{group}'") + search_params: Final = { + key: value + for source in (params, kwargs) + for key, value in source.items() + if key not in _SEARCH_PROTECTED_PARAMS and value is not None + } # comprehension-ok: the tool's params, then the request's + final: Final = { + "search_provider": search_provider, + "api_key": params.get("api_key"), + "api_base": params.get("api_base"), + **search_params, + } + return AttemptCall(plan.host_op, final, kwargs, group, group) + + +def unrouted_fields(plan: Plan) -> Mapping[str, object]: + """A search tool name is no model group: the Rust core hands every attempt to the host as is.""" + return {"unrouted": True} if plan.kind == "search" else {} + + +def failover_fields(plan: Plan) -> Mapping[str, object]: + """An endpoint whose attempt stamps no ``failed_deployment_id`` never runs weighted failover.""" + return {} if plan.stamps_failed_deployment else {"skip_weighted_failover": True} + + +# The Python router's handler names for the plans whose `function_name` differs from it. +_PYTHON_HANDLER_NAMES: Final[Mapping[str, str]] = MappingProxyType({"acreate_file": "_acreate_file"}) + + +def scoped_fallback_fields(plan: Plan, request: Mapping[str, object]) -> Mapping[str, object]: + """A request that creates or names a resource only one provider holds (a file upload, a batch + on an uploaded file, a fine-tuning job) falls back within its own model group only + (``run_async_fallback``'s ``same_model_group_only``), decided by the Python router's own rules.""" + from litellm.router_utils.fallback_event_handlers import ( + creates_provider_scoped_resource, + references_provider_scoped_resource, + ) + + function_name: Final = plan.function_name or plan.host_op + handlers: Final = { + "original_function": _Handler(_PYTHON_HANDLER_NAMES.get(function_name, function_name)), + "original_generic_function": _Handler(plan.host_op), + } + kwargs: Final = {**request, **handlers} + scoped: Final = references_provider_scoped_resource(kwargs) or creates_provider_scoped_resource(kwargs) + return {"same_group_fallbacks_only": True} if scoped else {} + + +@dataclasses.dataclass(frozen=True, slots=True) +class _Handler: + """A stand-in for a handler function: the scoped-resource rules read only its name.""" + + __name__: str + + +def _session_with_model(session: object, model_name: str) -> Mapping[str, object]: + if not isinstance(session, Mapping): + return {} + typed: Final = as_dict(session) # pyright: ignore[reportUnknownArgumentType] # litellm's untyped helpers and request values + if "model" not in typed and typed.get("type") != "transcription": + return {} + from litellm.realtime_api.main import with_resolved_session_model + + return {"session": with_resolved_session_model(typed, model_name)} + + +def _generic_provider(data: Mapping[str, object]) -> str | None: + from litellm.litellm_core_utils.routing_common_utils import provider_for_generic_call + + return provider_for_generic_call(data) # pyright: ignore[reportArgumentType] # deployment params + + +def model_client( + core: EndpointCore, plan: Plan, deployment: Mapping[str, object], kwargs: Mapping[str, object] +) -> object: + """The SDK client cached for the deployment (`_get_async_openai_model_client`; `_get_client` + with the sync keys for sync completion and embedding), dropped when the request carries + another api key.""" + from litellm.litellm_core_utils.core_helpers import get_parent_otel_span_from_kwargs + + if as_dict(deployment.get("litellm_params")).get("kubernetes_pod_discovery") is True: + return None + model_id: Final = as_dict(deployment.get("model_info")).get("id") + sync_keys: Final = plan.sync and plan.kind in ("chat", "embedding") + stream: Final = "stream_" if kwargs.get("stream") is True else "" + key: Final = f"{model_id}_{stream}client" if sync_keys else f"{model_id}_{stream}async_client" + span: Final = get_parent_otel_span_from_kwargs(dict(kwargs)) + client: Final = _cached_client(core, key, span, local_only=not sync_keys) + dynamic_api_key: Final = kwargs.get("api_key") + if dynamic_api_key is not None and client is not None and dynamic_api_key != getattr(client, "api_key", None): + return None + return client + + +def _cached_client(core: EndpointCore, key: str, span: object, *, local_only: bool) -> object: + cache: Final = core.routing_cache + with service_target(ROUTER_CLIENTS_TARGET): + if local_only: + return cache.get_cache(key=key, local_only=True, parent_otel_span=span) # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType, reportArgumentType] # the cache holds untyped values + return cache.get_cache(key=key, parent_otel_span=span) # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType, reportArgumentType] # see above + + +def _provider(core: EndpointCore, data: Mapping[str, object]) -> object: + """The deployment's declared provider, else the one its model names (`get_llm_provider`).""" + declared: Final = data.get("custom_llm_provider") + _, inferred, _, _ = core.host.invoke("get_llm_provider")(model=data["model"], custom_llm_provider=declared) + return declared or inferred + + +def _resolved_credentials(core: EndpointCore, deployment: Mapping[str, object], requested: str) -> Mapping[str, object]: + """`get_deployment_credentials_with_provider` for the picked deployment (`_acancel_batch`).""" + registry: Final = core.registry() + model_id: Final = as_dict(deployment.get("model_info")).get("id") + resolved: Final = registry.credentials(str(model_id or requested), form="with_provider") + return as_dict(resolved) # pyright: ignore[reportUnknownArgumentType] # litellm's untyped helpers and request values + + +def embedding_batch_refusal( + plan: Plan, deployment: Mapping[str, object], input: object, kwargs: Mapping[str, object], model: str +) -> Exception | None: + """``_embedding``'s check, past the deployment's kwargs shaping: an input the deployment admits + drops the mark an earlier rejected attempt left on the shared logging object.""" + if plan.kind != "embedding": + return None + from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging + from litellm.router_utils.embedding_batch_size_check import ( + clear_rejected_before_provider_call, + embedding_batch_too_large_error, + ) + + error: Final = embedding_batch_too_large_error( + deployment=deployment, + input=input, + request_kwargs=kwargs, + model=model, # pyright: ignore[reportArgumentType] # opaque values + ) + logging_obj: Final = kwargs.get("litellm_logging_obj") + if error is None and isinstance(logging_obj, LiteLLMLogging): + clear_rejected_before_provider_call(logging_obj) + return error + + +def log_refused_call_sync(call: AttemptCall, refused: Exception) -> None: + """``_embedding`` logs its refusal before raising it: no provider call will.""" + from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging + from litellm.router_utils.embedding_batch_size_check import log_rejection_before_provider_call + + logging_obj: Final = call.kwargs.get("litellm_logging_obj") + if isinstance(logging_obj, LiteLLMLogging): + log_rejection_before_provider_call(logging_obj, refused) + + +async def log_refused_call(call: AttemptCall, refused: Exception) -> None: + """``log_refused_call_sync`` for ``_aembedding``.""" + from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging + from litellm.router_utils.embedding_batch_size_check import async_log_rejection_before_provider_call + + logging_obj: Final = call.kwargs.get("litellm_logging_obj") + if isinstance(logging_obj, LiteLLMLogging): + await async_log_rejection_before_provider_call(logging_obj, refused) + + +def stamp_error( + plan: Plan, error: BaseException, deployment: Mapping[str, object], kwargs: Mapping[str, object] +) -> None: + """What each attempt helper stamps on the error it raises.""" + match plan.stamp: + case "deployment": + stamp_deployment_num_retries(error, deployment) + stamp_failed_deployment_id(error, kwargs.get("model_info") or deployment.get("model_info")) + case "retry_skip": + stamp_retry_skip_deployment_id(error, kwargs) + case "none": + pass + if plan.stamp != "deployment" and plan.stamps_failed_deployment: + # `_stamp_failed_attempt`: the deployment `_update_kwargs_with_deployment` put on the kwargs. + stamp_failed_deployment_id(error, kwargs.get("model_info")) + + +# --- silent experiments --------------------------------------------------------------------- + + +_SilentModel: TypeAlias = str | list[str] +_SILENT_MODEL_ADAPTER: Final[TypeAdapter[_SilentModel]] = TypeAdapter(_SilentModel) + + +def _silent_targets(silent_model: object) -> tuple[str, ...]: + """`_silent_experiment_targets`: none inside a mirrored call or for a malformed value.""" + from litellm._internal_context import in_mirrored_call + + if silent_model is None or in_mirrored_call(): + return () + try: + targets: Final = _SILENT_MODEL_ADAPTER.validate_python(silent_model) + except ValidationError: + verbose_router_logger.warning( + "silent_model must be a model name or a list of model names, got %r; skipping shadow traffic", + silent_model, + ) + return () + return (targets,) if isinstance(targets, str) else tuple(targets) + + +def _silent_kwargs(kwargs: Mapping[str, object]) -> dict[str, object]: # mutable-ok: the experiment's own kwargs + """`_get_silent_experiment_kwargs`.""" + from litellm.litellm_core_utils.core_helpers import ( + safe_deep_copy, # pyright: ignore[reportUnknownVariableType] # litellm's untyped helpers and request values + ) + + silent: Final = safe_deep_copy(dict(kwargs)) # pyright: ignore[reportUnknownVariableType] # litellm's untyped helpers and request values + original_metadata: Final = kwargs.get("metadata") + if original_metadata is not None and silent.get("metadata") is original_metadata: # pyright: ignore[reportUnknownMemberType] # litellm's untyped helpers and request values + silent["metadata"] = dict(original_metadata) # pyright: ignore[reportArgumentType, reportCallIssue] # the request's metadata dict + if "metadata" not in silent: + silent["metadata"] = {} + metadata: Final = silent["metadata"] # pyright: ignore[reportUnknownVariableType] # litellm's untyped helpers and request values + if isinstance(metadata, dict): + metadata.pop("litellm_parent_otel_span", None) # pyright: ignore[reportUnknownMemberType] # the request's metadata dict + metadata["is_silent_experiment"] = True + silent.pop("litellm_call_id", None) # pyright: ignore[reportUnknownMemberType] # litellm's untyped helpers and request values + silent.pop("litellm_logging_obj", None) # pyright: ignore[reportUnknownMemberType] # litellm's untyped helpers and request values + silent.pop("standard_logging_object", None) # pyright: ignore[reportUnknownMemberType] # litellm's untyped helpers and request values + return silent # pyright: ignore[reportUnknownVariableType] # litellm's untyped helpers and request values + + +def _snapshot(kwargs: Mapping[str, object]) -> dict[str, object]: # mutable-ok: a snapshot handed to a background task + snapshot: Final = dict(kwargs) + metadata: Final = snapshot.get("metadata") + if isinstance(metadata, Mapping): + snapshot["metadata"] = dict(as_dict(metadata)) # pyright: ignore[reportUnknownArgumentType] # litellm's untyped helpers and request values + return snapshot + + +async def _run_silent(core: EndpointCore, silent_model: str, messages: object, kwargs: Mapping[str, object]) -> None: + from litellm._internal_context import mirrored_call + + remaining: Final = without(kwargs, "stream") + with mirrored_call(): + response: Final = await core.acompletion( + model=silent_model, messages=messages, stream=bool(kwargs.get("stream", False)), **remaining + ) + from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper + + if not isinstance(response, CustomStreamWrapper): + return + async for _ in response: + pass + + +async def _silent_async(core: EndpointCore, silent_model: str, messages: object, kwargs: Mapping[str, object]) -> None: + try: + if as_dict(kwargs.get("metadata")).get("is_silent_experiment", False): + return + silent: Final = _silent_kwargs(kwargs) + as_dict_metadata: Final = silent["metadata"] + if isinstance(as_dict_metadata, dict): + as_dict_metadata["model_group"] = silent_model + await _run_silent(core, silent_model, copy.deepcopy(messages), silent) + except Exception: # noqa: BLE001, S110 # a silent experiment never fails the primary call + pass + + +def _silent_thread(core: EndpointCore, silent_model: str, messages: object, kwargs: Mapping[str, object]) -> None: + try: + if as_dict(kwargs.get("metadata")).get("is_silent_experiment", False): + return + silent: Final = _silent_kwargs(kwargs) + metadata: Final = silent["metadata"] + if isinstance(metadata, dict): + metadata["model_group"] = silent_model + loop: Final = asyncio.new_event_loop() + asyncio.set_event_loop(loop) + try: + + async def run() -> None: + await _run_silent(core, silent_model, copy.deepcopy(messages), silent) + pending: Final = asyncio.all_tasks() + current: Final = asyncio.current_task() + if current is not None: + pending.discard(current) + if pending: + await asyncio.gather(*pending, return_exceptions=True) + + loop.run_until_complete(run()) + finally: + loop.close() + except Exception: # noqa: BLE001, S110 # a silent experiment never fails the primary call + pass + + +def start_silent_experiments( + core: EndpointCore, plan: Plan, silent_model: object, messages: object, kwargs: Mapping[str, object] +) -> None: + """Mirror the call to each silent model in the background: a task on the async path, a daemon + thread on the sync one.""" + for target in _silent_targets(silent_model): + if plan.sync: + threading.Thread( + target=_silent_thread, args=(core, target, messages, _snapshot(kwargs)), daemon=True + ).start() + else: + _ = asyncio.create_task( + _silent_async(core, target, messages, _snapshot(kwargs)) + ) # fire and forget, as the Python router + + +# --- the endpoints ---------------------------------------------------------------------------- + +_KEEPS_CLIENT: Final = frozenset( + { + "acreate_realtime_client_secret", + "arealtime_calls", + "acreate_realtime_transcription_session", + "acreate_realtime_translation_client_secret", + "arealtime_translation_calls", + } +) +_KEEPS_PROVIDER_AND_CLIENT: Final = frozenset({"afile_delete", "afile_content"}) +_RESPONSE_ID_ENDPOINTS: Final = frozenset( + {"aget_responses", "acancel_responses", "acompact_responses", "adelete_responses", "alist_input_items"} +) +_CONTAINER_ENDPOINTS: Final = frozenset( + {"acreate_container", "alist_containers", "aretrieve_container", "adelete_container", *_CONTAINER_FILE_ROUTED} +) +_VECTOR_STORE_ENDPOINTS: Final = frozenset( + { + *(f"avector_store_{name}" for name in ("search", "create", "retrieve", "list", "update", "delete")), + *(f"avector_store_file_{name}" for name in ("create", "list", "retrieve", "content", "update", "delete")), + } +) +_INTERACTION_ENDPOINTS: Final = frozenset( + {"acreate_interaction", "aget_interaction", "adelete_interaction", "acancel_interaction"} +) +_MISSING: Final = object() + + +def pinned_video_deployment(plan: Plan, request: Mapping[str, object]) -> Mapping[str, object]: + """`_video_call_pin`: the deployment a video id names, for a call on that video only.""" + from litellm.types.videos.utils import VIDEO_ID_CALL_TYPES, decode_video_id_with_provider + + video_id: Final = request.get("video_id") + if plan.endpoint not in VIDEO_ID_CALL_TYPES or not isinstance(video_id, str): + return {} + deployment_id: Final = decode_video_id_with_provider(video_id).get("deployment_id") + return {"pinned_deployment_id": deployment_id} if deployment_id else {} + + +async def _await(value: object) -> object: + """``await value``, raising the TypeError Python raises for a value that is not awaitable.""" + return await value # pyright: ignore[reportGeneralTypeIssues, reportUnknownVariableType] # deliberately awaited as the Python router does + + +def _missing_model() -> TypeError: + return TypeError("_ageneric_api_call_with_fallbacks() missing 1 required positional argument: 'model'") + + +class RustEndpoints: + """The public endpoint methods of one Rust core, beyond the routed call loop the core drives.""" + + def __init__(self, core: EndpointCore, config: "RouterConfig") -> None: + from litellm.litellm_core_utils.request_timeout_resolver import get_configured_request_timeout + from litellm.router_utils.request_priority import InvalidPriority, parse_default_priority + + self._core: Final = core + self.counters: Final = Counters() + explicit: Final = config.get("timeout") + self._explicit_timeout: Final = explicit + self._request_timeout: Final = get_configured_request_timeout() if explicit is not None else None + self._stream_timeout: Final = config.get("stream_timeout") + # apply() flips it after construction, as on the Python core + self._cache_responses: bool = bool(config.get("cache_responses")) + self._assistants_config: Final = ( + as_dict(config.get("assistants_config")) if config.get("assistants_config") else None + ) + # These three are assigned by apply() after construction, as on the Python core. + tools: Final = list(_as_items(config.get("search_tools"))) + self._search_tools: list[object] = tools # mutable-ok: the router's live list + self._guardrails: list[object] = list(_as_items(config.get("guardrail_list"))) # mutable-ok: same + parsed: Final = parse_default_priority(config.get("default_priority")) + self._default_priority: Final = None if isinstance(parsed, InvalidPriority) else parsed + self._polling_interval: Final = config.get("polling_interval") + self._caching_groups: Final = config.get("caching_groups") + self._default_max_parallel_requests: object = config.get("default_max_parallel_requests") + self._scheduler: object = None + + @property + def cache_responses(self) -> bool: + return self._cache_responses + + def set_cache_responses(self, value: bool) -> None: + """``cache_responses`` assigned after construction (the proxy turns it on once its cache + exists): later calls pass ``caching=value``, as the Python core's flag does.""" + self._cache_responses = value + + def assign_setting(self, name: str, value: object) -> None: + """``search_tools``, ``guardrail_list`` or ``default_max_parallel_requests`` assigned after + construction, kept as given: later searches, guardrail picks and slot limits built from + then on read it, as the Python core's attributes are read.""" + if name == "search_tools": + self._search_tools = value # pyright: ignore[reportAttributeAccessIssue] # stored as given, as the Python core stores it + elif name == "guardrail_list": + self._guardrails = value # pyright: ignore[reportAttributeAccessIssue] # same + elif name == "default_max_parallel_requests": + self._default_max_parallel_requests = value + else: + raise ValueError(f"not an endpoint setting: {name}") + + @property + def default_max_parallel_requests(self) -> object: + """The setting as given: ``settings()`` reports it and the slot limits read it.""" + return self._default_max_parallel_requests + + @property + def search_tools(self) -> list[object]: # mutable-ok: the router's live list, as the Python router reports it + return self._search_tools + + @property + def guardrail_list(self) -> list[object]: # mutable-ok: same + return self._guardrails + + @staticmethod + def prepare_host(config: "RouterConfig", host: "PythonHost") -> None: + """`cache_responses` sets the process cache, unless one is set already (the proxy's).""" + if not config.get("cache_responses") or host.global_cache() is not None: + return + from litellm.caching.caching import Cache + + redis_url: Final = config.get("redis_url") + redis_host: Final = config.get("redis_host") + redis_port: Final = config.get("redis_port") + cache_config: Final[dict[str, object]] = {} # mutable-ok: the request kwargs the Python router mutates in place + cache_type: str = "local" # rebind-ok: redis when the router is given one + if redis_url is not None or (redis_host is not None and redis_port is not None): + cache_type = "redis" # rebind-ok: rebound as the Python router rebinds it + if redis_url is not None: + cache_config["url"] = redis_url + if redis_host is not None: + cache_config["host"] = redis_host + if redis_port is not None: + cache_config["port"] = str(redis_port) + if config.get("redis_password") is not None: + cache_config["password"] = config.get("redis_password") + if config.get("redis_db") is not None: + cache_config["db"] = str(config.get("redis_db")) + cache_config.update(as_dict(config.get("cache_kwargs"))) + host.set_global_cache(Cache(type=cache_type, **cache_config)) # pyright: ignore[reportArgumentType] # the Python router's cache config + + def shaping(self, resolved: "Resolved") -> Shaping: + return Shaping( + default_litellm_params=self.default_litellm_params(resolved), + timeout=resolved.timeout, + explicit_timeout=self._explicit_timeout, + request_timeout=self._request_timeout, + stream_timeout=self._stream_timeout, + cache_responses=self._cache_responses, + ) + + def default_litellm_params(self, resolved: "Resolved") -> Mapping[str, object]: + """The defaults the Python constructor completes before any call: the router ``timeout`` + argument, no client retries, and the caching groups in the default metadata.""" + configured: Final = resolved.default_litellm_params + metadata: Final = configured.get("metadata") + return MappingProxyType( + { + "timeout": self._explicit_timeout, + "max_retries": 0, + **configured, + "metadata": { + **(as_dict(metadata) if isinstance(metadata, Mapping) else {}), # pyright: ignore[reportUnknownArgumentType] # the configured metadata + "caching_groups": self._caching_groups, + }, + } + ) + + def prepare( + self, plan: Plan, request: Mapping[str, object], model: str + ) -> dict[str, object]: # mutable-ok: the call's kwargs + return prepare_request(plan, request, model, self._core.resolved.model_group_alias) + + # -- host and registry --------------------------------------------------------------------- + + def _invoke(self, op: str) -> Callable[..., object]: + return self._core.host.invoke(op) + + async def _direct(self, op: str, kwargs: Mapping[str, object]) -> object: + return await _await(self._invoke(op)(**kwargs)) + + def _group(self, name: str) -> list[Mapping[str, object]]: # mutable-ok: a fresh list + """`_core_get_model_list(model_name=name)`: the deployments of the group ``name``.""" + records: Final = self._core.handle.deployments() + return [record for record in (as_dict(r) for r in records) if record.get("model_name") == name] + + def _is_model_group(self, name: str) -> bool: + """`_core_get_model_list(model_name=name)`: an alias, a routing group or a wildcard counts.""" + return bool(self._core.group_deployments(name)) + + # -- members ------------------------------------------------------------------------------ + + def member(self, name: str) -> Callable[..., object] | None: + """The public endpoint ``name``, bound to this core, or None when the core has none.""" + if name in _UNBRANCHED: + return _answers_none + if name in _ASSISTANTS: + return self._assistants(name) + if name in _MANAGED_AGENTS: + return self._managed_agents(name) + if name in _SYNC_VECTOR_STORE_FILES: + return self._sync_vector_store_files(name) + if name == "create_interaction": + return self._factory("create_interaction", "interactions.create") + direct: Final = { + "avector_store_create": self.avector_store_create, + "acreate_file": self.acreate_file, + "acreate_batch": self.acreate_batch, + "acancel_batch": self.acancel_batch, + "aretrieve_batch": self.aretrieve_batch, + "alist_batches": self.alist_batches, + "aadapter_completion": self.aadapter_completion, + "abatch_completion": self.abatch_completion, + "abatch_completion_one_model_multiple_requests": self.abatch_completion_one_model_multiple_requests, + "abatch_completion_fastest_response": self.abatch_completion_fastest_response, + "schedule_acompletion": self.schedule_acompletion, + "priority_header_update": self.priority_header_update, + "aguardrail": self.aguardrail, + "asearch": self.asearch, + }.get(name) + if direct is not None: + return direct + if name == "responses": + return self._sync_factory(name, PLANS["responses"]) + plan: Final = generic_plan(name) + if plan is None or name in PLANS: + return None + if plan.sync: + return self._sync_factory(name, plan) + return self._factory(name, plan.host_op) + + # -- the factory wrappers ------------------------------------------------------------------ + + def _factory(self, name: str, host_op: str) -> Callable[..., Awaitable[object]]: + async def async_wrapper( + custom_llm_provider: str | None = None, + client: object = None, + **kwargs: object, # kwargs-ok: the endpoint's request + ) -> object: + return await self._wrapped(name, host_op, custom_llm_provider, client, dict(kwargs)) + + async_wrapper.__name__ = name + return async_wrapper + + async def _wrapped( + self, + name: str, + host_op: str, + custom_llm_provider: str | None, + client: object, + kwargs: dict[str, object], # mutable-ok: the wrapper's own kwargs + ) -> object: + plan: Final = generic_plan(name) or _generic(name, host_op) + if name in _CONTAINER_ENDPOINTS: + return await self._containers(plan, custom_llm_provider, kwargs) + if name in _RESPONSE_ID_ENDPOINTS: + from litellm.responses.utils import ResponsesAPIRequestUtils + + model_id: Final = ResponsesAPIRequestUtils.get_model_id_from_response_id(kwargs.get("response_id")) # pyright: ignore[reportArgumentType] # any id the caller sends + if model_id is not None: + kwargs[ # rebind-ok: the request kwargs are shaped in place, as the Python router shapes them + "model" + ] = model_id + return await self._ageneric(plan, kwargs) + if name in _VECTOR_STORE_ENDPOINTS: + return await self._vector_stores(name, plan, custom_llm_provider, kwargs) + if name in _INTERACTION_ENDPOINTS or name == "create_interaction": + return await self._interactions(plan, custom_llm_provider, kwargs) + if name in _KEEPS_CLIENT: + return await self._ageneric(plan, {**kwargs, "client": client}) + if name in _KEEPS_PROVIDER_AND_CLIENT: + return await self._ageneric(plan, {**kwargs, "custom_llm_provider": custom_llm_provider, "client": client}) + if name == "allm_passthrough_route": + if client: + kwargs[ # rebind-ok: the request kwargs are shaped in place, as the Python router shapes them + "client" + ] = client + return await self._ageneric(plan, kwargs, passthrough=True) + return await self._ageneric(plan, kwargs) + + async def generic(self, plan: Plan, request: Mapping[str, object]) -> object: + """A dispatched generic endpoint called with its whole request (the wrapper's + ``custom_llm_provider`` and ``client`` are consumed by the wrapper).""" + kwargs: Final = without(request, "custom_llm_provider", "client") + return await self._ageneric(plan, kwargs) + + async def _ageneric( + self, + plan: Plan, + kwargs: Mapping[str, object], + *, + passthrough: bool = False, + ) -> object: + """`_ageneric_api_call_with_fallbacks`.""" + if "model" not in kwargs: + raise _missing_model() + + async def passthrough_call(prepared: Mapping[str, object], name: str) -> object: + sent: Final = without(prepared, *ROUTING_KEYS) + return await self._direct(plan.host_op, {"model": name, **sent}) + + return await self._core.drive(plan, kwargs, passthrough=passthrough_call if passthrough else None) + + def _sync_factory(self, name: str, plan: Plan) -> Callable[..., object]: + is_vector_store: Final = name.startswith("vector_store_") + + def sync_wrapper( + custom_llm_provider: str | None = None, + client: object = None, + **kwargs: object, # kwargs-ok: the endpoint's request + ) -> object: + if not is_vector_store: + if "model" not in kwargs: + raise TypeError( + "_generic_api_call_with_fallbacks() missing 1 required positional argument: 'model'" + ) + return self._core.call_sync(plan, kwargs, _SINGLE_ATTEMPT) + provider_kwargs: dict[ # mutable-ok: the request kwargs the Python router mutates in place # rebind-ok: rebound as the Python router rebinds it + str, object + ] = ( + {**kwargs, "custom_llm_provider": custom_llm_provider} + if custom_llm_provider and "custom_llm_provider" not in kwargs + else dict(kwargs) + ) + if name == "vector_store_search": + provider_kwargs = { # rebind-ok: rebound as the Python router rebinds it + **provider_kwargs, + "_direct_vector_store_embedding_executor": self._embedding_executor(kwargs), + } + if provider_kwargs.get("model"): + return self._core.call_sync(plan, provider_kwargs, _SINGLE_ATTEMPT) + if name == "vector_store_search": + return self._invoke(plan.host_op)(**{**provider_kwargs, "router": self._core.router}) + return self._invoke(plan.host_op)(**provider_kwargs) + + sync_wrapper.__name__ = name + return sync_wrapper + + def _embedding_executor(self, kwargs: Mapping[str, object]) -> object: + from litellm.llms.base_llm.vector_store.transformation import ( + RouterVectorStoreEmbeddingExecutor, + vector_store_request_metadata, + ) + + return RouterVectorStoreEmbeddingExecutor( + router=self._core.router, # pyright: ignore[reportArgumentType] # the router serves aembedding + metadata=vector_store_request_metadata(kwargs), + ) + + async def _containers( + self, + plan: Plan, + custom_llm_provider: str | None, + kwargs: dict[str, object], # mutable-ok: the request kwargs the Python router mutates in place + ) -> object: + """`_init_containers_api_endpoints`.""" + from litellm.responses.utils import ResponsesAPIRequestUtils + + if custom_llm_provider and "custom_llm_provider" not in kwargs: + kwargs[ # rebind-ok: the request kwargs are shaped in place, as the Python router shapes them + "custom_llm_provider" + ] = custom_llm_provider + container_id: Final = kwargs.get("container_id") + forwarded: Final = kwargs.get("model_id") + if isinstance(container_id, str): + decoded: Final = ResponsesAPIRequestUtils.decode_container_id(container_id) + original_id: Final = decoded.get("response_id", container_id) + if original_id != container_id: + kwargs[ # rebind-ok: the request kwargs are shaped in place, as the Python router shapes them + "container_id" + ] = original_id + decoded_provider: Final = decoded.get("custom_llm_provider") + if decoded_provider and kwargs.get("custom_llm_provider") == "openai": + kwargs[ # rebind-ok: the request kwargs are shaped in place, as the Python router shapes them + "custom_llm_provider" + ] = decoded_provider + model_id: Final = decoded.get("model_id") or ( + forwarded.strip() if isinstance(forwarded, str) and forwarded.strip() else None + ) + if model_id: + kwargs[ # rebind-ok: the request kwargs are shaped in place, as the Python router shapes them + "model" + ] = model_id + return await self._ageneric(plan, kwargs) + requested: Final = kwargs.get("model") + if isinstance(requested, str) and requested.strip(): + return await self._ageneric(plan, kwargs, passthrough=True) + return await self._direct(plan.host_op, kwargs) + + async def _vector_stores( + self, + name: str, + plan: Plan, + custom_llm_provider: str | None, + kwargs: dict[str, object], # mutable-ok: see _wrapped + ) -> object: + """`_init_vector_store_api_endpoints`.""" + if name == "avector_store_search": + kwargs = { # rebind-ok: the request kwargs are shaped in place, as the Python router shapes them + **kwargs, + "_direct_vector_store_embedding_executor": self._embedding_executor(kwargs), + } + if custom_llm_provider and "custom_llm_provider" not in kwargs: + kwargs[ # rebind-ok: the request kwargs are shaped in place, as the Python router shapes them + "custom_llm_provider" + ] = custom_llm_provider + if kwargs.get("model"): + return await self._ageneric(plan, kwargs) + if name == "avector_store_search": + return await self._direct(plan.host_op, {**kwargs, "router": self._core.router}) + return await self._direct(plan.host_op, kwargs) + + async def _interactions( + self, + plan: Plan, + custom_llm_provider: str | None, + kwargs: dict[str, object], # mutable-ok: the request kwargs the Python router mutates in place + ) -> object: + """`_init_interactions_api_endpoints`.""" + if custom_llm_provider and "custom_llm_provider" not in kwargs: + kwargs[ # rebind-ok: the request kwargs are shaped in place, as the Python router shapes them + "custom_llm_provider" + ] = custom_llm_provider + if "custom_llm_provider" not in kwargs: + kwargs[ # rebind-ok: the request kwargs are shaped in place, as the Python router shapes them + "custom_llm_provider" + ] = "gemini" + if kwargs.get("agent") and kwargs.get("model") == kwargs.get("agent"): + kwargs["model"] = ( # rebind-ok: the request kwargs are shaped in place, as the Python router shapes them + None + ) + if kwargs.get("model"): + return await self._ageneric(plan, kwargs) + return await self._direct(plan.host_op, kwargs) + + def _managed_agents(self, name: str) -> Callable[..., object]: + op: Final = _MANAGED_AGENTS[name] + + def prepared( + custom_llm_provider: str | None, + kwargs: dict[str, object], # mutable-ok: the request kwargs the Python router mutates in place + ) -> dict[str, object]: # mutable-ok: the call's kwargs + if custom_llm_provider and "custom_llm_provider" not in kwargs: + kwargs[ # rebind-ok: the request kwargs are shaped in place, as the Python router shapes them + "custom_llm_provider" + ] = custom_llm_provider + if "custom_llm_provider" not in kwargs: + kwargs[ # rebind-ok: the request kwargs are shaped in place, as the Python router shapes them + "custom_llm_provider" + ] = "gemini" + return kwargs + + if name.startswith("a"): + + async def managed_agents_wrapper( + custom_llm_provider: str | None = None, + client: object = None, + **kwargs: object, # kwargs-ok: the request + ) -> object: + return await self._direct(op, prepared(custom_llm_provider, dict(kwargs))) + + return managed_agents_wrapper + + def managed_agents_sync_wrapper( + custom_llm_provider: str | None = None, + client: object = None, + **kwargs: object, # kwargs-ok: the request + ) -> object: + return self._invoke(op)(**prepared(custom_llm_provider, dict(kwargs))) + + return managed_agents_sync_wrapper + + def _sync_vector_store_files(self, name: str) -> Callable[..., object]: + op: Final = _SYNC_VECTOR_STORE_FILES[name] + + def vector_store_file_sync_wrapper( + custom_llm_provider: str | None = None, + client: object = None, + **kwargs: object, # kwargs-ok: the request + ) -> object: + return self._invoke(op)(custom_llm_provider=custom_llm_provider, client=client, **kwargs) + + return vector_store_file_sync_wrapper + + def _assistants(self, name: str) -> Callable[..., Awaitable[object]]: + async def assistants_wrapper( + custom_llm_provider: str | None = None, + client: object = None, + **kwargs: object, # kwargs-ok: the request + ) -> object: + """`_pass_through_assistants_endpoint_factory`.""" + provider: object = custom_llm_provider # rebind-ok: the assistants config's when none is given + sent: Final = dict(kwargs) + if provider is None: + if self._assistants_config is None: + raise Exception( # noqa: TRY002 # the Python router raises a bare Exception + "'custom_llm_provider' must be set. Either via:\n `Router(assistants_config={'custom_llm_provider': ..})` \nor\n `router.arun_thread(custom_llm_provider=..)`" + ) + provider = self._assistants_config[ # rebind-ok: rebound as the Python router rebinds it + "custom_llm_provider" + ] + sent.update(as_dict(self._assistants_config.get("litellm_params"))) + return await self._direct(name, {"custom_llm_provider": provider, "client": client, **sent}) + + return assistants_wrapper + + # -- completions, priority and prompt management ---------------------------------------- + + def _drops_params(self, model: str, kwargs: Mapping[str, object]) -> bool: + from litellm.litellm_core_utils.routing_common_utils import get_request_team_id + from litellm.router_utils.request_priority import request_drops_params + + def deployment_params() -> tuple[Mapping[str, object], ...]: + names: Final = {as_dict(r).get("model_name") for r in self._core.handle.deployments()} + pinned: Final = ( + self._core.handle.deployment(model) if model not in names and isinstance(model, str) else None # pyright: ignore[reportUnnecessaryIsInstance] # a caller's value, typed for the common case + ) + if pinned is not None: + return (as_dict(as_dict(pinned).get("litellm_params")),) + if kwargs.get("specific_deployment") is True: + candidates = [ # rebind-ok: rebound as the Python router rebinds it + as_dict(r) + for r in self._core.handle.deployments() + if as_dict(as_dict(r).get("litellm_params")).get("model") == model + ] + else: + # `_core_get_model_list`: an alias's target, a routing group's members and a + # wildcard's deployments list here too. + candidates = [ # rebind-ok: rebound as the Python router rebinds it + as_dict(record) + for record in self._core.registry().deployments(model, team_id=get_request_team_id(kwargs)) or () + ] + return tuple(as_dict(candidate.get("litellm_params")) for candidate in candidates) + + return request_drops_params( + kwargs, + self._core.resolved.default_litellm_params, + deployment_params(), + ) + + def _split_priority( + self, model: str, kwargs: Mapping[str, object] + ) -> tuple[int | None, dict[str, object]]: # mutable-ok: the call's kwargs + """`_core_split_request_priority`.""" + from litellm.router_utils.request_priority import InvalidPriority, priority_error, resolve_request_priority + + priority: Final = resolve_request_priority( + requested=kwargs.get("priority"), + default_priority=self._default_priority, + drops_params=lambda: self._drops_params(model=model, kwargs=kwargs), + ) + if isinstance(priority, InvalidPriority): + raise priority_error(model=model, message=priority.message) + return priority, without(kwargs, "priority") + + def priority_header_update( + self, model: str, kwargs: Mapping[str, object], header: str | None + ) -> Mapping[str, int | str]: + from litellm.router_utils.request_priority import priority_header_fields + + return priority_header_fields(header, drops_params=lambda: self._drops_params(model=model, kwargs=kwargs)) + + def set_scheduler(self, scheduler: object) -> None: + """``router.scheduler = ...``, which the next prioritized call waits on.""" + self._scheduler = scheduler + + def delete_scheduler(self) -> None: + """``del router.scheduler``: a prioritized call raises from then on, as on the Python core.""" + self._scheduler = _NO_SCHEDULER + + def _get_scheduler(self) -> object: + if self._scheduler is _NO_SCHEDULER: + raise AttributeError("'RustCore' object has no attribute 'scheduler'", name="scheduler") + if self._scheduler is None: + from litellm.scheduler import Scheduler + + self._scheduler = ( + Scheduler(polling_interval=self._polling_interval) # pyright: ignore[reportArgumentType] # as configured + if self._polling_interval is not None + else Scheduler() + ) + return self._scheduler + + async def _schedule( + self, model: str, priority: int, run: Callable[[], Awaitable[object]], request: Mapping[str, object] + ) -> object: + """`_schedule_factory`: wait for the scheduler's turn, run, mark the response.""" + from litellm.litellm_core_utils.hidden_params import get_hidden_params + from litellm.scheduler import FlowItem, Scheduler + + scheduler: Final = self._get_scheduler() + assert isinstance(scheduler, Scheduler) + + async def healthy_deployments() -> Sequence[object]: + healthy: Final = await self._core.select(model, dict(request), kind="healthy") + return list(healthy) if isinstance(healthy, Sequence) else [healthy] + + await scheduler.wait_for_turn( + request=FlowItem(priority=priority, request_id=str(uuid.uuid4()), model_name=model), + timeout=self._core.resolved.timeout, # pyright: ignore[reportArgumentType] # the router timeout + get_healthy_deployments=healthy_deployments, + ) + try: + response: Final = await run() + except Exception as error: + setattr(error, "priority", priority) # noqa: B010 # the Python router stamps the same attribute + raise + hidden: Final = get_hidden_params(response) + if hidden is not None: + headers: Final = hidden.setdefault("additional_headers", {}) + if isinstance(headers, dict): + headers.update({"x-litellm-request-prioritization-used": True}) # pyright: ignore[reportUnknownMemberType] # litellm's untyped helpers and request values + return response + + async def acompletion(self, request: Mapping[str, object]) -> object: + model: Final = request.get("model") + if model is None and "model" in request: + return await self._core.drive(PLANS["acompletion"], request) + if "model" not in request: + raise TypeError("acompletion() needs a model name") + if not isinstance(model, str): + _ = self._core.routed_name(model) + kwargs: Final = {**request, "stream": request.get("stream", False)} + start_time: Final = time.time() + response: Final = await self._routed_acompletion(model, kwargs) # pyright: ignore[reportArgumentType] # the name as given + self._core.service_logged("acompletion", start_time, time.time(), kwargs) + return response + + async def _routed_acompletion(self, model: str, kwargs: Mapping[str, object]) -> object: + plan: Final = PLANS["acompletion"] + priority, rest = self._split_priority(model, kwargs) + from litellm.router_rust.invoke.prompt_management import ( + is_prompt_management_model, + prompt_management_acompletion, + ) + + if is_prompt_management_model(self._core, model): + prepared: Final = self.prepare(plan, rest, model) + messages: Final = rest.get("messages") + + def factory() -> Awaitable[object]: + return prompt_management_acompletion(self._core, model, messages, prepared) + + if priority is None: + return await factory() + return await self._schedule(model, priority, factory, prepared) + if priority is None: + return await self._core.drive(plan, rest) + return await self._schedule(model, priority, lambda: self._core.drive(plan, rest), rest) + + async def atext_completion(self, request: Mapping[str, object]) -> object: + model: Final = request.get("model") + if model is None and "model" in request: + return await self._core.drive(PLANS["atext_completion"], request) + if "model" not in request: + raise TypeError("atext_completion() needs a model name") + if not isinstance(model, str): + _ = self._core.routed_name(model) + kwargs: Final = without(request, "is_retry", "is_fallback", "is_async") + plan: Final = PLANS["atext_completion"] + priority, rest = self._split_priority(model, kwargs) # pyright: ignore[reportArgumentType] # the name as given + if priority is None: + return await self._core.drive(plan, rest) + return await self._schedule(model, priority, lambda: self._core.drive(plan, rest), rest) # pyright: ignore[reportArgumentType] # see above + + async def schedule_acompletion( + self, + model: str, + messages: object, + priority: int | None = None, + stream: bool = False, + **kwargs: object, # kwargs-ok: the request + ) -> object: + from litellm.router_utils.request_priority import priority_error, unqueueable_priority_message + + queued, _ = self._split_priority(model, {**kwargs, "priority": priority}) + if queued is None: + raise priority_error(model=model, message=unqueueable_priority_message(priority)) + return await self._core.acompletion(model=model, messages=messages, stream=stream, priority=queued, **kwargs) + + # -- batch completions ------------------------------------------------------------------- + + async def _no_exceptions(self, **kwargs: object) -> object: # kwargs-ok: the request + try: + return await self._core.acompletion(**kwargs) + except Exception as error: # noqa: BLE001 # returned in place, as the Python router does + return error + + async def abatch_completion( + self, + models: list[str], # mutable-ok: the request kwargs the Python router mutates in place + messages: object, + **kwargs: object, # kwargs-ok: the endpoint's open request + ) -> object: + if isinstance(messages, list) and all(isinstance(m, dict) for m in messages): # pyright: ignore[reportUnknownVariableType] # checked item by item + return await asyncio.gather( + *(self._no_exceptions(model=model, messages=messages, **kwargs) for model in models) # pyright: ignore[reportUnknownArgumentType] # litellm's untyped helpers and request values + ) + if isinstance(messages, list) and all(isinstance(m, list) for m in messages): # pyright: ignore[reportUnknownVariableType] # same + rows: Final[list[object]] = list(messages) # pyright: ignore[reportUnknownArgumentType] # same # mutable-ok: the request kwargs the Python router mutates in place + + async def indexed(model: str, message: object, idx: int) -> tuple[object, int]: + return await self._no_exceptions(model=model, messages=message, **kwargs), idx + + responses: Final = await asyncio.gather( + *( + indexed(model, message, idx) for idx, message in enumerate(rows) for model in models + ) # comprehension-ok: a flat walk of nested rows + ) + final: Final[list[list[object]]] = [ # mutable-ok: the request kwargs the Python router mutates in place + [] for _ in rows + ] + for response, idx in responses: + final[idx].append(response) + return final + return None + + async def abatch_completion_one_model_multiple_requests( + self, + model: str, + messages: list[object], # mutable-ok: the request kwargs the Python router mutates in place + **kwargs: object, # kwargs-ok: the endpoint's open request + ) -> object: + return await asyncio.gather(*(self._no_exceptions(model=model, messages=m, **kwargs) for m in messages)) + + async def abatch_completion_fastest_response( + self, + model: str, + messages: object, + stream: bool = False, + **kwargs: object, # kwargs-ok: the request + ) -> object: + from litellm.litellm_core_utils.hidden_params import get_hidden_params + from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper + from litellm.types.utils import ModelResponse + + models: Final = [m.strip() for m in model.split(",")] + + async def no_exceptions(model_name: str) -> object: + try: + return await self._core.acompletion(model=model_name, messages=messages, stream=stream, **kwargs) + except asyncio.CancelledError: + raise + except Exception as error: # noqa: BLE001 # see _no_exceptions + return error + + pending: set[ # mutable-ok: the request kwargs the Python router mutates in place # rebind-ok: rebound as the Python router rebinds it + asyncio.Task[object] + ] = {asyncio.create_task(no_exceptions(m)) for m in models} + while pending: + done, pending = await asyncio.wait(pending, return_when=asyncio.FIRST_COMPLETED) + for task in done: + result = task.result() + if isinstance(result, (ModelResponse, CustomStreamWrapper)): + for other in pending: + other.cancel() + hidden = get_hidden_params(result) + if hidden is not None: + hidden["fastest_response_batch_completion"] = True + return result + raise Exception("All tasks failed") # noqa: TRY002 # the Python router raises a bare Exception + + # -- moderation, text completion and adapter -------------------------------------------- + + async def amoderation(self, request: Mapping[str, object]) -> object: + """`_pass_through_moderation_endpoint_factory`: one pick when the model is a group.""" + kwargs: Final = without(request, "custom_llm_provider", "client") + plan: Final = PLANS["amoderation"] + model: Final = kwargs.get("model", "") + if not isinstance(model, str): + _ = self._core.routed_name(model) + if model and isinstance(model, str) and self._is_model_group(model): + return await self._core.drive(plan, kwargs, _SINGLE_ATTEMPT) + prepared: Final = prepare_request(plan, kwargs, str(model), self._core.resolved.model_group_alias) + return await self._direct("amoderation", prepared) + + def text_completion(self, request: Mapping[str, object]) -> object: + """The sync `text_completion`: one pick, no retries, the group name as the model.""" + from litellm.litellm_core_utils.routing_common_utils import deployment_model_id_pin + + kwargs: Final = without(request, "is_retry", "is_fallback", "is_async") + model: Final = kwargs.get("model") + prompt: Final = kwargs.get("prompt") + if "model" not in kwargs: + raise TypeError("text_completion() missing 1 required positional argument: 'model'") + kwargs["num_retries"] = kwargs.get("num_retries", self._core.resolved.num_retries) + metadata_bucket(kwargs, "metadata").update({"model_group": model}) + with endpoint_selection("text_completion", model=model): + deployment: Final = self._core.router.get_available_deployment( + model=model, # pyright: ignore[reportArgumentType] # the name as given + messages=_user_message(prompt), + specific_deployment=kwargs.pop("specific_deployment", None), # pyright: ignore[reportArgumentType] # the request's value, as given + request_kwargs=kwargs, + ) + data: Final = as_dict(deployment.get("litellm_params")) + for key, value in self._core.resolved.default_litellm_params.items(): + if key not in kwargs: + kwargs[key] = value + elif key == "metadata": + bucket = kwargs[key] + if isinstance(bucket, dict): + bucket.update(as_dict(value)) # pyright: ignore[reportUnknownMemberType] # the request's metadata dict + pinned: Final = deployment_model_id_pin(data) + if pinned is not None: + kwargs["model_id"] = pinned + return self._invoke("text_completion")(**{**data, "prompt": prompt, "caching": self._cache_responses, **kwargs}) + + async def aadapter_completion( + self, + adapter_id: str, + model: str, + is_retry: bool | None = False, + is_fallback: bool | None = False, + is_async: bool | None = False, + **kwargs: object, # kwargs-ok: the request + ) -> object: + return await self._core.drive( + PLANS["aadapter_completion"], {**kwargs, "model": model, "adapter_id": adapter_id} + ) + + # -- files, batches and vector stores --------------------------------------------------- + + async def acreate_batch(self, model: str, **kwargs: object) -> object: # kwargs-ok: the request + return await self._core.drive(PLANS["acreate_batch"], {**kwargs, "model": model}) + + async def acancel_batch(self, model: str, **kwargs: object) -> object: # kwargs-ok: the request + return await self._core.drive(PLANS["acancel_batch"], {**kwargs, "model": model}) + + async def avector_store_create(self, model: str | None, **kwargs: object) -> object: # kwargs-ok: the request + if model is None: + return await self._factory("avector_store_create", "vector_stores.main.acreate")(**kwargs) + return await self._core.drive(PLANS["avector_store_create"], {**kwargs, "model": model}, _SINGLE_ATTEMPT) + + async def aretrieve_batch(self, model: str | None = None, **kwargs: object) -> object: # kwargs-ok: the request + """Ask every deployment of the group (or of the router) for the batch.""" + from litellm.litellm_core_utils.core_helpers import ( + safe_deep_copy, # pyright: ignore[reportUnknownVariableType] # litellm's untyped helpers and request values + ) + from litellm.types.utils import LiteLLMBatch + + shaping: Final = self._core.shaping + metadata_variable_name: Final = get_router_metadata_variable_name(function_name="aretrieve_batch") + if model is not None: + specific: Final = kwargs.pop("specific_deployment", None) + healthy = await self._core.select( # rebind-ok: rebound as the Python router rebinds it + model, + dict(kwargs), + kind="healthy", + specific_deployment=bool(specific) or None, + messages=_user_message("retrieve-api-fake-text"), + ) + else: + healthy = [as_dict(r) for r in self._core.handle.deployments()] # rebind-ok: as Python rebinds it + errors: Final[list[BaseException]] = [] # mutable-ok: the request kwargs the Python router mutates in place + + async def try_retrieve(deployment: Mapping[str, object]) -> object: + try: + data: Final = as_dict(deployment.get("litellm_params")) + deployment_model: Final = data.get("model") + provider: object = data.get("custom_llm_provider") # rebind-ok: inferred when not declared + if deployment_model is None: + raise Exception(f"Model not found in litellm_params for deployment: {deployment}") # noqa: TRY002 # as the Python router + if not provider: + lookup: Final = self._core.host.invoke("get_llm_provider") + _, provider, _, _ = lookup(model=deployment_model) # rebind-ok: as Python rebinds it + new_kwargs: Final = safe_deep_copy(dict(kwargs)) # pyright: ignore[reportUnknownVariableType] # litellm's untyped helpers and request values + update_kwargs_with_deployment( + shaping, + deployment, + new_kwargs, # pyright: ignore[reportUnknownArgumentType] # litellm's untyped helpers and request values + cache=self._core.routing_cache, + function_name="aretrieve_batch", + ) + bucket: Final = metadata_bucket(new_kwargs, metadata_variable_name) # pyright: ignore[reportUnknownArgumentType] # litellm's untyped helpers and request values + if not bucket.get("model_group"): + bucket["model_group"] = model or deployment.get("model_name") + new_kwargs.pop("custom_llm_provider", None) # pyright: ignore[reportUnknownMemberType] # litellm's untyped helpers and request values + data.pop("custom_llm_provider", None) + return await self._direct("aretrieve_batch", {**data, "custom_llm_provider": provider, **new_kwargs}) # pyright: ignore[reportUnknownArgumentType] # litellm's untyped helpers and request values + except Exception as error: # noqa: BLE001 # collected; the first one is raised when no deployment has the batch + errors.append(error) + return None + + deployments: Final = list(healthy) if isinstance(healthy, Sequence) else [healthy] + if not deployments: + raise Exception("No healthy deployments found.") # noqa: TRY002 # as the Python router + results: Final = await asyncio.gather(*(try_retrieve(as_dict(d)) for d in deployments), return_exceptions=True) + for result in results: + if isinstance(result, LiteLLMBatch): + return result + if errors: + raise errors[0] + raise Exception(f"Unable to find batch in any model. Received errors - {errors}") # noqa: TRY002 # as the Python router + + async def alist_batches(self, model: str, **kwargs: object) -> object: # kwargs-ok: the request + deployments: Final = self._group(model) + + async def try_list(deployment: Mapping[str, object]) -> object: + try: + return await self._direct("alist_batches", {**as_dict(deployment.get("litellm_params")), **kwargs}) + except Exception: # noqa: BLE001 # a failing deployment is left out, as the Python router does + return None + + results: Final = await asyncio.gather(*(try_list(d) for d in deployments)) + final: Final[dict[str, object]] = { # mutable-ok: the request kwargs the Python router mutates in place + "object": "list", + "data": [], + "first_id": None, + "last_id": None, + "has_more": False, + } + data: Final[list[object]] = [] # mutable-ok: the request kwargs the Python router mutates in place + final["data"] = data + for result in results: + if result is None: + continue + if final["first_id"] is None and hasattr(result, "first_id"): + final["first_id"] = getattr(result, "first_id") # noqa: B009 # as the Python router + final["last_id"] = getattr(result, "last_id") # noqa: B009 # same + data.extend(getattr(result, "data")) # noqa: B009 # same # pyright: ignore[reportAny] # litellm's untyped helpers and request values + if getattr(result, "has_more", False) is True: + final["has_more"] = True + return final + + async def acreate_file(self, model: str, **kwargs: object) -> object: # kwargs-ok: the request + return await self._core.drive(PLANS["acreate_file"], {**kwargs, "model": model}) + + async def create_file_everywhere(self, kwargs: Mapping[str, object], requested: str, specific: object) -> object: + """`_acreate_file`: create the file on every healthy deployment of the group.""" + from litellm.litellm_core_utils import batch_utils + from litellm.litellm_core_utils.core_helpers import ( + safe_deep_copy, # pyright: ignore[reportUnknownVariableType] # litellm's untyped helpers and request values + ) + from litellm.litellm_core_utils.hidden_params import get_hidden_params + from litellm.litellm_core_utils.routing_common_utils import ( + add_model_file_id_mappings, # pyright: ignore[reportUnknownVariableType] # litellm's untyped helpers and request values + ) + from litellm.types.utils import LlmProviders + + healthy: Final = await self._core.select( + requested, + dict(kwargs), + kind="healthy", + specific_deployment=bool(specific) or None, + messages=_user_message(_PICK_TEXTS["create_file"]), + ) + deployments: Final = list(healthy) if isinstance(healthy, Sequence) else [healthy] + + async def create_for(deployment: Mapping[str, object]) -> object: + kwargs_copy: Final = safe_deep_copy(dict(kwargs)) # pyright: ignore[reportUnknownVariableType] # litellm's untyped helpers and request values + update_kwargs_with_deployment( + self._core.shaping, + deployment, + kwargs_copy, # pyright: ignore[reportUnknownArgumentType] # litellm's untyped helpers and request values + cache=self._core.routing_cache, + function_name="acreate_file", + ) + data: Final = as_dict(deployment.get("litellm_params")) + model_name: Final = str(data.get("model")) + self.counters.total_calls[model_name] += 1 + declared: Final = data.get("custom_llm_provider") + stripped, inferred, _, _ = self._core.host.invoke("get_llm_provider")( + model=data["model"], custom_llm_provider=declared + ) + provider: Final = declared or inferred + purpose: Final = kwargs.get("purpose") + file: object = kwargs.get("file") # rebind-ok: rewritten for the deployment + if not file or not purpose: + raise Exception("file and file_purpose are required for create_file") # noqa: TRY002 # as the Python router + replaces: Final = batch_utils.should_replace_model_in_jsonl( + purpose=purpose, # pyright: ignore[reportArgumentType] # the caller's purpose + passthrough=kwargs.get("passthrough") is True, + ) + if replaces: + file = batch_utils.replace_model_in_jsonl( # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] # an untyped helper # rebind-ok: rewritten for the deployment + file_content=file, # pyright: ignore[reportArgumentType] # the caller's file + new_model_name=stripped, + keep_line_numbers=provider == LlmProviders.VERTEX_AI.value, + ) + kwargs_copy["file"] = file + if provider == LlmProviders.LITELLM_PROXY.value: + kwargs_copy["extra_body"] = MappingProxyType( + {**as_dict(kwargs_copy.get("extra_body")), "target_model_names": stripped} # pyright: ignore[reportUnknownArgumentType, reportUnknownMemberType] # litellm's untyped helpers and request values + ) + if "gcs_bucket_name" in data: + metadata_bucket(kwargs_copy, "litellm_metadata")["gcs_bucket_name"] = data["gcs_bucket_name"] # pyright: ignore[reportUnknownArgumentType] # litellm's untyped helpers and request values + from litellm.router_rust.invoke.slots import deployment_slot + + with deployment_slot(self._core, deployment, kwargs_copy): # pyright: ignore[reportUnknownArgumentType] # litellm's untyped helpers and request values + response: Final = await self._direct( + "acreate_file", + { + **data, + "custom_llm_provider": provider, + "caching": self._cache_responses, + "client": model_client(self._core, PLANS["acreate_file"], deployment, kwargs_copy), # pyright: ignore[reportUnknownArgumentType] # litellm's untyped helpers and request values + **kwargs_copy, + }, # pyright: ignore[reportUnknownArgumentType] # litellm's untyped helpers and request values + ) + self.counters.success_calls[model_name] += 1 + return response + + responses: Final = await asyncio.gather(*(create_for(as_dict(d)) for d in deployments)) + if len(responses) == 0: + raise Exception("No healthy deployments found.") # noqa: TRY002 # as the Python router + mapping: Final = add_model_file_id_mappings(healthy_deployments=deployments, responses=responses) # pyright: ignore[reportArgumentType, reportUnknownVariableType] # deployment records + returned: Final = responses[0] + hidden: Final = get_hidden_params(returned) + if hidden is not None: + hidden["model_file_id_mapping"] = mapping + return returned + + # -- search tools ------------------------------------------------------------------------- + + async def asearch( + self, + custom_llm_provider: str | None = None, + client: object = None, + **kwargs: object, # kwargs-ok: the request + ) -> object: + """`_asearch_with_fallbacks`: the search tool name (``search_tool_name``, else ``model``) + routes as the model; the factory wrapper drops ``custom_llm_provider`` and ``client``.""" + request: Final = dict(kwargs) + search_tool_name: Final = request.get("search_tool_name", request.get("model")) + if not search_tool_name: + raise ValueError("search_tool_name or model parameter is required for search") + request["model"] = search_tool_name + return await self._core.drive(PLANS["asearch"], request) + + # -- guardrails --------------------------------------------------------------------------- + + async def aguardrail( + self, + guardrail_name: str, + original_function: Callable[..., Awaitable[object]], + **kwargs: object, # kwargs-ok: the request + ) -> object: + """Pick one of the guardrails named ``guardrail_name`` (weighted) and run it. A guardrail name + is no model group, so the router's retries never apply.""" + request: Final = dict(kwargs) + request.setdefault("litellm_trace_id", str(uuid.uuid4())) + metadata_bucket(request, "litellm_metadata").update( + { + "model_group": guardrail_name, + "model_group_alias": guardrail_name + if guardrail_name in self._core.resolved.model_group_alias + else None, + } + ) + verdict: Final = await original_function(**request, selected_guardrail=self.available_guardrail(guardrail_name)) + # `make_call`, `async_function_with_retries` and `async_function_with_fallbacks` stamp the result. + headed: Final = await self._core.set_response_headers(verdict, guardrail_name, request) + return add_fallback_headers_to_response( + response=add_retry_headers_to_response(response=headed, attempted_retries=0, max_retries=None), + attempted_fallbacks=0, + ) + + def available_guardrail(self, guardrail_name: str) -> object: + """``_core_get_available_guardrail``: the one guardrail configured under the name, else one of + them picked by weight (``simple_shuffle``, the name resolved through the alias map).""" + from litellm.router_strategy.simple_shuffle import simple_shuffle + + candidates: Final = [g for g in self._guardrails if as_dict(g).get("guardrail_name") == guardrail_name] + if not candidates: + raise ValueError(f"No guardrail found with name: {guardrail_name}") + if len(candidates) == 1: + return candidates[0] + return simple_shuffle( # pyright: ignore[reportUnknownVariableType] # litellm's untyped helper + resolve_model_alias=self._core.resolve_alias, + healthy_deployments=candidates, # pyright: ignore[reportArgumentType] # guardrail configs + model=guardrail_name, + request_kwargs=None, + ) + + # -- attempts ----------------------------------------------------------------------------- + + def attempt_failed( + self, + plan: Plan, + error: BaseException, + deployment: Mapping[str, object], + call: AttemptCall | None, + requested: str, + ) -> None: + from litellm import exceptions as litellm_exceptions + from litellm.types.router import LiteLLM_Params + + # `_acompletion`'s Timeout branch counts no failure and adds the deployment's timeouts; the + # sync `_completion` has no such branch. + timed_out: Final = plan.kind == "chat" and not plan.sync and isinstance(error, litellm_exceptions.Timeout) + if not timed_out: + self.counters.failed(plan, call.model_name if call is not None else None, requested) + elif isinstance(error, litellm_exceptions.Timeout) and self._core.host.dynamic().expose_router_debug_in_errors: + params: Final = as_dict(deployment.get("litellm_params")) + try: + typed: Mapping[str, object] = ( # rebind-ok: rebound as the Python router rebinds it + LiteLLM_Params.model_validate(params).model_dump() + ) + except Exception: # noqa: BLE001 # the params as configured, when they do not validate + typed = params # rebind-ok: rebound as the Python router rebinds it + error.message += ( # rebind-ok: the request kwargs are shaped in place, as the Python router shapes them + f"\n\nDeployment Info: request_timeout: {typed.get('request_timeout')}" + f"\ntimeout: {typed.get('timeout')}" + f"\nstream_timeout: {typed.get('stream_timeout')}" + ) + + def check_response(self, plan: Plan, response: object, call: AttemptCall) -> None: + """What the attempt helper checks of a response before it counts it as served.""" + + def sync_pre_call_checks(self, plan: Plan, deployment: Mapping[str, object], request: Mapping[str, object]) -> None: + """`routing_strategy_pre_call_checks` on the sync path: the process callbacks' pre-call checks. + Sync completion runs them for a group but not for a deployment id.""" + from litellm.integrations.custom_logger import CustomLogger + + if plan.kind == "chat": + model: Final = str(request.get("model")) + names: Final = {as_dict(r).get("model_name") for r in self._core.handle.deployments()} + if model not in names and self._core.handle.deployment(model) is not None: + return + for callback in tuple(self._core.routing_callbacks()): + if isinstance(callback, CustomLogger): + callback.pre_call_check(dict(deployment)) # pyright: ignore[reportUnknownMemberType] # litellm's untyped helpers and request values + + +_SINGLE_ATTEMPT: Final[Mapping[str, object]] = MappingProxyType({"num_retries": 0, "disable_fallbacks": True}) + + +async def _answers_none(*_args: object, **_kwargs: object) -> None: # kwargs-ok: any request + """An endpoint the Python router's async wrapper has no branch for.""" + return + + +def _as_items(value: object) -> list[object]: # mutable-ok: a fresh list + return list(value) if isinstance(value, list) else [] # pyright: ignore[reportUnknownArgumentType] # a config list + + +if TYPE_CHECKING: + from litellm.router_config import RouterConfig + from litellm.router_python_host import PythonHost + from litellm.router_rust.core import Resolved + + +__all__ = [ + "PLANS", + "ROUTING_KEYS", + "AttemptCall", + "Counters", + "Plan", + "RustEndpoints", + "Shaping", + "arm_compaction_state", + "attempt_call", + "clean_realtime_session_deployment", + "compacted_attempt", + "failover_fields", + "generic_plan", + "log_refused_call", + "log_refused_call_sync", + "pinned_video_deployment", + "plan_for", + "prepare_request", + "reports_own_success", + "selection_endpoint", + "stamp_error", + "unrouted_fields", +] diff --git a/litellm/router_rust/invoke/prompt_management.py b/litellm/router_rust/invoke/prompt_management.py new file mode 100644 index 00000000000..62bd77e1b48 --- /dev/null +++ b/litellm/router_rust/invoke/prompt_management.py @@ -0,0 +1,129 @@ +""" +Prompt-management groups on the Rust router core (``_is_prompt_management_model`` and +``_prompt_management_factory``): a group of exactly one deployment whose model names a prompt +management integration (``langfuse/...``) or ``litellm_agent/...``. The request is not routed with +retries and fallbacks: the deployment is picked, its prompt rendered by the request's logging +object, and the rendered model called directly, or routed when it names a group. +""" + +from collections.abc import Mapping +from typing import Final + +from litellm.litellm_core_utils.core_helpers import without +from litellm.router_rust.invoke.endpoints import PLANS, EndpointCore +from litellm.router_rust.invoke.shaping import as_dict, update_kwargs_with_deployment + +_PROMPT_MANAGEMENT_PARAMS: Final = frozenset( + {"bitbucket_config", "dotprompt_config", "prompt_id", "prompt_variables", "prompt_label", "prompt_version"} +) + + +def is_prompt_management_model(core: EndpointCore, model: str) -> bool: + """`_is_prompt_management_model`: one deployment, whose model's provider prefix is a known + custom-logger-compatible callback.""" + model_list: Final = core.group_deployments(model_name=model) + if model_list is None or len(model_list) != 1: + return False + litellm_model: Final = as_dict(model_list[0].get("litellm_params")).get("model") + if not isinstance(litellm_model, str) or "/" not in litellm_model: + return False + known: Final = core.host.dynamic().known_custom_logger_compatible_callbacks + return litellm_model.split("/")[0] in known + + +def _logging_and_kwargs( + kwargs: dict[str, object], # mutable-ok: the request kwargs the Python router mutates in place +) -> tuple[object, dict[str, object]]: # mutable-ok: same + """The request's logging object, else the one `function_setup` builds (with the kwargs it returns).""" + from litellm.utils import Rules, function_setup, get_utc_datetime + + logging_obj: Final = kwargs.get("litellm_logging_obj") + if logging_obj is not None: + return logging_obj, kwargs + return function_setup( # pyright: ignore[reportUnknownVariableType] # litellm's untyped helpers + **{ # pyright: ignore[reportArgumentType] # the request's keyword payload + "original_function": "acompletion", + "rules_obj": Rules(), + "start_time": get_utc_datetime(), + **kwargs, + } + ) + + +async def prompt_management_acompletion( + core: EndpointCore, + model: str, + messages: object, + request: Mapping[str, object], +) -> object: + """`_prompt_management_factory`, on the request as `acompletion` prepared it.""" + from litellm.utils import get_non_default_completion_params + + # `acompletion` hands the routing function along; function_setup consumes it. + logging_obj, kwargs = _logging_and_kwargs({**request, "original_function": core.router.acompletion}) + caller_timeout: Final = kwargs.get("timeout") or kwargs.get("request_timeout") + caller_stream_timeout: Final = kwargs.get("stream_timeout") + specific_deployment: Final = kwargs.pop("specific_deployment", None) + picked: Final = await core.select( + model, + kwargs, # pyright: ignore[reportUnknownArgumentType] # the request kwargs + kind="pick", + specific_deployment=bool(specific_deployment), + messages=messages, + ) + deployment: Final = as_dict(picked) + update_kwargs_with_deployment( + core.shaping, + deployment, + kwargs, # pyright: ignore[reportUnknownArgumentType] # the request kwargs + cache=core.routing_cache, + splits_stream_timeout=True, + ) + params: Final = as_dict(deployment.get("litellm_params")) + data: Final = dict(params) + litellm_model: Final = data.get("model") + is_litellm_agent_model: Final = isinstance(litellm_model, str) and litellm_model.startswith("litellm_agent/") + prompt_id: Final = kwargs.get("prompt_id") or params.get("prompt_id") + prompt_variables: Final = kwargs.get("prompt_variables") or params.get("prompt_variables") + prompt_label: Final = kwargs.get("prompt_label") or params.get("prompt_label") + if not is_litellm_agent_model and (prompt_id is None or not isinstance(prompt_id, str)): + raise ValueError(f"Prompt ID is not set or not a string. Got={prompt_id}, type={type(prompt_id)}") + if prompt_variables is not None and not isinstance(prompt_variables, dict): + raise ValueError( + f"Prompt variables is set but not a dictionary. Got={prompt_variables}, type={type(prompt_variables)}" + ) + rendered_model, rendered_messages, optional_params = logging_obj.get_chat_completion_prompt( # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType, reportAttributeAccessIssue, reportOptionalMemberAccess] # the request's Logging object + model=litellm_model, # pyright: ignore[reportArgumentType] # a request value, passed as the Python router passes it + messages=messages, # pyright: ignore[reportArgumentType] # a request value, passed as the Python router passes it + non_default_params=get_non_default_completion_params(kwargs=kwargs), # pyright: ignore[reportArgumentType] # the request kwargs + prompt_id=prompt_id, # pyright: ignore[reportArgumentType] # a request value, passed as the Python router passes it + prompt_variables=prompt_variables, + prompt_label=prompt_label, # pyright: ignore[reportArgumentType] # a request value, passed as the Python router passes it + request_kwargs=kwargs, + injected_for_every_deployment=True, + ) + filtered_data: Final = without(data, *_PROMPT_MANAGEMENT_PARAMS) + final: Final[dict[str, object]] = {**filtered_data, **kwargs, **optional_params} # pyright: ignore[reportUnknownVariableType] # mutable-ok: the call's kwargs + final["model"] = rendered_model + final["messages"] = rendered_messages + final["litellm_logging_obj"] = logging_obj + final["prompt_id"] = prompt_id + final["prompt_variables"] = prompt_variables + final["prompt_label"] = prompt_label + target_list: Final = core.group_deployments(model_name=str(rendered_model)) # pyright: ignore[reportUnknownArgumentType] # the core's host registry + if target_list is None or len(target_list) == 0: # a direct call to the rendered model + _ = final.pop("original_function") + return await core.host.invoke("acompletion")(**final) # pyright: ignore[reportUnknownMemberType] # the core's host + timeouts: Final = { + key: value + for key, value in ( + ("timeout", caller_timeout if caller_timeout is not None else filtered_data.get("timeout")), + ( + "stream_timeout", + caller_stream_timeout if caller_stream_timeout is not None else filtered_data.get("stream_timeout"), + ), + ) + if value is not None + } + routed: Final = without(final, "timeout", "stream_timeout") + return await core.drive(PLANS["acompletion"], {**routed, **timeouts}, prepared_already=True) diff --git a/litellm/router_rust/invoke/shaping.py b/litellm/router_rust/invoke/shaping.py new file mode 100644 index 00000000000..258f919c9f1 --- /dev/null +++ b/litellm/router_rust/invoke/shaping.py @@ -0,0 +1,635 @@ +""" +The keyword arguments one router attempt sends to one deployment, as the Python router builds them. + +The Rust router core picks the deployment; the host shapes the request for it here +(``_update_kwargs_with_deployment``, ``_update_kwargs_with_default_litellm_params``, +``_merge_tools_from_deployment``, the timeout resolvers and the reasoning-effort override), with +the same precedence and the same oddities. Nothing here reads a Python router. +""" + +import dataclasses +from collections.abc import Mapping, MutableMapping, Sequence +from functools import partial +from types import MappingProxyType +from typing import TYPE_CHECKING, Final + +from pydantic import TypeAdapter, ValidationError + +from litellm import exceptions as litellm_exceptions +from litellm.constants import ( + BAGGAGE_MAPPED_TAGS_METADATA_KEY, + INTERNAL_CALL_ORIGIN_METADATA_KEY, + KUBERNETES_POD_ROUTING_KEY, + ROUTING_REQUEST_TAGS_METADATA_KEY, +) +from litellm.litellm_core_utils.batch_utils import get_router_metadata_variable_name +from litellm.litellm_core_utils.clientside_credential_handler import ( + get_dynamic_litellm_params, # pyright: ignore[reportUnknownVariableType] # litellm's untyped helpers and request values + is_clientside_credential, # pyright: ignore[reportUnknownVariableType] # litellm's untyped helpers and request values +) +from litellm.litellm_core_utils.completion_timeout import CompletionTimeout +from litellm.litellm_core_utils.core_helpers import without +from litellm.litellm_core_utils.get_litellm_params import SUPPLEMENTAL_KWARGS_KEYS +from litellm.litellm_core_utils.get_llm_provider_logic import declared_authenticating_provider, inferred_provider +from litellm.litellm_core_utils.request_timeout_resolver import get_service_tier_timeout +from litellm.litellm_core_utils.routing_common_utils import deployment_model_id_pin +from litellm.litellm_core_utils.tool_fields import TOOL_FIELDS, drop_policy +from litellm.responses.retrieve_stream_query import requests_streamed_retrieval, without_retrieve_stream_defaults +from litellm.router_api import deployment_id as generate_deployment_id +from litellm.router_strategy.tag_based_routing import ( + baggage_stamp_for_serving_deployment, + get_tags_from_request_kwargs, + tags_with_serving_deployment_tags, +) +from litellm.router_utils.pre_call_checks.io_token_rate_limit_check import ( + deployment_has_io_token_limits, # pyright: ignore[reportUnknownVariableType] # litellm's untyped helpers and request values + refund_stale_reservation_before_retry, + set_io_token_rate_limit_request_kwargs, +) +from litellm.router_utils.stream_usage_injection import router_config_requests_stream_usage +from litellm.types.litellm_params import STREAM_USAGE_INJECTED_FIELD +from litellm.types.router import Deployment, GenericLiteLLMParams, LiteLLM_Params, validate_litellm_params_strictly +from litellm.types.utils import AUTOROUTER_CLASSIFIER_CALL_ORIGIN + +if TYPE_CHECKING: + from litellm.caching.dual_cache import DualCache + from litellm.router_python_host import PythonHost + +# The SDK-client and max_parallel_requests reads of the routing cache: process-local, so the family +# only names a Redis span when a read misses through to Redis. +ROUTER_CLIENTS_TARGET: Final = "router_clients" +_MAPPING: Final = TypeAdapter(Mapping[str, object]) +_FORWARDED_CREDENTIAL_ISOLATION_PROVIDERS: Final = frozenset({"chatgpt", "anthropic"}) +_ALIAS_MARKER_FORWARDED_PARAMS_KWARG: Final = "_alias_marker_forwarded_params" +# Generic handlers whose stream timeout is split from the request timeout. +STREAM_TIMEOUT_AWARE_GENERIC_HANDLERS: Final = frozenset({"anthropic_messages", "aresponses", "responses"}) + + +@dataclasses.dataclass(frozen=True, slots=True) +class Shaping: + """What the request shaping reads from the router, resolved once when the core is built. + + ``timeout`` is the router's ``timeout`` (its own setting, else the process ``request_timeout``), + ``request_timeout`` the explicitly configured process ``request_timeout``, kept only when the + router has a timeout of its own, as the Python router keeps them.""" + + default_litellm_params: Mapping[str, object] + timeout: object + explicit_timeout: object + request_timeout: float | None + stream_timeout: object + cache_responses: bool + + +def as_dict(value: object) -> dict[str, object]: # mutable-ok: a fresh dict of a request mapping + try: + return dict(_MAPPING.validate_python(value)) + except ValidationError: + return {} + + +def _mapping(value: object) -> Mapping[str, object] | None: + try: + return _MAPPING.validate_python(value) + except ValidationError: + return None + + +def metadata_bucket( + kwargs: MutableMapping[str, object], # mutable-ok: the request kwargs the Python router mutates in place + name: str, +) -> dict[str, object]: # mutable-ok: the request's own bucket + """``kwargs.setdefault(name, {})``: the request's own metadata dict, created when absent.""" + bucket: Final = kwargs.get(name) + if isinstance(bucket, dict): + return bucket # pyright: ignore[reportUnknownVariableType] # the request's untyped metadata dict + if name not in kwargs: + fresh: Final[dict[str, object]] = {} # mutable-ok: the request kwargs the Python router mutates in place + kwargs[name] = fresh # rebind-ok: the request kwargs are shaped in place, as the Python router shapes them + return fresh + # A non-dict value under the key: `setdefault` returns it as is, and `.update` would raise. + raise AttributeError(f"'{type(bucket).__name__}' object has no attribute 'update'") + + +# --- request params the deployment sets -------------------------------------------------------- + + +def _drops_malformed_request_tools( + deployment_params: Mapping[str, object], defaults: Mapping[str, object], request: Mapping[str, object] +) -> bool: + if isinstance(request.get("tools") or [], (list, tuple)) or not deployment_params.get("tools"): + return False + set_defaults: Final = {name: value for name, value in defaults.items() if value is not None} + return drop_policy({**deployment_params, **set_defaults, **request}).drops("tools") + + +def _malformed_request_params_the_deployment_sets( + deployment_params: Mapping[str, object], defaults: Mapping[str, object], request: Mapping[str, object] +) -> frozenset[str]: + overrides: Final = { + key: value + for key, value in request.items() + if key in SUPPLEMENTAL_KWARGS_KEYS and deployment_params.get(key) is not None + } + if not overrides: + return frozenset() + try: + validate_litellm_params_strictly(GenericLiteLLMParams, overrides) + except ValidationError as error: + set_defaults: Final = {name: value for name, value in defaults.items() if value is not None} + policy: Final = drop_policy({**deployment_params, **set_defaults, **request}) + malformed: Final = (str(detail["loc"][0]) for detail in error.errors() if detail["loc"]) + return frozenset(key for key in malformed if policy.drops(key)) + return frozenset() + + +def _merged_tools( + deployment_tools: Sequence[object], request_tools: object +) -> list[object] | None: # mutable-ok: sent as the request's tools + if not isinstance(request_tools, (list, tuple)): + return None + items: Final[Sequence[object]] = request_tools # pyright: ignore[reportUnknownVariableType] # a request's untyped tools + return [*deployment_tools, *items] if deployment_tools or items else None + + +def merge_tools_from_deployment( + deployment_params: Mapping[str, object], + kwargs: MutableMapping[str, object], # mutable-ok: the request kwargs the Python router mutates in place + defaults: Mapping[str, object], +) -> None: + """``_merge_tools_from_deployment``: deployment tools first, then the request's; the request's + ``tool_choice`` else the deployment's.""" + if _drops_malformed_request_tools(deployment_params, defaults, kwargs): + for name in TOOL_FIELDS: + kwargs.pop(name, None) + deployment_tools: Final = deployment_params.get("tools") or [] # pyright: ignore[reportUnknownVariableType] # litellm's untyped helpers and request values + merged: Final = _merged_tools( + deployment_tools if isinstance(deployment_tools, list) else [], # pyright: ignore[reportUnknownArgumentType] # the deployment's untyped tools + kwargs.get("tools") or [], + ) + if merged is not None: + kwargs["tools"] = merged # rebind-ok: the request kwargs are shaped in place, as the Python router shapes them + if "tool_choice" not in kwargs and deployment_params.get("tool_choice") is not None: + kwargs["tool_choice"] = ( # rebind-ok: the request kwargs are shaped in place, as the Python router shapes them + deployment_params["tool_choice"] + ) + + +def _deployment_sets_litellm_param(deployment_params: Mapping[str, object], key: str) -> bool: + value: Final = deployment_params.get(key) + if value is None: + return False + field: Final = LiteLLM_Params.model_fields.get(key) + return field is None or value != field.default # pyright: ignore[reportAny] # litellm's untyped helpers and request values + + +def _forwarded_alias_marker_keys_the_deployment_sets( + deployment_params: Mapping[str, object], forwarded_keys: object +) -> tuple[str, ...]: + if not isinstance(forwarded_keys, tuple): + return () + keys: Final[tuple[object, ...]] = forwarded_keys # pyright: ignore[reportUnknownVariableType] # the router's own marker + return tuple(key for key in keys if isinstance(key, str) and _deployment_sets_litellm_param(deployment_params, key)) + + +# --- reasoning effort ------------------------------------------------------------------------- + + +def _pop_effort_from_nested_carrier(params: dict[str, object], carrier: str) -> None: # mutable-ok: copy-on-write + nested: Final = _mapping(params.get(carrier)) if isinstance(params.get(carrier), Mapping) else None + if nested is None: + return + sanitized: Final = without(nested, "effort") + if sanitized: + params[carrier] = ( # rebind-ok: the request kwargs are shaped in place, as the Python router shapes them + sanitized + ) + else: + params.pop(carrier, None) + + +def deployment_params_with_request_reasoning_override( + deployment_params: Mapping[str, object], request_kwargs: Mapping[str, object] +) -> dict[str, object]: # mutable-ok: the call's own copy of the deployment params + """``_deployment_params_with_request_reasoning_override``: a request ``reasoning_effort`` drops + the deployment's ``thinking`` and nested ``effort`` carriers.""" + sanitized: Final = dict(deployment_params) + if request_kwargs.get("reasoning_effort") is None: + return sanitized + sanitized.pop("thinking", None) + _pop_effort_from_nested_carrier(sanitized, "output_config") + _pop_effort_from_nested_carrier(sanitized, "reasoning") + extra_body: Final = sanitized.get("extra_body") + if isinstance(extra_body, Mapping): + sanitized_extra_body: Final = as_dict(extra_body) # pyright: ignore[reportUnknownArgumentType] # litellm's untyped helpers and request values + sanitized_extra_body.pop("reasoning_effort", None) + sanitized_extra_body.pop("thinking", None) + _pop_effort_from_nested_carrier(sanitized_extra_body, "output_config") + _pop_effort_from_nested_carrier(sanitized_extra_body, "reasoning") + if sanitized_extra_body: + sanitized["extra_body"] = sanitized_extra_body + else: + sanitized.pop("extra_body", None) + return sanitized + + +def _is_classifier_internal_call(kwargs: Mapping[str, object]) -> bool: + return any( + isinstance(candidate, Mapping) + and (_mapping(candidate) or {}).get(INTERNAL_CALL_ORIGIN_METADATA_KEY) == AUTOROUTER_CLASSIFIER_CALL_ORIGIN # pyright: ignore[reportUnknownArgumentType] # litellm's untyped helpers and request values + for candidate in (kwargs.get("metadata"), kwargs.get("litellm_metadata")) + ) + + +def _declared_param_allowlist(params: Mapping[str, object]) -> frozenset[str]: + declared: Final = params.get("allowed_openai_params") + if not isinstance(declared, (list, tuple, set, frozenset)): + return frozenset() + entries: Final[Sequence[object]] = list(declared) # pyright: ignore[reportUnknownArgumentType] # the deployment's untyped allowlist + return frozenset(entry for entry in entries if isinstance(entry, str)) + + +def _deployment_supports_param(deployment: Mapping[str, object], group: str, param: str, host: "PythonHost") -> bool: + params: Final = as_dict(deployment.get("litellm_params")) + if not params: + return True + model_info: Final = as_dict(deployment.get("model_info")) + base_model: Final = model_info.get("base_model") or params.get("base_model") + try: + model_label: Final = str(params.get("model") or group) + provider_param: Final = params.get("custom_llm_provider") + authenticating: Final = declared_authenticating_provider( + model_label, provider_param if isinstance(provider_param, str) else None + ) + if authenticating is not None: + model, custom_llm_provider = model_label.removeprefix(f"{authenticating}/"), authenticating + else: + model, custom_llm_provider, _, _ = host.invoke( # rebind-ok: rebound as the Python router rebinds it + "get_llm_provider" + )( + model=model_label, + custom_llm_provider=provider_param, + litellm_params=LiteLLM_Params.model_validate(params), + ) + supported: Final = host.invoke("get_supported_openai_params")( + model=model, + custom_llm_provider=custom_llm_provider, + base_model=base_model if isinstance(base_model, str) else None, + ) + except Exception: # noqa: BLE001 # best-effort filter: an unresolvable provider must not narrow the request + return True + return supported is None or param in supported + + +def deployment_accepts_param(deployment: Mapping[str, object], group: str, param: str, host: "PythonHost") -> bool: + params: Final = as_dict(deployment.get("litellm_params")) + if not params: + return True + if param in _declared_param_allowlist(params): + return True + provider_param: Final = params.get("custom_llm_provider") + if declared_authenticating_provider( + str(params.get("model") or ""), provider_param if isinstance(provider_param, str) else None + ): + return True + return _deployment_supports_param(deployment, group, param, host) + + +def drop_unsupported_classifier_reasoning_effort( + deployment: Mapping[str, object], + model: str, + kwargs: MutableMapping[str, object], # mutable-ok: the request kwargs the Python router mutates in place + host: "PythonHost", +) -> None: + """``_drop_unsupported_classifier_reasoning_effort``.""" + if kwargs.get("reasoning_effort") is None or not _is_classifier_internal_call(kwargs): + return + if deployment_accepts_param(deployment, model, "reasoning_effort", host): + return + kwargs.pop("reasoning_effort", None) + proxy_server_request: Final = kwargs.get("proxy_server_request") + if not isinstance(proxy_server_request, dict): + return + body: Final[object] = proxy_server_request.get("body") # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] # the proxy's untyped request + if isinstance(body, dict): + body.pop("reasoning_effort", None) # pyright: ignore[reportUnknownMemberType] # same + + +# --- timeouts --------------------------------------------------------------------------------- + + +def stream_timeout(shaping: Shaping, kwargs: Mapping[str, object], data: Mapping[str, object]) -> object: + """``_get_stream_timeout``.""" + return ( + kwargs.get("stream_timeout", None) + or get_service_tier_timeout(kwargs.get("service_tier"), data, "stream_timeout_by_service_tier") + or data.get("stream_timeout", None) + or shaping.stream_timeout + or shaping.request_timeout + or shaping.default_litellm_params.get("stream_timeout", None) + ) + + +def non_stream_timeout(shaping: Shaping, kwargs: Mapping[str, object], data: Mapping[str, object]) -> object: + """``_get_non_stream_timeout``.""" + return ( + kwargs.get("timeout", None) + or kwargs.get("request_timeout", None) + or get_service_tier_timeout(kwargs.get("service_tier"), data, "timeout_by_service_tier") + or data.get("timeout", None) + or data.get("request_timeout", None) + or shaping.request_timeout + or shaping.timeout + or shaping.default_litellm_params.get("timeout", None) + ) + + +def call_timeout(shaping: Shaping, kwargs: Mapping[str, object], data: Mapping[str, object]) -> object: + """``_get_timeout``.""" + stream: Final = stream_timeout(shaping, kwargs, data) if kwargs.get("stream") else None + return stream if stream is not None else non_stream_timeout(shaping, kwargs, data) + + +# --- default litellm params ------------------------------------------------------------------- + + +def update_kwargs_with_default_litellm_params( + shaping: Shaping, + kwargs: MutableMapping[str, object], # mutable-ok: the request kwargs the Python router mutates in place + metadata_variable_name: str, +) -> None: + """``_update_kwargs_with_default_litellm_params``: unset non-None defaults fill the request; + the default metadata is merged over the request's.""" + defaults: Final = dict(without_retrieve_stream_defaults(shaping.default_litellm_params, kwargs)) + metadata_defaults: Final = defaults.pop("metadata", {}) or {} # pyright: ignore[reportUnknownVariableType] # litellm's untyped helpers and request values + for key, value in defaults.items(): + if value is None: + continue + kwargs.setdefault(key, value) + metadata_bucket(kwargs, metadata_variable_name).update(as_dict(metadata_defaults)) # pyright: ignore[reportUnknownArgumentType] # litellm's untyped helpers and request values + + +# --- clientside credentials ------------------------------------------------------------------- + + +def _clientside_deployment( + deployment: Mapping[str, object], kwargs: Mapping[str, object], function_name: str | None +) -> tuple[Mapping[str, object], str, object]: + """``_handle_clientside_credential``: the request-scoped model_info, model and api_base.""" + model_info: Final = as_dict(deployment.get("model_info")) + litellm_params: Final = as_dict(deployment.get("litellm_params")) + dynamic_litellm_params: Final = get_dynamic_litellm_params( # pyright: ignore[reportUnknownVariableType] # litellm's untyped helpers and request values + litellm_params=litellm_params, request_kwargs=dict(kwargs) + ) + metadata_variable_name: Final = get_router_metadata_variable_name(function_name=function_name) + model_group: Final = as_dict(kwargs.get(metadata_variable_name)).get("model_group") + model_id: Final = generate_deployment_id(model_group, dynamic_litellm_params) # pyright: ignore[reportArgumentType, reportUnknownArgumentType] # the group as the request carries it + original_model_id: Final = model_info.get("id") + model_info["id"] = model_id + model_info["original_model_id"] = original_model_id + built: Final = Deployment( + model_name=model_group, # pyright: ignore[reportArgumentType] # the group as the request carries it + litellm_params=LiteLLM_Params.model_validate(dynamic_litellm_params), + model_info=model_info, # pyright: ignore[reportArgumentType] # model_info as configured + ) + return built.model_info.model_dump(), built.litellm_params.model, built.litellm_params.api_base + + +# --- the deployment -------------------------------------------------------------------------- + + +def update_kwargs_with_deployment( + shaping: Shaping, + deployment: Mapping[str, object], + kwargs: MutableMapping[str, object], # mutable-ok: the request kwargs the Python router mutates in place + *, + cache: "DualCache | None", + function_name: str | None = None, + splits_stream_timeout: bool = False, + opens_responses_socket: bool = False, +) -> None: + """``_update_kwargs_with_deployment``, in place on the attempt's kwargs.""" + deployment_params: Final = as_dict(deployment.get("litellm_params")) + for key in _forwarded_alias_marker_keys_the_deployment_sets( + deployment_params, kwargs.pop(_ALIAS_MARKER_FORWARDED_PARAMS_KWARG, ()) + ): + kwargs.pop(key, None) + kwargs.pop("daily_reset_time", None) + defaults: Final = shaping.default_litellm_params + merge_tools_from_deployment(deployment_params, kwargs, defaults) + + model_info: Mapping[str, object] = as_dict( # rebind-ok: rebound as the Python router rebinds it + deployment.get("model_info") + ) + deployment_model: object = deployment_params.get("model") # rebind-ok: same + deployment_api_base: object = deployment_params.get("api_base") # rebind-ok: same + deployment_model_name: Final = deployment.get("model_name") + from litellm.litellm_core_utils.get_provider_specific_headers import ProviderSpecificHeaderUtils + + for key in _malformed_request_params_the_deployment_sets(deployment_params, defaults, kwargs): + kwargs.pop(key) + declared_provider: Final = deployment_params.get("custom_llm_provider") + declared_model: Final = deployment_params.get("model") + deployment_provider: Final = ( + declared_provider + if isinstance(declared_provider, str) + else inferred_provider(declared_model if isinstance(declared_model, str) else None) + ) + forwarded_authorization: Final = ( + ProviderSpecificHeaderUtils.forwarded_authorization(kwargs.get("provider_specific_header"), deployment_provider) + if deployment_provider in _FORWARDED_CREDENTIAL_ISOLATION_PROVIDERS + else None + ) + if is_clientside_credential(request_kwargs=dict(kwargs)) or forwarded_authorization is not None: + model_info, deployment_model, deployment_api_base = ( # rebind-ok: rebound as the Python router rebinds it + _clientside_deployment( + deployment, + {**kwargs, **({"api_key": forwarded_authorization} if forwarded_authorization is not None else {})}, + function_name, + ) + ) + + metadata_variable_name: Final = get_router_metadata_variable_name(function_name=function_name) + routing: Final = deployment.get(KUBERNETES_POD_ROUTING_KEY) + metadata: Final = metadata_bucket(kwargs, metadata_variable_name) + metadata.update( + { + "deployment": deployment_model, + "model_info": model_info, + "api_base": deployment_api_base, + "deployment_model_name": deployment_model_name, + KUBERNETES_POD_ROUTING_KEY: routing, + } + ) + other_bucket: Final = "metadata" if metadata_variable_name == "litellm_metadata" else "litellm_metadata" + other_bucket_metadata: Final = kwargs.get(other_bucket) + if isinstance(other_bucket_metadata, Mapping) and KUBERNETES_POD_ROUTING_KEY in other_bucket_metadata: + kwargs[other_bucket] = {**as_dict(other_bucket_metadata), KUBERNETES_POD_ROUTING_KEY: routing} # pyright: ignore[reportUnknownArgumentType] # litellm's untyped helpers and request values # rebind-ok: the request kwargs are shaped in place, as the Python router shapes them + + if cache is not None: + refund_stale_reservation_before_retry(cache, kwargs) # pyright: ignore[reportArgumentType] # the request's kwargs dict + set_io_token_rate_limit_request_kwargs( + kwargs, # pyright: ignore[reportArgumentType] # the request's kwargs dict + store_in_context=deployment_has_io_token_limits(deployment), # pyright: ignore[reportArgumentType] # a deployment record + opens_responses_socket=opens_responses_socket, + metadata_channel=metadata_variable_name, + ) + metadata.setdefault( + ROUTING_REQUEST_TAGS_METADATA_KEY, + tuple(get_tags_from_request_kwargs(kwargs, metadata_variable_name=metadata_variable_name)), # pyright: ignore[reportArgumentType] # the request's kwargs dict + ) + deployment_tags: Final = deployment_params.get("tags") + merged_tags: Final = tags_with_serving_deployment_tags(metadata, deployment_tags) + if merged_tags is not None: + metadata["tags"] = list(merged_tags) + baggage_stamp: Final = baggage_stamp_for_serving_deployment(metadata, deployment_tags) + if baggage_stamp is not None: + metadata[BAGGAGE_MAPPED_TAGS_METADATA_KEY] = baggage_stamp + + credential_name: Final = deployment_params.get("litellm_credential_name") + if credential_name: + credential_tag: Final = f"Credential: {credential_name}" + existing_tags: Final = metadata.get("tags") or [] # pyright: ignore[reportUnknownVariableType] # litellm's untyped helpers and request values + if isinstance(existing_tags, list) and credential_tag not in existing_tags: + existing_tags.append(credential_tag) # pyright: ignore[reportUnknownMemberType] # the request's own tags list + metadata["tags"] = existing_tags + + from litellm.llms.github_copilot.per_user_auth import github_copilot_per_user_credential_name + + configured_per_user_name: Final = github_copilot_per_user_credential_name(deployment_params) + caller_credential_name: Final = kwargs.get("litellm_credential_name") + if ( + configured_per_user_name is not None + and isinstance(caller_credential_name, str) + and caller_credential_name != configured_per_user_name + ): + raise litellm_exceptions.BadRequestError( + message="litellm_credential_name cannot be overridden on a deployment that uses per-user GitHub OAuth", + model=str(deployment_model_name), + llm_provider="", + ) + + kwargs["model_info"] = ( # rebind-ok: the request kwargs are shaped in place, as the Python router shapes them + model_info + ) + _set_timeouts(shaping, deployment_params, kwargs, function_name, splits_stream_timeout) + + if router_config_requests_stream_usage(kwargs, (deployment_params, defaults)): + kwargs[ # rebind-ok: the request kwargs are shaped in place, as the Python router shapes them + STREAM_USAGE_INJECTED_FIELD + ] = True + update_kwargs_with_default_litellm_params(shaping, kwargs, metadata_variable_name) + pinned_model_id: Final = deployment_model_id_pin(deployment_params) + if pinned_model_id is not None: + kwargs["model_id"] = ( # rebind-ok: the request kwargs are shaped in place, as the Python router shapes them + pinned_model_id + ) + + +def _set_timeouts( + shaping: Shaping, + deployment_params: Mapping[str, object], + kwargs: MutableMapping[str, object], # mutable-ok: the request kwargs the Python router mutates in place + function_name: str | None, + splits_stream_timeout: bool, +) -> None: + if function_name == "_ageneric_api_call_with_fallbacks": + from litellm.passthrough.timeout_utils import resolve_llm_passthrough_timeout + + router_timeout: Final = ( + shaping.request_timeout + if shaping.request_timeout is not None + else float(shaping.explicit_timeout) + if isinstance(shaping.explicit_timeout, (int, float)) + else None + ) + router_stream_timeout: Final = ( + shaping.stream_timeout + if shaping.stream_timeout is not None + else shaping.request_timeout + if shaping.request_timeout is not None + else shaping.default_litellm_params.get("stream_timeout") + ) + passthrough_timeout: Final = partial( + resolve_llm_passthrough_timeout, + litellm_params=deployment_params, + router_timeout=router_timeout, + router_stream_timeout=router_stream_timeout, # pyright: ignore[reportArgumentType] # a configured default, passed as the Python router passes it + ) + timeout_request: Final = {**kwargs, "stream": True} if requests_streamed_retrieval(kwargs) else dict(kwargs) + kwargs.update( + CompletionTimeout.stream_request_kwargs( + passthrough_timeout(kwargs={**kwargs, "stream": False}), passthrough_timeout(kwargs=timeout_request) + ) + if splits_stream_timeout and timeout_request.get("stream") + else {"timeout": passthrough_timeout(kwargs=timeout_request)} + ) + elif splits_stream_timeout: + kwargs.update( + CompletionTimeout.stream_request_kwargs( + non_stream_timeout(shaping, kwargs, deployment_params), + stream_timeout(shaping, kwargs, deployment_params) if kwargs.get("stream") else None, + ) + ) + else: + kwargs["timeout"] = ( # rebind-ok: the request kwargs are shaped in place, as the Python router shapes them + call_timeout(shaping, kwargs, deployment_params) + ) + + +def stamp_deployment_num_retries(error: BaseException, deployment: Mapping[str, object]) -> None: + """``_set_deployment_num_retries_on_exception``.""" + if getattr(error, "num_retries", None) is not None: + return + value: Final = as_dict(deployment.get("litellm_params")).get("num_retries") + if value is None: + return + try: + error.num_retries = int(str(value)) if not isinstance(value, int) else value # pyright: ignore[reportAttributeAccessIssue] # the router stamps the same attribute # rebind-ok: the request kwargs are shaped in place, as the Python router shapes them + except (ValueError, TypeError): + pass + + +def stamp_failed_deployment_id(error: BaseException, model_info: object) -> None: + """``_set_failed_deployment_id_on_exception`` with the effective model_info.""" + if getattr(error, "failed_deployment_id", None): + return + deployment_id: Final = as_dict(model_info).get("id") + if deployment_id: + try: + error.failed_deployment_id = deployment_id # pyright: ignore[reportAttributeAccessIssue] # the router stamps the same attribute # rebind-ok: the request kwargs are shaped in place, as the Python router shapes them + except Exception: # noqa: BLE001, S110 # as the Python router: a read-only error keeps its attributes + pass + + +def stamp_retry_skip_deployment_id(error: BaseException, kwargs: Mapping[str, object]) -> None: + """``_stamp_retry_skip_deployment_id``.""" + model_info: Final = kwargs.get("model_info") + if not isinstance(model_info, Mapping): + return + info: Final = as_dict(model_info) # pyright: ignore[reportUnknownArgumentType] # litellm's untyped helpers and request values + configured: Final = ( + info.get("original_model_id") if isinstance(error, litellm_exceptions.UnsupportedEndpointError) else None + ) + deployment_id: Final = configured or info.get("id") + if isinstance(deployment_id, str) and deployment_id: + error.retry_skip_deployment_id = deployment_id # pyright: ignore[reportAttributeAccessIssue] # dynamic stamp, read on retry # rebind-ok: the request kwargs are shaped in place, as the Python router shapes them + + +NO_PARAMS: Final[Mapping[str, object]] = MappingProxyType({}) + +__all__ = [ + "STREAM_TIMEOUT_AWARE_GENERIC_HANDLERS", + "Shaping", + "as_dict", + "call_timeout", + "deployment_params_with_request_reasoning_override", + "drop_unsupported_classifier_reasoning_effort", + "merge_tools_from_deployment", + "metadata_bucket", + "stamp_deployment_num_retries", + "stamp_failed_deployment_id", + "stamp_retry_skip_deployment_id", + "update_kwargs_with_default_litellm_params", + "update_kwargs_with_deployment", +] diff --git a/litellm/router_rust/invoke/slots.py b/litellm/router_rust/invoke/slots.py new file mode 100644 index 00000000000..08e88b5b2f1 --- /dev/null +++ b/litellm/router_rust/invoke/slots.py @@ -0,0 +1,107 @@ +""" +A deployment's ``max_parallel_requests`` slot on the Rust router core (``_deployment_slot`` and +``_get_client(client_type="max_parallel_requests")``): the attempt holds one of the deployment's +slots around its provider call, and an open chat stream keeps it until the stream is exhausted, +closed, or falls back. A caller arriving while every slot is taken gets a 429, which the core +records as that attempt's failure. + +The limit object is the one the Python router caches (``MaxParallelRequestsLimit``), kept in the +core's routing cache under the same key, so its in-flight count is shared by every attempt. +""" + +import contextlib +from collections.abc import Mapping +from typing import TYPE_CHECKING, Final + +from litellm._internal_context import with_service_target +from litellm.router_rust.invoke.shaping import ROUTER_CLIENTS_TARGET, as_dict + +if TYPE_CHECKING: + from litellm.router_rust.invoke.endpoints import EndpointCore, Plan + +# The attempt helpers that run inside `_deployment_slot` (`_acompletion` holds the same slot +# itself). Sync helpers, rerank, moderation and the search helper hold none. +_SLOTTED_KINDS: Final = frozenset( + { + "chat", + "text", + "embedding", + "image", + "transcription", + "speech", + "adapter", + "generic", + "batch", + "cancel_batch", + "vector_store_create", + } +) + + +@with_service_target(ROUTER_CLIENTS_TARGET) +def max_parallel_requests_limit( + core: "EndpointCore", deployment: Mapping[str, object], kwargs: Mapping[str, object] +) -> object: + """`_get_client(client_type="max_parallel_requests")`: the cached limit, built on first use + (`InitalizeCachedClient.set_max_parallel_requests_client`).""" + from litellm.litellm_core_utils.core_helpers import get_parent_otel_span_from_kwargs + from litellm.router_utils.client_initalization_utils import MaxParallelRequestsLimit + from litellm.utils import calculate_max_parallel_requests + + model_id: Final = as_dict(deployment.get("model_info")).get("id") + key: Final = f"{model_id}_max_parallel_requests_client" + span: Final = get_parent_otel_span_from_kwargs(dict(kwargs)) + cache: Final = core.routing_cache + cached: Final[object] = cache.get_cache(key=key, local_only=True, parent_otel_span=span) # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType, reportArgumentType] # the cache holds untyped values + if cached is not None: + return cached # pyright: ignore[reportUnknownVariableType] # the cache holds untyped values + params: Final = as_dict(deployment.get("litellm_params")) + calculated: Final = calculate_max_parallel_requests( + rpm=params.get("rpm"), # pyright: ignore[reportArgumentType] # the configured value, as the Python router reads it + max_parallel_requests=params.get("max_parallel_requests"), # pyright: ignore[reportArgumentType] # same + tpm=params.get("tpm"), # pyright: ignore[reportArgumentType] # same + default_max_parallel_requests=core.endpoints.default_max_parallel_requests, # pyright: ignore[reportArgumentType] # the setting as given, as the Python router passes it + ) + if calculated: + cache.set_cache( # pyright: ignore[reportUnknownMemberType] # the routing cache + key=key, + value=MaxParallelRequestsLimit( + max_parallel_requests=calculated, + model_id=str(model_id), + model_group=str(deployment.get("model_name", "")), + ), + local_only=True, + ) + return cache.get_cache(key=key, local_only=True, parent_otel_span=span) # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType, reportArgumentType] # see above + + +def enter_deployment_slot( + core: "EndpointCore", plan: "Plan", deployment: Mapping[str, object], kwargs: Mapping[str, object] +) -> contextlib.ExitStack: + """The attempt's slot, entered (a full deployment raises its 429 here); an empty stack when + the plan's helper holds none or the deployment has no limit.""" + from litellm.router_utils.client_initalization_utils import MaxParallelRequestsLimit + + slot: Final = contextlib.ExitStack() + if plan.sync or plan.kind not in _SLOTTED_KINDS: + return slot + limit: Final = max_parallel_requests_limit(core, deployment, kwargs) + if isinstance(limit, MaxParallelRequestsLimit): + slot.enter_context(limit) + return slot + + +def deployment_slot( + core: "EndpointCore", deployment: Mapping[str, object], kwargs: Mapping[str, object] +) -> contextlib.ExitStack: + """`_deployment_slot` outside a plan's attempt (a file created on every deployment).""" + from litellm.router_utils.client_initalization_utils import MaxParallelRequestsLimit + + slot: Final = contextlib.ExitStack() + limit: Final = max_parallel_requests_limit(core, deployment, kwargs) + if isinstance(limit, MaxParallelRequestsLimit): + slot.enter_context(limit) + return slot + + +__all__ = ["deployment_slot", "enter_deployment_slot", "max_parallel_requests_limit"] diff --git a/litellm/router_rust/invoke/streaming.py b/litellm/router_rust/invoke/streaming.py new file mode 100644 index 00000000000..ca70143072f --- /dev/null +++ b/litellm/router_rust/invoke/streaming.py @@ -0,0 +1,1839 @@ +""" +The stream wrappers of the Rust router core: host-side objects around a streamed response. + +Every attempt the Rust core runs that returns a stream is wrapped here before the core sees the +outcome, the way the Python router wraps each attempt's stream: + +- a chat stream (``CustomStreamWrapper``) has its first chunk read inside the attempt, so a stream + that fails before any token fails the attempt and the core retries or falls back as for any error +- a chat, ``/v1/messages`` or Responses stream that fails later, before the caller received content, + hands the failure back to the core: the core decides whether a same-group retry or a fallback takes + over (``RouterCoreHandle.stream_policy`` / ``retry_stream`` / ``fallback_stream``), and the wrapper + splices what that call returns into the stream the caller already holds + +Each wrapper also reports its stream's outcome to the core when the stream ends, as litellm's stream +logging reports it to the Python router's callbacks: the attempt that returned the stream counted +only its request (the attempt answer is ``deferred``), so the success with the stream's final tokens, +or a failure the stream raised after it opened, is recorded then (``StreamAttempt.report_success`` / +``report_failure``). A stream closed before it ends reports nothing. + +The wrappers decide nothing about routing: what may take over, the retry budget and every +deployment pick come from the core. They hold live Python iterators, frames and hidden params, +which is why they live in the host. Ported from the Python router's ``_acompletion_streaming_iterator``, +``_completion_streaming_iterator``, ``_aanthropic_messages_streaming_iterator`` and +``_aresponses_streaming_iterator``, without importing it. +""" + +import asyncio +import contextlib +import dataclasses +import json +import weakref +from collections.abc import ( + AsyncGenerator, + AsyncIterable, + AsyncIterator, + Awaitable, + Callable, + Coroutine, + Generator, + Iterable, + Iterator, + Mapping, + Sequence, +) +from typing import ( + TYPE_CHECKING, + Final, + Generic, + NoReturn, + TypeAlias, + TypeVar, + cast, # noqa: TID251 # untyped stream iterators and provider chunk lists +) + +import anyio + +import litellm.exceptions as litellm_exceptions +from litellm._logging import verbose_router_logger +from litellm.exceptions import MidStreamFallbackError +from litellm.litellm_core_utils.add_retry_fallback_headers import ( + add_retry_headers_to_response, + get_hidden_params_dict, + replace_complexity_router_headers, +) +from litellm.litellm_core_utils.core_helpers import as_str_mapping, without +from litellm.litellm_core_utils.hidden_params import HIDDEN_PARAMS_ATTR, get_hidden_params, set_hidden_params +from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper +from litellm.router_rust.surface.health import failure_already_logged, logged_call +from litellm.types.utils import ModelResponseStream, Usage + +if TYPE_CHECKING: + from litellm.responses.streaming_iterator import BaseResponsesAPIStreamingIterator + from litellm.router_python_host import PythonHost + from litellm.router_rust.core import RustCore + from litellm.router_utils.pre_call_checks.continuation_prefill_check import ContinuationOutput, ContinuedPartial + from litellm.types.llms.openai import ResponseAPIUsage, ResponseIncompleteEvent + +# A stream that never emits content or an error could grow a held-frame buffer without bound, so +# hitting either cap commits the stream instead. +MAX_BUFFERED_PRE_CONTENT_ANTHROPIC_CHUNKS: Final = 200 +MAX_HELD_PRE_OUTPUT_RESPONSES_EVENTS: Final = 200 + +# The hidden params and headers a fallback's response carries, as dicts the proxy writes into. +_Prepared: TypeAlias = "tuple[dict[str, object], dict[str, object]]" +_RestT: Final = TypeVar("_RestT", Iterator[object], AsyncIterator[object]) + + +# --- the attempt a stream came from ------------------------------------------------------------- + + +@dataclasses.dataclass(slots=True) +class StreamAttempt: + """The attempt that returned a stream, and the core that can still recover it. + + ``request`` is the call's prepared request, the one every attempt of the call is built from; + ``kwargs`` are this attempt's own kwargs. The origin fields are what the core reads to decide + what may take over a failure of this stream. + """ + + core: "RustCore" + op: str + request: Mapping[str, object] + kwargs: Mapping[str, object] + group: str + order: int | None + depth: int + level: int + deployment_id: str | None + attempted_retries: int + committed_budget: int | None = None + # The deployment's max_parallel_requests slot the attempt holds; an open chat stream takes it. + slot: contextlib.ExitStack | None = None + # Set by the wrapper that reports the stream's outcome when it ends; the attempt's answer is then + # deferred and counts only the request and the tokens read so far (``counted_tokens``). + reports_outcome: bool = False + counted_tokens: int = 0 + outcome_reported: bool = False + # The background recording of the stream's failure, which a same-group retry waits for so it + # picks with the failed deployment already cooled down. + failure_recording: "asyncio.Task[None] | None" = None + + @property + def model_group(self) -> str: + return self.group + + def origin(self) -> Mapping[str, object]: + return { + "group": self.group, + "order": self.order, + "depth": self.depth, + "level": self.level, + "deployment_id": self.deployment_id, + "attempted_retries": self.attempted_retries, + "committed_budget": self.committed_budget, + } + + def policy(self, trigger: BaseException | None = None) -> Mapping[str, object]: + return self.core.stream_policy(self.request, self.origin(), trigger) + + def can_retry(self) -> bool: + return self.policy().get("can_retry") is True + + def can_fall_back(self) -> bool: + return self.policy().get("can_fall_back") is True + + def refusal_fallback_available(self) -> bool: + return self.policy().get("refusal_fallback") is True + + def retry_budget(self, trigger: BaseException) -> tuple[int, bool]: + policy: Final = self.policy(trigger) + budget: Final = policy.get("retry_budget") + return (budget if isinstance(budget, int) else 0), policy.get("policy_applies") is True + + def should_retry_status(self, status_code: int) -> bool: + return bool(self.core._host.invoke("should_retry")(status_code)) # pyright: ignore[reportPrivateUsage, reportUnknownMemberType] # the wrapper serves its own core's host + + async def fallback(self, trigger: BaseException, request: Mapping[str, object] | None = None) -> object: + """The fallback walk for this stream's failure; raises when no fallback answers.""" + return await self.core.recover_stream("fallback", self.op, request or self.request, self.origin(), trigger) + + def fallback_sync(self, trigger: BaseException) -> object: + return self.core.recover_stream_sync("fallback", self.op, self.request, self.origin(), trigger) + + async def retry(self, trigger: BaseException) -> "tuple[object, int, int] | BaseException": + """A same-group retry: the retried response with the retry count and budget, or the newest failure.""" + return await self.core.retry_stream(self.op, self.request, self.origin(), trigger) + + # -- the stream's outcome, reported once when it ends -------------------------------------- + + def logged(self) -> dict[str, object]: # mutable-ok: litellm's logging kwargs + """The logging kwargs litellm's stream logging hands the outcome callbacks, for this attempt.""" + from litellm.constants import ROUTER_USAGE_COUNTED_TOKENS_METADATA_KEY + + kwargs: Final = self.kwargs + metadata: Final = as_str_mapping(kwargs.get("litellm_metadata")) or as_str_mapping(kwargs.get("metadata")) or {} + return { + "call_type": self.op, + "litellm_params": { + "metadata": {**metadata, ROUTER_USAGE_COUNTED_TOKENS_METADATA_KEY: self.counted_tokens}, + "model_info": kwargs.get("model_info") or {}, + "api_base": kwargs.get("api_base"), + "client_side_timeout": kwargs.get("client_side_timeout"), + "timeout": kwargs.get("timeout"), + **({"cooldown_time": kwargs["cooldown_time"]} if "cooldown_time" in kwargs else {}), + }, + } + + def _failure_logged(self, error: BaseException) -> dict[str, object]: # mutable-ok: litellm's logging kwargs + return {**logged_call(self.op, self.kwargs), **self.logged(), "exception": _provider_error(error)} + + def _callbacks(self, *, sync: bool) -> Mapping[str, bool]: + """The outcome callbacks that run for this attempt when its stream ends.""" + return self.core.outcome_callbacks(self.kwargs, sync=sync) + + def _claim(self) -> bool: + if not self.reports_outcome or self.outcome_reported: + return False + self.outcome_reported = True + return True + + def report_success(self, total_tokens: int, last: object = None) -> None: + """The stream ended: count its success and the tokens the attempt did not count, and re-pin + the session on ``last``, the stream's final item (``deployment_callback_on_success``), in the + background.""" + if self._claim(): + logged: Final = {**self.logged(), "standard_logging_object": {"total_tokens": total_tokens}} + _record_in_background(self.core.report_stream_success(logged, last, self._callbacks(sync=False))) + + def report_failure(self, error: BaseException) -> None: + """The stream failed after it opened: count the failure, cool down by the rules and count one + more request (``async_deployment_callback_on_failure``), in the background. A logging object + shared with an earlier failure blocks it: the stream logs through ``attempt_scoped_copy``, read + here, when the stream fails.""" + if self._claim() and not failure_already_logged(self.kwargs, "async_failure"): + self.failure_recording = _record_in_background( + self.core.report_stream_failure(self._failure_logged(error), self._callbacks(sync=False)) + ) + + async def failure_recorded(self) -> None: + """Wait for the stream's failure to be recorded, if one is being recorded.""" + recording: Final = self.failure_recording + if recording is not None and not recording.done(): + _ = await asyncio.wait({recording}) + + def report_first_chunk_failure(self, error: BaseException) -> None: + """The stream failed before its attempt returned it. ``_dispatch_stream_failure`` logs that in + a task of its own, so the record reads the error as the router left it by then, with its + routing note when no real suspension came first. The attempt checked the logging object.""" + _record_in_background(self.core.report_stream_failure(self._failure_logged(error), self._callbacks(sync=False))) + + def report_success_sync(self) -> None: + """A sync stream ended: a sync call's success callback counts the success alone.""" + if self._claim(): + _record_now(lambda: self.core.report_stream_success_sync(self.logged(), self._callbacks(sync=True))) + + def report_failure_sync(self, error: BaseException) -> None: + """A sync stream failed after it opened: counted and cooled down, with no request counted.""" + if self._claim() and not failure_already_logged(self.kwargs, "sync_failure"): + logged: Final = self._failure_logged(error) + _record_now(lambda: self.core.report_stream_failure_sync(logged, self._callbacks(sync=True))) + + +# --- reporting a stream's outcome --------------------------------------------------------------- + +# The background recordings still running, held so none is collected before it finishes. +_RECORDINGS: Final[set["asyncio.Task[None]"]] = set() # mutable-ok: tasks leave it when they finish + + +def _provider_error(error: BaseException) -> BaseException: + """The provider's own error, which the failure rules classify, rather than the router's envelope.""" + if isinstance(error, MidStreamFallbackError) and error.original_exception is not None: + return error.original_exception + return error + + +def _record_in_background(recording: Coroutine[object, object, object]) -> "asyncio.Task[None]": + """Run a recording through the core's async cache driver without holding up the stream, as litellm + schedules its stream logging.""" + + async def run() -> None: + try: + _ = await recording + except Exception as error: # noqa: BLE001 # a recording never fails the stream + verbose_router_logger.warning("a stream's outcome was not recorded: %s", error) + + task: Final = asyncio.get_running_loop().create_task(run()) + _RECORDINGS.add(task) + task.add_done_callback(_RECORDINGS.discard) + return task + + +def _record_now(recording: Callable[[], object]) -> None: + try: + _ = recording() + except Exception as error: # noqa: BLE001 # see _record_in_background + verbose_router_logger.warning("a stream's outcome was not recorded: %s", error) + + +_ItemT: Final = TypeVar("_ItemT") + + +async def _areported( + attempt: StreamAttempt, + items: AsyncIterable[_ItemT], + total_tokens: Callable[[object], int], + observe: Callable[[object], None] | None = None, +) -> AsyncGenerator[_ItemT, None]: + """``items``, reporting the attempt's stream once they end: a success with the tokens + ``total_tokens`` reads off the last item, or the failure they raise. Closed early, nothing is + reported. The wrapper marks the attempt (``reports_outcome``) when it wraps the stream, before + the attempt answers.""" + last: object = None # rebind-ok: the newest item + try: + async for item in items: + last = item + if observe is not None: + observe(item) + yield item + except Exception as error: + attempt.report_failure(error) + raise + attempt.report_success(total_tokens(last), last) + + +def _reported(attempt: StreamAttempt, items: Iterable[_ItemT]) -> Generator[_ItemT, None, None]: + """``_areported`` for a sync stream, whose success counts no tokens.""" + try: + yield from items + except Exception as error: + attempt.report_failure_sync(error) + raise + attempt.report_success_sync() + + +def _usage_total_tokens(usage: object) -> int: + mapping: Final = as_str_mapping(usage) + tokens: Final[object] = mapping.get("total_tokens") if mapping is not None else getattr(usage, "total_tokens", None) + return tokens if isinstance(tokens, int) and not isinstance(tokens, bool) and tokens > 0 else 0 + + +def _chat_total_tokens(last: object) -> int: + """The total tokens of a finished chat stream: its usage chunk's, else the final usage its last + chunk's hidden params carry once the stream assembled its response.""" + return _usage_total_tokens(getattr(last, "usage", None)) or _usage_total_tokens( + (get_hidden_params(last) or {}).get("usage") + ) + + +def _responses_total_tokens(last: object, source: object) -> int: + """The total tokens of a finished Responses stream: its terminal event's response usage.""" + for event in (last, getattr(source, "completed_response", None)): + tokens = _usage_total_tokens(getattr(getattr(event, "response", None), "usage", None)) + if tokens: + return tokens + return 0 + + +class _AnthropicStreamUsage: + """The usage a ``/v1/messages`` stream reports in its frames, as litellm logs it: the prompt's input, + cache-read and cache-write tokens on ``message_start`` (a ``message_delta`` may restate them), and the + cumulative output tokens on each ``message_delta``. A frame split across byte chunks is read once its + line completes.""" + + __slots__ = ("_open_line", "_output", "_prompt") + + def __init__(self) -> None: + self._open_line = b"" + self._prompt: dict[str, int] = {} # mutable-ok: the newest count per prompt field + self._output = 0 + + def observe(self, chunk: object) -> None: + if not isinstance(chunk, (bytes, bytearray)): + return + *lines, self._open_line = (self._open_line + bytes(chunk)).split(b"\n") + for line in lines: + if line.startswith(b"data:") and b'"usage"' in line: + self._read(line[len(b"data:") :]) + + def _read(self, payload: bytes) -> None: + try: + event: Final[object] = json.loads(payload) # pyright: ignore[reportAny] # plain JSON + except ValueError: + return + frame: Final = as_str_mapping(event) or {} + kind: Final = frame.get("type") + usage: Final = as_str_mapping( + (as_str_mapping(frame.get("message")) or {}).get("usage") if kind == "message_start" else frame.get("usage") + ) + if usage is None or kind not in ("message_start", "message_delta"): + return + for field in ("input_tokens", "cache_creation_input_tokens", "cache_read_input_tokens"): + count = usage.get(field) + if isinstance(count, int) and not isinstance(count, bool): + self._prompt[field] = count + output: Final = usage.get("output_tokens") + if isinstance(output, int) and not isinstance(output, bool): + self._output = output + + def total_tokens(self, _last: object = None) -> int: + return sum(self._prompt.values()) + self._output + + +# --- the first chunk ---------------------------------------------------------------------------- + + +@dataclasses.dataclass(frozen=True, slots=True) +class _FirstChunk(Generic[_RestT]): + """The chunk an attempt already read off a stream (None: the provider closed it empty) and the + iterator it was read from, which carries on after it.""" + + chunk: object + rest: _RestT + + +def _mid_stream_retry_trigger(error: MidStreamFallbackError) -> Exception: + """The provider's own error, which is what the retry policy and the retry rules classify.""" + return error.original_exception if error.original_exception is not None else error + + +async def _aread_first_chunk(stream: CustomStreamWrapper) -> _FirstChunk[AsyncIterator[object]]: + """Read the first chunk while the attempt still holds the request, so a stream that fails before + any token fails the attempt. Iterates the way ``async for`` does: one ``__aiter__`` call, then + ``__anext__`` on what it returned.""" + rest: Final = stream.__aiter__() + try: + return _FirstChunk(chunk=await rest.__anext__(), rest=rest) + except StopAsyncIteration: + return _FirstChunk(chunk=None, rest=rest) + except MidStreamFallbackError as e: + raise _mid_stream_retry_trigger(e) from e + + +def _read_first_chunk(stream: CustomStreamWrapper) -> _FirstChunk[Iterator[object]]: + rest: Final = iter(stream) + try: + return _FirstChunk(chunk=next(rest), rest=rest) + except StopIteration: + return _FirstChunk(chunk=None, rest=rest) + except MidStreamFallbackError as e: + raise _mid_stream_retry_trigger(e) from e + + +def _chunks_with_first(model_response: Iterable[object], taken: _FirstChunk[Iterator[object]]) -> Iterator[object]: + if taken.chunk is None: + return + yield taken.chunk + yield from taken.rest + + +async def _achunks_with_first( + model_response: AsyncIterable[object], taken: _FirstChunk[AsyncIterator[object]] +) -> AsyncIterator[object]: + if taken.chunk is None: + return + yield taken.chunk + async for item in taken.rest: + yield item + + +# --- shared helpers ----------------------------------------------------------------------------- + + +_CONTENT_DELTA_KEYS: Final = ( + "content", + "tool_calls", + "function_call", + "reasoning_content", + "thinking_blocks", + "reasoning_items", + "audio", + "images", + "annotations", +) + + +def _delta_has_content(delta: object) -> bool: + get: Final[object] = getattr(delta, "get", None) + if not callable(get): + return False + for key in _CONTENT_DELTA_KEYS: + value: object = get(key) + if value: + return True + return False + + +def _stream_chunks_have_generated_content(chunks: Sequence[object]) -> bool: + for chunk in chunks: + choices: object = getattr(chunk, "choices", None) + if not isinstance(choices, Sequence) or not choices: + continue + delta: object = getattr(choices[0], "delta", None) # pyright: ignore[reportUnknownArgumentType] # a chunk choice + if _delta_has_content(delta): + return True + return False + + +def _raise_original_mid_stream_error(e: MidStreamFallbackError) -> NoReturn: + """Surface the provider error the stream wrapper carried instead of the internal MidStreamFallbackError.""" + if e.original_exception is not None: + raise e.original_exception from e + raise e + + +def _raise_fallback_failure(error: Exception) -> NoReturn: + """What a stream's recovery ends on: the error the fallback's own stream raised, or the failure + a walk that served nothing was handed; never the internal MidStreamFallbackError.""" + if isinstance(error, MidStreamFallbackError) and error.original_exception is not None: + raise error.original_exception from error + raise error + + +def _stream_pre_content_error(error: Exception, model: str, llm_provider: str = "anthropic") -> MidStreamFallbackError: + """The envelope the fallback chain judges a failure by when the client has received no content yet.""" + return MidStreamFallbackError( + message=str(error), + model=model, + llm_provider=llm_provider, + original_exception=error, + is_pre_first_chunk=True, + ) + + +def _context_window_stream_fallback_error( + error: "MidStreamFallbackError | litellm_exceptions.ContextWindowExceededError", + model: object, + llm_provider: str, + has_output: bool, +) -> MidStreamFallbackError: + if isinstance(error, litellm_exceptions.ContextWindowExceededError): + if has_output: + raise error + return _stream_pre_content_error(error, model if isinstance(model, str) else "", llm_provider) + if isinstance(error.original_exception, litellm_exceptions.ContextWindowExceededError) and ( + has_output or error.generated_content or error.emitted_disqualifying_content + ): + raise error.original_exception + return error + + +def _is_context_window(error: "MidStreamFallbackError | litellm_exceptions.ContextWindowExceededError") -> bool: + return isinstance(error, litellm_exceptions.ContextWindowExceededError) or isinstance( + error.original_exception, litellm_exceptions.ContextWindowExceededError + ) + + +def _fallback_trigger(e: MidStreamFallbackError, unwrapped: tuple[type[Exception], ...]) -> Exception: + """The error the fallback walk dispatches on: the provider's own for the kinds with a list of their own.""" + return e.original_exception if isinstance(e.original_exception, unwrapped) else e + + +def _prepare_fallback_hidden_params( + fallback_response: object, +) -> _Prepared: + fallback_hidden_params: Final = get_hidden_params_dict(fallback_response) + fallback_headers: Final = fallback_hidden_params.get("additional_headers") + if not isinstance(fallback_headers, dict): + return fallback_hidden_params, {} + return fallback_hidden_params, cast("dict[str, object]", fallback_headers) # cast-ok: a dict of headers + + +def _apply_fallback_hidden_params_to_item( + fallback_item: object, + prepared_fallback_hidden_params: _Prepared, +) -> None: + if fallback_item is None or not hasattr(fallback_item, HIDDEN_PARAMS_ATTR): + return + fallback_hidden_params, fallback_headers = prepared_fallback_hidden_params + item_hidden_params: Final = get_hidden_params_dict(fallback_item) + item_headers: object = item_hidden_params.get("additional_headers") # rebind-ok: replaced when not a dict + if not isinstance(item_headers, dict): + item_headers = {} # rebind-ok: rebound as the stream advances + set_hidden_params( + fallback_item, + { + **item_hidden_params, + **fallback_hidden_params, + "additional_headers": {**cast("dict[str, object]", item_headers), **fallback_headers}, # cast-ok: checked + }, + ) + + +def _combine_fallback_usage(fallback_item: ModelResponseStream, complete_response_object_usage: Usage | None) -> None: + """Merge partial-stream usage with fallback-stream usage on the chunk. A usage given as a plain + dict comes out all-zero: ``combine_usage_objects`` reads token counts as attributes.""" + from litellm.cost_calculator import BaseTokenUsageProcessor + + usage: Final = cast("Usage | None", getattr(fallback_item, "usage", None)) # cast-ok: a chunk's usage + usage_objects: Final = [usage] if usage is not None else [] + if complete_response_object_usage is not None: + usage_objects.append(complete_response_object_usage) + fallback_item.usage = BaseTokenUsageProcessor.combine_usage_objects(usage_objects=usage_objects) # pyright: ignore[reportAttributeAccessIssue] # usage is not a declared chunk field # rebind-ok: the Python router rewrites the chunk's usage in place + + +# --- mid-stream continuation -------------------------------------------------------------------- + +CONTINUATION_FLAG: Final = "enable_mid_stream_fallback_continuation" +NO_FLAG: Final = object() +# The routing-body key naming the deployments a continuation may not land on (``continuation_deployments``). +CONTINUATION_EXCLUDED_IDS: Final = "continuation_excluded_ids" + + +def continuation_flag(value: object) -> bool: + """``_is_flag_enabled``: True, or a string ``str_to_bool`` reads as True.""" + from litellm.secret_managers.main import str_to_bool + + return value is True or (isinstance(value, str) and str_to_bool(value) is True) + + +def register_continuation_prefill_check(host: "PythonHost") -> None: + """``_register_continuation_prefill_check``: process-global on purpose, as in the Python router.""" + from litellm.router_utils.pre_call_checks.continuation_prefill_check import ContinuationPrefillDeploymentCheck + + host.callbacks.add_litellm_callback(ContinuationPrefillDeploymentCheck()) + + +def continuation_routing(core: "RustCore", request: Mapping[str, object]) -> Mapping[str, object]: + """The routing fields of a continuation request: the deployments that cannot continue it, which the + core leaves out of selection the way ``continuation_deployments`` filters the healthy list.""" + from litellm.litellm_core_utils.routing_common_utils import MID_STREAM_CONTINUATION_KWARG + from litellm.router_utils.pre_call_checks.continuation_prefill_check import ( + MID_STREAM_CONTINUATION_MARKER, + deployment_can_continue, + ) + + if request.get(MID_STREAM_CONTINUATION_KWARG) is not MID_STREAM_CONTINUATION_MARKER: + return {} + registry: Final = core.registry() + excluded: Final = [ + deployment.model_info.id + for record in registry.all() + if (deployment := registry.deployment(str(_record_id(record)))) is not None + and not deployment_can_continue(deployment.model_dump(exclude_none=True)) + ] + return {"body": {CONTINUATION_EXCLUDED_IDS: excluded}} + + +def _record_id(record: Mapping[str, object]) -> object: + model_info: Final = record.get("model_info") + return model_info.get("id") if isinstance(model_info, Mapping) else None # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] # a plain record + + +def _names_deployment_that_cannot_continue(core: "RustCore", target: object) -> bool: + from litellm.router_utils.fallback_event_handlers import fallback_target_name + from litellm.router_utils.pre_call_checks.continuation_prefill_check import deployment_can_continue + + name: Final = fallback_target_name(target) + registry: Final = core.registry() + model_names: Final = {record.get("model_name") for record in registry.all()} + if name is None or name in model_names or not registry.has_deployment(name): + return False + deployment: Final = registry.deployment(name) + return deployment is not None and not deployment_can_continue(deployment.model_dump(exclude_none=True)) + + +def _mid_stream_continuation_eligible( + attempt: StreamAttempt, e: MidStreamFallbackError, cut_deployment_id: object +) -> bool: + """``_mid_stream_continuation_eligible``: whether a stream that broke after plain assistant text may + continue via prefill.""" + from litellm.integrations.custom_guardrail import CustomGuardrail + from litellm.router_utils.fallback_event_handlers import ( + continuation_fallback_targets, + fallback_lookup_groups, + fallback_target_overrides_request_params, + ) + from litellm.router_utils.pre_call_checks.continuation_prefill_check import constrains_continuation + from litellm.types.guardrails import GuardrailEventHooks + + core: Final = attempt.core + if not core.enable_mid_stream_fallback_continuation: + return False + if not e.generated_content or e.emitted_disqualifying_content: + return False + if isinstance(e.original_exception, litellm_exceptions.ContentPolicyViolationError): + return False + if isinstance(cut_deployment_id, str) and not core.registry().has_deployment(cut_deployment_id): + return False + if constrains_continuation(_params_sent_to_deployment(attempt, cut_deployment_id)): + return False + request: Final = attempt.request + requested_group: Final = request.get("model") + lookup_groups: Final = fallback_lookup_groups( + request, requested_group if isinstance(requested_group, str) else None + ) + targets: Final = continuation_fallback_targets(core.settings().fallbacks, lookup_groups) # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType] # fallbacks as Python spells them + if any(fallback_target_overrides_request_params(target) for target in targets): + return False + if any(_names_deployment_that_cannot_continue(core, target) for target in targets): + return False + return not any( + isinstance(callback, CustomGuardrail) + and callback._event_hook_is_event_type( # pyright: ignore[reportPrivateUsage] # the only check of a guardrail's hooks + GuardrailEventHooks.post_call + ) + for callback in core._host.callback_list("callbacks") # pyright: ignore[reportPrivateUsage, reportArgumentType] # the core's own host; read as getattr(litellm, name) + ) + + +def _extend_trailing_text_block(blocks: Sequence[Mapping[str, object]], text: str) -> Sequence[Mapping[str, object]]: + """Continue the last text block in place so the prefill reads as the text the caller already received.""" + last: Final = blocks[-1] if blocks else None + if last is not None and last.get("type") == "text" and isinstance(last_text := last.get("text"), str): + return [*blocks[:-1], {**last, "text": f"{last_text}{text}"}] + return [*blocks, {"type": "text", "text": text}] + + +def _build_completion_continuation_input( + messages: Sequence[Mapping[str, object]], generated_content: str +) -> Sequence[Mapping[str, object]] | None: + """The partial output as an assistant prefill, or a trailing ``prefix: True`` assistant turn extended in + place. None when the request ends in a plain assistant turn, in a prefill turn whose content this cannot + extend, or when the cut came mid-echo of the prefill.""" + from pydantic import TypeAdapter, ValidationError + + from litellm.router_utils.pre_call_checks.continuation_prefill_check import continued_text + + last: Final = messages[-1] if messages else None + if last is None or last.get("role") != "assistant": + return [*messages, {"role": "assistant", "content": generated_content, "prefix": True}] + if last.get("prefix") is not True: + return None + new_text: Final = continued_text(last, generated_content) + if new_text is None: + return None + content: Final = last.get("content") + if content is None or isinstance(content, str): + return [*messages[:-1], {**last, "content": f"{content or ''}{new_text}", "prefix": True}] + try: + blocks: Final = TypeAdapter(list[dict[str, object]]).validate_python(content) + except ValidationError: + return None + return [*messages[:-1], {**last, "content": _extend_trailing_text_block(blocks, new_text), "prefix": True}] + + +def _price_cut_partial(cut_stream: CustomStreamWrapper, partial_response: object) -> "ContinuedPartial | None": + """Price the cut stream's partial, plus any partial it inherited, before a fallback can reuse its logging.""" + from litellm.cost_calculator import BaseTokenUsageProcessor + from litellm.router_utils.pre_call_checks.continuation_prefill_check import ContinuedPartial + + cut_logging: Final = cut_stream.logging_obj + partial_usage: Final = getattr(partial_response, "usage", None) + if partial_response is None or not isinstance(partial_usage, Usage): + return None + details: Final = cut_logging.model_call_details + inherited_usage: Final = details.get("continued_partial_usage") + inherited_cost: Final = details.get("continued_partial_cost") + inherited_id: Final = details.get("continued_response_id") + priced_partial: Final = partial_response.model_copy( # pyright: ignore[reportAttributeAccessIssue, reportUnknownMemberType, reportUnknownVariableType] # a ModelResponse + update={"model": cut_stream.model} if cut_stream.model else {} # pyright: ignore[reportUnknownMemberType] # untyped attribute + ) + priced_partial._hidden_params = without( # pyright: ignore[reportUnknownMemberType] # no public accessor + get_hidden_params_dict(partial_response), "response_cost" + ) + return ContinuedPartial( + usage=BaseTokenUsageProcessor.combine_usage_objects( + usage_objects=[partial_usage, inherited_usage] if isinstance(inherited_usage, Usage) else [partial_usage] + ), + response_cost=(cut_logging.response_cost_calculator(result=priced_partial) or 0.0) # pyright: ignore[reportUnknownArgumentType] # a ModelResponse + + (inherited_cost if isinstance(inherited_cost, float) else 0.0), + response_id=inherited_id if isinstance(inherited_id, str) else str(getattr(partial_response, "id", "")), + ) + + +def _carry_partial_spend_to_fallback(fallback_response: object, cut_partial: "ContinuedPartial | None") -> None: + """Bill what the cut legs consumed on the continuation's success log.""" + from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging + + fallback_logging: Final = getattr(fallback_response, "logging_obj", None) + if cut_partial is None or not isinstance(fallback_logging, LiteLLMLogging): + return + fallback_logging.record_continued_partial( + usage=cut_partial.usage, # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType] # a Usage + response_cost=cut_partial.response_cost, + response_id=cut_partial.response_id, + ) + + +@dataclasses.dataclass(frozen=True, slots=True) +class _Continuation: + """What a chat continuation sends the fallback and splices back.""" + + request: Mapping[str, object] + output: "ContinuationOutput" + partial: "ContinuedPartial | None" + usage: Usage | None + + +async def _chat_continuation( + attempt: StreamAttempt, e: MidStreamFallbackError, model_response: CustomStreamWrapper +) -> _Continuation: + """The request a fallback continues a chat stream from, once it broke after content. Raises the + provider's error when the continuation is declined or the output ceiling is spent.""" + from litellm.constants import MID_STREAM_FAILURE_LOG_WAIT_SECONDS + from litellm.litellm_core_utils.routing_common_utils import MID_STREAM_CONTINUATION_KWARG + from litellm.main import ( + stream_chunk_builder, # pyright: ignore[reportUnknownVariableType] # untyped chunk list argument + ) + from litellm.router_utils.pre_call_checks.continuation_prefill_check import ( + MID_STREAM_CONTINUATION_MARKER, + ContinuationOutput, + continuation_output_ceilings, + ) + + request: Final = attempt.request + raw_messages: Final = request.get("messages") + messages: Final = cast("list[Mapping[str, object]]", raw_messages or []) # cast-ok: chat messages + complete: Final = stream_chunk_builder(chunks=model_response.chunks, messages=messages) # pyright: ignore[reportUnknownArgumentType] # provider chunk list + usage: Final = cast("Usage | None", getattr(complete, "usage", None)) # cast-ok: a built response's usage + continuation_messages: Final = _build_completion_continuation_input(messages, e.generated_content) + output: Final = ( + ContinuationOutput( + e.generated_content, + continuation_messages[-1], + str(request.get("model") or ""), + complete, + checks_post_call_rules=request.get(MID_STREAM_CONTINUATION_KWARG) is not MID_STREAM_CONTINUATION_MARKER, + ) + if continuation_messages is not None + else None + ) + partial: Final = _price_cut_partial(model_response, complete) if output is not None else None + if partial is not None: + model_response.logging_obj.mark_stash_billed_by_continuation() + emitted_tokens: Final = int(getattr(usage, "completion_tokens", 0) or 0) + reduced: Final = continuation_output_ceilings(request, emitted_tokens) + if reduced is None or continuation_messages is None or output is None: + _raise_original_mid_stream_error(e) + await model_response.wait_for_failure_logging(timeout=MID_STREAM_FAILURE_LOG_WAIT_SECONDS) + return _Continuation( + request={ + **request, + **reduced, + "messages": continuation_messages, + MID_STREAM_CONTINUATION_KWARG: MID_STREAM_CONTINUATION_MARKER, + }, + output=output, + partial=partial, + usage=usage, + ) + + +# --- chat completions --------------------------------------------------------------------------- + + +class FallbackAwareStreamWrapper(CustomStreamWrapper): + """A chat stream built around the attempt the core picked first, which repoints itself when a + fallback takes over.""" + + fallback_headers_adopted: bool = False + + def adopt_fallback_response_headers( + self, + fallback_response: object, + prepared_fallback_hidden_params: _Prepared, + ) -> None: + """Repoint this wrapper at the deployment that served the stream. Replaces rather than merges, + so the failed attempt's request id, counters, ``model_id`` and ``api_base`` cannot leak.""" + self._response_headers = getattr(fallback_response, "_response_headers", None) + fallback_hidden_params, fallback_headers = prepared_fallback_hidden_params + if fallback_hidden_params: + self._hidden_params = {**fallback_hidden_params, "additional_headers": dict(fallback_headers)} + self._base_hidden_params = {**self._hidden_params, "response_cost": None} + self.fallback_headers_adopted = True + + +def _adopt_fallback_response_headers( + wrapper_ref: "weakref.ref[FallbackAwareStreamWrapper]", fallback_response: object +) -> _Prepared: + prepared: Final = _prepare_fallback_hidden_params(fallback_response) + adopting_wrapper: Final = wrapper_ref() + if adopting_wrapper is not None: + adopting_wrapper.adopt_fallback_response_headers(fallback_response, prepared) + return prepared + + +def _chat_wrapper_init( + wrapper: FallbackAwareStreamWrapper, stream: object, model_response: CustomStreamWrapper +) -> None: + CustomStreamWrapper.__init__( # pyright: ignore[reportUnknownMemberType] # the base constructor is untyped + wrapper, + completion_stream=stream, + model=model_response.model, # pyright: ignore[reportUnknownMemberType] # an untyped attribute + custom_llm_provider=model_response.custom_llm_provider, + logging_obj=model_response.logging_obj, + _response_headers=getattr(model_response, "_response_headers", None), + ) + inner_chunks: Final[object] = getattr(model_response, "chunks", None) + if isinstance(inner_chunks, list): + chunk_list: Final = cast("list[ModelResponseStream]", inner_chunks) # cast-ok: the inner stream's chunks + wrapper.chunks = chunk_list # rebind-ok: initializes the wrapper being built + hidden_params: Final = get_hidden_params(model_response) + if hidden_params is not None: + wrapper._hidden_params = dict(hidden_params) # pyright: ignore[reportPrivateUsage] # the wrapper's own field # rebind-ok: initializes the wrapper being built + + +class _AsyncChatStream(FallbackAwareStreamWrapper): + def __init__(self, generator: AsyncGenerator[object, None], model_response: CustomStreamWrapper) -> None: # pyright: ignore[reportMissingSuperCall] # _chat_wrapper_init calls it + _chat_wrapper_init(self, generator, model_response) + self._async_generator: Final = generator + + def __aiter__(self) -> "_AsyncChatStream": # pyright: ignore[reportIncompatibleMethodOverride] # iterates the generator + return self + + async def __anext__(self) -> object: # pyright: ignore[reportIncompatibleMethodOverride] # yields what the stream yields + return await self._async_generator.__anext__() + + +class _SyncChatStream(FallbackAwareStreamWrapper): + def __init__(self, generator: Generator[object, None, None], model_response: CustomStreamWrapper) -> None: # pyright: ignore[reportMissingSuperCall] # _chat_wrapper_init calls it + _chat_wrapper_init(self, generator, model_response) + self.chunks = [] # the sync twin keeps no inner chunk list + hidden_params: Final = get_hidden_params(model_response) + if hidden_params is not None: + self._hidden_params = dict(hidden_params) + self._sync_generator: Final = generator + + def __iter__(self) -> "_SyncChatStream": # pyright: ignore[reportIncompatibleMethodOverride] # iterates the generator + return self + + def __next__(self) -> object: # pyright: ignore[reportIncompatibleMethodOverride] # yields what the stream yields + return next(self._sync_generator) + + +def _chat_has_output(model_response: CustomStreamWrapper) -> bool: + return bool(getattr(model_response, "response_uptil_now", "")) or _stream_chunks_have_generated_content( + model_response.chunks + ) + + +def wrap_chat_stream( # noqa: C901 # one function per Python wrapper it ports + attempt: StreamAttempt, + model_response: CustomStreamWrapper, + first_chunk: _FirstChunk[AsyncIterator[object]], + held_slot: contextlib.ExitStack | None = None, +) -> FallbackAwareStreamWrapper: + """``_acompletion_streaming_iterator``: a mid-stream error before content falls back through the core; + after content, the provider's error is raised. ``held_slot`` (the deployment's max_parallel_requests + slot) is released when the stream is exhausted, closed, or falls back.""" + slot: Final = held_slot if held_slot is not None else contextlib.ExitStack() + attempt.reports_outcome = True # rebind-ok: the wrapper reports this attempt's outcome + + async def close_model_response() -> None: + if not hasattr(model_response, "aclose"): + return + try: + await model_response.aclose() + except BaseException as e: # noqa: BLE001 # closing is best-effort + verbose_router_logger.debug("stream_with_fallbacks: error closing model_response: %s", e) + + async def stream_with_fallbacks() -> AsyncGenerator[object, None]: + fallback_response: object = None # rebind-ok: tracked for cleanup + try: + async for item in _areported(attempt, _achunks_with_first(model_response, first_chunk), _chat_total_tokens): + yield item + except (MidStreamFallbackError, litellm_exceptions.ContextWindowExceededError) as stream_error: + e: Final = ( + _context_window_stream_fallback_error( + stream_error, + getattr(model_response, "model", ""), + model_response.custom_llm_provider or "", + _chat_has_output(model_response), + ) + if _is_context_window(stream_error) + else cast("MidStreamFallbackError", stream_error) # cast-ok: not a context window error + ) + with anyio.CancelScope(shield=True): + await close_model_response() + slot.close() + committed: Final = bool( + not e.is_pre_first_chunk + and (e.generated_content or _stream_chunks_have_generated_content(model_response.chunks)) + ) + cut_deployment_id: Final = (get_hidden_params(model_response) or {}).get("model_id") + if committed and not _mid_stream_continuation_eligible(attempt, e, cut_deployment_id): + _raise_original_mid_stream_error(e) + trigger: Final = _fallback_trigger(e, (litellm_exceptions.ContextWindowExceededError,)) + try: + continuation: Final = await _chat_continuation(attempt, e, model_response) if committed else None + fallback_response = await attempt.fallback( # rebind-ok: rebound as the stream advances + trigger, continuation.request if continuation is not None else None + ) + if hasattr(fallback_response, "__aiter__"): + if continuation is not None: + _carry_partial_spend_to_fallback(fallback_response, continuation.partial) + prepared = _adopt_fallback_response_headers(wrapper_ref, fallback_response) # rebind-ok: re-adopted + settled = False # rebind-ok: rebound as the stream advances + fallback_stream: Final = cast("AsyncIterable[object]", fallback_response) # cast-ok: checked + fallback_items: Final = ( + continuation.output.splice(fallback_stream) if continuation is not None else fallback_stream + ) + async for fallback_item in fallback_items: + if not settled: + settled = True + # a fallback that failed over again only repoints itself once it yields + prepared = _adopt_fallback_response_headers(wrapper_ref, fallback_response) + _apply_fallback_hidden_params_to_item(fallback_item, prepared) + if ( + fallback_item + and isinstance(fallback_item, ModelResponseStream) + and hasattr(fallback_item, "usage") + ): + _combine_fallback_usage( + fallback_item, + continuation.usage if continuation is not None else None, # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType] # a Usage + ) + yield fallback_item + else: + # A fallback that answers without a stream yields None (an oddity kept as is). + yield None + except Exception as fallback_error: # noqa: BLE001 # as the Python wrapper: logged, then the stream's error raises + verbose_router_logger.error("Fallback also failed: %s", fallback_error) + _raise_fallback_failure(fallback_error) + finally: + with anyio.CancelScope(shield=True): + await close_model_response() + slot.close() + if fallback_response is not None and hasattr(fallback_response, "aclose"): + try: + await fallback_response.aclose() # pyright: ignore[reportAttributeAccessIssue, reportUnknownMemberType] # checked above + except BaseException as close_error: # noqa: BLE001 # closing is best-effort + verbose_router_logger.debug( + "stream_with_fallbacks: error closing fallback_response: %s", close_error + ) + + wrapped_response: Final = _AsyncChatStream(stream_with_fallbacks(), model_response) + # weak, so the generator closing over it does not keep the wrapper alive + wrapper_ref: Final[weakref.ref[FallbackAwareStreamWrapper]] = weakref.ref(wrapped_response) + return wrapped_response + + +def wrap_chat_stream_sync( + attempt: StreamAttempt, model_response: CustomStreamWrapper, first_chunk: _FirstChunk[Iterator[object]] +) -> FallbackAwareStreamWrapper: + """``_completion_streaming_iterator``: the sync twin, which never continues after content.""" + attempt.reports_outcome = True # rebind-ok: the wrapper reports this attempt's outcome + + def stream_with_fallbacks() -> Generator[object, None, None]: + fallback_response: object = None # rebind-ok: tracked for cleanup + try: + yield from _reported(attempt, _chunks_with_first(model_response, first_chunk)) + except (MidStreamFallbackError, litellm_exceptions.ContextWindowExceededError) as stream_error: + e: Final = ( + _context_window_stream_fallback_error( + stream_error, + getattr(model_response, "model", ""), + model_response.custom_llm_provider or "", + _chat_has_output(model_response), + ) + if _is_context_window(stream_error) + else cast("MidStreamFallbackError", stream_error) # cast-ok: not a context window error + ) + if attempt.request.get("disable_fallbacks") is True or ( + not e.is_pre_first_chunk + and (e.generated_content or _stream_chunks_have_generated_content(model_response.chunks)) + ): + _raise_original_mid_stream_error(e) + trigger: Final = _fallback_trigger(e, (litellm_exceptions.ContextWindowExceededError,)) + try: + fallback_response = attempt.fallback_sync(trigger) # rebind-ok: rebound as the stream advances + if hasattr(fallback_response, "__iter__"): + prepared = _adopt_fallback_response_headers(wrapper_ref, fallback_response) # rebind-ok: re-adopted + settled = False # rebind-ok: rebound as the stream advances + for fallback_item in cast("Iterable[object]", fallback_response): # cast-ok: checked + if not settled: + settled = True + prepared = _adopt_fallback_response_headers(wrapper_ref, fallback_response) + _apply_fallback_hidden_params_to_item(fallback_item, prepared) + if ( + fallback_item + and isinstance(fallback_item, ModelResponseStream) + and hasattr(fallback_item, "usage") + ): + _combine_fallback_usage(fallback_item, None) + yield fallback_item + else: + yield None + except Exception as fallback_error: # noqa: BLE001 # as the Python wrapper: logged, then the stream's error raises + verbose_router_logger.error("Fallback also failed: %s", fallback_error) + _raise_fallback_failure(fallback_error) + finally: + for closable in (model_response, fallback_response): + close = getattr(closable, "close", None) + if callable(close): + try: + close() + except BaseException as close_error: # noqa: BLE001 # closing is best-effort + verbose_router_logger.debug("stream_with_fallbacks: error closing: %s", close_error) + + wrapped_response: Final = _SyncChatStream(stream_with_fallbacks(), model_response) + wrapper_ref: Final[weakref.ref[FallbackAwareStreamWrapper]] = weakref.ref(wrapped_response) + return wrapped_response + + +# --- /v1/messages ------------------------------------------------------------------------------- + + +class FallbackAwareAnthropicMessagesStream: + """A ``/v1/messages`` byte stream that carries the ``_hidden_params`` the proxy reads its response + headers off, through retries and fallbacks of the stream it wraps.""" + + def __init__( + self, async_generator: AsyncGenerator[bytes, None], source_iterator: object, attempt: StreamAttempt + ) -> None: + self._async_generator: Final = async_generator + self._source_iterator: object = source_iterator + self.attempt: Final = attempt + self.fallback_headers_adopted = False + self._hidden_params: dict[str, object] = dict( # mutable-ok: the proxy writes headers into it + getattr(source_iterator, "_hidden_params", None) or {} + ) + self._followed_source_params: object = None + + @property + def has_buffered_provider_output(self) -> bool: + return getattr(self._source_iterator, "has_buffered_provider_output", False) is True + + @property + def chunks(self) -> object: + return getattr(self._source_iterator, "chunks", None) + + @property + def messages(self) -> object: + return getattr(self._source_iterator, "messages", None) + + @property + def model(self) -> object: + return getattr(self._source_iterator, "model", None) + + def adopt_fallback_source(self, fallback_response: object) -> None: + self._source_iterator = fallback_response + self.fallback_headers_adopted = True + + def partial_response_on_disconnect(self) -> "ResponseIncompleteEvent | None": + from litellm.responses.streaming_iterator import partial_responses_incomplete_event_on_disconnect + + return partial_responses_incomplete_event_on_disconnect(self._source_iterator) + + def follow_source_attribution(self) -> None: + """Mirror the source's current attribution onto this wrapper whenever the source has adopted a + new one: a retry's or a fallback's stream carries a wrapper of its own.""" + source: Final = self._source_iterator + if not getattr(source, "fallback_headers_adopted", False): + return + source_params: Final = getattr(source, "_hidden_params", None) + if source_params is None or source_params is self._followed_source_params: + return + self._followed_source_params = source_params + hidden_params, headers = _prepare_fallback_hidden_params(source) + self.merge_fallback_hidden_params(hidden_params, headers) + + def __aiter__(self) -> "FallbackAwareAnthropicMessagesStream": + return self + + async def __anext__(self) -> bytes: + return await self._async_generator.__anext__() + + async def aclose(self) -> None: + await self._async_generator.aclose() + + def merge_fallback_hidden_params( + self, fallback_hidden_params: Mapping[str, object], fallback_headers: Mapping[str, object] + ) -> None: + existing_headers: Final = cast( # cast-ok: additional_headers is a dict when present + "dict[str, object]", self._hidden_params.get("additional_headers") or {} + ) + self._hidden_params = { + **self._hidden_params, + **fallback_hidden_params, + "additional_headers": dict(replace_complexity_router_headers(existing_headers, fallback_headers)), + } + + +def _anthropic_stream_forwards_ping_live(chunk: object, has_generated_content: bool) -> bool: + from litellm.llms.anthropic.messages.native.streaming_iterator import is_anthropic_ping_chunk + + return not has_generated_content and is_anthropic_ping_chunk(chunk) + + +def _is_retriable_anthropic_status(status_code: int) -> bool: + return status_code == 429 or status_code >= 500 + + +def _anthropic_stream_error_is_gateway_verdict(chunk: object) -> bool: + from litellm.llms.anthropic.messages.native.agentic_streaming_iterator import is_server_fulfilled_tool_leak_error + + return is_server_fulfilled_tool_leak_error(chunk) + + +def _anthropic_stream_should_decline_fallback(has_generated_content: bool, error: MidStreamFallbackError) -> bool: + return has_generated_content or not error.is_pre_first_chunk + + +def _stream_raised_error_status(error: Exception) -> int | None: + raw_status: Final = getattr(error, "status_code", None) + if isinstance(raw_status, int): + return raw_status + if isinstance(raw_status, str) and raw_status.isdigit(): + return int(raw_status) + response_status: Final = getattr(getattr(error, "response", None), "status_code", None) + return response_status if isinstance(response_status, int) else None + + +def _stream_fallback_error_for_raised( + attempt: StreamAttempt, error: Exception, model: str, has_generated_content: bool, llm_provider: str = "anthropic" +) -> MidStreamFallbackError | None: + """The pre-stream retry rule (408, 409, 429, 5xx); None means the raise propagates unchanged.""" + if has_generated_content: + return None + status_code: Final = _stream_raised_error_status(error) + if status_code is None or attempt.should_retry_status(status_code): + return _stream_pre_content_error(error, model, llm_provider) + return None + + +def _stream_commits_now(chunk: object, has_generated_content: bool, buffered_chunk_count: int) -> bool: + from litellm.llms.anthropic.messages.native.streaming_iterator import is_anthropic_content_delta_chunk + + if has_generated_content: + return False + return is_anthropic_content_delta_chunk(chunk) or buffered_chunk_count >= MAX_BUFFERED_PRE_CONTENT_ANTHROPIC_CHUNKS + + +def _bridged_stream_failure(source: object) -> Exception | None: + """The provider's error behind an ``event: error`` frame the completion or Responses bridge wrote + in its place, which litellm's stream logging logged. A native frame has none: litellm logs what the + frame itself maps to (``anthropic_error_frame_exception``), once a takeover closes the stream.""" + failure: Final[object] = getattr(source, "stream_failure", None) + return failure if isinstance(failure, Exception) else None + + +def _anthropic_recoverable_frame_error( + attempt: StreamAttempt, + error_event: tuple[str, str, int] | None, + chunk: object, + has_generated_content: bool, + model_group: str, +) -> Exception | None: + """The exception a provider ``event: error`` frame before content recovers through when a retry of + its class or a fallback can still take over; None forwards the frame as the provider sent it.""" + from litellm.llms.anthropic.common_utils import anthropic_error_frame_exception + + if has_generated_content or error_event is None: + return None + error_type, message, status_code = error_event + if not _is_retriable_anthropic_status(status_code) or _anthropic_stream_error_is_gateway_verdict(chunk): + return None + frame_error: Final = anthropic_error_frame_exception(error_type, message, status_code, model_group) + budget, _ = attempt.retry_budget(frame_error) + if budget > attempt.attempted_retries: + return frame_error + if attempt.can_fall_back(): + return frame_error + return None + + +def wrap_anthropic_stream( + attempt: StreamAttempt, response: AsyncIterator[bytes] +) -> FallbackAwareAnthropicMessagesStream: + """``_aanthropic_messages_streaming_iterator``: lifecycle frames are held while a retry or a fallback + can still take over; an error frame, a raised error or a safeguard refusal before content is + retried in the group, then falls back, as the core decides.""" + from litellm.llms.anthropic.messages.native.fake_stream_iterator import FakeAnthropicMessagesStreamIterator + from litellm.llms.anthropic.messages.native.streaming_iterator import ( + aclose_if_supported, + aclose_superseded_attempt, + parse_anthropic_error_event, + parse_anthropic_refusal_stop_details, + ) + from litellm.llms.anthropic.messages.native.utils import safeguard_refusal_error + + source_iterator: Final = response + replays_a_finished_agentic_turn: Final = isinstance(source_iterator, FakeAnthropicMessagesStreamIterator) + model: Final = attempt.group + usage: Final = _AnthropicStreamUsage() + attempt.reports_outcome = True # rebind-ok: the wrapper reports this attempt's outcome + + async def close_source() -> None: + with anyio.CancelScope(shield=True), contextlib.suppress(BaseException): + await aclose_if_supported(source_iterator) + + async def close_superseded() -> None: + with anyio.CancelScope(shield=True), contextlib.suppress(BaseException): + await aclose_superseded_attempt(source_iterator) + + async def stream_with_fallbacks() -> AsyncGenerator[bytes, None]: + policy: Final = attempt.policy() + has_generated_content = not ( # rebind-ok: set once content is seen, the buffer cap is hit, or nothing can take over + policy.get("can_retry") is True or policy.get("can_fall_back") is True + ) + buffered: tuple[bytes, ...] = () # rebind-ok: flushed once committed or on decline + try: + async for chunk in _areported(attempt, source_iterator, usage.total_tokens, usage.observe): + if _anthropic_stream_forwards_ping_live(chunk, has_generated_content): + yield chunk + continue + if _stream_commits_now(chunk, has_generated_content, len(buffered)): + has_generated_content = True + # A transport can split one SSE data line across byte chunks, so pre-content detection + # parses the held frames plus the current chunk. + parse_window = ( + b"".join(c for c in (*buffered, chunk) if isinstance(c, (bytes, bytearray))) # pyright: ignore[reportUnnecessaryIsInstance] # bridge-path chunks are not always bytes + if not has_generated_content and isinstance(chunk, (bytes, bytearray)) # pyright: ignore[reportUnnecessaryIsInstance] # see above + else chunk + ) + error_event = parse_anthropic_error_event(parse_window) + bridged_failure = _bridged_stream_failure(source_iterator) if error_event is not None else None + if bridged_failure is not None: + attempt.report_failure(bridged_failure) + recoverable_frame_error = _anthropic_recoverable_frame_error( + attempt, error_event, chunk, has_generated_content, model + ) + refusal_stop_details = ( + parse_anthropic_refusal_stop_details(parse_window) + if not has_generated_content and error_event is None + else None + ) + if ( + refusal_stop_details is not None + and not replays_a_finished_agentic_turn + and attempt.refusal_fallback_available() + ): + refusal_error = safeguard_refusal_error(model=model, stop_details=refusal_stop_details) + raise MidStreamFallbackError( + message=refusal_error.message, + model=model, + llm_provider="anthropic", + original_exception=refusal_error, + is_pre_first_chunk=True, + ) + if not has_generated_content and error_event is None: + buffered = (*buffered, chunk) + continue + if recoverable_frame_error is not None and error_event is not None: + raise MidStreamFallbackError( + message=error_event[1], + model=model, + llm_provider="anthropic", + original_exception=recoverable_frame_error, + is_pre_first_chunk=True, + ) + for buffered_chunk in buffered: + yield buffered_chunk + buffered = () + yield chunk + for buffered_chunk in buffered: + yield buffered_chunk + except Exception as stream_error: # noqa: BLE001 # any raised provider error must reach the fallback gate + async for item in _anthropic_recover_stream_error( + attempt, stream_error, has_generated_content, buffered, model, wrapper, close_superseded + ): + yield item + finally: + await close_source() + + wrapper: Final = FallbackAwareAnthropicMessagesStream(stream_with_fallbacks(), source_iterator, attempt) + return wrapper + + +async def _anthropic_recover_stream_error( + attempt: StreamAttempt, + stream_error: Exception, + has_generated_content: bool, + buffered: tuple[bytes, ...], + model: str, + wrapper: FallbackAwareAnthropicMessagesStream, + close_superseded: Callable[[], Awaitable[None]], +) -> AsyncGenerator[bytes, None]: + """Turns a source-iterator failure into a retry, a fallback, or the error reaching the caller.""" + if isinstance(stream_error, MidStreamFallbackError) and _anthropic_stream_should_decline_fallback( + has_generated_content, stream_error + ): + for buffered_chunk in buffered: + yield buffered_chunk + if stream_error.original_exception is not None: + raise stream_error.original_exception from stream_error + raise stream_error + fallback_error: Final = ( + stream_error + if isinstance(stream_error, MidStreamFallbackError) + else _stream_fallback_error_for_raised(attempt, stream_error, model, has_generated_content) + ) + if fallback_error is None: + raise stream_error + # A native error frame or a refusal the wrapper raised fails the attempt too (a raised error or a + # bridged frame already did). + attempt.report_failure(fallback_error) + await close_superseded() + await attempt.failure_recorded() + retried: Final = await attempt.retry(_mid_stream_retry_trigger(fallback_error)) + if isinstance(retried, BaseException): + last_error: Final = ( + fallback_error + if retried is _mid_stream_retry_trigger(fallback_error) + else _stream_pre_content_error( + retried if isinstance(retried, Exception) else fallback_error, model, fallback_error.llm_provider + ) + ) + async for item in _anthropic_fallback_attempt(attempt, last_error, wrapper): + yield item + return + response, attempted, budget = retried + inner: Final = response + if isinstance(inner, FallbackAwareAnthropicMessagesStream): + inner.attempt.attempted_retries = attempted + inner.attempt.committed_budget = budget + async for item in _anthropic_yield_recovered(inner, wrapper, (attempted, budget)): + yield item + + +async def _anthropic_yield_recovered( + recovered: object, + wrapper: FallbackAwareAnthropicMessagesStream, + retry_counters: tuple[int, int] | None = None, +) -> AsyncGenerator[bytes, None]: + """Hands a retry's or a fallback's response to the client through the wrapper, closing it after.""" + from litellm.llms.anthropic.messages.native.streaming_iterator import ( + aclose_if_supported, + anthropic_messages_response_as_sse_events, + ) + + hidden_params, headers = _prepare_fallback_hidden_params(recovered) + wrapper.merge_fallback_hidden_params(hidden_params, headers) + wrapper.adopt_fallback_source(recovered) + if retry_counters is not None: + add_retry_headers_to_response( + response=wrapper, attempted_retries=retry_counters[0], max_retries=retry_counters[1] + ) + try: + if hasattr(recovered, "__aiter__"): + async for item in cast("AsyncIterator[bytes]", recovered): # cast-ok: __aiter__ checked above + wrapper.follow_source_attribution() + yield item + return + # A complete response for a streaming request is replayed as the SSE lifecycle a stream sends. + for event in anthropic_messages_response_as_sse_events(recovered): # pyright: ignore[reportArgumentType] # the non-streaming shape by elimination + yield event + finally: + with anyio.CancelScope(shield=True), contextlib.suppress(BaseException): + await aclose_if_supported(recovered) + + +async def _anthropic_fallback_attempt( + attempt: StreamAttempt, e: MidStreamFallbackError, wrapper: FallbackAwareAnthropicMessagesStream +) -> AsyncGenerator[bytes, None]: + """Re-enters the fallback walk for a ``/v1/messages`` stream error and yields what the fallback produces.""" + trigger: Final = _fallback_trigger( + e, (litellm_exceptions.ContentPolicyViolationError, litellm_exceptions.ContextWindowExceededError) + ) + try: + fallback_response: Final = await attempt.fallback(trigger) + async for fallback_item in _anthropic_yield_recovered(fallback_response, wrapper): + yield fallback_item + except Exception as fallback_error: # noqa: BLE001 # as the Python wrapper: logged, then the stream's error raises + verbose_router_logger.error("Anthropic messages streaming fallback also failed: %s", fallback_error) + _raise_fallback_failure(fallback_error) + + +# --- Responses API ------------------------------------------------------------------------------ + + +def _responses_stream_holds_event(item: object, held_event_count: int) -> bool: + from litellm.responses.streaming_iterator import PRE_OUTPUT_LIFECYCLE_EVENT_TYPES + + if held_event_count >= MAX_HELD_PRE_OUTPUT_RESPONSES_EVENTS: + return False + return getattr(item, "type", None) in PRE_OUTPUT_LIFECYCLE_EVENT_TYPES + + +def _responses_event_announces_background_run(item: object) -> bool: + return getattr(getattr(item, "response", None), "background", None) is True + + +def _mid_stream_error_before_output(error: MidStreamFallbackError) -> bool: + return not error.generated_content or error.is_pre_first_chunk + + +def _extract_partial_responses_usage(source_iterator: "BaseResponsesAPIStreamingIterator") -> "ResponseAPIUsage | None": + """Partial token usage of a Responses stream that failed, normalized to ``ResponseAPIUsage``.""" + from litellm.responses.litellm_completion_transformation.streaming_iterator import ( + LiteLLMCompletionStreamingIterator, + ) + from litellm.types.llms.openai import ( + ResponseAPIUsage, + ResponseCompletedEvent, + ResponseFailedEvent, + ResponseIncompleteEvent, + ) + + if isinstance(source_iterator, LiteLLMCompletionStreamingIterator): + chunks: Final = source_iterator.collected_chat_completion_chunks + if chunks: + try: + from litellm.main import ( + stream_chunk_builder, # pyright: ignore[reportUnknownVariableType] # untyped chunk list + ) + + built: Final = stream_chunk_builder(chunks=chunks) + chat: Final[object | None] = getattr(built, "usage", None) if built is not None else None + if chat is not None: + prompt: Final = int(getattr(chat, "prompt_tokens", 0) or 0) + completion: Final = int(getattr(chat, "completion_tokens", 0) or 0) + total: Final = int(getattr(chat, "total_tokens", prompt + completion) or (prompt + completion)) + return ResponseAPIUsage(input_tokens=prompt, output_tokens=completion, total_tokens=total) + except Exception as error: # noqa: BLE001 # the builder is best-effort + verbose_router_logger.debug("partial usage of a failed Responses stream: %s", error) + completed: Final = source_iterator.completed_response + if isinstance(completed, (ResponseCompletedEvent, ResponseFailedEvent, ResponseIncompleteEvent)): + return completed.response.usage + return None + + +def _combine_responses_fallback_usage(fallback_item: object, partial_usage: "ResponseAPIUsage") -> None: + from litellm.types.llms.openai import ( + ResponseAPIUsage, + ResponseCompletedEvent, + ResponseFailedEvent, + ResponseIncompleteEvent, + ) + + if not isinstance(fallback_item, (ResponseCompletedEvent, ResponseFailedEvent, ResponseIncompleteEvent)): + return + response: Final = fallback_item.response + if response.usage is None: + return + fb: Final = response.usage + response.usage = ResponseAPIUsage( + input_tokens=(partial_usage.input_tokens or 0) + (fb.input_tokens or 0), + output_tokens=(partial_usage.output_tokens or 0) + (fb.output_tokens or 0), + total_tokens=(partial_usage.total_tokens or 0) + (fb.total_tokens or 0), + ) + + +def _params_sent_to_deployment(attempt: StreamAttempt, deployment_id: object) -> Mapping[str, object]: + """The request as the deployment received it, merged in the order the attempt applies.""" + return attempt.core.params_sent_to_deployment(attempt.request, deployment_id) + + +def wrap_responses_stream( # noqa: C901 # one function per Python wrapper it ports + attempt: StreamAttempt, response: "BaseResponsesAPIStreamingIterator" +) -> "BaseResponsesAPIStreamingIterator": + """``_aresponses_streaming_iterator``: pre-output lifecycle events are held while a fallback can take + over; a mid-stream error falls back through the core, continuing after text.""" + import time + from datetime import datetime, timezone + + from litellm.responses.streaming_iterator import BaseResponsesAPIStreamingIterator, get_openai_response_types + + source_iterator: Final = response + attempt.reports_outcome = True # rebind-ok: the wrapper reports this attempt's outcome + openai_types: Final = get_openai_response_types() + terminal_event_types: Final = ( + openai_types.ResponsesAPIStreamEvents.RESPONSE_COMPLETED, + openai_types.ResponsesAPIStreamEvents.RESPONSE_INCOMPLETE, + openai_types.ResponsesAPIStreamEvents.RESPONSE_FAILED, + ) + + class FallbackResponsesStreamWrapper(BaseResponsesAPIStreamingIterator): + """Subclasses the base iterator only for isinstance checks; iteration is the generator's.""" + + fallback_headers_adopted: bool = False + + def __init__(self, async_generator: AsyncGenerator[object, None]) -> None: # pyright: ignore[reportMissingSuperCall] # no httpx response of its own + self._async_generator: Final = async_generator + self._billing_source: object = source_iterator + self._last_sequence_number = -1 + self.response = getattr(source_iterator, "response", None) # pyright: ignore[reportAttributeAccessIssue] # copied from the source + self.model = getattr(source_iterator, "model", None) # pyright: ignore[reportAttributeAccessIssue] # see above + self.logging_obj = getattr( + source_iterator, "logging_obj", getattr(source_iterator, "litellm_logging_obj", None) + ) # pyright: ignore[reportAttributeAccessIssue] # see above + self.finished = False + self.responses_api_provider_config = getattr(source_iterator, "responses_api_provider_config", None) # pyright: ignore[reportAttributeAccessIssue] # see above + self.completed_response = None + self.start_time = getattr(source_iterator, "start_time", None) or datetime.now(timezone.utc).replace( + tzinfo=None + ) + self._failure_handled = False + self._yielded_first_chunk = False + self._generated_content = "" + self._completed_response_cached = False + self._completed_response_logged = False + self._completed_response_cache_hit = None + self._persist_completed_response_before_logging = True + self._stream_created_time = time.time() + self.litellm_metadata = getattr(source_iterator, "litellm_metadata", None) # pyright: ignore[reportAttributeAccessIssue] # see above + self.custom_llm_provider = getattr(source_iterator, "custom_llm_provider", None) # pyright: ignore[reportAttributeAccessIssue] # see above + self.request_data = getattr(source_iterator, "request_data", {}) or {} + self.call_type = getattr(source_iterator, "call_type", None) # pyright: ignore[reportAttributeAccessIssue] # see above + self._hidden_params = dict(getattr(source_iterator, "_hidden_params", None) or {}) + + def adopt_fallback_headers(self, fallback_response: object) -> _Prepared: + prepared: Final = _prepare_fallback_hidden_params(fallback_response) + self._billing_source = fallback_response + self._hidden_params = {**prepared[0], "additional_headers": prepared[1]} + self.fallback_headers_adopted = True + return prepared + + def partial_response_on_disconnect(self) -> "ResponseIncompleteEvent | None": + from litellm.responses.streaming_iterator import partial_responses_incomplete_event_on_disconnect + + current: Final = partial_responses_incomplete_event_on_disconnect(self._billing_source) + if self._billing_source is source_iterator: + return current + primary: Final = partial_responses_incomplete_event_on_disconnect(source_iterator) + if current is None or primary is None: + return current or primary + if primary.response.usage is not None: + _combine_responses_fallback_usage(current, primary.response.usage) + return current + + def __aiter__(self) -> "FallbackResponsesStreamWrapper": + return self + + @property + def yielded_first_chunk(self) -> bool: + return self._yielded_first_chunk + + async def __anext__(self) -> object: # pyright: ignore[reportIncompatibleMethodOverride] # yields what the stream yields + try: + chunk: Final = await self._async_generator.__anext__() + except StopAsyncIteration: + if self.completed_response is None: + self.completed_response = getattr(source_iterator, "completed_response", None) + raise + if self.completed_response is None and getattr(chunk, "type", None) in terminal_event_types: + self.completed_response = chunk # pyright: ignore[reportAttributeAccessIssue] # the terminal event + self._yielded_first_chunk = True + sequence_number: Final = getattr(chunk, "sequence_number", None) + if not isinstance(sequence_number, int): + return chunk + if not self.fallback_headers_adopted: + self._last_sequence_number = max(self._last_sequence_number, sequence_number) + return chunk + self._last_sequence_number = max(self._last_sequence_number + 1, sequence_number) + if self._last_sequence_number == sequence_number: + return chunk + return chunk.model_copy(update={"sequence_number": self._last_sequence_number}) # pyright: ignore[reportAttributeAccessIssue, reportUnknownMemberType, reportUnknownVariableType] # a response event, the only item with a sequence_number + + async def aclose(self) -> None: + await self._async_generator.aclose() + + def background_stream_can_release(item: object) -> bool: + from litellm.responses.streaming_iterator import ResponsesAPIStreamingIterator + + if ( + type(source_iterator) is not ResponsesAPIStreamingIterator + or source_iterator.custom_llm_provider != "openai" + ): + return False + if not _responses_event_announces_background_run(item): + return False + return not attempt.can_fall_back() + + async def stream_with_fallbacks() -> AsyncGenerator[object, None]: + held: tuple[object, ...] = () # rebind-ok: flushed at first output, dropped on fallback + response_id_exposed = False # rebind-ok: set once, before an early-released event is yielded + try: + async for item in _areported( + attempt, + cast("AsyncIterator[object]", source_iterator), # cast-ok: an async stream + lambda last: _responses_total_tokens(last, source_iterator), + ): + if not response_id_exposed and background_stream_can_release(item): + response_id_exposed = True + if not response_id_exposed and _responses_stream_holds_event(item, len(held)): + held = (*held, item) + continue + for held_event in held: + yield held_event + held = () + yield item + for held_event in held: + yield held_event + except (MidStreamFallbackError, litellm_exceptions.ContextWindowExceededError) as stream_error: + e: Final = ( + _context_window_stream_fallback_error( + stream_error, + source_iterator.model or "", + source_iterator.custom_llm_provider or "", + wrapper.yielded_first_chunk, + ) + if _is_context_window(stream_error) + else cast("MidStreamFallbackError", stream_error) # cast-ok: not a context window error + ) + if response_id_exposed: + _raise_original_mid_stream_error(e) + async with contextlib.aclosing( + _responses_fallback_attempt(attempt, e, source_iterator, wrapper.adopt_fallback_headers, held) + ) as fallback_stream: + async for fallback_item in fallback_stream: + yield fallback_item + except Exception: + for held_event in held: + yield held_event + raise + finally: + with anyio.CancelScope(shield=True): + if hasattr(source_iterator, "aclose"): + try: + await source_iterator.aclose() # pyright: ignore[reportAttributeAccessIssue, reportUnknownMemberType] # checked above + except Exception as exc: # noqa: BLE001 # closing is best-effort + verbose_router_logger.debug("stream_with_fallbacks(aresponses): error closing source: %s", exc) + + wrapper: Final = FallbackResponsesStreamWrapper(stream_with_fallbacks()) + return wrapper + + +def _responses_continuation_request( + attempt: StreamAttempt, e: MidStreamFallbackError, source_iterator: "BaseResponsesAPIStreamingIterator" +) -> Mapping[str, object] | None: + """The request a fallback continues from once a Responses stream broke after text, or None before + output. Raises the primary's error when the output is constrained, a tool call began, a content + filter cut it, or the output ceiling is spent.""" + from litellm.litellm_core_utils.token_counter import ( + token_counter, # pyright: ignore[reportUnknownVariableType] # untyped tokenizer argument + ) + from litellm.router_utils.pre_call_checks.continuation_prefill_check import ( + RESPONSES_CONTINUATION_KWARG, + ResponsesContinuation, + constrains_responses_continuation, + continuation_output_ceilings, + responses_continuation_input, + ) + + if _mid_stream_error_before_output(e): + return None + if source_iterator.started_tool_call or isinstance( + e.original_exception, litellm_exceptions.ContentPolicyViolationError + ): + _raise_original_mid_stream_error(e) + cut_deployment_id: Final = (get_hidden_params(source_iterator) or {}).get("model_id") + if constrains_responses_continuation(_params_sent_to_deployment(attempt, cut_deployment_id)): + _raise_original_mid_stream_error(e) + request: Final = attempt.request + earlier: Final = request.get(RESPONSES_CONTINUATION_KWARG) + continuation: Final = ( + earlier.extended(e.generated_content) + if isinstance(earlier, ResponsesContinuation) + else ResponsesContinuation(request.get("input"), e.generated_content) + ) + emitted_tokens: Final = token_counter(model=e.model or "", text=e.generated_content) + reduced: Final = continuation_output_ceilings(request, emitted_tokens) + if reduced is None: + _raise_original_mid_stream_error(e) + return { + **request, + **reduced, + RESPONSES_CONTINUATION_KWARG: continuation, + "input": responses_continuation_input(continuation), + } + + +async def _responses_fallback_attempt( + attempt: StreamAttempt, + e: MidStreamFallbackError, + source_iterator: "BaseResponsesAPIStreamingIterator", + adopt_headers: Callable[[object], _Prepared], + held: tuple[object, ...], +) -> AsyncGenerator[object, None]: + """Re-enters the fallback walk for a mid-stream Responses error. The primary's held lifecycle events + reach the client only when no fallback lands.""" + from litellm.router_utils.pre_call_checks.continuation_prefill_check import ResponsesContinuationSeam + + continued: Final = _responses_continuation_request(attempt, e, source_iterator) + seam: Final = ResponsesContinuationSeam(e.generated_content if continued is not None else "") + partial_usage: Final = _extract_partial_responses_usage(source_iterator) + trigger: Final = _fallback_trigger( + e, (litellm_exceptions.ContentPolicyViolationError, litellm_exceptions.ContextWindowExceededError) + ) + fallback_response: object = None # rebind-ok: set once a fallback answered, closed at the end + fallback_yielded = False # rebind-ok: flipped on the first fallback item + try: + fallback_response = await attempt.fallback(trigger, continued) # rebind-ok: rebound as the stream advances + prepared: Final = adopt_headers(fallback_response) + if hasattr(fallback_response, "__aiter__"): + spliced: Final = cast( # cast-ok: the fallback's events + "AsyncIterator[object]", + seam.splice(fallback_response), # pyright: ignore[reportArgumentType] # an async stream + ) + async for fallback_item in spliced: + _apply_fallback_hidden_params_to_item(fallback_item, prepared) + if partial_usage is not None: + _combine_responses_fallback_usage(fallback_item, partial_usage) + fallback_yielded = True + yield fallback_item + else: + fallback_yielded = True # rebind-ok: rebound as the stream advances + yield fallback_response + except Exception as fallback_error: # noqa: BLE001 # as the Python wrapper: logged, then the stream's error raises + verbose_router_logger.error("Responses streaming fallback also failed: %s", fallback_error) + if not fallback_yielded: + for held_event in held: + yield held_event + _raise_fallback_failure(fallback_error) + finally: + if fallback_response is not None and hasattr(fallback_response, "aclose"): + with anyio.CancelScope(shield=True): + try: + await fallback_response.aclose() # pyright: ignore[reportAttributeAccessIssue, reportUnknownMemberType] # checked above + except Exception as exc: # noqa: BLE001 # closing is best-effort + verbose_router_logger.debug("stream_with_fallbacks(aresponses): error closing fallback: %s", exc) + + +# --- the attempt hook --------------------------------------------------------------------------- + + +async def open_attempt_stream(attempt: StreamAttempt, response: object) -> object: + """What an attempt returns once its stream is wrapped: a chat stream's first chunk is read here, so + a stream that fails before any token raises inside the attempt.""" + from litellm.responses.streaming_iterator import BaseResponsesAPIStreamingIterator + + if attempt.op == "acompletion" and isinstance(response, CustomStreamWrapper): + if response.completion_stream is None and response.make_call is not None: # pyright: ignore[reportUnknownMemberType] # untyped attributes + await response.fetch_stream() + first_chunk: Final = await _aread_first_chunk(response) + held_slot: Final = attempt.slot.pop_all() if attempt.slot is not None else contextlib.ExitStack() + return wrap_chat_stream(attempt, response, first_chunk, held_slot) + if not attempt.kwargs.get("stream"): + return response + if attempt.op == "aresponses" and isinstance(response, BaseResponsesAPIStreamingIterator): + return wrap_responses_stream(attempt, response) + if attempt.op in ("aanthropic_messages", "anthropic_messages") and hasattr(response, "__aiter__"): + return wrap_anthropic_stream(attempt, cast("AsyncIterator[bytes]", response)) # cast-ok: a byte stream + return response + + +def open_attempt_stream_sync(attempt: StreamAttempt, response: object) -> object: + if attempt.op == "completion" and isinstance(response, CustomStreamWrapper): + if response.completion_stream is None and response.make_call is not None: # pyright: ignore[reportUnknownMemberType] # untyped attributes + response.fetch_sync_stream() + return wrap_chat_stream_sync(attempt, response, _read_first_chunk(response)) + return response + + +def shape_attempt_kwargs(kwargs: Mapping[str, object], deployment: Mapping[str, object]) -> Mapping[str, object]: + """A Responses continuation is shaped for the deployment the fallback landed on.""" + from litellm.router_utils.pre_call_checks.continuation_prefill_check import ( + RESPONSES_CONTINUATION_KWARG, + responses_continuation_input_for, + ) + + if RESPONSES_CONTINUATION_KWARG not in kwargs: + return kwargs + shaped: Final = without(kwargs, RESPONSES_CONTINUATION_KWARG) + shaped["input"] = responses_continuation_input_for(kwargs, deployment) + return shaped + + +__all__ = [ + "MAX_BUFFERED_PRE_CONTENT_ANTHROPIC_CHUNKS", + "MAX_HELD_PRE_OUTPUT_RESPONSES_EVENTS", + "FallbackAwareAnthropicMessagesStream", + "FallbackAwareStreamWrapper", + "StreamAttempt", + "open_attempt_stream", + "open_attempt_stream_sync", + "shape_attempt_kwargs", +] diff --git a/litellm/router_rust/invoke/walk.py b/litellm/router_rust/invoke/walk.py new file mode 100644 index 00000000000..4dab8edc07f --- /dev/null +++ b/litellm/router_rust/invoke/walk.py @@ -0,0 +1,398 @@ +""" +The request metadata the Python router stamps as it walks retries and fallbacks, kept for a +``RustCore`` call. router-core decides the walk and reports it (each attempt's ``scope``, +``max_retries`` and ``model_group_size``, and the ``fallback`` and ``retry_logged`` events); this +module only holds what each level's metadata dict holds in the Python router and lays it over the +attempt's kwargs. + +In the Python router one metadata dict serves a level's retries, and every fallback hop copies +its parent's (``run_async_fallback``'s ``_hop_metadata``), so a hop's breadcrumbs never reach +its parent. Each scope here is one such dict. The root level's dict is the caller's own (a proxy +reads a failed request's walk off it), so a fresh call's root writes reach that dict as well. +""" + +from collections.abc import Callable, Coroutine, Mapping, Sequence +from typing import TYPE_CHECKING, Final + +from litellm.litellm_core_utils.core_helpers import as_str_mapping + +if TYPE_CHECKING: + from opentelemetry.trace import Span + + from litellm.router_api import RoutingContext + from litellm.router_utils.fallback_event_handlers import AttemptedFallbackTargets + +# `RETRY_BREADCRUMB_LIMIT`: how many failures `previous_models` keeps. +RETRY_BREADCRUMB_LIMIT: Final = 4 + +# The keys the walk stamps on each attempt's metadata. +_WALK_KEYS: Final = ( + "original_model_group", + "attempted_fallbacks", + "attempted_retries", + "max_retries", + "model_group_size", + "previous_models", + "request_retry_count", +) +# What a breadcrumb reads from the last attempt a scope made. +_SERVED_KEYS: Final = ("model_info", "deployment") +# The walk state `run_async_fallback` hands every hop as kwargs. +_HOP_KWARGS: Final = ("fallback_depth", "max_fallbacks", "attempted_targets", "include_fallback_errors") + + +class _Scope: + __slots__ = ("kwargs", "metadata") + + def __init__( + self, + metadata: dict[str, object], # mutable-ok: the level's metadata, stamped as the walk runs + kwargs: Mapping[str, object], + ) -> None: + self.metadata: Final = metadata + self.kwargs: Final = kwargs + + +def _breadcrumb( + metadata: Mapping[str, object], model: str, error: BaseException, params_base_model: Callable[[str], str | None] +) -> dict[str, object]: # mutable-ok: a fresh RetryAttemptRecord + """``_core_log_retry``'s record of the failure.""" + from litellm.litellm_core_utils.get_litellm_params import get_base_model_from_litellm_call_metadata + + model_info: Final = as_str_mapping(metadata.get("model_info")) or {} + deployment_id: Final = model_info.get("id") + known_id: Final = deployment_id if isinstance(deployment_id, str) else None + base_model: Final = get_base_model_from_litellm_call_metadata(metadata=metadata) or ( + params_base_model(known_id) if known_id is not None else None + ) + deployment: Final = metadata.get("deployment") + attempted: Final = metadata.get("attempted_retries") + return { + "model_group": model, + "deployment_id": known_id, + "model": deployment if isinstance(deployment, str) else None, + "base_model": base_model, + "exception_type": type(error).__name__, + "exception_string": str(error), + "attempted_retries": attempted if type(attempted) is int else None, + } + + +def _kept_breadcrumbs(earlier: object) -> tuple[object, ...]: + """The breadcrumbs a new one joins: the latest ``RETRY_BREADCRUMB_LIMIT - 1``.""" + if not isinstance(earlier, list | tuple): + return () + stamped: Final[tuple[object, ...]] = tuple(earlier) # pyright: ignore[reportUnknownArgumentType, reportUnknownVariableType] # the breadcrumbs as stamped + return stamped[-(RETRY_BREADCRUMB_LIMIT - 1) :] + + +def log_retry_into( + kwargs: dict[str, object], # mutable-ok: the request kwargs whose metadata the record joins + error: Exception, + params_base_model: Callable[[str], str | None], +) -> dict[str, object]: # mutable-ok: the same kwargs, as ``_core_log_retry`` returns them + """ + ``_core_log_retry`` (``RoutingContext.log_retry``): the failure joins the request metadata's + ``previous_models`` and counts toward ``request_retry_count``, with a ``litellm.routing.retry`` + phase event when a pre-routing strategy picked the group. + """ + from litellm.integrations.otel.runtime import phase_event + + field: Final = "litellm_metadata" if "litellm_metadata" in kwargs else "metadata" + metadata: Final = kwargs[field] + view: Final = as_str_mapping(metadata) + if view is None: + # Python reads it with ``.get``. + raise AttributeError(f"{type(metadata).__name__!r} object has no attribute 'get'") + model_group: Final = kwargs.get("model") + record: Final = _breadcrumb(view, "", error, params_base_model) + record["model_group"] = model_group if isinstance(model_group, str) else None + deployment_id: Final = record["deployment_id"] + breadcrumbs: Final = (*_kept_breadcrumbs(view.get("previous_models")), record) + earlier: Final = view.get("request_retry_count") + request_retry_count: Final = (earlier if type(earlier) is int and 0 <= earlier else 0) + 1 + decision: Final = view.get("routing_decision") + router_name: Final = (as_str_mapping(decision) or {}).get("router_model_name") + if isinstance(router_name, str): + phase_event( + "litellm.routing.retry", + { + "litellm.routing.router_model_name": router_name, + "litellm.retry.count": request_retry_count, + "error.type": type(error).__name__, + **({"litellm.deployment.model_group": model_group} if isinstance(model_group, str) else {}), + **({"litellm.deployment.id": deployment_id} if isinstance(deployment_id, str) else {}), + }, + ) + metadata["previous_models"] = breadcrumbs # pyright: ignore[reportIndexIssue] # the logging object already holds this dict + metadata["request_retry_count"] = request_retry_count # pyright: ignore[reportIndexIssue] # same dict, read by the cap + return kwargs + + +class WalkMetadata: + """The metadata each level of one call's walk holds, by scope.""" + + __slots__ = ("_attempted", "_field", "_live", "_params_base_model", "_request", "_scopes") + + def __init__( + self, + request: Mapping[str, object], + field: str, + params_base_model: Callable[[str], str | None], + *, + live: bool = False, + ) -> None: + model: Final = request.get("model") + root: Final[dict[str, object]] = {"attempted_fallbacks": 0} # mutable-ok: the root level's metadata + if model is not None: + root["original_model_group"] = model + self._request: Final = request + self._field: Final = field + self._params_base_model: Final = params_base_model + self._scopes: Final[dict[int, _Scope]] = {0: _Scope(root, {})} # mutable-ok: filled as the walk hops + self._attempted: AttemptedFallbackTargets | None = None + # A call whose retry and fallback wrappers stamp the request itself; a stream's recovery + # re-enters below them. + self._live: Final = live + if live: + self._stamp_fallback_entry() + + def _stamp_fallback_entry(self) -> None: + """The head of ``async_function_with_fallbacks``, in the request's own buckets. ``run_async_fallback``'s + hops carry their ``attempted_targets`` and never reach it.""" + from litellm.router_utils.fallback_event_handlers import AttemptedFallbackTargets + + request: Final = self._request + if isinstance(request.get("attempted_targets"), AttemptedFallbackTargets): + return + sibling: Final = request.get("metadata" if self._field == "litellm_metadata" else "litellm_metadata") + if isinstance(sibling, dict): + sibling.pop("attempted_fallbacks", None) # pyright: ignore[reportUnknownMemberType] # the caller's metadata dict + sibling.pop("original_model_group", None) # pyright: ignore[reportUnknownMemberType] # see above + bucket: Final = request.get(self._field) + if not isinstance(bucket, dict): + return + bucket["attempted_fallbacks"] = 0 + model: Final = request.get("model") + if model is not None: + bucket["original_model_group"] = model + + def _retry_bucket(self, scope: _Scope) -> dict[str, object] | None: # mutable-ok: the request's own dict + """``_metadata`` of ``async_function_with_retries`` when it is a dict of the request's: the + root level's, or a hop's when the bucket is not the one ``run_async_fallback`` copies.""" + if not self._live: + return None + name: Final = "litellm_metadata" if "litellm_metadata" in self._request else "metadata" + if scope is not self._scopes[0] and name == self._field: + return None + bucket: Final = self._request.get(name) + # `... or {}`: an empty bucket leaves the stamps on a dict of their own. + return bucket if isinstance(bucket, dict) and bucket else None # pyright: ignore[reportUnknownVariableType] # the request's untyped metadata dict + + def _write_through(self, scope: _Scope, values: Mapping[str, object]) -> None: + bucket: Final = self._retry_bucket(scope) + if bucket is not None: + bucket.update(values) + + def retries_opened(self, event: Mapping[str, object]) -> None: + """The head of ``async_function_with_retries``: the group's size and the retry tracking, before any pick.""" + scope: Final = self._scope(event.get("scope")) + stamped: Final = {"attempted_retries": 0, "max_retries": event.get("max_retries")} + scope.metadata.update(stamped) + self._write_through(scope, stamped) + bucket: Final = self._retry_bucket(scope) + if bucket is not None and isinstance(bucket.get("model_group"), str): + bucket["model_group_size"] = event.get("model_group_size") + + def retry_budget(self, event: Mapping[str, object]) -> None: + """``_metadata["max_retries"] = num_retries``, once a level's first failure settled its overrides.""" + scope: Final = self._scope(event.get("scope")) + stamped: Final = {"max_retries": event.get("max_retries")} + scope.metadata.update(stamped) + self._write_through(scope, stamped) + + def _scope(self, value: object) -> _Scope: + found: Final = self._scopes.get(value) if isinstance(value, int) else None + return found if found is not None else self._scopes[0] + + def targets(self) -> "AttemptedFallbackTargets": + """The walk's ``attempted_targets``: the targets tried and the ones skipped, and why.""" + from litellm.router_utils.fallback_event_handlers import AttemptedFallbackTargets + + if self._attempted is None: + self._attempted = AttemptedFallbackTargets() + return self._attempted + + def hop(self, event: Mapping[str, object]) -> None: + """A ``fallback`` event: the new level copies its parent's metadata and kwargs.""" + parent: Final = self._scope(event.get("parent")) + attempted: Final = self.targets() + for key in (event.get("from"), event.get("to")): + if isinstance(key, str): + attempted.record(key) + metadata: Final = dict(parent.metadata) + original: Final = metadata.pop("original_model_group", self._request.get("model")) + metadata.pop("model_group", None) + metadata.pop("attempted_fallbacks", None) + metadata["original_model_group"] = original + # A dict target's non-string `model` is what `kwargs["model"]` holds on the hop. + metadata["model_group"] = event["unroutable"] if "unroutable" in event else event.get("to") + metadata["attempted_fallbacks"] = event.get("depth") + kwargs: Final = { + **parent.kwargs, + "fallback_depth": event.get("depth"), + "max_fallbacks": event.get("max_fallbacks"), + "attempted_targets": attempted, + } + if self._request.get("include_fallback_errors") is True: + kwargs["include_fallback_errors"] = True + scope: Final = event.get("scope") + if isinstance(scope, int): + self._scopes[scope] = _Scope(metadata, kwargs) + + def hop_kwargs(self, event: Mapping[str, object]) -> dict[str, object]: # mutable-ok: a fresh dict for the loggers + """The kwargs the level held after a hop (``kwargs["model"] = target``), as the + fallback loggers see them.""" + scope: Final = self._scope(event.get("scope")) + model: Final = scope.metadata.get("model_group", event.get("to")) + return {**self._request, **scope.kwargs, "model": model, self._field: dict(scope.metadata)} + + def log_retry(self, event: Mapping[str, object], error: BaseException | None) -> None: + """A ``retry_logged`` event: the failure joins the scope's ``previous_models``.""" + scope: Final = self._scope(event.get("scope")) + metadata: Final = scope.metadata + attempted: Final = event.get("attempted_retries") + if isinstance(attempted, int): + metadata["attempted_retries"] = attempted + self._write_through(scope, {"attempted_retries": attempted}) + if error is None: + return + record: Final = _breadcrumb(metadata, str(event.get("model")), error, self._params_base_model) + earlier: Final = metadata.get("previous_models") + kept: Final = _kept_breadcrumbs(earlier) + count: Final = metadata.get("request_retry_count") + metadata["previous_models"] = (*kept, record) + metadata["request_retry_count"] = (count if type(count) is int and count >= 0 else 0) + 1 + self._write_through( + scope, + {"previous_models": metadata["previous_models"], "request_retry_count": metadata["request_retry_count"]}, + ) + + def deployment_that_just_failed(self, step: Mapping[str, object]) -> str | None: + """``_deployment_that_just_failed``: the deployment the metadata an attempt inherits names, + before the attempt stamps its own; ``None`` when the walk's level has run none yet.""" + model_info: Final = as_str_mapping(self._scope(step.get("scope")).metadata.get("model_info")) + model_id: Final = model_info.get("id") if model_info is not None else None + return model_id if isinstance(model_id, str) else None + + def request_retry_count(self, step: Mapping[str, object]) -> object: + """The level's ``request_retry_count``, as the retry loop has counted it so far.""" + return self._scope(step.get("scope")).metadata.get("request_retry_count") + + def stamp(self, step: Mapping[str, object], kwargs: Mapping[str, object]) -> Mapping[str, object]: + """The attempt's kwargs with its level's walk metadata and hop kwargs laid over them.""" + scope: Final = self._scope(step.get("scope")) + metadata: Final = scope.metadata + metadata["attempted_retries"] = step.get("retry") + metadata["max_retries"] = step.get("max_retries") + metadata["model_group_size"] = step.get("model_group_size") + self._write_through( + scope, {"attempted_retries": metadata["attempted_retries"], "max_retries": metadata["max_retries"]} + ) + served: Final = dict(as_str_mapping(kwargs.get(self._field)) or {}) + served.update({key: metadata[key] for key in _WALK_KEYS if key in metadata}) + metadata.update({key: served.get(key) for key in _SERVED_KEYS}) + return {**kwargs, **{key: scope.kwargs[key] for key in _HOP_KWARGS if key in scope.kwargs}, self._field: served} + + +async def _log_failed_fallbacks(ctx: "RoutingContext", detail: str, parent_otel_span: "Span | None") -> None: + from litellm._logging import verbose_router_logger + from litellm.router_utils.cooldown_handlers import async_get_cooldown_deployments_with_debug_info + + cooldown_info: Final = await async_get_cooldown_deployments_with_debug_info( + ctx=ctx, parent_otel_span=parent_otel_span + ) + verbose_router_logger.error( + "litellm.router.py::async_function_with_fallbacks() - " + "Error occurred while trying to do fallbacks - %s\n" + "Debug Information:\nCooldown Deployments=%s", + detail, + cooldown_info, + ) + + +class FallbackOutcomes: + """What ``run_async_fallback`` does when a hop returns: it logs the hop to the process's + custom loggers and, under ``include_fallback_errors``, lists the level's errors (its own + first, then each hop that failed) in the response's ``x-litellm-fallback-errors``.""" + + __slots__ = ("_failed", "_header_errors", "_include", "_pending") + + def __init__(self, request: Mapping[str, object]) -> None: + self._include: Final = request.get("include_fallback_errors", False) is True + self._failed: Final[dict[object, list[object]]] = {} # mutable-ok: each level's failed hops, in order + self._header_errors: Final[list[list[object]]] = [] # mutable-ok: each successful hop's errors, innermost first + self._pending: Final[list[Coroutine[object, object, None]]] = [] # mutable-ok: logger calls not yet awaited + + def done( + self, + event: Mapping[str, object], + level_error: BaseException | None, + hop_error: BaseException | None, + kwargs: Mapping[str, object], + ) -> None: + from litellm.litellm_core_utils.add_retry_fallback_headers import get_fallback_error_info + from litellm.router_utils import fallback_event_handlers as handlers + + if level_error is None or not isinstance(level_error, Exception): + return + parent: Final = event.get("parent") + failed: Final = self._failed.setdefault(parent, []) + original: Final = str(event.get("from")) + if event.get("succeeded") is True: + if self._include: + self._header_errors.append([get_fallback_error_info(level_error), *failed]) + self._pending.append( + handlers.log_success_fallback_event( # pyright: ignore[reportUnknownMemberType] # takes the kwargs as an untyped dict + original_model_group=original, kwargs=dict(kwargs), original_exception=level_error + ) + ) + return + if isinstance(hop_error, Exception): + failed.append(get_fallback_error_info(hop_error)) + self._pending.append( + handlers.log_failure_fallback_event( # pyright: ignore[reportUnknownMemberType] # see above + original_model_group=original, kwargs=dict(kwargs), original_exception=level_error + ) + ) + + def failed_fallbacks(self, ctx: "RoutingContext", detail: str, parent_otel_span: "Span | None") -> None: + """The ``except`` around the regular fallbacks of ``async_function_with_fallbacks_common_utils``: + ``detail`` is the failure's text as it read then, the cooldowns are read when the pending calls run.""" + self._pending.append(_log_failed_fallbacks(ctx, detail, parent_otel_span)) + + def take_pending(self) -> tuple[Coroutine[object, object, None], ...]: + pending: Final = tuple(self._pending) + self._pending.clear() + return pending + + def add_headers(self, response: object, attempted_fallbacks: int) -> object: + """``add_fallback_headers_to_response`` for each successful hop, innermost first.""" + from litellm.litellm_core_utils.add_retry_fallback_headers import add_fallback_headers_to_response + + headed: object = response # rebind-ok: gains each level's errors + for errors in self._header_errors: + headed = add_fallback_headers_to_response( + response=headed, + attempted_fallbacks=attempted_fallbacks, + fallback_errors=errors, # pyright: ignore[reportArgumentType] # FallbackErrorInfo dicts + ) + return headed + + +async def await_all(pending: Sequence[Coroutine[object, object, None]]) -> None: + for call in pending: + await call + + +__all__ = ["RETRY_BREADCRUMB_LIMIT", "FallbackOutcomes", "WalkMetadata", "await_all"] diff --git a/litellm/router_rust/invoke/websocket.py b/litellm/router_rust/invoke/websocket.py new file mode 100644 index 00000000000..e7baeb92345 --- /dev/null +++ b/litellm/router_rust/invoke/websocket.py @@ -0,0 +1,127 @@ +""" +Responses websocket turns on the Rust router core (``async_admit_deployment_turn``): the socket is +opened on one deployment, and the core hands the handler a ``turn_admission`` callable that re-runs +that deployment's admission before every later turn: the io-token turn budget, the deployment +budget limiter, the routing strategy (on that deployment alone) and the pre-call checks. + +Which strategy and selector apply is the Rust core's routing context (``routing_context``); the +selectors, limiters and checks are the live Python objects the host keeps. +""" + +from collections.abc import Callable, Mapping +from functools import partial +from types import MappingProxyType +from typing import TYPE_CHECKING, Final + +from litellm.litellm_core_utils.core_helpers import without +from litellm.router_rust.invoke.shaping import as_dict + +if TYPE_CHECKING: + from litellm.router_rust.invoke.endpoints import EndpointCore + + +def turn_admission( + core: "EndpointCore", model_group: str, deployment: Mapping[str, object], kwargs: Mapping[str, object] +) -> Callable[[], object]: + """The ``turn_admission`` the handler calls before each later turn.""" + return partial(admit_deployment_turn, core, model_group, deployment, kwargs) + + +def _deployment_as_configured_now(core: "EndpointCore", deployment: Mapping[str, object]) -> Mapping[str, object]: + model_info: Final = deployment.get("model_info") + model_id: Final = as_dict(model_info).get("id") if isinstance(model_info, dict) else None # pyright: ignore[reportUnknownArgumentType] # the configured model_info dict + current: Final = core.configured_deployment(model_id) if isinstance(model_id, str) else None + return deployment if current is None else current.model_dump(exclude_none=True) + + +def _budget_limiter(core: "EndpointCore") -> object: + """The router's deployment-budget callback; the Rust core declines every config that has one.""" + from litellm.router_rust.core import RustCoreUnsupported + + try: + return core.registry().plugins("budget_limiter") + except RustCoreUnsupported: + return None + + +async def admit_deployment_turn( + core: "EndpointCore", model_group: str, opened_on: Mapping[str, object], request_kwargs: Mapping[str, object] +) -> Mapping[str, object]: + """`async_admit_deployment_turn`.""" + from litellm import exceptions as litellm_exceptions + from litellm.router_utils.pre_call_checks.io_token_rate_limit_check import ( + async_admit_io_token_turn, + without_io_token_limits, + ) + from litellm.router_utils.pre_call_checks.model_rate_limit_check import ModelRateLimitingCheck + from litellm.utils import get_utc_datetime + + deployment: Final = _deployment_as_configured_now(core, opened_on) + rate_limit_check: Final = next( + (callback for callback in core.routing_callbacks() if isinstance(callback, ModelRateLimitingCheck)), + None, + ) + io_token_admission: Final[Mapping[str, object]] = ( + MappingProxyType({}) + if rate_limit_check is None + else await async_admit_io_token_turn(rate_limit_check.dual_cache, deployment, get_utc_datetime()) # pyright: ignore[reportArgumentType] # the configured deployment dict + ) + budget_limiter: Final[object] = _budget_limiter(core) + if budget_limiter is not None: + try: + await budget_limiter.async_filter_deployments( # pyright: ignore[reportAttributeAccessIssue, reportUnknownMemberType] # RouterBudgetLimiting + model=model_group, + healthy_deployments=[dict(deployment)], + messages=None, + request_kwargs=dict(request_kwargs), + ) + except ValueError as e: + raise litellm_exceptions.RateLimitError( + message=str(e), + llm_provider="", + model=model_group, + category=litellm_exceptions.RateLimitErrorCategory.LITELLM_RATE_LIMIT, + ) from e + await _admit_turn_through_routing_strategy(core, model_group, deployment, request_kwargs) + await core.router.async_routing_strategy_pre_call_checks( + deployment=dict(without_io_token_limits(deployment)), parent_otel_span=None + ) + return io_token_admission + + +async def _admit_turn_through_routing_strategy( + core: "EndpointCore", model_group: str, deployment: Mapping[str, object], request_kwargs: Mapping[str, object] +) -> None: + """`_admit_turn_through_routing_strategy`: the group's strategy picks among this deployment alone.""" + from litellm import exceptions as litellm_exceptions + from litellm.router_api import RoutingContext + from litellm.router_rust.steps.selection import pick_async + from litellm.router_utils.handle_error import async_raise_no_deployment_exception + + routing_kwargs: Final = without(request_kwargs, "litellm_logging_obj") + override: Final = routing_kwargs.get("routing_strategy") + strategy, scope, group = core.handle.routing_context(model_group, override if isinstance(override, str) else None) + selectors: Final = core.selectors + selector: Final = selectors.selector_for(strategy, scope, group, routing_kwargs) + if strategy == "simple-shuffle" or selector is None: + return + selected: Final = await pick_async( + strategy, selector, model_group, [dict(deployment)], None, None, dict(routing_kwargs) + ) + if selected is None: + no_deployment: Final = await async_raise_no_deployment_exception( + ctx=RoutingContext(core), # pyright: ignore[reportArgumentType] # the error reads registry() and health() + model=model_group, + parent_otel_span=None, + ) + raise litellm_exceptions.RateLimitError( + message=str(no_deployment), + llm_provider="", + model=model_group, + category=litellm_exceptions.RateLimitErrorCategory.LITELLM_RATE_LIMIT, + ) from no_deployment + if selectors.is_override(strategy, selector): + await selector.async_pre_call_check(dict(deployment), None) # pyright: ignore[reportAttributeAccessIssue, reportUnknownMemberType] # every selector is a CustomLogger + + +__all__ = ["admit_deployment_turn", "turn_admission"] diff --git a/litellm/router_rust/steps/README.md b/litellm/router_rust/steps/README.md new file mode 100644 index 00000000000..9b5c68bdba9 --- /dev/null +++ b/litellm/router_rust/steps/README.md @@ -0,0 +1,12 @@ +# steps + +Answers to router-core's host steps that need live Python objects. router-core makes the routing decision; these modules supply the Python-only facts or run the Python-only objects it asks about. + +- `selection`: stateful strategies (latency, tpm/rpm, cost, usage) on `select` steps +- `prerouting`: auto, complexity, adaptive and quality routers, and session binding, on `pre_route` steps +- `precall`: token counts, context windows and param support for pre-call checks +- `plugins`: `Router(plugins=[...])` objects +- `callbacks`: the optional pre-call checks (budgets, affinity, rate limits) as `CustomLogger`s +- `access`: the caller facts read from auth metadata + +A Rust host answers these steps natively, or declines the config at build. diff --git a/litellm/router_rust/steps/__init__.py b/litellm/router_rust/steps/__init__.py new file mode 100644 index 00000000000..c47505e7468 --- /dev/null +++ b/litellm/router_rust/steps/__init__.py @@ -0,0 +1 @@ +"""Answers to router-core host steps that need live Python objects.""" diff --git a/litellm/router_rust/steps/access.py b/litellm/router_rust/steps/access.py new file mode 100644 index 00000000000..087a87e2c6b --- /dev/null +++ b/litellm/router_rust/steps/access.py @@ -0,0 +1,85 @@ +"""The caller facts the Rust core routes on, read from the request's auth metadata: whether the +caller is a proxy admin, the models its key and team allow (which narrow an access-group-only +caller to its groups' deployments) and, for a proxy admin, the model the request names.""" + +from __future__ import annotations + +from collections.abc import Mapping +from typing import Final + +from pydantic import TypeAdapter + +from litellm.litellm_core_utils.core_helpers import as_str_mapping +from litellm.litellm_core_utils.routing_common_utils import is_proxy_admin_request, request_user_api_key_auth + +_NAMES: Final = TypeAdapter(tuple[object, ...]) + + +def _names(value: object) -> tuple[str, ...]: + names: Final = _NAMES.validate_python(value) if isinstance(value, list | tuple | set | frozenset) else () + return tuple(str(name) for name in names) + + +def _admin_model(request: Mapping[str, object]) -> object: + """The model ``filter_team_based_models`` compares team deployments against.""" + metadata: Final = as_str_mapping(request.get("metadata")) or {} + litellm_metadata: Final = as_str_mapping(request.get("litellm_metadata")) or {} + return request.get("model") or metadata.get("model_group") or litellm_metadata.get("model_group") + + +def caller_fields(request: Mapping[str, object]) -> Mapping[str, object]: + """``proxy_admin``, ``allowed_models`` and ``admin_model`` for the Rust request document.""" + auth: Final[object] = request_user_api_key_auth(request) + if auth is None: + return {} + allowed: Final = sorted({*_names(getattr(auth, "models", None)), *_names(getattr(auth, "team_models", None))}) + admin: Final = is_proxy_admin_request(request) + model: Final = _admin_model(request) if admin else None + return { + **({"allowed_models": allowed} if allowed else {}), + **({"proxy_admin": True} if admin else {}), + **({"admin_model": model} if isinstance(model, str) else {}), + } + + +def _router_weights_error(raw: object) -> BaseException: + """The error ``validate_router_weights`` raises for weights the Rust core refused.""" + from pydantic import ValidationError + + from litellm.types.router_weights import validate_router_weights + + try: + _ = validate_router_weights(raw) + except ValidationError as error: + return error + return RuntimeError(f"router-core refused router weights the validator accepts: {raw!r}") + + +def selection_refusal(step: Mapping[str, object], model: str, request: Mapping[str, object]) -> BaseException | None: + """The error for a pick the Rust core refused over who the caller is or what it asked for: + a public name several teams share, deployments reserved for other teams, invalid request + router weights or a non-numeric deployment weight. None for any other failure.""" + import litellm.exceptions as litellm_exceptions + + kind: Final = step.get("error_kind") + named: Final = str(step.get("error_model") or model) + if kind == "multiple_teams": + return litellm_exceptions.BadRequestError( + message=( + f"Model name '{named}' matches deployments from multiple teams. " + "Specify the deployment ID directly to disambiguate." + ), + model=named, + llm_provider="", + ) + if kind == "reserved": + return litellm_exceptions.BadRequestError( + message=f"Deployment {named} is reserved for another team until {step.get('error_detail')}", + model=named, + llm_provider="", + ) + if kind == "invalid_router_weights": + return _router_weights_error(request.get("_router_weights")) + if kind == "non_numeric_weight": + return TypeError(f"Deployment {step.get('error_detail')} must be numeric") + return None diff --git a/litellm/router_rust/steps/callbacks.py b/litellm/router_rust/steps/callbacks.py new file mode 100644 index 00000000000..b906a496550 --- /dev/null +++ b/litellm/router_rust/steps/callbacks.py @@ -0,0 +1,650 @@ +""" +Host-side optional pre-call checks for ``RustCore``: the ``CustomLogger`` checks a router builds +from ``optional_pre_call_checks``, ``provider_budget_config``, per-deployment budgets and +``model_group_affinity_config``, registered in the host's ``callbacks`` list exactly as the Python +router's ``add_optional_pre_call_checks`` registers them, since they are features users configure. + +The Rust core asks the host to run every routing callback's ``async_filter_deployments`` through a +``callback_filters`` select stage (``async_callback_filter_deployments``) and decides itself what +the surviving deployments or the raised error mean: a pick with nothing left, a retry, a fallback. + +This module also keeps the process-wide holder rules (``litellm/router_utils/global_callback_holders``): +a retiring router hands the checks it vacates to a survivor, a runtime removal keeps a check +another live router still has, and ``observe(request_scoped=True)`` takes every routing callback +off the host lists while the router keeps running them (``_routing_callbacks``). It also answers +which outcome callbacks the Python router would still have running for a request +(``OptionalChecks.outcome_callbacks``), which router-core records by. It never reaches a Python +router core. +""" + +from collections.abc import Iterable, Mapping, Sequence +from itertools import chain +from typing import TYPE_CHECKING, Final +from weakref import WeakSet + +from litellm._logging import verbose_router_logger +from litellm.caching.dual_cache import DualCache +from litellm.caching.in_memory_cache import InMemoryCache +from litellm.constants import PROMPT_CACHE_PIN_CACHE_MAX_SIZE +from litellm.integrations.custom_logger import CustomLogger +from litellm.litellm_core_utils.core_helpers import as_str_mapping +from litellm.litellm_core_utils.routing_common_utils import get_request_team_id +from litellm.router_strategy.base_routing_strategy import BaseRoutingStrategy +from litellm.router_strategy.budget_limiter import RouterBudgetLimiting +from litellm.router_strategy.least_busy import LeastBusyLoggingHandler +from litellm.router_utils.global_callback_holders import ( + live_check_owners, + routers_holding_global_callbacks, + routers_holding_global_callbacks_lock, +) +from litellm.router_utils.pre_call_checks.deployment_affinity_check import ( + DeploymentAffinityCheck, + completed_affinity_response, + warn_on_unknown_model_group_affinity_flags, +) +from litellm.router_utils.pre_call_checks.encrypted_content_affinity_check import EncryptedContentAffinityCheck +from litellm.router_utils.pre_call_checks.model_rate_limit_check import ModelRateLimitingCheck +from litellm.router_utils.pre_call_checks.prompt_caching_deployment_check import PromptCachingDeploymentCheck + +if TYPE_CHECKING: + from litellm.litellm_core_utils.litellm_logging import Logging + from litellm.router_python_host import PythonHost + from litellm.router_rust.steps.selection import StrategySelectors + from litellm.router_rust.surface.registry import HostRegistry + from litellm.router_utils.cooldown_cache import CooldownCache + from litellm.router_utils.global_callback_holders import GlobalCallbackHolder + from litellm.types.router import Deployment + +_AFFINITY_CHECKS: Final = ( + "deployment_affinity", + "responses_api_deployment_check", + "session_affinity", + "encrypted_content_affinity", +) +# `_RUNTIME_TOGGLEABLE_PRE_CALL_CHECKS`: the checks a runtime list drops when it leaves them out. +_RUNTIME_TOGGLEABLE: Final[Mapping[str, type[CustomLogger]]] = { + "prompt_caching": PromptCachingDeploymentCheck, + "enforce_model_rate_limits": ModelRateLimitingCheck, + "encrypted_content_affinity": EncryptedContentAffinityCheck, +} + + +def _callbacks(host: "PythonHost") -> list[object]: # mutable-ok: the process list itself + return host.callback_list("callbacks") # pyright: ignore[reportArgumentType] # not in the CallbackList literal + + +def register_router_selector(host: "PythonHost", selector: CustomLogger) -> None: + """``_register_router_selector``: least-busy also counts calls as they start.""" + if isinstance(selector, LeastBusyLoggingHandler): + if isinstance(host.callback_list("input_callback"), list): # pyright: ignore[reportUnnecessaryIsInstance] # a process list a user may have reassigned + host.callbacks.add_litellm_input_callback(selector) + else: + host.set_callback_list("input_callback", [selector]) + if isinstance(_callbacks(host), list): # pyright: ignore[reportUnnecessaryIsInstance] # see above + host.callbacks.add_litellm_callback(selector) + + +def _move_before_deployment_affinity( + callback_list: list[object], # mutable-ok: reordered in place + callback: object, +) -> None: + if callback not in callback_list: + return + callback_list.remove(callback) + index: Final = next( + (i for i, entry in enumerate(callback_list) if isinstance(entry, DeploymentAffinityCheck)), + len(callback_list), + ) + callback_list.insert(index, callback) + + +class OptionalChecks: + """One router's optional pre-call checks (``optional_callbacks``) and the stores they read.""" + + def __init__( + self, + owner: "GlobalCallbackHolder", + host: "PythonHost", + cache: DualCache, + hosted: "HostRegistry", + ctx: object, + arguments: Mapping[str, object], + ) -> None: + self._owner: Final = owner + self._host: Final = host + self._cache: Final = cache + self._hosted: Final = hosted + self._ctx: Final = ctx + self.optional_callbacks: list[object] | None = None # mutable-ok: the router's live list, as Python keeps it + self.router_budget_logger: RouterBudgetLimiting | None = None + self.provider_budget_config: object = arguments.get("provider_budget_config") + ttl: Final = arguments.get("deployment_affinity_ttl_seconds") + self.deployment_affinity_ttl_seconds: object = 3600 if ttl is None else ttl + self.session_affinity_repin_on_failover: object = bool(arguments.get("session_affinity_repin_on_failover")) + affinity_config: Final = arguments.get("model_group_affinity_config") + self.model_group_affinity_config: dict[str, list[str]] | None = ( # mutable-ok: the argument as passed + affinity_config if isinstance(affinity_config, dict) else None # pyright: ignore[reportAttributeAccessIssue] # the argument as passed + ) + warn_on_unknown_model_group_affinity_flags(self.model_group_affinity_config) + self.prompt_caching_pin_cache: DualCache = DualCache( + redis_cache=cache.redis_cache, + in_memory_cache=InMemoryCache(max_size_in_memory=PROMPT_CACHE_PIN_CACHE_MAX_SIZE), + ) + self._request_scoped_selectors: tuple[CustomLogger, ...] = () + self._request_scoped_policies: tuple[CustomLogger, ...] = () + # ``_slack_alerting_loggers``: the Slack alerting logger ``alerting_config`` builds, if any. + self.slack_alerting_loggers: tuple[CustomLogger, ...] = () + # Whether the outcome callbacks would still be on the host lists (``_release_global_callbacks``). + self._outcome_callbacks_listed: bool = True + # The logging objects ``_core_bind_to_request_logging`` would have bound them to. + self._bound_requests: Final[WeakSet[Logging]] = WeakSet() + + # -- construction --------------------------------------------------------------------------- + + def build(self, checks: object, model_list: object) -> None: + """The constructor's checks: a budget appends ``router_budget_limiting`` to the list given (in + place, as Python does), the list is added, a per-group affinity config adds the affinity + check, and ``apply_default_settings`` adds the empty default list.""" + listed: list[object] | None = checks if isinstance(checks, list) else None # pyright: ignore[reportUnknownVariableType] # mutable-ok: appended in place, as Python does # rebind-ok: the argument as passed + if RouterBudgetLimiting.should_init_router_budget_limiter( # pyright: ignore[reportUnknownMemberType] # untyped model_list + model_list=model_list if isinstance(model_list, list) else None, # pyright: ignore[reportUnknownArgumentType] # the argument as passed + provider_budget_config=self.provider_budget_config, # pyright: ignore[reportArgumentType] # the argument as passed + ): + if listed is not None: + listed.append("router_budget_limiting") + else: + listed = ["router_budget_limiting"] # rebind-ok: the list Python builds + if listed is not None: + self.add(listed) # pyright: ignore[reportUnknownArgumentType] # names as given + if self.model_group_affinity_config: + self.ensure_deployment_affinity_callback() + self.add([]) + + def _is_priority_group(self, model: str) -> bool: + """``_is_priority_routing_group``.""" + hosted: Final = self._hosted + resolved: Final = hosted._core_get_model_from_alias(model=model) or model # pyright: ignore[reportPrivateUsage] # the host registry's alias lookup + group: Final = hosted.get_routing_group(resolved) + return group is not None and group.routing_strategy == "priority" + + async def repin_session_after_success( + self, litellm_params: Mapping[str, object], response: object, cooldowns: "CooldownCache" + ) -> None: + """``_repin_session_after_success``: with ``session_affinity_repin_on_failover`` on, a completed + success moves the session's pin to the deployment that served it, among the group's + deployments not cooling down.""" + if not self.session_affinity_repin_on_failover: + return + try: + if not completed_affinity_response(response): + return + callback: Final = next( + (logger for logger in self.optional_callbacks or () if isinstance(logger, DeploymentAffinityCheck)), + None, + ) + if callback is None: + return + model_group: Final = callback.get_repin_model_group(litellm_params) + if model_group is None or self._is_priority_group(model_group): + return + registered_ids: Final = self._hosted._core_get_candidate_model_ids_for_route( # pyright: ignore[reportPrivateUsage] # the host registry's candidate lookup + model=model_group, team_id=get_request_team_id(litellm_params) + ) + cooldown_ids: Final = frozenset( + model_id + for model_id, _ in await cooldowns.async_get_active_cooldowns( + model_ids=list(registered_ids), parent_otel_span=None + ) + ) + await callback.async_repin_session_after_success(litellm_params, model_group, registered_ids - cooldown_ids) + except Exception as error: # noqa: BLE001 # optional pin updates must not interrupt success accounting + verbose_router_logger.debug("Session re-pinning failed: %s", error) + + def _affinity_check(self, user_key: bool, responses_api: bool, session_id: bool) -> DeploymentAffinityCheck: + return DeploymentAffinityCheck( + cache=self._cache, + ttl_seconds=self.deployment_affinity_ttl_seconds, # pyright: ignore[reportArgumentType] # the argument as passed + enable_user_key_affinity=user_key, + enable_responses_api_affinity=responses_api, + enable_session_id_affinity=session_id, + model_group_affinity_config=self.model_group_affinity_config, + is_priority_group=self._is_priority_group, + session_affinity_repin_on_failover=self.session_affinity_repin_on_failover, # pyright: ignore[reportArgumentType] # same + ) + + def _own(self, callback: CustomLogger) -> None: + if self.optional_callbacks is None: + self.optional_callbacks = [] + self.optional_callbacks.append(callback) + self._host.callbacks.add_litellm_callback(callback) + + def ensure_deployment_affinity_callback(self) -> None: + """``_ensure_deployment_affinity_callback``: the affinity check with every global flag off.""" + if any(isinstance(cb, DeploymentAffinityCheck) for cb in self.optional_callbacks or ()): + return + self._own(self._affinity_check(False, False, False)) + + def _add_affinity(self, checks: Sequence[object]) -> None: + user_key: Final = "deployment_affinity" in checks + responses_api: Final = "responses_api_deployment_check" in checks + session_id: Final = "session_affinity" in checks + if not (user_key or responses_api or session_id): + return + if self.optional_callbacks is None: + self.optional_callbacks = [] + existing: Final = next( + (cb for cb in self.optional_callbacks if isinstance(cb, DeploymentAffinityCheck)), + None, + ) + if existing is None: + self._own(self._affinity_check(user_key, responses_api, session_id)) + return + existing.enable_user_key_affinity = existing.enable_user_key_affinity or user_key + existing.enable_responses_api_affinity = existing.enable_responses_api_affinity or responses_api + existing.enable_session_id_affinity = existing.enable_session_id_affinity or session_id + existing.ttl_seconds = self.deployment_affinity_ttl_seconds # pyright: ignore[reportAttributeAccessIssue] # the argument as passed + if self.model_group_affinity_config: + existing.model_group_affinity_config = self.model_group_affinity_config + + def _add_encrypted_content_affinity(self, enable_global_affinity: bool) -> None: + """``_add_encrypted_content_affinity_check``.""" + if not ( + enable_global_affinity + or EncryptedContentAffinityCheck.has_model_group_affinity_enabled(self.model_group_affinity_config) + ): + return + if self.optional_callbacks is None: + self.optional_callbacks = [] + existing: Final = next( + (cb for cb in self.optional_callbacks if isinstance(cb, EncryptedContentAffinityCheck)), None + ) + if existing is not None: + existing.ctx = self._ctx # pyright: ignore[reportAttributeAccessIssue] # the core's RoutingContext serves registry() + existing.enable_global_affinity = existing.enable_global_affinity or enable_global_affinity + existing.model_group_affinity_config = self.model_group_affinity_config or {} + check: EncryptedContentAffinityCheck = existing # rebind-ok: the check to move + else: + check = EncryptedContentAffinityCheck( # rebind-ok: the check to move + ctx=self._ctx, # pyright: ignore[reportArgumentType] # see above + enable_global_affinity=enable_global_affinity, + model_group_affinity_config=self.model_group_affinity_config, + ) + self._own(check) + _move_before_deployment_affinity(self.optional_callbacks, check) + _move_before_deployment_affinity(_callbacks(self._host), check) + + def add(self, checks: Sequence[object] | None) -> None: + """``add_optional_pre_call_checks``.""" + if checks is None: + return + self._add_affinity(checks) + self._add_encrypted_content_affinity("encrypted_content_affinity" in checks) + for name in checks: + if name in _AFFINITY_CHECKS: + continue + callback: CustomLogger | None = None + if name == "prompt_caching": + callback = PromptCachingDeploymentCheck( + cache=self.prompt_caching_pin_cache, is_priority_group=self._is_priority_group + ) + elif name == "router_budget_limiting": + if self.budget_limiter() is not None: + continue + callback = RouterBudgetLimiting( + dual_cache=self._cache, + provider_budget_config=self.provider_budget_config, # pyright: ignore[reportArgumentType] # the argument as passed + model_list=self._hosted._core_model_list, # pyright: ignore[reportPrivateUsage, reportArgumentType] # the live model_list + ) + self.router_budget_logger = callback + elif name == "enforce_model_rate_limits": + callback = ModelRateLimitingCheck(dual_cache=self._cache) + if callback is None: + continue + if self.optional_callbacks is not None and any( + isinstance(existing, type(callback)) for existing in self.optional_callbacks + ): + continue + self._own(callback) + + def set(self, checks: object) -> None: + """``set_optional_pre_call_checks``: drop the toggleable checks the list leaves out, then add it.""" + if checks is None: + return + listed: Final[tuple[object, ...]] = tuple(checks) if isinstance(checks, list | tuple) else () # pyright: ignore[reportUnknownArgumentType, reportUnknownVariableType] # names as given + requested: Final = frozenset(listed) + for name, callback_cls in _RUNTIME_TOGGLEABLE.items(): + if name not in requested: + self._remove_of_type(callback_cls) + self.add(listed) + + def _remove_of_type(self, callback_cls: type[CustomLogger]) -> None: + """``_remove_optional_callbacks_of_type``: kept on the host while any other live router of + the process has a check of that type.""" + if self.optional_callbacks is None or not any(type(cb) is callback_cls for cb in self.optional_callbacks): + return + self.optional_callbacks = [cb for cb in self.optional_callbacks if type(cb) is not callback_cls] + others: Final = tuple(owner for owner in live_check_owners if id(owner) != id(self._owner)) + if any(any(type(cb) is callback_cls for cb in (owner.optional_callbacks or ())) for owner in others): + return + callbacks: Final = _callbacks(self._host) + for cb in tuple(callbacks): + if type(cb) is callback_cls: + self._host.callbacks.remove_callback_from_list_by_object(callbacks, cb, require_self=False) # pyright: ignore[reportUnknownMemberType] # the manager's untyped helper + + # -- budgets -------------------------------------------------------------------------------- + + def budget_limiter(self) -> RouterBudgetLimiting | None: + """``_core_get_router_deployment_budget_limiter``: exact type, so a proxy subclass is not taken.""" + if self.router_budget_logger is not None: + return self.router_budget_logger + for cb in self.optional_callbacks or (): + if isinstance(cb, RouterBudgetLimiting) and type(cb) is RouterBudgetLimiting: + self.router_budget_logger = cb + return cb + return None + + def sync_deployment_budget(self, deployment: "Deployment") -> None: + """``_sync_deployment_budget_config``, after a deployment is added.""" + model_id: Final = deployment.model_info.id + if model_id is None: + return + limiter: RouterBudgetLimiting | None = self.budget_limiter() # rebind-ok: built when first needed + params: Final = deployment.litellm_params + if params.max_budget is None or params.budget_duration is None: + if limiter is not None: + limiter.unregister_deployment_budget(model_id=model_id) + return + if limiter is None: + self.add(["router_budget_limiting"]) + limiter = self.budget_limiter() # rebind-ok: the one just built + if limiter is not None: + limiter.register_deployment_budget(deployment=deployment.to_json(exclude_none=True)) # pyright: ignore[reportUnknownMemberType] # pydantic's untyped dump + + def unregister_deployment_budget(self, model_id: str) -> None: + limiter: Final = self.budget_limiter() + if limiter is not None: + limiter.unregister_deployment_budget(model_id=model_id) + + # -- the routing callbacks ------------------------------------------------------------------ + + def loggers(self) -> tuple[CustomLogger, ...]: + """``_optional_loggers``.""" + return tuple(cb for cb in self.optional_callbacks or () if isinstance(cb, CustomLogger)) + + def _policies_not_set_globally(self) -> tuple[CustomLogger, ...]: + keys: Final = self._host.callbacks.get_custom_logger_key + global_keys: Final = frozenset(keys(cb) for cb in _callbacks(self._host) if isinstance(cb, CustomLogger)) + return tuple(policy for policy in self._request_scoped_policies if keys(policy) not in global_keys) + + def bind_to_request_logging(self, logging_obj: "Logging") -> None: + """``_core_bind_to_request_logging``'s routing half: the request-scoped selectors replace the + host's same-keyed callbacks for the request, and the policies the host does not run are + added to it. The outcome callbacks are not bound: the Rust core records every attempt's + outcome itself, so the request is noted for ``outcome_callbacks`` instead.""" + self._bound_requests.add(logging_obj) + for selector in self._request_scoped_selectors: + logging_obj.replace_global_callback(selector) + for policy in self._policies_not_set_globally(): + logging_obj.add_dynamic_callback(policy) + + def outcome_callbacks( + self, request: Mapping[str, object], *, sync: bool, make_call: bool = True + ) -> Mapping[str, bool]: + """Router-core's ``OutcomeCallbacks`` for ``request``, empty (its default) while the outcome + callbacks are listed. Past that only the request's bound logging object runs them, and + ``_core_bind_to_request_logging`` binds neither ``sync_deployment_callback_on_success`` nor any + success callback a sync call runs. ``make_call``: the attempt runs through ``make_call``.""" + if self._outcome_callbacks_listed: + return {} + bound: Final = request.get("litellm_logging_obj") in self._bound_requests + return {"success": bound and not sync, "failure": bound, "make_call": make_call} + + def routing_callbacks(self) -> Iterable[object]: + """``_routing_callbacks``: the host's callbacks, or after ``observe(request_scoped=True)`` the + host's callbacks with the router's own selectors and policies in place of theirs.""" + if not (self._request_scoped_selectors or self._request_scoped_policies): + return tuple(_callbacks(self._host)) + keys: Final = self._host.callbacks.get_custom_logger_key + replaced: Final = frozenset(map(keys, self._request_scoped_selectors)) + return ( + *self._host.callbacks.without_custom_logger_keys(_callbacks(self._host), replaced), # pyright: ignore[reportArgumentType] # the process list holds any callback + *self._request_scoped_selectors, + *self._policies_not_set_globally(), + ) + + async def filter_deployments( + self, + model: str, + deployments: list[dict[str, object]], # mutable-ok: CustomLogger deployment-list contract + messages: object, + request: Mapping[str, object], + ) -> list[object]: # mutable-ok: as the filters return it + """``async_callback_filter_deployments``: each routing callback narrows the deployments in + turn; a filter's error is raised for the Rust core to decide on.""" + remaining: list[object] = list(deployments) # mutable-ok: filter contract # rebind-ok: narrowed + base: Final[object] = CustomLogger.async_filter_deployments # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] # untyped deployment dicts + request_kwargs: Final = dict(request) + for callback in tuple(self.routing_callbacks()): + # The base filter returns its input unchanged: only overriding filters are run. + if not isinstance(callback, CustomLogger) or type(callback).async_filter_deployments is base: # pyright: ignore[reportUnknownMemberType] # see above + continue + remaining = await callback.async_filter_deployments( # pyright: ignore[reportUnknownMemberType, reportAssignmentType, reportUnknownVariableType] # see above + model=model, + healthy_deployments=remaining, + messages=messages, # pyright: ignore[reportArgumentType] # the request's messages, as given + request_kwargs=request_kwargs, + parent_otel_span=None, + ) + return remaining + + # -- release -------------------------------------------------------------------------------- + + def serve_through_request_logging(self, selectors: Sequence[CustomLogger], dropped: Sequence[object]) -> None: + """``_core_serve_through_request_logging``: keep running ``selectors`` (and request-scoped + adaptive hooks) and the optional checks, but take everything off the host lists.""" + from litellm.router_strategy.adaptive_router.hooks import AdaptiveRouterPostCallHook + + adaptive: Final = tuple( + AdaptiveRouterPostCallHook(adaptive_router=tagged.strategy) # pyright: ignore[reportArgumentType] # an adaptive router + for tagged in chain.from_iterable(self._hosted._core_adaptive_routers.values()) # pyright: ignore[reportPrivateUsage] # the host registry's live registry + ) + self._request_scoped_selectors = (*selectors, *adaptive) + self._request_scoped_policies = (*self.loggers(), *self.slack_alerting_loggers) + self.release(dropped) + + def own_adaptive_hooks(self) -> tuple[CustomLogger, ...]: + """``_adaptive_router_hooks``: the host's adaptive post-call hooks of this router's routers.""" + from litellm.router_strategy.adaptive_router.hooks import AdaptiveRouterPostCallHook + + own: Final = frozenset( + id(tagged.strategy) + for tagged in chain.from_iterable(self._hosted._core_adaptive_routers.values()) # pyright: ignore[reportPrivateUsage] # see above + ) + return tuple( + hook + for hook in self._host.callbacks.get_custom_loggers_for_type(AdaptiveRouterPostCallHook) + if isinstance(hook, AdaptiveRouterPostCallHook) and id(hook.adaptive_router) in own + ) + + def retire_state(self, selectors: Sequence[object]) -> None: + """``retire_routing_state``: stop the background work of the selectors and the budget limiter.""" + for strategy in (*selectors, *self.loggers()): + if isinstance(strategy, BaseRoutingStrategy | RouterBudgetLimiting): + strategy.retire() + + def release(self, dropped: Sequence[object]) -> None: + """``_release_global_callbacks``: drop ``dropped``, the optional checks and the adaptive hooks + from the host lists, then hand each vacated check to a surviving router's same-keyed one.""" + self._outcome_callbacks_listed = False + with routers_holding_global_callbacks_lock: + routers_holding_global_callbacks.discard(self._owner) + adaptive: Final = self.own_adaptive_hooks() + slack: Final = self.slack_alerting_loggers + vacated: Final = self._host.callbacks.drop_callbacks( + (*dropped, *(self.optional_callbacks or ()), *slack, *adaptive), method_owners=(self._owner, *slack) + ) + with routers_holding_global_callbacks_lock: + survivors: Final = tuple(routers_holding_global_callbacks) + candidates: Final = chain.from_iterable(router._global_callback_candidates() for router in survivors) # pyright: ignore[reportPrivateUsage] # the holder protocol + for candidate in self._host.callbacks.with_custom_logger_keys(candidates, vacated): + register_router_selector(self._host, candidate) + adaptive_keys: Final = frozenset(map(self._host.callbacks.get_custom_logger_key, adaptive)) + survivor: Final = next((router for router in survivors if router._core_adaptive_routers), None) # pyright: ignore[reportPrivateUsage] # see above + if survivor is not None and vacated & adaptive_keys: + survivor._sync_adaptive_router_hooks() # pyright: ignore[reportPrivateUsage] # see above + + +_ENCRYPTED_PIN: Final = "_encrypted_content_affinity_pinned" + + +def _request_with_failed_ids(request: Mapping[str, object], step: Mapping[str, object]) -> Mapping[str, object]: + """The request as the filters read it in the Python router: the Rust walk holds what the Python + one keeps in the attempt's kwargs, ``_excluded_deployment_ids`` from a weighted failover hop, + ``_retry_skipped_deployment_ids`` from the retry loop (``deployment_ids_failed_in_request``) and + ``_target_order`` from an order fallback.""" + excluded: Final = tuple(str(id) for id in _ids_of(step.get("excluded_ids"))) + skipped: Final = tuple(str(id) for id in _ids_of(step.get("skipped_ids"))) + target_order: Final = step.get("target_order") + if not excluded and not skipped and target_order is None: + return request + return { + **request, + **({"_excluded_deployment_ids": list(excluded)} if excluded else {}), + **({"_retry_skipped_deployment_ids": skipped} if skipped else {}), + **({"_target_order": target_order} if target_order is not None else {}), + } + + +def _ids_of(value: object) -> Sequence[object]: + return value if isinstance(value, list | tuple) else () # pyright: ignore[reportUnknownVariableType] # the step's id list + + +def filtered_ids(remaining: Iterable[object], candidates: Sequence[str]) -> list[str]: # mutable-ok: a JSON list + """The ids of the deployments the filters left, in their order, among the candidates.""" + ids: Final[list[str]] = [] # mutable-ok: built here + for deployment in remaining: + found = (as_str_mapping((as_str_mapping(deployment) or {}).get("model_info")) or {}).get("id") + if isinstance(found, str) and found in candidates and found not in ids: + ids.append(found) + return ids + + +class RustOptionalChecks: + """``RustCore`` members for its optional pre-call checks and routing callbacks.""" + + # RustCore sets these as Final attributes in its constructor. + @property + def _opt(self) -> OptionalChecks: + checks: Final[OptionalChecks] = self.__dict__["_checks"] # pyright: ignore[reportAny] # RustCore's own attribute + return checks + + @property + def _sel(self) -> "StrategySelectors": + selectors: Final[StrategySelectors] = self.__dict__["_selectors"] # pyright: ignore[reportAny] # same + return selectors + + @property + def _reg(self) -> "HostRegistry": + hosted: Final[HostRegistry] = self.__dict__["_hosted"] # pyright: ignore[reportAny] # same + return hosted + + @property + def optional_callbacks(self) -> "list[object] | None": # mutable-ok: the router's live list + return self._opt.optional_callbacks + + @property + def router_budget_logger(self) -> RouterBudgetLimiting | None: + return self._opt.router_budget_logger + + @property + def deployment_affinity_ttl_seconds(self) -> object: + return self._opt.deployment_affinity_ttl_seconds + + @property + def session_affinity_repin_on_failover(self) -> object: + return self._opt.session_affinity_repin_on_failover + + @property + def model_group_affinity_config(self) -> "dict[str, list[str]] | None": # mutable-ok: the argument as passed + return self._opt.model_group_affinity_config + + def add_optional_pre_call_checks(self, optional_pre_call_checks: "Sequence[object] | None") -> None: + self._opt.add(optional_pre_call_checks) + + def apply_default_settings(self) -> None: + """Add the (empty) default checks; an enabled continuation registers the prefill check.""" + from litellm.router_rust.invoke.streaming import register_continuation_prefill_check + + self._opt.add([]) + if getattr(self, "enable_mid_stream_fallback_continuation", False): + register_continuation_prefill_check(self._opt._host) # pyright: ignore[reportPrivateUsage] # the checks' host + + # The members another router's retire reads (`global_callback_holders`). + + def _global_callback_candidates(self) -> tuple[CustomLogger, ...]: + return (*self._sel.callback_selectors(), *self._opt.loggers(), *self._opt.slack_alerting_loggers) + + @property + def _core_adaptive_routers(self) -> Mapping[str, object]: + return self._reg._core_adaptive_routers # pyright: ignore[reportPrivateUsage] # the host registry's live registry + + def _sync_adaptive_router_hooks(self) -> None: + strategies: Final = self._reg.strategies + if strategies is not None: + strategies._sync_adaptive_router_hooks() # pyright: ignore[reportPrivateUsage] # the strategies' own hook sync + + def _join_callback_holders(self) -> None: + with routers_holding_global_callbacks_lock: + routers_holding_global_callbacks.add(self) + live_check_owners.add(self) + + def outcome_callbacks( + self, request: Mapping[str, object], *, sync: bool, make_call: bool = True + ) -> Mapping[str, bool]: + return self._opt.outcome_callbacks(request, sync=sync, make_call=make_call) + + def _serve_through_request_logging(self) -> None: + selectors: Final = self._sel + self._opt.serve_through_request_logging( + selectors.callback_selectors(), (*selectors.callback_selectors(), *selectors.override_selectors()) + ) + + def _retire_callbacks(self) -> None: + """``_core_discard``'s callback half: retire the routing state, release the callbacks. The + selectors stay: a retired router still picks with them.""" + live_check_owners.discard(self) + selectors: Final = (*self._sel.callback_selectors(), *self._sel.override_selectors()) + self._opt.retire_state(selectors) + self._opt.release(selectors) + + async def _answer_callback_filters( + self, step: Mapping[str, object], request: Mapping[str, object], messages: object + ) -> tuple[object, BaseException | None]: + """A ``callback_filters`` step: the remaining ids, or a filter's error as an attempt + failure with the error itself, which the caller keeps to raise.""" + from litellm.router_rust.core import attempt_failure + + candidates: Final = [str(id) for id in step.get("candidates") or ()] # pyright: ignore[reportGeneralTypeIssues, reportUnknownVariableType, reportUnknownArgumentType] # the step's id list + records: Final = self._records_of(candidates) # pyright: ignore[reportAttributeAccessIssue, reportUnknownMemberType, reportUnknownVariableType] # RustCore's deployment dicts + narrowed: Final = _request_with_failed_ids(request, step) + try: + remaining: Final = await self._opt.filter_deployments( + str(step.get("model_group")), + records, # pyright: ignore[reportUnknownArgumentType] # see above + messages, + narrowed, + ) + except Exception as error: # noqa: BLE001 # handed to the Rust core as the try's failure + return attempt_failure(error), error + # A filter's pin lands on the request kwargs the pick goes on to read. + if narrowed is not request and narrowed.get(_ENCRYPTED_PIN) and isinstance(request, dict): + request[_ENCRYPTED_PIN] = True # rebind-ok: the request dict the pick goes on to read + return filtered_ids(remaining, candidates), None # pyright: ignore[reportUnknownArgumentType] # see above + + +__all__ = [ + "OptionalChecks", + "RustOptionalChecks", + "filtered_ids", + "register_router_selector", +] diff --git a/litellm/router_rust/steps/plugins.py b/litellm/router_rust/steps/plugins.py new file mode 100644 index 00000000000..4af73313838 --- /dev/null +++ b/litellm/router_rust/steps/plugins.py @@ -0,0 +1,87 @@ +""" +Host side of the Rust core's routing-plugin pipeline (``Router(plugins=[...])``). + +The Rust core decides where the pipeline applies (only the async pick runs it; the sync pick +refuses a router with plugins) and which deployments the plugins' candidate models keep +(``router-core/src/routing_plugins.rs``). The plugins themselves are live Python objects, so this +module runs them on a ``routing_plugins`` select step, the way the Python router core's +``_run_routing_plugins`` does, and reports the candidate models left on the request. +""" + +from collections.abc import Mapping, Sequence +from typing import Final, Protocol + +from litellm.litellm_core_utils.core_helpers import as_str_mapping, get_metadata_variable_name_from_kwargs +from litellm.litellm_core_utils.routing_common_utils import resolve_model_group_alias +from litellm.types.router import RoutingContext, RoutingPlugin + +SYNC_PATH_MESSAGE: Final = ( + "Router(plugins=[...]) is configured, but this call resolved to the synchronous " + "deployment-selection path, which never runs the routing-plugin pipeline. This " + "happens for sync Router methods (e.g. Router.completion()) and for async calls " + "with a routing_strategy that has no async-native selector (e.g. legacy " + "'usage-based-routing', v1). Silently skipping " + "configured plugins would let a policy plugin (e.g. a deny-all rule) be bypassed. " + "Use an async Router method with a supported routing_strategy (simple-shuffle, " + "usage-based-routing-v2, cost-based-routing, latency-based-routing, least-busy), " + "or remove `plugins` from the Router config." +) +_CANDIDATES_KEY: Final = "_routing_plugin_candidate_models" + + +class PluginRegistry(Protocol): + def _core_get_model_list( + self, model_name: str | None = None, team_id: str | None = None + ) -> Sequence[Mapping[str, object]] | None: ... + + +def _resolved_litellm_models( + registry: PluginRegistry, model: str +) -> list[str]: # mutable-ok: the plugins' candidate list + """``_core_resolved_litellm_models``: the provider models a call to ``model`` is served by.""" + return [ + litellm_model + for deployment in registry._core_get_model_list(model_name=model) or () # pyright: ignore[reportPrivateUsage] # the host registry's listing + if isinstance(litellm_model := (as_str_mapping(deployment.get("litellm_params")) or {}).get("model"), str) + and litellm_model + ] + + +def candidate_models(request: Mapping[str, object]) -> Mapping[str, object]: + """The answer to a ``routing_plugins`` step: the candidate models the plugins left on the + request, ``None`` when none narrowed it.""" + metadata: Final = as_str_mapping(request.get(get_metadata_variable_name_from_kwargs(dict(request)))) + return {"models": metadata.get(_CANDIDATES_KEY) if metadata is not None else None} + + +async def run_plugins( + plugins: Sequence[RoutingPlugin], + registry: PluginRegistry, + model_group_alias: Mapping[str, object], + model: str, + request: dict[str, object], # mutable-ok: the plugins stamp their outcome on the caller's request + messages: object, +) -> None: + """``_run_routing_plugins``: run ``plugins`` in order over ``model``'s candidate models, then + stamp their signals (and a narrowed candidate list) on the request's metadata.""" + from litellm.litellm_core_utils.prompt_templates.factory import resolve_structured_messages + + registered: Final = resolve_model_group_alias(model_group_alias, model) or model + models: Final = _resolved_litellm_models(registry, registered) + metadata_key: Final = get_metadata_variable_name_from_kwargs(request) + metadata: Final = request.setdefault(metadata_key, {}) + raw_messages: Final[object] = messages or list[object]() + context = RoutingContext( # rebind-ok: each plugin returns the context the next one reads + raw_messages=raw_messages, # pyright: ignore[reportArgumentType] # the request's messages + structured_messages=resolve_structured_messages(messages=messages, request_kwargs=request) or [], # pyright: ignore[reportArgumentType] # the request's messages + candidate_models=list(models), + metadata=metadata, # pyright: ignore[reportArgumentType] # the request's metadata dict + ) + for plugin in plugins: + context = await plugin.run(context) + metadata["routing_plugin_signals"] = context.signals # pyright: ignore[reportIndexIssue] # the request's metadata dict + if len(context.candidate_models) < len(models): + metadata[_CANDIDATES_KEY] = context.candidate_models # pyright: ignore[reportIndexIssue] # see above + + +__all__ = ["SYNC_PATH_MESSAGE", "PluginRegistry", "candidate_models", "run_plugins"] diff --git a/litellm/router_rust/steps/precall.py b/litellm/router_rust/steps/precall.py new file mode 100644 index 00000000000..156c7b5e20f --- /dev/null +++ b/litellm/router_rust/steps/precall.py @@ -0,0 +1,494 @@ +""" +Host-side facts for the Rust core's ``enable_pre_call_checks`` filter. + +The Rust core decides which deployments the pre-call checks drop and which error a group with +nothing left raises (``router-core/src/precall.rs``). What needs live Python objects is answered +here, on a ``pre_call`` select step: the prompt's token count (messages, Responses input and +instructions, tools and an Anthropic ``system`` block), each deployment's context window from +its router model info, and which requested params each deployment takes. Ported from the Python +router core's ``_pre_call_checks`` helpers; it never reaches a Python router core. + +``pre_call_checks`` is the one decision made here: ``RoutingContext.pre_call_checks`` hands over +deployment records rather than a step, and those need not be in the Rust registry. +""" + +from collections.abc import Callable, Mapping, Sequence +from functools import partial +from types import MappingProxyType +from typing import TYPE_CHECKING, Final, Protocol + +from litellm._logging import verbose_router_logger +from litellm.litellm_core_utils.core_helpers import as_str_mapping +from litellm.litellm_core_utils.deployment_pricing import deployment_backend_models +from litellm.litellm_core_utils.get_llm_provider_logic import declared_authenticating_provider +from litellm.litellm_core_utils.token_counter import ( + offload_token_count, + served_openai_models, + single_encoding_openai_model, + split_embedding_input_tokens, +) +from litellm.router_python_host import python_host +from litellm.router_strategy.complexity_router.context_compaction import compaction_pending +from litellm.types.router import LiteLLM_Params + +if TYPE_CHECKING: + from litellm.router_python_host import PythonHost + from litellm.types.utils import LlmProviders + +# The chat params a deployment must take when the request sends them (`_PRE_CALL_CHECK_FILTERED_PARAMS`). +_FILTERED_PARAMS: Final = frozenset({"response_format"}) + + +class PreCallRegistry(Protocol): + """The host registry reads the facts need.""" + + def _core_get_model_info(self, id: str) -> Mapping[str, object] | None: ... # the registry's name + + def _core_deployments_named(self, model_name: str) -> Sequence[Mapping[str, object]] | None: ... + + def _core_get_router_model_info( + self, deployment: Mapping[str, object], received_model_name: str + ) -> Mapping[str, object]: ... + + +def count_tokens( + host: "PythonHost", + messages: object, + input: object, # the request field's name + request_kwargs: Mapping[str, object] | None, + tokenizer_model: str, + deployments: Sequence[Mapping[str, object]] = (), +) -> int: + """``_count_pre_call_check_tokens``: messages, or Responses input normalized to messages, plus + the tools and Anthropic ``system`` block the message list never carries. Embedding media an + item of ``deployments`` reads as media is left out of the input: it never counts as text.""" + from litellm.llms.anthropic.messages.native.utils import anthropic_system_to_openai_message + + extras: Final[Mapping[str, object]] = request_kwargs if request_kwargs is not None else MappingProxyType({}) + raw_instructions: Final = extras.get("instructions") + instructions: Final = raw_instructions if isinstance(raw_instructions, str) else None + raw_tools: Final = extras.get("tools") + tools: Final = raw_tools if _is_list(raw_tools) and raw_tools else None + system_message: Final = anthropic_system_to_openai_message(extras.get("system")) + counter: Final = host.invoke("token_counter") # pyright: ignore[reportUnknownMemberType] # the host's overloads name litellm's untyped returns + if messages is not None: + counted_messages: Final = (system_message, *_items(messages)) if system_message is not None else messages + return counter(model=tokenizer_model, messages=counted_messages, tools=tools) + if input is not None: + from litellm.responses.litellm_completion_transformation.transformation import ( + LiteLLMCompletionResponsesConfig, + ) + + text_input, _ = split_embedding_input_tokens(input, partial(deployment_backend_models, deployments)) + input_messages: Final = LiteLLMCompletionResponsesConfig.transform_responses_api_input_to_messages( # pyright: ignore[reportUnknownMemberType] # its input type is untyped + input=text_input, # pyright: ignore[reportArgumentType] # str | list, as the request carries it + responses_api_request={"instructions": instructions} if instructions is not None else {}, + only_since_compaction=True, + ) + return counter(model=tokenizer_model, messages=input_messages, tools=tools) + raise ValueError("Either messages or input must be provided to count tokens") + + +def _max_input_tokens(registry: PreCallRegistry, model: str, deployment: Mapping[str, object]) -> int | None: + """``_deployment_max_input_tokens``: the declared context window, or None when there is none or + the deployment cannot be resolved.""" + try: + model_info: Final = registry._core_get_router_model_info( # pyright: ignore[reportPrivateUsage] # the host registry's lookup + deployment=deployment, received_model_name=model + ) + except Exception as e: # noqa: BLE001 # best-effort: an unmappable deployment must not hide the others + verbose_router_logger.debug("litellm.router.py::_deployment_max_input_tokens: skipping deployment. Got - %s", e) + return None + max_input_tokens: Final = model_info.get("max_input_tokens") + return max_input_tokens if isinstance(max_input_tokens, int) else None + + +def _items(value: object) -> tuple[object, ...]: + """The items of a list, tuple or set; nothing for any other value.""" + return tuple(value) if isinstance(value, list | tuple | set | frozenset) else () # pyright: ignore[reportUnknownArgumentType, reportUnknownVariableType] # the items are read as objects + + +def _is_list(value: object) -> bool: + return isinstance(value, list) + + +def _declared_param_allowlist(params: Mapping[str, object]) -> frozenset[str]: + declared: Final = params.get("allowed_openai_params") + return frozenset(entry for entry in _items(declared) if isinstance(entry, str)) + + +def _capability_model_and_provider(deployment_params: Mapping[str, object], group: str) -> tuple[str, str]: + model_label: Final = str(deployment_params.get("model") or group) + custom_llm_provider: Final = deployment_params.get("custom_llm_provider") + authenticating_provider: Final = declared_authenticating_provider( + model_label, + custom_llm_provider, # pyright: ignore[reportArgumentType] # passed as configured, as the Python core does + ) + if authenticating_provider is not None: + return model_label.removeprefix(f"{authenticating_provider}/"), authenticating_provider + model, provider, _, _ = python_host.invoke("get_llm_provider")( # pyright: ignore[reportUnknownMemberType] # see count_tokens + model=model_label, + custom_llm_provider=custom_llm_provider, + litellm_params=LiteLLM_Params.model_validate(dict(deployment_params)), + ) + return model, provider + + +def _supports_param(deployment: Mapping[str, object], group: str, param: str) -> bool: + """``_deployment_supports_param``.""" + deployment_params: Final = as_str_mapping(deployment.get("litellm_params")) + if not deployment_params: + return True + deployment_model_info: Final = as_str_mapping(deployment.get("model_info")) + base_model: Final = ( + deployment_model_info.get("base_model") if deployment_model_info else None + ) or deployment_params.get("base_model") + try: + model, custom_llm_provider = _capability_model_and_provider(deployment_params, group) + supported: Final = python_host.invoke("get_supported_openai_params")( # pyright: ignore[reportUnknownMemberType] # see count_tokens + model=model, + custom_llm_provider=custom_llm_provider, + base_model=base_model if isinstance(base_model, str) else None, + ) + except Exception as e: # noqa: BLE001 # best-effort filter: an unresolvable provider must not narrow the request + verbose_router_logger.debug( + "litellm.router.py::_deployment_supports_param: keeping %s for model=%s. Got - %s", param, group, e + ) + return True + return supported is None or param in supported + + +def _lacks_params(deployment: Mapping[str, object], group: str, params: frozenset[str]) -> bool: + return any(not _supports_param(deployment, group, param) for param in params) + + +def _rejects_params( + deployment: Mapping[str, object], group: str, params: frozenset[str], request_kwargs: Mapping[str, object] +) -> bool: + """``_deployment_rejects_params``: lacks a param neither the deployment nor the request drops or allows.""" + call_settings: Final[Mapping[str, object]] = { + **(as_str_mapping(deployment.get("litellm_params")) or {}), + **request_kwargs, + } + if call_settings.get("drop_params") is True: + return False + dropped: Final = call_settings.get("additional_drop_params") + dropped_names: Final = frozenset( + entry for entry in (_items(dropped) if _is_list(dropped) else ()) if isinstance(entry, str) + ) + return _lacks_params(deployment, group, params - _declared_param_allowlist(call_settings) - dropped_names) + + +def _requested_params(request_kwargs: Mapping[str, object]) -> frozenset[str]: + """``_pre_call_check_requested_chat_params``.""" + from litellm import utils as litellm_utils + + return _filtered_params( + litellm_utils.get_non_default_params(passed_params=dict(request_kwargs)) # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType] # an untyped helper + ) + + +def _filtered_params(requested: Mapping[object, object]) -> frozenset[str]: + return frozenset( + param + for param, value in requested.items() + if isinstance(param, str) and param in _FILTERED_PARAMS and not isinstance(value, str) + ) + + +class PreCallFacts: + """The facts of one ``pre_call`` step.""" + + def __init__( + self, + host: "PythonHost", + registry: PreCallRegistry, + step: Mapping[str, object], + records: Callable[[str], Mapping[str, object]], + request: Mapping[str, object] | None, + messages: object, + input: object, # the request field's name + ) -> None: + self._host: Final = host + self._registry: Final = registry + self._model: Final = str(step.get("model_group")) + self._name: Final = str(step.get("model")) + self._candidates: Final = [str(id) for id in _items(step.get("candidates"))] + self._pool: Final = [str(id) for id in _items(step.get("pool"))] + self._records: Final = records + self._request: Final = request + self._messages: Final = messages + self._input: Final = input + self._has_input: Final = messages is not None or input is not None + self._windows: Final = {id: _max_input_tokens(registry, self._model, records(id)) for id in self._candidates} + + def _record(self, id: str) -> Mapping[str, object]: # a deployment id + return self._records(id) + + def _needs_count(self) -> bool: + return self._has_input and any(limit is not None for limit in self._windows.values()) + + def _tokenizer_model(self) -> str: + named: Final = self._registry._core_deployments_named(self._model) or () # pyright: ignore[reportPrivateUsage] # the host registry's index + return ( + single_encoding_openai_model( + served_openai_models(named) or served_openai_models(self._record(id) for id in self._candidates) + ) + or "" + ) + + def _candidate_records(self) -> Sequence[Mapping[str, object]]: + return tuple(self._record(id) for id in self._candidates) + + def _count(self) -> int | None: + try: + return count_tokens( + self._host, + self._messages, + self._input, + self._request, + self._tokenizer_model(), + self._candidate_records(), + ) + except Exception as e: # noqa: BLE001 # best-effort: an uncountable prompt must not fail the request + verbose_router_logger.error( + "litellm.router.py::_acount_pre_call_check_tokens: failed to count tokens. Got - %s", e + ) + return None + + async def _count_async(self) -> int | None: + try: + return await offload_token_count(count_tokens)( + self._host, + self._messages, + self._input, + self._request, + self._tokenizer_model(), + self._candidate_records(), + ) + except Exception as e: # noqa: BLE001 # best-effort: an uncountable prompt must not fail the request + verbose_router_logger.error( + "litellm.router.py::_acount_pre_call_check_tokens: failed to count tokens. Got - %s", e + ) + return None + + def _answer(self, input_tokens: int | None) -> Mapping[str, object]: + if not self._has_input: + return {"has_input": False} + request: Final[Mapping[str, object]] = self._request if self._request is not None else {} + params: Final = ( + _requested_params(request) + if self._request is not None and self._host.dynamic().drop_params is False + else frozenset[str]() + ) + deployments: Final[dict[str, dict[str, object]]] = {} # mutable-ok: built here, then returned + for id in dict.fromkeys([*self._candidates, *self._pool]): + record = self._record(id) + deployments[id] = { + "max_input_tokens": self._windows.get(id), + "lacks_params": bool(params) and id in self._candidates and _lacks_params(record, self._model, params), + "rejects_params": bool(params) and _rejects_params(record, self._model, params, request), + } + allowed_region: Final = request.get("allowed_model_region") + return { + "has_input": True, + "countable": not compaction_pending(self._request), + "input_tokens": input_tokens, + "allowed_model_region": allowed_region, + "requested_params": sorted(params), + "deployments": deployments, + } + + def answer(self) -> Mapping[str, object]: + """The facts on the sync path: the prompt is counted inline.""" + return self._answer(self._count() if self._needs_count() else None) + + async def answer_async(self) -> Mapping[str, object]: + """The facts on the async path: the prompt is counted off the event loop.""" + return self._answer(await self._count_async() if self._needs_count() else None) + + +def pre_call_checks( + host: "PythonHost", + registry: PreCallRegistry, + model: str, + deployments: Sequence[Mapping[str, object]], + messages: object, + input: object, # the request field's name + request_kwargs: Mapping[str, object] | None, + input_token_count: int | None = None, + skip_inline_token_count: bool = False, + request_deployments: Sequence[Mapping[str, object]] | None = None, +) -> Sequence[Mapping[str, object]]: + """``_pre_call_checks`` over the given deployments, as ``router-core/src/precall.rs`` decides it.""" + import litellm.exceptions as litellm_exceptions + + pool: Final = ( + tuple(request_deployments) + if request_deployments is not None + else (tuple(registry._core_deployments_named(model) or ()) or tuple(deployments)) # pyright: ignore[reportPrivateUsage] # the host registry's index + ) + candidates: Final = {f"candidate-{index}": record for index, record in enumerate(deployments)} + group: Final = {f"pool-{index}": record for index, record in enumerate(pool)} + records: Final = {**candidates, **group} + step: Final = {"model_group": model, "model": model, "candidates": list(candidates), "pool": list(group)} + facts: Final = PreCallFacts(host, registry, step, records.__getitem__, request_kwargs, messages, input) + counted: Final = ( + input_token_count + if input_token_count is not None or skip_inline_token_count or not facts._needs_count() # pyright: ignore[reportPrivateUsage] # this module's own class + else facts._count() # pyright: ignore[reportPrivateUsage] # see above + ) + answer: Final = facts._answer(counted) # pyright: ignore[reportPrivateUsage] # see above + rows: Final = as_str_mapping(answer.get("deployments")) or {} + allowed_region: Final = answer.get("allowed_model_region") + + def row(id: str) -> Mapping[str, object]: # a candidate or pool key + return as_str_mapping(rows.get(id)) or {} + + def over_window(id: str) -> bool: + limit: Final = row(id).get("max_input_tokens") + return isinstance(limit, int) and answer.get("countable") is True and counted is not None and counted > limit + + def out_of_region(id: str) -> bool: + params: Final = as_str_mapping(records[id].get("litellm_params")) or {} + return allowed_region is not None and params.get("region_name") != allowed_region + + too_small: Final = tuple(id for id in candidates if over_window(id)) + invalid: Final = frozenset(too_small) | frozenset(id for id in candidates if out_of_region(id)) + param_candidates: Final = tuple(id for id in candidates if id not in invalid) + lacking: Final = tuple(row(id).get("lacks_params") is True for id in param_candidates) + rejecting: Final = tuple(row(id).get("rejects_params") is True for id in param_candidates) + failing: Final = rejecting if lacking and all(lacking) else lacking + excluded: Final = invalid | frozenset(id for id, fails in zip(param_candidates, failing, strict=True) if fails) + if len(excluded) == len(candidates): + requested: Final = [str(param) for param in _items(answer.get("requested_params"))] + if not too_small and group and all(row(id).get("rejects_params") is True for id in group): + raise litellm_exceptions.UnsupportedParamsError( + message=( + f"no deployment of model={model} supports {requested}. " + "To drop these, set `litellm.drop_params=True` or for proxy `litellm_settings: drop_params: true`, " + f"or send allowed_openai_params={requested} to pass them through" + ), + model=model, + ) + if too_small: + detail: Final = "".join( + f"Model={_error_model(records[id])}, Max Input Tokens={row(id).get('max_input_tokens')}, Got={counted}" + for id in too_small + ) + raise litellm_exceptions.ContextWindowExceededError( + message="litellm._pre_call_checks: Context Window exceeded for given call. No models have context " + f"window large enough for this call.\n{detail}", + model=model, + llm_provider="", + ) + return tuple(record for id, record in candidates.items() if id not in excluded) + + +def _error_model(deployment: Mapping[str, object]) -> object: + """The name a context-window error gives a deployment: its base model, else its model.""" + params: Final = as_str_mapping(deployment.get("litellm_params")) or {} + model_info: Final = as_str_mapping(deployment.get("model_info")) or {} + return model_info.get("base_model") or params.get("base_model") or params.get("model") + + +def request_carries_image(messages: object, request: Mapping[str, object]) -> bool: + """``_request_carries_image``: the messages (or the request's own messages or input) hold an image.""" + from litellm.litellm_core_utils.prompt_templates.common_utils import request_contains_image_content + + return request_contains_image_content(messages or request.get("messages") or request.get("input")) + + +def declares_no_vision(deployment: Mapping[str, object], group: str) -> bool: + """``_deployment_declares_no_vision``: ``model_info.supports_vision`` is ``False``, or, when it + is unset, the cost map says the model has vision explicitly disabled.""" + from litellm import utils as litellm_utils + + model_info: Final = as_str_mapping(deployment.get("model_info")) or {} + declared: Final = model_info.get("supports_vision") + if declared is not None: + return declared is False + params: Final = as_str_mapping(deployment.get("litellm_params")) or {} + if not params: + return False + base_model: Final = model_info.get("base_model") or params.get("base_model") + try: + model, custom_llm_provider = _capability_model_and_provider(params, group) + return litellm_utils.is_vision_explicitly_disabled( + base_model if isinstance(base_model, str) else model, custom_llm_provider + ) + except Exception as e: # noqa: BLE001 # best-effort filter: an unresolvable provider must not narrow the request + verbose_router_logger.debug( + "litellm.router.py::_deployment_declares_no_vision: keeping deployment for model=%s. Got - %s", group, e + ) + return False + + +__all__ = [ + "PreCallFacts", + "PreCallRegistry", + "count_tokens", + "declares_no_vision", + "pre_call_checks", + "request_carries_image", +] + + +def _route_capability_resolvers() -> Mapping[str, Callable[[str, "LlmProviders"], object | None]]: + """``_ROUTE_CAPABILITY_RESOLVERS``: the calls only some providers serve, by handler name.""" + from litellm.utils import ProviderConfigManager + + return { + "adecisions": ProviderConfigManager.get_provider_decisions_config, # pyright: ignore[reportUnknownMemberType] # litellm's config lookup + "decisions": ProviderConfigManager.get_provider_decisions_config, # pyright: ignore[reportUnknownMemberType] # see above + } + + +# ``_DECISIONS_MEDIA_CALL_TYPES``: the calls whose audio or videos some deployments refuse. +DECISIONS_MEDIA_CALL_TYPES: Final = frozenset({"adecisions", "decisions", "asystemone", "systemone"}) + + +def serves_route_capability(handler: str | None) -> bool: + """Whether only some providers serve the call ``handler`` names.""" + return handler is not None and handler in _route_capability_resolvers() + + +def cannot_serve_route(deployment: Mapping[str, object], group: str, handler: str) -> bool: + """Not ``_deployment_serves_route``: the deployment's provider has no config for the call.""" + from litellm.types.utils import LlmProviders + + resolver: Final = _route_capability_resolvers().get(handler) + params: Final = as_str_mapping(deployment.get("litellm_params")) + if resolver is None or params is None: + return False + try: + model, provider_name = _capability_model_and_provider(params, group) + provider: Final = next((member for member in LlmProviders if member.value == provider_name), None) + return provider is None or resolver(model, provider) is None + except Exception as e: # noqa: BLE001 # best-effort capability filter, as the Python core keeps it + verbose_router_logger.debug( + "litellm.router.py::_deployment_serves_route: keeping deployment for model=%s. Got - %s", group, e + ) + return False + + +def refuses_decisions_media( + deployment: Mapping[str, object], group: str, sent_kwargs: Mapping[str, object], *, audio: bool, videos: bool +) -> bool: + """``_decisions_deployment_refuses_media``.""" + from litellm.decisions.call import deployment_refuses_media + + model_info: Final = as_str_mapping(deployment.get("model_info")) or {} + params: Final = as_str_mapping(deployment.get("litellm_params")) or {} + try: + model, custom_llm_provider = _capability_model_and_provider(params, group) + return deployment_refuses_media( + model, custom_llm_provider, model_info, {**params, **sent_kwargs}, audio=audio, videos=videos + ) + except Exception as e: # noqa: BLE001 # best-effort filter: a deployment it cannot resolve answers for itself + verbose_router_logger.debug( + "litellm.router.py::_decisions_deployment_refuses_media: keeping a deployment of model=%s. Got - %s", + group, + e, + ) + return False diff --git a/litellm/router_rust/steps/prerouting.py b/litellm/router_rust/steps/prerouting.py new file mode 100644 index 00000000000..a83fbe29996 --- /dev/null +++ b/litellm/router_rust/steps/prerouting.py @@ -0,0 +1,1107 @@ +# pyright: reportUnknownMemberType=false, reportUnknownVariableType=false, reportUnknownArgumentType=false, reportAny=false, reportExplicitAny=false +# (the ported core code reads the untyped strategy routers, LiteLLM_Params dict fields and request dicts) +""" +Host-side pre-routing for ``RustCore``. + +Before an async pick resolves its name, the Rust core asks the host to pre-route it (a +``pre_route`` step, ``router-core/src/prerouting.rs``): the complexity, auto, adaptive and quality +routers registered from ``auto_router/...`` marker deployments, and the Claude Code session +binding, may send the request to another model group. They are live Python objects, so they stay +host-side for now (bridged, to port): this module registers them from the host registry's +deployments the way the Python router's ``init_*_router_deployment``, +``_finalize_adaptive_router_if_configured`` and ``_unregister_pre_routing_strategy_for_deployment`` +do, and answers the step the way ``_core_async_pre_routing_hook`` and ``_pin_tier_params_onto_request`` +do. It never reaches a Python router core. +""" + +import json +import logging +import re +from collections.abc import Awaitable, Callable, Iterable, Mapping, MutableMapping, Sequence +from itertools import chain +from types import MappingProxyType +from typing import TYPE_CHECKING, Any, Final, TypeAlias, cast # noqa: TID251 # the strategy routers are untyped + +from litellm._internal_context import with_service_target +from litellm._logging import verbose_router_logger +from litellm.caching.caching import DualCache, InMemoryCache +from litellm.caching.redis_cache import log_redis_failure +from litellm.constants import ( + CLIENT_OUTPUT_CEILING_METADATA_KEY, + CONSUMED_REQUEST_TAGS_METADATA_KEY, + DEFAULT_AUTO_ROUTER_MAX_INPUT_CHARS, + OUTPUT_TOKEN_CEILING_PARAMS, + SESSION_DEPLOYMENT_AFFINITY_TTL_METADATA_KEY, +) +from litellm.integrations.otel.routing import routing_decision_attributes +from litellm.integrations.otel.runtime import phase_attributes +from litellm.litellm_core_utils.core_helpers import get_or_create_metadata_bucket, without +from litellm.litellm_core_utils.credential_accessor import reject_api_key_secret_reference +from litellm.router_api import deployment_id +from litellm.router_host import router_host +from litellm.router_strategy.auto_router.auto_router_model_naming import ( + AUTO_ROUTER_MODEL_PREFIX, + STRATEGY_ROUTER_PARAM_FIELDS, + GatedAutoRouterCapability, + capability_limit_violation, + claimed_capabilities, + classify_strategy_router_model, + count_capability_routers, +) +from litellm.router_strategy.complexity_router.context_compaction import ( + arm_compaction, + is_native_compaction_call, + reject_recursive_compactor, +) +from litellm.router_strategy.tag_based_routing import get_routing_tags_from_request_kwargs, is_valid_deployment_tag +from litellm.router_utils.fallback_event_handlers import record_pre_routing_selection +from litellm.types.router import ConsumedRequestTagsStamp, Deployment, LiteLLM_Params, TaggedPreRoutingStrategy +from litellm.types.utils import ( + PROMPT_QUOTING_ROUTING_DECISION_FIELDS, + CustomPricingLiteLLMParams, + StandardLoggingRoutingDecision, + all_litellm_params, +) +from litellm.utils import provider_rejectable_params + +if TYPE_CHECKING: + from litellm.router_python_host import PythonHost + from litellm.router_rust.surface.registry import HostRegistry + from litellm.router_strategy.auto_router.auto_router import PreRoutingHookResponse + +# The live registry dicts: model_name -> the strategies registered under it, one per tag set. +StrategyRegistry: TypeAlias = "dict[str, list[TaggedPreRoutingStrategy[Any]]]" + +_ALIAS_PARAMS_NEVER_FORWARDED: Final = frozenset( + {"model", "api_base", "api_key", "api_version", "responses_api_version"} +) +_ALIAS_MARKER_FORWARDED_PARAMS_KWARG: Final = "_alias_marker_forwarded_params" +_CLAUDE_CODE_SESSION_ID_RE: Final = re.compile(r"^[a-zA-Z0-9_\-]{8,}$") +_CLAUDE_CODE_SESSION_ROUTER_TTL_SECONDS: Final = 3600 +CLAUDE_CODE_SESSION_ROUTER_BINDING_TARGET: Final = "claude_code_session_router_binding" +# `_async_get_available_deployment` runs the pre-routing hook only for the strategies with an async +# pick; any other router strategy takes the sync path, which never pre-routes. +_ASYNC_PICK_STRATEGIES: Final = frozenset( + {"usage-based-routing-v2", "simple-shuffle", "cost-based-routing", "latency-based-routing", "least-busy"} +) + + +def picks_async(router_strategy: str | None) -> bool: + """Whether the router's strategy has an async selector: the others pick on the sync path even + for async callers (``_async_get_available_deployment``).""" + return router_strategy in _ASYNC_PICK_STRATEGIES + + +TIER_PARAMS_NEVER_DROPPED: Final = frozenset(all_litellm_params) | frozenset( + { + "additional_drop_params", + "drop_params", + "messages", + "model", + "extra_headers", + "max_tokens", + "max_completion_tokens", + } +) + + +def _json_default_stable_id(value: object) -> str: + return f"{type(value).__module__}.{type(value).__qualname__}" + + +def _deployment_tags(deployment: Deployment) -> tuple[str, ...]: + return tuple(deployment.litellm_params.tags or ()) + + +def _has_registered_strategy(registry: StrategyRegistry, model_name: str, tags: tuple[str, ...]) -> bool: + return any(existing.tags == tags for existing in registry.get(model_name, [])) + + +def _routing_definition_fingerprint(deployment: Deployment) -> str | None: + """Identify configured routing fields, excluding external files and live strategy state.""" + try: + definition: Final = json.dumps( + deployment.litellm_params.model_dump( + include=set(STRATEGY_ROUTER_PARAM_FIELDS | {"model", "tags"}), exclude_none=True, warnings=False + ), + sort_keys=True, + default=_json_default_stable_id, + ) + except (TypeError, ValueError): + return None + return deployment_id(deployment.model_name, {"routing_definition": definition}) + + +def _unregister(registry: StrategyRegistry, model_name: str, tags: tuple[str, ...]) -> bool: + existing: Final = registry.get(model_name, []) + remaining: Final = [entry for entry in existing if entry.tags != tags] + if len(remaining) == len(existing): + return False + if remaining: + registry[model_name] = remaining # rebind-ok: the live registry, edited in place as the Python router does + else: + registry.pop(model_name, None) + return True + + +def is_strategy_marker(deployment: Mapping[str, object]) -> bool: + params: Final = deployment.get("litellm_params") + if not isinstance(params, Mapping): + return False + model: Final = params.get("model") + return isinstance(model, str) and classify_strategy_router_model(model) is not None + + +def _participates_in_adaptive_routing(params: LiteLLM_Params) -> bool: + kind: Final = classify_strategy_router_model(params.model) + if kind == "adaptive": + return True + if kind != "complexity": + return False + config: Final = params.complexity_router_config + return bool(config.get("adaptive")) if config else False + + +class PreRoutingStrategies: + """The pre-routing strategies of one router, kept on its host registry (``registry().plugins`` + reads ``_core_auto_routers`` and its siblings there).""" + + def __init__( + self, hosted: "HostRegistry", host: "PythonHost", ctx: object, capability_limit: object, cache: DualCache + ) -> None: + self._hosted: Final = hosted + self._host: Final = host + self._ctx: Final = ctx + self._capability_limit: object = capability_limit + self._cache: Final = cache + self._affinity: object = None + # The router's own `_ensure_deployment_affinity_callback` (its optional checks), when set. + self.ensure_affinity: Callable[[], None] | None = None + + # -- the registries ------------------------------------------------------------------------ + + @property + def auto(self) -> StrategyRegistry: + return self._hosted._core_auto_routers # pyright: ignore[reportPrivateUsage] # the registries live on the host registry + + @property + def complexity(self) -> StrategyRegistry: + return self._hosted._core_complexity_routers # pyright: ignore[reportPrivateUsage] # see above + + @property + def adaptive(self) -> StrategyRegistry: + return self._hosted._core_adaptive_routers # pyright: ignore[reportPrivateUsage] # see above + + @property + def quality(self) -> StrategyRegistry: + return self._hosted._core_quality_routers # pyright: ignore[reportPrivateUsage] # see above + + def all(self) -> tuple[StrategyRegistry, ...]: + """In the order ``_select_pre_routing_strategy`` searches them.""" + return (self.auto, self.complexity, self.adaptive, self.quality) + + def any_registered(self) -> bool: + return any(self.all()) + + # -- registry hooks (HostRegistry calls these) ----------------------------------------------- + + def reset(self) -> None: + """``set_model_list`` drops the auto, complexity and quality routers (never the adaptive ones).""" + self._hosted._core_quality_routers = {} # pyright: ignore[reportPrivateUsage] # see above + self._hosted._core_complexity_routers = {} # pyright: ignore[reportPrivateUsage] # see above + self._hosted._core_auto_routers = {} # pyright: ignore[reportPrivateUsage] # see above + + def register(self, deployment: Deployment) -> None: + """The end of ``_add_deployment``: an auto, complexity or quality marker registers its router + (adaptive ones wait for ``finalize``).""" + match classify_strategy_router_model(deployment.litellm_params.model): + case "semantic": + self._init_auto_router(deployment) + case "complexity": + self._init_complexity_router(deployment) + case "quality": + self._init_quality_router(deployment) + case _: + pass + + def release(self, deployment: Deployment) -> None: + """``_unregister_pre_routing_strategy_for_deployment``.""" + if not deployment.litellm_params.model.startswith("auto_router/"): + return + model_name: Final = deployment.model_name + tags: Final = _deployment_tags(deployment) + for registry in (self.auto, self.complexity, self.quality): + _ = _unregister(registry, model_name, tags) + if _unregister(self.adaptive, model_name, tags): + self._sync_adaptive_router_hooks() + + def release_entry(self, entry: Mapping[str, object]) -> None: + """``release`` of a deleted model_list entry; a failure only logs, as the delete has happened.""" + try: + self.release(Deployment(**entry)) # pyright: ignore[reportArgumentType] # a model_list entry + except Exception: # noqa: BLE001 # the delete already happened; a failed release only logs + verbose_router_logger.exception( + "delete_deployment: could not release pre-routing strategies for model_id=%s; " + "the deployment is out of the model_list and its indices are repaired", + as_dict(entry.get("model_info")).get("id"), + ) + + def upserted(self, deployment: Deployment, previous: Deployment | None) -> None: + """After an upsert's re-add: rebuild the adaptive routers when it touched one.""" + if _participates_in_adaptive_routing(deployment.litellm_params) or ( + previous is not None and _participates_in_adaptive_routing(previous.litellm_params) + ): + self.finalize() + + def without_capability_limit(self, action: Callable[[], object]) -> None: + """Run ``action`` without judging the capability ceiling: a rollback re-admits a serving router.""" + resolver: Final = self._capability_limit + self._capability_limit = None + try: + _ = action() + finally: + self._capability_limit = resolver + + def finalize(self) -> None: + """``_finalize_adaptive_router_if_configured``: an adaptive router for every adaptive marker, + and the adaptive companion of every complexity router configured with ``adaptive``.""" + for entry in self._hosted._core_model_list or []: # pyright: ignore[reportPrivateUsage] # the host registry's live list + params = entry.get("litellm_params") + lp_model = params.get("model") if isinstance(params, dict) else None + if not (isinstance(lp_model, str) and lp_model.startswith("auto_router/adaptive_router")): + continue + model_name = entry.get("model_name") + if not model_name or not params: + continue + deployment = Deployment( + model_name=str(model_name), + litellm_params=LiteLLM_Params.model_validate(params), + model_info=entry.get("model_info"), # pyright: ignore[reportArgumentType] # the entry's model_info + ) + if _has_registered_strategy(self.adaptive, str(model_name), _deployment_tags(deployment)): + continue + self._init_adaptive_router(deployment) + for model_name, tagged_complexity_routers in self.complexity.items(): + for tagged in tagged_complexity_routers: + complexity_router = tagged.strategy + if not complexity_router.config.adaptive: + continue + if _has_registered_strategy(self.adaptive, model_name, tagged.tags): + continue + adaptive_router = complexity_router._ensure_adaptive_router() + if adaptive_router is not None: + self.adaptive[model_name] = [ + *self.adaptive.get(model_name, []), + TaggedPreRoutingStrategy( + tags=tagged.tags, + strategy=adaptive_router, + definition_fingerprint=tagged.definition_fingerprint, + deployment=tagged.deployment, + ), + ] + self._sync_adaptive_router_hooks() + + def retire(self) -> None: + """Release the adaptive routers' post-call hooks this router registered.""" + from litellm.router_strategy.adaptive_router.hooks import AdaptiveRouterPostCallHook + + own: Final = frozenset(id(tagged.strategy) for tagged in chain.from_iterable(self.adaptive.values())) + hooks: Final[list[object]] = [ # mutable-ok: built here + hook + for hook in self._host.callbacks.get_custom_loggers_for_type(AdaptiveRouterPostCallHook) + if isinstance(hook, AdaptiveRouterPostCallHook) and id(hook.adaptive_router) in own + ] + if self._affinity is not None: + hooks.append(self._affinity) + self._affinity = None + if hooks: + self._host.callbacks.drop_callbacks(hooks) + + # -- building -------------------------------------------------------------------------------- + + def _register( + self, registry: StrategyRegistry, deployment: Deployment, strategy: object, strategy_label: str + ) -> None: + tags: Final = _deployment_tags(deployment) + if _has_registered_strategy(registry, deployment.model_name, tags): + raise ValueError( + f"{strategy_label} deployment {deployment.model_name} with tags {list(tags)} already exists. " + "Please use a different model name or set different tags." + ) + registry[deployment.model_name] = [ # rebind-ok: the live registry, edited in place as the Python router does + *registry.get(deployment.model_name, []), + TaggedPreRoutingStrategy( + tags=tags, + strategy=strategy, + definition_fingerprint=_routing_definition_fingerprint(deployment), + deployment=deployment, + ), + ] + + def _capability_violation(self, capability: GatedAutoRouterCapability) -> str | None: + resolver: Final = self._capability_limit + # `self.auto_router_capability_limit() if ... is not None`: a value that cannot be called raises. + limit: Final = resolver() if resolver is not None else None # pyright: ignore[reportCallIssue] # called as given, as Python calls it + others: Final = count_capability_routers( + (entry for entry in self._hosted._core_model_list if isinstance(entry, Mapping)), # pyright: ignore[reportPrivateUsage, reportUnnecessaryIsInstance] # the live list + capability=capability, + ) + return capability_limit_violation(capability=capability, held=others + 1, limit=limit) # pyright: ignore[reportArgumentType] # the resolver's limit + + def _init_auto_router(self, deployment: Deployment) -> None: + from litellm.router_strategy.auto_router.auto_router import AutoRouter + + params: Final = deployment.litellm_params + if params.auto_router_config_path is None and params.auto_router_config is None: + raise ValueError( + "auto_router_config_path or auto_router_config is required for auto-router deployments. Please set it in the litellm_params" + ) + if params.auto_router_default_model is None: + raise ValueError( + "auto_router_default_model is required for auto-router deployments. Please set it in the litellm_params" + ) + if params.auto_router_embedding_model is None: + raise ValueError( + "auto_router_embedding_model is required for auto-router deployments. Please set it in the litellm_params" + ) + max_input_chars: Final = params.auto_router_max_input_chars + router: Final = AutoRouter( + model_name=deployment.model_name, + auto_router_config_path=params.auto_router_config_path, + auto_router_config=params.auto_router_config, + default_model=params.auto_router_default_model, + embedding_model=params.auto_router_embedding_model, + ctx=self._ctx, # pyright: ignore[reportArgumentType] # the core's RoutingContext + max_input_chars=max_input_chars if max_input_chars is not None else DEFAULT_AUTO_ROUTER_MAX_INPUT_CHARS, + ) + self._register(self.auto, deployment, router, "Auto-router") + + def _init_complexity_router(self, deployment: Deployment) -> None: + from litellm.router_strategy.complexity_router.complexity_router import ComplexityRouter + from litellm.router_strategy.complexity_router.config import ComplexityRouterConfig + + config: Final = deployment.litellm_params.complexity_router_config + for capability in claimed_capabilities(config): + if (violation := self._capability_violation(capability)) is not None: + raise ValueError(violation) + default_model: Final = ( + ComplexityRouterConfig.model_validate(config).resolve_default_model( + deployment.litellm_params.complexity_router_default_model + ) + if config + else deployment.litellm_params.complexity_router_default_model + ) + if default_model is None: + raise ValueError( + "complexity_router_default_model is required for complexity-router deployments, " + "or configure tiers in complexity_router_config. Please set it in the litellm_params" + ) + router: Final = ComplexityRouter( + model_name=deployment.model_name, + default_model=default_model, + ctx=self._ctx, # pyright: ignore[reportArgumentType] # the core's RoutingContext + complexity_router_config=config, + ) + self._register(self.complexity, deployment, router, "Complexity-router") + if router._uses_deployment_pin: # pyright: ignore[reportPrivateUsage] # the Python router reads the same flag + self._ensure_deployment_affinity_callback() + + def _ensure_deployment_affinity_callback(self) -> None: + """``_ensure_deployment_affinity_callback``: the affinity check (every global flag off) a + complexity router's session pin needs, registered once.""" + from litellm.router_utils.pre_call_checks.deployment_affinity_check import DeploymentAffinityCheck + + if self.ensure_affinity is not None: + self.ensure_affinity() + return + if self._affinity is not None: + return + hosted: Final = self._hosted + + def is_priority_group(model: str) -> bool: + resolved: Final = hosted._core_get_model_from_alias(model=model) or model # pyright: ignore[reportPrivateUsage] # the host registry's alias lookup + group: Final = hosted.get_routing_group(resolved) + return group is not None and group.routing_strategy == "priority" + + affinity: Final = DeploymentAffinityCheck( + cache=self._cache, + ttl_seconds=3600, + enable_user_key_affinity=False, + enable_responses_api_affinity=False, + enable_session_id_affinity=False, + model_group_affinity_config=None, + is_priority_group=is_priority_group, + session_affinity_repin_on_failover=False, + ) + self._affinity = affinity + self._host.callbacks.add_litellm_callback(affinity) + + def _init_quality_router(self, deployment: Deployment) -> None: + from litellm.router_strategy.quality_router.quality_router import QualityRouter + + config: Final = deployment.litellm_params.quality_router_config + configured: Final = deployment.litellm_params.quality_router_default_model + default_model: Final = configured if configured is not None or not config else config.get("default_model") + if default_model is None: + raise ValueError( + "quality_router_default_model is required for quality-router deployments, " + "or set default_model in quality_router_config. Please configure it in the litellm_params" + ) + router: Final = QualityRouter( + model_name=deployment.model_name, + default_model=default_model, + ctx=self._ctx, # pyright: ignore[reportArgumentType] # the core's RoutingContext + quality_router_config=config, + ) + self._register(self.quality, deployment, router, "Quality-router") + + def _init_adaptive_router(self, deployment: Deployment) -> None: + from litellm.router_strategy.adaptive_router.adaptive_router import AdaptiveRouter + from litellm.router_strategy.adaptive_router.hooks import AdaptiveRouterPostCallHook + from litellm.types.router import AdaptiveRouterConfig, AdaptiveRouterPreferences + + raw_config: Final = deployment.litellm_params.adaptive_router_config + if raw_config is None: + raise ValueError("adaptive_router_config is required for adaptive-router deployments.") + config: Final = AdaptiveRouterConfig(**raw_config) + model_to_prefs: Final[dict[str, AdaptiveRouterPreferences]] = {} # mutable-ok: built here + model_to_cost: Final[dict[str, float]] = {} # mutable-ok: built here + for name in config.available_models: + deployments = self._hosted._core_deployments_named(name) # pyright: ignore[reportPrivateUsage] # the host registry's name index + if not deployments: + continue + first = deployments[0] + model_info = first.get("model_info") + info: Mapping[str, object] = model_info if isinstance(model_info, dict) else {} + prefs_raw = info.get("adaptive_router_preferences") + if isinstance(prefs_raw, dict): + model_to_prefs[name] = AdaptiveRouterPreferences(**prefs_raw) + cost = _adaptive_cost(as_dict(first.get("litellm_params")), info) + if isinstance(cost, int | float | str): + model_to_cost[name] = float(cost) + router: Final = AdaptiveRouter( + router_name=deployment.model_name, + config=config, + model_to_prefs=model_to_prefs, + model_to_cost=model_to_cost, + ) + self._register(self.adaptive, deployment, router, "Adaptive-router") + self._host.callbacks.add_litellm_callback(AdaptiveRouterPostCallHook(adaptive_router=router)) + verbose_router_logger.info( + "AdaptiveRouter[%s] initialized with %d models", deployment.model_name, len(config.available_models) + ) + + def _sync_adaptive_router_hooks(self) -> None: + """One post-call hook per registered adaptive router, plus the header hook while any exists.""" + from litellm.router_strategy.adaptive_router.hooks import ( + AdaptiveRouterPostCallHook, + AdaptiveRouterResponseHeaderHook, + ) + + callbacks: Final = self._host.callbacks + for callback in callbacks.get_custom_loggers_for_type(AdaptiveRouterPostCallHook): + callbacks.remove_callback_from_all_lists(callback) + for tagged in chain.from_iterable(self.adaptive.values()): + callbacks.add_litellm_callback(AdaptiveRouterPostCallHook(adaptive_router=tagged.strategy)) + if self.adaptive: + callbacks.add_litellm_callback(AdaptiveRouterResponseHeaderHook()) + + +# -- the request side -------------------------------------------------------------------------- + + +def _stamp_or_clear_metadata_key( + request_kwargs: dict[str, object], # mutable-ok: the request kwargs, stamped in place + key: str, + value: object | None, +) -> None: + """Write a proxy-internal metadata key for this routing attempt, or clear it from both buckets.""" + if value is None: + for bucket in (request_kwargs.get("metadata"), request_kwargs.get("litellm_metadata")): + if isinstance(bucket, dict): + bucket.pop(key, None) + return + _, metadata_bucket = get_or_create_metadata_bucket(request_kwargs) + metadata_bucket[key] = value + + +def _redact_prompt_text_if_needed( + request_kwargs: Mapping[str, object], routing_decision: StandardLoggingRoutingDecision +) -> StandardLoggingRoutingDecision: + from litellm.litellm_core_utils.redact_messages import should_redact_message_logging + + if not should_redact_message_logging( + { + "litellm_params": request_kwargs, + "standard_callback_dynamic_params": request_kwargs.get("standard_callback_dynamic_params"), + } + ): + return routing_decision + kept: Final = without(routing_decision, *PROMPT_QUOTING_ROUTING_DECISION_FIELDS) + return cast(StandardLoggingRoutingDecision, kept) # cast-ok: dropping optional keys preserves the type + + +def _record_routing_decision( + request_kwargs: dict[str, object], # mutable-ok: the request kwargs + routing_decision: StandardLoggingRoutingDecision | None, +) -> None: + """Make the request's metadata describe this routing attempt, and only this one.""" + from litellm.router_strategy.auto_router.baseline_request import capture_baseline_parameters + from litellm.types.router import BaselineRouteStamp + + phase_attributes(routing_decision_attributes(routing_decision)) + baseline_model: Final = routing_decision.get("savings_baseline_model") if routing_decision else None + baseline_id: Final = routing_decision.get("savings_baseline_deployment_id") if routing_decision else None + router_name: Final = routing_decision.get("router_model_name") if routing_decision else None + caller_parameters: Final = capture_baseline_parameters(request_kwargs) if router_name and baseline_model else None + _stamp_or_clear_metadata_key( + request_kwargs, + "_autorouter_baseline_route", + BaselineRouteStamp(router_name, baseline_model, baseline_id, caller_parameters) + if router_name and baseline_model and baseline_id + else None, + ) + _stamp_or_clear_metadata_key( + request_kwargs, + "routing_decision", + None if routing_decision is None else _redact_prompt_text_if_needed(request_kwargs, routing_decision), + ) + + +def _as_output_cap(value: object) -> int | None: + if isinstance(value, bool) or not isinstance(value, int | float | str): + return None + try: + cap: Final = int(float(value)) + except (ValueError, OverflowError): + return None + return cap if cap >= 0 else None + + +def _stamped_client_ceiling(request_kwargs: Mapping[str, object]) -> Mapping[str, int] | None: + stamped: Final = next( + ( + bucket.get(CLIENT_OUTPUT_CEILING_METADATA_KEY) + for bucket in (request_kwargs.get("metadata"), request_kwargs.get("litellm_metadata")) + if isinstance(bucket, dict) and CLIENT_OUTPUT_CEILING_METADATA_KEY in bucket + ), + None, + ) + if not isinstance(stamped, dict): + return None + return MappingProxyType( + { + str(carrier): cap + for carrier, value in stamped.items() + if carrier in OUTPUT_TOKEN_CEILING_PARAMS and (cap := _as_output_cap(value)) is not None + } + ) + + +def _restore_client_output_ceiling( + request_kwargs: MutableMapping[str, object], # mutable-ok: the request kwargs, restored in place +) -> None: + callers_ceiling: Final = _stamped_client_ceiling(request_kwargs) + if callers_ceiling is None: + return + for carrier in OUTPUT_TOKEN_CEILING_PARAMS: + request_kwargs.pop(carrier, None) + request_kwargs.update(callers_ceiling) + + +def _stamp_client_output_ceiling(request_kwargs: dict[str, object]) -> None: # mutable-ok: the request kwargs + _, metadata_bucket = get_or_create_metadata_bucket(request_kwargs) + metadata_bucket.setdefault( + CLIENT_OUTPUT_CEILING_METADATA_KEY, + {carrier: request_kwargs[carrier] for carrier in OUTPUT_TOKEN_CEILING_PARAMS if carrier in request_kwargs}, + ) + + +def _pop_effort_from_nested_carrier(request_kwargs: dict[str, object], carrier: str) -> None: # mutable-ok: same + nested: Final = request_kwargs.get(carrier) + if not isinstance(nested, Mapping): + return + sanitized: Final = without(nested, "effort") + if sanitized: + request_kwargs[carrier] = sanitized # rebind-ok: copy-on-write of the request's nested carrier + else: + request_kwargs.pop(carrier, None) + + +def _tier_ceiling_under_the_surface_name( + tier_litellm_params: Mapping[str, object], responses_call: bool +) -> Mapping[str, object]: + surface_key: Final = "max_output_tokens" if responses_call else "max_tokens" + carried: Final = tuple( + key + for key in (surface_key, "max_tokens", "max_completion_tokens", "max_output_tokens") + if key in tier_litellm_params + ) + if not carried: + return tier_litellm_params + return MappingProxyType( + { + **without(tier_litellm_params, *OUTPUT_TOKEN_CEILING_PARAMS), + surface_key: tier_litellm_params[carried[0]], + } + ) + + +def _drop_client_carriers_a_tier_pin_supersedes( + request_kwargs: dict[str, object], # mutable-ok: the request kwargs + tier_litellm_params: Mapping[str, object], +) -> None: + if not OUTPUT_TOKEN_CEILING_PARAMS.isdisjoint(tier_litellm_params): + _stamp_client_output_ceiling(request_kwargs) + for carrier in OUTPUT_TOKEN_CEILING_PARAMS: + request_kwargs.pop(carrier, None) + if "reasoning_effort" not in tier_litellm_params: + return + request_kwargs.pop("thinking", None) + _pop_effort_from_nested_carrier(request_kwargs, "output_config") + _pop_effort_from_nested_carrier(request_kwargs, "reasoning") + + +def _request_header(request_kwargs: Mapping[str, object], header_name: str) -> str | None: + proxy_server_request: Final = request_kwargs.get("proxy_server_request") + if not isinstance(proxy_server_request, Mapping): + return None + headers: Final = proxy_server_request.get("headers") + if not isinstance(headers, Mapping): + return None + return next( + ( + value + for key, value in headers.items() + if isinstance(key, str) and key.lower() == header_name and isinstance(value, str) + ), + None, + ) + + +def _session_router_cache_key(request_kwargs: Mapping[str, object]) -> str | None: + session_id: Final = _request_header(request_kwargs, "x-claude-code-session-id") + if session_id is None or _CLAUDE_CODE_SESSION_ID_RE.fullmatch(session_id) is None: + return None + metadata_name: Final = "litellm_metadata" if "litellm_metadata" in request_kwargs else "metadata" + metadata: Final = request_kwargs.get(metadata_name) + if not isinstance(metadata, Mapping): + return None + caller_scope: Final = metadata.get("user_api_key_hash") + if not isinstance(caller_scope, str) or not caller_scope: + return None + return f"claude_code_session_router:v1:{caller_scope}:{session_id}" + + +class PreRouter: + """Answers the Rust core's ``pre_route`` steps (``_core_async_pre_routing_hook``).""" + + def __init__( + self, + strategies: PreRoutingStrategies, + hosted: "HostRegistry", + host: "PythonHost", + ctx: object, + routing_cache: DualCache, + ) -> None: + self._strategies: Final = strategies + self._hosted: Final = hosted + self._host: Final = host + self._ctx: Final = ctx + self._session_cache: Final = DualCache(redis_cache=routing_cache.redis_cache, in_memory_cache=InMemoryCache()) + + # -- the step -------------------------------------------------------------------------------- + + async def pre_route( + self, + model: str, + request_kwargs: dict[str, object], # mutable-ok: the request kwargs the hook stamps + *, + messages: object, + input: object, + fallback_depth: int, + router_strategy: str | None, + enable_tag_filtering: bool, + match_any: bool, + public_router: object, + run_plugins: Callable[[str], Awaitable[None]] | None = None, + pass_through: bool = False, + ) -> str | None: + """The start of ``_async_get_available_deployment``: run the hook, record its pick and pin the + tier's litellm_params. The model group to resolve, or ``None`` to resolve ``model`` itself. + ``pass_through`` is ``_async_get_available_deployment_for_pass_through``, which runs the hook + whatever the router's strategy.""" + if not pass_through and not picks_async(router_strategy): + return None + responses_call: Final = input is not None and messages is None + response: Final = await self._hook( + model, + request_kwargs, + messages=messages, + input=input, + fallback_depth=fallback_depth, + enable_tag_filtering=enable_tag_filtering, + match_any=match_any, + public_router=public_router, + run_plugins=run_plugins, + ) + routed: Final = response.model if response is not None else model + if response is not None: + record_pre_routing_selection(request_kwargs, routed) + self._pin_tier_params( + routed, response.litellm_params if response is not None else None, request_kwargs, responses_call + ) + return routed if response is not None else None + + # -- the hook -------------------------------------------------------------------------------- + + async def _hook( + self, + model: str, + request_kwargs: dict[str, object], # mutable-ok: see pre_route + *, + messages: object, + input: object, + fallback_depth: int, + enable_tag_filtering: bool, + match_any: bool, + public_router: object, + run_plugins: Callable[[str], Awaitable[None]] | None, + ) -> "PreRoutingHookResponse | None": + _restore_client_output_ceiling(request_kwargs) + _record_routing_decision(request_kwargs, None) + requested: Final = self._hosted._core_get_model_from_alias(model=model) or model # pyright: ignore[reportPrivateUsage] # the host registry's alias lookup + registered: Final = await self._resolve_session_router( + model, requested, request_kwargs, fallback_depth, enable_tag_filtering, match_any + ) + # The routing plugins run on the session-resolved name, before a strategy is selected, so a + # strategy reads their signals; router-core narrows the pick to what they left. + if run_plugins is not None: + await run_plugins(registered) + selected: Final = self._select_strategy(registered, request_kwargs, enable_tag_filtering, match_any) + if selected is None: + await arm_compaction(request_kwargs, None) + _stamp_or_clear_metadata_key(request_kwargs, SESSION_DEPLOYMENT_AFFINITY_TTL_METADATA_KEY, None) + _stamp_or_clear_metadata_key(request_kwargs, CONSUMED_REQUEST_TAGS_METADATA_KEY, None) + return None + from litellm.router_strategy.complexity_router.complexity_router import ComplexityRouter + + marker: Final = self._selected_marker(registered, selected.tags, request_kwargs) + definition: Final = selected.deployment + config_id: Final = definition.model_info.id if definition else None + updated_at: Final = definition.model_info.updated_at if definition else None + router_kind: Final = classify_strategy_router_model(definition.litellm_params.model) if definition else None + provenance: Final[StandardLoggingRoutingDecision] = { + "router_model_name": definition.model_name if definition else registered, + **({"router_type": router_kind} if router_kind is not None else {}), + **({"router_config_id": config_id} if config_id is not None else {}), + **({"router_config_updated_at": updated_at.isoformat()} if updated_at is not None else {}), + **( + {"router_config_fingerprint": selected.definition_fingerprint} + if selected.definition_fingerprint is not None + else {} + ), + } # pyright: ignore[reportAssignmentType] # the optional keys, as the Python router builds them + phase_attributes(routing_decision_attributes(provenance)) + reject_recursive_compactor(registered) + strategy: Final = selected.strategy + complexity: Final = strategy if isinstance(strategy, ComplexityRouter) else None + await arm_compaction( + request_kwargs, + complexity.config.context_compaction if complexity is not None else None, + _tier_members(complexity.config.tiers) if complexity is not None else (), # pyright: ignore[reportArgumentType] # the tiers' model names + parent_model=model, + ctx=self._ctx, # pyright: ignore[reportArgumentType] # the core's RoutingContext + allow_escalation=complexity is not None and complexity.config.enable_context_window_escalation, + messages=messages, # pyright: ignore[reportArgumentType] # the request's messages, as given + ) + await router_host().authorize_strategy_inference( + deployment=marker, + request_kwargs=request_kwargs, + llm_router=public_router, # pyright: ignore[reportArgumentType] # the litellm.Router this core serves + ) + compression_policy: Final = router_host().compression_policy( + llm_router=public_router, # pyright: ignore[reportArgumentType] # same + model_alias=registered, + request_kwargs=request_kwargs, + request_tags=get_routing_tags_from_request_kwargs(request_kwargs), + ) + independent_compression: Final = compression_policy is not None and not ( + compression_policy.is_same + and compression_policy.model is not None + and router_host().model_hop_compression_armed() + ) + routing_messages: Final = ( + await router_host().messages_for_routing( + policy=compression_policy, + messages=messages, # pyright: ignore[reportArgumentType] # same + request_kwargs=request_kwargs, + ) + if independent_compression + else None + ) + routed: Final = await strategy.async_pre_routing_hook( + model=registered, + request_kwargs=request_kwargs, + messages=routing_messages if routing_messages is not None else messages, # pyright: ignore[reportArgumentType] # the request's messages, as given + input=input, # pyright: ignore[reportArgumentType] # the request's input, as given + specific_deployment=False, + ) + hooked: Final[PreRoutingHookResponse | None] = routed + traced: Final = ( + hooked.model_copy(update={"routing_decision": {**provenance, **hooked.routing_decision}}) + if hooked is not None and hooked.routing_decision is not None + else hooked + ) + response: Final = ( + traced.model_copy(update={"messages": messages}) + if traced is not None and routing_messages is not None and traced.messages == routing_messages + else traced + ) + _record_routing_decision(request_kwargs, response.routing_decision if response else None) + _stamp_or_clear_metadata_key( + request_kwargs, + SESSION_DEPLOYMENT_AFFINITY_TTL_METADATA_KEY, + response.session_affinity_ttl_seconds if response else None, + ) + request_tags: Final = get_routing_tags_from_request_kwargs(request_kwargs) + consumed: Final = ( + ConsumedRequestTagsStamp(model_group=response.model, tags=selected.tags) + if response is not None + and selected.tags + and request_tags + and is_valid_deployment_tag(selected.tags, request_tags, match_any) # pyright: ignore[reportArgumentType] # the tags as the Python router passes them + else None + ) + _stamp_or_clear_metadata_key(request_kwargs, CONSUMED_REQUEST_TAGS_METADATA_KEY, consumed) + marker_params: Final = ( + self._forwardable_marker_params(registered, selected.tags, request_kwargs) if response is not None else () + ) + tier_param_keys: Final = tuple(response.litellm_params or ()) if response is not None else () + overridden: Final = frozenset(tier_param_keys) | ( + OUTPUT_TOKEN_CEILING_PARAMS if not OUTPUT_TOKEN_CEILING_PARAMS.isdisjoint(tier_param_keys) else frozenset() + ) + newly_forwarded: Final = tuple( + (key, value) for key, value in marker_params if key not in request_kwargs and key not in overridden + ) + if response is not None: + _stamp_client_output_ceiling(request_kwargs) + request_kwargs.pop(_ALIAS_MARKER_FORWARDED_PARAMS_KWARG, None) + request_kwargs.update(newly_forwarded) + if newly_forwarded: + request_kwargs.update({_ALIAS_MARKER_FORWARDED_PARAMS_KWARG: tuple(key for key, _ in newly_forwarded)}) + return response + + # -- strategy selection ---------------------------------------------------------------------- + + def _deployments_for_request( + self, model: str, request_kwargs: Mapping[str, object] + ) -> Sequence[Mapping[str, object]]: + return self._hosted._core_deployments_for_request(model, request_kwargs) # pyright: ignore[reportPrivateUsage] # the host registry's lookup + + def _select_strategy( + self, model: str, request_kwargs: Mapping[str, object], enable_tag_filtering: bool, match_any: bool + ) -> "TaggedPreRoutingStrategy[Any] | None": + """``_select_pre_routing_strategy``.""" + registries: Final = self._strategies.all() + if not any(registries): + return None + deployments: Final = self._deployments_for_request(model, request_kwargs) + registered_names: Final = tuple(dict.fromkeys(str(d.get("model_name")) for d in deployments)) or (model,) + candidates: Final = ( + tuple( # comprehension-ok: registry by registry, name by name, as the Python router searches + tagged for registry in registries for name in registered_names for tagged in registry.get(name, []) + ) + ) + if not candidates: + return None + request_tags: Final = get_routing_tags_from_request_kwargs(request_kwargs) + if request_tags: + for tagged in candidates: + if tagged.tags and is_valid_deployment_tag(list(tagged.tags), request_tags, match_any): + return tagged + for tagged in candidates: + if "default" in tagged.tags: + return tagged + request_scoped: Final = request_kwargs.get("enable_tag_filtering") is True + if ( + (enable_tag_filtering or request_scoped) + and all(tagged.tags for tagged in candidates) + and any(not is_strategy_marker(d) for d in deployments) + ): + return None + return candidates[0] + + def _selected_marker( + self, model: str, strategy_tags: tuple[str, ...], request_kwargs: Mapping[str, object] + ) -> Mapping[str, object] | None: + markers: Final = tuple( + deployment + for deployment in self._deployments_for_request(model, request_kwargs) + if str(as_dict(deployment.get("litellm_params")).get("model", "")).startswith(AUTO_ROUTER_MODEL_PREFIX) + and "model" in as_dict(deployment.get("litellm_params")) + ) + tag_matched: Final = tuple( + deployment + for deployment in markers + if tuple(_tags_of(as_dict(deployment.get("litellm_params")))) == strategy_tags + ) + return tag_matched[0] if tag_matched else (markers[0] if markers else None) + + def _forwardable_marker_params( + self, model: str, strategy_tags: tuple[str, ...], request_kwargs: Mapping[str, object] + ) -> tuple[tuple[str, object], ...]: + marker: Final = self._selected_marker(model, strategy_tags, request_kwargs) + if marker is None: + return () + return tuple( + (key, value) + for key, value in as_dict(marker.get("litellm_params")).items() + if key not in _ALIAS_PARAMS_NEVER_FORWARDED + and key not in CustomPricingLiteLLMParams.model_fields + and value is not None + ) + + # -- the Claude Code session binding --------------------------------------------------------- + + @with_service_target(CLAUDE_CODE_SESSION_ROUTER_BINDING_TARGET) + async def _delete_binding(self, cache_key: str) -> None: + try: + await self._session_cache.async_delete_cache(key=cache_key) + except Exception as e: # noqa: BLE001 # cache cleanup must not fail an otherwise routable request + verbose_router_logger.warning( + "Failed to delete Claude Code session router binding; the binding may remain until its TTL expires: %s", + e, + ) + + @with_service_target(CLAUDE_CODE_SESSION_ROUTER_BINDING_TARGET) + async def _get_binding(self, cache_key: str) -> object: + try: + if self._session_cache.redis_cache is None: + return await self._session_cache.async_get_cache(key=cache_key) + return await self._session_cache.redis_cache.async_get_cache(key=cache_key) + except Exception as e: # noqa: BLE001 # an optional binding must not make routing depend on Redis + log_redis_failure( + verbose_router_logger, + logging.WARNING, + "Failed to read Claude Code session router binding; using the requested model", + e, + ) + return None + + @with_service_target(CLAUDE_CODE_SESSION_ROUTER_BINDING_TARGET) + async def _resolve_session_router( + self, + model: str, + registered: str, + request_kwargs: dict[str, object], # mutable-ok: see pre_route + fallback_depth: int, + enable_tag_filtering: bool, + match_any: bool, + ) -> str: + """``_resolve_claude_code_session_router``.""" + if is_native_compaction_call() or not self._strategies.any_registered(): + return registered + cache_key: Final = _session_router_cache_key(request_kwargs) + if cache_key is None: + return registered + depth: Final = fallback_depth or request_kwargs.get("fallback_depth") + if depth not in (None, 0): + return registered + agent_id: Final = _request_header(request_kwargs, "x-claude-code-agent-id") + if agent_id is not None: + bound_model: Final = await self._get_binding(cache_key) + if not isinstance(bound_model, str): + return registered + bound_registered: Final = self._hosted._core_get_model_from_alias(model=bound_model) or bound_model # pyright: ignore[reportPrivateUsage] # see _hook + if self._select_strategy(bound_registered, request_kwargs, enable_tag_filtering, match_any) is None: + await self._delete_binding(cache_key) + return registered + await self._session_cache.async_set_cache( + key=cache_key, value=bound_model, ttl=_CLAUDE_CODE_SESSION_ROUTER_TTL_SECONDS + ) + _stamp_or_clear_metadata_key(request_kwargs, "model_group", bound_model) + return bound_registered + if self._select_strategy(registered, request_kwargs, enable_tag_filtering, match_any) is None: + return registered + await self._session_cache.async_set_cache( + key=cache_key, value=model, ttl=_CLAUDE_CODE_SESSION_ROUTER_TTL_SECONDS + ) + return registered + + # -- tier litellm_params --------------------------------------------------------------------- + + def _tier_params_the_target_accepts( + self, model: str, tier_params: Mapping[str, object], request_kwargs: Mapping[str, object] + ) -> Mapping[str, object]: + from litellm.router_rust.invoke.shaping import deployment_accepts_param + + deployments: Final = self._hosted._core_get_model_list(model_name=model) or () # pyright: ignore[reportPrivateUsage] # the host registry's lookup + if not deployments: + return tier_params + allowlisted: Final = _allowlist(tier_params) | _allowlist(request_kwargs) + candidates: Final = provider_rejectable_params(tier_params) - TIER_PARAMS_NEVER_DROPPED - allowlisted + unsupported: Final = frozenset( + param + for param in candidates + if not any(deployment_accepts_param(deployment, model, param, self._host) for deployment in deployments) + ) + if not unsupported: + return tier_params + verbose_router_logger.warning( + "litellm.router.py: dropping tier params %s for model=%s, no deployment behind it declares them", + ", ".join(sorted(unsupported)), + model, + ) + return MappingProxyType(without(tier_params, *unsupported)) + + def _pin_tier_params( + self, + model: str, + tier_litellm_params: Mapping[str, object] | None, + request_kwargs: dict[str, object], # mutable-ok: see pre_route + responses_call: bool, + ) -> None: + """``_pin_tier_params_onto_request``.""" + if not tier_litellm_params: + return + reject_api_key_secret_reference(tier_litellm_params) + accepted: Final = self._tier_params_the_target_accepts(model, tier_litellm_params, request_kwargs) + surface: Final = _tier_ceiling_under_the_surface_name(accepted, responses_call=responses_call) + _drop_client_carriers_a_tier_pin_supersedes(request_kwargs, surface) + request_kwargs.update(surface) + + +def as_dict(value: object) -> Mapping[str, object]: + if not isinstance(value, Mapping): + return {} + return cast(Mapping[str, object], value) # cast-ok: a str-keyed config or request dict + + +def _adaptive_cost(params: Mapping[str, object], info: Mapping[str, object]) -> object: + """A model's input price for the adaptive router: litellm_params win over model_info.""" + cost: Final = params.get("input_cost_per_token") + return cost if cost is not None else info.get("input_cost_per_token") + + +def _tier_members(tiers: Mapping[str, object]) -> tuple[object, ...]: + """Every model the tier pools name, once, in order.""" + pools: Final = ( + (pool,) if isinstance(pool, str) or not isinstance(pool, Iterable) else pool for pool in tiers.values() + ) + return tuple(dict.fromkeys(chain.from_iterable(pools))) + + +def _tags_of(params: Mapping[str, object]) -> Sequence[object]: + tags: Final = params.get("tags") + return tags if isinstance(tags, list | tuple) else () + + +def _allowlist(params: Mapping[str, object]) -> frozenset[str]: + declared: Final = params.get("allowed_openai_params") + if not isinstance(declared, list | tuple | set | frozenset): + return frozenset() + return frozenset(entry for entry in declared if isinstance(entry, str)) + + +__all__ = ["PreRouter", "PreRoutingStrategies", "is_strategy_marker"] diff --git a/litellm/router_rust/steps/selection.py b/litellm/router_rust/steps/selection.py new file mode 100644 index 00000000000..183350bc05f --- /dev/null +++ b/litellm/router_rust/steps/selection.py @@ -0,0 +1,445 @@ +""" +Host-side strategy selectors for ``RustCore``. + +The Rust core decides which strategy picks and whose state it reads (``strategy::routing_context``: +a request override, then the routing group, then the router's own strategy) and asks the host to +run the stateful strategies it does not implement natively through a ``select`` step. The +strategies themselves are the standalone ``litellm/router_strategy`` loggers, which double as +``CustomLogger`` callbacks that record traffic in the host's routing cache. This module builds, +registers and retires them the way the Python router's ``routing_strategy_init``, +``_init_routing_groups`` and ``_get_override_strategy_selector`` do. It never reaches a Python +router core. +""" + +import threading +from collections.abc import Iterable, Mapping, Sequence +from itertools import chain +from types import MappingProxyType +from typing import TYPE_CHECKING, Final, TypeAlias + +from litellm.integrations.custom_logger import CustomLogger +from litellm.litellm_core_utils.core_helpers import as_str_mapping +from litellm.router_strategy.base_routing_strategy import BaseRoutingStrategy +from litellm.router_strategy.least_busy import ( + CALL_IDENTITY, + SELECTION_RESERVATION_MAX_KEYS, + LeastBusyLoggingHandler, + call_id_of, +) +from litellm.router_strategy.lowest_cost import LowestCostLoggingHandler +from litellm.router_strategy.lowest_latency import LowestLatencyLoggingHandler +from litellm.router_strategy.lowest_tpm_rpm import LowestTPMLoggingHandler +from litellm.router_strategy.lowest_tpm_rpm_v2 import LowestTPMLoggingHandler_v2 +from litellm.types.router import RoutingGroup, RoutingStrategy + +if TYPE_CHECKING: + from litellm.caching.dual_cache import DualCache + from litellm.caching.in_memory_cache import InMemoryCache + from litellm.router_python_host import PythonHost + +Selector: TypeAlias = ( + LeastBusyLoggingHandler + | LowestCostLoggingHandler + | LowestLatencyLoggingHandler + | LowestTPMLoggingHandler + | LowestTPMLoggingHandler_v2 +) + +# The legacy router attribute that holds the default group's selector, per strategy. +SELECTOR_ATTRIBUTES: Final = MappingProxyType( + { + "least-busy": "leastbusy_logger", + "usage-based-routing": "lowesttpm_logger", + "usage-based-routing-v2": "lowesttpm_logger_v2", + "latency-based-routing": "lowestlatency_logger", + "cost-based-routing": "lowestcost_logger", + } +) + + +def normalize_strategy(strategy: object) -> str | None: + """A strategy as the string Python compares against (``_normalize_strategy``).""" + if isinstance(strategy, RoutingStrategy): + return strategy.value + return strategy if isinstance(strategy, str) else None + + +def build_selector( + strategy: str | None, + args: Mapping[str, object], + cache: "DualCache", + reservations: "InMemoryCache | None" = None, +) -> Selector | None: + """``_build_strategy_selector`` without registering: ``None`` for simple-shuffle and every + strategy without a selector. Every least-busy selector of a router keeps its pick + reservations in the router's one ``reservations`` store.""" + routing_args: Final = dict(args) + match strategy: + case "least-busy": + return LeastBusyLoggingHandler(router_cache=cache, reservation_store=reservations) + case "usage-based-routing": + return LowestTPMLoggingHandler(router_cache=cache, routing_args=routing_args) + case "usage-based-routing-v2": + return LowestTPMLoggingHandler_v2(router_cache=cache, routing_args=routing_args) + case "latency-based-routing": + return LowestLatencyLoggingHandler(router_cache=cache, routing_args=routing_args) + case "cost-based-routing": + return LowestCostLoggingHandler(router_cache=cache, routing_args={}) + case _: + return None + + +def _callbacks(host: "PythonHost") -> list[object]: # mutable-ok: the process list itself + """``litellm.callbacks``, which the host serves like its other callback lists.""" + return host.callback_list("callbacks") # pyright: ignore[reportArgumentType] # not in the CallbackList literal + + +class StrategySelectors: + """The selectors of one router: the default group's, each routing group's, and the lazily + built per-request override ones.""" + + def __init__(self, host: "PythonHost", cache: "DualCache") -> None: + from litellm.caching.in_memory_cache import InMemoryCache + + self._host: Final = host + self._cache: Final = cache + # `_core_least_busy_reservations`: the picks least-busy reserved, shared by its selectors. + self.reservations: Final = InMemoryCache(max_size_in_memory=SELECTION_RESERVATION_MAX_KEYS) + self._strategy: str | None = None + # The router's `routing_strategy`, which a direct `routing_strategy_init` leaves as it was. + self._configured: str | None = None + self._default: Selector | None = None + self._groups: dict[str, Selector] = {} # mutable-ok: replaced as routing groups change + self._group_strategies: frozenset[str] = frozenset() + self._overrides: dict[str, Selector | None] = {} # mutable-ok: filled lazily per strategy + self._lock: Final = threading.Lock() + + # -- building ------------------------------------------------------------------------------ + + def _register(self, selector: CustomLogger) -> None: + """``_register_router_selector``: least-busy also counts calls as they start.""" + host: Final = self._host + if isinstance(selector, LeastBusyLoggingHandler): + if isinstance(host.callback_list("input_callback"), list): # pyright: ignore[reportUnnecessaryIsInstance] # a process list a user may have reassigned + host.callbacks.add_litellm_input_callback(selector) + else: + host.set_callback_list("input_callback", [selector]) + if isinstance(_callbacks(host), list): # pyright: ignore[reportUnnecessaryIsInstance] # see above + host.callbacks.add_litellm_callback(selector) + + def _unregister(self, selectors: Sequence[object]) -> None: + live: Final = [selector for selector in selectors if selector is not None] + for selector in live: + if isinstance(selector, BaseRoutingStrategy): + selector.retire() + if live: + self._host.callbacks.drop_callbacks(live) + + def init_default(self, strategy: object, args: Mapping[str, object], *, configured: bool = True) -> None: + """``routing_strategy_init``: drop the old default and override selectors, build the new one. + ``configured``: ``strategy`` is the router's ``routing_strategy`` from now on.""" + self._unregister([self._default, *self._overrides.values()]) + self._overrides = {} + self._default = None + self._strategy = normalize_strategy(strategy) + if configured: + self._configured = self._strategy + self._default = build_selector(self._strategy, args, self._cache, self.reservations) + if self._default is not None: + self._register(self._default) + + def rebuild_default(self, args: Mapping[str, object]) -> None: + """``_apply_updated_routing_strategy_args``: relink the default selector to new args, keeping + the old one when they are invalid.""" + from pydantic import ValidationError + + from litellm._logging import verbose_router_logger + + if self._default is None or self._strategy != self._configured: + return + try: + rebuilt: Final = build_selector(self._strategy, args, self._cache, self.reservations) + except (TypeError, ValidationError): + verbose_router_logger.exception( + "Invalid routing_strategy_args %s for '%s'; keeping the previous ones", dict(args), self._strategy + ) + return + self._unregister([self._default]) + self._default = rebuilt + if rebuilt is not None: + self._register(rebuilt) + + def init_groups(self, groups: Sequence[tuple[str, str, Mapping[str, object]]]) -> None: + """``_replace_routing_groups``: one selector per group, from (name, strategy, args).""" + self._unregister(list(self._groups.values())) + built: Final = { + name: selector + for name, strategy, args in groups + if ( + selector := build_selector( + None if strategy == "priority" else strategy, args, self._cache, self.reservations + ) + ) + is not None + } + for selector in built.values(): + self._register(selector) + self._groups = built + self._group_strategies = frozenset(strategy for _, strategy, _ in groups) + + def retire(self) -> None: + self._unregister([self._default, *self._groups.values(), *self._overrides.values()]) + self.clear() + + def clear(self) -> None: + """Forget every selector, after the caller released them from the host lists.""" + self._default = None + self._groups = {} + self._overrides = {} + + def callback_selectors(self) -> tuple[Selector, ...]: + """``_callback_selectors``: the default group's selector and each routing group's.""" + return tuple(selector for selector in (self._default, *self._groups.values()) if selector is not None) + + def override_selectors(self) -> tuple[Selector, ...]: + """The per-request override selectors built so far.""" + return tuple(selector for selector in self._overrides.values() if selector is not None) + + @property + def default(self) -> Selector | None: + """The default selector of the router's ``routing_strategy``, if it has one.""" + return self._default if self._strategy == self._configured else None + + def attribute(self, name: str) -> Selector | None: + """A legacy ``*_logger`` attribute: the default selector when it is that strategy's.""" + return self._default if SELECTOR_ATTRIBUTES.get(self._strategy or "") == name else None + + def latency_args(self) -> Mapping[str, object] | None: + """The default latency selector's full routing args, as ``get_settings`` reports them.""" + if not isinstance(self._default, LowestLatencyLoggingHandler): + return None + reported: Final[object] = self._default.routing_args.json() # pyright: ignore[reportUnknownMemberType] # RoutingArgs.json is untyped + return as_str_mapping(reported) + + # -- picking ------------------------------------------------------------------------------- + + def selector_for( + self, strategy: str, scope: str, group: str | None, request: Mapping[str, object] | None + ) -> Selector | None: + """The selector a Rust ``select`` step names (``_get_routing_context``).""" + if scope == "group": + return self._groups.get(group or "") + if scope == "override": + return self._override(strategy, request) + return self._default if strategy == self._strategy else None + + def _override(self, strategy: str, request: Mapping[str, object] | None) -> Selector | None: + """``_get_override_strategy_selector`` and ``_bind_override_selector_to_request``.""" + if strategy == self._configured: + return self.default + with self._lock: + if strategy not in self._overrides: + self._overrides[strategy] = build_selector(strategy, {}, self._cache, self.reservations) + selector: Final = self._overrides[strategy] + if selector is not None and request is not None and strategy not in self._registered_strategies(): + from litellm.litellm_core_utils.litellm_logging import Logging + + logging_obj: Final = request.get("litellm_logging_obj") + if isinstance(logging_obj, Logging): + logging_obj.add_dynamic_callback(selector) + return selector + + def _registered_strategies(self) -> frozenset[str]: + return frozenset({*((self._configured,) if self._configured else ()), *self._group_strategies}) + + def is_override(self, strategy: str, selector: object) -> bool: + """Whether ``selector`` is a per-request override one, whose pre-call check the router runs + itself because it is not in the callback lists.""" + return selector is not None and self._overrides.get(strategy) is selector + + +async def pick_async( + strategy: str, + selector: Selector, + model: str, + deployments: list[dict[str, object]], # mutable-ok: the selectors take a list of deployment dicts + messages: object, + input: object, + request_kwargs: dict[str, object] | None, # mutable-ok: the latency selector reads it as given + reserve: bool = False, +) -> object: + """``_select_deployment_async``; ``reserve``: least-busy holds the pick for the request's call.""" + match selector: + case LeastBusyLoggingHandler() if strategy == "least-busy": + return await selector.async_get_available_deployments( + model_group=model, + healthy_deployments=deployments, + reserve=reserve, + call_id=call_id_of(CALL_IDENTITY.validate_python(request_kwargs or {})), + ) # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] # the selectors return untyped deployment dicts + case LowestTPMLoggingHandler() if strategy == "usage-based-routing": + return selector.get_available_deployments( # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] # the selectors return untyped deployment dicts + model_group=model, + healthy_deployments=deployments, + messages=messages, # pyright: ignore[reportArgumentType] # the request's messages, as given + input=input, # pyright: ignore[reportArgumentType] # same + ) + case LowestTPMLoggingHandler_v2() | LowestCostLoggingHandler() if strategy in ( + "usage-based-routing-v2", + "cost-based-routing", + ): + return await selector.async_get_available_deployments( # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] # the selectors return untyped deployment dicts + model_group=model, + healthy_deployments=deployments, + messages=messages, # pyright: ignore[reportArgumentType] # same + input=input, # pyright: ignore[reportArgumentType] # same + ) + case LowestLatencyLoggingHandler() if strategy == "latency-based-routing": + return await selector.async_get_available_deployments( # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] # the selectors return untyped deployment dicts + model_group=model, + healthy_deployments=deployments, + messages=messages, # pyright: ignore[reportArgumentType] # same + input=input, # pyright: ignore[reportArgumentType] # same + request_kwargs=request_kwargs, + ) + case _: + return None + + +def pick_sync( + strategy: str, + selector: Selector, + model: str, + deployments: list[dict[str, object]], # mutable-ok: see pick_async + messages: object, + input: object, + request_kwargs: dict[str, object] | None, # mutable-ok: see pick_async + reserve: bool = False, +) -> object: + """``_select_deployment_sync``: cost-based routing has no sync pick.""" + match selector: + case LeastBusyLoggingHandler() if strategy == "least-busy": + return selector.get_available_deployments( + model_group=model, + healthy_deployments=deployments, + reserve=reserve, + call_id=call_id_of(CALL_IDENTITY.validate_python(request_kwargs or {})), + ) # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] # the selectors return untyped deployment dicts + case LowestTPMLoggingHandler() | LowestTPMLoggingHandler_v2() if strategy in ( + "usage-based-routing", + "usage-based-routing-v2", + ): + return selector.get_available_deployments( # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] # the selectors return untyped deployment dicts + model_group=model, + healthy_deployments=deployments, + messages=messages, # pyright: ignore[reportArgumentType] # same + input=input, # pyright: ignore[reportArgumentType] # same + ) + case LowestLatencyLoggingHandler() if strategy == "latency-based-routing": + return selector.get_available_deployments( # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] # the selectors return untyped deployment dicts + model_group=model, + healthy_deployments=deployments, + messages=messages, # pyright: ignore[reportArgumentType] # same + input=input, # pyright: ignore[reportArgumentType] # same + request_kwargs=request_kwargs, + ) + case _: + return None + + +async def routing_pre_call_checks(callbacks: Iterable[object], deployment: Mapping[str, object]) -> None: + """``async_routing_strategy_pre_call_checks``: every ``CustomLogger`` among the router's routing + callbacks (``_routing_callbacks``) checks the picked deployment before the call + (usage-based-routing-v2 counts its rpm here). A check's error fails the attempt; the core + records that failure like any other.""" + for callback in tuple(callbacks): + if isinstance(callback, CustomLogger): + await callback.async_pre_call_check(dict(deployment), None) # pyright: ignore[reportUnknownMemberType] # CustomLogger takes an untyped deployment dict + + +def routing_pre_call_checks_sync(callbacks: Iterable[object], deployment: Mapping[str, object]) -> None: + """``routing_strategy_pre_call_checks``: the sync check of each ``CustomLogger`` callback.""" + for callback in tuple(callbacks): + if isinstance(callback, CustomLogger): + callback.pre_call_check(dict(deployment)) # pyright: ignore[reportUnknownMemberType] # see above + + +def picked_id(picked: object) -> str | None: + """The id of the deployment a selector returned.""" + info: Final = as_str_mapping((as_str_mapping(picked) or {}).get("model_info")) or {} + found: Final = info.get("id") + return str(found) if found is not None else None + + +__all__ = [ + "SELECTOR_ATTRIBUTES", + "Selector", + "StrategySelectors", + "build_selector", + "group_documents", + "group_settings", + "member_map", + "normalize_strategy", + "parse_groups", + "pick_async", + "pick_sync", + "picked_id", + "routing_pre_call_checks", + "routing_pre_call_checks_sync", +] + + +def parse_groups( + groups_input: object, model_names: frozenset[str], aliases: frozenset[str] +) -> tuple[RoutingGroup, ...]: + """``_init_routing_groups``' validation: ``parse_routing_groups``, then a priority group may not + share its name with a model group or alias (a plain group only warns: its name is shadowed).""" + from litellm._logging import verbose_router_logger + from litellm.litellm_core_utils.routing_groups import ( + parse_routing_groups, # pyright: ignore[reportUnknownVariableType] # takes untyped dicts + ) + + if not groups_input or not isinstance(groups_input, list | tuple): + return () + groups: Final = parse_routing_groups( # pyright: ignore[reportUnknownVariableType] # returns RoutingGroup models + list(groups_input), # pyright: ignore[reportUnknownArgumentType] # validated by parse_routing_groups + known_model_names=model_names, + ) + for group in groups: + if group.group_name in model_names or group.group_name in aliases: + if group.routing_strategy == "priority": + raise ValueError("Priority routing group names must not shadow a model or alias") + verbose_router_logger.warning( + "routing_groups: group_name '%s' is shadowed by an existing model_name or model_group_alias; " + "the group's strategy still applies to its members, but the name is not callable until renamed.", + group.group_name, + ) + return groups + + +def group_documents(groups: Sequence[RoutingGroup]) -> list[Mapping[str, object]]: # mutable-ok: a JSON list + """The routing groups as the Rust core reads them.""" + return [ + { + "group_name": group.group_name, + "models": list(group.models), + "routing_strategy": group.routing_strategy, + **({"model_priorities": dict(group.model_priorities)} if group.model_priorities is not None else {}), + } + for group in groups + ] + + +def group_settings(groups: Sequence[RoutingGroup]) -> list[object]: # mutable-ok: the settings snapshot's list + """``get_settings()['routing_groups']``.""" + return [ + group.model_dump(exclude={"model_priorities"}) if group.model_priorities is None else group.model_dump() + for group in groups + ] + + +def member_map(groups: Sequence[RoutingGroup]) -> Mapping[str, str]: + """``_model_to_group``: each member model group of a non-priority routing group, to the group + (a model listed by two groups maps to the last).""" + kept: Final = [group for group in groups if group.routing_strategy != "priority"] + return dict(chain.from_iterable(dict.fromkeys(group.models, group.group_name).items() for group in kept)) diff --git a/litellm/router_rust/surface/README.md b/litellm/router_rust/surface/README.md new file mode 100644 index 00000000000..ad9c0c62ed3 --- /dev/null +++ b/litellm/router_rust/surface/README.md @@ -0,0 +1,14 @@ +# surface + +The Python `Router` API, kept over router-core's state so Python callers see today's behavior. router-core holds the state and the rules; these modules are views and adapters. + +- `registry`, `model_list`, `indexes`, `members`: the deployment list, indexes and lookups; writes become changesets +- `legacy_members`: the old public members outside the entry points +- `model_info`: model and group info, answering the `model_info` step's lookups +- `replay`: re-registers pricing after a price data reload +- `settings`: `update_settings` rules and the settings view +- `health`, `health_routing`, `retries`: outcome callbacks, health-check routing and retry answers +- `fallbacks`: fallback config validation +- `alerting`: Slack alerting + +A Rust host needs none of this. diff --git a/litellm/router_rust/surface/__init__.py b/litellm/router_rust/surface/__init__.py new file mode 100644 index 00000000000..74704783ecb --- /dev/null +++ b/litellm/router_rust/surface/__init__.py @@ -0,0 +1 @@ +"""The Python ``Router`` API kept over router-core's state, for Python callers.""" diff --git a/litellm/router_rust/surface/alerting.py b/litellm/router_rust/surface/alerting.py new file mode 100644 index 00000000000..760eecb35f4 --- /dev/null +++ b/litellm/router_rust/surface/alerting.py @@ -0,0 +1,33 @@ +"""The Rust core's Slack alerting: ``alerting_config`` built into a ``SlackAlerting`` logger on the +host, as the Python core's ``_initialize_alerting`` builds it.""" + +from typing import TYPE_CHECKING, Final + +from litellm._logging import verbose_router_logger + +if TYPE_CHECKING: + from litellm.integrations.SlackAlerting.slack_alerting import SlackAlerting + from litellm.router_python_host import PythonHost + from litellm.types.router import AlertingConfig + + +def initialize_alerting(host: "PythonHost", alerting_config: "AlertingConfig | None") -> "SlackAlerting | None": + """``_initialize_alerting``: a ``SlackAlerting`` logger with the config's threshold and webhook, + added to the host's callbacks, and its slow-response check to the host's success callbacks. + No config builds nothing.""" + from litellm.integrations.SlackAlerting.slack_alerting import SlackAlerting + + if alerting_config is None: + return None + logger: Final = SlackAlerting( + alerting_threshold=alerting_config.alerting_threshold, + alerting=["slack"], + default_webhook_url=alerting_config.webhook_url, + ) + host.callbacks.add_litellm_callback(logger) + host.callbacks.add_litellm_success_callback(logger.response_taking_too_long_callback) # pyright: ignore[reportUnknownArgumentType, reportUnknownMemberType] # SlackAlerting's untyped callback + verbose_router_logger.info("\033[94m\nInitialized Alerting for litellm.Router\033[0m\n") + return logger + + +__all__ = ["initialize_alerting"] diff --git a/litellm/router_rust/surface/fallbacks.py b/litellm/router_rust/surface/fallbacks.py new file mode 100644 index 00000000000..48fbc972b43 --- /dev/null +++ b/litellm/router_rust/surface/fallbacks.py @@ -0,0 +1,129 @@ +""" +Fallback configuration checks ``RustCore`` runs on the host: ``validate_fallbacks``, with the +Python router's error texts. router-core walks the fallbacks; this module only refuses a list it +could not walk the way the Python router would, before the core is built. +""" + +from collections.abc import Mapping, Sequence +from types import MappingProxyType +from typing import Final + +from litellm.router_utils.fallback_event_handlers import check_non_standard_fallback_format + + +def refuse_malformed_fallbacks(fallback_param: Sequence[object] | None) -> None: + """``Router.validate_fallbacks``: every entry is a one-key dict whose targets are a list (or a + client-side dict target).""" + if fallback_param is None: + return + for fallback_dict in fallback_param: + if not isinstance(fallback_dict, dict): + raise ValueError(f"Item '{fallback_dict}' is not a dictionary.") # noqa: TRY004 # the Python router's error class + if len(fallback_dict) != 1: # pyright: ignore[reportUnknownArgumentType] # only the size is read + raise ValueError( + f"Dictionary '{fallback_dict}' must have exactly one key, but has {len(fallback_dict)} keys." # pyright: ignore[reportUnknownArgumentType] # see above + ) + ((model_group, targets),) = fallback_dict.items() # pyright: ignore[reportUnknownVariableType, reportUnknownMemberType] # an entry as configured + entry: list[object] = [fallback_dict] # mutable-ok: the one-entry list the check reads + if isinstance(targets, list) or check_non_standard_fallback_format(fallbacks=entry): + continue + raise TypeError( + f"Fallbacks for '{model_group}' must be a list of model names, got '{targets}'. " + f"Use {{'{model_group}': ['{targets}']}}" + ) + + +def rejected_fallback_target(targets: Sequence[object]) -> Exception | None: + """The error ``run_async_fallback`` raises before its walk when a chain, from any list and of + any level, has a dict target a request may not set (a server-owned federation field, an + ``oidc/`` api key): every target of that walk is refused.""" + from litellm.litellm_core_utils.credential_accessor import reject_api_key_secret_reference + from litellm.types.router import reject_server_owned_wif_params + + for target in targets: + if isinstance(target, dict): + try: + reject_server_owned_wif_params(target) # pyright: ignore[reportUnknownArgumentType] # a chain as configured + reject_api_key_secret_reference(target) # pyright: ignore[reportUnknownArgumentType] # see above + except Exception as error: # noqa: BLE001 # the walk's own refusal, reported as its failure + return error + return None + + +# The `mock_testing_*` flags the router reads, and the mock failure each asks router-core for. +MOCK_FLAGS: Final = MappingProxyType( + { + "mock_testing_fallbacks": "fallbacks", + "mock_testing_context_fallbacks": "context_window", + "mock_testing_content_policy_fallbacks": "content_policy", + "mock_testing_rate_limit_error": "rate_limit", + } +) + + +def mock_failure(request: Mapping[str, object]) -> str | None: + """The mock failure a request asks for (``MockRouterTestingParams.from_kwargs``, then + ``_handle_mock_testing_rate_limit_error``): the first fallback flag that is ``True`` (a string + is read as a bool), else the rate limit flag when it is ``True`` itself.""" + from litellm.secret_managers.main import str_to_bool + + for flag in ("mock_testing_fallbacks", "mock_testing_context_fallbacks", "mock_testing_content_policy_fallbacks"): + value = request.get(flag) + if (str_to_bool(value) if isinstance(value, str) else value) is True: + return MOCK_FLAGS[flag] + return "rate_limit" if request.get("mock_testing_rate_limit_error") is True else None + + +def mock_error( + mock: str, model: str, num_retries: object, request: Mapping[str, object], configured: Mapping[str, object] +) -> Exception: + """The exception the Python router raises for a mock failure; ``configured`` holds the + router's fallback lists, which the request's own replace.""" + import litellm.exceptions as litellm_exceptions + + def listed(name: str) -> object: + return request.get(name, configured.get(name)) + + if mock == "rate_limit": + return litellm_exceptions.RateLimitError( + model=model, + llm_provider="", + message=f"This is a mock exception for model={model}, to trigger a rate limit error.", + num_retries=num_retries if isinstance(num_retries, int) else None, + ) + if mock == "context_window": + return litellm_exceptions.ContextWindowExceededError( + model=model, + llm_provider="", + message=f"This is a mock exception for model={model}, to trigger a fallback. \ + Context_Window_Fallbacks={listed('context_window_fallbacks')}", + ) + if mock == "content_policy": + return litellm_exceptions.ContentPolicyViolationError( + model=model, + llm_provider="", + message=f"This is a mock exception for model={model}, to trigger a fallback. \ + Context_Policy_Fallbacks={listed('content_policy_fallbacks')}", + ) + return litellm_exceptions.InternalServerError( + model=model, + llm_provider="", + message=f"This is a mock exception for model={model}, to trigger a fallback. Fallbacks={listed('fallbacks')}", + ) + + +class RustFallbacks: + """Fallback members served on the host.""" + + def validate_fallbacks(self, fallback_param: Sequence[object] | None) -> None: + refuse_malformed_fallbacks(fallback_param) + + +__all__ = [ + "MOCK_FLAGS", + "RustFallbacks", + "mock_error", + "mock_failure", + "refuse_malformed_fallbacks", + "rejected_fallback_target", +] diff --git a/litellm/router_rust/surface/health.py b/litellm/router_rust/surface/health.py new file mode 100644 index 00000000000..5547471855a --- /dev/null +++ b/litellm/router_rust/surface/health.py @@ -0,0 +1,558 @@ +""" +The deployment outcome callbacks ``RustCore`` serves as members: ``deployment_callback_on_failure``, +``deployment_callback_on_success`` and ``sync_deployment_callback_on_success``. router-core counts +the outcome and decides the cooldown; this module reads litellm's logging kwargs, drops the +failures the Python router does not hold against the deployment, and drives the Rust call. The +stream wrappers report a stream's outcome through the same members when the stream ends. +""" + +from collections.abc import Mapping +from datetime import datetime +from typing import TYPE_CHECKING, Final, Literal + +from litellm.litellm_core_utils.core_helpers import as_str_mapping + +if TYPE_CHECKING: + from litellm.router_rust.steps.callbacks import OptionalChecks + from litellm.router_rust.surface.registry import HostRegistry + from litellm.router_utils.cooldown_cache import CooldownCache + from litellm.rust_bridge._native import ( + RouterCall, + RouterCoreHandle, + _FailureReportSteps, # pyright: ignore[reportPrivateUsage] # the stub's step shapes + _SuccessReportSteps, # pyright: ignore[reportPrivateUsage] # see above + ) + from litellm.types.router import Deployment + + +def _cast_status(exception_status: object) -> int | None: + if isinstance(exception_status, int) and not isinstance(exception_status, bool): + return exception_status + if isinstance(exception_status, str): + try: + return int(exception_status) + except ValueError: + return 500 + return None + + +def failure_status(exception_status: object) -> int | None: + """``cast_exception_status_to_int``: a number as given, a numeric string read as one, any + other string 500. A number no status can hold reads as no status.""" + status: Final = _cast_status(exception_status) + return status if status is not None and 0 <= status <= 65535 else None + + +def _logged_deployment_id(kwargs: Mapping[str, object]) -> object: + """``litellm_params.model_info.id`` of a logged call.""" + litellm_params: Final = as_str_mapping(kwargs.get("litellm_params")) or {} + return (as_str_mapping(litellm_params.get("model_info")) or {}).get("id") + + +def first_cooldown(*sources: Mapping[str, object] | None) -> int | float | None: + """``_first_present(*sources, key="cooldown_time")``, the number as it was written.""" + for source in sources: + value = source.get("cooldown_time") if source is not None else None + if isinstance(value, int | float) and not isinstance(value, bool): + return value + return None + + +def _configured_cooldown(model_info: object, litellm_params: Mapping[str, object]) -> float | None: + """``first_cooldown`` of a logged call, as the float router-core takes.""" + seconds: Final = first_cooldown(as_str_mapping(model_info), litellm_params) + return None if seconds is None else float(seconds) + + +def python_cooldown_time(seconds: float, *sources: float | None) -> float: + """``seconds`` as ``deployment_callback_on_failure`` hands them on. Router-core's duration is + always a float; Python's is the number that won its precedence, an int or a float as it was + written. ``sources`` run in that precedence, and only the type is read from them.""" + return next((source for source in sources if source is not None and source == seconds), seconds) + + +def _fallback_cooldown(deployment: "Deployment | None") -> float | None: + """``first_cooldown`` of the deployment a failed hop tried, as ``_core_get_model_info`` lists it.""" + if deployment is None: + return None + listed: Final = deployment.model_dump() + return first_cooldown(as_str_mapping(listed.get("model_info")), as_str_mapping(listed.get("litellm_params"))) + + +def _deployment_cooldown(deployment: Mapping[str, object]) -> float | None: + """``_configured_cooldown`` of a deployment as listed.""" + return _configured_cooldown(deployment.get("model_info"), as_str_mapping(deployment.get("litellm_params")) or {}) + + +def logged_call(op: str, kwargs: Mapping[str, object], started: datetime | None = None) -> Mapping[str, object]: + """The logging kwargs ``deployment_callback_on_failure`` would read for a call made with + ``kwargs``: the call's logging object's (a proxy request carries one), else what the request says. + Read before the failure logger stamps the call's ``end_time``, so it is stamped here.""" + ended: Final = datetime.now() # noqa: DTZ005 # naive to match the logging pipeline's api_call_start_time + details: Final = as_str_mapping(getattr(kwargs.get("litellm_logging_obj"), "model_call_details", None)) + call_type: Final = op.rsplit(".", 1)[-1] + if details is not None: + return {"call_type": call_type, **details, "end_time": ended} + fields: Final = ("metadata", "litellm_metadata", "client_side_timeout", "timeout") + return { + "call_type": call_type, + "litellm_params": {name: kwargs.get(name) for name in fields}, + "start_time": started, + "end_time": ended, + } + + +def _authenticates_per_caller(litellm_params: object) -> bool: + """``github_copilot_per_user_credential_name`` in ``_should_run_cooldown_logic``. Per-user params + with no credential raise there, which ends Python's failure callback short of any cooldown.""" + from litellm.exceptions import BadRequestError + from litellm.llms.github_copilot.per_user_auth import github_copilot_per_user_credential_name + + try: + return github_copilot_per_user_credential_name(litellm_params) is not None + except BadRequestError: + return True + + +def failure_already_logged(kwargs: Mapping[str, object], event: Literal["async_failure", "sync_failure"]) -> bool: + """``should_run_logging``: a logging object logs one failure a request. The attempts of a request + that shares one (``litellm_logging_obj``) after the first logged failure run no failure callback; + without one, every litellm call logs on its own.""" + from litellm.litellm_core_utils.litellm_logging import Logging + + logging_obj: Final = kwargs.get("litellm_logging_obj") + return isinstance(logging_obj, Logging) and logging_obj.model_call_details.get(f"has_logged_{event}") is True + + +def _success_target(kwargs: Mapping[str, object]) -> tuple[str, Mapping[str, object]] | None: + """The deployment id and metadata a success callback counts under, or None.""" + params: Final = as_str_mapping(kwargs.get("litellm_params")) or {} + metadata: Final = as_str_mapping(params.get("metadata")) + if metadata is None: + return None + model_info: Final = as_str_mapping(params.get("model_info")) or {} + deployment_id: Final = model_info.get("id") + if metadata.get("model_group") is None or deployment_id is None: + return None + return str(deployment_id), metadata + + +class RustHealthCallbacks: + """Outcome callbacks read into the Rust core's health state.""" + + _handle: "RouterCoreHandle" + _cooldown_given: object + + def _drive_sync_steps(self, call: "RouterCall[_FailureReportSteps | _SuccessReportSteps]") -> Mapping[str, object]: + raise NotImplementedError + + async def _drive_cache_steps( + self, call: "RouterCall[_FailureReportSteps | _SuccessReportSteps]" + ) -> Mapping[str, object]: + raise NotImplementedError + + def _sync_allowed_fails(self) -> None: + raise NotImplementedError + + def _core_get_allowed_fails_from_policy(self, exception: Exception) -> int | None: + """The router ``allowed_fails_policy``'s allowance for this exception, decided by router-core.""" + from litellm.router_rust.core import attempt_failure + + return self._handle.allowed_fails_for(dict(attempt_failure(exception))) + + def _deployment_of(self, deployment_id: object) -> "Deployment | None": + # RustCore sets `_hosted` as an attribute in its constructor. + hosted: Final[HostRegistry] = self.__dict__["_hosted"] # pyright: ignore[reportAny] # RustCore's own attribute + return hosted._core_get_deployment(model_id=deployment_id) if isinstance(deployment_id, str) else None # pyright: ignore[reportPrivateUsage] # the host registry's id lookup + + def deployment_tracks_usage(self, deployment_id: object, model_group: object) -> bool: + """``_cached_deployment_tracks_usage``. A deployment whose model is not mapped raises there, + which ends Python's usage write, so it reads as not tracked.""" + # RustCore sets `_hosted` as an attribute in its constructor. + hosted: Final[HostRegistry] = self.__dict__["_hosted"] # pyright: ignore[reportAny] # RustCore's own attribute + if not isinstance(deployment_id, str) or not isinstance(model_group, str): + return False + try: + return hosted.cached_deployment_tracks_usage(deployment_id, model_group) + except Exception: # noqa: BLE001 # as `deployment_callback_on_success` swallows it + return False + + def cooldown_event(self, deployment_id: str, exception_status: object, cooldown_time: float | None) -> None: + """The end of ``set_cooldown_deployments``, once the cooldown is written: the cooldown event + Prometheus reads, as a task of the running loop. Outside one this raises ``RuntimeError``.""" + import asyncio + + from litellm.router_api import RoutingContext + from litellm.router_utils import cooldown_handlers + + _ = asyncio.get_running_loop() + _ = asyncio.create_task( # fire and forget, as set_cooldown_deployments schedules it + cooldown_handlers.router_cooldown_event_callback( + ctx=RoutingContext(self), # pyright: ignore[reportArgumentType] # the core's RoutingContext + deployment_id=deployment_id, + exception_status=exception_status, # pyright: ignore[reportArgumentType] # the status as raised + cooldown_time=cooldown_time, + ) + ) + + def cooldown_event_on_own_loop( + self, deployment_id: str, exception_status: object, cooldown_time: float | None + ) -> None: + """``cooldown_event`` for a sync walk ``function_with_fallbacks`` would run on its private + loop: ``set_cooldown_deployments`` finds that loop however the caller runs, and the event's + task keeps what it raises to itself.""" + import asyncio + + from litellm.litellm_core_utils import asyncify + + async def fire() -> None: + self.cooldown_event(deployment_id, exception_status, cooldown_time) + await asyncio.sleep(0) + + asyncify.run_async_function(fire) # pyright: ignore[reportUnknownMemberType] # the untyped sync bridge + + def python_cooldown(self, seconds: float, configured: float | None = None, error: object = None) -> float: + """The ``time_to_cooldown`` Python hands ``set_cooldown_deployments``, for router-core's + ``seconds``. A failure callback's precedence is the ``configured`` number, then the ``retry-after`` + header of ``error``, then the router's own; a path passing neither hands on the router's own.""" + from litellm.router_rust.core import retry_after_seconds + + retry_after: Final = retry_after_seconds(error) if isinstance(error, BaseException) else None + default: Final = self._cooldown_given + return python_cooldown_time( + seconds, + configured, + None if retry_after is None else int(retry_after), + default if isinstance(default, int | float) else None, + ) + + def cooled_down( + self, + done: Mapping[str, object], + deployment_id: object, + exception_status: object, + *, + own_loop: bool = False, + configured: float | None = None, + error: object = None, + ) -> bool: + """Whether a failure report cooled the deployment down, with its cooldown event, which + carries ``python_cooldown`` of ``configured`` and ``error``.""" + cooled: Final = bool(done.get("cooled")) + if cooled and isinstance(deployment_id, str): + seconds: Final = done.get("cooldown_time") + event: Final = self.cooldown_event_on_own_loop if own_loop else self.cooldown_event + event( + deployment_id, + exception_status, + self.python_cooldown(seconds, configured, error) if isinstance(seconds, float) else None, + ) + return cooled + + def per_user_deployment(self, deployment_id: object) -> bool: + """The failed deployment authenticates per caller (``_should_run_cooldown_logic``), which is + all ``set_cooldown_deployments`` reads when it is called without the request's kwargs.""" + deployment: Final = self._deployment_of(deployment_id) + return deployment is not None and _authenticates_per_caller(deployment.litellm_params) + + def failure_exits( + self, exception: object, logged: Mapping[str, object], deployment_id: object + ) -> Mapping[str, bool]: + """What the Python router does not hold against the deployment, for every failure it logs. + ``outside_deployment``: ``deployment_callback_on_failure``'s early exits, counting nothing. + ``caller_scoped``: ``_should_run_cooldown_logic``'s per-caller exits, counted but not cooled; + its exits on the exception's class are router-core's (``FailureKind``).""" + from litellm.router_utils.cooldown_handlers import ( + _is_per_user_provider_request, # pyright: ignore[reportPrivateUsage] # the cooldown check's own rule + is_background_response_poll_not_found, + is_caller_scoped_auth_failure, + is_caller_timeout_408, + is_failure_outside_selected_deployment, + is_video_lookup_not_found, + ) + + status: Final[object] = getattr(exception, "status_code", "") + litellm_params: Final = as_str_mapping(logged.get("litellm_params")) or {} + outside: Final = ( + is_failure_outside_selected_deployment(exception if isinstance(exception, BaseException) else None) + or (isinstance(exception, Exception) and is_background_response_poll_not_found(exception, litellm_params)) + or is_video_lookup_not_found(exception, logged) + or is_caller_timeout_408(logged, status) # pyright: ignore[reportArgumentType] # the status as raised + or is_caller_scoped_auth_failure(self._deployment_of(deployment_id), status) # pyright: ignore[reportArgumentType] # see above + ) + return { + "outside_deployment": outside, + "caller_scoped": _is_per_user_provider_request(logged) or self.per_user_deployment(deployment_id), + } + + def trigger_fallback_cooldown( + self, kwargs: Mapping[str, object], exception: object, generic: bool, *, own_loop: bool = False + ) -> None: + """The failed fallback hop's ``_trigger_cooldown_for_failed_deployment``, which runs only when + the request's logging object blocked the hop's failure callbacks. ``generic``: the call's kwargs + carry ``original_generic_function``. ``own_loop``: see ``cooldown_event_on_own_loop``.""" + from litellm._logging import verbose_router_logger + from litellm.router_utils.fallback_event_handlers import ( + _failure_callbacks_blocked_logging_object, # pyright: ignore[reportPrivateUsage] # the fallback path's own rule + ) + + blocked: Final = _failure_callbacks_blocked_logging_object(kwargs) + if blocked is None or not isinstance(exception, Exception): + return + try: + call: Final = self._fallback_failure_call(kwargs, exception, blocked.model_call_details, generic) + if call is not None: + _ = self.cooled_down( + self._drive_sync_steps(call), + getattr(exception, "failed_deployment_id", None), + getattr(exception, "status_code", ""), + own_loop=own_loop, + configured=_fallback_cooldown( + self._deployment_of(getattr(exception, "failed_deployment_id", None)) + ), + error=exception, + ) + except Exception as e: # noqa: BLE001 # best effort, as the Python helper: it never breaks the fallback response + verbose_router_logger.debug("Error triggering cooldown for fallback deployment: %s", e) + + def _fallback_failure_call( + self, + kwargs: Mapping[str, object], + exception: Exception, + model_call_details: Mapping[str, object], + generic: bool, + ) -> "RouterCall[_FailureReportSteps] | None": + """The Rust call recording a failed hop, or None past the exits of + ``_trigger_cooldown_for_failed_deployment``: another list than ``failure_exits``, checked per + failed hop. The selection tag is not on it; a stream stall and a generic call's 404 are.""" + from litellm.router_rust.core import attempt_failure + from litellm.router_utils.cooldown_handlers import ( + _is_per_user_provider_request, # pyright: ignore[reportPrivateUsage] # the cooldown check's own rule + cast_exception_status_to_int, + is_advisor_orchestration_failure, + is_caller_scoped_auth_failure, + is_caller_timeout_408, + is_stream_stall, + ) + from litellm.router_utils.fallback_event_handlers import ( + _REQUEST_SCOPED_STATUS_CODES, # pyright: ignore[reportPrivateUsage] # the fallback path's own rule + ) + + self._sync_allowed_fails() + exception_status: Final[object] = getattr(exception, "status_code", "") + if ( + is_advisor_orchestration_failure(exception) + or is_stream_stall(exception) + or (generic and cast_exception_status_to_int(exception_status) in _REQUEST_SCOPED_STATUS_CODES) # pyright: ignore[reportArgumentType] # the status as raised + or is_caller_timeout_408(model_call_details, exception_status, ended=datetime.now()) # pyright: ignore[reportArgumentType] # noqa: DTZ005 # see logged_call + ): + return None + deployment_id: Final[object] = getattr(exception, "failed_deployment_id", None) + if not isinstance(deployment_id, str): + return None + deployment: Final = self._deployment_of(deployment_id) + if is_caller_scoped_auth_failure(deployment, exception_status): # pyright: ignore[reportArgumentType] # see above + return None + return self._handle.report_failure( + deployment_id, + { + **attempt_failure(exception), + "status": failure_status(exception_status), + "caller_scoped": _is_per_user_provider_request(kwargs) or self.per_user_deployment(deployment_id), + }, + deployment_cooldown=True, + count_only=exception_status == "", + cooldown_time=None if deployment is None else _deployment_cooldown(deployment.model_dump()), + ) + + def deployment_callback_on_failure( + self, + kwargs: Mapping[str, object], + completion_response: object, + start_time: object, + end_time: object, + ) -> bool: + """Count the failure and cool the deployment down when the rules say so.""" + call: Final = self._failure_call(kwargs) + litellm_params: Final = as_str_mapping(kwargs.get("litellm_params")) or {} + return call is not None and self.cooled_down( + self._drive_sync_steps(call), + _logged_deployment_id(kwargs), + getattr(kwargs.get("exception"), "status_code", ""), + configured=first_cooldown(as_str_mapping(litellm_params.get("model_info")), litellm_params), + error=kwargs.get("exception"), + ) + + async def report_stream_failure( + self, kwargs: Mapping[str, object], callbacks: Mapping[str, bool] | None = None + ) -> None: + """A stream's failure after its attempt returned it, as an async call's failure callbacks + record it: counted, cooled down by the rules, and one more request this minute + (``async_deployment_callback_on_failure``). No cooldown event: the stream logs its failure + through ``dispatch_failure_handlers``, whose sync callbacks run in an executor thread, where + ``set_cooldown_deployments`` has no loop to schedule one on.""" + call: Final = self._failure_call(kwargs, counts_request=True, callbacks=callbacks) + if call is not None: + _ = await self._drive_cache_steps(call) + + def report_stream_failure_sync( + self, kwargs: Mapping[str, object], callbacks: Mapping[str, bool] | None = None + ) -> None: + """``report_stream_failure`` for a sync stream: a sync call's failure callback counts no request. + No cooldown event: the stream logs its failure in a ``threading.Thread`` of its own, which + has no running loop whatever the consumer's thread has.""" + call: Final = self._failure_call(kwargs, callbacks=callbacks) + if call is not None: + _ = self._drive_sync_steps(call) + + def _failure_call( + self, + kwargs: Mapping[str, object], + counts_request: bool = False, + callbacks: Mapping[str, bool] | None = None, + ) -> "RouterCall[_FailureReportSteps] | None": + """The Rust call recording a logged failure, or None when it records nothing: one the Python + router does not hold against the deployment is still a request of the minute + (``async_deployment_callback_on_failure`` makes none of the failure callback's exits).""" + from litellm.litellm_core_utils import core_helpers + from litellm.router_rust.core import attempt_failure + + self._sync_allowed_fails() + exception: Final = kwargs.get("exception") + litellm_params: Final = as_str_mapping(kwargs.get("litellm_params")) or {} + model_info: Final = litellm_params.get("model_info", {}) + info: Final = as_str_mapping(model_info) + deployment_id: Final = info.get("id") if info is not None else None + if deployment_id is None: + return None + exits: Final = self.failure_exits(exception, kwargs, deployment_id) + if exits["outside_deployment"] and not counts_request: + return None + exception_status: Final[object] = getattr(exception, "status_code", "") + described: Final[Mapping[str, object]] = ( + attempt_failure(exception) if isinstance(exception, BaseException) else {"ok": False, "kind": "APIError"} + ) + failure: Final = {**described, **exits, "status": failure_status(exception_status)} + metadata: Final[object] = core_helpers.get_litellm_metadata_from_kwargs(dict(kwargs)) # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] # an untyped helper's dict + group: Final = (as_str_mapping(metadata) or {}).get("model_group") # pyright: ignore[reportUnknownArgumentType] # read as a mapping + usage_model: Final = (as_str_mapping(metadata) or {}).get("deployment") # pyright: ignore[reportUnknownArgumentType] # see above + return self._handle.report_failure( + str(deployment_id), + failure, + deployment_cooldown=True, + count_only=exception_status == "", + requested_group=group if isinstance(group, str) else None, + cooldown_time=_configured_cooldown(model_info, litellm_params), + counts_request=counts_request, + usage_model=usage_model if isinstance(usage_model, str) else None, + callbacks=callbacks, + ) + + def sync_deployment_callback_on_success( + self, + kwargs: Mapping[str, object], + completion_response: object, + start_time: object, + end_time: object, + ) -> str | None: + """Count a success this minute; returns the counter's key.""" + return self.report_stream_success_sync(kwargs) + + def report_stream_success_sync( + self, kwargs: Mapping[str, object], callbacks: Mapping[str, bool] | None = None + ) -> str | None: + """``sync_deployment_callback_on_success``, for a sync stream that ended.""" + from litellm.litellm_core_utils.batch_utils import is_batch_retrieve_call_type + + if is_batch_retrieve_call_type(kwargs.get("call_type")): + return None + target: Final = _success_target(kwargs) + if target is None: + return None + key: Final = self._drive_sync_steps(self._handle.count_success(target[0], callbacks)).get("key") + return key if isinstance(key, str) else None + + async def repin_after_success(self, kwargs: Mapping[str, object], completion_response: object) -> None: + """The session re-pin ``deployment_callback_on_success`` runs for a success logged against a + deployment the router serves (``_repin_session_after_success``). The Rust core runs it from the + attempt itself, or from the stream wrapper when a stream ends.""" + from litellm.litellm_core_utils.batch_utils import is_batch_retrieve_call_type + + call_type: Final = kwargs.get("call_type") + if call_type == "_arealtime" or is_batch_retrieve_call_type(call_type): + return + target: Final = _success_target(kwargs) + # RustCore sets these as attributes in its constructor. + hosted: Final[HostRegistry] = self.__dict__["_hosted"] # pyright: ignore[reportAny] # RustCore's own attribute + if target is None or hosted._core_get_deployment(model_id=target[0]) is None: # pyright: ignore[reportPrivateUsage] # the host registry's id lookup + return + checks: Final[OptionalChecks] = self.__dict__["_checks"] # pyright: ignore[reportAny] # see above + cooldowns: Final[CooldownCache] = self.__dict__["_cooldowns"] # pyright: ignore[reportAny] # see above + await checks.repin_session_after_success( + as_str_mapping(kwargs.get("litellm_params")) or {}, completion_response, cooldowns + ) + + async def deployment_callback_on_success( + self, + kwargs: Mapping[str, object], + completion_response: object, + start_time: object, + end_time: object, + ) -> None: + """Count a success and the call's tokens and request against a tracked deployment.""" + await self.report_stream_success(kwargs, completion_response) + + async def report_stream_success( + self, + kwargs: Mapping[str, object], + completion_response: object, + callbacks: Mapping[str, bool] | None = None, + ) -> None: + """``deployment_callback_on_success``, for a stream that ended. The re-pin is the host's half + of that callback, so it reads ``callbacks`` here.""" + from litellm.litellm_core_utils.batch_utils import is_batch_retrieve_call_type + from litellm.router_utils.router_callbacks.track_deployment_metrics import get_counted_usage_tokens + + call_type: Final = kwargs.get("call_type") + if call_type == "_arealtime" or is_batch_retrieve_call_type(call_type): + return + logged: Final = as_str_mapping(kwargs.get("standard_logging_object")) + if logged is None: + return + target: Final = _success_target(kwargs) + if target is None: + return + if (callbacks or {}).get("success", True): + await self.repin_after_success(kwargs, completion_response) + tokens: Final[object] = logged.get("total_tokens", 0) + total: Final = float(tokens) if isinstance(tokens, int | float) else 0.0 + counted: Final = get_counted_usage_tokens(as_str_mapping(kwargs.get("litellm_params")) or {}) + usage_model: Final = target[1].get("deployment") + await self._drive_cache_steps( + self._handle.report_success( + target[0], + total if counted is None else max(0.0, total - counted), + 1.0 if counted is None else 0.0, + usage_model if isinstance(usage_model, str) else None, + self.deployment_tracks_usage(target[0], target[1].get("model_group")), + callbacks, + ) + ) + + async def count_realtime_session_success(self, deployment_id: str) -> None: + """``increment_deployment_successes_for_current_minute`` for a realtime session that ended + cleanly: a success alone, as ``deployment_callback_on_success`` skips ``_arealtime``. A + counter that cannot be written does not fail the served session.""" + from litellm._logging import verbose_router_logger + + try: + await self._drive_cache_steps(self._handle.count_success(deployment_id)) + except Exception as e: # noqa: BLE001 # as the Rust call loop treats a success it could not record + verbose_router_logger.debug("count_realtime_session_success: not recorded. Got - %s", e) + + +__all__ = [ + "RustHealthCallbacks", + "failure_status", + "logged_call", +] diff --git a/litellm/router_rust/surface/health_routing.py b/litellm/router_rust/surface/health_routing.py new file mode 100644 index 00000000000..46542bf9af2 --- /dev/null +++ b/litellm/router_rust/surface/health_routing.py @@ -0,0 +1,160 @@ +""" +Health-check routing, allowed-fails policies and the failover flags on the host side of the Rust +core: how the constructor resolves them, the static document router-core reads them from, and the +``health()`` members that read the registry's live model list or the background health checks' +results. router-core makes every routing decision (the health-check filter, the cooldown rules, +weighted failover); this module only reads Python values into its documents and back. + +The rules are the Python router core's, ported (never imported). +""" + +import dataclasses +from collections.abc import Mapping, Sequence +from typing import TYPE_CHECKING, Final + +from litellm.litellm_core_utils.core_helpers import as_str_mapping +from litellm.router_utils.health_state_cache import DeploymentHealthCache +from litellm.types.router import AllowedFailsPolicy + +if TYPE_CHECKING: + from litellm.caching.caching import DualCache + from litellm.router_config import RouterConfig + from litellm.types.router import Deployment, DeploymentTypedDict + +HEALTH_STATE_KEY: Final = DeploymentHealthCache.CACHE_KEY +# The failover flags the Python core keeps as plain router attributes, with their defaults. +FAILOVER_FLAGS: Final = ("enable_weighted_failover", "fallback_on_unsupported_params") + + +def resolved_allowed_fails_policy(value: object) -> AllowedFailsPolicy | None: + """The router's ``allowed_fails_policy`` as the constructor keeps it: a dict is read into an + ``AllowedFailsPolicy`` (a bad field raises the model's own error), a policy is kept, anything + else is no policy.""" + if isinstance(value, AllowedFailsPolicy): + return value + if isinstance(value, dict): + return AllowedFailsPolicy.model_validate(value) + return None + + +def _staleness(threshold: object) -> float: + """``health_check_staleness_threshold or DEFAULT_HEALTH_CHECK_INTERVAL * multiplier``.""" + from litellm.constants import DEFAULT_HEALTH_CHECK_INTERVAL, DEFAULT_HEALTH_CHECK_STALENESS_MULTIPLIER + + given: Final = threshold if isinstance(threshold, int | float) and not isinstance(threshold, bool) else None + return float(given or (DEFAULT_HEALTH_CHECK_INTERVAL * DEFAULT_HEALTH_CHECK_STALENESS_MULTIPLIER)) + + +@dataclasses.dataclass(frozen=True, slots=True) +class HealthRouting: + """The health-check routing settings, as the constructor resolves them.""" + + enable_health_check_routing: bool = False + health_check_staleness: float = 600.0 + health_check_ignore_transient_errors: bool = False + background_health_check_model_groups: frozenset[str] | None = None + + @classmethod + def of(cls, config: "RouterConfig") -> "HealthRouting": + groups: Final = config.get("background_health_check_model_groups") + return cls( + enable_health_check_routing=bool(config.get("enable_health_check_routing")), + health_check_staleness=_staleness(config.get("health_check_staleness_threshold")), + health_check_ignore_transient_errors=bool(config.get("health_check_ignore_transient_errors")), + background_health_check_model_groups=frozenset(groups) if groups is not None else None, # pyright: ignore[reportArgumentType, reportUnknownArgumentType] # any iterable of names, as the Python core takes it + ) + + def static_document(self) -> Mapping[str, object]: + groups: Final = self.background_health_check_model_groups + return { + "enable_health_check_routing": True if self.enable_health_check_routing else None, + "health_check_staleness_threshold": self.health_check_staleness, + # Only a string can name a model group, so anything else in scope matches nothing. + "background_health_check_model_groups": sorted(name for name in groups if isinstance(name, str)) # pyright: ignore[reportUnnecessaryIsInstance] # `of` keeps any iterable the constructor was given + if groups is not None + else None, + } + + def health_checks(self, cache: "DualCache") -> DeploymentHealthCache: + return DeploymentHealthCache(cache=cache, staleness_threshold=self.health_check_staleness) + + +def failover_document(plain: Mapping[str, object]) -> Mapping[str, object]: + """The failover flags router-core reads, from the router's plain attributes (read for + truthiness, as the Python core's checks read them).""" + return { + "enable_weighted_failover": True if plain.get("enable_weighted_failover") else None, + "fallback_on_unsupported_params": False if not plain.get("fallback_on_unsupported_params", True) else None, + } + + +def _model_info(deployment: Mapping[str, object]) -> Mapping[str, object]: + return as_str_mapping(deployment.get("model_info")) or {} + + +def is_blocked(deployment: "Deployment") -> bool: + """``_core_is_deployment_blocked``.""" + model_info: Final[object | None] = getattr(deployment, "model_info", None) + if model_info is None: + return False + return getattr(model_info, "blocked", None) is True + + +def all_blocked(deployments: "Sequence[DeploymentTypedDict]") -> bool: + """``_core_are_all_deployments_blocked``.""" + return len(deployments) > 0 and all( + _model_info(deployment).get("blocked") is True # pyright: ignore[reportArgumentType] # a deployment dict + for deployment in deployments + ) + + +def fully_blocked_model_names( + model_list: Sequence[Mapping[str, object]], +) -> set[str]: # mutable-ok: the core's type, a fresh set + """``_core_get_fully_blocked_model_names``: groups whose every deployment is blocked.""" + blocked_by_name: Final[dict[str, bool]] = {} # mutable-ok: built here + for deployment in model_list: + name = deployment.get("model_name") or "" + if not isinstance(name, str) or not name: + continue + blocked = _model_info(deployment).get("blocked") is True + blocked_by_name[name] = blocked_by_name.get(name, True) and blocked + return {name for name, fully in blocked_by_name.items() if fully} + + +async def fully_unhealthy_model_names( + model_list: Sequence[Mapping[str, object]], health_checks: DeploymentHealthCache, policy: object +) -> set[str]: # mutable-ok: the core's type, a fresh set + """``_core_async_get_fully_unhealthy_model_names``: groups (and team public names) whose + every deployment failed its health check; nothing under an allowed_fails_policy.""" + if policy is not None: + return set() + unhealthy_ids: Final = await health_checks.async_get_unhealthy_deployment_ids() + if not unhealthy_ids: + return set() + unhealthy_by_name: Final[dict[str, bool]] = {} # mutable-ok: built here + for deployment in model_list: + model_info = _model_info(deployment) + names = [deployment.get("model_name") or ""] + team_public_model_name = model_info.get("team_public_model_name") + if team_public_model_name: + names.append(team_public_model_name) + unhealthy = model_info.get("id") in unhealthy_ids + for name in names: + if not isinstance(name, str) or not name: + continue + unhealthy_by_name[name] = unhealthy_by_name.get(name, True) and unhealthy + return {name for name, fully in unhealthy_by_name.items() if fully} + + +__all__ = [ + "FAILOVER_FLAGS", + "HEALTH_STATE_KEY", + "HealthRouting", + "all_blocked", + "failover_document", + "fully_blocked_model_names", + "fully_unhealthy_model_names", + "is_blocked", + "resolved_allowed_fails_policy", +] diff --git a/litellm/router_rust/surface/indexes.py b/litellm/router_rust/surface/indexes.py new file mode 100644 index 00000000000..f5aca3e1bad --- /dev/null +++ b/litellm/router_rust/surface/indexes.py @@ -0,0 +1,514 @@ +# pyright: reportIncompatibleMethodOverride=false +# (each view overrides its builtin's mutators with the same behaviour, as a changeset) +""" +``HostedIndexes``: the registry indexes the Rust core hands out, read back from router-core. + +router-core keeps the indexes the Python router keeps by history (``host_index.rs``): the id, +group and team position maps, the group names a ``set_model_list`` was given, the deployment +names that only ever grow, the team public names, the provider default ids and the wildcard +pattern routers. ``registry().index()``, ``names(kind=...)``, ``patterns()`` and the deprecated +attributes return the views here: a ``dict``, ``set`` or ``list`` (so ``isinstance`` and ``json`` +keep working), one object per router, whose contents are router-core's, read again when an entry +point hands a view out (or before a write to it) once the registry moved, not on every registry +write. A write to any of them, a nested one included (``index("model_name")["gpt"].append(3)``), +is a ``host_index`` changeset that replaces that index in router-core, and the view reads it back. +The lookups over these indexes run in router-core (``lookup.rs``), not over the views. +""" + +import copy +from collections.abc import Callable, Iterable, Iterator, Mapping, Sequence +from functools import partial +from typing import ( + TYPE_CHECKING, + Final, + Literal, + SupportsIndex, + TypeAlias, + cast, # noqa: TID251 # any value a caller wrote into a view +) + +from litellm.constants import DEFAULT_MAX_RECURSE_DEPTH +from litellm.litellm_core_utils.core_helpers import as_str_mapping +from litellm.litellm_core_utils.pattern_match_deployments import PatternMatchRouter + +if TYPE_CHECKING: + from litellm.rust_bridge._native import RouterCoreHandle + +IndexField: TypeAlias = Literal[ + "id_map", + "model_name", + "team_model", + "model_names", + "deployment_names", + "team_public_model_names", + "provider_default_ids", + "patterns", + "team_patterns", +] +_Commit: TypeAlias = Callable[[object], None] +# Deployment copies a pattern router routes to, fresh for the caller. +Routed: TypeAlias = list[dict[str, object]] +_Pairs: TypeAlias = list[tuple[str, list[object]]] + + +def _stay() -> None: + """A view with no contents of its own to read again before a write.""" + + +def _as_set(value: object) -> set[object]: # mutable-ok: set's own in-place operand + return cast(set[object], value) # cast-ok: an in-place set operand, checked by set itself + + +def _as_mapping(value: object) -> Mapping[object, object]: + if not isinstance(value, Mapping): + return {} + return cast(Mapping[object, object], value) # cast-ok: any mapping a caller wrote + + +def _json(value: object, depth: int = 0) -> object: + """A view's contents as router-core takes them: plain JSON, a set as a list. A caller wrote them, + so a cycle is possible: anything nested past ``DEFAULT_MAX_RECURSE_DEPTH`` is refused.""" + if depth > DEFAULT_MAX_RECURSE_DEPTH: + raise ValueError(f"a router index is nested deeper than {DEFAULT_MAX_RECURSE_DEPTH} levels") + mapping: Final = _as_mapping(value) + if mapping or isinstance(value, Mapping): + return {str(key): _json(item, depth + 1) for key, item in mapping.items()} + if isinstance(value, list | tuple | set | frozenset): + items: Final = cast(Iterable[object], value) # cast-ok: any collection a caller wrote + return [_json(item, depth + 1) for item in items] + return value + + +class _Positions(list[object]): + """One index's position list (or a pattern's deployments): a write commits the whole list.""" + + def __init__(self, items: Iterable[object], commit: _Commit, fresh: Callable[[], None] = _stay) -> None: + super().__init__(items) + self._commit: Final = commit + self._fresh: Final = fresh + + def _load(self, items: Iterable[object]) -> None: + """Show router-core's contents, without a write.""" + super().__setitem__(slice(None), list(items)) + + def _write(self, change: Callable[[list[object]], object]) -> object: # mutable-ok: the copy it changes + self._fresh() # a view held across a write elsewhere changes router-core's contents, not stale ones + changed: Final = list(self) + result: Final = change(changed) + self._commit(changed) + return result + + def append(self, value: object) -> None: + _ = self._write(lambda items: items.append(value)) + + def extend(self, values: Iterable[object]) -> None: + added: Final = tuple(values) + _ = self._write(lambda items: items.extend(added)) + + def insert(self, index: SupportsIndex, value: object) -> None: + _ = self._write(lambda items: items.insert(index, value)) + + def remove(self, value: object) -> None: + _ = self._write(lambda items: items.remove(value)) + + def pop(self, index: SupportsIndex = -1) -> object: + return self._write(lambda items: items.pop(index)) + + def clear(self) -> None: + _ = self._write(lambda items: items.clear()) + + def sort(self, *args: object, **kwargs: object) -> None: # kwargs-ok: list.sort's own arguments + _ = self._write(lambda items: items.sort(*args, **kwargs)) # pyright: ignore[reportCallIssue, reportArgumentType] # list.sort's own arguments + + def reverse(self) -> None: + _ = self._write(lambda items: items.reverse()) # pyright: ignore[reportUnknownLambdaType] # list.reverse returns None + + def __setitem__(self, index: object, value: object) -> None: + _ = self._write(lambda items: items.__setitem__(index, value)) # pyright: ignore[reportCallIssue, reportArgumentType, reportUnknownLambdaType] # list's own item and slice assignment + + def __delitem__(self, index: object) -> None: + _ = self._write(lambda items: items.__delitem__(index)) # pyright: ignore[reportArgumentType] # list's own item and slice deletion + + def __iadd__(self, values: Iterable[object]) -> "_Positions": + self.extend(values) + return self + + def __imul__(self, count: SupportsIndex) -> "_Positions": + _ = self._write(lambda items: items.__imul__(count)) + return self + + def __reduce_ex__(self, protocol: SupportsIndex) -> tuple[object, ...]: + return (list, (list(self),)) + + +class _IndexDict(dict[object, object]): + """A key-to-positions index: a write commits the whole index.""" + + def __init__(self, commit: _Commit, fresh: Callable[[], None] = _stay) -> None: + super().__init__() + self._commit: Final = commit + self._fresh: Final = fresh + + def _load(self, items: Mapping[object, object]) -> None: + """Show router-core's contents, without a write.""" + super().clear() + super().update(items) + + def _write(self, change: Callable[[dict[object, object]], object]) -> object: # mutable-ok: the copy it changes + self._fresh() # a view held across a write elsewhere changes router-core's contents, not stale ones + changed: Final = dict(self) + result: Final = change(changed) + self._commit(changed) + return result + + def __setitem__(self, key: object, value: object) -> None: + _ = self._write(lambda items: items.__setitem__(key, value)) + + def __delitem__(self, key: object) -> None: + _ = self._write(lambda items: items.__delitem__(key)) + + def pop(self, key: object, *default: object) -> object: + return self._write(lambda items: items.pop(key, *default)) + + def popitem(self) -> tuple[object, object]: + return cast(tuple[object, object], self._write(lambda items: items.popitem())) # cast-ok: dict.popitem's pair + + def clear(self) -> None: + _ = self._write(lambda items: items.clear()) + + def update(self, *args: object, **kwargs: object) -> None: # kwargs-ok: dict.update's own arguments + _ = self._write(lambda items: items.update(*args, **kwargs)) # pyright: ignore[reportCallIssue, reportArgumentType] # dict.update's own arguments + + def setdefault(self, key: object, default: object = None) -> object: + if key in self: + return self[key] + self[key] = default + return self[key] + + def __ior__(self, other: object) -> "_IndexDict": + self.update(other) + return self + + def __reduce_ex__(self, protocol: SupportsIndex) -> tuple[object, ...]: + return (dict, (dict(self),)) + + +class _NameSet(set[object]): + """``model_names``: a write commits the whole set.""" + + def __init__(self, commit: _Commit, fresh: Callable[[], None] = _stay) -> None: + super().__init__() + self._commit: Final = commit + self._fresh: Final = fresh + + def _load(self, items: Iterable[object]) -> None: + """Show router-core's contents, without a write.""" + super().clear() + super().update(items) + + def _write(self, change: Callable[[set[object]], object]) -> object: # mutable-ok: the copy it changes + self._fresh() # a view held across a write elsewhere changes router-core's contents, not stale ones + changed: Final = set(self) + result: Final = change(changed) + self._commit(changed) + return result + + def add(self, value: object) -> None: + _ = self._write(lambda items: items.add(value)) + + def discard(self, value: object) -> None: + _ = self._write(lambda items: items.discard(value)) + + def remove(self, value: object) -> None: + _ = self._write(lambda items: items.remove(value)) + + def pop(self) -> object: + return self._write(lambda items: items.pop()) + + def clear(self) -> None: + _ = self._write(lambda items: items.clear()) + + def update(self, *others: Iterable[object]) -> None: + _ = self._write(lambda items: items.update(*others)) + + def difference_update(self, *others: Iterable[object]) -> None: + _ = self._write(lambda items: items.difference_update(*others)) + + def intersection_update(self, *others: Iterable[object]) -> None: + _ = self._write(lambda items: items.intersection_update(*others)) + + def symmetric_difference_update(self, other: Iterable[object]) -> None: + _ = self._write(lambda items: items.symmetric_difference_update(other)) + + def __ior__(self, other: object) -> "_NameSet": + _ = self._write(lambda items: items.__ior__(_as_set(other))) + return self + + def __iand__(self, other: object) -> "_NameSet": + _ = self._write(lambda items: items.__iand__(_as_set(other))) + return self + + def __isub__(self, other: object) -> "_NameSet": + _ = self._write(lambda items: items.__isub__(_as_set(other))) + return self + + def __ixor__(self, other: object) -> "_NameSet": + _ = self._write(lambda items: items.__ixor__(_as_set(other))) + return self + + def __reduce_ex__(self, protocol: SupportsIndex) -> tuple[object, ...]: + return (set, (set(self),)) + + +class HostedPatternRouter(PatternMatchRouter): + """A wildcard pattern router whose patterns are router-core's (``pattern_router``, or one team's + router): routing runs in router-core, and ``add_pattern`` / ``remove_deployment`` / a write to + ``patterns`` are changesets.""" + + def __init__(self, owner: "HostedIndexes", team: str | None) -> None: + super().__init__() # its empty `patterns` write is ignored: router-core holds them + self._owner: Final = owner + self._team: Final = team + self._view: Final = _IndexDict(self._commit, owner.refresh) + + @property + def patterns(self) -> dict[str, list[object]]: # pyright: ignore[reportIncompatibleVariableOverride] # the base class's attribute, read from router-core # mutable-ok: the live patterns, writes are changesets + self._owner.refresh() + return cast(dict[str, list[object]], self._view) # cast-ok: regex keys, deployment lists + + @patterns.setter + def patterns(self, value: Mapping[str, object]) -> None: # pyright: ignore[reportIncompatibleVariableOverride] # see the getter + if not hasattr(self, "_owner"): + return # the base constructor's empty dict: router-core holds this router's patterns + self._commit(value) + + def _show(self, pairs: Iterable[tuple[str, Sequence[object]]]) -> None: + shown: Final = { + regex: _Positions(deployments, partial(self._write_pattern, regex)) for regex, deployments in pairs + } + self._view._load(cast(Mapping[object, object], shown)) # pyright: ignore[reportPrivateUsage] # cast-ok: regex to deployments + + def _write_pattern(self, regex: str, deployments: object) -> None: + self._commit({**self.patterns, regex: deployments}) + + def _commit(self, value: object) -> None: + pairs: Final = [ + [str(regex), _json(deployments)] for regex, deployments in (as_str_mapping(value) or {}).items() + ] + self._owner.write_router(self._team, pairs) + + def route(self, request: str | None, filtered_model_names: Sequence[str] | None = None) -> Routed | None: # pyright: ignore[reportIncompatibleMethodOverride] # the base's list types, read only here + if not isinstance(request, str): + # The base class matches its regexes against it, which fails and routes nowhere. + return super().route(request, None if filtered_model_names is None else list(filtered_model_names)) # pyright: ignore[reportReturnType, reportUnknownMemberType, reportUnknownVariableType] # None, or the base's copies of the deployments + return self._owner.route(self._team, request, filtered_model_names) + + def get_models_by_pattern(self, model: str, custom_llm_provider: str | None = None) -> frozenset[str]: + from litellm.litellm_core_utils.get_llm_provider_logic import inferred_provider + + if not isinstance(model, str): # pyright: ignore[reportUnnecessaryIsInstance] # a caller's non-string model goes to the base class + return super().get_models_by_pattern(model, custom_llm_provider) + provider: Final = custom_llm_provider or inferred_provider(model) + return self._owner.models_by_pattern(self._team, [model, f"{provider}/{model}"]) + + +class _TeamRouters(dict[str, HostedPatternRouter]): + """``team_pattern_routers``: team id to its router; removing a team's router is a changeset.""" + + def __init__(self, owner: "HostedIndexes") -> None: + super().__init__() + self._owner: Final = owner + + def _load(self, routers: Mapping[str, HostedPatternRouter]) -> None: + super().clear() + super().update(routers) + + def __delitem__(self, team: str) -> None: + self._owner.drop_team_router(team) + + def pop(self, team: str, *default: object) -> object: + if team not in self: + if default: + return default[0] + raise KeyError(team) + router: Final = self[team] + self._owner.drop_team_router(team) + return router + + def clear(self) -> None: + self._owner.write({"team_patterns": []}) + + def setdefault(self, team: str, default: HostedPatternRouter | None = None) -> HostedPatternRouter: + if team not in self: + self._owner.write_router(team, []) + return self[team] + + +class HostedIndexes: + """Router-core's host indexes, as the live objects the registry hands out.""" + + def __init__( + self, + handle: Callable[[], "RouterCoreHandle | None"], + apply: Callable[[Mapping[str, object]], None], + ) -> None: + self._handle: Final = handle + self._apply: Final = apply + self._read_at: tuple[object, int] | None = None + self._document: Mapping[str, object] = {} + # The views are read again when an entry point hands them out, and before a write to them. + self.id_map: Final = _IndexDict(lambda value: self.write({"id_map": value}), self.refresh) + self.model_name: Final = _IndexDict(lambda value: self.write({"model_name": value}), self.refresh) + self.team_model: Final = _IndexDict(lambda value: self._write_team_model(value), self.refresh) + self.model_names: Final = _NameSet(lambda value: self.write({"model_names": value}), self.refresh) + self.deployment_names: Final = _Positions( + (), lambda value: self.write({"deployment_names": value}), self.refresh + ) + self.provider_default_ids: Final = _Positions( + (), lambda value: self.write({"provider_default_ids": value}), self.refresh + ) + self.pattern_router: Final = HostedPatternRouter(self, None) + self.team_pattern_routers: Final = _TeamRouters(self) + self._team_routers: dict[str, HostedPatternRouter] = {} # mutable-ok: one router object per team + self._team_public: frozenset[str] = frozenset() + + # -- reading ---------------------------------------------------------------------------------- + + def refresh(self) -> None: + """Read router-core's indexes again when its registry moved since the last read.""" + handle: Final = self._handle() + if handle is None: + return + version: Final = handle.registry_version() + if self._read_at is not None and self._read_at[0] is handle and self._read_at[1] == version: + return + self._read_at = (handle, version) + self._document = as_str_mapping(handle.host_indexes()) or {} + self._show() + + def _field(self, name: str) -> object: + return self._document.get(name) + + def _show(self) -> None: + ids: Final = _as_mapping(self._field("id_map")) + self.id_map._load(ids) # pyright: ignore[reportPrivateUsage] # a view this object owns + named: Final[dict[object, object]] = {} # mutable-ok: built here + for name, positions in (as_str_mapping(self._field("model_name")) or {}).items(): + named[name] = self._positions(positions, partial(self._write_named, name)) + self.model_name._load(named) # pyright: ignore[reportPrivateUsage] # a view this object owns + teams: Final[dict[object, object]] = {} # mutable-ok: built here + for team, public, positions in _triples(self._field("team_model")): + teams[(team, public)] = self._positions(positions, partial(self._write_team_key, (team, public))) + self.team_model._load(teams) # pyright: ignore[reportPrivateUsage] # a view this object owns + self.model_names._load(_items(self._field("model_names"))) # pyright: ignore[reportPrivateUsage] # a view this object owns + self.deployment_names._load(_items(self._field("deployment_names"))) # pyright: ignore[reportPrivateUsage] # a view this object owns + self.provider_default_ids._load(_items(self._field("provider_default_ids"))) # pyright: ignore[reportPrivateUsage] # a view this object owns + self._team_public = frozenset(str(name) for name in _items(self._field("team_public_model_names"))) + self.pattern_router._show(_pattern_pairs(self._field("patterns"))) # pyright: ignore[reportPrivateUsage] # the router this object owns + routers: Final[dict[str, HostedPatternRouter]] = {} # mutable-ok: rebuilt each read + for team, pairs in _team_pairs(self._field("team_patterns")): + router = self._team_routers.get(team) or HostedPatternRouter(self, team) + router._show(pairs) # pyright: ignore[reportPrivateUsage] # the router this object owns + routers[team] = router + self._team_routers = routers + self.team_pattern_routers._load(routers) # pyright: ignore[reportPrivateUsage] # a view this object owns + + def _write_named(self, name: str, items: object) -> None: + self.write({"model_name": {**self.model_name, name: items}}) + + def _write_team_key(self, key: tuple[str, str], items: object) -> None: + self._write_team_model({**self.team_model, key: items}) + + def _positions(self, positions: object, commit: _Commit) -> object: + """A position list as a write-through list; anything else a caller wrote, as written.""" + items: Final = _items(positions) + return _Positions(items, commit) if isinstance(positions, list) else positions + + @property + def team_public_model_names(self) -> frozenset[str]: + self.refresh() + return self._team_public + + # -- writing ---------------------------------------------------------------------------------- + + def write(self, fields: Mapping[str, object]) -> None: + """Replace the named indexes in router-core, then read them back.""" + self._apply({"host_index": {name: _json(value) for name, value in fields.items()}}) + self.refresh() + + def assign(self, field: IndexField, value: object) -> None: + """A deprecated attribute assignment of one index (``router.model_names = {...}``).""" + if field == "team_model": + self._write_team_model(value) + else: + self.write({field: value}) + + def _write_team_model(self, value: object) -> None: + mapping: Final = _as_mapping(value) + triples: Final[list[list[object]]] = [] # mutable-ok: JSON triples + for key, positions in mapping.items(): + parts = cast(tuple[object, ...], key) if isinstance(key, tuple) else () # cast-ok: a key a caller wrote + if len(parts) == 2: + triples.append([str(parts[0]), str(parts[1]), _json(positions)]) + self.write({"team_model": triples}) + + def write_router(self, team: str | None, pairs: list[list[object]]) -> None: # mutable-ok: JSON pairs + """Replace one pattern router's patterns.""" + if team is None: + self.write({"patterns": pairs}) + return + teams: Final = [[t, _json(_pairs_of(router))] for t, router in self._team_routers.items() if t != team] + self.write({"team_patterns": [*teams, [team, pairs]]}) + + def drop_team_router(self, team: str) -> None: + if team not in self._team_routers: + raise KeyError(team) + self.write({"team_patterns": [[t, _json(_pairs_of(r))] for t, r in self._team_routers.items() if t != team]}) + + # -- routing ---------------------------------------------------------------------------------- + + def route(self, team: str | None, request: str, filtered: Sequence[str] | None) -> Routed | None: + handle: Final = self._handle() + routed: Final = ( + None + if handle is None + else handle.pattern_route(team, request, None if filtered is None else list(filtered)) + ) + return None if routed is None else [dict(deployment) for deployment in routed] + + def models_by_pattern(self, team: str | None, requests: list[str]) -> frozenset[str]: # mutable-ok: read only + handle: Final = self._handle() + found: Final = None if handle is None else handle.models_by_pattern(team, requests) + return frozenset(found or ()) + + +def _items(value: object) -> list[object]: # mutable-ok: a fresh list + return list(cast(list[object], value)) if isinstance(value, list) else [] # cast-ok: a JSON list + + +def _triples(value: object) -> Iterator[tuple[str, str, object]]: + for item in _items(value): + parts = _items(item) + if len(parts) == 3: + yield str(parts[0]), str(parts[1]), parts[2] + + +def _pattern_pairs(value: object) -> _Pairs: + pairs: Final[_Pairs] = [] + for item in _items(value): + parts = _items(item) + if len(parts) == 2: + pairs.append((str(parts[0]), _items(parts[1]))) + return pairs + + +def _team_pairs(value: object) -> Iterator[tuple[str, _Pairs]]: + for item in _items(value): + parts = _items(item) + if len(parts) == 2: + yield str(parts[0]), _pattern_pairs(parts[1]) + + +def _pairs_of(router: HostedPatternRouter) -> list[list[object]]: # mutable-ok: JSON pairs + view: Final[Mapping[object, object]] = router._view # pyright: ignore[reportPrivateUsage] # a router this module owns + return [[regex, copy.deepcopy(deployments)] for regex, deployments in view.items()] diff --git a/litellm/router_rust/surface/legacy_members.py b/litellm/router_rust/surface/legacy_members.py new file mode 100644 index 00000000000..6c30557eb4d --- /dev/null +++ b/litellm/router_rust/surface/legacy_members.py @@ -0,0 +1,547 @@ +""" +``RustLegacyMembers``: the old public ``Router`` members outside the entry points, on the Rust +core: the routing-state lifecycle, the routing callbacks' checks, one provider call's accounting +and the pre-routing strategy builders. Each one is the Python router core's, ported (never +imported), over the state the Rust core and its host keep. + +Not served, and refused by ``RustCore.__getattr__``: the Python fallback walk +(``async_function_with_fallbacks``, ``function_with_fallbacks``, ``async_function_with_retries``, +``async_function_with_fallbacks_common_utils``) and ``factory_function``. They run a caller's +``original_function``, which picks its own deployment; router-core makes each pick itself and +attempts one host op, named by the endpoint's plan. +""" + +import asyncio +import contextlib +import inspect +import traceback +import uuid +from collections.abc import Awaitable, Callable, Iterable, Mapping, Sequence +from types import MappingProxyType +from typing import TYPE_CHECKING, Final + +from typing_extensions import NotRequired, ReadOnly, TypedDict + +from litellm._internal_context import service_target +from litellm._logging import verbose_router_logger +from litellm.constants import ROUTER_USAGE_COUNTED_TOKENS_METADATA_KEY +from litellm.integrations.custom_logger import CustomLogger +from litellm.litellm_core_utils.add_retry_fallback_headers import response_total_token_count +from litellm.litellm_core_utils.batch_utils import is_batch_retrieve_call_type +from litellm.litellm_core_utils.core_helpers import as_str_mapping, get_parent_otel_span_from_kwargs +from litellm.litellm_core_utils.coroutine_checker import coroutine_checker +from litellm.litellm_core_utils.routing_groups import validate_routing_strategy +from litellm.litellm_core_utils.token_counter import served_openai_models as _served_openai_models +from litellm.router_api import RoutingContext +from litellm.router_rust.invoke.websocket import admit_deployment_turn +from litellm.router_rust.steps.callbacks import RustOptionalChecks +from litellm.router_rust.steps.prerouting import TIER_PARAMS_NEVER_DROPPED as _TIER_PARAMS_NEVER_DROPPED +from litellm.router_rust.surface.members import RustRegistryMembers +from litellm.router_strategy.auto_router.auto_router_model_naming import ( + GatedAutoRouterCapability, + capability_limit_violation, + count_capability_routers, +) +from litellm.router_utils.cooldown_handlers import set_cooldown_deployments +from litellm.router_utils.endpoint_eligibility import endpoint_selection +from litellm.router_utils.pre_call_checks.io_token_rate_limit_check import async_io_token_refund_failure +from litellm.router_utils.router_callbacks.track_deployment_metrics import find_deployment_metadata +from litellm.router_utils.routing_read_batch import ROUTER_USAGE_TARGET +from litellm.types.router import RouterCacheEnum +from litellm.utils import get_utc_datetime + +if TYPE_CHECKING: + from opentelemetry.trace import Span + + from litellm.caching.dual_cache import DualCache + from litellm.litellm_core_utils.litellm_logging import Logging + from litellm.router_rust.invoke.endpoints import EndpointCore + from litellm.types.router import CustomRoutingStrategyBase, Deployment, RoutingGroup + +# `RoutingArgs.ttl`: a usage counter lives for its minute. +_USAGE_TTL: Final = 60 + +# What `initialize_assistants_endpoint` binds. +_ASSISTANTS_ENDPOINTS: Final = ( + "acreate_assistants", + "adelete_assistant", + "aget_assistants", + "acreate_thread", + "aget_thread", + "a_add_message", + "aget_messages", + "arun_thread", +) + +# What `initialize_router_endpoints` binds, with the container file endpoints it reads from +# `litellm/containers/endpoints.json`. +_ROUTER_ENDPOINTS: Final = frozenset( + { + "_arealtime", + "_aresponses_websocket", + "aadapter_generate_content", + "aanthropic_messages", + "acancel_fine_tuning_job", + "acancel_interaction", + "acancel_responses", + "acompact_responses", + "acreate_agent", + "acreate_container", + "acreate_fine_tuning_job", + "acreate_interaction", + "acreate_realtime_client_secret", + "acreate_realtime_transcription_session", + "acreate_realtime_translation_client_secret", + "acreate_skill", + "adecisions", + "adelete_agent", + "adelete_container", + "adelete_interaction", + "adelete_responses", + "adelete_skill", + "aextraction", + "afile_content", + "afile_delete", + "afile_list", + "agenerate_content", + "agenerate_content_stream", + "aget_agent", + "aget_interaction", + "aget_responses", + "aget_skill", + "aimage_edit", + "alist_agent_versions", + "alist_agents", + "alist_containers", + "alist_fine_tuning_jobs", + "alist_input_items", + "alist_skills", + "allm_passthrough_route", + "amoderation", + "anthropic_messages", + "aocr", + "arealtime_calls", + "arealtime_translation_calls", + "aresponses", + "aretrieve_container", + "aretrieve_fine_tuning_job", + "asearch", + "asystemone", + "avector_store_delete", + "avector_store_file_content", + "avector_store_file_create", + "avector_store_file_delete", + "avector_store_file_list", + "avector_store_file_retrieve", + "avector_store_file_update", + "avector_store_list", + "avector_store_retrieve", + "avector_store_search", + "avector_store_update", + "avideo_content", + "avideo_create_character", + "avideo_edit", + "avideo_extension", + "avideo_generation", + "avideo_get_character", + "avideo_list", + "avideo_remix", + "avideo_status", + "cancel_interaction", + "create_agent", + "create_container", + "create_interaction", + "decisions", + "delete_agent", + "delete_container", + "delete_interaction", + "extraction", + "generate_content", + "generate_content_stream", + "get_agent", + "get_interaction", + "list_agent_versions", + "list_agents", + "list_containers", + "ocr", + "responses", + "retrieve_container", + "search", + "systemone", + "vector_store_create", + "vector_store_delete", + "vector_store_file_content", + "vector_store_file_create", + "vector_store_file_delete", + "vector_store_file_list", + "vector_store_file_retrieve", + "vector_store_file_update", + "vector_store_list", + "vector_store_retrieve", + "vector_store_search", + "vector_store_update", + "video_content", + "video_create_character", + "video_edit", + "video_extension", + "video_generation", + "video_get_character", + "video_list", + "video_remix", + "video_status", + } +) + + +def _dispatch_failure(logging_obj: "Logging | None", error: Exception) -> None: + """The ``## LOG FAILURE EVENT`` of a routing callback that raised.""" + if logging_obj is not None: + asyncio.create_task( + logging_obj.dispatch_failure_handlers( + exception=error, + traceback_exception=traceback.format_exc(), + prefer_async_handlers=True, + ) + ) + + +class _FailedCallParams(TypedDict): + """The ``litellm_params`` of a failed call, as the failure callback reads them.""" + + metadata: ReadOnly[Mapping[str, object]] + model_info: ReadOnly[NotRequired[Mapping[str, object] | None]] + + +async def _deployments_filtered_by( + callback: CustomLogger, + model: str, + deployments: object, + messages: object, + request_kwargs: dict[str, object] | None, # mutable-ok: CustomLogger takes the request dict + parent_otel_span: "Span | None", +) -> object: + return await callback.async_filter_deployments( # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] # untyped deployment dicts + model=model, + healthy_deployments=deployments, # pyright: ignore[reportArgumentType] # what the last callback returned + messages=messages, # pyright: ignore[reportArgumentType] # the request's messages, as given + request_kwargs=request_kwargs, + parent_otel_span=parent_otel_span, + ) + + +def _metadata_model_id(model_info: object) -> object: + return model_info.get("id") if isinstance(model_info, dict) else None # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] # the metadata's model_info dict + + +class RustLegacyMembers(RustRegistryMembers, RustOptionalChecks): + """Mixed into ``RustCore``, which supplies its host registry, selectors and checks.""" + + # RustCore sets these in its constructor. + _cache: "DualCache" + _cooldown_given: object + _custom_strategy: "CustomRoutingStrategyBase | None" + # The constructor's argument, or the attribute as assigned. + auto_router_capability_limit: "Callable[[], int | None] | None" + + TIER_PARAMS_NEVER_DROPPED: Final = _TIER_PARAMS_NEVER_DROPPED + + # -- registry lookups -------------------------------------------------------------------------- + + def get_routing_group(self, model_name: str) -> "RoutingGroup | None": + """The routing group callable as ``model_name`` (``get_routing_group``).""" + return self._registry_state.get_routing_group(model_name) + + def team_public_model_names_of(self, team_id: str) -> frozenset[str]: + """The public names ``team_id`` serves its own deployments under (``team_public_model_names_of``).""" + return self._registry_state.team_public_model_names_of(team_id) + + def served_openai_models(self, model_name: str) -> tuple[str | None, ...]: + """The OpenAI models ``model_name`` serves (``_served_openai_models``). The Python core keeps + an ``lru_cache`` over it, emptied with every model list change.""" + return _served_openai_models(self._registry_state._core_deployments_named(model_name) or ()) # pyright: ignore[reportPrivateUsage] # the core reads its own host registry + + def auto_router_capability_violation(self, capability: GatedAutoRouterCapability) -> str | None: + """Why one more router claiming ``capability`` cannot join this router, or None when it can + (``auto_router_capability_violation``); the limit is resolved on every call.""" + resolver: Final = self.auto_router_capability_limit + limit: Final = resolver() if resolver is not None else None + others: Final = count_capability_routers( + (deployment for deployment in self._registry_state._core_model_list if isinstance(deployment, Mapping)), # pyright: ignore[reportPrivateUsage, reportUnnecessaryIsInstance] # the live list + capability=capability, + ) + return capability_limit_violation(capability=capability, held=others + 1, limit=limit) + + # -- routing state ----------------------------------------------------------------------------- + + def retire_routing_state(self) -> None: + """Stop the background work of the selectors and the budget limiter (``retire_routing_state``).""" + self._opt.retire_state((*self._sel.callback_selectors(), *self._sel.override_selectors())) + + def routing_strategy_init(self, routing_strategy: object, routing_strategy_args: Mapping[str, object]) -> None: + """``routing_strategy_init``: the default selector only. As on the Python core the router's + ``routing_strategy`` setting stays, so picks keep following it; ``apply`` changes it.""" + verbose_router_logger.info("Routing strategy: %s", routing_strategy) + validate_routing_strategy(routing_strategy) # pyright: ignore[reportArgumentType] # the strategy as given + # `_reset_custom_routing_strategy`. + self._custom_strategy = None + _ = self.__dict__.pop("get_available_deployment", None) + _ = self.__dict__.pop("_core_async_get_available_deployment", None) + self._sel.init_default(routing_strategy, routing_strategy_args, configured=False) + + def set_optional_pre_call_checks(self, optional_pre_call_checks: Iterable[object] | None) -> None: + """``set_optional_pre_call_checks``: drop the toggleable checks the list leaves out, then add it.""" + self._opt.set(None if optional_pre_call_checks is None else tuple(optional_pre_call_checks)) + + def initialize_assistants_endpoint(self) -> None: + """``initialize_assistants_endpoint`` binds the endpoints; this core answers them from + ``__getattr__``, so what is left to do is to drop an assigned replacement.""" + for name in _ASSISTANTS_ENDPOINTS: + _ = self.__dict__.pop(name, None) + + def initialize_router_endpoints(self) -> None: + """``initialize_router_endpoints``, as ``initialize_assistants_endpoint`` above.""" + from litellm.containers.endpoint_factory import ( + _generated_endpoints, # pyright: ignore[reportPrivateUsage, reportUnknownVariableType] # the names the Python core registers + ) + + for name in (*_ROUTER_ENDPOINTS, *_generated_endpoints): + _ = self.__dict__.pop(name, None) + + # -- pre-routing strategies -------------------------------------------------------------------- + + def init_auto_router_deployment(self, deployment: "Deployment") -> None: + """Register the deployment's auto-router (``init_auto_router_deployment``).""" + self._registry_state.strategies._init_auto_router(deployment) # pyright: ignore[reportOptionalMemberAccess, reportPrivateUsage] # RustCore installs its strategies in its constructor + + def init_complexity_router_deployment(self, deployment: "Deployment") -> None: + """Register the deployment's complexity-router (``init_complexity_router_deployment``).""" + self._registry_state.strategies._init_complexity_router(deployment) # pyright: ignore[reportOptionalMemberAccess, reportPrivateUsage] # see above + + def init_quality_router_deployment(self, deployment: "Deployment") -> None: + """Register the deployment's quality-router (``init_quality_router_deployment``).""" + self._registry_state.strategies._init_quality_router(deployment) # pyright: ignore[reportOptionalMemberAccess, reportPrivateUsage] # see above + + def init_adaptive_router_deployment(self, deployment: "Deployment") -> None: + """Register the deployment's adaptive-router and its post-call hook (``init_adaptive_router_deployment``).""" + self._registry_state.strategies._init_adaptive_router(deployment) # pyright: ignore[reportOptionalMemberAccess, reportPrivateUsage] # see above + + # -- routing callbacks ------------------------------------------------------------------------- + + def routing_strategy_pre_call_checks( + self, + deployment: dict[str, object], # mutable-ok: CustomLogger takes the deployment dict + ) -> None: + """``routing_strategy_pre_call_checks``: the sync check of each ``CustomLogger`` callback.""" + for callback in self._opt.routing_callbacks(): + if isinstance(callback, CustomLogger): + callback.pre_call_check(deployment) # pyright: ignore[reportUnknownMemberType] # CustomLogger takes an untyped deployment dict + + async def async_routing_strategy_pre_call_checks( + self, + deployment: dict[str, object], # mutable-ok: CustomLogger takes the deployment dict + parent_otel_span: "Span | None", + logging_obj: "Logging | None" = None, + ) -> None: + """``async_routing_strategy_pre_call_checks``: a check's rate limit cools the deployment down + under the Python rules (``set_cooldown_deployments``, cooldown event included).""" + from litellm import exceptions as litellm_exceptions + + for callback in self._opt.routing_callbacks(): + if isinstance(callback, CustomLogger): + try: + await callback.async_pre_call_check(deployment, parent_otel_span) # pyright: ignore[reportUnknownMemberType] # see above + except litellm_exceptions.RateLimitError as e: + _dispatch_failure(logging_obj, e) + _ = set_cooldown_deployments( + ctx=RoutingContext(self), # pyright: ignore[reportArgumentType] # the core's RoutingContext + exception_status=e.status_code, + original_exception=e, + deployment=deployment["model_info"]["id"], # pyright: ignore[reportIndexIssue, reportUnknownArgumentType] # read as the Python core reads it + time_to_cooldown=self._cooldown_given, # pyright: ignore[reportArgumentType] # the value as given, as the Python core stores it + ) + raise + except Exception as e: + _dispatch_failure(logging_obj, e) + raise + + async def async_callback_filter_deployments( + self, + model: str, + healthy_deployments: list[dict[str, object]], # mutable-ok: CustomLogger deployment-list contract + messages: object, + parent_otel_span: "Span | None", + request_kwargs: dict[str, object] | None = None, # mutable-ok: CustomLogger takes the request dict + logging_obj: "Logging | None" = None, + ) -> object: + """``async_callback_filter_deployments``: each routing callback narrows the deployments in turn.""" + remaining: object = healthy_deployments # rebind-ok: narrowed by each callback + for callback in self._opt.routing_callbacks(): + if isinstance(callback, CustomLogger): + try: + remaining = await _deployments_filtered_by( + callback, model, remaining, messages, request_kwargs, parent_otel_span + ) + except Exception as e: + _dispatch_failure(logging_obj, e) + raise + return remaining + + async def async_admit_deployment_turn( + self, model_group: str, opened_on: Mapping[str, object], request_kwargs: Mapping[str, object] + ) -> Mapping[str, object]: + """Admit a later turn of an open websocket (``async_admit_deployment_turn``).""" + port: Final[EndpointCore] = self.__dict__["_port"] # pyright: ignore[reportAny] # RustCore's own attribute + return await admit_deployment_turn(port, model_group, opened_on, request_kwargs) + + # -- one provider call ------------------------------------------------------------------------- + + async def make_call( + self, + original_function: Callable[..., object], + *args: object, + **kwargs: object, # kwargs-ok: the endpoint's request + ) -> object: + """``make_call``: call the endpoint, count its usage and stamp the response's headers. A + failed call hands back the io tokens it reserved.""" + model_group: Final = kwargs.get("model") + generic_function: Final = kwargs.get("original_generic_function") + endpoint_function: Final = generic_function if callable(generic_function) else original_function + try: + with endpoint_selection( + getattr(endpoint_function, "__name__", None), + getattr(endpoint_function, "__module__", None), + model=model_group, + ): + response: object = original_function(*args, **kwargs) # rebind-ok: awaited below + if coroutine_checker.is_async_callable(response) or inspect.isawaitable(response): + awaitable: Final[Awaitable[object]] = response # pyright: ignore[reportAssignmentType] # the endpoint's awaitable + response = await awaitable # rebind-ok: awaited here + except Exception: + with contextlib.suppress(Exception): + await async_io_token_refund_failure( + self._cache, kwargs, parent_otel_span=get_parent_otel_span_from_kwargs(kwargs) + ) + raise + await self.increment_deployment_usage_for_response(response=response, request_kwargs=kwargs) + return await self.set_response_headers( + response=response, model_group=model_group if isinstance(model_group, str) else None, request_kwargs=kwargs + ) + + async def increment_deployment_usage_for_response( + self, response: object, request_kwargs: Mapping[str, object] + ) -> None: + """Count a response's tokens and request against its deployment, once the call returned, and + note them on the request for the success callback (``increment_deployment_usage_for_response``).""" + if response is None: + return + try: + deployment_metadata: Final = find_deployment_metadata(request_kwargs) + model_group: Final = request_kwargs.get("model") + if deployment_metadata is None or not isinstance(model_group, str): + return + deployment_id: Final = _metadata_model_id(deployment_metadata["model_info"]) + if deployment_id is None: + return + total_tokens: Final = response_total_token_count(response) + deployment_name: Final = deployment_metadata.get("deployment") + deployment_metadata[ROUTER_USAGE_COUNTED_TOKENS_METADATA_KEY] = total_tokens + try: + await self._increment_deployment_usage( + deployment_id=str(deployment_id), + deployment_name=deployment_name if isinstance(deployment_name, str) else None, + model_group=model_group, + total_tokens=total_tokens, + parent_otel_span=get_parent_otel_span_from_kwargs(dict(request_kwargs)), + ) + except Exception: + _ = deployment_metadata.pop(ROUTER_USAGE_COUNTED_TOKENS_METADATA_KEY, None) + raise + except Exception as e: # noqa: BLE001 # as the Python core: usage accounting never fails a served call + verbose_router_logger.debug( + "litellm.router.Router::increment_deployment_usage_for_response(): Exception occured - %s", e + ) + + async def _increment_deployment_usage( + self, + *, + deployment_id: str, + deployment_name: str | None, + model_group: str, + total_tokens: float, + parent_otel_span: "Span | None", + ) -> None: + """``_increment_deployment_usage`` with its one request, in one pipeline; this worker's memory + then holds what Redis answered.""" + from litellm.types.caching import RedisPipelineIncrementOperation + + if not self._registry_state.cached_deployment_tracks_usage(deployment_id, model_group): + return + current_minute: Final = get_utc_datetime().strftime("%H-%M") + tpm_key: Final = RouterCacheEnum.TPM.value.format( + id=deployment_id, current_minute=current_minute, model=deployment_name + ) + rpm_key: Final = RouterCacheEnum.RPM.value.format( + id=deployment_id, current_minute=current_minute, model=deployment_name + ) + operations: Final[Sequence[RedisPipelineIncrementOperation]] = [ + RedisPipelineIncrementOperation(key=key, increment_value=increment_value, ttl=_USAGE_TTL) + for key, increment_value in ((tpm_key, total_tokens), (rpm_key, 1)) + ] + values: Final = await self._cache.async_increment_cache_pipeline( + increment_list=list(operations), parent_otel_span=parent_otel_span + ) + if values is not None and self._cache.redis_cache is not None: + for operation, value in zip(operations, values): + await self._cache.async_set_cache(operation["key"], int(value), local_only=True, ttl=_USAGE_TTL) # pyright: ignore[reportUnknownMemberType] # DualCache's writes are untyped + + async def async_deployment_callback_on_failure( + self, kwargs: Mapping[str, object], completion_response: object, start_time: object, end_time: object + ) -> None: + """Count a failed call as one more request of the minute (``async_deployment_callback_on_failure``).""" + if is_batch_retrieve_call_type(kwargs.get("call_type")): + return + litellm_params: Final[_FailedCallParams] = kwargs["litellm_params"] # pyright: ignore[reportAssignmentType] # read as the Python core reads it, a missing key included + deployment_name: Final = litellm_params["metadata"].get("deployment", None) + model_group: Final = litellm_params["metadata"].get("model_group", None) + model_info: Final = litellm_params.get("model_info", {}) or {} + deployment_id: Final = model_info.get("id", None) + if model_group is None or deployment_id is None: + return + rpm_key: Final = RouterCacheEnum.RPM.value.format( + id=deployment_id, + current_minute=get_utc_datetime().strftime("%H-%M"), + model=deployment_name, + ) + with service_target(ROUTER_USAGE_TARGET): + await self._cache.async_increment_cache( # pyright: ignore[reportUnknownMemberType] # DualCache's writes are untyped + key=rpm_key, + value=1, + parent_otel_span=get_parent_otel_span_from_kwargs(dict(kwargs)), + ttl=_USAGE_TTL, + ) + + # -- the plugin context ------------------------------------------------------------------------ + + def _update_kwargs_before_fallbacks( + self, + model: str, + kwargs: dict[str, object], # mutable-ok: the core updates the request kwargs in place + metadata_variable_name: str | None = "metadata", + ) -> None: + """``_update_kwargs_before_fallbacks``: a trace id and the requested group, in place.""" + _ = kwargs.setdefault("litellm_trace_id", str(uuid.uuid4())) + alias: Final = model if self._registry_state._core_get_model_from_alias(model=model) else None # pyright: ignore[reportPrivateUsage] # the core reads its own host registry + kwargs.setdefault(metadata_variable_name, {}).update({"model_group": model, "model_group_alias": alias}) # pyright: ignore[reportArgumentType, reportAttributeAccessIssue, reportUnknownMemberType] # the bucket, as the Python core writes it + + def _params_sent_to_deployment( + self, request_kwargs: Mapping[str, object], deployment: "Deployment | None" + ) -> Mapping[str, object]: + """``_params_sent_to_deployment`` of a ``Deployment``, which ``RoutingContext.deployment_params`` hands in.""" + params: Final = deployment.litellm_params.model_dump(exclude_none=True) if deployment else {} + configured: Final = as_str_mapping(RoutingContext(self).settings().default_litellm_params) # pyright: ignore[reportArgumentType, reportUnknownMemberType, reportUnknownArgumentType] # the core's RoutingContext + defaults: Final = {k: v for k, v in (configured or {}).items() if v is not None} + tool_choice: Final = {"tool_choice": params["tool_choice"]} if params.get("tool_choice") is not None else {} + return MappingProxyType({**params, **defaults, **tool_choice, **request_kwargs}) + + +__all__ = ["RustLegacyMembers"] diff --git a/litellm/router_rust/surface/members.py b/litellm/router_rust/surface/members.py new file mode 100644 index 00000000000..0f7dc23d966 --- /dev/null +++ b/litellm/router_rust/surface/members.py @@ -0,0 +1,163 @@ +""" +``RustRegistryMembers``: the Rust core's public registry members outside ``registry()`` and +``apply``: ``add_deployment``, ``set_model_list``, the ``model_list`` attribute, the environment +gate, the deprecated id hash, ``print_deployment`` and the pre-call checks the plugin context reads. + +Each one works on the core's host registry (``litellm/router_rust/surface/registry.py``), which hands every +change to the Rust core. The rules are the Python router core's, ported (never imported). +""" + +import copy +from collections.abc import Mapping, Sequence +from typing import TYPE_CHECKING, Final + +from litellm._logging import verbose_router_logger +from litellm.litellm_core_utils.sensitive_data_masker import SensitiveDataMasker +from litellm.router_api import deployment_id +from litellm.router_rust.steps.precall import count_tokens, pre_call_checks +from litellm.router_rust.surface.registry import HostRegistry, model_info_is_active_for_environment +from litellm.types.router import Deployment + +if TYPE_CHECKING: + from litellm.router_python_host import PythonHost + from litellm.router_rust.core import RustCore + + +class _CallCounter: + """``total_calls``, ``success_calls`` or ``fail_calls``: plain attributes on the Python core, + which a write replaces and a ``del`` removes; every call then reads what is there.""" + + def __init__(self, name: str) -> None: + self._name: Final = name + + def __get__(self, core: "RustCore | None", owner: type | None = None) -> object: + if core is None: + return self + try: + return getattr(core._endpoints.counters, self._name) # pyright: ignore[reportPrivateUsage] # the core's own counters + except AttributeError: + raise core._no_attribute(self._name) from None # pyright: ignore[reportPrivateUsage] # see above + + def __set__(self, core: "RustCore", value: object) -> None: + setattr(core._endpoints.counters, self._name, value) # pyright: ignore[reportPrivateUsage] # see __get__ + + def __delete__(self, core: "RustCore") -> None: + try: + delattr(core._endpoints.counters, self._name) # pyright: ignore[reportPrivateUsage] # see __get__ + except AttributeError: + raise core._no_attribute(self._name) from None # pyright: ignore[reportPrivateUsage] # see __get__ + + +class RustRegistryMembers: + """Mixed into ``RustCore``, which supplies its host registry and host.""" + + total_calls: Final = _CallCounter("total_calls") + success_calls: Final = _CallCounter("success_calls") + fail_calls: Final = _CallCounter("fail_calls") + + @property + def _registry_state(self) -> HostRegistry: + raise NotImplementedError + + @property + def _python_host(self) -> "PythonHost": + raise NotImplementedError + + def add_deployment(self, deployment: Deployment) -> Deployment | None: + """Add a deployment unless its id exists (``add_deployment``).""" + return self._registry_state.add_deployment(deployment) + + def set_model_list(self, model_list: Sequence[object]) -> None: + """Replace every deployment (``set_model_list``); aliases and wildcard routes stay.""" + self._registry_state.set_model_list(model_list) + + @property + def model_list(self) -> Sequence[object]: + """The live model list (the deprecated ``model_list`` attribute).""" + return self._registry_state._core_model_list # pyright: ignore[reportPrivateUsage] # the core reads its own host registry + + def _assign_model_list(self, model_list: Sequence[object]) -> None: + """``router.model_list = [...]``: the list is swapped as given, without reindexing, and the + registry version moves once.""" + self._registry_state.replace_model_list(model_list) # pyright: ignore[reportArgumentType] # the list as assigned, as the core keeps it + + @staticmethod + def deployment_is_active_for_environment(deployment: Deployment) -> bool: + """Whether LITELLM_ENVIRONMENT admits the deployment (``deployment_is_active_for_environment``).""" + return model_info_is_active_for_environment(model_info=deployment.model_info) # pyright: ignore[reportArgumentType] # ModelInfo reads like a mapping + + @staticmethod + def generate_model_id(model_group: str, litellm_params: Mapping[str, object]) -> str: + """The deprecated id hash (``generate_model_id``).""" + from litellm.router_deprecated import ( + _note, # pyright: ignore[reportPrivateUsage] # the router's deprecation notice + ) + + _note("generate_model_id", "litellm.router_api.deployment_id(model_group, litellm_params)") + return deployment_id(model_group, litellm_params) + + async def set_response_headers( + self, response: object, model_group: str | None = None, request_kwargs: Mapping[str, object] | None = None + ) -> object: + """Stamp the group's remaining usage and the model id on a response (``set_response_headers``).""" + headed: Final[object] = await self._registry_state.set_response_headers( # pyright: ignore[reportAny] # the response as the core hands it back + response, model_group, None if request_kwargs is None else dict(request_kwargs) + ) + return headed + + def replay_model_cost_registrations(self) -> None: + """Re-register the served deployments' cost-map entries (``replay_model_cost_registrations``).""" + self._registry_state.replay_model_cost_registrations() + + def _count_pre_call_check_tokens( + self, + messages: Sequence[Mapping[str, object]] | None, + input: str | list[object] | None, # mutable-ok: the request field, as the Python core takes it + request_kwargs: Mapping[str, object] | None = None, + tokenizer_model: str = "", + deployments: Sequence[Mapping[str, object]] = (), + ) -> int: + """``_count_pre_call_check_tokens``, which context compaction counts with.""" + return count_tokens(self._python_host, messages, input, request_kwargs, tokenizer_model, deployments) + + def _pre_call_checks( + self, + model: str, + healthy_deployments: Sequence[Mapping[str, object]], + messages: object = None, + input: object = None, + request_kwargs: Mapping[str, object] | None = None, + input_token_count: int | None = None, + skip_inline_token_count: bool = False, + request_deployments: Sequence[Mapping[str, object]] | None = None, + ) -> Sequence[Mapping[str, object]]: + """``_pre_call_checks``, which context compaction admits a client-managed request with.""" + return pre_call_checks( + self._python_host, + self._registry_state, # pyright: ignore[reportArgumentType] # the host registry serves the reads + model, + healthy_deployments, + messages, + input, + request_kwargs, + input_token_count, + skip_inline_token_count, + request_deployments, + ) + + def print_deployment(self, deployment: Mapping[str, object]) -> Mapping[str, object]: + """A copy of the deployment, of the type it was given, with its api key masked + (``print_deployment``). A ``Deployment`` and its params read and write like dicts.""" + try: + printed: Final = copy.deepcopy(deployment) + litellm_params: Final = printed["litellm_params"] + if self._python_host.dynamic().redact_user_api_key_info: + masker: Final = SensitiveDataMasker(visible_prefix=2, visible_suffix=0) + printed["litellm_params"] = masker.mask_dict(litellm_params) # pyright: ignore[reportArgumentType, reportIndexIssue] # the deployment's params, written as the Python core writes them + elif "api_key" in litellm_params: # pyright: ignore[reportOperatorIssue] # see the docstring + litellm_params["api_key"] = litellm_params["api_key"][:2] + "*" * 10 # pyright: ignore[reportIndexIssue] # see the docstring + except Exception as e: + verbose_router_logger.debug("Error occurred while printing deployment - %s", e) + raise + else: + return printed diff --git a/litellm/router_rust/surface/model_info.py b/litellm/router_rust/surface/model_info.py new file mode 100644 index 00000000000..f742225e2f4 --- /dev/null +++ b/litellm/router_rust/surface/model_info.py @@ -0,0 +1,871 @@ +# pyright: reportUnknownMemberType=false, reportUnknownVariableType=false, reportUnknownArgumentType=false, reportUnknownParameterType=false, reportMissingParameterType=false, reportAny=false, reportExplicitAny=false, reportMissingTypeArgument=false, reportAttributeAccessIssue=false, reportTypedDictNotRequiredAccess=false +# (ported from the Python router core as it is written there, over untyped litellm helpers) +""" +``HostModelInfo``: model info, group info, usage reads and response headers for the Rust core's +host registry. + +router-core holds the rules: a deployment's model info (``get_router_model_info``, +``get_deployment_model_info``: ``router-core/src/model_info.rs``), a model group's info and the +configured-value reads (``_set_model_group_info``, ``get_model_listing_info``, the configured token +limits, mode, display name and service tiers: ``router-core/src/group_info.rs``), and the usage +counters' keys (``RouterCoreHandle.usage_key``). What lives in this process is answered here as a +query asks for it, at query time: the cost map (``get_model_info``, ``litellm.model_cost``), +provider resolution, the wildcard pattern routes, what the background refresh discovered, and a +group deployment's validated params, supported params, endpoints, fast mode and reasoning efforts. +``ModelGroupInfo`` is built here from the group's first deployment (its validation is the model's) +and takes what router-core's merge assigned. + +Still the Python router core's rules here: the model group alias a group-info query resolves +first (the alias map is the registry's live Python objects), the cache of discovered limits, the +configured ``max_input_tokens`` read (a value echoed from the loaded cost map is skipped), and the +response headers. +""" + +import builtins +import dataclasses +import math +from collections.abc import Callable, Mapping, Sequence +from functools import partial +from types import MappingProxyType +from typing import TYPE_CHECKING, Any, Final, Union, cast # noqa: TID251 # the core's own signatures, ported unchanged + +from pydantic import BaseModel, TypeAdapter +from typing_extensions import overload + +import litellm +from litellm._internal_context import service_target +from litellm._logging import verbose_router_logger +from litellm.litellm_core_utils.add_retry_fallback_headers import ( + apply_quality_router_decision_headers, + apply_remaining_usage_headers, + apply_response_model_id, + complexity_router_decision_headers, + ensure_response_additional_headers, + prepare_response_for_header_attachment, +) +from litellm.litellm_core_utils.batch_utils import deployment_supported_endpoints +from litellm.litellm_core_utils.core_helpers import as_str_mapping, coerce_token_limit +from litellm.litellm_core_utils.get_model_cost_map import GetModelCostMap +from litellm.litellm_core_utils.reasoning_effort_capability import ( + deployment_is_catalog_mapped, + intersect_supported_reasoning_efforts, + resolve_supported_reasoning_efforts, +) +from litellm.litellm_core_utils.routing_common_utils import team_may_use_deployment +from litellm.llms.anthropic.common_utils import AnthropicModelInfo +from litellm.llms.bedrock.common_utils import get_bedrock_base_model +from litellm.router_python_host import python_host +from litellm.router_utils.pre_call_checks.io_token_rate_limit_check import build_io_token_rate_limit_headers +from litellm.router_utils.router_callbacks.track_deployment_metrics import find_deployment_metadata +from litellm.router_utils.routing_read_batch import ROUTER_USAGE_TARGET +from litellm.types.router import ( + Deployment, + DeploymentModelListingInfo, + DeploymentTypedDict, + DiscoveredDeploymentModelInfo, + LiteLLM_Params, + ModelGroupInfo, +) +from litellm.types.utils import COST_MAP_LOOKUP_KEY, ModelInfo, echoed_cost_map_fields +from litellm.types.utils import ModelInfo as ModelMapInfo +from litellm.utils import get_utc_datetime, is_generalized_model_info + +if TYPE_CHECKING: + from litellm.rust_bridge._native import ( + RouterCall, + RouterCoreHandle, + _ModelInfoSteps, # pyright: ignore[reportPrivateUsage] # the stub's step shapes + ) + +_MODEL_INFO_ADAPTER: Final = TypeAdapter(Mapping[str, object]) + +# The key router-core's merge reads as an opaque value (``router-core/src/model_info.rs``). +_OPAQUE_KEY: Final = "__litellm_python_value__" + + +class _Wire: + """Values as router-core's model-info merge reads them: plain JSON, with what JSON cannot spell + (a datetime, a tuple, a non-finite float) passed as an opaque token the merge moves between + layers without reading, and restored to the same object on the way back.""" + + def __init__(self) -> None: + self._values: Final[list[object]] = [] # mutable-ok: the tokens of one query + + def to_wire(self, value: object) -> object: + if value is None or isinstance(value, (bool, str)): + return value + if isinstance(value, int): + return value if -(2**63) <= value < 2**64 else self._token(value) + if isinstance(value, float): + return value if math.isfinite(value) else self._token(value) + if isinstance(value, list): + return [self.to_wire(item) for item in value] + if isinstance(value, dict) and all(isinstance(key, str) for key in value): + return {str(key): self.to_wire(item) for key, item in value.items()} + return self._token(value) + + def from_wire(self, value: object) -> object: + if isinstance(value, dict): + if len(value) == 1 and _OPAQUE_KEY in value: + return self._values[int(value[_OPAQUE_KEY])] + return {str(key): self.from_wire(item) for key, item in value.items()} + if isinstance(value, list): + return [self.from_wire(item) for item in value] + return value + + def _token(self, value: object) -> Mapping[str, int]: + self._values.append(value) + return {_OPAQUE_KEY: len(self._values) - 1} + + +def _drive_model_info( + call: "RouterCall[_ModelInfoSteps]", answer: Callable[[Mapping[str, object]], object] +) -> Mapping[str, object]: + """Drive one model-info merge to its end, answering each lookup with ``answer``. A lookup that + raises raises out of the query, as it raised in the Python router.""" + outcome: object = None # rebind-ok: each lookup's answer + while True: + step = call.next(outcome) + if step["kind"] == "done": + return step + if step["kind"] == "failed": + if str(step.get("error_kind")) == "model_info": + exception: object = getattr(builtins, str(step.get("error_exception")), ValueError) + raise (exception if isinstance(exception, type) else ValueError)(step.get("error_detail")) + raise RuntimeError(f"router-core ({step.get('error_kind')}): {step.get('error')}") + outcome = answer(step) + + +def _configured_model_info(deployment: Mapping[str, object]) -> Mapping[str, object]: + return _MODEL_INFO_ADAPTER.validate_python(deployment.get("model_info") or {}) + + +def _operator_max_input_tokens(deployment: Mapping[str, object]) -> int | None: + """``_operator_max_input_tokens``: the deployment's own limit, skipping one a saved + ``/model/info`` response copied from the cost map.""" + model_info: Final = _configured_model_info(deployment) + catalog_key: Final = model_info.get(COST_MAP_LOOKUP_KEY) + catalog_entry: Final = ( + GetModelCostMap.loaded_model_cost_map().get(catalog_key) if isinstance(catalog_key, str) else None + ) + if catalog_entry is not None and "max_input_tokens" in echoed_cost_map_fields(model_info, catalog_entry): + return None + return coerce_token_limit(model_info.get("max_input_tokens")) + + +# The ModelGroupInfo fields router-core's group merge compares against, as the model validated them. +_GROUP_COMPARED_FIELDS: Final = frozenset( + { + "providers", + "max_input_tokens", + "max_output_tokens", + "input_cost_per_token", + "output_cost_per_token", + "supports_fast_mode", + "supported_endpoints", + "supported_reasoning_efforts", + } +) +# The ModelGroupInfo fields the Python router assigns tuples to. +_GROUP_TUPLE_FIELDS: Final = frozenset({"supported_endpoints", "supported_reasoning_efforts"}) + + +def _group_field(field: str, value: object) -> object: + """A merged group field as the Python router assigns it: a tuple field takes the list as a tuple.""" + if field in _GROUP_TUPLE_FIELDS and isinstance(value, list): + return tuple(cast(list[object], value)) # cast-ok: the merge's list + return value + + +@dataclasses.dataclass +class _GroupBuild: + """One group-info merge as the host answers it: the deployments it reads, each one's validated + params once its first lookup builds them, and the group built from the first deployment.""" + + model_list: "Sequence[dict[str, Any]]" # mutable-ok: the stored entries, read only + team_id: str | None + # Filled as the lookups arrive. + params: dict[int, LiteLLM_Params] = dataclasses.field(default_factory=dict) # mutable-ok: see above + info: ModelGroupInfo | None = None + + +class HostModelInfo: + """Mixed into ``HostRegistry``, which holds the model list, the indexes and the host.""" + + def _live_handle(self) -> "RouterCoreHandle | None": ... # the registry's: the handle ``RustCore`` set + + def _core_handle(self) -> "RouterCoreHandle": + handle: Final = self._live_handle() + if handle is None: + raise RuntimeError("the host registry reads model info through its Rust core") + return handle + + def _core_get_model_list_of_routed_group( + self, model_group: str, team_id: str | None = None + ) -> list: # mutable-ok: a fresh list, as the core returns it + if team_id is not None and self._team_deployments_named(team_id, model_group): + return self._core_get_all_deployments(model_name=model_group, team_id=team_id) + named: Final = [ + *self._core_get_all_deployments(model_name=model_group), + *self._core_get_model_list_from_routing_groups(model_name=model_group), + ] + return named or self._get_wildcard_deployments(model_name=model_group) + + def _get_wildcard_deployments( + self, model_name: str, team_id: str | None = None + ) -> list[dict[str, object]]: # mutable-ok: a fresh list + """The deployments a wildcard pattern routes ``model_name`` to, renamed to it (the routing + runs in router-core, through the registry's pattern router views).""" + team_router: Final = self._core_team_pattern_routers.get(team_id) if team_id is not None else None + matches: Final = [ + *self._core_pattern_router.get_deployments_by_pattern(model=model_name), + *(team_router.get_deployments_by_pattern(model=model_name) if team_router is not None else ()), + ] + return [{**DeploymentTypedDict(**m), "model_name": model_name} for m in matches] # pyright: ignore[reportCallIssue] # a stored entry + + def _core_get_model_group_info( + self, model_group: str, *, include_hidden: bool = False, team_id: str | None = None + ) -> ModelGroupInfo | None: + """ + For a given model group name, return the combined model info + + Returns: + - ModelGroupInfo if able to construct a model group + - None if error constructing model group info or hidden model group (unless include_hidden) + """ + ## Check if model group alias + if model_group in self._core_model_group_alias: + item: Final = self._core_model_group_alias[model_group] + if isinstance(item, str): + _router_model_group = item # rebind-ok: ported from the Python core unchanged + elif isinstance(item, dict): + if item["hidden"] is True and not include_hidden: + return None + else: + _router_model_group = item["model"] # rebind-ok: ported from the Python core unchanged + else: + return None + + return self._set_model_group_info( + model_group=_router_model_group, + user_facing_model_group_name=model_group, + team_id=team_id, + ) + + ## Check if actual model + return self._set_model_group_info( + model_group=model_group, user_facing_model_group_name=model_group, team_id=team_id + ) + + def _set_model_group_info( + self, model_group: str, user_facing_model_group_name: str, team_id: str | None = None + ) -> ModelGroupInfo | None: + """ + For a given model group name, return the combined model info + + Returns: + - ModelGroupInfo if able to construct a model group + - None if error constructing model group info + + router-core merges the group (``router-core/src/group_info.rs``); this picks its + deployments, answers what needs this process (validated params, model info, provider, + capabilities), builds ``ModelGroupInfo`` from the first deployment and sets what the merge + assigned on it. + """ + serves_team_deployments: Final = team_id is not None and bool( + self._team_deployments_named(team_id, model_group) + ) + model_list: Final = self._core_get_model_list_of_routed_group(model_group, team_id) + description: Final[str | None] = next( + (d for d in (_configured_model_info(m).get("description") for m in model_list) if isinstance(d, str)), + None, + ) + # The wildcard check is asked only for a deployment that is not the group's by name. + pattern_routes: Final = ( + not serves_team_deployments + and any("model_name" in m and m["model_name"] != model_group for m in model_list) + and self._core_pattern_router.route(model_group) is not None + ) + wire: Final = _Wire() + group: Final = _GroupBuild(model_list=model_list, team_id=team_id) + call: Final = self._core_handle().model_group_info( + { + "model_group": model_group, + "user_facing_model_group_name": user_facing_model_group_name, + "serves_team_deployments": serves_team_deployments, + "pattern_routes": pattern_routes, + "description": description, + "deployments": [wire.to_wire(model) for model in model_list], + } + ) + done: Final = _drive_model_info(call, partial(self._answer_group_lookup, group=group, wire=wire)) + answer: Final = as_str_mapping(done["value"]) + if answer is None or group.info is None: + return None + for field, value in (as_str_mapping(answer["assignments"]) or {}).items(): + setattr(group.info, field, _group_field(field, wire.from_wire(value))) + return group.info + + def _answer_group_lookup(self, step: Mapping[str, object], group: "_GroupBuild", wire: "_Wire") -> object: + """One lookup of a group-info merge: a deployment's facts, or the group's first entry.""" + match step["lookup"]: + case "group_params": + index: Final = int(str(step["index"])) + litellm_params: Final = LiteLLM_Params(**group.model_list[index]["litellm_params"]) + group.params[index] = litellm_params # rebind-ok: the merge's state, which each lookup advances + return { + "model": litellm_params.model, + "configurable_clientside_auth_params": wire.to_wire( + litellm_params.configurable_clientside_auth_params + ), + } + case "group_deployment": + return self._group_deployment(group, int(str(step["index"])), wire) + case "group_capabilities": + deployment: Final = group.model_list[int(str(step["index"]))] + merged_info: Final[object] = wire.from_wire(step["model_info"]) + model_info: Final = cast(Mapping[str, object], merged_info) # cast-ok: the merged model info + litellm_model: Final = str(step["litellm_model"]) + llm_provider: Final = str(step["llm_provider"]) + efforts: Final = resolve_supported_reasoning_efforts( + model_info, deployment_is_mapped=step["mapped"] is True + ) + return { + "supported_endpoints": wire.to_wire( + list( + deployment_supported_endpoints( + model_info.get("supported_endpoints"), deployment, llm_provider + ) + ) # pyright: ignore[reportArgumentType] # a stored entry, as the core passes it + if team_may_use_deployment(_configured_model_info(deployment).get("team_id"), group.team_id) # pyright: ignore[reportArgumentType] # the configured team id + else [] + ), + "supports_fast_mode": AnthropicModelInfo.supports_fast_mode(litellm_model, llm_provider), + "reasoning_efforts": None if efforts is None else list(efforts), + } + case "group_seed": + entry: Final[object] = wire.from_wire(step["info"]) + seed: Final = cast(dict[str, Any], entry) # cast-ok: the first deployment's entry + group.info = ModelGroupInfo(**seed) # rebind-ok: the merge's state, built from the first deployment + return group.info.model_dump(mode="json", include=set(_GROUP_COMPARED_FIELDS)) + case "reasoning_intersect": + current: Final = cast("Sequence[str] | None", step["current"]) # cast-ok: plain JSON efforts + resolved: Final = cast("Sequence[str] | None", step["resolved"]) # cast-ok: plain JSON efforts + intersected: Final = intersect_supported_reasoning_efforts(current, resolved) + return None if intersected is None else list(intersected) + case _: + return self._answer_model_info_lookup(step, wire) + + def _group_deployment(self, group: "_GroupBuild", index: int, wire: "_Wire") -> Mapping[str, object]: + """A group deployment's model info (None when it fails), whether the catalog maps it, and its + provider; with no model info, the provider's supported params.""" + model: Final = group.model_list[index] + litellm_params: Final = group.params[index] + model_info_dict: Final = model.get("model_info", {}) + model_info: object + try: + model_id: Final = model_info_dict.get("id", None) + model_info = ( # rebind-ok: one branch binds it + self._core_get_deployment_model_info(model_id=model_id, model_name=litellm_params.model) + if model_id is not None + else None + ) + except Exception: # noqa: BLE001 # ported as the core writes it + model_info = None # rebind-ok: see above + deployment_is_mapped: Final = deployment_is_catalog_mapped(model_info, model_info_dict) # pyright: ignore[reportArgumentType] # the merged model info, or None + litellm_model, llm_provider = self._model_group_llm_provider(litellm_params) + answer: Final[dict[str, object]] = { # mutable-ok: built here + "model_info": wire.to_wire(model_info), + "mapped": deployment_is_mapped, + "litellm_model": litellm_model, + "llm_provider": llm_provider, + } + if model_info is None: + answer["supported_openai_params"] = wire.to_wire( + self._host.invoke("get_supported_openai_params")(model=litellm_model, custom_llm_provider=llm_provider) + ) + return answer + + def _cached_get_model_group_info(self, model_group: str) -> ModelGroupInfo | None: + """ + Cached version of get_model_group_info, uses @lru_cache wrapper + + This is a speed optimization, since set_response_headers makes a call to get_model_group_info on every request + """ + return self._core_get_model_group_info(model_group) + + def _core_cached_model_group_info(self, model_group: str) -> ModelGroupInfo | None: + return self._cached_get_model_group_info(model_group) + + @overload + def _core_get_router_model_info( + self, + deployment: Union[dict, "Deployment"], # mutable-ok: ported from the Python core unchanged + received_model_name: str, + id: None = None, + ) -> ModelMapInfo: + pass + + @overload + def _core_get_router_model_info(self, deployment: None, received_model_name: str, id: str) -> ModelMapInfo: + pass + + def _core_get_router_model_info( + self, + deployment: Union[dict, "Deployment"] | None, # mutable-ok: ported from the Python core unchanged + received_model_name: str, + id: str | None = None, + ) -> ModelMapInfo: + """ + For a given model id, return the model info (max tokens, input cost, output cost, etc.). + + Augment litellm info with additional params set in `model_info`. + + For azure models, ignore the `model:`. Only set max tokens, cost values if base_model is set. + For every other provider, a base_model naming the same model as `model:` can only widen its context window. + + router-core merges it (``router-core/src/model_info.rs``); this checks the deployment and + answers the cost-map lookups the merge asks for while it runs. + + Returns + - ModelInfo - If found -> typed dict with max tokens, input cost, etc. + + Raises: + - ValueError -> If model is not mapped yet + """ + if id is not None: + _deployment: Final = self._core_get_deployment(model_id=id) + if _deployment is not None: + deployment = _deployment # rebind-ok: ported from the Python core unchanged + + if deployment is None: + raise ValueError("Deployment not found") + + model_info_base_model: Final = (deployment.get("model_info") or {}).get("base_model", None) + params_base_model: Final = ( + (deployment.get("litellm_params") or {}).get("base_model", None) if model_info_base_model is None else None + ) + + ## GET PROVIDER - reuse LiteLLM_Params if already constructed + litellm_params_data: Final = deployment.get("litellm_params") + litellm_params: LiteLLM_Params + if isinstance(litellm_params_data, LiteLLM_Params): + litellm_params = litellm_params_data # rebind-ok: ported from the Python core unchanged + elif isinstance(litellm_params_data, dict) and "model" in litellm_params_data: + litellm_params = LiteLLM_Params.model_validate( # rebind-ok: ported from the Python core unchanged + litellm_params_data + ) + else: + raise ValueError( + f"Deployment missing valid litellm_params. " + f"Got: {type(litellm_params_data).__name__}, " + f"deployment_id: {(deployment.get('model_info') or {}).get('id', 'unknown')}" + ) + + wire: Final = _Wire() + user_model_info: dict[str, object] = {} # mutable-ok: filled when the merge asks for it # rebind-ok: see above + + def read_user_model_info() -> object: + nonlocal user_model_info + raw_user_model_info: Final = deployment.get("model_info") + user_model_info = dict( # rebind-ok: see above + _MODEL_INFO_ADAPTER.validate_python( + raw_user_model_info.model_dump(exclude_none=True) + if isinstance(raw_user_model_info, BaseModel) + else raw_user_model_info or {} + ) + ) + return wire.to_wire(user_model_info) + + def answer(step: Mapping[str, object]) -> object: + if step["lookup"] == "provider": + _model, custom_llm_provider, _, _ = self._host.invoke("get_llm_provider")( + model=litellm_params.model, + litellm_params=litellm_params, + ) + return [_model, custom_llm_provider] + if step["lookup"] == "user_model_info": + return read_user_model_info() + return self._answer_model_info_lookup(step, wire) + + query: Final = { + "deployment_id": wire.to_wire((deployment.get("model_info") or {}).get("id")), + "model_info_base_model": wire.to_wire(model_info_base_model), + "params_base_model": wire.to_wire(params_base_model), + "params_model": litellm_params.model, + "extra_base_model": wire.to_wire((litellm_params.model_extra or {}).get("base_model")), + "received_model_name": received_model_name, + } + done: Final = _drive_model_info(self._core_handle().router_model_info(query), answer) + merged: Final = wire.from_wire(done["info"]) + if not isinstance(merged, dict): + return merged # pyright: ignore[reportReturnType] # the core returns get_model_info's answer as is + # The deployment's own values are its objects, as the Python router merges them. + for key in cast(list[str], done["user_keys"]): # cast-ok: router-core's key list + merged[key] = user_model_info[key] + return merged # pyright: ignore[reportReturnType] # the merged map, as the core builds it + + def _core_get_deployment_model_info(self, model_id: str, model_name: str) -> ModelInfo | None: + """ + For a given model id, return the model info + + 1. Check if model_id is in model info + 2. If not, check if litellm model name is in model info + 3. If not, return None + + router-core merges it (``router-core/src/model_info.rs``); this answers its lookups. + """ + wire: Final = _Wire() + call: Final = self._core_handle().deployment_model_info(model_id, wire.to_wire(model_name)) + done: Final = _drive_model_info(call, partial(self._answer_model_info_lookup, wire=wire)) + return wire.from_wire(done["info"]) # pyright: ignore[reportReturnType] # the merged map, as the core builds it + + def _answer_model_info_lookup(self, step: Mapping[str, object], wire: "_Wire") -> object: + """One cost-map lookup of a model-info merge, answered from this process.""" + match step["lookup"]: + case "model_info": + return self._model_info_row(wire.from_wire(step["model"]), wire, step) + case "cost_entry": + key: Final = str(step["key"]) + model_cost: Final = litellm.model_cost if step["process"] else self._host.dynamic().model_cost + return {"entry": wire.to_wire(model_cost.get(key))} if key in model_cost else None + case "pattern_models": + return self._pattern_models(str(step["name"]), wire) + case "bedrock_base": + return get_bedrock_base_model(str(step["model_id"])) + case "discovered": + return wire.to_wire(dict(self._core_get_discovered_model_info(wire.from_wire(step["id"])))) # pyright: ignore[reportArgumentType] # the deployment's id, as configured + case "log": + log: Final = verbose_router_logger.error if step["level"] == "error" else verbose_router_logger.debug + log("%s", step["message"]) + return None + case other: + raise ValueError(f"router-core asked an unknown model-info lookup {other!r}") + + def _model_info_row(self, model: object, wire: "_Wire", step: Mapping[str, object]) -> object: + """``get_model_info(model=...)`` as a lookup answers it: the row, ``{"failed": True}`` when a + tolerated lookup raises, and for an exact lookup whether a generalization rule made it.""" + # The base model's own window is read through the process host, as the Python router reads it. + host: Final = python_host if step["exact"] else self._host + try: + row: Final = host.invoke("get_model_info")(model=model) + except Exception: # noqa: BLE001 # get_model_info raises ModelNotMappedError or a bare Exception for an unmapped model + if step["tolerate"]: + return {"failed": True} + raise + if step["exact"]: + return {"row": wire.to_wire(row), "generalized": is_generalized_model_info(row)} + return wire.to_wire(row) + + def _pattern_models(self, name: str, wire: "_Wire") -> object: + """The ``{id, model}`` of each deployment a wildcard pattern routes ``name`` to.""" + potential_models: Final = self._core_pattern_router.route(name) + if potential_models is None: + return None + routed: Final[list[object]] = [] # mutable-ok: built here + for potential_model in potential_models: + try: + routed.append( + { + "id": wire.to_wire((potential_model.get("model_info") or {}).get("id")), + "model": wire.to_wire((potential_model.get("litellm_params") or {}).get("model")), + } + ) + except Exception: # noqa: BLE001 # ported as the core writes it + verbose_router_logger.debug("model info lookup failed; ignored as the core ignores it") + return routed + + def _core_get_discovered_model_info(self, model_id: str | None) -> Mapping[str, int]: + cached: Final[object] = self._discovered_model_info_cache.get_cache(model_id) + if ( + model_id is not None + and isinstance(cached, DiscoveredDeploymentModelInfo) + and cached.deployment is self._core_get_model_info(model_id) + ): + configured: Final = TypeAdapter(Mapping[str, object]).validate_python(cached.deployment["model_info"]) + return MappingProxyType({key: value for key, value in cached.limits.items() if configured.get(key) is None}) + return MappingProxyType({}) + + def _core_get_model_listing_info(self, model_name: str) -> DeploymentModelListingInfo | None: + """ + Return what the concrete deployments behind model_name contribute to its + /v1/models entry: the cost-map keys for their underlying models, plus the widest + configured or discovered token limits. Resolved via O(1) index + lookup. + + Returns None for wildcard-expanded or unknown names, where the listed name is the + real model name and no deployment-specific information exists, and treats a + malformed configured limit as absent rather than failing the listing. + + The whole group is read rather than just its first deployment, so a group that + mixes models does not advertise a window that depends on config order; the widest + one is reported, which is what get_model_group_info already shows the Admin UI. + Keys are deduplicated, so the ordinary group of interchangeable deployments of one + model still costs the caller a single cost-map lookup. Unlike get_model_group_info, + this never triggers pattern matching or deep copies, so it is safe to call per + listed model on the /v1/models hot path. + + router-core combines the deployments (``router-core/src/group_info.rs``). + """ + deployments: Final = self._core_deployments_named(model_name) + if not deployments: + return None + wire: Final = _Wire() + call: Final = self._core_handle().model_listing_info([wire.to_wire(deployment) for deployment in deployments]) + done: Final = _drive_model_info(call, partial(self._answer_model_info_lookup, wire=wire)) + listing: Final = as_str_mapping(done["value"]) or {} + return DeploymentModelListingInfo( + cost_map_keys=tuple(cast(list[str], listing["cost_map_keys"])), # cast-ok: router-core's key list + max_input_tokens=cast("int | None", listing["max_input_tokens"]), # cast-ok: router-core's limit + max_output_tokens=cast("int | None", listing["max_output_tokens"]), # cast-ok: router-core's limit + ) + + def _core_get_configured_token_limits(self, model_name: str) -> "tuple[int | None, int | None]": + """ + Return (max_input_tokens, max_output_tokens) configured or discovered for a concrete + deployment of model_name, via O(1) index lookup. + + Returns (None, None) for wildcard-expanded or unknown names, and treats a + malformed configured value as absent rather than failing the caller. + + Deliberately reads one deployment rather than aggregating the group the way + get_model_listing_info does: its caller truncates an embedding input to this + value, so the widest window in a mixed group would be the wrong answer there. + """ + deployment: Final = self._core_get_deployment_by_model_group_name(model_group_name=model_name) + if deployment is None: + return (None, None) + wire: Final = _Wire() + configured: Final = { + key: wire.to_wire(value) for key, value in deployment.model_info.model_dump(exclude_none=True).items() + } + call: Final = self._core_handle().configured_token_limits(wire.to_wire(deployment.model_info.id), configured) + done: Final = _drive_model_info(call, partial(self._answer_model_info_lookup, wire=wire)) + limits: Final = cast("list[int | None]", done["value"]) # cast-ok: router-core's two limits + return (limits[0], limits[1]) + + def _core_get_configured_mode(self, model_name: str) -> "str | None": + """Return the mode explicitly configured for a concrete deployment.""" + deployment: Final = self._core_get_deployment_by_model_group_name(model_group_name=model_name) + if deployment is None: + return None + return self._core_handle().configured_text(_Wire().to_wire(deployment.model_info.get("mode"))) + + def _core_get_configured_display_name(self, model_name: str) -> "str | None": + """ + Return the display_name explicitly configured in a concrete deployment's + model_info for model_name, via O(1) index lookup. + + Returns None for wildcard-expanded or unknown names, and treats a + non-string or empty configured value as absent rather than failing the + listing. Like get_configured_token_limits, this never triggers pattern + matching or deep copies, so it is safe to call per listed model on the + /v1/models hot path. + """ + deployment: Final = self._core_get_deployment_by_model_group_name(model_group_name=model_name) + if deployment is None: + return None + return self._core_handle().configured_text(_Wire().to_wire(deployment.model_info.get("display_name"))) + + def _core_get_configured_service_tiers(self, model_name: str, team_id: str | None = None) -> tuple[object, ...]: + """ + Return the service_tiers value each deployment a request from team_id + to model_name can route to (see _routable_deployments) configures in + its model_info, unvalidated and in model_list order, None for a + deployment that sets none; the caller validates the shape it expects + and decides how the deployments combine. + + Returns an empty tuple for wildcard-expanded or unknown names. + """ + wire: Final = _Wire() + tiers: Final = self._core_handle().configured_service_tiers( + [ + wire.to_wire(dict(_configured_model_info(deployment))) + for deployment in self._routable_deployments(model_name, team_id) + ] + ) + return tuple(wire.from_wire(tier) for tier in tiers) + + def _core_get_configured_max_input_tokens( + self, model_name: str, team_id: str | None = None + ) -> tuple[int | None, ...]: + """ + Return the max_input_tokens each deployment a request from team_id to + model_name can route to (see _routable_deployments) sets in its own + model_info, in model_list order, None for a deployment that sets none + or a malformed one. Unlike get_model_listing_info, neither the cost map + nor a discovered limit is read, so only an operator's own value counts, + and a value a saved /model/info response copied from the cost map is + skipped too. + + Returns an empty tuple for wildcard-expanded or unknown names. + """ + return tuple( + _operator_max_input_tokens(deployment) for deployment in self._routable_deployments(model_name, team_id) + ) + + async def _core_get_model_group_usage(self, model_group: str) -> tuple[int | None, int | None]: + """ + Returns current tpm/rpm usage for model group + + Parameters: + - model_group: str - the received model name from the user (can be a wildcard route). + + Returns: + - usage: Tuple[tpm, rpm] + """ + dt: Final = get_utc_datetime() + current_minute: Final = dt.strftime("%H-%M") # use the same timezone regardless of system clock + tpm_keys: Final[list[str]] = [] # mutable-ok: ported from the Python core unchanged + rpm_keys: Final[list[str]] = [] # mutable-ok: ported from the Python core unchanged + + model_list: Final = self._core_get_model_list(model_name=model_group) + if model_list is None: # no matching deployments + return None, None + + for model in model_list: + id: str | None = model.get("model_info", {}).get("id") + litellm_model: str | None = model["litellm_params"].get( + "model" + ) # USE THE MODEL SENT TO litellm.completion() - consistent with how global_router cache is written. + if id is None or litellm_model is None: + continue + tpm_keys.append(self._core_handle().usage_key(id, litellm_model, "tpm", current_minute)) + rpm_keys.append(self._core_handle().usage_key(id, litellm_model, "rpm", current_minute)) + combined_tpm_rpm_keys: Final = tpm_keys + rpm_keys + + with service_target(ROUTER_USAGE_TARGET): + combined_tpm_rpm_values: Final = await self._core_cache.async_batch_get_cache(keys=combined_tpm_rpm_keys) + if combined_tpm_rpm_values is None: + return None, None + + rpm_start: Final = len(tpm_keys) + tpm_usage_list: Final[list | None] = combined_tpm_rpm_values[:rpm_start] # mutable-ok: ported from the Python core unchanged # fmt: skip + rpm_usage_list: Final[list | None] = combined_tpm_rpm_values[rpm_start:] # mutable-ok: ported from the Python core unchanged # fmt: skip + + ## TPM + tpm_usage: int | None = None # rebind-ok: ported from the Python core unchanged + if tpm_usage_list is not None: + for t in tpm_usage_list: + if isinstance(t, int): + if tpm_usage is None: + tpm_usage = 0 + tpm_usage += t + ## RPM + rpm_usage: int | None = None # rebind-ok: ported from the Python core unchanged + if rpm_usage_list is not None: + for t in rpm_usage_list: + if isinstance(t, int): + if rpm_usage is None: + rpm_usage = 0 + rpm_usage += t + return tpm_usage, rpm_usage + + async def get_model_group_io_token_usage(self, model_group: str) -> tuple[int | None, int | None]: + """ + Returns current ITPM/OTPM usage for a model group (sum across deployments). + """ + dt: Final = get_utc_datetime() + current_minute: Final = dt.strftime("%H-%M") + itpm_keys: Final[list[str]] = [] # mutable-ok: ported from the Python core unchanged + otpm_keys: Final[list[str]] = [] # mutable-ok: ported from the Python core unchanged + + model_list: Final = self._core_get_model_list(model_name=model_group) + if model_list is None: + return None, None + + for model in model_list: + model_id: str | None = model.get("model_info", {}).get("id") + litellm_model: str | None = model["litellm_params"].get("model") + if model_id is None or litellm_model is None: + continue + itpm_keys.append(self._core_handle().usage_key(model_id, litellm_model, "itpm", current_minute)) + otpm_keys.append(self._core_handle().usage_key(model_id, litellm_model, "otpm", current_minute)) + + with service_target(ROUTER_USAGE_TARGET): + combined_values: Final = await self._core_cache.async_batch_get_cache(keys=itpm_keys + otpm_keys) + if combined_values is None: + return None, None + + itpm_values: Final = combined_values[: len(itpm_keys)] + otpm_values: Final = combined_values[len(itpm_keys) :] + + total_itpm: int | None = None # rebind-ok: ported from the Python core unchanged + for value in itpm_values: + if isinstance(value, int): + total_itpm = (total_itpm or 0) + value + + total_otpm: int | None = None # rebind-ok: ported from the Python core unchanged + for value in otpm_values: + if isinstance(value, int): + total_otpm = (total_otpm or 0) + value + + return total_itpm, total_otpm + + async def _core_get_remaining_model_group_usage( + self, model_group: str + ) -> dict[str, int]: # mutable-ok: ported from the Python core unchanged + model_group_info: Final = self._cached_get_model_group_info(model_group) + + returned_dict: Final[dict[str, int]] = {} # mutable-ok: ported from the Python core unchanged + + # ITPM/OTPM groups emit input/output token headers, but they may also set + # tpm/rpm, so build both sets rather than returning early - clients and + # prometheus gauges that read the standard headers still get data. + if model_group_info is not None and (model_group_info.itpm is not None or model_group_info.otpm is not None): + current_itpm, current_otpm = await self.get_model_group_io_token_usage(model_group) + returned_dict.update( + build_io_token_rate_limit_headers( + itpm_limit=model_group_info.itpm, + otpm_limit=model_group_info.otpm, + current_itpm=current_itpm, + current_otpm=current_otpm, + ) + ) + + tpm_limit: Final = model_group_info.tpm if model_group_info is not None else None + rpm_limit: Final = model_group_info.rpm if model_group_info is not None else None + + if tpm_limit is None and rpm_limit is None: + return returned_dict + + current_tpm, current_rpm = await self._core_get_model_group_usage(model_group) + + if tpm_limit is not None: + returned_dict["x-ratelimit-remaining-tokens"] = tpm_limit - (current_tpm or 0) + returned_dict["x-ratelimit-limit-tokens"] = tpm_limit + if rpm_limit is not None: + returned_dict["x-ratelimit-remaining-requests"] = rpm_limit - (current_rpm or 0) + returned_dict["x-ratelimit-limit-requests"] = rpm_limit + + return returned_dict + + async def set_response_headers( + self, + response: object, + model_group: str | None = None, + request_kwargs: dict[str, object] | None = None, # mutable-ok: ported from the Python core unchanged + ) -> Any: # noqa: ANN401 # the core's signature, ported unchanged + """ + Add the most accurate rate limit headers for a given model response. + + ## TODO: add model group rate limit headers + # - if healthy_deployments > 1, return model group rate limit headers + # - else return the model's rate limit headers + """ + response = prepare_response_for_header_attachment(response) # rebind-ok: ported from the Python core unchanged + if response is None: + return response + + additional_headers: Final = ensure_response_additional_headers(response) + apply_response_model_id( + response, find_deployment_metadata(request_kwargs) if request_kwargs is not None else None + ) + additional_headers["x-litellm-model-group"] = model_group + apply_quality_router_decision_headers(additional_headers, request_kwargs) + additional_headers.update(complexity_router_decision_headers(request_kwargs)) + + if model_group is not None: + remaining_usage: Final = await self._core_get_remaining_model_group_usage(model_group) + apply_remaining_usage_headers(additional_headers, remaining_usage) + return response diff --git a/litellm/router_rust/surface/model_list.py b/litellm/router_rust/surface/model_list.py new file mode 100644 index 00000000000..405a4b47d7b --- /dev/null +++ b/litellm/router_rust/surface/model_list.py @@ -0,0 +1,375 @@ +""" +``HostedModelList``: the Rust core's live ``model_list``, read back from router-core. + +router-core holds the deployments. This list is what ``registry().all()`` and the deprecated +``model_list`` attribute hand out: a ``list`` (so callers that check for one, serialize it or +validate it keep working) whose contents are router-core's deployment records, read through the +handle after every write. Every write to it is a changeset instead of a Python-side mutation: + +- ``append`` / ``extend`` / an ``insert`` at the end add at the end (``append``) +- removing one deployment whose id no other entry shares is a ``delete`` +- any other write (an assignment, a slice, a removal of one of several entries sharing an id, + ``clear``, ``sort``, ...) hands router-core the resulting list (``replace``) + +A build serves the deployments it added once it ends (or fails), and the list reads them back +then. After writes that only appended, it reads just the appended records +(``RouterCoreHandle.deployments_from``), so adding deployments one at a time costs each one its own +record rather than the whole list. These writes are the raw ones a caller makes: like the Python +router's list, they reindex nothing. + +An entry read back from router-core is a dict whose writes (its own, or a nested dict's) hand +router-core the list again, so an edit in place reaches it as the Python router's list would show +it. Each deployment keeps one Python object while router-core's record for it is unchanged, keyed by +its id and its place among entries sharing that id, so identity checks on stored entries (the +model info caches) mean what they did: the same object until a write changes that deployment. +""" + +import copy +import datetime +from collections.abc import Callable, Iterable, Mapping, Sequence +from types import MappingProxyType +from typing import ( + TYPE_CHECKING, + Final, + SupportsIndex, + TypeAlias, + TypeVar, + cast, # noqa: TID251 # any value a caller wrote into the list + overload, +) + +from pydantic import TypeAdapter, ValidationError + +from litellm.litellm_core_utils.core_helpers import as_str_mapping +from litellm.types.router import ModelAccessWindow + +if TYPE_CHECKING: + from litellm.rust_bridge._native import RouterCoreHandle + +Entry: TypeAlias = dict[str, object] +# (id, its place among the entries sharing it) -> (the record it was read as, the entry object) +_Stable: TypeAlias = dict[tuple[str, int], tuple[object, Entry]] + + +_DATETIME: Final = TypeAdapter(datetime.datetime) +_Field = TypeVar("_Field") # rebind-ok: a TypeVar cannot be Final + + +def _dumped_as(adapter: TypeAdapter[_Field]) -> Callable[[object], object]: + """A value validated by ``adapter``, then dumped back to its Python form.""" + + def dump(value: object) -> object: + return adapter.dump_python(adapter.validate_python(value), exclude_none=True) + + return dump + + +# The ``ModelInfo`` fields whose Python value is no JSON value, which router-core holds as strings. +_TYPED_MODEL_INFO: Final[Mapping[str, Callable[[object], object]]] = MappingProxyType( + { + "created_at": _dumped_as(_DATETIME), + "updated_at": _dumped_as(_DATETIME), + "ptu_effective_from": _dumped_as(_DATETIME), + "ptu_effective_to": _dumped_as(_DATETIME), + "access_windows": _dumped_as(TypeAdapter(tuple[ModelAccessWindow, ...])), + } +) + + +def typed_deployment(record: Mapping[str, object]) -> Mapping[str, object]: + """A deployment record as ``Deployment.to_json`` dumped it: its ``model_info`` typed again. One + with none of those fields comes back as it is, and so does a value a raw write left unreadable.""" + info: Final = as_str_mapping(record.get("model_info")) + if info is None or info.keys().isdisjoint(_TYPED_MODEL_INFO): + return record + typed: Final[dict[str, object]] = {} # mutable-ok: filled here + for key, dumped in _TYPED_MODEL_INFO.items(): + value = info.get(key) + if value is None: + continue + try: + typed[key] = dumped(value) + except ValidationError: + continue + return {**record, "model_info": {**info, **typed}} + + +def _record(value: object) -> Mapping[str, object]: + return cast(Mapping[str, object], value) # cast-ok: a JSON object router-core wrote, or a dict inside one + + +def _entry_id(entry: object) -> str | None: + info: Final = as_str_mapping((as_str_mapping(entry) or {}).get("model_info")) or {} + value: Final = info.get("id") + return None if value is None else str(value) + + +def _keys( + entries: Iterable[object], seen: Mapping[str, int] = MappingProxyType({}) +) -> tuple[list[tuple[str, int] | None], Mapping[str, int]]: # mutable-ok: a fresh list + """Each entry's identity key (its id and how many earlier entries share it, ``seen`` counting + the ids of entries before these), and the counts once past them.""" + counts: Final[dict[str, int]] = dict(seen) # mutable-ok: counted here + keys: Final[list[tuple[str, int] | None]] = [] # mutable-ok: filled here + for entry in entries: + entry_id = _entry_id(entry) + if entry_id is None: + keys.append(None) + continue + nth = counts.get(entry_id, 0) + counts[entry_id] = nth + 1 + keys.append((entry_id, nth)) + return keys, counts + + +class _HostedEntry(dict[str, object]): + """A deployment entry read back from router-core: a write to it, or to a dict inside it, hands + router-core the list again (``replace``), so it learns what the entry now says. A copy is a + plain dict.""" + + def __init__(self, record: Mapping[str, object], written: Callable[[], None]) -> None: + super().__init__( + (key, value if (nested := as_str_mapping(value)) is None else _HostedEntry(nested, written)) + for key, value in record.items() + ) + self._written: Final = written + + def __setitem__(self, key: str, value: object) -> None: + super().__setitem__(key, value) + self._written() + + def __delitem__(self, key: str) -> None: + super().__delitem__(key) + self._written() + + def pop(self, key: str, *default: object) -> object: + value: Final = super().pop(key, *default) + self._written() + return value + + def popitem(self) -> tuple[str, object]: + item: Final = super().popitem() + self._written() + return item + + def clear(self) -> None: + super().clear() + self._written() + + def update(self, *args: object, **kwargs: object) -> None: # kwargs-ok: dict.update's own arguments + super().update(*args, **kwargs) # pyright: ignore[reportCallIssue, reportArgumentType] # dict.update's own arguments + self._written() + + def setdefault(self, key: str, default: object = None) -> object: + if key in self: + return self[key] + self[key] = default + return default + + def __ior__(self, other: object) -> "_HostedEntry": # pyright: ignore[reportIncompatibleMethodOverride] # dict's in-place merge, as a write + self.update(other) + return self + + def __copy__(self) -> dict[str, object]: # mutable-ok: a plain copy + return dict(self) + + def __deepcopy__(self, memo: dict[int, object]) -> dict[str, object]: # mutable-ok: a plain copy + return copy.deepcopy(dict(self), memo) + + def __reduce_ex__(self, protocol: SupportsIndex) -> tuple[object, ...]: + return (dict, (copy.deepcopy(dict(self)),)) + + +class HostedModelList(list[Entry]): + """The live ``model_list``: router-core's deployments, every write a changeset.""" + + def __init__( + self, + handle: Callable[[], "RouterCoreHandle | None"], + apply: Callable[[Mapping[str, object]], None], + ) -> None: + super().__init__() + self._handle: Final = handle + self._apply: Final = apply + # The handle and registry version the contents were last read at. + self._read_at: tuple[object, int] | None = None + self._stable: _Stable = {} + # How many of the shown entries hold each id, for the keys of entries appended next. + self._id_counts: Mapping[str, int] = {} + # While a build runs: the length, id counts and read to go back to when it ends. + self._building: tuple[int, Mapping[str, int], tuple[object, int] | None] | None = None + + # -- reading ---------------------------------------------------------------------------------- + + def refresh(self) -> None: + """Read router-core's deployments again when its registry moved since the last read.""" + handle: Final = self._handle() + if handle is None or self._building is not None: + # While a build runs, the list shows what it accepted until router-core serves it. + return + version: Final = handle.registry_version() + if self._read_at is not None and self._read_at[0] is handle and self._read_at[1] == version: + return + read_at: Final = self._read_at + self._read_at = (handle, version) + if read_at is not None and read_at[0] is handle and handle.registry_rewritten() <= read_at[1]: + # Only appends since the last read: the shown entries stand, and the new ones follow. + super().extend(self._kept(handle.deployments_from(len(self)), appended=True)) + return + self._show(self._kept(handle.deployments())) + + def _kept(self, records: Sequence[object], *, appended: bool = False) -> list[Entry]: # mutable-ok: a fresh list + """``records`` as entries, each the object it was while its record is unchanged; when + ``appended``, they follow the shown entries, which keep theirs.""" + stable: Final[_Stable] = self._stable if appended else {} + keys, counts = _keys(records, self._id_counts if appended else {}) + entries: Final[list[Entry]] = [] # mutable-ok: filled here + for key, record in zip(keys, records, strict=True): + typed = typed_deployment(_record(record)) + held = self._stable.get(key) if key is not None else None + kept = held is not None and held[0] == typed + entry = held[1] if held is not None and kept else _HostedEntry(typed, self._entry_written) + if key is not None: + stable[key] = (held[0] if held is not None and kept else copy.deepcopy(typed), entry) + entries.append(entry) + self._stable = stable + self._id_counts = counts + return entries + + # -- a build ---------------------------------------------------------------------------------- + + def begin_build(self, *, replacing: bool) -> None: + """A build starts: router-core serves what it adds once it ends, so until then the list + shows each deployment it accepts here (``built``), after the entries it keeps, or alone + when it replaces them, as the Python router's list grows while it adds them.""" + self.refresh() + if replacing: + self._show([]) + self._id_counts = {} + self._building = (len(self), self._id_counts, self._read_at) + + def built(self, entry: Mapping[str, object]) -> None: + """A deployment the build accepted, shown until the build ends.""" + if self._building is not None: + super().extend(self._kept([entry], appended=True)) + + def end_build(self) -> None: + """The build ended (or failed): router-core serves what it added, read back in place of + the entries shown meanwhile, each kept as the object it was shown as while unchanged.""" + if self._building is None: + return + start, counts, read_at = self._building + self._building = None + super().__delitem__(slice(start, None)) + self._id_counts = counts + self._read_at = read_at + self.refresh() + + def _seed(self, entries: Sequence[object]) -> None: + """Keep ``entries``, the list router-core now holds, as the objects its records for them + read back as.""" + for key, entry in zip(_keys(entries)[0], entries, strict=True): + self._seed_one(key, entry) + + def _seed_appended(self, entries: Sequence[object]) -> None: + """Keep ``entries``, appended after the shown ones, as the objects they read back as.""" + for key, entry in zip(_keys(entries, self._id_counts)[0], entries, strict=True): + self._seed_one(key, entry) + + def _seed_one(self, key: tuple[str, int] | None, entry: object) -> None: + if key is not None and isinstance(entry, dict): + stored: Final = cast(Entry, entry) # cast-ok: a model_list entry is a dict + self._stable[key] = (copy.deepcopy(stored), stored) + + def _entry_written(self) -> None: + """A write inside an entry: router-core takes the list as it now reads.""" + self._replace(list(self)) + + def _show(self, entries: list[Entry]) -> None: # mutable-ok: the entries shown + super().__setitem__(slice(None), entries) + + # -- writing ---------------------------------------------------------------------------------- + + def _replace(self, entries: list[Entry]) -> None: # mutable-ok: the list router-core takes + # Seeded first: router-core's records for them read back as these objects however soon + # the list is read again. + self._seed(entries) + self._apply({"replace": entries}) + self.refresh() + + def replace(self, entries: Iterable[object]) -> None: + """Every deployment at once (``router.model_list = [...]``).""" + self._replace(cast(list[Entry], list(entries))) # cast-ok: the entries as assigned, which router-core validates + + def append(self, value: Entry) -> None: + self.refresh() + self._seed_appended([value]) + self._apply({"append": [value]}) + self.refresh() + + def extend(self, values: Iterable[Entry]) -> None: + for value in tuple(values): + self.append(value) + + def __iadd__(self, values: Iterable[Entry]) -> "HostedModelList": # pyright: ignore[reportIncompatibleMethodOverride] # list's in-place add, as a changeset + self.extend(values) + return self + + def insert(self, index: SupportsIndex, value: Entry) -> None: + if int(index) >= len(self): + self.append(value) + return + changed: Final = list(self) + changed.insert(index, value) + self._replace(changed) + + @overload + def __setitem__(self, index: SupportsIndex, value: Entry) -> None: ... + @overload + def __setitem__(self, index: slice, value: Iterable[Entry]) -> None: ... + def __setitem__(self, index: SupportsIndex | slice, value: object) -> None: + changed: Final = list(self) + changed[index] = value # pyright: ignore[reportArgumentType, reportCallIssue] # list's own item and slice assignment + self._replace(changed) + + def __delitem__(self, index: SupportsIndex | slice) -> None: + changed: Final = list(self) + if isinstance(index, slice): + del changed[index] + self._replace(changed) + return + removed: Final = changed.pop(index) + removed_id: Final = _entry_id(removed) + if removed_id is None or any(_entry_id(entry) == removed_id for entry in changed): + # A delete removes every deployment with the id; the list removes only this one. + self._replace(changed) + return + self._apply({"delete": [removed_id]}) + self.refresh() + + def pop(self, index: SupportsIndex = -1) -> Entry: + entry: Final = self[index] + del self[index] + return entry + + def remove(self, value: Entry) -> None: + del self[self.index(value)] + + def clear(self) -> None: + self._replace([]) + + def sort(self, *args: object, **kwargs: object) -> None: # kwargs-ok: list.sort's own arguments + changed: Final = list(self) + changed.sort(*args, **kwargs) # pyright: ignore[reportCallIssue, reportArgumentType] # list.sort's own arguments + self._replace(changed) + + def reverse(self) -> None: + self._replace(list(reversed(self))) + + def __imul__(self, count: SupportsIndex) -> "HostedModelList": # pyright: ignore[reportIncompatibleMethodOverride] # list's in-place repeat, as a changeset + self._replace(list(self) * int(count)) + return self + + def __reduce_ex__(self, protocol: SupportsIndex) -> tuple[object, ...]: + # A copy or pickle of the live list is a plain list of its entries. + return (list, (list(self),)) diff --git a/litellm/router_rust/surface/registry.py b/litellm/router_rust/surface/registry.py new file mode 100644 index 00000000000..a55a1cfc6d6 --- /dev/null +++ b/litellm/router_rust/surface/registry.py @@ -0,0 +1,2191 @@ +# pyright: reportUnknownMemberType=false, reportUnknownVariableType=false, reportUnknownArgumentType=false, reportAny=false +# (the ported core code reads untyped litellm helpers: PatternMatchRouter, Deployment.to_json, host.invoke) +""" +``HostRegistry``: the Rust core's live ``model_list`` and the indexes ``registry()`` reads. + +router-core builds the deployments (``litellm-rust/crates/router-core/src/build.rs``): it walks +the list, decides which values are ``os.environ/`` secret references, assigns the id of a +deployment that declares none and adds one deployment per ``organization``. What needs live Python +objects is a host step this registry answers, one deployment at a time (``_answer_build_step``): +``secrets`` (``get_secret``), ``validate`` (the PTU, access-window and strict param checks, and the +pricing each deployment writes into ``litellm.model_cost``), ``provider`` (``get_llm_provider``) +and ``added`` (pass-through credentials, pre-routing strategies, the patterns and indexes below). + +router-core holds the deployments. ``_core_model_list`` is ``HostedModelList`` +(litellm/router_rust/surface/model_list.py): router-core's records read back through the handle, every +write a changeset (a build serves what it added once it ends, a removal is a delete, any other +write a replace), each deployment one Python object while its record is unchanged. + +The indexes are router-core's too (``host_index.rs``), kept by the Python router's history rules: +the id, group and team position maps, ``model_names``, the only-growing ``deployment_names``, +the team public names, the provider default ids and the wildcard pattern routers. The registry +hands them out as the live views of litellm/router_rust/surface/indexes.py, whose writes are changesets. + +The lookups run in router-core (``lookup.rs``): the deployments by name (team public name first), +the model list with its alias, routing group and wildcard rows, the ids, names and access groups +it lists, the candidate ids for a route, the routable deployments, the credential deployment picks +and a team model's alternatives. Each answers with rows, positions in the list: this registry +reads its stored entries at them (``_rows``), so identity and in-place reads stay its own, and a +position router-core could not read (an index a caller wrote junk into) is read here and raises +as the Python router's read raised. What stays here needs live Python objects: the alias map and +routing groups a lookup is scoped by (``_list_scope``), the provider a wildcard request is tried +under, pydantic ``Deployment`` values, the credentials, the PTU ceilings, search tool support, a +team model's provider pick (``provider_for_generic_call``) and the SDK clients. + +Model info and group info are router-core's rules too (``model_info.rs``, ``group_info.rs``), +answered from this process through ``HostModelInfo`` (litellm/router_rust/surface/model_info.py). + +The attribute and method names follow the router core's (``_core_model_list``, +``_core_get_model_list`` ...), so ``litellm.router_api.RouterRegistry`` reads this object as its +router and every ``registry()`` query dispatches the way it does on the Python core. +""" + +import copy +import dataclasses +import os +import re +import threading +from collections.abc import Callable, Iterable, Iterator, Mapping, Sequence +from functools import lru_cache, partial +from types import MappingProxyType +from typing import TYPE_CHECKING, Final, Literal, TypeAlias, cast # noqa: TID251 # entries as a caller wrote them + +from pydantic import TypeAdapter + +from litellm import exceptions as litellm_exceptions +from litellm import get_secret_str +from litellm._internal_context import service_target +from litellm._logging import verbose_router_logger +from litellm.caching.dual_cache import DualCache +from litellm.caching.in_memory_cache import InMemoryCache +from litellm.constants import DEFAULT_MAX_LRU_CACHE_SIZE +from litellm.litellm_core_utils.access_windows import access_windows_config_error +from litellm.litellm_core_utils.core_helpers import as_str_mapping, without +from litellm.litellm_core_utils.credential_accessor import CredentialAccessor +from litellm.litellm_core_utils.get_llm_provider_logic import ( + declared_authenticating_provider, + is_registered_custom_provider, + resolve_model_and_provider_for_metadata, +) +from litellm.litellm_core_utils.llm_cost_calc.utils import SERVICE_TIER_COST_KEY_SUFFIXES +from litellm.litellm_core_utils.pattern_match_deployments import PatternMatchRouter +from litellm.litellm_core_utils.ptu_pricing import ( + PTU_COST_ATTRIBUTION_ENV_VAR, + declares_ptu, + declares_ptu_shares, + is_ptu_cost_attribution_enabled, + ptu_config_error, + ptu_identity_error, + ptu_terms, + zeroed_ptu_pricing, +) +from litellm.litellm_core_utils.ptu_shares import ptu_capacity_warning +from litellm.litellm_core_utils.routing_common_utils import ( + provider_for_generic_call, + resolve_model_group_alias, + team_may_use_deployment, + warn_on_ignored_model_capability_flags, + warn_on_provider_credential_mismatch, + with_admin_api_key_secret_reference, + without_model_capability_flags, +) +from litellm.litellm_core_utils.routing_groups import apply_routing_group_priority +from litellm.llms.openai_like.json_loader import JSONProviderRegistry +from litellm.llms.openai_like.model_info import MODEL_INFO_REFRESH_SECONDS +from litellm.responses.bridge_check import keep_registered_mode +from litellm.router_api import RouterRegistry +from litellm.router_rust.invoke.shaping import ROUTER_CLIENTS_TARGET +from litellm.router_rust.surface.indexes import HostedIndexes +from litellm.router_rust.surface.model_info import HostModelInfo +from litellm.router_rust.surface.model_list import HostedModelList +from litellm.router_strategy.auto_router.auto_router_model_naming import classify_strategy_router_model +from litellm.router_utils.deployment_numbers import refuse_unservable_numbers, refuse_wide_numbers +from litellm.router_utils.endpoint_eligibility import warn_on_unenforced_endpoint_flag +from litellm.router_utils.pre_call_checks.io_token_rate_limit_check import deployment_has_io_token_limits +from litellm.types.router import ( + VALID_LITELLM_ENVIRONMENTS, + CredentialLiteLLMParams, + Deployment, + DeploymentTypedDict, + LiteLLM_Params, + ModelGroupInfo, + RouterModelGroupAliasItem, + RoutingGroup, + RoutingPlugin, + TaggedPreRoutingStrategy, + validate_litellm_params_strictly, +) +from litellm.types.utils import ( + DEPLOYMENT_SCOPED_TOKEN_LIMIT_FIELDS, + CustomPricingLiteLLMParams, + LlmProviders, + LlmProvidersSet, + is_server_derived_pricing_key, + shared_backend_model_info, +) +from litellm.utils import ( + ModelInfoLookupMiss, + catalog_with_runtime_registrations, + get_potential_model_names, + get_secret, + lookup_model_info_for_registration, + model_info_with_pricing_overrides, +) + +if TYPE_CHECKING: + from litellm.litellm_core_utils.ptu_shares import PTUTeamCeiling + from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler + from litellm.router_python_host import PythonHost + from litellm.router_rust.steps.callbacks import OptionalChecks + from litellm.router_rust.steps.prerouting import PreRoutingStrategies + from litellm.router_strategy.budget_limiter import RouterBudgetLimiting + from litellm.rust_bridge._native import RouterCoreHandle + +# The PTU share teams and the team ceilings read so far. +_PTUMemo: TypeAlias = tuple[ + frozenset[str], + dict[tuple[str, str], "PTUTeamCeiling | None"], +] +_MODEL_INFO_ADAPTER: Final = TypeAdapter(Mapping[str, object]) +_LITELLM_PARAMS_ADAPTER: Final = TypeAdapter(Mapping[str, object]) + + +_CACHE_PRICING_FIELDS: Final = ( + "cache_creation_input_token_cost", + "cache_creation_input_token_cost_above_1hr", + "cache_creation_input_token_cost_above_200k_tokens", + "cache_read_input_token_cost", + "cache_read_input_token_cost_above_200k_tokens", +) +_OFF_PEAK_BASE_RATE_FIELDS: Final = ( + "input_cost_per_token", + "input_cost_per_second", + "cost_per_second", + "tiered_pricing", +) + +# A stored model_list entry: the live dict the router core keeps. +UntypedEntry: TypeAlias = dict[str, object] +_TeamIndex: TypeAlias = dict[tuple[str, str], list[int]] + +# How a build's deployments join router-core's registry (`build::Install`). +BuildInstall: TypeAlias = Literal["replace", "append", "upsert"] +# Runs one router-core build (`RouterCoreHandle.build`): the entries, the path, whether it resolves +# secrets, how it installs, and the registry's answer to each step. ``RustCore`` supplies it. +BuildRunner: TypeAlias = Callable[ + [Sequence[object], str, bool, BuildInstall, Callable[[Mapping[str, object]], object]], None +] + + +@dataclasses.dataclass(slots=True) +class _BuildEntry: + """One ``set_model_list`` entry as its loop holds it (the router core pops ``model_name``, + ``litellm_params`` and ``model_info`` off the entry and keeps the rest), which router-core's + build steps name by index.""" + + extras: UntypedEntry + model_name: str + litellm_params: UntypedEntry + model_info: UntypedEntry + declared_id: str | None + # `organization` is a list, which router-core adds one deployment for each of. + per_organization: bool = False + # A retried entry keeps whether its id was declared twice in the list it came from. + duplicate: bool | None = None + + +@dataclasses.dataclass(slots=True) +class _Build: + """The state of one router-core build, which ``_answer_build_step`` reads and writes.""" + + path: Literal["create", "add"] + # `create`: the entries (a deep copy, which the build consumes), and each one once router-core + # reaches it. + raw: Sequence[UntypedEntry] = () + # Filled as the build reaches each entry. + entries: dict[int, _BuildEntry] = dataclasses.field(default_factory=dict) # mutable-ok: filled in step order + # `add`: the deployment to add, and the entry it is stored as. + deployment: Deployment | None = None + stored: UntypedEntry | None = None + # The deployment the build is on, its provider and its data-source values. + pending: Deployment | None = None + provider: str | None = None + data_source_values: tuple[tuple[int, str, str], ...] = () + # `add`: the deployment it added. + result: Deployment | None = None + # `create`: the error a deployment raised. The Python core's loop stops there, so every later + # step is answered as invalid and the build ends, then this is raised. + failure: Exception | None = None + + +def model_info_is_active_for_environment(model_info: Mapping[str, object] | None) -> bool: + """The environment-gating rule (``model_info_is_active_for_environment``): a deployment whose + model_info names ``supported_environments`` loads only where LITELLM_ENVIRONMENT is in it.""" + if model_info is None: + return True + supported_environments: Final = model_info.get("supported_environments") + if supported_environments is None: + return True + if not isinstance(supported_environments, list | tuple): + raise ValueError( # noqa: TRY004 # the core raises ValueError here + f"supported_environments must be a list of {VALID_LITELLM_ENVIRONMENTS}. " + f"but set as: {supported_environments} for model_info: {model_info}" + ) + litellm_environment: Final = get_secret_str(secret_name="LITELLM_ENVIRONMENT") + if litellm_environment is None: + raise ValueError("Set 'supported_environments' for model but not 'LITELLM_ENVIRONMENT' set in .env") + if litellm_environment not in VALID_LITELLM_ENVIRONMENTS: + raise ValueError( + f"LITELLM_ENVIRONMENT must be one of {VALID_LITELLM_ENVIRONMENTS}. but set as: {litellm_environment}" + ) + for environment in supported_environments: # pyright: ignore[reportUnknownVariableType] # a configured list + if environment not in VALID_LITELLM_ENVIRONMENTS: + raise ValueError( + f"supported_environments must be one of {VALID_LITELLM_ENVIRONMENTS}. but set as: {environment} " + f"for model_info: {model_info}" + ) + return litellm_environment in supported_environments + + +# -- pricing written into litellm.model_cost ----------------------------------------------------- + + +def _backend_model_info( + host: "PythonHost", backend_model: str, custom_llm_provider: str | None +) -> Mapping[str, object]: + """``get_model_info`` for the backend model; raises like it for an unmapped model.""" + info: Final = host.invoke("get_model_info")(model=backend_model, custom_llm_provider=custom_llm_provider) + return _MODEL_INFO_ADAPTER.validate_python(info) + + +def _backend_cost_entry( + host: "PythonHost", backend_model: str, custom_llm_provider: str | None +) -> Mapping[str, object] | None: + """The backend model's raw ``litellm.model_cost`` entry, or None when it has none.""" + try: + info: Final = _backend_model_info(host, backend_model, custom_llm_provider) + except Exception: # noqa: BLE001 # get_model_info raises plain Exception for an unmapped backend model + return None + entry: Final[object] = host.dynamic().model_cost.get(str(info.get("key") or "")) + return _MODEL_INFO_ADAPTER.validate_python(entry) if isinstance(entry, dict) else None + + +def _inherited_cache_pricing( + host: "PythonHost", model_info: Mapping[str, object], backend_model: str, provider: str | None +) -> Mapping[str, object]: + """The cache prices a custom-priced entry inherits (``_inherit_builtin_cache_pricing``).""" + inherited_fields: Final = ( + _CACHE_PRICING_FIELDS + if model_info.get("cache_read_input_token_cost") is not None + else (*_CACHE_PRICING_FIELDS, "explicit_cache_read_input_token_cost") + ) + missing_fields: Final = tuple(field for field in inherited_fields if model_info.get(field) is None) + if not missing_fields: + return {} + try: + backend_info: Final = _backend_model_info(host, backend_model, provider) + except Exception: # noqa: BLE001 # an unmapped backend model contributes nothing + return {} + return {field: backend_info[field] for field in missing_fields if backend_info.get(field) is not None} + + +def _inherited_service_tier_pricing( + host: "PythonHost", model_info: Mapping[str, object], backend_model: str, provider: str | None +) -> Mapping[str, object]: + """The service-tier rates a custom-priced entry inherits (``_inherit_builtin_service_tier_pricing``).""" + if ptu_terms(model_info) is not None and is_ptu_cost_attribution_enabled(): + return {} + if all(model_info.get(field) is None for field in ("input_cost_per_token", "output_cost_per_token")): + return {} + backend_entry: Final = _backend_cost_entry(host, backend_model, provider) + if backend_entry is None: + return {} + return { + field: copy.deepcopy(value) + for field, value in backend_entry.items() + if field.endswith(SERVICE_TIER_COST_KEY_SUFFIXES) and model_info.get(field) is None and value is not None + } + + +def _inherited_off_peak_base_rates( + host: "PythonHost", model_info: Mapping[str, object], backend_model: str, provider: str | None +) -> Mapping[str, object]: + """The whole backend entry an off-peak-only entry copies (``_inherit_builtin_base_rates_for_off_peak``).""" + if not model_info.get("off_peak_pricing"): + return {} + if any(model_info.get(field) is not None for field in _OFF_PEAK_BASE_RATE_FIELDS): + return {} + backend_entry: Final = _backend_cost_entry(host, backend_model, provider) + if backend_entry is None: + return {} + return { + field: copy.deepcopy(value) + for field, value in backend_entry.items() + if model_info.get(field) is None and value is not None + } + + +def _inherited_tiered_output_rate( + host: "PythonHost", model_info: Mapping[str, object], backend_model: str, provider: str | None +) -> Mapping[str, object]: + """The flat output rate a tier table without one inherits (``_inherit_builtin_tiered_output_rate``).""" + tiers: Final = model_info.get("tiered_pricing") + if not isinstance(tiers, list) or not tiers: + return {} + if model_info.get("output_cost_per_token") is not None: + return {} + if all(isinstance(tier, dict) and "output_cost_per_token" in tier for tier in tiers): # pyright: ignore[reportUnknownVariableType] # a configured tier table + return {} + try: + backend_info: Final = _backend_model_info(host, backend_model, provider) + except Exception: # noqa: BLE001 # an unmapped backend model contributes nothing + return {} + backend_rate: Final = backend_info.get("output_cost_per_token") + return {"output_cost_per_token": backend_rate} if backend_rate else {} + + +def _bedrock_profile(params: Mapping[str, object], model: str) -> bool: + return params.get("custom_llm_provider") == "bedrock" or model.startswith("bedrock/") + + +def _search_tool_support(deployment: Mapping[str, object]) -> bool | None: + """A deployment's answer to Codex's client ``tool_search``.""" + from litellm.responses.litellm_completion_transformation.tool_search import chat_config_builds_anthropic_request + from litellm.utils import ProviderConfigManager + + configured: Final = _model_info_of(deployment).get("supports_search_tool") + if isinstance(configured, bool): + return configured + params: Final = _params_of(deployment) + model: Final = params.get("model") + provider: Final = params.get("custom_llm_provider") + if not isinstance(model, str): + return None + try: + resolved_model, resolved_provider = resolve_model_and_provider_for_metadata( + model=model, custom_llm_provider=provider if isinstance(provider, str) else None + ) + chat_config: Final = ProviderConfigManager.get_provider_chat_config( + model=resolved_model, provider=LlmProviders(resolved_provider) + ) + except (litellm_exceptions.BadRequestError, ValueError): + return None + return True if chat_config_builds_anthropic_request(chat_config) else None + + +def cost_map_backend_model(deployment: Deployment) -> str: + """The model a deployment's prices are inherited from (``_cost_map_backend_model``).""" + model_info_base_model: Final = deployment.model_info.base_model + if isinstance(model_info_base_model, str) and model_info_base_model: + return model_info_base_model + params: Final = _LITELLM_PARAMS_ADAPTER.validate_python(deployment.litellm_params.model_dump(exclude_none=True)) + params_base_model: Final = params.get("base_model") + if isinstance(params_base_model, str) and params_base_model: + return params_base_model + model: Final = deployment.litellm_params.model + invoked_model: Final = params.get("model_id") + if isinstance(invoked_model, str) and invoked_model and _bedrock_profile(params, model): + profile_model: Final = invoked_model if invoked_model.startswith("bedrock/") else f"bedrock/{invoked_model}" + profile_info: Final = lookup_model_info_for_registration(model=profile_model, custom_llm_provider="bedrock") + if not isinstance(profile_info, ModelInfoLookupMiss): + return profile_model + return model + + +def deployment_model_cost_payload(host: "PythonHost", deployment: Deployment) -> Mapping[str, object]: + """The ``model_info`` a deployment contributes to ``litellm.model_cost`` + (``_deployment_model_cost_payload``): its model_info, its custom pricing, then the inherited rates.""" + declared: Final = dict(_MODEL_INFO_ADAPTER.validate_python(deployment.model_info.model_dump(exclude_none=True))) + custom_pricing: Final = { + field: value + for field in CustomPricingLiteLLMParams.model_fields + if (value := deployment.litellm_params.get(field)) is not None + } + backend_model: Final = cost_map_backend_model(deployment) + provider: Final = deployment.model_info.pricing_provider or deployment.litellm_params.custom_llm_provider + params: Final = _LITELLM_PARAMS_ADAPTER.validate_python(deployment.litellm_params.model_dump(exclude_none=True)) + priced: Final = {**declared, **custom_pricing} + profile_pricing: Final[Mapping[str, object]] = ( + _bedrock_profile_pricing(backend_model, priced) + if params.get("model_id") + and backend_model != deployment.litellm_params.model + and _bedrock_profile(params, deployment.litellm_params.model) + else {} + ) + with_profile: Final = {**priced, **profile_pricing} + off_peak: Final = {**with_profile, **_inherited_off_peak_base_rates(host, with_profile, backend_model, provider)} + cached: Final = ( + {**off_peak, **_inherited_cache_pricing(host, off_peak, backend_model, provider)} + if off_peak.get("input_cost_per_token") is not None + else off_peak + ) + tiers: Final = {**cached, **_inherited_service_tier_pricing(host, cached, backend_model, provider)} + return {**tiers, **_inherited_tiered_output_rate(host, tiers, backend_model, provider)} + + +def _bedrock_profile_pricing(backend_model: str, priced: Mapping[str, object]) -> Mapping[str, object]: + profile_info: Final = lookup_model_info_for_registration(model=backend_model, custom_llm_provider="bedrock") + if isinstance(profile_info, ModelInfoLookupMiss): + return {} + explicit_info: Final = _MODEL_INFO_ADAPTER.validate_python(priced) + priced_info: Final = model_info_with_pricing_overrides(profile_info, explicit_info) # pyright: ignore[reportArgumentType] # the explicit model_info, as the core passes it + return {key: value for key, value in priced_info.items() if is_server_derived_pricing_key(key)} + + +def _backend_cost_map_keys(model: str, custom_llm_provider: str | None) -> tuple[str, ...]: + """The ``litellm.model_cost`` keys a deployment's shared backend info is registered under.""" + backend_key: Final = model if custom_llm_provider is None else f"{custom_llm_provider}/{model}" + if "responses/" in backend_key: + return (backend_key, backend_key.replace("responses/", "")) + return (backend_key,) + + +def _catalog_entry_has_token_limit(catalog_entry: Mapping[str, object], field: str) -> bool: + return catalog_entry.get(field) is not None or ( + field == "max_tokens" and catalog_entry.get("max_output_tokens") is not None + ) + + +def _shared_backend_model_info_for_keys( + catalog: Mapping[str, Mapping[str, object]], + backend_keys: tuple[str, ...], + shared_model_info: Mapping[str, object], +) -> Mapping[str, object]: + """``shared_model_info`` without a deployment's token limit the catalog already sets on one of + ``backend_keys`` (``_shared_backend_model_info_for_keys``).""" + catalog_token_limit_fields: Final = frozenset( + field + for field in DEPLOYMENT_SCOPED_TOKEN_LIMIT_FIELDS + if any(_catalog_entry_has_token_limit(catalog.get(key) or {}, field) for key in backend_keys) + ) + return {field: value for field, value in shared_model_info.items() if field not in catalog_token_limit_fields} + + +def _catalog_backend_keys( + model: str, custom_llm_provider: str | None, backend_keys: tuple[str, ...] +) -> tuple[str, ...]: + """Every name the catalog may list the backend model under. The provider is the declared one: + resolving it would reach provider auth (``catalog_backend_keys``).""" + model_provider_prefix: Final = model.partition("/")[0] + catalog_provider: Final = ( + custom_llm_provider + if custom_llm_provider is not None + else model_provider_prefix + if model_provider_prefix in LlmProvidersSet + else None + ) + if catalog_provider is None: + return (*backend_keys, model) + potential_names: Final = get_potential_model_names(model, catalog_provider) + return ( + *backend_keys, + model, + potential_names["split_model"], + potential_names["stripped_model_name"], + potential_names["combined_model_name"], + potential_names["combined_stripped_model_name"], + potential_names["provider_prefixed_model_name"], + potential_names["region_free_combined_model_name"], + ) + + +def register_deployment_in_model_cost( + host: "PythonHost", + *, + model_id: str | None, + model_info: Mapping[str, object], + model: str, + custom_llm_provider: str | None, + pricing_provider: str | None = None, +) -> None: + """Write a deployment's metadata into ``litellm.model_cost`` (``_register_deployment_in_model_cost``): + the whole entry under its id, the schema fields under its backend model key(s).""" + registered: Final = ( + without(model_info, *CustomPricingLiteLLMParams.model_fields) + if classify_strategy_router_model(model) is not None + else dict(model_info) + ) + register_model: Final = host.invoke("register_model") + if model_id is not None: + register_model( + model_cost={model_id: registered}, + persist_across_reloads=False, + warning_display_name=model, + custom_llm_provider=pricing_provider or custom_llm_provider, + add_to_provider_model_sets=False, + ) + backend_keys: Final = _backend_cost_map_keys(model=model, custom_llm_provider=custom_llm_provider) + shared_model_info: Final = _shared_backend_model_info_for_keys( + catalog=catalog_with_runtime_registrations(), + backend_keys=_catalog_backend_keys(model, custom_llm_provider, backend_keys), + shared_model_info=shared_backend_model_info(registered), + ) + register_model( + model_cost={key: keep_registered_mode(key, shared_model_info) for key in backend_keys}, + persist_across_reloads=False, + ) + + +def register_deployment_pricing(host: "PythonHost", deployment: Deployment) -> None: + """Register a deployment's custom and inherited pricing (``_register_deployment_pricing``).""" + register_deployment_in_model_cost( + host, + model_id=deployment.model_info.id, + model_info=deployment_model_cost_payload(host, deployment), + model=deployment.litellm_params.model, + custom_llm_provider=deployment.litellm_params.custom_llm_provider, + pricing_provider=deployment.model_info.pricing_provider, + ) + + +# -- the registry -------------------------------------------------------------------------------- + + +def _model_info_of(entry: Mapping[str, object]) -> Mapping[str, object]: + model_info: Final[object] = entry.get("model_info") + return _MODEL_INFO_ADAPTER.validate_python(model_info) if isinstance(model_info, dict) else {} + + +def _params_of(entry: Mapping[str, object]) -> Mapping[str, object]: + params: Final[object] = entry.get("litellm_params") + return _LITELLM_PARAMS_ADAPTER.validate_python(params) if isinstance(params, dict) else {} + + +def _renamed(entry: UntypedEntry, model_name: str) -> UntypedEntry: + return {**entry, "model_name": model_name} + + +def _as_routing_group_row(deployment: Mapping[str, object]) -> UntypedEntry: + """A member deployment under its group's name, without the member's ``access_groups``.""" + model_info: Final = without(_model_info_of(deployment), "access_groups") + return {**deployment, "model_info": model_info} + + +def _ref_env_only(ref: object) -> bool: + return bool((as_str_mapping(ref) or {}).get("env_only")) + + +def _duplicate_of(step: Mapping[str, object], entry: _BuildEntry) -> bool: + """Whether the entry's declared id is declared twice in its list: a retried entry keeps the + answer its first build had.""" + return entry.duplicate if entry.duplicate is not None else step.get("duplicate_id") is True + + +def _index_key(value: object) -> str | None: + """``value`` as a key of the router's name and id indexes: a dict lookup hashes any key and only + a string can match, so another hashable key finds nothing (``None``), and an unhashable one + raises as the lookup raises. The native lookups take strings alone.""" + _ = hash(value) + return value if isinstance(value, str) else None + + +class HostRegistry(HostModelInfo): + """The live model_list and its indexes, kept the way the Python router core keeps them.""" + + def __init__( + self, + host: "PythonHost", + *, + model_group_alias: "dict[str, str | RouterModelGroupAliasItem]", # mutable-ok: the router's live alias map + ignore_invalid_deployments: bool, + on_change: Callable[[], None], + ) -> None: + self._host: Final = host + self._on_change: Final = on_change + self.ignore_invalid_deployments: bool = ignore_invalid_deployments + self._registry_version: int = 0 + self._deployment_index_lock: Final = threading.RLock() + self._core_model_group_alias: dict[str, str | RouterModelGroupAliasItem] = model_group_alias # mutable-ok: the live alias map # fmt: skip + # The live model_list: router-core's deployments read back, every write a changeset. + self._core_model_list: Final[HostedModelList] = HostedModelList(self._live_handle, self._hand_changes) + # The indexes the registry hands out: router-core's, read back, every write a changeset. + self._indexes: Final = HostedIndexes(self._live_handle, self._hand_changes) + self._core_deployment_latency_map: dict[str, float] = {} # mutable-ok: the live latency map + self._access_groups_cache: Mapping[str, Sequence[str]] | None = None + # The routing groups by name; RustCore sets them from its resolved settings. + self._core_routing_groups: dict[str, RoutingGroup] = {} # mutable-ok: the routing_groups attribute + self._routing_group_rows: tuple[UntypedEntry, ...] | None = None + self._provider_unresolved_deployments: tuple[Callable[[], None], ...] = () + # Runs router-core's deployment build; RustCore sets it before the first build. + self.build_runner: BuildRunner | None = None + self.missing_attribute: Callable[[str], AttributeError] | None = None + # Hands router-core a write its build did not make (a removal, a swapped list) as a + # changeset, and reads its current handle; RustCore sets both before the first build. + self.changes_runner: Callable[[Mapping[str, object]], None] | None = None + self.handle_lookup: Callable[[], RouterCoreHandle | None] | None = None + # The pre-routing strategy registries (``registry().plugins``) and the hooks that keep them; + # RustCore installs its ``PreRoutingStrategies`` (litellm/router_rust/steps/prerouting.py). + self._core_auto_routers: dict[str, list[TaggedPreRoutingStrategy[object]]] = {} # mutable-ok: the live registry + self._core_complexity_routers: dict[str, list[TaggedPreRoutingStrategy[object]]] = {} # mutable-ok: same + self._core_adaptive_routers: dict[str, list[TaggedPreRoutingStrategy[object]]] = {} # mutable-ok: same + self._core_quality_routers: dict[str, list[TaggedPreRoutingStrategy[object]]] = {} # mutable-ok: same + self.strategies: PreRoutingStrategies | None = None + # The router's optional pre-call checks (deployment budgets, prompt caching pins); + # RustCore installs its ``OptionalChecks`` (litellm/router_rust/steps/callbacks.py). + self.checks: OptionalChecks | None = None + # The Rust core's resolution for a request (`_core_get_routable_deployments`), set by the core. + self.routable_lookup: ( + Callable[[str, Mapping[str, object], str | None], Sequence[Mapping[str, object]]] | None + ) = None + # A deployment's cached `max_parallel_requests` limit (litellm/router_rust/invoke/slots.py), set by the core. + self.max_parallel_requests_lookup: Callable[[Mapping[str, object], Mapping[str, object]], object] | None = None + # `plugins("routing")`: the router's routing plugins, set by the core. + self._core_routing_plugins: list[RoutingPlugin] = [] # mutable-ok: the Python core hands out its live list + # `zero_cost()`: the per-deployment zero-cost lookup the auth checks fill, emptied with the + # group info caches. + self._zero_cost_cache: dict[str, bool] = {} # mutable-ok: the auth checks write into it + # The routing cache usage reads come from; RustCore sets it to the one its Rust core writes. + self._core_cache: DualCache = DualCache() + self._discovered_model_info_cache: InMemoryCache = InMemoryCache( + max_size_in_memory=1, default_ttl=2 * MODEL_INFO_REFRESH_SECONDS + ) + # The PTU share teams and team ceilings, until the model list next changes. + self._ptu_cache: _PTUMemo | None = None + self._cached_get_model_group_info: Callable[[str], ModelGroupInfo | None] = lru_cache( # pyright: ignore[reportIncompatibleMethodOverride] # the per-router cache shadows the uncached method, as the core keeps it + maxsize=DEFAULT_MAX_LRU_CACHE_SIZE + )(self._core_get_model_group_info) + self._core_cached_deployment_model_info: Callable[[str, str], object] = lru_cache( + maxsize=DEFAULT_MAX_LRU_CACHE_SIZE + )(self._core_get_deployment_model_info) + self.cached_deployment_tracks_usage: Callable[[str, str], bool] = lru_cache(maxsize=DEFAULT_MAX_LRU_CACHE_SIZE)( + self._deployment_tracks_usage + ) + + def __getattr__(self, name: str) -> object: + from litellm.router_rust.core import RustCoreUnsupported + + missing: Final = self.__dict__.get("missing_attribute") + if name == "_core_model_group_alias" and missing is not None: + # `del router.model_group_alias`: every read raises, as on the Python core. + raise missing(name) + raise RustCoreUnsupported(f"registry().{name.removeprefix('_core_')} is not served by the Rust router core") + + # -- building --------------------------------------------------------------------------------- + + def build(self, model_list: Sequence[object] | None) -> None: + """The constructor's registry: ``set_model_list`` and the latency map, as the core builds them.""" + if model_list is None: + return + self._discovered_model_info_cache = InMemoryCache( + max_size_in_memory=max(len(model_list), 1), default_ttl=2 * MODEL_INFO_REFRESH_SECONDS + ) + self.set_model_list(model_list) + for entry in model_list: + params = _params_of(_LITELLM_PARAMS_ADAPTER.validate_python(entry)) + if "model" in params: + self._core_deployment_latency_map[str(params["model"])] = 0 + + async def arefresh_model_info(self, client: "AsyncHTTPHandler | None") -> None: + """``_core_arefresh_model_info``: token limits the served OpenAI-compatible deployments + advertise on their models endpoint, read with a few workers over a snapshot of the list.""" + import asyncio + + from litellm.llms.openai_like.model_info import MODEL_INFO_REFRESH_CONCURRENCY + + deployments: Final = iter(tuple(self._core_model_list)) + + async def refresh_worker() -> None: + for raw_deployment in deployments: + try: + await self._arefresh_deployment_model_info(raw_deployment, client=client) + except Exception: # noqa: BLE001 # one invalid deployment must not prevent refreshing the others + verbose_router_logger.debug("Could not refresh deployment model info") + + _ = await asyncio.gather(*(refresh_worker() for _ in range(MODEL_INFO_REFRESH_CONCURRENCY))) + self._invalidate_model_group_info_cache() + + async def _arefresh_deployment_model_info( + self, raw_deployment: Mapping[str, object], *, client: "AsyncHTTPHandler | None" + ) -> None: + from litellm.llms.custom_httpx.http_handler import get_async_httpx_client + from litellm.llms.openai_like.model_info import ( + MODEL_INFO_DISCOVERY_PROVIDERS, + get_openai_compatible_model_info, + ) + from litellm.types.router import DiscoveredDeploymentModelInfo + from litellm.utils import aresolve_api_key_secret_reference + + deployment: Final = Deployment.model_validate(raw_deployment) + params: Final = LiteLLM_Params.model_validate( + MappingProxyType( + { + **deployment.litellm_params.model_dump(exclude_none=True), + **( + self._core_get_deployment_credentials_with_provider(deployment.model_info.id or "") + or MappingProxyType({}) + ), + } + ) + ) + model, provider, dynamic_api_key, api_base = self._host.invoke("get_llm_provider")( # pyright: ignore[reportAny] # litellm's untyped tuple + model=params.model, litellm_params=params + ) + if provider not in MODEL_INFO_DISCOVERY_PROVIDERS: + return + if api_base is None or "*" in model or params.get("use_clientside_credentials"): # pyright: ignore[reportUnknownMemberType] # the params' untyped get + return + api_key: Final = await aresolve_api_key_secret_reference( + params.api_key or dynamic_api_key, # pyright: ignore[reportAny] # see above + model=model, # pyright: ignore[reportAny] # see above + llm_provider=provider, # pyright: ignore[reportAny] # see above + upstream=MappingProxyType({"api_base": params.api_base}), + ) + headers: Final = TypeAdapter(Mapping[str, str]).validate_python( + params.get("extra_headers") or params.get("headers") or MappingProxyType({}) # pyright: ignore[reportUnknownMemberType] # see above + ) + auth_headers: Final = ( + MappingProxyType({"authorization": f"Bearer {api_key}"}) if api_key else MappingProxyType({}) + ) + limits: Final = await get_openai_compatible_model_info( + model=model, # pyright: ignore[reportAny] # see above + api_base=api_base, # pyright: ignore[reportAny] # see above + headers=MappingProxyType( + {**auth_headers, **MappingProxyType({key.lower(): value for key, value in headers.items()})} + ), + client=client or get_async_httpx_client(llm_provider=LlmProviders.OPENAI), + cache=self._core_cache.in_memory_cache, + ) + model_id: Final = deployment.model_info.id + if not limits or model_id is None or self._core_get_model_info(model_id) is not raw_deployment: + return + self._discovered_model_info_cache.max_size_in_memory = max(len(self._core_model_list), 1) + self._discovered_model_info_cache.delete_cache(model_id) + self._discovered_model_info_cache.set_cache( + model_id, + DiscoveredDeploymentModelInfo(deployment=raw_deployment, limits=limits), # pyright: ignore[reportArgumentType] # the model_list entry, as the core keeps it + ) + self._invalidate_model_group_info_cache() + + def _deployment_tracks_usage(self, deployment_id: str, model_group: str) -> bool: + """``_deployment_tracks_usage``: a limit on the deployment or in its merged model info, where + the cost map's own ``tpm`` and ``rpm`` count too.""" + deployment_info: Final = self._core_get_deployment(model_id=deployment_id) + if deployment_info is None: + return False + deployment_model_info: Final = self._core_get_router_model_info( + deployment=deployment_info, + received_model_name=model_group, + ) + configured_limits: Final[tuple[object, ...]] = ( + deployment_info.get("tpm", None), + deployment_info.get("rpm", None), + deployment_info.litellm_params.tpm, + deployment_info.litellm_params.rpm, + deployment_model_info.get("tpm", None), + deployment_model_info.get("rpm", None), + ) + return any(limit is not None for limit in configured_limits) or deployment_has_io_token_limits( + deployment_info.model_dump() + ) + + def _invalidate_model_group_info_cache(self) -> None: + self._ptu_cache = None + self._cached_get_model_group_info.cache_clear() # pyright: ignore[reportFunctionMemberAccess] # an lru_cache + self._core_cached_deployment_model_info.cache_clear() # pyright: ignore[reportFunctionMemberAccess] # an lru_cache + self.cached_deployment_tracks_usage.cache_clear() # pyright: ignore[reportFunctionMemberAccess] # an lru_cache + self._zero_cost_cache.clear() + self._routing_group_rows = None + + def set_routing_groups(self, groups: Sequence[RoutingGroup]) -> None: + """``_init_routing_groups``' registry side: the groups by name, and the caches they feed.""" + self._core_routing_groups = {group.group_name: group for group in groups} + self._invalidate_model_group_info_cache() + self._invalidate_access_groups_cache() + + def _invalidate_access_groups_cache(self) -> None: + self._access_groups_cache = None + self._registry_version += 1 + + def set_model_list(self, model_list: Sequence[object]) -> None: + """Replace every deployment (``set_model_list``). All or nothing: when an entry raises, + the deployments served before the call are put back and the error is raised.""" + previous: Final = [copy.deepcopy(dict(entry)) for entry in self._core_model_list] + try: + self._replace_model_list(model_list) + except Exception: + self._replace_model_list(previous) + raise + + def _replace_model_list(self, model_list: Sequence[object]) -> None: + self._discovered_model_info_cache.flush_cache() + original_model_list: Final[list[UntypedEntry]] = copy.deepcopy(list(model_list)) # pyright: ignore[reportAssignmentType] # model_list entries # mutable-ok: a deep copy the build consumes + if self.strategies is not None: + self.strategies.reset() + self._provider_unresolved_deployments = () + self._invalidate_model_group_info_cache() + self._invalidate_access_groups_cache() + ptu_declared: Final = tuple( + str(entry.get("model_name")) + for entry in original_model_list + if (ptu_info := _model_info_of(entry)).get("db_model") is not True and declares_ptu(ptu_info) + ) + if ptu_declared and not is_ptu_cost_attribution_enabled(): + verbose_router_logger.warning( + "PTU fields are set on config.yaml deployment(s) %s, but PTU cost attribution is disabled, so no " + "flat cost accrues and this traffic is billed per token. Set %s=True to enable it", + ", ".join(ptu_declared), + PTU_COST_ATTRIBUTION_ENV_VAR, + ) + # router-core walks the list (one deployment per organization, secrets resolved first, + # ids assigned), asking this registry about each deployment in turn; it starts the list + # and its id, name and team indexes over, and takes the given names as `model_names`. + self._run_build( + _Build(path="create", raw=original_model_list), + original_model_list, + resolve_secrets=True, + install="replace", + ) + if self.strategies is not None: + self.strategies.finalize() + self._on_change() + + # -- the build's steps (router-core's `build::build`) ------------------------------------------ + + def _run_build( + self, build: _Build, entries: Sequence[object], *, resolve_secrets: bool, install: BuildInstall + ) -> None: + """Have router-core build ``entries``, answering its steps from this registry.""" + if self.build_runner is None: + raise RuntimeError("the host registry builds deployments through its Rust core") + for entry in entries: + refuse_wide_numbers(entry) + self._core_model_list.begin_build(replacing=install == "replace") + try: + self.build_runner(entries, build.path, resolve_secrets, install, partial(self._answer_shown, build)) + finally: + self._core_model_list.end_build() + if build.failure is not None: + # router-core served what it built before the entry that raised (`set_model_list` + # cleared the list first), so the router is left as the Python core leaves it. + raise build.failure + + def _answer_shown(self, build: _Build, step: Mapping[str, object]) -> object: + """A build step's answer; a deployment the ``added`` step hands router-core shows in the + live list from then on (router-core serves it when the build ends).""" + answer: Final = self._answer_build_step(build, step) + stored: Final = as_str_mapping(answer).get("entry") if step["kind"] == "added" else None # pyright: ignore[reportOptionalMemberAccess] # an added step answers a mapping + if isinstance(stored, Mapping): + self._core_model_list.built(cast(Mapping[str, object], stored)) # cast-ok: the entry it stores + return answer + + def _answer_build_step(self, build: _Build, step: Mapping[str, object]) -> object: + """The answer to one of router-core's build steps. A ``set_model_list`` deployment that + fails is left out when the router ignores invalid deployments (``_create_deployment``), + and retried on the next price data reload when its provider did not resolve.""" + if build.failure is not None: + return {"invalid": True} + kind: Final = step["kind"] + index: Final = int(str(step["index"])) + refs: Final = step.get("refs") + if kind == "secrets" and isinstance(refs, list) and not any(_ref_env_only(ref) for ref in refs): + return self._build_secrets(build, index, refs) + if kind == "validate": + # A new deployment: the entry is popped where the loop pops it, outside the error + # handling, and nothing of the last deployment carries over. + build.pending = None # rebind-ok: the build's state, which each step advances + build.provider = None # rebind-ok: the build's state, which each step advances + build.data_source_values = () # rebind-ok: the build's state, which each step advances + if build.path == "create": + _ = self._build_entry(build, index) + try: + if kind == "secrets": + return self._build_data_sources(build, refs if isinstance(refs, list) else []) + if kind == "validate": + return self._build_validate(build, index, step) + if kind == "provider": + build.provider = self._build_provider(build) # rebind-ok: the build's state, which each step advances + return {"provider": build.provider} + return self._build_added(build) + except Exception as e: + if build.path != "create": + raise + if not self.ignore_invalid_deployments: + build.failure = e # rebind-ok: raised once router-core ends the build + return {"invalid": True} + if isinstance(e, litellm_exceptions.BadRequestError): + entry: Final = build.entries[index] + retried: Final = dataclasses.replace(entry, duplicate=_duplicate_of(step, entry)) + self._provider_unresolved_deployments = ( + *self._provider_unresolved_deployments, + partial(self._retry_create, retried), + ) + verbose_router_logger.exception( + "Error creating deployment: %s, ignoring and continuing with other deployments.", e + ) + return {"invalid": True} + + def _build_entry(self, build: _Build, index: int) -> _BuildEntry: + """Entry ``index`` as ``set_model_list``'s loop holds it, popped off the first time a step + names it (so a malformed entry raises where the loop raises).""" + held: Final = build.entries.get(index) + if held is not None: + return held + model: Final = build.raw[index] + model_name: Final = model.pop("model_name") + raw_params: Final = model.pop("litellm_params") + litellm_params: Final[UntypedEntry] = raw_params if isinstance(raw_params, dict) else {} # pyright: ignore[reportAssignmentType] # configured params + model_info: Final = model.pop("model_info", {}) + info_dict: Final[UntypedEntry] = model_info if isinstance(model_info, dict) else {} # pyright: ignore[reportAssignmentType] # a configured model_info + entry: Final = _BuildEntry( + extras=model, + model_name=str(model_name), + litellm_params=litellm_params, + model_info=info_dict, + declared_id=None if info_dict.get("id") is None else str(info_dict["id"]), + per_organization=isinstance(litellm_params.get("organization"), list), + ) + build.entries[index] = entry # rebind-ok: the build's state, which each step advances + return entry + + def _build_secrets(self, build: _Build, index: int, refs: Sequence[object]) -> Sequence[object]: + """``set_model_list``'s ``get_secret`` for each top-level ``os.environ/`` value router-core + found, written into the entry's params.""" + entry: Final = self._build_entry(build, index) + references: Final = tuple(as_str_mapping(ref) or {} for ref in refs) + values: Final = tuple(get_secret(str(reference["name"])) for reference in references) + entry.litellm_params.update( + {str(reference["key"]): value for reference, value in zip(references, values, strict=True)} + ) + return values + + def _build_data_sources(self, build: _Build, refs: Sequence[object]) -> Sequence[object]: + """``_add_deployment``'s environment reads for ``dataSources[].parameters``, applied to the + deployment where it applied them (``_build_added``).""" + references: Final = tuple(as_str_mapping(ref) or {} for ref in refs) + build.data_source_values = tuple( # rebind-ok: the build's state, which each step advances + ( + int(str(reference["source"])), + str(reference["key"]), + os.environ.get(str(reference["name"]).replace("os.environ/", ""), ""), + ) + for reference in references + ) + return tuple(value for _, _, value in build.data_source_values) + + def _build_validate(self, build: _Build, index: int, step: Mapping[str, object]) -> Mapping[str, object]: + if build.path == "add": + return self._validate_added_deployment(build) + entry: Final = self._build_entry(build, index) + routed: Final = as_str_mapping(step["entry"]) or {} + routed_info: Final = as_str_mapping(routed.get("model_info")) or {} + if "id" not in entry.model_info: + entry.model_info["id"] = routed_info.get("id") + if entry.per_organization: + entry.litellm_params["organization"] = (as_str_mapping(routed.get("litellm_params")) or {}).get( + "organization" + ) + duplicate: Final = _duplicate_of(step, entry) + deployment: Final = self._validated_deployment(entry, taken=duplicate) + if deployment is None: + return {"skip": True} + served: Final = with_admin_api_key_secret_reference(deployment) + build.pending = served # rebind-ok: the build's state, which each step advances + # A dump of this core's own: a `mock_response` ModelResponse warns, and only Python's dumps may. + return {"entry": build.pending.to_json(exclude_none=True, warnings=False)} + + def _validated_deployment(self, entry: _BuildEntry, *, taken: bool) -> Deployment | None: + """``_create_deployment`` up to ``_add_deployment``: the config checks, the deployment, its + pricing, and None when the environment gate leaves it out.""" + model_name: Final = entry.model_name + litellm_params: Final = entry.litellm_params + model_info: Final = entry.model_info + config_sourced: Final = model_info.get("db_model") is not True + identity_error: Final = ( + ptu_identity_error( + declared_id=entry.declared_id, + taken=taken, + current_id=model_info.get("id"), # pyright: ignore[reportArgumentType] # the configured id + model_name=model_name, + ) + if config_sourced and ptu_terms(model_info) is not None + else None + ) + ptu_error: Final = ( + (ptu_config_error(model_info, model_name=model_name) or identity_error) if config_sourced else None + ) + if ptu_error is not None and (declares_ptu_shares(model_info) or is_ptu_cost_attribution_enabled()): + raise ValueError(ptu_error) + access_windows_error: Final = access_windows_config_error(model_info, model_name=model_name) + if access_windows_error is not None: + raise ValueError(access_windows_error) + capacity_warning: Final = ( + ptu_capacity_warning( + model_name, MappingProxyType({"model_info": model_info, "litellm_params": litellm_params}) + ) + if is_ptu_cost_attribution_enabled() + else None + ) + if capacity_warning is not None: + verbose_router_logger.warning(capacity_warning) + zeroed_pricing: Final = zeroed_ptu_pricing(model_info, litellm_params) if config_sourced else None + merged_params: Final[Mapping[str, object]] = ( + litellm_params if zeroed_pricing is None else MappingProxyType({**litellm_params, **zeroed_pricing}) + ) + params: Final[LiteLLM_Params] = validate_litellm_params_strictly(LiteLLM_Params, dict(**merged_params)) + declared_params: Final = _LITELLM_PARAMS_ADAPTER.validate_python(litellm_params) + warn_on_provider_credential_mismatch(model_name=model_name, litellm_params=declared_params) + warn_on_ignored_model_capability_flags(model_name=model_name, litellm_params=declared_params) + warn_on_unenforced_endpoint_flag(model_name=model_name, model_info=model_info) + deployment: Final = without_model_capability_flags( + Deployment(**entry.extras, model_name=model_name, litellm_params=params, model_info=model_info) # pyright: ignore[reportArgumentType] # the entry's extra fields, as the core passes them + ) + # The values as typed: a cost written as "inf", a share written as 1e20. + refuse_unservable_numbers(deployment.to_json(exclude_none=True, warnings=False)) + register_deployment_pricing(self._host, deployment) + if model_info_is_active_for_environment(model_info=deployment.model_info) is not True: # pyright: ignore[reportArgumentType] # ModelInfo reads like a mapping + verbose_router_logger.warning( + "Ignoring deployment %s as it is not active for environment %s", + deployment.model_name, + deployment.model_info["supported_environments"], + ) + return None + for pattern in deployment.litellm_params.get("tag_regex") or []: + try: + re.compile(pattern) + except re.error as exc: + raise ValueError( + f"Invalid regex in tag_regex for model '{deployment.model_name}': {pattern!r} — {exc}" + ) from exc + return deployment + + def _validate_added_deployment(self, build: _Build) -> Mapping[str, object]: + """``add_deployment`` up to ``_add_deployment``: an id the registry holds is skipped.""" + deployment: Final = build.deployment + if deployment is None: + raise RuntimeError("an add build holds its deployment") + if deployment.model_info.id and self._core_has_model_id(deployment.model_info.id): + return {"skip": True} + declared_params: Final = deployment.litellm_params.model_dump(exclude_none=True) + warn_on_provider_credential_mismatch(model_name=deployment.model_name, litellm_params=declared_params) + warn_on_ignored_model_capability_flags(model_name=deployment.model_name, litellm_params=declared_params) + warn_on_unenforced_endpoint_flag( + model_name=deployment.model_name, model_info=deployment.model_info.model_dump() + ) + served: Final = with_admin_api_key_secret_reference(without_model_capability_flags(deployment)) + build.pending = served # rebind-ok: the build's state, which each step advances + build.stored = served.to_json(exclude_none=True) # rebind-ok: the build's state, which each step advances + refuse_unservable_numbers(build.stored) + return {"entry": build.stored} + + def _build_provider(self, build: _Build) -> str | None: + """``_add_deployment``'s provider: None for a prompt-management model, else the provider + ``get_llm_provider`` resolves, refused when litellm does not serve it.""" + deployment: Final = build.pending + if deployment is None: + raise RuntimeError("a provider step follows a validated deployment") + litellm_model: Final = deployment.litellm_params.model + if isinstance(deployment.litellm_params.drop_params, str): + verbose_router_logger.warning( + "model=%s drop_params=%r is not a flag value, treating it as unset", + deployment.model_name, + deployment.litellm_params.drop_params, + ) + dynamic: Final = self._host.dynamic() + if "/" in litellm_model and litellm_model.split("/")[0] in dynamic.known_custom_logger_compatible_callbacks: + return None + if declared_authenticating_provider(litellm_model, deployment.litellm_params.custom_llm_provider) is not None: + _, provider = self._model_group_llm_provider(deployment.litellm_params) + return provider + _, custom_llm_provider, _, _ = self._host.invoke("get_llm_provider")( + model=deployment.litellm_params.model, + custom_llm_provider=deployment.litellm_params.get("custom_llm_provider", None), + api_base=deployment.litellm_params.api_base, + litellm_params=deployment.litellm_params, + ) + if ( + custom_llm_provider not in dynamic.provider_list + and not JSONProviderRegistry.exists(custom_llm_provider) + and not is_registered_custom_provider(custom_llm_provider) + ): + raise Exception(f"Unsupported provider - {custom_llm_provider}") # noqa: TRY002 # the core raises a plain Exception + return custom_llm_provider + + def _build_added(self, build: _Build) -> Mapping[str, object]: + """The rest of ``_add_deployment``: the entry router-core stores, which joins the list and + the indexes (``host_index.rs``) when the step returns.""" + deployment: Final = build.pending + if deployment is None: + raise RuntimeError("an added step follows a validated deployment") + self._register_deployment(deployment, build.provider, build.data_source_values) + build.pending = None # rebind-ok: the build's state, which each step advances + build.data_source_values = () # rebind-ok: the build's state, which each step advances + self._invalidate_model_group_info_cache() + self._invalidate_access_groups_cache() + if build.path == "create": + return {"entry": deployment.to_json(exclude_none=True)} + entry: Final = build.stored if build.stored is not None else deployment.to_json(exclude_none=True) + register_deployment_pricing(self._host, deployment) + build.result = deployment # rebind-ok: the build's state, which each step advances + return {"entry": entry} + + def _retry_create(self, entry: _BuildEntry) -> None: + """A ``set_model_list`` deployment left out when its provider did not resolve, built again + and added at the end.""" + build: Final = _Build(path="create") + build.entries[0] = entry + self._run_build( + build, + [ + { + **entry.extras, + "model_name": entry.model_name, + "litellm_params": entry.litellm_params, + "model_info": entry.model_info, + } + ], + resolve_secrets=False, + install="append", + ) + + @staticmethod + def _model_group_llm_provider(litellm_params: LiteLLM_Params) -> tuple[str, str]: + try: + return resolve_model_and_provider_for_metadata( + model=litellm_params.model, custom_llm_provider=litellm_params.custom_llm_provider + ) + except litellm_exceptions.BadRequestError as e: + verbose_router_logger.error("litellm.router.py::get_model_group_info() - %s", e) + return "", "" + + def _register_deployment( + self, + deployment: Deployment, + custom_llm_provider: str | None, + data_source_values: Sequence[tuple[int, str, str]], + ) -> None: + """``_add_deployment`` after the provider: its data-source values, pass-through + credentials and pre-routing strategies (router-core registers its names and patterns).""" + if deployment.litellm_params.rpm is None and getattr(deployment, "rpm", None) is not None: + deployment.litellm_params.rpm = getattr(deployment, "rpm") # noqa: B009 # an extra field of the entry # rebind-ok: the core sets it on the deployment it adds + if deployment.litellm_params.tpm is None and getattr(deployment, "tpm", None) is not None: + deployment.litellm_params.tpm = getattr(deployment, "tpm") # noqa: B009 # an extra field of the entry # rebind-ok: see above + data_sources: Final = deployment.litellm_params.get("dataSources", []) or [] + for source, key, value in data_source_values: + data_sources[source]["parameters"][key] = value + if custom_llm_provider is not None: + self._initialize_deployment_for_pass_through(deployment, custom_llm_provider) + if self.strategies is not None: + self.strategies.register(deployment) + + def _initialize_deployment_for_pass_through(self, deployment: Deployment, custom_llm_provider: str) -> None: + """Register vertex credentials for a pass-through deployment (``_initialize_deployment_for_pass_through``).""" + if deployment.litellm_params.use_in_pass_through is not True or custom_llm_provider != "vertex_ai": + return + credential_name: Final = deployment.litellm_params.litellm_credential_name + credential_values: Final = ( + CredentialAccessor.get_credential_values(credential_name) if credential_name is not None else {} + ) + project: Final = credential_values.get("vertex_project") or deployment.litellm_params.vertex_project + location: Final = credential_values.get("vertex_location") or deployment.litellm_params.vertex_location + credentials: Final = credential_values.get("vertex_credentials") or deployment.litellm_params.vertex_credentials + if project is None or location is None: + raise ValueError( + "vertex_project, and vertex_location must be set in litellm_params for pass-through endpoints." + ) + from litellm.router_host import router_host + + router_host().add_vertex_pass_through_credentials( + project_id=str(project), + location=str(location), + vertex_credentials=credentials, # pyright: ignore[reportArgumentType] # the configured credentials + ) + + # -- writes ----------------------------------------------------------------------------------- + + def add_deployment(self, deployment: Deployment) -> Deployment | None: + """Add one deployment unless its id exists (``add_deployment``), built by router-core.""" + build: Final = _Build(path="add", deployment=deployment) + self._run_build(build, [deployment.to_json(exclude_none=True)], resolve_secrets=False, install="upsert") + if build.result is None: + return None + if self.checks is not None: + self.checks.sync_deployment_budget(build.result) + self._on_change() + return build.result + + def _core_upsert_deployment(self, deployment: Deployment) -> Deployment | None: + """Add or replace one deployment (``_upsert_deployment``).""" + model_id: Final = deployment.model_info.id or "" + on_router: Final = self._core_get_deployment(model_id=model_id) + served: Final = without_model_capability_flags(deployment) + try: + if on_router is not None: + if ( + deployment.model_name == on_router.model_name + and served.litellm_params == on_router.litellm_params + and deployment.model_info == on_router.model_info + ): + return None + with self._deployment_index_lock: + removal_idx: Final = self.model_id_to_deployment_index_map.get(model_id) + if removal_idx is not None: + _ = self._core_model_list[removal_idx] # pyright: ignore[reportArgumentType] # raises IndexError for a position past the list, as the pop does + self._remove(model_id) + if self.strategies is not None: + self.strategies.release(on_router) + self._on_change() + self.add_deployment(deployment=deployment) + if self.strategies is not None: + self.strategies.upserted(deployment, on_router) + except Exception as e: + if not self.ignore_invalid_deployments: + raise + verbose_router_logger.warning( + "Error upserting deployment %s (id=%s): %s. Dropping it and continuing with other deployments.", + deployment.model_name, + deployment.model_info.id, + e, + ) + self._restore_deployment_after_failed_upsert(on_router, model_id) + return None + else: + return deployment + + def _restore_deployment_after_failed_upsert(self, previous: Deployment | None, model_id: str) -> None: + if previous is None or self._core_has_model_id(model_id): + return + try: + if self.strategies is None: + self.add_deployment(deployment=previous) + else: + self.strategies.without_capability_limit(lambda: self.add_deployment(deployment=previous)) + verbose_router_logger.info( + "Restored deployment %s (id=%s); it keeps serving its previous configuration.", + previous.model_name, + model_id, + ) + except Exception as restore_error: # noqa: BLE001 # best-effort restore: a second failure must not abort the reload + verbose_router_logger.warning( + "Could not restore previously served deployment %s (id=%s) after the failed upsert: %s", + previous.model_name, + model_id, + restore_error, + ) + + def _live_handle(self) -> "RouterCoreHandle | None": + return None if self.handle_lookup is None else self.handle_lookup() + + def _hand_changes(self, changes: Mapping[str, object]) -> None: + if self.changes_runner is None: + raise RuntimeError("the host registry writes deployments through its Rust core") + self.changes_runner(changes) + self._core_model_list.refresh() + + def refresh(self) -> None: + """Read router-core's list again after a write (only what was appended, when that is all + that changed). The index views are read again when an entry point hands them out, which + keeps one ``add_deployment`` from reading every index back.""" + self._core_model_list.refresh() + + def _remove(self, model_id: str) -> None: + """Remove the deployment at the id map's position for ``model_id``, every index shifted + (``_update_deployment_indices_after_removal``, in router-core).""" + self._hand_changes({"remove": model_id}) + self._invalidate_model_group_info_cache() + self._invalidate_access_groups_cache() + + # -- the indexes (router-core's, read back as live views) ------------------------------------- + + @property + def model_id_to_deployment_index_map(self) -> dict[str, int]: # mutable-ok: the live id index + self._indexes.refresh() + return cast(dict[str, int], self._indexes.id_map) # cast-ok: id to position, as router-core keeps it + + @property + def _core_model_name_to_deployment_indices(self) -> dict[str, list[int]]: # mutable-ok: the live name index + self._indexes.refresh() + return cast(dict[str, list[int]], self._indexes.model_name) # cast-ok: name to positions + + @_core_model_name_to_deployment_indices.setter + def _core_model_name_to_deployment_indices(self, value: object) -> None: + self._indexes.assign("model_name", value) + + @property + def _core_team_model_to_deployment_indices( + self, + ) -> dict[tuple[str, str], list[int]]: # mutable-ok: the live team index + self._indexes.refresh() + return cast(_TeamIndex, self._indexes.team_model) # cast-ok: (team, public name) to positions + + @_core_team_model_to_deployment_indices.setter + def _core_team_model_to_deployment_indices(self, value: object) -> None: + self._indexes.assign("team_model", value) + + @property + def _core_model_names(self) -> set[str]: # mutable-ok: the live group set + self._indexes.refresh() + return cast(set[str], self._indexes.model_names) # cast-ok: group names + + @_core_model_names.setter + def _core_model_names(self, value: object) -> None: + self._indexes.assign("model_names", value) + + @property + def _core_deployment_names(self) -> list[str]: # mutable-ok: the live name list + self._indexes.refresh() + return cast(list[str], self._indexes.deployment_names) # cast-ok: litellm model names + + @_core_deployment_names.setter + def _core_deployment_names(self, value: object) -> None: + self._indexes.assign("deployment_names", value) + + @property + def _core_team_public_model_names(self) -> frozenset[str]: + self._indexes.refresh() + return self._indexes.team_public_model_names + + @_core_team_public_model_names.setter + def _core_team_public_model_names(self, value: object) -> None: + self._indexes.assign("team_public_model_names", value) + + @property + def _core_provider_default_deployment_ids(self) -> list[str]: # mutable-ok: the live id list + self._indexes.refresh() + return cast(list[str], self._indexes.provider_default_ids) # cast-ok: deployment ids + + @_core_provider_default_deployment_ids.setter + def _core_provider_default_deployment_ids(self, value: object) -> None: + self._indexes.assign("provider_default_ids", value) + + @property + def _core_pattern_router(self) -> PatternMatchRouter: + self._indexes.refresh() + return self._indexes.pattern_router + + @_core_pattern_router.setter + def _core_pattern_router(self, value: object) -> None: + self._indexes.pattern_router.patterns = getattr(value, "patterns", {}) # pyright: ignore[reportAttributeAccessIssue] # a router a caller assigned + + @property + def _core_team_pattern_routers(self) -> dict[str, PatternMatchRouter]: # mutable-ok: the live per-team routers + self._indexes.refresh() + return cast(dict[str, PatternMatchRouter], self._indexes.team_pattern_routers) # cast-ok: team id to its router + + @_core_team_pattern_routers.setter + def _core_team_pattern_routers(self, value: object) -> None: + routers: Final = as_str_mapping(value) or {} + self._indexes.write( + { + "team_patterns": [ + [ + str(team), + [ + [regex, deployments] + for regex, deployments in (as_str_mapping(getattr(router, "patterns", None)) or {}).items() + ], + ] + for team, router in routers.items() + ] + } + ) + + def replace_model_list(self, model_list: "list[UntypedEntry]") -> None: # mutable-ok: the list as assigned + """``router.model_list = [...]``: the list is swapped as given, without reindexing, and + router-core takes it whole.""" + self._core_model_list.replace(model_list) + self._registry_version += 1 + self._on_change() + + def _core_delete_deployment(self, id: str) -> UntypedEntry | None: + """Remove one deployment and return its stored entry (``_delete_deployment``).""" + try: + with self._deployment_index_lock: + deployment_idx: Final = self.model_id_to_deployment_index_map.get(id) + if deployment_idx is None: + return None + item: Final = self._core_model_list[deployment_idx] # pyright: ignore[reportArgumentType] # raises IndexError for a position past the list, as the pop does + self._remove(id) + if self.checks is not None: + self.checks.unregister_deployment_budget(id) + if self.strategies is not None: + self.strategies.release_entry(item) + except Exception: # noqa: BLE001 # the core answers None for any failed delete + return None + self._on_change() + return item + + def _core_get_router_deployment_budget_limiter(self) -> "RouterBudgetLimiting | None": + """The router's deployment-budget callback (``registry().plugins("budget_limiter")``).""" + return self.checks.budget_limiter() if self.checks is not None else None + + @property + def _core_prompt_caching_pin_cache(self) -> DualCache: + """The prompt caching pins (``registry().plugins("prompt_caching_pins")``).""" + if self.checks is None: + return self.__getattr__("_core_prompt_caching_pin_cache") # pyright: ignore[reportReturnType] # raises RustCoreUnsupported + return self.checks.prompt_caching_pin_cache + + def replay_model_cost_registrations(self) -> None: + """Re-register every served deployment's cost-map entries (``replay_model_cost_registrations``).""" + provider_unresolved: Final = self._provider_unresolved_deployments + self._provider_unresolved_deployments = () + for create_deployment in provider_unresolved: + create_deployment() + for entry in tuple(self._core_model_list): + try: + deployment = Deployment(**entry) # pyright: ignore[reportArgumentType] # a stored entry + except Exception: # noqa: BLE001 # a malformed entry must not abort the rest of the rebuild + verbose_router_logger.exception( + "Router: could not rebuild cost-map entry for a deployment during a price data reload" + ) + continue + register_deployment_pricing(self._host, deployment) + self._invalidate_model_group_info_cache() + if provider_unresolved: + self._on_change() + + # -- one deployment --------------------------------------------------------------------------- + + def _core_has_model_id(self, candidate_id: str) -> bool: + handle: Final = self._live_handle() + key: Final = _index_key(candidate_id) + return handle is not None and key is not None and handle.has_model_id(key) + + def _core_get_model_info(self, id: str) -> UntypedEntry | None: + handle: Final = self._live_handle() + key: Final = _index_key(id) + idx: Final[object] = None if handle is None or key is None else handle.id_position(key) + return None if idx is None else self._core_model_list[idx] # pyright: ignore[reportArgumentType, reportCallIssue] # a position, or what a caller wrote, read as Python reads it + + def _core_get_deployment(self, model_id: str) -> Deployment | None: + entry: Final = self._core_get_model_info(model_id) + return None if entry is None else Deployment(**entry) # pyright: ignore[reportArgumentType] # a stored entry + + def _core_get_deployment_by_model_group_name(self, model_group_name: str) -> Deployment | None: + handle: Final = self._live_handle() + key: Final = _index_key(model_group_name) + indices: Final[object] = None if handle is None or key is None else handle.named_positions(key) + model: Final = self._core_model_list[indices[0]] if indices else None # pyright: ignore[reportIndexIssue, reportArgumentType, reportCallIssue] # positions, or what a caller wrote, read as Python reads them + return None if model is None else Deployment(**model) # pyright: ignore[reportArgumentType] # a stored entry + + def _at(self, indices: object) -> tuple[UntypedEntry, ...]: + """The entries at ``indices`` (what an index holds), read as Python's list reads them.""" + return tuple(self._core_model_list[idx] for idx in indices) # pyright: ignore[reportGeneralTypeIssues, reportArgumentType, reportCallIssue] # positions, or what a caller wrote + + def _core_deployments_named(self, model_name: str) -> tuple[UntypedEntry, ...] | None: + handle: Final = self._live_handle() + key: Final = _index_key(model_name) + indices: Final[object] = None if handle is None or key is None else handle.named_positions(key) + return None if indices is None else self._at(indices) + + def team_public_model_names_of(self, team_id: object) -> frozenset[str]: + handle: Final = self._live_handle() + if handle is None or not isinstance(team_id, str): + return frozenset() + return frozenset(handle.team_public_names_of(team_id)) + + def _team_deployments_named(self, team_id: str, public_model_name: str) -> tuple[UntypedEntry, ...] | None: + handle: Final = self._live_handle() + team: Final = _index_key(team_id) + public: Final = _index_key(public_model_name) + indices: Final[object] = ( + None if handle is None or team is None or public is None else handle.team_positions(team, public) + ) + return None if indices is None else self._at(indices) + + # -- lookups (router-core `lookup.rs`): rows read at the stored entries -------------------- + + def _rows(self, rows: object) -> list[UntypedEntry]: # mutable-ok: a fresh list + """The stored entries router-core's lookup ``rows`` name, read as Python's list reads them: + a row it could not read is indexed here, and raises as the Python router's read raised.""" + entries: Final[list[UntypedEntry]] = [] # mutable-ok: filled per row + for row in cast(Sequence[Mapping[str, object]], rows): # cast-ok: router-core's lookup rows + if "raw" in row: + entries.append(self._core_model_list[row["raw"]]) # pyright: ignore[reportArgumentType, reportCallIssue] # what a caller wrote, read as Python reads it + elif "raw_positions" in row: + entries.extend(self._at(row["raw_positions"])) + elif "pattern" in row: + entries.append({**DeploymentTypedDict(**row["pattern"]), "model_name": row["name"]}) # pyright: ignore[reportCallIssue] # a pattern router's stored entry + else: + entries.append(self._row_entry(row)) + return entries + + def _row_entry(self, row: Mapping[str, object]) -> UntypedEntry: + held: Final[UntypedEntry] = self._core_model_list[int(str(row["at"]))] + name: Final = row.get("name") + named: Final = held if name is None else _renamed(held, str(name)) + group: Final = row.get("group") + if group is None: + return named + return _as_routing_group_row( + apply_routing_group_priority(self._core_routing_groups[str(group)], str(row["member"]), named) # pyright: ignore[reportArgumentType] # a stored entry + ) + + def _raise_at(self, raw: object) -> None: + """Read the list at the row router-core could not read it at, which raises.""" + if raw is not None: + _ = self._rows([raw]) + + def _list_scope(self, model_name: str | None, *, only: str | None = None) -> Mapping[str, object]: + """What ``get_model_list`` lists beyond the deployments: the aliases (read here, where an + alias item that is no mapping or lacks ``hidden`` raises as the Python router's read + does), the callable routing groups and the wildcard requests for ``model_name``.""" + from litellm.litellm_core_utils.get_llm_provider_logic import inferred_provider + + aliases: Final[list[Mapping[str, str]]] = [] # mutable-ok: filled per alias + alias_map: Final = self._core_model_group_alias or {} + if model_name is None or model_name in alias_map: + items: Final[Sequence[tuple[str, object]]] = ( + [(model_name, alias_map[model_name])] if model_name is not None else list(alias_map.items()) + ) + for model_alias, model_value in items: + if isinstance(model_value, str): + aliases.append({"alias": model_alias, "target": model_value}) + elif isinstance(model_value, dict): + item = RouterModelGroupAliasItem(**model_value) # pyright: ignore[reportCallIssue, reportArgumentType] # an alias item dict + if item["hidden"] is True and model_name is None: + continue + aliases.append({"alias": model_alias, "target": item["model"]}) + groups: Final = [ + {"name": name, "members": list(group.models)} + for name in (self._core_routing_groups or {}) + if (group := self.get_routing_group(name)) is not None + ] + wildcard: Final = ( + [model_name, f"{inferred_provider(model_name)}/{model_name}"] if model_name is not None else [] + ) + return {"aliases": aliases, "groups": groups, "wildcard": wildcard, "only": only} + + def _lookup(self) -> "RouterCoreHandle": + handle: Final = self._live_handle() + if handle is None: + raise RuntimeError("the host registry's lookups run in its Rust core") + return handle + + # -- deployments by name ---------------------------------------------------------------------- + + @staticmethod + def _core_should_include_deployment( + model_name: str, model: Mapping[str, object], team_id: str | None = None + ) -> bool: + info: Final = _model_info_of(model) + if team_id is not None and info.get("team_id") == team_id and model_name == info.get("team_public_model_name"): + return True + if model_name is not None and model.get("model_name") == model_name: # pyright: ignore[reportUnnecessaryComparison] # the core checks it + model_team_id: Final = info.get("team_id") + return team_id is None or model_team_id is None or model_team_id == team_id + return False + + @staticmethod + def _core_deployment_usable_by_team(model: Mapping[str, object] | Deployment, team_id: str | None) -> bool: + info: Final[object] = ( + model.model_info.model_dump() if isinstance(model, Deployment) else model.get("model_info") + ) + owner: Final = _MODEL_INFO_ADAPTER.validate_python(info).get("team_id") if isinstance(info, Mapping) else None + return team_may_use_deployment(owner, team_id) + + def _core_deployments_for_request( + self, model: str, request_kwargs: Mapping[str, object] + ) -> Sequence[Mapping[str, object]]: + """The deployments ``model`` names for this caller: alias, then team-first, then global, then + admin-across-teams (``_core_deployments_for_request``).""" + from litellm.litellm_core_utils.routing_common_utils import get_request_team_id, is_proxy_admin_request + + registered: Final = self._core_get_model_from_alias(model=model) or model + team_id: Final = get_request_team_id(request_kwargs) + deployments: Final = self._core_get_all_deployments(model_name=registered, team_id=team_id) + if deployments or team_id is not None or not is_proxy_admin_request(request_kwargs): + return deployments + return self._team_deployments_across_teams(registered) + + def _team_deployments_across_teams(self, model: str) -> list[UntypedEntry]: # mutable-ok: a fresh list + """Every team's deployments under public name ``model``, for a proxy admin calling without a team.""" + rows, teams = self._lookup().lookup_across_teams(model) + team_deployments: Final = self._rows(rows) + if teams > 1: + raise litellm_exceptions.BadRequestError( + message=( + f"Model name '{model}' matches deployments from multiple teams. " + "Specify the deployment ID directly to disambiguate." + ), + model=model, + llm_provider="", + ) + return team_deployments + + def _core_get_all_deployments( + self, model_name: str, model_alias: str | None = None, team_id: str | None = None + ) -> list[UntypedEntry]: # mutable-ok: a fresh list, as the core returns it + key: Final = _index_key(model_name) + return [] if key is None else self._rows(self._lookup().lookup_all(key, team_id, model_alias)) + + def get_routing_group(self, model_name: str) -> RoutingGroup | None: + """The routing group callable as ``model_name``: a deployment model_name or an alias of the + same name shadows it.""" + if not self._core_routing_groups: + return None + group: Final = self._core_routing_groups.get(model_name) + handle: Final = self._live_handle() + if ( + group is None + or (handle is not None and handle.named_positions(model_name) is not None) + or model_name in (self._core_model_group_alias or {}) + ): + return None + return group + + def _core_routing_group_has_alternatives(self, model_group: str | None) -> bool: + """A callable routing group whose member union spans more than one deployment.""" + if model_group is None: + return False + resolved: Final = self._core_get_model_from_alias(model=model_group) or model_group + group: Final = self.get_routing_group(resolved) + if group is None: + return False + handle: Final = self._lookup() + + def held(member: str) -> Sequence[object]: + return cast(Sequence[object], handle.named_positions(member) or ()) # cast-ok: positions, or junk + + return sum(len(held(member)) for member in group.models) > 1 + + def _core_get_model_list_from_model_alias( + self, model_name: str | None = None + ) -> list[UntypedEntry]: # mutable-ok: a fresh list + return self._rows( + self._lookup().lookup_model_list(model_name, None, self._list_scope(model_name, only="aliases")) + ) + + def _core_get_model_list_from_routing_groups(self, model_name: str | None = None) -> Sequence[UntypedEntry]: + """Callable routing groups as model-list rows: each member deployment under the group's name.""" + if model_name is None and self._routing_group_rows is not None: + return self._routing_group_rows + rows: Final = tuple( + self._rows(self._lookup().lookup_model_list(model_name, None, self._list_scope(model_name, only="groups"))) + ) + if model_name is None: + self._routing_group_rows = rows + return rows + + def _core_get_model_list( + self, model_name: str | None = None, team_id: str | None = None + ) -> list[UntypedEntry] | None: # mutable-ok: a fresh list, as get_model_list returns + if model_name is not None and _index_key(model_name) is None: + # No index holds it, so `_get_wildcard_deployments` reads it: what the pattern routers raise. + team_router: Final = self._core_team_pattern_routers.get(team_id) if team_id is not None else None + return [ + *self._core_pattern_router.get_deployments_by_pattern(model=model_name), + *(team_router.get_deployments_by_pattern(model=model_name) if team_router is not None else ()), + ] + return self._rows(self._lookup().lookup_model_list(model_name, team_id, self._list_scope(model_name))) + + def _core_wildcard_route_deployments(self, model: str, team_id: str | None) -> Sequence[Mapping[str, object]]: + from litellm.litellm_core_utils.get_llm_provider_logic import inferred_provider + + if not isinstance(model, str): # pyright: ignore[reportUnnecessaryIsInstance] # a caller's non-string model reads the base routers + team_router: Final = self._core_team_pattern_routers.get(team_id) if team_id is not None else None + team_owned: Final = team_router.get_deployments_by_pattern(model=model) if team_router is not None else () + return team_owned or self._core_pattern_router.get_deployments_by_pattern(model=model) + rows: Final = self._lookup().lookup_wildcard_route( + model, team_id, [model, f"{inferred_provider(model)}/{model}"] + ) + routed: Final = cast(Sequence[Mapping[str, object]], rows) # cast-ok: router-core's wildcard rows + return [cast(Mapping[str, object], row["pattern"]) for row in routed] # cast-ok: each carries its routed copy + + def _core_wildcard_routed_models(self, model: str, team_id: str | None) -> frozenset[str]: + team_router: Final = self._core_team_pattern_routers.get(team_id) if team_id is not None else None + team_owned: Final = ( + team_router.get_models_by_pattern(model=model) if team_router is not None else frozenset[str]() + ) + return team_owned or self._core_pattern_router.get_models_by_pattern(model=model) + + def _core_config_deployments(self) -> Iterator[Mapping[str, object]]: + """A generator, as the Python core's: the rows are read when iteration starts.""" + yield from self._rows(self._lookup().lookup_config()) + + # -- names ------------------------------------------------------------------------------------ + + def _core_get_model_names(self, team_id: str | None = None) -> list[str]: # mutable-ok: a fresh list + names, raw = self._lookup().lookup_model_names(team_id, self._list_scope(None)) + self._raise_at(raw) + return names + + def _core_get_model_ids( + self, model_name: str | None = None, exclude_team_models: bool = False + ) -> list[str]: # mutable-ok: a fresh list + if model_name is not None and _index_key(model_name) is None: + return [] + ids, raw = self._lookup().lookup_model_ids(model_name, exclude_team_models) + self._raise_at(raw) + return ids + + # -- resolving a name ------------------------------------------------------------------------- + + def _core_get_model_from_alias(self, model: str) -> str | None: + return resolve_model_group_alias(self._core_model_group_alias, model) + + def _core_routable_model_group(self, model_name: str) -> str: + target: Final = self._core_get_model_from_alias(model_name) + return target if target is not None else model_name + + def _core_is_recognized_model(self, model: str) -> bool: + handle: Final = self._live_handle() + key: Final = _index_key(model) + return ( + (handle is not None and key is not None and handle.has_model_name(key)) + or self._core_has_model_id(model) + or model in self._core_model_group_alias + or self.get_routing_group(model) is not None + ) + + def _core_get_model_group(self, id: str) -> list[UntypedEntry] | None: # mutable-ok: a fresh list + model_info: Final = self._core_get_model_info(id=id) + if model_info is None: + return None + return self._core_get_model_list(model_name=str(model_info["model_name"])) + + def _core_resolve_model_name_from_model_id( + self, model_id: str | None, custom_llm_provider: str | None = None + ) -> str | None: + if not model_id: + return None + handle: Final = self._lookup() + if handle.has_model_name(model_id): + return model_id + if self._core_has_model_id(model_id): + deployment: Final = self._core_get_deployment(model_id=model_id) + if deployment is not None and deployment.model_name: + return deployment.model_name + return model_id + listed, serving, raw = handle.lookup_name_serving(model_id, self._list_scope(None)) + self._raise_at(raw) + if not listed: + return None + if serving is not None: + return serving + if self._core_is_recognized_model(model_id): + return model_id + if custom_llm_provider is None: + return None + if self._core_pattern_router.routes_to_provider(model_id, custom_llm_provider): + return model_id + provider_prefixed: Final = f"{custom_llm_provider}/{model_id}" + if not model_id.startswith(f"{custom_llm_provider}/") and self._core_pattern_router.routes_to_provider( + provider_prefixed, custom_llm_provider + ): + return provider_prefixed + return None + + def _get_deployment_by_litellm_model(self, model: object) -> list[UntypedEntry]: # mutable-ok: a fresh list + handle: Final = self._live_handle() + if handle is None or not isinstance(model, str): + return [] + return list(self._at(handle.litellm_model_positions(model))) + + # -- credentials ------------------------------------------------------------------------------ + + @staticmethod + def _core_is_deployment_blocked(deployment: Deployment) -> bool: + model_info: Final[object] = getattr(deployment, "model_info", None) + return model_info is not None and getattr(model_info, "blocked", None) is True + + def _core_get_deployment_credentials(self, model_id: str) -> UntypedEntry | None: + deployment: Final = self._core_get_deployment(model_id=model_id) + if deployment is None or self._core_is_deployment_blocked(deployment): + return None + return CredentialLiteLLMParams.model_validate( + deployment.litellm_params.model_dump(exclude_none=True) + ).model_dump(exclude_none=True) + + def _get_model_group_deployment_usable_by_team( + self, model_group_name: str, team_id: str | None + ) -> Deployment | None: + if _index_key(model_group_name) is None: + return None + position, raw = self._lookup().lookup_group_usable_by_team(model_group_name, team_id) + self._raise_at(raw) + return None if position is None else Deployment(**self._core_model_list[position]) # pyright: ignore[reportArgumentType] # a stored entry + + def _get_team_public_name_deployment(self, model_id: str, team_id: str | None) -> Deployment | None: + if team_id is None or _index_key(model_id) is None: + return None + position, raw = self._lookup().lookup_team_public_name(model_id, team_id) + self._raise_at(raw) + return None if position is None else Deployment(**self._core_model_list[position]) # pyright: ignore[reportArgumentType] # a stored entry + + def _get_wildcard_deployment_usable_by_team(self, model_id: str, team_id: str | None) -> Deployment | None: + team_router: Final = self._core_team_pattern_routers.get(team_id) if team_id is not None else None + team_models: Final = team_router.route(model_id) if team_router else None + global_models: Final = tuple( + model + for model in (self._core_pattern_router.route(model_id) or ()) + if self._core_deployment_usable_by_team(model, team_id) + ) + candidates: Final = team_models or global_models + if not candidates: + return None + return Deployment(**candidates[0]) # a pattern router holds stored entry dicts + + def _core_get_credential_deployment(self, model_id: str, team_id: str | None = None) -> Deployment | None: + deployment: Final = ( + self._core_get_deployment(model_id=model_id) + or self._get_model_group_deployment_usable_by_team(model_group_name=model_id, team_id=team_id) + or self._get_team_public_name_deployment(model_id=model_id, team_id=team_id) + or self._get_wildcard_deployment_usable_by_team(model_id=model_id, team_id=team_id) + ) + if deployment is None or self._core_is_deployment_blocked(deployment): + return None + return deployment + + def _core_get_deployment_credentials_with_provider( + self, model_id: str, team_id: str | None = None + ) -> UntypedEntry | None: + deployment: Final = self._core_get_credential_deployment(model_id=model_id, team_id=team_id) + if deployment is None: + return None + base: Final = CredentialLiteLLMParams.model_validate( + deployment.litellm_params.model_dump(exclude_none=True) + ).model_dump(exclude_none=True) + credential_name: Final = deployment.litellm_params.litellm_credential_name + resolved: Final[UntypedEntry] = dict(base) + if credential_name is not None: + credential_values: Final = CredentialAccessor.get_credential_values(credential_name) + if not credential_values: + verbose_router_logger.warning("Credential '%s' not found in credential_list", credential_name) + resolved.update(credential_values) + _ = resolved.pop("litellm_credential_name", None) + model: Final = deployment.litellm_params.model + resolved["model"] = model + resolved["custom_llm_provider"] = ( + deployment.litellm_params.custom_llm_provider or (model.split("/")[0] if "/" in model else None) or "openai" + ) + return resolved + + # -- routable names --------------------------------------------------------------------------- + + def _routable_deployments(self, model_name: str, team_id: str | None) -> tuple[UntypedEntry, ...]: + return tuple(self._rows(self._lookup().lookup_routable(self._core_routable_model_group(model_name), team_id))) + + def _core_team_ptu_ceiling(self, team_id: str, model_group: str) -> "PTUTeamCeiling | None": + """The per-minute ceiling ``team_id``'s PTU shares set on the group serving ``model_group``, + else None (``team_ptu_ceiling``), kept until the model list next changes.""" + from litellm.litellm_core_utils.ptu_shares import ptu_share_teams, team_ptu_ceiling + + listed: Final = tuple(self._core_get_model_list() or ()) + if self._ptu_cache is None: + self._ptu_cache = (ptu_share_teams(listed), {}) + share_teams, ceilings = self._ptu_cache + if team_id not in share_teams: + return None + key: Final = (team_id, model_group) + if key not in ceilings: + ceilings[key] = team_ptu_ceiling( + listed, + self._core_model_list, + team_id, + model_group, + lambda name: self._core_wildcard_route_deployments(name, team_id), + ) + return ceilings[key] + + def _get_client( + self, deployment: Mapping[str, object], kwargs: Mapping[str, object], client_type: str | None = None + ) -> object: + """The SDK client cached in the routing cache for a deployment (``_get_client``); a pod + discovery deployment has none. ``max_parallel_requests`` is the deployment's slot limit, + built on first use.""" + from litellm.litellm_core_utils.core_helpers import get_parent_otel_span_from_kwargs + from litellm.router_rust.core import RustCoreUnsupported + + # `deployment["model_info"]["id"]`: the Python router reads the id first, raising for a + # deployment without one whatever the client type. + _ = deployment["model_info"]["id"] # pyright: ignore[reportIndexIssue] # read as the Python router reads it + if client_type == "max_parallel_requests": + limit_lookup: Final = self.max_parallel_requests_lookup + if limit_lookup is None: + raise RustCoreUnsupported( + "registry().client(client_type='max_parallel_requests') is not served by the Rust router core" + ) + return limit_lookup(deployment, kwargs) + if _params_of(deployment).get("kubernetes_pod_discovery") is True: + return None + model_id: Final = _model_info_of(deployment).get("id") + span: Final = get_parent_otel_span_from_kwargs(dict(kwargs)) + stream: Final = "_stream" if kwargs.get("stream") is True else "" + with service_target(ROUTER_CLIENTS_TARGET): + if client_type == "async": + return self._core_cache.get_cache( + key=f"{model_id}{stream}_async_client", local_only=True, parent_otel_span=span + ) # pyright: ignore[reportUnknownMemberType] # DualCache's reads are untyped + return self._core_cache.get_cache(key=f"{model_id}{stream}_client", parent_otel_span=span) # pyright: ignore[reportUnknownMemberType] # see above + + def _core_get_search_tool_support(self, model_name: str, team_id: str | None = None) -> tuple[bool | None, ...]: + """Per routable deployment: its ``supports_search_tool`` when a bool, else True when its route + builds an Anthropic Messages request, else None (``get_search_tool_support``).""" + return tuple(_search_tool_support(deployment) for deployment in self._routable_deployments(model_name, team_id)) + + def _core_get_candidate_model_ids_for_route(self, model: str, team_id: str | None = None) -> frozenset[str]: + """The deployment ids that could serve ``model`` for ``team_id``: the alias target, then a + routing group, then a team public name or wildcard route for a name that is no model + group, then the model group (``get_candidate_model_ids_for_route``). Read-only.""" + from litellm.litellm_core_utils.get_llm_provider_logic import inferred_provider + + resolved: Final = self._core_get_model_from_alias(model=model) or model + group: Final = self.get_routing_group(resolved) if self._core_routing_groups else None + ids, raw = self._lookup().lookup_candidate_ids( + resolved, + team_id, + None if group is None else list(group.models), + [resolved, f"{inferred_provider(resolved)}/{resolved}"], + ) + self._raise_at(raw) + return frozenset(ids) + + def _core_team_model_has_alternatives(self, deployment_id: str) -> bool: + """Whether a team deployment's public name has another unpaused deployment of that team.""" + found, raw = self._lookup().lookup_team_model_has_alternatives(deployment_id) + self._raise_at(raw) + return found + + def _core_resolve_served_model_name(self, model_name: str, served_model: str) -> str: + """``served_model`` when ``model_name`` routes to it (the alias target, a routing group + member, or the name itself), else ``model_name``.""" + target: Final = self._core_get_model_from_alias(model_name) or model_name + group: Final = self.get_routing_group(target) + members: Final = group.models if group is not None else () + return served_model if served_model == target or served_model in members else model_name + + def _core_resolved_litellm_models( + self, model_name: str, team_id: str | None = None, *, use_base_model: bool = False + ) -> tuple[str, ...]: + """The provider model strings a call to ``model_name`` can be served by (``resolved_litellm_models``).""" + listed: Final = self._core_get_model_list(model_name=model_name, team_id=team_id) or () + if use_base_model: + return tuple( + backend for entry in listed if (backend := cost_map_backend_model(Deployment.model_validate(entry))) + ) + return tuple(model for entry in listed if isinstance(model := _params_of(entry).get("model"), str) and model) + + def _core_get_routable_deployments( + self, model: str, request_kwargs: Mapping[str, object] = MappingProxyType({}), call_type: str | None = None + ) -> tuple[Mapping[str, object], ...]: + """The deployments routing would consider for ``request_kwargs``, as the Rust core resolves + them, the endpoint of ``call_type`` included.""" + lookup: Final = self.routable_lookup + if lookup is None: + from litellm.router_rust.core import RustCoreUnsupported + + raise RustCoreUnsupported("registry().get_routable_deployments is not served by the Rust router core") + return tuple(lookup(model, request_kwargs, call_type)) + + def _core_get_routable_upstream_model(self, model_name: str, team_id: str | None = None) -> str | None: + deployment: Final = next(iter(self._routable_deployments(model_name, team_id)), None) + if deployment is None: + return None + model: Final = _params_of(deployment).get("model") + return None if model is None else str(model) + + def _core_map_team_model( + self, team_model_name: str | None, team_id: str, custom_llm_provider: str | None = None + ) -> str | None: + """A team's public name to its group, or a deployment for a model-less call (``map_team_model``).""" + models: Final = self._core_get_model_list(model_name=team_model_name, team_id=team_id) + if not models: + return None + if team_model_name is None and custom_llm_provider is not None: + return self._team_model_for_provider(models, team_id, custom_llm_provider) + for model in models: + info = _model_info_of(model) + if info.get("team_id") == team_id: + if team_model_name is None: + public = info.get("team_public_model_name") or model.get("model_name") + return None if public is None else str(public) + return team_model_name + return None + + @staticmethod + def _team_model_for_provider(models: Sequence[UntypedEntry], team_id: str, custom_llm_provider: str) -> str | None: + def public_name(deployment: UntypedEntry) -> str | None: + info: Final = _model_info_of(deployment) + if info.get("team_id") != team_id: + return None + name: Final = info.get("team_public_model_name") + return name if isinstance(name, str) and name else str(deployment["model_name"]) + + def provider(deployment: UntypedEntry) -> object: + return provider_for_generic_call(_params_of(deployment)) # pyright: ignore[reportArgumentType] # the deployment's params + + matching: Final = next( + (m for m in models if public_name(m) is not None and provider(m) == custom_llm_provider), None + ) + if matching is None: + return None + name: Final = public_name(matching) + if name is None: + return None + if all(provider(m) == custom_llm_provider for m in models if public_name(m) == name): + return name + matched_id: Final = _model_info_of(matching).get("id") + return matched_id if isinstance(matched_id, str) else None + + # -- access groups ---------------------------------------------------------------------------- + + def _core_get_model_access_groups( + self, model_name: str | None = None, model_access_group: str | None = None, team_id: str | None = None + ) -> Mapping[str, Sequence[str]]: + use_cache: Final = model_name is None and model_access_group is None and team_id is None + if use_cache and self._access_groups_cache is not None: + return self._access_groups_cache + pairs, raw = self._lookup().lookup_access_groups( + model_name, model_access_group, team_id, self._list_scope(model_name) + ) + self._raise_at(raw) + groups: Final[dict[str, list[str]]] = {g: list(n) for g, n in pairs} # mutable-ok: as a defaultdict + if use_cache: + self._access_groups_cache = groups + return groups + + def _core_is_model_access_group_for_wildcard_route(self, model_access_group: str) -> bool: + groups: Final = self._core_get_model_access_groups(model_access_group=model_access_group) + return any( + self._core_pattern_router.route(request=model) is not None for model in groups.get(model_access_group, []) + ) + + +# -- the registry() view ------------------------------------------------------------------------- + + +class RustRegistry(RouterRegistry): + """``registry()`` of the Rust core: ``RouterRegistry``'s queries over the host registry.""" + + __slots__ = () + + def includes( + self, + deployment: "DeploymentTypedDict | dict[str, object]", # mutable-ok: the RouterRegistry signature + *, + model_name: str | None = None, + team_id: str | None = None, + ) -> bool: + hosted: Final = self._router + if not isinstance(hosted, HostRegistry): + return super().includes(deployment, model_name=model_name, team_id=team_id) + if model_name is None: + return HostRegistry._core_deployment_usable_by_team(deployment, team_id) # pyright: ignore[reportPrivateUsage, reportArgumentType] # the view reads its own registry + return HostRegistry._core_should_include_deployment(model_name=model_name, model=deployment, team_id=team_id) # pyright: ignore[reportPrivateUsage, reportArgumentType] # see above + + +def hosted_registry_view(hosted: HostRegistry) -> RouterRegistry: + """The ``registry()`` view over ``hosted``.""" + return RustRegistry(hosted) # pyright: ignore[reportArgumentType] # the host registry answers every _core_ query the view makes + + +# The registry indexes a deprecated attribute write replaces (``router.model_names = ...`` reaches +# the host registry as ``_core_model_names``), by the field router-core reads each one as. +INDEX_ATTRIBUTES: Final[Mapping[str, str]] = MappingProxyType( + { + "_core_model_names": "model_names", + "_core_model_name_to_deployment_indices": "model_name", + "_core_team_model_to_deployment_indices": "team_model", + } +) + + +def _positions(indices: object, length: int) -> list[int]: # mutable-ok: a JSON list + """The model-list positions ``indices`` read, as Python's ``model_list[index]`` reads them: a + negative one counts from the end, and one before the start becomes ``length``, a position past + the end, so the lookup raises ``IndexError`` as Python's does.""" + if not isinstance(indices, Iterable) or isinstance(indices, str | bytes): + return [] + return [ + index if index >= 0 else (index + length if index + length >= 0 else length) + for index in indices + if isinstance(index, int) and not isinstance(index, bool) + ] + + +def index_document(hosted: HostRegistry, written: Iterable[str]) -> Mapping[str, object]: + """The written registry indexes, as the host registry holds them now, in the shape router-core + reads (``IndexOverrides``): names a lookup finds, and the deployment positions each model name + (or ``(team_id, public name)`` pair) indexes. Entries Python could not read as such are left out.""" + document: Final[dict[str, object]] = {} # mutable-ok: built here + length: Final = len(hosted._core_model_list) # pyright: ignore[reportPrivateUsage] # positions read against the live list + for attribute in written: + value: object = getattr(hosted, attribute, None) + field: str = INDEX_ATTRIBUTES[attribute] + if field == "model_names": + names = value if isinstance(value, Iterable) and not isinstance(value, str | bytes | Mapping) else () + document[field] = sorted({name for name in names if isinstance(name, str)}) + elif field == "model_name": + entries = value if isinstance(value, Mapping) else {} + document[field] = { + name: _positions(indices, length) for name, indices in entries.items() if isinstance(name, str) + } + else: + pairs = value if isinstance(value, Mapping) else {} + document[field] = [ + [key[0], key[1], _positions(indices, length)] + for key, indices in pairs.items() + if isinstance(key, tuple) and len(key) == 2 and all(isinstance(part, str) for part in key) + ] + return document + + +DEFAULT_DEPLOYMENT_FAILING_ID: Final = "__litellm_default_deployment_failing_call__" +"""The id router-core serves a ``default_deployment`` under when every call to it fails.""" + +DEFAULT_DEPLOYMENT_COPIED_ID: Final = "__litellm_default_deployment_copied__" +"""The id router-core serves a callable ``default_deployment`` under when its ``model_info.id`` is +not a string: the host picks and calls Python's copy of it, so its id stays what was assigned.""" + + +def default_deployment_copy(value: object, model: str) -> object: + """``_try_early_resolve_deployments_for_model_not_in_names``' copy of the default deployment + for ``model``, raising what it raises for a value it cannot copy.""" + updated: Final = value.copy() # pyright: ignore[reportAttributeAccessIssue, reportUnknownMemberType, reportUnknownVariableType] # any assigned value, read as Python reads it + updated["litellm_params"] = value["litellm_params"].copy() # pyright: ignore[reportIndexIssue, reportUnknownMemberType] # same + updated["litellm_params"]["model"] = model # pyright: ignore[reportUnknownMemberType] # same + return updated # pyright: ignore[reportUnknownVariableType] # same + + +def default_deployment_copy_error(value: object, model: str) -> BaseException: + """The error the default deployment's copy raises for ``model``.""" + try: + _ = default_deployment_copy(value, model) + except Exception as raised: # noqa: BLE001 # the error the pick raises + return raised + return RuntimeError(f"router-core: the default deployment copies for {model!r}") + + +def default_deployment_call_error(value: object) -> BaseException | None: + """The error every call to the default deployment raises, from its first reads + (``_update_kwargs_with_deployment``'s ``deployment["model_name"]``, ``_get_client``'s + ``deployment["model_info"]["id"]``); None when they read.""" + try: + served: Final = default_deployment_copy(value, "") + _ = served["model_name"] # pyright: ignore[reportIndexIssue] # read as Python reads it + _ = served["model_info"]["id"] # pyright: ignore[reportIndexIssue] # same + except Exception as raised: # noqa: BLE001 # the error the call raises + return raised + return None + + +def served_default_deployment(value: object) -> tuple[Mapping[str, object] | None, bool]: + """The ``default_deployment`` as router-core serves it, and whether resolving to it raises + (a value Python cannot copy). A deployment every call to fails is served under + ``DEFAULT_DEPLOYMENT_FAILING_ID``, one whose id is not a string under + ``DEFAULT_DEPLOYMENT_COPIED_ID``; names that are not strings are served as text.""" + if value is None: + return None, False + try: + copied: Final = default_deployment_copy(value, "") + except Exception: # noqa: BLE001 # any copy error: the pick raises it + return None, True + entry: Final = as_str_mapping(copied) or {} + model_info: Final = dict(as_str_mapping(entry.get("model_info")) or {}) + failing: Final = default_deployment_call_error(value) is not None + assigned_id: Final = model_info.get("id") + model_info["id"] = ( + DEFAULT_DEPLOYMENT_FAILING_ID + if failing + else assigned_id + if isinstance(assigned_id, str) + else DEFAULT_DEPLOYMENT_COPIED_ID + ) + name: Final = entry.get("model_name") + return { + **entry, + "model_name": name if isinstance(name, str) else str(name or ""), + "litellm_params": dict(as_str_mapping(entry.get("litellm_params")) or {}), + "model_info": model_info, + }, False diff --git a/litellm/router_rust/surface/replay.py b/litellm/router_rust/surface/replay.py new file mode 100644 index 00000000000..6c770e07883 --- /dev/null +++ b/litellm/router_rust/surface/replay.py @@ -0,0 +1,30 @@ +"""The Rust core's live routers, whose deployments a price data reload re-registers in +``litellm.model_cost`` until they retire (``_live_routers`` / ``_replay_live_router_model_cost``).""" + +from __future__ import annotations + +import weakref +from typing import TYPE_CHECKING, Final + +from litellm.utils import add_live_deployment_replay + +if TYPE_CHECKING: + from litellm.router_rust.surface.registry import HostRegistry + +_live_registries: Final[weakref.WeakSet[HostRegistry]] = weakref.WeakSet() + + +def track_live_registry(registry: HostRegistry) -> None: + _live_registries.add(registry) + + +def untrack_live_registry(registry: HostRegistry) -> None: + _live_registries.discard(registry) + + +def _replay_live_registries() -> None: + for registry in tuple(_live_registries): + registry.replay_model_cost_registrations() + + +add_live_deployment_replay(_replay_live_registries) diff --git a/litellm/router_rust/surface/retries.py b/litellm/router_rust/surface/retries.py new file mode 100644 index 00000000000..4a9c6429ffa --- /dev/null +++ b/litellm/router_rust/surface/retries.py @@ -0,0 +1,295 @@ +""" +The retry members ``RustCore`` serves: ``should_retry_this_error`` and +``get_num_retries_from_retry_policy``. router-core decides both; this module only describes the +Python exception the way an attempt failure is described and reads the answer back. +""" + +import dataclasses +import operator +from collections.abc import Iterable, Mapping, Sequence +from typing import TYPE_CHECKING, Final + +from litellm.types.router import RetryPolicy + +# The most retries a request's `num_retries` asks router-core for (its count is a u32). +_MOST_RETRIES: Final = 2**32 - 1 + +if TYPE_CHECKING: + from litellm.router_rust.core import Resolved + from litellm.rust_bridge._native import ( # pyright: ignore[reportPrivateUsage] # the stub's handle + RouterCoreHandle, + ) + + +def resolved_retry_policy(value: object) -> RetryPolicy | None: + """The router's ``retry_policy`` as the Python router's constructor keeps it: a dict is read + into a ``RetryPolicy``, a ``RetryPolicy`` is kept, anything else is no policy.""" + if isinstance(value, RetryPolicy): + return value + if isinstance(value, Mapping): + fields: Final[Mapping[str, int | None]] = value # pyright: ignore[reportUnknownVariableType] # the policy's fields + return RetryPolicy(**fields) # pyright: ignore[reportArgumentType] # validated by the model + return None + + +def _is_policy(entry: object) -> bool: + return isinstance(entry, RetryPolicy | Mapping) + + +def reads_as_policy(entry: object) -> bool: + """Whether ``resolve_retry_policy`` reads ``entry`` without raising: no entry, a ``RetryPolicy``, + or a mapping ``RetryPolicy(**entry)`` accepts.""" + if entry is None or isinstance(entry, RetryPolicy): + return True + if not isinstance(entry, Mapping): + return False + try: + _ = RetryPolicy(**entry) # pyright: ignore[reportUnknownArgumentType] # validated by the model + except Exception: # noqa: BLE001 # any error the lookup would raise + return False + return True + + +def is_policy_per_group(policies: object) -> bool: + """Whether ``model_group_retry_policy`` is what router-core's settings hold: a mapping of + group names to policies (or nothing).""" + if policies is None: + return True + if not isinstance(policies, Mapping): + return False + entries: Final[Mapping[object, object]] = policies # pyright: ignore[reportUnknownVariableType] # read as objects + return all(isinstance(group, str) and _is_policy(entry) for group, entry in entries.items()) + + +def readable_retry_policy(value: object) -> RetryPolicy | None: + """A router ``retry_policy`` written as a plain attribute, as router-core's settings hold it. + The attribute keeps what was written; a value the lookup raises for reaches router-core as + ``retry_policy_raises`` instead.""" + return resolved_retry_policy(value) if reads_as_policy(value) else None + + +@dataclasses.dataclass(frozen=True, slots=True) +class GroupPolicies: + """A ``model_group_retry_policy`` as router-core reads it, resolved the way + ``resolve_retry_policy`` reads it for each group the call may evaluate. + + ``document`` holds a group's policy where the lookup finds one, and an empty policy (which + grants nothing) where it finds ``None``; a group it does not find is absent, so the router's + ``retry_policy`` applies. ``raising`` names the groups whose lookup raises (an ``in``/``[]`` + the value does not support, or an entry that is not a policy, whose fields + ``get_num_retries_from_retry_policy`` then fails to read): the Python retry loop raises that + error in place of the attempt's on a hop's first failure. + """ + + document: Mapping[str, object] + raising: frozenset[str] + + +def _looked_up(policies: object, group: str) -> tuple[bool, object]: + """``(found, entry)`` as ``resolve_retry_policy`` looks ``group`` up; raises as it raises.""" + found: Final = group in policies # pyright: ignore[reportOperatorIssue] # any value, as the Python router reads it + return (True, policies[group]) if found else (False, None) # pyright: ignore[reportIndexIssue, reportUnknownVariableType] # see above + + +def _policy_document(entry: object) -> object: + """A readable entry as router-core reads it: no entry grants nothing, a mapping is read the way + ``RetryPolicy(**entry)`` reads it.""" + if entry is None: + return {} + if isinstance(entry, Mapping): + return RetryPolicy(**entry) # pyright: ignore[reportUnknownArgumentType] # checked by reads_as_policy + return entry + + +def group_policies(policies: object, groups: Iterable[str]) -> GroupPolicies: + """``policies`` resolved for ``groups`` (a mapping is read by its own keys).""" + if policies is None: + return GroupPolicies({}, raising=frozenset()) + if isinstance(policies, Mapping): + entries: Final[Mapping[object, object]] = policies # pyright: ignore[reportUnknownVariableType] # read as objects + # A key that is no string never matches a group name. + named: Final = {group: entry for group, entry in entries.items() if isinstance(group, str)} + readable: Final = {group: entry for group, entry in named.items() if reads_as_policy(entry)} + return GroupPolicies( + {group: _policy_document(entry) for group, entry in readable.items()}, + raising=frozenset(named) - frozenset(readable), + ) + document: Final[dict[str, object]] = {} # mutable-ok: built here; router-core reads it as plain JSON + raising: Final[set[str]] = set() # mutable-ok: the groups whose lookup raises + for group in groups: + try: + found, entry = _looked_up(policies, group) + except Exception: # noqa: BLE001 # the Python retry loop raises it on a hop's first failure + raising.add(group) + continue + if not found: + continue + if not reads_as_policy(entry): + raising.add(group) + continue + document[group] = _policy_document(entry) + return GroupPolicies(document, raising=frozenset(raising)) + + +def static_group_policies(policies: object) -> object: + """The router's ``model_group_retry_policy`` as router-core's static settings hold it. A value + that is not a mapping of policies is resolved per request instead (``group_policies``).""" + return group_policies(policies, ()).document if isinstance(policies, Mapping) else None # pyright: ignore[reportUnknownArgumentType] # read as objects + + +def raising_groups(request: Mapping[str, object], resolved: GroupPolicies) -> Sequence[str]: + """The groups whose ``model_group_retry_policy`` lookup raises on a hop's first failure, unless + the retry loop never reads it: the request opted out with ``num_retries=0``. A streamed Anthropic + messages call is no exception: only the stream's own retry rules read such an entry as no policy + (``_anthropic_messages_resolved_retry_policy``); a failure before the stream opened reaches the + retry loop.""" + if not resolved.raising or request.get("num_retries") == 0: + return () + return sorted(resolved.raising) + + +def policy_lookup_error(original: BaseException, model_group: str, retry_policy: object, policies: object) -> Exception: + """The error ``get_num_retries_from_retry_policy`` raises for ``model_group`` after ``original`` + failed (the Python retry loop raises it in the attempt's place, chained to it); a lookup that + unexpectedly reads is reported as such.""" + from litellm.litellm_core_utils.get_retry_from_policy import get_num_retries_from_retry_policy + + traceback: Final = original.__traceback__ + try: + # Looked up while `original` is being handled, as in the retry loop's `except` block, so + # the error chains to it the way Python chains it. + try: + raise original + except BaseException: # noqa: BLE001 # re-raised only to be handled here + _ = get_num_retries_from_retry_policy( + exception=original, # pyright: ignore[reportArgumentType] # any exception the hop raised + retry_policy=retry_policy, # pyright: ignore[reportArgumentType] # the router's resolved policy + model_group=model_group, + model_group_retry_policy=policies, # pyright: ignore[reportArgumentType] # read as the Python router reads it + ) + except Exception as raised: # noqa: BLE001 # the error the hop ends with + return raised + finally: + original.__traceback__ = traceback # rebind-ok: put back as the hop left it + return RuntimeError(f"the model_group_retry_policy lookup for {model_group!r} did not raise") + + +def router_policy_retries( + exception: Exception, retry_policy: object, policies: object, model_group: str | None +) -> int | None: + """``get_num_retries_from_retry_policy`` for a router ``model_group_retry_policy`` that is not + a mapping of policies, read as the Python router reads it (raising as it raises).""" + from litellm.litellm_core_utils.get_retry_from_policy import get_num_retries_from_retry_policy + + return get_num_retries_from_retry_policy( + exception=exception, + retry_policy=retry_policy, # pyright: ignore[reportArgumentType] # the router's resolved policy + model_group=model_group, + model_group_retry_policy=policies, # pyright: ignore[reportArgumentType] # read as the Python router reads it + ) + + +def request_retry_fields(value: object) -> Mapping[str, object]: + """A request's ``num_retries`` as router-core reads it. ``async_function_with_retries`` never + casts it: it compares it with 0, and counts with ``range`` only once it is positive.""" + if value is None: + return {} + if isinstance(value, int): + count: Final = int(value) + return {"num_retries": min(count, _MOST_RETRIES)} if count >= 0 else {"uncounted_retries": "negative"} + try: + if value == 0: + return {"num_retries": 0} + if not value > 0: # pyright: ignore[reportOperatorIssue] # any value, compared as the retry loop compares it + return {"uncounted_retries": "not_positive"} + return {"num_retries": min(operator.index(value), _MOST_RETRIES)} # pyright: ignore[reportArgumentType] # see above + except Exception: # noqa: BLE001 # the retry loop raises it on a hop's first failure + return {"uncounted_retries": "raises"} + + +def retry_count_error(original: BaseException, value: object) -> Exception: + """The error the retry loop raises reading a request's ``num_retries`` after ``original`` + failed, chained to it; a value that unexpectedly reads is reported as such.""" + traceback: Final = original.__traceback__ + try: + try: + raise original + except BaseException: # noqa: BLE001 # re-raised only to be handled here + if value > 0: # pyright: ignore[reportOperatorIssue] # `if num_retries > 0:` + _ = range(value) # pyright: ignore[reportArgumentType] # `for current_attempt in range(num_retries):` + except Exception as raised: # noqa: BLE001 # the error the hop ends with + return raised + finally: + original.__traceback__ = traceback # rebind-ok: put back as the hop left it + return RuntimeError(f"the request's num_retries {value!r} did not raise") + + +def _count(deployments: object) -> int: + return len(deployments) if isinstance(deployments, list) else 0 # pyright: ignore[reportUnknownArgumentType] # only the length is read + + +def _filled(fallbacks: object) -> bool: + return isinstance(fallbacks, Sequence) and len(fallbacks) > 0 # pyright: ignore[reportUnknownArgumentType] # only the length is read + + +class RustRetries: + """Retry members read from the Rust core's handle.""" + + _handle: "RouterCoreHandle" + _resolved: "Resolved" + + @property + def _written_retry_policy(self) -> object: + """The ``retry_policy`` attribute as written (``Resolved`` holds what router-core reads).""" + raise NotImplementedError + + def should_retry_this_error( + self, + error: Exception, + healthy_deployments: list[object] | None = None, # mutable-ok: the Python router's signature + all_deployments: list[object] | None = None, # mutable-ok: same + context_window_fallbacks: list[object] | None = None, # mutable-ok: same + content_policy_fallbacks: list[object] | None = None, # mutable-ok: same + regular_fallbacks: list[object] | None = None, # mutable-ok: same + ) -> bool: + """``True`` when the error is retried in the same group; otherwise the error is raised.""" + from litellm.router_rust.core import attempt_failure + + retried: Final = self._handle.should_retry( + dict(attempt_failure(error)), + (_count(healthy_deployments), _count(all_deployments)), + (context_window_fallbacks is not None, content_policy_fallbacks is not None, _filled(regular_fallbacks)), + ) + if not retried: + raise error + return True + + def get_num_retries_from_retry_policy(self, exception: Exception, model_group: str | None = None) -> int | None: + from litellm.router_rust.core import attempt_failure + + resolved: Final = self._resolved + written: Final = self._written_retry_policy + if not is_policy_per_group(resolved.model_group_retry_policy) or not reads_as_policy(written): + return router_policy_retries(exception, written, resolved.model_group_retry_policy, model_group) + # A dict lookup hashes any key and only a string can match; the native call takes strings alone. + _ = hash(model_group) + return self._handle.retry_policy_retries( + dict(attempt_failure(exception)), model_group if isinstance(model_group, str) else None + ) + + +__all__ = [ + "GroupPolicies", + "RustRetries", + "group_policies", + "is_policy_per_group", + "policy_lookup_error", + "raising_groups", + "readable_retry_policy", + "reads_as_policy", + "request_retry_fields", + "resolved_retry_policy", + "retry_count_error", + "router_policy_retries", + "static_group_policies", +] diff --git a/litellm/router_rust/surface/settings.py b/litellm/router_rust/surface/settings.py new file mode 100644 index 00000000000..c8635593de5 --- /dev/null +++ b/litellm/router_rust/surface/settings.py @@ -0,0 +1,236 @@ +""" +The Rust core's settings semantics on the host side: how ``ChangeSet.settings`` is read (the rules +of ``update_settings``), the ``settings().as_dict()`` view (``get_settings``), the router attributes +the Python core keeps as plain attributes, deleted settings, and the construction-time side effects +of the logging arguments. + +The rules are the Python router core's, ported (never imported). Settings change only through a +changeset; a direct attribute write is turned into one by ``RustCore.__setattr__``. +""" + +import logging +from collections.abc import Mapping +from dataclasses import dataclass, field +from types import MappingProxyType +from typing import TYPE_CHECKING, Final + +from litellm._logging import verbose_router_logger +from litellm.constants import RUNTIME_UPDATABLE_ROUTER_SETTINGS +from litellm.litellm_core_utils.routing_groups import validate_routing_strategy +from litellm.types.router import RetryPolicy, RoutingStrategy + +if TYPE_CHECKING: + from litellm.router_config import RouterConfig + from litellm.router_python_host import PythonHost + +MISSING: Final = object() + +# `update_settings` casts these with `int(value)`. +_INT_SETTINGS: Final = frozenset({"timeout", "num_retries", "retry_after", "allowed_fails", "cooldown_time"}) +CONTINUATION_SETTING: Final = "enable_mid_stream_fallback_continuation" + +# Settings the Python core keeps as plain router attributes and `update_settings` sets as given, +# with their constructor defaults. ``max_retries`` has none: it exists once a changeset sets it. +PLAIN_DEFAULTS: Final[Mapping[str, object]] = MappingProxyType( + { + "retry_policy": None, + "model_group_retry_policy": {}, + "enable_weighted_failover": False, + "fallback_on_unsupported_params": True, + } +) +# `get_settings`: the settings it lists, in its order (`routing_groups` comes last). +_VIEW_ORDER: Final = ( + "routing_strategy_args", + "routing_strategy", + "allowed_fails", + "cooldown_time", + "num_retries", + "timeout", + "max_retries", + "retry_after", + "fallbacks", + "context_window_fallbacks", + "model_group_retry_policy", + "retry_policy", + "model_group_alias", + "enable_weighted_failover", + "enable_mid_stream_fallback_continuation", + "enable_tag_filtering", + "tag_routing_prefix", + "fallback_on_unsupported_params", +) + +# Router attributes that are their constructor argument as given, else this default. Arguments +# the Rust core declines are only ever at their default here. +_ARGUMENT_ATTRIBUTES: Final[Mapping[str, object]] = MappingProxyType( + { + "set_verbose": False, + "debug_level": "INFO", + "client_ttl": 3600, + "stream_timeout": None, + "assistants_config": None, + "ignore_invalid_deployments": False, + "alerting_config": None, + "auto_router_capability_limit": None, + "deployment_affinity_ttl_seconds": 3600, + "session_affinity_repin_on_failover": False, + "model_group_affinity_config": None, + } +) +# Class-level and constant attributes of the Python core. +_CONSTANT_ATTRIBUTES: Final[Mapping[str, object]] = MappingProxyType( + {"default_cache_time_seconds": 1 * 60 * 60, "tenacity": None, "access_groups": None} +) + + +@dataclass(frozen=True) +class SettingsUpdate: + """What a ``ChangeSet.settings`` mapping changes, read with ``update_settings``' rules.""" + + values: Mapping[str, object] = field(default_factory=lambda: MappingProxyType({})) # fields and plain settings + plain: Mapping[str, object] = field(default_factory=lambda: MappingProxyType({})) # plain attributes, as stored + continuation: object = MISSING # the continuation flag, when given + routing_groups: object = MISSING # the routing groups input, when given + optional_pre_call_checks: object = MISSING # the optional pre-call checks, when given + failure: Exception | None = None # what stopped the reading; the settings before it stay applied + + +def read_settings_update(settings: Mapping[str, object], routing_strategy: object) -> SettingsUpdate: + """``_core_update_settings``' reading of each setting, in order: names that are not + runtime-updatable are ignored, the integer settings are cast with ``int``, a retry policy dict + becomes a ``RetryPolicy`` (any other type but ``None`` is dropped), and a strategy other than + ``routing_strategy``, the router's, is validated. The loop has no rollback: a setting that + raises ends it, and the ones before it are applied.""" + values: Final[dict[str, object]] = {} # mutable-ok: built here, then frozen into the update + plain: Final[dict[str, object]] = {} # mutable-ok: same + continuation: object = MISSING # rebind-ok: set when the flag is given + routing_groups: object = MISSING # rebind-ok: set when routing groups are given + optional_checks: object = MISSING # rebind-ok: set when the optional pre-call checks are given + failure: Exception | None = None # rebind-ok: set when a setting raises + try: + for name, value in settings.items(): + if name not in RUNTIME_UPDATABLE_ROUTER_SETTINGS: + verbose_router_logger.debug("Setting %s is not allowed", name) + elif name in _INT_SETTINGS: + values[name] = int(value) # pyright: ignore[reportArgumentType] # `int(kwargs[var])`, as given: a bad value raises int's own error + elif name == "routing_groups": + routing_groups = value + elif name == "optional_pre_call_checks": + # `set_optional_pre_call_checks` reads a list it is given as a set. + _ = frozenset[object](value) if value is not None else None # pyright: ignore[reportArgumentType] # as given: a bad value raises frozenset's own error + optional_checks = value + elif name == CONTINUATION_SETTING: + continuation = value + elif name == "retry_policy": + policy = RetryPolicy.model_validate(value) if isinstance(value, dict) else value + if policy is None or isinstance(policy, RetryPolicy): + plain[name] = policy + elif name in PLAIN_DEFAULTS or name == "max_retries": + plain[name] = value + else: + if name == "routing_strategy" and _normalized(value) != routing_strategy: + validate_routing_strategy(value) # pyright: ignore[reportArgumentType] # validated as given + values[name] = value + except Exception as raised: # noqa: BLE001 # raised by `apply` once the settings before it are applied + failure = raised # rebind-ok: None until a setting fails + return SettingsUpdate( + values=MappingProxyType(values), + plain=MappingProxyType(plain), + continuation=continuation, + routing_groups=routing_groups, + optional_pre_call_checks=optional_checks, + failure=failure, + ) + + +def _normalized(strategy: object) -> object: + """``_normalize_strategy``, which leaves anything but the enum as given.""" + return strategy.value if isinstance(strategy, RoutingStrategy) else strategy + + +def initial_plain(config: "RouterConfig") -> dict[str, object]: # mutable-ok: the core's own plain attributes + """The plain attributes as the constructor sets them: a ``retry_policy`` dict is read into a + ``RetryPolicy`` and any other type but a policy is no policy.""" + from litellm.router_rust.surface.retries import resolved_retry_policy + + plain: Final = {name: config.get(name, default) for name, default in PLAIN_DEFAULTS.items()} + plain["retry_policy"] = resolved_retry_policy(plain["retry_policy"]) + return plain + + +def settings_view( + values: Mapping[str, object], plain: Mapping[str, object], deleted: frozenset[str] +) -> dict[str, object]: # mutable-ok: the settings view is built fresh per call, as `get_settings` builds it + """``get_settings``: the listed settings the router has, then the routing groups. ``values`` + holds the settings fields and plain attributes the Rust core resolves.""" + view: Final[dict[str, object]] = {} # mutable-ok: see above + for name in _VIEW_ORDER: + if name in deleted: + continue + if name in plain: + view[name] = plain[name] + elif name in values: + view[name] = values[name] + view["routing_groups"] = values.get("routing_groups", []) + return view + + +def argument_attribute(config: "RouterConfig", name: str) -> object: + """A router attribute that is its constructor argument (or a constant), else ``MISSING``.""" + if name in _ARGUMENT_ATTRIBUTES: + return config.get(name, _ARGUMENT_ATTRIBUTES[name]) + if name in _CONSTANT_ATTRIBUTES: + return _CONSTANT_ATTRIBUTES[name] + return MISSING + + +def default_litellm_params(config: "RouterConfig") -> dict[str, object]: # mutable-ok: the router's own dict + """``default_litellm_params`` as the constructor completes it: the router timeout, no client + retries and the caching groups.""" + given: Final = _as_dict(config.get("default_litellm_params")) or {} + params: Final[dict[str, object]] = dict(given) # mutable-ok: see above + params.setdefault("timeout", config.get("timeout")) + params.setdefault("max_retries", 0) + metadata: Final = _as_dict(params.setdefault("metadata", {})) + if metadata is not None: + params["metadata"] = {**metadata, "caching_groups": config.get("caching_groups")} + return params + + +def _as_dict(value: object) -> Mapping[str, object] | None: + return value if isinstance(value, dict) else None # pyright: ignore[reportUnknownVariableType] # the argument dict, as passed + + +def request_timeout(config: "RouterConfig") -> float | None: + """The per-attempt ``request_timeout``: only stored when a router timeout is also set.""" + from litellm.litellm_core_utils.request_timeout_resolver import get_configured_request_timeout + + return get_configured_request_timeout() if config.get("timeout") is not None else None + + +def apply_construction_effects(config: "RouterConfig", host: "PythonHost") -> None: + """What the constructor's logging arguments do: the host stops its debug info, and + ``set_verbose`` sets the router logger to ``debug_level``.""" + host.suppress_debug_info() + if config.get("set_verbose", False) is True: + level: Final = config.get("debug_level", "INFO") + if level == "INFO": + verbose_router_logger.setLevel(logging.INFO) + elif level == "DEBUG": + verbose_router_logger.setLevel(logging.DEBUG) + + +__all__ = [ + "CONTINUATION_SETTING", + "MISSING", + "PLAIN_DEFAULTS", + "SettingsUpdate", + "apply_construction_effects", + "argument_attribute", + "default_litellm_params", + "initial_plain", + "read_settings_update", + "request_timeout", + "settings_view", +] diff --git a/litellm/router_utils/cooldown_handlers.py b/litellm/router_utils/cooldown_handlers.py index 7a8a645ad13..03485ebd5f8 100644 --- a/litellm/router_utils/cooldown_handlers.py +++ b/litellm/router_utils/cooldown_handlers.py @@ -125,12 +125,13 @@ _EXCEPTION_POLICY_FIELDS: Final[tuple[tuple[type, str], ...]] = ( def _first_present(*sources: Mapping[str, Any] | None, key: str) -> int | float | None: """Return *key* from the first source mapping where it's set, so callers can support a setting living in more than one deployment config location. Sources - are checked in order from most to least specific to that setting.""" + are checked in order from most to least specific to that setting. Only a number + counts: a bool or a string is skipped, so the router default applies.""" for source in sources: if source is None: continue value = source.get(key) - if value is not None: + if isinstance(value, int | float) and not isinstance(value, bool): return value return None @@ -811,7 +812,7 @@ def _is_allowed_fails_set_on_router( return False -def cast_exception_status_to_int(exception_status: str | int) -> int: +def cast_exception_status_to_int(exception_status: object) -> int: if isinstance(exception_status, str): try: exception_status = int(exception_status) @@ -820,6 +821,13 @@ def cast_exception_status_to_int(exception_status: str | int) -> int: "Unable to cast exception status to int %s. Defaulting to status=500.", exception_status ) exception_status = 500 + # A bool or a number no HTTP status can be is read as a 500, as the Rust router core reads it. + if ( + isinstance(exception_status, bool) + or not isinstance(exception_status, int) + or not 0 <= exception_status <= 65535 + ): + return 500 return exception_status diff --git a/litellm/router_utils/deployment_numbers.py b/litellm/router_utils/deployment_numbers.py new file mode 100644 index 00000000000..7cee96f6a64 --- /dev/null +++ b/litellm/router_utils/deployment_numbers.py @@ -0,0 +1,90 @@ +# pyright: reportUnknownVariableType=false, reportUnknownArgumentType=false +"""Numbers a router deployment cannot be served with, refused the same way by both router cores. + +A whole number wider than 64 bits (a `ptu_shares` of `1e20`) does not fit the router's JSON +deployment record, and an infinite or NaN cost would read back as null there, so the call +would silently be priced from the model cost map. Both are refused when the deployment is +built or added, with the error naming the field. +""" + +import math +import re +from collections.abc import Iterator, Mapping +from typing import Final + +from litellm.litellm_core_utils.core_helpers import as_str_mapping + +# The whole numbers a deployment record holds (serde_json's i64 and u64). +SMALLEST_NUMBER: Final = -(2**63) +LARGEST_NUMBER: Final = 2**64 - 1 +# A pricing field: `input_cost_per_token`, `cache_read_input_token_cost`, ... +_COST_FIELD: Final = re.compile(r"cost_per|_cost$") + + +def _reads_as_non_finite(value: object) -> bool: + """Whether a cost, as written, is an infinity or a NaN (``float()`` reads ``"inf"`` and ``"1e400"``).""" + if isinstance(value, float): + return not math.isfinite(value) + if not isinstance(value, str): + return False + try: + return not math.isfinite(float(value)) + except ValueError: + return False + + +def _is_cost(key: object) -> bool: + return isinstance(key, str) and _COST_FIELD.search(key) is not None + + +def _number_leaves(value: object) -> Iterator[tuple[str, object, bool]]: + """Every value under ``value`` that is no container, in order: its path, and whether it sits in a + cost field. Each container is walked once, so an entry that holds itself ends.""" + pending: Final[list[tuple[object, str, bool]]] = [(value, "", False)] # mutable-ok: the walk's work list + walked: Final[set[int]] = set() # mutable-ok: the ids walked so far + while pending: + current, path, cost = pending.pop() + if not isinstance(current, Mapping | list | tuple): + yield path, current, cost + elif id(current) not in walked: + walked.add(id(current)) + pending.extend( + reversed( + [(item, f"{path}.{key}".removeprefix("."), cost or _is_cost(key)) for key, item in current.items()] + if isinstance(current, Mapping) + else [(item, f"{path}[{index}]", cost) for index, item in enumerate(current)] + ) + ) + + +def _refusal(entry: Mapping[str, object], reason: str) -> ValueError: + return ValueError(f"Deployment {entry.get('model_name')!r} cannot be served: {reason}") + + +def refuse_wide_numbers(entry: object) -> None: + """Refuse a deployment holding a whole number wider than 64 bits, by name.""" + deployment: Final = as_str_mapping(entry) + if deployment is None: + return + for path, value, _ in _number_leaves(deployment): + if isinstance(value, int) and not SMALLEST_NUMBER <= value <= LARGEST_NUMBER: + raise _refusal( + deployment, + f"{path} is {value}, outside what the router can hold " + f"({SMALLEST_NUMBER} to {LARGEST_NUMBER}). Use a smaller value.", + ) + + +def refuse_unservable_numbers(entry: object) -> None: + """``refuse_wide_numbers``, and a deployment with an infinite or NaN cost.""" + refuse_wide_numbers(entry) + deployment: Final = as_str_mapping(entry) + if deployment is None: + return + for path, value, cost in _number_leaves(deployment): + if cost and _reads_as_non_finite(value): + raise _refusal( + deployment, + f"{path} is {value!r}, and the router only prices with finite costs. " + "Set a finite cost, or remove it to use the model cost map's price.", + ) diff --git a/litellm/router_utils/endpoint_eligibility.py b/litellm/router_utils/endpoint_eligibility.py index 4a87e582c97..7f58c138bba 100644 --- a/litellm/router_utils/endpoint_eligibility.py +++ b/litellm/router_utils/endpoint_eligibility.py @@ -57,10 +57,14 @@ class _EndpointDeclaration(BaseModel): return (value,) if isinstance(value, str) else value +def endpoint_for_call(call_type: object, module: object = None) -> str | None: + qualified_call: Final = f"{module}.{call_type}" if module == "litellm.interactions.main" else call_type + return CHAT_ENDPOINT_BY_CALL.get(qualified_call) if isinstance(qualified_call, str) else None + + @contextmanager def endpoint_selection(call_type: object, module: object = None, model: object = None) -> Generator[None]: - qualified_call: Final = f"{module}.{call_type}" if module == "litellm.interactions.main" else call_type - endpoint: Final = CHAT_ENDPOINT_BY_CALL.get(qualified_call) if isinstance(qualified_call, str) else None + endpoint: Final = endpoint_for_call(call_type, module) token: Final = _REQUEST_SELECTION.set((endpoint, model if isinstance(model, str) else None)) try: yield @@ -68,6 +72,10 @@ def endpoint_selection(call_type: object, module: object = None, model: object = _REQUEST_SELECTION.reset(token) +def selected_endpoint() -> str | None: + return _REQUEST_SELECTION.get()[0] + + def _enforced_endpoints(model_info: object) -> tuple[str, ...] | None: try: declaration: Final = _EndpointDeclaration.model_validate(model_info) diff --git a/litellm/router_utils/fallback_event_handlers.py b/litellm/router_utils/fallback_event_handlers.py index e8f832b4b2d..20f9e114747 100644 --- a/litellm/router_utils/fallback_event_handlers.py +++ b/litellm/router_utils/fallback_event_handlers.py @@ -632,6 +632,8 @@ def get_fallback_model_group( ## check for specific model group-specific fallbacks for idx, item in enumerate(fallbacks): if isinstance(item, dict): + if not item: + continue # an empty entry names no group: the next entries still apply fallback_key = next(iter(item)) if fallback_key == model_group: # check exact match fallback_model_group = item[model_group] diff --git a/litellm/router_utils/global_callback_holders.py b/litellm/router_utils/global_callback_holders.py new file mode 100644 index 00000000000..53a7c20cff3 --- /dev/null +++ b/litellm/router_utils/global_callback_holders.py @@ -0,0 +1,53 @@ +""" +The routers of this process that hold optional pre-call checks and routing callbacks, whatever core +serves them. + +A router that retires hands the checks it vacates in its host's ``callbacks`` list to a surviving +router's same-keyed check (``routers_holding_global_callbacks``), and a router that drops a check at +runtime keeps it registered while another live router still has one of that type +(``live_check_owners``). Both rules look at every router in the process, so the sets live here, +outside any one router core. + +A member of ``routers_holding_global_callbacks`` serves ``_global_callback_candidates()``, +``_core_adaptive_routers`` and ``_sync_adaptive_router_hooks()``; a member of ``live_check_owners`` +serves ``optional_callbacks``. +""" + +import threading +import weakref +from collections.abc import Mapping, Sequence +from typing import TYPE_CHECKING, Final, Protocol + +if TYPE_CHECKING: + from litellm.integrations.custom_logger import CustomLogger + + +class GlobalCallbackHolder(Protocol): + """A router whose callbacks a retiring router may hand its vacated checks to.""" + + @property + def _core_adaptive_routers(self) -> Mapping[str, object]: ... + + def _global_callback_candidates(self) -> "tuple[CustomLogger, ...]": ... + + def _sync_adaptive_router_hooks(self) -> None: ... + + +class CheckOwner(Protocol): + """A live router with optional pre-call checks.""" + + @property + def optional_callbacks(self) -> "Sequence[object] | None": ... + + +routers_holding_global_callbacks: Final["weakref.WeakSet[GlobalCallbackHolder]"] = weakref.WeakSet() +routers_holding_global_callbacks_lock: Final = threading.Lock() +live_check_owners: Final["weakref.WeakSet[CheckOwner]"] = weakref.WeakSet() + +__all__ = [ + "CheckOwner", + "GlobalCallbackHolder", + "live_check_owners", + "routers_holding_global_callbacks", + "routers_holding_global_callbacks_lock", +] diff --git a/litellm/router_utils/health_state_cache.py b/litellm/router_utils/health_state_cache.py index c759af8fdb8..1c9cb10a48b 100644 --- a/litellm/router_utils/health_state_cache.py +++ b/litellm/router_utils/health_state_cache.py @@ -6,6 +6,7 @@ and exposes it for router candidate filtering. """ import time +from collections.abc import Mapping from typing import TYPE_CHECKING, Any, Final from typing_extensions import TypedDict @@ -78,7 +79,7 @@ class DeploymentHealthCache: merged: Final = { model_id: state for model_id, state in {**existing, **states}.items() - if isinstance(state, dict) and (now - state.get("timestamp", 0)) < expiry_seconds + if isinstance(state, dict) and (now - _timestamp_of(state)) < expiry_seconds # pyright: ignore[reportUnknownArgumentType] # a cached record } self.cache.set_cache( key=self.CACHE_KEY, @@ -101,7 +102,7 @@ class DeploymentHealthCache: for model_id, state in raw.items() if isinstance(state, dict) and not state.get("is_healthy", True) - and (now - state.get("timestamp", 0)) < self.staleness_threshold + and (now - _timestamp_of(state)) < self.staleness_threshold # pyright: ignore[reportUnknownArgumentType] # a cached record } @with_service_target(HEALTH_CHECKS_TARGET) @@ -129,3 +130,10 @@ class DeploymentHealthCache: str(e), ) return set() + + +def _timestamp_of(state: Mapping[str, object]) -> float: + """A health record's timestamp; one that is no number reads as 0 (stale), so a bad record + drops out alone instead of turning health routing off for every deployment.""" + timestamp: Final = state.get("timestamp", 0) + return float(timestamp) if isinstance(timestamp, int | float) and not isinstance(timestamp, bool) else 0.0 diff --git a/litellm/rust_bridge/_native.pyi b/litellm/rust_bridge/_native.pyi index 954815f6a4f..dc170116089 100644 --- a/litellm/rust_bridge/_native.pyi +++ b/litellm/rust_bridge/_native.pyi @@ -1,6 +1,6 @@ from asyncio import Future from collections.abc import Mapping, Sequence -from typing import Never, final +from typing import Generic, Literal, Never, NotRequired, TypedDict, TypeVar, final import httpx from pydantic import JsonValue @@ -218,6 +218,8 @@ __all__ = [ "NativeDiagnosticProcessor", "ProcessReservedForForking", "ResponsesWebSocketConnection", + "RouterCall", + "RouterCoreHandle", "RustBridgeDeclined", "RustUpstreamError", "TokenCounter", @@ -296,3 +298,493 @@ class _SecretManagerRuntime: timeout: float | httpx.Timeout | None = None, primary_secret_name: str | None = None, ) -> Future[JsonValue]: ... + +# The FailureKind names an attempt outcome's `kind` takes (router-core `FailureKind`). +_FailureKind = Literal[ + "RateLimitError", + "ContextWindowExceededError", + "ContentPolicyViolationError", + "AuthenticationError", + "PermissionDeniedError", + "NotFoundError", + "BadRequestError", + "Timeout", + "APIConnectionError", + "InternalServerError", + "ServiceUnavailableError", + "BadGatewayError", + "GuardrailIntervention", + "Declined", + "CallerCredentialAuthenticationError", + "CallerCredentialRateLimitError", + "APIError", +] + +# One event a call reported since the last step (router-core `RouterEvent`), carried on the next step. +class _Event(TypedDict): + kind: str + +# Every step `RouterCall.next` returns carries the events reported since the last one. +class _Step(TypedDict): + events: list[_Event] + +class _AttemptStep(_Step): + kind: Literal["attempt"] + deployment_id: str + # The deployment as a `model_list` entry, with its id in `model_info`. + deployment: dict[str, JsonValue] + # The name this attempt serves: the requested model or the fallback target. + model: str + model_group: str + target_order: int | None + candidate_ids: list[str] + retry: int + fallback_depth: int + fallback_level: int + timeout: float | None + # Request params the fallback walk's dict targets merged in. + params: dict[str, JsonValue] | None + # The walk level whose request metadata the attempt carries. + scope: int + max_retries: int + model_group_size: int + +class _SleepStep(_Step): + kind: Literal["sleep"] + seconds: float + +class _AllowFallbackStep(_Step): + kind: Literal["allow_fallback"] + target: str + +class _RefuseFallbacksStep(_Step): + kind: Literal["refuse_fallbacks"] + # The walk level whose chain the dict targets are. + level: int + targets: list[dict[str, JsonValue]] + +class _SelectStep(_Step): + kind: Literal["select"] + stage: Literal[ + "plugins", + "strategy", + "tags", + "pre_call", + "text_only", + "callback_filters", + "routing_plugins", + "provider_request", + ] + strategy: str | None + model_group: str + # The name the request resolved from, which strategies key their state by. + model: str + # Whose selector a strategy reads: "default", "group" (with `group`) or "override". + scope: str + group: str | None + # The deployments the name resolved to, before health filters. + pool: list[str] + candidates: list[str] + # `routing_plugins` only: whether the plugins run. + run: NotRequired[bool] + # `callback_filters` only: the ids the request already failed on. + excluded_ids: NotRequired[list[str]] + skipped_ids: NotRequired[list[str]] + +# Pre-routing before a pick resolves `model`: answered with the model group to resolve instead +# (a string), a custom strategy's deployment document (a mapping), or None. +class _PreRouteStep(_Step): + kind: Literal["pre_route"] + model: str + fallback_depth: int + retry: int + +# One RouterCache op against the router's DualCache. `key` is set for every op but `batch_get`, +# and `increment_usage`, which set `keys`; `value` is the JSON to store (`set`), the amount +# (`increment`) or one amount per key (`increment_usage`); `ttl` is in seconds; `local` reads or +# writes this worker's in-memory tier only. Outcomes: get -> the JSON value or None; batch_get -> +# one JSON value or None per key, in order; increment -> the new int or float; set, delete and +# increment_usage -> None. +class _CacheStep(_Step): + kind: Literal["cache"] + op: Literal["get", "set", "increment", "increment_usage", "batch_get", "delete"] + key: str | None + keys: list[str] | None + value: JsonValue + ttl: float | None + local: bool + +# One lookup of a model-info query (router-core `ModelInfoLookup`): its `lookup` name and that +# lookup's arguments, answered as the lookup documents. +class _ModelInfoStep(_Step): + kind: Literal["model_info"] + lookup: Literal[ + "model_info", + "cost_entry", + "provider", + "pattern_models", + "bedrock_base", + "discovered", + "user_model_info", + "log", + "group_params", + "group_deployment", + "group_capabilities", + "group_seed", + "reasoning_intersect", + ] + model: NotRequired[JsonValue] + tolerate: NotRequired[bool] + exact: NotRequired[bool] + key: NotRequired[str] + process: NotRequired[bool] + name: NotRequired[str] + model_id: NotRequired[str] + id: NotRequired[JsonValue] + level: NotRequired[Literal["debug", "error"]] + message: NotRequired[str] + index: NotRequired[int] + model_info: NotRequired[JsonValue] + litellm_model: NotRequired[JsonValue] + llm_provider: NotRequired[JsonValue] + mapped: NotRequired[bool] + info: NotRequired[dict[str, JsonValue]] + current: NotRequired[JsonValue] + resolved: NotRequired[JsonValue] + +class _FailedStep(_Step): + kind: Literal["failed"] + error: str + error_kind: Literal[ + "config", + "unknown_model", + "no_healthy_deployments", + "attempt", + "cache", + "retired", + "paused", + "not_pass_through", + "no_pass_through_deployments", + "only_markers", + "ptu_reserved", + "context_window_exceeded", + "unsupported_params", + "routing_plugins_left_nothing", + "routing_plugins_need_async", + "multiple_teams", + "reserved", + "invalid_router_weights", + "non_numeric_weight", + "missing_routing_group", + "default_deployment_raises", + "index_out_of_range", + "model_info", + ] + attempt_failed: bool + status: int | None + failure_kind: _FailureKind | None + error_model: NotRequired[str] + error_detail: NotRequired[str] + error_params: NotRequired[list[str]] + # `model_info` only: the builtin exception the lookup raised. + error_exception: NotRequired[str] + +# The end of a `RouterCoreHandle.start` (and the stream recovery calls). +class _DoneStep(_Step): + kind: Literal["done"] + deployment_id: str + model_group: str + attempted_retries: int + max_retries: int + attempted_fallbacks: int + total_tokens: int + +# The end of a `RouterCoreHandle.select`. +class _SelectedStep(_Step): + kind: Literal["done"] + deployment_id: str + model_group: str + deployment: dict[str, JsonValue] + +# The end of a `RouterCoreHandle.healthy`. +class _HealthyStep(_Step): + kind: Literal["done"] + # The name was a deployment id. + single: bool + deployment_ids: list[str] + deployments: list[dict[str, JsonValue]] + +# The end of a `RouterCoreHandle.report_failure`. +class _ReportedStep(_Step): + kind: Literal["done"] + cooled: bool + # The cooldown applied, in seconds; None when the deployment was not cooled down. + cooldown_time: float | None + +# The end of a `RouterCoreHandle.report_success` or `count_success`. +class _CountedStep(_Step): + kind: Literal["done"] + key: str + +# One `os.environ/` reference a `secrets` step asks for: a top-level param (`key`), or an +# env-only `dataSources[source].parameters[key]` (`env_only`). +class _SecretRef(TypedDict): + name: str + key: str + env_only: bool + source: NotRequired[int] + +# The deployment build's `os.environ/` references: answered with one value per ref, in order. +class _SecretsStep(_Step): + kind: Literal["secrets"] + index: int + refs: list[_SecretRef] + +# The deployment build's steps for one deployment. `validate` and `added` are answered with +# `{"entry": }`, `{"skip": True}` or `{"invalid": True}`; `provider` with +# `{"provider": str | None}` or `{"invalid": True}`. +class _BuildStep(_Step): + kind: Literal["validate", "provider", "added"] + path: Literal["create", "add"] + index: int + entry: dict[str, JsonValue] + declared_id: str | None + duplicate_id: bool + provider: str | None + +class _BuiltDeployment(TypedDict): + index: int + id: str | None + provider: str | None + +# The end of a `RouterCoreHandle.build`. +class _BuiltStep(_Step): + kind: Literal["done"] + built: list[_BuiltDeployment] + +# The end of a model-info query: `router_model_info` sets `info` (the merged map or None) and +# `user_keys`; `deployment_model_info` sets `info`; the others set `value`. +class _ModelInfoDoneStep(_Step): + kind: Literal["done"] + info: NotRequired[dict[str, JsonValue] | None] + user_keys: NotRequired[list[str]] + value: NotRequired[JsonValue] + +# `RouterCoreHandle.resolve`: what a name resolves to. +class _Resolution(TypedDict): + model_group: str + kind: str + ids: list[str] + +# The steps each kind of `RouterCall` returns, by the handle method that starts it. +_CallSteps = ( + _AttemptStep + | _SleepStep + | _AllowFallbackStep + | _RefuseFallbacksStep + | _SelectStep + | _PreRouteStep + | _CacheStep + | _DoneStep + | _FailedStep +) +_SelectSteps = _SelectStep | _PreRouteStep | _CacheStep | _SelectedStep | _FailedStep +_HealthySteps = _SelectStep | _PreRouteStep | _CacheStep | _HealthyStep | _FailedStep +_BuildSteps = _SecretsStep | _BuildStep | _BuiltStep | _FailedStep +_ModelInfoSteps = _ModelInfoStep | _ModelInfoDoneStep | _FailedStep +_FailureReportSteps = _CacheStep | _ReportedStep | _FailedStep +_SuccessReportSteps = _CacheStep | _CountedStep | _FailedStep + +class _AttemptSucceeded(TypedDict): + ok: Literal[True] + total_tokens: NotRequired[int] + blocked: NotRequired[Literal["content_filter", "refusal"] | None] + # A stream whose success or failure the host reports when it ends (`report_success`, + # `report_failure`): the attempt counts its request and `total_tokens` only. + deferred: NotRequired[bool] + # Whether the attempt counts against its deployment's usage; unset, its own limits decide. + tracks_usage: NotRequired[bool | None] + +class _AttemptFailed(TypedDict): + ok: Literal[False] + kind: _FailureKind + status: NotRequired[int | None] + retry_after: NotRequired[float | None] + num_retries: NotRequired[int | None] + # The deployment the exception names, which a retry skips over the attempt's own. + failed_deployment_id: NotRequired[str | None] + detail: NotRequired[str | None] + agentic_followup: NotRequired[bool] + # The exception has no `status_code`, which no cooldown rule reads. + status_absent: NotRequired[bool] + # The failure counts against no per-minute requests. + uncounted_request: NotRequired[bool] + +# Each document is a mapping or its JSON text. +_Document = Mapping[str, object] | str + +@final +class RouterCoreHandle: + def __new__(cls, config: _Document, dynamic: _Document) -> RouterCoreHandle: ... + def start(self, op: str, request: _Document) -> RouterCall[_CallSteps]: ... + # Driven like `start`; its `done` step is a `_SelectedStep`. + def select(self, model: str, request: _Document) -> RouterCall[_SelectSteps]: ... + # Driven like `select`; its `done` step is a `_HealthyStep`. + def healthy(self, model: str, request: _Document) -> RouterCall[_HealthySteps]: ... + # `{"replace" | "append" | "upsert": [deployments], "delete": [ids], "settings": {dynamic document}, "index": {...}}`. + def apply(self, changes: _Document) -> int: ... + # Build model_list entries through the host, driven like `start`: `_SecretsStep` and + # `_BuildStep` steps, then a `_BuiltStep`. `path`: "create" (`set_model_list`) or "add"; + # `install`: "replace", "append" or "upsert". + def build( + self, entries: Sequence[object], path: str, resolve_secrets: bool, install: str + ) -> RouterCall[_BuildSteps]: ... + # Keep the providers `other`'s builds resolved for the same deployments. + def adopt_providers(self, other: RouterCoreHandle) -> None: ... + # `get_router_model_info`, driven like `start`: `model_info` lookup steps (a `lookup` field + # and its arguments), then a done step with `info` (the merged map or None) and `user_keys`. + def router_model_info(self, query: _Document) -> RouterCall[_ModelInfoSteps]: ... + # `get_deployment_model_info`, driven like `router_model_info`; its done step has `info`. + def deployment_model_info(self, model_id: str, model_name: object) -> RouterCall[_ModelInfoSteps]: ... + # `_set_model_group_info` for a `GroupInfoQuery` document, driven like `router_model_info`; its + # done step has `value`: None, or `{"assignments": {field: value}}`. + def model_group_info(self, query: _Document) -> RouterCall[_ModelInfoSteps]: ... + # `get_configured_token_limits`, driven like `router_model_info`; `value` is `[input, output]`. + def configured_token_limits(self, id: object, configured: _Document) -> RouterCall[_ModelInfoSteps]: ... + # `get_model_listing_info`, driven like `router_model_info`; `value` has `cost_map_keys`, + # `max_input_tokens` and `max_output_tokens`. + def model_listing_info(self, deployments: list[object]) -> RouterCall[_ModelInfoSteps]: ... + # A configured string that is not blank (`get_configured_mode`, `get_configured_display_name`). + @staticmethod + def configured_text(value: object) -> str | None: ... + # Each model_info's `service_tiers`, unvalidated. + @staticmethod + def configured_service_tiers(model_infos: list[object]) -> list[object]: ... + # The routing-cache key a deployment counts `metric` (tpm, rpm, itpm, otpm) under in `minute`. + def usage_key(self, id: str, model: str, metric: str, minute: str) -> str: ... + def registry_version(self) -> int: ... + # The version the deployments before the newest appends last changed at. + def registry_rewritten(self) -> int: ... + # A failure the router did not see itself, driven like `start`: cache steps, then a + # `_ReportedStep`. + def report_failure( + self, + id: str, + failure: Mapping[str, object], + deployment_cooldown: bool = False, + count_only: bool = False, + requested_group: str | None = None, + cooldown_time: float | None = None, + uncounted: bool = False, + counts_request: bool = False, + usage_model: str | None = None, + callbacks: Mapping[str, bool] | None = None, + ) -> RouterCall[_FailureReportSteps]: ... + # A success a logging callback saw, driven like `report_failure`; `done` carries `key`. + def report_success( + self, + id: str, + total_tokens: float, + requests: float, + usage_model: str | None = None, + tracks_usage: bool | None = None, + callbacks: Mapping[str, bool] | None = None, + ) -> RouterCall[_SuccessReportSteps]: ... + # Count a success against any id, driven like `report_failure`; `done` carries `key`. + def count_success(self, id: str, callbacks: Mapping[str, bool] | None = None) -> RouterCall[_SuccessReportSteps]: ... + # [(id, tpm_key, rpm_key)] for this minute, or None for a group with no deployment. + def usage_keys(self, model_group: str) -> list[tuple[str, str, str]] | None: ... + # Every deployment as a `model_list` entry, in config order. + def deployments(self) -> list[dict[str, JsonValue]]: ... + def deployments_from(self, start: int) -> list[dict[str, JsonValue]]: ... + def deployment(self, id: str) -> dict[str, JsonValue] | None: ... + # Lookups over the registry indexes, answering what each index holds as written. + # Registry lookups (router-core `lookup.rs`): rows the host reads its stored entries at + # (`{"at": position, "name"?, "group"?, "member"?}`, `{"pattern": entry, "name"}`, or a + # `{"raw": value}` / `{"raw_positions": value}` its own read raises at). + def lookup_all( + self, model_name: str, team_id: str | None = None, rename: str | None = None + ) -> list[dict[str, JsonValue]]: ... + def lookup_across_teams(self, model: str) -> tuple[list[dict[str, JsonValue]], int]: ... + def lookup_config(self) -> list[dict[str, JsonValue]]: ... + def lookup_model_ids( + self, model_name: str | None = None, exclude_team_models: bool = False + ) -> tuple[list[str], dict[str, JsonValue] | None]: ... + def lookup_model_list( + self, model_name: str | None, team_id: str | None, scope: Mapping[str, object] + ) -> list[dict[str, JsonValue]]: ... + def lookup_model_names( + self, team_id: str | None, scope: Mapping[str, object] + ) -> tuple[list[str], dict[str, JsonValue] | None]: ... + def lookup_access_groups( + self, model_name: str | None, access_group: str | None, team_id: str | None, scope: Mapping[str, object] + ) -> tuple[list[tuple[str, list[str]]], dict[str, JsonValue] | None]: ... + def lookup_name_serving( + self, model_id: str, scope: Mapping[str, object] + ) -> tuple[bool, str | None, dict[str, JsonValue] | None]: ... + def has_model_name(self, name: str) -> bool: ... + def lookup_team_model_has_alternatives(self, id: str) -> tuple[bool, dict[str, JsonValue] | None]: ... + def lookup_group_usable_by_team( + self, model_group: str, team_id: str | None + ) -> tuple[int | None, dict[str, JsonValue] | None]: ... + def lookup_team_public_name( + self, model_id: str, team_id: str + ) -> tuple[int | None, dict[str, JsonValue] | None]: ... + def lookup_routable(self, model_group: str, team_id: str | None) -> list[dict[str, JsonValue]]: ... + def lookup_wildcard_route( + self, model: str, team_id: str | None, requests: list[str] + ) -> list[dict[str, JsonValue]]: ... + def lookup_candidate_ids( + self, resolved: str, team_id: str | None, members: list[str] | None, requests: list[str] + ) -> tuple[list[str], dict[str, JsonValue] | None]: ... + def has_model_id(self, id: str) -> bool: ... + def id_position(self, id: str) -> JsonValue | None: ... + def id_positions(self) -> list[tuple[str, JsonValue]]: ... + def named_positions(self, name: str) -> JsonValue | None: ... + def team_positions(self, team: str, public: str) -> JsonValue | None: ... + def team_public_names_of(self, team: str) -> list[str]: ... + def litellm_model_positions(self, model: str) -> list[int]: ... + # The registry indexes the host hands out (`HostIndexes`), as JSON. + def host_indexes(self) -> dict[str, JsonValue]: ... + def provider_default_ids(self) -> list[str]: ... + # `PatternMatchRouter.route` over the wildcard router, or `team`'s. + def pattern_route( + self, team: str | None, request: str, filtered: list[str] | None = None + ) -> list[dict[str, JsonValue]] | None: ... + # `get_models_by_pattern` over the wildcard router, or `team`'s. + def models_by_pattern(self, team: str | None, requests: list[str]) -> list[str] | None: ... + # The provider the host resolved when it built deployment `id` (its `provider` build step). + def provider(self, id: str) -> str | None: ... + def resolve(self, name: str, specific_deployment: bool = False) -> _Resolution | None: ... + def routable(self, model: str, request: _Document) -> list[dict[str, JsonValue]]: ... + # `_get_routing_context` for a pick of `model`: (strategy, scope, group). + def routing_context(self, model: str, routing_strategy: str | None = None) -> tuple[str, str, str | None]: ... + def retire(self) -> None: ... + # A stream that failed before content: what may still take over it, and the retry budget for + # a failure of `kind`: {"can_retry", "can_fall_back", "refusal_fallback": bool, + # "retry_budget": int | None, "policy_applies": bool}. + def stream_policy( + self, request: _Document, origin: _Document, kind: str | None = None + ) -> dict[str, bool | int | None]: ... + # The same-group retry and the fallback walk for such a stream, driven like `start`. + def retry_stream( + self, op: str, request: _Document, origin: _Document, failure: Mapping[str, object] + ) -> RouterCall[_CallSteps]: ... + def fallback_stream( + self, op: str, request: _Document, origin: _Document, failure: Mapping[str, object] + ) -> RouterCall[_CallSteps]: ... + def retry_policy_retries(self, failure: Mapping[str, object], model_group: str | None = None) -> int | None: ... + def allowed_fails_for(self, failure: Mapping[str, object]) -> int | None: ... + def should_retry( + self, failure: Mapping[str, object], counts: tuple[int, int], fallbacks: tuple[bool, bool, bool] + ) -> bool: ... + +_S = TypeVar("_S", covariant=True) + +@final +class RouterCall(Generic[_S]): + # The answer to the last step, which Rust validates against it: attempt: _AttemptSucceeded | + # _AttemptFailed; sleep: None; allow_fallback: bool; refuse_fallbacks: bool; select: the remaining + # candidate ids, pick first; pre_route: see _PreRouteStep; cache: see _CacheStep; secrets, build and model_info: + # see those steps; first call: None. + def next(self, outcome: object) -> _S: ... diff --git a/litellm/rust_bridge/catalog.py b/litellm/rust_bridge/catalog.py index 9ba8cbc35a7..a12f79c6c8a 100644 --- a/litellm/rust_bridge/catalog.py +++ b/litellm/rust_bridge/catalog.py @@ -23,6 +23,7 @@ class Route(str, Enum): OCR = "ocr" TOKEN_COUNTER = "token_counter" TOKENIZER = "tokenizer" + ROUTER = "router" class Target(str, Enum): @@ -224,6 +225,7 @@ RULES: Final[Rules] = ( RouteRule(Route.RESPONSES, Rollout.PYTHON_ONLY), RouteRule(Route.TOKEN_COUNTER, Rollout.PYTHON_ONLY), RouteRule(Route.TOKENIZER, Rollout.PYTHON_ONLY), + RouteRule(Route.ROUTER, Rollout.RUST_OPT_IN), RouteRule(Route.TRANSCRIPTION, Rollout.RUST_REQUIRED, providers=frozenset({"bedrock"})), ) diff --git a/litellm/types/router.py b/litellm/types/router.py index 4e3e4d92b3c..3363d113fc7 100644 --- a/litellm/types/router.py +++ b/litellm/types/router.py @@ -1154,9 +1154,9 @@ class CustomRoutingStrategyBase: self, model: str, messages: list[dict[str, str]] | None = None, - input: str | list | None = None, + input: str | list[object] | None = None, specific_deployment: bool | None = False, - request_kwargs: dict | None = None, + request_kwargs: dict[str, object] | None = None, ) -> None: """ Asynchronously retrieves the available deployment based on the given parameters. @@ -1177,9 +1177,9 @@ class CustomRoutingStrategyBase: self, model: str, messages: list[dict[str, str]] | None = None, - input: str | list | None = None, + input: str | list[object] | None = None, specific_deployment: bool | None = False, - request_kwargs: dict | None = None, + request_kwargs: dict[str, object] | None = None, ) -> None: """ Synchronously retrieves the available deployment based on the given parameters. diff --git a/litellm/utils.py b/litellm/utils.py index 7f81ece8314..53b02ebddac 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -4305,6 +4305,13 @@ class _LiveDeploymentReplay: """ callback: Callable[[], None] | None = None + # Replays other router cores add beside the slot, run after it. + extra: tuple[Callable[[], None], ...] = () + + +def add_live_deployment_replay(replay: Callable[[], None]) -> None: + """Add a callback that re-asserts another router core's live deployments after a refresh.""" + _LiveDeploymentReplay.extra = (*_LiveDeploymentReplay.extra, replay) def set_live_deployment_replay(replay: Callable[[], None]) -> None: @@ -4338,6 +4345,8 @@ def reapply_runtime_model_cost_registrations() -> None: """ if _LiveDeploymentReplay.callback is not None: _LiveDeploymentReplay.callback() + for extra_replay in _LiveDeploymentReplay.extra: + extra_replay() if _runtime_registered_model_cost: register_model( model_cost=dict( @@ -6338,7 +6347,15 @@ def get_deployment_order(deployment: Mapping[str, object]) -> int | None: resolved_order: Final = ( order if order is not None else cast(Mapping[str, object], deployment.get("model_info", {})).get("order") ) - return cast(int | None, resolved_order) + # One number for every spelling (`1`, `1.0`, `"1"`, `True`), so orders written different ways + # sort and compare together; anything else counts as no order. + if isinstance(resolved_order, bool | int): + return int(resolved_order) + if isinstance(resolved_order, float): + return int(resolved_order) if resolved_order.is_integer() else None + if isinstance(resolved_order, str) and re.fullmatch(r"[+-]?[0-9]+", resolved_order): + return int(resolved_order) + return None _get_deployment_order = get_deployment_order diff --git a/tests/code_coverage_tests/recursive_detector.py b/tests/code_coverage_tests/recursive_detector.py index a84a40c67c1..f0275c34c06 100644 --- a/tests/code_coverage_tests/recursive_detector.py +++ b/tests/code_coverage_tests/recursive_detector.py @@ -87,6 +87,9 @@ IGNORE_FUNCTIONS = [ "scrub_json_strings", # max depth set (MAX_SCRUB_DEPTH); fails closed by returning "[Filtered]" for anything nested past the cap. "_without_requester_ip_address", # walks only the fresh tree _sanitize_request_body_for_spend_logs_payload returns (mapping cycles already cut), so it never goes deeper than that sanitizer did. "_merge_session_fields", # max depth set (_SESSION_MERGE_MAX_DEPTH); deeper mappings are replaced wholesale by the later session.update. + "_json", # router_rust index views: max depth set (DEFAULT_MAX_RECURSE_DEPTH); fails closed by raising ValueError at the cap. + "to_wire", # router_rust model info: bounded by the nesting depth of a stored deployment or cost-map entry (finite JSON trees: pydantic refuses a cycle when it dumps a deployment). + "from_wire", # router_rust model info: walks only what router-core merged from to_wire's output, so it never goes deeper than to_wire did. ] diff --git a/tests/code_coverage_tests/router_litellm_globals_check.py b/tests/code_coverage_tests/router_litellm_globals_check.py index ff9fac0e3dc..172a3754000 100644 --- a/tests/code_coverage_tests/router_litellm_globals_check.py +++ b/tests/code_coverage_tests/router_litellm_globals_check.py @@ -21,6 +21,7 @@ CHECKED = ( "litellm/router_core.py", "litellm/router_deprecated.py", "litellm/router_python_core.py", + "litellm/router_rust/core.py", ) diff --git a/tests/code_coverage_tests/router_surface_check.py b/tests/code_coverage_tests/router_surface_check.py index a8e0d6fa0d7..1c41526101a 100644 --- a/tests/code_coverage_tests/router_surface_check.py +++ b/tests/code_coverage_tests/router_surface_check.py @@ -11,7 +11,8 @@ modules listed there. The lists hold what is used today and may only shrink: also fails, so the entry is removed in the same change. "Outside" means every non-test file under litellm/ and enterprise/ except -litellm/router.py and the router's own modules. Uses from the router family +litellm/router.py and the router's own modules (including every +module under litellm/router_rust/, the Rust core's host-side glue). Uses from the router family (litellm/router_utils, litellm/router_strategy) are listed separately, because those modules become private parts of the router core in a later step. @@ -42,9 +43,12 @@ ROUTER_FILES = { "litellm/router_core.py", "litellm/router_python_core.py", "litellm/router_python_host.py", + "litellm/router_rust/core.py", "litellm/router_deprecated.py", "litellm/router_host.py", } +# The Rust core's host-side glue modules are part of the router. +ROUTER_FILE_PREFIXES = ("litellm/router_rust/",) ROUTER_CLASS_FILES = ["litellm/router.py", "litellm/router_python_core.py", "litellm/router_deprecated.py"] ROUTER_CLASS_NAMES = {"Router", "PythonCore", "DeprecatedRouterMembers"} @@ -128,7 +132,7 @@ def scan() -> Dict[str, Dict[str, Set[str]]]: uses: Dict[str, Dict[str, Set[str]]] = {"outside": {}, "router_family": {}} imports: Dict[str, Set[str]] = {} for rel, path in _python_files(): - if rel in ROUTER_FILES or _is_test(rel): + if rel in ROUTER_FILES or rel.startswith(ROUTER_FILE_PREFIXES) or _is_test(rel): continue with open(path, encoding="utf-8") as f: src = f.read() diff --git a/tests/unit/conftest.py b/tests/unit/conftest.py index 914eb1f2508..4a07c97b709 100644 --- a/tests/unit/conftest.py +++ b/tests/unit/conftest.py @@ -252,6 +252,12 @@ def socket_tempdir(socket_dir: Path, monkeypatch: pytest.MonkeyPatch) -> Path: return socket_dir +@pytest.fixture +def python_router_core(monkeypatch: pytest.MonkeyPatch) -> None: + """Keep ``Router`` on ``PythonCore``, for a test that patches a member the Rust core never reads.""" + monkeypatch.setenv("LITELLM_RUST", "0") + + @pytest.fixture(autouse=True, scope="session") def bundled_tiktoken_cache() -> None: importlib.import_module("litellm.litellm_core_utils.default_encoding") diff --git a/tests/unit/litellm_core_utils/llm_response_utils/test_invalid_response_traceback.py b/tests/unit/litellm_core_utils/llm_response_utils/test_invalid_response_traceback.py new file mode 100644 index 00000000000..a2beb60665e --- /dev/null +++ b/tests/unit/litellm_core_utils/llm_response_utils/test_invalid_response_traceback.py @@ -0,0 +1,54 @@ +""" +The "Invalid response object" error carries only the conversion's own traceback. A conversion that +runs while another error is being handled (a router retry inside the previous attempt's ``except``) +must not copy that error's text into its message, where exception mapping would read it. +""" + +import pytest + +from litellm.litellm_core_utils.llm_response_utils.convert_dict_to_response import convert_to_model_response_object # pyright: ignore[reportUnknownVariableType] # its untyped start_time, end_time and dict parameters +from litellm.types.utils import ModelResponse + +# A text_completion body handed to the chat conversion: ``choices[0]`` has no ``message``. +_MALFORMED_CHAT_BODY = { + "id": "cmpl-1", + "object": "text_completion", + "created": 1, + "model": "gpt-4o", + "choices": [{"text": "ok", "index": 0, "finish_reason": "stop"}], + "usage": {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2}, +} + +_OUTER_TEXT = "outer attempt failed: rate limit exceeded" + + +class _OuterAttemptError(Exception): + pass + + +def _convert_while_handling_another_error() -> Exception: + try: + raise _OuterAttemptError(_OUTER_TEXT) + except _OuterAttemptError: + with pytest.raises(Exception, match="Invalid response object") as exc_info: + _ = convert_to_model_response_object( + response_object=dict(_MALFORMED_CHAT_BODY), + model_response_object=ModelResponse(), + response_type="completion", + ) + return exc_info.value + + +def test_the_message_leaves_out_the_error_being_handled_when_the_conversion_ran() -> None: + message = str(_convert_while_handling_another_error()) + + assert _OUTER_TEXT not in message + assert "_OuterAttemptError" not in message + assert "During handling of the above exception" not in message + + +def test_the_message_names_the_conversion_error_itself() -> None: + message = str(_convert_while_handling_another_error()) + + assert "KeyError: 'message'" in message + assert "Traceback (most recent call last)" in message diff --git a/tests/unit/litellm_core_utils/test_core_helpers.py b/tests/unit/litellm_core_utils/test_core_helpers.py index b21e9eb8b5a..5d6d14a4ca8 100644 --- a/tests/unit/litellm_core_utils/test_core_helpers.py +++ b/tests/unit/litellm_core_utils/test_core_helpers.py @@ -34,6 +34,7 @@ from litellm.litellm_core_utils.core_helpers import ( response_cost_in_hidden_params, set_provider_response_headers_in_hidden_params, unbind_budget_reservation_from_callbacks, + without, ) from litellm.proxy._types import UserAPIKeyAuth from litellm.types.utils import ( @@ -852,3 +853,21 @@ def test_billed_usage_cost_keeps_the_upstream_charge_after_litellm_wrote_its_bil stamp_litellm_bill_on_usage(unreported, unreported.usage, 0.0015) assert unreported.usage.cost == 0.0015 assert billed_usage_cost(unreported, "openai") is None + + +def test_without_drops_the_named_keys_and_keeps_the_rest_in_order() -> None: + source: Final = {"a": 1, "b": 2, "c": 3, "d": 4} + + assert list(without(source, "b", "d").items()) == [("a", 1), ("c", 3)] + assert without(source, "missing") == source + assert without(source) == source + + +def test_without_returns_a_new_dict_and_leaves_the_source_unchanged() -> None: + source: Final = {"a": 1, "b": 2} + + result: Final = without(source, "a") + result["c"] = 3 + + assert source == {"a": 1, "b": 2} + assert result == {"b": 2, "c": 3} diff --git a/tests/unit/proxy/auth/test_fallback_model_access.py b/tests/unit/proxy/auth/test_fallback_model_access.py index feeb2d26ff3..5c6a98d9e42 100644 --- a/tests/unit/proxy/auth/test_fallback_model_access.py +++ b/tests/unit/proxy/auth/test_fallback_model_access.py @@ -1,4 +1,4 @@ -from typing import Final +from typing import Final, NoReturn import pytest @@ -50,14 +50,15 @@ async def test_is_model_authorized_for_token_follows_the_key_access_groups(): assert await is_model_authorized_for_token(model="secret-model", valid_token=token, llm_router=router) is False -class _RouterWithBrokenAccessGroupLookup(Router): - def get_model_access_groups(self, *args, **kwargs): - raise RuntimeError("access group store unavailable") +def _broken_access_group_lookup(*args: object, **kwargs: object) -> NoReturn: + raise RuntimeError("access group store unavailable") @pytest.mark.asyncio -async def test_is_model_authorized_for_token_fails_closed_when_the_lookup_breaks(): - router = _RouterWithBrokenAccessGroupLookup(model_list=_router().model_list) +@pytest.mark.usefixtures("python_router_core") # the Rust core keeps this lookup on its host registry +async def test_is_model_authorized_for_token_fails_closed_when_the_lookup_breaks(monkeypatch: pytest.MonkeyPatch): + router = _router() + monkeypatch.setattr(router, "_core_get_model_access_groups", _broken_access_group_lookup) assert ( await is_model_authorized_for_token( diff --git a/tests/unit/proxy/auth/test_model_access_group_budgets.py b/tests/unit/proxy/auth/test_model_access_group_budgets.py index 7506bd031d9..d1b381907fa 100644 --- a/tests/unit/proxy/auth/test_model_access_group_budgets.py +++ b/tests/unit/proxy/auth/test_model_access_group_budgets.py @@ -8,7 +8,8 @@ unioned rather than ranked. """ from types import SimpleNamespace -from unittest.mock import patch +from typing import NoReturn +from unittest.mock import MagicMock, patch import pytest @@ -66,18 +67,6 @@ class _ExplodingPrismaClient: raise AssertionError(f"unexpected database access: {name}") -class _CountingRouter(Router): - """Counts access-group lookups, so a test can prove the registry gate skipped them.""" - - def __init__(self, **kwargs) -> None: - super().__init__(**kwargs) - self.access_group_lookups = 0 - - def get_model_access_groups(self, *args, **kwargs): - self.access_group_lookups += 1 - return super().get_model_access_groups(*args, **kwargs) - - async def _cache( budgeted_groups: tuple[str, ...] = BUDGETED_GROUPS, member_allowed_models: tuple[str, ...] = (), @@ -233,14 +222,17 @@ async def test_group_without_a_budget_is_not_attributed(): @pytest.mark.asyncio -async def test_empty_registry_skips_the_access_group_matching_entirely(): - router = _CountingRouter(model_list=MODEL_LIST) +@pytest.mark.usefixtures("python_router_core") # the Rust core keeps this lookup on its host registry +async def test_empty_registry_skips_the_access_group_matching_entirely(monkeypatch: pytest.MonkeyPatch): + router = Router(model_list=MODEL_LIST) + access_group_lookups = MagicMock(wraps=router._core_get_model_access_groups) + monkeypatch.setattr(router, "_core_get_model_access_groups", access_group_lookups) assert await _matched(key_models=["tier-a"], cache=await _cache(budgeted_groups=()), llm_router=router) == () - assert router.access_group_lookups == 0 + assert access_group_lookups.call_count == 0 assert await _matched(key_models=["tier-a"], llm_router=router) == ("tier-a",) - assert router.access_group_lookups == 1 + assert access_group_lookups.call_count == 1 @pytest.mark.asyncio @@ -262,22 +254,24 @@ async def test_stamp_records_the_matched_groups_on_the_auth_object(): assert valid_token.matched_model_access_groups == ["tier-a", "tier-b"] -class _BrokenRouter(Router): - def get_model_access_groups(self, *args, **kwargs): - raise RuntimeError("access group store unavailable") +def _broken_access_group_lookup(*args: object, **kwargs: object) -> NoReturn: + raise RuntimeError("access group store unavailable") @pytest.mark.asyncio -async def test_stamp_does_not_break_auth_when_the_access_group_lookup_fails(): +@pytest.mark.usefixtures("python_router_core") # the Rust core keeps this lookup on its host registry +async def test_stamp_does_not_break_auth_when_the_access_group_lookup_fails(monkeypatch: pytest.MonkeyPatch): cache = await _cache() valid_token = UserAPIKeyAuth(api_key="hashed", models=["tier-a"], team_id=TEAM_ID, user_id=USER_ID) + router = Router(model_list=MODEL_LIST) + monkeypatch.setattr(router, "_core_get_model_access_groups", _broken_access_group_lookup) await stamp_matched_model_access_groups( model="gpt-4o", valid_token=valid_token, team_object=None, project_object=None, - llm_router=_BrokenRouter(model_list=MODEL_LIST), + llm_router=router, prisma_client=_ExplodingPrismaClient(), user_api_key_cache=cache, proxy_logging_obj=ProxyLogging(user_api_key_cache=cache), diff --git a/tests/unit/proxy/management_endpoints/test_auto_router_endpoints.py b/tests/unit/proxy/management_endpoints/test_auto_router_endpoints.py index c1505d3c8ec..59e12a97899 100644 --- a/tests/unit/proxy/management_endpoints/test_auto_router_endpoints.py +++ b/tests/unit/proxy/management_endpoints/test_auto_router_endpoints.py @@ -75,25 +75,25 @@ def _router() -> Router: return Router(model_list=ROUTER_MODEL_LIST) -class RecordingRouter(Router): +def _recording_router(classified_tier: str) -> tuple[Router, list[dict]]: """A real router that records the classifier calls the endpoint makes instead of sending them. Injected at the same `proxy_server.llm_router` boundary the endpoint reads, so model resolution and the key's model-access checks still run against a genuine Router. """ + recorded_calls: Final[list[dict]] = [] - def __init__(self, classified_tier: str) -> None: - super().__init__(model_list=ROUTER_MODEL_LIST) - self.classified_tier = classified_tier - self.recorded_calls: list[dict] = [] - - async def acompletion(self, model, messages, stream=False, **kwargs): - self.recorded_calls.append({"model": model, "messages": messages, **kwargs}) + async def record(model, messages, stream=False, **kwargs): + recorded_calls.append({"model": model, "messages": messages, **kwargs}) return ModelResponse( - choices=[Choices(message=Message(content=f'{{"tier": "{self.classified_tier}"}}'))], + choices=[Choices(message=Message(content=f'{{"tier": "{classified_tier}"}}'))], model=model, ) + router: Final = Router(model_list=ROUTER_MODEL_LIST) + router.acompletion = AsyncMock(side_effect=record) + return router, recorded_calls + def _request_from(body: Mapping[str, object], **config_overrides: object) -> AutoRouterRoutingTestRequest: return AutoRouterRoutingTestRequest.model_validate( @@ -137,7 +137,7 @@ async def _classifier_user_payload(body: Mapping[str, object], monkeypatch: pyte """The variable half of the classifier call this body produces.""" from litellm.proxy import proxy_server - router = RecordingRouter("SIMPLE") + router, recorded_calls = _recording_router("SIMPLE") monkeypatch.setattr(proxy_server, "llm_router", router) await preview_auto_router_routing( @@ -145,7 +145,7 @@ async def _classifier_user_payload(body: Mapping[str, object], monkeypatch: pyte data=_request_from(body, classifier_type="llm", classifier_llm_config={"model": "classifier-model"}), user_api_key_dict=ADMIN, ) - return router.recorded_calls[0]["messages"][1]["content"] + return recorded_calls[0]["messages"][1]["content"] @pytest.mark.parametrize( @@ -2710,7 +2710,7 @@ async def test_list_shadow_eval_jobs_collapses_legs_into_jobs_newest_first(monke async def test_jev_test_routing_authorizes_paid_evaluation_before_contacting_typesafe( monkeypatch: pytest.MonkeyPatch, denial: str | None ) -> None: - router: Final = RecordingRouter("SIMPLE") + router, recorded_calls = _recording_router("SIMPLE") monkeypatch.setattr(proxy_server, "llm_router", router) monkeypatch.setenv("TYPESAFE_API_KEY", "test") monkeypatch.setenv("TYPESAFE_API_BASE", "https://typesafe.test") @@ -2765,7 +2765,7 @@ async def test_jev_test_routing_authorizes_paid_evaluation_before_contacting_typ assert response.routing_decision["cause"] == "jev_classifier" assert response.routed_model == "cheap-model" assert evaluation.call_count == 1 - assert router.recorded_calls == [] + assert recorded_calls == [] await handler.client.aclose() @@ -2775,7 +2775,7 @@ async def test_jev_test_routing_authorizes_paid_evaluation_before_contacting_typ ["allowed", "credential-free", "team-admin-environment", "missing", "blocked", "key", "budget", "team", "not-router"], ) async def test_saved_jev_probe_uses_authorized_server_configuration(monkeypatch: pytest.MonkeyPatch, case: str) -> None: - router: Final = RecordingRouter("SIMPLE") + router, recorded_calls = _recording_router("SIMPLE") stored_key: Final = "synthetic-server-jev-key" monkeypatch.setenv("SAVED_JEV_KEY", stored_key) monkeypatch.setenv("SAVED_JEV_BASE", "https://saved-jev.test") @@ -2876,7 +2876,7 @@ async def test_saved_jev_probe_uses_authorized_server_configuration(monkeypatch: assert evaluation.calls.last.request.headers["authorization"] == f"Bearer {stored_key}" assert stored_key not in result.model_dump_json() assert evaluation.call_count == (1 if case in ("allowed", "credential-free", "team-admin-environment") else 0) - assert router.recorded_calls == [] + assert recorded_calls == [] await handler.client.aclose() @@ -3513,7 +3513,7 @@ async def test_member_billable_preview_checks_and_charges_destination_team( auth_module: Final = importlib.import_module("litellm.proxy.auth.user_api_key_auth") actor: Final = _configure_member_preview(monkeypatch).model_copy(update={"metadata": {"tags": ["key-tag"]}}) - router: Final = RecordingRouter("SIMPLE") + router, recorded_calls = _recording_router("SIMPLE") monkeypatch.setattr(proxy_server, "llm_router", router) async def check_and_tag( @@ -3547,13 +3547,13 @@ async def test_member_billable_preview_checks_and_charges_destination_team( if over_budget: with pytest.raises(litellm.BudgetExceededError): await preview_auto_router_routing(data, actor, http_request) - assert router.recorded_calls == [] + assert recorded_calls == [] else: await preview_auto_router_routing(data, actor, http_request) - assert len(router.recorded_calls) == 1 - assert router.recorded_calls[0]["metadata"]["user_api_key_team_id"] == "member-preview-team" - assert router.recorded_calls[0]["metadata"]["user_api_key_user_id"] == "preview-member" - assert set(router.recorded_calls[0]["metadata"]["tags"]) == {"key-tag", "header-tag"} + assert len(recorded_calls) == 1 + assert recorded_calls[0]["metadata"]["user_api_key_team_id"] == "member-preview-team" + assert recorded_calls[0]["metadata"]["user_api_key_user_id"] == "preview-member" + assert set(recorded_calls[0]["metadata"]["tags"]) == {"key-tag", "header-tag"} checks.assert_awaited_once() assert checks.await_args.kwargs["user_api_key_auth_obj"].team_id == "member-preview-team" assert checks.await_args.kwargs["route"] == "/auto_router/test_routing" diff --git a/tests/unit/router_rust/__init__.py b/tests/unit/router_rust/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/router_rust/test_surface_fallbacks.py b/tests/unit/router_rust/test_surface_fallbacks.py new file mode 100644 index 00000000000..dec6e864dae --- /dev/null +++ b/tests/unit/router_rust/test_surface_fallbacks.py @@ -0,0 +1,11 @@ +from dataclasses import fields + +from litellm.router_rust.surface.fallbacks import MOCK_FLAGS +from litellm.types.router import MockRouterTestingParams + + +def test_mock_flags_cover_every_mock_testing_param(): + """The Rust core serves exactly the ``mock_testing_*`` flags the Python core reads: the + ``MockRouterTestingParams`` fields and ``mock_testing_rate_limit_error``. A new flag fails + here until ``MOCK_FLAGS`` serves it.""" + assert set(MOCK_FLAGS) == {f.name for f in fields(MockRouterTestingParams)} | {"mock_testing_rate_limit_error"} diff --git a/tests/unit/router_utils/test_fallback_event_handlers.py b/tests/unit/router_utils/test_fallback_event_handlers.py index e9dd91d4362..a386542d61b 100644 --- a/tests/unit/router_utils/test_fallback_event_handlers.py +++ b/tests/unit/router_utils/test_fallback_event_handlers.py @@ -257,23 +257,22 @@ _FIRST_HOP_FAILURE_MARKER: Final = datetime(2026, 1, 2) _FINAL_HOP_FAILURE_MARKER: Final = datetime(2026, 1, 3) -class _TwoFailureFallbackRouter(Router): - def __init__(self) -> None: - super().__init__( - model_list=[ - { - "model_name": "fallback-model", - "litellm_params": {"model": "openai/gpt-5.6", "api_key": "fake-key"}, - "model_info": {"id": f"fallback-dep-{index}"}, - } - for index in (1, 2) - ], - allowed_fails=0, - cooldown_time=60, - num_retries=0, - ) +def _two_failure_fallback_router() -> Router: + router: Final = Router( + model_list=[ + { + "model_name": "fallback-model", + "litellm_params": {"model": "openai/gpt-5.6", "api_key": "fake-key"}, + "model_info": {"id": f"fallback-dep-{index}"}, + } + for index in (1, 2) + ], + allowed_fails=0, + cooldown_time=60, + num_retries=0, + ) - async def async_function_with_fallbacks(self, *args: object, **kwargs: object) -> NoReturn: + async def fail_both_hops(*args: object, **kwargs: object) -> NoReturn: logging_obj: Final[object] = kwargs["litellm_logging_obj"] assert isinstance(logging_obj, Logging) assert logging_obj.model_call_details.get("has_logged_async_failure", False) is False @@ -287,6 +286,9 @@ class _TwoFailureFallbackRouter(Router): logging_obj.model_call_details["current_attempt_api_call_start_time"] = _FINAL_HOP_FAILURE_MARKER raise _RateLimitedHop("fallback-dep-2") + router.async_function_with_fallbacks = fail_both_hops + return router + def _two_deployment_router() -> Router: return Router( @@ -397,8 +399,9 @@ async def test_run_async_fallback_cools_down_a_failed_hop_once(failure_already_l @pytest.mark.asyncio +@pytest.mark.usefixtures("python_router_core") # the Rust core walks fallbacks in Rust, so it has no such member async def test_run_async_fallback_cools_down_the_final_deployment_after_retry_failure_logging() -> None: - router: Final = _TwoFailureFallbackRouter() + router: Final = _two_failure_fallback_router() logging_obj: Final = _request_logging_obj() assert logging_obj.model_call_details.get("has_logged_async_failure", False) is False diff --git a/tests/unit/test_private_usage_aliases.py b/tests/unit/test_private_usage_aliases.py index 3c0aab8a4b4..a3168ad6d36 100644 --- a/tests/unit/test_private_usage_aliases.py +++ b/tests/unit/test_private_usage_aliases.py @@ -5242,7 +5242,6 @@ def test_public_aliases( resolved_owner: Final = _get_owner(module_path, owner_name) alias_owner: Final = _get_instance(resolved_owner, new_name) if use_instance else resolved_owner old_value, new_value = getattr(alias_owner, old_name), getattr(alias_owner, new_name) - # Router's private aliases wrap the public method so they skip a subclass override of it. if isinstance(old_value, MethodType) and isinstance(new_value, MethodType): assert old_value.__func__ is new_value.__func__ or inspect.unwrap(old_value.__func__) is new_value.__func__ else: diff --git a/tests/unit/test_router/core/conftest.py b/tests/unit/test_router/core/conftest.py index dffa938c441..3477cfc0c3d 100644 --- a/tests/unit/test_router/core/conftest.py +++ b/tests/unit/test_router/core/conftest.py @@ -1,4 +1,6 @@ -from collections.abc import Callable +from collections.abc import Callable, Generator +from pathlib import Path +from typing import Final import pytest @@ -15,12 +17,53 @@ def host() -> ScriptedHost: def make_core(request: pytest.FixtureRequest, host: ScriptedHost) -> Callable[..., RouterCore]: """``make_core(**router_arguments)``: the core under test, on the ``host`` fixture. Every test that uses it runs once per core in ``CORE_BUILDERS`` (``-k python`` or ``-k rust`` picks one).""" - builder = CORE_BUILDERS[request.param] + name: Final[str] = str(request.param) # pyright: ignore[reportAny] # the fixture parameter + builder: Final = CORE_BUILDERS[name] def build(**arguments: object) -> RouterCore: # kwargs-ok: the router's constructor arguments try: return builder(host, **arguments) except CoreDeclined as declined: - pytest.skip(f"{request.param} core declines this configuration: {declined}") + pytest.skip(f"{name} core declines this configuration: {declined}") return build + + +_RUST_PARITY_PENDING: Final = frozenset( + line.strip() + for line in (Path(__file__).parent / "rust_parity_pending.txt").read_text().splitlines() + if line.strip() and not line.startswith("#") +) + + +def pytest_collection_modifyitems(items: list[pytest.Item]) -> None: + """Characterization tests the Rust core does not match yet are listed in + ``rust_parity_pending.txt`` and expected to fail (strictly: one that starts passing must leave the list).""" + here: Final = Path(__file__).parent + for item in items: + path: Final = Path(str(item.path)) + if not path.is_relative_to(here): + continue # the hook sees every collected item, not only this folder's + name: Final = "::".join((path.relative_to(here).as_posix(), *item.nodeid.split("::")[1:])) + if name in _RUST_PARITY_PENDING: + item.add_marker(pytest.mark.xfail(reason="Rust core parity pending", strict=True)) + + +@pytest.hookimpl(hookwrapper=True) +def pytest_runtest_call(item: pytest.Item) -> Generator[None, object, None]: + """A test that reaches a member or request the Rust core does not serve yet is skipped with + that member named, so the run separates "not served yet" from a wrong answer.""" + outcome: Final = yield + excinfo: Final[object] = getattr(outcome, "excinfo", None) + if not isinstance(excinfo, tuple): + return + info: Final[tuple[object, ...]] = excinfo # pyright: ignore[reportUnknownVariableType] # pluggy's (type, value, traceback) + if len(info) < 2: + return + from litellm.router_rust.core import RustCoreUnsupported + + error: Final[object] = info[1] + cause: Final[object] = getattr(error, "__cause__", None) + force_exception: Final[object] = getattr(outcome, "force_exception", None) + if (isinstance(error, RustCoreUnsupported) or isinstance(cause, RustCoreUnsupported)) and callable(force_exception): + force_exception(pytest.skip.Exception(f"not served by the Rust core yet: {error}", _use_item_location=True)) diff --git a/tests/unit/test_router/core/construction_precall_helpers.py b/tests/unit/test_router/core/construction_precall_helpers.py new file mode 100644 index 00000000000..c83bd37d692 --- /dev/null +++ b/tests/unit/test_router/core/construction_precall_helpers.py @@ -0,0 +1,98 @@ +""" +Helpers for the construction / pre-call characterization tests. + +``CallbackHost`` is a ``ScriptedHost`` whose callback manager also has +``remove_callback_from_list_by_object`` (litellm's ``LoggingCallbackManager`` has it; ``HostCallbacks`` +does not), so the path that switches optional pre-call checks off at runtime can run. +""" + +import copy +import gc +import uuid +from collections.abc import Generator, Mapping +from typing import Final + +import litellm +from tests.unit.test_router.core.scripted_host import HostCallbacks, ScriptedHost, deployment + + +class RemovingHostCallbacks(HostCallbacks): + def remove_callback_from_list_by_object( + self, + callback_list: object, + obj: object, + require_self: bool = True, + ) -> None: + """``LoggingCallbackManager.remove_callback_from_list_by_object``, on the host's lists.""" + if not isinstance(callback_list, list): + return + if require_self: + doomed = [c for c in callback_list if getattr(c, "__self__", None) is obj] + else: + doomed = [c for c in callback_list if c == obj] + for callback in doomed: + callback_list.remove(callback) + + +class CallbackHost(ScriptedHost): + __slots__ = () + + @property + def callbacks(self) -> RemovingHostCallbacks: # pyright: ignore[reportIncompatibleVariableOverride] + return RemovingHostCallbacks(self) + + +def callback_names(callbacks: list[object]) -> list[str]: # mutable-ok: a fresh list + """Each callback by name: a bound method by its function name, anything else by its class name.""" + names: Final[list[str]] = [] + for callback in callbacks: + function = getattr(callback, "__func__", None) + names.append(function.__name__ if function is not None else type(callback).__name__) + return names + + +def host_callback_names(host: ScriptedHost) -> dict[str, list[str]]: # mutable-ok: a fresh dict + return {name: callback_names(callbacks) for name, callbacks in host.lists.items()} + + +def windowed( + group: str, deployment_id: str, max_input_tokens: int, **params: object +) -> dict[str, object]: # mutable-ok: a model_list entry + """A deployment declaring ``max_input_tokens`` under a model name no other test uses: a deployment's + model_info is written into the process-wide cost map under its model name and id.""" + model_info: Final = {"max_input_tokens": max_input_tokens, **dict(params.pop("model_info", None) or {})} # pyright: ignore[reportArgumentType] + return deployment( + group, + deployment_id, + model=f"openai/cw-{deployment_id}-{uuid.uuid4().hex[:8]}", + model_info=model_info, + **params, + ) + + +def collect_other_routers() -> None: + """Drop unreachable routers from earlier tests: retire() and runtime pre-call-check changes look + at every live router in the process, not only the ones on this host.""" + gc.collect() + + +def restoring_model_cost(*names: str) -> Generator[None, None, None]: + """Undo what routers built meanwhile wrote into the process cost map: keys they added are removed + and the named entries get their old value back.""" + added_before: Final = frozenset(litellm.model_cost) + saved: Final = {name: copy.deepcopy(litellm.model_cost[name]) for name in names if name in litellm.model_cost} + yield + for key in frozenset(litellm.model_cost) - added_before: + litellm.model_cost.pop(key, None) + litellm.model_cost.update(saved) + + +__all__ = [ + "CallbackHost", + "RemovingHostCallbacks", + "callback_names", + "collect_other_routers", + "host_callback_names", + "restoring_model_cost", + "windowed", +] diff --git a/tests/unit/test_router/core/endpoints_helpers.py b/tests/unit/test_router/core/endpoints_helpers.py new file mode 100644 index 00000000000..cc79e05a686 --- /dev/null +++ b/tests/unit/test_router/core/endpoints_helpers.py @@ -0,0 +1,212 @@ +""" +Helpers for the endpoint characterization tests (``test_endpoints_characterization*.py``). + +``EndpointHost`` is a ``ScriptedHost`` that also scripts the host ops ``ScriptedHost`` leaves to +litellm (assistants, skills, realtime, fine-tuning, batches, ...), so every ``host.invoke`` the core +makes is answered and recorded. + +Some endpoints do not go through ``host.invoke``: the Python core binds litellm's own functions for +them when it is built (vector stores, Google GenAI, OCR, search, videos, containers, interactions, +managed agents). ``patch_direct_endpoints`` replaces those litellm functions, before the core is +built, with recording stand-ins that answer from the same host script, under the same ``__name__`` +(the core routes on the handler's name). A call to one of them is recorded with the function's name +as its ``op``. +""" + +import importlib +from collections.abc import Callable, Mapping +from typing import Final + +import pytest + +from tests.unit.test_router.core.scripted_host import INFERENCE_OPS, ScriptedHost + +EXTRA_HOST_OPS: Final = frozenset( + { + "aadapter_completion", + "acreate_assistants", + "adelete_assistant", + "aget_assistants", + "acreate_thread", + "aget_thread", + "a_add_message", + "aget_messages", + "arun_thread", + "acompact_responses", + "_arealtime", + "acreate_realtime_client_secret", + "arealtime_calls", + "acreate_realtime_transcription_session", + "acreate_realtime_translation_client_secret", + "arealtime_translation_calls", + "_aresponses_websocket", + "acancel_fine_tuning_job", + "alist_fine_tuning_jobs", + "aretrieve_fine_tuning_job", + "afile_list", + "aimage_edit", + "allm_passthrough_route", + "acreate_skill", + "alist_skills", + "aget_skill", + "adelete_skill", + "aretrieve_batch", + "acancel_batch", + "alist_batches", + } +) + + +class EndpointHost(ScriptedHost): + """``ScriptedHost`` that scripts every host op the core's endpoints invoke.""" + + __slots__ = () + + def invoke(self, op: str) -> Callable[..., object]: + if op in INFERENCE_OPS or op not in EXTRA_HOST_OPS: + return super().invoke(op) + return recording_endpoint(self, op, is_async=True) + + +def recording_endpoint(host: ScriptedHost, name: str, *, is_async: bool) -> Callable[..., object]: + """A stand-in for litellm's ``name`` that records the call on ``host`` and answers from its script.""" + if is_async: + + async def endpoint(*args: object, **kwargs: object) -> object: # kwargs-ok: the endpoint's request + return await host._answer_async(name, args, kwargs) # pyright: ignore[reportPrivateUsage] # the test host + + endpoint.__name__ = name + endpoint.__qualname__ = name + return endpoint + + def sync_endpoint(*args: object, **kwargs: object) -> object: # kwargs-ok: the endpoint's request + return host._answer_sync(name, args, kwargs) # pyright: ignore[reportPrivateUsage] # the test host + + sync_endpoint.__name__ = name + sync_endpoint.__qualname__ = name + return sync_endpoint + + +# module -> {attribute on the module: the function's own name} +DIRECT_ENDPOINTS: Final[Mapping[str, Mapping[str, str]]] = { + "litellm.vector_stores.main": { + name: name + for name in ( + "asearch", + "search", + "acreate", + "create", + "aretrieve", + "retrieve", + "alist", + "list", + "aupdate", + "update", + "adelete", + "delete", + ) + }, + "litellm.vector_stores": {"acreate": "acreate"}, + "litellm.vector_store_files.main": { + name: name + for name in ( + "acreate", + "create", + "alist", + "list", + "aretrieve", + "retrieve", + "aretrieve_content", + "retrieve_content", + "aupdate", + "update", + "adelete", + "delete", + ) + }, + "litellm.google_genai": { + name: name + for name in ("agenerate_content", "generate_content", "agenerate_content_stream", "generate_content_stream") + }, + "litellm.ocr": {"aocr": "aocr", "ocr": "ocr"}, + "litellm.search": {"asearch": "asearch", "search": "search"}, + "litellm.decisions": {"adecisions": "adecisions", "decisions": "decisions"}, + "litellm.systemone": {"asystemone": "asystemone", "systemone": "systemone"}, + "litellm.extraction": {"aextraction": "aextraction", "extraction": "extraction"}, + "litellm.videos": { + f"{prefix}video_{name}": f"{prefix}video_{name}" + for prefix in ("a", "") + for name in ( + "generation", + "list", + "status", + "content", + "remix", + "create_character", + "get_character", + "edit", + "extension", + ) + }, + "litellm.containers": { + f"{prefix}{name}": f"{prefix}{name}" + for prefix in ("a", "") + for name in ("create_container", "list_containers", "retrieve_container", "delete_container") + }, + "litellm.interactions": { + f"{prefix}{name}": f"{prefix}{name}" for prefix in ("a", "") for name in ("create", "get", "delete", "cancel") + }, + "litellm.interactions.agents": { + f"{prefix}{name}": f"{prefix}{name}" + for prefix in ("a", "") + for name in ("create", "get", "delete", "list", "list_versions") + }, +} + +# Recorded op names for the direct endpoints whose bare names collide across modules. +RECORDED_NAMES: Final[Mapping[tuple[str, str], str]] = { + **{("litellm.vector_stores.main", name): f"vs.{name}" for name in DIRECT_ENDPOINTS["litellm.vector_stores.main"]}, + ("litellm.vector_stores", "acreate"): "vs.acreate", + **{ + ("litellm.vector_store_files.main", name): f"vsf.{name}" + for name in DIRECT_ENDPOINTS["litellm.vector_store_files.main"] + }, + **{("litellm.interactions", name): f"interactions.{name}" for name in DIRECT_ENDPOINTS["litellm.interactions"]}, + **{ + ("litellm.interactions.agents", name): f"agents.{name}" + for name in DIRECT_ENDPOINTS["litellm.interactions.agents"] + }, +} + + +def patch_direct_endpoints(monkeypatch: pytest.MonkeyPatch, host: ScriptedHost) -> None: + """Replace litellm's directly bound endpoint functions with recording stand-ins on ``host``. + + The stand-in keeps the function's ``__name__`` (the core routes on it); the call is recorded + under ``RECORDED_NAMES`` when the bare name is ambiguous, else under the function's name.""" + for module_name, attributes in DIRECT_ENDPOINTS.items(): + module: Final = importlib.import_module(module_name) + for attribute, function_name in attributes.items(): + original: Final[object] = getattr(module, attribute) + is_async: Final = function_name.startswith("a") + recorded: Final = RECORDED_NAMES.get((module_name, attribute), function_name) + stand_in: Final = recording_endpoint(host, recorded, is_async=is_async) + stand_in.__name__ = getattr(original, "__name__", function_name) + monkeypatch.setattr(module, attribute, stand_in) + from litellm.containers import endpoint_factory + + for name in list(endpoint_factory._generated_endpoints): # pyright: ignore[reportPrivateUsage] # litellm's registry + monkeypatch.setitem( + endpoint_factory._generated_endpoints, # pyright: ignore[reportPrivateUsage] # litellm's registry + name, + recording_endpoint(host, name, is_async=name.startswith("a")), + ) + + +__all__ = [ + "DIRECT_ENDPOINTS", + "EXTRA_HOST_OPS", + "EndpointHost", + "patch_direct_endpoints", + "recording_endpoint", +] diff --git a/tests/unit/test_router/core/fallbacks_helpers.py b/tests/unit/test_router/core/fallbacks_helpers.py new file mode 100644 index 00000000000..c88c31a9371 --- /dev/null +++ b/tests/unit/test_router/core/fallbacks_helpers.py @@ -0,0 +1,105 @@ +"""Helpers for the fallbacks characterization tests (``test_fallbacks_characterization*.py``).""" + +from collections.abc import Callable, Mapping +from typing import Final + +from litellm.router_core import RouterCore +from litellm.router_python_host import Endpoint +from tests.unit.test_router.core.scripted_host import Outcome, ScriptedHost, deployment + +MakeCore = Callable[..., RouterCore] + +MESSAGES: Final = [{"role": "user", "content": "hi"}] +GROUPS_ABC: Final = [deployment("a", "a"), deployment("b", "b"), deployment("c", "c")] +GROUP_A_OF_THREE: Final = [deployment("a", "a1"), deployment("a", "a2"), deployment("a", "a3"), deployment("b", "b")] +TIERED_A: Final = [ + deployment("a", "a1", order=1), + deployment("a", "a2", order=2), + deployment("b", "b"), + deployment("c", "c"), +] + + +def build(make_core: MakeCore, **arguments: object) -> RouterCore: # kwargs-ok: the router's constructor arguments + return make_core(**{"model_list": GROUPS_ABC, "num_retries": 0, **arguments}) + + +async def complete(core: RouterCore, model: str = "a", **request: object) -> object: # kwargs-ok: request kwargs + return await core.acompletion(model=model, messages=MESSAGES, **request) # pyright: ignore[reportAttributeAccessIssue] # endpoint method + + +def complete_sync(core: RouterCore, model: str = "a", **request: object) -> object: # kwargs-ok: request kwargs + return core.completion(model=model, messages=MESSAGES, **request) # pyright: ignore[reportAttributeAccessIssue] # endpoint method + + +def hidden(response: object) -> Mapping[str, object]: + return getattr(response, "_hidden_params") + + +def headers(response: object) -> Mapping[str, object]: + found: Final = hidden(response)["additional_headers"] + assert isinstance(found, Mapping) + return found + + +def metadata_of(kwargs: Mapping[str, object]) -> Mapping[str, object]: + found: Final = kwargs.get("metadata") or kwargs.get("litellm_metadata") or {} + assert isinstance(found, Mapping) + return found + + +def show_router_debug(host: ScriptedHost) -> None: + host.set_dynamic(expose_router_debug_in_errors=True) + + +def fails_once_then_answers(error: BaseException) -> Callable[..., Outcome]: + """One scripted outcome shared by several models: the first call to any of them raises + ``error``, every later call answers the default response.""" + calls: Final = [0] + + def outcome(**_request: object) -> Outcome: # kwargs-ok: the scripted request + calls[0] += 1 + return error if calls[0] == 1 else None + + return outcome + + +class NamedEndpointHost(ScriptedHost): + """A ``ScriptedHost`` whose ``anthropic_messages`` endpoint carries that name, as litellm's + own function does. The core recognises a /v1/messages call (and its safeguard-refusal gate) + by the endpoint's ``__name__``; ``ScriptedHost``'s endpoints are all named ``endpoint``.""" + + __slots__ = () + + def invoke(self, op: str) -> Endpoint: + endpoint: Final = super().invoke(op) + if op == "anthropic_messages": + endpoint.__name__ = op # pyright: ignore[reportFunctionMemberAccess] # a plain function + return endpoint + + +def anthropic_refusal(*, with_stop_details: bool = True) -> dict[str, object]: # mutable-ok: a provider response + response: Final[dict[str, object]] = { + "id": "msg_1", + "type": "message", + "role": "assistant", + "model": "claude", + "content": [], + "stop_reason": "refusal", + "usage": {"input_tokens": 1, "output_tokens": 1}, + } + if with_stop_details: + response["stop_details"] = {"type": "refusal", "category": None, "explanation": "no"} + return response + + +def anthropic_message(text: str = "fine") -> dict[str, object]: # mutable-ok: a provider response + return { + "id": "msg_2", + "type": "message", + "role": "assistant", + "model": "claude", + "content": [{"type": "text", "text": text}], + "stop_reason": "end_turn", + "usage": {"input_tokens": 1, "output_tokens": 1}, + } diff --git a/tests/unit/test_router/core/gaps_l11100_helpers.py b/tests/unit/test_router/core/gaps_l11100_helpers.py new file mode 100644 index 00000000000..38fb2006503 --- /dev/null +++ b/tests/unit/test_router/core/gaps_l11100_helpers.py @@ -0,0 +1,112 @@ +"""Shared helpers for ``test_gaps_l11100*.py``: strategy-router marker deployments and a host whose +callback registration can also remove a callback from every list (what the adaptive router's hook +sync asks of litellm's ``LoggingCallbackManager``).""" + +from collections.abc import Callable, Mapping +from typing import Final + +from litellm.router_core import RouterCore +from tests.unit.test_router.core.scripted_host import HostCallbacks, ScriptedHost + +MakeCore = Callable[..., RouterCore] + +TIERS: Final = {"SIMPLE": "small", "MEDIUM": "big", "COMPLEX": "big", "REASONING": "big"} + + +class RemovingHostCallbacks(HostCallbacks): + """``HostCallbacks`` plus ``remove_callback_from_all_lists``, as litellm's callback manager has it.""" + + def __init__(self, host: "RemovingHost") -> None: + super().__init__(host) + self._owner: Final = host + + def remove_callback_from_all_lists(self, obj: object, require_self: bool = False) -> None: + del require_self + for name, callbacks in self._owner.lists.items(): + self._owner.lists[name] = [callback for callback in callbacks if callback is not obj] + + +class RemovingHost(ScriptedHost): + """A ``ScriptedHost`` whose callbacks can be removed from every list.""" + + __slots__ = () + + @property + def callbacks(self) -> RemovingHostCallbacks: # pyright: ignore[reportIncompatibleVariableOverride] # same methods plus one + return RemovingHostCallbacks(self) + + +def complexity_marker( + name: str = "smart", + deployment_id: str = "marker", + *, + config: Mapping[str, object] | None = None, + **params: object, # kwargs-ok: extra marker litellm_params +) -> dict[str, object]: # mutable-ok: a model_list entry + """An ``auto_router/complexity_router`` marker with the shipped tiers and default model ``big``.""" + return { + "model_name": name, + "litellm_params": { + "model": "auto_router/complexity_router", + "complexity_router_config": dict(config) if config is not None else {"tiers": dict(TIERS)}, + "complexity_router_default_model": "big", + **params, + }, + "model_info": {"id": deployment_id}, + } + + +def adaptive_marker( + name: str, deployment_id: str, available_models: list[str], **params: object +) -> dict[str, object]: # mutable-ok: a model_list entry # kwargs-ok: extra marker litellm_params + """An ``auto_router/adaptive_router`` marker over ``available_models``.""" + return { + "model_name": name, + "litellm_params": { + "model": "auto_router/adaptive_router", + "adaptive_router_config": {"available_models": available_models}, + **params, + }, + "model_info": {"id": deployment_id}, + } + + +def marker( + name: str, deployment_id: str, model: str, **params: object +) -> dict[str, object]: # mutable-ok: a model_list entry # kwargs-ok: marker litellm_params + """Any strategy-router marker: ``model`` is its ``auto_router/...`` litellm model.""" + return {"model_name": name, "litellm_params": {"model": model, **params}, "model_info": {"id": deployment_id}} + + +def registered_tags(core: RouterCore, kind: str) -> dict[str, list[tuple[str, ...]]]: # mutable-ok: a fresh dict + """``{router name: [tags of each registered strategy]}`` for a pre-routing strategy registry.""" + registry: Final = core.registry().plugins(kind) # pyright: ignore[reportCallIssue, reportArgumentType] # one of the strategy kinds + assert isinstance(registry, dict) + return {name: [entry.tags for entry in entries] for name, entries in registry.items()} # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] # TaggedPreRoutingStrategy lists + + +def model_ids(core: RouterCore) -> list[str]: # mutable-ok: a fresh list + """The ``model_info.id`` of every deployment on the model list, in order.""" + return [str(entry["model_info"]["id"]) for entry in core.registry().all()] # pyright: ignore[reportAny] # model_list rows + + +class ModelsEndpoint: + """A stand-in for the ``AsyncHTTPHandler`` that model-info discovery reads ``/v1/models`` with: + answers ``body`` (or raises ``error``) and records each request's url and headers.""" + + def __init__(self, body: Mapping[str, object] | None = None, error: BaseException | None = None) -> None: + self.body: Final = body or {"data": []} + self.error: Final = error + self.requests: Final[list[tuple[str, dict[str, str]]]] = [] + + async def get( + self, url: str, headers: dict[str, str] | None = None, **_: object + ) -> object: # kwargs-ok: httpx options + import json + + import httpx + + self.requests.append((url, dict(headers or {}))) + if self.error is not None: + raise self.error + return httpx.Response(200, content=json.dumps(self.body).encode(), request=httpx.Request("GET", url)) diff --git a/tests/unit/test_router/core/registry_helpers.py b/tests/unit/test_router/core/registry_helpers.py new file mode 100644 index 00000000000..aa154454b8d --- /dev/null +++ b/tests/unit/test_router/core/registry_helpers.py @@ -0,0 +1,47 @@ +"""Shared helpers for the registry characterization tests (``test_registry_characterization*.py``).""" + +import uuid +from collections.abc import Callable, Iterable, Mapping + +from litellm.router_core import RouterCore +from litellm.types.router import Deployment +from tests.unit.test_router.core.scripted_host import deployment + +MakeCore = Callable[..., RouterCore] + +MESSAGES = [{"role": "user", "content": "hi"}] + + +def ids(deployments: Iterable[Mapping[str, object]] | None) -> list[str]: # mutable-ok: a fresh list + """The ``model_info.id`` of each deployment dict, in order.""" + return [str(entry["model_info"]["id"]) for entry in deployments or ()] # pyright: ignore[reportIndexIssue] # deployment dicts + + +def names_and_ids( + deployments: Iterable[Mapping[str, object]] | None, +) -> list[tuple[str, str]]: # mutable-ok: a fresh list + """``(model_name, model_info.id)`` of each deployment dict, in order.""" + return [(str(entry["model_name"]), str(entry["model_info"]["id"])) for entry in deployments or ()] # pyright: ignore[reportIndexIssue] # deployment dicts + + +def typed(model_name: str, deployment_id: str, **params: object) -> Deployment: # kwargs-ok: litellm_params + """A ``Deployment`` built from the same entry ``scripted_host.deployment`` returns.""" + return Deployment(**deployment(model_name, deployment_id, **params)) # pyright: ignore[reportArgumentType] # a model_list entry + + +def team_deployment( + model_name: str, deployment_id: str, team_id: str, public_name: str, **params: object +) -> dict[str, object]: # mutable-ok: a model_list entry # kwargs-ok: litellm_params + """A team-owned deployment: ``model_info.team_id`` and ``team_public_model_name`` set.""" + extra_info: dict[str, object] = dict(params.pop("model_info", None) or {}) # pyright: ignore[reportArgumentType] # a model_info dict + return deployment( + model_name, + deployment_id, + model_info={"team_id": team_id, "team_public_model_name": public_name, **extra_info}, + **params, + ) + + +def unique_id(prefix: str = "reg") -> str: + """A deployment id no other test uses, for tests that write into the process-wide ``litellm.model_cost``.""" + return f"{prefix}-{uuid.uuid4().hex}" diff --git a/tests/unit/test_router/core/retries_health_helpers.py b/tests/unit/test_router/core/retries_health_helpers.py new file mode 100644 index 00000000000..eae0b836530 --- /dev/null +++ b/tests/unit/test_router/core/retries_health_helpers.py @@ -0,0 +1,198 @@ +""" +Helpers for the retries / deployment-health characterization tests (``test_retries_health_characterization*.py``). + +Backoff values carry litellm's random jitter (up to ``JITTER`` seconds), so sleeps are asserted inside +``[base, base + JITTER)``. +""" + +import asyncio +import copy +from collections.abc import Callable, Mapping +from typing import Final + +import httpx + +import litellm +from litellm.router_core import RouterCore +from tests.unit.test_router.core.scripted_host import Call, ScriptedHost + +MESSAGES: Final = [{"role": "user", "content": "hi"}] + +MakeCore = Callable[..., RouterCore] +ErrorBuilder = Callable[[], Exception] + +_REQUEST: Final = httpx.Request("POST", "https://api.openai.com/v1/chat/completions") + + +def fresh(builder: ErrorBuilder) -> Callable[..., Exception]: + """A scripted outcome that raises a new error per call.""" + return lambda **_: builder() + + +def http_response(status: int, headers: Mapping[str, str] | None = None) -> httpx.Response: + return httpx.Response(status, headers=dict(headers or {}), request=_REQUEST) + + +def not_found() -> litellm.NotFoundError: + return litellm.NotFoundError(message="no such model", llm_provider="openai", model="gpt-4o") + + +def permission_denied() -> litellm.PermissionDeniedError: + return litellm.PermissionDeniedError( + message="forbidden", llm_provider="openai", model="gpt-4o", response=http_response(403) + ) + + +def unprocessable() -> litellm.UnprocessableEntityError: + return litellm.UnprocessableEntityError( + message="unprocessable", llm_provider="openai", model="gpt-4o", response=http_response(422) + ) + + +def service_unavailable() -> litellm.ServiceUnavailableError: + return litellm.ServiceUnavailableError(message="unavailable", llm_provider="openai", model="gpt-4o") + + +def bad_gateway() -> litellm.BadGatewayError: + return litellm.BadGatewayError(message="bad gateway", llm_provider="openai", model="gpt-4o") + + +def api_connection() -> litellm.APIConnectionError: + return litellm.APIConnectionError(message="connection reset", llm_provider="openai", model="gpt-4o") + + +def api_status(status: int) -> ErrorBuilder: + return lambda: litellm.APIError( + status_code=status, message=f"status {status}", llm_provider="openai", model="gpt-4o" + ) + + +def plain_error() -> ValueError: + return ValueError("not a provider error") + + +def rate_limited(retry_after: str | None = None) -> litellm.RateLimitError: + headers: Final = {"retry-after": retry_after} if retry_after is not None else {} + return litellm.RateLimitError( + message="slow down", llm_provider="openai", model="gpt-4o", response=httpx.Response(429, headers=headers) + ) + + +async def attempt(core: RouterCore, model: str = "g", **kwargs: object) -> object: # kwargs-ok: request kwargs + """The response, or the error the core raised. A request the Rust core does not serve yet propagates, so the + test is skipped rather than failed.""" + from litellm.router_rust.core import RustCoreUnsupported + + try: + return await core.acompletion(model=model, messages=MESSAGES, **kwargs) # pyright: ignore[reportAttributeAccessIssue] # every core serves acompletion + except RustCoreUnsupported: + raise + except Exception as error: + return error + + +def attempt_sync(core: RouterCore, model: str = "g", **kwargs: object) -> object: # kwargs-ok: request kwargs + from litellm.router_rust.core import RustCoreUnsupported + + try: + return core.completion(model=model, messages=MESSAGES, **kwargs) # pyright: ignore[reportAttributeAccessIssue] # every core serves completion + except RustCoreUnsupported: + raise + except Exception as error: + return error + + +def cooled_down(core: RouterCore, ids: list[str]) -> dict[str, Mapping[str, object]]: # mutable-ok: a fresh dict + return dict(core.health().cooldowns.get_active_cooldowns(ids, None)) + + +def headers_of(response: object) -> Mapping[str, object]: + hidden: Final = getattr(response, "_hidden_params", None) or {} + return hidden.get("additional_headers") or {} + + +class MetadataRecorder: + """A scripted outcome factory that snapshots the request metadata each call carried, then answers.""" + + KEYS: Final = ( + "attempted_retries", + "max_retries", + "model_group", + "model_group_size", + "request_retry_count", + "previous_models", + ) + + def __init__(self) -> None: + self.snapshots: Final[list[dict[str, object]]] = [] + + def then(self, builder: ErrorBuilder | None) -> Callable[..., object]: + def outcome(**kwargs: object) -> object: # kwargs-ok: the request + metadata: Final = kwargs.get("metadata") or kwargs.get("litellm_metadata") or {} + assert isinstance(metadata, Mapping) + self.snapshots.append({key: copy.deepcopy(metadata.get(key)) for key in self.KEYS}) + return builder() if builder is not None else None + + return outcome + + +class BackgroundSuccessLoggingHost(ScriptedHost): + """``ScriptedHost`` whose async success callbacks run in a background task after the response is returned, + as litellm's logging worker runs them (``ScriptedHost`` runs them before returning). Await ``drain()`` + before reading what they wrote.""" + + __slots__ = ("_pending",) + + def __init__(self) -> None: + super().__init__() + self._pending: list[asyncio.Task[None]] = [] + + async def _answer_async(self, op: str, args: tuple[object, ...], kwargs: Mapping[str, object]) -> object: + call, outcome = self._outcome(op, args, kwargs) + error, response = self._finish(op, call, outcome) + if error is not None: + logged_failure: Final = self._logging_kwargs(call, error, response) + for callback in list(self.lists.get("failure_callback", ())): + if callable(callback): + callback(logged_failure, None, 0.0, 0.0) + for callback in list(self.lists.get("_async_failure_callback", ())): + if callable(callback): + await callback(logged_failure, None, 0.0, 0.0) + raise error + + async def log_success(logged_call: Call) -> None: + logged: Final = self._logging_kwargs(logged_call, None, response) + for callback in list(self.lists.get("_async_success_callback", ())): + if callable(callback): + await callback(logged, response, 0.0, 0.0) + + self._pending.append(asyncio.ensure_future(log_success(call))) + return response + + async def drain(self) -> None: + while self._pending: + await self._pending.pop() + + +__all__ = [ + "MESSAGES", + "BackgroundSuccessLoggingHost", + "ErrorBuilder", + "MakeCore", + "MetadataRecorder", + "api_connection", + "api_status", + "attempt", + "attempt_sync", + "bad_gateway", + "cooled_down", + "fresh", + "headers_of", + "http_response", + "not_found", + "permission_denied", + "plain_error", + "rate_limited", + "service_unavailable", + "unprocessable", +] diff --git a/tests/unit/test_router/core/rust_parity_pending.txt b/tests/unit/test_router/core/rust_parity_pending.txt new file mode 100644 index 00000000000..8cfcebad287 --- /dev/null +++ b/tests/unit/test_router/core/rust_parity_pending.txt @@ -0,0 +1,2 @@ +# Characterization tests the Rust core does not match yet, one node id per line (`file.py::test[rust]`). +# conftest.py marks each one xfail(strict=True), so a test that starts passing must leave the list. diff --git a/tests/unit/test_router/core/scripted_host.py b/tests/unit/test_router/core/scripted_host.py index af28160b620..f673ab98610 100644 --- a/tests/unit/test_router/core/scripted_host.py +++ b/tests/unit/test_router/core/scripted_host.py @@ -18,6 +18,7 @@ from typing import Final import litellm from litellm.integrations.custom_logger import CustomLogger +from litellm.litellm_core_utils.agentic_followup_failure import mark_agentic_followup_failure from litellm.litellm_core_utils.logging_callback_manager import LoggingCallbackManager from litellm.router_config import RouterDynamicConfig, build_router_config from litellm.router_core import RouterCore @@ -138,6 +139,12 @@ def server_error(model: str = "gpt-4o") -> litellm.InternalServerError: return litellm.InternalServerError(message="boom", llm_provider="openai", model=model) +def agentic_followup_failure(error: Exception) -> Exception: + """``error`` as an agentic follow-up raises it, after the first call was served.""" + mark_agentic_followup_failure(error, ()) + return error + + def timeout(model: str = "gpt-4o") -> litellm.Timeout: return litellm.Timeout(message="timed out", llm_provider="openai", model=model) @@ -324,6 +331,9 @@ class ScriptedHost(PythonHost): return done() + def sleep_sync(self, seconds: float) -> None: + self.sleeps.append(seconds) + def dynamic(self) -> RouterDynamicConfig: return self._dynamic @@ -398,9 +408,9 @@ class ScriptedHost(PythonHost): return outcome, None return None, _default_response(op, call.model) if outcome is None else outcome - async def _answer_async(self, op: str, args: tuple[object, ...], kwargs: Mapping[str, object]) -> object: - call, outcome = self._outcome(op, args, kwargs) - error, response = self._finish(op, call, outcome) + async def log_async(self, call: Call, error: BaseException | None, response: object) -> None: + """Run the callbacks an async call's logging runs: the failure callbacks for ``error``, else + the async success callbacks for ``response``.""" logged: Final = self._logging_kwargs(call, error, response) if error is not None: for callback in list(self.lists.get("failure_callback", ())): @@ -409,24 +419,50 @@ class ScriptedHost(PythonHost): for callback in list(self.lists.get("_async_failure_callback", ())): if callable(callback): await callback(logged, None, 0.0, 0.0) - raise error + return for callback in list(self.lists.get("_async_success_callback", ())): if callable(callback): await callback(logged, response, 0.0, 0.0) - return response - def _answer_sync(self, op: str, args: tuple[object, ...], kwargs: Mapping[str, object]) -> object: - call, outcome = self._outcome(op, args, kwargs) - error, response = self._finish(op, call, outcome) + def log_sync(self, call: Call, error: BaseException | None, response: object) -> None: + """``log_async`` for a sync call: the failure callbacks, else the sync success callbacks.""" logged: Final = self._logging_kwargs(call, error, response) if error is not None: for callback in list(self.lists.get("failure_callback", ())): if callable(callback): callback(logged, None, 0.0, 0.0) - raise error + return for callback in list(self.lists.get("success_callback", ())): if callable(callback): callback(logged, response, 0.0, 0.0) + + def _logs_at_stream_end(self, call: Call, response: object) -> bool: + """A stream whose logging object logs through this host (``streaming_helpers.HostLogging``) + logs when it ends, as litellm's stream logging does, so the call logs nothing when it opens.""" + bind: Final[object] = getattr(getattr(response, "logging_obj", None), "bind_host", None) + if not callable(bind): + return False + bind(self, call) + return True + + async def _answer_async(self, op: str, args: tuple[object, ...], kwargs: Mapping[str, object]) -> object: + call, outcome = self._outcome(op, args, kwargs) + error, response = self._finish(op, call, outcome) + if error is not None: + await self.log_async(call, error, response) + raise error + if not self._logs_at_stream_end(call, response): + await self.log_async(call, None, response) + return response + + def _answer_sync(self, op: str, args: tuple[object, ...], kwargs: Mapping[str, object]) -> object: + call, outcome = self._outcome(op, args, kwargs) + error, response = self._finish(op, call, outcome) + if error is not None: + self.log_sync(call, error, response) + raise error + if not self._logs_at_stream_end(call, response): + self.log_sync(call, None, response) return response @@ -452,25 +488,46 @@ def python_core(host: ScriptedHost, **arguments: object) -> PythonCore: # kwarg return PythonCore(**config.arguments, _host=host) # pyright: ignore[reportArgumentType] # the same arguments Router binds -CoreBuilder = Callable[..., RouterCore] -"""``builder(host, **router_arguments)``: a router core on ``host``.""" - -CORE_BUILDERS: Final[dict[str, CoreBuilder]] = {"python": python_core} -"""Every core the suite runs on. Another core adds its builder here; a builder that cannot take a -configuration raises ``CoreDeclined``, and the test is skipped for that core.""" - - class CoreDeclined(Exception): """The core does not serve this configuration (the router would stay on the Python core).""" +CoreBuilder = Callable[..., RouterCore] +"""``builder(host, **router_arguments)``: a router core on ``host``.""" + + +def rust_core(host: ScriptedHost, **arguments: object) -> RouterCore: # kwargs-ok: the router's constructor arguments + """A ``RustCore`` on ``host``, built the way ``make_core`` builds one when the Rust router is on.""" + from litellm.router_rust.core import RustCore + from litellm.rust_bridge._native import RustBridgeDeclined + + config: Final = build_router_config(arguments) # an invalid config raises here, as for the Python core + try: + return RustCore.build(config, host) + except RustBridgeDeclined as declined: # make_core keeps such a router on the Python core + raise CoreDeclined(str(declined)) from declined + + +CORE_BUILDERS: Final[dict[str, CoreBuilder]] = {"python": python_core, "rust": rust_core} + + +def records_outcomes_through_callbacks(core: RouterCore) -> bool: + """Whether ``core`` learns attempt outcomes from callbacks it registers on its host. The Python + core does; the Rust core records every outcome from the attempt itself and registers none.""" + return isinstance(core, PythonCore) + + +"""Every core the suite runs on. Another core adds its builder here; a builder that cannot take a +configuration raises ``CoreDeclined``, and the test is skipped for that core.""" + + __all__ = [ - "Call", - "HostCallbacks", - "INFERENCE_OPS", "CORE_BUILDERS", + "INFERENCE_OPS", + "Call", "CoreBuilder", "CoreDeclined", + "HostCallbacks", "ScriptedHost", "auth_error", "bad_request", @@ -482,6 +539,7 @@ __all__ = [ "embedding_response", "python_core", "rate_limit", + "rust_core", "server_error", "timeout", ] diff --git a/tests/unit/test_router/core/selection_helpers.py b/tests/unit/test_router/core/selection_helpers.py new file mode 100644 index 00000000000..b38cfec1a19 --- /dev/null +++ b/tests/unit/test_router/core/selection_helpers.py @@ -0,0 +1,88 @@ +""" +Helpers for the selection characterization tests. + +``TrafficHost`` is a ``ScriptedHost`` whose async inference calls also feed the routing strategy +loggers the core registered on the host, the way litellm's logging does around a real call: input +callbacks see the call start (``log_pre_api_call``), and the ``CustomLogger`` instances in the +``callbacks`` list see it succeed or fail with real start and end times. A call can be held open +on a gate (an ``asyncio.Event``) or delayed for real, so in-flight counts and latencies are real +for any core. +""" + +import asyncio +import datetime +from collections.abc import Callable, Mapping +from typing import Final + +from litellm.integrations.custom_logger import CustomLogger +from litellm.router_core import RouterCore +from tests.unit.test_router.core.scripted_host import Call, ScriptedHost + +MakeCore = Callable[..., RouterCore] + +MESSAGES: Final = [{"role": "user", "content": "hi"}] + + +class TrafficHost(ScriptedHost): + """``ScriptedHost`` plus the strategy-logger traffic litellm's logging produces around a call.""" + + def __init__(self) -> None: + super().__init__() + self.gates: dict[str, asyncio.Event] = {} + self.delays: dict[str, float] = {} + self.started: list[str] = [] # mutable-ok: the models whose calls have started, in order + + def _loggers(self, name: str) -> tuple[CustomLogger, ...]: + return tuple(callback for callback in self.lists.get(name, ()) if isinstance(callback, CustomLogger)) + + async def _answer_async(self, op: str, args: tuple[object, ...], kwargs: Mapping[str, object]) -> object: + model: Final = str(kwargs.get("model") or (args[0] if args else "")) + logged: Final = self._logging_kwargs(Call(op=op, model=model, kwargs=dict(kwargs)), None, None) + for logger in self._loggers("input_callback"): + logger.log_pre_api_call(model=model, messages=kwargs.get("messages"), kwargs=logged) + self.started.append(model) + start: Final = datetime.datetime.now() + gate: Final = self.gates.get(model) + if gate is not None: + await gate.wait() + delay: Final = self.delays.get(model) + if delay: + await asyncio.sleep(delay) + try: + response = await super()._answer_async(op, args, kwargs) + except Exception as error: + failed = {**logged, "exception": error} + for logger in self._loggers("callbacks"): + await logger.async_log_failure_event(failed, None, start, datetime.datetime.now()) + raise + succeeded: Final = self._logging_kwargs(Call(op=op, model=model, kwargs=dict(kwargs)), None, response) + for logger in self._loggers("callbacks"): + await logger.async_log_success_event(succeeded, response, start, datetime.datetime.now()) + return response + + +async def picked_ids( + core: RouterCore, + model: str, + times: int = 20, + request_kwargs: Mapping[str, object] | None = None, + **select_arguments: object, # kwargs-ok: select's keyword options +) -> set[str]: + """The deployment ids ``select`` picks over ``times`` picks, each with a fresh copy of ``request_kwargs``.""" + picked: Final[set[str]] = set() + for _ in range(times): + chosen = await core.select(model, dict(request_kwargs or {}), **select_arguments) # pyright: ignore[reportCallIssue] # PythonCore's select takes keyword options + assert isinstance(chosen, Mapping) + picked.add(str(chosen["model_info"]["id"])) + return picked + + +def ids_of(deployments: object) -> list[str]: # mutable-ok: a fresh list + """The ids of a ``select(kind="healthy")`` result, in order (a single dict gives one id).""" + if isinstance(deployments, Mapping): + return [str(deployments["model_info"]["id"])] + assert isinstance(deployments, list) + return [str(deployment["model_info"]["id"]) for deployment in deployments] + + +__all__ = ["MESSAGES", "MakeCore", "TrafficHost", "ids_of", "picked_ids"] diff --git a/tests/unit/test_router/core/streaming_helpers.py b/tests/unit/test_router/core/streaming_helpers.py new file mode 100644 index 00000000000..3e436103e60 --- /dev/null +++ b/tests/unit/test_router/core/streaming_helpers.py @@ -0,0 +1,474 @@ +""" +Scripted streams for the streaming characterization tests. + +``ScriptedHost`` outcomes may be callables, so each builder here returns a callable that the host +calls with the request kwargs and that answers with a fresh stream: + +- ``chat_stream``: a ``CustomStreamWrapper`` over chat chunks, optionally raising ``error`` after + ``fail_after`` chunks, and ending with a ``finish_reason`` chunk unless ``finish`` is False. +- ``anthropic_stream``: an async iterator of raw ``/v1/messages`` SSE byte frames, optionally raising. +- ``responses_stream``: a ``BaseResponsesAPIStreamingIterator`` over Responses API events, + optionally raising. + +A stream built with ``logged=True`` logs through the ``ScriptedHost`` that served it when it ends +(``HostLogging``), as litellm's stream logging does, rather than when the call returns it. +""" + +import datetime +import json +import sys +import types +from collections.abc import AsyncIterator, Callable, Coroutine, Iterator, Mapping, Sequence +from typing import TYPE_CHECKING, Final + +from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging +from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper +from litellm.responses.streaming_iterator import BaseResponsesAPIStreamingIterator +from litellm.types.utils import ModelResponseStream + +if TYPE_CHECKING: + from tests.unit.test_router.core.scripted_host import Call, ScriptedHost + +MESSAGES: Final = [{"role": "user", "content": "hi"}] + + +def install_orjson_stub_if_missing() -> None: + """ + The /v1/messages stream wrapper imports proxy modules that import ``orjson`` and ``soundfile`` at load + time. Where either is not installed (some dev containers), a stand-in is registered so the wrapper can + load (json-backed for ``orjson``, empty for ``soundfile``, which only audio passthrough logging uses); + where they are installed (CI), nothing changes. + """ + try: + import orjson # noqa: F401 # pyright: ignore[reportUnusedImport] # only probing for the module + except ImportError: + stub: Final = types.ModuleType("orjson") + stub.loads = json.loads # pyright: ignore[reportAttributeAccessIssue] # module attribute + stub.dumps = lambda value, **_options: json.dumps(value).encode() # pyright: ignore[reportAttributeAccessIssue] # module attribute + stub.JSONDecodeError = json.JSONDecodeError # pyright: ignore[reportAttributeAccessIssue] # module attribute + sys.modules["orjson"] = stub + try: + import soundfile # noqa: F401 # pyright: ignore[reportUnusedImport, reportMissingImports] # only probing for the module + except (ImportError, OSError): + sys.modules["soundfile"] = types.ModuleType("soundfile") + + +def _logging_obj(model: str, call_type: str = "acompletion", logged: bool = False) -> LiteLLMLogging: + logging_obj: Final = (HostLogging if logged else LiteLLMLogging)( + model=model, + messages=MESSAGES, + stream=True, + call_type=call_type, + start_time=datetime.datetime(2026, 1, 1, tzinfo=datetime.timezone.utc), + litellm_call_id=f"stream-{model}", + function_id=f"stream-{model}", + ) + logging_obj.update_environment_variables(model=model, user=None, optional_params={}, litellm_params={}) + return logging_obj + + +def request_logging(model: str = "a") -> LiteLLMLogging: + """The logging object a request shares across its attempts (``litellm_logging_obj``).""" + return _logging_obj(model) + + +class HostLogging(LiteLLMLogging): + """A stream's logging object that runs the serving ``ScriptedHost``'s callbacks where litellm's + stream logging runs litellm's: once with the assembled response when the stream ends, and once + per failure the stream raises (``dispatch_attempt_failure_handlers``). The host binds itself and + the call when it returns the stream (``ScriptedHost._logs_at_stream_end``).""" + + def bind_host(self, host: "ScriptedHost", call: "Call") -> None: + self._scripted: tuple[ScriptedHost, Call] | None = (host, call) # pyright: ignore[reportUninitializedInstanceVariable] # set when the host binds + self._logged_success = False + + def _bound(self) -> "tuple[ScriptedHost, Call] | None": + return getattr(self, "_scripted", None) + + async def dispatch_success_handlers( # pyright: ignore[reportIncompatibleMethodOverride] # the arguments the stream passes + self, result: object = None, *_args: object, **_kwargs: object + ) -> None: + bound: Final = self._bound() + if bound is None or self._logged_success: + return + self._logged_success = True + await bound[0].log_async(bound[1], None, result) + + def dispatch_attempt_failure_handlers( + self, exception: Exception, traceback_exception: str + ) -> Coroutine[None, None, None]: + bound: Final = self._bound() + + async def dispatch() -> None: + if bound is not None: + await bound[0].log_async(bound[1], exception, None) + + return dispatch() + + def success_handler(self, result: object = None, *_args: object, **_kwargs: object) -> None: # pyright: ignore[reportIncompatibleMethodOverride] # see dispatch_success_handlers + bound: Final = self._bound() + if bound is None or self._logged_success or not getattr(result, "usage", None): + return + self._logged_success = True + bound[0].log_sync(bound[1], None, result) + + def failure_handler(self, exception: Exception, *_args: object, **_kwargs: object) -> None: # pyright: ignore[reportIncompatibleMethodOverride] # see dispatch_success_handlers + bound: Final = self._bound() + if bound is not None: + bound[0].log_sync(bound[1], exception, None) + + +class _EndLogged: + """What a hand-written scripted stream logs once when it ends: its failure, or a response carrying + ``total_tokens`` (the usage the provider's own stream logging would assemble).""" + + def __init__(self, model: str, call_type: str, logged: bool, total_tokens: int) -> None: + self.logging_obj: Final[LiteLLMLogging | None] = _logging_obj(model, call_type, True) if logged else None + self._total_tokens: Final = total_tokens + self._logged = False + + async def log_end(self, error: BaseException | None) -> None: + if not isinstance(self.logging_obj, HostLogging) or self._logged: + return + self._logged = True + if error is not None: + # The provider's own error, as a stream logs it before wrapping it for the router. + original: Final = getattr(error, "original_exception", None) + logged: Final = original if isinstance(original, Exception) else error + await self.logging_obj.dispatch_attempt_failure_handlers( + logged if isinstance(logged, Exception) else Exception(logged), "" + ) + return + usage: Final = types.SimpleNamespace(total_tokens=self._total_tokens) + await self.logging_obj.dispatch_success_handlers(types.SimpleNamespace(usage=usage)) + + +def text_chunk(text: str) -> ModelResponseStream: + return ModelResponseStream(choices=[{"index": 0, "delta": {"role": "assistant", "content": text}}]) + + +def role_chunk() -> ModelResponseStream: + return ModelResponseStream(choices=[{"index": 0, "delta": {"role": "assistant", "content": ""}}]) + + +def finish_chunk(reason: str = "stop", usage: Mapping[str, int] | None = None) -> ModelResponseStream: + chunk: Final = ModelResponseStream(choices=[{"index": 0, "delta": {}, "finish_reason": reason}]) + if usage is not None: + setattr(chunk, "usage", dict(usage)) + return chunk + + +def chat_stream( + *texts: str, + error: BaseException | None = None, + fail_after: int | None = None, + finish: bool = True, + leading_role_chunk: bool = False, + model: str = "d1", + hidden_params: Mapping[str, object] | None = None, + sync: bool = False, + usage: Mapping[str, int] | None = None, + logged: bool = False, +) -> Callable[..., CustomStreamWrapper]: + """ + A stream of ``texts`` (one chunk each). With ``error``, raises it after ``fail_after`` chunks + (default: after every text chunk, before the finish chunk). ``leading_role_chunk`` sends a + role-only, empty-content chunk first, which counts as the first chunk but carries no content. + ``usage`` rides on the finish chunk; ``logged`` logs through the host when the stream ends. + """ + + def answer(**request: object) -> CustomStreamWrapper: # kwargs-ok: the host passes the request kwargs + produced: list[ModelResponseStream] = [] + if leading_role_chunk: + produced.append(role_chunk()) + produced.extend(text_chunk(text) for text in texts) + cut_at: Final = len(produced) if fail_after is None else fail_after + + def items() -> Iterator[ModelResponseStream]: + for index, chunk in enumerate(produced): + if error is not None and index == cut_at: + raise error + yield chunk + if error is not None: + raise error + if finish: + yield finish_chunk(usage=usage) + + async def achunks() -> AsyncIterator[ModelResponseStream]: + for chunk in items(): + yield chunk + + wrapper: Final = CustomStreamWrapper( + completion_stream=items() if sync else achunks(), + model=model, + custom_llm_provider="openai", + logging_obj=_logging_obj(model, "completion" if sync else "acompletion", logged), + ) + if hidden_params is not None: + wrapper._hidden_params.update(hidden_params) # pyright: ignore[reportPrivateUsage] # test seam + return wrapper + + return answer + + +async def collect_texts(stream: object) -> str: + """The text content of every chunk of an async chat stream, joined.""" + texts: Final[list[str]] = [] + async for chunk in stream: # pyright: ignore[reportGeneralTypeIssues] # an async stream + choices = getattr(chunk, "choices", None) + if choices: + content = choices[0].delta.content + if content: + texts.append(content) + return "".join(texts) + + +def collect_texts_sync(stream: object) -> str: + texts: Final[list[str]] = [] + for chunk in stream: # pyright: ignore[reportGeneralTypeIssues] # a sync stream + choices = getattr(chunk, "choices", None) + if choices: + content = choices[0].delta.content + if content: + texts.append(content) + return "".join(texts) + + +# --- /v1/messages ------------------------------------------------------------------------------ + + +def sse(event: str, data: Mapping[str, object]) -> bytes: + return f"event: {event}\ndata: {json.dumps(data)}\n\n".encode() + + +MESSAGE_START: Final = sse( + "message_start", + {"type": "message_start", "message": {"id": "msg_1", "type": "message", "role": "assistant", "content": []}}, +) +BLOCK_START: Final = sse( + "content_block_start", {"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}} +) +PING: Final = sse("ping", {"type": "ping"}) +MESSAGE_STOP: Final = sse("message_stop", {"type": "message_stop"}) + + +def text_delta(text: str) -> bytes: + return sse( + "content_block_delta", + {"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": text}}, + ) + + +def error_frame(error_type: str = "overloaded_error", message: str = "Overloaded") -> bytes: + return sse("error", {"type": "error", "error": {"type": error_type, "message": message}}) + + +def refusal_delta(stop_details: Mapping[str, object] | None = None) -> bytes: + delta: Final[dict[str, object]] = {"stop_reason": "refusal"} + if stop_details is not None: + delta["stop_details"] = dict(stop_details) + return sse("message_delta", {"type": "message_delta", "delta": delta, "usage": {"output_tokens": 1}}) + + +class ByteStream: + """An async iterator of SSE frames that records whether it was closed. ``logged`` logs through the host + when it ends, with ``total_tokens`` (``_EndLogged``).""" + + def __init__( + self, + frames: Sequence[bytes], + error: BaseException | None, + hidden_params: Mapping[str, object] | None, + logged: bool = False, + total_tokens: int = 0, + ) -> None: + self._frames: Final = list(frames) + self._error: Final = error + self._index = 0 + self.closed = False + self._end: Final = _EndLogged("d1", "anthropic_messages", logged, total_tokens) + if logged: + self.logging_obj = self._end.logging_obj + if hidden_params is not None: + self._hidden_params = dict(hidden_params) + + def __aiter__(self) -> "ByteStream": + return self + + async def __anext__(self) -> bytes: + if self._index < len(self._frames): + self._index += 1 + return self._frames[self._index - 1] + await self._end.log_end(self._error) + if self._error is not None: + raise self._error + raise StopAsyncIteration + + async def aclose(self) -> None: + self.closed = True + + +def anthropic_stream( + *frames: bytes, + error: BaseException | None = None, + hidden_params: Mapping[str, object] | None = None, + opened: list[ByteStream] | None = None, + logged: bool = False, + total_tokens: int = 0, +) -> Callable[..., ByteStream]: + def answer(**_request: object) -> ByteStream: # kwargs-ok: the host passes the request kwargs + stream: Final = ByteStream(frames, error, hidden_params, logged, total_tokens) + if opened is not None: + opened.append(stream) + return stream + + return answer + + +async def collect_bytes(stream: object) -> list[bytes]: + return [frame async for frame in stream] # pyright: ignore[reportGeneralTypeIssues] # an async stream + + +# --- Responses API ----------------------------------------------------------------------------- + + +class Event: + """A Responses API stream event: only ``type`` (and optionally ``response``) is read.""" + + def __init__(self, event_type: str, response: object = None, delta: str | None = None) -> None: + self.type = event_type + self.response = response + self.delta = delta + self._hidden_params: dict[str, object] = {} + + def __repr__(self) -> str: + return f"Event({self.type!r})" + + +class BackgroundResponse: + background = True + id = "resp_bg" + + +class ScriptedResponsesStream(BaseResponsesAPIStreamingIterator): + """A Responses stream over ``events`` that bypasses the provider plumbing of the base class.""" + + def __init__( # pyright: ignore[reportMissingSuperCall] # the base needs an httpx response + self, + events: Sequence[object], + error: BaseException | None, + model: str, + custom_llm_provider: str, + hidden_params: Mapping[str, object] | None, + started_tool_call: bool = False, + completed_response: object = None, + logged: bool = False, + total_tokens: int = 0, + ) -> None: + self._end: Final = _EndLogged(model, "aresponses", logged, total_tokens) + self._events: Final = list(events) + self._started_tool_call = started_tool_call + self._error: Final = error + self._index = 0 + self._hidden_params = dict(hidden_params or {}) + self.model = model + self.custom_llm_provider = custom_llm_provider + self.logging_obj = self._end.logging_obj or _logging_obj(model, "aresponses") + self.litellm_metadata = None + self.responses_api_provider_config = None + self.finished = False + self.completed_response = completed_response + self.response = None + self.start_time = datetime.datetime(2026, 1, 1, tzinfo=datetime.timezone.utc) + self.request_data = {} + self.call_type = None + self.closed = False + + def __aiter__(self) -> "ScriptedResponsesStream": + return self + + async def __anext__(self) -> object: + if self._index < len(self._events): + self._index += 1 + return self._events[self._index - 1] + await self._end.log_end(self._error) + if self._error is not None: + raise self._error + raise StopAsyncIteration + + async def aclose(self) -> None: + self.closed = True + + +def responses_stream( + *events: object, + error: BaseException | None = None, + model: str = "d1", + custom_llm_provider: str = "anthropic", + hidden_params: Mapping[str, object] | None = None, + opened: list[ScriptedResponsesStream] | None = None, + started_tool_call: bool = False, + completed_response: object = None, + logged: bool = False, + total_tokens: int = 0, +) -> Callable[..., ScriptedResponsesStream]: + """``completed_response`` is the terminal event the stream latched before failing, if any. ``logged`` + logs through the host when the stream ends, with ``total_tokens``.""" + + def answer(**_request: object) -> ScriptedResponsesStream: # kwargs-ok: the host passes the request kwargs + stream: Final = ScriptedResponsesStream( + events, + error, + model, + custom_llm_provider, + hidden_params, + started_tool_call, + completed_response, + logged, + total_tokens, + ) + if opened is not None: + opened.append(stream) + return stream + + return answer + + +async def drain_into(stream: object, seen: list[object]) -> None: + """Append every item of an async stream to ``seen``, letting an error the stream raises propagate.""" + async for item in stream: # pyright: ignore[reportGeneralTypeIssues] # an async stream + seen.append(item) + + +async def collect_event_types(stream: object) -> list[str]: + return [getattr(event, "type", repr(event)) async for event in stream] # pyright: ignore[reportGeneralTypeIssues] # an async stream + + +__all__ = [ + "BLOCK_START", + "MESSAGES", + "MESSAGE_START", + "MESSAGE_STOP", + "PING", + "BackgroundResponse", + "ByteStream", + "Event", + "HostLogging", + "ScriptedResponsesStream", + "anthropic_stream", + "chat_stream", + "collect_bytes", + "collect_event_types", + "collect_texts", + "collect_texts_sync", + "drain_into", + "error_frame", + "finish_chunk", + "install_orjson_stub_if_missing", + "refusal_delta", + "responses_stream", + "role_chunk", + "sse", + "text_chunk", + "text_delta", +] diff --git a/tests/unit/test_router/core/test_construction_precall_characterization.py b/tests/unit/test_router/core/test_construction_precall_characterization.py new file mode 100644 index 00000000000..a56cd78ccac --- /dev/null +++ b/tests/unit/test_router/core/test_construction_precall_characterization.py @@ -0,0 +1,1302 @@ +""" +Characterization of a router core's construction and lifecycle: what each constructor argument sets +(and the process settings it falls back to), the validation errors, the caches and callbacks it +creates, alerting, optional pre-call check registration, ``observe``, ``register_plugin``, +``retire``, ``flush_cache``/``reset``, attribute writes, ``get_valid_args`` and ``print_deployment``. + +Each test asserts what the Python core does today; oddities are named in the test's docstring. +""" + +import logging +from collections.abc import Callable, Generator, Mapping +from typing import Final + +import pytest + +import litellm +from litellm._logging import verbose_router_logger +from litellm.caching.dual_cache import DualCache +from litellm.caching.redis_cache import RedisCache +from litellm.router_api import ChangeSet, router_init_args +from litellm.router_core import RouterCore +from litellm.router_python_core import PythonCore +from litellm.router_strategy.budget_limiter import RouterBudgetLimiting +from litellm.router_utils.pre_call_checks.prompt_caching_deployment_check import PromptCachingDeploymentCheck +from litellm.types.router import ( + AlertingConfig, + AllowedFailsPolicy, + CustomRoutingStrategyBase, + RetryPolicy, + RoutingStrategy, +) +from tests.unit.test_router.core.construction_precall_helpers import ( + CallbackHost, + collect_other_routers, + host_callback_names, +) +from tests.unit.test_router.core.scripted_host import ScriptedHost, deployment, records_outcomes_through_callbacks + +MakeCore = Callable[..., RouterCore] +MESSAGES: Final = [{"role": "user", "content": "hi"}] + + +@pytest.fixture +def host() -> ScriptedHost: + return CallbackHost() + + +def _one(make_core: MakeCore, **arguments: object) -> RouterCore: # kwargs-ok: constructor arguments + return make_core(model_list=[deployment("g", "a")], **arguments) + + +# --- defaults -------------------------------------------------------------------------------- + +SETTING_DEFAULTS: Final = { + "routing_strategy": "simple-shuffle", + "fallbacks": None, + "context_window_fallbacks": None, + "content_policy_fallbacks": None, + "default_deployment": None, + "default_max_parallel_requests": None, + "cooldown_time": 5, + "allowed_fails": 3, + "allowed_fails_policy": None, + "disable_cooldowns": None, + "enable_pre_call_checks": False, + "enable_health_check_routing": False, + "health_check_ignore_transient_errors": False, + "background_health_check_model_groups": None, + "enable_tag_filtering": False, + "tag_filtering_match_any": True, + "tag_routing_prefix": "", + "provider_budget_config": None, + "provider_default_deployment_ids": [], + "fallback_access_check": None, + "fallback_budget_check": None, + "search_tools": [], + "guardrail_list": [], + "cache_responses": False, +} + + +@pytest.mark.parametrize(("field", "expected"), sorted(SETTING_DEFAULTS.items())) +def test_settings_default_of_a_router_built_with_only_a_model_list( + make_core: MakeCore, field: str, expected: object +) -> None: + """Note: ``fallbacks`` defaults to ``[]`` in the signature, yet reads ``None``: an empty list falls + through ``fallbacks or ``.""" + assert getattr(_one(make_core).settings(), field) == expected + + +def test_default_litellm_params_default_to_no_timeout_no_client_retries_and_no_caching_groups( + make_core: MakeCore, +) -> None: + assert _one(make_core).settings().default_litellm_params == { + "timeout": None, + "max_retries": 0, + "metadata": {"caching_groups": None}, + } + + +def test_router_general_settings_default_to_async_only_and_pass_through_off(make_core: MakeCore) -> None: + general: Final = _one(make_core).settings().router_general_settings + assert (general.async_only_mode, general.pass_through_all_models) == (False, False) + + +def test_settings_as_dict_of_a_default_router(make_core: MakeCore) -> None: + assert dict(_one(make_core).settings().as_dict()) == { + "routing_strategy_args": {}, + "routing_strategy": "simple-shuffle", + "allowed_fails": 3, + "cooldown_time": 5, + "num_retries": 2, + "timeout": 6000.0, + "retry_after": 0, + "fallbacks": None, + "context_window_fallbacks": None, + "model_group_retry_policy": {}, + "retry_policy": None, + "model_group_alias": {}, + "enable_weighted_failover": False, + "enable_mid_stream_fallback_continuation": False, + "enable_tag_filtering": False, + "tag_routing_prefix": "", + "fallback_on_unsupported_params": True, + "routing_groups": [], + } + + +ATTRIBUTE_DEFAULTS: Final = { + "num_retries": 2, + "max_fallbacks": 5, + "timeout": 6000.0, + "stream_timeout": None, + "request_timeout": None, + "retry_after": 0, + "client_ttl": 3600, + "set_verbose": False, + "debug_level": "INFO", + "retry_policy": None, + "model_group_retry_policy": {}, + "alerting_config": None, + "routing_strategy_args": {}, + "deployment_affinity_ttl_seconds": 3600, + "default_priority": None, + "optional_callbacks": None, + "router_budget_logger": None, + "fallback_on_unsupported_params": True, + "enable_weighted_failover": False, + "enable_mid_stream_fallback_continuation": False, + "explicit_allowed_fails": None, + "ignore_invalid_deployments": False, + "auto_router_capability_limit": None, + "assistants_config": None, + "session_affinity_repin_on_failover": False, + "model_group_affinity_config": None, + "access_groups": None, + "leastbusy_logger": None, + "lowesttpm_logger": None, +} + + +@pytest.mark.parametrize(("name", "expected"), sorted(ATTRIBUTE_DEFAULTS.items())) +def test_attribute_default_of_a_router_built_with_only_a_model_list( + make_core: MakeCore, name: str, expected: object +) -> None: + assert getattr(_one(make_core), name) == expected + + +def test_a_router_without_a_model_list_has_no_deployments(make_core: MakeCore) -> None: + core: Final = make_core() + assert list(core.registry().all()) == [] + + +def test_construction_turns_off_the_host_debug_info(make_core: MakeCore, host: ScriptedHost) -> None: + _one(make_core) + assert host.suppressed is True + + +def test_model_group_alias_none_reads_as_an_empty_alias_map(make_core: MakeCore) -> None: + core: Final = _one(make_core, model_group_alias=None) + assert dict(core.settings().as_dict()["model_group_alias"]) == {} # pyright: ignore[reportArgumentType] + + +# --- process settings each argument falls back to ----------------------------------------- + + +@pytest.mark.parametrize(("argument", "host_value", "expected"), [(5, 4, 5), (0, 4, 0), (None, 4, 4), (None, None, 2)]) +def test_num_retries_comes_from_the_argument_then_the_process_then_the_openai_default( + make_core: MakeCore, host: ScriptedHost, argument: int | None, host_value: int | None, expected: int +) -> None: + host.set_dynamic(num_retries=host_value) + core: Final = _one(make_core, **({"num_retries": argument} if argument is not None else {})) + assert core.settings().as_dict()["num_retries"] == expected + + +@pytest.mark.asyncio +async def test_num_retries_from_the_process_settings_governs_the_attempts( + make_core: MakeCore, host: ScriptedHost +) -> None: + host.set_dynamic(num_retries=1) + host.answer("openai/a", litellm.InternalServerError(message="boom", llm_provider="openai", model="a")) + core: Final = _one(make_core) + with pytest.raises(litellm.InternalServerError): + await core.acompletion(model="g", messages=MESSAGES) # pyright: ignore[reportAttributeAccessIssue] + assert len(host.calls) == 2 + + +@pytest.mark.parametrize( + ("argument", "host_max", "host_router_max", "expected"), + [(3, 9, 11, 3), (0, 9, 11, 0), (None, 9, 11, 9), (None, None, 11, 11)], +) +def test_max_fallbacks_comes_from_the_argument_then_litellm_max_fallbacks_then_router_max_fallbacks( + make_core: MakeCore, + host: ScriptedHost, + argument: int | None, + host_max: int | None, + host_router_max: int, + expected: int, +) -> None: + host.set_dynamic(max_fallbacks=host_max, router_max_fallbacks=host_router_max) + core: Final = _one(make_core, **({"max_fallbacks": argument} if argument is not None else {})) + assert core.max_fallbacks == expected # pyright: ignore[reportAttributeAccessIssue] + + +@pytest.mark.parametrize(("argument", "expected", "explicit"), [(7, 7, 7), (0, 0, 0), (None, 8, None)]) +def test_allowed_fails_comes_from_the_argument_then_the_process_and_only_an_argument_counts_as_explicit( + make_core: MakeCore, host: ScriptedHost, argument: int | None, expected: int, explicit: int | None +) -> None: + host.set_dynamic(allowed_fails=8) + core: Final = _one(make_core, **({"allowed_fails": argument} if argument is not None else {})) + assert core.settings().allowed_fails == expected + assert core.explicit_allowed_fails == explicit # pyright: ignore[reportAttributeAccessIssue] + + +@pytest.mark.parametrize(("argument", "expected"), [(12.5, 12.5), (None, 33.0), (0, 33.0)]) +def test_timeout_comes_from_the_argument_else_the_process_request_timeout( + make_core: MakeCore, host: ScriptedHost, argument: float | None, expected: float +) -> None: + """Oddity: ``timeout=0`` is falsy and reads as the process ``request_timeout``.""" + host.set_dynamic(request_timeout=33.0) + core: Final = _one(make_core, **({"timeout": argument} if argument is not None else {})) + assert core.settings().as_dict()["timeout"] == expected + + +def test_timeout_zero_still_reaches_default_litellm_params_as_zero(make_core: MakeCore) -> None: + core: Final = _one(make_core, timeout=0) + assert core.settings().default_litellm_params["timeout"] == 0 + + +@pytest.mark.parametrize(("argument", "expected"), [(30, 30), (0.5, 0.5), (0, 5), (None, 5)]) +def test_cooldown_time_defaults_to_five_seconds_and_zero_reads_as_the_default( + make_core: MakeCore, argument: float | None, expected: float +) -> None: + """Oddity: ``cooldown_time=0`` is replaced by the 5 second default (``cooldown_time or DEFAULT``).""" + core: Final = _one(make_core, **({"cooldown_time": argument} if argument is not None else {})) + assert core.settings().cooldown_time == expected + + +def test_request_timeout_is_unset_without_a_router_timeout( + make_core: MakeCore, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(litellm, "request_timeout", 77.0) + assert _one(make_core).request_timeout is None # pyright: ignore[reportAttributeAccessIssue] + + +def test_request_timeout_reads_litellm_request_timeout_directly_when_a_router_timeout_is_set( + make_core: MakeCore, host: ScriptedHost, monkeypatch: pytest.MonkeyPatch +) -> None: + """Oddity: this one setting is read from ``litellm.request_timeout`` itself, not from the host's + process settings (which say 33 here).""" + host.set_dynamic(request_timeout=33.0) + monkeypatch.setattr(litellm, "request_timeout", 77.0) + assert _one(make_core, timeout=10).request_timeout == 77.0 # pyright: ignore[reportAttributeAccessIssue] + + +def test_request_timeout_is_unset_when_litellm_request_timeout_is_its_untouched_default( + make_core: MakeCore, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(litellm, "request_timeout", 6000) + monkeypatch.setattr(litellm, "request_timeout_explicitly_set", False, raising=False) + assert _one(make_core, timeout=10).request_timeout is None # pyright: ignore[reportAttributeAccessIssue] + + +def test_stream_timeout_is_kept_as_given(make_core: MakeCore) -> None: + assert _one(make_core, stream_timeout=3).stream_timeout == 3 # pyright: ignore[reportAttributeAccessIssue] + + +@pytest.mark.parametrize("field", ["fallbacks", "context_window_fallbacks", "content_policy_fallbacks"]) +def test_fallback_lists_come_from_the_process_when_the_argument_is_empty( + make_core: MakeCore, host: ScriptedHost, field: str +) -> None: + host.set_dynamic(**{field: [{"g": ["process"]}]}) + assert getattr(_one(make_core).settings(), field) == [{"g": ["process"]}] + + +@pytest.mark.parametrize("field", ["fallbacks", "context_window_fallbacks", "content_policy_fallbacks"]) +def test_fallback_list_arguments_win_over_the_process_lists( + make_core: MakeCore, host: ScriptedHost, field: str +) -> None: + host.set_dynamic(**{field: [{"g": ["process"]}]}) + assert getattr(_one(make_core, **{field: [{"g": ["argument"]}]}).settings(), field) == [{"g": ["argument"]}] + + +def test_default_fallbacks_are_appended_to_the_fallbacks_as_a_star_entry(make_core: MakeCore) -> None: + core: Final = _one(make_core, fallbacks=[{"g": ["b"]}], default_fallbacks=["c"]) + assert core.settings().fallbacks == [{"g": ["b"]}, {"*": ["c"]}] + + +def test_default_fallbacks_alone_become_the_only_star_entry(make_core: MakeCore) -> None: + assert _one(make_core, default_fallbacks=["c"]).settings().fallbacks == [{"*": ["c"]}] + + +def test_default_fallbacks_come_from_the_process_when_not_given(make_core: MakeCore, host: ScriptedHost) -> None: + host.set_dynamic(default_fallbacks=["d"]) + assert _one(make_core).settings().fallbacks == [{"*": ["d"]}] + + +def test_default_fallbacks_given_as_a_string_raise_type_error(make_core: MakeCore) -> None: + with pytest.raises(TypeError, match=r"^default_fallbacks must be a list of model names, got 'b'\. Use \['b'\]$"): + _one(make_core, default_fallbacks="b") + + +def test_default_fallbacks_are_appended_to_the_process_fallbacks_list_itself( + make_core: MakeCore, host: ScriptedHost +) -> None: + """Oddity: with no ``fallbacks`` argument the router adopts the process list object and appends the + star entry to it, so the process setting is changed too.""" + process_fallbacks: Final[list[object]] = [{"g": ["e"]}] + host.set_dynamic(fallbacks=process_fallbacks) + core: Final = _one(make_core, default_fallbacks=["c"]) + assert core.settings().fallbacks == [{"g": ["e"]}, {"*": ["c"]}] + assert process_fallbacks == [{"g": ["e"]}, {"*": ["c"]}] + + +# --- validation -------------------------------------------------------------------------------- + + +@pytest.mark.parametrize("field", ["fallbacks", "content_policy_fallbacks"]) +def test_a_fallback_entry_that_is_not_a_dict_raises_value_error(make_core: MakeCore, field: str) -> None: + with pytest.raises(ValueError, match=r"^Item 'x' is not a dictionary\.$"): + _one(make_core, **{field: ["x"]}) + + +@pytest.mark.parametrize("field", ["fallbacks", "content_policy_fallbacks"]) +def test_a_fallback_entry_with_two_keys_raises_value_error(make_core: MakeCore, field: str) -> None: + with pytest.raises(ValueError, match=r"must have exactly one key, but has 2 keys\.$"): + _one(make_core, **{field: [{"a": ["b"], "c": ["d"]}]}) + + +@pytest.mark.parametrize("field", ["fallbacks", "content_policy_fallbacks"]) +def test_a_fallback_target_given_as_a_string_raises_type_error(make_core: MakeCore, field: str) -> None: + with pytest.raises(TypeError) as error: + _one(make_core, **{field: [{"a": "b"}]}) + assert str(error.value) == "Fallbacks for 'a' must be a list of model names, got 'b'. Use {'a': ['b']}" + + +@pytest.mark.parametrize("targets", [[], [{"model": "b"}], ["b", "c"]]) +def test_fallback_targets_as_a_list_of_names_or_dicts_are_accepted(make_core: MakeCore, targets: list[object]) -> None: + assert _one(make_core, fallbacks=[{"g": targets}]).settings().fallbacks == [{"g": targets}] + + +@pytest.mark.parametrize("value", [["x"], [{"a": "b"}]]) +def test_context_window_fallbacks_are_not_validated(make_core: MakeCore, value: list[object]) -> None: + """Oddity: unlike ``fallbacks`` and ``content_policy_fallbacks``, malformed context window + fallbacks are accepted as given.""" + assert _one(make_core, context_window_fallbacks=value).settings().context_window_fallbacks == value + + +def test_malformed_process_fallbacks_raise_at_construction(make_core: MakeCore, host: ScriptedHost) -> None: + host.set_dynamic(fallbacks=[{"a": "e"}]) + with pytest.raises(TypeError, match=r"^Fallbacks for 'a' must be a list of model names"): + _one(make_core) + + +def test_validate_fallbacks_accepts_none(make_core: MakeCore) -> None: + assert _one(make_core).validate_fallbacks(fallback_param=None) is None # pyright: ignore[reportAttributeAccessIssue] + + +def test_an_unknown_routing_strategy_raises_value_error(make_core: MakeCore) -> None: + with pytest.raises(ValueError, match=r"^Invalid routing_strategy: 'bogus'\.") as error: + _one(make_core, routing_strategy="bogus") + assert str(error.value).startswith("Invalid routing_strategy: 'bogus'. Valid options: [") + assert str(error.value).endswith( + "Check 'router_settings.routing_strategy' in your config.yaml " + "or the 'routing_strategy' parameter if using the Router SDK directly." + ) + + +def test_a_routing_strategy_enum_is_stored_as_its_string_value(make_core: MakeCore) -> None: + assert _one(make_core, routing_strategy=RoutingStrategy.LEAST_BUSY).settings().routing_strategy == "least-busy" + + +@pytest.mark.parametrize( + "strategy", + [ + "simple-shuffle", + "least-busy", + "usage-based-routing", + "usage-based-routing-v2", + "latency-based-routing", + "cost-based-routing", + "lar1", + ], +) +def test_each_routing_strategy_is_reported_in_settings(make_core: MakeCore, strategy: str) -> None: + assert _one(make_core, routing_strategy=strategy).settings().routing_strategy == strategy + + +# --- default_litellm_params ---------------------------------------------------------------- + + +def test_default_litellm_params_keep_given_values_and_gain_max_retries_and_caching_groups(make_core: MakeCore) -> None: + core: Final = _one( + make_core, + timeout=12.5, + default_litellm_params={"timeout": 99, "temperature": 0.1, "metadata": {"x": 1}}, + caching_groups=[("a", "b")], + ) + assert core.settings().default_litellm_params == { + "timeout": 99, + "temperature": 0.1, + "metadata": {"x": 1, "caching_groups": [("a", "b")]}, + "max_retries": 0, + } + + +def test_default_litellm_params_timeout_defaults_to_the_router_timeout_argument(make_core: MakeCore) -> None: + assert _one(make_core, timeout=12.5).settings().default_litellm_params["timeout"] == 12.5 + + +def test_default_litellm_params_argument_is_not_changed_by_construction(make_core: MakeCore) -> None: + given: Final = {"temperature": 0.1} + _one(make_core, default_litellm_params=given) + assert given == {"temperature": 0.1} + + +# --- policies, priority, scheduler ------------------------------------------------------- + + +def test_retry_policy_given_as_a_dict_becomes_a_retry_policy(make_core: MakeCore) -> None: + policy: Final = _one(make_core, retry_policy={"RateLimitErrorRetries": 3}).retry_policy # pyright: ignore[reportAttributeAccessIssue] + assert isinstance(policy, RetryPolicy) + assert policy.RateLimitErrorRetries == 3 + + +def test_retry_policy_given_as_a_retry_policy_is_kept(make_core: MakeCore) -> None: + policy: Final = RetryPolicy(TimeoutErrorRetries=1) + assert _one(make_core, retry_policy=policy).retry_policy == policy # pyright: ignore[reportAttributeAccessIssue] + + +def test_retry_policy_of_another_type_is_ignored(make_core: MakeCore) -> None: + assert _one(make_core, retry_policy="nope").retry_policy is None # pyright: ignore[reportAttributeAccessIssue] + + +def test_allowed_fails_policy_given_as_a_dict_becomes_an_allowed_fails_policy(make_core: MakeCore) -> None: + policy: Final = ( + _one(make_core, allowed_fails_policy={"RateLimitErrorAllowedFails": 7}).settings().allowed_fails_policy + ) + assert isinstance(policy, AllowedFailsPolicy) + assert policy.RateLimitErrorAllowedFails == 7 + + +def test_allowed_fails_policy_of_another_type_is_ignored(make_core: MakeCore) -> None: + assert _one(make_core, allowed_fails_policy="nope").settings().allowed_fails_policy is None + + +def test_model_group_retry_policy_is_kept_as_given(make_core: MakeCore) -> None: + policy: Final = {"g": RetryPolicy(TimeoutErrorRetries=4)} + assert _one(make_core, model_group_retry_policy=policy).model_group_retry_policy == policy # pyright: ignore[reportAttributeAccessIssue] + + +@pytest.mark.parametrize( + ("given", "expected"), [(1, 1), ("2", 2), (-1, -1), (None, None), ("x", None), (1.5, None), (True, None)] +) +def test_default_priority_keeps_an_integer_and_ignores_anything_else( + make_core: MakeCore, given: object, expected: int | None +) -> None: + assert _one(make_core, default_priority=given).default_priority == expected # pyright: ignore[reportAttributeAccessIssue] + + +@pytest.mark.parametrize(("given", "expected"), [(None, 0.03), (0.5, 0.5)]) +def test_scheduler_polling_interval_defaults_to_three_hundredths_of_a_second( + make_core: MakeCore, given: float | None, expected: float +) -> None: + core: Final = _one(make_core, **({"polling_interval": given} if given is not None else {})) + assert core.scheduler.polling_interval == expected # pyright: ignore[reportAttributeAccessIssue] + + +def test_client_ttl_is_kept_as_given(make_core: MakeCore) -> None: + assert _one(make_core, client_ttl=10).client_ttl == 10 # pyright: ignore[reportAttributeAccessIssue] + + +@pytest.mark.parametrize( + ("name", "value"), + [ + ("routing_strategy_args", {"ttl": 5}), + ("deployment_affinity_ttl_seconds", 60), + ("session_affinity_repin_on_failover", True), + ("enable_weighted_failover", True), + ("fallback_on_unsupported_params", False), + ("ignore_invalid_deployments", True), + ("retry_after", 4), + ], +) +def test_plain_arguments_are_kept_as_attributes(make_core: MakeCore, name: str, value: object) -> None: + assert getattr(_one(make_core, **{name: value}), name) == value + + +@pytest.mark.parametrize( + ("name", "value"), + [ + ("enable_tag_filtering", True), + ("tag_filtering_match_any", False), + ("tag_routing_prefix", "team-"), + ("enable_pre_call_checks", True), + ("disable_cooldowns", True), + ("default_max_parallel_requests", 4), + ("enable_health_check_routing", True), + ("health_check_ignore_transient_errors", True), + ], +) +def test_setting_arguments_are_reported_in_settings(make_core: MakeCore, name: str, value: object) -> None: + assert getattr(_one(make_core, **{name: value}).settings(), name) == value + + +def test_background_health_check_model_groups_become_a_frozenset(make_core: MakeCore) -> None: + core: Final = _one(make_core, background_health_check_model_groups=["g", "g", "h"]) + assert core.settings().background_health_check_model_groups == frozenset({"g", "h"}) + + +@pytest.mark.parametrize(("given", "expected"), [(True, True), ("true", True), ("false", False), (False, False)]) +def test_mid_stream_fallback_continuation_accepts_a_bool_or_a_bool_string( + make_core: MakeCore, given: object, expected: bool +) -> None: + assert ( + _one(make_core, enable_mid_stream_fallback_continuation=given).enable_mid_stream_fallback_continuation + is expected + ) # pyright: ignore[reportAttributeAccessIssue] + + +# --- logging level ------------------------------------------------------------------------- + + +@pytest.fixture +def router_log_level() -> Generator[int, None, None]: + level: Final = verbose_router_logger.level + yield level + verbose_router_logger.setLevel(level) + + +@pytest.mark.parametrize( + ("set_verbose", "debug_level", "expected"), + [(True, "DEBUG", logging.DEBUG), (True, "INFO", logging.INFO)], +) +def test_set_verbose_sets_the_router_logger_to_the_debug_level( + make_core: MakeCore, router_log_level: object, set_verbose: bool, debug_level: str, expected: int +) -> None: + verbose_router_logger.setLevel(logging.WARNING) + _one(make_core, set_verbose=set_verbose, debug_level=debug_level) + assert verbose_router_logger.level == expected + + +def test_debug_level_without_set_verbose_leaves_the_router_logger_alone( + make_core: MakeCore, router_log_level: object +) -> None: + verbose_router_logger.setLevel(logging.WARNING) + _one(make_core, debug_level="DEBUG") + assert verbose_router_logger.level == logging.WARNING + + +# --- redis and response caching ---------------------------------------------------------- + + +def test_redis_host_and_port_give_the_routing_state_a_redis_cache(make_core: MakeCore) -> None: + """The port is passed on as a string.""" + core: Final = _one( + make_core, redis_host="localhost", redis_port=6399, redis_password="pw", cache_kwargs={"ssl": False} + ) + redis_cache: Final = core.health().store.redis_cache + assert isinstance(redis_cache, RedisCache) + assert {k: v for k, v in redis_cache.redis_kwargs.items() if k != "socket_timeout"} == { + "host": "localhost", + "port": "6399", + "password": "pw", + "ssl": False, + } + + +def test_redis_url_gives_the_routing_state_a_redis_cache(make_core: MakeCore) -> None: + redis_cache: Final = _one(make_core, redis_url="redis://localhost:6399").health().store.redis_cache + assert isinstance(redis_cache, RedisCache) + assert redis_cache.redis_kwargs["url"] == "redis://localhost:6399" + + +def test_redis_db_is_passed_on_as_a_string(make_core: MakeCore) -> None: + redis_cache: Final = _one(make_core, redis_host="localhost", redis_port=6399, redis_db=3).health().store.redis_cache + assert isinstance(redis_cache, RedisCache) + assert redis_cache.redis_kwargs["db"] == "3" + + +def test_redis_host_without_a_port_creates_no_redis_cache(make_core: MakeCore) -> None: + assert _one(make_core, redis_host="localhost").health().store.redis_cache is None + + +def test_without_redis_arguments_the_routing_state_is_in_memory_only(make_core: MakeCore) -> None: + assert _one(make_core).health().store.redis_cache is None + + +def test_the_routing_state_and_prompt_caching_pins_share_one_redis_cache(make_core: MakeCore) -> None: + core: Final = _one(make_core, redis_host="localhost", redis_port=6399) + assert core.registry().plugins("prompt_caching_pins").redis_cache is core.health().store.redis_cache + + +def test_applying_a_redis_cache_attaches_it_to_the_routing_state_and_the_pins(make_core: MakeCore) -> None: + core: Final = _one(make_core) + redis_cache: Final = RedisCache(host="localhost", port=6399) + core.apply(ChangeSet(redis_cache=redis_cache)) + assert core.health().store.redis_cache is redis_cache + assert core.registry().plugins("prompt_caching_pins").redis_cache is redis_cache + + +def test_cache_responses_installs_a_local_response_cache_on_the_host(make_core: MakeCore, host: ScriptedHost) -> None: + core: Final = _one(make_core, cache_responses=True) + assert isinstance(host.cache, litellm.Cache) + assert host.cache.type == "local" + assert core.settings().cache_responses is True + + +def test_cache_responses_with_redis_installs_a_redis_response_cache(make_core: MakeCore, host: ScriptedHost) -> None: + _one(make_core, cache_responses=True, redis_host="localhost", redis_port=6399) + assert isinstance(host.cache, litellm.Cache) + assert host.cache.type == "redis" + + +def test_cache_responses_keeps_a_response_cache_the_host_already_has(make_core: MakeCore, host: ScriptedHost) -> None: + existing: Final = object() + host.cache = existing + _one(make_core, cache_responses=True) + assert host.cache is existing + + +def test_without_cache_responses_the_host_response_cache_is_untouched(make_core: MakeCore, host: ScriptedHost) -> None: + _one(make_core, redis_host="localhost", redis_port=6399) + assert host.cache is None + + +# --- callback registration ---------------------------------------------------------------- + + +@pytest.mark.parametrize( + ("strategy", "callbacks", "input_callbacks"), + [ + ("simple-shuffle", [], []), + ("lar1", [], []), + ("least-busy", ["LeastBusyLoggingHandler"], ["LeastBusyLoggingHandler"]), + ("usage-based-routing", ["LowestTPMLoggingHandler"], []), + ("usage-based-routing-v2", ["LowestTPMLoggingHandler_v2"], []), + ("latency-based-routing", ["LowestLatencyLoggingHandler"], []), + ("cost-based-routing", ["LowestCostLoggingHandler"], []), + ], +) +def test_each_routing_strategy_registers_its_selector_on_the_host( + make_core: MakeCore, host: ScriptedHost, strategy: str, callbacks: list[str], input_callbacks: list[str] +) -> None: + _one(make_core, routing_strategy=strategy) + names: Final = host_callback_names(host) + assert (names["callbacks"], names["input_callback"]) == (callbacks, input_callbacks) + + +def test_the_least_busy_selector_is_the_router_leastbusy_logger(make_core: MakeCore, host: ScriptedHost) -> None: + core: Final = _one(make_core, routing_strategy="least-busy") + assert host.lists["callbacks"] == [core.leastbusy_logger] # pyright: ignore[reportAttributeAccessIssue] + + +# --- alerting ---------------------------------------------------------------------------------- + + +def test_alerting_config_registers_slack_alerting_and_its_slow_response_callback( + make_core: MakeCore, host: ScriptedHost +) -> None: + core: Final = _one( + make_core, alerting_config=AlertingConfig(webhook_url="https://example.invalid/hook", alerting_threshold=10) + ) + names: Final = host_callback_names(host) + assert names["callbacks"] == ["SlackAlerting"] + assert "response_taking_too_long_callback" in names["success_callback"] + logger: Final = core.observe().slack_alerting_logger + assert host.lists["callbacks"] == [logger] + assert logger is not None + assert (logger.alerting_threshold, logger.alerting, logger.default_webhook_url) == ( + 10.0, + ["slack"], + "https://example.invalid/hook", + ) + + +def test_the_python_core_registers_its_sync_outcome_callback_before_the_slow_response_callback( + host: ScriptedHost, +) -> None: + """Python core only: it learns sync outcomes from a success callback on its host. The Rust core + records outcomes from the attempt itself and registers no outcome callbacks.""" + from tests.unit.test_router.core.scripted_host import python_core + + python_core( + host, + model_list=[deployment("g", "a")], + alerting_config=AlertingConfig(webhook_url="https://example.invalid/hook", alerting_threshold=10), + ) + assert host_callback_names(host)["success_callback"] == [ + "sync_deployment_callback_on_success", + "response_taking_too_long_callback", + ] + + +def test_observe_reports_no_slack_alerting_without_an_alerting_config(make_core: MakeCore) -> None: + assert _one(make_core).observe().slack_alerting_logger is None + + +# --- optional pre-call checks: registration ----------------------------------------------- + + +@pytest.mark.parametrize( + ("checks", "expected"), + [ + (["prompt_caching"], ["PromptCachingDeploymentCheck"]), + (["enforce_model_rate_limits"], ["ModelRateLimitingCheck"]), + (["deployment_affinity"], ["DeploymentAffinityCheck"]), + (["responses_api_deployment_check"], ["DeploymentAffinityCheck"]), + (["session_affinity"], ["DeploymentAffinityCheck"]), + (["deployment_affinity", "responses_api_deployment_check", "session_affinity"], ["DeploymentAffinityCheck"]), + (["encrypted_content_affinity"], ["EncryptedContentAffinityCheck"]), + (["prompt_caching", "prompt_caching"], ["PromptCachingDeploymentCheck"]), + (["prompt_caching", "enforce_model_rate_limits"], ["PromptCachingDeploymentCheck", "ModelRateLimitingCheck"]), + ], +) +def test_each_optional_pre_call_check_registers_one_callback_on_the_router_and_the_host( + make_core: MakeCore, host: ScriptedHost, checks: list[str], expected: list[str] +) -> None: + core: Final = _one(make_core, optional_pre_call_checks=checks) + assert [type(callback).__name__ for callback in core.optional_callbacks] == expected # pyright: ignore[reportAttributeAccessIssue] + assert host.lists["callbacks"] == core.optional_callbacks # pyright: ignore[reportAttributeAccessIssue] + + +@pytest.mark.parametrize( + ("checks", "flags"), + [ + (["deployment_affinity"], (True, False, False)), + (["responses_api_deployment_check"], (False, True, False)), + (["session_affinity"], (False, False, True)), + (["deployment_affinity", "session_affinity"], (True, False, True)), + ], +) +def test_affinity_checks_share_one_callback_with_one_flag_each( + make_core: MakeCore, checks: list[str], flags: tuple[bool, bool, bool] +) -> None: + (callback,) = _one( + make_core, optional_pre_call_checks=checks, deployment_affinity_ttl_seconds=60 + ).optional_callbacks # pyright: ignore[reportAttributeAccessIssue] + assert ( + callback.enable_user_key_affinity, + callback.enable_responses_api_affinity, + callback.enable_session_id_affinity, + ) == flags + assert callback.ttl_seconds == 60 + + +def test_encrypted_content_affinity_is_placed_before_deployment_affinity( + make_core: MakeCore, host: ScriptedHost +) -> None: + core: Final = _one(make_core, optional_pre_call_checks=["deployment_affinity", "encrypted_content_affinity"]) + expected: Final = ["EncryptedContentAffinityCheck", "DeploymentAffinityCheck"] + assert [type(callback).__name__ for callback in core.optional_callbacks] == expected # pyright: ignore[reportAttributeAccessIssue] + assert host_callback_names(host)["callbacks"] == expected + + +def test_forward_client_headers_by_model_group_registers_nothing(make_core: MakeCore, host: ScriptedHost) -> None: + """Oddity: a listed optional pre-call check that the router ignores.""" + core: Final = _one(make_core, optional_pre_call_checks=["forward_client_headers_by_model_group"]) + assert core.optional_callbacks is None # pyright: ignore[reportAttributeAccessIssue] + assert host.lists["callbacks"] == [] + + +def test_an_empty_optional_pre_call_check_list_registers_nothing(make_core: MakeCore, host: ScriptedHost) -> None: + assert _one(make_core, optional_pre_call_checks=[]).optional_callbacks is None # pyright: ignore[reportAttributeAccessIssue] + assert host.lists["callbacks"] == [] + + +def test_router_budget_limiting_cannot_be_built_outside_an_event_loop(make_core: MakeCore) -> None: + """Oddity: the budget limiter starts a background task in its constructor, so a router that needs + one cannot be built from synchronous code with no running loop.""" + with pytest.raises(RuntimeError, match="no running event loop"): + _one(make_core, optional_pre_call_checks=["router_budget_limiting"]) + + +@pytest.mark.asyncio +async def test_router_budget_limiting_registers_the_budget_limiter(make_core: MakeCore, host: ScriptedHost) -> None: + core: Final = _one(make_core, optional_pre_call_checks=["router_budget_limiting"]) + limiter: Final = core.router_budget_logger # pyright: ignore[reportAttributeAccessIssue] + assert isinstance(limiter, RouterBudgetLimiting) + assert core.registry().plugins("budget_limiter") is limiter + assert host.lists["callbacks"] == [limiter] + core.retire() + + +@pytest.mark.asyncio +async def test_a_provider_budget_config_adds_the_budget_limiter_by_itself(make_core: MakeCore) -> None: + core: Final = _one(make_core, provider_budget_config={"openai": {"budget_limit": 1, "time_period": "1d"}}) + assert [type(callback).__name__ for callback in core.optional_callbacks] == ["RouterBudgetLimiting"] # pyright: ignore[reportAttributeAccessIssue] + core.retire() + + +@pytest.mark.asyncio +async def test_a_deployment_budget_adds_the_budget_limiter_by_itself(make_core: MakeCore) -> None: + core: Final = make_core(model_list=[deployment("g", "a", max_budget=1, budget_duration="1d")]) + assert isinstance(core.registry().plugins("budget_limiter"), RouterBudgetLimiting) + core.retire() + + +def test_model_group_affinity_config_adds_an_affinity_callback_with_every_global_flag_off(make_core: MakeCore) -> None: + (callback,) = _one(make_core, model_group_affinity_config={"g": ["deployment_affinity"]}).optional_callbacks # pyright: ignore[reportAttributeAccessIssue] + assert type(callback).__name__ == "DeploymentAffinityCheck" + assert ( + callback.enable_user_key_affinity, + callback.enable_responses_api_affinity, + callback.enable_session_id_affinity, + ) == (False, False, False) + assert callback.model_group_affinity_config == {"g": ["deployment_affinity"]} + + +def test_model_group_encrypted_content_affinity_adds_both_affinity_callbacks( + make_core: MakeCore, host: ScriptedHost +) -> None: + _one(make_core, model_group_affinity_config={"g": ["encrypted_content_affinity"]}) + assert host_callback_names(host)["callbacks"] == ["EncryptedContentAffinityCheck", "DeploymentAffinityCheck"] + + +def test_mid_stream_fallback_continuation_registers_a_global_prefill_check( + make_core: MakeCore, host: ScriptedHost +) -> None: + core: Final = _one(make_core, enable_mid_stream_fallback_continuation=True) + assert host_callback_names(host)["callbacks"] == ["ContinuationPrefillDeploymentCheck"] + assert core.optional_callbacks is None # pyright: ignore[reportAttributeAccessIssue] + + +def test_add_optional_pre_call_checks_after_construction_registers_the_new_check( + make_core: MakeCore, host: ScriptedHost +) -> None: + core: Final = _one(make_core, optional_pre_call_checks=["prompt_caching"]) + core.add_optional_pre_call_checks(["enforce_model_rate_limits", "prompt_caching"]) # pyright: ignore[reportAttributeAccessIssue] + assert host_callback_names(host)["callbacks"] == ["PromptCachingDeploymentCheck", "ModelRateLimitingCheck"] + + +def test_add_optional_pre_call_checks_with_none_and_apply_default_settings_change_nothing( + make_core: MakeCore, host: ScriptedHost +) -> None: + core: Final = _one(make_core, optional_pre_call_checks=["prompt_caching"]) + before: Final = list(host.lists["callbacks"]) + core.add_optional_pre_call_checks(None) # pyright: ignore[reportAttributeAccessIssue] + core.apply_default_settings() # pyright: ignore[reportAttributeAccessIssue] + assert host.lists["callbacks"] == before + + +def test_runtime_optional_pre_call_checks_drop_unlisted_toggleable_checks_but_keep_affinity( + make_core: MakeCore, host: ScriptedHost +) -> None: + collect_other_routers() + core: Final = _one(make_core, optional_pre_call_checks=["prompt_caching", "deployment_affinity"]) + core.apply(ChangeSet(settings={"optional_pre_call_checks": ["enforce_model_rate_limits"]})) + expected: Final = ["DeploymentAffinityCheck", "ModelRateLimitingCheck"] + assert [type(callback).__name__ for callback in core.optional_callbacks] == expected # pyright: ignore[reportAttributeAccessIssue] + assert host_callback_names(host)["callbacks"] == expected + + +def test_runtime_optional_pre_call_checks_none_changes_nothing(make_core: MakeCore, host: ScriptedHost) -> None: + core: Final = _one(make_core, optional_pre_call_checks=["prompt_caching"]) + core.apply(ChangeSet(settings={"optional_pre_call_checks": None})) + assert host_callback_names(host)["callbacks"] == ["PromptCachingDeploymentCheck"] + + +def test_runtime_removal_keeps_the_host_callback_while_another_live_router_uses_that_check( + make_core: MakeCore, host: ScriptedHost +) -> None: + """Oddity: the check is left in this host's list because some other live router (here on another + host) still has one of that type.""" + collect_other_routers() + other_host: Final = CallbackHost() + from tests.unit.test_router.core.scripted_host import python_core + + other: Final = python_core( + other_host, model_list=[deployment("g", "o")], optional_pre_call_checks=["prompt_caching"] + ) + core: Final = _one(make_core, optional_pre_call_checks=["prompt_caching"]) + core.apply(ChangeSet(settings={"optional_pre_call_checks": []})) + assert core.optional_callbacks == [] # pyright: ignore[reportAttributeAccessIssue] + assert host_callback_names(host)["callbacks"] == ["PromptCachingDeploymentCheck"] + other.retire() + + +def test_enable_pre_call_checks_is_not_a_runtime_updatable_setting(make_core: MakeCore) -> None: + core: Final = _one(make_core) + core.apply(ChangeSet(settings={"enable_pre_call_checks": True})) + assert core.settings().enable_pre_call_checks is False + + +def test_enable_pre_call_checks_can_be_assigned(make_core: MakeCore) -> None: + core: Final = _one(make_core) + core.apply(ChangeSet(assign={"enable_pre_call_checks": True})) + assert core.settings().enable_pre_call_checks is True + + +# --- retire ------------------------------------------------------------------------------ + + +def test_retire_removes_selectors_optional_checks_and_alerting_but_keeps_foreign_callbacks( + make_core: MakeCore, host: ScriptedHost +) -> None: + collect_other_routers() + host.lists["callbacks"].append("foreign") + host.lists["success_callback"].append("foreign-success") + core: Final = _one( + make_core, + routing_strategy="least-busy", + optional_pre_call_checks=["enforce_model_rate_limits"], + alerting_config=AlertingConfig(webhook_url="https://example.invalid/hook"), + ) + core.retire() + assert {name: callbacks for name, callbacks in host.lists.items() if callbacks} == { + "callbacks": ["foreign"], + "success_callback": ["foreign-success"], + } + + +def test_retire_twice_is_harmless(make_core: MakeCore, host: ScriptedHost) -> None: + core: Final = _one(make_core) + core.retire() + core.retire() + assert {name: callbacks for name, callbacks in host.lists.items() if callbacks} == {} + + +def test_retiring_one_router_keeps_the_callbacks_of_another_on_the_same_host( + make_core: MakeCore, host: ScriptedHost +) -> None: + collect_other_routers() + first: Final = _one(make_core, routing_strategy="least-busy") + second: Final = _one(make_core, routing_strategy="least-busy") + first.retire() + assert host.lists["callbacks"] == [second.leastbusy_logger] # pyright: ignore[reportAttributeAccessIssue] + assert host.lists["input_callback"] == [second.leastbusy_logger] # pyright: ignore[reportAttributeAccessIssue] + + +def test_retiring_one_python_core_keeps_the_outcome_callback_of_another_on_the_same_host( + host: ScriptedHost, +) -> None: + """Python core only: each one registers an async outcome callback bound to itself; the Rust core + records outcomes from the attempt itself and registers none.""" + from tests.unit.test_router.core.scripted_host import python_core + + collect_other_routers() + first: Final = python_core(host, model_list=[deployment("g", "a")], routing_strategy="least-busy") + second: Final = python_core(host, model_list=[deployment("g", "a")], routing_strategy="least-busy") + first.retire() + assert host_callback_names(host)["_async_success_callback"] == ["deployment_callback_on_success"] + assert all(callback.__self__ is second for callback in host.lists["_async_success_callback"]) # pyright: ignore[reportAttributeAccessIssue] + + +def test_retire_hands_a_vacated_check_to_another_live_router_on_a_different_host( + make_core: MakeCore, host: ScriptedHost +) -> None: + """Oddity: retire re-registers the same-keyed check of any other live router in the process into + this host's list, even one built on a different host.""" + collect_other_routers() + from tests.unit.test_router.core.scripted_host import python_core + + other: Final = python_core( + CallbackHost(), model_list=[deployment("g", "o")], optional_pre_call_checks=["prompt_caching"] + ) + core: Final = _one(make_core, optional_pre_call_checks=["prompt_caching"]) + core.retire() + assert host.lists["callbacks"] == other.optional_callbacks + other.retire() + + +# --- observe ----------------------------------------------------------------------------- + + +class RecordingLogging: + def __init__(self) -> None: + self.calls: list[tuple[str, object]] = [] + + def add_dynamic_function_callbacks(self, **callbacks: object) -> None: # kwargs-ok: the logging object's API + self.calls.append(("functions", sorted(callbacks))) + + def replace_global_callback(self, callback: object) -> None: + self.calls.append(("replace", type(callback).__name__)) + + def add_dynamic_callback(self, callback: object) -> None: + self.calls.append(("add", type(callback).__name__)) + + +def _bound_outcome_callbacks(core: RouterCore) -> list[tuple[str, object]]: # mutable-ok: a fresh list + return ( + [("functions", ["async_failure", "async_success", "failure"])] + if records_outcomes_through_callbacks(core) + else [] + ) + + +def test_observe_request_scoped_removes_every_callback_the_router_registered_globally( + make_core: MakeCore, host: ScriptedHost +) -> None: + collect_other_routers() + core: Final = _one(make_core, routing_strategy="least-busy", optional_pre_call_checks=["prompt_caching"]) + observers: Final = core.observe(request_scoped=True) + assert observers.slack_alerting_logger is None + assert {name: callbacks for name, callbacks in host.lists.items() if callbacks} == {} + + +def test_observe_with_a_logging_object_binds_the_router_outcome_callbacks( + make_core: MakeCore, host: ScriptedHost +) -> None: + """The Python core binds its three outcome callbacks; the Rust core records every attempt's + outcome itself and binds none, as it registers none on the host.""" + core: Final = _one(make_core) + before: Final = host_callback_names(host) + logging_obj: Final = RecordingLogging() + core.observe(logging_obj) # pyright: ignore[reportArgumentType] + assert logging_obj.calls == _bound_outcome_callbacks(core) + assert host_callback_names(host) == before + + +def test_observe_request_scoped_then_binding_hands_selectors_and_policies_to_the_request(make_core: MakeCore) -> None: + collect_other_routers() + core: Final = _one(make_core, routing_strategy="least-busy", optional_pre_call_checks=["prompt_caching"]) + core.observe(request_scoped=True) + logging_obj: Final = RecordingLogging() + core.observe(logging_obj) # pyright: ignore[reportArgumentType] + assert logging_obj.calls == [ + *_bound_outcome_callbacks(core), + ("replace", "LeastBusyLoggingHandler"), + ("add", "PromptCachingDeploymentCheck"), + ] + + +def test_binding_without_request_scoping_hands_the_request_no_selector_or_policy(make_core: MakeCore) -> None: + """Without ``observe(request_scoped=True)`` the selectors and checks stay on the host lists, so a + bound request gets none of them.""" + collect_other_routers() + core: Final = _one(make_core, routing_strategy="least-busy", optional_pre_call_checks=["prompt_caching"]) + logging_obj: Final = RecordingLogging() + core.observe(logging_obj) # pyright: ignore[reportCallIssue] # observe takes a logging object on every core + assert logging_obj.calls == _bound_outcome_callbacks(core) + + +def test_binding_skips_a_policy_the_host_already_runs(make_core: MakeCore, host: ScriptedHost) -> None: + """A request-scoped policy whose key is back on the host's ``callbacks`` list (here re-added by + hand) is not added to the request again; the selector still replaces the host's.""" + collect_other_routers() + core: Final = _one(make_core, routing_strategy="least-busy", optional_pre_call_checks=["prompt_caching"]) + core.observe(request_scoped=True) # pyright: ignore[reportCallIssue] # see above + host.lists["callbacks"].append(PromptCachingDeploymentCheck(cache=DualCache())) + logging_obj: Final = RecordingLogging() + core.observe(logging_obj) # pyright: ignore[reportCallIssue] # observe takes a logging object on every core + assert logging_obj.calls == [*_bound_outcome_callbacks(core), ("replace", "LeastBusyLoggingHandler")] + + +# --- register_plugin --------------------------------------------------------------------- + + +class PickB(CustomRoutingStrategyBase): + def __init__(self) -> None: + self.asked: list[str] = [] + + @staticmethod + def _b() -> dict[str, object]: + return deployment("g", "b") + + async def async_get_available_deployment( # pyright: ignore[reportIncompatibleMethodOverride] + self, + model: str, + messages: object = None, + input: object = None, + specific_deployment: bool = False, + request_kwargs: object = None, + ) -> dict[str, object]: + self.asked.append(model) + return self._b() + + def get_available_deployment( # pyright: ignore[reportIncompatibleMethodOverride] + self, + model: str, + messages: object = None, + input: object = None, + specific_deployment: bool = False, + request_kwargs: object = None, + ) -> dict[str, object]: + self.asked.append(model) + return self._b() + + +@pytest.mark.asyncio +async def test_a_registered_routing_plugin_picks_for_select(make_core: MakeCore) -> None: + core: Final = make_core(model_list=[deployment("g", "a"), deployment("g", "b")]) + plugin: Final = PickB() + core.register_plugin(plugin) + picked: Final = {(await core.select("g"))["model_info"]["id"] for _ in range(5)} # pyright: ignore[reportIndexIssue] + assert picked == {"b"} + assert plugin.asked == ["g"] * 5 + + +@pytest.mark.asyncio +async def test_a_registered_routing_plugin_picks_for_async_and_sync_endpoints( + make_core: MakeCore, host: ScriptedHost +) -> None: + core: Final = make_core(model_list=[deployment("g", "a"), deployment("g", "b")]) + core.register_plugin(PickB()) + for _ in range(3): + await core.acompletion(model="g", messages=MESSAGES) # pyright: ignore[reportAttributeAccessIssue] + core.completion(model="g", messages=MESSAGES) # pyright: ignore[reportAttributeAccessIssue] + assert set(host.deployments_called()) == {"b"} + + +@pytest.mark.asyncio +async def test_changing_the_routing_strategy_drops_a_registered_routing_plugin(make_core: MakeCore) -> None: + core: Final = make_core(model_list=[deployment("g", "a"), deployment("g", "b")]) + plugin: Final = PickB() + core.register_plugin(plugin) + core.apply(ChangeSet(settings={"routing_strategy": "least-busy"})) + for _ in range(5): + await core.select("g") + assert plugin.asked == [] + + +# --- flush_cache / reset ------------------------------------------------------------------ + + +def test_flush_cache_drops_the_response_cache_routing_state_and_prompt_caching_pins( + make_core: MakeCore, host: ScriptedHost +) -> None: + core: Final = _one(make_core) + host.cache = object() + core.health().store.set_cache("k", 1) + core.registry().plugins("prompt_caching_pins").set_cache("p", 1) + before: Final = host_callback_names(host) + core.flush_cache() # pyright: ignore[reportAttributeAccessIssue] + assert host.cache is None + assert core.health().store.get_cache("k") is None + assert core.registry().plugins("prompt_caching_pins").get_cache("p") is None + assert host_callback_names(host) == before + + +def test_reset_empties_every_success_and_failure_list_and_the_retry_policy( + make_core: MakeCore, host: ScriptedHost +) -> None: + """Note: reset clears the host's success/failure lists outright, foreign callbacks included, but + leaves ``callbacks`` and ``input_callback`` alone.""" + core: Final = _one(make_core, retry_policy={"RateLimitErrorRetries": 3}, routing_strategy="least-busy") + host.lists["success_callback"].append("foreign") + core.health().store.set_cache("k", 1) + core.reset() # pyright: ignore[reportAttributeAccessIssue] + assert {name: callbacks for name, callbacks in host_callback_names(host).items() if callbacks} == { + "callbacks": ["LeastBusyLoggingHandler"], + "input_callback": ["LeastBusyLoggingHandler"], + } + assert core.retry_policy is None # pyright: ignore[reportAttributeAccessIssue] + assert core.health().store.get_cache("k") is None + + +# --- attribute writes ---------------------------------------------------------------------- + + +def test_assigning_num_retries_after_construction_is_reported_by_settings(make_core: MakeCore) -> None: + core: Final = _one(make_core) + core.num_retries = 7 # pyright: ignore[reportAttributeAccessIssue] + assert core.settings().as_dict()["num_retries"] == 7 + + +def test_assigning_a_setting_attribute_replaces_the_settings_snapshot(make_core: MakeCore) -> None: + core: Final = _one(make_core) + before: Final = core.settings() + core.enable_pre_call_checks = True # pyright: ignore[reportAttributeAccessIssue] + assert core.settings() is not before + assert core.settings().enable_pre_call_checks is True + assert before.enable_pre_call_checks is False + + +def test_assigning_a_plain_setting_keeps_the_settings_snapshot(make_core: MakeCore) -> None: + core: Final = _one(make_core) + before: Final = core.settings() + core.num_retries = 7 # pyright: ignore[reportAttributeAccessIssue] + assert core.settings() is before + + +def test_assigning_an_unrelated_attribute_keeps_the_settings_snapshot(make_core: MakeCore) -> None: + core: Final = _one(make_core) + before: Final = core.settings() + core.some_custom_attribute = 1 # pyright: ignore[reportAttributeAccessIssue] + assert core.settings() is before + assert core.some_custom_attribute == 1 # pyright: ignore[reportAttributeAccessIssue] + + +def test_assigning_a_new_model_list_advances_the_registry_version_once(make_core: MakeCore) -> None: + core: Final = _one(make_core) + version: Final = core.registry().version + core.model_list = [deployment("g", "z")] # pyright: ignore[reportAttributeAccessIssue] + assert core.registry().version == version + 1 + assert [d["model_info"]["id"] for d in core.registry().all()] == ["z"] + + +def test_deleting_a_setting_attribute_makes_it_unset_in_settings(make_core: MakeCore) -> None: + core: Final = _one(make_core) + del core.routing_strategy # pyright: ignore[reportAttributeAccessIssue] + assert core.settings().routing_strategy is None + assert core.settings().is_set("routing_strategy") is False + + +def test_registry_settings_and_health_views_are_built_once(make_core: MakeCore) -> None: + core: Final = _one(make_core) + assert core.registry() is core.registry() + assert core.settings() is core.settings() + assert core.health() is core.health() + + +# --- get_valid_args ----------------------------------------------------------------------- + + +def test_get_valid_args_lists_every_constructor_argument_but_self_and_the_host() -> None: + names: Final = PythonCore.get_valid_args() + assert names[0] == "model_list" + assert "self" not in names + assert "_host" not in names + assert {"routing_strategy", "optional_pre_call_checks", "enable_pre_call_checks", "redis_url"} <= set(names) + assert len(names) == len(set(names)) + + +def test_get_valid_args_matches_the_router_init_args() -> None: + assert tuple(PythonCore.get_valid_args()) == tuple(router_init_args()) + + +def test_get_valid_args_on_a_built_router(make_core: MakeCore) -> None: + assert _one(make_core).get_valid_args() == PythonCore.get_valid_args() # pyright: ignore[reportAttributeAccessIssue] + + +# --- print_deployment ---------------------------------------------------------------------- + +SECRET_DEPLOYMENT: Final = { + "model_name": "g", + "litellm_params": { + "model": "openai/x", + "api_key": "sk-abcdef", + "api_base": "https://h", + "aws_secret_access_key": "secretvalue", + }, + "model_info": {"id": "x"}, +} + + +def test_print_deployment_shows_two_characters_of_the_api_key_then_ten_stars(make_core: MakeCore) -> None: + printed: Final = _one(make_core).print_deployment(SECRET_DEPLOYMENT) # pyright: ignore[reportAttributeAccessIssue] + assert printed["litellm_params"] == { + "model": "openai/x", + "api_key": "sk**********", + "api_base": "https://h", + "aws_secret_access_key": "secretvalue", + } + assert printed["model_info"] == {"id": "x"} + + +def test_print_deployment_leaves_the_given_deployment_unchanged(make_core: MakeCore) -> None: + _one(make_core).print_deployment(SECRET_DEPLOYMENT) # pyright: ignore[reportAttributeAccessIssue] + assert SECRET_DEPLOYMENT["litellm_params"]["api_key"] == "sk-abcdef" # pyright: ignore[reportIndexIssue] + + +def test_print_deployment_without_an_api_key_returns_an_equal_copy(make_core: MakeCore) -> None: + given: Final = {"litellm_params": {"model": "m"}} + printed: Final = _one(make_core).print_deployment(given) # pyright: ignore[reportAttributeAccessIssue] + assert printed == given + assert printed is not given + + +def test_print_deployment_without_litellm_params_raises_key_error(make_core: MakeCore) -> None: + with pytest.raises(KeyError, match="litellm_params"): + _one(make_core).print_deployment({"model_name": "x"}) # pyright: ignore[reportAttributeAccessIssue] + + +def test_print_deployment_with_redaction_masks_every_secret_keeping_its_length( + make_core: MakeCore, host: ScriptedHost +) -> None: + core: Final = _one(make_core) + host.set_dynamic(redact_user_api_key_info=True) + printed: Final = core.print_deployment(SECRET_DEPLOYMENT) # pyright: ignore[reportAttributeAccessIssue] + assert printed["litellm_params"] == { + "model": "openai/x", + "api_key": "sk*******", + "api_base": "https://h", + "aws_secret_access_key": "se*********", + } + + +def test_settings_snapshot_reads_constructor_mappings_as_given(make_core: MakeCore) -> None: + alias: Final[Mapping[str, object]] = {"alias": "g"} + core: Final = _one(make_core, model_group_alias=dict(alias)) + assert core.registry().aliases() == alias diff --git a/tests/unit/test_router/core/test_construction_precall_characterization_precall.py b/tests/unit/test_router/core/test_construction_precall_characterization_precall.py new file mode 100644 index 00000000000..96d4d1879ce --- /dev/null +++ b/tests/unit/test_router/core/test_construction_precall_characterization_precall.py @@ -0,0 +1,707 @@ +""" +Characterization of the pre-call checks on a router core: with ``enable_pre_call_checks`` the +context window (token counting over messages, Responses input, instructions, tools and an Anthropic +system block), the allowed region and unsupported params (``response_format``) narrow a group's +deployments before one is picked; image requests skip text-only deployments; budget limits drop +deployments over budget. Covers the errors raised when nothing is left (type and message), which +deployment ``select`` and the endpoints reach, and that no provider call is made then. + +Each test asserts what the Python core does today; oddities are named in the test's docstring. +""" + +import base64 +import re +from collections.abc import Callable, Generator, Mapping +from typing import Final + +import pytest + +import litellm +from litellm.router_core import RouterCore +from litellm.types.router import RouterRateLimitError +from tests.unit.test_router.core.construction_precall_helpers import CallbackHost, restoring_model_cost, windowed +from tests.unit.test_router.core.scripted_host import ScriptedHost, deployment, python_core + +MakeCore = Callable[..., RouterCore] +PICKS: Final = 60 # three-way shuffles miss a deployment in 20 picks about once in a thousand runs +SHORT: Final = [{"role": "user", "content": "hi"}] +LONG_TEXT: Final = " ".join(["word"] * 100) +LONG: Final = [{"role": "user", "content": LONG_TEXT}] +EMBEDDING_IMAGE: Final = f"data:image/png;base64,{base64.b64encode(bytes(range(256)) * 1024).decode()}" +IMAGE: Final = [ + { + "role": "user", + "content": [ + {"type": "text", "text": "what is this"}, + {"type": "image_url", "image_url": {"url": "https://example.com/a.png"}}, + ], + } +] +UNCOUNTABLE: Final = [ + {"role": "user", "content": [{"type": "image_url", "image_url": {"detail": "not-a-real-detail-value"}}]} +] +JSON_MODE: Final = {"response_format": {"type": "json_object"}} +NO_JSON_MODE_MODEL: Final = "openai/gpt-3.5-turbo-16k" +UNMAPPED_MODEL: Final = "openai/a-model-the-catalog-does-not-map" +JSON_MODE_MODEL: Final = "openai/gpt-4o" +CONTEXT_WINDOW_PREFIX: Final = ( + "litellm.ContextWindowExceededError: litellm.BadRequestError: litellm._pre_call_checks: Context Window exceeded " + "for given call. No models have context window large enough for this call.\n" +) +UNSUPPORTED_MESSAGE: Final = ( + "litellm.UnsupportedParamsError: no deployment of model=g supports ['response_format']. To drop these, set " + "`litellm.drop_params=True` or for proxy `litellm_settings: drop_params: true`, or send " + "allowed_openai_params=['response_format'] to pass them through" +) + + +@pytest.fixture +def host() -> ScriptedHost: + return CallbackHost() + + +@pytest.fixture(autouse=True) +def isolated_model_cost() -> Generator[None, None, None]: + """Routers write their deployments' model info into the process cost map; undo it after each test.""" + yield from restoring_model_cost("gpt-3.5-turbo-16k") + + +async def _picked( + core: RouterCore, model: str = "g", times: int = PICKS, **select_arguments: object +) -> set[str]: # kwargs-ok: select's arguments + picked: Final[set[str]] = set() + for _ in range(times): + chosen = await core.select(model, **select_arguments) # pyright: ignore[reportCallIssue] + assert isinstance(chosen, Mapping) + picked.add(str(chosen["model_info"]["id"])) + return picked + + +def _small_and_big( + make_core: MakeCore, small: int = 5, **arguments: object +) -> RouterCore: # kwargs-ok: constructor arguments + return make_core( + model_list=[windowed("g", "small", small), windowed("g", "big", 10000)], + **{"enable_pre_call_checks": True, **arguments}, + ) + + +def _json_mode_group(make_core: MakeCore, **arguments: object) -> RouterCore: # kwargs-ok: constructor arguments + return make_core( + model_list=[deployment("g", "old", model=NO_JSON_MODE_MODEL), deployment("g", "new", model=JSON_MODE_MODEL)], + **{"enable_pre_call_checks": True, **arguments}, + ) + + +# --- context window ---------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_without_enable_pre_call_checks_a_too_small_context_window_is_still_picked(make_core: MakeCore) -> None: + core: Final = _small_and_big(make_core, enable_pre_call_checks=False) + assert await _picked(core, messages=LONG) == {"small", "big"} + + +@pytest.mark.asyncio +async def test_without_enable_pre_call_checks_a_too_long_prompt_reaches_the_provider( + make_core: MakeCore, host: ScriptedHost +) -> None: + core: Final = make_core(model_list=[windowed("g", "small", 5)]) + await core.acompletion(model="g", messages=LONG) # pyright: ignore[reportAttributeAccessIssue] + assert host.deployments_called() == ["small"] + + +@pytest.mark.asyncio +async def test_select_skips_a_deployment_whose_context_window_the_prompt_exceeds(make_core: MakeCore) -> None: + assert await _picked(_small_and_big(make_core), messages=LONG) == {"big"} + + +@pytest.mark.asyncio +async def test_a_short_prompt_may_pick_either_deployment(make_core: MakeCore) -> None: + assert await _picked(_small_and_big(make_core, small=50), messages=SHORT) == {"small", "big"} + + +@pytest.mark.asyncio +async def test_async_completion_reaches_only_a_deployment_whose_context_window_fits( + make_core: MakeCore, host: ScriptedHost +) -> None: + core: Final = _small_and_big(make_core) + for _ in range(8): + await core.acompletion(model="g", messages=LONG) # pyright: ignore[reportAttributeAccessIssue] + assert set(host.deployments_called()) == {"big"} + + +def test_sync_completion_reaches_only_a_deployment_whose_context_window_fits( + make_core: MakeCore, host: ScriptedHost +) -> None: + core: Final = _small_and_big(make_core) + for _ in range(8): + core.completion(model="g", messages=LONG) # pyright: ignore[reportAttributeAccessIssue] + assert set(host.deployments_called()) == {"big"} + + +@pytest.mark.asyncio +async def test_pass_through_selection_skips_a_deployment_whose_context_window_the_prompt_exceeds( + make_core: MakeCore, +) -> None: + core: Final = make_core( + model_list=[ + windowed("g", "small", 5, use_in_pass_through=True), + windowed("g", "big", 10000, use_in_pass_through=True), + ], + enable_pre_call_checks=True, + ) + assert await _picked(core, kind="pass_through", messages=LONG) == {"big"} + + +@pytest.mark.asyncio +async def test_select_without_messages_or_input_runs_no_context_window_check(make_core: MakeCore) -> None: + assert await _picked(_small_and_big(make_core)) == {"small", "big"} + + +@pytest.mark.asyncio +async def test_when_no_context_window_fits_select_raises_context_window_exceeded(make_core: MakeCore) -> None: + core: Final = make_core(model_list=[windowed("g", "small", 5)], enable_pre_call_checks=True) + with pytest.raises(litellm.ContextWindowExceededError) as error: + await core.select("g", messages=LONG) + message: Final = str(error.value) + assert message.startswith(CONTEXT_WINDOW_PREFIX) + assert re.fullmatch( + r"Model=openai/cw-small-[0-9a-f]{8}, Max Input Tokens=5, Got=\d+", message.removeprefix(CONTEXT_WINDOW_PREFIX) + ) + assert error.value.model == "g" + assert error.value.status_code == 400 + + +@pytest.mark.asyncio +async def test_when_no_context_window_fits_async_completion_raises_without_calling_the_provider( + make_core: MakeCore, host: ScriptedHost +) -> None: + core: Final = make_core(model_list=[windowed("g", "small", 5)], enable_pre_call_checks=True) + with pytest.raises(litellm.ContextWindowExceededError, match="Context Window exceeded for given call"): + await core.acompletion(model="g", messages=LONG) # pyright: ignore[reportAttributeAccessIssue] + assert host.calls == [] + assert host.sleeps == [] + + +def test_when_no_context_window_fits_sync_completion_raises_without_calling_the_provider( + make_core: MakeCore, host: ScriptedHost +) -> None: + core: Final = make_core(model_list=[windowed("g", "small", 5)], enable_pre_call_checks=True) + with pytest.raises(litellm.ContextWindowExceededError): + core.completion(model="g", messages=LONG) # pyright: ignore[reportAttributeAccessIssue] + assert host.calls == [] + + +@pytest.mark.asyncio +async def test_the_context_window_error_lists_each_deployment_back_to_back(make_core: MakeCore) -> None: + """Oddity: the per-deployment entries are concatenated with no separator between them.""" + core: Final = make_core(model_list=[windowed("g", "one", 5), windowed("g", "two", 6)], enable_pre_call_checks=True) + with pytest.raises(litellm.ContextWindowExceededError) as error: + await core.select("g", messages=LONG) + assert re.fullmatch( + r"Model=openai/cw-one-[0-9a-f]{8}, Max Input Tokens=5, Got=(\d+)Model=openai/cw-two-[0-9a-f]{8}, Max Input Tokens=6, Got=\1", + str(error.value).removeprefix(CONTEXT_WINDOW_PREFIX), + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("where", "expected"), + [("model_info", "azure/gpt-4-mi-base"), ("litellm_params", "azure/gpt-4-lp-base")], +) +async def test_the_context_window_error_names_a_deployment_by_its_base_model( + make_core: MakeCore, where: str, expected: str +) -> None: + extra: Final[dict[str, object]] = ( + {"model_info": {"base_model": expected}} if where == "model_info" else {"base_model": expected} + ) + core: Final = make_core(model_list=[windowed("g", f"base-{where}", 5, **extra)], enable_pre_call_checks=True) + with pytest.raises(litellm.ContextWindowExceededError) as error: + await core.select("g", messages=LONG) + assert f"Model={expected}, Max Input Tokens=5, Got=" in str(error.value) + + +@pytest.mark.asyncio +async def test_when_no_context_window_fits_the_context_window_fallback_group_answers( + make_core: MakeCore, host: ScriptedHost +) -> None: + core: Final = make_core( + model_list=[windowed("small-group", "small", 5), deployment("big-group", "big")], + enable_pre_call_checks=True, + context_window_fallbacks=[{"small-group": ["big-group"]}], + num_retries=0, + ) + await core.acompletion(model="small-group", messages=LONG) # pyright: ignore[reportAttributeAccessIssue] + assert host.deployments_called() == ["big"] + + +@pytest.mark.asyncio +async def test_a_deployment_without_a_known_context_window_is_kept_for_any_prompt(make_core: MakeCore) -> None: + core: Final = make_core( + model_list=[windowed("g", "small", 5), deployment("g", "unknown-window-model")], enable_pre_call_checks=True + ) + assert await _picked(core, messages=LONG) == {"unknown-window-model"} + + +@pytest.mark.asyncio +async def test_the_context_window_comes_from_the_cost_map_when_the_deployment_declares_none( + make_core: MakeCore, +) -> None: + core: Final = make_core( + model_list=[deployment("g", "gpt4", model="openai/gpt-4"), deployment("g", "gpt4o", model="openai/gpt-4o")], + enable_pre_call_checks=True, + ) + assert await _picked(core, times=10, messages=[{"role": "user", "content": "word " * 10000}]) == {"gpt4o"} + + +@pytest.mark.asyncio +async def test_an_uncountable_prompt_skips_the_context_window_check(make_core: MakeCore) -> None: + core: Final = make_core(model_list=[windowed("g", "tiny", 1)], enable_pre_call_checks=True) + assert await _picked(core, times=2, messages=UNCOUNTABLE) == {"tiny"} + + +@pytest.mark.asyncio +async def test_responses_input_text_is_counted_against_the_context_window(make_core: MakeCore) -> None: + assert await _picked(_small_and_big(make_core), times=10, input=LONG_TEXT) == {"big"} + + +@pytest.mark.asyncio +async def test_an_embedding_image_stays_out_of_the_context_window_on_a_backend_that_reads_it_as_media( + make_core: MakeCore, +) -> None: + core: Final = make_core( + model_list=[ + deployment("g", "image", model="vertex_ai/gemini-embedding-2", model_info={"max_input_tokens": 8192}) + ], + enable_pre_call_checks=True, + ) + assert await _picked(core, times=2, input=EMBEDDING_IMAGE) == {"image"} + + +@pytest.mark.asyncio +async def test_an_embedding_image_counts_as_text_on_a_text_only_backend(make_core: MakeCore) -> None: + core: Final = make_core( + model_list=[ + deployment("g", "text", model="openai/text-embedding-3-small", model_info={"max_input_tokens": 10}) + ], + enable_pre_call_checks=True, + ) + with pytest.raises(litellm.ContextWindowExceededError): + await _picked(core, times=1, input=EMBEDDING_IMAGE) + + +@pytest.mark.asyncio +async def test_a_short_responses_input_may_pick_either_deployment(make_core: MakeCore) -> None: + assert await _picked(_small_and_big(make_core, small=50), input="hi") == {"small", "big"} + + +@pytest.mark.asyncio +async def test_responses_instructions_are_counted_against_the_context_window(make_core: MakeCore) -> None: + core: Final = _small_and_big(make_core, small=50) + assert await _picked(core, request_kwargs={"instructions": LONG_TEXT}, input="hi") == {"big"} + + +@pytest.mark.asyncio +async def test_tool_definitions_are_counted_against_the_context_window(make_core: MakeCore) -> None: + tool: Final = { + "type": "function", + "function": {"name": "f", "description": LONG_TEXT, "parameters": {"type": "object", "properties": {}}}, + } + core: Final = _small_and_big(make_core, small=50) + assert await _picked(core, request_kwargs={"tools": [tool]}, messages=SHORT) == {"big"} + + +@pytest.mark.asyncio +async def test_an_anthropic_system_block_is_counted_against_the_context_window(make_core: MakeCore) -> None: + core: Final = _small_and_big(make_core, small=50) + assert await _picked(core, request_kwargs={"system": LONG_TEXT}, messages=SHORT) == {"big"} + + +@pytest.mark.asyncio +async def test_a_deployment_context_window_is_written_into_the_process_cost_map(make_core: MakeCore) -> None: + """Oddity: for a model the catalog does not map, a deployment's declared ``max_input_tokens`` lands on + the model's process-wide entry, so it outlives the router and narrows another router's deployment.""" + first: Final = make_core( + model_list=[deployment("g", "declares-five", model=UNMAPPED_MODEL, model_info={"max_input_tokens": 5})] + ) + first.retire() + other: Final = python_core( + CallbackHost(), + model_list=[deployment("g", "declares-nothing", model=UNMAPPED_MODEL)], + enable_pre_call_checks=True, + ) + with pytest.raises(litellm.ContextWindowExceededError, match="Max Input Tokens=5"): + await other.select("g", messages=SHORT) + other.retire() + + +@pytest.mark.asyncio +async def test_a_deployment_context_window_stays_off_a_mapped_models_process_entry(make_core: MakeCore) -> None: + """The catalog's token limits win on the shared entry: the declared limit binds its own deployment.""" + first: Final = make_core( + model_list=[deployment("g", "declares-five", model=NO_JSON_MODE_MODEL, model_info={"max_input_tokens": 5})] + ) + first.retire() + other: Final = python_core( + CallbackHost(), + model_list=[deployment("g", "declares-nothing", model=NO_JSON_MODE_MODEL)], + enable_pre_call_checks=True, + ) + assert await _picked(other, messages=SHORT) == {"declares-nothing"} + other.retire() + + +# --- region ------------------------------------------------------------------------------ + + +def _regional(make_core: MakeCore, **arguments: object) -> RouterCore: # kwargs-ok: constructor arguments + return make_core( + model_list=[ + deployment("g", "eu", region_name="eu"), + deployment("g", "us", region_name="us"), + deployment("g", "none"), + ], + **{"enable_pre_call_checks": True, **arguments}, + ) + + +@pytest.mark.asyncio +async def test_allowed_model_region_keeps_only_deployments_in_that_region(make_core: MakeCore) -> None: + core: Final = _regional(make_core) + assert await _picked(core, request_kwargs={"allowed_model_region": "eu"}, messages=SHORT) == {"eu"} + + +@pytest.mark.asyncio +async def test_allowed_model_region_is_ignored_without_enable_pre_call_checks(make_core: MakeCore) -> None: + core: Final = _regional(make_core, enable_pre_call_checks=False) + assert await _picked(core, request_kwargs={"allowed_model_region": "eu"}, messages=SHORT) == {"eu", "us", "none"} + + +@pytest.mark.asyncio +async def test_allowed_model_region_is_ignored_without_messages_or_input(make_core: MakeCore) -> None: + core: Final = _regional(make_core) + assert await _picked(core, request_kwargs={"allowed_model_region": "eu"}) == {"eu", "us", "none"} + + +@pytest.mark.asyncio +async def test_no_deployment_in_the_allowed_region_raises_router_rate_limit_error(make_core: MakeCore) -> None: + core: Final = _regional(make_core) + with pytest.raises(RouterRateLimitError) as error: + await core.select("g", {"allowed_model_region": "ap"}, messages=SHORT) + assert str(error.value) == ( + "No deployments available for selected model, Try again in 5 seconds. Passed model=g. " + "pre-call-checks=True, cooldown_list=[]" + ) + + +@pytest.mark.asyncio +async def test_no_deployment_in_the_allowed_region_never_calls_the_provider( + make_core: MakeCore, host: ScriptedHost +) -> None: + core: Final = _regional(make_core) + with pytest.raises(RouterRateLimitError): + await core.acompletion(model="g", messages=SHORT, allowed_model_region="ap") # pyright: ignore[reportAttributeAccessIssue] + assert host.calls == [] + + +@pytest.mark.asyncio +async def test_async_completion_with_an_allowed_region_reaches_only_that_region( + make_core: MakeCore, host: ScriptedHost +) -> None: + core: Final = _regional(make_core) + for _ in range(6): + await core.acompletion(model="g", messages=SHORT, allowed_model_region="us") # pyright: ignore[reportAttributeAccessIssue] + assert set(host.deployments_called()) == {"us"} + + +# --- unsupported params ------------------------------------------------------------------ + + +@pytest.mark.asyncio +async def test_with_drop_params_off_a_deployment_that_lacks_response_format_is_skipped( + make_core: MakeCore, host: ScriptedHost +) -> None: + host.set_dynamic(drop_params=False) + assert await _picked(_json_mode_group(make_core), request_kwargs=dict(JSON_MODE), messages=SHORT) == {"new"} + + +@pytest.mark.asyncio +@pytest.mark.parametrize("drop_params", [None, True]) +async def test_unless_drop_params_is_exactly_false_unsupported_params_are_not_checked( + make_core: MakeCore, host: ScriptedHost, drop_params: bool | None +) -> None: + """Oddity: the unset process default (``None``) skips the check as ``drop_params=True`` does.""" + host.set_dynamic(drop_params=drop_params) + assert await _picked(_json_mode_group(make_core), request_kwargs=dict(JSON_MODE), messages=SHORT) == {"old", "new"} + + +@pytest.mark.asyncio +async def test_unsupported_params_are_not_checked_without_enable_pre_call_checks( + make_core: MakeCore, host: ScriptedHost +) -> None: + host.set_dynamic(drop_params=False) + core: Final = _json_mode_group(make_core, enable_pre_call_checks=False) + assert await _picked(core, request_kwargs=dict(JSON_MODE), messages=SHORT) == {"old", "new"} + + +@pytest.mark.asyncio +async def test_unsupported_params_are_not_checked_without_messages_or_input( + make_core: MakeCore, host: ScriptedHost +) -> None: + host.set_dynamic(drop_params=False) + assert await _picked(_json_mode_group(make_core), request_kwargs=dict(JSON_MODE)) == {"old", "new"} + + +@pytest.mark.asyncio +async def test_a_string_response_format_is_not_checked(make_core: MakeCore, host: ScriptedHost) -> None: + host.set_dynamic(drop_params=False) + core: Final = _json_mode_group(make_core) + assert await _picked(core, request_kwargs={"response_format": "b64_json"}, messages=SHORT) == {"old", "new"} + + +@pytest.mark.asyncio +async def test_when_no_deployment_supports_the_param_select_raises_unsupported_params( + make_core: MakeCore, host: ScriptedHost +) -> None: + host.set_dynamic(drop_params=False) + core: Final = make_core(model_list=[deployment("g", "old", model=NO_JSON_MODE_MODEL)], enable_pre_call_checks=True) + with pytest.raises(litellm.UnsupportedParamsError) as error: + await core.select("g", dict(JSON_MODE), messages=SHORT) + assert str(error.value) == UNSUPPORTED_MESSAGE + assert error.value.status_code == 400 + + +@pytest.mark.asyncio +async def test_when_no_deployment_supports_the_param_async_completion_raises_without_calling_or_retrying( + make_core: MakeCore, host: ScriptedHost +) -> None: + host.set_dynamic(drop_params=False) + core: Final = make_core( + model_list=[deployment("g", "old", model=NO_JSON_MODE_MODEL)], enable_pre_call_checks=True, num_retries=2 + ) + with pytest.raises(litellm.UnsupportedParamsError, match="allowed_openai_params"): + await core.acompletion(model="g", messages=SHORT, **JSON_MODE) # pyright: ignore[reportAttributeAccessIssue] + assert host.calls == [] + assert host.sleeps == [] + + +def test_when_no_deployment_supports_the_param_sync_completion_raises_without_calling( + make_core: MakeCore, host: ScriptedHost +) -> None: + host.set_dynamic(drop_params=False) + core: Final = make_core(model_list=[deployment("g", "old", model=NO_JSON_MODE_MODEL)], enable_pre_call_checks=True) + with pytest.raises(litellm.UnsupportedParamsError): + core.completion(model="g", messages=SHORT, **JSON_MODE) # pyright: ignore[reportAttributeAccessIssue] + assert host.calls == [] + + +REQUEST_ESCAPES: Final = [ + {"allowed_openai_params": ["response_format"]}, + {"drop_params": True}, + {"additional_drop_params": ["response_format"]}, +] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("escape", REQUEST_ESCAPES, ids=["allowlist", "drop", "additional-drop"]) +async def test_a_request_that_allows_or_drops_the_param_keeps_the_only_deployment( + make_core: MakeCore, host: ScriptedHost, escape: dict[str, object] +) -> None: + host.set_dynamic(drop_params=False) + core: Final = make_core(model_list=[deployment("g", "old", model=NO_JSON_MODE_MODEL)], enable_pre_call_checks=True) + assert await _picked(core, times=2, request_kwargs={**JSON_MODE, **escape}, messages=SHORT) == {"old"} + + +@pytest.mark.asyncio +@pytest.mark.parametrize("escape", REQUEST_ESCAPES, ids=["allowlist", "drop", "additional-drop"]) +async def test_a_deployment_that_takes_the_param_is_preferred_over_dropping_it( + make_core: MakeCore, host: ScriptedHost, escape: dict[str, object] +) -> None: + host.set_dynamic(drop_params=False) + core: Final = _json_mode_group(make_core) + assert await _picked(core, request_kwargs={**JSON_MODE, **escape}, messages=SHORT) == {"new"} + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "deployment_escape", + [{"drop_params": True}, {"allowed_openai_params": ["response_format"]}], + ids=["drop", "allowlist"], +) +async def test_a_deployment_that_allows_or_drops_the_param_is_kept( + make_core: MakeCore, host: ScriptedHost, deployment_escape: dict[str, object] +) -> None: + host.set_dynamic(drop_params=False) + core: Final = make_core( + model_list=[deployment("g", "old", model=NO_JSON_MODE_MODEL, **deployment_escape)], enable_pre_call_checks=True + ) + assert await _picked(core, times=2, request_kwargs=dict(JSON_MODE), messages=SHORT) == {"old"} + + +@pytest.mark.asyncio +async def test_a_request_drop_params_false_overrides_the_deployment_drop_params( + make_core: MakeCore, host: ScriptedHost +) -> None: + host.set_dynamic(drop_params=False) + core: Final = make_core( + model_list=[deployment("g", "old", model=NO_JSON_MODE_MODEL, drop_params=True)], enable_pre_call_checks=True + ) + with pytest.raises(litellm.UnsupportedParamsError): + await core.select("g", {**JSON_MODE, "drop_params": False}, messages=SHORT) + + +@pytest.mark.asyncio +async def test_a_context_window_error_wins_over_an_unsupported_param(make_core: MakeCore, host: ScriptedHost) -> None: + host.set_dynamic(drop_params=False) + core: Final = make_core( + model_list=[deployment("g", "small-no-json", model=NO_JSON_MODE_MODEL, model_info={"max_input_tokens": 5})], + enable_pre_call_checks=True, + ) + with pytest.raises(litellm.ContextWindowExceededError): + await core.select("g", dict(JSON_MODE), messages=LONG) + core.retire() + + +@pytest.mark.asyncio +async def test_a_capable_deployment_outside_the_region_gives_no_deployments_rather_than_a_400( + make_core: MakeCore, host: ScriptedHost +) -> None: + host.set_dynamic(drop_params=False) + core: Final = make_core( + model_list=[ + deployment("g", "old", model=NO_JSON_MODE_MODEL, region_name="eu"), + deployment("g", "new", model=JSON_MODE_MODEL, region_name="us"), + ], + enable_pre_call_checks=True, + ) + with pytest.raises(RouterRateLimitError, match=r"^No deployments available for selected model"): + await core.select("g", {**JSON_MODE, "allowed_model_region": "eu"}, messages=SHORT) + + +@pytest.mark.asyncio +async def test_dropping_the_param_keeps_the_in_region_deployment_that_lacks_it( + make_core: MakeCore, host: ScriptedHost +) -> None: + host.set_dynamic(drop_params=False) + core: Final = make_core( + model_list=[ + deployment("g", "old", model=NO_JSON_MODE_MODEL, region_name="eu"), + deployment("g", "new", model=JSON_MODE_MODEL, region_name="us"), + ], + enable_pre_call_checks=True, + ) + request: Final = {**JSON_MODE, "allowed_model_region": "eu", "drop_params": True} + assert await _picked(core, times=5, request_kwargs=request, messages=SHORT) == {"old"} + + +@pytest.mark.asyncio +async def test_a_capable_deployment_in_cooldown_gives_no_deployments_rather_than_a_400( + make_core: MakeCore, host: ScriptedHost +) -> None: + host.set_dynamic(drop_params=False) + core: Final = _json_mode_group(make_core) + core.health().cooldowns.add_deployment_to_cooldown( + "new", original_exception=Exception("down"), exception_status=429, cooldown_time=60 + ) + with pytest.raises(RouterRateLimitError) as error: + await core.select("g", dict(JSON_MODE), messages=SHORT) + assert str(error.value).startswith( + "No deployments available for selected model, Try again in 60 seconds. Passed model=g. pre-call-checks=True" + ) + + +# --- images -------------------------------------------------------------------------------- + + +def _vision_group(make_core: MakeCore) -> RouterCore: + return make_core( + model_list=[ + deployment("g", "text", model_info={"supports_vision": False}), + deployment("g", "vision", model_info={"supports_vision": True}), + ] + ) + + +@pytest.mark.asyncio +async def test_an_image_request_skips_a_declared_text_only_deployment_even_without_pre_call_checks( + make_core: MakeCore, +) -> None: + assert await _picked(_vision_group(make_core), messages=IMAGE) == {"vision"} + + +@pytest.mark.asyncio +async def test_an_image_in_the_request_messages_counts_as_an_image_request(make_core: MakeCore) -> None: + assert await _picked(_vision_group(make_core), request_kwargs={"messages": IMAGE}) == {"vision"} + + +@pytest.mark.asyncio +async def test_a_text_request_may_pick_a_text_only_deployment(make_core: MakeCore) -> None: + assert await _picked(_vision_group(make_core), messages=SHORT) == {"text", "vision"} + + +@pytest.mark.asyncio +async def test_an_image_request_keeps_every_deployment_when_all_are_text_only(make_core: MakeCore) -> None: + core: Final = make_core( + model_list=[ + deployment("g", "text", model_info={"supports_vision": False}), + deployment("g", "text2", model_info={"supports_vision": False}), + ] + ) + assert await _picked(core, messages=IMAGE) == {"text", "text2"} + + +@pytest.mark.asyncio +async def test_async_completion_with_an_image_reaches_the_vision_deployment( + make_core: MakeCore, host: ScriptedHost +) -> None: + core: Final = _vision_group(make_core) + for _ in range(6): + await core.acompletion(model="g", messages=IMAGE) # pyright: ignore[reportAttributeAccessIssue] + assert set(host.deployments_called()) == {"vision"} + + +# --- budgets ------------------------------------------------------------------------------- + +PROVIDER_BUDGET: Final = {"openai": {"budget_limit": 1, "time_period": "1d"}} + + +@pytest.mark.asyncio +async def test_a_provider_over_budget_leaves_no_deployment_and_select_raises_value_error(make_core: MakeCore) -> None: + core: Final = make_core(model_list=[deployment("g", "a")], provider_budget_config=PROVIDER_BUDGET) + await core.health().store.async_set_cache("provider_spend:openai:1d", 5.0) + with pytest.raises(ValueError, match=r"^No deployments available - crossed budget: ") as error: + await core.select("g") + assert ( + str(error.value) + == "No deployments available - crossed budget: Exceeded budget for provider openai: 5.0 >= 1.0\n" + ) + core.retire() + + +@pytest.mark.asyncio +async def test_a_provider_over_budget_raises_from_async_completion_without_calling_the_provider( + make_core: MakeCore, host: ScriptedHost +) -> None: + core: Final = make_core(model_list=[deployment("g", "a")], provider_budget_config=PROVIDER_BUDGET) + await core.health().store.async_set_cache("provider_spend:openai:1d", 5.0) + with pytest.raises(ValueError, match=r"^No deployments available - crossed budget"): + await core.acompletion(model="g", messages=SHORT) # pyright: ignore[reportAttributeAccessIssue] + assert host.calls == [] + core.retire() + + +@pytest.mark.asyncio +async def test_a_provider_under_budget_is_picked(make_core: MakeCore) -> None: + core: Final = make_core(model_list=[deployment("g", "a")], provider_budget_config=PROVIDER_BUDGET) + await core.health().store.async_set_cache("provider_spend:openai:1d", 0.5) + assert await _picked(core, times=2) == {"a"} + core.retire() + + +@pytest.mark.asyncio +async def test_a_deployment_over_its_own_budget_is_skipped(make_core: MakeCore) -> None: + core: Final = make_core(model_list=[deployment("g", "a", max_budget=1, budget_duration="1d"), deployment("g", "b")]) + await core.health().store.async_set_cache("deployment_spend:a:1d", 5.0) + assert await _picked(core, times=10) == {"b"} + core.retire() diff --git a/tests/unit/test_router/core/test_deprecated_members.py b/tests/unit/test_router/core/test_deprecated_members.py new file mode 100644 index 00000000000..b73dda78ecd --- /dev/null +++ b/tests/unit/test_router/core/test_deprecated_members.py @@ -0,0 +1,665 @@ +""" +The deprecated Router members (``DeprecatedRouterMembers``, litellm/router_deprecated.py) on a Router +over each core: every old attribute reads what the Python core's ``_core_`` holds, through the +entry point named in its DeprecationWarning, and the old writes have the Python core's effect. +""" + +from collections.abc import Callable +from datetime import datetime +from typing import Final + +import pytest + +import litellm +from litellm.caching.dual_cache import DualCache +from litellm.litellm_core_utils.core_helpers import as_str_mapping +from litellm.router import Router +from litellm.router_api import SETTING_FIELDS, ChangeSet +from litellm.router_core import RouterCore +from litellm.router_utils.cooldown_cache import CooldownCache +from litellm.router_utils.kubernetes_pod_discovery import KubernetesPodDiscovery +from litellm.types.router import AllowedFailsPolicy, Deployment, RouterGeneralSettings, RoutingContext +from tests.unit.test_router.core.registry_helpers import MESSAGES, MakeCore, ids, team_deployment +from tests.unit.test_router.core.scripted_host import deployment + + +def as_router(core: RouterCore) -> Router: + """A ``litellm.Router`` facade over ``core``, as ``Router.from_config`` attaches the core it builds.""" + router: Final = Router.__new__(Router) + router._attach(core) # pyright: ignore[reportPrivateUsage, reportAttributeAccessIssue, reportUnknownMemberType] # the facade's own attach, as from_config calls it + return router + + +def picked_id(picked: object) -> object: + """The ``model_info.id`` of a picked deployment dict.""" + return (as_str_mapping((as_str_mapping(picked) or {}).get("model_info")) or {}).get("id") + + +def read(router: Router, name: str) -> object: + with pytest.warns(DeprecationWarning, match=rf"Router\.{name} is deprecated"): + return getattr(router, name) # pyright: ignore[reportAny] # each old member's own type + + +def write(router: Router, name: str, value: object) -> None: + with pytest.warns(DeprecationWarning, match=rf"Router\.{name} is deprecated"): + setattr(router, name, value) + + +_CONFIGURED: Final[dict[str, object]] = { # mutable-ok: constructor arguments, copied per router + "model_list": [ + deployment("gpt", "a"), + deployment("gpt", "b"), + deployment("claude", "c", model="anthropic/claude-x"), + team_deployment("int-t", "t1", "TA", "pub", model="openai/team-model"), + deployment("openai/*", "w", model="openai/*"), + ], + "fallbacks": [{"gpt": ["claude"]}], + "context_window_fallbacks": [{"gpt": ["claude"]}], + "content_policy_fallbacks": [{"claude": ["gpt"]}], + "cooldown_time": 7, + "allowed_fails": 2, + "allowed_fails_policy": AllowedFailsPolicy(RateLimitErrorAllowedFails=4), + "disable_cooldowns": True, + "enable_pre_call_checks": True, + "enable_tag_filtering": True, + "tag_filtering_match_any": False, + "tag_routing_prefix": "x-", + "default_litellm_params": {"metadata": {"k": "v"}}, + "routing_strategy": "least-busy", + "default_max_parallel_requests": 5, + "guardrail_list": [{"guardrail_name": "g", "litellm_params": {"guardrail": "x"}}], + "search_tools": [{"search_tool_name": "s", "litellm_params": {"search_provider": "tavily"}}], + "health_check_ignore_transient_errors": True, + "enable_health_check_routing": True, + "background_health_check_model_groups": ["gpt"], + "model_group_alias": {"fast": "gpt"}, + "cache_responses": True, +} + +# The Python core's values for _CONFIGURED, by setting (the objects it builds are compared by type). +_CONFIGURED_SETTINGS: Final[dict[str, object]] = { # mutable-ok: expected values + "routing_strategy": "least-busy", + "fallbacks": [{"gpt": ["claude"]}], + "context_window_fallbacks": [{"gpt": ["claude"]}], + "content_policy_fallbacks": [{"claude": ["gpt"]}], + "default_deployment": None, + "default_litellm_params": {"timeout": None, "max_retries": 0, "metadata": {"k": "v", "caching_groups": None}}, + "default_max_parallel_requests": 5, + "cooldown_time": 7, + "allowed_fails": 2, + "allowed_fails_policy": AllowedFailsPolicy(RateLimitErrorAllowedFails=4), + "disable_cooldowns": True, + "enable_pre_call_checks": True, + "enable_health_check_routing": True, + "health_check_ignore_transient_errors": True, + "background_health_check_model_groups": frozenset({"gpt"}), + "enable_tag_filtering": True, + "tag_filtering_match_any": False, + "tag_routing_prefix": "x-", + "router_general_settings": RouterGeneralSettings(), + "provider_budget_config": None, + "provider_default_deployment_ids": ["w"], + "fallback_access_check": None, + "fallback_budget_check": None, + "search_tools": [{"search_tool_name": "s", "litellm_params": {"search_provider": "tavily"}}], + "guardrail_list": [{"guardrail_name": "g", "litellm_params": {"guardrail": "x"}}], + "cache_responses": True, +} + + +def _configured(make_core: MakeCore) -> Router: + return as_router(make_core(**{name: _copy(value) for name, value in _CONFIGURED.items()})) + + +def _copy(value: object) -> object: + import copy + + return copy.deepcopy(value) + + +# -- settings ------------------------------------------------------------------------------------ + + +@pytest.mark.parametrize("name", sorted(set(SETTING_FIELDS) - {"kubernetes_pod_discovery"})) +def test_a_setting_attribute_reads_the_configured_setting(make_core: MakeCore, name: str) -> None: + router: Final = _configured(make_core) + value: Final = read(router, name) + assert value == _CONFIGURED_SETTINGS[name] + assert value is getattr(router.settings(), name) + + +def test_the_pod_discovery_setting_is_the_routers_pod_lookup(make_core: MakeCore) -> None: + router: Final = as_router(make_core(model_list=[deployment("gpt", "a")])) + discovery: Final = read(router, "kubernetes_pod_discovery") + assert isinstance(discovery, KubernetesPodDiscovery) + assert discovery is router.settings().kubernetes_pod_discovery + + +def test_default_settings_read_as_the_python_cores_defaults(make_core: MakeCore) -> None: + router: Final = as_router(make_core(model_list=[deployment("gpt", "a")])) + assert read(router, "routing_strategy") == "simple-shuffle" + assert read(router, "cooldown_time") == 5 + assert read(router, "allowed_fails") == 3 + assert read(router, "fallbacks") is None + assert read(router, "default_litellm_params") == { + "timeout": None, + "max_retries": 0, + "metadata": {"caching_groups": None}, + } + assert read(router, "tag_filtering_match_any") is True + assert read(router, "tag_routing_prefix") == "" + assert read(router, "cache_responses") is False + assert read(router, "search_tools") == [] + assert read(router, "guardrail_list") == [] + + +def test_a_setting_attribute_write_is_a_changeset(make_core: MakeCore) -> None: + router: Final = as_router(make_core(model_list=[deployment("gpt", "a")])) + write(router, "cooldown_time", 9) + write(router, "fallbacks", [{"gpt": ["gpt"]}]) + assert read(router, "cooldown_time") == 9 + assert router.settings().cooldown_time == 9 + assert read(router, "fallbacks") == [{"gpt": ["gpt"]}] + + +def test_a_deleted_setting_attribute_no_longer_reads(make_core: MakeCore) -> None: + router: Final = as_router(make_core(model_list=[deployment("gpt", "a")])) + with pytest.warns(DeprecationWarning, match=r"Router\.fallbacks is deprecated"): + del router.fallbacks + assert not router.settings().is_set("fallbacks") + with pytest.raises(AttributeError, match="'fallbacks'"): + read(router, "fallbacks") + + +# -- registry state ------------------------------------------------------------------------------ + +_REGISTRY_READS: Final[dict[str, Callable[[Router], object]]] = { + "model_list": lambda r: r.registry().all(), + "model_names": lambda r: r.registry().names(kind="group_set"), + "deployment_names": lambda r: r.registry().names(kind="deployment_names"), + "team_public_model_names": lambda r: r.registry().names(kind="team_public"), + "model_name_to_deployment_indices": lambda r: r.registry().index("model_name"), + "team_model_to_deployment_indices": lambda r: r.registry().index("team_model"), + "cached_deployment_model_info": lambda r: r.registry().deployment_model_info_cache(), + "model_group_alias": lambda r: r.registry().aliases(), + "pattern_router": lambda r: r.registry().patterns(), + "team_pattern_routers": lambda r: r.registry().patterns(team=True), + "routing_groups": lambda r: r.registry().routing_groups(), + "auto_routers": lambda r: r.registry().plugins("auto"), + "complexity_routers": lambda r: r.registry().plugins("complexity"), + "quality_routers": lambda r: r.registry().plugins("quality"), + "adaptive_routers": lambda r: r.registry().plugins("adaptive"), + "routing_plugins": lambda r: r.registry().plugins("routing"), + "prompt_caching_pin_cache": lambda r: r.registry().plugins("prompt_caching_pins"), + "cache": lambda r: r.health().store, + "cooldown_cache": lambda r: r.health().cooldowns, + "health_state_cache": lambda r: r.health().health_checks, + "deployment_latency_map": lambda r: r.health().latency, +} + + +@pytest.mark.parametrize("name", sorted(_REGISTRY_READS)) +def test_a_state_attribute_is_the_live_object_its_entry_point_returns(make_core: MakeCore, name: str) -> None: + router: Final = _configured(make_core) + assert read(router, name) is _REGISTRY_READS[name](router) + + +def test_the_registry_index_attributes_hold_the_python_cores_values(make_core: MakeCore) -> None: + router: Final = _configured(make_core) + assert read(router, "model_names") == {"gpt", "claude", "int-t", "openai/*"} + assert read(router, "deployment_names") == [ + "openai/a", + "openai/b", + "anthropic/claude-x", + "openai/team-model", + "openai/*", + ] + assert read(router, "team_public_model_names") == frozenset({"pub"}) + assert read(router, "model_name_to_deployment_indices") == { + "gpt": [0, 1], + "claude": [2], + "int-t": [3], + "openai/*": [4], + } + assert read(router, "team_model_to_deployment_indices") == {("TA", "pub"): [3]} + assert read(router, "model_group_alias") == {"fast": "gpt"} + assert read(router, "routing_groups") == {} + with pytest.warns(DeprecationWarning, match=r"Router\.routing_groups is deprecated"): + assert router._routing_groups == {} # pyright: ignore[reportPrivateUsage] # the old private alias + assert read(router, "routing_plugins") == [] + assert read(router, "deployment_latency_map") == { + "openai/a": 0, + "openai/b": 0, + "anthropic/claude-x": 0, + "openai/team-model": 0, + "openai/*": 0, + } + assert ids(read(router, "model_list")) == ["a", "b", "c", "t1", "w"] # pyright: ignore[reportArgumentType] # the live model list + + +def test_the_cached_model_info_attribute_is_the_lru_cache(make_core: MakeCore) -> None: + router: Final = as_router(make_core(model_list=[deployment("gpt", "a")])) + with pytest.warns(DeprecationWarning, match=r"Router\.cached_deployment_model_info is deprecated"): + cached: Final = router.cached_deployment_model_info + first: Final = cached("a", "gpt") + assert cached("a", "gpt") is first + assert cached.cache_info().hits == 1 + assert first == router.registry().deployment_model_info("a", "gpt") + + +def test_a_configured_routing_plugin_is_listed(make_core: MakeCore) -> None: + class Keep: + async def run(self, context: RoutingContext) -> RoutingContext: + return context + + plugin: Final = Keep() + router: Final = as_router(make_core(model_list=[deployment("gpt", "a")], plugins=[plugin])) + assert read(router, "routing_plugins") == [plugin] + + +def test_a_routing_plugin_appended_to_the_live_list_runs(make_core: MakeCore) -> None: + class Narrow: + async def run(self, context: RoutingContext) -> RoutingContext: + context.candidate_models = [m for m in context.candidate_models if m == "openai/b"] + return context + + class Keep: + async def run(self, context: RoutingContext) -> RoutingContext: + return context + + router: Final = as_router(make_core(model_list=[deployment("g", "a"), deployment("g", "b")], plugins=[Keep()])) + with pytest.warns(DeprecationWarning, match=r"Router\.routing_plugins is deprecated"): + plugins: Final = router.routing_plugins + plugins.append(Narrow()) + + import asyncio + + for _ in range(4): + picked = asyncio.run(router.select("g", {}, messages=MESSAGES)) # pyright: ignore[reportCallIssue] + assert picked_id(picked) == "b" + + +# -- state writes -------------------------------------------------------------------------------- + + +def test_a_registry_state_write_is_what_the_registry_reads(make_core: MakeCore) -> None: + router: Final = as_router(make_core(model_list=[deployment("gpt", "a")])) + write(router, "model_names", {"x"}) + write(router, "deployment_names", ["openai/x"]) + write(router, "model_name_to_deployment_indices", {"x": [0]}) + write(router, "auto_routers", {}) + assert read(router, "model_names") == {"x"} + assert router.registry().names(kind="group_set") == {"x"} + assert router.registry().names(kind="deployment_names") == ["openai/x"] + assert router.registry().index("model_name") == {"x": [0]} + assert router.registry().plugins("auto") == {} + + +def test_a_routing_state_write_is_what_health_reads(make_core: MakeCore) -> None: + router: Final = as_router(make_core(model_list=[deployment("gpt", "a")])) + store: Final = DualCache() + cooldowns: Final = CooldownCache(cache=store, default_cooldown_time=1) + latency: Final[dict[str, float]] = {"openai/a": 0.5} + write(router, "cache", store) + write(router, "cooldown_cache", cooldowns) + write(router, "deployment_latency_map", latency) + assert read(router, "cache") is store + assert router.health().store is store + assert router.health().cooldowns is cooldowns + assert router.health().latency is latency + + +def test_an_assigned_slack_alerting_logger_is_the_routers_observer(make_core: MakeCore) -> None: + from litellm.integrations.SlackAlerting.slack_alerting import SlackAlerting + + router: Final = as_router(make_core(model_list=[deployment("gpt", "a")])) + with ( + pytest.raises(AttributeError, match="slack_alerting_logger"), + pytest.warns(DeprecationWarning, match=r"Router\.slack_alerting_logger is deprecated"), + ): + _ = router.slack_alerting_logger + logger: Final = SlackAlerting(alerting=["slack"], default_webhook_url="https://example.invalid/hook") + write(router, "slack_alerting_logger", logger) + assert read(router, "slack_alerting_logger") is logger + assert router.observe().slack_alerting_logger is logger + + +@pytest.mark.asyncio +async def test_an_assigned_async_pick_answers_select(make_core: MakeCore) -> None: + router: Final = as_router(make_core(model_list=[deployment("gpt", "a"), deployment("gpt", "b")])) + calls: Final[list[str]] = [] # mutable-ok: the models the pick was asked for + + async def pick( + model: str, + request_kwargs: dict[str, object], # mutable-ok: the old pick's signature + messages: object = None, + input: object = None, + specific_deployment: object = False, + ) -> dict[str, object]: # mutable-ok: a deployment dict + calls.append(model) + return deployment("gpt", "b") + + write(router, "async_get_available_deployment", pick) + assert read(router, "async_get_available_deployment") is pick + picked: Final = await router.select("gpt", {}) + assert picked_id(picked) == "b" + assert calls == ["gpt"] + + +@pytest.mark.asyncio +async def test_the_async_pick_attribute_picks_like_select(make_core: MakeCore) -> None: + router: Final = as_router(make_core(model_list=[deployment("gpt", "a")])) + with pytest.warns(DeprecationWarning, match=r"Router\.async_get_available_deployment is deprecated"): + pick: Final = router.async_get_available_deployment + picked: Final = await pick(model="gpt", request_kwargs={}) + assert picked_id(picked) == "a" + + +# -- methods ------------------------------------------------------------------------------------- + + +def _call(router: Router, name: str, *args: object, **kwargs: object) -> object: # kwargs-ok: the old member's + with pytest.warns(DeprecationWarning, match=rf"Router\.{name} is deprecated"): + return getattr(router, name)(*args, **kwargs) # pyright: ignore[reportAny] # each old member's own type + + +def test_registry_methods_answer_through_the_registry(make_core: MakeCore) -> None: + router: Final = _configured(make_core) + registry: Final = router.registry() + found: Final = _call(router, "get_deployment", "a") + assert isinstance(found, Deployment) + assert found.model_info.id == "a" + assert _call(router, "has_model_id", "a") is True + assert _call(router, "get_model_info", "zzz") is None + assert ids(_call(router, "get_model_list", "gpt")) == ["a", "b"] # pyright: ignore[reportArgumentType] + assert ids(_call(router, "get_all_deployments", "gpt")) == ["a", "b"] # pyright: ignore[reportArgumentType] + assert ids(_call(router, "get_model_list_from_model_alias", "fast")) == ["a", "b"] # pyright: ignore[reportArgumentType] + assert _call(router, "get_model_names") == registry.names() + assert _call(router, "get_model_ids", "gpt") == ["a", "b"] + assert _call(router, "get_model_from_alias", "fast") == "gpt" + assert _call(router, "routable_model_group", "fast") == "gpt" + assert _call(router, "is_recognized_model", "claude") is True + assert _call(router, "get_model_group", "a") == registry.resolve("a", kind="group_of_id") + assert _call(router, "map_team_model", "pub", "TA") == "pub" + assert _call(router, "get_model_access_groups") == registry.access_groups() + assert _call(router, "get_deployment_credentials", "a") == {"api_key": "sk-test"} + assert _call(router, "get_router_deployment_budget_limiter") is None + assert _call(router, "should_include_deployment", "gpt", deployment("gpt", "a")) is True + assert _call(router, "get_router_model_info", None, "gpt", id="a") == registry.router_model_info( + None, "gpt", id="a" + ) + + +def test_settings_methods_answer_through_settings(make_core: MakeCore) -> None: + router: Final = _configured(make_core) + assert _call(router, "get_settings") == dict(router.settings().as_dict()) + assert _call(router, "get_available_guardrail", "g") == { + "guardrail_name": "g", + "litellm_params": {"guardrail": "x"}, + } + with ( + pytest.raises(ValueError, match="No guardrail found with name: nope"), + pytest.warns(DeprecationWarning, match=r"Router\.get_available_guardrail is deprecated"), + ): + router.get_available_guardrail("nope") + + +@pytest.mark.asyncio +async def test_health_and_selection_methods_answer_through_their_entry_points(make_core: MakeCore) -> None: + router: Final = as_router( + make_core(model_list=[deployment("gpt", "a"), deployment("pt", "p", use_in_pass_through=True)]) + ) + assert _call(router, "get_fully_blocked_model_names") == set() + with pytest.warns(DeprecationWarning, match=r"Router\.async_get_fully_unhealthy_model_names is deprecated"): + unhealthy: Final = await router.async_get_fully_unhealthy_model_names() + assert unhealthy == set() + with pytest.warns(DeprecationWarning, match=r"Router\.get_model_group_usage is deprecated"): + usage: Final = await router.get_model_group_usage("gpt") + assert usage == (None, None) + with pytest.warns(DeprecationWarning, match=r"Router\.get_remaining_model_group_usage is deprecated"): + remaining: Final = await router.get_remaining_model_group_usage("gpt") + assert remaining == {} + with pytest.warns(DeprecationWarning, match=r"Router\.async_get_healthy_deployments is deprecated"): + healthy = await router.async_get_healthy_deployments("gpt", {}) # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] # the old untyped dicts + assert isinstance(healthy, list) + assert [picked_id(entry) for entry in healthy] == ["a"] # pyright: ignore[reportUnknownArgumentType, reportUnknownVariableType] # see above + assert picked_id(_call(router, "get_available_deployment_for_pass_through", "pt")) == "p" + + +def test_write_methods_go_through_apply(make_core: MakeCore) -> None: + router: Final = as_router(make_core(model_list=[deployment("gpt", "a")])) + _call(router, "update_settings", allowed_fails=6) + assert router.settings().allowed_fails == 6 + router.apply(ChangeSet()) + assert _call(router, "delete_deployment", "a") is not None + assert router.registry().all() == [] + + +def test_the_priority_split_answers_through_health(make_core: MakeCore) -> None: + router: Final = as_router(make_core(model_list=[deployment("gpt", "a")], default_priority=2)) + assert _call(router, "split_request_priority", "gpt", {"priority": 3, "x": 1}) == (3, {"x": 1}) + assert _call(router, "split_request_priority", "gpt", {"x": 1}) == (2, {"x": 1}) + assert router.health().priority_header_fields("gpt", {}, "5") == {"priority": 5} + assert router.health().priority_header_fields("gpt", {}, None) == {} + with pytest.raises(litellm.BadRequestError, match="priority must be an integer, got str 'bad'"): + router.health().split_priority("gpt", {"priority": "bad"}) + + +# -- plugin context, request logging ------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_the_pre_routing_hook_answers_through_the_routing_context(make_core: MakeCore) -> None: + router: Final = as_router( + make_core(model_list=[deployment("small", "s"), deployment("big", "b"), _complexity_marker()]) + ) + request: Final[dict[str, object]] = {} + with pytest.warns(DeprecationWarning, match=r"Router\.async_pre_routing_hook is deprecated"): + response = await router.async_pre_routing_hook("smart", request, messages=MESSAGES) # pyright: ignore[reportUnknownMemberType] # the old signature's bare ``list`` + assert response is not None + assert response.model == "small" + decision: Final = as_str_mapping((as_str_mapping(request.get("metadata")) or {}).get("routing_decision")) or {} + assert decision.get("router_model_name") == "smart" + assert decision.get("routed_model") == "small" + plain: Final[dict[str, object]] = {} + with pytest.warns(DeprecationWarning, match="async_pre_routing_hook"): + hooked: Final = await router.async_pre_routing_hook("small", plain) # pyright: ignore[reportUnknownMemberType] # see above + assert hooked is None + assert plain == {} + + +def _complexity_marker() -> dict[str, object]: # mutable-ok: a model_list entry + return { + "model_name": "smart", + "litellm_params": { + "model": "auto_router/complexity_router", + "complexity_router_config": {"tiers": {"SIMPLE": "small", "MEDIUM": "big", "COMPLEX": "big"}}, + "complexity_router_default_model": "big", + }, + "model_info": {"id": "marker"}, + } + + +def test_log_retry_records_the_failure_on_the_request_metadata(make_core: MakeCore) -> None: + router: Final = as_router(make_core(model_list=[deployment("gpt", "a", base_model="bm")])) + kwargs: Final[dict[str, object]] = { + "model": "gpt", + "metadata": {"model_info": {"id": "a"}, "deployment": "openai/a", "attempted_retries": 1}, + } + assert _call(router, "log_retry", kwargs, ValueError("boom")) is kwargs + assert kwargs["metadata"] == { + "model_info": {"id": "a"}, + "deployment": "openai/a", + "attempted_retries": 1, + "previous_models": ( + { + "model_group": "gpt", + "deployment_id": "a", + "model": "openai/a", + "base_model": "bm", + "exception_type": "ValueError", + "exception_string": "boom", + "attempted_retries": 1, + }, + ), + "request_retry_count": 1, + } + _call(router, "log_retry", kwargs, KeyError("again")) + metadata: Final = as_str_mapping(kwargs["metadata"]) or {} + assert metadata["request_retry_count"] == 2 + previous: Final = metadata["previous_models"] + assert isinstance(previous, tuple) + assert len(previous) == 2 # pyright: ignore[reportUnknownArgumentType] # an untyped tuple of breadcrumbs + + +def test_arming_the_read_prefetch_without_a_request_batch_does_nothing(make_core: MakeCore) -> None: + router: Final = as_router(make_core(model_list=[deployment("gpt", "a")], routing_strategy="usage-based-routing-v2")) + assert _call(router, "arm_routing_read_prefetch", "gpt", {}) is None + assert _call(router, "arm_routing_read_prefetch", "unknown-model") is None + + +def test_binding_to_request_logging_binds_through_observe(make_core: MakeCore) -> None: + from litellm.litellm_core_utils.litellm_logging import Logging + + router: Final = as_router(make_core(model_list=[deployment("gpt", "a")])) + logging_obj: Final = Logging( + model="gpt", + messages=MESSAGES, + stream=False, + call_type="acompletion", + start_time=datetime.now(), + litellm_call_id="call-1", + function_id="fn-1", + ) + assert _call(router, "bind_to_request_logging", logging_obj) is None + with ( + pytest.raises(AttributeError, match="add_dynamic_function_callbacks"), + pytest.warns(DeprecationWarning, match=r"Router\.bind_to_request_logging is deprecated"), + ): + router.bind_to_request_logging(None) # pyright: ignore[reportArgumentType] + + +# -- writes with a routing effect ---------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_assigned_routing_groups_route_from_the_next_call(make_core: MakeCore) -> None: + from litellm.types.router import RoutingGroup + + router: Final = as_router(make_core(model_list=[deployment("gpt", "a"), deployment("claude", "c")])) + with pytest.raises(Exception, match="grp"): + await router.select("grp", {}) + write( + router, + "routing_groups", + {"grp": RoutingGroup(group_name="grp", models=["claude"], routing_strategy="simple-shuffle")}, + ) + assert picked_id(await router.select("grp", {})) == "c" + assert router.registry().routing_groups().keys() == {"grp"} + + +class _NarrowTo: + def __init__(self, keep: str) -> None: + self.keep: Final = keep + + async def run(self, context: RoutingContext) -> RoutingContext: + context.candidate_models = [m for m in context.candidate_models if m == self.keep] + return context + + +@pytest.mark.asyncio +async def test_a_plugin_added_to_an_empty_plugin_list_runs(make_core: MakeCore) -> None: + router: Final = as_router(make_core(model_list=[deployment("g", "a"), deployment("g", "b")])) + with pytest.warns(DeprecationWarning, match=r"Router\.routing_plugins is deprecated"): + plugins: Final = router.routing_plugins + plugins.append(_NarrowTo("openai/b")) + for _ in range(4): + assert picked_id(await router.select("g", {}, messages=MESSAGES)) == "b" # pyright: ignore[reportCallIssue] + plugins.clear() + assert {picked_id(await router.select("g", {}, messages=MESSAGES)) for _ in range(40)} == {"a", "b"} # pyright: ignore[reportCallIssue] + + +@pytest.mark.asyncio +async def test_an_assigned_plugin_list_runs(make_core: MakeCore) -> None: + router: Final = as_router(make_core(model_list=[deployment("g", "a"), deployment("g", "b")])) + write(router, "routing_plugins", [_NarrowTo("openai/a")]) + for _ in range(4): + assert picked_id(await router.select("g", {}, messages=MESSAGES)) == "a" # pyright: ignore[reportCallIssue] + + +def test_an_assigned_prompt_caching_pin_cache_is_the_one_flushed(make_core: MakeCore) -> None: + router: Final = as_router(make_core(model_list=[deployment("gpt", "a")])) + pins: Final = DualCache() + write(router, "prompt_caching_pin_cache", pins) + assert read(router, "prompt_caching_pin_cache") is pins + pins.in_memory_cache.set_cache("k", "v") # pyright: ignore[reportUnknownMemberType] # the in-memory cache is untyped + router.flush_cache() # pyright: ignore[reportAttributeAccessIssue] + assert pins.in_memory_cache.get_cache("k") is None # pyright: ignore[reportUnknownMemberType] # see above + + +# -- pod resolution and cooldown_time -------------------------------------------------------------- + + +class _TaggingDiscovery: + """Stands in for KubernetesPodDiscovery: stamps which discovery resolved the pick.""" + + def __init__(self, name: str) -> None: + self.name: Final = name + + def resolve_deployment(self, deployment: object, request_kwargs: object) -> object: + return {**(as_str_mapping(deployment) or {}), "resolved_by": self.name} + + async def async_resolve_deployment(self, deployment: object, request_kwargs: object) -> object: + return self.resolve_deployment(deployment, request_kwargs) + + +@pytest.mark.asyncio +async def test_a_custom_strategy_resolves_pods_through_the_discovery_bound_when_registered( + make_core: MakeCore, +) -> None: + from litellm.types.router import CustomRoutingStrategyBase + + class Pick(CustomRoutingStrategyBase): + async def async_get_available_deployment( # pyright: ignore[reportIncompatibleMethodOverride] + self, + model: str, + messages: object = None, + input: object = None, + specific_deployment: object = False, + request_kwargs: object = None, + ) -> object: + return deployment("gpt", "b") + + def get_available_deployment( # pyright: ignore[reportIncompatibleMethodOverride] + self, + model: str, + messages: object = None, + input: object = None, + specific_deployment: object = False, + request_kwargs: object = None, + ) -> object: + return deployment("gpt", "b") + + router: Final = as_router(make_core(model_list=[deployment("gpt", "a"), deployment("gpt", "b")])) + router.apply(ChangeSet(assign={"kubernetes_pod_discovery": _TaggingDiscovery("first")})) + assert (as_str_mapping(await router.select("gpt", {})) or {}).get("resolved_by") == "first" + router.register_plugin(Pick()) + router.apply(ChangeSet(assign={"kubernetes_pod_discovery": _TaggingDiscovery("second")})) + picked: Final = as_str_mapping(await router.select("gpt", {})) or {} + assert picked_id(picked) == "b" + assert picked.get("resolved_by") == "first" + + +def test_cooldown_time_reads_back_as_given(make_core: MakeCore) -> None: + router: Final = as_router(make_core(model_list=[deployment("gpt", "a")], cooldown_time=7)) + value: Final = read(router, "cooldown_time") + assert value == 7 and type(value) is int + assert router.settings().as_dict()["cooldown_time"] == 7 + assert type(router.settings().as_dict()["cooldown_time"]) is int + write(router, "cooldown_time", 2.5) + assert router.settings().cooldown_time == 2.5 + _call(router, "update_settings", cooldown_time="3") + updated: Final = router.settings().cooldown_time + assert updated == 3 and type(updated) is int + default: Final = as_router(make_core(model_list=[deployment("gpt", "a")])).settings().cooldown_time + assert default == 5 and type(default) is int diff --git a/tests/unit/test_router/core/test_endpoints.py b/tests/unit/test_router/core/test_endpoints.py index c0c4c8cf6dd..ce56ae403ba 100644 --- a/tests/unit/test_router/core/test_endpoints.py +++ b/tests/unit/test_router/core/test_endpoints.py @@ -22,6 +22,7 @@ from tests.unit.test_router.core.scripted_host import ( completion_response, deployment, embedding_response, + records_outcomes_through_callbacks, ) MESSAGES: Final = [{"role": "user", "content": "hi"}] @@ -45,11 +46,16 @@ ENDPOINTS: Final[dict[str, tuple[str, EndpointCall]]] = { "anthropic_messages", lambda core: core.aanthropic_messages(model="g", messages=MESSAGES, max_tokens=5), ), + # The proxy's /v1/messages route calls the router by this name (route_request). + "anthropic_messages": ( + "anthropic_messages", + lambda core: core.anthropic_messages(model="g", messages=MESSAGES, max_tokens=5), + ), "call-acompletion": ("acompletion", lambda core: core.call("acompletion", {"model": "g", "messages": MESSAGES})), "call-completion": ("completion", lambda core: core.call("completion", {"model": "g", "messages": MESSAGES})), } SYNC_ENDPOINTS: Final = frozenset({"completion", "embedding", "image_generation"}) -LITELLM_METADATA_ENDPOINTS: Final = frozenset({"aresponses", "aanthropic_messages"}) +LITELLM_METADATA_ENDPOINTS: Final = frozenset({"aresponses", "aanthropic_messages", "anthropic_messages"}) EMBEDDING_ENDPOINTS: Final = frozenset({"aembedding", "embedding"}) @@ -276,14 +282,15 @@ async def test_request_level_api_key_replaces_the_deployment_key_on_the_host_cal async def test_construction_registers_one_success_and_one_failure_callback_per_mode( make_core: MakeCore, host: ScriptedHost ): - _one_group(make_core) + core: Final = _one_group(make_core) registered: Final = {name: len(callbacks) for name, callbacks in host.lists.items()} + one: Final = 1 if records_outcomes_through_callbacks(core) else 0 assert registered == { "callbacks": 0, - "success_callback": 1, - "_async_success_callback": 1, - "failure_callback": 1, - "_async_failure_callback": 1, + "success_callback": one, + "_async_success_callback": one, + "failure_callback": one, + "_async_failure_callback": one, "input_callback": 0, } assert all(callable(callbacks[0]) for callbacks in host.lists.values() if callbacks) @@ -416,7 +423,6 @@ def _scripted_stream(*texts: str) -> Callable[..., CustomStreamWrapper]: yield ModelResponseStream(choices=[{"index": 0, "delta": {"role": "assistant", "content": text}}]) # A stream that ends without a finish reason fails since #47665. yield ModelResponseStream(choices=[{"index": 0, "delta": {}, "finish_reason": "stop"}]) - yield ModelResponseStream(choices=[{"index": 0, "delta": {}, "finish_reason": "stop"}]) logging_obj: Final = LiteLLMLogging( model="d1", diff --git a/tests/unit/test_router/core/test_endpoints_characterization_batch.py b/tests/unit/test_router/core/test_endpoints_characterization_batch.py new file mode 100644 index 00000000000..6ba298accfe --- /dev/null +++ b/tests/unit/test_router/core/test_endpoints_characterization_batch.py @@ -0,0 +1,525 @@ +""" +Characterization of the fan-out endpoints and request shapers: the files API (one file per +deployment), batches (create, retrieve across deployments, cancel, list), batch completions, +fastest-response completion, the scheduler (priority) and prompt-management model groups. +""" + +import re +import datetime +import io +import json +from collections.abc import Callable, Mapping +from types import SimpleNamespace +from typing import Final + +import pytest + +import litellm +from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging +from litellm.router_core import RouterCore +from litellm.types.utils import LiteLLMBatch, ModelResponse +from tests.unit.test_router.core.endpoints_helpers import EndpointHost +from tests.unit.test_router.core.scripted_host import ScriptedHost, bad_request, deployment + +MESSAGES: Final = [{"role": "user", "content": "hi"}] +MakeCore = Callable[..., RouterCore] +TWO_DEPLOYMENTS: Final = [deployment("g", "d1"), deployment("g", "d2")] + + +@pytest.fixture +def host() -> EndpointHost: + return EndpointHost() + + +def _bucket(kwargs: Mapping[str, object], name: str) -> Mapping[str, object]: + bucket: Final = kwargs[name] + assert isinstance(bucket, Mapping) + return bucket + + +def _batch(batch_id: str = "b1") -> LiteLLMBatch: + return LiteLLMBatch( + id=batch_id, + completion_window="24h", + created_at=1, + endpoint="/v1/chat/completions", + input_file_id="f", + object="batch", + status="completed", + ) + + +def _jsonl() -> io.BytesIO: + return io.BytesIO(b'{"custom_id": "1", "body": {"model": "g", "messages": []}}\n') + + +# --- files ---------------------------------------------------------------------------------- + + +async def test_acreate_file_creates_the_file_on_every_deployment_of_the_group(make_core: MakeCore, host: EndpointHost): + core: Final = make_core(model_list=TWO_DEPLOYMENTS) + await core.acreate_file(model="g", file=_jsonl(), purpose="batch") + assert sorted(host.deployments_called("acreate_file")) == ["d1", "d2"] + + +async def test_acreate_file_returns_the_first_response_with_the_deployment_file_id_mapping( + make_core: MakeCore, host: EndpointHost +): + first: Final = ModelResponse(id="file-1") + second: Final = ModelResponse(id="file-2") + host.answer("openai/d1", first) + host.answer("openai/d2", second) + core: Final = make_core(model_list=TWO_DEPLOYMENTS) + response: Final = await core.acreate_file(model="g", file=_jsonl(), purpose="batch") + assert response in (first, second) + assert response._hidden_params["model_file_id_mapping"] == {"d1": "file-1", "d2": "file-2"} + assert "model_id" not in response._hidden_params + + +async def test_acreate_file_rewrites_the_jsonl_model_to_each_deployment_model(make_core: MakeCore, host: EndpointHost): + core: Final = make_core(model_list=TWO_DEPLOYMENTS) + await core.acreate_file(model="g", file=_jsonl(), purpose="batch") + models: Final = {} + for call in host.calls: + sent_file: Final = call.kwargs["file"] + assert isinstance(sent_file, io.BytesIO) + models[call.deployment_id] = json.loads(sent_file.getvalue().decode().strip())["body"]["model"] + assert models == {"d1": "d1", "d2": "d2"} + + +async def test_acreate_file_keeps_a_non_batch_file_as_sent(make_core: MakeCore, host: EndpointHost): + upload: Final = _jsonl() + core: Final = make_core(model_list=[deployment("g", "d1")]) + await core.acreate_file(model="g", file=upload, purpose="assistants") + assert host.calls[0].kwargs["file"] is not None + sent: Final = host.calls[0].kwargs["file"] + assert isinstance(sent, io.BytesIO) + assert json.loads(sent.getvalue().decode().strip())["body"]["model"] == "g" + + +async def test_acreate_file_puts_the_model_group_in_metadata_and_the_deployment_in_litellm_metadata( + make_core: MakeCore, host: EndpointHost +): + core: Final = make_core(model_list=[deployment("g", "d1")]) + await core.acreate_file(model="g", file=_jsonl(), purpose="batch") + sent: Final = host.calls[0].kwargs + assert _bucket(sent, "metadata")["model_group"] == "g" + assert "deployment" not in _bucket(sent, "metadata") + assert _bucket(sent, "litellm_metadata")["deployment"] == "openai/d1" + assert "model_group" not in _bucket(sent, "litellm_metadata") + assert sent["custom_llm_provider"] == "openai" + + +async def test_acreate_file_without_a_purpose_raises_before_calling_the_host(make_core: MakeCore, host: EndpointHost): + core: Final = make_core(model_list=[deployment("g", "d1")], num_retries=0) + with pytest.raises(Exception, match="file and file_purpose are required for create_file"): + await core.acreate_file(model="g", file=_jsonl()) + assert host.calls == [] + + +async def test_acreate_file_sends_target_model_names_to_a_litellm_proxy_deployment( + make_core: MakeCore, host: EndpointHost +): + core: Final = make_core(model_list=[deployment("g", "d1", model="litellm_proxy/remote-model")]) + await core.acreate_file(model="g", file=_jsonl(), purpose="batch", extra_body={"keep": 1}) + assert dict(host.calls[0].kwargs["extra_body"]) == {"keep": 1, "target_model_names": "remote-model"} # pyright: ignore[reportArgumentType] # a mapping + + +async def test_acreate_file_sends_the_deployment_gcs_bucket_in_litellm_metadata( + make_core: MakeCore, host: EndpointHost +): + core: Final = make_core(model_list=[deployment("g", "d1", gcs_bucket_name="bucket-1")]) + await core.acreate_file(model="g", file=_jsonl(), purpose="batch") + assert _bucket(host.calls[0].kwargs, "litellm_metadata")["gcs_bucket_name"] == "bucket-1" + + +# --- batches -------------------------------------------------------------------------------- + + +async def test_acreate_batch_routes_with_litellm_metadata_only(make_core: MakeCore, host: EndpointHost): + core: Final = make_core(model_list=[deployment("g", "d1")]) + response: Final = await core.acreate_batch( + model="g", input_file_id="f", endpoint="/v1/chat/completions", completion_window="24h" + ) + call: Final = host.calls[0] + assert (call.op, call.model, call.deployment_id, call.kwargs["custom_llm_provider"]) == ( + "acreate_batch", + "openai/d1", + "d1", + "openai", + ) + assert "metadata" not in call.kwargs + assert _bucket(call.kwargs, "litellm_metadata")["model_group"] == "g" + assert response._hidden_params["model_id"] == "d1" + + +async def test_acreate_batch_is_retried(make_core: MakeCore, host: EndpointHost): + host.answer("openai/d1", litellm.InternalServerError(message="boom", llm_provider="openai", model="d1"), _batch()) + core: Final = make_core(model_list=[deployment("g", "d1")], num_retries=1) + await core.acreate_batch(model="g", input_file_id="f", endpoint="/v1/chat/completions", completion_window="24h") + assert len(host.calls_to("acreate_batch")) == 2 + + +async def test_acreate_batch_never_fails_over_to_another_deployment(make_core: MakeCore, host: EndpointHost): + # The batch belongs to the deployment that created it, so its attempt stamps no failed deployment. + host.answer("*", litellm.InternalServerError(message="boom", llm_provider="openai", model="d")) + core: Final = make_core(model_list=TWO_DEPLOYMENTS, num_retries=0, enable_weighted_failover=True) + with pytest.raises(litellm.InternalServerError): + await core.acreate_batch(model="g", input_file_id="f", endpoint="/v1/chat/completions", completion_window="24h") + assert len(host.calls_to("acreate_batch")) == 1 + + +async def test_aretrieve_batch_asks_every_deployment_and_returns_the_batch_found( + make_core: MakeCore, host: EndpointHost +): + found: Final = _batch() + host.answer("openai/d1", bad_request("d1", "no such batch")) + host.answer("openai/d2", found) + core: Final = make_core(model_list=TWO_DEPLOYMENTS) + assert await core.aretrieve_batch(model="g", batch_id="b1") is found + assert sorted(host.deployments_called("aretrieve_batch")) == ["d1", "d2"] + for call in host.calls: + assert (call.kwargs["batch_id"], call.kwargs["custom_llm_provider"], call.model_group) == ("b1", "openai", "g") + assert "metadata" not in call.kwargs + + +async def test_aretrieve_batch_raises_the_first_error_when_no_deployment_has_the_batch( + make_core: MakeCore, host: EndpointHost +): + host.answer("openai/d1", bad_request("d1", "first missing")) + host.answer("openai/d2", bad_request("d2", "second missing")) + core: Final = make_core(model_list=TWO_DEPLOYMENTS) + with pytest.raises(litellm.BadRequestError, match="missing"): + await core.aretrieve_batch(model="g", batch_id="b1") + assert len(host.calls) == 2 + + +async def test_aretrieve_batch_ignores_a_non_batch_answer_and_raises_a_generic_error( + make_core: MakeCore, host: EndpointHost +): + core: Final = make_core(model_list=[deployment("g", "d1")]) + with pytest.raises(Exception, match="Unable to find batch in any model"): + await core.aretrieve_batch(model="g", batch_id="b1") + + +async def test_aretrieve_batch_without_a_model_asks_every_deployment_with_its_own_group( + make_core: MakeCore, host: EndpointHost +): + host.answer("*", bad_request(message="missing")) + core: Final = make_core(model_list=[*TWO_DEPLOYMENTS, deployment("h", "d3")]) + with pytest.raises(litellm.BadRequestError): + await core.aretrieve_batch(batch_id="b1") + assert sorted((call.deployment_id, call.model_group) for call in host.calls) == [ + ("d1", "g"), + ("d2", "g"), + ("d3", "h"), + ] + + +async def test_aretrieve_batch_is_not_retried(make_core: MakeCore, host: EndpointHost): + host.answer("openai/d1", litellm.InternalServerError(message="boom", llm_provider="openai", model="d1")) + core: Final = make_core(model_list=[deployment("g", "d1")], num_retries=3) + with pytest.raises(litellm.InternalServerError): + await core.aretrieve_batch(model="g", batch_id="b1") + assert len(host.calls) == 1 + + +async def test_acancel_batch_routes_with_litellm_metadata(make_core: MakeCore, host: EndpointHost): + host.answer("openai/d1", _batch()) + core: Final = make_core(model_list=[deployment("g", "d1", api_base="https://d1.example")]) + response: Final = await core.acancel_batch(model="g", batch_id="b1") + call: Final = host.calls[0] + assert (call.op, call.deployment_id, call.kwargs["batch_id"], call.kwargs["api_base"]) == ( + "acancel_batch", + "d1", + "b1", + "https://d1.example", + ) + assert "metadata" not in call.kwargs + assert response._hidden_params["model_id"] == "d1" + + +async def test_alist_batches_merges_the_lists_of_every_deployment(make_core: MakeCore, host: EndpointHost): + host.answer("openai/d1", SimpleNamespace(first_id="a", last_id="b", data=[1, 2], has_more=False)) + host.answer("openai/d2", SimpleNamespace(first_id="c", last_id="d", data=[3], has_more=True)) + core: Final = make_core(model_list=TWO_DEPLOYMENTS) + merged: Final = await core.alist_batches(model="g", limit=2) + assert sorted(merged["data"]) == [1, 2, 3] + assert (merged["object"], merged["has_more"]) == ("list", True) + assert merged["first_id"] in ("a", "c") + assert merged["last_id"] in ("b", "d") + + +async def test_alist_batches_sends_the_deployment_params_and_request_without_routing_metadata( + make_core: MakeCore, host: EndpointHost +): + host.answer("*", SimpleNamespace(first_id=None, last_id=None, data=[], has_more=False)) + core: Final = make_core(model_list=[deployment("g", "d1", api_base="https://d1.example")]) + await core.alist_batches(model="g", limit=2, api_base="https://request.example") + call: Final = host.calls[0] + assert (call.model, call.kwargs["limit"], call.kwargs["api_base"], call.deployment_id) == ( + "openai/d1", + 2, + "https://request.example", + None, + ) + assert "metadata" not in call.kwargs and "litellm_metadata" not in call.kwargs + + +async def test_alist_batches_drops_failing_deployments(make_core: MakeCore, host: EndpointHost): + host.answer("openai/d1", bad_request("d1")) + host.answer("openai/d2", SimpleNamespace(first_id="c", last_id="d", data=[3], has_more=False)) + core: Final = make_core(model_list=TWO_DEPLOYMENTS) + merged: Final = await core.alist_batches(model="g") + assert merged == {"object": "list", "data": [3], "first_id": "c", "last_id": "d", "has_more": False} + + +async def test_alist_batches_for_an_unknown_group_returns_an_empty_list(make_core: MakeCore, host: EndpointHost): + core: Final = make_core(model_list=TWO_DEPLOYMENTS) + assert await core.alist_batches(model="nope") == { + "object": "list", + "data": [], + "first_id": None, + "last_id": None, + "has_more": False, + } + assert host.calls == [] + + +# --- batch completions ------------------------------------------------------------------------ + + +async def test_abatch_completion_sends_one_request_to_each_model_and_returns_errors_in_place( + make_core: MakeCore, host: EndpointHost +): + error: Final = bad_request("d1", "a failed") + host.answer("openai/d1", error) + core: Final = make_core(model_list=[deployment("a", "d1"), deployment("b", "d2")], num_retries=0) + results: Final = await core.abatch_completion(models=["a", "b"], messages=MESSAGES) + assert isinstance(results[0], litellm.BadRequestError) + assert isinstance(results[1], ModelResponse) and results[1]._hidden_params["model_id"] == "d2" + + +async def test_abatch_completion_with_message_lists_groups_results_per_request(make_core: MakeCore, host: EndpointHost): + core: Final = make_core(model_list=[deployment("a", "d1"), deployment("b", "d2")]) + second: Final = [{"role": "user", "content": "second"}] + results: Final = await core.abatch_completion(models=["a", "b"], messages=[MESSAGES, second]) + assert [[response._hidden_params["model_id"] for response in row] for row in results] == [ + ["d1", "d2"], + ["d1", "d2"], + ] + assert sorted(str(call.kwargs["messages"]) for call in host.calls) == sorted( + [str(MESSAGES), str(MESSAGES), str(second), str(second)] + ) + + +async def test_abatch_completion_forwards_extra_kwargs_to_every_request(make_core: MakeCore, host: EndpointHost): + core: Final = make_core(model_list=[deployment("a", "d1"), deployment("b", "d2")]) + await core.abatch_completion(models=["a", "b"], messages=MESSAGES, max_tokens=15) + assert [call.kwargs["max_tokens"] for call in host.calls] == [15, 15] + + +async def test_abatch_completion_one_model_multiple_requests_returns_one_result_per_request( + make_core: MakeCore, host: EndpointHost +): + host.answer("openai/d1", ModelResponse(id="r1"), bad_request("d1"), ModelResponse(id="r3")) + core: Final = make_core(model_list=[deployment("a", "d1")], num_retries=0) + results: Final = await core.abatch_completion_one_model_multiple_requests(model="a", messages=[MESSAGES] * 3) + assert len(results) == 3 + assert sum(isinstance(result, litellm.BadRequestError) for result in results) == 1 + assert len(host.calls) == 3 + + +async def test_fastest_response_returns_the_first_success_marked_in_hidden_params( + make_core: MakeCore, host: EndpointHost +): + host.answer("openai/d1", bad_request("d1")) + core: Final = make_core(model_list=[deployment("a", "d1"), deployment("b", "d2")], num_retries=0) + response: Final = await core.abatch_completion_fastest_response(model="a, b", messages=MESSAGES) + assert response._hidden_params["model_id"] == "d2" + assert response._hidden_params["fastest_response_batch_completion"] is True + + +async def test_fastest_response_raises_when_every_model_fails(make_core: MakeCore, host: EndpointHost): + host.answer("*", bad_request()) + core: Final = make_core(model_list=[deployment("a", "d1"), deployment("b", "d2")], num_retries=0) + with pytest.raises(Exception, match="All tasks failed"): + await core.abatch_completion_fastest_response(model="a,b", messages=MESSAGES) + + +async def test_fastest_response_with_an_unknown_model_counts_it_as_failed(make_core: MakeCore, host: EndpointHost): + core: Final = make_core(model_list=[deployment("a", "d1")], num_retries=0) + response: Final = await core.abatch_completion_fastest_response(model="nope, a", messages=MESSAGES) + assert response._hidden_params["model_id"] == "d1" + + +# --- scheduler / priority ------------------------------------------------------------------- + + +async def test_schedule_acompletion_marks_the_response_as_prioritized_and_does_not_send_priority( + make_core: MakeCore, host: EndpointHost +): + core: Final = make_core(model_list=[deployment("g", "d1")]) + response: Final = await core.schedule_acompletion(model="g", messages=MESSAGES, priority=0) + assert response._hidden_params["additional_headers"]["x-litellm-request-prioritization-used"] is True + assert "priority" not in host.calls[0].kwargs + + +async def test_schedule_acompletion_without_a_priority_or_default_raises(make_core: MakeCore, host: EndpointHost): + core: Final = make_core(model_list=[deployment("g", "d1")]) + with pytest.raises( + litellm.BadRequestError, + match=re.escape("priority is required to queue a request, send one or set router_settings.default_priority"), + ): + await core.schedule_acompletion(model="g", messages=MESSAGES) + assert host.calls == [] + + +async def test_schedule_acompletion_without_a_priority_uses_the_default_priority( + make_core: MakeCore, host: EndpointHost +): + core: Final = make_core(model_list=[deployment("g", "d1")], default_priority=5) + response: Final = await core.schedule_acompletion(model="g", messages=MESSAGES) + assert response._hidden_params["additional_headers"]["x-litellm-request-prioritization-used"] is True + + +@pytest.mark.parametrize("endpoint", ["schedule_acompletion", "acompletion"]) +async def test_a_non_integer_priority_raises_bad_request(make_core: MakeCore, host: EndpointHost, endpoint: str): + core: Final = make_core(model_list=[deployment("g", "d1")]) + with pytest.raises(litellm.BadRequestError, match="priority must be an integer, got str 'high'"): + await core.call(endpoint, {"model": "g", "messages": MESSAGES, "priority": "high"}) + assert host.calls == [] + + +async def test_a_non_integer_priority_with_drop_params_is_served_unscheduled(make_core: MakeCore, host: EndpointHost): + core: Final = make_core(model_list=[deployment("g", "d1")]) + response: Final = await core.acompletion(model="g", messages=MESSAGES, priority="high", drop_params=True) + assert "x-litellm-request-prioritization-used" not in response._hidden_params["additional_headers"] + assert "priority" not in host.calls[0].kwargs + + +async def test_a_non_integer_priority_with_drop_params_queues_at_the_default_priority( + make_core: MakeCore, host: EndpointHost +): + core: Final = make_core(model_list=[deployment("g", "d1")], default_priority=5) + response: Final = await core.acompletion(model="g", messages=MESSAGES, priority="high", drop_params=True) + assert response._hidden_params["additional_headers"]["x-litellm-request-prioritization-used"] is True + + +@pytest.mark.parametrize( + ("endpoint", "request_body"), + [("acompletion", {"messages": MESSAGES}), ("atext_completion", {"prompt": "p"})], +) +async def test_a_request_priority_schedules_the_call( + make_core: MakeCore, host: EndpointHost, endpoint: str, request_body: dict[str, object] +): + core: Final = make_core(model_list=[deployment("g", "d1")]) + response: Final = await core.call(endpoint, {"model": "g", "priority": 3, **request_body}) + assert response._hidden_params["additional_headers"]["x-litellm-request-prioritization-used"] is True # pyright: ignore[reportAttributeAccessIssue] # the scripted response + assert "priority" not in host.calls[0].kwargs + + +async def test_default_priority_schedules_every_acompletion(make_core: MakeCore, host: EndpointHost): + core: Final = make_core(model_list=[deployment("g", "d1")], default_priority=5) + response: Final = await core.acompletion(model="g", messages=MESSAGES) + assert response._hidden_params["additional_headers"]["x-litellm-request-prioritization-used"] is True + + +async def test_a_scheduled_call_error_carries_its_priority(make_core: MakeCore, host: EndpointHost): + host.answer("openai/d1", bad_request("d1")) + core: Final = make_core(model_list=[deployment("g", "d1")], num_retries=0) + with pytest.raises(litellm.BadRequestError) as raised: + await core.acompletion(model="g", messages=MESSAGES, priority=4) + assert raised.value.priority == 4 # pyright: ignore[reportAttributeAccessIssue] # stamped by the scheduler + + +@pytest.mark.parametrize(("header", "expected"), [("7", {"priority": 7}), (None, {}), ("", {})]) +async def test_priority_header_update_parses_the_priority_header( + make_core: MakeCore, host: EndpointHost, header: str | None, expected: dict[str, int] +): + core: Final = make_core(model_list=[deployment("g", "d1")]) + assert dict(core.priority_header_update("g", {}, header)) == expected + + +# --- prompt management models ----------------------------------------------------------------- + + +def _logging_obj() -> LiteLLMLogging: + return LiteLLMLogging( + model="pm", + messages=MESSAGES, + stream=False, + call_type="acompletion", + start_time=datetime.datetime(2026, 1, 1, tzinfo=datetime.timezone.utc), + litellm_call_id="c1", + function_id="f1", + ) + + +async def test_prompt_management_group_without_a_prompt_id_raises_value_error(make_core: MakeCore, host: EndpointHost): + core: Final = make_core(model_list=[deployment("pm", "p1", model="langfuse/gpt-4o")]) + with pytest.raises(ValueError, match=re.escape("Prompt ID is not set or not a string. Got=None")): + await core.acompletion(model="pm", messages=MESSAGES, litellm_call_id="c1") + assert host.calls == [] + + +async def test_prompt_management_group_with_non_dict_prompt_variables_raises_value_error( + make_core: MakeCore, host: EndpointHost +): + core: Final = make_core( + model_list=[deployment("pm", "p1", model="langfuse/gpt-4o", prompt_id="p", prompt_variables="x")] + ) + with pytest.raises(ValueError, match=re.escape("Prompt variables is set but not a dictionary. Got=x")): + await core.acompletion(model="pm", messages=MESSAGES, litellm_call_id="c1") + + +async def test_prompt_management_group_without_a_call_id_or_logging_object_raises_key_error( + make_core: MakeCore, host: EndpointHost +): + core: Final = make_core(model_list=[deployment("pm", "p1", model="langfuse/gpt-4o")]) + with pytest.raises(KeyError, match="litellm_call_id"): + await core.acompletion(model="pm", messages=MESSAGES) + + +async def test_litellm_agent_group_without_a_logging_object_raises_key_error(make_core: MakeCore, host: EndpointHost): + """function_setup consumes original_function, so the direct-call path cannot pop it.""" + core: Final = make_core(model_list=[deployment("pm", "p1", model="litellm_agent/openai/gpt-4o")]) + with pytest.raises(KeyError, match="original_function"): + await core.acompletion(model="pm", messages=MESSAGES, litellm_call_id="c1") + assert host.calls == [] + + +async def test_litellm_agent_group_calls_the_stripped_model_directly_with_the_deployment_params( + make_core: MakeCore, host: EndpointHost +): + core: Final = make_core( + model_list=[deployment("pm", "p1", model="litellm_agent/openai/gpt-4o", timeout=9, temperature=0.2)] + ) + response: Final = await core.acompletion(model="pm", messages=MESSAGES, litellm_logging_obj=_logging_obj()) + call: Final = host.calls[0] + assert (call.op, call.model, call.deployment_id, call.model_group) == ("acompletion", "openai/gpt-4o", "p1", "pm") + assert (call.kwargs["timeout"], call.kwargs["temperature"], "original_function" in call.kwargs) == (9.0, 0.2, False) + assert response._hidden_params == {} + + +async def test_litellm_agent_group_naming_another_group_routes_into_it_with_the_agent_timeout( + make_core: MakeCore, host: EndpointHost +): + core: Final = make_core( + model_list=[deployment("pm", "p1", model="litellm_agent/g", timeout=9), deployment("g", "d1")] + ) + response: Final = await core.acompletion(model="pm", messages=MESSAGES, litellm_logging_obj=_logging_obj()) + call: Final = host.calls[0] + assert (call.model, call.model_group, call.kwargs["timeout"]) == ("openai/d1", "pm", 9.0) + # The agent deployment's api_key travels as a request credential, so the target gets a request-scoped id. + assert call.deployment_id != "d1" + assert response._hidden_params["additional_headers"]["x-litellm-model-group"] == "g" + + +async def test_a_group_with_two_prompt_management_deployments_routes_normally(make_core: MakeCore, host: ScriptedHost): + core: Final = make_core( + model_list=[deployment("pm", "p1", model="langfuse/gpt-4o"), deployment("pm", "p2", model="langfuse/gpt-4o")] + ) + await core.acompletion(model="pm", messages=MESSAGES) + call: Final = host.calls[0] + assert (call.op, call.model, call.model_group) == ("acompletion", "langfuse/gpt-4o", "pm") diff --git a/tests/unit/test_router/core/test_endpoints_characterization_generic.py b/tests/unit/test_router/core/test_endpoints_characterization_generic.py new file mode 100644 index 00000000000..b1a356cf8eb --- /dev/null +++ b/tests/unit/test_router/core/test_endpoints_characterization_generic.py @@ -0,0 +1,850 @@ +""" +Characterization of the generic endpoints (Responses API and its get/delete/cancel/compact/input +items, Google GenAI, OCR, search, decisions, video, containers, skills, interactions, managed +agents, realtime, fine-tuning, files, image edit, LLM pass-through, vector stores, assistants, +guardrails): which host op each reaches, the kwargs the core builds, and how an id in the request +pins the deployment. + +The Python core binds litellm's own functions for some of these endpoints instead of going through +``host.invoke``; the ``host`` fixture replaces those functions with recording stand-ins (see +``endpoints_helpers.py``), so every call shows up in ``host.calls``. +""" + +import re +import inspect +import uuid +from collections.abc import Callable, Mapping +from datetime import datetime +from typing import Final + +import pytest + +import litellm +from litellm.constants import REALTIME_SESSION_ENDED_ABNORMALLY_KEY, REALTIME_SESSION_SUCCESS_LOGGED_KEY +from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging +from litellm.responses.utils import ResponsesAPIRequestUtils +from litellm.router_core import RouterCore +from litellm.types.videos.utils import encode_video_id_with_provider +from tests.unit.test_router.core.endpoints_helpers import EndpointHost, patch_direct_endpoints +from tests.unit.test_router.core.scripted_host import ScriptedHost, deployment + +MESSAGES: Final = [{"role": "user", "content": "hi"}] +MakeCore = Callable[..., RouterCore] + + +@pytest.fixture +def host(monkeypatch: pytest.MonkeyPatch) -> EndpointHost: + endpoint_host: Final = EndpointHost() + patch_direct_endpoints(monkeypatch, endpoint_host) + return endpoint_host + + +async def _call(core: RouterCore, endpoint: str, request: Mapping[str, object]) -> object: + result: Final = getattr(core, endpoint)(**request) + return await result if inspect.isawaitable(result) else result + + +def _bucket(kwargs: Mapping[str, object], name: str) -> Mapping[str, object]: + bucket: Final = kwargs[name] + assert isinstance(bucket, Mapping) + return bucket + + +# endpoint -> (recorded op, request without model) +ROUTED_ASYNC: Final[dict[str, tuple[str, dict[str, object]]]] = { + "aresponses": ("aresponses", {"input": "hi"}), + "agenerate_content": ("agenerate_content", {"contents": []}), + "agenerate_content_stream": ("agenerate_content_stream", {"contents": []}), + "afile_delete": ("afile_delete", {"file_id": "f1"}), + "afile_content": ("afile_content", {"file_id": "f1"}), + "afile_list": ("afile_list", {}), + "aget_responses": ("aget_responses", {"response_id": "resp_plain"}), + "acancel_responses": ("acancel_responses", {"response_id": "resp_plain"}), + "acompact_responses": ("acompact_responses", {"response_id": "resp_plain"}), + "adelete_responses": ("adelete_responses", {"response_id": "resp_plain"}), + "alist_input_items": ("alist_input_items", {"response_id": "resp_plain"}), + "acreate_realtime_client_secret": ("acreate_realtime_client_secret", {}), + "arealtime_calls": ("arealtime_calls", {}), + "acreate_realtime_transcription_session": ("acreate_realtime_transcription_session", {}), + "acreate_realtime_translation_client_secret": ("acreate_realtime_translation_client_secret", {}), + "arealtime_translation_calls": ("arealtime_translation_calls", {}), + "acreate_fine_tuning_job": ("acreate_fine_tuning_job", {"training_file": "f"}), + "acancel_fine_tuning_job": ("acancel_fine_tuning_job", {"fine_tuning_job_id": "j"}), + "alist_fine_tuning_jobs": ("alist_fine_tuning_jobs", {}), + "aretrieve_fine_tuning_job": ("aretrieve_fine_tuning_job", {"fine_tuning_job_id": "j"}), + "aimage_edit": ("aimage_edit", {"image": b"x", "prompt": "p"}), + "allm_passthrough_route": ("allm_passthrough_route", {"method": "POST", "endpoint": "/v1/x"}), + "aocr": ("aocr", {"document": {}}), + "adecisions": ("adecisions", {}), + "asystemone": ("asystemone", {}), + "aextraction": ("aextraction", {}), + "avideo_generation": ("avideo_generation", {"prompt": "p"}), + "avideo_list": ("avideo_list", {}), + "avideo_status": ("avideo_status", {"video_id": "v"}), + "avideo_content": ("avideo_content", {"video_id": "v"}), + "avideo_remix": ("avideo_remix", {"video_id": "v", "prompt": "p"}), + "avideo_create_character": ("avideo_create_character", {}), + "avideo_get_character": ("avideo_get_character", {}), + "avideo_edit": ("avideo_edit", {"video_id": "v"}), + "avideo_extension": ("avideo_extension", {"video_id": "v"}), + "acreate_container": ("acreate_container", {"name": "c"}), + "alist_containers": ("alist_containers", {}), + "aretrieve_container": ("aretrieve_container", {"container_id": "cntr_native"}), + "adelete_container": ("adelete_container", {"container_id": "cntr_native"}), + "aupload_container_file": ("aupload_container_file", {"container_id": "cntr_native", "file": b"x"}), + "alist_container_files": ("alist_container_files", {"container_id": "cntr_native"}), + "aretrieve_container_file": ("aretrieve_container_file", {"container_id": "cntr_native", "file_id": "f"}), + "adelete_container_file": ("adelete_container_file", {"container_id": "cntr_native", "file_id": "f"}), + "aretrieve_container_file_content": ( + "aretrieve_container_file_content", + {"container_id": "cntr_native", "file_id": "f"}, + ), + "acreate_skill": ("acreate_skill", {}), + "alist_skills": ("alist_skills", {}), + "aget_skill": ("aget_skill", {"skill_id": "s"}), + "adelete_skill": ("adelete_skill", {"skill_id": "s"}), + "acreate_interaction": ("interactions.acreate", {"input": "x"}), + "avector_store_search": ("vs.asearch", {"vector_store_id": "v", "query": "q"}), + "avector_store_retrieve": ("vs.aretrieve", {"vector_store_id": "v"}), + "avector_store_list": ("vs.alist", {}), + "avector_store_update": ("vs.aupdate", {"vector_store_id": "v"}), + "avector_store_delete": ("vs.adelete", {"vector_store_id": "v"}), + "avector_store_file_create": ("vsf.acreate", {"vector_store_id": "v", "file_id": "f"}), + "avector_store_file_list": ("vsf.alist", {"vector_store_id": "v"}), + "avector_store_file_retrieve": ("vsf.aretrieve", {"vector_store_id": "v", "file_id": "f"}), + "avector_store_file_content": ("vsf.aretrieve_content", {"vector_store_id": "v", "file_id": "f"}), + "avector_store_file_update": ("vsf.aupdate", {"vector_store_id": "v", "file_id": "f"}), + "avector_store_file_delete": ("vsf.adelete", {"vector_store_id": "v", "file_id": "f"}), +} + +ROUTED_SYNC: Final[dict[str, tuple[str, dict[str, object]]]] = { + "responses": ("responses", {"input": "hi"}), + "generate_content": ("generate_content", {"contents": []}), + "generate_content_stream": ("generate_content_stream", {"contents": []}), + "ocr": ("ocr", {"document": {}}), + "decisions": ("decisions", {}), + "systemone": ("systemone", {}), + "extraction": ("extraction", {}), + "video_generation": ("video_generation", {"prompt": "p"}), + "video_list": ("video_list", {}), + "video_status": ("video_status", {"video_id": "v"}), + "video_content": ("video_content", {"video_id": "v"}), + "video_remix": ("video_remix", {"video_id": "v", "prompt": "p"}), + "create_container": ("create_container", {"name": "c"}), + "list_containers": ("list_containers", {}), + "retrieve_container": ("retrieve_container", {"container_id": "c"}), + "delete_container": ("delete_container", {"container_id": "c"}), + "vector_store_search": ("vs.search", {"vector_store_id": "v", "query": "q"}), + "vector_store_create": ("vs.create", {"name": "v"}), + "vector_store_retrieve": ("vs.retrieve", {"vector_store_id": "v"}), + "vector_store_list": ("vs.list", {}), + "vector_store_update": ("vs.update", {"vector_store_id": "v"}), + "vector_store_delete": ("vs.delete", {"vector_store_id": "v"}), +} + + +@pytest.mark.parametrize("endpoint", sorted(ROUTED_ASYNC)) +async def test_async_generic_endpoint_routes_to_the_group_deployment_with_litellm_metadata( + make_core: MakeCore, host: EndpointHost, endpoint: str +): + op, request = ROUTED_ASYNC[endpoint] + core: Final = make_core(model_list=[deployment("g", "d1", api_base="https://d1.example")]) + response: Final = await _call(core, endpoint, {"model": "g", **request}) + assert [call.op for call in host.calls] == [op] + call: Final = host.calls[0] + assert (call.model, call.deployment_id, call.kwargs["api_base"], call.kwargs["custom_llm_provider"]) == ( + "openai/d1", + "d1", + "https://d1.example", + "openai", + ) + metadata: Final = _bucket(call.kwargs, "litellm_metadata") + assert (metadata["model_group"], metadata["deployment"]) == ("g", "openai/d1") + assert "metadata" not in call.kwargs + assert response._hidden_params["model_id"] == "d1" # pyright: ignore[reportAttributeAccessIssue] # the scripted response + assert response._hidden_params["additional_headers"]["x-litellm-model-group"] == "g" # pyright: ignore[reportAttributeAccessIssue] # the scripted response + + +@pytest.mark.parametrize("endpoint", sorted(ROUTED_SYNC)) +async def test_sync_generic_endpoint_routes_to_the_group_deployment_and_returns_the_response( + make_core: MakeCore, host: EndpointHost, endpoint: str +): + op, request = ROUTED_SYNC[endpoint] + core: Final = make_core(model_list=[deployment("g", "d1", api_base="https://d1.example")]) + response: Final = getattr(core, endpoint)(model="g", **request) + assert not inspect.isawaitable(response) + assert [call.op for call in host.calls] == [op] + call: Final = host.calls[0] + assert (call.model, call.deployment_id, call.kwargs["api_base"], call.kwargs["custom_llm_provider"]) == ( + "openai/d1", + "d1", + "https://d1.example", + "openai", + ) + assert _bucket(call.kwargs, "litellm_metadata")["model_group"] == "g" + assert response is not None + + +@pytest.mark.parametrize("endpoint", ["ocr", "generate_content", "video_generation"]) +async def test_sync_generic_endpoints_set_no_routing_hidden_params( + make_core: MakeCore, host: EndpointHost, endpoint: str +): + _, request = ROUTED_SYNC[endpoint] + core: Final = make_core(model_list=[deployment("g", "d1")]) + response: Final = getattr(core, endpoint)(model="g", **request) + assert response._hidden_params == {} + + +@pytest.mark.parametrize("endpoint", ["ocr", "generate_content"]) +async def test_sync_generic_endpoints_are_not_retried(make_core: MakeCore, host: EndpointHost, endpoint: str): + host.answer("openai/d1", litellm.InternalServerError(message="boom", llm_provider="openai", model="d1")) + _, request = ROUTED_SYNC[endpoint] + core: Final = make_core(model_list=[deployment("g", "d1")], num_retries=3) + with pytest.raises(litellm.InternalServerError): + getattr(core, endpoint)(model="g", **request) + assert len(host.calls) == 1 + + +async def test_async_generic_endpoint_is_retried_with_the_router_num_retries(make_core: MakeCore, host: EndpointHost): + host.answer("openai/d1", litellm.InternalServerError(message="boom", llm_provider="openai", model="d1")) + core: Final = make_core(model_list=[deployment("g", "d1")], num_retries=2) + with pytest.raises(litellm.InternalServerError): + await core.aocr(model="g", document={}) + assert len(host.calls) == 3 + + +@pytest.mark.parametrize( + "endpoint", + [ + "aadapter_generate_content", + "video_create_character", + "video_get_character", + "video_edit", + "video_extension", + "upload_container_file", + "list_container_files", + "retrieve_container_file", + "delete_container_file", + "retrieve_container_file_content", + "get_interaction", + "delete_interaction", + "cancel_interaction", + ], +) +async def test_endpoints_without_a_dispatch_branch_resolve_to_none_without_calling_anything( + make_core: MakeCore, host: EndpointHost, endpoint: str +): + """These are wrapped by factory_function's async wrapper, which has no branch for them.""" + core: Final = make_core(model_list=[deployment("g", "d1")]) + result: Final = getattr(core, endpoint)(model="g", container_id="c", video_id="v", interaction_id="i") + assert inspect.isawaitable(result) + assert await result is None + assert host.calls == [] + + +async def test_sync_create_interaction_is_awaited_and_fails_with_type_error(make_core: MakeCore, host: EndpointHost): + core: Final = make_core(model_list=[deployment("g", "d1")], num_retries=0) + with pytest.raises(TypeError, match="can't be used in 'await' expression"): + await core.create_interaction(model="g", input="x") + assert [call.op for call in host.calls] == ["interactions.create"] + + +# --- request params, provider and timeouts ------------------------------------------------------ + + +@pytest.mark.parametrize("endpoint", ["aocr", "ocr"]) +async def test_generic_request_params_win_over_deployment_params( + make_core: MakeCore, host: EndpointHost, endpoint: str +): + core: Final = make_core(model_list=[deployment("g", "d1", temperature=0.1)]) + await _call(core, endpoint, {"model": "g", "document": {}, "temperature": 0.9}) + assert host.calls[0].kwargs["temperature"] == 0.9 + + +async def test_generic_call_always_sends_the_deployment_model(make_core: MakeCore, host: EndpointHost): + core: Final = make_core(model_list=[deployment("g", "d1")]) + await core.aocr(model="g", document={}) + assert host.calls[0].model == "openai/d1" + + +async def test_generic_call_sends_the_declared_custom_llm_provider(make_core: MakeCore, host: EndpointHost): + core: Final = make_core(model_list=[deployment("g", "d1", model="my-model", custom_llm_provider="hosted_vllm")]) + await core.aocr(model="g", document={}) + assert (host.calls[0].model, host.calls[0].kwargs["custom_llm_provider"]) == ("my-model", "hosted_vllm") + + +@pytest.mark.parametrize( + ("endpoint", "expected"), + [("aocr", 600.0), ("aresponses", 600.0), ("ocr", 6000.0), ("generate_content", 6000.0)], +) +async def test_generic_default_timeout_is_the_passthrough_default_for_async_and_the_process_setting_for_sync( + make_core: MakeCore, host: EndpointHost, endpoint: str, expected: float +): + core: Final = make_core(model_list=[deployment("g", "d1")]) + request: Final = {**(ROUTED_ASYNC.get(endpoint) or ROUTED_SYNC[endpoint])[1], "model": "g"} + await _call(core, endpoint, request) + assert host.calls[0].kwargs["timeout"] == expected + + +@pytest.mark.parametrize("endpoint", ["aocr", "ocr"]) +async def test_generic_timeout_takes_the_router_timeout(make_core: MakeCore, host: EndpointHost, endpoint: str): + core: Final = make_core(model_list=[deployment("g", "d1")], timeout=42) + await _call(core, endpoint, {"model": "g", "document": {}}) + assert host.calls[0].kwargs["timeout"] == 42 + + +@pytest.mark.parametrize( + ("endpoint", "request_body"), [("aresponses", {"input": "hi"}), ("agenerate_content", {"contents": []})] +) +async def test_streaming_generic_call_sends_the_deployment_stream_timeout_as_timeout( + make_core: MakeCore, host: EndpointHost, endpoint: str, request_body: dict[str, object] +): + core: Final = make_core(model_list=[deployment("g", "d1", timeout=11, stream_timeout=3)]) + await _call(core, endpoint, {"model": "g", "stream": True, **request_body}) + sent: Final = host.calls[0].kwargs + assert (sent["timeout"], sent["stream_timeout"]) == (3.0, 3.0) + + +async def test_non_streaming_aresponses_sends_the_deployment_timeout(make_core: MakeCore, host: EndpointHost): + core: Final = make_core(model_list=[deployment("g", "d1", timeout=11, stream_timeout=3)]) + await core.aresponses(model="g", input="hi") + assert host.calls[0].kwargs["timeout"] == 11.0 + + +async def test_realtime_client_secret_session_model_follows_the_routed_deployment( + make_core: MakeCore, host: EndpointHost +): + core: Final = make_core(model_list=[deployment("g", "d1")]) + await core.acreate_realtime_client_secret(model="g", session={"model": "g", "type": "realtime"}) + assert host.calls[0].kwargs["session"] == {"model": "openai/d1", "type": "realtime"} + + +@pytest.mark.parametrize( + ("session", "sent"), + [ + ( + {"type": "transcription"}, + {"type": "transcription", "audio": {"input": {"transcription": {"model": "openai/d1"}}}}, + ), + ( + {"type": "transcription", "input_audio_transcription": {"language": "en"}}, + {"type": "transcription", "input_audio_transcription": {"language": "en", "model": "openai/d1"}}, + ), + ({"type": "realtime"}, {"type": "realtime"}), + ], + ids=["transcription-no-model", "transcription-flat-model", "realtime-no-model"], +) +async def test_a_transcription_session_with_no_model_follows_the_routed_deployment( + make_core: MakeCore, host: EndpointHost, session: dict[str, object], sent: dict[str, object] +): + core: Final = make_core(model_list=[deployment("g", "d1")]) + await core.acreate_realtime_translation_client_secret(model="g", session=session) + assert host.calls[0].kwargs["session"] == sent + + +@pytest.mark.parametrize("endpoint", ["acreate_realtime_translation_client_secret", "arealtime_translation_calls"]) +async def test_a_realtime_translation_endpoint_keeps_the_callers_client( + make_core: MakeCore, host: EndpointHost, endpoint: str +): + core: Final = make_core(model_list=[deployment("g", "d1")]) + client: Final = object() + await _call(core, endpoint, {"model": "g", "client": client}) + assert host.calls[0].kwargs["client"] is client + + +def _realtime_session_logging(details: Mapping[str, object]) -> LiteLLMLogging: + session_logging: Final = LiteLLMLogging( + model="gpt-realtime-2.1", + messages=[], + stream=False, + call_type="_arealtime", + start_time=datetime.now(), + litellm_call_id=str(uuid.uuid4()), + function_id="realtime-session", + ) + session_logging.model_call_details.update(details) + return session_logging + + +@pytest.mark.parametrize( + ("details", "counted"), + [ + ({REALTIME_SESSION_SUCCESS_LOGGED_KEY: True}, 1), + ({}, None), + ({REALTIME_SESSION_SUCCESS_LOGGED_KEY: True, REALTIME_SESSION_ENDED_ABNORMALLY_KEY: True}, None), + ], + ids=["clean-session", "no-success-logged", "ended-abnormally"], +) +async def test_a_realtime_session_counts_a_deployment_success_only_when_it_ended_cleanly( + make_core: MakeCore, host: EndpointHost, details: dict[str, object], counted: int | None +): + core: Final = make_core(model_list=[deployment("g", "d1")], num_retries=0) + await core._arealtime(model="g", litellm_logging_obj=_realtime_session_logging(details)) + assert core.health().store.get_cache("d1:successes", local_only=True) == counted + + +async def test_generic_endpoint_with_an_unknown_group_raises_without_calling_the_host( + make_core: MakeCore, host: EndpointHost +): + core: Final = make_core(model_list=[deployment("g", "d1")]) + with pytest.raises(litellm.BadRequestError, match="model=nope"): + await core.agenerate_content(model="nope", contents=[]) + assert host.calls == [] + + +# --- Responses API ids pin the deployment ----------------------------------------------------- + +TWO_DEPLOYMENTS: Final = [deployment("g", "d1"), deployment("g", "d2")] +ENCODED_RESPONSE_ID: Final = ResponsesAPIRequestUtils._build_responses_api_response_id("openai", "d2", "resp_abc") # pyright: ignore[reportPrivateUsage] # litellm's encoder + + +@pytest.mark.parametrize( + "endpoint", ["aget_responses", "acancel_responses", "acompact_responses", "adelete_responses", "alist_input_items"] +) +async def test_a_litellm_response_id_pins_the_deployment_it_encodes( + make_core: MakeCore, host: EndpointHost, endpoint: str +): + core: Final = make_core(model_list=TWO_DEPLOYMENTS) + for _ in range(5): + await _call(core, endpoint, {"response_id": ENCODED_RESPONSE_ID}) + assert host.deployments_called() == ["d2"] * 5 + + +async def test_a_litellm_response_id_overrides_the_request_model(make_core: MakeCore, host: EndpointHost): + core: Final = make_core(model_list=TWO_DEPLOYMENTS) + for _ in range(5): + await core.adelete_responses(response_id=ENCODED_RESPONSE_ID, model="g") + assert host.deployments_called() == ["d2"] * 5 + assert _bucket(host.calls[0].kwargs, "litellm_metadata")["model_group"] == "d2" + + +async def test_the_response_id_reaches_the_host_still_encoded(make_core: MakeCore, host: EndpointHost): + core: Final = make_core(model_list=TWO_DEPLOYMENTS) + await core.aget_responses(response_id=ENCODED_RESPONSE_ID) + assert host.calls[0].kwargs["response_id"] == ENCODED_RESPONSE_ID + + +async def test_a_plain_response_id_without_a_model_fails_with_type_error(make_core: MakeCore, host: EndpointHost): + core: Final = make_core(model_list=TWO_DEPLOYMENTS) + with pytest.raises(TypeError, match="missing 1 required positional argument: 'model'"): + await core.aget_responses(response_id="resp_plain") + assert host.calls == [] + + +# --- container ids ------------------------------------------------------------------------ + +ENCODED_CONTAINER_ID: Final = ResponsesAPIRequestUtils.build_container_id("azure", "d2", "cntr_upstream") + + +async def test_a_litellm_container_id_pins_its_deployment_and_is_sent_decoded(make_core: MakeCore, host: EndpointHost): + core: Final = make_core(model_list=TWO_DEPLOYMENTS) + for _ in range(5): + await core.aretrieve_container(container_id=ENCODED_CONTAINER_ID) + assert host.deployments_called() == ["d2"] * 5 + assert {call.kwargs["container_id"] for call in host.calls} == {"cntr_upstream"} + + +async def test_a_routed_container_call_sends_the_deployment_provider_not_the_decoded_one( + make_core: MakeCore, host: EndpointHost +): + core: Final = make_core(model_list=TWO_DEPLOYMENTS) + await core.aretrieve_container(container_id=ENCODED_CONTAINER_ID, custom_llm_provider="openai") + assert host.calls[0].kwargs["custom_llm_provider"] == "openai" + + +async def test_a_native_container_id_routes_by_the_forwarded_model_id(make_core: MakeCore, host: EndpointHost): + core: Final = make_core(model_list=TWO_DEPLOYMENTS) + for _ in range(5): + await core.aretrieve_container(container_id="cntr_native", model_id=" d1 ") + assert host.deployments_called() == ["d1"] * 5 + assert host.calls[0].kwargs["container_id"] == "cntr_native" + + +async def test_a_native_container_id_without_a_model_calls_the_handler_directly( + make_core: MakeCore, host: EndpointHost +): + core: Final = make_core(model_list=TWO_DEPLOYMENTS) + await core.aretrieve_container(container_id="cntr_native", custom_llm_provider="azure") + call: Final = host.calls[0] + assert (call.deployment_id, call.kwargs["container_id"], call.kwargs["custom_llm_provider"]) == ( + None, + "cntr_native", + "azure", + ) + assert "litellm_metadata" not in call.kwargs + + +async def test_create_container_with_an_unknown_model_passes_through_with_that_model( + make_core: MakeCore, host: EndpointHost +): + core: Final = make_core(model_list=TWO_DEPLOYMENTS) + await core.acreate_container(name="c", model="not-a-group") + call: Final = host.calls[0] + assert (call.op, call.model, call.deployment_id) == ("acreate_container", "not-a-group", None) + + +async def test_create_container_without_a_model_calls_the_handler_directly(make_core: MakeCore, host: EndpointHost): + core: Final = make_core(model_list=TWO_DEPLOYMENTS) + await core.acreate_container(name="c") + call: Final = host.calls[0] + assert (call.op, call.deployment_id, "model" in call.kwargs) == ("acreate_container", None, False) + + +# --- video ids ----------------------------------------------------------------------------- + +ENCODED_VIDEO_ID: Final = encode_video_id_with_provider("video_up", "openai", model_id="d2", deployment_id="d2") + + +@pytest.mark.parametrize( + "endpoint", ["avideo_status", "avideo_content", "avideo_remix", "avideo_edit", "avideo_extension"] +) +async def test_a_litellm_video_id_pins_the_deployment_that_created_it( + make_core: MakeCore, host: EndpointHost, endpoint: str +): + core: Final = make_core(model_list=TWO_DEPLOYMENTS) + for _ in range(6): + await _call(core, endpoint, {"model": "g", "video_id": ENCODED_VIDEO_ID, "prompt": "p"}) + assert host.deployments_called() == ["d2"] * 6 + + +async def test_a_video_id_on_a_video_generation_call_does_not_pin(make_core: MakeCore, host: EndpointHost): + core: Final = make_core(model_list=TWO_DEPLOYMENTS) + for _ in range(30): + await core.avideo_generation(model="g", prompt="p", video_id=ENCODED_VIDEO_ID) + assert set(host.deployments_called()) == {"d1", "d2"} + + +async def test_sync_video_status_pins_the_deployment_of_the_video_id(make_core: MakeCore, host: EndpointHost): + core: Final = make_core(model_list=TWO_DEPLOYMENTS) + for _ in range(6): + core.video_status(model="g", video_id=ENCODED_VIDEO_ID) + assert host.deployments_called() == ["d2"] * 6 + + +# --- LLM pass-through ---------------------------------------------------------------------- + + +async def test_llm_passthrough_with_an_unknown_model_calls_the_handler_with_that_model( + make_core: MakeCore, host: EndpointHost +): + core: Final = make_core(model_list=TWO_DEPLOYMENTS) + await core.allm_passthrough_route(model="not-a-group", method="POST", endpoint="/model/not-a-group/invoke") + call: Final = host.calls[0] + assert (call.model, call.deployment_id, call.kwargs["endpoint"]) == ( + "not-a-group", + None, + "/model/not-a-group/invoke", + ) + + +async def test_llm_passthrough_rewrites_the_group_name_in_the_endpoint_to_the_provider_model( + make_core: MakeCore, host: EndpointHost +): + core: Final = make_core(model_list=[deployment("g", "d1", model="bedrock/us.anthropic.claude")]) + await core.allm_passthrough_route(model="g", method="POST", endpoint="/model/g/invoke") + call: Final = host.calls[0] + assert (call.model, call.kwargs["endpoint"], call.kwargs["custom_llm_provider"]) == ( + "bedrock/us.anthropic.claude", + "/model/us.anthropic.claude/invoke", + "bedrock", + ) + + +async def test_llm_passthrough_forwards_the_request_client(make_core: MakeCore, host: EndpointHost): + client: Final = object() + core: Final = make_core(model_list=[deployment("g", "d1")]) + await core.allm_passthrough_route(model="g", method="POST", endpoint="/v1/x", client=client) + assert host.calls[0].kwargs["client"] is client + + +# --- interactions and managed agents ---------------------------------------------------------- + + +async def test_interactions_get_without_a_model_calls_the_handler_with_the_given_provider( + make_core: MakeCore, host: EndpointHost +): + core: Final = make_core(model_list=TWO_DEPLOYMENTS) + await core.aget_interaction(interaction_id="i", custom_llm_provider="vertex_ai") + call: Final = host.calls[0] + assert (call.op, call.deployment_id, call.kwargs["custom_llm_provider"]) == ("interactions.aget", None, "vertex_ai") + + +@pytest.mark.parametrize("endpoint", ["aget_interaction", "adelete_interaction", "acancel_interaction"]) +async def test_interactions_default_to_the_gemini_provider(make_core: MakeCore, host: EndpointHost, endpoint: str): + core: Final = make_core(model_list=TWO_DEPLOYMENTS) + await _call(core, endpoint, {"interaction_id": "i"}) + assert host.calls[0].kwargs["custom_llm_provider"] == "gemini" + + +async def test_interaction_with_an_agent_named_as_the_model_is_not_routed(make_core: MakeCore, host: EndpointHost): + core: Final = make_core(model_list=[deployment("ag", "d1")]) + await core.acreate_interaction(agent="ag", model="ag", input="x") + call: Final = host.calls[0] + assert (call.deployment_id, call.kwargs["model"], call.kwargs["custom_llm_provider"]) == (None, None, "gemini") + + +@pytest.mark.parametrize( + ("endpoint", "op"), + [ + ("acreate_agent", "agents.acreate"), + ("alist_agents", "agents.alist"), + ("aget_agent", "agents.aget"), + ("adelete_agent", "agents.adelete"), + ("alist_agent_versions", "agents.alist_versions"), + ("create_agent", "agents.create"), + ("list_agents", "agents.list"), + ("get_agent", "agents.get"), + ("delete_agent", "agents.delete"), + ("list_agent_versions", "agents.list_versions"), + ], +) +async def test_managed_agent_endpoints_call_the_handler_directly_with_gemini_by_default( + make_core: MakeCore, host: EndpointHost, endpoint: str, op: str +): + core: Final = make_core(model_list=[deployment("g", "d1")]) + await _call(core, endpoint, {"name": "a", "model": "g"}) + call: Final = host.calls[0] + assert (call.op, call.deployment_id, call.kwargs["custom_llm_provider"], call.kwargs["model"]) == ( + op, + None, + "gemini", + "g", + ) + + +@pytest.mark.parametrize("endpoint", ["acreate_agent", "create_agent"]) +async def test_managed_agent_endpoints_keep_the_given_provider(make_core: MakeCore, host: EndpointHost, endpoint: str): + core: Final = make_core(model_list=[deployment("g", "d1")]) + await _call(core, endpoint, {"name": "a", "custom_llm_provider": "vertex_ai"}) + assert host.calls[0].kwargs["custom_llm_provider"] == "vertex_ai" + + +# --- assistants ---------------------------------------------------------------------------- + +ASSISTANT_ENDPOINTS: Final = ( + "acreate_assistants", + "adelete_assistant", + "aget_assistants", + "acreate_thread", + "aget_thread", + "a_add_message", + "aget_messages", + "arun_thread", +) + + +@pytest.mark.parametrize("endpoint", ASSISTANT_ENDPOINTS) +async def test_assistants_endpoints_without_a_provider_or_config_raise( + make_core: MakeCore, host: EndpointHost, endpoint: str +): + core: Final = make_core(model_list=[deployment("g", "d1")]) + with pytest.raises(Exception, match=re.escape("'custom_llm_provider' must be set. Either via:")): + await getattr(core, endpoint)() + assert host.calls == [] + + +@pytest.mark.parametrize("endpoint", ASSISTANT_ENDPOINTS) +async def test_assistants_endpoints_use_the_assistants_config(make_core: MakeCore, host: EndpointHost, endpoint: str): + core: Final = make_core( + model_list=[deployment("g", "d1")], + assistants_config={ + "custom_llm_provider": "azure", + "litellm_params": {"api_key": "k", "api_base": "https://az"}, + }, + ) + await getattr(core, endpoint)(thread_id="t") + call: Final = host.calls[0] + assert (call.op, call.kwargs["custom_llm_provider"], call.kwargs["api_key"], call.kwargs["api_base"]) == ( + endpoint, + "azure", + "k", + "https://az", + ) + assert (call.kwargs["client"], call.kwargs["thread_id"]) == (None, "t") + + +async def test_assistants_endpoint_with_an_explicit_provider_ignores_the_config( + make_core: MakeCore, host: EndpointHost +): + client: Final = object() + core: Final = make_core( + model_list=[deployment("g", "d1")], + assistants_config={"custom_llm_provider": "azure", "litellm_params": {"api_key": "k"}}, + ) + await core.arun_thread(custom_llm_provider="openai", client=client, thread_id="t") + call: Final = host.calls[0] + assert (call.kwargs["custom_llm_provider"], call.kwargs["client"], "api_key" in call.kwargs) == ( + "openai", + client, + False, + ) + + +# --- vector stores --------------------------------------------------------------------------- + + +async def test_vector_store_search_without_a_model_gets_the_router_and_overrides_a_client_router( + make_core: MakeCore, host: EndpointHost +): + core: Final = make_core(model_list=TWO_DEPLOYMENTS) + await core.avector_store_search(vector_store_id="v", query="q", router="client-value") + call: Final = host.calls[0] + assert (call.op, call.deployment_id, call.kwargs["router"]) == ("vs.asearch", None, core) + assert "_direct_vector_store_embedding_executor" in call.kwargs + + +async def test_sync_vector_store_search_without_a_model_gets_the_router(make_core: MakeCore, host: EndpointHost): + core: Final = make_core(model_list=TWO_DEPLOYMENTS) + core.vector_store_search(vector_store_id="v", query="q") + assert host.calls[0].kwargs["router"] is core + + +async def test_routed_vector_store_search_gets_the_embedding_executor_but_not_the_router( + make_core: MakeCore, host: EndpointHost +): + core: Final = make_core(model_list=[deployment("g", "d1")]) + await core.avector_store_search(model="g", vector_store_id="v", query="q") + sent: Final = host.calls[0].kwargs + assert ("router" in sent, "_direct_vector_store_embedding_executor" in sent) == (False, True) + + +async def test_vector_store_retrieve_without_a_model_calls_the_handler_directly( + make_core: MakeCore, host: EndpointHost +): + core: Final = make_core(model_list=TWO_DEPLOYMENTS) + await core.avector_store_retrieve(vector_store_id="v", custom_llm_provider="openai") + call: Final = host.calls[0] + assert (call.deployment_id, call.kwargs["custom_llm_provider"], "router" in call.kwargs) == (None, "openai", False) + + +async def test_sync_vector_store_create_without_a_model_passes_the_provider(make_core: MakeCore, host: EndpointHost): + core: Final = make_core(model_list=TWO_DEPLOYMENTS) + core.vector_store_create(name="v", custom_llm_provider="openai") + call: Final = host.calls[0] + assert (call.op, call.deployment_id, call.kwargs["custom_llm_provider"]) == ("vs.create", None, "openai") + + +@pytest.mark.parametrize( + ("endpoint", "op"), + [ + ("vector_store_file_create", "vsf.create"), + ("vector_store_file_list", "vsf.list"), + ("vector_store_file_retrieve", "vsf.retrieve"), + ("vector_store_file_content", "vsf.retrieve_content"), + ("vector_store_file_update", "vsf.update"), + ("vector_store_file_delete", "vsf.delete"), + ], +) +async def test_sync_vector_store_file_endpoints_never_route_even_with_a_model( + make_core: MakeCore, host: EndpointHost, endpoint: str, op: str +): + core: Final = make_core(model_list=[deployment("g", "d1")]) + getattr(core, endpoint)(model="g", vector_store_id="v", file_id="f") + call: Final = host.calls[0] + assert (call.op, call.model, call.deployment_id, call.kwargs["custom_llm_provider"], call.kwargs["client"]) == ( + op, + "g", + None, + None, + None, + ) + + +async def test_avector_store_create_with_a_model_routes_with_metadata_and_no_model_group( + make_core: MakeCore, host: EndpointHost +): + core: Final = make_core(model_list=[deployment("g", "d1", api_base="https://d1.example")]) + await core.avector_store_create(model="g", name="n") + call: Final = host.calls[0] + assert (call.op, call.model, call.deployment_id, call.kwargs["custom_llm_provider"]) == ( + "vs.acreate", + "openai/d1", + "d1", + "openai", + ) + metadata: Final = _bucket(call.kwargs, "metadata") + assert (metadata["deployment"], "model_group" in metadata) == ("openai/d1", False) + + +async def test_avector_store_create_without_a_model_calls_the_handler_directly(make_core: MakeCore, host: EndpointHost): + core: Final = make_core(model_list=[deployment("g", "d1")]) + await core.avector_store_create(model=None, name="n") + call: Final = host.calls[0] + assert (call.op, call.deployment_id) == ("vs.acreate", None) + + +async def test_avector_store_create_is_not_retried(make_core: MakeCore, host: EndpointHost): + host.answer("openai/d1", litellm.InternalServerError(message="boom", llm_provider="openai", model="d1")) + core: Final = make_core(model_list=[deployment("g", "d1")], num_retries=3) + with pytest.raises(litellm.InternalServerError): + await core.avector_store_create(model="g", name="n") + assert len(host.calls) == 1 + + +# --- search tools --------------------------------------------------------------------------- + +SEARCH_TOOLS: Final = [{"search_tool_name": "s", "litellm_params": {"search_provider": "tavily", "api_key": "tk"}}] + + +async def test_asearch_calls_the_named_search_tool_with_its_litellm_params(make_core: MakeCore, host: EndpointHost): + core: Final = make_core(model_list=[deployment("g", "d1")], search_tools=SEARCH_TOOLS) + response: Final = await core.asearch(model="s", query="q") + call: Final = host.calls[0] + assert (call.op, call.kwargs["search_provider"], call.kwargs["api_key"], call.kwargs["query"]) == ( + "asearch", + "tavily", + "tk", + "q", + ) + assert _bucket(call.kwargs, "litellm_metadata")["model_group"] == "s" + assert response._hidden_params["additional_headers"]["x-litellm-model-group"] == "s" # pyright: ignore[reportAttributeAccessIssue] # the scripted response + + +async def test_asearch_with_an_unknown_tool_raises_value_error(make_core: MakeCore, host: EndpointHost): + core: Final = make_core(model_list=[deployment("g", "d1")], search_tools=SEARCH_TOOLS, num_retries=0) + with pytest.raises(ValueError, match=re.escape("Search tool 'g' not found in router.search_tools")): + await core.asearch(model="g", query="q") + assert host.calls == [] + + +# --- guardrails --------------------------------------------------------------------------- + +GUARDRAILS: Final = [{"guardrail_name": "gr", "litellm_params": {"guardrail": "presidio", "mode": "pre_call"}}] + + +async def test_aguardrail_runs_the_function_with_the_selected_guardrail(make_core: MakeCore, host: ScriptedHost): + seen: Final[list[Mapping[str, object]]] = [] + + async def guard(**kwargs: object) -> str: # kwargs-ok: the guardrail's keyword payload + seen.append(kwargs) + return "verdict" + + core: Final = make_core(model_list=[deployment("g", "d1")], guardrail_list=GUARDRAILS) + assert await core.aguardrail(guardrail_name="gr", original_function=guard, data={"x": 1}) == "verdict" + assert sorted(seen[0]) == ["data", "litellm_metadata", "litellm_trace_id", "selected_guardrail"] + assert seen[0]["selected_guardrail"] == GUARDRAILS[0] + assert _bucket(seen[0], "litellm_metadata")["model_group"] == "gr" + assert host.calls == [] + + +async def test_aguardrail_with_an_unknown_name_raises_value_error(make_core: MakeCore, host: ScriptedHost): + async def guard(**_kwargs: object) -> str: # kwargs-ok: the guardrail's keyword payload + return "verdict" + + core: Final = make_core(model_list=[deployment("g", "d1")], guardrail_list=GUARDRAILS, num_retries=0) + with pytest.raises(ValueError, match="No guardrail found with name: nope"): + await core.aguardrail(guardrail_name="nope", original_function=guard) + + +async def test_aguardrail_does_not_retry_a_failing_guardrail_function(make_core: MakeCore, host: ScriptedHost): + """A guardrail name is no model group (model_group_size 0), so the router num_retries do not apply.""" + attempts: Final[list[int]] = [] + + async def guard(**_kwargs: object) -> str: # kwargs-ok: the guardrail's keyword payload + attempts.append(1) + raise litellm.InternalServerError(message="boom", llm_provider="", model="gr") + + core: Final = make_core(model_list=[deployment("g", "d1")], guardrail_list=GUARDRAILS, num_retries=2) + with pytest.raises(litellm.InternalServerError, match="boom"): + await core.aguardrail(guardrail_name="gr", original_function=guard) + assert len(attempts) == 1 diff --git a/tests/unit/test_router/core/test_endpoints_characterization_inference.py b/tests/unit/test_router/core/test_endpoints_characterization_inference.py new file mode 100644 index 00000000000..8db0d6bea12 --- /dev/null +++ b/tests/unit/test_router/core/test_endpoints_characterization_inference.py @@ -0,0 +1,648 @@ +""" +Characterization of the inference endpoints (completion, text completion, embedding, image, +speech, transcription, rerank, adapter completion, moderation): which host op each invokes and +which keyword arguments the core builds for it. + +Every test builds its core through ``make_core`` on the ``host`` fixture, so it runs once per core. +""" + +import asyncio +import io +from collections.abc import Callable, Mapping +from typing import Final + +import pytest + +import litellm +from litellm.router_core import RouterCore +from tests.unit.test_router.core.endpoints_helpers import EndpointHost +from tests.unit.test_router.core.scripted_host import ScriptedHost, deployment, timeout + +MESSAGES: Final = [{"role": "user", "content": "hi"}] +MakeCore = Callable[..., RouterCore] + + +@pytest.fixture +def host() -> EndpointHost: + return EndpointHost() + + +def _sent(host: ScriptedHost, index: int = -1) -> Mapping[str, object]: + return host.calls[index].kwargs + + +def _bucket(kwargs: Mapping[str, object], name: str = "metadata") -> Mapping[str, object]: + bucket: Final = kwargs[name] + assert isinstance(bucket, Mapping) + return bucket + + +# --- host op per endpoint ------------------------------------------------------------------- + + +async def test_aadapter_completion_invokes_the_adapter_host_op_with_the_adapter_id( + make_core: MakeCore, host: EndpointHost +): + core: Final = make_core(model_list=[deployment("g", "d1")]) + await core.aadapter_completion(adapter_id="my-adapter", model="g", messages=MESSAGES) + call: Final = host.calls[0] + assert (call.op, call.model, call.kwargs["adapter_id"], call.deployment_id, call.model_group) == ( + "aadapter_completion", + "openai/d1", + "my-adapter", + "d1", + "g", + ) + + +async def test_amoderation_routes_a_known_group_through_its_deployment(make_core: MakeCore, host: EndpointHost): + core: Final = make_core(model_list=[deployment("g", "d1", api_base="https://d1.example")]) + await core.amoderation(model="g", input="x") + call: Final = host.calls[0] + assert (call.op, call.model, call.kwargs["api_base"], call.deployment_id, call.model_group) == ( + "amoderation", + "openai/d1", + "https://d1.example", + "d1", + "g", + ) + assert "litellm_metadata" not in call.kwargs + + +async def test_amoderation_calls_the_host_directly_when_the_model_is_not_a_group( + make_core: MakeCore, host: EndpointHost +): + core: Final = make_core(model_list=[deployment("g", "d1")]) + await core.amoderation(model="omni-moderation-latest", input="x") + call: Final = host.calls[0] + assert (call.op, call.model, call.deployment_id) == ("amoderation", "omni-moderation-latest", None) + assert _bucket(call.kwargs)["model_group"] == "omni-moderation-latest" + + +async def test_amoderation_deployment_params_overwrite_the_request_params(make_core: MakeCore, host: EndpointHost): + """The deployment's litellm_params are applied after the request's (kwargs.update(data)).""" + core: Final = make_core(model_list=[deployment("g", "d1", temperature=0.5)]) + await core.amoderation(model="g", input="x", temperature=1.0) + assert _sent(host)["temperature"] == 0.5 + + +async def test_amoderation_with_a_request_api_key_keeps_the_deployment_key_under_a_request_scoped_id( + make_core: MakeCore, host: EndpointHost +): + core: Final = make_core(model_list=[deployment("g", "d1")]) + await core.amoderation(model="g", input="x", api_key="sk-request") + sent: Final = _sent(host) + model_info: Final = _bucket(sent, "model_info") + assert sent["api_key"] == "sk-test" + assert (model_info["original_model_id"], model_info["id"] != "d1") == ("d1", True) + + +async def test_amoderation_is_not_retried(make_core: MakeCore, host: EndpointHost): + host.answer("openai/d1", litellm.InternalServerError(message="boom", llm_provider="openai", model="d1")) + core: Final = make_core(model_list=[deployment("g", "d1")], num_retries=3) + with pytest.raises(litellm.InternalServerError): + await core.amoderation(model="g", input="x") + assert len(host.calls) == 1 + assert host.sleeps == [] + + +async def test_arerank_sends_the_builtin_input_function_as_input(make_core: MakeCore, host: EndpointHost): + """arerank sets kwargs["input"] = input, the Python builtin, and it reaches the host.""" + core: Final = make_core(model_list=[deployment("g", "d1")]) + await core.arerank(model="g", query="q", documents=["a"]) + sent: Final = _sent(host) + assert (sent["input"], sent["query"], sent["documents"]) == (input, "q", ["a"]) + + +@pytest.mark.parametrize(("request_voice", "expected"), [(None, "deployment-voice"), ("alloy", "alloy")]) +async def test_aspeech_voice_defaults_to_the_deployment_voice( + make_core: MakeCore, host: EndpointHost, request_voice: str | None, expected: str +): + core: Final = make_core(model_list=[deployment("g", "d1", voice="deployment-voice")]) + await core.aspeech(model="g", input="hello", voice=request_voice) + assert (_sent(host)["voice"], _sent(host)["input"]) == (expected, "hello") + + +async def test_aspeech_sends_no_caching_flag(make_core: MakeCore, host: EndpointHost): + core: Final = make_core(model_list=[deployment("g", "d1")], cache_responses=True) + await core.aspeech(model="g", input="hello", voice="alloy") + assert "caching" not in _sent(host) + + +async def test_atranscription_forwards_the_same_file_object(make_core: MakeCore, host: EndpointHost): + audio: Final = io.BytesIO(b"audio") + core: Final = make_core(model_list=[deployment("g", "d1")]) + await core.atranscription(model="g", file=audio) + assert _sent(host)["file"] is audio + + +@pytest.mark.parametrize( + "endpoint", + ["acompletion", "atext_completion", "aembedding", "aimage_generation", "atranscription", "arerank", "completion"], +) +async def test_cache_responses_reaches_the_host_as_caching(make_core: MakeCore, host: EndpointHost, endpoint: str): + core: Final = make_core(model_list=[deployment("g", "d1")], cache_responses=True) + requests: Final[dict[str, dict[str, object]]] = { + "acompletion": {"messages": MESSAGES}, + "completion": {"messages": MESSAGES}, + "atext_completion": {"prompt": "hi"}, + "aembedding": {"input": ["a"]}, + "aimage_generation": {"prompt": "a cat"}, + "atranscription": {"file": io.BytesIO(b"x")}, + "arerank": {"query": "q", "documents": ["a"]}, + } + await core.call(endpoint, {"model": "g", **requests[endpoint]}) + assert _sent(host)["caching"] is True + + +@pytest.mark.parametrize("endpoint", ["acompletion", "completion"]) +async def test_cache_responses_turned_on_after_construction_reaches_the_host( + make_core: MakeCore, host: EndpointHost, endpoint: str +): + """The proxy turns ``cache_responses`` on through ``apply`` once its cache exists.""" + from litellm.router_api import ChangeSet + + core: Final = make_core(model_list=[deployment("g", "d1")]) + _ = core.apply(ChangeSet(assign={"cache_responses": True})) + await core.call(endpoint, {"model": "g", "messages": MESSAGES}) + assert core.settings().cache_responses is True + assert _sent(host)["caching"] is True + + +# --- sync text completion and image generation ---------------------------------------------- + + +async def test_sync_text_completion_sends_the_router_num_retries_and_no_routing_fields( + make_core: MakeCore, host: EndpointHost +): + core: Final = make_core(model_list=[deployment("g", "d1")], num_retries=4) + core.text_completion(model="g", prompt="p") + sent: Final = _sent(host) + assert sent["num_retries"] == 4 + assert "model_info" not in sent + assert "litellm_trace_id" not in sent + assert sent.get("timeout") is None + metadata: Final = _bucket(sent) + assert metadata["model_group"] == "g" + assert "deployment" not in metadata + + +async def test_sync_text_completion_request_num_retries_wins(make_core: MakeCore, host: EndpointHost): + core: Final = make_core(model_list=[deployment("g", "d1")], num_retries=4) + core.text_completion(model="g", prompt="p", num_retries=1) + assert _sent(host)["num_retries"] == 1 + + +async def test_sync_text_completion_merges_default_metadata_into_request_metadata( + make_core: MakeCore, host: EndpointHost +): + core: Final = make_core( + model_list=[deployment("g", "d1")], default_litellm_params={"metadata": {"x": 1}, "user": "u"} + ) + core.text_completion(model="g", prompt="p", metadata={"y": 2}) + sent: Final = _sent(host) + metadata: Final = _bucket(sent) + assert (metadata["x"], metadata["y"], metadata["model_group"], sent["user"]) == (1, 2, "g", "u") + + +async def test_sync_text_completion_sends_the_model_group_name_not_the_deployment_model( + make_core: MakeCore, host: EndpointHost +): + """The request's model="g" is spread after the deployment params, so it replaces "openai/d1".""" + core: Final = make_core(model_list=[deployment("g", "d1", api_base="https://d1.example")]) + core.text_completion(model="g", prompt="p") + call: Final = host.calls[0] + assert (call.model, call.kwargs["api_base"]) == ("g", "https://d1.example") + + +async def test_sync_text_completion_is_not_retried(make_core: MakeCore, host: EndpointHost): + host.answer("g", litellm.InternalServerError(message="boom", llm_provider="openai", model="d1")) + core: Final = make_core(model_list=[deployment("g", "d1")], num_retries=3) + with pytest.raises(litellm.InternalServerError): + core.text_completion(model="g", prompt="p") + assert len(host.calls) == 1 + assert host.sleeps == [] + + +async def test_sync_image_generation_sends_routing_metadata_without_a_trace_id(make_core: MakeCore, host: EndpointHost): + core: Final = make_core(model_list=[deployment("g", "d1")]) + core.image_generation(model="g", prompt="p") + sent: Final = _sent(host) + assert "litellm_trace_id" not in sent + assert (_bucket(sent)["model_group"], _bucket(sent)["deployment"], sent["prompt"]) == ("g", "openai/d1", "p") + + +@pytest.mark.parametrize("endpoint", ["acompletion", "atext_completion", "aembedding", "aresponses"]) +async def test_async_endpoints_generate_a_trace_id_unless_the_request_sends_one( + make_core: MakeCore, host: EndpointHost, endpoint: str +): + core: Final = make_core(model_list=[deployment("g", "d1")]) + request: Final[dict[str, object]] = ( + {"messages": MESSAGES} + if endpoint == "acompletion" + else {"prompt": "p"} + if endpoint == "atext_completion" + else {"input": ["a"]} + if endpoint == "aembedding" + else {"input": "hi"} + ) + await core.call(endpoint, {"model": "g", **request}) + await core.call(endpoint, {"model": "g", "litellm_trace_id": "trace-1", **request}) + generated: Final = _sent(host, 0)["litellm_trace_id"] + assert isinstance(generated, str) and len(generated) == 36 + assert _sent(host, 1)["litellm_trace_id"] == "trace-1" + + +# --- routing metadata --------------------------------------------------------------------- + + +async def test_a_model_group_alias_is_recorded_as_the_model_group_and_its_alias( + make_core: MakeCore, host: EndpointHost +): + core: Final = make_core(model_list=[deployment("g", "d1")], model_group_alias={"al": "g"}) + await core.acompletion(model="al", messages=MESSAGES) + metadata: Final = _bucket(_sent(host)) + assert (metadata["model_group"], metadata["model_group_alias"], metadata["deployment_model_name"]) == ( + "al", + "al", + "g", + ) + + +async def test_a_request_for_a_deployment_id_routes_to_that_deployment(make_core: MakeCore, host: EndpointHost): + core: Final = make_core(model_list=[deployment("g", "d1"), deployment("g", "d2")]) + for _ in range(4): + await core.acompletion(model="d2", messages=MESSAGES) + assert host.deployments_called() == ["d2"] * 4 + assert _bucket(_sent(host))["model_group"] == "d2" + + +async def test_deployment_tags_and_credential_name_join_the_request_tags(make_core: MakeCore, host: EndpointHost): + core: Final = make_core(model_list=[deployment("g", "d1", tags=["t1"], litellm_credential_name="cred")]) + await core.acompletion(model="g", messages=MESSAGES, metadata={"tags": ["r1"]}) + sent: Final = _sent(host) + assert _bucket(sent)["tags"] == ["r1", "t1", "Credential: cred"] + assert sent["litellm_credential_name"] == "cred" + + +async def test_daily_reset_time_is_not_sent_to_the_host(make_core: MakeCore, host: EndpointHost): + core: Final = make_core(model_list=[deployment("g", "d1")]) + await core.acompletion(model="g", messages=MESSAGES, daily_reset_time="00:00") + assert "daily_reset_time" not in _sent(host) + + +async def test_deployment_model_id_param_is_sent_as_model_id(make_core: MakeCore, host: EndpointHost): + core: Final = make_core(model_list=[deployment("g", "d1", model_id="arn:pinned")]) + await core.acompletion(model="g", messages=MESSAGES) + core.text_completion(model="g", prompt="p") + assert [call.kwargs["model_id"] for call in host.calls] == ["arn:pinned", "arn:pinned"] + + +async def test_deployment_model_id_reaches_the_host_even_when_clientside_auth_allows_model_id( + make_core: MakeCore, host: EndpointHost +): + """The pin is skipped, but the deployment's litellm_params (with model_id) are still spread into the call.""" + core: Final = make_core( + model_list=[deployment("g", "d1", model_id="arn:pinned", configurable_clientside_auth_params=["model_id"])] + ) + await core.acompletion(model="g", messages=MESSAGES) + assert _sent(host)["model_id"] == "arn:pinned" + + +async def test_request_model_id_wins_over_a_clientside_configurable_deployment_model_id( + make_core: MakeCore, host: EndpointHost +): + core: Final = make_core( + model_list=[deployment("g", "d1", model_id="arn:pinned", configurable_clientside_auth_params=["model_id"])] + ) + await core.acompletion(model="g", messages=MESSAGES, model_id="arn:request") + assert _sent(host)["model_id"] == "arn:request" + + +# --- default litellm params ------------------------------------------------------------------- + + +async def test_default_litellm_params_metadata_merges_into_request_metadata(make_core: MakeCore, host: EndpointHost): + core: Final = make_core(model_list=[deployment("g", "d1")], default_litellm_params={"metadata": {"team": "t"}}) + await core.acompletion(model="g", messages=MESSAGES, metadata={"trace": "x"}) + metadata: Final = _bucket(_sent(host)) + assert (metadata["team"], metadata["trace"], metadata["model_group"]) == ("t", "x", "g") + + +async def test_default_litellm_params_metadata_overwrites_the_same_request_metadata_key( + make_core: MakeCore, host: EndpointHost +): + core: Final = make_core(model_list=[deployment("g", "d1")], default_litellm_params={"metadata": {"team": "t"}}) + await core.acompletion(model="g", messages=MESSAGES, metadata={"team": "request"}) + assert _bucket(_sent(host))["team"] == "t" + + +async def test_default_litellm_params_metadata_goes_to_litellm_metadata_for_generic_endpoints( + make_core: MakeCore, host: EndpointHost +): + core: Final = make_core(model_list=[deployment("g", "d1")], default_litellm_params={"metadata": {"team": "t"}}) + await core.aresponses(model="g", input="hi") + sent: Final = _sent(host) + assert _bucket(sent, "litellm_metadata")["team"] == "t" + assert "metadata" not in sent + + +async def test_default_litellm_params_with_none_values_are_not_sent(make_core: MakeCore, host: EndpointHost): + core: Final = make_core(model_list=[deployment("g", "d1")], default_litellm_params={"user": None}) + await core.acompletion(model="g", messages=MESSAGES) + assert "user" not in _sent(host) + + +# --- tools ---------------------------------------------------------------------------------- + +DEPLOYMENT_TOOL: Final = {"type": "function", "function": {"name": "dep", "parameters": {}}} +REQUEST_TOOL: Final = {"type": "function", "function": {"name": "req", "parameters": {}}} + + +async def test_deployment_tools_come_before_request_tools(make_core: MakeCore, host: EndpointHost): + core: Final = make_core(model_list=[deployment("g", "d1", tools=[DEPLOYMENT_TOOL])]) + await core.acompletion(model="g", messages=MESSAGES, tools=[REQUEST_TOOL]) + assert _sent(host)["tools"] == [DEPLOYMENT_TOOL, REQUEST_TOOL] + + +async def test_deployment_tools_are_sent_when_the_request_has_none(make_core: MakeCore, host: EndpointHost): + core: Final = make_core(model_list=[deployment("g", "d1", tools=[DEPLOYMENT_TOOL])]) + await core.acompletion(model="g", messages=MESSAGES) + assert _sent(host)["tools"] == [DEPLOYMENT_TOOL] + + +@pytest.mark.parametrize(("request_choice", "expected"), [(None, "required"), ("none", "none")]) +async def test_tool_choice_comes_from_the_request_else_the_deployment( + make_core: MakeCore, host: EndpointHost, request_choice: str | None, expected: str +): + core: Final = make_core(model_list=[deployment("g", "d1", tools=[DEPLOYMENT_TOOL], tool_choice="required")]) + extra: Final = {} if request_choice is None else {"tool_choice": request_choice} + await core.acompletion(model="g", messages=MESSAGES, **extra) + assert _sent(host)["tool_choice"] == expected + + +async def test_merging_tools_does_not_change_the_deployment_tools(make_core: MakeCore, host: EndpointHost): + core: Final = make_core(model_list=[deployment("g", "d1", tools=[DEPLOYMENT_TOOL])]) + await core.acompletion(model="g", messages=MESSAGES, tools=[REQUEST_TOOL]) + await core.acompletion(model="g", messages=MESSAGES) + assert _sent(host)["tools"] == [DEPLOYMENT_TOOL] + + +# --- reasoning effort ----------------------------------------------------------------------- + + +async def test_request_reasoning_effort_drops_the_deployment_thinking_param(make_core: MakeCore, host: EndpointHost): + core: Final = make_core(model_list=[deployment("g", "d1", thinking={"type": "enabled", "budget_tokens": 100})]) + await core.acompletion(model="g", messages=MESSAGES, reasoning_effort="low") + sent: Final = _sent(host) + assert (sent["reasoning_effort"], "thinking" in sent) == ("low", False) + + +async def test_deployment_thinking_is_kept_without_a_request_reasoning_effort(make_core: MakeCore, host: EndpointHost): + thinking: Final = {"type": "enabled", "budget_tokens": 100} + core: Final = make_core(model_list=[deployment("g", "d1", thinking=thinking)]) + await core.acompletion(model="g", messages=MESSAGES) + assert _sent(host)["thinking"] == thinking + + +async def test_request_reasoning_effort_strips_effort_from_deployment_extra_body( + make_core: MakeCore, host: EndpointHost +): + core: Final = make_core( + model_list=[ + deployment( + "g", + "d1", + extra_body={"reasoning_effort": "high", "thinking": {"type": "enabled"}, "other": 1}, + reasoning={"effort": "high", "summary": "auto"}, + ) + ] + ) + await core.acompletion(model="g", messages=MESSAGES, reasoning_effort="low") + sent: Final = _sent(host) + assert sent["extra_body"] == {"other": 1} + assert sent["reasoning"] == {"summary": "auto"} + + +async def test_request_reasoning_effort_drops_an_extra_body_left_empty(make_core: MakeCore, host: EndpointHost): + core: Final = make_core(model_list=[deployment("g", "d1", extra_body={"reasoning_effort": "high"})]) + await core.acompletion(model="g", messages=MESSAGES, reasoning_effort="low") + assert "extra_body" not in _sent(host) + + +async def test_reasoning_override_does_not_change_the_deployment_config(make_core: MakeCore, host: EndpointHost): + thinking: Final = {"type": "enabled", "budget_tokens": 100} + core: Final = make_core(model_list=[deployment("g", "d1", thinking=thinking)]) + await core.acompletion(model="g", messages=MESSAGES, reasoning_effort="low") + await core.acompletion(model="g", messages=MESSAGES) + assert _sent(host)["thinking"] == thinking + + +CLASSIFIER_METADATA: Final = {"internal_call_origin": "autorouter_classifier"} + + +async def test_classifier_call_drops_reasoning_effort_the_deployment_does_not_support( + make_core: MakeCore, host: EndpointHost +): + core: Final = make_core(model_list=[deployment("g", "d1", model="openai/gpt-3.5-turbo")]) + body: Final[dict[str, object]] = {"reasoning_effort": "low", "keep": 1} + await core.acompletion( + model="g", + messages=MESSAGES, + reasoning_effort="low", + metadata=dict(CLASSIFIER_METADATA), + proxy_server_request={"body": body}, + ) + sent: Final = _sent(host) + assert "reasoning_effort" not in sent + assert sent["proxy_server_request"] == {"body": {"keep": 1}} + + +async def test_non_classifier_call_keeps_reasoning_effort_on_a_deployment_without_support( + make_core: MakeCore, host: EndpointHost +): + core: Final = make_core(model_list=[deployment("g", "d1", model="openai/gpt-3.5-turbo")]) + await core.acompletion(model="g", messages=MESSAGES, reasoning_effort="low") + assert _sent(host)["reasoning_effort"] == "low" + + +async def test_classifier_call_keeps_reasoning_effort_on_a_deployment_with_support( + make_core: MakeCore, host: EndpointHost +): + core: Final = make_core(model_list=[deployment("g", "d1", model="openai/o3-mini")]) + await core.acompletion( + model="g", messages=MESSAGES, reasoning_effort="low", litellm_metadata=dict(CLASSIFIER_METADATA) + ) + assert _sent(host)["reasoning_effort"] == "low" + + +# --- silent experiments --------------------------------------------------------------------- + + +async def _settle() -> None: + for _ in range(20): + await asyncio.sleep(0) + + +async def test_silent_model_mirrors_acompletion_to_the_silent_group(make_core: MakeCore, host: EndpointHost): + core: Final = make_core(model_list=[deployment("g", "d1", silent_model="shadow"), deployment("shadow", "s1")]) + await core.acompletion(model="g", messages=MESSAGES, metadata={"trace": "t"}) + await _settle() + assert sorted(host.deployments_called("acompletion")) == ["d1", "s1"] + primary: Final = next(call for call in host.calls if call.deployment_id == "d1") + shadow: Final = next(call for call in host.calls if call.deployment_id == "s1") + assert "silent_model" not in primary.kwargs + assert "is_silent_experiment" not in _bucket(primary.kwargs) + shadow_metadata: Final = _bucket(shadow.kwargs) + assert (shadow_metadata["is_silent_experiment"], shadow_metadata["model_group"], shadow_metadata["trace"]) == ( + True, + "shadow", + "t", + ) + + +async def test_silent_model_list_mirrors_to_every_listed_group(make_core: MakeCore, host: EndpointHost): + core: Final = make_core( + model_list=[ + deployment("g", "d1", silent_model=["shadow-a", "shadow-b"]), + deployment("shadow-a", "sa"), + deployment("shadow-b", "sb"), + ] + ) + await core.acompletion(model="g", messages=MESSAGES) + await _settle() + assert sorted(host.deployments_called("acompletion")) == ["d1", "sa", "sb"] + + +async def test_a_failing_silent_experiment_does_not_fail_the_primary_call(make_core: MakeCore, host: EndpointHost): + host.answer("openai/s1", litellm.BadRequestError(message="shadow broke", llm_provider="openai", model="s1")) + core: Final = make_core(model_list=[deployment("g", "d1", silent_model="shadow"), deployment("shadow", "s1")]) + response: Final = await core.acompletion(model="g", messages=MESSAGES) + await _settle() + assert response is not None + assert "s1" in host.deployments_called("acompletion") + + +async def test_a_malformed_silent_model_is_ignored(make_core: MakeCore, host: EndpointHost): + core: Final = make_core(model_list=[deployment("g", "d1", silent_model=5)]) + await core.acompletion(model="g", messages=MESSAGES) + await _settle() + assert host.deployments_called() == ["d1"] + assert "silent_model" not in _sent(host) + + +# --- timeouts ------------------------------------------------------------------------------- + + +@pytest.mark.parametrize(("stream", "expected"), [(False, (99.0, None)), (True, (99.0, 33.0))]) +async def test_service_tier_timeouts_come_from_the_deployment_tier_maps( + make_core: MakeCore, host: EndpointHost, stream: bool, expected: tuple[float, float | None] +): + core: Final = make_core( + model_list=[ + deployment( + "g", "d1", timeout=11, timeout_by_service_tier={"flex": 99}, stream_timeout_by_service_tier={"flex": 33} + ) + ] + ) + await core.acompletion(model="g", messages=MESSAGES, service_tier="flex", stream=stream) + sent: Final = _sent(host) + assert (sent["timeout"], sent.get("stream_timeout")) == expected + + +async def test_router_stream_timeout_applies_to_streaming_acompletion(make_core: MakeCore, host: EndpointHost): + core: Final = make_core(model_list=[deployment("g", "d1", timeout=11)], stream_timeout=5) + await core.acompletion(model="g", messages=MESSAGES, stream=True) + sent: Final = _sent(host) + assert (sent["timeout"], sent["stream_timeout"]) == (11.0, 5.0) + + +async def test_non_streaming_acompletion_sends_no_stream_timeout(make_core: MakeCore, host: EndpointHost): + core: Final = make_core(model_list=[deployment("g", "d1", stream_timeout=3)]) + await core.acompletion(model="g", messages=MESSAGES) + assert "stream_timeout" not in _sent(host) or _sent(host)["stream_timeout"] == 3 + + +async def test_aembedding_uses_the_request_timeout_over_the_deployment_timeout(make_core: MakeCore, host: EndpointHost): + core: Final = make_core(model_list=[deployment("g", "d1", timeout=11)]) + await core.aembedding(model="g", input=["a"], request_timeout=4) + assert _sent(host)["timeout"] == 4 + + +async def test_timeout_message_carries_deployment_timeouts_when_router_debug_is_exposed( + make_core: MakeCore, host: EndpointHost +): + host.set_dynamic(expose_router_debug_in_errors=True) + host.answer("openai/d1", timeout("d1")) + core: Final = make_core(model_list=[deployment("g", "d1", timeout=11, stream_timeout=3)], num_retries=0) + with pytest.raises(litellm.Timeout) as raised: + await core.acompletion(model="g", messages=MESSAGES) + assert "Deployment Info: request_timeout: None\ntimeout: 11.0\nstream_timeout: 3.0" in str(raised.value) + + +async def test_timeout_message_has_no_deployment_info_by_default(make_core: MakeCore, host: EndpointHost): + host.answer("openai/d1", timeout("d1")) + core: Final = make_core(model_list=[deployment("g", "d1", timeout=11)], num_retries=0) + with pytest.raises(litellm.Timeout) as raised: + await core.acompletion(model="g", messages=MESSAGES) + assert "Deployment Info" not in str(raised.value) + + +# --- clientside credentials ------------------------------------------------------------------- + + +@pytest.mark.parametrize("credential", [{"api_key": "sk-client"}, {"api_base": "https://client.example"}]) +async def test_clientside_credentials_get_a_request_scoped_deployment_id( + make_core: MakeCore, host: EndpointHost, credential: dict[str, str] +): + core: Final = make_core(model_list=[deployment("g", "d1")]) + await core.acompletion(model="g", messages=MESSAGES, **credential) + sent: Final = _sent(host) + model_info: Final = _bucket(sent, "model_info") + assert model_info["original_model_id"] == "d1" + assert isinstance(model_info["id"], str) and len(model_info["id"]) == 64 + assert _bucket(_bucket(sent), "model_info")["id"] == model_info["id"] + + +async def test_clientside_credential_id_is_stable_for_the_same_credentials(make_core: MakeCore, host: EndpointHost): + core: Final = make_core(model_list=[deployment("g", "d1")]) + await core.acompletion(model="g", messages=MESSAGES, api_key="sk-a") + await core.acompletion(model="g", messages=MESSAGES, api_key="sk-a") + await core.acompletion(model="g", messages=MESSAGES, api_key="sk-b") + ids: Final = [call.deployment_id for call in host.calls] + assert ids[0] == ids[1] != ids[2] + + +async def test_clientside_credentials_do_not_register_a_deployment(make_core: MakeCore, host: EndpointHost): + core: Final = make_core(model_list=[deployment("g", "d1")]) + await core.acompletion(model="g", messages=MESSAGES, api_key="sk-client") + await core.acompletion(model="g", messages=MESSAGES) + assert host.deployments_called()[-1] == "d1" + assert [d["model_info"]["id"] for d in core.registry().deployments("g")] == ["d1"] # pyright: ignore[reportIndexIssue, reportUnknownMemberType] # registry rows are deployment dicts + + +# --- embedding batch size -------------------------------------------------------------------- + + +@pytest.mark.parametrize("endpoint", ["aembedding", "embedding"]) +async def test_embedding_over_the_deployment_batch_size_raises_without_calling_the_host( + make_core: MakeCore, host: EndpointHost, endpoint: str +): + core: Final = make_core( + model_list=[deployment("g", "d1", model_info={"max_embedding_batch_size": 2})], num_retries=0 + ) + with pytest.raises( + litellm.BadRequestError, + match="Embedding request has 3 inputs, which exceeds max_embedding_batch_size=2 for model=g", + ): + await core.call(endpoint, {"model": "g", "input": ["a", "b", "c"]}) + assert host.calls == [] + + +@pytest.mark.parametrize("endpoint", ["aembedding", "embedding"]) +async def test_embedding_within_the_deployment_batch_size_is_sent( + make_core: MakeCore, host: EndpointHost, endpoint: str +): + core: Final = make_core(model_list=[deployment("g", "d1", model_info={"max_embedding_batch_size": 2})]) + await core.call(endpoint, {"model": "g", "input": ["a", "b"]}) + assert host.calls[0].kwargs["input"] == ["a", "b"] diff --git a/tests/unit/test_router/core/test_fallbacks.py b/tests/unit/test_router/core/test_fallbacks.py index bc6c70543c8..ef51634b920 100644 --- a/tests/unit/test_router/core/test_fallbacks.py +++ b/tests/unit/test_router/core/test_fallbacks.py @@ -9,8 +9,11 @@ import litellm from litellm.exceptions import BlockedPiiEntityError, GuardrailRaisedException from litellm.router_api import ChangeSet from litellm.router_core import RouterCore +from litellm.types.utils import Usage +from tests.unit.test_router.core.construction_precall_helpers import windowed from tests.unit.test_router.core.scripted_host import ( ScriptedHost, + agentic_followup_failure, bad_request, content_policy, context_window, @@ -19,6 +22,7 @@ from tests.unit.test_router.core.scripted_host import ( server_error, timeout, ) +from tests.unit.test_router.core.streaming_helpers import request_logging MakeCore = Callable[..., RouterCore] @@ -476,6 +480,179 @@ async def test_a_context_window_error_skips_the_order_tiers_for_its_own_chain( assert host.deployments_called("acompletion") == ["a1", "b"] +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("window_key", "tiered", "client_side", "expected"), + [ + (None, True, False, "a2"), + ("other", True, False, "a2"), + ("other", False, False, "b"), + ("other", False, True, "b"), + ("a", True, False, "c"), + ("*", True, False, "c"), + ("a", True, True, "c"), + ("a", False, True, "c"), + ], + ids=[ + "no-list-takes-the-order-tier", + "unmatched-list-takes-the-order-tier", + "unmatched-list-falls-through-to-the-general-chain", + "unmatched-list-leaves-the-client-side-list", + "matched-list-beats-the-order-tier", + "matched-star-beats-the-order-tier", + "matched-list-beats-the-order-tier-and-the-client-side-list", + "matched-list-beats-the-client-side-list", + ], +) +async def test_a_context_window_error_takes_its_own_chain_only_when_the_group_matches( + make_core: MakeCore, + host: ScriptedHost, + window_key: str | None, + tiered: bool, + client_side: bool, + expected: str, +) -> None: + core: Final = build( + make_core, + model_list=[ + deployment("a", "a1", order=1), + deployment("a" if tiered else "unused", "a2", order=2), + deployment("b", "b"), + deployment("c", "c"), + ], + fallbacks=[{"a": ["b"]}], + context_window_fallbacks=None if window_key is None else [{window_key: ["c"]}], + ) + host.answer("openai/a1", context_window()) + + await complete(core, **({"fallbacks": ["b"]} if client_side else {})) + + assert host.deployments_called("acompletion") == ["a1", expected] + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("first", "raised"), + [(context_window, litellm.ContextWindowExceededError), (server_error, litellm.Timeout)], + ids=["context-window-keeps-the-original-error", "any-other-error-raises-the-last-hop"], +) +async def test_a_failed_order_chain_raises_by_the_error_that_started_it( + make_core: MakeCore, host: ScriptedHost, first: Callable[[], Exception], raised: type[Exception] +) -> None: + core: Final = build( + make_core, model_list=TIERED_A, fallbacks=[{"a": ["b"]}], context_window_fallbacks=[{"other": ["c"]}] + ) + host.answer("openai/a1", first()) + host.answer("openai/a2", rate_limit()) + host.answer("openai/b", timeout()) + + with pytest.raises(raised): + await complete(core) + + assert host.deployments_called("acompletion") == ["a1", "a2", "b"] + + +@pytest.mark.asyncio +async def test_a_failed_order_chain_appends_no_outcome_message(make_core: MakeCore, host: ScriptedHost) -> None: + show_router_debug(host) + core: Final = build(make_core, model_list=TIERED_A) + host.answer("openai/a1", context_window()) + host.answer("openai/a2", rate_limit()) + + with pytest.raises(litellm.ContextWindowExceededError) as raised: + await complete(core) + + assert "LiteLLM: model group" not in str(raised.value) + + +@pytest.mark.asyncio +async def test_an_unmatched_context_window_list_reports_the_general_chain_it_fell_through_to( + make_core: MakeCore, host: ScriptedHost +) -> None: + show_router_debug(host) + core: Final = build(make_core, fallbacks=[{"a": ["b"]}], context_window_fallbacks=[{"other": ["c"]}]) + host.answer("openai/a", context_window()) + host.answer("openai/b", rate_limit()) + + with pytest.raises(litellm.ContextWindowExceededError) as raised: + await complete(core) + + assert "No fallback was attempted" not in str(raised.value) + assert "context_window_fallbacks=None" not in str(raised.value) + assert host.deployments_called("acompletion") == ["a", "b"] + + +@pytest.mark.asyncio +async def test_a_pre_call_context_window_error_is_kept_when_its_order_chain_fails( + make_core: MakeCore, host: ScriptedHost +) -> None: + core: Final = build( + make_core, + model_list=[windowed("a", "a1", 1, order=1), windowed("a", "a2", 1, order=2), deployment("b", "b")], + enable_pre_call_checks=True, + fallbacks=[{"a": ["b"]}], + context_window_fallbacks=[{"other": ["c"]}], + ) + host.answer("openai/b", rate_limit()) + + with pytest.raises(litellm.ContextWindowExceededError): + await complete(core) + + assert host.deployments_called("acompletion") == ["b"] + + +@pytest.mark.asyncio +async def test_weighted_failover_skips_a_context_window_error_with_an_unmatched_list( + make_core: MakeCore, host: ScriptedHost +) -> None: + core: Final = build( + make_core, + model_list=[deployment("a", "a1"), deployment("a", "a2"), deployment("b", "b")], + enable_weighted_failover=True, + fallbacks=[{"a": ["b"]}], + context_window_fallbacks=[{"other": ["c"]}], + ) + host.answer("openai/a1", context_window()) + host.answer("openai/a2", context_window()) + + await complete(core) + + called: Final = host.deployments_called("acompletion") + assert len(called) == 2 + assert called[-1] == "b" + + +# --- a failed agentic follow-up ------------------------------------------------------------------ + + +@pytest.mark.asyncio +async def test_a_failed_agentic_followup_still_falls_back(make_core: MakeCore, host: ScriptedHost) -> None: + core: Final = build(make_core, fallbacks=[{"a": ["b"]}], num_retries=3) + host.answer("openai/a", lambda **_: agentic_followup_failure(server_error())) + + response: Final = await complete(core) + + assert host.deployments_called("acompletion") == ["a", "b"] + assert headers(response)["x-litellm-model-group"] == "b" + + +@pytest.mark.asyncio +async def test_a_fallback_hop_carries_the_superseded_attempt_usage(make_core: MakeCore, host: ScriptedHost) -> None: + logging_obj: Final = request_logging() + usage: Final = Usage(prompt_tokens=1, completion_tokens=1, total_tokens=2) + + def failed_after_usage(**_: object) -> Exception: + logging_obj.model_call_details["combined_usage_object"] = usage + return server_error() + + core: Final = build(make_core, fallbacks=[{"a": ["b"]}]) + host.answer("openai/a", failed_after_usage) + + await complete(core, litellm_logging_obj=logging_obj) + + assert logging_obj.model_call_details["failed_attempts_usage"] == ((usage, 0.0),) + + # --- when no fallback runs ----------------------------------------------------------------------- diff --git a/tests/unit/test_router/core/test_fallbacks_characterization.py b/tests/unit/test_router/core/test_fallbacks_characterization.py new file mode 100644 index 00000000000..41a751358a9 --- /dev/null +++ b/tests/unit/test_router/core/test_fallbacks_characterization.py @@ -0,0 +1,584 @@ +""" +Characterization of a router core's fallbacks (non-streaming): chain sources and formats, +client-side fallback lists, per-fallback params, validation, mock-testing flags, depth bounds, +default fallbacks and wildcard lookup. Complements ``test_fallbacks.py``. +""" + +import pytest + +import litellm +from tests.unit.test_router.core.fallbacks_helpers import ( + TIERED_A, + MakeCore, + build, + complete, + complete_sync, + headers, + show_router_debug, +) +from tests.unit.test_router.core.scripted_host import ( + ScriptedHost, + content_policy, + context_window, + deployment, + rate_limit, +) + +# --- validation of the configured chains --------------------------------------------------------- + + +@pytest.mark.parametrize( + ("field", "value", "error", "message"), + [ + ("fallbacks", ["b"], ValueError, "Item 'b' is not a dictionary."), + ( + "fallbacks", + [{"a": ["b"], "c": ["b"]}], + ValueError, + "Dictionary '{'a': ['b'], 'c': ['b']}' must have exactly one key, but has 2 keys.", + ), + ( + "fallbacks", + [{"a": "b"}], + TypeError, + "Fallbacks for 'a' must be a list of model names, got 'b'. Use {'a': ['b']}", + ), + ( + "content_policy_fallbacks", + [{"a": "b"}], + TypeError, + "Fallbacks for 'a' must be a list of model names, got 'b'. Use {'a': ['b']}", + ), + ], + ids=["plain-string-item", "two-keys", "string-target", "content-policy-string-target"], +) +def test_a_malformed_fallbacks_configuration_is_refused_at_construction( + make_core: MakeCore, field: str, value: object, error: type[Exception], message: str +) -> None: + with pytest.raises(error) as raised: + build(make_core, **{field: value}) + + assert str(raised.value) == message + + +@pytest.mark.asyncio +async def test_context_window_fallbacks_are_not_validated_at_construction( + make_core: MakeCore, host: ScriptedHost +) -> None: + core = build(make_core, context_window_fallbacks=[{"a": "b"}]) + + await complete(core) + + assert host.deployments_called("acompletion") == ["a"] + + +@pytest.mark.asyncio +async def test_a_constructor_chain_of_parameter_dicts_passes_validation_and_applies_to_any_group( + make_core: MakeCore, host: ScriptedHost +) -> None: + core = build(make_core, fallbacks=[{"model": "b"}]) + host.answer("openai/a", rate_limit()) + + await complete(core) + + assert host.deployments_called("acompletion") == ["a", "b"] + + +@pytest.mark.asyncio +async def test_an_empty_constructor_chain_defers_to_the_process_setting( + make_core: MakeCore, host: ScriptedHost +) -> None: + host.set_dynamic(fallbacks=[{"a": ["c"]}]) + core = build(make_core, fallbacks=[]) + host.answer("openai/a", rate_limit()) + + await complete(core) + + assert host.deployments_called("acompletion") == ["a", "c"] + + +# --- client-side fallback lists ------------------------------------------------------------------ + + +@pytest.mark.asyncio +async def test_a_plain_request_list_is_walked_in_order(make_core: MakeCore, host: ScriptedHost) -> None: + core = build(make_core) + host.answer("openai/a", rate_limit()) + host.answer("openai/b", rate_limit()) + + response = await complete(core, fallbacks=["b", "c"]) + + assert host.deployments_called("acompletion") == ["a", "b", "c"] + assert headers(response)["x-litellm-model-group"] == "c" + + +@pytest.mark.asyncio +async def test_the_context_window_chain_beats_a_plain_request_list(make_core: MakeCore, host: ScriptedHost) -> None: + core = build(make_core, context_window_fallbacks=[{"a": ["c"]}]) + host.answer("openai/a", context_window()) + + await complete(core, fallbacks=["b"]) + + assert host.deployments_called("acompletion") == ["a", "c"] + + +@pytest.mark.asyncio +async def test_a_dict_target_with_params_overrides_the_request_for_that_attempt( + make_core: MakeCore, host: ScriptedHost +) -> None: + core = build(make_core) + host.answer("openai/a", rate_limit()) + override = [{"role": "user", "content": "shorter"}] + + await complete(core, fallbacks=[{"model": "b", "temperature": 0.3, "messages": override}]) + + primary, fallback = host.calls_to("acompletion") + assert (primary.deployment_id, primary.kwargs.get("temperature")) == ("a", None) + assert primary.kwargs["messages"] == [{"role": "user", "content": "hi"}] + assert (fallback.deployment_id, fallback.kwargs["temperature"], fallback.kwargs["messages"]) == ("b", 0.3, override) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "arguments", + [ + {"request": [{"model": "b", "temperature": 0.3}, {"model": "c"}]}, + {"constructor": [{"a": [{"model": "b", "temperature": 0.3}, "c"]}]}, + ], + ids=["request-list", "constructor-chain"], +) +async def test_a_dict_targets_params_carry_over_to_the_next_target( + make_core: MakeCore, host: ScriptedHost, arguments: dict[str, list[object]] +) -> None: + core = build(make_core, **({"fallbacks": arguments["constructor"]} if "constructor" in arguments else {})) + host.answer("openai/a", rate_limit()) + host.answer("openai/b", rate_limit()) + + await complete(core, **({"fallbacks": arguments["request"]} if "request" in arguments else {})) + + assert [(c.deployment_id, c.kwargs.get("temperature")) for c in host.calls_to("acompletion")] == [ + ("a", None), + ("b", 0.3), + ("c", 0.3), + ] + + +@pytest.mark.asyncio +async def test_a_dict_target_naming_the_failed_group_with_params_retries_that_group( + make_core: MakeCore, host: ScriptedHost +) -> None: + core = build(make_core) + host.answer("openai/a", rate_limit(), None) + + response = await complete(core, fallbacks=[{"model": "a", "temperature": 0}]) + + assert [(c.deployment_id, c.kwargs.get("temperature")) for c in host.calls_to("acompletion")] == [ + ("a", None), + ("a", 0), + ] + assert headers(response)["x-litellm-attempted-fallbacks"] == 1 + + +@pytest.mark.asyncio +async def test_a_dict_target_with_its_own_api_key_runs_on_a_client_credential_deployment( + make_core: MakeCore, host: ScriptedHost +) -> None: + core = build(make_core) + host.answer("openai/a", rate_limit()) + + response = await complete(core, fallbacks=[{"model": "b", "api_key": "sk-other"}]) + + primary, fallback = host.calls_to("acompletion") + assert (fallback.model, fallback.kwargs["api_key"]) == ("openai/b", "sk-other") + assert fallback.deployment_id not in {None, "b"} + assert headers(response)["x-litellm-model-group"] == "b" + assert primary.deployment_id == "a" + + +@pytest.mark.asyncio +async def test_a_dict_target_carrying_an_oidc_api_key_is_refused_and_the_primary_error_raised( + make_core: MakeCore, host: ScriptedHost +) -> None: + core = build(make_core) + show_router_debug(host) + host.answer("openai/a", rate_limit()) + + with pytest.raises(litellm.RateLimitError) as raised: + await complete(core, fallbacks=[{"model": "b", "api_key": "oidc/google/x"}]) + + assert host.deployments_called("acompletion") == ["a"] + assert str(raised.value) == ( + "litellm.RateLimitError: slow down\n\nLiteLLM: model group 'a' failed with the error above. " + "Fallback to {'model': 'b', 'api_key': 'oidc*********'} also failed: Rejected Request: api_key holds an " + "oidc/ secret reference, which only a deployment or credential a proxy admin configures can carry." + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("fallbacks", "expected"), + [(["c", {"a": ["b"]}], ["a", "b", "c"]), (["b", {"x": ["c"]}], ["a", "b"])], + ids=["keyed-entry-wins-then-string-catches-the-next-hop", "string-entry-is-a-chain-for-every-group"], +) +async def test_a_mixed_list_of_names_and_keyed_entries_is_resolved_as_a_chain_map( + make_core: MakeCore, host: ScriptedHost, fallbacks: list[object], expected: list[str] +) -> None: + core = build(make_core) + host.answer("*", rate_limit()) + + with pytest.raises(litellm.RateLimitError): + await complete(core, fallbacks=fallbacks) + + assert host.deployments_called("acompletion") == expected + + +@pytest.mark.asyncio +async def test_the_request_fallbacks_kwarg_is_not_forwarded_to_the_provider( + make_core: MakeCore, host: ScriptedHost +) -> None: + core = build(make_core) + host.answer("openai/a", rate_limit()) + + await complete(core, fallbacks=["b"]) + + assert [("fallbacks" in c.kwargs) for c in host.calls_to("acompletion")] == [False, False] + + +# --- disable_fallbacks ------------------------------------------------------------------------------ + + +@pytest.mark.asyncio +async def test_disable_fallbacks_given_as_a_string_does_not_disable_them( + make_core: MakeCore, host: ScriptedHost +) -> None: + core = build(make_core, fallbacks=[{"a": ["b"]}]) + host.answer("openai/a", rate_limit()) + + await complete(core, disable_fallbacks="true") + + assert host.deployments_called("acompletion") == ["a", "b"] + + +@pytest.mark.asyncio +async def test_disable_fallbacks_also_stops_the_order_tier_walk(make_core: MakeCore, host: ScriptedHost) -> None: + core = build(make_core, model_list=TIERED_A) + host.answer("*", rate_limit()) + + with pytest.raises(litellm.RateLimitError): + await complete(core, disable_fallbacks=True) + + assert host.deployments_called("acompletion") == ["a1"] + + +@pytest.mark.asyncio +async def test_disable_fallbacks_is_not_forwarded_to_the_provider(make_core: MakeCore, host: ScriptedHost) -> None: + core = build(make_core) + + await complete(core, disable_fallbacks=True) + + assert "disable_fallbacks" not in host.calls_to("acompletion")[0].kwargs + + +# --- depth bounds ----------------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_a_request_max_fallbacks_overrides_the_routers(make_core: MakeCore, host: ScriptedHost) -> None: + core = build(make_core, fallbacks=[{"a": ["b"]}, {"b": ["c"]}]) + host.answer("*", rate_limit()) + + with pytest.raises(litellm.RateLimitError): + await complete(core, max_fallbacks=1) + + assert host.deployments_called("acompletion") == ["a", "b"] + + +@pytest.mark.asyncio +async def test_max_fallbacks_zero_also_stops_the_order_tier_walk(make_core: MakeCore, host: ScriptedHost) -> None: + core = build(make_core, model_list=TIERED_A, max_fallbacks=0) + host.answer("*", rate_limit()) + + with pytest.raises(litellm.RateLimitError): + await complete(core) + + assert host.deployments_called("acompletion") == ["a1"] + + +@pytest.mark.asyncio +async def test_max_fallbacks_one_still_tries_every_target_of_the_first_level( + make_core: MakeCore, host: ScriptedHost +) -> None: + core = build(make_core, fallbacks=[{"a": ["b", "c"]}, {"b": ["c"]}], max_fallbacks=1) + host.answer("*", rate_limit()) + + with pytest.raises(litellm.RateLimitError): + await complete(core) + + assert host.deployments_called("acompletion") == ["a", "b", "c"] + + +@pytest.mark.asyncio +async def test_a_fallback_group_spends_its_own_retries(make_core: MakeCore, host: ScriptedHost) -> None: + core = build(make_core, fallbacks=[{"a": ["b"]}], num_retries=1) + host.answer("*", rate_limit()) + + with pytest.raises(litellm.RateLimitError) as raised: + await complete(core) + + assert host.deployments_called("acompletion") == ["a", "a", "b", "b"] + assert str(raised.value) == "litellm.RateLimitError: slow down LiteLLM Retried: 1 times, LiteLLM Max Retries: 1" + + +@pytest.mark.asyncio +async def test_a_fallback_hop_resumes_the_original_groups_chain(make_core: MakeCore, host: ScriptedHost) -> None: + core = build(make_core, fallbacks=[{"a": ["b", "c"]}], max_fallbacks=5) + host.answer("openai/a", rate_limit()) + host.answer("openai/b", rate_limit()) + + response = await complete(core, include_fallback_errors=True) + + assert host.deployments_called("acompletion") == ["a", "b", "c"] + # c answered inside b's hop (b resumed a's chain), so b's error is listed before a's. + assert headers(response)["x-litellm-fallback-errors"] == ( + '[{"message": "litellm.RateLimitError: slow down", "type": "throttling_error", "param": null, "code": "429"}, ' + '{"message": "litellm.RateLimitError: slow down", "type": "throttling_error", "param": null, "code": "429"}]' + ) + + +# --- order tiers with other chains --------------------------------------------------------------- + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "arguments", + [{"request": {"fallbacks": ["c"]}}, {"constructor": {"default_fallbacks": ["c"]}}], + ids=["plain-request-list", "default-fallbacks"], +) +async def test_order_tiers_run_before_a_plain_list_or_the_default_chain( + make_core: MakeCore, host: ScriptedHost, arguments: dict[str, dict[str, object]] +) -> None: + core = build(make_core, model_list=TIERED_A, **arguments.get("constructor", {})) + host.answer("*", rate_limit()) + host.answer("openai/c", None) + + response = await complete(core, **arguments.get("request", {})) + + assert host.deployments_called("acompletion") == ["a1", "a2", "c"] + assert headers(response)["x-litellm-model-group"] == "c" + + +# --- default fallbacks --------------------------------------------------------------------------- + + +def test_default_fallbacks_are_appended_to_the_process_fallbacks_list_in_place( + make_core: MakeCore, host: ScriptedHost +) -> None: + process_fallbacks: list[object] = [{"a": ["b"]}] + host.set_dynamic(fallbacks=process_fallbacks) + + core = build(make_core, default_fallbacks=["c"]) + + assert core.settings().fallbacks == [{"a": ["b"]}, {"*": ["c"]}] + assert process_fallbacks == [{"a": ["b"]}, {"*": ["c"]}] + + +def test_default_fallbacks_leave_the_constructor_fallbacks_list_unchanged(make_core: MakeCore) -> None: + fallbacks: list[object] = [{"a": ["b"]}] + + core = build(make_core, fallbacks=fallbacks, default_fallbacks=["c"]) + + assert core.settings().fallbacks == [{"a": ["b"]}, {"*": ["c"]}] + assert fallbacks == [{"a": ["b"]}] + + +@pytest.mark.asyncio +async def test_an_unknown_group_is_served_by_the_first_default_fallback_without_a_fallback_hop( + make_core: MakeCore, host: ScriptedHost +) -> None: + core = build(make_core, default_fallbacks=["b", "c"]) + + response = await complete(core, model="unknown") + + assert host.deployments_called("acompletion") == ["b"] + assert headers(response)["x-litellm-attempted-fallbacks"] == 0 + assert headers(response)["x-litellm-model-group"] == "unknown" + assert host.calls_to("acompletion")[0].model_group == "unknown" + + +@pytest.mark.asyncio +async def test_an_unknown_group_with_its_own_chain_falls_back_through_it( + make_core: MakeCore, host: ScriptedHost +) -> None: + core = build(make_core, fallbacks=[{"unknown": ["b"]}]) + + response = await complete(core, model="unknown") + + assert host.deployments_called("acompletion") == ["b"] + assert headers(response)["x-litellm-attempted-fallbacks"] == 1 + assert headers(response)["x-litellm-model-group"] == "b" + + +@pytest.mark.asyncio +async def test_default_fallbacks_catch_a_content_policy_error_without_a_content_policy_chain( + make_core: MakeCore, host: ScriptedHost +) -> None: + core = build(make_core, default_fallbacks=["c"]) + host.answer("openai/a", content_policy()) + + await complete(core) + + assert host.deployments_called("acompletion") == ["a", "c"] + + +@pytest.mark.asyncio +async def test_a_request_context_window_fallbacks_of_none_falls_through_to_the_general_chain( + make_core: MakeCore, host: ScriptedHost +) -> None: + core = build(make_core, context_window_fallbacks=[{"a": ["c"]}], fallbacks=[{"a": ["b"]}]) + host.answer("openai/a", context_window()) + + await complete(core, context_window_fallbacks=None) + + assert host.deployments_called("acompletion") == ["a", "b"] + + +# --- wildcard groups ----------------------------------------------------------------------------- + +WILDCARD = [deployment("openai/*", "w", model="openai/*"), deployment("b", "b")] + + +@pytest.mark.asyncio +async def test_a_wildcard_routed_request_uses_the_chain_keyed_by_the_bare_model_name( + make_core: MakeCore, host: ScriptedHost +) -> None: + core = build(make_core, model_list=WILDCARD, fallbacks=[{"gpt-x": ["b"]}]) + host.answer("openai/gpt-x", rate_limit()) + + response = await complete(core, model="openai/gpt-x") + + assert host.deployments_called("acompletion") == ["w", "b"] + assert [c.model for c in host.calls_to("acompletion")] == ["openai/gpt-x", "openai/b"] + assert headers(response)["x-litellm-model-group"] == "b" + + +@pytest.mark.asyncio +async def test_a_chain_keyed_by_the_wildcard_pattern_does_not_match_a_routed_request( + make_core: MakeCore, host: ScriptedHost +) -> None: + core = build(make_core, model_list=WILDCARD, fallbacks=[{"openai/*": ["b"]}]) + host.answer("openai/gpt-x", rate_limit()) + + with pytest.raises(litellm.RateLimitError): + await complete(core, model="openai/gpt-x") + + assert host.deployments_called("acompletion") == ["w"] + + +# --- mock testing flags ------------------------------------------------------------------------ + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("flag", "field", "expected"), + [ + ("mock_testing_fallbacks", "fallbacks", "b"), + ("mock_testing_context_fallbacks", "context_window_fallbacks", "c"), + ("mock_testing_content_policy_fallbacks", "content_policy_fallbacks", "c"), + ], + ids=["general", "context-window", "content-policy"], +) +async def test_a_mock_testing_flag_skips_the_primary_and_walks_the_matching_chain( + make_core: MakeCore, host: ScriptedHost, flag: str, field: str, expected: str +) -> None: + core = build(make_core, num_retries=2, **{"fallbacks": [{"a": ["b"]}], field: [{"a": [expected]}]}) + + response = await complete(core, **{flag: True}) + + assert host.deployments_called("acompletion") == [expected] + assert headers(response)["x-litellm-attempted-fallbacks"] == 1 + assert headers(response)["x-litellm-attempted-retries"] == 0 + assert host.sleeps == [] + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("flag", "error", "message"), + [ + ( + "mock_testing_fallbacks", + litellm.InternalServerError, + "litellm.InternalServerError: This is a mock exception for model=a, to trigger a fallback. Fallbacks=None", + ), + ( + "mock_testing_context_fallbacks", + litellm.ContextWindowExceededError, + "litellm.ContextWindowExceededError: litellm.BadRequestError: This is a mock exception for model=a, " + "to trigger a fallback. Context_Window_Fallbacks=None", + ), + ( + "mock_testing_content_policy_fallbacks", + litellm.ContentPolicyViolationError, + "litellm.BadRequestError: litellm.ContentPolicyViolationError: This is a mock exception for model=a, " + "to trigger a fallback. Context_Policy_Fallbacks=None", + ), + ], + ids=["general", "context-window", "content-policy"], +) +async def test_a_mock_testing_flag_without_a_chain_raises_the_mock_error_without_calling_the_provider( + make_core: MakeCore, host: ScriptedHost, flag: str, error: type[Exception], message: str +) -> None: + core = build(make_core) + + with pytest.raises(error) as raised: + await complete(core, **{flag: True}) + + assert str(raised.value) == message + assert host.calls == [] + + +@pytest.mark.asyncio +async def test_the_mock_error_names_the_configured_chain_and_the_fallback_outcome( + make_core: MakeCore, host: ScriptedHost +) -> None: + core = build(make_core, fallbacks=[{"a": ["b"]}]) + show_router_debug(host) + host.answer("*", rate_limit()) + + with pytest.raises(litellm.InternalServerError) as raised: + await complete(core, mock_testing_fallbacks=True) + + assert str(raised.value) == ( + "litellm.InternalServerError: This is a mock exception for model=a, to trigger a fallback. " + "Fallbacks=[{'a': ['b']}]\n\nLiteLLM: model group 'a' failed with the error above. " + "Fallback to b also failed: litellm.RateLimitError: slow down" + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize(("value", "expected"), [("true", ["b"]), ("false", ["a"]), (False, ["a"])]) +async def test_a_mock_testing_flag_given_as_a_string_is_parsed_as_a_bool( + make_core: MakeCore, host: ScriptedHost, value: object, expected: list[str] +) -> None: + core = build(make_core, fallbacks=[{"a": ["b"]}]) + + await complete(core, mock_testing_fallbacks=value) + + assert host.deployments_called("acompletion") == expected + + +@pytest.mark.asyncio +async def test_the_mock_testing_flag_is_not_forwarded_to_the_fallback(make_core: MakeCore, host: ScriptedHost) -> None: + core = build(make_core, fallbacks=[{"a": ["b"]}]) + + await complete(core, mock_testing_fallbacks=True) + + assert "mock_testing_fallbacks" not in host.calls_to("acompletion")[0].kwargs + + +def test_sync_completion_honours_the_mock_testing_flag(make_core: MakeCore, host: ScriptedHost) -> None: + core = build(make_core, fallbacks=[{"a": ["b"]}]) + + response = complete_sync(core, mock_testing_fallbacks=True) + + assert host.deployments_called("completion") == ["b"] + assert headers(response)["x-litellm-attempted-fallbacks"] == 1 diff --git a/tests/unit/test_router/core/test_fallbacks_characterization_outcomes.py b/tests/unit/test_router/core/test_fallbacks_characterization_outcomes.py new file mode 100644 index 00000000000..b5589e24cf7 --- /dev/null +++ b/tests/unit/test_router/core/test_fallbacks_characterization_outcomes.py @@ -0,0 +1,690 @@ +""" +Characterization of what a router core's fallbacks leave behind (non-streaming): response headers +and hidden params, the metadata each attempt carries, the error text when fallbacks give up, +fallback access and budget checks, fallback events, content-filter responses and weighted failover. +""" + +from collections.abc import Mapping + +import pytest + +import litellm +from litellm.integrations.custom_logger import CustomLogger +from tests.unit.test_router.core.fallbacks_helpers import ( + GROUP_A_OF_THREE, + MakeCore, + build, + complete, + complete_sync, + fails_once_then_answers, + headers, + hidden, + metadata_of, + show_router_debug, +) +from tests.unit.test_router.core.scripted_host import ( + ScriptedHost, + completion_response, + context_window, + deployment, + rate_limit, + server_error, + timeout, +) + +ERROR_429 = '{"message": "litellm.RateLimitError: slow down", "type": "throttling_error", "param": null, "code": "429"}' +ERROR_500 = ( + '{"message": "litellm.InternalServerError: boom", "type": "internal_server_error", "param": null, "code": "500"}' +) + +# --- headers and hidden params ------------------------------------------------------------------ + + +@pytest.mark.asyncio +async def test_a_two_hop_chain_reports_one_attempted_fallback(make_core: MakeCore, host: ScriptedHost) -> None: + core = build(make_core, fallbacks=[{"a": ["b"]}, {"b": ["c"]}]) + host.answer("openai/a", rate_limit()) + host.answer("openai/b", rate_limit()) + + response = await complete(core) + + assert host.deployments_called("acompletion") == ["a", "b", "c"] + assert headers(response)["x-litellm-attempted-fallbacks"] == 1 + assert headers(response)["x-litellm-model-group"] == "c" + + +@pytest.mark.asyncio +async def test_the_answering_fallback_deployment_is_the_responses_model_id( + make_core: MakeCore, host: ScriptedHost +) -> None: + core = build(make_core, fallbacks=[{"a": ["b"]}]) + host.answer("openai/a", rate_limit()) + + response = await complete(core) + + assert hidden(response)["model_id"] == "b" + + +@pytest.mark.asyncio +async def test_success_headers_after_a_fallback(make_core: MakeCore, host: ScriptedHost) -> None: + core = build(make_core, fallbacks=[{"a": ["b"]}]) + host.answer("openai/a", rate_limit()) + + response = await complete(core) + + assert headers(response) == { + "x-litellm-model-group": "b", + "x-litellm-attempted-retries": 0, + "x-litellm-attempted-fallbacks": 1, + } + + +@pytest.mark.asyncio +async def test_include_fallback_errors_lists_the_primary_error_in_a_header( + make_core: MakeCore, host: ScriptedHost +) -> None: + core = build(make_core, fallbacks=[{"a": ["b"]}]) + host.answer("openai/a", rate_limit()) + + response = await complete(core, include_fallback_errors=True) + + assert headers(response)["x-litellm-fallback-errors"] == f"[{ERROR_429}]" + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "fallbacks", + [[{"a": ["b", "c"]}], [{"a": ["b"]}, {"b": ["c"]}]], + ids=["one-list", "chained"], +) +async def test_include_fallback_errors_lists_the_inner_hops_errors_first( + make_core: MakeCore, host: ScriptedHost, fallbacks: list[object] +) -> None: + core = build(make_core, fallbacks=fallbacks) + host.answer("openai/a", rate_limit()) + host.answer("openai/b", server_error()) + + response = await complete(core, include_fallback_errors=True) + + assert host.deployments_called("acompletion") == ["a", "b", "c"] + assert headers(response)["x-litellm-fallback-errors"] == f"[{ERROR_500}, {ERROR_429}]" + + +@pytest.mark.asyncio +async def test_without_include_fallback_errors_no_errors_header_is_set(make_core: MakeCore, host: ScriptedHost) -> None: + core = build(make_core, fallbacks=[{"a": ["b"]}]) + host.answer("openai/a", rate_limit()) + + response = await complete(core) + + assert "x-litellm-fallback-errors" not in headers(response) + + +# --- what each attempt carries to the host ------------------------------------------------------------ + + +@pytest.mark.asyncio +async def test_each_attempt_carries_the_fallback_count_and_original_group_in_its_metadata( + make_core: MakeCore, host: ScriptedHost +) -> None: + core = build(make_core, fallbacks=[{"a": ["b"]}, {"b": ["c"]}]) + host.answer("openai/a", rate_limit()) + host.answer("openai/b", rate_limit()) + + await complete(core) + + assert [ + (meta["model_group"], meta["original_model_group"], meta["attempted_fallbacks"]) + for meta in (metadata_of(call.kwargs) for call in host.calls_to("acompletion")) + ] == [("a", "a", 0), ("b", "a", 1), ("c", "a", 2)] + + +@pytest.mark.asyncio +async def test_a_fallback_attempt_also_receives_the_routers_walk_state_as_kwargs( + make_core: MakeCore, host: ScriptedHost +) -> None: + core = build(make_core, fallbacks=[{"a": ["b"]}]) + host.answer("openai/a", rate_limit()) + + await complete(core) + + primary, fallback = host.calls_to("acompletion") + assert {"fallback_depth", "max_fallbacks", "attempted_targets"} & set(primary.kwargs) == set() + assert (fallback.kwargs["fallback_depth"], fallback.kwargs["max_fallbacks"]) == (1, 5) + assert "attempted_targets" in fallback.kwargs + + +@pytest.mark.asyncio +async def test_the_fallback_attempt_sees_the_failed_attempt_in_previous_models( + make_core: MakeCore, host: ScriptedHost +) -> None: + core = build(make_core, fallbacks=[{"a": ["b"]}]) + host.answer("openai/a", rate_limit()) + + await complete(core) + + previous = metadata_of(host.calls_to("acompletion")[1].kwargs)["previous_models"] + assert [ + (entry["model_group"], entry["deployment_id"], entry["exception_type"]) + for entry in previous # pyright: ignore[reportGeneralTypeIssues] # a tuple of dicts + ] == [("a", "a", "RateLimitError")] + + +# --- the error when fallbacks give up ------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_router_debug_on_a_context_window_error_without_its_chain_names_the_missing_setting( + make_core: MakeCore, host: ScriptedHost +) -> None: + core = build(make_core) + show_router_debug(host) + host.answer("openai/a", context_window()) + + with pytest.raises(litellm.ContextWindowExceededError) as raised: + await complete(core) + + assert str(raised.value) == ( + "litellm.ContextWindowExceededError: litellm.BadRequestError: too long\n" + "model=a. context_window_fallbacks=None. fallbacks=None.\n\n" + "Set 'context_window_fallback' - https://docs.litellm.ai/docs/routing#fallbacks\n\n" + "LiteLLM: model group 'a' failed with the error above. No fallback was attempted." + ) + + +@pytest.mark.asyncio +async def test_router_debug_on_a_context_window_error_that_also_fails_its_fallback_nests_the_hint( + make_core: MakeCore, host: ScriptedHost +) -> None: + core = build(make_core, fallbacks=[{"a": ["b"]}]) + show_router_debug(host) + host.answer("*", context_window()) + + with pytest.raises(litellm.ContextWindowExceededError) as raised: + await complete(core) + + hint = "context_window_fallbacks=None. fallbacks=[{'a': ['b']}].\n\nSet 'context_window_fallback' - https://docs.litellm.ai/docs/routing#fallbacks" + assert str(raised.value) == ( + f"litellm.ContextWindowExceededError: litellm.BadRequestError: too long\nmodel=a. {hint}\n\n" + "LiteLLM: model group 'a' failed with the error above. Fallback to b also failed: " + f"litellm.ContextWindowExceededError: litellm.BadRequestError: too long\nmodel=b. {hint}" + ) + + +@pytest.mark.asyncio +async def test_router_debug_names_a_plain_request_list_that_failed(make_core: MakeCore, host: ScriptedHost) -> None: + core = build(make_core) + show_router_debug(host) + host.answer("*", rate_limit()) + + with pytest.raises(litellm.RateLimitError) as raised: + await complete(core, fallbacks=["b", "c"]) + + assert str(raised.value) == ( + "litellm.RateLimitError: slow down\n\nLiteLLM: model group 'a' failed with the error above. " + "Fallback to b, c also failed: litellm.RateLimitError: slow down" + ) + + +@pytest.mark.asyncio +async def test_router_debug_with_max_fallbacks_one_names_only_the_first_level( + make_core: MakeCore, host: ScriptedHost +) -> None: + core = build(make_core, fallbacks=[{"a": ["b"]}, {"b": ["c"]}], max_fallbacks=1) + show_router_debug(host) + host.answer("*", rate_limit()) + + with pytest.raises(litellm.RateLimitError) as raised: + await complete(core) + + assert host.deployments_called("acompletion") == ["a", "b"] + assert str(raised.value).endswith("Fallback to b also failed: litellm.RateLimitError: slow down") + + +@pytest.mark.asyncio +async def test_router_debug_reports_the_first_fallback_errors_message_not_the_last( + make_core: MakeCore, host: ScriptedHost +) -> None: + core = build(make_core, fallbacks=[{"a": ["b"]}, {"b": ["c"]}]) + show_router_debug(host) + host.answer("openai/a", rate_limit()) + host.answer("openai/b", server_error()) + host.answer("openai/c", timeout()) + + with pytest.raises(litellm.RateLimitError) as raised: + await complete(core) + + assert str(raised.value) == ( + "litellm.RateLimitError: slow down\n\nLiteLLM: model group 'a' failed with the error above. " + "Fallback to b also failed: litellm.InternalServerError: boom" + ) + + +@pytest.mark.asyncio +async def test_router_debug_truncates_a_long_fallback_error(make_core: MakeCore, host: ScriptedHost) -> None: + core = build(make_core, fallbacks=[{"a": ["b"]}]) + show_router_debug(host) + host.answer("openai/a", rate_limit()) + host.answer("openai/b", rate_limit(message="x" * 3000)) + + with pytest.raises(litellm.RateLimitError) as raised: + await complete(core) + + detail = "litellm.RateLimitError: " + "x" * 3000 + assert str(raised.value) == ( + "litellm.RateLimitError: slow down\n\nLiteLLM: model group 'a' failed with the error above. " + f"Fallback to b also failed: {detail[:2000]}... [truncated {len(detail) - 2000} characters]" + ) + + +@pytest.mark.asyncio +async def test_router_debug_on_an_unmatched_content_policy_chain_says_no_fallback_was_attempted( + make_core: MakeCore, host: ScriptedHost +) -> None: + from tests.unit.test_router.core.scripted_host import content_policy + + core = build(make_core, content_policy_fallbacks=[{"other": ["b"]}]) + show_router_debug(host) + host.answer("openai/a", content_policy()) + + with pytest.raises(litellm.ContentPolicyViolationError) as raised: + await complete(core) + + assert str(raised.value) == ( + "litellm.BadRequestError: litellm.ContentPolicyViolationError: blocked\n\n" + "LiteLLM: model group 'a' failed with the error above. No fallback was attempted." + ) + + +@pytest.mark.asyncio +async def test_router_debug_no_chain_message_without_any_configured_keys( + make_core: MakeCore, host: ScriptedHost +) -> None: + core = build(make_core) + show_router_debug(host) + host.answer("openai/a", rate_limit()) + + with pytest.raises(litellm.RateLimitError) as raised: + await complete(core, fallbacks=[{"x": ["b"]}, {"y": ["c"]}]) + + assert str(raised.value).endswith( + "LiteLLM: model group 'a' failed with the error above and no fallback model group was found for it, " + "so the request was not retried on another model. Fallbacks are configured for: x, y. " + "Add a fallbacks entry for that model group (Router fallbacks or proxy router_settings.fallbacks) " + "to retry on another model." + ) + + +# --- fallback access and budget checks ------------------------------------------------------------- + + +def denying(denied: str, seen: list[str]): # noqa: ANN201 # a FallbackAccessCheck / FallbackBudgetCheck + async def check(*, model: str, request_kwargs: Mapping[str, object], ctx: object) -> bool: + seen.append(model) + return model != denied + + return check + + +@pytest.mark.asyncio +@pytest.mark.parametrize("field", ["fallback_access_check", "fallback_budget_check"]) +async def test_a_fallback_target_the_check_rejects_is_skipped( + make_core: MakeCore, host: ScriptedHost, field: str +) -> None: + seen: list[str] = [] + core = build(make_core, fallbacks=[{"a": ["b", "c"]}], **{field: denying("b", seen)}) + host.answer("openai/a", rate_limit()) + + response = await complete(core) + + assert host.deployments_called("acompletion") == ["a", "c"] + assert headers(response)["x-litellm-model-group"] == "c" + assert seen == ["b", "c"] + + +def test_sync_completion_runs_the_fallback_access_check(make_core: MakeCore, host: ScriptedHost) -> None: + seen: list[str] = [] + core = build(make_core, fallbacks=[{"a": ["b", "c"]}], fallback_access_check=denying("b", seen)) + host.answer("openai/a", rate_limit()) + + complete_sync(core) + + assert host.deployments_called("completion") == ["a", "c"] + assert seen == ["b", "c"] + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("field", "sentence"), + [ + ("fallback_access_check", "Fallback to b was skipped because the caller is not authorized to call it."), + ("fallback_budget_check", "Fallback to b was skipped because the caller is over budget for it."), + ], +) +async def test_router_debug_explains_a_skipped_fallback_target( + make_core: MakeCore, host: ScriptedHost, field: str, sentence: str +) -> None: + core = build(make_core, fallbacks=[{"a": ["b"]}], **{field: denying("b", [])}) + show_router_debug(host) + host.answer("openai/a", rate_limit()) + + with pytest.raises(litellm.RateLimitError) as raised: + await complete(core) + + assert host.deployments_called("acompletion") == ["a"] + assert str(raised.value) == ( + f"litellm.RateLimitError: slow down\n\nLiteLLM: model group 'a' failed with the error above. {sentence}" + ) + + +@pytest.mark.asyncio +async def test_router_debug_names_both_the_failed_and_the_skipped_targets( + make_core: MakeCore, host: ScriptedHost +) -> None: + seen: list[str] = [] + core = build( + make_core, + fallbacks=[{"a": ["b", "c"]}], + fallback_access_check=denying("b", seen), + fallback_budget_check=denying("b", seen), + ) + show_router_debug(host) + host.answer("*", rate_limit()) + + with pytest.raises(litellm.RateLimitError) as raised: + await complete(core) + + assert host.deployments_called("acompletion") == ["a", "c"] + assert str(raised.value) == ( + "litellm.RateLimitError: slow down\n\nLiteLLM: model group 'a' failed with the error above. " + "Fallback to c also failed: litellm.RateLimitError: slow down " + "Fallback to b was skipped because the caller is not authorized to call it." + ) + # The access check runs before the budget check; a skipped target is asked again in c's hop. + assert seen == ["b", "c", "c", "b"] + + +# --- fallback events --------------------------------------------------------------------------------- + + +class FallbackEvents(CustomLogger): + def __init__(self) -> None: + super().__init__() + self.events: list[tuple[str, str, object, str]] = [] + + async def log_success_fallback_event( + self, original_model_group: str, kwargs: dict, original_exception: Exception + ) -> None: + self.events.append(("success", original_model_group, kwargs.get("model"), type(original_exception).__name__)) + + async def log_failure_fallback_event( + self, original_model_group: str, kwargs: dict, original_exception: Exception + ) -> None: + self.events.append(("failure", original_model_group, kwargs.get("model"), type(original_exception).__name__)) + + +@pytest.mark.asyncio +async def test_fallback_events_reach_the_process_custom_loggers_innermost_hop_first( + make_core: MakeCore, host: ScriptedHost, monkeypatch: pytest.MonkeyPatch +) -> None: + events = FallbackEvents() + monkeypatch.setattr(litellm, "callbacks", [events]) + core = build(make_core, fallbacks=[{"a": ["b", "c"]}]) + host.answer("openai/a", rate_limit()) + host.answer("openai/b", server_error()) + + await complete(core) + + assert events.events == [ + ("success", "b", "c", "InternalServerError"), + ("success", "a", "b", "RateLimitError"), + ] + + +@pytest.mark.asyncio +async def test_a_fallback_that_fails_everywhere_logs_failure_events( + make_core: MakeCore, host: ScriptedHost, monkeypatch: pytest.MonkeyPatch +) -> None: + events = FallbackEvents() + monkeypatch.setattr(litellm, "callbacks", [events]) + core = build(make_core, fallbacks=[{"a": ["b"]}]) + host.answer("*", rate_limit()) + + with pytest.raises(litellm.RateLimitError): + await complete(core) + + assert events.events == [("failure", "a", "b", "RateLimitError")] + + +# --- content-filter responses ------------------------------------------------------------------------ + + +def content_filtered() -> object: + response = completion_response() + response.choices[0].finish_reason = "content_filter" # pyright: ignore[reportAttributeAccessIssue] # Choices + return response + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("arguments", "expected"), + [ + ({"content_policy_fallbacks": [{"a": ["b"]}]}, ["a", "b"]), + ({"default_fallbacks": ["c"]}, ["a", "c"]), + ({"fallbacks": [{"a": ["b"]}]}, ["a"]), + ({"content_policy_fallbacks": [{"other": ["b"]}], "default_fallbacks": ["c"]}, ["a"]), + ], + ids=["content-policy-chain", "default-fallbacks", "general-chain-only", "unmatched-content-policy-chain"], +) +async def test_a_content_filtered_response_falls_back_only_when_a_content_policy_fallback_resolves( + make_core: MakeCore, host: ScriptedHost, arguments: dict[str, object], expected: list[str] +) -> None: + core = build(make_core, **arguments) + host.answer("openai/a", content_filtered()) + + response = await complete(core) + + assert host.deployments_called("acompletion") == expected + finish_reason = response.choices[0].finish_reason # pyright: ignore[reportAttributeAccessIssue] # ModelResponse + assert finish_reason == ("content_filter" if expected == ["a"] else "stop") + + +def test_sync_completion_falls_back_on_a_content_filtered_response(make_core: MakeCore, host: ScriptedHost) -> None: + core = build(make_core, content_policy_fallbacks=[{"a": ["b"]}]) + host.answer("openai/a", content_filtered()) + + complete_sync(core) + + assert host.deployments_called("completion") == ["a", "b"] + + +@pytest.mark.asyncio +async def test_a_content_filtered_response_with_fallbacks_disabled_raises_content_policy_violation( + make_core: MakeCore, host: ScriptedHost +) -> None: + core = build(make_core, content_policy_fallbacks=[{"a": ["b"]}]) + host.answer("openai/a", content_filtered()) + + with pytest.raises(litellm.ContentPolicyViolationError) as raised: + await complete(core, disable_fallbacks=True) + + assert ( + str(raised.value) + == "litellm.BadRequestError: litellm.ContentPolicyViolationError: Response output was blocked." + ) + assert host.deployments_called("acompletion") == ["a"] + + +@pytest.mark.asyncio +async def test_the_last_content_filtered_response_is_returned_when_every_fallback_is_filtered( + make_core: MakeCore, host: ScriptedHost +) -> None: + core = build(make_core, content_policy_fallbacks=[{"a": ["b"]}]) + host.answer("*", content_filtered()) + + response = await complete(core) + + assert host.deployments_called("acompletion") == ["a", "b"] + assert response.choices[0].finish_reason == "content_filter" # pyright: ignore[reportAttributeAccessIssue] # ModelResponse + assert headers(response)["x-litellm-model-group"] == "b" + + +# --- weighted failover --------------------------------------------------------------------------- + + +def a_deployments_fail(host: ScriptedHost, error: object) -> None: + for deployment_id in ("a1", "a2", "a3"): + host.answer(f"openai/{deployment_id}", error) + + +@pytest.mark.asyncio +async def test_weighted_failover_tries_every_deployment_of_the_group_before_the_external_chain( + make_core: MakeCore, host: ScriptedHost +) -> None: + core = build(make_core, model_list=GROUP_A_OF_THREE, enable_weighted_failover=True, fallbacks=[{"a": ["b"]}]) + a_deployments_fail(host, rate_limit()) + + response = await complete(core) + + called = host.deployments_called("acompletion") + assert sorted(called[:3]) == ["a1", "a2", "a3"] + assert called[3:] == ["b"] + assert headers(response)["x-litellm-attempted-fallbacks"] == 1 + + +@pytest.mark.asyncio +async def test_without_weighted_failover_one_deployment_fails_over_to_the_external_chain( + make_core: MakeCore, host: ScriptedHost +) -> None: + core = build(make_core, model_list=GROUP_A_OF_THREE, fallbacks=[{"a": ["b"]}]) + a_deployments_fail(host, rate_limit()) + + await complete(core) + + called = host.deployments_called("acompletion") + assert len(called) == 2 + assert called[0] in {"a1", "a2", "a3"} + assert called[1] == "b" + + +@pytest.mark.asyncio +async def test_weighted_failover_moves_an_embedding_to_another_deployment_of_the_group( + make_core: MakeCore, host: ScriptedHost +) -> None: + core = build(make_core, model_list=GROUP_A_OF_THREE, enable_weighted_failover=True) + host.answer("*", fails_once_then_answers(rate_limit())) + + await core.aembedding(model="a", input=["x"]) # pyright: ignore[reportAttributeAccessIssue] # endpoint method + + calls = host.calls_to("aembedding") + assert len(calls) == 2 + assert calls[0].deployment_id != calls[1].deployment_id + assert [metadata_of(call.kwargs)["attempted_fallbacks"] for call in calls] == [0, 1] + + +@pytest.mark.asyncio +async def test_a_weighted_failover_that_answers_counts_as_a_fallback_inside_the_group( + make_core: MakeCore, host: ScriptedHost +) -> None: + core = build(make_core, model_list=GROUP_A_OF_THREE, enable_weighted_failover=True) + host.answer("*", fails_once_then_answers(rate_limit())) + + response = await complete(core) + + calls = host.calls_to("acompletion") + assert len(calls) == 2 + assert calls[0].deployment_id != calls[1].deployment_id + assert [call.model_group for call in calls] == ["a", "a"] + assert [metadata_of(call.kwargs)["attempted_fallbacks"] for call in calls] == [0, 1] + assert headers(response)["x-litellm-model-group"] == "a" + assert headers(response)["x-litellm-attempted-fallbacks"] == 1 + + +@pytest.mark.asyncio +@pytest.mark.parametrize(("max_fallbacks", "group_attempts"), [(0, 1), (1, 2), (2, 3)]) +async def test_weighted_failover_is_bounded_by_max_fallbacks( + make_core: MakeCore, host: ScriptedHost, max_fallbacks: int, group_attempts: int +) -> None: + core = build( + make_core, + model_list=GROUP_A_OF_THREE, + enable_weighted_failover=True, + fallbacks=[{"a": ["b"]}], + max_fallbacks=max_fallbacks, + ) + a_deployments_fail(host, rate_limit()) + + if max_fallbacks == 0: + with pytest.raises(litellm.RateLimitError): + await complete(core) + else: + await complete(core) + + called = host.deployments_called("acompletion") + assert len(set(called[:group_attempts])) == group_attempts + assert called[group_attempts:] == ([] if max_fallbacks == 0 else ["b"]) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("arguments", "error"), + [ + ({"routing_strategy": "latency-based-routing"}, rate_limit), + ({}, context_window), + ], + ids=["not-simple-shuffle", "context-window-error"], +) +async def test_weighted_failover_does_not_apply( + make_core: MakeCore, host: ScriptedHost, arguments: dict[str, object], error: object +) -> None: + core = build( + make_core, model_list=GROUP_A_OF_THREE, enable_weighted_failover=True, fallbacks=[{"a": ["b"]}], **arguments + ) + a_deployments_fail(host, error()) # pyright: ignore[reportCallIssue] # an error factory + + await complete(core) + + called = host.deployments_called("acompletion") + assert len(called) == 2 + assert called[1] == "b" + + +@pytest.mark.asyncio +async def test_weighted_failover_with_fallbacks_disabled_tries_one_deployment( + make_core: MakeCore, host: ScriptedHost +) -> None: + core = build(make_core, model_list=GROUP_A_OF_THREE, enable_weighted_failover=True, fallbacks=[{"a": ["b"]}]) + host.answer("*", rate_limit()) + + with pytest.raises(litellm.RateLimitError): + await complete(core, disable_fallbacks=True) + + assert len(host.deployments_called("acompletion")) == 1 + + +def test_sync_completion_also_runs_weighted_failover(make_core: MakeCore, host: ScriptedHost) -> None: + core = build(make_core, model_list=GROUP_A_OF_THREE, enable_weighted_failover=True, fallbacks=[{"a": ["b"]}]) + a_deployments_fail(host, rate_limit()) + + complete_sync(core) + + called = host.deployments_called("completion") + assert sorted(called[:3]) == ["a1", "a2", "a3"] + assert called[3:] == ["b"] + + +@pytest.mark.asyncio +async def test_weighted_failover_on_a_single_deployment_group_goes_straight_to_the_external_chain( + make_core: MakeCore, host: ScriptedHost +) -> None: + core = build( + make_core, + model_list=[deployment("a", "a1"), deployment("b", "b")], + enable_weighted_failover=True, + fallbacks=[{"a": ["b"]}], + ) + host.answer("openai/a1", rate_limit()) + + await complete(core) + + assert host.deployments_called("acompletion") == ["a1", "b"] diff --git a/tests/unit/test_router/core/test_fallbacks_characterization_refusal.py b/tests/unit/test_router/core/test_fallbacks_characterization_refusal.py new file mode 100644 index 00000000000..d7c893a36f1 --- /dev/null +++ b/tests/unit/test_router/core/test_fallbacks_characterization_refusal.py @@ -0,0 +1,101 @@ +""" +Characterization of a router core's Anthropic safeguard-refusal fallbacks on /v1/messages +(non-streaming): a ``stop_reason: "refusal"`` response carrying ``stop_details`` re-enters the +fallback chain only when a fallback would resolve for it. + +The core recognises a /v1/messages call by its endpoint's name, so this module runs on +``NamedEndpointHost`` (a ``ScriptedHost`` whose ``anthropic_messages`` endpoint is named so). +""" + +import pytest + +from litellm.router_core import RouterCore +from tests.unit.test_router.core.fallbacks_helpers import ( + MESSAGES, + MakeCore, + NamedEndpointHost, + anthropic_message, + anthropic_refusal, + build, +) + + +@pytest.fixture +def host() -> NamedEndpointHost: + return NamedEndpointHost() + + +async def messages(core: RouterCore, **request: object) -> object: # kwargs-ok: request kwargs + return await core.aanthropic_messages(model="a", messages=MESSAGES, max_tokens=10, **request) # pyright: ignore[reportAttributeAccessIssue] # endpoint method + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("arguments", "expected"), + [ + ({"content_policy_fallbacks": [{"a": ["b"]}]}, ["a", "b"]), + ({"fallbacks": [{"a": ["b"]}]}, ["a", "b"]), + ({"default_fallbacks": ["c"]}, ["a", "c"]), + ], + ids=["content-policy-chain", "general-chain", "default-fallbacks"], +) +async def test_a_safeguard_refusal_falls_back_when_a_chain_resolves( + make_core: MakeCore, host: NamedEndpointHost, arguments: dict[str, object], expected: list[str] +) -> None: + core = build(make_core, **arguments) + host.answer("openai/a", anthropic_refusal()) + host.answer(f"openai/{expected[-1]}", anthropic_message()) + + response = await messages(core) + + assert host.deployments_called("anthropic_messages") == expected + assert isinstance(response, dict) + assert response["stop_reason"] == "end_turn" + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("arguments", "request_kwargs", "refusal"), + [ + ({}, {}, anthropic_refusal()), + ({"content_policy_fallbacks": [{"other": ["b"]}], "fallbacks": [{"a": ["b"]}]}, {}, anthropic_refusal()), + ({"fallbacks": [{"a": ["b"]}]}, {"disable_fallbacks": True}, anthropic_refusal()), + ({"fallbacks": [{"a": ["b"]}]}, {}, anthropic_refusal(with_stop_details=False)), + ], + ids=[ + "no-chain", + "unmatched-content-policy-chain-is-authoritative", + "fallbacks-disabled", + "plain-refusal-without-stop-details", + ], +) +async def test_a_refusal_is_returned_to_the_caller_when_no_fallback_would_resolve( + make_core: MakeCore, + host: NamedEndpointHost, + arguments: dict[str, object], + request_kwargs: dict[str, object], + refusal: dict[str, object], +) -> None: + core = build(make_core, **arguments) + host.answer("openai/a", refusal) + + response = await messages(core, **request_kwargs) + + assert host.deployments_called("anthropic_messages") == ["a"] + assert isinstance(response, dict) + assert response["stop_reason"] == "refusal" + assert response["_hidden_params"]["additional_headers"]["x-litellm-attempted-fallbacks"] == 0 + + +@pytest.mark.asyncio +async def test_the_last_refusal_is_returned_when_every_fallback_refuses( + make_core: MakeCore, host: NamedEndpointHost +) -> None: + core = build(make_core, fallbacks=[{"a": ["b"]}]) + host.answer("*", anthropic_refusal()) + + response = await messages(core) + + assert host.deployments_called("anthropic_messages") == ["a", "b"] + assert isinstance(response, dict) + assert response["stop_reason"] == "refusal" diff --git a/tests/unit/test_router/core/test_fallbacks_walk_parity.py b/tests/unit/test_router/core/test_fallbacks_walk_parity.py new file mode 100644 index 00000000000..47c035d096f --- /dev/null +++ b/tests/unit/test_router/core/test_fallbacks_walk_parity.py @@ -0,0 +1,353 @@ +""" +The fallback walk over what the Python router walks and a smaller walk would miss: order levels +behind any name ``_core_get_model_list`` lists, errors raised while picking, fallback checks that +raise, dict targets naming no group, weighted failover past an affinity pin, and the encrypted +reasoning a fallback hop strips. +""" + +from collections.abc import Mapping +from typing import Final + +import pytest + +import litellm +from litellm.router_utils.pre_call_checks.deployment_affinity_check import DeploymentAffinityCheck +from litellm.types.router import RoutingContext as PluginContext +from tests.unit.test_router.core.fallbacks_helpers import ( + MakeCore, + build, + complete, + headers, + show_router_debug, +) +from tests.unit.test_router.core.scripted_host import ( + ScriptedHost, + completion_response, + deployment, + rate_limit, + server_error, +) + +TEAM_METADATA: Final = {"user_api_key_team_id": "T1"} + +# --- order levels behind a name ------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_order_levels_walk_through_a_wildcard(make_core: MakeCore, host: ScriptedHost) -> None: + core: Final = build( + make_core, + model_list=[ + deployment("openai/*", "w1", model="openai/*", order=1), + deployment("openai/*", "w2", model="openai/*", order=2), + ], + ) + host.answer("openai/gpt-4o", rate_limit(), completion_response()) + + await complete(core, "openai/gpt-4o") + + assert host.deployments_called("acompletion") == ["w1", "w2"] + + +@pytest.mark.asyncio +async def test_order_levels_walk_through_an_alias(make_core: MakeCore, host: ScriptedHost) -> None: + core: Final = build( + make_core, + model_list=[deployment("gg", "d1", order=1), deployment("gg", "d2", order=2)], + model_group_alias={"al": "gg"}, + ) + host.answer("openai/d1", rate_limit()) + + await complete(core, "al") + + assert host.deployments_called("acompletion") == ["d1", "d2"] + + +@pytest.mark.asyncio +async def test_order_levels_walk_through_a_team_public_name(make_core: MakeCore, host: ScriptedHost) -> None: + public: Final = {"team_id": "T1", "team_public_model_name": "tp"} + core: Final = build( + make_core, + model_list=[ + deployment("t1a", "ta1", order=1, model_info=public), + deployment("t1b", "ta2", order=2, model_info=public), + ], + ) + host.answer("openai/ta1", rate_limit()) + + await complete(core, "tp", metadata=dict(TEAM_METADATA)) + + assert host.deployments_called("acompletion") == ["ta1", "ta2"] + + +# --- a team wildcard ------------------------------------------------------------------------------ + + +@pytest.mark.asyncio +async def test_a_team_wildcard_public_name_routes_with_its_wildcard_filled_in( + make_core: MakeCore, host: ScriptedHost +) -> None: + core: Final = build( + make_core, + model_list=[ + deployment( + "t1-w", + "tw1", + model="openai/*", + model_info={"team_id": "T1", "team_public_model_name": "tw-*"}, + ), + deployment("g", "gp"), + ], + ) + + await complete(core, "tw-gpt-4o", metadata=dict(TEAM_METADATA)) + + assert host.deployments_called("acompletion") == ["tw1"] + assert host.calls[0].kwargs["model"] == "openai/gpt-4o" + + +# --- errors raised while picking ------------------------------------------------------------------ + + +@pytest.mark.asyncio +async def test_a_usage_refusal_takes_the_fallback(make_core: MakeCore, host: ScriptedHost) -> None: + core: Final = build( + make_core, + model_list=[deployment("one", "o", rpm=1), deployment("safe", "s")], + routing_strategy="usage-based-routing-v2", + fallbacks=[{"one": ["safe"]}], + ) + await complete(core, "one") + + response: Final = await complete(core, "one") + + assert host.deployments_called("acompletion") == ["o", "s"] + assert headers(response)["x-litellm-attempted-fallbacks"] == 1 + + +@pytest.mark.asyncio +async def test_a_usage_refusal_carries_the_routing_note(make_core: MakeCore, host: ScriptedHost) -> None: + core: Final = make_core( + model_list=[deployment("g", "d1", rpm=1)], routing_strategy="usage-based-routing-v2", num_retries=0 + ) + show_router_debug(host) + await complete(core, "g") + + with pytest.raises(litellm.RateLimitError) as raised: + await complete(core, "g") + + assert "No fallback" in str(raised.value) + + +class _RaisingPlugin: + """Refuses every pick of ``refused`` in the pre-routing hook.""" + + def __init__(self, refused: str) -> None: + self.refused: Final = refused + self.runs: int = 0 + + async def run(self, context: PluginContext) -> PluginContext: + if self.refused in context.candidate_models: + self.runs += 1 + raise rate_limit() + return context + + +@pytest.mark.asyncio +async def test_a_pre_routing_error_runs_the_pipeline_once_per_retry(make_core: MakeCore, host: ScriptedHost) -> None: + plugin: Final = _RaisingPlugin("openai/a") + core: Final = build(make_core, num_retries=1, retry_after=0, plugins=[plugin]) + + with pytest.raises(litellm.RateLimitError) as raised: + await complete(core) + + assert plugin.runs == 2 + assert host.calls == [] + assert getattr(raised.value, "max_retries", None) == 1 + + +@pytest.mark.asyncio +async def test_a_pre_routing_error_takes_the_fallback(make_core: MakeCore, host: ScriptedHost) -> None: + core: Final = build(make_core, plugins=[_RaisingPlugin("openai/a")], fallbacks=[{"a": ["b"]}]) + + await complete(core) + + assert host.deployments_called("acompletion") == ["b"] + + +# --- fallback checks that raise -------------------------------------------------------------------- + + +def _sync_check(*, model: str, request_kwargs: Mapping[str, object], ctx: object) -> bool: + return True + + +@pytest.mark.asyncio +async def test_a_sync_access_check_ends_the_chain_with_the_levels_own_error( + make_core: MakeCore, host: ScriptedHost +) -> None: + core: Final = build(make_core, fallbacks=[{"a": ["b"]}], fallback_access_check=_sync_check) + show_router_debug(host) + host.answer("openai/a", server_error()) + + with pytest.raises(litellm.InternalServerError) as raised: + await complete(core) + + assert "also failed" in str(raised.value) + assert host.deployments_called("acompletion") == ["a"] + + +@pytest.mark.asyncio +async def test_a_sync_access_check_on_the_order_path_raises_its_own_error( + make_core: MakeCore, host: ScriptedHost +) -> None: + core: Final = build( + make_core, + model_list=[deployment("a", "a1", order=1), deployment("a", "a2", order=2), deployment("b", "b")], + fallbacks=[{"a": ["b"]}], + fallback_access_check=_sync_check, + ) + host.answer("openai/a1", server_error()) + host.answer("openai/a2", server_error()) + + with pytest.raises(TypeError): + await complete(core) + + assert host.deployments_called("acompletion") == ["a1", "a2"] + + +# --- dict targets naming no group ------------------------------------------------------------------- + + +@pytest.mark.asyncio +@pytest.mark.parametrize("model", [5, None], ids=["int", "none"]) +async def test_a_dict_target_with_a_non_string_model_is_a_hop_that_calls_nothing( + make_core: MakeCore, host: ScriptedHost, model: object +) -> None: + core: Final = build(make_core, fallbacks=[{"a": [{"model": model}, "b"]}]) + host.answer("openai/a", server_error()) + + response: Final = await complete(core) + + assert host.deployments_called("acompletion") == ["a", "b"] + assert headers(response)["x-litellm-attempted-fallbacks"] == 2 + + +@pytest.mark.asyncio +async def test_a_chain_ending_on_a_non_string_model_notes_the_hops_error( + make_core: MakeCore, host: ScriptedHost +) -> None: + core: Final = build(make_core, fallbacks=[{"a": [{"model": 5}]}]) + show_router_debug(host) + host.answer("openai/a", server_error()) + + with pytest.raises(litellm.InternalServerError) as raised: + await complete(core) + + assert "argument of type 'int'" in str(raised.value) + assert host.deployments_called("acompletion") == ["a"] + + +# --- weighted failover --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +@pytest.mark.parametrize("repin", [True, False]) +async def test_weighted_failover_skips_an_affinity_pin_on_the_failed_deployment( + make_core: MakeCore, host: ScriptedHost, repin: bool +) -> None: + core: Final = build( + make_core, + model_list=[deployment("g", "a"), deployment("g", "b")], + enable_weighted_failover=True, + optional_pre_call_checks=["session_affinity"], + session_affinity_repin_on_failover=repin, + ) + check: Final = next( + callback + for callback in core.optional_callbacks or () # pyright: ignore[reportAttributeAccessIssue, reportUnknownVariableType, reportUnknownMemberType] + if isinstance(callback, DeploymentAffinityCheck) + ) + key: Final = DeploymentAffinityCheck.get_session_affinity_cache_key("g", "s", "c") + await check.cache.async_set_cache(key, {"model_id": "a"}) # pyright: ignore[reportUnknownMemberType] + host.answer("openai/a", server_error()) + + response: Final = await complete(core, "g", metadata={"session_id": "s", "user_api_key_hash": "c"}) + + assert host.deployments_called("acompletion") == ["a", "b"] + assert headers(response)["x-litellm-attempted-fallbacks"] == 1 + + +def _fails_then_raises(raised: BaseException): # noqa: ANN202 # a scripted outcome + """One outcome both deployments of a group share: a provider failure, then ``raised``.""" + calls: Final = [0] + + def outcome(**_request: object) -> object: # kwargs-ok: the scripted request + calls[0] += 1 + return server_error() if calls[0] == 1 else raised + + return outcome + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("hop_error", "error"), + [(TypeError("a bug"), TypeError), (rate_limit(), litellm.InternalServerError)], + ids=["a-bug-propagates", "a-provider-error-defers"], +) +async def test_weighted_failover_swallows_only_api_and_rate_limit_errors( + make_core: MakeCore, host: ScriptedHost, hop_error: BaseException, error: type[Exception] +) -> None: + core: Final = build( + make_core, + model_list=[deployment("a", "a1", model="openai/same"), deployment("a", "a2", model="openai/same")], + enable_weighted_failover=True, + ) + host.answer("openai/same", _fails_then_raises(hop_error)) + + with pytest.raises(error): + await complete(core) + + assert len(host.deployments_called("acompletion")) == 2 + + +# --- refusals the walk falls back from --------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_a_paused_deployment_named_by_id_falls_back(make_core: MakeCore, host: ScriptedHost) -> None: + core: Final = build( + make_core, + model_list=[deployment("a", "a1", model_info={"blocked": True}), deployment("b", "b")], + fallbacks=[{"a1": ["b"]}], + ) + + await complete(core, "a1") + + assert host.deployments_called("acompletion") == ["b"] + + +# --- encrypted reasoning on a fallback hop -------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_a_fallback_hop_strips_the_reasoning_the_target_cannot_decrypt( + make_core: MakeCore, host: ScriptedHost +) -> None: + core: Final = build( + make_core, + model_list=[deployment("a", "a", api_key="k-a"), deployment("b", "b", api_key="k-b")], + fallbacks=[{"a": ["b"]}], + ) + host.answer("openai/a", server_error()) + reasoning: Final = { + "type": "reasoning", + "encrypted_content": "opaque", + "summary": [{"type": "summary_text", "text": "thought"}], + } + + await core.aresponses(model="a", input=[reasoning, {"role": "user", "content": "hi"}]) # pyright: ignore[reportAttributeAccessIssue] # endpoint method + + sent: Final = host.calls_to("aresponses")[-1].kwargs["input"] + assert isinstance(sent, list) + assert not any(isinstance(item, dict) and item.get("encrypted_content") for item in sent) diff --git a/tests/unit/test_router/core/test_gaps_adapter_timeout.py b/tests/unit/test_router/core/test_gaps_adapter_timeout.py new file mode 100644 index 00000000000..4b895673481 --- /dev/null +++ b/tests/unit/test_router/core/test_gaps_adapter_timeout.py @@ -0,0 +1,34 @@ +""" +Characterization gap: a streaming adapter completion splits its timeouts like a streaming chat +completion, sending the request timeout as ``timeout`` and the stream timeout beside it, so a +fake-streaming provider stays bounded by the request timeout rather than the stream timeout. + +Every test builds its core through ``make_core`` on the ``host`` fixture, so it runs once per core. +""" + +from collections.abc import Callable +from typing import Final + +import pytest + +from litellm.router_core import RouterCore +from tests.unit.test_router.core.endpoints_helpers import EndpointHost +from tests.unit.test_router.core.scripted_host import deployment + +MESSAGES: Final = [{"role": "user", "content": "hi"}] +MakeCore = Callable[..., RouterCore] + + +@pytest.fixture +def host() -> EndpointHost: + return EndpointHost() + + +@pytest.mark.parametrize(("stream", "expected"), [(False, (11.0, None)), (True, (11.0, 5.0))]) +async def test_adapter_completion_sends_the_request_timeout_and_the_stream_timeout_beside_it( + make_core: MakeCore, host: EndpointHost, stream: bool, expected: tuple[float, float | None] +): + core: Final = make_core(model_list=[deployment("g", "d1", timeout=11)], stream_timeout=5) + await core.aadapter_completion(adapter_id="my-adapter", model="g", messages=MESSAGES, stream=stream) + sent: Final = host.calls[-1].kwargs + assert (sent["timeout"], sent.get("stream_timeout")) == expected diff --git a/tests/unit/test_router/core/test_gaps_fuzz_access.py b/tests/unit/test_router/core/test_gaps_fuzz_access.py new file mode 100644 index 00000000000..3adec609033 --- /dev/null +++ b/tests/unit/test_router/core/test_gaps_fuzz_access.py @@ -0,0 +1,267 @@ +""" +Access and policy gaps differential fuzzing found between the cores: dict fallback targets that +may not be set from any list, which metadata bucket names the caller's team, the routing plugins on +the async pass-through pick, and a name several teams share falling back. +""" + +from typing import Final + +import pytest + +import litellm +from litellm.proxy._types import UserAPIKeyAuth +from litellm.types.router import CustomRoutingStrategyBase, RouterRateLimitError +from tests.unit.test_router.core.fallbacks_helpers import ( + MakeCore, + build, + complete, + complete_sync, + show_router_debug, +) +from tests.unit.test_router.core.scripted_host import ScriptedHost, content_policy, deployment, rate_limit +from tests.unit.test_router.core.test_selection_characterization_prerouting import Narrowing + +OIDC_TARGET: Final = {"model": "b", "api_key": "oidc/google/x"} + +# --- dict fallback targets from a router's lists --------------------------------------------------- + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "target", + [OIDC_TARGET, {"model": "b", "anthropic_issuer_url": "https://issuer.example"}], + ids=["oidc-api-key", "federation-field"], +) +async def test_a_router_level_dict_target_refuses_the_whole_chain( + make_core: MakeCore, host: ScriptedHost, target: dict[str, str] +) -> None: + core: Final = build(make_core, fallbacks=[{"a": [target, "c"]}]) + host.answer("openai/a", rate_limit()) + + with pytest.raises(litellm.RateLimitError): + await complete(core) + + assert host.deployments_called("acompletion") == ["a"] + + +@pytest.mark.asyncio +async def test_a_refused_router_level_chain_explains_itself_under_router_debug( + make_core: MakeCore, host: ScriptedHost +) -> None: + core: Final = build(make_core, fallbacks=[{"a": [OIDC_TARGET]}]) + show_router_debug(host) + host.answer("openai/a", rate_limit()) + + with pytest.raises(litellm.RateLimitError) as raised: + await complete(core) + + assert "also failed: Rejected Request: api_key holds an oidc/ secret reference" in str(raised.value) + + +@pytest.mark.asyncio +async def test_a_content_policy_list_with_a_refused_target_tries_nothing( + make_core: MakeCore, host: ScriptedHost +) -> None: + core: Final = build(make_core, content_policy_fallbacks=[{"a": [OIDC_TARGET]}]) + host.answer("openai/a", content_policy()) + + with pytest.raises(litellm.ContentPolicyViolationError): + await complete(core) + + assert host.deployments_called("acompletion") == ["a"] + + +@pytest.mark.asyncio +async def test_a_refused_chain_on_the_order_path_raises_the_refusal_itself( + make_core: MakeCore, host: ScriptedHost +) -> None: + # The order path's `run_async_fallback` sits outside the `try` that re-raises the original error. + core: Final = make_core( + model_list=[deployment("a", "a1", order=1), deployment("a", "a2", order=2), deployment("b", "b")], + num_retries=0, + fallbacks=[{"a": [OIDC_TARGET]}], + ) + host.answer("openai/a1", rate_limit()) + + with pytest.raises(ValueError, match="oidc/ secret reference"): + await complete(core) + + assert host.deployments_called("acompletion") == ["a1"] + + +@pytest.mark.asyncio +async def test_a_refused_chain_fails_only_its_own_level(make_core: MakeCore, host: ScriptedHost) -> None: + core: Final = build(make_core, fallbacks=[{"a": ["b", "c"]}, {"b": [{"model": "c", "api_key": "oidc/google/x"}]}]) + host.answer("openai/a", rate_limit()) + host.answer("openai/b", rate_limit()) + + await complete(core) + + assert host.deployments_called("acompletion") == ["a", "b", "c"] + + +def test_a_refused_router_level_chain_is_refused_on_the_sync_path(make_core: MakeCore, host: ScriptedHost) -> None: + core: Final = build(make_core, fallbacks=[{"a": [OIDC_TARGET]}]) + host.answer("openai/a", rate_limit()) + + with pytest.raises(litellm.RateLimitError): + complete_sync(core) + + assert host.deployments_called("completion") == ["a"] + + +# --- which bucket names the team ------------------------------------------------------------------- + +_TEAM: Final = {"user_api_key_team_id": "t1"} +_OWN_OR_PUBLIC: Final = [ + deployment("g", "own", weight=1, model_info={"team_id": "t1"}), + deployment("g", "public", weight=0), +] + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("buckets", "expected"), + [ + ({"litellm_metadata": _TEAM}, "own"), + ({"metadata": {"trace": "x"}, "litellm_metadata": _TEAM}, "own"), + ({"metadata": {"user_api_key_team_id": ""}, "litellm_metadata": _TEAM}, "own"), + ({"metadata": {"user_api_key_team_id": "t2"}, "litellm_metadata": _TEAM}, "public"), + ], + ids=["only-litellm-metadata", "metadata-without-a-team", "metadata-with-an-empty-team", "metadata-wins"], +) +async def test_the_first_bucket_holding_a_team_routes_the_call( + make_core: MakeCore, host: ScriptedHost, buckets: dict[str, dict[str, str]], expected: str +) -> None: + core: Final = make_core(model_list=list(_OWN_OR_PUBLIC), num_retries=0) + + await complete(core, "g", **{name: dict(bucket) for name, bucket in buckets.items()}) + + assert host.deployments_called("acompletion") == [expected] + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("bucket", "raised"), + [("metadata", litellm.RateLimitError), ("litellm_metadata", RouterRateLimitError)], + ids=["one-level-listed-so-the-level-raises-its-own-error", "no-team-listed-so-the-second-level-raises"], +) +async def test_the_order_levels_of_a_failed_group_read_only_the_metadata_bucket( + make_core: MakeCore, host: ScriptedHost, bucket: str, raised: type[Exception] +) -> None: + core: Final = make_core( + model_list=[ + deployment("g", "own", order=1, model_info={"team_id": "t1"}), + deployment("g", "theirs", order=2, model_info={"team_id": "t2"}), + ], + num_retries=0, + ) + host.answer("openai/own", rate_limit()) + + with pytest.raises(raised): + await complete(core, "g", **{bucket: dict(_TEAM)}) + + assert host.deployments_called("acompletion") == ["own"] + + +# --- the routing plugins on the async pass-through pick -------------------------------------------- + +_PASS_THROUGH_PAIR: Final = [ + deployment("g", "a", use_in_pass_through=True, weight=1), + deployment("g", "b", use_in_pass_through=True, weight=0), +] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("strategy", ["simple-shuffle", "usage-based-routing"]) +async def test_the_async_pass_through_pick_runs_the_routing_plugins(make_core: MakeCore, strategy: str) -> None: + plugin: Final = Narrowing(frozenset({"openai/b"})) + core: Final = make_core(model_list=list(_PASS_THROUGH_PAIR), plugins=[plugin], routing_strategy=strategy) + + chosen: Final = await core.async_get_available_deployment_for_pass_through("g", {}) # pyright: ignore[reportAttributeAccessIssue] + + assert chosen["model_info"]["id"] == "b" + assert plugin.seen == [["openai/a", "openai/b"]] + + +@pytest.mark.asyncio +async def test_a_custom_strategy_does_not_answer_the_async_pass_through_pick(make_core: MakeCore) -> None: + class Refuses(CustomRoutingStrategyBase): + async def async_get_available_deployment(self, *args: object, **kwargs: object) -> object: # pyright: ignore[reportIncompatibleMethodOverride] # kwargs-ok: never called + raise AssertionError("a pass-through pick never reaches the custom strategy") + + core: Final = make_core(model_list=[deployment("g", "a", use_in_pass_through=True)]) + core.register_plugin(Refuses()) + + chosen: Final = await core.async_get_available_deployment_for_pass_through("g", {}) # pyright: ignore[reportAttributeAccessIssue] + + assert chosen["model_info"]["id"] == "a" + + +# --- a name several teams share -------------------------------------------------------------------- + +_SHARED_NAME: Final = [ + deployment("g", "gp"), + deployment("t1-pub", "p1", model_info={"team_id": "T1", "team_public_model_name": "pub"}), + deployment("t2-pub", "p2", model_info={"team_id": "T2", "team_public_model_name": "pub"}), +] +_REFUSAL: Final = ( + "litellm.BadRequestError: Model name 'pub' matches deployments from multiple teams. " + "Specify the deployment ID directly to disambiguate." +) + + +def _admin() -> dict[str, object]: # mutable-ok: a request's metadata + return {"user_api_key_auth": UserAPIKeyAuth(user_role="proxy_admin")} # pyright: ignore[reportCallIssue] + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("configured", "requested"), + [({"fallbacks": [{"pub": ["g"]}]}, {}), ({}, {"fallbacks": ["g"]})], + ids=["router-list", "request-list"], +) +async def test_a_name_several_teams_share_falls_back( + make_core: MakeCore, host: ScriptedHost, configured: dict[str, object], requested: dict[str, object] +) -> None: + core: Final = make_core(model_list=list(_SHARED_NAME), num_retries=0, **configured) + + await complete(core, "pub", metadata=_admin(), **requested) + + assert host.deployments_called("acompletion") == ["gp"] + + +def test_a_name_several_teams_share_falls_back_on_the_sync_path(make_core: MakeCore, host: ScriptedHost) -> None: + core: Final = make_core(model_list=list(_SHARED_NAME), num_retries=0) + + complete_sync(core, "pub", metadata=_admin(), fallbacks=["g"]) + + assert host.deployments_called("completion") == ["gp"] + + +@pytest.mark.asyncio +async def test_a_shared_name_raises_its_refusal_when_the_fallback_fails( + make_core: MakeCore, host: ScriptedHost +) -> None: + core: Final = make_core(model_list=list(_SHARED_NAME), num_retries=0) + host.answer("openai/gp", rate_limit()) + + with pytest.raises(litellm.BadRequestError) as raised: + await complete(core, "pub", metadata=_admin(), fallbacks=["g"]) + + assert raised.value.status_code == 400 + assert str(raised.value) == _REFUSAL + + +@pytest.mark.asyncio +async def test_a_shared_name_explains_the_failed_fallback_under_router_debug( + make_core: MakeCore, host: ScriptedHost +) -> None: + core: Final = make_core(model_list=list(_SHARED_NAME), num_retries=0) + show_router_debug(host) + host.answer("openai/gp", rate_limit()) + + with pytest.raises(litellm.BadRequestError) as raised: + await complete(core, "pub", metadata=_admin(), fallbacks=["g"]) + + assert str(raised.value).startswith(f"{_REFUSAL}\n\nLiteLLM: model group 'pub' failed with the error above.") diff --git a/tests/unit/test_router/core/test_gaps_l0001.py b/tests/unit/test_router/core/test_gaps_l0001.py new file mode 100644 index 00000000000..bb6a4fe6c44 --- /dev/null +++ b/tests/unit/test_router/core/test_gaps_l0001.py @@ -0,0 +1,445 @@ +""" +Gap tests for the first part of ``litellm/router_python_core.py`` (module helpers, the entry points, +``__init__``, callback registration, strategy setup, routing groups and the strategy dispatch): +behaviors no other characterization test pinned down. + +Each test asserts what the Python core does today; oddities are named in the test's docstring. +""" + +import datetime +import re +from typing import Final + +import pytest + +import litellm +from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging +from litellm.router_api import ChangeSet +from litellm.types.router import Deployment, RouterRateLimitError +from litellm.utils import reapply_runtime_model_cost_registrations +from tests.unit.test_router.core.construction_precall_helpers import collect_other_routers, host_callback_names +from tests.unit.test_router.core.registry_helpers import unique_id +from tests.unit.test_router.core.scripted_host import ScriptedHost, deployment +from tests.unit.test_router.core.selection_helpers import MESSAGES, MakeCore, picked_ids + +NO_DEPLOYMENTS_MESSAGE: Final = ( + "No deployments available for selected model, Try again in 5 seconds. Passed model={model}. " + "pre-call-checks=False, cooldown_list=[]" +) + + +def _group( + name: str, models: list[str], strategy: str, **extra: object +) -> dict[str, object]: # mutable-ok: a routing_groups entry # kwargs-ok: extra RoutingGroup fields + return {"group_name": name, "models": models, "routing_strategy": strategy, **extra} + + +def _logging_obj() -> LiteLLMLogging: + return LiteLLMLogging( + model="g", + messages=MESSAGES, + stream=False, + call_type="acompletion", + start_time=datetime.datetime(2026, 1, 1, tzinfo=datetime.timezone.utc), + litellm_call_id="c1", + function_id="f1", + ) + + +def _dynamic_callback_names(logging_obj: LiteLLMLogging) -> list[str]: # mutable-ok: a fresh list + return [type(callback).__name__ for callback in logging_obj.dynamic_async_success_callbacks or ()] + + +# --- construction: routing_strategy=None and legacy attributes ------------------------------------- + + +@pytest.mark.asyncio +async def test_routing_strategy_none_is_accepted_but_leaves_every_pick_without_a_deployment( + make_core: MakeCore, +) -> None: + """Oddity: ``routing_strategy=None`` passes validation and is stored as None (not the + ``simple-shuffle`` default), so no strategy arm matches and every pick rate-limits, while the + healthy listing still lists the deployment.""" + core: Final = make_core(model_list=[deployment("g", "a")], routing_strategy=None) + assert core.settings().routing_strategy is None + assert core.settings().as_dict()["routing_strategy"] is None + with pytest.raises(RouterRateLimitError) as raised: + await core.select("g") + assert str(raised.value) == NO_DEPLOYMENTS_MESSAGE.format(model="g") + healthy: Final = await core.select("g", kind="healthy") + assert isinstance(healthy, list) + assert [entry["model_info"]["id"] for entry in healthy] == ["a"] + + +LEGACY_ATTRIBUTE_DEFAULTS: Final = { + "total_calls": {}, + "fail_calls": {}, + "success_calls": {}, + "deployment_stats": {}, + "default_cache_time_seconds": 3600, + "tenacity": None, + "lowesttpm_logger_v2": None, + "lowestlatency_logger": None, + "lowestcost_logger": None, +} + + +@pytest.mark.parametrize(("name", "expected"), sorted(LEGACY_ATTRIBUTE_DEFAULTS.items())) +def test_legacy_attribute_default(make_core: MakeCore, name: str, expected: object) -> None: + assert getattr(make_core(model_list=[deployment("g", "a")]), name) == expected + + +def test_healthy_deployments_is_the_live_deployment_list(make_core: MakeCore) -> None: + core: Final = make_core(model_list=[deployment("g", "a")]) + assert core.healthy_deployments is core.registry().all() # pyright: ignore[reportAttributeAccessIssue] + + +def test_chat_completions_create_calls_the_router_with_the_default_litellm_params( + make_core: MakeCore, host: ScriptedHost +) -> None: + core: Final = make_core(model_list=[deployment("g", "a")], default_litellm_params={"temperature": 0.3}) + core.chat.completions.create(model="g", messages=MESSAGES) # pyright: ignore[reportAttributeAccessIssue] + assert [(call.op, call.deployment_id, call.kwargs.get("temperature")) for call in host.calls] == [ + ("completion", "a", 0.3) + ] + + +@pytest.mark.asyncio +async def test_a_provider_budget_config_adds_the_budget_limiter_after_the_listed_checks(make_core: MakeCore) -> None: + """A provider budget appends ``router_budget_limiting`` to the optional pre-call checks given.""" + core: Final = make_core( + model_list=[deployment("g", "a")], + provider_budget_config={"openai": {"budget_limit": 1, "time_period": "1d"}}, + optional_pre_call_checks=["prompt_caching"], + ) + assert [type(callback).__name__ for callback in core.optional_callbacks] == [ # pyright: ignore[reportAttributeAccessIssue] + "PromptCachingDeploymentCheck", + "RouterBudgetLimiting", + ] + core.retire() + + +# --- price data reloads --------------------------------------------------------------------------- + + +def test_a_price_data_reload_restores_a_live_routers_deployment_pricing_until_it_retires( + make_core: MakeCore, +) -> None: + model_id: Final = unique_id("l0001") + core: Final = make_core(model_list=[deployment("gpt", model_id, model="openai/gpt-4o", input_cost_per_token=1e-3)]) + try: + litellm.model_cost.pop(model_id) + reapply_runtime_model_cost_registrations() + assert litellm.model_cost[model_id]["input_cost_per_token"] == 1e-3 + core.retire() + litellm.model_cost.pop(model_id) + reapply_runtime_model_cost_registrations() + assert model_id not in litellm.model_cost + finally: + litellm.model_cost.pop(model_id, None) + + +# --- runtime routing_strategy_args ---------------------------------------------------------------- + + +def test_new_routing_strategy_args_rebuild_the_selector_and_replace_it_on_the_host( + make_core: MakeCore, host: ScriptedHost +) -> None: + core: Final = make_core( + model_list=[deployment("g", "a")], routing_strategy="latency-based-routing", routing_strategy_args={"ttl": 5} + ) + (before,) = host.lists["callbacks"] + core.apply(ChangeSet(settings={"routing_strategy_args": {"ttl": 9}})) + (after,) = host.lists["callbacks"] + assert after is not before + assert after is core.lowestlatency_logger # pyright: ignore[reportAttributeAccessIssue] + assert after.routing_args.ttl == 9 # pyright: ignore[reportAttributeAccessIssue] + assert core.settings().as_dict()["routing_strategy_args"]["ttl"] == 9 + + +def test_invalid_routing_strategy_args_at_runtime_keep_the_previous_selector( + make_core: MakeCore, host: ScriptedHost +) -> None: + """Oddity: the invalid args are still stored as ``routing_strategy_args``, while the selector (and + so ``as_dict``, which reads the selector) keeps the previous ones.""" + core: Final = make_core( + model_list=[deployment("g", "a")], routing_strategy="latency-based-routing", routing_strategy_args={"ttl": 5} + ) + (before,) = host.lists["callbacks"] + core.apply(ChangeSet(settings={"routing_strategy_args": {"ttl": "abc"}})) + assert host.lists["callbacks"] == [before] + assert core.lowestlatency_logger is before # pyright: ignore[reportAttributeAccessIssue] + assert core.settings().as_dict()["routing_strategy_args"]["ttl"] == 5 + assert core.routing_strategy_args == {"ttl": "abc"} # pyright: ignore[reportAttributeAccessIssue] + + +def test_new_least_busy_args_replace_the_selector_in_callbacks_and_input_callbacks( + make_core: MakeCore, host: ScriptedHost +) -> None: + core: Final = make_core(model_list=[deployment("g", "a")], routing_strategy="least-busy") + before: Final = core.leastbusy_logger # pyright: ignore[reportAttributeAccessIssue] + core.apply(ChangeSet(settings={"routing_strategy_args": {"x": 1}})) + after: Final = core.leastbusy_logger # pyright: ignore[reportAttributeAccessIssue] + assert after is not before + assert (host.lists["callbacks"], host.lists["input_callback"]) == ([after], [after]) + + +def test_new_routing_strategy_args_under_simple_shuffle_register_nothing( + make_core: MakeCore, host: ScriptedHost +) -> None: + core: Final = make_core(model_list=[deployment("g", "a")]) + core.apply(ChangeSet(settings={"routing_strategy_args": {"ttl": 9}})) + assert core.settings().as_dict()["routing_strategy_args"] == {"ttl": 9} + assert (host.lists["callbacks"], host.lists["input_callback"]) == ([], []) + + +# --- routing groups: selectors on the host, shadowing --------------------------------------------- + + +def test_routing_group_selectors_are_registered_on_the_host_and_replaced_when_groups_change( + make_core: MakeCore, host: ScriptedHost +) -> None: + core: Final = make_core( + model_list=[deployment("m1", "a"), deployment("m2", "b")], + routing_strategy="least-busy", + routing_groups=[_group("g1", ["m1"], "latency-based-routing", routing_strategy_args={"ttl": 7})], + ) + assert host_callback_names(host)["callbacks"] == ["LeastBusyLoggingHandler", "LowestLatencyLoggingHandler"] + (latency,) = ( + callback for callback in host.lists["callbacks"] if type(callback).__name__ == "LowestLatencyLoggingHandler" + ) + assert latency.routing_args.ttl == 7 # pyright: ignore[reportAttributeAccessIssue] + core.apply(ChangeSet(settings={"routing_groups": [_group("g2", ["m2"], "usage-based-routing-v2")]})) + assert host_callback_names(host)["callbacks"] == ["LeastBusyLoggingHandler", "LowestTPMLoggingHandler_v2"] + core.apply(ChangeSet(settings={"routing_groups": []})) + assert host_callback_names(host)["callbacks"] == ["LeastBusyLoggingHandler"] + assert host_callback_names(host)["input_callback"] == ["LeastBusyLoggingHandler"] + + +@pytest.mark.asyncio +async def test_a_group_shadowed_by_a_model_name_still_routes_its_members_with_its_strategy( + make_core: MakeCore, +) -> None: + """The shadowed group's name is not callable (calling it reaches the model), but its least-busy + strategy still serves its member: with no traffic least-busy always picks the first deployment.""" + core: Final = make_core( + model_list=[deployment("m1", "a"), deployment("m2", "b"), deployment("m2", "c")], + routing_groups=[_group("m1", ["m2"], "least-busy")], + ) + assert await picked_ids(core, "m2", times=30, messages=MESSAGES) == {"b"} + assert core.registry().routing_group("m1") is None + + +@pytest.mark.asyncio +async def test_a_group_named_like_an_alias_is_shadowed_by_the_alias(make_core: MakeCore) -> None: + core: Final = make_core( + model_list=[deployment("m1", "a"), deployment("m2", "b")], + model_group_alias={"al": "m1"}, + routing_groups=[_group("al", ["m2"], "simple-shuffle")], + ) + assert await picked_ids(core, "al") == {"a"} + assert core.registry().routing_group("al") is None + + +def test_a_priority_group_named_like_an_alias_is_rejected(make_core: MakeCore) -> None: + with pytest.raises(ValueError, match=re.escape("Priority routing group names must not shadow a model or alias")): + make_core( + model_list=[deployment("m1", "a")], + model_group_alias={"al": "m1"}, + routing_groups=[_group("al", ["m1"], "priority", model_priorities={"m1": 1})], + ) + + +@pytest.mark.asyncio +async def test_a_deployment_added_under_a_group_name_shadows_the_group(make_core: MakeCore) -> None: + core: Final = make_core( + model_list=[deployment("m1", "a"), deployment("m2", "b")], + routing_groups=[_group("grp", ["m1", "m2"], "simple-shuffle")], + ) + core.apply(ChangeSet(upsert=[Deployment(**deployment("grp", "z"))])) # pyright: ignore[reportArgumentType] + assert await picked_ids(core, "grp") == {"z"} + assert core.registry().routing_group("grp") is None + + +# --- registry().resolve kinds backed by this range ------------------------------------------------ + + +def test_routing_group_alternatives_count_the_deployments_of_every_member(make_core: MakeCore) -> None: + """Only a callable routing group (or an alias to one) can have alternatives; a plain model group + with two deployments has none by this rule.""" + core: Final = make_core( + model_list=[deployment("m1", "a"), deployment("m1", "a2"), deployment("m2", "b"), deployment("m3", "c")], + model_group_alias={"al": "grp"}, + routing_groups=[ + _group("grp", ["m2", "m3"], "simple-shuffle"), + _group("one", ["m2"], "priority", model_priorities={"m2": 1}), + ], + ) + registry: Final = core.registry() + answers: Final = { + name: registry.resolve(name, kind="routing_group_alternatives") # pyright: ignore[reportCallIssue, reportArgumentType] + for name in (None, "grp", "one", "al", "m1", "zzz") + } + assert answers == {None: False, "grp": True, "one": False, "al": True, "m1": False, "zzz": False} + + +TEAM_MODELS: Final = [ + deployment("pub", "t1", model_info={"team_id": "T", "team_public_model_name": "pubm"}), + deployment("pub2", "t2", model_info={"team_id": "T", "team_public_model_name": "pubm"}), + deployment("solo", "t3", model_info={"team_id": "T", "team_public_model_name": "solo"}), + deployment("plain", "p1"), + deployment("half", "h1", model_info={"team_id": "T"}), +] + + +@pytest.mark.parametrize( + ("deployment_id", "expected"), + [("t1", True), ("t3", False), ("p1", False), ("h1", False), ("zzz", False)], + ids=["two-siblings", "only-one", "no-team", "no-public-name", "unknown-id"], +) +def test_team_model_alternatives_need_another_routable_deployment_under_the_same_team_name( + make_core: MakeCore, deployment_id: str, expected: bool +) -> None: + core: Final = make_core(model_list=list(TEAM_MODELS)) + assert core.registry().resolve(deployment_id, kind="team_model_alternatives") is expected # pyright: ignore[reportCallIssue, reportArgumentType] + + +def test_a_blocked_team_sibling_is_no_alternative(make_core: MakeCore) -> None: + core: Final = make_core( + model_list=[ + deployment("pub", "t1", model_info={"team_id": "T", "team_public_model_name": "pubm"}), + deployment("pub2", "t2", model_info={"team_id": "T", "team_public_model_name": "pubm", "blocked": True}), + ] + ) + assert core.registry().resolve("t1", kind="team_model_alternatives") is False # pyright: ignore[reportCallIssue, reportArgumentType] + + +def test_search_tool_support_comes_from_model_info_else_an_anthropic_request_shape(make_core: MakeCore) -> None: + """Per routable deployment: a configured ``supports_search_tool`` wins, a deployment whose chat + config builds an Anthropic request answers True, anything else None; an unknown model has none.""" + core: Final = make_core( + model_list=[ + deployment("g", "o1"), + deployment("g", "an1", model="anthropic/claude-3-5-sonnet-20240620"), + deployment( + "g", "cfg", model="anthropic/claude-3-5-sonnet-20240620", model_info={"supports_search_tool": False} + ), + ] + ) + registry: Final = core.registry() + assert registry.resolve("g", kind="search_tool_support") == (None, True, False) # pyright: ignore[reportCallIssue, reportArgumentType] + assert registry.resolve("zzz", kind="search_tool_support") == () # pyright: ignore[reportCallIssue, reportArgumentType] + + +# --- per-request strategy overrides --------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_an_override_selector_is_bound_to_the_request_logging_object_and_reused(make_core: MakeCore) -> None: + """An override strategy the router does not register globally is bound to the request's logging + object; one selector per strategy serves every request. An override matching the router's own + strategy reuses the global selector and binds nothing.""" + core: Final = make_core(model_list=[deployment("g", "a"), deployment("g", "b")], routing_strategy="least-busy") + own: Final = _logging_obj() + await core.select("g", {"routing_strategy": "least-busy", "litellm_logging_obj": own}) + first: Final = _logging_obj() + await core.select("g", {"routing_strategy": "usage-based-routing-v2", "litellm_logging_obj": first}) + second: Final = _logging_obj() + await core.select("g", {"routing_strategy": "usage-based-routing-v2", "litellm_logging_obj": second}) + assert _dynamic_callback_names(own) == [] + assert _dynamic_callback_names(first) == ["LowestTPMLoggingHandler_v2"] + assert (first.dynamic_async_success_callbacks or [None])[0] is (second.dynamic_async_success_callbacks or [])[0] + + +@pytest.mark.asyncio +async def test_an_override_matching_a_routing_group_strategy_is_not_bound_to_the_request(make_core: MakeCore) -> None: + core: Final = make_core( + model_list=[deployment("g", "a"), deployment("m", "b")], + routing_groups=[_group("gg", ["m"], "usage-based-routing-v2")], + ) + logging_obj: Final = _logging_obj() + await core.select("g", {"routing_strategy": "usage-based-routing-v2", "litellm_logging_obj": logging_obj}) + assert _dynamic_callback_names(logging_obj) == [] + + +def test_a_sync_usage_based_v2_override_counts_rpm_before_the_call(make_core: MakeCore, host: ScriptedHost) -> None: + """The override selector is on no callback list, so the sync pick runs its pre-call check itself: + the first call is counted against ``rpm=1`` and the next pick finds no deployment. Oddity: the sync + path raises a bare ``ValueError`` where the async one raises ``RateLimitError``.""" + core: Final = make_core(model_list=[deployment("g", "a", rpm=1)], num_retries=0) + core.completion(model="g", messages=MESSAGES, routing_strategy="usage-based-routing-v2") # pyright: ignore[reportAttributeAccessIssue] + with pytest.raises(ValueError, match=re.escape("No deployments available for selected model. Passed model=g.")): + core.completion(model="g", messages=MESSAGES, routing_strategy="usage-based-routing-v2") # pyright: ignore[reportAttributeAccessIssue] + assert host.deployments_called() == ["a"] + + +@pytest.mark.asyncio +async def test_an_async_usage_based_v2_override_counts_rpm_before_the_call( + make_core: MakeCore, host: ScriptedHost +) -> None: + core: Final = make_core(model_list=[deployment("g", "a", rpm=1)], num_retries=0) + await core.acompletion(model="g", messages=MESSAGES, routing_strategy="usage-based-routing-v2") # pyright: ignore[reportAttributeAccessIssue] + with pytest.raises(litellm.RateLimitError, match=re.escape("No deployments available for selected model.")): + await core.acompletion(model="g", messages=MESSAGES, routing_strategy="usage-based-routing-v2") # pyright: ignore[reportAttributeAccessIssue] + assert host.deployments_called() == ["a"] + + +@pytest.mark.asyncio +async def test_request_scoped_routing_still_runs_the_strategy_rpm_check( + make_core: MakeCore, host: ScriptedHost +) -> None: + """After ``observe(request_scoped=True)`` the selector is off the host's lists, but the router's + own pre-call checks still reach it, so ``rpm=1`` still holds.""" + collect_other_routers() + core: Final = make_core( + model_list=[deployment("g", "a", rpm=1)], routing_strategy="usage-based-routing-v2", num_retries=0 + ) + core.observe(request_scoped=True) + assert host.lists["callbacks"] == [] + await core.acompletion(model="g", messages=MESSAGES) # pyright: ignore[reportAttributeAccessIssue] + with pytest.raises(litellm.RateLimitError): + await core.acompletion(model="g", messages=MESSAGES) # pyright: ignore[reportAttributeAccessIssue] + assert host.deployments_called() == ["a"] + + +# --- strategy dispatch arms ----------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_async_pick_under_usage_based_routing_v1_picks_the_first_configured(make_core: MakeCore) -> None: + """The async pick calls v1's sync selector inline; with no traffic it picks the first deployment.""" + core: Final = make_core( + model_list=[deployment("g", "a"), deployment("g", "b")], routing_strategy="usage-based-routing" + ) + assert await picked_ids(core, "g", times=10, messages=MESSAGES) == {"a"} + + +def test_sync_completion_under_latency_based_routing_reaches_the_group(make_core: MakeCore, host: ScriptedHost) -> None: + core: Final = make_core( + model_list=[deployment("g", "a"), deployment("g", "b")], routing_strategy="latency-based-routing" + ) + for _ in range(10): + core.completion(model="g", messages=MESSAGES) # pyright: ignore[reportAttributeAccessIssue] + assert set(host.deployments_called()) <= {"a", "b"} + assert len(host.calls) == 10 + + +def test_sync_completion_under_cost_based_routing_finds_no_deployment(make_core: MakeCore, host: ScriptedHost) -> None: + """Oddity: the cost-based selector has no sync pick, so every sync pick under cost-based routing + rate-limits without calling a provider.""" + core: Final = make_core( + model_list=[deployment("g", "a"), deployment("g", "b")], routing_strategy="cost-based-routing", num_retries=0 + ) + with pytest.raises(RouterRateLimitError) as raised: + core.completion(model="g", messages=MESSAGES) # pyright: ignore[reportAttributeAccessIssue] + assert str(raised.value) == NO_DEPLOYMENTS_MESSAGE.format(model="g") + assert host.calls == [] + + +@pytest.mark.asyncio +async def test_pass_through_pick_under_cost_based_routing_finds_no_deployment(make_core: MakeCore) -> None: + core: Final = make_core( + model_list=[deployment("g", "a", use_in_pass_through=True)], routing_strategy="cost-based-routing" + ) + with pytest.raises(RouterRateLimitError): + await core.select("g", kind="pass_through") diff --git a/tests/unit/test_router/core/test_gaps_l11100_registry.py b/tests/unit/test_router/core/test_gaps_l11100_registry.py new file mode 100644 index 00000000000..808c12a693c --- /dev/null +++ b/tests/unit/test_router/core/test_gaps_l11100_registry.py @@ -0,0 +1,732 @@ +""" +Gap tests for the deployment registry of ``router_python_core.py`` (about lines 11550-14060): +provider validation in ``_add_deployment``, the os.environ references of Azure ``dataSources``, +vertex pass-through registration, config-sourced rows, index repair after a delete, cost-map +registration (``responses/`` keys, strategy-marker pricing, replay), deployment budgets kept in +step with upserts and deletes, credentials with provider, router / deployment model info, and how +model group info merges its deployments. +""" + +import os +import uuid +from collections.abc import Generator, Iterator +from typing import Final + +import pytest + +import litellm +from litellm.router_api import ChangeSet +from litellm.types.router import Deployment +from tests.unit.test_router.core.construction_precall_helpers import restoring_model_cost +from tests.unit.test_router.core.gaps_l11100_helpers import MakeCore, ModelsEndpoint, complexity_marker, model_ids +from tests.unit.test_router.core.registry_helpers import team_deployment +from tests.unit.test_router.core.scripted_host import ScriptedHost, deployment + + +def _unique(prefix: str = "l11100") -> str: + return f"{prefix}-{uuid.uuid4().hex}" + + +@pytest.fixture +def data_source_env(monkeypatch: pytest.MonkeyPatch) -> Iterator[None]: + monkeypatch.setenv("L11100_ENDPOINT", "https://search.example") + monkeypatch.delenv("L11100_MISSING", raising=False) + yield + + +# --- provider validation ------------------------------------------------------------------------------ + + +def test_a_provider_missing_from_the_hosts_provider_list_is_unsupported( + make_core: MakeCore, host: ScriptedHost +) -> None: + host.set_dynamic(provider_list=["anthropic"]) + + with pytest.raises(Exception, match=r"^Unsupported provider - openai$"): + make_core(model_list=[deployment("gpt", "a")]) + + +def test_ignore_invalid_deployments_drops_a_deployment_of_an_unsupported_provider( + make_core: MakeCore, host: ScriptedHost +) -> None: + host.set_dynamic(provider_list=["anthropic"]) + + core: Final = make_core( + model_list=[deployment("gpt", "a"), deployment("claude", "c", model="anthropic/claude-opus-4-6")], + ignore_invalid_deployments=True, + ) + + assert model_ids(core) == ["c"] + + +def test_a_prompt_management_model_skips_provider_validation(make_core: MakeCore, host: ScriptedHost) -> None: + """A ``/...`` model (here langfuse) is resolved from the prompt at call time, so it is + accepted even when no provider is supported at all.""" + host.set_dynamic(provider_list=[]) + + core: Final = make_core(model_list=[deployment("prompted", "p", model="langfuse/my-prompt")]) + + assert model_ids(core) == ["p"] + assert core.registry().names(kind="deployment_names") == ["langfuse/my-prompt"] + + +# --- Azure dataSources and vertex pass-through --------------------------------------------------------- + + +@pytest.mark.usefixtures("data_source_env") +def test_data_source_endpoint_and_key_env_references_are_resolved(make_core: MakeCore) -> None: + """Only ``endpoint`` and ``key`` are read from the environment; a missing variable becomes ``""``.""" + core: Final = make_core( + model_list=[ + deployment( + "gpt", + "a", + dataSources=[ + { + "type": "AzureCognitiveSearch", + "parameters": { + "endpoint": "os.environ/L11100_ENDPOINT", + "key": "os.environ/L11100_MISSING", + "indexName": "os.environ/L11100_ENDPOINT", + }, + } + ], + ) + ] + ) + + stored: Final = core.registry().all()[0]["litellm_params"]["dataSources"] + assert stored == [ + { + "type": "AzureCognitiveSearch", + "parameters": { + "endpoint": "https://search.example", + "key": "", + "indexName": "os.environ/L11100_ENDPOINT", + }, + } + ] + + +def test_a_vertex_pass_through_deployment_without_project_and_location_is_rejected(make_core: MakeCore) -> None: + with pytest.raises(ValueError, match="vertex_project") as raised: + make_core(model_list=[deployment("gem", "g", model="vertex_ai/gemini-pro", use_in_pass_through=True)]) + assert str(raised.value) == ( + "vertex_project, and vertex_location must be set in litellm_params for pass-through endpoints." + ) + + +def test_pass_through_needs_no_registration_for_other_providers(make_core: MakeCore) -> None: + core: Final = make_core(model_list=[deployment("gpt", "a", use_in_pass_through=True)]) + + assert model_ids(core) == ["a"] + + +# --- config rows and index repair ----------------------------------------------------------------------- + + +def test_config_deployments_leave_out_rows_that_came_from_the_db(make_core: MakeCore) -> None: + core: Final = make_core( + model_list=[ + deployment("gpt", "a"), + deployment("gpt", "b", model_info={"db_model": True}), + deployment("other", "c", model_info={"db_model": False}), + ] + ) + + assert [entry["model_info"]["id"] for entry in core.registry().deployments(source="config")] == ["a", "c"] + + +def test_deleting_a_later_deployment_keeps_earlier_team_indices_and_shifts_later_ones(make_core: MakeCore) -> None: + core: Final = make_core( + model_list=[ + team_deployment("internal-1", "t1", "TA", "pub"), + deployment("gpt", "b"), + team_deployment("internal-2", "t2", "TA", "other"), + ] + ) + + core.apply(ChangeSet(delete=["b"])) + + assert dict(core.registry().index("team_model")) == {("TA", "pub"): [0], ("TA", "other"): [1]} + + +def test_deleting_a_teams_only_wildcard_deployment_drops_its_team_pattern_router(make_core: MakeCore) -> None: + core: Final = make_core( + model_list=[team_deployment("internal", "tw", "TA", "openai/*"), deployment("gpt", "b")], + ) + assert list(core.registry().patterns(team=True)) == ["TA"] + + core.apply(ChangeSet(delete=["tw"])) + + assert dict(core.registry().patterns(team=True)) == {} + + +# --- cost-map registration --------------------------------------------------------------------------------- + + +@pytest.fixture +def catalog_backend_keys() -> Generator[None, None, None]: + yield from restoring_model_cost("gpt-4o-mini", "openai/gpt-4o-mini") + + +def _limits(key: str) -> tuple[object, object]: + entry: Final = litellm.model_cost.get(key) or {} + return entry.get("max_input_tokens"), entry.get("max_output_tokens") + + +@pytest.mark.usefixtures("catalog_backend_keys") +def test_a_deployments_token_limits_stay_off_a_shared_key_the_catalog_sets(make_core: MakeCore) -> None: + shared_keys: Final = ("gpt-4o-mini", "openai/gpt-4o-mini") + catalog_limits: Final = {key: _limits(key) for key in shared_keys} + override_id: Final = _unique("l11100-limits") + + make_core( + model_list=[ + deployment( + "g", + override_id, + model="openai/gpt-4o-mini", + model_info={"max_input_tokens": 4321, "max_output_tokens": 987, "max_tokens": 987}, + ) + ] + ) + + assert {key: _limits(key) for key in shared_keys} == catalog_limits + assert _limits(override_id) == (4321, 987) + + +def test_a_deployments_token_limits_fill_a_shared_key_the_catalog_does_not_list(make_core: MakeCore) -> None: + backend: Final = f"hosted_vllm/{_unique()}" + + make_core( + model_list=[ + deployment( + "g", + _unique("l11100-limits"), + model=backend, + model_info={"max_input_tokens": 4321, "max_output_tokens": 987}, + ) + ] + ) + + assert _limits(backend) == (4321, 987) + + +def test_a_responses_model_registers_its_backend_key_with_and_without_the_responses_segment( + make_core: MakeCore, +) -> None: + """The shared backend keys carry no custom pricing: that stays under the deployment id.""" + name: Final = _unique() + deployment_id: Final = _unique("l11100-id") + make_core(model_list=[deployment("gpt", deployment_id, model=f"openai/responses/{name}", input_cost_per_token=0.1)]) + + assert litellm.model_cost[f"openai/responses/{name}"].get("input_cost_per_token") is None + assert litellm.model_cost[f"openai/{name}"].get("input_cost_per_token") is None + assert litellm.model_cost[deployment_id]["input_cost_per_token"] == 0.1 + + +def test_custom_pricing_on_a_strategy_marker_is_not_registered_as_a_price(make_core: MakeCore) -> None: + """A marker is never the deployment that is billed, so even an explicit zero price stays out of its + cost-map entry.""" + marker_id: Final = _unique() + make_core( + model_list=[ + deployment("small", "s"), + deployment("big", "b"), + complexity_marker("smart", marker_id, input_cost_per_token=0.0, output_cost_per_token=0.0), + ] + ) + + entry: Final = litellm.model_cost[marker_id] + assert "input_cost_per_token" not in entry + assert "output_cost_per_token" not in entry + + +def test_replaying_cost_registrations_restores_the_entries_of_deployments_still_served(make_core: MakeCore) -> None: + kept: Final = _unique() + deleted: Final = _unique() + core: Final = make_core( + model_list=[ + deployment("gpt", kept, input_cost_per_token=0.25), + deployment("gpt", deleted, input_cost_per_token=0.5), + ] + ) + core.apply(ChangeSet(delete=[deleted])) + litellm.model_cost.pop(kept, None) + litellm.model_cost.pop(deleted, None) + + core.replay_model_cost_registrations() # pyright: ignore[reportAttributeAccessIssue] # public on the core + + assert litellm.model_cost[kept]["input_cost_per_token"] == 0.25 + assert deleted not in litellm.model_cost + + +# --- deployment budgets -------------------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_deployment_budgets_from_the_model_list_are_registered_at_construction(make_core: MakeCore) -> None: + core: Final = make_core( + model_list=[ + deployment("gpt", "a", max_budget=1.0, budget_duration="1d"), + deployment("gpt", "b", max_budget=2.0, budget_duration="1d"), + deployment("gpt", "c"), + ] + ) + + limiter: Final = core.registry().plugins("budget_limiter") + assert limiter is not None + assert limiter.deployment_budget_config is not None + assert {key: config.max_budget for key, config in limiter.deployment_budget_config.items()} == {"a": 1.0, "b": 2.0} + + +@pytest.mark.asyncio +async def test_deleting_and_upserting_deployments_keep_the_budget_limiter_in_step(make_core: MakeCore) -> None: + core: Final = make_core( + model_list=[ + deployment("gpt", "a", max_budget=1.0, budget_duration="1d"), + deployment("gpt", "b", max_budget=2.0, budget_duration="1d"), + ] + ) + limiter: Final = core.registry().plugins("budget_limiter") + assert limiter is not None + + core.apply(ChangeSet(delete=["a"])) + assert list(limiter.deployment_budget_config or {}) == ["b"] + + core.apply(ChangeSet(upsert=[Deployment(**deployment("gpt", "b"))])) # pyright: ignore[reportArgumentType] + assert limiter.deployment_budget_config is None + + core.apply(ChangeSet(upsert=[Deployment(**deployment("gpt", "b", max_budget=3.0, budget_duration="1d"))])) # pyright: ignore[reportArgumentType] + assert {key: config.max_budget for key, config in (limiter.deployment_budget_config or {}).items()} == {"b": 3.0} + + +@pytest.mark.asyncio +@pytest.mark.parametrize("params", [{"max_budget": 1.0}, {"budget_duration": "1d"}], ids=["no-duration", "no-budget"]) +async def test_adding_a_deployment_with_half_a_budget_installs_no_budget_limiter( + make_core: MakeCore, params: dict[str, object] +) -> None: + core: Final = make_core(model_list=[deployment("gpt", "a")]) + + core.add_deployment(Deployment(**deployment("gpt", "half", **params))) # pyright: ignore[reportAttributeAccessIssue, reportArgumentType] + + assert core.registry().plugins("budget_limiter") is None + + +# --- credentials with provider ------------------------------------------------------------------------------- + + +@pytest.mark.parametrize( + ("entry", "expected_provider"), + [ + (deployment("gpt", "a", model="gpt-4o"), "openai"), + (deployment("gpt", "a", model="gpt-4o", custom_llm_provider="azure", api_base="https://x"), "azure"), + (deployment("gpt", "a", model="anthropic/claude-opus-4-6"), "anthropic"), + ], + ids=["bare-model-defaults-to-openai", "explicit-provider-wins", "provider-prefix"], +) +def test_credentials_with_provider_name_the_provider( + make_core: MakeCore, entry: dict[str, object], expected_provider: str +) -> None: + found: Final = make_core(model_list=[entry]).registry().credentials("a", form="with_provider") + + assert found is not None + assert found["custom_llm_provider"] == expected_provider + + +def test_credentials_with_provider_drop_an_unknown_credential_name(make_core: MakeCore) -> None: + found: Final = ( + make_core(model_list=[deployment("gpt", "a", litellm_credential_name="l11100-no-such-credential")]) + .registry() + .credentials("a", form="with_provider") + ) + + assert found == {"api_key": "sk-test", "model": "openai/a", "custom_llm_provider": "openai"} + + +# --- router and deployment model info ------------------------------------------------------------------------ + + +@pytest.mark.parametrize( + ("model", "model_info", "expected_key"), + [ + ("azure/gpt-4o", {}, "azure/gpt-4o"), + ("azure/my-deployment", {"base_model": "azure/gpt-4o"}, "azure/gpt-4o"), + ], + ids=["mapped-deployment-name", "base-model"], +) +def test_router_model_info_of_an_azure_deployment( + make_core: MakeCore, model: str, model_info: dict[str, object], expected_key: str +) -> None: + core: Final = make_core( + model_list=[ + deployment("az", "az", model=model, api_base="https://x", api_version="2024", model_info=model_info) + ] + ) + + info: Final = core.registry().router_model_info(None, "az", id="az") + builtin: Final = litellm.get_model_info("azure/gpt-4o") + assert info["key"] == expected_key + assert info["max_input_tokens"] == builtin["max_input_tokens"] + assert info["input_cost_per_token"] == builtin["input_cost_per_token"] + + +def test_router_model_info_of_an_unmapped_azure_deployment_is_its_zeroed_stub(make_core: MakeCore) -> None: + core: Final = make_core( + model_list=[deployment("az", "az", model="azure/l11100-unmapped", api_base="https://x", api_version="2024")] + ) + + info: Final = core.registry().router_model_info(None, "az", id="az") + assert info["key"] == "azure/l11100-unmapped" + assert info["max_input_tokens"] is None + assert info["input_cost_per_token"] == 0 + + +def test_deployment_model_info_takes_defaults_from_a_custom_base_model(make_core: MakeCore) -> None: + """The deployment's own pricing wins; the rest comes from its ``base_model``'s built-in entry.""" + name: Final = _unique() + core: Final = make_core( + model_list=[ + deployment( + "gpt", name, model=f"openai/{name}", input_cost_per_token=0.5, model_info={"base_model": "gpt-4o"} + ) + ] + ) + + info: Final = core.registry().deployment_model_info(name, f"openai/{name}") + builtin: Final = litellm.get_model_info("gpt-4o") + assert info is not None + assert info["input_cost_per_token"] == 0.5 + assert info["output_cost_per_token"] == builtin["output_cost_per_token"] + assert info["max_input_tokens"] == builtin["max_input_tokens"] + + +def test_deployment_model_info_of_an_unpriced_unmapped_model_reports_no_prices(make_core: MakeCore) -> None: + """``get_model_info`` reads a missing price as 0; the deployment reports ``None`` instead.""" + name: Final = _unique() + core: Final = make_core(model_list=[deployment("gpt", name, model=f"openai/{name}")]) + + info: Final = core.registry().deployment_model_info(name, f"openai/{name}") + assert info is not None + assert (info["input_cost_per_token"], info["output_cost_per_token"]) == (None, None) + + +# --- model group info ------------------------------------------------------------------------------------------ + + +def test_group_info_takes_the_widest_limits_highest_prices_and_any_capability(make_core: MakeCore) -> None: + core: Final = make_core( + model_list=[ + deployment("g", "a", model="openai/gpt-3.5-turbo"), + deployment("g", "b", model="openai/gpt-4o"), + deployment("g", "c", model="openai/o3-mini"), + deployment("g", "d", model="anthropic/claude-opus-4-6"), + ] + ) + + info: Final = core.registry().group("g") + opus: Final = litellm.get_model_info("anthropic/claude-opus-4-6") + assert info is not None + assert info.providers == ["openai", "anthropic"] + assert (info.max_input_tokens, info.max_output_tokens) == (opus["max_input_tokens"], opus["max_output_tokens"]) + assert (info.input_cost_per_token, info.output_cost_per_token) == ( + opus["input_cost_per_token"], + opus["output_cost_per_token"], + ) + assert info.supports_vision is True + assert info.supports_function_calling is True + assert info.supports_parallel_function_calling is True + assert info.supports_web_search is True + assert info.supports_reasoning is True + assert info.supports_url_context is False + assert info.supports_fast_mode is False + + +def test_group_info_supports_url_context_when_any_deployment_does(make_core: MakeCore) -> None: + core: Final = make_core( + model_list=[deployment("g", "a", model="openai/gpt-4o"), deployment("g", "b", model="gemini/gemini-2.5-flash")] + ) + + info: Final = core.registry().group("g") + assert info is not None + assert info.supports_url_context is True + + +def test_group_info_supports_continuation_when_any_deployment_does(make_core: MakeCore) -> None: + core: Final = make_core( + model_list=[ + deployment("g", "a", model="openai/gpt-4o"), + deployment("g", "b", model="openai/gpt-4o-mini", model_info={"supports_continuation": True}), + deployment("plain", "c", model="openai/gpt-4o"), + ] + ) + + grouped: Final = core.registry().group("g") + plain: Final = core.registry().group("plain") + assert grouped is not None and grouped.supports_continuation is True + assert plain is not None and plain.supports_continuation is False + + +@pytest.mark.parametrize("value", ["maybe", "", ["x"], "true", 1]) +def test_group_info_counts_only_a_boolean_true_supports_continuation(make_core: MakeCore, value: object) -> None: + core: Final = make_core( + model_list=[deployment("g", "a", model="openai/gpt-4o", model_info={"supports_continuation": value})] + ) + + info: Final = core.registry().group("g") + assert info is not None + assert info.supports_continuation is False + + +@pytest.mark.parametrize( + ("models", "fast"), + [ + (["anthropic/claude-opus-4-8"], True), + (["anthropic/claude-opus-4-8", "anthropic/claude-opus-4-6"], False), + ], + ids=["every-deployment-fast", "one-not-fast"], +) +def test_group_info_supports_fast_mode_only_when_every_deployment_does( + make_core: MakeCore, models: list[str], fast: bool +) -> None: + core: Final = make_core(model_list=[deployment("g", f"d{i}", model=m) for i, m in enumerate(models)]) + + info: Final = core.registry().group("g") + assert info is not None + assert info.supports_fast_mode is fast + + +@pytest.mark.parametrize( + ("models", "efforts"), + [ + (["anthropic/claude-opus-4-8"], ("none", "minimal", "low", "medium", "high", "xhigh", "max")), + ( + ["anthropic/claude-opus-4-8", "anthropic/claude-opus-4-6"], + ("none", "minimal", "low", "medium", "high", "max"), + ), + (["anthropic/claude-sonnet-4-5-20250929", "anthropic/claude-opus-4-6"], None), + (["anthropic/claude-opus-4-6", "anthropic/claude-sonnet-4-5-20250929"], None), + (["openai/gpt-4o", "openai/gpt-4o"], ()), + ], + ids=["one", "intersection", "unknown-first-stays-unknown", "unknown-later-clears", "none-supported"], +) +def test_group_info_reasoning_efforts_are_the_intersection_or_unknown( + make_core: MakeCore, models: list[str], efforts: tuple[str, ...] | None +) -> None: + core: Final = make_core(model_list=[deployment("g", f"d{i}", model=m) for i, m in enumerate(models)]) + + info: Final = core.registry().group("g") + assert info is not None + assert info.supported_reasoning_efforts == efforts + + +def test_group_info_supported_endpoints_are_the_union_in_deployment_order(make_core: MakeCore) -> None: + core: Final = make_core( + model_list=[ + deployment("g", "a", model="openai/gpt-5"), + deployment("g", "b", model="openai/gpt-4o", model_info={"supported_endpoints": ["/v1/embeddings"]}), + ] + ) + + info: Final = core.registry().group("g") + assert info is not None + assert info.supported_endpoints == ("/v1/chat/completions", "/v1/batch", "/v1/responses", "/v1/embeddings") + + +def test_group_info_of_a_prompt_management_model_names_its_logger_as_provider(make_core: MakeCore) -> None: + core: Final = make_core(model_list=[deployment("prompted", "p", model="langfuse/my-prompt")]) + + info: Final = core.registry().group("prompted") + assert info is not None + assert info.providers == ["langfuse"] + + +def test_group_info_of_an_alias_to_neither_a_name_nor_a_dict_is_none(make_core: MakeCore) -> None: + core: Final = make_core(model_list=[deployment("gpt", "a")]) + core.apply(ChangeSet(assign={"model_group_alias": {"odd": 5}})) + + assert core.registry().group("odd") is None + + +def test_group_info_of_an_unpriced_unmapped_model_has_no_prices_or_limits(make_core: MakeCore) -> None: + name: Final = _unique() + core: Final = make_core(model_list=[deployment("g", name, model=f"openai/{name}")]) + + info: Final = core.registry().group("g") + assert info is not None + assert (info.input_cost_per_token, info.output_cost_per_token, info.max_input_tokens) == (None, None, None) + + +# --- lookups --------------------------------------------------------------------------------------------------- + + +def test_from_id_of_a_deployment_id_is_its_group_name(make_core: MakeCore) -> None: + core: Final = make_core(model_list=[deployment("gpt", "a")]) + + assert core.registry().resolve("a", kind="from_id") == "gpt" + + +def test_from_id_on_an_empty_core_or_of_an_empty_id_is_none(make_core: MakeCore) -> None: + assert make_core(model_list=[]).registry().resolve("x", kind="from_id") is None + assert make_core(model_list=[deployment("gpt", "a")]).registry().resolve("", kind="from_id") is None + + +def test_route_candidates_of_a_routing_group_are_its_members_deployments(make_core: MakeCore) -> None: + core: Final = make_core( + model_list=[deployment("g", "a"), deployment("h", "b"), deployment("other", "c")], + routing_groups=[{"group_name": "rg", "models": ["g", "h"], "routing_strategy": "simple-shuffle"}], + ) + + assert core.registry().names("rg", kind="route_candidates") == frozenset({"a", "b"}) + + +def test_search_tool_support_reads_model_info_then_the_anthropic_route(make_core: MakeCore) -> None: + core: Final = make_core( + model_list=[ + deployment("g", "a", model="anthropic/claude-opus-4-6"), + deployment("g", "b", model="openai/gpt-4o", model_info={"supports_search_tool": False}), + deployment("g", "c", model="openai/gpt-4o"), + ] + ) + + assert core.registry().resolve("g", kind="search_tool_support") == (True, False, None) + + +# --- discovered model info (aapply refresh_model_info) ---------------------------------------------------- + +VLLM_CARDS: Final = {"data": [{"id": "my-model", "max_model_len": 8192}, {"id": "other", "max_model_len": 4096}]} + + +def _vllm( + deployment_id: str = "v", **params: object +) -> dict[str, object]: # mutable-ok: a model_list entry # kwargs-ok: litellm_params + return deployment( + "vllm", deployment_id, model="hosted_vllm/my-model", api_base="http://vllm.local", api_key="sk-vllm", **params + ) + + +@pytest.mark.asyncio +async def test_a_refresh_reads_the_upstream_models_endpoint_with_the_deployments_key(make_core: MakeCore) -> None: + endpoint: Final = ModelsEndpoint(VLLM_CARDS) + core: Final = make_core(model_list=[_vllm(extra_headers={"X-Team": "a"})]) + + await core.aapply(ChangeSet(refresh_model_info=True, refresh_client=endpoint)) # pyright: ignore[reportAttributeAccessIssue, reportArgumentType] + + assert endpoint.requests == [("http://vllm.local/v1/models", {"authorization": "Bearer sk-vllm", "x-team": "a"})] + assert dict(core.registry().discovered_model_info("v")) == { + "max_tokens": 8192, + "max_input_tokens": 8192, + "max_output_tokens": 8192, + } + + +@pytest.mark.asyncio +async def test_discovered_limits_never_override_configured_ones_and_feed_the_group_limits(make_core: MakeCore) -> None: + endpoint: Final = ModelsEndpoint(VLLM_CARDS) + core: Final = make_core(model_list=[_vllm(model_info={"max_output_tokens": 1000})]) + + await core.aapply(ChangeSet(refresh_model_info=True, refresh_client=endpoint)) # pyright: ignore[reportAttributeAccessIssue, reportArgumentType] + + registry: Final = core.registry() + assert dict(registry.discovered_model_info("v")) == {"max_tokens": 8192, "max_input_tokens": 8192} + assert registry.group("vllm", field="token_limits") == (8192, 1000) + listing: Final = registry.group("vllm", field="listing_info") + assert listing is not None + assert (listing.max_input_tokens, listing.max_output_tokens) == (8192, 1000) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "entry", + [ + deployment("plain", "p", model="anthropic/claude-opus-4-6", api_base="http://vllm.local"), + deployment("wild", "w", model="hosted_vllm/*", api_base="http://vllm.local"), + deployment("nobase", "n", model="openai/gpt-4o"), + ], + ids=["provider-without-discovery", "wildcard-model", "no-api-base"], +) +async def test_a_refresh_skips_deployments_it_cannot_discover(make_core: MakeCore, entry: dict[str, object]) -> None: + endpoint: Final = ModelsEndpoint(VLLM_CARDS) + core: Final = make_core(model_list=[entry]) + + await core.aapply(ChangeSet(refresh_model_info=True, refresh_client=endpoint)) # pyright: ignore[reportAttributeAccessIssue, reportArgumentType] + + assert endpoint.requests == [] + + +@pytest.mark.asyncio +async def test_an_unreachable_models_endpoint_discovers_nothing(make_core: MakeCore) -> None: + endpoint: Final = ModelsEndpoint(error=ConnectionError("down")) + core: Final = make_core(model_list=[_vllm()]) + + await core.aapply(ChangeSet(refresh_model_info=True, refresh_client=endpoint)) # pyright: ignore[reportAttributeAccessIssue, reportArgumentType] + + assert len(endpoint.requests) == 1 + assert dict(core.registry().discovered_model_info("v")) == {} + + +@pytest.mark.asyncio +async def test_discovered_limits_are_forgotten_when_the_deployment_is_replaced(make_core: MakeCore) -> None: + endpoint: Final = ModelsEndpoint(VLLM_CARDS) + core: Final = make_core(model_list=[_vllm()]) + await core.aapply(ChangeSet(refresh_model_info=True, refresh_client=endpoint)) # pyright: ignore[reportAttributeAccessIssue, reportArgumentType] + + core.apply(ChangeSet(upsert=[Deployment(**_vllm(rpm=5))])) # pyright: ignore[reportArgumentType] + + assert dict(core.registry().discovered_model_info("v")) == {} + + +# --- failed upserts and replays under ignore_invalid_deployments --------------------------------------------- + + +def test_a_failing_upsert_of_a_new_deployment_under_ignore_invalid_deployments_adds_nothing( + make_core: MakeCore, host: ScriptedHost +) -> None: + core: Final = make_core(model_list=[deployment("gpt", "a")], ignore_invalid_deployments=True) + host.set_dynamic(provider_list=["anthropic"]) + + result: Final = core.apply(ChangeSet(upsert=[Deployment(**deployment("gpt", "new"))])) # pyright: ignore[reportArgumentType] + + assert list(result.upserted) == [None] + assert model_ids(core) == ["a"] + + +def test_a_replay_retries_deployments_whose_provider_did_not_resolve_and_keeps_them_out_while_it_still_fails( + make_core: MakeCore, +) -> None: + core: Final = make_core( + model_list=[deployment("gpt", "a"), deployment("odd", "u", model="l11100-no-provider/x")], + ignore_invalid_deployments=True, + ) + assert model_ids(core) == ["a"] + + core.replay_model_cost_registrations() # pyright: ignore[reportAttributeAccessIssue] # public on the core + + assert model_ids(core) == ["a"] + + +# --- group info without any model info ------------------------------------------------------------------------ + + +def test_group_info_without_a_cost_map_entry_falls_back_to_the_configured_mode_and_costs(make_core: MakeCore) -> None: + """With the entries construction registered withdrawn from the cost map (as a price data reload + does before the replay) and a model litellm does not map, the group is described from the + deployment's own model_info: its mode, its costs (a string cost read as a float, a missing one as + None) and the provider's supported params, with no token limits.""" + name: Final = _unique() + core: Final = make_core( + model_list=[ + deployment( + "g", name, model=f"openai/{name}", model_info={"mode": "embedding", "input_cost_per_token": "0.5"} + ) + ] + ) + litellm.model_cost.pop(name, None) + litellm.model_cost.pop(f"openai/{name}", None) + + info: Final = core.registry().group("g") + assert info is not None + assert (info.mode, info.input_cost_per_token, info.output_cost_per_token) == ("embedding", 0.5, None) + assert (info.max_input_tokens, info.max_output_tokens) == (None, None) + assert info.supported_openai_params diff --git a/tests/unit/test_router/core/test_gaps_l11100_strategy.py b/tests/unit/test_router/core/test_gaps_l11100_strategy.py new file mode 100644 index 00000000000..2459c577b9c --- /dev/null +++ b/tests/unit/test_router/core/test_gaps_l11100_strategy.py @@ -0,0 +1,438 @@ +""" +Gap tests for the strategy-router (pre-routing) deployments of ``router_python_core.py`` (about lines +11114-11550, plus the strategy hooks of set_model_list / upsert / delete): construction-time +validation of auto, adaptive and quality router markers, the (model_name, tags) slot each marker +holds and when it is released, the licensed capability ceiling, and the adaptive router's derived +inputs and post-call hooks. +""" + +from typing import Final + +import pytest + +from litellm.router_api import ChangeSet +from litellm.types.router import Deployment +from tests.unit.test_router.core.gaps_l11100_helpers import ( + TIERS, + MakeCore, + RemovingHost, + adaptive_marker, + complexity_marker, + marker, + model_ids, + registered_tags, +) +from tests.unit.test_router.core.scripted_host import ScriptedHost, deployment + + +@pytest.fixture +def host() -> ScriptedHost: + """The adaptive router's hook sync removes callbacks from every list, so the host must support it.""" + return RemovingHost() + + +PLAIN: Final = [deployment("small", "s"), deployment("big", "b")] + +HEURISTIC_V2: Final = {"classifier_type": "heuristic_v2", "tiers": dict(TIERS)} + +HEURISTIC_V2_LIMIT_MESSAGE: Final = ( + "At most 1 auto-router(s) using the heuristic_v2 classifier can be registered but this would make 2. " + "Choose the heuristic v1 classifier for this router or remove an existing heuristic_v2 router." +) + + +# --- marker validation at construction ------------------------------------------------------------- + + +@pytest.mark.parametrize( + ("params", "message"), + [ + ( + {}, + "auto_router_config_path or auto_router_config is required for auto-router deployments. " + "Please set it in the litellm_params", + ), + ( + {"auto_router_config": "{}"}, + "auto_router_default_model is required for auto-router deployments. Please set it in the litellm_params", + ), + ( + {"auto_router_config": "{}", "auto_router_default_model": "big"}, + "auto_router_embedding_model is required for auto-router deployments. Please set it in the litellm_params", + ), + ], + ids=["no-config", "no-default-model", "no-embedding-model"], +) +def test_an_auto_router_marker_missing_a_required_param_is_rejected( + make_core: MakeCore, params: dict[str, object], message: str +) -> None: + with pytest.raises(ValueError, match="is required") as raised: + make_core(model_list=[*PLAIN, marker("auto", "m", "auto_router/semantic", **params)]) + assert str(raised.value) == message + + +def test_an_adaptive_router_marker_without_its_config_is_rejected(make_core: MakeCore) -> None: + with pytest.raises(ValueError, match="is required") as raised: + make_core(model_list=[*PLAIN, marker("ad", "m", "auto_router/adaptive_router")]) + assert str(raised.value) == "adaptive_router_config is required for adaptive-router deployments." + + +def test_a_quality_router_marker_without_a_default_model_is_rejected(make_core: MakeCore) -> None: + with pytest.raises(ValueError, match="is required") as raised: + make_core(model_list=[*PLAIN, marker("q", "m", "auto_router/quality_router")]) + assert str(raised.value) == ( + "quality_router_default_model is required for quality-router deployments, or set default_model in " + "quality_router_config. Please configure it in the litellm_params" + ) + + +def test_a_quality_router_takes_its_default_model_from_its_config(make_core: MakeCore) -> None: + core: Final = make_core( + model_list=[ + *PLAIN, + marker("q", "m", "auto_router/quality_router", quality_router_config={"default_model": "big"}), + ] + ) + registry: Final = core.registry().plugins("quality") + + assert registered_tags(core, "quality") == {"q": [()]} + assert registry["q"][0].strategy.config.default_model == "big" + + +def test_quality_router_default_model_param_beats_the_one_in_its_config(make_core: MakeCore) -> None: + core: Final = make_core( + model_list=[ + *PLAIN, + marker( + "q", + "m", + "auto_router/quality_router", + quality_router_default_model="small", + quality_router_config={"default_model": "big"}, + ), + ] + ) + + assert core.registry().plugins("quality")["q"][0].strategy.config.default_model == "small" + + +# --- the (model_name, tags) slot -------------------------------------------------------------------- + + +@pytest.mark.parametrize(("tags", "shown"), [(None, "[]"), (["x"], "['x']")], ids=["untagged", "tagged"]) +def test_a_second_marker_with_the_same_name_and_tags_is_rejected( + make_core: MakeCore, tags: list[str] | None, shown: str +) -> None: + extra: Final = {} if tags is None else {"tags": tags} + with pytest.raises(ValueError, match="already exists") as raised: + make_core( + model_list=[*PLAIN, complexity_marker("smart", "m1", **extra), complexity_marker("smart", "m2", **extra)] + ) + assert str(raised.value) == ( + f"Complexity-router deployment smart with tags {shown} already exists. " + "Please use a different model name or set different tags." + ) + + +def test_ignore_invalid_deployments_drops_the_repeated_marker_and_keeps_the_first(make_core: MakeCore) -> None: + core: Final = make_core( + model_list=[*PLAIN, complexity_marker("smart", "m1"), complexity_marker("smart", "m2")], + ignore_invalid_deployments=True, + ) + + assert model_ids(core) == ["s", "b", "m1"] + assert registered_tags(core, "complexity") == {"smart": [()]} + + +def test_markers_of_one_name_with_different_tags_each_register(make_core: MakeCore) -> None: + core: Final = make_core( + model_list=[*PLAIN, complexity_marker("smart", "m1", tags=["a"]), complexity_marker("smart", "m2", tags=["b"])] + ) + + assert registered_tags(core, "complexity") == {"smart": [("a",), ("b",)]} + + +def test_adding_a_marker_whose_slot_is_taken_raises_and_adds_nothing(make_core: MakeCore) -> None: + core: Final = make_core(model_list=[*PLAIN, complexity_marker("smart", "m1")]) + + with pytest.raises(ValueError, match="already exists"): + core.add_deployment(Deployment(**complexity_marker("smart", "m2"))) # pyright: ignore[reportAttributeAccessIssue, reportArgumentType] + assert model_ids(core) == ["s", "b", "m1"] + assert registered_tags(core, "complexity") == {"smart": [()]} + + +def test_deleting_a_marker_frees_its_slot_for_a_new_marker(make_core: MakeCore) -> None: + core: Final = make_core(model_list=[*PLAIN, complexity_marker("smart", "m1")]) + + core.apply(ChangeSet(delete=["m1"])) + assert registered_tags(core, "complexity") == {} + result: Final = core.apply(ChangeSet(upsert=[Deployment(**complexity_marker("smart", "m2"))])) # pyright: ignore[reportArgumentType] + + assert result.upserted[0] is not None + assert registered_tags(core, "complexity") == {"smart": [()]} + assert model_ids(core) == ["s", "b", "m2"] + + +def test_deleting_one_tagged_marker_keeps_the_other_tags_of_its_name(make_core: MakeCore) -> None: + core: Final = make_core( + model_list=[*PLAIN, complexity_marker("smart", "m1", tags=["a"]), complexity_marker("smart", "m2", tags=["b"])] + ) + + core.apply(ChangeSet(delete=["m1"])) + + assert registered_tags(core, "complexity") == {"smart": [("b",)]} + + +def test_deleting_a_plain_deployment_sharing_a_markers_name_keeps_the_marker(make_core: MakeCore) -> None: + core: Final = make_core(model_list=[*PLAIN, complexity_marker("smart", "m1"), deployment("smart", "p")]) + + core.apply(ChangeSet(delete=["p"])) + + assert registered_tags(core, "complexity") == {"smart": [()]} + + +def test_an_upsert_that_retags_a_marker_moves_it_to_the_new_slot(make_core: MakeCore) -> None: + core: Final = make_core(model_list=[*PLAIN, complexity_marker("smart", "m1")]) + + core.apply(ChangeSet(upsert=[Deployment(**complexity_marker("smart", "m1", tags=["t"]))])) # pyright: ignore[reportArgumentType] + + assert registered_tags(core, "complexity") == {"smart": [("t",)]} + assert model_ids(core) == ["s", "b", "m1"] + + +def test_set_model_list_drops_the_strategy_routers_of_the_old_list(make_core: MakeCore) -> None: + core: Final = make_core( + model_list=[ + *PLAIN, + complexity_marker("smart", "m1"), + marker("q", "m2", "auto_router/quality_router", quality_router_default_model="big"), + ] + ) + + core.set_model_list(list(PLAIN)) # pyright: ignore[reportAttributeAccessIssue] + + assert registered_tags(core, "complexity") == {} + assert registered_tags(core, "quality") == {} + + +# --- the licensed capability ceiling ----------------------------------------------------------------- + + +def test_a_router_beyond_the_capability_limit_is_rejected(make_core: MakeCore) -> None: + with pytest.raises(ValueError, match="At most 1") as raised: + make_core( + model_list=[ + *PLAIN, + complexity_marker("one", "m1", config=HEURISTIC_V2), + complexity_marker("two", "m2", config=HEURISTIC_V2), + ], + auto_router_capability_limit=lambda: 1, + ) + assert str(raised.value) == HEURISTIC_V2_LIMIT_MESSAGE + + +def test_routers_within_the_capability_limit_or_without_the_capability_register(make_core: MakeCore) -> None: + core: Final = make_core( + model_list=[ + *PLAIN, + complexity_marker("one", "m1", config=HEURISTIC_V2), + complexity_marker("two", "m2"), + complexity_marker("three", "m3"), + ], + auto_router_capability_limit=lambda: 1, + ) + + assert registered_tags(core, "complexity") == {"one": [()], "two": [()], "three": [()]} + + +def test_without_a_capability_limit_any_number_of_gated_routers_register(make_core: MakeCore) -> None: + core: Final = make_core( + model_list=[ + *PLAIN, + complexity_marker("one", "m1", config=HEURISTIC_V2), + complexity_marker("two", "m2", config=HEURISTIC_V2), + ] + ) + + assert sorted(registered_tags(core, "complexity")) == ["one", "two"] + + +def test_the_capability_limit_is_read_again_on_every_registration(make_core: MakeCore) -> None: + limit: Final = {"value": 2} + core: Final = make_core( + model_list=[ + *PLAIN, + complexity_marker("one", "m1", config=HEURISTIC_V2), + complexity_marker("two", "m2", config=HEURISTIC_V2), + ], + auto_router_capability_limit=lambda: limit["value"], + ) + limit["value"] = 1 + + with pytest.raises(ValueError, match="At most 1") as raised: + core.add_deployment(Deployment(**complexity_marker("three", "m3", config=HEURISTIC_V2))) # pyright: ignore[reportAttributeAccessIssue, reportArgumentType] + assert str(raised.value) == HEURISTIC_V2_LIMIT_MESSAGE.replace("would make 2", "would make 3") + + +def test_a_failed_edit_rolls_back_to_the_serving_router_even_above_a_tightened_limit(make_core: MakeCore) -> None: + """The edited router is popped first, so the edit itself counts 1 other + itself = 2 > 1 and fails; + the rollback re-adds the previous router with the ceiling switched off.""" + limit: Final = {"value": 2} + core: Final = make_core( + model_list=[ + *PLAIN, + complexity_marker("one", "m1", config=HEURISTIC_V2), + complexity_marker("two", "m2", config=HEURISTIC_V2), + ], + auto_router_capability_limit=lambda: limit["value"], + ignore_invalid_deployments=True, + ) + limit["value"] = 1 + edited: Final = complexity_marker("two", "m2", config=HEURISTIC_V2, tags=["edited"]) + + result: Final = core.apply(ChangeSet(upsert=[Deployment(**edited)])) # pyright: ignore[reportArgumentType] + + assert list(result.upserted) == [None] + assert model_ids(core) == ["s", "b", "m1", "m2"] + assert registered_tags(core, "complexity") == {"one": [()], "two": [()]} + restored: Final = core.registry().deployment("m2") + assert restored is not None + assert restored.litellm_params.tags is None + + +def test_a_failed_edit_without_ignore_invalid_deployments_raises_and_leaves_the_router_removed( + make_core: MakeCore, +) -> None: + limit: Final = {"value": 2} + core: Final = make_core( + model_list=[ + *PLAIN, + complexity_marker("one", "m1", config=HEURISTIC_V2), + complexity_marker("two", "m2", config=HEURISTIC_V2), + ], + auto_router_capability_limit=lambda: limit["value"], + ) + limit["value"] = 1 + + with pytest.raises(ValueError, match="At most 1") as raised: + core.apply( + ChangeSet(upsert=[Deployment(**complexity_marker("two", "m2", config=HEURISTIC_V2, tags=["edited"]))]) # pyright: ignore[reportArgumentType] + ) + assert str(raised.value) == HEURISTIC_V2_LIMIT_MESSAGE + assert model_ids(core) == ["s", "b", "m1"] + assert registered_tags(core, "complexity") == {"one": [()]} + + +# --- the adaptive router ------------------------------------------------------------------------------- + + +def _adaptive_list() -> list[dict[str, object]]: # mutable-ok: a model_list + return [ + adaptive_marker("ar", "ar1", ["x", "y", "z"]), + deployment( + "x", + "x", + input_cost_per_token=0.5, + model_info={"input_cost_per_token": 0.9, "adaptive_router_preferences": {"quality_tier": 2}}, + ), + deployment("y", "y", model_info={"input_cost_per_token": 0.25}), + ] + + +def _hook_counts(host: ScriptedHost) -> tuple[int, int]: + from litellm.router_strategy.adaptive_router.hooks import ( + AdaptiveRouterPostCallHook, + AdaptiveRouterResponseHeaderHook, + ) + + callbacks: Final = host.lists["callbacks"] + return ( + sum(isinstance(callback, AdaptiveRouterPostCallHook) for callback in callbacks), + sum(isinstance(callback, AdaptiveRouterResponseHeaderHook) for callback in callbacks), + ) + + +def test_an_adaptive_router_is_built_after_the_whole_list_from_the_models_it_names(make_core: MakeCore) -> None: + """The marker is listed before x and y, yet sees them: litellm_params pricing wins over model_info + pricing, preferences come from model_info, and z (not configured) gets neither.""" + core: Final = make_core(model_list=_adaptive_list()) + registry: Final = core.registry().plugins("adaptive") + + assert registered_tags(core, "adaptive") == {"ar": [()]} + strategy: Final = registry["ar"][0].strategy + assert strategy.model_to_cost == {"x": 0.5, "y": 0.25} + assert {name: prefs.quality_tier for name, prefs in strategy.model_to_prefs.items()} == {"x": 2} + + +def test_an_adaptive_router_registers_one_post_call_hook_and_a_header_hook( + make_core: MakeCore, host: ScriptedHost +) -> None: + core: Final = make_core(model_list=_adaptive_list()) + if not registered_tags(core, "adaptive"): + pytest.fail("no adaptive router was registered") + + post_call, header = _hook_counts(host) + assert post_call == 1 + assert header >= 1 + + +def test_deleting_an_adaptive_router_releases_it_and_its_post_call_hook( + make_core: MakeCore, host: ScriptedHost +) -> None: + core: Final = make_core(model_list=_adaptive_list()) + + core.apply(ChangeSet(delete=["ar1"])) + + assert registered_tags(core, "adaptive") == {} + assert _hook_counts(host)[0] == 0 + + +def test_a_complexity_router_with_adaptive_enabled_also_registers_an_adaptive_router(make_core: MakeCore) -> None: + core: Final = make_core( + model_list=[*PLAIN, complexity_marker("smart", "m1", config={"tiers": TIERS, "adaptive": True})] + ) + + assert registered_tags(core, "complexity") == {"smart": [()]} + assert registered_tags(core, "adaptive") == {"smart": [()]} + + +def test_deleting_an_adaptive_complexity_router_releases_both_registrations(make_core: MakeCore) -> None: + core: Final = make_core( + model_list=[*PLAIN, complexity_marker("smart", "m1", config={"tiers": TIERS, "adaptive": True})] + ) + + core.apply(ChangeSet(delete=["m1"])) + + assert registered_tags(core, "complexity") == {} + assert registered_tags(core, "adaptive") == {} + + +def test_upserting_an_adaptive_complexity_router_registers_its_adaptive_companion(make_core: MakeCore) -> None: + core: Final = make_core(model_list=list(PLAIN)) + + core.apply( + ChangeSet(upsert=[Deployment(**complexity_marker("smart", "m1", config={"tiers": TIERS, "adaptive": True}))]) # pyright: ignore[reportArgumentType] + ) + + assert registered_tags(core, "adaptive") == {"smart": [()]} + + +def test_an_upsert_that_turns_adaptive_off_releases_the_adaptive_companion(make_core: MakeCore) -> None: + core: Final = make_core( + model_list=[*PLAIN, complexity_marker("smart", "m1", config={"tiers": TIERS, "adaptive": True})] + ) + + core.apply(ChangeSet(upsert=[Deployment(**complexity_marker("smart", "m1"))])) # pyright: ignore[reportArgumentType] + + assert registered_tags(core, "complexity") == {"smart": [()]} + assert registered_tags(core, "adaptive") == {} + + +def test_editing_an_adaptive_router_rebuilds_it(make_core: MakeCore) -> None: + core: Final = make_core(model_list=_adaptive_list()) + + core.apply(ChangeSet(upsert=[Deployment(**adaptive_marker("ar", "ar1", ["y"]))])) # pyright: ignore[reportArgumentType] + + registry: Final = core.registry().plugins("adaptive") + assert registered_tags(core, "adaptive") == {"ar": [()]} + assert registry["ar"][0].strategy.model_to_cost == {"y": 0.25} diff --git a/tests/unit/test_router/core/test_gaps_l14000.py b/tests/unit/test_router/core/test_gaps_l14000.py new file mode 100644 index 00000000000..51c47bbf802 --- /dev/null +++ b/tests/unit/test_router/core/test_gaps_l14000.py @@ -0,0 +1,299 @@ +""" +Gaps in the characterization of ``router_python_core.py`` lines 14000-17669 that the views answer: +``settings().allowed_fails_for`` under an allowed_fails_policy, ``settings().as_dict()`` for the +latency strategy and after a ``routing_strategy_args`` update, the ``registry()`` deployment +sources ``routable`` / ``alias`` / ``routing_groups`` / ``routed_group`` / ``for_request`` and team +aliases of ``all``, cached clients of a pod-discovery deployment, and PTU-share ceilings. +""" + +from typing import Final + +import pytest + +import litellm +from litellm.litellm_core_utils.ptu_shares import PTUTeamCeiling +from litellm.proxy._types import UserAPIKeyAuth +from litellm.router_api import ChangeSet +from litellm.types.router import AllowedFailsPolicy +from tests.unit.test_router.core.scripted_host import ( + auth_error, + bad_request, + content_policy, + context_window, + deployment, + rate_limit, + server_error, + timeout, +) +from tests.unit.test_router.core.selection_helpers import MakeCore, picked_ids + +# --- settings().allowed_fails_for ---------------------------------------------------------------- + + +def _bad_gateway() -> litellm.BadGatewayError: + return litellm.BadGatewayError(message="bad gateway", llm_provider="openai", model="gpt-4o") + + +def _connection_error() -> litellm.APIConnectionError: + return litellm.APIConnectionError(message="no route", llm_provider="openai", model="gpt-4o") + + +FULL_POLICY: Final = AllowedFailsPolicy( + ContentPolicyViolationErrorAllowedFails=1, + BadRequestErrorAllowedFails=2, + BadGatewayErrorAllowedFails=3, + RateLimitErrorAllowedFails=4, +) + + +@pytest.mark.parametrize( + ("error", "expected"), + [ + (content_policy(), 1), + (bad_request(), 2), + # A context-window error is a BadRequestError, so it takes the bad-request allowance. + (context_window(), 2), + (_bad_gateway(), 3), + (rate_limit(), 4), + # Types the policy names no allowance for answer None, never the router's allowed_fails. + (server_error(), None), + (timeout(), None), + (auth_error(), None), + (_connection_error(), None), + ], + ids=[ + "content-policy", + "bad-request", + "context-window", + "bad-gateway", + "rate-limit", + "server-error", + "timeout", + "auth", + "connection", + ], +) +def test_allowed_fails_for_reads_the_policy_field_of_the_error_type( + make_core: MakeCore, error: Exception, expected: int | None +) -> None: + core: Final = make_core(model_list=[deployment("g", "a")], allowed_fails_policy=FULL_POLICY) + assert core.settings().allowed_fails_for(error) == expected + + +def test_allowed_fails_for_a_content_policy_error_falls_back_to_the_bad_request_allowance(make_core: MakeCore) -> None: + core: Final = make_core( + model_list=[deployment("g", "a")], allowed_fails_policy=AllowedFailsPolicy(BadRequestErrorAllowedFails=2) + ) + assert core.settings().allowed_fails_for(content_policy()) == 2 + + +def test_allowed_fails_for_without_a_policy_is_none(make_core: MakeCore) -> None: + core: Final = make_core(model_list=[deployment("g", "a")], allowed_fails=7) + assert core.settings().allowed_fails_for(rate_limit()) is None + + +# --- settings().as_dict() and routing_strategy_args ---------------------------------------------- + +LATENCY_ARGS_DEFAULTS: Final = {"lowest_latency_buffer": 0, "max_latency_list_size": 10, "ttft_percentile": None} + + +def test_latency_routing_reports_the_strategys_full_routing_args(make_core: MakeCore) -> None: + core: Final = make_core( + model_list=[deployment("g", "a")], routing_strategy="latency-based-routing", routing_strategy_args={"ttl": 30} + ) + assert core.settings().as_dict()["routing_strategy_args"] == {"ttl": 30.0, **LATENCY_ARGS_DEFAULTS} + + +def test_updating_routing_strategy_args_reaches_the_latency_strategy(make_core: MakeCore) -> None: + core: Final = make_core( + model_list=[deployment("g", "a")], routing_strategy="latency-based-routing", routing_strategy_args={"ttl": 30} + ) + core.apply(ChangeSet(settings={"routing_strategy_args": {"ttl": 99}})) + assert core.settings().as_dict()["routing_strategy_args"] == {"ttl": 99.0, **LATENCY_ARGS_DEFAULTS} + + +def test_routing_strategy_args_under_simple_shuffle_are_reported_as_given(make_core: MakeCore) -> None: + core: Final = make_core(model_list=[deployment("g", "a")]) + assert core.settings().as_dict()["routing_strategy_args"] == {} + core.apply(ChangeSet(settings={"routing_strategy_args": {"ttl": 99}})) + assert core.settings().as_dict()["routing_strategy_args"] == {"ttl": 99} + + +def test_switching_to_latency_routing_takes_the_args_sent_with_it(make_core: MakeCore) -> None: + core: Final = make_core(model_list=[deployment("g", "a")]) + core.apply(ChangeSet(settings={"routing_strategy": "latency-based-routing", "routing_strategy_args": {"ttl": 7}})) + settings: Final = core.settings().as_dict() + assert settings["routing_strategy"] == "latency-based-routing" + assert settings["routing_strategy_args"] == {"ttl": 7.0, **LATENCY_ARGS_DEFAULTS} + + +# --- registry() deployment sources --------------------------------------------------------------- + +TEAMS_AND_GROUPS: Final = [ + deployment("internal-x", "tx", model_info={"team_id": "team-x", "team_public_model_name": "pub"}), + deployment("internal-y", "ty", model_info={"team_id": "team-y", "team_public_model_name": "pub"}), + deployment("g", "a"), + deployment("g", "b", model_info={"access_groups": ["ag"]}), +] + +ROUTING_GROUP: Final = {"group_name": "grp", "models": ["g"], "routing_strategy": "simple-shuffle"} + + +def _ids(deployments: object) -> list[str]: # mutable-ok: a fresh list + return [deployment_id for _, deployment_id in _rows(deployments)] + + +def _rows(deployments: object) -> list[tuple[str, str]]: # mutable-ok: a fresh list + assert isinstance(deployments, (list, tuple)) + return [(str(row["model_name"]), str(row["model_info"]["id"])) for row in deployments] + + +def test_all_deployments_of_a_team_public_name_are_renamed_to_the_alias(make_core: MakeCore) -> None: + core: Final = make_core(model_list=list(TEAMS_AND_GROUPS)) + rows: Final = core.registry().deployments("pub", source="all", team_id="team-x", model_alias="nick") + assert _rows(rows) == [("nick", "tx")] + + +def test_an_alias_to_a_routing_group_lists_the_members_under_the_alias_without_access_groups( + make_core: MakeCore, +) -> None: + core: Final = make_core( + model_list=list(TEAMS_AND_GROUPS), routing_groups=[dict(ROUTING_GROUP)], model_group_alias={"al": "grp"} + ) + rows: Final = core.registry().deployments("al") + assert _rows(rows) == [("al", "a"), ("al", "b")] + assert isinstance(rows, list) + assert all("access_groups" not in row["model_info"] for row in rows) + + +def test_the_alias_source_lists_every_alias_row(make_core: MakeCore) -> None: + core: Final = make_core( + model_list=list(TEAMS_AND_GROUPS), + routing_groups=[dict(ROUTING_GROUP)], + model_group_alias={"al": "grp", "al2": "g"}, + ) + rows: Final = core.registry().deployments(source="alias") + assert _rows(rows) == [("al", "a"), ("al", "b"), ("al2", "a"), ("al2", "b")] + # A plain alias row keeps the member's access groups; only routing-group rows drop them. + assert isinstance(rows, list) + assert rows[3]["model_info"].get("access_groups") == ["ag"] + + +def test_the_routing_groups_source_lists_members_under_the_group_without_access_groups(make_core: MakeCore) -> None: + core: Final = make_core(model_list=list(TEAMS_AND_GROUPS), routing_groups=[dict(ROUTING_GROUP)]) + rows: Final = core.registry().deployments(source="routing_groups") + assert _rows(rows) == [("grp", "a"), ("grp", "b")] + assert isinstance(rows, tuple) + assert all("access_groups" not in row["model_info"] for row in rows) + + +def test_the_routed_group_of_a_team_public_name_is_the_teams_own_deployments(make_core: MakeCore) -> None: + core: Final = make_core(model_list=list(TEAMS_AND_GROUPS)) + assert _rows(core.registry().deployments("pub", source="routed_group", team_id="team-y")) == [("internal-y", "ty")] + assert core.registry().deployments("pub", source="routed_group") == [] + + +def test_the_routable_source_lists_the_deployments_routing_would_consider(make_core: MakeCore) -> None: + core: Final = make_core(model_list=list(TEAMS_AND_GROUPS)) + registry: Final = core.registry() + assert _ids(registry.deployments("g", source="routable")) == ["a", "b"] + assert registry.deployments("nope", source="routable") == () + assert _ids(registry.deployments("a", source="routable")) == ["a"] + team_request: Final = {"metadata": {"user_api_key_team_id": "team-x"}} + assert _ids(registry.deployments("pub", source="routable", request=team_request)) == ["tx"] + + +def test_the_routable_source_answers_nothing_for_a_name_several_teams_share_for_an_admin(make_core: MakeCore) -> None: + core: Final = make_core(model_list=list(TEAMS_AND_GROUPS)) + admin: Final = {"metadata": {"user_api_key_auth": UserAPIKeyAuth(user_role="proxy_admin")}} # pyright: ignore[reportCallIssue] + assert core.registry().deployments("pub", source="routable", request=admin) == () + + +def test_the_for_request_source_finds_one_teams_public_name_for_a_proxy_admin(make_core: MakeCore) -> None: + core: Final = make_core(model_list=[TEAMS_AND_GROUPS[0], *TEAMS_AND_GROUPS[2:]]) + admin: Final = {"metadata": {"user_api_key_auth": UserAPIKeyAuth(user_role="proxy_admin")}} # pyright: ignore[reportCallIssue] + assert _ids(core.registry().deployments("pub", source="for_request", request=admin)) == ["tx"] + assert core.registry().deployments("pub", source="for_request", request={}) == [] + + +def test_the_for_request_source_refuses_a_name_several_teams_share_for_a_proxy_admin(make_core: MakeCore) -> None: + core: Final = make_core(model_list=list(TEAMS_AND_GROUPS)) + admin: Final = {"metadata": {"user_api_key_auth": UserAPIKeyAuth(user_role="proxy_admin")}} # pyright: ignore[reportCallIssue] + with pytest.raises(litellm.BadRequestError) as raised: + core.registry().deployments("pub", source="for_request", request=admin) + assert str(raised.value) == ( + "litellm.BadRequestError: Model name 'pub' matches deployments from multiple teams. " + "Specify the deployment ID directly to disambiguate." + ) + + +@pytest.mark.parametrize("client_type", ["async", None]) +def test_a_pod_discovery_deployment_has_no_cached_client(make_core: MakeCore, client_type: str | None) -> None: + core: Final = make_core(model_list=[deployment("g", "a")]) + pod_discovered: Final = deployment("g", "k", kubernetes_pod_discovery=True) + assert core.registry().client(pod_discovered, client_type=client_type) is None + + +# --- PTU shares ---------------------------------------------------------------------------------- + + +def _ptu_shared(model_name: str, deployment_id: str) -> dict[str, object]: # mutable-ok: a model_list entry + return deployment( + model_name, + deployment_id, + model=f"azure/{deployment_id}", + api_base="https://example.openai.azure.com", + model_info={ + "ptu_shares": {"TA": 50}, + "ptu_count": 50, + "cost_per_ptu_per_hour": 1.0, + "ptu_effective_from": "2026-01-01T00:00:00Z", + "base_model": "gpt-4o", + }, + ) + + +PTU_MODELS: Final = [_ptu_shared("gpt", "shared"), _ptu_shared("solo", "s2"), deployment("solo", "open")] + + +def test_a_share_holding_team_has_a_ptu_ceiling_on_its_group(make_core: MakeCore) -> None: + core: Final = make_core(model_list=list(PTU_MODELS)) + expected: Final = PTUTeamCeiling( + model_group="gpt", tpm_limit=125000, output_to_input_ratio=4.0, cached_input_ratio=0.0 + ) + assert core.health().team_ptu_ceiling("TA", "gpt") == expected + assert core.health().team_ptu_ceiling("TA", "solo") == PTUTeamCeiling( + model_group="solo", tpm_limit=125000, output_to_input_ratio=4.0, cached_input_ratio=0.0 + ) + assert core.health().team_ptu_ceiling("TB", "gpt") is None + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("request_kwargs", "group", "expected"), + [ + ({"metadata": {"user_api_key_team_id": "TA"}}, "gpt", "shared"), + ({"metadata": {"user_api_key_team_id": "TA"}}, "solo", "s2"), + ({"metadata": {"user_api_key_team_id": "TB"}}, "solo", "open"), + ({}, "solo", "open"), + ], + ids=["holder-shared", "holder-prefers-its-share", "other-team-open", "no-team-open"], +) +async def test_ptu_shares_route_holders_to_their_share_and_others_to_open_deployments( + make_core: MakeCore, request_kwargs: dict[str, object], group: str, expected: str +) -> None: + core: Final = make_core(model_list=list(PTU_MODELS)) + assert await picked_ids(core, group, request_kwargs=request_kwargs) == {expected} + + +@pytest.mark.asyncio +@pytest.mark.parametrize("request_kwargs", [{"metadata": {"user_api_key_team_id": "TB"}}, {}], ids=["other", "none"]) +async def test_a_group_only_share_holders_may_use_refuses_everyone_else( + make_core: MakeCore, request_kwargs: dict[str, object] +) -> None: + core: Final = make_core(model_list=list(PTU_MODELS)) + with pytest.raises(litellm.BadRequestError) as raised: + await core.select("gpt", dict(request_kwargs)) + assert str(raised.value) == ( + "litellm.BadRequestError: Deployment gpt is reserved for the teams holding a PTU share of it" + ) diff --git a/tests/unit/test_router/core/test_gaps_l14000_selection.py b/tests/unit/test_router/core/test_gaps_l14000_selection.py new file mode 100644 index 00000000000..8cff6bcdbe5 --- /dev/null +++ b/tests/unit/test_router/core/test_gaps_l14000_selection.py @@ -0,0 +1,491 @@ +""" +Gaps in the characterization of ``router_python_core.py`` lines 14000-17669 on the selection path: +health-check routing (the cooldown bypass under an allowed_fails_policy, health-check probes, the +sync filter the pass-through pick uses), pass-through picks (paused deployments by id, priority +groups, cooling, other strategies, the async entry point), the failure dispatch to a request's +logging object, ``model_alias_map``, access groups and litellm model names, tier ``litellm_params`` +(dropped params, effort carriers, restored output ceilings), routing-decision redaction and the +Claude Code session binding's guards. +""" + +import asyncio +import time +from collections.abc import Mapping +from typing import Final + +import pytest + +import litellm +from litellm.proxy._types import UserAPIKeyAuth +from litellm.types.router import AllowedFailsPolicy, RouterRateLimitError +from tests.unit.test_router.core.scripted_host import ScriptedHost, auth_error, deployment, rate_limit +from tests.unit.test_router.core.selection_helpers import MESSAGES, MakeCore, ids_of, picked_ids +from tests.unit.test_router.core.test_selection_characterization_prerouting import ( + HARD_MESSAGES, + SESSION_ID, + TIERS, + _complexity_router, + _session_request, + _tiered, +) + +SPREAD_PICKS: Final = 40 + + +def _unhealthy(*ids: str) -> dict[str, dict[str, object]]: # mutable-ok: health states + return {i: {"is_healthy": False, "timestamp": time.time(), "reason": "probe failed"} for i in ids} + + +def _metadata(request: Mapping[str, object], bucket: str = "metadata") -> Mapping[str, object]: + metadata: Final = request.get(bucket) + assert isinstance(metadata, Mapping) + return metadata + + +# --- health-check routing ------------------------------------------------------------------------ + + +@pytest.mark.asyncio +async def test_health_check_routing_with_an_allowed_fails_policy_ignores_cooldowns_when_all_cool( + make_core: MakeCore, +) -> None: + core: Final = make_core( + model_list=[deployment("g", "a"), deployment("g", "b")], + enable_health_check_routing=True, + allowed_fails_policy=AllowedFailsPolicy(RateLimitErrorAllowedFails=0), + cooldown_time=60, + ) + assert [core.health().record_failure(i, rate_limit(), 429) for i in ("a", "b")] == [True, True] + assert await picked_ids(core, "g", times=SPREAD_PICKS) == {"a", "b"} + + +@pytest.mark.asyncio +async def test_health_check_routing_without_a_policy_keeps_honoring_cooldowns(make_core: MakeCore) -> None: + core: Final = make_core( + model_list=[deployment("g", "a"), deployment("g", "b")], enable_health_check_routing=True, cooldown_time=60 + ) + for deployment_id in ("a", "b"): + core.health().record_failure(deployment_id, rate_limit(), 429) + with pytest.raises(RouterRateLimitError) as raised: + await core.select("g") + assert str(raised.value) == ( + "No deployments available for selected model, Try again in 60 seconds. All deployments for selected " + "model are in cooldown. Passed model=g. pre-call-checks=False, cooldown_list=['a', 'b']" + ) + + +@pytest.mark.asyncio +async def test_a_health_check_probe_gets_no_deployments_when_all_are_unhealthy(make_core: MakeCore) -> None: + core: Final = make_core(model_list=[deployment("g", "a"), deployment("g", "b")], enable_health_check_routing=True) + core.health().health_checks.set_deployment_health_states(_unhealthy("a", "b")) # pyright: ignore[reportArgumentType] # health states + assert ids_of(await core.select("g", kind="healthy")) == ["a", "b"] + with pytest.raises(RouterRateLimitError) as raised: + await core.select("g", kind="healthy", health_check_probe=True) # pyright: ignore[reportCallIssue] + assert str(raised.value) == ( + "No deployments available for selected model, Try again in 5 seconds. Passed model=g. " + "pre-call-checks=False, cooldown_list=[]" + ) + + +@pytest.mark.asyncio +async def test_a_health_check_probe_still_drops_only_the_unhealthy_deployments(make_core: MakeCore) -> None: + core: Final = make_core(model_list=[deployment("g", "a"), deployment("g", "b")], enable_health_check_routing=True) + core.health().health_checks.set_deployment_health_states(_unhealthy("a")) # pyright: ignore[reportArgumentType] # health states + assert ids_of(await core.select("g", kind="healthy", health_check_probe=True)) == ["b"] # pyright: ignore[reportCallIssue] + + +PASS_THROUGH_PAIR: Final = [ + deployment("g", "pt", use_in_pass_through=True), + deployment("g", "pt2", use_in_pass_through=True), +] + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("arguments", "expected"), + [ + ({}, {"pt2"}), + ({"allowed_fails_policy": {"RateLimitErrorAllowedFails": 1}}, {"pt", "pt2"}), + ({"background_health_check_model_groups": ["other"]}, {"pt", "pt2"}), + ( + {"background_health_check_model_groups": ["g"], "allowed_fails_policy": {"RateLimitErrorAllowedFails": 1}}, + {"pt2"}, + ), + ], + ids=["filtered", "policy-disables", "group-not-listed", "listed-group-with-policy"], +) +async def test_pass_through_selection_applies_the_health_check_filter( + make_core: MakeCore, arguments: dict[str, object], expected: set[str] +) -> None: + core: Final = make_core(model_list=list(PASS_THROUGH_PAIR), enable_health_check_routing=True, **arguments) + core.health().health_checks.set_deployment_health_states(_unhealthy("pt")) # pyright: ignore[reportArgumentType] # health states + assert await picked_ids(core, "g", times=SPREAD_PICKS, kind="pass_through") == expected + + +@pytest.mark.asyncio +async def test_pass_through_selection_uses_every_deployment_when_all_are_unhealthy(make_core: MakeCore) -> None: + core: Final = make_core(model_list=list(PASS_THROUGH_PAIR), enable_health_check_routing=True) + core.health().health_checks.set_deployment_health_states(_unhealthy("pt", "pt2")) # pyright: ignore[reportArgumentType] # health states + assert await picked_ids(core, "g", times=SPREAD_PICKS, kind="pass_through") == {"pt", "pt2"} + + +# --- pass-through selection ---------------------------------------------------------------------- + +PAUSED_PASS_THROUGH: Final = [ + deployment("g", "pt", use_in_pass_through=True, model_info={"blocked": True}), + deployment("g", "pt2", use_in_pass_through=True), +] + +PAUSED_MESSAGE: Final = ( + # Oddity: the message names the deployment's litellm model, not the id that was asked for. + "litellm.ServiceUnavailableError: Model 'openai/pt' is currently paused and cannot accept requests." +) + + +@pytest.mark.asyncio +async def test_pass_through_selection_of_a_paused_deployment_by_id_is_unavailable(make_core: MakeCore) -> None: + core: Final = make_core(model_list=list(PAUSED_PASS_THROUGH)) + with pytest.raises(litellm.ServiceUnavailableError) as raised: + await core.select("pt", kind="pass_through") + assert str(raised.value) == PAUSED_MESSAGE + + +@pytest.mark.asyncio +async def test_async_pass_through_selection_of_a_paused_deployment_by_id_is_unavailable(make_core: MakeCore) -> None: + core: Final = make_core(model_list=list(PAUSED_PASS_THROUGH)) + with pytest.raises(litellm.ServiceUnavailableError) as raised: + await core.async_get_available_deployment_for_pass_through("pt", {}) # pyright: ignore[reportAttributeAccessIssue] + assert str(raised.value) == PAUSED_MESSAGE + + +@pytest.mark.asyncio +async def test_async_pass_through_selection_by_id_returns_a_pass_through_deployment(make_core: MakeCore) -> None: + core: Final = make_core(model_list=list(PAUSED_PASS_THROUGH)) + chosen: Final = await core.async_get_available_deployment_for_pass_through("pt2", {}) # pyright: ignore[reportAttributeAccessIssue] + assert chosen["model_info"]["id"] == "pt2" + + +@pytest.mark.asyncio +async def test_async_pass_through_selection_by_id_refuses_a_plain_deployment(make_core: MakeCore) -> None: + core: Final = make_core(model_list=[deployment("g", "pt", use_in_pass_through=True), deployment("g", "plain")]) + with pytest.raises(litellm.BadRequestError) as raised: + await core.async_get_available_deployment_for_pass_through("plain", {}) # pyright: ignore[reportAttributeAccessIssue] + assert str(raised.value) == ( + "litellm.BadRequestError: Deployment plain does not support pass-through endpoint (use_in_pass_through=False)" + ) + + +@pytest.mark.asyncio +async def test_async_pass_through_selection_picks_only_pass_through_deployments(make_core: MakeCore) -> None: + core: Final = make_core(model_list=[deployment("g", "pt", use_in_pass_through=True), deployment("g", "plain")]) + for _ in range(10): + chosen = await core.async_get_available_deployment_for_pass_through("g", {}) # pyright: ignore[reportAttributeAccessIssue] + assert chosen["model_info"]["id"] == "pt" + + +@pytest.mark.asyncio +async def test_pass_through_selection_with_a_selector_strategy_picks_through_it(make_core: MakeCore) -> None: + core: Final = make_core(model_list=list(PASS_THROUGH_PAIR), routing_strategy="least-busy") + assert await picked_ids(core, "g", kind="pass_through") == {"pt"} + chosen: Final = await core.async_get_available_deployment_for_pass_through("g", {}) # pyright: ignore[reportAttributeAccessIssue] + assert chosen["model_info"]["id"] == "pt" + + +PRIORITY_PASS_THROUGH: Final = [ + deployment("m1", "a", use_in_pass_through=True), + deployment("m2", "b", use_in_pass_through=True), +] +PRIORITY_GROUP: Final = { + "group_name": "pg", + "models": ["m1", "m2"], + "routing_strategy": "priority", + "model_priorities": {"m1": 1, "m2": 2}, +} + + +@pytest.mark.asyncio +async def test_pass_through_selection_of_a_priority_group_keeps_the_lowest_priority(make_core: MakeCore) -> None: + core: Final = make_core(model_list=list(PRIORITY_PASS_THROUGH), routing_groups=[dict(PRIORITY_GROUP)]) + assert await picked_ids(core, "pg", kind="pass_through") == {"a"} + assert await picked_ids(core, "pg", request_kwargs={"_target_order": 2}, kind="pass_through") == {"b"} + + +@pytest.mark.asyncio +async def test_pass_through_selection_of_a_priority_order_nothing_has_rate_limits(make_core: MakeCore) -> None: + core: Final = make_core(model_list=list(PRIORITY_PASS_THROUGH), routing_groups=[dict(PRIORITY_GROUP)]) + with pytest.raises(RouterRateLimitError) as raised: + await core.select("pg", {"_target_order": 3}, kind="pass_through") + assert str(raised.value) == ( + "No deployments available for selected model, Try again in 5 seconds. Passed model=pg. " + "pre-call-checks=False, cooldown_list=[]" + ) + + +@pytest.mark.asyncio +async def test_pass_through_selection_with_every_pass_through_deployment_cooling_rate_limits( + make_core: MakeCore, +) -> None: + core: Final = make_core( + model_list=[deployment("g", "pt", use_in_pass_through=True), deployment("g", "plain")], cooldown_time=60 + ) + assert core.health().record_failure("pt", auth_error(), 401) is True + with pytest.raises(RouterRateLimitError) as raised: + await core.select("g", kind="pass_through") + assert str(raised.value) == ( + "No deployments available for selected model, Try again in 60 seconds. All deployments for selected " + "model are in cooldown. Passed model=g. pre-call-checks=False, cooldown_list=['pt']" + ) + + +# --- failure dispatch to the request's logging object -------------------------------------------- + + +class _RecordingLoggingObject: + """Stands in for a request's ``litellm_logging_obj``: records the failure dispatches.""" + + def __init__(self) -> None: + self.dispatched: list[dict[str, object]] = [] # mutable-ok: the dispatches, in order + + async def dispatch_failure_handlers(self, **kwargs: object) -> None: # kwargs-ok: the dispatch arguments + self.dispatched.append(kwargs) + + +async def _let_tasks_run() -> None: + for _ in range(3): + await asyncio.sleep(0) + + +@pytest.mark.asyncio +async def test_a_refused_pick_is_dispatched_to_the_requests_logging_object(make_core: MakeCore) -> None: + core: Final = make_core(model_list=[deployment("g", "a")]) + logging_obj: Final = _RecordingLoggingObject() + with pytest.raises(litellm.BadRequestError) as raised: + await core.select("nope", {"litellm_logging_obj": logging_obj}) + await _let_tasks_run() + (dispatched,) = logging_obj.dispatched + assert dispatched["exception"] is raised.value + assert dispatched["prefer_async_handlers"] is True + assert "Traceback" in str(dispatched["traceback_exception"]) + + +@pytest.mark.asyncio +async def test_a_refused_async_pass_through_pick_is_dispatched_to_the_requests_logging_object( + make_core: MakeCore, +) -> None: + core: Final = make_core(model_list=[deployment("g", "a")]) + logging_obj: Final = _RecordingLoggingObject() + with pytest.raises(litellm.BadRequestError) as raised: + await core.async_get_available_deployment_for_pass_through("g", {"litellm_logging_obj": logging_obj}) # pyright: ignore[reportAttributeAccessIssue] + await _let_tasks_run() + assert [entry["exception"] for entry in logging_obj.dispatched] == [raised.value] + + +# --- model_alias_map, access groups and litellm model names -------------------------------------- + + +@pytest.mark.asyncio +async def test_the_process_model_alias_map_renames_the_model_in_a_no_deployments_error( + make_core: MakeCore, host: ScriptedHost +) -> None: + host.set_dynamic(model_alias_map={"g": "renamed"}) + core: Final = make_core(model_list=[deployment("g", "a")], cooldown_time=60) + assert (await core.select("g"))["model_info"]["id"] == "a" # pyright: ignore[reportIndexIssue] + core.health().record_failure("a", auth_error(), 401) + with pytest.raises(RouterRateLimitError) as raised: + await core.select("g") + assert "Passed model=renamed." in str(raised.value) + + +ACCESS_GROUPED: Final = [ + deployment("g", "in_a", model_info={"access_groups": ["grpA"]}), + deployment("g", "in_b", model_info={"access_groups": ["grpB"]}), + deployment("h", "hh"), +] + + +@pytest.mark.asyncio +async def test_an_access_group_only_caller_reaches_any_deployment_by_its_litellm_model(make_core: MakeCore) -> None: + # Oddity: the access-group narrowing looks up access groups by the requested name, and a + # litellm model string names no group, so a grpA-only caller reaches the grpB deployment. + core: Final = make_core(model_list=list(ACCESS_GROUPED), default_fallbacks=["h"]) + auth: Final = {"metadata": {"user_api_key_auth": UserAPIKeyAuth(models=["grpA"])}} # pyright: ignore[reportCallIssue] + assert (await core.select("openai/in_b", dict(auth)))["model_info"]["id"] == "in_b" # pyright: ignore[reportIndexIssue] + assert (await core.select("openai/in_a", dict(auth)))["model_info"]["id"] == "in_a" # pyright: ignore[reportIndexIssue] + assert (await core.select("unknown", dict(auth)))["model_info"]["id"] == "hh" # pyright: ignore[reportIndexIssue] + + +# --- tier litellm_params ------------------------------------------------------------------------- + + +def _pinned(**tier_params: object) -> dict[str, object]: # mutable-ok: tiers # kwargs-ok: tier litellm_params + return {**TIERS, "SIMPLE": {"model_name": "small", "litellm_params": tier_params}} + + +@pytest.mark.asyncio +async def test_a_tier_param_no_deployment_supports_is_dropped(make_core: MakeCore) -> None: + core: Final = make_core( + model_list=[ + deployment("small", "s", model="anthropic/claude-3-5-sonnet-20240620"), + deployment("big", "b"), + _complexity_router(tiers=_pinned(logit_bias={"1": 2}, temperature=0.2)), + ] + ) + request: Final[dict[str, object]] = {} + chosen: Final = await core.select("smart", request, messages=MESSAGES) # pyright: ignore[reportCallIssue] + assert chosen["model_info"]["id"] == "s" # pyright: ignore[reportIndexIssue] + assert request["temperature"] == 0.2 + assert "logit_bias" not in request + + +@pytest.mark.asyncio +async def test_a_tier_param_the_request_allowlists_is_kept(make_core: MakeCore) -> None: + core: Final = make_core( + model_list=[ + deployment("small", "s", model="anthropic/claude-3-5-sonnet-20240620"), + deployment("big", "b"), + _complexity_router(tiers=_pinned(logit_bias={"1": 2})), + ] + ) + request: Final[dict[str, object]] = {"allowed_openai_params": ["logit_bias"]} + await core.select("smart", request, messages=MESSAGES) # pyright: ignore[reportCallIssue] + assert request["logit_bias"] == {"1": 2} + + +@pytest.mark.asyncio +async def test_a_tier_reasoning_effort_removes_the_callers_other_effort_carriers(make_core: MakeCore) -> None: + core: Final = make_core(model_list=_tiered(tiers=_pinned(reasoning_effort="low"))) + request: Final[dict[str, object]] = { + "thinking": {"type": "enabled"}, + "output_config": {"effort": "high", "format": "json"}, + "reasoning": {"effort": "high"}, + "max_completion_tokens": 100, + } + await core.select("smart", request, messages=MESSAGES) # pyright: ignore[reportCallIssue] + assert request["reasoning_effort"] == "low" + assert "thinking" not in request + assert request["output_config"] == {"format": "json"} + assert "reasoning" not in request + assert request["max_completion_tokens"] == 100 + + +@pytest.mark.asyncio +async def test_a_tier_without_an_output_ceiling_leaves_the_callers_ceiling(make_core: MakeCore) -> None: + core: Final = make_core(model_list=_tiered(tiers=_pinned(temperature=0.1))) + request: Final[dict[str, object]] = {"max_completion_tokens": 100} + await core.select("smart", request, messages=MESSAGES) # pyright: ignore[reportCallIssue] + assert request["temperature"] == 0.1 + assert request["max_completion_tokens"] == 100 + assert "max_tokens" not in request + # The caller's ceiling is stamped on every routed pick, pinned or not. + assert _metadata(request)["_client_output_ceiling"] == {"max_completion_tokens": 100} + + +@pytest.mark.asyncio +async def test_a_later_pick_restores_the_callers_ceiling_a_tier_pin_replaced(make_core: MakeCore) -> None: + core: Final = make_core(model_list=_tiered(tiers=_pinned(max_tokens=7))) + request: Final[dict[str, object]] = {"max_completion_tokens": 100} + await core.select("smart", request, messages=MESSAGES) # pyright: ignore[reportCallIssue] + assert request["max_tokens"] == 7 + assert "max_completion_tokens" not in request + assert _metadata(request)["_client_output_ceiling"] == {"max_completion_tokens": 100} + + chosen: Final = await core.select("smart", request, messages=HARD_MESSAGES) # pyright: ignore[reportCallIssue] + assert chosen["model_info"]["id"] == "b" # pyright: ignore[reportIndexIssue] + assert request["max_completion_tokens"] == 100 + assert "max_tokens" not in request + + +@pytest.mark.asyncio +async def test_a_second_pick_leaves_the_first_picks_forwarded_marker_params_untracked(make_core: MakeCore) -> None: + # Oddity: forwarded marker params already on the request are not "newly forwarded" on the + # next attempt, so the tracking key is dropped while the params themselves stay behind. + core: Final = make_core(model_list=_tiered(temperature=0.3)) + request: Final[dict[str, object]] = {} + await core.select("smart", request, messages=MESSAGES) # pyright: ignore[reportCallIssue] + forwarded: Final = request["_alias_marker_forwarded_params"] + assert isinstance(forwarded, tuple) + assert "temperature" in forwarded + assert request["temperature"] == 0.3 + + await core.select("smart", request, messages=MESSAGES) # pyright: ignore[reportCallIssue] + assert "_alias_marker_forwarded_params" not in request + assert request["temperature"] == 0.3 + + +# --- routing decision redaction ------------------------------------------------------------------ + + +@pytest.mark.asyncio +async def test_the_routing_decision_keeps_its_prompt_signals_by_default(make_core: MakeCore) -> None: + core: Final = make_core(model_list=_tiered()) + request: Final[dict[str, object]] = {} + await core.select("smart", request, messages=MESSAGES) # pyright: ignore[reportCallIssue] + decision: Final = _metadata(request)["routing_decision"] + assert isinstance(decision, Mapping) + assert "signals" in decision + + +@pytest.mark.asyncio +async def test_redacted_message_logging_drops_the_prompt_quoting_routing_decision_fields(make_core: MakeCore) -> None: + core: Final = make_core(model_list=_tiered()) + request: Final[dict[str, object]] = {"standard_callback_dynamic_params": {"turn_off_message_logging": True}} + await core.select("smart", request, messages=MESSAGES) # pyright: ignore[reportCallIssue] + decision: Final = _metadata(request)["routing_decision"] + assert isinstance(decision, Mapping) + assert "signals" not in decision + assert (decision["router_model_name"], decision["tier"], decision["routed_model"]) == ("smart", "SIMPLE", "small") + + +@pytest.mark.asyncio +async def test_routing_metadata_goes_to_litellm_metadata_when_the_request_has_it(make_core: MakeCore) -> None: + core: Final = make_core(model_list=_tiered()) + request: Final[dict[str, object]] = {"litellm_metadata": {}} + await core.select("smart", request, messages=MESSAGES) # pyright: ignore[reportCallIssue] + assert "routing_decision" in _metadata(request, "litellm_metadata") + assert "metadata" not in request + + +# --- the Claude Code session binding's guards ---------------------------------------------------- + + +async def _subagent_pick(core: object) -> str: + chosen: Final = await core.select("big", _session_request(agent_id="agent-1"), messages=MESSAGES) # pyright: ignore[reportAttributeAccessIssue] + return str(chosen["model_info"]["id"]) + + +@pytest.mark.asyncio +@pytest.mark.parametrize(("depth", "expected"), [(1, "b"), (0, "s")], ids=["fallback-attempt", "first-attempt"]) +async def test_only_a_first_attempt_binds_its_session(make_core: MakeCore, depth: int, expected: str) -> None: + core: Final = make_core(model_list=_tiered()) + first: Final = _session_request() + first["fallback_depth"] = depth + await core.select("smart", first, messages=MESSAGES) # pyright: ignore[reportCallIssue] + assert await _subagent_pick(core) == expected + + +@pytest.mark.asyncio +async def test_a_request_without_a_caller_hash_never_binds(make_core: MakeCore) -> None: + core: Final = make_core(model_list=_tiered()) + first: Final = _session_request() + first["metadata"] = {} + await core.select("smart", first, messages=MESSAGES) # pyright: ignore[reportCallIssue] + assert await _subagent_pick(core) == "b" + + +@pytest.mark.asyncio +async def test_the_caller_hash_is_read_from_litellm_metadata_when_present(make_core: MakeCore) -> None: + core: Final = make_core(model_list=_tiered()) + first: Final = _session_request() + first["metadata"] = {} + first["litellm_metadata"] = {"user_api_key_hash": "caller-hash"} + await core.select("smart", first, messages=MESSAGES) # pyright: ignore[reportCallIssue] + assert await _subagent_pick(core) == "s" + + +@pytest.mark.asyncio +async def test_the_session_header_name_is_matched_case_insensitively(make_core: MakeCore) -> None: + core: Final = make_core(model_list=_tiered()) + first: Final = _session_request() + first["proxy_server_request"] = {"headers": {"X-Claude-Code-Session-Id": SESSION_ID}} + await core.select("smart", first, messages=MESSAGES) # pyright: ignore[reportCallIssue] + assert await _subagent_pick(core) == "s" diff --git a/tests/unit/test_router/core/test_gaps_l2800.py b/tests/unit/test_router/core/test_gaps_l2800.py new file mode 100644 index 00000000000..78e237f0696 --- /dev/null +++ b/tests/unit/test_router/core/test_gaps_l2800.py @@ -0,0 +1,431 @@ +""" +Gap tests for ``router_python_core.py`` lines 2800-5600: optional pre-call check re-registration, +silent experiments, per-deployment ``num_retries`` stamped on errors, and the request kwargs +``_update_kwargs_with_deployment`` builds (Kubernetes pod routing, baggage stamps, stream usage +injection, malformed request tools, per-user GitHub Copilot credentials). + +Each test asserts what the Python core does today; oddities are named in the test's docstring. +""" + +import asyncio +import threading +import time +from collections.abc import Callable, Mapping +from typing import Final + +import pytest + +import litellm +from litellm.constants import BAGGAGE_MAPPED_TAGS_METADATA_KEY, KUBERNETES_POD_ROUTING_KEY +from litellm.router_api import ChangeSet +from litellm.router_core import RouterCore +from litellm.router_strategy.budget_limiter import RouterBudgetLimiting +from litellm.types.litellm_params import STREAM_USAGE_INJECTED_FIELD +from litellm.types.router import BaggageMappedTagsStamp +from tests.unit.test_router.core.construction_precall_helpers import CallbackHost, host_callback_names +from tests.unit.test_router.core.scripted_host import ScriptedHost, deployment, rate_limit + +MakeCore = Callable[..., RouterCore] +MESSAGES: Final = [{"role": "user", "content": "hi"}] + + +@pytest.fixture +def host() -> ScriptedHost: + return CallbackHost() + + +def _metadata(kwargs: Mapping[str, object], name: str = "metadata") -> Mapping[str, object]: + bucket: Final = kwargs[name] + assert isinstance(bucket, Mapping) + return bucket + + +async def _settle() -> None: + for _ in range(30): + await asyncio.sleep(0) + + +# --- optional pre-call checks added again at runtime (add_optional_pre_call_checks) ------------ + + +def test_runtime_affinity_checks_merge_into_the_existing_affinity_callback( + make_core: MakeCore, host: ScriptedHost +) -> None: + core: Final = make_core(model_list=[deployment("g", "a")], optional_pre_call_checks=["deployment_affinity"]) + (existing,) = core.optional_callbacks # pyright: ignore[reportAttributeAccessIssue] + core.apply(ChangeSet(settings={"optional_pre_call_checks": ["deployment_affinity", "session_affinity"]})) + (merged,) = core.optional_callbacks # pyright: ignore[reportAttributeAccessIssue] + assert merged is existing + assert ( + merged.enable_user_key_affinity, + merged.enable_responses_api_affinity, + merged.enable_session_id_affinity, + ) == ( + True, + False, + True, + ) + assert host_callback_names(host)["callbacks"] == ["DeploymentAffinityCheck"] + + +def test_runtime_affinity_flags_are_only_ever_turned_on(make_core: MakeCore) -> None: + """Oddity: a runtime list without ``deployment_affinity`` leaves the flag it set on, since the + existing callback's flags are OR-ed with the new list's.""" + core: Final = make_core(model_list=[deployment("g", "a")], optional_pre_call_checks=["deployment_affinity"]) + core.apply(ChangeSet(settings={"optional_pre_call_checks": ["responses_api_deployment_check"]})) + (merged,) = core.optional_callbacks # pyright: ignore[reportAttributeAccessIssue] + assert ( + merged.enable_user_key_affinity, + merged.enable_responses_api_affinity, + merged.enable_session_id_affinity, + ) == ( + True, + True, + False, + ) + + +def test_runtime_affinity_merge_takes_the_model_group_affinity_config(make_core: MakeCore) -> None: + core: Final = make_core( + model_list=[deployment("g", "a")], + optional_pre_call_checks=["deployment_affinity"], + model_group_affinity_config={"g": ["session_affinity"]}, + ) + core.apply(ChangeSet(settings={"optional_pre_call_checks": ["session_affinity"]})) + (merged,) = [cb for cb in core.optional_callbacks if type(cb).__name__ == "DeploymentAffinityCheck"] # pyright: ignore[reportAttributeAccessIssue] + assert merged.model_group_affinity_config == {"g": ["session_affinity"]} + assert merged.enable_session_id_affinity is True + + +def test_runtime_encrypted_content_affinity_reuses_the_existing_callback( + make_core: MakeCore, host: ScriptedHost +) -> None: + core: Final = make_core(model_list=[deployment("g", "a")], optional_pre_call_checks=["encrypted_content_affinity"]) + (existing,) = core.optional_callbacks # pyright: ignore[reportAttributeAccessIssue] + core.apply(ChangeSet(settings={"optional_pre_call_checks": ["encrypted_content_affinity"]})) + (kept,) = core.optional_callbacks # pyright: ignore[reportAttributeAccessIssue] + assert kept is existing + assert kept.enable_global_affinity is True + assert host_callback_names(host)["callbacks"] == ["EncryptedContentAffinityCheck"] + + +def test_runtime_encrypted_content_affinity_turns_global_affinity_on_for_a_group_scoped_callback( + make_core: MakeCore, host: ScriptedHost +) -> None: + core: Final = make_core( + model_list=[deployment("g", "a")], model_group_affinity_config={"g": ["encrypted_content_affinity"]} + ) + (existing,) = [cb for cb in core.optional_callbacks if type(cb).__name__ == "EncryptedContentAffinityCheck"] # pyright: ignore[reportAttributeAccessIssue] + assert existing.enable_global_affinity is False + core.add_optional_pre_call_checks(["encrypted_content_affinity"]) # pyright: ignore[reportAttributeAccessIssue] + assert existing.enable_global_affinity is True + assert host_callback_names(host)["callbacks"] == ["EncryptedContentAffinityCheck", "DeploymentAffinityCheck"] + + +async def test_router_budget_limiting_is_not_added_twice_beside_a_provider_budget_config(make_core: MakeCore) -> None: + core: Final = make_core( + model_list=[deployment("g", "a")], + provider_budget_config={"openai": {"budget_limit": 1, "time_period": "1d"}}, + optional_pre_call_checks=["router_budget_limiting"], + ) + limiters: Final = [cb for cb in core.optional_callbacks if isinstance(cb, RouterBudgetLimiting)] # pyright: ignore[reportAttributeAccessIssue] + assert len(limiters) == 1 + assert core.registry().plugins("budget_limiter") is limiters[0] + core.retire() + + +# --- silent experiments ----------------------------------------------------------------------- + + +def _wait_for_calls(host: ScriptedHost, count: int) -> None: + deadline: Final = time.monotonic() + 5 + while len(host.calls) < count and time.monotonic() < deadline: + time.sleep(0.01) + + +def test_sync_completion_mirrors_to_the_silent_model_on_a_background_thread( + make_core: MakeCore, host: ScriptedHost +) -> None: + core: Final = make_core(model_list=[deployment("g", "d1", silent_model="shadow"), deployment("shadow", "s1")]) + core.completion(model="g", messages=MESSAGES, metadata={"trace": "t"}) # pyright: ignore[reportAttributeAccessIssue] + _wait_for_calls(host, 2) + # the thread starts before the primary call, so the two calls can land in either order + assert sorted((call.op, call.deployment_id) for call in host.calls) == [("acompletion", "s1"), ("completion", "d1")] + primary: Final = next(call for call in host.calls if call.deployment_id == "d1") + shadow: Final = _metadata(next(call for call in host.calls if call.deployment_id == "s1").kwargs) + assert (shadow["is_silent_experiment"], shadow["model_group"], shadow["trace"]) == (True, "shadow", "t") + assert "is_silent_experiment" not in _metadata(primary.kwargs) + + +async def test_a_request_already_marked_silent_does_not_start_another_experiment( + make_core: MakeCore, host: ScriptedHost +) -> None: + core: Final = make_core(model_list=[deployment("g", "d1", silent_model="shadow"), deployment("shadow", "s1")]) + await core.acompletion(model="g", messages=MESSAGES, metadata={"is_silent_experiment": True}) # pyright: ignore[reportAttributeAccessIssue] + await _settle() + assert host.deployments_called() == ["d1"] + + +def test_a_sync_request_already_marked_silent_does_not_start_another_experiment( + make_core: MakeCore, host: ScriptedHost +) -> None: + core: Final = make_core(model_list=[deployment("g", "d1", silent_model="shadow"), deployment("shadow", "s1")]) + core.completion(model="g", messages=MESSAGES, metadata={"is_silent_experiment": True}) # pyright: ignore[reportAttributeAccessIssue] + time.sleep(0.2) + assert host.deployments_called() == ["d1"] + + +async def test_silent_experiment_marks_its_own_metadata_when_the_request_metadata_cannot_be_deep_copied( + make_core: MakeCore, host: ScriptedHost +) -> None: + lock: Final = threading.Lock() + core: Final = make_core(model_list=[deployment("g", "d1", silent_model="shadow"), deployment("shadow", "s1")]) + await core.acompletion(model="g", messages=MESSAGES, metadata={"lock": lock}) # pyright: ignore[reportAttributeAccessIssue] + await _settle() + primary: Final = next(call for call in host.calls if call.deployment_id == "d1") + shadow: Final = next(call for call in host.calls if call.deployment_id == "s1") + assert "is_silent_experiment" not in _metadata(primary.kwargs) + assert _metadata(primary.kwargs)["model_group"] == "g" + assert (_metadata(shadow.kwargs)["is_silent_experiment"], _metadata(shadow.kwargs)["lock"]) == (True, lock) + + +async def test_a_streaming_request_mirrors_a_streaming_silent_experiment( + make_core: MakeCore, host: ScriptedHost +) -> None: + from tests.unit.test_router.core.streaming_helpers import chat_stream, collect_texts + + host.answer("openai/d1", chat_stream("A", model="d1")) + host.answer("openai/s1", chat_stream("S", model="s1")) + core: Final = make_core(model_list=[deployment("g", "d1", silent_model="shadow"), deployment("shadow", "s1")]) + stream: Final = await core.acompletion(model="g", messages=MESSAGES, stream=True) # pyright: ignore[reportAttributeAccessIssue] + assert await collect_texts(stream) == "A" + await _settle() + shadow: Final = next(call for call in host.calls if call.deployment_id == "s1") + assert shadow.kwargs["stream"] is True + + +# --- per-deployment num_retries stamped on the error ------------------------------------------ + + +async def test_a_num_retries_already_on_the_error_wins_over_the_deployment_num_retries( + make_core: MakeCore, host: ScriptedHost +) -> None: + error: Final = rate_limit() + error.num_retries = 0 # pyright: ignore[reportAttributeAccessIssue] + host.answer("openai/d1", error) + core: Final = make_core(model_list=[deployment("g", "d1", num_retries=2)]) + with pytest.raises(litellm.RateLimitError): + await core.acompletion(model="g", messages=MESSAGES) # pyright: ignore[reportAttributeAccessIssue] + assert len(host.calls) == 1 + + +async def test_a_string_deployment_num_retries_is_used_as_an_int(make_core: MakeCore, host: ScriptedHost) -> None: + host.answer("openai/d1", rate_limit()) + core: Final = make_core(model_list=[deployment("g", "d1", num_retries="1")]) + with pytest.raises(litellm.RateLimitError) as raised: + await core.acompletion(model="g", messages=MESSAGES) # pyright: ignore[reportAttributeAccessIssue] + assert len(host.calls) == 2 + assert getattr(raised.value, "num_retries") == 1 + + +async def test_a_non_numeric_deployment_num_retries_is_ignored(make_core: MakeCore, host: ScriptedHost) -> None: + """The router's own default (2 retries) applies when the deployment value cannot be read as an int.""" + host.answer("openai/d1", rate_limit()) + core: Final = make_core(model_list=[deployment("g", "d1", num_retries="abc")]) + with pytest.raises(litellm.RateLimitError): + await core.acompletion(model="g", messages=MESSAGES) # pyright: ignore[reportAttributeAccessIssue] + assert len(host.calls) == 3 + + +# --- request kwargs built for the deployment -------------------------------------------------- + + +async def test_responses_replace_kubernetes_pod_routing_in_the_request_metadata( + make_core: MakeCore, host: ScriptedHost +) -> None: + """On the Responses path the routing goes to ``litellm_metadata``; a pod routing key the caller put + in ``metadata`` is overwritten with the serving deployment's (none here).""" + core: Final = make_core(model_list=[deployment("g", "d1")]) + await core.aresponses( # pyright: ignore[reportAttributeAccessIssue] + model="g", input="hi", metadata={KUBERNETES_POD_ROUTING_KEY: {"pod": "stale"}, "keep": 1} + ) + sent: Final = host.calls[0].kwargs + assert _metadata(sent) == {KUBERNETES_POD_ROUTING_KEY: None, "keep": 1} + assert _metadata(sent, "litellm_metadata")[KUBERNETES_POD_ROUTING_KEY] is None + + +async def test_a_baggage_stamp_records_the_serving_deployment_tags(make_core: MakeCore, host: ScriptedHost) -> None: + core: Final = make_core(model_list=[deployment("g", "d1", tags=["t1", "t2"])]) + stamp: Final = BaggageMappedTagsStamp(tags=("b1",)) + await core.acompletion( # pyright: ignore[reportAttributeAccessIssue] + model="g", messages=MESSAGES, metadata={BAGGAGE_MAPPED_TAGS_METADATA_KEY: stamp, "tags": ["b1"]} + ) + metadata: Final = _metadata(host.calls[0].kwargs) + assert metadata[BAGGAGE_MAPPED_TAGS_METADATA_KEY] == BaggageMappedTagsStamp( + tags=("b1",), serving_deployment_tags=("t1", "t2") + ) + assert metadata["tags"] == ["t1", "t2", "b1"] + + +@pytest.mark.parametrize( + ("request_extra", "injected"), + [ + ({"stream": True}, True), + ({"stream": True, "stream_options": {"include_usage": False}}, False), + ({}, False), + ], +) +async def test_a_deployment_stream_options_include_usage_marks_the_streaming_request_as_injected( + make_core: MakeCore, host: ScriptedHost, request_extra: dict[str, object], injected: bool +) -> None: + from tests.unit.test_router.core.streaming_helpers import chat_stream, collect_texts + + host.answer("openai/d1", chat_stream("A")) + core: Final = make_core(model_list=[deployment("g", "d1", stream_options={"include_usage": True})]) + response: Final = await core.acompletion(model="g", messages=MESSAGES, **request_extra) # pyright: ignore[reportAttributeAccessIssue] + if request_extra.get("stream"): + await collect_texts(response) + assert host.calls[0].kwargs.get(STREAM_USAGE_INJECTED_FIELD, False) is injected + + +async def test_malformed_request_tools_are_dropped_for_the_deployment_tools_when_drop_params_is_set( + make_core: MakeCore, host: ScriptedHost +) -> None: + tool: Final = {"type": "function", "function": {"name": "dep_tool", "parameters": {}}} + core: Final = make_core(model_list=[deployment("g", "d1", tools=[tool], drop_params=True)]) + await core.acompletion(model="g", messages=MESSAGES, tools="not-a-list", tool_choice="auto") # pyright: ignore[reportAttributeAccessIssue] + sent: Final = host.calls[0].kwargs + assert sent["tools"] == [tool] + assert "tool_choice" not in sent + + +async def test_a_per_user_github_copilot_deployment_rejects_another_credential_name( + make_core: MakeCore, host: ScriptedHost +) -> None: + core: Final = make_core( + model_list=[ + deployment( + "g", + "d1", + model="github_copilot/gpt-4o", + litellm_credential_name="copilot-cred", + github_copilot_auth_type="per_user_oauth", + ) + ], + num_retries=0, + ) + with pytest.raises(litellm.BadRequestError) as raised: + await core.acompletion(model="g", messages=MESSAGES, litellm_credential_name="someone-else") # pyright: ignore[reportAttributeAccessIssue] + assert "litellm_credential_name cannot be overridden on a deployment that uses per-user GitHub OAuth" in str( + raised.value + ) + assert host.calls == [] + + +# --- timeouts --------------------------------------------------------------------------------- + + +async def test_the_process_request_timeout_outranks_the_router_timeout_once_a_router_timeout_is_set( + make_core: MakeCore, host: ScriptedHost, monkeypatch: pytest.MonkeyPatch +) -> None: + """Oddity: this per-attempt timeout is read from ``litellm.request_timeout`` itself (not the host's + process settings), and only when the router has a ``timeout`` of its own.""" + monkeypatch.setattr(litellm, "request_timeout", 33.0) + core: Final = make_core(model_list=[deployment("g", "d1")], timeout=42) + await core.acompletion(model="g", messages=MESSAGES) # pyright: ignore[reportAttributeAccessIssue] + assert host.calls[0].kwargs["timeout"] == 33.0 + + +async def test_the_process_request_timeout_is_the_stream_timeout_without_a_stream_timeout( + make_core: MakeCore, host: ScriptedHost, monkeypatch: pytest.MonkeyPatch +) -> None: + from tests.unit.test_router.core.streaming_helpers import chat_stream, collect_texts + + monkeypatch.setattr(litellm, "request_timeout", 33.0) + host.answer("openai/d1", chat_stream("A")) + core: Final = make_core(model_list=[deployment("g", "d1")], timeout=42) + await collect_texts(await core.acompletion(model="g", messages=MESSAGES, stream=True)) # pyright: ignore[reportAttributeAccessIssue] + assert (host.calls[0].kwargs["stream_timeout"], host.calls[0].kwargs["timeout"]) == (33.0, 33.0) + + +async def test_default_litellm_params_timeout_applies_when_no_router_or_process_timeout_is_set( + make_core: MakeCore, host: ScriptedHost +) -> None: + host.set_dynamic(request_timeout=None) + core: Final = make_core(model_list=[deployment("g", "d1")], default_litellm_params={"timeout": 9}) + await core.acompletion(model="g", messages=MESSAGES) # pyright: ignore[reportAttributeAccessIssue] + assert host.calls[0].kwargs["timeout"] == 9 + + +async def test_default_litellm_params_stream_timeout_is_the_last_stream_timeout_fallback( + make_core: MakeCore, host: ScriptedHost +) -> None: + from tests.unit.test_router.core.streaming_helpers import chat_stream, collect_texts + + host.answer("openai/d1", chat_stream("A")) + core: Final = make_core(model_list=[deployment("g", "d1")], default_litellm_params={"stream_timeout": 4}) + await collect_texts(await core.acompletion(model="g", messages=MESSAGES, stream=True)) # pyright: ignore[reportAttributeAccessIssue] + assert host.calls[0].kwargs["stream_timeout"] == 4 + + +# --- sync completion runs the routing pre-call checks for a model group only ------------------- + + +def test_sync_completion_runs_routing_pre_call_checks_for_a_group_but_not_a_deployment_id( + make_core: MakeCore, host: ScriptedHost +) -> None: + from litellm.integrations.custom_logger import CustomLogger + + checked: Final[list[object]] = [] + + class Recorder(CustomLogger): + def pre_call_check(self, deployment: dict) -> dict | None: # pyright: ignore[reportIncompatibleMethodOverride] + checked.append(deployment["model_info"]["id"]) + return None + + host.lists["callbacks"].append(Recorder()) + core: Final = make_core(model_list=[deployment("g", "d1")]) + core.completion(model="g", messages=MESSAGES) # pyright: ignore[reportAttributeAccessIssue] + core.completion(model="d1", messages=MESSAGES) # pyright: ignore[reportAttributeAccessIssue] + assert checked == ["d1"] + assert host.deployments_called() == ["d1", "d1"] + + +def test_a_sync_routing_pre_call_check_error_stops_the_call(make_core: MakeCore, host: ScriptedHost) -> None: + from litellm.integrations.custom_logger import CustomLogger + + class OverLimit(CustomLogger): + def pre_call_check(self, deployment: dict) -> dict | None: # pyright: ignore[reportIncompatibleMethodOverride] + raise litellm.RateLimitError(message="over the rpm limit", llm_provider="openai", model="d1") + + host.lists["callbacks"].append(OverLimit()) + core: Final = make_core(model_list=[deployment("g", "d1")], num_retries=0) + with pytest.raises(litellm.RateLimitError, match="over the rpm limit"): + core.completion(model="g", messages=MESSAGES) # pyright: ignore[reportAttributeAccessIssue] + assert host.calls == [] + + +# --- per-model call counters ------------------------------------------------------------------ + + +async def test_acompletion_counts_total_success_and_failed_calls_per_deployment_model( + make_core: MakeCore, host: ScriptedHost +) -> None: + core: Final = make_core(model_list=[deployment("g", "d1")], num_retries=0) + await core.acompletion(model="g", messages=MESSAGES) # pyright: ignore[reportAttributeAccessIssue] + host.answer("openai/d1", rate_limit()) + with pytest.raises(litellm.RateLimitError): + await core.acompletion(model="g", messages=MESSAGES) # pyright: ignore[reportAttributeAccessIssue] + counts: Final = ( + core.total_calls["openai/d1"], # pyright: ignore[reportAttributeAccessIssue] + core.success_calls["openai/d1"], # pyright: ignore[reportAttributeAccessIssue] + core.fail_calls["openai/d1"], # pyright: ignore[reportAttributeAccessIssue] + ) + assert counts == (2, 1, 1) + + +def test_sync_completion_does_not_count_calls(make_core: MakeCore, host: ScriptedHost) -> None: + """Oddity: only the async path updates ``total_calls``/``success_calls``/``fail_calls``.""" + core: Final = make_core(model_list=[deployment("g", "d1")]) + core.completion(model="g", messages=MESSAGES) # pyright: ignore[reportAttributeAccessIssue] + assert core.total_calls["openai/d1"] == 0 # pyright: ignore[reportAttributeAccessIssue] diff --git a/tests/unit/test_router/core/test_gaps_l2800_streaming.py b/tests/unit/test_router/core/test_gaps_l2800_streaming.py new file mode 100644 index 00000000000..42c95c6f894 --- /dev/null +++ b/tests/unit/test_router/core/test_gaps_l2800_streaming.py @@ -0,0 +1,177 @@ +""" +Gap tests for the stream wrappers of ``router_python_core.py`` (lines 2800-5600): what a mid-stream +fallback yields when the fallback answers without a stream, the usage on fallback chunks, and the +Responses stream's handling of an error that is not a mid-stream fallback error. + +Each test asserts what the Python core does today; oddities are named in the test's docstring. +""" + +from collections.abc import AsyncIterator, Callable +from typing import Final + +import pytest + +import litellm +from litellm.router_core import RouterCore +from litellm.types.utils import ModelResponseStream, Usage +from tests.unit.test_router.core.scripted_host import ( + ScriptedHost, + bad_request, + completion_response, + deployment, + rate_limit, +) +from tests.unit.test_router.core.streaming_helpers import ( + MESSAGES, + Event, + chat_stream, + drain_into, + finish_chunk, + responses_stream, + text_chunk, +) + +MakeCore = Callable[..., RouterCore] +GROUPS_AB: Final = [deployment("a", "a"), deployment("b", "b")] + + +def build(make_core: MakeCore, **arguments: object) -> RouterCore: # kwargs-ok: the router's constructor arguments + return make_core(**{"model_list": GROUPS_AB, "fallbacks": [{"a": ["b"]}], "num_retries": 0, **arguments}) + + +def stream_with_usage(sync: bool = False) -> Callable[..., object]: + """A bare chunk iterator (not a ``CustomStreamWrapper``, which would strip the usage) of one text + chunk and a finish chunk that carries usage.""" + + def answer(**_request: object) -> object: # kwargs-ok: the host passes the request kwargs + chunks: Final = [ + text_chunk("B"), + finish_chunk(usage={"prompt_tokens": 3, "completion_tokens": 4, "total_tokens": 7}), + ] + + async def achunks() -> AsyncIterator[ModelResponseStream]: + for chunk in chunks: + yield chunk + + return iter(chunks) if sync else achunks() + + return answer + + +# --- a fallback that answers without a stream -------------------------------------------------- + + +async def test_a_mid_stream_fallback_that_answers_without_a_stream_yields_one_none_item( + make_core: MakeCore, host: ScriptedHost +) -> None: + """Oddity: the fallback's whole response is dropped and the stream yields ``None`` in its place.""" + host.answer("openai/a", chat_stream(error=rate_limit(), leading_role_chunk=True)) + host.answer("openai/b", completion_response("not streamed", model="b")) + core: Final = build(make_core) + stream: Final = await core.acompletion(model="a", messages=MESSAGES, stream=True) # pyright: ignore[reportAttributeAccessIssue] + seen: list[object] = [] + await drain_into(stream, seen) + assert seen[-1] is None + assert host.deployments_called() == ["a", "b"] + + +def test_a_sync_mid_stream_fallback_that_answers_without_a_stream_yields_the_response_fields( + make_core: MakeCore, host: ScriptedHost +) -> None: + """Oddity: a ``ModelResponse`` is iterable (pydantic yields ``(field, value)`` pairs), so the sync + wrapper streams the fallback response's fields instead of chunks.""" + host.answer("openai/a", chat_stream(error=rate_limit(), leading_role_chunk=True, sync=True)) + host.answer("openai/b", completion_response("not streamed", model="b")) + core: Final = build(make_core) + stream: Final = core.completion(model="a", messages=MESSAGES, stream=True) # pyright: ignore[reportAttributeAccessIssue] + items: Final = list(stream) # pyright: ignore[reportArgumentType] + fields: Final = [item for item in items if isinstance(item, tuple)] + assert ("model", "b") in fields + assert {name for name, _ in fields} >= {"id", "choices", "usage"} + assert host.deployments_called() == ["a", "b"] + + +# --- usage on fallback chunks ------------------------------------------------------------------ + + +async def test_a_mid_stream_fallback_chunk_dict_usage_becomes_an_all_zero_usage( + make_core: MakeCore, host: ScriptedHost +) -> None: + """Oddity: a fallback chunk's usage is rebuilt through ``combine_usage_objects``, which reads token + counts as attributes, so a usage given as a plain dict comes out as an all-zero ``Usage``.""" + host.answer("openai/a", chat_stream(error=rate_limit(), leading_role_chunk=True)) + host.answer("openai/b", stream_with_usage()) + core: Final = build(make_core) + stream: Final = await core.acompletion(model="a", messages=MESSAGES, stream=True) # pyright: ignore[reportAttributeAccessIssue] + seen: list[object] = [] + await drain_into(stream, seen) + usages: Final = [ + getattr(item, "usage") + for item in seen + if isinstance(item, ModelResponseStream) and getattr(item, "usage", None) is not None + ] + assert usages + assert isinstance(usages[-1], Usage) + assert (usages[-1].prompt_tokens, usages[-1].completion_tokens, usages[-1].total_tokens) == (0, 0, 0) + + +def test_a_sync_mid_stream_fallback_chunk_dict_usage_becomes_an_all_zero_usage( + make_core: MakeCore, host: ScriptedHost +) -> None: + """Oddity: as in the async wrapper, a plain-dict usage comes out as an all-zero ``Usage``.""" + host.answer("openai/a", chat_stream(error=rate_limit(), leading_role_chunk=True, sync=True)) + host.answer("openai/b", stream_with_usage(sync=True)) + core: Final = build(make_core) + stream: Final = core.completion(model="a", messages=MESSAGES, stream=True) # pyright: ignore[reportAttributeAccessIssue] + usages: Final = [ + getattr(item, "usage") + for item in stream # pyright: ignore[reportGeneralTypeIssues] + if isinstance(item, ModelResponseStream) and getattr(item, "usage", None) is not None + ] + assert usages + assert isinstance(usages[-1], Usage) + assert (usages[-1].prompt_tokens, usages[-1].completion_tokens, usages[-1].total_tokens) == (0, 0, 0) + + +# --- Responses streams ------------------------------------------------------------------------- + + +async def test_a_responses_error_that_is_not_a_mid_stream_error_flushes_held_events_then_raises( + make_core: MakeCore, host: ScriptedHost +) -> None: + host.answer( + "openai/a", + responses_stream(Event("response.created"), Event("response.in_progress"), error=bad_request(message="plain")), + ) + host.answer("openai/b", responses_stream(Event("response.created"), Event("response.completed"))) + core: Final = build(make_core) + stream: Final = await core.aresponses(model="a", input="hi", stream=True) # pyright: ignore[reportAttributeAccessIssue] + seen: list[object] = [] + with pytest.raises(litellm.BadRequestError, match="plain"): + await drain_into(stream, seen) + assert [getattr(event, "type") for event in seen] == ["response.created", "response.in_progress"] + assert host.deployments_called() == ["a"] + + +async def test_a_responses_stream_fallback_that_answers_without_a_stream_yields_that_response( + make_core: MakeCore, host: ScriptedHost +) -> None: + from litellm.exceptions import MidStreamFallbackError + + cut: Final = MidStreamFallbackError( + message="stream cut", + model="a", + llm_provider="anthropic", + original_exception=rate_limit(message="stream cut"), + generated_content="", + is_pre_first_chunk=True, + ) + plain: Final = completion_response("not streamed", model="b") + host.answer("openai/a", responses_stream(Event("response.created"), error=cut)) + host.answer("openai/b", plain) + core: Final = build(make_core) + stream: Final = await core.aresponses(model="a", input="hi", stream=True) # pyright: ignore[reportAttributeAccessIssue] + seen: list[object] = [] + await drain_into(stream, seen) + assert seen == [plain] + assert host.deployments_called() == ["a", "b"] diff --git a/tests/unit/test_router/core/test_gaps_l5600.py b/tests/unit/test_router/core/test_gaps_l5600.py new file mode 100644 index 00000000000..83ab8e943b9 --- /dev/null +++ b/tests/unit/test_router/core/test_gaps_l5600.py @@ -0,0 +1,479 @@ +""" +Gap characterization for ``router_python_core.py`` lines 5600-8400: batch completions, the scheduler, +prompt-management groups, the image/audio/rerank/text/adapter endpoints, generic calls and their +pass-through endpoint rewrite, guardrail selection, embeddings, files and batches. + +Each test pins a behavior the earlier characterization files leave unasserted; oddities are pinned as +Python does them today and called out in the test's docstring. +""" + +import datetime +import io +from collections.abc import Callable, Mapping +from types import SimpleNamespace +from typing import Final + +import pytest + +import litellm +from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging +from litellm.router_core import RouterCore +from litellm.types.utils import CredentialItem, LiteLLMBatch +from tests.unit.test_router.core.endpoints_helpers import EndpointHost, patch_direct_endpoints +from tests.unit.test_router.core.scripted_host import ( + ScriptedHost, + auth_error, + deployment, + rate_limit, + server_error, +) + +MESSAGES: Final = [{"role": "user", "content": "hi"}] +MakeCore = Callable[..., RouterCore] +TWO_DEPLOYMENTS: Final = [deployment("g", "d1"), deployment("g", "d2")] + + +@pytest.fixture +def host() -> EndpointHost: + return EndpointHost() + + +def _batch() -> LiteLLMBatch: + return LiteLLMBatch( + id="b1", + completion_window="24h", + created_at=1, + endpoint="/v1/chat/completions", + input_file_id="f", + object="batch", + status="completed", + ) + + +def _hidden(response: object) -> Mapping[str, object]: + return getattr(response, "_hidden_params") + + +def _headers(response: object) -> Mapping[str, object]: + found: Final = _hidden(response)["additional_headers"] + assert isinstance(found, Mapping) + return found + + +# The routed endpoints that answer through ``async_function_with_fallbacks`` with their own attempt helper. +ROUTED_ENDPOINTS: Final[dict[str, dict[str, object]]] = { + "aimage_generation": {"prompt": "p"}, + "atranscription": {"file": b"audio"}, + "aspeech": {"input": "x", "voice": "alloy"}, + "arerank": {"query": "q", "documents": ["a"]}, + "aadapter_completion": {"adapter_id": "ad", "messages": MESSAGES}, + "atext_completion": {"prompt": "p"}, + "aembedding": {"input": ["a"]}, + "acreate_batch": {"input_file_id": "f", "endpoint": "/v1/chat/completions", "completion_window": "24h"}, + "acancel_batch": {"batch_id": "b1"}, +} + + +# --- batch completions (5565-5790) ---------------------------------------------------------- + + +async def test_abatch_completion_with_message_lists_keeps_a_failed_request_in_its_own_row( + make_core: MakeCore, host: EndpointHost +): + host.answer("openai/d1", rate_limit("d1", "a failed")) + core: Final = make_core(model_list=[deployment("a", "d1"), deployment("b", "d2")], num_retries=0) + rows: Final = await core.abatch_completion(models=["a", "b"], messages=[MESSAGES, MESSAGES]) + assert [[type(result).__name__ for result in row] for row in rows] == [ + ["RateLimitError", "ModelResponse"], + ["RateLimitError", "ModelResponse"], + ] + + +async def test_abatch_completion_with_mixed_message_shapes_returns_none_without_calling( + make_core: MakeCore, host: EndpointHost +): + """Neither every element a dict nor every element a list: no branch matches and None comes back.""" + core: Final = make_core(model_list=[deployment("a", "d1")]) + assert await core.abatch_completion(models=["a"], messages=[MESSAGES, {"role": "user", "content": "x"}]) is None + assert host.calls == [] + + +async def test_fastest_response_forwards_the_stream_flag(make_core: MakeCore, host: EndpointHost): + core: Final = make_core(model_list=[deployment("a", "d1")]) + await core.abatch_completion_fastest_response(model="a", messages=MESSAGES, stream=True) + assert host.calls[0].kwargs["stream"] is True + + +# --- scheduler (5792-5912) ------------------------------------------------------------------ + + +async def test_schedule_acompletion_forwards_the_stream_flag(make_core: MakeCore, host: EndpointHost): + core: Final = make_core(model_list=[deployment("g", "d1")]) + await core.schedule_acompletion(model="g", messages=MESSAGES, priority=1, stream=True) + assert host.calls[0].kwargs["stream"] is True + + +async def test_a_prioritized_request_for_a_deployment_id_is_queued_and_served_by_it( + make_core: MakeCore, host: EndpointHost +): + core: Final = make_core(model_list=TWO_DEPLOYMENTS) + for _ in range(3): + response = await core.acompletion(model="d1", messages=MESSAGES, priority=2) + assert _headers(response)["x-litellm-request-prioritization-used"] is True + assert host.deployments_called() == ["d1", "d1", "d1"] + + +async def test_a_bad_priority_for_a_deployment_id_follows_that_deployments_drop_params( + make_core: MakeCore, host: EndpointHost +): + core: Final = make_core(model_list=[deployment("g", "d1", drop_params=True), deployment("g", "d2")]) + response: Final = await core.acompletion(model="d1", messages=MESSAGES, priority="high") + assert "x-litellm-request-prioritization-used" not in _headers(response) + assert "priority" not in host.calls[0].kwargs + with pytest.raises(litellm.BadRequestError, match="priority must be an integer, got str 'high'"): + await core.acompletion(model="d2", messages=MESSAGES, priority="high") + assert host.deployments_called() == ["d1"] + + +async def test_a_bad_priority_is_dropped_when_any_deployment_of_the_group_drops_params( + make_core: MakeCore, host: EndpointHost +): + core: Final = make_core(model_list=[deployment("g", "d1", drop_params=True), deployment("g", "d2")]) + response: Final = await core.acompletion(model="g", messages=MESSAGES, priority="high") + assert "x-litellm-request-prioritization-used" not in _headers(response) + assert len(host.calls) == 1 + + +# --- prompt-management groups (5893-6034) --------------------------------------------------- + + +class _PromptLogging(LiteLLMLogging): + """A logging object whose prompt template renders to ``rendered_model`` with a system message.""" + + rendered_model: str = "openai/rendered" + seen: Mapping[str, object] = {} + + def get_chat_completion_prompt(self, **kwargs: object): # kwargs-ok: Logging's own keyword signature + self.seen = kwargs + return self.rendered_model, [{"role": "system", "content": "rendered"}], {"temperature": 0.7} + + +def _prompt_logging(rendered_model: str = "openai/rendered") -> _PromptLogging: + logging_obj: Final = _PromptLogging( + model="pm", + messages=MESSAGES, + stream=False, + call_type="acompletion", + start_time=datetime.datetime(2026, 1, 1, tzinfo=datetime.timezone.utc), + litellm_call_id="c1", + function_id="f1", + ) + logging_obj.rendered_model = rendered_model + return logging_obj + + +async def test_a_prompt_management_group_renders_the_deployment_prompt_and_calls_the_rendered_model( + make_core: MakeCore, host: EndpointHost +): + """The deployment's prompt_id/variables/label reach the template; the prompt-management config keys + (bitbucket_config, ...) are stripped from what is sent; the rendered model, messages and optional + params are sent directly (no routing hidden params) since the rendered model is no group.""" + core: Final = make_core( + model_list=[ + deployment( + "pm", + "p1", + model="langfuse/gpt-4o", + prompt_id="dep-pid", + prompt_variables={"a": 1}, + prompt_label="lbl", + bitbucket_config={"x": 1}, + timeout=7, + ) + ] + ) + logging_obj: Final = _prompt_logging() + response: Final = await core.acompletion(model="pm", messages=MESSAGES, litellm_logging_obj=logging_obj) + call: Final = host.calls[0] + assert (call.op, call.model, call.deployment_id, call.model_group) == ("acompletion", "openai/rendered", "p1", "pm") + sent: Final = call.kwargs + assert (sent["prompt_id"], sent["prompt_variables"], sent["prompt_label"]) == ("dep-pid", {"a": 1}, "lbl") + assert "bitbucket_config" not in sent + assert (sent["temperature"], sent["timeout"], sent["messages"]) == ( + 0.7, + 7.0, + [{"role": "system", "content": "rendered"}], + ) + assert sent["litellm_logging_obj"] is logging_obj + assert (logging_obj.seen["model"], logging_obj.seen["messages"], logging_obj.seen["prompt_id"]) == ( + "langfuse/gpt-4o", + MESSAGES, + "dep-pid", + ) + assert logging_obj.seen["injected_for_every_deployment"] is True + assert _hidden(response) == {} + + +async def test_a_prompt_rendered_to_another_group_routes_into_it_with_the_request_prompt_id_and_timeout( + make_core: MakeCore, host: EndpointHost +): + core: Final = make_core( + model_list=[ + deployment("pm", "p1", model="langfuse/gpt-4o", prompt_id="dep-pid", timeout=9), + deployment("rendered-g", "r1"), + ] + ) + response: Final = await core.acompletion( + model="pm", + messages=MESSAGES, + litellm_logging_obj=_prompt_logging("rendered-g"), + prompt_id="req-pid", + timeout=3, + ) + call: Final = host.calls[0] + assert (call.model, call.model_group, call.kwargs["prompt_id"], call.kwargs["timeout"]) == ( + "openai/r1", + "pm", + "req-pid", + 3, + ) + assert _headers(response)["x-litellm-model-group"] == "rendered-g" + + +async def test_a_single_deployment_without_a_provider_prefix_is_not_a_prompt_management_group( + make_core: MakeCore, host: EndpointHost +): + core: Final = make_core(model_list=[deployment("g", "d1", model="gpt-4o")]) + response: Final = await core.acompletion(model="g", messages=MESSAGES) + assert (host.calls[0].model, host.calls[0].deployment_id) == ("gpt-4o", "d1") + assert _hidden(response)["model_id"] == "d1" + + +# --- routed endpoints: retries and counters (6036-6603, 7646-7789, 8026-8315) ------------- + + +@pytest.mark.parametrize("endpoint", sorted(ROUTED_ENDPOINTS)) +async def test_a_routed_endpoint_retries_a_401_on_the_other_deployment( + make_core: MakeCore, host: EndpointHost, endpoint: str +): + for _ in range(4): + host.calls.clear() + host.answer("*", auth_error(), None) + core = make_core(model_list=TWO_DEPLOYMENTS, num_retries=1, disable_cooldowns=True) + await core.call(endpoint, {"model": "g", **ROUTED_ENDPOINTS[endpoint]}) + assert sorted(host.deployments_called()) == ["d1", "d2"] + + +@pytest.mark.parametrize("endpoint", sorted(ROUTED_ENDPOINTS)) +async def test_a_routed_endpoint_ignores_the_deployment_num_retries( + make_core: MakeCore, host: EndpointHost, endpoint: str +): + """Unlike acompletion and the generic endpoints, these attempt helpers never stamp the deployment's + num_retries on the error, so the router's num_retries=0 governs.""" + host.answer("*", server_error()) + core: Final = make_core(model_list=[deployment("g", "d1", num_retries=2)], num_retries=0) + with pytest.raises(litellm.InternalServerError): + await core.call(endpoint, {"model": "g", **ROUTED_ENDPOINTS[endpoint]}) + assert len(host.calls) == 1 + assert host.raised[-1].num_retries is None # pyright: ignore[reportAttributeAccessIssue] # stamped by the retry loop + + +async def test_a_generic_endpoint_retries_with_the_deployment_num_retries( + make_core: MakeCore, monkeypatch: pytest.MonkeyPatch, host: EndpointHost +): + patch_direct_endpoints(monkeypatch, host) + host.answer("*", server_error()) + core: Final = make_core(model_list=[deployment("g", "d1", num_retries=2)], num_retries=0) + with pytest.raises(litellm.InternalServerError) as raised: + await core.aresponses(model="g", input="hi") + assert len(host.calls) == 3 + assert raised.value.num_retries == 2 # pyright: ignore[reportAttributeAccessIssue] # stamped from the deployment + + +async def test_sync_image_generation_is_retried_with_the_router_num_retries(make_core: MakeCore, host: EndpointHost): + host.answer("*", rate_limit()) + core: Final = make_core(model_list=[deployment("g", "d1")], num_retries=2) + with pytest.raises(litellm.RateLimitError): + core.image_generation(model="g", prompt="p") + assert len(host.calls_to("image_generation")) == 3 + + +@pytest.mark.parametrize("endpoint", ["aimage_generation", "image_generation"]) +async def test_image_generation_for_an_unknown_group_raises_without_calling_the_host( + make_core: MakeCore, host: EndpointHost, endpoint: str +): + core: Final = make_core(model_list=[deployment("g", "d1")], num_retries=0) + with pytest.raises(litellm.BadRequestError, match="There are no healthy deployments for this model"): + await core.call(endpoint, {"model": "nope", "prompt": "p"}) + assert host.calls == [] + + +@pytest.mark.parametrize( + ("endpoint", "total_key", "fail_key"), + [ + ("aimage_generation", "openai/d1", "openai/d1"), + ("atranscription", "g", "g"), + ("aspeech", "g", "g"), + ("arerank", "openai/d1", "openai/d1"), + ("aadapter_completion", "openai/d1", "g"), + ("atext_completion", "openai/d1", "g"), + ("aembedding", "openai/d1", "openai/d1"), + ("acreate_batch", "openai/d1", "g"), + ("acancel_batch", "openai/d1", "g"), + ("image_generation", "", ""), + ], +) +async def test_call_counters_are_keyed_differently_per_endpoint( + make_core: MakeCore, host: EndpointHost, endpoint: str, total_key: str, fail_key: str +): + """total_calls/fail_calls are keyed by the deployment model for some endpoints and by the model group + for others (fail_calls even differs from total_calls for some); sync image_generation counts under "".""" + host.answer("*", rate_limit()) + core: Final = make_core(model_list=[deployment("g", "d1")], num_retries=1) + request: Final = ROUTED_ENDPOINTS.get(endpoint, {"prompt": "p"}) + with pytest.raises(litellm.RateLimitError): + await core.call(endpoint, {"model": "g", **request}) + assert dict(core.total_calls) == {total_key: 2} # pyright: ignore[reportAttributeAccessIssue] # a public counter + assert dict(core.fail_calls) == {fail_key: 2} # pyright: ignore[reportAttributeAccessIssue] # a public counter + assert dict(core.success_calls) == {} # pyright: ignore[reportAttributeAccessIssue] # a public counter + + +async def test_success_counters_count_the_deployment_model(make_core: MakeCore, host: EndpointHost): + core: Final = make_core(model_list=[deployment("g", "d1")]) + await core.aspeech(model="g", input="x", voice="alloy") + await core.aembedding(model="g", input=["a"]) + assert dict(core.total_calls) == {"g": 1, "openai/d1": 1} # pyright: ignore[reportAttributeAccessIssue] # a public counter + assert dict(core.success_calls) == {"g": 1, "openai/d1": 1} # pyright: ignore[reportAttributeAccessIssue] # a public counter + + +# --- guardrails (6692-6721) ----------------------------------------------------------------- + + +async def test_aguardrail_picks_among_same_named_guardrails_by_weight(make_core: MakeCore, host: ScriptedHost): + picked: Final[list[object]] = [] + + async def guard(**kwargs: object) -> str: # kwargs-ok: the guardrail's keyword payload + selected: Final = kwargs["selected_guardrail"] + assert isinstance(selected, Mapping) + picked.append(selected["litellm_params"]["guardrail"]) + return "verdict" + + guardrails: Final = [ + {"guardrail_name": "gr", "litellm_params": {"guardrail": "never", "weight": 0}}, + {"guardrail_name": "gr", "litellm_params": {"guardrail": "always", "weight": 1}}, + ] + core: Final = make_core(model_list=[deployment("g", "d1")], guardrail_list=guardrails) + for _ in range(10): + await core.aguardrail(guardrail_name="gr", original_function=guard) + assert set(picked) == {"always"} + + +async def test_aguardrail_without_weights_spreads_over_same_named_guardrails(make_core: MakeCore, host: ScriptedHost): + picked: Final[set[object]] = set() + + async def guard(**kwargs: object) -> str: # kwargs-ok: the guardrail's keyword payload + selected: Final = kwargs["selected_guardrail"] + assert isinstance(selected, Mapping) + picked.add(selected["litellm_params"]["guardrail"]) + return "verdict" + + guardrails: Final = [ + {"guardrail_name": "gr", "litellm_params": {"guardrail": "a"}}, + {"guardrail_name": "gr", "litellm_params": {"guardrail": "b"}}, + ] + core: Final = make_core(model_list=[deployment("g", "d1")], guardrail_list=guardrails) + for _ in range(60): + await core.aguardrail(guardrail_name="gr", original_function=guard) + assert picked == {"a", "b"} + + +# --- generic calls: pass-through endpoint rewrite (6754-6781, 7561-7627) -------------------- + + +async def test_async_passthrough_rewrites_only_whole_path_segments_naming_the_group( + make_core: MakeCore, monkeypatch: pytest.MonkeyPatch, host: EndpointHost +): + patch_direct_endpoints(monkeypatch, host) + core: Final = make_core(model_list=[deployment("g", "d1", model="bedrock/us.anthropic.claude")]) + await core.allm_passthrough_route(model="g", method="POST", endpoint="/model/g/invoke/gg/g:stream") + assert host.calls[0].kwargs["endpoint"] == "/model/us.anthropic.claude/invoke/gg/us.anthropic.claude:stream" + + +async def test_sync_generic_call_replaces_every_substring_of_the_group_name_in_the_endpoint( + make_core: MakeCore, monkeypatch: pytest.MonkeyPatch, host: EndpointHost +): + """The sync path uses str.replace with the full deployment model (provider prefix kept), so a group + name inside other path segments is rewritten too, unlike the async path's segment-bounded rewrite.""" + patch_direct_endpoints(monkeypatch, host) + core: Final = make_core(model_list=[deployment("g", "d1")]) + core.responses(model="g", input="hi", endpoint="/gg/g") + call: Final = host.calls[0] + assert (call.model, call.kwargs["endpoint"]) == ("openai/d1", "/openai/d1openai/d1/openai/d1") + + +# --- batches (8112-8362) -------------------------------------------------------------------- + + +async def test_aretrieve_batch_sends_the_deployment_provider_over_the_request_one( + make_core: MakeCore, host: EndpointHost +): + host.answer("*", _batch()) + core: Final = make_core(model_list=[deployment("g", "d1"), deployment("h", "d2", custom_llm_provider="azure")]) + await core.aretrieve_batch(model="g", batch_id="b1", custom_llm_provider="azure") + await core.aretrieve_batch(model="h", batch_id="b1") + assert [call.kwargs["custom_llm_provider"] for call in host.calls] == ["openai", "azure"] + + +async def test_alist_batches_takes_the_first_set_first_id_and_the_last_answer_last_id( + make_core: MakeCore, host: EndpointHost +): + """last_id is overwritten by every answer in deployment order, even with None.""" + host.answer("openai/d1", SimpleNamespace(first_id=None, last_id="b", data=[1], has_more=False)) + host.answer("openai/d2", SimpleNamespace(first_id="c", last_id=None, data=[2], has_more=False)) + core: Final = make_core(model_list=TWO_DEPLOYMENTS) + assert await core.alist_batches(model="g") == { + "object": "list", + "data": [1, 2], + "first_id": "c", + "last_id": None, + "has_more": False, + } + + +async def test_acancel_batch_resolves_the_deployment_credential_but_acreate_batch_does_not( + make_core: MakeCore, monkeypatch: pytest.MonkeyPatch, host: EndpointHost +): + monkeypatch.setattr( + litellm, + "credential_list", + [ + CredentialItem( + credential_name="cred1", + credential_values={"api_key": "sk-cred", "api_base": "https://cred.example"}, + credential_info={}, + ) + ], + ) + host.answer("*", _batch()) + core: Final = make_core(model_list=[deployment("g", "d1", litellm_credential_name="cred1")]) + await core.acancel_batch(model="g", batch_id="b1") + await core.acreate_batch(model="g", **ROUTED_ENDPOINTS["acreate_batch"]) + cancel, create = host.calls[0].kwargs, host.calls[1].kwargs + assert (cancel["api_key"], cancel["api_base"], "litellm_credential_name" in cancel) == ( + "sk-cred", + "https://cred.example", + False, + ) + assert (create["api_key"], create.get("api_base"), create["litellm_credential_name"]) == ("sk-test", None, "cred1") + + +# --- files (7814-7934) ---------------------------------------------------------------------- + + +async def test_acreate_file_passthrough_keeps_the_batch_file_model(make_core: MakeCore, host: EndpointHost): + upload: Final = io.BytesIO(b'{"custom_id": "1", "body": {"model": "g", "messages": []}}\n') + core: Final = make_core(model_list=[deployment("g", "d1")]) + await core.acreate_file(model="g", file=upload, purpose="batch", passthrough=True) + sent: Final = host.calls[0].kwargs["file"] + assert isinstance(sent, io.BytesIO) + assert b'"model": "g"' in sent.getvalue() diff --git a/tests/unit/test_router/core/test_gaps_l5600_anthropic.py b/tests/unit/test_router/core/test_gaps_l5600_anthropic.py new file mode 100644 index 00000000000..e2b19058a50 --- /dev/null +++ b/tests/unit/test_router/core/test_gaps_l5600_anthropic.py @@ -0,0 +1,103 @@ +""" +Gap characterization for /v1/messages streaming recovery (``router_python_core.py`` 7106-7455): a +``MidStreamFallbackError`` raised by the source stream itself, a same-group stream retry whose own +attempt raises before opening, and a request-level retry policy that is not a ``RetryPolicy``. +""" + +from collections.abc import Callable +from typing import Final + +import pytest + +import litellm +from litellm.exceptions import MidStreamFallbackError +from litellm.router_core import RouterCore +from tests.unit.test_router.core.scripted_host import ScriptedHost, deployment, rate_limit, server_error +from tests.unit.test_router.core.streaming_helpers import ( + MESSAGE_START, + MESSAGE_STOP, + MESSAGES, + anthropic_stream, + collect_bytes, + drain_into, + error_frame, + install_orjson_stub_if_missing, + text_delta, +) + +install_orjson_stub_if_missing() + +MakeCore = Callable[..., RouterCore] +GROUPS_AB: Final = [deployment("a", "a"), deployment("b", "b")] + + +async def _open(core: RouterCore, **request: object) -> object: # kwargs-ok: request kwargs + return await core.aanthropic_messages(model="a", messages=MESSAGES, max_tokens=5, stream=True, **request) # pyright: ignore[reportAttributeAccessIssue] # endpoint method + + +def _mid_stream(*, pre_first_chunk: bool) -> MidStreamFallbackError: + return MidStreamFallbackError( + message="bridge dropped", + model="a", + llm_provider="openai", + original_exception=rate_limit("a", "bridge original"), + is_pre_first_chunk=pre_first_chunk, + ) + + +async def test_a_raised_mid_stream_error_not_before_the_first_chunk_flushes_held_frames_and_raises_the_original( + make_core: MakeCore, host: ScriptedHost +): + host.answer("openai/a", anthropic_stream(MESSAGE_START, error=_mid_stream(pre_first_chunk=False))) + host.answer("openai/b", anthropic_stream(MESSAGE_START, text_delta("B"), MESSAGE_STOP)) + core: Final = make_core(model_list=GROUPS_AB, fallbacks=[{"a": ["b"]}], num_retries=0) + stream: Final = await _open(core) + seen: Final[list[object]] = [] + with pytest.raises(litellm.RateLimitError, match="bridge original"): + await drain_into(stream, seen) + assert seen == [MESSAGE_START] + assert host.deployments_called() == ["a"] + + +async def test_a_raised_mid_stream_error_before_the_first_chunk_falls_back(make_core: MakeCore, host: ScriptedHost): + host.answer("openai/a", anthropic_stream(MESSAGE_START, error=_mid_stream(pre_first_chunk=True))) + host.answer("openai/b", anthropic_stream(MESSAGE_START, text_delta("B"), MESSAGE_STOP)) + core: Final = make_core(model_list=GROUPS_AB, fallbacks=[{"a": ["b"]}], num_retries=0) + frames: Final = await collect_bytes(await _open(core)) + assert frames == [MESSAGE_START, text_delta("B"), MESSAGE_STOP] + assert host.deployments_called() == ["a", "b"] + + +async def test_a_stream_retry_that_raises_before_opening_spends_a_retry_and_the_next_one_serves( + make_core: MakeCore, host: ScriptedHost +): + host.answer( + "openai/a", + anthropic_stream(MESSAGE_START, error_frame()), + server_error("a"), + anthropic_stream(MESSAGE_START, text_delta("A"), MESSAGE_STOP), + ) + core: Final = make_core(model_list=[deployment("a", "a")], num_retries=2) + stream: Final = await _open(core) + frames: Final = await collect_bytes(stream) + assert frames == [MESSAGE_START, text_delta("A"), MESSAGE_STOP] + assert host.deployments_called() == ["a", "a", "a"] + assert len(host.sleeps) == 2 + headers: Final = getattr(stream, "_hidden_params")["additional_headers"] + assert (headers["x-litellm-attempted-retries"], headers["x-litellm-max-retries"]) == (2, 2) + + +async def test_a_request_retry_policy_that_is_not_a_retry_policy_streams_with_the_plain_budget( + make_core: MakeCore, host: ScriptedHost +): + host.answer( + "openai/a", + anthropic_stream(MESSAGE_START, error_frame("rate_limit_error", "slow")), + anthropic_stream(MESSAGE_START, text_delta("A"), MESSAGE_STOP), + ) + core: Final = make_core(model_list=[deployment("a", "a")], num_retries=1) + stream: Final = await _open(core, model_group_retry_policy={"a": 5}) + frames: Final = await collect_bytes(stream) + assert frames == [MESSAGE_START, text_delta("A"), MESSAGE_STOP] + assert host.deployments_called() == ["a", "a"] + assert getattr(stream, "_hidden_params")["additional_headers"]["x-litellm-max-retries"] == 1 diff --git a/tests/unit/test_router/core/test_gaps_l8400_precall.py b/tests/unit/test_router/core/test_gaps_l8400_precall.py new file mode 100644 index 00000000000..1e64257bac8 --- /dev/null +++ b/tests/unit/test_router/core/test_gaps_l8400_precall.py @@ -0,0 +1,178 @@ +""" +Gaps in the routing-strategy pre-call checks and the Responses-websocket turn admission of ``PythonCore`` +(router_python_core.py ``routing_strategy_pre_call_checks``, ``async_routing_strategy_pre_call_checks``, +``async_admit_deployment_turn``). A ``CustomLogger`` on the host's ``callbacks`` list (``litellm.callbacks``) +takes part in those checks through ``pre_call_check`` / ``async_pre_call_check``. +""" + +from collections.abc import Callable +from typing import Final + +import pytest + +import litellm +from litellm.integrations.custom_logger import CustomLogger +from litellm.router_core import RouterCore +from tests.unit.test_router.core.endpoints_helpers import EndpointHost +from tests.unit.test_router.core.retries_health_helpers import attempt, attempt_sync, cooled_down +from tests.unit.test_router.core.scripted_host import ScriptedHost, deployment, rate_limit + +MakeCore = Callable[..., RouterCore] + + +class PreCallGate(CustomLogger): + """Records the deployment ids its pre-call checks saw and raises ``error`` (when set) from them.""" + + def __init__(self, error: Exception | None = None) -> None: + super().__init__() + self.error = error + self.seen: Final[list[tuple[str, str]]] = [] + + async def async_pre_call_check(self, deployment: dict, parent_otel_span: object) -> dict | None: # pyright: ignore[reportIncompatibleMethodOverride] + self.seen.append(("async", deployment["model_info"]["id"])) + if self.error is not None: + raise self.error + return None + + def pre_call_check(self, deployment: dict) -> dict | None: + self.seen.append(("sync", deployment["model_info"]["id"])) + if self.error is not None: + raise self.error + return None + + +def gated(host: ScriptedHost, error: Exception | None = None) -> PreCallGate: + gate: Final = PreCallGate(error) + host.lists["callbacks"].append(gate) + return gate + + +GROUP: Final = [deployment("g", "a"), deployment("g", "b")] + + +# --- async pre-call checks ---------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_a_rate_limit_from_a_pre_call_check_cools_the_deployment_down_and_is_raised( + make_core: MakeCore, host: ScriptedHost +) -> None: + core: Final = make_core(model_list=GROUP, num_retries=0, cooldown_time=33) + gate: Final = gated(host, rate_limit(message="gate says no")) + + error: Final = await attempt(core) + + assert isinstance(error, litellm.RateLimitError) + assert str(error) == "litellm.RateLimitError: gate says no" + assert host.calls == [] + ((kind, checked),) = gate.seen + assert kind == "async" + cooled: Final = cooled_down(core, ["a", "b"]) + assert list(cooled) == [checked] + assert cooled[checked]["cooldown_time"] == 33 + + +@pytest.mark.asyncio +async def test_any_other_error_from_a_pre_call_check_is_raised_without_a_cooldown( + make_core: MakeCore, host: ScriptedHost +) -> None: + core: Final = make_core(model_list=GROUP, num_retries=0, cooldown_time=33) + gate: Final = gated(host, ValueError("gate broke")) + + error: Final = await attempt(core) + + assert type(error) is ValueError + assert str(error) == "gate broke" + assert host.calls == [] + assert len(gate.seen) == 1 + assert cooled_down(core, ["a", "b"]) == {} + + +@pytest.mark.asyncio +async def test_a_passing_pre_call_check_sees_the_deployment_that_is_then_called( + make_core: MakeCore, host: ScriptedHost +) -> None: + core: Final = make_core(model_list=GROUP, num_retries=0) + gate: Final = gated(host) + + await attempt(core) + + assert gate.seen == [("async", host.deployments_called("acompletion")[0])] + + +def test_a_rate_limit_from_a_sync_pre_call_check_is_raised_without_a_cooldown( + make_core: MakeCore, host: ScriptedHost +) -> None: + """Oddity: the sync twin raises the gate's error but, unlike the async path, sets no cooldown.""" + core: Final = make_core(model_list=GROUP, num_retries=0, cooldown_time=33) + gate: Final = gated(host, rate_limit(message="gate says no")) + + error: Final = attempt_sync(core) + + assert isinstance(error, litellm.RateLimitError) + assert host.calls == [] + assert [kind for kind, _ in gate.seen] == ["sync"] + assert cooled_down(core, ["a", "b"]) == {} + + +# --- Responses websocket turn admission ---------------------------------------------------------- + + +class TestResponsesWebsocketTurnAdmission: + """The core hands the websocket handler a ``turn_admission`` callable that re-runs the deployment's + admission (budget limiter, routing strategy, pre-call checks) before every later turn.""" + + @pytest.fixture + def host(self) -> EndpointHost: + return EndpointHost() + + @staticmethod + async def open_socket(core: RouterCore, host: EndpointHost) -> Callable[[], object]: + await core.call("_aresponses_websocket", {"model": "g"}) + (call,) = host.calls_to("_aresponses_websocket") + admission: Final = call.kwargs["turn_admission"] + assert callable(admission) + return admission + + @pytest.mark.asyncio + async def test_simple_shuffle_admits_every_turn_and_runs_the_pre_call_checks( + self, make_core: MakeCore, host: EndpointHost + ) -> None: + core: Final = make_core(model_list=[deployment("g", "a", rpm=1)]) + gate: Final = gated(host) + admit: Final = await self.open_socket(core, host) + + admitted: Final = [dict(await admit()) for _ in range(3)] # pyright: ignore[reportGeneralTypeIssues] + + assert admitted == [{}, {}, {}] + assert gate.seen == [("async", "a")] * 4 # the socket opening, then each turn + + @pytest.mark.asyncio + async def test_a_rate_limit_from_a_pre_call_check_refuses_the_turn( + self, make_core: MakeCore, host: EndpointHost + ) -> None: + core: Final = make_core(model_list=[deployment("g", "a")]) + gate: Final = gated(host) + admit: Final = await self.open_socket(core, host) + gate.error = rate_limit(message="turn refused") + + with pytest.raises(litellm.RateLimitError, match="turn refused"): + await admit() # pyright: ignore[reportGeneralTypeIssues] + + @pytest.mark.asyncio + async def test_usage_based_routing_admits_turns_until_the_deployment_rpm_is_used_up( + self, make_core: MakeCore, host: EndpointHost + ) -> None: + """Opening the socket already counts two requests (its pre-call check and its success), so an + rpm of 5 leaves room for three turns.""" + core: Final = make_core(model_list=[deployment("g", "a", rpm=5)], routing_strategy="usage-based-routing-v2") + admit: Final = await self.open_socket(core, host) + + admitted: Final = [dict(await admit()) for _ in range(3)] # pyright: ignore[reportGeneralTypeIssues] + with pytest.raises(litellm.RateLimitError) as refused: + await admit() # pyright: ignore[reportGeneralTypeIssues] + + assert admitted == [{}, {}, {}] + assert str(refused.value).startswith( + "litellm.RateLimitError: No deployments available for selected model. Passed model=g. Deployments=" + ) diff --git a/tests/unit/test_router/core/test_gaps_l8400_pricing.py b/tests/unit/test_router/core/test_gaps_l8400_pricing.py new file mode 100644 index 00000000000..4c0df5a17b6 --- /dev/null +++ b/tests/unit/test_router/core/test_gaps_l8400_pricing.py @@ -0,0 +1,240 @@ +""" +Gaps in the built-in pricing a custom-priced deployment inherits when the core registers it in +``litellm.model_cost`` (router_python_core.py ``_inherit_builtin_cache_pricing``, ``_cost_map_backend_model``, +``_inherit_builtin_service_tier_pricing``, ``_inherit_builtin_base_rates_for_off_peak``, +``_inherit_builtin_tiered_output_rate``), and the deployment checks of ``_create_deployment`` +(access windows, PTU configuration). Expected rates are read from the built-in cost map, not hardcoded. +""" + +from collections.abc import Mapping +from typing import Final + +import pytest + +import litellm +from tests.unit.test_router.core.registry_helpers import MakeCore, unique_id +from tests.unit.test_router.core.scripted_host import deployment + +OFF_PEAK: Final = {"hours_utc": "01:00-02:00", "input_cost_per_token": 1e-9} + + +def registered(make_core: MakeCore, model: str, **params: object) -> Mapping[str, object]: # kwargs-ok: litellm_params + """The ``litellm.model_cost`` entry registered under a fresh deployment id.""" + model_id: Final = unique_id("l8400") + make_core(model_list=[deployment("g", model_id, model=model, **params)]) + return litellm.model_cost[model_id] + + +def builtin(key: str) -> Mapping[str, object]: + return litellm.model_cost[key] + + +# --- service-tier rates --------------------------------------------------------------------------- + + +def test_custom_input_pricing_inherits_the_backend_service_tier_rates(make_core: MakeCore) -> None: + entry: Final = registered(make_core, "openai/gpt-4o", input_cost_per_token=1e-3) + + for field in ( + "input_cost_per_token_priority", + "output_cost_per_token_priority", + "cache_read_input_token_cost_priority", + ): + assert entry[field] == builtin("gpt-4o")[field] + assert entry.get("input_cost_per_token_batches") is None # batch rates are not a service tier + + +def test_a_user_set_service_tier_rate_wins_over_the_inherited_one(make_core: MakeCore) -> None: + entry: Final = registered(make_core, "openai/gpt-4o", input_cost_per_token=1e-3, input_cost_per_token_priority=9e-3) + + assert entry["input_cost_per_token_priority"] == 9e-3 + assert entry["output_cost_per_token_priority"] == builtin("gpt-4o")["output_cost_per_token_priority"] + + +def test_output_only_custom_pricing_inherits_service_tiers_but_not_cache_pricing(make_core: MakeCore) -> None: + entry: Final = registered(make_core, "openai/gpt-4o", output_cost_per_token=1e-3) + + assert entry["input_cost_per_token_priority"] == builtin("gpt-4o")["input_cost_per_token_priority"] + assert entry.get("cache_read_input_token_cost") is None + + +def test_a_backend_missing_from_the_cost_map_contributes_nothing(make_core: MakeCore) -> None: + entry: Final = registered(make_core, "openai/l8400-not-a-real-model", input_cost_per_token=1e-3) + + assert entry["input_cost_per_token"] == 1e-3 + assert entry.get("cache_read_input_token_cost") is None + assert entry.get("input_cost_per_token_priority") is None + + +# --- which backend model the rates come from ------------------------------------------------------ + + +@pytest.mark.parametrize( + "params", + [{"model_info": {"base_model": "gpt-5"}}, {"base_model": "gpt-5"}], + ids=["model-info-base-model", "litellm-params-base-model"], +) +def test_the_base_model_names_the_backend_the_rates_are_inherited_from( + make_core: MakeCore, params: dict[str, object] +) -> None: + entry: Final = registered(make_core, "openai/my-fine-tune", input_cost_per_token=1e-3, **params) + + assert entry["cache_read_input_token_cost"] == builtin("gpt-5")["cache_read_input_token_cost"] + assert entry["input_cost_per_token_flex"] == builtin("gpt-5")["input_cost_per_token_flex"] + + +# --- explicit cache-read rate ----------------------------------------------------------------------- + + +def test_explicit_cache_read_rate_is_inherited_only_without_a_user_cache_read_rate(make_core: MakeCore) -> None: + backend: Final = builtin("dashscope/qwen3.7-max") + inherited: Final = registered(make_core, "dashscope/qwen3.7-max", input_cost_per_token=1e-3) + user_set: Final = registered( + make_core, "dashscope/qwen3.7-max", input_cost_per_token=1e-3, cache_read_input_token_cost=1e-8 + ) + + assert inherited["explicit_cache_read_input_token_cost"] == backend["explicit_cache_read_input_token_cost"] + assert inherited["cache_read_input_token_cost"] == backend["cache_read_input_token_cost"] + assert user_set.get("explicit_cache_read_input_token_cost") is None + assert user_set["cache_read_input_token_cost"] == 1e-8 + + +# --- off-peak only ---------------------------------------------------------------------------------- + + +def test_an_off_peak_only_deployment_copies_the_whole_backend_entry(make_core: MakeCore) -> None: + backend: Final = builtin("deepseek-v4-flash") + entry: Final = registered(make_core, "deepseek/deepseek-v4-flash", model_info={"off_peak_pricing": OFF_PEAK}) + + assert entry["off_peak_pricing"] == OFF_PEAK + for field in ("input_cost_per_token", "output_cost_per_token", "cache_read_input_token_cost", "max_input_tokens"): + assert entry[field] == backend[field] + assert entry["source"] == backend["source"] + + +def test_an_off_peak_deployment_keeps_its_own_fields(make_core: MakeCore) -> None: + entry: Final = registered( + make_core, "deepseek/deepseek-v4-flash", model_info={"off_peak_pricing": OFF_PEAK, "max_input_tokens": 7} + ) + + assert entry["max_input_tokens"] == 7 + assert entry["input_cost_per_token"] == builtin("deepseek-v4-flash")["input_cost_per_token"] + + +def test_an_off_peak_deployment_with_a_base_rate_copies_nothing_more(make_core: MakeCore) -> None: + entry: Final = registered( + make_core, + "deepseek/deepseek-v4-flash", + model_info={"off_peak_pricing": OFF_PEAK}, + input_cost_per_token=5e-3, + ) + + assert entry["input_cost_per_token"] == 5e-3 + assert entry.get("output_cost_per_token") is None + assert entry.get("max_input_tokens") is None + + +def test_off_peak_pricing_given_in_litellm_params_is_not_registered(make_core: MakeCore) -> None: + """Oddity: only ``model_info.off_peak_pricing`` counts; the same key in litellm_params is dropped.""" + entry: Final = registered(make_core, "deepseek/deepseek-v4-flash", off_peak_pricing=OFF_PEAK) + + assert entry.get("off_peak_pricing") is None + assert entry.get("input_cost_per_token") is None + + +# --- tiered pricing output rate --------------------------------------------------------------------- + + +def test_a_tier_table_without_output_rates_inherits_the_backend_flat_output_rate(make_core: MakeCore) -> None: + entry: Final = registered( + make_core, "openai/gpt-4o", tiered_pricing=[{"input_cost_per_token": 1e-9, "range": [0, 1000]}] + ) + + assert entry["output_cost_per_token"] == builtin("gpt-4o")["output_cost_per_token"] + + +@pytest.mark.parametrize( + ("model", "params"), + [ + ("dashscope/qwen-flash", {"tiered_pricing": [{"input_cost_per_token": 1e-9, "range": [0, 1000]}]}), + ( + "openai/gpt-4o", + {"tiered_pricing": [{"input_cost_per_token": 1e-9, "output_cost_per_token": 2e-9, "range": [0, 1000]}]}, + ), + ( + "openai/gpt-4o", + {"tiered_pricing": [{"input_cost_per_token": 1e-9, "range": [0, 1000]}], "output_cost_per_token": 3e-9}, + ), + ], + ids=["tiered-only-backend", "every-tier-has-an-output-rate", "user-output-rate"], +) +def test_the_tiered_output_rate_is_not_inherited(make_core: MakeCore, model: str, params: dict[str, object]) -> None: + entry: Final = registered(make_core, model, **params) + + assert entry.get("output_cost_per_token") == params.get("output_cost_per_token") + + +# --- _create_deployment checks ------------------------------------------------------------------------ + + +def test_an_invalid_access_windows_config_raises_naming_the_model(make_core: MakeCore) -> None: + with pytest.raises(ValueError, match=r"^model 'gpt': invalid model_info\.access_windows: "): + make_core(model_list=[deployment("gpt", "a", model_info={"access_windows": "bad"})]) + + +@pytest.mark.parametrize( + "model_info", + [ + {"access_windows": "bad"}, + {"ptu_shares": {"team-a": 1}}, + ], + ids=["access-windows", "ptu-shares"], +) +def test_ignore_invalid_deployments_drops_a_deployment_with_an_invalid_config( + make_core: MakeCore, model_info: dict[str, object] +) -> None: + core: Final = make_core( + model_list=[deployment("gpt", "a", model_info=model_info), deployment("gpt", "b")], + ignore_invalid_deployments=True, + ) + + assert core.registry().names(kind="ids") == ["b"] + + +def test_ptu_shares_without_a_count_and_rate_raise(make_core: MakeCore) -> None: + with pytest.raises( + ValueError, + match=r"^PTU configuration on model 'gpt' is invalid: ptu_count and cost_per_ptu_per_hour are required " + r"when ptu_shares is set", + ): + make_core(model_list=[deployment("gpt", "a", model_info={"ptu_shares": {"team-a": 1}})]) + + +def test_an_inverted_ptu_window_with_shares_raises(make_core: MakeCore) -> None: + with pytest.raises( + ValueError, + match=r"^PTU configuration on model 'gpt' is invalid: ptu_effective_to must be after ptu_effective_from", + ): + make_core( + model_list=[ + deployment( + "gpt", + "a", + model_info={ + "ptu_effective_from": "2026-01-02T00:00:00Z", + "ptu_effective_to": "2026-01-01T00:00:00Z", + "ptu_shares": {"team-a": 1}, + "ptu_count": 1, + "cost_per_ptu_per_hour": 1, + }, + ) + ] + ) + + +def test_a_db_sourced_deployment_skips_the_ptu_config_check(make_core: MakeCore) -> None: + core: Final = make_core( + model_list=[deployment("gpt", "a", model_info={"ptu_shares": {"team-a": 1}, "db_model": True})] + ) + + assert core.registry().names(kind="ids") == ["a"] diff --git a/tests/unit/test_router/core/test_gaps_l8400_retries.py b/tests/unit/test_router/core/test_gaps_l8400_retries.py new file mode 100644 index 00000000000..34373f6e03e --- /dev/null +++ b/tests/unit/test_router/core/test_gaps_l8400_retries.py @@ -0,0 +1,138 @@ +""" +Gaps in the retry loop and the per-attempt bookkeeping of ``PythonCore`` (router_python_core.py +lines 8400-11100): the ``UnsupportedEndpointError`` skip-and-stop rule, the retry breadcrumb cap and +``base_model``, and the request the async failure callback counts toward group usage. +""" + +from typing import Final + +import pytest + +import litellm +from tests.unit.test_router.core.retries_health_helpers import ( + BackgroundSuccessLoggingHost, + MakeCore, + MetadataRecorder, + attempt, +) +from tests.unit.test_router.core.scripted_host import ScriptedHost, deployment, rate_limit, server_error + + +def unsupported_endpoint() -> litellm.UnsupportedEndpointError: + return litellm.UnsupportedEndpointError(message="endpoint not supported", llm_provider="openai", model="gpt-4o") + + +# --- UnsupportedEndpointError: retried on the other deployments, stopped once every one refused --------- + + +@pytest.mark.asyncio +async def test_unsupported_endpoint_is_retried_once_per_deployment_then_raised( + make_core: MakeCore, host: ScriptedHost +) -> None: + """A 400 is normally never retried; ``UnsupportedEndpointError`` is, but each deployment that refused + is skipped and the loop stops as soon as every deployment of the group refused (num_retries=3 unused).""" + core: Final = make_core(model_list=[deployment("g", "a"), deployment("g", "b")], num_retries=3) + host.answer("*", lambda **_: unsupported_endpoint()) + + error: Final = await attempt(core) + + assert isinstance(error, litellm.UnsupportedEndpointError) + called: Final = host.deployments_called("acompletion") + assert sorted(called) == ["a", "b"] + assert host.sleeps == [0] + + +@pytest.mark.asyncio +async def test_unsupported_endpoint_on_a_lone_deployment_is_raised_without_a_retry( + make_core: MakeCore, host: ScriptedHost +) -> None: + core: Final = make_core(model_list=[deployment("g", "a")], num_retries=3) + host.answer("*", lambda **_: unsupported_endpoint()) + + error: Final = await attempt(core) + + assert isinstance(error, litellm.UnsupportedEndpointError) + assert host.deployments_called("acompletion") == ["a"] + assert host.sleeps == [] + + +# --- retry breadcrumbs ---------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_previous_models_keeps_only_the_four_most_recent_failures( + make_core: MakeCore, host: ScriptedHost +) -> None: + recorder: Final = MetadataRecorder() + core: Final = make_core(model_list=[deployment("g", "a")], num_retries=6) + host.answer("*", *([recorder.then(server_error)] * 6), recorder.then(None)) + + response: Final = await attempt(core) + + assert not isinstance(response, Exception) + trails: Final = [snapshot["previous_models"] for snapshot in recorder.snapshots] + assert trails[0] is None + assert [len(trail) for trail in trails[1:]] == [1, 2, 3, 4, 4, 4] # pyright: ignore[reportArgumentType] + assert [entry["attempted_retries"] for entry in trails[-1]] == [2, 3, 4, 5] # pyright: ignore[reportGeneralTypeIssues, reportOptionalIterable] + assert [snapshot["request_retry_count"] for snapshot in recorder.snapshots] == [None, 1, 2, 3, 4, 5, 6] + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("params", "expected"), + [ + ({"base_model": "azure/gpt-4o"}, "azure/gpt-4o"), + ({"model_info": {"base_model": "azure/gpt-4o-mini"}}, "azure/gpt-4o-mini"), + ({}, None), + ], + ids=["litellm-params-base-model", "model-info-base-model", "none"], +) +async def test_a_retry_breadcrumb_names_the_deployments_base_model( + make_core: MakeCore, host: ScriptedHost, params: dict[str, object], expected: str | None +) -> None: + recorder: Final = MetadataRecorder() + core: Final = make_core(model_list=[deployment("g", "a", **params)], num_retries=1) + host.answer("*", recorder.then(server_error), recorder.then(None)) + + await attempt(core) + + (entry,) = recorder.snapshots[1]["previous_models"] # pyright: ignore[reportGeneralTypeIssues] + assert entry["base_model"] == expected # pyright: ignore[reportIndexIssue] + + +# --- the async failure callback counts the failed request --------------------------------------- + + +@pytest.mark.asyncio +async def test_a_failed_call_counts_a_request_but_no_tokens_in_group_usage( + make_core: MakeCore, host: ScriptedHost +) -> None: + core: Final = make_core(model_list=[deployment("g", "a", rpm=5)], num_retries=0) + host.answer("*", rate_limit()) + + error: Final = await attempt(core) + + assert isinstance(error, litellm.RateLimitError) + assert await core.health().usage("g") == (None, 1) + + +class TestUsageWhenTheHostLogsSuccessAfterTheResponse: + """litellm runs the async success callbacks after the endpoint returned (``ScriptedHost`` runs them + before, which lets the success callback and the router's own count both add the same response).""" + + @pytest.fixture + def host(self) -> BackgroundSuccessLoggingHost: + return BackgroundSuccessLoggingHost() + + @pytest.mark.asyncio + @pytest.mark.parametrize(("failures", "requests"), [(0, 1), (1, 2), (2, 3)]) + async def test_every_failed_attempt_and_the_final_success_count_as_requests( + self, make_core: MakeCore, host: BackgroundSuccessLoggingHost, failures: int, requests: int + ) -> None: + core: Final = make_core(model_list=[deployment("g", "a", rpm=50)], num_retries=2) + host.answer("*", *([server_error()] * failures), None) + + await attempt(core) + await host.drain() + + assert await core.health().usage("g") == (10, requests) diff --git a/tests/unit/test_router/core/test_gaps_pick_input_and_compaction.py b/tests/unit/test_router/core/test_gaps_pick_input_and_compaction.py new file mode 100644 index 00000000000..ded4b50c45f --- /dev/null +++ b/tests/unit/test_router/core/test_gaps_pick_input_and_compaction.py @@ -0,0 +1,253 @@ +""" +Parity gaps differential fuzzing found (regress5 R2, U4, N10): the ``messages`` and ``input`` each +endpoint hands its pick, which decide whether the pre-call checks run and what they count; +``print_deployment`` on a ``Deployment``; and the context compaction checks of an auto-router. +""" + +import io +from collections.abc import Callable, Mapping +from typing import Final + +import pytest + +import litellm +from litellm.router_api import RoutingContext +from litellm.router_core import RouterCore +from litellm.router_strategy.complexity_router.context_compaction import CompactionState +from litellm.types.router import Deployment, LiteLLM_Params, RouterRateLimitError +from tests.unit.test_router.core.construction_precall_helpers import windowed +from tests.unit.test_router.core.endpoints_helpers import EndpointHost +from tests.unit.test_router.core.scripted_host import deployment + +MakeCore = Callable[..., RouterCore] + +MESSAGES: Final = [{"role": "user", "content": "hi"}] +LONG_TEXT: Final = "word " * 200 +REGIONAL: Final = [deployment("g", "eu", region_name="eu"), deployment("g", "ap", region_name="ap")] +EU_ONLY: Final = [deployment("g", "eu", region_name="eu")] +BATCH: Final = {"input_file_id": "f", "endpoint": "/v1/chat/completions", "completion_window": "24h"} + + +@pytest.fixture +def host() -> EndpointHost: + return EndpointHost() + + +# --- R2: the pre-call checks of a prompt-only request --------------------------------------- + + +@pytest.mark.asyncio +async def test_async_text_completion_is_refused_outside_the_allowed_region( + make_core: MakeCore, host: EndpointHost +) -> None: + core: Final = make_core(model_list=EU_ONLY, enable_pre_call_checks=True, num_retries=0) + with pytest.raises(RouterRateLimitError): + await core.call("atext_completion", {"model": "g", "prompt": "hi", "allowed_model_region": "ap"}) + assert host.calls_to("atext_completion") == [] + + +def test_sync_text_completion_is_refused_outside_the_allowed_region(make_core: MakeCore, host: EndpointHost) -> None: + core: Final = make_core(model_list=EU_ONLY, enable_pre_call_checks=True, num_retries=0) + with pytest.raises(RouterRateLimitError): + core.text_completion(model="g", prompt="hi", allowed_model_region="ap") # pyright: ignore[reportAttributeAccessIssue] + assert host.calls_to("text_completion") == [] + + +@pytest.mark.asyncio +async def test_async_text_completion_keeps_the_deployment_in_the_allowed_region( + make_core: MakeCore, host: EndpointHost +) -> None: + core: Final = make_core(model_list=REGIONAL, enable_pre_call_checks=True, num_retries=0) + await core.call("atext_completion", {"model": "g", "prompt": "hi", "allowed_model_region": "ap"}) + assert host.deployments_called("atext_completion") == ["ap"] + + +def test_sync_text_completion_keeps_the_deployment_in_the_allowed_region( + make_core: MakeCore, host: EndpointHost +) -> None: + core: Final = make_core(model_list=REGIONAL, enable_pre_call_checks=True, num_retries=0) + core.text_completion(model="g", prompt="hi", allowed_model_region="ap") # pyright: ignore[reportAttributeAccessIssue] + # The sync text path sends no ``model_info``: the params name the deployment. + assert [call.kwargs["region_name"] for call in host.calls_to("text_completion")] == ["ap"] + + +@pytest.mark.asyncio +async def test_async_text_completion_counts_the_prompt_against_the_context_window( + make_core: MakeCore, host: EndpointHost +) -> None: + core: Final = make_core(model_list=[windowed("g", "small", 5)], enable_pre_call_checks=True, num_retries=0) + with pytest.raises(litellm.ContextWindowExceededError, match="Context Window exceeded for given call"): + await core.call("atext_completion", {"model": "g", "prompt": LONG_TEXT}) + assert host.calls_to("atext_completion") == [] + + +def test_sync_text_completion_counts_the_prompt_against_the_context_window( + make_core: MakeCore, host: EndpointHost +) -> None: + core: Final = make_core(model_list=[windowed("g", "small", 5)], enable_pre_call_checks=True, num_retries=0) + with pytest.raises(litellm.ContextWindowExceededError, match="Context Window exceeded for given call"): + core.text_completion(model="g", prompt=LONG_TEXT) # pyright: ignore[reportAttributeAccessIssue] + assert host.calls_to("text_completion") == [] + + +@pytest.mark.parametrize( + ("endpoint", "request_fields"), + [ + ("aimage_generation", {"prompt": "a cat"}), + ("image_generation", {"prompt": "a cat"}), + ("atranscription", {"file": io.BytesIO(b"x")}), + ("aspeech", {"input": "hi", "voice": "alloy"}), + ("aadapter_completion", {"adapter_id": "a", "messages": MESSAGES}), + ("acreate_batch", BATCH), + ], +) +@pytest.mark.asyncio +async def test_endpoints_picking_with_a_stand_in_message_run_the_region_check( + make_core: MakeCore, host: EndpointHost, endpoint: str, request_fields: Mapping[str, object] +) -> None: + core: Final = make_core(model_list=EU_ONLY, enable_pre_call_checks=True, num_retries=0) + with pytest.raises(RouterRateLimitError): + await core.call(endpoint, {"model": "g", "allowed_model_region": "ap", **request_fields}) + assert host.calls == [] + + +@pytest.mark.parametrize( + ("endpoint", "request_fields"), + [("arerank", {"query": "q", "documents": ["a"]}), ("amoderation", {"input": "text"})], +) +@pytest.mark.asyncio +async def test_endpoints_picking_with_no_input_skip_the_region_check( + make_core: MakeCore, host: EndpointHost, endpoint: str, request_fields: Mapping[str, object] +) -> None: + core: Final = make_core(model_list=EU_ONLY, enable_pre_call_checks=True, num_retries=0) + await core.call(endpoint, {"model": "g", "allowed_model_region": "ap", **request_fields}) + assert host.deployments_called(endpoint) == ["eu"] + + +@pytest.mark.asyncio +async def test_speech_counts_its_stand_in_message_and_not_its_input(make_core: MakeCore, host: EndpointHost) -> None: + core: Final = make_core(model_list=[windowed("g", "small", 50)], enable_pre_call_checks=True, num_retries=0) + await core.call("aspeech", {"model": "g", "input": LONG_TEXT, "voice": "alloy"}) + assert host.deployments_called("aspeech") == ["small"] + + +# --- U4: print_deployment -------------------------------------------------------------------- + + +def test_print_deployment_masks_the_key_of_a_deployment_and_returns_a_deployment(make_core: MakeCore) -> None: + given: Final = Deployment(model_name="g", litellm_params=LiteLLM_Params(model="openai/x", api_key="sk-obj-123456")) + core: Final = make_core(model_list=[deployment("g", "a")]) + printed: Final = core.print_deployment(given) # pyright: ignore[reportAttributeAccessIssue] + assert isinstance(printed, Deployment) + assert printed.litellm_params.api_key == "sk**********" + assert given.litellm_params.api_key == "sk-obj-123456" + + +def test_print_deployment_refuses_a_list_of_pairs(make_core: MakeCore) -> None: + core: Final = make_core(model_list=[deployment("g", "a")]) + with pytest.raises(TypeError): + core.print_deployment([["litellm_params", {"api_key": "sk-pairs-1234"}]]) # pyright: ignore[reportAttributeAccessIssue] + + +# --- N10: context compaction ----------------------------------------------------------------- + + +def _complexity_router_over_a_small_window() -> list[dict[str, object]]: # mutable-ok: a model_list + return [ + windowed("mini", "mini", 60), + deployment("o47", "o47", model="anthropic/claude-opus-4-7"), + { + "model_name": "smart", + "litellm_params": { + "model": "auto_router/complexity_router", + "complexity_router_config": { + "tiers": {"SIMPLE": "mini", "REASONING": "o47"}, + "enable_context_window_escalation": True, + }, + "complexity_router_default_model": "mini", + }, + "model_info": {"id": "smart"}, + }, + ] + + +@pytest.mark.asyncio +async def test_an_auto_router_refuses_a_prompt_compaction_cannot_fit(make_core: MakeCore, host: EndpointHost) -> None: + core: Final = make_core(model_list=_complexity_router_over_a_small_window(), num_retries=0) + messages: Final = [ + {"role": "user", "content": "context line about the project history and many details " * 30}, + {"role": "assistant", "content": "earlier answer"}, + {"role": "user", "content": "and next?"}, + ] + with pytest.raises(litellm.ContextWindowExceededError, match="Context compaction"): + await core.call("acompletion", {"model": "smart", "messages": messages}) + assert host.calls_to("acompletion") == [] + + +@pytest.mark.asyncio +async def test_the_async_chat_attempt_carries_the_compaction_state(make_core: MakeCore, host: EndpointHost) -> None: + core: Final = make_core(model_list=[deployment("g", "a")], num_retries=0) + await core.call("acompletion", {"model": "g", "messages": MESSAGES}) + state: Final = host.calls_to("acompletion")[0].kwargs["_context_compaction_state"] + assert isinstance(state, CompactionState) + assert state.surface == "chat" + assert state.config is None + + +@pytest.mark.parametrize( + ("endpoint", "request_fields"), + [("completion", {"messages": MESSAGES}), ("atext_completion", {"prompt": "hi"})], +) +@pytest.mark.asyncio +async def test_other_attempts_carry_no_compaction_state( + make_core: MakeCore, host: EndpointHost, endpoint: str, request_fields: Mapping[str, object] +) -> None: + core: Final = make_core(model_list=[deployment("g", "a")], num_retries=0) + await core.call(endpoint, {"model": "g", **request_fields}) + assert "_context_compaction_state" not in host.calls_to(endpoint)[0].kwargs + + +def test_the_routing_context_counts_messages_and_tools(make_core: MakeCore) -> None: + ctx: Final = RoutingContext(make_core(model_list=[deployment("g", "a")])) # pyright: ignore[reportArgumentType] + tool: Final = {"type": "function", "function": {"name": "lookup", "description": "look it up", "parameters": {}}} + plain: Final = ctx.count_pre_call_check_tokens(messages=MESSAGES, input=None) + with_tools: Final = ctx.count_pre_call_check_tokens(messages=MESSAGES, input=None, request_kwargs={"tools": [tool]}) + assert 0 < plain < with_tools + + +def test_the_routing_context_counts_responses_input(make_core: MakeCore) -> None: + ctx: Final = RoutingContext(make_core(model_list=[deployment("g", "a")])) # pyright: ignore[reportArgumentType] + assert ctx.count_pre_call_check_tokens(messages=None, input="hello there") > 0 + + +def test_the_routing_context_needs_messages_or_input_to_count(make_core: MakeCore) -> None: + ctx: Final = RoutingContext(make_core(model_list=[deployment("g", "a")])) # pyright: ignore[reportArgumentType] + with pytest.raises(ValueError, match="Either messages or input must be provided to count tokens"): + ctx.count_pre_call_check_tokens(messages=None, input=None) + + +def test_the_routing_context_pre_call_checks_refuse_a_count_over_the_window(make_core: MakeCore) -> None: + small: Final = windowed("g", "small", 5) + ctx: Final = RoutingContext(make_core(model_list=[small])) # pyright: ignore[reportArgumentType] + with pytest.raises(litellm.ContextWindowExceededError, match="Context Window exceeded for given call"): + ctx.pre_call_checks( + model="g", + healthy_deployments=[small], + messages=MESSAGES, # pyright: ignore[reportArgumentType] + request_kwargs={}, + input_token_count=100, + skip_inline_token_count=True, + ) + + +def test_the_routing_context_pre_call_checks_drop_a_deployment_outside_the_region(make_core: MakeCore) -> None: + ctx: Final = RoutingContext(make_core(model_list=REGIONAL)) # pyright: ignore[reportArgumentType] + kept: Final = ctx.pre_call_checks( + model="g", + healthy_deployments=list(REGIONAL), + messages=MESSAGES, # pyright: ignore[reportArgumentType] + request_kwargs={"allowed_model_region": "ap"}, + input_token_count=1, + skip_inline_token_count=True, + ) + assert [record["model_info"]["id"] for record in kept] == ["ap"] # pyright: ignore[reportIndexIssue, reportCallIssue, reportArgumentType] diff --git a/tests/unit/test_router/core/test_gaps_request_overrides.py b/tests/unit/test_router/core/test_gaps_request_overrides.py new file mode 100644 index 00000000000..3793dd5f898 --- /dev/null +++ b/tests/unit/test_router/core/test_gaps_request_overrides.py @@ -0,0 +1,227 @@ +""" +Per-request routing overrides a router core reads (or does not read) on each call: + +- ``default_fallbacks`` and ``retry_policy`` in a request are no routing input: the router neither + falls back on the request's ``default_fallbacks`` nor retries on its ``retry_policy``. Both reach + the deployment call as given (litellm's own ``completion`` reads ``retry_policy`` there), on every + attempt, sync and async, streamed or not. +- ``model_group_retry_policy`` that is not a mapping of policies, in a request or as the router's + setting: the retry loop looks the model group up in it (``group in value``, then ``value[group]``), + so a value that does not hold the group leaves the router's ``retry_policy`` in charge, a ``None`` + entry is no policy at all, and an explicit request ``None`` replaces the router's per-group policies. +""" + +from collections.abc import Callable, Mapping +from typing import Final + +import pytest + +import litellm +from litellm.router_api import ChangeSet +from litellm.router_core import RouterCore +from litellm.types.router import RetryPolicy +from tests.unit.test_router.core.retries_health_helpers import attempt, attempt_sync +from tests.unit.test_router.core.scripted_host import ScriptedHost, deployment, rate_limit +from tests.unit.test_router.core.streaming_helpers import chat_stream, collect_texts, collect_texts_sync + +MakeCore = Callable[..., RouterCore] +MESSAGES: Final = [{"role": "user", "content": "hi"}] +GROUPS: Final = [deployment("a", "a"), deployment("b", "b"), deployment("c", "c")] +OVERRIDES: Final[Mapping[str, object]] = {"default_fallbacks": ["c"], "retry_policy": {"RateLimitErrorRetries": 3}} + + +def sent_overrides(host: ScriptedHost) -> list[tuple[str | None, object, object]]: # mutable-ok: a fresh list + return [ + (call.deployment_id, call.kwargs.get("default_fallbacks"), call.kwargs.get("retry_policy")) + for call in host.calls_to() + ] + + +# --- default_fallbacks and retry_policy: passed to the deployment, never routed on ---------------- + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("endpoint", "request_body"), + [ + ("acompletion", {"messages": MESSAGES}), + ("atext_completion", {"prompt": "p"}), + ("aembedding", {"input": ["x"]}), + ("aresponses", {"input": "hi"}), + ], +) +async def test_request_default_fallbacks_and_retry_policy_reach_the_deployment_as_given( + make_core: MakeCore, host: ScriptedHost, endpoint: str, request_body: Mapping[str, object] +) -> None: + core: Final = make_core(model_list=GROUPS, num_retries=0) + + await core.call(endpoint, {"model": "a", **request_body, **OVERRIDES}) + + assert sent_overrides(host) == [("a", ["c"], {"RateLimitErrorRetries": 3})] + + +@pytest.mark.asyncio +async def test_a_failed_call_neither_falls_back_on_nor_retries_by_the_request_overrides( + make_core: MakeCore, host: ScriptedHost +) -> None: + core: Final = make_core(model_list=GROUPS, num_retries=0) + host.answer("openai/a", rate_limit()) + + error: Final = await attempt(core, "a", **OVERRIDES) + + assert isinstance(error, litellm.RateLimitError) + assert sent_overrides(host) == [("a", ["c"], {"RateLimitErrorRetries": 3})] + + +def test_a_failed_sync_call_neither_falls_back_on_nor_retries_by_the_request_overrides( + make_core: MakeCore, host: ScriptedHost +) -> None: + core: Final = make_core(model_list=GROUPS, num_retries=0) + host.answer("openai/a", rate_limit()) + + error: Final = attempt_sync(core, "a", **OVERRIDES) + + assert isinstance(error, litellm.RateLimitError) + assert sent_overrides(host) == [("a", ["c"], {"RateLimitErrorRetries": 3})] + + +@pytest.mark.asyncio +async def test_the_router_retries_and_fallbacks_carry_the_request_overrides_on_every_attempt( + make_core: MakeCore, host: ScriptedHost +) -> None: + core: Final = make_core(model_list=GROUPS, num_retries=1, fallbacks=[{"a": ["b"]}]) + host.answer("openai/a", rate_limit(), rate_limit()) + + response: Final = await attempt(core, "a", **OVERRIDES) + + assert not isinstance(response, Exception) + assert sent_overrides(host) == [("a", ["c"], {"RateLimitErrorRetries": 3})] * 2 + [ + ("b", ["c"], {"RateLimitErrorRetries": 3}) + ] + + +@pytest.mark.asyncio +async def test_a_stream_that_fails_before_content_falls_back_with_the_request_overrides( + make_core: MakeCore, host: ScriptedHost +) -> None: + core: Final = make_core(model_list=GROUPS, num_retries=0, fallbacks=[{"a": ["b"]}]) + host.answer("openai/a", chat_stream("x", error=rate_limit(), fail_after=0)) + host.answer("openai/b", chat_stream("B")) + + stream: Final = await core.call("acompletion", {"model": "a", "messages": MESSAGES, "stream": True, **OVERRIDES}) + + assert await collect_texts(stream) == "B" + assert sent_overrides(host) == [ + ("a", ["c"], {"RateLimitErrorRetries": 3}), + ("b", ["c"], {"RateLimitErrorRetries": 3}), + ] + + +def test_a_sync_stream_carries_the_request_overrides(make_core: MakeCore, host: ScriptedHost) -> None: + core: Final = make_core(model_list=GROUPS, num_retries=0) + host.answer("openai/a", chat_stream("A", sync=True)) + + completion: Final[Callable[..., object]] = getattr(core, "completion") # pyright: ignore[reportAny] # an endpoint method + stream: Final = completion(model="a", messages=MESSAGES, stream=True, **OVERRIDES) + + assert collect_texts_sync(stream) == "A" + assert sent_overrides(host) == [("a", ["c"], {"RateLimitErrorRetries": 3})] + + +# --- model_group_retry_policy that is not a mapping of policies -------------------------------------- + + +def calls_after_a_rate_limit( + core: RouterCore, host: ScriptedHost, **request: object +) -> int: # kwargs-ok: request kwargs + host.answer("openai/a", rate_limit()) + error: Final = attempt_sync(core, "a", **request) + assert isinstance(error, litellm.RateLimitError) + return len(host.calls_to()) + + +@pytest.mark.parametrize( + "policies", + [None, {"a": None}, "xyz", [], ("b",)], + ids=["none", "none-entry", "string-without-the-group", "empty-list", "tuple-without-the-group"], +) +def test_a_request_policy_that_holds_no_policy_for_the_group_keeps_the_plain_retries( + make_core: MakeCore, host: ScriptedHost, policies: object +) -> None: + core: Final = make_core(model_list=GROUPS, num_retries=2) + + assert calls_after_a_rate_limit(core, host, model_group_retry_policy=policies) == 3 + + +@pytest.mark.parametrize( + ("policies", "calls"), + [ + ("xyz", 2), # the group is not in it: the router's retry_policy applies + ({"a": None}, 3), # a None entry is no policy: the plain retries apply + (None, 2), # an explicit None replaces the router's per-group policies + ], + ids=["string-without-the-group", "none-entry", "none"], +) +def test_a_request_policy_without_the_group_leaves_the_router_retry_policy_in_charge( + make_core: MakeCore, host: ScriptedHost, policies: object, calls: int +) -> None: + core: Final = make_core( + model_list=GROUPS, + num_retries=2, + retry_policy=RetryPolicy(RateLimitErrorRetries=1), + model_group_retry_policy={"a": RetryPolicy(RateLimitErrorRetries=0)}, + ) + + assert calls_after_a_rate_limit(core, host, model_group_retry_policy=policies) == calls + + +@pytest.mark.parametrize("policies", [5, "abc", {"a": 3}], ids=["int", "string-holding-the-group", "int-entry"]) +def test_a_request_that_opts_out_of_retries_never_reads_its_policy( + make_core: MakeCore, host: ScriptedHost, policies: object +) -> None: + core: Final = make_core(model_list=GROUPS, num_retries=2) + + assert calls_after_a_rate_limit(core, host, model_group_retry_policy=policies, num_retries=0) == 1 + + +@pytest.mark.parametrize( + "policies", [{"a": None}, "xyz", []], ids=["none-entry", "string-without-the-group", "empty-list"] +) +@pytest.mark.parametrize("how", ["constructor", "apply"]) +def test_a_router_policy_that_holds_no_policy_for_the_group_keeps_the_plain_retries( + make_core: MakeCore, host: ScriptedHost, policies: object, how: str +) -> None: + if how == "constructor": + core = make_core(model_list=GROUPS, num_retries=2, model_group_retry_policy=policies) + else: + core = make_core(model_list=GROUPS, num_retries=2) + core.apply(ChangeSet(settings={"model_group_retry_policy": policies})) + + stored: Final[object] = getattr(core, "model_group_retry_policy") # pyright: ignore[reportAny] # the router attribute + assert stored == policies + assert calls_after_a_rate_limit(core, host) == 3 + + +@pytest.mark.parametrize("how", ["constructor", "apply"]) +def test_a_router_policy_that_is_no_mapping_is_not_read_when_the_request_opts_out( + make_core: MakeCore, host: ScriptedHost, how: str +) -> None: + if how == "constructor": + core = make_core(model_list=GROUPS, num_retries=2, model_group_retry_policy=5) + else: + core = make_core(model_list=GROUPS, num_retries=2) + core.apply(ChangeSet(settings={"model_group_retry_policy": 5})) + + assert calls_after_a_rate_limit(core, host, num_retries=0) == 1 + + +def test_get_num_retries_from_retry_policy_reads_a_router_policy_that_is_no_mapping(make_core: MakeCore) -> None: + core: Final = make_core( + model_list=GROUPS, retry_policy=RetryPolicy(RateLimitErrorRetries=4), model_group_retry_policy="xyz" + ) + error: Final = rate_limit() + retries_for: Final[Callable[..., object]] = getattr(core, "get_num_retries_from_retry_policy") # pyright: ignore[reportAny] # the router member + + assert retries_for(error, model_group="a") == 4 + with pytest.raises(TypeError, match="string indices must be integers"): + retries_for(error, model_group="x") diff --git a/tests/unit/test_router/core/test_gaps_router_core_crate.py b/tests/unit/test_router/core/test_gaps_router_core_crate.py new file mode 100644 index 00000000000..7f00d26c944 --- /dev/null +++ b/tests/unit/test_router/core/test_gaps_router_core_crate.py @@ -0,0 +1,283 @@ +""" +Gaps the Rust core closes inside router-core, pinned to ``PythonCore``'s behaviour (oddities included): +``default_deployment``, a ``model_group_retry_policy`` whose lookup raises, the deprecated registry +index writes, a raw ``routing_groups`` write against the stale member map, and ``cooldown_time=0`` +assigned after construction. +""" + +import warnings +from typing import Final + +import pytest +from pydantic import ValidationError + +import litellm +from litellm.litellm_core_utils.core_helpers import as_str_mapping +from litellm.router import Router +from litellm.router_api import ChangeSet +from litellm.router_core import RouterCore +from litellm.types.router import RoutingGroup +from tests.unit.test_router.core.registry_helpers import team_deployment +from tests.unit.test_router.core.retries_health_helpers import MakeCore, attempt, cooled_down, headers_of +from tests.unit.test_router.core.scripted_host import ScriptedHost, completion_response, deployment, rate_limit + +DEFAULT: Final[dict[str, object]] = { # mutable-ok: a deployment dict, copied per test + "model_name": "dm", + "litellm_params": {"model": "openai/other-model"}, + "model_info": {"id": "dd"}, +} + + +def as_router(core: RouterCore) -> Router: + router: Final = Router.__new__(Router) + router._attach(core) # pyright: ignore[reportPrivateUsage, reportAttributeAccessIssue, reportUnknownMemberType] # the facade's own attach, as from_config calls it + return router + + +def write(router: Router, name: str, value: object) -> None: + with warnings.catch_warnings(): + warnings.simplefilter("ignore", DeprecationWarning) + setattr(router, name, value) + + +async def picked(router: Router, model: str) -> str: + """The id ``select`` picks, or the error's class and message.""" + try: + chosen: Final[object] = await router.select(model, {}) + except Exception as error: + return f"{type(error).__name__}: {error}" + return str((as_str_mapping((as_str_mapping(chosen) or {}).get("model_info")) or {}).get("id")) + + +# --- default_deployment -------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_an_unknown_model_routes_to_the_default_deployment_without_retries_or_cooldown( + make_core: MakeCore, host: ScriptedHost +) -> None: + """Oddity: the default deployment skips cooldown and the strategy, and the retry count reads no healthy + deployment, so ``num_retries`` is unused; its id is never cooled while it names no configured deployment.""" + core: Final = make_core(model_list=[deployment("g", "a")], num_retries=1, allowed_fails=0) + _ = core.apply(ChangeSet(assign={"default_deployment": dict(DEFAULT)})) + host.answer("other-model", rate_limit(), completion_response()) + + failed: Final = await attempt(core, model="other-model") + assert isinstance(failed, litellm.RateLimitError) + assert cooled_down(core, ["dd", "a"]) == {} + served: Final = await attempt(core, model="other-model") + + assert headers_of(served)["x-litellm-model-group"] == "other-model" + assert [call.deployment_id for call in host.calls] == ["dd", "dd"] + assert host.calls[0].kwargs["model"] == "other-model" + assert host.sleeps == [] + assert core.settings().default_deployment == DEFAULT + + +@pytest.mark.asyncio +async def test_a_known_model_never_reaches_the_default_deployment(make_core: MakeCore, host: ScriptedHost) -> None: + core: Final = make_core(model_list=[deployment("g", "a")]) + _ = core.apply(ChangeSet(assign={"default_deployment": dict(DEFAULT)})) + + _ = await attempt(core) + + assert host.deployments_called("acompletion") == ["a"] + + +# --- model_group_retry_policy lookups that raise ----------------------------------------------------- + + +RAISING_POLICIES: Final = [ + (3, TypeError, "argument of type 'int' is not iterable"), + ("g", TypeError, "string indices must be integers"), + (["g"], TypeError, "list indices must be integers or slices"), + ({"g": 3}, AttributeError, "'int' object has no attribute 'DefaultRetries'"), + ({"g": {"DefaultRetries": "x"}}, ValidationError, "validation error for RetryPolicy"), +] +"""Each policy, the error its lookup raises for group ``g``, and a fragment of that error's message.""" + + +@pytest.mark.asyncio +@pytest.mark.parametrize(("policy", "raised", "message"), RAISING_POLICIES) +async def test_a_raising_retry_policy_lookup_replaces_the_error_and_stops_the_hop( + make_core: MakeCore, host: ScriptedHost, policy: object, raised: type[Exception], message: str +) -> None: + """Oddity: the lookup runs on the hop's first failure; its error replaces the attempt's, and no retry + follows. The error is raised while the attempt's is handled, so it chains to it, except pydantic's + ``ValidationError``, which pydantic raises with no context.""" + core: Final = make_core( + model_list=[deployment("g", "a"), deployment("g", "b")], num_retries=2, model_group_retry_policy=policy + ) + host.answer("*", rate_limit()) + + error: Final = await attempt(core) + + assert isinstance(error, raised) + assert message in str(error) + if raised is ValidationError: + assert error.__context__ is None + else: + assert isinstance(error.__context__, litellm.RateLimitError) + assert len(host.calls) == 1 + assert host.sleeps == [] + + +@pytest.mark.asyncio +@pytest.mark.parametrize(("policy", "raised", "message"), RAISING_POLICIES) +async def test_a_raising_retry_policy_lookup_takes_the_generic_fallbacks( + make_core: MakeCore, host: ScriptedHost, policy: object, raised: type[Exception], message: str +) -> None: + del raised, message + core: Final = make_core( + model_list=[deployment("g", "a"), deployment("g", "b"), deployment("h", "c")], + num_retries=2, + model_group_retry_policy=policy, + fallbacks=[{"g": ["h"]}], + ) + host.answer("openai/a", rate_limit()) + host.answer("openai/b", rate_limit()) + + served: Final = await attempt(core) + + assert headers_of(served)["x-litellm-attempted-fallbacks"] == 1 + assert host.deployments_called("acompletion")[-1] == "c" + assert len(host.calls) == 2 + + +@pytest.mark.asyncio +async def test_a_policy_naming_another_group_leaves_this_group_retrying( + make_core: MakeCore, host: ScriptedHost +) -> None: + core: Final = make_core( + model_list=[deployment("g", "a"), deployment("g", "b")], num_retries=2, model_group_retry_policy={"h": 3} + ) + host.answer("*", rate_limit()) + + error: Final = await attempt(core) + + assert isinstance(error, litellm.RateLimitError) + assert len(host.calls) == 2 + + +# --- deprecated registry index writes ---------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_a_written_model_name_index_decides_which_names_resolve(make_core: MakeCore) -> None: + router: Final = as_router(make_core(model_list=[deployment("gpt", "a"), deployment("claude", "c")])) + write(router, "model_name_to_deployment_indices", {"gpt": [0], "claude": []}) + + assert await picked(router, "gpt") == "a" + assert (await picked(router, "claude")).startswith("BadRequestError") + + +@pytest.mark.asyncio +async def test_a_model_name_index_entry_pointing_at_another_name_resolves_nothing(make_core: MakeCore) -> None: + """Each indexed deployment is still checked against the requested name.""" + router: Final = as_router(make_core(model_list=[deployment("gpt", "a"), deployment("claude", "c")])) + write(router, "model_name_to_deployment_indices", {"gpt": [1], "claude": [1]}) + + assert (await picked(router, "gpt")).startswith("BadRequestError") + assert await picked(router, "claude") == "c" + + +@pytest.mark.asyncio +async def test_a_written_model_names_set_decides_which_names_reach_the_default_deployment( + make_core: MakeCore, +) -> None: + core: Final = make_core(model_list=[deployment("gpt", "a"), deployment("claude", "c")]) + _ = core.apply(ChangeSet(assign={"default_deployment": dict(DEFAULT)})) + router: Final = as_router(core) + + write(router, "model_names", set[str]()) + assert [await picked(router, name) for name in ("gpt", "claude", "other-model")] == ["dd", "dd", "dd"] + + write(router, "model_names", {"gpt", "claude", "other-model"}) + assert await picked(router, "gpt") == "a" + assert (await picked(router, "other-model")).startswith("BadRequestError") + + +@pytest.mark.asyncio +async def test_an_emptied_team_index_still_finds_team_deployments_by_scanning(make_core: MakeCore) -> None: + router: Final = as_router( + make_core(model_list=[deployment("gpt", "a"), team_deployment("int-t", "t1", "TA", "pub")]) + ) + write(router, "team_model_to_deployment_indices", {}) + write(router, "model_name_to_deployment_indices", {"gpt": [0]}) + + chosen: Final = await router.select("pub", {"metadata": {"user_api_key_team_id": "TA"}}) + assert chosen["model_info"]["id"] == "t1" + assert (await picked(router, "int-t")).startswith("BadRequestError") + + +# --- a raw routing_groups write: the member map stays as constructed ------------------------------------ + + +def _grouped(make_core: MakeCore) -> Router: + return as_router( + make_core( + model_list=[deployment("gpt", "a"), deployment("gpt", "b"), deployment("claude", "c")], + routing_groups=[{"group_name": "G", "models": ["gpt"], "routing_strategy": "least-busy"}], + ) + ) + + +@pytest.mark.asyncio +async def test_dropping_a_routing_group_raw_leaves_its_members_raising_a_key_error(make_core: MakeCore) -> None: + """Oddity: the member-to-group map is built once, so a member of a dropped group looks it up and fails.""" + router: Final = _grouped(make_core) + write(router, "routing_groups", {}) + + assert await picked(router, "gpt") == "KeyError: 'G'" + assert await picked(router, "claude") == "c" + + +@pytest.mark.asyncio +async def test_a_routing_group_added_raw_leaves_its_members_on_the_default_strategy(make_core: MakeCore) -> None: + router: Final = _grouped(make_core) + write( + router, + "routing_groups", + { + "G": RoutingGroup(group_name="G", models=["gpt"], routing_strategy="least-busy"), + "K": RoutingGroup(group_name="K", models=["claude"], routing_strategy="latency-based-routing"), + }, + ) + + assert await picked(router, "gpt") == "a" + assert await picked(router, "claude") == "c" + assert (await picked(router, "K")).startswith("RouterRateLimitError: No deployments available") + + +# --- cooldown_time assigned after construction ---------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_an_assigned_zero_cooldown_time_never_keeps_a_deployment_cooled( + make_core: MakeCore, host: ScriptedHost +) -> None: + """Oddity: an assigned 0 is a zero cooldown ttl (only the constructor reads 0 as the default).""" + core: Final = make_core(model_list=[deployment("g", "a"), deployment("g", "b")], num_retries=0, allowed_fails=0) + _ = core.apply(ChangeSet(assign={"cooldown_time": 0})) + host.answer("*", rate_limit()) + + for _ in range(3): + assert isinstance(await attempt(core), litellm.RateLimitError) + assert cooled_down(core, ["a", "b"]) == {} + + assert core.settings().cooldown_time == 0 + assert len(host.calls) == 3 + + +@pytest.mark.asyncio +async def test_an_assigned_none_cooldown_time_cools_for_the_default(make_core: MakeCore, host: ScriptedHost) -> None: + core: Final = make_core(model_list=[deployment("g", "a"), deployment("g", "b")], num_retries=0, allowed_fails=0) + _ = core.apply(ChangeSet(assign={"cooldown_time": None})) + host.answer("*", rate_limit()) + + _ = await attempt(core) + _ = await attempt(core) + + assert sorted(cooled_down(core, ["a", "b"])) == ["a", "b"] + assert {entry["cooldown_time"] for entry in cooled_down(core, ["a", "b"]).values()} == {5} + assert isinstance(await attempt(core), litellm.RouterRateLimitError) diff --git a/tests/unit/test_router/core/test_gaps_router_core_crate_picks.py b/tests/unit/test_router/core/test_gaps_router_core_crate_picks.py new file mode 100644 index 00000000000..255d731142d --- /dev/null +++ b/tests/unit/test_router/core/test_gaps_router_core_crate_picks.py @@ -0,0 +1,372 @@ +""" +Pick errors the Rust core serves inside router-core, pinned to ``PythonCore``'s behaviour (oddities +included): a pick's error retried by the retry loop, a ``default_deployment`` Python cannot serve, a +written index position past the model list, a stale routing group at a deeper fallback level or in +weighted failover, and a strategy change through ``update_settings`` after a raw ``routing_groups`` write. +""" + +import warnings +from typing import Final + +import pytest +from pydantic import ValidationError + +import litellm +from litellm.litellm_core_utils.core_helpers import as_str_mapping +from litellm.router import Router +from litellm.router_api import ChangeSet +from litellm.router_core import RouterCore +from litellm.types.router import RoutingGroup +from tests.unit.test_router.core.registry_helpers import team_deployment +from tests.unit.test_router.core.retries_health_helpers import MakeCore, attempt, attempt_sync +from tests.unit.test_router.core.scripted_host import ScriptedHost, deployment, rate_limit +from tests.unit.test_router.core.test_fallbacks_characterization_outcomes import FallbackEvents + +TEAM: Final = {"metadata": {"user_api_key_team_id": "TA"}} + + +def as_router(core: RouterCore) -> Router: + router: Final = Router.__new__(Router) + router._attach(core) # pyright: ignore[reportPrivateUsage, reportAttributeAccessIssue, reportUnknownMemberType] # the facade's own attach, as from_config calls it + return router + + +def write(router: Router, name: str, value: object) -> None: + with warnings.catch_warnings(): + warnings.simplefilter("ignore", DeprecationWarning) + setattr(router, name, value) + + +async def selected(router: Router, model: str, **request: object) -> object: # kwargs-ok: the request kwargs + """The deployment ``select`` picks, or the error it raised.""" + try: + return await router.select(model, dict(request)) + except Exception as error: + return error + + +def picked_id(picked: object) -> object: + return (as_str_mapping((as_str_mapping(picked) or {}).get("model_info")) or {}).get("id") + + +# --- a pick's error in the retry loop ----------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_a_retry_policy_retries_a_group_whose_deployments_all_cool_down( + make_core: MakeCore, host: ScriptedHost +) -> None: + """The pick's ``RouterRateLimitError`` is a plain ``ValueError``: ``DefaultRetries`` applies to it.""" + core: Final = make_core( + model_list=[deployment("g", "a"), deployment("g", "b")], + num_retries=2, + allowed_fails=0, + model_group_retry_policy={"g": {"DefaultRetries": 1}}, + ) + host.answer("*", rate_limit()) + for _ in range(2): + _ = await attempt(core, num_retries=0) + host.calls.clear() + host.sleeps.clear() + + error: Final = await attempt(core) + + assert isinstance(error, litellm.RouterRateLimitError) + assert isinstance(error.__context__, litellm.RouterRateLimitError) + assert host.calls == [] + assert len(host.sleeps) == 2 + assert all(sleep > 0 for sleep in host.sleeps) + + +@pytest.mark.asyncio +async def test_a_retry_policy_retries_an_unknown_model_and_stamps_the_error( + make_core: MakeCore, host: ScriptedHost +) -> None: + core: Final = make_core(model_list=[deployment("g", "a")], num_retries=2, retry_policy={"DefaultRetries": 1}) + + error: Final = await attempt(core, model="zz") + + assert isinstance(error, litellm.BadRequestError) + assert (getattr(error, "num_retries", None), getattr(error, "max_retries", None)) == (1, 1) + assert isinstance(error.__context__, litellm.BadRequestError) + assert len(host.sleeps) == 2 + + +@pytest.mark.asyncio +async def test_invalid_router_weights_are_retried_while_a_deployment_is_healthy( + make_core: MakeCore, host: ScriptedHost +) -> None: + """No policy: ``should_retry_this_error`` retries a plain exception while a deployment is healthy.""" + core: Final = make_core(model_list=[deployment("g", "a"), deployment("g", "b")], num_retries=2) + + error: Final = await attempt(core, _router_weights={"a": "x"}) + + assert isinstance(error, ValidationError) + assert isinstance(error.__context__, ValidationError) + assert host.calls == [] + assert host.sleeps == [0, 0, 0] + assert all(type(sleep) is int for sleep in host.sleeps) + + +@pytest.mark.asyncio +async def test_without_a_policy_an_unknown_model_is_raised_at_once(make_core: MakeCore, host: ScriptedHost) -> None: + core: Final = make_core(model_list=[deployment("g", "a")], num_retries=2) + + error: Final = await attempt(core, model="zz") + + assert isinstance(error, litellm.BadRequestError) + assert error.__context__ is None + assert host.sleeps == [] + + +@pytest.mark.asyncio +async def test_the_last_retry_failure_chains_to_the_hops_first_failure(make_core: MakeCore, host: ScriptedHost) -> None: + core: Final = make_core( + model_list=[deployment("g", "a"), deployment("g", "b")], + num_retries=2, + model_group_retry_policy={"g": {"DefaultRetries": 1}}, + ) + host.answer("*", rate_limit()) + + error: Final = await attempt(core) + + assert isinstance(error, litellm.RateLimitError) + assert isinstance(error.__context__, litellm.RateLimitError) + assert error.__context__ is not error + assert host.sleeps[0] == 0 + assert type(host.sleeps[0]) is int + + +# --- a default_deployment Python cannot serve --------------------------------------------------------- + + +UNCOPYABLE_DEFAULTS: Final[list[tuple[object, type[Exception], str]]] = [ # mutable-ok: pytest parameters + ({}, KeyError, "'litellm_params'"), + ({"model_name": "dm"}, KeyError, "'litellm_params'"), + ({"model_name": "dm", "litellm_params": "s", "model_info": {"id": "dd"}}, AttributeError, "no attribute 'copy'"), + ("str", AttributeError, "'str' object has no attribute 'copy'"), + (5, AttributeError, "'int' object has no attribute 'copy'"), + (["a"], TypeError, "list indices must be integers or slices, not str"), +] +"""Each value, and the error its copy raises for any name.""" + + +@pytest.mark.asyncio +@pytest.mark.parametrize(("value", "raised", "message"), UNCOPYABLE_DEFAULTS) +async def test_a_default_deployment_python_cannot_copy_fails_each_pick_of_an_unknown_name( + make_core: MakeCore, host: ScriptedHost, value: object, raised: type[Exception], message: str +) -> None: + """Oddity: assigning it succeeds; every request for a name nothing else serves raises the copy's error.""" + core: Final = make_core(model_list=[deployment("g", "a")], num_retries=1) + _ = core.apply(ChangeSet(assign={"default_deployment": value})) + + for error in ( + await attempt(core, model="zz"), + attempt_sync(core, model="zz"), + await selected(as_router(core), "zz"), + ): + assert isinstance(error, raised) + assert message in str(error) + assert host.deployments_called("acompletion") == [] + _ = await attempt(core) + assert host.deployments_called("acompletion") == ["a"] + + +FAILING_DEFAULTS: Final[list[tuple[object, str, object]]] = [ # mutable-ok: pytest parameters + ({"litellm_params": {"model": "x"}}, "'model_name'", None), + ({"model_name": "dm", "litellm_params": {"model": "x"}}, "'model_info'", None), + ({"model_name": "dm", "litellm_params": {"model": "x"}, "model_info": {}}, "'id'", {}), +] +"""Each value, the ``KeyError`` every call to it raises, and the ``model_info`` its pick returns.""" + + +@pytest.mark.asyncio +@pytest.mark.parametrize(("value", "missing", "model_info"), FAILING_DEFAULTS) +async def test_a_default_deployment_missing_a_field_is_picked_but_every_call_fails( + make_core: MakeCore, host: ScriptedHost, value: object, missing: str, model_info: object +) -> None: + """Oddity: the pick returns Python's copy; the call's first reads of it raise ``KeyError``, with no retry.""" + core: Final = make_core(model_list=[deployment("g", "a")], num_retries=1) + _ = core.apply(ChangeSet(assign={"default_deployment": value})) + + called: Final = await attempt(core, model="zz") + called_sync: Final = attempt_sync(core, model="zz") + chosen: Final = as_str_mapping(await selected(as_router(core), "zz")) or {} + + assert isinstance(called, KeyError) + assert isinstance(called_sync, KeyError) + assert str(called) == missing + assert chosen.get("model_info") == model_info + assert (as_str_mapping(chosen.get("litellm_params")) or {}).get("model") == "zz" + assert host.calls == [] + assert host.sleeps == [] + + +# --- a written index position past the model list ----------------------------------------------------- + + +def _teams(make_core: MakeCore) -> tuple[RouterCore, Router]: + core: Final = make_core( + model_list=[deployment("gpt", "a"), deployment("claude", "c"), team_deployment("int-t", "t1", "TA", "pub")], + num_retries=1, + fallbacks=[{"gpt": ["claude"]}, {"pub": ["claude"]}], + ) + return core, as_router(core) + + +@pytest.mark.asyncio +async def test_a_name_index_position_past_the_list_raises_index_error_for_that_name( + make_core: MakeCore, host: ScriptedHost +) -> None: + core, router = _teams(make_core) + write(router, "model_name_to_deployment_indices", {"gpt": [0, 7], "claude": [1], "int-t": [2]}) + + picked: Final = await selected(router, "gpt") + called: Final = await attempt(core, model="gpt") + + assert isinstance(picked, IndexError) + assert isinstance(called, IndexError) + assert str(called) == "list index out of range" + assert picked_id(await selected(router, "claude")) == "c" + assert host.calls == [] + + +@pytest.mark.asyncio +async def test_a_team_index_position_past_the_list_raises_from_the_fallback_handling( + make_core: MakeCore, host: ScriptedHost +) -> None: + """Oddity: the level's fallback handling lists the group again and raises a second ``IndexError`` + (chained to the pick's), before any fallback.""" + core, router = _teams(make_core) + write(router, "team_model_to_deployment_indices", {("TA", "pub"): [9]}) + + picked: Final = await selected(router, "pub", **TEAM) + called: Final = await attempt(core, model="pub", **TEAM) + + assert isinstance(picked, IndexError) + assert isinstance(called, IndexError) + assert isinstance(called.__context__, IndexError) + assert host.calls == [] + assert picked_id(await selected(router, "gpt")) == "a" + + +# --- a stale routing group deeper in the walk --------------------------------------------------------- + + +def _stale(make_core: MakeCore, **arguments: object) -> RouterCore: # kwargs-ok: extra router arguments + core: Final = make_core( + model_list=[deployment("gpt", "a"), deployment("gpt", "b"), deployment("claude", "c"), deployment("x", "x1")], + routing_groups=[{"group_name": "G", "models": ["gpt"], "routing_strategy": "simple-shuffle"}], + num_retries=0, + fallbacks=[{"x": ["gpt"]}, {"gpt": ["claude"]}], + **arguments, + ) + write(as_router(core), "routing_groups", {}) + return core + + +@pytest.mark.asyncio +async def test_a_stale_routing_group_at_a_fallback_level_is_that_levels_error( + make_core: MakeCore, host: ScriptedHost, monkeypatch: pytest.MonkeyPatch +) -> None: + events: Final = FallbackEvents() + monkeypatch.setattr(litellm, "callbacks", [events]) + core: Final = _stale(make_core) + host.answer("*", rate_limit()) + + error: Final = await attempt(core, model="x") + + assert isinstance(error, litellm.RateLimitError) + assert events.events == [("failure", "gpt", "claude", "KeyError"), ("failure", "x", "gpt", "RateLimitError")] + + +@pytest.mark.asyncio +async def test_weighted_failover_raises_the_stale_groups_key_error_before_any_fallback( + make_core: MakeCore, host: ScriptedHost +) -> None: + """Oddity: ``_maybe_run_weighted_failover`` reads the routing context first, so the stale member map's + ``KeyError`` escapes the fallback handling (chained to the pick's) and the regular fallback never runs.""" + core: Final = _stale(make_core, enable_weighted_failover=True) + + error: Final = await attempt(core, model="gpt") + + assert isinstance(error, KeyError) + assert str(error) == "'G'" + assert isinstance(error.__context__, KeyError) + assert host.calls == [] + + +# --- update_settings(routing_strategy) after a raw routing_groups write ----------------------------------- + + +@pytest.mark.asyncio +async def test_a_strategy_change_re_inits_the_routing_groups_from_their_last_input(make_core: MakeCore) -> None: + core: Final = make_core( + model_list=[deployment("gpt", "a"), deployment("gpt", "b"), deployment("claude", "c")], + routing_groups=[{"group_name": "G", "models": ["gpt"], "routing_strategy": "least-busy"}], + ) + router: Final = as_router(core) + write( + router, "routing_groups", {"K": RoutingGroup(group_name="K", models=["claude"], routing_strategy="least-busy")} + ) + assert isinstance(await selected(router, "gpt"), KeyError) + + _ = core.apply(ChangeSet(settings={"routing_strategy": "usage-based-routing-v2"})) + + assert sorted(core.registry().routing_groups()) == ["G"] + assert picked_id(await selected(router, "G")) == "a" + assert picked_id(await selected(router, "gpt")) == "a" + assert isinstance(await selected(router, "K"), litellm.BadRequestError) + + +@pytest.mark.asyncio +async def test_an_unchanged_strategy_keeps_the_raw_routing_groups(make_core: MakeCore) -> None: + core: Final = make_core( + model_list=[deployment("gpt", "a"), deployment("claude", "c")], + routing_groups=[{"group_name": "G", "models": ["gpt"], "routing_strategy": "least-busy"}], + ) + write( + as_router(core), + "routing_groups", + {"K": RoutingGroup(group_name="K", models=["claude"], routing_strategy="least-busy")}, + ) + + _ = core.apply(ChangeSet(settings={"routing_strategy": "simple-shuffle"})) + + assert sorted(core.registry().routing_groups()) == ["K"] + + +@pytest.mark.asyncio +async def test_a_negative_index_position_counts_from_the_end_of_the_list( + make_core: MakeCore, host: ScriptedHost +) -> None: + """As Python's ``model_list[index]`` reads it; one before the start raises ``IndexError``.""" + core, router = _teams(make_core) + write(router, "model_name_to_deployment_indices", {"gpt": [-3], "claude": [-2], "int-t": [-4]}) + + assert picked_id(await selected(router, "gpt")) == "a" + assert picked_id(await selected(router, "claude")) == "c" + before_the_start: Final = await attempt(core, model="int-t") + assert isinstance(before_the_start, IndexError) + assert str(before_the_start) == "list index out of range" + assert host.calls == [] + + +@pytest.mark.asyncio +async def test_a_default_deployment_whose_id_is_not_a_string_is_picked_and_called_as_python_copies_it( + make_core: MakeCore, host: ScriptedHost +) -> None: + core: Final = make_core(model_list=[deployment("g", "a")]) + value: Final = { + "model_name": "dm", + "litellm_params": {"model": "openai/dm", "api_key": "k"}, + "model_info": {"id": 5}, + } + _ = core.apply(ChangeSet(assign={"default_deployment": value})) + + served: Final = await attempt(core, model="zz") + chosen: Final = await selected(as_router(core), "zz") + + assert isinstance(served, litellm.ModelResponse) + assert host.calls[0].kwargs["model"] == "zz" + assert chosen == {"model_name": "dm", "litellm_params": {"model": "zz", "api_key": "k"}, "model_info": {"id": 5}} diff --git a/tests/unit/test_router/core/test_health_registry_views.py b/tests/unit/test_router/core/test_health_registry_views.py new file mode 100644 index 00000000000..0180f991ccb --- /dev/null +++ b/tests/unit/test_router/core/test_health_registry_views.py @@ -0,0 +1,174 @@ +""" +The ``health()`` priority helpers and the ``registry()`` members that read router state beyond the +deployment list: ``client(client_type="max_parallel_requests")``, ``plugins("routing")`` and +``zero_cost()``. Each test pins what the Python core does; every core must answer the same. +""" + +from typing import Final + +import pytest + +import litellm +from litellm.router_api import ChangeSet +from litellm.router_utils.client_initalization_utils import MaxParallelRequestsLimit +from litellm.types.router import Deployment, LiteLLM_Params, RoutingContext +from tests.unit.test_router.core.scripted_host import deployment +from tests.unit.test_router.core.selection_helpers import MESSAGES, MakeCore + +# --- health().split_priority ---------------------------------------------------------------------- + + +def test_split_priority_gives_the_request_priority_and_the_kwargs_without_it(make_core: MakeCore) -> None: + core: Final = make_core(model_list=[deployment("g", "a")]) + assert core.health().split_priority("g", {"priority": 5, "user": "u"}) == (5, {"user": "u"}) + + +def test_split_priority_falls_back_to_the_default_priority_when_the_request_sets_none(make_core: MakeCore) -> None: + core: Final = make_core(model_list=[deployment("g", "a")], default_priority=2) + assert core.health().split_priority("g", {"user": "u"}) == (2, {"user": "u"}) + + +def test_split_priority_without_a_default_leaves_the_request_unqueued(make_core: MakeCore) -> None: + core: Final = make_core(model_list=[deployment("g", "a")]) + assert core.health().split_priority("g", {"priority": None}) == (None, {}) + + +def test_split_priority_refuses_a_priority_that_is_no_integer(make_core: MakeCore) -> None: + core: Final = make_core(model_list=[deployment("g", "a")], default_priority=2) + with pytest.raises(litellm.BadRequestError, match="priority must be an integer, got str 'high'"): + core.health().split_priority("g", {"priority": "high"}) + + +def test_split_priority_drops_a_bad_priority_for_a_group_that_drops_params(make_core: MakeCore) -> None: + """A deployment of the group with ``drop_params=True`` turns the bad priority into the default.""" + core: Final = make_core(model_list=[deployment("g", "a", drop_params=True)], default_priority=2) + assert core.health().split_priority("g", {"priority": "high"}) == (2, {}) + + +def test_split_priority_reads_a_pinned_deployment_s_params_for_its_id(make_core: MakeCore) -> None: + """A deployment id as the model reads that deployment's ``drop_params``, not a group's.""" + core: Final = make_core( + model_list=[deployment("g", "a", drop_params=True), deployment("g", "b")], default_priority=1 + ) + assert core.health().split_priority("a", {"priority": "high"}) == (1, {}) + with pytest.raises(litellm.BadRequestError): + core.health().split_priority("b", {"priority": "high"}) + + +# --- health().priority_header_fields -------------------------------------------------------------- + + +def test_priority_header_fields_read_an_integer_header(make_core: MakeCore) -> None: + core: Final = make_core(model_list=[deployment("g", "a")]) + assert core.health().priority_header_fields("g", {}, "3") == {"priority": 3} + + +@pytest.mark.parametrize("header", [None, ""]) +def test_priority_header_fields_are_empty_without_a_header(make_core: MakeCore, header: str | None) -> None: + core: Final = make_core(model_list=[deployment("g", "a")]) + assert core.health().priority_header_fields("g", {}, header) == {} + + +def test_priority_header_fields_pass_a_bad_header_on_for_the_call_to_refuse(make_core: MakeCore) -> None: + core: Final = make_core(model_list=[deployment("g", "a")]) + assert core.health().priority_header_fields("g", {}, "urgent") == {"priority": "urgent"} + + +def test_priority_header_fields_drop_a_bad_header_when_the_request_drops_params(make_core: MakeCore) -> None: + core: Final = make_core(model_list=[deployment("g", "a")]) + assert core.health().priority_header_fields("g", {"drop_params": True}, "urgent") == {} + + +# --- registry().client(client_type="max_parallel_requests") --------------------------------------- + + +def test_the_max_parallel_requests_client_is_the_deployment_s_cached_limit(make_core: MakeCore) -> None: + core: Final = make_core(model_list=[deployment("g", "a", max_parallel_requests=2)]) + entry: Final = deployment("g", "a", max_parallel_requests=2) + limit: Final = core.registry().client(entry, client_type="max_parallel_requests") + assert isinstance(limit, MaxParallelRequestsLimit) + assert (limit.max_parallel_requests, limit.model_id, limit.model_group, limit.in_flight) == (2, "a", "g", 0) + assert core.registry().client(entry, client_type="max_parallel_requests") is limit + + +def test_the_max_parallel_requests_limit_comes_from_rpm_or_the_router_default(make_core: MakeCore) -> None: + core: Final = make_core( + model_list=[deployment("g", "a", rpm=7), deployment("g", "b")], default_max_parallel_requests=4 + ) + by_rpm: Final = core.registry().client(deployment("g", "a", rpm=7), client_type="max_parallel_requests") + by_default: Final = core.registry().client(deployment("g", "b"), client_type="max_parallel_requests") + assert isinstance(by_rpm, MaxParallelRequestsLimit) + assert isinstance(by_default, MaxParallelRequestsLimit) + assert (by_rpm.max_parallel_requests, by_default.max_parallel_requests) == (7, 4) + + +def test_a_deployment_without_a_limit_has_no_max_parallel_requests_client(make_core: MakeCore) -> None: + core: Final = make_core(model_list=[deployment("g", "a")]) + assert core.registry().client(deployment("g", "a"), client_type="max_parallel_requests") is None + + +def test_the_client_lookup_reads_the_deployment_id_first(make_core: MakeCore) -> None: + """``deployment["model_info"]["id"]``: a deployment without model_info raises, whatever the client type.""" + core: Final = make_core(model_list=[deployment("g", "a")]) + entry: Final = {"model_name": "g", "litellm_params": {"model": "openai/a"}} + for client_type in ("max_parallel_requests", "async", None): + with pytest.raises(KeyError, match="model_info"): + core.registry().client(entry, client_type=client_type) + + +@pytest.mark.asyncio +async def test_a_full_max_parallel_requests_client_turns_the_next_call_away(make_core: MakeCore) -> None: + """The limit ``client`` returns is the one calls take their slot from.""" + core: Final = make_core(model_list=[deployment("g", "a", max_parallel_requests=1)], num_retries=0) + entry: Final = deployment("g", "a", max_parallel_requests=1) + limit: Final = core.registry().client(entry, client_type="max_parallel_requests") + assert isinstance(limit, MaxParallelRequestsLimit) + with limit, pytest.raises(litellm.RateLimitError, match="already has max_parallel_requests=1"): + await core.acompletion(model="g", messages=MESSAGES) # pyright: ignore[reportAttributeAccessIssue, reportUnknownMemberType] # every core serves acompletion + assert limit.in_flight == 0 + + +# --- registry().plugins("routing") and zero_cost() ------------------------------------------------ + + +class _Keep: + async def run(self, context: RoutingContext) -> RoutingContext: + return context + + +def test_the_routing_plugins_are_the_ones_the_router_was_built_with(make_core: MakeCore) -> None: + plugin: Final = _Keep() + core: Final = make_core(model_list=[deployment("g", "a")], plugins=[plugin]) + assert core.registry().plugins("routing") == [plugin] + + +def test_a_router_without_plugins_lists_no_routing_plugins(make_core: MakeCore) -> None: + assert make_core(model_list=[deployment("g", "a")]).registry().plugins("routing") == [] + + +def test_zero_cost_is_an_empty_lookup_the_auth_checks_fill(make_core: MakeCore) -> None: + core: Final = make_core(model_list=[deployment("g", "a")]) + lookup: Final = core.registry().zero_cost() + assert lookup == {} + assert isinstance(lookup, dict) + lookup["g"] = True + assert core.registry().zero_cost() == {"g": True} + + +def test_a_deployment_change_empties_the_zero_cost_lookup(make_core: MakeCore) -> None: + core: Final = make_core(model_list=[deployment("g", "a")]) + lookup: Final = core.registry().zero_cost() + assert isinstance(lookup, dict) + lookup["g"] = True + core.apply( + ChangeSet( + upsert=[ + Deployment( + model_name="g", + litellm_params=LiteLLM_Params(model="openai/b", api_key="sk-test"), + model_info={"id": "b"}, + ) + ] + ) + ) + assert core.registry().zero_cost() == {} diff --git a/tests/unit/test_router/core/test_host_settings_changes_characterization.py b/tests/unit/test_router/core/test_host_settings_changes_characterization.py new file mode 100644 index 00000000000..fa1446aa51a --- /dev/null +++ b/tests/unit/test_router/core/test_host_settings_changes_characterization.py @@ -0,0 +1,428 @@ +""" +Characterization of the settings only the host reads, changed after construction: ``search_tools``, +``guardrail_list``, ``default_max_parallel_requests``, ``provider_budget_config``, +``health_check_ignore_transient_errors`` and ``router_general_settings``. + +``apply(ChangeSet(assign=...))`` stores each as given, ``settings()`` reports it, and whatever reads +it afterwards sees it from then on. None of them is runtime-updatable, so ``ChangeSet(settings=...)`` +(``update_settings``) ignores them. +""" + +import asyncio +import time +from collections.abc import Awaitable, Callable, Mapping +from typing import Final + +import pytest + +import litellm +from litellm.litellm_core_utils.core_helpers import as_str_mapping +from litellm.router_api import ChangeSet +from litellm.router_core import RouterCore +from litellm.router_strategy.budget_limiter import RouterBudgetLimiting +from litellm.router_utils.kubernetes_pod_discovery import KubernetesPodDiscovery +from litellm.types.router import Deployment, RouterGeneralSettings +from tests.unit.test_router.core.endpoints_helpers import EndpointHost, patch_direct_endpoints +from tests.unit.test_router.core.retries_health_helpers import cooled_down +from tests.unit.test_router.core.scripted_host import ScriptedHost, deployment, rate_limit +from tests.unit.test_router.core.selection_helpers import ids_of +from tests.unit.test_router.core.streaming_helpers import chat_stream, collect_texts + +MakeCore = Callable[..., RouterCore] +MESSAGES: Final = [{"role": "user", "content": "hi"}] +SEARCH_TOOLS: Final = [{"search_tool_name": "s", "litellm_params": {"search_provider": "tavily", "api_key": "tk"}}] +GUARDRAILS: Final = [{"guardrail_name": "gr", "litellm_params": {"guardrail": "presidio", "mode": "pre_call"}}] +PROVIDER_BUDGET: Final = {"openai": {"budget_limit": 1, "time_period": "1d"}} + +ASSIGNED: Final[list[tuple[str, object]]] = [ # mutable-ok: pytest parameters + ("search_tools", SEARCH_TOOLS), + ("guardrail_list", GUARDRAILS), + ("default_max_parallel_requests", 3), + ("provider_budget_config", PROVIDER_BUDGET), + ("health_check_ignore_transient_errors", True), + ("router_general_settings", RouterGeneralSettings(pass_through_all_models=True)), + ("provider_default_deployment_ids", ["d1"]), + ("enable_health_check_routing", True), + ("background_health_check_model_groups", ["g"]), +] + + +@pytest.fixture +def host(monkeypatch: pytest.MonkeyPatch) -> EndpointHost: + endpoint_host: Final = EndpointHost() + patch_direct_endpoints(monkeypatch, endpoint_host) + return endpoint_host + + +def _endpoint(core: RouterCore, name: str) -> Callable[..., Awaitable[object]]: + """An async endpoint method, which the core serves by name.""" + method: Final[Callable[..., Awaitable[object]]] = getattr(core, name) # pyright: ignore[reportAny] # endpoint methods are served by name + return method + + +async def _open_stream(core: RouterCore) -> object: + return await _endpoint(core, "acompletion")(model="a", messages=MESSAGES, stream=True) + + +# --- settings() --------------------------------------------------------------------------------- + + +@pytest.mark.parametrize(("name", "value"), ASSIGNED) +def test_an_assigned_host_setting_is_reported_by_settings(make_core: MakeCore, name: str, value: object) -> None: + core: Final = make_core(model_list=[deployment("g", "d1")]) + before: Final = core.settings() + _ = core.apply(ChangeSet(assign={name: value})) + assert getattr(core.settings(), name) == value + assert core.settings() is not before + + +@pytest.mark.parametrize(("name", "value"), [("search_tools", SEARCH_TOOLS), ("guardrail_list", GUARDRAILS)]) +def test_an_assigned_host_setting_is_copied_from_the_callers_list( + make_core: MakeCore, name: str, value: list[dict[str, object]] +) -> None: + given: Final = [dict(entry) for entry in value] + core: Final = make_core(model_list=[deployment("g", "d1")]) + _ = core.apply(ChangeSet(assign={name: given})) + given.clear() + assert getattr(core.settings(), name) == value # pyright: ignore[reportAny] # a settings field named by the parameter + + +@pytest.mark.parametrize(("name", "value"), ASSIGNED) +def test_update_settings_ignores_a_host_setting(make_core: MakeCore, name: str, value: object) -> None: + core: Final = make_core(model_list=[deployment("g", "d1")]) + before: Final[object] = getattr(core.settings(), name) # pyright: ignore[reportAny] # a settings field named by the parameter + _ = core.apply(ChangeSet(settings={name: value})) + assert getattr(core.settings(), name) == before + + +def test_a_router_attribute_write_assigns_a_host_setting(make_core: MakeCore) -> None: + core: Final = make_core(model_list=[deployment("g", "d1")]) + core.search_tools = SEARCH_TOOLS # pyright: ignore[reportAttributeAccessIssue] # the deprecated attribute write + assert core.settings().search_tools == SEARCH_TOOLS + + +# --- router_general_settings as a constructor argument ------------------------------------------ + + +def test_router_general_settings_passed_to_the_constructor_are_reported(make_core: MakeCore) -> None: + general: Final = RouterGeneralSettings(async_only_mode=True, pass_through_all_models=True) + core: Final = make_core(model_list=[deployment("g", "d1")], router_general_settings=general) + assert core.settings().router_general_settings == general + + +def test_router_general_settings_of_none_become_the_defaults(make_core: MakeCore) -> None: + core: Final = make_core(model_list=[deployment("g", "d1")], router_general_settings=None) + assert core.settings().router_general_settings == RouterGeneralSettings() + + +def test_routers_built_without_general_settings_share_one_default_object(make_core: MakeCore) -> None: + """The Python core's default argument is one ``RouterGeneralSettings()`` every router shares.""" + first: Final = make_core(model_list=[deployment("g", "d1")]) + second: Final = make_core(model_list=[deployment("g", "d1")]) + assert first.settings().router_general_settings is second.settings().router_general_settings + + +async def test_router_general_settings_leave_routing_unchanged(make_core: MakeCore, host: EndpointHost) -> None: + general: Final = RouterGeneralSettings(async_only_mode=True, pass_through_all_models=True) + core: Final = make_core(model_list=[deployment("g", "d1")], router_general_settings=general) + _ = await _endpoint(core, "acompletion")(model="g", messages=MESSAGES) + completion: Final[Callable[..., object]] = getattr(core, "completion") # pyright: ignore[reportAny] # see _endpoint + _ = completion(model="g", messages=MESSAGES) + assert host.deployments_called() == ["d1", "d1"] + + +# --- readers ------------------------------------------------------------------------------------ + + +async def test_asearch_reads_search_tools_assigned_after_construction(make_core: MakeCore, host: EndpointHost): + core: Final = make_core(model_list=[deployment("g", "d1")], num_retries=0) + _ = core.apply(ChangeSet(assign={"search_tools": SEARCH_TOOLS})) + _ = await _endpoint(core, "asearch")(model="s", query="q") + call: Final = host.calls[0] + assert (call.op, call.kwargs["search_provider"], call.kwargs["api_key"]) == ("asearch", "tavily", "tk") + + +async def test_asearch_no_longer_finds_a_tool_dropped_by_a_changeset(make_core: MakeCore, host: EndpointHost): + core: Final = make_core(model_list=[deployment("g", "d1")], search_tools=SEARCH_TOOLS, num_retries=0) + _ = core.apply(ChangeSet(assign={"search_tools": []})) + with pytest.raises(ValueError, match=r"Search tool 's' not found in router\.search_tools"): + _ = await _endpoint(core, "asearch")(model="s", query="q") + assert host.calls == [] + + +async def test_aguardrail_reads_a_guardrail_list_assigned_after_construction(make_core: MakeCore, host: ScriptedHost): + seen: Final[list[Mapping[str, object]]] = [] + + async def guard(**kwargs: object) -> str: # kwargs-ok: the guardrail's keyword payload + seen.append(kwargs) + return "verdict" + + core: Final = make_core(model_list=[deployment("g", "d1")], num_retries=0) + _ = core.apply(ChangeSet(assign={"guardrail_list": GUARDRAILS})) + assert await _endpoint(core, "aguardrail")(guardrail_name="gr", original_function=guard) == "verdict" + assert seen[0]["selected_guardrail"] == GUARDRAILS[0] + + +async def test_a_default_max_parallel_requests_assigned_later_limits_the_next_requests( + make_core: MakeCore, host: ScriptedHost +) -> None: + host.answer("openai/a", chat_stream("x")) + core: Final = make_core(model_list=[deployment("a", "a")], num_retries=0) + _ = core.apply(ChangeSet(assign={"default_max_parallel_requests": 1})) + first: Final = await _open_stream(core) + with pytest.raises(litellm.RateLimitError): + await asyncio.wait_for(_open_stream(core), 5) + assert await collect_texts(first) == "x" + + +async def test_a_deployment_keeps_the_slot_limit_it_was_given_before_a_change( + make_core: MakeCore, host: ScriptedHost +) -> None: + """The limit is cached per deployment on first use: a later default does not replace it.""" + host.answer("openai/a", chat_stream("x")) + core: Final = make_core(model_list=[deployment("a", "a")], num_retries=0, default_max_parallel_requests=1) + assert await collect_texts(await _open_stream(core)) == "x" + _ = core.apply(ChangeSet(assign={"default_max_parallel_requests": 5})) + first: Final = await _open_stream(core) + with pytest.raises(litellm.RateLimitError): + await asyncio.wait_for(_open_stream(core), 5) + assert await collect_texts(first) == "x" + + +async def test_a_budget_config_assigned_beside_the_budget_check_builds_the_limiter_with_it(make_core: MakeCore) -> None: + """``apply_changes`` assigns before ``update_settings`` adds the check.""" + core: Final = make_core(model_list=[deployment("g", "a")]) + _ = core.apply( + ChangeSet( + assign={"provider_budget_config": PROVIDER_BUDGET}, + settings={"optional_pre_call_checks": ["router_budget_limiting"]}, + ) + ) + limiter: Final = core.registry().plugins("budget_limiter") + assert isinstance(limiter, RouterBudgetLimiting) + assert limiter.provider_budget_config is core.settings().provider_budget_config + assert list(limiter.provider_budget_config or {}) == ["openai"] + core.retire() + + +async def test_a_budget_limiter_built_already_keeps_its_config(make_core: MakeCore) -> None: + core: Final = make_core(model_list=[deployment("g", "a")], provider_budget_config=PROVIDER_BUDGET) + limiter: Final = core.registry().plugins("budget_limiter") + assert isinstance(limiter, RouterBudgetLimiting) + _ = core.apply(ChangeSet(assign={"provider_budget_config": {"azure": {"budget_limit": 2, "time_period": "1d"}}})) + assert core.registry().plugins("budget_limiter") is limiter + assert list(limiter.provider_budget_config or {}) == ["openai"] + assert list(core.settings().provider_budget_config or {}) == ["azure"] + core.retire() + + +# --- settings().guardrail(name) ----------------------------------------------------------------- + + +def test_settings_guardrail_returns_the_one_configured_under_the_name(make_core: MakeCore) -> None: + core: Final = make_core(model_list=[deployment("g", "d1")], guardrail_list=GUARDRAILS) + assert core.settings().guardrail("gr") == GUARDRAILS[0] + + +def test_settings_guardrail_with_an_unknown_name_raises_value_error(make_core: MakeCore) -> None: + core: Final = make_core(model_list=[deployment("g", "d1")], guardrail_list=GUARDRAILS) + with pytest.raises(ValueError, match="No guardrail found with name: nope"): + _ = core.settings().guardrail("nope") + + +def test_settings_guardrail_picks_among_same_named_guardrails_by_weight(make_core: MakeCore) -> None: + guardrails: Final = [ + {"guardrail_name": "gr", "litellm_params": {"guardrail": "never", "weight": 0}}, + {"guardrail_name": "gr", "litellm_params": {"guardrail": "always", "weight": 1}}, + ] + core: Final = make_core(model_list=[deployment("g", "d1")], guardrail_list=guardrails) + assert {core.settings().guardrail("gr")["litellm_params"]["guardrail"] for _ in range(20)} == {"always"} + + +def test_settings_guardrail_reads_a_guardrail_list_assigned_after_construction(make_core: MakeCore) -> None: + core: Final = make_core(model_list=[deployment("g", "d1")]) + _ = core.apply(ChangeSet(assign={"guardrail_list": GUARDRAILS})) + assert core.settings().guardrail("gr") == GUARDRAILS[0] + + +# --- a default_max_parallel_requests that is no int --------------------------------------------- + + +@pytest.mark.parametrize("assigned", [False, True]) +async def test_a_float_default_max_parallel_requests_limits_the_slots( + make_core: MakeCore, host: ScriptedHost, assigned: bool +) -> None: + host.answer("openai/a", chat_stream("x")) + given: Final = {} if assigned else {"default_max_parallel_requests": 1.0} + core: Final = make_core(model_list=[deployment("a", "a")], num_retries=0, **given) + if assigned: + _ = core.apply(ChangeSet(assign={"default_max_parallel_requests": 1.0})) + first: Final = await _open_stream(core) + with pytest.raises(litellm.RateLimitError): + await asyncio.wait_for(_open_stream(core), 5) + assert await collect_texts(first) == "x" + + +# --- health-check routing assigned after construction ------------------------------------------- + + +async def _healthy_ids(core: RouterCore, model: str) -> list[str]: # mutable-ok: a fresh list + select: Final[Callable[..., Awaitable[object]]] = getattr(core, "select") # pyright: ignore[reportAny] # the core's select, which takes kind + return ids_of(await select(model, kind="healthy")) + + +def _unhealthy(*ids: str) -> dict[str, dict[str, object]]: # mutable-ok: health states + return {i: {"is_healthy": False, "timestamp": time.time(), "reason": "probe failed"} for i in ids} + + +async def test_health_check_routing_turned_on_later_drops_unhealthy_deployments(make_core: MakeCore) -> None: + core: Final = make_core(model_list=[deployment("g", "a"), deployment("g", "b")]) + core.health().health_checks.set_deployment_health_states(_unhealthy("a")) # pyright: ignore[reportArgumentType] # health states + assert await _healthy_ids(core, "g") == ["a", "b"] + _ = core.apply(ChangeSet(assign={"enable_health_check_routing": True})) + assert await _healthy_ids(core, "g") == ["b"] + _ = core.apply(ChangeSet(assign={"enable_health_check_routing": False})) + assert await _healthy_ids(core, "g") == ["a", "b"] + + +async def test_background_health_check_model_groups_assigned_later_scope_the_filter(make_core: MakeCore) -> None: + core: Final = make_core( + model_list=[deployment("g", "a"), deployment("g", "b"), deployment("h", "c"), deployment("h", "d")], + enable_health_check_routing=True, + ) + core.health().health_checks.set_deployment_health_states(_unhealthy("a", "c")) # pyright: ignore[reportArgumentType] # health states + assert await _healthy_ids(core, "h") == ["d"] + _ = core.apply(ChangeSet(assign={"background_health_check_model_groups": ["g"]})) + assert await _healthy_ids(core, "g") == ["b"] + assert await _healthy_ids(core, "h") == ["c", "d"] + + +# --- provider_default_deployment_ids assigned after construction -------------------------------- + + +@pytest.mark.parametrize(("default_ids", "cooled"), [([], ["a"]), (["a"], [])]) +async def test_an_assigned_provider_default_deployment_never_cools_down( + make_core: MakeCore, host: ScriptedHost, default_ids: list[str], cooled: list[str] +) -> None: + host.answer("openai/a", rate_limit()) + core: Final = make_core(model_list=[deployment("g", "a"), deployment("g", "b")], num_retries=0) + _ = core.apply(ChangeSet(assign={"provider_default_deployment_ids": default_ids})) + with pytest.raises(litellm.RateLimitError): + await core.call("acompletion", {"model": "a", "messages": MESSAGES}) + assert sorted(cooled_down(core, ["a", "b"])) == cooled + + +async def test_wildcard_deployments_never_cool_down(make_core: MakeCore, host: ScriptedHost) -> None: + """Each is a provider default deployment, which the cooldown check skips.""" + host.answer("openai/x", rate_limit()) + core: Final = make_core( + model_list=[deployment("openai/*", "w", model="openai/*"), deployment("openai/*", "w2", model="openai/*")], + num_retries=0, + ) + for _ in range(3): + with pytest.raises(litellm.RateLimitError): + await core.call("acompletion", {"model": "openai/x", "messages": MESSAGES}) + assert cooled_down(core, ["w", "w2"]) == {} + assert len(host.calls) == 3 + + +async def test_an_added_wildcard_deployment_never_cools_down(make_core: MakeCore, host: ScriptedHost) -> None: + host.answer("openai/x", rate_limit()) + core: Final = make_core(model_list=[deployment("openai/*", "w", model="openai/*")], num_retries=0) + _ = core.apply(ChangeSet(upsert=[Deployment.model_validate(deployment("openai/*", "w2", model="openai/*"))])) + assert core.settings().provider_default_deployment_ids == ["w", "w2"] + for _ in range(3): + with pytest.raises(litellm.RateLimitError): + await core.call("acompletion", {"model": "openai/x", "messages": MESSAGES}) + assert cooled_down(core, ["w", "w2"]) == {} + + +# --- Kubernetes pod discovery ------------------------------------------------------------------- + + +class _RecordingDiscovery(KubernetesPodDiscovery): + """Resolves every pod-discovery deployment to one pod IP, and records the retry counts seen.""" + + def __init__(self) -> None: + super().__init__(proxy_environment=lambda: {}) + self.seen: list[object] = [] # mutable-ok: filled as picks resolve + + def _pod(self, deployment: object) -> object: + record: Final = dict(as_str_mapping(deployment) or {}) + params: Final = as_str_mapping(record.get("litellm_params")) or {} + if params.get("kubernetes_pod_discovery") is not True: + return deployment + record["litellm_params"] = {**params, "api_base": "http://10.0.0.7:4000"} + return record + + def resolve_deployment(self, deployment: object, request_kwargs: Mapping[str, object] | None = None) -> object: # pyright: ignore[reportIncompatibleMethodOverride] # the resolver's generic signature, read as objects + self.seen.append(self._request_retry_count(request_kwargs)) + return self._pod(deployment) + + async def async_resolve_deployment( # pyright: ignore[reportIncompatibleMethodOverride] # see resolve_deployment + self, deployment: object, request_kwargs: Mapping[str, object] | None = None + ) -> object: + self.seen.append(self._request_retry_count(request_kwargs)) + return self._pod(deployment) + + +def _pod_deployment(deployment_id: str) -> dict[str, object]: # mutable-ok: a model_list entry + return deployment("g", deployment_id, api_base="http://svc.ns:4000", kubernetes_pod_discovery=True) + + +def test_settings_report_the_routers_pod_discovery(make_core: MakeCore) -> None: + settings: Final = make_core(model_list=[deployment("g", "d1")]).settings() + assert settings.is_set("kubernetes_pod_discovery") + assert isinstance(settings.kubernetes_pod_discovery, KubernetesPodDiscovery) + + +async def test_an_assigned_pod_discovery_resolves_each_attempt(make_core: MakeCore, host: ScriptedHost) -> None: + host.answer("openai/a", rate_limit()) + discovery: Final = _RecordingDiscovery() + core: Final = make_core(model_list=[_pod_deployment("a")], num_retries=1) + _ = core.apply(ChangeSet(assign={"kubernetes_pod_discovery": discovery})) + assert core.settings().kubernetes_pod_discovery is discovery + with pytest.raises(litellm.RateLimitError): + await core.call("acompletion", {"model": "g", "messages": MESSAGES}) + assert [call.kwargs["api_base"] for call in host.calls] == ["http://10.0.0.7:4000"] * 2 + assert discovery.seen == [0, 1] + + +def test_an_assigned_pod_discovery_resolves_a_sync_attempt(make_core: MakeCore, host: ScriptedHost) -> None: + discovery: Final = _RecordingDiscovery() + core: Final = make_core(model_list=[_pod_deployment("a")]) + _ = core.apply(ChangeSet(assign={"kubernetes_pod_discovery": discovery})) + completion: Final[Callable[..., object]] = getattr(core, "completion") # pyright: ignore[reportAny] # see _endpoint + _ = completion(model="g", messages=MESSAGES) + assert host.calls[0].kwargs["api_base"] == "http://10.0.0.7:4000" + + +def _api_base(record: object) -> object: + return (as_str_mapping((as_str_mapping(record) or {}).get("litellm_params")) or {}).get("api_base") + + +async def test_a_pick_is_resolved_to_a_pod_and_a_healthy_listing_is_not(make_core: MakeCore) -> None: + core: Final = make_core(model_list=[_pod_deployment("a")]) + _ = core.apply(ChangeSet(assign={"kubernetes_pod_discovery": _RecordingDiscovery()})) + assert _api_base(await core.select("g")) == "http://10.0.0.7:4000" + select: Final[Callable[..., Awaitable[object]]] = getattr(core, "select") # pyright: ignore[reportAny] # see _healthy_ids + listed: Final[object] = await select("g", kind="healthy") + assert isinstance(listed, list) + assert [_api_base(record) for record in listed] == ["http://svc.ns:4000"] # pyright: ignore[reportUnknownVariableType, reportUnknownArgumentType] # the listing + + +def test_a_sync_pick_is_resolved_to_a_pod(make_core: MakeCore) -> None: + core: Final = make_core(model_list=[_pod_deployment("a")]) + _ = core.apply(ChangeSet(assign={"kubernetes_pod_discovery": _RecordingDiscovery()})) + pick: Final[Callable[..., object]] = getattr(core, "get_available_deployment") # pyright: ignore[reportAny] # the deprecated sync pick + assert _api_base(pick(model="g")) == "http://10.0.0.7:4000" + + +async def test_the_default_pod_discovery_resolves_a_service_host(make_core: MakeCore, host: ScriptedHost) -> None: + """``localhost`` resolves to a loopback address, which takes the service host's place.""" + core: Final = make_core( + model_list=[deployment("g", "a", api_base="http://localhost:4000", kubernetes_pod_discovery=True)] + ) + _ = core.apply(ChangeSet(assign={"kubernetes_pod_discovery": KubernetesPodDiscovery(proxy_environment=lambda: {})})) + await core.call("acompletion", {"model": "g", "messages": MESSAGES}) + assert str(host.calls[0].kwargs["api_base"]) in ("http://127.0.0.1:4000", "http://[::1]:4000") diff --git a/tests/unit/test_router/core/test_legacy_members.py b/tests/unit/test_router/core/test_legacy_members.py new file mode 100644 index 00000000000..1158a44e7fd --- /dev/null +++ b/tests/unit/test_router/core/test_legacy_members.py @@ -0,0 +1,711 @@ +""" +The old public Router members outside the entry points, on each core: registry lookups, the +routing-state lifecycle, the routing callbacks' checks, one provider call's accounting and the +pre-routing strategy builders (``RustLegacyMembers``, litellm/router_rust/surface/legacy_members.py). +""" + +import asyncio +import inspect +from collections.abc import Mapping +from typing import Final + +import pytest + +import litellm +from litellm.constants import ROUTER_USAGE_COUNTED_TOKENS_METADATA_KEY +from litellm.containers.endpoint_factory import ( + _generated_endpoints, # pyright: ignore[reportPrivateUsage] # the names the routers register +) +from litellm.integrations.custom_logger import CustomLogger +from litellm.litellm_core_utils.token_counter import served_openai_models +from litellm.router_api import RoutingContext +from litellm.router_core import RouterCore +from litellm.router_python_core import PythonCore +from litellm.router_rust.surface.legacy_members import ( + _ASSISTANTS_ENDPOINTS, # pyright: ignore[reportPrivateUsage] # pinned against the Python core below + _ROUTER_ENDPOINTS, # pyright: ignore[reportPrivateUsage] # see above +) +from litellm.router_strategy.adaptive_router.hooks import AdaptiveRouterPostCallHook +from litellm.router_strategy.auto_router.auto_router_model_naming import HEURISTIC_V2_CAPABILITY +from litellm.router_strategy.base_routing_strategy import BaseRoutingStrategy +from litellm.router_strategy.least_busy import LeastBusyLoggingHandler +from litellm.router_strategy.lowest_tpm_rpm_v2 import LowestTPMLoggingHandler_v2 +from litellm.router_utils.pre_call_checks.prompt_caching_deployment_check import PromptCachingDeploymentCheck +from litellm.types.router import CustomRoutingStrategyBase, Deployment +from tests.unit.test_router.core.construction_precall_helpers import CallbackHost +from tests.unit.test_router.core.gaps_l11100_helpers import ( + TIERS, + adaptive_marker, + complexity_marker, + marker, + registered_tags, +) +from tests.unit.test_router.core.registry_helpers import MESSAGES, MakeCore, team_deployment +from tests.unit.test_router.core.retries_health_helpers import cooled_down +from tests.unit.test_router.core.scripted_host import ( + ScriptedHost, + completion_response, + deployment, + python_core, + rate_limit, + rust_core, + server_error, +) + +GROUP: Final = [deployment("g", "a"), deployment("g", "b")] +PLAIN: Final = [deployment("small", "s"), deployment("big", "b")] + + +@pytest.fixture +def host() -> ScriptedHost: + """Switching an optional pre-call check off removes it from the host's lists.""" + return CallbackHost() + + +class Gate(CustomLogger): + """Records what its pre-call checks saw and raises ``error`` (when set) from them.""" + + def __init__(self, error: Exception | None = None) -> None: + super().__init__() + self.error = error + self.seen: Final[list[tuple[object, object]]] = [] + + def pre_call_check(self, deployment: dict) -> dict | None: + self.seen.append((deployment, None)) + if self.error is not None: + raise self.error + return None + + async def async_pre_call_check(self, deployment: dict, parent_otel_span: object) -> dict | None: # pyright: ignore[reportIncompatibleMethodOverride] + self.seen.append((deployment, parent_otel_span)) + if self.error is not None: + raise self.error + return None + + +class Narrow(CustomLogger): + """Keeps the deployments other than ``dropped``, or raises ``error``.""" + + def __init__(self, dropped: str = "", error: Exception | None = None) -> None: + super().__init__() + self.dropped = dropped + self.error = error + self.calls: Final[list[Mapping[str, object]]] = [] + + async def async_filter_deployments( + self, + model: str, + healthy_deployments: list, + messages: object, + request_kwargs: object = None, + parent_otel_span: object = None, + ) -> list: + self.calls.append( + { + "model": model, + "messages": messages, + "request_kwargs": request_kwargs, + "parent_otel_span": parent_otel_span, + } + ) + if self.error is not None: + raise self.error + return [d for d in healthy_deployments if d["model_info"]["id"] != self.dropped] + + +class RecordingLogging: + """The logging object's failure dispatch, recorded.""" + + def __init__(self) -> None: + self.dispatched: Final[list[Mapping[str, object]]] = [] + + async def dispatch_failure_handlers(self, **kwargs: object) -> None: # kwargs-ok: recorded as passed + self.dispatched.append(kwargs) + + +def counters(core: RouterCore, deployment_id: str, kind: str) -> list[object]: # mutable-ok: a fresh list + """The per-minute ``kind`` (``tpm``, ``rpm``) counters of ``deployment_id`` in the routing cache.""" + memory: Final = core.health().store.in_memory_cache + return [ + memory.get_cache(key) + for key in tuple(memory.cache_dict) + if key.startswith(f"global_router:{deployment_id}:") and f":{kind}:" in key + ] + + +def call_metadata(deployment_id: str = "a") -> dict[str, object]: # mutable-ok: the request's own bucket + return {"model_info": {"id": deployment_id}, "deployment": f"openai/{deployment_id}", "model_group": "g"} + + +# --- registry lookups -------------------------------------------------------------------------- + + +def test_get_routing_group_names_a_group_a_deployment_does_not_shadow(make_core: MakeCore) -> None: + core: Final = make_core( + model_list=[deployment("m1", "a"), deployment("m2", "b")], + routing_groups=[{"group_name": "grp", "models": ["m1", "m2"], "routing_strategy": "simple-shuffle"}], + ) + group: Final = core.get_routing_group("grp") # pyright: ignore[reportAttributeAccessIssue] # an old public member + + assert group is not None + assert group.group_name == "grp" + assert core.get_routing_group("m1") is None # pyright: ignore[reportAttributeAccessIssue] # see above + assert core.get_routing_group("nope") is None # pyright: ignore[reportAttributeAccessIssue] # see above + + +def test_team_public_model_names_of_lists_one_teams_public_names(make_core: MakeCore) -> None: + core: Final = make_core( + model_list=[ + team_deployment("int-1", "t1", "TA", "pub-a"), + team_deployment("int-2", "t2", "TB", "pub-b"), + ] + ) + + assert core.team_public_model_names_of("TA") == frozenset({"pub-a"}) # pyright: ignore[reportAttributeAccessIssue] # an old public member + assert core.team_public_model_names_of("nobody") == frozenset() # pyright: ignore[reportAttributeAccessIssue] # see above + + +def test_served_openai_models_reads_the_group_and_follows_a_model_list_change(make_core: MakeCore) -> None: + core: Final = make_core(model_list=[deployment("g", "a", model="openai/gpt-4o")]) + before: Final = core.served_openai_models("g") # pyright: ignore[reportAttributeAccessIssue] # an old public member + + assert before == served_openai_models(core.registry().deployments("g", source="named") or ()) + assert core.served_openai_models("nope") == () # pyright: ignore[reportAttributeAccessIssue] # see above + + core.set_model_list([deployment("g", "b", model="openai/gpt-4o-mini")]) # pyright: ignore[reportAttributeAccessIssue] # an old public member + + assert core.served_openai_models("g") == served_openai_models( + core.registry().deployments("g", source="named") or () + ) # pyright: ignore[reportAttributeAccessIssue] # see above + + +def test_the_tier_params_never_dropped_constant_matches(make_core: MakeCore) -> None: + core: Final = make_core(model_list=list(GROUP)) + + assert core.TIER_PARAMS_NEVER_DROPPED == PythonCore.TIER_PARAMS_NEVER_DROPPED # pyright: ignore[reportAttributeAccessIssue] # a class constant + + +def test_auto_router_capability_violation_resolves_the_limit_on_every_call(make_core: MakeCore) -> None: + limits: Final[list[int | None]] = [None] # mutable-ok: what the resolver answers + + core: Final = make_core( + model_list=[*PLAIN, complexity_marker(config={"classifier_type": "heuristic_v2", "tiers": dict(TIERS)})], + auto_router_capability_limit=lambda: limits[0], + ) + + assert core.auto_router_capability_violation(HEURISTIC_V2_CAPABILITY) is None # pyright: ignore[reportAttributeAccessIssue] # an old public member + limits[0] = 1 + violation: Final = core.auto_router_capability_violation(HEURISTIC_V2_CAPABILITY) # pyright: ignore[reportAttributeAccessIssue] # see above + assert violation is not None + assert violation.startswith("At most 1 auto-router(s) using the heuristic_v2 classifier can be registered") + + +# --- routing state ------------------------------------------------------------------------------ + + +def test_retire_routing_state_retires_the_selectors_and_keeps_them_registered( + make_core: MakeCore, host: ScriptedHost, monkeypatch: pytest.MonkeyPatch +) -> None: + core: Final = make_core(model_list=list(GROUP), routing_strategy="usage-based-routing-v2") + selector: Final = core.lowesttpm_logger_v2 # pyright: ignore[reportAttributeAccessIssue] # an old public attribute + assert isinstance(selector, BaseRoutingStrategy) + retired: Final[list[object]] = [] # mutable-ok: what retire() saw + monkeypatch.setattr(selector, "retire", lambda: retired.append(selector)) + + core.retire_routing_state() # pyright: ignore[reportAttributeAccessIssue] # an old public member + + assert retired == [selector] + assert core.lowesttpm_logger_v2 is selector # pyright: ignore[reportAttributeAccessIssue] # see above + assert selector in host.lists["callbacks"] + + +def test_routing_strategy_init_refuses_an_unknown_strategy_and_changes_nothing(make_core: MakeCore) -> None: + core: Final = make_core(model_list=list(GROUP), routing_strategy="least-busy") + selector: Final = core.leastbusy_logger # pyright: ignore[reportAttributeAccessIssue] # an old public attribute + + with pytest.raises(ValueError, match="not-a-real-strategy"): + core.routing_strategy_init(routing_strategy="not-a-real-strategy", routing_strategy_args={}) # pyright: ignore[reportAttributeAccessIssue] # an old public member + + assert core.leastbusy_logger is selector # pyright: ignore[reportAttributeAccessIssue] # see above + + +def test_routing_strategy_init_replaces_the_selector_and_leaves_the_setting( + make_core: MakeCore, host: ScriptedHost +) -> None: + core: Final = make_core(model_list=list(GROUP), routing_strategy="least-busy") + old: Final = core.leastbusy_logger # pyright: ignore[reportAttributeAccessIssue] # an old public attribute + + core.routing_strategy_init("usage-based-routing-v2", {}) # pyright: ignore[reportAttributeAccessIssue] # an old public member + + assert core.leastbusy_logger is None # pyright: ignore[reportAttributeAccessIssue] # see above + selector: Final = core.lowesttpm_logger_v2 # pyright: ignore[reportAttributeAccessIssue] # see above + assert isinstance(selector, LowestTPMLoggingHandler_v2) + assert selector in host.lists["callbacks"] + assert old not in host.lists["callbacks"] + assert core.settings().routing_strategy == "least-busy" + + +@pytest.mark.asyncio +async def test_a_request_override_of_the_initialized_strategy_is_not_the_routers_own(make_core: MakeCore) -> None: + """``_get_override_strategy_selector`` compares with the router's ``routing_strategy``, which + ``routing_strategy_init`` leaves, so the request gets its own selector.""" + core: Final = make_core(model_list=list(GROUP)) + core.routing_strategy_init("least-busy", {}) # pyright: ignore[reportAttributeAccessIssue] # an old public member + initialized: Final = core.leastbusy_logger # pyright: ignore[reportAttributeAccessIssue] # an old public attribute + assert isinstance(initialized, LeastBusyLoggingHandler) + + picked: Final = await core.select("g", {"routing_strategy": "least-busy", "messages": list(MESSAGES)}) + + assert picked["model_info"]["id"] in {"a", "b"} # pyright: ignore[reportIndexIssue] # the pick + assert core.leastbusy_logger is initialized # pyright: ignore[reportAttributeAccessIssue] # see above + + +@pytest.mark.asyncio +async def test_routing_strategy_init_drops_a_custom_strategy(make_core: MakeCore) -> None: + class Custom(CustomRoutingStrategyBase): + async def async_get_available_deployment(self, *args: object, **kwargs: object) -> dict: # kwargs-ok: any pick + return dict(GROUP[0]) + + def get_available_deployment(self, *args: object, **kwargs: object) -> dict: # kwargs-ok: any pick + return dict(GROUP[0]) + + core: Final = make_core(model_list=list(GROUP)) + core.register_plugin(Custom()) + assert {(await core.select("g", {"messages": list(MESSAGES)}))["model_info"]["id"] for _ in range(20)} == {"a"} # pyright: ignore[reportIndexIssue] # the pick + + core.routing_strategy_init("simple-shuffle", {}) # pyright: ignore[reportAttributeAccessIssue] # an old public member + + picks: Final = {(await core.select("g", {"messages": list(MESSAGES)}))["model_info"]["id"] for _ in range(60)} # pyright: ignore[reportIndexIssue] # the pick + assert picks == {"a", "b"} + + +def test_set_optional_pre_call_checks_adds_and_drops_the_toggleable_checks( + make_core: MakeCore, host: ScriptedHost +) -> None: + core: Final = make_core(model_list=list(GROUP)) + + core.set_optional_pre_call_checks(["prompt_caching"]) # pyright: ignore[reportAttributeAccessIssue] # an old public member + assert [type(callback) for callback in core.optional_callbacks] == [PromptCachingDeploymentCheck] # pyright: ignore[reportAttributeAccessIssue] # see above + assert any(isinstance(callback, PromptCachingDeploymentCheck) for callback in host.lists["callbacks"]) + + core.set_optional_pre_call_checks(None) # pyright: ignore[reportAttributeAccessIssue] # see above + assert [type(callback) for callback in core.optional_callbacks] == [PromptCachingDeploymentCheck] # pyright: ignore[reportAttributeAccessIssue] # see above + + core.set_optional_pre_call_checks([]) # pyright: ignore[reportAttributeAccessIssue] # see above + assert core.optional_callbacks == [] # pyright: ignore[reportAttributeAccessIssue] # see above + assert not any(isinstance(callback, PromptCachingDeploymentCheck) for callback in host.lists["callbacks"]) + + +def test_initialize_assistants_endpoint_replaces_an_assigned_endpoint(make_core: MakeCore) -> None: + core: Final = make_core(model_list=list(GROUP)) + assigned: Final = object() + core.aget_assistants = assigned # pyright: ignore[reportAttributeAccessIssue] # an old public endpoint + core.amoderation = assigned # pyright: ignore[reportAttributeAccessIssue] # see above + + core.initialize_assistants_endpoint() # pyright: ignore[reportAttributeAccessIssue] # an old public member + + assert core.aget_assistants is not assigned # pyright: ignore[reportAttributeAccessIssue] # see above + assert callable(core.aget_assistants) # pyright: ignore[reportAttributeAccessIssue] # see above + assert core.amoderation is assigned # pyright: ignore[reportAttributeAccessIssue] # see above + + +def test_initialize_router_endpoints_replaces_only_what_it_binds(make_core: MakeCore) -> None: + core: Final = make_core(model_list=list(GROUP)) + assigned: Final = object() + for name in ("amoderation", "aresponses", "aupload_container_file", "aget_assistants", "acreate_file"): + setattr(core, name, assigned) + + core.initialize_router_endpoints() # pyright: ignore[reportAttributeAccessIssue] # an old public member + + assert core.amoderation is not assigned # pyright: ignore[reportAttributeAccessIssue] # see above + assert core.aresponses is not assigned # pyright: ignore[reportAttributeAccessIssue] # see above + assert core.aupload_container_file is not assigned # pyright: ignore[reportAttributeAccessIssue] # see above + # `initialize_router_endpoints` binds neither of these: the assistants have their own, and + # `acreate_file` is a method an assignment shadows. + assert core.aget_assistants is assigned # pyright: ignore[reportAttributeAccessIssue] # see above + assert core.acreate_file is assigned # pyright: ignore[reportAttributeAccessIssue] # see above + + +def test_the_endpoint_names_are_the_ones_the_python_core_binds() -> None: + bound: Final = vars(python_core(ScriptedHost(), model_list=[])) + listed: Final = {*_ROUTER_ENDPOINTS, *_ASSISTANTS_ENDPOINTS, *_generated_endpoints} + + assert listed <= set(bound) + wrappers: Final = { + name + for name, value in bound.items() + if inspect.isfunction(value) and value.__qualname__.startswith("PythonCore.factory_function.") + } + assert wrappers <= listed + + +# --- pre-routing strategies ---------------------------------------------------------------------- + + +def test_init_complexity_router_deployment_registers_it_once(make_core: MakeCore) -> None: + core: Final = make_core(model_list=list(PLAIN)) + router: Final = Deployment(**complexity_marker()) + + core.init_complexity_router_deployment(router) # pyright: ignore[reportAttributeAccessIssue] # an old public member + + assert registered_tags(core, "complexity") == {"smart": [()]} + with pytest.raises(ValueError, match="already exists"): + core.init_complexity_router_deployment(router) # pyright: ignore[reportAttributeAccessIssue] # see above + + +def test_init_quality_router_deployment_needs_a_default_model(make_core: MakeCore) -> None: + core: Final = make_core(model_list=list(PLAIN)) + + with pytest.raises(ValueError, match="quality_router_default_model is required"): + core.init_quality_router_deployment(Deployment(**marker("q", "m", "auto_router/quality_router"))) # pyright: ignore[reportAttributeAccessIssue] # an old public member + + core.init_quality_router_deployment( # pyright: ignore[reportAttributeAccessIssue] # see above + Deployment(**marker("q", "m", "auto_router/quality_router", quality_router_default_model="big")) + ) + assert registered_tags(core, "quality") == {"q": [()]} + + +def test_init_adaptive_router_deployment_registers_it_and_its_hook(make_core: MakeCore, host: ScriptedHost) -> None: + core: Final = make_core(model_list=list(PLAIN)) + + core.init_adaptive_router_deployment(Deployment(**adaptive_marker("ad", "m", ["big"]))) # pyright: ignore[reportAttributeAccessIssue] # an old public member + + assert registered_tags(core, "adaptive") == {"ad": [()]} + assert any(isinstance(callback, AdaptiveRouterPostCallHook) for callback in host.lists["callbacks"]) + + +def test_init_auto_router_deployment_needs_its_config(make_core: MakeCore) -> None: + core: Final = make_core(model_list=list(PLAIN)) + + with pytest.raises(ValueError, match="auto_router_config_path or auto_router_config is required"): + core.init_auto_router_deployment(Deployment(**marker("auto", "m", "auto_router/semantic"))) # pyright: ignore[reportAttributeAccessIssue] # an old public member + + +# --- routing callbacks --------------------------------------------------------------------------- + + +def test_routing_strategy_pre_call_checks_hands_each_callback_the_deployment( + make_core: MakeCore, host: ScriptedHost +) -> None: + gate: Final = Gate() + host.lists["callbacks"].append(gate) + core: Final = make_core(model_list=list(GROUP)) + picked: Final = dict(GROUP[0]) + + core.routing_strategy_pre_call_checks(picked) # pyright: ignore[reportAttributeAccessIssue] # an old public member + + assert gate.seen[0][0] is picked + + gate.error = rate_limit(message="sync gate") + with pytest.raises(litellm.RateLimitError, match="sync gate"): + core.routing_strategy_pre_call_checks(picked) # pyright: ignore[reportAttributeAccessIssue] # see above + assert cooled_down(core, ["a"]) == {} + + +@pytest.mark.asyncio +async def test_a_rate_limited_async_pre_call_check_cools_down_and_logs_the_failure( + make_core: MakeCore, host: ScriptedHost +) -> None: + gate: Final = Gate(rate_limit(message="gate says no")) + host.lists["callbacks"].append(gate) + core: Final = make_core(model_list=list(GROUP), cooldown_time=33) + logging_obj: Final = RecordingLogging() + span: Final = object() + + with pytest.raises(litellm.RateLimitError, match="gate says no") as raised: + await core.async_routing_strategy_pre_call_checks( # pyright: ignore[reportAttributeAccessIssue] # an old public member + deployment=dict(GROUP[0]), parent_otel_span=span, logging_obj=logging_obj + ) + await asyncio.sleep(0) + + assert gate.seen[0][1] is span + assert logging_obj.dispatched[0]["exception"] is raised.value + assert logging_obj.dispatched[0]["prefer_async_handlers"] is True + assert set(cooled_down(core, ["a", "b"])) == {"a"} + + +@pytest.mark.asyncio +async def test_another_async_pre_call_check_error_logs_the_failure_without_a_cooldown( + make_core: MakeCore, host: ScriptedHost +) -> None: + host.lists["callbacks"].append(Gate(ValueError("no"))) + core: Final = make_core(model_list=list(GROUP)) + logging_obj: Final = RecordingLogging() + + with pytest.raises(ValueError, match="no"): + await core.async_routing_strategy_pre_call_checks( # pyright: ignore[reportAttributeAccessIssue] # an old public member + deployment=dict(GROUP[0]), parent_otel_span=None, logging_obj=logging_obj + ) + await asyncio.sleep(0) + + assert len(logging_obj.dispatched) == 1 + assert cooled_down(core, ["a", "b"]) == {} + + +@pytest.mark.asyncio +async def test_async_callback_filter_deployments_narrows_through_each_callback( + make_core: MakeCore, host: ScriptedHost +) -> None: + first: Final = Narrow(dropped="a") + second: Final = Narrow() + host.lists["callbacks"].extend((first, second)) + core: Final = make_core(model_list=list(GROUP)) + span: Final = object() + request: Final = {"metadata": {}} + + remaining: Final = await core.async_callback_filter_deployments( # pyright: ignore[reportAttributeAccessIssue] # an old public member + model="g", + healthy_deployments=[dict(entry) for entry in GROUP], + messages=list(MESSAGES), + parent_otel_span=span, + request_kwargs=request, + ) + + assert [entry["model_info"]["id"] for entry in remaining] == ["b"] + assert second.calls[0]["parent_otel_span"] is span + assert second.calls[0]["request_kwargs"] is request + + +@pytest.mark.asyncio +async def test_a_failing_callback_filter_logs_the_failure_and_raises(make_core: MakeCore, host: ScriptedHost) -> None: + host.lists["callbacks"].append(Narrow(error=ValueError("filter refused"))) + core: Final = make_core(model_list=list(GROUP)) + logging_obj: Final = RecordingLogging() + + with pytest.raises(ValueError, match="filter refused"): + await core.async_callback_filter_deployments( # pyright: ignore[reportAttributeAccessIssue] # an old public member + model="g", + healthy_deployments=[dict(entry) for entry in GROUP], + messages=None, + parent_otel_span=None, + logging_obj=logging_obj, + ) + await asyncio.sleep(0) + + assert len(logging_obj.dispatched) == 1 + + +@pytest.mark.asyncio +async def test_async_admit_deployment_turn_runs_the_checks_on_the_deployment_as_configured_now( + make_core: MakeCore, host: ScriptedHost +) -> None: + gate: Final = Gate() + host.lists["callbacks"].append(gate) + core: Final = make_core(model_list=list(GROUP)) + + admission: Final = await core.async_admit_deployment_turn("g", {"model_info": {"id": "a"}}, {}) # pyright: ignore[reportAttributeAccessIssue] # an old public member + + assert dict(admission) == {} + assert gate.seen[0][0]["litellm_params"]["model"] == "openai/a" # pyright: ignore[reportIndexIssue] # the deployment dict + + +@pytest.mark.asyncio +async def test_a_rate_limited_turn_cools_the_deployment_down(make_core: MakeCore, host: ScriptedHost) -> None: + host.lists["callbacks"].append(Gate(rate_limit(message="turn refused"))) + core: Final = make_core(model_list=list(GROUP)) + + with pytest.raises(litellm.RateLimitError, match="turn refused"): + await core.async_admit_deployment_turn("g", {"model_info": {"id": "a"}}, {}) # pyright: ignore[reportAttributeAccessIssue] # an old public member + + assert set(cooled_down(core, ["a", "b"])) == {"a"} + + +# --- one provider call --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_make_call_counts_a_tracked_deployments_usage_and_notes_it_on_the_request(make_core: MakeCore) -> None: + core: Final = make_core(model_list=[deployment("g", "a", model="openai/gpt-4o", rpm=5)]) + metadata: Final = call_metadata() + + async def endpoint(**kwargs: object) -> object: # kwargs-ok: the request, as make_call passes it + return completion_response(total_tokens=10) + + await core.make_call(endpoint, model="g", metadata=metadata) # pyright: ignore[reportAttributeAccessIssue] # an old public member + + assert counters(core, "a", "tpm") == [10] + assert counters(core, "a", "rpm") == [1] + assert metadata[ROUTER_USAGE_COUNTED_TOKENS_METADATA_KEY] == 10 + + +@pytest.mark.asyncio +async def test_make_call_runs_a_sync_endpoint_and_stamps_the_response(make_core: MakeCore) -> None: + core: Final = make_core(model_list=[deployment("g", "a", model="openai/gpt-4o")]) + answer: Final = completion_response() + + def endpoint(**kwargs: object) -> object: # kwargs-ok: the request, as make_call passes it + return answer + + response: Final = await core.make_call(endpoint, model="g", metadata=call_metadata()) # pyright: ignore[reportAttributeAccessIssue] # an old public member + + assert response is answer + assert answer._hidden_params["model_id"] == "a" # pyright: ignore[reportPrivateUsage] # what set_response_headers stamps + assert counters(core, "a", "tpm") == [] + + +@pytest.mark.asyncio +async def test_make_call_raises_the_endpoints_error(make_core: MakeCore) -> None: + core: Final = make_core(model_list=[deployment("g", "a", model="openai/gpt-4o", rpm=5)]) + metadata: Final = call_metadata() + + async def endpoint(**kwargs: object) -> object: # kwargs-ok: the request, as make_call passes it + raise server_error() + + with pytest.raises(litellm.InternalServerError): + await core.make_call(endpoint, model="g", metadata=metadata) # pyright: ignore[reportAttributeAccessIssue] # an old public member + + assert counters(core, "a", "rpm") == [] + assert ROUTER_USAGE_COUNTED_TOKENS_METADATA_KEY not in metadata + + +@pytest.mark.asyncio +async def test_increment_usage_notes_tokens_even_for_a_deployment_it_does_not_track(make_core: MakeCore) -> None: + core: Final = make_core(model_list=[deployment("g", "a", model="openai/gpt-4o")]) + metadata: Final = call_metadata() + + await core.increment_deployment_usage_for_response( # pyright: ignore[reportAttributeAccessIssue] # an old public member + response=completion_response(total_tokens=7), request_kwargs={"model": "g", "metadata": metadata} + ) + + assert metadata[ROUTER_USAGE_COUNTED_TOKENS_METADATA_KEY] == 7 + assert counters(core, "a", "rpm") == [] + + +@pytest.mark.asyncio +async def test_increment_usage_takes_the_note_back_when_the_count_fails( + make_core: MakeCore, monkeypatch: pytest.MonkeyPatch +) -> None: + core: Final = make_core(model_list=[deployment("g", "a", model="openai/gpt-4o", rpm=5)]) + metadata: Final = call_metadata() + + async def refused(**kwargs: object) -> None: # kwargs-ok: the pipeline's arguments + raise ConnectionError("the usage store is down") + + monkeypatch.setattr(core.health().store, "async_increment_cache_pipeline", refused) + + await core.increment_deployment_usage_for_response(completion_response(), {"model": "g", "metadata": metadata}) # pyright: ignore[reportAttributeAccessIssue] # an old public member + + assert ROUTER_USAGE_COUNTED_TOKENS_METADATA_KEY not in metadata + assert counters(core, "a", "rpm") == [] + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "request_kwargs", + [{"model": "g", "metadata": {}}, {"metadata": call_metadata()}, {"model": "g"}], + ids=["no-model-info", "no-model", "no-metadata"], +) +async def test_increment_usage_skips_a_request_with_nothing_to_count_under( + make_core: MakeCore, request_kwargs: dict[str, object] +) -> None: + core: Final = make_core(model_list=[deployment("g", "a", model="openai/gpt-4o", rpm=5)]) + + await core.increment_deployment_usage_for_response(completion_response(), dict(request_kwargs)) # pyright: ignore[reportAttributeAccessIssue] # an old public member + + assert counters(core, "a", "rpm") == [] + + +@pytest.mark.asyncio +async def test_increment_usage_counts_nothing_for_no_response(make_core: MakeCore) -> None: + core: Final = make_core(model_list=[deployment("g", "a", model="openai/gpt-4o", rpm=5)]) + metadata: Final = call_metadata() + + await core.increment_deployment_usage_for_response(None, {"model": "g", "metadata": metadata}) # pyright: ignore[reportAttributeAccessIssue] # an old public member + + assert ROUTER_USAGE_COUNTED_TOKENS_METADATA_KEY not in metadata + + +@pytest.mark.asyncio +async def test_async_deployment_callback_on_failure_counts_one_request(make_core: MakeCore) -> None: + core: Final = make_core(model_list=list(GROUP)) + logged: Final = {"litellm_params": {"metadata": call_metadata(), "model_info": {"id": "a"}}} + + await core.async_deployment_callback_on_failure(logged, None, None, None) # pyright: ignore[reportAttributeAccessIssue] # an old public member + + assert counters(core, "a", "rpm") == [1] + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "logged", + [ + {"call_type": "aretrieve_batch", "litellm_params": {"metadata": call_metadata(), "model_info": {"id": "a"}}}, + {"litellm_params": {"metadata": {"deployment": "x"}, "model_info": {"id": "a"}}}, + {"litellm_params": {"metadata": call_metadata()}}, + ], + ids=["batch-retrieve", "no-group", "no-id"], +) +async def test_async_deployment_callback_on_failure_skips(make_core: MakeCore, logged: dict[str, object]) -> None: + core: Final = make_core(model_list=list(GROUP)) + + await core.async_deployment_callback_on_failure(logged, None, None, None) # pyright: ignore[reportAttributeAccessIssue] # an old public member + + assert counters(core, "a", "rpm") == [] + + +@pytest.mark.asyncio +async def test_async_deployment_callback_on_failure_needs_the_params(make_core: MakeCore) -> None: + core: Final = make_core(model_list=list(GROUP)) + + with pytest.raises(KeyError, match="litellm_params"): + await core.async_deployment_callback_on_failure({}, None, None, None) # pyright: ignore[reportAttributeAccessIssue] # an old public member + + +# --- the plugin context -------------------------------------------------------------------------- + + +def test_deployment_params_merge_in_the_order_the_deployment_receives_them(make_core: MakeCore) -> None: + core: Final = make_core( + model_list=[deployment("g", "a", temperature=0.1, tool_choice="auto", top_p=0.5)], + default_litellm_params={"temperature": 0.9, "tool_choice": "none"}, + ) + target: Final = core.registry().deployment("a") + + params: Final = RoutingContext(core).deployment_params({"top_p": 0.7}, target) # pyright: ignore[reportArgumentType] # the core a context wraps + + assert params["model"] == "openai/a" + assert params["temperature"] == 0.9 + assert params["tool_choice"] == "auto" + assert params["top_p"] == 0.7 + assert RoutingContext(core).deployment_params({}, None).get("temperature") == 0.9 # pyright: ignore[reportArgumentType] # see above + + +def test_update_kwargs_before_fallbacks_writes_the_group_into_the_named_bucket(make_core: MakeCore) -> None: + core: Final = make_core(model_list=list(GROUP), model_group_alias={"al": "g"}) + request: Final[dict[str, object]] = {"litellm_trace_id": "t", "litellm_metadata": {"k": "v"}} + + RoutingContext(core).update_kwargs_before_fallbacks("al", request, "litellm_metadata") # pyright: ignore[reportArgumentType] # see above + + assert request == { + "litellm_trace_id": "t", + "litellm_metadata": {"k": "v", "model_group": "al", "model_group_alias": "al"}, + } + + plain: Final[dict[str, object]] = {} # mutable-ok: the request the core writes into + RoutingContext(core).update_kwargs_before_fallbacks("g", plain) # pyright: ignore[reportArgumentType] # see above + + assert isinstance(plain["litellm_trace_id"], str) + assert plain["metadata"] == {"model_group": "g", "model_group_alias": None} + + +# --- what stays refused -------------------------------------------------------------------------- + + +@pytest.mark.parametrize( + "name", + [ + "async_function_with_fallbacks", + "async_function_with_fallbacks_common_utils", + "async_function_with_retries", + "function_with_fallbacks", + "factory_function", + "model_id_to_deployment_index_map", + ], +) +def test_the_rust_core_refuses_what_it_cannot_serve(name: str) -> None: + from litellm.router_rust.core import RustCoreUnsupported + + core: Final = rust_core(ScriptedHost(), model_list=list(GROUP)) + + with pytest.raises(RustCoreUnsupported): + getattr(core, name) diff --git a/tests/unit/test_router/core/test_parity_cooldown_service_model.py b/tests/unit/test_router/core/test_parity_cooldown_service_model.py new file mode 100644 index 00000000000..4204ebeeca3 --- /dev/null +++ b/tests/unit/test_router/core/test_parity_cooldown_service_model.py @@ -0,0 +1,901 @@ +""" +Parity of the Rust core with ``router_python_core.py`` on three things the logging pipeline and the +fuzzer read: the cooldown event Prometheus counts (``router_cooldown_event_callback``, scheduled by +``set_cooldown_deployments``), the ``ServiceTypes.ROUTER`` service events, and a ``model`` that is no +string, which the native calls cannot take. + +``tests/unit/test_router/core/conftest.py`` skips a test that reaches ``RustCoreUnsupported``. These +members are served, so ``served`` turns that refusal into a failure. + +The scripted host runs a failure callback where the test puts it. ``litellm`` runs it on the loop's +thread for an async call, inline for a sync one, and on a thread with no loop for a stream; in every +case ``CustomLogger.log_event`` swallows what the callback raises. ``logged_like_litellm`` and +``streams_log_off_the_loop`` give the host those two behaviours. +""" + +import asyncio +import contextlib +import functools +import inspect +import re +import time +from collections.abc import Awaitable, Callable, Coroutine, Generator, Mapping +from dataclasses import dataclass, field +from typing import Final, TypeVar + +import pydantic +import pytest + +import litellm +from litellm._service_logger import ServiceLogging +from litellm.constants import DEFAULT_COOLDOWN_TIME_SECONDS +from litellm.integrations.custom_logger import CustomLogger +from litellm.router_core import RouterCore +from litellm.router_utils import cooldown_callbacks, cooldown_handlers +from litellm.types.services import ServiceTypes +from tests.unit.test_router.core.retries_health_helpers import ( + MESSAGES, + api_connection, + attempt, + attempt_sync, + cooled_down, + fresh, + not_found, + rate_limited, + unprocessable, +) +from tests.unit.test_router.core.scripted_host import Call, ScriptedHost, deployment, rate_limit, server_error +from tests.unit.test_router.core.streaming_helpers import HostLogging, chat_stream, collect_texts, collect_texts_sync + +MakeCore = Callable[..., RouterCore] +_Test = TypeVar("_Test", bound=Callable[..., object]) + +ACQUIRE: Final = "async_get_available_deployment" +STRATEGY: Final = ".async_get_available_deployments" + + +def _refused(error: BaseException) -> None: + from litellm.router_rust.core import RustCoreUnsupported + + if isinstance(error, RustCoreUnsupported) or isinstance(error.__cause__, RustCoreUnsupported): + pytest.fail(f"the Rust core stopped serving this: {error}") + + +def served(test: _Test) -> _Test: + """Fail, rather than skip, when the test reaches ``RustCoreUnsupported``.""" + + @functools.wraps(test) + def run(*args: object, **kwargs: object) -> object: # kwargs-ok: the test's fixtures + try: + return test(*args, **kwargs) + except Exception as error: + _refused(error) + raise + + @functools.wraps(test) + async def arun(*args: object, **kwargs: object) -> object: # kwargs-ok: the test's fixtures + try: + return await test(*args, **kwargs) # pyright: ignore[reportGeneralTypeIssues] # an async test + except Exception as error: + _refused(error) + raise + + return arun if inspect.iscoroutinefunction(test) else run # pyright: ignore[reportReturnType] # the test, wrapped + + +async def settle() -> None: + """Let the tasks a call scheduled run.""" + for _ in range(20): + await asyncio.sleep(0) + + +# --- cooldown events ------------------------------------------------------------------------------ + + +@dataclass +class CooldownEvents: + """``router_cooldown_event_callback`` calls, as ``(deployment id, status, cooldown time)``.""" + + fired: list[tuple[str, object, object]] = field(default_factory=list) # mutable-ok: appended as they fire + + +@pytest.fixture +def events(monkeypatch: pytest.MonkeyPatch) -> CooldownEvents: + recorded: Final = CooldownEvents() + + async def record(ctx: object, deployment_id: str, exception_status: object, cooldown_time: object) -> None: + recorded.fired.append((deployment_id, exception_status, cooldown_time)) + + monkeypatch.setattr(cooldown_handlers, "router_cooldown_event_callback", record) + return recorded + + +@pytest.fixture +def logged_like_litellm(monkeypatch: pytest.MonkeyPatch) -> None: + """``CustomLogger.log_event``: what a sync failure callback raises stops at the logging.""" + log_sync: Final = ScriptedHost.log_sync + + def swallowed(host: ScriptedHost, call: Call, error: BaseException | None, response: object) -> None: + with contextlib.suppress(RuntimeError): + log_sync(host, call, error, response) + + monkeypatch.setattr(ScriptedHost, "log_sync", swallowed) + + +@pytest.fixture +def streams_log_off_the_loop(monkeypatch: pytest.MonkeyPatch, logged_like_litellm: None) -> None: + """``dispatch_failure_handlers``: a stream's sync callbacks run in ``executor.submit``, off the loop.""" + + def dispatch(logging: HostLogging, exception: Exception, _traceback: str) -> Coroutine[None, None, None]: + async def run() -> None: + await asyncio.to_thread(logging.failure_handler, exception) + + return run() + + monkeypatch.setattr(HostLogging, "dispatch_attempt_failure_handlers", dispatch) + + +def fired_exactly(events: CooldownEvents) -> list[str]: # mutable-ok: a fresh list + """The events as written: ``5 == 5.0``, and the cooldown's type is part of the event.""" + return [repr(fired) for fired in events.fired] + + +def lone(**arguments: object) -> Mapping[str, object]: # kwargs-ok: the router's constructor arguments + return {"model_list": [deployment("g", "a")], "num_retries": 0, **arguments} + + +def ordered(**arguments: object) -> Mapping[str, object]: # kwargs-ok: the router's constructor arguments + return { + "model_list": [deployment("g", "a", rpm=100, order=1), deployment("g", "b", order=2)], + "num_retries": 0, + "allowed_fails": 0, + **arguments, + } + + +@served +async def test_a_failed_attempt_fires_one_event_with_the_status_and_the_cooldown( + make_core: MakeCore, host: ScriptedHost, events: CooldownEvents +) -> None: + core: Final = make_core(**lone(cooldown_time=7)) + host.answer("openai/a", fresh(not_found)) + + assert isinstance(await attempt(core), litellm.NotFoundError) + await settle() + + assert fired_exactly(events) == [repr(("a", 404, 7))] + + +@served +async def test_a_second_call_refused_by_the_cooldown_fires_nothing_more( + make_core: MakeCore, host: ScriptedHost, events: CooldownEvents +) -> None: + core: Final = make_core(**lone()) + host.answer("openai/a", fresh(not_found)) + + _ = await attempt(core) + assert isinstance(await attempt(core), litellm.RouterRateLimitError) + await settle() + + assert [fired[0] for fired in events.fired] == ["a"] + + +@served +async def test_each_retry_that_cools_down_fires_its_own_event( + make_core: MakeCore, host: ScriptedHost, events: CooldownEvents +) -> None: + core: Final = make_core( + model_list=[deployment("g", "a"), deployment("g", "b")], num_retries=1, retry_after=0, allowed_fails=0 + ) + host.answer("*", fresh(server_error)) + + _ = await attempt(core) + await settle() + + assert sorted(fired[0] for fired in events.fired) == ["a", "b"] + + +@pytest.mark.parametrize("builder", [unprocessable, api_connection]) +@served +async def test_a_failure_that_sets_no_cooldown_fires_nothing( + make_core: MakeCore, host: ScriptedHost, events: CooldownEvents, builder: Callable[[], Exception] +) -> None: + core: Final = make_core(**lone()) + host.answer("openai/a", fresh(builder)) + + _ = await attempt(core) + await settle() + + assert cooled_down(core, ["a"]) == {} + assert events.fired == [] + + +@served +async def test_disabled_cooldowns_fire_nothing(make_core: MakeCore, host: ScriptedHost, events: CooldownEvents) -> None: + core: Final = make_core(**lone(disable_cooldowns=True)) + host.answer("openai/a", fresh(not_found)) + + _ = await attempt(core) + await settle() + + assert cooled_down(core, ["a"]) == {} + assert events.fired == [] + + +@served +async def test_a_sync_call_on_a_running_loop_fires_on_that_loop( + make_core: MakeCore, host: ScriptedHost, events: CooldownEvents +) -> None: + core: Final = make_core(**lone()) + host.answer("openai/a", fresh(not_found)) + + assert isinstance(attempt_sync(core), litellm.NotFoundError) + await settle() + + assert fired_exactly(events) == [repr(("a", 404, DEFAULT_COOLDOWN_TIME_SECONDS))] + + +@served +def test_a_sync_call_with_no_loop_cools_down_and_fires_on_its_own_loop( + make_core: MakeCore, host: ScriptedHost, events: CooldownEvents, logged_like_litellm: None +) -> None: + """``function_with_fallbacks`` runs the sync walk on a private loop, which the event's task lands on.""" + core: Final = make_core(**lone()) + host.answer("openai/a", fresh(not_found)) + + assert isinstance(attempt_sync(core), litellm.NotFoundError) + + assert set(cooled_down(core, ["a"])) == {"a"} + assert fired_exactly(events) == [repr(("a", 404, DEFAULT_COOLDOWN_TIME_SECONDS))] + + +SYNC_ENDPOINTS: Final[Mapping[str, Callable[[RouterCore], object]]] = { + "completion": lambda core: core.completion(model="g", messages=MESSAGES), # pyright: ignore[reportAttributeAccessIssue, reportUnknownMemberType, reportUnknownLambdaType] # endpoint method + "embedding": lambda core: core.embedding(model="g", input=["hi"]), # pyright: ignore[reportAttributeAccessIssue, reportUnknownMemberType, reportUnknownLambdaType] # endpoint method + "image_generation": lambda core: core.image_generation(model="g", prompt="hi"), # pyright: ignore[reportAttributeAccessIssue, reportUnknownMemberType, reportUnknownLambdaType] # endpoint method +} + + +@served +@pytest.mark.parametrize("endpoint", sorted(SYNC_ENDPOINTS)) +def test_every_sync_endpoint_fires_the_default_cooldown_as_the_int_it_is( + make_core: MakeCore, host: ScriptedHost, events: CooldownEvents, logged_like_litellm: None, endpoint: str +) -> None: + host.answer("openai/a", fresh(not_found)) + + with pytest.raises(litellm.NotFoundError, match="no such model"): + _ = SYNC_ENDPOINTS[endpoint](make_core(**lone())) + + assert fired_exactly(events) == [repr(("a", 404, DEFAULT_COOLDOWN_TIME_SECONDS))] + + +@served +@pytest.mark.parametrize( + ("arguments", "request_fields", "expected"), + [ + pytest.param({"cooldown_time": 7}, {}, 7, id="router-int"), + pytest.param({"cooldown_time": 7.0}, {}, 7.0, id="router-whole-float"), + pytest.param({"cooldown_time": 7.5}, {}, 7.5, id="router-float"), + pytest.param({"cooldown_time": 0}, {}, DEFAULT_COOLDOWN_TIME_SECONDS, id="router-zero-is-unset"), + pytest.param({"model_list": [deployment("g", "a", cooldown_time=7)]}, {}, 7, id="deployment-int"), + pytest.param({"model_list": [deployment("g", "a", cooldown_time=7.0)]}, {}, 7.0, id="deployment-float"), + pytest.param( + {"cooldown_time": 7, "model_list": [deployment("g", "a", cooldown_time=7.0)]}, + {}, + 7.0, + id="deployment-float-over-router-int", + ), + pytest.param( + {"cooldown_time": 7.0, "model_list": [deployment("g", "a", cooldown_time=7)]}, + {}, + 7, + id="deployment-int-over-router-float", + ), + pytest.param({}, {"cooldown_time": 7}, 7, id="request-int"), + pytest.param({}, {"cooldown_time": 7.0}, 7.0, id="request-float"), + pytest.param( + {"model_list": [deployment("g", "a", cooldown_time=7.0)]}, + {"cooldown_time": 7}, + 7, + id="request-int-over-deployment-float", + ), + pytest.param( + {"model_list": [deployment("g", "a", model_info={"cooldown_time": 7})]}, + {"cooldown_time": 7.0}, + 7, + id="model-info-int-over-request-float", + ), + ], +) +def test_the_event_carries_the_number_python_chose_as_it_was_written( + make_core: MakeCore, + host: ScriptedHost, + events: CooldownEvents, + logged_like_litellm: None, + arguments: Mapping[str, object], + request_fields: Mapping[str, object], + expected: float, +) -> None: + host.answer("openai/a", fresh(not_found)) + + assert isinstance(attempt_sync(make_core(**lone(**arguments)), **request_fields), litellm.NotFoundError) + + assert fired_exactly(events) == [repr(("a", 404, expected))] + + +@served +@pytest.mark.parametrize( + ("arguments", "expected"), + [ + pytest.param({}, 9, id="over-the-default"), + pytest.param({"cooldown_time": 9.0}, 9, id="an-int-over-an-equal-router-float"), + ], +) +def test_a_retry_after_header_fires_as_the_int_python_reads( + make_core: MakeCore, + host: ScriptedHost, + events: CooldownEvents, + logged_like_litellm: None, + arguments: Mapping[str, object], + expected: int, +) -> None: + for name in ("a", "b"): + host.answer(f"openai/{name}", fresh(lambda: rate_limited("9"))) + core: Final = make_core( + model_list=[deployment("g", "a"), deployment("g", "b")], num_retries=0, allowed_fails=0, **arguments + ) + + assert isinstance(attempt_sync(core), litellm.RateLimitError) + + assert {repr(fired[1:]) for fired in events.fired} == {repr((429, expected))} + + +@served +async def test_an_async_stream_failure_cools_down_and_fires_nothing( + make_core: MakeCore, host: ScriptedHost, events: CooldownEvents, streams_log_off_the_loop: None +) -> None: + host.answer("openai/a", chat_stream("partial", error=rate_limit(message="mid stream"), logged=True)) + core: Final = make_core(**ordered()) + + stream: Final = await core.acompletion(model="g", messages=MESSAGES, stream=True) # pyright: ignore[reportAttributeAccessIssue, reportUnknownMemberType, reportUnknownVariableType] # endpoint method + with pytest.raises(litellm.RateLimitError, match="mid stream"): + await collect_texts(stream) + await asyncio.sleep(0.05) + await settle() + + assert set(cooled_down(core, ["a", "b"])) == {"a"} + assert events.fired == [] + + +@served +async def test_a_sync_stream_failure_on_a_running_loop_fires_nothing( + make_core: MakeCore, host: ScriptedHost, events: CooldownEvents, logged_like_litellm: None +) -> None: + host.answer("openai/a", chat_stream("partial", error=rate_limit(message="mid stream"), logged=True, sync=True)) + core: Final = make_core(**ordered()) + + stream: Final = core.completion(model="g", messages=MESSAGES, stream=True) # pyright: ignore[reportAttributeAccessIssue, reportUnknownMemberType, reportUnknownVariableType] # endpoint method + with pytest.raises(litellm.RateLimitError, match="mid stream"): + collect_texts_sync(stream) + await asyncio.sleep(0.05) + await settle() + + assert events.fired == [] + + +class _RefusingCheck(CustomLogger): + async def async_pre_call_check(self, deployment: dict, parent_otel_span: object) -> dict | None: # pyright: ignore[reportIncompatibleMethodOverride, reportMissingTypeArgument, reportUnknownParameterType] # CustomLogger's own signature + raise rate_limit(message="refused") + + +@served +async def test_a_refusing_pre_call_check_fires_the_cooldown_it_sets( + make_core: MakeCore, host: ScriptedHost, events: CooldownEvents +) -> None: + host.lists["callbacks"].append(_RefusingCheck()) + core: Final = make_core( + model_list=[deployment("g", "a"), deployment("g", "b")], num_retries=0, allowed_fails=0, cooldown_time=7 + ) + + assert isinstance(await attempt(core), litellm.RateLimitError) + await settle() + + assert len(events.fired) == 1 + assert repr(events.fired[0][1:]) == repr((429, 7)) + assert events.fired[0][0] in cooled_down(core, ["a", "b"]) + + +@served +@pytest.mark.parametrize( + ("params", "cooldown_time", "expected"), + [ + pytest.param({}, 7.0, 7.0, id="router-float"), + pytest.param({}, None, DEFAULT_COOLDOWN_TIME_SECONDS, id="default"), + pytest.param({"cooldown_time": 9}, 7, 7, id="the-deployment-is-not-read"), + ], +) +async def test_a_refusing_pre_call_check_fires_the_routers_own_cooldown( + make_core: MakeCore, + host: ScriptedHost, + events: CooldownEvents, + params: Mapping[str, object], + cooldown_time: float | None, + expected: float, +) -> None: + host.lists["callbacks"].append(_RefusingCheck()) + core: Final = make_core( + model_list=[deployment("g", "a", **params), deployment("g", "b", **params)], + num_retries=0, + allowed_fails=0, + cooldown_time=cooldown_time, + ) + + assert isinstance(await attempt(core), litellm.RateLimitError) + await settle() + + assert [repr(fired[1:]) for fired in events.fired] == [repr((429, expected))] + + +@served +async def test_a_recorded_failure_fires_with_the_status_it_was_given( + make_core: MakeCore, events: CooldownEvents +) -> None: + core: Final = make_core(model_list=[deployment("g", "a"), deployment("g", "b")], allowed_fails=0) + + assert core.health().record_failure("a", server_error(), 500) is True + await settle() + + assert fired_exactly(events) == [repr(("a", 500, DEFAULT_COOLDOWN_TIME_SECONDS))] + + +@served +@pytest.mark.parametrize( + ("params", "cooldown_time", "error", "expected"), + [ + pytest.param({}, 7, server_error, 7, id="router-int"), + pytest.param({}, 7.0, server_error, 7.0, id="router-float"), + pytest.param({"cooldown_time": 9}, 7.0, server_error, 7.0, id="the-deployment-is-not-read"), + pytest.param({}, 7.0, lambda: rate_limited("9"), 7.0, id="the-header-is-not-read"), + ], +) +async def test_a_recorded_failure_fires_the_routers_own_cooldown( + make_core: MakeCore, + events: CooldownEvents, + params: Mapping[str, object], + cooldown_time: float, + error: Callable[[], Exception], + expected: float, +) -> None: + core: Final = make_core( + model_list=[deployment("g", "a", **params), deployment("g", "b", **params)], + allowed_fails=0, + cooldown_time=cooldown_time, + ) + + assert core.health().record_failure("a", error(), 500) is True + await settle() + + assert fired_exactly(events) == [repr(("a", 500, expected))] + + +@served +@pytest.mark.parametrize( + ("model_info", "litellm_params", "cooldown_time", "error", "expected"), + [ + pytest.param({}, {}, 7, server_error, 7, id="router-int"), + pytest.param({}, {}, 7.0, server_error, 7.0, id="router-float"), + pytest.param({}, {"cooldown_time": 7}, 7.0, server_error, 7, id="params-int-over-router-float"), + pytest.param({}, {"cooldown_time": 7.0}, 7, server_error, 7.0, id="params-float-over-router-int"), + pytest.param({"cooldown_time": 7}, {"cooldown_time": 7.0}, None, server_error, 7, id="info-over-params"), + pytest.param({"cooldown_time": -1}, {"cooldown_time": 7}, 7.0, server_error, 7.0, id="a-negative-is-unset"), + pytest.param({}, {}, 9.0, lambda: rate_limited("9"), 9, id="header-int-over-router-float"), + pytest.param({}, {"cooldown_time": 9.0}, None, lambda: rate_limited("9"), 9.0, id="params-over-header"), + ], +) +async def test_a_logged_failure_fires_the_callbacks_precedence( + make_core: MakeCore, + events: CooldownEvents, + model_info: Mapping[str, object], + litellm_params: Mapping[str, object], + cooldown_time: float | None, + error: Callable[[], Exception], + expected: float, +) -> None: + core: Final = make_core( + model_list=[deployment("g", "a"), deployment("g", "b")], allowed_fails=0, cooldown_time=cooldown_time + ) + logged: Final = { + "exception": error(), + "litellm_params": { + "metadata": {"model_group": "g"}, + "model_info": {"id": "a", **model_info}, + **litellm_params, + }, + } + + assert core.deployment_callback_on_failure(logged, None, 0.0, 0.0) is True # pyright: ignore[reportAttributeAccessIssue, reportUnknownMemberType] # the failure callback + await settle() + + assert [repr(fired[2]) for fired in events.fired] == [repr(expected)] + + +@served +def test_a_recorded_failure_with_no_loop_raises_and_fires_nothing(make_core: MakeCore, events: CooldownEvents) -> None: + core: Final = make_core(model_list=[deployment("g", "a"), deployment("g", "b")], allowed_fails=0) + + with pytest.raises(RuntimeError, match="no running event loop"): + core.health().record_failure("a", server_error(), 500) + + assert events.fired == [] + + +class _Prometheus: + def __init__(self) -> None: + self.cooled: Final[list[Mapping[str, object]]] = [] # mutable-ok: appended as the event runs + self.outages: Final[list[Mapping[str, object]]] = [] # mutable-ok: same + + def set_deployment_complete_outage(self, **labels: object) -> None: # kwargs-ok: the metric's labels + self.outages.append(labels) + + def increment_deployment_cooled_down(self, **labels: object) -> None: # kwargs-ok: the metric's labels + self.cooled.append(labels) + + +@served +async def test_the_event_reaches_prometheus_with_the_deployments_labels( + make_core: MakeCore, host: ScriptedHost, monkeypatch: pytest.MonkeyPatch +) -> None: + prometheus: Final = _Prometheus() + monkeypatch.setattr(cooldown_callbacks, "get_prometheus_logger_from_callbacks", lambda: prometheus) + core: Final = make_core(**lone()) + host.answer("openai/a", fresh(not_found)) + + _ = await attempt(core) + await settle() + + assert [(labels["model_id"], labels["exception_status"]) for labels in prometheus.cooled] == [("a", "404")] + assert [labels["model_id"] for labels in prometheus.outages] == ["a"] + + +# --- service events ------------------------------------------------------------------------------- + + +class _ServiceEvents: + """Stands in for ``service_logger_obj``.""" + + def __init__(self) -> None: + self.seen: Final[list[Mapping[str, object]]] = [] # mutable-ok: appended as they run + + async def async_service_success_hook(self, **event: object) -> None: # kwargs-ok: the hook's arguments + self.seen.append(event) + + def call_types(self) -> list[object]: # mutable-ok: a fresh list + return [event["call_type"] for event in self.seen] + + +def service_events(core: RouterCore) -> _ServiceEvents: + recorded: Final = _ServiceEvents() + core.service_logger_obj = recorded # pyright: ignore[reportAttributeAccessIssue] # the attribute both cores keep + return recorded + + +@served +def test_the_core_serves_its_service_logger(make_core: MakeCore) -> None: + core: Final = make_core(**lone()) + + assert isinstance(core.service_logger_obj, ServiceLogging) # pyright: ignore[reportAttributeAccessIssue] # see service_events + + +@served +async def test_a_served_acompletion_logs_the_pick_and_the_call(make_core: MakeCore) -> None: + core: Final = make_core(**lone()) + recorded: Final = service_events(core) + span: Final = object() + + _ = await attempt(core, metadata={"litellm_parent_otel_span": span}) + await settle() + + assert sorted(map(str, recorded.call_types())) == sorted([ACQUIRE, "acompletion"]) + for event in recorded.seen: + assert event["service"] is ServiceTypes.ROUTER + assert event["parent_otel_span"] is span + duration: object = event["duration"] + start_time: object = event["start_time"] + end_time: object = event["end_time"] + assert isinstance(duration, float) and isinstance(start_time, float) and isinstance(end_time, float) + assert duration == pytest.approx(end_time - start_time) + assert abs(end_time - time.time()) < 60 + + +@served +async def test_a_retry_logs_a_pick_per_attempt_and_one_call(make_core: MakeCore, host: ScriptedHost) -> None: + core: Final = make_core(model_list=[deployment("g", "a"), deployment("g", "b")], num_retries=1, retry_after=0) + host.answer("*", server_error(), None) + recorded: Final = service_events(core) + + _ = await attempt(core) + await settle() + + assert recorded.call_types().count(ACQUIRE) == 2 + assert recorded.call_types().count("acompletion") == 1 + + +@served +async def test_a_failed_acompletion_logs_the_pick_alone(make_core: MakeCore, host: ScriptedHost) -> None: + core: Final = make_core(**lone()) + host.answer("openai/a", fresh(not_found)) + recorded: Final = service_events(core) + + _ = await attempt(core) + await settle() + + assert recorded.call_types() == [ACQUIRE] + + +@served +async def test_a_refused_pick_logs_nothing(make_core: MakeCore) -> None: + core: Final = make_core(**lone()) + recorded: Final = service_events(core) + + assert isinstance(await attempt(core, model="nope"), litellm.BadRequestError) + await settle() + + assert recorded.call_types() == [] + + +@served +async def test_a_strategy_pick_logs_its_own_event(make_core: MakeCore) -> None: + core: Final = make_core( + model_list=[deployment("g", "a"), deployment("g", "b")], num_retries=0, routing_strategy="least-busy" + ) + recorded: Final = service_events(core) + + _ = await attempt(core) + await settle() + + assert sorted(map(str, recorded.call_types())) == sorted([STRATEGY, ACQUIRE, "acompletion"]) + + +@served +async def test_another_endpoint_logs_only_the_strategy_pick(make_core: MakeCore) -> None: + core: Final = make_core( + model_list=[deployment("g", "a"), deployment("g", "b")], num_retries=0, routing_strategy="least-busy" + ) + recorded: Final = service_events(core) + + _ = await core.aembedding(model="g", input=["hi"]) # pyright: ignore[reportAttributeAccessIssue, reportUnknownMemberType] # endpoint method + await settle() + + assert recorded.call_types() == [STRATEGY] + + +@served +async def test_a_shuffled_endpoint_logs_nothing(make_core: MakeCore) -> None: + core: Final = make_core(**lone()) + recorded: Final = service_events(core) + + _ = await core.aembedding(model="g", input=["hi"]) # pyright: ignore[reportAttributeAccessIssue, reportUnknownMemberType] # endpoint method + await settle() + + assert recorded.call_types() == [] + + +@served +async def test_a_sync_call_logs_nothing(make_core: MakeCore) -> None: + core: Final = make_core(**lone()) + recorded: Final = service_events(core) + + _ = attempt_sync(core) + await settle() + + assert recorded.call_types() == [] + + +@served +async def test_a_pass_through_strategy_pick_runs_on_the_perf_counter(make_core: MakeCore) -> None: + core: Final = make_core( + model_list=[ + deployment("g", "a", use_in_pass_through=True), + deployment("g", "b", use_in_pass_through=True), + ], + num_retries=0, + routing_strategy="least-busy", + ) + recorded: Final = service_events(core) + + _ = await core.async_get_available_deployment_for_pass_through(model="g", request_kwargs={}) # pyright: ignore[reportAttributeAccessIssue, reportUnknownMemberType] # the legacy member + await settle() + + assert recorded.call_types() == [STRATEGY] + start_time: object = recorded.seen[0]["start_time"] + assert isinstance(start_time, float) and abs(start_time - time.perf_counter()) < 60 + + +# --- a model that is no string -------------------------------------------------------------------- + +Endpoint = Callable[[RouterCore, object], Awaitable[object]] + + +def _call(name: str, **request: object) -> Endpoint: # kwargs-ok: the endpoint's request + def call(core: RouterCore, model: object) -> Awaitable[object]: + answer: Final[object] = getattr(core, name)(model=model, **request) # pyright: ignore[reportAny] # an endpoint + + async def awaited() -> object: + return await answer if inspect.isawaitable(answer) else answer + + return awaited() + + return call + + +ENDPOINTS: Final = { + "acompletion": _call("acompletion", messages=MESSAGES), + "atext_completion": _call("atext_completion", prompt="hi"), + "aembedding": _call("aembedding", input=["hi"]), + "aimage_generation": _call("aimage_generation", prompt="hi"), + "aresponses": _call("aresponses", input="hi"), + "aanthropic_messages": _call("aanthropic_messages", messages=MESSAGES, max_tokens=5), + "completion": _call("completion", messages=MESSAGES), + "embedding": _call("embedding", input=["hi"]), + "image_generation": _call("image_generation", prompt="hi"), +} + + +def search_error(model: object) -> str: + """What ``"/" in model`` raises (``declared_provider``).""" + with pytest.raises(TypeError) as raised: + _ = "/" in model # pyright: ignore[reportOperatorIssue] # raising is the point + return str(raised.value) + + +def hash_error(lookup: Callable[[], object]) -> str: + with pytest.raises(TypeError) as raised: + _ = lookup() + return str(raised.value) + + +SEARCH_FAILS: Final = (b"bytes", 1, 1.5, True) +ROUTES_NOWHERE: Final = (0, 0.0, False, b"", (), ("a",)) + + +@pytest.mark.parametrize("value", SEARCH_FAILS, ids=repr) +@pytest.mark.parametrize("endpoint", sorted(ENDPOINTS)) +@served +async def test_a_model_no_search_can_read_raises_that_search(make_core: MakeCore, endpoint: str, value: object) -> None: + core: Final = make_core(**lone()) + + with pytest.raises(TypeError, match=f"^{re.escape(search_error(value))}$"): + await ENDPOINTS[endpoint](core, value) + + +@pytest.mark.parametrize( + "mock", + [ + {"mock_testing_fallbacks": True}, + {"mock_testing_rate_limit_error": True}, + {"mock_response": "litellm.RateLimitError"}, + ], + ids=["fallbacks", "rate_limit", "response"], +) +@served +async def test_bytes_raise_before_any_mock(make_core: MakeCore, mock: Mapping[str, object]) -> None: + core: Final = make_core(**lone()) + + with pytest.raises(TypeError, match=r"^a bytes-like object is required, not 'str'$"): + await core.aimage_generation(model=b"bytes", prompt="hi", **mock) # pyright: ignore[reportAttributeAccessIssue, reportUnknownMemberType] # endpoint method + + +@pytest.mark.parametrize("fallbacks", [None, [{"*": ["other"]}]], ids=["plain", "default_fallback"]) +@served +async def test_bytes_raise_whatever_the_fallbacks(make_core: MakeCore, fallbacks: object) -> None: + core: Final = make_core(**lone(fallbacks=fallbacks)) + + with pytest.raises(TypeError, match=r"^a bytes-like object is required, not 'str'$"): + await core.acompletion(model=b"bytes", messages=MESSAGES) # pyright: ignore[reportAttributeAccessIssue, reportUnknownMemberType] # endpoint method + + +@pytest.mark.parametrize("value", [["g"], {"a": 1}], ids=["list", "dict"]) +@pytest.mark.parametrize("endpoint", sorted(ENDPOINTS)) +@served +async def test_an_unhashable_model_raises_as_the_alias_lookup_raises( + make_core: MakeCore, endpoint: str, value: object +) -> None: + core: Final = make_core(**lone()) + expected: Final = hash_error(lambda: {}.get(value)) # pyright: ignore[reportUnknownLambdaType, reportArgumentType, reportUnknownMemberType, reportCallIssue] # raising is the point + + with pytest.raises(TypeError, match=f"^{re.escape(expected)}$"): + await ENDPOINTS[endpoint](core, value) + + +@served +def test_a_sync_pick_reads_the_names_first(make_core: MakeCore) -> None: + core: Final = make_core(**lone()) + expected: Final = hash_error(lambda: ["g"] in set[object]()) # pyright: ignore[reportUnhashable] # raising is the point + + with pytest.raises(TypeError, match=f"^{re.escape(expected)}$"): + core.text_completion(model=["g"], prompt="hi") # pyright: ignore[reportAttributeAccessIssue, reportUnknownMemberType] # endpoint method + + +@pytest.mark.parametrize("value", ROUTES_NOWHERE, ids=repr) +@pytest.mark.parametrize("endpoint", sorted(ENDPOINTS)) +@served +async def test_a_model_the_lookups_accept_is_a_name_no_deployment_serves( + make_core: MakeCore, endpoint: str, value: object +) -> None: + core: Final = make_core(**lone()) + + with pytest.raises(litellm.BadRequestError, match=re.escape(f"You passed in model={value}. There are no healthy")): + await ENDPOINTS[endpoint](core, value) + + +@pytest.mark.parametrize("value", ROUTES_NOWHERE, ids=repr) +@served +async def test_a_mock_failure_names_the_model_as_given(make_core: MakeCore, value: object) -> None: + core: Final = make_core(**lone()) + + with pytest.raises(litellm.InternalServerError, match=re.escape(f"This is a mock exception for model={value},")): + await core.aimage_generation(model=value, prompt="hi", mock_testing_fallbacks=True) # pyright: ignore[reportAttributeAccessIssue, reportUnknownMemberType] # endpoint method + with pytest.raises(litellm.RateLimitError, match=re.escape(f"This is a mock exception for model={value},")): + await core.acompletion(model=value, messages=MESSAGES, mock_testing_rate_limit_error=True) # pyright: ignore[reportAttributeAccessIssue, reportUnknownMemberType] # endpoint method + + +@pytest.mark.parametrize("value", ROUTES_NOWHERE, ids=repr) +@served +async def test_a_default_fallback_serves_such_a_model(make_core: MakeCore, host: ScriptedHost, value: object) -> None: + core: Final = make_core(**lone(fallbacks=[{"*": ["g"]}])) + + _ = await core.acompletion(model=value, messages=MESSAGES) # pyright: ignore[reportAttributeAccessIssue, reportUnknownMemberType] # endpoint method + + assert host.deployments_called("acompletion") == ["a"] + + +@pytest.mark.parametrize("value", [0, b"", ()], ids=repr) +@served +async def test_a_catch_all_wildcard_serves_such_a_model(make_core: MakeCore, host: ScriptedHost, value: object) -> None: + core: Final = make_core(model_list=[deployment("*", "w")], num_retries=0) + + _ = await core.acompletion(model=value, messages=MESSAGES) # pyright: ignore[reportAttributeAccessIssue, reportUnknownMemberType] # endpoint method + + assert host.deployments_called("acompletion") == ["w"] + + +@pytest.mark.parametrize("kind", ["pick", "healthy", "pass_through"]) +@served +async def test_a_selection_raises_as_a_call_does(make_core: MakeCore, kind: str) -> None: + core: Final = make_core(**lone()) + + with pytest.raises(TypeError, match=r"^a bytes-like object is required, not 'str'$"): + await core.select(b"bytes", kind=kind) # pyright: ignore[reportArgumentType] # raising is the point + with pytest.raises(litellm.BadRequestError, match=r"You passed in model=0\. There are no healthy"): + await core.select(0, kind=kind) # pyright: ignore[reportArgumentType] # see above + + +@served +def test_the_sync_pick_and_text_completion_raise_as_a_call_does(make_core: MakeCore) -> None: + core: Final = make_core(**lone()) + + with pytest.raises(TypeError, match=r"^a bytes-like object is required, not 'str'$"): + core.get_available_deployment(model=b"bytes") # pyright: ignore[reportAttributeAccessIssue, reportUnknownMemberType] # the legacy member + with pytest.raises(TypeError, match=r"^a bytes-like object is required, not 'str'$"): + core.text_completion(model=b"bytes", prompt="hi") # pyright: ignore[reportAttributeAccessIssue, reportUnknownMemberType] # endpoint method + with pytest.raises(litellm.BadRequestError, match=r"You passed in model=0\. There are no healthy"): + core.text_completion(model=0, prompt="hi") # pyright: ignore[reportAttributeAccessIssue, reportUnknownMemberType] # endpoint method + + +@served +async def test_moderation_raises_or_goes_straight_to_the_host(make_core: MakeCore, host: ScriptedHost) -> None: + core: Final = make_core(**lone()) + + with pytest.raises(TypeError, match=r"^a bytes-like object is required, not 'str'$"): + await core.amoderation(model=b"bytes", input="hi") # pyright: ignore[reportAttributeAccessIssue, reportUnknownMemberType] # endpoint method + _ = await core.amoderation(model=("a",), input="hi") # pyright: ignore[reportAttributeAccessIssue, reportUnknownMemberType] # endpoint method + + assert len(host.calls_to("amoderation")) == 1 + + +@served +async def test_a_prioritized_call_fails_building_its_queue_item(make_core: MakeCore) -> None: + core: Final = make_core(**lone()) + + with pytest.raises(pydantic.ValidationError, match="model_name"): + await core.acompletion(model=0, messages=MESSAGES, priority=1) # pyright: ignore[reportAttributeAccessIssue, reportUnknownMemberType] # endpoint method diff --git a/tests/unit/test_router/core/test_parity_decisions_routing.py b/tests/unit/test_router/core/test_parity_decisions_routing.py new file mode 100644 index 00000000000..7a939baa833 --- /dev/null +++ b/tests/unit/test_router/core/test_parity_decisions_routing.py @@ -0,0 +1,235 @@ +""" +Parity of the Rust core with ``router_python_core.py`` on routing a Decisions call: a group's +deployments whose provider has no Decisions config are dropped before any later stage, unless that +leaves none (``_without_deployments_that_cannot_serve_route``), and a Decisions or SystemOne call +carrying audio or videos skips the deployments that refuse them, unless that leaves none +(``_without_decisions_deployments_refusing_media``). + +The weights put the deployment that must be skipped first, so a core that does not narrow the +group picks it every time. +""" + +import asyncio +from collections.abc import Callable, Coroutine, Mapping +from typing import Final, cast + +import pytest + +import litellm +from litellm.router_core import RouterCore +from litellm.types.router import RouterRateLimitError +from tests.unit.test_router.core.endpoints_helpers import EndpointHost, patch_direct_endpoints +from tests.unit.test_router.core.scripted_host import auth_error, deployment + +MakeCore = Callable[..., RouterCore] + +QUESTION: Final[Mapping[str, object]] = { + "input": "Is it?", + "questions": [{"type": "predicate", "name": "is_it", "instructions": "Is it true?"}], +} +CLIP: Final = [{"url": "https://media.example/clip"}] + + +@pytest.fixture +def host(monkeypatch: pytest.MonkeyPatch) -> EndpointHost: + """``litellm.decisions`` and ``litellm.systemone`` are bound directly, so they answer from this host.""" + endpoint_host: Final = EndpointHost() + patch_direct_endpoints(monkeypatch, endpoint_host) + return endpoint_host + + +def _sync(core: object, name: str) -> Callable[..., object]: + """An endpoint both cores serve outside the ``RouterCore`` protocol.""" + return cast("Callable[..., object]", getattr(core, name)) # cast-ok: the member is a plain function + + +def _async(core: object, name: str) -> Callable[..., Coroutine[object, object, object]]: + """``_sync`` for a coroutine endpoint.""" + return cast( + "Callable[..., Coroutine[object, object, object]]", getattr(core, name) + ) # cast-ok: the member is a coroutine function + + +def _call(core: RouterCore, op: str, **request: object) -> None: # kwargs-ok: the endpoint's request + """One call on ``op``: an ``a``-prefixed op is awaited, the other one called.""" + if op.startswith("a"): + _ = asyncio.run(_async(core, op)(**request)) + else: + _ = _sync(core, op)(**request) + + +# --- deployments whose provider has no Decisions config ------------------------------------------- + +_OPENAI: Final = deployment("e", "p7-openai", model="openai/gpt-4o", weight=0) +_ANTHROPIC: Final = deployment("e", "p7-anthropic", model="anthropic/claude-sonnet-4-5-20250929", weight=1) + + +def test_adecisions_skips_a_heavier_deployment_whose_provider_has_no_decisions_config( + make_core: MakeCore, host: EndpointHost +) -> None: + core: Final = make_core(model_list=[_OPENAI, _ANTHROPIC], num_retries=0) + + for _ in range(5): + _call(core, "adecisions", model="e", **QUESTION) + + assert host.deployments_called("adecisions") == ["p7-openai"] * 5 + + +def test_sync_decisions_skips_a_heavier_deployment_whose_provider_has_no_decisions_config( + make_core: MakeCore, host: EndpointHost +) -> None: + core: Final = make_core(model_list=[_OPENAI, _ANTHROPIC], num_retries=0) + + for _ in range(5): + _call(core, "decisions", model="e", **QUESTION) + + assert host.deployments_called("decisions") == ["p7-openai"] * 5 + + +@pytest.mark.parametrize("op", ["adecisions", "decisions"]) +def test_a_group_no_provider_of_which_serves_decisions_still_picks( + make_core: MakeCore, host: EndpointHost, op: str +) -> None: + other: Final = deployment("e", "p7-anthropic-2", model="anthropic/claude-opus-4-1", weight=0) + core: Final = make_core(model_list=[_ANTHROPIC, other], num_retries=0) + + _call(core, op, model="e", **QUESTION) + + assert host.deployments_called(op) == ["p7-anthropic"] + + +@pytest.mark.parametrize("op", ["adecisions", "decisions"]) +def test_a_decisions_pick_by_id_is_not_narrowed(make_core: MakeCore, host: EndpointHost, op: str) -> None: + core: Final = make_core(model_list=[_OPENAI, _ANTHROPIC], num_retries=0) + + _call(core, op, model="p7-anthropic", **QUESTION) + + assert host.deployments_called(op) == ["p7-anthropic"] + + +@pytest.mark.parametrize("op", ["adecisions", "decisions"]) +def test_the_route_filter_runs_before_order(make_core: MakeCore, host: EndpointHost, op: str) -> None: + first: Final = deployment("e", "p7-anthropic", model="anthropic/claude-sonnet-4-5-20250929", order=1) + second: Final = deployment("e", "p7-openai", model="openai/gpt-4o", order=2) + core: Final = make_core(model_list=[first, second], num_retries=0) + + _call(core, op, model="e", **QUESTION) + + assert host.deployments_called(op) == ["p7-openai"] + + +def test_the_route_filter_runs_before_cooldown(make_core: MakeCore, host: EndpointHost) -> None: + """With the only Decisions deployment cooling down, the group has nothing to serve: the + deployment the route filter dropped is not brought back.""" + host.answer("openai/gpt-4o", auth_error()) + core: Final = make_core(model_list=[_OPENAI, _ANTHROPIC], num_retries=0) + + with pytest.raises(litellm.AuthenticationError): + _call(core, "adecisions", model="e", **QUESTION) + with pytest.raises(RouterRateLimitError, match=r"(?i)no deployments available"): + _call(core, "adecisions", model="e", **QUESTION) + + assert host.deployments_called("adecisions") == ["p7-openai"] + + +def test_the_retry_count_reads_the_group_before_the_route_filter(make_core: MakeCore, host: EndpointHost) -> None: + """An authentication error is retried only when the group holds another deployment, counted + without the route filter: both deployments count, so the Decisions deployment is retried.""" + host.answer("openai/gpt-4o", auth_error()) + core: Final = make_core(model_list=[_OPENAI, _ANTHROPIC], num_retries=2, retry_after=0, disable_cooldowns=True) + + with pytest.raises(litellm.AuthenticationError): + _call(core, "adecisions", model="e", **QUESTION) + + assert host.deployments_called("adecisions") == ["p7-openai"] * 3 + + +# --- deployments refusing a call's audio or videos ------------------------------------------------ + + +def _media_group(**changes: object) -> list[dict[str, object]]: # mutable-ok: a fresh model_list + """``refuse`` declares no audio or video input and weighs more; ``take`` declares both.""" + flags: Final = ("supports_audio_input", "supports_video_input") + return [ + deployment( + "m", + "refuse", + model="hosted_vllm/refuse-m", + api_base="https://vllm.example", + weight=1, + model_info=dict.fromkeys(flags, False), + **changes, + ), + deployment( + "m", + "take", + model="hosted_vllm/take-m", + api_base="https://vllm.example", + weight=0, + model_info=dict.fromkeys(flags, True), + **changes, + ), + ] + + +@pytest.mark.parametrize("media", ["audio", "videos"]) +@pytest.mark.parametrize("op", ["adecisions", "decisions", "asystemone", "systemone"]) +def test_a_call_carrying_media_skips_the_deployment_refusing_it( + make_core: MakeCore, host: EndpointHost, op: str, media: str +) -> None: + core: Final = make_core(model_list=_media_group(), num_retries=0) + + for _ in range(3): + _call(core, op, model="m", **{media: CLIP}, **QUESTION) + + assert host.deployments_called(op) == ["take"] * 3 + + +@pytest.mark.parametrize("op", ["adecisions", "decisions"]) +def test_a_call_without_media_keeps_the_deployment_refusing_media( + make_core: MakeCore, host: EndpointHost, op: str +) -> None: + core: Final = make_core(model_list=_media_group(), num_retries=0) + + _call(core, op, model="m", **QUESTION) + + assert host.deployments_called(op) == ["refuse"] + + +@pytest.mark.parametrize( + ("router_arguments", "request_extra"), + [({"default_litellm_params": {"drop_params": True}}, {}), ({}, {"drop_params": True})], + ids=["router-default", "request"], +) +def test_drop_params_keeps_the_deployment_refusing_media( + make_core: MakeCore, + host: EndpointHost, + router_arguments: dict[str, object], + request_extra: dict[str, object], +) -> None: + """A deployment that drops the params it cannot send takes the call, media dropped.""" + core: Final = make_core(model_list=_media_group(), num_retries=0, **router_arguments) + + _call(core, "adecisions", model="m", audio=CLIP, **request_extra, **QUESTION) + + assert host.deployments_called("adecisions") == ["refuse"] + + +def test_the_media_filter_runs_before_order(make_core: MakeCore, host: EndpointHost) -> None: + refuse, take = _media_group() + cast("dict[str, object]", refuse["litellm_params"])["order"] = 1 # cast-ok: a model_list entry + cast("dict[str, object]", take["litellm_params"])["order"] = 2 # cast-ok: a model_list entry + core: Final = make_core(model_list=[refuse, take], num_retries=0) + + _call(core, "adecisions", model="m", audio=CLIP, **QUESTION) + + assert host.deployments_called("adecisions") == ["take"] + + +def test_every_deployment_refusing_media_still_picks(make_core: MakeCore, host: EndpointHost) -> None: + refuse, _ = _media_group() + core: Final = make_core(model_list=[refuse, deployment("m", "refuse-2", model="openai/gpt-4o", weight=0)]) + + _call(core, "adecisions", model="m", videos=CLIP, **QUESTION) + + assert host.deployments_called("adecisions") == ["refuse"] diff --git a/tests/unit/test_router/core/test_parity_failure_logging.py b/tests/unit/test_router/core/test_parity_failure_logging.py new file mode 100644 index 00000000000..59173b6ad69 --- /dev/null +++ b/tests/unit/test_router/core/test_parity_failure_logging.py @@ -0,0 +1,647 @@ +""" +Parity of the Rust core with ``router_python_core.py`` on what a failed request leaves behind: the +failure dispatched to the request's logging object (``_async_get_available_deployment``'s ``except``) +and the walk stamped on the caller's own metadata dict, which a proxy's failure SpendLogs row reads +(``async_function_with_fallbacks``, ``async_function_with_retries``). Also the decision headers of a +routed response, typed ``model_info`` read back, ``routing_strategy_args`` as a plain dict and the +``cooldown_list`` of a no-deployments error. +""" + +import asyncio +from collections.abc import Mapping +from datetime import datetime, time, timezone +from typing import Final + +import pytest +from pydantic_core import to_jsonable_python + +import litellm +from litellm.constants import PROXY_REJECTED_BEFORE_ROUTING_KEY +from litellm.integrations.custom_logger import CustomLogger +from litellm.litellm_core_utils.litellm_logging import Logging +from litellm.proxy.spend_tracking.spend_tracking_utils import ( + _get_spend_logs_metadata, # pyright: ignore[reportPrivateUsage] # what a failure row keeps +) +from litellm.router_api import ChangeSet +from litellm.router_core import RouterCore +from litellm.types.router import AllowedFailsPolicy, RouterRateLimitError +from tests.unit.test_router.core.fallbacks_helpers import ( + MESSAGES, + MakeCore, + complete, + complete_sync, + headers, + show_router_debug, +) +from tests.unit.test_router.core.scripted_host import ScriptedHost, deployment, rate_limit, server_error + +GROUP_OF_TWO: Final = [deployment("a", "a1"), deployment("a", "a2")] +UTC_2030: Final = datetime(2030, 1, 1, tzinfo=timezone.utc) + + +class _RecordingLoggingObject: + """Stands in for a request's ``litellm_logging_obj``: records the failure dispatches.""" + + def __init__(self) -> None: + self.dispatched: list[dict[str, object]] = [] # mutable-ok: the dispatches, in order + self.messages: list[object] = [] # mutable-ok: each error's message when the dispatch ran + + async def dispatch_failure_handlers(self, **kwargs: object) -> None: # kwargs-ok: the dispatch arguments + self.dispatched.append(kwargs) + self.messages.append(getattr(kwargs.get("exception"), "message", None)) + + +class _RedisAtARoundTrip: + """The shared tier of the cooldown store, down to the one thing these tests need of it: a read + gives the event loop a turn, as a real Redis round trip does.""" + + def __init__(self) -> None: + self.batches: list[list[str]] = [] # mutable-ok: the batch reads, in order + + async def async_batch_get_cache( + self, key_list: list[str], **_: object + ) -> dict[str, object]: # mutable-ok: as RedisCache + self.batches.append(key_list) + await asyncio.sleep(0) + return dict.fromkeys(key_list) + + +class _RecordingLogging(Logging): + """A real logging object whose failure handlers only record that they ran.""" + + def __init__(self, call_type: str) -> None: + super().__init__( + model="e", + messages=[], + stream=False, + call_type=call_type, + start_time=datetime.now(tz=timezone.utc), + litellm_call_id="c", + function_id="f", + ) + self.update_environment_variables(model="e", user=None, optional_params={}, litellm_params={}) + self.failures: list[tuple[str, object]] = [] # mutable-ok: the handlers that ran, in order + + def failure_handler(self, exception: Exception, traceback_exception: str, *_: object, **__: object) -> None: + self.failures.append(("sync", exception)) + + async def async_failure_handler( + self, exception: Exception, traceback_exception: str, *_: object, **__: object + ) -> None: + self.failures.append(("async", exception)) + + +async def _let_tasks_run() -> None: + for _ in range(10): + await asyncio.sleep(0) + + +async def _failed(core: RouterCore, model: str, **request: object) -> BaseException: # kwargs-ok: request kwargs + with pytest.raises(Exception) as raised: # noqa: PT011 # any failure: each caller asserts on what it gets back + await complete(core, model, **request) + await _let_tasks_run() + return raised.value + + +# --- router-raised errors ------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_an_unknown_model_is_dispatched_once_with_the_error_raised(make_core: MakeCore) -> None: + core: Final = make_core(model_list=list(GROUP_OF_TWO), num_retries=0) + logging_obj: Final = _RecordingLoggingObject() + + error: Final = await _failed(core, "nope", litellm_logging_obj=logging_obj, metadata={}) + + assert isinstance(error, litellm.BadRequestError) + (dispatched,) = logging_obj.dispatched + assert dispatched["exception"] is error + assert dispatched["prefer_async_handlers"] is True + assert isinstance(dispatched["traceback_exception"], str) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("behind_redis", [True, False], ids=["redis", "memory"]) +async def test_a_refused_pick_is_logged_before_the_note_when_the_retry_read_yields( + make_core: MakeCore, host: ScriptedHost, behind_redis: bool +) -> None: + """The retry handler reads the cooldowns even for a name with no group: over Redis that hands + the loop to the failure log, which so reads the error before the no-fallback note lands on it.""" + show_router_debug(host) + core: Final = make_core(model_list=list(GROUP_OF_TWO), num_retries=0, fallbacks=[{"a": ["a"]}]) + redis: Final = _RedisAtARoundTrip() + if behind_redis: + core.health().store.redis_cache = redis # pyright: ignore[reportAttributeAccessIssue] # a double + logging_obj: Final = _RecordingLoggingObject() + + error: Final = await _failed(core, "nope", litellm_logging_obj=logging_obj) + + note: Final = "no fallback model group was found" + assert note in str(error) + (logged,) = logging_obj.messages + assert (note in str(logged)) is not behind_redis + assert redis.batches == ([["deployment:a1:cooldown", "deployment:a2:cooldown"]] if behind_redis else []) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("endpoint", "handlers"), + [("aembedding", ["sync", "async"]), ("embedding", ["sync"])], + ids=["async", "sync"], +) +async def test_an_oversized_embedding_is_stamped_logged_and_counted_as_its_deployments_failure( + make_core: MakeCore, host: ScriptedHost, endpoint: str, handlers: list[str] +) -> None: + core: Final = make_core( + model_list=[deployment("e", "e1", model_info={"max_embedding_batch_size": 1})], num_retries=0 + ) + logging_obj: Final = _RecordingLogging(endpoint) + + with pytest.raises(litellm.BadRequestError) as raised: + await core.call(endpoint, {"model": "e", "input": ["a", "b"], "litellm_logging_obj": logging_obj}) + + error: Final = raised.value + assert getattr(error, "failed_deployment_id", None) == "e1" + assert getattr(error, "retry_skip_deployment_id", None) == "e1" + assert [(kind, exception is error) for kind, exception in logging_obj.failures] == [(k, True) for k in handlers] + assert logging_obj.model_call_details["litellm_params"][PROXY_REJECTED_BEFORE_ROUTING_KEY] is True + assert host.calls == [] + assert dict(core.fail_calls) == {"openai/e1": 1} + assert dict(core.total_calls) == {} + + +@pytest.mark.asyncio +@pytest.mark.parametrize("endpoint", ["aembedding", "embedding"]) +async def test_an_admitted_embedding_drops_the_rejection_an_earlier_attempt_left( + make_core: MakeCore, endpoint: str +) -> None: + core: Final = make_core( + model_list=[deployment("e", "e1", model_info={"max_embedding_batch_size": 2})], num_retries=0 + ) + logging_obj: Final = _RecordingLogging(endpoint) + logging_obj.model_call_details["litellm_params"][PROXY_REJECTED_BEFORE_ROUTING_KEY] = True + + await core.call(endpoint, {"model": "e", "input": ["a", "b"], "litellm_logging_obj": logging_obj}) + + assert PROXY_REJECTED_BEFORE_ROUTING_KEY not in logging_obj.model_call_details["litellm_params"] + + +@pytest.mark.asyncio +async def test_an_unknown_model_stamps_the_walk_on_the_callers_metadata(make_core: MakeCore) -> None: + core: Final = make_core(model_list=list(GROUP_OF_TWO), num_retries=0) + metadata: Final[dict[str, object]] = {"user_api_key_alias": "k"} # mutable-ok: the caller's own dict + + _ = await _failed(core, "nope", metadata=metadata) + + assert metadata["attempted_retries"] == 0 + assert metadata["max_retries"] == 0 + assert metadata["attempted_fallbacks"] == 0 + assert metadata["original_model_group"] == "nope" + assert metadata["model_group"] == "nope" + assert metadata["model_group_size"] == 0 + assert metadata["user_api_key_alias"] == "k" + + +@pytest.mark.asyncio +async def test_a_refused_pick_leaves_every_field_of_the_failure_row(make_core: MakeCore) -> None: + """A real logging object, read the way ``get_logging_payload`` reads it: the walk from the + request's metadata, the rest from the payload the dispatched failure built.""" + core: Final = make_core(model_list=list(GROUP_OF_TWO), num_retries=0) + metadata: Final[dict[str, object]] = {} # mutable-ok: the caller's own dict + logging_obj: Final = Logging( + model="nope", + messages=MESSAGES, + stream=False, + call_type="acompletion", + start_time=datetime.now(), # noqa: DTZ005 # SDK logging uses naive local timestamps + litellm_call_id="call", + function_id="fn", + applied_guardrails=["g1"], + ) + logging_obj.update_environment_variables( + model="nope", user="", optional_params={}, litellm_params={"metadata": metadata} + ) + + _ = await _failed(core, "nope", litellm_logging_obj=logging_obj, metadata=metadata) + + payload: Final = logging_obj.model_call_details["standard_logging_object"] + assert payload["status"] == "failure" + row: Final = _get_spend_logs_metadata( + metadata, + applied_guardrails=payload["metadata"].get("applied_guardrails"), + usage_object=payload["metadata"].get("usage_object"), + model_map_information=payload["model_map_information"], + ) + assert row["attempted_retries"] == 0 + assert row["max_retries"] == 0 + assert row["attempted_fallbacks"] == 0 + assert row["original_model_group"] == "nope" + assert row["applied_guardrails"] == ["g1"] + assert row["usage_object"] == payload["metadata"]["usage_object"] + assert row["model_map_information"] == payload["model_map_information"] + + +@pytest.mark.asyncio +async def test_all_cooled_deployments_are_dispatched_with_their_cooldown_details(make_core: MakeCore) -> None: + core: Final = make_core( + model_list=[ + deployment("a", "a1", api_base="https://a1.example.com"), + deployment("a", "a2", api_base="https://a2.example.com"), + ], + num_retries=0, + cooldown_time=60, + ) + for deployment_id in ("a1", "a2"): + core.health().record_failure(deployment_id, rate_limit(), 429) + logging_obj: Final = _RecordingLoggingObject() + metadata: Final[dict[str, object]] = {} # mutable-ok: the caller's own dict + + error: Final = await _failed(core, "a", litellm_logging_obj=logging_obj, metadata=metadata) + + assert isinstance(error, RouterRateLimitError) + assert error.cooldown_list == ["a1", "a2"] + assert error.cooldown_deployments == ( + {"model_id": "a1", "api_base": "https://a1.example.com"}, + {"model_id": "a2", "api_base": "https://a2.example.com"}, + ) + (dispatched,) = logging_obj.dispatched + assert dispatched["exception"] is error + assert metadata["model_group_size"] == 2 + + +@pytest.mark.asyncio +async def test_a_refused_pick_with_retries_settles_the_retry_budget(make_core: MakeCore) -> None: + core: Final = make_core(model_list=list(GROUP_OF_TWO), num_retries=3) + metadata: Final[dict[str, object]] = {} # mutable-ok: the caller's own dict + + error: Final = await _failed(core, "nope", metadata=metadata) + + assert isinstance(error, litellm.BadRequestError) + # A 400 is not retried: `should_retry_this_error` raises before the budget is written again. + assert metadata["attempted_retries"] == 0 + assert metadata["max_retries"] == 3 + + +@pytest.mark.asyncio +async def test_a_strategy_with_no_async_selector_does_not_dispatch(make_core: MakeCore) -> None: + core: Final = make_core(model_list=list(GROUP_OF_TWO), num_retries=0, routing_strategy="usage-based-routing") + logging_obj: Final = _RecordingLoggingObject() + + _ = await _failed(core, "nope", litellm_logging_obj=logging_obj, metadata={}) + + assert logging_obj.dispatched == [] + + +def test_the_sync_path_stamps_the_walk_and_does_not_dispatch(make_core: MakeCore) -> None: + core: Final = make_core(model_list=list(GROUP_OF_TWO), num_retries=0) + logging_obj: Final = _RecordingLoggingObject() + metadata: Final[dict[str, object]] = {} # mutable-ok: the caller's own dict + + with pytest.raises(litellm.BadRequestError): + complete_sync(core, "nope", litellm_logging_obj=logging_obj, metadata=metadata) + + assert logging_obj.dispatched == [] + assert metadata["attempted_retries"] == 0 + assert metadata["attempted_fallbacks"] == 0 + assert metadata["original_model_group"] == "nope" + + +@pytest.mark.asyncio +async def test_a_refused_image_pick_with_no_logging_object_builds_one_to_dispatch( + make_core: MakeCore, monkeypatch: pytest.MonkeyPatch +) -> None: + dispatched: Final[list[BaseException]] = [] # mutable-ok: the dispatches, in order + + async def record( + self: Logging, exception: Exception, **_arguments: object + ) -> None: # kwargs-ok: dispatch arguments + dispatched.append(exception) + + monkeypatch.setattr(Logging, "dispatch_failure_handlers", record) + core: Final = make_core(model_list=list(GROUP_OF_TWO), num_retries=0) + + with pytest.raises(litellm.BadRequestError) as raised: + await core.aimage_generation(prompt="a cat", model="nope") # pyright: ignore[reportAttributeAccessIssue, reportUnknownMemberType] # an endpoint + await _let_tasks_run() + + assert dispatched == [raised.value] + + +@pytest.mark.asyncio +async def test_a_refused_pre_call_check_is_dispatched_on_chat_only(make_core: MakeCore, host: ScriptedHost) -> None: + class Gate(CustomLogger): + async def async_pre_call_check(self, deployment: dict, parent_otel_span: object) -> dict | None: # pyright: ignore[reportIncompatibleMethodOverride] + raise rate_limit(message="gate says no") + + host.lists["callbacks"].append(Gate()) + core: Final = make_core(model_list=list(GROUP_OF_TWO), num_retries=0) + logging_obj: Final = _RecordingLoggingObject() + + error: Final = await _failed(core, "a", litellm_logging_obj=logging_obj, metadata={}) + + assert isinstance(error, litellm.RateLimitError) + (dispatched,) = logging_obj.dispatched + assert dispatched["exception"] is error + + embedding_logging: Final = _RecordingLoggingObject() + with pytest.raises(litellm.RateLimitError): + await core.aembedding(model="a", input="hi", litellm_logging_obj=embedding_logging) # pyright: ignore[reportAttributeAccessIssue, reportUnknownMemberType] # an endpoint + await _let_tasks_run() + assert embedding_logging.dispatched == [] + + +# --- provider failures ---------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_a_provider_failure_stamps_each_field_and_the_router_dispatches_nothing( + make_core: MakeCore, host: ScriptedHost +) -> None: + host.answer("*", server_error()) + core: Final = make_core(model_list=list(GROUP_OF_TWO), num_retries=1) + logging_obj: Final = _RecordingLoggingObject() + metadata: Final[dict[str, object]] = {} # mutable-ok: the caller's own dict + + error: Final = await _failed(core, "a", litellm_logging_obj=logging_obj, metadata=metadata) + + assert isinstance(error, litellm.InternalServerError) + assert metadata["attempted_retries"] == 1 + assert metadata["max_retries"] == 1 + assert metadata["attempted_fallbacks"] == 0 + assert metadata["original_model_group"] == "a" + assert metadata["model_group_size"] == 2 + assert metadata["request_retry_count"] == 2 + breadcrumbs: Final = metadata["previous_models"] + assert isinstance(breadcrumbs, tuple) + assert len(breadcrumbs) == 2 + # The SDK's own wrapper logs a provider failure, on the same logging object. + assert logging_obj.dispatched == [] + + +@pytest.mark.asyncio +async def test_a_deployments_retry_budget_replaces_the_routers_in_the_stamp( + make_core: MakeCore, host: ScriptedHost +) -> None: + host.answer("*", server_error()) + core: Final = make_core( + model_list=[deployment("a", "a1", num_retries=0), deployment("a", "a2", num_retries=0)], num_retries=2 + ) + metadata: Final[dict[str, object]] = {} # mutable-ok: the caller's own dict + + _ = await _failed(core, "a", metadata=metadata) + + assert metadata["max_retries"] == 0 + assert metadata["attempted_retries"] == 0 + + +@pytest.mark.asyncio +async def test_a_fallback_hop_works_on_its_own_copy_of_the_metadata(make_core: MakeCore, host: ScriptedHost) -> None: + host.answer("*", server_error()) + core: Final = make_core( + model_list=[deployment("a", "a1"), deployment("b", "b1")], num_retries=0, fallbacks=[{"a": ["b"]}] + ) + metadata: Final[dict[str, object]] = {} # mutable-ok: the caller's own dict + + _ = await _failed(core, "a", metadata=metadata) + + assert metadata["attempted_fallbacks"] == 0 + assert metadata["attempted_retries"] == 0 + assert metadata["max_retries"] == 0 + assert metadata["original_model_group"] == "a" + assert metadata["model_group"] == "a" + hop: Final = host.calls_to("acompletion")[1].kwargs["metadata"] + assert isinstance(hop, Mapping) + assert hop["attempted_fallbacks"] == 1 + assert hop["original_model_group"] == "a" + assert hop["model_group"] == "b" + + +@pytest.mark.asyncio +async def test_the_fallbacks_mock_raises_before_the_retry_loop_stamps_anything(make_core: MakeCore) -> None: + core: Final = make_core(model_list=list(GROUP_OF_TWO), num_retries=2) + logging_obj: Final = _RecordingLoggingObject() + metadata: Final[dict[str, object]] = {} # mutable-ok: the caller's own dict + + _ = await _failed(core, "a", mock_testing_fallbacks=True, litellm_logging_obj=logging_obj, metadata=metadata) + + assert metadata["attempted_fallbacks"] == 0 + assert metadata["original_model_group"] == "a" + assert "attempted_retries" not in metadata + assert "max_retries" not in metadata + assert logging_obj.dispatched == [] + + +@pytest.mark.asyncio +async def test_the_rate_limit_mock_raises_inside_the_retry_loop(make_core: MakeCore) -> None: + core: Final = make_core(model_list=list(GROUP_OF_TWO), num_retries=0) + logging_obj: Final = _RecordingLoggingObject() + metadata: Final[dict[str, object]] = {} # mutable-ok: the caller's own dict + + _ = await _failed(core, "a", mock_testing_rate_limit_error=True, litellm_logging_obj=logging_obj, metadata=metadata) + + assert metadata["attempted_retries"] == 0 + assert metadata["max_retries"] == 0 + assert metadata["model_group_size"] == 2 + assert logging_obj.dispatched == [] + + +@pytest.mark.asyncio +async def test_a_generic_endpoint_stamps_its_bucket_and_clears_the_siblings_walk( + make_core: MakeCore, host: ScriptedHost +) -> None: + host.answer("*", server_error()) + core: Final = make_core(model_list=list(GROUP_OF_TWO), num_retries=0) + sibling: Final[dict[str, object]] = { # mutable-ok: the caller's own dict + "attempted_fallbacks": 7, + "original_model_group": "stale", + "kept": True, + } + bucket: Final[dict[str, object]] = {} # mutable-ok: the caller's own dict + + with pytest.raises(litellm.InternalServerError): + await core.aresponses(model="a", input="hi", metadata=sibling, litellm_metadata=bucket) # pyright: ignore[reportAttributeAccessIssue, reportUnknownMemberType] # an endpoint + + assert "attempted_fallbacks" not in sibling + assert "original_model_group" not in sibling + assert sibling["kept"] is True + assert bucket["attempted_fallbacks"] == 0 + assert bucket["original_model_group"] == "a" + assert bucket["attempted_retries"] == 0 + assert bucket["max_retries"] == 0 + + +# --- R7: the decision headers --------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_a_stamped_quality_decision_reaches_the_response_headers(make_core: MakeCore) -> None: + core: Final = make_core(model_list=list(GROUP_OF_TWO), num_retries=0) + decision: Final = {"routed_model": "t1", "quality_tier": 2, "routed_via": "default", "matched_keyword": None} + + response: Final = await complete(core, "a", metadata={"quality_router_decision": decision}) + + found: Final = headers(response) + assert found["x-litellm-quality-router-model"] == "t1" + assert found["x-litellm-quality-router-tier"] == "2" + assert found["x-litellm-quality-router-via"] == "default" + assert "x-litellm-quality-router-keyword" not in found + + +@pytest.mark.asyncio +async def test_a_quality_router_stamps_its_decision_headers(make_core: MakeCore) -> None: + core: Final = make_core( + model_list=[ + deployment( + "t1", + "t1-21", + model="anthropic/claude-3-5-sonnet-20240620", + model_info={"litellm_routing_preferences": {"quality_tier": 1}}, + ), + deployment("qr", "qr0", model="auto_router/quality_router", quality_router_default_model="t1"), + ], + num_retries=0, + ) + + response: Final = await complete(core, "qr") + + assert headers(response)["x-litellm-quality-router-model"] == "t1" + + +# --- P9, N9: typed model_info --------------------------------------------------------------------- + + +def _typed_model_list() -> list[dict[str, object]]: # mutable-ok: a model_list + return [ + deployment( + "p", + "p1", + model="azure/gpt-4o", + api_base="https://example.openai.azure.com", + model_info={ + "ptu_shares": {"a": 2}, + "ptu_count": 2, + "cost_per_ptu_per_hour": 2, + "ptu_effective_from": "2026-01-01T00:00:00Z", + }, + ), + deployment( + "h", + "w3", + model_info={ + "access_windows": [ + {"start": "03:00:00", "end": "03:00:01", "timezone": "Asia/Tokyo", "team_ids": ["t1"]} + ] + }, + ), + deployment("d", "db1", model_info={"created_at": UTC_2030, "updated_at": UTC_2030, "db_model": True}), + ] + + +def _model_info(core: RouterCore, deployment_id: str) -> Mapping[str, object]: + for entry in core.registry().all(): + info = entry["model_info"] + assert isinstance(info, Mapping) + if info["id"] == deployment_id: + return info + raise AssertionError(f"no deployment {deployment_id}") + + +def test_model_info_reads_back_a_datetime(make_core: MakeCore) -> None: + core: Final = make_core(model_list=_typed_model_list()) + effective_from: Final = _model_info(core, "p1")["ptu_effective_from"] + assert isinstance(effective_from, datetime) + assert effective_from == datetime(2026, 1, 1, tzinfo=timezone.utc) + + +def test_model_info_reads_back_access_windows_typed(make_core: MakeCore) -> None: + core: Final = make_core(model_list=_typed_model_list()) + windows: Final = _model_info(core, "w3")["access_windows"] + assert isinstance(windows, tuple) + (window,) = windows + assert window["start"] == time(3) + assert window["end"] == time(3, 0, 1) + assert window["team_ids"] == ("t1",) + + +def test_db_timestamps_read_back_as_the_datetimes_stamped(make_core: MakeCore) -> None: + core: Final = make_core(model_list=_typed_model_list()) + info: Final = _model_info(core, "db1") + assert info["created_at"] == UTC_2030 + assert info["updated_at"] == UTC_2030 + created_at: Final = info["created_at"] + assert isinstance(created_at, datetime) + assert created_at.isoformat() == "2030-01-01T00:00:00+00:00" + + +@pytest.mark.asyncio +async def test_a_picked_deployment_carries_its_typed_model_info(make_core: MakeCore) -> None: + core: Final = make_core(model_list=_typed_model_list()) + picked: Final = await core.select("d") + assert isinstance(picked, Mapping) + model_info: Final = picked["model_info"] + assert model_info["created_at"] == UTC_2030 + + +# --- PV1: routing_strategy_args ------------------------------------------------------------------- + + +def test_routing_strategy_args_serialize_as_a_plain_dict(make_core: MakeCore) -> None: + core: Final = make_core(model_list=list(GROUP_OF_TWO), routing_strategy_args={"ttl": 10}) + args: Final = core.routing_strategy_args # pyright: ignore[reportAttributeAccessIssue, reportUnknownMemberType, reportUnknownVariableType] # the deprecated attribute + assert isinstance(args, dict) + assert to_jsonable_python(args) == {"ttl": 10} + + +def test_unset_routing_strategy_args_serialize_too(make_core: MakeCore) -> None: + core: Final = make_core(model_list=list(GROUP_OF_TWO)) + assert to_jsonable_python(core.routing_strategy_args) == {} # pyright: ignore[reportAttributeAccessIssue, reportUnknownMemberType] # the deprecated attribute + + +# --- U13: the cooldown list ----------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_a_routing_group_with_nothing_left_lists_the_cooling_deployment( + make_core: MakeCore, host: ScriptedHost +) -> None: + host.answer("openai/rc1", rate_limit()) + core: Final = make_core( + model_list=[deployment("g1", "rc1"), deployment("g2", "rc2")], + num_retries=0, + routing_groups=[ + { + "group_name": "pg", + "models": ["g1", "g2"], + "routing_strategy": "priority", + "model_priorities": {"g1": 1, "g2": 2}, + } + ], + ) + _ = await complete(core, "pg") + core.apply(ChangeSet(delete=("rc2",))) + + error: Final = await _failed(core, "pg") + + assert isinstance(error, RouterRateLimitError) + assert error.cooldown_list == ["rc1"] + + +@pytest.mark.asyncio +async def test_the_cooldown_list_names_every_cooling_deployment_of_the_router(make_core: MakeCore) -> None: + core: Final = make_core( + model_list=[deployment("a", "a1"), deployment("b", "b1")], + num_retries=0, + cooldown_time=60, + allowed_fails_policy=AllowedFailsPolicy( + RateLimitErrorAllowedFails=0 + ), # a lone deployment stays warm without it + ) + for deployment_id in ("a1", "b1"): + core.health().record_failure(deployment_id, rate_limit(), 429) + + error: Final = await _failed(core, "a") + + assert isinstance(error, RouterRateLimitError) + assert error.cooldown_list == ["a1", "b1"] + assert error.all_deployments_in_cooldown is True + assert error.cooldown_deployments == ({"model_id": "a1", "api_base": None},) diff --git a/tests/unit/test_router/core/test_parity_malformed_after_rate_limit.py b/tests/unit/test_router/core/test_parity_malformed_after_rate_limit.py new file mode 100644 index 00000000000..7181beaa39a --- /dev/null +++ b/tests/unit/test_router/core/test_parity_malformed_after_rate_limit.py @@ -0,0 +1,86 @@ +""" +Parity of the Rust core with ``router_python_core.py`` on a retry whose answer fails conversion +after a rate limit (fuzz case ff68399265ba9aee, ruled 2026-11-05: a 500 on both cores). + +The Python core runs a retry inside the ``except`` of the attempt before it, so the retry's +conversion error is raised while the first attempt's RateLimitError is being handled. The +"Invalid response object" message once printed that whole chain, and exception mapping read the +rate-limit phrase in it and answered 429. The retry here runs the real conversion and the real +mapping, as ``OpenAIChatCompletion.acompletion`` and ``atext_completion`` do, at the moment the core +calls the deployment, so the error being handled then is whatever the core leaves open. +""" + +from collections.abc import Awaitable, Callable +from typing import Final, cast + +import pytest + +import litellm +from litellm.litellm_core_utils.exception_mapping_utils import exception_type # pyright: ignore[reportUnknownVariableType] # its untyped parameters +from litellm.litellm_core_utils.llm_response_utils.convert_dict_to_response import convert_to_model_response_object # pyright: ignore[reportUnknownVariableType] # its untyped start_time, end_time and dict parameters +from litellm.llms.openai.common_utils import OpenAIError +from litellm.router_core import RouterCore +from litellm.types.utils import ModelResponse +from tests.unit.test_router.core.scripted_host import ScriptedHost, deployment, rate_limit + +MakeCore = Callable[..., RouterCore] + +# A text_completion body handed to the chat conversion (gpt-4o is a chat model): no ``message``. +_TEXT_COMPLETION_BODY: Final = { + "id": "cmpl-b6", + "object": "text_completion", + "created": 1, + "model": "gpt-4o", + "choices": [{"text": "ok", "index": 0, "finish_reason": "stop"}], + "usage": {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2}, +} + + +def _malformed_answer(**_request: object) -> BaseException: # kwargs-ok: the scripted call's request + """The error the OpenAI handler and ``atext_completion`` raise for ``_TEXT_COMPLETION_BODY``: + the conversion fails, the handler wraps it in a status-500 ``OpenAIError``, and ``exception_type`` + maps that. Returned, so the host raises and logs it as any scripted error.""" + try: + _ = convert_to_model_response_object( + response_object=dict(_TEXT_COMPLETION_BODY), model_response_object=ModelResponse() + ) + except Exception as conversion_error: # noqa: BLE001 # the handler wraps every conversion failure + try: + raise exception_type( # exception_type raises the mapped error itself + model="gpt-4o", + original_exception=OpenAIError(status_code=500, message=str(conversion_error)), + custom_llm_provider="openai", + ) + except Exception as mapped: # noqa: BLE001 # the mapped error is the scripted answer + return mapped + raise AssertionError("the malformed body converted") + + +def _atext_completion(core: RouterCore) -> Callable[..., Awaitable[object]]: + """``atext_completion``, which both cores serve outside the ``RouterCore`` protocol.""" + return cast( + "Callable[..., Awaitable[object]]", getattr(core, "atext_completion") + ) # cast-ok: the member is a coroutine function + + +@pytest.mark.asyncio +async def test_a_retry_whose_answer_fails_conversion_after_a_rate_limit_is_a_500( + make_core: MakeCore, host: ScriptedHost +) -> None: + core: Final = make_core( + model_list=[deployment("g", "b6-d1", weight=1), deployment("g", "b6-d2", weight=1)], + num_retries=1, + retry_after=0, + ) + host.answer("*", rate_limit(message="rate limit exceeded"), _malformed_answer) + + with pytest.raises(litellm.InternalServerError) as raised: + _ = await _atext_completion(core)(model="g", prompt="hi") + + error: Final = raised.value + assert error.status_code == 500 + assert "Invalid response object" in str(error) + assert "KeyError: 'message'" in str(error) + assert "rate limit exceeded" not in str(error) + assert len(host.calls_to("atext_completion")) == 2 + assert type(host.raised[0]) is litellm.RateLimitError diff --git a/tests/unit/test_router/core/test_parity_pick_count_and_fallback_log.py b/tests/unit/test_router/core/test_parity_pick_count_and_fallback_log.py new file mode 100644 index 00000000000..01b1c1fe0e1 --- /dev/null +++ b/tests/unit/test_router/core/test_parity_pick_count_and_fallback_log.py @@ -0,0 +1,316 @@ +""" +Parity of the Rust core with ``router_python_core.py`` on two traces a call leaves on the proxy: +the per-deployment pick count in the router cache (``_track_deployment_metrics``, ``_update_usage``) +and the error log of a level whose fallbacks served nothing (the ``except`` around the regular +fallbacks of ``async_function_with_fallbacks_common_utils``). +""" + +import logging +from collections.abc import Generator +from typing import Final + +import openai +import pytest + +import litellm +from litellm._logging import verbose_router_logger +from litellm.router_api import ChangeSet +from litellm.router_core import RouterCore +from tests.unit.test_router.core.fallbacks_helpers import ( + TIERED_A, + MakeCore, + build, + complete, + complete_sync, + show_router_debug, +) +from tests.unit.test_router.core.scripted_host import ( + ScriptedHost, + auth_error, + content_policy, + deployment, + embedding_response, + rate_limit, + server_error, +) + +LOG_HEAD: Final = "litellm.router.py::async_function_with_fallbacks() - Error occurred while trying to do fallbacks - " +NO_COOLDOWNS: Final = "\nDebug Information:\nCooldown Deployments=[]" + + +class _Records(logging.Handler): + def __init__(self) -> None: + super().__init__(logging.ERROR) + self.messages: Final[list[str]] = [] + + def emit(self, record: logging.LogRecord) -> None: + self.messages.append(record.getMessage()) + + +@pytest.fixture +def router_errors() -> Generator[list[str], None, None]: + """The router logger's error messages, in order.""" + records: Final = _Records() + verbose_router_logger.addHandler(records) + yield records.messages + verbose_router_logger.removeHandler(records) + + +def _picks(core: RouterCore, deployment_id: object) -> object: + return core.health().store.in_memory_cache.get_cache(deployment_id) + + +def _fallback_logs(messages: list[str]) -> list[str]: # mutable-ok: a fresh list + return [message for message in messages if message.startswith(LOG_HEAD)] + + +async def _fails(core: RouterCore, model: str = "a", **request: object) -> BaseException: # kwargs-ok: request kwargs + with pytest.raises(openai.APIStatusError) as raised: + await complete(core, model, **request) + return raised.value + + +# --- the pick count ------------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_each_async_chat_pick_counts_in_the_router_cache(make_core: MakeCore) -> None: + core: Final = build(make_core) + + _ = await complete(core) + _ = await complete(core) + + assert _picks(core, "a") == 2 + + +@pytest.mark.asyncio +async def test_a_pick_counts_before_its_attempt_fails_and_a_retry_picks_again( + make_core: MakeCore, host: ScriptedHost +) -> None: + host.answer("openai/a", server_error()) + core: Final = build(make_core, num_retries=2, retry_after=0) + + _ = await _fails(core) + + assert _picks(core, "a") == 3 + + +@pytest.mark.asyncio +async def test_every_hop_of_a_fallback_chain_counts_its_pick(make_core: MakeCore, host: ScriptedHost) -> None: + host.answer("openai/a", server_error()) + core: Final = build(make_core, fallbacks=[{"a": ["b"]}]) + + _ = await complete(core) + + assert _picks(core, "a") == 1 + assert _picks(core, "b") == 1 + assert _picks(core, "c") is None + + +def test_a_sync_chat_pick_does_not_count(make_core: MakeCore) -> None: + core: Final = build(make_core) + + _ = complete_sync(core) + + assert _picks(core, "a") is None + + +@pytest.mark.asyncio +async def test_an_embedding_pick_does_not_count(make_core: MakeCore, host: ScriptedHost) -> None: + host.answer("openai/a", embedding_response()) + core: Final = build(make_core) + + _ = await core.aembedding(model="a", input="hi") # pyright: ignore[reportAttributeAccessIssue] # endpoint method + + assert _picks(core, "a") is None + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("assigned", "counted"), + [("fallback-id", "fallback-id"), (7, 7), (None, None)], + ids=["string-id", "other-id", "no-id"], +) +async def test_a_default_deployment_pick_counts_under_the_id_it_holds( + make_core: MakeCore, assigned: object, counted: object +) -> None: + core: Final = build(make_core) + served: Final = deployment("*", "x", model="openai/x") + _ = core.apply(ChangeSet(assign={"default_deployment": {**served, "model_info": {"id": assigned}}})) + + _ = await complete(core, "not-a-group") + + if counted is not None: + assert _picks(core, counted) == 1 + assert _picks(core, "x") is None + + +# --- the fallback failure log --------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_a_chain_that_fails_logs_every_level_innermost_first( + make_core: MakeCore, host: ScriptedHost, router_errors: list[str] +) -> None: + for name in ("a", "b", "c"): + host.answer(f"openai/{name}", rate_limit(message=f"{name} down")) + core: Final = build(make_core, fallbacks=[{"a": ["b"]}, {"b": ["c"]}], disable_cooldowns=True) + + _ = await _fails(core) + + # Each level logs what its `try` raised: `c` has no chain, so its own error, and `b`'s hop to it + # raised that; `a`'s hop to `b` raised `b`'s own error. + logged: Final = _fallback_logs(router_errors) + assert ["c down" in logged[0], "c down" in logged[1], "b down" in logged[2], len(logged)] == [True, True, True, 3] + assert all(entry.endswith(NO_COOLDOWNS) for entry in logged) + + +@pytest.mark.asyncio +async def test_a_failed_hop_is_logged_with_the_hops_error_not_the_levels( + make_core: MakeCore, host: ScriptedHost, router_errors: list[str] +) -> None: + host.answer("openai/a", server_error()) + host.answer("openai/b", auth_error()) + core: Final = build(make_core, fallbacks=[{"a": ["b"]}], disable_cooldowns=True) + + error: Final = await _fails(core) + + assert type(error).__name__ == "InternalServerError" + logged: Final = _fallback_logs(router_errors) + assert len(logged) == 2 + assert all("AuthenticationError" in entry for entry in logged) + + +@pytest.mark.asyncio +async def test_no_fallback_list_logs_nothing(make_core: MakeCore, host: ScriptedHost, router_errors: list[str]) -> None: + host.answer("openai/a", server_error()) + core: Final = build(make_core) + + _ = await _fails(core) + + assert _fallback_logs(router_errors) == [] + + +@pytest.mark.asyncio +async def test_a_fallback_list_that_does_not_name_the_group_logs_the_levels_error( + make_core: MakeCore, host: ScriptedHost, router_errors: list[str] +) -> None: + host.answer("openai/a", server_error()) + core: Final = build(make_core, fallbacks=[{"b": ["c"]}], disable_cooldowns=True) + + error: Final = await _fails(core) + + assert _fallback_logs(router_errors) == [f"{LOG_HEAD}{error}{NO_COOLDOWNS}"] + + +@pytest.mark.asyncio +async def test_the_missing_group_explanation_is_in_the_log_and_on_the_error_once( + make_core: MakeCore, host: ScriptedHost, router_errors: list[str] +) -> None: + show_router_debug(host) + host.answer("openai/a", server_error()) + core: Final = build(make_core, fallbacks=[{"b": ["c"]}], disable_cooldowns=True) + + error: Final = await _fails(core) + + assert str(error).count("no fallback model group was found") == 1 + assert _fallback_logs(router_errors) == [f"{LOG_HEAD}{error}{NO_COOLDOWNS}"] + + +@pytest.mark.asyncio +async def test_the_outcome_message_of_a_failed_chain_stays_out_of_the_log( + make_core: MakeCore, host: ScriptedHost, router_errors: list[str] +) -> None: + show_router_debug(host) + host.answer("openai/a", server_error()) + host.answer("openai/b", server_error()) + core: Final = build(make_core, fallbacks=[{"a": ["b"]}], disable_cooldowns=True) + + error: Final = await _fails(core) + + logged: Final = _fallback_logs(router_errors) + assert len(logged) == 2 + assert "Fallback to b also failed" in str(error) + assert not any("also failed" in entry for entry in logged) + + +@pytest.mark.asyncio +async def test_a_content_policy_list_without_the_group_logs_before_its_outcome_message( + make_core: MakeCore, host: ScriptedHost, router_errors: list[str] +) -> None: + show_router_debug(host) + host.answer("openai/a", content_policy()) + core: Final = build(make_core, content_policy_fallbacks=[{"b": ["c"]}], disable_cooldowns=True) + + error: Final = await _fails(core) + + (logged,) = _fallback_logs(router_errors) + assert logged.endswith(NO_COOLDOWNS) + assert len(logged) < len(f"{LOG_HEAD}{error}{NO_COOLDOWNS}") + + +@pytest.mark.asyncio +async def test_max_fallbacks_of_zero_logs_the_levels_error( + make_core: MakeCore, host: ScriptedHost, router_errors: list[str] +) -> None: + host.answer("openai/a", server_error()) + core: Final = build(make_core, fallbacks=[{"a": ["b"]}], max_fallbacks=0, disable_cooldowns=True) + + error: Final = await _fails(core) + + assert _fallback_logs(router_errors) == [f"{LOG_HEAD}{error}{NO_COOLDOWNS}"] + + +@pytest.mark.asyncio +async def test_the_order_path_logs_only_where_its_last_tier_reaches_the_regular_fallbacks( + make_core: MakeCore, host: ScriptedHost, router_errors: list[str] +) -> None: + host.answer("openai/a1", server_error()) + host.answer("openai/a2", server_error()) + core: Final = make_core(model_list=list(TIERED_A), num_retries=0, fallbacks=[{"b": ["c"]}], disable_cooldowns=True) + + _ = await _fails(core) + + assert len(_fallback_logs(router_errors)) == 1 + + +@pytest.mark.asyncio +async def test_disabled_fallbacks_log_nothing( + make_core: MakeCore, host: ScriptedHost, router_errors: list[str] +) -> None: + host.answer("openai/a", server_error()) + core: Final = build(make_core, fallbacks=[{"a": ["b"]}]) + + _ = await _fails(core, disable_fallbacks=True) + + assert _fallback_logs(router_errors) == [] + + +def test_a_sync_call_logs_through_the_same_handling( + make_core: MakeCore, host: ScriptedHost, router_errors: list[str] +) -> None: + host.answer("openai/a", server_error()) + core: Final = build(make_core, fallbacks=[{"b": ["c"]}], disable_cooldowns=True) + + with pytest.raises(litellm.InternalServerError) as raised: + _ = complete_sync(core) + + assert _fallback_logs(router_errors) == [f"{LOG_HEAD}{raised.value}{NO_COOLDOWNS}"] + + +@pytest.mark.asyncio +async def test_the_log_lists_the_deployments_cooling_down( + make_core: MakeCore, host: ScriptedHost, router_errors: list[str] +) -> None: + host.answer("openai/a1", auth_error()) + host.answer("openai/a2", auth_error()) + core: Final = make_core( + model_list=[deployment("a", "a1"), deployment("a", "a2"), deployment("b", "b")], + num_retries=0, + fallbacks=[{"b": ["c"]}], + ) + + _ = await _fails(core) + + (logged,) = _fallback_logs(router_errors) + assert "Cooldown Deployments=[('a" in logged diff --git a/tests/unit/test_router/core/test_parity_provider_scoped_fallbacks.py b/tests/unit/test_router/core/test_parity_provider_scoped_fallbacks.py new file mode 100644 index 00000000000..5bd48bfa4ea --- /dev/null +++ b/tests/unit/test_router/core/test_parity_provider_scoped_fallbacks.py @@ -0,0 +1,182 @@ +""" +Parity of the Rust core with ``router_python_core.py`` on fallbacks for a request that creates or +names a resource only one provider holds (``run_async_fallback``'s ``same_model_group_only``): a file +upload, a batch on an uploaded file, a batch cancel and the generic file and fine-tuning ops fall +back within their own model group only, while the same op naming no such resource still falls back +to another group. +""" + +import io +from collections.abc import Awaitable, Callable +from typing import Final, cast + +import pytest + +from litellm.router_core import RouterCore +from litellm.router_python_host import Endpoint +from tests.unit.test_router.core.endpoints_helpers import EndpointHost +from tests.unit.test_router.core.scripted_host import deployment, server_error + +MakeCore = Callable[..., RouterCore] + +_BATCH: Final = {"endpoint": "/v1/chat/completions", "completion_window": "24h"} +_TWO_GROUPS: Final = [deployment("a", "a1"), deployment("b", "b1")] +_TO_B: Final = [{"a": ["b"]}] + + +class _NamedHost(EndpointHost): + """``EndpointHost`` whose endpoints carry the op's name, as litellm's own functions do: the Python + core reads a generic op's handler name to decide it names a provider-scoped resource.""" + + __slots__ = () + + def invoke(self, op: str) -> Endpoint: + endpoint: Final = super().invoke(op) + endpoint.__name__ = op.rsplit(".", 1)[-1] + return endpoint + + +@pytest.fixture +def host() -> _NamedHost: + return _NamedHost() + + +def _async(core: object, name: str) -> Callable[..., Awaitable[object]]: + """An endpoint both cores serve outside the ``RouterCore`` protocol.""" + return cast("Callable[..., Awaitable[object]]", getattr(core, name)) # cast-ok: the member is a coroutine function + + +def _upload() -> dict[str, object]: # mutable-ok: a fresh request + return {"file": io.BytesIO(b'{"custom_id": "1", "body": {"model": "a", "messages": []}}\n'), "purpose": "batch"} + + +def _failing_a( + make_core: MakeCore, host: EndpointHost, model_list: list[dict[str, object]] | None = None +) -> RouterCore: + host.answer("openai/a1", server_error()) + return make_core(model_list=list(model_list or _TWO_GROUPS), num_retries=0) + + +# --- no fallback to another group ------------------------------------------------------------------ + + +@pytest.mark.asyncio +async def test_a_file_upload_does_not_fall_back_to_another_group(make_core: MakeCore, host: _NamedHost) -> None: + core: Final = _failing_a(make_core, host) + + with pytest.raises(type(server_error())): + await _async(core, "acreate_file")(model="a", fallbacks=_TO_B, **_upload()) + + assert host.deployments_called("acreate_file") == ["a1"] + + +@pytest.mark.asyncio +async def test_a_batch_on_an_uploaded_file_does_not_fall_back_to_another_group( + make_core: MakeCore, host: _NamedHost +) -> None: + core: Final = _failing_a(make_core, host) + + with pytest.raises(type(server_error())): + await _async(core, "acreate_batch")(model="a", input_file_id="file-1", fallbacks=_TO_B, **_BATCH) + + assert host.deployments_called("acreate_batch") == ["a1"] + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("op", "request_fields"), + [ + ("acancel_batch", {"batch_id": "batch-1"}), + ("afile_content", {"file_id": "file-1"}), + ("afile_delete", {"file_id": "file-1"}), + ("acreate_fine_tuning_job", {"training_file": "file-1"}), + ("acancel_fine_tuning_job", {"fine_tuning_job_id": "ftjob-1"}), + ], + ids=["cancel-batch", "file-content", "file-delete", "create-fine-tuning-job", "cancel-fine-tuning-job"], +) +async def test_an_op_naming_a_provider_resource_does_not_fall_back_to_another_group( + make_core: MakeCore, host: _NamedHost, op: str, request_fields: dict[str, str] +) -> None: + core: Final = _failing_a(make_core, host) + + with pytest.raises(type(server_error())): + await _async(core, op)(model="a", fallbacks=_TO_B, **request_fields) + + assert host.deployments_called(op) == ["a1"] + + +@pytest.mark.asyncio +async def test_a_dict_fallback_without_a_model_is_skipped_for_a_scoped_batch( + make_core: MakeCore, host: _NamedHost +) -> None: + core: Final = _failing_a(make_core, host) + + with pytest.raises(type(server_error())): + await _async(core, "acreate_batch")( + model="a", input_file_id="file-1", fallbacks=[{"a": [{"user": "u"}]}], **_BATCH + ) + + assert host.deployments_called("acreate_batch") == ["a1"] + + +# --- what still falls back ------------------------------------------------------------------------ + + +@pytest.mark.asyncio +async def test_a_batch_without_an_input_file_still_falls_back_to_another_group( + make_core: MakeCore, host: _NamedHost +) -> None: + core: Final = _failing_a(make_core, host) + + await _async(core, "acreate_batch")(model="a", fallbacks=_TO_B, **_BATCH) + + assert host.deployments_called("acreate_batch") == ["a1", "b1"] + + +@pytest.mark.asyncio +async def test_a_file_op_naming_no_file_still_falls_back_to_another_group( + make_core: MakeCore, host: _NamedHost +) -> None: + core: Final = _failing_a(make_core, host) + + await _async(core, "afile_content")(model="a", fallbacks=_TO_B) + + assert host.deployments_called("afile_content") == ["a1", "b1"] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("op", ["acreate_file", "acreate_batch"], ids=["file-upload", "batch-on-file"]) +async def test_a_scoped_request_still_falls_back_within_its_own_group( + make_core: MakeCore, host: _NamedHost, op: str +) -> None: + core: Final = _failing_a(make_core, host) + request: Final = _upload() if op == "acreate_file" else {"input_file_id": "file-1", **_BATCH} + # A dict target naming the failed group with a param of its own is a distinct target, not the failed one. + same_group: Final = [{"a": [{"model": "a", "user": "u"}, "b"]}] + + with pytest.raises(type(server_error())): + await _async(core, op)(model="a", fallbacks=same_group, **request) + + assert host.deployments_called(op) == ["a1", "a1"] + + +@pytest.mark.asyncio +async def test_a_batch_on_an_uploaded_file_still_falls_back_to_the_next_order_level( + make_core: MakeCore, host: _NamedHost +) -> None: + tiered: Final = [deployment("a", "a1", order=1), deployment("a", "a2", order=2), deployment("b", "b1")] + core: Final = _failing_a(make_core, host, tiered) + + await _async(core, "acreate_batch")(model="a", input_file_id="file-1", fallbacks=_TO_B, **_BATCH) + + assert host.deployments_called("acreate_batch") == ["a1", "a2"] + + +@pytest.mark.asyncio +async def test_a_file_upload_moves_to_the_next_order_level(make_core: MakeCore, host: _NamedHost) -> None: + tiered: Final = [deployment("a", "a1", order=1), deployment("a", "a2", order=2)] + core: Final = _failing_a(make_core, host, tiered) + + await _async(core, "acreate_file")(model="a", **_upload()) + + assert host.deployments_called("acreate_file") == ["a1", "a2"] diff --git a/tests/unit/test_router/core/test_parity_regress5.py b/tests/unit/test_router/core/test_parity_regress5.py new file mode 100644 index 00000000000..3337b7abdf6 --- /dev/null +++ b/tests/unit/test_router/core/test_parity_regress5.py @@ -0,0 +1,482 @@ +""" +Parity of the Rust core with ``router_python_core.py`` on the selection causes of the regress5 pack: +LAR-1 set after build, endpoint eligibility, web search, ``order`` in ``model_info``, ``priority`` +through an alias, old-style custom strategies, wildcard filling and cache flags, ``specific_deployment`` +read as a strict True, pass-through picks, affinity narrowing and a health scope given as a dict. +""" + +import time +from collections.abc import Mapping, Sequence +from typing import Final + +import pytest + +import litellm +from litellm.exceptions import EndpointNotSupportedByModelError +from litellm.router_api import ChangeSet +from litellm.router_core import RouterCore +from litellm.types.router import CustomRoutingStrategyBase, RouterRateLimitError +from tests.unit.test_router.core.scripted_host import ScriptedHost, deployment +from tests.unit.test_router.core.selection_helpers import MESSAGES, MakeCore, ids_of, picked_ids + +SPREAD_PICKS: Final = 40 +WEB_SEARCH_TOOLS: Final = [{"type": "web_search"}] +EMBEDDINGS_ONLY: Final = {"enforce_supported_endpoints": True, "supported_endpoints": ["/v1/embeddings"]} + + +def _embedding_group() -> Sequence[Mapping[str, object]]: + return ( + deployment("g", "embonly", weight=1, model_info=dict(EMBEDDINGS_ONLY)), + deployment("g", "chat", weight=0), + ) + + +def _web_group() -> Sequence[Mapping[str, object]]: + return ( + deployment("g", "noweb", weight=1, model_info={"supports_web_search": False}), + deployment("g", "web", weight=0), + ) + + +# --- V1: lar1 set after build --------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_lar1_set_after_build_serves_the_next_async_call(make_core: MakeCore, host: ScriptedHost) -> None: + core: Final = make_core(model_list=[deployment("a", "a1")], num_retries=0) + core.apply(ChangeSet(settings={"routing_strategy": "lar1"})) + await core.acompletion(model="a", messages=MESSAGES) # pyright: ignore[reportAttributeAccessIssue, reportUnknownMemberType] # an endpoint + assert host.deployments_called("acompletion") == ["a1"] + assert core.settings().routing_strategy == "lar1" + + +@pytest.mark.asyncio +async def test_lar1_set_after_build_picks_without_calling(make_core: MakeCore) -> None: + core: Final = make_core(model_list=[deployment("a", "a1"), deployment("a", "a2")], num_retries=0) + core.apply(ChangeSet(settings={"routing_strategy": "lar1"})) + assert await picked_ids(core, "a", times=SPREAD_PICKS, request_kwargs={}) <= {"a1", "a2"} + + +def test_lar1_set_after_build_refuses_the_sync_path(make_core: MakeCore) -> None: + core: Final = make_core(model_list=[deployment("a", "a1")], num_retries=0) + core.apply(ChangeSet(settings={"routing_strategy": "lar1"})) + with pytest.raises(NotImplementedError, match="LAR-1 routing only supports async routing"): + core.completion(model="a", messages=MESSAGES) # pyright: ignore[reportAttributeAccessIssue, reportUnknownMemberType] # an endpoint + + +@pytest.mark.asyncio +async def test_leaving_lar1_goes_back_to_the_router_strategy(make_core: MakeCore, host: ScriptedHost) -> None: + core: Final = make_core(model_list=[deployment("a", "a1")], num_retries=0) + core.apply(ChangeSet(settings={"routing_strategy": "lar1"})) + core.apply(ChangeSet(settings={"routing_strategy": "simple-shuffle"})) + core.completion(model="a", messages=MESSAGES) # pyright: ignore[reportAttributeAccessIssue, reportUnknownMemberType] # an endpoint + assert host.deployments_called("completion") == ["a1"] + + +# --- R4: endpoint eligibility --------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_acompletion_skips_a_deployment_that_enforces_other_endpoints( + make_core: MakeCore, host: ScriptedHost +) -> None: + core: Final = make_core(model_list=list(_embedding_group()), num_retries=0) + await core.acompletion(model="g", messages=MESSAGES) # pyright: ignore[reportAttributeAccessIssue, reportUnknownMemberType] # an endpoint + assert host.deployments_called("acompletion") == ["chat"] + + +@pytest.mark.asyncio +async def test_atext_completion_skips_a_deployment_that_enforces_other_endpoints( + make_core: MakeCore, host: ScriptedHost +) -> None: + core: Final = make_core(model_list=list(_embedding_group()), num_retries=0) + await core.atext_completion(model="g", prompt="hi") # pyright: ignore[reportAttributeAccessIssue, reportUnknownMemberType] # an endpoint + assert host.deployments_called("atext_completion") == ["chat"] + + +def test_sync_text_completion_skips_a_deployment_that_enforces_other_endpoints( + make_core: MakeCore, host: ScriptedHost +) -> None: + core: Final = make_core(model_list=list(_embedding_group()), num_retries=0) + core.text_completion(model="g", prompt="hi") # pyright: ignore[reportAttributeAccessIssue, reportUnknownMemberType] # an endpoint + # The sync text path sends the group as ``model`` and no ``model_info``: the params name the deployment. + assert [call.kwargs["weight"] for call in host.calls_to("text_completion")] == [0] + + +def test_sync_completion_skips_a_deployment_that_enforces_other_endpoints( + make_core: MakeCore, host: ScriptedHost +) -> None: + """``function_with_fallbacks`` is the async walk, whose ``endpoint_selection`` covers the sync attempt.""" + core: Final = make_core(model_list=list(_embedding_group()), num_retries=0) + core.completion(model="g", messages=MESSAGES) # pyright: ignore[reportAttributeAccessIssue, reportUnknownMemberType] # an endpoint + assert host.deployments_called("completion") == ["chat"] + + +@pytest.mark.asyncio +async def test_an_endpoint_outside_the_chat_table_keeps_every_deployment( + make_core: MakeCore, host: ScriptedHost +) -> None: + core: Final = make_core(model_list=list(_embedding_group()), num_retries=0) + await core.aembedding(model="g", input="hi") # pyright: ignore[reportAttributeAccessIssue, reportUnknownMemberType] # an endpoint + assert host.deployments_called("aembedding") == ["embonly"] + + +@pytest.mark.asyncio +async def test_a_group_that_serves_no_chat_endpoint_names_what_it_serves(make_core: MakeCore) -> None: + core: Final = make_core( + model_list=[deployment("g", "e1", model_info=dict(EMBEDDINGS_ONLY))], + num_retries=0, + ) + with pytest.raises(EndpointNotSupportedByModelError) as refused: + await core.acompletion(model="g", messages=MESSAGES) # pyright: ignore[reportAttributeAccessIssue, reportUnknownMemberType] # an endpoint + assert ( + refused.value.message + == "/v1/chat/completions is not supported by model=g. Supported endpoints: /v1/embeddings" + ) + + +@pytest.mark.asyncio +async def test_the_refusal_falls_back_like_any_bad_request_naming_no_provider( + make_core: MakeCore, host: ScriptedHost +) -> None: + core: Final = make_core( + model_list=[deployment("g", "e1", model_info=dict(EMBEDDINGS_ONLY)), deployment("h", "h1")], + num_retries=0, + fallbacks=[{"g": ["h"]}], + ) + await core.acompletion(model="g", messages=MESSAGES) # pyright: ignore[reportAttributeAccessIssue, reportUnknownMemberType] # an endpoint + assert host.deployments_called("acompletion") == ["h1"] + + +@pytest.mark.asyncio +async def test_a_direct_pick_runs_no_endpoint_selection(make_core: MakeCore) -> None: + core: Final = make_core(model_list=[deployment("g", "e1", model_info=dict(EMBEDDINGS_ONLY))], num_retries=0) + assert await picked_ids(core, "g", times=1) == {"e1"} + + +def test_routable_deployments_read_the_call_type(make_core: MakeCore) -> None: + core: Final = make_core(model_list=list(_embedding_group()), num_retries=0) + routable: Final = core.registry().deployments("g", source="routable", call_type="completion") + assert ids_of(list(routable)) == ["chat"] + + +# --- R5: web search ------------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_a_web_search_request_skips_deployments_that_declare_no_support( + make_core: MakeCore, host: ScriptedHost +) -> None: + core: Final = make_core(model_list=list(_web_group()), num_retries=0) + await core.acompletion(model="g", messages=MESSAGES, tools=WEB_SEARCH_TOOLS) # pyright: ignore[reportAttributeAccessIssue, reportUnknownMemberType] # an endpoint + assert host.deployments_called("acompletion") == ["web"] + + +def test_the_sync_path_runs_no_web_search_filter(make_core: MakeCore, host: ScriptedHost) -> None: + core: Final = make_core(model_list=list(_web_group()), num_retries=0) + core.completion(model="g", messages=MESSAGES, tools=WEB_SEARCH_TOOLS) # pyright: ignore[reportAttributeAccessIssue, reportUnknownMemberType] # an endpoint + assert host.deployments_called("completion") == ["noweb"] + + +@pytest.mark.asyncio +async def test_no_deployment_with_web_search_is_a_no_deployments_error(make_core: MakeCore) -> None: + core: Final = make_core( + model_list=[deployment("g", "noweb", model_info={"supports_web_search": False})], + num_retries=0, + ) + with pytest.raises(RouterRateLimitError): + await core.acompletion(model="g", messages=MESSAGES, tools=WEB_SEARCH_TOOLS) # pyright: ignore[reportAttributeAccessIssue, reportUnknownMemberType] # an endpoint + + +@pytest.mark.asyncio +async def test_the_healthy_listing_runs_the_web_search_filter(make_core: MakeCore) -> None: + core: Final = make_core(model_list=list(_web_group()), num_retries=0) + listed: Final = await core.select("g", {"tools": WEB_SEARCH_TOOLS}, kind="healthy") # pyright: ignore[reportCallIssue, reportArgumentType] # PythonCore's select takes keyword options + assert ids_of(listed) == ["web"] + + +# --- R6: order ------------------------------------------------------------------------------------ + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("first", "second"), + [ + ({"model_info": {"order": 1}}, {"model_info": {"order": 2}}), + ({"order": "1"}, {"order": "2"}), + ({"model_info": {"order": "1"}}, {"model_info": {"order": "2"}}), + ({"order": None, "model_info": {"order": 1}}, {"model_info": {"order": 2}}), + ], + ids=["model-info", "numeric-string", "model-info-numeric-string", "null-param-reads-model-info"], +) +async def test_order_is_read_where_get_deployment_order_reads_it( + make_core: MakeCore, first: dict[str, object], second: dict[str, object] +) -> None: + core: Final = make_core( + model_list=[deployment("g", "first", weight=0, **first), deployment("g", "second", weight=1, **second)], + num_retries=0, + ) + assert await picked_ids(core, "g", times=SPREAD_PICKS) == {"first"} + + +# --- R10: priority through an alias --------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_a_priority_through_an_alias_reads_the_target_group(make_core: MakeCore, host: ScriptedHost) -> None: + core: Final = make_core( + model_list=[deployment("g", "d1", drop_params=True)], + num_retries=0, + model_group_alias={"al": "g"}, + ) + await core.acompletion(model="al", messages=MESSAGES, priority="high") # pyright: ignore[reportAttributeAccessIssue, reportUnknownMemberType] # an endpoint + assert host.deployments_called("acompletion") == ["d1"] + + +# --- R11: old-style custom strategies ------------------------------------------------------------- + + +class _OldStyleStrategy(CustomRoutingStrategyBase): + """Written before picks took ``request_kwargs``.""" + + def __init__(self, picked: Mapping[str, object]) -> None: + self.picked: Final = picked + + async def async_get_available_deployment( # pyright: ignore[reportIncompatibleMethodOverride] # the old signature is the point + self, model: str, messages: object = None, input: object = None, specific_deployment: bool | None = False + ) -> Mapping[str, object]: + return self.picked + + def get_available_deployment( # pyright: ignore[reportIncompatibleMethodOverride] # the old signature is the point + self, model: str, messages: object = None, input: object = None, specific_deployment: bool | None = False + ) -> Mapping[str, object]: + return self.picked + + +def _with_old_style_strategy(make_core: MakeCore) -> RouterCore: + picked: Final = deployment("b", "db") + core: Final = make_core(model_list=[deployment("a", "da"), picked], num_retries=0) + core.register_plugin(_OldStyleStrategy(picked)) + return core + + +@pytest.mark.asyncio +async def test_an_old_style_strategy_answers_a_pick_that_names_only_the_model(make_core: MakeCore) -> None: + core: Final = _with_old_style_strategy(make_core) + picked: Final = await core._core_async_get_available_deployment("b") # pyright: ignore[reportAttributeAccessIssue, reportUnknownMemberType, reportPrivateUsage] # the member the deprecated pick returns + assert ids_of(picked) == ["db"] + + +def test_an_old_style_strategy_answers_a_sync_pick_that_names_only_the_model(make_core: MakeCore) -> None: + core: Final = _with_old_style_strategy(make_core) + assert ids_of(core.get_available_deployment("b")) == ["db"] # pyright: ignore[reportAttributeAccessIssue, reportUnknownMemberType] # the sync pick + + +@pytest.mark.asyncio +async def test_an_endpoint_always_hands_the_strategy_its_request(make_core: MakeCore) -> None: + core: Final = _with_old_style_strategy(make_core) + with pytest.raises(TypeError, match="request_kwargs"): + await core.acompletion(model="b", messages=MESSAGES) # pyright: ignore[reportAttributeAccessIssue, reportUnknownMemberType] # an endpoint + + +# --- N1, N15: wildcards --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_an_alias_to_a_wildcard_fills_the_model_from_the_alias_target( + make_core: MakeCore, host: ScriptedHost +) -> None: + core: Final = make_core( + model_list=[deployment("anthropic/*", "w", model="anthropic/*")], + num_retries=0, + model_group_alias={"al": "anthropic/claude-3-haiku"}, + ) + await core.acompletion(model="al", messages=MESSAGES) # pyright: ignore[reportAttributeAccessIssue, reportUnknownMemberType] # an endpoint + assert host.calls_to("acompletion")[0].model == "anthropic/claude-3-haiku" + + +@pytest.mark.asyncio +async def test_a_name_no_pattern_matches_is_tried_under_its_provider(make_core: MakeCore) -> None: + core: Final = make_core(model_list=[deployment("openai/*", "w0", model="openai/*")], num_retries=0) + listed: Final = await core.select("gpt-4o", {"user": "u"}, kind="healthy") # pyright: ignore[reportCallIssue, reportArgumentType] # PythonCore's select takes keyword options + assert ids_of(listed) == ["w0"] + assert isinstance(listed, list) + assert listed[0]["litellm_params"]["model"] == "openai/gpt-4o" + + +@pytest.mark.asyncio +async def test_a_call_through_the_provider_spelling_is_served(make_core: MakeCore, host: ScriptedHost) -> None: + core: Final = make_core(model_list=[deployment("openai/*", "w0", model="openai/*")], num_retries=0) + await core.acompletion(model="gpt-4o", messages=MESSAGES) # pyright: ignore[reportAttributeAccessIssue, reportUnknownMemberType] # an endpoint + assert host.calls_to("acompletion")[0].model == "openai/gpt-4o" + + +# --- #49685: wildcard cache flags ----------------------------------------------------------------- + +CACHE_FLAGS: Final = {"supports_prompt_cache_breakpoint": True, "supports_prompt_caching": True} + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("pattern", "request_model", "kept"), + [ + ("openai/*", "openai/gpt-4o", {}), + ("azure_ai/*", "azure_ai/gpt-4o", {"supports_prompt_cache_breakpoint": True}), + ], + ids=["any-wildcard", "azure-ai-keeps-the-breakpoint-flag"], +) +async def test_a_wildcard_leaves_the_cache_flags_to_the_catalog( + make_core: MakeCore, host: ScriptedHost, pattern: str, request_model: str, kept: Mapping[str, object] +) -> None: + core: Final = make_core( + model_list=[deployment(pattern, "w", model=pattern, model_info=dict(CACHE_FLAGS))], + num_retries=0, + ) + await core.acompletion(model=request_model, messages=MESSAGES) # pyright: ignore[reportAttributeAccessIssue, reportUnknownMemberType] # an endpoint + model_info: Final = host.calls_to("acompletion")[0].kwargs["model_info"] + assert isinstance(model_info, Mapping) + assert {flag: model_info[flag] for flag in CACHE_FLAGS if flag in model_info} == kept + + +@pytest.mark.asyncio +async def test_a_wildcard_group_with_a_fixed_model_keeps_its_cache_flags( + make_core: MakeCore, host: ScriptedHost +) -> None: + core: Final = make_core( + model_list=[deployment("openai/*", "w", model="openai/gpt-4o", model_info=dict(CACHE_FLAGS))], + num_retries=0, + ) + await core.acompletion(model="openai/whatever", messages=MESSAGES) # pyright: ignore[reportAttributeAccessIssue, reportUnknownMemberType] # an endpoint + model_info: Final = host.calls_to("acompletion")[0].kwargs["model_info"] + assert isinstance(model_info, Mapping) + assert model_info["supports_prompt_caching"] is True + + +def test_the_wildcard_route_lookup_leaves_the_cache_flags_to_the_catalog(make_core: MakeCore) -> None: + core: Final = make_core( + model_list=[deployment("openai/*", "w", model="openai/*", model_info=dict(CACHE_FLAGS))], + num_retries=0, + ) + routed: Final = core.registry().deployments("openai/gpt-4o", source="wildcard") + assert not set(CACHE_FLAGS) & set(routed[0]["model_info"]) + + +# --- N12: specific_deployment --------------------------------------------------------------------- + + +def test_only_true_names_a_litellm_model(make_core: MakeCore) -> None: + core: Final = make_core(model_list=[deployment("gpt-4o", "d1")], num_retries=0) + with pytest.raises(litellm.BadRequestError): + core.get_available_deployment( # pyright: ignore[reportAttributeAccessIssue, reportUnknownMemberType] # the sync pick + model="gpt-4o ", specific_deployment="yes", request_kwargs={"input": "hi"} + ) + + +# --- N14: pass-through picks ---------------------------------------------------------------------- + +COMPLEXITY_GROUP: Final = [ + deployment("g1", "b-1", model="openai/gpt-4o", use_in_pass_through=True), + deployment("g2", "b-2", model="openai/gpt-4o-mini", use_in_pass_through=True), + { + "model_name": "smart", + "litellm_params": { + "model": "auto_router/complexity_router", + "complexity_router_config": {"tiers": {"SIMPLE": "g2", "REASONING": "g1"}}, + "complexity_router_default_model": "g2", + }, + }, +] + + +@pytest.mark.asyncio +async def test_a_listed_pass_through_pick_names_the_group_the_alias_resolved_to(make_core: MakeCore) -> None: + core: Final = make_core( + model_list=[deployment("team-gpt", "d1")], + num_retries=0, + model_group_alias={"alias-b": "team-gpt"}, + ) + with pytest.raises(litellm.BadRequestError, match="Model team-gpt has no deployment"): + await core.select("alias-b", {}, kind="pass_through") # pyright: ignore[reportCallIssue] # PythonCore's select takes keyword options + + +@pytest.mark.asyncio +async def test_an_async_pass_through_pick_of_a_strategy_marker_group_takes_the_tier(make_core: MakeCore) -> None: + core: Final = make_core( + model_list=list(COMPLEXITY_GROUP), num_retries=0, routing_strategy="usage-based-routing-v2" + ) + picked: Final = await core.async_get_available_deployment_for_pass_through( # pyright: ignore[reportAttributeAccessIssue, reportUnknownMemberType] # the async pass-through pick + "smart", {}, messages=[{"role": "user", "content": "what is up"}] + ) + assert ids_of(picked) == ["b-2"] + + +# --- N16: affinity -------------------------------------------------------------------------------- + +ENCRYPTED_INPUT: Final = [ + {"role": "user", "content": "hi"}, + { + "type": "reasoning", + "summary": [{"type": "summary_text", "text": "thought"}], + "encrypted_content": "litellm_enc:bW9kZWxfaWQ6bTktMg==;gAAAAABenc", + }, + {"type": "reasoning", "encrypted_content": "litellm_enc:bW9kZWxfaWQ6bTktMQ==;gAAAAABenc", "summary": []}, +] + + +@pytest.mark.asyncio +async def test_encrypted_affinity_reads_the_ids_the_pick_excludes(make_core: MakeCore) -> None: + core: Final = make_core( + model_list=[ + deployment("h", "m9-0", order=1), + deployment("h", "m9-1", order=1), + deployment("g", "m9-2"), + ], + num_retries=0, + model_group_affinity_config={"h": ["encrypted_content_affinity"]}, + ) + request: Final = { + "input": list(ENCRYPTED_INPUT), + "_excluded_deployment_ids": ["m9-1"], + "_retry_skipped_deployment_ids": ["m9-2", "m9-0"], + } + assert ids_of(await core.select("h", request, input=list(ENCRYPTED_INPUT))) == ["m9-0"] # pyright: ignore[reportCallIssue, reportArgumentType] # PythonCore's select takes keyword options + + +@pytest.mark.asyncio +async def test_a_target_order_on_a_call_narrows_its_pick(make_core: MakeCore) -> None: + core: Final = make_core( + model_list=[deployment("gpt", "id-a"), deployment("gpt", "id-b")], + num_retries=0, + optional_pre_call_checks=["session_affinity", "deployment_affinity"], + ) + with pytest.raises(RouterRateLimitError): + await core.acompletion(model="gpt", messages=MESSAGES, _target_order=1) # pyright: ignore[reportAttributeAccessIssue, reportUnknownMemberType] # an endpoint + + +@pytest.mark.asyncio +async def test_a_target_order_on_a_call_serves_that_level(make_core: MakeCore, host: ScriptedHost) -> None: + core: Final = make_core( + model_list=[deployment("gpt", "low", order=1), deployment("gpt", "high", order=2)], + num_retries=0, + ) + await core.acompletion(model="gpt", messages=MESSAGES, _target_order=2) # pyright: ignore[reportAttributeAccessIssue, reportUnknownMemberType] # an endpoint + assert host.deployments_called("acompletion") == ["high"] + assert "_target_order" not in host.calls_to("acompletion")[0].kwargs + + +# --- S7: a health scope as a dict ----------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_a_health_scope_given_as_a_dict_reads_its_keys(make_core: MakeCore) -> None: + core: Final = make_core( + model_list=[deployment("g", "d1", weight=1), deployment("g", "d2", weight=0), deployment("h", "h1")], + num_retries=0, + enable_health_check_routing=True, + background_health_check_model_groups={"h": True}, + ) + assert core.settings().background_health_check_model_groups == frozenset({"h"}) + core.health().health_checks.set_deployment_health_states( # pyright: ignore[reportArgumentType] # health states + {"d1": {"is_healthy": False, "timestamp": time.time(), "reason": "probe failed"}} + ) + assert await picked_ids(core, "g", times=SPREAD_PICKS) == {"d1"} diff --git a/tests/unit/test_router/core/test_parity_regress5_errors.py b/tests/unit/test_router/core/test_parity_regress5_errors.py new file mode 100644 index 00000000000..a0122f8eda6 --- /dev/null +++ b/tests/unit/test_router/core/test_parity_regress5_errors.py @@ -0,0 +1,501 @@ +""" +Parity of the Rust core with ``router_python_core.py`` on the error causes of the regress5 pack: when +the routing note is added and what it says, the error an order chain raises, what a blocked response +and a failed attempt stamp and record, the names the errors carry, and non-string names at the lookup +members. +""" + +import asyncio +import datetime +from collections.abc import Mapping +from typing import Final + +import pytest + +import litellm +from litellm.exceptions import GuardrailRaisedException, UnsupportedEndpointError +from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging +from litellm.router import Router +from litellm.router_core import RouterCore +from tests.unit.test_router.core.endpoints_helpers import EndpointHost, patch_direct_endpoints +from tests.unit.test_router.core.fallbacks_helpers import ( + MESSAGES, + MakeCore, + build, + complete, + complete_sync, + show_router_debug, +) +from tests.unit.test_router.core.scripted_host import ( + auth_error, + bad_request, + completion_response, + deployment, + rate_limit, + server_error, + timeout, +) +from tests.unit.test_router.core.streaming_helpers import chat_stream, collect_texts, install_orjson_stub_if_missing + +install_orjson_stub_if_missing() + +NOTE: Final = "LiteLLM: model group" +NO_FALLBACK: Final = "\n\nLiteLLM: model group 'a' failed with the error above. No fallback was attempted." +LONG_MESSAGE: Final = "x" * 900 + + +@pytest.fixture +def host(monkeypatch: pytest.MonkeyPatch) -> EndpointHost: + endpoint_host: Final = EndpointHost() + patch_direct_endpoints(monkeypatch, endpoint_host) + return endpoint_host + + +def as_router(core: RouterCore) -> Router: + router: Final = Router.__new__(Router) + router._attach(core) # pyright: ignore[reportPrivateUsage, reportAttributeAccessIssue, reportUnknownMemberType] # the facade's own attach + return router + + +def content_filtered() -> object: + response: Final = completion_response() + response.choices[0].finish_reason = "content_filter" # pyright: ignore[reportAttributeAccessIssue] # Choices + return response + + +def cooled(core: RouterCore, *ids: str) -> Mapping[str, Mapping[str, object]]: + return dict(core.health().cooldowns.get_active_cooldowns(list(ids), None)) # pyright: ignore[reportReturnType] # a mapping of entries + + +async def settle() -> None: + """Let the tasks a failed stream scheduled run.""" + for _ in range(20): + await asyncio.sleep(0) + + +async def failed_with(core: RouterCore, **request: object) -> BaseException: # kwargs-ok: request kwargs + with pytest.raises(Exception) as raised: # noqa: PT011 # any failure: each caller asserts on what it gets back + await complete(core, **request) + return raised.value + + +# --- N2: when the note is added ------------------------------------------------------------------- + + +async def test_disable_fallbacks_raises_with_no_note(make_core: MakeCore, host: EndpointHost) -> None: + core: Final = build(make_core, fallbacks=[{"a": ["b"]}]) + show_router_debug(host) + host.answer("openai/a", server_error()) + error: Final = await failed_with(core, disable_fallbacks=True) + assert NOTE not in str(error) + assert host.deployments_called() == ["a"] + + +async def test_a_guardrail_block_raises_with_no_note(make_core: MakeCore, host: EndpointHost) -> None: + core: Final = build(make_core, fallbacks=[{"a": ["b"]}]) + show_router_debug(host) + host.answer("openai/a", GuardrailRaisedException(guardrail_name="g", message="blocked")) + error: Final = await failed_with(core) + assert isinstance(error, GuardrailRaisedException) + assert NOTE not in str(error) + + +async def test_a_providers_bad_request_falls_back_and_notes_the_failed_fallback( + make_core: MakeCore, host: EndpointHost +) -> None: + core: Final = build(make_core, fallbacks=[{"a": ["b"]}]) + show_router_debug(host) + host.answer("openai/a", bad_request()) + host.answer("openai/b", bad_request()) + error: Final = await failed_with(core) + assert f"{NOTE} 'a' failed with the error above. Fallback to b also failed" in str(error) + assert host.deployments_called() == ["a", "b"] + + +def test_a_single_attempt_endpoint_raises_with_no_note(make_core: MakeCore, host: EndpointHost) -> None: + core: Final = build(make_core, fallbacks=[{"a": ["b"]}]) + show_router_debug(host) + host.answer("openai/a", server_error()) + with pytest.raises(litellm.InternalServerError) as raised: + core.responses(model="a", input="hi") # pyright: ignore[reportAttributeAccessIssue, reportUnknownMemberType] # an endpoint + assert NOTE not in str(raised.value) + + +async def test_an_empty_request_list_says_no_fallbacks_are_configured(make_core: MakeCore, host: EndpointHost) -> None: + core: Final = build(make_core) + show_router_debug(host) + host.answer("openai/a", rate_limit()) + message: Final = str(await failed_with(core, fallbacks=[])) + assert "and no fallback model group was found for it" in message + assert "No fallbacks are configured." in message + + +def test_a_strategy_marker_group_with_an_empty_list_says_no_fallbacks_are_configured( + make_core: MakeCore, host: EndpointHost +) -> None: + core: Final = make_core( + model_list=[ + { + "model_name": "qr", + "litellm_params": { + "model": "auto_router/quality_router", + "quality_router_default_model": "t1", + "quality_router_config": { + "available_models": ["t1", "t3"], + "complexity_to_quality": {"SIMPLE": 2, "REASONING": 3}, + }, + }, + "model_info": {"id": "qr"}, + }, + deployment("t1", "t1", model_info={"litellm_routing_preferences": {"quality_tier": 1}}), + deployment("t3", "t3", model_info={"litellm_routing_preferences": {"quality_tier": 3}}), + ], + num_retries=0, + ) + show_router_debug(host) + with pytest.raises( + Exception, match="model group 'qr' failed with the error above and no fallback model group was found" + ) as raised: + complete_sync(core, "qr", fallbacks=[]) + assert "No fallbacks are configured." in str(raised.value) + + +async def test_a_chain_that_only_names_the_failed_group_reads_the_primary_error( + make_core: MakeCore, host: EndpointHost +) -> None: + core: Final = build(make_core, fallbacks=[{"a": ["a"]}]) + show_router_debug(host) + host.answer("openai/a", rate_limit()) + assert str(await failed_with(core)).endswith( + "LiteLLM: model group 'a' failed with the error above. " + "Fallback to a also failed: litellm.RateLimitError: slow down" + ) + + +async def test_a_depth_spent_before_any_hop_reads_the_primary_error(make_core: MakeCore, host: EndpointHost) -> None: + core: Final = build(make_core, fallbacks=[{"a": ["b"]}], max_fallbacks=0) + show_router_debug(host) + host.answer("openai/a", rate_limit()) + assert str(await failed_with(core)).endswith("Fallback to b also failed: litellm.RateLimitError: slow down") + assert host.deployments_called() == ["a"] + + +# --- U11: an order chain --------------------------------------------------------------------------- + + +async def test_every_order_level_failing_raises_the_last_levels_error_with_no_note( + make_core: MakeCore, host: EndpointHost +) -> None: + core: Final = make_core( + model_list=[deployment("bb", "bb1", order=1), deployment("bb", "bb2", order=2)], num_retries=0 + ) + show_router_debug(host) + host.answer("openai/bb1", rate_limit(message="first")) + host.answer("openai/bb2", rate_limit(message="second")) + with pytest.raises(litellm.RateLimitError) as raised: + await complete(core, "bb") + assert str(raised.value) == "litellm.RateLimitError: second" + assert getattr(raised.value, "failed_deployment_id", None) == "bb2" + + +# --- R13: a blocked response ------------------------------------------------------------------------ + + +async def test_a_blocked_response_stamps_the_deployment(make_core: MakeCore, host: EndpointHost) -> None: + core: Final = build(make_core, content_policy_fallbacks=[{"a": ["a"]}]) + host.answer("openai/a", content_filtered()) + error: Final = await failed_with(core) + assert isinstance(error, litellm.ContentPolicyViolationError) + assert getattr(error, "failed_deployment_id", None) == "a" + + +def test_a_blocked_sync_response_stamps_the_deployment(make_core: MakeCore, host: EndpointHost) -> None: + core: Final = build(make_core, content_policy_fallbacks=[{"a": ["a"]}]) + host.answer("openai/a", content_filtered()) + with pytest.raises(litellm.ContentPolicyViolationError) as raised: + complete_sync(core) + assert getattr(raised.value, "failed_deployment_id", None) == "a" + + +async def test_a_blocked_response_counts_as_a_failure(make_core: MakeCore, host: EndpointHost) -> None: + core: Final = build(make_core, content_policy_fallbacks=[{"a": ["a"]}]) + host.answer("openai/a", content_filtered()) + await failed_with(core) + assert core.fail_calls["openai/a"] == 1 # pyright: ignore[reportAttributeAccessIssue] # the router's live counters + assert core.success_calls["openai/a"] == 0 # pyright: ignore[reportAttributeAccessIssue] # see above + + +# --- R12, S11, U6: what a cooldown keeps and what a timeout says -------------------------------------- + + +async def test_a_cooldown_keeps_the_timeout_as_it_was_raised(make_core: MakeCore, host: EndpointHost) -> None: + core: Final = make_core( + model_list=[deployment("a", "a1", weight=1), deployment("a", "a2", weight=0)], num_retries=0, allowed_fails=0 + ) + show_router_debug(host) + host.answer("openai/a1", timeout()) + error: Final = await failed_with(core) + assert "Deployment Info" in str(error) + assert len(str(cooled(core, "a1")["a1"]["exception_received"])) == len(str(timeout())) + + +async def test_a_cooldown_keeps_the_whole_error(make_core: MakeCore, host: EndpointHost) -> None: + core: Final = make_core( + model_list=[deployment("a", "a1", weight=1), deployment("a", "a2", weight=0)], num_retries=0, allowed_fails=0 + ) + host.answer("openai/a1", rate_limit(message=LONG_MESSAGE)) + await failed_with(core) + expected: Final = str(rate_limit(message=LONG_MESSAGE)) + assert len(str(cooled(core, "a1")["a1"]["exception_received"])) == len(expected) + + +async def test_a_stream_failing_on_its_first_chunk_cools_with_the_routing_note( + make_core: MakeCore, host: EndpointHost +) -> None: + core: Final = make_core(model_list=[deployment("a", "a")], num_retries=2, retry_after=0) + show_router_debug(host) + host.answer("openai/a", chat_stream(error=auth_error(), fail_after=0, logged=True)) + error: Final = await failed_with(core, stream=True) + await settle() + assert NOTE in str(error) + assert len(str(cooled(core, "a")["a"]["exception_received"])) == len(str(error)) + + +def test_a_sync_timeout_says_nothing_of_the_deployment(make_core: MakeCore, host: EndpointHost) -> None: + core: Final = build(make_core) + show_router_debug(host) + host.answer("openai/a", timeout()) + with pytest.raises(litellm.Timeout) as raised: + complete_sync(core) + assert "Deployment Info" not in str(raised.value) + + +async def test_an_async_timeout_counts_no_failure(make_core: MakeCore, host: EndpointHost) -> None: + core: Final = build(make_core) + host.answer("openai/a", timeout()) + await failed_with(core) + assert core.fail_calls["openai/a"] == 0 # pyright: ignore[reportAttributeAccessIssue] # the router's live counters + + +# --- R14, N19: a name nothing serves ------------------------------------------------------------------ + + +async def test_a_default_fallback_with_no_deployment_is_the_name_the_error_gives( + make_core: MakeCore, host: EndpointHost +) -> None: + core: Final = make_core(model_list=[deployment("g", "d1")], num_retries=0, default_fallbacks=["missing"]) + with pytest.raises(litellm.BadRequestError) as raised: + await complete(core, "nope") + assert raised.value.model == "missing" + assert "You passed in model=missing." in str(raised.value) + + +def test_an_empty_router_names_the_first_default_fallback(make_core: MakeCore) -> None: + core: Final = make_core(model_list=[], num_retries=1, retry_after=0, default_fallbacks=["team-gpt", "gpt-4o-mini"]) + with pytest.raises(litellm.BadRequestError) as raised: + complete_sync(core, "gpt-4o") + assert raised.value.model == "team-gpt" + + +async def test_an_unknown_alias_target_is_the_name_the_error_gives(make_core: MakeCore, host: EndpointHost) -> None: + core: Final = make_core(model_list=[deployment("g", "d1")], num_retries=0, model_group_alias={"al": "nowhere"}) + with pytest.raises(litellm.BadRequestError) as raised: + await complete(core, "al") + assert raised.value.model == "nowhere" + + +async def test_a_missing_message_list_through_a_wildcard_alias_reaches_the_call( + make_core: MakeCore, host: EndpointHost +) -> None: + core: Final = make_core( + model_list=[deployment("openai/*", "w", model="openai/*"), deployment("g", "d1")], + num_retries=0, + model_group_alias={"al": {"hidden": False, "model": "gpt-4o-mini"}}, + ) + await core.acompletion(model="al", messages=None) # pyright: ignore[reportAttributeAccessIssue, reportUnknownMemberType] # an endpoint + assert host.deployments_called("acompletion") == ["w"] + + +async def test_moderation_of_a_name_only_a_wildcard_serves_uses_the_wildcard( + make_core: MakeCore, host: EndpointHost +) -> None: + core: Final = make_core(model_list=[deployment("openai/*", "w", model="openai/*")], num_retries=0) + await core.amoderation(model="openai/omni-moderation-latest", input="hello") # pyright: ignore[reportAttributeAccessIssue, reportUnknownMemberType] # an endpoint + assert host.deployments_called("amoderation") == ["w"] + + +# --- R15, T7: pass-through ------------------------------------------------------------------------ + + +async def test_a_listed_pick_of_an_id_names_its_litellm_model(make_core: MakeCore) -> None: + core: Final = make_core( + model_list=[deployment("g", "p1", model="openai/b"), deployment("g", "p4", use_in_pass_through=True)], + num_retries=0, + ) + with pytest.raises(litellm.BadRequestError, match="Deployment p1 does not support pass-through") as raised: + await core.select("p1", {}, kind="pass_through") # pyright: ignore[reportCallIssue] # PythonCore's select takes keyword options + assert raised.value.model == "openai/b" + + +async def test_an_async_pick_of_an_id_names_the_id(make_core: MakeCore) -> None: + core: Final = make_core( + model_list=[deployment("g", "p1", model="openai/b"), deployment("g", "p4", use_in_pass_through=True)], + num_retries=0, + ) + with pytest.raises(litellm.BadRequestError, match="Deployment p1 does not support pass-through") as raised: + await core.async_get_available_deployment_for_pass_through("p1", {}, messages=MESSAGES) # pyright: ignore[reportAttributeAccessIssue, reportUnknownMemberType] # the async pass-through pick + assert raised.value.model == "p1" + + +def unknown_provider(model: str = "nope") -> litellm.BadRequestError: + return litellm.BadRequestError(message="LLM Provider NOT provided", model=model, llm_provider="") + + +async def test_a_pass_through_to_an_unknown_model_carries_the_note(make_core: MakeCore, host: EndpointHost) -> None: + core: Final = make_core(model_list=[deployment("g", "d1")], num_retries=0) + show_router_debug(host) + host.answer("nope", unknown_provider()) + with pytest.raises(litellm.BadRequestError) as raised: + await core.allm_passthrough_route(model="nope", method="GET", endpoint="/models") # pyright: ignore[reportAttributeAccessIssue, reportUnknownMemberType] # an endpoint + assert str(raised.value).endswith( + "\n\nLiteLLM: model group 'nope' failed with the error above. No fallback was attempted." + ) + + +async def test_a_pass_through_with_no_deployment_calls_the_handler_before_any_fallback( + make_core: MakeCore, host: EndpointHost +) -> None: + core: Final = make_core(model_list=[deployment("g", "d1")], num_retries=0, fallbacks=[{"nope": ["g"]}]) + host.answer("nope", unknown_provider()) + await core.allm_passthrough_route(model="nope", method="GET", endpoint="/models") # pyright: ignore[reportAttributeAccessIssue, reportUnknownMemberType] # an endpoint + assert [call.model for call in host.calls_to("allm_passthrough_route")] == ["nope", "openai/d1"] + + +# --- S9: the deployment an exception names -------------------------------------------------------- + + +def unsupported() -> UnsupportedEndpointError: + return UnsupportedEndpointError(message="not here", llm_provider="openai", model="gpt-4o") + + +async def test_an_exception_naming_a_retried_deployment_exhausts_the_retries( + make_core: MakeCore, host: EndpointHost +) -> None: + core: Final = make_core( + model_list=[deployment("g", "d1", weight=1), deployment("g", "d2", weight=0)], num_retries=1, retry_after=0 + ) + reused: Final = unsupported() + reused.failed_deployment_id = "d1" # pyright: ignore[reportAttributeAccessIssue] # what the first failure stamped + host.answer("openai/d1", unsupported()) + host.answer("openai/d2", reused) + with pytest.raises(UnsupportedEndpointError) as raised: + await complete(core, "g") + assert (raised.value.num_retries, raised.value.max_retries) == (1, 1) + assert "LiteLLM Retried: 1 times" in str(raised.value) + assert host.deployments_called("acompletion") == ["d1", "d2"] + + +async def test_every_deployment_refusing_raises_before_the_retries_are_exhausted( + make_core: MakeCore, host: EndpointHost +) -> None: + core: Final = make_core( + model_list=[deployment("g", "d1", weight=1), deployment("g", "d2", weight=0)], num_retries=1, retry_after=0 + ) + host.answer("*", unsupported()) + with pytest.raises(UnsupportedEndpointError) as raised: + await complete(core, "g") + assert raised.value.max_retries is None + assert host.deployments_called("acompletion") == ["d1", "d2"] + + +# --- N5: model=None ----------------------------------------------------------------------------------- + + +async def test_model_none_is_a_bad_request(make_core: MakeCore, host: EndpointHost) -> None: + core: Final = make_core(model_list=[deployment("g", "d1")], num_retries=0) + show_router_debug(host) + with pytest.raises(litellm.BadRequestError, match="You passed in model=None") as raised: + await core.acompletion(model=None, messages=MESSAGES) # pyright: ignore[reportAttributeAccessIssue, reportUnknownMemberType] # an endpoint + assert raised.value.status_code == 400 + assert NOTE not in str(raised.value) + + +def test_sync_model_none_is_a_bad_request(make_core: MakeCore) -> None: + core: Final = make_core(model_list=[deployment("g", "d1")], num_retries=0) + with pytest.raises(litellm.BadRequestError, match="You passed in model=None"): + core.completion(model=None, messages=MESSAGES) # pyright: ignore[reportAttributeAccessIssue, reportUnknownMemberType] # an endpoint + + +# --- N4: lookup members given a non-string name --------------------------------------------------- + + +def test_a_non_string_name_is_not_recognized(make_core: MakeCore) -> None: + core: Final = make_core(model_list=[deployment("g", "d1")], num_retries=0) + assert core.registry().resolve(True, kind="recognized") is False + + +def test_an_unhashable_name_raises_as_the_index_lookup_raises(make_core: MakeCore) -> None: + core: Final = make_core(model_list=[deployment("g", "d1")], num_retries=0) + with pytest.raises(TypeError, match="unhashable type: 'list'"): + core.registry().resolve(["g"], kind="recognized") # pyright: ignore[reportArgumentType] # raising is the point + + +@pytest.mark.parametrize( + ("member", "expected"), + [ + ("get_configured_display_name", None), + ("get_configured_mode", None), + ("get_deployment", None), + ("get_model_ids", []), + ], +) +def test_a_lookup_given_an_int_finds_nothing(make_core: MakeCore, member: str, expected: object) -> None: + router: Final = as_router(make_core(model_list=[deployment("g", "d1")], num_retries=0)) + with pytest.warns(DeprecationWarning, match="is deprecated"): + assert getattr(router, member)(1) == expected # pyright: ignore[reportAny] # each old member's own type + + +def test_a_model_list_lookup_given_an_int_raises_as_the_wildcard_step_raises(make_core: MakeCore) -> None: + router: Final = as_router(make_core(model_list=[deployment("g", "d1")], num_retries=0)) + with ( + pytest.warns(DeprecationWarning, match="is deprecated"), + pytest.raises(TypeError, match="argument of type 'int' is not"), + ): + router.get_model_list(1) # pyright: ignore[reportArgumentType] # raising is the point + + +# --- V6: a base_model that is no string ----------------------------------------------------------- + + +def test_an_int_base_model_fails_as_startswith_fails(make_core: MakeCore) -> None: + core: Final = make_core( + model_list=[deployment("g", "az1", model="azure/my-dep", api_base="https://x", base_model=7)], num_retries=0 + ) + with pytest.raises(AttributeError, match=r"^'int' object has no attribute 'startswith'$"): + core.registry().router_model_info(None, "g", id="az1") + + +# --- a shared logging object ---------------------------------------------------------------------- + + +async def test_a_failure_already_logged_blocks_a_streams_count(make_core: MakeCore, host: EndpointHost) -> None: + core: Final = make_core( + model_list=[deployment("a", "a1", weight=1), deployment("a", "a2", weight=0)], num_retries=0, allowed_fails=0 + ) + logging_obj: Final = LiteLLMLogging( + model="a", + messages=MESSAGES, + stream=True, + call_type="acompletion", + start_time=datetime.datetime(2026, 1, 1, tzinfo=datetime.timezone.utc), + litellm_call_id="c", + function_id="f", + ) + logging_obj.update_environment_variables(model="a", user=None, optional_params={}, litellm_params={}) + logging_obj.model_call_details["has_logged_async_failure"] = True + host.answer("openai/a1", chat_stream("hi", error=rate_limit(), fail_after=1)) + stream: Final = await complete(core, stream=True, litellm_logging_obj=logging_obj) + with pytest.raises(litellm.RateLimitError): + await collect_texts(stream) + await settle() + assert cooled(core, "a1", "a2") == {} diff --git a/tests/unit/test_router/core/test_parity_regress5_settings.py b/tests/unit/test_router/core/test_parity_regress5_settings.py new file mode 100644 index 00000000000..6d3c8dea7ce --- /dev/null +++ b/tests/unit/test_router/core/test_parity_regress5_settings.py @@ -0,0 +1,473 @@ +""" +Parity of the Rust core with ``router_python_core.py`` on the settings causes of the regress5 pack: +attributes written after build, ``update_settings`` applied key by key, a request ``num_retries`` +that is no count, calls after ``retire``, deleted attributes, and a default deployment that fails +every call. The refusals at the end are ruled to differ, so they run on the Rust core only. +""" + +import inspect +from typing import Final + +import pytest + +import litellm +from litellm.router_api import ChangeSet +from litellm.router_core import RouterCore +from litellm.types.router import Deployment, LiteLLM_Params, RetryPolicy +from tests.unit.test_router.core.gaps_l11100_helpers import TIERS, complexity_marker +from tests.unit.test_router.core.retries_health_helpers import MakeCore, attempt +from tests.unit.test_router.core.scripted_host import ( + CORE_BUILDERS, + ScriptedHost, + deployment, + rate_limit, + server_error, +) + +HEURISTIC_V2: Final = {"classifier_type": "heuristic_v2", "tiers": dict(TIERS)} + + +def _one_group(make_core: MakeCore, **arguments: object) -> RouterCore: # kwargs-ok: the router's constructor arguments + return make_core(model_list=[deployment("g", "d1")], **{"num_retries": 0, **arguments}) + + +# --- U7: retry policies written after build --------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_retry_policy_assigned_after_build_governs_the_next_failure( + make_core: MakeCore, host: ScriptedHost +) -> None: + core: Final = _one_group(make_core, num_retries=1, retry_after=0, disable_cooldowns=True) + core.retry_policy = RetryPolicy(RateLimitErrorRetries=0) # pyright: ignore[reportAttributeAccessIssue] # a plain attribute of both cores + host.answer("*", lambda **_: rate_limit()) + + error: Final = await attempt(core) + + assert isinstance(error, litellm.RateLimitError) + assert host.deployments_called("acompletion") == ["d1"] + + +def test_model_group_retry_policy_assigned_after_build_answers_the_lookup(make_core: MakeCore) -> None: + core: Final = _one_group(make_core) + policies: Final = {"g": RetryPolicy(RateLimitErrorRetries=1)} + core.model_group_retry_policy = policies # pyright: ignore[reportAttributeAccessIssue] # see above + + assert core.model_group_retry_policy == policies # pyright: ignore[reportAttributeAccessIssue] # see above + assert core.get_num_retries_from_retry_policy(rate_limit(), "g") == 1 # pyright: ignore[reportAttributeAccessIssue] # see above + + +@pytest.mark.asyncio +async def test_a_retry_policy_written_as_a_dict_stays_a_dict_and_applies( + make_core: MakeCore, host: ScriptedHost +) -> None: + core: Final = _one_group(make_core, num_retries=1, retry_after=0, disable_cooldowns=True) + core.retry_policy = {"RateLimitErrorRetries": 0} # pyright: ignore[reportAttributeAccessIssue] # stored as given + host.answer("*", lambda **_: rate_limit()) + + error: Final = await attempt(core) + + assert isinstance(error, litellm.RateLimitError) + assert core.retry_policy == {"RateLimitErrorRetries": 0} # pyright: ignore[reportAttributeAccessIssue] # see above + assert host.deployments_called("acompletion") == ["d1"] + + +@pytest.mark.asyncio +async def test_a_retry_policy_written_as_no_policy_raises_on_the_first_failure_only( + make_core: MakeCore, host: ScriptedHost +) -> None: + core: Final = _one_group(make_core, num_retries=1, retry_after=0, disable_cooldowns=True) + core.retry_policy = 5 # pyright: ignore[reportAttributeAccessIssue] # stored as given + + served: Final = await attempt(core) + host.answer("*", lambda **_: rate_limit()) + error: Final = await attempt(core) + + assert not isinstance(served, Exception) + assert isinstance(error, AttributeError) + assert "DefaultRetries" in str(error) + + +@pytest.mark.parametrize( + ("name", "value"), [("enable_weighted_failover", True), ("fallback_on_unsupported_params", False)] +) +def test_other_plain_settings_read_back_what_was_written(make_core: MakeCore, name: str, value: object) -> None: + core: Final = _one_group(make_core) + setattr(core, name, value) + assert getattr(core, name) is value + + +# --- N8: update_settings applies its keys in order, with no rollback -------------------------------- + + +def test_update_settings_keeps_the_keys_before_the_one_that_raises(make_core: MakeCore) -> None: + core: Final = _one_group(make_core) + with pytest.raises(TypeError): + core.update_settings( # pyright: ignore[reportAttributeAccessIssue] # the deprecated member + context_window_fallbacks=[{"g": ["h"]}], allowed_fails=["not a number"], timeout=3 + ) + assert core.settings().context_window_fallbacks == [{"g": ["h"]}] + assert core.timeout != 3 # pyright: ignore[reportAttributeAccessIssue] # a plain attribute of both cores + + +@pytest.mark.parametrize("strategy", [1, "bogus"]) +def test_update_settings_refuses_an_unknown_strategy(make_core: MakeCore, strategy: object) -> None: + core: Final = _one_group(make_core) + with pytest.raises(ValueError, match="Invalid routing_strategy"): + core.update_settings(fallbacks=[{"g": ["h"]}], routing_strategy=strategy, timeout=3) # pyright: ignore[reportAttributeAccessIssue] # see above + assert core.settings().routing_strategy == "simple-shuffle" + assert core.settings().fallbacks == [{"g": ["h"]}] + + +def test_update_settings_keeps_no_strategy_as_given(make_core: MakeCore) -> None: + core: Final = _one_group(make_core) + core.update_settings(routing_strategy=None) # pyright: ignore[reportAttributeAccessIssue] # see above + assert core.settings().routing_strategy is None + + +def test_a_failed_setting_leaves_the_groups_passed_before_it_unread(make_core: MakeCore) -> None: + core: Final = _one_group(make_core) + groups: Final = [{"group_name": "z", "models": ["g"], "routing_strategy": "least-busy"}] + with pytest.raises(ValueError, match="invalid literal"): + core.update_settings(routing_groups=groups, timeout="soon") # pyright: ignore[reportAttributeAccessIssue] # see above + assert core.settings().as_dict()["routing_groups"] == [] + + +def test_groups_that_do_not_parse_raise_after_every_other_key(make_core: MakeCore) -> None: + core: Final = _one_group(make_core) + with pytest.raises(Exception, match="group_name"): + core.update_settings(timeout=6, routing_groups=[{"bad": 1}], retry_after=5) # pyright: ignore[reportAttributeAccessIssue] # see above + assert (core.timeout, core.retry_after) == (6, 5) # pyright: ignore[reportAttributeAccessIssue] # see above + + +def test_optional_pre_call_checks_that_cannot_be_read_stop_the_keys_after_them(make_core: MakeCore) -> None: + core: Final = _one_group(make_core) + with pytest.raises(TypeError, match="not iterable"): + core.update_settings(timeout=5, optional_pre_call_checks=1, retry_after=4) # pyright: ignore[reportAttributeAccessIssue] # see above + assert core.timeout == 5 # pyright: ignore[reportAttributeAccessIssue] # see above + assert core.retry_after != 4 # pyright: ignore[reportAttributeAccessIssue] # see above + + +def test_a_changeset_stops_at_the_setting_that_raises(make_core: MakeCore) -> None: + core: Final = make_core(model_list=[deployment("g", "d1"), deployment("h", "d2")], num_retries=0) + with pytest.raises(TypeError): + core.apply(ChangeSet(settings={"fallbacks": [{"g": ["h"]}], "allowed_fails": [1]}, delete=["d1"])) + assert core.settings().fallbacks == [{"g": ["h"]}] + assert sorted(core.registry().names(kind="ids")) == ["d1", "d2"] + + +# --- N7: a request num_retries is compared, never cast ---------------------------------------------- + + +@pytest.mark.asyncio +@pytest.mark.parametrize("value", ["2", 1.5, [1], -1, float("nan"), True, False, 0.0]) +async def test_a_request_that_succeeds_never_reads_its_num_retries(make_core: MakeCore, value: object) -> None: + served: Final = await attempt(_one_group(make_core), num_retries=value) + assert not isinstance(served, Exception) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("value", "message"), + [ + ("2", "'>' not supported between instances of 'str' and 'int'"), + ([1], "'>' not supported between instances of 'list' and 'int'"), + (1.5, "'float' object cannot be interpreted as an integer"), + ], +) +async def test_a_failure_reads_a_num_retries_that_is_no_count( + make_core: MakeCore, host: ScriptedHost, value: object, message: str +) -> None: + core: Final = make_core(model_list=[deployment("g", "a"), deployment("g", "b")], num_retries=0) + host.answer("*", lambda **_: server_error()) + + error: Final = await attempt(core, num_retries=value) + + assert isinstance(error, TypeError) + assert str(error) == message + assert isinstance(error.__context__, litellm.InternalServerError) + assert len(host.deployments_called("acompletion")) == 1 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("value", [-1, float("nan"), False, 0.0]) +async def test_a_num_retries_below_one_leaves_the_failure_standing( + make_core: MakeCore, host: ScriptedHost, value: object +) -> None: + core: Final = make_core(model_list=[deployment("g", "a"), deployment("g", "b")], num_retries=3) + host.answer("*", lambda **_: server_error()) + + error: Final = await attempt(core, num_retries=value) + + assert isinstance(error, litellm.InternalServerError) + assert len(host.deployments_called("acompletion")) == 1 + + +@pytest.mark.asyncio +async def test_a_retry_policy_outranks_a_num_retries_that_is_no_count( + make_core: MakeCore, host: ScriptedHost +) -> None: + core: Final = make_core( + model_list=[deployment("g", "a"), deployment("g", "b")], + num_retries=0, + retry_policy=RetryPolicy(InternalServerErrorRetries=1), + retry_after=0, + ) + host.answer("*", server_error(), lambda **_: None) + + served: Final = await attempt(core, num_retries="2") + + assert not isinstance(served, Exception) + assert len(host.deployments_called("acompletion")) == 2 + + +@pytest.mark.asyncio +async def test_num_retries_true_counts_as_one_retry(make_core: MakeCore, host: ScriptedHost) -> None: + core: Final = make_core( + model_list=[deployment("g", "a"), deployment("g", "b")], num_retries=0, retry_after=0, disable_cooldowns=True + ) + host.answer("*", lambda **_: server_error()) + + error: Final = await attempt(core, num_retries=True) + + assert isinstance(error, litellm.InternalServerError) + assert len(host.deployments_called("acompletion")) == 2 + + +# --- N17: retire releases registrations, not the router ---------------------------------------------- + + +@pytest.mark.asyncio +async def test_a_retired_core_keeps_serving(make_core: MakeCore, host: ScriptedHost) -> None: + core: Final = _one_group(make_core) + core.retire() + + served: Final = await attempt(core) + + assert not isinstance(served, Exception) + assert host.deployments_called("acompletion") == ["d1"] + + +# --- N18: attributes the Python core keeps as plain attributes --------------------------------------- + + +@pytest.mark.asyncio +async def test_a_deleted_scheduler_is_gone_for_reads_and_for_prioritized_calls(make_core: MakeCore) -> None: + core: Final = _one_group(make_core) + del core.scheduler # pyright: ignore[reportAttributeAccessIssue] # a plain attribute of both cores + + with pytest.raises(AttributeError, match="scheduler"): + _ = core.scheduler # pyright: ignore[reportAttributeAccessIssue] # see above + served: Final = await attempt(core) + refused: Final = await attempt(core, priority=1) + + assert not isinstance(served, Exception) + assert isinstance(refused, AttributeError) + assert "scheduler" in str(refused) + + +@pytest.mark.parametrize("name", ["total_calls", "success_calls", "fail_calls"]) +def test_call_counters_are_replaced_and_deleted_like_attributes(make_core: MakeCore, name: str) -> None: + core: Final = _one_group(make_core) + replacement: Final[dict[str, int]] = {"openai/d1": 7} # mutable-ok: the value under test + + setattr(core, name, replacement) + assert getattr(core, name) is replacement + + delattr(core, name) + with pytest.raises(AttributeError, match=name): + getattr(core, name) + + +@pytest.mark.asyncio +async def test_a_success_counter_replaced_with_a_plain_dict_fails_the_attempt(make_core: MakeCore) -> None: + core: Final = _one_group(make_core) + core.success_calls = {} # pyright: ignore[reportAttributeAccessIssue] # see above + + error: Final = await attempt(core) + + assert isinstance(error, KeyError) + assert getattr(error, "failed_deployment_id", None) == "d1" + assert core.fail_calls == {"openai/d1": 1} # pyright: ignore[reportAttributeAccessIssue] # see above + + +def test_embedding_names_its_parameters(make_core: MakeCore) -> None: + signature: Final = inspect.signature(_one_group(make_core).embedding) # pyright: ignore[reportAttributeAccessIssue] # an endpoint + assert list(signature.parameters) == ["model", "input", "is_async", "kwargs"] + assert signature.parameters["is_async"].default is False + + +def test_embedding_takes_its_arguments_by_position(make_core: MakeCore, host: ScriptedHost) -> None: + core: Final = _one_group(make_core) + core.embedding("g", "hi") # pyright: ignore[reportAttributeAccessIssue] # an endpoint + assert host.deployments_called("embedding") == ["d1"] + + +# --- O2: deleted aliases and retry budget ------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_a_deleted_alias_map_fails_every_read_until_it_is_assigned_again(make_core: MakeCore) -> None: + core: Final = _one_group(make_core, model_group_alias={"al": "g"}) + del core.model_group_alias # pyright: ignore[reportAttributeAccessIssue] # the deprecated attribute + + assert "model_group_alias" not in core.settings().as_dict() + with pytest.raises(AttributeError, match="model_group_alias"): + core.registry().aliases() + with pytest.raises(AttributeError, match="model_group_alias"): + core.registry().names() + for model in ("al", "g"): + refused = await attempt(core, model=model) + assert isinstance(refused, AttributeError) + assert "model_group_alias" in str(refused) + with pytest.raises(AttributeError, match="model_group_alias"): + del core.model_group_alias # pyright: ignore[reportAttributeAccessIssue] # see above + + core.apply(ChangeSet(assign={"model_group_alias": {"x": "g"}})) + + assert not isinstance(await attempt(core, model="x"), Exception) + + +@pytest.mark.asyncio +async def test_a_deleted_num_retries_fails_calls_that_name_none(make_core: MakeCore, host: ScriptedHost) -> None: + core: Final = make_core(model_list=[deployment("g", "d1"), deployment("h", "d2")], fallbacks=[{"g": ["h"]}]) + del core.num_retries # pyright: ignore[reportAttributeAccessIssue] # a plain attribute of both cores + + assert "num_retries" not in core.settings().as_dict() + with pytest.raises(AttributeError, match="num_retries"): + _ = core.num_retries # pyright: ignore[reportAttributeAccessIssue] # see above + refused: Final = await attempt(core) + assert isinstance(refused, AttributeError) + assert "num_retries" in str(refused) + assert host.deployments_called("acompletion") == [] + + assert not isinstance(await attempt(core, num_retries=0), Exception) + with pytest.raises(AttributeError, match="num_retries"): + del core.num_retries # pyright: ignore[reportAttributeAccessIssue] # see above + + core.num_retries = 1 # pyright: ignore[reportAttributeAccessIssue] # see above + + assert core.num_retries == 1 # pyright: ignore[reportAttributeAccessIssue] # see above + assert not isinstance(await attempt(core), Exception) + + +# --- T5: a default deployment every call to fails ------------------------------------------------------ + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("model_info", "stamped"), [(None, None), ({"id": "dd"}, "dd")], ids=["no_model_info", "its_own_id"] +) +async def test_a_failing_default_deployment_stamps_only_its_own_id( + make_core: MakeCore, model_info: dict[str, object] | None, stamped: str | None +) -> None: + core: Final = _one_group(make_core) + assigned: Final = { + "litellm_params": {"model": "openai/gpt-4o", "api_key": "k"}, + **({"model_info": model_info} if model_info is not None else {}), + } + core.apply(ChangeSet(assign={"default_deployment": assigned})) + + error: Final = await attempt(core, model="unknown") + + assert isinstance(error, KeyError) + assert error.args == ("model_name",) + assert getattr(error, "failed_deployment_id", None) == stamped + + +# --- S10: the capability limit is called as given -------------------------------------------------------- + + +def test_a_capability_limit_that_cannot_be_called_fails_the_build(make_core: MakeCore) -> None: + with pytest.raises(TypeError, match="not callable"): + make_core( + model_list=[ + deployment("big", "big"), + deployment("small", "small"), + complexity_marker("smart", "m1", config=HEURISTIC_V2), + ], + auto_router_capability_limit=0, + ) + + +# --- ruled to differ: the Rust core only --------------------------------------------------------------- + + +@pytest.mark.parametrize( + ("params", "named"), + [ + ({"model_info": {"input_cost_per_token": "inf"}}, "model_info.input_cost_per_token"), + ({"model_info": {"output_cost_per_token": "nan"}}, "model_info.output_cost_per_token"), + ({"model_info": {"input_cost_per_token": "1e400"}}, "model_info.input_cost_per_token"), + ({"model_info": {"input_cost_per_token": float("inf")}}, "model_info.input_cost_per_token"), + ({"input_cost_per_token": "inf"}, "litellm_params.input_cost_per_token"), + ], + ids=["inf", "nan", "overflow", "float", "litellm_params"], +) +def test_the_rust_core_refuses_a_cost_that_is_not_finite( + host: ScriptedHost, params: dict[str, object], named: str +) -> None: + with pytest.raises(ValueError, match=f"{named} is .*Set a finite cost"): + CORE_BUILDERS["rust"](host, model_list=[deployment("g", "m", model="openai/gpt-4o", **params)]) + + +def test_the_rust_core_refuses_a_cost_that_is_not_finite_on_add(host: ScriptedHost) -> None: + core: Final = CORE_BUILDERS["rust"](host, model_list=[deployment("g", "m")]) + added: Final = Deployment( + model_name="g", + litellm_params=LiteLLM_Params(model="openai/gpt-4o", api_key="k"), + model_info={"id": "late", "input_cost_per_token": float("inf")}, + ) + with pytest.raises(ValueError, match="Set a finite cost"): + core.add_deployment(added) # pyright: ignore[reportAttributeAccessIssue] # the deprecated member + assert core.registry().names(kind="ids") == ["m"] + + +def test_the_rust_core_leaves_out_a_bad_cost_under_ignore_invalid_deployments(host: ScriptedHost) -> None: + core: Final = CORE_BUILDERS["rust"]( + host, + model_list=[deployment("g", "m", model_info={"input_cost_per_token": "inf"}), deployment("g", "n")], + ignore_invalid_deployments=True, + ) + assert core.registry().names(kind="ids") == ["n"] + + +@pytest.mark.parametrize("share", [10**20, 1e20], ids=["int", "float"]) +def test_the_rust_core_refuses_a_share_wider_than_it_can_hold(host: ScriptedHost, share: object) -> None: + listed: Final = deployment("g", "sh", model_info={"db_model": True, "ptu_shares": {"T1": share}}) + with pytest.raises(ValueError, match=r"model_info\.ptu_shares\.T1 is 100000000000000000000, outside"): + CORE_BUILDERS["rust"](host, model_list=[listed]) + + +def test_the_rust_core_refuses_a_share_wider_than_it_can_hold_on_add(host: ScriptedHost) -> None: + core: Final = CORE_BUILDERS["rust"](host, model_list=[deployment("g", "free")]) + added: Final = Deployment( + model_name="g", + litellm_params=LiteLLM_Params(model="openai/gpt-4o", api_key="k"), + model_info={"id": "late", "db_model": True, "ptu_shares": {"T1": 10**20}}, + ) + with pytest.raises(ValueError, match=r"model_info\.ptu_shares\.T1 is 100000000000000000000, outside"): + core.add_deployment(added) # pyright: ignore[reportAttributeAccessIssue] # the deprecated member + assert core.registry().names(kind="ids") == ["free"] + + +# --- S5: in-place edits never reach the Rust walk, or its notes ------------------------------------------- + + +@pytest.mark.asyncio +async def test_the_rust_error_note_reads_the_fallbacks_the_walk_used(host: ScriptedHost) -> None: + host.set_dynamic(expose_router_debug_in_errors=True) + core: Final = CORE_BUILDERS["rust"]( + host, + model_list=[deployment("g", "d1"), deployment("h", "h1")], + num_retries=0, + fallbacks=[{"x": ["h"]}], + ) + core.settings().fallbacks.append({"g": ["h"]}) # pyright: ignore[reportOptionalMemberAccess] # the live list + host.answer("*", lambda **_: server_error()) + + error: Final = await attempt(core) + + assert isinstance(error, litellm.InternalServerError) + assert "configured: h" not in error.message + assert host.deployments_called("acompletion") == ["d1"] diff --git a/tests/unit/test_router/core/test_parity_regress5_streams.py b/tests/unit/test_router/core/test_parity_regress5_streams.py new file mode 100644 index 00000000000..173cb2c5db8 --- /dev/null +++ b/tests/unit/test_router/core/test_parity_regress5_streams.py @@ -0,0 +1,297 @@ +""" +Parity of the Rust core with ``router_python_core.py`` on the stream causes of the regress5 pack: +request-level fallback lists on a ``/v1/messages`` stream, what a stream's failure records, a +malformed group retry policy, the error a failed recovery raises, and the hint on a stream's failure +event. +""" + +import asyncio +from collections.abc import Callable, Mapping +from typing import Final + +import pytest +from pydantic import ValidationError + +import litellm +from litellm.llms.anthropic.messages.adapters.streaming_iterator import AnthropicSSEStream, AnthropicStreamWrapper +from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper +from litellm.router_core import RouterCore +from litellm.types.router import RetryPolicy +from tests.unit.test_router.core.scripted_host import ( + ScriptedHost, + auth_error, + content_policy, + context_window, + deployment, + rate_limit, + server_error, +) +from tests.unit.test_router.core.streaming_helpers import ( + MESSAGE_START, + MESSAGE_STOP, + MESSAGES, + anthropic_stream, + chat_stream, + collect_bytes, + collect_texts, + collect_texts_sync, + error_frame, + install_orjson_stub_if_missing, + refusal_delta, + text_delta, +) + +install_orjson_stub_if_missing() + +MakeCore = Callable[..., RouterCore] +GROUPS_AB: Final = [deployment("a", "a"), deployment("b", "b")] +PREFILL_MODEL: Final = "anthropic/claude-sonnet-4-5" +FALLBACK_FRAMES: Final = [MESSAGE_START, text_delta("B"), MESSAGE_STOP] +REFUSAL_DETAILS: Final = {"type": "refusal", "category": "cyber"} +PER_USER_COPILOT: Final = { + "model": "github_copilot/gpt-4o", + "litellm_credential_name": "ghc", + "github_copilot_auth_type": "per_user_oauth", +} + + +async def open_messages(core: RouterCore, **request: object) -> object: # kwargs-ok: request kwargs + return await core.aanthropic_messages(model="a", messages=MESSAGES, max_tokens=5, stream=True, **request) # pyright: ignore[reportAttributeAccessIssue] # endpoint method + + +async def settle() -> None: + """Let the logging and recording tasks a failed stream scheduled run.""" + for _ in range(20): + await asyncio.sleep(0) + + +def cooled(core: RouterCore, *ids: str) -> Mapping[str, Mapping[str, object]]: + return dict(core.health().cooldowns.get_active_cooldowns(list(ids), None)) # pyright: ignore[reportReturnType] # a mapping of entries + + +class BridgedStream(AnthropicSSEStream): + """The byte stream the completion bridge returns for ``/v1/messages``: a chat stream's failure is + logged by that stream and reaches the caller as an ``event: error`` frame.""" + + def __init__(self, stream: CustomStreamWrapper) -> None: + super().__init__( + AnthropicStreamWrapper(completion_stream=stream, model="d1", litellm_logging_obj=stream.logging_obj) + ) + self.logging_obj: Final = stream.logging_obj + + +def bridged(error: Exception) -> Callable[..., BridgedStream]: + chat: Final = chat_stream(error=error, fail_after=0, logged=True) + + def answer(**request: object) -> BridgedStream: # kwargs-ok: the host passes the request kwargs + return BridgedStream(chat(**request)) + + return answer + + +# --- T1: the request's fallback lists on a /v1/messages stream ------------------------------------- + + +@pytest.mark.parametrize("fallbacks", [["b"], [{"a": ["b"]}], [{"*": ["b"]}]], ids=["names", "keyed", "star"]) +async def test_request_fallbacks_rescue_a_messages_stream_before_content( + make_core: MakeCore, host: ScriptedHost, fallbacks: list[object] +) -> None: + host.answer("openai/a", anthropic_stream(MESSAGE_START, error_frame())) + host.answer("openai/b", anthropic_stream(*FALLBACK_FRAMES)) + core: Final = make_core(model_list=GROUPS_AB, num_retries=0) + assert await collect_bytes(await open_messages(core, fallbacks=fallbacks)) == FALLBACK_FRAMES + assert host.deployments_called() == ["a", "b"] + + +async def test_request_content_policy_fallbacks_rescue_a_refusal_before_content( + make_core: MakeCore, host: ScriptedHost +) -> None: + host.answer("openai/a", anthropic_stream(MESSAGE_START, refusal_delta(REFUSAL_DETAILS))) + host.answer("openai/b", anthropic_stream(*FALLBACK_FRAMES)) + core: Final = make_core(model_list=GROUPS_AB, num_retries=0) + stream: Final = await open_messages(core, content_policy_fallbacks=[{"a": ["b"]}]) + assert await collect_bytes(stream) == FALLBACK_FRAMES + assert host.deployments_called() == ["a", "b"] + + +async def test_a_request_without_fallbacks_replaces_the_router_list_on_a_messages_stream( + make_core: MakeCore, host: ScriptedHost +) -> None: + host.answer("openai/a", anthropic_stream(MESSAGE_START, error_frame())) + core: Final = make_core(model_list=GROUPS_AB, num_retries=0, fallbacks=[{"a": ["b"]}]) + assert await collect_bytes(await open_messages(core, fallbacks=None)) == [MESSAGE_START, error_frame()] + assert host.deployments_called() == ["a"] + + +# --- T2: what a bridged stream's failure records -------------------------------------------------- + + +async def test_a_bridged_401_before_content_cools_the_deployment_down(make_core: MakeCore, host: ScriptedHost) -> None: + host.answer("openai/a", bridged(auth_error())) + host.answer("openai/b", anthropic_stream(*FALLBACK_FRAMES)) + core: Final = make_core(model_list=GROUPS_AB, num_retries=0, fallbacks=[{"a": ["b"]}]) + assert await collect_bytes(await open_messages(core)) == FALLBACK_FRAMES + await settle() + assert sorted(cooled(core, "a", "b")) == ["a"] + assert str(cooled(core, "a")["a"]["status_code"]) == "401" + + +async def test_a_bridged_rate_limit_cools_with_the_provider_error_not_the_frame( + make_core: MakeCore, host: ScriptedHost +) -> None: + host.answer("openai/a", bridged(rate_limit(message="provider text"))) + host.answer("openai/b", anthropic_stream(*FALLBACK_FRAMES)) + core: Final = make_core(model_list=GROUPS_AB, num_retries=0, allowed_fails=0, fallbacks=[{"a": ["b"]}]) + assert await collect_bytes(await open_messages(core)) == FALLBACK_FRAMES + await settle() + received: Final = str(cooled(core, "a")["a"]["exception_received"]) + assert "provider text" in received + assert "AnthropicException" not in received + + +async def test_a_caller_scoped_401_on_a_bridged_stream_never_cools_the_deployment( + make_core: MakeCore, host: ScriptedHost +) -> None: + host.answer( + "github_copilot/gpt-4o", + bridged( + litellm.CallerCredentialAuthenticationError( + message="not connected", llm_provider="github_copilot", model="gpt-4o" + ) + ), + ) + host.answer("openai/b", anthropic_stream(*FALLBACK_FRAMES)) + core: Final = make_core( + model_list=[deployment("a", "a", **PER_USER_COPILOT), deployment("b", "b")], + num_retries=0, + allowed_fails=0, + fallbacks=[{"a": ["b"]}], + ) + assert await collect_bytes(await open_messages(core)) == FALLBACK_FRAMES + await settle() + assert cooled(core, "a", "b") == {} + + +# --- T3: a malformed group retry policy on a messages stream --------------------------------------- + + +async def test_a_malformed_group_policy_raises_on_a_messages_stream_that_never_opened( + make_core: MakeCore, host: ScriptedHost +) -> None: + host.answer("*", rate_limit()) + core: Final = make_core( + model_list=[deployment("a", "a1"), deployment("a", "a2")], + num_retries=1, + model_group_retry_policy={"a": {"RateLimitErrorRetries": "many"}}, + ) + with pytest.raises(ValidationError): + await open_messages(core) + assert len(host.deployments_called()) == 1 + + +@pytest.mark.parametrize("retry_policy", [None, RetryPolicy(InternalServerErrorRetries=0)], ids=["alone", "router"]) +async def test_a_malformed_group_policy_governs_nothing_on_a_stream_that_opened( + make_core: MakeCore, host: ScriptedHost, retry_policy: RetryPolicy | None +) -> None: + host.answer("openai/a", anthropic_stream(MESSAGE_START, error=server_error()), anthropic_stream(*FALLBACK_FRAMES)) + core: Final = make_core( + model_list=[deployment("a", "a")], + num_retries=1, + retry_policy=retry_policy, + model_group_retry_policy={"a": {"RateLimitErrorRetries": "many"}}, + ) + assert await collect_bytes(await open_messages(core)) == FALLBACK_FRAMES + assert host.deployments_called() == ["a", "a"] + + +# --- S8: the error a failed recovery raises -------------------------------------------------------- + + +async def test_a_continuation_whose_fallback_stream_fails_raises_the_fallback_error( + make_core: MakeCore, host: ScriptedHost +) -> None: + host.answer("openai/a", chat_stream("he", error=rate_limit(message="primary"))) + host.answer(PREFILL_MODEL, chat_stream(" wor", "ld", error=server_error(), fail_after=1)) + core: Final = make_core( + model_list=[deployment("a", "a"), deployment("b", "b", model=PREFILL_MODEL)], + num_retries=0, + fallbacks=[{"a": ["b"]}], + enable_mid_stream_fallback_continuation=True, + ) + stream: Final = await core.acompletion(model="a", messages=MESSAGES, stream=True) # pyright: ignore[reportAttributeAccessIssue] # endpoint method + with pytest.raises(litellm.InternalServerError, match="boom"): + await collect_texts(stream) + assert host.deployments_called() == ["a", "b"] + + +async def test_a_fallback_stream_that_fails_raises_its_own_error(make_core: MakeCore, host: ScriptedHost) -> None: + host.answer("openai/a", chat_stream(error=rate_limit(message="primary"), leading_role_chunk=True)) + host.answer("openai/b", chat_stream("B", error=server_error())) + core: Final = make_core(model_list=GROUPS_AB, num_retries=0, fallbacks=[{"a": ["b"]}]) + stream: Final = await core.acompletion(model="a", messages=MESSAGES, stream=True) # pyright: ignore[reportAttributeAccessIssue] # endpoint method + with pytest.raises(litellm.InternalServerError, match="boom"): + await collect_texts(stream) + + +def test_a_sync_fallback_stream_that_fails_raises_its_own_error(make_core: MakeCore, host: ScriptedHost) -> None: + host.answer("openai/a", chat_stream(error=rate_limit(message="primary"), leading_role_chunk=True, sync=True)) + host.answer("openai/b", chat_stream("B", error=server_error(), sync=True)) + core: Final = make_core(model_list=GROUPS_AB, num_retries=0, fallbacks=[{"a": ["b"]}]) + stream: Final = core.completion(model="a", messages=MESSAGES, stream=True) # pyright: ignore[reportAttributeAccessIssue] # endpoint method + with pytest.raises(litellm.InternalServerError, match="boom"): + collect_texts_sync(stream) + + +async def test_a_messages_fallback_stream_that_fails_raises_its_own_error( + make_core: MakeCore, host: ScriptedHost +) -> None: + host.answer("openai/a", anthropic_stream(MESSAGE_START, error_frame())) + host.answer("openai/b", anthropic_stream(MESSAGE_START, text_delta("B"), error=server_error())) + core: Final = make_core(model_list=GROUPS_AB, num_retries=0, fallbacks=[{"a": ["b"]}]) + stream: Final = await open_messages(core) + with pytest.raises(litellm.InternalServerError, match="boom"): + await collect_bytes(stream) + + +async def test_an_order_chain_that_serves_nothing_raises_its_last_hop_error_from_a_stream( + make_core: MakeCore, host: ScriptedHost +) -> None: + """The order path of the fallback walk raises its last hop's error, not the failure it was handed.""" + host.answer("openai/a1", chat_stream(error=rate_limit(message="primary"), leading_role_chunk=True)) + host.answer("openai/a2", server_error()) + core: Final = make_core(model_list=[deployment("a", "a1", order=1), deployment("a", "a2", order=2)], num_retries=0) + stream: Final = await core.acompletion(model="a", messages=MESSAGES, stream=True) # pyright: ignore[reportAttributeAccessIssue] # endpoint method + with pytest.raises(litellm.InternalServerError, match="boom"): + await collect_texts(stream) + assert host.deployments_called() == ["a1", "a2"] + + +# --- V3: the hint on a stream's failure event ------------------------------------------------------ + + +@pytest.mark.parametrize( + ("error", "raised", "hint"), + [ + (context_window, litellm.ContextWindowExceededError, "Set 'context_window_fallback'"), + (content_policy, litellm.ContentPolicyViolationError, "content_policy_fallback"), + ], + ids=["context_window", "content_policy"], +) +async def test_a_failure_event_of_a_stream_failing_before_its_first_chunk_carries_the_hint( + make_core: MakeCore, host: ScriptedHost, error: Callable[[], Exception], raised: type[Exception], hint: str +) -> None: + host.set_dynamic(expose_router_debug_in_errors=True) + host.answer("openai/a", chat_stream(error=error(), fail_after=0, logged=True)) + host.answer("openai/b", server_error()) + core: Final = make_core(model_list=GROUPS_AB, num_retries=0, fallbacks=[{"a": ["b"]}]) + logged: Final[list[str]] = [] + + def record(kwargs: Mapping[str, object], *_rest: object) -> None: + logged.append(str(getattr(kwargs.get("exception"), "message", ""))) + + host.callback_list("failure_callback").append(record) + with pytest.raises(raised): + await core.acompletion(model="a", messages=MESSAGES, stream=True) # pyright: ignore[reportAttributeAccessIssue] # endpoint method + await settle() + # The fallback's own failure may be logged first. + assert any(hint in message for message in logged) diff --git a/tests/unit/test_router/core/test_parity_response_headers.py b/tests/unit/test_router/core/test_parity_response_headers.py new file mode 100644 index 00000000000..22794eccbef --- /dev/null +++ b/tests/unit/test_router/core/test_parity_response_headers.py @@ -0,0 +1,205 @@ +""" +Parity of the Rust core with ``router_python_core.py`` on the router's response headers, what a sync +pick filters, and the names the no-deployments errors carry. +""" + +import asyncio +import warnings +from collections.abc import Callable, Mapping +from typing import Final + +import pytest +from pydantic import ValidationError + +import litellm +from litellm.exceptions import BadRequestError, EndpointNotSupportedByModelError +from litellm.integrations.custom_logger import CustomLogger +from litellm.types.router import RouterRateLimitError +from tests.unit.test_router.core.fallbacks_helpers import MESSAGES, MakeCore, complete, complete_sync, headers +from tests.unit.test_router.core.scripted_host import CORE_BUILDERS, ScriptedHost, completion_response, deployment, server_error + +GROUP_LIMITS: Final = {"rpm": 5, "tpm": 1000} +TEAM_OWNED: Final = {"team_id": "T1"} +UNREADABLE_LIMIT: Final = {"rpm": "abc"} +RETRY_COUNT: Final = "x-litellm-attempted-retries" +FALLBACK_COUNT: Final = "x-litellm-attempted-fallbacks" +GUARDRAILS: Final = [{"guardrail_name": "gr", "litellm_params": {"guardrail": "aporia", "mode": "pre_call"}}] +BLOCKED_RESPONSE: Final = { + "id": "chatcmpl-blocked", + "object": "chat.completion", + "model": "gpt-4o", + "created": 1, + "choices": [{"index": 0, "finish_reason": "content_filter", "message": {"role": "assistant", "content": "no"}}], +} +# The dumps of a deployment `_create_deployment` makes, each of which warns about a mocked ModelResponse. +PYTHON_DUMPS_OF_A_DEPLOYMENT: Final = 4 + + +# --- x-ratelimit-* and the counts ------------------------------------------------------------------ + + +def test_a_sync_call_reads_the_remaining_usage_of_its_group(make_core: MakeCore) -> None: + core: Final = make_core(model_list=[deployment("g", "d1", **GROUP_LIMITS)], num_retries=0) + stamped: Final = headers(complete_sync(core, "g")) + assert stamped["x-ratelimit-limit-requests"] == 5 + assert stamped["x-ratelimit-remaining-requests"] == 4 + assert stamped["x-ratelimit-limit-tokens"] == 1000 + assert stamped[RETRY_COUNT] == 0 + assert stamped[FALLBACK_COUNT] == 0 + + +async def test_an_unreadable_group_limit_leaves_the_counts_off_the_response( + make_core: MakeCore, host: ScriptedHost +) -> None: + answered: Final = completion_response() + host.answer("openai/d1", answered) + core: Final = make_core(model_list=[deployment("g", "d1", model_info=UNREADABLE_LIMIT)], num_retries=0) + with pytest.raises(ValidationError): + await complete(core, "g") + stamped: Final = headers(answered) + assert stamped["x-litellm-model-group"] == "g" + assert RETRY_COUNT not in stamped + assert FALLBACK_COUNT not in stamped + + +def test_a_sync_call_reads_the_group_limit_too(make_core: MakeCore) -> None: + core: Final = make_core(model_list=[deployment("g", "d1", model_info=UNREADABLE_LIMIT)], num_retries=0) + with pytest.raises(ValidationError): + complete_sync(core, "g") + + +def test_a_sync_call_opens_at_most_one_event_loop(make_core: MakeCore, monkeypatch: pytest.MonkeyPatch) -> None: + """Each loop of its own hands the GIL to the logging thread before the headers are on the response.""" + opened: Final[list[object]] = [] # mutable-ok: filled by the patched factory + real: Final[Callable[[], asyncio.AbstractEventLoop]] = asyncio.new_event_loop + + def counted() -> asyncio.AbstractEventLoop: + loop: Final = real() + opened.append(loop) + return loop + + core: Final = make_core(model_list=[deployment("g", "d1", **GROUP_LIMITS)], num_retries=0) + monkeypatch.setattr(asyncio, "new_event_loop", counted) + _ = complete_sync(core, "g") + assert len(opened) <= 1 + + +class _FallbackCountAtLog(CustomLogger): + def __init__(self, response: object) -> None: + super().__init__() + self._response: Final = response + self.seen: Final[list[object]] = [] # mutable-ok: filled as the events log + + async def log_success_fallback_event( + self, original_model_group: str, kwargs: dict, original_exception: Exception + ) -> None: + self.seen.append(headers(self._response).get(FALLBACK_COUNT)) + + +async def test_a_fallback_counts_itself_before_its_success_is_logged( + make_core: MakeCore, host: ScriptedHost, monkeypatch: pytest.MonkeyPatch +) -> None: + """``run_async_fallback`` stamps the response, then awaits ``log_success_fallback_event``.""" + answered: Final = completion_response() + logger: Final = _FallbackCountAtLog(answered) + monkeypatch.setattr(litellm, "callbacks", [logger]) + host.answer("openai/a1", server_error()) + host.answer("openai/b1", answered) + core: Final = make_core( + model_list=[deployment("a", "a1"), deployment("b", "b1")], num_retries=0, fallbacks=[{"a": ["b"]}] + ) + _ = await complete(core, "a") + assert logger.seen == [1] + + +# --- guardrails ------------------------------------------------------------------------------------ + + +async def test_a_guardrail_result_carries_the_router_headers(make_core: MakeCore) -> None: + answered: Final = completion_response() + + async def guard(**_kwargs: object) -> object: # kwargs-ok: the guardrail's keyword payload + return answered + + core: Final = make_core(model_list=[deployment("g", "d1")], guardrail_list=GUARDRAILS) + result: Final = await core.aguardrail(guardrail_name="gr", original_function=guard) # pyright: ignore[reportAttributeAccessIssue, reportUnknownMemberType] # an endpoint + assert headers(result) == {"x-litellm-model-group": "gr", RETRY_COUNT: 0, FALLBACK_COUNT: 0} + + +# --- the team filter ------------------------------------------------------------------------------- + + +# S2 and Y4 are ruled to differ: the Rust core filters every pick, so these run on it only. +def test_the_rust_core_filters_a_team_owned_deployment_on_a_sync_pick(host: ScriptedHost) -> None: + core: Final = CORE_BUILDERS["rust"](host, model_list=[deployment("own", "o1", model_info=TEAM_OWNED)], num_retries=0) + with pytest.raises(RouterRateLimitError): + _ = complete_sync(core, "own") + assert host.deployments_called("completion") == [] + + +async def test_an_async_pick_runs_the_team_filter(make_core: MakeCore) -> None: + core: Final = make_core(model_list=[deployment("own", "o1", model_info=TEAM_OWNED)], num_retries=0) + with pytest.raises(RouterRateLimitError): + await complete(core, "own") + + +@pytest.mark.parametrize( + "request_fields", + [ + pytest.param({}, id="no-team"), + pytest.param({"metadata": {"user_api_key_team_id": "T2"}}, id="another-team"), + ], +) +async def test_the_rust_core_refuses_a_team_owned_deployment_named_by_its_id( + host: ScriptedHost, request_fields: Mapping[str, object] +) -> None: + core: Final = CORE_BUILDERS["rust"](host, model_list=[deployment("own", "o1", model_info=TEAM_OWNED)], num_retries=0) + with pytest.raises(RouterRateLimitError): + _ = await complete(core, "o1", **request_fields) + assert host.deployments_called("acompletion") == [] + + +# --- what the errors name -------------------------------------------------------------------------- + + +async def test_a_wildcard_refused_for_a_team_names_the_requested_model(make_core: MakeCore) -> None: + core: Final = make_core( + model_list=[deployment("anthropic/*", "w1", "anthropic/*", model_info=TEAM_OWNED)], num_retries=0 + ) + with pytest.raises(RouterRateLimitError) as refused: + await complete(core, "anthropic/claude-x", metadata={"user_api_key_team_id": "T2"}) + assert refused.value.model == "anthropic/claude-x" + + +def test_a_pass_through_refusal_names_the_group_an_alias_resolved_to(make_core: MakeCore) -> None: + core: Final = make_core( + model_list=[deployment("target", "t1")], num_retries=0, model_group_alias={"alias": "target"} + ) + with pytest.raises(BadRequestError) as refused: + core.get_available_deployment_for_pass_through("alias", []) # pyright: ignore[reportAttributeAccessIssue, reportUnknownMemberType] # the legacy member + assert refused.value.model == "target" + + +def test_a_sync_completion_is_refused_by_a_deployment_that_serves_no_chat(make_core: MakeCore) -> None: + core: Final = make_core( + model_list=[ + deployment( + "g", "e1", model_info={"enforce_supported_endpoints": True, "supported_endpoints": ["/v1/embeddings"]} + ) + ], + num_retries=0, + ) + with pytest.raises(EndpointNotSupportedByModelError): + core.completion(model="g", messages=MESSAGES) # pyright: ignore[reportAttributeAccessIssue, reportUnknownMemberType] # an endpoint + + +# --- construction ---------------------------------------------------------------------------------- + + +def test_a_mocked_model_response_warns_once_per_dump_python_makes(make_core: MakeCore) -> None: + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter("always") + _ = make_core(model_list=[deployment("g", "d1", mock_response=BLOCKED_RESPONSE)], num_retries=0) + serializer: Final = [w for w in caught if "Pydantic serializer warnings" in str(w.message)] + assert len(serializer) == PYTHON_DUMPS_OF_A_DEPLOYMENT + diff --git a/tests/unit/test_router/core/test_parity_rulings_2026_11.py b/tests/unit/test_router/core/test_parity_rulings_2026_11.py new file mode 100644 index 00000000000..dd78224774c --- /dev/null +++ b/tests/unit/test_router/core/test_parity_rulings_2026_11.py @@ -0,0 +1,286 @@ +""" +Parity of the Rust core with ``router_python_core.py`` on the router differences ruled on +2026-11-05, where both cores now take the side that was ruled correct: team ownership on the sync +and by-id paths, settings lists that change only through ``apply()``, one unreadable health record, +a failed ``set_model_list`` that keeps the old list, odd ``cooldown_time`` and status values, +numbers a deployment cannot be served with, an empty fallback entry, and a ``/v1/messages`` +stream that raises before its first frame. +""" + +import asyncio +import datetime +import time +from collections.abc import Awaitable, Callable, Mapping +from typing import Final, cast + +import pytest + +import litellm +from litellm.router_core import RouterCore +from litellm.utils import get_deployment_order +from tests.unit.test_router.core.scripted_host import ScriptedHost, deployment, rate_limit +from tests.unit.test_router.core.streaming_helpers import ( + MESSAGE_START, + MESSAGE_STOP, + MESSAGES, + anthropic_stream, + collect_bytes, + install_orjson_stub_if_missing, + text_delta, +) + +install_orjson_stub_if_missing() + +MakeCore = Callable[..., RouterCore] + + +def _cooled(core: RouterCore, *ids: str) -> dict[str, Mapping[str, object]]: # mutable-ok: a fresh dict + return dict(core.health().cooldowns.get_active_cooldowns(list(ids), None)) + + +def _sync(core: object, name: str) -> Callable[..., object]: + """A member both cores serve outside the ``RouterCore`` protocol (an endpoint or a deprecated name).""" + return cast("Callable[..., object]", getattr(core, name)) # cast-ok: the member is a plain function + + +def _async(core: object, name: str) -> Callable[..., Awaitable[object]]: + """``_sync`` for a coroutine member.""" + return cast("Callable[..., Awaitable[object]]", getattr(core, name)) # cast-ok: the member is a coroutine function + + +async def _on_failure(core: RouterCore, logged: Mapping[str, object]) -> None: + """``deployment_callback_on_failure``, awaited when the core answers a coroutine.""" + now: Final = datetime.datetime(2026, 1, 1) + outcome: Final = _sync(core, "deployment_callback_on_failure")(logged, None, now, now) + if asyncio.iscoroutine(outcome): + _ = cast("object", await outcome) # cast-ok: the callback's verdict, unused + await _settle() + + +def _ids(core: RouterCore) -> list[str]: # mutable-ok: a fresh list + return [str(entry["model_info"]["id"]) for entry in core.registry().all()] # pyright: ignore[reportAny] # model_list rows + + +async def _settle() -> None: + for _ in range(20): + await asyncio.sleep(0) + + +# --- team ownership -------------------------------------------------------------------------------- + +_TEAM_OWNED: Final = [deployment("own", "o1", model_info={"team_id": "T1"})] + + +@pytest.mark.parametrize("metadata", [{}, {"user_api_key_team_id": "T2"}], ids=["no-team", "another-team"]) +def test_a_sync_call_is_refused_a_deployment_another_team_owns( + make_core: MakeCore, host: ScriptedHost, metadata: dict[str, str] +) -> None: + core: Final = make_core(model_list=list(_TEAM_OWNED), num_retries=0) + + with pytest.raises(Exception, match=r"(?i)no (healthy )?deployments"): + _sync(core, "completion")(model="own", messages=MESSAGES, metadata=dict(metadata)) + + assert host.deployments_called("completion") == [] + + +def test_a_sync_pick_skips_a_deployment_another_team_owns(make_core: MakeCore) -> None: + core: Final = make_core(model_list=list(_TEAM_OWNED), num_retries=0) + + with pytest.raises(Exception, match=r"(?i)no (healthy )?deployments"): + _sync(core, "get_available_deployment")(model="own", messages=MESSAGES, request_kwargs={}) + + +@pytest.mark.asyncio +async def test_a_pick_by_id_is_refused_another_teams_deployment(make_core: MakeCore, host: ScriptedHost) -> None: + core: Final = make_core(model_list=[deployment("g", "t1-dep", model_info={"team_id": "T1"})], num_retries=0) + + with pytest.raises(Exception, match="g"): + await _async(core, "acompletion")(model="t1-dep", messages=MESSAGES, metadata={"user_api_key_team_id": "T2"}) + + assert host.deployments_called("acompletion") == [] + + +# --- settings change only through apply() --------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_an_entry_appended_to_the_fallbacks_list_is_ignored(make_core: MakeCore, host: ScriptedHost) -> None: + listed: Final[list[object]] = [] # mutable-ok: the caller's list, edited after construction + core: Final = make_core(model_list=[deployment("g", "d1"), deployment("fb", "f1")], num_retries=0, fallbacks=listed) + host.answer("openai/d1", rate_limit()) + listed.append({"g": ["fb"]}) + + with pytest.raises(litellm.RateLimitError): + await _async(core, "acompletion")(model="g", messages=MESSAGES) + + assert host.deployments_called("acompletion") == ["d1"] + + +# --- one unreadable health record ------------------------------------------------------------------ + + +@pytest.mark.asyncio +@pytest.mark.parametrize("timestamp", ["soon", True, None], ids=["string", "bool", "none"]) +async def test_an_unreadable_health_timestamp_drops_only_that_record( + make_core: MakeCore, host: ScriptedHost, timestamp: object +) -> None: + core: Final = make_core( + model_list=[deployment("g", "d1", weight=1), deployment("g", "d2", weight=0), deployment("x", "x1")], + num_retries=0, + enable_health_check_routing=True, + ) + states: Final = { + "d1": {"is_healthy": False, "timestamp": time.time()}, + "x1": {"is_healthy": False, "timestamp": timestamp}, + } + health_checks: Final = core.health().health_checks + # As another writer left it in the shared cache, then merged by this router's own writer. + _sync(health_checks.cache, "set_cache")(key=health_checks.CACHE_KEY, value=states) + await _async(core, "acompletion")(model="g", messages=MESSAGES) + health_checks.set_deployment_health_states(states) # pyright: ignore[reportArgumentType] # a record with a bad timestamp, on purpose + await _async(core, "acompletion")(model="g", messages=MESSAGES) + + assert host.deployments_called("acompletion") == ["d2", "d2"] + + +# --- a failed set_model_list ------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_a_failed_set_model_list_keeps_the_old_list(make_core: MakeCore, host: ScriptedHost) -> None: + core: Final = make_core(model_list=[deployment("a", "d-a"), deployment("b", "d-b")], num_retries=0) + unservable: Final = deployment("e", "d-e", input_cost_per_token=float("inf")) + + with pytest.raises(ValueError, match="cannot be served"): + _sync(core, "set_model_list")([deployment("c", "d-c"), unservable]) + + assert sorted(_ids(core)) == ["d-a", "d-b"] + await _async(core, "acompletion")(model="a", messages=MESSAGES) + assert host.deployments_called("acompletion") == ["d-a"] + + +# --- odd cooldown_time and status values ------------------------------------------------------------ + + +@pytest.mark.parametrize("value", [True, False], ids=["true", "false"]) +def test_a_bool_cooldown_time_is_refused_at_construction(make_core: MakeCore, value: bool) -> None: + with pytest.raises(TypeError, match="cooldown_time"): + make_core(model_list=[deployment("g", "d1")], cooldown_time=value) + + +# A group of one is never cooled down, so these groups hold two. +_TWO: Final = [deployment("g", "d1"), deployment("g", "d2")] + + +def _failure( + deployment_id: str, exception: Exception, **litellm_params: object +) -> dict[str, object]: # mutable-ok: logging kwargs + return { + "litellm_params": {"metadata": {"model_group": "g"}, "model_info": {"id": deployment_id}, **litellm_params}, + "call_type": "acompletion", + "exception": exception, + } + + +@pytest.mark.asyncio +async def test_a_bool_cooldown_time_on_the_call_cools_for_the_router_default(make_core: MakeCore) -> None: + core: Final = make_core(model_list=list(_TWO), num_retries=0, allowed_fails=0, cooldown_time=60) + + await _on_failure(core, _failure("d1", rate_limit(), cooldown_time=True)) + + assert _cooled(core, "d1")["d1"]["cooldown_time"] == 60 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("status", [70000, True, -1], ids=["too-wide", "bool", "negative"]) +async def test_a_status_outside_http_is_recorded_as_500(make_core: MakeCore, status: object) -> None: + core: Final = make_core(model_list=list(_TWO), num_retries=0, allowed_fails=0) + error: Final = rate_limit() + error.status_code = status # pyright: ignore[reportAttributeAccessIssue] # a provider error with an odd status + + await _on_failure(core, _failure("d1", error)) + + assert str(_cooled(core, "d1")["d1"]["status_code"]) == "500" + + +# --- numbers a deployment cannot be served with ----------------------------------------------------- + + +@pytest.mark.parametrize( + ("field", "value"), + [ + ("input_cost_per_token", float("inf")), + ("output_cost_per_token", float("nan")), + ("input_cost_per_token", "1e400"), + ], + ids=["inf", "nan", "overflowing-string"], +) +def test_a_non_finite_cost_is_refused_by_name(make_core: MakeCore, field: str, value: object) -> None: + with pytest.raises(ValueError, match=f"Deployment 'g' cannot be served: .*{field}.*finite costs"): + make_core(model_list=[deployment("g", "d1", model_info={field: value})]) + + +def test_a_ptu_share_wider_than_64_bits_is_refused_at_build_and_on_add(make_core: MakeCore) -> None: + wide: Final = deployment("g", "sh", model_info={"db_model": True, "ptu_shares": {"T1": 10**20}}) + with pytest.raises(ValueError, match=r"model_info\.ptu_shares\.T1 is 100000000000000000000, outside"): + make_core(model_list=[wide]) + + core: Final = make_core(model_list=[deployment("g", "free")]) + with pytest.raises(ValueError, match=r"ptu_shares\.T1 .* outside"): + _sync(core, "add_deployment")( + litellm.Deployment.model_validate( + deployment("g", "late", model_info={"db_model": True, "ptu_shares": {"T1": 10**20}}) + ) + ) + assert _ids(core) == ["free"] + + +# --- an empty fallback entry -------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_an_empty_fallback_entry_is_skipped(make_core: MakeCore, host: ScriptedHost) -> None: + core: Final = make_core(model_list=[deployment("g", "d1"), deployment("fb", "f1")], num_retries=0) + host.answer("openai/d1", rate_limit()) + + await _async(core, "acompletion")(model="g", messages=MESSAGES, fallbacks=[{}, {"g": ["fb"]}]) + + assert host.deployments_called("acompletion") == ["d1", "f1"] + + +# --- a deployment's order -------------------------------------------------------------------------- + + +@pytest.mark.parametrize( + ("order", "expected"), + [(1, 1), ("2", 2), (" 3", None), (2.0, 2), (2.5, None), (True, 1), ("x", None), (None, None)], + ids=["int", "digits", "padded", "whole-float", "fraction", "bool", "word", "unset"], +) +def test_a_deployment_order_reads_as_a_whole_number(order: object, expected: int | None) -> None: + assert get_deployment_order({"litellm_params": {"order": order}}) == expected + + +# --- a /v1/messages stream that raises before its first frame ---------------------------------------- + + +@pytest.mark.asyncio +async def test_a_messages_stream_raising_before_its_first_frame_cools_and_retries( + make_core: MakeCore, host: ScriptedHost +) -> None: + frames: Final = [MESSAGE_START, text_delta("B"), MESSAGE_STOP] + host.answer("openai/d1", anthropic_stream(error=rate_limit(message="provider text"))) + host.answer("openai/d2", anthropic_stream(*frames)) + core: Final = make_core( + model_list=[deployment("g", "d1", weight=1), deployment("g", "d2", weight=0)], + num_retries=1, + retry_after=0, + allowed_fails=0, + ) + + stream: Final = await _async(core, "aanthropic_messages")(model="g", messages=MESSAGES, max_tokens=5, stream=True) + assert await collect_bytes(stream) == frames + await _settle() + + assert host.deployments_called("anthropic_messages") == ["d1", "d2"] + assert sorted(_cooled(core, "d1", "d2")) == ["d1"] + assert "provider text" in str(_cooled(core, "d1")["d1"]["exception_received"]) diff --git a/tests/unit/test_router/core/test_registry_changesets.py b/tests/unit/test_router/core/test_registry_changesets.py index d4aa4d14a75..c889255a5a6 100644 --- a/tests/unit/test_router/core/test_registry_changesets.py +++ b/tests/unit/test_router/core/test_registry_changesets.py @@ -636,7 +636,8 @@ def _registered(host: ScriptedHost) -> Mapping[str, int]: def test_retire_releases_every_callback_the_core_registered_on_its_host( make_core: MakeCore, host: ScriptedHost ) -> None: - core = make_core(model_list=_model_list()) + # least-busy registers its selector on the host on both cores; the Python core adds its outcome callbacks. + core = make_core(model_list=_model_list(), routing_strategy="least-busy") assert _registered(host) != {} core.retire() @@ -645,8 +646,8 @@ def test_retire_releases_every_callback_the_core_registered_on_its_host( def test_retire_leaves_another_cores_callbacks_registered(make_core: MakeCore, host: ScriptedHost) -> None: - kept = make_core(model_list=_model_list()) - retired = make_core(model_list=[deployment("other", "o")]) + kept = make_core(model_list=_model_list(), routing_strategy="least-busy") + retired = make_core(model_list=[deployment("other", "o")], routing_strategy="least-busy") registered_by_both = _registered(host) retired.retire() diff --git a/tests/unit/test_router/core/test_registry_characterization_construction.py b/tests/unit/test_router/core/test_registry_characterization_construction.py new file mode 100644 index 00000000000..7b80f6c3b71 --- /dev/null +++ b/tests/unit/test_router/core/test_registry_characterization_construction.py @@ -0,0 +1,379 @@ +""" +How the core builds its registry from ``model_list``: deployment ids, provider checks, environment +gating, ``os.environ/`` references, per-organization fan-out, wildcard and team registration, and +the pricing each deployment writes into ``litellm.model_cost``. +""" + +import pytest + +import litellm +from litellm.router_api import deployment_id +from tests.unit.test_router.core.registry_helpers import MakeCore, ids, team_deployment, unique_id +from tests.unit.test_router.core.scripted_host import ScriptedHost, deployment + + +def _bare( + model_name: str, model: str, **params: object +) -> dict[str, object]: # mutable-ok: a model_list entry # kwargs-ok: litellm_params + """A model_list entry with no ``model_info`` at all.""" + return {"model_name": model_name, "litellm_params": {"model": model, "api_key": "sk-test", **params}} + + +# -- deployment ids ----------------------------------------------------------------------------- + + +def test_a_deployment_without_an_id_gets_the_stable_hash_of_its_group_and_params(make_core: MakeCore) -> None: + core = make_core(model_list=[_bare("gpt", "openai/gpt-4o")]) + + expected = deployment_id("gpt", {"model": "openai/gpt-4o", "api_key": "sk-test"}) + assert core.registry().names(kind="ids") == [expected] + assert core.registry().has_deployment(expected) + + +def test_generated_ids_are_the_same_across_cores_built_from_the_same_list(make_core: MakeCore) -> None: + first = make_core(model_list=[_bare("gpt", "openai/gpt-4o")]) + second = make_core(model_list=[_bare("gpt", "openai/gpt-4o")]) + + assert first.registry().names(kind="ids") == second.registry().names(kind="ids") + + +def test_generated_ids_differ_when_the_group_or_params_differ(make_core: MakeCore) -> None: + core = make_core( + model_list=[_bare("gpt", "openai/gpt-4o"), _bare("other", "openai/gpt-4o"), _bare("gpt", "openai/gpt-4o-mini")] + ) + + assert len(set(core.registry().names(kind="ids"))) == 3 + + +def test_a_declared_non_string_id_is_kept_as_a_string(make_core: MakeCore) -> None: + core = make_core( + model_list=[ + {"model_name": "gpt", "litellm_params": {"model": "openai/x", "api_key": "k"}, "model_info": {"id": 1}} + ] + ) + + assert core.registry().names(kind="ids") == ["1"] + found = core.registry().deployment("1") + assert found is not None + assert found.model_info.id == "1" + + +def test_two_deployments_declaring_one_id_are_both_kept_and_the_id_finds_the_last(make_core: MakeCore) -> None: + core = make_core(model_list=[deployment("gpt", "d"), deployment("emb", "d")]) + + assert [entry["model_name"] for entry in core.registry().all()] == ["gpt", "emb"] + assert core.registry().index("model_name") == {"gpt": [0], "emb": [1]} + found = core.registry().deployment("d") + assert found is not None + assert found.model_name == "emb" + + +# -- the stored entries ------------------------------------------------------------------------- + + +def test_the_stored_entry_carries_litellm_param_defaults_and_model_info_flags(make_core: MakeCore) -> None: + core = make_core(model_list=[deployment("gpt", "a")]) + + (entry,) = core.registry().all() + assert entry["model_name"] == "gpt" + assert entry["litellm_params"]["model"] == "openai/a" + assert entry["litellm_params"]["api_key"] == "sk-test" + assert entry["model_info"] == {"id": "a", "db_model": False, "member_auto_router": False} + + +def test_top_level_rpm_and_tpm_are_copied_into_litellm_params(make_core: MakeCore) -> None: + core = make_core( + model_list=[ + { + "model_name": "gpt", + "litellm_params": {"model": "openai/x", "api_key": "k"}, + "model_info": {"id": "r"}, + "rpm": 7, + "tpm": 9, + } + ] + ) + + (entry,) = core.registry().all() + assert (entry["litellm_params"]["rpm"], entry["litellm_params"]["tpm"]) == (7, 9) + assert (entry["rpm"], entry["tpm"]) == (7, 9) + + +def test_os_environ_references_in_litellm_params_are_resolved( + make_core: MakeCore, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setenv("REGISTRY_CHAR_BASE", "https://env.example") + monkeypatch.setenv("REGISTRY_CHAR_KEY", "sk-from-env") + core = make_core( + model_list=[ + deployment("gpt", "a", api_key="os.environ/REGISTRY_CHAR_KEY", api_base="os.environ/REGISTRY_CHAR_BASE") + ] + ) + + params = core.registry().all()[0]["litellm_params"] + assert (params["api_key"], params["api_base"]) == ("sk-from-env", "https://env.example") + + +@pytest.mark.asyncio +async def test_resolved_environment_credentials_are_sent_to_the_provider( + make_core: MakeCore, host: ScriptedHost, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setenv("REGISTRY_CHAR_BASE", "https://env.example") + core = make_core(model_list=[deployment("gpt", "a", api_base="os.environ/REGISTRY_CHAR_BASE")]) + + await core.acompletion(model="gpt", messages=[{"role": "user", "content": "hi"}]) # pyright: ignore[reportAttributeAccessIssue] # an endpoint method + + assert [call.kwargs.get("api_base") for call in host.calls_to("acompletion")] == ["https://env.example"] + + +def test_a_list_of_organizations_fans_out_into_one_deployment_per_organization(make_core: MakeCore) -> None: + core = make_core(model_list=[_bare("gpt", "openai/x", organization=["org-1", "org-2"])]) + + entries = core.registry().all() + assert [entry["litellm_params"]["organization"] for entry in entries] == ["org-1", "org-2"] + assert len(set(ids(entries))) == 1 + assert core.registry().index("model_name") == {"gpt": [0, 1]} + + +def test_deployment_names_list_each_deployments_litellm_model(make_core: MakeCore) -> None: + core = make_core(model_list=[deployment("gpt", "a"), deployment("gpt", "b"), deployment("emb", "e")]) + + assert core.registry().names(kind="deployment_names") == ["openai/a", "openai/b", "openai/e"] + + +def test_every_configured_litellm_model_starts_with_zero_latency(make_core: MakeCore) -> None: + core = make_core(model_list=[deployment("gpt", "a"), deployment("gpt", "b")]) + + assert core.health().latency == {"openai/a": 0, "openai/b": 0} + + +@pytest.mark.parametrize("model_list", [None, []], ids=["none", "empty"]) +def test_a_core_without_deployments_has_an_empty_registry(make_core: MakeCore, model_list: object) -> None: + core = make_core(model_list=model_list) + + assert core.registry().all() == [] + assert core.registry().names(kind="ids") == [] + assert core.registry().names(kind="group_set") == set() + assert core.registry().deployments() == [] + + +# -- invalid deployments ------------------------------------------------------------------------ + + +def test_an_unknown_provider_raises_bad_request_at_construction(make_core: MakeCore) -> None: + with pytest.raises(litellm.BadRequestError, match="LLM Provider NOT provided"): + make_core(model_list=[deployment("gpt", "a", model="notaprovider/model")]) + + +def test_ignore_invalid_deployments_drops_the_invalid_one_and_keeps_the_rest(make_core: MakeCore) -> None: + core = make_core( + model_list=[deployment("gpt", "a", model="notaprovider/model"), deployment("gpt", "b")], + ignore_invalid_deployments=True, + ) + + assert core.registry().names(kind="ids") == ["b"] + + +def test_an_invalid_tag_regex_raises_value_error_naming_the_model(make_core: MakeCore) -> None: + with pytest.raises(ValueError, match=r"Invalid regex in tag_regex for model 'gpt': '\('"): + make_core(model_list=[deployment("gpt", "a", tag_regex=["("])]) + + +def test_ignore_invalid_deployments_drops_a_deployment_with_an_invalid_tag_regex(make_core: MakeCore) -> None: + core = make_core( + model_list=[deployment("gpt", "a", tag_regex=["("]), deployment("gpt", "b", tag_regex=["^ok$"])], + ignore_invalid_deployments=True, + ) + + assert core.registry().names(kind="ids") == ["b"] + + +# -- environment-gated deployments -------------------------------------------------------------- + + +def test_a_deployment_for_another_environment_is_left_out(make_core: MakeCore, monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv("LITELLM_ENVIRONMENT", "staging") + core = make_core( + model_list=[ + deployment("prod-only", "a", model_info={"supported_environments": ["production"]}), + deployment("both", "b", model_info={"supported_environments": ["staging", "production"]}), + ] + ) + + assert core.registry().names(kind="ids") == ["b"] + assert core.registry().deployments("prod-only") == [] + + +def test_a_left_out_environment_deployment_still_names_its_group( + make_core: MakeCore, monkeypatch: pytest.MonkeyPatch +) -> None: + """Odd but current: the group set is read from the given model_list, not from what was added.""" + monkeypatch.setenv("LITELLM_ENVIRONMENT", "staging") + core = make_core( + model_list=[ + deployment("prod-only", "a", model_info={"supported_environments": ["production"]}), + deployment("h", "b"), + ] + ) + + assert core.registry().names(kind="group_set") == {"prod-only", "h"} + assert core.registry().index("model_name") == {"h": [0]} + + +def test_supported_environments_without_litellm_environment_raises( + make_core: MakeCore, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.delenv("LITELLM_ENVIRONMENT", raising=False) + + with pytest.raises(ValueError, match="Set 'supported_environments' for model but not 'LITELLM_ENVIRONMENT' set"): + make_core(model_list=[deployment("gpt", "a", model_info={"supported_environments": ["production"]})]) + + +@pytest.mark.parametrize( + ("environment", "supported", "message"), + [ + ("bogus", ["production"], "LITELLM_ENVIRONMENT must be one of"), + ("staging", ["prod"], "supported_environments must be one of"), + ("staging", "production", "supported_environments must be a list of"), + ], + ids=["unknown_environment", "unknown_supported_value", "not_a_list"], +) +def test_invalid_environment_configuration_raises( + make_core: MakeCore, monkeypatch: pytest.MonkeyPatch, environment: str, supported: object, message: str +) -> None: + monkeypatch.setenv("LITELLM_ENVIRONMENT", environment) + + with pytest.raises(ValueError, match=message): + make_core(model_list=[deployment("gpt", "a", model_info={"supported_environments": supported})]) + + +def test_ignore_invalid_deployments_drops_an_invalid_environment_deployment( + make_core: MakeCore, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.delenv("LITELLM_ENVIRONMENT", raising=False) + core = make_core( + model_list=[ + deployment("gpt", "a", model_info={"supported_environments": ["production"]}), + deployment("gpt", "b"), + ], + ignore_invalid_deployments=True, + ) + + assert core.registry().names(kind="ids") == ["b"] + + +@pytest.mark.parametrize( + ("model_info", "active"), + [ + ({}, True), + ({"supported_environments": ["development"]}, True), + ({"supported_environments": ["production"]}, False), + ], + ids=["unset", "matching", "other"], +) +def test_deployment_is_active_for_environment_follows_litellm_environment( + make_core: MakeCore, monkeypatch: pytest.MonkeyPatch, model_info: dict[str, object], active: bool +) -> None: + core = make_core(model_list=[deployment("gpt", "a")]) + monkeypatch.setenv("LITELLM_ENVIRONMENT", "development") + candidate = litellm.types.router.Deployment(**deployment("gpt", "z", model_info=model_info)) # pyright: ignore[reportArgumentType] # a model_list entry + + assert core.deployment_is_active_for_environment(candidate) is active # pyright: ignore[reportAttributeAccessIssue] # public on the core + + +# -- wildcard and team registration ------------------------------------------------------------- + + +def test_a_wildcard_group_registers_a_pattern_and_a_provider_default_deployment(make_core: MakeCore) -> None: + core = make_core( + model_list=[ + deployment("openai/*", "w", model="openai/*"), + deployment("anthropic/*", "x", model="anthropic/*"), + deployment("gpt", "a"), + ] + ) + + assert core.settings().provider_default_deployment_ids == ["w", "x"] + assert ids(core.registry().deployments("openai/gpt-4o")) == ["w"] + assert ids(core.registry().deployments("anthropic/claude-x")) == ["x"] + + +def test_team_deployments_are_indexed_by_team_and_public_name(make_core: MakeCore) -> None: + core = make_core( + model_list=[ + team_deployment("internal-1", "t1", "team-a", "pub"), + team_deployment("internal-2", "t2", "team-a", "pub"), + team_deployment("internal-3", "t3", "team-b", "other"), + deployment("gpt", "g"), + ] + ) + + assert core.registry().index("team_model") == {("team-a", "pub"): [0, 1], ("team-b", "other"): [2]} + assert core.registry().names(kind="team_public") == frozenset({"pub", "other"}) + assert core.registry().names(kind="team_public", team_id="team-a") == frozenset({"pub"}) + assert core.registry().names(kind="team_public", team_id="nobody") == frozenset() + + +def test_a_team_deployment_needs_both_team_id_and_public_name_to_be_team_indexed(make_core: MakeCore) -> None: + core = make_core( + model_list=[ + deployment("gpt", "a", model_info={"team_id": "team-a"}), + deployment("gpt", "b", model_info={"team_public_model_name": "pub"}), + ] + ) + + assert core.registry().index("team_model") == {} + assert core.registry().names(kind="team_public") == frozenset() + + +# -- pricing written into litellm.model_cost ----------------------------------------------------- + + +def test_custom_pricing_is_registered_under_the_deployment_id(make_core: MakeCore) -> None: + model_id = unique_id() + make_core( + model_list=[ + deployment("gpt", model_id, model="openai/gpt-4o", input_cost_per_token=1e-3, output_cost_per_token=2e-3) + ] + ) + + entry = litellm.model_cost[model_id] + assert (entry["input_cost_per_token"], entry["output_cost_per_token"]) == (1e-3, 2e-3) + + +def test_custom_input_pricing_inherits_the_backend_models_cache_read_price(make_core: MakeCore) -> None: + model_id = unique_id() + make_core(model_list=[deployment("gpt", model_id, model="openai/gpt-4o", input_cost_per_token=1e-3)]) + + builtin = litellm.get_model_info("openai/gpt-4o") + assert litellm.model_cost[model_id]["cache_read_input_token_cost"] == builtin["cache_read_input_token_cost"] + + +def test_a_user_set_cache_read_price_is_not_overwritten_by_inheritance(make_core: MakeCore) -> None: + model_id = unique_id() + make_core( + model_list=[ + deployment( + "gpt", model_id, model="openai/gpt-4o", input_cost_per_token=1e-3, cache_read_input_token_cost=7e-9 + ) + ] + ) + + assert litellm.model_cost[model_id]["cache_read_input_token_cost"] == 7e-9 + + +def test_a_deployment_without_custom_pricing_registers_no_prices_under_its_id(make_core: MakeCore) -> None: + model_id = unique_id() + make_core(model_list=[deployment("gpt", model_id, model="openai/gpt-4o")]) + + entry = litellm.model_cost[model_id] + assert entry.get("input_cost_per_token") is None + assert entry.get("output_cost_per_token") is None + + +def test_model_info_fields_are_registered_under_the_deployment_id(make_core: MakeCore) -> None: + model_id = unique_id() + make_core(model_list=[deployment("gpt", model_id, model_info={"max_input_tokens": 1000, "mode": "embedding"})]) + + entry = litellm.model_cost[model_id] + assert (entry["max_input_tokens"], entry["mode"]) == (1000, "embedding") diff --git a/tests/unit/test_router/core/test_registry_characterization_lookups.py b/tests/unit/test_router/core/test_registry_characterization_lookups.py new file mode 100644 index 00000000000..19d58ff10d4 --- /dev/null +++ b/tests/unit/test_router/core/test_registry_characterization_lookups.py @@ -0,0 +1,682 @@ +""" +Registry reads: deployments by group, alias, wildcard and team; names and ids; name resolution; +access groups; blocked and unhealthy groups; credentials lookups. +""" + +import time + +import pytest + +from litellm.router_core import RouterCore +from tests.unit.test_router.core.registry_helpers import MakeCore, ids, names_and_ids, team_deployment, typed +from tests.unit.test_router.core.scripted_host import deployment + + +def _core(make_core: MakeCore) -> RouterCore: + """ + gpt (a, b), emb (e), the openai/* wildcard (w), team TA's ``pub`` (t1) and ``openai/*`` (tw), + team TB's ``pub`` (t2). Aliases: fast -> gpt, hid -> gpt (hidden), gpt -> emb. + """ + return make_core( + model_list=[ + deployment("gpt", "a", model_info={"access_groups": ["ag1", "ag2"]}), + deployment("gpt", "b", model_info={"access_groups": ["ag1"]}), + deployment("emb", "e", model="openai/text-embedding-3-small", model_info={"access_groups": ["ag2"]}), + deployment("openai/*", "w", model="openai/*", model_info={"access_groups": ["wild"]}), + team_deployment("int-a", "t1", "TA", "pub", model="openai/team-model"), + team_deployment("int-b", "t2", "TB", "pub", model="anthropic/claude-x"), + team_deployment("team-wild", "tw", "TA", "openai/*", model="openai/*"), + ], + model_group_alias={"fast": "gpt", "hid": {"model": "gpt", "hidden": True}, "gpt": "emb"}, + ) + + +# -- registry().deployments(): the group listing ------------------------------------------------ + + +def test_listing_every_deployment_puts_visible_alias_rows_before_the_model_list(make_core: MakeCore) -> None: + listed = _core(make_core).registry().deployments() + + assert names_and_ids(listed) == [ + ("fast", "a"), + ("fast", "b"), + ("gpt", "e"), + ("gpt", "a"), + ("gpt", "b"), + ("emb", "e"), + ("openai/*", "w"), + ("int-a", "t1"), + ("int-b", "t2"), + ("team-wild", "tw"), + ] + + +def test_an_alias_lists_its_targets_deployments_under_the_alias_name(make_core: MakeCore) -> None: + assert names_and_ids(_core(make_core).registry().deployments("fast")) == [("fast", "a"), ("fast", "b")] + + +def test_a_hidden_alias_is_listed_when_named(make_core: MakeCore) -> None: + assert names_and_ids(_core(make_core).registry().deployments("hid")) == [("hid", "a"), ("hid", "b")] + + +def test_a_group_that_is_also_an_alias_lists_its_own_deployments_then_the_alias_targets(make_core: MakeCore) -> None: + assert names_and_ids(_core(make_core).registry().deployments("gpt")) == [("gpt", "a"), ("gpt", "b"), ("gpt", "e")] + + +def test_an_alias_row_is_a_copy_and_leaves_the_configured_entry_named_as_configured(make_core: MakeCore) -> None: + core = _core(make_core) + + core.registry().deployments("fast") + + assert [entry["model_name"] for entry in core.registry().all()[:2]] == ["gpt", "gpt"] + + +def test_an_unknown_name_lists_nothing(make_core: MakeCore) -> None: + assert _core(make_core).registry().deployments("nope") == [] + + +def test_a_name_matching_only_a_wildcard_lists_the_wildcard_deployment_under_the_requested_name( + make_core: MakeCore, +) -> None: + assert names_and_ids(_core(make_core).registry().deployments("openai/gpt-x")) == [("openai/gpt-x", "w")] + + +def test_a_team_caller_gets_its_own_wildcard_deployment_after_the_global_one(make_core: MakeCore) -> None: + listed = _core(make_core).registry().deployments("openai/gpt-x", team_id="TA") + + assert names_and_ids(listed) == [("openai/gpt-x", "w"), ("openai/gpt-x", "tw")] + + +@pytest.mark.parametrize(("team_id", "expected"), [("TA", [("int-a", "t1")]), ("TB", [("int-b", "t2")]), (None, [])]) +def test_a_team_public_name_lists_only_that_teams_deployment( + make_core: MakeCore, team_id: str | None, expected: list[tuple[str, str]] +) -> None: + assert names_and_ids(_core(make_core).registry().deployments("pub", team_id=team_id)) == expected + + +@pytest.mark.parametrize(("team_id", "expected"), [(None, ["t1"]), ("TA", ["t1"]), ("TB", [])]) +def test_a_team_deployments_internal_name_lists_it_for_no_team_or_its_own_team_only( + make_core: MakeCore, team_id: str | None, expected: list[str] +) -> None: + assert ids(_core(make_core).registry().deployments("int-a", team_id=team_id)) == expected + + +def test_all_deployments_renames_rows_to_the_given_model_alias(make_core: MakeCore) -> None: + listed = _core(make_core).registry().deployments("gpt", source="all", model_alias="renamed") + + assert names_and_ids(listed) == [("renamed", "a"), ("renamed", "b")] + + +def test_all_deployments_follow_no_alias_and_no_wildcard(make_core: MakeCore) -> None: + registry = _core(make_core).registry() + + assert registry.deployments("fast", source="all") == [] + assert registry.deployments("openai/gpt-x", source="all") == [] + + +def test_all_deployments_for_a_team_public_name_reads_the_team_index(make_core: MakeCore) -> None: + assert ids(_core(make_core).registry().deployments("pub", source="all", team_id="TA")) == ["t1"] + + +def test_named_deployments_are_the_group_index_entry_or_none(make_core: MakeCore) -> None: + registry = _core(make_core).registry() + + assert ids(registry.deployments("gpt", source="named")) == ["a", "b"] + assert registry.deployments("fast", source="named") is None + + +def test_first_in_group_is_the_first_configured_deployment_of_the_group(make_core: MakeCore) -> None: + registry = _core(make_core).registry() + + first = registry.deployments("gpt", source="first_in_group") + assert first is not None + assert first.model_info.id == "a" + assert registry.deployments("fast", source="first_in_group") is None + + +@pytest.mark.parametrize("source", ["all", "named", "wildcard", "routable", "first_in_group"]) +def test_a_source_that_needs_a_model_name_raises_without_one(make_core: MakeCore, source: str) -> None: + with pytest.raises(TypeError, match="needs a model_name"): + _core(make_core).registry().deployments(source=source) # pyright: ignore[reportCallIssue, reportArgumentType] # source is one of the literals + + +def test_an_unknown_deployment_source_raises(make_core: MakeCore) -> None: + with pytest.raises(ValueError, match="unknown deployment source 'bogus'"): + _core(make_core).registry().deployments("gpt", source="bogus") # pyright: ignore[reportCallIssue, reportArgumentType] # an invalid source on purpose + + +# -- wildcards ------------------------------------------------------------------------------------ + + +@pytest.mark.parametrize(("team_id", "expected"), [("TA", ["tw"]), ("TB", ["w"]), (None, ["w"])]) +def test_wildcard_route_deployments_prefer_the_teams_own_pattern_over_the_global_one( + make_core: MakeCore, team_id: str | None, expected: list[str] +) -> None: + assert ids(_core(make_core).registry().deployments("openai/gpt-x", source="wildcard", team_id=team_id)) == expected + + +def test_wildcard_routed_models_name_the_requested_model(make_core: MakeCore) -> None: + registry = _core(make_core).registry() + + assert registry.names("openai/gpt-x", kind="wildcard") == frozenset({"openai/gpt-x"}) + assert registry.names("anthropic/x", kind="wildcard") == frozenset() + + +# -- names ------------------------------------------------------------------------------------------ + + +def test_group_names_repeat_per_deployment_and_include_visible_aliases(make_core: MakeCore) -> None: + assert _core(make_core).registry().names() == ["fast", "fast", "gpt", "gpt", "gpt", "emb", "openai/*"] + + +def test_group_names_for_a_team_add_its_public_names(make_core: MakeCore) -> None: + assert _core(make_core).registry().names(team_id="TA") == [ + "fast", + "fast", + "gpt", + "gpt", + "gpt", + "emb", + "openai/*", + "pub", + "openai/*", + ] + + +def test_ids_can_leave_out_team_deployments(make_core: MakeCore) -> None: + registry = _core(make_core).registry() + + assert registry.names(kind="ids") == ["a", "b", "e", "w", "t1", "t2", "tw"] + assert registry.names(kind="ids", exclude_team_models=True) == ["a", "b", "e", "w"] + + +def test_ids_of_a_group_read_only_the_group_index(make_core: MakeCore) -> None: + registry = _core(make_core).registry() + + assert registry.names("fast", kind="ids") == [] + assert registry.names("pub", kind="ids") == [] + assert registry.names("int-a", kind="ids") == ["t1"] + assert registry.names("int-a", kind="ids", exclude_team_models=True) == [] + + +def test_litellm_models_follow_aliases_and_wildcards(make_core: MakeCore) -> None: + registry = _core(make_core).registry() + + assert registry.names("fast", kind="litellm_models") == ("openai/a", "openai/b") + assert registry.names("openai/gpt-x", kind="litellm_models") == ("openai/gpt-x",) + assert registry.names("nope", kind="litellm_models") == () + + +def test_litellm_base_models_use_a_deployments_base_model(make_core: MakeCore) -> None: + core = make_core( + model_list=[ + deployment( + "gpt", "a", model="azure/my-deployment", api_base="https://x", model_info={"base_model": "azure/gpt-4o"} + ), + deployment("gpt", "b", model="openai/gpt-4o-mini"), + ] + ) + + assert core.registry().names("gpt", kind="litellm_base_models") == ("azure/gpt-4o", "openai/gpt-4o-mini") + assert core.registry().names("gpt", kind="litellm_models") == ("azure/my-deployment", "openai/gpt-4o-mini") + + +@pytest.mark.parametrize( + ("name", "team_id", "expected"), + [("fast", None, {"a", "b"}), ("openai/gpt-x", None, {"w"}), ("pub", "TA", {"t1"}), ("nope", None, set())], +) +def test_route_candidates_are_the_ids_a_request_for_the_name_may_reach( + make_core: MakeCore, name: str, team_id: str | None, expected: set[str] +) -> None: + assert _core(make_core).registry().names(name, kind="route_candidates", team_id=team_id) == expected + + +@pytest.mark.parametrize("kind", ["litellm_models", "wildcard", "route_candidates"]) +def test_a_name_kind_that_needs_a_model_name_raises_without_one(make_core: MakeCore, kind: str) -> None: + with pytest.raises(TypeError, match="needs a model_name"): + _core(make_core).registry().names(kind=kind) # pyright: ignore[reportCallIssue, reportArgumentType] # kind is one of the literals + + +# -- resolve ---------------------------------------------------------------------------------------- + + +@pytest.mark.parametrize( + ("name", "expected"), + [("fast", "gpt"), ("hid", "gpt"), ("gpt", "emb"), ("emb", "emb"), ("nope", "nope")], +) +def test_resolving_a_group_follows_one_alias_hop(make_core: MakeCore, name: str, expected: str) -> None: + assert _core(make_core).registry().resolve(name) == expected + + +@pytest.mark.parametrize(("name", "expected"), [("fast", "gpt"), ("hid", "gpt"), ("emb", None), ("pub", None)]) +def test_resolving_an_alias_reads_plain_and_hidden_aliases( + make_core: MakeCore, name: str, expected: str | None +) -> None: + assert _core(make_core).registry().resolve(name, kind="alias") == expected + + +@pytest.mark.parametrize( + ("name", "expected"), + [ + ("fast", True), + ("hid", True), + ("gpt", True), + ("a", True), + ("pub", False), + ("openai/gpt-x", False), + ("nope", False), + ], +) +def test_recognized_names_are_groups_aliases_and_deployment_ids(make_core: MakeCore, name: str, expected: bool) -> None: + assert _core(make_core).registry().resolve(name, kind="recognized") is expected + + +@pytest.mark.parametrize( + ("name", "team_id", "expected"), + [ + ("fast", None, "openai/a"), + ("gpt", None, "openai/text-embedding-3-small"), + ("pub", "TA", "openai/team-model"), + ("pub", None, None), + ("openai/gpt-x", None, None), + ("a", None, None), + ], +) +def test_upstream_is_the_first_routable_deployments_litellm_model( + make_core: MakeCore, name: str, team_id: str | None, expected: str | None +) -> None: + assert _core(make_core).registry().resolve(name, kind="upstream", team_id=team_id) == expected + + +def test_upstream_skips_a_blocked_deployment(make_core: MakeCore) -> None: + core = make_core(model_list=[deployment("half", "x2", model_info={"blocked": True}), deployment("half", "x3")]) + + assert core.registry().resolve("half", kind="upstream") == "openai/x3" + + +@pytest.mark.parametrize( + ("served", "expected"), [("gpt", "gpt"), ("emb", "fast"), (None, "fast")], ids=["target", "other", "none"] +) +def test_the_served_name_is_the_served_model_only_when_the_name_routes_to_it( + make_core: MakeCore, served: str | None, expected: str +) -> None: + assert _core(make_core).registry().resolve("fast", kind="served", served_model=served) == expected + + +def test_group_of_id_lists_the_deployments_group_including_alias_rows(make_core: MakeCore) -> None: + registry = _core(make_core).registry() + + assert names_and_ids(registry.resolve("a", kind="group_of_id")) == [("gpt", "a"), ("gpt", "b"), ("gpt", "e")] + assert registry.resolve("gpt", kind="group_of_id") is None + + +@pytest.mark.parametrize( + ("name", "expected"), + [ + ("gpt", "gpt"), + ("a", "gpt"), + ("fast", "fast"), + ("nope", None), + (None, None), + ("", None), + ], +) +def test_from_id_resolves_a_group_name_or_deployment_id( + make_core: MakeCore, name: str | None, expected: str | None +) -> None: + assert _core(make_core).registry().resolve(name, kind="from_id") == expected + + +def test_from_id_matches_a_litellm_model_suffix_and_the_first_matching_row_wins(make_core: MakeCore) -> None: + """Odd but current: the alias row ``gpt -> emb`` is listed first, so ``emb``'s model resolves to ``gpt``.""" + registry = _core(make_core).registry() + + assert registry.resolve("text-embedding-3-small", kind="from_id") == "gpt" + assert registry.resolve("openai/text-embedding-3-small", kind="from_id") == "gpt" + + +def test_from_id_suffix_match_on_a_plain_core(make_core: MakeCore) -> None: + core = make_core( + model_list=[deployment("veo", "v", model="vertex_ai/veo-2.0", vertex_project="p", vertex_location="l")] + ) + + assert core.registry().resolve("veo-2.0", kind="from_id") == "veo" + + +def test_from_id_with_a_provider_returns_the_provider_prefixed_name_a_wildcard_serves(make_core: MakeCore) -> None: + registry = _core(make_core).registry() + + assert registry.resolve("zzz", kind="from_id", custom_llm_provider="openai") == "openai/zzz" + assert registry.resolve("zzz", kind="from_id", custom_llm_provider="anthropic") is None + + +@pytest.mark.parametrize( + ("name", "team_id", "provider", "expected"), + [ + ("pub", "TA", None, "pub"), + ("gpt", "TA", None, None), + (None, "TA", None, "pub"), + (None, "TA", "openai", "pub"), + (None, "TB", "openai", None), + (None, "TB", "anthropic", "pub"), + (None, "nobody", None, None), + ], +) +def test_team_model_maps_a_teams_public_name_or_picks_its_first_deployment( + make_core: MakeCore, name: str | None, team_id: str, provider: str | None, expected: str | None +) -> None: + registry = _core(make_core).registry() + + assert registry.resolve(name, kind="team_model", team_id=team_id, custom_llm_provider=provider) == expected + + +def test_team_model_without_a_team_id_raises(make_core: MakeCore) -> None: + with pytest.raises(TypeError, match="needs a team_id"): + _core(make_core).registry().resolve("pub", kind="team_model") + + +def test_an_unknown_resolve_kind_raises(make_core: MakeCore) -> None: + with pytest.raises(ValueError, match="unknown resolve kind 'bogus'"): + _core(make_core).registry().resolve("gpt", kind="bogus") # pyright: ignore[reportCallIssue, reportArgumentType] # an invalid kind on purpose + + +# -- includes ------------------------------------------------------------------------------------- + + +@pytest.mark.parametrize(("team_id", "expected"), [(None, False), ("TA", True), ("TB", False)]) +def test_a_team_deployment_is_usable_only_by_its_own_team( + make_core: MakeCore, team_id: str | None, expected: bool +) -> None: + registry = _core(make_core).registry() + team_entry = registry.all()[4] + + assert registry.includes(team_entry, team_id=team_id) is expected + + +def test_a_shared_deployment_is_usable_by_any_team(make_core: MakeCore) -> None: + registry = _core(make_core).registry() + + assert all(registry.includes(registry.all()[0], team_id=team) for team in (None, "TA", "TB")) + + +@pytest.mark.parametrize( + ("model_name", "team_id", "expected"), + [("pub", "TA", True), ("pub", "TB", False), ("int-a", None, True), ("int-a", "TB", False), ("gpt", "TA", False)], +) +def test_includes_with_a_model_name_matches_the_public_or_internal_name( + make_core: MakeCore, model_name: str, team_id: str | None, expected: bool +) -> None: + registry = _core(make_core).registry() + + assert registry.includes(registry.all()[4], model_name=model_name, team_id=team_id) is expected + + +# -- access groups -------------------------------------------------------------------------------- + + +def test_access_groups_map_each_group_to_the_model_names_that_carry_it(make_core: MakeCore) -> None: + assert _core(make_core).registry().access_groups() == { + "ag1": ["fast", "fast", "gpt", "gpt"], + "ag2": ["fast", "gpt", "gpt", "emb"], + "wild": ["openai/*"], + } + + +def test_access_groups_of_one_model_read_its_listing(make_core: MakeCore) -> None: + registry = _core(make_core).registry() + + assert dict(registry.access_groups(model_name="gpt")) == {"ag1": ["gpt", "gpt"], "ag2": ["gpt", "gpt"]} + assert dict(registry.access_groups(model_name="openai/gpt-x")) == {"wild": ["openai/gpt-x"]} + assert dict(registry.access_groups(model_name="nope")) == {} + + +def test_access_groups_filtered_to_one_group(make_core: MakeCore) -> None: + assert dict(_core(make_core).registry().access_groups(model_access_group="ag2")) == { + "ag2": ["fast", "gpt", "gpt", "emb"] + } + + +def test_access_groups_follow_a_deployment_added_later(make_core: MakeCore) -> None: + core = _core(make_core) + assert "late" not in core.registry().access_groups() + + core.add_deployment( # pyright: ignore[reportAttributeAccessIssue] # public on the core + typed("new", "n", model_info={"access_groups": ["late"]}) + ) + + assert core.registry().access_groups()["late"] == ["new"] + + +@pytest.mark.parametrize(("group", "expected"), [("wild", True), ("ag1", False), ("nope", False)]) +def test_an_access_group_is_a_wildcard_route_when_one_of_its_models_is_a_pattern( + make_core: MakeCore, group: str, expected: bool +) -> None: + assert _core(make_core).registry().resolve(group, kind="wildcard_access_group") is expected + + +# -- blocked and unhealthy groups ----------------------------------------------------------------- + + +def _blocked_core(make_core: MakeCore) -> RouterCore: + return make_core( + model_list=[ + deployment("blk", "x1", model_info={"blocked": True}), + deployment("half", "x2", model_info={"blocked": True}), + deployment("half", "x3"), + deployment("ok", "x4"), + ], + model_group_alias={"blk-alias": "blk"}, + ) + + +def test_fully_blocked_names_are_groups_whose_every_deployment_is_blocked(make_core: MakeCore) -> None: + assert _blocked_core(make_core).health().fully_blocked_model_names() == {"blk", "blk-alias"} + + +def test_is_blocked_and_all_blocked_read_model_info_blocked(make_core: MakeCore) -> None: + core = _blocked_core(make_core) + health = core.health() + entries = core.registry().all() + + blocked = core.registry().deployment("x1") + open_ = core.registry().deployment("x3") + assert blocked is not None + assert open_ is not None + assert (health.is_blocked(blocked), health.is_blocked(open_)) == (True, False) + assert health.all_blocked(entries[:1]) is True + assert health.all_blocked(entries[1:3]) is False + assert health.all_blocked([]) is False + + +@pytest.mark.asyncio +async def test_fully_unhealthy_names_are_groups_whose_every_deployment_failed_its_health_check( + make_core: MakeCore, +) -> None: + core = make_core( + model_list=[ + deployment("sick", "s1"), + deployment("half", "h1"), + deployment("half", "h2"), + team_deployment("int", "t1", "TA", "team-sick"), + ] + ) + now = time.time() + core.health().health_checks.set_deployment_health_states( + { + "s1": {"is_healthy": False, "timestamp": now, "reason": "down"}, + "h1": {"is_healthy": False, "timestamp": now, "reason": "down"}, + "h2": {"is_healthy": True, "timestamp": now, "reason": ""}, + "t1": {"is_healthy": False, "timestamp": now, "reason": "down"}, + } + ) + + assert await core.health().fully_unhealthy_model_names() == {"sick", "int", "team-sick"} + + +@pytest.mark.asyncio +async def test_no_health_state_hides_nothing(make_core: MakeCore) -> None: + core = make_core(model_list=[deployment("gpt", "a")]) + + assert await core.health().fully_unhealthy_model_names() == set() + + +@pytest.mark.asyncio +async def test_a_stale_unhealthy_state_hides_nothing(make_core: MakeCore) -> None: + core = make_core(model_list=[deployment("gpt", "a")]) + stale = time.time() - 10 * core.health().health_checks.staleness_threshold + core.health().health_checks.set_deployment_health_states( + {"a": {"is_healthy": False, "timestamp": stale, "reason": "x"}} + ) + + assert await core.health().fully_unhealthy_model_names() == set() + + +@pytest.mark.asyncio +async def test_an_allowed_fails_policy_hides_no_unhealthy_group(make_core: MakeCore) -> None: + core = make_core(model_list=[deployment("gpt", "a")], allowed_fails_policy={"RateLimitErrorAllowedFails": 2}) + core.health().health_checks.set_deployment_health_states( + {"a": {"is_healthy": False, "timestamp": time.time(), "reason": "x"}} + ) + + assert await core.health().fully_unhealthy_model_names() == set() + + +# -- credentials ------------------------------------------------------------------------------------ + + +def _credential_core(make_core: MakeCore) -> RouterCore: + return make_core( + model_list=[ + deployment("gpt", "a", api_base="https://a.example", temperature=0.2), + deployment("plain", "p", model="gpt-4o-mini"), + deployment("claude", "c", model="claude-3-haiku", custom_llm_provider="anthropic"), + deployment("blk", "x", model_info={"blocked": True}), + deployment("openai/*", "w", model="openai/*"), + team_deployment("shared", "ta", "TA", "team-name", api_key="sk-team-a"), + team_deployment("team-wild", "tw", "TA", "openai/*", model="openai/*", api_key="sk-team-wild"), + ] + ) + + +def test_credentials_are_the_credential_params_only(make_core: MakeCore) -> None: + assert _credential_core(make_core).registry().credentials("a") == { + "api_key": "sk-test", + "api_base": "https://a.example", + } + + +def test_credentials_of_an_unknown_or_blocked_deployment_are_none(make_core: MakeCore) -> None: + registry = _credential_core(make_core).registry() + + assert registry.credentials("missing") is None + assert registry.credentials("x") is None + assert registry.credentials("x", form="with_provider") is None + assert registry.credentials("x", form="deployment") is None + + +@pytest.mark.parametrize( + ("lookup", "provider"), + [("a", "openai"), ("p", "openai"), ("c", "anthropic")], + ids=["prefix", "default", "explicit"], +) +def test_credentials_with_provider_add_the_model_and_its_provider( + make_core: MakeCore, lookup: str, provider: str +) -> None: + registry = _credential_core(make_core).registry() + + found = registry.credentials(lookup, form="with_provider") + deployment_found = registry.deployment(lookup) + assert found is not None + assert deployment_found is not None + assert (found["model"], found["custom_llm_provider"]) == (deployment_found.litellm_params.model, provider) + + +@pytest.mark.parametrize( + ("lookup", "team_id", "expected"), + [ + ("a", None, "a"), + ("gpt", None, "a"), + ("team-name", "TA", "ta"), + ("team-name", None, None), + ("shared", "TB", None), + ("shared", "TA", "ta"), + ("openai/gpt-x", None, "w"), + ("openai/gpt-x", "TA", "tw"), + ("openai/gpt-x", "TB", "w"), + ("nope", None, None), + ], +) +def test_the_credential_deployment_is_found_by_id_group_team_name_then_wildcard( + make_core: MakeCore, lookup: str, team_id: str | None, expected: str | None +) -> None: + found = _credential_core(make_core).registry().credentials(lookup, form="deployment", team_id=team_id) + + assert (None if found is None else found.model_info.id) == expected + + +def test_credentials_with_provider_for_a_wildcard_name_carry_the_requested_model(make_core: MakeCore) -> None: + found = _credential_core(make_core).registry().credentials("openai/gpt-x", form="with_provider", team_id="TA") + + assert found is not None + assert (found["api_key"], found["model"], found["custom_llm_provider"]) == ( + "sk-team-wild", + "openai/gpt-x", + "openai", + ) + + +def test_an_unknown_credential_form_raises(make_core: MakeCore) -> None: + with pytest.raises(ValueError, match="unknown credential form 'bogus'"): + _credential_core(make_core).registry().credentials("a", form="bogus") # pyright: ignore[reportCallIssue, reportArgumentType] # an invalid form on purpose + + +# -- single deployments and the raw list ------------------------------------------------------------ + + +def test_deployment_dict_is_the_live_model_list_entry(make_core: MakeCore) -> None: + registry = _core(make_core).registry() + + assert registry.deployment_dict("a") is registry.all()[0] + assert registry.deployment_dict("missing") is None + + +def test_a_looked_up_deployment_is_a_copy(make_core: MakeCore) -> None: + core = _core(make_core) + found = core.registry().deployment("a") + assert found is not None + + found.litellm_params.api_base = "https://mutated.example" + + assert core.registry().all()[0]["litellm_params"].get("api_base") is None + + +def test_a_visible_dict_alias_is_listed_like_a_plain_alias(make_core: MakeCore) -> None: + core = make_core(model_list=[deployment("gpt", "a")], model_group_alias={"vis": {"model": "gpt", "hidden": False}}) + + assert names_and_ids(core.registry().deployments()) == [("vis", "a"), ("gpt", "a")] + assert core.registry().names() == ["vis", "gpt"] + + +def test_the_deprecated_generate_model_id_is_the_registry_id_hash(make_core: MakeCore) -> None: + from litellm.router_api import deployment_id + + core = _core(make_core) + + with pytest.warns(DeprecationWarning, match="generate_model_id is deprecated"): + generated = core.generate_model_id("g", {"model": "x"}) # pyright: ignore[reportAttributeAccessIssue] # deprecated, still public + + assert generated == deployment_id("g", {"model": "x"}) + + +def test_assigning_model_list_swaps_the_list_without_reindexing(make_core: MakeCore) -> None: + """Odd but current: the id index still maps ``a`` to position 0, which now holds the new entry.""" + core = make_core(model_list=[deployment("gpt", "a")]) + version = core.registry().version + + core.model_list = [deployment("n", "n")] # pyright: ignore[reportAttributeAccessIssue] # deprecated attribute, still assignable + + assert core.registry().version == version + 1 + assert core.registry().names(kind="ids") == ["a"] + found = core.registry().deployment("a") + assert found is not None + assert found.model_name == "n" diff --git a/tests/unit/test_router/core/test_registry_characterization_model_info.py b/tests/unit/test_router/core/test_registry_characterization_model_info.py new file mode 100644 index 00000000000..8d287ab46a2 --- /dev/null +++ b/tests/unit/test_router/core/test_registry_characterization_model_info.py @@ -0,0 +1,580 @@ +""" +What the registry knows about models: group info, deployment and router model info, configured +limits/mode/display name/service tiers/listing info, group usage this minute, remaining usage and +the rate limit headers ``set_response_headers`` attaches. +""" + +import pytest + +import litellm +from litellm.litellm_core_utils.get_model_cost_map import GetModelCostMap +from litellm.router_core import RouterCore +from litellm.types.router import Deployment, RouterCacheEnum +from litellm.utils import get_utc_datetime +from tests.unit.test_router.core.registry_helpers import MakeCore, typed, unique_id +from tests.unit.test_router.core.scripted_host import completion_response, deployment + + +def _usage_core(make_core: MakeCore) -> RouterCore: + """gpt: a (tpm 100, rpm 10) and b (top-level tpm 1000, model_info tpm 50 / rpm 5); io: i (itpm 300, otpm 40).""" + return make_core( + model_list=[ + deployment("gpt", "a", model="openai/gpt-4o", tpm=100, rpm=10), + { + "model_name": "gpt", + "litellm_params": {"model": "anthropic/claude-3-haiku-20240307", "api_key": "k"}, + "model_info": {"id": "b", "tpm": 50, "rpm": 5, "description": "desc b"}, + "tpm": 1000, + }, + deployment("io", "i", model="openai/gpt-4o", itpm=300, otpm=40), + deployment("plain", "p", model="openai/gpt-4o"), + ], + model_group_alias={"hid": {"model": "gpt", "hidden": True}, "al": "gpt"}, + ) + + +def _minute() -> str: + return get_utc_datetime().strftime("%H-%M") + + +async def _count(core: RouterCore, kind: RouterCacheEnum, deployment_id: str, model: str, value: int) -> None: + key = kind.value.format(id=deployment_id, model=model, current_minute=_minute()) + await core.health().store.async_set_cache(key, value) + + +# -- model group info ----------------------------------------------------------------------------- + + +def test_group_info_merges_providers_limits_and_description_across_deployments(make_core: MakeCore) -> None: + info = _usage_core(make_core).registry().group("gpt") + + assert info is not None + assert info.model_group == "gpt" + assert info.providers == ["openai", "anthropic"] + assert info.description == "desc b" + assert info.mode == "chat" + builtin = litellm.get_model_info("openai/gpt-4o") + assert info.max_input_tokens == builtin["max_input_tokens"] + assert info.input_cost_per_token == builtin["input_cost_per_token"] + + +def test_group_info_sums_tpm_and_rpm_with_top_level_values_before_model_info(make_core: MakeCore) -> None: + info = _usage_core(make_core).registry().group("gpt") + + assert info is not None + assert (info.tpm, info.rpm) == (1100, 15) + + +def test_group_info_sums_itpm_and_otpm(make_core: MakeCore) -> None: + info = _usage_core(make_core).registry().group("io") + + assert info is not None + assert (info.itpm, info.otpm, info.tpm, info.rpm) == (300, 40, None, None) + + +def test_group_info_through_an_alias_is_reported_under_the_alias(make_core: MakeCore) -> None: + info = _usage_core(make_core).registry().group("al") + + assert info is not None + assert (info.model_group, info.tpm) == ("al", 1100) + + +def test_group_info_of_a_hidden_alias_needs_include_hidden(make_core: MakeCore) -> None: + registry = _usage_core(make_core).registry() + + assert registry.group("hid") is None + shown = registry.group("hid", include_hidden=True) + assert shown is not None + assert shown.model_group == "hid" + + +def test_group_info_of_an_unknown_name_is_none(make_core: MakeCore) -> None: + assert _usage_core(make_core).registry().group("nope") is None + + +def test_group_info_of_a_name_a_wildcard_serves_reads_the_requested_model(make_core: MakeCore) -> None: + core = make_core(model_list=[deployment("openai/*", "w", model="openai/*")]) + + info = core.registry().group("openai/gpt-4o-mini") + + assert info is not None + assert (info.model_group, info.providers) == ("openai/gpt-4o-mini", ["openai"]) + assert info.max_input_tokens == litellm.get_model_info("openai/gpt-4o-mini")["max_input_tokens"] + + +def test_group_info_of_an_unmapped_model_uses_its_configured_mode_and_costs(make_core: MakeCore) -> None: + core = make_core( + model_list=[ + deployment( + "unm", + unique_id(), + model="openai/totally-unmapped-registry-model", + model_info={"mode": "embedding", "input_cost_per_token": "0.5"}, + ) + ] + ) + + info = core.registry().group("unm") + + assert info is not None + assert (info.mode, info.input_cost_per_token, info.providers) == ("embedding", 0.5, ["openai"]) + assert info.max_input_tokens is None + + +def test_group_info_reports_configurable_clientside_auth_params(make_core: MakeCore) -> None: + core = make_core(model_list=[deployment("gpt", "a", configurable_clientside_auth_params=["api_base"])]) + + info = core.registry().group("gpt") + + assert info is not None + assert info.configurable_clientside_auth_params == ["api_base"] + + +def test_cached_group_info_equals_group_info(make_core: MakeCore) -> None: + registry = _usage_core(make_core).registry() + + assert registry.group("gpt", field="cached_info") == registry.group("gpt") + + +def test_cached_group_info_follows_a_deployment_added_later(make_core: MakeCore) -> None: + core = _usage_core(make_core) + before = core.registry().group("gpt", field="cached_info") + assert before is not None + + core.add_deployment(typed("gpt", "c", tpm=1)) # pyright: ignore[reportAttributeAccessIssue] # public on the core + + after = core.registry().group("gpt", field="cached_info") + assert after is not None + assert (before.tpm, after.tpm) == (1100, 1101) + + +def test_an_unknown_group_field_raises(make_core: MakeCore) -> None: + with pytest.raises(ValueError, match="unknown group field 'bogus'"): + _usage_core(make_core).registry().group("gpt", field="bogus") # pyright: ignore[reportCallIssue, reportArgumentType] # an invalid field on purpose + + +# -- configured values of a group ----------------------------------------------------------------- + + +def _configured_core(make_core: MakeCore) -> RouterCore: + return make_core( + model_list=[ + deployment( + "cfg", + "c1", + model="openai/gpt-4o", + model_info={ + "max_input_tokens": 1000, + "mode": "embedding", + "display_name": "Nice", + "service_tiers": ["flex"], + }, + ), + deployment( + "cfg", "c2", model="openai/gpt-4o-mini", model_info={"max_input_tokens": 5000, "max_output_tokens": 77} + ), + deployment("blank", "b1", model_info={"mode": " ", "display_name": ""}), + deployment("bad", "x1", model_info={"max_input_tokens": "lots"}), + deployment("openai/*", "w", model="openai/*"), + ], + model_group_alias={"cfg-alias": "cfg"}, + ) + + +def test_token_limits_come_from_the_first_deployment_of_the_group(make_core: MakeCore) -> None: + assert _configured_core(make_core).registry().group("cfg", field="token_limits") == (1000, None) + + +@pytest.mark.parametrize("name", ["nope", "openai/gpt-x", "cfg-alias"]) +def test_token_limits_of_an_unknown_wildcard_or_alias_name_are_none(make_core: MakeCore, name: str) -> None: + assert _configured_core(make_core).registry().group(name, field="token_limits") == (None, None) + + +def test_a_malformed_token_limit_reads_as_absent(make_core: MakeCore) -> None: + assert _configured_core(make_core).registry().group("bad", field="token_limits") == (None, None) + + +@pytest.mark.parametrize( + ("name", "mode", "display_name"), + [("cfg", "embedding", "Nice"), ("blank", None, None), ("nope", None, None), ("cfg-alias", None, None)], +) +def test_mode_and_display_name_are_the_first_deployments_non_blank_values( + make_core: MakeCore, name: str, mode: str | None, display_name: str | None +) -> None: + registry = _configured_core(make_core).registry() + + assert registry.group(name, field="mode") == mode + assert registry.group(name, field="display_name") == display_name + + +def test_listing_info_reports_the_widest_limits_and_each_backend_model_once(make_core: MakeCore) -> None: + info = _configured_core(make_core).registry().group("cfg", field="listing_info") + + assert info is not None + assert info.cost_map_keys == ("openai/gpt-4o", "openai/gpt-4o-mini") + assert (info.max_input_tokens, info.max_output_tokens) == (5000, 77) + + +@pytest.mark.parametrize("name", ["nope", "openai/gpt-x", "cfg-alias"]) +def test_listing_info_of_an_unknown_wildcard_or_alias_name_is_none(make_core: MakeCore, name: str) -> None: + assert _configured_core(make_core).registry().group(name, field="listing_info") is None + + +def test_service_tiers_are_each_routable_deployments_value_in_order(make_core: MakeCore) -> None: + registry = _configured_core(make_core).registry() + + assert registry.group("cfg", field="service_tiers") == (["flex"], None) + assert registry.group("cfg-alias", field="service_tiers") == (["flex"], None) + assert registry.group("nope", field="service_tiers") == () + + +def test_service_tiers_leave_out_blocked_deployments_and_other_teams(make_core: MakeCore) -> None: + core = make_core( + model_list=[ + deployment("st", "s1", model_info={"service_tiers": ["a"], "blocked": True}), + deployment("st", "s2", model_info={"service_tiers": ["b"]}), + deployment("st", "s3", model_info={"service_tiers": ["c"], "team_id": "TA"}), + ] + ) + + assert core.registry().group("st", field="service_tiers") == (["b"],) + assert core.registry().group("st", field="service_tiers", team_id="TA") == (["b"], ["c"]) + assert core.registry().group("st", field="service_tiers", team_id="TB") == (["b"],) + + +def test_max_input_tokens_are_each_routable_deployments_own_value_in_order(make_core: MakeCore) -> None: + registry = _configured_core(make_core).registry() + + assert registry.group("cfg", field="max_input_tokens") == (1000, 5000) + assert registry.group("cfg-alias", field="max_input_tokens") == (1000, 5000) + assert registry.group("bad", field="max_input_tokens") == (None,) + assert registry.group("blank", field="max_input_tokens") == (None,) + assert registry.group("nope", field="max_input_tokens") == () + + +def test_max_input_tokens_leave_out_blocked_deployments_and_other_teams(make_core: MakeCore) -> None: + core = make_core( + model_list=[ + deployment("mt", "m1", model_info={"max_input_tokens": 1, "blocked": True}), + deployment("mt", "m2", model_info={"max_input_tokens": 2}), + deployment("mt", "m3", model_info={"max_input_tokens": 3, "team_id": "TA"}), + ] + ) + + assert core.registry().group("mt", field="max_input_tokens") == (2,) + assert core.registry().group("mt", field="max_input_tokens", team_id="TA") == (2, 3) + + +def test_max_input_tokens_skip_a_value_a_saved_model_info_copied_from_the_cost_map(make_core: MakeCore) -> None: + catalog_limit = GetModelCostMap.loaded_model_cost_map()["gpt-4o"]["max_input_tokens"] + core = make_core( + model_list=[ + deployment( + "echo", "e1", model="openai/gpt-4o", model_info={"key": "gpt-4o", "max_input_tokens": catalog_limit} + ), + deployment( + "echo", "e2", model="openai/gpt-4o", model_info={"key": "gpt-4o", "max_input_tokens": catalog_limit - 1} + ), + ] + ) + + assert core.registry().group("echo", field="max_input_tokens") == (None, catalog_limit - 1) + + +# -- deployment model info -------------------------------------------------------------------------- + + +def test_deployment_model_info_merges_custom_pricing_over_the_builtin_entry(make_core: MakeCore) -> None: + model_id = unique_id() + core = make_core(model_list=[deployment("gpt", model_id, model="openai/gpt-4o", input_cost_per_token=1e-3)]) + + info = core.registry().deployment_model_info(model_id, "openai/gpt-4o") + + assert info is not None + builtin = litellm.get_model_info("openai/gpt-4o") + assert info["input_cost_per_token"] == 1e-3 + assert info["max_input_tokens"] == builtin["max_input_tokens"] + + +def test_deployment_model_info_without_custom_pricing_is_the_builtin_entry(make_core: MakeCore) -> None: + model_id = unique_id() + core = make_core(model_list=[deployment("gpt", model_id, model="openai/gpt-4o")]) + + info = core.registry().deployment_model_info(model_id, "openai/gpt-4o") + + assert info is not None + builtin = litellm.get_model_info("openai/gpt-4o") + assert (info["input_cost_per_token"], info["output_cost_per_token"]) == ( + builtin["input_cost_per_token"], + builtin["output_cost_per_token"], + ) + + +def test_deployment_model_info_of_an_unmapped_model_is_its_custom_entry(make_core: MakeCore) -> None: + model_id = unique_id() + core = make_core( + model_list=[ + deployment( + "unm", + model_id, + model="openai/unmapped-registry-thing", + input_cost_per_token=5e-6, + model_info={"max_input_tokens": 1000}, + ) + ] + ) + + info = core.registry().deployment_model_info(model_id, "openai/unmapped-registry-thing") + + assert info is not None + assert (info["input_cost_per_token"], info["max_input_tokens"]) == (5e-6, 1000) + + +def test_deployment_model_info_of_an_unknown_id_and_model_is_none(make_core: MakeCore) -> None: + core = make_core(model_list=[deployment("gpt", "a")]) + + assert core.registry().deployment_model_info("missing-id", "nope/nothing-at-all") is None + + +def test_cached_deployment_model_info_equals_the_uncached_one(make_core: MakeCore) -> None: + model_id = unique_id() + core = make_core(model_list=[deployment("gpt", model_id, model="openai/gpt-4o", input_cost_per_token=1e-3)]) + + assert core.registry().deployment_model_info(model_id, "openai/gpt-4o", cached=True) == ( + core.registry().deployment_model_info(model_id, "openai/gpt-4o") + ) + + +def test_discovered_model_info_is_empty_without_a_refresh(make_core: MakeCore) -> None: + core = make_core(model_list=[deployment("gpt", "a")]) + + assert dict(core.registry().discovered_model_info("a")) == {} + assert dict(core.registry().discovered_model_info(None)) == {} + + +# -- router model info ------------------------------------------------------------------------------ + + +def test_router_model_info_by_id_is_the_builtin_entry_with_the_deployments_model_info_on_top( + make_core: MakeCore, +) -> None: + model_id = unique_id() + core = make_core( + model_list=[ + deployment( + "gpt", model_id, model="openai/gpt-4o", input_cost_per_token=1e-3, model_info={"max_input_tokens": 99} + ) + ] + ) + + info = core.registry().router_model_info(None, "gpt", id=model_id) + + builtin = litellm.get_model_info("openai/gpt-4o") + assert info["key"] == builtin["key"] + assert info["max_input_tokens"] == 99 + assert info["max_output_tokens"] == builtin["max_output_tokens"] + + +def test_router_model_info_of_a_passed_deployment_dict(make_core: MakeCore) -> None: + core = make_core(model_list=[deployment("gpt", "a")]) + + info = core.registry().router_model_info(deployment("x", "x", model="openai/gpt-4o-mini"), "x") + + assert info["key"] == litellm.get_model_info("openai/gpt-4o-mini")["key"] + + +def test_router_model_info_reads_the_base_model(make_core: MakeCore) -> None: + core = make_core( + model_list=[ + deployment( + "az", "z", model="azure/my-deployment", api_base="https://x", model_info={"base_model": "azure/gpt-4o"} + ) + ] + ) + + info = core.registry().router_model_info(None, "az", id="z") + + assert info["max_input_tokens"] == litellm.get_model_info("azure/gpt-4o")["max_input_tokens"] + + +def test_router_model_info_of_a_wildcard_deployment_reads_the_requested_model(make_core: MakeCore) -> None: + core = make_core(model_list=[deployment("openai/*", "w", model="openai/*")]) + + info = core.registry().router_model_info(None, "openai/gpt-4o-mini", id="w") + + assert info["key"] == litellm.get_model_info("openai/gpt-4o-mini")["key"] + + +def test_router_model_info_for_an_unknown_id_and_no_deployment_raises(make_core: MakeCore) -> None: + core = make_core(model_list=[deployment("gpt", "a")]) + + with pytest.raises(ValueError, match="Deployment not found"): + core.registry().router_model_info(None, "gpt", id="missing") + + +def test_router_model_info_of_a_deployment_without_litellm_params_raises(make_core: MakeCore) -> None: + core = make_core(model_list=[deployment("gpt", "a")]) + + with pytest.raises(ValueError, match="Deployment missing valid litellm_params"): + core.registry().router_model_info({"model_name": "x", "litellm_params": {}, "model_info": {"id": "q"}}, "x") + + +def test_router_model_info_of_an_unmapped_model_is_the_zeroed_stub_construction_registered(make_core: MakeCore) -> None: + """Construction registers every backend model in the cost map, so an unmapped one reads as free, not missing.""" + core = make_core(model_list=[deployment("gpt", "a", model="openai/not-a-known-registry-model")]) + + info = core.registry().router_model_info(None, "gpt", id="a") + + assert info["key"] == "openai/not-a-known-registry-model" + assert (info["input_cost_per_token"], info["output_cost_per_token"], info["max_input_tokens"]) == (0, 0, None) + assert info["litellm_provider"] == "openai" + + +# -- usage this minute ------------------------------------------------------------------------------ + + +@pytest.mark.asyncio +async def test_group_usage_is_none_until_something_is_counted(make_core: MakeCore) -> None: + core = _usage_core(make_core) + + assert await core.health().usage("gpt") == (None, None) + assert await core.health().usage("nope") == (None, None) + + +@pytest.mark.asyncio +async def test_group_usage_sums_each_deployments_counters_for_this_minute(make_core: MakeCore) -> None: + core = _usage_core(make_core) + await _count(core, RouterCacheEnum.TPM, "a", "openai/gpt-4o", 30) + await _count(core, RouterCacheEnum.TPM, "b", "anthropic/claude-3-haiku-20240307", 4) + await _count(core, RouterCacheEnum.RPM, "b", "anthropic/claude-3-haiku-20240307", 2) + + assert await core.health().usage("gpt") == (34, 2) + assert await core.health().usage("al") == (34, 2) + + +@pytest.mark.asyncio +async def test_io_token_usage_sums_input_and_output_counters(make_core: MakeCore) -> None: + core = _usage_core(make_core) + assert await core.get_model_group_io_token_usage("io") == (None, None) # pyright: ignore[reportAttributeAccessIssue] # public on the core + await _count(core, RouterCacheEnum.ITPM, "i", "openai/gpt-4o", 7) + + assert await core.get_model_group_io_token_usage("io") == (7, None) # pyright: ignore[reportAttributeAccessIssue] # public on the core + + +@pytest.mark.asyncio +async def test_remaining_usage_is_each_limit_minus_this_minutes_usage(make_core: MakeCore) -> None: + core = _usage_core(make_core) + await _count(core, RouterCacheEnum.TPM, "a", "openai/gpt-4o", 30) + await _count(core, RouterCacheEnum.RPM, "b", "anthropic/claude-3-haiku-20240307", 2) + + assert await core.health().usage("gpt", kind="remaining") == { + "x-ratelimit-remaining-tokens": 1070, + "x-ratelimit-limit-tokens": 1100, + "x-ratelimit-remaining-requests": 13, + "x-ratelimit-limit-requests": 15, + } + + +@pytest.mark.asyncio +async def test_remaining_usage_of_an_io_limited_group_reports_input_and_output_token_headers( + make_core: MakeCore, +) -> None: + core = _usage_core(make_core) + await _count(core, RouterCacheEnum.ITPM, "i", "openai/gpt-4o", 7) + + remaining = await core.health().usage("io", kind="remaining") + + assert remaining["x-ratelimit-limit-input-tokens"] == 300 + assert remaining["x-ratelimit-remaining-input-tokens"] == 293 + assert remaining["x-ratelimit-limit-output-tokens"] == 40 + assert remaining["x-ratelimit-remaining-output-tokens"] == 40 + assert "x-ratelimit-limit-tokens" not in remaining + + +@pytest.mark.asyncio +@pytest.mark.parametrize("group", ["plain", "nope"]) +async def test_remaining_usage_of_a_group_without_limits_is_empty(make_core: MakeCore, group: str) -> None: + assert await _usage_core(make_core).health().usage(group, kind="remaining") == {} + + +@pytest.mark.asyncio +async def test_an_unknown_usage_kind_raises(make_core: MakeCore) -> None: + with pytest.raises(ValueError, match="unknown usage kind 'bogus'"): + await _usage_core(make_core).health().usage("gpt", kind="bogus") # pyright: ignore[reportCallIssue, reportArgumentType] # an invalid kind on purpose + + +# -- set_response_headers --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_response_headers_name_the_group_and_its_remaining_limits(make_core: MakeCore) -> None: + core = _usage_core(make_core) + await _count(core, RouterCacheEnum.TPM, "a", "openai/gpt-4o", 30) + response = completion_response() + + returned = await core.set_response_headers(response, model_group="gpt") # pyright: ignore[reportAttributeAccessIssue] # public on the core + + assert returned is response + assert response._hidden_params["additional_headers"] == { + "x-litellm-model-group": "gpt", + "x-ratelimit-remaining-tokens": 1070, + "x-ratelimit-limit-tokens": 1100, + "x-ratelimit-remaining-requests": 15, + "x-ratelimit-limit-requests": 15, + } + + +@pytest.mark.asyncio +async def test_response_headers_without_a_group_carry_only_a_none_group(make_core: MakeCore) -> None: + returned = await _usage_core(make_core).set_response_headers(completion_response(), model_group=None) # pyright: ignore[reportAttributeAccessIssue] # public on the core + + assert returned._hidden_params["additional_headers"] == {"x-litellm-model-group": None} + + +@pytest.mark.asyncio +async def test_response_headers_of_a_group_without_limits_carry_only_the_group(make_core: MakeCore) -> None: + returned = await _usage_core(make_core).set_response_headers(completion_response(), model_group="plain") # pyright: ignore[reportAttributeAccessIssue] # public on the core + + assert returned._hidden_params["additional_headers"] == {"x-litellm-model-group": "plain"} + + +@pytest.mark.asyncio +async def test_response_headers_of_no_response_are_none(make_core: MakeCore) -> None: + assert await _usage_core(make_core).set_response_headers(None, model_group="gpt") is None # pyright: ignore[reportAttributeAccessIssue] # public on the core + + +@pytest.mark.asyncio +async def test_response_headers_stamp_the_model_id_of_the_deployment_in_the_request(make_core: MakeCore) -> None: + returned = await _usage_core(make_core).set_response_headers( # pyright: ignore[reportAttributeAccessIssue] # public on the core + completion_response(), model_group="plain", request_kwargs={"metadata": {"model_info": {"id": "p"}}} + ) + + assert returned._hidden_params["model_id"] == "p" + + +# -- budgets per deployment ------------------------------------------------------------------------- + + +def test_no_deployment_budget_means_no_budget_limiter(make_core: MakeCore) -> None: + core = make_core(model_list=[deployment("gpt", "a")]) + + assert core.registry().plugins("budget_limiter") is None + + +@pytest.mark.asyncio +async def test_adding_a_deployment_with_a_budget_installs_the_budget_limiter(make_core: MakeCore) -> None: + core = make_core(model_list=[deployment("gpt", "a")]) + + core.add_deployment( # pyright: ignore[reportAttributeAccessIssue] # public on the core + Deployment(**deployment("gpt", "budgeted", max_budget=1.0, budget_duration="1d")) # pyright: ignore[reportArgumentType] # a model_list entry + ) + + limiter = core.registry().plugins("budget_limiter") + assert limiter is not None + assert type(limiter).__name__ == "RouterBudgetLimiting" + + +def test_a_team_without_ptu_shares_has_no_ptu_ceiling(make_core: MakeCore) -> None: + core = make_core(model_list=[deployment("gpt", "a")]) + + assert core.health().team_ptu_ceiling("TA", "gpt") is None diff --git a/tests/unit/test_router/core/test_registry_characterization_writes.py b/tests/unit/test_router/core/test_registry_characterization_writes.py new file mode 100644 index 00000000000..0558a16fe9d --- /dev/null +++ b/tests/unit/test_router/core/test_registry_characterization_writes.py @@ -0,0 +1,516 @@ +""" +Registry and settings writes beyond the version bookkeeping in ``test_registry_changesets.py``: +what ``add_deployment``, ``set_model_list``, upserts and deletes leave behind in every index, how +``update_settings`` / ``ChangeSet(settings=...)`` coerce and filter values, and the ``settings()`` +and ``registry()`` views. +""" + +import warnings + +import pytest + +import litellm +from litellm.router_api import ChangeSet +from litellm.router_core import RouterCore +from litellm.types.router import Deployment, RetryPolicy +from tests.unit.test_router.core.registry_helpers import ( + MESSAGES, + MakeCore, + ids, + names_and_ids, + team_deployment, + typed, + unique_id, +) +from tests.unit.test_router.core.scripted_host import ScriptedHost, deployment + + +def _core(make_core: MakeCore, **arguments: object) -> RouterCore: # kwargs-ok: the router's constructor arguments + """gpt (a, b), emb (e), the openai/* wildcard (w), team TA's ``pub`` (t1, internal name ``int``).""" + return make_core( + model_list=[ + deployment("gpt", "a"), + deployment("gpt", "b"), + deployment("emb", "e"), + deployment("openai/*", "w", model="openai/*"), + team_deployment("int", "t1", "TA", "pub"), + ], + **arguments, + ) + + +# -- add_deployment --------------------------------------------------------------------------------- + + +def test_add_deployment_returns_the_deployment_it_serves(make_core: MakeCore) -> None: + core = _core(make_core) + + added = core.add_deployment(typed("new", "n")) # pyright: ignore[reportAttributeAccessIssue] # public on the core + + assert isinstance(added, Deployment) + assert (added.model_name, added.model_info.id, added.litellm_params.model) == ("new", "n", "openai/n") + + +def test_an_added_deployment_is_appended_and_names_a_new_group(make_core: MakeCore) -> None: + core = _core(make_core) + + core.add_deployment(typed("new", "n")) # pyright: ignore[reportAttributeAccessIssue] # public on the core + + assert core.registry().names(kind="ids") == ["a", "b", "e", "w", "t1", "n"] + assert "new" in core.registry().names(kind="group_set") + assert core.registry().index("model_name")["new"] == [5] + assert core.registry().names(kind="deployment_names")[-1] == "openai/n" + + +def test_an_added_deployment_without_an_id_gets_the_id_its_model_info_generated(make_core: MakeCore) -> None: + core = _core(make_core) + candidate = Deployment(model_name="new", litellm_params={"model": "openai/q", "api_key": "k"}) # pyright: ignore[reportArgumentType] # litellm_params as a dict + + added = core.add_deployment(candidate) # pyright: ignore[reportAttributeAccessIssue] # public on the core + + assert added is not None + assert added.model_info.id == candidate.model_info.id + assert core.registry().has_deployment(str(candidate.model_info.id)) + + +@pytest.mark.asyncio +async def test_an_added_deployment_serves_requests_for_its_group(make_core: MakeCore, host: ScriptedHost) -> None: + core = _core(make_core) + + core.add_deployment(typed("new", "n")) # pyright: ignore[reportAttributeAccessIssue] # public on the core + await core.acompletion(model="new", messages=MESSAGES) # pyright: ignore[reportAttributeAccessIssue] # an endpoint method + + assert host.deployments_called("acompletion") == ["n"] + + +def test_an_added_wildcard_deployment_serves_matching_names(make_core: MakeCore) -> None: + core = make_core(model_list=[deployment("gpt", "a")]) + + core.add_deployment(typed("anthropic/*", "x", model="anthropic/*")) # pyright: ignore[reportAttributeAccessIssue] # public on the core + + assert names_and_ids(core.registry().deployments("anthropic/claude-x")) == [("anthropic/claude-x", "x")] + assert core.settings().provider_default_deployment_ids == ["x"] + + +def test_an_added_team_deployment_is_team_indexed(make_core: MakeCore) -> None: + core = _core(make_core) + + core.add_deployment(Deployment(**team_deployment("int-2", "t2", "TB", "pub-b"))) # pyright: ignore[reportAttributeAccessIssue, reportArgumentType] # public on the core + + assert core.registry().index("team_model")[("TB", "pub-b")] == [5] + assert core.registry().names(kind="team_public") == frozenset({"pub", "pub-b"}) + assert ids(core.registry().deployments("pub-b", team_id="TB")) == ["t2"] + + +def test_adding_a_deployment_with_an_unknown_provider_raises_and_adds_nothing(make_core: MakeCore) -> None: + core = _core(make_core) + version = core.registry().version + + with pytest.raises(litellm.BadRequestError, match="LLM Provider NOT provided"): + core.add_deployment(typed("bad", "bad", model="notaprovider/x")) # pyright: ignore[reportAttributeAccessIssue] # public on the core + + assert core.registry().names(kind="ids") == ["a", "b", "e", "w", "t1"] + assert "bad" not in core.registry().names(kind="group_set") + assert core.registry().version == version + + +def test_an_added_deployment_registers_its_custom_pricing(make_core: MakeCore) -> None: + core = _core(make_core) + model_id = unique_id() + + core.add_deployment(typed("priced", model_id, model="openai/gpt-4o", input_cost_per_token=3e-3)) # pyright: ignore[reportAttributeAccessIssue] # public on the core + + assert litellm.model_cost[model_id]["input_cost_per_token"] == 3e-3 + + +# -- upsert ----------------------------------------------------------------------------------------- + + +def test_a_changed_upsert_moves_the_deployment_to_the_end_of_the_list(make_core: MakeCore) -> None: + core = _core(make_core) + + core.apply(ChangeSet(upsert=[typed("gpt", "a", api_base="https://changed.example")])) + + assert core.registry().names(kind="ids") == ["b", "e", "w", "t1", "a"] + assert core.registry().index("model_name") == {"gpt": [0, 4], "emb": [1], "openai/*": [2], "int": [3]} + assert ids(core.registry().deployments("gpt")) == ["b", "a"] + + +def test_an_upsert_that_renames_the_group_moves_the_deployment_between_groups(make_core: MakeCore) -> None: + core = _core(make_core) + + core.apply(ChangeSet(upsert=[typed("emb", "a")])) + + assert ids(core.registry().deployments("gpt")) == ["b"] + assert ids(core.registry().deployments("emb")) == ["e", "a"] + + +def test_an_upsert_that_empties_a_group_drops_the_group_name(make_core: MakeCore) -> None: + core = _core(make_core) + + core.apply(ChangeSet(upsert=[typed("emb2", "e")])) + + assert core.registry().names(kind="group_set") == {"gpt", "openai/*", "int", "emb2"} + assert core.registry().deployments("emb") == [] + + +def test_a_failing_upsert_of_an_existing_deployment_raises_and_leaves_it_removed(make_core: MakeCore) -> None: + """Odd but current: without ignore_invalid_deployments the old deployment is popped and not restored.""" + core = _core(make_core) + + with pytest.raises(litellm.BadRequestError): + core.apply(ChangeSet(upsert=[typed("gpt", "a", model="notaprovider/x")])) + + assert core.registry().names(kind="ids") == ["b", "e", "w", "t1"] + + +def test_a_failing_upsert_under_ignore_invalid_deployments_restores_the_previous_deployment( + make_core: MakeCore, +) -> None: + core = _core(make_core, ignore_invalid_deployments=True) + + result = core.apply(ChangeSet(upsert=[typed("gpt", "a", model="notaprovider/x")])) + + assert list(result.upserted) == [None] + restored = core.registry().deployment("a") + assert restored is not None + assert restored.litellm_params.model == "openai/a" + assert core.registry().names(kind="ids") == ["b", "e", "w", "t1", "a"] + + +def test_the_deprecated_upsert_deployment_warns_and_returns_the_deployment(make_core: MakeCore) -> None: + core = _core(make_core) + + with pytest.warns(DeprecationWarning, match=r"upsert_deployment is deprecated; use .*apply\(ChangeSet\(upsert="): + upserted = core.upsert_deployment(typed("gpt", "c")) # pyright: ignore[reportAttributeAccessIssue] # deprecated, still public + + assert upserted is not None + assert upserted.model_info.id == "c" + + +# -- delete ----------------------------------------------------------------------------------------- + + +def test_delete_returns_the_removed_model_list_entry(make_core: MakeCore) -> None: + core = _core(make_core) + + result = core.apply(ChangeSet(delete=["e"])) + + (removed,) = result.deleted + assert isinstance(removed, dict) + assert (removed["model_name"], removed["model_info"]["id"]) == ("emb", "e") + + +def test_deleting_the_last_deployment_of_a_group_drops_the_group(make_core: MakeCore) -> None: + core = _core(make_core) + + core.apply(ChangeSet(delete=["e"])) + + assert "emb" not in core.registry().names(kind="group_set") + assert core.registry().deployments("emb") == [] + assert core.registry().index("model_name") == {"gpt": [0, 1], "openai/*": [2], "int": [3]} + + +def test_deleting_a_deployment_reindexes_the_ones_after_it(make_core: MakeCore) -> None: + core = _core(make_core) + + core.apply(ChangeSet(delete=["a"])) + + for deployment_id in ("b", "e", "w", "t1"): + found = core.registry().deployment(deployment_id) + assert found is not None + assert found.model_info.id == deployment_id + assert core.registry().index("team_model") == {("TA", "pub"): [3]} + + +def test_deleting_a_team_deployment_drops_its_team_index_and_public_name(make_core: MakeCore) -> None: + core = _core(make_core) + + core.apply(ChangeSet(delete=["t1"])) + + assert core.registry().index("team_model") == {} + assert core.registry().names(kind="team_public") == frozenset() + assert core.registry().deployments("pub", team_id="TA") == [] + + +def test_deleting_a_wildcard_deployment_stops_it_serving_matching_names(make_core: MakeCore) -> None: + core = _core(make_core) + + core.apply(ChangeSet(delete=["w"])) + + assert core.registry().deployments("openai/gpt-x") == [] + + +def test_deleting_a_wildcard_deployment_keeps_it_among_provider_default_deployment_ids(make_core: MakeCore) -> None: + """Odd but current: the provider default id list is append-only.""" + core = _core(make_core) + + core.apply(ChangeSet(delete=["w"])) + + assert core.settings().provider_default_deployment_ids == ["w"] + + +@pytest.mark.asyncio +async def test_a_deleted_deployment_is_no_longer_called(make_core: MakeCore, host: ScriptedHost) -> None: + core = _core(make_core) + + core.apply(ChangeSet(delete=["a"])) + for _ in range(3): + await core.acompletion(model="gpt", messages=MESSAGES) # pyright: ignore[reportAttributeAccessIssue] # an endpoint method + + assert host.deployments_called("acompletion") == ["b", "b", "b"] + + +def test_the_deprecated_delete_deployment_warns_and_returns_none_for_an_unknown_id(make_core: MakeCore) -> None: + core = _core(make_core) + + with pytest.warns(DeprecationWarning, match=r"delete_deployment is deprecated; use .*apply\(ChangeSet\(delete="): + assert core.delete_deployment("missing") is None # pyright: ignore[reportAttributeAccessIssue] # deprecated, still public + + +# -- set_model_list --------------------------------------------------------------------------------- + + +def test_set_model_list_replaces_every_deployment_and_group(make_core: MakeCore) -> None: + core = _core(make_core) + + core.set_model_list([deployment("new", "n")]) # pyright: ignore[reportAttributeAccessIssue] # public on the core + + assert core.registry().names(kind="ids") == ["n"] + assert core.registry().names(kind="group_set") == {"new"} + assert core.registry().index("model_name") == {"new": [0]} + assert core.registry().index("team_model") == {} + assert core.registry().names(kind="team_public") == frozenset() + + +def test_set_model_list_keeps_the_aliases(make_core: MakeCore) -> None: + core = _core(make_core, model_group_alias={"fast": "new"}) + + core.set_model_list([deployment("new", "n")]) # pyright: ignore[reportAttributeAccessIssue] # public on the core + + assert ids(core.registry().deployments("fast")) == ["n"] + + +def test_set_model_list_leaves_the_given_list_unchanged(make_core: MakeCore) -> None: + core = _core(make_core) + given = [{"model_name": "new", "litellm_params": {"model": "openai/n", "api_key": "k"}}] + + core.set_model_list(given) # pyright: ignore[reportAttributeAccessIssue] # public on the core + + assert given == [{"model_name": "new", "litellm_params": {"model": "openai/n", "api_key": "k"}}] + + +def test_set_model_list_keeps_the_old_wildcard_routes(make_core: MakeCore) -> None: + """Odd but current: the wildcard pattern router is not reset, so a replaced wildcard still matches.""" + core = _core(make_core) + + core.set_model_list([deployment("new", "n")]) # pyright: ignore[reportAttributeAccessIssue] # public on the core + + assert names_and_ids(core.registry().deployments("openai/gpt-x")) == [("openai/gpt-x", "w")] + assert core.settings().provider_default_deployment_ids == ["w"] + + +def test_set_model_list_appends_to_the_deployment_names(make_core: MakeCore) -> None: + """Odd but current: deployment_names keep the replaced deployments' models.""" + core = make_core(model_list=[deployment("gpt", "a")]) + + core.set_model_list([deployment("new", "n")]) # pyright: ignore[reportAttributeAccessIssue] # public on the core + + assert core.registry().names(kind="deployment_names") == ["openai/a", "openai/n"] + + +@pytest.mark.asyncio +async def test_after_set_model_list_requests_route_to_the_new_deployments( + make_core: MakeCore, host: ScriptedHost +) -> None: + core = _core(make_core) + + core.set_model_list([deployment("gpt", "n")]) # pyright: ignore[reportAttributeAccessIssue] # public on the core + await core.acompletion(model="gpt", messages=MESSAGES) # pyright: ignore[reportAttributeAccessIssue] # an endpoint method + + assert host.deployments_called("acompletion") == ["n"] + + +# -- update_settings / ChangeSet(settings=...) ------------------------------------------------------ + + +@pytest.mark.parametrize( + ("name", "given", "stored"), + [ + ("timeout", "5", 5), + ("num_retries", "3", 3), + ("retry_after", 2.9, 2), + ("allowed_fails", "4", 4), + ("cooldown_time", 7.5, 7), + ], +) +def test_integer_settings_are_cast_to_int(make_core: MakeCore, name: str, given: object, stored: int) -> None: + core = _core(make_core) + + core.apply(ChangeSet(settings={name: given})) + + assert core.settings().as_dict()[name] == stored + + +def test_a_non_integer_value_for_an_integer_setting_raises(make_core: MakeCore) -> None: + core = _core(make_core) + + with pytest.raises(ValueError, match="invalid literal for int"): + core.apply(ChangeSet(settings={"timeout": "abc"})) + + +def test_a_setting_that_is_not_runtime_updatable_is_ignored(make_core: MakeCore) -> None: + core = _core(make_core) + before = dict(core.settings().as_dict()) + + core.apply(ChangeSet(settings={"bogus": 1, "model_list": []})) + + assert dict(core.settings().as_dict()) == before + assert core.registry().names(kind="ids") == ["a", "b", "e", "w", "t1"] + + +def test_a_retry_policy_dict_is_stored_as_a_retry_policy(make_core: MakeCore) -> None: + core = _core(make_core) + + core.apply(ChangeSet(settings={"retry_policy": {"RateLimitErrorRetries": 2}})) + + stored = core.settings().as_dict()["retry_policy"] + assert isinstance(stored, RetryPolicy) + assert stored.RateLimitErrorRetries == 2 + + +@pytest.mark.parametrize(("given", "stored"), [("true", True), (True, True), ("false", False), ("yes-ish", False)]) +def test_mid_stream_fallback_continuation_reads_flag_strings(make_core: MakeCore, given: object, stored: bool) -> None: + core = _core(make_core) + + core.apply(ChangeSet(settings={"enable_mid_stream_fallback_continuation": given})) + + assert core.settings().as_dict()["enable_mid_stream_fallback_continuation"] is stored + + +def test_model_group_alias_from_settings_is_visible_to_group_info(make_core: MakeCore) -> None: + core = _core(make_core) + assert core.registry().group("quick") is None + + core.apply(ChangeSet(settings={"model_group_alias": {"quick": "emb"}})) + + info = core.registry().group("quick") + assert info is not None + assert info.model_group == "quick" + + +def test_the_deprecated_update_settings_warns_and_applies(make_core: MakeCore) -> None: + core = _core(make_core) + + with pytest.warns(DeprecationWarning, match=r"update_settings is deprecated"): + core.update_settings(cooldown_time=13) # pyright: ignore[reportAttributeAccessIssue] # deprecated, still public + + assert core.settings().cooldown_time == 13 + + +# -- the settings() and registry() views ------------------------------------------------------------ + + +def test_settings_as_dict_lists_the_runtime_updatable_settings(make_core: MakeCore) -> None: + core = _core(make_core, model_group_alias={"fast": "gpt"}, num_retries=2, cooldown_time=5) + + settings = core.settings().as_dict() + + assert set(settings) == { + "routing_strategy_args", + "routing_strategy", + "allowed_fails", + "cooldown_time", + "num_retries", + "timeout", + "retry_after", + "fallbacks", + "context_window_fallbacks", + "model_group_retry_policy", + "retry_policy", + "model_group_alias", + "enable_weighted_failover", + "enable_mid_stream_fallback_continuation", + "enable_tag_filtering", + "tag_routing_prefix", + "fallback_on_unsupported_params", + "routing_groups", + } + assert (settings["routing_strategy"], settings["num_retries"], settings["cooldown_time"]) == ( + "simple-shuffle", + 2, + 5, + ) + assert settings["model_group_alias"] == {"fast": "gpt"} + assert settings["routing_groups"] == [] + + +def test_settings_defaults_on_a_fresh_core(make_core: MakeCore) -> None: + settings = _core(make_core).settings() + + assert settings.routing_strategy == "simple-shuffle" + assert settings.allowed_fails == 3 + assert settings.fallbacks is None + assert settings.enable_tag_filtering is False + assert settings.provider_default_deployment_ids == ["w"] + + +@pytest.mark.parametrize(("name", "expected"), [("cooldown_time", True), ("bogus", False), ("num_retries", False)]) +def test_settings_is_set_names_only_settings_fields(make_core: MakeCore, name: str, expected: bool) -> None: + assert _core(make_core).settings().is_set(name) is expected + + +def test_registry_view_is_reused(make_core: MakeCore) -> None: + core = _core(make_core) + + assert core.registry() is core.registry() + + +def test_registry_all_is_the_live_model_list(make_core: MakeCore) -> None: + core = _core(make_core) + listed = core.registry().all() + + core.add_deployment(typed("new", "n")) # pyright: ignore[reportAttributeAccessIssue] # public on the core + + assert ids(listed) == ["a", "b", "e", "w", "t1", "n"] + + +def test_an_unknown_index_name_raises(make_core: MakeCore) -> None: + with pytest.raises(ValueError, match="unknown index 'bogus'"): + _core(make_core).registry().index("bogus") # pyright: ignore[reportCallIssue, reportArgumentType] # an invalid name on purpose + + +# -- aapply(refresh_model_info) --------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_refreshing_model_info_with_nothing_to_discover_changes_nothing(make_core: MakeCore) -> None: + core = _core(make_core) + version = core.registry().version + + result = await core.aapply(ChangeSet(refresh_model_info=True)) # pyright: ignore[reportAttributeAccessIssue] # aapply is on every core + + assert (result.version, list(result.upserted), list(result.deleted)) == (version, [], []) + assert core.registry().version == version + assert dict(core.registry().discovered_model_info("a")) == {} + + +@pytest.mark.asyncio +async def test_aapply_applies_upserts_and_deletes_like_apply(make_core: MakeCore) -> None: + core = _core(make_core) + + result = await core.aapply(ChangeSet(upsert=[typed("gpt", "c")], delete=["e"])) # pyright: ignore[reportAttributeAccessIssue] # aapply is on every core + + assert [None if d is None else d.model_info.id for d in result.upserted] == ["c"] + assert core.registry().names(kind="ids") == ["a", "b", "w", "t1", "c"] + + +@pytest.mark.asyncio +async def test_the_deprecated_arefresh_model_info_warns(make_core: MakeCore) -> None: + core = _core(make_core) + + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter("always") + await core.arefresh_model_info() # pyright: ignore[reportAttributeAccessIssue] # deprecated, still public + + assert any("arefresh_model_info is deprecated" in str(w.message) for w in caught) diff --git a/tests/unit/test_router/core/test_request_scoped_outcomes.py b/tests/unit/test_router/core/test_request_scoped_outcomes.py new file mode 100644 index 00000000000..ebe162dc090 --- /dev/null +++ b/tests/unit/test_router/core/test_request_scoped_outcomes.py @@ -0,0 +1,466 @@ +""" +What a router records once its outcome callbacks left the process lists, on both cores: +``observe(request_scoped=True)``, the old ``serve_through_request_logging()`` and ``retire()`` take +them off, and ``observe(logging_obj)`` puts them back on one request's logging object. What the +Python router writes outside those callbacks goes on: ``make_call``'s usage, a refused pre-call +check's cooldown and a failed fallback hop's cooldown. +""" + +import asyncio +from collections.abc import Callable, Iterable, Mapping +from typing import Final + +import pytest + +import litellm +from litellm.integrations.custom_logger import CustomLogger +from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging +from litellm.router import Router +from litellm.router_core import RouterCore +from tests.unit.test_router.core.endpoints_helpers import EndpointHost +from tests.unit.test_router.core.retries_health_helpers import ( + MESSAGES, + MakeCore, + attempt, + attempt_sync, + cooled_down, + fresh, +) +from tests.unit.test_router.core.scripted_host import Call, deployment, rate_limit, server_error +from tests.unit.test_router.core.streaming_helpers import chat_stream, collect_texts, request_logging + +USAGE: Final = {"prompt_tokens": 4, "completion_tokens": 8, "total_tokens": 12} + + +class RequestLoggingHost(EndpointHost): + """``EndpointHost`` that logs on a request's logging object as litellm does: its first failure + alone runs the failure callbacks, and what ``_core_bind_to_request_logging`` put on the object + runs beside the process lists. A stream logs a failure through its own copy of the object, so + only an attempt's failure marks it.""" + + __slots__ = () + + async def _answer_async(self, op: str, args: tuple[object, ...], kwargs: Mapping[str, object]) -> object: + logging_obj: Final = kwargs.get("litellm_logging_obj") + if not isinstance(logging_obj, LiteLLMLogging): + return await super()._answer_async(op, args, kwargs) + details: Final = logging_obj.model_call_details + details["current_attempt_api_call_start_time"] = object() + try: + return await super()._answer_async(op, args, kwargs) + except Exception: + if details.get("has_logged_async_failure") is not True: + details["has_logged_async_failure"] = True + details["async_failure_logged_attempt"] = details.get("current_attempt_api_call_start_time") + raise + + async def log_async(self, call: Call, error: BaseException | None, response: object) -> None: + logging_obj: Final = call.kwargs.get("litellm_logging_obj") + if not isinstance(logging_obj, LiteLLMLogging): + await super().log_async(call, error, response) + return + if error is not None and logging_obj.model_call_details.get("has_logged_async_failure") is True: + return + await super().log_async(call, error, response) + logged: Final = self._logging_kwargs(call, error, response) + if error is not None: + for callback in _functions(logging_obj.dynamic_failure_callbacks): + _ = callback(logged, None, 0.0, 0.0) + for callback in _functions(logging_obj.dynamic_async_failure_callbacks): + await callback(logged, None, 0.0, 0.0) # pyright: ignore[reportGeneralTypeIssues] # an async callback + return + for callback in _functions(logging_obj.dynamic_async_success_callbacks): + await callback(logged, response, 0.0, 0.0) # pyright: ignore[reportGeneralTypeIssues] # an async callback + + def log_sync(self, call: Call, error: BaseException | None, response: object) -> None: + super().log_sync(call, error, response) + logging_obj: Final = call.kwargs.get("litellm_logging_obj") + if error is None or not isinstance(logging_obj, LiteLLMLogging): + return + logged: Final = self._logging_kwargs(call, error, response) + for callback in _functions(logging_obj.dynamic_failure_callbacks): + _ = callback(logged, None, 0.0, 0.0) + + +def _functions(callbacks: Iterable[object] | None) -> tuple[Callable[..., object], ...]: + """The function callbacks of a logging object's list; a ``CustomLogger`` logs on its own hooks.""" + return tuple( + callback for callback in callbacks or () if callable(callback) and not isinstance(callback, CustomLogger) + ) + + +@pytest.fixture +def host() -> RequestLoggingHost: + return RequestLoggingHost() + + +def observed(core: RouterCore) -> None: + _ = core.observe(request_scoped=True) # pyright: ignore[reportCallIssue] # every core takes it + + +def through_the_old_member(core: RouterCore) -> None: + router: Final = Router.__new__(Router) + router._attach(core) # pyright: ignore[reportPrivateUsage, reportAttributeAccessIssue, reportUnknownMemberType] # as Router.from_config attaches it + with pytest.warns(DeprecationWarning, match=r"Router\.serve_through_request_logging is deprecated"): + router.serve_through_request_logging() + + +def retired(core: RouterCore) -> None: + core.retire() + + +RELEASES: Final = pytest.mark.parametrize( + "release", [observed, through_the_old_member, retired], ids=["observe", "old-member", "retire"] +) + + +def weighted_pair(make_core: MakeCore, **limits: object) -> RouterCore: # kwargs-ok: both deployments' params + """``d1`` takes every pick while it is not cooling down, and one failure cools it.""" + return make_core( + model_list=[deployment("g", "d1", weight=1, **limits), deployment("g", "d2", weight=0, **limits)], + num_retries=0, + allowed_fails=0, + ) + + +def successes(core: RouterCore, deployment_id: str) -> object: + return core.health().store.get_cache(f"{deployment_id}:successes", local_only=True) + + +async def settle() -> None: + """Let the recording tasks a stream's end scheduled run.""" + for _ in range(20): + await asyncio.sleep(0) + + +# --- a router still on the process lists ------------------------------------------------------ + + +async def test_a_router_on_the_lists_cools_a_failing_deployment_and_moves_on( + make_core: MakeCore, host: RequestLoggingHost +) -> None: + core: Final = weighted_pair(make_core) + host.answer("openai/d1", fresh(server_error)) + + first: Final = await attempt(core, disable_fallbacks=True) + second: Final = await attempt(core, disable_fallbacks=True) + + assert isinstance(first, litellm.InternalServerError) + assert not isinstance(second, Exception) + assert host.deployments_called() == ["d1", "d2"] + assert list(cooled_down(core, ["d1", "d2"])) == ["d1"] + assert core.health().failures_this_minute("d1") == 1 + assert successes(core, "d2") == 1 + + +async def test_a_binding_changes_nothing_on_a_router_on_the_lists( + make_core: MakeCore, host: RequestLoggingHost +) -> None: + core: Final = weighted_pair(make_core) + bound: Final = request_logging() + _ = core.observe(bound) # pyright: ignore[reportCallIssue] # every core takes it + + _ = attempt_sync(core, litellm_logging_obj=bound) + + assert successes(core, "d1") == 1 + + +# --- released, with no binding ------------------------------------------------------------------ + + +@RELEASES +async def test_a_failing_deployment_never_cools_and_the_next_call_returns_to_it( + make_core: MakeCore, host: RequestLoggingHost, release: Callable[[RouterCore], None] +) -> None: + core: Final = weighted_pair(make_core) + host.answer("openai/d1", fresh(server_error)) + release(core) + + first: Final = await attempt(core, disable_fallbacks=True) + second: Final = await attempt(core, disable_fallbacks=True) + + assert isinstance(first, litellm.InternalServerError) + assert isinstance(second, litellm.InternalServerError) + assert host.deployments_called() == ["d1", "d1"] + assert cooled_down(core, ["d1", "d2"]) == {} + assert core.health().failures_this_minute("d1") == 0 + assert await core.health().usage("g") == (None, None) + + +@RELEASES +async def test_a_retry_goes_back_to_the_failing_deployment( + make_core: MakeCore, host: RequestLoggingHost, release: Callable[[RouterCore], None] +) -> None: + core: Final = make_core( + model_list=[deployment("g", "d1", weight=1), deployment("g", "d2", weight=0)], num_retries=1 + ) + host.answer("openai/d1", fresh(rate_limit)) + release(core) + + answered: Final = await attempt(core) + + assert isinstance(answered, litellm.RateLimitError) + assert host.deployments_called() == ["d1", "d1"] + assert cooled_down(core, ["d1", "d2"]) == {} + + +@RELEASES +def test_a_sync_failure_records_nothing( + make_core: MakeCore, host: RequestLoggingHost, release: Callable[[RouterCore], None] +) -> None: + core: Final = weighted_pair(make_core) + host.answer("openai/d1", fresh(server_error)) + release(core) + + answered: Final = attempt_sync(core) + + assert isinstance(answered, litellm.InternalServerError) + assert cooled_down(core, ["d1", "d2"]) == {} + assert core.health().failures_this_minute("d1") == 0 + + +@RELEASES +async def test_a_success_counts_its_usage_and_no_success( + make_core: MakeCore, host: RequestLoggingHost, release: Callable[[RouterCore], None] +) -> None: + core: Final = weighted_pair(make_core, rpm=100) + release(core) + + answered: Final = await attempt(core) + + assert not isinstance(answered, Exception) + assert successes(core, "d1") is None + assert await core.health().usage("g") == (10, 1) + + +@RELEASES +async def test_a_sync_success_counts_its_usage_and_no_success( + make_core: MakeCore, host: RequestLoggingHost, release: Callable[[RouterCore], None] +) -> None: + core: Final = weighted_pair(make_core, rpm=100) + release(core) + + answered: Final = attempt_sync(core) + + assert not isinstance(answered, Exception) + assert successes(core, "d1") is None + assert await core.health().usage("g") == (10, 1) + + +async def test_an_endpoint_outside_make_call_counts_no_usage(make_core: MakeCore, host: RequestLoggingHost) -> None: + """``amoderation`` makes one attempt, outside ``make_call``: only the success callback counts it.""" + core: Final = make_core(model_list=[deployment("g", "d1", rpm=100)], num_retries=0) + observed(core) + + _ = await core.amoderation(model="g", input="x") # pyright: ignore[reportAttributeAccessIssue] # every core serves it + + assert await core.health().usage("g") == (None, None) + + +@RELEASES +async def test_a_stream_counts_its_request_when_it_opens_and_nothing_when_it_ends( + make_core: MakeCore, host: RequestLoggingHost, release: Callable[[RouterCore], None] +) -> None: + host.answer("openai/d1", chat_stream("hello", usage=USAGE, logged=True)) + core: Final = weighted_pair(make_core, rpm=100) + release(core) + + stream: Final = await core.acompletion(model="g", messages=MESSAGES, stream=True) # pyright: ignore[reportAttributeAccessIssue] # every core serves acompletion + _ = await collect_texts(stream) + await settle() + + assert successes(core, "d1") is None + assert await core.health().usage("g") == (0, 1) + + +@RELEASES +async def test_a_stream_failing_after_content_records_nothing( + make_core: MakeCore, host: RequestLoggingHost, release: Callable[[RouterCore], None] +) -> None: + host.answer("openai/d1", chat_stream("partial", error=rate_limit(message="mid stream"), logged=True)) + core: Final = weighted_pair(make_core) + release(core) + + stream: Final = await core.acompletion(model="g", messages=MESSAGES, stream=True) # pyright: ignore[reportAttributeAccessIssue] # every core serves acompletion + with pytest.raises(litellm.RateLimitError, match="mid stream"): + _ = await collect_texts(stream) + await settle() + + assert cooled_down(core, ["d1", "d2"]) == {} + assert core.health().failures_this_minute("d1") == 0 + + +# --- what goes on outside the callbacks ------------------------------------------------------------ + + +class RefusesEveryDeployment(CustomLogger): + async def async_pre_call_check(self, deployment: object, parent_otel_span: object) -> None: # pyright: ignore[reportIncompatibleMethodOverride] # any deployment + raise rate_limit() + + +@RELEASES +async def test_a_refused_pre_call_check_still_cools_the_deployment( + make_core: MakeCore, host: RequestLoggingHost, release: Callable[[RouterCore], None] +) -> None: + core: Final = weighted_pair(make_core) + release(core) + host.callback_list("callbacks").append(RefusesEveryDeployment()) + + _ = await attempt(core, disable_fallbacks=True) + + assert list(cooled_down(core, ["d1", "d2"])) == ["d1"] + assert core.health().failures_this_minute("d1") == 0 + + +async def test_a_failed_fallback_hop_on_a_shared_logging_object_is_still_counted( + make_core: MakeCore, host: RequestLoggingHost +) -> None: + """The first failure logs through no callback, yet blocks the object for the hops, and a blocked + hop records through ``_trigger_cooldown_for_failed_deployment``.""" + core: Final = make_core( + model_list=[deployment("a", "a"), deployment("b", "b"), deployment("c", "c")], + fallbacks=[{"a": ["b"]}, {"b": ["c"]}], + num_retries=0, + ) + observed(core) + host.answer("*", fresh(server_error)) + + _ = await attempt(core, "a", litellm_logging_obj=request_logging()) + + assert core.health().failures_this_minute("a") == 0 + assert core.health().failures_this_minute("b") == 1 + assert core.health().failures_this_minute("c") == 1 + + +async def test_failed_fallback_hops_with_no_logging_object_are_not_counted( + make_core: MakeCore, host: RequestLoggingHost +) -> None: + core: Final = make_core( + model_list=[deployment("a", "a"), deployment("b", "b")], fallbacks=[{"a": ["b"]}], num_retries=0 + ) + observed(core) + host.answer("*", fresh(server_error)) + + _ = await attempt(core, "a") + + assert core.health().failures_this_minute("a") == 0 + assert core.health().failures_this_minute("b") == 0 + + +# --- a binding, for its request alone ----------------------------------------------------------- + + +@pytest.mark.parametrize("release", [observed, retired], ids=["observe", "retire"]) +async def test_a_binding_records_a_failure_for_that_request_only( + make_core: MakeCore, host: RequestLoggingHost, release: Callable[[RouterCore], None] +) -> None: + core: Final = weighted_pair(make_core) + host.answer("openai/d1", fresh(server_error)) + release(core) + + _ = await attempt(core, disable_fallbacks=True, litellm_logging_obj=request_logging()) + assert cooled_down(core, ["d1", "d2"]) == {} + assert core.health().failures_this_minute("d1") == 0 + + bound: Final = request_logging() + _ = core.observe(bound) # pyright: ignore[reportCallIssue] # every core takes it + _ = await attempt(core, disable_fallbacks=True, litellm_logging_obj=bound) + + assert list(cooled_down(core, ["d1", "d2"])) == ["d1"] + assert core.health().failures_this_minute("d1") == 1 + assert await core.health().usage("g") == (None, 1) + + +async def test_a_binding_records_an_async_success(make_core: MakeCore, host: RequestLoggingHost) -> None: + core: Final = weighted_pair(make_core) + observed(core) + bound: Final = request_logging() + _ = core.observe(bound) # pyright: ignore[reportCallIssue] # every core takes it + + _ = await attempt(core, litellm_logging_obj=bound) + + assert successes(core, "d1") == 1 + + +async def test_a_binding_leaves_a_sync_success_uncounted(make_core: MakeCore, host: RequestLoggingHost) -> None: + """``_core_bind_to_request_logging`` does not bind ``sync_deployment_callback_on_success``.""" + core: Final = weighted_pair(make_core, rpm=100) + observed(core) + bound: Final = request_logging() + _ = core.observe(bound) # pyright: ignore[reportCallIssue] # every core takes it + + _ = attempt_sync(core, litellm_logging_obj=bound) + + assert successes(core, "d1") is None + assert await core.health().usage("g") == (10, 1) + + +async def test_a_binding_records_a_sync_failure_with_no_request(make_core: MakeCore, host: RequestLoggingHost) -> None: + core: Final = weighted_pair(make_core) + host.answer("openai/d1", fresh(server_error)) + observed(core) + bound: Final = request_logging() + _ = core.observe(bound) # pyright: ignore[reportCallIssue] # every core takes it + + _ = attempt_sync(core, litellm_logging_obj=bound) + + assert list(cooled_down(core, ["d1", "d2"])) == ["d1"] + assert core.health().failures_this_minute("d1") == 1 + assert await core.health().usage("g") == (None, None) + + +async def test_a_binding_records_a_stream_that_ends(make_core: MakeCore, host: RequestLoggingHost) -> None: + host.answer("openai/d1", chat_stream("hello", usage=USAGE, logged=True)) + core: Final = weighted_pair(make_core, rpm=100) + observed(core) + bound: Final = request_logging() + _ = core.observe(bound) # pyright: ignore[reportCallIssue] # every core takes it + + stream: Final = await core.acompletion( # pyright: ignore[reportAttributeAccessIssue] # every core serves acompletion + model="g", messages=MESSAGES, stream=True, litellm_logging_obj=bound + ) + _ = await collect_texts(stream) + await settle() + + assert successes(core, "d1") == 1 + assert await core.health().usage("g") == (12, 1) + + +async def test_a_binding_records_a_stream_failing_after_content(make_core: MakeCore, host: RequestLoggingHost) -> None: + host.answer("openai/d1", chat_stream("partial", error=rate_limit(message="mid stream"), logged=True)) + core: Final = weighted_pair(make_core) + observed(core) + bound: Final = request_logging() + _ = core.observe(bound) # pyright: ignore[reportCallIssue] # every core takes it + + stream: Final = await core.acompletion( # pyright: ignore[reportAttributeAccessIssue] # every core serves acompletion + model="g", messages=MESSAGES, stream=True, litellm_logging_obj=bound + ) + with pytest.raises(litellm.RateLimitError, match="mid stream"): + _ = await collect_texts(stream) + await settle() + + assert list(cooled_down(core, ["d1", "d2"])) == ["d1"] + assert core.health().failures_this_minute("d1") == 1 + # No limit on the deployment, so the open counted nothing: the failure is the one request. + assert await core.health().usage("g") == (None, 1) + + +# --- retire ---------------------------------------------------------------------------------------- + + +@pytest.mark.parametrize("strategy", ["latency-based-routing", "least-busy", "usage-based-routing-v2"]) +async def test_a_retired_router_still_picks_with_its_strategy( + make_core: MakeCore, host: RequestLoggingHost, strategy: str +) -> None: + core: Final = make_core( + model_list=[deployment("g", "d1", rpm=100), deployment("g", "d2", rpm=100)], + routing_strategy=strategy, + num_retries=0, + ) + core.retire() + + answered: Final = await attempt(core) + + assert not isinstance(answered, Exception) + assert len(host.deployments_called()) == 1 diff --git a/tests/unit/test_router/core/test_retries_health_characterization_cooldowns.py b/tests/unit/test_router/core/test_retries_health_characterization_cooldowns.py new file mode 100644 index 00000000000..4aeb11c58ec --- /dev/null +++ b/tests/unit/test_router/core/test_retries_health_characterization_cooldowns.py @@ -0,0 +1,845 @@ +""" +Characterization of deployment health in a router core: when a failure cools a deployment down (by status, +group size, allowed_fails and allowed_fails_policy at router and deployment level), what the failure and +success callbacks count, blocked deployments and health-check-driven filtering. +""" + +import time +from typing import Final + +import pytest + +import litellm +from litellm.constants import DEFAULT_COOLDOWN_TIME_SECONDS, DEFAULT_FAILURE_THRESHOLD_MINIMUM_REQUESTS +from litellm.router_utils.cooldown_handlers import mark_advisor_orchestration_failure +from litellm.types.router import AllowedFailsPolicy, Deployment +from tests.unit.test_router.core.retries_health_helpers import ( + BackgroundSuccessLoggingHost, + ErrorBuilder, + MakeCore, + api_connection, + api_status, + attempt, + cooled_down, + fresh, + not_found, + permission_denied, + service_unavailable, +) +from tests.unit.test_router.core.scripted_host import ( + ScriptedHost, + auth_error, + bad_request, + deployment, + rate_limit, + server_error, + timeout, +) + + +def group( + *ids: str, **params: object +) -> list[dict[str, object]]: # kwargs-ok: litellm_params # mutable-ok: a model_list + return [deployment("g", deployment_id, **params) for deployment_id in ids] + + +def unhealthy(*ids: str, age: float = 0.0) -> dict[str, dict[str, object]]: # mutable-ok: health states + return {i: {"is_healthy": False, "timestamp": time.time() - age, "reason": "probe failed"} for i in ids} + + +def failure_kwargs( + model_info: object, exception: Exception, **extra: object +) -> dict[str, object]: # kwargs-ok: logging kwargs # mutable-ok: litellm's logging kwargs + litellm_params: Final = {"metadata": {"model_group": "g"}, "model_info": model_info} + litellm_params.update(extra.pop("litellm_params", {})) # pyright: ignore[reportCallIssue, reportArgumentType] # test extras + return {"exception": exception, "litellm_params": litellm_params, **extra} + + +# --- which failures cool a deployment down ------------------------------------------------------- + + +@pytest.mark.parametrize("builder", [auth_error, not_found]) +@pytest.mark.asyncio +async def test_a_lone_deployment_cools_down_on_a_401_or_404( + make_core: MakeCore, host: ScriptedHost, builder: ErrorBuilder +) -> None: + core: Final = make_core(model_list=group("a"), num_retries=0) + host.answer("openai/a", fresh(builder)) + + first: Final = await attempt(core) + second: Final = await attempt(core) + + assert type(first) is type(builder()) + assert set(cooled_down(core, ["a"])) == {"a"} + assert isinstance(second, litellm.RouterRateLimitError) + assert host.deployments_called("acompletion") == ["a"] + + +@pytest.mark.asyncio +async def test_a_404_cools_a_group_deployment_down_at_once(make_core: MakeCore, host: ScriptedHost) -> None: + core: Final = make_core(model_list=group("a", "b"), num_retries=0) + host.answer("*", fresh(not_found)) + + await attempt(core) + + assert list(cooled_down(core, ["a", "b"])) == host.deployments_called("acompletion") + + +@pytest.mark.parametrize("builder", [permission_denied, bad_request, api_status(422), api_status(409)]) +@pytest.mark.asyncio +async def test_other_4xx_errors_never_cool_down(make_core: MakeCore, host: ScriptedHost, builder: ErrorBuilder) -> None: + core: Final = make_core(model_list=group("a", "b"), num_retries=0, allowed_fails=0) + host.answer("*", fresh(builder)) + + for _ in range(4): + await attempt(core) + + assert cooled_down(core, ["a", "b"]) == {} + + +@pytest.mark.parametrize("builder", [server_error, timeout, service_unavailable]) +@pytest.mark.asyncio +async def test_by_default_a_group_deployment_cools_down_on_its_fifth_failure_this_minute( + make_core: MakeCore, host: ScriptedHost, builder: ErrorBuilder +) -> None: + core: Final = make_core(model_list=group("a", "b"), num_retries=0) + host.answer("*", fresh(builder)) + + while not cooled_down(core, ["a", "b"]): + await attempt(core) + + called: Final = host.deployments_called("acompletion") + (cooled,) = cooled_down(core, ["a", "b"]) + assert called[-1] == cooled + assert called.count(cooled) == DEFAULT_FAILURE_THRESHOLD_MINIMUM_REQUESTS + + +@pytest.mark.asyncio +async def test_successes_this_minute_keep_the_failure_rate_under_the_default_threshold(make_core: MakeCore) -> None: + core: Final = make_core(model_list=group("a", "b")) + for _ in range(5): + core.sync_deployment_callback_on_success( # pyright: ignore[reportAttributeAccessIssue] # PythonCore's public member + {"litellm_params": {"metadata": {"model_group": "g"}, "model_info": {"id": "a"}}}, None, 0.0, 0.0 + ) + + outcomes: Final = [core.health().record_failure("a", server_error(), 500) for _ in range(6)] + + assert outcomes == [False] * 5 + [True] + + +@pytest.mark.parametrize("allowed_fails", [3, 1]) +@pytest.mark.asyncio +async def test_an_explicit_allowed_fails_counts_failures_even_at_the_default_value( + make_core: MakeCore, allowed_fails: int +) -> None: + core: Final = make_core(model_list=group("a", "b"), allowed_fails=allowed_fails) + + outcomes: Final = [core.health().record_failure("a", server_error(), 500) for _ in range(allowed_fails + 1)] + + assert outcomes == [False] * allowed_fails + [True] + + +@pytest.mark.asyncio +async def test_a_connection_error_never_cools_down(make_core: MakeCore, host: ScriptedHost) -> None: + core: Final = make_core(model_list=group("a", "b"), num_retries=0, allowed_fails=0) + host.answer("*", fresh(api_connection)) + + for _ in range(3): + await attempt(core) + + assert cooled_down(core, ["a", "b"]) == {} + + +@pytest.mark.parametrize(("ids", "cools"), [(("a", "b"), True), (("a",), False)], ids=["group", "lone"]) +@pytest.mark.asyncio +async def test_a_402_cools_a_group_deployment_but_not_a_lone_one( + make_core: MakeCore, host: ScriptedHost, ids: tuple[str, ...], cools: bool +) -> None: + core: Final = make_core(model_list=group(*ids), num_retries=0) + host.answer("*", fresh(api_status(402))) + + await attempt(core) + + assert bool(cooled_down(core, list(ids))) is cools + + +@pytest.mark.asyncio +async def test_a_lone_deployment_with_a_fallback_group_cools_down_under_explicit_allowed_fails( + make_core: MakeCore, host: ScriptedHost +) -> None: + core: Final = make_core( + model_list=[deployment("g", "a"), deployment("f", "z")], + num_retries=0, + allowed_fails=0, + fallbacks=[{"g": ["f"]}], + ) + host.answer("openai/a", fresh(server_error)) + + await attempt(core) + + assert set(cooled_down(core, ["a", "z"])) == {"a"} + assert host.deployments_called("acompletion") == ["a", "z"] + + +@pytest.mark.asyncio +async def test_a_lone_deployment_with_a_fallback_group_is_not_cooled_by_a_429_under_defaults( + make_core: MakeCore, host: ScriptedHost +) -> None: + core: Final = make_core( + model_list=[deployment("g", "a"), deployment("f", "z")], num_retries=0, fallbacks=[{"g": ["f"]}] + ) + host.answer("openai/a", fresh(rate_limit)) + + await attempt(core) + + assert cooled_down(core, ["a", "z"]) == {} + + +@pytest.mark.asyncio +async def test_record_failure_never_cools_down_when_cooldowns_are_disabled(make_core: MakeCore) -> None: + core: Final = make_core(model_list=group("a", "b"), allowed_fails=0, disable_cooldowns=True) + + cooled: Final = core.health().record_failure("a", rate_limit(), 429) + + assert cooled is False + assert cooled_down(core, ["a", "b"]) == {} + assert core.health().failures_this_minute("a") == 1 + + +@pytest.mark.parametrize( + ("deployment_id", "status"), [("missing", 429), ("a", ""), ("a", 400)], ids=["unknown", "empty", "400"] +) +@pytest.mark.asyncio +async def test_record_failure_without_a_cooldown(make_core: MakeCore, deployment_id: str, status: object) -> None: + core: Final = make_core(model_list=group("a", "b"), allowed_fails=0) + + assert core.health().record_failure(deployment_id, rate_limit(), status) is False # pyright: ignore[reportArgumentType] # status as given + assert core.health().failures_this_minute(deployment_id) == 1 + + +@pytest.mark.asyncio +async def test_a_router_cooldown_time_of_zero_falls_back_to_the_default( + make_core: MakeCore, host: ScriptedHost +) -> None: + core: Final = make_core(model_list=group("a", "b"), num_retries=0, cooldown_time=0) + host.answer("*", fresh(rate_limit)) + + await attempt(core) + + assert core.settings().cooldown_time == DEFAULT_COOLDOWN_TIME_SECONDS + (entry,) = cooled_down(core, ["a", "b"]).values() + assert entry["cooldown_time"] == DEFAULT_COOLDOWN_TIME_SECONDS + + +# --- allowed_fails_policy ------------------------------------------------------------------------ + + +@pytest.mark.asyncio +async def test_router_allowed_fails_policy_counts_every_error_type_in_one_counter(make_core: MakeCore) -> None: + core: Final = make_core( + model_list=group("a", "b"), + allowed_fails=5, + allowed_fails_policy=AllowedFailsPolicy(RateLimitErrorAllowedFails=1), + ) + + outcomes: Final = [ + core.health().record_failure("a", server_error(), 500), + core.health().record_failure("a", server_error(), 500), + core.health().record_failure("a", rate_limit(), 429), + ] + + assert outcomes == [False, False, True] + + +@pytest.mark.parametrize( + ("field", "error", "status"), + [ + ("AuthenticationErrorAllowedFails", auth_error, 401), + ("TimeoutErrorAllowedFails", timeout, 408), + ("RateLimitErrorAllowedFails", rate_limit, 429), + ("InternalServerErrorAllowedFails", server_error, 500), + ("ServiceUnavailableErrorAllowedFails", service_unavailable, 503), + ("NotFoundErrorAllowedFails", not_found, 404), + ], +) +@pytest.mark.asyncio +async def test_router_allowed_fails_policy_threshold_per_error_type( + make_core: MakeCore, field: str, error: ErrorBuilder, status: int +) -> None: + core: Final = make_core( + model_list=group("a", "b"), allowed_fails=10, allowed_fails_policy=AllowedFailsPolicy(**{field: 2}) + ) + + outcomes: Final = [core.health().record_failure("a", error(), status) for _ in range(3)] + + assert outcomes == [False, False, True] + + +@pytest.mark.asyncio +async def test_router_allowed_fails_policy_for_bad_requests_does_not_cool_down(make_core: MakeCore) -> None: + core: Final = make_core( + model_list=group("a", "b"), allowed_fails_policy=AllowedFailsPolicy(BadRequestErrorAllowedFails=0) + ) + + outcomes: Final = [core.health().record_failure("a", bad_request(), 400) for _ in range(3)] + + assert outcomes == [False, False, False] + + +@pytest.mark.asyncio +async def test_deployment_allowed_fails_policy_counts_each_error_type_separately(make_core: MakeCore) -> None: + core: Final = make_core( + model_list=group( + "a", + "b", + model_info={ + "allowed_fails_policy": {"RateLimitErrorAllowedFails": 1, "InternalServerErrorAllowedFails": 1} + }, + ) + ) + + outcomes: Final = [ + core.health().record_failure("a", rate_limit(), 429), + core.health().record_failure("a", server_error(), 500), + core.health().record_failure("a", rate_limit(), 429), + ] + + assert outcomes == [False, False, True] + + +@pytest.mark.asyncio +async def test_deployment_allowed_fails_policy_wins_over_the_router_policy(make_core: MakeCore) -> None: + core: Final = make_core( + model_list=[ + deployment("g", "a", model_info={"allowed_fails_policy": {"RateLimitErrorAllowedFails": 2}}), + deployment("g", "b"), + ], + allowed_fails_policy=AllowedFailsPolicy(RateLimitErrorAllowedFails=0), + ) + + a_outcomes: Final = [core.health().record_failure("a", rate_limit(), 429) for _ in range(3)] + b_outcome: Final = core.health().record_failure("b", rate_limit(), 429) + + assert a_outcomes == [False, False, True] + assert b_outcome is True + + +@pytest.mark.asyncio +async def test_deployment_policy_not_naming_the_error_falls_back_to_router_allowed_fails(make_core: MakeCore) -> None: + core: Final = make_core( + model_list=group("a", "b", model_info={"allowed_fails_policy": {"RateLimitErrorAllowedFails": 9}}) + ) + + outcomes: Final = [core.health().record_failure("a", server_error(), 500) for _ in range(4)] + + assert outcomes == [False, False, False, True] + + +@pytest.mark.parametrize(("ids", "cools"), [(("a", "b"), True), (("a",), False)], ids=["group", "lone"]) +@pytest.mark.asyncio +async def test_deployment_allowed_fails_in_model_info(make_core: MakeCore, ids: tuple[str, ...], cools: bool) -> None: + core: Final = make_core(model_list=group(*ids, model_info={"allowed_fails": 1})) + + outcomes: Final = [core.health().record_failure("a", server_error(), 500) for _ in range(2)] + + assert outcomes == [False, cools] + + +@pytest.mark.asyncio +async def test_deployment_allowed_fails_policy_cools_down_a_lone_deployment(make_core: MakeCore) -> None: + core: Final = make_core(model_list=group("a", model_info={"allowed_fails_policy": {"TimeoutErrorAllowedFails": 0}})) + + assert core.health().record_failure("a", timeout(), 408) is True + + +@pytest.mark.asyncio +async def test_deployment_cooldown_time_in_model_info_sets_the_cooldown( + make_core: MakeCore, host: ScriptedHost +) -> None: + core: Final = make_core( + model_list=group( + "a", "b", model_info={"cooldown_time": 17, "allowed_fails_policy": {"RateLimitErrorAllowedFails": 0}} + ), + num_retries=0, + ) + host.answer("*", fresh(rate_limit)) + + await attempt(core) + + (entry,) = cooled_down(core, ["a", "b"]).values() + assert entry["cooldown_time"] == 17 + + +# --- deployment_callback_on_failure ---------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_failure_callback_counts_the_failure_and_cools_down(make_core: MakeCore) -> None: + core: Final = make_core(model_list=group("a", "b"), allowed_fails=0) + + cooled: Final = core.deployment_callback_on_failure(failure_kwargs({"id": "a"}, server_error()), None, 0.0, 0.0) # pyright: ignore[reportAttributeAccessIssue] # PythonCore's public member + + assert cooled is True + assert core.health().failures_this_minute("a") == 1 + assert set(cooled_down(core, ["a", "b"])) == {"a"} + + +@pytest.mark.parametrize("model_info", [None, "a", {}, {"name": "no id"}], ids=["none", "not-a-dict", "empty", "no-id"]) +@pytest.mark.asyncio +async def test_failure_callback_without_a_deployment_id_does_nothing(make_core: MakeCore, model_info: object) -> None: + core: Final = make_core(model_list=group("a", "b"), allowed_fails=0) + + cooled: Final = core.deployment_callback_on_failure(failure_kwargs(model_info, server_error()), None, 0.0, 0.0) # pyright: ignore[reportAttributeAccessIssue] # PythonCore's public member + + assert cooled is False + assert core.health().failures_this_minute("a") == 0 + + +@pytest.mark.asyncio +async def test_failure_callback_ignores_an_advisor_orchestration_failure(make_core: MakeCore) -> None: + core: Final = make_core(model_list=group("a", "b"), allowed_fails=0) + error: Final = server_error() + mark_advisor_orchestration_failure(error) + + cooled: Final = core.deployment_callback_on_failure(failure_kwargs({"id": "a"}, error), None, 0.0, 0.0) # pyright: ignore[reportAttributeAccessIssue] # PythonCore's public member + + assert cooled is False + assert core.health().failures_this_minute("a") == 0 + + +@pytest.mark.parametrize( + "extra", + [ + { + "litellm_params": { + "metadata": {"model_group": "g", "internal_call_origin": "background_response_status_poll"} + } + }, + {"call_type": "avideo_status"}, + ], + ids=["background-poll", "video-lookup"], +) +@pytest.mark.asyncio +async def test_failure_callback_ignores_a_404_that_is_not_about_the_deployment( + make_core: MakeCore, extra: dict[str, object] +) -> None: + core: Final = make_core(model_list=group("a", "b"), allowed_fails=0) + + cooled: Final = core.deployment_callback_on_failure( # pyright: ignore[reportAttributeAccessIssue] # PythonCore's public member + failure_kwargs({"id": "a"}, not_found(), **extra), None, 0.0, 0.0 + ) + + assert cooled is False + assert cooled_down(core, ["a", "b"]) == {} + + +@pytest.mark.asyncio +async def test_failure_callback_ignores_a_408_from_the_callers_own_timeout(make_core: MakeCore) -> None: + core: Final = make_core(model_list=group("a", "b"), allowed_fails=0) + started: Final = time.time() + from datetime import datetime, timedelta + + start: Final = datetime.fromtimestamp(started) + kwargs: Final = failure_kwargs( + {"id": "a"}, + timeout(), + litellm_params={"client_side_timeout": True, "timeout": 1.0}, + start_time=start, + end_time=start + timedelta(seconds=2), + ) + + cooled: Final = core.deployment_callback_on_failure(kwargs, None, 0.0, 0.0) # pyright: ignore[reportAttributeAccessIssue] # PythonCore's public member + + assert cooled is False + assert core.health().failures_this_minute("a") == 0 + + +@pytest.mark.parametrize( + ("model_info", "litellm_params", "expected"), + [ + ({"id": "a", "cooldown_time": 3}, {"cooldown_time": 4}, 3), + ({"id": "a"}, {"cooldown_time": 4}, 4), + ({"id": "a", "cooldown_time": -1}, {}, 60), + ], + ids=["model-info-first", "litellm-params", "negative-ignored"], +) +@pytest.mark.asyncio +async def test_failure_callback_cooldown_time_resolution( + make_core: MakeCore, model_info: dict[str, object], litellm_params: dict[str, object], expected: float +) -> None: + core: Final = make_core(model_list=group("a", "b"), allowed_fails=0, cooldown_time=60) + + core.deployment_callback_on_failure( # pyright: ignore[reportAttributeAccessIssue] # PythonCore's public member + failure_kwargs(model_info, server_error(), litellm_params=litellm_params), None, 0.0, 0.0 + ) + + (entry,) = cooled_down(core, ["a", "b"]).values() + assert entry["cooldown_time"] == expected + + +# --- success callbacks and usage counters -------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_sync_success_callback_counts_a_success_and_returns_its_key(make_core: MakeCore) -> None: + core: Final = make_core(model_list=group("a", "b")) + kwargs: Final = {"litellm_params": {"metadata": {"model_group": "g"}, "model_info": {"id": "a"}}} + + key: Final = core.sync_deployment_callback_on_success(kwargs, None, 0.0, 0.0) # pyright: ignore[reportAttributeAccessIssue] # PythonCore's public member + core.sync_deployment_callback_on_success(kwargs, None, 0.0, 0.0) # pyright: ignore[reportAttributeAccessIssue] # PythonCore's public member + + assert key == "a:successes" + assert core.health().store.get_cache("a:successes", local_only=True) == 2 + + +@pytest.mark.asyncio +async def test_sync_success_callback_turns_an_integer_id_into_a_string_key(make_core: MakeCore) -> None: + core: Final = make_core(model_list=group("a")) + + key: Final = core.sync_deployment_callback_on_success( # pyright: ignore[reportAttributeAccessIssue] # PythonCore's public member + {"litellm_params": {"metadata": {"model_group": "g"}, "model_info": {"id": 7}}}, None, 0.0, 0.0 + ) + + assert key == "7:successes" + + +@pytest.mark.parametrize( + "kwargs", + [ + {"litellm_params": {"metadata": None, "model_info": {"id": "a"}}}, + {"litellm_params": {"metadata": {}, "model_info": {"id": "a"}}}, + {"litellm_params": {"metadata": {"model_group": "g"}, "model_info": None}}, + { + "call_type": "aretrieve_batch", + "litellm_params": {"metadata": {"model_group": "g"}, "model_info": {"id": "a"}}, + }, + ], + ids=["no-metadata", "no-group", "no-model-info", "batch-retrieve"], +) +@pytest.mark.asyncio +async def test_sync_success_callback_counts_nothing_without_a_group_and_id( + make_core: MakeCore, kwargs: dict[str, object] +) -> None: + core: Final = make_core(model_list=group("a")) + + assert core.sync_deployment_callback_on_success(kwargs, None, 0.0, 0.0) is None # pyright: ignore[reportAttributeAccessIssue] # PythonCore's public member + assert core.health().store.get_cache("a:successes", local_only=True) is None + + +@pytest.mark.parametrize( + "kwargs", + [ + {"call_type": "_arealtime"}, + {"call_type": "aretrieve_batch"}, + {"litellm_params": {"metadata": {"model_group": "g"}, "model_info": {"id": "a"}}}, + ], + ids=["realtime", "batch-retrieve", "no-standard-logging-object"], +) +@pytest.mark.asyncio +async def test_async_success_callback_skips_calls_it_does_not_track( + make_core: MakeCore, kwargs: dict[str, object] +) -> None: + core: Final = make_core(model_list=group("a", rpm=100)) + full: Final = { + "litellm_params": {"metadata": {"model_group": "g"}, "model_info": {"id": "a"}}, + "standard_logging_object": {"total_tokens": 10}, + **kwargs, + } + if "litellm_params" in kwargs: + del full["standard_logging_object"] + + result: Final = await core.deployment_callback_on_success(full, None, 0.0, 0.0) # pyright: ignore[reportAttributeAccessIssue] # PythonCore's public member + + assert result is None + assert core.health().store.get_cache("a:successes", local_only=True) is None + assert await core.health().usage("g") == (None, None) + + +@pytest.mark.asyncio +async def test_async_success_callback_for_an_unknown_deployment_counts_nothing(make_core: MakeCore) -> None: + core: Final = make_core(model_list=group("a", rpm=100)) + + await core.deployment_callback_on_success( # pyright: ignore[reportAttributeAccessIssue] # PythonCore's public member + { + "litellm_params": {"metadata": {"model_group": "g"}, "model_info": {"id": "zz"}}, + "standard_logging_object": {"total_tokens": 10}, + }, + None, + 0.0, + 0.0, + ) + + assert core.health().store.get_cache("zz:successes", local_only=True) is None + + +@pytest.mark.asyncio +async def test_async_success_callback_writes_tokens_and_a_request_for_a_tracked_deployment( + make_core: MakeCore, +) -> None: + core: Final = make_core(model_list=group("a", rpm=100)) + + await core.deployment_callback_on_success( # pyright: ignore[reportAttributeAccessIssue] # PythonCore's public member + { + "litellm_params": { + "metadata": {"model_group": "g", "deployment": "openai/a"}, + "model_info": {"id": "a"}, + }, + "standard_logging_object": {"total_tokens": 12}, + }, + None, + 0.0, + 0.0, + ) + + assert core.health().store.get_cache("a:successes", local_only=True) == 1 + assert await core.health().usage("g") == (12, 1) + + +@pytest.mark.asyncio +async def test_async_success_callback_adds_only_tokens_not_yet_counted(make_core: MakeCore) -> None: + core: Final = make_core(model_list=group("a", rpm=100)) + + await core.deployment_callback_on_success( # pyright: ignore[reportAttributeAccessIssue] # PythonCore's public member + { + "litellm_params": { + "metadata": { + "model_group": "g", + "deployment": "openai/a", + "_litellm_router_usage_counted_tokens": 8, + }, + "model_info": {"id": "a"}, + }, + "standard_logging_object": {"total_tokens": 12}, + }, + None, + 0.0, + 0.0, + ) + + assert await core.health().usage("g") == (4, 0) + + +@pytest.mark.asyncio +async def test_async_success_callback_counts_successes_for_an_untracked_deployment(make_core: MakeCore) -> None: + core: Final = make_core(model_list=group("a")) + + await core.deployment_callback_on_success( # pyright: ignore[reportAttributeAccessIssue] # PythonCore's public member + { + "litellm_params": {"metadata": {"model_group": "g"}, "model_info": {"id": "a"}}, + "standard_logging_object": {"total_tokens": 12}, + }, + None, + 0.0, + 0.0, + ) + + assert core.health().store.get_cache("a:successes", local_only=True) == 1 + assert await core.health().usage("g") == (None, None) + + +@pytest.mark.parametrize("limit", [{"rpm": 100}, {"tpm": 1000}], ids=["rpm", "tpm"]) +@pytest.mark.asyncio +async def test_a_successful_request_counts_its_tokens_and_one_request_once( + make_core: MakeCore, host: ScriptedHost, limit: dict[str, int] +) -> None: + core: Final = make_core(model_list=group("a", **limit)) + + await attempt(core) + if isinstance(host, BackgroundSuccessLoggingHost): + await host.drain() + + assert await core.health().usage("g") == (10, 1) + + +@pytest.fixture +def host(request: pytest.FixtureRequest) -> ScriptedHost: + """Background success logging for the usage test above, the plain host for the rest.""" + if request.node.name.startswith("test_a_successful_request_counts_its_tokens"): # pyright: ignore[reportAny] # pytest node + return BackgroundSuccessLoggingHost() + return ScriptedHost() + + +@pytest.mark.asyncio +async def test_a_failed_request_counts_one_request_and_no_tokens(make_core: MakeCore, host: ScriptedHost) -> None: + core: Final = make_core(model_list=group("a", rpm=100), num_retries=0) + host.answer("openai/a", fresh(bad_request)) + + await attempt(core) + + assert await core.health().usage("g") == (None, 1) + assert core.health().failures_this_minute("a") == 1 + + +@pytest.mark.asyncio +async def test_an_untracked_deployment_reports_no_usage(make_core: MakeCore, host: ScriptedHost) -> None: + core: Final = make_core(model_list=group("a")) + + await attempt(core) + + assert await core.health().usage("g") == (None, None) + + +@pytest.mark.asyncio +async def test_a_successful_request_counts_one_success_this_minute(make_core: MakeCore, host: ScriptedHost) -> None: + core: Final = make_core(model_list=group("a")) + + await attempt(core) + await attempt(core) + + assert core.health().store.get_cache("a:successes", local_only=True) == 2 + assert core.health().failures_this_minute("a") == 0 + + +# --- blocked deployments ------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_a_blocked_deployment_is_never_called(make_core: MakeCore, host: ScriptedHost) -> None: + core: Final = make_core(model_list=[deployment("g", "a", model_info={"blocked": True}), deployment("g", "b")]) + + for _ in range(4): + await attempt(core) + + assert host.deployments_called("acompletion") == ["b"] * 4 + assert core.health().fully_blocked_model_names() == set() + + +@pytest.mark.asyncio +async def test_a_group_with_every_deployment_blocked_raises_no_deployments_available( + make_core: MakeCore, host: ScriptedHost +) -> None: + core: Final = make_core(model_list=group("a", "b", model_info={"blocked": True})) + + error: Final = await attempt(core) + + assert isinstance(error, litellm.RouterRateLimitError) + assert "No deployments available for selected model" in str(error) + assert host.calls == [] + assert core.health().fully_blocked_model_names() == {"g"} + + +def test_is_blocked_and_all_blocked_read_model_info(make_core: MakeCore) -> None: + core: Final = make_core(model_list=group("a")) + blocked: Final = Deployment(**deployment("g", "a", model_info={"blocked": True})) # pyright: ignore[reportArgumentType] # a model_list entry + open_: Final = Deployment(**deployment("g", "b")) # pyright: ignore[reportArgumentType] # a model_list entry + + assert core.health().is_blocked(blocked) is True + assert core.health().is_blocked(open_) is False + assert core.health().all_blocked([blocked.model_dump()]) is True # pyright: ignore[reportArgumentType] # deployment dicts + assert core.health().all_blocked([blocked.model_dump(), open_.model_dump()]) is False # pyright: ignore[reportArgumentType] # deployment dicts + + +# --- health-check-driven filtering --------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_health_check_routing_skips_deployments_marked_unhealthy(make_core: MakeCore, host: ScriptedHost) -> None: + core: Final = make_core(model_list=group("a", "b"), enable_health_check_routing=True) + core.health().health_checks.set_deployment_health_states(unhealthy("a")) # pyright: ignore[reportArgumentType] # health states + + for _ in range(4): + await attempt(core) + + assert host.deployments_called("acompletion") == ["b"] * 4 + assert await core.health().fully_unhealthy_model_names() == set() + + +@pytest.mark.asyncio +async def test_health_check_routing_uses_every_deployment_when_all_are_unhealthy( + make_core: MakeCore, host: ScriptedHost +) -> None: + core: Final = make_core(model_list=group("a", "b"), enable_health_check_routing=True) + core.health().health_checks.set_deployment_health_states(unhealthy("a", "b")) # pyright: ignore[reportArgumentType] # health states + + for _ in range(12): + await attempt(core) + + assert set(host.deployments_called("acompletion")) == {"a", "b"} + assert await core.health().fully_unhealthy_model_names() == {"g"} + + +@pytest.mark.parametrize( + ("arguments", "age"), + [ + ({}, 0.0), + ( + { + "enable_health_check_routing": True, + "allowed_fails_policy": AllowedFailsPolicy(TimeoutErrorAllowedFails=1), + }, + 0.0, + ), + ({"enable_health_check_routing": True, "health_check_staleness_threshold": 10}, 30.0), + ({"enable_health_check_routing": True, "background_health_check_model_groups": ["other"]}, 0.0), + ], + ids=["routing-off", "allowed-fails-policy", "stale", "group-not-listed"], +) +@pytest.mark.asyncio +async def test_health_check_state_is_ignored( + make_core: MakeCore, host: ScriptedHost, arguments: dict[str, object], age: float +) -> None: + core: Final = make_core(model_list=group("a", "b"), **arguments) + core.health().health_checks.set_deployment_health_states(unhealthy("a", age=age)) # pyright: ignore[reportArgumentType] # health states + + for _ in range(12): + await attempt(core) + + assert set(host.deployments_called("acompletion")) == {"a", "b"} + + +@pytest.mark.asyncio +async def test_listed_health_check_group_is_filtered_even_with_an_allowed_fails_policy( + make_core: MakeCore, host: ScriptedHost +) -> None: + core: Final = make_core( + model_list=group("a", "b"), + enable_health_check_routing=True, + background_health_check_model_groups=["g"], + allowed_fails_policy=AllowedFailsPolicy(TimeoutErrorAllowedFails=1), + ) + core.health().health_checks.set_deployment_health_states(unhealthy("a")) # pyright: ignore[reportArgumentType] # health states + + for _ in range(4): + await attempt(core) + + assert host.deployments_called("acompletion") == ["b"] * 4 + + +@pytest.mark.asyncio +async def test_an_unhealthy_deployment_stays_excluded_when_the_healthy_one_is_cooling_down( + make_core: MakeCore, host: ScriptedHost +) -> None: + core: Final = make_core(model_list=group("a", "b"), enable_health_check_routing=True, num_retries=0) + core.health().health_checks.set_deployment_health_states(unhealthy("a")) # pyright: ignore[reportArgumentType] # health states + host.answer("openai/b", fresh(rate_limit)) + + await attempt(core) + error: Final = await attempt(core) + + assert set(cooled_down(core, ["a", "b"])) == {"b"} + assert isinstance(error, litellm.RouterRateLimitError) + assert "cooldown_list=['b']" in str(error) + assert host.deployments_called("acompletion") == ["b"] + + +def test_a_cooldown_recorded_outside_an_event_loop_raises_after_writing_it(make_core: MakeCore) -> None: + core: Final = make_core(model_list=group("a", "b"), allowed_fails=0) + + with pytest.raises(RuntimeError, match="no running event loop"): + core.health().record_failure("a", server_error(), 500) + + assert set(cooled_down(core, ["a", "b"])) == {"a"} + + +@pytest.mark.asyncio +async def test_a_request_cooldown_time_reaches_the_call_and_sets_the_cooldown( + make_core: MakeCore, host: ScriptedHost +) -> None: + core: Final = make_core(model_list=group("a", "b"), num_retries=0) + host.answer("*", fresh(rate_limit)) + + await attempt(core, cooldown_time=11) + + assert host.calls[0].kwargs["cooldown_time"] == 11 + (entry,) = cooled_down(core, ["a", "b"]).values() + assert entry["cooldown_time"] == 11 diff --git a/tests/unit/test_router/core/test_retries_health_characterization_failure_exits.py b/tests/unit/test_router/core/test_retries_health_characterization_failure_exits.py new file mode 100644 index 00000000000..63f261f1f09 --- /dev/null +++ b/tests/unit/test_router/core/test_retries_health_characterization_failure_exits.py @@ -0,0 +1,365 @@ +""" +Characterization of the failures a router core does not hold against a deployment, on every path a +failure is recorded from: the early exits of ``deployment_callback_on_failure`` (nothing counted, no +cooldown), the per-caller exits of ``_should_run_cooldown_logic`` (counted, no cooldown), and what else +decides a cooldown: the ``retry-after`` header, ``litellm.allowed_fails`` and a team deployment's siblings. +""" + +import asyncio +import email.utils +from datetime import datetime, timedelta, timezone +from typing import Final + +import httpx +import pytest + +import litellm +from litellm.constants import INTERNAL_CALL_ORIGIN_METADATA_KEY +from litellm.router_core import RouterCore +from litellm.router_utils.cooldown_handlers import mark_deployment_selection_failure +from litellm.types.router import AllowedFailsPolicy +from litellm.types.utils import BACKGROUND_RESPONSE_STATUS_POLL_CALL_ORIGIN +from tests.unit.test_router.core.retries_health_helpers import ( + ErrorBuilder, + MakeCore, + attempt, + attempt_sync, + cooled_down, + fresh, + not_found, +) +from tests.unit.test_router.core.scripted_host import ( + ScriptedHost, + auth_error, + deployment, + rate_limit, + server_error, +) +from tests.unit.test_router.core.streaming_helpers import chat_stream, collect_texts, collect_texts_sync + +MESSAGES: Final = [{"role": "user", "content": "hi"}] +TOKEN_EXCHANGE: Final = "http://127.0.0.1:9/token" +PER_USER_COPILOT: Final = { + "model": "github_copilot/gpt-4o", + "litellm_credential_name": "ghc", + "github_copilot_auth_type": "per_user_oauth", +} + + +def group( + *ids: str, **params: object +) -> list[dict[str, object]]: # kwargs-ok: litellm_params # mutable-ok: a model_list + return [deployment("g", deployment_id, **params) for deployment_id in ids] + + +def selection_failure() -> Exception: + error: Final = rate_limit() + mark_deployment_selection_failure(error) + return error + + +def caller_auth_error() -> Exception: + return litellm.CallerCredentialAuthenticationError( + message="not connected", llm_provider="github_copilot", model="gpt-4o" + ) + + +def caller_rate_limit() -> Exception: + return litellm.CallerCredentialRateLimitError( + message="token exchange limited", llm_provider="github_copilot", model="gpt-4o" + ) + + +def rate_limited_with(headers: dict[str, str]) -> litellm.RateLimitError: # mutable-ok: httpx's header mapping + return litellm.RateLimitError( + message="slow down", llm_provider="openai", model="gpt-4o", response=httpx.Response(429, headers=headers) + ) + + +async def settle() -> None: + """Let the logging and recording tasks a finished stream scheduled run.""" + for _ in range(20): + await asyncio.sleep(0) + + +# --- U1: deployment_callback_on_failure's early exits ----------------------------------------------- + + +@pytest.mark.parametrize( + ("builder", "params"), + [(selection_failure, {}), (auth_error, {"token_exchange_endpoint": TOKEN_EXCHANGE})], + ids=["selection-failure", "caller-scoped-oauth"], +) +@pytest.mark.asyncio +async def test_an_attempt_failure_outside_the_deployment_is_neither_counted_nor_cooled( + make_core: MakeCore, host: ScriptedHost, builder: ErrorBuilder, params: dict[str, object] +) -> None: + core: Final = make_core(model_list=group("a", "b", **params), num_retries=0, allowed_fails=0, cooldown_time=60) + host.answer("*", fresh(builder)) + + await attempt(core) + + assert cooled_down(core, ["a", "b"]) == {} + assert core.health().failures_this_minute("a") + core.health().failures_this_minute("b") == 0 + + +@pytest.mark.parametrize( + ("builder", "params"), + [(selection_failure, {}), (auth_error, {"token_exchange_endpoint": TOKEN_EXCHANGE})], + ids=["selection-failure", "caller-scoped-oauth"], +) +def test_a_sync_attempt_failure_outside_the_deployment_is_neither_counted_nor_cooled( + make_core: MakeCore, host: ScriptedHost, builder: ErrorBuilder, params: dict[str, object] +) -> None: + core: Final = make_core(model_list=group("a", "b", **params), num_retries=0, allowed_fails=0, cooldown_time=60) + host.answer("*", fresh(builder)) + + attempt_sync(core) + + assert cooled_down(core, ["a", "b"]) == {} + assert core.health().failures_this_minute("a") + core.health().failures_this_minute("b") == 0 + + +@pytest.mark.asyncio +async def test_a_background_poll_404_on_an_attempt_is_neither_counted_nor_cooled( + make_core: MakeCore, host: ScriptedHost +) -> None: + core: Final = make_core(model_list=group("a", "b"), num_retries=0, allowed_fails=0) + host.answer("*", fresh(not_found)) + + await attempt(core, metadata={INTERNAL_CALL_ORIGIN_METADATA_KEY: BACKGROUND_RESPONSE_STATUS_POLL_CALL_ORIGIN}) + + assert cooled_down(core, ["a", "b"]) == {} + assert core.health().failures_this_minute("a") + core.health().failures_this_minute("b") == 0 + + +@pytest.mark.asyncio +async def test_a_stream_failure_outside_the_deployment_still_counts_its_request( + make_core: MakeCore, host: ScriptedHost +) -> None: + host.answer("openai/a", chat_stream("partial", error=selection_failure(), logged=True)) + core: Final = make_core( + model_list=[deployment("g", "a", rpm=100, order=1), deployment("g", "b", order=2)], + num_retries=0, + allowed_fails=0, + ) + + stream: Final = await core.acompletion(model="g", messages=MESSAGES, stream=True) # pyright: ignore[reportAttributeAccessIssue] # every core serves acompletion + with pytest.raises(litellm.RateLimitError): + await collect_texts(stream) + await settle() + + assert cooled_down(core, ["a", "b"]) == {} + assert core.health().failures_this_minute("a") == 0 + # `async_deployment_callback_on_failure` makes none of the failure callback's exits. + assert await core.health().usage("g") == (0, 2) + + +def test_a_sync_stream_failure_outside_the_deployment_is_neither_counted_nor_cooled( + make_core: MakeCore, host: ScriptedHost +) -> None: + host.answer("openai/a", chat_stream("partial", error=selection_failure(), logged=True, sync=True)) + core: Final = make_core( + model_list=[deployment("g", "a", order=1), deployment("g", "b", order=2)], num_retries=0, allowed_fails=0 + ) + + stream: Final = core.completion(model="g", messages=MESSAGES, stream=True) # pyright: ignore[reportAttributeAccessIssue] # every core serves completion + with pytest.raises(litellm.RateLimitError): + collect_texts_sync(stream) + + assert cooled_down(core, ["a", "b"]) == {} + assert core.health().failures_this_minute("a") == 0 + + +@pytest.mark.asyncio +async def test_the_failure_hook_counts_a_failure_outside_the_deployment_as_nothing( + make_core: MakeCore, host: ScriptedHost +) -> None: + core: Final = make_core(model_list=group("a", "b"), num_retries=0, allowed_fails=0) + logged: Final = { + "exception": selection_failure(), + "litellm_params": {"metadata": {"model_group": "g"}, "model_info": {"id": "a"}}, + } + + assert core.deployment_callback_on_failure(logged, None, None, None) is False # pyright: ignore[reportAttributeAccessIssue] # every core serves it + assert core.health().failures_this_minute("a") == 0 + + +# --- N13: _should_run_cooldown_logic's per-caller exits --------------------------------------------- + + +@pytest.mark.parametrize("builder", [caller_auth_error, caller_rate_limit], ids=["authentication", "rate-limit"]) +@pytest.mark.asyncio +async def test_a_caller_credential_error_is_counted_but_never_cools_a_shared_deployment( + make_core: MakeCore, host: ScriptedHost, builder: ErrorBuilder +) -> None: + core: Final = make_core(model_list=group("a", "b"), num_retries=0, allowed_fails=0) + host.answer("*", fresh(builder)) + + await attempt(core) + + assert cooled_down(core, ["a", "b"]) == {} + assert core.health().failures_this_minute("a") + core.health().failures_this_minute("b") == 1 + + +@pytest.mark.parametrize("builder", [auth_error, server_error], ids=["401", "500"]) +@pytest.mark.asyncio +async def test_a_per_user_copilot_deployment_is_counted_but_never_cools_down( + make_core: MakeCore, host: ScriptedHost, builder: ErrorBuilder +) -> None: + core: Final = make_core( + model_list=[deployment("g", "gh", weight=1, **PER_USER_COPILOT), deployment("g", "other", weight=0)], + num_retries=0, + allowed_fails=0, + ) + host.answer("*", fresh(builder)) + + await attempt(core) + + assert cooled_down(core, ["gh", "other"]) == {} + assert core.health().failures_this_minute("gh") == 1 + + +# --- S3: the retry-after header ---------------------------------------------------------------------- + + +@pytest.mark.parametrize( + ("headers", "expected_seconds"), + [ + ({"retry-after-ms": "3000", "retry-after": "8"}, 8), + ({"retry-after-ms": "7000"}, 5), + ({"retry-after": "2.5"}, 5), + ({"retry-after": "soon"}, 5), + ({"retry-after": "-3"}, 5), + ({"retry-after": " 7 "}, 7), + ], + ids=["seconds-over-ms", "ms-ignored", "float-ignored", "unreadable", "negative", "padded-int"], +) +@pytest.mark.asyncio +async def test_a_cooldown_takes_its_time_from_retry_after_read_as_an_int_only( + make_core: MakeCore, host: ScriptedHost, headers: dict[str, str], expected_seconds: int +) -> None: + core: Final = make_core(model_list=group("a", "b"), num_retries=0, cooldown_time=5) + host.answer("*", lambda **_: rate_limited_with(headers)) + + await attempt(core) + + (entry,) = cooled_down(core, ["a", "b"]).values() + assert entry["cooldown_time"] == expected_seconds + + +@pytest.mark.asyncio +async def test_a_cooldown_reads_a_retry_after_date(make_core: MakeCore, host: ScriptedHost) -> None: + core: Final = make_core(model_list=group("a", "b"), num_retries=0, cooldown_time=5) + until: Final = email.utils.format_datetime(datetime.now(timezone.utc) + timedelta(seconds=120), usegmt=True) + host.answer("*", lambda **_: rate_limited_with({"retry-after": until})) + + await attempt(core) + + (entry,) = cooled_down(core, ["a", "b"]).values() + cooldown_time: Final = entry["cooldown_time"] + assert isinstance(cooldown_time, int | float) and 100 <= cooldown_time <= 120 + + +# --- U5: litellm.allowed_fails ----------------------------------------------------------------------- + + +def weighted() -> list[dict[str, object]]: # mutable-ok: a model_list + return [deployment("g", "a", weight=1), deployment("g", "b", weight=0)] + + +@pytest.mark.asyncio +async def test_a_write_to_the_process_allowed_fails_starts_the_count_at_the_copied_value( + make_core: MakeCore, host: ScriptedHost, monkeypatch: pytest.MonkeyPatch +) -> None: + core: Final = make_core(model_list=weighted(), num_retries=0) + monkeypatch.setattr(litellm, "allowed_fails", 0) + host.answer("*", fresh(rate_limit)) + + await attempt(core) + + assert core.settings().allowed_fails == 3 + assert cooled_down(core, ["a", "b"]) == {} + + +@pytest.mark.asyncio +async def test_a_process_allowed_fails_put_back_ends_the_count( + make_core: MakeCore, host: ScriptedHost, monkeypatch: pytest.MonkeyPatch +) -> None: + core: Final = make_core(model_list=weighted(), num_retries=0) + host.answer("*", fresh(rate_limit)) + monkeypatch.setattr(litellm, "allowed_fails", 0) + await attempt(core) + assert cooled_down(core, ["a", "b"]) == {} + + monkeypatch.setattr(litellm, "allowed_fails", 3) + await attempt(core) + + assert list(cooled_down(core, ["a", "b"])) == ["a"] + + +@pytest.mark.asyncio +async def test_the_allowed_fails_count_is_the_process_value_copied_at_init( + make_core: MakeCore, host: ScriptedHost, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(litellm, "allowed_fails", 1) + host.set_dynamic(allowed_fails=1) + core: Final = make_core( + model_list=weighted(), num_retries=0, allowed_fails_policy=AllowedFailsPolicy(TimeoutErrorAllowedFails=0) + ) + host.answer("*", fresh(rate_limit)) + + await attempt(core) + assert cooled_down(core, ["a", "b"]) == {} + + await attempt(core) + assert list(cooled_down(core, ["a", "b"])) == ["a"] + + +# --- U10: a team deployment's siblings -------------------------------------------------------------- + + +def team_deployment( + name: str, deployment_id: str, **info: object +) -> dict[str, object]: # kwargs-ok: model_info # mutable-ok: a model_list entry + return deployment( + name, + deployment_id, + weight=1 if deployment_id == "ta1" else 0, + model_info={"team_id": "T1", "team_public_model_name": "tp", **info}, + ) + + +async def call_team_model(core: RouterCore) -> object: + return await attempt(core, "tp", metadata={"user_api_key_team_id": "T1"}) + + +@pytest.mark.asyncio +async def test_a_team_deployment_with_a_sibling_under_its_public_name_cools_down( + make_core: MakeCore, host: ScriptedHost +) -> None: + core: Final = make_core(model_list=[team_deployment("t1a", "ta1"), team_deployment("t1b", "ta2")], num_retries=0) + host.answer("*", fresh(rate_limit)) + + await call_team_model(core) + + assert list(cooled_down(core, ["ta1", "ta2"])) == ["ta1"] + + +@pytest.mark.parametrize("others", [(), (("t1b", "ta2", {"blocked": True}),)], ids=["alone", "blocked-sibling"]) +@pytest.mark.asyncio +async def test_a_team_deployment_with_no_routable_sibling_stays_a_lone_deployment( + make_core: MakeCore, host: ScriptedHost, others: tuple[tuple[str, str, dict[str, object]], ...] +) -> None: + core: Final = make_core( + model_list=[ + team_deployment("t1a", "ta1"), + *(team_deployment(name, deployment_id, **info) for name, deployment_id, info in others), + ], + num_retries=0, + ) + host.answer("*", fresh(rate_limit)) + + await call_team_model(core) + + assert cooled_down(core, ["ta1", "ta2"]) == {} diff --git a/tests/unit/test_router/core/test_retries_health_characterization_retries.py b/tests/unit/test_router/core/test_retries_health_characterization_retries.py new file mode 100644 index 00000000000..5c6c885a59f --- /dev/null +++ b/tests/unit/test_router/core/test_retries_health_characterization_retries.py @@ -0,0 +1,683 @@ +""" +Characterization of a router core's retry loop: which errors retry, how many times, how long it sleeps, +what it records on the request and the final error, and the public retry helpers. +""" + +from typing import Final + +import pytest + +import litellm +from litellm.constants import INITIAL_RETRY_DELAY, JITTER, MAX_RETRY_DELAY +from litellm.exceptions import GuardrailRaisedException +from litellm.litellm_core_utils.agentic_followup_failure import is_agentic_followup_failure +from litellm.types.router import RetryPolicy +from litellm.types.utils import Usage +from tests.unit.test_router.core.retries_health_helpers import ( + ErrorBuilder, + MakeCore, + MetadataRecorder, + api_connection, + api_status, + attempt, + attempt_sync, + bad_gateway, + cooled_down, + fresh, + headers_of, + http_response, + not_found, + permission_denied, + plain_error, + rate_limited, + service_unavailable, + unprocessable, +) +from tests.unit.test_router.core.scripted_host import ( + ScriptedHost, + agentic_followup_failure, + auth_error, + bad_request, + content_policy, + context_window, + deployment, + rate_limit, + server_error, + timeout, +) +from tests.unit.test_router.core.streaming_helpers import request_logging + + +def group( + *ids: str, **params: object +) -> list[dict[str, object]]: # kwargs-ok: litellm_params # mutable-ok: a model_list + return [deployment("g", deployment_id, **params) for deployment_id in ids] + + +def within(slept: float, base: float) -> bool: + return base <= slept < base + JITTER + + +# --- which errors retry, by error class and status ------------------------------------------------ + + +@pytest.mark.parametrize( + ("builder", "lone_calls", "pair_calls"), + [ + (not_found, 1, 1), + (permission_denied, 1, 3), + (unprocessable, 1, 1), + (service_unavailable, 3, 3), + (bad_gateway, 3, 3), + (api_connection, 3, 3), + (api_status(409), 3, 3), + (plain_error, 3, 3), + ], + ids=["404", "403", "422", "503", "502", "connection", "409", "no-status"], +) +@pytest.mark.parametrize("size", ["lone", "pair"]) +@pytest.mark.asyncio +async def test_retry_decision_per_error_class_and_group_size( + make_core: MakeCore, host: ScriptedHost, builder: ErrorBuilder, lone_calls: int, pair_calls: int, size: str +) -> None: + core: Final = make_core(model_list=group("a") if size == "lone" else group("a", "b"), num_retries=2) + host.answer("*", fresh(builder)) + + error: Final = await attempt(core) + + assert type(error) is type(builder()) + assert len(host.calls_to("acompletion")) == (lone_calls if size == "lone" else pair_calls) + + +@pytest.mark.asyncio +async def test_an_error_without_a_litellm_type_is_retried_but_not_stamped( + make_core: MakeCore, host: ScriptedHost +) -> None: + core: Final = make_core(model_list=group("a"), num_retries=2) + host.answer("openai/a", fresh(plain_error)) + + error: Final = await attempt(core) + + assert isinstance(error, ValueError) + assert getattr(error, "num_retries", None) is None + assert getattr(error, "max_retries", None) is None + assert "LiteLLM Retried" not in str(error) + + +@pytest.mark.asyncio +async def test_a_403_retry_moves_to_a_deployment_that_has_not_refused_yet( + make_core: MakeCore, host: ScriptedHost +) -> None: + core: Final = make_core(model_list=group("a", "b", "c"), num_retries=4) + host.answer("*", fresh(permission_denied)) + + error: Final = await attempt(core) + + called: Final = host.deployments_called("acompletion") + assert isinstance(error, litellm.PermissionDeniedError) + assert len(called) == 5 + assert sorted(called[:3]) == ["a", "b", "c"] + assert cooled_down(core, ["a", "b", "c"]) == {} + + +@pytest.mark.asyncio +async def test_a_401_on_each_deployment_cools_each_down_and_stops_when_none_is_left( + make_core: MakeCore, host: ScriptedHost +) -> None: + core: Final = make_core(model_list=group("a", "b", "c"), num_retries=4) + host.answer("*", fresh(auth_error)) + + error: Final = await attempt(core) + + assert isinstance(error, litellm.AuthenticationError) + assert sorted(host.deployments_called("acompletion")) == ["a", "b", "c"] + assert set(cooled_down(core, ["a", "b", "c"])) == {"a", "b", "c"} + + +@pytest.mark.asyncio +async def test_a_non_retryable_error_mid_loop_stops_the_retries(make_core: MakeCore, host: ScriptedHost) -> None: + core: Final = make_core(model_list=group("a"), num_retries=4) + host.answer("openai/a", fresh(server_error), fresh(bad_request)) + + error: Final = await attempt(core) + + assert isinstance(error, litellm.BadRequestError) + assert len(host.calls_to("acompletion")) == 2 + assert getattr(error, "num_retries", None) is None + + +@pytest.mark.asyncio +async def test_a_retry_policy_keeps_retrying_through_a_later_non_retryable_error( + make_core: MakeCore, host: ScriptedHost +) -> None: + core: Final = make_core(model_list=group("a"), retry_policy=RetryPolicy(InternalServerErrorRetries=3)) + host.answer("openai/a", fresh(server_error), fresh(bad_request)) + + error: Final = await attempt(core) + + assert isinstance(error, litellm.BadRequestError) + assert len(host.calls_to("acompletion")) == 4 + assert getattr(error, "num_retries", None) == 3 + + +@pytest.mark.asyncio +async def test_a_guardrail_block_is_raised_without_retrying(make_core: MakeCore, host: ScriptedHost) -> None: + core: Final = make_core(model_list=group("a", "b"), num_retries=3) + host.answer("*", lambda **_: GuardrailRaisedException(guardrail_name="pii", message="blocked by guardrail")) + + error: Final = await attempt(core) + + assert isinstance(error, GuardrailRaisedException) + assert len(host.calls_to("acompletion")) == 1 + assert host.sleeps == [] + + +@pytest.mark.asyncio +async def test_a_failed_agentic_followup_is_raised_without_retrying(make_core: MakeCore, host: ScriptedHost) -> None: + core: Final = make_core(model_list=group("a", "b"), num_retries=3) + host.answer("*", lambda **_: agentic_followup_failure(server_error())) + + error: Final = await attempt(core) + + assert isinstance(error, litellm.InternalServerError) + assert len(host.calls_to("acompletion")) == 1 + assert host.sleeps == [] + + +@pytest.mark.asyncio +async def test_a_retry_that_fails_its_agentic_followup_ends_the_retries( + make_core: MakeCore, host: ScriptedHost +) -> None: + core: Final = make_core(model_list=group("a", "b"), num_retries=3) + host.answer("*", server_error(), lambda **_: agentic_followup_failure(server_error()), None) + + error: Final = await attempt(core) + + assert isinstance(error, Exception) + assert is_agentic_followup_failure(error) + assert len(host.calls_to("acompletion")) == 2 + + +@pytest.mark.asyncio +async def test_a_retry_carries_the_superseded_attempt_usage(make_core: MakeCore, host: ScriptedHost) -> None: + logging_obj: Final = request_logging() + usage: Final = Usage(prompt_tokens=1, completion_tokens=1, total_tokens=2) + + def failed_after_usage(**_: object) -> Exception: + logging_obj.model_call_details["combined_usage_object"] = usage + return server_error() + + core: Final = make_core(model_list=group("a", "b"), num_retries=1) + host.answer("*", failed_after_usage, None) + + await attempt(core, litellm_logging_obj=logging_obj) + + assert logging_obj.model_call_details["failed_attempts_usage"] == ((usage, 0.0),) + + +# --- retry policies, more error types -------------------------------------------------------------- + + +@pytest.mark.parametrize( + ("policy", "builder", "expected_calls"), + [ + (RetryPolicy(ServiceUnavailableErrorRetries=2), service_unavailable, 3), + (RetryPolicy(InternalServerErrorRetries=2), service_unavailable, 1 + 2), + (RetryPolicy(DefaultRetries=0), api_status(409), 1), + (RetryPolicy(DefaultRetries=2), unprocessable, 3), + (RetryPolicy(BadRequestErrorRetries=1), content_policy, 2), + (RetryPolicy(BadRequestErrorRetries=1, ContentPolicyViolationErrorRetries=0), content_policy, 1), + ], + ids=[ + "service-unavailable", + "503-falls-back-to-internal-server-error", + "default-for-409", + "default-for-422", + "content-policy-inherits-bad-request", + "content-policy-specific-wins", + ], +) +@pytest.mark.asyncio +async def test_retry_policy_field_chosen_for_each_error( + make_core: MakeCore, host: ScriptedHost, policy: RetryPolicy, builder: ErrorBuilder, expected_calls: int +) -> None: + core: Final = make_core(model_list=group("a"), retry_policy=policy) + host.answer("openai/a", fresh(builder)) + + error: Final = await attempt(core) + + assert type(error) is type(builder()) + assert len(host.calls_to("acompletion")) == expected_calls + + +@pytest.mark.parametrize( + ("policy", "builder"), + [ + (RetryPolicy(AuthenticationErrorRetries=2), auth_error), + (RetryPolicy(NotFoundErrorRetries=2), not_found), + (RetryPolicy(DefaultRetries=1), not_found), + ], + ids=["auth", "not-found", "default-for-404"], +) +@pytest.mark.asyncio +async def test_retry_policy_retries_a_401_or_404_into_the_cooled_lone_deployment( + make_core: MakeCore, host: ScriptedHost, policy: RetryPolicy, builder: ErrorBuilder +) -> None: + core: Final = make_core(model_list=group("a"), retry_policy=policy) + host.answer("openai/a", fresh(builder)) + + error: Final = await attempt(core) + + assert isinstance(error, litellm.RouterRateLimitError) + assert "All deployments for selected model are in cooldown" in str(error) + assert len(host.calls_to("acompletion")) == 1 + + +@pytest.mark.parametrize( + ("policy", "builder"), + [(RetryPolicy(AuthenticationErrorRetries=2), auth_error), (RetryPolicy(NotFoundErrorRetries=2), not_found)], + ids=["auth", "not-found"], +) +@pytest.mark.asyncio +async def test_retry_policy_retries_a_401_or_404_when_cooldowns_are_off( + make_core: MakeCore, host: ScriptedHost, policy: RetryPolicy, builder: ErrorBuilder +) -> None: + core: Final = make_core(model_list=group("a"), retry_policy=policy, disable_cooldowns=True) + host.answer("openai/a", fresh(builder)) + + error: Final = await attempt(core) + + assert type(error) is type(builder()) + assert len(host.calls_to("acompletion")) == 3 + + +@pytest.mark.asyncio +async def test_retry_policy_wins_over_the_deployment_num_retries(make_core: MakeCore, host: ScriptedHost) -> None: + core: Final = make_core( + model_list=group("a", num_retries=4), retry_policy=RetryPolicy(InternalServerErrorRetries=1) + ) + host.answer("openai/a", fresh(server_error)) + + await attempt(core) + + assert len(host.calls_to("acompletion")) == 2 + + +@pytest.mark.asyncio +async def test_request_model_group_retry_policy_overrides_the_router_one( + make_core: MakeCore, host: ScriptedHost +) -> None: + core: Final = make_core( + model_list=group("a"), model_group_retry_policy={"g": RetryPolicy(InternalServerErrorRetries=0)} + ) + host.answer("openai/a", fresh(server_error)) + + await attempt(core, model_group_retry_policy={"g": RetryPolicy(InternalServerErrorRetries=3)}) + + assert len(host.calls_to("acompletion")) == 4 + + +@pytest.mark.asyncio +async def test_model_group_retry_policy_given_as_a_dict(make_core: MakeCore, host: ScriptedHost) -> None: + core: Final = make_core( + model_list=group("a"), num_retries=0, model_group_retry_policy={"g": {"InternalServerErrorRetries": 2}} + ) + host.answer("openai/a", fresh(server_error)) + + await attempt(core) + + assert len(host.calls_to("acompletion")) == 3 + + +@pytest.mark.parametrize( + ("error", "model_group", "expected"), + [ + (server_error(), "g", 1), + (server_error(), "h", 4), + (server_error(), None, 4), + (rate_limit(), "g", None), + (rate_limit(), "h", 2), + (bad_request(), "h", 7), + (context_window(), "h", 7), + (not_found(), "h", 7), + ], +) +def test_get_num_retries_from_retry_policy( + make_core: MakeCore, error: Exception, model_group: str | None, expected: int | None +) -> None: + core: Final = make_core( + model_list=group("a"), + retry_policy=RetryPolicy(InternalServerErrorRetries=4, RateLimitErrorRetries=2, DefaultRetries=7), + model_group_retry_policy={"g": RetryPolicy(InternalServerErrorRetries=1)}, + ) + + assert core.get_num_retries_from_retry_policy(error, model_group=model_group) == expected # pyright: ignore[reportAttributeAccessIssue] # PythonCore's public member + + +def test_get_num_retries_from_retry_policy_without_any_policy_is_none(make_core: MakeCore) -> None: + core: Final = make_core(model_list=group("a"), num_retries=3) + + assert core.get_num_retries_from_retry_policy(server_error(), model_group="g") is None # pyright: ignore[reportAttributeAccessIssue] # PythonCore's public member + + +# --- should_retry_this_error --------------------------------------------------------------------- + +HEALTHY_ONE: Final = [{"model_info": {"id": "a"}}] +ALL_TWO: Final = [{"model_info": {"id": "a"}}, {"model_info": {"id": "b"}}] + + +@pytest.mark.parametrize( + ("error", "arguments", "retries"), + [ + (server_error(), {"healthy_deployments": HEALTHY_ONE, "all_deployments": HEALTHY_ONE}, True), + (server_error(), {"healthy_deployments": [], "all_deployments": HEALTHY_ONE}, False), + (server_error(), {}, False), + (timeout(), {"healthy_deployments": HEALTHY_ONE}, True), + (plain_error(), {"healthy_deployments": HEALTHY_ONE}, True), + (plain_error(), {"healthy_deployments": None}, False), + (bad_request(), {"healthy_deployments": HEALTHY_ONE}, False), + (context_window(), {"healthy_deployments": HEALTHY_ONE, "context_window_fallbacks": []}, False), + (content_policy(), {"healthy_deployments": HEALTHY_ONE, "content_policy_fallbacks": []}, False), + (not_found(), {"healthy_deployments": ALL_TWO, "all_deployments": ALL_TWO}, False), + (unprocessable(), {"healthy_deployments": ALL_TWO, "all_deployments": ALL_TWO}, False), + (auth_error(), {"healthy_deployments": HEALTHY_ONE, "all_deployments": HEALTHY_ONE}, False), + (auth_error(), {"healthy_deployments": HEALTHY_ONE, "all_deployments": ALL_TWO}, True), + (auth_error(), {"healthy_deployments": [], "all_deployments": ALL_TWO}, False), + (permission_denied(), {"healthy_deployments": HEALTHY_ONE, "all_deployments": HEALTHY_ONE}, False), + (permission_denied(), {"healthy_deployments": HEALTHY_ONE, "all_deployments": ALL_TWO}, True), + (rate_limit(), {"healthy_deployments": HEALTHY_ONE, "regular_fallbacks": [{"g": ["h"]}]}, True), + (rate_limit(), {"healthy_deployments": [], "regular_fallbacks": [{"g": ["h"]}]}, False), + (rate_limit(), {"healthy_deployments": [], "regular_fallbacks": []}, False), + (api_status(409)(), {"healthy_deployments": HEALTHY_ONE}, True), + (service_unavailable(), {"healthy_deployments": HEALTHY_ONE}, True), + ], +) +def test_should_retry_this_error_returns_true_or_raises_the_error( + make_core: MakeCore, error: Exception, arguments: dict[str, object], retries: bool +) -> None: + core: Final = make_core(model_list=group("a", "b")) + + if retries: + assert core.should_retry_this_error(error=error, **arguments) is True # pyright: ignore[reportAttributeAccessIssue] # PythonCore's public member + else: + with pytest.raises(type(error)) as raised: + core.should_retry_this_error(error=error, **arguments) # pyright: ignore[reportAttributeAccessIssue] # PythonCore's public member + assert raised.value is error + + +# --- mock rate limit errors ---------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_mock_rate_limit_error_raises_without_calling_the_deployment( + make_core: MakeCore, host: ScriptedHost +) -> None: + core: Final = make_core(model_list=group("a"), num_retries=0) + + error: Final = await attempt(core, mock_testing_rate_limit_error=True) + + assert isinstance(error, litellm.RateLimitError) + assert "This is a mock exception for model=g, to trigger a rate limit error." in str(error) + assert host.calls == [] + + +@pytest.mark.asyncio +async def test_mock_rate_limit_error_is_retried_and_the_retry_reaches_the_deployment( + make_core: MakeCore, host: ScriptedHost +) -> None: + core: Final = make_core(model_list=group("a"), num_retries=2) + + response: Final = await attempt(core, mock_testing_rate_limit_error=True) + + assert not isinstance(response, Exception) + assert host.deployments_called("acompletion") == ["a"] + assert "mock_testing_rate_limit_error" not in host.calls[0].kwargs + assert headers_of(response)["x-litellm-attempted-retries"] == 1 + assert headers_of(response)["x-litellm-max-retries"] == 2 + + +@pytest.mark.asyncio +async def test_mock_rate_limit_error_carries_a_lone_deployments_num_retries( + make_core: MakeCore, host: ScriptedHost +) -> None: + core: Final = make_core(model_list=group("a", num_retries=1), num_retries=5) + + response: Final = await attempt(core, mock_testing_rate_limit_error=True) + + assert headers_of(response)["x-litellm-max-retries"] == 1 + + +@pytest.mark.parametrize("value", [False, None]) +@pytest.mark.asyncio +async def test_mock_rate_limit_error_only_fires_when_true( + make_core: MakeCore, host: ScriptedHost, value: object +) -> None: + core: Final = make_core(model_list=group("a"), num_retries=0) + + response: Final = await attempt(core, mock_testing_rate_limit_error=value) + + assert not isinstance(response, Exception) + assert host.deployments_called("acompletion") == ["a"] + + +# --- sleeps -------------------------------------------------------------------------------------- + + +@pytest.mark.parametrize("num_retries", [1, 2, 3]) +@pytest.mark.asyncio +async def test_lone_deployment_sleeps_once_more_than_it_retries( + make_core: MakeCore, host: ScriptedHost, num_retries: int +) -> None: + core: Final = make_core(model_list=group("a"), num_retries=num_retries) + host.answer("openai/a", fresh(server_error)) + + await attempt(core) + + bases: Final = [INITIAL_RETRY_DELAY * 2**n for n in range(num_retries + 1)] + assert len(host.sleeps) == num_retries + 1 + assert all(within(slept, base) for slept, base in zip(host.sleeps, bases)) + + +@pytest.mark.asyncio +async def test_backoff_is_capped_at_max_retry_delay(make_core: MakeCore, host: ScriptedHost) -> None: + core: Final = make_core(model_list=group("a"), num_retries=6) + host.answer("openai/a", fresh(server_error)) + + await attempt(core) + + expected: Final = [min(INITIAL_RETRY_DELAY * 2**n, MAX_RETRY_DELAY) for n in range(7)] + assert all(within(slept, base) for slept, base in zip(host.sleeps, expected)) + assert within(host.sleeps[-1], MAX_RETRY_DELAY) + + +@pytest.mark.asyncio +async def test_retry_after_setting_above_max_retry_delay_is_capped(make_core: MakeCore, host: ScriptedHost) -> None: + core: Final = make_core(model_list=group("a"), num_retries=1, retry_after=20) + host.answer("openai/a", fresh(server_error)) + + await attempt(core) + + assert all(within(slept, MAX_RETRY_DELAY) for slept in host.sleeps) + + +@pytest.mark.asyncio +async def test_a_retry_after_header_of_zero_falls_back_to_backoff(make_core: MakeCore, host: ScriptedHost) -> None: + core: Final = make_core(model_list=group("a"), num_retries=1) + host.answer("openai/a", lambda **_: rate_limited(retry_after="0")) + + await attempt(core) + + assert within(host.sleeps[0], INITIAL_RETRY_DELAY) + + +@pytest.mark.asyncio +async def test_retry_after_header_on_any_error_type_sets_the_sleep(make_core: MakeCore, host: ScriptedHost) -> None: + core: Final = make_core(model_list=group("a"), num_retries=1) + host.answer( + "openai/a", + lambda **_: litellm.InternalServerError( + message="boom", + llm_provider="openai", + model="gpt-4o", + response=http_response(500, {"retry-after": "4"}), + ), + ) + + await attempt(core) + + assert all(within(slept, 4) for slept in host.sleeps) + + +@pytest.mark.asyncio +async def test_litellm_response_headers_on_the_error_set_the_sleep(make_core: MakeCore, host: ScriptedHost) -> None: + def failure(**_: object) -> Exception: + error: Final = server_error() + error.litellm_response_headers = {"retry-after": "6"} # pyright: ignore[reportAttributeAccessIssue] # what provider handlers attach + return error + + core: Final = make_core(model_list=group("a"), num_retries=1) + host.answer("openai/a", failure) + + await attempt(core) + + assert all(within(slept, 6) for slept in host.sleeps) + + +@pytest.mark.asyncio +async def test_retry_after_header_is_ignored_while_siblings_are_healthy( + make_core: MakeCore, host: ScriptedHost +) -> None: + core: Final = make_core(model_list=group("a", "b"), num_retries=2, allowed_fails=10) + host.answer("*", fresh(server_error)) + + await attempt(core) + + assert host.sleeps == [0, 0, 0] + + +# --- what the retry loop writes on the request ------------------------------------------------- + + +@pytest.mark.asyncio +async def test_each_attempt_carries_its_retry_count_and_the_failures_before_it( + make_core: MakeCore, host: ScriptedHost +) -> None: + recorder: Final = MetadataRecorder() + core: Final = make_core(model_list=group("a", "b"), num_retries=2, allowed_fails=10) + host.answer("*", recorder.then(server_error), recorder.then(server_error), recorder.then(None)) + + response: Final = await attempt(core) + + assert not isinstance(response, Exception) + snapshots: Final = recorder.snapshots + assert [s["attempted_retries"] for s in snapshots] == [0, 1, 2] + assert [s["max_retries"] for s in snapshots] == [2, 2, 2] + assert [s["model_group_size"] for s in snapshots] == [2, 2, 2] + assert [s["request_retry_count"] for s in snapshots] == [None, 1, 2] + assert snapshots[0]["previous_models"] is None + trail: Final = snapshots[2]["previous_models"] + assert isinstance(trail, tuple) + assert [entry["attempted_retries"] for entry in trail] == [0, 1] + assert [entry["exception_type"] for entry in trail] == ["InternalServerError"] * 2 + assert all(entry["model_group"] == "g" for entry in trail) + called: Final = host.deployments_called("acompletion") + assert [entry["deployment_id"] for entry in trail] == called[:2] + assert [entry["model"] for entry in trail] == [f"openai/{d}" for d in called[:2]] + assert all(entry["exception_string"] == "litellm.InternalServerError: boom" for entry in trail) + + +@pytest.mark.asyncio +async def test_max_retries_on_the_request_follows_the_deployment_num_retries( + make_core: MakeCore, host: ScriptedHost +) -> None: + recorder: Final = MetadataRecorder() + core: Final = make_core(model_list=group("a", num_retries=1), num_retries=3) + host.answer("openai/a", recorder.then(server_error), recorder.then(None)) + + response: Final = await attempt(core) + + assert [s["max_retries"] for s in recorder.snapshots] == [3, 1] + assert headers_of(response)["x-litellm-max-retries"] == 1 + + +@pytest.mark.asyncio +async def test_exhausted_deployment_num_retries_are_stamped_on_the_error( + make_core: MakeCore, host: ScriptedHost +) -> None: + core: Final = make_core(model_list=group("a", num_retries=2), num_retries=5) + host.answer("openai/a", fresh(server_error)) + + error: Final = await attempt(core) + + assert getattr(error, "num_retries", None) == 2 + assert getattr(error, "max_retries", None) == 2 + + +@pytest.mark.parametrize("bad_value", ["many", None]) +@pytest.mark.asyncio +async def test_a_deployment_num_retries_that_is_not_a_number_falls_back_to_the_router( + make_core: MakeCore, host: ScriptedHost, bad_value: object +) -> None: + core: Final = make_core(model_list=group("a", num_retries=bad_value), num_retries=1) + host.answer("openai/a", fresh(server_error)) + + await attempt(core) + + assert len(host.calls_to("acompletion")) == 2 + + +# --- sync and other endpoints ------------------------------------------------------------------ + + +@pytest.mark.asyncio +async def test_sync_completion_retries_like_the_async_one(make_core: MakeCore, host: ScriptedHost) -> None: + core: Final = make_core(model_list=group("a"), num_retries=2) + host.answer("openai/a", fresh(server_error)) + + error: Final = attempt_sync(core) + + assert isinstance(error, litellm.InternalServerError) + assert host.deployments_called("completion") == ["a", "a", "a"] + assert len(host.sleeps) == 3 + assert getattr(error, "num_retries", None) == 2 + + +@pytest.mark.asyncio +async def test_sync_completion_success_after_a_retry_reports_retry_headers( + make_core: MakeCore, host: ScriptedHost +) -> None: + core: Final = make_core(model_list=group("a"), num_retries=2) + host.answer("openai/a", fresh(server_error), None) + + response: Final = attempt_sync(core) + + assert headers_of(response)["x-litellm-attempted-retries"] == 1 + assert headers_of(response)["x-litellm-max-retries"] == 2 + + +@pytest.mark.parametrize("op", ["aembedding", "embedding"]) +@pytest.mark.asyncio +async def test_embeddings_retry_with_the_router_num_retries(make_core: MakeCore, host: ScriptedHost, op: str) -> None: + core: Final = make_core(model_list=group("a"), num_retries=2) + host.answer("openai/a", fresh(server_error)) + + async def embed() -> object: + if op == "aembedding": + return await core.aembedding(model="g", input=["x"]) # pyright: ignore[reportAttributeAccessIssue] # every core serves aembedding + return core.embedding(model="g", input=["x"]) # pyright: ignore[reportAttributeAccessIssue] # every core serves embedding + + with pytest.raises(litellm.InternalServerError): + await embed() + + assert host.deployments_called(op) == ["a", "a", "a"] + + +@pytest.mark.asyncio +async def test_embedding_success_after_a_retry_reports_retry_headers(make_core: MakeCore, host: ScriptedHost) -> None: + core: Final = make_core(model_list=group("a"), num_retries=2) + host.answer("openai/a", fresh(timeout), None) + + response: Final = await core.aembedding(model="g", input=["x"]) # pyright: ignore[reportAttributeAccessIssue] # every core serves aembedding + + assert headers_of(response)["x-litellm-attempted-retries"] == 1 + assert headers_of(response)["x-litellm-max-retries"] == 2 diff --git a/tests/unit/test_router/core/test_selection.py b/tests/unit/test_router/core/test_selection.py index 4b4c320c428..b33316ff380 100644 --- a/tests/unit/test_router/core/test_selection.py +++ b/tests/unit/test_router/core/test_selection.py @@ -26,7 +26,7 @@ from tests.unit.test_router.core.scripted_host import ( ) MESSAGES: Final = [{"role": "user", "content": "hello there, router"}] -PICKS: Final = 20 +PICKS: Final = 60 # a shuffle misses a deployment in 20 picks about once in a thousand runs MakeCore = Callable[..., RouterCore] @@ -622,3 +622,40 @@ async def test_cooling_down_one_group_leaves_other_groups_pickable(make_core: Ma with pytest.raises(RouterRateLimitError): await core.select("g") assert await _picked_ids(core, "h") <= {"c", "d"} + + +# --- custom strategies --------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_a_registered_strategy_picks_every_deployment_an_endpoint_call_reaches( + make_core: MakeCore, host: ScriptedHost +) -> None: + from litellm.types.router import CustomRoutingStrategyBase + + core: Final = make_core(model_list=[deployment("g", "a"), deployment("g", "b"), deployment("g", "c")]) + seen: Final[list[Mapping[str, object]]] = [] # mutable-ok: the request each pick was asked for + + class HighestId(CustomRoutingStrategyBase): + async def async_get_available_deployment( + self, + model: str, + messages: list[dict[str, str]] | None = None, # mutable-ok: the strategy interface's own type + input: str | list | None = None, # mutable-ok: the strategy interface's own type + specific_deployment: bool | None = False, + request_kwargs: dict | None = None, # mutable-ok: the strategy interface's own type + ) -> dict: # mutable-ok: the strategy interface's own type + seen.append(dict(request_kwargs or {})) + deployments: Final = core.get_model_list(model_name=model) # pyright: ignore[reportAttributeAccessIssue] # both cores serve the Python core's lookup + return max(deployments, key=lambda d: d["model_info"]["id"]) + + def get_available_deployment(self, model: str) -> dict: # mutable-ok: the strategy interface's own type + raise NotImplementedError + + core.register_plugin(HighestId()) + for _ in range(PICKS): + await core.acompletion(model="g", messages=MESSAGES) # pyright: ignore[reportAttributeAccessIssue] # endpoint method + assert set(host.deployments_called("acompletion")) == {"c"} + # The strategy picks from the whole group: no other core narrowed the pick before it. + assert len(seen) == PICKS + assert not any("_retry_skipped_deployment_ids" in request for request in seen) diff --git a/tests/unit/test_router/core/test_selection_characterization_groups.py b/tests/unit/test_router/core/test_selection_characterization_groups.py new file mode 100644 index 00000000000..465118e1c65 --- /dev/null +++ b/tests/unit/test_router/core/test_selection_characterization_groups.py @@ -0,0 +1,747 @@ +""" +Characterization of which deployments a router core considers for a request before a strategy +picks: routing groups and priority routing groups, tag filtering, region filtering, model access +groups, team-only models, default deployments and default fallbacks, paused (blocked) +deployments, pass-through selection, the healthy listing and the request-scoped narrowing keys. +""" + +import datetime +import re +from collections.abc import Mapping +from typing import Final + +import pydantic +import pytest + +import litellm +from litellm.proxy._types import UserAPIKeyAuth +from litellm.router_api import ChangeSet +from litellm.types.router import RouterRateLimitError +from tests.unit.test_router.core.scripted_host import ScriptedHost, deployment, rate_limit +from tests.unit.test_router.core.selection_helpers import MESSAGES, MakeCore, ids_of, picked_ids + +SPREAD_PICKS: Final = 60 + +NO_DEPLOYMENTS_MESSAGE: Final = ( + "No deployments available for selected model, Try again in 5 seconds. Passed model={model}. " + "pre-call-checks={checks}, cooldown_list=[]" +) + +NO_HEALTHY_MESSAGE: Final = ( + "litellm.BadRequestError: You passed in model={model}. There are no healthy deployments for this model" +) + +GROUP_MODELS: Final = [deployment("m1", "a"), deployment("m2", "b"), deployment("m3", "c")] + + +def _group( + name: str, models: list[str], strategy: str, **extra: object +) -> dict[str, object]: # mutable-ok: a routing_groups entry # kwargs-ok: extra RoutingGroup fields + return {"group_name": name, "models": models, "routing_strategy": strategy, **extra} + + +def _priority(name: str, priorities: dict[str, int]) -> dict[str, object]: # mutable-ok: a routing_groups entry + return _group(name, list(priorities), "priority", model_priorities=priorities) + + +# --- routing groups ------------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_a_routing_group_name_is_callable_and_unions_its_members(make_core: MakeCore) -> None: + core: Final = make_core( + model_list=list(GROUP_MODELS), routing_groups=[_group("grp", ["m1", "m2"], "simple-shuffle")] + ) + assert await picked_ids(core, "grp", times=SPREAD_PICKS) == {"a", "b"} + assert ids_of(await core.select("grp", kind="healthy")) == ["a", "b"] + + +@pytest.mark.asyncio +async def test_a_routing_group_call_uses_the_groups_strategy(make_core: MakeCore) -> None: + core: Final = make_core(model_list=list(GROUP_MODELS), routing_groups=[_group("grp", ["m1", "m2"], "least-busy")]) + assert await picked_ids(core, "grp", messages=MESSAGES) == {"a"} + + +@pytest.mark.asyncio +async def test_a_member_called_by_name_uses_its_groups_strategy(make_core: MakeCore) -> None: + core: Final = make_core( + model_list=[deployment("m1", "a"), deployment("m1", "b")], + routing_groups=[_group("grp", ["m1"], "least-busy")], + ) + assert await picked_ids(core, "m1", messages=MESSAGES) == {"a"} + + +@pytest.mark.asyncio +async def test_models_outside_every_group_keep_the_top_level_strategy(make_core: MakeCore) -> None: + core: Final = make_core( + model_list=[deployment("m1", "a"), deployment("free", "x"), deployment("free", "y")], + routing_groups=[_group("grp", ["m1"], "least-busy")], + ) + assert await picked_ids(core, "free", times=SPREAD_PICKS, messages=MESSAGES) == {"x", "y"} + + +@pytest.mark.asyncio +async def test_an_alias_to_a_routing_group_resolves_to_the_group(make_core: MakeCore) -> None: + core: Final = make_core( + model_list=list(GROUP_MODELS), + model_group_alias={"al": "grp"}, + routing_groups=[_group("grp", ["m1", "m2"], "simple-shuffle")], + ) + assert await picked_ids(core, "al", times=SPREAD_PICKS) == {"a", "b"} + + +@pytest.mark.asyncio +async def test_a_group_named_like_a_model_is_shadowed_by_the_model(make_core: MakeCore) -> None: + core: Final = make_core(model_list=list(GROUP_MODELS), routing_groups=[_group("m1", ["m2"], "least-busy")]) + assert await picked_ids(core, "m1") == {"a"} + + +@pytest.mark.asyncio +async def test_a_group_whose_members_have_no_deployments_has_no_healthy_deployments(make_core: MakeCore) -> None: + core: Final = make_core(model_list=list(GROUP_MODELS), routing_groups=[_group("empty", ["zz"], "least-busy")]) + with pytest.raises(litellm.BadRequestError) as raised: + await core.select("empty") + assert str(raised.value) == NO_HEALTHY_MESSAGE.format(model="empty") + + +@pytest.mark.asyncio +async def test_a_routing_group_completion_reports_the_group_as_model_group( + make_core: MakeCore, host: ScriptedHost +) -> None: + core: Final = make_core(model_list=list(GROUP_MODELS), routing_groups=[_group("grp", ["m1"], "simple-shuffle")]) + response: Final = await core.acompletion(model="grp", messages=MESSAGES) # pyright: ignore[reportAttributeAccessIssue] + assert host.deployments_called() == ["a"] + assert host.calls[-1].model_group == "grp" + assert response._hidden_params["additional_headers"]["x-litellm-model-group"] == "grp" + + +def test_routing_group_names_are_listed_among_the_routable_names(make_core: MakeCore) -> None: + core: Final = make_core(model_list=list(GROUP_MODELS), routing_groups=[_group("grp", ["m1", "m2"], "least-busy")]) + # The group name is listed twice today. + assert sorted(core.registry().names(kind="groups")) == ["grp", "grp", "m1", "m2", "m3"] + + +@pytest.mark.asyncio +async def test_routing_groups_applied_at_runtime_become_callable(make_core: MakeCore) -> None: + core: Final = make_core(model_list=list(GROUP_MODELS)) + with pytest.raises(litellm.BadRequestError): + await core.select("grp") + core.apply(ChangeSet(settings={"routing_groups": [_group("grp", ["m1", "m2"], "least-busy")]})) + assert await picked_ids(core, "grp", messages=MESSAGES) == {"a"} + + +@pytest.mark.parametrize( + ("groups", "message"), + [ + ( + [_group("x", ["m1"], "least-busy"), _group("x", ["m2"], "least-busy")], + "routing_groups: group names must be unique, duplicate group_name 'x'.", + ), + ( + [_group("default", ["m1"], "least-busy")], + "routing_groups: 'default' is reserved for the implicit fallback group.", + ), + ([_group("", ["m1"], "least-busy")], "routing_groups: group_name must be non-empty."), + ( + [_group("x", ["m1"], "least-busy"), _group("y", ["m1"], "simple-shuffle")], + "routing_groups: model_name 'm1' appears in 'x' and 'y'. Each model may belong to at most one group.", + ), + ([_group("x", ["m1"], "nope")], "Invalid routing_strategy: 'nope'."), + ([_priority("m1", {"m2": 1})], "Priority routing group names must not shadow a model or alias"), + ], + ids=["duplicate", "reserved-default", "empty-name", "overlap", "bad-strategy", "priority-shadow"], +) +def test_invalid_routing_groups_are_rejected_at_construction( + make_core: MakeCore, groups: list[dict[str, object]], message: str +) -> None: + with pytest.raises(ValueError, match=re.escape(message)): + make_core(model_list=list(GROUP_MODELS), routing_groups=groups) + + +@pytest.mark.parametrize( + ("group", "message"), + [ + (_group("p", ["m1", "m2"], "priority", model_priorities={"m1": 1}), "model_priorities must contain exactly"), + ( + _group("p", ["m1"], "priority", model_priorities={"m1": 1}, routing_strategy_args={"x": 1}), + "Priority routing groups use model_priorities, not routing_strategy_args", + ), + (_group("p", ["m1"], "least-busy", model_priorities={"m1": 1}), "model_priorities requires routing_strategy"), + (_group("p", ["m1"], "priority", model_priorities={"m1": 0}), "greater than or equal to 1"), + ( + _group("p", ["m1", "m1"], "priority", model_priorities={"m1": 1}), + "Priority routing groups require nonempty, distinct models", + ), + ], + ids=["missing-priority", "strategy-args", "priorities-without-priority", "zero-priority", "duplicate-member"], +) +def test_invalid_priority_groups_fail_validation_at_construction( + make_core: MakeCore, group: dict[str, object], message: str +) -> None: + with pytest.raises(pydantic.ValidationError, match=re.escape(message)): + make_core(model_list=list(GROUP_MODELS), routing_groups=[group]) + + +# --- priority routing groups --------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_a_priority_group_picks_only_its_lowest_priority_members(make_core: MakeCore) -> None: + core: Final = make_core( + model_list=[*GROUP_MODELS, deployment("m2", "b2")], + routing_groups=[_priority("p", {"m1": 2, "m2": 1, "m3": 2})], + ) + assert await picked_ids(core, "p", times=SPREAD_PICKS) == {"b", "b2"} + assert sorted(ids_of(await core.select("p", kind="healthy"))) == ["b", "b2"] + + +@pytest.mark.asyncio +async def test_a_priority_group_pick_carries_the_member_priority_as_order(make_core: MakeCore) -> None: + core: Final = make_core(model_list=list(GROUP_MODELS), routing_groups=[_priority("p", {"m1": 3, "m2": 1})]) + chosen: Final = await core.select("p") + assert chosen["model_info"]["id"] == "b" # pyright: ignore[reportIndexIssue] + assert chosen["model_name"] == "m2" # pyright: ignore[reportIndexIssue] + assert chosen["litellm_params"]["order"] == 1 # pyright: ignore[reportIndexIssue] + + +@pytest.mark.asyncio +async def test_a_priority_group_member_called_directly_ignores_the_priority(make_core: MakeCore) -> None: + core: Final = make_core(model_list=list(GROUP_MODELS), routing_groups=[_priority("p", {"m1": 3, "m2": 1})]) + chosen: Final = await core.select("m1") + assert chosen["model_info"]["id"] == "a" # pyright: ignore[reportIndexIssue] + assert "order" not in chosen["litellm_params"] # pyright: ignore[reportIndexIssue] + + +@pytest.mark.asyncio +async def test_a_priority_group_fails_over_to_the_next_priority_and_stays_there_while_the_first_cools_down( + make_core: MakeCore, host: ScriptedHost +) -> None: + core: Final = make_core( + model_list=list(GROUP_MODELS), + routing_groups=[_priority("p", {"m1": 1, "m2": 2})], + allowed_fails=0, + num_retries=0, + cooldown_time=60, + ) + host.answer("openai/a", rate_limit()) + await core.acompletion(model="p", messages=MESSAGES) # pyright: ignore[reportAttributeAccessIssue] + assert host.deployments_called() == ["a", "b"] + assert await picked_ids(core, "p") == {"b"} + + +# --- tag filtering ----------------------------------------------------------------------------------- + +TAGGED: Final = [ + deployment("g", "free", tags=["free"]), + deployment("g", "paid", tags=["paid"]), + deployment("g", "dflt", tags=["default"]), + deployment("g", "untagged"), +] + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("tags", "expected"), + [(["free"], {"free"}), (["paid"], {"paid"}), (["free", "paid"], {"free", "paid"})], +) +async def test_tag_filtering_picks_deployments_carrying_a_request_tag( + make_core: MakeCore, tags: list[str], expected: set[str] +) -> None: + core: Final = make_core(model_list=list(TAGGED), enable_tag_filtering=True) + assert await picked_ids(core, "g", times=SPREAD_PICKS, request_kwargs={"metadata": {"tags": tags}}) == expected + + +@pytest.mark.asyncio +@pytest.mark.parametrize("request_kwargs", [{}, {"metadata": {"tags": ["unknown"]}}], ids=["untagged", "unmatched"]) +async def test_tag_filtering_sends_untagged_or_unmatched_requests_to_default_tagged_deployments( + make_core: MakeCore, request_kwargs: dict[str, object] +) -> None: + core: Final = make_core(model_list=list(TAGGED), enable_tag_filtering=True) + assert await picked_ids(core, "g", request_kwargs=request_kwargs) == {"dflt"} + + +@pytest.mark.asyncio +async def test_tag_filtering_without_a_default_rejects_an_unmatched_tag(make_core: MakeCore) -> None: + core: Final = make_core(model_list=TAGGED[:2], enable_tag_filtering=True) + with pytest.raises(ValueError, match="Not allowed to access model") as raised: + await core.select("g", {"metadata": {"tags": ["unknown"]}}) + assert str(raised.value) == ( + "Not allowed to access model due to tags configuration. Passed model=g and tags=['unknown']" + ) + + +@pytest.mark.asyncio +async def test_tag_filtering_without_a_default_serves_an_untagged_request_from_every_deployment( + make_core: MakeCore, +) -> None: + core: Final = make_core(model_list=TAGGED[:2], enable_tag_filtering=True) + assert await picked_ids(core, "g", times=SPREAD_PICKS) == {"free", "paid"} + + +@pytest.mark.asyncio +async def test_request_tags_are_ignored_without_tag_filtering(make_core: MakeCore) -> None: + core: Final = make_core(model_list=list(TAGGED)) + picked: Final = await picked_ids(core, "g", times=SPREAD_PICKS, request_kwargs={"metadata": {"tags": ["free"]}}) + assert picked == {"free", "paid", "dflt", "untagged"} + + +@pytest.mark.asyncio +async def test_a_request_can_enable_tag_filtering_for_itself(make_core: MakeCore) -> None: + core: Final = make_core(model_list=list(TAGGED)) + request: Final = {"metadata": {"tags": ["free"]}, "enable_tag_filtering": True} + assert await picked_ids(core, "g", request_kwargs=request) == {"free"} + + +@pytest.mark.asyncio +@pytest.mark.parametrize(("match_any", "expected"), [(True, {"ab", "a"}), (False, {"ab"})]) +async def test_tag_filtering_match_any_decides_whether_one_shared_tag_is_enough( + make_core: MakeCore, match_any: bool, expected: set[str] +) -> None: + core: Final = make_core( + model_list=[deployment("g", "ab", tags=["a", "b"]), deployment("g", "a", tags=["a"])], + enable_tag_filtering=True, + tag_filtering_match_any=match_any, + ) + picked: Final = await picked_ids(core, "g", times=SPREAD_PICKS, request_kwargs={"metadata": {"tags": ["a", "b"]}}) + assert picked == expected + + +@pytest.mark.asyncio +async def test_tags_in_litellm_metadata_are_read_too(make_core: MakeCore) -> None: + core: Final = make_core(model_list=list(TAGGED), enable_tag_filtering=True) + assert await picked_ids(core, "g", request_kwargs={"litellm_metadata": {"tags": ["paid"]}}) == {"paid"} + + +@pytest.mark.asyncio +async def test_completion_with_tags_calls_the_tagged_deployment(make_core: MakeCore, host: ScriptedHost) -> None: + core: Final = make_core(model_list=list(TAGGED), enable_tag_filtering=True) + for _ in range(5): + await core.acompletion(model="g", messages=MESSAGES, metadata={"tags": ["paid"]}) # pyright: ignore[reportAttributeAccessIssue] + assert host.deployments_called() == ["paid"] * 5 + + +# --- region filtering -------------------------------------------------------------------------------- + +REGIONAL: Final = [ + deployment("g", "eu", region_name="eu"), + deployment("g", "us", region_name="us"), + deployment("g", "unknown"), +] + + +@pytest.mark.asyncio +async def test_allowed_model_region_keeps_only_deployments_in_that_region(make_core: MakeCore) -> None: + core: Final = make_core(model_list=list(REGIONAL), enable_pre_call_checks=True) + assert await picked_ids(core, "g", request_kwargs={"allowed_model_region": "eu"}, messages=MESSAGES) == {"eu"} + + +@pytest.mark.asyncio +async def test_allowed_model_region_is_ignored_without_request_input(make_core: MakeCore) -> None: + core: Final = make_core(model_list=list(REGIONAL), enable_pre_call_checks=True) + picked: Final = await picked_ids(core, "g", times=SPREAD_PICKS, request_kwargs={"allowed_model_region": "eu"}) + assert picked == {"eu", "us", "unknown"} + + +@pytest.mark.asyncio +async def test_allowed_model_region_is_ignored_without_pre_call_checks(make_core: MakeCore) -> None: + core: Final = make_core(model_list=list(REGIONAL)) + picked: Final = await picked_ids( + core, "g", times=SPREAD_PICKS, request_kwargs={"allowed_model_region": "eu"}, messages=MESSAGES + ) + assert picked == {"eu", "us", "unknown"} + + +@pytest.mark.asyncio +async def test_an_allowed_model_region_no_deployment_is_in_rate_limits(make_core: MakeCore) -> None: + core: Final = make_core(model_list=list(REGIONAL), enable_pre_call_checks=True) + with pytest.raises(RouterRateLimitError) as raised: + await core.select("g", {"allowed_model_region": "ap"}, messages=MESSAGES) # pyright: ignore[reportCallIssue] + assert str(raised.value) == NO_DEPLOYMENTS_MESSAGE.format(model="g", checks=True) + + +# --- model access groups ------------------------------------------------------------------------------- + +ACCESS_GROUPED: Final = [ + deployment("g", "in_a", model_info={"access_groups": ["grpA"]}), + deployment("g", "in_b", model_info={"access_groups": ["grpB"]}), +] + + +def _auth( + bucket: str = "metadata", **fields: object +) -> dict[str, object]: # mutable-ok: request kwargs # kwargs-ok: UserAPIKeyAuth fields + return {bucket: {"user_api_key_auth": UserAPIKeyAuth(**fields)}} # pyright: ignore[reportArgumentType] + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("request_kwargs", "expected"), + [ + (_auth(models=["grpA"]), {"in_a"}), + (_auth(models=["grpB"]), {"in_b"}), + (_auth(team_models=["grpB"]), {"in_b"}), + (_auth("litellm_metadata", models=["grpA"]), {"in_a"}), + ], + ids=["key-group-a", "key-group-b", "team-group-b", "litellm-metadata"], +) +async def test_an_access_group_only_caller_is_narrowed_to_its_groups_deployments( + make_core: MakeCore, request_kwargs: dict[str, object], expected: set[str] +) -> None: + core: Final = make_core(model_list=list(ACCESS_GROUPED)) + assert await picked_ids(core, "g", request_kwargs=request_kwargs) == expected + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "models", [["g"], ["*"], ["all-proxy-models"], [], ["unrelated"]], ids=["direct", "star", "all", "none", "other"] +) +async def test_a_caller_not_limited_to_access_groups_can_reach_every_deployment( + make_core: MakeCore, models: list[str] +) -> None: + core: Final = make_core(model_list=list(ACCESS_GROUPED)) + picked: Final = await picked_ids(core, "g", times=SPREAD_PICKS, request_kwargs=_auth(models=models)) + assert picked == {"in_a", "in_b"} + + +# --- team-only models ------------------------------------------------------------------------------------ + +TEAM_ONLY: Final = [ + deployment("internal-x", "tx", model_info={"team_id": "team-x", "team_public_model_name": "pub"}), + deployment("internal-y", "ty", model_info={"team_id": "team-y", "team_public_model_name": "pub"}), +] + + +@pytest.mark.asyncio +async def test_a_team_id_in_litellm_metadata_selects_the_team_deployment(make_core: MakeCore) -> None: + core: Final = make_core(model_list=list(TEAM_ONLY)) + assert await picked_ids(core, "pub", request_kwargs={"litellm_metadata": {"user_api_key_team_id": "team-y"}}) == { + "ty" + } + + +@pytest.mark.asyncio +async def test_a_proxy_admin_without_a_team_is_refused_a_public_name_several_teams_share(make_core: MakeCore) -> None: + core: Final = make_core(model_list=list(TEAM_ONLY)) + with pytest.raises(litellm.BadRequestError) as raised: + await core.select("pub", _auth(user_role="proxy_admin")) + assert str(raised.value) == ( + "litellm.BadRequestError: Model name 'pub' matches deployments from multiple teams. " + "Specify the deployment ID directly to disambiguate." + ) + + +@pytest.mark.asyncio +async def test_a_proxy_admin_without_a_team_finds_one_teams_public_name_but_cannot_use_it(make_core: MakeCore) -> None: + core: Final = make_core(model_list=TEAM_ONLY[:1]) + with pytest.raises(RouterRateLimitError) as raised: + await core.select("pub", _auth(user_role="proxy_admin")) + assert str(raised.value) == NO_DEPLOYMENTS_MESSAGE.format(model="pub", checks=False) + + +# --- default deployment, default fallbacks ------------------------------------------------------------------ + + +@pytest.mark.asyncio +async def test_an_unknown_model_goes_to_a_catch_all_deployment_under_its_own_name(make_core: MakeCore) -> None: + core: Final = make_core(model_list=[deployment("g", "a"), deployment("*", "catch_all", model="openai/*")]) + chosen: Final = await core.select("brand-new") + assert chosen["model_info"]["id"] == "catch_all" # pyright: ignore[reportIndexIssue] + assert chosen["litellm_params"]["model"] == "openai/brand-new" # pyright: ignore[reportIndexIssue] + + +@pytest.mark.asyncio +async def test_an_unknown_model_goes_to_the_first_default_fallback(make_core: MakeCore, host: ScriptedHost) -> None: + core: Final = make_core(model_list=[deployment("g", "a"), deployment("h", "b")], default_fallbacks=["h", "g"]) + chosen: Final = await core.select("unknown") + assert chosen["model_info"]["id"] == "b" # pyright: ignore[reportIndexIssue] + await core.acompletion(model="unknown", messages=MESSAGES) # pyright: ignore[reportAttributeAccessIssue] + assert host.deployments_called() == ["b"] + + +@pytest.mark.asyncio +async def test_an_unknown_model_goes_to_the_process_default_fallback(make_core: MakeCore, host: ScriptedHost) -> None: + host.set_dynamic(default_fallbacks=["g"]) + core: Final = make_core(model_list=[deployment("g", "a")]) + chosen: Final = await core.select("unknown") + assert chosen["model_info"]["id"] == "a" # pyright: ignore[reportIndexIssue] + + +# --- paused (blocked) deployments ------------------------------------------------------------------------------ + + +@pytest.mark.asyncio +async def test_a_blocked_deployment_is_never_picked_for_its_group(make_core: MakeCore) -> None: + core: Final = make_core(model_list=[deployment("g", "a", model_info={"blocked": True}), deployment("g", "b")]) + assert await picked_ids(core, "g") == {"b"} + assert ids_of(await core.select("g", kind="healthy")) == ["b"] + + +@pytest.mark.asyncio +async def test_selecting_a_blocked_deployment_by_id_raises_service_unavailable(make_core: MakeCore) -> None: + core: Final = make_core(model_list=[deployment("g", "a", model_info={"blocked": True}), deployment("g", "b")]) + with pytest.raises(litellm.ServiceUnavailableError) as raised: + await core.select("a") + # The message names the deployment's litellm model, not the id that was asked for. + assert str(raised.value) == ( + "litellm.ServiceUnavailableError: Model 'openai/a' is currently paused and cannot accept requests." + ) + assert raised.value.status_code == 503 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("entry", ["select", "acompletion", "completion"]) +async def test_a_group_whose_every_deployment_is_blocked_rate_limits( + make_core: MakeCore, host: ScriptedHost, entry: str +) -> None: + core: Final = make_core(model_list=[deployment("g", "a", model_info={"blocked": True})]) + call: Final = core.select("g") if entry == "select" else core.call(entry, {"model": "g", "messages": MESSAGES}) + with pytest.raises(RouterRateLimitError) as raised: + await call + assert str(raised.value) == NO_DEPLOYMENTS_MESSAGE.format(model="g", checks=False) + assert host.calls == [] + + +# --- pass-through selection ------------------------------------------------------------------------------------ + +PASS_THROUGH: Final = [deployment("g", "pt", use_in_pass_through=True), deployment("g", "plain")] + + +@pytest.mark.asyncio +async def test_pass_through_selection_picks_only_pass_through_deployments(make_core: MakeCore) -> None: + core: Final = make_core(model_list=list(PASS_THROUGH)) + assert await picked_ids(core, "g", kind="pass_through") == {"pt"} + + +@pytest.mark.asyncio +async def test_pass_through_selection_by_id_returns_a_pass_through_deployment(make_core: MakeCore) -> None: + core: Final = make_core(model_list=list(PASS_THROUGH)) + chosen: Final = await core.select("pt", kind="pass_through") + assert chosen["model_info"]["id"] == "pt" # pyright: ignore[reportIndexIssue] + + +@pytest.mark.asyncio +async def test_pass_through_selection_by_id_refuses_a_plain_deployment(make_core: MakeCore) -> None: + core: Final = make_core(model_list=list(PASS_THROUGH)) + with pytest.raises(litellm.BadRequestError) as raised: + await core.select("plain", kind="pass_through") + assert str(raised.value) == ( + "litellm.BadRequestError: Deployment plain does not support pass-through endpoint (use_in_pass_through=False)" + ) + + +@pytest.mark.asyncio +async def test_pass_through_selection_without_pass_through_deployments_raises_bad_request(make_core: MakeCore) -> None: + core: Final = make_core(model_list=[deployment("g", "a")]) + with pytest.raises(litellm.BadRequestError) as raised: + await core.select("g", kind="pass_through") + assert str(raised.value) == ( + "litellm.BadRequestError: Model g has no deployment configured with use_in_pass_through=True. " + "Please add use_in_pass_through: true in the deployment configuration" + ) + + +@pytest.mark.asyncio +async def test_async_pass_through_selection_words_the_missing_deployment_error_differently( + make_core: MakeCore, +) -> None: + core: Final = make_core(model_list=[deployment("g", "a")]) + with pytest.raises(litellm.BadRequestError) as raised: + await core.async_get_available_deployment_for_pass_through("g", {}) # pyright: ignore[reportAttributeAccessIssue] + assert str(raised.value) == ( + "litellm.BadRequestError: Model g has no deployments configured with use_in_pass_through=True. " + "Please add use_in_pass_through: true to the deployment configuration" + ) + + +@pytest.mark.asyncio +async def test_pass_through_selection_of_an_unknown_model_has_no_healthy_deployments(make_core: MakeCore) -> None: + core: Final = make_core(model_list=list(PASS_THROUGH)) + with pytest.raises(litellm.BadRequestError) as raised: + await core.select("nope", kind="pass_through") + assert str(raised.value) == NO_HEALTHY_MESSAGE.format(model="nope") + + +@pytest.mark.asyncio +async def test_pass_through_selection_skips_a_cooling_pass_through_deployment( + make_core: MakeCore, host: ScriptedHost +) -> None: + core: Final = make_core( + model_list=[ + deployment("g", "pt1", use_in_pass_through=True), + deployment("g", "pt2", use_in_pass_through=True), + ], + allowed_fails=0, + num_retries=0, + cooldown_time=60, + ) + host.answer("openai/pt1", rate_limit()) + with pytest.raises(litellm.RateLimitError): + await core.acompletion(model="pt1", messages=MESSAGES) # pyright: ignore[reportAttributeAccessIssue] + assert await picked_ids(core, "g", kind="pass_through") == {"pt2"} + + +def test_deprecated_pass_through_pick_matches_select(make_core: MakeCore) -> None: + core: Final = make_core(model_list=list(PASS_THROUGH)) + assert core.get_available_deployment_for_pass_through("g")["model_info"]["id"] == "pt" # pyright: ignore[reportAttributeAccessIssue] + + +# --- the healthy listing -------------------------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_the_healthy_listing_keeps_only_the_lowest_order(make_core: MakeCore) -> None: + core: Final = make_core(model_list=[deployment("g", "a", order=1), deployment("g", "b", order=2)]) + assert ids_of(await core.select("g", kind="healthy")) == ["a"] + + +@pytest.mark.asyncio +async def test_the_healthy_listing_for_a_deployment_id_is_that_single_deployment(make_core: MakeCore) -> None: + core: Final = make_core(model_list=[deployment("g", "a"), deployment("g", "b")]) + healthy: Final = await core.select("b", kind="healthy") + assert isinstance(healthy, Mapping) + assert healthy["model_info"]["id"] == "b" + + +@pytest.mark.asyncio +async def test_the_healthy_listing_of_an_unknown_model_raises_bad_request(make_core: MakeCore) -> None: + core: Final = make_core(model_list=[deployment("g", "a")]) + with pytest.raises(litellm.BadRequestError) as raised: + await core.select("nope", kind="healthy") + assert str(raised.value) == NO_HEALTHY_MESSAGE.format(model="nope") + + +@pytest.mark.asyncio +async def test_the_healthy_listing_leaves_out_cooling_deployments(make_core: MakeCore, host: ScriptedHost) -> None: + core: Final = make_core( + model_list=[deployment("g", "a"), deployment("g", "b")], allowed_fails=0, num_retries=0, cooldown_time=60 + ) + host.answer("openai/a", rate_limit()) + with pytest.raises(litellm.RateLimitError): + await core.acompletion(model="a", messages=MESSAGES) # pyright: ignore[reportAttributeAccessIssue] + assert ids_of(await core.select("g", kind="healthy")) == ["b"] + + +# --- request-scoped narrowing keys ----------------------------------------------------------------------------------- + +THREE: Final = [deployment("g", "a"), deployment("g", "b"), deployment("g", "c")] + + +@pytest.mark.asyncio +async def test_target_order_picks_that_order_instead_of_the_lowest(make_core: MakeCore) -> None: + core: Final = make_core(model_list=[deployment("g", "a", order=1), deployment("g", "b", order=2)]) + assert await picked_ids(core, "g", request_kwargs={"_target_order": 2}) == {"b"} + assert ids_of(await core.select("g", {"_target_order": 2}, kind="healthy")) == ["b"] + + +@pytest.mark.asyncio +async def test_a_target_order_no_deployment_has_rate_limits(make_core: MakeCore) -> None: + core: Final = make_core(model_list=[deployment("g", "a", order=1)]) + with pytest.raises(RouterRateLimitError) as raised: + await core.select("g", {"_target_order": 3}) + assert str(raised.value) == NO_DEPLOYMENTS_MESSAGE.format(model="g", checks=False) + + +@pytest.mark.asyncio +async def test_excluded_deployment_ids_are_never_picked(make_core: MakeCore) -> None: + core: Final = make_core(model_list=list(THREE)) + assert await picked_ids(core, "g", request_kwargs={"_excluded_deployment_ids": ["a", "b"]}) == {"c"} + + +@pytest.mark.asyncio +async def test_excluding_every_deployment_rate_limits(make_core: MakeCore) -> None: + core: Final = make_core(model_list=list(THREE)) + with pytest.raises(RouterRateLimitError) as raised: + await core.select("g", {"_excluded_deployment_ids": ["a", "b", "c"]}) + assert str(raised.value) == NO_DEPLOYMENTS_MESSAGE.format(model="g", checks=False) + + +@pytest.mark.asyncio +async def test_retry_skipped_deployment_ids_are_not_picked_while_others_remain(make_core: MakeCore) -> None: + core: Final = make_core(model_list=list(THREE)) + assert await picked_ids(core, "g", request_kwargs={"_retry_skipped_deployment_ids": ("a", "b")}) == {"c"} + + +@pytest.mark.asyncio +async def test_skipping_every_deployment_on_retry_keeps_them_all(make_core: MakeCore) -> None: + core: Final = make_core(model_list=list(THREE)) + healthy: Final = await core.select("g", {"_retry_skipped_deployment_ids": ("a", "b", "c")}, kind="healthy") + assert ids_of(healthy) == ["a", "b", "c"] + + +@pytest.mark.asyncio +async def test_retry_skipped_deployment_ids_given_as_a_list_are_ignored(make_core: MakeCore) -> None: + core: Final = make_core(model_list=list(THREE)) + healthy: Final = await core.select("g", {"_retry_skipped_deployment_ids": ["a", "b"]}, kind="healthy") + assert ids_of(healthy) == ["a", "b", "c"] + + +@pytest.mark.asyncio +async def test_select_pops_the_narrowing_keys_from_the_request(make_core: MakeCore) -> None: + core: Final = make_core(model_list=list(THREE)) + request: Final[dict[str, object]] = { + "_target_order": None, + "_excluded_deployment_ids": ["a"], + "_retry_skipped_deployment_ids": ("b",), + } + assert ids_of(await core.select("g", request, kind="healthy")) == ["c"] + assert request == {} + + +# --- access windows (reserved deployments) ---------------------------------------------------------------------------- + + +def _window_for(team_ids: list[str]) -> dict[str, object]: # mutable-ok: an access_windows entry + now: Final = datetime.datetime.now(datetime.timezone.utc) + return { + "start": (now - datetime.timedelta(hours=1)).strftime("%H:%M"), + "end": (now + datetime.timedelta(hours=1)).strftime("%H:%M"), + "timezone": "UTC", + "team_ids": team_ids, + } + + +def _reserved(window: dict[str, object]) -> dict[str, object]: # mutable-ok: a model_list entry + return deployment("g", "reserved", model_info={"access_windows": [window]}) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("request_kwargs", [{"metadata": {"user_api_key_team_id": "team-b"}}, {}]) +async def test_a_deployment_reserved_for_another_team_is_hidden_during_its_window( + make_core: MakeCore, request_kwargs: dict[str, object] +) -> None: + core: Final = make_core(model_list=[_reserved(_window_for(["team-a"])), deployment("g", "open")]) + assert await picked_ids(core, "g", request_kwargs=request_kwargs) == {"open"} + + +@pytest.mark.asyncio +async def test_the_reserving_team_reaches_its_reserved_deployment(make_core: MakeCore) -> None: + core: Final = make_core(model_list=[_reserved(_window_for(["team-a"]))]) + assert await picked_ids(core, "g", request_kwargs={"metadata": {"user_api_key_team_id": "team-a"}}) == {"reserved"} + + +@pytest.mark.asyncio +@pytest.mark.parametrize(("model", "named"), [("g", "g"), ("reserved", "openai/reserved")], ids=["group", "id"]) +async def test_only_reserved_deployments_left_raises_bad_request_naming_the_window_end( + make_core: MakeCore, model: str, named: str +) -> None: + window: Final = _window_for(["team-a"]) + core: Final = make_core(model_list=[_reserved(window)]) + with pytest.raises(litellm.BadRequestError) as raised: + await core.select(model, {"metadata": {"user_api_key_team_id": "team-b"}}) + assert str(raised.value) == ( + f"litellm.BadRequestError: Deployment {named} is reserved for another team until {window['end']} UTC" + ) + + +# --- specific_deployment ------------------------------------------------------------------------------------------------ + + +@pytest.mark.asyncio +@pytest.mark.parametrize("model", ["openai/unknown", "g"], ids=["unknown-litellm-model", "group-name"]) +async def test_specific_deployment_for_a_name_no_litellm_model_matches_rate_limits( + make_core: MakeCore, model: str +) -> None: + core: Final = make_core(model_list=[deployment("g", "a")]) + with pytest.raises(RouterRateLimitError) as raised: + await core.select(model, specific_deployment=True) # pyright: ignore[reportCallIssue] + assert str(raised.value) == NO_DEPLOYMENTS_MESSAGE.format(model=model, checks=False) diff --git a/tests/unit/test_router/core/test_selection_characterization_prerouting.py b/tests/unit/test_router/core/test_selection_characterization_prerouting.py new file mode 100644 index 00000000000..317fbf4be46 --- /dev/null +++ b/tests/unit/test_router/core/test_selection_characterization_prerouting.py @@ -0,0 +1,505 @@ +""" +Characterization of what runs before a router core's strategy picks: pre-routing strategies (the +complexity router, chosen among tag-scoped markers), the strategy marker deployments they are +registered from, tier ``litellm_params`` pinned onto the request, alias-marker params forwarded to +the routed call, the Claude Code session router binding, and the ``Router(plugins=[...])`` +routing-plugin pipeline. +""" + +import asyncio +import re +from collections.abc import Mapping +from typing import Final + +import pytest + +import litellm +from litellm.router_core import RouterCore +from litellm.types.router import RoutingContext +from tests.unit.test_router.core.scripted_host import ScriptedHost, deployment +from tests.unit.test_router.core.selection_helpers import MESSAGES, MakeCore, picked_ids + +HARD_MESSAGES: Final = [ + { + "role": "user", + "content": ( + "Design a distributed consensus algorithm, prove its correctness step by step, analyze the time " + "complexity, and implement it in Rust with a kubernetes deployment and database sharding strategy. " + "Think step by step and reason about edge cases." + ), + } +] + +TIERS: Final = {"SIMPLE": "small", "MEDIUM": "big", "COMPLEX": "big", "REASONING": "big"} +ALL_BIG: Final = {"SIMPLE": "big", "MEDIUM": "big", "COMPLEX": "big", "REASONING": "big"} + +ONLY_MARKERS_MESSAGE: Final = ( + "litellm.BadRequestError: You passed in model={model}. Every deployment for it is a strategy router marker " + "(auto_router/...), which is not a callable model, and no pre-routing strategy selected a deployment for " + "this request" +) + +SESSION_ID: Final = "0f8fad5b-d9cb-469f-a165-70867728950e" + + +def _complexity_router( + name: str = "smart", + tiers: Mapping[str, object] = TIERS, + deployment_id: str = "marker", + **params: object, # kwargs-ok: extra marker litellm_params +) -> dict[str, object]: # mutable-ok: a model_list entry + return { + "model_name": name, + "litellm_params": { + "model": "auto_router/complexity_router", + "complexity_router_config": {"tiers": dict(tiers)}, + "complexity_router_default_model": "big", + **params, + }, + "model_info": {"id": deployment_id}, + } + + +def _tiered( + **marker_params: object, +) -> list[dict[str, object]]: # mutable-ok: a model_list # kwargs-ok: marker litellm_params + return [deployment("small", "s"), deployment("big", "b"), _complexity_router(**marker_params)] + + +def _sync_entry(core: RouterCore, entry: str, model: str) -> object: + if entry == "completion": + return core.completion(model=model, messages=MESSAGES) # pyright: ignore[reportAttributeAccessIssue] + return core.get_available_deployment(model, messages=MESSAGES) # pyright: ignore[reportAttributeAccessIssue] + + +def _metadata(request: Mapping[str, object]) -> Mapping[str, object]: + metadata: Final = request.get("metadata") + assert isinstance(metadata, Mapping) + return metadata + + +# --- the complexity router ------------------------------------------------------------------------ + + +@pytest.mark.asyncio +@pytest.mark.parametrize(("messages", "expected"), [(MESSAGES, "s"), (HARD_MESSAGES, "b")], ids=["simple", "hard"]) +async def test_a_complexity_router_picks_the_tier_deployment_for_the_prompt( + make_core: MakeCore, messages: list[dict[str, str]], expected: str +) -> None: + core: Final = make_core(model_list=_tiered()) + chosen: Final = await core.select("smart", {}, messages=messages) # pyright: ignore[reportCallIssue] + assert chosen["model_info"]["id"] == expected # pyright: ignore[reportIndexIssue] + + +@pytest.mark.asyncio +async def test_a_complexity_router_pick_stamps_the_routing_decision_on_the_request(make_core: MakeCore) -> None: + core: Final = make_core(model_list=_tiered()) + request: Final[dict[str, object]] = {} + await core.select("smart", request, messages=MESSAGES) # pyright: ignore[reportCallIssue] + metadata: Final = _metadata(request) + decision: Final = metadata["routing_decision"] + assert isinstance(decision, Mapping) + assert decision["router_model_name"] == "smart" + assert decision["router_type"] == "complexity" + assert decision["router_config_id"] == "marker" + assert decision["routed_model"] == "small" + assert decision["tier"] == "SIMPLE" + assert metadata["pre_routing_selected_model"] == "small" + + +@pytest.mark.asyncio +async def test_completion_through_a_complexity_router_reports_the_router_name_as_model_group( + make_core: MakeCore, host: ScriptedHost +) -> None: + core: Final = make_core(model_list=_tiered()) + await core.acompletion(model="smart", messages=MESSAGES) # pyright: ignore[reportAttributeAccessIssue] + assert host.deployments_called() == ["s"] + assert host.calls[-1].model == "openai/s" + assert host.calls[-1].model_group == "smart" + + +@pytest.mark.asyncio +async def test_an_alias_to_a_complexity_router_routes_through_it(make_core: MakeCore, host: ScriptedHost) -> None: + core: Final = make_core(model_list=_tiered(), model_group_alias={"al": "smart"}) + request: Final[dict[str, object]] = {} + chosen: Final = await core.select("al", request, messages=MESSAGES) # pyright: ignore[reportCallIssue] + assert chosen["model_info"]["id"] == "s" # pyright: ignore[reportIndexIssue] + assert _metadata(request)["routing_decision"]["router_model_name"] == "smart" # pyright: ignore[reportIndexIssue] + await core.acompletion(model="al", messages=MESSAGES) # pyright: ignore[reportAttributeAccessIssue] + assert host.calls[-1].model_group == "al" + + +@pytest.mark.asyncio +async def test_a_complexity_router_takes_precedence_over_plain_deployments_of_its_name(make_core: MakeCore) -> None: + core: Final = make_core(model_list=[*_tiered(), deployment("smart", "plain")]) + assert await picked_ids(core, "smart", times=10, messages=MESSAGES) == {"s"} + + +@pytest.mark.asyncio +async def test_a_complexity_router_tier_naming_an_unknown_model_falls_back_to_the_default_model( + make_core: MakeCore, +) -> None: + core: Final = make_core( + model_list=[deployment("big", "b"), _complexity_router(tiers={**TIERS, "SIMPLE": "missing"})] + ) + request: Final[dict[str, object]] = {} + chosen: Final = await core.select("smart", request, messages=MESSAGES) # pyright: ignore[reportCallIssue] + assert chosen["model_info"]["id"] == "b" # pyright: ignore[reportIndexIssue] + decision: Final = _metadata(request)["routing_decision"] + assert isinstance(decision, Mapping) + assert decision["cause"] == "health_default_fallback" + assert decision["routed_model"] == "big" + + +def test_a_complexity_router_without_tiers_or_a_default_model_is_rejected(make_core: MakeCore) -> None: + with pytest.raises(ValueError, match="complexity_router_default_model is required for complexity-router"): + make_core(model_list=[deployment("smart", "marker", model="auto_router/complexity_router")]) + + +# --- strategy markers are never callable ---------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_the_healthy_listing_of_a_router_name_refuses_the_marker(make_core: MakeCore) -> None: + core: Final = make_core(model_list=_tiered()) + with pytest.raises(litellm.BadRequestError) as raised: + await core.select("smart", {}, kind="healthy") + assert str(raised.value) == ONLY_MARKERS_MESSAGE.format(model="smart") + + +@pytest.mark.asyncio +async def test_pass_through_selection_of_a_router_name_refuses_the_marker(make_core: MakeCore) -> None: + core: Final = make_core(model_list=_tiered()) + with pytest.raises(litellm.BadRequestError) as raised: + await core.select("smart", {}, kind="pass_through", messages=MESSAGES) # pyright: ignore[reportCallIssue] + assert str(raised.value) == ONLY_MARKERS_MESSAGE.format(model="smart") + + +@pytest.mark.asyncio +async def test_selecting_a_marker_by_its_id_refuses_it_naming_its_litellm_model(make_core: MakeCore) -> None: + core: Final = make_core(model_list=_tiered()) + with pytest.raises(litellm.BadRequestError) as raised: + await core.select("marker", {}, messages=MESSAGES) # pyright: ignore[reportCallIssue] + assert str(raised.value) == ONLY_MARKERS_MESSAGE.format(model="auto_router/complexity_router") + + +@pytest.mark.parametrize("entry", ["get_available_deployment", "completion"]) +def test_the_sync_path_never_runs_pre_routing_and_refuses_the_marker( + make_core: MakeCore, host: ScriptedHost, entry: str +) -> None: + core: Final = make_core(model_list=_tiered()) + with pytest.raises(litellm.BadRequestError) as raised: + _sync_entry(core, entry, "smart") + assert str(raised.value) == ONLY_MARKERS_MESSAGE.format(model="smart") + assert host.calls == [] + + +# --- choosing among tag-scoped markers ----------------------------------------------------------------- + +TAG_SCOPED: Final = [ + deployment("small", "s"), + deployment("big", "b"), + _complexity_router(deployment_id="cheap_marker", tags=["cheap"]), + _complexity_router(tiers=ALL_BIG, deployment_id="premium_marker", tags=["premium"]), +] + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("request_kwargs", "expected"), + [({"metadata": {"tags": ["cheap"]}}, "s"), ({"metadata": {"tags": ["premium"]}}, "b"), ({}, "s")], + ids=["cheap-tag", "premium-tag", "untagged-takes-the-first"], +) +async def test_request_tags_pick_among_markers_of_one_name( + make_core: MakeCore, request_kwargs: dict[str, object], expected: str +) -> None: + core: Final = make_core(model_list=list(TAG_SCOPED)) + chosen: Final = await core.select("smart", request_kwargs, messages=MESSAGES) # pyright: ignore[reportCallIssue] + assert chosen["model_info"]["id"] == expected # pyright: ignore[reportIndexIssue] + + +@pytest.mark.asyncio +async def test_an_untagged_request_takes_the_default_tagged_marker(make_core: MakeCore) -> None: + core: Final = make_core( + model_list=[ + deployment("small", "s"), + deployment("big", "b"), + _complexity_router(deployment_id="cheap_marker", tags=["cheap"]), + _complexity_router(tiers=ALL_BIG, deployment_id="default_marker", tags=["default"]), + ] + ) + chosen: Final = await core.select("smart", {}, messages=MESSAGES) # pyright: ignore[reportCallIssue] + assert chosen["model_info"]["id"] == "b" # pyright: ignore[reportIndexIssue] + + +SCOPED_WITH_PLAIN: Final = [ + deployment("small", "s"), + deployment("big", "b"), + _complexity_router(deployment_id="cheap_marker", tags=["cheap"]), + deployment("smart", "plain"), +] + + +@pytest.mark.asyncio +async def test_with_tag_filtering_an_unmatched_tag_skips_the_markers_for_plain_tag_selection( + make_core: MakeCore, +) -> None: + core: Final = make_core(model_list=list(SCOPED_WITH_PLAIN), enable_tag_filtering=True) + with pytest.raises( + ValueError, match=re.escape("Not allowed to access model due to tags configuration. Passed model=smart") + ): + await core.select("smart", {"metadata": {"tags": ["other"]}}, messages=MESSAGES) # pyright: ignore[reportCallIssue] + chosen: Final = await core.select("smart", {"metadata": {"tags": ["cheap"]}}, messages=MESSAGES) # pyright: ignore[reportCallIssue] + assert chosen["model_info"]["id"] == "s" # pyright: ignore[reportIndexIssue] + + +@pytest.mark.asyncio +async def test_without_tag_filtering_an_unmatched_tag_still_takes_the_first_marker(make_core: MakeCore) -> None: + core: Final = make_core(model_list=list(SCOPED_WITH_PLAIN)) + chosen: Final = await core.select("smart", {"metadata": {"tags": ["other"]}}, messages=MESSAGES) # pyright: ignore[reportCallIssue] + assert chosen["model_info"]["id"] == "s" # pyright: ignore[reportIndexIssue] + + +# --- tier litellm_params and alias-marker params ------------------------------------------------------------- + +PINNED_TIERS: Final = { + **TIERS, + "SIMPLE": {"model_name": "small", "litellm_params": {"temperature": 0.1, "max_tokens": 7}}, +} + + +@pytest.mark.asyncio +async def test_tier_litellm_params_override_the_callers_params_on_the_call( + make_core: MakeCore, host: ScriptedHost +) -> None: + core: Final = make_core(model_list=_tiered(tiers=PINNED_TIERS)) + await core.acompletion(model="smart", messages=MESSAGES, temperature=0.9, max_completion_tokens=100) # pyright: ignore[reportAttributeAccessIssue] + sent: Final = host.calls[-1].kwargs + assert host.deployments_called() == ["s"] + assert sent["temperature"] == 0.1 + assert sent["max_tokens"] == 7 + assert "max_completion_tokens" not in sent + + +@pytest.mark.asyncio +async def test_tier_litellm_params_are_pinned_onto_the_request_by_select(make_core: MakeCore) -> None: + core: Final = make_core(model_list=_tiered(tiers=PINNED_TIERS)) + request: Final[dict[str, object]] = {"temperature": 0.9, "max_completion_tokens": 100} + await core.select("smart", request, messages=MESSAGES) # pyright: ignore[reportCallIssue] + assert request["temperature"] == 0.1 + assert request["max_tokens"] == 7 + assert "max_completion_tokens" not in request + + +@pytest.mark.asyncio +async def test_tier_litellm_params_apply_only_to_their_tier(make_core: MakeCore, host: ScriptedHost) -> None: + core: Final = make_core(model_list=_tiered(tiers=PINNED_TIERS)) + await core.acompletion(model="smart", messages=HARD_MESSAGES, temperature=0.9) # pyright: ignore[reportAttributeAccessIssue] + assert host.deployments_called() == ["b"] + assert host.calls[-1].kwargs["temperature"] == 0.9 + assert "max_tokens" not in host.calls[-1].kwargs + + +@pytest.mark.asyncio +async def test_marker_params_fill_gaps_on_the_routed_call_but_never_override_the_caller( + make_core: MakeCore, host: ScriptedHost +) -> None: + core: Final = make_core(model_list=_tiered(temperature=0.3)) + await core.acompletion(model="smart", messages=MESSAGES) # pyright: ignore[reportAttributeAccessIssue] + assert host.calls[-1].kwargs["temperature"] == 0.3 + await core.acompletion(model="smart", messages=MESSAGES, temperature=0.9) # pyright: ignore[reportAttributeAccessIssue] + assert host.calls[-1].kwargs["temperature"] == 0.9 + + +@pytest.mark.asyncio +async def test_marker_credentials_and_pricing_are_never_forwarded(make_core: MakeCore, host: ScriptedHost) -> None: + core: Final = make_core(model_list=_tiered(api_base="http://marker.example", input_cost_per_token=0.5)) + await core.acompletion(model="smart", messages=HARD_MESSAGES) # pyright: ignore[reportAttributeAccessIssue] + sent: Final = host.calls[-1].kwargs + assert host.deployments_called() == ["b"] + assert sent.get("api_base") is None + assert "input_cost_per_token" not in sent + + +@pytest.mark.asyncio +async def test_a_deployments_own_param_beats_the_forwarded_marker_param( + make_core: MakeCore, host: ScriptedHost +) -> None: + core: Final = make_core( + model_list=[ + deployment("small", "s", aws_region_name="deployment-region"), + deployment("big", "b"), + _complexity_router(aws_region_name="marker-region"), + ] + ) + await core.acompletion(model="smart", messages=MESSAGES) # pyright: ignore[reportAttributeAccessIssue] + assert host.calls[-1].kwargs["aws_region_name"] == "deployment-region" + await core.acompletion(model="smart", messages=HARD_MESSAGES) # pyright: ignore[reportAttributeAccessIssue] + assert host.calls[-1].kwargs["aws_region_name"] == "marker-region" + + +# --- the Claude Code session router binding -------------------------------------------------------------------- + + +def _session_request( + agent_id: str | None = None, caller: str = "caller-hash" +) -> dict[str, object]: # mutable-ok: request kwargs + headers: Final[dict[str, str]] = {"x-claude-code-session-id": SESSION_ID} + if agent_id is not None: + headers["x-claude-code-agent-id"] = agent_id + return {"proxy_server_request": {"headers": headers}, "metadata": {"user_api_key_hash": caller}} + + +@pytest.mark.asyncio +async def test_a_subagent_request_follows_the_router_its_session_started_on(make_core: MakeCore) -> None: + core: Final = make_core(model_list=_tiered()) + await core.select("smart", _session_request(), messages=MESSAGES) # pyright: ignore[reportCallIssue] + subagent: Final = _session_request(agent_id="agent-1") + chosen: Final = await core.select("big", subagent, messages=MESSAGES) # pyright: ignore[reportCallIssue] + assert chosen["model_info"]["id"] == "s" # pyright: ignore[reportIndexIssue] + assert _metadata(subagent)["model_group"] == "smart" + + +@pytest.mark.asyncio +async def test_a_session_binding_is_scoped_to_the_caller(make_core: MakeCore) -> None: + core: Final = make_core(model_list=_tiered()) + await core.select("smart", _session_request(), messages=MESSAGES) # pyright: ignore[reportCallIssue] + other_caller: Final = _session_request(agent_id="agent-1", caller="other-hash") + chosen: Final = await core.select("big", other_caller, messages=MESSAGES) # pyright: ignore[reportCallIssue] + assert chosen["model_info"]["id"] == "b" # pyright: ignore[reportIndexIssue] + + +@pytest.mark.asyncio +async def test_a_main_agent_request_is_not_redirected_by_its_session_binding(make_core: MakeCore) -> None: + core: Final = make_core(model_list=_tiered()) + await core.select("smart", _session_request(), messages=MESSAGES) # pyright: ignore[reportCallIssue] + chosen: Final = await core.select("big", _session_request(), messages=MESSAGES) # pyright: ignore[reportCallIssue] + assert chosen["model_info"]["id"] == "b" # pyright: ignore[reportIndexIssue] + + +@pytest.mark.asyncio +async def test_a_subagent_request_without_a_session_binding_keeps_its_model(make_core: MakeCore) -> None: + core: Final = make_core(model_list=_tiered()) + chosen: Final = await core.select("big", _session_request(agent_id="agent-1"), messages=MESSAGES) # pyright: ignore[reportCallIssue] + assert chosen["model_info"]["id"] == "b" # pyright: ignore[reportIndexIssue] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("session_id", ["short", "has spaces in it"]) +async def test_an_invalid_session_id_never_binds(make_core: MakeCore, session_id: str) -> None: + core: Final = make_core(model_list=_tiered()) + first: Final = _session_request() + first["proxy_server_request"] = {"headers": {"x-claude-code-session-id": session_id}} + await core.select("smart", first, messages=MESSAGES) # pyright: ignore[reportCallIssue] + subagent: Final = _session_request(agent_id="agent-1") + subagent["proxy_server_request"] = { + "headers": {"x-claude-code-session-id": session_id, "x-claude-code-agent-id": "agent-1"} + } + chosen: Final = await core.select("big", subagent, messages=MESSAGES) # pyright: ignore[reportCallIssue] + assert chosen["model_info"]["id"] == "b" # pyright: ignore[reportIndexIssue] + + +# --- the routing-plugin pipeline ------------------------------------------------------------------------------ + + +class Narrowing: + """Keeps only ``keep`` among the candidate models and records what it was handed.""" + + def __init__(self, keep: frozenset[str]) -> None: + self.keep: Final = keep + self.seen: list[list[str]] = [] # mutable-ok: the candidate lists each run saw + + async def run(self, context: RoutingContext) -> RoutingContext: + self.seen.append(list(context.candidate_models)) + context.candidate_models = [model for model in context.candidate_models if model in self.keep] + context.signals["narrowing"] = {"kept": sorted(self.keep)} + return context + + +TWO: Final = [deployment("g", "a"), deployment("g", "b")] + + +@pytest.mark.asyncio +async def test_a_routing_plugin_narrows_the_candidates_to_the_models_it_keeps(make_core: MakeCore) -> None: + plugin: Final = Narrowing(frozenset({"openai/b"})) + core: Final = make_core(model_list=list(TWO), plugins=[plugin]) + assert await picked_ids(core, "g", messages=MESSAGES) == {"b"} + assert plugin.seen[0] == ["openai/a", "openai/b"] + + +@pytest.mark.asyncio +async def test_routing_plugin_signals_and_narrowed_candidates_are_stamped_on_the_request(make_core: MakeCore) -> None: + core: Final = make_core(model_list=list(TWO), plugins=[Narrowing(frozenset({"openai/b"}))]) + request: Final[dict[str, object]] = {} + await core.select("g", request, messages=MESSAGES) # pyright: ignore[reportCallIssue] + assert request == { + "metadata": { + "routing_plugin_signals": {"narrowing": {"kept": ["openai/b"]}}, + "_routing_plugin_candidate_models": ["openai/b"], + } + } + + +@pytest.mark.asyncio +async def test_a_routing_plugin_that_keeps_every_candidate_stamps_only_its_signals(make_core: MakeCore) -> None: + core: Final = make_core(model_list=list(TWO), plugins=[Narrowing(frozenset({"openai/a", "openai/b"}))]) + request: Final[dict[str, object]] = {} + await core.select("g", request, messages=MESSAGES) # pyright: ignore[reportCallIssue] + assert request == {"metadata": {"routing_plugin_signals": {"narrowing": {"kept": ["openai/a", "openai/b"]}}}} + + +@pytest.mark.asyncio +async def test_routing_plugins_run_in_order_each_seeing_the_previous_narrowing(make_core: MakeCore) -> None: + first: Final = Narrowing(frozenset({"openai/b", "openai/c"})) + second: Final = Narrowing(frozenset({"openai/c"})) + core: Final = make_core(model_list=[*TWO, deployment("g", "c")], plugins=[first, second]) + assert await picked_ids(core, "g", times=1, messages=MESSAGES) == {"c"} + assert second.seen == [["openai/b", "openai/c"]] + + +@pytest.mark.asyncio +async def test_a_routing_plugin_that_narrows_to_nothing_fails_the_pick(make_core: MakeCore, host: ScriptedHost) -> None: + core: Final = make_core(model_list=list(TWO), plugins=[Narrowing(frozenset())]) + with pytest.raises(ValueError, match="No deployments left") as raised: + await core.select("g", {}, messages=MESSAGES) # pyright: ignore[reportCallIssue] + assert str(raised.value) == "No deployments left after routing-plugin filtering. candidate_models=[]" + + +SYNC_PLUGIN_MESSAGE_START: Final = ( + "Router(plugins=[...]) is configured, but this call resolved to the synchronous deployment-selection path" +) + + +@pytest.mark.parametrize("entry", ["get_available_deployment", "completion"]) +def test_routing_plugins_refuse_the_sync_selection_path(make_core: MakeCore, host: ScriptedHost, entry: str) -> None: + core: Final = make_core(model_list=list(TWO), plugins=[Narrowing(frozenset({"openai/a"}))]) + with pytest.raises(ValueError, match=re.escape(SYNC_PLUGIN_MESSAGE_START)) as raised: + _sync_entry(core, entry, "g") + assert str(raised.value).startswith(SYNC_PLUGIN_MESSAGE_START) + assert host.calls == [] + + +@pytest.mark.asyncio +async def test_routing_plugins_refuse_usage_based_routing_v1_even_from_async_callers(make_core: MakeCore) -> None: + core: Final = make_core( + model_list=list(TWO), plugins=[Narrowing(frozenset({"openai/a"}))], routing_strategy="usage-based-routing" + ) + with pytest.raises(ValueError, match=re.escape(SYNC_PLUGIN_MESSAGE_START)) as raised: + await core.select("g", {}, messages=MESSAGES) # pyright: ignore[reportCallIssue] + assert str(raised.value).startswith(SYNC_PLUGIN_MESSAGE_START) + + +@pytest.mark.asyncio +async def test_a_routing_plugin_does_not_narrow_a_pick_by_deployment_id(make_core: MakeCore) -> None: + core: Final = make_core(model_list=list(TWO), plugins=[Narrowing(frozenset({"openai/b"}))]) + chosen: Final = await core.select("a", {}, messages=MESSAGES) # pyright: ignore[reportCallIssue] + assert chosen["model_info"]["id"] == "a" # pyright: ignore[reportIndexIssue] + + +@pytest.mark.asyncio +async def test_completion_with_a_routing_plugin_calls_a_kept_deployment( + make_core: MakeCore, host: ScriptedHost +) -> None: + core: Final = make_core(model_list=list(TWO), plugins=[Narrowing(frozenset({"openai/a"}))]) + await asyncio.gather(*(core.acompletion(model="g", messages=MESSAGES) for _ in range(5))) # pyright: ignore[reportAttributeAccessIssue] + assert host.deployments_called() == ["a"] * 5 diff --git a/tests/unit/test_router/core/test_selection_characterization_session_repin.py b/tests/unit/test_router/core/test_selection_characterization_session_repin.py new file mode 100644 index 00000000000..f022e267586 --- /dev/null +++ b/tests/unit/test_router/core/test_selection_characterization_session_repin.py @@ -0,0 +1,119 @@ +""" +Characterization of ``session_affinity_repin_on_failover``: a session pinned to a deployment that +fails is served by another deployment of the group, and a completed success there moves the +session's pin to it when the flag is on. With it off, the pin stays where it was. +""" + +import asyncio +from collections.abc import Callable + +import pytest + +from litellm.router_core import RouterCore +from litellm.router_utils.pre_call_checks.deployment_affinity_check import DeploymentAffinityCheck +from tests.unit.test_router.core.scripted_host import ScriptedHost, completion_response, deployment, rate_limit +from tests.unit.test_router.core.streaming_helpers import ( + MESSAGES, + chat_stream, + collect_texts, + install_orjson_stub_if_missing, +) + +install_orjson_stub_if_missing() + +MakeCore = Callable[..., RouterCore] + +_SESSION = "session-1" +_CALLER = "caller-1" +_METADATA = {"session_id": _SESSION, "user_api_key_hash": _CALLER} + + +def _affinity_check(core: RouterCore) -> DeploymentAffinityCheck: + return next( + callback + for callback in core.optional_callbacks or () # pyright: ignore[reportAttributeAccessIssue, reportUnknownVariableType, reportUnknownMemberType] + if isinstance(callback, DeploymentAffinityCheck) + ) + + +async def _pinned(check: DeploymentAffinityCheck) -> object: + key = DeploymentAffinityCheck.get_session_affinity_cache_key("g", _SESSION, _CALLER) + pinned: object = await check.cache.async_get_cache(key) # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] # DualCache's reads are untyped + return pinned # pyright: ignore[reportUnknownVariableType] # see above + + +async def _complete(core: RouterCore) -> object: + response: object = await core.acompletion( # pyright: ignore[reportAttributeAccessIssue, reportUnknownMemberType, reportUnknownVariableType] # endpoint method + model="g", messages=[{"role": "user", "content": "hi"}], metadata=dict(_METADATA) + ) + return response # pyright: ignore[reportUnknownVariableType] # see above + + +@pytest.mark.asyncio +@pytest.mark.parametrize(("repin", "next_served_by"), [(True, "b"), (False, "a")]) +async def test_a_failover_success_repins_the_session_only_with_repin_on_failover( + make_core: MakeCore, host: ScriptedHost, repin: bool, next_served_by: str +) -> None: + core = make_core( + model_list=[deployment("g", "a"), deployment("g", "b")], + optional_pre_call_checks=["session_affinity"], + session_affinity_repin_on_failover=repin, + num_retries=1, + ) + check = _affinity_check(core) + key = DeploymentAffinityCheck.get_session_affinity_cache_key("g", _SESSION, _CALLER) + await check.cache.async_set_cache(key, {"model_id": "a"}) # pyright: ignore[reportUnknownMemberType] + host.answer("openai/a", rate_limit(), completion_response()) + host.answer("openai/b", completion_response()) + + await _complete(core) + + # The pinned deployment fails and cools down; the group's other deployment serves the retry. + assert host.deployments_called("acompletion") == ["a", "b"] + assert await _pinned(check) == {"model_id": "b" if repin else "a"} + + # Once the first deployment is healthy again, the session goes where it is pinned. + core.health().cooldowns.cooldown_store.flush_cache() # pyright: ignore[reportAttributeAccessIssue, reportUnknownMemberType] + await _complete(core) + + assert host.deployments_called("acompletion")[-1] == next_served_by + + +async def _settle() -> None: + """Let the logging and recording tasks a finished stream scheduled run.""" + for _ in range(20): + await asyncio.sleep(0) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("repin", [True, False]) +async def test_a_failover_stream_repins_the_session_when_it_ends( + make_core: MakeCore, host: ScriptedHost, repin: bool +) -> None: + core = make_core( + model_list=[deployment("g", "a"), deployment("g", "b")], + optional_pre_call_checks=["session_affinity"], + session_affinity_repin_on_failover=repin, + num_retries=1, + ) + check = _affinity_check(core) + key = DeploymentAffinityCheck.get_session_affinity_cache_key("g", _SESSION, _CALLER) + await check.cache.async_set_cache(key, {"model_id": "a"}) # pyright: ignore[reportUnknownMemberType] + host.answer("openai/a", rate_limit()) + host.answer( + "openai/b", + chat_stream("hello", usage={"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2}, logged=True), + ) + + stream: object = await core.acompletion( # pyright: ignore[reportAttributeAccessIssue, reportUnknownMemberType, reportUnknownVariableType] # endpoint method + model="g", messages=MESSAGES, stream=True, metadata=dict(_METADATA) + ) + await _settle() + # The pin moves only once the stream has been read to its end. + assert await _pinned(check) == {"model_id": "a"} + + assert await collect_texts(stream) == "hello" # pyright: ignore[reportUnknownArgumentType] # see above + await _settle() + + assert host.deployments_called("acompletion") == ["a", "b"] + assert await _pinned(check) == {"model_id": "b" if repin else "a"} diff --git a/tests/unit/test_router/core/test_selection_characterization_strategies.py b/tests/unit/test_router/core/test_selection_characterization_strategies.py new file mode 100644 index 00000000000..3999c24857f --- /dev/null +++ b/tests/unit/test_router/core/test_selection_characterization_strategies.py @@ -0,0 +1,394 @@ +""" +Characterization of how a router core's routing strategies pick a deployment: strategy validation, +simple-shuffle weighting, the stateful strategies without traffic, per-request strategy overrides, +runtime strategy changes and custom strategies registered with ``register_plugin``. +""" + +import re +from typing import Final + +import pydantic +import pytest + +import litellm +from litellm.router_api import ChangeSet +from litellm.types.router import CustomRoutingStrategyBase, RouterRateLimitError, RoutingStrategy +from tests.unit.test_router.core.scripted_host import ScriptedHost, deployment +from tests.unit.test_router.core.selection_helpers import MESSAGES, MakeCore, picked_ids + +SPREAD_PICKS: Final = 60 +"""Enough uniform picks between two deployments that missing one is a 2**-59 event.""" + +INVALID_STRATEGY_MESSAGE: Final = ( + "Invalid routing_strategy: '{name}'. Valid options: ['simple-shuffle', 'lar1', 'least-busy', " + "'latency-based-routing', 'cost-based-routing', 'usage-based-routing-v2', 'usage-based-routing', " + "'provider-budget-routing']. Check 'router_settings.routing_strategy' in your config.yaml or the " + "'routing_strategy' parameter if using the Router SDK directly." +) + +NO_DEPLOYMENTS_MESSAGE: Final = ( + "No deployments available for selected model, Try again in 5 seconds. Passed model=g. " + "pre-call-checks=False, cooldown_list=[]" +) + + +def _two(**params: object) -> list[dict[str, object]]: # mutable-ok: a model_list # kwargs-ok: shared litellm_params + return [deployment("g", "a", **params), deployment("g", "b", **params)] + + +# --- strategy validation --------------------------------------------------------------------- + + +@pytest.mark.parametrize("routing_strategy", ["bogus", "priority", "Simple-Shuffle"]) +def test_an_unknown_routing_strategy_is_rejected_at_construction(make_core: MakeCore, routing_strategy: str) -> None: + with pytest.raises(ValueError, match="Invalid routing_strategy") as raised: + make_core(model_list=_two(), routing_strategy=routing_strategy) + assert str(raised.value) == INVALID_STRATEGY_MESSAGE.format(name=routing_strategy) + + +@pytest.mark.asyncio +async def test_lar1_is_an_accepted_strategy_that_picks_within_the_group(make_core: MakeCore) -> None: + core: Final = make_core(model_list=[*_two(), deployment("other", "c")], routing_strategy="lar1") + assert await picked_ids(core, "g", messages=MESSAGES) <= {"a", "b"} + + +@pytest.mark.asyncio +async def test_provider_budget_routing_without_budgets_has_no_selector_and_rate_limits_every_pick( + make_core: MakeCore, host: ScriptedHost +) -> None: + core: Final = make_core(model_list=_two(), routing_strategy="provider-budget-routing") + with pytest.raises(RouterRateLimitError) as raised: + await core.select("g", messages=MESSAGES) # pyright: ignore[reportCallIssue] + assert str(raised.value) == NO_DEPLOYMENTS_MESSAGE + with pytest.raises(RouterRateLimitError): + await core.acompletion(model="g", messages=MESSAGES) # pyright: ignore[reportAttributeAccessIssue] + assert host.calls == [] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("ids", [("a", "b"), ("b", "a")]) +async def test_a_routing_strategy_enum_member_is_accepted_like_its_string( + make_core: MakeCore, ids: tuple[str, str] +) -> None: + core: Final = make_core( + model_list=[deployment("g", deployment_id) for deployment_id in ids], + routing_strategy=RoutingStrategy.LEAST_BUSY, + ) + assert await picked_ids(core, "g", messages=MESSAGES) == {ids[0]} + + +@pytest.mark.asyncio +async def test_select_with_an_unknown_kind_raises_value_error(make_core: MakeCore) -> None: + core: Final = make_core(model_list=_two()) + with pytest.raises(ValueError, match=re.escape("unknown select kind 'nope'")): + await core.select("g", kind="nope") # pyright: ignore[reportCallIssue, reportArgumentType] + + +# --- simple-shuffle -------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_simple_shuffle_spreads_picks_over_unweighted_deployments(make_core: MakeCore) -> None: + core: Final = make_core(model_list=_two()) + assert await picked_ids(core, "g", times=SPREAD_PICKS) == {"a", "b"} + + +@pytest.mark.asyncio +async def test_simple_shuffle_spreads_picks_when_every_weight_is_zero(make_core: MakeCore) -> None: + core: Final = make_core(model_list=_two(weight=0, rpm=0, tpm=0)) + assert await picked_ids(core, "g", times=SPREAD_PICKS) == {"a", "b"} + + +@pytest.mark.asyncio +async def test_simple_shuffle_falls_through_all_zero_weights_to_rpm(make_core: MakeCore) -> None: + core: Final = make_core(model_list=[deployment("g", "a", weight=0, rpm=5), deployment("g", "b", weight=0, rpm=0)]) + assert await picked_ids(core, "g") == {"a"} + + +@pytest.mark.asyncio +async def test_simple_shuffle_falls_through_unset_weight_and_rpm_to_tpm(make_core: MakeCore) -> None: + core: Final = make_core(model_list=[deployment("g", "a", tpm=0), deployment("g", "b", tpm=100)]) + assert await picked_ids(core, "g") == {"b"} + + +@pytest.mark.asyncio +async def test_rpm_takes_precedence_over_tpm(make_core: MakeCore) -> None: + core: Final = make_core(model_list=[deployment("g", "a", rpm=1, tpm=0), deployment("g", "b", rpm=0, tpm=1000)]) + assert await picked_ids(core, "g") == {"a"} + + +@pytest.mark.asyncio +async def test_simple_shuffle_with_a_non_numeric_weight_raises_type_error_at_pick(make_core: MakeCore) -> None: + core: Final = make_core(model_list=[deployment("g", "a", weight="heavy"), deployment("g", "b")]) + with pytest.raises(TypeError, match="Deployment weight must be numeric"): + await core.select("g") + + +@pytest.mark.asyncio +async def test_request_router_weights_are_keyed_by_the_alias_target_group(make_core: MakeCore) -> None: + core: Final = make_core(model_list=_two(), model_group_alias={"alias": "g"}) + assert await picked_ids(core, "alias", request_kwargs={"_router_weights": {"g": {"b": 1.0}}}) == {"b"} + + +@pytest.mark.asyncio +async def test_request_router_weights_for_another_group_are_ignored(make_core: MakeCore) -> None: + core: Final = make_core(model_list=[deployment("g", "a", weight=1), deployment("g", "b", weight=0)]) + assert await picked_ids(core, "g", request_kwargs={"_router_weights": {"other": {"b": 1.0}}}) == {"a"} + + +@pytest.mark.asyncio +@pytest.mark.parametrize("weights", [{"g": {"a": 0.0, "b": 0.0}}, {"g": {"a": -1.0}}, {"g": {"a": "1"}}]) +async def test_invalid_request_router_weights_fail_the_pick_with_a_validation_error( + make_core: MakeCore, weights: dict[str, object] +) -> None: + core: Final = make_core(model_list=_two()) + with pytest.raises(pydantic.ValidationError): + await core.select("g", {"_router_weights": weights}) + + +@pytest.mark.asyncio +async def test_an_empty_request_router_weights_group_falls_back_to_deployment_weights(make_core: MakeCore) -> None: + core: Final = make_core(model_list=[deployment("g", "a", weight=0), deployment("g", "b", weight=2)]) + assert await picked_ids(core, "g", request_kwargs={"_router_weights": {"g": {}}}) == {"b"} + + +@pytest.mark.asyncio +async def test_request_router_weights_give_unlisted_deployments_zero_weight(make_core: MakeCore) -> None: + core: Final = make_core(model_list=[*_two(), deployment("g", "c")]) + assert await picked_ids(core, "g", request_kwargs={"_router_weights": {"g": {"c": 0.5}}}) == {"c"} + + +# --- stateful strategies without traffic --------------------------------------------------- + + +@pytest.mark.asyncio +@pytest.mark.parametrize("routing_strategy", ["latency-based-routing", "usage-based-routing-v2"]) +async def test_strategy_without_traffic_spreads_picks(make_core: MakeCore, routing_strategy: str) -> None: + core: Final = make_core(model_list=_two(), routing_strategy=routing_strategy) + assert await picked_ids(core, "g", times=SPREAD_PICKS, messages=MESSAGES) == {"a", "b"} + + +@pytest.mark.asyncio +@pytest.mark.parametrize("ids", [("a", "b"), ("b", "a")]) +async def test_cost_based_routing_without_known_costs_picks_the_first_configured( + make_core: MakeCore, ids: tuple[str, str] +) -> None: + core: Final = make_core( + model_list=[deployment("g", deployment_id) for deployment_id in ids], routing_strategy="cost-based-routing" + ) + assert await picked_ids(core, "g", messages=MESSAGES) == {ids[0]} + + +def test_sync_pick_under_usage_based_routing_v1_picks_the_first_configured(make_core: MakeCore) -> None: + core: Final = make_core(model_list=_two(), routing_strategy="usage-based-routing") + for _ in range(5): + assert core.get_available_deployment("g", messages=MESSAGES)["model_info"]["id"] == "a" # pyright: ignore[reportAttributeAccessIssue] + + +def test_sync_completion_under_least_busy_calls_the_first_configured(make_core: MakeCore, host: ScriptedHost) -> None: + core: Final = make_core(model_list=_two(), routing_strategy="least-busy") + for _ in range(3): + core.completion(model="g", messages=MESSAGES) # pyright: ignore[reportAttributeAccessIssue] + assert host.deployments_called("completion") == ["a", "a", "a"] + + +# --- per-request routing strategy override -------------------------------------------------- + + +@pytest.mark.asyncio +@pytest.mark.parametrize("override", ["least-busy", RoutingStrategy.LEAST_BUSY], ids=["string", "enum"]) +async def test_a_request_routing_strategy_overrides_simple_shuffle(make_core: MakeCore, override: object) -> None: + core: Final = make_core(model_list=_two()) + assert await picked_ids(core, "g", request_kwargs={"routing_strategy": override}, messages=MESSAGES) == {"a"} + + +@pytest.mark.asyncio +async def test_a_simple_shuffle_request_override_spreads_picks_on_a_least_busy_router(make_core: MakeCore) -> None: + core: Final = make_core(model_list=_two(), routing_strategy="least-busy") + picked: Final = await picked_ids( + core, "g", times=SPREAD_PICKS, request_kwargs={"routing_strategy": "simple-shuffle"}, messages=MESSAGES + ) + assert picked == {"a", "b"} + + +@pytest.mark.asyncio +async def test_a_cost_based_request_override_picks_the_cheapest(make_core: MakeCore) -> None: + core: Final = make_core( + model_list=[ + deployment("g", "pricey", input_cost_per_token=0.00001, output_cost_per_token=0.00002), + deployment("g", "cheap", input_cost_per_token=0.000001, output_cost_per_token=0.000002), + ] + ) + picked: Final = await picked_ids( + core, "g", request_kwargs={"routing_strategy": "cost-based-routing"}, messages=MESSAGES + ) + assert picked == {"cheap"} + + +@pytest.mark.asyncio +@pytest.mark.parametrize("override", ["bogus", "lar1", "provider-budget-routing", 5]) +async def test_an_unsupported_request_routing_strategy_is_ignored(make_core: MakeCore, override: object) -> None: + core: Final = make_core(model_list=_two()) + picked: Final = await picked_ids( + core, "g", times=SPREAD_PICKS, request_kwargs={"routing_strategy": override}, messages=MESSAGES + ) + assert picked == {"a", "b"} + + +@pytest.mark.asyncio +async def test_a_request_routing_strategy_overrides_the_routing_group_strategy(make_core: MakeCore) -> None: + core: Final = make_core( + model_list=_two(), routing_groups=[{"group_name": "grp", "models": ["g"], "routing_strategy": "least-busy"}] + ) + assert await picked_ids(core, "g", messages=MESSAGES) == {"a"} + picked: Final = await picked_ids( + core, "g", times=SPREAD_PICKS, request_kwargs={"routing_strategy": "simple-shuffle"}, messages=MESSAGES + ) + assert picked == {"a", "b"} + + +@pytest.mark.asyncio +async def test_completion_with_a_request_routing_strategy_calls_the_overridden_pick( + make_core: MakeCore, host: ScriptedHost +) -> None: + core: Final = make_core(model_list=_two()) + for _ in range(5): + await core.acompletion(model="g", messages=MESSAGES, routing_strategy="least-busy") # pyright: ignore[reportAttributeAccessIssue] + assert host.deployments_called("acompletion") == ["a"] * 5 + + +# --- runtime strategy changes ------------------------------------------------------------------ + + +@pytest.mark.asyncio +async def test_applying_a_routing_strategy_switches_how_picks_are_made(make_core: MakeCore) -> None: + core: Final = make_core(model_list=_two()) + core.apply(ChangeSet(settings={"routing_strategy": "least-busy"})) + assert core.settings().routing_strategy == "least-busy" + assert await picked_ids(core, "g", messages=MESSAGES) == {"a"} + + +@pytest.mark.asyncio +async def test_applying_an_invalid_routing_strategy_raises_and_keeps_the_previous_one(make_core: MakeCore) -> None: + core: Final = make_core(model_list=_two(), routing_strategy="least-busy") + with pytest.raises(ValueError, match="Invalid routing_strategy") as raised: + core.apply(ChangeSet(settings={"routing_strategy": "bogus"})) + assert str(raised.value) == INVALID_STRATEGY_MESSAGE.format(name="bogus") + assert core.settings().routing_strategy == "least-busy" + assert await picked_ids(core, "g", messages=MESSAGES) == {"a"} + + +# --- custom routing strategies ----------------------------------------------------------------- + + +def _fixed(deployment_id: str) -> dict[str, object]: # mutable-ok: the strategy interface's own type + return { + "model_name": "x", + "litellm_params": {"model": f"openai/{deployment_id}"}, + "model_info": {"id": deployment_id}, + } + + +class FixedStrategy(CustomRoutingStrategyBase): + """Always answers ``custom`` from the async pick and ``sync-custom`` from the sync pick.""" + + def __init__(self) -> None: + self.asked: list[str] = [] # mutable-ok: the models each pick was asked for + + async def async_get_available_deployment( + self, + model: str, + messages: list[dict[str, str]] | None = None, # mutable-ok: the strategy interface's own type + input: str | list | None = None, # mutable-ok: the strategy interface's own type + specific_deployment: bool | None = False, + request_kwargs: dict | None = None, # mutable-ok: the strategy interface's own type + ) -> dict: # mutable-ok: the strategy interface's own type + self.asked.append(model) + return _fixed("custom") + + def get_available_deployment( + self, + model: str, + messages: list[dict[str, str]] | None = None, # mutable-ok: the strategy interface's own type + input: str | list | None = None, # mutable-ok: the strategy interface's own type + specific_deployment: bool | None = False, + request_kwargs: dict | None = None, # mutable-ok: the strategy interface's own type + ) -> dict: # mutable-ok: the strategy interface's own type + self.asked.append(model) + return _fixed("sync-custom") + + +@pytest.mark.asyncio +@pytest.mark.parametrize("model", ["g", "unknown-model"]) +async def test_a_custom_strategy_answer_is_returned_as_is_for_any_model(make_core: MakeCore, model: str) -> None: + core: Final = make_core(model_list=_two()) + strategy: Final = FixedStrategy() + core.register_plugin(strategy) + chosen: Final = await core.select(model) + assert chosen == _fixed("custom") + assert strategy.asked == [model] + + +@pytest.mark.asyncio +async def test_a_custom_strategy_does_not_change_the_healthy_listing(make_core: MakeCore) -> None: + core: Final = make_core(model_list=_two()) + core.register_plugin(FixedStrategy()) + healthy: Final = await core.select("g", kind="healthy") + assert sorted(d["model_info"]["id"] for d in healthy) == ["a", "b"] # pyright: ignore[reportIndexIssue, reportArgumentType] + + +@pytest.mark.asyncio +async def test_completion_calls_the_deployment_a_custom_strategy_returns( + make_core: MakeCore, host: ScriptedHost +) -> None: + core: Final = make_core(model_list=_two()) + core.register_plugin(FixedStrategy()) + await core.acompletion(model="g", messages=MESSAGES) # pyright: ignore[reportAttributeAccessIssue] + assert host.deployments_called("acompletion") == ["custom"] + assert host.calls[-1].model == "openai/custom" + + +def test_sync_completion_uses_the_custom_strategys_sync_pick(make_core: MakeCore, host: ScriptedHost) -> None: + core: Final = make_core(model_list=_two()) + core.register_plugin(FixedStrategy()) + core.completion(model="g", messages=MESSAGES) # pyright: ignore[reportAttributeAccessIssue] + assert host.deployments_called("completion") == ["sync-custom"] + assert core.get_available_deployment("g") == _fixed("sync-custom") # pyright: ignore[reportAttributeAccessIssue] + + +@pytest.mark.asyncio +async def test_set_custom_routing_strategy_registers_like_register_plugin(make_core: MakeCore) -> None: + core: Final = make_core(model_list=_two()) + core.set_custom_routing_strategy(FixedStrategy()) # pyright: ignore[reportAttributeAccessIssue] + assert await core.select("g") == _fixed("custom") + + +@pytest.mark.asyncio +async def test_applying_a_different_routing_strategy_drops_the_custom_strategy(make_core: MakeCore) -> None: + core: Final = make_core(model_list=_two()) + core.register_plugin(FixedStrategy()) + core.apply(ChangeSet(settings={"routing_strategy": "least-busy"})) + assert await picked_ids(core, "g", messages=MESSAGES) == {"a"} + + +@pytest.mark.asyncio +async def test_applying_the_current_routing_strategy_keeps_the_custom_strategy(make_core: MakeCore) -> None: + core: Final = make_core(model_list=_two()) + core.register_plugin(FixedStrategy()) + core.apply(ChangeSet(settings={"routing_strategy": "simple-shuffle"})) + assert await core.select("g") == _fixed("custom") + + +@pytest.mark.asyncio +async def test_a_custom_strategy_error_reaches_the_caller(make_core: MakeCore, host: ScriptedHost) -> None: + class Refusing(FixedStrategy): + async def async_get_available_deployment( + self, model: str, *args: object, **kwargs: object + ) -> dict: # mutable-ok: the strategy interface's own type # kwargs-ok: the strategy interface + raise litellm.BadRequestError(message="refused by strategy", model=model, llm_provider="") + + core: Final = make_core(model_list=_two()) + core.register_plugin(Refusing()) + with pytest.raises(litellm.BadRequestError, match="refused by strategy"): + await core.acompletion(model="g", messages=MESSAGES, num_retries=0) # pyright: ignore[reportAttributeAccessIssue] + assert host.calls == [] diff --git a/tests/unit/test_router/core/test_selection_characterization_traffic.py b/tests/unit/test_router/core/test_selection_characterization_traffic.py new file mode 100644 index 00000000000..7793b14ef6e --- /dev/null +++ b/tests/unit/test_router/core/test_selection_characterization_traffic.py @@ -0,0 +1,145 @@ +""" +Characterization of the stateful routing strategies once traffic has gone through the core: +least-busy follows in-flight requests, latency-based follows measured latency, and usage-based +routing follows the tokens and requests each deployment served. The host here is a +``TrafficHost``, which feeds the strategy loggers the way litellm's logging does around a call. +""" + +import asyncio +import re +from typing import Final + +import pytest + +import litellm +from tests.unit.test_router.core.scripted_host import bad_request, deployment +from tests.unit.test_router.core.selection_helpers import MESSAGES, MakeCore, TrafficHost, picked_ids + + +@pytest.fixture +def host() -> TrafficHost: + return TrafficHost() + + +async def _until_called(host: TrafficHost, count: int) -> None: + for _ in range(200): + if len(host.started) >= count: + return + await asyncio.sleep(0.005) + raise AssertionError(f"expected {count} started calls, saw {len(host.started)}") + + +# --- least-busy ------------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_least_busy_avoids_a_deployment_with_a_request_in_flight(make_core: MakeCore, host: TrafficHost) -> None: + core: Final = make_core(model_list=[deployment("g", "a"), deployment("g", "b")], routing_strategy="least-busy") + host.gates["openai/a"] = asyncio.Event() + in_flight: Final = asyncio.create_task(core.acompletion(model="g", messages=MESSAGES)) # pyright: ignore[reportAttributeAccessIssue] + await _until_called(host, 1) + assert host.started == ["openai/a"] + assert await picked_ids(core, "g", times=5, messages=MESSAGES) == {"b"} + host.gates["openai/a"].set() + await in_flight + assert await picked_ids(core, "g", times=5, messages=MESSAGES) == {"a"} + + +@pytest.mark.asyncio +async def test_least_busy_releases_a_deployment_whose_request_failed(make_core: MakeCore, host: TrafficHost) -> None: + core: Final = make_core(model_list=[deployment("g", "a"), deployment("g", "b")], routing_strategy="least-busy") + host.gates["openai/a"] = asyncio.Event() + host.answer("openai/a", bad_request()) + in_flight: Final = asyncio.create_task(core.acompletion(model="g", messages=MESSAGES, num_retries=0)) # pyright: ignore[reportAttributeAccessIssue] + await _until_called(host, 1) + assert await picked_ids(core, "g", times=5, messages=MESSAGES) == {"b"} + host.gates["openai/a"].set() + with pytest.raises(litellm.BadRequestError): + await in_flight + assert await picked_ids(core, "g", times=5, messages=MESSAGES) == {"a"} + + +@pytest.mark.asyncio +async def test_least_busy_spreads_concurrent_requests_over_idle_deployments( + make_core: MakeCore, host: TrafficHost +) -> None: + core: Final = make_core( + model_list=[deployment("g", "a"), deployment("g", "b"), deployment("g", "c")], routing_strategy="least-busy" + ) + gate: Final = asyncio.Event() + for model in ("openai/a", "openai/b", "openai/c"): + host.gates[model] = gate + calls: Final = [] + for count in range(1, 4): + calls.append(asyncio.create_task(core.acompletion(model="g", messages=MESSAGES))) # pyright: ignore[reportAttributeAccessIssue] + await _until_called(host, count) + assert host.started == ["openai/a", "openai/b", "openai/c"] + gate.set() + await asyncio.gather(*calls) + + +# --- latency-based ------------------------------------------------------------------------------ + + +@pytest.mark.asyncio +async def test_latency_based_routing_moves_to_the_faster_deployment(make_core: MakeCore, host: TrafficHost) -> None: + core: Final = make_core( + model_list=[deployment("g", "slow"), deployment("g", "fast")], routing_strategy="latency-based-routing" + ) + host.delays["openai/slow"] = 0.15 + for _ in range(6): + await core.acompletion(model="g", messages=MESSAGES) # pyright: ignore[reportAttributeAccessIssue] + assert "fast" in host.deployments_called() + assert await picked_ids(core, "g", times=10, messages=MESSAGES) == {"fast"} + + +# --- usage-based --------------------------------------------------------------------------------- + + +@pytest.mark.asyncio +@pytest.mark.parametrize("routing_strategy", ["usage-based-routing", "usage-based-routing-v2"]) +async def test_usage_based_routing_alternates_between_deployments_by_tokens_served( + make_core: MakeCore, host: TrafficHost, routing_strategy: str +) -> None: + core: Final = make_core(model_list=[deployment("g", "a"), deployment("g", "b")], routing_strategy=routing_strategy) + for _ in range(4): + await core.acompletion(model="g", messages=MESSAGES) # pyright: ignore[reportAttributeAccessIssue] + called: Final = host.deployments_called() + assert sorted(called[:2]) == ["a", "b"] + assert sorted(called[2:]) == ["a", "b"] + + +@pytest.mark.asyncio +async def test_usage_based_v2_stops_picking_a_deployment_once_its_tpm_is_used( + make_core: MakeCore, host: TrafficHost +) -> None: + core: Final = make_core( + model_list=[deployment("g", "a", tpm=15), deployment("g", "b", tpm=1000)], + routing_strategy="usage-based-routing-v2", + ) + for _ in range(4): + await core.acompletion(model="g", messages=MESSAGES) # pyright: ignore[reportAttributeAccessIssue] + called: Final = host.deployments_called() + assert called.count("a") <= 1 + assert called[-2:] == ["b", "b"] + + +@pytest.mark.asyncio +async def test_usage_based_v2_rate_limits_once_every_deployment_used_its_rpm( + make_core: MakeCore, host: TrafficHost +) -> None: + core: Final = make_core( + model_list=[deployment("g", "a", rpm=1), deployment("g", "b", rpm=1)], + routing_strategy="usage-based-routing-v2", + ) + await core.acompletion(model="g", messages=MESSAGES, num_retries=0) # pyright: ignore[reportAttributeAccessIssue] + await core.acompletion(model="g", messages=MESSAGES, num_retries=0) # pyright: ignore[reportAttributeAccessIssue] + assert sorted(host.deployments_called()) == ["a", "b"] + with pytest.raises(litellm.RateLimitError) as raised: + await core.acompletion(model="g", messages=MESSAGES, num_retries=0) # pyright: ignore[reportAttributeAccessIssue] + assert re.search( + r"No deployments available for selected model\. Passed model=g\. Deployments=\{'[ab]': \{'current_tpm': 10, " + r"'tpm_limit': inf, 'current_rpm': 1, 'rpm_limit': 1\}", + str(raised.value), + ) + assert len(host.calls) == 2 diff --git a/tests/unit/test_router/core/test_streaming_characterization_anthropic.py b/tests/unit/test_router/core/test_streaming_characterization_anthropic.py new file mode 100644 index 00000000000..6328767f96e --- /dev/null +++ b/tests/unit/test_router/core/test_streaming_characterization_anthropic.py @@ -0,0 +1,517 @@ +""" +Characterization of /v1/messages streaming through a router core: lifecycle frames held back while a +retry or fallback can still take over, provider ``event: error`` frames, raised stream errors, safeguard +refusals, same-group stream retries and their budget, and the hidden params the returned stream carries. + +Scripted streams are async iterators of raw SSE byte frames (``streaming_helpers.anthropic_stream``). +""" + +from collections.abc import Callable, Mapping +from typing import Final + +import pytest + +import litellm +from litellm.router_core import RouterCore +from litellm.types.router import RetryPolicy +from litellm.types.utils import Usage +from tests.unit.test_router.core.scripted_host import ( + ScriptedHost, + agentic_followup_failure, + bad_request, + deployment, + rate_limit, + server_error, +) +from tests.unit.test_router.core.streaming_helpers import ( + BLOCK_START, + MESSAGE_START, + MESSAGE_STOP, + MESSAGES, + PING, + ByteStream, + anthropic_stream, + collect_bytes, + error_frame, + install_orjson_stub_if_missing, + refusal_delta, + request_logging, + text_delta, +) + +install_orjson_stub_if_missing() + +MakeCore = Callable[..., RouterCore] +GROUPS_AB: Final = [deployment("a", "a"), deployment("b", "b")] +REFUSAL_DETAILS: Final = {"type": "refusal", "category": "cyber"} + + +def event_names(frames: list[bytes]) -> list[bytes]: + return [frame.split(b"\n", 1)[0].removeprefix(b"event: ") for frame in frames] + + +def fallback_ok(opened: list[ByteStream] | None = None) -> Callable[..., ByteStream]: + return anthropic_stream( + MESSAGE_START, + text_delta("B"), + MESSAGE_STOP, + hidden_params={"model_id": "b-stream", "additional_headers": {"x-provider-request-id": "req-b"}}, + opened=opened, + ) + + +def answered_ok(opened: list[ByteStream] | None = None) -> Callable[..., ByteStream]: + return anthropic_stream(MESSAGE_START, text_delta("A"), MESSAGE_STOP, opened=opened) + + +async def open_stream(core: RouterCore, **request: object) -> object: # kwargs-ok: request kwargs + return await core.aanthropic_messages(model="a", messages=MESSAGES, max_tokens=5, stream=True, **request) # pyright: ignore[reportAttributeAccessIssue] # endpoint method + + +def hidden(stream: object) -> Mapping[str, object]: + return getattr(stream, "_hidden_params") + + +def headers(stream: object) -> Mapping[str, object]: + return hidden(stream)["additional_headers"] # pyright: ignore[reportReturnType] # a dict of headers + + +# --- provider error frames before content ------------------------------------------------------ + + +async def test_an_error_frame_before_content_falls_back_and_drops_the_primarys_held_frames( + make_core: MakeCore, host: ScriptedHost +) -> None: + host.answer("openai/a", anthropic_stream(MESSAGE_START, BLOCK_START, error_frame())) + host.answer("openai/b", fallback_ok()) + core = make_core(model_list=GROUPS_AB, fallbacks=[{"a": ["b"]}], num_retries=0) + frames = await collect_bytes(await open_stream(core)) + assert frames == [MESSAGE_START, text_delta("B"), MESSAGE_STOP] + assert host.deployments_called() == ["a", "b"] + + +async def test_an_error_frame_before_content_with_nothing_to_take_over_is_forwarded_verbatim( + make_core: MakeCore, host: ScriptedHost +) -> None: + host.answer("openai/a", anthropic_stream(MESSAGE_START, error_frame())) + core = make_core(model_list=GROUPS_AB, num_retries=0) + frames = await collect_bytes(await open_stream(core)) + assert frames == [MESSAGE_START, error_frame()] + assert host.deployments_called() == ["a"] + + +async def test_an_error_frame_split_across_byte_chunks_is_still_recognised( + make_core: MakeCore, host: ScriptedHost +) -> None: + frame = error_frame() + host.answer("openai/a", anthropic_stream(MESSAGE_START, frame[:25], frame[25:])) + host.answer("openai/b", fallback_ok()) + core = make_core(model_list=GROUPS_AB, fallbacks=[{"a": ["b"]}], num_retries=0) + frames = await collect_bytes(await open_stream(core)) + assert frames == [MESSAGE_START, text_delta("B"), MESSAGE_STOP] + assert host.deployments_called() == ["a", "b"] + + +async def test_an_error_frame_after_content_is_forwarded_without_fallback( + make_core: MakeCore, host: ScriptedHost +) -> None: + host.answer("openai/a", anthropic_stream(MESSAGE_START, text_delta("x"), error_frame())) + host.answer("openai/b", fallback_ok()) + core = make_core(model_list=GROUPS_AB, fallbacks=[{"a": ["b"]}], num_retries=0) + frames = await collect_bytes(await open_stream(core)) + assert frames == [MESSAGE_START, text_delta("x"), error_frame()] + assert host.deployments_called() == ["a"] + + +async def test_a_non_retriable_error_frame_is_forwarded_even_with_a_fallback( + make_core: MakeCore, host: ScriptedHost +) -> None: + frame = error_frame("invalid_request_error", "bad request") + host.answer("openai/a", anthropic_stream(MESSAGE_START, frame)) + host.answer("openai/b", fallback_ok()) + core = make_core(model_list=GROUPS_AB, fallbacks=[{"a": ["b"]}], num_retries=0) + frames = await collect_bytes(await open_stream(core)) + assert frames == [MESSAGE_START, frame] + assert host.deployments_called() == ["a"] + + +async def test_disable_fallbacks_forwards_a_retriable_error_frame(make_core: MakeCore, host: ScriptedHost) -> None: + host.answer("openai/a", anthropic_stream(MESSAGE_START, error_frame())) + host.answer("openai/b", fallback_ok()) + core = make_core(model_list=GROUPS_AB, fallbacks=[{"a": ["b"]}], num_retries=0) + frames = await collect_bytes(await open_stream(core, disable_fallbacks=True)) + assert frames == [MESSAGE_START, error_frame()] + assert host.deployments_called() == ["a"] + + +async def test_when_the_fallback_after_an_error_frame_fails_the_frames_error_is_raised( + make_core: MakeCore, host: ScriptedHost +) -> None: + host.answer("openai/a", anthropic_stream(MESSAGE_START, error_frame("overloaded_error", "primary overloaded"))) + host.answer("openai/b", server_error()) + core = make_core(model_list=GROUPS_AB, fallbacks=[{"a": ["b"]}], num_retries=0) + stream = await open_stream(core) + with pytest.raises(litellm.InternalServerError, match="primary overloaded"): + await collect_bytes(stream) + assert host.deployments_called() == ["a", "b"] + + +# --- held lifecycle frames -------------------------------------------------------------------- + + +async def test_a_ping_is_forwarded_live_while_lifecycle_frames_are_held( + make_core: MakeCore, host: ScriptedHost +) -> None: + host.answer("openai/a", anthropic_stream(MESSAGE_START, PING, error_frame())) + host.answer("openai/b", fallback_ok()) + core = make_core(model_list=GROUPS_AB, fallbacks=[{"a": ["b"]}], num_retries=0) + frames = await collect_bytes(await open_stream(core)) + assert frames == [PING, MESSAGE_START, text_delta("B"), MESSAGE_STOP] + + +async def test_with_nothing_to_take_over_every_frame_is_forwarded_live_in_order( + make_core: MakeCore, host: ScriptedHost +) -> None: + host.answer("openai/a", anthropic_stream(MESSAGE_START, PING, BLOCK_START, error_frame())) + core = make_core(model_list=GROUPS_AB, num_retries=0) + frames = await collect_bytes(await open_stream(core)) + assert frames == [MESSAGE_START, PING, BLOCK_START, error_frame()] + + +async def test_held_frames_are_flushed_in_order_when_content_arrives(make_core: MakeCore, host: ScriptedHost) -> None: + host.answer("openai/a", anthropic_stream(MESSAGE_START, BLOCK_START, text_delta("A"), MESSAGE_STOP)) + core = make_core(model_list=GROUPS_AB, fallbacks=[{"a": ["b"]}], num_retries=0) + frames = await collect_bytes(await open_stream(core)) + assert frames == [MESSAGE_START, BLOCK_START, text_delta("A"), MESSAGE_STOP] + assert host.deployments_called() == ["a"] + + +async def test_held_frames_are_flushed_when_the_stream_ends_without_content( + make_core: MakeCore, host: ScriptedHost +) -> None: + host.answer("openai/a", anthropic_stream(MESSAGE_START, BLOCK_START, MESSAGE_STOP)) + core = make_core(model_list=GROUPS_AB, fallbacks=[{"a": ["b"]}], num_retries=0) + frames = await collect_bytes(await open_stream(core)) + assert frames == [MESSAGE_START, BLOCK_START, MESSAGE_STOP] + assert host.deployments_called() == ["a"] + + +async def test_two_hundred_held_frames_commit_the_stream_so_a_later_error_frame_is_forwarded( + make_core: MakeCore, host: ScriptedHost +) -> None: + host.answer("openai/a", anthropic_stream(*([BLOCK_START] * 201), error_frame())) + host.answer("openai/b", fallback_ok()) + core = make_core(model_list=GROUPS_AB, fallbacks=[{"a": ["b"]}], num_retries=0) + frames = await collect_bytes(await open_stream(core)) + assert frames == [*([BLOCK_START] * 201), error_frame()] + assert host.deployments_called() == ["a"] + + +# --- raised stream errors ---------------------------------------------------------------------- + + +async def test_a_raised_rate_limit_before_content_falls_back(make_core: MakeCore, host: ScriptedHost) -> None: + host.answer("openai/a", anthropic_stream(MESSAGE_START, error=rate_limit())) + host.answer("openai/b", fallback_ok()) + core = make_core(model_list=GROUPS_AB, fallbacks=[{"a": ["b"]}], num_retries=0) + frames = await collect_bytes(await open_stream(core)) + assert frames == [MESSAGE_START, text_delta("B"), MESSAGE_STOP] + assert host.deployments_called() == ["a", "b"] + + +async def test_a_raised_bad_request_before_content_is_raised_without_fallback( + make_core: MakeCore, host: ScriptedHost +) -> None: + host.answer("openai/a", anthropic_stream(MESSAGE_START, error=bad_request(message="raised 400"))) + host.answer("openai/b", fallback_ok()) + core = make_core(model_list=GROUPS_AB, fallbacks=[{"a": ["b"]}], num_retries=0) + stream = await open_stream(core) + with pytest.raises(litellm.BadRequestError, match="raised 400"): + await collect_bytes(stream) + assert host.deployments_called() == ["a"] + + +async def test_a_raised_error_after_content_is_raised_without_fallback(make_core: MakeCore, host: ScriptedHost) -> None: + host.answer("openai/a", anthropic_stream(MESSAGE_START, text_delta("x"), error=rate_limit(message="after"))) + host.answer("openai/b", fallback_ok()) + core = make_core(model_list=GROUPS_AB, fallbacks=[{"a": ["b"]}], num_retries=0) + stream = await open_stream(core) + with pytest.raises(litellm.RateLimitError, match="after"): + await collect_bytes(stream) + assert host.deployments_called() == ["a"] + + +async def test_a_raised_server_error_before_content_is_retried_within_the_group( + make_core: MakeCore, host: ScriptedHost +) -> None: + host.answer("*", anthropic_stream(MESSAGE_START, error=server_error()), answered_ok()) + core = make_core(model_list=[deployment("a", "a1"), deployment("a", "a2")], num_retries=1) + frames = await collect_bytes(await open_stream(core)) + assert frames == [MESSAGE_START, text_delta("A"), MESSAGE_STOP] + assert len(host.deployments_called()) == 2 + assert len(host.sleeps) == 1 + + +# --- same-group stream retries ----------------------------------------------------------------- + + +async def test_a_failed_agentic_followup_before_content_skips_the_retries_for_the_fallback( + make_core: MakeCore, host: ScriptedHost +) -> None: + failed: Final = agentic_followup_failure(server_error()) + host.answer("openai/a", anthropic_stream(MESSAGE_START, error=failed), answered_ok()) + host.answer("openai/b", fallback_ok()) + core = make_core(model_list=GROUPS_AB, fallbacks=[{"a": ["b"]}], num_retries=2) + frames = await collect_bytes(await open_stream(core)) + assert frames == [MESSAGE_START, text_delta("B"), MESSAGE_STOP] + assert host.deployments_called() == ["a", "b"] + assert host.sleeps == [] + + +async def test_a_stream_retry_carries_the_superseded_attempt_usage(make_core: MakeCore, host: ScriptedHost) -> None: + logging_obj: Final = request_logging() + usage: Final = Usage(prompt_tokens=1, completion_tokens=1, total_tokens=2) + failing: Final = anthropic_stream(MESSAGE_START, error=server_error()) + + def fails_after_usage(**kwargs: object) -> ByteStream: # kwargs-ok: the request + logging_obj.model_call_details["combined_usage_object"] = usage + return failing(**kwargs) + + host.answer("openai/a", fails_after_usage, answered_ok()) + core = make_core(model_list=[deployment("a", "a")], num_retries=1) + await collect_bytes(await open_stream(core, litellm_logging_obj=logging_obj)) + assert host.deployments_called() == ["a", "a"] + assert logging_obj.model_call_details["failed_attempts_usage"] == ((usage, 0.0),) + + +async def test_an_error_frame_before_content_is_retried_within_the_group( + make_core: MakeCore, host: ScriptedHost +) -> None: + host.answer("openai/a", anthropic_stream(MESSAGE_START, error_frame()), answered_ok()) + core = make_core(model_list=[deployment("a", "a")], num_retries=2) + stream = await open_stream(core) + frames = await collect_bytes(stream) + assert frames == [MESSAGE_START, text_delta("A"), MESSAGE_STOP] + assert host.deployments_called() == ["a", "a"] + assert len(host.sleeps) == 1 + assert headers(stream)["x-litellm-attempted-retries"] == 1 + assert headers(stream)["x-litellm-max-retries"] == 2 + + +async def test_a_retry_policy_grants_stream_retries_when_num_retries_is_zero( + make_core: MakeCore, host: ScriptedHost +) -> None: + host.answer("openai/a", anthropic_stream(MESSAGE_START, error_frame("rate_limit_error", "slow")), answered_ok()) + core = make_core( + model_list=[deployment("a", "a")], num_retries=0, retry_policy=RetryPolicy(RateLimitErrorRetries=2) + ) + stream = await open_stream(core) + frames = await collect_bytes(stream) + assert frames == [MESSAGE_START, text_delta("A"), MESSAGE_STOP] + assert host.deployments_called() == ["a", "a"] + assert headers(stream)["x-litellm-max-retries"] == 2 + + +async def test_request_num_retries_zero_opts_out_of_the_retry_policy_for_a_stream( + make_core: MakeCore, host: ScriptedHost +) -> None: + frame = error_frame("rate_limit_error", "slow") + host.answer("openai/a", anthropic_stream(MESSAGE_START, frame), answered_ok()) + core = make_core( + model_list=[deployment("a", "a")], num_retries=0, retry_policy=RetryPolicy(RateLimitErrorRetries=2) + ) + frames = await collect_bytes(await open_stream(core, num_retries=0)) + assert frames == [MESSAGE_START, frame] + assert host.deployments_called() == ["a"] + + +async def test_the_deployment_num_retries_is_the_stream_retry_budget(make_core: MakeCore, host: ScriptedHost) -> None: + host.answer("openai/a", anthropic_stream(MESSAGE_START, error_frame()), answered_ok()) + core = make_core(model_list=[deployment("a", "a", num_retries=1)], num_retries=0) + stream = await open_stream(core) + frames = await collect_bytes(stream) + assert frames == [MESSAGE_START, text_delta("A"), MESSAGE_STOP] + assert headers(stream)["x-litellm-max-retries"] == 1 + + +async def test_stream_retries_stop_at_the_budget_and_forward_the_last_error_frame( + make_core: MakeCore, host: ScriptedHost +) -> None: + host.answer("openai/a", anthropic_stream(MESSAGE_START, error_frame())) + core = make_core(model_list=[deployment("a", "a")], num_retries=2) + frames = await collect_bytes(await open_stream(core)) + assert frames == [MESSAGE_START, error_frame()] + assert host.deployments_called() == ["a", "a", "a"] + assert len(host.sleeps) == 2 + + +async def test_spent_stream_retries_hand_over_to_the_fallback_chain(make_core: MakeCore, host: ScriptedHost) -> None: + host.answer("openai/a", anthropic_stream(MESSAGE_START, error_frame())) + host.answer("openai/b", fallback_ok()) + core = make_core(model_list=GROUPS_AB, fallbacks=[{"a": ["b"]}], num_retries=1) + stream = await open_stream(core) + frames = await collect_bytes(stream) + assert frames == [MESSAGE_START, text_delta("B"), MESSAGE_STOP] + assert host.deployments_called() == ["a", "a", "b"] + assert len(host.sleeps) == 1 + + +async def test_a_superseded_attempt_and_every_later_stream_are_closed(make_core: MakeCore, host: ScriptedHost) -> None: + opened: list[ByteStream] = [] + host.answer("openai/a", anthropic_stream(MESSAGE_START, error_frame(), opened=opened)) + host.answer("openai/b", fallback_ok(opened)) + core = make_core(model_list=GROUPS_AB, fallbacks=[{"a": ["b"]}], num_retries=1) + await collect_bytes(await open_stream(core)) + assert len(opened) == 3 + assert all(stream.closed for stream in opened) + + +# --- safeguard refusals ------------------------------------------------------------------------ + + +async def test_a_refusal_with_stop_details_before_content_falls_back(make_core: MakeCore, host: ScriptedHost) -> None: + host.answer("openai/a", anthropic_stream(MESSAGE_START, refusal_delta(REFUSAL_DETAILS), MESSAGE_STOP)) + host.answer("openai/b", fallback_ok()) + core = make_core(model_list=GROUPS_AB, fallbacks=[{"a": ["b"]}], num_retries=0) + frames = await collect_bytes(await open_stream(core)) + assert frames == [MESSAGE_START, text_delta("B"), MESSAGE_STOP] + assert host.deployments_called() == ["a", "b"] + + +async def test_a_refusal_falls_back_through_the_content_policy_chain(make_core: MakeCore, host: ScriptedHost) -> None: + host.answer("openai/a", anthropic_stream(MESSAGE_START, refusal_delta(REFUSAL_DETAILS), MESSAGE_STOP)) + host.answer("openai/b", fallback_ok()) + core = make_core(model_list=GROUPS_AB, content_policy_fallbacks=[{"a": ["b"]}], num_retries=0) + frames = await collect_bytes(await open_stream(core)) + assert frames == [MESSAGE_START, text_delta("B"), MESSAGE_STOP] + assert host.deployments_called() == ["a", "b"] + + +async def test_a_refusal_without_stop_details_is_forwarded(make_core: MakeCore, host: ScriptedHost) -> None: + host.answer("openai/a", anthropic_stream(MESSAGE_START, refusal_delta(), MESSAGE_STOP)) + host.answer("openai/b", fallback_ok()) + core = make_core(model_list=GROUPS_AB, fallbacks=[{"a": ["b"]}], num_retries=0) + frames = await collect_bytes(await open_stream(core)) + assert frames == [MESSAGE_START, refusal_delta(), MESSAGE_STOP] + assert host.deployments_called() == ["a"] + + +@pytest.mark.parametrize( + ("router_arguments", "request_kwargs"), + [({"fallbacks": [{"a": ["b"]}]}, {"disable_fallbacks": True}), ({}, {})], + ids=["fallbacks-disabled", "no-fallbacks"], +) +async def test_a_refusal_with_no_reachable_fallback_is_forwarded( + make_core: MakeCore, host: ScriptedHost, router_arguments: dict[str, object], request_kwargs: dict[str, object] +) -> None: + host.answer("openai/a", anthropic_stream(MESSAGE_START, refusal_delta(REFUSAL_DETAILS), MESSAGE_STOP)) + host.answer("openai/b", fallback_ok()) + core = make_core(model_list=GROUPS_AB, num_retries=0, **router_arguments) + frames = await collect_bytes(await open_stream(core, **request_kwargs)) + assert frames == [MESSAGE_START, refusal_delta(REFUSAL_DETAILS), MESSAGE_STOP] + assert host.deployments_called() == ["a"] + + +async def test_a_refusal_after_content_is_forwarded(make_core: MakeCore, host: ScriptedHost) -> None: + host.answer("openai/a", anthropic_stream(MESSAGE_START, text_delta("x"), refusal_delta(REFUSAL_DETAILS))) + host.answer("openai/b", fallback_ok()) + core = make_core(model_list=GROUPS_AB, fallbacks=[{"a": ["b"]}], num_retries=0) + frames = await collect_bytes(await open_stream(core)) + assert frames == [MESSAGE_START, text_delta("x"), refusal_delta(REFUSAL_DETAILS)] + assert host.deployments_called() == ["a"] + + +# --- what the fallback hands back -------------------------------------------------------------- + + +async def test_a_fallback_that_answers_without_streaming_is_replayed_as_an_sse_lifecycle( + make_core: MakeCore, host: ScriptedHost +) -> None: + host.answer("openai/a", anthropic_stream(MESSAGE_START, error_frame())) + host.answer( + "openai/b", + { + "id": "msg_b", + "type": "message", + "role": "assistant", + "model": "b", + "content": [{"type": "text", "text": "B"}], + "stop_reason": "end_turn", + "stop_sequence": None, + "usage": {"input_tokens": 1, "output_tokens": 1}, + }, + ) + core = make_core(model_list=GROUPS_AB, fallbacks=[{"a": ["b"]}], num_retries=0) + frames = await collect_bytes(await open_stream(core)) + assert event_names(frames) == [ + b"message_start", + b"content_block_start", + b"content_block_delta", + b"content_block_stop", + b"message_delta", + b"message_stop", + ] + assert b'"B"' in frames[2] + + +# --- hidden params on the returned stream ------------------------------------------------------ + + +async def test_a_stream_without_recovery_reports_the_serving_deployment( + make_core: MakeCore, host: ScriptedHost +) -> None: + host.answer("openai/a", answered_ok()) + core = make_core(model_list=GROUPS_AB, fallbacks=[{"a": ["b"]}], num_retries=0) + stream = await open_stream(core) + await collect_bytes(stream) + assert hidden(stream)["model_id"] == "a" + assert headers(stream) == { + "x-litellm-model-group": "a", + "x-litellm-attempted-retries": 0, + "x-litellm-attempted-fallbacks": 0, + } + + +async def test_after_a_stream_fallback_the_stream_carries_the_fallbacks_hidden_params_and_headers( + make_core: MakeCore, host: ScriptedHost +) -> None: + host.answer("openai/a", anthropic_stream(MESSAGE_START, error_frame())) + host.answer("openai/b", fallback_ok()) + core = make_core(model_list=GROUPS_AB, fallbacks=[{"a": ["b"]}], num_retries=0) + stream = await open_stream(core) + await collect_bytes(stream) + assert hidden(stream)["model_id"] == "b-stream" + assert headers(stream) == { + "x-litellm-model-group": "b", + "x-litellm-attempted-retries": 0, + "x-litellm-attempted-fallbacks": 1, + "x-provider-request-id": "req-b", + } + + +async def test_after_retries_then_a_fallback_the_retry_budget_header_lingers( + make_core: MakeCore, host: ScriptedHost +) -> None: + host.answer("openai/a", anthropic_stream(MESSAGE_START, error_frame())) + host.answer("openai/b", answered_ok()) + core = make_core(model_list=GROUPS_AB, fallbacks=[{"a": ["b"]}], num_retries=1) + stream = await open_stream(core) + await collect_bytes(stream) + assert headers(stream) == { + "x-litellm-model-group": "b", + "x-litellm-attempted-retries": 0, + "x-litellm-attempted-fallbacks": 1, + "x-litellm-max-retries": 1, + } + + +async def test_a_stream_retry_keeps_reporting_the_first_deployment_when_the_retry_has_no_hidden_params( + make_core: MakeCore, host: ScriptedHost +) -> None: + host.answer("*", anthropic_stream(MESSAGE_START, error=server_error()), answered_ok()) + core = make_core(model_list=[deployment("a", "a1"), deployment("a", "a2")], num_retries=1) + stream = await open_stream(core) + await collect_bytes(stream) + first, _retried = host.deployments_called() + assert hidden(stream)["model_id"] == first + assert headers(stream)["x-litellm-attempted-retries"] == 1 diff --git a/tests/unit/test_router/core/test_streaming_characterization_chat.py b/tests/unit/test_router/core/test_streaming_characterization_chat.py new file mode 100644 index 00000000000..260ad8534cf --- /dev/null +++ b/tests/unit/test_router/core/test_streaming_characterization_chat.py @@ -0,0 +1,737 @@ +""" +Characterization of chat-completion streaming through a router core: the first chunk read inside the +attempt, mid-stream errors and their fallback, mid-stream continuation, the sync twin, and the stream +timeout sent to the host. + +Each scripted stream is a ``CustomStreamWrapper`` built by ``streaming_helpers.chat_stream``; a stream +that raises before its first chunk fails inside the attempt (retries and fallbacks run as for any error), +while one that raises later fails while the caller iterates. +""" + +import asyncio +from collections.abc import Callable, Mapping +from typing import Final + +import pytest + +import litellm +from litellm.exceptions import MidStreamFallbackError +from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging +from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper +from litellm.router_core import RouterCore +from litellm.types.utils import Usage +from tests.unit.test_router.core.scripted_host import ( + ScriptedHost, + bad_request, + content_policy, + context_window, + deployment, + rate_limit, + server_error, +) +from tests.unit.test_router.core.streaming_helpers import ( + MESSAGES, + chat_stream, + collect_texts, + collect_texts_sync, + drain_into, +) + +MakeCore = Callable[..., RouterCore] +PREFILL_MODEL: Final = "anthropic/claude-sonnet-4-5" +"""A model whose deployments accept an assistant prefill, so a continuation may land on them.""" +GROUPS_AB: Final = [deployment("a", "a"), deployment("b", "b")] + + +def build(make_core: MakeCore, **arguments: object) -> RouterCore: # kwargs-ok: the router's constructor arguments + return make_core(**{"model_list": GROUPS_AB, "fallbacks": [{"a": ["b"]}], "num_retries": 0, **arguments}) + + +async def open_stream(core: RouterCore, **request: object) -> object: # kwargs-ok: request kwargs + return await core.acompletion(model="a", messages=MESSAGES, stream=True, **request) # pyright: ignore[reportAttributeAccessIssue] # endpoint method + + +def headers(response: object) -> Mapping[str, object]: + return getattr(response, "_hidden_params")["additional_headers"] + + +def model_id(response: object) -> object: + return getattr(response, "_hidden_params").get("model_id") + + +# --- the first chunk is read inside the attempt ----------------------------------------------- + + +async def test_a_stream_that_fails_on_its_first_chunk_is_retried_inside_the_attempt( + make_core: MakeCore, host: ScriptedHost +) -> None: + host.answer("openai/a", chat_stream("x", error=rate_limit(), fail_after=0), chat_stream("ok")) + core = make_core(model_list=[deployment("a", "a")], num_retries=2) + stream = await open_stream(core) + assert await collect_texts(stream) == "ok" + assert host.deployments_called() == ["a", "a"] + assert len(host.sleeps) == 1 + assert headers(stream)["x-litellm-attempted-retries"] == 1 + + +async def test_a_stream_that_fails_on_its_first_chunk_falls_back_after_its_retries( + make_core: MakeCore, host: ScriptedHost +) -> None: + host.answer("openai/a", chat_stream("x", error=rate_limit(), fail_after=0)) + host.answer("openai/b", chat_stream("B")) + core = build(make_core, num_retries=2) + stream = await open_stream(core) + assert await collect_texts(stream) == "B" + assert host.deployments_called() == ["a", "a", "a", "b"] + assert len(host.sleeps) == 3 + assert model_id(stream) == "b" + assert headers(stream)["x-litellm-attempted-fallbacks"] == 1 + + +async def test_a_first_chunk_failure_with_nowhere_to_go_raises_the_provider_error_from_the_call( + make_core: MakeCore, host: ScriptedHost +) -> None: + host.answer("openai/a", chat_stream("x", error=rate_limit(message="first chunk"), fail_after=0)) + core = make_core(model_list=[deployment("a", "a")], num_retries=0) + with pytest.raises(litellm.RateLimitError, match="first chunk"): + await open_stream(core) + + +async def test_a_first_chunk_bad_request_is_raised_from_the_call_without_retry( + make_core: MakeCore, host: ScriptedHost +) -> None: + host.answer("openai/a", chat_stream("x", error=bad_request(message="nope"), fail_after=0)) + core = make_core(model_list=GROUPS_AB, num_retries=2) + with pytest.raises(litellm.BadRequestError, match="nope"): + await open_stream(core) + assert host.deployments_called() == ["a"] + + +async def test_the_first_chunk_already_read_is_still_delivered_to_the_caller( + make_core: MakeCore, host: ScriptedHost +) -> None: + host.answer("openai/a", chat_stream("first", "second", "third")) + core = make_core(model_list=[deployment("a", "a")]) + stream = await open_stream(core) + assert await collect_texts(stream) == "firstsecondthird" + + +async def test_a_stream_that_ends_without_any_chunk_or_finish_reason_ends_quietly_without_text( + make_core: MakeCore, host: ScriptedHost +) -> None: + host.answer("openai/a", chat_stream(finish=False)) + host.answer("openai/b", chat_stream("B")) + core = build(make_core) + stream = await open_stream(core) + assert await collect_texts(stream) == "" + assert host.deployments_called() == ["a"] + + +async def test_a_stream_that_ends_after_content_without_a_finish_reason_raises_while_iterating( + make_core: MakeCore, host: ScriptedHost +) -> None: + host.answer("openai/a", chat_stream("partial", finish=False)) + core = build(make_core) + stream = await open_stream(core) + with pytest.raises(litellm.APIConnectionError, match="Stream ended before a finish_reason was received"): + await collect_texts(stream) + assert host.deployments_called() == ["a"] + + +# --- mid-stream errors ------------------------------------------------------------------------- + + +async def test_a_mid_stream_error_after_content_raises_the_provider_error_not_mid_stream_fallback_error( + make_core: MakeCore, host: ScriptedHost +) -> None: + host.answer("openai/a", chat_stream("he", "llo", error=rate_limit(message="cut after content"))) + host.answer("openai/b", chat_stream("B")) + core = build(make_core) + stream = await open_stream(core) + with pytest.raises(litellm.RateLimitError, match="cut after content") as raised: + await collect_texts(stream) + assert not isinstance(raised.value, MidStreamFallbackError) + assert host.deployments_called() == ["a"] + + +async def test_a_mid_stream_error_after_content_delivers_the_text_before_it( + make_core: MakeCore, host: ScriptedHost +) -> None: + host.answer("openai/a", chat_stream("he", "llo", error=rate_limit())) + core = build(make_core) + stream = await open_stream(core) + received: list[object] = [] + with pytest.raises(litellm.RateLimitError): + await drain_into(stream, received) + assert "".join(getattr(chunk, "choices")[0].delta.content or "" for chunk in received) == "hello" + + +async def test_a_mid_stream_error_before_content_falls_back_to_the_next_group( + make_core: MakeCore, host: ScriptedHost +) -> None: + host.answer("openai/a", chat_stream(error=rate_limit(), leading_role_chunk=True)) + host.answer("openai/b", chat_stream("B")) + core = build(make_core) + stream = await open_stream(core) + assert await collect_texts(stream) == "B" + assert host.deployments_called() == ["a", "b"] + + +async def test_a_mid_stream_error_before_content_does_not_retry_the_same_group( + make_core: MakeCore, host: ScriptedHost +) -> None: + host.answer("openai/a", chat_stream(error=rate_limit(), leading_role_chunk=True)) + host.answer("openai/b", chat_stream("B")) + core = build(make_core, num_retries=2) + stream = await open_stream(core) + assert await collect_texts(stream) == "B" + assert host.deployments_called() == ["a", "b"] + assert host.sleeps == [] + + +async def test_a_mid_stream_fallback_sends_the_original_messages_to_the_fallback( + make_core: MakeCore, host: ScriptedHost +) -> None: + host.answer("openai/a", chat_stream(error=rate_limit(), leading_role_chunk=True)) + host.answer("openai/b", chat_stream("B")) + core = build(make_core) + stream = await open_stream(core) + await collect_texts(stream) + fallback_call = host.calls_to("acompletion")[1] + assert fallback_call.kwargs["messages"] == MESSAGES + assert fallback_call.model_group == "b" + assert fallback_call.kwargs["stream"] is True + + +async def test_a_mid_stream_error_before_content_without_fallbacks_raises_the_provider_error( + make_core: MakeCore, host: ScriptedHost +) -> None: + host.answer("openai/a", chat_stream(error=rate_limit(message="no fallback"), leading_role_chunk=True)) + core = make_core(model_list=GROUPS_AB, num_retries=0) + stream = await open_stream(core) + with pytest.raises(litellm.RateLimitError, match="no fallback") as raised: + await collect_texts(stream) + assert not isinstance(raised.value, MidStreamFallbackError) + assert host.deployments_called() == ["a"] + + +async def test_disable_fallbacks_on_the_request_raises_a_mid_stream_error_before_content( + make_core: MakeCore, host: ScriptedHost +) -> None: + host.answer("openai/a", chat_stream(error=rate_limit(message="disabled"), leading_role_chunk=True)) + host.answer("openai/b", chat_stream("B")) + core = build(make_core) + stream = await open_stream(core, disable_fallbacks=True) + with pytest.raises(litellm.RateLimitError, match="disabled"): + await collect_texts(stream) + assert host.deployments_called() == ["a"] + + +async def test_a_per_request_fallback_list_is_used_for_a_mid_stream_fallback( + make_core: MakeCore, host: ScriptedHost +) -> None: + host.answer("openai/a", chat_stream(error=rate_limit(), leading_role_chunk=True)) + host.answer("openai/c", chat_stream("C")) + core = build(make_core, model_list=[*GROUPS_AB, deployment("c", "c")]) + stream = await open_stream(core, fallbacks=["c"]) + assert await collect_texts(stream) == "C" + assert host.deployments_called() == ["a", "c"] + + +async def test_a_mid_stream_fallback_walks_the_chain_until_a_group_answers( + make_core: MakeCore, host: ScriptedHost +) -> None: + host.answer("openai/a", chat_stream(error=rate_limit(), leading_role_chunk=True)) + host.answer("openai/b", server_error()) + host.answer("openai/c", chat_stream("C")) + core = build(make_core, model_list=[*GROUPS_AB, deployment("c", "c")], fallbacks=[{"a": ["b", "c"]}]) + stream = await open_stream(core) + assert await collect_texts(stream) == "C" + assert host.deployments_called() == ["a", "b", "c"] + + +async def test_when_the_mid_stream_fallback_also_fails_the_primary_provider_error_is_raised( + make_core: MakeCore, host: ScriptedHost +) -> None: + host.answer("openai/a", chat_stream(error=rate_limit(message="primary"), leading_role_chunk=True)) + host.answer("openai/b", server_error()) + core = build(make_core) + stream = await open_stream(core) + with pytest.raises(litellm.RateLimitError, match="primary"): + await collect_texts(stream) + assert host.deployments_called() == ["a", "b"] + + +async def test_a_mid_stream_server_error_before_content_falls_through_order_tiers( + make_core: MakeCore, host: ScriptedHost +) -> None: + host.answer("openai/a1", chat_stream(error=server_error(), leading_role_chunk=True)) + host.answer("openai/a2", chat_stream("A2")) + core = make_core(model_list=[deployment("a", "a1", order=1), deployment("a", "a2", order=2)], num_retries=0) + stream = await open_stream(core) + assert await collect_texts(stream) == "A2" + assert host.deployments_called() == ["a1", "a2"] + assert model_id(stream) == "a2" + assert headers(stream)["x-litellm-model-group"] == "a" + + +async def test_a_mid_stream_bad_request_before_content_is_raised_without_fallback( + make_core: MakeCore, host: ScriptedHost +) -> None: + host.answer("openai/a", chat_stream(error=bad_request(message="client error"), leading_role_chunk=True)) + host.answer("openai/b", chat_stream("B")) + core = build(make_core) + stream = await open_stream(core) + with pytest.raises(litellm.BadRequestError, match="client error"): + await collect_texts(stream) + assert host.deployments_called() == ["a"] + + +async def test_a_mid_stream_content_policy_error_before_content_is_raised_without_fallback( + make_core: MakeCore, host: ScriptedHost +) -> None: + host.answer("openai/a", chat_stream(error=content_policy(), leading_role_chunk=True)) + host.answer("openai/b", chat_stream("B")) + core = build(make_core, fallbacks=None, content_policy_fallbacks=[{"a": ["b"]}]) + stream = await open_stream(core) + with pytest.raises(litellm.ContentPolicyViolationError): + await collect_texts(stream) + assert host.deployments_called() == ["a"] + + +async def test_a_mid_stream_context_window_error_before_content_uses_the_context_window_chain( + make_core: MakeCore, host: ScriptedHost +) -> None: + host.answer("openai/a", chat_stream(error=context_window(), leading_role_chunk=True)) + host.answer("openai/b", chat_stream("B")) + core = build(make_core, fallbacks=None, context_window_fallbacks=[{"a": ["b"]}]) + stream = await open_stream(core) + assert await collect_texts(stream) == "B" + assert host.deployments_called() == ["a", "b"] + + +async def test_a_mid_stream_context_window_error_after_content_is_raised( + make_core: MakeCore, host: ScriptedHost +) -> None: + host.answer("openai/a", chat_stream("x", error=context_window())) + host.answer("openai/b", chat_stream("B")) + core = build(make_core, context_window_fallbacks=[{"a": ["b"]}]) + stream = await open_stream(core) + with pytest.raises(litellm.ContextWindowExceededError): + await collect_texts(stream) + assert host.deployments_called() == ["a"] + + +# --- hidden params on a streamed response ----------------------------------------------------- + + +async def test_a_stream_without_fallback_reports_the_serving_deployment( + make_core: MakeCore, host: ScriptedHost +) -> None: + host.answer("openai/a", chat_stream("x")) + core = build(make_core) + stream = await open_stream(core) + await collect_texts(stream) + assert model_id(stream) == "a" + assert headers(stream)["x-litellm-model-group"] == "a" + assert headers(stream)["x-litellm-attempted-fallbacks"] == 0 + + +async def test_after_a_mid_stream_fallback_the_stream_reports_the_fallback_deployment( + make_core: MakeCore, host: ScriptedHost +) -> None: + host.answer("openai/a", chat_stream(error=rate_limit(), leading_role_chunk=True)) + host.answer("openai/b", chat_stream("B")) + core = build(make_core) + stream = await open_stream(core) + await collect_texts(stream) + assert model_id(stream) == "b" + assert headers(stream) == { + "x-litellm-model-group": "b", + "x-litellm-attempted-retries": 0, + "x-litellm-attempted-fallbacks": 1, + } + + +async def test_after_a_mid_stream_fallback_each_fallback_chunk_carries_the_fallback_deployment( + make_core: MakeCore, host: ScriptedHost +) -> None: + host.answer("openai/a", chat_stream(error=rate_limit(), leading_role_chunk=True)) + host.answer("openai/b", chat_stream("B")) + core = build(make_core) + stream = await open_stream(core) + chunks = [chunk async for chunk in stream] # pyright: ignore[reportGeneralTypeIssues] # an async stream + fallback_chunks = chunks[1:] + assert fallback_chunks + assert all(chunk._hidden_params.get("model_id") == "b" for chunk in fallback_chunks) + assert fallback_chunks[-1]._hidden_params["additional_headers"]["x-litellm-model-group"] == "b" + + +async def test_when_the_mid_stream_fallback_fails_the_stream_still_reports_the_primary( + make_core: MakeCore, host: ScriptedHost +) -> None: + host.answer("openai/a", chat_stream(error=rate_limit(), leading_role_chunk=True)) + host.answer("openai/b", server_error()) + core = build(make_core) + stream = await open_stream(core) + with pytest.raises(litellm.RateLimitError): + await collect_texts(stream) + assert model_id(stream) == "a" + assert headers(stream)["x-litellm-model-group"] == "a" + + +# --- mid-stream continuation ------------------------------------------------------------------- + + +def continuation_core(make_core: MakeCore, **arguments: object) -> RouterCore: # kwargs-ok: constructor arguments + return build( + make_core, + model_list=[deployment("a", "a"), deployment("b", "b", model=PREFILL_MODEL)], + enable_mid_stream_fallback_continuation=True, + **arguments, + ) + + +async def test_continuation_off_by_default_raises_after_content(make_core: MakeCore, host: ScriptedHost) -> None: + host.answer("openai/a", chat_stream("he", "llo", error=rate_limit(message="cut"))) + host.answer(PREFILL_MODEL, chat_stream(" world")) + core = build(make_core, model_list=[deployment("a", "a"), deployment("b", "b", model=PREFILL_MODEL)]) + stream = await open_stream(core) + with pytest.raises(litellm.RateLimitError, match="cut"): + await collect_texts(stream) + assert host.deployments_called() == ["a"] + + +async def test_continuation_continues_after_content_on_a_prefill_capable_fallback( + make_core: MakeCore, host: ScriptedHost +) -> None: + host.answer("openai/a", chat_stream("he", "llo", error=rate_limit())) + host.answer(PREFILL_MODEL, chat_stream(" world")) + core = continuation_core(make_core) + stream = await open_stream(core) + assert await collect_texts(stream) == "hello world" + assert host.deployments_called() == ["a", "b"] + + +async def test_continuation_sends_the_partial_text_as_an_assistant_prefill( + make_core: MakeCore, host: ScriptedHost +) -> None: + host.answer("openai/a", chat_stream("he", "llo", error=rate_limit())) + host.answer(PREFILL_MODEL, chat_stream(" world")) + core = continuation_core(make_core) + stream = await open_stream(core) + await collect_texts(stream) + assert host.calls[1].kwargs["messages"] == [*MESSAGES, {"role": "assistant", "content": "hello", "prefix": True}] + + +async def test_continuation_extends_a_trailing_prefill_turn_in_place(make_core: MakeCore, host: ScriptedHost) -> None: + host.answer("openai/a", chat_stream("Once", " upon", error=rate_limit())) + host.answer(PREFILL_MODEL, chat_stream(" a time")) + core = continuation_core(make_core) + messages = [*MESSAGES, {"role": "assistant", "content": "Once", "prefix": True}] + stream = await core.acompletion(model="a", messages=messages, stream=True) # pyright: ignore[reportAttributeAccessIssue] # endpoint method + await collect_texts(stream) + sent = host.calls[1].kwargs["messages"] + assert sent[-1] == {"role": "assistant", "content": "Once upon", "prefix": True} + assert len(sent) == 2 + + +async def test_continuation_declines_when_the_request_ends_in_a_plain_assistant_turn( + make_core: MakeCore, host: ScriptedHost +) -> None: + host.answer("openai/a", chat_stream("x", error=rate_limit(message="cut"))) + host.answer(PREFILL_MODEL, chat_stream("y")) + core = continuation_core(make_core) + messages = [*MESSAGES, {"role": "assistant", "content": "earlier"}] + stream = await core.acompletion(model="a", messages=messages, stream=True) # pyright: ignore[reportAttributeAccessIssue] # endpoint method + with pytest.raises(litellm.RateLimitError, match="cut"): + await collect_texts(stream) + assert host.deployments_called() == ["a"] + + +async def test_continuation_reduces_the_output_ceiling_by_the_tokens_already_emitted( + make_core: MakeCore, host: ScriptedHost +) -> None: + host.answer("openai/a", chat_stream("hello", " there", error=rate_limit())) + host.answer(PREFILL_MODEL, chat_stream(" friend")) + core = continuation_core(make_core) + stream = await open_stream(core, max_tokens=50) + await collect_texts(stream) + assert [call.kwargs["max_tokens"] for call in host.calls] == [50, 48] + + +async def test_continuation_raises_the_provider_error_when_the_output_ceiling_is_spent( + make_core: MakeCore, host: ScriptedHost +) -> None: + host.answer("openai/a", chat_stream("hello", " there", error=rate_limit(message="spent"))) + host.answer(PREFILL_MODEL, chat_stream(" friend")) + core = continuation_core(make_core) + stream = await open_stream(core, max_tokens=2) + with pytest.raises(litellm.RateLimitError, match="spent"): + await collect_texts(stream) + assert host.deployments_called() == ["a"] + + +async def test_continuation_never_lands_on_a_fallback_that_cannot_take_a_prefill( + make_core: MakeCore, host: ScriptedHost +) -> None: + host.answer("openai/a", chat_stream("he", "llo", error=rate_limit(message="no prefill target"))) + host.answer("openai/b", chat_stream("B")) + core = build(make_core, enable_mid_stream_fallback_continuation=True) + stream = await open_stream(core) + with pytest.raises(litellm.RateLimitError, match="no prefill target"): + await collect_texts(stream) + assert host.deployments_called() == ["a"] + + +@pytest.mark.parametrize( + "request_kwargs", + [ + {"response_format": {"type": "json_object"}}, + {"tool_choice": "required", "tools": [{"type": "function", "function": {"name": "f", "parameters": {}}}]}, + ], + ids=["json-response-format", "forced-tool-choice"], +) +async def test_continuation_declines_a_constrained_request( + make_core: MakeCore, host: ScriptedHost, request_kwargs: dict[str, object] +) -> None: + host.answer("openai/a", chat_stream("he", "llo", error=rate_limit(message="constrained"))) + host.answer(PREFILL_MODEL, chat_stream(" world")) + core = continuation_core(make_core) + stream = await open_stream(core, **request_kwargs) + with pytest.raises(litellm.RateLimitError, match="constrained"): + await collect_texts(stream) + assert host.deployments_called() == ["a"] + + +async def test_continuation_declines_a_fallback_target_that_overrides_request_params( + make_core: MakeCore, host: ScriptedHost +) -> None: + host.answer("openai/a", chat_stream("he", "llo", error=rate_limit(message="override"))) + host.answer(PREFILL_MODEL, chat_stream(" world")) + core = continuation_core(make_core, fallbacks=[{"a": [{"model": "b", "temperature": 0.1}]}]) + stream = await open_stream(core) + with pytest.raises(litellm.RateLimitError, match="override"): + await collect_texts(stream) + assert host.deployments_called() == ["a"] + + +async def test_continuation_runs_a_plain_fallback_for_an_error_before_content( + make_core: MakeCore, host: ScriptedHost +) -> None: + host.answer("openai/a", chat_stream(error=rate_limit(), leading_role_chunk=True)) + host.answer(PREFILL_MODEL, chat_stream("B")) + core = continuation_core(make_core) + stream = await open_stream(core) + assert await collect_texts(stream) == "B" + assert host.calls[1].kwargs["messages"] == MESSAGES + + +def recorded( + answer: Callable[..., CustomStreamWrapper], opened: list[CustomStreamWrapper] +) -> Callable[..., CustomStreamWrapper]: + """``answer``, keeping the stream it opens so a test can read that stream's logging object.""" + + def record(**request: object) -> CustomStreamWrapper: # kwargs-ok: the host passes the request kwargs + stream: Final = answer(**request) + opened.append(stream) + return stream + + return record + + +def owe_a_served_round(logging_obj: LiteLLMLogging) -> None: + """What an agentic follow-up that failed leaves on the attempt: a stash its success log must also bill.""" + logging_obj.record_served_rounds_for_failure(((Usage(prompt_tokens=1, completion_tokens=1, total_tokens=2), 0.5),)) + + +async def test_continuation_bills_the_priced_partial_on_the_fallbacks_log( + make_core: MakeCore, host: ScriptedHost, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(LiteLLMLogging, "response_cost_calculator", lambda *_args, **_kwargs: 0.25) + opened: Final[list[CustomStreamWrapper]] = [] + host.answer("openai/a", chat_stream("he", "llo", error=rate_limit())) + host.answer(PREFILL_MODEL, recorded(chat_stream(" world"), opened)) + core = continuation_core(make_core) + assert await collect_texts(await open_stream(core)) == "hello world" + details: Final = opened[0].logging_obj.model_call_details + partial_usage: Final = details["continued_partial_usage"] + assert isinstance(partial_usage, Usage) and partial_usage.completion_tokens > 0 + assert details["continued_partial_cost"] == 0.25 + assert details["continued_response_id"] + + +async def test_continuation_leaves_the_cut_streams_stash_for_the_failure_log_only( + make_core: MakeCore, host: ScriptedHost +) -> None: + opened: Final[list[CustomStreamWrapper]] = [] + host.answer("openai/a", recorded(chat_stream("he", "llo", error=rate_limit()), opened)) + host.answer(PREFILL_MODEL, chat_stream(" world")) + core = continuation_core(make_core) + stream = await open_stream(core) + cut_logging: Final = opened[0].logging_obj + owe_a_served_round(cut_logging) + assert await collect_texts(stream) == "hello world" + cut_logging.carry_failed_attempt_usage() + assert cut_logging._failed_attempts_owed_on_success == () # pyright: ignore[reportPrivateUsage] # the continuation bills it, so the success log must not + assert len(cut_logging.model_call_details["failed_attempts_usage"]) == 1 + + +async def test_a_plain_fallback_before_content_leaves_the_stash_owed_on_success( + make_core: MakeCore, host: ScriptedHost +) -> None: + opened: Final[list[CustomStreamWrapper]] = [] + host.answer("openai/a", recorded(chat_stream(error=rate_limit(), leading_role_chunk=True), opened)) + host.answer(PREFILL_MODEL, chat_stream("B")) + core = continuation_core(make_core) + stream = await open_stream(core) + cut_logging: Final = opened[0].logging_obj + owe_a_served_round(cut_logging) + assert await collect_texts(stream) == "B" + cut_logging.carry_failed_attempt_usage() + assert len(cut_logging._failed_attempts_owed_on_success) == 1 # pyright: ignore[reportPrivateUsage] # nothing else bills it + + +# --- the deployment slot held by an open stream ----------------------------------------------- + + +async def test_an_open_stream_holds_its_max_parallel_requests_slot_until_consumed( + make_core: MakeCore, host: ScriptedHost +) -> None: + host.answer("openai/a", chat_stream("x")) + core = make_core(model_list=[deployment("a", "a", max_parallel_requests=1)], num_retries=0) + first = await open_stream(core) + with pytest.raises(litellm.RateLimitError): + await asyncio.wait_for(open_stream(core), 5) + await collect_texts(first) + second = await asyncio.wait_for(open_stream(core), 5) + assert await collect_texts(second) == "x" + + +# --- the sync twin ----------------------------------------------------------------------------- + + +def test_sync_stream_falls_back_on_a_mid_stream_error_before_content(make_core: MakeCore, host: ScriptedHost) -> None: + host.answer("openai/a", chat_stream(error=rate_limit(), leading_role_chunk=True, sync=True)) + host.answer("openai/b", chat_stream("B", sync=True)) + core = build(make_core) + stream = core.completion(model="a", messages=MESSAGES, stream=True) # pyright: ignore[reportAttributeAccessIssue] # endpoint method + assert collect_texts_sync(stream) == "B" + assert [(call.op, call.deployment_id) for call in host.calls] == [("completion", "a"), ("completion", "b")] + assert model_id(stream) == "b" + + +def test_sync_stream_raises_the_provider_error_after_content(make_core: MakeCore, host: ScriptedHost) -> None: + host.answer("openai/a", chat_stream("x", error=rate_limit(message="sync cut"), sync=True)) + host.answer("openai/b", chat_stream("B", sync=True)) + core = build(make_core) + stream = core.completion(model="a", messages=MESSAGES, stream=True) # pyright: ignore[reportAttributeAccessIssue] # endpoint method + with pytest.raises(litellm.RateLimitError, match="sync cut"): + collect_texts_sync(stream) + assert host.deployments_called() == ["a"] + + +def test_sync_stream_never_continues_after_content_even_with_continuation_on( + make_core: MakeCore, host: ScriptedHost +) -> None: + host.answer("openai/a", chat_stream("he", "llo", error=rate_limit(message="sync no continuation"), sync=True)) + host.answer(PREFILL_MODEL, chat_stream(" world", sync=True)) + core = continuation_core(make_core) + stream = core.completion(model="a", messages=MESSAGES, stream=True) # pyright: ignore[reportAttributeAccessIssue] # endpoint method + with pytest.raises(litellm.RateLimitError, match="sync no continuation"): + collect_texts_sync(stream) + assert host.deployments_called() == ["a"] + + +def test_sync_stream_disable_fallbacks_raises_a_mid_stream_error_before_content( + make_core: MakeCore, host: ScriptedHost +) -> None: + host.answer("openai/a", chat_stream(error=rate_limit(message="sync disabled"), leading_role_chunk=True, sync=True)) + host.answer("openai/b", chat_stream("B", sync=True)) + core = build(make_core) + stream = core.completion(model="a", messages=MESSAGES, stream=True, disable_fallbacks=True) # pyright: ignore[reportAttributeAccessIssue] # endpoint method + with pytest.raises(litellm.RateLimitError, match="sync disabled"): + collect_texts_sync(stream) + assert host.deployments_called() == ["a"] + + +def test_sync_stream_first_chunk_failure_falls_back_inside_the_call(make_core: MakeCore, host: ScriptedHost) -> None: + host.answer("openai/a", chat_stream("x", error=rate_limit(), fail_after=0, sync=True)) + host.answer("openai/b", chat_stream("B", sync=True)) + core = build(make_core) + stream = core.completion(model="a", messages=MESSAGES, stream=True) # pyright: ignore[reportAttributeAccessIssue] # endpoint method + assert collect_texts_sync(stream) == "B" + assert host.deployments_called() == ["a", "b"] + + +def test_sync_stream_when_the_fallback_fails_raises_the_primary_error(make_core: MakeCore, host: ScriptedHost) -> None: + host.answer("openai/a", chat_stream(error=rate_limit(message="sync primary"), leading_role_chunk=True, sync=True)) + host.answer("openai/b", server_error()) + core = build(make_core) + stream = core.completion(model="a", messages=MESSAGES, stream=True) # pyright: ignore[reportAttributeAccessIssue] # endpoint method + with pytest.raises(litellm.RateLimitError, match="sync primary"): + collect_texts_sync(stream) + assert host.deployments_called() == ["a", "b"] + + +# --- the stream timeout sent to the host ------------------------------------------------------- + + +@pytest.mark.parametrize( + ("router_arguments", "deployment_params", "request_kwargs", "expected"), + [ + ({"stream_timeout": 4}, {}, {}, 4), + ({"stream_timeout": 4}, {"stream_timeout": 3}, {}, 3.0), + ({"stream_timeout": 4}, {"stream_timeout": 3}, {"stream_timeout": 2}, 2), + ({}, {"stream_timeout_by_service_tier": {"flex": 9}, "stream_timeout": 3}, {"service_tier": "flex"}, 9.0), + ({}, {"stream_timeout_by_service_tier": {"flex": 9}, "stream_timeout": 3}, {"service_tier": "auto"}, 3.0), + ( + {}, + {"stream_timeout_by_service_tier": {"flex": 9}}, + {"service_tier": "flex", "stream_timeout": 2}, + 2, + ), + ({"default_litellm_params": {"stream_timeout": 8}}, {}, {}, 8), + ], + ids=[ + "router", + "deployment-over-router", + "request-over-deployment", + "service-tier-over-deployment", + "unmatched-service-tier-uses-deployment", + "request-over-service-tier", + "default-litellm-params", + ], +) +async def test_the_stream_timeout_comes_from_request_then_tier_then_deployment_then_router( + make_core: MakeCore, + host: ScriptedHost, + router_arguments: dict[str, object], + deployment_params: dict[str, object], + request_kwargs: dict[str, object], + expected: float, +) -> None: + host.answer("openai/a", chat_stream("x")) + core = make_core(model_list=[deployment("a", "a", **deployment_params)], **router_arguments) + stream = await open_stream(core, **request_kwargs) + await collect_texts(stream) + assert host.calls[0].kwargs["stream_timeout"] == expected + + +@pytest.mark.parametrize("router_arguments", [{}, {"timeout": 42}], ids=["no-router-timeout", "router-timeout"]) +async def test_without_any_stream_timeout_none_is_sent_and_the_plain_timeout_applies( + make_core: MakeCore, host: ScriptedHost, router_arguments: dict[str, object] +) -> None: + host.answer("openai/a", chat_stream("x")) + core = make_core(model_list=[deployment("a", "a")], **router_arguments) + stream = await open_stream(core) + await collect_texts(stream) + sent = host.calls[0].kwargs + assert "stream_timeout" not in sent + assert sent["timeout"] == router_arguments.get("timeout", 6000.0) + + +async def test_a_non_streaming_request_sends_no_stream_timeout(make_core: MakeCore, host: ScriptedHost) -> None: + core = make_core(model_list=[deployment("a", "a")], stream_timeout=4) + await core.acompletion(model="a", messages=MESSAGES) # pyright: ignore[reportAttributeAccessIssue] # endpoint method + assert "stream_timeout" not in host.calls[0].kwargs diff --git a/tests/unit/test_router/core/test_streaming_characterization_outcomes.py b/tests/unit/test_router/core/test_streaming_characterization_outcomes.py new file mode 100644 index 00000000000..99b3add153b --- /dev/null +++ b/tests/unit/test_router/core/test_streaming_characterization_outcomes.py @@ -0,0 +1,382 @@ +""" +Characterization of what a streamed attempt records against its deployment: the success, tokens and +request a stream counts once it is fully consumed, and the failure and cooldown a stream that breaks +after its first chunk counts, for chat, Responses and /v1/messages streams. + +Every scripted stream here logs through the host when it ends (``logged=True``), as litellm's stream +logging does, so a core that learns outcomes from logging callbacks sees the stream's outcome only +then, with the assembled response's usage. +""" + +import asyncio +import time +from collections.abc import Callable +from typing import Final + +import pytest + +import litellm +from litellm.router_core import RouterCore +from litellm.types.llms.openai import ResponseAPIUsage, ResponseCompletedEvent, ResponsesAPIResponse +from tests.unit.test_router.core.scripted_host import ScriptedHost, deployment, rate_limit +from tests.unit.test_router.core.streaming_helpers import ( + MESSAGE_STOP, + MESSAGES, + Event, + anthropic_stream, + chat_stream, + collect_bytes, + collect_event_types, + collect_texts, + collect_texts_sync, + install_orjson_stub_if_missing, + responses_stream, + sse, + text_delta, +) + +install_orjson_stub_if_missing() + +MakeCore = Callable[..., RouterCore] +USAGE: Final = {"prompt_tokens": 4, "completion_tokens": 8, "total_tokens": 12} + + +def group( + *ids: str, **params: object +) -> list[dict[str, object]]: # kwargs-ok: litellm_params # mutable-ok: a model_list + return [deployment("g", deployment_id, None, **params) for deployment_id in ids] + + +async def open_chat(core: RouterCore, model: str = "g") -> object: + stream: Final[object] = await core.acompletion(model=model, messages=MESSAGES, stream=True) # pyright: ignore[reportAttributeAccessIssue, reportUnknownMemberType, reportUnknownVariableType] # endpoint method + return stream # pyright: ignore[reportUnknownVariableType] # see above + + +def open_chat_sync(core: RouterCore) -> object: + stream: Final[object] = core.completion(model="g", messages=MESSAGES, stream=True) # pyright: ignore[reportAttributeAccessIssue, reportUnknownMemberType, reportUnknownVariableType] # endpoint method + return stream # pyright: ignore[reportUnknownVariableType] # see above + + +async def open_responses(core: RouterCore) -> object: + stream: Final[object] = await core.aresponses(model="g", input="hi", stream=True) # pyright: ignore[reportAttributeAccessIssue, reportUnknownMemberType, reportUnknownVariableType] # endpoint method + return stream # pyright: ignore[reportUnknownVariableType] # see above + + +async def open_messages(core: RouterCore) -> object: + stream: Final[object] = await core.aanthropic_messages(model="g", messages=MESSAGES, max_tokens=5, stream=True) # pyright: ignore[reportAttributeAccessIssue, reportUnknownMemberType, reportUnknownVariableType] # endpoint method + return stream # pyright: ignore[reportUnknownVariableType] # see above + + +async def settle() -> None: + """Let the logging and recording tasks a finished stream scheduled run.""" + for _ in range(20): + await asyncio.sleep(0) + + +def eventually(check: Callable[[], bool]) -> None: + """Wait for ``check``: litellm logs a sync stream's outcome from a worker thread.""" + deadline: Final = time.monotonic() + 5.0 + while not check() and time.monotonic() < deadline: + time.sleep(0.01) + + +def successes(core: RouterCore, deployment_id: str) -> object: + count: Final[object] = core.health().store.get_cache(f"{deployment_id}:successes", local_only=True) # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] # DualCache's reads are untyped + return count # pyright: ignore[reportUnknownVariableType] # see above + + +def cooled(core: RouterCore, *ids: str) -> list[str]: # mutable-ok: a fresh list + return sorted(dict(core.health().cooldowns.get_active_cooldowns(list(ids), None))) + + +# --- a stream consumed to its end ---------------------------------------------------------------- + + +async def test_a_consumed_stream_counts_its_final_tokens_and_one_request( + make_core: MakeCore, host: ScriptedHost +) -> None: + host.answer("openai/a", chat_stream("hello", " world", usage=USAGE, logged=True)) + core: Final = make_core(model_list=group("a", rpm=100)) + + stream: Final = await open_chat(core) + assert await collect_texts(stream) == "hello world" + await settle() + + assert await core.health().usage("g") == (12, 1) + assert successes(core, "a") == 1 + + +async def test_a_stream_counts_only_its_request_before_it_is_consumed(make_core: MakeCore, host: ScriptedHost) -> None: + host.answer("openai/a", chat_stream("hello", usage=USAGE, logged=True)) + core: Final = make_core(model_list=group("a", rpm=100)) + + await open_chat(core) + await settle() + + # The request counts when the stream opens, with the tokens it has read so far (none). + assert await core.health().usage("g") == (0, 1) + assert successes(core, "a") is None + + +def test_a_sync_stream_counts_its_success_only_once_consumed(make_core: MakeCore, host: ScriptedHost) -> None: + host.answer("openai/a", chat_stream("hello", usage=USAGE, logged=True, sync=True)) + core: Final = make_core(model_list=group("a", rpm=100)) + + stream: Final = open_chat_sync(core) + assert successes(core, "a") is None + assert collect_texts_sync(stream) == "hello" + eventually(lambda: successes(core, "a") is not None) + + assert successes(core, "a") == 1 + # A sync call's success callback counts the success alone: the request counted when the stream + # opened, and no tokens. + assert asyncio.run(core.health().usage("g")) == (0, 1) + + +def test_a_sync_stream_that_breaks_after_content_counts_one_failure_and_cools_down( + make_core: MakeCore, host: ScriptedHost +) -> None: + host.answer("openai/a", chat_stream("partial", error=rate_limit(message="mid stream"), logged=True, sync=True)) + core: Final = make_core( + model_list=[deployment("g", "a", rpm=100, order=1), deployment("g", "b", order=2)], + num_retries=0, + allowed_fails=0, + ) + + stream: Final = open_chat_sync(core) + with pytest.raises(litellm.RateLimitError, match="mid stream"): + collect_texts_sync(stream) + eventually(lambda: bool(cooled(core, "a", "b"))) + + assert core.health().failures_this_minute("a") == 1 + assert successes(core, "a") is None + assert cooled(core, "a", "b") == ["a"] + assert asyncio.run(core.health().usage("g")) == (0, 1) + + +# --- a stream that breaks after its first chunk -------------------------------------------------- + + +async def test_a_stream_that_breaks_after_content_counts_one_failure_and_cools_down( + make_core: MakeCore, host: ScriptedHost +) -> None: + host.answer("openai/a", chat_stream("partial", error=rate_limit(message="mid stream"), logged=True)) + core: Final = make_core( + model_list=[deployment("g", "a", rpm=100, order=1), deployment("g", "b", order=2)], + num_retries=0, + allowed_fails=0, + ) + + stream: Final = await open_chat(core) + with pytest.raises(litellm.RateLimitError, match="mid stream"): + await collect_texts(stream) + await settle() + + assert host.deployments_called() == ["a"] + assert core.health().failures_this_minute("a") == 1 + assert successes(core, "a") is None + assert cooled(core, "a", "b") == ["a"] + # One request when the stream opened, one more for the failure. + assert await core.health().usage("g") == (0, 2) + + +async def test_a_stream_that_breaks_after_content_within_the_allowed_fails_does_not_cool_down( + make_core: MakeCore, host: ScriptedHost +) -> None: + host.answer("openai/a", chat_stream("partial", error=rate_limit(message="mid stream"), logged=True)) + core: Final = make_core(model_list=group("a"), num_retries=0, allowed_fails=3) + + stream: Final = await open_chat(core) + with pytest.raises(litellm.RateLimitError, match="mid stream"): + await collect_texts(stream) + await settle() + + assert core.health().failures_this_minute("a") == 1 + assert successes(core, "a") is None + assert cooled(core, "a") == [] + + +async def test_a_stream_that_breaks_before_content_and_falls_back_counts_one_failure_and_one_success( + make_core: MakeCore, host: ScriptedHost +) -> None: + host.answer( + "openai/a", + chat_stream("x", error=rate_limit(message="cut"), fail_after=1, leading_role_chunk=True, logged=True), + ) + host.answer("openai/b", chat_stream("B", usage=USAGE, logged=True)) + core: Final = make_core( + model_list=[deployment("a", "a", rpm=100), deployment("b", "b", rpm=100)], + fallbacks=[{"a": ["b"]}], + num_retries=0, + allowed_fails=3, + ) + + stream: Final = await open_chat(core, model="a") + assert await collect_texts(stream) == "B" + await settle() + + assert host.deployments_called() == ["a", "b"] + assert core.health().failures_this_minute("a") == 1 + assert successes(core, "a") is None + assert await core.health().usage("a") == (0, 2) + assert successes(core, "b") == 1 + assert await core.health().usage("b") == (12, 1) + + +async def test_a_stream_that_fails_on_its_first_chunk_counts_its_failure_once( + make_core: MakeCore, host: ScriptedHost +) -> None: + host.answer( + "openai/a", + chat_stream("x", error=rate_limit(), fail_after=0, logged=True), + chat_stream("ok", usage=USAGE, logged=True), + ) + core: Final = make_core(model_list=group("a", rpm=100), num_retries=1, allowed_fails=3) + + stream: Final = await open_chat(core) + assert await collect_texts(stream) == "ok" + await settle() + + assert host.deployments_called() == ["a", "a"] + assert core.health().failures_this_minute("a") == 1 + assert successes(core, "a") == 1 + assert await core.health().usage("g") == (12, 2) + + +# --- Responses streams --------------------------------------------------------------------------- + + +def completed_event(input_tokens: int, output_tokens: int) -> ResponseCompletedEvent: + response: Final = ResponsesAPIResponse.model_construct( + id="resp", + background=None, + usage=ResponseAPIUsage( + input_tokens=input_tokens, output_tokens=output_tokens, total_tokens=input_tokens + output_tokens + ), + ) + return ResponseCompletedEvent(type="response.completed", response=response) # pyright: ignore[reportArgumentType] # the enum accepts its value + + +async def test_a_consumed_responses_stream_counts_its_final_tokens_and_one_request( + make_core: MakeCore, host: ScriptedHost +) -> None: + host.answer( + "openai/a", + responses_stream( + Event("response.created"), + Event("response.output_text.delta"), + completed_event(4, 8), + logged=True, + total_tokens=12, + ), + ) + core: Final = make_core(model_list=group("a", rpm=100)) + + stream: Final = await open_responses(core) + assert (await collect_event_types(stream))[-1] == "response.completed" + await settle() + + assert await core.health().usage("g") == (12, 1) + assert successes(core, "a") == 1 + + +async def test_a_responses_stream_that_breaks_after_output_counts_one_failure_and_cools_down( + make_core: MakeCore, host: ScriptedHost +) -> None: + host.answer( + "openai/a", + responses_stream( + Event("response.created"), + Event("response.output_text.delta", delta="partial"), + error=rate_limit(message="mid stream"), + logged=True, + ), + ) + core: Final = make_core( + model_list=[deployment("g", "a", rpm=100, order=1), deployment("g", "b", order=2)], + num_retries=0, + allowed_fails=0, + ) + + stream: Final = await open_responses(core) + with pytest.raises(litellm.RateLimitError, match="mid stream"): + await collect_event_types(stream) + await settle() + + assert core.health().failures_this_minute("a") == 1 + assert successes(core, "a") is None + assert cooled(core, "a", "b") == ["a"] + assert await core.health().usage("g") == (0, 2) + + +# --- /v1/messages streams ------------------------------------------------------------------------ + +MESSAGE_START_WITH_USAGE: Final = sse( + "message_start", + { + "type": "message_start", + "message": { + "id": "msg_1", + "type": "message", + "role": "assistant", + "content": [], + "usage": {"input_tokens": 3, "cache_read_input_tokens": 1, "output_tokens": 1}, + }, + }, +) +MESSAGE_DELTA_WITH_USAGE: Final = sse( + "message_delta", + {"type": "message_delta", "delta": {"stop_reason": "end_turn"}, "usage": {"output_tokens": 8}}, +) + + +async def test_a_consumed_messages_stream_counts_its_final_tokens_and_one_request( + make_core: MakeCore, host: ScriptedHost +) -> None: + host.answer( + "openai/a", + anthropic_stream( + MESSAGE_START_WITH_USAGE, + text_delta("hello"), + # the delta frame split across two byte chunks + MESSAGE_DELTA_WITH_USAGE[:30], + MESSAGE_DELTA_WITH_USAGE[30:], + MESSAGE_STOP, + logged=True, + total_tokens=12, + ), + ) + core: Final = make_core(model_list=group("a", rpm=100)) + + stream: Final = await open_messages(core) + assert len(await collect_bytes(stream)) == 5 + await settle() + + assert await core.health().usage("g") == (12, 1) + assert successes(core, "a") == 1 + + +async def test_a_messages_stream_that_breaks_after_content_counts_one_failure_and_cools_down( + make_core: MakeCore, host: ScriptedHost +) -> None: + host.answer( + "openai/a", + anthropic_stream( + MESSAGE_START_WITH_USAGE, text_delta("x"), error=rate_limit(message="mid stream"), logged=True + ), + ) + core: Final = make_core( + model_list=[deployment("g", "a", rpm=100, order=1), deployment("g", "b", order=2)], + num_retries=0, + allowed_fails=0, + ) + + stream: Final = await open_messages(core) + with pytest.raises(litellm.RateLimitError, match="mid stream"): + await collect_bytes(stream) + await settle() + + assert core.health().failures_this_minute("a") == 1 + assert successes(core, "a") is None + assert cooled(core, "a", "b") == ["a"] + assert await core.health().usage("g") == (0, 2) diff --git a/tests/unit/test_router/core/test_streaming_characterization_responses.py b/tests/unit/test_router/core/test_streaming_characterization_responses.py new file mode 100644 index 00000000000..4b601705324 --- /dev/null +++ b/tests/unit/test_router/core/test_streaming_characterization_responses.py @@ -0,0 +1,458 @@ +""" +Characterization of Responses API streaming through a router core: pre-output lifecycle events held +back while a fallback can still take over, mid-stream errors and their fallback, the continuation a +fallback receives after text was streamed, usage combined across a stream fallback, and the hidden +params the returned stream and its events carry. + +Scripted streams are ``BaseResponsesAPIStreamingIterator`` subclasses +(``streaming_helpers.responses_stream``) that raise a ``MidStreamFallbackError`` the way the +completion-bridge iterator does. +""" + +from collections.abc import Callable, Mapping +from typing import Final + +import pytest + +import litellm +from litellm.exceptions import MidStreamFallbackError +from litellm.responses.streaming_iterator import BaseResponsesAPIStreamingIterator +from litellm.router_core import RouterCore +from litellm.types.llms.openai import ( + OutputTextDeltaEvent, + ResponseAPIUsage, + ResponseCompletedEvent, + ResponsesAPIResponse, + ResponsesAPIStreamEvents, +) +from tests.unit.test_router.core.scripted_host import ( + ScriptedHost, + content_policy, + context_window, + deployment, + rate_limit, + server_error, +) +from tests.unit.test_router.core.streaming_helpers import ( + Event, + ScriptedResponsesStream, + collect_event_types, + drain_into, + responses_stream, +) + +MakeCore = Callable[..., RouterCore] +GROUPS_AB: Final = [deployment("a", "a"), deployment("b", "b")] +CREATED: Final = "response.created" +IN_PROGRESS: Final = "response.in_progress" +TEXT_DELTA: Final = "response.output_text.delta" +COMPLETED: Final = "response.completed" + + +def cut(generated_content: str = "", original: Exception | None = None) -> MidStreamFallbackError: + """The error the bridge iterator raises when its stream breaks (before output when nothing was generated).""" + return MidStreamFallbackError( + message="stream cut", + model="d1", + llm_provider="anthropic", + original_exception=original if original is not None else rate_limit(message="stream cut"), + generated_content=generated_content, + is_pre_first_chunk=not generated_content, + ) + + +def fallback_ok( + *events: object, opened: list[ScriptedResponsesStream] | None = None +) -> Callable[..., ScriptedResponsesStream]: + return responses_stream( + *(events or (Event(CREATED), Event(TEXT_DELTA), Event(COMPLETED))), + hidden_params={"model_id": "b-stream", "additional_headers": {"x-provider-request-id": "req-b"}}, + opened=opened, + ) + + +def build(make_core: MakeCore, **arguments: object) -> RouterCore: # kwargs-ok: the router's constructor arguments + return make_core(**{"model_list": GROUPS_AB, "fallbacks": [{"a": ["b"]}], "num_retries": 0, **arguments}) + + +async def open_stream(core: RouterCore, **request: object) -> object: # kwargs-ok: request kwargs + return await core.aresponses(model="a", input="hi", stream=True, **request) # pyright: ignore[reportAttributeAccessIssue] # endpoint method + + +def hidden(stream: object) -> Mapping[str, object]: + return getattr(stream, "_hidden_params") + + +def completed_event(input_tokens: int, output_tokens: int) -> ResponseCompletedEvent: + response = ResponsesAPIResponse.model_construct( + id="resp", + background=None, + usage=ResponseAPIUsage( + input_tokens=input_tokens, output_tokens=output_tokens, total_tokens=input_tokens + output_tokens + ), + ) + return ResponseCompletedEvent(type=COMPLETED, response=response) # pyright: ignore[reportArgumentType] # the enum accepts its value + + +# --- the returned stream ------------------------------------------------------------------------ + + +async def test_a_streaming_response_is_returned_as_a_responses_streaming_iterator( + make_core: MakeCore, host: ScriptedHost +) -> None: + host.answer("openai/a", responses_stream(Event(CREATED), Event(TEXT_DELTA), Event(COMPLETED))) + core = build(make_core) + stream = await open_stream(core) + assert isinstance(stream, BaseResponsesAPIStreamingIterator) + assert await collect_event_types(stream) == [CREATED, TEXT_DELTA, COMPLETED] + assert hidden(stream)["model_id"] == "a" + + +async def test_the_returned_stream_latches_the_terminal_event_it_forwarded( + make_core: MakeCore, host: ScriptedHost +) -> None: + completed = Event(COMPLETED) + host.answer("openai/a", responses_stream(Event(CREATED), Event(TEXT_DELTA), completed)) + core = build(make_core) + stream = await open_stream(core) + await collect_event_types(stream) + assert getattr(stream, "completed_response") is completed + + +async def test_lifecycle_events_are_held_until_output_and_then_forwarded_in_order( + make_core: MakeCore, host: ScriptedHost +) -> None: + host.answer("openai/a", responses_stream(Event(CREATED), Event(IN_PROGRESS), Event(TEXT_DELTA), Event(COMPLETED))) + core = build(make_core) + assert await collect_event_types(await open_stream(core)) == [CREATED, IN_PROGRESS, TEXT_DELTA, COMPLETED] + + +async def test_held_lifecycle_events_are_flushed_when_the_stream_ends_without_output( + make_core: MakeCore, host: ScriptedHost +) -> None: + host.answer("openai/a", responses_stream(Event(CREATED), Event(IN_PROGRESS))) + core = build(make_core) + assert await collect_event_types(await open_stream(core)) == [CREATED, IN_PROGRESS] + + +async def test_the_source_stream_is_closed_when_the_stream_ends(make_core: MakeCore, host: ScriptedHost) -> None: + opened: list[ScriptedResponsesStream] = [] + host.answer("openai/a", responses_stream(Event(CREATED), Event(COMPLETED), opened=opened)) + core = build(make_core) + await collect_event_types(await open_stream(core)) + assert [stream.closed for stream in opened] == [True] + + +# --- mid-stream errors before output ----------------------------------------------------------- + + +async def test_a_mid_stream_error_before_output_falls_back_and_drops_the_primarys_held_events( + make_core: MakeCore, host: ScriptedHost +) -> None: + host.answer("openai/a", responses_stream(Event(CREATED), Event(IN_PROGRESS), error=cut())) + host.answer("openai/b", fallback_ok()) + core = build(make_core) + stream = await open_stream(core) + assert await collect_event_types(stream) == [CREATED, TEXT_DELTA, COMPLETED] + assert host.deployments_called() == ["a", "b"] + assert host.calls[1].kwargs["input"] == "hi" + assert host.calls[1].model_group == "b" + + +async def test_a_mid_stream_error_before_output_without_fallbacks_flushes_held_events_then_raises( + make_core: MakeCore, host: ScriptedHost +) -> None: + host.answer("openai/a", responses_stream(Event(CREATED), Event(IN_PROGRESS), error=cut())) + core = make_core(model_list=GROUPS_AB, num_retries=0) + stream = await open_stream(core) + seen: list[object] = [] + with pytest.raises(litellm.RateLimitError, match="stream cut") as raised: + await drain_into(stream, seen) + assert not isinstance(raised.value, MidStreamFallbackError) + assert [getattr(event, "type") for event in seen] == [CREATED, IN_PROGRESS] + assert host.deployments_called() == ["a"] + + +async def test_disable_fallbacks_raises_a_mid_stream_error_before_output( + make_core: MakeCore, host: ScriptedHost +) -> None: + host.answer("openai/a", responses_stream(Event(CREATED), error=cut())) + host.answer("openai/b", fallback_ok()) + core = build(make_core) + stream = await open_stream(core, disable_fallbacks=True) + with pytest.raises(litellm.RateLimitError, match="stream cut"): + await collect_event_types(stream) + assert host.deployments_called() == ["a"] + + +async def test_when_the_fallback_fails_the_held_events_are_replayed_before_the_primary_error( + make_core: MakeCore, host: ScriptedHost +) -> None: + host.answer("openai/a", responses_stream(Event(CREATED), error=cut())) + host.answer("openai/b", server_error()) + core = build(make_core) + stream = await open_stream(core) + seen: list[object] = [] + with pytest.raises(litellm.RateLimitError, match="stream cut"): + await drain_into(stream, seen) + assert [getattr(event, "type") for event in seen] == [CREATED] + assert host.deployments_called() == ["a", "b"] + + +async def test_more_than_two_hundred_lifecycle_events_are_forwarded_before_a_fallback( + make_core: MakeCore, host: ScriptedHost +) -> None: + host.answer("openai/a", responses_stream(*(Event(IN_PROGRESS) for _ in range(201)), error=cut())) + host.answer("openai/b", fallback_ok()) + core = build(make_core) + types = await collect_event_types(await open_stream(core)) + assert types == [*([IN_PROGRESS] * 201), CREATED, TEXT_DELTA, COMPLETED] + + +async def test_a_context_window_error_before_output_uses_the_context_window_chain( + make_core: MakeCore, host: ScriptedHost +) -> None: + host.answer("openai/a", responses_stream(Event(CREATED), error=context_window())) + host.answer("openai/b", fallback_ok()) + core = build(make_core, fallbacks=None, context_window_fallbacks=[{"a": ["b"]}]) + assert await collect_event_types(await open_stream(core)) == [CREATED, TEXT_DELTA, COMPLETED] + assert host.deployments_called() == ["a", "b"] + + +async def test_a_context_window_error_after_output_is_raised(make_core: MakeCore, host: ScriptedHost) -> None: + host.answer("openai/a", responses_stream(Event(CREATED), Event(TEXT_DELTA), error=context_window())) + host.answer("openai/b", fallback_ok()) + core = build(make_core, context_window_fallbacks=[{"a": ["b"]}]) + stream = await open_stream(core) + with pytest.raises(litellm.ContextWindowExceededError): + await collect_event_types(stream) + assert host.deployments_called() == ["a"] + + +async def test_a_content_policy_cut_before_output_uses_the_content_policy_chain( + make_core: MakeCore, host: ScriptedHost +) -> None: + host.answer("openai/a", responses_stream(Event(CREATED), error=cut(original=content_policy()))) + host.answer("openai/b", fallback_ok()) + core = build(make_core, fallbacks=None, content_policy_fallbacks=[{"a": ["b"]}]) + assert await collect_event_types(await open_stream(core)) == [CREATED, TEXT_DELTA, COMPLETED] + assert host.deployments_called() == ["a", "b"] + + +# --- continuation after output ----------------------------------------------------------------- + + +async def test_after_text_the_fallback_receives_the_partial_output_and_a_request_to_continue( + make_core: MakeCore, host: ScriptedHost +) -> None: + host.answer("openai/a", responses_stream(Event(CREATED), Event(TEXT_DELTA), error=cut("hel"))) + host.answer("openai/b", fallback_ok()) + core = build(make_core) + await collect_event_types(await open_stream(core)) + sent = host.calls[1].kwargs["input"] + assert isinstance(sent, list) + assert sent[0] == {"type": "message", "role": "user", "content": [{"type": "input_text", "text": "hi"}]} + assert sent[1] == {"type": "message", "role": "assistant", "content": [{"type": "output_text", "text": "hel"}]} + assert sent[2]["role"] == "user" + assert 'It ended with "hel"' in sent[2]["content"][0]["text"] + + +async def test_after_text_the_output_ceiling_is_reduced_by_the_emitted_tokens( + make_core: MakeCore, host: ScriptedHost +) -> None: + host.answer("openai/a", responses_stream(Event(CREATED), Event(TEXT_DELTA), error=cut("hel"))) + host.answer("openai/b", fallback_ok()) + core = build(make_core) + await collect_event_types(await open_stream(core, max_output_tokens=100)) + assert [call.kwargs["max_output_tokens"] for call in host.calls] == [100, 99] + + +async def test_after_text_a_spent_output_ceiling_raises_the_primary_error( + make_core: MakeCore, host: ScriptedHost +) -> None: + host.answer("openai/a", responses_stream(Event(CREATED), Event(TEXT_DELTA), error=cut("hello there friend"))) + host.answer("openai/b", fallback_ok()) + core = build(make_core) + stream = await open_stream(core, max_output_tokens=1) + with pytest.raises(litellm.RateLimitError, match="stream cut"): + await collect_event_types(stream) + assert host.deployments_called() == ["a"] + + +async def test_after_text_a_content_policy_cut_raises_instead_of_continuing( + make_core: MakeCore, host: ScriptedHost +) -> None: + host.answer("openai/a", responses_stream(Event(CREATED), Event(TEXT_DELTA), error=cut("hel", content_policy()))) + host.answer("openai/b", fallback_ok()) + core = build(make_core, content_policy_fallbacks=[{"a": ["b"]}]) + stream = await open_stream(core) + with pytest.raises(litellm.ContentPolicyViolationError): + await collect_event_types(stream) + assert host.deployments_called() == ["a"] + + +async def test_after_a_tool_call_began_the_primary_error_is_raised(make_core: MakeCore, host: ScriptedHost) -> None: + host.answer( + "openai/a", + responses_stream(Event(CREATED), Event(TEXT_DELTA), error=cut("hel"), started_tool_call=True), + ) + host.answer("openai/b", fallback_ok()) + core = build(make_core) + stream = await open_stream(core) + with pytest.raises(litellm.RateLimitError, match="stream cut"): + await collect_event_types(stream) + assert host.deployments_called() == ["a"] + + +async def test_after_text_a_json_schema_output_format_raises_the_primary_error( + make_core: MakeCore, host: ScriptedHost +) -> None: + host.answer("openai/a", responses_stream(Event(CREATED), Event(TEXT_DELTA), error=cut("hel"))) + host.answer("openai/b", fallback_ok()) + core = build(make_core) + stream = await open_stream(core, text={"format": {"type": "json_schema", "name": "x", "schema": {}}}) + with pytest.raises(litellm.RateLimitError, match="stream cut"): + await collect_event_types(stream) + assert host.deployments_called() == ["a"] + + +async def test_after_text_the_primarys_events_stay_and_the_fallbacks_follow( + make_core: MakeCore, host: ScriptedHost +) -> None: + host.answer("openai/a", responses_stream(Event(CREATED), Event(TEXT_DELTA), error=cut("hel"))) + host.answer("openai/b", fallback_ok()) + core = build(make_core) + types = await collect_event_types(await open_stream(core)) + assert types[:2] == [CREATED, TEXT_DELTA] + assert types[-1] == COMPLETED + assert host.deployments_called() == ["a", "b"] + + +# --- usage across a stream fallback -------------------------------------------------------------- + + +async def test_the_fallbacks_completed_usage_adds_the_usage_the_primary_latched( + make_core: MakeCore, host: ScriptedHost +) -> None: + host.answer("openai/a", responses_stream(Event(CREATED), error=cut(), completed_response=completed_event(3, 4))) + host.answer("openai/b", fallback_ok(Event(CREATED), completed_event(10, 20))) + core = build(make_core) + stream = await open_stream(core) + events = [event async for event in stream] # pyright: ignore[reportGeneralTypeIssues] # an async stream + usage = events[-1].response.usage + assert (usage.input_tokens, usage.output_tokens, usage.total_tokens) == (13, 24, 37) + + +async def test_without_primary_usage_the_fallbacks_completed_usage_is_unchanged( + make_core: MakeCore, host: ScriptedHost +) -> None: + host.answer("openai/a", responses_stream(Event(CREATED), error=cut())) + host.answer("openai/b", fallback_ok(Event(CREATED), completed_event(10, 20))) + core = build(make_core) + stream = await open_stream(core) + events = [event async for event in stream] # pyright: ignore[reportGeneralTypeIssues] # an async stream + usage = events[-1].response.usage + assert (usage.input_tokens, usage.output_tokens, usage.total_tokens) == (10, 20, 30) + + +# --- sequence numbers across a stream fallback --------------------------------------------------- + + +def delta(text: str, sequence_number: int | None = None) -> OutputTextDeltaEvent: + event: Final = OutputTextDeltaEvent( + type=ResponsesAPIStreamEvents.OUTPUT_TEXT_DELTA, item_id="msg_1", output_index=0, content_index=0, delta=text + ) + return event if sequence_number is None else event.model_copy(update={"sequence_number": sequence_number}) + + +async def numbered(stream: object) -> list[tuple[str, object]]: # mutable-ok: a fresh list + return [(event.delta, event.get("sequence_number")) async for event in stream] # pyright: ignore[reportGeneralTypeIssues, reportUnknownMemberType, reportUnknownVariableType] # an async stream of deltas + + +async def test_sequence_numbers_keep_increasing_across_a_stream_fallback( + make_core: MakeCore, host: ScriptedHost +) -> None: + host.answer("openai/a", responses_stream(delta("a", 0), delta("b", 1), error=cut("ab"))) + host.answer("openai/b", fallback_ok(delta("c", 0), delta("d", 1))) + core = build(make_core) + assert await numbered(await open_stream(core)) == [("a", 0), ("b", 1), ("c", 2), ("d", 3)] + + +async def test_a_stream_fallback_continues_above_the_highest_number_the_primary_sent( + make_core: MakeCore, host: ScriptedHost +) -> None: + host.answer("openai/a", responses_stream(delta("a", 0), delta("b", 3), delta("c", 1), error=cut("abc"))) + host.answer("openai/b", fallback_ok(delta("d", 0))) + core = build(make_core) + assert await numbered(await open_stream(core)) == [("a", 0), ("b", 3), ("c", 1), ("d", 4)] + + +async def test_a_stream_fallback_keeps_higher_numbers_and_passes_unnumbered_events( + make_core: MakeCore, host: ScriptedHost +) -> None: + host.answer("openai/a", responses_stream(delta("a", 0), delta("b", 1), error=cut("ab"))) + host.answer("openai/b", fallback_ok(delta("c", 7), delta("d"), delta("e", 8))) + core = build(make_core) + assert await numbered(await open_stream(core)) == [("a", 0), ("b", 1), ("c", 7), ("d", None), ("e", 8)] + + +async def test_sequence_numbers_pass_through_untouched_without_a_fallback( + make_core: MakeCore, host: ScriptedHost +) -> None: + host.answer("openai/a", responses_stream(delta("a", 0), delta("b", 1), delta("c", 1), delta("d", 0))) + core = build(make_core) + assert await numbered(await open_stream(core)) == [("a", 0), ("b", 1), ("c", 1), ("d", 0)] + + +# --- hidden params after a stream fallback ------------------------------------------------------- + + +async def test_after_a_stream_fallback_the_stream_carries_the_fallbacks_hidden_params_and_headers( + make_core: MakeCore, host: ScriptedHost +) -> None: + host.answer("openai/a", responses_stream(Event(CREATED), error=cut())) + host.answer("openai/b", fallback_ok()) + core = build(make_core) + stream = await open_stream(core) + await collect_event_types(stream) + assert hidden(stream)["model_id"] == "b-stream" + assert hidden(stream)["additional_headers"] == { + "x-provider-request-id": "req-b", + "x-litellm-model-group": "b", + "x-litellm-attempted-retries": 0, + "x-litellm-attempted-fallbacks": 1, + } + + +async def test_after_a_stream_fallback_each_fallback_event_carries_the_fallbacks_hidden_params( + make_core: MakeCore, host: ScriptedHost +) -> None: + host.answer("openai/a", responses_stream(Event(CREATED), error=cut())) + host.answer("openai/b", fallback_ok()) + core = build(make_core) + stream = await open_stream(core) + events = [event async for event in stream] # pyright: ignore[reportGeneralTypeIssues] # an async stream + assert all(event._hidden_params["model_id"] == "b-stream" for event in events) + assert all(event._hidden_params["additional_headers"]["x-litellm-model-group"] == "b" for event in events) + + +async def test_when_the_stream_fallback_fails_the_stream_still_reports_the_primary( + make_core: MakeCore, host: ScriptedHost +) -> None: + host.answer("openai/a", responses_stream(Event(CREATED), error=cut())) + host.answer("openai/b", server_error()) + core = build(make_core) + stream = await open_stream(core) + with pytest.raises(litellm.RateLimitError): + await collect_event_types(stream) + assert hidden(stream)["model_id"] == "a" + assert hidden(stream)["additional_headers"]["x-litellm-model-group"] == "a" + + +async def test_the_primary_and_fallback_streams_are_closed(make_core: MakeCore, host: ScriptedHost) -> None: + opened: list[ScriptedResponsesStream] = [] + host.answer("openai/a", responses_stream(Event(CREATED), error=cut(), opened=opened)) + host.answer("openai/b", fallback_ok(opened=opened)) + core = build(make_core) + await collect_event_types(await open_stream(core)) + assert [stream.closed for stream in opened] == [True, True] diff --git a/tests/unit/test_router/core/test_usage_and_failure_counting_parity.py b/tests/unit/test_router/core/test_usage_and_failure_counting_parity.py new file mode 100644 index 00000000000..674f5496b84 --- /dev/null +++ b/tests/unit/test_router/core/test_usage_and_failure_counting_parity.py @@ -0,0 +1,455 @@ +""" +Parity of the Rust core with ``router_python_core.py`` on what a call counts: which deployments track +per-minute usage, when a failure's counters are written on a logging object the attempts share, the +exits of the failed fallback hop's cooldown, the reads and writes that reach Redis, and the keys a +wildcard stream counts under. +""" + +import asyncio +from collections.abc import Callable, Mapping +from datetime import datetime, timezone +from typing import Final + +import pytest + +import litellm +from litellm.exceptions import MidStreamFallbackError +from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging +from litellm.router_core import RouterCore +from litellm.router_utils import cooldown_handlers +from litellm.router_utils.cooldown_handlers import mark_deployment_selection_failure +from tests.unit.test_router.core.endpoints_helpers import EndpointHost +from tests.unit.test_router.core.retries_health_helpers import ( + MESSAGES, + ErrorBuilder, + MakeCore, + attempt, + attempt_sync, + cooled_down, + fresh, + not_found, + plain_error, + rate_limited, +) +from tests.unit.test_router.core.scripted_host import Call, deployment, server_error, timeout +from tests.unit.test_router.core.streaming_helpers import chat_stream, collect_texts, request_logging + +USAGE: Final = {"prompt_tokens": 4, "completion_tokens": 8, "total_tokens": 12} +PER_USER_COPILOT: Final = { + "litellm_credential_name": "ghc", + "github_copilot_auth_type": "per_user_oauth", +} + + +async def settle() -> None: + """Let the recording tasks a stream's end scheduled run.""" + for _ in range(20): + await asyncio.sleep(0) + + +class SharedLoggingHost(EndpointHost): + """``EndpointHost`` that logs on a shared logging object as litellm does: every attempt stamps its + marker, and only the first failure runs the failure callbacks. ``_arealtime`` answers ``None``.""" + + __slots__ = ("timeout_cause",) + + def __init__(self) -> None: + super().__init__() + self.timeout_cause: BaseException | None = None + + async def _answer_async(self, op: str, args: tuple[object, ...], kwargs: Mapping[str, object]) -> object: + logging_obj: Final = kwargs.get("litellm_logging_obj") + if isinstance(logging_obj, LiteLLMLogging): + logging_obj.model_call_details["current_attempt_api_call_start_time"] = object() + return await super()._answer_async(op, args, kwargs) + + def _finish(self, op: str, call: Call, outcome: object) -> tuple[BaseException | None, object]: + error, response = super()._finish(op, call, outcome) + if isinstance(error, litellm.Timeout) and self.timeout_cause is not None: + error.__cause__ = self.timeout_cause # rebind-ok: a scripted error is rebuilt, which drops its cause + return error, None if op == "_arealtime" else response + + async def log_async(self, call: Call, error: BaseException | None, response: object) -> None: + logging_obj: Final = call.kwargs.get("litellm_logging_obj") + if error is None or not isinstance(logging_obj, LiteLLMLogging): + await super().log_async(call, error, response) + return + details: Final = logging_obj.model_call_details + if details.get("has_logged_async_failure") is True: + return + details["has_logged_async_failure"] = True + details["async_failure_logged_attempt"] = details.get("current_attempt_api_call_start_time") + await super().log_async(call, error, response) + + +class RedisDouble: + """The shared tier the routing caches talk to, holding the totals every pod counted.""" + + def __init__(self) -> None: + self.totals: Final[dict[str, float]] = {} # mutable-ok: the double's own store + self.gets: Final[list[str]] = [] # mutable-ok: what was read + self.batches: Final[list[list[str]]] = [] # mutable-ok: what was read + + async def async_get_cache(self, key: str, **_: object) -> object: + self.gets.append(key) + return self.totals.get(key) + + def get_cache(self, key: str, **_: object) -> object: + self.gets.append(key) + return self.totals.get(key) + + async def async_batch_get_cache( + self, key_list: list[str], **_: object + ) -> dict[str, object]: # mutable-ok: as RedisCache + self.batches.append(list(key_list)) + return {key: self.totals.get(key) for key in key_list} + + def batch_get_cache(self, key_list: list[str], **_: object) -> dict[str, object]: # mutable-ok: as RedisCache + self.batches.append(list(key_list)) + return {key: self.totals.get(key) for key in key_list} + + async def async_increment(self, key: str, value: float, **_: object) -> float: + self.totals[key] = self.totals.get(key, 0) + value + return self.totals[key] + + async def async_increment_pipeline( + self, increment_list: list[Mapping[str, object]], **_: object + ) -> list[float]: # mutable-ok: as RedisCache + return [ + await self.async_increment(str(row["key"]), float(str(row["increment_value"]))) for row in increment_list + ] + + async def async_set_cache(self, key: str, value: object, **_: object) -> None: + return None + + def set_cache(self, key: str, value: object, **_: object) -> None: + return None + + +@pytest.fixture +def host() -> SharedLoggingHost: + return SharedLoggingHost() + + +def minute() -> str: + return datetime.now(timezone.utc).strftime("%H-%M") + + +def chain(make_core: MakeCore, **arguments: object) -> RouterCore: # kwargs-ok: router arguments + return make_core( + model_list=[deployment("a", "a"), deployment("b", "b"), deployment("c", "c")], + fallbacks=[{"a": ["b"]}, {"b": ["c"]}], + num_retries=0, + **arguments, + ) + + +def with_redis(core: RouterCore) -> RedisDouble: + redis: Final = RedisDouble() + core.health().store.redis_cache = redis # pyright: ignore[reportAttributeAccessIssue] # a double + return redis + + +# --- PV4: which deployments track usage ---------------------------------------------------------- + + +@pytest.mark.parametrize( + ("entry", "expected"), + [ + (deployment("g", "d1"), (None, None)), + (deployment("g", "d1", rpm=100), (10, 1)), + (deployment("g", "d1", itpm=100), (10, 1)), + (deployment("g", "d1", model_info={"rpm": 100}), (10, 1)), + (deployment("g", "d1", model_info={"tpm": 100}), (10, 1)), + (deployment("g", "d1", model_info={"otpm": 100}), (10, 1)), + ({**deployment("g", "d1"), "tpm": 100}, (10, 1)), + (deployment("g", "d1", model="gemini/gemini-3.1-flash-lite-preview"), (10, 1)), + ], + ids=[ + "no-limit", + "params-rpm", + "params-itpm", + "model-info-rpm", + "model-info-tpm", + "model-info-otpm", + "top-level-tpm", + "catalog-rpm", + ], +) +async def test_usage_is_counted_for_a_limit_anywhere_the_python_core_reads_one( + make_core: MakeCore, entry: dict[str, object], expected: tuple[int | None, int | None] +) -> None: + core: Final = make_core(model_list=[entry]) + + # The scripted host logs a sync call's success apart from the usage `make_call` writes. + _ = attempt_sync(core, "g") + + assert await core.health().usage("g") == expected + + +async def test_a_stream_counts_usage_for_a_limit_in_model_info(make_core: MakeCore, host: SharedLoggingHost) -> None: + host.answer("openai/d1", chat_stream("hello", usage=USAGE, logged=True)) + core: Final = make_core(model_list=[deployment("g", "d1", model_info={"rpm": 100})]) + + stream: Final = await core.acompletion(model="g", messages=MESSAGES, stream=True) # pyright: ignore[reportAttributeAccessIssue] # every core serves acompletion + _ = await collect_texts(stream) + await settle() + + assert await core.health().usage("g") == (12, 1) + + +# --- realtime ------------------------------------------------------------------------------------ + + +@pytest.mark.parametrize("limit", ["tpm", "rpm", "itpm", "otpm"]) +async def test_a_realtime_session_counts_no_request_against_a_limited_deployment( + make_core: MakeCore, limit: str +) -> None: + core: Final = make_core(model_list=[deployment("g", "d1", **{limit: 100})], num_retries=0) + + await core._arealtime(model="g") # pyright: ignore[reportAttributeAccessIssue] # every core serves it + + assert await core.health().usage("g") == (None, None) + + +# --- P3: what a shared logging object lets through ------------------------------------------------- + + +async def test_only_the_first_failure_on_a_shared_logging_object_counts_a_request( + make_core: MakeCore, host: SharedLoggingHost +) -> None: + core: Final = chain(make_core) + host.answer("*", fresh(server_error)) + + _ = await attempt(core, "a", litellm_logging_obj=request_logging()) + + assert await core.health().usage("a") == (None, 1) + assert await core.health().usage("b") == (None, None) + assert await core.health().usage("c") == (None, None) + + +async def test_every_failure_counts_a_request_when_each_call_logs_on_its_own( + make_core: MakeCore, host: SharedLoggingHost +) -> None: + core: Final = chain(make_core) + host.answer("*", fresh(server_error)) + + _ = await attempt(core, "a") + + assert await core.health().usage("a") == (None, 1) + assert await core.health().usage("b") == (None, 1) + assert await core.health().usage("c") == (None, 1) + + +async def test_a_blocked_retry_in_the_same_group_records_nothing(make_core: MakeCore, host: SharedLoggingHost) -> None: + core: Final = make_core(model_list=[deployment("g", "a"), deployment("g", "b")], num_retries=1, allowed_fails=5) + host.answer("*", fresh(server_error)) + + _ = await attempt(core, "g", litellm_logging_obj=request_logging()) + + assert core.health().failures_this_minute("a") + core.health().failures_this_minute("b") == 1 + + +def test_a_sync_failure_counts_no_request(make_core: MakeCore, host: SharedLoggingHost) -> None: + core: Final = make_core(model_list=[deployment("g", "d1")], num_retries=0) + host.answer("*", fresh(server_error)) + + with pytest.raises(litellm.InternalServerError): + core.completion(model="g", messages=MESSAGES) # pyright: ignore[reportAttributeAccessIssue] # see above + + assert core.health().failures_this_minute("d1") == 1 + assert core.health().store.get_cache(f"global_router:d1:openai/d1:rpm:{minute()}", local_only=True) is None + + +# --- failure counting ------------------------------------------------------------------------------ + + +async def test_a_failure_raised_before_the_litellm_call_records_nothing( + make_core: MakeCore, host: SharedLoggingHost +) -> None: + core: Final = make_core( + model_list=[deployment("g", "d1", model="github_copilot/gpt-4o", **PER_USER_COPILOT)], num_retries=0 + ) + + raised: Final = await attempt(core, "g", litellm_credential_name="another") + + assert isinstance(raised, litellm.BadRequestError) + assert host.calls == [] + assert core.health().failures_this_minute("d1") == 0 + assert await core.health().usage("g") == (None, None) + + +async def test_an_exception_with_no_status_is_counted_and_never_cooled( + make_core: MakeCore, host: SharedLoggingHost +) -> None: + core: Final = make_core( + model_list=[deployment("g", "a"), deployment("g", "b")], num_retries=0, allowed_fails=0, cooldown_time=60 + ) + host.answer("*", fresh(plain_error)) + + _ = await attempt(core, "g") + + assert cooled_down(core, ["a", "b"]) == {} + assert core.health().failures_this_minute("a") + core.health().failures_this_minute("b") == 1 + + +# --- _trigger_cooldown_for_failed_deployment --------------------------------------------------------- + + +def selection_failure() -> Exception: + error: Final = rate_limited() + mark_deployment_selection_failure(error) + return error + + +@pytest.mark.parametrize( + ("builder", "stalled", "counted"), + [(server_error, False, 1), (selection_failure, False, 1), (timeout, True, 0), (timeout, False, 1)], + ids=["server-error", "selection-tag-not-exempt", "stream-stall-exempt", "plain-timeout"], +) +async def test_a_failed_hop_past_the_first_logged_failure_takes_the_fallback_paths_exits( + make_core: MakeCore, host: SharedLoggingHost, builder: ErrorBuilder, stalled: bool, counted: int +) -> None: + core: Final = make_core( + model_list=[deployment("a", "a"), deployment("b", "b")], fallbacks=[{"a": ["b"]}], num_retries=0 + ) + host.answer("openai/a", fresh(server_error)) + host.answer("openai/b", fresh(builder)) + if stalled: + host.timeout_cause = MidStreamFallbackError(message="stalled", model="b", llm_provider="openai") + + _ = await attempt(core, "a", litellm_logging_obj=request_logging()) + + assert core.health().failures_this_minute("b") == counted + assert await core.health().usage("b") == (None, None) + + +@pytest.mark.parametrize(("logged", "counted"), [(True, 0), (False, 1)], ids=["generic", "chat"]) +async def test_a_failed_hops_404_is_the_callers_only_on_a_generic_call( + make_core: MakeCore, host: SharedLoggingHost, logged: bool, counted: int +) -> None: + core: Final = make_core( + model_list=[deployment("a", "a"), deployment("b", "b")], fallbacks=[{"a": ["b"]}], num_retries=0 + ) + host.answer("openai/a", fresh(server_error)) + host.answer("openai/b", fresh(not_found)) + + logging_obj: Final = request_logging() + if logged: + with pytest.raises(litellm.InternalServerError): + await core.agenerate_content(model="a", contents=[], litellm_logging_obj=logging_obj) # pyright: ignore[reportAttributeAccessIssue] # see above + else: + _ = await attempt(core, "a", litellm_logging_obj=logging_obj) + + assert core.health().failures_this_minute("b") == counted + + +@pytest.mark.parametrize( + ("params", "router_cooldown", "error", "expected"), + [ + pytest.param({}, 7, server_error, 7, id="router-int"), + pytest.param({}, 7.0, server_error, 7.0, id="router-float"), + pytest.param({"cooldown_time": 7}, 7.0, server_error, 7, id="deployment-int-over-router-float"), + pytest.param({"cooldown_time": 7.0}, 7, server_error, 7.0, id="deployment-float-over-router-int"), + pytest.param({}, 9.0, lambda: rate_limited("9"), 9, id="header-int-over-router-float"), + ], +) +async def test_a_failed_hop_fires_the_number_python_chose_as_it_was_written( + make_core: MakeCore, + host: SharedLoggingHost, + monkeypatch: pytest.MonkeyPatch, + params: Mapping[str, object], + router_cooldown: float, + error: Callable[[], Exception], + expected: float, +) -> None: + fired: Final[list[tuple[str, object]]] = [] # mutable-ok: appended as they fire + + async def record(ctx: object, deployment_id: str, exception_status: object, cooldown_time: object) -> None: + fired.append((deployment_id, cooldown_time)) + + monkeypatch.setattr(cooldown_handlers, "router_cooldown_event_callback", record) + core: Final = make_core( + model_list=[deployment("a", "a"), deployment("b", "b", **params), deployment("b", "c", **params)], + fallbacks=[{"a": ["b"]}], + num_retries=0, + allowed_fails=0, + cooldown_time=router_cooldown, + ) + host.answer("openai/a", fresh(server_error)) + host.answer("openai/b", fresh(error)) + host.answer("openai/c", fresh(error)) + + _ = await attempt(core, "a", litellm_logging_obj=request_logging()) + for _ in range(20): + await asyncio.sleep(0) + + assert [repr(seconds) for name, seconds in fired if name != "a"] == [repr(expected)] + + +# --- P6, PV7: what reaches Redis --------------------------------------------------------------------- + + +async def test_cooldowns_are_read_as_one_batch_of_every_deployment_with_no_key_of_their_own( + make_core: MakeCore, +) -> None: + core: Final = make_core( + model_list=[deployment("g", "a"), deployment("g", "b"), deployment("other", "c")], + num_retries=0, + disable_cooldowns=True, + ) + redis: Final = with_redis(core) + + _ = await attempt(core, "g") + + assert [key for key in redis.gets if key.endswith(":cooldown")] == [] + assert redis.batches[0] == ["deployment:a:cooldown", "deployment:b:cooldown", "deployment:c:cooldown"] + + +async def test_a_second_request_within_the_read_interval_reads_no_cooldown_from_redis(make_core: MakeCore) -> None: + core: Final = make_core(model_list=[deployment("g", "a")], num_retries=0) + redis: Final = with_redis(core) + + _ = await attempt(core, "g") + _ = await attempt(core, "g") + + assert len([batch for batch in redis.batches if batch[0].endswith(":cooldown")]) == 1 + + +async def test_a_usage_write_leaves_the_shared_total_in_memory(make_core: MakeCore) -> None: + core: Final = make_core(model_list=[deployment("g", "d1", rpm=1000000)], num_retries=0) + redis: Final = with_redis(core) + rpm_key: Final = f"global_router:d1:openai/d1:rpm:{minute()}" + redis.totals[rpm_key] = 100000 + + _ = attempt_sync(core, "g") + + assert core.health().store.get_cache(rpm_key, local_only=True) == 100001 + + +async def test_a_usage_read_on_a_pod_sees_what_the_other_pods_counted(make_core: MakeCore) -> None: + core: Final = make_core(model_list=[deployment("g", "d1", rpm=1000000)], num_retries=0) + redis: Final = with_redis(core) + redis.totals[f"global_router:d1:openai/d1:rpm:{minute()}"] = 41 + + _ = attempt_sync(core, "g") + + assert (await core.health().usage("g"))[1] == 42 + + +# --- wildcard streams -------------------------------------------------------------------------------- + + +async def test_a_wildcard_stream_counts_its_tokens_under_the_filled_model( + make_core: MakeCore, host: SharedLoggingHost +) -> None: + host.answer("openai/gpt-4o", chat_stream("hello", usage=USAGE, logged=True)) + core: Final = make_core(model_list=[deployment("openai/*", "w", model="openai/*", rpm=100)]) + + stream: Final = await core.acompletion(model="openai/gpt-4o", messages=MESSAGES, stream=True) # pyright: ignore[reportAttributeAccessIssue] # see above + _ = await collect_texts(stream) + await settle() + + store: Final = core.health().store + assert store.get_cache(f"global_router:w:openai/gpt-4o:tpm:{minute()}", local_only=True) == 12 + assert store.get_cache(f"global_router:w:openai/*:tpm:{minute()}", local_only=True) is None diff --git a/tests/unit/test_router/test_router.py b/tests/unit/test_router/test_router.py index f5962cb4933..a21777fbee6 100644 --- a/tests/unit/test_router/test_router.py +++ b/tests/unit/test_router/test_router.py @@ -63,7 +63,6 @@ from litellm.router import ( _responses_stream_holds_event, _without_line_breaks, _search_tool_support, - Span, ) from litellm.proxy.response_polling.polling_handler import should_use_polling_for_request from litellm.router_strategy import simple_shuffle @@ -6663,13 +6662,14 @@ def _make_router_with_fallback(primary="gpt-4", secondary="gpt-3.5-turbo"): ) -class _InjectedFallbackRouter(Router): - def __init__(self, fallback_response: object) -> None: - super().__init__(model_list=[], fallbacks=[{"primary": ["fallback"]}]) - self._fallback_response: Final = fallback_response +# The Rust core walks fallbacks and runs pre-routing in Rust, past the members these tests patch. +_PYTHON_CORE_ONLY: Final = pytest.mark.usefixtures("python_router_core") - async def async_function_with_fallbacks_common_utils( - self, + +def _injected_fallback_router(fallback_response: object) -> Router: + router: Final = Router(model_list=[], fallbacks=[{"primary": ["fallback"]}]) + + async def inject_fallback( e: Exception, disable_fallbacks: bool | None, fallbacks: list | None, @@ -6680,7 +6680,10 @@ class _InjectedFallbackRouter(Router): kwargs: dict[str, object], include_fallback_errors: bool = False, ) -> object: - return self._fallback_response + return fallback_response + + router.async_function_with_fallbacks_common_utils = inject_fallback + return router @pytest.mark.asyncio @@ -7043,6 +7046,7 @@ async def test_aresponses_streaming_content_policy_error_event_routes_to_content @pytest.mark.asyncio +@_PYTHON_CORE_ONLY async def test_aresponses_stream_keeps_numbering_after_a_mid_stream_fallback(): from litellm.types.llms.openai import OutputTextDeltaEvent, ResponsesAPIStreamEvents @@ -7058,7 +7062,7 @@ async def test_aresponses_stream_keeps_numbering_after_a_mid_stream_fallback(): from litellm.exceptions import MidStreamFallbackError - router: Final = _InjectedFallbackRouter(_AsyncList([delta(0, "c"), delta(1, "d")])) + router: Final = _injected_fallback_router(_AsyncList([delta(0, "c"), delta(1, "d")])) primary: Final = _make_responses_iterator( chunks=[delta(0, "a"), delta(1, "b")], error=MidStreamFallbackError( @@ -7082,6 +7086,7 @@ async def test_aresponses_stream_keeps_numbering_after_a_mid_stream_fallback(): @pytest.mark.asyncio +@_PYTHON_CORE_ONLY async def test_aresponses_stream_fallback_keeps_higher_numbers_and_passes_unnumbered_events(): from litellm.exceptions import MidStreamFallbackError from litellm.types.llms.openai import OutputTextDeltaEvent, ResponsesAPIStreamEvents @@ -7094,7 +7099,7 @@ async def test_aresponses_stream_fallback_keeps_higher_numbers_and_passes_unnumb def numbered(sequence_number: int, text: str) -> OutputTextDeltaEvent: return delta(text).model_copy(update={"sequence_number": sequence_number}) - router: Final = _InjectedFallbackRouter(_AsyncList([numbered(7, "c"), delta("d"), numbered(8, "e")])) + router: Final = _injected_fallback_router(_AsyncList([numbered(7, "c"), delta("d"), numbered(8, "e")])) primary: Final = _make_responses_iterator( chunks=[numbered(0, "a"), numbered(1, "b")], error=MidStreamFallbackError( @@ -7123,6 +7128,7 @@ async def test_aresponses_stream_fallback_keeps_higher_numbers_and_passes_unnumb @pytest.mark.asyncio +@_PYTHON_CORE_ONLY async def test_aresponses_stream_without_fallback_passes_upstream_numbers_unchanged(): from litellm.types.llms.openai import OutputTextDeltaEvent, ResponsesAPIStreamEvents @@ -7136,7 +7142,7 @@ async def test_aresponses_stream_without_fallback_passes_upstream_numbers_unchan sequence_number=sequence_number, ) - router: Final = _InjectedFallbackRouter(_AsyncList([])) + router: Final = _injected_fallback_router(_AsyncList([])) primary: Final = _make_responses_iterator( chunks=[delta(0, "a"), delta(1, "b"), delta(1, "c"), delta(0, "d")], model="primary" ) @@ -7151,6 +7157,7 @@ async def test_aresponses_stream_without_fallback_passes_upstream_numbers_unchan @pytest.mark.asyncio +@_PYTHON_CORE_ONLY async def test_aresponses_stream_fallback_continues_above_the_highest_number_the_primary_sent(): from litellm.exceptions import MidStreamFallbackError from litellm.types.llms.openai import OutputTextDeltaEvent, ResponsesAPIStreamEvents @@ -7165,7 +7172,7 @@ async def test_aresponses_stream_fallback_continues_above_the_highest_number_the sequence_number=sequence_number, ) - router: Final = _InjectedFallbackRouter(_AsyncList([delta(0, "d")])) + router: Final = _injected_fallback_router(_AsyncList([delta(0, "d")])) primary: Final = _make_responses_iterator( chunks=[delta(0, "a"), delta(3, "b"), delta(1, "c")], error=MidStreamFallbackError( @@ -7234,6 +7241,7 @@ async def test_aresponses_streaming_iterator_unwraps_content_policy_trigger_for_ ], ids=["plain-fallback", "complexity-tier-fallback"], ) +@_PYTHON_CORE_ONLY async def test_aresponses_streaming_iterator_replaces_complexity_headers_before_fallback_output( fallback_headers: dict[str, str], ) -> None: @@ -7257,7 +7265,7 @@ async def test_aresponses_streaming_iterator_replaces_complexity_headers_before_ "model_id": "fallback-deployment", "additional_headers": fallback_headers, } - router: Final = _InjectedFallbackRouter(fallback) + router: Final = _injected_fallback_router(fallback) wrapped: Final = await router._aresponses_streaming_iterator( response=source, @@ -19126,6 +19134,7 @@ async def test_anthropic_messages_fallback_merges_fallback_hidden_params(): ], ids=["plain-fallback", "complexity-tier-fallback"], ) +@_PYTHON_CORE_ONLY async def test_anthropic_messages_fallback_replaces_complexity_headers_before_output( fallback_headers: dict[str, str], ) -> None: @@ -19145,7 +19154,7 @@ async def test_anthropic_messages_fallback_replaces_complexity_headers_before_ou "additional_headers": fallback_headers, }, ) - router: Final = _InjectedFallbackRouter(fallback) + router: Final = _injected_fallback_router(fallback) wrapped: Final = await router._aanthropic_messages_streaming_iterator( response=source, @@ -22020,24 +22029,23 @@ class TestRequestReasoningEffortOverride: assert ("reasoning_effort" not in body) is should_drop +@_PYTHON_CORE_ONLY class TestPreRoutingTierDrivesFallbacks: """#38832: a complexity/auto router picks a tier behind the router name, but fallback lookup stayed on the router name, so the tier's configured chain never ran and a provider failure on the tier's first hop was returned to the client.""" - class _TierRouter(litellm.Router): - async def async_pre_routing_hook( - self, model, request_kwargs, messages=None, input=None, specific_deployment=False - ): - from litellm.types.router import PreRoutingHookResponse + @staticmethod + async def _select_tier(model, request_kwargs, messages=None, input=None, specific_deployment=False): + from litellm.types.router import PreRoutingHookResponse - if model == "smart-router": - return PreRoutingHookResponse(model="tier1", messages=messages) - return None + if model == "smart-router": + return PreRoutingHookResponse(model="tier1", messages=messages) + return None @classmethod def _router(cls, fallbacks) -> "litellm.Router": - return cls._TierRouter( + router = litellm.Router( model_list=[ { "model_name": "smart-router", @@ -22087,6 +22095,8 @@ class TestPreRoutingTierDrivesFallbacks: fallbacks=fallbacks, num_retries=0, ) + router._core_async_pre_routing_hook = cls._select_tier + return router @pytest.mark.asyncio async def test_the_selected_tier_fallback_chain_runs(self): @@ -24545,50 +24555,6 @@ def test_a_failed_routing_read_prefetch_logs_the_request_model_without_its_line_ ] -@pytest.mark.asyncio -@pytest.mark.parametrize( - "routing_strategy", - ["simple-shuffle", "usage-based-routing-v2", "least-busy", "latency-based-routing"], -) -async def test_router_subclass_overriding_async_get_healthy_deployments_with_the_old_signature_still_routes( - routing_strategy: str, -) -> None: - class OldSignatureRouter(litellm.Router): - async def async_get_healthy_deployments( - self, - model: str, - request_kwargs: dict, - messages: list[dict[str, str]] | None = None, - input: str | list | None = None, - specific_deployment: bool | None = False, - parent_otel_span: Span | None = None, - health_check_probe: bool = False, - ): - return await super().async_get_healthy_deployments( - model=model, - request_kwargs=request_kwargs, - messages=messages, - input=input, - specific_deployment=specific_deployment, - parent_otel_span=parent_otel_span, - health_check_probe=health_check_probe, - ) - - router: Final = OldSignatureRouter( - model_list=[ - { - "model_name": "m", - "litellm_params": {"model": "openai/gpt-4o", "api_key": "x", "mock_response": "hi"}, - } - ], - routing_strategy=routing_strategy, - ) - - response: Final = await router.acompletion(model="m", messages=[{"role": "user", "content": "x"}]) - - assert response.choices[0].message.content == "hi" - - def test_get_deployment_credentials_with_provider_preserves_anthropic_wif_params(): """ Test that get_deployment_credentials_with_provider preserves a litellm_params-configured diff --git a/tests/unit/test_router/test_router_entry_points.py b/tests/unit/test_router/test_router_entry_points.py index 81c1982b987..25dd8fe47cf 100644 --- a/tests/unit/test_router/test_router_entry_points.py +++ b/tests/unit/test_router/test_router_entry_points.py @@ -334,8 +334,10 @@ def test_settings_snapshot_is_reused_until_a_setting_is_written(router): assert router.settings().fallbacks == [{"gpt": ["team-gpt"]}] router._core_fallbacks = [{"gpt": ["other"]}] # the router's own writes invalidate it too assert router.settings().fallbacks == [{"gpt": ["other"]}] - router.settings().fallbacks.append({"x": ["y"]}) # the snapshot holds the router's own objects - assert router._core_fallbacks[-1] == {"x": ["y"]} + router.settings().fallbacks.append( + {"x": ["y"]} + ) # the snapshot's lists are copies: settings change only via apply() + assert router._core_fallbacks == [{"gpt": ["other"]}] unrelated = router.settings() router._core_model_list = list(router._core_model_list) assert router.settings() is unrelated @@ -446,84 +448,30 @@ def test_generate_model_id_is_the_module_function(): assert Router.generate_model_id("g", params) == deployment_id("g", params) -def test_subclass_override_of_an_old_member_is_still_called(): - calls = [] +def test_router_cannot_be_subclassed(): + with pytest.raises(TypeError, match=r"^litellm\.Router cannot be subclassed"): - class Sub(litellm.Router): - async def async_pre_routing_hook( - self, model, request_kwargs, messages=None, input=None, specific_deployment=False - ): - calls.append(model) - return await super().async_pre_routing_hook(model, request_kwargs, messages, input, specific_deployment) - - r = Sub( - model_list=[ - { - "model_name": "m", - "litellm_params": {"model": "openai/gpt-4o-mini", "api_key": "sk-x", "mock_response": "ok"}, - } - ] - ) - response = asyncio.run(r.acompletion(model="m", messages=[{"role": "user", "content": "hi"}])) - assert response.choices[0].message.content == "ok" - assert calls == ["m"] + class Sub(litellm.Router): + pass -def test_direct_call_of_a_subclass_override_runs_it_once(): - calls = [] +def test_core_classes_can_still_be_subclassed(): + from litellm.router_python_core import PythonCore - class Sub(litellm.Router): - def get_model_access_groups(self, model_name=None, model_access_group=None, team_id=None): - calls.append(model_name) - return super().get_model_access_groups(model_name, model_access_group, team_id) + class PythonSub(PythonCore): + pass - r = Sub(model_list=[{"model_name": "m", "litellm_params": {"model": "openai/gpt-4o-mini", "api_key": "sk-x"}}]) - with warnings.catch_warnings(): - warnings.simplefilter("ignore") - assert r.get_model_access_groups("m") == {} - assert calls == ["m"] - assert asyncio.run(asyncio.to_thread(r.registry().access_groups, model_name="m")) == {} - assert calls == ["m", "m"] + assert isinstance(object.__new__(PythonSub), Router) + try: + from litellm.router_rust.core import RustCore + except ImportError: + pytest.skip("the native lib does not load") -def test_subclass_staticmethod_override_is_still_consulted(): - calls = [] + class RustSub(RustCore): + pass - class Sub(litellm.Router): - @staticmethod - def is_deployment_blocked(deployment): - calls.append(deployment.model_info.id) - return True - - r = Sub( - model_list=[ - { - "model_name": "m", - "litellm_params": {"model": "openai/gpt-4o-mini", "api_key": "sk-x"}, - "model_info": {"id": "d1"}, - } - ] - ) - with warnings.catch_warnings(): - warnings.simplefilter("ignore") - assert r.get_deployment_credentials(model_id="d1") is None - assert calls == ["d1"] - - -def test_private_alias_skips_a_subclass_override_of_the_public_name(): - calls = [] - - class Sub(litellm.Router): - def get_all_deployments(self, *args, **kwargs): - calls.append("get_all_deployments") - return super().get_all_deployments(*args, **kwargs) - - r = Sub(model_list=[{"model_name": "m", "litellm_params": {"model": "openai/gpt-4o-mini", "api_key": "sk-x"}}]) - calls.clear() - with warnings.catch_warnings(): - warnings.simplefilter("ignore") - assert len(r._get_all_deployments(model_name="m")) == 1 - assert calls == [] + assert issubclass(RustSub, RustCore) def test_old_name_writes_and_deletes_warn(router): @@ -779,23 +727,6 @@ def test_old_members_refuse_objects_that_are_not_routers(): _as_router(object()) # pyright: ignore[reportArgumentType] -def test_direct_call_of_an_async_subclass_override_runs_it_once(): - calls = [] - - class Sub(litellm.Router): - async def async_pre_routing_hook( - self, model, request_kwargs, messages=None, input=None, specific_deployment=False - ): - calls.append(model) - return await super().async_pre_routing_hook(model, request_kwargs, messages, input, specific_deployment) - - r = Sub(model_list=[{"model_name": "m", "litellm_params": {"model": "openai/gpt-4o-mini", "api_key": "sk-x"}}]) - with warnings.catch_warnings(): - warnings.simplefilter("ignore") - asyncio.run(r.async_pre_routing_hook("m", {})) - assert calls == ["m"] - - def test_routing_context_forwards_to_the_router(router, monkeypatch): from unittest.mock import AsyncMock, MagicMock diff --git a/tests/unit/test_router_model_cost_isolation.py b/tests/unit/test_router_model_cost_isolation.py index 931fabd85ad..724c07f8c31 100644 --- a/tests/unit/test_router_model_cost_isolation.py +++ b/tests/unit/test_router_model_cost_isolation.py @@ -174,20 +174,26 @@ async def test_discovery_discards_metadata_for_a_replaced_deployment(monkeypatch _invalidate_model_cost_lowercase_map() -async def test_discovery_sends_the_secret_behind_an_oidc_api_key(monkeypatch: pytest.MonkeyPatch, tmp_path: Path) -> None: +async def test_discovery_sends_the_secret_behind_an_oidc_api_key( + monkeypatch: pytest.MonkeyPatch, tmp_path: Path +) -> None: monkeypatch.setattr(litellm, "model_cost", copy.deepcopy(litellm.model_cost)) monkeypatch.setenv("LITELLM_OIDC_ALLOWED_CREDENTIAL_DIRS", str(tmp_path)) token: Final = tmp_path / "token" token.write_text("local-secret") - router: Final = Router(model_list=[{ - "model_name": "local", - "litellm_params": { - "model": "hosted_vllm/local-model", - "api_base": "https://discovery.test/v1", - "api_key": f"oidc/file/{token}", - }, - "model_info": {"id": "oidc-deployment"}, - }]) + router: Final = Router( + model_list=[ + { + "model_name": "local", + "litellm_params": { + "model": "hosted_vllm/local-model", + "api_base": "https://discovery.test/v1", + "api_key": f"oidc/file/{token}", + }, + "model_info": {"id": "oidc-deployment"}, + } + ] + ) authorizations: Final[list[str | None]] = [] def respond(request: httpx.Request) -> httpx.Response: @@ -1074,8 +1080,12 @@ def test_bedrock_deployment_keeps_custom_prices_when_profile_lookup_fails( tier_rates: Final = inherit_configured and service_tier != "default" input_rate: Final = configured_prices[f"input_cost_per_token_{service_tier}"] if tier_rates else 0.01 output_rate: Final = configured_prices[f"output_cost_per_token_{service_tier}"] if tier_rates else 0.02 - effective_input: Final = 0.0 if override == "input_cost_per_token_priority" and service_tier == "priority" else input_rate - cache_rate: Final = 0.0 if override == "cache_read_input_token_cost" else (0.001 if inherit_configured else effective_input) + effective_input: Final = ( + 0.0 if override == "input_cost_per_token_priority" and service_tier == "priority" else input_rate + ) + cache_rate: Final = ( + 0.0 if override == "cache_read_input_token_cost" else (0.001 if inherit_configured else effective_input) + ) assert completion_cost( model=model, completion_response=ModelResponse( @@ -1565,6 +1575,8 @@ def test_ptu_service_tier_pricing_is_disabled_only_when_attribution_is_enabled( response: Final = router.completion( model=f"ptu-tier-model-{ptu_enabled}", messages=[{"role": "user", "content": "ptu service tier pricing"}], + # The deployment belongs to team-alpha, and a sync call is team filtered like an async one. + metadata={"user_api_key_team_id": "team-alpha"}, service_tier="priority", mock_response=litellm.ModelResponse( model=_TIER_BACKEND_MODEL, @@ -2179,8 +2191,7 @@ def test_token_limits_stay_on_the_deployment_that_set_them(override_first: bool) override_id: Final = "token-limit-override-deployment" plain_id: Final = "token-limit-plain-deployment" model_keys: Final = { - key: copy.deepcopy(litellm.model_cost.get(key)) - for key in (*shared_keys, override_id, plain_id) + key: copy.deepcopy(litellm.model_cost.get(key)) for key in (*shared_keys, override_id, plain_id) } expected_cost_map_limits: Final = { key: { @@ -2217,13 +2228,15 @@ def test_token_limits_stay_on_the_deployment_that_set_them(override_first: bool) assert { field: (litellm.model_cost.get(key) or {}).get(field) for field in token_fields } == expected_cost_map_limits[key] - assert { - field: litellm.get_model_info(model=key).get(field) for field in token_fields - } == catalog_limits[key] + assert {field: litellm.get_model_info(model=key).get(field) for field in token_fields} == catalog_limits[ + key + ] - assert { - field: litellm.get_model_info(model=override_id).get(field) for field in token_fields - } == {"max_input_tokens": 4321, "max_output_tokens": 987, "max_tokens": 987} + assert {field: litellm.get_model_info(model=override_id).get(field) for field in token_fields} == { + "max_input_tokens": 4321, + "max_output_tokens": 987, + "max_tokens": 987, + } plain_listing: Final = create_model_info_response( model_id="plain-gpt-4o-mini", provider="openai", llm_router=router @@ -2280,9 +2293,7 @@ def test_unknown_backend_token_limits_still_fill_the_shared_key() -> None: def test_unknown_backend_token_limit_updates_when_deployment_is_upserted() -> None: backend_model: Final = "hosted_vllm/somelab.lit4544-token-limit-update" model_id: Final = "unknown-token-limit-updated-deployment" - model_keys: Final = { - key: copy.deepcopy(litellm.model_cost.get(key)) for key in (backend_model, model_id) - } + model_keys: Final = {key: copy.deepcopy(litellm.model_cost.get(key)) for key in (backend_model, model_id)} try: deployment: Final = Deployment( model_name="unknown-token-limit-update", @@ -2999,7 +3010,8 @@ def test_a_github_copilot_chat_pin_opts_a_dual_endpoint_model_out_of_the_respons @pytest.mark.parametrize("registered_mode, pinned_mode", [("chat", "responses"), ("responses", "chat")]) @pytest.mark.parametrize( - "model, custom_llm_provider", [("github_copilot/routing-test-model", None), ("routing-test-model", "github_copilot")] + "model, custom_llm_provider", + [("github_copilot/routing-test-model", None), ("routing-test-model", "github_copilot")], ) def test_a_github_copilot_responses_pin_still_opts_the_model_into_the_responses_api( monkeypatch: pytest.MonkeyPatch, @@ -4181,7 +4193,9 @@ def test_price_data_reload_refreshes_the_cached_model_group_and_deployment_info( assert router.cached_deployment_model_info("dep-a", "openai/gpt-4o")["input_cost_per_token"] == new_price -def test_price_data_reload_keeps_geo_deployments_priced_while_requests_register_pricing(monkeypatch: pytest.MonkeyPatch) -> None: +def test_price_data_reload_keeps_geo_deployments_priced_while_requests_register_pricing( + monkeypatch: pytest.MonkeyPatch, +) -> None: monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") catalog: Final = litellm.get_model_cost_map(url="") monkeypatch.setattr(litellm, "model_cost", dict(catalog)) diff --git a/tests/unit/test_router_order_fallback.py b/tests/unit/test_router_order_fallback.py index 47e0be28635..6dd445e2b5d 100644 --- a/tests/unit/test_router_order_fallback.py +++ b/tests/unit/test_router_order_fallback.py @@ -95,8 +95,9 @@ class TestGetOrderFilteredDeployments: assert len(result) == 2 -def test_get_deployment_order_returns_unvalidated_order(): - assert get_deployment_order({"litellm_params": {"order": "first"}}) == "first" +def test_get_deployment_order_reads_a_non_number_as_unordered(): + assert get_deployment_order({"litellm_params": {"order": "first"}}) is None + assert get_deployment_order({"litellm_params": {"order": "2"}}) == 2 # --------------------------------------------------------------------------- @@ -646,7 +647,6 @@ async def test_text_completion_order_fallback_hop_does_not_send_target_order_ups assert all("_target_order" not in body for body in upstream_bodies) - _OPENAI_RESPONSES_URL: Final = "https://api.openai.com/v1/responses" _MANTLE_RESPONSES_URL: Final = "https://bedrock-mantle.us-east-1.api.aws/openai/v1/responses" _OVERLOADED_UPSTREAM: Final = {"error": {"message": "overloaded", "type": "server_error", "code": "server_error"}}