mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-13 08:52:02 +00:00
chore(alerts): reconcile activity views with main
This commit is contained in:
commit
fb0c3e1ec5
843 changed files with 49716 additions and 14134 deletions
10
.github/workflows/test-unit.yml
vendored
10
.github/workflows/test-unit.yml
vendored
|
|
@ -1035,9 +1035,9 @@ jobs:
|
|||
run: |
|
||||
diff -u model_prices_and_context_window.json litellm/model_prices_and_context_window_backup.json
|
||||
if [ -z "$RUST_BRIDGE_ARTIFACT" ]; then
|
||||
.github/scripts/uv_sync_with_retries.sh --frozen --group ci --group proxy-dev --extra google --extra proxy --extra semantic-router --extra saml --extra extra_proxy --extra proxy-runtime --extra utils
|
||||
.github/scripts/uv_sync_with_retries.sh --frozen --group ci --group proxy-dev --extra google --extra proxy --extra semantic-router --extra switchyard --extra saml --extra extra_proxy --extra proxy-runtime --extra utils
|
||||
else
|
||||
.github/scripts/uv_sync_with_retries.sh --frozen --group ci --group proxy-dev --extra google --extra proxy --extra semantic-router --extra saml --extra extra_proxy --extra proxy-runtime --extra utils --no-install-project --no-install-package litellm-proxy
|
||||
.github/scripts/uv_sync_with_retries.sh --frozen --group ci --group proxy-dev --extra google --extra proxy --extra semantic-router --extra switchyard --extra saml --extra extra_proxy --extra proxy-runtime --extra utils --no-install-project --no-install-package litellm-proxy
|
||||
uv pip install --no-deps --python .venv/bin/python rust-bridge-dist/*.whl
|
||||
cp rust-bridge-dist/litellm/rust_bridge/_native.abi3.so litellm/rust_bridge/_native.abi3.so
|
||||
uv run --no-sync python -c "import importlib.metadata; import litellm.rust_bridge._native; print(importlib.metadata.version('litellm'))"
|
||||
|
|
@ -1294,7 +1294,7 @@ jobs:
|
|||
- name: Install dependencies
|
||||
if: steps.changes.outputs.decision != 'skip'
|
||||
run: |
|
||||
.github/scripts/uv_sync_with_retries.sh --frozen --group ci --group proxy-dev --extra google --extra proxy --extra semantic-router
|
||||
.github/scripts/uv_sync_with_retries.sh --frozen --group ci --group proxy-dev --extra google --extra proxy --extra semantic-router --extra switchyard
|
||||
|
||||
- name: Cache Prisma binaries
|
||||
if: steps.changes.outputs.decision != 'skip'
|
||||
|
|
@ -1497,9 +1497,9 @@ jobs:
|
|||
RUST_BRIDGE_ARTIFACT: ${{ needs.rust-bridge.outputs.artifact }}
|
||||
run: |
|
||||
if [ -z "$RUST_BRIDGE_ARTIFACT" ]; then
|
||||
.github/scripts/uv_sync_with_retries.sh --frozen --group ci --group proxy-dev --extra google --extra proxy --extra semantic-router --extra saml --extra extra_proxy --extra proxy-runtime --extra utils
|
||||
.github/scripts/uv_sync_with_retries.sh --frozen --group ci --group proxy-dev --extra google --extra proxy --extra semantic-router --extra switchyard --extra saml --extra extra_proxy --extra proxy-runtime --extra utils
|
||||
else
|
||||
.github/scripts/uv_sync_with_retries.sh --frozen --group ci --group proxy-dev --extra google --extra proxy --extra semantic-router --extra saml --extra extra_proxy --extra proxy-runtime --extra utils --no-install-project --no-install-package litellm-proxy
|
||||
.github/scripts/uv_sync_with_retries.sh --frozen --group ci --group proxy-dev --extra google --extra proxy --extra semantic-router --extra switchyard --extra saml --extra extra_proxy --extra proxy-runtime --extra utils --no-install-project --no-install-package litellm-proxy
|
||||
uv pip install --no-deps --python .venv/bin/python rust-bridge-dist/*.whl
|
||||
cp rust-bridge-dist/litellm/rust_bridge/_native.abi3.so litellm/rust_bridge/_native.abi3.so
|
||||
fi
|
||||
|
|
|
|||
|
|
@ -52,7 +52,6 @@ If you ever write any human-facing text (pull requests, issues, commit messages,
|
|||
- don't use "—". Don't use ";" or ":" unless it would be nonsensical not to. Instead, reach for ",", ".", conjunction words, etc. in descending order of preference: vary among them, weighted toward the front of the list, and skip "," where it would cause a comma splice or the sentence is getting long. Overusing any one of them also feels AI-y. A word cap does not penalize you for adding more sentences
|
||||
- don't use the pattern "It's not X, it's Y", "You're not X, you're Y", etc.
|
||||
- unless explicitly asked, don't use bulleted or numbered lists unless it would be nonsensical not to. Instead, prefer prose
|
||||
- don't add a trailing "." at the end of paragraphs (just like this file). That means every paragraph, not just the last one (of the markdown file, PR description, GitHub comment, etc.). Rule of thumb: if you're adding new line(s) before the next sentence, don't add a ".". Exception is that docs (i.e., litellm-docs) should have trailing "."
|
||||
- don't use →. Instead, prefer not to use arrows, and if need be, use -> instead
|
||||
- use plain, simple, everyday engineering language: the common phrase engineers actually say over rare compact phrasing, in grammatically complete sentences. When explicitly asked to use bullets or ordered lists and structure legitimately helps the reader, prefer nested bullets (any depth is fine) over dense lines in a flat structure
|
||||
|
||||
|
|
|
|||
|
|
@ -46,7 +46,43 @@ For a managed workspace rollout, use OpenAI's [gateway configuration guide](http
|
|||
|
||||
## Native compatibility and Pi
|
||||
|
||||
For a native Codex model catalog and request format, change Codex's `base_url` to `https://litellm.example.com/backend-api/codex`. This route uses the same credentials but does not need a model deployment or resolve model aliases. Restricted keys must allow the upstream model name.
|
||||
For a native Codex model catalog and request format, change Codex's `base_url` to `https://litellm.example.com/backend-api/codex`. This route uses the same credentials. A ChatGPT model keeps its upstream name and needs no deployment of its own.
|
||||
|
||||
### Mix ChatGPT and LiteLLM models
|
||||
|
||||
On the native route, Codex's model picker shows OpenAI's catalog first and then the other chat and responses models the virtual key can use. Pick one of those and Codex's requests go through LiteLLM's `/v1/responses` with the key's usual routing, budgets, guardrails, and logging. The proxy removes the ChatGPT bearer and account header from that request, so the other provider and your callbacks never receive them. A name is routed this way when another provider serves it as an exact `model_name`, through a router `model_group_alias`, or as a team's public model name. A ChatGPT model, a name that only a wildcard deployment matches, or a name OpenAI's catalog lists for the caller's ChatGPT account keeps the passthrough, so a model OpenAI adds later still uses the subscription even when a deployment of another provider has the same name. The proxy learns the account's list from the catalog Codex fetches, or asks OpenAI with the caller's login. When it cannot get the list, for example because OpenAI is down, the request times out, or the client is not Codex, every request keeps the passthrough for the next minute, as it did before LiteLLM models were added to the picker. When the key can use many models, the picker lists the ones that fit in Codex's catalog size limit, and the proxy logs a warning with how many it left out.
|
||||
|
||||
Codex's picker sends a `none` reasoning effort for a LiteLLM model whose catalog entry lists no effort levels. With `litellm_settings: drop_params: true`, the proxy drops it for a model that does not support `none`, and the provider applies its own default. A model whose cost map entry or `model_info` sets `supports_none_reasoning_effort: true` keeps it. Without `drop_params`, the effort reaches the provider unchanged.
|
||||
|
||||
Codex also sends its ChatGPT app tools with every turn, which on a LiteLLM model adds a few hundred thousand prompt tokens to each turn. Set `apps = false` under `[features]` in Codex's `config.toml` to turn them off.
|
||||
|
||||
Only a key that may call `/v1/responses` gets these models. A key whose `allowed_routes` covers only the native routes sees OpenAI's catalog unchanged, and every request it sends uses the subscription.
|
||||
|
||||
Codex sends earlier reasoning items back with each turn. The proxy tags the reasoning that a LiteLLM model returns on this route, OpenAI and Azure deployments included. When a session switches to a ChatGPT model, the proxy drops those items before forwarding the request to OpenAI, which would reject them. Messages, tool calls, and ChatGPT's own reasoning are kept, and a LiteLLM model still gets its own reasoning back.
|
||||
|
||||
When the catalog adds LiteLLM models for a key, its `ETag` is a digest of OpenAI's ETag and those models, and that key's responses carry the same value in `X-Models-Etag`. Codex keeps its cached catalog while the two match and fetches it again when OpenAI's catalog or the key's models change. A key whose catalog adds no LiteLLM models keeps OpenAI's ETags.
|
||||
|
||||
### Restrict native models with an access group
|
||||
|
||||
A key with a model allowlist needs a rule for ChatGPT names that OpenAI adds over time. Put the subscription wildcards in an access group under `chatgpt/` names:
|
||||
|
||||
```yaml
|
||||
model_list:
|
||||
- model_name: chatgpt/gpt-*
|
||||
litellm_params:
|
||||
model: chatgpt/gpt-*
|
||||
model_info:
|
||||
access_groups: ["chatgpt-native"]
|
||||
- model_name: chatgpt/codex-*
|
||||
litellm_params:
|
||||
model: chatgpt/codex-*
|
||||
model_info:
|
||||
access_groups: ["chatgpt-native"]
|
||||
```
|
||||
|
||||
Then add `chatgpt-native` to the key's models, next to any other model names it should reach. On the native route, a key that holds only the `chatgpt/` form of a name is checked as that form, so a request for `gpt-5.6-sol` is allowed when the key may call `chatgpt/gpt-5.6-sol`. The key can call every ChatGPT model that matches the wildcards, including ones the proxy configuration does not name.
|
||||
|
||||
On `/v1` these wildcards match only a request that asks for a `chatgpt/` name, so a request for a plain `gpt-` or `codex-` name still goes where it went before, such as a `*` catch-all deployment. A `/v1` request for a `chatgpt/` name without a ChatGPT login falls back to the proxy's own saved login if it has one, so keep that login unset as described above.
|
||||
|
||||
### Configure Pi
|
||||
|
||||
|
|
@ -71,7 +107,7 @@ Keep Pi's subscription login as the provider credential. The header value names
|
|||
|
||||
Send a prompt and check the completed request in LiteLLM's Logs page. Its gateway user/key and token counts should match the caller. Give each person their own LiteLLM virtual key for separate attribution. The gateway credential and ChatGPT login are independently authenticated; LiteLLM does not establish that both identities belong to the same person.
|
||||
|
||||
The native routes are `GET /backend-api/codex/models` and `POST /backend-api/codex/responses`. Model requests use the upstream model name, so a restricted LiteLLM key must allow that name. OpenAI's model catalog is returned in its native format; seeing a model there does not override gateway permissions.
|
||||
The native routes are `GET /backend-api/codex/models` and `POST /backend-api/codex/responses`. Model requests use the upstream model name, so a restricted LiteLLM key needs the `chatgpt-native` access group above or that exact name. OpenAI's model catalog is returned in its native format; seeing a model there does not override gateway permissions.
|
||||
|
||||
The proxy forwards the ChatGPT bearer and account header to the fixed OpenAI Codex endpoint. It removes gateway credentials and cookies from the upstream request and redacts credential headers in logging callbacks. The client remains responsible for refresh and reauthentication. Native passthrough returns authentication and quota errors without switching subscriptions. Standard model routes retain the caller login across ChatGPT retries; explicitly configured cross-provider fallbacks use their own credentials. An explicit workspace header is preserved for OpenAI to validate.
|
||||
|
||||
|
|
|
|||
|
|
@ -1,6 +1,6 @@
|
|||
[project]
|
||||
name = "litellm-enterprise"
|
||||
version = "0.1.81"
|
||||
version = "0.1.83"
|
||||
description = "Package for LiteLLM Enterprise features"
|
||||
readme = "README.md"
|
||||
requires-python = ">=3.9"
|
||||
|
|
@ -26,7 +26,7 @@ required-version = ">=0.10.9"
|
|||
module-root = ""
|
||||
|
||||
[tool.commitizen]
|
||||
version = "0.1.81"
|
||||
version = "0.1.83"
|
||||
version_files = [
|
||||
"pyproject.toml:^version",
|
||||
"../pyproject.toml:litellm-enterprise==",
|
||||
|
|
|
|||
|
|
@ -1,6 +1,6 @@
|
|||
[project]
|
||||
name = "litellm-proxy-extras"
|
||||
version = "0.4.112"
|
||||
version = "0.4.114"
|
||||
description = "Additional files for the LiteLLM Proxy. Reduces the size of the main litellm package."
|
||||
readme = "README.md"
|
||||
requires-python = ">=3.9"
|
||||
|
|
@ -30,7 +30,7 @@ required-version = ">=0.10.9"
|
|||
module-root = ""
|
||||
|
||||
[tool.commitizen]
|
||||
version = "0.4.112"
|
||||
version = "0.4.114"
|
||||
version_files = [
|
||||
"pyproject.toml:^version",
|
||||
"../pyproject.toml:litellm-proxy-extras==",
|
||||
|
|
|
|||
2
litellm-rust/Cargo.lock
generated
2
litellm-rust/Cargo.lock
generated
|
|
@ -4381,6 +4381,7 @@ dependencies = [
|
|||
"criterion",
|
||||
"fancy-regex 0.19.2",
|
||||
"futures-util",
|
||||
"indexmap 2.14.0",
|
||||
"litellm-auth",
|
||||
"litellm-auth-aws",
|
||||
"litellm-cache",
|
||||
|
|
@ -5773,6 +5774,7 @@ version = "0.29.2"
|
|||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "4688ddedf473e32662b9b067670129a8afb8c18e351482c70d62ba4a88171e8b"
|
||||
dependencies = [
|
||||
"indexmap 2.14.0",
|
||||
"libc",
|
||||
"once_cell",
|
||||
"portable-atomic",
|
||||
|
|
|
|||
|
|
@ -60,7 +60,8 @@ litellm-router-core.workspace = true
|
|||
litellm-host-python.workspace = true
|
||||
litellm-model-catalog.workspace = true
|
||||
litellm-token-counter = { path = "../token-counter", default-features = false }
|
||||
pyo3.workspace = true
|
||||
indexmap = "2.14.0"
|
||||
pyo3 = { workspace = true, features = ["indexmap"] }
|
||||
prost.workspace = true
|
||||
pyo3-async-runtimes.workspace = true
|
||||
reqwest.workspace = true
|
||||
|
|
|
|||
|
|
@ -42,9 +42,24 @@ mod _native {
|
|||
#[pymodule_export]
|
||||
use crate::redis::client::native_redis_breaker_states;
|
||||
#[pymodule_export]
|
||||
use crate::redis::store::auth_cache::{NativeAuthCache, PyLocalSizes};
|
||||
#[pymodule_export]
|
||||
use crate::redis::store::rate_limits::{
|
||||
NativeRateLimits, PyAtomicApplied, PyAtomicOverLimit, PyGaugeOverLimit, PyGaugesAcquired,
|
||||
PyGuardedReply, PyWindowCount,
|
||||
};
|
||||
#[pymodule_export]
|
||||
use crate::redis::store::request::{NativeRequestRedis, NativeSend};
|
||||
#[pymodule_export]
|
||||
use crate::redis::store::{NativePending, NativeRedisStoreError};
|
||||
use crate::redis::store::router_usage::NativeRouterUsage;
|
||||
#[pymodule_export]
|
||||
use crate::redis::store::sensitive_route_pins::NativeSensitiveRoutePins;
|
||||
#[pymodule_export]
|
||||
use crate::redis::store::spend_counters::{NativeSpendCounters, PySpendRead};
|
||||
#[pymodule_export]
|
||||
use crate::redis::store::{
|
||||
NativePending, NativeRedisStoreError, router_cooldowns::NativeRouterCooldowns,
|
||||
};
|
||||
#[pymodule_export]
|
||||
use crate::router::{RouterCall, RouterCoreHandle};
|
||||
#[pymodule_export]
|
||||
|
|
@ -123,8 +138,22 @@ mod tests {
|
|||
"NativeRedisClient",
|
||||
"NativeRequestRedis",
|
||||
"NativeSend",
|
||||
"NativeAuthCache",
|
||||
"NativeRateLimits",
|
||||
"WindowCount",
|
||||
"AtomicApplied",
|
||||
"AtomicOverLimit",
|
||||
"GuardedReply",
|
||||
"GaugesAcquired",
|
||||
"GaugeOverLimit",
|
||||
"NativePending",
|
||||
"NativeRouterCooldowns",
|
||||
"NativeRouterUsage",
|
||||
"NativeSensitiveRoutePins",
|
||||
"NativeRedisStoreError",
|
||||
"NativeSpendCounters",
|
||||
"SpendRead",
|
||||
"LocalSizes",
|
||||
"native_redis_breaker_states",
|
||||
"NativeClickHouseSpendConfig",
|
||||
"NativeClickHouseSpendStorage",
|
||||
|
|
|
|||
|
|
@ -13,6 +13,8 @@ pub(crate) enum PythonSettings {
|
|||
UrlPolicy,
|
||||
#[strum(serialize = "provider_defaults")]
|
||||
ProviderDefaults,
|
||||
#[strum(serialize = "router_tiers")]
|
||||
RouterTiers,
|
||||
#[strum(serialize = "secret_manager")]
|
||||
SecretManager,
|
||||
#[strum(serialize = "secret_manager_binding")]
|
||||
|
|
|
|||
|
|
@ -18,7 +18,7 @@ use pyo3::{
|
|||
exceptions::{PyOverflowError, PyTypeError, PyValueError},
|
||||
prelude::*,
|
||||
sync::MutexExt,
|
||||
types::{PyBool, PyDict, PyFloat, PyInt, PyList, PyString, PyTuple},
|
||||
types::{PyBool, PyDict, PyFloat, PyInt, PyList, PyMapping, PyString, PyTuple},
|
||||
};
|
||||
|
||||
use self::credentials::PythonCredentialSource;
|
||||
|
|
@ -40,16 +40,18 @@ impl<'a, 'py> FromPyObject<'a, 'py> for RedisSpec {
|
|||
|
||||
fn extract(kwargs: Borrowed<'a, 'py, PyAny>) -> PyResult<Self> {
|
||||
let py = kwargs.py();
|
||||
let source = PyDict::new(py);
|
||||
let mapping = kwargs.cast::<PyMapping>()?;
|
||||
source.update(&mapping)?;
|
||||
let plain = PyDict::new(py);
|
||||
let mut host_objects = Vec::new();
|
||||
for (name, value) in kwargs.cast::<PyDict>()?.iter() {
|
||||
for (name, value) in source.iter() {
|
||||
match value.is_none() || is_plain(&value) {
|
||||
true => plain.set_item(&name, value)?,
|
||||
false => host_objects.push(name.extract::<String>()?),
|
||||
}
|
||||
}
|
||||
let credential_provider = kwargs
|
||||
.cast::<PyDict>()?
|
||||
let credential_provider = source
|
||||
.get_item("credential_provider")?
|
||||
.filter(|provider| !provider.is_none() && !is_plain(provider))
|
||||
.map(Bound::unbind);
|
||||
|
|
@ -65,27 +67,44 @@ impl<'a, 'py> FromPyObject<'a, 'py> for RedisSpec {
|
|||
}
|
||||
}
|
||||
|
||||
/// `litellm.constants`' `REDIS_CIRCUIT_BREAKER_*`, which Python reads from the environment once.
|
||||
pub(crate) struct RedisBreaker(BreakerConfig);
|
||||
/// `litellm.redis.store_settings.RedisBreakerSettings`, which `RedisCache` fills from `litellm.constants`.
|
||||
#[derive(FromPyObject)]
|
||||
pub(crate) struct RedisBreaker {
|
||||
#[pyo3(attribute)]
|
||||
failure_threshold: u32,
|
||||
#[pyo3(attribute)]
|
||||
recovery_timeout: f64,
|
||||
#[pyo3(attribute)]
|
||||
timeout_min_duration: f64,
|
||||
}
|
||||
|
||||
impl<'a, 'py> FromPyObject<'a, 'py> for RedisBreaker {
|
||||
type Error = PyErr;
|
||||
|
||||
fn extract(breaker: Borrowed<'a, 'py, PyAny>) -> PyResult<Self> {
|
||||
let (failure_threshold, recovery_timeout, timeout_min_duration): (u32, f64, f64) =
|
||||
breaker.extract()?;
|
||||
impl RedisBreaker {
|
||||
fn config(&self) -> PyResult<BreakerConfig> {
|
||||
let seconds = |seconds: f64| {
|
||||
Duration::try_from_secs_f64(seconds)
|
||||
.map_err(|error| pyo3::exceptions::PyValueError::new_err(error.to_string()))
|
||||
.map_err(|error| PyValueError::new_err(error.to_string()))
|
||||
};
|
||||
Ok(Self(BreakerConfig {
|
||||
failure_threshold,
|
||||
recovery_timeout: seconds(recovery_timeout)?,
|
||||
timeout_min_duration: seconds(timeout_min_duration)?,
|
||||
}))
|
||||
Ok(BreakerConfig {
|
||||
failure_threshold: self.failure_threshold,
|
||||
recovery_timeout: seconds(self.recovery_timeout)?,
|
||||
timeout_min_duration: seconds(self.timeout_min_duration)?,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
/// `litellm.redis.store_settings.RedisStoreSettings`, which a `RedisCache` captures when it is built.
|
||||
#[derive(FromPyObject)]
|
||||
pub(crate) struct StoreSettings {
|
||||
#[pyo3(attribute)]
|
||||
redis_kwargs: RedisSpec,
|
||||
#[pyo3(attribute)]
|
||||
namespace: Option<String>,
|
||||
#[pyo3(attribute)]
|
||||
default_ttl: f64,
|
||||
#[pyo3(attribute)]
|
||||
breaker: Option<RedisBreaker>,
|
||||
}
|
||||
|
||||
fn is_plain(value: &Bound<'_, PyAny>) -> bool {
|
||||
value.is_instance_of::<PyString>()
|
||||
|| value.is_instance_of::<PyBool>()
|
||||
|
|
@ -119,7 +138,7 @@ fn shared_client(
|
|||
py: Python<'_>,
|
||||
spec: RedisSpec,
|
||||
namespace: Option<String>,
|
||||
breaker: Option<RedisBreaker>,
|
||||
breaker: Option<BreakerConfig>,
|
||||
) -> PyResult<Option<RedisClient>> {
|
||||
let environment: SharedEnvironment = Arc::new(ProcessEnvironment);
|
||||
let RedisSpec {
|
||||
|
|
@ -176,7 +195,7 @@ fn shared_client(
|
|||
}
|
||||
let connected = {
|
||||
let _runtime = runtime.enter();
|
||||
litellm_redis_connect::connect(resolved.clone(), host, breaker.map(|breaker| breaker.0))
|
||||
litellm_redis_connect::connect(resolved.clone(), host, breaker)
|
||||
};
|
||||
let client = match connected {
|
||||
Ok(client) => client,
|
||||
|
|
@ -205,15 +224,15 @@ pub(crate) struct NativeRedisClient {
|
|||
impl NativeRedisClient {
|
||||
/// `None` when redis-py would fail only once a command connects, which leaves every store given it on memory.
|
||||
#[staticmethod]
|
||||
#[pyo3(signature = (redis_kwargs, *, default_ttl_seconds, namespace = None, breaker = None))]
|
||||
fn connect(
|
||||
py: Python<'_>,
|
||||
redis_kwargs: RedisSpec,
|
||||
default_ttl_seconds: f64,
|
||||
namespace: Option<String>,
|
||||
breaker: Option<RedisBreaker>,
|
||||
) -> PyResult<Option<Self>> {
|
||||
let default_ttl = number::seconds(default_ttl_seconds).map_err(seconds_error)?;
|
||||
fn connect(py: Python<'_>, settings: StoreSettings) -> PyResult<Option<Self>> {
|
||||
let StoreSettings {
|
||||
redis_kwargs,
|
||||
namespace,
|
||||
default_ttl,
|
||||
breaker,
|
||||
} = settings;
|
||||
let default_ttl = number::seconds(default_ttl).map_err(seconds_error)?;
|
||||
let breaker = breaker.as_ref().map(RedisBreaker::config).transpose()?;
|
||||
Ok(
|
||||
shared_client(py, redis_kwargs, namespace, breaker)?.map(|client| Self {
|
||||
client,
|
||||
|
|
|
|||
|
|
@ -1,11 +1,3 @@
|
|||
#[expect(
|
||||
dead_code,
|
||||
reason = "the native stores that use these helpers land in the next PR"
|
||||
)]
|
||||
pub(crate) mod client;
|
||||
mod events;
|
||||
#[expect(
|
||||
dead_code,
|
||||
reason = "the native stores that use these helpers land in the next PR"
|
||||
)]
|
||||
pub(crate) mod store;
|
||||
|
|
|
|||
265
litellm-rust/crates/python-bridge/src/redis/store/auth_cache.rs
Normal file
265
litellm-rust/crates/python-bridge/src/redis/store/auth_cache.rs
Normal file
|
|
@ -0,0 +1,265 @@
|
|||
use std::sync::Arc;
|
||||
|
||||
use litellm_python_compat::Value;
|
||||
use litellm_redis_store::auth_cache::{AuthCache, Settings};
|
||||
use pyo3::prelude::*;
|
||||
|
||||
use super::{
|
||||
NativePending, PyNumber, kind_error, pending, redis_async, redis_sync, remote,
|
||||
request::{NativeRequestRedis, NativeSend},
|
||||
ttl,
|
||||
value::PyValue,
|
||||
};
|
||||
use crate::redis::client::NativeRedisClient;
|
||||
|
||||
fn values(values: Vec<Option<Value>>) -> Vec<Option<PyValue>> {
|
||||
values.into_iter().map(|value| value.map(PyValue)).collect()
|
||||
}
|
||||
|
||||
/// `UserApiKeyCache` with `enable_redis_auth_cache` on. The sync methods back `DualCache`'s sync
|
||||
/// API and block with the GIL released, like redis-py's sync client.
|
||||
#[pyclass(frozen, name = "NativeAuthCache")]
|
||||
pub(crate) struct NativeAuthCache(Arc<AuthCache>);
|
||||
|
||||
#[pymethods]
|
||||
impl NativeAuthCache {
|
||||
#[new]
|
||||
#[pyo3(signature = (client, *, memory_entries, memory_ttl_seconds = None))]
|
||||
fn new(
|
||||
client: Option<PyRef<'_, NativeRedisClient>>,
|
||||
memory_entries: usize,
|
||||
memory_ttl_seconds: Option<f64>,
|
||||
) -> PyResult<Self> {
|
||||
Ok(Self(Arc::new(AuthCache::new(
|
||||
client.map(|client| remote(&client)),
|
||||
Settings {
|
||||
memory_entries,
|
||||
memory_ttl: ttl(memory_ttl_seconds)?,
|
||||
},
|
||||
))))
|
||||
}
|
||||
|
||||
#[getter]
|
||||
fn has_redis(&self) -> bool {
|
||||
self.0.has_redis()
|
||||
}
|
||||
|
||||
/// `request` answers a key its pre-call reads already found absent without another `GET`.
|
||||
#[pyo3(signature = (key, request = None))]
|
||||
fn get<'py>(
|
||||
&self,
|
||||
py: Python<'py>,
|
||||
key: String,
|
||||
request: Option<&NativeRequestRedis>,
|
||||
) -> PyResult<Bound<'py, PyAny>> {
|
||||
let store = self.0.clone();
|
||||
let request = request.map(|request| request.0.clone());
|
||||
redis_async(py, "async_get_cache", kind_error, async move {
|
||||
store
|
||||
.get(request.as_ref(), &key)
|
||||
.await
|
||||
.map(|value| value.map(PyValue))
|
||||
})
|
||||
}
|
||||
|
||||
fn get_sync(&self, py: Python<'_>, key: String) -> PyResult<Option<PyValue>> {
|
||||
let store = self.0.clone();
|
||||
redis_sync(py, "get_cache", kind_error, async move {
|
||||
store.get(None, &key).await.map(|value| value.map(PyValue))
|
||||
})
|
||||
}
|
||||
|
||||
/// `update_cache_ttl` and `update_in_memory_max_size`: both partitions keep what they hold.
|
||||
#[pyo3(signature = (*, memory_entries, memory_ttl_seconds = None))]
|
||||
fn reconfigure(&self, memory_entries: usize, memory_ttl_seconds: Option<f64>) -> PyResult<()> {
|
||||
self.0.reconfigure(Settings {
|
||||
memory_entries,
|
||||
memory_ttl: ttl(memory_ttl_seconds)?,
|
||||
});
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[pyo3(signature = (key, value, ttl_seconds = None))]
|
||||
fn set<'py>(
|
||||
&self,
|
||||
py: Python<'py>,
|
||||
key: String,
|
||||
value: PyValue,
|
||||
ttl_seconds: Option<f64>,
|
||||
) -> PyResult<Bound<'py, PyAny>> {
|
||||
let store = self.0.clone();
|
||||
let ttl = ttl(ttl_seconds)?;
|
||||
redis_async(py, "async_set_cache", kind_error, async move {
|
||||
store.set(&key, value.0, ttl).await
|
||||
})
|
||||
}
|
||||
|
||||
#[pyo3(signature = (key, value, ttl_seconds = None))]
|
||||
fn set_sync(
|
||||
&self,
|
||||
py: Python<'_>,
|
||||
key: String,
|
||||
value: PyValue,
|
||||
ttl_seconds: Option<f64>,
|
||||
) -> PyResult<()> {
|
||||
let store = self.0.clone();
|
||||
let ttl = ttl(ttl_seconds)?;
|
||||
redis_sync(py, "set_cache", kind_error, async move {
|
||||
store.set_str(&key, value.0, ttl).await
|
||||
})
|
||||
}
|
||||
|
||||
/// Memory now and the `SET` on the batch `send` picks.
|
||||
#[pyo3(signature = (key, value, ttl_seconds, send))]
|
||||
fn declare_set(
|
||||
&self,
|
||||
key: String,
|
||||
value: PyValue,
|
||||
ttl_seconds: Option<f64>,
|
||||
send: &NativeSend,
|
||||
) -> PyResult<NativePending> {
|
||||
let written = self
|
||||
.0
|
||||
.declare_set(send.send(), &key, value.0, ttl(ttl_seconds)?)
|
||||
.map_err(kind_error)?;
|
||||
Ok(pending(
|
||||
"async_set_cache",
|
||||
kind_error,
|
||||
written,
|
||||
std::convert::identity,
|
||||
))
|
||||
}
|
||||
|
||||
fn declare_delete(&self, key: String, send: &NativeSend) -> PyResult<NativePending> {
|
||||
let deleted = self
|
||||
.0
|
||||
.declare_delete(send.send(), &key)
|
||||
.map_err(kind_error)?;
|
||||
Ok(pending(
|
||||
"async_delete_cache",
|
||||
kind_error,
|
||||
deleted,
|
||||
std::convert::identity,
|
||||
))
|
||||
}
|
||||
|
||||
#[pyo3(signature = (entries, ttl_seconds = None))]
|
||||
fn set_many<'py>(
|
||||
&self,
|
||||
py: Python<'py>,
|
||||
entries: Vec<(String, PyValue)>,
|
||||
ttl_seconds: Option<f64>,
|
||||
) -> PyResult<Bound<'py, PyAny>> {
|
||||
let store = self.0.clone();
|
||||
let entries = entries
|
||||
.into_iter()
|
||||
.map(|(key, value)| (key, value.0))
|
||||
.collect();
|
||||
let ttl = ttl(ttl_seconds)?;
|
||||
redis_async(py, "async_set_cache_pipeline", kind_error, async move {
|
||||
store.set_many(entries, ttl).await
|
||||
})
|
||||
}
|
||||
|
||||
/// On a request batch it returns before anything is sent, and `Now` is a pipeline of its own.
|
||||
fn write_back<'py>(
|
||||
&self,
|
||||
py: Python<'py>,
|
||||
entries: Vec<(String, PyValue, Option<f64>)>,
|
||||
send: &NativeSend,
|
||||
) -> PyResult<Bound<'py, PyAny>> {
|
||||
let entries = entries
|
||||
.into_iter()
|
||||
.map(|(key, value, seconds)| Ok((key, value.0, ttl(seconds)?)))
|
||||
.collect::<PyResult<_>>()?;
|
||||
let written = self
|
||||
.0
|
||||
.write_back(send.send(), entries)
|
||||
.map_err(kind_error)?;
|
||||
redis_async(
|
||||
py,
|
||||
"async_set_cache_pipeline_with_ttls",
|
||||
kind_error,
|
||||
written,
|
||||
)
|
||||
}
|
||||
|
||||
fn delete<'py>(&self, py: Python<'py>, keys: Vec<String>) -> PyResult<Bound<'py, PyAny>> {
|
||||
let store = self.0.clone();
|
||||
redis_async(py, "async_delete_cache", kind_error, async move {
|
||||
store.delete(&keys).await
|
||||
})
|
||||
}
|
||||
|
||||
fn delete_sync(&self, py: Python<'_>, keys: Vec<String>) -> PyResult<()> {
|
||||
let store = self.0.clone();
|
||||
redis_sync(py, "delete_cache", kind_error, async move {
|
||||
store.delete(&keys).await
|
||||
})
|
||||
}
|
||||
|
||||
/// One `MGET` of keys the caller found missing in memory, on the batch `send` picks; `None`
|
||||
/// without Redis.
|
||||
fn get_many(&self, keys: Vec<String>, send: &NativeSend) -> Option<NativePending> {
|
||||
let read = self.0.get_many(send.send(), &keys)?;
|
||||
Some(pending("async_batch_get_cache", kind_error, read, values))
|
||||
}
|
||||
|
||||
fn increment_if_present<'py>(
|
||||
&self,
|
||||
py: Python<'py>,
|
||||
key: String,
|
||||
amount: PyNumber,
|
||||
) -> PyResult<Bound<'py, PyAny>> {
|
||||
let store = self.0.clone();
|
||||
redis_async(py, "async_increment_if_present", kind_error, async move {
|
||||
store.increment_if_present(&key, amount.0).await
|
||||
})
|
||||
}
|
||||
|
||||
fn get_local(&self, key: &str) -> PyResult<Option<PyValue>> {
|
||||
Ok(self.0.get_local(key).map_err(kind_error)?.map(PyValue))
|
||||
}
|
||||
|
||||
#[pyo3(signature = (key, value, ttl_seconds = None))]
|
||||
fn set_local(&self, key: &str, value: PyValue, ttl_seconds: Option<f64>) -> PyResult<()> {
|
||||
self.0
|
||||
.set_local(key, value.0, ttl(ttl_seconds)?)
|
||||
.map_err(kind_error)
|
||||
}
|
||||
|
||||
fn delete_local(&self, key: &str) -> PyResult<()> {
|
||||
self.0.delete_local(key).map_err(kind_error)
|
||||
}
|
||||
|
||||
fn flush_local(&self) -> PyResult<()> {
|
||||
self.0.flush_local().map_err(kind_error)
|
||||
}
|
||||
|
||||
fn len_local(&self) -> PyResult<PyLocalSizes> {
|
||||
let (objects, key_objects) = self.0.len_local().map_err(kind_error)?;
|
||||
Ok(PyLocalSizes {
|
||||
objects,
|
||||
key_objects,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
/// Entries the management-object and the user-key partitions hold in memory.
|
||||
#[pyclass(frozen, eq, get_all, name = "LocalSizes")]
|
||||
#[derive(PartialEq)]
|
||||
pub(crate) struct PyLocalSizes {
|
||||
objects: usize,
|
||||
key_objects: usize,
|
||||
}
|
||||
|
||||
#[pymethods]
|
||||
impl PyLocalSizes {
|
||||
#[new]
|
||||
fn new(objects: usize, key_objects: usize) -> Self {
|
||||
Self {
|
||||
objects,
|
||||
key_objects,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
|
@ -1,4 +1,11 @@
|
|||
pub(crate) mod auth_cache;
|
||||
mod policy;
|
||||
pub(crate) mod rate_limits;
|
||||
pub(crate) mod request;
|
||||
pub(crate) mod router_cooldowns;
|
||||
pub(crate) mod router_usage;
|
||||
pub(crate) mod sensitive_route_pins;
|
||||
pub(crate) mod spend_counters;
|
||||
mod value;
|
||||
|
||||
use value::PyValue;
|
||||
|
|
@ -159,6 +166,7 @@ where
|
|||
}
|
||||
|
||||
/// A Python `int` or `float` across the bridge, keeping which one it is.
|
||||
#[derive(Clone, Copy, PartialEq)]
|
||||
pub(crate) struct PyNumber(Number);
|
||||
|
||||
impl<'a, 'py> FromPyObject<'a, 'py> for PyNumber {
|
||||
|
|
@ -188,6 +196,21 @@ impl<'py> IntoPyObject<'py> for PyNumber {
|
|||
}
|
||||
}
|
||||
|
||||
/// A store enum read from its wire spelling, so the bridge names no variant.
|
||||
pub(crate) struct Spelled<T>(pub(crate) T);
|
||||
|
||||
impl<'a, 'py, T: std::str::FromStr> FromPyObject<'a, 'py> for Spelled<T> {
|
||||
type Error = PyErr;
|
||||
|
||||
fn extract(value: Borrowed<'a, 'py, PyAny>) -> PyResult<Self> {
|
||||
let spelling: String = value.extract()?;
|
||||
spelling
|
||||
.parse()
|
||||
.map(Self)
|
||||
.map_err(|_| PyValueError::new_err(format!("unknown spelling {spelling:?}")))
|
||||
}
|
||||
}
|
||||
|
||||
/// A store type whose Python side is plain data, converted through its `FromValue` and `IntoValue`
|
||||
/// so the bridge never names a field. A shape the store refuses is `ValueError`.
|
||||
pub(crate) struct PyData<T>(pub(crate) T);
|
||||
|
|
|
|||
151
litellm-rust/crates/python-bridge/src/redis/store/policy.rs
Normal file
151
litellm-rust/crates/python-bridge/src/redis/store/policy.rs
Normal file
|
|
@ -0,0 +1,151 @@
|
|||
use std::time::Duration;
|
||||
|
||||
use litellm_redis_store::TierPolicy;
|
||||
use pyo3::prelude::*;
|
||||
|
||||
use crate::{
|
||||
coercion::{Field, FieldSpec, ProjectionError},
|
||||
python_settings::{PythonSettings, Snapshot},
|
||||
};
|
||||
|
||||
const COOLDOWN_REFRESH: FieldSpec<Duration> =
|
||||
FieldSpec::new("cooldown_redis_read_interval", seconds);
|
||||
const USAGE_REFRESH: FieldSpec<Duration> = FieldSpec::new("usage_redis_read_interval", seconds);
|
||||
const TRACKED_READS: FieldSpec<usize> = FieldSpec::new("max_redis_batch_cache_size", count);
|
||||
const MEMORY_TTL: FieldSpec<Option<Duration>> =
|
||||
FieldSpec::new("default_in_memory_ttl", optional_seconds);
|
||||
|
||||
/// The tier values this Python process runs with, from its environment and `litellm` globals.
|
||||
#[derive(Debug)]
|
||||
pub(crate) struct RouterTierPolicies {
|
||||
pub(crate) cooldowns: TierPolicy,
|
||||
pub(crate) usage: TierPolicy,
|
||||
}
|
||||
|
||||
pub(crate) fn router_tier_policies(py: Python<'_>) -> PyResult<RouterTierPolicies> {
|
||||
Ok(project(&PythonSettings::RouterTiers.read(py)?)?)
|
||||
}
|
||||
|
||||
fn project(snapshot: &Snapshot<'_>) -> Result<RouterTierPolicies, ProjectionError> {
|
||||
let tracked_reads = snapshot.read(&TRACKED_READS)?;
|
||||
let default_in_memory_ttl = snapshot.read(&MEMORY_TTL)?;
|
||||
Ok(RouterTierPolicies {
|
||||
cooldowns: TierPolicy {
|
||||
refresh_after: snapshot.read(&COOLDOWN_REFRESH)?,
|
||||
tracked_reads,
|
||||
default_in_memory_ttl,
|
||||
},
|
||||
usage: TierPolicy {
|
||||
refresh_after: snapshot.read(&USAGE_REFRESH)?,
|
||||
tracked_reads,
|
||||
default_in_memory_ttl,
|
||||
},
|
||||
})
|
||||
}
|
||||
|
||||
fn seconds(field: &Field<'_>) -> Result<Duration, ProjectionError> {
|
||||
match field.value().extract::<f64>() {
|
||||
Ok(seconds) if seconds.is_finite() && seconds >= 0.0 => {
|
||||
Ok(Duration::from_secs_f64(seconds))
|
||||
}
|
||||
_ => Err(field.invalid("a non-negative number of seconds")),
|
||||
}
|
||||
}
|
||||
|
||||
fn optional_seconds(field: &Field<'_>) -> Result<Option<Duration>, ProjectionError> {
|
||||
match field.value().is_none() {
|
||||
true => Ok(None),
|
||||
false => seconds(field).map(Some),
|
||||
}
|
||||
}
|
||||
|
||||
fn count(field: &Field<'_>) -> Result<usize, ProjectionError> {
|
||||
field
|
||||
.value()
|
||||
.extract::<usize>()
|
||||
.map_err(|_| field.invalid("a non-negative integer"))
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use std::time::Duration;
|
||||
|
||||
use pyo3::{prelude::*, types::PyDict};
|
||||
|
||||
use super::project;
|
||||
use crate::python_settings::PythonSettings;
|
||||
|
||||
fn projected(py: Python<'_>, fields: &str) -> Result<super::RouterTierPolicies, String> {
|
||||
let locals = PyDict::new(py);
|
||||
let source = std::ffi::CString::new(format!(
|
||||
"import types\nsettings = types.SimpleNamespace({fields})"
|
||||
))
|
||||
.unwrap();
|
||||
py.run(&source, Some(&locals), Some(&locals)).unwrap();
|
||||
let snapshot =
|
||||
PythonSettings::RouterTiers.snapshot(locals.get_item("settings").unwrap().unwrap());
|
||||
project(&snapshot).map_err(|error| PyErr::from(error).to_string())
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn the_host_values_are_each_familys_refresh_and_tracked_reads() {
|
||||
Python::initialize();
|
||||
Python::attach(|py| {
|
||||
let policies = projected(
|
||||
py,
|
||||
"cooldown_redis_read_interval=2.5, usage_redis_read_interval=30, max_redis_batch_cache_size=42, default_in_memory_ttl=None",
|
||||
)
|
||||
.unwrap();
|
||||
assert_eq!(
|
||||
policies.cooldowns.refresh_after,
|
||||
Duration::from_millis(2500)
|
||||
);
|
||||
assert_eq!(policies.usage.refresh_after, Duration::from_secs(30));
|
||||
assert_eq!(policies.cooldowns.tracked_reads, 42);
|
||||
assert_eq!(policies.usage.tracked_reads, 42);
|
||||
assert_eq!(policies.cooldowns.default_in_memory_ttl, None);
|
||||
});
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn a_set_default_in_memory_ttl_reaches_both_familys_backfill() {
|
||||
Python::initialize();
|
||||
Python::attach(|py| {
|
||||
let policies = projected(
|
||||
py,
|
||||
"cooldown_redis_read_interval=1, usage_redis_read_interval=10, max_redis_batch_cache_size=1000, default_in_memory_ttl=7.5",
|
||||
)
|
||||
.unwrap();
|
||||
assert_eq!(
|
||||
policies.cooldowns.default_in_memory_ttl,
|
||||
Some(Duration::from_millis(7500))
|
||||
);
|
||||
assert_eq!(
|
||||
policies.usage.default_in_memory_ttl,
|
||||
Some(Duration::from_millis(7500))
|
||||
);
|
||||
});
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn a_value_python_would_not_run_with_is_a_configuration_error() {
|
||||
Python::initialize();
|
||||
Python::attach(|py| {
|
||||
let negative = projected(
|
||||
py,
|
||||
"cooldown_redis_read_interval=-1, usage_redis_read_interval=10, max_redis_batch_cache_size=1000, default_in_memory_ttl=None",
|
||||
)
|
||||
.unwrap_err();
|
||||
let text = projected(
|
||||
py,
|
||||
"cooldown_redis_read_interval=1, usage_redis_read_interval=10, max_redis_batch_cache_size='many', default_in_memory_ttl=None",
|
||||
)
|
||||
.unwrap_err();
|
||||
assert!(
|
||||
negative.contains("cooldown_redis_read_interval"),
|
||||
"{negative}"
|
||||
);
|
||||
assert!(text.contains("max_redis_batch_cache_size"), "{text}");
|
||||
});
|
||||
}
|
||||
}
|
||||
591
litellm-rust/crates/python-bridge/src/redis/store/rate_limits.rs
Normal file
591
litellm-rust/crates/python-bridge/src/redis/store/rate_limits.rs
Normal file
|
|
@ -0,0 +1,591 @@
|
|||
use std::{sync::Arc, time::Duration};
|
||||
|
||||
use indexmap::IndexMap;
|
||||
use litellm_redis_store::{
|
||||
Pending,
|
||||
rate_limits::{
|
||||
AcquireReply, Applied, AtomicCounter, AtomicReply, Counter, Entry, Gauge, GuardedIncrement,
|
||||
GuardedReply, Increment, Kind, Local, RateLimits, Scope, WindowCount,
|
||||
},
|
||||
};
|
||||
use pyo3::{exceptions::PyValueError, prelude::*};
|
||||
|
||||
use super::{
|
||||
NativePending, PyNumber, Spelled, kind_error, redis_async, remote,
|
||||
request::{NativeRequestRedis, NativeSend},
|
||||
ttl,
|
||||
};
|
||||
use crate::redis::client::NativeRedisClient;
|
||||
|
||||
/// The v3 rate limiter's keyspace. Memory reads and writes are synchronous like `InMemoryCache`;
|
||||
/// every Redis call is a future.
|
||||
#[pyclass(frozen, name = "NativeRateLimits")]
|
||||
pub(crate) struct NativeRateLimits(Arc<RateLimits>);
|
||||
|
||||
/// An `int` or `float` counter, a `str` window start, or a `{slot_id: acquired_at}` registry.
|
||||
#[derive(FromPyObject)]
|
||||
enum PyLocal {
|
||||
Number(PyNumber),
|
||||
Text(String),
|
||||
Slots(IndexMap<String, f64>),
|
||||
}
|
||||
|
||||
impl From<PyLocal> for Local {
|
||||
fn from(local: PyLocal) -> Self {
|
||||
match local {
|
||||
PyLocal::Number(number) => Self::Number(number.0),
|
||||
PyLocal::Text(text) => Self::Text(text),
|
||||
PyLocal::Slots(slots) => Self::Slots(slots.into_iter().collect()),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// What the memory tier holds, back as the Python value it was written as.
|
||||
struct LocalValue(Local);
|
||||
|
||||
impl<'py> IntoPyObject<'py> for LocalValue {
|
||||
type Target = PyAny;
|
||||
type Output = Bound<'py, PyAny>;
|
||||
type Error = PyErr;
|
||||
|
||||
fn into_pyobject(self, py: Python<'py>) -> PyResult<Self::Output> {
|
||||
Ok(match self.0 {
|
||||
Local::Number(number) => PyNumber(number).into_pyobject(py)?,
|
||||
Local::Text(text) => text.into_pyobject(py)?.into_any(),
|
||||
Local::Slots(slots) => slots
|
||||
.into_iter()
|
||||
.collect::<IndexMap<_, _>>()
|
||||
.into_pyobject(py)?
|
||||
.into_any(),
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
/// `litellm.redis.stores.rate_limits.RateLimitScope`.
|
||||
#[derive(FromPyObject)]
|
||||
struct PyScope {
|
||||
#[pyo3(attribute)]
|
||||
descriptor_key: String,
|
||||
#[pyo3(attribute)]
|
||||
descriptor_value: String,
|
||||
}
|
||||
|
||||
impl From<PyScope> for Scope {
|
||||
fn from(scope: PyScope) -> Self {
|
||||
Self {
|
||||
descriptor: scope.descriptor_key,
|
||||
value: scope.descriptor_value,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// `RateLimitCounter`.
|
||||
#[derive(FromPyObject)]
|
||||
struct PyCounter {
|
||||
#[pyo3(attribute)]
|
||||
scope: PyScope,
|
||||
#[pyo3(attribute)]
|
||||
kind: Spelled<Kind>,
|
||||
}
|
||||
|
||||
impl From<PyCounter> for Counter {
|
||||
fn from(counter: PyCounter) -> Self {
|
||||
Self {
|
||||
scope: counter.scope.into(),
|
||||
kind: counter.kind.0,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// `RateLimitWindow`, `RateLimitCounter` or `RateLimitGauge`, told apart by the key suffix each names.
|
||||
#[derive(FromPyObject)]
|
||||
struct PyEntry {
|
||||
#[pyo3(attribute)]
|
||||
suffix: String,
|
||||
#[pyo3(attribute)]
|
||||
scope: PyScope,
|
||||
}
|
||||
|
||||
impl TryFrom<PyEntry> for Entry {
|
||||
type Error = PyErr;
|
||||
|
||||
fn try_from(entry: PyEntry) -> PyResult<Self> {
|
||||
Self::new(&entry.suffix, entry.scope.into())
|
||||
.map_err(|_| PyValueError::new_err(format!("unknown key suffix {:?}", entry.suffix)))
|
||||
}
|
||||
}
|
||||
|
||||
/// `CounterIncrement`.
|
||||
#[derive(FromPyObject)]
|
||||
struct PyIncrement {
|
||||
#[pyo3(attribute)]
|
||||
counter: PyCounter,
|
||||
#[pyo3(attribute)]
|
||||
amount: PyNumber,
|
||||
#[pyo3(attribute)]
|
||||
ttl: Option<u64>,
|
||||
}
|
||||
|
||||
impl From<PyIncrement> for Increment {
|
||||
fn from(increment: PyIncrement) -> Self {
|
||||
Self {
|
||||
counter: increment.counter.into(),
|
||||
amount: increment.amount.0,
|
||||
ttl: increment.ttl.map(Duration::from_secs),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// `AtomicCounter`.
|
||||
#[derive(FromPyObject)]
|
||||
struct PyAtomicCounter {
|
||||
#[pyo3(attribute)]
|
||||
counter: PyCounter,
|
||||
#[pyo3(attribute)]
|
||||
limit: i64,
|
||||
#[pyo3(attribute)]
|
||||
increment: i64,
|
||||
#[pyo3(attribute)]
|
||||
ttl: i64,
|
||||
#[pyo3(attribute)]
|
||||
window_size: i64,
|
||||
}
|
||||
|
||||
impl From<PyAtomicCounter> for AtomicCounter {
|
||||
fn from(counter: PyAtomicCounter) -> Self {
|
||||
Self {
|
||||
counter: counter.counter.into(),
|
||||
limit: counter.limit,
|
||||
increment: counter.increment,
|
||||
ttl_seconds: counter.ttl,
|
||||
window_size_seconds: counter.window_size,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn all<P, T: From<P>>(records: Vec<P>) -> Vec<T> {
|
||||
records.into_iter().map(T::from).collect()
|
||||
}
|
||||
|
||||
#[pyclass(frozen, eq, get_all, name = "WindowCount")]
|
||||
#[derive(PartialEq)]
|
||||
pub(crate) struct PyWindowCount {
|
||||
window_start: String,
|
||||
count: i64,
|
||||
}
|
||||
|
||||
#[pymethods]
|
||||
impl PyWindowCount {
|
||||
#[new]
|
||||
fn new(window_start: String, count: i64) -> Self {
|
||||
Self {
|
||||
window_start,
|
||||
count,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Per counter, the count after the increment and the start of the window it landed in.
|
||||
#[pyclass(frozen, eq, get_all, name = "AtomicApplied")]
|
||||
#[derive(PartialEq)]
|
||||
pub(crate) struct PyAtomicApplied {
|
||||
counters: Vec<(i64, i64)>,
|
||||
}
|
||||
|
||||
#[pymethods]
|
||||
impl PyAtomicApplied {
|
||||
#[new]
|
||||
fn new(counters: Vec<(i64, i64)>) -> Self {
|
||||
Self { counters }
|
||||
}
|
||||
}
|
||||
|
||||
#[pyclass(frozen, eq, get_all, name = "AtomicOverLimit")]
|
||||
#[derive(PartialEq)]
|
||||
pub(crate) struct PyAtomicOverLimit {
|
||||
index: usize,
|
||||
current: i64,
|
||||
limit: i64,
|
||||
}
|
||||
|
||||
#[pymethods]
|
||||
impl PyAtomicOverLimit {
|
||||
#[new]
|
||||
fn new(index: usize, current: i64, limit: i64) -> Self {
|
||||
Self {
|
||||
index,
|
||||
current,
|
||||
limit,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(IntoPyObject)]
|
||||
enum PyAtomicReply {
|
||||
Applied(PyAtomicApplied),
|
||||
OverLimit(PyAtomicOverLimit),
|
||||
}
|
||||
|
||||
#[pyclass(frozen, eq, get_all, name = "GuardedReply")]
|
||||
#[derive(PartialEq)]
|
||||
pub(crate) struct PyGuardedReply {
|
||||
applied: bool,
|
||||
count: i64,
|
||||
}
|
||||
|
||||
#[pymethods]
|
||||
impl PyGuardedReply {
|
||||
#[new]
|
||||
fn new(applied: bool, count: i64) -> Self {
|
||||
Self { applied, count }
|
||||
}
|
||||
}
|
||||
|
||||
#[pyclass(frozen, eq, get_all, name = "GaugesAcquired")]
|
||||
#[derive(PartialEq)]
|
||||
pub(crate) struct PyGaugesAcquired {
|
||||
in_flight: Vec<i64>,
|
||||
}
|
||||
|
||||
#[pymethods]
|
||||
impl PyGaugesAcquired {
|
||||
#[new]
|
||||
fn new(in_flight: Vec<i64>) -> Self {
|
||||
Self { in_flight }
|
||||
}
|
||||
}
|
||||
|
||||
#[pyclass(frozen, eq, get_all, name = "GaugeOverLimit")]
|
||||
#[derive(PartialEq)]
|
||||
pub(crate) struct PyGaugeOverLimit {
|
||||
index: usize,
|
||||
in_flight: i64,
|
||||
limit: i64,
|
||||
}
|
||||
|
||||
#[pymethods]
|
||||
impl PyGaugeOverLimit {
|
||||
#[new]
|
||||
fn new(index: usize, in_flight: i64, limit: i64) -> Self {
|
||||
Self {
|
||||
index,
|
||||
in_flight,
|
||||
limit,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(IntoPyObject)]
|
||||
enum PyAcquireReply {
|
||||
Acquired(PyGaugesAcquired),
|
||||
OverLimit(PyGaugeOverLimit),
|
||||
}
|
||||
|
||||
fn window_counts(counts: Vec<WindowCount>) -> Vec<PyWindowCount> {
|
||||
counts
|
||||
.into_iter()
|
||||
.map(
|
||||
|WindowCount {
|
||||
window_start,
|
||||
count,
|
||||
}| PyWindowCount {
|
||||
window_start,
|
||||
count,
|
||||
},
|
||||
)
|
||||
.collect()
|
||||
}
|
||||
|
||||
fn atomic(reply: AtomicReply) -> PyAtomicReply {
|
||||
match reply {
|
||||
AtomicReply::Applied(applied) => PyAtomicReply::Applied(PyAtomicApplied {
|
||||
counters: applied
|
||||
.into_iter()
|
||||
.map(
|
||||
|Applied {
|
||||
count,
|
||||
window_start,
|
||||
}| (count, window_start),
|
||||
)
|
||||
.collect(),
|
||||
}),
|
||||
AtomicReply::OverLimit {
|
||||
index,
|
||||
current,
|
||||
limit,
|
||||
} => PyAtomicReply::OverLimit(PyAtomicOverLimit {
|
||||
index,
|
||||
current,
|
||||
limit,
|
||||
}),
|
||||
}
|
||||
}
|
||||
|
||||
fn guarded(GuardedReply { applied, count }: GuardedReply) -> PyGuardedReply {
|
||||
PyGuardedReply { applied, count }
|
||||
}
|
||||
|
||||
fn acquire(reply: AcquireReply) -> PyAcquireReply {
|
||||
match reply {
|
||||
AcquireReply::Acquired(in_flight) => {
|
||||
PyAcquireReply::Acquired(PyGaugesAcquired { in_flight })
|
||||
}
|
||||
AcquireReply::OverLimit {
|
||||
index,
|
||||
in_flight,
|
||||
limit,
|
||||
} => PyAcquireReply::OverLimit(PyGaugeOverLimit {
|
||||
index,
|
||||
in_flight,
|
||||
limit,
|
||||
}),
|
||||
}
|
||||
}
|
||||
|
||||
/// One script call declared on the batch the caller's `NativeSend` picked; `None` without Redis.
|
||||
fn script<T, R>(
|
||||
call_type: &'static str,
|
||||
declared: Option<Pending<T>>,
|
||||
reply: fn(T) -> R,
|
||||
) -> Option<NativePending>
|
||||
where
|
||||
T: Send + 'static,
|
||||
R: for<'a> IntoPyObject<'a> + Send + 'static,
|
||||
{
|
||||
Some(super::pending(call_type, kind_error, declared?, reply))
|
||||
}
|
||||
|
||||
#[pymethods]
|
||||
impl NativeRateLimits {
|
||||
#[new]
|
||||
#[pyo3(signature = (client, backfill_ttl_seconds = None))]
|
||||
fn new(
|
||||
client: Option<PyRef<'_, NativeRedisClient>>,
|
||||
backfill_ttl_seconds: Option<f64>,
|
||||
) -> PyResult<Self> {
|
||||
Ok(Self(Arc::new(RateLimits::new(
|
||||
client.map(|client| remote(&client)),
|
||||
ttl(backfill_ttl_seconds)?,
|
||||
))))
|
||||
}
|
||||
|
||||
#[getter]
|
||||
fn has_redis(&self) -> bool {
|
||||
self.0.has_redis()
|
||||
}
|
||||
|
||||
fn slot_groups(&self, scopes: Vec<PyScope>, force: bool) -> Vec<Vec<usize>> {
|
||||
self.0.slot_groups(&all(scopes), force)
|
||||
}
|
||||
|
||||
fn get_local(&self, entry: PyEntry) -> PyResult<Option<LocalValue>> {
|
||||
Ok(self
|
||||
.0
|
||||
.get_local(&entry.try_into()?)
|
||||
.map_err(kind_error)?
|
||||
.map(LocalValue))
|
||||
}
|
||||
|
||||
#[pyo3(signature = (entry, value, ttl_seconds))]
|
||||
fn set_local(&self, entry: PyEntry, value: PyLocal, ttl_seconds: Option<f64>) -> PyResult<()> {
|
||||
self.0
|
||||
.set_local(&entry.try_into()?, value.into(), ttl(ttl_seconds)?)
|
||||
.map_err(kind_error)
|
||||
}
|
||||
|
||||
/// `request` answers a key its pre-call reads already found absent without another `GET`.
|
||||
#[pyo3(signature = (entry, *, backfill_ttl_seconds = None, request = None))]
|
||||
fn get<'py>(
|
||||
&self,
|
||||
py: Python<'py>,
|
||||
entry: PyEntry,
|
||||
backfill_ttl_seconds: Option<f64>,
|
||||
request: Option<&NativeRequestRedis>,
|
||||
) -> PyResult<Bound<'py, PyAny>> {
|
||||
let store = self.0.clone();
|
||||
let entry = Entry::try_from(entry)?;
|
||||
let backfill_ttl = ttl(backfill_ttl_seconds)?;
|
||||
let request = request.map(|request| request.0.clone());
|
||||
redis_async(py, "async_get_cache", kind_error, async move {
|
||||
store
|
||||
.get(request.as_ref(), &entry, backfill_ttl)
|
||||
.await
|
||||
.map(|local| local.map(LocalValue))
|
||||
})
|
||||
}
|
||||
|
||||
fn sliding_window(
|
||||
&self,
|
||||
counters: Vec<PyCounter>,
|
||||
now: i64,
|
||||
window_size: i64,
|
||||
send: &NativeSend,
|
||||
) -> Option<NativePending> {
|
||||
script(
|
||||
"batch_rate_limiter_script",
|
||||
self.0
|
||||
.sliding_window(send.send(), &all(counters), now, window_size),
|
||||
window_counts,
|
||||
)
|
||||
}
|
||||
|
||||
fn read_counters(&self, counters: Vec<PyCounter>, send: &NativeSend) -> Option<NativePending> {
|
||||
script(
|
||||
"batch_counter_read_script",
|
||||
self.0.read_counters(send.send(), &all(counters)),
|
||||
std::convert::identity,
|
||||
)
|
||||
}
|
||||
|
||||
fn check_and_increment(
|
||||
&self,
|
||||
counters: Vec<PyAtomicCounter>,
|
||||
send: &NativeSend,
|
||||
) -> Option<NativePending> {
|
||||
script(
|
||||
"check_and_increment_by_n_script",
|
||||
self.0.check_and_increment(send.send(), &all(counters)),
|
||||
atomic,
|
||||
)
|
||||
}
|
||||
|
||||
fn guarded_increment(
|
||||
&self,
|
||||
counter: PyCounter,
|
||||
expected_window_start: String,
|
||||
amount: PyNumber,
|
||||
ttl_seconds: Option<u64>,
|
||||
send: &NativeSend,
|
||||
) -> Option<NativePending> {
|
||||
let increment = GuardedIncrement {
|
||||
counter: counter.into(),
|
||||
expected_window_start,
|
||||
amount: amount.0,
|
||||
ttl: ttl_seconds.map(Duration::from_secs),
|
||||
};
|
||||
script(
|
||||
"window_guarded_token_increment_script",
|
||||
self.0.guarded_increment(send.send(), &increment),
|
||||
guarded,
|
||||
)
|
||||
}
|
||||
|
||||
fn parallel_acquire(
|
||||
&self,
|
||||
gauges: Vec<(PyScope, i64)>,
|
||||
slot_ttl_seconds: u64,
|
||||
slot_id: &str,
|
||||
send: &NativeSend,
|
||||
) -> Option<NativePending> {
|
||||
let gauges: Vec<Gauge> = gauges
|
||||
.into_iter()
|
||||
.map(|(scope, limit)| Gauge {
|
||||
scope: scope.into(),
|
||||
limit,
|
||||
})
|
||||
.collect();
|
||||
script(
|
||||
"parallel_acquire_script",
|
||||
self.0.parallel_acquire(
|
||||
send.send(),
|
||||
&gauges,
|
||||
Duration::from_secs(slot_ttl_seconds),
|
||||
slot_id,
|
||||
),
|
||||
acquire,
|
||||
)
|
||||
}
|
||||
|
||||
fn parallel_release(
|
||||
&self,
|
||||
gauges: Vec<PyScope>,
|
||||
slot_id: &str,
|
||||
send: &NativeSend,
|
||||
) -> Option<NativePending> {
|
||||
script(
|
||||
"parallel_release_script",
|
||||
self.0.parallel_release(send.send(), &all(gauges), slot_id),
|
||||
std::convert::identity,
|
||||
)
|
||||
}
|
||||
|
||||
fn parallel_count(
|
||||
&self,
|
||||
gauges: Vec<PyScope>,
|
||||
slot_ttl_seconds: u64,
|
||||
send: &NativeSend,
|
||||
) -> Option<NativePending> {
|
||||
script(
|
||||
"parallel_count_script",
|
||||
self.0.parallel_count(
|
||||
send.send(),
|
||||
&all(gauges),
|
||||
Duration::from_secs(slot_ttl_seconds),
|
||||
),
|
||||
std::convert::identity,
|
||||
)
|
||||
}
|
||||
|
||||
fn token_increment(
|
||||
&self,
|
||||
increments: Vec<PyIncrement>,
|
||||
send: &NativeSend,
|
||||
) -> Option<NativePending> {
|
||||
script(
|
||||
"token_increment_script",
|
||||
self.0.token_increment(send.send(), &all(increments)),
|
||||
std::convert::identity,
|
||||
)
|
||||
}
|
||||
|
||||
/// `False` without Redis.
|
||||
fn token_increment_post_call(
|
||||
&self,
|
||||
request: &NativeRequestRedis,
|
||||
increments: Vec<PyIncrement>,
|
||||
) -> bool {
|
||||
self.0
|
||||
.token_increment_post_call(&request.0, all(increments))
|
||||
}
|
||||
|
||||
/// `False` without Redis.
|
||||
fn parallel_release_post_call(
|
||||
&self,
|
||||
request: &NativeRequestRedis,
|
||||
gauges: Vec<PyScope>,
|
||||
slot_id: &str,
|
||||
) -> bool {
|
||||
self.0
|
||||
.parallel_release_post_call(&request.0, &all(gauges), slot_id)
|
||||
}
|
||||
|
||||
fn increment<'py>(
|
||||
&self,
|
||||
py: Python<'py>,
|
||||
increment: PyIncrement,
|
||||
) -> PyResult<Bound<'py, PyAny>> {
|
||||
let store = self.0.clone();
|
||||
let increment = Increment::from(increment);
|
||||
redis_async(py, "async_increment", kind_error, async move {
|
||||
Ok(store.increment(&increment).await.map(PyNumber))
|
||||
})
|
||||
}
|
||||
|
||||
fn increment_many<'py>(
|
||||
&self,
|
||||
py: Python<'py>,
|
||||
increments: Vec<PyIncrement>,
|
||||
) -> PyResult<Bound<'py, PyAny>> {
|
||||
let store = self.0.clone();
|
||||
let increments: Vec<Increment> = all(increments);
|
||||
redis_async(py, "async_increment_pipeline", kind_error, async move {
|
||||
store.increment_many(&increments).await;
|
||||
Ok(())
|
||||
})
|
||||
}
|
||||
|
||||
fn increment_many_post_call(&self, request: &NativeRequestRedis, increments: Vec<PyIncrement>) {
|
||||
self.0
|
||||
.increment_many_post_call(&request.0, &all(increments));
|
||||
}
|
||||
}
|
||||
|
|
@ -0,0 +1,183 @@
|
|||
use std::sync::Arc;
|
||||
|
||||
use litellm_redis_store::{
|
||||
Send,
|
||||
router_cooldowns::{Cooldown, RouterCooldowns},
|
||||
};
|
||||
use pyo3::{
|
||||
prelude::*,
|
||||
types::{PyDict, PyTuple},
|
||||
};
|
||||
|
||||
use super::{
|
||||
NativePending, PyData, PyNumber, duration, kind_error, pending, policy::router_tier_policies,
|
||||
redis_async, redis_sync, remote, request::NativeSend,
|
||||
};
|
||||
use crate::redis::client::NativeRedisClient;
|
||||
|
||||
/// One router's cooldown store. The router's sync completion path calls `put`, `active_sync`,
|
||||
/// `minimum` and `increment_allowed_fails`, so those block with the GIL released like redis-py; the
|
||||
/// `async_` twins return awaitables for async callers.
|
||||
#[pyclass(frozen, name = "NativeRouterCooldowns")]
|
||||
pub(crate) struct NativeRouterCooldowns(Arc<RouterCooldowns>);
|
||||
|
||||
/// `[(model_id, cooldown), ...]` the way `RedisCache` hands the router its batch read.
|
||||
struct PyActive(Vec<(String, Cooldown)>);
|
||||
|
||||
impl<'py> IntoPyObject<'py> for PyActive {
|
||||
type Target = PyAny;
|
||||
type Output = Bound<'py, PyAny>;
|
||||
type Error = PyErr;
|
||||
|
||||
fn into_pyobject(self, py: Python<'py>) -> PyResult<Self::Output> {
|
||||
active_to_py(py, self.0)?
|
||||
.into_pyobject(py)
|
||||
.map(Bound::into_any)
|
||||
}
|
||||
}
|
||||
|
||||
fn active_to_py(
|
||||
py: Python<'_>,
|
||||
active: Vec<(String, Cooldown)>,
|
||||
) -> PyResult<Vec<Bound<'_, PyTuple>>> {
|
||||
active
|
||||
.into_iter()
|
||||
.map(|(model_id, cooldown)| {
|
||||
PyTuple::new(
|
||||
py,
|
||||
[
|
||||
model_id.into_pyobject(py)?.into_any(),
|
||||
PyData(cooldown).into_pyobject(py)?,
|
||||
],
|
||||
)
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
|
||||
#[pymethods]
|
||||
impl NativeRouterCooldowns {
|
||||
#[new]
|
||||
fn new(py: Python<'_>, client: Option<PyRef<'_, NativeRedisClient>>) -> PyResult<Self> {
|
||||
let remote = client.map(|client| remote(&client));
|
||||
let policy = router_tier_policies(py)?.cooldowns;
|
||||
Ok(Self(Arc::new(RouterCooldowns::new(remote, policy))))
|
||||
}
|
||||
|
||||
/// Adds the Redis tier and keeps memory; `false` when one was already attached.
|
||||
fn attach(&self, client: PyRef<'_, NativeRedisClient>) -> bool {
|
||||
self.0.attach(remote(&client))
|
||||
}
|
||||
|
||||
#[getter]
|
||||
fn has_redis(&self) -> bool {
|
||||
self.0.has_redis()
|
||||
}
|
||||
|
||||
fn put(
|
||||
&self,
|
||||
py: Python<'_>,
|
||||
model_id: String,
|
||||
cooldown: PyData<Cooldown>,
|
||||
ttl_seconds: f64,
|
||||
) -> PyResult<()> {
|
||||
let store = self.0.clone();
|
||||
let ttl = duration(ttl_seconds)?;
|
||||
redis_sync(py, "set_cache", kind_error, async move {
|
||||
store.put_str(&model_id, &cooldown.0, ttl).await
|
||||
})
|
||||
}
|
||||
|
||||
fn increment_allowed_fails(
|
||||
&self,
|
||||
py: Python<'_>,
|
||||
model_id: String,
|
||||
suffix: Option<String>,
|
||||
ttl_seconds: f64,
|
||||
) -> PyResult<i64> {
|
||||
let store = self.0.clone();
|
||||
let ttl = duration(ttl_seconds)?;
|
||||
redis_sync(py, "increment_cache", kind_error, async move {
|
||||
store
|
||||
.increment_allowed_fails_int(&model_id, suffix.as_deref(), ttl)
|
||||
.await
|
||||
})
|
||||
}
|
||||
|
||||
fn async_put<'py>(
|
||||
&self,
|
||||
py: Python<'py>,
|
||||
model_id: String,
|
||||
cooldown: PyData<Cooldown>,
|
||||
ttl_seconds: f64,
|
||||
) -> PyResult<Bound<'py, PyAny>> {
|
||||
let store = self.0.clone();
|
||||
let ttl = duration(ttl_seconds)?;
|
||||
redis_async(py, "async_set_cache", kind_error, async move {
|
||||
store.put(&model_id, &cooldown.0, ttl).await
|
||||
})
|
||||
}
|
||||
|
||||
fn async_increment_allowed_fails<'py>(
|
||||
&self,
|
||||
py: Python<'py>,
|
||||
model_id: String,
|
||||
suffix: Option<String>,
|
||||
ttl_seconds: f64,
|
||||
) -> PyResult<Bound<'py, PyAny>> {
|
||||
let store = self.0.clone();
|
||||
let ttl = duration(ttl_seconds)?;
|
||||
redis_async(py, "async_increment", kind_error, async move {
|
||||
store
|
||||
.increment_allowed_fails(&model_id, suffix.as_deref(), ttl)
|
||||
.await
|
||||
})
|
||||
}
|
||||
|
||||
/// The deployments still cooling down, read on the batch `send` picks.
|
||||
fn active(&self, model_ids: Vec<String>, send: &NativeSend) -> PyResult<NativePending> {
|
||||
let active = self.0.active(send.send(), &model_ids).map_err(kind_error)?;
|
||||
Ok(pending(
|
||||
"async_batch_get_cache",
|
||||
kind_error,
|
||||
active,
|
||||
PyActive,
|
||||
))
|
||||
}
|
||||
|
||||
/// `active` on its own with the GIL released, for the router's sync completion path.
|
||||
fn active_sync<'py>(
|
||||
&self,
|
||||
py: Python<'py>,
|
||||
model_ids: Vec<String>,
|
||||
) -> PyResult<Vec<Bound<'py, PyTuple>>> {
|
||||
let active = self.0.active(Send::Now, &model_ids).map_err(kind_error)?;
|
||||
let active = redis_sync(py, "batch_get_cache", kind_error, active)?;
|
||||
active_to_py(py, active)
|
||||
}
|
||||
|
||||
fn minimum(&self, py: Python<'_>, model_ids: Vec<String>) -> PyResult<Option<PyNumber>> {
|
||||
let store = self.0.clone();
|
||||
redis_sync(py, "batch_get_cache", kind_error, async move {
|
||||
store
|
||||
.minimum(&model_ids)
|
||||
.await
|
||||
.map(|minimum| minimum.map(PyNumber))
|
||||
})
|
||||
}
|
||||
|
||||
fn delete_local(&self, model_id: &str) -> PyResult<()> {
|
||||
self.0.delete_local(model_id).map_err(kind_error)
|
||||
}
|
||||
|
||||
fn flush_local_counters(&self) -> PyResult<()> {
|
||||
self.0.flush_local_counters().map_err(kind_error)
|
||||
}
|
||||
|
||||
fn snapshot<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, PyDict>> {
|
||||
let snapshot = PyDict::new(py);
|
||||
for (key, cooldown) in self.0.snapshot().map_err(kind_error)? {
|
||||
snapshot.set_item(key, PyData(cooldown))?;
|
||||
}
|
||||
Ok(snapshot)
|
||||
}
|
||||
}
|
||||
|
|
@ -0,0 +1,184 @@
|
|||
use std::{sync::Arc, time::Duration};
|
||||
|
||||
use litellm_redis_store::router_usage::{Counter, Kind, RouterUsage};
|
||||
use pyo3::{prelude::*, types::PyDict};
|
||||
|
||||
use super::{
|
||||
PyNumber, Spelled, kind_error, policy::router_tier_policies, redis_async, redis_sync, remote,
|
||||
};
|
||||
use crate::redis::client::NativeRedisClient;
|
||||
|
||||
/// `litellm.redis.stores.router_usage.UsageCounter`.
|
||||
#[derive(FromPyObject)]
|
||||
struct PyCounter {
|
||||
#[pyo3(attribute)]
|
||||
deployment_id: String,
|
||||
#[pyo3(attribute)]
|
||||
model: Option<String>,
|
||||
#[pyo3(attribute)]
|
||||
minute: String,
|
||||
#[pyo3(attribute)]
|
||||
kind: Spelled<Kind>,
|
||||
}
|
||||
|
||||
impl From<PyCounter> for Counter {
|
||||
fn from(counter: PyCounter) -> Self {
|
||||
Self {
|
||||
deployment_id: counter.deployment_id,
|
||||
model: counter.model,
|
||||
minute: counter.minute,
|
||||
kind: counter.kind.0,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn counters(counters: Vec<PyCounter>) -> Vec<Counter> {
|
||||
counters.into_iter().map(Counter::from).collect()
|
||||
}
|
||||
|
||||
/// One router's global TPM/RPM counters. The `_sync` methods block with the GIL released, for the
|
||||
/// Rust router's sync calls, which have no event loop to await on.
|
||||
#[pyclass(frozen, name = "NativeRouterUsage")]
|
||||
pub(crate) struct NativeRouterUsage(Arc<RouterUsage>);
|
||||
|
||||
#[pymethods]
|
||||
impl NativeRouterUsage {
|
||||
#[new]
|
||||
fn new(py: Python<'_>, client: Option<PyRef<'_, NativeRedisClient>>) -> PyResult<Self> {
|
||||
let remote = client.map(|client| remote(&client));
|
||||
let policy = router_tier_policies(py)?.usage;
|
||||
Ok(Self(Arc::new(RouterUsage::new(remote, policy))))
|
||||
}
|
||||
|
||||
/// Adds the Redis tier and keeps memory; `false` when one was already attached.
|
||||
fn attach(&self, client: PyRef<'_, NativeRedisClient>) -> bool {
|
||||
self.0.attach(remote(&client))
|
||||
}
|
||||
|
||||
#[getter]
|
||||
fn has_redis(&self) -> bool {
|
||||
self.0.has_redis()
|
||||
}
|
||||
|
||||
fn increment<'py>(
|
||||
&self,
|
||||
py: Python<'py>,
|
||||
counter: PyCounter,
|
||||
value: PyNumber,
|
||||
ttl_seconds: u64,
|
||||
) -> PyResult<Bound<'py, PyAny>> {
|
||||
let store = self.0.clone();
|
||||
let counter = Counter::from(counter);
|
||||
redis_async(py, "async_increment", kind_error, async move {
|
||||
Ok(store
|
||||
.increment(&counter, value.0, Duration::from_secs(ttl_seconds))
|
||||
.await
|
||||
.map(PyNumber))
|
||||
})
|
||||
}
|
||||
|
||||
/// `_increment_deployment`'s `async_increment_pipeline`, always on its own round trip.
|
||||
fn increment_deployment<'py>(
|
||||
&self,
|
||||
py: Python<'py>,
|
||||
deployment_id: String,
|
||||
model: Option<String>,
|
||||
minute: String,
|
||||
tokens: PyNumber,
|
||||
requests: i64,
|
||||
) -> PyResult<Bound<'py, PyAny>> {
|
||||
let store = self.0.clone();
|
||||
redis_async(py, "async_increment_pipeline", kind_error, async move {
|
||||
store
|
||||
.increment_deployment(
|
||||
&deployment_id,
|
||||
model.as_deref(),
|
||||
&minute,
|
||||
tokens.0,
|
||||
requests,
|
||||
)
|
||||
.await
|
||||
})
|
||||
}
|
||||
|
||||
fn get_many<'py>(
|
||||
&self,
|
||||
py: Python<'py>,
|
||||
counters: Vec<PyCounter>,
|
||||
) -> PyResult<Bound<'py, PyAny>> {
|
||||
let store = self.0.clone();
|
||||
let counters = self::counters(counters);
|
||||
redis_async(py, "async_batch_get_cache", kind_error, async move {
|
||||
let read = store.get_many(&counters).await?;
|
||||
Ok(read
|
||||
.into_iter()
|
||||
.map(|count| count.map(PyNumber))
|
||||
.collect::<Vec<_>>())
|
||||
})
|
||||
}
|
||||
|
||||
fn increment_sync(
|
||||
&self,
|
||||
py: Python<'_>,
|
||||
counter: PyCounter,
|
||||
value: PyNumber,
|
||||
ttl_seconds: u64,
|
||||
) -> PyResult<Option<PyNumber>> {
|
||||
let store = self.0.clone();
|
||||
let counter = Counter::from(counter);
|
||||
redis_sync(py, "increment_cache", kind_error, async move {
|
||||
Ok(store
|
||||
.increment(&counter, value.0, Duration::from_secs(ttl_seconds))
|
||||
.await
|
||||
.map(PyNumber))
|
||||
})
|
||||
}
|
||||
|
||||
fn increment_deployment_sync(
|
||||
&self,
|
||||
py: Python<'_>,
|
||||
deployment_id: String,
|
||||
model: Option<String>,
|
||||
minute: String,
|
||||
tokens: PyNumber,
|
||||
requests: i64,
|
||||
) -> PyResult<()> {
|
||||
let store = self.0.clone();
|
||||
redis_sync(py, "increment_pipeline", kind_error, async move {
|
||||
store
|
||||
.increment_deployment(
|
||||
&deployment_id,
|
||||
model.as_deref(),
|
||||
&minute,
|
||||
tokens.0,
|
||||
requests,
|
||||
)
|
||||
.await
|
||||
})
|
||||
}
|
||||
|
||||
fn get_many_sync(
|
||||
&self,
|
||||
py: Python<'_>,
|
||||
counters: Vec<PyCounter>,
|
||||
) -> PyResult<Vec<Option<PyNumber>>> {
|
||||
let store = self.0.clone();
|
||||
let counters = self::counters(counters);
|
||||
redis_sync(py, "batch_get_cache", kind_error, async move {
|
||||
let read = store.get_many(&counters).await?;
|
||||
Ok(read.into_iter().map(|count| count.map(PyNumber)).collect())
|
||||
})
|
||||
}
|
||||
|
||||
fn flush_local_counters(&self) -> PyResult<()> {
|
||||
self.0.flush_local_counters().map_err(kind_error)
|
||||
}
|
||||
|
||||
fn snapshot<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, PyDict>> {
|
||||
let snapshot = PyDict::new(py);
|
||||
for (key, count) in self.0.snapshot().map_err(kind_error)? {
|
||||
snapshot.set_item(key, PyNumber(count))?;
|
||||
}
|
||||
Ok(snapshot)
|
||||
}
|
||||
}
|
||||
|
|
@ -0,0 +1,49 @@
|
|||
use std::sync::Arc;
|
||||
|
||||
use litellm_redis_store::sensitive_route_pins::SensitiveRoutePins;
|
||||
use pyo3::prelude::*;
|
||||
|
||||
use super::{duration, kind_error, redis_async, remote};
|
||||
use crate::redis::client::NativeRedisClient;
|
||||
|
||||
#[pyclass(frozen, name = "NativeSensitiveRoutePins")]
|
||||
pub(crate) struct NativeSensitiveRoutePins(Arc<SensitiveRoutePins>);
|
||||
|
||||
#[pymethods]
|
||||
impl NativeSensitiveRoutePins {
|
||||
#[new]
|
||||
fn new(client: Option<PyRef<'_, NativeRedisClient>>) -> Self {
|
||||
Self(Arc::new(SensitiveRoutePins::new(
|
||||
client.map(|client| remote(&client)),
|
||||
)))
|
||||
}
|
||||
|
||||
fn get<'py>(
|
||||
&self,
|
||||
py: Python<'py>,
|
||||
tenant: String,
|
||||
session_id: String,
|
||||
default_ttl_seconds: f64,
|
||||
) -> PyResult<Bound<'py, PyAny>> {
|
||||
let store = self.0.clone();
|
||||
let default_ttl = duration(default_ttl_seconds)?;
|
||||
redis_async(py, "async_get_cache", kind_error, async move {
|
||||
store.get(&tenant, &session_id, default_ttl).await
|
||||
})
|
||||
}
|
||||
|
||||
fn set<'py>(
|
||||
&self,
|
||||
py: Python<'py>,
|
||||
tenant: String,
|
||||
session_id: String,
|
||||
model: String,
|
||||
ttl_seconds: f64,
|
||||
) -> PyResult<Bound<'py, PyAny>> {
|
||||
let store = self.0.clone();
|
||||
let ttl = duration(ttl_seconds)?;
|
||||
redis_async(py, "async_set_cache", kind_error, async move {
|
||||
store.set(&tenant, &session_id, &model, ttl).await
|
||||
})
|
||||
}
|
||||
}
|
||||
|
|
@ -0,0 +1,256 @@
|
|||
use std::sync::Arc;
|
||||
|
||||
use litellm_redis_store::{
|
||||
Number,
|
||||
spend_counters::{LostSeed, Reading, SpendCounter, SpendCounters},
|
||||
};
|
||||
use pyo3::{prelude::*, types::PyDict};
|
||||
|
||||
use super::{
|
||||
NativePending, PyData, PyNumber, kind_error, pending, redis_async, remote,
|
||||
request::{NativeRequestRedis, NativeSend},
|
||||
ttl, unwrap_all,
|
||||
};
|
||||
use crate::redis::client::NativeRedisClient;
|
||||
|
||||
/// A counter's value and whether Redis answered it, a clean miss included.
|
||||
#[pyclass(frozen, eq, get_all, name = "SpendRead")]
|
||||
#[derive(PartialEq)]
|
||||
pub(crate) struct PySpendRead {
|
||||
value: Option<PyNumber>,
|
||||
from_redis: bool,
|
||||
}
|
||||
|
||||
#[pymethods]
|
||||
impl PySpendRead {
|
||||
#[new]
|
||||
fn new(value: Option<PyNumber>, from_redis: bool) -> Self {
|
||||
Self { value, from_redis }
|
||||
}
|
||||
}
|
||||
|
||||
fn increments(increments: Vec<(PyData<SpendCounter>, PyNumber)>) -> Vec<(SpendCounter, Number)> {
|
||||
increments
|
||||
.into_iter()
|
||||
.map(|(counter, amount)| (counter.0, amount.0))
|
||||
.collect()
|
||||
}
|
||||
|
||||
fn numbers(values: Option<Vec<Option<Number>>>) -> Option<Vec<Option<PyNumber>>> {
|
||||
values.map(|values| {
|
||||
values
|
||||
.into_iter()
|
||||
.map(|value| value.map(PyNumber))
|
||||
.collect()
|
||||
})
|
||||
}
|
||||
|
||||
/// The proxy's spend counters, Redis first with this worker's memory copy behind it.
|
||||
#[pyclass(frozen, name = "NativeSpendCounters")]
|
||||
pub(crate) struct NativeSpendCounters(Arc<SpendCounters>);
|
||||
|
||||
#[pymethods]
|
||||
impl NativeSpendCounters {
|
||||
#[new]
|
||||
fn new(client: Option<PyRef<'_, NativeRedisClient>>) -> Self {
|
||||
Self(Arc::new(SpendCounters::new(
|
||||
client.map(|client| remote(&client)),
|
||||
)))
|
||||
}
|
||||
|
||||
/// Adds the Redis tier and keeps memory; `false` when one was already attached.
|
||||
fn attach(&self, client: PyRef<'_, NativeRedisClient>) -> bool {
|
||||
self.0.attach(remote(&client))
|
||||
}
|
||||
|
||||
#[getter]
|
||||
fn has_redis(&self) -> bool {
|
||||
self.0.has_redis()
|
||||
}
|
||||
|
||||
#[getter]
|
||||
fn redis_ttl_seconds(&self) -> Option<f64> {
|
||||
self.0.redis_ttl().map(|ttl| ttl.as_secs_f64())
|
||||
}
|
||||
|
||||
/// `("redis" | "memory", value)`: where the answer came from, a Redis miss included.
|
||||
fn read<'py>(
|
||||
&self,
|
||||
py: Python<'py>,
|
||||
counter: PyData<SpendCounter>,
|
||||
) -> PyResult<Bound<'py, PyAny>> {
|
||||
let store = self.0.clone();
|
||||
redis_async(py, "async_get_cache", kind_error, async move {
|
||||
let reading = store.read(&counter.0).await?;
|
||||
Ok(PySpendRead {
|
||||
value: reading.value().map(PyNumber),
|
||||
from_redis: matches!(reading, Reading::Redis(_)),
|
||||
})
|
||||
})
|
||||
}
|
||||
|
||||
/// One `MGET` on the batch `send` picks; `None` without Redis. It resolves to `None` when the
|
||||
/// read failed, so each caller reads on its own.
|
||||
fn read_many(
|
||||
&self,
|
||||
counters: Vec<PyData<SpendCounter>>,
|
||||
send: &NativeSend,
|
||||
) -> Option<NativePending> {
|
||||
let read = self.0.read_many(send.send(), &unwrap_all(counters))?;
|
||||
Some(pending("async_batch_get_cache", kind_error, read, numbers))
|
||||
}
|
||||
|
||||
fn increment<'py>(
|
||||
&self,
|
||||
py: Python<'py>,
|
||||
counter: PyData<SpendCounter>,
|
||||
amount: PyNumber,
|
||||
) -> PyResult<Bound<'py, PyAny>> {
|
||||
let store = self.0.clone();
|
||||
redis_async(py, "async_increment", kind_error, async move {
|
||||
store.increment(&counter.0, amount.0).await.map(PyNumber)
|
||||
})
|
||||
}
|
||||
|
||||
fn increment_many<'py>(
|
||||
&self,
|
||||
py: Python<'py>,
|
||||
increments: Vec<(PyData<SpendCounter>, PyNumber)>,
|
||||
) -> PyResult<Bound<'py, PyAny>> {
|
||||
let store = self.0.clone();
|
||||
let increments = self::increments(increments);
|
||||
redis_async(py, "async_increment_pipeline", kind_error, async move {
|
||||
let totals = store.increment_many(&increments).await?;
|
||||
Ok(totals.into_iter().map(PyNumber).collect::<Vec<_>>())
|
||||
})
|
||||
}
|
||||
|
||||
/// Declares the increments on the request's post-call batch; `False` when there is no Redis.
|
||||
fn increment_many_post_call(
|
||||
&self,
|
||||
request: &NativeRequestRedis,
|
||||
increments: Vec<(PyData<SpendCounter>, PyNumber)>,
|
||||
) -> bool {
|
||||
self.0
|
||||
.increment_many_post_call(&request.0, &self::increments(increments))
|
||||
}
|
||||
|
||||
fn set_max<'py>(
|
||||
&self,
|
||||
py: Python<'py>,
|
||||
counter: PyData<SpendCounter>,
|
||||
spend: PyNumber,
|
||||
) -> PyResult<Bound<'py, PyAny>> {
|
||||
let store = self.0.clone();
|
||||
redis_async(py, "async_set_max", kind_error, async move {
|
||||
store
|
||||
.set_max(&counter.0, spend.0)
|
||||
.await
|
||||
.map(|stored| stored.map(PyNumber))
|
||||
})
|
||||
}
|
||||
|
||||
#[pyo3(signature = (counter, spend, *, increment_if_lost))]
|
||||
fn seed<'py>(
|
||||
&self,
|
||||
py: Python<'py>,
|
||||
counter: PyData<SpendCounter>,
|
||||
spend: PyNumber,
|
||||
increment_if_lost: bool,
|
||||
) -> PyResult<Bound<'py, PyAny>> {
|
||||
let store = self.0.clone();
|
||||
let lost = if increment_if_lost {
|
||||
LostSeed::Increment
|
||||
} else {
|
||||
LostSeed::Answer
|
||||
};
|
||||
redis_async(py, "async_set_cache", kind_error, async move {
|
||||
store.seed(&counter.0, spend.0, lost).await.map(PyNumber)
|
||||
})
|
||||
}
|
||||
|
||||
#[pyo3(signature = (counter, spend, ttl_seconds = None))]
|
||||
fn set<'py>(
|
||||
&self,
|
||||
py: Python<'py>,
|
||||
counter: PyData<SpendCounter>,
|
||||
spend: PyNumber,
|
||||
ttl_seconds: Option<f64>,
|
||||
) -> PyResult<Bound<'py, PyAny>> {
|
||||
let store = self.0.clone();
|
||||
let ttl = ttl(ttl_seconds)?;
|
||||
redis_async(py, "async_set_cache", kind_error, async move {
|
||||
store.set(&counter.0, spend.0, ttl).await
|
||||
})
|
||||
}
|
||||
|
||||
fn delete<'py>(
|
||||
&self,
|
||||
py: Python<'py>,
|
||||
counters: Vec<PyData<SpendCounter>>,
|
||||
) -> PyResult<Bound<'py, PyAny>> {
|
||||
let store = self.0.clone();
|
||||
let counters = unwrap_all(counters);
|
||||
redis_async(py, "async_delete_cache", kind_error, async move {
|
||||
store.delete(&counters).await
|
||||
})
|
||||
}
|
||||
|
||||
fn refresh_ttl<'py>(
|
||||
&self,
|
||||
py: Python<'py>,
|
||||
counter: PyData<SpendCounter>,
|
||||
) -> PyResult<Bound<'py, PyAny>> {
|
||||
let store = self.0.clone();
|
||||
redis_async(py, "async_refresh_ttl", kind_error, async move {
|
||||
store.refresh_ttl(&counter.0).await
|
||||
})
|
||||
}
|
||||
|
||||
fn get_local(&self, counter: PyData<SpendCounter>) -> PyResult<Option<PyNumber>> {
|
||||
Ok(self
|
||||
.0
|
||||
.get_local(&counter.0)
|
||||
.map_err(kind_error)?
|
||||
.map(PyNumber))
|
||||
}
|
||||
|
||||
#[pyo3(signature = (counter, spend, ttl_seconds = None))]
|
||||
fn set_local(
|
||||
&self,
|
||||
counter: PyData<SpendCounter>,
|
||||
spend: PyNumber,
|
||||
ttl_seconds: Option<f64>,
|
||||
) -> PyResult<()> {
|
||||
self.0
|
||||
.set_local(&counter.0, spend.0, ttl(ttl_seconds)?)
|
||||
.map_err(kind_error)
|
||||
}
|
||||
|
||||
fn delete_local(&self, counter: PyData<SpendCounter>) -> PyResult<()> {
|
||||
self.0.delete_local(&counter.0).map_err(kind_error)
|
||||
}
|
||||
|
||||
fn increment_local(
|
||||
&self,
|
||||
counter: PyData<SpendCounter>,
|
||||
amount: PyNumber,
|
||||
) -> PyResult<PyNumber> {
|
||||
self.0
|
||||
.increment_local(&counter.0, amount.0)
|
||||
.map(PyNumber)
|
||||
.map_err(kind_error)
|
||||
}
|
||||
|
||||
fn flush_local(&self) -> PyResult<()> {
|
||||
self.0.flush_local().map_err(kind_error)
|
||||
}
|
||||
|
||||
fn snapshot<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, PyDict>> {
|
||||
let snapshot = PyDict::new(py);
|
||||
for (key, spend) in self.0.snapshot().map_err(kind_error)? {
|
||||
snapshot.set_item(key, PyNumber(spend))?;
|
||||
}
|
||||
Ok(snapshot)
|
||||
}
|
||||
}
|
||||
|
|
@ -20,10 +20,10 @@ use std::{
|
|||
|
||||
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,
|
||||
Attempt, AttemptError, Blocked, BuildAsk, BuildPath, Built, CallbackFiltered, 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};
|
||||
|
|
@ -298,8 +298,9 @@ pub(crate) enum HostCallResponse {
|
|||
PreCall(PreCallFacts),
|
||||
/// The candidate models a `routing_plugins` select step reported.
|
||||
Models(Option<Vec<String>>),
|
||||
/// The ids a `callback_filters` select step left, or the error a filter raised.
|
||||
Filtered(Result<Vec<String>, AttemptError>),
|
||||
/// The ids a `callback_filters` select step left with the one a filter prefers, or the error
|
||||
/// a filter raised.
|
||||
Filtered(Result<(Vec<String>, Option<String>), 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),
|
||||
|
|
@ -352,8 +353,10 @@ impl HostCallResponse {
|
|||
got: self.kind(),
|
||||
});
|
||||
}
|
||||
if let (Self::Select(ids) | Self::Filtered(Ok(ids)), HostCall::Select { candidates, .. }) =
|
||||
(self, call)
|
||||
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() });
|
||||
|
|
@ -890,7 +893,7 @@ impl Host for StepHost {
|
|||
&self,
|
||||
ctx: &RoutingContext<'_>,
|
||||
candidates: Vec<Arc<Deployment>>,
|
||||
) -> Result<Vec<Arc<Deployment>>, Error> {
|
||||
) -> Result<CallbackFiltered, Error> {
|
||||
let (scope, group) = ctx.scope.wire();
|
||||
let joined = |own: &[String], walked: &[String]| [own, walked].concat();
|
||||
let call = HostCall::Select {
|
||||
|
|
@ -909,8 +912,11 @@ impl Host for StepHost {
|
|||
let HostCallResponse::Filtered(result) = call_host(call).await else {
|
||||
unreachable!("{CHECKED}");
|
||||
};
|
||||
let ids = result.map_err(Error::Attempt)?;
|
||||
Ok(answered_by_id(&candidates, &ids))
|
||||
let (ids, preferred) = result.map_err(Error::Attempt)?;
|
||||
Ok(CallbackFiltered {
|
||||
kept: answered_by_id(&candidates, &ids),
|
||||
preferred,
|
||||
})
|
||||
}
|
||||
|
||||
async fn allow_fallback(
|
||||
|
|
|
|||
|
|
@ -1181,7 +1181,8 @@ fn host_call_response(
|
|||
stage: Stage::CallbackFilters { .. },
|
||||
..
|
||||
} => {
|
||||
// The remaining ids, or a failure mapping as an attempt reports one.
|
||||
// The remaining ids, alone or paired with the id a filter prefers, or a failure
|
||||
// mapping as an attempt reports one.
|
||||
let outcome = outcome.ok_or_else(missing)?;
|
||||
if let Ok(failure) = outcome.cast::<PyDict>() {
|
||||
return match attempt_outcome(failure)? {
|
||||
|
|
@ -1191,7 +1192,10 @@ fn host_call_response(
|
|||
)),
|
||||
};
|
||||
}
|
||||
Ok(HostCallResponse::Filtered(Ok(outcome.extract()?)))
|
||||
if let Ok(answer) = outcome.cast::<PyTuple>() {
|
||||
return Ok(HostCallResponse::Filtered(Ok(answer.extract()?)));
|
||||
}
|
||||
Ok(HostCallResponse::Filtered(Ok((outcome.extract()?, None))))
|
||||
}
|
||||
HostCall::Select {
|
||||
stage: Stage::Strategy { records, .. },
|
||||
|
|
|
|||
|
|
@ -1,11 +1,11 @@
|
|||
- One folder per Redis family (`router_cooldowns/`, `router_usage/`, `spend_counters/`, `auth_cache/`, `rate_limits/`), each with `mod.rs` (the store and its ops), `keys.rs` (key types and spellings) and `types.rs` (Python-facing values with their `FromValue` / `IntoValue`). It is the only place that family's keys, value codec and scripts are written, and it takes domain ids, never keys. Shared pieces live once at the crate root: `Error` and `ErrorKind` (`error.rs`), `Memory` and `Clock` (`memory.rs`), `Pending` (`request.rs`), `Remote` and its attach-once `RemoteSlot`, which knows nothing of memory (`remote.rs`), `DualCache` (`dual_cache/`) and the Python codecs (`codec.rs`, `value.rs`)
|
||||
- One folder per Redis family (`router_cooldowns/`, `router_usage/`, `spend_counters/`, `auth_cache/`, `rate_limits/`), each with `mod.rs` (the store and its ops), `keys.rs` (key types and spellings) and `types.rs` (the family's values, with the `FromValue` / `IntoValue` codecs for what it stores in Redis). It is the only place that family's keys, value codec and scripts are written, and it takes domain ids, never keys. Shared pieces live once at the crate root: `Error` and `ErrorKind` (`error.rs`), `Memory` and `Clock` (`memory.rs`), `Pending` (`request.rs`), `Remote` and its attach-once `RemoteSlot`, which knows nothing of memory (`remote.rs`), `DualCache` (`dual_cache/`) and the Python codecs (`codec.rs`, `value.rs`)
|
||||
- An op declared on a `Send` returns `Pending<T>`, a lazy future built when the op is declared. An op that only means something on Redis (an `MGET` the caller falls back from, a script) returns `None` without one, and an op that replaced a `DualCache` call answers from memory. A post-call method returns `false` when there was no Redis to declare on
|
||||
- A store that sends a pipeline of its own uses `RedisBatch::pipelined`, never `Pipeline` and raw `Value`s
|
||||
- Python and Rust pods share these keys during a rolling deploy, so a module writes the bytes `main`'s Python writes and reads everything `main` reads. Python formats go through `litellm-python-compat` (`repr`, `literal_eval`, `json`), never a hand-rolled codec
|
||||
- A family holds one `DualCache` per `DualCache` its Python held on `main`, and each `DualCache` method is the Python method it names, so a family op that was a `DualCache` call is one call here. Where `main`'s Python reached past `DualCache` to `in_memory_cache` or `redis_cache`, the family does the same through `memory()` and `redis()`, which keeps that reach visible. Caches a family's Python attached to one `RedisCache` share one `Arc<RemoteSlot>`. Only a family that batch-reads builds its cache with `DualCache::batch_reading` from a `TierPolicy` the host passes in, and the Python host fills it from what `main` reads from the environment or `litellm` globals (`litellm/rust_bridge/host/settings.py`, `router_tiers`). Callers pass a `Remote` (or none), never a cache. Redis-only families are plain functions over `&RedisClient`
|
||||
- Every family works with no Redis, answering from memory. `attach` adds the Redis tier at most once and keeps memory, as `DualCache.attach_redis_cache` does, so a host builds the store at startup and never rebuilds it to connect
|
||||
- Where `main` swallows a Redis failure and serves memory, the family does the same and logs. It never turns a failure into a miss that `main` would not
|
||||
- Key suffixes and wire spellings derive from strum on the family's enums, and a type whose Python side is plain data implements `FromValue` / `IntoValue` over `litellm_python_compat::Value`, so the bridge never names a field or a spelling
|
||||
- Key suffixes and wire spellings derive from strum on the family's enums. A value the family stores in Redis implements `FromValue` / `IntoValue` over `litellm_python_compat::Value`. An argument or reply is a typed record, never a positional tuple: a frozen Python dataclass read in `python-bridge` by a `#[derive(FromPyObject)]` struct with `#[pyo3(attribute)]` fields and converted into the family's type with `From`, and a reply returned as a frozen `#[pyclass]`. The bridge reads an enum through its strum spelling, never a hand-written match
|
||||
- Wire contracts are hand-written cases with expected values, inlined as the same literals in this crate's test and in `tests/unit/redis/stores/`, checked here against real Redis and there against Python's own codec. Change both sides together
|
||||
- `rate_limits.rs` keeps the v3 limiter's Lua byte for byte (`rate_limits/lua.rs`) so pods on either stack share one script SHA; the admission decisions stay in the Python limiter
|
||||
- `RequestRedis` is one request's batches: a pre-call and a post-call `RedisBatch` per client, the post one flushed by `flush_post` or its deadline. Each op exists once and takes a `Send` (`Now`, `PreCall`, `PostCall`) picked by the caller, and the store calls `send.batch(&client)`. A store never encodes a phase in a method name. A post-call method exists only where a follow-up runs after the write settles (`token_increment_post_call`), as a thin wrapper over the shared op. `Send::Now` is an alone batch, so a script call outside a request goes out alone, like a redis-py `Script` outside a pipeline, so `NOSCRIPT` costs `SCRIPT LOAD` and one more `EVALSHA` the way it does on main. Only what main's Python puts on its request pipeline rides these batches. An op main sends on its own round trip, like a deployment's tpm and rpm increment, stays on its own round trip here
|
||||
|
|
|
|||
8
litellm-rust/crates/redis-store/src/auth_cache/keys.rs
Normal file
8
litellm-rust/crates/redis-store/src/auth_cache/keys.rs
Normal file
|
|
@ -0,0 +1,8 @@
|
|||
/// `is_user_key_cache_key`: user-key objects are cached under a bare `hash_token` digest, every
|
||||
/// other object under a prefixed key.
|
||||
pub fn is_key_object(key: &str) -> bool {
|
||||
key.len() == 64
|
||||
&& key
|
||||
.bytes()
|
||||
.all(|byte| byte.is_ascii_digit() || (b'a'..=b'f').contains(&byte))
|
||||
}
|
||||
258
litellm-rust/crates/redis-store/src/auth_cache/mod.rs
Normal file
258
litellm-rust/crates/redis-store/src/auth_cache/mod.rs
Normal file
|
|
@ -0,0 +1,258 @@
|
|||
mod keys;
|
||||
mod types;
|
||||
|
||||
use std::{
|
||||
sync::{Arc, Mutex},
|
||||
time::Duration,
|
||||
};
|
||||
|
||||
use litellm_python_compat::{Value, json};
|
||||
use litellm_redis_client::{Declared, RedisBatch, commands};
|
||||
|
||||
pub use self::keys::is_key_object;
|
||||
pub use self::types::Settings;
|
||||
use crate::{
|
||||
Error, Number, Pending, RequestRedis, Send,
|
||||
dual_cache::DualCache,
|
||||
memory::{Clock, lock, system_clock},
|
||||
remote::{Remote, RemoteSlot},
|
||||
request::ready,
|
||||
};
|
||||
|
||||
/// `RedisCache.async_increment_if_present`, byte for byte, sent as `EVAL` like it.
|
||||
const INCREMENT_IF_PRESENT: &str = "if redis.call('EXISTS', KEYS[1]) == 0 then return false end return redis.call('INCRBYFLOAT', KEYS[1], ARGV[1])";
|
||||
|
||||
/// The proxy's auth object cache with `enable_redis_auth_cache` on: `UserApiKeyCache`'s two
|
||||
/// `DualCache` partitions, so churn in management objects cannot evict user-key objects, in front of
|
||||
/// one Redis shared by every pod. Keys are the logical cache keys `UserApiKeyCache` callers already
|
||||
/// use; values are what `json.loads` returns, written as `json.dumps`, like `RedisCache.async_set_cache`.
|
||||
pub struct AuthCache {
|
||||
key_objects: DualCache<Value>,
|
||||
objects: DualCache<Value>,
|
||||
redis: Arc<RemoteSlot>,
|
||||
settings: Mutex<Settings>,
|
||||
}
|
||||
|
||||
impl AuthCache {
|
||||
pub fn new(remote: Option<Remote>, settings: Settings) -> Self {
|
||||
Self::with_clock(remote, settings, system_clock())
|
||||
}
|
||||
|
||||
pub fn with_clock(remote: Option<Remote>, settings: Settings, now: Clock) -> Self {
|
||||
let redis = Arc::new(RemoteSlot::new(remote));
|
||||
let partition = || DualCache::new(redis.clone(), settings.memory(), now.clone());
|
||||
Self {
|
||||
key_objects: partition(),
|
||||
objects: partition(),
|
||||
redis,
|
||||
settings: Mutex::new(settings),
|
||||
}
|
||||
}
|
||||
|
||||
pub fn has_redis(&self) -> bool {
|
||||
self.redis.get().is_some()
|
||||
}
|
||||
|
||||
pub fn settings(&self) -> Settings {
|
||||
*lock(&self.settings)
|
||||
}
|
||||
|
||||
/// `UserApiKeyCache.update_cache_ttl` and `update_in_memory_max_size` on both partitions, which
|
||||
/// keep what they hold.
|
||||
pub fn reconfigure(&self, settings: Settings) {
|
||||
*lock(&self.settings) = settings;
|
||||
self.key_objects.reconfigure(settings.memory());
|
||||
self.objects.reconfigure(settings.memory());
|
||||
}
|
||||
|
||||
pub fn get_local(&self, key: &str) -> Result<Option<Value>, Error> {
|
||||
Ok(self.partition(key).memory().get(key)?)
|
||||
}
|
||||
|
||||
/// A memory write with `ttl`, or `default_in_memory_ttl` when it names none.
|
||||
pub fn set_local(&self, key: &str, value: Value, ttl: Option<Duration>) -> Result<(), Error> {
|
||||
let partition = self.partition(key);
|
||||
Ok(partition
|
||||
.memory()
|
||||
.set(key, value, ttl.or(partition.default_in_memory_ttl()))?)
|
||||
}
|
||||
|
||||
pub fn delete_local(&self, key: &str) -> Result<(), Error> {
|
||||
Ok(self.partition(key).memory().delete(key)?)
|
||||
}
|
||||
|
||||
pub fn flush_local(&self) -> Result<(), Error> {
|
||||
self.key_objects.memory().flush()?;
|
||||
Ok(self.objects.memory().flush()?)
|
||||
}
|
||||
|
||||
/// Entries held by the management-object and the user-key partitions.
|
||||
pub fn len_local(&self) -> Result<(usize, usize), Error> {
|
||||
Ok((
|
||||
self.objects.memory().len()?,
|
||||
self.key_objects.memory().len()?,
|
||||
))
|
||||
}
|
||||
|
||||
/// `DualCache.async_get_cache`. A key `request` already read as missing from Redis is not read
|
||||
/// again.
|
||||
pub async fn get(
|
||||
&self,
|
||||
request: Option<&RequestRedis>,
|
||||
key: &str,
|
||||
) -> Result<Option<Value>, Error> {
|
||||
self.partition(key).get(request, key, None).await
|
||||
}
|
||||
|
||||
/// `DualCache.async_set_cache`, writing `json.dumps(value)`.
|
||||
pub async fn set(&self, key: &str, value: Value, ttl: Option<Duration>) -> Result<(), Error> {
|
||||
self.partition(key).set(key, value, ttl).await
|
||||
}
|
||||
|
||||
/// `DualCache.set_cache`, whose `RedisCache.set_cache` writes `str(value)`.
|
||||
pub async fn set_str(
|
||||
&self,
|
||||
key: &str,
|
||||
value: Value,
|
||||
ttl: Option<Duration>,
|
||||
) -> Result<(), Error> {
|
||||
self.partition(key).set_str(key, value, ttl).await
|
||||
}
|
||||
|
||||
/// `DualCache.async_set_cache_pre_call`: memory now, the `SET` on the batch `send` picks.
|
||||
pub fn declare_set(
|
||||
&self,
|
||||
send: Send<'_>,
|
||||
key: &str,
|
||||
value: Value,
|
||||
ttl: Option<Duration>,
|
||||
) -> Result<Pending<()>, Error> {
|
||||
self.partition(key).declare_set(send, key, value, ttl)
|
||||
}
|
||||
|
||||
/// `DualCache.async_delete_cache_pre_call`, memory now and the `DEL` on the batch `send` picks.
|
||||
pub fn declare_delete(&self, send: Send<'_>, key: &str) -> Result<Pending<()>, Error> {
|
||||
self.partition(key).declare_delete(send, key)
|
||||
}
|
||||
|
||||
/// `UserApiKeyCache.async_set_cache_pipeline`: the key-object entries in one pipeline, then
|
||||
/// the rest in another, each with one TTL.
|
||||
pub async fn set_many(
|
||||
&self,
|
||||
entries: Vec<(String, Value)>,
|
||||
ttl: Option<Duration>,
|
||||
) -> Result<(), Error> {
|
||||
let (key_objects, others): (Vec<_>, Vec<_>) =
|
||||
entries.into_iter().partition(|(key, _)| is_key_object(key));
|
||||
self.key_objects.set_pipeline(key_objects, ttl).await?;
|
||||
self.objects.set_pipeline(others, ttl).await
|
||||
}
|
||||
|
||||
/// The auth prefetch write-back: memory first, an entry without a TTL taking
|
||||
/// `default_in_memory_ttl`, then `RedisCache.async_set_cache_pipeline_with_ttls`, which encodes
|
||||
/// every value before it sends anything and gives an entry without a TTL the connection's
|
||||
/// default. On a request batch the `SET`s are declared and nothing is left to wait for.
|
||||
/// [`Send::Now`] is one pipeline of its own rather than one round trip per entry.
|
||||
pub fn write_back(
|
||||
&self,
|
||||
send: Send<'_>,
|
||||
entries: Vec<(String, Value, Option<Duration>)>,
|
||||
) -> Result<Pending<()>, Error> {
|
||||
let encoded = self
|
||||
.redis
|
||||
.get()
|
||||
.map(|_| {
|
||||
entries
|
||||
.iter()
|
||||
.map(|(_, value, _)| json::dumps(value))
|
||||
.collect::<Result<Vec<_>, _>>()
|
||||
})
|
||||
.transpose();
|
||||
let writes = entries
|
||||
.into_iter()
|
||||
.map(|(key, value, ttl)| {
|
||||
self.set_local(&key, value, ttl)?;
|
||||
Ok((key, ttl))
|
||||
})
|
||||
.collect::<Result<Vec<_>, Error>>()?;
|
||||
let (Some(remote), Some(encoded)) = (self.redis.get(), encoded?) else {
|
||||
return Ok(ready(()));
|
||||
};
|
||||
let sets = writes.iter().zip(encoded).map(|((key, ttl), value)| {
|
||||
commands::set(key, value, ttl.unwrap_or(remote.default_ttl))
|
||||
});
|
||||
let Send::Now = send else {
|
||||
let batch = send.batch(&remote.client);
|
||||
sets.for_each(|set| drop(batch.add(set)));
|
||||
return Ok(ready(()));
|
||||
};
|
||||
let pipeline = RedisBatch::pipelined(remote.client.clone());
|
||||
let written: Vec<_> = sets.map(|set| pipeline.add(set)).collect();
|
||||
Ok(Box::pin(async move {
|
||||
Declared::all(written)
|
||||
.await
|
||||
.into_iter()
|
||||
.try_for_each(|reply| Ok(reply?))
|
||||
}))
|
||||
}
|
||||
|
||||
/// `UserApiKeyCache.async_delete_cache_keys`: the key-object partition and the rest as two
|
||||
/// concurrent `DualCache.async_delete_cache_keys`. Both run even when one fails.
|
||||
pub async fn delete(&self, keys: &[String]) -> Result<(), Error> {
|
||||
let (key_objects, others): (Vec<String>, Vec<String>) =
|
||||
keys.iter().cloned().partition(|key| is_key_object(key));
|
||||
let (key_objects, others) = tokio::join!(
|
||||
self.key_objects.delete_keys(&key_objects),
|
||||
self.objects.delete_keys(&others)
|
||||
);
|
||||
key_objects.and(others)
|
||||
}
|
||||
|
||||
/// `RedisCache.async_batch_get_cache` for keys the caller found missing in memory, one `MGET` on
|
||||
/// the batch `send` picks, noting the keys Redis lacks as the request's misses. `None` without
|
||||
/// Redis. Any failed read or undecodable value fails the whole read, as a failed `MGET` does.
|
||||
pub fn get_many(&self, send: Send<'_>, keys: &[String]) -> Option<Pending<Vec<Option<Value>>>> {
|
||||
let client = self.redis.client()?;
|
||||
let misses = send.misses(client);
|
||||
let read = send.batch(client).add(commands::mget::<Vec<u8>>(
|
||||
&keys.iter().map(String::as_str).collect::<Vec<_>>(),
|
||||
));
|
||||
let keys = keys.to_vec();
|
||||
Some(Box::pin(async move {
|
||||
let raw = read.resolve().await?;
|
||||
if let Some(misses) = misses {
|
||||
misses.note(
|
||||
keys.into_iter()
|
||||
.zip(&raw)
|
||||
.filter(|(_, raw)| raw.is_none())
|
||||
.map(|(key, _)| key),
|
||||
);
|
||||
}
|
||||
raw.into_iter().map(|raw| Ok(types::decode(raw)?)).collect()
|
||||
}))
|
||||
}
|
||||
|
||||
/// `RedisCache.async_increment_if_present`: `INCRBYFLOAT` only an entry Redis still holds.
|
||||
/// `None` without Redis or for an absent entry.
|
||||
pub async fn increment_if_present(
|
||||
&self,
|
||||
key: &str,
|
||||
amount: Number,
|
||||
) -> Result<Option<f64>, Error> {
|
||||
let Some(client) = self.redis.client() else {
|
||||
return Ok(None);
|
||||
};
|
||||
let amount = litellm_redis_client::Amount::from(amount);
|
||||
Ok(client
|
||||
.send(commands::eval_source(INCREMENT_IF_PRESENT, &[key], amount))
|
||||
.await?)
|
||||
}
|
||||
|
||||
fn partition(&self, key: &str) -> &DualCache<Value> {
|
||||
if is_key_object(key) {
|
||||
&self.key_objects
|
||||
} else {
|
||||
&self.objects
|
||||
}
|
||||
}
|
||||
}
|
||||
35
litellm-rust/crates/redis-store/src/auth_cache/types.rs
Normal file
35
litellm-rust/crates/redis-store/src/auth_cache/types.rs
Normal file
|
|
@ -0,0 +1,35 @@
|
|||
use std::time::Duration;
|
||||
|
||||
use litellm_python_compat::Value;
|
||||
|
||||
use crate::{DecodeError, codec::python_value, dual_cache::MemorySettings};
|
||||
|
||||
/// `UserApiKeyCache`'s memory configuration: each partition's size and `default_in_memory_ttl`.
|
||||
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
|
||||
pub struct Settings {
|
||||
pub memory_entries: usize,
|
||||
pub memory_ttl: Option<Duration>,
|
||||
}
|
||||
|
||||
impl Default for Settings {
|
||||
/// `InMemoryCache`'s 200 entries.
|
||||
fn default() -> Self {
|
||||
Self {
|
||||
memory_entries: 200,
|
||||
memory_ttl: None,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl Settings {
|
||||
pub(super) fn memory(self) -> MemorySettings {
|
||||
MemorySettings {
|
||||
entries: Some(self.memory_entries),
|
||||
ttl: self.memory_ttl,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub(super) fn decode(raw: Option<Vec<u8>>) -> Result<Option<Value>, DecodeError> {
|
||||
raw.as_deref().map(python_value).transpose()
|
||||
}
|
||||
|
|
@ -1,23 +1,17 @@
|
|||
#![expect(
|
||||
dead_code,
|
||||
reason = "the store families built on these tiers and codecs land in the next PR"
|
||||
)]
|
||||
#![cfg_attr(
|
||||
not(test),
|
||||
expect(
|
||||
unused_imports,
|
||||
reason = "the store families built on these tiers and codecs land in the next PR"
|
||||
)
|
||||
)]
|
||||
|
||||
pub mod auth_cache;
|
||||
mod codec;
|
||||
mod dual_cache;
|
||||
mod error;
|
||||
mod failure_log;
|
||||
mod memory;
|
||||
mod number;
|
||||
pub mod rate_limits;
|
||||
mod remote;
|
||||
mod request;
|
||||
pub mod router_cooldowns;
|
||||
pub mod router_usage;
|
||||
pub mod sensitive_route_pins;
|
||||
pub mod spend_counters;
|
||||
mod value;
|
||||
|
||||
pub use codec::DecodeError;
|
||||
|
|
|
|||
74
litellm-rust/crates/redis-store/src/rate_limits/keys.rs
Normal file
74
litellm-rust/crates/redis-store/src/rate_limits/keys.rs
Normal file
|
|
@ -0,0 +1,74 @@
|
|||
/// A rate-limit descriptor, `{descriptor:value}`, the hash tag every key of its window shares.
|
||||
#[derive(Clone, Debug, PartialEq, Eq, Hash)]
|
||||
pub struct Scope {
|
||||
pub descriptor: String,
|
||||
pub value: String,
|
||||
}
|
||||
|
||||
impl Scope {
|
||||
pub(super) fn key(&self, suffix: &str) -> String {
|
||||
format!("{{{}:{}}}:{suffix}", self.descriptor, self.value)
|
||||
}
|
||||
|
||||
pub(super) fn window_key(&self) -> String {
|
||||
self.key(WINDOW)
|
||||
}
|
||||
|
||||
pub(super) fn gauge_key(&self) -> String {
|
||||
self.key(GAUGE)
|
||||
}
|
||||
}
|
||||
|
||||
const WINDOW: &str = "window";
|
||||
const GAUGE: &str = "max_parallel_requests";
|
||||
|
||||
/// A counter's key suffix, which is also how the Python limiter names it.
|
||||
#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, strum::EnumString, strum::IntoStaticStr)]
|
||||
pub enum Kind {
|
||||
#[strum(serialize = "requests")]
|
||||
Requests,
|
||||
#[strum(serialize = "tokens")]
|
||||
Tokens,
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, PartialEq, Eq, Hash)]
|
||||
pub struct Counter {
|
||||
pub scope: Scope,
|
||||
pub kind: Kind,
|
||||
}
|
||||
|
||||
impl Counter {
|
||||
pub(super) fn key(&self) -> String {
|
||||
self.scope.key(self.kind.into())
|
||||
}
|
||||
}
|
||||
|
||||
/// One value of the limiter's keyspace, as both tiers address it.
|
||||
#[derive(Clone, Debug, PartialEq, Eq, Hash)]
|
||||
pub enum Entry {
|
||||
Window(Scope),
|
||||
Counter(Counter),
|
||||
Gauge(Scope),
|
||||
}
|
||||
|
||||
impl Entry {
|
||||
/// The entry the Python limiter names by its key suffix: `window`, `max_parallel_requests`, or a
|
||||
/// counter's `Kind`.
|
||||
pub fn new(suffix: &str, scope: Scope) -> Result<Self, strum::ParseError> {
|
||||
match suffix {
|
||||
WINDOW => Ok(Self::Window(scope)),
|
||||
GAUGE => Ok(Self::Gauge(scope)),
|
||||
kind => kind
|
||||
.parse()
|
||||
.map(|kind| Self::Counter(Counter { scope, kind })),
|
||||
}
|
||||
}
|
||||
|
||||
pub(super) fn key(&self) -> String {
|
||||
match self {
|
||||
Self::Window(scope) => scope.window_key(),
|
||||
Self::Counter(counter) => counter.key(),
|
||||
Self::Gauge(scope) => scope.gauge_key(),
|
||||
}
|
||||
}
|
||||
}
|
||||
260
litellm-rust/crates/redis-store/src/rate_limits/lua.rs
Normal file
260
litellm-rust/crates/redis-store/src/rate_limits/lua.rs
Normal file
|
|
@ -0,0 +1,260 @@
|
|||
//! The limiter's Lua, byte for byte as `main`'s Python registers it, so both share one script SHA.
|
||||
|
||||
pub(super) const BATCH_RATE_LIMITER: &str = r#"
|
||||
local results = {}
|
||||
local now = tonumber(ARGV[1])
|
||||
local window_size = tonumber(ARGV[2])
|
||||
|
||||
-- Process each window/counter pair
|
||||
for i = 1, #KEYS, 2 do
|
||||
local window_key = KEYS[i]
|
||||
local counter_key = KEYS[i + 1]
|
||||
local increment_value = 1
|
||||
|
||||
-- Check if window exists and is valid
|
||||
local window_start = redis.call('GET', window_key)
|
||||
if not window_start or (now - tonumber(window_start)) >= window_size then
|
||||
-- Reset window and counter
|
||||
local prefix = string.sub(window_key, 1, -(#':window') - 1)
|
||||
redis.call('DEL', prefix .. ':requests', prefix .. ':tokens')
|
||||
redis.call('SET', window_key, tostring(now))
|
||||
redis.call('SET', counter_key, increment_value)
|
||||
redis.call('EXPIRE', window_key, window_size)
|
||||
redis.call('EXPIRE', counter_key, window_size)
|
||||
table.insert(results, tostring(now)) -- window_start
|
||||
table.insert(results, increment_value) -- counter
|
||||
else
|
||||
local counter = redis.call('INCR', counter_key)
|
||||
-- This happens when window_key exists but counter_key doesn't (e.g., tokens key
|
||||
-- created after requests key when both share the same window_key)
|
||||
local current_ttl = redis.call('TTL', counter_key)
|
||||
if current_ttl == -1 then
|
||||
redis.call('EXPIRE', counter_key, window_size)
|
||||
end
|
||||
table.insert(results, window_start) -- window_start
|
||||
table.insert(results, counter) -- counter
|
||||
end
|
||||
end
|
||||
|
||||
return results
|
||||
"#;
|
||||
|
||||
pub(super) const BATCH_COUNTER_READ: &str = r#"return redis.call('MGET', unpack(KEYS))"#;
|
||||
|
||||
pub(super) const CHECK_AND_INCREMENT_BY_N: &str = r#"
|
||||
-- Atomic check-and-increment-by-N across one or more descriptors.
|
||||
-- All-or-nothing: if any descriptor would exceed its limit, no counter is
|
||||
-- modified.
|
||||
--
|
||||
-- Uses Redis server time (`redis.call('TIME')`) instead of a client-supplied
|
||||
-- timestamp so that window resets are deterministic across replicas with
|
||||
-- skewed wall-clocks. This prevents a clock-skew-induced reopening of the
|
||||
-- TOCTOU window across multi-replica deployments.
|
||||
--
|
||||
-- KEYS layout: pairs of (window_key, counter_key), one pair per descriptor.
|
||||
-- ARGV layout: per-descriptor 4-tuple, starting at ARGV[1]:
|
||||
-- ARGV[(i-1)*4 + 1] = limit
|
||||
-- ARGV[(i-1)*4 + 2] = increment
|
||||
-- ARGV[(i-1)*4 + 3] = ttl_seconds (counter TTL when window resets)
|
||||
-- ARGV[(i-1)*4 + 4] = window_size_seconds (sliding-window length)
|
||||
--
|
||||
-- Return on success:
|
||||
-- { 0, new_counter_1, window_start_1, new_counter_2, window_start_2, ... }
|
||||
-- Return on over-limit: { 1, descriptor_index, current_counter, limit }
|
||||
local time_reply = redis.call('TIME')
|
||||
local now = tonumber(time_reply[1])
|
||||
local descriptor_count = #KEYS / 2
|
||||
local reset_windows = {}
|
||||
|
||||
-- Pass 1: read state, validate. Abort without writing if any over limit.
|
||||
local descriptor_state = {}
|
||||
for i = 1, descriptor_count do
|
||||
local window_key = KEYS[(i - 1) * 2 + 1]
|
||||
local counter_key = KEYS[(i - 1) * 2 + 2]
|
||||
local arg_base = (i - 1) * 4 + 1
|
||||
local limit = tonumber(ARGV[arg_base])
|
||||
local increment = tonumber(ARGV[arg_base + 1])
|
||||
local window_size = tonumber(ARGV[arg_base + 3])
|
||||
|
||||
local window_start = redis.call('GET', window_key)
|
||||
local window_expired = (not window_start) or
|
||||
((now - tonumber(window_start)) >= window_size)
|
||||
|
||||
local current_counter
|
||||
if window_expired then
|
||||
current_counter = 0
|
||||
else
|
||||
current_counter = tonumber(redis.call('GET', counter_key) or 0)
|
||||
end
|
||||
|
||||
local blocked
|
||||
if increment > 0 then
|
||||
blocked = current_counter + increment > limit
|
||||
else
|
||||
blocked = current_counter >= limit
|
||||
end
|
||||
if blocked then
|
||||
return { 1, i, current_counter, limit }
|
||||
end
|
||||
|
||||
descriptor_state[i] = { window_expired, current_counter, window_start }
|
||||
end
|
||||
|
||||
-- Pass 2: all checks passed. Apply increments.
|
||||
local results = { 0 }
|
||||
for i = 1, descriptor_count do
|
||||
local window_key = KEYS[(i - 1) * 2 + 1]
|
||||
local counter_key = KEYS[(i - 1) * 2 + 2]
|
||||
local arg_base = (i - 1) * 4 + 1
|
||||
local increment = tonumber(ARGV[arg_base + 1])
|
||||
local ttl = tonumber(ARGV[arg_base + 2])
|
||||
local window_size = tonumber(ARGV[arg_base + 3])
|
||||
|
||||
local window_expired = descriptor_state[i][1]
|
||||
local active_window_start
|
||||
|
||||
if window_expired then
|
||||
active_window_start = now
|
||||
if not reset_windows[window_key] then
|
||||
local prefix = string.sub(window_key, 1, -(#':window') - 1)
|
||||
redis.call('DEL', prefix .. ':requests', prefix .. ':tokens')
|
||||
reset_windows[window_key] = true
|
||||
end
|
||||
redis.call('SET', window_key, tostring(now))
|
||||
redis.call('SET', counter_key, increment)
|
||||
redis.call('EXPIRE', window_key, window_size)
|
||||
if ttl > 0 then
|
||||
redis.call('EXPIRE', counter_key, ttl)
|
||||
end
|
||||
table.insert(results, increment)
|
||||
else
|
||||
active_window_start = tonumber(descriptor_state[i][3])
|
||||
local new_counter = redis.call('INCRBY', counter_key, increment)
|
||||
local current_ttl = redis.call('TTL', counter_key)
|
||||
if current_ttl == -1 and ttl > 0 then
|
||||
redis.call('EXPIRE', counter_key, ttl)
|
||||
end
|
||||
table.insert(results, new_counter)
|
||||
end
|
||||
table.insert(results, active_window_start)
|
||||
end
|
||||
|
||||
return results
|
||||
"#;
|
||||
|
||||
pub(super) const WINDOW_GUARDED_TOKEN_INCREMENT: &str = r#"
|
||||
local results = {}
|
||||
for i = 1, #KEYS, 2 do
|
||||
local window_key = KEYS[i]
|
||||
local counter_key = KEYS[i + 1]
|
||||
local arg_base = ((i - 1) / 2) * 3 + 1
|
||||
local expected_window_start = ARGV[arg_base]
|
||||
local increment = tonumber(ARGV[arg_base + 1])
|
||||
local ttl = tonumber(ARGV[arg_base + 2])
|
||||
local active_window_start = redis.call('GET', window_key)
|
||||
|
||||
if active_window_start and active_window_start == expected_window_start then
|
||||
local new_counter = redis.call('INCRBY', counter_key, increment)
|
||||
local current_ttl = redis.call('TTL', counter_key)
|
||||
if current_ttl == -1 and ttl > 0 then
|
||||
redis.call('EXPIRE', counter_key, ttl)
|
||||
end
|
||||
table.insert(results, 1)
|
||||
table.insert(results, new_counter)
|
||||
else
|
||||
table.insert(results, 0)
|
||||
table.insert(results, tonumber(redis.call('GET', counter_key) or 0))
|
||||
end
|
||||
end
|
||||
return results
|
||||
"#;
|
||||
|
||||
pub(super) const PARALLEL_ACQUIRE: &str = r#"
|
||||
-- Atomic check-and-acquire for the max_parallel_requests concurrency gauge.
|
||||
-- Each gauge key is a sorted set of per-request slot ids scored by acquire
|
||||
-- time (Redis server clock). In-flight requests are counted by ZCARD after
|
||||
-- pruning slots older than the slot TTL, so unlike the windowed RPM/TPM
|
||||
-- counters the gauge is never reset while requests are in flight, a
|
||||
-- rejected request never occupies a slot, and a slot leaked by a crashed
|
||||
-- worker self-heals after the slot TTL even under continuous traffic.
|
||||
--
|
||||
-- KEYS: one gauge zset key per descriptor.
|
||||
-- ARGV: per-key triples (limit, slot_ttl_seconds, slot_id).
|
||||
-- Success: { 0, in_flight_1, ... }. Over-limit: { 1, key_index, in_flight, limit }.
|
||||
local time_reply = redis.call('TIME')
|
||||
local now = tonumber(time_reply[1])
|
||||
for i = 1, #KEYS do
|
||||
local limit = tonumber(ARGV[(i - 1) * 3 + 1])
|
||||
local slot_ttl = tonumber(ARGV[(i - 1) * 3 + 2])
|
||||
redis.call('ZREMRANGEBYSCORE', KEYS[i], '-inf', now - slot_ttl)
|
||||
local in_flight = redis.call('ZCARD', KEYS[i])
|
||||
if in_flight + 1 > limit then
|
||||
return { 1, i, in_flight, limit }
|
||||
end
|
||||
end
|
||||
local results = { 0 }
|
||||
for i = 1, #KEYS do
|
||||
local slot_ttl = tonumber(ARGV[(i - 1) * 3 + 2])
|
||||
local slot_id = ARGV[(i - 1) * 3 + 3]
|
||||
redis.call('ZADD', KEYS[i], now, slot_id)
|
||||
redis.call('EXPIRE', KEYS[i], slot_ttl)
|
||||
table.insert(results, redis.call('ZCARD', KEYS[i]))
|
||||
end
|
||||
return results
|
||||
"#;
|
||||
|
||||
pub(super) const PARALLEL_RELEASE: &str = r#"
|
||||
-- Release one slot per gauge key by removing this request's slot id.
|
||||
-- ZREM of an absent member (or key) is a no-op, so a release without a
|
||||
-- matching acquire (proxy-side rejection, double-fired callback, slot
|
||||
-- already expired) can never free a slot owned by another request.
|
||||
-- KEYS: gauge zset keys. ARGV: per-key slot_id.
|
||||
-- Returns the remaining in-flight count per key.
|
||||
local results = {}
|
||||
for i = 1, #KEYS do
|
||||
redis.call('ZREM', KEYS[i], ARGV[i])
|
||||
table.insert(results, redis.call('ZCARD', KEYS[i]))
|
||||
end
|
||||
return results
|
||||
"#;
|
||||
|
||||
pub(super) const PARALLEL_COUNT: &str = r#"
|
||||
-- Read the current in-flight count per gauge key (prunes expired slots
|
||||
-- first so leaked slots do not inflate the reading).
|
||||
-- KEYS: gauge zset keys. ARGV: per-key slot_ttl_seconds.
|
||||
local time_reply = redis.call('TIME')
|
||||
local now = tonumber(time_reply[1])
|
||||
local results = {}
|
||||
for i = 1, #KEYS do
|
||||
redis.call('ZREMRANGEBYSCORE', KEYS[i], '-inf', now - tonumber(ARGV[i]))
|
||||
table.insert(results, redis.call('ZCARD', KEYS[i]))
|
||||
end
|
||||
return results
|
||||
"#;
|
||||
|
||||
pub(super) const TOKEN_INCREMENT: &str = r#"
|
||||
local results = {}
|
||||
|
||||
-- Process each key/increment_value/ttl triplet
|
||||
for i = 1, #KEYS do
|
||||
local key = KEYS[i]
|
||||
local increment_value = tonumber(ARGV[i * 2 - 1])
|
||||
local ttl_seconds = tonumber(ARGV[i * 2])
|
||||
|
||||
-- Increment the value
|
||||
local new_value = redis.call('INCRBYFLOAT', key, increment_value)
|
||||
|
||||
-- Handle TTL: only set expire if ttl_seconds > 0 and key has no current TTL
|
||||
-- ttl_seconds can be 0 (no TTL) or positive (set TTL)
|
||||
if ttl_seconds and ttl_seconds > 0 then
|
||||
local current_ttl = redis.call('TTL', key)
|
||||
if current_ttl == -1 then
|
||||
redis.call('EXPIRE', key, ttl_seconds)
|
||||
end
|
||||
end
|
||||
|
||||
table.insert(results, new_value)
|
||||
end
|
||||
|
||||
return results
|
||||
"#;
|
||||
304
litellm-rust/crates/redis-store/src/rate_limits/mod.rs
Normal file
304
litellm-rust/crates/redis-store/src/rate_limits/mod.rs
Normal file
|
|
@ -0,0 +1,304 @@
|
|||
mod keys;
|
||||
mod lua;
|
||||
mod scripts;
|
||||
mod types;
|
||||
|
||||
use std::{sync::Arc, time::Duration};
|
||||
|
||||
use litellm_redis_client::{Declared, RedisBatch, RedisClient, Slot};
|
||||
|
||||
pub use self::keys::{Counter, Entry, Kind, Scope};
|
||||
pub use self::types::{
|
||||
AcquireReply, Applied, AtomicCounter, AtomicReply, Gauge, GuardedIncrement, GuardedReply,
|
||||
Increment, Local, WindowCount,
|
||||
};
|
||||
use crate::{
|
||||
Error, Number, Pending, RequestRedis, Send,
|
||||
dual_cache::{DualCache, MemorySettings, PipelineIncrement},
|
||||
memory::{Clock, system_clock},
|
||||
remote::{Remote, RemoteSlot},
|
||||
request::pending,
|
||||
};
|
||||
|
||||
/// The v3 rate limiter's keyspace: its memory tier, its Lua scripts and the plain counter writes
|
||||
/// `DualCache` made for it on `main`. The admission logic stays with the caller. A script needs
|
||||
/// Redis, so without one it declares nothing and answers `None`.
|
||||
pub struct RateLimits {
|
||||
cache: DualCache<Local>,
|
||||
}
|
||||
|
||||
impl RateLimits {
|
||||
/// `backfill_ttl` is the `default_in_memory_ttl` of the `DualCache` this tier replaces. An
|
||||
/// increment without a TTL takes the remote's default.
|
||||
pub fn new(remote: Option<Remote>, backfill_ttl: Option<Duration>) -> Self {
|
||||
Self::with_clock(remote, backfill_ttl, system_clock())
|
||||
}
|
||||
|
||||
pub fn with_clock(remote: Option<Remote>, backfill_ttl: Option<Duration>, now: Clock) -> Self {
|
||||
Self {
|
||||
cache: DualCache::new(
|
||||
Arc::new(RemoteSlot::new(remote)),
|
||||
MemorySettings {
|
||||
entries: None,
|
||||
ttl: backfill_ttl,
|
||||
},
|
||||
now,
|
||||
),
|
||||
}
|
||||
}
|
||||
|
||||
pub fn has_redis(&self) -> bool {
|
||||
self.cache.redis().get().is_some()
|
||||
}
|
||||
|
||||
/// `_group_keys_by_hash_tag`: indices of `scopes` sharing a cluster slot, in first-seen order,
|
||||
/// or one group when the client is not a cluster and grouping is not forced.
|
||||
pub fn slot_groups(&self, scopes: &[Scope], force: bool) -> Vec<Vec<usize>> {
|
||||
let by_slot = force
|
||||
|| self
|
||||
.cache
|
||||
.redis()
|
||||
.client()
|
||||
.is_some_and(RedisClient::is_cluster);
|
||||
if !by_slot {
|
||||
return vec![(0..scopes.len()).collect()];
|
||||
}
|
||||
scopes
|
||||
.iter()
|
||||
.enumerate()
|
||||
.fold(
|
||||
Vec::<(Slot, Vec<usize>)>::new(),
|
||||
|mut groups, (index, scope)| {
|
||||
let slot = Slot::for_key(scope.window_key());
|
||||
match groups.iter_mut().find(|(seen, _)| *seen == slot) {
|
||||
Some((_, members)) => members.push(index),
|
||||
None => groups.push((slot, vec![index])),
|
||||
}
|
||||
groups
|
||||
},
|
||||
)
|
||||
.into_iter()
|
||||
.map(|(_, members)| members)
|
||||
.collect()
|
||||
}
|
||||
|
||||
pub fn get_local(&self, entry: &Entry) -> Result<Option<Local>, Error> {
|
||||
Ok(self.cache.memory().get(&entry.key())?)
|
||||
}
|
||||
|
||||
pub fn set_local(
|
||||
&self,
|
||||
entry: &Entry,
|
||||
value: Local,
|
||||
ttl: Option<Duration>,
|
||||
) -> Result<(), Error> {
|
||||
Ok(self.cache.memory().set(&entry.key(), value, ttl)?)
|
||||
}
|
||||
|
||||
/// `DualCache.async_get_cache`, backfilling memory for `backfill_ttl` or the store's default.
|
||||
/// A key `request` already read as missing from Redis is not read again.
|
||||
pub async fn get(
|
||||
&self,
|
||||
request: Option<&RequestRedis>,
|
||||
entry: &Entry,
|
||||
backfill_ttl: Option<Duration>,
|
||||
) -> Result<Option<Local>, Error> {
|
||||
self.cache.get(request, &entry.key(), backfill_ttl).await
|
||||
}
|
||||
|
||||
pub fn sliding_window(
|
||||
&self,
|
||||
send: Send<'_>,
|
||||
counters: &[Counter],
|
||||
now: i64,
|
||||
window_size: i64,
|
||||
) -> Option<Pending<Vec<WindowCount>>> {
|
||||
Some(pending(scripts::sliding_window(
|
||||
&self.batch(send)?,
|
||||
counters,
|
||||
now,
|
||||
window_size,
|
||||
)))
|
||||
}
|
||||
|
||||
pub fn read_counters(
|
||||
&self,
|
||||
send: Send<'_>,
|
||||
counters: &[Counter],
|
||||
) -> Option<Pending<Vec<Option<String>>>> {
|
||||
Some(pending(scripts::read_counters(
|
||||
&self.batch(send)?,
|
||||
counters,
|
||||
)))
|
||||
}
|
||||
|
||||
pub fn check_and_increment(
|
||||
&self,
|
||||
send: Send<'_>,
|
||||
counters: &[AtomicCounter],
|
||||
) -> Option<Pending<AtomicReply>> {
|
||||
Some(pending(scripts::check_and_increment(
|
||||
&self.batch(send)?,
|
||||
counters,
|
||||
)))
|
||||
}
|
||||
|
||||
pub fn guarded_increment(
|
||||
&self,
|
||||
send: Send<'_>,
|
||||
increment: &GuardedIncrement,
|
||||
) -> Option<Pending<GuardedReply>> {
|
||||
Some(pending(scripts::guarded_increment(
|
||||
&self.batch(send)?,
|
||||
increment,
|
||||
)))
|
||||
}
|
||||
|
||||
pub fn parallel_acquire(
|
||||
&self,
|
||||
send: Send<'_>,
|
||||
gauges: &[Gauge],
|
||||
slot_ttl: Duration,
|
||||
slot_id: &str,
|
||||
) -> Option<Pending<AcquireReply>> {
|
||||
Some(pending(scripts::parallel_acquire(
|
||||
&self.batch(send)?,
|
||||
gauges,
|
||||
slot_ttl,
|
||||
slot_id,
|
||||
)))
|
||||
}
|
||||
|
||||
pub fn parallel_release(
|
||||
&self,
|
||||
send: Send<'_>,
|
||||
gauges: &[Scope],
|
||||
slot_id: &str,
|
||||
) -> Option<Pending<Vec<i64>>> {
|
||||
Some(pending(self.declare_release(send, gauges, slot_id)?))
|
||||
}
|
||||
|
||||
pub fn parallel_count(
|
||||
&self,
|
||||
send: Send<'_>,
|
||||
gauges: &[Scope],
|
||||
slot_ttl: Duration,
|
||||
) -> Option<Pending<Vec<i64>>> {
|
||||
Some(pending(scripts::parallel_count(
|
||||
&self.batch(send)?,
|
||||
gauges,
|
||||
slot_ttl,
|
||||
)))
|
||||
}
|
||||
|
||||
pub fn token_increment(
|
||||
&self,
|
||||
send: Send<'_>,
|
||||
increments: &[Increment],
|
||||
) -> Option<Pending<Vec<f64>>> {
|
||||
Some(pending(self.declare_token_increment(send, increments)?))
|
||||
}
|
||||
|
||||
/// `TOKEN_INCREMENT_SCRIPT` on the request's post-call batch. A failed script falls back to
|
||||
/// `increment_many` for these increments once the batch has gone out. `false` without Redis.
|
||||
pub fn token_increment_post_call(
|
||||
self: &Arc<Self>,
|
||||
request: &RequestRedis,
|
||||
increments: Vec<Increment>,
|
||||
) -> bool {
|
||||
let Some(declared) = self.declare_token_increment(Send::PostCall(request), &increments)
|
||||
else {
|
||||
return false;
|
||||
};
|
||||
let store = Arc::clone(self);
|
||||
request.spawn(async move {
|
||||
let Err(error) = declared.settled().await else {
|
||||
return;
|
||||
};
|
||||
crate::failure_log::redis_failure!(
|
||||
&error,
|
||||
"TTL preservation failed, falling back to regular pipeline"
|
||||
);
|
||||
store.increment_many(&increments).await;
|
||||
});
|
||||
true
|
||||
}
|
||||
|
||||
/// `PARALLEL_RELEASE_SCRIPT` on the request's post-call batch; the caller has already freed
|
||||
/// the slot in memory. `false` without Redis.
|
||||
pub fn parallel_release_post_call(
|
||||
&self,
|
||||
request: &RequestRedis,
|
||||
gauges: &[Scope],
|
||||
slot_id: &str,
|
||||
) -> bool {
|
||||
let Some(declared) = self.declare_release(Send::PostCall(request), gauges, slot_id) else {
|
||||
return false;
|
||||
};
|
||||
request.spawn(async move {
|
||||
if let Err(error) = declared.settled().await {
|
||||
crate::failure_log::redis_failure!(
|
||||
&error,
|
||||
"parallel_release_script failed, the slot stays released in memory only"
|
||||
);
|
||||
}
|
||||
});
|
||||
true
|
||||
}
|
||||
|
||||
/// `DualCache.async_increment_cache`. A Redis failure answers with the memory count, and an
|
||||
/// entry memory cannot add to with `None`.
|
||||
pub async fn increment(&self, increment: &Increment) -> Option<Number> {
|
||||
self.cache
|
||||
.increment(&increment.counter.key(), increment.amount, increment.ttl)
|
||||
.await
|
||||
}
|
||||
|
||||
/// `DualCache.async_increment_cache_pipeline`, whose failures only log.
|
||||
pub async fn increment_many(&self, increments: &[Increment]) {
|
||||
self.cache.increment_pipeline(&pipeline(increments)).await;
|
||||
}
|
||||
|
||||
/// `DualCache.async_increment_cache_pipeline_post_call`: memory now, the Redis increments on
|
||||
/// the request's post-call batch.
|
||||
pub fn increment_many_post_call(&self, request: &RequestRedis, increments: &[Increment]) {
|
||||
self.cache
|
||||
.increment_pipeline_post_call(request, &pipeline(increments));
|
||||
}
|
||||
|
||||
fn declare_release(
|
||||
&self,
|
||||
send: Send<'_>,
|
||||
gauges: &[Scope],
|
||||
slot_id: &str,
|
||||
) -> Option<Declared<Vec<i64>>> {
|
||||
Some(scripts::parallel_release(
|
||||
&self.batch(send)?,
|
||||
gauges,
|
||||
slot_id,
|
||||
))
|
||||
}
|
||||
|
||||
fn declare_token_increment(
|
||||
&self,
|
||||
send: Send<'_>,
|
||||
increments: &[Increment],
|
||||
) -> Option<Declared<Vec<f64>>> {
|
||||
Some(scripts::token_increment(&self.batch(send)?, increments))
|
||||
}
|
||||
|
||||
fn batch(&self, send: Send<'_>) -> Option<RedisBatch> {
|
||||
self.cache.redis().batch(send)
|
||||
}
|
||||
}
|
||||
|
||||
fn pipeline(increments: &[Increment]) -> Vec<PipelineIncrement> {
|
||||
increments
|
||||
.iter()
|
||||
.map(|increment| PipelineIncrement {
|
||||
key: increment.counter.key(),
|
||||
amount: increment.amount,
|
||||
ttl: increment.ttl,
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
303
litellm-rust/crates/redis-store/src/rate_limits/scripts.rs
Normal file
303
litellm-rust/crates/redis-store/src/rate_limits/scripts.rs
Normal file
|
|
@ -0,0 +1,303 @@
|
|||
use std::{sync::LazyLock, time::Duration};
|
||||
|
||||
use litellm_python_compat::repr::float_repr;
|
||||
use litellm_redis_client::{
|
||||
Command, Declared, Error, RedisBatch, Script, Value,
|
||||
commands::{self, unexpected_reply},
|
||||
};
|
||||
|
||||
use super::{
|
||||
keys::{Counter, Scope},
|
||||
lua,
|
||||
types::{
|
||||
AcquireReply, Applied, AtomicCounter, AtomicReply, Gauge, GuardedIncrement, GuardedReply,
|
||||
Increment, WindowCount,
|
||||
},
|
||||
};
|
||||
use crate::Number;
|
||||
|
||||
static BATCH_RATE_LIMITER: LazyLock<Script> =
|
||||
LazyLock::new(|| Script::new(lua::BATCH_RATE_LIMITER));
|
||||
static BATCH_COUNTER_READ: LazyLock<Script> =
|
||||
LazyLock::new(|| Script::new(lua::BATCH_COUNTER_READ));
|
||||
static CHECK_AND_INCREMENT_BY_N: LazyLock<Script> =
|
||||
LazyLock::new(|| Script::new(lua::CHECK_AND_INCREMENT_BY_N));
|
||||
static WINDOW_GUARDED_TOKEN_INCREMENT: LazyLock<Script> =
|
||||
LazyLock::new(|| Script::new(lua::WINDOW_GUARDED_TOKEN_INCREMENT));
|
||||
static PARALLEL_ACQUIRE: LazyLock<Script> = LazyLock::new(|| Script::new(lua::PARALLEL_ACQUIRE));
|
||||
static PARALLEL_RELEASE: LazyLock<Script> = LazyLock::new(|| Script::new(lua::PARALLEL_RELEASE));
|
||||
static PARALLEL_COUNT: LazyLock<Script> = LazyLock::new(|| Script::new(lua::PARALLEL_COUNT));
|
||||
static TOKEN_INCREMENT: LazyLock<Script> = LazyLock::new(|| Script::new(lua::TOKEN_INCREMENT));
|
||||
|
||||
pub(super) fn sliding_window(
|
||||
batch: &RedisBatch,
|
||||
counters: &[Counter],
|
||||
now: i64,
|
||||
window_size: i64,
|
||||
) -> Declared<Vec<WindowCount>> {
|
||||
let keys = counters
|
||||
.iter()
|
||||
.flat_map(|counter| [counter.scope.window_key(), counter.key()])
|
||||
.collect::<Vec<_>>();
|
||||
batch.add(
|
||||
eval(
|
||||
&BATCH_RATE_LIMITER,
|
||||
&keys,
|
||||
vec![now.to_string(), window_size.to_string()],
|
||||
)
|
||||
.map(window_counts),
|
||||
)
|
||||
}
|
||||
|
||||
/// `BATCH_COUNTER_READ_SCRIPT` over each counter's window and counter keys, in that order.
|
||||
pub(super) fn read_counters(
|
||||
batch: &RedisBatch,
|
||||
counters: &[Counter],
|
||||
) -> Declared<Vec<Option<String>>> {
|
||||
let keys = counters
|
||||
.iter()
|
||||
.flat_map(|counter| [counter.scope.window_key(), counter.key()])
|
||||
.collect::<Vec<_>>();
|
||||
batch.add(eval(&BATCH_COUNTER_READ, &keys, Vec::new()).map(optional_texts))
|
||||
}
|
||||
|
||||
pub(super) fn check_and_increment(
|
||||
batch: &RedisBatch,
|
||||
counters: &[AtomicCounter],
|
||||
) -> Declared<AtomicReply> {
|
||||
let keys = counters
|
||||
.iter()
|
||||
.flat_map(|counter| [counter.counter.scope.window_key(), counter.counter.key()])
|
||||
.collect::<Vec<_>>();
|
||||
let args = counters
|
||||
.iter()
|
||||
.flat_map(|counter| {
|
||||
[
|
||||
counter.limit,
|
||||
counter.increment,
|
||||
counter.ttl_seconds,
|
||||
counter.window_size_seconds,
|
||||
]
|
||||
.map(|arg| arg.to_string())
|
||||
})
|
||||
.collect();
|
||||
batch.add(eval(&CHECK_AND_INCREMENT_BY_N, &keys, args).map(atomic_reply))
|
||||
}
|
||||
|
||||
pub(super) fn guarded_increment(
|
||||
batch: &RedisBatch,
|
||||
increment: &GuardedIncrement,
|
||||
) -> Declared<GuardedReply> {
|
||||
let keys = [
|
||||
increment.counter.scope.window_key(),
|
||||
increment.counter.key(),
|
||||
];
|
||||
let args = vec![
|
||||
increment.expected_window_start.clone(),
|
||||
python_arg(increment.amount),
|
||||
ttl_arg(increment.ttl),
|
||||
];
|
||||
batch.add(eval(&WINDOW_GUARDED_TOKEN_INCREMENT, &keys, args).map(guarded_reply))
|
||||
}
|
||||
|
||||
pub(super) fn parallel_acquire(
|
||||
batch: &RedisBatch,
|
||||
gauges: &[Gauge],
|
||||
slot_ttl: Duration,
|
||||
slot_id: &str,
|
||||
) -> Declared<AcquireReply> {
|
||||
let keys = gauges
|
||||
.iter()
|
||||
.map(|gauge| gauge.scope.gauge_key())
|
||||
.collect::<Vec<_>>();
|
||||
let args = gauges
|
||||
.iter()
|
||||
.flat_map(|gauge| {
|
||||
[
|
||||
gauge.limit.to_string(),
|
||||
slot_ttl.as_secs().to_string(),
|
||||
slot_id.to_owned(),
|
||||
]
|
||||
})
|
||||
.collect();
|
||||
batch.add(eval(&PARALLEL_ACQUIRE, &keys, args).map(acquire_reply))
|
||||
}
|
||||
|
||||
pub(super) fn parallel_release(
|
||||
batch: &RedisBatch,
|
||||
gauges: &[Scope],
|
||||
slot_id: &str,
|
||||
) -> Declared<Vec<i64>> {
|
||||
let keys = gauges.iter().map(Scope::gauge_key).collect::<Vec<_>>();
|
||||
let args = gauges.iter().map(|_| slot_id.to_owned()).collect();
|
||||
batch.add(eval(&PARALLEL_RELEASE, &keys, args).map(integers))
|
||||
}
|
||||
|
||||
pub(super) fn parallel_count(
|
||||
batch: &RedisBatch,
|
||||
gauges: &[Scope],
|
||||
slot_ttl: Duration,
|
||||
) -> Declared<Vec<i64>> {
|
||||
let keys = gauges.iter().map(Scope::gauge_key).collect::<Vec<_>>();
|
||||
let args = gauges
|
||||
.iter()
|
||||
.map(|_| slot_ttl.as_secs().to_string())
|
||||
.collect();
|
||||
batch.add(eval(&PARALLEL_COUNT, &keys, args).map(integers))
|
||||
}
|
||||
|
||||
pub(super) fn token_increment(batch: &RedisBatch, increments: &[Increment]) -> Declared<Vec<f64>> {
|
||||
let keys = increments
|
||||
.iter()
|
||||
.map(|increment| increment.counter.key())
|
||||
.collect::<Vec<_>>();
|
||||
let args = increments
|
||||
.iter()
|
||||
.flat_map(|increment| [python_arg(increment.amount), ttl_arg(increment.ttl)])
|
||||
.collect();
|
||||
batch.add(eval(&TOKEN_INCREMENT, &keys, args).map(floats))
|
||||
}
|
||||
|
||||
/// `RedisBatch`'s `_Increment`: `INCRBYFLOAT`, then `EXPIRE` when the increment carries a TTL. On
|
||||
/// its own it keeps a TTL the key already has, like `async_increment`.
|
||||
fn eval(script: &Script, keys: &[String], args: Vec<String>) -> Command<Value> {
|
||||
let keys = keys.iter().map(String::as_str).collect::<Vec<_>>();
|
||||
commands::eval(script, &keys, args)
|
||||
}
|
||||
|
||||
/// The bytes redis-py sends for a Python `int` or `float` argument.
|
||||
fn python_arg(amount: Number) -> String {
|
||||
match amount {
|
||||
Number::Int(value) => value.to_string(),
|
||||
Number::Float(value) => float_repr(value),
|
||||
}
|
||||
}
|
||||
|
||||
/// `ttl or 0`: the scripts read a zero TTL as "leave the expiry alone".
|
||||
fn ttl_arg(ttl: Option<Duration>) -> String {
|
||||
ttl.map_or(0, |ttl| ttl.as_secs()).to_string()
|
||||
}
|
||||
|
||||
fn items(reply: Value) -> Result<Vec<Value>, Error> {
|
||||
match reply {
|
||||
Value::Array(items) => Ok(items),
|
||||
_ => Err(unexpected_reply("rate limiter script reply is not a list")),
|
||||
}
|
||||
}
|
||||
|
||||
fn integer(value: &Value) -> Result<i64, Error> {
|
||||
match value {
|
||||
Value::Int(value) => Ok(*value),
|
||||
_ => Err(unexpected_reply(
|
||||
"rate limiter script reply holds a non-integer where a count belongs",
|
||||
)),
|
||||
}
|
||||
}
|
||||
|
||||
fn text(value: &Value) -> Result<Option<String>, Error> {
|
||||
match value {
|
||||
Value::Nil => Ok(None),
|
||||
Value::BulkString(bytes) => String::from_utf8(bytes.clone()).map(Some).map_err(|_| {
|
||||
unexpected_reply("rate limiter script reply holds a value that is not UTF-8")
|
||||
}),
|
||||
Value::SimpleString(text) => Ok(Some(text.clone())),
|
||||
Value::Int(value) => Ok(Some(value.to_string())),
|
||||
_ => Err(unexpected_reply(
|
||||
"rate limiter script reply holds a nested value",
|
||||
)),
|
||||
}
|
||||
}
|
||||
|
||||
fn float(reply: Value) -> Result<f64, Error> {
|
||||
text(&reply)?
|
||||
.and_then(|text| text.parse().ok())
|
||||
.ok_or_else(|| unexpected_reply("rate limiter script reply is not a float"))
|
||||
}
|
||||
|
||||
fn index(raw: &Value) -> Result<usize, Error> {
|
||||
usize::try_from(integer(raw)? - 1)
|
||||
.map_err(|_| unexpected_reply("rate limiter script reply names an index below 1"))
|
||||
}
|
||||
|
||||
fn window_counts(reply: Value) -> Result<Vec<WindowCount>, Error> {
|
||||
items(reply)?
|
||||
.chunks(2)
|
||||
.map(|pair| match pair {
|
||||
[window_start, count] => Ok(WindowCount {
|
||||
window_start: text(window_start)?.ok_or_else(|| {
|
||||
unexpected_reply("rate limiter script reply holds a nil window start")
|
||||
})?,
|
||||
count: integer(count)?,
|
||||
}),
|
||||
_ => Err(unexpected_reply(
|
||||
"rate limiter script reply has an odd length",
|
||||
)),
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
|
||||
fn optional_texts(reply: Value) -> Result<Vec<Option<String>>, Error> {
|
||||
items(reply)?.iter().map(text).collect()
|
||||
}
|
||||
|
||||
fn integers(reply: Value) -> Result<Vec<i64>, Error> {
|
||||
items(reply)?.iter().map(integer).collect()
|
||||
}
|
||||
|
||||
fn floats(reply: Value) -> Result<Vec<f64>, Error> {
|
||||
items(reply)?.into_iter().map(float).collect()
|
||||
}
|
||||
|
||||
fn atomic_reply(reply: Value) -> Result<AtomicReply, Error> {
|
||||
let items = items(reply)?;
|
||||
match items.split_first() {
|
||||
Some((Value::Int(0), rest)) => rest
|
||||
.chunks(2)
|
||||
.map(|pair| match pair {
|
||||
[count, window_start] => Ok(Applied {
|
||||
count: integer(count)?,
|
||||
window_start: integer(window_start)?,
|
||||
}),
|
||||
_ => Err(unexpected_reply(
|
||||
"rate limiter script reply has an odd length",
|
||||
)),
|
||||
})
|
||||
.collect::<Result<_, _>>()
|
||||
.map(AtomicReply::Applied),
|
||||
Some((Value::Int(1), [index_value, current, limit])) => Ok(AtomicReply::OverLimit {
|
||||
index: index(index_value)?,
|
||||
current: integer(current)?,
|
||||
limit: integer(limit)?,
|
||||
}),
|
||||
_ => Err(unexpected_reply("rate limiter script reply has no status")),
|
||||
}
|
||||
}
|
||||
|
||||
fn guarded_reply(reply: Value) -> Result<GuardedReply, Error> {
|
||||
match items(reply)?.as_slice() {
|
||||
[applied, count] => Ok(GuardedReply {
|
||||
applied: integer(applied)? == 1,
|
||||
count: integer(count)?,
|
||||
}),
|
||||
_ => Err(unexpected_reply(
|
||||
"rate limiter script reply is not an (applied, count) pair",
|
||||
)),
|
||||
}
|
||||
}
|
||||
|
||||
fn acquire_reply(reply: Value) -> Result<AcquireReply, Error> {
|
||||
let items = items(reply)?;
|
||||
match items.split_first() {
|
||||
Some((Value::Int(0), rest)) => rest
|
||||
.iter()
|
||||
.map(integer)
|
||||
.collect::<Result<_, _>>()
|
||||
.map(AcquireReply::Acquired),
|
||||
Some((Value::Int(1), [index_value, in_flight, limit])) => Ok(AcquireReply::OverLimit {
|
||||
index: index(index_value)?,
|
||||
in_flight: integer(in_flight)?,
|
||||
limit: integer(limit)?,
|
||||
}),
|
||||
_ => Err(unexpected_reply("rate limiter script reply has no status")),
|
||||
}
|
||||
}
|
||||
117
litellm-rust/crates/redis-store/src/rate_limits/types.rs
Normal file
117
litellm-rust/crates/redis-store/src/rate_limits/types.rs
Normal file
|
|
@ -0,0 +1,117 @@
|
|||
use std::time::Duration;
|
||||
|
||||
use litellm_python_compat::Value;
|
||||
|
||||
use super::keys::{Counter, Scope};
|
||||
use crate::{DecodeError, FromValue, Number};
|
||||
|
||||
/// `BATCH_RATE_LIMITER_SCRIPT`'s pair for one counter: the window start it saw or opened and the
|
||||
/// counter after its `+1`.
|
||||
#[derive(Clone, Debug, PartialEq, Eq)]
|
||||
pub struct WindowCount {
|
||||
pub window_start: String,
|
||||
pub count: i64,
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, PartialEq)]
|
||||
pub struct AtomicCounter {
|
||||
pub counter: Counter,
|
||||
pub limit: i64,
|
||||
pub increment: i64,
|
||||
pub ttl_seconds: i64,
|
||||
pub window_size_seconds: i64,
|
||||
}
|
||||
|
||||
/// `CHECK_AND_INCREMENT_BY_N_SCRIPT`'s reply. `index` is 0-based where the Lua's is 1-based.
|
||||
#[derive(Clone, Debug, PartialEq, Eq)]
|
||||
pub enum AtomicReply {
|
||||
Applied(Vec<Applied>),
|
||||
OverLimit {
|
||||
index: usize,
|
||||
current: i64,
|
||||
limit: i64,
|
||||
},
|
||||
}
|
||||
|
||||
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
|
||||
pub struct Applied {
|
||||
pub count: i64,
|
||||
pub window_start: i64,
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, PartialEq)]
|
||||
pub struct GuardedIncrement {
|
||||
pub counter: Counter,
|
||||
pub expected_window_start: String,
|
||||
pub amount: Number,
|
||||
pub ttl: Option<Duration>,
|
||||
}
|
||||
|
||||
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
|
||||
pub struct GuardedReply {
|
||||
pub applied: bool,
|
||||
pub count: i64,
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, PartialEq, Eq)]
|
||||
pub struct Gauge {
|
||||
pub scope: Scope,
|
||||
pub limit: i64,
|
||||
}
|
||||
|
||||
/// `PARALLEL_ACQUIRE_SCRIPT`'s reply. `index` is 0-based where the Lua's is 1-based.
|
||||
#[derive(Clone, Debug, PartialEq, Eq)]
|
||||
pub enum AcquireReply {
|
||||
Acquired(Vec<i64>),
|
||||
OverLimit {
|
||||
index: usize,
|
||||
in_flight: i64,
|
||||
limit: i64,
|
||||
},
|
||||
}
|
||||
|
||||
/// What the memory tier holds: counters and mirrored gauge counts, window starts, and the per-slot
|
||||
/// registry a gauge keeps while Redis is unreachable.
|
||||
#[derive(Clone, Debug, PartialEq)]
|
||||
pub enum Local {
|
||||
Number(Number),
|
||||
Text(String),
|
||||
Slots(Vec<(String, f64)>),
|
||||
}
|
||||
|
||||
impl From<Number> for Local {
|
||||
fn from(number: Number) -> Self {
|
||||
Self::Number(number)
|
||||
}
|
||||
}
|
||||
|
||||
impl crate::memory::Counter for Local {
|
||||
fn as_number(&self) -> Option<Number> {
|
||||
match self {
|
||||
Self::Number(number) => Some(*number),
|
||||
Self::Text(_) | Self::Slots(_) => None,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// A `RedisPipelineIncrementOperation`.
|
||||
#[derive(Clone, Debug, PartialEq)]
|
||||
pub struct Increment {
|
||||
pub counter: Counter,
|
||||
pub amount: Number,
|
||||
pub ttl: Option<Duration>,
|
||||
}
|
||||
|
||||
/// A rate-limit value as `RedisCache._get_cache_logic` reads it.
|
||||
impl FromValue for Local {
|
||||
fn from_value(value: &Value) -> Result<Self, DecodeError> {
|
||||
match value {
|
||||
Value::Str(text) => Ok(Self::Text(text.clone())),
|
||||
value => Number::from_value(value)
|
||||
.map(Local::Number)
|
||||
.ok_or(DecodeError::Shape(
|
||||
"a rate-limit value is a number or a string",
|
||||
)),
|
||||
}
|
||||
}
|
||||
}
|
||||
17
litellm-rust/crates/redis-store/src/router_cooldowns/keys.rs
Normal file
17
litellm-rust/crates/redis-store/src/router_cooldowns/keys.rs
Normal file
|
|
@ -0,0 +1,17 @@
|
|||
pub(super) fn cooldown(model_id: &str) -> String {
|
||||
format!("deployment:{model_id}:cooldown")
|
||||
}
|
||||
|
||||
pub(super) fn cooldowns(model_ids: &[String]) -> Vec<String> {
|
||||
model_ids
|
||||
.iter()
|
||||
.map(|model_id| cooldown(model_id))
|
||||
.collect()
|
||||
}
|
||||
|
||||
pub(super) fn allowed_fails(model_id: &str, suffix: Option<&str>) -> String {
|
||||
match suffix.filter(|suffix| !suffix.is_empty()) {
|
||||
Some(suffix) => format!("deployment:{model_id}:allowed_fails:{suffix}"),
|
||||
None => format!("deployment:{model_id}:allowed_fails"),
|
||||
}
|
||||
}
|
||||
205
litellm-rust/crates/redis-store/src/router_cooldowns/mod.rs
Normal file
205
litellm-rust/crates/redis-store/src/router_cooldowns/mod.rs
Normal file
|
|
@ -0,0 +1,205 @@
|
|||
mod keys;
|
||||
mod types;
|
||||
|
||||
use std::{sync::Arc, time::Duration};
|
||||
|
||||
use litellm_redis_client::{Amount, Command, Expiry, commands};
|
||||
|
||||
pub use self::types::Cooldown;
|
||||
use crate::{
|
||||
Error, Number, Pending, Send,
|
||||
dual_cache::{BatchReads, DualCache, TierPolicy},
|
||||
memory::{Clock, system_clock},
|
||||
remote::{Remote, RemoteSlot},
|
||||
};
|
||||
|
||||
const MAX_CORRECTED_TTL: Duration = Duration::from_secs(60);
|
||||
const CORRECTION_SLACK: Duration = Duration::from_secs(5);
|
||||
|
||||
/// One router's cooldowns and allowed-fails counters: memory per router, Redis shared by every pod.
|
||||
pub struct RouterCooldowns {
|
||||
cooldowns: DualCache<Cooldown, BatchReads>,
|
||||
fails: DualCache<i64>,
|
||||
redis: Arc<RemoteSlot>,
|
||||
}
|
||||
|
||||
impl RouterCooldowns {
|
||||
pub fn new(remote: Option<Remote>, policy: TierPolicy) -> Self {
|
||||
Self::with_clock(remote, policy, system_clock())
|
||||
}
|
||||
|
||||
pub fn with_clock(remote: Option<Remote>, policy: TierPolicy, now: Clock) -> Self {
|
||||
let redis = Arc::new(RemoteSlot::new(remote));
|
||||
Self {
|
||||
fails: DualCache::new(redis.clone(), policy.memory(), now.clone()),
|
||||
cooldowns: DualCache::batch_reading(redis.clone(), policy, now),
|
||||
redis,
|
||||
}
|
||||
}
|
||||
|
||||
/// `DualCache.attach_redis_cache`: memory stays, and `false` means a Redis was already attached.
|
||||
pub fn attach(&self, remote: Remote) -> bool {
|
||||
self.redis.attach(remote)
|
||||
}
|
||||
|
||||
pub fn has_redis(&self) -> bool {
|
||||
self.redis.get().is_some()
|
||||
}
|
||||
|
||||
/// `DualCache.async_set_cache`: memory first, then a `SET` of `json.dumps(value)` with
|
||||
/// `EX int(ttl)` whose failure only logs.
|
||||
pub async fn put(
|
||||
&self,
|
||||
model_id: &str,
|
||||
cooldown: &Cooldown,
|
||||
ttl: Duration,
|
||||
) -> Result<(), Error> {
|
||||
self.cooldowns
|
||||
.set(&keys::cooldown(model_id), cooldown.clone(), Some(ttl))
|
||||
.await
|
||||
}
|
||||
|
||||
/// `DualCache.set_cache`: memory first, then a `SET` of `str(value)` with `EX int(ttl)` whose
|
||||
/// failure only logs.
|
||||
pub async fn put_str(
|
||||
&self,
|
||||
model_id: &str,
|
||||
cooldown: &Cooldown,
|
||||
ttl: Duration,
|
||||
) -> Result<(), Error> {
|
||||
self.cooldowns
|
||||
.set_str(&keys::cooldown(model_id), cooldown.clone(), Some(ttl))
|
||||
.await
|
||||
}
|
||||
|
||||
/// `DualCache.async_increment_cache`: the fleet-wide fail count through `INCRBYFLOAT`. A Redis
|
||||
/// failure falls back to this router's own count.
|
||||
pub async fn increment_allowed_fails(
|
||||
&self,
|
||||
model_id: &str,
|
||||
suffix: Option<&str>,
|
||||
ttl: Duration,
|
||||
) -> Result<i64, Error> {
|
||||
let key = keys::allowed_fails(model_id, suffix);
|
||||
let command =
|
||||
commands::incr_by_float(&key, Amount::Int(1), Some(Expiry::KeepExisting(ttl)));
|
||||
self.count_fail(&key, ttl, command, Number::Float).await
|
||||
}
|
||||
|
||||
/// `DualCache.increment_cache`: the fleet-wide fail count through `INCRBY`. A Redis failure
|
||||
/// falls back to this router's own count.
|
||||
pub async fn increment_allowed_fails_int(
|
||||
&self,
|
||||
model_id: &str,
|
||||
suffix: Option<&str>,
|
||||
ttl: Duration,
|
||||
) -> Result<i64, Error> {
|
||||
let key = keys::allowed_fails(model_id, suffix);
|
||||
let command = commands::incr_by(&key, 1, ttl);
|
||||
self.count_fail(&key, ttl, command, Number::Int).await
|
||||
}
|
||||
|
||||
async fn count_fail<T>(
|
||||
&self,
|
||||
key: &str,
|
||||
ttl: Duration,
|
||||
command: Command<T>,
|
||||
total: fn(T) -> Number,
|
||||
) -> Result<i64, Error> {
|
||||
let local = self
|
||||
.fails
|
||||
.memory()
|
||||
.update(key, Some(ttl), |count| Ok(count.unwrap_or(0) + 1))?;
|
||||
let Some(client) = self.redis.client() else {
|
||||
return Ok(local);
|
||||
};
|
||||
match client.send(command).await {
|
||||
Ok(count) => Ok(total(count).truncate()),
|
||||
Err(error) => {
|
||||
crate::failure_log::redis_failure!(
|
||||
&error,
|
||||
"allowed_fails counter fell back to this worker's in-memory count"
|
||||
);
|
||||
Ok(local)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// The deployments still cooling down, read on the batch `send` picks. An entry whose cooldown
|
||||
/// already ended is dropped from memory, and one memory would keep past its end is shortened, at
|
||||
/// most to 60 s.
|
||||
pub fn active(
|
||||
self: &Arc<Self>,
|
||||
send: Send<'_>,
|
||||
model_ids: &[String],
|
||||
) -> Result<Pending<Vec<(String, Cooldown)>>, Error> {
|
||||
let read = self
|
||||
.cooldowns
|
||||
.declare_batch_get(send, &keys::cooldowns(model_ids))?;
|
||||
let model_ids = model_ids.to_vec();
|
||||
let store = Arc::clone(self);
|
||||
Ok(Box::pin(async move {
|
||||
let read = store.cooldowns.resolve_batch_get(read).await?;
|
||||
let now = store.cooldowns.now();
|
||||
let active = model_ids
|
||||
.into_iter()
|
||||
.zip(read)
|
||||
.filter_map(|(model_id, cooldown)| cooldown.map(|cooldown| (model_id, cooldown)))
|
||||
.map(|(model_id, cooldown)| {
|
||||
Ok(store
|
||||
.correct(&model_id, &cooldown, now)?
|
||||
.then_some((model_id, cooldown)))
|
||||
})
|
||||
.collect::<Result<Vec<_>, Error>>()?;
|
||||
Ok(active.into_iter().flatten().collect())
|
||||
}))
|
||||
}
|
||||
|
||||
/// The shortest stored cooldown, ended ones included, as `main`'s `get_min_cooldown` reads it.
|
||||
pub async fn minimum(&self, model_ids: &[String]) -> Result<Option<Number>, Error> {
|
||||
let read = self
|
||||
.cooldowns
|
||||
.batch_get(&keys::cooldowns(model_ids))
|
||||
.await?;
|
||||
Ok(Number::min(
|
||||
read.into_iter()
|
||||
.flatten()
|
||||
.map(|cooldown| cooldown.cooldown_time),
|
||||
))
|
||||
}
|
||||
|
||||
pub fn delete_local(&self, model_id: &str) -> Result<(), Error> {
|
||||
Ok(self.cooldowns.memory().delete(&keys::cooldown(model_id))?)
|
||||
}
|
||||
|
||||
/// Drops this router's allowed-fails counts, leaving its cooldowns in place.
|
||||
pub fn flush_local_counters(&self) -> Result<(), Error> {
|
||||
Ok(self.fails.memory().flush()?)
|
||||
}
|
||||
|
||||
/// The memory tier keyed as Python's `InMemoryCache` was, for the fuzz probe that reads
|
||||
/// in-memory router state.
|
||||
pub fn snapshot(&self) -> Result<Vec<(String, Cooldown)>, Error> {
|
||||
Ok(self.cooldowns.memory().entries()?)
|
||||
}
|
||||
|
||||
fn correct(&self, model_id: &str, cooldown: &Cooldown, now: Duration) -> Result<bool, Error> {
|
||||
let key = keys::cooldown(model_id);
|
||||
let memory = self.cooldowns.memory();
|
||||
let remaining = cooldown.remaining(now);
|
||||
if remaining <= 0.0 {
|
||||
memory.delete(&key)?;
|
||||
return Ok(false);
|
||||
}
|
||||
let remaining = Duration::from_secs_f64(remaining);
|
||||
if memory.expires_at(&key)? > Some(now + remaining + CORRECTION_SLACK) {
|
||||
memory.delete(&key)?;
|
||||
memory.set(
|
||||
&key,
|
||||
cooldown.clone(),
|
||||
Some(remaining.min(MAX_CORRECTED_TTL)),
|
||||
)?;
|
||||
}
|
||||
Ok(true)
|
||||
}
|
||||
}
|
||||
|
|
@ -0,0 +1,61 @@
|
|||
use std::time::Duration;
|
||||
|
||||
use litellm_python_compat::Value;
|
||||
|
||||
use crate::{DecodeError, FromValue, IntoValue, Number};
|
||||
|
||||
#[derive(Clone, Debug, PartialEq)]
|
||||
pub struct Cooldown {
|
||||
pub exception_received: String,
|
||||
pub status_code: String,
|
||||
pub timestamp: Number,
|
||||
pub cooldown_time: Number,
|
||||
}
|
||||
|
||||
impl Cooldown {
|
||||
pub(super) fn remaining(&self, now: Duration) -> f64 {
|
||||
self.timestamp.as_f64() + self.cooldown_time.as_f64() - now.as_secs_f64()
|
||||
}
|
||||
}
|
||||
|
||||
/// The `CooldownCacheValue` dict the router reads and writes.
|
||||
impl IntoValue for Cooldown {
|
||||
fn into_value(self) -> Value {
|
||||
let field = |name: &str, value: Value| (Value::Str(name.to_owned()), value);
|
||||
Value::Dict(vec![
|
||||
field("exception_received", Value::Str(self.exception_received)),
|
||||
field("status_code", Value::Str(self.status_code)),
|
||||
field("timestamp", self.timestamp.to_value()),
|
||||
field("cooldown_time", self.cooldown_time.to_value()),
|
||||
])
|
||||
}
|
||||
}
|
||||
|
||||
impl FromValue for Cooldown {
|
||||
fn from_value(value: &Value) -> Result<Self, DecodeError> {
|
||||
let Value::Dict(fields) = value else {
|
||||
return Err(DecodeError::Shape("a cooldown is a dict"));
|
||||
};
|
||||
let field = |name: &str| {
|
||||
fields
|
||||
.iter()
|
||||
.find(|(key, _)| matches!(key, Value::Str(key) if key == name))
|
||||
.map(|(_, value)| value)
|
||||
};
|
||||
let text = |name: &'static str| match field(name) {
|
||||
Some(Value::Str(text)) => Ok(text.clone()),
|
||||
_ => Err(DecodeError::Shape(name)),
|
||||
};
|
||||
let number = |name: &'static str| {
|
||||
field(name)
|
||||
.and_then(Number::from_value)
|
||||
.ok_or(DecodeError::Shape(name))
|
||||
};
|
||||
Ok(Self {
|
||||
exception_received: text("exception_received")?,
|
||||
status_code: text("status_code")?,
|
||||
timestamp: number("timestamp")?,
|
||||
cooldown_time: number("cooldown_time")?,
|
||||
})
|
||||
}
|
||||
}
|
||||
28
litellm-rust/crates/redis-store/src/router_usage/keys.rs
Normal file
28
litellm-rust/crates/redis-store/src/router_usage/keys.rs
Normal file
|
|
@ -0,0 +1,28 @@
|
|||
/// A counter's key segment, which is also how the Python router names it.
|
||||
#[derive(Clone, Copy, Debug, PartialEq, Eq, strum::EnumString, strum::IntoStaticStr)]
|
||||
pub enum Kind {
|
||||
#[strum(serialize = "tpm")]
|
||||
Tpm,
|
||||
#[strum(serialize = "rpm")]
|
||||
Rpm,
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, PartialEq, Eq)]
|
||||
pub struct Counter {
|
||||
pub deployment_id: String,
|
||||
pub model: Option<String>,
|
||||
pub minute: String,
|
||||
pub kind: Kind,
|
||||
}
|
||||
|
||||
impl Counter {
|
||||
/// `global_router:{id}:{model}:{tpm|rpm}:{HH-MM}`, where Python formats a missing model as `None`.
|
||||
pub(super) fn key(&self) -> String {
|
||||
let kind: &str = self.kind.into();
|
||||
let model = self.model.as_deref().unwrap_or("None");
|
||||
format!(
|
||||
"global_router:{}:{model}:{kind}:{}",
|
||||
self.deployment_id, self.minute
|
||||
)
|
||||
}
|
||||
}
|
||||
110
litellm-rust/crates/redis-store/src/router_usage/mod.rs
Normal file
110
litellm-rust/crates/redis-store/src/router_usage/mod.rs
Normal file
|
|
@ -0,0 +1,110 @@
|
|||
mod keys;
|
||||
|
||||
use std::{sync::Arc, time::Duration};
|
||||
|
||||
pub use self::keys::{Counter, Kind};
|
||||
use crate::{
|
||||
Error, Number,
|
||||
dual_cache::{BatchReads, DualCache, PipelineIncrement, TierPolicy},
|
||||
memory::{Clock, system_clock},
|
||||
remote::{Remote, RemoteSlot},
|
||||
};
|
||||
|
||||
/// `RoutingArgs.ttl`, the window of a per-minute counter.
|
||||
pub const COUNTER_TTL: Duration = Duration::from_secs(60);
|
||||
|
||||
/// One router's global TPM/RPM counters: memory per router, Redis shared by every pod.
|
||||
pub struct RouterUsage {
|
||||
usage: DualCache<Number, BatchReads>,
|
||||
}
|
||||
|
||||
impl RouterUsage {
|
||||
pub fn new(remote: Option<Remote>, policy: TierPolicy) -> Self {
|
||||
Self::with_clock(remote, policy, system_clock())
|
||||
}
|
||||
|
||||
pub fn with_clock(remote: Option<Remote>, policy: TierPolicy, now: Clock) -> Self {
|
||||
Self {
|
||||
usage: DualCache::batch_reading(Arc::new(RemoteSlot::new(remote)), policy, now),
|
||||
}
|
||||
}
|
||||
|
||||
/// `DualCache.attach_redis_cache`: memory stays, and `false` means a Redis was already attached.
|
||||
pub fn attach(&self, remote: Remote) -> bool {
|
||||
self.usage.redis().attach(remote)
|
||||
}
|
||||
|
||||
pub fn has_redis(&self) -> bool {
|
||||
self.usage.redis().get().is_some()
|
||||
}
|
||||
|
||||
/// `DualCache.async_increment_cache`. A Redis failure answers with the memory count, and an
|
||||
/// entry memory cannot add to with `None`.
|
||||
pub async fn increment(
|
||||
&self,
|
||||
counter: &Counter,
|
||||
value: Number,
|
||||
ttl: Duration,
|
||||
) -> Option<Number> {
|
||||
self.usage.increment(&counter.key(), value, Some(ttl)).await
|
||||
}
|
||||
|
||||
/// One deployment's tokens and requests, `async_increment_pipeline` on its own round trip, with
|
||||
/// the totals mirrored back into memory as `int`s.
|
||||
pub async fn increment_deployment(
|
||||
&self,
|
||||
deployment_id: &str,
|
||||
model: Option<&str>,
|
||||
minute: &str,
|
||||
tokens: Number,
|
||||
requests: i64,
|
||||
) -> Result<(), Error> {
|
||||
let increment = |kind, amount| PipelineIncrement {
|
||||
key: Counter {
|
||||
deployment_id: deployment_id.to_owned(),
|
||||
model: model.map(str::to_owned),
|
||||
minute: minute.to_owned(),
|
||||
kind,
|
||||
}
|
||||
.key(),
|
||||
amount,
|
||||
ttl: Some(COUNTER_TTL),
|
||||
};
|
||||
let increments = [
|
||||
increment(Kind::Tpm, tokens),
|
||||
increment(Kind::Rpm, Number::Int(requests)),
|
||||
];
|
||||
let Some(totals) = self.usage.increment_pipeline(&increments).await else {
|
||||
return Ok(());
|
||||
};
|
||||
if self.usage.redis().get().is_none() {
|
||||
return Ok(());
|
||||
}
|
||||
increments
|
||||
.iter()
|
||||
.zip(totals)
|
||||
.try_for_each(|(increment, total)| {
|
||||
self.usage.memory().set(
|
||||
&increment.key,
|
||||
Number::Int(total.truncate()),
|
||||
Some(COUNTER_TTL),
|
||||
)
|
||||
})?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub async fn get_many(&self, counters: &[Counter]) -> Result<Vec<Option<Number>>, Error> {
|
||||
let keys = counters.iter().map(Counter::key).collect::<Vec<_>>();
|
||||
self.usage.batch_get(&keys).await
|
||||
}
|
||||
|
||||
pub fn flush_local_counters(&self) -> Result<(), Error> {
|
||||
Ok(self.usage.memory().flush()?)
|
||||
}
|
||||
|
||||
/// The memory tier keyed as Python's `InMemoryCache` was, for the fuzz probe that reads
|
||||
/// in-memory router state.
|
||||
pub fn snapshot(&self) -> Result<Vec<(String, Number)>, Error> {
|
||||
Ok(self.usage.memory().entries()?)
|
||||
}
|
||||
}
|
||||
|
|
@ -0,0 +1,3 @@
|
|||
pub(super) fn pin(tenant: &str, session_id: &str) -> String {
|
||||
format!("{{sensitive_route:{tenant}:{session_id}}}:model")
|
||||
}
|
||||
111
litellm-rust/crates/redis-store/src/sensitive_route_pins/mod.rs
Normal file
111
litellm-rust/crates/redis-store/src/sensitive_route_pins/mod.rs
Normal file
|
|
@ -0,0 +1,111 @@
|
|||
mod keys;
|
||||
|
||||
use std::{sync::Arc, time::Duration};
|
||||
|
||||
use litellm_python_compat::{Value, json};
|
||||
use litellm_redis_client::{Ttl, commands};
|
||||
|
||||
use crate::{
|
||||
DecodeError, Error,
|
||||
codec::python_value,
|
||||
dual_cache::{DualCache, MemorySettings},
|
||||
memory::{Clock, system_clock},
|
||||
remote::{Remote, RemoteSlot},
|
||||
value::text,
|
||||
};
|
||||
|
||||
/// The model each session is pinned to: memory per worker, Redis shared by every pod. Like
|
||||
/// `SensitiveDataRoutingHandler` on `main`, it reaches past `DualCache` to Redis first.
|
||||
pub struct SensitiveRoutePins {
|
||||
pins: DualCache<String>,
|
||||
}
|
||||
|
||||
impl SensitiveRoutePins {
|
||||
pub fn new(remote: Option<Remote>) -> Self {
|
||||
Self::with_clock(remote, system_clock())
|
||||
}
|
||||
|
||||
pub fn with_clock(remote: Option<Remote>, now: Clock) -> Self {
|
||||
Self {
|
||||
pins: DualCache::new(
|
||||
Arc::new(RemoteSlot::new(remote)),
|
||||
MemorySettings::default(),
|
||||
now,
|
||||
),
|
||||
}
|
||||
}
|
||||
|
||||
/// Redis first, kept in memory for the key's remaining TTL. A Redis miss or failure serves memory.
|
||||
pub async fn get(
|
||||
&self,
|
||||
tenant: &str,
|
||||
session_id: &str,
|
||||
default_ttl: Duration,
|
||||
) -> Result<Option<String>, Error> {
|
||||
let key = keys::pin(tenant, session_id);
|
||||
if let Some(model) = self.read_remote(&key, default_ttl).await? {
|
||||
return Ok(Some(model));
|
||||
}
|
||||
Ok(self.pins.memory().get(&key)?)
|
||||
}
|
||||
|
||||
/// `SET ... EX` whose failure only logs, then memory.
|
||||
pub async fn set(
|
||||
&self,
|
||||
tenant: &str,
|
||||
session_id: &str,
|
||||
model: &str,
|
||||
ttl: Duration,
|
||||
) -> Result<(), Error> {
|
||||
let key = keys::pin(tenant, session_id);
|
||||
if let Some(client) = self.pins.redis().client() {
|
||||
let encoded = json::dumps(&Value::Str(model.to_owned()))?;
|
||||
if let Err(error) = client.send(commands::set(&key, encoded, ttl)).await {
|
||||
crate::failure_log::redis_failure!(
|
||||
&error,
|
||||
"writing a sensitive route pin to Redis failed"
|
||||
);
|
||||
}
|
||||
}
|
||||
self.pins.memory().set(&key, model.to_owned(), Some(ttl))?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// `GET`, then on a hit its own `TTL` trip, as `RedisCache.async_get_cache` and `async_get_ttl`
|
||||
/// make them. A key without a TTL, or a failed `TTL`, keeps `default_ttl`.
|
||||
async fn read_remote(&self, key: &str, default_ttl: Duration) -> Result<Option<String>, Error> {
|
||||
let Some(client) = self.pins.redis().client() else {
|
||||
return Ok(None);
|
||||
};
|
||||
let read = client
|
||||
.send(commands::get::<Vec<u8>>(key))
|
||||
.await
|
||||
.map_err(Error::from)
|
||||
.and_then(|raw| raw.as_deref().map(decode).transpose().map_err(Error::from));
|
||||
let model = match read {
|
||||
Ok(Some(model)) => model,
|
||||
Ok(None) => return Ok(None),
|
||||
Err(error) => {
|
||||
crate::failure_log::redis_failure!(
|
||||
&error,
|
||||
"reading a sensitive route pin from Redis failed, serving memory"
|
||||
);
|
||||
return Ok(None);
|
||||
}
|
||||
};
|
||||
let ttl = match client.send(commands::ttl(key)).await {
|
||||
Ok(Ttl::Expires(remaining)) => remaining,
|
||||
Ok(Ttl::Missing | Ttl::Persistent) => default_ttl,
|
||||
Err(error) => {
|
||||
tracing::debug!(%error, "Redis TTL Error");
|
||||
default_ttl
|
||||
}
|
||||
};
|
||||
self.pins.memory().set(key, model.clone(), Some(ttl))?;
|
||||
Ok(Some(model))
|
||||
}
|
||||
}
|
||||
|
||||
fn decode(raw: &[u8]) -> Result<String, DecodeError> {
|
||||
text(&python_value(raw)?, "a pinned model name")
|
||||
}
|
||||
85
litellm-rust/crates/redis-store/src/spend_counters/keys.rs
Normal file
85
litellm-rust/crates/redis-store/src/spend_counters/keys.rs
Normal file
|
|
@ -0,0 +1,85 @@
|
|||
/// One budget's spend counter. `parse` and `key` round-trip `main`'s `spend:*` keys exactly.
|
||||
#[derive(Clone, Debug, PartialEq, Eq, Hash)]
|
||||
pub enum SpendCounter {
|
||||
Key {
|
||||
token: String,
|
||||
window: Option<String>,
|
||||
},
|
||||
Team {
|
||||
team_id: String,
|
||||
window: Option<String>,
|
||||
},
|
||||
TeamMember {
|
||||
user_id: String,
|
||||
team_id: String,
|
||||
},
|
||||
User(String),
|
||||
EndUser(String),
|
||||
Org(String),
|
||||
Tag(String),
|
||||
Project(String),
|
||||
ModelAccessGroup(String),
|
||||
ShadowEval(String),
|
||||
}
|
||||
|
||||
#[derive(Debug, thiserror::Error)]
|
||||
#[error("{0:?} is not a spend counter")]
|
||||
pub struct NotASpendCounter(pub String);
|
||||
|
||||
impl SpendCounter {
|
||||
pub fn key(&self) -> String {
|
||||
let windowed = |scope: &str, id: &str, window: &Option<String>| match window {
|
||||
Some(window) => format!("spend:{scope}:{id}:window:{window}"),
|
||||
None => format!("spend:{scope}:{id}"),
|
||||
};
|
||||
match self {
|
||||
Self::Key { token, window } => windowed("key", token, window),
|
||||
Self::Team { team_id, window } => windowed("team", team_id, window),
|
||||
Self::TeamMember { user_id, team_id } => {
|
||||
format!("spend:team_member:{user_id}:{team_id}")
|
||||
}
|
||||
Self::User(id) => format!("spend:user:{id}"),
|
||||
Self::EndUser(id) => format!("spend:end_user:{id}"),
|
||||
Self::Org(id) => format!("spend:org:{id}"),
|
||||
Self::Tag(name) => format!("spend:tag:{name}"),
|
||||
Self::Project(id) => format!("spend:project:{id}"),
|
||||
Self::ModelAccessGroup(name) => format!("spend:model_access_group:{name}"),
|
||||
Self::ShadowEval(job_id) => format!("spend:shadow_eval:{job_id}"),
|
||||
}
|
||||
}
|
||||
|
||||
pub fn parse(key: &str) -> Result<Self, NotASpendCounter> {
|
||||
let invalid = || NotASpendCounter(key.to_owned());
|
||||
let rest = key.strip_prefix("spend:").ok_or_else(invalid)?;
|
||||
let (scope, id) = rest.split_once(':').ok_or_else(invalid)?;
|
||||
let windowed = || match id.rsplit_once(":window:") {
|
||||
Some((id, window)) if !window.is_empty() => (id.to_owned(), Some(window.to_owned())),
|
||||
_ => (id.to_owned(), None),
|
||||
};
|
||||
Ok(match scope {
|
||||
"key" => {
|
||||
let (token, window) = windowed();
|
||||
Self::Key { token, window }
|
||||
}
|
||||
"team" => {
|
||||
let (team_id, window) = windowed();
|
||||
Self::Team { team_id, window }
|
||||
}
|
||||
"team_member" => {
|
||||
let (user_id, team_id) = id.rsplit_once(':').ok_or_else(invalid)?;
|
||||
Self::TeamMember {
|
||||
user_id: user_id.to_owned(),
|
||||
team_id: team_id.to_owned(),
|
||||
}
|
||||
}
|
||||
"user" => Self::User(id.to_owned()),
|
||||
"end_user" => Self::EndUser(id.to_owned()),
|
||||
"org" => Self::Org(id.to_owned()),
|
||||
"tag" => Self::Tag(id.to_owned()),
|
||||
"project" => Self::Project(id.to_owned()),
|
||||
"model_access_group" => Self::ModelAccessGroup(id.to_owned()),
|
||||
"shadow_eval" => Self::ShadowEval(id.to_owned()),
|
||||
_ => return Err(invalid()),
|
||||
})
|
||||
}
|
||||
}
|
||||
380
litellm-rust/crates/redis-store/src/spend_counters/mod.rs
Normal file
380
litellm-rust/crates/redis-store/src/spend_counters/mod.rs
Normal file
|
|
@ -0,0 +1,380 @@
|
|||
mod keys;
|
||||
mod types;
|
||||
|
||||
use std::{sync::Arc, time::Duration};
|
||||
|
||||
use litellm_python_compat::repr::to_str;
|
||||
use litellm_redis_client::{Declared, Expiry, RedisBatch, commands};
|
||||
|
||||
pub use self::keys::{NotASpendCounter, SpendCounter};
|
||||
pub use self::types::{LostSeed, Reading};
|
||||
use crate::{
|
||||
Error, Number, Pending, RequestRedis, Send,
|
||||
dual_cache::{DualCache, MemorySettings},
|
||||
memory::{Clock, system_clock},
|
||||
remote::{Remote, RemoteSlot},
|
||||
};
|
||||
|
||||
/// `RedisCache.async_set_max`, byte for byte, sent as `EVAL` like it: a counter only ever rises to
|
||||
/// the authoritative spend.
|
||||
const SET_MAX: &str = concat!(
|
||||
"local cur = redis.call('GET', KEYS[1]) ",
|
||||
"if cur == false or tonumber(cur) < tonumber(ARGV[1]) then ",
|
||||
"redis.call('SET', KEYS[1], ARGV[1]) ",
|
||||
"if tonumber(ARGV[2]) > 0 then redis.call('EXPIRE', KEYS[1], ARGV[2]) end ",
|
||||
"return ARGV[1] end ",
|
||||
"return cur"
|
||||
);
|
||||
|
||||
/// The proxy's spend counters: Redis shared by every pod, and this pod's memory copy that answers
|
||||
/// only when Redis cannot. Like `spend_counter_cache`'s callers on `main`, most ops reach past
|
||||
/// `DualCache` to its Redis and its memory.
|
||||
pub struct SpendCounters {
|
||||
cache: DualCache<Number>,
|
||||
}
|
||||
|
||||
impl SpendCounters {
|
||||
pub fn new(remote: Option<Remote>) -> Self {
|
||||
Self::with_clock(remote, system_clock())
|
||||
}
|
||||
|
||||
pub fn with_clock(remote: Option<Remote>, now: Clock) -> Self {
|
||||
Self {
|
||||
cache: DualCache::new(
|
||||
Arc::new(RemoteSlot::new(remote)),
|
||||
MemorySettings::default(),
|
||||
now,
|
||||
),
|
||||
}
|
||||
}
|
||||
|
||||
/// `DualCache.attach_redis_cache`: memory stays, and `false` means a Redis was already attached.
|
||||
pub fn attach(&self, remote: Remote) -> bool {
|
||||
self.cache.redis().attach(remote)
|
||||
}
|
||||
|
||||
pub fn has_redis(&self) -> bool {
|
||||
self.cache.redis().get().is_some()
|
||||
}
|
||||
|
||||
/// The TTL every Redis write gives a counter, `None` without Redis.
|
||||
pub fn redis_ttl(&self) -> Option<Duration> {
|
||||
self.cache.redis().get().map(|remote| remote.default_ttl)
|
||||
}
|
||||
|
||||
/// `read_spend_counter_cache_value`: Redis first. `RedisCache.async_get_cache` swallows an unreachable
|
||||
/// `GET` as a miss, so an open breaker, a transport error, or a value `float()` refuses reads memory instead.
|
||||
pub async fn read(&self, counter: &SpendCounter) -> Result<Reading, Error> {
|
||||
let Some(client) = self.cache.redis().client() else {
|
||||
return Ok(Reading::Memory(self.get_local(counter)?));
|
||||
};
|
||||
match client.send(commands::get::<Vec<u8>>(&counter.key())).await {
|
||||
Ok(raw) => match types::decode(raw) {
|
||||
Ok(value) => Ok(Reading::Redis(value)),
|
||||
Err(error) => {
|
||||
tracing::debug!(%error, "spend counter Redis value did not decode, falling back to memory");
|
||||
Ok(Reading::Memory(self.get_local(counter)?))
|
||||
}
|
||||
},
|
||||
Err(litellm_redis_client::Error::BreakerOpen) => {
|
||||
Ok(Reading::Memory(self.get_local(counter)?))
|
||||
}
|
||||
Err(error) => {
|
||||
crate::failure_log::redis_failure!(&error, "spend counter Redis read failed");
|
||||
Ok(Reading::Memory(self.get_local(counter)?))
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// One `MGET` for many counters on the batch `send` picks; `None` without Redis. The read
|
||||
/// resolves to `None` when it failed, so each caller reads on its own.
|
||||
pub fn read_many(
|
||||
&self,
|
||||
send: Send<'_>,
|
||||
counters: &[SpendCounter],
|
||||
) -> Option<Pending<Option<Vec<Option<Number>>>>> {
|
||||
let batch = self.cache.redis().batch(send)?;
|
||||
let keys: Vec<String> = counters.iter().map(SpendCounter::key).collect();
|
||||
let keys: Vec<&str> = keys.iter().map(String::as_str).collect();
|
||||
let read = batch.add(commands::mget::<Vec<u8>>(&keys));
|
||||
Some(Box::pin(async move {
|
||||
let decoded = read
|
||||
.resolve()
|
||||
.await
|
||||
.map_err(Error::from)
|
||||
.and_then(|replies| {
|
||||
replies
|
||||
.into_iter()
|
||||
.map(|raw| Ok(types::decode(raw)?))
|
||||
.collect::<Result<Vec<_>, Error>>()
|
||||
});
|
||||
Ok(decoded
|
||||
.inspect_err(|error| {
|
||||
tracing::debug!(%error, "spend counter batch read failed, falling back to per-key reads");
|
||||
})
|
||||
.ok())
|
||||
}))
|
||||
}
|
||||
|
||||
/// `increment_spend_counter_cache`: `INCRBYFLOAT` refreshing the TTL, mirrored into memory. A
|
||||
/// failed increment may or may not have landed, so the counter is dropped everywhere before the
|
||||
/// error returns. Without Redis it is a memory increment.
|
||||
pub async fn increment(&self, counter: &SpendCounter, amount: Number) -> Result<Number, Error> {
|
||||
let Some(remote) = self.cache.redis().get() else {
|
||||
return self.increment_local(counter, amount);
|
||||
};
|
||||
let key = counter.key();
|
||||
let increment = commands::incr_by_float(
|
||||
&key,
|
||||
amount.into(),
|
||||
Some(Expiry::Refresh(remote.default_ttl)),
|
||||
);
|
||||
match remote.client.send(increment).await {
|
||||
Ok(total) => {
|
||||
let total = Number::Float(total);
|
||||
self.cache.memory().set(&key, total, None)?;
|
||||
Ok(total)
|
||||
}
|
||||
Err(error) => {
|
||||
self.invalidate(std::slice::from_ref(counter)).await;
|
||||
Err(error.into())
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// `run_spend_counter_pipeline`: every increment in one round trip, totals in order. The caller
|
||||
/// decides what a failure means for counters whose increment may or may not have landed.
|
||||
pub async fn increment_many(
|
||||
&self,
|
||||
increments: &[(SpendCounter, Number)],
|
||||
) -> Result<Vec<Number>, Error> {
|
||||
let Some(remote) = self.cache.redis().get() else {
|
||||
return increments
|
||||
.iter()
|
||||
.map(|(counter, amount)| self.increment_local(counter, *amount))
|
||||
.collect();
|
||||
};
|
||||
let pipeline = RedisBatch::pipelined(remote.client.clone());
|
||||
let totals = increments.iter().map(|(counter, amount)| {
|
||||
pipeline.add(commands::incr_by_float(
|
||||
&counter.key(),
|
||||
(*amount).into(),
|
||||
Some(Expiry::Refresh(remote.default_ttl)),
|
||||
))
|
||||
});
|
||||
let totals = Declared::all(totals.collect::<Vec<_>>()).await;
|
||||
increments
|
||||
.iter()
|
||||
.zip(totals)
|
||||
.map(|((counter, _), total)| {
|
||||
let total = Number::Float(total?);
|
||||
self.cache.memory().set(&counter.key(), total, None)?;
|
||||
Ok(total)
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
|
||||
/// `increment_many` on the request's post-call batch. Each total reaches memory once the batch
|
||||
/// goes out, and a counter whose increment failed is dropped everywhere. `false` without Redis.
|
||||
pub fn increment_many_post_call(
|
||||
self: &Arc<Self>,
|
||||
request: &RequestRedis,
|
||||
increments: &[(SpendCounter, Number)],
|
||||
) -> bool {
|
||||
let Some(remote) = self.cache.redis().get() else {
|
||||
return false;
|
||||
};
|
||||
let batch = Send::PostCall(request).batch(&remote.client);
|
||||
let declared = increments
|
||||
.iter()
|
||||
.map(|(counter, amount)| {
|
||||
let total = batch.add(commands::incr_by_float(
|
||||
&counter.key(),
|
||||
(*amount).into(),
|
||||
Some(Expiry::KeepExisting(remote.default_ttl)),
|
||||
));
|
||||
(counter.clone(), total)
|
||||
})
|
||||
.collect::<Vec<_>>();
|
||||
let store = Arc::clone(self);
|
||||
request.spawn(async move {
|
||||
for (counter, total) in declared {
|
||||
match total.settled().await {
|
||||
Ok(total) => {
|
||||
if let Err(error) = store.cache.memory().set(&counter.key(), Number::Float(total), None) {
|
||||
tracing::warn!(%error, "mirroring a spend counter total into memory failed");
|
||||
}
|
||||
}
|
||||
Err(error) => {
|
||||
let key = counter.key();
|
||||
crate::failure_log::redis_failure!(
|
||||
&error,
|
||||
&format!("spend counter {key} increment did not land in the post-call pipeline; invalidating it"),
|
||||
);
|
||||
store.invalidate(std::slice::from_ref(&counter)).await;
|
||||
}
|
||||
}
|
||||
}
|
||||
});
|
||||
true
|
||||
}
|
||||
|
||||
/// `_repair_stale_spend_counter`: raise memory and Redis to `spend`, never lowering either.
|
||||
/// `None` without Redis, else what Redis holds afterwards.
|
||||
pub async fn set_max(
|
||||
&self,
|
||||
counter: &SpendCounter,
|
||||
spend: Number,
|
||||
) -> Result<Option<Number>, Error> {
|
||||
let key = counter.key();
|
||||
{
|
||||
let _guard = self.cache.memory().guard();
|
||||
if self
|
||||
.cache
|
||||
.memory()
|
||||
.get(&key)?
|
||||
.is_none_or(|cached| cached.as_f64() < spend.as_f64())
|
||||
{
|
||||
self.cache.memory().set(&key, spend, None)?;
|
||||
}
|
||||
}
|
||||
let Some(remote) = self.cache.redis().get() else {
|
||||
return Ok(None);
|
||||
};
|
||||
let args = (
|
||||
to_str(&spend.to_value()),
|
||||
remote.default_ttl.as_secs().to_string(),
|
||||
);
|
||||
let stored: Option<f64> = remote
|
||||
.client
|
||||
.send(commands::eval_source(SET_MAX, &[&key], args))
|
||||
.await?;
|
||||
Ok(stored.map(Number::Float))
|
||||
}
|
||||
|
||||
/// `SpendCounterReseed.coalesced(_window)`: seed a cold counter with `SET NX` so one pod wins,
|
||||
/// answer with whatever the counter holds, and mirror it into memory. Without Redis, memory
|
||||
/// keeps the larger of the seed and what it held.
|
||||
pub async fn seed(
|
||||
&self,
|
||||
counter: &SpendCounter,
|
||||
spend: Number,
|
||||
lost: LostSeed,
|
||||
) -> Result<Number, Error> {
|
||||
let key = counter.key();
|
||||
let Some(Remote {
|
||||
client,
|
||||
default_ttl,
|
||||
}) = self.cache.redis().get()
|
||||
else {
|
||||
return Ok(self.cache.memory().update(&key, None, |cached| {
|
||||
Ok(match cached {
|
||||
Some(cached) if cached.as_f64() > spend.as_f64() => {
|
||||
Number::Float(cached.as_f64())
|
||||
}
|
||||
_ => spend,
|
||||
})
|
||||
})?);
|
||||
};
|
||||
let current = if client
|
||||
.send(commands::set_nx(&key, spend.to_json(), *default_ttl))
|
||||
.await?
|
||||
{
|
||||
spend
|
||||
} else {
|
||||
match types::decode(client.send(commands::get::<Vec<u8>>(&key)).await?)? {
|
||||
Some(winner) => Number::Float(winner.as_f64()),
|
||||
None => match lost {
|
||||
LostSeed::Answer => spend,
|
||||
LostSeed::Increment => {
|
||||
let increment = commands::incr_by_float(
|
||||
&key,
|
||||
spend.into(),
|
||||
Some(Expiry::KeepExisting(*default_ttl)),
|
||||
);
|
||||
Number::Float(client.send(increment).await?)
|
||||
}
|
||||
},
|
||||
}
|
||||
};
|
||||
self.cache.memory().set(&key, current, None)?;
|
||||
Ok(current)
|
||||
}
|
||||
|
||||
/// A spend reset: memory, then Redis, each for `ttl` or its own default.
|
||||
pub async fn set(
|
||||
&self,
|
||||
counter: &SpendCounter,
|
||||
spend: Number,
|
||||
ttl: Option<Duration>,
|
||||
) -> Result<(), Error> {
|
||||
let key = counter.key();
|
||||
self.cache.memory().set(&key, spend, ttl)?;
|
||||
let Some(remote) = self.cache.redis().get() else {
|
||||
return Ok(());
|
||||
};
|
||||
let ttl = ttl.unwrap_or(remote.default_ttl);
|
||||
remote
|
||||
.client
|
||||
.send(commands::set(&key, spend.to_json(), ttl))
|
||||
.await?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Drop counters from memory, then from Redis in one round trip.
|
||||
pub async fn delete(&self, counters: &[SpendCounter]) -> Result<(), Error> {
|
||||
let keys: Vec<String> = counters.iter().map(SpendCounter::key).collect();
|
||||
self.cache.delete_keys(&keys).await
|
||||
}
|
||||
|
||||
/// `refresh_spend_counter_ttl`: `false` without Redis or for a counter Redis no longer holds.
|
||||
pub async fn refresh_ttl(&self, counter: &SpendCounter) -> Result<bool, Error> {
|
||||
let Some(remote) = self.cache.redis().get() else {
|
||||
return Ok(false);
|
||||
};
|
||||
Ok(remote
|
||||
.client
|
||||
.send(commands::expire(&counter.key(), remote.default_ttl))
|
||||
.await?)
|
||||
}
|
||||
|
||||
pub fn get_local(&self, counter: &SpendCounter) -> Result<Option<Number>, Error> {
|
||||
Ok(self.cache.memory().get(&counter.key())?)
|
||||
}
|
||||
|
||||
pub fn set_local(
|
||||
&self,
|
||||
counter: &SpendCounter,
|
||||
spend: Number,
|
||||
ttl: Option<Duration>,
|
||||
) -> Result<(), Error> {
|
||||
Ok(self.cache.memory().set(&counter.key(), spend, ttl)?)
|
||||
}
|
||||
|
||||
pub fn delete_local(&self, counter: &SpendCounter) -> Result<(), Error> {
|
||||
Ok(self.cache.memory().delete(&counter.key())?)
|
||||
}
|
||||
|
||||
/// `DualCache.async_increment_cache(local_only=True)`, which passes `InMemoryCache` no TTL.
|
||||
pub fn increment_local(&self, counter: &SpendCounter, amount: Number) -> Result<Number, Error> {
|
||||
Ok(self
|
||||
.cache
|
||||
.memory()
|
||||
.increment(&counter.key(), amount, None)?)
|
||||
}
|
||||
|
||||
pub fn flush_local(&self) -> Result<(), Error> {
|
||||
Ok(self.cache.memory().flush()?)
|
||||
}
|
||||
|
||||
pub fn snapshot(&self) -> Result<Vec<(String, Number)>, Error> {
|
||||
Ok(self.cache.memory().entries()?)
|
||||
}
|
||||
|
||||
/// `invalidate_spend_counter`: memory always, Redis best effort.
|
||||
pub async fn invalidate(&self, counters: &[SpendCounter]) {
|
||||
if let Err(error) = self.delete(counters).await {
|
||||
tracing::debug!(%error, "dropping spend counters from Redis failed");
|
||||
}
|
||||
}
|
||||
}
|
||||
48
litellm-rust/crates/redis-store/src/spend_counters/types.rs
Normal file
48
litellm-rust/crates/redis-store/src/spend_counters/types.rs
Normal file
|
|
@ -0,0 +1,48 @@
|
|||
use litellm_python_compat::Value;
|
||||
|
||||
use super::keys::SpendCounter;
|
||||
use crate::{DecodeError, FromValue, Number, codec};
|
||||
|
||||
/// The `spend:*` key Python names a counter by.
|
||||
impl FromValue for SpendCounter {
|
||||
fn from_value(value: &Value) -> Result<Self, DecodeError> {
|
||||
let Value::Str(key) = value else {
|
||||
return Err(DecodeError::Shape("a spend counter is a spend:* key"));
|
||||
};
|
||||
Self::parse(key).map_err(|_| DecodeError::Shape("a spend counter is a spend:* key"))
|
||||
}
|
||||
}
|
||||
|
||||
/// Where a read was answered, named the way the Python store reports it. Redis is cross-pod truth,
|
||||
/// including a clean miss; memory only holds this pod's writes and answers only when there is no
|
||||
/// Redis or it failed.
|
||||
#[derive(Clone, Copy, Debug, PartialEq, strum::IntoStaticStr)]
|
||||
pub enum Reading {
|
||||
#[strum(serialize = "redis")]
|
||||
Redis(Option<Number>),
|
||||
#[strum(serialize = "memory")]
|
||||
Memory(Option<Number>),
|
||||
}
|
||||
|
||||
impl Reading {
|
||||
pub fn value(self) -> Option<Number> {
|
||||
match self {
|
||||
Self::Redis(value) | Self::Memory(value) => value,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Who wins when a seed's `SET NX` lost and the winner's value is gone by the time it is read.
|
||||
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
|
||||
pub enum LostSeed {
|
||||
/// `SpendCounterReseed.coalesced`: answer with the seed without writing it.
|
||||
Answer,
|
||||
/// `SpendCounterReseed.coalesced_window`: `INCRBYFLOAT` the seed, keeping any TTL.
|
||||
Increment,
|
||||
}
|
||||
|
||||
pub(super) fn decode(raw: Option<Vec<u8>>) -> Result<Option<Number>, DecodeError> {
|
||||
raw.as_deref()
|
||||
.map(|raw| codec::number(raw, "a spend counter is a number"))
|
||||
.transpose()
|
||||
}
|
||||
|
|
@ -2,12 +2,12 @@ use litellm_python_compat::Value;
|
|||
|
||||
use crate::{DecodeError, Number};
|
||||
|
||||
/// A store type read from the plain Python data its host sends, so the host never names a field.
|
||||
/// A store value read from the plain Python data it is held as.
|
||||
pub trait FromValue: Sized {
|
||||
fn from_value(value: &Value) -> Result<Self, DecodeError>;
|
||||
}
|
||||
|
||||
/// A store type written as the plain Python data its host hands back.
|
||||
/// A store value written as the plain Python data it is held as.
|
||||
pub trait IntoValue {
|
||||
fn into_value(self) -> Value;
|
||||
}
|
||||
|
|
@ -36,20 +36,6 @@ impl IntoValue for Number {
|
|||
}
|
||||
}
|
||||
|
||||
/// The fields of a Python tuple of exactly `N` items.
|
||||
pub(crate) fn fields<'a, const N: usize>(
|
||||
value: &'a Value,
|
||||
shape: &'static str,
|
||||
) -> Result<&'a [Value; N], DecodeError> {
|
||||
match value {
|
||||
Value::Tuple(items) | Value::List(items) => items
|
||||
.as_slice()
|
||||
.try_into()
|
||||
.map_err(|_| DecodeError::Shape(shape)),
|
||||
_ => Err(DecodeError::Shape(shape)),
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn text(value: &Value, shape: &'static str) -> Result<String, DecodeError> {
|
||||
match value {
|
||||
Value::Str(text) => Ok(text.clone()),
|
||||
|
|
@ -57,23 +43,6 @@ pub(crate) fn text(value: &Value, shape: &'static str) -> Result<String, DecodeE
|
|||
}
|
||||
}
|
||||
|
||||
pub(crate) fn optional_text(
|
||||
value: &Value,
|
||||
shape: &'static str,
|
||||
) -> Result<Option<String>, DecodeError> {
|
||||
match value {
|
||||
Value::None => Ok(None),
|
||||
value => text(value, shape).map(Some),
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn integer(value: &Value, shape: &'static str) -> Result<i64, DecodeError> {
|
||||
match value {
|
||||
Value::Int(value) => i64::try_from(value).map_err(|_| DecodeError::Shape(shape)),
|
||||
_ => Err(DecodeError::Shape(shape)),
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn number(value: &Value, shape: &'static str) -> Result<Number, DecodeError> {
|
||||
Number::from_value(value).ok_or(DecodeError::Shape(shape))
|
||||
}
|
||||
|
|
|
|||
741
litellm-rust/crates/redis-store/tests/auth_cache.rs
Normal file
741
litellm-rust/crates/redis-store/tests/auth_cache.rs
Normal file
|
|
@ -0,0 +1,741 @@
|
|||
mod support;
|
||||
|
||||
use std::time::Duration;
|
||||
|
||||
use litellm_python_compat::Value;
|
||||
use litellm_redis_client::commands;
|
||||
use litellm_redis_client::{Operation, RedisClient, Ttl};
|
||||
use litellm_redis_store::Send as RedisSend;
|
||||
use litellm_redis_store::{
|
||||
Number, Remote, RequestRedis,
|
||||
auth_cache::{self, AuthCache, Settings},
|
||||
router_cooldowns::RouterCooldowns,
|
||||
};
|
||||
use litellm_redis_testing::{RedisDeployment, Topology};
|
||||
use rstest::rstest;
|
||||
use support::{COOLDOWN_POLICY, connect, connect_observed, remote, within};
|
||||
use tokio::runtime::Handle;
|
||||
|
||||
const LONG_DEADLINE: Duration = Duration::from_secs(60);
|
||||
const KEY_OBJECT: &str = "88dc28d0f030c55ed4ab77ed8faf098196cb1c05df778539800c9f1243fe6b4b";
|
||||
const TEAM_JSON: &str = concat!(
|
||||
r#"{"team_id": "team-\u00e9", "models": ["gpt-5", "claude-opus-5"], "max_budget": null, "#,
|
||||
r#""spend": 0.30000000000000004, "blocked": false, "tpm_limit": 10000000000, "#,
|
||||
r#""metadata": {"tags": [], "quoted": "it's \"x\""}}"#,
|
||||
);
|
||||
const TEAM_STR: &str = concat!(
|
||||
"{'team_id': 'team-é', 'models': ['gpt-5', 'claude-opus-5'], 'max_budget': None, ",
|
||||
"'spend': 0.30000000000000004, 'blocked': False, 'tpm_limit': 10000000000, ",
|
||||
r#"'metadata': {'tags': [], 'quoted': 'it\'s "x"'}}"#,
|
||||
);
|
||||
|
||||
fn dict<const N: usize>(entries: [(&str, Value); N]) -> Value {
|
||||
Value::Dict(
|
||||
entries
|
||||
.into_iter()
|
||||
.map(|(key, value)| (Value::from(key), value))
|
||||
.collect(),
|
||||
)
|
||||
}
|
||||
|
||||
fn strings(items: &[&str]) -> Value {
|
||||
Value::List(items.iter().copied().map(Value::from).collect())
|
||||
}
|
||||
|
||||
fn team_record() -> Value {
|
||||
dict([
|
||||
("team_id", Value::from("team-é")),
|
||||
("models", strings(&["gpt-5", "claude-opus-5"])),
|
||||
("max_budget", Value::None),
|
||||
("spend", Value::Float(0.30000000000000004)),
|
||||
("blocked", Value::Bool(false)),
|
||||
("tpm_limit", Value::from(10000000000)),
|
||||
(
|
||||
"metadata",
|
||||
dict([
|
||||
("tags", strings(&[])),
|
||||
("quoted", Value::from("it's \"x\"")),
|
||||
]),
|
||||
),
|
||||
])
|
||||
}
|
||||
|
||||
/// The writes `tests/unit/redis/stores/test_auth_cache.py` checks against Python's own
|
||||
/// `json.dumps` and `str`.
|
||||
fn written() -> [(Value, &'static str, &'static str); 7] {
|
||||
[
|
||||
(team_record(), TEAM_JSON, TEAM_STR),
|
||||
(
|
||||
dict([
|
||||
("token", Value::from(KEY_OBJECT)),
|
||||
("spend", Value::Float(1e-05)),
|
||||
("rpm_limit", Value::from(5)),
|
||||
("expires", Value::from("2026-11-03T00:00:00Z")),
|
||||
]),
|
||||
concat!(
|
||||
r#"{"token": "88dc28d0f030c55ed4ab77ed8faf098196cb1c05df778539800c9f1243fe6b4b", "spend": 1e-05, "#,
|
||||
r#""rpm_limit": 5, "expires": "2026-11-03T00:00:00Z"}"#,
|
||||
),
|
||||
concat!(
|
||||
"{'token': '88dc28d0f030c55ed4ab77ed8faf098196cb1c05df778539800c9f1243fe6b4b', 'spend': 1e-05, ",
|
||||
"'rpm_limit': 5, 'expires': '2026-11-03T00:00:00Z'}",
|
||||
),
|
||||
),
|
||||
(
|
||||
Value::from("__no_team_membership__"),
|
||||
"\"__no_team_membership__\"",
|
||||
"__no_team_membership__",
|
||||
),
|
||||
(
|
||||
strings(&["mcp-a", "mcp-b"]),
|
||||
r#"["mcp-a", "mcp-b"]"#,
|
||||
"['mcp-a', 'mcp-b']",
|
||||
),
|
||||
(Value::Float(12.5), "12.5", "12.5"),
|
||||
(Value::from(7), "7", "7"),
|
||||
(Value::Bool(true), "true", "True"),
|
||||
]
|
||||
}
|
||||
|
||||
/// The reads `tests/unit/redis/stores/test_auth_cache.py` checks against Python's own
|
||||
/// `RedisCache._get_cache_logic`.
|
||||
fn read_back() -> [(&'static str, Value); 5] {
|
||||
[
|
||||
(TEAM_JSON, team_record()),
|
||||
(TEAM_STR, team_record()),
|
||||
("['mcp-a', 'mcp-b']", strings(&["mcp-a", "mcp-b"])),
|
||||
(
|
||||
"{'user_id': 'u1', 'teams': ['t1'], 'active': True, 'budget': None}",
|
||||
dict([
|
||||
("user_id", Value::from("u1")),
|
||||
("teams", strings(&["t1"])),
|
||||
("active", Value::Bool(true)),
|
||||
("budget", Value::None),
|
||||
]),
|
||||
),
|
||||
("\"__NO_MAPPING__\"", Value::from("__NO_MAPPING__")),
|
||||
]
|
||||
}
|
||||
|
||||
fn settings(memory_ttl: Option<u64>) -> Settings {
|
||||
Settings {
|
||||
memory_ttl: memory_ttl.map(Duration::from_secs),
|
||||
..Settings::default()
|
||||
}
|
||||
}
|
||||
|
||||
fn store(client: Option<RedisClient>) -> AuthCache {
|
||||
AuthCache::new(client.map(remote), settings(Some(60)))
|
||||
}
|
||||
|
||||
fn team(id: &str) -> Value {
|
||||
Value::Dict(vec![(Value::from("team_id"), Value::from(id))])
|
||||
}
|
||||
|
||||
async fn redis_ttl(raw: &RedisClient, key: &str) -> Duration {
|
||||
let Ttl::Expires(remaining) = raw.send(commands::ttl(key)).await.unwrap() else {
|
||||
panic!("{key} has no TTL");
|
||||
};
|
||||
remaining
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case(KEY_OBJECT, true)]
|
||||
#[case(
|
||||
"88DC28D0F030C55ED4AB77ED8FAF098196CB1C05DF778539800C9F1243FE6B4B",
|
||||
false
|
||||
)]
|
||||
#[case(
|
||||
"88dc28d0f030c55ed4ab77ed8faf098196cb1c05df778539800c9f1243fe6b4",
|
||||
false
|
||||
)]
|
||||
#[case(
|
||||
"team_id:88dc28d0f030c55ed4ab77ed8faf098196cb1c05df778539800c9f1243fe6b4b",
|
||||
false
|
||||
)]
|
||||
#[case("team_membership:user-1:team-1", false)]
|
||||
#[case("end_user_restricted_registry", false)]
|
||||
fn only_a_bare_lowercase_sha256_digest_is_a_user_key_object(
|
||||
#[case] key: &str,
|
||||
#[case] key_object: bool,
|
||||
) {
|
||||
assert_eq!(auth_cache::is_key_object(key), key_object, "{key}");
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn writes_mains_json_and_str_bytes(
|
||||
#[values(Topology::Standalone, Topology::Cluster)] topology: Topology,
|
||||
) {
|
||||
let deployment = RedisDeployment::start(topology).await.unwrap();
|
||||
let raw = connect(&deployment);
|
||||
let store = store(Some(connect(&deployment)));
|
||||
|
||||
for (index, (value, json, printed)) in written().into_iter().enumerate() {
|
||||
let key = format!("team_id:t{index}");
|
||||
store.set(&key, value.clone(), None).await.unwrap();
|
||||
assert_eq!(
|
||||
raw.send(commands::get::<String>(&key))
|
||||
.await
|
||||
.unwrap()
|
||||
.as_deref(),
|
||||
Some(json)
|
||||
);
|
||||
store.set_str(&key, value, None).await.unwrap();
|
||||
assert_eq!(
|
||||
raw.send(commands::get::<String>(&key))
|
||||
.await
|
||||
.unwrap()
|
||||
.as_deref(),
|
||||
Some(printed)
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn reads_every_encoding_main_writes_and_backfills_memory(
|
||||
#[values(Topology::Standalone, Topology::Cluster)] topology: Topology,
|
||||
) {
|
||||
let deployment = RedisDeployment::start(topology).await.unwrap();
|
||||
let raw = connect(&deployment);
|
||||
let store = store(Some(connect(&deployment)));
|
||||
|
||||
for (index, (bytes, expected)) in read_back().into_iter().enumerate() {
|
||||
let key = format!("user_id:u{index}");
|
||||
raw.send(commands::set(&key, bytes, Duration::from_secs(60)))
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(store.get(None, &key).await.unwrap(), Some(expected.clone()));
|
||||
raw.send(commands::del(&[&key])).await.unwrap();
|
||||
assert_eq!(
|
||||
store.get(None, &key).await.unwrap(),
|
||||
Some(expected),
|
||||
"{key}"
|
||||
);
|
||||
}
|
||||
assert_eq!(store.get(None, "user_id:absent").await.unwrap(), None);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn ttls_follow_dual_cache_and_the_write_back_rule(
|
||||
#[values(Topology::Standalone, Topology::Cluster)] topology: Topology,
|
||||
) {
|
||||
let deployment = RedisDeployment::start(topology).await.unwrap();
|
||||
let raw = connect(&deployment);
|
||||
let store = AuthCache::new(
|
||||
Some(Remote {
|
||||
client: connect(&deployment),
|
||||
default_ttl: Duration::from_secs(45),
|
||||
}),
|
||||
Settings {
|
||||
memory_ttl: Some(Duration::from_secs(20)),
|
||||
..Settings::default()
|
||||
},
|
||||
);
|
||||
|
||||
store
|
||||
.set(
|
||||
"team_id:given",
|
||||
team("a"),
|
||||
Some(Duration::from_secs_f64(30.9)),
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
store.set("team_id:default", team("b"), None).await.unwrap();
|
||||
store
|
||||
.write_back(
|
||||
RedisSend::Now,
|
||||
vec![
|
||||
("user-1".to_owned(), team("c"), None),
|
||||
("user-2".to_owned(), team("d"), Some(Duration::from_secs(5))),
|
||||
],
|
||||
)
|
||||
.unwrap()
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
assert!(within(redis_ttl(&raw, "team_id:given").await, 30));
|
||||
assert!(within(redis_ttl(&raw, "team_id:default").await, 20));
|
||||
assert!(within(redis_ttl(&raw, "user-1").await, 45));
|
||||
assert!(within(redis_ttl(&raw, "user-2").await, 5));
|
||||
assert_eq!(store.get_local("user-1").unwrap(), Some(team("c")));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn user_key_objects_survive_churn_in_the_management_object_partition() {
|
||||
let store = AuthCache::new(
|
||||
None,
|
||||
Settings {
|
||||
memory_entries: 2,
|
||||
..Settings::default()
|
||||
},
|
||||
);
|
||||
store.set(KEY_OBJECT, team("key"), None).await.unwrap();
|
||||
|
||||
for index in 0..5 {
|
||||
store
|
||||
.set(&format!("team_id:t{index}"), team("t"), None)
|
||||
.await
|
||||
.unwrap();
|
||||
}
|
||||
|
||||
assert_eq!(
|
||||
store.get(None, KEY_OBJECT).await.unwrap(),
|
||||
Some(team("key"))
|
||||
);
|
||||
assert_eq!(store.len_local().unwrap(), (2, 1));
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn a_value_json_refuses_stays_in_memory_and_never_reaches_redis(
|
||||
#[values(Topology::Standalone, Topology::Cluster)] topology: Topology,
|
||||
) {
|
||||
let deployment = RedisDeployment::start(topology).await.unwrap();
|
||||
let raw = connect(&deployment);
|
||||
let store = store(Some(connect(&deployment)));
|
||||
let members = Value::Set(vec![Value::from("a")]);
|
||||
|
||||
store
|
||||
.set("team_id:set", members.clone(), None)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(store.get_local("team_id:set").unwrap(), Some(members));
|
||||
assert_eq!(
|
||||
raw.send(commands::get::<String>("team_id:set"))
|
||||
.await
|
||||
.unwrap(),
|
||||
None
|
||||
);
|
||||
}
|
||||
|
||||
// RedisCache._pipeline_helper runs json.dumps inside the pipeline's try, so one value it refuses
|
||||
// aborts every SET of that pipeline, while DualCache has already written them all to memory
|
||||
#[tokio::test]
|
||||
async fn a_value_json_refuses_keeps_its_whole_pipeline_out_of_redis() {
|
||||
let deployment = RedisDeployment::start(Topology::Standalone).await.unwrap();
|
||||
let raw = connect(&deployment);
|
||||
let store = store(Some(connect(&deployment)));
|
||||
let members = Value::Set(vec![Value::from("a")]);
|
||||
|
||||
store
|
||||
.set_many(
|
||||
vec![
|
||||
("team_id:a".to_owned(), team("a")),
|
||||
("team_id:set".to_owned(), members.clone()),
|
||||
],
|
||||
None,
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(store.get_local("team_id:a").unwrap(), Some(team("a")));
|
||||
assert_eq!(store.get_local("team_id:set").unwrap(), Some(members));
|
||||
assert_eq!(
|
||||
raw.send(commands::get::<String>("team_id:a"))
|
||||
.await
|
||||
.unwrap(),
|
||||
None
|
||||
);
|
||||
}
|
||||
|
||||
// auth_object_prefetch writes every entry to memory before
|
||||
// RedisCache.async_set_cache_pipeline_with_ttls encodes them, and that encodes all before sending
|
||||
#[tokio::test]
|
||||
async fn a_write_back_value_json_refuses_fails_after_memory_and_before_redis() {
|
||||
let deployment = RedisDeployment::start(Topology::Standalone).await.unwrap();
|
||||
let raw = connect(&deployment);
|
||||
let store = store(Some(connect(&deployment)));
|
||||
let members = Value::Set(vec![Value::from("a")]);
|
||||
|
||||
let written = store.write_back(
|
||||
RedisSend::Now,
|
||||
vec![
|
||||
("team_id:set".to_owned(), members.clone(), None),
|
||||
("team_id:a".to_owned(), team("a"), None),
|
||||
],
|
||||
);
|
||||
|
||||
assert!(written.is_err());
|
||||
assert_eq!(store.get_local("team_id:set").unwrap(), Some(members));
|
||||
assert_eq!(store.get_local("team_id:a").unwrap(), Some(team("a")));
|
||||
assert_eq!(
|
||||
raw.send(commands::get::<String>("team_id:a"))
|
||||
.await
|
||||
.unwrap(),
|
||||
None
|
||||
);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn pre_call_writes_and_reads_ride_one_pipeline_with_the_cooldown_read(
|
||||
#[values(Topology::Standalone, Topology::Cluster)] topology: Topology,
|
||||
) {
|
||||
let deployment = RedisDeployment::start(topology).await.unwrap();
|
||||
let raw = connect(&deployment);
|
||||
let (client, trips) = connect_observed(&deployment);
|
||||
let cooldowns = std::sync::Arc::new(RouterCooldowns::new(
|
||||
Some(remote(client.clone())),
|
||||
COOLDOWN_POLICY,
|
||||
));
|
||||
let store = store(Some(client));
|
||||
let request = RequestRedis::new(Handle::current(), LONG_DEADLINE);
|
||||
raw.send(commands::set(
|
||||
KEY_OBJECT,
|
||||
"{\"token\": \"t\"}",
|
||||
Duration::from_secs(60),
|
||||
))
|
||||
.await
|
||||
.unwrap();
|
||||
raw.send(commands::set(
|
||||
"team_alias:alpha",
|
||||
"\"stale\"",
|
||||
Duration::from_secs(60),
|
||||
))
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
let reads = store
|
||||
.get_many(
|
||||
RedisSend::PreCall(&request),
|
||||
&[KEY_OBJECT.to_owned(), "end_user_id:absent".to_owned()],
|
||||
)
|
||||
.unwrap();
|
||||
let active = cooldowns
|
||||
.active(RedisSend::PreCall(&request), &["dep-a".to_owned()])
|
||||
.unwrap();
|
||||
let written = store
|
||||
.declare_set(RedisSend::PreCall(&request), "team_id:t1", team("t1"), None)
|
||||
.unwrap();
|
||||
let deleted = store
|
||||
.declare_delete(RedisSend::PreCall(&request), "team_alias:alpha")
|
||||
.unwrap();
|
||||
|
||||
assert!(trips.seen().is_empty());
|
||||
assert_eq!(store.get_local("team_id:t1").unwrap(), Some(team("t1")));
|
||||
request.flush_pre().await;
|
||||
|
||||
let mut seen = trips.seen();
|
||||
seen.sort_by_key(|operation| format!("{operation:?}"));
|
||||
let expected = match topology {
|
||||
Topology::Standalone => vec![Operation::Pipeline { ops: 4 }],
|
||||
Topology::Cluster => vec![
|
||||
Operation::Command("DEL"),
|
||||
Operation::Command("MGET"),
|
||||
Operation::Command("MGET"),
|
||||
Operation::Command("SET"),
|
||||
],
|
||||
};
|
||||
assert_eq!(seen, expected);
|
||||
assert_eq!(
|
||||
reads.await.unwrap(),
|
||||
vec![
|
||||
Some(Value::Dict(vec![(Value::from("token"), Value::from("t"))])),
|
||||
None
|
||||
]
|
||||
);
|
||||
assert_eq!(active.await.unwrap(), vec![]);
|
||||
written.await.unwrap();
|
||||
deleted.await.unwrap();
|
||||
assert_eq!(
|
||||
raw.send(commands::get::<String>("team_id:t1"))
|
||||
.await
|
||||
.unwrap(),
|
||||
Some("{\"team_id\": \"t1\"}".to_owned())
|
||||
);
|
||||
assert!(within(redis_ttl(&raw, "team_id:t1").await, 60));
|
||||
assert_eq!(
|
||||
raw.send(commands::get::<String>("team_alias:alpha"))
|
||||
.await
|
||||
.unwrap(),
|
||||
None
|
||||
);
|
||||
assert_eq!(trips.seen().len(), expected.len());
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn a_write_back_on_the_request_waits_for_the_next_flush(
|
||||
#[values(Topology::Standalone, Topology::Cluster)] topology: Topology,
|
||||
) {
|
||||
let deployment = RedisDeployment::start(topology).await.unwrap();
|
||||
let raw = connect(&deployment);
|
||||
let (client, trips) = connect_observed(&deployment);
|
||||
let store = store(Some(client));
|
||||
let request = RequestRedis::new(Handle::current(), LONG_DEADLINE);
|
||||
|
||||
store
|
||||
.write_back(
|
||||
RedisSend::PreCall(&request),
|
||||
vec![(
|
||||
"user-1".to_owned(),
|
||||
team("u"),
|
||||
Some(Duration::from_secs(42)),
|
||||
)],
|
||||
)
|
||||
.unwrap()
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
assert!(trips.seen().is_empty());
|
||||
assert_eq!(store.get_local("user-1").unwrap(), Some(team("u")));
|
||||
request.flush_pre().await;
|
||||
let write = match topology {
|
||||
Topology::Standalone => Operation::Pipeline { ops: 1 },
|
||||
Topology::Cluster => Operation::Command("SET"),
|
||||
};
|
||||
assert_eq!(trips.seen(), [write]);
|
||||
assert!(within(redis_ttl(&raw, "user-1").await, 42));
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn an_eviction_broadcast_drops_only_this_pods_copy_and_the_next_read_refills_from_redis(
|
||||
#[values(Topology::Standalone, Topology::Cluster)] topology: Topology,
|
||||
) {
|
||||
let deployment = RedisDeployment::start(topology).await.unwrap();
|
||||
let raw = connect(&deployment);
|
||||
let this_pod = store(Some(connect(&deployment)));
|
||||
let other_pod = store(Some(connect(&deployment)));
|
||||
this_pod.set("team_id:t1", team("old"), None).await.unwrap();
|
||||
other_pod
|
||||
.set("team_id:t1", team("new"), None)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(
|
||||
this_pod.get(None, "team_id:t1").await.unwrap(),
|
||||
Some(team("old"))
|
||||
);
|
||||
this_pod.delete_local("team_id:t1").unwrap();
|
||||
|
||||
assert_eq!(
|
||||
this_pod.get(None, "team_id:t1").await.unwrap(),
|
||||
Some(team("new"))
|
||||
);
|
||||
assert!(
|
||||
raw.send(commands::get::<String>("team_id:t1"))
|
||||
.await
|
||||
.unwrap()
|
||||
.is_some()
|
||||
);
|
||||
}
|
||||
|
||||
// RedisBatch.read_as_missing: once the request's MGET found a key absent, or the request deleted it,
|
||||
// a GET of it is answered without a round trip until the request sets it again
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn a_key_the_request_read_as_missing_is_not_read_again(
|
||||
#[values(Topology::Standalone, Topology::Cluster)] topology: Topology,
|
||||
) {
|
||||
let deployment = RedisDeployment::start(topology).await.unwrap();
|
||||
let raw = connect(&deployment);
|
||||
let store = store(Some(connect(&deployment)));
|
||||
let request = RequestRedis::new(Handle::current(), LONG_DEADLINE);
|
||||
let other_pod_writes = |key: &'static str| {
|
||||
raw.send(commands::set(
|
||||
key,
|
||||
TEAM_JSON.to_owned(),
|
||||
Duration::from_secs(60),
|
||||
))
|
||||
};
|
||||
let read = store
|
||||
.get_many(RedisSend::PreCall(&request), &["team_id:read".to_owned()])
|
||||
.unwrap();
|
||||
assert_eq!(read.await.unwrap(), vec![None]);
|
||||
store
|
||||
.declare_delete(RedisSend::PreCall(&request), "team_id:deleted")
|
||||
.unwrap()
|
||||
.await
|
||||
.unwrap();
|
||||
other_pod_writes("team_id:read").await.unwrap();
|
||||
other_pod_writes("team_id:deleted").await.unwrap();
|
||||
|
||||
assert_eq!(
|
||||
store.get(Some(&request), "team_id:read").await.unwrap(),
|
||||
None
|
||||
);
|
||||
assert_eq!(
|
||||
store.get(Some(&request), "team_id:deleted").await.unwrap(),
|
||||
None
|
||||
);
|
||||
assert_eq!(
|
||||
store.get(None, "team_id:read").await.unwrap(),
|
||||
Some(team_record())
|
||||
);
|
||||
|
||||
store
|
||||
.declare_set(
|
||||
RedisSend::PreCall(&request),
|
||||
"team_id:deleted",
|
||||
team("t"),
|
||||
None,
|
||||
)
|
||||
.unwrap()
|
||||
.await
|
||||
.unwrap();
|
||||
store.delete_local("team_id:deleted").unwrap();
|
||||
assert_eq!(
|
||||
store.get(Some(&request), "team_id:deleted").await.unwrap(),
|
||||
Some(team("t"))
|
||||
);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn delete_drops_memory_and_redis(
|
||||
#[values(Topology::Standalone, Topology::Cluster)] topology: Topology,
|
||||
) {
|
||||
let deployment = RedisDeployment::start(topology).await.unwrap();
|
||||
let raw = connect(&deployment);
|
||||
let store = store(Some(connect(&deployment)));
|
||||
let keys = [KEY_OBJECT.to_owned(), "team_id:t1".to_owned()];
|
||||
store
|
||||
.set_many(
|
||||
keys.iter().map(|key| (key.clone(), team("t"))).collect(),
|
||||
None,
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
assert!(
|
||||
raw.send(commands::get::<String>("team_id:t1"))
|
||||
.await
|
||||
.unwrap()
|
||||
.is_some()
|
||||
);
|
||||
|
||||
store.delete(&keys).await.unwrap();
|
||||
|
||||
for key in &keys {
|
||||
assert_eq!(store.get_local(key).unwrap(), None);
|
||||
assert_eq!(raw.send(commands::get::<String>(key)).await.unwrap(), None);
|
||||
}
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn increment_if_present_only_raises_an_entry_redis_holds(
|
||||
#[values(Topology::Standalone, Topology::Cluster)] topology: Topology,
|
||||
) {
|
||||
let deployment = RedisDeployment::start(topology).await.unwrap();
|
||||
let raw = connect(&deployment);
|
||||
let store = store(Some(connect(&deployment)));
|
||||
|
||||
assert_eq!(
|
||||
store
|
||||
.increment_if_present("litellm-proxy-budget:spend", Number::Float(2.5))
|
||||
.await
|
||||
.unwrap(),
|
||||
None
|
||||
);
|
||||
assert_eq!(
|
||||
raw.send(commands::get::<String>("litellm-proxy-budget:spend"))
|
||||
.await
|
||||
.unwrap(),
|
||||
None
|
||||
);
|
||||
store
|
||||
.set("litellm-proxy-budget:spend", Value::Float(1.5), None)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(
|
||||
store
|
||||
.increment_if_present("litellm-proxy-budget:spend", Number::Float(2.5))
|
||||
.await
|
||||
.unwrap(),
|
||||
Some(4.0)
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn without_redis_every_operation_is_memory_only() {
|
||||
let store = store(None);
|
||||
let request = RequestRedis::new(Handle::current(), LONG_DEADLINE);
|
||||
|
||||
store.set("team_id:t1", team("t"), None).await.unwrap();
|
||||
|
||||
store
|
||||
.declare_set(RedisSend::PreCall(&request), "team_id:t2", team("t"), None)
|
||||
.unwrap()
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(store.get_local("team_id:t2").unwrap(), Some(team("t")));
|
||||
store
|
||||
.declare_delete(RedisSend::PreCall(&request), "team_id:t2")
|
||||
.unwrap()
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(store.get_local("team_id:t2").unwrap(), None);
|
||||
assert!(
|
||||
store
|
||||
.get_many(RedisSend::PreCall(&request), &["team_id:t1".to_owned()])
|
||||
.is_none()
|
||||
);
|
||||
assert!(
|
||||
store
|
||||
.get_many(RedisSend::Now, &["team_id:t1".to_owned()])
|
||||
.is_none()
|
||||
);
|
||||
assert_eq!(
|
||||
store
|
||||
.increment_if_present("k", Number::Float(1.0))
|
||||
.await
|
||||
.unwrap(),
|
||||
None
|
||||
);
|
||||
assert_eq!(
|
||||
store.get(None, "team_id:t1").await.unwrap(),
|
||||
Some(team("t"))
|
||||
);
|
||||
assert_eq!(store.get_local("team_id:t2").unwrap(), None);
|
||||
}
|
||||
|
||||
// UserApiKeyCache writes the key-object partition, then the rest, and deletes both partitions
|
||||
// concurrently, while async_increment_if_present sends its script as a plain EVAL
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn the_auth_cache_sends_the_commands_user_api_key_cache_sends(
|
||||
#[values(Topology::Standalone, Topology::Cluster)] topology: Topology,
|
||||
) {
|
||||
let deployment = RedisDeployment::start(topology).await.unwrap();
|
||||
let (client, trips) = connect_observed(&deployment);
|
||||
let store = store(Some(client));
|
||||
deployment
|
||||
.cli_on_every_node(&["CONFIG", "RESETSTAT"])
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
store
|
||||
.set_many(
|
||||
vec![
|
||||
(KEY_OBJECT.to_owned(), team("t0")),
|
||||
("team_id:t1".to_owned(), team("t1")),
|
||||
],
|
||||
None,
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
let written = trips.seen();
|
||||
store
|
||||
.delete(&[KEY_OBJECT.to_owned(), "team_id:t1".to_owned()])
|
||||
.await
|
||||
.unwrap();
|
||||
store
|
||||
.increment_if_present("litellm-proxy-budget:spend", Number::Int(5))
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(written, vec![Operation::Pipeline { ops: 1 }; 2]);
|
||||
assert_eq!(
|
||||
trips.seen()[written.len()..written.len() + 2],
|
||||
[Operation::Command("DEL"), Operation::Command("DEL")]
|
||||
);
|
||||
assert_eq!(support::calls(&deployment, "eval").await, 1);
|
||||
assert_eq!(support::calls(&deployment, "evalsha").await, 0);
|
||||
}
|
||||
1083
litellm-rust/crates/redis-store/tests/rate_limits.rs
Normal file
1083
litellm-rust/crates/redis-store/tests/rate_limits.rs
Normal file
File diff suppressed because it is too large
Load diff
512
litellm-rust/crates/redis-store/tests/router_cooldowns.rs
Normal file
512
litellm-rust/crates/redis-store/tests/router_cooldowns.rs
Normal file
|
|
@ -0,0 +1,512 @@
|
|||
mod support;
|
||||
|
||||
use std::time::Duration;
|
||||
|
||||
use litellm_redis_client::commands;
|
||||
use litellm_redis_client::{Operation, Ttl};
|
||||
use litellm_redis_store::Send as RedisSend;
|
||||
use litellm_redis_store::{
|
||||
Number, RequestRedis,
|
||||
router_cooldowns::{Cooldown, RouterCooldowns},
|
||||
};
|
||||
use litellm_redis_testing::{RedisDeployment, Topology};
|
||||
use rstest::rstest;
|
||||
use support::{ALLOWED_FAILS, COOLDOWN_POLICY as POLICY, connect, cooldowns, remote, within};
|
||||
use tokio::runtime::Handle;
|
||||
|
||||
async fn read_active(
|
||||
store: &std::sync::Arc<RouterCooldowns>,
|
||||
model_ids: &[String],
|
||||
) -> Vec<(String, Cooldown)> {
|
||||
store
|
||||
.active(RedisSend::Now, model_ids)
|
||||
.unwrap()
|
||||
.await
|
||||
.unwrap()
|
||||
}
|
||||
|
||||
/// A clock fixed just after the fixture timestamps, so every fixture cooldown is still active.
|
||||
fn fixture_clock() -> litellm_redis_store::Clock {
|
||||
std::sync::Arc::new(|| Duration::from_secs_f64(1_730_000_000.75))
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn put_writes_the_bytes_mains_set_path_writes_with_whole_second_ttl(
|
||||
#[values(Topology::Standalone, Topology::Cluster)] topology: Topology,
|
||||
#[values("json", "str_dict")] format: &str,
|
||||
) {
|
||||
let deployment = RedisDeployment::start(topology).await.unwrap();
|
||||
let raw = connect(&deployment);
|
||||
let store = std::sync::Arc::new(RouterCooldowns::new(
|
||||
Some(remote(connect(&deployment))),
|
||||
POLICY,
|
||||
));
|
||||
|
||||
for case in cooldowns() {
|
||||
let ttl = Duration::from_secs_f64(case.ttl_seconds);
|
||||
let (written_by_put, expected) = match format {
|
||||
"json" => (
|
||||
store.put(case.model_id, &case.cooldown, ttl).await,
|
||||
case.json,
|
||||
),
|
||||
_ => (
|
||||
store.put_str(case.model_id, &case.cooldown, ttl).await,
|
||||
case.str_dict,
|
||||
),
|
||||
};
|
||||
written_by_put.unwrap();
|
||||
let key = case.key;
|
||||
|
||||
let written: Option<String> = raw.send(commands::get(key)).await.unwrap();
|
||||
let Ttl::Expires(remaining) = raw.send(commands::ttl(key)).await.unwrap() else {
|
||||
panic!("{key} has no TTL");
|
||||
};
|
||||
assert_eq!(written.as_deref(), Some(expected));
|
||||
assert!(within(remaining, case.redis_ttl_seconds), "{remaining:?}");
|
||||
}
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn reads_cooldowns_python_wrote_as_str_dict_or_json(
|
||||
#[values(Topology::Standalone, Topology::Cluster)] topology: Topology,
|
||||
#[values("str_dict", "json")] format: &str,
|
||||
) {
|
||||
let deployment = RedisDeployment::start(topology).await.unwrap();
|
||||
let raw = connect(&deployment);
|
||||
let store = std::sync::Arc::new(RouterCooldowns::with_clock(
|
||||
Some(remote(connect(&deployment))),
|
||||
POLICY,
|
||||
fixture_clock(),
|
||||
));
|
||||
|
||||
for case in cooldowns() {
|
||||
let bytes = match format {
|
||||
"json" => case.json,
|
||||
_ => case.str_dict,
|
||||
};
|
||||
raw.send(commands::set(case.key, bytes, Duration::from_secs(60)))
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(
|
||||
read_active(&store, &[case.model_id.to_owned()]).await,
|
||||
vec![(case.model_id.to_owned(), case.cooldown)]
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn an_ended_cooldown_is_not_active_but_still_counts_toward_the_minimum(
|
||||
#[values(Topology::Standalone, Topology::Cluster)] topology: Topology,
|
||||
) {
|
||||
let deployment = RedisDeployment::start(topology).await.unwrap();
|
||||
let raw = connect(&deployment);
|
||||
let late = std::sync::Arc::new(|| Duration::from_secs(1_730_000_100));
|
||||
let store = std::sync::Arc::new(RouterCooldowns::with_clock(
|
||||
Some(remote(connect(&deployment))),
|
||||
POLICY,
|
||||
late,
|
||||
));
|
||||
let models = cooldowns()
|
||||
.iter()
|
||||
.map(|case| case.model_id.to_owned())
|
||||
.collect::<Vec<_>>();
|
||||
for case in cooldowns() {
|
||||
raw.send(commands::set(
|
||||
case.key,
|
||||
case.str_dict,
|
||||
Duration::from_secs(60),
|
||||
))
|
||||
.await
|
||||
.unwrap();
|
||||
}
|
||||
|
||||
assert_eq!(
|
||||
store.minimum(&models).await.unwrap(),
|
||||
Some(Number::Float(0.5))
|
||||
);
|
||||
assert_eq!(read_active(&store, &models).await, vec![]);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn allowed_fails_count_across_async_and_sync_routers_and_keep_the_first_ttl(
|
||||
#[values(Topology::Standalone, Topology::Cluster)] topology: Topology,
|
||||
) {
|
||||
let deployment = RedisDeployment::start(topology).await.unwrap();
|
||||
let raw = connect(&deployment);
|
||||
let first = std::sync::Arc::new(RouterCooldowns::new(
|
||||
Some(remote(connect(&deployment))),
|
||||
POLICY,
|
||||
));
|
||||
let second = std::sync::Arc::new(RouterCooldowns::new(
|
||||
Some(remote(connect(&deployment))),
|
||||
POLICY,
|
||||
));
|
||||
|
||||
for case in ALLOWED_FAILS {
|
||||
let key = case.key;
|
||||
raw.send(commands::del(&[key])).await.unwrap();
|
||||
let suffix = case.suffix;
|
||||
|
||||
let counts = [
|
||||
first
|
||||
.increment_allowed_fails("dep-a", suffix, Duration::from_secs_f64(30.9))
|
||||
.await
|
||||
.unwrap(),
|
||||
second
|
||||
.increment_allowed_fails_int("dep-a", suffix, Duration::from_secs(120))
|
||||
.await
|
||||
.unwrap(),
|
||||
];
|
||||
let Ttl::Expires(remaining) = raw.send(commands::ttl(key)).await.unwrap() else {
|
||||
panic!("{key} has no TTL");
|
||||
};
|
||||
|
||||
assert_eq!(counts, [1, 2], "{key}");
|
||||
assert!(within(remaining, 30), "{remaining:?}");
|
||||
}
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::async_increment_cache_sends_incrbyfloat("INCRBYFLOAT")]
|
||||
#[case::increment_cache_sends_incrby("INCRBY")]
|
||||
#[tokio::test]
|
||||
async fn allowed_fails_send_the_increment_mains_path_sends(
|
||||
#[case] command: &'static str,
|
||||
#[values(Topology::Standalone, Topology::Cluster)] topology: Topology,
|
||||
) {
|
||||
let deployment = RedisDeployment::start(topology).await.unwrap();
|
||||
let (client, trips) = support::connect_observed(&deployment);
|
||||
let store = std::sync::Arc::new(RouterCooldowns::new(Some(remote(client)), POLICY));
|
||||
let ttl = Duration::from_secs(60);
|
||||
|
||||
let count = match command {
|
||||
"INCRBYFLOAT" => store.increment_allowed_fails("dep-a", None, ttl).await,
|
||||
_ => store.increment_allowed_fails_int("dep-a", None, ttl).await,
|
||||
};
|
||||
|
||||
assert_eq!(count.unwrap(), 1);
|
||||
assert_eq!(trips.seen().first(), Some(&Operation::Command(command)));
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn python_written_fail_counts_continue_in_rust(
|
||||
#[values(Topology::Standalone, Topology::Cluster)] topology: Topology,
|
||||
) {
|
||||
let deployment = RedisDeployment::start(topology).await.unwrap();
|
||||
let raw = connect(&deployment);
|
||||
let store = std::sync::Arc::new(RouterCooldowns::new(
|
||||
Some(remote(connect(&deployment))),
|
||||
POLICY,
|
||||
));
|
||||
raw.send(commands::incr_by(
|
||||
"deployment:dep-a:allowed_fails",
|
||||
4,
|
||||
Duration::from_secs(60),
|
||||
))
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(
|
||||
store
|
||||
.increment_allowed_fails("dep-a", None, Duration::from_secs(60))
|
||||
.await
|
||||
.unwrap(),
|
||||
5
|
||||
);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn a_router_without_redis_keeps_its_cooldowns_and_counts_in_memory() {
|
||||
let store = std::sync::Arc::new(RouterCooldowns::with_clock(None, POLICY, fixture_clock()));
|
||||
let [case, _] = cooldowns();
|
||||
|
||||
store
|
||||
.put(case.model_id, &case.cooldown, Duration::from_secs(90))
|
||||
.await
|
||||
.unwrap();
|
||||
let fails = [
|
||||
store
|
||||
.increment_allowed_fails("dep-a", None, Duration::from_secs(60))
|
||||
.await
|
||||
.unwrap(),
|
||||
store
|
||||
.increment_allowed_fails("dep-a", None, Duration::from_secs(60))
|
||||
.await
|
||||
.unwrap(),
|
||||
];
|
||||
store.flush_local_counters().unwrap();
|
||||
|
||||
assert_eq!(
|
||||
read_active(&store, &[case.model_id.to_owned()]).await,
|
||||
vec![(case.model_id.to_owned(), case.cooldown)]
|
||||
);
|
||||
assert_eq!(fails, [1, 2]);
|
||||
assert_eq!(
|
||||
store
|
||||
.increment_allowed_fails("dep-a", None, Duration::from_secs(60))
|
||||
.await
|
||||
.unwrap(),
|
||||
1
|
||||
);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn a_redis_outage_serves_memory_and_counts_fails_locally() {
|
||||
let deployment = RedisDeployment::start(Topology::Standalone).await.unwrap();
|
||||
let store = std::sync::Arc::new(RouterCooldowns::with_clock(
|
||||
Some(remote(connect(&deployment))),
|
||||
POLICY,
|
||||
fixture_clock(),
|
||||
));
|
||||
let [case, _] = cooldowns();
|
||||
store
|
||||
.put(case.model_id, &case.cooldown, Duration::from_secs(90))
|
||||
.await
|
||||
.unwrap();
|
||||
deployment.stop().await.unwrap();
|
||||
|
||||
let active = read_active(
|
||||
&store,
|
||||
&[case.model_id.to_owned(), "dep-missing".to_owned()],
|
||||
)
|
||||
.await;
|
||||
let fails = store
|
||||
.increment_allowed_fails("dep-a", None, Duration::from_secs(60))
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(active, vec![(case.model_id.to_owned(), case.cooldown)]);
|
||||
assert_eq!(fails, 1);
|
||||
}
|
||||
|
||||
struct StepClock(std::sync::Mutex<Duration>);
|
||||
|
||||
impl StepClock {
|
||||
fn at(seconds: f64) -> std::sync::Arc<Self> {
|
||||
std::sync::Arc::new(Self(std::sync::Mutex::new(Duration::from_secs_f64(
|
||||
seconds,
|
||||
))))
|
||||
}
|
||||
|
||||
fn advance(&self, by: Duration) {
|
||||
*self.0.lock().unwrap() += by;
|
||||
}
|
||||
|
||||
fn clock(self: &std::sync::Arc<Self>) -> litellm_redis_store::Clock {
|
||||
let this = self.clone();
|
||||
std::sync::Arc::new(move || *this.0.lock().unwrap())
|
||||
}
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn a_miss_is_read_from_redis_again_only_after_the_one_second_refresh(
|
||||
#[values(Topology::Standalone, Topology::Cluster)] topology: Topology,
|
||||
) {
|
||||
let deployment = RedisDeployment::start(topology).await.unwrap();
|
||||
let raw = connect(&deployment);
|
||||
let time = StepClock::at(1_730_000_000.75);
|
||||
let store = std::sync::Arc::new(RouterCooldowns::with_clock(
|
||||
Some(remote(connect(&deployment))),
|
||||
POLICY,
|
||||
time.clock(),
|
||||
));
|
||||
let [case, _] = cooldowns();
|
||||
|
||||
let before = read_active(&store, &[case.model_id.to_owned()]).await;
|
||||
raw.send(commands::set(
|
||||
case.key,
|
||||
case.str_dict,
|
||||
Duration::from_secs(60),
|
||||
))
|
||||
.await
|
||||
.unwrap();
|
||||
let throttled = read_active(&store, &[case.model_id.to_owned()]).await;
|
||||
time.advance(Duration::from_secs(1));
|
||||
let refreshed = read_active(&store, &[case.model_id.to_owned()]).await;
|
||||
|
||||
assert_eq!(before, vec![]);
|
||||
assert_eq!(throttled, vec![]);
|
||||
assert_eq!(
|
||||
refreshed,
|
||||
vec![(case.model_id.to_owned(), case.cooldown.clone())]
|
||||
);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn a_breaker_refusal_does_not_hold_back_the_next_read() {
|
||||
let deployment = RedisDeployment::start(Topology::Standalone).await.unwrap();
|
||||
let raw = connect(&deployment);
|
||||
let breaker = litellm_circuit_breaker::Breaker::new(
|
||||
litellm_circuit_breaker::Config {
|
||||
failure_threshold: 1,
|
||||
recovery_timeout: Duration::from_millis(200),
|
||||
..litellm_circuit_breaker::Config::default()
|
||||
},
|
||||
Box::new(|_| {}),
|
||||
);
|
||||
let time = StepClock::at(1_730_000_000.75);
|
||||
let guarded = support::connect_guarded(&deployment, breaker);
|
||||
let store = std::sync::Arc::new(RouterCooldowns::with_clock(
|
||||
Some(remote(guarded.clone())),
|
||||
POLICY,
|
||||
time.clock(),
|
||||
));
|
||||
let [case, _] = cooldowns();
|
||||
deployment.stop().await.unwrap();
|
||||
read_active(&store, &[case.model_id.to_owned()]).await;
|
||||
time.advance(Duration::from_secs(1));
|
||||
let refused = read_active(&store, &[case.model_id.to_owned()]).await;
|
||||
deployment.resume().await.unwrap();
|
||||
litellm_redis_testing::eventually(async || guarded.ping().await.ok()).await;
|
||||
raw.send(commands::set(
|
||||
case.key,
|
||||
case.str_dict,
|
||||
Duration::from_secs(60),
|
||||
))
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
let recovered = read_active(&store, &[case.model_id.to_owned()]).await;
|
||||
|
||||
assert_eq!(refused, vec![]);
|
||||
assert_eq!(
|
||||
recovered,
|
||||
vec![(case.model_id.to_owned(), case.cooldown.clone())]
|
||||
);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn a_declared_active_read_rides_the_request_batch_and_answers_what_active_does(
|
||||
#[values(Topology::Standalone, Topology::Cluster)] topology: Topology,
|
||||
) {
|
||||
let deployment = RedisDeployment::start(topology).await.unwrap();
|
||||
let raw = connect(&deployment);
|
||||
let (client, trips) = support::connect_observed(&deployment);
|
||||
let store = std::sync::Arc::new(RouterCooldowns::with_clock(
|
||||
Some(remote(client)),
|
||||
POLICY,
|
||||
fixture_clock(),
|
||||
));
|
||||
let models = cooldowns().map(|case| case.model_id.to_owned());
|
||||
for case in cooldowns() {
|
||||
raw.send(commands::set(
|
||||
case.key,
|
||||
case.str_dict,
|
||||
Duration::from_secs(60),
|
||||
))
|
||||
.await
|
||||
.unwrap();
|
||||
}
|
||||
let request = RequestRedis::new(Handle::current(), Duration::from_secs(60));
|
||||
|
||||
let declared = store.active(RedisSend::PreCall(&request), &models).unwrap();
|
||||
assert!(trips.seen().is_empty());
|
||||
request.flush_pre().await;
|
||||
let active = declared.await.unwrap();
|
||||
|
||||
let one_mget = match topology {
|
||||
Topology::Standalone => Operation::Pipeline { ops: 1 },
|
||||
Topology::Cluster => Operation::Command("MGET"),
|
||||
};
|
||||
assert_eq!(trips.seen(), [one_mget]);
|
||||
assert_eq!(
|
||||
active,
|
||||
cooldowns()
|
||||
.map(|case| (case.model_id.to_owned(), case.cooldown))
|
||||
.to_vec()
|
||||
);
|
||||
assert_eq!(read_active(&store, &models).await, active);
|
||||
assert_eq!(trips.seen().len(), 1, "the declared read backfilled memory");
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn a_cooldown_recorded_after_the_declaration_wins_over_the_redis_miss(
|
||||
#[values(Topology::Standalone, Topology::Cluster)] topology: Topology,
|
||||
) {
|
||||
let deployment = RedisDeployment::start(topology).await.unwrap();
|
||||
let store = std::sync::Arc::new(RouterCooldowns::with_clock(
|
||||
Some(remote(connect(&deployment))),
|
||||
POLICY,
|
||||
fixture_clock(),
|
||||
));
|
||||
let [case, _] = cooldowns();
|
||||
let request = RequestRedis::new(Handle::current(), Duration::from_secs(60));
|
||||
|
||||
let declared = store
|
||||
.active(RedisSend::PreCall(&request), &[case.model_id.to_owned()])
|
||||
.unwrap();
|
||||
request.flush_pre().await;
|
||||
store
|
||||
.put(case.model_id, &case.cooldown, Duration::from_secs(90))
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(
|
||||
declared.await.unwrap(),
|
||||
vec![(case.model_id.to_owned(), case.cooldown)]
|
||||
);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn a_breaker_refusal_of_a_declared_read_does_not_hold_back_the_next_read() {
|
||||
let deployment = RedisDeployment::start(Topology::Standalone).await.unwrap();
|
||||
let raw = connect(&deployment);
|
||||
let breaker = litellm_circuit_breaker::Breaker::new(
|
||||
litellm_circuit_breaker::Config {
|
||||
failure_threshold: 1,
|
||||
recovery_timeout: Duration::from_millis(200),
|
||||
..litellm_circuit_breaker::Config::default()
|
||||
},
|
||||
Box::new(|_| {}),
|
||||
);
|
||||
let guarded = support::connect_guarded(&deployment, breaker);
|
||||
let store = std::sync::Arc::new(RouterCooldowns::with_clock(
|
||||
Some(remote(guarded.clone())),
|
||||
POLICY,
|
||||
fixture_clock(),
|
||||
));
|
||||
let [case, _] = cooldowns();
|
||||
let models = [case.model_id.to_owned()];
|
||||
deployment.stop().await.unwrap();
|
||||
let opened = RequestRedis::new(Handle::current(), Duration::from_secs(60));
|
||||
let failed = store
|
||||
.active(RedisSend::PreCall(&opened), &["dep-other".to_owned()])
|
||||
.unwrap();
|
||||
opened.flush_pre().await;
|
||||
failed.await.unwrap();
|
||||
let refused_request = RequestRedis::new(Handle::current(), Duration::from_secs(60));
|
||||
|
||||
let refused = store
|
||||
.active(RedisSend::PreCall(&refused_request), &models)
|
||||
.unwrap();
|
||||
refused_request.flush_pre().await;
|
||||
let refused = refused.await.unwrap();
|
||||
deployment.resume().await.unwrap();
|
||||
litellm_redis_testing::eventually(async || guarded.ping().await.ok()).await;
|
||||
raw.send(commands::set(
|
||||
case.key,
|
||||
case.str_dict,
|
||||
Duration::from_secs(60),
|
||||
))
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(refused, vec![]);
|
||||
assert_eq!(
|
||||
read_active(&store, &models).await,
|
||||
vec![(case.model_id.to_owned(), case.cooldown.clone())]
|
||||
);
|
||||
}
|
||||
329
litellm-rust/crates/redis-store/tests/router_usage.rs
Normal file
329
litellm-rust/crates/redis-store/tests/router_usage.rs
Normal file
|
|
@ -0,0 +1,329 @@
|
|||
mod support;
|
||||
|
||||
use std::{
|
||||
sync::{
|
||||
Arc,
|
||||
atomic::{AtomicU64, Ordering},
|
||||
},
|
||||
time::Duration,
|
||||
};
|
||||
|
||||
use litellm_redis_client::commands;
|
||||
use litellm_redis_client::{Amount, Expiry, Operation, Ttl};
|
||||
use litellm_redis_store::{
|
||||
Number, TierPolicy,
|
||||
router_usage::{COUNTER_TTL, Counter, Kind, RouterUsage},
|
||||
};
|
||||
use litellm_redis_testing::{RedisDeployment, Topology};
|
||||
use rstest::rstest;
|
||||
use support::{USAGE_POLICY as POLICY, connect, remote, usage, within};
|
||||
|
||||
fn tpm_and_rpm(model: Option<&str>) -> [Counter; 2] {
|
||||
[Kind::Tpm, Kind::Rpm].map(|kind| Counter {
|
||||
deployment_id: "dep".to_owned(),
|
||||
model: model.map(str::to_owned),
|
||||
minute: "12-34".to_owned(),
|
||||
kind,
|
||||
})
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::tpm(Kind::Tpm, "tpm")]
|
||||
#[case::rpm(Kind::Rpm, "rpm")]
|
||||
fn a_counter_kind_round_trips_through_its_key_segment(#[case] kind: Kind, #[case] segment: &str) {
|
||||
assert_eq!(<&str>::from(kind), segment);
|
||||
assert_eq!(segment.parse::<Kind>().unwrap(), kind);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::unknown("tokens")]
|
||||
#[case::wrong_case("TPM")]
|
||||
fn other_text_is_not_a_counter_kind(#[case] segment: &str) {
|
||||
assert!(segment.parse::<Kind>().is_err());
|
||||
}
|
||||
|
||||
// DualCache._backfill_kwargs: a router's batch read keeps a Redis value in memory for
|
||||
// litellm.default_in_memory_ttl when one is set, else for InMemoryCache's own 600 s
|
||||
#[rstest]
|
||||
#[case::configured(Some(Duration::from_secs(5)), None)]
|
||||
#[case::unset(None, Some(Number::Int(3)))]
|
||||
#[tokio::test]
|
||||
async fn a_batch_read_backfills_memory_for_default_in_memory_ttl(
|
||||
#[case] default_in_memory_ttl: Option<Duration>,
|
||||
#[case] after_six_seconds: Option<Number>,
|
||||
) {
|
||||
let deployment = RedisDeployment::start(Topology::Standalone).await.unwrap();
|
||||
let raw = connect(&deployment);
|
||||
let seconds = Arc::new(AtomicU64::new(1_000));
|
||||
let now = {
|
||||
let seconds = seconds.clone();
|
||||
Arc::new(move || Duration::from_secs(seconds.load(Ordering::SeqCst)))
|
||||
};
|
||||
let policy = TierPolicy {
|
||||
refresh_after: Duration::from_secs(3600),
|
||||
default_in_memory_ttl,
|
||||
..POLICY
|
||||
};
|
||||
let store = RouterUsage::with_clock(Some(remote(connect(&deployment))), policy, now);
|
||||
let [case, _] = usage();
|
||||
raw.send(commands::incr_by_float(
|
||||
case.key,
|
||||
Amount::Float(3.0),
|
||||
Some(Expiry::Refresh(COUNTER_TTL)),
|
||||
))
|
||||
.await
|
||||
.unwrap();
|
||||
store
|
||||
.get_many(std::slice::from_ref(&case.counter))
|
||||
.await
|
||||
.unwrap();
|
||||
raw.send(commands::del(&[case.key])).await.unwrap();
|
||||
seconds.fetch_add(6, Ordering::SeqCst);
|
||||
|
||||
let read = store.get_many(&[case.counter]).await.unwrap();
|
||||
|
||||
assert_eq!(read, vec![after_six_seconds]);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn reads_counters_python_incremented_under_mains_keys(
|
||||
#[values(Topology::Standalone, Topology::Cluster)] topology: Topology,
|
||||
) {
|
||||
let deployment = RedisDeployment::start(topology).await.unwrap();
|
||||
let raw = connect(&deployment);
|
||||
let store = RouterUsage::new(Some(remote(connect(&deployment))), POLICY);
|
||||
let cases = usage();
|
||||
for (case, delta) in cases.iter().zip([7.75, 3.0]) {
|
||||
raw.send(commands::incr_by_float(
|
||||
case.key,
|
||||
Amount::Float(delta),
|
||||
Some(Expiry::Refresh(COUNTER_TTL)),
|
||||
))
|
||||
.await
|
||||
.unwrap();
|
||||
}
|
||||
|
||||
assert_eq!(
|
||||
store
|
||||
.get_many(&cases.map(|case| case.counter))
|
||||
.await
|
||||
.unwrap(),
|
||||
vec![Some(Number::Float(7.75)), Some(Number::Int(3))]
|
||||
);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn deployment_increments_write_incrbyfloat_totals_with_a_refreshed_minute_window(
|
||||
#[values(Topology::Standalone, Topology::Cluster)] topology: Topology,
|
||||
) {
|
||||
let deployment = RedisDeployment::start(topology).await.unwrap();
|
||||
let raw = connect(&deployment);
|
||||
let writer = RouterUsage::new(Some(remote(connect(&deployment))), POLICY);
|
||||
let reader = RouterUsage::new(Some(remote(connect(&deployment))), POLICY);
|
||||
let [case, _] = usage();
|
||||
let model = case.counter.model.as_deref();
|
||||
let [tpm_key, rpm_key] =
|
||||
["tpm", "rpm"].map(|kind| case.key.replace(":tpm:", &format!(":{kind}:")));
|
||||
|
||||
writer
|
||||
.increment_deployment("dep", model, "12-34", Number::Float(5.5), 1)
|
||||
.await
|
||||
.unwrap();
|
||||
raw.send(commands::expire(&tpm_key, Duration::from_secs(5)))
|
||||
.await
|
||||
.unwrap();
|
||||
writer
|
||||
.increment_deployment("dep", model, "12-34", Number::Int(2), 1)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
let written: [Option<String>; 2] = [
|
||||
raw.send(commands::get(&tpm_key)).await.unwrap(),
|
||||
raw.send(commands::get(&rpm_key)).await.unwrap(),
|
||||
];
|
||||
let Ttl::Expires(remaining) = raw.send(commands::ttl(&tpm_key)).await.unwrap() else {
|
||||
panic!("{tpm_key} has no TTL");
|
||||
};
|
||||
assert_eq!(written, [Some("7.5".to_owned()), Some("2".to_owned())]);
|
||||
assert!(within(remaining, 60), "{remaining:?}");
|
||||
assert_eq!(
|
||||
reader.get_many(&tpm_and_rpm(model)).await.unwrap(),
|
||||
vec![Some(Number::Float(7.5)), Some(Number::Int(2))]
|
||||
);
|
||||
assert_eq!(
|
||||
writer.get_many(&tpm_and_rpm(model)).await.unwrap(),
|
||||
vec![Some(Number::Int(7)), Some(Number::Int(2))]
|
||||
);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn a_request_increment_keeps_the_window_python_started(
|
||||
#[values(Topology::Standalone, Topology::Cluster)] topology: Topology,
|
||||
) {
|
||||
let deployment = RedisDeployment::start(topology).await.unwrap();
|
||||
let raw = connect(&deployment);
|
||||
let store = RouterUsage::new(Some(remote(connect(&deployment))), POLICY);
|
||||
let [_, case] = usage();
|
||||
let key = case.key;
|
||||
raw.send(commands::incr_by_float(
|
||||
key,
|
||||
Amount::Float(1.0),
|
||||
Some(Expiry::KeepExisting(Duration::from_secs(20))),
|
||||
))
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
let total = store
|
||||
.increment(&case.counter, Number::Int(1), COUNTER_TTL)
|
||||
.await
|
||||
.unwrap();
|
||||
let Ttl::Expires(remaining) = raw.send(commands::ttl(key)).await.unwrap() else {
|
||||
panic!("{key} has no TTL");
|
||||
};
|
||||
|
||||
assert_eq!(total, Number::Float(2.0));
|
||||
assert!(within(remaining, 20), "{remaining:?}");
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn without_redis_counters_accumulate_in_memory_until_flushed() {
|
||||
let store = RouterUsage::new(None, POLICY);
|
||||
let counters = tpm_and_rpm(Some("model"));
|
||||
|
||||
store
|
||||
.increment_deployment("dep", Some("model"), "12-34", Number::Int(2), 1)
|
||||
.await
|
||||
.unwrap();
|
||||
store
|
||||
.increment_deployment("dep", Some("model"), "12-34", Number::Float(5.5), 1)
|
||||
.await
|
||||
.unwrap();
|
||||
let counted = store.get_many(&counters).await.unwrap();
|
||||
store.flush_local_counters().unwrap();
|
||||
|
||||
assert_eq!(
|
||||
counted,
|
||||
vec![Some(Number::Float(7.5)), Some(Number::Int(2))]
|
||||
);
|
||||
assert_eq!(store.get_many(&counters).await.unwrap(), vec![None, None]);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn a_late_attach_keeps_memory_and_later_attaches_are_ignored(
|
||||
#[values(Topology::Standalone, Topology::Cluster)] topology: Topology,
|
||||
) {
|
||||
let attached = RedisDeployment::start(topology).await.unwrap();
|
||||
let ignored = RedisDeployment::start(topology).await.unwrap();
|
||||
let store = RouterUsage::new(None, POLICY);
|
||||
let [case, _] = usage();
|
||||
store
|
||||
.increment(&case.counter, Number::Int(2), COUNTER_TTL)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
assert!(store.attach(remote(connect(&attached))));
|
||||
assert!(!store.attach(remote(connect(&ignored))));
|
||||
let total = store
|
||||
.increment(&case.counter, Number::Int(3), COUNTER_TTL)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(total, Number::Float(3.0));
|
||||
assert_eq!(
|
||||
store
|
||||
.get_many(std::slice::from_ref(&case.counter))
|
||||
.await
|
||||
.unwrap(),
|
||||
vec![Some(Number::Int(5))]
|
||||
);
|
||||
let ignored_total: Option<String> = connect(&ignored)
|
||||
.send(commands::get(case.key))
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(ignored_total, None);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn a_failed_pipeline_keeps_the_memory_counts_as_integers() {
|
||||
let deployment = RedisDeployment::start(Topology::Standalone).await.unwrap();
|
||||
let store = RouterUsage::new(Some(remote(connect(&deployment))), POLICY);
|
||||
deployment.stop().await.unwrap();
|
||||
|
||||
store
|
||||
.increment_deployment("dep", Some("model"), "12-34", Number::Float(5.75), 1)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(
|
||||
store.get_many(&tpm_and_rpm(Some("model"))).await.unwrap(),
|
||||
vec![Some(Number::Int(5)), Some(Number::Int(1))]
|
||||
);
|
||||
}
|
||||
|
||||
// main's _increment_deployment sends async_increment_pipeline on its own round trip, never on the
|
||||
// request's post-call pipeline
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn a_deployment_increment_is_its_own_pipeline_and_mirrors_the_redis_totals(
|
||||
#[values(Topology::Standalone, Topology::Cluster)] topology: Topology,
|
||||
) {
|
||||
let deployment = RedisDeployment::start(topology).await.unwrap();
|
||||
let raw = connect(&deployment);
|
||||
let (client, trips) = support::connect_observed(&deployment);
|
||||
let store = RouterUsage::new(Some(remote(client)), POLICY);
|
||||
let counters = tpm_and_rpm(Some("model"));
|
||||
let [tpm_key, rpm_key] =
|
||||
["tpm", "rpm"].map(|kind| format!("global_router:dep:model:{kind}:12-34"));
|
||||
raw.send(commands::incr_by_float(
|
||||
&tpm_key,
|
||||
Amount::Float(10.0),
|
||||
Some(Expiry::Refresh(COUNTER_TTL)),
|
||||
))
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
store
|
||||
.increment_deployment("dep", Some("model"), "12-34", Number::Float(5.5), 1)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(trips.seen(), [Operation::Pipeline { ops: 2 }]);
|
||||
assert_eq!(
|
||||
store.get_many(&counters).await.unwrap(),
|
||||
vec![Some(Number::Int(15)), Some(Number::Int(1))]
|
||||
);
|
||||
let written: [Option<String>; 2] = [
|
||||
raw.send(commands::get(&tpm_key)).await.unwrap(),
|
||||
raw.send(commands::get(&rpm_key)).await.unwrap(),
|
||||
];
|
||||
assert_eq!(written, [Some("15.5".to_owned()), Some("1".to_owned())]);
|
||||
let Ttl::Expires(remaining) = raw.send(commands::ttl(&rpm_key)).await.unwrap() else {
|
||||
panic!("{rpm_key} has no TTL");
|
||||
};
|
||||
assert!(within(remaining, 60), "{remaining:?}");
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn a_failed_deployment_pipeline_keeps_the_memory_counts_as_integers() {
|
||||
let deployment = RedisDeployment::start(Topology::Standalone).await.unwrap();
|
||||
let store = RouterUsage::new(Some(remote(connect(&deployment))), POLICY);
|
||||
let counters = tpm_and_rpm(Some("model"));
|
||||
deployment.stop().await.unwrap();
|
||||
|
||||
store
|
||||
.increment_deployment("dep", Some("model"), "12-34", Number::Float(5.75), 1)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(
|
||||
store.get_many(&counters).await.unwrap(),
|
||||
[Some(Number::Int(5)), Some(Number::Int(1))]
|
||||
);
|
||||
}
|
||||
134
litellm-rust/crates/redis-store/tests/sensitive_route_pins.rs
Normal file
134
litellm-rust/crates/redis-store/tests/sensitive_route_pins.rs
Normal file
|
|
@ -0,0 +1,134 @@
|
|||
mod support;
|
||||
|
||||
use std::{
|
||||
sync::{Arc, Mutex},
|
||||
time::Duration,
|
||||
};
|
||||
|
||||
use litellm_redis_client::{Ttl, commands};
|
||||
use litellm_redis_store::{Clock, sensitive_route_pins::SensitiveRoutePins};
|
||||
use litellm_redis_testing::{RedisDeployment, Topology};
|
||||
use rstest::rstest;
|
||||
use support::{connect, remote, within};
|
||||
|
||||
const PIN_KEY: &str = "{sensitive_route:tenant-a:session-1}:model";
|
||||
const HOUR: Duration = Duration::from_secs(3600);
|
||||
|
||||
struct StepClock(Mutex<Duration>);
|
||||
|
||||
impl StepClock {
|
||||
fn start() -> Arc<Self> {
|
||||
Arc::new(Self(Mutex::new(Duration::from_secs(1_000))))
|
||||
}
|
||||
|
||||
fn advance(&self, by: Duration) {
|
||||
*self.0.lock().unwrap() += by;
|
||||
}
|
||||
|
||||
fn clock(self: &Arc<Self>) -> Clock {
|
||||
let this = self.clone();
|
||||
Arc::new(move || *this.0.lock().unwrap())
|
||||
}
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn a_pin_one_pod_writes_is_read_by_another_under_mains_key_and_encoding(
|
||||
#[values(Topology::Standalone, Topology::Cluster)] topology: Topology,
|
||||
) {
|
||||
let deployment = RedisDeployment::start(topology).await.unwrap();
|
||||
let raw = connect(&deployment);
|
||||
let writer = SensitiveRoutePins::new(Some(remote(connect(&deployment))));
|
||||
let reader = SensitiveRoutePins::new(Some(remote(connect(&deployment))));
|
||||
|
||||
writer
|
||||
.set("tenant-a", "session-1", "on-prem-model", HOUR)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
let written: Option<String> = raw.send(commands::get(PIN_KEY)).await.unwrap();
|
||||
let Ttl::Expires(remaining) = raw.send(commands::ttl(PIN_KEY)).await.unwrap() else {
|
||||
panic!("{PIN_KEY} has no TTL");
|
||||
};
|
||||
assert_eq!(written.as_deref(), Some("\"on-prem-model\""));
|
||||
assert!(within(remaining, 3600), "{remaining:?}");
|
||||
assert_eq!(
|
||||
reader.get("tenant-a", "session-1", HOUR).await.unwrap(),
|
||||
Some("on-prem-model".to_owned())
|
||||
);
|
||||
assert_eq!(
|
||||
reader.get("tenant-b", "session-1", HOUR).await.unwrap(),
|
||||
None
|
||||
);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn a_pin_read_from_redis_stays_in_memory_for_its_remaining_ttl_only() {
|
||||
let deployment = RedisDeployment::start(Topology::Standalone).await.unwrap();
|
||||
let raw = connect(&deployment);
|
||||
let time = StepClock::start();
|
||||
let reader = SensitiveRoutePins::with_clock(Some(remote(connect(&deployment))), time.clock());
|
||||
raw.send(commands::set(
|
||||
PIN_KEY,
|
||||
"\"on-prem-model\"",
|
||||
Duration::from_secs(30),
|
||||
))
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
let read = reader.get("tenant-a", "session-1", HOUR).await.unwrap();
|
||||
deployment.stop().await.unwrap();
|
||||
time.advance(Duration::from_secs(20));
|
||||
let before_expiry = reader.get("tenant-a", "session-1", HOUR).await.unwrap();
|
||||
time.advance(Duration::from_secs(15));
|
||||
let after_expiry = reader.get("tenant-a", "session-1", HOUR).await.unwrap();
|
||||
|
||||
assert_eq!(read, Some("on-prem-model".to_owned()));
|
||||
assert_eq!(before_expiry, read);
|
||||
assert_eq!(after_expiry, None);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn a_redis_outage_keeps_pinning_sessions_in_memory() {
|
||||
let deployment = RedisDeployment::start(Topology::Standalone).await.unwrap();
|
||||
let store = SensitiveRoutePins::new(Some(remote(connect(&deployment))));
|
||||
store
|
||||
.set("tenant-a", "session-1", "on-prem-model", HOUR)
|
||||
.await
|
||||
.unwrap();
|
||||
deployment.stop().await.unwrap();
|
||||
|
||||
let pinned_before = store.get("tenant-a", "session-1", HOUR).await.unwrap();
|
||||
store
|
||||
.set("tenant-a", "session-2", "other-model", HOUR)
|
||||
.await
|
||||
.unwrap();
|
||||
let pinned_during = store.get("tenant-a", "session-2", HOUR).await.unwrap();
|
||||
|
||||
assert_eq!(pinned_before, Some("on-prem-model".to_owned()));
|
||||
assert_eq!(pinned_during, Some("other-model".to_owned()));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn without_redis_pins_live_in_memory_until_they_expire() {
|
||||
let time = StepClock::start();
|
||||
let store = SensitiveRoutePins::with_clock(None, time.clock());
|
||||
store
|
||||
.set(
|
||||
"tenant-a",
|
||||
"session-1",
|
||||
"on-prem-model",
|
||||
Duration::from_secs(10),
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
let pinned = store.get("tenant-a", "session-1", HOUR).await.unwrap();
|
||||
time.advance(Duration::from_secs(11));
|
||||
let expired = store.get("tenant-a", "session-1", HOUR).await.unwrap();
|
||||
|
||||
assert_eq!(pinned, Some("on-prem-model".to_owned()));
|
||||
assert_eq!(expired, None);
|
||||
}
|
||||
648
litellm-rust/crates/redis-store/tests/spend_counters.rs
Normal file
648
litellm-rust/crates/redis-store/tests/spend_counters.rs
Normal file
|
|
@ -0,0 +1,648 @@
|
|||
mod support;
|
||||
|
||||
use std::{sync::Arc, time::Duration};
|
||||
|
||||
use litellm_redis_client::commands;
|
||||
use litellm_redis_client::{Operation, RedisClient, Ttl};
|
||||
use litellm_redis_store::Send as RedisSend;
|
||||
use litellm_redis_store::{
|
||||
Number, RequestRedis,
|
||||
router_cooldowns::RouterCooldowns,
|
||||
spend_counters::{LostSeed, Reading, SpendCounter, SpendCounters},
|
||||
};
|
||||
use litellm_redis_testing::{RedisDeployment, Topology};
|
||||
use rstest::rstest;
|
||||
use support::{COOLDOWN_POLICY, connect, connect_observed, remote, within};
|
||||
use tokio::runtime::Handle;
|
||||
|
||||
const TTL: Duration = Duration::from_secs(60);
|
||||
const LONG_DEADLINE: Duration = Duration::from_secs(60);
|
||||
|
||||
const READS: [(&str, Number); 4] = [
|
||||
("7.25", Number::Float(7.25)),
|
||||
("3", Number::Int(3)),
|
||||
("0.0", Number::Float(0.0)),
|
||||
("1e-05", Number::Float(1e-05)),
|
||||
];
|
||||
|
||||
const SETS: [(Number, &str); 4] = [
|
||||
(Number::Float(0.0), "0.0"),
|
||||
(Number::Int(12), "12"),
|
||||
(Number::Float(1e-05), "1e-05"),
|
||||
(Number::Float(0.30000000000000004), "0.30000000000000004"),
|
||||
];
|
||||
|
||||
const SET_MAXES: [(Number, &str); 3] = [
|
||||
(Number::Float(2.5), "2.5"),
|
||||
(Number::Int(100), "100"),
|
||||
(Number::Float(1e-05), "1e-05"),
|
||||
];
|
||||
|
||||
fn team(id: &str) -> SpendCounter {
|
||||
SpendCounter::Team {
|
||||
team_id: id.to_owned(),
|
||||
window: None,
|
||||
}
|
||||
}
|
||||
|
||||
fn store(client: Option<RedisClient>) -> Arc<SpendCounters> {
|
||||
Arc::new(SpendCounters::new(client.map(remote)))
|
||||
}
|
||||
|
||||
/// Mirrored by `tests/unit/redis/stores/test_spend_counters.py`, which checks the same keys and
|
||||
/// bytes against Python's own key builders and `RedisCache` codec.
|
||||
#[rstest]
|
||||
#[case::redis(Reading::Redis(Some(Number::Int(3))), "redis", Some(Number::Int(3)))]
|
||||
#[case::redis_miss(Reading::Redis(None), "redis", None)]
|
||||
#[case::memory(
|
||||
Reading::Memory(Some(Number::Float(2.5))),
|
||||
"memory",
|
||||
Some(Number::Float(2.5))
|
||||
)]
|
||||
fn a_reading_names_where_it_was_answered(
|
||||
#[case] reading: Reading,
|
||||
#[case] source: &str,
|
||||
#[case] value: Option<Number>,
|
||||
) {
|
||||
assert_eq!(<&str>::from(reading), source);
|
||||
assert_eq!(reading.value(), value);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case(
|
||||
SpendCounter::Key {
|
||||
token: "88dc28d0f030c55ed4ab77ed8faf098196cb1c05df778539800c9f1243fe6b4b".to_owned(),
|
||||
window: None,
|
||||
},
|
||||
"spend:key:88dc28d0f030c55ed4ab77ed8faf098196cb1c05df778539800c9f1243fe6b4b"
|
||||
)]
|
||||
#[case(
|
||||
SpendCounter::Key {
|
||||
token: "88dc28d0".to_owned(),
|
||||
window: Some("30d".to_owned()),
|
||||
},
|
||||
"spend:key:88dc28d0:window:30d"
|
||||
)]
|
||||
#[case(
|
||||
SpendCounter::Team {
|
||||
team_id: "team-a".to_owned(),
|
||||
window: None,
|
||||
},
|
||||
"spend:team:team-a"
|
||||
)]
|
||||
#[case(
|
||||
SpendCounter::Team {
|
||||
team_id: "team-a".to_owned(),
|
||||
window: Some("1h".to_owned()),
|
||||
},
|
||||
"spend:team:team-a:window:1h"
|
||||
)]
|
||||
#[case(
|
||||
SpendCounter::TeamMember {
|
||||
user_id: "user:with:colons".to_owned(),
|
||||
team_id: "team-a".to_owned(),
|
||||
},
|
||||
"spend:team_member:user:with:colons:team-a"
|
||||
)]
|
||||
#[case(SpendCounter::User("user@example.com".to_owned()), "spend:user:user@example.com")]
|
||||
#[case(SpendCounter::EndUser("customer é 1".to_owned()), "spend:end_user:customer é 1")]
|
||||
#[case(SpendCounter::Org("org-1".to_owned()), "spend:org:org-1")]
|
||||
#[case(SpendCounter::Tag("prod:eu".to_owned()), "spend:tag:prod:eu")]
|
||||
#[case(SpendCounter::Project("proj-1".to_owned()), "spend:project:proj-1")]
|
||||
#[case(
|
||||
SpendCounter::ModelAccessGroup("gpt-group".to_owned()),
|
||||
"spend:model_access_group:gpt-group"
|
||||
)]
|
||||
#[case(SpendCounter::ShadowEval("job-7".to_owned()), "spend:shadow_eval:job-7")]
|
||||
fn every_counter_formats_and_parses_back_to_mains_key(
|
||||
#[case] counter: SpendCounter,
|
||||
#[case] key: &str,
|
||||
) {
|
||||
assert_eq!(counter.key(), key);
|
||||
assert_eq!(SpendCounter::parse(key).unwrap(), counter, "{key}");
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case("spend_db_floor:spend:key:abc")]
|
||||
#[case("virtual_key_spend:abc:gpt:1d")]
|
||||
#[case("spend:unknown:abc")]
|
||||
#[case("spend:team_member:no-team")]
|
||||
#[case("spend:")]
|
||||
fn keys_outside_the_spend_counter_keyspace_are_refused(#[case] key: &str) {
|
||||
assert!(SpendCounter::parse(key).is_err(), "{key}");
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn reads_values_python_wrote_and_a_clean_miss_is_a_redis_answer(
|
||||
#[values(Topology::Standalone, Topology::Cluster)] topology: Topology,
|
||||
) {
|
||||
let deployment = RedisDeployment::start(topology).await.unwrap();
|
||||
let raw = connect(&deployment);
|
||||
let store = store(Some(connect(&deployment)));
|
||||
let counters = (0..READS.len())
|
||||
.map(|index| team(&format!("t{index}")))
|
||||
.collect::<Vec<_>>();
|
||||
for ((bytes, _), counter) in READS.iter().zip(&counters) {
|
||||
raw.send(commands::set(&counter.key(), *bytes, TTL))
|
||||
.await
|
||||
.unwrap();
|
||||
}
|
||||
let expected = READS
|
||||
.iter()
|
||||
.map(|(_, spend)| Some(*spend))
|
||||
.collect::<Vec<_>>();
|
||||
|
||||
for (counter, value) in counters.iter().zip(&expected) {
|
||||
assert_eq!(store.read(counter).await.unwrap(), Reading::Redis(*value));
|
||||
}
|
||||
assert_eq!(
|
||||
store
|
||||
.read_many(RedisSend::Now, &counters)
|
||||
.unwrap()
|
||||
.await
|
||||
.unwrap(),
|
||||
Some(expected)
|
||||
);
|
||||
assert_eq!(
|
||||
store.read(&team("absent")).await.unwrap(),
|
||||
Reading::Redis(None)
|
||||
);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn resets_and_repairs_write_mains_bytes_with_the_redis_ttl(
|
||||
#[values(Topology::Standalone, Topology::Cluster)] topology: Topology,
|
||||
) {
|
||||
let deployment = RedisDeployment::start(topology).await.unwrap();
|
||||
let raw = connect(&deployment);
|
||||
let store = store(Some(connect(&deployment)));
|
||||
let counter = team("t");
|
||||
|
||||
for (spend, bytes) in SETS {
|
||||
store.set(&counter, spend, None).await.unwrap();
|
||||
assert_eq!(
|
||||
raw.send(commands::get::<String>(&counter.key()))
|
||||
.await
|
||||
.unwrap(),
|
||||
Some(bytes.to_owned())
|
||||
);
|
||||
}
|
||||
for (spend, bytes) in SET_MAXES {
|
||||
raw.send(commands::del(&[&counter.key()])).await.unwrap();
|
||||
let repaired = store.set_max(&counter, spend).await.unwrap();
|
||||
assert_eq!(
|
||||
raw.send(commands::get::<String>(&counter.key()))
|
||||
.await
|
||||
.unwrap(),
|
||||
Some(bytes.to_owned())
|
||||
);
|
||||
assert_eq!(repaired, Some(Number::Float(spend.as_f64())));
|
||||
}
|
||||
store
|
||||
.set(&counter, Number::Int(5), Some(Duration::from_secs(30)))
|
||||
.await
|
||||
.unwrap();
|
||||
let Ttl::Expires(remaining) = raw.send(commands::ttl(&counter.key())).await.unwrap() else {
|
||||
panic!("a reset counter has no TTL");
|
||||
};
|
||||
assert!(within(remaining, 30), "{remaining:?}");
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn set_max_only_ever_raises_a_counter(
|
||||
#[values(Topology::Standalone, Topology::Cluster)] topology: Topology,
|
||||
) {
|
||||
let deployment = RedisDeployment::start(topology).await.unwrap();
|
||||
let raw = connect(&deployment);
|
||||
let store = store(Some(connect(&deployment)));
|
||||
let counter = SpendCounter::Key {
|
||||
token: "hashed".to_owned(),
|
||||
window: None,
|
||||
};
|
||||
store.increment(&counter, Number::Float(9.5)).await.unwrap();
|
||||
|
||||
let lower = store.set_max(&counter, Number::Float(4.0)).await.unwrap();
|
||||
let higher = store.set_max(&counter, Number::Float(12.0)).await.unwrap();
|
||||
|
||||
assert_eq!(lower, Some(Number::Float(9.5)));
|
||||
assert_eq!(higher, Some(Number::Float(12.0)));
|
||||
assert_eq!(
|
||||
raw.send(commands::get::<String>(&counter.key()))
|
||||
.await
|
||||
.unwrap(),
|
||||
Some("12.0".to_owned())
|
||||
);
|
||||
let Ttl::Expires(remaining) = raw.send(commands::ttl(&counter.key())).await.unwrap() else {
|
||||
panic!("a repaired counter has no TTL");
|
||||
};
|
||||
assert!(within(remaining, 60), "{remaining:?}");
|
||||
assert_eq!(
|
||||
store.get_local(&counter).unwrap(),
|
||||
Some(Number::Float(12.0))
|
||||
);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn increments_refresh_the_ttl_and_mirror_the_shared_total_into_memory(
|
||||
#[values(Topology::Standalone, Topology::Cluster)] topology: Topology,
|
||||
) {
|
||||
let deployment = RedisDeployment::start(topology).await.unwrap();
|
||||
let raw = connect(&deployment);
|
||||
let (client, trips) = connect_observed(&deployment);
|
||||
let other_pod = store(Some(connect(&deployment)));
|
||||
let store = store(Some(client));
|
||||
let [first, second] = [team("a"), team("b")];
|
||||
other_pod
|
||||
.increment(&first, Number::Float(1.25))
|
||||
.await
|
||||
.unwrap();
|
||||
raw.send(commands::expire(&first.key(), Duration::from_secs(5)))
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
let totals = store
|
||||
.increment_many(&[
|
||||
(first.clone(), Number::Float(0.5)),
|
||||
(second.clone(), Number::Int(2)),
|
||||
])
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(totals, [Number::Float(1.75), Number::Float(2.0)]);
|
||||
assert_eq!(trips.seen(), [Operation::Pipeline { ops: 2 }]);
|
||||
let Ttl::Expires(remaining) = raw.send(commands::ttl(&first.key())).await.unwrap() else {
|
||||
panic!("an incremented counter has no TTL");
|
||||
};
|
||||
assert!(within(remaining, 60), "{remaining:?}");
|
||||
assert_eq!(
|
||||
raw.send(commands::get::<String>(&first.key()))
|
||||
.await
|
||||
.unwrap(),
|
||||
Some("1.75".to_owned())
|
||||
);
|
||||
assert_eq!(store.get_local(&first).unwrap(), Some(Number::Float(1.75)));
|
||||
assert_eq!(store.get_local(&second).unwrap(), Some(Number::Float(2.0)));
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn a_post_call_increment_from_a_thread_without_a_runtime_lands_in_redis_and_memory(
|
||||
#[values(Topology::Standalone, Topology::Cluster)] topology: Topology,
|
||||
) {
|
||||
let deployment = RedisDeployment::start(topology).await.unwrap();
|
||||
let raw = connect(&deployment);
|
||||
let store = store(Some(connect(&deployment)));
|
||||
let request = RequestRedis::new(Handle::current(), LONG_DEADLINE);
|
||||
let counter = team("t");
|
||||
|
||||
let declared = std::thread::spawn({
|
||||
let (store, request, counter) = (Arc::clone(&store), request.clone(), counter.clone());
|
||||
move || store.increment_many_post_call(&request, &[(counter, Number::Float(3.5))])
|
||||
})
|
||||
.join()
|
||||
.expect("a Python thread has no Tokio context, and declaring there must not panic");
|
||||
request.flush_post().await;
|
||||
|
||||
assert!(declared);
|
||||
let mirrored =
|
||||
litellm_redis_testing::eventually(async || store.get_local(&counter).unwrap()).await;
|
||||
assert_eq!(mirrored, Number::Float(3.5));
|
||||
assert_eq!(
|
||||
raw.send(commands::get::<String>(&counter.key()))
|
||||
.await
|
||||
.unwrap(),
|
||||
Some("3.5".to_owned())
|
||||
);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn post_call_increments_wait_for_the_post_batch_then_land_in_memory(
|
||||
#[values(Topology::Standalone, Topology::Cluster)] topology: Topology,
|
||||
) {
|
||||
let deployment = RedisDeployment::start(topology).await.unwrap();
|
||||
let raw = connect(&deployment);
|
||||
let (client, trips) = connect_observed(&deployment);
|
||||
let store = store(Some(client));
|
||||
let request = RequestRedis::new(Handle::current(), LONG_DEADLINE);
|
||||
let counter = team("t");
|
||||
|
||||
assert!(store.increment_many_post_call(&request, &[(counter.clone(), Number::Float(3.5))]));
|
||||
request.flush_pre().await;
|
||||
|
||||
assert!(trips.seen().is_empty());
|
||||
assert_eq!(
|
||||
raw.send(commands::get::<String>(&counter.key()))
|
||||
.await
|
||||
.unwrap(),
|
||||
None
|
||||
);
|
||||
|
||||
request.flush_post().await;
|
||||
tokio::time::sleep(Duration::from_millis(100)).await;
|
||||
|
||||
let increment = match topology {
|
||||
Topology::Standalone => Operation::Pipeline { ops: 1 },
|
||||
Topology::Cluster => Operation::Command("INCRBYFLOAT"),
|
||||
};
|
||||
assert_eq!(trips.seen(), [increment]);
|
||||
assert_eq!(
|
||||
raw.send(commands::get::<String>(&counter.key()))
|
||||
.await
|
||||
.unwrap(),
|
||||
Some("3.5".to_owned())
|
||||
);
|
||||
assert_eq!(store.get_local(&counter).unwrap(), Some(Number::Float(3.5)));
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn the_budget_read_rides_the_same_pre_call_pipeline_as_the_cooldown_read(
|
||||
#[values(Topology::Standalone, Topology::Cluster)] topology: Topology,
|
||||
) {
|
||||
let deployment = RedisDeployment::start(topology).await.unwrap();
|
||||
let raw = connect(&deployment);
|
||||
let (client, trips) = connect_observed(&deployment);
|
||||
let cooldowns = std::sync::Arc::new(RouterCooldowns::new(
|
||||
Some(remote(client.clone())),
|
||||
COOLDOWN_POLICY,
|
||||
));
|
||||
let spend = store(Some(client));
|
||||
let request = RequestRedis::new(Handle::current(), LONG_DEADLINE);
|
||||
let counters = [
|
||||
SpendCounter::Key {
|
||||
token: "hashed".to_owned(),
|
||||
window: None,
|
||||
},
|
||||
team("t"),
|
||||
];
|
||||
raw.send(commands::set(&counters[0].key(), "4.5", TTL))
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
let active = cooldowns
|
||||
.active(RedisSend::PreCall(&request), &["dep-a".to_owned()])
|
||||
.unwrap();
|
||||
let budget = spend
|
||||
.read_many(RedisSend::PreCall(&request), &counters)
|
||||
.unwrap();
|
||||
request.flush_pre().await;
|
||||
|
||||
let expected = match topology {
|
||||
Topology::Standalone => vec![Operation::Pipeline { ops: 2 }],
|
||||
Topology::Cluster => vec![Operation::Command("MGET"); 2],
|
||||
};
|
||||
assert_eq!(trips.seen(), expected);
|
||||
assert_eq!(active.await.unwrap(), vec![]);
|
||||
assert_eq!(
|
||||
budget.await.unwrap(),
|
||||
Some(vec![Some(Number::Float(4.5)), None])
|
||||
);
|
||||
assert_eq!(trips.seen().len(), expected.len());
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn a_seed_lets_the_first_pod_win_and_others_answer_with_its_value(
|
||||
#[values(Topology::Standalone, Topology::Cluster)] topology: Topology,
|
||||
) {
|
||||
let deployment = RedisDeployment::start(topology).await.unwrap();
|
||||
let raw = connect(&deployment);
|
||||
let first = store(Some(connect(&deployment)));
|
||||
let second = store(Some(connect(&deployment)));
|
||||
let counter = SpendCounter::Org("org".to_owned());
|
||||
|
||||
let won = first
|
||||
.seed(&counter, Number::Float(10.0), LostSeed::Answer)
|
||||
.await
|
||||
.unwrap();
|
||||
let lost = second
|
||||
.seed(&counter, Number::Float(99.0), LostSeed::Increment)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
assert_eq!([won, lost], [Number::Float(10.0), Number::Float(10.0)]);
|
||||
assert_eq!(
|
||||
raw.send(commands::get::<String>(&counter.key()))
|
||||
.await
|
||||
.unwrap(),
|
||||
Some("10.0".to_owned())
|
||||
);
|
||||
let Ttl::Expires(remaining) = raw.send(commands::ttl(&counter.key())).await.unwrap() else {
|
||||
panic!("a seeded counter has no TTL");
|
||||
};
|
||||
assert!(within(remaining, 60), "{remaining:?}");
|
||||
assert_eq!(
|
||||
second.get_local(&counter).unwrap(),
|
||||
Some(Number::Float(10.0))
|
||||
);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn delete_drops_memory_and_redis_and_refresh_ttl_reports_whether_redis_held_it(
|
||||
#[values(Topology::Standalone, Topology::Cluster)] topology: Topology,
|
||||
) {
|
||||
let deployment = RedisDeployment::start(topology).await.unwrap();
|
||||
let raw = connect(&deployment);
|
||||
let store = store(Some(connect(&deployment)));
|
||||
let counters = [team("a"), team("b")];
|
||||
for counter in &counters {
|
||||
store.increment(counter, Number::Int(1)).await.unwrap();
|
||||
}
|
||||
raw.send(commands::expire(&counters[0].key(), Duration::from_secs(5)))
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
assert!(store.refresh_ttl(&counters[0]).await.unwrap());
|
||||
let Ttl::Expires(remaining) = raw.send(commands::ttl(&counters[0].key())).await.unwrap() else {
|
||||
panic!("a refreshed counter has no TTL");
|
||||
};
|
||||
assert!(within(remaining, 60), "{remaining:?}");
|
||||
|
||||
store.delete(&counters).await.unwrap();
|
||||
|
||||
for counter in &counters {
|
||||
assert_eq!(
|
||||
raw.send(commands::get::<String>(&counter.key()))
|
||||
.await
|
||||
.unwrap(),
|
||||
None
|
||||
);
|
||||
assert_eq!(store.get_local(counter).unwrap(), None);
|
||||
}
|
||||
assert!(!store.refresh_ttl(&counters[0]).await.unwrap());
|
||||
}
|
||||
|
||||
// main's read_spend_counter_cache_value: RedisCache.async_get_cache swallows a failed GET and
|
||||
// returns None, an authoritative miss. Only the breaker guard's refusal raises into the fallback
|
||||
// that reads this pod's memory
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn a_failed_read_serves_the_local_counter_and_a_failed_increment_drops_the_counter() {
|
||||
let deployment = RedisDeployment::start(Topology::Standalone).await.unwrap();
|
||||
let store = store(Some(connect(&deployment)));
|
||||
let [kept, failed] = [team("kept"), team("failed")];
|
||||
store.increment(&kept, Number::Float(2.0)).await.unwrap();
|
||||
store.increment(&failed, Number::Float(3.0)).await.unwrap();
|
||||
deployment.stop().await.unwrap();
|
||||
|
||||
let read = store.read(&kept).await.unwrap();
|
||||
let increment = store.increment(&failed, Number::Float(1.0)).await;
|
||||
|
||||
assert_eq!(read, Reading::Memory(Some(Number::Float(2.0))));
|
||||
assert!(increment.is_err());
|
||||
assert_eq!(store.get_local(&failed).unwrap(), None);
|
||||
assert_eq!(
|
||||
store
|
||||
.read_many(RedisSend::Now, &[kept])
|
||||
.unwrap()
|
||||
.await
|
||||
.unwrap(),
|
||||
None
|
||||
);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn a_read_with_an_open_breaker_serves_the_local_counter() {
|
||||
let deployment = RedisDeployment::start(Topology::Standalone).await.unwrap();
|
||||
let breaker = litellm_circuit_breaker::Breaker::new(
|
||||
litellm_circuit_breaker::Config {
|
||||
failure_threshold: 1,
|
||||
recovery_timeout: Duration::from_secs(60),
|
||||
..litellm_circuit_breaker::Config::default()
|
||||
},
|
||||
Box::new(|_| {}),
|
||||
);
|
||||
let store = store(Some(support::connect_guarded(&deployment, breaker)));
|
||||
let counter = team("t");
|
||||
store.increment(&counter, Number::Float(2.0)).await.unwrap();
|
||||
deployment.stop().await.unwrap();
|
||||
|
||||
let failed = store.read(&counter).await.unwrap();
|
||||
let refused = store.read(&counter).await.unwrap();
|
||||
|
||||
assert_eq!(failed, Reading::Memory(Some(Number::Float(2.0))));
|
||||
assert_eq!(refused, Reading::Memory(Some(Number::Float(2.0))));
|
||||
}
|
||||
|
||||
// increment_in_memory calls async_increment_cache(local_only=True), which hands InMemoryCache no
|
||||
// TTL, so the entry takes InMemoryCache's own default rather than spend_counter_cache's 60 s
|
||||
#[test]
|
||||
fn a_local_increment_outlives_the_default_in_memory_ttl() {
|
||||
let seconds = Arc::new(std::sync::atomic::AtomicU64::new(1_000));
|
||||
let now = {
|
||||
let seconds = seconds.clone();
|
||||
Arc::new(move || Duration::from_secs(seconds.load(std::sync::atomic::Ordering::SeqCst)))
|
||||
};
|
||||
let store = SpendCounters::with_clock(None, now);
|
||||
let counter = team("t");
|
||||
|
||||
store.increment_local(&counter, Number::Int(1)).unwrap();
|
||||
seconds.fetch_add(61, std::sync::atomic::Ordering::SeqCst);
|
||||
|
||||
assert_eq!(store.get_local(&counter).unwrap(), Some(Number::Int(1)));
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn a_failed_post_call_increment_drops_the_counter_from_memory() {
|
||||
let deployment = RedisDeployment::start(Topology::Standalone).await.unwrap();
|
||||
let store = store(Some(connect(&deployment)));
|
||||
let request = RequestRedis::new(Handle::current(), LONG_DEADLINE);
|
||||
let counter = team("t");
|
||||
store.increment(&counter, Number::Float(1.0)).await.unwrap();
|
||||
assert!(store.increment_many_post_call(&request, &[(counter.clone(), Number::Float(1.0))]));
|
||||
deployment.stop().await.unwrap();
|
||||
|
||||
request.flush_post().await;
|
||||
tokio::time::sleep(Duration::from_millis(200)).await;
|
||||
|
||||
assert_eq!(store.get_local(&counter).unwrap(), None);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn without_redis_counters_live_in_memory() {
|
||||
let store = store(None);
|
||||
let counter = team("t");
|
||||
let request = RequestRedis::new(Handle::current(), LONG_DEADLINE);
|
||||
|
||||
let totals = store
|
||||
.increment_many(&[
|
||||
(counter.clone(), Number::Int(2)),
|
||||
(counter.clone(), Number::Float(0.5)),
|
||||
])
|
||||
.await
|
||||
.unwrap();
|
||||
let seeded = store
|
||||
.seed(&counter, Number::Float(1.0), LostSeed::Answer)
|
||||
.await
|
||||
.unwrap();
|
||||
let raised = store.set_max(&counter, Number::Float(7.0)).await.unwrap();
|
||||
|
||||
assert_eq!(totals, [Number::Int(2), Number::Float(2.5)]);
|
||||
assert_eq!(seeded, Number::Float(2.5));
|
||||
assert_eq!(raised, None);
|
||||
assert_eq!(
|
||||
store.read(&counter).await.unwrap(),
|
||||
Reading::Memory(Some(Number::Float(7.0)))
|
||||
);
|
||||
assert!(!store.increment_many_post_call(&request, &[(counter.clone(), Number::Int(1))]));
|
||||
assert!(
|
||||
store
|
||||
.read_many(RedisSend::PreCall(&request), &[counter])
|
||||
.is_none()
|
||||
);
|
||||
}
|
||||
|
||||
// main's RedisCache: async_increment(refresh_ttl=True), MGET for a batched read, async_set_max as
|
||||
// a plain EVAL, SET EX NX for a seed and one multi-key DEL per chunk
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn spend_counters_send_the_commands_main_sends(
|
||||
#[values(Topology::Standalone, Topology::Cluster)] topology: Topology,
|
||||
) {
|
||||
let deployment = RedisDeployment::start(topology).await.unwrap();
|
||||
let store = store(Some(connect(&deployment)));
|
||||
let counters = [team("a"), team("b")];
|
||||
deployment
|
||||
.cli_on_every_node(&["CONFIG", "RESETSTAT"])
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
store
|
||||
.increment(&counters[0], Number::Float(1.5))
|
||||
.await
|
||||
.unwrap();
|
||||
store
|
||||
.read_many(RedisSend::Now, &counters)
|
||||
.unwrap()
|
||||
.await
|
||||
.unwrap()
|
||||
.unwrap();
|
||||
store
|
||||
.set_max(&counters[0], Number::Float(9.0))
|
||||
.await
|
||||
.unwrap();
|
||||
store
|
||||
.seed(&counters[1], Number::Float(2.0), LostSeed::Answer)
|
||||
.await
|
||||
.unwrap();
|
||||
store.delete(&counters).await.unwrap();
|
||||
|
||||
// commandstats also counts the GET, SET and EXPIRE the set_max script runs inside Redis
|
||||
let calls = async |command| support::calls(&deployment, command).await;
|
||||
assert_eq!(calls("incrbyfloat").await, 1);
|
||||
assert_eq!(calls("expire").await, 2);
|
||||
assert_eq!(calls("eval").await, 1);
|
||||
assert_eq!(calls("evalsha").await, 0);
|
||||
assert_eq!(calls("set").await, 2);
|
||||
assert!(calls("mget").await > 0);
|
||||
assert_eq!(calls("get").await, 1);
|
||||
assert!(calls("del").await > 0);
|
||||
}
|
||||
|
|
@ -10,9 +10,27 @@ use litellm_redis_client::{
|
|||
CallEvent, Cluster, Credentials, Namespace, Node, Operation, RedisClient, RedisConfig,
|
||||
RedisEvents, RedisUrl, Target, Timeouts, current_tag,
|
||||
};
|
||||
use litellm_redis_store::Remote;
|
||||
use litellm_redis_store::{
|
||||
Number, Remote, TierPolicy,
|
||||
router_cooldowns::Cooldown,
|
||||
router_usage::{Counter, Kind},
|
||||
};
|
||||
use litellm_redis_testing::RedisDeployment;
|
||||
|
||||
/// `DEFAULT_COOLDOWN_REDIS_READ_INTERVAL_SECONDS` and `DEFAULT_MAX_REDIS_BATCH_CACHE_SIZE`.
|
||||
pub const COOLDOWN_POLICY: TierPolicy = TierPolicy {
|
||||
refresh_after: Duration::from_secs(1),
|
||||
tracked_reads: 1000,
|
||||
default_in_memory_ttl: None,
|
||||
};
|
||||
|
||||
/// `DualCache`'s default `redis_batch_cache_expiry` and `DEFAULT_MAX_REDIS_BATCH_CACHE_SIZE`.
|
||||
pub const USAGE_POLICY: TierPolicy = TierPolicy {
|
||||
refresh_after: Duration::from_secs(10),
|
||||
tracked_reads: 1000,
|
||||
default_in_memory_ttl: None,
|
||||
};
|
||||
|
||||
pub fn connect(deployment: &RedisDeployment) -> RedisClient {
|
||||
connect_with(deployment, None, None)
|
||||
}
|
||||
|
|
@ -116,6 +134,103 @@ fn connect_with(
|
|||
.unwrap()
|
||||
}
|
||||
|
||||
/// What main's `RedisCache` writes for a cooldown: `str(value)` with `EX int(ttl)`, read back with
|
||||
/// `json.loads` then `ast.literal_eval` (`litellm/caching/redis_cache.py` `set_cache` and
|
||||
/// `_get_cache_logic`). Keys come from main's `CooldownCache`. The same literals sit in
|
||||
/// `tests/unit/redis/stores/test_router_cooldowns.py`, checked against main's codec.
|
||||
pub struct CooldownCase {
|
||||
pub model_id: &'static str,
|
||||
pub cooldown: Cooldown,
|
||||
pub ttl_seconds: f64,
|
||||
pub key: &'static str,
|
||||
pub redis_ttl_seconds: u64,
|
||||
pub str_dict: &'static str,
|
||||
pub json: &'static str,
|
||||
}
|
||||
|
||||
pub fn cooldowns() -> [CooldownCase; 2] {
|
||||
[
|
||||
CooldownCase {
|
||||
model_id: "dep-a",
|
||||
cooldown: Cooldown {
|
||||
exception_received: "RateLimitError: limited é 'quoted'".to_owned(),
|
||||
status_code: "429".to_owned(),
|
||||
timestamp: Number::Float(1_730_000_000.25),
|
||||
cooldown_time: Number::Int(60),
|
||||
},
|
||||
ttl_seconds: 90.5,
|
||||
key: "deployment:dep-a:cooldown",
|
||||
redis_ttl_seconds: 90,
|
||||
str_dict: "{'exception_received': \"RateLimitError: limited é 'quoted'\", 'status_code': '429', 'timestamp': 1730000000.25, 'cooldown_time': 60}",
|
||||
json: "{\"exception_received\": \"RateLimitError: limited \\u00e9 'quoted'\", \"status_code\": \"429\", \"timestamp\": 1730000000.25, \"cooldown_time\": 60}",
|
||||
},
|
||||
CooldownCase {
|
||||
model_id: "dep-b",
|
||||
cooldown: Cooldown {
|
||||
exception_received: "Timeout".to_owned(),
|
||||
status_code: "408".to_owned(),
|
||||
timestamp: Number::Float(1_730_000_000.5),
|
||||
cooldown_time: Number::Float(0.5),
|
||||
},
|
||||
ttl_seconds: 30.0,
|
||||
key: "deployment:dep-b:cooldown",
|
||||
redis_ttl_seconds: 30,
|
||||
str_dict: "{'exception_received': 'Timeout', 'status_code': '408', 'timestamp': 1730000000.5, 'cooldown_time': 0.5}",
|
||||
json: "{\"exception_received\": \"Timeout\", \"status_code\": \"408\", \"timestamp\": 1730000000.5, \"cooldown_time\": 0.5}",
|
||||
},
|
||||
]
|
||||
}
|
||||
|
||||
/// main's allowed-fails keys for `dep-a`: an empty suffix is the same key as none.
|
||||
pub struct AllowedFailsCase {
|
||||
pub suffix: Option<&'static str>,
|
||||
pub key: &'static str,
|
||||
}
|
||||
|
||||
pub const ALLOWED_FAILS: [AllowedFailsCase; 3] = [
|
||||
AllowedFailsCase {
|
||||
suffix: None,
|
||||
key: "deployment:dep-a:allowed_fails",
|
||||
},
|
||||
AllowedFailsCase {
|
||||
suffix: Some(""),
|
||||
key: "deployment:dep-a:allowed_fails",
|
||||
},
|
||||
AllowedFailsCase {
|
||||
suffix: Some("timeout"),
|
||||
key: "deployment:dep-a:allowed_fails:timeout",
|
||||
},
|
||||
];
|
||||
|
||||
/// main's `RouterCacheEnum` usage keys, including the `None` a missing model prints into the key.
|
||||
pub struct UsageCase {
|
||||
pub counter: Counter,
|
||||
pub key: &'static str,
|
||||
}
|
||||
|
||||
pub fn usage() -> [UsageCase; 2] {
|
||||
[
|
||||
UsageCase {
|
||||
counter: Counter {
|
||||
deployment_id: "dep".to_owned(),
|
||||
model: Some("anthropic/claude-sonnet-5-5".to_owned()),
|
||||
minute: "12-34".to_owned(),
|
||||
kind: Kind::Tpm,
|
||||
},
|
||||
key: "global_router:dep:anthropic/claude-sonnet-5-5:tpm:12-34",
|
||||
},
|
||||
UsageCase {
|
||||
counter: Counter {
|
||||
deployment_id: "dep".to_owned(),
|
||||
model: None,
|
||||
minute: "09-05".to_owned(),
|
||||
kind: Kind::Rpm,
|
||||
},
|
||||
key: "global_router:dep:None:rpm:09-05",
|
||||
},
|
||||
]
|
||||
}
|
||||
|
||||
/// A TTL Redis reports for a key written `whole_seconds` ago at most a few seconds back.
|
||||
pub fn within(ttl: Duration, whole_seconds: u64) -> bool {
|
||||
ttl <= Duration::from_secs(whole_seconds)
|
||||
|
|
|
|||
|
|
@ -284,6 +284,7 @@ impl Error {
|
|||
strum::EnumString,
|
||||
strum::IntoStaticStr,
|
||||
strum::VariantArray,
|
||||
serde_with::SerializeDisplay,
|
||||
)]
|
||||
pub enum FailureKind {
|
||||
#[strum(serialize = "RateLimitError")]
|
||||
|
|
@ -471,6 +472,17 @@ mod tests {
|
|||
assert_eq!(<&'static str>::from(kind), wire);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::rate_limit(FailureKind::RateLimit, "RateLimitError")]
|
||||
#[case::timeout(FailureKind::Timeout, "Timeout")]
|
||||
#[case::api_connection(FailureKind::ApiConnection, "APIConnectionError")]
|
||||
fn failure_kind_serializes_as_its_exception_name(
|
||||
#[case] kind: FailureKind,
|
||||
#[case] wire: &str,
|
||||
) {
|
||||
assert_eq!(serde_json::to_value(kind).unwrap(), serde_json::json!(wire));
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::empty("")]
|
||||
#[case::lowercase("ratelimiterror")]
|
||||
|
|
|
|||
|
|
@ -342,7 +342,10 @@ pub struct Output {
|
|||
}
|
||||
|
||||
/// How a provider blocked a response the call otherwise served.
|
||||
#[derive(Clone, Copy, Debug, PartialEq, Eq, strum::IntoStaticStr, strum::EnumString)]
|
||||
#[derive(
|
||||
Clone, Copy, Debug, PartialEq, Eq, serde::Serialize, strum::IntoStaticStr, strum::EnumString,
|
||||
)]
|
||||
#[serde(rename_all = "snake_case")]
|
||||
#[strum(serialize_all = "snake_case")]
|
||||
pub enum Blocked {
|
||||
/// A chat response whose first choice was content-filtered (or that has no choices).
|
||||
|
|
@ -402,8 +405,16 @@ pub enum SelectStage {
|
|||
OverrideCheck(RoutingStrategy),
|
||||
}
|
||||
|
||||
/// Routing events, delivered through [`Host::observe`].
|
||||
#[derive(Clone, Debug, PartialEq)]
|
||||
#[derive(Debug)]
|
||||
pub struct CallbackFiltered {
|
||||
pub kept: Vec<Arc<Deployment>>,
|
||||
pub preferred: Option<String>,
|
||||
}
|
||||
|
||||
/// Routing events, delivered through [`Host::observe`]. They serialize as JSON objects tagged by
|
||||
/// `event` (the variant name), durations in seconds.
|
||||
#[derive(Clone, Debug, PartialEq, serde::Serialize)]
|
||||
#[serde(tag = "event")]
|
||||
pub enum RouterEvent {
|
||||
/// A pick for `model` (the level's group) begins: its route phase opens.
|
||||
PickStarted {
|
||||
|
|
@ -413,6 +424,7 @@ pub enum RouterEvent {
|
|||
model_group: String,
|
||||
deployment_id: String,
|
||||
/// How long the pick took (`_acompletion`'s `async_get_available_deployment` service event).
|
||||
#[serde(serialize_with = "seconds")]
|
||||
pick: Duration,
|
||||
},
|
||||
Succeeded {
|
||||
|
|
@ -425,11 +437,13 @@ pub enum RouterEvent {
|
|||
},
|
||||
CooldownSet {
|
||||
deployment_id: String,
|
||||
#[serde(serialize_with = "seconds")]
|
||||
cooldown: Duration,
|
||||
},
|
||||
Retry {
|
||||
model_group: String,
|
||||
retry: u32,
|
||||
#[serde(serialize_with = "seconds")]
|
||||
delay: Duration,
|
||||
},
|
||||
Fallback {
|
||||
|
|
@ -584,6 +598,13 @@ impl PickError {
|
|||
}
|
||||
}
|
||||
|
||||
/// A pick's error serializes as its message.
|
||||
impl serde::Serialize for PickError {
|
||||
fn serialize<S: serde::Serializer>(&self, serializer: S) -> Result<S::Ok, S::Error> {
|
||||
serializer.collect_str(&self.0)
|
||||
}
|
||||
}
|
||||
|
||||
impl PartialEq for PickError {
|
||||
fn eq(&self, other: &Self) -> bool {
|
||||
format!("{:?}", self.0) == format!("{:?}", other.0)
|
||||
|
|
@ -591,7 +612,8 @@ impl PartialEq for PickError {
|
|||
}
|
||||
|
||||
/// Which error a [`RouterEvent::RetryLogged`] records.
|
||||
#[derive(Clone, Debug, PartialEq)]
|
||||
#[derive(Clone, Debug, PartialEq, serde::Serialize)]
|
||||
#[serde(tag = "cause", content = "error", rename_all = "snake_case")]
|
||||
pub enum RetryCause {
|
||||
/// The last attempt's failure.
|
||||
Attempt,
|
||||
|
|
@ -698,14 +720,19 @@ pub trait Host: Send + Sync {
|
|||
}
|
||||
|
||||
/// `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.
|
||||
/// turn, and may name one deployment they prefer, which the pick takes only when the filters
|
||||
/// after the callbacks keep it. 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 and prefers none.
|
||||
async fn callback_filters(
|
||||
&self,
|
||||
_ctx: &RoutingContext<'_>,
|
||||
candidates: Vec<Arc<Deployment>>,
|
||||
) -> Result<Vec<Arc<Deployment>>, Error> {
|
||||
Ok(candidates)
|
||||
) -> Result<CallbackFiltered, Error> {
|
||||
Ok(CallbackFiltered {
|
||||
kept: candidates,
|
||||
preferred: None,
|
||||
})
|
||||
}
|
||||
/// 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
|
||||
|
|
@ -745,3 +772,83 @@ pub trait Host: Send + Sync {
|
|||
Ok(Value::Null)
|
||||
}
|
||||
}
|
||||
|
||||
fn seconds<S: serde::Serializer>(duration: &Duration, serializer: S) -> Result<S::Ok, S::Error> {
|
||||
serializer.serialize_f64(duration.as_secs_f64())
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use rstest::rstest;
|
||||
use serde_json::{Value, json};
|
||||
|
||||
use super::*;
|
||||
|
||||
#[rstest]
|
||||
#[case::selected(
|
||||
RouterEvent::Selected {
|
||||
model_group: "gpt-4o".to_owned(),
|
||||
deployment_id: "azure-eastus".to_owned(),
|
||||
pick: Duration::from_millis(1500),
|
||||
},
|
||||
json!({"event": "Selected", "model_group": "gpt-4o", "deployment_id": "azure-eastus", "pick": 1.5})
|
||||
)]
|
||||
#[case::attempt_failed(
|
||||
RouterEvent::AttemptFailed {
|
||||
deployment_id: "openai".to_owned(),
|
||||
kind: FailureKind::RateLimit,
|
||||
status: Some(429),
|
||||
},
|
||||
json!({"event": "AttemptFailed", "deployment_id": "openai", "kind": "RateLimitError", "status": 429})
|
||||
)]
|
||||
#[case::cooldown_set(
|
||||
RouterEvent::CooldownSet { deployment_id: "openai".to_owned(), cooldown: Duration::from_secs(30) },
|
||||
json!({"event": "CooldownSet", "deployment_id": "openai", "cooldown": 30.0})
|
||||
)]
|
||||
#[case::retry(
|
||||
RouterEvent::Retry { model_group: "gpt-4o".to_owned(), retry: 1, delay: Duration::ZERO },
|
||||
json!({"event": "Retry", "model_group": "gpt-4o", "retry": 1, "delay": 0.0})
|
||||
)]
|
||||
#[case::blocked(
|
||||
RouterEvent::Blocked { level: 0, model: "claude".to_owned(), blocked: Blocked::ContentFilter },
|
||||
json!({"event": "Blocked", "level": 0, "model": "claude", "blocked": "content_filter"})
|
||||
)]
|
||||
#[case::retry_logged_after_an_attempt(
|
||||
RouterEvent::RetryLogged {
|
||||
scope: 0,
|
||||
level: 0,
|
||||
model: "gpt-4o".to_owned(),
|
||||
attempted_retries: None,
|
||||
cause: RetryCause::Attempt,
|
||||
},
|
||||
json!({
|
||||
"event": "RetryLogged", "scope": 0, "level": 0, "model": "gpt-4o",
|
||||
"attempted_retries": null, "cause": {"cause": "attempt"},
|
||||
})
|
||||
)]
|
||||
#[case::retry_logged_after_a_pick(
|
||||
RouterEvent::RetryLogged {
|
||||
scope: 1,
|
||||
level: 1,
|
||||
model: "backup".to_owned(),
|
||||
attempted_retries: Some(2),
|
||||
cause: RetryCause::Pick(PickError(Arc::new(Error::UnknownModel { model: "backup".to_owned() }))),
|
||||
},
|
||||
json!({
|
||||
"event": "RetryLogged", "scope": 1, "level": 1, "model": "backup", "attempted_retries": 2,
|
||||
"cause": {"cause": "pick", "error": "no deployment serves model `backup`"},
|
||||
})
|
||||
)]
|
||||
#[case::mock_failed(
|
||||
RouterEvent::MockFailed {
|
||||
mock: MockFailure::RateLimit,
|
||||
level: 0,
|
||||
model: "gpt-4o".to_owned(),
|
||||
num_retries: None,
|
||||
},
|
||||
json!({"event": "MockFailed", "mock": "rate_limit", "level": 0, "model": "gpt-4o", "num_retries": null})
|
||||
)]
|
||||
fn an_event_serializes_tagged_by_its_name(#[case] event: RouterEvent, #[case] expected: Value) {
|
||||
assert_eq!(serde_json::to_value(&event).unwrap(), expected);
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -124,8 +124,8 @@ pub use crate::{
|
|||
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,
|
||||
Attempt, Blocked, BuildAsk, BuildPath, Built, CallbackFiltered, Host, Output, Provided,
|
||||
RetryCause, RouterEvent, RouterRequest, RoutingContext, SecretAt, SecretRef, SelectStage,
|
||||
},
|
||||
outcome::{OutcomeCallbacks, OutcomeWrite},
|
||||
prerouting::{PreRoute, PreRouted},
|
||||
|
|
|
|||
|
|
@ -7,7 +7,15 @@ use crate::error::{AttemptError, FailureKind};
|
|||
|
||||
/// Which mock failure a request asks for.
|
||||
#[derive(
|
||||
Clone, Copy, Debug, PartialEq, Eq, Deserialize, strum::IntoStaticStr, strum::EnumString,
|
||||
Clone,
|
||||
Copy,
|
||||
Debug,
|
||||
PartialEq,
|
||||
Eq,
|
||||
Deserialize,
|
||||
serde::Serialize,
|
||||
strum::IntoStaticStr,
|
||||
strum::EnumString,
|
||||
)]
|
||||
#[serde(rename_all = "snake_case")]
|
||||
#[strum(serialize_all = "snake_case")]
|
||||
|
|
|
|||
|
|
@ -456,7 +456,7 @@ impl Core {
|
|||
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
|
||||
let preferred = if settings.callback_filters
|
||||
&& !inputs.count_only
|
||||
&& !resolution.kind.is_single()
|
||||
&& request.pass_through != Some(PassThrough::Listed)
|
||||
|
|
@ -472,8 +472,12 @@ impl Core {
|
|||
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?;
|
||||
}
|
||||
let filtered = self.host.callback_filters(&ctx, healthy).await?;
|
||||
healthy = filtered.kept;
|
||||
filtered.preferred
|
||||
} else {
|
||||
None
|
||||
};
|
||||
if inputs.count_only {
|
||||
return Ok(Candidates {
|
||||
name: name.to_owned(),
|
||||
|
|
@ -648,8 +652,13 @@ impl Core {
|
|||
healthy.retain(|d| !d.cannot_run_batches());
|
||||
}
|
||||
|
||||
if let Some(pinned) = request.pinned_deployment_id.as_deref()
|
||||
&& healthy.iter().any(|d| d.id == pinned)
|
||||
if let Some(pinned) = [
|
||||
request.pinned_deployment_id.as_deref(),
|
||||
preferred.as_deref(),
|
||||
]
|
||||
.into_iter()
|
||||
.flatten()
|
||||
.find(|id| healthy.iter().any(|d| d.id == *id))
|
||||
{
|
||||
healthy.retain(|d| d.id == pinned);
|
||||
}
|
||||
|
|
|
|||
|
|
@ -17,7 +17,7 @@ use litellm_router_core::{
|
|||
};
|
||||
use parking_lot::Mutex;
|
||||
use rstest::{fixture, rstest};
|
||||
use serde_json::json;
|
||||
use serde_json::{Value, json};
|
||||
|
||||
#[derive(Default)]
|
||||
struct ScriptedHost {
|
||||
|
|
@ -1809,3 +1809,55 @@ fn request_for(model: &str) -> RouterRequest {
|
|||
..RouterRequest::default()
|
||||
}
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn every_event_of_a_retried_and_fallen_back_call_serializes_tagged_by_its_name(
|
||||
config: RouterConfig,
|
||||
) {
|
||||
let host = Arc::new(ScriptedHost::failing(
|
||||
"primary",
|
||||
vec![failure(FailureKind::RateLimit, 429); 3],
|
||||
));
|
||||
let router = router(
|
||||
config,
|
||||
json!({"num_retries": 1, "fallbacks": [{"primary": ["backup"]}]}),
|
||||
host.clone(),
|
||||
);
|
||||
router
|
||||
.call(Op::from("acompletion"), request())
|
||||
.await
|
||||
.unwrap();
|
||||
let tags: Value = host
|
||||
.events
|
||||
.lock()
|
||||
.iter()
|
||||
.map(|event| serde_json::to_value(event).unwrap()["event"].clone())
|
||||
.collect();
|
||||
assert_eq!(
|
||||
tags,
|
||||
json!([
|
||||
"RetriesOpened",
|
||||
"PickStarted",
|
||||
"Selected",
|
||||
"AttemptFailed",
|
||||
"CooldownSet",
|
||||
"RetryBudget",
|
||||
"RetryLogged",
|
||||
"Retry",
|
||||
"PickStarted",
|
||||
"Selected",
|
||||
"AttemptFailed",
|
||||
"CooldownSet",
|
||||
"RetryLogged",
|
||||
"FallingBack",
|
||||
"RetryLogged",
|
||||
"Fallback",
|
||||
"RetriesOpened",
|
||||
"PickStarted",
|
||||
"Selected",
|
||||
"Succeeded",
|
||||
"FallbackDone",
|
||||
])
|
||||
);
|
||||
}
|
||||
|
|
|
|||
|
|
@ -105,8 +105,15 @@ def current_service_target() -> str | None:
|
|||
return _service_target.get()
|
||||
|
||||
|
||||
def with_service_target(target: str) -> Callable[[Callable[_P, _R]], Callable[_P, _R]]:
|
||||
"""Run every call of the decorated function, coroutine functions included, under ``service_target(target)``."""
|
||||
def with_service_target(
|
||||
target: str, *, preserve_existing: bool = False
|
||||
) -> Callable[[Callable[_P, _R]], Callable[_P, _R]]:
|
||||
"""Run calls under the target, or the inherited target when ``preserve_existing`` is true."""
|
||||
|
||||
def call_target() -> str:
|
||||
if preserve_existing:
|
||||
return current_service_target() or target
|
||||
return target
|
||||
|
||||
def decorate(fn: Callable[_P, _R]) -> Callable[_P, _R]:
|
||||
if inspect.iscoroutinefunction(fn):
|
||||
|
|
@ -116,14 +123,14 @@ def with_service_target(target: str) -> Callable[[Callable[_P, _R]], Callable[_P
|
|||
|
||||
@wraps(fn)
|
||||
async def run_async(*args: _P.args, **kwargs: _P.kwargs) -> object:
|
||||
with service_target(target):
|
||||
with service_target(call_target()):
|
||||
return await awaitable_fn(*args, **kwargs)
|
||||
|
||||
return cast("Callable[_P, _R]", run_async) # cast-ok: same coroutine-returning signature as ``fn``
|
||||
|
||||
@wraps(fn)
|
||||
def run(*args: _P.args, **kwargs: _P.kwargs) -> _R:
|
||||
with service_target(target):
|
||||
with service_target(call_target()):
|
||||
return fn(*args, **kwargs)
|
||||
|
||||
return run
|
||||
|
|
|
|||
|
|
@ -518,7 +518,7 @@ def _or_secrets(value: object, read: Callable[[str], object], *names: str) -> ob
|
|||
return reduce(lambda found, name: found or read(name), names, value)
|
||||
|
||||
|
||||
def _merge_redis_kwargs(
|
||||
def merge_redis_kwargs(
|
||||
env_overrides: Mapping[str, object],
|
||||
) -> dict[str, object]: # mutable-ok: _get_redis_client_logic swaps settings for client objects in place
|
||||
overrides: Final = {
|
||||
|
|
@ -580,7 +580,7 @@ def _get_redis_client_logic(**env_overrides):
|
|||
"""
|
||||
Common functionality across sync + async redis client implementations
|
||||
"""
|
||||
redis_kwargs: Final = _merge_redis_kwargs(env_overrides)
|
||||
redis_kwargs: Final = merge_redis_kwargs(env_overrides)
|
||||
|
||||
_startup_nodes: Final = redis_kwargs.get("startup_nodes") or None
|
||||
|
||||
|
|
|
|||
|
|
@ -4,13 +4,14 @@ Custom A2A Card Resolver for LiteLLM.
|
|||
Extends the A2A SDK's card resolver to support multiple well-known paths.
|
||||
"""
|
||||
|
||||
from collections.abc import Mapping
|
||||
from collections.abc import Mapping, Sequence
|
||||
from types import MappingProxyType
|
||||
from typing import TYPE_CHECKING, Any, Final, Protocol, runtime_checkable
|
||||
|
||||
from litellm._logging import verbose_logger
|
||||
from litellm.a2a_protocol.exceptions import A2AAgentCardDiscoveryError
|
||||
from litellm.constants import LOCALHOST_URL_PATTERNS
|
||||
from litellm.litellm_core_utils.core_helpers import has_attributes
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from a2a.types import AgentCard
|
||||
|
|
@ -130,15 +131,34 @@ def normalize_agent_card_interfaces(agent_card: "AgentCard", protocol_version: s
|
|||
return normalized
|
||||
|
||||
|
||||
class _UrlCarrier(Protocol):
|
||||
url: str
|
||||
|
||||
|
||||
class _InterfacesCarrier(Protocol):
|
||||
@property
|
||||
def supported_interfaces(self) -> Sequence[_UrlCarrier]: ...
|
||||
|
||||
|
||||
def _url_of(value: object) -> str | None:
|
||||
return value.url if has_attributes(value, _UrlCarrier, "url") else None
|
||||
|
||||
|
||||
def _supported_interfaces(agent_card: object) -> Sequence[_UrlCarrier] | None:
|
||||
if has_attributes(agent_card, _InterfacesCarrier, "supported_interfaces"):
|
||||
return agent_card.supported_interfaces
|
||||
return None
|
||||
|
||||
|
||||
def get_agent_card_url(agent_card: "AgentCard") -> str | None:
|
||||
"""Return the agent endpoint URL from the resolved SDK card."""
|
||||
url: Final = getattr(agent_card, "url", None)
|
||||
url: Final = _url_of(agent_card)
|
||||
if url:
|
||||
return url
|
||||
|
||||
interfaces: Final = getattr(agent_card, "supported_interfaces", None)
|
||||
interfaces: Final = _supported_interfaces(agent_card)
|
||||
if interfaces:
|
||||
return getattr(interfaces[0], "url", None)
|
||||
return _url_of(interfaces[0])
|
||||
return None
|
||||
|
||||
|
||||
|
|
@ -168,16 +188,16 @@ def fix_agent_card_url(agent_card: "AgentCard", base_url: str) -> "AgentCard":
|
|||
Returns:
|
||||
The agent card with the URL fixed if necessary
|
||||
"""
|
||||
card_url: Final = getattr(agent_card, "url", None)
|
||||
card_url: Final = _url_of(agent_card)
|
||||
|
||||
if card_url and is_localhost_or_internal_url(card_url):
|
||||
# Normalize base_url to ensure it ends with /
|
||||
fixed_url: Final = base_url.rstrip("/") + "/"
|
||||
agent_card.url = fixed_url
|
||||
|
||||
interfaces: Final = getattr(agent_card, "supported_interfaces", None)
|
||||
interfaces: Final = _supported_interfaces(agent_card)
|
||||
if interfaces:
|
||||
interface_url: Final = getattr(interfaces[0], "url", None)
|
||||
interface_url: Final = _url_of(interfaces[0])
|
||||
if interface_url and is_localhost_or_internal_url(interface_url):
|
||||
interfaces[0].url = base_url.rstrip("/") + "/"
|
||||
|
||||
|
|
@ -194,6 +214,9 @@ class LiteLLMA2ACardResolver(_A2ACardResolver):
|
|||
- /agentCard/v1.0
|
||||
"""
|
||||
|
||||
if TYPE_CHECKING:
|
||||
base_url: str
|
||||
|
||||
async def get_agent_card(
|
||||
self,
|
||||
relative_card_path: str | None = None,
|
||||
|
|
|
|||
|
|
@ -4,7 +4,7 @@ A2A Protocol Exception Mapping Utils.
|
|||
Maps A2A SDK exceptions to LiteLLM A2A exception types.
|
||||
"""
|
||||
|
||||
from typing import TYPE_CHECKING, Any, Final
|
||||
from typing import TYPE_CHECKING, Final
|
||||
|
||||
from litellm._logging import verbose_logger
|
||||
from litellm.a2a_protocol.card_resolver import (
|
||||
|
|
@ -21,6 +21,7 @@ from litellm.constants import CONNECTION_ERROR_PATTERNS
|
|||
|
||||
if TYPE_CHECKING:
|
||||
from a2a.client import Client as A2AClientType
|
||||
from a2a.types import AgentCard
|
||||
|
||||
|
||||
_A2A_KIND_DIALECT_HINT: Final = (
|
||||
|
|
@ -165,7 +166,7 @@ def map_a2a_exception(
|
|||
|
||||
async def handle_a2a_localhost_retry(
|
||||
error: A2ALocalhostURLError,
|
||||
agent_card: Any,
|
||||
agent_card: "AgentCard | None",
|
||||
a2a_client: "A2AClientType",
|
||||
is_streaming: bool = False,
|
||||
) -> "A2AClientType":
|
||||
|
|
|
|||
|
|
@ -10,7 +10,7 @@ A2A Streaming Events (in order):
|
|||
4. Status update (kind: "status-update") - Final status "completed" with final=true
|
||||
"""
|
||||
|
||||
from collections.abc import AsyncIterator, Callable, Coroutine, Mapping
|
||||
from collections.abc import AsyncIterator, Mapping
|
||||
from types import MappingProxyType
|
||||
from typing import Any, Final
|
||||
|
||||
|
|
@ -126,10 +126,7 @@ class A2ACompletionBridgeHandler:
|
|||
|
||||
@staticmethod
|
||||
async def _acompletion(completion_params: Mapping[str, object]) -> ModelResponse | CustomStreamWrapper:
|
||||
acompletion_fn: Final[Callable[..., Coroutine[object, object, ModelResponse | CustomStreamWrapper]]] = vars(
|
||||
litellm
|
||||
)["acompletion"]
|
||||
return await acompletion_fn(**completion_params)
|
||||
return await litellm.acompletion(**completion_params)
|
||||
|
||||
@staticmethod
|
||||
async def handle_non_streaming(
|
||||
|
|
|
|||
|
|
@ -14,7 +14,7 @@ import datetime
|
|||
import uuid
|
||||
from collections.abc import AsyncIterator, Coroutine, Mapping
|
||||
from types import MappingProxyType, ModuleType
|
||||
from typing import TYPE_CHECKING, Final, Optional, cast
|
||||
from typing import TYPE_CHECKING, Final, Optional, Protocol, cast
|
||||
|
||||
from typing_extensions import TypedDict, Unpack
|
||||
|
||||
|
|
@ -24,6 +24,7 @@ from litellm.a2a_protocol.streaming_iterator import A2AStreamingIterator
|
|||
from litellm.a2a_protocol.utils import A2ARequestUtils
|
||||
from litellm.constants import DEFAULT_A2A_AGENT_TIMEOUT
|
||||
from litellm.litellm_core_utils.asyncify import asyncify
|
||||
from litellm.litellm_core_utils.core_helpers import require_attributes
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging
|
||||
from litellm.llms.custom_httpx.http_handler import (
|
||||
get_async_httpx_client,
|
||||
|
|
@ -36,7 +37,6 @@ if TYPE_CHECKING:
|
|||
from a2a.client import Client as A2AClientType
|
||||
from a2a.client import ClientCallContext as A2ACallContextType
|
||||
from a2a.compat.v0_3.types import (
|
||||
AgentCard,
|
||||
Message,
|
||||
SendMessageRequest,
|
||||
SendMessageResponse,
|
||||
|
|
@ -45,6 +45,7 @@ if TYPE_CHECKING:
|
|||
SendStreamingMessageSuccessResponse,
|
||||
Task,
|
||||
)
|
||||
from a2a.types import AgentCard
|
||||
from a2a.types.a2a_pb2 import SendMessageRequest as CoreSendMessageRequest
|
||||
from a2a.types.a2a_pb2 import StreamResponse as CoreStreamResponse
|
||||
|
||||
|
|
@ -95,6 +96,12 @@ class _SendMessageOptions(TypedDict, total=False, extra_items=object):
|
|||
pass
|
||||
|
||||
|
||||
class _StreamingRequestConverter(Protocol):
|
||||
def to_core_send_message_request(
|
||||
self, compat_req: "SendStreamingMessageRequest", /
|
||||
) -> "CoreSendMessageRequest": ...
|
||||
|
||||
|
||||
def _set_usage_on_logging_obj(
|
||||
kwargs: Mapping[str, object],
|
||||
prompt_tokens: int,
|
||||
|
|
@ -375,7 +382,9 @@ async def _stream_messages(
|
|||
"The 'a2a' package is required for A2A agent invocation. Install it with: pip install a2a-sdk"
|
||||
)
|
||||
|
||||
pb_request: Final[CoreSendMessageRequest] = _a2a_conversions.to_core_send_message_request(request)
|
||||
pb_request: Final[CoreSendMessageRequest] = require_attributes(
|
||||
_a2a_conversions, _StreamingRequestConverter, "to_core_send_message_request"
|
||||
).to_core_send_message_request(request)
|
||||
async for event in a2a_client.send_message(pb_request, context=_get_a2a_call_context(a2a_client)):
|
||||
compat_chunk = _to_compat_stream_response(event, request_id=request.id)
|
||||
yield SendStreamingMessageResponse(root=compat_chunk)
|
||||
|
|
|
|||
|
|
@ -427,6 +427,7 @@ def update_headers_with_filtered_beta(
|
|||
provider: str,
|
||||
api_base: str | None = None,
|
||||
unlisted: UnlistedBetas = "drop",
|
||||
keep: frozenset[str] = frozenset(),
|
||||
) -> dict:
|
||||
"""
|
||||
Update headers dict by filtering and transforming anthropic-beta header values.
|
||||
|
|
@ -437,6 +438,7 @@ def update_headers_with_filtered_beta(
|
|||
provider: Provider name
|
||||
api_base: Deployment api_base, used to pick the allowlist row (Foundry hosts use ``azure_ai``)
|
||||
unlisted: Whether a beta missing from the row is dropped or forwarded as sent
|
||||
keep: Betas the routed model accepts, forwarded as sent whatever the row says
|
||||
|
||||
Returns:
|
||||
Updated headers dict
|
||||
|
|
@ -446,13 +448,20 @@ def update_headers_with_filtered_beta(
|
|||
return headers
|
||||
|
||||
# Parse existing beta headers
|
||||
beta_values: Final = [b.strip() for b in existing_beta.split(",") if b.strip()]
|
||||
beta_values: Final[tuple[str, ...]] = tuple(b.strip() for b in existing_beta.split(",") if b.strip())
|
||||
|
||||
# Filter and transform based on provider
|
||||
filtered_beta_values: Final = filter_and_transform_beta_headers(
|
||||
beta_headers=beta_values,
|
||||
provider=resolve_beta_provider(provider, api_base),
|
||||
unlisted=unlisted,
|
||||
filtered_beta_values: Final = tuple(
|
||||
dict.fromkeys(
|
||||
[
|
||||
*filter_and_transform_beta_headers(
|
||||
beta_headers=[b for b in beta_values if b not in keep],
|
||||
provider=resolve_beta_provider(provider, api_base),
|
||||
unlisted=unlisted,
|
||||
),
|
||||
*(b for b in beta_values if b in keep),
|
||||
]
|
||||
)
|
||||
)
|
||||
|
||||
# Update or remove the header
|
||||
|
|
|
|||
|
|
@ -13,10 +13,10 @@ import json
|
|||
import logging
|
||||
import time
|
||||
import traceback
|
||||
from collections.abc import Generator, Mapping
|
||||
from collections.abc import Awaitable, Callable, Generator, Mapping
|
||||
from contextlib import contextmanager
|
||||
from enum import Enum
|
||||
from typing import Any, Final, Literal, cast
|
||||
from typing import Any, Final, Literal, Protocol, cast
|
||||
|
||||
from pydantic import BaseModel, ConfigDict, TypeAdapter
|
||||
|
||||
|
|
@ -25,6 +25,7 @@ from litellm._internal_context import current_service_target, service_target
|
|||
from litellm._logging import verbose_logger
|
||||
from litellm.constants import CACHED_STREAMING_CHUNK_DELAY
|
||||
from litellm.integrations.otel.runtime import phase_span
|
||||
from litellm.litellm_core_utils.core_helpers import require_attributes
|
||||
from litellm.litellm_core_utils.model_param_helper import (
|
||||
ModelParamHelper,
|
||||
embeds_thought_signature_in_tool_call_id,
|
||||
|
|
@ -52,6 +53,16 @@ _REDIS_STARTUP_NODES_ADAPTER: Final = TypeAdapter(
|
|||
)
|
||||
|
||||
|
||||
class _SupportsAsyncPing(Protocol):
|
||||
@property
|
||||
def ping(self) -> Callable[[], Awaitable[object]] | None: ...
|
||||
|
||||
|
||||
class _SupportsAsyncDeleteCacheKeys(Protocol):
|
||||
@property
|
||||
def delete_cache_keys(self) -> Callable[[object], Awaitable[object]] | None: ...
|
||||
|
||||
|
||||
def _keyed_provider_params_in_extra_body(kwargs: Mapping[str, object]) -> str:
|
||||
raw_extra_body: Final = kwargs.get("extra_body")
|
||||
if not isinstance(raw_extra_body, dict):
|
||||
|
|
@ -1032,14 +1043,16 @@ class Cache:
|
|||
cache_key, cached_data, kwargs = self._add_cache_logic(result=result, **kwargs)
|
||||
await self.cache.batch_cache_write(cache_key, cached_data, **kwargs)
|
||||
|
||||
async def ping(self):
|
||||
cache_ping: Final = getattr(self.cache, "ping")
|
||||
async def ping(self) -> object:
|
||||
cache_ping: Final = require_attributes(self.cache, _SupportsAsyncPing, "ping").ping
|
||||
if cache_ping:
|
||||
return await cache_ping()
|
||||
return None
|
||||
|
||||
async def delete_cache_keys(self, keys):
|
||||
cache_delete_cache_keys: Final = getattr(self.cache, "delete_cache_keys")
|
||||
async def delete_cache_keys(self, keys: object) -> object:
|
||||
cache_delete_cache_keys: Final = require_attributes(
|
||||
self.cache, _SupportsAsyncDeleteCacheKeys, "delete_cache_keys"
|
||||
).delete_cache_keys
|
||||
if cache_delete_cache_keys:
|
||||
return await cache_delete_cache_keys(keys)
|
||||
return None
|
||||
|
|
|
|||
|
|
@ -22,6 +22,7 @@ from collections.abc import AsyncGenerator, AsyncIterator, Awaitable, Callable,
|
|||
from typing import TYPE_CHECKING, Any, Final, Optional, TypeVar
|
||||
|
||||
from pydantic import ConfigDict, SkipValidation, TypeAdapter, ValidationError
|
||||
from typing_extensions import TypedDict
|
||||
|
||||
import litellm
|
||||
from litellm._internal_context import is_internal_call, post_response_phase
|
||||
|
|
@ -90,6 +91,10 @@ class CachingHandlerResponse(LiteLLMBaseModel):
|
|||
in_memory_cache_obj: Final = InMemoryCache()
|
||||
|
||||
|
||||
class _CacheCallKwargs(TypedDict, total=False, extra_items=object):
|
||||
pass
|
||||
|
||||
|
||||
def _drop_logging_obj_from_kwargs(request_kwargs: dict[str, object]) -> dict[str, object]:
|
||||
"""
|
||||
The caching handler is stored on the Logging object
|
||||
|
|
@ -455,7 +460,7 @@ class LLMCachingHandler:
|
|||
if litellm.cache is not None and self._is_call_type_supported_by_cache(original_function=original_function):
|
||||
args = args or ()
|
||||
# Now that we confirmed caching will happen, prepare kwargs
|
||||
new_kwargs: Final = kwargs.copy()
|
||||
new_kwargs: Final = dict[str, object](kwargs)
|
||||
new_kwargs.update(
|
||||
convert_args_to_kwargs(
|
||||
self.original_function,
|
||||
|
|
@ -470,7 +475,7 @@ class LLMCachingHandler:
|
|||
print_verbose("Checking Sync Cache")
|
||||
logging_obj.reset_caching_details()
|
||||
with response_cache_phase("get"):
|
||||
cached_result = litellm.cache.get_cache(**new_kwargs)
|
||||
cached_result = litellm.cache.get_cache(**_CacheCallKwargs(**new_kwargs))
|
||||
if is_response_without_output(cached_result):
|
||||
verbose_logger.debug("LiteLLM Cache: cached response has no output, treating it as a miss")
|
||||
return CachingHandlerResponse(cached_result=None)
|
||||
|
|
@ -480,7 +485,7 @@ class LLMCachingHandler:
|
|||
pass
|
||||
else:
|
||||
call_type = original_function.__name__
|
||||
cached_result = self._convert_cached_result_to_model_response(
|
||||
converted_result: Final = self._convert_cached_result_to_model_response(
|
||||
cached_result=cached_result,
|
||||
call_type=call_type,
|
||||
kwargs=kwargs,
|
||||
|
|
@ -508,16 +513,16 @@ class LLMCachingHandler:
|
|||
logging_obj=logging_obj,
|
||||
model=f"{custom_llm_provider}/{model}",
|
||||
kwargs=kwargs,
|
||||
cached_result=cached_result,
|
||||
cached_result=converted_result,
|
||||
is_async=False,
|
||||
custom_llm_provider=custom_llm_provider,
|
||||
)
|
||||
|
||||
if not is_internal_call.get() and not _should_defer_streaming_cache_hit_callbacks(
|
||||
cached_result=cached_result
|
||||
cached_result=converted_result
|
||||
):
|
||||
logging_obj.handle_sync_cache_hit_success_callbacks(
|
||||
result=cached_result,
|
||||
result=converted_result,
|
||||
start_time=start_time,
|
||||
end_time=end_time,
|
||||
cache_hit=cache_hit,
|
||||
|
|
@ -527,8 +532,8 @@ class LLMCachingHandler:
|
|||
or self.request_kwargs.get("cache_key")
|
||||
or litellm.cache.get_cache_key(**self.request_kwargs)
|
||||
)
|
||||
_set_cached_hidden_param(cached_result, "cache_key", cache_key)
|
||||
return CachingHandlerResponse(cached_result=cached_result)
|
||||
_set_cached_hidden_param(converted_result, "cache_key", cache_key)
|
||||
return CachingHandlerResponse(cached_result=converted_result)
|
||||
return CachingHandlerResponse(cached_result=cached_result)
|
||||
|
||||
_sync_get_cache = sync_get_cache
|
||||
|
|
@ -950,22 +955,14 @@ class LLMCachingHandler:
|
|||
|
||||
def _convert_cached_result_to_model_response(
|
||||
self,
|
||||
cached_result: Any,
|
||||
cached_result: object,
|
||||
call_type: str,
|
||||
kwargs: dict[str, object],
|
||||
logging_obj: LiteLLMLoggingObj,
|
||||
model: str,
|
||||
args: tuple[object, ...],
|
||||
custom_llm_provider: str | None = None,
|
||||
) -> (
|
||||
ModelResponse
|
||||
| TextCompletionResponse
|
||||
| EmbeddingResponse
|
||||
| RerankResponse
|
||||
| TranscriptionResponse
|
||||
| CustomStreamWrapper
|
||||
| None
|
||||
):
|
||||
) -> object:
|
||||
"""
|
||||
Internal method to process the cached result
|
||||
|
||||
|
|
@ -1141,7 +1138,7 @@ class LLMCachingHandler:
|
|||
self,
|
||||
result: object,
|
||||
original_function: Callable,
|
||||
kwargs: dict[str, Any],
|
||||
kwargs: Mapping[str, object],
|
||||
args: tuple[object, ...] | None = None,
|
||||
):
|
||||
"""
|
||||
|
|
@ -1169,7 +1166,7 @@ class LLMCachingHandler:
|
|||
return
|
||||
cache: Final = litellm.cache
|
||||
|
||||
new_kwargs: Final = kwargs.copy()
|
||||
new_kwargs: Final = dict(kwargs)
|
||||
new_kwargs.update(
|
||||
convert_args_to_kwargs(
|
||||
original_function,
|
||||
|
|
@ -1206,7 +1203,7 @@ class LLMCachingHandler:
|
|||
)
|
||||
)
|
||||
else:
|
||||
create_cache_write_task(lambda: cache.async_add_cache(result, **new_kwargs))
|
||||
create_cache_write_task(lambda: cache.async_add_cache(result, **_CacheCallKwargs(**new_kwargs)))
|
||||
|
||||
def sync_set_cache(
|
||||
self,
|
||||
|
|
|
|||
|
|
@ -12,14 +12,14 @@ import asyncio
|
|||
import itertools
|
||||
import logging
|
||||
import time
|
||||
from collections.abc import Mapping, Sequence
|
||||
from collections.abc import Awaitable, Mapping, Sequence
|
||||
from dataclasses import dataclass
|
||||
from threading import Lock
|
||||
from typing import TYPE_CHECKING, Any, Final
|
||||
|
||||
import litellm
|
||||
from litellm._logging import print_verbose, verbose_logger
|
||||
from litellm.constants import DEFAULT_MAX_REDIS_BATCH_CACHE_SIZE
|
||||
from litellm.constants import DEFAULT_MAX_REDIS_BATCH_CACHE_SIZE, DEFAULT_REDIS_BATCH_CACHE_EXPIRY_SECONDS
|
||||
from litellm.types.caching import RedisPipelineIncrementOperation
|
||||
|
||||
from .base_cache import BaseCache
|
||||
|
|
@ -91,7 +91,9 @@ class DualCache(BaseCache):
|
|||
self.last_redis_batch_access_time = LimitedSizeOrderedDict(max_size=default_max_redis_batch_cache_size)
|
||||
self._last_redis_batch_access_time_lock = Lock()
|
||||
self.redis_batch_cache_expiry = (
|
||||
default_redis_batch_cache_expiry or litellm.default_redis_batch_cache_expiry or 10
|
||||
default_redis_batch_cache_expiry
|
||||
or litellm.default_redis_batch_cache_expiry
|
||||
or DEFAULT_REDIS_BATCH_CACHE_EXPIRY_SECONDS
|
||||
)
|
||||
self.default_in_memory_ttl = default_in_memory_ttl or litellm.default_in_memory_ttl
|
||||
self.default_redis_ttl = default_redis_ttl or litellm.default_redis_ttl
|
||||
|
|
@ -497,7 +499,7 @@ class DualCache(BaseCache):
|
|||
batch: Final = None if self.redis_cache is None else active_request_redis_batch(self.redis_cache)
|
||||
return None if batch is None else await self._set_on_batch(batch, key, value, ttl)
|
||||
|
||||
async def async_delete_cache_pre_call(self, key: str) -> BatchResult[None] | None:
|
||||
async def async_delete_cache_pre_call(self, key: str) -> Awaitable[None] | None:
|
||||
"""Memory now, the Redis DEL on the request's pipeline; None when no pipeline is open, so the caller
|
||||
takes its direct path."""
|
||||
batch: Final = None if self.redis_cache is None else active_request_redis_batch(self.redis_cache)
|
||||
|
|
|
|||
|
|
@ -622,9 +622,11 @@ class RequestRedisBatches:
|
|||
self._deadline_flush = asyncio.ensure_future(self.flush_post_call())
|
||||
|
||||
async def flush_all(self) -> None:
|
||||
"""Batch whatever is still declared (write-backs nobody awaits) before the request scope closes."""
|
||||
"""Batch whatever is still declared (write-backs nobody awaits) before the request scope closes. The pre-call
|
||||
flush hooks are skipped: they declare reads for later readers, and none are left."""
|
||||
await asyncio.gather(
|
||||
*(batch.flush() for batch in self._batches.values() if batch.pending), self.flush_native_pre()
|
||||
*(batch.flush() for batch in self._batches.values() if batch.pending),
|
||||
*(() if self._native is None else (self._native.flush_pre(),)),
|
||||
)
|
||||
|
||||
async def flush_post_call(self) -> None:
|
||||
|
|
|
|||
|
|
@ -40,9 +40,9 @@ from litellm.constants import (
|
|||
REDIS_CIRCUIT_BREAKER_TIMEOUT_MIN_DURATION,
|
||||
REDIS_TIMEOUT_LOG_INTERVAL,
|
||||
)
|
||||
from litellm.litellm_core_utils.core_helpers import get_parent_otel_span_from_kwargs
|
||||
from litellm.litellm_core_utils.core_helpers import get_parent_otel_span_from_kwargs, has_attributes
|
||||
from litellm.litellm_core_utils.coroutine_checker import coroutine_checker
|
||||
from litellm.redis.store_settings import RedisStoreSettings
|
||||
from litellm.redis.store_settings import RedisStoreSettings, breaker_settings
|
||||
from litellm.types.caching import (
|
||||
RedisPipelineIncrementOperation,
|
||||
RedisPipelineLpopOperation,
|
||||
|
|
@ -97,6 +97,22 @@ class _AsyncRedisCommands(Protocol):
|
|||
def eval(self, script: str, numkeys: int, *keys_and_args: str | bytes | float) -> Awaitable[object]: ...
|
||||
|
||||
|
||||
class _AsyncLuaScript(Protocol):
|
||||
def __call__(
|
||||
self, keys: Sequence[str], args: Sequence[str | bytes | int | float], client: object
|
||||
) -> Awaitable[object]: ...
|
||||
|
||||
|
||||
class _SupportsRegisterScript(Protocol):
|
||||
def register_script(self, script: str, /) -> _AsyncLuaScript: ...
|
||||
|
||||
|
||||
class _SupportsScriptLoad(Protocol):
|
||||
def script_load(self, script: str, /) -> object: ...
|
||||
|
||||
def evalsha(self, sha: object, numkeys: int, /, *keys_and_args: str | bytes | float) -> Awaitable[object]: ...
|
||||
|
||||
|
||||
_GENERIC_CALLER_MODULES: Final = frozenset(
|
||||
{
|
||||
__name__,
|
||||
|
|
@ -661,7 +677,7 @@ class RedisCache(BaseCache):
|
|||
):
|
||||
from litellm._service_logger import ServiceLogging
|
||||
|
||||
from .._redis import get_redis_client, get_redis_connection_pool
|
||||
from .._redis import get_redis_client, get_redis_connection_pool, merge_redis_kwargs
|
||||
|
||||
redis_kwargs: Final = {}
|
||||
if host is not None:
|
||||
|
|
@ -715,9 +731,10 @@ class RedisCache(BaseCache):
|
|||
else:
|
||||
super().__init__() # defaults to 60s
|
||||
self.store_settings = RedisStoreSettings(
|
||||
redis_kwargs=MappingProxyType(dict(redis_kwargs)),
|
||||
redis_kwargs=MappingProxyType(merge_redis_kwargs(redis_kwargs)),
|
||||
namespace=namespace,
|
||||
default_ttl=self.default_ttl,
|
||||
breaker=breaker_settings(),
|
||||
)
|
||||
|
||||
def _admit_async_call(self, call: Callable[[], Awaitable[_RedisCallResult]]) -> Awaitable[_RedisCallResult]:
|
||||
|
|
@ -843,6 +860,9 @@ class RedisCache(BaseCache):
|
|||
def _async_commands(self) -> _AsyncRedisCommands:
|
||||
return self.init_async_client()
|
||||
|
||||
def _async_client_object(self) -> object:
|
||||
return self.init_async_client()
|
||||
|
||||
def check_and_fix_namespace(self, key: str) -> str:
|
||||
"""
|
||||
Make sure each key starts with the given namespace
|
||||
|
|
@ -1116,8 +1136,8 @@ class RedisCache(BaseCache):
|
|||
Kept separate from async_register_script so each loop caches its own
|
||||
executor; see that method for why the binding must be per loop.
|
||||
"""
|
||||
_redis_client: Final[Any] = self.init_async_client()
|
||||
if hasattr(_redis_client, "register_script"):
|
||||
_redis_client: Final = self._async_client_object()
|
||||
if has_attributes(_redis_client, _SupportsRegisterScript, "register_script"):
|
||||
registered_script: Final = _redis_client.register_script(script)
|
||||
|
||||
async def standalone_executor(
|
||||
|
|
@ -1130,7 +1150,7 @@ class RedisCache(BaseCache):
|
|||
|
||||
return standalone_executor
|
||||
|
||||
if hasattr(_redis_client, "script_load"):
|
||||
if has_attributes(_redis_client, _SupportsScriptLoad, "script_load"):
|
||||
script_sha: Final = _redis_client.script_load(script)
|
||||
|
||||
async def cluster_executor(
|
||||
|
|
|
|||
|
|
@ -43,6 +43,7 @@ if TYPE_CHECKING:
|
|||
_CACHE_RESPONSE_ADAPTER: Final[TypeAdapter[str | bytes | None]] = TypeAdapter(
|
||||
str | bytes | None, config=ConfigDict(strict=True, hide_input_in_errors=True)
|
||||
)
|
||||
_DECODED_CACHE_VALUE: Final = TypeAdapter(object)
|
||||
|
||||
|
||||
class RedisSemanticCache(BaseCache):
|
||||
|
|
@ -407,10 +408,10 @@ class RedisSemanticCache(BaseCache):
|
|||
|
||||
# Convert string representation to Python object
|
||||
try:
|
||||
return json.loads(response_text)
|
||||
return _DECODED_CACHE_VALUE.validate_python(json.loads(response_text))
|
||||
except json.JSONDecodeError:
|
||||
try:
|
||||
return ast.literal_eval(response_text)
|
||||
return _DECODED_CACHE_VALUE.validate_python(ast.literal_eval(response_text))
|
||||
except (ValueError, SyntaxError) as e:
|
||||
print_verbose(f"Error parsing cached response: {e}")
|
||||
return None
|
||||
|
|
|
|||
|
|
@ -39,7 +39,7 @@ from litellm.litellm_core_utils.allowed_tools_tool_choice import (
|
|||
parse_chat_allowed_tools_tool_choice,
|
||||
responses_allowed_tools_tool_choice,
|
||||
)
|
||||
from litellm.litellm_core_utils.core_helpers import drops_unsupported_params
|
||||
from litellm.litellm_core_utils.core_helpers import drops_unsupported_params, iter_items
|
||||
from litellm.litellm_core_utils.hidden_params import get_hidden_params, get_or_create_hidden_params
|
||||
from litellm.litellm_core_utils.prompt_templates.prompt_cache_breakpoint_utils import (
|
||||
prompt_cache_breakpoint_with_readable_mode,
|
||||
|
|
@ -54,6 +54,7 @@ from litellm.llms.base_llm.bridges.completion_transformation import (
|
|||
CompletionTransformationBridge,
|
||||
)
|
||||
from litellm.llms.custom_httpx.sse_lines import split_sse_lines
|
||||
from litellm.responses.additional_tools import additional_tools_in_input
|
||||
from litellm.responses.litellm_completion_transformation.tool_search import is_anthropic_tool_search_chat_tool
|
||||
from litellm.responses.sse_output_recovery import (
|
||||
parse_sse_json_chunk,
|
||||
|
|
@ -348,6 +349,9 @@ _TOOL_SEARCH_REPLAY_ITEM_TYPES: Final = frozenset({"reasoning", *_TOOL_SEARCH_IT
|
|||
_TOOL_SEARCH_ITEMS_KEY: Final = "tool_search_items"
|
||||
|
||||
|
||||
_DECODED_JSON: Final = TypeAdapter(object)
|
||||
|
||||
|
||||
def _json_field(value: object, key: str) -> object:
|
||||
return cast(Mapping[str, object], value).get(key) if isinstance(value, dict) else None # cast-ok: untyped JSON
|
||||
|
||||
|
|
@ -426,6 +430,17 @@ def replayed_tool_search_tool_names(request: object) -> tuple[str, ...]:
|
|||
return tuple(chain.from_iterable(_tool_definition_name(tool) for tool in replayed_tool_search_tools(request)))
|
||||
|
||||
|
||||
def _top_level_namespaces(request: object) -> Iterator[object]:
|
||||
tools: Final = _json_list(_json_field(request, "tools")) or ()
|
||||
return (tool for tool in tools if _json_field(tool, "type") == "namespace")
|
||||
|
||||
|
||||
def nested_tool_member_names(request: object) -> tuple[str, ...]:
|
||||
tools: Final = (*_top_level_namespaces(request), *additional_tools_in_input(_json_field(request, "input")))
|
||||
members: Final = chain.from_iterable(_namespace_members(tool) for tool in tools)
|
||||
return tuple(chain.from_iterable(_tool_definition_name(member) for member in members))
|
||||
|
||||
|
||||
def _loads_any_tool(message: object, names: frozenset[str]) -> bool:
|
||||
loaded_names: Final = chain.from_iterable(_tool_definition_name(tool) for tool in _message_loaded_tools(message))
|
||||
return any(name in names for name in loaded_names)
|
||||
|
|
@ -845,8 +860,8 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge):
|
|||
|
||||
# Handle message items with output_text content
|
||||
if item_type == "message":
|
||||
content_list: Final = item.get("content", [])
|
||||
for content_item in content_list:
|
||||
content_list: Final = _DECODED_JSON.validate_python(item.get("content", []))
|
||||
for content_item in iter_items(content_list):
|
||||
if isinstance(content_item, dict):
|
||||
content_type = content_item.get("type")
|
||||
if content_type == "output_text":
|
||||
|
|
|
|||
|
|
@ -159,6 +159,7 @@ DEFAULT_ALLOWED_FAILS: Final = get_env_int("DEFAULT_ALLOWED_FAILS", 3)
|
|||
DEFAULT_REDIS_SYNC_INTERVAL: Final = get_env_int("DEFAULT_REDIS_SYNC_INTERVAL", 1)
|
||||
DEFAULT_COOLDOWN_TIME_SECONDS: Final = get_env_int("DEFAULT_COOLDOWN_TIME_SECONDS", 5)
|
||||
DEFAULT_COOLDOWN_REDIS_READ_INTERVAL_SECONDS: Final = get_env_float("DEFAULT_COOLDOWN_REDIS_READ_INTERVAL_SECONDS", 1.0)
|
||||
DEFAULT_REDIS_BATCH_CACHE_EXPIRY_SECONDS: Final = 10
|
||||
DEFAULT_REPLICATE_POLLING_RETRIES: Final = get_env_int("DEFAULT_REPLICATE_POLLING_RETRIES", 5)
|
||||
DEFAULT_REPLICATE_POLLING_DELAY_SECONDS: Final = get_env_int("DEFAULT_REPLICATE_POLLING_DELAY_SECONDS", 1)
|
||||
DEFAULT_IMAGE_TOKEN_COUNT: Final = get_env_int("DEFAULT_IMAGE_TOKEN_COUNT", 250)
|
||||
|
|
|
|||
|
|
@ -8,7 +8,7 @@ from datetime import datetime
|
|||
from functools import lru_cache
|
||||
from itertools import accumulate
|
||||
from types import MappingProxyType
|
||||
from typing import TYPE_CHECKING, Any, Final, Literal, TypeAlias, cast
|
||||
from typing import TYPE_CHECKING, Any, Final, Literal, Protocol, TypeAlias, cast
|
||||
|
||||
from httpx import Response
|
||||
from pydantic import BaseModel, ConfigDict, TypeAdapter, ValidationError
|
||||
|
|
@ -24,7 +24,12 @@ from litellm.constants import (
|
|||
REALTIME_METERED_USAGE_KEY,
|
||||
)
|
||||
from litellm.litellm_core_utils.audio_utils.utils import normalized_audio_duration_seconds
|
||||
from litellm.litellm_core_utils.core_helpers import as_str_mapping, attribute_or_none, billed_usage_cost
|
||||
from litellm.litellm_core_utils.core_helpers import (
|
||||
as_str_mapping,
|
||||
attribute_or_none,
|
||||
billed_usage_cost,
|
||||
require_attributes,
|
||||
)
|
||||
from litellm.litellm_core_utils.cost_discount import resolve_cost_discount
|
||||
from litellm.litellm_core_utils.decisions_pricing import (
|
||||
decisions_pricing_fields,
|
||||
|
|
@ -1188,6 +1193,10 @@ def _model_contains_known_llm_provider(model: str) -> bool:
|
|||
return _provider_prefix in LlmProvidersSet
|
||||
|
||||
|
||||
class _SupportsGet(Protocol):
|
||||
def get(self, key: str, /) -> object: ...
|
||||
|
||||
|
||||
def _get_response_model(completion_response: object) -> str | None:
|
||||
"""
|
||||
Extract the model name from a completion response object.
|
||||
|
|
@ -1219,7 +1228,7 @@ _GEMINI_TRAFFIC_TYPE_TO_SERVICE_TIER: Final[dict] = {
|
|||
}
|
||||
|
||||
|
||||
def _map_traffic_type_to_service_tier(traffic_type: str | None) -> str | None:
|
||||
def _map_traffic_type_to_service_tier(traffic_type: object) -> str | None:
|
||||
"""
|
||||
Map a Gemini usageMetadata.trafficType value to a LiteLLM service_tier string.
|
||||
|
||||
|
|
@ -1273,11 +1282,13 @@ def _served_service_tier(completion_response: object, usage_object: Usage | None
|
|||
usage_tier: Final = _extract_service_tier(usage_object)
|
||||
if isinstance(usage_tier, str):
|
||||
return usage_tier
|
||||
hidden_params: Final = getattr(completion_response, "_hidden_params", None)
|
||||
hidden_params: Final[object] = getattr(completion_response, "_hidden_params", None)
|
||||
if hidden_params is None:
|
||||
return None
|
||||
provider_specific: Final = hidden_params.get("provider_specific_fields") or {}
|
||||
raw_traffic_type: Final = provider_specific.get("traffic_type")
|
||||
provider_specific: Final[object] = (
|
||||
require_attributes(hidden_params, _SupportsGet, "get").get("provider_specific_fields") or {}
|
||||
)
|
||||
raw_traffic_type: Final = require_attributes(provider_specific, _SupportsGet, "get").get("traffic_type")
|
||||
if not raw_traffic_type:
|
||||
return None
|
||||
return _map_traffic_type_to_service_tier(raw_traffic_type) or "default"
|
||||
|
|
@ -1341,8 +1352,16 @@ def get_usage_object(
|
|||
return None
|
||||
|
||||
|
||||
class _DumpedUsageFields(TypedDict, total=False, extra_items=object):
|
||||
pass
|
||||
|
||||
|
||||
def _dumped_usage_fields(usage_model: BaseModel) -> Mapping[str, object]:
|
||||
return usage_model.model_dump()
|
||||
|
||||
|
||||
def _usage_from_model(usage_model: BaseModel) -> Usage:
|
||||
return Usage(**usage_model.model_dump())
|
||||
return Usage(**_DumpedUsageFields(**_dumped_usage_fields(usage_model)))
|
||||
|
||||
|
||||
def _decisions_prompt_tokens_details(usage: DecisionsUsage | OpenAIDecisionUsage) -> PromptTokensDetailsWrapper:
|
||||
|
|
|
|||
|
|
@ -12,10 +12,11 @@
|
|||
import datetime
|
||||
import enum
|
||||
from collections.abc import Mapping, Sequence
|
||||
from typing import Any, Final
|
||||
from typing import Any, Final, Protocol, SupportsIndex, SupportsInt, runtime_checkable
|
||||
|
||||
import httpx
|
||||
import openai
|
||||
from typing_extensions import Buffer
|
||||
|
||||
import litellm
|
||||
from litellm.types.utils import LiteLLMCommonStrings
|
||||
|
|
@ -1250,6 +1251,19 @@ class BlockedPiiEntityError(Exception):
|
|||
super().__init__(self.message)
|
||||
|
||||
|
||||
@runtime_checkable
|
||||
class _SupportsTrunc(Protocol):
|
||||
def __trunc__(self) -> int: ...
|
||||
|
||||
|
||||
def _status_code_int(value: object) -> int:
|
||||
if isinstance(value, (str, Buffer, SupportsInt, SupportsIndex, _SupportsTrunc)):
|
||||
return int(value)
|
||||
raise TypeError(
|
||||
f"int() argument must be a string, a bytes-like object or a real number, not '{type(value).__name__}'"
|
||||
)
|
||||
|
||||
|
||||
class MidStreamFallbackError(ServiceUnavailableError):
|
||||
def __init__(
|
||||
self,
|
||||
|
|
@ -1265,8 +1279,8 @@ class MidStreamFallbackError(ServiceUnavailableError):
|
|||
is_pre_first_chunk: bool = False,
|
||||
emitted_disqualifying_content: bool = False,
|
||||
):
|
||||
original_status: Final = getattr(original_exception, "status_code", None)
|
||||
self.status_code = int(original_status) if original_status is not None else 503
|
||||
original_status: Final[object] = getattr(original_exception, "status_code", None)
|
||||
self.status_code = _status_code_int(original_status) if original_status is not None else 503
|
||||
self.message = f"litellm.MidStreamFallbackError: {message}"
|
||||
self.model = model
|
||||
self.llm_provider = llm_provider
|
||||
|
|
@ -1309,7 +1323,7 @@ class MidStreamFallbackError(ServiceUnavailableError):
|
|||
)
|
||||
|
||||
# Restore the propagated status and original response/request objects
|
||||
self.status_code = int(original_status) if original_status is not None else 503
|
||||
self.status_code = _status_code_int(original_status) if original_status is not None else 503
|
||||
self.response = _saved_response
|
||||
self.request = _saved_request
|
||||
self.message = _saved_message
|
||||
|
|
|
|||
|
|
@ -14,10 +14,11 @@ import contextlib
|
|||
import itertools
|
||||
import json
|
||||
import logging
|
||||
import math
|
||||
import secrets
|
||||
from collections.abc import AsyncIterable, AsyncIterator, Mapping
|
||||
from collections.abc import AsyncIterable, AsyncIterator, Callable, Mapping
|
||||
from dataclasses import dataclass
|
||||
from types import MappingProxyType, ModuleType
|
||||
from types import MappingProxyType
|
||||
from typing import TYPE_CHECKING, Any, Final, Protocol
|
||||
|
||||
import httpx
|
||||
|
|
@ -36,6 +37,7 @@ from litellm.constants import (
|
|||
from litellm.harness.context import GatewayTarget
|
||||
from litellm.harness.errors import HarnessError, HarnessInstallFailed
|
||||
from litellm.harness.types import Harness, Usage
|
||||
from litellm.litellm_core_utils.core_helpers import has_attributes
|
||||
from litellm.llms.custom_httpx.http_handler import get_async_httpx_client
|
||||
from litellm.types.litellm_params import STREAM_USAGE_INJECTED_FIELD
|
||||
from litellm.types.llms.custom_http import httpxSpecialProvider
|
||||
|
|
@ -45,7 +47,7 @@ if TYPE_CHECKING:
|
|||
from starlette.requests import Request
|
||||
from starlette.responses import JSONResponse, Response, StreamingResponse
|
||||
from starlette.routing import Route
|
||||
from uvicorn import Server
|
||||
from uvicorn import Config, Server
|
||||
|
||||
verbose_logger: Final = logging.getLogger("LiteLLM")
|
||||
_JSON_VALUE_ADAPTER: Final = TypeAdapter(object, config=ConfigDict(strict=True))
|
||||
|
|
@ -110,9 +112,22 @@ class _ResponsesModule(Protocol):
|
|||
StreamingResponse: type[StreamingResponse]
|
||||
|
||||
|
||||
class _UvicornModule(Protocol):
|
||||
Config: type[Config]
|
||||
Server: type[Server]
|
||||
|
||||
|
||||
class _CapturesSignals(Protocol):
|
||||
capture_signals: Callable[[], contextlib.AbstractContextManager[None]]
|
||||
|
||||
|
||||
class _InstallsSignalHandlers(Protocol):
|
||||
install_signal_handlers: Callable[[], None]
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class _ServerDeps:
|
||||
uvicorn: ModuleType
|
||||
uvicorn: _UvicornModule
|
||||
applications: _ApplicationsModule
|
||||
routing: _RoutingModule
|
||||
responses: _ResponsesModule
|
||||
|
|
@ -438,6 +453,14 @@ def _noop() -> None:
|
|||
return None
|
||||
|
||||
|
||||
def _signal_capture(server: object) -> _CapturesSignals | None:
|
||||
return server if has_attributes(server, _CapturesSignals, "capture_signals") else None
|
||||
|
||||
|
||||
def _signal_handler_installer(server: object) -> _InstallsSignalHandlers | None:
|
||||
return server if has_attributes(server, _InstallsSignalHandlers, "install_signal_handlers") else None
|
||||
|
||||
|
||||
class ModelEndpoint:
|
||||
"""Local HTTP endpoint for one harness session. Use as an async context manager."""
|
||||
|
||||
|
|
@ -521,7 +544,7 @@ class ModelEndpoint:
|
|||
await asyncio.sleep(DEFAULT_POLLING_INTERVAL)
|
||||
|
||||
def _build_server(self, deps: _ServerDeps) -> Server:
|
||||
config = deps.uvicorn.Config(
|
||||
config: Final = deps.uvicorn.Config(
|
||||
self._build_app(deps),
|
||||
host=HARNESS_ENDPOINT_HOST,
|
||||
port=0,
|
||||
|
|
@ -529,14 +552,16 @@ class ModelEndpoint:
|
|||
log_level="warning",
|
||||
access_log=False,
|
||||
lifespan="off",
|
||||
timeout_graceful_shutdown=HARNESS_PROCESS_KILL_GRACE_SECONDS,
|
||||
timeout_graceful_shutdown=math.ceil(HARNESS_PROCESS_KILL_GRACE_SECONDS),
|
||||
)
|
||||
server = deps.uvicorn.Server(config)
|
||||
server: Final = deps.uvicorn.Server(config)
|
||||
# Never touch the host process's signal handlers.
|
||||
if hasattr(server, "capture_signals"):
|
||||
server.capture_signals = contextlib.nullcontext
|
||||
if hasattr(server, "install_signal_handlers"):
|
||||
server.install_signal_handlers = _noop
|
||||
signal_capture: Final = _signal_capture(server)
|
||||
if signal_capture is not None:
|
||||
signal_capture.capture_signals = contextlib.nullcontext
|
||||
handler_installer: Final = _signal_handler_installer(server)
|
||||
if handler_installer is not None:
|
||||
handler_installer.install_signal_handlers = _noop
|
||||
return server
|
||||
|
||||
def _build_app(self, deps: _ServerDeps) -> Starlette:
|
||||
|
|
|
|||
|
|
@ -6,12 +6,25 @@ from pathlib import Path
|
|||
from typing import Any, Final, Protocol, cast, get_type_hints, runtime_checkable
|
||||
|
||||
import litellm
|
||||
from litellm.litellm_core_utils.core_helpers import has_attributes
|
||||
from litellm.litellm_core_utils.token_counter import get_image_type
|
||||
from litellm.llms.base_llm.image_edit.transformation import BaseImageEditConfig
|
||||
from litellm.types.files import FILE_MIME_TYPES, FileType
|
||||
from litellm.types.images.main import ImageEditOptionalRequestParams
|
||||
|
||||
|
||||
class _ReadsBytes(Protocol):
|
||||
def read(self, size: int, /) -> bytes: ...
|
||||
|
||||
|
||||
class _Tells(Protocol):
|
||||
def tell(self) -> int: ...
|
||||
|
||||
|
||||
class _Seeks(Protocol):
|
||||
def seek(self, offset: int, /) -> object: ...
|
||||
|
||||
|
||||
class ImageEditRequestUtils:
|
||||
@staticmethod
|
||||
def get_optional_params_image_edit(
|
||||
|
|
@ -114,12 +127,12 @@ class ImageEditRequestUtils:
|
|||
bytes_data = image_data[:100]
|
||||
else:
|
||||
# For other types, try to read if possible
|
||||
if hasattr(image_data, "read"):
|
||||
current_pos = getattr(image_data, "tell", lambda: 0)()
|
||||
if hasattr(image_data, "seek"):
|
||||
if has_attributes(image_data, _ReadsBytes, "read"):
|
||||
current_pos = image_data.tell() if has_attributes(image_data, _Tells, "tell") else 0
|
||||
if has_attributes(image_data, _Seeks, "seek"):
|
||||
image_data.seek(0)
|
||||
bytes_data = image_data.read(100)
|
||||
if hasattr(image_data, "seek"):
|
||||
if has_attributes(image_data, _Seeks, "seek"):
|
||||
image_data.seek(current_pos)
|
||||
else:
|
||||
return FILE_MIME_TYPES[FileType.PNG] # Default fallback
|
||||
|
|
|
|||
|
|
@ -53,11 +53,13 @@ async def fetch_user_spend_rows(
|
|||
today: datetime.date,
|
||||
baseline_days: int,
|
||||
) -> tuple[UserSpendRow, ...]:
|
||||
from litellm.proxy.utils import query_raw_rows
|
||||
|
||||
today_str: Final = today.strftime("%Y-%m-%d")
|
||||
month_start_str: Final = today.replace(day=1).strftime("%Y-%m-%d")
|
||||
baseline_start_str: Final = (today - datetime.timedelta(days=max(baseline_days, 1))).strftime("%Y-%m-%d")
|
||||
return USER_SPEND_ROWS_ADAPTER.validate_python(
|
||||
await prisma_client.db.query_raw(USER_SPEND_QUERY, today_str, month_start_str, baseline_start_str)
|
||||
await query_raw_rows(prisma_client.db, USER_SPEND_QUERY, today_str, month_start_str, baseline_start_str)
|
||||
)
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -7,11 +7,12 @@ import json
|
|||
import os
|
||||
import random
|
||||
import types
|
||||
from collections.abc import Mapping
|
||||
from typing import Final
|
||||
from collections.abc import Mapping, Sequence
|
||||
from typing import Final, Protocol
|
||||
|
||||
import httpx
|
||||
from pydantic import BaseModel
|
||||
from typing_extensions import ReadOnly, TypedDict
|
||||
|
||||
import litellm
|
||||
from litellm._logging import verbose_logger
|
||||
|
|
@ -29,6 +30,22 @@ from litellm.types.integrations.argilla import (
|
|||
from litellm.types.utils import StandardLoggingPayload
|
||||
|
||||
|
||||
class _ArgillaDataset(TypedDict):
|
||||
id: ReadOnly[str]
|
||||
|
||||
|
||||
class _ArgillaDatasetsBody(TypedDict, total=False):
|
||||
items: ReadOnly[Sequence[_ArgillaDataset]]
|
||||
|
||||
|
||||
class _ArgillaDatasetsResponse(Protocol):
|
||||
def json(self) -> _ArgillaDatasetsBody: ...
|
||||
|
||||
|
||||
def _datasets_body(response: _ArgillaDatasetsResponse) -> _ArgillaDatasetsBody:
|
||||
return response.json()
|
||||
|
||||
|
||||
def is_serializable(value):
|
||||
non_serializable_types: Final = (
|
||||
types.CoroutineType,
|
||||
|
|
@ -102,7 +119,7 @@ class ArgillaLogger(CustomBatchLogger):
|
|||
url=f"{_credentials_base_url}/api/v1/me/datasets?name={_credentials_dataset_name}",
|
||||
headers={"X-Argilla-Api-Key": _credentials_api_key},
|
||||
)
|
||||
json_response: Final = dataset_response.json()
|
||||
json_response: Final = _datasets_body(dataset_response)
|
||||
if (
|
||||
"items" in json_response
|
||||
and isinstance(json_response["items"], list)
|
||||
|
|
|
|||
|
|
@ -1,6 +1,6 @@
|
|||
import json
|
||||
from collections.abc import Mapping
|
||||
from typing import TYPE_CHECKING, Any, Final
|
||||
from typing import TYPE_CHECKING, Final, Protocol
|
||||
|
||||
from pydantic import ConfigDict, TypeAdapter
|
||||
from typing_extensions import ReadOnly, TypedDict, override
|
||||
|
|
@ -10,10 +10,12 @@ from litellm.integrations.opentelemetry_utils.base_otel_llm_obs_attributes impor
|
|||
BaseLLMObsOTELAttributes,
|
||||
safe_set_attribute,
|
||||
)
|
||||
from litellm.litellm_core_utils.core_helpers import iter_items, require_attributes, subscript
|
||||
from litellm.litellm_core_utils.redact_messages import (
|
||||
should_redact_message_logging,
|
||||
)
|
||||
from litellm.litellm_core_utils.safe_json_dumps import safe_dumps
|
||||
from litellm.llms.anthropic.common_utils import tool_calls_with_server_tool_calls
|
||||
from litellm.types.utils import CallTypes, StandardLoggingMCPToolCall, StandardLoggingPayload
|
||||
|
||||
if TYPE_CHECKING:
|
||||
|
|
@ -33,32 +35,40 @@ from litellm.integrations._types.open_inference import (
|
|||
_JSON_OBJECT: Final = TypeAdapter(dict[str, object], config=ConfigDict(hide_input_in_errors=True))
|
||||
|
||||
|
||||
class _SupportsGetWithDefault(Protocol):
|
||||
def get(self, key: str, default: object = None, /) -> object: ...
|
||||
|
||||
|
||||
class ArizeOTELAttributes(BaseLLMObsOTELAttributes):
|
||||
@staticmethod
|
||||
@override
|
||||
def set_messages(span: "Span", kwargs: dict[str, Any]):
|
||||
def set_messages(span: "Span", kwargs: Mapping[str, object]):
|
||||
messages: Final = kwargs.get("messages")
|
||||
|
||||
# for /chat/completions
|
||||
# https://docs.arize.com/arize/large-language-models/tracing/semantic-conventions
|
||||
if messages:
|
||||
last_message: Final = messages[-1]
|
||||
last_message: Final = subscript(messages, -1)
|
||||
safe_set_attribute(
|
||||
span,
|
||||
SpanAttributes.INPUT_VALUE,
|
||||
last_message.get("content", ""),
|
||||
require_attributes(last_message, _SupportsGetWithDefault, "get").get("content", ""),
|
||||
)
|
||||
|
||||
# LLM_INPUT_MESSAGES shows up under `input_messages` tab on the span page.
|
||||
for idx, msg in enumerate(messages):
|
||||
for idx, msg in enumerate(iter_items(messages)):
|
||||
prefix = f"{SpanAttributes.LLM_INPUT_MESSAGES}.{idx}"
|
||||
# Set the role per message.
|
||||
safe_set_attribute(span, f"{prefix}.{MessageAttributes.MESSAGE_ROLE}", msg.get("role"))
|
||||
safe_set_attribute(
|
||||
span,
|
||||
f"{prefix}.{MessageAttributes.MESSAGE_ROLE}",
|
||||
require_attributes(msg, _SupportsGetWithDefault, "get").get("role"),
|
||||
)
|
||||
# Set the content per message.
|
||||
safe_set_attribute(
|
||||
span,
|
||||
f"{prefix}.{MessageAttributes.MESSAGE_CONTENT}",
|
||||
msg.get("content", ""),
|
||||
require_attributes(msg, _SupportsGetWithDefault, "get").get("content", ""),
|
||||
)
|
||||
|
||||
# Additive: emit structured tool_calls / multimodal content
|
||||
|
|
@ -663,12 +673,11 @@ def _to_plain_dict(value):
|
|||
|
||||
|
||||
def _get_tool_calls(message) -> list | None:
|
||||
"""Return ``message.tool_calls`` only when it's a non-empty list.
|
||||
|
||||
Works for dicts and Pydantic message objects via ``_safe_get``.
|
||||
"""Return the message's tool calls, Claude's server tool calls first, or
|
||||
``None`` when there are none. Works for dicts and Pydantic message objects.
|
||||
"""
|
||||
tool_calls: Final = _safe_get(message, "tool_calls")
|
||||
return tool_calls if isinstance(tool_calls, list) and tool_calls else None
|
||||
tool_calls: Final = tool_calls_with_server_tool_calls(message)
|
||||
return list(tool_calls) if tool_calls else None
|
||||
|
||||
|
||||
class _NormalizedToolCallFunction(TypedDict):
|
||||
|
|
|
|||
|
|
@ -18,10 +18,11 @@ import time
|
|||
import traceback
|
||||
from collections.abc import Mapping, Sequence
|
||||
from types import MappingProxyType
|
||||
from typing import Final, TypeVar
|
||||
from typing import Final, Protocol, TypeVar
|
||||
from urllib.parse import urlparse
|
||||
|
||||
import httpx
|
||||
from typing_extensions import ReadOnly, TypedDict
|
||||
|
||||
from litellm._logging import verbose_logger
|
||||
from litellm.integrations.batch_utils import (
|
||||
|
|
@ -52,6 +53,19 @@ MONITOR_SCOPE_BY_AUTHORITY_HOST: Final[Mapping[str, str]] = MappingProxyType(
|
|||
)
|
||||
|
||||
|
||||
class _OAuthTokenBody(TypedDict, total=False):
|
||||
access_token: ReadOnly[str]
|
||||
expires_in: ReadOnly[int]
|
||||
|
||||
|
||||
class _OAuthTokenResponse(Protocol):
|
||||
def json(self) -> _OAuthTokenBody: ...
|
||||
|
||||
|
||||
def _oauth_token_body(response: _OAuthTokenResponse) -> _OAuthTokenBody:
|
||||
return response.json()
|
||||
|
||||
|
||||
class AzureSentinelLogger(CustomBatchLogger):
|
||||
"""
|
||||
Logger that sends LiteLLM logs to Azure Sentinel via Azure Monitor Logs Ingestion API
|
||||
|
|
@ -226,7 +240,7 @@ class AzureSentinelLogger(CustomBatchLogger):
|
|||
if response.status_code != 200:
|
||||
raise Exception(f"Failed to get OAuth2 token: {response.status_code} - {response.text}")
|
||||
|
||||
token_response: Final = response.json()
|
||||
token_response: Final = _oauth_token_body(response)
|
||||
self.oauth_token = token_response.get("access_token")
|
||||
expires_in: Final = token_response.get("expires_in", 3600)
|
||||
|
||||
|
|
|
|||
|
|
@ -5,7 +5,7 @@ BitBucket API client for fetching .prompt files from BitBucket repositories.
|
|||
import base64
|
||||
import urllib.parse
|
||||
from collections.abc import Mapping
|
||||
from typing import Final, TypedDict
|
||||
from typing import Final, Protocol, TypedDict
|
||||
|
||||
from typing_extensions import NotRequired, ReadOnly
|
||||
|
||||
|
|
@ -36,6 +36,30 @@ class BitBucketFileMetadata(TypedDict):
|
|||
last_modified: ReadOnly[str | None]
|
||||
|
||||
|
||||
class _SrcListingResponse(Protocol):
|
||||
def json(self) -> BitBucketSrcListing: ...
|
||||
|
||||
|
||||
class _RepositoryInfoResponse(Protocol):
|
||||
def json(self) -> Mapping[str, object]: ...
|
||||
|
||||
|
||||
class _BranchListingResponse(Protocol):
|
||||
def json(self) -> BitBucketBranchListing: ...
|
||||
|
||||
|
||||
def _src_listing(response: _SrcListingResponse) -> BitBucketSrcListing:
|
||||
return response.json()
|
||||
|
||||
|
||||
def _repository_info(response: _RepositoryInfoResponse) -> Mapping[str, object]:
|
||||
return response.json()
|
||||
|
||||
|
||||
def _branch_listing(response: _BranchListingResponse) -> BitBucketBranchListing:
|
||||
return response.json()
|
||||
|
||||
|
||||
def _sanitize_file_path(file_path: str) -> str:
|
||||
"""Reject path traversal and URL-encode each path segment."""
|
||||
if "#" in file_path or "?" in file_path:
|
||||
|
|
@ -166,7 +190,7 @@ class BitBucketClient:
|
|||
response: Final = self.http_handler.get(url, headers=self.headers)
|
||||
response.raise_for_status()
|
||||
|
||||
data: Final[BitBucketSrcListing] = response.json()
|
||||
data: Final = _src_listing(response)
|
||||
return [
|
||||
file_path
|
||||
for item in data.get("values", [])
|
||||
|
|
@ -201,7 +225,7 @@ class BitBucketClient:
|
|||
try:
|
||||
response: Final = self.http_handler.get(url, headers=self.headers)
|
||||
response.raise_for_status()
|
||||
return response.json()
|
||||
return _repository_info(response)
|
||||
except Exception as e:
|
||||
raise Exception(f"Failed to get repository info: {e}")
|
||||
|
||||
|
|
@ -231,7 +255,7 @@ class BitBucketClient:
|
|||
response: Final = self.http_handler.get(url, headers=self.headers)
|
||||
response.raise_for_status()
|
||||
|
||||
data: Final[BitBucketBranchListing] = response.json()
|
||||
data: Final = _branch_listing(response)
|
||||
return data.get("values", [])
|
||||
except Exception as e:
|
||||
raise Exception(f"Failed to get branches: {e}")
|
||||
|
|
|
|||
|
|
@ -4,12 +4,13 @@ Fetches .prompt files from BitBucket repositories and provides team-based access
|
|||
"""
|
||||
|
||||
from collections.abc import Mapping
|
||||
from typing import TYPE_CHECKING, Any, Final
|
||||
from typing import TYPE_CHECKING, Any, Final, Protocol
|
||||
|
||||
from jinja2 import DictLoader, select_autoescape
|
||||
from jinja2.sandbox import ImmutableSandboxedEnvironment
|
||||
|
||||
from litellm.integrations.custom_prompt_management import CustomPromptManagement
|
||||
from litellm.litellm_core_utils.core_helpers import require_attributes
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
|
||||
|
|
@ -26,6 +27,10 @@ from litellm.types.utils import StandardCallbackDynamicParams
|
|||
from .bitbucket_client import BitBucketClient
|
||||
|
||||
|
||||
class _SupportsGetWithDefault(Protocol):
|
||||
def get(self, key: str, default: object = None, /) -> object: ...
|
||||
|
||||
|
||||
class BitBucketPromptTemplate:
|
||||
"""
|
||||
Represents a prompt template loaded from BitBucket.
|
||||
|
|
@ -35,7 +40,7 @@ class BitBucketPromptTemplate:
|
|||
self,
|
||||
template_id: str,
|
||||
content: str,
|
||||
metadata: dict[str, Any],
|
||||
metadata: Mapping[str, object],
|
||||
model: str | None = None,
|
||||
):
|
||||
self.template_id = template_id
|
||||
|
|
@ -44,7 +49,9 @@ class BitBucketPromptTemplate:
|
|||
self.model = model or metadata.get("model")
|
||||
self.temperature = metadata.get("temperature")
|
||||
self.max_tokens = metadata.get("max_tokens")
|
||||
self.input_schema = metadata.get("input", {}).get("schema", {})
|
||||
self.input_schema = require_attributes(metadata.get("input", {}), _SupportsGetWithDefault, "get").get(
|
||||
"schema", {}
|
||||
)
|
||||
self.optional_params = {k: v for k, v in metadata.items() if k not in ["model", "input", "content"]}
|
||||
|
||||
def __repr__(self):
|
||||
|
|
|
|||
|
|
@ -7,10 +7,10 @@ so `response_id` is always the raw provider response id (cache-hit suffix stripp
|
|||
|
||||
import json
|
||||
import re
|
||||
from collections.abc import Iterator, Mapping
|
||||
from collections.abc import Iterable, Iterator, Mapping
|
||||
from math import isfinite
|
||||
from types import MappingProxyType
|
||||
from typing import Any, Final
|
||||
from typing import Any, Final, Protocol
|
||||
|
||||
from pydantic import JsonValue, TypeAdapter, ValidationError
|
||||
|
||||
|
|
@ -19,6 +19,7 @@ from litellm._logging import verbose_logger
|
|||
from litellm.integrations.clickhouse.clickhouse_batch_logger import ClickHouseBatchLogger
|
||||
from litellm.integrations.clickhouse.context import is_lens_analysis
|
||||
from litellm.integrations.clickhouse.schema import SPEND_LOGS_TABLE
|
||||
from litellm.litellm_core_utils.core_helpers import require_attributes
|
||||
from litellm.litellm_core_utils.internal_call_metadata import logged_event_key
|
||||
from litellm.litellm_core_utils.llm_response_utils.get_headers import get_provider_request_id
|
||||
from litellm.litellm_core_utils.sensitive_data_masker import redact_credentials_in_payload
|
||||
|
|
@ -40,6 +41,14 @@ _INTERNAL_METADATA_KEYS: Final = frozenset(
|
|||
)
|
||||
|
||||
|
||||
class _SupportsGet(Protocol):
|
||||
def get(self, key: str, /) -> object: ...
|
||||
|
||||
|
||||
class _SupportsItems(Protocol):
|
||||
def items(self) -> Iterable[tuple[object, object]]: ...
|
||||
|
||||
|
||||
def _request_metadata_fields(value: object) -> Iterator[tuple[str, JsonValue]]:
|
||||
if value is None:
|
||||
return
|
||||
|
|
@ -85,22 +94,22 @@ def _json_mapping(value: Mapping[str, object]) -> str:
|
|||
return _json(dict(value))
|
||||
|
||||
|
||||
def _find_traceparent(metadata: Mapping[str, Any], kwargs: Mapping[str, Any]) -> tuple[str, str]:
|
||||
def _find_traceparent(metadata: Mapping[str, object], kwargs: Mapping[str, object]) -> tuple[str, str]:
|
||||
custom_headers = metadata.get("requester_custom_headers") or MappingProxyType({})
|
||||
proxy_request = (kwargs.get("litellm_params") or MappingProxyType({})).get(
|
||||
proxy_request = require_attributes(kwargs.get("litellm_params") or MappingProxyType({}), _SupportsGet, "get").get(
|
||||
"proxy_server_request"
|
||||
) or MappingProxyType({})
|
||||
request_headers = proxy_request.get("headers") or MappingProxyType({})
|
||||
request_headers = require_attributes(proxy_request, _SupportsGet, "get").get("headers") or MappingProxyType({})
|
||||
for headers in (custom_headers, request_headers):
|
||||
for name, value in headers.items():
|
||||
for name, value in require_attributes(headers, _SupportsItems, "items").items():
|
||||
if str(name).lower() == "traceparent":
|
||||
return parse_traceparent(value)
|
||||
return "", ""
|
||||
|
||||
|
||||
def _cache_tokens(usage: Mapping[str, Any]) -> tuple[int, int]:
|
||||
def _cache_tokens(usage: Mapping[str, object]) -> tuple[int, int]:
|
||||
"""(cache_read, cache_write) from a Usage dict: OpenAI prompt_tokens_details first, Anthropic fields as fallback."""
|
||||
details = usage.get("prompt_tokens_details") or MappingProxyType({})
|
||||
details = require_attributes(usage.get("prompt_tokens_details") or MappingProxyType({}), _SupportsGet, "get")
|
||||
cache_read = _int(details.get("cached_tokens")) or _int(usage.get("cache_read_input_tokens"))
|
||||
cache_write = (
|
||||
_int(details.get("cache_write_tokens"))
|
||||
|
|
@ -116,10 +125,17 @@ def _request_tags(value: object) -> list[str]:
|
|||
return [str(tag) for tag in value]
|
||||
|
||||
|
||||
def _session_id(payload: StandardLoggingPayload | SpendLogPayload, kwargs: Mapping[str, Any]) -> str:
|
||||
def _session_id(payload: StandardLoggingPayload | SpendLogPayload, kwargs: Mapping[str, object]) -> str:
|
||||
"""Mirrors proxy `_get_session_id_for_spend_log`: explicit session id, else the payload trace id."""
|
||||
request_metadata = (kwargs.get("litellm_params") or MappingProxyType({})).get("metadata") or MappingProxyType({})
|
||||
return str(payload.get("session_id") or request_metadata.get("session_id") or payload.get("trace_id") or "")
|
||||
request_metadata = require_attributes(
|
||||
kwargs.get("litellm_params") or MappingProxyType({}), _SupportsGet, "get"
|
||||
).get("metadata") or MappingProxyType({})
|
||||
return str(
|
||||
payload.get("session_id")
|
||||
or require_attributes(request_metadata, _SupportsGet, "get").get("session_id")
|
||||
or payload.get("trace_id")
|
||||
or ""
|
||||
)
|
||||
|
||||
|
||||
def _is_trace_ingest(payload: StandardLoggingPayload) -> bool:
|
||||
|
|
@ -128,7 +144,7 @@ def _is_trace_ingest(payload: StandardLoggingPayload) -> bool:
|
|||
|
||||
|
||||
def spend_log_row_from_payload(
|
||||
payload: StandardLoggingPayload | SpendLogPayload, kwargs: Mapping[str, Any]
|
||||
payload: StandardLoggingPayload | SpendLogPayload, kwargs: Mapping[str, object]
|
||||
) -> SpendLogRecord:
|
||||
metadata: Mapping[str, Any] = payload.get("metadata") or MappingProxyType({})
|
||||
hidden_params: Mapping[str, Any] = payload.get("hidden_params") or MappingProxyType({})
|
||||
|
|
|
|||
|
|
@ -1,16 +1,17 @@
|
|||
import contextvars
|
||||
import copy
|
||||
import hashlib
|
||||
import json
|
||||
import os
|
||||
import secrets
|
||||
from collections.abc import Iterable, Mapping, Sequence
|
||||
from datetime import datetime
|
||||
from itertools import chain
|
||||
from types import MappingProxyType
|
||||
from typing import TYPE_CHECKING, Any, ClassVar, Final, Literal, Optional, cast, get_args
|
||||
from typing import TYPE_CHECKING, Any, ClassVar, Final, Literal, Optional, Protocol, cast, get_args
|
||||
|
||||
import httpx
|
||||
from pydantic import TypeAdapter
|
||||
from pydantic import BaseModel, TypeAdapter
|
||||
|
||||
from litellm._internal_context import with_service_target
|
||||
from litellm._logging import verbose_logger
|
||||
|
|
@ -20,6 +21,7 @@ from litellm.litellm_core_utils.core_helpers import (
|
|||
get_metadata_variable_name_from_kwargs,
|
||||
get_or_create_metadata_bucket,
|
||||
redact_nested_match_and_regex_keys,
|
||||
require_attributes,
|
||||
)
|
||||
from litellm.secret_managers.main import str_to_bool
|
||||
from litellm.types.guardrails import (
|
||||
|
|
@ -73,6 +75,7 @@ GUARDRAIL_SESSIONS_TARGET: Final = "guardrail_sessions"
|
|||
_PRE_CALL_EXECUTED_TOKEN: Final = secrets.token_hex(16)
|
||||
_PRE_CALL_RAISED_TOKEN: Final = secrets.token_hex(16)
|
||||
_DEPLOYMENT_HOOK_RAN_TOKEN: Final = secrets.token_hex(16)
|
||||
_POST_CALL_CHECKED_TOKEN: Final = secrets.token_hex(16)
|
||||
|
||||
_GUARDRAIL_BLOCK_STATUS_CODES: Final = frozenset({400, 403, 422})
|
||||
|
||||
|
|
@ -167,18 +170,22 @@ def _strict_guardrail_modes_enabled() -> bool:
|
|||
return True if parsed is None else parsed
|
||||
|
||||
|
||||
def get_session_id_from_request_data(request_data: dict[str, Any]) -> str | None:
|
||||
class _SupportsGet(Protocol):
|
||||
def get(self, key: str, /) -> object: ...
|
||||
|
||||
|
||||
def get_session_id_from_request_data(request_data: Mapping[str, object]) -> str | None:
|
||||
"""Extract session_id from request data (litellm_session_id or metadata)."""
|
||||
session_id = request_data.get("litellm_session_id")
|
||||
if session_id:
|
||||
return str(session_id)
|
||||
|
||||
metadata: Final = request_data.get("metadata") or {}
|
||||
metadata: Final = require_attributes(request_data.get("metadata") or {}, _SupportsGet, "get")
|
||||
session_id = metadata.get("session_id")
|
||||
if session_id:
|
||||
return str(session_id)
|
||||
|
||||
litellm_metadata: Final = request_data.get("litellm_metadata") or {}
|
||||
litellm_metadata: Final = require_attributes(request_data.get("litellm_metadata") or {}, _SupportsGet, "get")
|
||||
session_id = litellm_metadata.get("session_id")
|
||||
if session_id:
|
||||
return str(session_id)
|
||||
|
|
@ -281,6 +288,95 @@ def _request_carries_marker(data: Mapping[str, object], marker: str | None) -> b
|
|||
return False
|
||||
|
||||
|
||||
def _record_guardrail_marker(
|
||||
data: dict[str, object], # mutable-ok: the request payload the marker is recorded on
|
||||
marker: str | None,
|
||||
) -> None:
|
||||
if marker is None:
|
||||
return
|
||||
for meta_key in ("metadata", "litellm_metadata"):
|
||||
meta = data.get(meta_key)
|
||||
if isinstance(meta, dict):
|
||||
executed = meta.get(PRE_CALL_EXECUTED_GUARDRAILS_KEY)
|
||||
if isinstance(executed, list):
|
||||
if marker not in executed:
|
||||
executed.append(marker)
|
||||
else:
|
||||
meta[PRE_CALL_EXECUTED_GUARDRAILS_KEY] = [marker]
|
||||
return
|
||||
data["metadata"] = {PRE_CALL_EXECUTED_GUARDRAILS_KEY: [marker]}
|
||||
|
||||
|
||||
_RESPONSE_METADATA_KEYS: Final = frozenset(
|
||||
{
|
||||
"id",
|
||||
"created",
|
||||
"created_at",
|
||||
"model",
|
||||
"usage",
|
||||
"usage_info",
|
||||
"system_fingerprint",
|
||||
"_hidden_params",
|
||||
"_response_ms",
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
def _response_fingerprint(response: object) -> str | None:
|
||||
"""Digest of everything but metadata set between layers, so an unknown field counts as content."""
|
||||
try:
|
||||
fields: Final[Mapping[str, object]] = (
|
||||
response.model_dump(mode="json")
|
||||
if isinstance(response, BaseModel)
|
||||
else _REQUEST_FIELDS.validate_python(response)
|
||||
)
|
||||
content: Final = {key: value for key, value in fields.items() if key not in _RESPONSE_METADATA_KEYS}
|
||||
if not content:
|
||||
return None
|
||||
return hashlib.sha256(
|
||||
json.dumps(content, sort_keys=True, separators=(",", ":"), allow_nan=False).encode()
|
||||
).hexdigest()
|
||||
except (TypeError, ValueError):
|
||||
return None
|
||||
|
||||
|
||||
_RESPONSE_GUARDRAIL_CHECKS_KEY: Final = "guardrail_checks"
|
||||
|
||||
|
||||
def _str_keyed_dict(value: object) -> Mapping[str, object] | None:
|
||||
if not isinstance(value, dict):
|
||||
return None
|
||||
return cast("Mapping[str, object]", value) # cast-ok: response bodies and _hidden_params are keyed by field name
|
||||
|
||||
|
||||
def _response_hidden_params(response: object) -> Mapping[str, object] | None:
|
||||
as_dict: Final = _str_keyed_dict(response)
|
||||
return _str_keyed_dict(
|
||||
as_dict.get("_hidden_params") if as_dict is not None else getattr(response, "_hidden_params", None)
|
||||
)
|
||||
|
||||
|
||||
def _write_response_marker(response: object, marker: str | None) -> None:
|
||||
"""Record a check on the response itself, so it survives the router copying request metadata per hop."""
|
||||
if marker is None:
|
||||
return
|
||||
hidden: Final = _response_hidden_params(response)
|
||||
if not isinstance(hidden, dict):
|
||||
return
|
||||
hidden[_RESPONSE_GUARDRAIL_CHECKS_KEY] = tuple(
|
||||
_str_collection(hidden.get(_RESPONSE_GUARDRAIL_CHECKS_KEY)) | {marker}
|
||||
)
|
||||
|
||||
|
||||
def _response_carries_marker(response: object, marker: str | None) -> bool:
|
||||
hidden: Final = _response_hidden_params(response)
|
||||
return (
|
||||
marker is not None
|
||||
and hidden is not None
|
||||
and marker in _str_collection(hidden.get(_RESPONSE_GUARDRAIL_CHECKS_KEY))
|
||||
)
|
||||
|
||||
|
||||
_REQUEST_FIELDS: Final = TypeAdapter(Mapping[str, object])
|
||||
|
||||
|
||||
|
|
@ -291,6 +387,7 @@ def _guardrails_listed_in(metadata: object) -> list[str | dict[str, DynamicGuard
|
|||
|
||||
|
||||
class CustomGuardrail(CustomLogger):
|
||||
mcp_block_behavior: Literal["tool_result", "raise"] | None = None
|
||||
# If True, during_call runs async_moderation_hook instead of the unified apply_guardrail path.
|
||||
use_native_during_call_hook: ClassVar[bool] = False
|
||||
|
||||
|
|
@ -912,20 +1009,7 @@ class CustomGuardrail(CustomLogger):
|
|||
With ``raised``, records instead that the hook raised as a policy pipeline
|
||||
step the pipeline then moved past. The deployment-level hook still runs it.
|
||||
"""
|
||||
marker: Final = self._pre_call_marker(raised=raised)
|
||||
if marker is None:
|
||||
return
|
||||
for meta_key in ("metadata", "litellm_metadata"):
|
||||
meta = data.get(meta_key)
|
||||
if isinstance(meta, dict):
|
||||
executed = meta.get(PRE_CALL_EXECUTED_GUARDRAILS_KEY)
|
||||
if isinstance(executed, list):
|
||||
if marker not in executed:
|
||||
executed.append(marker)
|
||||
else:
|
||||
meta[PRE_CALL_EXECUTED_GUARDRAILS_KEY] = [marker]
|
||||
return
|
||||
data["metadata"] = {PRE_CALL_EXECUTED_GUARDRAILS_KEY: [marker]}
|
||||
_record_guardrail_marker(data, self._pre_call_marker(raised=raised))
|
||||
|
||||
def pre_call_hook_ran_or_raised(self, data: Mapping[str, object]) -> bool:
|
||||
"""Whether this guardrail's ``async_pre_call_hook`` ran for this request, whatever it decided."""
|
||||
|
|
@ -1123,10 +1207,26 @@ class CustomGuardrail(CustomLogger):
|
|||
cast("dict[str, object]", request_data) # cast-ok: proxy request payloads are str-keyed
|
||||
)
|
||||
|
||||
if not self._is_valid_response_type(result):
|
||||
return None
|
||||
returned: Final = result if self._is_valid_response_type(result) else None
|
||||
if target is self:
|
||||
return returned
|
||||
checked: Final = returned if returned is not None else response
|
||||
marker: Final = self._post_call_check_marker(checked)
|
||||
_record_guardrail_marker(
|
||||
cast("dict[str, object]", request_data), # cast-ok: proxy request payloads are str-keyed
|
||||
marker,
|
||||
)
|
||||
_write_response_marker(checked, marker)
|
||||
return returned
|
||||
|
||||
return result
|
||||
def _post_call_check_marker(self, response: object) -> str | None:
|
||||
fingerprint: Final = _response_fingerprint(response)
|
||||
return self._marker(f"{_POST_CALL_CHECKED_TOKEN}:{fingerprint}") if fingerprint is not None else None
|
||||
|
||||
def post_call_already_checked(self, data: Mapping[str, object], response: object) -> bool:
|
||||
"""Whether a deployment-level hook of this request already ran this guardrail on this exact response."""
|
||||
marker: Final = self._post_call_check_marker(response)
|
||||
return _request_carries_marker(data, marker) or _response_carries_marker(response, marker)
|
||||
|
||||
async def async_logging_hook(
|
||||
self,
|
||||
|
|
|
|||
|
|
@ -44,6 +44,7 @@ from litellm.litellm_core_utils.prompt_templates.common_utils import (
|
|||
from litellm.litellm_core_utils.redact_messages import should_redact_message_logging
|
||||
from litellm.litellm_core_utils.safe_json_dumps import safe_dumps
|
||||
from litellm.litellm_core_utils.safe_json_loads import safe_json_loads
|
||||
from litellm.llms.anthropic.common_utils import anthropic_server_tool_call_params
|
||||
from litellm.llms.custom_httpx.http_handler import (
|
||||
get_async_httpx_client,
|
||||
httpxSpecialProvider,
|
||||
|
|
@ -315,7 +316,10 @@ def _to_dd_tool_calls(message: Mapping[str, object]) -> tuple[ToolCall, ...]:
|
|||
tool_id=_text_field(tool_call, "id"),
|
||||
type=_text_field(tool_call, "type", "function"),
|
||||
)
|
||||
for tool_call in (raw_tool_calls if isinstance(raw_tool_calls, list) else ())
|
||||
for tool_call in (
|
||||
*anthropic_server_tool_call_params(message),
|
||||
*(raw_tool_calls if isinstance(raw_tool_calls, list) else ()),
|
||||
)
|
||||
if isinstance(tool_call, dict)
|
||||
for function in [_mapping_field(tool_call, "function")]
|
||||
)
|
||||
|
|
|
|||
|
|
@ -3,17 +3,29 @@ Functions for sending Email Alerts
|
|||
"""
|
||||
|
||||
import os
|
||||
from typing import Final
|
||||
from collections.abc import Awaitable, Mapping, Sequence
|
||||
from typing import TYPE_CHECKING, Final
|
||||
|
||||
from litellm._logging import verbose_logger, verbose_proxy_logger
|
||||
from litellm.repositories.team_repository import TeamRepository
|
||||
from litellm.types.integrations.slack_alerting import WebhookEvent
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from litellm.proxy.utils import PrismaClient
|
||||
|
||||
# we use this for the email header, please send a test email if you change this. verify it looks good on email
|
||||
LITELLM_LOGO_URL: Final = "https://litellm-listing.s3.amazonaws.com/litellm_logo.png"
|
||||
LITELLM_SUPPORT_CONTACT: Final = "support@berri.ai"
|
||||
|
||||
|
||||
def _query_raw_rows(
|
||||
client: "PrismaClient", query: str, *args: object
|
||||
) -> Awaitable[Sequence[Mapping[str, object]] | None]:
|
||||
from litellm.proxy.utils import query_raw_rows
|
||||
|
||||
return query_raw_rows(client.db, query, *args)
|
||||
|
||||
|
||||
async def get_all_team_member_emails(team_id: str | None = None) -> list:
|
||||
verbose_logger.debug("Email Alerting: Getting all team members for team_id=%s", team_id)
|
||||
if team_id is None:
|
||||
|
|
@ -51,17 +63,18 @@ async def get_all_team_member_emails(team_id: str | None = None) -> list:
|
|||
WHERE user_id = ANY($1::TEXT[]);
|
||||
"""
|
||||
|
||||
_result: Final = await prisma_client.db.query_raw(sql_query, _team_member_user_ids)
|
||||
_result: Final = await _query_raw_rows(prisma_client, sql_query, _team_member_user_ids)
|
||||
|
||||
verbose_logger.debug("Email Alerting: Got all Emails for team, emails=%s", _result)
|
||||
|
||||
if _result is None:
|
||||
return []
|
||||
|
||||
emails: Final = []
|
||||
for user in _result:
|
||||
if user and isinstance(user, dict) and user.get("user_email", None) is not None:
|
||||
emails.append(user.get("user_email"))
|
||||
emails: Final = [
|
||||
user.get("user_email")
|
||||
for user in _result
|
||||
if user and isinstance(user, dict) and user.get("user_email", None) is not None
|
||||
]
|
||||
return emails
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -123,7 +123,9 @@ class FocusLiteLLMDatabase:
|
|||
ORDER BY ordinal_position;
|
||||
"""
|
||||
try:
|
||||
columns_response: Final = await client.db.query_raw(info_query)
|
||||
from litellm.proxy.utils import query_raw_rows
|
||||
|
||||
columns_response: Final = await query_raw_rows(client.db, info_query)
|
||||
return {"columns": columns_response, "table_name": "LiteLLM_DailyUserSpend"}
|
||||
except Exception as exc:
|
||||
raise RuntimeError(f"Error getting table info: {exc}") from exc
|
||||
|
|
|
|||
|
|
@ -15,6 +15,7 @@ from typing_extensions import ReadOnly, TypedDict
|
|||
import litellm
|
||||
from litellm._logging import verbose_logger
|
||||
from litellm.integrations.custom_logger import CustomLogger
|
||||
from litellm.litellm_core_utils.core_helpers import require_attributes, subscript
|
||||
from litellm.litellm_core_utils.prompt_templates.common_utils import (
|
||||
convert_content_list_to_str,
|
||||
get_content_from_model_response,
|
||||
|
|
@ -60,6 +61,10 @@ class _JsonResponse(Protocol):
|
|||
def json(self) -> object: ...
|
||||
|
||||
|
||||
class _SupportsGet(Protocol):
|
||||
def get(self, key: str, /) -> object: ...
|
||||
|
||||
|
||||
def _login_access_token(response: _GalileoLoginResponse) -> str:
|
||||
"""Read the bearer token out of a Galileo login response body."""
|
||||
return response.json()["access_token"]
|
||||
|
|
@ -463,13 +468,13 @@ class GalileoObserve(CustomLogger):
|
|||
pass
|
||||
|
||||
@staticmethod
|
||||
def _build_prompt(kwargs: Mapping[str, Any]) -> dict[str, object]:
|
||||
optional_params: Final[Mapping[str, object]] = kwargs.get("optional_params", {}) or {}
|
||||
def _build_prompt(kwargs: Mapping[str, object]) -> dict[str, object]:
|
||||
optional_params: Final = require_attributes(kwargs.get("optional_params", {}) or {}, _SupportsGet, "get")
|
||||
prompt: Final[dict[str, object]] = {"messages": kwargs.get("messages")}
|
||||
if optional_params.get("functions") is not None:
|
||||
prompt["functions"] = optional_params["functions"]
|
||||
prompt["functions"] = subscript(optional_params, "functions")
|
||||
if optional_params.get("tools") is not None:
|
||||
prompt["tools"] = optional_params["tools"]
|
||||
prompt["tools"] = subscript(optional_params, "tools")
|
||||
return prompt
|
||||
|
||||
@staticmethod
|
||||
|
|
|
|||
|
|
@ -5,7 +5,7 @@ import os
|
|||
import time
|
||||
from collections.abc import Mapping
|
||||
from datetime import datetime, timedelta, timezone
|
||||
from typing import TYPE_CHECKING, Any, Final
|
||||
from typing import TYPE_CHECKING, Any, Final, Protocol
|
||||
|
||||
from litellm._logging import verbose_logger
|
||||
from litellm._uuid import uuid
|
||||
|
|
@ -16,7 +16,7 @@ from litellm.integrations.gcs_bucket.gcs_bucket_base import GCSBucketBase
|
|||
from litellm.litellm_core_utils.cloud_storage_security import (
|
||||
sanitize_cloud_object_component,
|
||||
)
|
||||
from litellm.litellm_core_utils.core_helpers import as_str_mapping
|
||||
from litellm.litellm_core_utils.core_helpers import as_str_mapping, require_attributes
|
||||
from litellm.litellm_core_utils.internal_call_metadata import logged_event_key
|
||||
from litellm.types.integrations.base_health_check import IntegrationHealthCheckStatus
|
||||
from litellm.types.integrations.gcs_bucket import *
|
||||
|
|
@ -28,6 +28,10 @@ else:
|
|||
VertexBase = Any
|
||||
|
||||
|
||||
class _SupportsGetWithDefault(Protocol):
|
||||
def get(self, key: str, default: str, /) -> object: ...
|
||||
|
||||
|
||||
class GCSBucketLogger(GCSBucketBase, AdditionalLoggingUtils):
|
||||
def __init__(self, bucket_name: str | None = None) -> None:
|
||||
from litellm.proxy.proxy_server import premium_user
|
||||
|
|
@ -298,7 +302,7 @@ class GCSBucketLogger(GCSBucketBase, AdditionalLoggingUtils):
|
|||
await self._send_individual_logs(items_to_process)
|
||||
|
||||
def _get_object_name(
|
||||
self, kwargs: Mapping[str, object], logging_payload: StandardLoggingPayload, response_obj: Any
|
||||
self, kwargs: Mapping[str, object], logging_payload: StandardLoggingPayload, response_obj: object
|
||||
) -> str:
|
||||
"""
|
||||
Get the object name to use for the current payload
|
||||
|
|
@ -313,7 +317,7 @@ class GCSBucketLogger(GCSBucketBase, AdditionalLoggingUtils):
|
|||
return f"{current_date}/custom-{uuid.uuid4().hex}-{safe_log_id}"
|
||||
if logging_payload.get("error_str", None) is not None:
|
||||
return self._generate_failure_object_name(request_date_str=current_date)
|
||||
raw_response_id: Final[object] = response_obj.get("id", "")
|
||||
raw_response_id: Final[object] = require_attributes(response_obj, _SupportsGetWithDefault, "get").get("id", "")
|
||||
response_id: Final = raw_response_id if isinstance(raw_response_id, str) else ""
|
||||
return self._generate_success_object_name(
|
||||
request_date_str=current_date,
|
||||
|
|
@ -360,7 +364,7 @@ class GCSBucketLogger(GCSBucketBase, AdditionalLoggingUtils):
|
|||
def _generate_success_object_name(
|
||||
self,
|
||||
request_date_str: str,
|
||||
response_id: str,
|
||||
response_id: object,
|
||||
) -> str:
|
||||
return f"{request_date_str}/{response_id}"
|
||||
|
||||
|
|
|
|||
|
|
@ -132,7 +132,7 @@ class GenericPromptManager(CustomPromptManagement):
|
|||
|
||||
async def async_fetch_prompt_from_api(
|
||||
self, prompt_id: str | None, prompt_spec: PromptSpec | None
|
||||
) -> dict[str, Any]:
|
||||
) -> dict[str, object]:
|
||||
"""
|
||||
Fetch a prompt from the API asynchronously.
|
||||
"""
|
||||
|
|
|
|||
|
|
@ -10,6 +10,23 @@ from litellm.integrations.helicone_mock_client import (
|
|||
create_mock_helicone_client,
|
||||
should_use_helicone_mock,
|
||||
)
|
||||
from litellm.llms.anthropic.common_utils import tool_calls_with_server_tool_calls
|
||||
from litellm.types.utils import ChatCompletionMessageCustomToolCall, ChatCompletionMessageToolCall
|
||||
|
||||
|
||||
def _tool_use_block(
|
||||
tool_call: ChatCompletionMessageToolCall | ChatCompletionMessageCustomToolCall,
|
||||
) -> dict[str, object]:
|
||||
match tool_call:
|
||||
case ChatCompletionMessageCustomToolCall(custom=custom):
|
||||
return {"type": "tool_use", "id": tool_call.id, "name": custom.name, "input": custom.input}
|
||||
case _:
|
||||
return {
|
||||
"type": "tool_use",
|
||||
"id": tool_call.id,
|
||||
"name": tool_call.function.name,
|
||||
"input": tool_call.function.arguments,
|
||||
}
|
||||
|
||||
|
||||
class HeliconeLogger:
|
||||
|
|
@ -57,23 +74,8 @@ class HeliconeLogger:
|
|||
message = choice["message"]
|
||||
|
||||
content = []
|
||||
if "tool_calls" in message and message["tool_calls"]:
|
||||
for tool_call in message["tool_calls"]:
|
||||
function = tool_call.get("function")
|
||||
custom = tool_call.get("custom")
|
||||
if not function and not custom:
|
||||
continue
|
||||
name, tool_input = (
|
||||
(function["name"], function["arguments"]) if function else (custom["name"], custom["input"])
|
||||
)
|
||||
content.append(
|
||||
{
|
||||
"type": "tool_use",
|
||||
"id": tool_call["id"],
|
||||
"name": name,
|
||||
"input": tool_input,
|
||||
}
|
||||
)
|
||||
if tool_calls := tool_calls_with_server_tool_calls(message):
|
||||
content = [_tool_use_block(tool_call) for tool_call in tool_calls]
|
||||
elif "content" in message and message["content"]:
|
||||
content = [{"type": "text", "text": message["content"]}]
|
||||
|
||||
|
|
|
|||
|
|
@ -7,10 +7,12 @@ https://humanloop.com/
|
|||
from typing import Final, cast
|
||||
|
||||
import httpx
|
||||
from pydantic import TypeAdapter
|
||||
from typing_extensions import TypedDict
|
||||
|
||||
import litellm
|
||||
from litellm.caching import DualCache
|
||||
from litellm.litellm_core_utils.core_helpers import as_object_list, as_str_mapping, decoded_json_object, subscript
|
||||
from litellm.llms.custom_httpx.http_handler import get_httpx_client
|
||||
from litellm.secret_managers.main import get_secret_str
|
||||
from litellm.types.llms.openai import AllMessageValues
|
||||
|
|
@ -19,6 +21,8 @@ from litellm.types.utils import StandardCallbackDynamicParams
|
|||
|
||||
from .custom_logger import CustomLogger
|
||||
|
||||
_DECODED_JSON: Final = TypeAdapter(object)
|
||||
|
||||
|
||||
class PromptManagementClient(TypedDict):
|
||||
prompt_id: str
|
||||
|
|
@ -79,18 +83,20 @@ class HumanLoopPromptManager(DualCache):
|
|||
raise Exception(f"Error getting prompt from Humanloop: {e.response.text}")
|
||||
|
||||
json_response: Final = response.json()
|
||||
template_message: Final = json_response["template"]
|
||||
if isinstance(template_message, dict):
|
||||
template_messages = [template_message]
|
||||
elif isinstance(template_message, list):
|
||||
template_messages = template_message
|
||||
decoded_response: Final = _DECODED_JSON.validate_python(json_response)
|
||||
template_message: Final = subscript(decoded_response, "template")
|
||||
template_mapping: Final = as_str_mapping(template_message)
|
||||
template_list: Final = as_object_list(template_message)
|
||||
if template_mapping is not None:
|
||||
template_messages = [template_mapping]
|
||||
elif template_list is not None:
|
||||
template_messages = template_list
|
||||
else:
|
||||
raise ValueError(f"Invalid template message type: {type(template_message)}")
|
||||
template_model: Final = json_response["model"]
|
||||
optional_params: Final = {}
|
||||
for k, v in json_response.items():
|
||||
if k in litellm.OPENAI_CHAT_COMPLETION_PARAMS:
|
||||
optional_params[k] = v
|
||||
optional_params: Final = {
|
||||
k: v for k, v in decoded_json_object(decoded_response).items() if k in litellm.OPENAI_CHAT_COMPLETION_PARAMS
|
||||
}
|
||||
return PromptManagementClient(
|
||||
prompt_id=humanloop_prompt_id,
|
||||
prompt_template=cast(list[AllMessageValues], template_messages),
|
||||
|
|
|
|||
|
|
@ -3,7 +3,7 @@
|
|||
|
||||
import json
|
||||
import os
|
||||
from typing import Final, Literal
|
||||
from typing import Final, Literal, Protocol
|
||||
|
||||
import httpx
|
||||
|
||||
|
|
@ -11,6 +11,7 @@ import litellm
|
|||
from litellm._logging import verbose_logger
|
||||
from litellm._uuid import uuid
|
||||
from litellm.integrations.custom_logger import CustomLogger
|
||||
from litellm.litellm_core_utils.core_helpers import require_attributes
|
||||
from litellm.llms.custom_httpx.http_handler import (
|
||||
HTTPHandler,
|
||||
get_async_httpx_client,
|
||||
|
|
@ -18,6 +19,10 @@ from litellm.llms.custom_httpx.http_handler import (
|
|||
)
|
||||
|
||||
|
||||
class _SupportsGet(Protocol):
|
||||
def get(self, key: str, default: object = None, /) -> object: ...
|
||||
|
||||
|
||||
def get_utc_datetime():
|
||||
import datetime as dt
|
||||
from datetime import datetime
|
||||
|
|
@ -70,10 +75,11 @@ class LagoLogger(CustomLogger):
|
|||
if (
|
||||
isinstance(response_obj, litellm.ModelResponse) or isinstance(response_obj, litellm.EmbeddingResponse)
|
||||
) and hasattr(response_obj, "usage"):
|
||||
response_usage: Final = require_attributes(response_obj["usage"], _SupportsGet, "get")
|
||||
usage = {
|
||||
"prompt_tokens": response_obj["usage"].get("prompt_tokens", 0),
|
||||
"completion_tokens": response_obj["usage"].get("completion_tokens", 0),
|
||||
"total_tokens": response_obj["usage"].get("total_tokens"),
|
||||
"prompt_tokens": response_usage.get("prompt_tokens", 0),
|
||||
"completion_tokens": response_usage.get("completion_tokens", 0),
|
||||
"total_tokens": response_usage.get("total_tokens"),
|
||||
}
|
||||
|
||||
litellm_params: Final = kwargs.get("litellm_params", {}) or {}
|
||||
|
|
|
|||
|
|
@ -7,6 +7,7 @@ from itertools import chain
|
|||
from typing import TYPE_CHECKING, Any, Final, Optional, Protocol, runtime_checkable
|
||||
|
||||
from opentelemetry import trace
|
||||
from pydantic import TypeAdapter
|
||||
|
||||
import litellm
|
||||
from litellm._logging import verbose_logger
|
||||
|
|
@ -26,12 +27,15 @@ from litellm.integrations.opentelemetry import (
|
|||
from litellm.litellm_core_utils.core_helpers import as_str_mapping
|
||||
from litellm.litellm_core_utils.safe_json_dumps import safe_dumps
|
||||
from litellm.litellm_core_utils.safe_json_loads import safe_json_loads
|
||||
from litellm.llms.anthropic.common_utils import tool_calls_with_server_tool_calls
|
||||
from litellm.secret_managers.main import str_to_bool
|
||||
from litellm.types.integrations.langfuse_otel import (
|
||||
LangfuseSpanAttributes,
|
||||
)
|
||||
from litellm.types.utils import StandardCallbackDynamicParams
|
||||
|
||||
_DECODED_JSON: Final = TypeAdapter(object)
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from opentelemetry.trace import Span as _Span
|
||||
|
||||
|
|
@ -605,8 +609,8 @@ def _extract_choices_output(response_obj: _Gettable) -> str | None:
|
|||
message: Final = _as_gettable(first_choice.get("message", {}))
|
||||
if message is None:
|
||||
return None
|
||||
tool_calls: Final = message.get("tool_calls")
|
||||
if isinstance(tool_calls, list) and tool_calls:
|
||||
tool_calls: Final = tool_calls_with_server_tool_calls(message)
|
||||
if tool_calls:
|
||||
transformed_tool_calls: Final = [
|
||||
entry for tool_call in tool_calls if (entry := _transformed_tool_call(response_obj, tool_call)) is not None
|
||||
]
|
||||
|
|
@ -642,7 +646,7 @@ def _tool_call_arguments(arguments: object) -> object:
|
|||
if not isinstance(arguments, str):
|
||||
return arguments
|
||||
try:
|
||||
return json.loads(arguments)
|
||||
return _DECODED_JSON.validate_python(json.loads(arguments))
|
||||
except json.JSONDecodeError:
|
||||
return {}
|
||||
|
||||
|
|
|
|||
|
|
@ -7,7 +7,7 @@ Relevant Issue: https://github.com/BerriAI/litellm/issues/13764
|
|||
import json
|
||||
from collections.abc import Mapping
|
||||
from datetime import datetime
|
||||
from typing import TYPE_CHECKING, Any, Final
|
||||
from typing import TYPE_CHECKING, Final, Protocol
|
||||
|
||||
from pydantic import BaseModel
|
||||
from typing_extensions import override
|
||||
|
|
@ -16,6 +16,7 @@ from litellm.integrations.opentelemetry_utils.base_otel_llm_obs_attributes impor
|
|||
BaseLLMObsOTELAttributes,
|
||||
safe_set_attribute,
|
||||
)
|
||||
from litellm.litellm_core_utils.core_helpers import require_attributes
|
||||
from litellm.types.llms.openai import HttpxBinaryResponseContent, ResponsesAPIResponse
|
||||
from litellm.types.utils import (
|
||||
EmbeddingResponse,
|
||||
|
|
@ -30,6 +31,10 @@ if TYPE_CHECKING:
|
|||
from opentelemetry.trace import Span
|
||||
|
||||
|
||||
class _SupportsGet(Protocol):
|
||||
def get(self, key: str, /) -> object: ...
|
||||
|
||||
|
||||
def serialize_langfuse_datetime(value: datetime) -> str:
|
||||
return json.dumps(value.astimezone().isoformat())
|
||||
|
||||
|
|
@ -87,11 +92,11 @@ def get_output_content_by_type(
|
|||
class LangfuseLLMObsOTELAttributes(BaseLLMObsOTELAttributes):
|
||||
@staticmethod
|
||||
@override
|
||||
def set_messages(span: "Span", kwargs: dict[str, Any]):
|
||||
def set_messages(span: "Span", kwargs: Mapping[str, object]):
|
||||
prompt: Final = {"messages": kwargs.get("messages")}
|
||||
optional_params: Final = kwargs.get("optional_params", {})
|
||||
functions: Final = optional_params.get("functions")
|
||||
tools: Final = optional_params.get("tools")
|
||||
functions: Final = require_attributes(optional_params, _SupportsGet, "get").get("functions")
|
||||
tools: Final = require_attributes(optional_params, _SupportsGet, "get").get("tools")
|
||||
if functions is not None:
|
||||
prompt["functions"] = functions
|
||||
if tools is not None:
|
||||
|
|
|
|||
|
|
@ -6,6 +6,7 @@ from datetime import datetime, timezone
|
|||
from typing import TYPE_CHECKING, Final
|
||||
|
||||
import packaging
|
||||
from litellm.llms.anthropic.common_utils import tool_calls_with_server_tool_calls
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from openai.types.chat.chat_completion_message_custom_tool_call import Custom
|
||||
|
|
@ -67,8 +68,8 @@ def parse_messages(input):
|
|||
}
|
||||
|
||||
# Only add tool_calls and function_call to res if they are set
|
||||
if message.get("tool_calls"):
|
||||
serialized["tool_calls"] = parse_tool_calls(message.get("tool_calls"))
|
||||
if tool_calls := tool_calls_with_server_tool_calls(message):
|
||||
serialized["tool_calls"] = parse_tool_calls(tool_calls)
|
||||
|
||||
return serialized
|
||||
|
||||
|
|
|
|||
|
|
@ -1,9 +1,11 @@
|
|||
import json
|
||||
import threading
|
||||
from typing import TYPE_CHECKING, Any, Final
|
||||
from collections.abc import Sequence
|
||||
from typing import TYPE_CHECKING, Any, Final, Protocol
|
||||
|
||||
from litellm._logging import verbose_logger
|
||||
from litellm.integrations.custom_logger import CustomLogger
|
||||
from litellm.litellm_core_utils.core_helpers import attribute_or_none, iter_items, require_attributes
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from litellm.types.utils import StandardLoggingPayload
|
||||
|
|
@ -11,6 +13,25 @@ else:
|
|||
StandardLoggingPayload = Any
|
||||
|
||||
|
||||
class _ChoiceWithMessage(Protocol):
|
||||
@property
|
||||
def message(self) -> object: ...
|
||||
|
||||
|
||||
class _ExcludeNoneDumpable(Protocol):
|
||||
def model_dump(self, *, exclude_none: bool) -> object: ...
|
||||
|
||||
|
||||
def _dumped_choice_message(choice: object) -> object:
|
||||
message: Final = require_attributes(choice, _ChoiceWithMessage, "message").message
|
||||
return require_attributes(message, _ExcludeNoneDumpable, "model_dump").model_dump(exclude_none=True)
|
||||
|
||||
|
||||
def _output_messages(response_obj: object) -> Sequence[object]:
|
||||
choices: Final = attribute_or_none(response_obj, "choices") if hasattr(response_obj, "choices") else ()
|
||||
return [_dumped_choice_message(c) for c in iter_items(choices)]
|
||||
|
||||
|
||||
class MlflowLogger(CustomLogger):
|
||||
def __init__(self):
|
||||
from mlflow.tracking import MlflowClient
|
||||
|
|
@ -62,7 +83,7 @@ class MlflowLogger(CustomLogger):
|
|||
|
||||
inputs: Final = self._construct_input(kwargs)
|
||||
input_messages: Final = inputs.get("messages", [])
|
||||
output_messages: Final = [c.message.model_dump(exclude_none=True) for c in getattr(response_obj, "choices", [])]
|
||||
output_messages: Final = _output_messages(response_obj)
|
||||
if messages := [*input_messages, *output_messages]:
|
||||
set_span_chat_messages(span, messages)
|
||||
if tools := inputs.get("tools"):
|
||||
|
|
|
|||
|
|
@ -56,6 +56,7 @@ from litellm._logging import verbose_logger
|
|||
from litellm.integrations.custom_logger import CustomLogger
|
||||
from litellm.litellm_core_utils.core_helpers import has_attributes
|
||||
from litellm.litellm_core_utils.redact_messages import should_redact_message_logging
|
||||
from litellm.llms.anthropic.common_utils import tool_calls_with_server_tool_calls
|
||||
from litellm.types.integrations.base_health_check import IntegrationHealthCheckStatus
|
||||
from litellm.types.integrations.newrelic import NewRelicInitParams
|
||||
from litellm.types.utils import Message, ModelResponse, StandardLoggingPayload
|
||||
|
|
@ -479,11 +480,11 @@ class NewRelicLogger(CustomLogger):
|
|||
content: Final = message.get("content")
|
||||
|
||||
# Handle tool calls
|
||||
if message.get("tool_calls"):
|
||||
if tool_calls := tool_calls_with_server_tool_calls(message):
|
||||
try:
|
||||
return json.dumps(message["tool_calls"])
|
||||
return json.dumps([tool_call.model_dump() for tool_call in tool_calls])
|
||||
except Exception:
|
||||
return str(message["tool_calls"])
|
||||
return str(tool_calls)
|
||||
|
||||
# Handle None or missing content
|
||||
if content is None:
|
||||
|
|
|
|||
|
|
@ -4,13 +4,13 @@
|
|||
import json
|
||||
import os
|
||||
from collections.abc import Mapping
|
||||
from typing import Final
|
||||
from typing import Final, Protocol
|
||||
|
||||
import httpx
|
||||
|
||||
import litellm
|
||||
from litellm.integrations.custom_logger import CustomLogger
|
||||
from litellm.litellm_core_utils.core_helpers import as_str_mapping
|
||||
from litellm.litellm_core_utils.core_helpers import as_str_mapping, require_attributes
|
||||
from litellm.litellm_core_utils.internal_call_metadata import logged_event_key
|
||||
from litellm.llms.custom_httpx.http_handler import (
|
||||
HTTPHandler,
|
||||
|
|
@ -19,6 +19,10 @@ from litellm.llms.custom_httpx.http_handler import (
|
|||
)
|
||||
|
||||
|
||||
class _SupportsGet(Protocol):
|
||||
def get(self, key: str, default: object = None, /) -> object: ...
|
||||
|
||||
|
||||
def get_utc_datetime():
|
||||
import datetime as dt
|
||||
from datetime import datetime
|
||||
|
|
@ -60,10 +64,11 @@ class OpenMeterLogger(CustomLogger):
|
|||
if (
|
||||
isinstance(response_obj, litellm.ModelResponse) or isinstance(response_obj, litellm.EmbeddingResponse)
|
||||
) and hasattr(response_obj, "usage"):
|
||||
response_usage: Final = require_attributes(response_obj["usage"], _SupportsGet, "get")
|
||||
usage = {
|
||||
"prompt_tokens": response_obj["usage"].get("prompt_tokens", 0),
|
||||
"completion_tokens": response_obj["usage"].get("completion_tokens", 0),
|
||||
"total_tokens": response_obj["usage"].get("total_tokens"),
|
||||
"prompt_tokens": response_usage.get("prompt_tokens", 0),
|
||||
"completion_tokens": response_usage.get("completion_tokens", 0),
|
||||
"total_tokens": response_usage.get("total_tokens"),
|
||||
}
|
||||
|
||||
# OPENMETER_TRUST_REQUEST_USER (default "true"): when set to "false",
|
||||
|
|
|
|||
|
|
@ -6,7 +6,7 @@ from concurrent.futures import ThreadPoolExecutor
|
|||
from dataclasses import dataclass, field
|
||||
from datetime import datetime
|
||||
from types import MappingProxyType
|
||||
from typing import TYPE_CHECKING, Any, Final, cast
|
||||
from typing import TYPE_CHECKING, Any, Final, Protocol, cast
|
||||
|
||||
from pydantic import ConfigDict, TypeAdapter
|
||||
from typing_extensions import ReadOnly, TypedDict
|
||||
|
|
@ -35,6 +35,7 @@ from litellm.integrations.otel.model.metadata import flatten_metadata
|
|||
from litellm.integrations.otel.model.semconv import LiteLLM, Metric
|
||||
from litellm.integrations.otel.plumbing.otlp_tls import resolve_otlp_http_tls
|
||||
from litellm.integrations.otel.routing import routing_decision_attributes
|
||||
from litellm.litellm_core_utils.core_helpers import iter_items, require_attributes
|
||||
from litellm.litellm_core_utils.internal_call_metadata import is_unbilled_non_inference_call_from_params
|
||||
from litellm.litellm_core_utils.safe_json_dumps import safe_dumps
|
||||
from litellm.litellm_core_utils.secret_redaction import redact_string
|
||||
|
|
@ -42,11 +43,11 @@ from litellm.litellm_core_utils.service_tier_utils import (
|
|||
get_requested_service_tier,
|
||||
get_served_service_tier,
|
||||
)
|
||||
from litellm.llms.anthropic.common_utils import anthropic_server_tool_call_params, tool_calls_with_server_tool_calls
|
||||
from litellm.secret_managers.main import get_secret_bool, str_to_bool
|
||||
from litellm.types.integrations.otel_span_attributes import SpanAttributes
|
||||
from litellm.types.services import ServiceLoggerPayload
|
||||
from litellm.types.utils import (
|
||||
ChatCompletionMessageToolCall,
|
||||
CostBreakdown,
|
||||
Function,
|
||||
LLMResponseTypes,
|
||||
|
|
@ -129,6 +130,22 @@ def _text_as_stored_on_span(span: Span, key: str, text: str) -> str:
|
|||
_MAX_DYNAMIC_TRACER_PROVIDERS: Final = 256
|
||||
|
||||
|
||||
class _SupportsGet(Protocol):
|
||||
def get(self, key: str, /) -> object: ...
|
||||
|
||||
|
||||
def _response_field(response_obj: object, key: str) -> object:
|
||||
if not response_obj:
|
||||
return None
|
||||
return require_attributes(response_obj, _SupportsGet, "get").get(key)
|
||||
|
||||
|
||||
def _billed_usage(response_obj: object, call_type: str | None, litellm_params: Mapping[str, object] | None) -> object:
|
||||
if not response_obj or is_unbilled_non_inference_call_from_params(call_type, litellm_params, response_obj):
|
||||
return None
|
||||
return _response_field(response_obj, "usage")
|
||||
|
||||
|
||||
class _CacheHitView(TypedDict, total=False):
|
||||
cache_hit: ReadOnly[object]
|
||||
saved_cache_cost: ReadOnly[object]
|
||||
|
|
@ -2406,17 +2423,17 @@ class OpenTelemetry(OTELGenAISemconvMixin, CustomLogger):
|
|||
|
||||
@staticmethod
|
||||
def _tool_calls_kv_pair(
|
||||
tool_calls: list[ChatCompletionMessageToolCall],
|
||||
tool_calls: Iterable[object],
|
||||
) -> dict[str, object]:
|
||||
kv_pairs: Final[dict[str, object]] = {}
|
||||
for idx, tool_call in enumerate(tool_calls):
|
||||
_function = tool_call.get("function")
|
||||
_function: object = require_attributes(tool_call, _SupportsGet, "get").get("function")
|
||||
if not _function:
|
||||
continue
|
||||
|
||||
keys = Function.__annotations__.keys()
|
||||
for key in keys:
|
||||
_value = _function.get(key)
|
||||
_value = require_attributes(_function, _SupportsGet, "get").get(key)
|
||||
if _value:
|
||||
kv_pairs[f"{SpanAttributes.LLM_COMPLETIONS.value}.{idx}.function_call.{key}"] = _value
|
||||
|
||||
|
|
@ -2583,7 +2600,7 @@ class OpenTelemetry(OTELGenAISemconvMixin, CustomLogger):
|
|||
# but Embeddings and Image-gen responses do not. Fall back to
|
||||
# the litellm call ID so every call type can be correlated
|
||||
# across LiteLLM UI, Phoenix traces, and provider logs (Issue #8).
|
||||
response_id = (response_obj.get("id") if response_obj else None) or standard_logging_payload.get("id")
|
||||
response_id = _response_field(response_obj, "id") or standard_logging_payload.get("id")
|
||||
if response_id:
|
||||
self.safe_set_attribute(
|
||||
span=span,
|
||||
|
|
@ -2600,42 +2617,37 @@ class OpenTelemetry(OTELGenAISemconvMixin, CustomLogger):
|
|||
)
|
||||
|
||||
# The model used to generate the response.
|
||||
if response_obj and response_obj.get("model"):
|
||||
response_model: Final = _response_field(response_obj, "model")
|
||||
if response_model:
|
||||
self.safe_set_attribute(
|
||||
span=span,
|
||||
key=SpanAttributes.LLM_RESPONSE_MODEL.value,
|
||||
value=response_obj.get("model"),
|
||||
value=response_model,
|
||||
)
|
||||
|
||||
self._set_service_tier_attributes(span=span, standard_logging_payload=standard_logging_payload)
|
||||
|
||||
usage: Final = (
|
||||
response_obj.get("usage")
|
||||
if response_obj
|
||||
and not is_unbilled_non_inference_call_from_params(
|
||||
kwargs.get("call_type"), litellm_params, response_obj
|
||||
)
|
||||
else None
|
||||
)
|
||||
usage: Final = _billed_usage(response_obj, kwargs.get("call_type"), litellm_params)
|
||||
if usage:
|
||||
usage_fields: Final = require_attributes(usage, _SupportsGet, "get")
|
||||
self.safe_set_attribute(
|
||||
span=span,
|
||||
key=SpanAttributes.GEN_AI_USAGE_TOTAL_TOKENS.value,
|
||||
value=usage.get("total_tokens"),
|
||||
value=usage_fields.get("total_tokens"),
|
||||
)
|
||||
|
||||
# The number of tokens used in the LLM response (completion).
|
||||
self.safe_set_attribute(
|
||||
span=span,
|
||||
key=SpanAttributes.GEN_AI_USAGE_OUTPUT_TOKENS.value,
|
||||
value=usage.get("completion_tokens"),
|
||||
value=usage_fields.get("completion_tokens"),
|
||||
)
|
||||
|
||||
# The number of tokens used in the LLM prompt.
|
||||
self.safe_set_attribute(
|
||||
span=span,
|
||||
key=SpanAttributes.GEN_AI_USAGE_INPUT_TOKENS.value,
|
||||
value=usage.get("prompt_tokens"),
|
||||
value=usage_fields.get("prompt_tokens"),
|
||||
)
|
||||
|
||||
########################################################################
|
||||
|
|
@ -2711,7 +2723,8 @@ class OpenTelemetry(OTELGenAISemconvMixin, CustomLogger):
|
|||
########## LLM Response Attributes ##########
|
||||
#############################################
|
||||
if response_obj is not None:
|
||||
if response_obj.get("choices"):
|
||||
response_view: Final = require_attributes(response_obj, _SupportsGet, "get")
|
||||
if response_view.get("choices"):
|
||||
transformed_choices: Final = self._transform_choices_to_otel_semantic_conventions(
|
||||
response_obj.get("choices")
|
||||
)
|
||||
|
|
@ -2721,10 +2734,13 @@ class OpenTelemetry(OTELGenAISemconvMixin, CustomLogger):
|
|||
value=safe_dumps(transformed_choices),
|
||||
)
|
||||
|
||||
finish_reasons: Final = []
|
||||
for idx, choice in enumerate(response_obj.get("choices")):
|
||||
if choice.get("finish_reason"):
|
||||
finish_reasons.append(choice.get("finish_reason"))
|
||||
choice_views: Final = [
|
||||
require_attributes(choice, _SupportsGet, "get")
|
||||
for choice in iter_items(response_view.get("choices"))
|
||||
]
|
||||
finish_reasons: Final = [
|
||||
choice.get("finish_reason") for choice in choice_views if choice.get("finish_reason")
|
||||
]
|
||||
|
||||
if finish_reasons:
|
||||
self.safe_set_attribute(
|
||||
|
|
@ -2733,10 +2749,9 @@ class OpenTelemetry(OTELGenAISemconvMixin, CustomLogger):
|
|||
value=safe_dumps(finish_reasons),
|
||||
)
|
||||
|
||||
for idx, choice in enumerate(response_obj.get("choices")):
|
||||
for choice in choice_views:
|
||||
if choice.get("finish_reason"):
|
||||
message = choice.get("message")
|
||||
tool_calls = message.get("tool_calls")
|
||||
tool_calls = tool_calls_with_server_tool_calls(choice.get("message"))
|
||||
if tool_calls:
|
||||
kv_pairs = OpenTelemetry._tool_calls_kv_pair(tool_calls)
|
||||
for key, value in kv_pairs.items():
|
||||
|
|
@ -2746,13 +2761,13 @@ class OpenTelemetry(OTELGenAISemconvMixin, CustomLogger):
|
|||
value=value,
|
||||
)
|
||||
|
||||
elif response_obj.get("output"):
|
||||
elif response_view.get("output"):
|
||||
# Responses API: ResponsesAPIResponse has an "output"
|
||||
# list instead of "choices". Each item with
|
||||
# type="message" contains a "content" list of
|
||||
# OutputText objects (type="output_text").
|
||||
output_items: Final = response_obj.get("output")
|
||||
output_messages: Final = self._transform_responses_api_output_to_otel(output_items)
|
||||
output_items: Final = response_view.get("output")
|
||||
output_messages: Final = self._transform_responses_api_output_to_otel(iter_items(output_items))
|
||||
if output_messages:
|
||||
self.safe_set_attribute(
|
||||
span=span,
|
||||
|
|
@ -2766,7 +2781,7 @@ class OpenTelemetry(OTELGenAISemconvMixin, CustomLogger):
|
|||
# ChatCompletionMessageToolCall format expected by
|
||||
# _tool_calls_kv_pair.
|
||||
tool_calls = []
|
||||
for out_item in output_items:
|
||||
for out_item in iter_items(output_items):
|
||||
item_d = self._to_dict(out_item)
|
||||
if item_d and item_d.get("type") == "function_call":
|
||||
tool_calls.append(
|
||||
|
|
@ -2787,7 +2802,7 @@ class OpenTelemetry(OTELGenAISemconvMixin, CustomLogger):
|
|||
)
|
||||
|
||||
# Extract finish reason from ResponsesAPIResponse.status
|
||||
status: Final = response_obj.get("status")
|
||||
status: Final = response_view.get("status")
|
||||
if status:
|
||||
self.safe_set_attribute(
|
||||
span=span,
|
||||
|
|
@ -2856,7 +2871,9 @@ class OpenTelemetry(OTELGenAISemconvMixin, CustomLogger):
|
|||
transformed_msg = {"role": role, "parts": parts}
|
||||
if "id" in msg:
|
||||
transformed_msg["id"] = msg["id"]
|
||||
if "tool_calls" in msg:
|
||||
if server_tool_calls := anthropic_server_tool_call_params(cast(object, msg)):
|
||||
transformed_msg["tool_calls"] = [*server_tool_calls, *(msg.get("tool_calls") or [])]
|
||||
elif "tool_calls" in msg:
|
||||
transformed_msg["tool_calls"] = msg["tool_calls"]
|
||||
if "tool_call_id" in msg:
|
||||
transformed_msg["tool_call_id"] = msg["tool_call_id"]
|
||||
|
|
@ -2905,7 +2922,7 @@ class OpenTelemetry(OTELGenAISemconvMixin, CustomLogger):
|
|||
return obj.model_dump()
|
||||
return None
|
||||
|
||||
def _transform_responses_api_output_to_otel(self, output: list) -> list[dict]:
|
||||
def _transform_responses_api_output_to_otel(self, output: Iterable[object]) -> list[dict]:
|
||||
"""
|
||||
Transform Responses API output items into OTEL GenAI 1.38 format.
|
||||
|
||||
|
|
|
|||
|
|
@ -1,4 +1,5 @@
|
|||
from abc import ABC
|
||||
from collections.abc import Mapping
|
||||
from typing import TYPE_CHECKING, Any, Final
|
||||
|
||||
if TYPE_CHECKING:
|
||||
|
|
@ -7,7 +8,7 @@ if TYPE_CHECKING:
|
|||
|
||||
class BaseLLMObsOTELAttributes(ABC):
|
||||
@staticmethod
|
||||
def set_messages(span: "Span", kwargs: dict[str, Any]):
|
||||
def set_messages(span: "Span", kwargs: Mapping[str, object]):
|
||||
pass
|
||||
|
||||
@staticmethod
|
||||
|
|
|
|||
|
|
@ -1034,10 +1034,12 @@ def emit_guardrail_span(entry: "StandardLoggingGuardrailInformation") -> None:
|
|||
pass
|
||||
|
||||
|
||||
def seed_request_identity(user_api_key_dict: object, model: str | None = None) -> None:
|
||||
def seed_request_identity(user_api_key_dict: object, model: str | None = None, request_data: object = None) -> None:
|
||||
logger: Final = _registered_v2_logger()
|
||||
if logger is not None:
|
||||
logger.seed_request_identity(user_api_key_dict, model=model)
|
||||
logger.seed_request_identity(
|
||||
user_api_key_dict, model=model, request_metadata=metadata_from_request_data(request_data)
|
||||
)
|
||||
|
||||
|
||||
def fan_out_provider() -> ApiTracerProvider:
|
||||
|
|
|
|||
Some files were not shown because too many files have changed in this diff Show more
Loading…
Add table
Reference in a new issue