mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-12 03:48:56 +00:00
feat(router): put fallback targets to the proxy's access and budget checks from rust
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01Aw1xktpUnix5tWhvSzey8t
This commit is contained in:
parent
86e7541cd5
commit
d458d55a01
11 changed files with 367 additions and 15 deletions
|
|
@ -4,7 +4,7 @@ use litellm_host_python::{
|
|||
};
|
||||
use litellm_router::{
|
||||
engine::Routed,
|
||||
host::{Attempt, Invoked, RouterHost},
|
||||
host::{Attempt, FallbackCheck, Invoked, RouterHost},
|
||||
};
|
||||
use pyo3::{
|
||||
exceptions::PyRuntimeError,
|
||||
|
|
@ -34,11 +34,18 @@ impl RouterHost for BridgeHost {
|
|||
.call(|reply| RouterHostCall::Sleep { seconds, reply })
|
||||
.await
|
||||
}
|
||||
|
||||
async fn allow_fallback(&self, check: FallbackCheck<PyObj>) -> Result<bool, BridgeError> {
|
||||
self.0
|
||||
.call(|reply| RouterHostCall::AllowFallback { check, reply })
|
||||
.await
|
||||
}
|
||||
}
|
||||
|
||||
enum Pending {
|
||||
Invoke(Reply<Invoked<PyObj, PyObj>>),
|
||||
Sleep(Reply<()>),
|
||||
AllowFallback(Reply<bool>),
|
||||
}
|
||||
|
||||
pub(super) struct RouterBinding {
|
||||
|
|
@ -67,6 +74,7 @@ impl RouterBinding {
|
|||
match self.pending.take() {
|
||||
Some(Pending::Invoke(reply)) => reply.send(python::invoked(py, result)?),
|
||||
Some(Pending::Sleep(reply)) => reply.send(()),
|
||||
Some(Pending::AllowFallback(reply)) => reply.send(result.extract()?),
|
||||
None => {
|
||||
return Err(PyRuntimeError::new_err(
|
||||
"router host reply without a pending op",
|
||||
|
|
@ -168,6 +176,15 @@ impl PythonHostCalls<RouterProtocol> for RouterBinding {
|
|||
pyo3::types::PyFloat::new(py, seconds).into_any().unbind(),
|
||||
Pending::Sleep(reply),
|
||||
),
|
||||
RouterHostCall::AllowFallback { check, reply } => (
|
||||
if self.asynchronous {
|
||||
"allow_fallback"
|
||||
} else {
|
||||
"allow_fallback_sync"
|
||||
},
|
||||
python::fallback_check(py, &check)?,
|
||||
Pending::AllowFallback(reply),
|
||||
),
|
||||
};
|
||||
self.pending = Some(pending);
|
||||
let result = driver.call_method1(method, (argument,));
|
||||
|
|
|
|||
|
|
@ -15,7 +15,7 @@ use litellm_host::{
|
|||
use litellm_host_python::{HookChain, enter_native, from_py_argument};
|
||||
use litellm_router::{
|
||||
engine::{Engine, Failed, Resume, RouteError, Routed, RouterCall},
|
||||
host::{Attempt, Invoked},
|
||||
host::{Attempt, FallbackCheck, Invoked},
|
||||
random::PythonRandom,
|
||||
settings::Settings,
|
||||
snapshot::{RoutedDeployment, Snapshot},
|
||||
|
|
@ -60,6 +60,10 @@ pub(crate) enum RouterHostCall {
|
|||
seconds: f64,
|
||||
reply: Reply<()>,
|
||||
},
|
||||
AllowFallback {
|
||||
check: FallbackCheck<PyObj>,
|
||||
reply: Reply<bool>,
|
||||
},
|
||||
}
|
||||
|
||||
/// A call to route, or to resume from a stream that failed after its content reached the caller.
|
||||
|
|
|
|||
|
|
@ -7,7 +7,7 @@ use litellm_host_python::from_py;
|
|||
use litellm_router::{
|
||||
engine::{Outcome, Override, Resume, RouterCall},
|
||||
failure::{Classified, Raised, Rejection},
|
||||
host::{Attempt, Invoked, MockFailure, Op, Target, TypedFallback},
|
||||
host::{Attempt, FallbackCheck, Invoked, MockFailure, Op, Target, TypedFallback},
|
||||
operation::Operation,
|
||||
pyrepr::PyNumber,
|
||||
settings::Fallbacks,
|
||||
|
|
@ -148,6 +148,16 @@ pub(super) fn attempt<'py>(py: Python<'py>, attempt: &Attempt<PyObj>) -> PyResul
|
|||
Ok(dict.into_any().unbind())
|
||||
}
|
||||
|
||||
pub(super) fn fallback_check(py: Python<'_>, check: &FallbackCheck<PyObj>) -> PyResult<Py<PyAny>> {
|
||||
let dict = PyDict::new(py);
|
||||
dict.set_item("target", &check.target)?;
|
||||
dict.set_item("model_group", &check.model_group)?;
|
||||
dict.set_item("bucket", check.bucket)?;
|
||||
dict.set_item("model", &check.model)?;
|
||||
dict.set_item("ops", ops(py, &check.ops)?)?;
|
||||
Ok(dict.into_any().unbind())
|
||||
}
|
||||
|
||||
pub(super) fn outcome(py: Python<'_>, outcome: &Outcome) -> PyResult<Py<PyAny>> {
|
||||
let dict = PyDict::new(py);
|
||||
dict.set_item("model_group", &outcome.model_group)?;
|
||||
|
|
|
|||
|
|
@ -13,7 +13,8 @@ use crate::{
|
|||
failure::{Classified, ExceptionClass, Failure, Raised, StreamFailure},
|
||||
fallback::{chain_for_groups, generic_targets},
|
||||
host::{
|
||||
Attempt, HopStamp, Invoked, MockFailure, Op, RetryStamp, RouterHost, Target, TypedFallback,
|
||||
Attempt, FallbackCheck, HopStamp, Invoked, MockFailure, Op, RetryStamp, RouterHost, Target,
|
||||
TypedFallback,
|
||||
},
|
||||
operation::Operation,
|
||||
random::PythonRandom,
|
||||
|
|
@ -877,7 +878,22 @@ impl<'a, H: RouterHost> Run<'a, H> {
|
|||
let mut model = hop.group.clone();
|
||||
let mut last = original.clone();
|
||||
for target in targets {
|
||||
if *target == hop.group || !self.attempted_targets.insert(target.clone()) {
|
||||
if *target == hop.group {
|
||||
continue;
|
||||
}
|
||||
if self.snapshot.settings.fallback_checks {
|
||||
let check = FallbackCheck {
|
||||
target: target.clone(),
|
||||
model_group: hop.group.clone(),
|
||||
bucket,
|
||||
model: model.clone(),
|
||||
ops: std::mem::take(&mut self.ops),
|
||||
};
|
||||
if !self.host.allow_fallback(check).await.map_err(Stop::Host)? {
|
||||
continue;
|
||||
}
|
||||
}
|
||||
if !self.attempted_targets.insert(target.clone()) {
|
||||
continue;
|
||||
}
|
||||
self.ops.push(Op::LogRetry {
|
||||
|
|
|
|||
|
|
@ -108,6 +108,19 @@ pub struct Attempt<E> {
|
|||
pub ops: Vec<Op<E>>,
|
||||
}
|
||||
|
||||
/// `run_async_fallback`'s access and budget checks for one fallback target, asked against
|
||||
/// `bucket` with `kwargs["model"] == model`, the hop the chain falls back from. `ops` are
|
||||
/// applied first, so the request the checks read is the one Python's would be.
|
||||
#[derive(Clone, Debug)]
|
||||
pub struct FallbackCheck<E> {
|
||||
pub target: String,
|
||||
/// The group whose chain this is, which the checks never reject.
|
||||
pub model_group: String,
|
||||
pub bucket: u32,
|
||||
pub model: String,
|
||||
pub ops: Vec<Op<E>>,
|
||||
}
|
||||
|
||||
pub enum Invoked<R, E> {
|
||||
Success(R),
|
||||
Failure { error: E, classified: Classified },
|
||||
|
|
@ -125,4 +138,11 @@ pub trait RouterHost: Send + Sync {
|
|||
) -> impl Future<Output = Result<Invoked<Self::Response, Self::Error>, Self::Fault>> + Send;
|
||||
|
||||
fn sleep(&self, seconds: f64) -> impl Future<Output = Result<(), Self::Fault>> + Send;
|
||||
|
||||
/// Whether the fallback chain may try `check.target`. Asked only when the settings say a
|
||||
/// check is configured.
|
||||
fn allow_fallback(
|
||||
&self,
|
||||
check: FallbackCheck<Self::Error>,
|
||||
) -> impl Future<Output = Result<bool, Self::Fault>> + Send;
|
||||
}
|
||||
|
|
|
|||
|
|
@ -127,6 +127,10 @@ pub struct Settings {
|
|||
pub cooldown_time: PyNumber,
|
||||
pub disable_cooldowns: bool,
|
||||
pub enable_pre_call_checks: bool,
|
||||
/// A `fallback_access_check` or `fallback_budget_check` is set, so each cross-group
|
||||
/// fallback target is put to the host before it is tried.
|
||||
#[serde(default)]
|
||||
pub fallback_checks: bool,
|
||||
#[serde(default)]
|
||||
pub tunables: Tunables,
|
||||
}
|
||||
|
|
|
|||
|
|
@ -6,7 +6,7 @@ use std::{
|
|||
use litellm_router::{
|
||||
engine::{Engine, Override, Resume, RouteError, RouterCall},
|
||||
failure::{Classified, ExceptionClass, Raised, StreamFailure},
|
||||
host::{Attempt, Invoked, Op, RouterHost, Target},
|
||||
host::{Attempt, FallbackCheck, Invoked, Op, RouterHost, Target},
|
||||
operation::Operation,
|
||||
pyrepr::PyNumber,
|
||||
random::PythonRandom,
|
||||
|
|
@ -30,6 +30,8 @@ struct ScriptedHost {
|
|||
attempts: Mutex<Vec<String>>,
|
||||
sleeps: Mutex<Vec<f64>>,
|
||||
ops: Mutex<Vec<String>>,
|
||||
denied: Vec<String>,
|
||||
checks: Mutex<Vec<String>>,
|
||||
}
|
||||
|
||||
impl ScriptedHost {
|
||||
|
|
@ -44,6 +46,15 @@ impl ScriptedHost {
|
|||
fn attempts(&self) -> Vec<String> {
|
||||
self.attempts.lock().unwrap().clone()
|
||||
}
|
||||
|
||||
fn deny(mut self, target: &str) -> Self {
|
||||
self.denied.push(target.into());
|
||||
self
|
||||
}
|
||||
|
||||
fn checks(&self) -> Vec<String> {
|
||||
self.checks.lock().unwrap().clone()
|
||||
}
|
||||
}
|
||||
|
||||
fn op_name(op: &Op<String>) -> String {
|
||||
|
|
@ -128,6 +139,18 @@ impl RouterHost for ScriptedHost {
|
|||
self.sleeps.lock().unwrap().push(seconds);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn allow_fallback(&self, check: FallbackCheck<String>) -> Result<bool, ()> {
|
||||
self.ops
|
||||
.lock()
|
||||
.unwrap()
|
||||
.extend(check.ops.iter().map(op_name));
|
||||
self.checks.lock().unwrap().push(format!(
|
||||
"{}->{} (bucket {}, model {})",
|
||||
check.model_group, check.target, check.bucket, check.model
|
||||
));
|
||||
Ok(!self.denied.contains(&check.target))
|
||||
}
|
||||
}
|
||||
|
||||
fn classified(status: i64, class: ExceptionClass) -> Classified {
|
||||
|
|
@ -635,6 +658,10 @@ impl RouterHost for ReplacingHost {
|
|||
async fn sleep(&self, seconds: f64) -> Result<(), ()> {
|
||||
self.inner.sleep(seconds).await
|
||||
}
|
||||
|
||||
async fn allow_fallback(&self, check: FallbackCheck<String>) -> Result<bool, ()> {
|
||||
self.inner.allow_fallback(check).await
|
||||
}
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
|
|
@ -660,3 +687,80 @@ async fn a_running_call_keeps_its_snapshot_while_the_next_call_sees_the_new_one(
|
|||
("a".to_owned(), "b".to_owned())
|
||||
);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::checks_off(false, &["a", "b"], &[])]
|
||||
#[case::checks_on(true, &["a", "c"], &["g->h (bucket 0, model g)", "g->i (bucket 0, model g)"])]
|
||||
#[tokio::test]
|
||||
async fn a_fallback_target_the_checks_reject_is_skipped(
|
||||
#[case] fallback_checks: bool,
|
||||
#[case] expected_attempts: &[&str],
|
||||
#[case] expected_checks: &[&str],
|
||||
) {
|
||||
let mut settings = settings();
|
||||
settings.fallbacks = Some(vec![chain("g", &["h", "i"])]);
|
||||
settings.num_retries = 0;
|
||||
settings.fallback_checks = fallback_checks;
|
||||
let engine = engine(
|
||||
vec![
|
||||
deployment("a", "g"),
|
||||
deployment("b", "h"),
|
||||
deployment("c", "i"),
|
||||
],
|
||||
settings,
|
||||
);
|
||||
let host = ScriptedHost::default().script("a", &[SERVER]).deny("h");
|
||||
|
||||
let routed = engine.route(&host, call("g")).await.ok().unwrap();
|
||||
|
||||
assert_eq!(host.attempts(), expected_attempts);
|
||||
assert_eq!(host.checks(), expected_checks);
|
||||
assert_eq!(routed.outcome.attempted_fallbacks, 1);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn a_failed_fallback_hop_checks_the_rest_of_the_chain_as_its_own_group() {
|
||||
let mut settings = settings();
|
||||
settings.fallbacks = Some(vec![chain("g", &["h", "i"])]);
|
||||
settings.num_retries = 0;
|
||||
settings.fallback_checks = true;
|
||||
let engine = engine(
|
||||
vec![
|
||||
deployment("a", "g"),
|
||||
deployment("b", "h"),
|
||||
deployment("c", "i"),
|
||||
],
|
||||
settings,
|
||||
);
|
||||
let host = ScriptedHost::default()
|
||||
.script("a", &[SERVER])
|
||||
.script("b", &[SERVER]);
|
||||
|
||||
engine.route(&host, call("g")).await.ok().unwrap();
|
||||
|
||||
assert_eq!(host.attempts(), ["a", "b", "c"]);
|
||||
assert_eq!(
|
||||
host.checks(),
|
||||
["g->h (bucket 0, model g)", "h->i (bucket 1, model h)"]
|
||||
);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn a_rejected_target_is_not_recorded_as_attempted() {
|
||||
let mut settings = settings();
|
||||
settings.fallbacks = Some(vec![chain("g", &["h", "h"])]);
|
||||
settings.num_retries = 0;
|
||||
settings.fallback_checks = true;
|
||||
let engine = engine(vec![deployment("a", "g"), deployment("b", "h")], settings);
|
||||
let host = ScriptedHost::default().script("a", &[SERVER]).deny("h");
|
||||
|
||||
let Err(RouteError::Failed(failed)) = engine.route(&host, call("g")).await else {
|
||||
panic!("every target was rejected");
|
||||
};
|
||||
|
||||
assert_eq!(host.attempts(), ["a"]);
|
||||
assert_eq!(host.checks().len(), 2);
|
||||
assert_eq!(raised(&failed.error), "a#1");
|
||||
}
|
||||
|
|
|
|||
|
|
@ -21,6 +21,7 @@ from pydantic import BaseModel, ConfigDict, TypeAdapter
|
|||
import litellm
|
||||
from litellm.exceptions import MidStreamFallbackError
|
||||
from litellm.integrations.custom_guardrail import is_guardrail_intervention
|
||||
from litellm.litellm_core_utils.asyncify import run_async_function
|
||||
from litellm.litellm_core_utils.exception_mapping_utils import (
|
||||
_get_response_headers, # pyright: ignore[reportPrivateUsage] # the header reader the cooldown callback uses
|
||||
)
|
||||
|
|
@ -60,6 +61,10 @@ from litellm.router_utils.cooldown_handlers import (
|
|||
is_background_response_cost_poll_not_found,
|
||||
is_caller_timeout_408,
|
||||
)
|
||||
from litellm.router_utils.fallback_event_handlers import (
|
||||
_is_fallback_target_authorized, # pyright: ignore[reportPrivateUsage] # run_async_fallback's access check
|
||||
_is_fallback_target_within_budget, # pyright: ignore[reportPrivateUsage] # run_async_fallback's budget check
|
||||
)
|
||||
from litellm.types.llms.openai import ResponseAPIUsage
|
||||
from litellm.types.router import RouterErrors, RouterRateLimitError
|
||||
from litellm.types.utils import ModelResponse
|
||||
|
|
@ -204,6 +209,16 @@ class Attempt(BaseModel):
|
|||
ops: tuple[Mapping[str, object], ...]
|
||||
|
||||
|
||||
class FallbackCheck(BaseModel):
|
||||
model_config = ConfigDict(frozen=True, extra="forbid")
|
||||
|
||||
target: str
|
||||
model_group: str
|
||||
bucket: int
|
||||
model: str
|
||||
ops: tuple[Mapping[str, object], ...]
|
||||
|
||||
|
||||
class Outcome(BaseModel):
|
||||
model_config = ConfigDict(frozen=True, extra="forbid")
|
||||
|
||||
|
|
@ -356,6 +371,27 @@ class RoutedCall:
|
|||
def sleep_sync(self, seconds: float) -> None:
|
||||
time.sleep(seconds)
|
||||
|
||||
async def allow_fallback(self, check: Mapping[str, object]) -> bool:
|
||||
"""`run_async_fallback`'s access and budget checks, asked of the request as the hop the
|
||||
chain falls back from left it."""
|
||||
parsed: Final = FallbackCheck.model_validate(check)
|
||||
self._apply(parsed.ops)
|
||||
kwargs: Final = {
|
||||
**self._kwargs,
|
||||
**self._hop_kwargs[parsed.bucket],
|
||||
"model": parsed.model,
|
||||
self._metadata_key: self._buckets[parsed.bucket],
|
||||
}
|
||||
router: Final = self._normalizer
|
||||
return await _is_fallback_target_authorized(
|
||||
router, parsed.target, parsed.model_group, kwargs
|
||||
) and await _is_fallback_target_within_budget(router, parsed.target, parsed.model_group, kwargs)
|
||||
|
||||
def allow_fallback_sync(self, check: Mapping[str, object]) -> bool:
|
||||
"""Python's sync calls run the async fallback loop, checks included, through `run_async_function`."""
|
||||
allowed: Final[object] = run_async_function(self.allow_fallback, check)
|
||||
return allowed is True
|
||||
|
||||
def success(self, response: object, outcome: Mapping[str, object], ops: Sequence[Mapping[str, object]]) -> object:
|
||||
self._apply(ops)
|
||||
parsed: Final = Outcome.model_validate(outcome)
|
||||
|
|
|
|||
|
|
@ -107,6 +107,8 @@ class NormalizerView(AttemptRouter, Protocol):
|
|||
cooldown_time: float
|
||||
disable_cooldowns: bool | None
|
||||
enable_pre_call_checks: bool
|
||||
fallback_access_check: object
|
||||
fallback_budget_check: object
|
||||
|
||||
def get_model_names(self) -> list[str]: ... # mutable-ok: PythonRouter's read surface
|
||||
|
||||
|
|
@ -189,6 +191,7 @@ def _settings(router: NormalizerView) -> Mapping[str, object]:
|
|||
"cooldown_time": router.cooldown_time,
|
||||
"disable_cooldowns": bool(router.disable_cooldowns),
|
||||
"enable_pre_call_checks": router.enable_pre_call_checks,
|
||||
"fallback_checks": router.fallback_access_check is not None or router.fallback_budget_check is not None,
|
||||
"tunables": {
|
||||
"initial_retry_delay": constants.INITIAL_RETRY_DELAY,
|
||||
"max_retry_delay": constants.MAX_RETRY_DELAY,
|
||||
|
|
|
|||
|
|
@ -53,6 +53,13 @@ SUPPORTED_ARGUMENTS: Final = frozenset(
|
|||
"cooldown_time",
|
||||
"disable_cooldowns",
|
||||
"max_retries",
|
||||
"router_general_settings",
|
||||
"search_tools",
|
||||
"assistants_config",
|
||||
"ignore_invalid_deployments",
|
||||
"fallback_access_check",
|
||||
"fallback_budget_check",
|
||||
"auto_router_capability_limit",
|
||||
}
|
||||
)
|
||||
_UNSET: Final = object()
|
||||
|
|
@ -79,7 +86,7 @@ class _RawDeployment(BaseModel):
|
|||
model_info: Mapping[str, object] | None = None
|
||||
|
||||
|
||||
_MODEL_LIST: Final = TypeAdapter(tuple[_RawDeployment, ...])
|
||||
_MODEL_LIST: Final = TypeAdapter(tuple[object, ...])
|
||||
|
||||
|
||||
def bind_arguments(args: tuple[object, ...], kwargs: Mapping[str, object]) -> Mapping[str, object]:
|
||||
|
|
@ -97,7 +104,7 @@ def unsupported_reason(arguments: Mapping[str, object]) -> str | None:
|
|||
for reason in (
|
||||
_argument_reason(arguments),
|
||||
_fallbacks_reason(arguments),
|
||||
_model_list_reason(arguments.get("model_list")),
|
||||
_model_list_reason(arguments.get("model_list"), arguments.get("ignore_invalid_deployments") is True),
|
||||
_global_reason(),
|
||||
)
|
||||
if reason is not None
|
||||
|
|
@ -132,10 +139,23 @@ def _fallbacks_reason(arguments: Mapping[str, object]) -> str | None:
|
|||
|
||||
|
||||
def _group_names(model_list: object) -> frozenset[str]:
|
||||
return frozenset(
|
||||
deployment.model_name for deployment in map(_readable, _entries(model_list)) if deployment is not None
|
||||
)
|
||||
|
||||
|
||||
def _entries(model_list: object) -> tuple[object, ...]:
|
||||
try:
|
||||
return frozenset(deployment.model_name for deployment in _MODEL_LIST.validate_python(model_list or ()))
|
||||
return _MODEL_LIST.validate_python(model_list or ())
|
||||
except ValidationError:
|
||||
return frozenset()
|
||||
return ()
|
||||
|
||||
|
||||
def _readable(entry: object) -> _RawDeployment | None:
|
||||
try:
|
||||
return _RawDeployment.model_validate(entry)
|
||||
except ValidationError:
|
||||
return None
|
||||
|
||||
|
||||
def _plain_chain(value: object, groups: frozenset[str]) -> bool:
|
||||
|
|
@ -158,15 +178,22 @@ def _plain_chain_of(key: str, targets: object, groups: frozenset[str]) -> bool:
|
|||
)
|
||||
|
||||
|
||||
def _model_list_reason(model_list: object) -> str | None:
|
||||
def _model_list_reason(model_list: object, ignore_invalid: bool) -> str | None:
|
||||
"""With `ignore_invalid_deployments`, Python drops an entry it cannot read instead of failing,
|
||||
so the Rust router only serves the ones it can."""
|
||||
if model_list is None:
|
||||
return None
|
||||
try:
|
||||
deployments: Final = _MODEL_LIST.validate_python(model_list)
|
||||
except ValidationError:
|
||||
if not isinstance(model_list, (list, tuple)):
|
||||
return "a model_list the Rust router cannot read"
|
||||
deployments: Final = tuple(map(_readable, _entries(model_list)))
|
||||
if None in deployments and not ignore_invalid:
|
||||
return "model_list entries the Rust router cannot read"
|
||||
return next(
|
||||
(reason for reason in map(_deployment_reason, deployments) if reason is not None),
|
||||
(
|
||||
reason
|
||||
for reason in (_deployment_reason(deployment) for deployment in deployments if deployment is not None)
|
||||
if reason is not None
|
||||
),
|
||||
None,
|
||||
)
|
||||
|
||||
|
|
|
|||
111
tests/test_litellm_rust/router/test_proxy_switchover.py
Normal file
111
tests/test_litellm_rust/router/test_proxy_switchover.py
Normal file
|
|
@ -0,0 +1,111 @@
|
|||
"""What the proxy needs from the Rust router backend: its constructor arguments, its fallback
|
||||
checks, its reads of the router, and a hand-over to Python for whatever Rust does not serve."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from collections.abc import Callable, Mapping
|
||||
from typing import Final
|
||||
|
||||
import pytest
|
||||
|
||||
from tests.test_litellm_rust.support.router_backends import BACKENDS, Backend, router_headers
|
||||
|
||||
pytestmark = pytest.mark.requires_rust_extension
|
||||
|
||||
MESSAGES: Final = ({"role": "user", "content": "hi"},)
|
||||
|
||||
|
||||
def _deployment(
|
||||
group: str, name: str, response: object = None
|
||||
) -> dict[str, object]: # mutable-ok: PythonRouter consumes model_list dicts
|
||||
return {
|
||||
"model_name": group,
|
||||
"litellm_params": {"model": f"openai/{name}", "api_key": "k", "mock_response": response or f"from {name}"},
|
||||
"model_info": {"id": name},
|
||||
}
|
||||
|
||||
|
||||
class _Check:
|
||||
"""A fallback access or budget check that records what it was asked and rejects `denied`."""
|
||||
|
||||
def __init__(self, denied: frozenset[str]) -> None:
|
||||
self.denied: Final = denied
|
||||
self.asked: Final[list[tuple[str, object, object]]] = [] # mutable-ok: one entry per question
|
||||
|
||||
async def __call__(self, *, model: str, request_kwargs: Mapping[str, object], llm_router: object) -> bool:
|
||||
metadata: Final = request_kwargs.get("metadata")
|
||||
self.asked.append(
|
||||
(
|
||||
model,
|
||||
request_kwargs.get("model"),
|
||||
metadata.get("user_api_key_hash") if isinstance(metadata, Mapping) else None,
|
||||
)
|
||||
)
|
||||
return model not in self.denied
|
||||
|
||||
|
||||
def _arguments(access: _Check, budget: _Check, h_answers: bool = False) -> Mapping[str, object]:
|
||||
return {
|
||||
"model_list": (
|
||||
_deployment("g", "a", Exception("primary broke")),
|
||||
_deployment("h", "b", None if h_answers else Exception("h broke")),
|
||||
_deployment("i", "c"),
|
||||
_deployment("j", "d"),
|
||||
),
|
||||
"fallbacks": [{"g": ["h", "i", "j"]}],
|
||||
"num_retries": 0,
|
||||
"fallback_access_check": access,
|
||||
"fallback_budget_check": budget,
|
||||
}
|
||||
|
||||
|
||||
Observed = tuple[
|
||||
str, Mapping[str, object], tuple[tuple[str, object, object], ...], tuple[tuple[str, object, object], ...]
|
||||
]
|
||||
|
||||
|
||||
async def _observe(
|
||||
build: Callable[[Mapping[str, object], int], Backend], denied: frozenset[str], over: frozenset[str]
|
||||
) -> Observed:
|
||||
access: Final = _Check(denied)
|
||||
budget: Final = _Check(over)
|
||||
response: Final = await build(_arguments(access, budget), 1).acompletion(
|
||||
"g", MESSAGES, metadata={"user_api_key_hash": "caller"}
|
||||
)
|
||||
content: Final = str(response.choices[0].message.content) # pyright: ignore[reportAttributeAccessIssue,reportUnknownMemberType,reportUnknownArgumentType] # ModelResponse
|
||||
return content, router_headers(response), tuple(access.asked), tuple(budget.asked)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("denied", "over", "expected"),
|
||||
(
|
||||
pytest.param(frozenset(), frozenset(), "from c", id="nothing-rejected"),
|
||||
pytest.param(frozenset({"h"}), frozenset(), "from c", id="an-unauthorized-target-is-skipped"),
|
||||
pytest.param(frozenset(), frozenset({"h", "i"}), "from d", id="over-budget-targets-are-skipped"),
|
||||
),
|
||||
)
|
||||
async def test_both_backends_put_each_fallback_target_to_the_proxys_checks(
|
||||
denied: frozenset[str], over: frozenset[str], expected: str
|
||||
) -> None:
|
||||
observed: Final = [await _observe(build, denied, over) for build in BACKENDS]
|
||||
|
||||
assert observed[1] == observed[0]
|
||||
assert observed[1][0] == expected
|
||||
|
||||
|
||||
async def test_the_checks_read_the_request_as_the_failed_hop_left_it() -> None:
|
||||
_, _, access, _ = await _observe(BACKENDS[1], frozenset(), frozenset())
|
||||
|
||||
assert access == (("h", "g", "caller"), ("i", "h", "caller"))
|
||||
|
||||
|
||||
@pytest.mark.parametrize("build", BACKENDS)
|
||||
def test_the_sync_path_runs_the_checks_too(build: Callable[[Mapping[str, object], int], Backend]) -> None:
|
||||
access: Final = _Check(frozenset({"h"}))
|
||||
budget: Final = _Check(frozenset({"h"}))
|
||||
|
||||
response: Final = build(_arguments(access, budget, h_answers=True), 1).completion("g", MESSAGES)
|
||||
|
||||
assert response.choices[0].message.content == "from c" # pyright: ignore[reportAttributeAccessIssue,reportUnknownMemberType] # ModelResponse
|
||||
assert [model for model, _, _ in access.asked] == ["h", "i"]
|
||||
assert [model for model, _, _ in budget.asked] == ["i"]
|
||||
Loading…
Add table
Reference in a new issue