scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-13 23:11:40 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 101

chore: merge current staging into JWT E2E infrastructure

Browse source
This commit is contained in:
Yuneng Jiang 2026-09-11 17:13:52 -07:00
commit a706dbbb7d
No known key found for this signature in database
52 changed files with 12333 additions and 316 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

103
.github/workflows/test-e2e-redis-chaos.yml vendored Normal file
View file

@ -0,0 +1,103 @@
name: "Redis Chaos E2E"
on:
workflow_dispatch:
workflow_call:
inputs:
ref:
description: "Commit SHA or ref to test. Defaults to the ref the workflow was triggered on"
required: false
type: string
permissions:
contents: read
jobs:
redis-chaos-e2e:
runs-on: ubuntu-latest-16-cores
timeout-minutes: 30
services:
postgres:
image: postgres:16.6@sha256:557fea37a744d5f4c8faab304b0a90858b53ab119735a88c131fd19dab802f36
env:
POSTGRES_USER: llmproxy
POSTGRES_PASSWORD: dbpassword9090
POSTGRES_DB: litellm
ports:
- 5432:5432
options: >-
--health-cmd "pg_isready -U llmproxy"
--health-interval 5s
--health-timeout 5s
--health-retries 10
valkey:
image: valkey/valkey:8.1.4@sha256:81db6d39e1bba3b3ff32bd3a1b19a6d69690f94a3954ec131277b9a26b95b3aa
ports:
- 6379:6379
options: >-
--health-cmd "valkey-cli ping"
--health-interval 5s
--health-timeout 5s
--health-retries 10
env:
DATABASE_URL: postgresql://llmproxy:dbpassword9090@localhost:5432/litellm
LITELLM_MASTER_KEY: sk-redis-chaos-e2e
LITELLM_LOG: WARNING
JSON_LOGS: "true"
steps:
- uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0
with:
persist-credentials: false
ref: ${{ inputs.ref || github.sha }}
- name: Set up Python
uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5.6.0
with:
python-version: "3.12"
- name: Set up uv
uses: ./.github/actions/setup-uv-with-retries
with:
version: "0.10.9"
- name: Cache the Rust build
uses: ./.github/actions/cache-cargo-build
- name: Install dependencies
run: |
.github/scripts/uv_sync_with_retries.sh --frozen --group ci --group proxy-dev --group e2e-dev --extra proxy
- name: Cache Prisma binaries
uses: ./.github/actions/cache-prisma-binaries
- name: Generate Prisma client
run: |
uv run --no-sync prisma generate --schema litellm/proxy/schema.prisma
- name: Start a multi-worker proxy on the chaos config
run: |
nohup uv run --no-sync litellm --config tests/e2e/gateway/redis_chaos_ci_config.yml --port 4000 --num_workers 4 > proxy.log 2>&1 &
echo "E2E_PROXY_PID=$!" >> "$GITHUB_ENV"
echo "E2E_PROXY_LOG=$(pwd)/proxy.log" >> "$GITHUB_ENV"
for _ in $(seq 1 90); do
if curl -fs http://localhost:4000/health/liveliness > /dev/null; then
exit 0
fi
sleep 2
done
echo "proxy never became live"
tail -n 100 proxy.log
exit 1
- name: Run the Redis chaos load test
env:
E2E_REDIS_CHAOS: "1"
LITELLM_PROXY_URL: http://localhost:4000
REDIS_HOST: 127.0.0.1
REDIS_PORT: "6379"
run: |
uv run --no-sync pytest tests/e2e/load/test_redis_chaos_e2e.py -v --tb=short -rA -s
- name: Show proxy log on failure
if: failure()
run: tail -n 300 proxy.log

136
enterprise/litellm_enterprise/enterprise_callbacks/secret_detection.py
View file

@ -11,10 +11,14 @@ import sys
sys.path.insert(
0, os.path.abspath("../..")
) # Adds the parent directory to the system path
import functools
import configparser
import contextlib
import itertools
import re
import tempfile
from collections.abc import Generator, Iterator, Sequence
from contextvars import ContextVar
from typing import TYPE_CHECKING, ClassVar, Literal, Optional
from typing import TYPE_CHECKING, ClassVar, Final, Literal, Optional
from litellm._logging import verbose_proxy_logger
from litellm.caching.caching import DualCache
@ -433,12 +437,101 @@ _default_detect_secrets_config = {
"name": "ZendeskSecretKeyDetector",
"path": _custom_plugins_path + "/zendesk_secret_key.py",
},
{
"name": "CredentialKeywordDetector",
"path": _custom_plugins_path + "/credential_keyword.py",
},
{"name": "Base64HighEntropyString", "limit": 4.5},
{"name": "HexHighEntropyString", "limit": 3.0},
],
}
_CONFIG_SECTION: Final = "litellm-prompt"
_ASSIGNMENT_LINE: Final = re.compile(r"[^\s\[#;:=][^:=]*[:=]")
_SHELL_ASSIGNMENT: Final = re.compile(r"(?P<key>[^\s\[#;:=](?:[^:=]*[^\s:=])?)=(?P<value>\S+)")
_SHELL_OPERATORS: Final = ";&|"
_SHELL_TRAILER: Final = re.compile(r"\\|#.*|-*\w[\w.-]*=\S*")
_SCAN_SUFFIX: Final = ".py"
@contextlib.contextmanager
def _temp_file(text: str) -> Generator[str, None, None]:
temp_file: Final = tempfile.NamedTemporaryFile(suffix=_SCAN_SUFFIX, delete=False)
try:
temp_file.write(text.encode("utf-8"))
temp_file.close()
yield temp_file.name
finally:
temp_file.close()
os.remove(temp_file.name)
def _scan_lines(lines: Sequence[str]) -> frozenset[tuple[str, str]]:
from detect_secrets import SecretsCollection
secrets: Final = SecretsCollection()
with _temp_file("\n".join(lines)) as path:
secrets.scan_file(path)
return frozenset(
(found_secret.secret_value, found_secret.type)
for file in secrets.files
for found_secret in secrets[file]
if found_secret.secret_value is not None
)
def _classify_line(state: tuple[bool, str | None], numbered: tuple[int, str]) -> tuple[bool, str | None]:
open_option: Final = state[0]
number, line = numbered
stripped: Final = line.strip()
if not stripped or stripped[0] in "#;":
return open_option, None
shell_assignment: Final = _SHELL_ASSIGNMENT.match(stripped)
if shell_assignment is not None:
return True, f"{shell_assignment['key']}_{number}={shell_assignment['value']}"
assignment: Final = _ASSIGNMENT_LINE.match(stripped)
if assignment is not None:
return True, f"{assignment.group()[:-1].strip()}_{number}{stripped[assignment.end() - 1 :]}"
if line[0].isspace() and open_option:
return True, line
return False, None
def _parseable_lines(text: str) -> Iterator[str]:
states: Final = itertools.accumulate(enumerate(text.splitlines()), _classify_line, initial=(False, None))
return (line for _, line in states if line is not None)
def _lone_value(line: str) -> str | None:
tokens: Final = line.split()
if not tokens or '"' in tokens[0]:
return None
value: Final = tokens[0].rstrip(_SHELL_OPERATORS)
if len(tokens) == 1 or value != tokens[0] or _SHELL_TRAILER.fullmatch(tokens[1]) is not None:
return value
return None
def _quoted_assignments(text: str) -> tuple[str, ...]:
parser: Final = configparser.ConfigParser(interpolation=None)
parser.optionxform = str # pyright: ignore[reportAttributeAccessIssue] # configparser types optionxform as a method
parser.read_string(f"[{_CONFIG_SECTION}]\n" + "\n".join(_parseable_lines(text)))
return tuple(
f'{key} = "{value}"'
for section in parser
for key, values in parser.items(section)
for line in values.splitlines()
if (value := _lone_value(line)) is not None
)
class _ENTERPRISE_SecretDetection(CustomGuardrail):
# Keeps proxied traffic on async_pre_call_hook (the unified apply_guardrail
# path skips should_run_check and never sees data["prompt"]).
@ -449,35 +542,21 @@ class _ENTERPRISE_SecretDetection(CustomGuardrail):
super().__init__(**kwargs)
def scan_message_for_secrets(self, message_content: str):
from detect_secrets import SecretsCollection
from detect_secrets.settings import transient_settings
temp_file = tempfile.NamedTemporaryFile(delete=False)
temp_file.write(message_content.encode("utf-8"))
temp_file.close()
secrets = SecretsCollection()
detect_secrets_config = (
self.user_defined_detect_secrets_config or _default_detect_secrets_config
)
with transient_settings(detect_secrets_config):
secrets.scan_file(temp_file.name)
os.remove(temp_file.name)
found: Final = _scan_lines(
(*message_content.splitlines(), *_quoted_assignments(message_content))
)
return [
{"type": found_secret.type, "value": found_secret.secret_value}
for file in sorted(secrets.files)
for found_secret in sorted(
secrets[file],
key=lambda secret: (
-len(secret.secret_value or ""),
secret.type,
secret.secret_value or "",
),
{"type": secret_type, "value": value}
for value, secret_type in sorted(
found, key=lambda pair: (-len(pair[0]), pair[1], pair[0])
)
if found_secret.secret_value is not None
]
def redact_text(self, text: str, source: str = "message") -> str:
@ -490,15 +569,16 @@ class _ENTERPRISE_SecretDetection(CustomGuardrail):
if counts is not None:
for secret in detected_secrets:
counts[secret["type"]] = counts.get(secret["type"], 0) + 1
secret_types = [secret["type"] for secret in detected_secrets]
secret_types: Final = sorted(
dict.fromkeys(secret["type"] for secret in detected_secrets)
)
verbose_proxy_logger.warning(
f"Detected and redacted secrets in {source}: {secret_types}"
"Detected and redacted secrets in %s: %s", source, secret_types
)
return functools.reduce(
lambda redacted, secret: redacted.replace(secret["value"], "[REDACTED]"),
detected_secrets,
text,
pattern: Final = re.compile(
"|".join(re.escape(secret["value"]) for secret in detected_secrets)
)
return pattern.sub("[REDACTED]", text)
async def should_run_check(self, user_api_key_dict: UserAPIKeyAuth) -> bool:
if user_api_key_dict.permissions is not None:

63
enterprise/litellm_enterprise/enterprise_callbacks/secrets_plugins/credential_keyword.py Normal file
View file

@ -0,0 +1,63 @@
import re
from collections.abc import Generator, Mapping
from string import punctuation
from typing import Final
from detect_secrets.plugins.keyword import (
QUOTES_REQUIRED_DENYLIST_REGEX_TO_GROUP,
KeywordDetector,
)
_CREDENTIAL_VALUE: Final = re.compile(r"[^\s()\[\]]+")
_ENVIRONMENT_REFERENCE: Final = re.compile(r"os\.environ/\w+", re.IGNORECASE)
_ENVIRONMENT_VARIABLE_NAME: Final = re.compile(r"[A-Z][A-Z0-9]*(?:_[A-Z0-9]+)+")
_LOWERCASE_WORD_SEQUENCE: Final = re.compile(r"[a-z]+(?:[-._/][a-z]+)+")
_ISO_8601_TIMESTAMP: Final = re.compile(
r"\d{4}-\d{2}-\d{2}(?:T\d{2}:\d{2}(?::\d{2}(?:\.\d+)?)?(?:Z|[+-]\d{2}:?\d{2})?)?"
)
_URL_WITHOUT_USERINFO_OR_QUERY: Final = re.compile(r"[A-Za-z][A-Za-z0-9+.-]*://[^\s@?]*")
_BENIGN_VALUES: Final = (
_ENVIRONMENT_REFERENCE,
_ENVIRONMENT_VARIABLE_NAME,
_LOWERCASE_WORD_SEQUENCE,
_ISO_8601_TIMESTAMP,
_URL_WITHOUT_USERINFO_OR_QUERY,
)
class CredentialKeywordDetector(KeywordDetector): # pyright: ignore[reportUntypedBaseClass] # detect_secrets ships no type information
secret_type = "Credential Keyword"
def __init__(self, minimum_length: int = 12, keyword_exclude: str | None = None) -> None:
if (
not isinstance(minimum_length, int) # pyright: ignore[reportUnnecessaryIsInstance] # the value comes from an operator's YAML
or minimum_length < 1
):
raise ValueError(f"minimum_length must be a positive integer, got {minimum_length!r}")
super().__init__(keyword_exclude=keyword_exclude)
self.minimum_length = minimum_length
def _is_credential(self, value: str) -> bool:
core: Final = value.strip(punctuation)
return (
len(value) >= self.minimum_length
and _CREDENTIAL_VALUE.fullmatch(value) is not None
and all(benign.fullmatch(core) is None for benign in _BENIGN_VALUES)
)
def analyze_string(
self,
string: str,
denylist_regex_to_group: Mapping[re.Pattern[str], int] | None = None,
) -> Generator[str, None, None]:
if self.keyword_exclude is not None and self.keyword_exclude.search(string):
return
regex_to_group: Final = (
QUOTES_REQUIRED_DENYLIST_REGEX_TO_GROUP if denylist_regex_to_group is None else denylist_regex_to_group
)
yield from (
match.group(group)
for regex, group in regex_to_group.items()
for match in regex.finditer(string)
if self._is_credential(match.group(group))
)

2
litellm-rust/Cargo.lock generated
View file

@ -2002,11 +2002,13 @@ dependencies = [
name = "litellm-token-counter"
version = "0.1.0"
dependencies = [
"base64 0.22.1",
"criterion",
"indexmap",
"itoa",
"rand 0.8.7",
"rstest",
"rustc-hash",
"serde",
"serde_json",
"thiserror 2.0.19",

32
litellm-rust/crates/python-bridge/src/token_counter.rs
View file

@ -30,12 +30,17 @@ struct TokenCounter {
impl TokenCounter {
#[new]
fn new(py: Python<'_>, tokenizer_json: &str) -> PyResult<Self> {
let inner = release_gil(py, || CoreTokenCounter::from_json(tokenizer_json))
.map_err(token_count_error_to_pyerr)?;
Ok(Self {
inner: Arc::new(inner),
encode_slots: Arc::new(Semaphore::new(encode_parallelism())),
})
Self::load(py, || CoreTokenCounter::from_json(tokenizer_json))
}
#[staticmethod]
fn from_cl100k_ranks(py: Python<'_>, rank_file: &str) -> PyResult<Self> {
Self::load(py, || CoreTokenCounter::from_cl100k_ranks(rank_file))
}
#[staticmethod]
fn from_o200k_ranks(py: Python<'_>, rank_file: &str) -> PyResult<Self> {
Self::load(py, || CoreTokenCounter::from_o200k_ranks(rank_file))
}
fn acount_request<'py>(&self, py: Python<'py>, body: &[u8]) -> PyResult<Bound<'py, PyAny>> {
@ -58,6 +63,19 @@ impl TokenCounter {
}
}
impl TokenCounter {
fn load(
py: Python<'_>,
load: impl FnOnce() -> Result<CoreTokenCounter, Error> + Send,
) -> PyResult<Self> {
let inner = release_gil(py, load).map_err(token_count_error_to_pyerr)?;
Ok(Self {
inner: Arc::new(inner),
encode_slots: Arc::new(Semaphore::new(encode_parallelism())),
})
}
}
fn encode_parallelism() -> usize {
available_parallelism().map_or(TOKEN_COUNT_FALLBACK_PARALLELISM, NonZero::get)
}
@ -70,7 +88,7 @@ fn count_body(counter: &CoreTokenCounter, body: &[u8]) -> Result<InputTokenCount
fn token_count_error_to_pyerr(error: Error) -> PyErr {
let message = error.to_string();
match error {
Error::Load(_) => PyValueError::new_err(message),
Error::Load(_) | Error::Ranks(_) | Error::UnicodeClasses => PyValueError::new_err(message),
Error::RequestParse(_)
| Error::MissingInput
| Error::FloatText

2
litellm-rust/crates/token-counter/Cargo.toml
View file

@ -6,8 +6,10 @@ license.workspace = true
repository.workspace = true
[dependencies]
base64.workspace = true
indexmap = { version = "2.14.0", features = ["serde"] }
itoa = "1.0"
rustc-hash = "2.1.3"
serde.workspace = true
serde_json.workspace = true
thiserror.workspace = true

29
litellm-rust/crates/token-counter/src/byte_level.rs
View file

@ -12,7 +12,7 @@ use tokenizers::pre_tokenizers::PreTokenizerWrapper;
use tokenizers::{Model, Tokenizer};
use unicode_normalization_alignments::{IsNormalized, UnicodeNormalization, is_nfkc_quick};
use super::unicode_classes::UnicodeClasses;
use super::unicode_classes::{Class, UnicodeClasses, class, run_len};
const CONTRACTIONS: [&str; 7] = ["'s", "'t", "'re", "'ve", "'m", "'ll", "'d"];
@ -110,27 +110,6 @@ fn mapped_len(piece: &str) -> usize {
.count()
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
enum Class {
Letter,
Number,
Space,
Other,
}
fn class(character: char, unicode_classes: &UnicodeClasses) -> Class {
match character {
'A'..='Z' | 'a'..='z' => Class::Letter,
'0'..='9' => Class::Number,
'\t'..='\r' | ' ' => Class::Space,
_ if character.is_ascii() => Class::Other,
_ if unicode_classes.is_letter(character) => Class::Letter,
_ if unicode_classes.is_number(character) => Class::Number,
_ if unicode_classes.is_space(character) => Class::Space,
_ => Class::Other,
}
}
/// The regex matches every character, so the pieces tile the text.
fn pieces<'a>(
text: &'a str,
@ -169,12 +148,6 @@ fn piece_len(text: &str, first: char, unicode_classes: &UnicodeClasses) -> usize
}
}
fn run_len(text: &str, run_class: Class, unicode_classes: &UnicodeClasses) -> usize {
text.char_indices()
.find(|(_, character)| class(*character, unicode_classes) != run_class)
.map_or(text.len(), |(index, _)| index)
}
/// `\s+(?!\S)|\s+`: whitespace followed by a non-space leaves its last
/// character to start the next piece (` ?` on the following alternatives).
fn space_run_len(text: &str, unicode_classes: &UnicodeClasses) -> usize {

125
litellm-rust/crates/token-counter/src/cl100k.rs Normal file
View file

@ -0,0 +1,125 @@
//! Scanner for tiktoken's `cl100k_base` split regex,
//! `'(?i:[sdmt]|ll|ve|re)|[^\r\n\p{L}\p{N}]?+\p{L}++|\p{N}{1,3}+| ?[^\s\p{L}\p{N}]++[\r\n]*+|\s++$|\s*[\r\n]|\s+(?!\S)|\s`.
use super::scanner::{contraction_len, digit_run_len, is_newline};
use super::unicode_classes::{Class, UnicodeClasses, class, run_len};
/// The alternatives in regex order; the possessive quantifiers mean an
/// alternative that starts matching and runs out of input fails as a whole.
pub(super) fn piece_len(text: &str, first: char, unicode_classes: &UnicodeClasses) -> usize {
if let Some(len) = contraction_len(text) {
return len;
}
let first_class = class(first, unicode_classes);
match first_class {
Class::Letter => return run_len(text, Class::Letter, unicode_classes),
Class::Number => return digit_run_len(text, unicode_classes),
Class::Space | Class::Other => {}
}
let rest = &text[first.len_utf8()..];
let second_class = rest
.chars()
.next()
.map(|character| class(character, unicode_classes));
if !is_newline(first) && second_class == Some(Class::Letter) {
return first.len_utf8() + run_len(rest, Class::Letter, unicode_classes);
}
if first_class == Class::Other {
return symbol_run_len(text, unicode_classes);
}
if first == ' ' && second_class == Some(Class::Other) {
return 1 + symbol_run_len(rest, unicode_classes);
}
space_run_len(text, unicode_classes)
}
/// `[^\s\p{L}\p{N}]++[\r\n]*+`
fn symbol_run_len(text: &str, unicode_classes: &UnicodeClasses) -> usize {
let symbols = run_len(text, Class::Other, unicode_classes);
symbols
+ text[symbols..]
.bytes()
.take_while(|byte| matches!(byte, b'\r' | b'\n'))
.count()
}
/// `\s++$|\s*[\r\n]|\s+(?!\S)|\s`: whitespace to the end of the text is one
/// piece; otherwise the piece ends at the last newline of the run, or leaves
/// the run's last character for the next piece's optional leading space.
fn space_run_len(text: &str, unicode_classes: &UnicodeClasses) -> usize {
let run = run_len(text, Class::Space, unicode_classes);
if run == text.len() {
return run;
}
if let Some(newline) = text[..run].rfind(['\r', '\n']) {
return newline + 1;
}
let last = text[..run].chars().next_back().map_or(0, char::len_utf8);
match run - last {
0 => run,
shorter => shorter,
}
}
#[cfg(test)]
mod tests {
use rstest::rstest;
use super::*;
use crate::scanner::pieces;
fn split(text: &str) -> Vec<&str> {
pieces(
text,
piece_len,
UnicodeClasses::get().expect("Oniguruma exposes Unicode classes"),
)
.collect()
}
#[rstest]
#[case("", &[])]
#[case("Hello world", &["Hello", " world"])]
#[case("don't I'LL you'Ve we'RE he'd I'm", &["don", "'t", " I", "'LL", " you", "'Ve", " we", "'RE", " he", "'d", " I", "'m"])]
#[case("IT'SOK it'Dbe 'Sx 'Tx", &["IT", "'S", "OK", " it", "'D", "be", " '", "Sx", " '", "Tx"])]
#[case("'Sx'Tx'Mx'LLx'VEx'REx'Dx", &["'S", "x", "'T", "x", "'M", "x", "'LL", "x", "'VE", "x", "'RE", "x", "'D", "x"])]
#[case("'ſ 'lx", &["'ſ", " '", "lx"])]
#[case("12345 6", &["123", "45", " ", "6"])]
#[case("!abc !!abc", &["!abc", " !!", "abc"])]
#[case(" !!!\r\n\r\nx", &[" !!!\r\n\r\n", "x"])]
#[case("a b \n\n c", &["a", " ", " b", " \n\n", " ", " c"])]
#[case("a\nb\r\nc\n\nd \n e", &["a", "\n", "b", "\r\n", "c", "\n\n", "d", " \n", " e"])]
#[case("x \t\n \t y\n", &["x", " \t\n", " \t", " y", "\n"])]
#[case("end ", &["end", " "])]
#[case("\u{a0}abc\u{a0}!", &["\u{a0}abc", "\u{a0}", "!"])]
#[case("<|endoftext|>", &["<|", "endoftext", "|>"])]
#[case("e\u{301}a", &["e", "\u{301}a"])]
#[case("日本語 ١٢٣٤", &["日本語", " ", "١٢٣", "٤"])]
fn scanner_splits_like_the_regex(#[case] text: &str, #[case] expected: &[&str]) {
assert_eq!(split(text), expected);
}
#[derive(serde::Deserialize)]
struct TextFixture {
text: String,
pieces: Vec<String>,
}
#[test]
fn scanner_splits_the_fixture_corpus_like_tiktoken_regex() {
let fixtures: Vec<TextFixture> = include_str!(concat!(
env!("CARGO_MANIFEST_DIR"),
"/tests/fixtures/cl100k/texts.jsonl"
))
.lines()
.map(|line| serde_json::from_str(line).expect("fixture line is json"))
.collect();
assert!(fixtures.len() > 3000);
let mismatches: Vec<_> = fixtures
.iter()
.filter(|fixture| split(&fixture.text) != fixture.pieces)
.map(|fixture| (&fixture.text, split(&fixture.text), &fixture.pieces))
.collect();
assert!(mismatches.is_empty(), "{mismatches:#?}");
}
}

69
litellm-rust/crates/token-counter/src/counter.rs
View file

@ -3,6 +3,7 @@ use serde::Serialize;
use crate::Error;
use crate::byte_level::ByteLevelCounter;
use crate::python_json;
use crate::scanner::{SplitPattern, TiktokenCounter};
use crate::tools::format_function_definitions;
use crate::types::{
ContentBlock, ContentItem, CountableRequest, Message, MessageContent, TextValue, ToolChoice,
@ -23,12 +24,19 @@ pub struct InputTokenCount {
pub input_tokens: usize,
}
/// A loaded HuggingFace tokenizer plus the message accounting Python applies on
/// top of it. Encoding is CPU-bound and synchronous; hosts run it off their
/// event loop.
enum Encoder {
HuggingFace {
tokenizer: Box<tokenizers::Tokenizer>,
byte_level: Option<ByteLevelCounter>,
},
Tiktoken(TiktokenCounter),
}
/// A loaded tokenizer plus the message accounting Python applies on top of
/// it. Encoding is CPU-bound and synchronous; hosts run it off their event
/// loop.
pub struct TokenCounter {
tokenizer: tokenizers::Tokenizer,
byte_level: Option<ByteLevelCounter>,
encoder: Encoder,
}
impl TokenCounter {
@ -39,23 +47,50 @@ impl TokenCounter {
.map_err(Error::Load)?;
let byte_level = ByteLevelCounter::detect(&tokenizer);
Ok(Self {
tokenizer,
byte_level,
encoder: Encoder::HuggingFace {
tokenizer: Box::new(tokenizer),
byte_level,
},
})
}
/// Load tiktoken's `cl100k_base` rank file (`base64(token) rank` lines).
/// The host reads the file.
pub fn from_cl100k_ranks(rank_file: &str) -> Result<Self, Error> {
Self::from_tiktoken_ranks(SplitPattern::Cl100k, rank_file)
}
/// Load tiktoken's `o200k_base` rank file (`base64(token) rank` lines).
/// The host reads the file.
pub fn from_o200k_ranks(rank_file: &str) -> Result<Self, Error> {
Self::from_tiktoken_ranks(SplitPattern::O200k, rank_file)
}
fn from_tiktoken_ranks(split: SplitPattern, rank_file: &str) -> Result<Self, Error> {
Ok(Self {
encoder: Encoder::Tiktoken(TiktokenCounter::from_ranks(split, rank_file)?),
})
}
pub fn count_text(&self, text: &str) -> Result<usize, Error> {
if let Some(count) = self
.byte_level
.as_ref()
.and_then(|counter| counter.count(&self.tokenizer, text))
{
return Ok(count);
match &self.encoder {
Encoder::Tiktoken(counter) => Ok(counter.count(text)),
Encoder::HuggingFace {
tokenizer,
byte_level,
} => {
if let Some(count) = byte_level
.as_ref()
.and_then(|counter| counter.count(tokenizer, text))
{
return Ok(count);
}
tokenizer
.encode_fast(text, true)
.map(|encoding| encoding.len())
.map_err(Error::Encode)
}
}
self.tokenizer
.encode_fast(text, true)
.map(|encoding| encoding.len())
.map_err(Error::Encode)
}
/// Mirrors the host's key precedence: `messages`, then `prompt`, then

4
litellm-rust/crates/token-counter/src/error.rs
View file

@ -6,6 +6,10 @@ use thiserror::Error as ThisError;
pub enum Error {
#[error("failed to load tokenizer: {0}")]
Load(#[source] tokenizers::Error),
#[error("failed to load tokenizer: tiktoken rank file: {0}")]
Ranks(String),
#[error("failed to load tokenizer: Unicode character classes are unavailable")]
UnicodeClasses,
#[error("unsupported by the rust token counter: request body could not be parsed: {0}")]
RequestParse(#[source] serde_json::Error),
#[error("unsupported by the rust token counter: request has no countable input")]

4
litellm-rust/crates/token-counter/src/lib.rs
View file

@ -5,9 +5,13 @@
#![forbid(unsafe_code)]
mod byte_level;
mod cl100k;
mod counter;
mod error;
mod o200k;
mod python_json;
mod scanner;
mod tiktoken;
mod tools;
mod types;
mod unicode_classes;

217
litellm-rust/crates/token-counter/src/o200k.rs Normal file
View file

@ -0,0 +1,217 @@
//! Scanner for tiktoken's `o200k_base` split regex,
//! `[^\r\n\p{L}\p{N}]?[\p{Lu}\p{Lt}\p{Lm}\p{Lo}\p{M}]*[\p{Ll}\p{Lm}\p{Lo}\p{M}]+(?i:'s|'t|'re|'ve|'m|'ll|'d)?|[^\r\n\p{L}\p{N}]?[\p{Lu}\p{Lt}\p{Lm}\p{Lo}\p{M}]+[\p{Ll}\p{Lm}\p{Lo}\p{M}]*(?i:'s|'t|'re|'ve|'m|'ll|'d)?|\p{N}{1,3}| ?[^\s\p{L}\p{N}]+[\r\n/]*|\s*[\r\n]+|\s+(?!\S)|\s+`.
//! tiktoken runs it with a backtracking engine, so the letter alternatives
//! below reproduce where the greedy quantifiers settle, not only what the
//! classes say.
use super::scanner::{contraction_len, digit_run_len, is_newline};
use super::unicode_classes::{Case, Class, UnicodeClasses, case, case_run_len, class, run_len};
/// The alternatives in regex order: a number is never a letter piece, a
/// letter always is, and only whitespace and symbols reach the last three.
pub(super) fn piece_len(text: &str, first: char, unicode_classes: &UnicodeClasses) -> usize {
let first_class = class(first, unicode_classes);
if first_class == Class::Number {
return digit_run_len(text, unicode_classes);
}
if let Some(len) = letter_piece_len(text, first, first_class, unicode_classes) {
return len;
}
if first_class == Class::Other {
return symbol_run_len(text, unicode_classes);
}
let rest = &text[first.len_utf8()..];
if first == ' '
&& rest
.chars()
.next()
.is_some_and(|character| class(character, unicode_classes) == Class::Other)
{
return 1 + symbol_run_len(rest, unicode_classes);
}
space_run_len(text, unicode_classes)
}
/// The two letter alternatives, each first with then without the optional
/// `[^\r\n\p{L}\p{N}]` prefix: the order the engine tries them in.
fn letter_piece_len(
text: &str,
first: char,
first_class: Class,
unicode_classes: &UnicodeClasses,
) -> Option<usize> {
let prefix = (!is_newline(first) && matches!(first_class, Class::Space | Class::Other))
.then(|| first.len_utf8());
let after_prefix = |shape: fn(&str, &UnicodeClasses) -> Option<usize>| {
prefix.and_then(|prefix| shape(&text[prefix..], unicode_classes).map(|len| prefix + len))
};
after_prefix(upper_then_lower_len)
.or_else(|| upper_then_lower_len(text, unicode_classes))
.or_else(|| after_prefix(upper_run_len))
.or_else(|| upper_run_len(text, unicode_classes))
}
/// `[\p{Lu}\p{Lt}\p{Lm}\p{Lo}\p{M}]*[\p{Ll}\p{Lm}\p{Lo}\p{M}]+(?i:'s|'t|'re|'ve|'m|'ll|'d)?`.
/// The upper run is greedy; when no lower character follows it, the engine
/// gives characters back until the one it just gave back is lower too, and
/// that single character is the lower run.
fn upper_then_lower_len(text: &str, unicode_classes: &UnicodeClasses) -> Option<usize> {
let upper = case_run_len(text, Case::is_upper, unicode_classes);
let lower = case_run_len(&text[upper..], Case::is_lower, unicode_classes);
let letters = if lower > 0 {
upper + lower
} else {
let (index, last_both) = text[..upper]
.char_indices()
.rev()
.find(|(_, character)| case(*character, unicode_classes).is_lower())?;
index + last_both.len_utf8()
};
Some(letters + contraction_len(&text[letters..]).unwrap_or(0))
}
/// `[\p{Lu}\p{Lt}\p{Lm}\p{Lo}\p{M}]+[\p{Ll}\p{Lm}\p{Lo}\p{M}]*(?i:'s|'t|'re|'ve|'m|'ll|'d)?`
fn upper_run_len(text: &str, unicode_classes: &UnicodeClasses) -> Option<usize> {
let upper = case_run_len(text, Case::is_upper, unicode_classes);
if upper == 0 {
return None;
}
let letters = upper + case_run_len(&text[upper..], Case::is_lower, unicode_classes);
Some(letters + contraction_len(&text[letters..]).unwrap_or(0))
}
/// `[^\s\p{L}\p{N}]+[\r\n/]*`
fn symbol_run_len(text: &str, unicode_classes: &UnicodeClasses) -> usize {
let symbols = run_len(text, Class::Other, unicode_classes);
symbols
+ text[symbols..]
.bytes()
.take_while(|byte| matches!(byte, b'\r' | b'\n' | b'/'))
.count()
}
/// `\s*[\r\n]+|\s+(?!\S)|\s+`: a run with a newline ends at its last newline,
/// even at the end of the text; otherwise whitespace to the end of the text
/// is one piece, or the run leaves its last character for the next piece's
/// optional leading space.
fn space_run_len(text: &str, unicode_classes: &UnicodeClasses) -> usize {
let run = run_len(text, Class::Space, unicode_classes);
if let Some(newline) = text[..run].rfind(['\r', '\n']) {
return newline + 1;
}
if run == text.len() {
return run;
}
let last = text[..run].chars().next_back().map_or(0, char::len_utf8);
match run - last {
0 => run,
shorter => shorter,
}
}
#[cfg(test)]
mod tests {
use rstest::rstest;
use tokenizers::utils::SysRegex;
use super::*;
use crate::scanner::pieces;
fn split(text: &str) -> Vec<&str> {
pieces(
text,
piece_len,
UnicodeClasses::get().expect("Oniguruma exposes Unicode classes"),
)
.collect()
}
#[rstest]
#[case("", &[])]
#[case("Hello world", &["Hello", " world"])]
#[case("camelCase PascalCase ABCdef ABCdeF ABC", &["camel", "Case", " Pascal", "Case", " ABCdef", " ABCde", "F", " ABC"])]
#[case("日本ABC ABC日本 日本語abc abc日本語", &["日本", "ABC", " ABC日本", " 日本語abc", " abc日本語"])]
#[case("\u{301}ABC \u{301}abc \u{301}\u{301}A A\u{301}\u{301} E\u{301}A aE\u{301}", &["\u{301}", "ABC", " \u{301}abc", " \u{301}\u{301}", "A", " A\u{301}\u{301}", " E\u{301}", "A", " a", "E\u{301}"])]
#[case("ᵃbc ᵃBC Aᵃbc Aᵃ ᵃ' ᵃ's", &["ᵃbc", " ᵃ", "BC", " Aᵃbc", " Aᵃ", " ᵃ", "'", " ᵃ's"])]
#[case("Džungla aDžB ADžB ADžb", &["Džungla", " a", "DžB", " ADžB", " ADžb"])]
#[case("don'tx ABC's abc'S abc'ſ ABC'ſx IT'SOK it'Dbe", &["don't", "x", " ABC's", " abc'S", " abc'ſ", " ABC'ſ", "x", " IT'S", "OK", " it'D", "be"])]
#[case("'sabc x's 's 'Sx'Tx 9'9 a'9 ' s", &["'sabc", " x's", " '", "s", " '", "Sx'T", "x", " ", "9", "'", "9", " a", "'", "9", " '", " s"])]
#[case("!ABC !AbC !!abc !!\u{301}a \u{a0}\u{301}A", &["!ABC", " !", "Ab", "C", " !!", "abc", " !!\u{301}", "a", " ", "\u{a0}\u{301}", "A"])]
#[case("!!/\n/x a/b !!\n/x /x //", &["!!/\n/", "x", " a", "/b", " !!\n/", "x", " ", " /", "x", " ", " //"])]
#[case("12345 6 1abc abc1", &["123", "45", " ", "6", " ", "1", "abc", " abc", "1"])]
#[case("x \n x \r\n \r\n y", &["x", " \n", " x", " \r\n \r\n", " y"])]
#[case("x \n ", &["x", " \n", " "])]
#[case("a b \n\n c", &["a", " ", " b", " \n\n", " ", " c"])]
#[case("x\t\ty x\t\t", &["x", "\t", "\ty", " x", "\t\t"])]
#[case("end ", &["end", " "])]
#[case("\u{a0}abc\u{a0}!", &["\u{a0}abc", "\u{a0}", "!"])]
#[case("<|endoftext|>", &["<|", "endoftext", "|>"])]
#[case("İstanbul ΣΊΣΥΦΟΣ Ελληνικά Русский", &["İstanbul", " ΣΊΣΥΦΟΣ", " Ελληνικά", " Русский"])]
#[case("日本語 ١٢٣٤", &["日本語", " ", "١٢٣", "٤"])]
fn scanner_splits_like_the_regex(#[case] text: &str, #[case] expected: &[&str]) {
assert_eq!(split(text), expected);
}
#[test]
fn every_scalar_alone_is_one_piece() {
for character in (0..=0x10FFFFu32).filter_map(char::from_u32) {
let text = character.to_string();
assert_eq!(
split(&text),
[text.as_str()],
"U+{:04X}",
u32::from(character)
);
}
}
#[test]
fn cases_match_oniguruma() {
let unicode_classes = UnicodeClasses::get().expect("Oniguruma exposes Unicode classes");
let upper = SysRegex::new(r"[\p{Lu}\p{Lt}\p{Lm}\p{Lo}\p{M}]").expect("regex");
let lower = SysRegex::new(r"[\p{Ll}\p{Lm}\p{Lo}\p{M}]").expect("regex");
let whole =
|regex: &SysRegex, text: &str| regex.find_iter(text).next() == Some((0, text.len()));
let mut text = String::new();
for character in (0..=0x10FFFFu32).filter_map(char::from_u32) {
text.clear();
text.push(character);
let expected = match (whole(&upper, &text), whole(&lower, &text)) {
(true, true) => Case::Both,
(true, false) => Case::Upper,
(false, true) => Case::Lower,
(false, false) => Case::Neither,
};
assert_eq!(
case(character, unicode_classes),
expected,
"U+{:04X}",
u32::from(character)
);
}
}
#[derive(serde::Deserialize)]
struct TextFixture {
text: String,
pieces: Vec<String>,
}
#[test]
fn scanner_splits_the_fixture_corpus_like_tiktoken_regex() {
let fixtures: Vec<TextFixture> = include_str!(concat!(
env!("CARGO_MANIFEST_DIR"),
"/tests/fixtures/o200k/texts.jsonl"
))
.lines()
.map(|line| serde_json::from_str(line).expect("fixture line is json"))
.collect();
assert!(fixtures.len() > 3000);
let mismatches: Vec<_> = fixtures
.iter()
.filter(|fixture| split(&fixture.text) != fixture.pieces)
.map(|fixture| (&fixture.text, split(&fixture.text), &fixture.pieces))
.collect();
assert!(mismatches.is_empty(), "{mismatches:#?}");
}
}

107
litellm-rust/crates/token-counter/src/scanner.rs Normal file
View file

@ -0,0 +1,107 @@
//! Exact token counting for tiktoken encodings. A hand-written scanner
//! reproduces the piece boundaries of the encoding's split regex, and each
//! piece is merged with the rank file. Special tokens are ordinary text, as
//! with `encode(text, disallowed_special=())`.
use std::iter;
use super::tiktoken::{MergeRanks, MergeScratch};
use super::unicode_classes::{Class, UnicodeClasses, class};
use super::{cl100k, o200k};
use crate::Error;
const MAX_DIGITS_PER_PIECE: usize = 3;
/// Byte length of the piece the split regex matches at the start of the
/// text, given the text's first character.
pub(super) type PieceLen = fn(&str, char, &UnicodeClasses) -> usize;
#[derive(Clone, Copy, Debug)]
pub(super) enum SplitPattern {
Cl100k,
O200k,
}
impl SplitPattern {
fn piece_len(self) -> PieceLen {
match self {
Self::Cl100k => cl100k::piece_len,
Self::O200k => o200k::piece_len,
}
}
}
pub(super) struct TiktokenCounter {
ranks: MergeRanks,
piece_len: PieceLen,
unicode_classes: &'static UnicodeClasses,
}
impl TiktokenCounter {
pub(super) fn from_ranks(split: SplitPattern, rank_file: &str) -> Result<Self, Error> {
Ok(Self {
ranks: MergeRanks::parse(rank_file)?,
piece_len: split.piece_len(),
unicode_classes: UnicodeClasses::get().ok_or(Error::UnicodeClasses)?,
})
}
pub(super) fn count(&self, text: &str) -> usize {
let mut scratch = MergeScratch::default();
pieces(text, self.piece_len, self.unicode_classes)
.map(|piece| self.ranks.count_piece(piece.as_bytes(), &mut scratch))
.sum()
}
}
/// The regex matches every character, so the pieces tile the text.
pub(super) fn pieces<'a>(
text: &'a str,
piece_len: PieceLen,
unicode_classes: &'static UnicodeClasses,
) -> impl Iterator<Item = &'a str> {
iter::successors(
split_piece(text, piece_len, unicode_classes),
move |(_, rest)| split_piece(rest, piece_len, unicode_classes),
)
.map(|(piece, _)| piece)
}
fn split_piece<'a>(
text: &'a str,
piece_len: PieceLen,
unicode_classes: &UnicodeClasses,
) -> Option<(&'a str, &'a str)> {
let first = text.chars().next()?;
Some(text.split_at(piece_len(text, first, unicode_classes)))
}
/// `'(?i:s|t|re|ve|m|ll|d)`, the contraction both encodings spell out. Simple
/// case folding also maps U+017F (long s) onto `s`.
pub(super) fn contraction_len(text: &str) -> Option<usize> {
let mut characters = text.chars();
if characters.next()? != '\'' {
return None;
}
let first = characters.next()?;
let len = match first {
's' | 'S' | '\u{17F}' | 'd' | 'D' | 'm' | 'M' | 't' | 'T' => first.len_utf8(),
'l' | 'L' => matches!(characters.next(), Some('l' | 'L')).then_some(2)?,
'v' | 'V' | 'r' | 'R' => matches!(characters.next(), Some('e' | 'E')).then_some(2)?,
_ => return None,
};
Some(1 + len)
}
pub(super) fn is_newline(character: char) -> bool {
matches!(character, '\r' | '\n')
}
/// `\p{N}{1,3}`
pub(super) fn digit_run_len(text: &str, unicode_classes: &UnicodeClasses) -> usize {
text.chars()
.take(MAX_DIGITS_PER_PIECE)
.take_while(|character| class(*character, unicode_classes) == Class::Number)
.map(char::len_utf8)
.sum()
}

215
litellm-rust/crates/token-counter/src/tiktoken.rs Normal file
View file

@ -0,0 +1,215 @@
//! tiktoken's byte-level BPE: a rank file of `base64(token) rank` lines and
//! the merge loop that turns one regex piece into tokens. The merge order is
//! tiktoken's (lowest rank first, leftmost pair on ties) so the token count is
//! identical, but pairs are tracked in a heap so a long piece costs
//! `O(n log n)` instead of tiktoken's `O(n^2)`.
use std::cmp::Reverse;
use std::collections::BinaryHeap;
use base64::Engine;
use base64::engine::general_purpose::STANDARD;
use rustc_hash::FxHashMap;
use crate::Error;
type Rank = u32;
const NO_RANK: Rank = Rank::MAX;
const END: usize = usize::MAX;
pub(super) struct MergeRanks(FxHashMap<Box<[u8]>, Rank>);
impl MergeRanks {
pub(super) fn parse(text: &str) -> Result<Self, Error> {
let ranks = text
.lines()
.filter(|line| !line.is_empty())
.map(parse_line)
.collect::<Result<FxHashMap<_, _>, _>>()?;
if let Some(byte) = (0..=u8::MAX).find(|byte| !ranks.contains_key(&[*byte][..])) {
return Err(Error::Ranks(format!("byte 0x{byte:02X} has no token")));
}
Ok(Self(ranks))
}
fn rank(&self, bytes: &[u8]) -> Rank {
self.0.get(bytes).copied().unwrap_or(NO_RANK)
}
/// Token count of one regex piece, as `encode_ordinary` would produce.
pub(super) fn count_piece(&self, piece: &[u8], scratch: &mut MergeScratch) -> usize {
if piece.len() < 2 || self.0.contains_key(piece) {
return 1;
}
scratch.reset(piece.len());
for start in 0..piece.len() - 1 {
scratch.set_rank(start, self.rank(&piece[start..start + 2]));
}
let mut parts = piece.len();
while let Some(Reverse((rank, start))) = scratch.heap.pop() {
if scratch.next[start] == END || scratch.rank[start] != rank {
continue;
}
let merged = scratch.next[start];
let after = scratch.next[merged];
scratch.next[merged] = END;
scratch.next[start] = after;
parts -= 1;
if after < piece.len() {
scratch.prev[after] = start;
scratch.set_rank(start, self.rank(&piece[start..scratch.end(after)]));
} else {
scratch.rank[start] = NO_RANK;
}
let before = scratch.prev[start];
if before != END {
scratch.set_rank(before, self.rank(&piece[before..scratch.end(start)]));
}
}
parts
}
}
fn parse_line(line: &str) -> Result<(Box<[u8]>, Rank), Error> {
let (token, rank) = line
.split_once(' ')
.ok_or_else(|| Error::Ranks(format!("line without a rank: {line:?}")))?;
let bytes = STANDARD
.decode(token)
.map_err(|error| Error::Ranks(format!("token is not base64: {error}")))?;
let rank = rank
.parse()
.map_err(|error| Error::Ranks(format!("rank is not an integer: {error}")))?;
Ok((bytes.into_boxed_slice(), rank))
}
/// Buffers reused across the pieces of one text. Parts are addressed by the
/// byte offset they start at, which also gives the leftmost-pair tie break.
#[derive(Default)]
pub(super) struct MergeScratch {
next: Vec<usize>,
prev: Vec<usize>,
rank: Vec<Rank>,
heap: BinaryHeap<Reverse<(Rank, usize)>>,
}
impl MergeScratch {
fn reset(&mut self, len: usize) {
self.next.clear();
self.next.extend(1..=len);
self.prev.clear();
self.prev.push(END);
self.prev.extend(0..len - 1);
self.rank.clear();
self.rank.resize(len, NO_RANK);
self.heap.clear();
}
fn end(&self, start: usize) -> usize {
self.next[start]
}
fn set_rank(&mut self, start: usize, rank: Rank) {
self.rank[start] = rank;
if rank != NO_RANK {
self.heap.push(Reverse((rank, start)));
}
}
}
#[cfg(test)]
mod tests {
use rand::rngs::StdRng;
use rand::{Rng, SeedableRng};
use super::*;
fn ranks() -> MergeRanks {
let path = concat!(
env!("CARGO_MANIFEST_DIR"),
"/../../../litellm/litellm_core_utils/tokenizers/9b5ad71b2ce5302211f9c61530b329a4922fc6a4"
);
MergeRanks::parse(&std::fs::read_to_string(path).expect("cl100k rank file is in the repo"))
.expect("rank file parses")
}
/// tiktoken's `_byte_pair_merge`, transcribed, as the reference.
fn reference_count(ranks: &MergeRanks, piece: &[u8]) -> usize {
if piece.len() < 2 || ranks.0.contains_key(piece) {
return 1;
}
let mut parts: Vec<(usize, Rank)> = (0..piece.len() - 1)
.map(|index| (index, ranks.rank(&piece[index..index + 2])))
.chain([(piece.len() - 1, NO_RANK), (piece.len(), NO_RANK)])
.collect();
let get_rank = |parts: &[(usize, Rank)], index: usize| {
if index + 3 < parts.len() {
ranks.rank(&piece[parts[index].0..parts[index + 3].0])
} else {
NO_RANK
}
};
loop {
let Some(index) = parts[..parts.len() - 1]
.iter()
.enumerate()
.filter(|(_, (_, rank))| *rank != NO_RANK)
.min_by_key(|(index, (_, rank))| (*rank, *index))
.map(|(index, _)| index)
else {
return parts.len() - 1;
};
if index > 0 {
parts[index - 1].1 = get_rank(&parts, index - 1);
}
parts[index].1 = get_rank(&parts, index);
parts.remove(index + 1);
}
}
#[test]
fn every_byte_is_a_token() {
let ranks = ranks();
assert_eq!(ranks.0.len(), 100_256);
assert!((0..=u8::MAX).all(|byte| ranks.rank(&[byte]) != NO_RANK));
}
#[test]
fn heap_merge_matches_tiktokens_merge_loop() {
let ranks = ranks();
let mut scratch = MergeScratch::default();
let mut rng = StdRng::seed_from_u64(99);
let alphabet = b" abcdeorstn.,'\n\xc3\xa9\xe2\x82\xac0123";
for _ in 0..20_000 {
let piece: Vec<u8> = (0..rng.gen_range(1..24))
.map(|_| alphabet[rng.gen_range(0..alphabet.len())])
.collect();
assert_eq!(
ranks.count_piece(&piece, &mut scratch),
reference_count(&ranks, &piece),
"piece {:?}",
String::from_utf8_lossy(&piece)
);
}
}
#[test]
fn long_repeated_runs_stay_cheap() {
let ranks = ranks();
let mut scratch = MergeScratch::default();
let piece = vec![b' '; 1 << 20];
let started = std::time::Instant::now();
let count = ranks.count_piece(&piece, &mut scratch);
assert!(count > 0);
assert!(started.elapsed().as_secs() < 5, "{:?}", started.elapsed());
}
#[test]
fn malformed_rank_files_are_rejected() {
assert!(MergeRanks::parse("IQ==").is_err());
assert!(MergeRanks::parse("IQ== x").is_err());
assert!(MergeRanks::parse("!!! 1").is_err());
assert!(MergeRanks::parse("IQ== 1").is_err());
}
}

97
litellm-rust/crates/token-counter/src/unicode_classes.rs
View file

@ -9,6 +9,8 @@ pub(super) struct UnicodeClasses {
letters: Ranges,
numbers: Ranges,
spaces: Ranges,
uppers: Ranges,
lowers: Ranges,
}
static CLASSES: LazyLock<Option<UnicodeClasses>> = LazyLock::new(|| {
@ -19,6 +21,8 @@ static CLASSES: LazyLock<Option<UnicodeClasses>> = LazyLock::new(|| {
letters: Ranges::load(r"\p{L}+", &scalars)?,
numbers: Ranges::load(r"\p{N}+", &scalars)?,
spaces: Ranges::load(r"\s+", &scalars)?,
uppers: Ranges::load(r"[\p{Lu}\p{Lt}\p{Lm}\p{Lo}\p{M}]+", &scalars)?,
lowers: Ranges::load(r"[\p{Ll}\p{Lm}\p{Lo}\p{M}]+", &scalars)?,
})
});
@ -59,15 +63,102 @@ impl UnicodeClasses {
CLASSES.as_ref()
}
pub(super) fn is_letter(&self, character: char) -> bool {
fn is_letter(&self, character: char) -> bool {
self.letters.contains(character)
}
pub(super) fn is_number(&self, character: char) -> bool {
fn is_number(&self, character: char) -> bool {
self.numbers.contains(character)
}
pub(super) fn is_space(&self, character: char) -> bool {
fn is_space(&self, character: char) -> bool {
self.spaces.contains(character)
}
fn is_upper(&self, character: char) -> bool {
self.uppers.contains(character)
}
fn is_lower(&self, character: char) -> bool {
self.lowers.contains(character)
}
}
/// `\p{L}`, `\p{N}`, `\s` and everything else, the character classes the
/// split regexes are written in.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub(super) enum Class {
Letter,
Number,
Space,
Other,
}
pub(super) fn class(character: char, unicode_classes: &UnicodeClasses) -> Class {
match character {
'A'..='Z' | 'a'..='z' => Class::Letter,
'0'..='9' => Class::Number,
'\t'..='\r' | ' ' => Class::Space,
_ if character.is_ascii() => Class::Other,
_ if unicode_classes.is_letter(character) => Class::Letter,
_ if unicode_classes.is_number(character) => Class::Number,
_ if unicode_classes.is_space(character) => Class::Space,
_ => Class::Other,
}
}
/// Byte length of the leading run of `run_class` characters.
pub(super) fn run_len(text: &str, run_class: Class, unicode_classes: &UnicodeClasses) -> usize {
text.char_indices()
.find(|(_, character)| class(*character, unicode_classes) != run_class)
.map_or(text.len(), |(index, _)| index)
}
/// Membership in the two letter classes of the o200k split regex,
/// `[\p{Lu}\p{Lt}\p{Lm}\p{Lo}\p{M}]` and `[\p{Ll}\p{Lm}\p{Lo}\p{M}]`; `Lm`,
/// `Lo` and `M` are in both.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub(super) enum Case {
Upper,
Lower,
Both,
Neither,
}
impl Case {
pub(super) fn is_upper(self) -> bool {
matches!(self, Case::Upper | Case::Both)
}
pub(super) fn is_lower(self) -> bool {
matches!(self, Case::Lower | Case::Both)
}
}
pub(super) fn case(character: char, unicode_classes: &UnicodeClasses) -> Case {
match character {
'A'..='Z' => Case::Upper,
'a'..='z' => Case::Lower,
_ if character.is_ascii() => Case::Neither,
_ => match (
unicode_classes.is_upper(character),
unicode_classes.is_lower(character),
) {
(true, true) => Case::Both,
(true, false) => Case::Upper,
(false, true) => Case::Lower,
(false, false) => Case::Neither,
},
}
}
/// Byte length of the leading run of characters whose case passes `in_class`.
pub(super) fn case_run_len(
text: &str,
in_class: fn(Case) -> bool,
unicode_classes: &UnicodeClasses,
) -> usize {
text.char_indices()
.find(|(_, character)| !in_class(case(*character, unicode_classes)))
.map_or(text.len(), |(index, _)| index)
}

10
litellm-rust/crates/token-counter/tests/fixtures/cl100k/requests.jsonl vendored Normal file
View file

File diff suppressed because one or more lines are too long

4053
litellm-rust/crates/token-counter/tests/fixtures/cl100k/texts.jsonl vendored Normal file
View file

File diff suppressed because one or more lines are too long

422
litellm-rust/crates/token-counter/tests/fixtures/generate.py vendored Normal file
View file

@ -0,0 +1,422 @@
"""Pin tiktoken reference counts for the Rust parity tests of one encoding.
Run from the repository root with the project environment, once per encoding:
uv run --no-sync python litellm-rust/crates/token-counter/tests/fixtures/generate.py cl100k_base
uv run --no-sync python litellm-rust/crates/token-counter/tests/fixtures/generate.py o200k_base
`<encoding>/texts.jsonl` holds `{"text", "tokens", "pieces"}` lines: `tokens`
counted with `tiktoken.get_encoding(name).encode(text, disallowed_special=())`,
the same call `litellm.token_counter` makes, and `pieces` the installed
encoding's split pattern applied with the `regex` module tiktoken itself uses,
so a scanner that splits differently fails even where BPE would count the same.
`<encoding>/requests.jsonl` holds `{"body", "input_tokens"}` lines, `body` being
the exact request bytes as a JSON string, counted with the proxy's admission
counter (`_count_input_tokens(body, model)`) for a model Python counts with that
encoding. Every message in the 50k-token body is shorter than the Python chunk
size so the chunked Python count equals the exact whole-text tiktoken count the
Rust counter produces.
"""
import itertools
import json
import random
import sys
from collections.abc import Iterator
from pathlib import Path
from typing import Final
import regex
import tiktoken
from litellm.constants import TIKTOKEN_ENCODE_CHUNK_SIZE_CHARS
from litellm.litellm_core_utils.token_counter import openai_tokenizer_encoding
from litellm.proxy.spend_tracking.budget_reservation import _count_input_tokens
HERE: Final = Path(__file__).resolve().parent
MODELS: Final = {"cl100k_base": "gpt-4", "o200k_base": "gpt-4o"}
ENCODING_NAME: Final = sys.argv[1]
MODEL: Final = MODELS[ENCODING_NAME]
ENCODING: Final = tiktoken.get_encoding(ENCODING_NAME)
assert openai_tokenizer_encoding(MODEL).name == ENCODING_NAME
SPLIT_PATTERN: Final = regex.compile(ENCODING._pat_str) # pyright: ignore[reportPrivateUsage] # tiktoken has no public accessor
OUT: Final = HERE / ENCODING_NAME.removesuffix("_base")
# Mirrors ALPHABET in src/byte_level.rs, plus the pieces the tiktoken patterns treat differently.
ALPHABET: Final = (
"a",
"Z",
"e",
"s",
"t",
"d",
"m",
"'",
"'s",
"'re",
"'ll",
"'S",
"0",
"9",
" ",
" ",
"\t",
"\n",
"\r\n",
"\x0b",
".",
",",
"!",
"-",
"(",
'"',
"\xa0",
"\x85",
"\u2028",
"\u3000",
"\u200b",
"\u200d",
"é",
"e\u0301",
"ß",
"漢",
"字",
"ع",
"३",
"½",
"Ⅳ",
"🙂",
"👍🏽",
"Ａ",
"fi",
"㍿",
"㋿",
"ꟲ",
"𐞁",
"a\u030a",
"\u1e0b\u0323",
"<",
">",
"EOT",
"<EOT>",
"<META_START>",
"'D",
"'M",
"'T",
"'VE",
"'Re",
"'ſ",
"ſ",
"12345678",
"٣٤٥٦",
"<|endoftext|>",
"<|fim_prefix|>",
"\r",
"\r\n\r\n",
" \n",
"!!",
"#$%",
"\u00ad",
"\u0301",
"\U0001f600\U0001f3fd",
"İ",
"Dž",
)
# The pieces the o200k case-shaped letter branch and slash-absorbing symbol branch split differently.
CASE_ALPHABET: Final = ALPHABET + (
"B",
"Ab",
"aB",
"ABC",
"ᵃ",
"camelCase",
"HTTPServer",
"iOS",
"Džungla",
"/",
"\n/",
"/\r\n",
" \n ",
"a/b",
)
CORPUS: Final = (
"",
"Hello, how are you today?",
"I'm sure they're right, we'll see. WE'LL SEE, I'M SURE THEY'RE RIGHT, IT'S HERS AND IT'D BE 'D",
"don't Don'T DON'T won'T i've I'VE i'Ve you'RE 'S 'T 'M 'D 'LL 'VE 'RE 'ſ 'x",
"1234567890 123 12 1 0000000 ٣٤٥٦٧٨ ३४५६ 1,234,567.89 2026-09-11T18:00:00Z",
"$abc %def &ghi @jkl _mno #pqr ~stu ^vwx |yz \\a /b :c ;d ?e !f (g )h [i ]j {k }l <m >n =o +p *q",
"foo bar baz \t qux\t\tquux \n\nline\r\nline\r\n\r\n \n\t\r\n x ",
"trailing spaces ",
"trailing tabs\t\t",
"trailing newline\n",
"\n\n\n",
"\r\n\r\n\r\n",
" ",
"😀😃😄 👍🏽 🇺🇸 👨‍👩‍👧‍👦 ✈️ ❤️‍🔥 ٭ ※ ⌘ ⏎",
"漢字かな交じり文、東京都千代田区。日本語のテキストです。中文测试。한국어 텍스트",
"مرحبا بالعالم، هذا نص عربي مع أرقام ١٢٣٤٥٦٧ و علامات ترقيم!",
"Zürich, façade, naïve, Ærøskøbing, Ελληνικά, Русский текст, עברית, हिन्दी, ไทย",
"e\u0301 a\u030a \u1e0b\u0323 \u0301\u0301 combining\u0308 marks\u0301!",
"ΣΊΣΥΦΟΣ Džungla İstanbul file flow Ａｂｃ ㍿ ㋿ ꟲ 𐞁",
"<|endoftext|> <|fim_prefix|>code<|fim_middle|>more<|fim_suffix|> <|endofprompt|> <|im_start|>",
"<EOT> <META_START> <s> </s> [INST] [/INST] <<SYS>>",
"def f(x):\n return {'a': x ** 2, \"b\": [1, 2, 3]} # comment\n\nprint(f(10))\n",
'{"model":"gpt-4","messages":[{"role":"user","content":"hi\\n"}],"temperature":0.7}',
"https://example.com/path?query=1&other=two#fragment user@example.com 192.168.0.1",
"a" * 3000,
" " * 3000,
"." * 3000,
"ab" * 1500,
"\n" * 3000,
"0" * 3000,
"!" * 3000,
"😀" * 1000,
"漢" * 1000,
"\u00a0abc\u00a0! \u2028x \u3000y \u200bz \u200d\u200d q",
"x\u0085y \x0b\x0c z",
"\x00\x01\x02 \x7f \ufffd",
"tab\tseparated\tvalues\n1\t2\t3\n",
"MiXeD cAsE wOrDs AND ACRONYMS like NASA, HTTP/2, gRPC, iOS, macOS",
"snake_case_identifier camelCaseIdentifier PascalCaseIdentifier SCREAMING_SNAKE_CASE kebab-case",
"x'sy x'ty x'rey x'vey x'my x'lly x'dy x'S x'T x'RE x'VE x'M x'LL x'D x'sS x'llL",
"IT'SOK it'Dbe x'Sy x'Ty x'My x'Dy x'LLy x'VEy x'REy x'Ly x'Vy x'Ry 'Sx'Tx'Mx'LLx'VEx'REx'Dx",
"'s't're've'm'll'd 'S'T'RE'VE'M'LL'D ''s '''s",
"9'9 9's a'9 '9 ' 's' ' 's",
"١٢٣٤ ½⅓¼ ⅣⅤ 𝟘𝟙𝟚𝟛𝟜𝟝𝟞𝟟𝟠𝟡 ①②③",
"camelCase PascalCase ABCdef ABCdeF ABC aB Ab ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyzABC",
"日本ABC ABC日本 日本語abc abc日本語 漢字Kanji kanji漢字 KANJI漢字kanji مرحباABC ABCمرحبا abcمرحبا",
"\u0301ABC \u0301abc \u0301\u0301A A\u0301\u0301 E\u0301A aE\u0301 !!\u0301a \u00a0\u0301A x\u0308Y X\u0308y",
"ᵃbc ᵃBC Aᵃbc Aᵃ ᵃ' ᵃ's Džungla aDžB ADžB ADžb DžDž Ljx İi ΣΊΣΥΦΟΣσ ΣσΣ",
"don'tx ABC's abc'S abc'ſ ABC'ſx IT'SOK it'Dbe 'sabc x's 's 'Sx'Tx x’s X'LLx X'Ll",
"!ABC !AbC !!abc #camelCase (ABCdef) \u00a0ABC\u00a0abc\u00a0Abc \tABC\tabc",
"!!/\n/x a/b !!\n/x /x // path/to/file.rs http://x.y/z?a=b/c \\/\\/ //\r\n//\n",
"x \n x \r\n \r\n y x \n a b \n\n c x\t\ty x\t\t end \n \n",
"12345 6 1abc abc1 ABC123abc 123ABC ١٢٣٤٥abc",
)
WORDS: Final = (
"the",
"quick",
"brown",
"fox",
"jumps",
"over",
"lazy",
"dog",
"while",
"counting",
"tokens",
"for",
"budget",
"reservation",
"before",
"admission",
"on",
"the",
"gateway",
"and",
"every",
"request",
"body",
"is",
"scanned",
"exactly",
"once",
"with",
"a",
"hand",
"written",
"piece",
"scanner",
"that",
"mirrors",
"tiktoken's",
"regex",
"boundaries",
"It's",
"faster",
"because",
"there's",
"no",
"backtracking",
"engine",
"involved",
"so",
"we'll",
"keep",
"it",
"that",
"way",
"Zürich",
"café",
"naïve",
"東京",
"مرحبا",
"🙂",
"42",
"1999",
"3.14159",
"$1,234.56",
"100%",
"user@example.com",
"https://example.com/a/b?c=d",
"C++",
"F#",
"node.js",
"v1.2.3",
"(parens)",
"[brackets]",
"{braces}",
"<tags>",
'"quotes"',
"'single'",
"don't",
"WON'T",
"I'M",
"They'RE",
)
def random_text(rng: random.Random, alphabet: tuple[str, ...]) -> str:
return "".join(rng.choice(alphabet) for _ in range(rng.randrange(0, 40)))
def paragraph(rng: random.Random, words: int) -> str:
return " ".join(rng.choice(WORDS) for _ in range(words))
def short_paragraphs(rng: random.Random) -> Iterator[str]:
while True:
content = paragraph(rng, rng.randrange(60, 140))
if len(content) < TIKTOKEN_ENCODE_CHUNK_SIZE_CHARS:
yield content
def chat_body(rng: random.Random, target_tokens: int) -> dict[str, object]:
candidates: Final = tuple(itertools.islice(short_paragraphs(rng), 2000))
running: Final = tuple(itertools.accumulate(len(ENCODING.encode(content)) + 3 for content in candidates))
turns: Final = next(index for index, total in enumerate(running) if total >= target_tokens) + 1
contents: Final = candidates[: turns + (turns % 2)]
return {
"model": MODEL,
"messages": [
{"role": "system", "content": "You are a helpful assistant. Answer precisely and cite sources."},
*(
{"role": "user" if index % 2 == 0 else "assistant", "content": content}
for index, content in enumerate(contents)
),
{"role": "user", "content": "Summarise the conversation so far in three sentences."},
],
}
TOOLS: Final = [
{
"type": "function",
"function": {
"name": "get_weather",
"description": "Get weather",
"parameters": {
"type": "object",
"properties": {
"location": {"type": "string", "description": "City name"},
"unit": {"type": "string", "enum": ["celsius", "fahrenheit"]},
"days": {"type": "integer"},
"tags": {"type": "array", "items": {"type": "string"}},
"opts": {
"type": "object",
"properties": {"verbose": {"type": "boolean"}, "level": {"type": "integer", "enum": [1, 2]}},
"required": ["verbose"],
},
"anything": {},
},
"required": ["location"],
},
},
},
{"type": "function", "function": {"name": "noop"}},
]
SMALL_REQUESTS: Final = (
{"model": MODEL, "messages": [{"role": "user", "content": "Hello, how are you today?"}]},
{
"model": MODEL,
"messages": [
{"role": "system", "content": "You are a terse assistant."},
{
"role": "user",
"name": "alice",
"content": [
{"type": "text", "text": "Summarise this paragraph about ships and harbours."},
"plain string item",
],
},
{"role": "assistant", "content": [{"type": "text", "text": "Sure."}]},
],
},
{
"model": MODEL,
"messages": [{"role": "user", "content": "weather?"}],
"tools": TOOLS,
"tool_choice": {"type": "function", "function": {"name": "get_weather"}},
},
{
"model": MODEL,
"messages": [{"role": "system", "content": "sys"}, {"role": "user", "content": "weather?"}],
"tools": TOOLS,
"tool_choice": "none",
},
{"model": MODEL, "prompt": "Write a haiku about ships."},
{"model": MODEL, "prompt": ["first prompt", "second prompt"]},
{
"model": MODEL,
"input": [
{
"role": "user",
"content": [
{"type": "input_text", "text": 'Summarise caf\u00e9 menus, na\u00efve \u2014 ok? "quoted"\n'}
],
},
{"role": "assistant", "content": "Sure."},
],
"instructions": "be terse",
},
{"model": MODEL, "input": [[101, 2023, 5], [7]], "encoding_format": "float"},
{
"model": MODEL,
"query": "best harbour",
"documents": [
"doc one",
{"text": "doc two", "title": "T", "n": 3, "ok": True, "none": None, "tags": ["a", "b"]},
],
},
)
def main() -> None:
rng: Final = random.Random(2026)
case_rng: Final = random.Random(200_000)
texts: Final = (
tuple(CORPUS)
+ tuple(random_text(rng, ALPHABET) for _ in range(3000))
+ tuple(random_text(case_rng, CASE_ALPHABET) for _ in range(1000))
)
OUT.mkdir(exist_ok=True)
with (OUT / "texts.jsonl").open("w", encoding="utf-8") as handle:
for text in texts:
tokens = len(ENCODING.encode(text, disallowed_special=()))
pieces = SPLIT_PATTERN.findall(text)
handle.write(json.dumps({"text": text, "tokens": tokens, "pieces": pieces}, ensure_ascii=False) + "\n")
bodies: Final = tuple(SMALL_REQUESTS) + (chat_body(rng, 50_000),)
with (OUT / "requests.jsonl").open("w", encoding="utf-8") as handle:
for body in bodies:
input_tokens = _count_input_tokens(dict(body), MODEL)
assert input_tokens is not None
handle.write(json.dumps({"body": json.dumps(body), "input_tokens": input_tokens}) + "\n")
if __name__ == "__main__":
main()

10
litellm-rust/crates/token-counter/tests/fixtures/o200k/requests.jsonl vendored Normal file
View file

File diff suppressed because one or more lines are too long

4053
litellm-rust/crates/token-counter/tests/fixtures/o200k/texts.jsonl vendored Normal file
View file

File diff suppressed because one or more lines are too long

145
litellm-rust/crates/token-counter/tests/token_counter.rs
View file

@ -1,4 +1,5 @@
use rstest::rstest;
use serde::Deserialize;
use litellm_token_counter::{CountableRequest, Error, InputTokenCount, TokenCounter};
@ -174,3 +175,147 @@ fn tool_choice_and_system_discount_change_the_count() {
fn loading_a_bad_tokenizer_is_a_load_error() {
assert!(matches!(TokenCounter::from_json("{}"), Err(Error::Load(_))));
}
/// A tiktoken encoding: its fixture directory, the vendored rank file Python
/// loads, the constructor, and the model `generate.py` counted the requests for.
#[derive(Clone, Copy)]
struct TiktokenEncoding {
fixtures: &'static str,
rank_file: &'static str,
load: fn(&str) -> Result<TokenCounter, Error>,
model: &'static str,
}
const CL100K: TiktokenEncoding = TiktokenEncoding {
fixtures: "cl100k",
rank_file: "9b5ad71b2ce5302211f9c61530b329a4922fc6a4",
load: TokenCounter::from_cl100k_ranks,
model: "gpt-4",
};
const O200K: TiktokenEncoding = TiktokenEncoding {
fixtures: "o200k",
rank_file: "fb374d419588a4632f3f557e76b4b70aebbca790",
load: TokenCounter::from_o200k_ranks,
model: "gpt-4o",
};
fn tiktoken_counter(encoding: TiktokenEncoding) -> TokenCounter {
let path = format!(
"{}/../../../litellm/litellm_core_utils/tokenizers/{}",
env!("CARGO_MANIFEST_DIR"),
encoding.rank_file
);
let ranks = std::fs::read_to_string(&path).expect("rank file is in the repo");
(encoding.load)(&ranks).expect("ranks load")
}
fn tiktoken_fixture(encoding: TiktokenEncoding, name: &str) -> String {
let path = format!(
"{}/tests/fixtures/{}/{name}",
env!("CARGO_MANIFEST_DIR"),
encoding.fixtures
);
std::fs::read_to_string(&path).expect("fixture generated by tests/fixtures/generate.py")
}
#[derive(Deserialize)]
struct TextFixture {
text: String,
tokens: usize,
}
#[derive(Deserialize)]
struct RequestFixture {
body: String,
input_tokens: usize,
}
/// Reference counts come from `tiktoken.get_encoding(name)`; see
/// `tests/fixtures/generate.py`.
#[rstest]
#[case::cl100k(CL100K)]
#[case::o200k(O200K)]
fn tiktoken_text_counts_match_tiktoken(#[case] encoding: TiktokenEncoding) {
let counter = tiktoken_counter(encoding);
let fixtures: Vec<TextFixture> = tiktoken_fixture(encoding, "texts.jsonl")
.lines()
.map(|line| serde_json::from_str(line).expect("fixture line is json"))
.collect();
assert!(fixtures.len() > 3000);
let mismatches: Vec<_> = fixtures
.iter()
.filter_map(|fixture| {
let count = counter.count_text(&fixture.text).expect("text counts");
(count != fixture.tokens).then(|| (fixture.text.clone(), fixture.tokens, count))
})
.collect();
assert!(
mismatches.is_empty(),
"(text, tiktoken, rust): {mismatches:?}"
);
}
/// Reference counts come from the proxy's admission counter
/// (`_count_input_tokens(body, model)`), so this pins the shared message,
/// tool and reply-priming accounting on the tiktoken paths as well.
#[rstest]
#[case::cl100k(CL100K)]
#[case::o200k(O200K)]
fn tiktoken_request_counts_match_python_admission_counter(#[case] encoding: TiktokenEncoding) {
let counter = tiktoken_counter(encoding);
let fixtures: Vec<RequestFixture> = tiktoken_fixture(encoding, "requests.jsonl")
.lines()
.map(|line| serde_json::from_str(line).expect("fixture line is json"))
.collect();
let counts: Vec<usize> = fixtures
.iter()
.map(|fixture| {
let request = CountableRequest::parse(fixture.body.as_bytes()).expect("fixture parses");
let count = counter.count_request(&request).expect("fixture counts");
assert_eq!(count.model.as_deref(), Some(encoding.model));
assert_eq!(count.input_tokens, fixture.input_tokens, "{}", fixture.body);
count.input_tokens
})
.collect();
assert!(counts.last().is_some_and(|tokens| *tokens >= 50_000));
}
#[rstest]
#[case::cl100k(CL100K)]
#[case::o200k(O200K)]
fn tiktoken_shares_the_message_accounting_with_the_anthropic_path(
#[case] encoding: TiktokenEncoding,
) {
let counter = tiktoken_counter(encoding);
let count = |body: &str| {
counter
.count_request(&CountableRequest::parse(body.as_bytes()).expect("parses"))
.expect("counts")
.input_tokens
};
let text = |text: &str| counter.count_text(text).expect("counts");
let base = count(r#"{"model":"m","messages":[{"role":"user","content":"hi"}]}"#);
assert_eq!(base, 3 + text("user") + text("hi") + 3);
assert_eq!(
count(r#"{"model":"m","messages":[{"role":"user","name":"al","content":"hi"}]}"#),
base + text("al") + 1
);
assert_eq!(
count(r#"{"model":"m","messages":[{"role":"user","content":"hi"}],"tool_choice":"none"}"#),
base + 1
);
}
#[rstest]
#[case::empty("")]
#[case::not_base64("!!!! 0")]
#[case::missing_rank("YQ==")]
#[case::rank_not_a_number("YQ== x")]
#[case::single_byte_tokens_missing("YWI= 0")]
fn loading_a_bad_rank_file_is_a_load_error(
#[case] rank_file: &str,
#[values(CL100K, O200K)] encoding: TiktokenEncoding,
) {
assert!(matches!((encoding.load)(rank_file), Err(Error::Ranks(_))));
}

92
litellm/integrations/custom_guardrail.py
View file

@ -1130,7 +1130,7 @@ class CustomGuardrail(CustomLogger):
def add_standard_logging_guardrail_information_to_request_data(
self,
guardrail_json_response: Exception | str | dict | list[dict],
guardrail_json_response: object,
request_data: dict,
guardrail_status: GuardrailStatus,
start_time: float | None = None,
@ -1275,17 +1275,10 @@ class CustomGuardrail(CustomLogger):
This gets logged on downsteam Langfuse, DataDog, etc.
"""
# Convert None to empty dict to satisfy type requirements
guardrail_response: dict[str, object] | str = {} if response is None else response
# For apply_guardrail functions in custom_code_guardrail scenario,
# simplify the logged response to "allow", "deny", or "mask"
if original_inputs is not None and isinstance(response, dict):
# Check if inputs were modified by comparing them
if self._inputs_were_modified(original_inputs, response):
guardrail_response = "mask"
else:
guardrail_response = "allow"
guardrail_response: Final = self._summarize_guardrail_response(
response=response,
original_inputs=original_inputs,
)
verbose_logger.debug("Guardrail response: %s", response)
@ -1300,6 +1293,27 @@ class CustomGuardrail(CustomLogger):
)
return response
def _summarize_guardrail_response(
self,
response: object,
original_inputs: Mapping[str, object] | None,
) -> object:
"""Reduce a hook's return value to what is safe to log as ``guardrail_response``.
``apply_guardrail`` returns the (possibly masked) inputs and ``async_pre_call_hook``
returns the (possibly modified) request payload. Neither is a provider verdict, and
logging them verbatim ships the user's prompt to every logging sink (OTEL spans,
Datadog, spend logs), so both collapse to ``"allow"`` / ``"mask"`` by comparing
against ``original_inputs``, a copy taken before the hook ran. A string result is the
hook's own rejection message (the proxy turns it into a 400), not user input, so it is
logged as is.
"""
if response is None:
return {}
if original_inputs is None or not isinstance(response, Mapping):
return response
return "mask" if self._inputs_were_modified(original_inputs, response) else "allow"
@staticmethod
def _is_guardrail_intervention(e: Exception) -> bool:
"""Retained spelling for existing callers; prefer ``is_guardrail_intervention``."""
@ -1339,24 +1353,9 @@ class CustomGuardrail(CustomLogger):
)
raise e
def _inputs_were_modified(self, original_inputs: dict, response: dict) -> bool:
"""
Compare original inputs with response to determine if content was modified.
Returns True if the inputs were modified (mask scenario), False otherwise (allow scenario).
"""
# Get all keys from both dictionaries
all_keys: Final = set(original_inputs.keys()) | set(response.keys())
# Compare each key's value
for key in all_keys:
original_value = original_inputs.get(key)
response_value = response.get(key)
if original_value != response_value:
return True
# No modifications detected
return False
def _inputs_were_modified(self, original_inputs: Mapping[str, object], response: Mapping[str, object]) -> bool:
"""True when any baseline key's value differs in ``response`` (mask), False otherwise (allow)."""
return any(response.get(key) != value for key, value in original_inputs.items())
def mask_content_in_string(
self,
@ -1463,6 +1462,31 @@ def _sync_guardrail_info_to_logging_obj(request_data: dict, logging_obj: object)
_append_slg_to_litellm_params(mcd.get("litellm_params"), entries)
_PRE_CALL_CONTENT_KEYS: Final = frozenset(
{"messages", "input", "prompt", "system", "instructions", "tools", "functions", "function_call", "tool_choice"}
)
def _original_inputs_for(
func_name: str,
kwargs: Mapping[str, object],
request_data: Mapping[str, object],
event_type: GuardrailEventHooks | None,
) -> dict | None: # mutable-ok: matches _process_response(original_inputs=) signature
"""Baseline the hook's return value is compared against to decide "allow" vs "mask".
``apply_guardrail`` masks a fresh ``inputs`` dict, so that dict is the baseline. Pre-call
hooks edit the request in place and return it, so the baseline is a deep copy of the
prompt-bearing keys taken before the hook runs.
"""
if func_name == "apply_guardrail":
inputs: Final = kwargs.get("inputs")
return inputs if isinstance(inputs, dict) else None
if event_type != GuardrailEventHooks.pre_call:
return None
return {key: copy.deepcopy(value) for key, value in request_data.items() if key in _PRE_CALL_CONTENT_KEYS}
def log_guardrail_information(func):
"""
Decorator to add standard logging guardrail information to any function
@ -1521,9 +1545,7 @@ def log_guardrail_information(func):
event_type: Final = _infer_event_type_from_function_name(func.__name__)
# Store original inputs for comparison (for apply_guardrail functions)
original_inputs = None
if func.__name__ == "apply_guardrail" and "inputs" in kwargs:
original_inputs = kwargs.get("inputs")
original_inputs: Final = _original_inputs_for(func.__name__, kwargs, request_data, event_type)
logging_obj: Final = kwargs.get("logging_obj") or request_data.get("litellm_logging_obj")
self_recorded_token: Final = _guardrail_self_recorded.set(False)
@ -1563,9 +1585,7 @@ def log_guardrail_information(func):
event_type: Final = _infer_event_type_from_function_name(func.__name__)
# Store original inputs for comparison (for apply_guardrail functions)
original_inputs = None
if func.__name__ == "apply_guardrail" and "inputs" in kwargs:
original_inputs = kwargs.get("inputs")
original_inputs: Final = _original_inputs_for(func.__name__, kwargs, request_data, event_type)
logging_obj: Final = kwargs.get("logging_obj") or request_data.get("litellm_logging_obj")
self_recorded_token: Final = _guardrail_self_recorded.set(False)

15
litellm/litellm_core_utils/default_encoding.py
View file

@ -1,4 +1,5 @@
import os
from pathlib import Path
from typing import Final
import litellm
@ -14,6 +15,20 @@ except (ImportError, AttributeError):
filename = pkg_resources.resource_filename(__name__, "litellm_core_utils/tokenizers")
CL100K_BASE_RANK_FILE: Final = "9b5ad71b2ce5302211f9c61530b329a4922fc6a4"
O200K_BASE_RANK_FILE: Final = "fb374d419588a4632f3f557e76b4b70aebbca790"
def cl100k_base_rank_file() -> str:
"""The vendored tiktoken `cl100k_base` rank file (`base64(token) rank` lines)."""
return Path(filename, CL100K_BASE_RANK_FILE).read_text(encoding="ascii")
def o200k_base_rank_file() -> str:
"""The vendored tiktoken `o200k_base` rank file (`base64(token) rank` lines)."""
return Path(filename, O200K_BASE_RANK_FILE).read_text(encoding="ascii")
# Always default TIKTOKEN_CACHE_DIR to the bundled tokenizers directory
# unless the user explicitly overrides it via CUSTOM_TIKTOKEN_CACHE_DIR.
# This keeps tiktoken fully offline-capable by default (see #1071).

33
litellm/litellm_core_utils/token_counter.py
View file

@ -381,7 +381,7 @@ class _MessageCountParams:
from litellm.utils import print_verbose
actual_model: Final = _fix_model_name(model)
if actual_model == "gpt-3.5-turbo-0301":
if uses_legacy_message_accounting(model):
self.tokens_per_message = 4 # every message follows <|start|>{role/name}\n{content}<|end|>\n
self.tokens_per_name = -1 # if there's a name, the role is omitted
elif actual_model in litellm.open_ai_chat_completion_models or actual_model in litellm.azure_llms:
@ -615,7 +615,7 @@ def _get_exact_count_function(
) -> TokenCounterFunction:
"""
Get the function to count tokens based on the model and custom tokenizer."""
from litellm.utils import _select_tokenizer, print_verbose
from litellm.utils import _select_tokenizer
if model is not None or custom_tokenizer is not None:
tokenizer_json: Final = custom_tokenizer or _select_tokenizer(model)
@ -627,15 +627,7 @@ def _get_exact_count_function(
return count_tokens
elif tokenizer_json["type"] == "openai_tokenizer":
model_to_use: Final = _fix_model_name(model)
try:
if "gpt-4o" in model_to_use:
encoding = tiktoken.get_encoding("o200k_base")
else:
encoding = tiktoken.encoding_for_model(model_to_use)
except KeyError:
print_verbose("Warning: model not found. Using cl100k_base encoding.")
encoding = tiktoken.get_encoding("cl100k_base")
encoding: Final = openai_tokenizer_encoding(model)
def encode_length(text: str) -> int:
return len(encoding.encode(text, disallowed_special=()))
@ -651,6 +643,25 @@ def _get_exact_count_function(
return _get_tiktoken_count_function(encode_length)
def openai_tokenizer_encoding(model: str) -> tiktoken.Encoding:
"""The tiktoken encoding `token_counter` uses for a model on the `openai_tokenizer` path."""
from litellm.utils import print_verbose
model_to_use: Final = _fix_model_name(model)
if "gpt-4o" in model_to_use:
return tiktoken.get_encoding("o200k_base")
try:
return tiktoken.encoding_for_model(model_to_use)
except KeyError:
print_verbose("Warning: model not found. Using cl100k_base encoding.")
return tiktoken.get_encoding("cl100k_base")
def uses_legacy_message_accounting(model: str) -> bool:
"""Whether `token_counter` prices messages with the `gpt-3.5-turbo-0301` constants (4 per message, -1 per name)."""
return _fix_model_name(model) == "gpt-3.5-turbo-0301"
def _fix_model_name(model: str) -> str:
"""We normalize some model names to others"""
if model in litellm.azure_llms:

7
litellm/proxy/guardrails/guardrail_hooks/azure/prompt_shield.py
View file

@ -338,10 +338,9 @@ class AzureContentSafetyPromptShieldGuardrail(AzureGuardrailBase, CustomGuardrai
estimated cost) and the ``azure`` provider label to the recorded guardrail
information. Follows the OpenAI moderation override pattern
(openai/moderations.py)."""
guardrail_response: Final[dict | str] = ( # mutable-ok: mirrors CustomGuardrail._process_response
("mask" if self._inputs_were_modified(original_inputs, response) else "allow")
if original_inputs is not None and isinstance(response, dict)
else ({} if response is None else response) # mutable-ok: empty placeholder, never mutated
guardrail_response: Final = self._summarize_guardrail_response(
response=response,
original_inputs=original_inputs,
)
self.add_standard_logging_guardrail_information_to_request_data(
guardrail_json_response=guardrail_response,

36
litellm/proxy/spend_tracking/budget_reservation.py
View file

@ -34,7 +34,7 @@ from litellm.proxy.common_utils.user_api_key_cache import (
)
from litellm.proxy.utils import PrismaClient, ProxyLogging
from litellm.router import Router
from litellm.rust_bridge.token_counter import count_anthropic_input_tokens, uses_anthropic_tokenizer
from litellm.rust_bridge.token_counter import RustTokenizer, count_input_tokens, rust_tokenizer
from litellm.types.proxy.model_access_group_budget import ModelAccessGroupBudget
from litellm.types.router import DeploymentTypedDict
@ -1365,8 +1365,9 @@ async def count_request_input_tokens(
Tokenizing is the reservation path's dominant CPU cost and is O(prompt), so
counting a large prompt inline stalls every other request on the worker.
Models on the Anthropic tokenizer are counted from the raw body by the Rust
bridge when it is enabled, which parses and tokenizes with the GIL released.
Models whose tokenizer the Rust bridge ports (Anthropic, tiktoken cl100k_base
and o200k_base) are counted from the raw body by the bridge when it is enabled, once per
distinct tokenizer, which parses and tokenizes with the GIL released.
Everything it declines is counted in Python, large prompts in a worker
thread. The counts are reused by both the max-cost and the input-cost
estimate.
@ -1374,23 +1375,31 @@ async def count_request_input_tokens(
models: Final = _get_request_models(request_body=request_body, route=route, llm_router=llm_router)
if not models:
return MappingProxyType({})
rust_count: Final = (
await count_anthropic_input_tokens(raw_body)
if raw_body is not None and any(uses_anthropic_tokenizer(model) for model in models)
else None
tokenizers: Final[Mapping[str, RustTokenizer | None]] = MappingProxyType(
{model: rust_tokenizer(model) for model in models}
)
distinct_tokenizers: Final[tuple[RustTokenizer, ...]] = tuple(
dict.fromkeys(tokenizer for tokenizer in tokenizers.values() if tokenizer is not None)
)
rust_counts_by_tokenizer: Final[Mapping[RustTokenizer, int]] = MappingProxyType(
{
tokenizer: count.input_tokens
for tokenizer in distinct_tokenizers
if raw_body is not None and (count := await count_input_tokens(raw_body, tokenizer)) is not None
}
)
rust_counts: Final = MappingProxyType(
{
model: rust_count.input_tokens
for model in models
if rust_count is not None and uses_anthropic_tokenizer(model)
model: rust_counts_by_tokenizer[tokenizer]
for model, tokenizer in tokenizers.items()
if tokenizer is not None and tokenizer in rust_counts_by_tokenizer
}
)
python_models: Final = tuple(model for model in models if model not in rust_counts)
if not python_models:
return rust_counts
python_counts: Final = (
_count_input_tokens_for_models(request_body=request_body, models=python_models)
MappingProxyType({})
if not python_models
else _count_input_tokens_for_models(request_body=request_body, models=python_models)
if _approximate_input_size(request_body) < TOKENIZE_OFF_EVENT_LOOP_MIN_CHARS
else await asyncio.to_thread(
_count_input_tokens_for_models,
@ -1398,6 +1407,7 @@ async def count_request_input_tokens(
models=python_models,
)
)
verbose_proxy_logger.debug("input token counts: rust=%s python=%s", dict(rust_counts), dict(python_counts))
return MappingProxyType({**rust_counts, **python_counts})

63
litellm/rust_bridge/token_counter.py
View file

@ -5,17 +5,20 @@ from __future__ import annotations
from collections.abc import Awaitable
from dataclasses import dataclass
from functools import lru_cache
from typing import Final, Protocol, cast # noqa: TID251 # native extension exposes dynamically typed callables
from typing import Final, Literal, Protocol, cast # noqa: TID251 # native extension exposes untyped callables
from pydantic import TypeAdapter
import litellm
from litellm._logging import verbose_logger
from litellm.litellm_core_utils.default_encoding import cl100k_base_rank_file, o200k_base_rank_file
from litellm.litellm_core_utils.token_counter import openai_tokenizer_encoding, uses_legacy_message_accounting
from litellm.rust_bridge.bindings import NativeBinding
from litellm.rust_bridge.configuration import rust_enabled
from litellm.rust_bridge.runtime import BridgeErrorContext, RustHandled, aattempt
from litellm.utils import claude_json_str
from litellm.utils import uses_anthropic_tokenizer as _python_uses_anthropic_tokenizer
from litellm.utils import claude_json_str, huggingface_tokenizer_kind
RustTokenizer = Literal["anthropic", "cl100k_base", "o200k_base"]
class RustTokenCounter(Protocol):
@ -27,6 +30,12 @@ class RustTokenCounterFactory(Protocol):
def __call__(self, tokenizer_json: str) -> RustTokenCounter:
raise NotImplementedError
def from_cl100k_ranks(self, rank_file: str) -> RustTokenCounter:
raise NotImplementedError
def from_o200k_ranks(self, rank_file: str) -> RustTokenCounter:
raise NotImplementedError
@dataclass(frozen=True, slots=True)
class InputTokenCount:
@ -50,18 +59,41 @@ def _as_factory(value: object) -> RustTokenCounterFactory | None:
TOKEN_COUNTER: Final = NativeBinding("TokenCounter", validate=_as_factory)
def uses_anthropic_tokenizer(model: str) -> bool:
if litellm.disable_token_counter is True or litellm.disable_hf_tokenizer_download is True:
return False
return _python_uses_anthropic_tokenizer(model)
def rust_tokenizer(model: str) -> RustTokenizer | None:
"""The Rust counter for the tokenizer `litellm.token_counter` selects for `model`, `None` when Python must count.
Mirrors `_select_tokenizer_helper`: the Anthropic tokenizer has a Rust port, the other HuggingFace
downloads do not, and of the tiktoken encodings `cl100k_base` and `o200k_base` do (p50k/r50k do not). Rust
prices every message with the default constants, so the legacy `gpt-3.5-turbo-0301` accounting stays in
Python."""
if litellm.disable_token_counter is True:
return None
kind: Final = None if litellm.disable_hf_tokenizer_download is True else huggingface_tokenizer_kind(model)
if kind == "anthropic":
return "anthropic"
if kind is not None or uses_legacy_message_accounting(model):
return None
match openai_tokenizer_encoding(model).name:
case "cl100k_base":
return "cl100k_base"
case "o200k_base":
return "o200k_base"
case _:
return None
@lru_cache(maxsize=4)
def _anthropic_counter(factory: RustTokenCounterFactory) -> RustTokenCounter:
return factory(claude_json_str)
def _counter(factory: RustTokenCounterFactory, tokenizer: RustTokenizer) -> RustTokenCounter:
match tokenizer:
case "anthropic":
return factory(claude_json_str)
case "cl100k_base":
return factory.from_cl100k_ranks(cl100k_base_rank_file())
case "o200k_base":
return factory.from_o200k_ranks(o200k_base_rank_file())
async def count_anthropic_input_tokens(body: bytes) -> InputTokenCount | None:
async def count_input_tokens(body: bytes, tokenizer: RustTokenizer) -> InputTokenCount | None:
if not rust_enabled():
return None
factory: Final = TOKEN_COUNTER.load()
@ -69,11 +101,14 @@ async def count_anthropic_input_tokens(body: bytes) -> InputTokenCount | None:
return None
try:
attempt: Final = await aattempt(
native_call=lambda: _anthropic_counter(factory).acount_request(body),
native_call=lambda: _counter(factory, tokenizer).acount_request(body),
adapt=_INPUT_TOKEN_COUNT.validate_python,
context=BridgeErrorContext(route="token_counter", provider="anthropic", model=""),
context=BridgeErrorContext(route="token_counter", provider=tokenizer, model=""),
)
except (RuntimeError, ValueError) as error:
verbose_logger.debug("Rust token counter failed, counting in Python: %s", error)
verbose_logger.debug("Rust token counter (%s) failed, counting in Python: %s", tokenizer, error)
return None
return attempt.value if isinstance(attempt, RustHandled) else None
if not isinstance(attempt, RustHandled):
return None
verbose_logger.debug("Rust token counter (%s) counted %d input tokens", tokenizer, attempt.value.input_tokens)
return attempt.value

48
litellm/utils.py
View file

@ -2227,25 +2227,39 @@ def uses_anthropic_tokenizer(model: str) -> bool:
return model in litellm.anthropic_models and "claude-3" not in model
def _return_huggingface_tokenizer(model: str) -> SelectTokenizerResponse | None:
HuggingFaceTokenizerKind = Literal["cohere", "anthropic", "llama2", "llama3"]
def huggingface_tokenizer_kind(model: str) -> HuggingFaceTokenizerKind | None:
"""Which HuggingFace tokenizer `token_counter` selects for a model; `None` means tiktoken."""
if model in litellm.cohere_models and "command-r" in model:
# cohere
cohere_tokenizer: Final = Tokenizer.from_pretrained("Xenova/c4ai-command-r-v01-tokenizer")
return {"type": "huggingface_tokenizer", "tokenizer": cohere_tokenizer}
# anthropic
elif uses_anthropic_tokenizer(model):
claude_tokenizer: Final = Tokenizer.from_str(claude_json_str)
return {"type": "huggingface_tokenizer", "tokenizer": claude_tokenizer}
# llama2
elif "llama-2" in model.lower() or "replicate" in model.lower():
tokenizer = Tokenizer.from_pretrained("hf-internal-testing/llama-tokenizer")
return {"type": "huggingface_tokenizer", "tokenizer": tokenizer}
# llama3
elif "llama-3" in model.lower():
tokenizer = Tokenizer.from_pretrained("Xenova/llama-3-tokenizer")
return {"type": "huggingface_tokenizer", "tokenizer": tokenizer}
else:
return "cohere"
if uses_anthropic_tokenizer(model):
return "anthropic"
if "llama-2" in model.lower() or "replicate" in model.lower():
return "llama2"
if "llama-3" in model.lower():
return "llama3"
return None
def _return_huggingface_tokenizer(model: str) -> SelectTokenizerResponse | None:
kind: Final = huggingface_tokenizer_kind(model)
if kind is None:
return None
return {"type": "huggingface_tokenizer", "tokenizer": _load_huggingface_tokenizer(kind)}
def _load_huggingface_tokenizer(kind: HuggingFaceTokenizerKind) -> Tokenizer:
match kind:
case "cohere":
return Tokenizer.from_pretrained("Xenova/c4ai-command-r-v01-tokenizer")
case "anthropic":
return Tokenizer.from_str(claude_json_str)
case "llama2":
return Tokenizer.from_pretrained("hf-internal-testing/llama-tokenizer")
case "llama3":
return Tokenizer.from_pretrained("Xenova/llama-3-tokenizer")
def encode(model="", text="", custom_tokenizer: dict | None = None):

1
pyproject.toml
View file

@ -220,6 +220,7 @@ e2e-dev = [
"playwright==1.61.0",
"websockets>=15.0.1,<16.0",
"locust==2.45.0",
"psutil==7.2.2",
"mcp>=1.28.1,<2.0",
]
proxy-dev = [

2
tests/e2e/CLAUDE.md
View file

@ -18,7 +18,7 @@ Each subdirectory under `tests/e2e/` is one suite, scoped to an endpoint family
- `logging/` - logging-integration delivery (datadog and friends)
- `security/` - secret handling and log-leak protection
- `router/` - routing and reliability behavior (fallbacks, cooldowns)
- `load/` - performance-category tests, kept OUT of the main suite: throughput/load SLO tests are a different testing category from functional e2e (variance-driven, historically flaky) and live outside this suite until re-implemented as their own pipeline (LIT-5163); do not add a live load test that runs in the default collection. What remains here: the weekly session-anomaly test (`test_weekly_session_anomaly_e2e.py`, Claude Code-shaped multi-turn sessions against real providers with ceilings on error rate, cache read/write, turn time, and spend; marked `weekly` and deselected unless `E2E_WEEKLY_ANOMALY` is set, driven by `.github/workflows/weekly_load_anomaly.yml`) and markerless harness unit tests for the Locust/session-anomaly aggregation logic
- `load/` - performance-category tests, kept OUT of the main suite: throughput/load SLO tests are a different testing category from functional e2e (variance-driven, historically flaky) and live outside this suite until re-implemented as their own pipeline (LIT-5163); do not add a live load test that runs in the default collection. What lives here: the weekly session-anomaly test (`test_weekly_session_anomaly_e2e.py`, Claude Code-shaped multi-turn sessions against real providers with ceilings on error rate, cache read/write, turn time, and spend; marked `weekly` and deselected unless `E2E_WEEKLY_ANOMALY` is set, driven by `.github/workflows/weekly_load_anomaly.yml`), the Redis chaos test (`test_redis_chaos_e2e.py`, locust load against mock deployments split round robin over `/chat/completions` and `/v1/messages`, one endpoint per simulated user, with `CLIENT PAUSE ALL` on the proxy's Redis mid-run to simulate it being down outright, asserting zero failed requests on every endpoint, budgeting RSS and CPU-per-request as ratios against the same run's healthy phase, and holding p50/p90/p99 latency and log-bytes-per-request to flat ceilings (a ratio cannot bound those two: an open breaker skips Redis instead of waiting on it, so the chaos phase can measure cheaper than baseline while still being far slower than a user should see); needs a proxy booted from `gateway/redis_chaos_ci_config.yml` on the same host with `E2E_PROXY_PID` and `E2E_PROXY_LOG` set, marked `redis_chaos`, deselected unless `E2E_REDIS_CHAOS` is set and excluded from the per-PR selector like the rest of `load/`, driven by `.github/workflows/test-e2e-redis-chaos.yml` and by the Buildkite `e2e-redis-chaos` step in project-releaser, which runs the proxy, Postgres and Valkey co-located with pytest in one pod and sets the opt-in), and markerless harness unit tests for the locust, process-usage, and session-anomaly aggregation logic
- `other/` - the holding-pen suite for the `other.*` registry cluster with no home of its own yet: the master-key auth gate, JWT auth (access tokens issued by a real Keycloak realm, `idp.py` plus `idp_realm.json`, whose JWKS the proxy's `JWT_PUBLIC_KEY_URL` points at; see CONTRIBUTING.md for the start command and config block), and the process-lifecycle health probes (liveness, public readiness, authenticated readiness diagnostics). Promote a cluster out once it is large/stable enough for its own suite
- `gateway/` - proxy configuration only (`litellm-config.yml`); no tests
- `claude_code/` - the Claude Code compatibility matrix: drives the real `claude` CLI (and HTTP probes) against a proxy for each feature x provider cell, reporting tagged-union outcomes via the `compat_result` fixture; ships its own driver/builder/publisher plus `_*_unit_tests/` trees. The HTTP probes ride the shared transport (`ProxyClient.count_tokens` / `ProxyClient.messages`); the CLI-driving path stays bespoke

2
tests/e2e/CONTRIBUTING.md
View file

@ -101,7 +101,7 @@ A couple of logging destinations are configured on the proxy rather than by the
### The pull request check
Every same-repository PR that adds, modifies, or renames a `tests/e2e/**/test_*.py` file runs those changed files three times. A change to the harness itself, meaning a root-level `tests/e2e/*.py` file or `pytest.ini`, `tests/e2e/gateway/`, `.github/e2e-stack/`, or the workflow, also runs the `access_control` suite and both JWT suites as canaries, because those files have no test of their own that exercises the stack. `.github/e2e-stack/select_tests.py` applies both rules. The stack config at `tests/e2e/gateway/stage_mirror_ci_config.yml` must declare every model the selected suites use; a missing one shows up as a failed test id in the public log. The suite's own single rerun for network errors and 5xx responses (see `pytest.ini`) applies on every pass, so a transport blip does not fail the check while a race inside a test still does. The stage-mirror stack has a control-plane backend, two gateways behind nginx, Postgres, Keycloak, Jaeger, and TLS cluster-mode Valkey. Realm-only edits also trigger these canaries. The stack exports every gateway address in `LITELLM_PROXY_REPLICA_URLS`, so model registration waits until each gateway lists the new model rather than whichever one the load balancer answered from. Documentation, deleted-file, and application-only changes do not start the stack or request environment approval. The `ui/`, `claude_code/`, and `load/` directories, `batches/test_managed_files_enforcement_e2e.py`, `llm_translation/realtime/test_realtime_pipecat_audio_e2e.py`, and `guardrails/test_presidio_masking_e2e.py` remain outside this check because they use separate tooling or need a differently configured stack: the pipecat audio suite skips itself at import time unless the NLTK `punkt_tab` data is installed, and the presidio suite fails without the analyzer and anonymizer services this stack does not start
Every same-repository PR that adds, modifies, or renames a `tests/e2e/**/test_*.py` file runs those changed files three times. A change to the harness itself, meaning a root-level `tests/e2e/*.py` file or `pytest.ini`, `tests/e2e/gateway/`, `.github/e2e-stack/`, or the workflow, also runs the `access_control` suite and both JWT suites as canaries, because those files have no test of their own that exercises the stack. `.github/e2e-stack/select_tests.py` applies both rules. The stack config at `tests/e2e/gateway/stage_mirror_ci_config.yml` must declare every model the selected suites use; a missing one shows up as a failed test id in the public log. The suite's own single rerun for network errors and 5xx responses (see `pytest.ini`) applies on every pass, so a transport blip does not fail the check while a race inside a test still does. The stage-mirror stack has a control-plane backend, two gateways behind nginx, Postgres, Keycloak, Jaeger, and TLS cluster-mode Valkey. Realm-only edits also trigger these canaries. The stack exports every gateway address in `LITELLM_PROXY_REPLICA_URLS`, so model registration waits until each gateway lists the new model rather than whichever one the load balancer answered from. Documentation, deleted-file, and application-only changes do not start the stack or request environment approval. The `ui/`, `claude_code/`, and `load/` directories, `batches/test_managed_files_enforcement_e2e.py`, `llm_translation/realtime/test_realtime_pipecat_audio_e2e.py`, and `guardrails/test_presidio_masking_e2e.py` remain outside this check because they use separate tooling or need a differently configured stack: the pipecat audio suite skips itself at import time unless the NLTK `punkt_tab` data is installed, and the presidio suite fails without the analyzer and anonymizer services this stack does not start. The Redis chaos test under `load/` needs a proxy it can pause the Redis of on the same host (`gateway/redis_chaos_ci_config.yml`), which `.github/workflows/test-e2e-redis-chaos.yml` boots, and which the Buildkite `e2e-redis-chaos` step in project-releaser runs co-located with Postgres and Valkey in one pod; it is deselected unless `E2E_REDIS_CHAOS` is set
Every selected file must execute at least one passing test in each pass, and any test failure, collection error, or entirely skipped or deselected file fails the check. A file whose tests are all marked skip therefore cannot pass this check, so unskip at least one of them, or add the file to `UNSUPPORTED` in `select_tests.py` with the reason, before changing one. A failed pass stops the run. The public log prints pytest's one-line summary for each pass, including the rerun count, and names each failed or errored test as `classname::name`, so a retried network error or a failing test is visible without the raw output. The final `e2e-changed-tests` job succeeds only when no supported test files changed or the approved run completed all three passes. Fork PRs with selected tests fail this gate until a maintainer brings the reviewed change onto a same-repository branch

5
tests/e2e/conftest.py
View file

@ -89,6 +89,11 @@ def pytest_configure(config: pytest.Config) -> None:
"markers",
"managed_files: needs a proxy running with require_managed_files enabled; deselected unless E2E_MANAGED_FILES_STACK is set",
)
config.addinivalue_line(
"markers",
"redis_chaos: load test that pauses the proxy's Redis outright mid-run; needs a proxy booted from "
"gateway/redis_chaos_ci_config.yml on the same host, and is deselected unless E2E_REDIS_CHAOS is set",
)
def pytest_sessionstart(session: pytest.Session) -> None:

1
tests/e2e/coverage_registry/reliability.yaml
View file

@ -31,6 +31,7 @@
- {id: reliability.cache.exact.returns_cached, module: reliability, tier: P1, behavior: cache, variant: exact, assertions: [returns_cached], exercised_on: [chat_completions, messages, embeddings], source: "litellm/caching/caching.py", rationale: "Response cache returns cached on exact match"}
- {id: reliability.cache.prompt_caching_model_select.returns_cached, module: reliability, tier: P1, behavior: cache, variant: prompt_caching_model_select, assertions: [returns_cached], exercised_on: [chat_completions], source: "router_utils/prompt_caching_cache.py", rationale: "Selects model supporting prompt caching for cacheable prefix"}
- {id: reliability.circuit_breaker.redis.trips_then_recovers, module: reliability, tier: P0, behavior: circuit_breaker, variant: redis, assertions: [trips_then_recovers], exercised_on: [chat_completions, messages, embeddings], source: "litellm/caching/redis_cache.py:99", rationale: "Redis breaker CLOSED->OPEN->HALF_OPEN; guards all cache/rate-limit ops"}
- {id: reliability.circuit_breaker.redis_timeout.stays_responsive, module: reliability, tier: P1, behavior: circuit_breaker, variant: redis_timeout, assertions: [stays_responsive], exercised_on: [chat_completions, messages], source: "litellm/proxy/hooks/proxy_track_cost_callback.py:386", fail_before_fix: proven, rationale: "Under locust load split round robin over /chat/completions and /v1/messages with every request retrying through failing mock deployments, holding Redis in CLIENT PAUSE ALL for the phase trips the breaker and every request still succeeds, with latency, RSS, and CPU reported as p50/p90/p99 against the pre-pause baseline; on v1.100.0 the failed-tracking alert body doubled per request until the worker OOMed (LIT-6780)"}
- {id: reliability.timeout.request_timeout.exceeds_deadline, module: reliability, tier: P1, behavior: timeout, variant: request_timeout, assertions: [exceeds_deadline], exercised_on: [chat_completions, messages], source: "litellm/router.py:545-551", rationale: "Per-request timeout raises Timeout"}
- {id: reliability.timeout.stream_timeout.exceeds_deadline, module: reliability, tier: P1, behavior: timeout, variant: stream_timeout, assertions: [exceeds_deadline], exercised_on: [chat_completions], source: "litellm/router.py:551", rationale: "Streaming chunk-delivery timeout"}
- {id: reliability.perf.throughput.under_slo, module: reliability, tier: P1, behavior: perf, variant: throughput, assertions: [under_slo], exercised_on: [chat_completions, messages], source: grammar, rationale: "Throughput SLO under load"}

5
tests/e2e/e2e_config.py
View file

@ -13,7 +13,6 @@ from pathlib import Path
from typing import Final
from dotenv import load_dotenv
from fixture_mode import deterministic_marker, parse_fixture_mode
from provider_edge import provider_edge_api_base
@ -144,6 +143,7 @@ LOAD_MIN_CONCURRENCY_EFFICIENCY = float(os.environ.get("E2E_LOAD_MIN_CONCURRENCY
WEEKLY_ANOMALY_OPT_IN_ENV = "E2E_WEEKLY_ANOMALY"
MANAGED_FILES_OPT_IN_ENV = "E2E_MANAGED_FILES_STACK"
REDIS_CHAOS_OPT_IN_ENV = "E2E_REDIS_CHAOS"
ANOMALY_SESSIONS = int(os.environ.get("E2E_ANOMALY_SESSIONS", "6"))
ANOMALY_TURNS_PER_SESSION = int(os.environ.get("E2E_ANOMALY_TURNS_PER_SESSION", "6"))
ANOMALY_TURN_ATTEMPTS = int(os.environ.get("E2E_ANOMALY_TURN_ATTEMPTS", "3"))
@ -181,8 +181,7 @@ def datadog_mcp_url(*, toolsets: str = "core") -> str:
site = (
os.environ.get("DD_SITE", DD_SITE) or "datadoghq.com"
).strip().removeprefix("https://").removeprefix("http://").rstrip("/")
if site.startswith("app."):
site = site[len("app.") :]
site = site.removeprefix("app.")
host = "mcp.datadoghq.com" if site in ("", "datadoghq.com") else f"mcp.{site}"
base = f"https://{host}/v1/mcp"
return f"{base}?toolsets={toolsets}" if toolsets else base

19
tests/e2e/gateway/redis_chaos_ci_config.yml Normal file
View file

@ -0,0 +1,19 @@
general_settings:
master_key: os.environ/LITELLM_MASTER_KEY
store_model_in_db: true
use_redis_transaction_buffer: true
litellm_settings:
callbacks: ["prometheus"]
require_auth_for_metrics_endpoint: false
enable_redis_auth_cache: true
cache: true
cache_params:
type: redis
host: 127.0.0.1
port: 6379
socket_timeout: 0.1
router_settings:
num_retries: 2
disable_cooldowns: true

23
tests/e2e/load/conftest.py
View file

@ -3,26 +3,23 @@ from __future__ import annotations
import os
import pytest
from e2e_config import WEEKLY_ANOMALY_OPT_IN_ENV
from e2e_config import REDIS_CHAOS_OPT_IN_ENV, WEEKLY_ANOMALY_OPT_IN_ENV
from load_client import LoadClient, build_client
from proxy_client import ProxyClient
_OPT_IN_MARKERS = (
("weekly", WEEKLY_ANOMALY_OPT_IN_ENV),
("redis_chaos", REDIS_CHAOS_OPT_IN_ENV),
)
def pytest_collection_modifyitems(
config: pytest.Config, items: list[pytest.Item]
) -> None:
if os.environ.get(WEEKLY_ANOMALY_OPT_IN_ENV):
return
deselected = [
item for item in items if item.get_closest_marker("weekly") is not None
]
def pytest_collection_modifyitems(config: pytest.Config, items: list[pytest.Item]) -> None:
opted_out = {marker for marker, opt_in_env in _OPT_IN_MARKERS if not os.environ.get(opt_in_env)}
deselected = [item for item in items if any(item.get_closest_marker(marker) is not None for marker in opted_out)]
if not deselected:
return
config.hook.pytest_deselected(items=deselected)
items[:] = [
item for item in items if item.get_closest_marker("weekly") is None
]
items[:] = [item for item in items if item not in deselected]
@pytest.fixture(scope="session")

150
tests/e2e/load/locust_load.py
View file

@ -1,17 +1,26 @@
from __future__ import annotations
import csv
import os
import subprocess
import sys
import tempfile
from collections.abc import Sequence
from dataclasses import dataclass
from itertools import accumulate
from pathlib import Path
from typing import Final
from pydantic import BaseModel, TypeAdapter
_LOCUSTFILE = Path(__file__).with_name("locustfile.py")
_CSV_PREFIX = "locust"
_GENERATOR_SATURATION_MARKER = "CPU usage above"
_MAX_REPORTED_ERRORS = 5
class LocustStatEntry(BaseModel):
name: str
num_requests: int
num_failures: int
start_time: float
@ -29,12 +38,25 @@ class LoadError:
occurrences: int
@dataclass(frozen=True, slots=True)
class EndpointLoad:
"""One route's share of a phase, so a run that silently drove only one of them is visible."""
name: str
requests: int
failures: int
p50_seconds: float
@dataclass(frozen=True, slots=True)
class LoadResult:
requests: int
failures: int
requests_per_second: float
median_response_seconds: float
p50_seconds: float
p90_seconds: float
p99_seconds: float
endpoints: tuple[EndpointLoad, ...]
errors: tuple[LoadError, ...]
generator_warnings: tuple[str, ...]
@ -53,33 +75,65 @@ class LoadResult:
lines.append("locust recorded no error breakdown")
return "; ".join((*lines, *self.generator_warnings))
def latency_summary(self) -> str:
return f"p50 {self.p50_seconds:.3f}s, p90 {self.p90_seconds:.3f}s, p99 {self.p99_seconds:.3f}s"
def median_seconds(entries: list[LocustStatEntry]) -> float:
samples = sorted(
(milliseconds, count) for entry in entries for milliseconds, count in entry.response_times.items()
)
def endpoint_summary(self) -> str:
return ", ".join(
f"{endpoint.name} {endpoint.requests} requests, {endpoint.failures} failures, "
f"p50 {endpoint.p50_seconds:.3f}s"
for endpoint in self.endpoints
)
def percentile_seconds(entries: Sequence[LocustStatEntry], fraction: float) -> float:
"""The response time at `fraction` of the merged histograms, in seconds.
Locust buckets response times by millisecond, so this reads the first bucket whose
running count reaches the rank, the same lower-sample convention locust's own
percentiles use.
"""
samples = sorted((milliseconds, count) for entry in entries for milliseconds, count in entry.response_times.items())
total = sum(count for _, count in samples)
if total == 0:
return 0.0
running = accumulate(count for _, count in samples)
return next(
milliseconds for (milliseconds, _), seen in zip(samples, running) if seen >= total / 2
) / 1000.0
rank: Final = total * fraction
return next(milliseconds for (milliseconds, _), seen in zip(samples, running) if seen >= rank) / 1000.0
def per_endpoint(entries: Sequence[LocustStatEntry]) -> tuple[EndpointLoad, ...]:
"""Each locust request name's own totals, in the order the names first appear."""
names: Final = tuple(dict.fromkeys(entry.name for entry in entries))
grouped: Final = ((name, tuple(entry for entry in entries if entry.name == name)) for name in names)
return tuple(
EndpointLoad(
name=name,
requests=sum(entry.num_requests for entry in group),
failures=sum(entry.num_failures for entry in group),
p50_seconds=percentile_seconds(group, 0.5),
)
for name, group in grouped
)
def aggregate_stats(
entries: list[LocustStatEntry],
entries: Sequence[LocustStatEntry],
errors: tuple[LoadError, ...],
generator_warnings: tuple[str, ...],
) -> LoadResult:
requests = sum(entry.num_requests for entry in entries)
failures = sum(entry.num_failures for entry in entries)
endpoints = per_endpoint(entries)
if not entries or requests == 0:
return LoadResult(
requests=requests,
failures=failures,
requests_per_second=0.0,
median_response_seconds=0.0,
p50_seconds=0.0,
p90_seconds=0.0,
p99_seconds=0.0,
endpoints=endpoints,
errors=errors,
generator_warnings=generator_warnings,
)
@ -88,7 +142,10 @@ def aggregate_stats(
requests=requests,
failures=failures,
requests_per_second=requests / elapsed if elapsed > 0 else 0.0,
median_response_seconds=median_seconds(entries),
p50_seconds=percentile_seconds(entries, 0.5),
p90_seconds=percentile_seconds(entries, 0.9),
p99_seconds=percentile_seconds(entries, 0.99),
endpoints=endpoints,
errors=errors,
generator_warnings=generator_warnings,
)
@ -121,3 +178,74 @@ def read_generator_warnings(stderr: str) -> tuple[str, ...]:
if _GENERATOR_SATURATION_MARKER in line
)
return tuple(dict.fromkeys(saturated))
def run_gateway_load(
*,
base_url: str,
api_keys: tuple[str, ...],
model: str,
endpoints: tuple[str, ...],
users: int,
spawn_rate: float,
duration_seconds: float,
) -> LoadResult:
"""Drive `endpoints` from headless locust and aggregate what it reported.
Each simulated user picks one of `api_keys`, so auth and budget lookups spread over a
pool of virtual keys instead of keeping one key's cache entry permanently warm, and one
of `endpoints` round robin, so the run covers every route the caller asked for.
"""
with tempfile.TemporaryDirectory(prefix="e2e-load-") as report_dir:
csv_prefix = Path(report_dir) / _CSV_PREFIX
completed = subprocess.run(
[
sys.executable,
"-m",
"locust",
"--headless",
"--json",
"--csv",
str(csv_prefix),
"--locustfile",
str(_LOCUSTFILE),
"--host",
base_url,
"--users",
str(users),
"--spawn-rate",
str(spawn_rate),
"--run-time",
f"{int(duration_seconds)}s",
"--exit-code-on-error",
"0",
],
env={
**os.environ,
"LOAD_API_KEYS": ",".join(api_keys),
"LOAD_MODEL": model,
"LOAD_ENDPOINTS": ",".join(endpoints),
},
capture_output=True,
text=True,
timeout=duration_seconds + 120,
check=False,
)
if completed.returncode != 0:
raise RuntimeError(
f"locust exited {completed.returncode} before it could report throughput "
f"(a startup failure, not request failures, which are folded into the JSON summary via "
f"--exit-code-on-error 0):\n{completed.stderr}"
)
try:
entries = _STATS_ADAPTER.validate_json(completed.stdout)
except ValueError as exc:
raise RuntimeError(
f"locust exited 0 but did not print a parseable --json throughput summary on stdout; "
f"got stdout={completed.stdout!r}, stderr={completed.stderr!r}"
) from exc
return aggregate_stats(
entries,
read_errors(csv_prefix.with_name(f"{_CSV_PREFIX}_failures.csv")),
read_generator_warnings(completed.stderr),
)

52
tests/e2e/load/locustfile.py Normal file
View file

@ -0,0 +1,52 @@
from __future__ import annotations
import os
import random
import uuid
from itertools import cycle
from typing import Final
from locust import FastHttpUser, constant, task
_MODEL: Final = os.environ["LOAD_MODEL"]
_API_KEYS: Final = tuple(os.environ["LOAD_API_KEYS"].split(","))
_NEXT_ENDPOINT: Final = cycle(os.environ["LOAD_ENDPOINTS"].split(","))
_FILLER: Final = "x" * 40_000
def _payload() -> dict[str, object]:
"""A prompt no other request sent, so the response cache never answers for the deployment.
Both endpoints take the same body: /v1/messages requires max_tokens, which /chat/completions
also accepts, so one payload serves the whole round robin. Padded to tens of KB so a
per-request bookkeeping cost that scales with body size (string formatting, hashing) shows
up in the CPU and log-size budgets instead of hiding behind a 40-byte prompt.
"""
return {
"model": _MODEL,
"messages": [{"role": "user", "content": f"load test ping {uuid.uuid4().hex} {_FILLER}"}],
"max_tokens": 16,
}
class GatewayUser(FastHttpUser):
"""One simulated user, pinned to one endpoint for its lifetime.
Endpoints are handed out round robin as users spawn, so a run spreads evenly over them
while each user's traffic stays on a single route, the way a real client behaves.
"""
wait_time = constant(0)
def on_start(self) -> None:
self.headers = {"Authorization": f"Bearer {random.choice(_API_KEYS)}"}
self.endpoint = next(_NEXT_ENDPOINT)
@task
def call(self) -> None:
self.client.post( # pyright: ignore[reportUnknownMemberType] # locust FastHttpSession.post types json/**kwargs as Any
self.endpoint,
json=_payload(),
headers=self.headers,
name=self.endpoint,
)

81
tests/e2e/load/phase_budget.py Normal file
View file

@ -0,0 +1,81 @@
"""Comparing one load phase against another, for tests that degrade a dependency mid-run.
Two shapes of ceiling, because the metrics divide into two kinds. RSS and CPU are
machine-shaped: RSS scales with worker count and CPU with core count, so an absolute number
calibrated on one runner means nothing on the next, and what travels is the ratio against a
healthy phase measured on the same machine in the same run. Latency and log volume are not:
a ratio there is actively misleading, because a dependency that fails fast once its breaker
opens can make the degraded phase look cheaper than the healthy one while still being far
slower or noisier than a user should ever see. Those get a flat ceiling, which is the promise
the test is actually making.
"""
from __future__ import annotations
from dataclasses import dataclass
from typing import Final, TypeAlias
def _rendered(value: float, unit: str, decimals: int) -> str:
return f"{value:.{decimals}f}{unit}"
@dataclass(frozen=True, slots=True)
class RatioBudget:
"""One metric's healthy value, its degraded value, and how much growth is allowed."""
name: str
baseline: float
degraded: float
ratio_ceiling: float
unit: str
decimals: int = 1
@property
def ratio(self) -> float | None:
"""How many times the baseline the degraded value is, or None if there is no baseline."""
return self.degraded / self.baseline if self.baseline > 0 else None
def violation(self) -> str | None:
"""Why this metric fails its budget, or None if it passes."""
ratio: Final = self.ratio
if ratio is None:
return (
f"{self.name} measured {_rendered(self.baseline, self.unit, self.decimals)} in the healthy phase, "
f"so there is nothing to compare the degraded phase against; the measurement did not happen"
)
if ratio > self.ratio_ceiling:
return (
f"{self.name} went from {_rendered(self.baseline, self.unit, self.decimals)} healthy to "
f"{_rendered(self.degraded, self.unit, self.decimals)} degraded, {ratio:.1f}x the baseline and past "
f"the {self.ratio_ceiling:.1f}x allowed"
)
return None
@dataclass(frozen=True, slots=True)
class AbsoluteBudget:
"""One metric's degraded value against a flat ceiling, for metrics a ratio cannot bound."""
name: str
measured: float
ceiling: float
unit: str
decimals: int = 1
def violation(self) -> str | None:
"""Why this metric fails its budget, or None if it passes."""
if self.measured > self.ceiling:
return (
f"{self.name} measured {_rendered(self.measured, self.unit, self.decimals)} in the degraded phase, "
f"past the {_rendered(self.ceiling, self.unit, self.decimals)} allowed"
)
return None
Budget: TypeAlias = RatioBudget | AbsoluteBudget
def violations(budgets: tuple[Budget, ...]) -> tuple[str, ...]:
"""Every budget the run blew, so one failure reports all of them instead of the first."""
return tuple(violation for budget in budgets if (violation := budget.violation()) is not None)

164
tests/e2e/load/proxy_usage.py Normal file
View file

@ -0,0 +1,164 @@
"""Resident memory and CPU of the proxy process tree, sampled on a background thread.
The proxy under load runs several worker processes, and `/metrics` cannot report their
memory: litellm sets PROMETHEUS_MULTIPROC_DIR when num_workers > 1, and the multiprocess
collector drops the process collector's `process_resident_memory_bytes` /
`process_cpu_seconds_total` entirely. So the test measures the tree itself through psutil,
which needs the proxy to run on the same host as the test.
"""
from __future__ import annotations
import math
import threading
import time
from dataclasses import dataclass
from typing import Final
import psutil
from pydantic import BaseModel, ConfigDict
class _MemoryInfo(BaseModel):
model_config = ConfigDict(from_attributes=True)
rss: int
@dataclass(frozen=True, slots=True)
class UsageSample:
elapsed_seconds: float
rss_bytes: int
cpu_seconds: float
@dataclass(frozen=True, slots=True)
class UsageWindow:
"""The samples taken across one phase, plus what they say about that phase."""
samples: tuple[UsageSample, ...]
def rss_percentile(self, fraction: float) -> int:
if not self.samples:
return 0
ordered: Final = sorted(sample.rss_bytes for sample in self.samples)
return ordered[_rank(len(ordered), fraction)]
def cpu_seconds_consumed(self) -> float:
"""CPU seconds the tree burned across the window, from its monotonic counter."""
if len(self.samples) < 2:
return 0.0
return self.samples[-1].cpu_seconds - self.samples[0].cpu_seconds
def cpu_seconds_per_request(self, requests: int) -> float:
"""CPU seconds the tree spent per request served.
The portable cost figure: cores-busy saturates at the worker count under enough load,
so it reads the same whether a request costs 10 ms of CPU or 40 ms. This does not.
"""
return self.cpu_seconds_consumed() / requests if requests else 0.0
def cpu_utilization_percentiles(self) -> tuple[float, float, float]:
"""Per-interval CPU utilization (cores busy) at p50, p90 and p99.
Derived from consecutive samples of the cumulative counter rather than
psutil's own cpu_percent, so it covers every process in the tree including
workers that came and went between samples.
"""
rates: Final = sorted(
(later.cpu_seconds - earlier.cpu_seconds) / (later.elapsed_seconds - earlier.elapsed_seconds)
for earlier, later in zip(self.samples, self.samples[1:])
if later.elapsed_seconds > earlier.elapsed_seconds
)
if not rates:
return 0.0, 0.0, 0.0
return (
rates[_rank(len(rates), 0.5)],
rates[_rank(len(rates), 0.9)],
rates[_rank(len(rates), 0.99)],
)
def summary(self) -> str:
p50_cpu, p90_cpu, p99_cpu = self.cpu_utilization_percentiles()
return (
f"RSS p50 {self.rss_percentile(0.5) / 2**20:.0f} MB, "
f"p90 {self.rss_percentile(0.9) / 2**20:.0f} MB, "
f"p99 {self.rss_percentile(0.99) / 2**20:.0f} MB; "
f"CPU cores busy p50 {p50_cpu:.2f}, p90 {p90_cpu:.2f}, p99 {p99_cpu:.2f}; "
f"{self.cpu_seconds_consumed():.1f} CPU seconds consumed"
)
def _rank(count: int, fraction: float) -> int:
"""Index of the sample at `fraction`, the same lower-sample convention as locust's percentiles."""
return min(count - 1, max(0, math.ceil(count * fraction) - 1))
def _read_process(process: psutil.Process) -> tuple[int, float] | None:
try:
with process.oneshot():
memory: Final = _MemoryInfo.model_validate(process.memory_info())
times: Final = process.cpu_times()
return memory.rss, times.user + times.system
except (psutil.NoSuchProcess, psutil.AccessDenied):
return None
class ProxyUsageSampler:
"""Samples the proxy process tree every `interval_seconds` until stopped.
`split()` returns the samples taken so far and starts a new window, so one sampler
covers a baseline phase and a chaos phase without a gap between them.
"""
def __init__(self, pid: int, interval_seconds: float = 1.0) -> None:
self._process: Final = psutil.Process(pid)
self._interval: Final = interval_seconds
self._stop: Final = threading.Event()
self._lock: Final = threading.Lock()
self._samples: list[UsageSample] = [] # mutable-ok: a sampling buffer the reader drains under a lock
self._started: Final = time.monotonic()
self._thread: Final = threading.Thread(target=self._run, name="proxy-usage-sampler", daemon=True)
def __enter__(self) -> ProxyUsageSampler:
self._thread.start()
return self
def __exit__(self, *_: object) -> None:
self._stop.set()
self._thread.join(timeout=self._interval * 5)
def _tree(self) -> tuple[psutil.Process, ...]:
try:
return (self._process, *self._process.children(recursive=True))
except psutil.NoSuchProcess:
return ()
def _sample(self) -> UsageSample | None:
readings: Final = tuple(reading for process in self._tree() if (reading := _read_process(process)) is not None)
if not readings:
return None
return UsageSample(
elapsed_seconds=time.monotonic() - self._started,
rss_bytes=sum(rss for rss, _ in readings),
cpu_seconds=sum(cpu for _, cpu in readings),
)
def _run(self) -> None:
while not self._stop.is_set():
sample = self._sample()
if sample is not None:
with self._lock:
self._samples.append(sample)
self._stop.wait(self._interval)
def split(self) -> UsageWindow:
"""The window that ends now; the next one starts from this window's last sample.
The boundary sample is carried into the next window so its CPU counter has a
starting point, which is what makes the two windows' utilization comparable.
"""
with self._lock:
taken = tuple(self._samples)
self._samples = [taken[-1]] if taken else [] # rebind-ok: drains the buffer under the lock
return UsageWindow(samples=taken)

91
tests/e2e/load/test_locust_load.py
View file

@ -1,13 +1,14 @@
from __future__ import annotations
from pathlib import Path
from typing import Final
from locust_load import (
LoadError,
LoadResult,
LocustStatEntry,
aggregate_stats,
median_seconds,
percentile_seconds,
read_errors,
read_generator_warnings,
)
@ -18,12 +19,14 @@ _FAILURES_HEADER = "Method,Name,Error,Occurrences,First Seen,Last Seen\n"
def _entry(
*,
num_requests: int,
name: str = "/chat/completions",
num_failures: int = 0,
start_time: float = 1000.0,
last_request_timestamp: float = 1010.0,
response_times: dict[int, int] | None = None,
) -> LocustStatEntry:
return LocustStatEntry(
name=name,
num_requests=num_requests,
num_failures=num_failures,
start_time=start_time,
@ -41,35 +44,48 @@ def _result(
requests=10,
failures=10,
requests_per_second=1.0,
median_response_seconds=0.05,
p50_seconds=0.05,
p90_seconds=0.08,
p99_seconds=0.1,
endpoints=(),
errors=errors,
generator_warnings=generator_warnings,
)
class TestSerialLatency:
class TestPercentiles:
def test_median_is_the_middle_sample_not_the_mean_a_slow_tail_would_drag(self) -> None:
# Nine fast requests and one very slow one: the mean is 1.99s, the median is 20ms.
entry = _entry(num_requests=10, response_times={20: 9, 20000: 1})
assert median_seconds([entry]) == 0.02
assert percentile_seconds([entry], 0.5) == 0.02
def test_median_merges_the_histograms_of_every_stats_entry(self) -> None:
def test_the_tail_percentiles_reach_the_slow_samples_the_median_hides(self) -> None:
# 100 samples: 89 fast, 10 slow, 1 very slow. p50 sits in the fast bucket, p90 in the
# slow one, and p99 lands on the single very slow sample.
entry = _entry(num_requests=100, response_times={20: 89, 500: 10, 20000: 1})
assert percentile_seconds([entry], 0.5) == 0.02
assert percentile_seconds([entry], 0.9) == 0.5
assert percentile_seconds([entry], 0.99) == 0.5
assert percentile_seconds([entry], 1.0) == 20.0
def test_percentiles_merge_the_histograms_of_every_stats_entry(self) -> None:
# Per entry the median would be 10ms and 90ms; merged, the middle of the five samples is 90ms.
entries = [
_entry(num_requests=2, response_times={10: 2}),
_entry(num_requests=3, response_times={90: 3}),
]
assert median_seconds(entries) == 0.09
assert percentile_seconds(entries, 0.5) == 0.09
def test_an_even_split_takes_the_lower_middle_sample_as_locust_itself_does(self) -> None:
entry = _entry(num_requests=4, response_times={10: 2, 90: 2})
assert median_seconds([entry]) == 0.01
assert percentile_seconds([entry], 0.5) == 0.01
def test_no_samples_reports_zero_rather_than_dividing_by_an_empty_histogram(self) -> None:
assert median_seconds([]) == 0.0
assert percentile_seconds([], 0.5) == 0.0
class TestAggregate:
@ -84,9 +100,20 @@ class TestAggregate:
result = aggregate_stats([entry], (), ())
assert result.requests_per_second == 3.0
assert result.median_response_seconds == 0.057
assert result.p50_seconds == 0.057
assert result.p99_seconds == 0.057
assert result.failure_ratio == 0.0
def test_tail_percentiles_come_from_the_slow_end_of_the_histogram(self) -> None:
entry = _entry(num_requests=100, response_times={20: 89, 500: 10, 3000: 1})
result = aggregate_stats([entry], (), ())
assert result.p50_seconds == 0.02
assert result.p90_seconds == 0.5
assert result.p99_seconds == 0.5
assert result.latency_summary() == "p50 0.020s, p90 0.500s, p99 0.500s"
def test_throughput_spans_from_the_earliest_start_when_locust_reports_several_entries(self) -> None:
entries = [
_entry(num_requests=60, start_time=1000.0, last_request_timestamp=1030.0),
@ -103,6 +130,49 @@ class TestAggregate:
assert result.requests == 0
assert result.requests_per_second == 0.0
assert result.failure_ratio == 1.0
assert result.endpoints == ()
class TestPerEndpoint:
def test_each_route_keeps_its_own_requests_failures_and_median(self) -> None:
entries: Final = (
_entry(name="/chat/completions", num_requests=100, response_times={20: 100}),
_entry(name="/v1/messages", num_requests=40, num_failures=3, response_times={900: 40}),
)
result: Final = aggregate_stats(entries, (), ())
assert tuple((one.name, one.requests, one.failures, one.p50_seconds) for one in result.endpoints) == (
("/chat/completions", 100, 0, 0.02),
("/v1/messages", 40, 3, 0.9),
)
def test_several_stats_entries_for_one_route_fold_into_a_single_row(self) -> None:
entries: Final = (
_entry(name="/v1/messages", num_requests=10, response_times={30: 10}),
_entry(name="/v1/messages", num_requests=30, num_failures=1, response_times={30: 30}),
)
result: Final = aggregate_stats(entries, (), ())
assert tuple((one.name, one.requests, one.failures) for one in result.endpoints) == (("/v1/messages", 40, 1),)
def test_a_route_that_never_ran_is_absent_so_a_one_sided_run_cannot_pass_unnoticed(self) -> None:
result: Final = aggregate_stats((_entry(name="/chat/completions", num_requests=10),), (), ())
assert tuple(one.name for one in result.endpoints) == ("/chat/completions",)
def test_the_summary_names_every_route_with_its_counts(self) -> None:
entries: Final = (
_entry(name="/chat/completions", num_requests=2, response_times={20: 2}),
_entry(name="/v1/messages", num_requests=1, num_failures=1, response_times={500: 1}),
)
result: Final = aggregate_stats(entries, (), ())
assert result.endpoint_summary() == (
"/chat/completions 2 requests, 0 failures, p50 0.020s, /v1/messages 1 requests, 1 failures, p50 0.500s"
)
class TestErrorBreakdown:
@ -133,8 +203,7 @@ class TestErrorBreakdown:
def test_diagnosis_caps_the_list_and_says_how_many_it_left_out(self) -> None:
result = _result(
errors=tuple(
LoadError(name="/chat/completions", error=f"error-{index}", occurrences=index)
for index in range(1, 9)
LoadError(name="/chat/completions", error=f"error-{index}", occurrences=index) for index in range(1, 9)
)
)

105
tests/e2e/load/test_phase_budget.py Normal file
View file

@ -0,0 +1,105 @@
from __future__ import annotations
from typing import Final
from phase_budget import AbsoluteBudget, RatioBudget, violations
def _budget(*, baseline: float, degraded: float, ceiling: float = 2.0) -> RatioBudget:
return RatioBudget(
name="p99 RSS", baseline=baseline, degraded=degraded, ratio_ceiling=ceiling, unit=" MB", decimals=0
)
class TestRatioBudget:
def test_growth_within_the_ceiling_is_not_a_violation(self) -> None:
assert _budget(baseline=100, degraded=199).violation() is None
def test_growth_exactly_at_the_ceiling_is_allowed(self) -> None:
assert _budget(baseline=100, degraded=200).violation() is None
def test_growth_past_the_ceiling_reports_both_values_and_the_ratio(self) -> None:
violation: Final = _budget(baseline=100, degraded=250).violation()
assert violation is not None
assert "100 MB" in violation
assert "250 MB" in violation
assert "2.5x" in violation
assert "2.0x allowed" in violation
def test_shrinking_is_never_a_violation(self) -> None:
assert _budget(baseline=100, degraded=10).violation() is None
def test_a_missing_baseline_is_a_violation_rather_than_a_silent_pass(self) -> None:
# The trap this guards: 0 as a baseline would make every ratio a division by zero, and
# treating it as "no growth" would pass a run that measured nothing at all.
violation: Final = _budget(baseline=0, degraded=4000).violation()
assert violation is not None
assert "nothing to compare" in violation
def test_the_unit_and_decimals_carry_into_the_message(self) -> None:
violation: Final = RatioBudget(
name="p99 latency", baseline=0.16, degraded=9.5, ratio_ceiling=8.0, unit="s", decimals=3
).violation()
assert violation is not None
assert "0.160s" in violation
assert "9.500s" in violation
class TestAbsoluteBudget:
def test_a_value_under_the_ceiling_is_not_a_violation(self) -> None:
assert AbsoluteBudget(name="p99 latency", measured=1.2, ceiling=5.0, unit="s", decimals=3).violation() is None
def test_a_value_exactly_at_the_ceiling_is_allowed(self) -> None:
assert AbsoluteBudget(name="p99 latency", measured=5.0, ceiling=5.0, unit="s", decimals=3).violation() is None
def test_a_value_past_the_ceiling_reports_the_measurement_and_the_ceiling(self) -> None:
violation: Final = AbsoluteBudget(
name="p99 latency", measured=9.5, ceiling=5.0, unit="s", decimals=3
).violation()
assert violation is not None
assert "9.500s" in violation
assert "5.000s allowed" in violation
def test_a_flat_ceiling_fails_a_degraded_phase_that_is_cheaper_than_its_baseline(self) -> None:
# The whole reason this shape exists: once the breaker opens, requests skip Redis instead
# of waiting on its socket timeout, so the chaos phase can measure faster than the healthy
# one. A ratio against that baseline passes; the user still waited 9.5s.
assert _budget(baseline=20.0, degraded=9.5, ceiling=2.0).violation() is None
assert AbsoluteBudget(name="p99 latency", measured=9.5, ceiling=5.0, unit="s").violation() is not None
def test_a_zero_measurement_is_not_a_violation(self) -> None:
assert AbsoluteBudget(name="log bytes per request", measured=0, ceiling=12_000, unit=" B").violation() is None
class TestViolations:
def test_every_blown_budget_is_reported_not_just_the_first(self) -> None:
blown: Final = violations(
(
_budget(baseline=100, degraded=500),
_budget(baseline=100, degraded=120),
RatioBudget(name="CPU per request", baseline=10, degraded=90, ratio_ceiling=6.0, unit=" ms"),
)
)
assert len(blown) == 2
assert blown[0].startswith("p99 RSS")
assert blown[1].startswith("CPU per request")
def test_both_budget_shapes_report_together(self) -> None:
blown: Final = violations(
(
_budget(baseline=100, degraded=500),
AbsoluteBudget(name="p99 latency", measured=9.5, ceiling=5.0, unit="s", decimals=3),
)
)
assert len(blown) == 2
assert blown[0].startswith("p99 RSS")
assert blown[1].startswith("p99 latency")
def test_a_run_inside_every_budget_reports_nothing(self) -> None:
assert violations((_budget(baseline=100, degraded=150),)) == ()

71
tests/e2e/load/test_proxy_usage.py Normal file
View file

@ -0,0 +1,71 @@
from __future__ import annotations
from typing import Final
from proxy_usage import UsageSample, UsageWindow
_MB: Final = 2**20
def _window(*points: tuple[float, int, float]) -> UsageWindow:
return UsageWindow(
samples=tuple(
UsageSample(elapsed_seconds=elapsed, rss_bytes=rss, cpu_seconds=cpu) for elapsed, rss, cpu in points
)
)
class TestRssPercentiles:
def test_the_tail_percentiles_reach_the_peak_the_median_hides(self) -> None:
# 100 one-second samples: 89 flat, 10 elevated, 1 spike. The median stays flat, p90 sees the
# elevated plateau, and only the max reaches the spike.
window: Final = _window(
*((float(i), 100 * _MB, float(i)) for i in range(89)),
*((float(89 + i), 300 * _MB, float(89 + i)) for i in range(10)),
(99.0, 900 * _MB, 99.0),
)
assert window.rss_percentile(0.5) == 100 * _MB
assert window.rss_percentile(0.9) == 300 * _MB
assert window.rss_percentile(0.99) == 300 * _MB
assert window.rss_percentile(1.0) == 900 * _MB
def test_an_empty_window_reports_zero_rather_than_indexing_nothing(self) -> None:
assert _window().rss_percentile(0.5) == 0
class TestCpuUtilization:
def test_utilization_is_the_counter_delta_over_the_interval_not_the_counter_itself(self) -> None:
# The counter climbs 0.5 CPU seconds per second, then 4.0 per second: half a core, then four.
window: Final = _window((0.0, _MB, 0.0), (1.0, _MB, 0.5), (2.0, _MB, 1.0), (3.0, _MB, 5.0))
p50, p90, p99 = window.cpu_utilization_percentiles()
assert (p50, p90, p99) == (0.5, 4.0, 4.0)
assert window.cpu_seconds_consumed() == 5.0
def test_a_single_sample_has_no_interval_and_reports_zero(self) -> None:
window: Final = _window((0.0, _MB, 3.0))
assert window.cpu_utilization_percentiles() == (0.0, 0.0, 0.0)
assert window.cpu_seconds_consumed() == 0.0
def test_cost_per_request_separates_runs_that_cores_busy_reports_identically(self) -> None:
# Both windows pin 4 cores for 10 seconds, so utilization cannot tell them apart. The
# second one served a tenth of the traffic for the same CPU, which is the regression shape.
window: Final = _window(*((float(i), _MB, 4.0 * i) for i in range(11)))
assert window.cpu_utilization_percentiles()[0] == 4.0
assert window.cpu_seconds_per_request(4000) == 0.01
assert window.cpu_seconds_per_request(400) == 0.1
def test_no_requests_reports_zero_cost_rather_than_dividing_by_zero(self) -> None:
assert _window((0.0, _MB, 0.0), (1.0, _MB, 1.0)).cpu_seconds_per_request(0) == 0.0
def test_summary_reports_every_percentile_in_human_units(self) -> None:
window: Final = _window((0.0, 200 * _MB, 0.0), (1.0, 200 * _MB, 1.5), (2.0, 200 * _MB, 3.0))
assert window.summary() == (
"RSS p50 200 MB, p90 200 MB, p99 200 MB; "
"CPU cores busy p50 1.50, p90 1.50, p99 1.50; 3.0 CPU seconds consumed"
)

454
tests/e2e/load/test_redis_chaos_e2e.py Normal file
View file

@ -0,0 +1,454 @@
"""Live e2e: the proxy under load keeps serving every request while Redis is down entirely.
Runs against a proxy booted from tests/e2e/gateway/redis_chaos_ci_config.yml, which points
cache_params at a real Redis with litellm's default socket_timeout. That one client backs all
three Redis touchpoints on the request path: the virtual-key auth cache, the response cache,
and the cross-pod spend counter the cost-tracking callback awaits.
The load runs in two phases against one model group of three mock deployments. The two at
order 1 raise InternalServerError and the one at order 2 serves, so every request burns its
retries on the failing pair (a 500 is retryable, so retries keep re-picking inside the lowest
order) and the router's order-based fallback then re-targets order 2. Every request is expected
to succeed, and each one carries retry breadcrumbs into cost tracking.
Traffic is split round robin between /chat/completions and /v1/messages, one endpoint per
simulated user: the Redis touchpoints and the cost-tracking callback are shared by both, but
the Anthropic Messages route reaches them through its own request path, so a regression that
only shows up there would not surface from chat completions alone.
Phase A is a baseline with Redis healthy; phase B holds Redis in CLIENT PAUSE ALL for the
length of the phase, simulating Redis being down outright rather than merely slow to write.
Every touchpoint times out: the auth cache read falls back to Postgres, the response cache
read and write both fail, and the spend counter increment times out and the callback
stringifies the request metadata, breadcrumbs included, into a failed-tracking alert. On
v1.100.0 that string doubled per request until the worker hung (LIT-6780), which is what the
per-phase RSS, CPU, and log-bytes budgets are here to catch.
Needs the proxy on the same host, since RSS and CPU come from psutil on its process tree:
a multi-worker proxy serves /metrics from the prometheus multiprocess collector, which drops
the process collector's memory and CPU series. Log bytes are read from the file the proxy's
stdout/stderr was redirected to, so the same host requirement covers that too. Deselected
unless E2E_REDIS_CHAOS is set.
"""
from __future__ import annotations
import os
import re
import time
from collections.abc import Iterator
from dataclasses import dataclass
from itertools import pairwise
from pathlib import Path
from typing import Final
import pytest
import redis
from e2e_config import PROXY_BASE_URL, unique_marker
from e2e_http import NoBody
from lifecycle import ResourceManager
from load_client import LoadClient
from locust_load import LoadResult, run_gateway_load
from models import KeyGenerateBody, LiteLLMParamsBody
from phase_budget import AbsoluteBudget, Budget, RatioBudget, violations
from proxy_client import ProxyClient
from proxy_usage import ProxyUsageSampler, UsageWindow
pytestmark: Final = pytest.mark.e2e
MODEL_GROUP: Final = f"redis-chaos-fable-{unique_marker()}"
MOCK_MODEL: Final = "anthropic/claude-fable-5-1"
FAILING_DEPLOYMENTS: Final = 2
SERVING_DEPLOYMENTS: Final = 1
FAILING_ORDER: Final = 1
SERVING_ORDER: Final = 2
KEY_POOL_SIZE: Final = 8
LOAD_ENDPOINTS: Final = ("/chat/completions", "/v1/messages")
LOCUST_USERS: Final = 50
LOCUST_SPAWN_RATE: Final = 50.0
BASELINE_SECONDS: Final = 60.0
CHAOS_SECONDS: Final = 90.0
REDIS_PAUSE_MS: Final = int(CHAOS_SECONDS * 1000)
# RSS and CPU are budgeted as a multiple of the same metric in the baseline phase, because both
# are machine-shaped: RSS scales with worker count and CPU with core count, so a number
# calibrated on one runner means nothing on another. RSS moved 0.91x-1.40x across three otherwise
# identical local runs, so it stays loose; CPU per request held steady at 1.33x-1.36x across the
# same runs, so it sits close to what is actually measured. That makes CPU the likeliest of these
# to flake first on a runner whose core count shifts how much of baseline CPU is fixed per-request
# work: loosen it rather than widening the others if a CI run trips it without a real cause.
CHAOS_RSS_RATIO_CEILING: Final = 2.0
CHAOS_CPU_PER_REQUEST_RATIO_CEILING: Final = 2.0
# Latency and log volume get flat ceilings instead, because a ratio cannot bound either one. Once
# the breaker opens, a request skips Redis rather than waiting on its socket timeout, so the chaos
# phase can come in faster than baseline (local runs measured p90 at 0.61x) and a ratio passes on a
# phase that was never slow. What a user actually cares about is the wall-clock number, which these
# hold directly. Calibrated from local runs whose worst chaos phase was p50 0.19s, p90 0.23s, p99
# 0.69s and 3.5 KB of log per request, with several times that left as slack for a shared CI runner.
CHAOS_P50_LATENCY_CEILING_SECONDS: Final = 1.0
CHAOS_P90_LATENCY_CEILING_SECONDS: Final = 2.0
CHAOS_P99_LATENCY_CEILING_SECONDS: Final = 3.0
CHAOS_LOG_BYTES_PER_REQUEST_CEILING: Final = 10_000.0
DRAIN_TIMEOUT_SECONDS: Final = 30.0
DRAIN_POLL_SECONDS: Final = 1.0
TIMEOUT_FAILURES_RE: Final = re.compile(
r'^litellm_redis_circuit_breaker_failures_total\{failure_class="timeout"\} ([0-9.e+]+)$', re.M
)
# The state gauge carries a pid label under the multiprocess collector, one series per worker,
# so this matches any label order rather than a bare {state="open"} that never appears.
BREAKER_OPEN_RE: Final = re.compile(
r'^litellm_redis_circuit_breaker_state\{[^}]*state="open"[^}]*\} ([0-9.e+]+)$', re.M
)
BREAKER_TRANSITIONS_RE: Final = re.compile(
r'^litellm_redis_circuit_breaker_transitions_total\{state="[a-z_]+"\} ([0-9.e+]+)$', re.M
)
def _deployment_metric_re(name: str, model_ids: tuple[str, ...]) -> re.Pattern[str]:
"""A per-deployment counter, narrowed to the deployments one run registered, so traffic
anything else sends the same proxy during the run cannot pad the retry count."""
ids: Final = "|".join(re.escape(model_id) for model_id in model_ids)
return re.compile(rf'^litellm_{name}\{{[^}}]*model_id="(?:{ids})"[^}}]*\}} ([0-9.e+]+)$', re.M)
@dataclass(frozen=True, slots=True)
class Phase:
"""One load phase's traffic and what the proxy's process tree did during it."""
name: str
load: LoadResult
usage: UsageWindow
redis_timeouts: float
log_bytes: int
@property
def timeouts_per_request(self) -> float:
return self.redis_timeouts / self.load.requests if self.load.requests else 0.0
@property
def cpu_seconds_per_request(self) -> float:
return self.usage.cpu_seconds_per_request(self.load.requests)
@property
def log_bytes_per_request(self) -> float:
return self.log_bytes / self.load.requests if self.load.requests else 0.0
def report(self) -> str:
return (
f"{self.name}: {self.load.requests} requests, {self.load.failures} failures, "
f"{self.load.requests_per_second:.0f} rps, {self.load.latency_summary()}; {self.usage.summary()}; "
f"{self.cpu_seconds_per_request * 1000:.1f} ms CPU per request; "
f"{self.log_bytes_per_request:.0f} log bytes per request; "
f"{self.timeouts_per_request:.2f} Redis timeouts per request; "
f"by endpoint: {self.load.endpoint_summary()}"
)
def _failing_params() -> LiteLLMParamsBody:
return LiteLLMParamsBody(
model=MOCK_MODEL,
api_key="sk-redis-chaos-not-used",
mock_response="litellm.InternalServerError",
order=FAILING_ORDER,
)
def _serving_params() -> LiteLLMParamsBody:
return LiteLLMParamsBody(
model=MOCK_MODEL,
api_key="sk-redis-chaos-not-used",
mock_response="redis chaos ok",
order=SERVING_ORDER,
)
@pytest.fixture
def proxy_pid() -> int:
"""The proxy's PID, which the workflow exports after starting it.
Required rather than discovered: picking a process out of the table by name would be
ambiguous on a developer machine running more than one proxy.
"""
pid: Final = os.environ.get("E2E_PROXY_PID")
assert pid and pid.isdigit(), (
"E2E_PROXY_PID must hold the PID of the proxy under test; RSS and CPU are read from "
"its process tree because a multi-worker proxy does not report them on /metrics"
)
return int(pid)
@pytest.fixture
def proxy_log() -> Path:
"""Path to the proxy's stdout/stderr log, which the workflow captures to a file.
Required rather than discovered for the same reason as proxy_pid: a developer machine may
have more than one proxy log around.
"""
path: Final = os.environ.get("E2E_PROXY_LOG")
assert path, "E2E_PROXY_LOG must hold the path the proxy's stdout/stderr was redirected to"
return Path(path)
def _log_bytes(path: Path) -> int:
return path.stat().st_size
@pytest.fixture
def redis_control() -> Iterator[redis.Redis[bytes]]:
"""A control connection to the proxy's Redis, which unpauses it in teardown as a safety net.
CLIENT PAUSE ALL freezes every connection including this one, so REDIS_PAUSE_MS is sized
to the chaos phase: by the time teardown runs, the pause has
already lapsed on its own and CLIENT UNPAUSE here returns immediately. It only actually
waits out a lapsed pause if the chaos phase itself overran that duration.
"""
host: Final = os.environ.get("REDIS_HOST")
port: Final = os.environ.get("REDIS_PORT")
assert host and port, "REDIS_HOST and REDIS_PORT must name the Redis the proxy under test uses"
control: Final = redis.Redis(host=host, port=int(port), socket_timeout=5)
try:
yield control
finally:
control.client_unpause() # pyright: ignore[reportUnknownMemberType] # redis-py stubs return Any
control.close()
def _scrape(proxy: ProxyClient) -> str:
"""One /metrics body, read once per checkpoint so every counter comes from the same instant."""
scrape: Final = proxy.probe("/metrics", params=NoBody())
assert scrape.status_code == 200, (
f"/metrics did not answer ({scrape.status_code}: {scrape.body[:200]}), so no counter can be read; "
f"a silent 0 here would turn every before-and-after difference negative"
)
return scrape.body
def _metric(scrape: str, pattern: re.Pattern[str]) -> float:
return sum(float(match.group(1)) for match in pattern.finditer(scrape))
def _scrape_after_drain(proxy: ProxyClient, pattern: re.Pattern[str]) -> str:
"""A /metrics body taken once `pattern`'s count has stopped moving.
`set_llm_deployment_failure_metrics` runs from the async logging callback queue, so a load
generator that just stopped sending traffic can still have thousands of failure increments
in flight, and a scrape taken the instant load stops undercounts them. Settling on the
counter rather than sleeping a fixed duration keeps the wait proportional to how backed up
the queue actually is.
"""
deadline: Final = time.monotonic() + DRAIN_TIMEOUT_SECONDS
def scrapes() -> Iterator[str]:
yield _scrape(proxy)
while time.monotonic() < deadline:
time.sleep(DRAIN_POLL_SECONDS)
yield _scrape(proxy)
settled: Final = next(
(later for earlier, later in pairwise(scrapes()) if _metric(earlier, pattern) == _metric(later, pattern)),
None,
)
return settled if settled is not None else _scrape(proxy)
def _register_deployments(proxy: ProxyClient, resources: ResourceManager) -> tuple[str, ...]:
"""The model ids this run registered, which scope its per-deployment metric reads."""
params: Final = (
*(_failing_params() for _ in range(FAILING_DEPLOYMENTS)),
*(_serving_params() for _ in range(SERVING_DEPLOYMENTS)),
)
model_ids: Final = tuple(proxy.create_model(MODEL_GROUP, one) for one in params)
for model_id in model_ids:
resources.defer(lambda doomed=model_id: proxy.delete_model(doomed))
return model_ids
def _generate_key_pool(proxy: ProxyClient, resources: ResourceManager) -> tuple[str, ...]:
"""A pool of virtual keys so auth and budget lookups are not one permanently warm
cache entry; each locust user picks one, so Redis auth reads actually happen."""
keys: Final = tuple(
proxy.generate_key(
KeyGenerateBody(models=[MODEL_GROUP], key_alias=f"e2e-redis-chaos-{unique_marker()}-{index}")
)
for index in range(KEY_POOL_SIZE)
)
for key in keys:
resources.defer(lambda doomed=key: proxy.delete_key(doomed))
return keys
def _drive(keys: tuple[str, ...], seconds: float) -> LoadResult:
return run_gateway_load(
base_url=PROXY_BASE_URL,
api_keys=keys,
model=MODEL_GROUP,
endpoints=LOAD_ENDPOINTS,
users=LOCUST_USERS,
spawn_rate=LOCUST_SPAWN_RATE,
duration_seconds=seconds,
)
def _latency_budget(percentile: str, measured: float, ceiling: float) -> Budget:
return AbsoluteBudget(name=f"{percentile} latency", measured=measured, ceiling=ceiling, unit="s", decimals=3)
def _rss_budget(percentile: str, baseline: UsageWindow, degraded: UsageWindow, fraction: float) -> Budget:
return RatioBudget(
name=f"{percentile} RSS",
baseline=baseline.rss_percentile(fraction) / 2**20,
degraded=degraded.rss_percentile(fraction) / 2**20,
ratio_ceiling=CHAOS_RSS_RATIO_CEILING,
unit=" MB",
decimals=0,
)
def _chaos_budgets(baseline: Phase, chaos: Phase) -> tuple[Budget, ...]:
"""What a Redis outage is allowed to cost.
Every request still succeeding is the headline assertion, but a proxy can answer every
request while leaking: the v1.100.0 regression (LIT-6780) served traffic the whole way up
to a 61 GB worker. These bound the cost of serving it. RSS and CPU are bounded against the
same run's healthy phase, latency and log bytes against a flat ceiling; see phase_budget
for why the two kinds of metric cannot share one shape.
Latency and RSS are budgeted at p50, p90 and p99 so a regression that only shows up in the
tail (or only in the median) cannot hide behind the other. RSS gets the tightest bound: the
failure path has no business allocating more per request. CPU and log bytes are each budgeted
once, as an amount per request rather than per percentile: cores-busy saturates at the worker
count under load, so its percentiles read the same whether a request costs 10 ms of CPU or
40, and cannot budget anything; per-request is the figure that actually moves. Log bytes
isolates the cost of the failed-tracking alert's own noisy error handling from the CPU it
burns doing useful retry work, since the two would otherwise be indistinguishable in one
CPU number.
"""
return (
_latency_budget("p50", chaos.load.p50_seconds, CHAOS_P50_LATENCY_CEILING_SECONDS),
_latency_budget("p90", chaos.load.p90_seconds, CHAOS_P90_LATENCY_CEILING_SECONDS),
_latency_budget("p99", chaos.load.p99_seconds, CHAOS_P99_LATENCY_CEILING_SECONDS),
_rss_budget("p50", baseline.usage, chaos.usage, 0.5),
_rss_budget("p90", baseline.usage, chaos.usage, 0.9),
_rss_budget("p99", baseline.usage, chaos.usage, 0.99),
RatioBudget(
name="CPU per request",
baseline=baseline.cpu_seconds_per_request * 1000,
degraded=chaos.cpu_seconds_per_request * 1000,
ratio_ceiling=CHAOS_CPU_PER_REQUEST_RATIO_CEILING,
unit=" ms",
),
AbsoluteBudget(
name="log bytes per request",
measured=chaos.log_bytes_per_request,
ceiling=CHAOS_LOG_BYTES_PER_REQUEST_CEILING,
unit=" B",
decimals=0,
),
)
@pytest.mark.redis_chaos
class TestRedisChaos:
@pytest.mark.covers(
"reliability.circuit_breaker.redis_timeout.stays_responsive",
exercised_on=("chat_completions", "messages"),
)
def test_load_survives_redis_being_down(
self,
client: LoadClient,
resources: ResourceManager,
proxy_pid: int,
proxy_log: Path,
redis_control: redis.Redis[bytes],
) -> None:
proxy: Final = client.proxy
model_ids: Final = _register_deployments(proxy, resources)
keys: Final = _generate_key_pool(proxy, resources)
retries_re: Final = _deployment_metric_re("deployment_failure_responses_total", model_ids)
cooldown_re: Final = _deployment_metric_re("deployment_cooled_down_total", model_ids)
at_start: Final = _scrape(proxy)
log_at_start: Final = _log_bytes(proxy_log)
with ProxyUsageSampler(proxy_pid) as sampler:
baseline_load: Final = _drive(keys, BASELINE_SECONDS)
baseline_usage: Final = sampler.split()
after_baseline: Final = _scrape(proxy)
log_after_baseline: Final = _log_bytes(proxy_log)
redis_control.client_pause(REDIS_PAUSE_MS, all=True) # pyright: ignore[reportUnknownMemberType] # redis-py stubs return Any
chaos_load: Final = _drive(keys, CHAOS_SECONDS)
chaos_usage: Final = sampler.split()
at_end: Final = _scrape_after_drain(proxy, retries_re)
log_at_end: Final = _log_bytes(proxy_log)
baseline: Final = Phase(
name="baseline",
load=baseline_load,
usage=baseline_usage,
redis_timeouts=_metric(after_baseline, TIMEOUT_FAILURES_RE) - _metric(at_start, TIMEOUT_FAILURES_RE),
log_bytes=log_after_baseline - log_at_start,
)
chaos: Final = Phase(
name="chaos",
load=chaos_load,
usage=chaos_usage,
redis_timeouts=_metric(at_end, TIMEOUT_FAILURES_RE) - _metric(after_baseline, TIMEOUT_FAILURES_RE),
log_bytes=log_at_end - log_after_baseline,
)
report: Final = f"{baseline.report()} | {chaos.report()}"
for phase in (baseline, chaos):
assert phase.load.requests > 0, (
f"{phase.name} drove no traffic at all, so it proved nothing: {phase.load.diagnosis()}. {report}"
)
assert frozenset(endpoint.name for endpoint in phase.load.endpoints) == frozenset(LOAD_ENDPOINTS), (
f"{phase.name} drove {tuple(endpoint.name for endpoint in phase.load.endpoints)} rather than every "
f"endpoint in {LOAD_ENDPOINTS}; the round robin hands one endpoint to each simulated user, so a "
f"missing one means a route never ran and its request path was never exercised. {report}"
)
assert phase.load.failures == 0, (
f"{phase.name} had {phase.load.failures} of {phase.load.requests} requests fail. Every request "
f"must succeed: the failing deployments sit at order {FAILING_ORDER} and the serving one at order "
f"{SERVING_ORDER}, so once the retries on order {FAILING_ORDER} are spent the order-based fallback "
f"lands on the serving deployment. Failures mean it was cooled down, the fallback did not run, or "
f"a Redis failure reached the response path. {phase.load.diagnosis()}. {report}"
)
cooldowns: Final = _metric(at_end, cooldown_re) - _metric(at_start, cooldown_re)
assert cooldowns == 0, (
f"{cooldowns:.0f} deployments were cooled down during the run; the failing deployments are supposed "
f"to stay in rotation so every request keeps exercising the retry path. {report}"
)
retries: Final = _metric(at_end, retries_re) - _metric(at_start, retries_re)
assert retries >= baseline.load.requests + chaos.load.requests, (
f"only {retries:.0f} deployment failures were counted across "
f"{baseline.load.requests + chaos.load.requests} requests; the mock deployments did not fail, so no "
f"request carried retry breadcrumbs into cost tracking and the regression path was never entered. "
f"{report}"
)
transitions: Final = _metric(at_end, BREAKER_TRANSITIONS_RE) - _metric(after_baseline, BREAKER_TRANSITIONS_RE)
breaker_open: Final = _metric(at_end, BREAKER_OPEN_RE) >= 1
assert transitions >= 1 or breaker_open, (
f"pausing Redis produced no circuit breaker state transitions and it ended closed; nothing on the "
f"request path ever saw Redis fail, so this run proved nothing. {report}"
)
blown: Final = violations(_chaos_budgets(baseline, chaos))
assert not blown, (
f"pausing Redis cost the proxy more than a Redis outage is allowed to: {'; '.join(blown)}. {report}"
)
rows: Final = proxy.poll_logs_for_key(keys[0], min_rows=1)
assert rows, (
f"no spend rows landed for the first key in the pool; a Redis outage must not cost the proxy its "
f"spend logs, which are written to Postgres through a queue rather than through Redis. {report}"
)
print(f"\nredis chaos load: {report}") # noqa: T201 # the numbers this test exists to report, read off the CI log

3
tests/e2e/models.py
View file

@ -922,10 +922,11 @@ class LiteLLMParamsBody(BaseModel):
auto_router_default_model: str | None = None
auto_router_embedding_model: str | None = None
tags: list[str] | None = None
mock_response: str | None = None
mock_response: str | list[float] | None = None
timeout: float | None = None
tpm: int | None = None
weight: int | None = None
order: int | None = None
ModelMode = Literal["batch", "realtime", "image_generation"]

1
tests/e2e/pytest.ini
View file

@ -9,3 +9,4 @@ markers =
load: heavy throughput/load test; collected last so it never perturbs latency-sensitive suites
weekly: real-provider anomaly load test that spends real money; deselected unless E2E_WEEKLY_ANOMALY is set
managed_files: needs a proxy running with require_managed_files enabled; deselected unless E2E_MANAGED_FILES_STACK is set
redis_chaos: load test that pauses the proxy's Redis outright mid-run; needs a proxy booted from gateway/redis_chaos_ci_config.yml on the same host, and is deselected unless E2E_REDIS_CHAOS is set

9
tests/e2e/router/conftest.py
View file

@ -13,10 +13,7 @@ from __future__ import annotations
from collections.abc import Iterator
import pytest
from requests import RequestException
from complexity_router_client import ComplexityRouterClient, build_client
from proxy_client import ProxyClient
from e2e_http import NoBody, Success
from lifecycle import ResourceManager
from models import (
@ -26,6 +23,8 @@ from models import (
LiteLLMParamsBody,
ModelsListResponse,
)
from proxy_client import ProxyClient
from requests import RequestException
ROUTER_MODEL = "complexity-smart-router"
ROUTER_PARAMS = LiteLLMParamsBody(
@ -120,8 +119,6 @@ def _ensure_complexity_smart_router( # pyright: ignore[reportUnusedFunction] #
@pytest.fixture
def complexity_key(resources: ResourceManager, client: ComplexityRouterClient) -> str:
"""Per-test key allowed to call the complexity router and its tier backends."""
key = client.proxy.generate_key(
KeyGenerateBody(models=ROUTER_KEY_MODELS, user_id="e2e-complexity-router")
)
key = client.proxy.generate_key(KeyGenerateBody(models=ROUTER_KEY_MODELS, user_id="e2e-complexity-router"))
resources.defer(lambda: client.proxy.delete_key(key))
return key

598
tests/test_litellm/enterprise/enterprise_callbacks/test_secret_detection.py
View file

@ -10,12 +10,15 @@ Covers the three defects from the ticket:
handling live only on the native path).
"""
import tempfile
import time
import pytest
from litellm_enterprise.enterprise_callbacks.secret_detection import (
_ENTERPRISE_SecretDetection,
_default_detect_secrets_config,
_masked_entity_count,
)
from litellm.caching.caching import DualCache
from litellm.proxy._types import UserAPIKeyAuth
@ -29,6 +32,13 @@ URL_ENCODED_KEY = "Bearer%20sk-Ab3dEf6Gh7Ij8Kl9Mn0Pq2Rs3Tu4Vw5X"
AWS_KEYS = [f"AKIAIOSFODNN7EXAMPL{suffix}" for suffix in "FEDCBA"]
@pytest.fixture(autouse=True)
def _isolate_masked_entity_count():
token = _masked_entity_count.set(None)
yield
_masked_entity_count.reset(token)
def _guardrail() -> _ENTERPRISE_SecretDetection:
return _ENTERPRISE_SecretDetection(guardrail_name="hide-secrets", event_hook="pre_call", default_on=True)
@ -58,6 +68,561 @@ def test_scan_message_preserves_quoted_benign_identifiers():
assert guardrail.redact_text(content) == content
@pytest.mark.parametrize(
"content,secret",
[
("REDIS_PASSWORD=aB3dE6gH9jK2mN5p", "aB3dE6gH9jK2mN5p"),
("SESSION_SECRET=Kp7Nq2Wz9Bt4Xr6Vm1Ls", "Kp7Nq2Wz9Bt4Xr6Vm1Ls"),
('{"db_password": "Tq8Zm2XpLv9KdNbRcYw3"}', "Tq8Zm2XpLv9KdNbRcYw3"),
("api_secret: Zx4Kp9Lm2Qr7Ns3Vt", "Zx4Kp9Lm2Qr7Ns3Vt"),
("password = hunter2brahms9x", "hunter2brahms9x"),
("client_secret=Hq7Zm3XkLp9Wd2Nb", "Hq7Zm3XkLp9Wd2Nb"),
('apiKey: "aB3dE6gH9jK2mN5p"', "aB3dE6gH9jK2mN5p"),
('{"clientSecret": "Kp7Nq2Wz9Bt4Xr6Vm1Ls"}', "Kp7Nq2Wz9Bt4Xr6Vm1Ls"),
('dbPassword = "Zx4Kp9Lm2Qr7Ns3Vt"', "Zx4Kp9Lm2Qr7Ns3Vt"),
("MY_APP_DB_PASSWORD=Kp7Nq2Wz9Bt4Xr6Vm1Ls", "Kp7Nq2Wz9Bt4Xr6Vm1Ls"),
("x_api_key: 8f3Kd9Lm2Qr7Ns3Vt", "8f3Kd9Lm2Qr7Ns3Vt"),
("password: Zm9vYmFyYmF6+abc/def123=", "Zm9vYmFyYmF6+abc/def123="),
("REDIS_PASSWORD=correcthorsebattery", "correcthorsebattery"),
('SECRET_KEY = "django-insecure-9v2xk4qw8z"', "django-insecure-9v2xk4qw8z"),
(
"aws_secret_access_key = wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY",
"wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY",
),
("password=aB3dE6gH9jK2", "aB3dE6gH9jK2"),
("api_key: hunter2!brahms", "hunter2!brahms"),
('db_password: "p@ssw0rd!2026"', "p@ssw0rd!2026"),
(
'url: "postgresql://user:s3cr3t@db-host:5432/app"',
"postgresql://user:s3cr3t@db-host:5432/app",
),
(
'db_password: "postgresql://user:s3cr3t@db-host:5432/app"',
"postgresql://user:s3cr3t@db-host:5432/app",
),
(
'signing_secret_url: "https://example.com/cb?sig=Zx4Kp9Lm2Qr7Ns3Vt"',
"https://example.com/cb?sig=Zx4Kp9Lm2Qr7Ns3Vt",
),
(
'redis_secret_url: "redis://:Zx4Kp9Lm2Qr7Ns3Vt@cache-host:6379/0"',
"Zx4Kp9Lm2Qr7Ns3Vt",
),
("password=2026-09-08T17:38:40Zbrahms", "2026-09-08T17:38:40Zbrahms"),
(
'{"password": "YOUR_API_KEY_HERE", "client_secret": "correcthorsebattery"}',
"correcthorsebattery",
),
("docker run -e REDIS_PASSWORD=aB3dE6gH9jK2mN5p \\\n -e REDIS_PORT=6379 redis", "aB3dE6gH9jK2mN5p"),
("DB_PASSWORD=Zx4Kp9Lm2Qr7Ns3Vt && echo done", "Zx4Kp9Lm2Qr7Ns3Vt"),
("password = Zx4Kp9Lm2Qr7Ns3Vt # rotate me", "Zx4Kp9Lm2Qr7Ns3Vt"),
("my db password: Zx4Kp9Lm2Qr7Ns3Vt.", "Zx4Kp9Lm2Qr7Ns3Vt"),
("export DB_PASSWORD=Zx4Kp9Lm2Qr7Ns3Vt DB_HOST=db.internal", "Zx4Kp9Lm2Qr7Ns3Vt"),
("DB_PASSWORD=Zx4Kp9Lm2Qr7Ns3Vt; systemctl restart app", "Zx4Kp9Lm2Qr7Ns3Vt"),
("DB_PASSWORD=Zx4Kp9Lm2Qr7Ns3Vt | tee creds.txt", "Zx4Kp9Lm2Qr7Ns3Vt"),
("DB_PASSWORD=Zx4Kp9Lm2Qr7Ns3Vt > setup.log", "Zx4Kp9Lm2Qr7Ns3Vt"),
("docker run -e DB_PASSWORD=Zx4Kp9Lm2Qr7Ns3Vt --name app postgres", "Zx4Kp9Lm2Qr7Ns3Vt"),
("password=correcthorsebattery please", "correcthorsebattery"),
("DB_PASSWORD = Zx4Kp9Lm2Qr7Ns3Vt; systemctl restart app", "Zx4Kp9Lm2Qr7Ns3Vt"),
("DB_PASSWORD = Zx4Kp9Lm2Qr7Ns3Vt \\", "Zx4Kp9Lm2Qr7Ns3Vt"),
("DB_PASSWORD = Zx4Kp9Lm2Qr7Ns3Vt DB_HOST=db.internal", "Zx4Kp9Lm2Qr7Ns3Vt"),
("DB_PASSWORD = Zx4Kp9Lm2Qr7Ns3Vt --db-host=db.internal", "Zx4Kp9Lm2Qr7Ns3Vt"),
("DB_PASSWORD = Zx4Kp9Lm2Qr7Ns3Vt DEBUG=", "Zx4Kp9Lm2Qr7Ns3Vt"),
],
ids=[
"env-password",
"env-secret",
"json-field",
"yaml-field",
"bare-assignment",
"client-secret",
"camel-case-key",
"camel-case-secret",
"camel-case-password",
"namespaced-env",
"underscored-header",
"base64-padding",
"digit-free-value",
"django-secret-key",
"slashed-aws-secret",
"shortest-accepted-value",
"punctuation-bearing-password",
"symbol-heavy-password",
"connection-string-under-a-url-key",
"connection-string-under-a-credential-key",
"signed-url-under-a-credential-key",
"password-only-url-under-a-credential-key",
"timestamp-prefixed-password",
"credential-after-a-rejected-placeholder",
"docker-flag-with-a-line-continuation",
"shell-command-after-the-value",
"inline-comment-after-the-value",
"sentence-ending-in-the-value",
"second-assignment-after-the-value",
"semicolon-after-the-value",
"pipe-after-the-value",
"redirect-after-the-value",
"docker-flag-after-the-value",
"prose-after-a-shell-assignment",
"spaced-assignment-then-a-shell-command",
"spaced-assignment-then-a-line-continuation",
"spaced-assignment-then-a-second-assignment",
"spaced-assignment-then-a-dashed-flag",
"spaced-assignment-then-an-empty-assignment",
],
)
def test_scan_message_redacts_credentials_assigned_to_credential_keys(content, secret):
guardrail = _guardrail()
assert secret not in guardrail.redact_text(content)
def test_scan_message_redacts_only_the_first_token_of_a_shell_assignment():
guardrail = _guardrail()
content = "docker run -e REDIS_PASSWORD=aB3dE6gH9jK2mN5p \\\n -e REDIS_PORT=6379 redis && echo done"
assert (
guardrail.redact_text(content)
== "docker run -e REDIS_PASSWORD=[REDACTED] \\\n -e REDIS_PORT=6379 redis && echo done"
)
@pytest.mark.parametrize("operator", [";", "&&", "|"])
def test_scan_message_keeps_a_shell_operator_glued_to_the_value(operator):
guardrail = _guardrail()
assert (
guardrail.redact_text(f"DB_PASSWORD=Zx4Kp9Lm2Qr7Ns3Vt{operator} systemctl restart app")
== f"DB_PASSWORD=[REDACTED]{operator} systemctl restart app"
)
def test_scan_message_closes_a_yaml_block_at_the_next_unindented_line():
guardrail = _guardrail()
content = "api_key: >\n aB3dE6gH9jK2mN5p\nSteps\n Rotate-Before-Friday please"
assert guardrail.redact_text(content) == "api_key: >\n [REDACTED]\nSteps\n Rotate-Before-Friday please"
def test_scan_message_redacts_every_credential_on_one_line():
guardrail = _guardrail()
content = '{"db_password": "Tq8Zm2XpLv9KdNbRcYw3", "client_secret": "correcthorsebattery"}'
assert guardrail.redact_text(content) == '{"db_password": "[REDACTED]", "client_secret": "[REDACTED]"}'
@pytest.mark.parametrize(
"content",
[
"The user forgot their password and asked for a reset link",
"Rotate the client secret every 90 days",
"The secret: keep it quiet",
"My password: correct horse battery staple",
"secretary: Maria Gonzalez",
"password_reset_email: Please click the link below to reset",
'config = {"api_key": "YOUR_API_KEY_HERE"}',
"api_key: <your-key-here>",
'{"max_tokens": 4096, "model": "gpt-4o-mini"}',
'def get_api_key():\n return os.environ["OPENAI_API_KEY"]',
' valid_token = UserAPIKeyAuth(user_id="u1")',
'password = get_password(user, "prod")',
"monkey=aB3dE6gH9jK2mN5p",
"idempotency_key: req_2026090712000000",
'cache_key = "u1_user_api_key_user_id"',
"the key: 2026-09-07T12:00:00Z",
"api_key: os.environ/E2B_API_KEY",
"langfuse_secret: os.environ/LANGFUSE_PROJECT1_SECRET",
"api_key = OPENAI_API_KEY",
"password = pwd12345678",
"model_key: gpt-4o-mini-2024-07-18",
"openrouter/anthropic/claude-3-5-sonnet-20240620",
'{"content-type": "application/json"}',
"passwordless_login: enabled-for-all-users",
'password: "I forgot mine, can you reset it"',
"secret_sauce: tomatoes-basil-garlic-oregano",
"user_secret_question: what-was-your-first-pet",
"password_reset_url: example.com/reset-password/flow",
"private_key_path: keys/prod/server-cert.pem",
"litellm.completion(model=model, api_key=openai_api_key)",
"params['aws_secret_access_key'] = aws_secret_access_key",
'api_key = "OPENAI_API_KEY"',
"model_list:\n - litellm_params:\n api_key: 'PERPLEXITY_API_KEY'",
'config = build(_provider("ve_missing", api_key_env="VE_MISSING_KEY"))',
"api_key = get_api_key_from_env()",
"api_key = get_secret_str(MISTRAL_OCR_API_KEY_ENV_VAR)",
"secret_manager = MagicMock(spec=BaseSecretManager)",
"api_key = self.resolve_server_api_key(",
"api_key = sys.argv[1]",
"password = credentials[environment]",
'api_key_created_at: "2026-09-08T17:38:40Z"',
'api_key_expires_at: "2026-09-08T17:38:40.123456+05:30"',
'password_reset_url: "https://example.com/reset-password/flow"',
'secret_docs_url: "https://example.com/reset-password/flow#step-2"',
'{"api_key_created_at": "2026-09-08T17:38:40Z", "password_reset_url": "https://example.com/reset/flow"}',
"secret_sauce: tomatoes-basil-garlic-oregano.",
"secret_docs_url: https://example.com/docs/keys, then rotate",
"api_key_created_at: 2026-09-08T17:38:40Z; api_key_env: OPENAI_API_KEY!",
"api_key: $OPENAI_API_KEY",
'api_key: "${OPENAI_API_KEY}"',
"private_key_path: /keys/prod/server-cert.pem",
"password_hint: your usual one followed by Ticket-LIT7049-Suffix",
"Translate this recipe note into French:\nsecret_sauce: Worcestershire sauce",
"api_key = Massachusetts (the state, not a key)",
"secret_sauce:Worcestershire sauce",
"password: correctHorseBattery != anotherValue",
],
ids=[
"prose-password",
"prose-secret",
"colon-prose-secret",
"colon-prose-password",
"secretary",
"sentence-after-keyword",
"uppercase-placeholder",
"templated-placeholder",
"max-tokens",
"code-paste",
"constructor-call",
"indirect-reference",
"word-ending-in-key",
"idempotency-key",
"cache-key",
"timestamp-after-key",
"env-reference",
"env-reference-nested",
"env-variable-name",
"below-minimum-length",
"model-name",
"namespaced-model-name",
"media-type",
"hyphenated-english",
"quoted-sentence-under-a-credential-key",
"hyphenated-phrase",
"hyphenated-question",
"url-under-credential-key",
"path-under-credential-key",
"snake-case-argument",
"snake-case-assignment",
"quoted-env-variable-name",
"quoted-env-name-in-a-config",
"quoted-env-name-in-a-code-paste",
"bare-call",
"call-with-an-argument",
"keyword-argument-call",
"unclosed-call",
"positional-subscript",
"keyed-subscript",
"timestamp-under-a-credential-key",
"offset-timestamp-under-a-credential-key",
"url-under-a-credential-key",
"fragment-url-under-a-credential-key",
"metadata-object-under-credential-keys",
"hyphenated-english-ending-a-sentence",
"url-followed-by-a-clause",
"timestamp-and-env-name-with-trailing-punctuation",
"shell-variable-reference",
"quoted-braced-shell-variable-reference",
"absolute-path-under-a-credential-key",
"sentence-holding-a-later-mixed-case-token",
"capitalized-word-starting-a-phrase",
"capitalized-word-before-a-parenthetical",
"yaml-scalar-without-a-space-after-the-colon",
"comparison-operator-after-the-value",
],
)
def test_scan_message_keeps_benign_values(content):
guardrail = _guardrail()
assert guardrail.scan_message_for_secrets(content) == []
assert guardrail.redact_text(content) == content
@pytest.mark.parametrize(
"value,redacted",
[("aB3dE6gH9jK2", True), ("aB3dE6gH9jK", False)],
ids=["at-minimum-length", "below-minimum-length"],
)
def test_credential_keyword_detector_honours_its_minimum_length(value, redacted):
guardrail = _guardrail()
assert (value not in guardrail.redact_text(f"password={value}")) is redacted
@pytest.mark.parametrize(
"value,redacted",
[("aB3dE6gH9jK2", True), ("aB3dE6gH9jK", False)],
ids=["at-default-minimum-length", "below-default-minimum-length"],
)
def test_credential_keyword_detector_defaults_its_minimum_length(value, redacted):
guardrail = _ENTERPRISE_SecretDetection(
guardrail_name="hide-secrets",
event_hook="pre_call",
default_on=True,
detect_secrets_config={
"plugins_used": [
{key: setting for key, setting in plugin.items() if key != "minimum_length"}
for plugin in _default_detect_secrets_config["plugins_used"]
]
},
)
assert (value not in guardrail.redact_text(f"password={value}")) is redacted
def test_credential_keyword_detector_honours_keyword_exclude():
guardrail = _ENTERPRISE_SecretDetection(
guardrail_name="hide-secrets",
event_hook="pre_call",
default_on=True,
detect_secrets_config={
"plugins_used": [
{**plugin, "keyword_exclude": "fixture_"} if plugin["name"] == "CredentialKeywordDetector" else plugin
for plugin in _default_detect_secrets_config["plugins_used"]
]
},
)
content = "fixture_password=aB3dE6gH9jK2mN5p\npassword=Kp7Nq2Wz9Bt4Xr6Vm1Ls"
assert guardrail.redact_text(content) == "fixture_password=aB3dE6gH9jK2mN5p\npassword=[REDACTED]"
@pytest.mark.parametrize("minimum_length", ["12", 0, -1, 1.5], ids=["string", "zero", "negative", "float"])
def test_credential_keyword_detector_rejects_an_unusable_minimum_length(minimum_length):
guardrail = _ENTERPRISE_SecretDetection(
guardrail_name="hide-secrets",
event_hook="pre_call",
default_on=True,
detect_secrets_config={
"plugins_used": [
{**plugin, "minimum_length": minimum_length}
if plugin["name"] == "CredentialKeywordDetector"
else plugin
for plugin in _default_detect_secrets_config["plugins_used"]
]
},
)
with pytest.raises(ValueError, match="minimum_length"):
guardrail.scan_message_for_secrets("password=aB3dE6gH9jK2mN5p")
@pytest.mark.parametrize(
"content",
[
"[db\npassword = Zx4Kp9Lm2Qr7Ns3Vt\n",
"[\npassword = Zx4Kp9Lm2Qr7Ns3Vt\n",
"[note] have a look\npassword = Zx4Kp9Lm2Qr7Ns3Vt\n",
"]\npassword = Zx4Kp9Lm2Qr7Ns3Vt\n",
"[]\npassword = Zx4Kp9Lm2Qr7Ns3Vt\n",
],
ids=["unclosed", "bare-bracket", "bracketed-prose", "stray-close", "empty-header"],
)
def test_scan_message_reads_a_config_with_a_broken_section_header(content):
guardrail = _guardrail()
assert "Zx4Kp9Lm2Qr7Ns3Vt" not in guardrail.redact_text(content)
@pytest.mark.parametrize(
"content",
[
"=orphan\npassword = Zx4Kp9Lm2Qr7Ns3Vt\n",
" indented before any key\npassword = Zx4Kp9Lm2Qr7Ns3Vt\n",
"greeting = %(name)s\npassword = Zx4Kp9Lm2Qr7Ns3Vt\n",
"token = a\x00b\npassword = Zx4Kp9Lm2Qr7Ns3Vt\n",
],
ids=["empty-key", "leading-continuation", "interpolation", "nul-byte"],
)
def test_scan_message_reads_lines_that_a_stock_ini_parser_rejects(content):
guardrail = _guardrail()
assert "Zx4Kp9Lm2Qr7Ns3Vt" not in guardrail.redact_text(content)
def test_scan_message_reads_a_config_that_repeats_a_section():
guardrail = _guardrail()
content = "[db]\nhost = localhost\n[db]\npassword = Zx4Kp9Lm2Qr7Ns3Vt\n"
assert "Zx4Kp9Lm2Qr7Ns3Vt" not in guardrail.redact_text(content)
def test_scan_message_keeps_every_value_when_a_config_repeats_a_key():
guardrail = _guardrail()
content = (
"model_list:\n"
" - model_name: gpt-4o\n litellm_params:\n api_key: aB3dE6gH9jK2mN5p\n"
" - model_name: claude\n litellm_params:\n api_key: Kp7Nq2Wz9Bt4Xr6Vm1Ls\n"
)
redacted = guardrail.redact_text(content)
assert "aB3dE6gH9jK2mN5p" not in redacted
assert "Kp7Nq2Wz9Bt4Xr6Vm1Ls" not in redacted
@pytest.mark.parametrize(
"content,secret",
[
(
f"api_key: {OPENAI_KEY}\nREDIS_PASSWORD=aB3dE6gH9jK2mN5p",
"aB3dE6gH9jK2mN5p",
),
(
f"OPENAI_API_KEY={OPENAI_KEY}\nDB_PASSWORD=Kp7Nq2Wz9Bt4Xr6Vm1Ls",
"Kp7Nq2Wz9Bt4Xr6Vm1Ls",
),
(
f"api_key: {OPENAI_KEY}\npassword =\n Zx4Kp9Lm2Qr7Ns3Vt",
"Zx4Kp9Lm2Qr7Ns3Vt",
),
(
"Here is my config, can you review it?\nREDIS_PASSWORD=aB3dE6gH9jK2mN5p",
"aB3dE6gH9jK2mN5p",
),
(
"REDIS_PASSWORD=aB3dE6gH9jK2mN5p\nCan you tell me what is wrong with it?",
"aB3dE6gH9jK2mN5p",
),
(
"Hi team\nplease rotate this before Friday\ndb_password=Zx4Kp9Lm2Qr7Ns3Vt\nthanks!",
"Zx4Kp9Lm2Qr7Ns3Vt",
),
(
"model_list:\n - model_name: gpt-4o\n litellm_params:\n api_key: aB3dE6gH9jK2mN5p\n",
"aB3dE6gH9jK2mN5p",
),
("api_key: >\n aB3dE6gH9jK2mN5p", "aB3dE6gH9jK2mN5p"),
("api_key: |-\n aB3dE6gH9jK2mN5p", "aB3dE6gH9jK2mN5p"),
("secret= \\\n aB3dE6gH9jK2mN5p", "aB3dE6gH9jK2mN5p"),
("password =\n# rotate me\n Zx4Kp9Lm2Qr7Ns3Vt", "Zx4Kp9Lm2Qr7Ns3Vt"),
("api_key =\n; rotate me\n aB3dE6gH9jK2mN5p", "aB3dE6gH9jK2mN5p"),
(" # pasted from the vault\napi_key=aB3dE6gH9jK2mN5p", "aB3dE6gH9jK2mN5p"),
(" [db]\napi_key=aB3dE6gH9jK2mN5p", "aB3dE6gH9jK2mN5p"),
(" pasted with a leading indent\napi_key=aB3dE6gH9jK2mN5p", "aB3dE6gH9jK2mN5p"),
],
ids=[
"flat-assignment",
"env-file",
"continuation-line",
"prose-before",
"prose-after",
"prose-both-sides",
"indented-config",
"yaml-folded-block",
"yaml-literal-block",
"backslash-continuation",
"comment-inside-a-value",
"semicolon-comment-inside-a-value",
"indented-comment-above",
"indented-section-header-above",
"indented-prose-above",
],
)
def test_scan_message_still_sees_assignments_sharing_a_message_with_a_vendor_key(content, secret):
guardrail = _guardrail()
redacted = guardrail.redact_text(content)
assert secret not in redacted
assert OPENAI_KEY not in redacted
def test_environment_reference_filter_only_drops_the_whole_value():
guardrail = _guardrail()
for reference in ("os.environ/OPENAI_API_KEY", "os.environ/e2b_api_key"):
assert guardrail.redact_text(f"password={reference}") == f"password={reference}"
assert guardrail.redact_text("password=notos.environ/OPENAI_API_KEY") == ("password=[REDACTED]")
def test_environment_variable_names_are_dropped_only_for_the_keyword_plugin():
guardrail = _guardrail()
assert guardrail.redact_text("password=REDIS_PASSWORD") == "password=REDIS_PASSWORD"
assert guardrail.scan_message_for_secrets('k = "ABCD1234_EFGH5678_IJKLMN"') == [
{"type": "Base64 High Entropy String", "value": "ABCD1234_EFGH5678_IJKLMN"}
]
def test_masked_entity_count_keeps_the_vendor_type_beside_the_entropy_type():
guardrail = _guardrail()
_masked_entity_count.set({})
guardrail.redact_text('k = "ghp_abcdefghijklmnopqrstuvwxyzABCDEF1234"')
assert _masked_entity_count.get() == {
"Base64 High Entropy String": 1,
"GitHub Token": 1,
}
@pytest.mark.parametrize(
"content",
[
f"api_key: '{OPENAI_KEY}'\n"
+ "a: &a ["
+ ", ".join(['"x"'] * 9)
+ "]\n"
+ "".join(f"{chr(98 + i)}: &{chr(98 + i)} [" + ", ".join([f"*{chr(97 + i)}"] * 9) + "]\n" for i in range(7)),
f"api_key: '{OPENAI_KEY}'\ndeep: " + "[" * 400 + "]" * 400,
f"api_key: '{OPENAI_KEY}'\nbroken: [unclosed",
],
ids=["anchor-expansion", "deep-nesting", "unparseable"],
)
def test_scan_message_contains_hostile_config_text(content, monkeypatch, tmp_path):
guardrail = _guardrail()
monkeypatch.setenv("TMPDIR", str(tmp_path))
monkeypatch.setattr(tempfile, "tempdir", None)
started = time.perf_counter()
found = guardrail.scan_message_for_secrets(content)
assert time.perf_counter() - started < 10.0
assert OPENAI_KEY in [secret["value"] for secret in found]
assert list(tmp_path.iterdir()) == []
@pytest.mark.parametrize(
"content",
[
f"api_key = '{OPENAI_KEY}'\nbase = abcdefghijkl\npassword = x\n %(base)sZZZZQQQQ\n",
"base = abcdefghijkl\npassword = x\n %(base)sZZZZQQQQ\n",
f"api_key = '{OPENAI_KEY}'\nbase = Kp7Nq2Wz9Bt4\npassword = x\n"
" %(base)s-primary\nnote = Kp7Nq2Wz9Bt4-primary is the hostname\n",
'base = "abcdefghijkl"\npassword = "%(base)sZZZZQQQQ"\n',
],
ids=[
"vendor-key-present",
"no-vendor-key",
"value-echoed-elsewhere",
"quoted-interpolation",
],
)
def test_scan_message_never_reports_a_value_the_message_does_not_hold(content):
guardrail = _guardrail()
for secret in guardrail.scan_message_for_secrets(content):
assert secret["value"] in content
def test_scan_message_leaves_unrelated_text_alone_when_a_value_is_echoed():
guardrail = _guardrail()
content = (
f"api_key = '{OPENAI_KEY}'\nbase = Kp7Nq2Wz9Bt4\npassword = x\n"
" %(base)s-primary\nnote = Kp7Nq2Wz9Bt4-primary is the hostname\n"
)
assert "note = Kp7Nq2Wz9Bt4-primary is the hostname" in guardrail.redact_text(content)
def test_masked_entity_count_counts_each_secret_once():
guardrail = _guardrail()
_masked_entity_count.set({})
guardrail.redact_text(f"first {OPENAI_KEY} second {OPENAI_KEY}")
assert _masked_entity_count.get() == {"Strict OpenAI API Key": 1}
def test_scan_message_redacts_every_openai_key_occurrence():
guardrail = _guardrail()
content = f"first {OPENAI_KEY}, second {OPENAI_KEY}"
@ -81,9 +646,7 @@ def test_scan_message_requires_ascii_digits_for_openai_like_values():
def test_scan_message_redacts_openai_key_after_separator():
guardrail = _guardrail()
assert guardrail.redact_text(f"openai_{OPENAI_KEY} key-{OPENAI_KEY}") == (
"openai_[REDACTED] key-[REDACTED]"
)
assert guardrail.redact_text(f"openai_{OPENAI_KEY} key-{OPENAI_KEY}") == ("openai_[REDACTED] key-[REDACTED]")
assert guardrail.redact_text(URL_ENCODED_KEY) == "Bearer%20[REDACTED]"
@ -102,6 +665,31 @@ def test_scan_message_stays_linear_on_repeated_sk_separators():
assert time.perf_counter() - started < 2.0
@pytest.mark.parametrize(
"content",
[
f"api_key: '{OPENAI_KEY}'\npassword=" + "a-" * 10_000 + "!",
f"api_key: '{OPENAI_KEY}'\npassword:" + '"' * 20_000,
f"api_key: '{OPENAI_KEY}'\n" + "api_key:" * 10_000,
f"api_key: '{OPENAI_KEY}'\nsecret=" + "aB3dE6gH9jK2mN5p " * 2_000,
f"api_key: '{OPENAI_KEY}'\n" + "\n".join(f"password{i}=aB3dE6gH9jK2mN5p{i}" for i in range(3_000)),
],
ids=[
"value-run",
"quote-run",
"keyword-run",
"value-repeat",
"assignment-flood",
],
)
def test_scan_message_stays_linear_on_adversarial_credential_lines(content):
guardrail = _guardrail()
started = time.perf_counter()
guardrail.redact_text(content)
assert time.perf_counter() - started < 10.0
def test_scan_message_redacts_whole_stripe_live_key():
guardrail = _guardrail()
@ -119,8 +707,8 @@ def test_scan_message_replaces_longest_overlapping_match_first():
guardrail = _guardrail()
content = f'token = "{OPENAI_KEY}/extra"'
detected = guardrail.scan_message_for_secrets(content)
assert [secret["value"] for secret in detected] == [f"{OPENAI_KEY}/extra", OPENAI_KEY]
values = [secret["value"] for secret in guardrail.scan_message_for_secrets(content)]
assert values == [f"{OPENAI_KEY}/extra", OPENAI_KEY]
assert guardrail.redact_text(content) == 'token = "[REDACTED]"'

128
tests/test_litellm/integrations/test_custom_guardrail.py
View file

@ -2829,3 +2829,131 @@ class TestCustomGuardrailPostCallSuccessDeploymentHook:
assert response.choices[0].message.content == "filtered response"
assert "guardrail_to_apply" not in request_data
assert len(_guardrail_entries(request_data)) == 1
class TestPreCallHookResponseIsNotLoggedVerbatim:
"""Regression for LIT-6935: a pre_call hook returning the request payload leaked the prompt
into ``guardrail_response`` and from there onto OTEL guardrail spans."""
@staticmethod
def _logged_response(request_data: dict[str, object]) -> object:
metadata = request_data["litellm_metadata"]
assert isinstance(metadata, dict)
entries = metadata["standard_logging_guardrail_information"]
assert len(entries) == 1
return entries[0]["guardrail_response"]
@staticmethod
def _request() -> dict[str, object]:
return {
"model": "gpt-4.1-mini",
"input": "SECRET_PROMPT",
"messages": [{"role": "user", "content": "SECRET_PROMPT"}],
"litellm_metadata": {},
}
@pytest.mark.asyncio
async def test_pre_call_hook_returning_request_logs_allow(self):
class PassthroughGuardrail(CustomGuardrail):
@log_guardrail_information
async def async_pre_call_hook(
self,
user_api_key_dict: UserAPIKeyAuth,
cache: object,
data: dict[str, object],
call_type: str,
) -> dict[str, object]:
return data
data = self._request()
await PassthroughGuardrail(guardrail_name="g").async_pre_call_hook(
user_api_key_dict=UserAPIKeyAuth(), cache=None, data=data, call_type="aresponses"
)
assert self._logged_response(data) == "allow"
@pytest.mark.asyncio
async def test_pre_call_hook_returning_modified_copy_logs_mask(self):
class MaskingGuardrail(CustomGuardrail):
@log_guardrail_information
async def async_pre_call_hook(
self,
user_api_key_dict: UserAPIKeyAuth,
cache: object,
data: dict[str, object],
call_type: str,
) -> dict[str, object]:
return {**data, "input": "[MASKED]"}
data = self._request()
await MaskingGuardrail(guardrail_name="g").async_pre_call_hook(
user_api_key_dict=UserAPIKeyAuth(), cache=None, data=data, call_type="aresponses"
)
assert self._logged_response(data) == "mask"
@pytest.mark.asyncio
async def test_pre_call_hook_mutating_request_in_place_logs_mask(self):
class InPlaceMaskingGuardrail(CustomGuardrail):
@log_guardrail_information
async def async_pre_call_hook(
self,
user_api_key_dict: UserAPIKeyAuth,
cache: object,
data: dict[str, object],
call_type: str,
) -> dict[str, object]:
messages = data["messages"]
assert isinstance(messages, list)
messages[0]["content"] = "[MASKED]"
return data
data = self._request()
await InPlaceMaskingGuardrail(guardrail_name="g").async_pre_call_hook(
user_api_key_dict=UserAPIKeyAuth(), cache=None, data=data, call_type="acompletion"
)
assert self._logged_response(data) == "mask"
@pytest.mark.asyncio
async def test_pre_call_hook_returning_rejection_string_logs_that_string(self):
class RejectingGuardrail(CustomGuardrail):
@log_guardrail_information
async def async_pre_call_hook(
self,
user_api_key_dict: UserAPIKeyAuth,
cache: object,
data: dict[str, object],
call_type: str,
) -> str:
return "Blocked by policy"
data = self._request()
result = await RejectingGuardrail(guardrail_name="g").async_pre_call_hook(
user_api_key_dict=UserAPIKeyAuth(), cache=None, data=data, call_type="acompletion"
)
assert result == "Blocked by policy"
assert self._logged_response(data) == "Blocked by policy"
@pytest.mark.asyncio
async def test_pre_call_hook_removing_legacy_functions_in_place_logs_mask(self):
class FunctionStrippingGuardrail(CustomGuardrail):
@log_guardrail_information
async def async_pre_call_hook(
self,
user_api_key_dict: UserAPIKeyAuth,
cache: object,
data: dict[str, object],
call_type: str,
) -> dict[str, object]:
data["functions"] = []
data["function_call"] = "none"
return data
data = {**self._request(), "functions": [{"name": "delete_db"}], "function_call": "auto"}
await FunctionStrippingGuardrail(guardrail_name="g").async_pre_call_hook(
user_api_key_dict=UserAPIKeyAuth(), cache=None, data=data, call_type="acompletion"
)
assert self._logged_response(data) == "mask"

176
tests/test_litellm/proxy/spend_tracking/test_budget_reservation.py
View file

@ -1,5 +1,8 @@
from __future__ import annotations
import json
import math
from types import MappingProxyType
from typing import Final
import pytest
@ -207,8 +210,13 @@ def test_deployment_pricing_update_invalidates_cached_estimate() -> None:
ANTHROPIC_TOKENIZER_MODEL: Final = "claude-sonnet-4-5-20250929"
CL100K_MODEL: Final = "gpt-4"
O200K_MODEL: Final = "gpt-4o"
RUST_COUNTED_BODY: Final = {"model": ANTHROPIC_TOKENIZER_MODEL, "max_tokens": 16, "messages": ANTHROPIC_MESSAGES}
RUST_INPUT_TOKENS: Final = 4_321
RUST_INPUT_TOKENS_BY_TOKENIZER: Final = MappingProxyType(
{"anthropic": RUST_INPUT_TOKENS, "cl100k_base": 1_234, "o200k_base": 2_345}
)
class _FakeDeclined(Exception):
@ -225,33 +233,57 @@ class _FakeNative:
class _RecordingCounter:
bodies: Final[list[bytes]] = []
"""Stands in for one native counter; records `(tokenizer, body)` on the shared factory."""
def __init__(self, tokenizer_json: str) -> None:
pass
def __init__(self, factory: _RecordingFactory, tokenizer: rust_token_counter.RustTokenizer) -> None:
self.factory = factory
self.tokenizer = tokenizer
async def acount_request(self, body: bytes) -> object:
self.bodies.append(body)
return {"model": ANTHROPIC_TOKENIZER_MODEL, "input_tokens": RUST_INPUT_TOKENS}
self.factory.calls.append((self.tokenizer, body))
return {"model": "", "input_tokens": RUST_INPUT_TOKENS_BY_TOKENIZER[self.tokenizer]}
class _RecordingFactory:
"""Stands in for the native `TokenCounter` class: called with tokenizer JSON, or `from_*_ranks`."""
def __init__(self) -> None:
self.calls: list[tuple[rust_token_counter.RustTokenizer, bytes]] = []
def __call__(self, tokenizer_json: str) -> _RecordingCounter:
return _RecordingCounter(self, "anthropic")
def from_cl100k_ranks(self, rank_file: str) -> _RecordingCounter:
return _RecordingCounter(self, "cl100k_base")
def from_o200k_ranks(self, rank_file: str) -> _RecordingCounter:
return _RecordingCounter(self, "o200k_base")
class _DecliningCounter:
def __init__(self, tokenizer_json: str) -> None:
pass
async def acount_request(self, body: bytes) -> object:
raise _FakeDeclined("unsupported content block")
class _DecliningFactory:
def __call__(self, tokenizer_json: str) -> _DecliningCounter:
return _DecliningCounter()
def from_cl100k_ranks(self, rank_file: str) -> _DecliningCounter:
return _DecliningCounter()
def from_o200k_ranks(self, rank_file: str) -> _DecliningCounter:
return _DecliningCounter()
@pytest.fixture
def rust_counter(monkeypatch: pytest.MonkeyPatch):
monkeypatch.setattr(bindings, "get_native_bridge", lambda: _FakeNative())
rust_token_counter._anthropic_counter.cache_clear()
rust_token_counter._counter.cache_clear()
configuration.reset_rust_configuration()
_RecordingCounter.bodies.clear()
yield
rust_token_counter.TOKEN_COUNTER.reset()
rust_token_counter._anthropic_counter.cache_clear()
rust_token_counter._counter.cache_clear()
configuration.reset_rust_configuration()
@ -270,8 +302,9 @@ def rust_counter(monkeypatch: pytest.MonkeyPatch):
async def test_rust_count_replaces_python_tokenizing_on_every_llm_route(
rust_counter: None, route: str, request_body: dict
) -> None:
factory: Final = _RecordingFactory()
litellm.rust(True)
rust_token_counter.TOKEN_COUNTER.override(_RecordingCounter)
rust_token_counter.TOKEN_COUNTER.override(factory)
raw_body: Final = json.dumps(request_body).encode()
counts: Final = await count_request_input_tokens(
@ -279,53 +312,136 @@ async def test_rust_count_replaces_python_tokenizing_on_every_llm_route(
)
assert dict(counts) == {ANTHROPIC_TOKENIZER_MODEL: RUST_INPUT_TOKENS}
assert _RecordingCounter.bodies == [raw_body]
assert factory.calls == [("anthropic", raw_body)]
@pytest.mark.asyncio
async def test_rust_decline_falls_back_to_python_count(rust_counter: None) -> None:
@pytest.mark.parametrize("model", (CL100K_MODEL, "azure/gpt-35-turbo", "gemini/gemini-2.5-pro", "my-router-alias"))
async def test_tiktoken_cl100k_models_are_counted_by_rust(rust_counter: None, model: str) -> None:
factory: Final = _RecordingFactory()
litellm.rust(True)
rust_token_counter.TOKEN_COUNTER.override(_DecliningCounter)
rust_token_counter.TOKEN_COUNTER.override(factory)
body: Final = {"model": model, "messages": ANTHROPIC_MESSAGES}
raw_body: Final = json.dumps(body).encode()
counts: Final = await count_request_input_tokens(
request_body=body, route="/v1/chat/completions", llm_router=None, raw_body=raw_body
)
assert dict(counts) == {model: RUST_INPUT_TOKENS_BY_TOKENIZER["cl100k_base"]}
assert factory.calls == [("cl100k_base", raw_body)]
@pytest.mark.asyncio
@pytest.mark.parametrize("model", (O200K_MODEL, "gpt-5", "o3", "gpt-4.1", "chatgpt-4o-latest"))
async def test_tiktoken_o200k_models_are_counted_by_rust(rust_counter: None, model: str) -> None:
factory: Final = _RecordingFactory()
litellm.rust(True)
rust_token_counter.TOKEN_COUNTER.override(factory)
body: Final = {"model": model, "messages": ANTHROPIC_MESSAGES}
raw_body: Final = json.dumps(body).encode()
counts: Final = await count_request_input_tokens(
request_body=body, route="/v1/chat/completions", llm_router=None, raw_body=raw_body
)
assert dict(counts) == {model: RUST_INPUT_TOKENS_BY_TOKENIZER["o200k_base"]}
assert factory.calls == [("o200k_base", raw_body)]
@pytest.mark.asyncio
async def test_multi_model_request_counts_once_per_tokenizer_and_python_for_the_rest(rust_counter: None) -> None:
factory: Final = _RecordingFactory()
litellm.rust(True)
rust_token_counter.TOKEN_COUNTER.override(factory)
models: Final = (
CL100K_MODEL,
ANTHROPIC_TOKENIZER_MODEL,
"gemini/gemini-2.5-pro",
O200K_MODEL,
"gpt-5",
"replicate/meta/llama-2-70b-chat",
)
body: Final = {"model": list(models), "messages": ANTHROPIC_MESSAGES}
raw_body: Final = json.dumps(body).encode()
python_counts: Final = await count_request_input_tokens(
request_body=RUST_COUNTED_BODY, route="/v1/messages", llm_router=None
request_body=body, route="/v1/chat/completions", llm_router=None
)
counts: Final = await count_request_input_tokens(
request_body=RUST_COUNTED_BODY,
request_body=body, route="/v1/chat/completions", llm_router=None, raw_body=raw_body
)
assert factory.calls == [("cl100k_base", raw_body), ("anthropic", raw_body), ("o200k_base", raw_body)]
assert dict(counts) == {
CL100K_MODEL: RUST_INPUT_TOKENS_BY_TOKENIZER["cl100k_base"],
"gemini/gemini-2.5-pro": RUST_INPUT_TOKENS_BY_TOKENIZER["cl100k_base"],
ANTHROPIC_TOKENIZER_MODEL: RUST_INPUT_TOKENS,
O200K_MODEL: RUST_INPUT_TOKENS_BY_TOKENIZER["o200k_base"],
"gpt-5": RUST_INPUT_TOKENS_BY_TOKENIZER["o200k_base"],
"replicate/meta/llama-2-70b-chat": python_counts["replicate/meta/llama-2-70b-chat"],
}
assert counts["replicate/meta/llama-2-70b-chat"] not in RUST_INPUT_TOKENS_BY_TOKENIZER.values()
@pytest.mark.asyncio
@pytest.mark.parametrize("model", (ANTHROPIC_TOKENIZER_MODEL, CL100K_MODEL, O200K_MODEL))
async def test_rust_decline_falls_back_to_python_count(rust_counter: None, model: str) -> None:
litellm.rust(True)
rust_token_counter.TOKEN_COUNTER.override(_DecliningFactory())
body: Final = {**RUST_COUNTED_BODY, "model": model}
python_counts: Final = await count_request_input_tokens(request_body=body, route="/v1/messages", llm_router=None)
counts: Final = await count_request_input_tokens(
request_body=body,
route="/v1/messages",
llm_router=None,
raw_body=json.dumps(RUST_COUNTED_BODY).encode(),
raw_body=json.dumps(body).encode(),
)
assert dict(counts) == dict(python_counts)
assert counts[ANTHROPIC_TOKENIZER_MODEL] != RUST_INPUT_TOKENS
assert counts[model] not in RUST_INPUT_TOKENS_BY_TOKENIZER.values()
@pytest.mark.asyncio
async def test_disabled_rust_never_sees_the_raw_body(rust_counter: None) -> None:
factory: Final = _RecordingFactory()
litellm.rust(False)
rust_token_counter.TOKEN_COUNTER.override(_RecordingCounter)
rust_token_counter.TOKEN_COUNTER.override(factory)
body: Final = {"model": [ANTHROPIC_TOKENIZER_MODEL, CL100K_MODEL, O200K_MODEL], "messages": ANTHROPIC_MESSAGES}
counts: Final = await count_request_input_tokens(
request_body=RUST_COUNTED_BODY,
route="/v1/messages",
request_body=body,
route="/v1/chat/completions",
llm_router=None,
raw_body=json.dumps(RUST_COUNTED_BODY).encode(),
raw_body=json.dumps(body).encode(),
)
assert _RecordingCounter.bodies == []
assert counts[ANTHROPIC_TOKENIZER_MODEL] != RUST_INPUT_TOKENS
assert factory.calls == []
assert set(counts) == {ANTHROPIC_TOKENIZER_MODEL, CL100K_MODEL, O200K_MODEL}
assert not set(counts.values()) & set(RUST_INPUT_TOKENS_BY_TOKENIZER.values())
@pytest.mark.asyncio
async def test_non_anthropic_tokenizer_models_stay_in_python(rust_counter: None) -> None:
@pytest.mark.parametrize("model", ("replicate/meta/llama-2-70b-chat", "meta-llama/Llama-3-8b", "text-davinci-003"))
async def test_models_without_a_rust_tokenizer_stay_in_python(
rust_counter: None, monkeypatch: pytest.MonkeyPatch, model: str
) -> None:
monkeypatch.setattr(
litellm, "open_ai_chat_completion_models", litellm.open_ai_chat_completion_models | {"text-davinci-003"}
)
factory: Final = _RecordingFactory()
litellm.rust(True)
rust_token_counter.TOKEN_COUNTER.override(_RecordingCounter)
body: Final = {"model": "gpt-4o", "messages": ANTHROPIC_MESSAGES}
rust_token_counter.TOKEN_COUNTER.override(factory)
body: Final = {"model": model, "messages": ANTHROPIC_MESSAGES}
python_counts: Final = await count_request_input_tokens(
request_body=body, route="/v1/chat/completions", llm_router=None
)
counts: Final = await count_request_input_tokens(
request_body=body, route="/v1/chat/completions", llm_router=None, raw_body=json.dumps(body).encode()
)
assert _RecordingCounter.bodies == []
assert counts["gpt-4o"] != RUST_INPUT_TOKENS
assert factory.calls == []
assert dict(counts) == dict(python_counts)
assert counts[model] not in RUST_INPUT_TOKENS_BY_TOKENIZER.values()

312
tests/test_litellm/rust_bridge/test_token_counter.py
View file

@ -8,16 +8,26 @@ cases need the extension and are skipped when it is not built.
from __future__ import annotations
import json
from types import MappingProxyType
from typing import Final
import pytest
import tiktoken
from tokenizers import Tokenizer
import litellm
from litellm.constants import TIKTOKEN_ENCODE_CHUNK_SIZE_CHARS
from litellm.litellm_core_utils.token_counter import openai_tokenizer_encoding
from litellm.proxy.spend_tracking.budget_reservation import _count_input_tokens
from litellm.rust_bridge import bindings, configuration
from litellm.rust_bridge import token_counter as bridge
from litellm.utils import claude_json_str
MODEL: Final = "claude-sonnet-4-5-20250929"
CL100K_MODEL: Final = "gpt-4"
O200K_MODEL: Final = "gpt-4o"
TOKENIZERS: Final[tuple[bridge.RustTokenizer, ...]] = ("anthropic", "cl100k_base", "o200k_base")
RANK_FILE_LINES: Final = MappingProxyType({"cl100k_base": 100_256, "o200k_base": 199_998})
BODY: Final = json.dumps({"model": MODEL, "messages": [{"role": "user", "content": "hello"}]}).encode()
@ -44,69 +54,122 @@ class _RecordingCounter:
return {"model": MODEL, "input_tokens": 42}
class _DecliningCounter:
def __init__(self, tokenizer_json: str) -> None:
pass
class _RecordingFactory:
"""Stands in for the native `TokenCounter` class: callable for tokenizer JSON, `from_*_ranks` for rank files."""
def __init__(self) -> None:
self.counters: list[_RecordingCounter] = []
self.rank_files: list[str] = []
def __call__(self, tokenizer_json: str) -> _RecordingCounter:
counter = _RecordingCounter(tokenizer_json)
self.counters.append(counter)
return counter
def from_cl100k_ranks(self, rank_file: str) -> _RecordingCounter:
self.rank_files.append(rank_file)
return self("cl100k_base")
def from_o200k_ranks(self, rank_file: str) -> _RecordingCounter:
self.rank_files.append(rank_file)
return self("o200k_base")
class _RaisingCounter:
def __init__(self, error: Exception) -> None:
self.error = error
async def acount_request(self, body: bytes) -> object:
raise _FakeDeclined("request has no messages")
raise self.error
class _FailingCounter:
def __init__(self, tokenizer_json: str) -> None:
pass
class _RaisingFactory:
"""Every counter it builds, for either tokenizer, raises `error` on count."""
async def acount_request(self, body: bytes) -> object:
raise RuntimeError("encode failed")
def __init__(self, error: Exception) -> None:
self.error = error
def __call__(self, tokenizer_json: str) -> _RaisingCounter:
return _RaisingCounter(self.error)
def from_cl100k_ranks(self, rank_file: str) -> _RaisingCounter:
return _RaisingCounter(self.error)
def from_o200k_ranks(self, rank_file: str) -> _RaisingCounter:
return _RaisingCounter(self.error)
@pytest.fixture(autouse=True)
def _reset_bridge(monkeypatch: pytest.MonkeyPatch):
bridge.TOKEN_COUNTER.reset()
bridge._anthropic_counter.cache_clear()
bridge._counter.cache_clear()
configuration.reset_rust_configuration()
monkeypatch.setattr(bindings, "get_native_bridge", lambda: _FakeNative())
yield
bridge.TOKEN_COUNTER.reset()
bridge._anthropic_counter.cache_clear()
bridge._counter.cache_clear()
configuration.reset_rust_configuration()
@pytest.mark.asyncio
async def test_disabled_bridge_never_constructs_a_counter() -> None:
constructed: list[str] = []
def factory(tokenizer_json: str) -> _RecordingCounter:
constructed.append(tokenizer_json)
return _RecordingCounter(tokenizer_json)
@pytest.mark.parametrize("tokenizer", TOKENIZERS)
async def test_disabled_bridge_never_constructs_a_counter(tokenizer: bridge.RustTokenizer) -> None:
factory: Final = _RecordingFactory()
litellm.rust(False)
bridge.TOKEN_COUNTER.override(factory)
assert await bridge.count_anthropic_input_tokens(BODY) is None
assert constructed == []
assert await bridge.count_input_tokens(BODY, tokenizer) is None
assert factory.counters == []
@pytest.mark.asyncio
async def test_enabled_bridge_returns_typed_count_and_reuses_one_counter() -> None:
counters: list[_RecordingCounter] = []
def factory(tokenizer_json: str) -> _RecordingCounter:
counter = _RecordingCounter(tokenizer_json)
counters.append(counter)
return counter
factory: Final = _RecordingFactory()
litellm.rust(True)
bridge.TOKEN_COUNTER.override(factory)
first: Final = await bridge.count_anthropic_input_tokens(BODY)
second: Final = await bridge.count_anthropic_input_tokens(BODY)
first: Final = await bridge.count_input_tokens(BODY, "anthropic")
second: Final = await bridge.count_input_tokens(BODY, "anthropic")
assert first == bridge.InputTokenCount(model=MODEL, input_tokens=42)
assert second == first
assert len(counters) == 1
assert counters[0].bodies == [BODY, BODY]
assert json.loads(counters[0].tokenizer_json)["model"]["type"] == "BPE"
assert len(factory.counters) == 1
assert factory.counters[0].bodies == [BODY, BODY]
assert json.loads(factory.counters[0].tokenizer_json)["model"]["type"] == "BPE"
@pytest.mark.asyncio
@pytest.mark.parametrize("tokenizer", ("cl100k_base", "o200k_base"))
async def test_tiktoken_counter_is_built_from_the_vendored_rank_file_once(tokenizer: bridge.RustTokenizer) -> None:
factory: Final = _RecordingFactory()
litellm.rust(True)
bridge.TOKEN_COUNTER.override(factory)
first: Final = await bridge.count_input_tokens(BODY, tokenizer)
second: Final = await bridge.count_input_tokens(BODY, tokenizer)
assert first == second == bridge.InputTokenCount(model=MODEL, input_tokens=42)
assert len(factory.rank_files) == 1
assert factory.rank_files[0].startswith("IQ== 0\n")
assert factory.rank_files[0].count("\n") == RANK_FILE_LINES[tokenizer]
assert factory.counters[0].tokenizer_json == tokenizer
assert factory.counters[0].bodies == [BODY, BODY]
@pytest.mark.asyncio
async def test_each_tokenizer_gets_its_own_cached_counter() -> None:
factory: Final = _RecordingFactory()
litellm.rust(True)
bridge.TOKEN_COUNTER.override(factory)
await bridge.count_input_tokens(BODY, "anthropic")
await bridge.count_input_tokens(BODY, "cl100k_base")
await bridge.count_input_tokens(BODY, "o200k_base")
await bridge.count_input_tokens(BODY, "anthropic")
await bridge.count_input_tokens(BODY, "o200k_base")
assert [counter.tokenizer_json for counter in factory.counters][1:] == ["cl100k_base", "o200k_base"]
assert [len(counter.bodies) for counter in factory.counters] == [2, 1, 2]
@pytest.mark.asyncio
@ -114,38 +177,131 @@ async def test_missing_native_module_falls_back(monkeypatch: pytest.MonkeyPatch)
litellm.rust(True)
monkeypatch.setattr(bindings, "get_native_bridge", lambda: None)
assert await bridge.count_anthropic_input_tokens(BODY) is None
assert [await bridge.count_input_tokens(BODY, tokenizer) for tokenizer in TOKENIZERS] == [None, None, None]
@pytest.mark.asyncio
async def test_declined_request_falls_back() -> None:
@pytest.mark.parametrize("tokenizer", TOKENIZERS)
async def test_declined_request_falls_back(tokenizer: bridge.RustTokenizer) -> None:
litellm.rust(True)
bridge.TOKEN_COUNTER.override(_DecliningCounter)
bridge.TOKEN_COUNTER.override(_RaisingFactory(_FakeDeclined("request has no messages")))
assert await bridge.count_anthropic_input_tokens(BODY) is None
assert await bridge.count_input_tokens(BODY, tokenizer) is None
@pytest.mark.asyncio
async def test_runtime_failure_falls_back() -> None:
@pytest.mark.parametrize("tokenizer", TOKENIZERS)
async def test_runtime_failure_falls_back(tokenizer: bridge.RustTokenizer) -> None:
litellm.rust(True)
bridge.TOKEN_COUNTER.override(_FailingCounter)
bridge.TOKEN_COUNTER.override(_RaisingFactory(RuntimeError("encode failed")))
assert await bridge.count_anthropic_input_tokens(BODY) is None
assert await bridge.count_input_tokens(BODY, tokenizer) is None
@pytest.mark.parametrize(
("model", "expected"),
((MODEL, True), ("claude-3-5-sonnet-20241022", False), ("gpt-4o", False), ("my-router-alias", False)),
(
(MODEL, "anthropic"),
("claude-3-5-sonnet-20241022", "cl100k_base"),
("gpt-4", "cl100k_base"),
("gpt-4-turbo", "cl100k_base"),
("gpt-3.5-turbo", "cl100k_base"),
("azure/gpt-35-turbo", "cl100k_base"),
("gemini/gemini-2.5-pro", "cl100k_base"),
("mistral/mistral-large-latest", "cl100k_base"),
("my-router-alias", "cl100k_base"),
("azure/gpt-4o", "cl100k_base"),
("command-r-plus", "cl100k_base"),
("gpt-4o", "o200k_base"),
("gpt-4o-mini", "o200k_base"),
("gpt-4o-2024-08-06", "o200k_base"),
("chatgpt-4o-latest", "o200k_base"),
("gpt-4.1", "o200k_base"),
("gpt-5", "o200k_base"),
("gpt-5-mini", "o200k_base"),
("o1", "o200k_base"),
("o3", "o200k_base"),
("o3-mini", "o200k_base"),
("o4-mini", "o200k_base"),
("replicate/meta/llama-2-70b-chat", None),
("meta-llama/Llama-3-8b", None),
),
)
def test_uses_anthropic_tokenizer_mirrors_python_tokenizer_selection(model: str, expected: bool) -> None:
assert bridge.uses_anthropic_tokenizer(model) is expected
def test_rust_tokenizer_mirrors_python_tokenizer_selection(model: str, expected: bridge.RustTokenizer | None) -> None:
assert bridge.rust_tokenizer(model) == expected
@pytest.mark.parametrize("flag", ("disable_hf_tokenizer_download", "disable_token_counter"))
def test_uses_anthropic_tokenizer_respects_python_opt_outs(monkeypatch: pytest.MonkeyPatch, flag: str) -> None:
monkeypatch.setattr(litellm, flag, True)
@pytest.mark.parametrize(
("model", "python_encoding"),
(("text-davinci-003", "p50k_base"), ("gpt-oss-120b", "o200k_harmony")),
)
def test_rust_tokenizer_declines_tiktoken_encodings_rust_does_not_have(
monkeypatch: pytest.MonkeyPatch, model: str, python_encoding: str
) -> None:
monkeypatch.setattr(litellm, "open_ai_chat_completion_models", litellm.open_ai_chat_completion_models | {model})
assert bridge.uses_anthropic_tokenizer(MODEL) is False
assert openai_tokenizer_encoding(model).name == python_encoding
assert bridge.rust_tokenizer(model) is None
def test_rust_tokenizer_declines_the_cohere_tokenizer_download(monkeypatch: pytest.MonkeyPatch) -> None:
monkeypatch.setattr(litellm, "cohere_models", litellm.cohere_models | {"command-r-plus"})
assert bridge.rust_tokenizer("command-r-plus") is None
@pytest.mark.parametrize("legacy_model", ("gpt-3.5-turbo-0301", "gpt-35-turbo-0301"))
def test_rust_tokenizer_declines_legacy_message_accounting_python_prices_differently(
monkeypatch: pytest.MonkeyPatch, legacy_model: str
) -> None:
monkeypatch.setattr(
litellm, "open_ai_chat_completion_models", litellm.open_ai_chat_completion_models | {"gpt-3.5-turbo-0301"}
)
monkeypatch.setattr(litellm, "azure_llms", {**litellm.azure_llms, "gpt-35-turbo-0301": "azure"})
messages: Final = [{"role": "user", "name": "bob", "content": "hello there"}]
assert litellm.token_counter(model=legacy_model, messages=messages) != litellm.token_counter(
model=CL100K_MODEL, messages=messages
)
assert bridge.rust_tokenizer(legacy_model) is None
assert bridge.rust_tokenizer(CL100K_MODEL) == "cl100k_base"
@pytest.mark.parametrize("model", (MODEL, CL100K_MODEL, O200K_MODEL, "gpt-5", "o3"))
def test_rust_tokenizer_names_the_encoding_python_actually_counts_with(model: str) -> None:
text: Final = (
"Hello, world! camelCase ABCdef \u00e9\u00e8 12345 \u3053\u3093\u306b\u3061\u306f <|endoftext|>\r\n" * 9
)
python_count: Final = litellm.token_counter(model=model, text=text)
cl100k_count: Final = len(tiktoken.get_encoding("cl100k_base").encode(text, disallowed_special=()))
o200k_count: Final = len(tiktoken.get_encoding("o200k_base").encode(text, disallowed_special=()))
assert cl100k_count != o200k_count
match bridge.rust_tokenizer(model):
case "cl100k_base":
assert python_count == cl100k_count
case "o200k_base":
assert python_count == o200k_count
case "anthropic":
assert python_count == len(Tokenizer.from_str(claude_json_str).encode(text).ids)
assert python_count not in {cl100k_count, o200k_count}
case None:
pytest.fail(f"{model} must have a Rust tokenizer")
def test_disabled_hf_download_routes_anthropic_models_to_cl100k_like_python(monkeypatch: pytest.MonkeyPatch) -> None:
monkeypatch.setattr(litellm, "disable_hf_tokenizer_download", True)
assert bridge.rust_tokenizer(MODEL) == "cl100k_base"
assert bridge.rust_tokenizer("meta-llama/Llama-3-8b") == "cl100k_base"
assert bridge.rust_tokenizer(O200K_MODEL) == "o200k_base"
def test_disabled_token_counter_declines_every_model(monkeypatch: pytest.MonkeyPatch) -> None:
monkeypatch.setattr(litellm, "disable_token_counter", True)
assert bridge.rust_tokenizer(MODEL) is None
assert bridge.rust_tokenizer(CL100K_MODEL) is None
assert bridge.rust_tokenizer(O200K_MODEL) is None
PARITY_REQUESTS: Final[tuple[dict[str, object], ...]] = (
@ -182,7 +338,16 @@ PARITY_REQUESTS: Final[tuple[dict[str, object], ...]] = (
},
{
"model": MODEL,
"messages": [{"role": "user", "content": "x " * 20_000}],
"messages": [{"role": "user", "content": "x " * 500}],
},
{
"model": MODEL,
"messages": [
{
"role": "user",
"content": "I'VE got 1234567 things; it's \"fine\"...\r\n\r\n caf\u00e9 \u0645\u0631\u062d\u0628\u0627 \U0001f600 <|endoftext|>",
}
],
},
{"model": MODEL, "prompt": "Write a haiku about ships.", "max_tokens": 20},
{"model": MODEL, "prompt": ["first prompt", "second prompt"]},
@ -190,7 +355,7 @@ PARITY_REQUESTS: Final[tuple[dict[str, object], ...]] = (
"model": MODEL,
"instructions": "be terse",
"input": [
{"role": "user", "content": [{"type": "input_text", "text": "Summarise caf\u00e9 menus \u2014 \"ok\"?\n"}]},
{"role": "user", "content": [{"type": "input_text", "text": 'Summarise caf\u00e9 menus \u2014 "ok"?\n'}]},
{"role": "assistant", "content": "Sure."},
],
},
@ -202,27 +367,63 @@ PARITY_REQUESTS: Final[tuple[dict[str, object], ...]] = (
)
PARITY_MODELS: Final[tuple[tuple[str, bridge.RustTokenizer], ...]] = (
(MODEL, "anthropic"),
(CL100K_MODEL, "cl100k_base"),
(O200K_MODEL, "o200k_base"),
("gpt-5", "o200k_base"),
)
@pytest.mark.asyncio
@pytest.mark.parametrize(("model", "tokenizer"), PARITY_MODELS)
@pytest.mark.parametrize("request_body", PARITY_REQUESTS)
async def test_native_count_matches_python_budget_counter(
monkeypatch: pytest.MonkeyPatch, request_body: dict[str, object]
monkeypatch: pytest.MonkeyPatch, request_body: dict[str, object], model: str, tokenizer: bridge.RustTokenizer
) -> None:
native: Final = pytest.importorskip("litellm.rust_bridge._native")
monkeypatch.setattr(bindings, "get_native_bridge", lambda: native)
litellm.rust(True)
body: Final = json.dumps(request_body).replace(MODEL, model)
rust_count: Final = await bridge.count_anthropic_input_tokens(json.dumps(request_body).encode())
python_count: Final = _count_input_tokens(request_body=request_body, model=MODEL)
rust_count: Final = await bridge.count_input_tokens(body.encode(), tokenizer)
python_count: Final = _count_input_tokens(request_body=json.loads(body), model=model)
assert rust_count is not None
assert rust_count.model == request_body.get("model")
assert rust_count.model == json.loads(body).get("model")
assert rust_count.input_tokens == python_count
@pytest.mark.asyncio
@pytest.mark.parametrize(("model", "tokenizer"), ((CL100K_MODEL, "cl100k_base"), (O200K_MODEL, "o200k_base")))
async def test_tiktoken_counts_long_text_exactly_where_python_chunks(
monkeypatch: pytest.MonkeyPatch, model: str, tokenizer: bridge.RustTokenizer
) -> None:
"""Python encodes tiktoken text in fixed-size chunks (drift of up to one token per chunk boundary); Rust does not."""
native: Final = pytest.importorskip("litellm.rust_bridge._native")
monkeypatch.setattr(bindings, "get_native_bridge", lambda: native)
litellm.rust(True)
text: Final = "x " * 20_000
body: Final = {"model": model, "messages": [{"role": "user", "content": text}]}
encoding: Final = tiktoken.get_encoding(tokenizer)
exact: Final = 3 + len(encoding.encode("user")) + len(encoding.encode(text)) + 3
chunks: Final = -(-len(text) // TIKTOKEN_ENCODE_CHUNK_SIZE_CHARS)
rust_count: Final = await bridge.count_input_tokens(json.dumps(body).encode(), tokenizer)
python_count: Final = _count_input_tokens(request_body=body, model=model)
assert rust_count is not None
assert rust_count.input_tokens == exact
assert python_count is not None
assert exact < python_count <= exact + chunks
DECLINED_REQUESTS: Final[tuple[dict[str, object], ...]] = (
{
"model": MODEL,
"messages": [{"role": "user", "content": [{"type": "image_url", "image_url": {"url": "data:image/png;base64,AA"}}]}],
"messages": [
{"role": "user", "content": [{"type": "image_url", "image_url": {"url": "data:image/png;base64,AA"}}]}
],
},
{"model": MODEL, "prompt": 1.5},
{"model": MODEL, "documents": [{"score": 0.5}]},
@ -231,12 +432,13 @@ DECLINED_REQUESTS: Final[tuple[dict[str, object], ...]] = (
@pytest.mark.asyncio
@pytest.mark.parametrize("tokenizer", TOKENIZERS)
@pytest.mark.parametrize("request_body", DECLINED_REQUESTS)
async def test_native_declines_shapes_python_prices_differently(
monkeypatch: pytest.MonkeyPatch, request_body: dict[str, object]
monkeypatch: pytest.MonkeyPatch, request_body: dict[str, object], tokenizer: bridge.RustTokenizer
) -> None:
native: Final = pytest.importorskip("litellm.rust_bridge._native")
monkeypatch.setattr(bindings, "get_native_bridge", lambda: native)
litellm.rust(True)
assert await bridge.count_anthropic_input_tokens(json.dumps(request_body).encode()) is None
assert await bridge.count_input_tokens(json.dumps(request_body).encode(), tokenizer) is None

4
uv.lock generated
View file

@ -10,7 +10,7 @@ resolution-markers = [
]
[options]
exclude-newer = "2026-09-07T23:09:03.362777Z"
exclude-newer = "2026-09-08T03:56:24.358378Z"
exclude-newer-span = "P3D"
[manifest]
@ -4559,6 +4559,7 @@ e2e-dev = [
{ name = "locust" },
{ name = "mcp" },
{ name = "playwright" },
{ name = "psutil" },
{ name = "websockets" },
]
healthcheck = [
@ -4747,6 +4748,7 @@ e2e-dev = [
{ name = "locust", specifier = "==2.45.0" },
{ name = "mcp", specifier = ">=1.28.1,<2.0" },
{ name = "playwright", specifier = "==1.61.0" },
{ name = "psutil", specifier = "==7.2.2" },
{ name = "websockets", specifier = ">=15.0.1,<16.0" },
]
healthcheck = [