From 7dcdc2ffc930fac68f8a7bf0c02780afe390c035 Mon Sep 17 00:00:00 2001 From: yujonglee Date: Thu, 29 Oct 2026 07:55:02 -0700 Subject: [PATCH 01/13] build(model-catalog): ship the cost map inside the native extension (#48907) Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .github/scripts/smoke_test_native_wheel.py | 9 +++++- backend/Dockerfile | 1 + gateway/Dockerfile | 1 + litellm-rust/Cargo.lock | 1 + litellm-rust/crates/python-bridge/AGENTS.md | 12 ++++++++ litellm-rust/crates/python-bridge/Cargo.toml | 1 + litellm-rust/crates/python-bridge/src/lib.rs | 4 +++ .../python-bridge/src/model_registry.rs | 30 +++++++++++++++++++ litellm/rust_bridge/_native.pyi | 2 ++ migrations/Dockerfile | 1 + packaging/litellm-core/pyproject.toml | 1 + pyproject.toml | 2 ++ scripts/build_core_distribution.py | 10 ++++++- .../base_sdk_tests/test_core_distribution.py | 4 ++- 14 files changed, 76 insertions(+), 3 deletions(-) create mode 100644 litellm-rust/crates/python-bridge/src/model_registry.rs diff --git a/.github/scripts/smoke_test_native_wheel.py b/.github/scripts/smoke_test_native_wheel.py index 577bb32fcf0..af4054ea92e 100644 --- a/.github/scripts/smoke_test_native_wheel.py +++ b/.github/scripts/smoke_test_native_wheel.py @@ -7,6 +7,8 @@ import zipfile from pathlib import Path from typing import Final +COST_MAP: Final = Path(__file__).resolve().parents[2] / "model_prices_and_context_window.json" + CHILD_SCRIPT: Final = """ from importlib.util import module_from_spec, spec_from_file_location from pathlib import Path @@ -34,6 +36,9 @@ else: after = module.gil_stats() if not isinstance(after.get("releases"), int): raise AssertionError(f"native module unusable after panic: {after!r}") + +if module.bundled_model_cost_map() != Path(sys.argv[2]).read_bytes(): + raise AssertionError("the wheel embeds a different model_prices_and_context_window.json than this checkout") """ @@ -55,7 +60,9 @@ def main() -> int: native_path: Final = Path(temporary_directory) / Path(native_members[0].filename).name native_path.write_bytes(archive.read(native_members[0])) - result: Final = subprocess.run((sys.executable, "-c", CHILD_SCRIPT, str(native_path)), check=False) + result: Final = subprocess.run( + (sys.executable, "-c", CHILD_SCRIPT, str(native_path), str(COST_MAP)), check=False + ) if result.returncode != 0: sys.stderr.write(f"native wheel smoke test exited with status {result.returncode}\n") diff --git a/backend/Dockerfile b/backend/Dockerfile index 1c22f914100..aad59ba17c6 100644 --- a/backend/Dockerfile +++ b/backend/Dockerfile @@ -24,6 +24,7 @@ ENV UV_PYTHON_DOWNLOADS=0 COPY pyproject.toml README.md rust-toolchain.toml ./ COPY .cargo ./.cargo COPY litellm-rust ./litellm-rust +COPY model_prices_and_context_window.json ./ RUN --mount=type=cache,target=/root/.cache/uv \ mkdir -p litellm/rust_bridge /wheels && \ diff --git a/gateway/Dockerfile b/gateway/Dockerfile index 6dc8bab45ce..9d228627543 100644 --- a/gateway/Dockerfile +++ b/gateway/Dockerfile @@ -40,6 +40,7 @@ ENV UV_PYTHON_DOWNLOADS=0 COPY pyproject.toml README.md rust-toolchain.toml ./ COPY .cargo ./.cargo COPY litellm-rust ./litellm-rust +COPY model_prices_and_context_window.json ./ RUN --mount=type=cache,target=/root/.cache/uv \ mkdir -p litellm/rust_bridge /wheels && \ diff --git a/litellm-rust/Cargo.lock b/litellm-rust/Cargo.lock index f7d7e5f4dab..e959f77917b 100644 --- a/litellm-rust/Cargo.lock +++ b/litellm-rust/Cargo.lock @@ -4345,6 +4345,7 @@ dependencies = [ "litellm-inference-transcription", "litellm-llms", "litellm-llms-types", + "litellm-model-catalog", "litellm-router-types", "litellm-secrets", "litellm-secrets-aws", diff --git a/litellm-rust/crates/python-bridge/AGENTS.md b/litellm-rust/crates/python-bridge/AGENTS.md index d90fe073b20..26858b2afe2 100644 --- a/litellm-rust/crates/python-bridge/AGENTS.md +++ b/litellm-rust/crates/python-bridge/AGENTS.md @@ -106,6 +106,18 @@ GIL handling to `litellm-host-python`. ## Data Handling +- Reuse large immutable bundled data and parsed catalogs across calls, including in + replacement APIs. Share Rust storage and cache repeated Python projections instead + of reparsing, serializing or copying the full payload on every access +- Cache owned Python objects with `PyOnceLock>`, then bind and clone references + while attached. Do not use blocking `OnceLock::get_or_init` while attached + ([PyO3 initialization](https://pyo3.rs/v0.29.2/faq.html), + [reference ownership](https://pyo3.rs/v0.29.2/types.html)) +- Never reuse Python objects across interpreters or interpreter reinitialization. + Once-only caches are for immutable data. Reloadable/custom data needs versioned + entries or invalidation, and caller-mutable Python containers must not be shared +- Test unchanged contents and allocation reuse, not just equal values. Keep this + coverage when replacing a bundled-data accessor with a catalog API - OCR payloads can contain personal data and large base64 images. Do not log payloads or provider responses. - Avoid copying large payloads more than needed. The current JSON round-trip is diff --git a/litellm-rust/crates/python-bridge/Cargo.toml b/litellm-rust/crates/python-bridge/Cargo.toml index 9d69cd8c14b..a4e1c57ab1e 100644 --- a/litellm-rust/crates/python-bridge/Cargo.toml +++ b/litellm-rust/crates/python-bridge/Cargo.toml @@ -49,6 +49,7 @@ litellm-secrets-types.workspace = true litellm-llms-types.workspace = true litellm-router-types.workspace = true litellm-host-python.workspace = true +litellm-model-catalog.workspace = true litellm-token-counter = { path = "../token-counter", default-features = false } pyo3.workspace = true prost.workspace = true diff --git a/litellm-rust/crates/python-bridge/src/lib.rs b/litellm-rust/crates/python-bridge/src/lib.rs index 12827b76ff3..683e90d02a2 100644 --- a/litellm-rust/crates/python-bridge/src/lib.rs +++ b/litellm-rust/crates/python-bridge/src/lib.rs @@ -8,6 +8,7 @@ mod execution; mod http; mod logger; mod marshal; +mod model_registry; mod preflight; mod python_settings; mod routes; @@ -30,6 +31,8 @@ mod _native { #[pymodule_export] use crate::logger::NativeDiagnosticProcessor; #[pymodule_export] + use crate::model_registry::bundled_model_cost_map; + #[pymodule_export] use crate::routes::audio_transcription::{atranscription, transcription}; #[pymodule_export] use crate::routes::chat_completions::{acompletion, completion}; @@ -110,6 +113,7 @@ mod tests { "NativeClickHouseSpendConfig", "NativeClickHouseSpendStorage", "trace_encode_error", + "bundled_model_cost_map", "TokenCounter", "Tokenizer", "gil_stats", diff --git a/litellm-rust/crates/python-bridge/src/model_registry.rs b/litellm-rust/crates/python-bridge/src/model_registry.rs new file mode 100644 index 00000000000..b7b7cd1fbb3 --- /dev/null +++ b/litellm-rust/crates/python-bridge/src/model_registry.rs @@ -0,0 +1,30 @@ +use pyo3::prelude::*; +use pyo3::sync::PyOnceLock; +use pyo3::types::PyBytes; + +#[pyfunction] +pub fn bundled_model_cost_map(py: Python<'_>) -> Bound<'_, PyBytes> { + static BYTES: PyOnceLock> = PyOnceLock::new(); + BYTES + .get_or_init(py, || { + PyBytes::new(py, litellm_model_catalog::BUNDLED_REGISTRY).unbind() + }) + .bind(py) + .clone() +} + +#[cfg(test)] +mod tests { + use super::*; + + #[rstest::rstest] + fn repeated_calls_reuse_the_bundled_registry_bytes() { + Python::initialize(); + let first = Python::attach(|py| bundled_model_cost_map(py).unbind()); + Python::attach(|py| { + let second = bundled_model_cost_map(py); + assert!(first.bind(py).is(&second)); + assert_eq!(second.as_bytes(), litellm_model_catalog::BUNDLED_REGISTRY); + }); + } +} diff --git a/litellm/rust_bridge/_native.pyi b/litellm/rust_bridge/_native.pyi index 7364a1d72f2..62b48621c05 100644 --- a/litellm/rust_bridge/_native.pyi +++ b/litellm/rust_bridge/_native.pyi @@ -14,6 +14,7 @@ class ForkedAfterNativeRuntimeStarted(RuntimeError): ... class ProcessReservedForForking(RuntimeError): ... def trace_encode_error(message: str) -> bytes: ... +def bundled_model_cost_map() -> bytes: ... @final class NativeClickHouseSpendConfig: @@ -220,6 +221,7 @@ __all__ = [ "aocr", "aresponses", "atranscription", + "bundled_model_cost_map", "completion", "embedding", "gil_stats", diff --git a/migrations/Dockerfile b/migrations/Dockerfile index 2b1fd029143..3ffa9140543 100644 --- a/migrations/Dockerfile +++ b/migrations/Dockerfile @@ -43,6 +43,7 @@ ENV UV_PYTHON_DOWNLOADS=0 COPY pyproject.toml README.md rust-toolchain.toml ./ COPY .cargo ./.cargo COPY litellm-rust ./litellm-rust +COPY model_prices_and_context_window.json ./ RUN --mount=type=cache,target=/root/.cache/uv \ mkdir -p litellm/rust_bridge /wheels && \ diff --git a/packaging/litellm-core/pyproject.toml b/packaging/litellm-core/pyproject.toml index 27681d9116c..db50589a675 100644 --- a/packaging/litellm-core/pyproject.toml +++ b/packaging/litellm-core/pyproject.toml @@ -56,6 +56,7 @@ editable-profile = "dev" include = [ { path = "rust-toolchain.toml", format = "sdist" }, { path = ".cargo/config.toml", format = "sdist" }, + { path = "model_prices_and_context_window.json", format = "sdist" }, "litellm/router_strategy/complexity_router/artifacts/*.json", "litellm/router_strategy/complexity_router/fuse_presets.json", "litellm/proxy/model_insights_tasks.json", diff --git a/pyproject.toml b/pyproject.toml index 200bb463313..e6d34a3d90c 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -335,6 +335,7 @@ include = [ "litellm/proxy/public_endpoints/whats_new.json", "litellm/proxy/common_utils/codex_base_instructions.md", "litellm/proxy/common_utils/codex_bundled_models_0.159.3.json", + { path = "model_prices_and_context_window.json", format = "sdist" }, ] exclude = [ "litellm/proxy/enterprise", @@ -357,6 +358,7 @@ cache-keys = [ { file = "litellm-rust/Cargo.lock" }, { file = "litellm-rust/Cargo.toml" }, { file = "litellm-rust/crates/**/*" }, + { file = "model_prices_and_context_window.json" }, ] constraint-dependencies = [ "tornado>=6.5.8", diff --git a/scripts/build_core_distribution.py b/scripts/build_core_distribution.py index 36ec20a9031..41d7208d2c3 100644 --- a/scripts/build_core_distribution.py +++ b/scripts/build_core_distribution.py @@ -9,7 +9,15 @@ from pathlib import Path from typing import Final ROOT: Final = Path(__file__).resolve().parents[1] -SOURCES: Final = ("litellm", "litellm-rust", ".cargo", "rust-toolchain.toml", "README.md", "LICENSE") +SOURCES: Final = ( + "litellm", + "litellm-rust", + ".cargo", + "rust-toolchain.toml", + "model_prices_and_context_window.json", + "README.md", + "LICENSE", +) def stage_core_distribution(source: Path, destination: Path) -> None: diff --git a/tests/base_sdk_tests/test_core_distribution.py b/tests/base_sdk_tests/test_core_distribution.py index ba68d08e4d1..c22fa4e54fc 100644 --- a/tests/base_sdk_tests/test_core_distribution.py +++ b/tests/base_sdk_tests/test_core_distribution.py @@ -132,7 +132,9 @@ def test_core_artifacts_exclude_local_configuration(distributions: tuple[Path, P assert not any(member.name.endswith(f"/{name}") for member in archive.getmembers()) -@pytest.mark.parametrize("relative_path", ["rust-toolchain.toml", ".cargo/config.toml"]) +@pytest.mark.parametrize( + "relative_path", ["rust-toolchain.toml", ".cargo/config.toml", "model_prices_and_context_window.json"] +) def test_core_sdist_preserves_native_build_configuration(distribution_directory: Path, relative_path: str) -> None: sdists: Final = tuple(distribution_directory.glob("litellm_core-*.tar.gz")) assert len(sdists) == 1 From 2a61f12f78d09c0cf31f7d57156d51760879d00e Mon Sep 17 00:00:00 2001 From: yujonglee Date: Thu, 29 Oct 2026 07:55:02 -0700 Subject: [PATCH 02/13] ci(model-catalog): check the cost map by loading it through model-catalog (#48908) Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .github/workflows/cost-map-guard.yml | 21 ++++ litellm-rust/crates/model-catalog/AGENTS.md | 2 + litellm-rust/crates/model-catalog/Cargo.toml | 2 +- .../model-catalog/src/bin/check-registry.rs | 26 ++++ .../crates/model-catalog/src/catalog.rs | 114 +++++++++++++++++- .../crates/model-catalog/src/error.rs | 10 +- .../crates/model-catalog/tests/catalog.rs | 18 ++- 7 files changed, 179 insertions(+), 14 deletions(-) create mode 100644 litellm-rust/crates/model-catalog/src/bin/check-registry.rs diff --git a/.github/workflows/cost-map-guard.yml b/.github/workflows/cost-map-guard.yml index 6519b65e71b..fb86237e504 100644 --- a/.github/workflows/cost-map-guard.yml +++ b/.github/workflows/cost-map-guard.yml @@ -46,3 +46,24 @@ jobs: HEAD_REF: ${{ github.event.pull_request.head.ref }} run: | uv run --no-project python ci_cd/cost_map_guard.py --base "$MERGE_BASE" --head "$HEAD_SHA" --head-ref "$HEAD_REF" + - name: Check whether the pull request changes the cost map + id: changed + env: + MERGE_BASE: ${{ steps.revisions.outputs.merge_base }} + HEAD_SHA: ${{ github.event.pull_request.head.sha }} + run: | + if git diff --quiet "$MERGE_BASE" "$HEAD_SHA" -- model_prices_and_context_window.json; then + echo "cost_map=false" >> "$GITHUB_OUTPUT" + else + echo "cost_map=true" >> "$GITHUB_OUTPUT" + fi + - if: steps.changed.outputs.cost_map == 'true' + run: rustup toolchain install --no-self-update + - name: Load the cost map through model-catalog + if: steps.changed.outputs.cost_map == 'true' + env: + HEAD_SHA: ${{ github.event.pull_request.head.sha }} + run: | + git show "${HEAD_SHA}:model_prices_and_context_window.json" > "$RUNNER_TEMP/model_prices_and_context_window.json" + cargo run --locked --quiet --manifest-path litellm-rust/Cargo.toml -p litellm-model-catalog --bin check-registry \ + -- "$RUNNER_TEMP/model_prices_and_context_window.json" diff --git a/litellm-rust/crates/model-catalog/AGENTS.md b/litellm-rust/crates/model-catalog/AGENTS.md index fc4b28d4ef0..ff24fb81143 100644 --- a/litellm-rust/crates/model-catalog/AGENTS.md +++ b/litellm-rust/crates/model-catalog/AGENTS.md @@ -25,3 +25,5 @@ A new `N`, image size or resolution in an existing key shape parses with no Rust ## Tests `BUNDLED_REGISTRY` embeds the root `model_prices_and_context_window.json`, so the tests over it fail a JSON change with a bad value or a key no type reads + +`Catalog::problems` is the one definition of a mergeable cost map. The bundled test asserts it is empty, and the cost map guard workflow runs it on a pull request's JSON through `cargo run --bin check-registry ` diff --git a/litellm-rust/crates/model-catalog/Cargo.toml b/litellm-rust/crates/model-catalog/Cargo.toml index 45d16515dcd..63b9ad14768 100644 --- a/litellm-rust/crates/model-catalog/Cargo.toml +++ b/litellm-rust/crates/model-catalog/Cargo.toml @@ -10,6 +10,7 @@ schema = ["dep:schemars", "litellm-cost/schema"] [dependencies] litellm-cost = { workspace = true, features = ["serde"] } +litellm-error.workspace = true indexmap = { version = "2.14.0", features = ["serde"] } jiff = { version = "0.2.37", default-features = false, features = ["std", "tzdb-bundle-always"] } @@ -24,5 +25,4 @@ time.workspace = true [dev-dependencies] jsonschema = { version = "0.55.1", default-features = false } -litellm-error.workspace = true rstest.workspace = true diff --git a/litellm-rust/crates/model-catalog/src/bin/check-registry.rs b/litellm-rust/crates/model-catalog/src/bin/check-registry.rs new file mode 100644 index 00000000000..6fdc7621a0a --- /dev/null +++ b/litellm-rust/crates/model-catalog/src/bin/check-registry.rs @@ -0,0 +1,26 @@ +use std::{env, fs, path::PathBuf, process::ExitCode}; + +use litellm_model_catalog::Catalog; + +fn main() -> ExitCode { + let Some(path) = env::args_os().nth(1).map(PathBuf::from) else { + eprintln!("usage: check-registry "); + return ExitCode::from(2); + }; + let body = match fs::read(&path) { + Ok(body) => body, + Err(error) => { + eprintln!("{}: {error}", path.display()); + return ExitCode::FAILURE; + } + }; + let problems = Catalog::problems(&body); + for problem in &problems { + eprintln!("{}", litellm_error::Report::new(problem)); + } + if problems.is_empty() { + ExitCode::SUCCESS + } else { + ExitCode::FAILURE + } +} diff --git a/litellm-rust/crates/model-catalog/src/catalog.rs b/litellm-rust/crates/model-catalog/src/catalog.rs index ac5f4ffcf25..98af785b17e 100644 --- a/litellm-rust/crates/model-catalog/src/catalog.rs +++ b/litellm-rust/crates/model-catalog/src/catalog.rs @@ -5,7 +5,7 @@ use indexmap::IndexMap; use serde::Deserialize; use serde_json::Value; -use crate::error::{Error, Loaded, Rejects, Repr}; +use crate::error::{Error, Loaded, Rejects, Repr, pointer_child, pointer_for_key}; use crate::fallback::{FallbackGeneralizations, FallbackRule}; use crate::model::Model; use crate::store::{Rows, Store}; @@ -211,6 +211,36 @@ impl Catalog { }) } + pub fn problems(body: &[u8]) -> Vec { + let Loaded { + value: catalog, + errors, + } = match Self::parse(body) { + Ok(loaded) => loaded, + Err(fatal) => return vec![fatal], + }; + let unread = catalog.rows().entries.iter().flat_map(|(model, entry)| { + EntryWire::read(entry.fields()) + .unknown + .into_iter() + .map(move |(key, _)| { + Error::new( + pointer_child(&pointer_for_key(model), &key), + Repr::UnreadField, + ) + }) + }); + let no_fallback_rules = catalog + .fallback_rules() + .is_none_or(<[_]>::is_empty) + .then(|| Error::from_key(FALLBACK_GENERALIZATIONS, Repr::NoFallbackRules)); + errors + .into_iter() + .chain(unread) + .chain(no_fallback_rules) + .collect() + } + pub fn row_errors(model: &str, fields: &Fields) -> Vec { let (_, _, errors) = read_row(fields); errors @@ -404,3 +434,85 @@ fn unknown_timezone(wire: &EntryWire) -> Option { ) }) } + +#[cfg(test)] +mod tests { + use rstest::rstest; + + use super::*; + + const RULES: &str = r#""fallback_generalizations":{"rules":[{"name":"family","pattern":"^new-","model_info":{"mode":"chat"}}]}"#; + + #[rstest] + #[case::clean(format!(r#"{{{RULES},"m":{{"litellm_provider":"test"}}}}"#), vec![])] + #[case::not_json("{".to_owned(), vec![("", "json")])] + #[case::row_error(format!(r#"{{{RULES},"m":{{}}}}"#), vec![("/m/litellm_provider", "missing_provider")])] + #[case::unread_key( + format!(r#"{{{RULES},"m":{{"litellm_provider":"test","typo_cost":1}}}}"#), + vec![("/m/typo_cost", "unread_field")] + )] + #[case::empty_model_key( + format!(r#"{{{RULES},"":{{"litellm_provider":"test","future_field":true}}}}"#), + vec![("//future_field", "unread_field")] + )] + #[case::unread_key_on_invalid_row( + format!(r#"{{{RULES},"m":{{"typo_cost":1}}}}"#), + vec![("/m/litellm_provider", "missing_provider"), ("/m/typo_cost", "unread_field")] + )] + #[case::escaped_unread_key_on_invalid_row( + format!(r#"{{{RULES},"open/router~v1":{{"typo/cost~":1}}}}"#), + vec![("/open~1router~0v1/litellm_provider", "missing_provider"), ("/open~1router~0v1/typo~1cost~0", "unread_field")] + )] + #[case::no_rules( + r#"{"m":{"litellm_provider":"test"}}"#.to_owned(), + vec![("/fallback_generalizations", "no_fallback_rules")] + )] + #[case::empty_rules( + r#"{"fallback_generalizations":{"rules":[]},"m":{"litellm_provider":"test"}}"#.to_owned(), + vec![("/fallback_generalizations", "no_fallback_rules")] + )] + fn problems_list_what_a_cost_map_change_must_fix( + #[case] body: String, + #[case] expected: Vec<(&str, &str)>, + ) { + let problems: Vec<_> = Catalog::problems(body.as_bytes()) + .iter() + .map(|problem| (problem.path().to_owned(), <&str>::from(problem.kind()))) + .collect(); + let expected: Vec<_> = expected + .into_iter() + .map(|(path, kind)| (path.to_owned(), kind)) + .collect(); + assert_eq!(problems, expected); + } + + #[rstest] + fn runtime_parsing_preserves_custom_fields_and_recovers_valid_rows() { + let body = br#"{ + "good":{"litellm_provider":"test","custom_metadata":true}, + "bad":{"litellm_provider":"test","input_cost_per_token":"wrong"} + }"#; + let loaded = Catalog::parse(body).unwrap(); + assert_eq!(loaded.errors.len(), 1); + assert_eq!(loaded.errors[0].path(), "/bad/input_cost_per_token"); + let (catalog, rejected) = loaded.value.without_rejected(&loaded.errors); + assert_eq!(rejected.len(), 1); + assert!(catalog.lookup("bad").is_none()); + let good = catalog.lookup("good").unwrap(); + assert!(good.entry.model().is_some()); + assert_eq!( + good.entry.field("custom_metadata"), + Some(&Value::Bool(true)) + ); + assert!( + Catalog::problems(body) + .iter() + .any(|problem| problem.kind() == crate::ErrorKind::UnreadField) + ); + assert!( + Catalog::problems(body) + .iter() + .any(|problem| problem.kind() == crate::ErrorKind::NoFallbackRules) + ); + } +} diff --git a/litellm-rust/crates/model-catalog/src/error.rs b/litellm-rust/crates/model-catalog/src/error.rs index 425ae47687f..89e27c05f46 100644 --- a/litellm-rust/crates/model-catalog/src/error.rs +++ b/litellm-rust/crates/model-catalog/src/error.rs @@ -109,6 +109,12 @@ pub(crate) enum Repr { #[source] cause: Option, }, + #[error("not a field litellm reads")] + #[strum_discriminants(strum(serialize = "unread_field"))] + UnreadField, + #[error("has no rules")] + #[strum_discriminants(strum(serialize = "no_fallback_rules"))] + NoFallbackRules, } /// What a problem takes out of the served map @@ -148,7 +154,7 @@ impl Error { | Repr::NoOffPeakHours | Repr::EmptyList | Repr::InvalidDate { .. } => Rejects::Row, - Repr::MissingProvider => Rejects::Nothing, + Repr::MissingProvider | Repr::UnreadField | Repr::NoFallbackRules => Rejects::Nothing, } } } @@ -334,6 +340,8 @@ mod tests { #[case::no_off_peak_hours(Repr::NoOffPeakHours, Rejects::Row)] #[case::empty_list(Repr::EmptyList, Rejects::Row)] #[case::invalid_date(Repr::InvalidDate { cause: None }, Rejects::Row)] + #[case::unread_field(Repr::UnreadField, Rejects::Nothing)] + #[case::no_fallback_rules(Repr::NoFallbackRules, Rejects::Nothing)] fn policy_classifies_each_error(#[case] error: Repr, #[case] expected: Rejects) { assert_eq!(Error::from(error).rejects(), expected); } diff --git a/litellm-rust/crates/model-catalog/tests/catalog.rs b/litellm-rust/crates/model-catalog/tests/catalog.rs index 30c070cfb52..31f907689d3 100644 --- a/litellm-rust/crates/model-catalog/tests/catalog.rs +++ b/litellm-rust/crates/model-catalog/tests/catalog.rs @@ -366,18 +366,14 @@ fn written_is_the_document_in_order_including_sections_and_bad_entries() { } #[rstest] -fn the_bundled_registry_parses_with_no_problems_and_every_key_understood() { - let loaded = Catalog::parse(BUNDLED_REGISTRY).unwrap(); - let catalog = loaded.value; - assert!(loaded.errors.is_empty(), "{:#?}", loaded.errors); - let unrecognized: Vec<_> = catalog.unrecognized().collect(); - assert!(unrecognized.is_empty(), "{unrecognized:#?}"); - assert!( - catalog - .fallback_rules() - .is_some_and(|rules| !rules.is_empty()) - ); +fn the_bundled_registry_has_no_problems() { + let problems: Vec<_> = Catalog::problems(BUNDLED_REGISTRY) + .iter() + .map(report) + .collect(); + assert!(problems.is_empty(), "{problems:#?}"); const REFERENCE_MODEL_COUNT: usize = 4303; + let catalog = Catalog::parse(BUNDLED_REGISTRY).unwrap().value; assert!(catalog.model_count() as f64 >= REFERENCE_MODEL_COUNT as f64 * 0.9); } From cd58014315ef584e505a05cb2ab4d4eb07004a6d Mon Sep 17 00:00:00 2001 From: Dan Lemon Date: Thu, 29 Oct 2026 15:58:38 +0100 Subject: [PATCH 03/13] feat(soniox): add realtime speech-to-text on /v1/realtime (#44822) * feat(soniox): add realtime speech-to-text on /v1/realtime Maps OpenAI realtime transcription sessions onto the Soniox websocket API (stt-rt-v5): server VAD uses Soniox endpoint detection, commit sends a manual finalize, deployment litellm_params pass Soniox options through, translation sessions transcribe the translated tokens, idle streams get keepalives, and cost is tracked from processed audio duration. * fix(soniox): bill realtime usage from integer milliseconds * fix(soniox): bound session-start queue and cover realtime backend edge paths * fix(soniox): bill realtime stream time, start sessions on connect and surface client protocol errors Soniox charges for the whole open stream, so session close now bills the connection's open time. The backend announces the session on connect, which flushes audio sent before any session.update, and a single upstream pump makes recv safe to cancel. Invalid client frames now get an OpenAI error event instead of a silent disconnect * fix(soniox): skip deferred setup so client frames reach Soniox in order Drops the connect marker frame, which counted as an upstream frame and hid refusals, and could start a buffered stream at the default rate before the client's session.update. Soniox now opts out of deferred setup, the handler announces a synthetic session.created whenever no setup went upstream, a session.update that would change a started stream gets an error event, and stream time is billed only once a stream started * fix(soniox): judge a follow-up session.update only on the settings it sends A partial follow-up session.update was rebuilt with defaults for the settings it left out, so restating part of a 16 kHz or VAD-off stream was rejected as a change * fix(soniox): treat an omitted rate in a follow-up session.update as unchanged * chore(soniox): drop comments that restate the code * fix(soniox): return speaker labels, log refused sessions as failures, and probe realtime health Deltas carry speaker and language when diarization or language identification is on, and the completed event lists segments per speaker and language. A session Soniox refuses before any audio is processed is logged as a failure with no spend. Billing starts at the start request. Deferred-setup providers (Meta, Vertex Chirp) send the error event and still end the session. Soniox realtime models get a websocket /health probe and cost rows for the stt-rt-v5 aliases. * fix(soniox): bill failed streams, probe health with deployment options, and keep token accumulation linear Only request rejections (400, 401, 402, 403, 429) before any processed audio count as a refused session, so a 408 or 5xx still bills the open stream. The /health probe starts Soniox with the deployment's options, sends 100 ms of silence, and passes only on Soniox's finished frame. A turn's final tokens append in constant time instead of copying the tuple per token. * fix(soniox): log every Soniox error before any processed audio as a failed session * fix(soniox): keep untranslated speech, support input_audio_buffer.clear, ignore an empty language * fix(soniox): reject frames without a type, keep strict language hints, queue finalizes in a deque * fix(soniox): answer a commit with no speech with an empty completed turn * fix(soniox): let a server VAD turn answer a commit that is still in flight * fix(soniox): let a server VAD turn answer only the oldest pending commit --------- Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> --- .../health_check_helpers.py | 8 +- .../litellm_core_utils/realtime_streaming.py | 26 +- .../llms/base_llm/realtime/transformation.py | 9 +- litellm/llms/custom_httpx/llm_http_handler.py | 6 +- litellm/llms/soniox/common_utils.py | 8 +- litellm/llms/soniox/realtime/__init__.py | 0 .../llms/soniox/realtime/transformation.py | 609 +++++++++++++ ...odel_prices_and_context_window_backup.json | 120 +++ .../provider_endpoints_support_backup.json | 2 +- .../provider_create_fields.json | 2 +- litellm/realtime_api/README.md | 2 +- litellm/realtime_api/main.py | 14 +- litellm/utils.py | 5 + model_prices_and_context_window.json | 120 +++ provider_endpoints_support.json | 2 +- .../test_health_check_helpers.py | 5 +- .../test_realtime_streaming.py | 125 ++- .../custom_httpx/test_llm_http_handler.py | 154 ++++ tests/unit/llms/soniox/realtime/__init__.py | 1 + .../test_soniox_realtime_transformation.py | 807 ++++++++++++++++++ .../test_get_complete_url_trailing_slash.py | 1 + tests/unit/realtime_api/test_main.py | 101 ++- 22 files changed, 2093 insertions(+), 34 deletions(-) create mode 100644 litellm/llms/soniox/realtime/__init__.py create mode 100644 litellm/llms/soniox/realtime/transformation.py create mode 100644 tests/unit/llms/soniox/realtime/__init__.py create mode 100644 tests/unit/llms/soniox/realtime/test_soniox_realtime_transformation.py diff --git a/litellm/litellm_core_utils/health_check_helpers.py b/litellm/litellm_core_utils/health_check_helpers.py index b19bcd941fa..940730eeb9f 100644 --- a/litellm/litellm_core_utils/health_check_helpers.py +++ b/litellm/litellm_core_utils/health_check_helpers.py @@ -31,7 +31,13 @@ EXTRACTION_HEALTH_CHECK_FIELDS: Final[Mapping[str, Mapping[str, str]]] = Mapping ) -def native_health_check_mode(model: str, custom_llm_provider: str | None) -> Literal["anthropic_messages"] | None: +def native_health_check_mode( + model: str, custom_llm_provider: str | None +) -> Literal["anthropic_messages", "realtime"] | None: + if custom_llm_provider == "soniox": + from litellm.llms.soniox.common_utils import soniox_health_check_mode + + return soniox_health_check_mode(model) if custom_llm_provider != "bedrock_mantle": return None from litellm.llms.bedrock_mantle.common_utils import mantle_health_check_mode diff --git a/litellm/litellm_core_utils/realtime_streaming.py b/litellm/litellm_core_utils/realtime_streaming.py index 927caec3350..484bacdc72c 100644 --- a/litellm/litellm_core_utils/realtime_streaming.py +++ b/litellm/litellm_core_utils/realtime_streaming.py @@ -20,6 +20,7 @@ from litellm.constants import ( REALTIME_SESSION_SUCCESS_LOGGED_KEY, ) from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER +from litellm.llms.base_llm.realtime.transcription_protocol import RealtimeTranscriptionProtocolError from litellm.llms.base_llm.realtime.transformation import BaseRealtimeConfig, RealtimeBackend from litellm.types.llms.openai import ( OpenAIRealtimeEvents, @@ -1281,24 +1282,29 @@ class RealTimeStreaming: verbose_logger.exception("Connection closed in backend to client send messages - %s", e) close: Final = backend_close_from(e) self._flush_unbilled_transcription_usage() - if self._backend_refused_session(close): - await self.log_backend_refusal(e) - else: - await self.log_messages() + await self._log_session_end(e if self._backend_refused_session(close) else None) return close except asyncio.CancelledError: self._flush_unbilled_transcription_usage() - await self.log_messages() + await self._log_session_end(None) raise except Exception as e: verbose_logger.exception("Error in backend to client send messages: %s", e) self._flush_unbilled_transcription_usage() - await self.log_messages() + await self._log_session_end(None) return BackendClose(code=1011, reason="proxy failed while relaying the upstream websocket") def _backend_refused_session(self, close: BackendClose) -> bool: return close.code != 1000 and not self._backend_sent_frames + async def _log_session_end(self, transport_refusal: Exception | None) -> None: + provider_refusal: Final = None if self.provider_config is None else self.provider_config.session_refusal() + refusal: Final = provider_refusal or transport_refusal + if refusal is None: + await self.log_messages() + return + await self.log_backend_refusal(refusal) + async def log_backend_refusal(self, error: Exception) -> None: if not self.logging_obj: return @@ -1691,7 +1697,13 @@ class RealTimeStreaming: # ``_maybe_send_guardrail_turn_detection_update`` would retry. if msg_type == "session.update" and self.provider_config is None: await self._send_session_gate_once() - sent = await self._send_to_backend(message) + try: + sent = await self._send_to_backend(message) + except RealtimeTranscriptionProtocolError as e: + await self.websocket.send_text(realtime_error_event(str(e), error_type="invalid_request_error")) + if self._uses_deferred_backend_setup(): + raise + continue if guardrail_turn_detection_injected and sent: self._guardrail_turn_detection_update_sent = True diff --git a/litellm/llms/base_llm/realtime/transformation.py b/litellm/llms/base_llm/realtime/transformation.py index 1f4ad29fa74..f1539d16b30 100644 --- a/litellm/llms/base_llm/realtime/transformation.py +++ b/litellm/llms/base_llm/realtime/transformation.py @@ -4,7 +4,6 @@ from types import TracebackType from typing import TYPE_CHECKING, Any, Protocol import httpx -from typing_extensions import Self from litellm.types.llms.openai import OpenAIRealtimeStreamSessionEvents from litellm.types.realtime import ( @@ -24,7 +23,7 @@ else: class RealtimeBackend(Protocol): - async def __aenter__(self) -> Self: ... + async def __aenter__(self) -> object: ... async def __aexit__( self, @@ -97,9 +96,15 @@ class BaseRealtimeConfig(ABC): def unbilled_usage_on_session_close(self, model: str) -> RealtimeInputAudioTranscriptionUsage | None: return None + def session_refusal(self) -> Exception | None: + return None + async def open_backend(self, url: str, headers: Mapping[str, str]) -> RealtimeBackend | None: return None + def wrap_backend(self, backend: RealtimeBackend) -> RealtimeBackend: + return backend + def transform_session_created_event( self, model: str, diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index 4c0f84a003d..8a9ff62cfae 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -6923,7 +6923,7 @@ class BaseLLMHTTPHandler: try: ssl_context: Final = realtime_ssl_for_url(url) provider_backend: Final = await provider_config.open_backend(url, headers) - backend_ws: Final = ( + backend_ws: Final = provider_config.wrap_backend( provider_backend if provider_backend is not None else await self._open_realtime_backend_ws(websockets, url, headers, ssl_context) @@ -6961,9 +6961,7 @@ class BaseLLMHTTPHandler: await backend_ws.send(_session_config) realtime_streaming.session_configuration_request = _session_config - # For providers that defer setup until client session.update, optionally - # send synthetic session.created to unblock clients waiting on connect. - if not provider_config.requires_session_configuration(): + if not _session_config: synthetic_session: Final = provider_config.transform_session_created_event( model=model, logging_session_id=logging_obj.litellm_trace_id, diff --git a/litellm/llms/soniox/common_utils.py b/litellm/llms/soniox/common_utils.py index 8d83b4c9218..796f6ca12e4 100644 --- a/litellm/llms/soniox/common_utils.py +++ b/litellm/llms/soniox/common_utils.py @@ -3,7 +3,7 @@ Shared utilities for the Soniox provider (https://soniox.com). """ from collections.abc import Mapping, Sequence -from typing import Final, TypeAlias +from typing import Final, Literal, TypeAlias from litellm.litellm_core_utils.audio_utils.subtitle_utils import ( SubtitleToken, @@ -17,6 +17,8 @@ SonioxToken: TypeAlias = Mapping[str, object] # Soniox API base URL. SONIOX_API_BASE: Final[str] = "https://api.soniox.com" +SONIOX_REALTIME_MODEL_PREFIX: Final = "stt-rt-" + # Default polling interval in seconds when waiting for an async transcription # to finish. Mirrors the Soniox SDK default. SONIOX_DEFAULT_POLL_INTERVAL: Final[float] = 1.0 @@ -54,6 +56,10 @@ class SonioxException(BaseLLMException): """Provider-specific exception class for Soniox.""" +def soniox_health_check_mode(model: str) -> Literal["realtime"] | None: + return "realtime" if model.startswith(SONIOX_REALTIME_MODEL_PREFIX) else None + + def get_soniox_api_key(api_key: str | None = None) -> str | None: """Resolve the Soniox API key from arg or env var.""" # Local import to avoid a circular import: litellm.secret_managers.main diff --git a/litellm/llms/soniox/realtime/__init__.py b/litellm/llms/soniox/realtime/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/litellm/llms/soniox/realtime/transformation.py b/litellm/llms/soniox/realtime/transformation.py new file mode 100644 index 00000000000..ed7590b6f39 --- /dev/null +++ b/litellm/llms/soniox/realtime/transformation.py @@ -0,0 +1,609 @@ +import asyncio +import time +from collections import deque +from collections.abc import Awaitable, Callable, Mapping +from itertools import chain, groupby +from types import MappingProxyType, TracebackType +from typing import Final +from urllib.parse import urlparse, urlunparse + +from pydantic import BaseModel, ConfigDict, JsonValue, ValidationError +from typing_extensions import NotRequired, ReadOnly, Required, Self, TypedDict + +from litellm import verbose_logger +from litellm._uuid import uuid +from litellm.constants import REALTIME_WEBSOCKET_MAX_MESSAGE_SIZE_BYTES +from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj +from litellm.llms.base_llm.realtime.transcription_protocol import ( + SESSION_UPDATE_EVENT_TYPES, + RealtimeTranscriptionProtocolError, + TranscriptionSessionUpdate, + completed_event, + decode_pcm16_append, + delta_event, + duration_usage, + error_event, + json_object, + parse_transcription_session_update, + speech_event, + transcription_session, + transcription_session_created_event, + uses_server_vad, +) +from litellm.llms.base_llm.realtime.transformation import BaseRealtimeConfig, RealtimeBackend +from litellm.llms.custom_httpx.http_handler import realtime_ssl_for_url +from litellm.llms.soniox.common_utils import SonioxException, get_soniox_api_key +from litellm.types.llms.openai import ( + OpenAIRealtimeErrorEvent, + OpenAIRealtimeEvents, + OpenAIRealtimeInputAudioTranscriptionCompleted, + OpenAIRealtimeInputAudioTranscriptionDelta, + OpenAIRealtimeTranscriptionSessionCreated, +) +from litellm.types.realtime import ( + RealtimeInputAudioTranscriptionUsage, + RealtimeResponseTransformInput, + RealtimeResponseTypedDict, +) + +DEFAULT_SONIOX_REALTIME_URL: Final = "wss://stt-rt.soniox.com/transcribe-websocket" +WEBSOCKET_SCHEMES: Final = MappingProxyType({"wss": "wss", "ws": "ws", "https": "wss", "http": "ws"}) +DEFAULT_SAMPLE_RATE: Final = 24_000 +KEEPALIVE_INTERVAL_SECONDS: Final = 5.0 +_MAX_PENDING_FRAMES: Final = 64 +_FINALIZE: Final = '{"type":"finalize"}' +_KEEPALIVE: Final = '{"type":"keepalive"}' +_END_STREAM: Final = "" +_HEALTH_CHECK_SILENCE: Final = bytes(DEFAULT_SAMPLE_RATE // 5) +_CONTROL_FRAMES: Final = frozenset((_FINALIZE, _KEEPALIVE, _END_STREAM)) +SESSION_STARTED_FRAME: Final = '{"type":"litellm.soniox.session_started"}' +_ENDPOINT_TOKEN: Final = "" +_FINALIZED_TOKEN: Final = "" + + +class SonioxProtocolError(RealtimeTranscriptionProtocolError): + pass + + +class SonioxRealtimeOptions(BaseModel): + model_config = ConfigDict(frozen=True, extra="ignore") + + language_hints: tuple[str, ...] | None = None + language_hints_strict: bool | None = None + enable_language_identification: bool | None = None + enable_speaker_diarization: bool | None = None + context: JsonValue | None = None + translation: JsonValue | None = None + client_reference_id: str | None = None + max_endpoint_delay_ms: int | None = None + endpoint_sensitivity: float | None = None + endpoint_latency_adjustment_level: int | None = None + + +class SonioxStartRequest(SonioxRealtimeOptions): + model: str + audio_format: str + sample_rate: int + num_channels: int + enable_endpoint_detection: bool + + +class SonioxSpeechLabels(TypedDict, total=False): + speaker: ReadOnly[str] + language: ReadOnly[str] + + +class SonioxTranscriptSegment(SonioxSpeechLabels): + text: Required[ReadOnly[str]] + + +class SonioxTranscriptionDelta(OpenAIRealtimeInputAudioTranscriptionDelta, SonioxSpeechLabels): + pass + + +class SonioxTranscriptionCompleted(OpenAIRealtimeInputAudioTranscriptionCompleted): + segments: NotRequired[ReadOnly[tuple[SonioxTranscriptSegment, ...]]] + + +class _SonioxToken(BaseModel): + model_config = ConfigDict(frozen=True) + + text: str + is_final: bool = False + translation_status: str | None = None + speaker: str | None = None + language: str | None = None + + @property + def labelled(self) -> bool: + return self.speaker is not None or self.language is not None + + +class _SonioxResponse(BaseModel): + model_config = ConfigDict(frozen=True) + + tokens: tuple[_SonioxToken, ...] = () + total_audio_proc_ms: int | None = None + error_code: int | None = None + error_message: str | None = None + finished: bool = False + + +def _sample_rate(update: TranscriptionSessionUpdate | None) -> int: + audio_format: Final = None if update is None else update.audio_format + if audio_format is None: + return DEFAULT_SAMPLE_RATE + if not audio_format.is_pcm16: + raise SonioxProtocolError("Soniox realtime requires pcm16 input audio") + if audio_format.channels not in (None, 1): + raise SonioxProtocolError("Soniox realtime requires mono input audio") + if audio_format.rate is None: + return DEFAULT_SAMPLE_RATE + if audio_format.rate <= 0: + raise SonioxProtocolError("sample rate must be positive") + return audio_format.rate + + +def _language_hints( + language: str | None, configured: tuple[str, ...] | None, strict: bool | None +) -> tuple[str, ...] | None: + if not language or (strict and configured): + return configured + return (language, *(hint for hint in configured or () if hint != language)) + + +def _uses_endpoint_detection(update: TranscriptionSessionUpdate, drop_params: bool) -> bool: + return uses_server_vad(update, "Soniox realtime", drop_params, SonioxProtocolError) + + +def _ignored_settings(update: TranscriptionSessionUpdate) -> tuple[str, ...]: + vad_tuning: Final = (f"turn_detection.{key}" for key in sorted(update.turn_detection or {}) if key != "type") + return (*update.unsupported_transcription_keys, *vad_tuning) + + +def build_start_request( + model: str, update: TranscriptionSessionUpdate | None, options: SonioxRealtimeOptions, drop_params: bool = False +) -> SonioxStartRequest: + return SonioxStartRequest( + **options.model_dump(exclude={"language_hints"}), + language_hints=_language_hints( + None if update is None else update.language, options.language_hints, options.language_hints_strict + ), + model=model, + audio_format="pcm_s16le", + sample_rate=_sample_rate(update), + num_channels=1, + enable_endpoint_detection=update is None or _uses_endpoint_detection(update, drop_params), + ) + + +def _session_created(session_id: str, start_request: SonioxStartRequest) -> OpenAIRealtimeTranscriptionSessionCreated: + session: Final = transcription_session( + session_id=session_id, + model=start_request.model, + sample_rate=start_request.sample_rate, + language=start_request.language_hints[0] if start_request.language_hints else None, + server_vad=start_request.enable_endpoint_detection, + ) + return transcription_session_created_event(session) + + +def _changes_started_rate(update: TranscriptionSessionUpdate, started: SonioxStartRequest) -> bool: + audio_format: Final = update.audio_format + if audio_format is None: + return False + requested_rate: Final = _sample_rate(update) + return audio_format.rate is not None and requested_rate != started.sample_rate + + +def _changes_started_stream(update: TranscriptionSessionUpdate, started: SonioxStartRequest, drop_params: bool) -> bool: + turn_detection_set: Final = update.turn_detection is not None or update.turn_detection_disabled + return ( + _changes_started_rate(update, started) + or ( + update.language is not None + and _language_hints(update.language, started.language_hints, started.language_hints_strict) + != started.language_hints + ) + or (turn_detection_set and _uses_endpoint_detection(update, drop_params) != started.enable_endpoint_detection) + ) + + +def build_soniox_realtime_url(api_base: str | None) -> str: + if api_base is None: + return DEFAULT_SONIOX_REALTIME_URL + parsed: Final = urlparse(api_base) + scheme: Final = WEBSOCKET_SCHEMES.get(parsed.scheme) + if scheme is None: + raise ValueError("Soniox realtime api_base must be a wss://, ws://, https://, or http:// URL") + return urlunparse(parsed._replace(scheme=scheme)) + + +def _labels(token: _SonioxToken) -> SonioxSpeechLabels: + speaker: Final[SonioxSpeechLabels] = {} if token.speaker is None else {"speaker": token.speaker} + language: Final[SonioxSpeechLabels] = {} if token.language is None else {"language": token.language} + return {**speaker, **language} + + +def _delta(item_id: str, token: _SonioxToken) -> SonioxTranscriptionDelta: + return {**delta_event(item_id, token.text), **_labels(token)} + + +def _segment(run: tuple[_SonioxToken, ...]) -> SonioxTranscriptSegment: + return {**_labels(run[0]), "text": "".join(token.text for token in run).strip()} + + +def _segments(tokens: tuple[_SonioxToken, ...]) -> tuple[SonioxTranscriptSegment, ...]: + return tuple(_segment(tuple(run)) for _, run in groupby(tokens, key=lambda token: (token.speaker, token.language))) + + +def _new_item_id() -> str: + return f"item_{uuid.uuid4().hex}" + + +def _completed( + item_id: str, tokens: tuple[_SonioxToken, ...], usage: RealtimeInputAudioTranscriptionUsage | None +) -> OpenAIRealtimeInputAudioTranscriptionCompleted: + completed: Final = completed_event(item_id, "".join(token.text for token in tokens).strip(), usage) + if not any(token.labelled for token in tokens): + return completed + labelled: Final[SonioxTranscriptionCompleted] = {**completed, "segments": _segments(tokens)} + return labelled + + +def _soniox_error(response: _SonioxResponse) -> SonioxException | None: + if response.error_code is None and response.error_message is None: + return None + return SonioxException( + status_code=response.error_code or 500, + message=f"Soniox realtime error {response.error_code}: {response.error_message}", + ) + + +class SonioxEventTransformer: + def __init__(self, *, drops_originals: bool) -> None: + self._drops_originals: Final = drops_originals + self._item_id: str | None = None + self._final_tokens: Final[list[_SonioxToken]] = [] # mutable-ok: a turn's tokens append in O(1) each + self._finalizes: Final[deque[bool]] = deque() # mutable-ok: pending finalizes queue and pop in O(1) each + self._oldest_commit_answered: bool = False + self._processed_ms: int = 0 + self._billed_ms: int = 0 + self._refusal: SonioxException | None = None + + @property + def refusal(self) -> SonioxException | None: + return self._refusal + + def transform(self, payload: str) -> tuple[OpenAIRealtimeEvents, ...]: + try: + response: Final = _SonioxResponse.model_validate_json(payload) + except ValidationError: + raise SonioxProtocolError("Soniox returned an invalid realtime response") from None + self._bill(response.total_audio_proc_ms) + error: Final = _soniox_error(response) + if error is not None: + return (self._report(error),) + return tuple(chain.from_iterable(self._token_events(token) for token in response.tokens)) + + def queue_finalize(self, *, discards_turn: bool) -> None: + if discards_turn and not self._finalizes: + self._item_id = None + self._final_tokens.clear() + self._finalizes.append(discards_turn) + + def take_unbilled_usage(self, stream_ms: int = 0) -> RealtimeInputAudioTranscriptionUsage | None: + billable_ms: Final = max(self._processed_ms, stream_ms) + unbilled_ms: Final = billable_ms - self._billed_ms + if unbilled_ms <= 0: + return None + self._billed_ms = billable_ms + return duration_usage(unbilled_ms / 1000) + + def _report(self, error: SonioxException) -> OpenAIRealtimeErrorEvent: + if self._processed_ms == 0: + self._refusal = error + return error_event(error.message) + + def _bill(self, total_audio_proc_ms: int | None) -> None: + if total_audio_proc_ms is not None and total_audio_proc_ms > self._processed_ms: + self._processed_ms = total_audio_proc_ms + + def _token_events(self, token: _SonioxToken) -> tuple[OpenAIRealtimeEvents, ...]: + if token.text == _FINALIZED_TOKEN: + return self._finalized() + if self._discarding: + return () + if token.text == _ENDPOINT_TOKEN: + return self._end_turn() + started: Final = self._start() + if not token.is_final or (self._drops_originals and token.translation_status == "original"): + return started + self._final_tokens.append(token) + return (*started, _delta(self._current_item_id(), token)) + + @property + def _discarding(self) -> bool: + return bool(self._finalizes) and self._finalizes[0] + + def _finalized(self) -> tuple[OpenAIRealtimeEvents, ...]: + discarded: Final = self._finalizes.popleft() if self._finalizes else None + answered: Final = self._oldest_commit_answered + self._oldest_commit_answered = False + if discarded: + return () + if self._item_id is not None: + return self._complete(speech_stopped=False) + if discarded is None or answered: + return () + return (_completed(_new_item_id(), (), self.take_unbilled_usage()),) + + def _end_turn(self) -> tuple[OpenAIRealtimeEvents, ...]: + events: Final = self._complete(speech_stopped=True) + if events and self._finalizes: + self._oldest_commit_answered = True + return events + + def _start(self) -> tuple[OpenAIRealtimeEvents, ...]: + if self._item_id is not None: + return () + item_id: Final = _new_item_id() + self._item_id = item_id + return (speech_event("input_audio_buffer.speech_started", item_id),) + + def _current_item_id(self) -> str: + assert self._item_id is not None + return self._item_id + + def _complete(self, *, speech_stopped: bool) -> tuple[OpenAIRealtimeEvents, ...]: + item_id: Final = self._item_id + if item_id is None: + return () + tokens: Final = tuple(self._final_tokens) + self._item_id = None + self._final_tokens.clear() + stopped: Final[tuple[OpenAIRealtimeEvents, ...]] = ( + (speech_event("input_audio_buffer.speech_stopped", item_id),) if speech_stopped else () + ) + return (*stopped, _completed(item_id, tokens, self.take_unbilled_usage())) + + +class SonioxRealtimeBackend: + def __init__( + self, + inner: RealtimeBackend, + *, + keepalive_interval: float = KEEPALIVE_INTERVAL_SECONDS, + clock: Callable[[], float] = time.monotonic, + sleep: Callable[[float], Awaitable[None]] = asyncio.sleep, + ) -> None: + self._inner: Final = inner + self._interval: Final = keepalive_interval + self._clock: Final = clock + self._sleep: Final = sleep + self._frames: Final[asyncio.Queue[str | bytes | Exception]] = asyncio.Queue(maxsize=_MAX_PENDING_FRAMES) + self._last_sent: float = clock() + self._started: bool = False + self._ended: bool = False + self._tasks: tuple[asyncio.Task[None], ...] = () + + async def __aenter__(self) -> Self: + await self._inner.__aenter__() + self._tasks = (asyncio.create_task(self._pump()), asyncio.create_task(self._keepalive())) + return self + + async def __aexit__( + self, + exc_type: type[BaseException] | None, + exc_value: BaseException | None, + traceback: TracebackType | None, + ) -> None: + self._stop() + await self._inner.__aexit__(exc_type, exc_value, traceback) + + async def send(self, message: str | bytes) -> None: + self._last_sent = self._clock() + self._ended = self._ended or message == _END_STREAM + await self._inner.send(message) + if not self._started and isinstance(message, str) and message not in _CONTROL_FRAMES: + self._started = True + await self._frames.put(SESSION_STARTED_FRAME) + + async def recv(self, decode: bool | None = None) -> str | bytes: + frame: Final = await self._frames.get() + if isinstance(frame, Exception): + raise frame + return frame + + async def close(self) -> None: + self._stop() + await self._inner.close() + + def _stop(self) -> None: + for task in self._tasks: + task.cancel() + + async def _pump(self) -> None: + while True: + try: + frame = await self._inner.recv() + except Exception as e: # noqa: BLE001 # recv() re-raises it, so an upstream failure never hangs the reader + await self._frames.put(e) + return + await self._frames.put(frame) + + async def _keepalive(self) -> None: + while not self._ended: + await self._sleep(self._interval) + if self._ended or self._clock() - self._last_sent < self._interval: + continue + try: + await self.send(_KEEPALIVE) + except Exception as e: # noqa: BLE001 # stops only keepalives; the pump reports the broken socket + verbose_logger.debug("Soniox realtime: stopping keepalive after send failed: %s", e) + return + + +class SonioxRealtimeConfig(BaseRealtimeConfig): + def __init__( + self, + *, + options: SonioxRealtimeOptions | None = None, + wrap: Callable[[RealtimeBackend], RealtimeBackend] = SonioxRealtimeBackend, + clock: Callable[[], float] = time.monotonic, + drop_params: bool = False, + ) -> None: + self._options: Final = options or SonioxRealtimeOptions() + self._drop_params: Final = drop_params + self._wrap: Final = wrap + self._clock: Final = clock + self._transformer: Final = SonioxEventTransformer(drops_originals=self._options.translation is not None) + self._start_request: SonioxStartRequest | None = None + self._session_id: str = f"sess_{uuid.uuid4().hex}" + self._started_at: float | None = None + + @classmethod + def from_litellm_params(cls, litellm_params: Mapping[str, object], *, drop_params: bool = False) -> Self: + return cls(options=SonioxRealtimeOptions.model_validate(dict(litellm_params)), drop_params=drop_params) + + def requires_session_configuration(self) -> bool: + return True + + def validate_environment( + self, + headers: dict[str, str], # mutable-ok: BaseRealtimeConfig contract + model: str, + api_key: str | None = None, + ) -> dict[str, str]: # mutable-ok: BaseRealtimeConfig contract + resolved_key: Final = get_soniox_api_key(api_key) + if not resolved_key: + raise ValueError("Missing Soniox API key. Set SONIOX_API_KEY or api_key in the deployment litellm_params") + return {**headers, "Authorization": f"Bearer {resolved_key}"} + + def get_complete_url(self, api_base: str | None, model: str, api_key: str | None = None) -> str: + return build_soniox_realtime_url(api_base) + + def wrap_backend(self, backend: RealtimeBackend) -> RealtimeBackend: + return self._wrap(backend) + + def transform_session_created_event( + self, + model: str, + logging_session_id: str, + session_configuration_request: str | None = None, + ) -> OpenAIRealtimeTranscriptionSessionCreated: + self._session_id = logging_session_id + return _session_created(logging_session_id, build_start_request(model, None, self._options)) + + def transform_realtime_request( + self, + message: str, + model: str, + session_configuration_request: str | None = None, + ) -> tuple[str | bytes, ...]: + request: Final = json_object(message, SonioxProtocolError) + event_type: Final = request.get("type") + if event_type is None: + raise SonioxProtocolError("Soniox realtime takes OpenAI realtime events, and this message has no type") + if event_type in SESSION_UPDATE_EVENT_TYPES: + update: Final = parse_transcription_session_update(message, SonioxProtocolError) + ignored: Final = _ignored_settings(update) + if ignored: + verbose_logger.warning("Soniox realtime: ignoring unsupported session settings %s", ignored) + return self._start(model, update) + if event_type == "input_audio_buffer.append": + audio: Final = decode_pcm16_append(request.get("audio"), error=SonioxProtocolError) + return (*self._start(model, None), audio) + if event_type == "input_audio_buffer.commit": + return self._finalize(model, discards_turn=False) + if event_type == "input_audio_buffer.clear": + return () if self._start_request is None else self._finalize(model, discards_turn=True) + if event_type == "input_audio_buffer.end": + return (*self._start(model, None), _END_STREAM) + verbose_logger.debug("Soniox realtime: dropping unsupported client event %s", event_type) + return () + + def _finalize(self, model: str, *, discards_turn: bool) -> tuple[str, ...]: + start: Final = self._start(model, None) + self._transformer.queue_finalize(discards_turn=discards_turn) + return (*start, _FINALIZE) + + def session_refusal(self) -> Exception | None: + return self._transformer.refusal + + def unbilled_usage_on_session_close(self, model: str) -> RealtimeInputAudioTranscriptionUsage | None: + started_at: Final = self._started_at + if self._transformer.refusal is not None: + return None + if started_at is None: + return self._transformer.take_unbilled_usage() + return self._transformer.take_unbilled_usage(int((self._clock() - started_at) * 1000)) + + def transform_realtime_response( + self, + message: str | bytes, + model: str, + logging_obj: LiteLLMLoggingObj, + realtime_response_transform_input: RealtimeResponseTransformInput, + ) -> RealtimeResponseTypedDict: + payload: Final = message.decode("utf-8") if isinstance(message, bytes) else message + result: Final[RealtimeResponseTypedDict] = { + "response": list(self._backend_events(payload)), + "current_output_item_id": realtime_response_transform_input.get("current_output_item_id"), + "current_response_id": realtime_response_transform_input.get("current_response_id"), + "current_delta_chunks": realtime_response_transform_input.get("current_delta_chunks"), + "current_conversation_id": realtime_response_transform_input.get("current_conversation_id"), + "current_item_chunks": realtime_response_transform_input.get("current_item_chunks"), + "current_delta_type": realtime_response_transform_input.get("current_delta_type"), + "session_configuration_request": realtime_response_transform_input.get("session_configuration_request"), + } + return result + + def _backend_events(self, payload: str) -> tuple[OpenAIRealtimeEvents, ...]: + if payload != SESSION_STARTED_FRAME: + return self._transformer.transform(payload) + assert self._start_request is not None + return (_session_created(self._session_id, self._start_request),) + + def _start(self, model: str, update: TranscriptionSessionUpdate | None) -> tuple[str, ...]: + started: Final = self._start_request + if started is not None: + if update is not None and _changes_started_stream(update, started, self._drop_params): + raise SonioxProtocolError("Soniox realtime can't change the session once the stream has started") + return () + start_request: Final = build_start_request(model, update, self._options, self._drop_params) + self._start_request = start_request + self._started_at = self._clock() + return (start_request.model_dump_json(exclude_none=True),) + + +def _finishes_stream(frame: str | bytes) -> bool: + response: Final = _SonioxResponse.model_validate_json(frame) + error: Final = _soniox_error(response) + if error is not None: + raise error + return response.finished + + +async def soniox_realtime_health_check( + model: str, api_base: str | None, api_key: str | None, options: SonioxRealtimeOptions +) -> bool: + import websockets + + config: Final = SonioxRealtimeConfig() + url: Final = config.get_complete_url(api_base, model) + headers: Final = config.validate_environment({}, model, api_key) + start_request: Final = build_start_request(model, None, options) + async with websockets.connect( + url, + additional_headers=headers, + max_size=REALTIME_WEBSOCKET_MAX_MESSAGE_SIZE_BYTES, + ssl=realtime_ssl_for_url(url), + ) as backend: + await backend.send(start_request.model_dump_json(exclude_none=True)) + await backend.send(_HEALTH_CHECK_SILENCE) + await backend.send(_END_STREAM) + async for frame in backend: + if _finishes_stream(frame): + return True + raise SonioxException(status_code=502, message="Soniox closed the realtime stream without finishing it") diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 4fad4857c60..1256fef9181 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -61177,6 +61177,126 @@ ], "supports_audio_input": true }, + "soniox/stt-rt-v5": { + "input_cost_per_second": 3.333e-05, + "litellm_provider": "soniox", + "metadata": { + "calculation": "$0.12/3600 seconds = $0.00003333 per second", + "original_pricing_per_hour": 0.12 + }, + "mode": "audio_transcription", + "source": "https://soniox.com/pricing", + "supported_endpoints": [ + "/v1/realtime" + ], + "supported_modalities": [ + "audio" + ], + "supported_output_modalities": [ + "text" + ], + "supports_audio_input": true + }, + "soniox/stt-rt-preview": { + "input_cost_per_second": 3.333e-05, + "litellm_provider": "soniox", + "metadata": { + "calculation": "$0.12/3600 seconds = $0.00003333 per second", + "original_pricing_per_hour": 0.12 + }, + "mode": "audio_transcription", + "source": "https://soniox.com/pricing", + "supported_endpoints": [ + "/v1/realtime" + ], + "supported_modalities": [ + "audio" + ], + "supported_output_modalities": [ + "text" + ], + "supports_audio_input": true + }, + "soniox/stt-rt-preview-v2": { + "input_cost_per_second": 3.333e-05, + "litellm_provider": "soniox", + "metadata": { + "calculation": "$0.12/3600 seconds = $0.00003333 per second", + "original_pricing_per_hour": 0.12 + }, + "mode": "audio_transcription", + "source": "https://soniox.com/pricing", + "supported_endpoints": [ + "/v1/realtime" + ], + "supported_modalities": [ + "audio" + ], + "supported_output_modalities": [ + "text" + ], + "supports_audio_input": true + }, + "soniox/stt-rt-v3": { + "input_cost_per_second": 3.333e-05, + "litellm_provider": "soniox", + "metadata": { + "calculation": "$0.12/3600 seconds = $0.00003333 per second", + "original_pricing_per_hour": 0.12 + }, + "mode": "audio_transcription", + "source": "https://soniox.com/pricing", + "supported_endpoints": [ + "/v1/realtime" + ], + "supported_modalities": [ + "audio" + ], + "supported_output_modalities": [ + "text" + ], + "supports_audio_input": true + }, + "soniox/stt-rt-v3-preview": { + "input_cost_per_second": 3.333e-05, + "litellm_provider": "soniox", + "metadata": { + "calculation": "$0.12/3600 seconds = $0.00003333 per second", + "original_pricing_per_hour": 0.12 + }, + "mode": "audio_transcription", + "source": "https://soniox.com/pricing", + "supported_endpoints": [ + "/v1/realtime" + ], + "supported_modalities": [ + "audio" + ], + "supported_output_modalities": [ + "text" + ], + "supports_audio_input": true + }, + "soniox/stt-rt-v4": { + "input_cost_per_second": 3.333e-05, + "litellm_provider": "soniox", + "metadata": { + "calculation": "$0.12/3600 seconds = $0.00003333 per second", + "original_pricing_per_hour": 0.12 + }, + "mode": "audio_transcription", + "source": "https://soniox.com/pricing", + "supported_endpoints": [ + "/v1/realtime" + ], + "supported_modalities": [ + "audio" + ], + "supported_output_modalities": [ + "text" + ], + "supports_audio_input": true + }, "tensormesh/Qwen/Qwen3.5-397B-A17B-FP8": { "litellm_provider": "tensormesh", "mode": "chat", diff --git a/litellm/provider_endpoints_support_backup.json b/litellm/provider_endpoints_support_backup.json index 4188866de08..9f5989e765c 100644 --- a/litellm/provider_endpoints_support_backup.json +++ b/litellm/provider_endpoints_support_backup.json @@ -3724,7 +3724,7 @@ "ocr": false, "search": false, "moderations": false, - "realtime": false, + "realtime": true, "video_generations": false, "batches": false, "messages": false, diff --git a/litellm/proxy/public_endpoints/provider_create_fields.json b/litellm/proxy/public_endpoints/provider_create_fields.json index 4237517508b..04e6c5d8ba9 100644 --- a/litellm/proxy/public_endpoints/provider_create_fields.json +++ b/litellm/proxy/public_endpoints/provider_create_fields.json @@ -3612,7 +3612,7 @@ "key": "api_key", "label": "Soniox API Key", "placeholder": null, - "tooltip": "Currently only the async Speech-to-Text REST API (api.soniox.com) is supported. Realtime STT (stt-rt.soniox.com) and TTS (tts-rt.soniox.com) are not yet available.", + "tooltip": "Supports async Speech-to-Text (api.soniox.com) on /v1/audio/transcriptions and realtime Speech-to-Text (stt-rt.soniox.com) on /v1/realtime. TTS (tts-rt.soniox.com) is not yet available.", "required": true, "field_type": "password", "options": null, diff --git a/litellm/realtime_api/README.md b/litellm/realtime_api/README.md index d810de2f24f..987e8be1ed1 100644 --- a/litellm/realtime_api/README.md +++ b/litellm/realtime_api/README.md @@ -4,6 +4,6 @@ Supported endpoints: - WebSocket: `/v1/realtime` (with `intent=transcription` for transcription-only sessions) - HTTP: `/v1/realtime/client_secrets`, `/v1/realtime/transcription_sessions` -Supported providers: OpenAI, Azure OpenAI, Bedrock, Vertex AI, xAI. +Supported providers: OpenAI, Azure OpenAI, Bedrock, Vertex AI, xAI, Gemini, Meta, Soniox. For user-facing documentation and usage examples, see the litellm-docs repo. \ No newline at end of file diff --git a/litellm/realtime_api/main.py b/litellm/realtime_api/main.py index e9b534bdd89..919fd54e7d6 100644 --- a/litellm/realtime_api/main.py +++ b/litellm/realtime_api/main.py @@ -7,6 +7,7 @@ from dataclasses import dataclass from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final, Literal, NoReturn, cast +from pydantic import TypeAdapter from typing_extensions import ReadOnly, TypedDict, assert_never import litellm @@ -55,6 +56,7 @@ from ..llms.openai.realtime.handler import ( confirm_session_started, ) from ..llms.openai.workload_identity import openai_static_api_key, resolve_openai_bearer_token_async +from ..llms.soniox.realtime.transformation import SonioxRealtimeOptions, soniox_realtime_health_check from ..llms.vertex_ai.audio_transcription.realtime_transformation import is_vertex_speech_to_text_model from ..llms.vertex_ai.realtime.transformation import VertexAIRealtimeConfig, vertex_realtime_config from ..llms.vertex_ai.vertex_llm_base import VertexBase @@ -75,6 +77,7 @@ vertex_llm_base: Final = VertexBase() base_llm_http_handler = BaseLLMHTTPHandler() _EMPTY_MODEL_PARAMS: Final[Mapping[str, object]] = MappingProxyType({}) _EMPTY_AUTH_HEADERS: Final[Mapping[str, str]] = MappingProxyType({}) +_DEPLOYMENT_PARAMS: Final = TypeAdapter(Mapping[str, object]) def _model_params_with_stored_credentials(model_params: Mapping[str, object]) -> Mapping[str, object]: @@ -527,6 +530,7 @@ async def _arealtime( custom_llm_provider=_custom_llm_provider, ) + deployment_params: Final = _DEPLOYMENT_PARAMS.validate_python(kwargs) if live_session_start is not None: await _alive_session( model=model, @@ -537,7 +541,7 @@ async def _arealtime( api_base=dynamic_api_base or litellm_params.api_base, api_key=dynamic_api_key, litellm_metadata=_build_litellm_metadata(kwargs), - litellm_params=kwargs, + litellm_params=deployment_params, ) return @@ -548,6 +552,7 @@ async def _arealtime( model=model, provider=LlmProviders(_custom_llm_provider), drop_params=drop_params, + litellm_params=deployment_params, ) if provider_config is not None: await base_llm_http_handler.async_realtime( @@ -882,6 +887,13 @@ async def realtime_health_check( ) elif custom_llm_provider == "openai": url = openai_realtime.construct_url(api_base=openai_api_base, query_params={"model": model}) + elif custom_llm_provider == "soniox": + return await soniox_realtime_health_check( + model=model, + api_base=resolved_api_base, + api_key=health_check_api_key, + options=SonioxRealtimeOptions.model_validate(dict(resolved_params)), + ) elif custom_llm_provider == "xai": url = xai_realtime.construct_url( api_base=resolved_api_base or "https://api.x.ai/v1", query_params={"model": model} diff --git a/litellm/utils.py b/litellm/utils.py index a656d8ebca3..342a2f12310 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -10759,6 +10759,7 @@ class ProviderConfigManager: model: str, provider: LlmProviders, drop_params: bool = False, + litellm_params: Mapping[str, object] = MappingProxyType({}), ) -> BaseRealtimeConfig | None: if LlmProviders.GEMINI == provider: from litellm.llms.gemini.realtime.transformation import GeminiRealtimeConfig @@ -10768,6 +10769,10 @@ class ProviderConfigManager: from litellm.llms.meta.realtime.transformation import MetaRealtimeConfig return MetaRealtimeConfig(drop_params=drop_params) + if LlmProviders.SONIOX == provider: + from litellm.llms.soniox.realtime.transformation import SonioxRealtimeConfig + + return SonioxRealtimeConfig.from_litellm_params(litellm_params, drop_params=drop_params) return None @staticmethod diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 4fad4857c60..1256fef9181 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -61177,6 +61177,126 @@ ], "supports_audio_input": true }, + "soniox/stt-rt-v5": { + "input_cost_per_second": 3.333e-05, + "litellm_provider": "soniox", + "metadata": { + "calculation": "$0.12/3600 seconds = $0.00003333 per second", + "original_pricing_per_hour": 0.12 + }, + "mode": "audio_transcription", + "source": "https://soniox.com/pricing", + "supported_endpoints": [ + "/v1/realtime" + ], + "supported_modalities": [ + "audio" + ], + "supported_output_modalities": [ + "text" + ], + "supports_audio_input": true + }, + "soniox/stt-rt-preview": { + "input_cost_per_second": 3.333e-05, + "litellm_provider": "soniox", + "metadata": { + "calculation": "$0.12/3600 seconds = $0.00003333 per second", + "original_pricing_per_hour": 0.12 + }, + "mode": "audio_transcription", + "source": "https://soniox.com/pricing", + "supported_endpoints": [ + "/v1/realtime" + ], + "supported_modalities": [ + "audio" + ], + "supported_output_modalities": [ + "text" + ], + "supports_audio_input": true + }, + "soniox/stt-rt-preview-v2": { + "input_cost_per_second": 3.333e-05, + "litellm_provider": "soniox", + "metadata": { + "calculation": "$0.12/3600 seconds = $0.00003333 per second", + "original_pricing_per_hour": 0.12 + }, + "mode": "audio_transcription", + "source": "https://soniox.com/pricing", + "supported_endpoints": [ + "/v1/realtime" + ], + "supported_modalities": [ + "audio" + ], + "supported_output_modalities": [ + "text" + ], + "supports_audio_input": true + }, + "soniox/stt-rt-v3": { + "input_cost_per_second": 3.333e-05, + "litellm_provider": "soniox", + "metadata": { + "calculation": "$0.12/3600 seconds = $0.00003333 per second", + "original_pricing_per_hour": 0.12 + }, + "mode": "audio_transcription", + "source": "https://soniox.com/pricing", + "supported_endpoints": [ + "/v1/realtime" + ], + "supported_modalities": [ + "audio" + ], + "supported_output_modalities": [ + "text" + ], + "supports_audio_input": true + }, + "soniox/stt-rt-v3-preview": { + "input_cost_per_second": 3.333e-05, + "litellm_provider": "soniox", + "metadata": { + "calculation": "$0.12/3600 seconds = $0.00003333 per second", + "original_pricing_per_hour": 0.12 + }, + "mode": "audio_transcription", + "source": "https://soniox.com/pricing", + "supported_endpoints": [ + "/v1/realtime" + ], + "supported_modalities": [ + "audio" + ], + "supported_output_modalities": [ + "text" + ], + "supports_audio_input": true + }, + "soniox/stt-rt-v4": { + "input_cost_per_second": 3.333e-05, + "litellm_provider": "soniox", + "metadata": { + "calculation": "$0.12/3600 seconds = $0.00003333 per second", + "original_pricing_per_hour": 0.12 + }, + "mode": "audio_transcription", + "source": "https://soniox.com/pricing", + "supported_endpoints": [ + "/v1/realtime" + ], + "supported_modalities": [ + "audio" + ], + "supported_output_modalities": [ + "text" + ], + "supports_audio_input": true + }, "tensormesh/Qwen/Qwen3.5-397B-A17B-FP8": { "litellm_provider": "tensormesh", "mode": "chat", diff --git a/provider_endpoints_support.json b/provider_endpoints_support.json index 4188866de08..9f5989e765c 100644 --- a/provider_endpoints_support.json +++ b/provider_endpoints_support.json @@ -3724,7 +3724,7 @@ "ocr": false, "search": false, "moderations": false, - "realtime": false, + "realtime": true, "video_generations": false, "batches": false, "messages": false, diff --git a/tests/unit/litellm_core_utils/test_health_check_helpers.py b/tests/unit/litellm_core_utils/test_health_check_helpers.py index 44b36a355c6..6979d72ff14 100644 --- a/tests/unit/litellm_core_utils/test_health_check_helpers.py +++ b/tests/unit/litellm_core_utils/test_health_check_helpers.py @@ -795,9 +795,12 @@ async def test_ahealth_check_keeps_the_noul_probe_for_other_decisions_providers( ("us.anthropic.claude-haiku-4-5-20251001-v1:0", "bedrock", None), ("claude-haiku-4-5", "anthropic", None), ("anthropic.claude-haiku-4-5", None, None), + ("stt-rt-v5", "soniox", "realtime"), + ("stt-rt-preview", "soniox", "realtime"), + ("stt-async-v5", "soniox", None), ), ) -def test_native_health_check_mode_is_messages_only_for_mantle_claude( +def test_native_health_check_mode_routes_mantle_claude_to_messages_and_soniox_realtime_models_to_realtime( model: str, custom_llm_provider: str | None, expected: str | None ) -> None: assert native_health_check_mode(model=model, custom_llm_provider=custom_llm_provider) == expected diff --git a/tests/unit/litellm_core_utils/test_realtime_streaming.py b/tests/unit/litellm_core_utils/test_realtime_streaming.py index 7d85da3f756..a8543ebcfc2 100644 --- a/tests/unit/litellm_core_utils/test_realtime_streaming.py +++ b/tests/unit/litellm_core_utils/test_realtime_streaming.py @@ -24,6 +24,7 @@ from litellm.litellm_core_utils.realtime_streaming import ( client_sent_openai_beta_realtime_header, ) from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj +from litellm.llms.base_llm.realtime.transcription_protocol import RealtimeTranscriptionProtocolError from litellm.llms.base_llm.realtime.transformation import BaseRealtimeConfig from litellm.llms.xai.realtime.transformation import XAIRealtimeNormalizer from litellm.llms.xai.realtime.usage_meter import XAIRealtimeUsageMeter @@ -827,9 +828,7 @@ async def test_usage_meter_snapshot_stored_in_model_call_details(): backend_ws.send = AsyncMock() backend_ws.recv = AsyncMock( side_effect=[ - json.dumps( - {"type": "session.created", "session": {"turn_detection": {"type": None}}} - ).encode(), + json.dumps({"type": "session.created", "session": {"turn_detection": {"type": None}}}).encode(), json.dumps( {"type": "response.output_audio.delta", "delta": base64.b64encode(b"x" * 24000).decode()} ).encode(), @@ -3187,6 +3186,7 @@ async def test_session_close_flushes_unbilled_transcription_usage(): provider_config: Final = MagicMock() provider_config.transform_realtime_response = MagicMock(return_value=transcript_frame) provider_config.unbilled_usage_on_session_close = MagicMock(return_value=usage) + provider_config.session_refusal = MagicMock(return_value=None) streaming: Final = RealTimeStreaming( client_ws, @@ -3320,10 +3320,14 @@ def _backend_ws_closing_with(*frames: bytes | Exception) -> MagicMock: return backend_ws -def _relay_session(client_ws: MagicMock, backend_ws: MagicMock) -> _RelaySession: +def _relay_session( + client_ws: MagicMock, backend_ws: MagicMock, provider_config: BaseRealtimeConfig | None = None +) -> _RelaySession: logging: Final = _RecordingLogging() worker: Final = _InlineLoggingWorker() - streaming: Final = RealTimeStreaming(client_ws, backend_ws, logging, model="gpt-realtime", logging_worker=worker) + streaming: Final = RealTimeStreaming( + client_ws, backend_ws, logging, provider_config=provider_config, model="gpt-realtime", logging_worker=worker + ) return _RelaySession(streaming=streaming, logging=logging, worker=worker) @@ -3408,6 +3412,103 @@ async def test_upstream_refusal_before_any_frame_logs_a_failure_not_a_success(): assert session.logging.logged_sessions == () +class _RefusingRealtimeConfig(_PassthroughRealtimeConfig): + def __init__(self, refusal: Exception | None) -> None: + self._refusal: Final = refusal + + def requires_session_configuration(self) -> bool: + return True + + def session_refusal(self) -> Exception | None: + return self._refusal + + +@pytest.mark.asyncio +async def test_a_session_the_provider_refused_in_band_logs_a_failure_even_after_a_normal_close(): + refusal: Final = ValueError("Soniox realtime error 401: Invalid API key.") + session: Final = _relay_session( + _client_ws_that_never_sends(), + _backend_ws_closing_with( + json.dumps({"type": "error", "error": {"message": str(refusal)}}).encode(), + ConnectionClosed(Close(1000, ""), None), + ), + _RefusingRealtimeConfig(refusal), + ) + + await session.run() + + assert session.logging.logged_failures == (refusal,) + assert session.logging.logged_sessions == () + + +@pytest.mark.asyncio +async def test_a_provider_session_without_a_refusal_logs_a_success_after_a_normal_close(): + session: Final = _relay_session( + _client_ws_that_never_sends(), + _backend_ws_closing_with( + json.dumps({"type": "session.created", "session": {"id": "sess_1"}}).encode(), + ConnectionClosed(Close(1000, ""), None), + ), + _RefusingRealtimeConfig(None), + ) + + await session.run() + + assert session.logging.logged_failures == () + assert len(session.logging.logged_sessions) == 1 + + +class _TranscriptionConfigRejectingSessionUpdates(_PassthroughRealtimeConfig): + def __init__(self, *, defers_setup: bool) -> None: + self._defers_setup: Final = defers_setup + + def requires_session_configuration(self) -> bool: + return not self._defers_setup + + def transform_realtime_request( + self, message: str, model: str, session_configuration_request: str | None = None + ) -> Sequence[str | bytes]: + if json.loads(message).get("type") == "session.update": + raise RealtimeTranscriptionProtocolError("requires pcm16 input audio") + return (message,) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "defers_setup,client_frames_read,audio_reaches_backend", + ((True, 1, False), (False, 3, True)), + ids=["deferred_setup_session_ends", "eager_setup_session_continues"], +) +async def test_a_rejected_transcription_session_update_sends_an_error_and_ends_only_a_deferred_setup_session( + defers_setup: bool, client_frames_read: int, audio_reaches_backend: bool +): + append: Final = json.dumps({"type": "input_audio_buffer.append", "audio": "AAE="}) + client_ws: Final = MagicMock() + client_ws.headers = {} + client_ws.receive_text = AsyncMock( + side_effect=[json.dumps({"type": "session.update", "session": {"type": "transcription"}}), append, RuntimeError] + ) + client_ws.send_text = AsyncMock() + backend_ws: Final = MagicMock() + backend_ws.send = AsyncMock() + streaming: Final = RealTimeStreaming( + client_ws, + backend_ws, + _RecordingLogging(), + provider_config=_TranscriptionConfigRejectingSessionUpdates(defers_setup=defers_setup), + model="chirp_3", + logging_worker=_InlineLoggingWorker(), + ) + + await asyncio.wait_for(streaming.client_ack_messages(), timeout=2) + + assert _error_events_sent_to(client_ws) == [ + {"type": "error", "error": {"type": "invalid_request_error", "message": "requires pcm16 input audio"}} + ] + assert client_ws.receive_text.await_count == client_frames_read + assert (append in [call.args[0] for call in backend_ws.send.await_args_list]) is audio_reaches_backend + + @pytest.mark.asyncio async def test_upstream_refusal_after_a_synthetic_session_created_still_logs_a_failure(): """LIT-6973: deferred Gemini Live setup stores a synthetic ``session.created`` before @@ -4056,9 +4157,7 @@ _XAI_TYPED_VAD_SESSION: Final = { pytest.param(({"type": "realtime", "instructions": "hi"},), id="no-turn-detection"), pytest.param(({"type": "realtime", "audio": {"input": {"turn_detection": None}}},), id="push-to-talk"), pytest.param((_XAI_TYPED_VAD_SESSION,), id="typed-server-vad"), - pytest.param( - ({"type": "realtime", "instructions": "hi"}, _XAI_TYPED_VAD_SESSION), id="later-typed-server-vad" - ), + pytest.param(({"type": "realtime", "instructions": "hi"}, _XAI_TYPED_VAD_SESSION), id="later-typed-server-vad"), ], ) async def test_xai_session_update_gate_leaves_out_the_vad_type( @@ -4070,9 +4169,9 @@ async def test_xai_session_update_gate_leaves_out_the_vad_type( monkeypatch, *client_sessions, event_normalizer=XAIRealtimeNormalizer() ) - assert [session["audio"]["input"]["turn_detection"] for session in forwarded] == [ - {"create_response": False} - ] * (len(client_sessions) + 1) + assert [session["audio"]["input"]["turn_detection"] for session in forwarded] == [{"create_response": False}] * ( + len(client_sessions) + 1 + ) @pytest.mark.asyncio @@ -4206,7 +4305,9 @@ async def test_client_never_sees_the_error_caused_by_the_guardrail_cancel(monkey await RealTimeStreaming(client_ws, backend_ws, logging_obj).backend_to_client_send_messages() client_errors: Final = [ - event["error"] for call in client_ws.send_text.await_args_list if (event := json.loads(call.args[0]))["type"] == "error" + event["error"] + for call in client_ws.send_text.await_args_list + if (event := json.loads(call.args[0]))["type"] == "error" ] assert [(error["code"], error.get("event_id")) for error in client_errors] == [ ("content_policy_violation", None), diff --git a/tests/unit/llms/custom_httpx/test_llm_http_handler.py b/tests/unit/llms/custom_httpx/test_llm_http_handler.py index d8326536d03..dca21db5b6f 100644 --- a/tests/unit/llms/custom_httpx/test_llm_http_handler.py +++ b/tests/unit/llms/custom_httpx/test_llm_http_handler.py @@ -4911,6 +4911,160 @@ async def test_async_realtime_bridges_a_transcription_session_through_the_provid assert [bytes(request.audio) for request in speech_client.requests[1:]] == [b"\x00\x01" * 800, b"\x00\x01" * 800] +class _FakeSonioxSocket: + def __init__(self) -> None: + self.sent: Final[list[str | bytes]] = [] + self._responses: Final[asyncio.Queue[str]] = asyncio.Queue() + + async def __aenter__(self) -> "_FakeSonioxSocket": + return self + + async def __aexit__(self, exc_type, exc_value, traceback) -> None: + return None + + async def send(self, message: str | bytes) -> None: + self.sent.append(message) + if message == '{"type":"finalize"}': + tokens = [{"text": "Hello", "is_final": True}, {"text": "", "is_final": True}] + await self._responses.put(json.dumps({"tokens": tokens, "total_audio_proc_ms": 1200})) + + async def recv(self, decode: bool | None = None) -> str: + return await self._responses.get() + + async def close(self) -> None: + return None + + +@pytest.mark.asyncio +async def test_async_realtime_bridges_a_soniox_transcription_through_the_wrapped_websocket(): + import websockets.exceptions # noqa: F401 # binds the submodule so async_realtime's except clause resolves, as in the proxy process + + from litellm.llms.soniox.realtime.transformation import SonioxRealtimeBackend, SonioxRealtimeConfig + + upstream = _FakeSonioxSocket() + wrapped: list[object] = [] + provider_config = SonioxRealtimeConfig( + wrap=lambda backend: wrapped.append(backend) or SonioxRealtimeBackend(backend) + ) + audio = b"\x00\x01" * 800 + client_ws = _ScriptedClientWebSocket( + [ + json.dumps( + { + "type": "session.update", + "session": { + "type": "transcription", + "audio": { + "input": { + "format": {"type": "audio/pcm", "rate": 16000}, + "transcription": {"model": "stt-rt-v5", "language": "en"}, + "turn_detection": None, + } + }, + }, + } + ), + json.dumps({"type": "input_audio_buffer.append", "audio": base64.b64encode(audio).decode()}), + json.dumps({"type": "input_audio_buffer.commit"}), + ], + last_event_type="conversation.item.input_audio_transcription.completed", + ) + logging_obj = Mock() + logging_obj.litellm_trace_id = "trace_1" + logging_obj.model_call_details = {} + logging_obj.dispatch_success_handlers = AsyncMock() + logging_obj.dispatch_failure_handlers = AsyncMock() + handler = BaseLLMHTTPHandler() + + with patch.object(handler, "_open_realtime_backend_ws", AsyncMock(return_value=upstream)) as dial: + await handler.async_realtime( + model="stt-rt-v5", + websocket=client_ws, + logging_obj=logging_obj, + provider_config=provider_config, + headers={}, + api_key="soniox-key", + query_params={"model": "stt-rt-v5", "intent": "transcription"}, + ) + + assert dial.await_args.args[1:3] == ( + "wss://stt-rt.soniox.com/transcribe-websocket", + {"Authorization": "Bearer soniox-key"}, + ) + assert wrapped == [upstream] + assert json.loads(upstream.sent[0]) == { + "model": "stt-rt-v5", + "audio_format": "pcm_s16le", + "sample_rate": 16000, + "num_channels": 1, + "enable_endpoint_detection": False, + "language_hints": ["en"], + } + assert upstream.sent[1:] == [audio, '{"type":"finalize"}'] + events = client_ws.sent_events() + assert [event["type"] for event in events] == [ + "session.created", + "session.updated", + "input_audio_buffer.speech_started", + "conversation.item.input_audio_transcription.delta", + "conversation.item.input_audio_transcription.completed", + ] + assert events[-1]["transcript"] == "Hello" + assert events[-1]["usage"] == {"type": "duration", "seconds": 1.2} + + +@pytest.mark.asyncio +async def test_async_realtime_soniox_rejects_a_bad_session_update_and_still_transcribes_audio_sent_without_one(): + import websockets.exceptions # noqa: F401 # binds the submodule so async_realtime's except clause resolves, as in the proxy process + + from litellm.llms.soniox.realtime.transformation import SonioxRealtimeConfig + + upstream = _FakeSonioxSocket() + audio = b"\x00\x01" * 800 + client_ws = _ScriptedClientWebSocket( + [ + json.dumps( + { + "type": "session.update", + "session": {"type": "transcription", "audio": {"input": {"format": {"type": "audio/pcmu"}}}}, + } + ), + json.dumps({"type": "input_audio_buffer.append", "audio": base64.b64encode(audio).decode()}), + json.dumps({"type": "input_audio_buffer.commit"}), + ], + last_event_type="conversation.item.input_audio_transcription.completed", + ) + logging_obj = Mock() + logging_obj.litellm_trace_id = "trace_1" + logging_obj.model_call_details = {} + logging_obj.dispatch_success_handlers = AsyncMock() + logging_obj.dispatch_failure_handlers = AsyncMock() + handler = BaseLLMHTTPHandler() + + with patch.object(handler, "_open_realtime_backend_ws", AsyncMock(return_value=upstream)): + await handler.async_realtime( + model="stt-rt-v5", + websocket=client_ws, + logging_obj=logging_obj, + provider_config=SonioxRealtimeConfig(), + headers={}, + api_key="soniox-key", + query_params={"model": "stt-rt-v5", "intent": "transcription"}, + ) + + assert json.loads(upstream.sent[0])["sample_rate"] == 24_000 + assert upstream.sent[1:] == [audio, '{"type":"finalize"}'] + events = client_ws.sent_events() + errors = [event for event in events if event["type"] == "error"] + assert errors == [ + { + "type": "error", + "error": {"type": "invalid_request_error", "message": "Soniox realtime requires pcm16 input audio"}, + } + ] + assert events[-1]["transcript"] == "Hello" + + @pytest.mark.asyncio async def test_responses_agentic_followup_does_not_repeat_request_params_from_plan_kwargs(monkeypatch): """A plan whose kwargs repeat a request param must not crash the Responses follow-up with a duplicate keyword""" diff --git a/tests/unit/llms/soniox/realtime/__init__.py b/tests/unit/llms/soniox/realtime/__init__.py new file mode 100644 index 00000000000..b503dff241f --- /dev/null +++ b/tests/unit/llms/soniox/realtime/__init__.py @@ -0,0 +1 @@ +"""Soniox realtime tests.""" diff --git a/tests/unit/llms/soniox/realtime/test_soniox_realtime_transformation.py b/tests/unit/llms/soniox/realtime/test_soniox_realtime_transformation.py new file mode 100644 index 00000000000..653f73ce22c --- /dev/null +++ b/tests/unit/llms/soniox/realtime/test_soniox_realtime_transformation.py @@ -0,0 +1,807 @@ +import asyncio +import base64 +import json +from typing import Final +from unittest.mock import MagicMock + +import pytest + +from litellm.llms.soniox.common_utils import SonioxException +from litellm.llms.soniox.realtime.transformation import ( + DEFAULT_SONIOX_REALTIME_URL, + SESSION_STARTED_FRAME, + SonioxEventTransformer, + SonioxRealtimeBackend, + SonioxProtocolError, + SonioxRealtimeConfig, + SonioxRealtimeOptions, +) +from litellm.types.llms.openai import OpenAIRealtimeEvents +from litellm.types.realtime import RealtimeResponseTransformInput +from litellm.types.utils import LlmProviders +from litellm.utils import ProviderConfigManager + +MODEL: Final = "stt-rt-v5" +EMPTY_TRANSFORM_INPUT: Final[RealtimeResponseTransformInput] = { + "session_configuration_request": None, + "current_output_item_id": None, + "current_response_id": None, + "current_delta_chunks": None, + "current_item_chunks": None, + "current_conversation_id": None, + "current_delta_type": None, +} +PCM: Final = b"\x01\x02" * 160 + + +def _event(event_type: str, **fields: object) -> str: + return json.dumps({"type": event_type, **fields}) + + +def _session_update(audio_input: dict[str, object]) -> str: + return _event("session.update", session={"type": "transcription", "audio": {"input": audio_input}}) + + +def _append() -> str: + return _event("input_audio_buffer.append", audio=base64.b64encode(PCM).decode()) + + +def _start_request(frames: tuple[str | bytes, ...]) -> dict[str, object]: + assert isinstance(frames[0], str) + return json.loads(frames[0]) + + +def _response(tokens: list[dict[str, object]], total_audio_proc_ms: int | None = None) -> str: + return json.dumps({"tokens": tokens, "total_audio_proc_ms": total_audio_proc_ms}) + + +def _token(text: str, is_final: bool = True, **fields: object) -> dict[str, object]: + return {"text": text, "is_final": is_final, **fields} + + +def _types(events: tuple[OpenAIRealtimeEvents, ...]) -> list[object]: + return [event["type"] for event in events] + + +def test_session_update_starts_the_stream_with_client_audio_and_deployment_options(): + config = SonioxRealtimeConfig( + options=SonioxRealtimeOptions(language_hints=("en", "fr"), enable_speaker_diarization=True) + ) + + frames = config.transform_realtime_request( + _session_update( + { + "format": {"type": "audio/pcm", "rate": 16_000}, + "transcription": {"model": "client-chosen-model", "language": "fr"}, + "turn_detection": None, + } + ), + MODEL, + ) + + assert _start_request(frames) == { + "model": MODEL, + "audio_format": "pcm_s16le", + "sample_rate": 16_000, + "num_channels": 1, + "enable_endpoint_detection": False, + "language_hints": ["fr", "en"], + "enable_speaker_diarization": True, + } + + +def test_an_empty_transcription_language_keeps_the_deployment_language_hints(): + config = SonioxRealtimeConfig(options=SonioxRealtimeOptions(language_hints=("fa",))) + + frames = config.transform_realtime_request(_session_update({"transcription": {"language": ""}}), MODEL) + + assert _start_request(frames)["language_hints"] == ["fa"] + + +def test_a_client_language_cannot_widen_strict_deployment_language_hints(): + config = SonioxRealtimeConfig(options=SonioxRealtimeOptions(language_hints=("fa",), language_hints_strict=True)) + + frames = config.transform_realtime_request(_session_update({"transcription": {"language": "en"}}), MODEL) + + assert _start_request(frames)["language_hints"] == ["fa"] + assert config.transform_realtime_request(_session_update({"transcription": {"language": "en"}}), MODEL) == () + + +def test_a_soniox_native_config_frame_is_rejected_without_starting_the_stream(): + config = SonioxRealtimeConfig() + + with pytest.raises(SonioxProtocolError): + config.transform_realtime_request(json.dumps({"model": "stt-rt-v5", "audio_format": "auto"}), MODEL) + assert _start_request(config.transform_realtime_request(_append(), MODEL))["sample_rate"] == 24_000 + + +def test_audio_before_session_update_starts_a_default_stream_once_and_rejects_a_later_change(): + config = SonioxRealtimeConfig() + + first = config.transform_realtime_request(_append(), MODEL) + second = config.transform_realtime_request(_append(), MODEL) + with pytest.raises(SonioxProtocolError, match="once the stream has started"): + config.transform_realtime_request(_session_update({"turn_detection": None}), MODEL) + + assert _start_request(first) == { + "model": MODEL, + "audio_format": "pcm_s16le", + "sample_rate": 24_000, + "num_channels": 1, + "enable_endpoint_detection": True, + } + assert first[1:] == (PCM,) + assert second == (PCM,) + assert config.transform_realtime_request(_append(), MODEL) == (PCM,) + + +@pytest.mark.parametrize( + ("event_type", "control_frame"), + [("input_audio_buffer.commit", '{"type":"finalize"}'), ("input_audio_buffer.end", "")], +) +def test_buffer_control_events_map_to_soniox_control_frames(event_type: str, control_frame: str): + config = SonioxRealtimeConfig() + config.transform_realtime_request(_append(), MODEL) + + assert config.transform_realtime_request(_event(event_type), MODEL) == (control_frame,) + + +@pytest.mark.parametrize("event_type", ["response.create", "input_audio_buffer.clear"]) +def test_events_with_nothing_to_send_upstream_are_dropped_without_starting_the_stream(event_type: str): + config = SonioxRealtimeConfig() + + assert config.transform_realtime_request(_event(event_type), MODEL) == () + assert _start_request(config.transform_realtime_request(_append(), MODEL))["model"] == MODEL + + +def _events(config: SonioxRealtimeConfig, tokens: list[dict[str, object]]) -> tuple[OpenAIRealtimeEvents, ...]: + return tuple( + config.transform_realtime_response(_response(tokens), MODEL, MagicMock(), EMPTY_TRANSFORM_INPUT)["response"] + ) + + +def test_clear_discards_the_open_turn_and_everything_soniox_finalizes_for_it(): + config = SonioxRealtimeConfig() + config.transform_realtime_request(_append(), MODEL) + _events(config, [_token("Hel")]) + + assert config.transform_realtime_request(_event("input_audio_buffer.clear"), MODEL) == ('{"type":"finalize"}',) + cleared = _events(config, [_token("lo"), _token(""), _token("")]) + config.transform_realtime_request(_event("input_audio_buffer.commit"), MODEL) + kept = _events(config, [_token("Bye"), _token("")]) + + assert cleared == () + assert kept[-1]["transcript"] == "Bye" + + +def test_clear_after_a_commit_keeps_the_committed_turn_and_drops_only_audio_sent_after_it(): + config = SonioxRealtimeConfig() + config.transform_realtime_request(_append(), MODEL) + _events(config, [_token("One")]) + config.transform_realtime_request(_event("input_audio_buffer.commit"), MODEL) + config.transform_realtime_request(_event("input_audio_buffer.clear"), MODEL) + + committed = _events(config, [_token("")]) + cleared = _events(config, [_token(" two"), _token("")]) + later = _events(config, [_token("Three"), _token("")]) + + assert committed[-1]["transcript"] == "One" + assert cleared == () + assert later[-1]["transcript"] == "Three" + + +@pytest.mark.parametrize( + "audio_input", + [ + {"format": {"type": "audio/pcmu"}}, + {"format": {"type": "audio/pcm", "rate": 24_000, "channels": 2}}, + {"format": {"type": "audio/pcm", "rate": 0}}, + ], +) +def test_unsupported_audio_formats_are_rejected_without_starting_the_stream(audio_input: dict[str, object]): + config = SonioxRealtimeConfig() + + with pytest.raises(SonioxProtocolError): + config.transform_realtime_request(_session_update(audio_input), MODEL) + assert _start_request(config.transform_realtime_request(_append(), MODEL))["sample_rate"] == 24_000 + + +@pytest.mark.parametrize( + ("audio_input", "sample_rate"), + [({"format": {"type": "audio/pcm"}}, 24_000), ({"input_audio_format": "pcm16"}, 24_000)], +) +def test_pcm16_without_a_rate_streams_at_the_openai_default_rate(audio_input: dict[str, object], sample_rate: int): + config = SonioxRealtimeConfig() + + frames = config.transform_realtime_request(_session_update(audio_input), MODEL) + + assert _start_request(frames)["sample_rate"] == sample_rate + + +def test_validate_environment_sends_the_key_as_a_bearer_header(monkeypatch): + monkeypatch.setenv("SONIOX_API_KEY", "env-key") + config = SonioxRealtimeConfig() + + assert config.validate_environment({"x": "y"}, MODEL, "deployment-key") == { + "x": "y", + "Authorization": "Bearer deployment-key", + } + assert config.validate_environment({}, MODEL) == {"Authorization": "Bearer env-key"} + + +def test_validate_environment_rejects_a_missing_key(monkeypatch): + monkeypatch.delenv("SONIOX_API_KEY", raising=False) + + with pytest.raises(ValueError, match="SONIOX_API_KEY"): + SonioxRealtimeConfig().validate_environment({}, MODEL) + + +def test_get_complete_url_defaults_to_soniox_and_opens_a_websocket_for_http_bases(): + config = SonioxRealtimeConfig() + + assert config.get_complete_url(None, MODEL) == DEFAULT_SONIOX_REALTIME_URL + assert config.get_complete_url("wss://proxy.example/ws", MODEL) == "wss://proxy.example/ws" + assert config.get_complete_url("https://proxy.example/ws?region=eu", MODEL) == "wss://proxy.example/ws?region=eu" + assert config.get_complete_url("http://localhost:8080/ws", MODEL) == "ws://localhost:8080/ws" + with pytest.raises(ValueError, match="wss://"): + config.get_complete_url("ftp://proxy.example/ws", MODEL) + + +def test_endpoint_detection_emits_one_openai_turn_billed_by_processed_audio(): + transformer = SonioxEventTransformer(drops_originals=False) + + provisional = transformer.transform(_response([_token("Hel", is_final=False)], total_audio_proc_ms=500)) + finals = transformer.transform( + _response([_token("Hello"), _token(" world"), _token(" next", is_final=False)], total_audio_proc_ms=1500) + ) + endpoint = transformer.transform(_response([_token("")], total_audio_proc_ms=2000)) + + assert _types(provisional) == ["input_audio_buffer.speech_started"] + item_id = provisional[0]["item_id"] + assert [(event["type"], event.get("delta")) for event in finals] == [ + ("conversation.item.input_audio_transcription.delta", "Hello"), + ("conversation.item.input_audio_transcription.delta", " world"), + ] + assert _types(endpoint) == [ + "input_audio_buffer.speech_stopped", + "conversation.item.input_audio_transcription.completed", + ] + assert {event["item_id"] for event in (*finals, *endpoint)} == {item_id} + assert all("speaker" not in event and "language" not in event for event in finals) + assert endpoint[1]["transcript"] == "Hello world" + assert "segments" not in endpoint[1] + assert endpoint[1]["usage"] == {"type": "duration", "seconds": 2.0} + assert transformer.take_unbilled_usage() is None + + +def test_manual_finalization_completes_the_turn_without_speech_stopped_and_starts_a_new_item(): + transformer = SonioxEventTransformer(drops_originals=False) + + first = transformer.transform(_response([_token("One"), _token("")])) + second = transformer.transform(_response([_token("Two"), _token("")])) + + assert _types(first) == [ + "input_audio_buffer.speech_started", + "conversation.item.input_audio_transcription.delta", + "conversation.item.input_audio_transcription.completed", + ] + assert first[-1]["transcript"] == "One" + assert "usage" not in first[-1] + assert second[-1]["transcript"] == "Two" + assert second[0]["item_id"] != first[0]["item_id"] + + +def test_a_commit_without_speech_answers_with_an_empty_turn_billed_for_the_silence(): + config = SonioxRealtimeConfig() + config.transform_realtime_request(_append(), MODEL) + config.transform_realtime_request(_event("input_audio_buffer.commit"), MODEL) + + events = tuple( + config.transform_realtime_response( + _response([_token("")], total_audio_proc_ms=2000), MODEL, MagicMock(), EMPTY_TRANSFORM_INPUT + )["response"] + ) + + assert _types(events) == ["conversation.item.input_audio_transcription.completed"] + assert events[0]["transcript"] == "" + assert events[0]["usage"] == {"type": "duration", "seconds": 2.0} + assert _events(config, [_token("Next"), _token("")])[0]["item_id"] != events[0]["item_id"] + + +def test_a_commit_soniox_answers_after_its_endpoint_already_ended_the_turn_adds_no_second_turn(): + config = SonioxRealtimeConfig() + config.transform_realtime_request(_append(), MODEL) + _events(config, [_token("Hello")]) + config.transform_realtime_request(_event("input_audio_buffer.commit"), MODEL) + + events = _events(config, [_token(""), _token("")]) + + assert [event["transcript"] for event in events if "transcript" in event] == ["Hello"] + + +def test_a_server_vad_turn_answers_only_the_oldest_pending_commit(): + config = SonioxRealtimeConfig() + config.transform_realtime_request(_append(), MODEL) + _events(config, [_token("Hello")]) + config.transform_realtime_request(_event("input_audio_buffer.commit"), MODEL) + config.transform_realtime_request(_append(), MODEL) + config.transform_realtime_request(_event("input_audio_buffer.commit"), MODEL) + + events = _events(config, [_token(""), _token(""), _token("")]) + + assert [event["transcript"] for event in events if "transcript" in event] == ["Hello", ""] + + +def test_a_clear_without_speech_still_answers_nothing(): + config = SonioxRealtimeConfig() + config.transform_realtime_request(_append(), MODEL) + config.transform_realtime_request(_event("input_audio_buffer.clear"), MODEL) + + assert _events(config, [_token("")]) == () + + +def test_translation_sessions_transcribe_only_the_translated_tokens(): + transformer = SonioxEventTransformer(drops_originals=True) + + events = transformer.transform( + _response( + [ + _token("Hola", translation_status="original"), + _token("Hello", translation_status="translation"), + _token(""), + ] + ) + ) + + assert [event.get("delta") for event in events if "delta" in event] == ["Hello"] + assert events[-1]["transcript"] == "Hello" + + +def test_translation_sessions_keep_speech_soniox_leaves_untranslated(): + transformer = SonioxEventTransformer(drops_originals=True) + + events = transformer.transform( + _response( + [_token("Hola", translation_status="none"), _token(" amigos", translation_status="none"), _token("")] + ) + ) + + assert events[-1]["transcript"] == "Hola amigos" + + +def test_diarized_tokens_label_each_delta_and_split_the_completed_turn_into_speaker_segments(): + transformer = SonioxEventTransformer(drops_originals=False) + + events = transformer.transform( + _response( + [ + _token("Hi", speaker="1", language="en"), + _token(" doctor", speaker="1", language="en"), + _token(" Salam", speaker="2", language="fa"), + _token(" Hi", speaker="1", language="en"), + _token(""), + ] + ) + ) + + assert [(event["delta"], event["speaker"], event["language"]) for event in events if "delta" in event] == [ + ("Hi", "1", "en"), + (" doctor", "1", "en"), + (" Salam", "2", "fa"), + (" Hi", "1", "en"), + ] + assert events[-1]["transcript"] == "Hi doctor Salam Hi" + assert events[-1]["segments"] == ( + {"speaker": "1", "language": "en", "text": "Hi doctor"}, + {"speaker": "2", "language": "fa", "text": "Salam"}, + {"speaker": "1", "language": "en", "text": "Hi"}, + ) + + +def test_translation_sessions_announce_speech_on_the_first_spoken_token_before_any_translation(): + transformer = SonioxEventTransformer(drops_originals=True) + + events = transformer.transform(_response([_token("Hol", is_final=False, translation_status="original")])) + + assert _types(events) == ["input_audio_buffer.speech_started"] + + +def test_provider_errors_surface_as_openai_error_events(): + transformer = SonioxEventTransformer(drops_originals=False) + + events = transformer.transform(json.dumps({"error_code": 401, "error_message": "Invalid API key."})) + + assert events == ( + {"type": "error", "error": {"type": "server_error", "message": "Soniox realtime error 401: Invalid API key."}}, + ) + + +def test_malformed_provider_frames_are_rejected(): + with pytest.raises(SonioxProtocolError): + SonioxEventTransformer(drops_originals=False).transform('{"tokens": "nope"}') + + +def test_audio_processed_after_the_last_turn_is_billed_on_session_close(): + config = SonioxRealtimeConfig() + config.transform_realtime_response( + _response([_token("Hi"), _token("")], total_audio_proc_ms=1000), MODEL, MagicMock(), EMPTY_TRANSFORM_INPUT + ) + config.transform_realtime_response( + _response([], total_audio_proc_ms=3500), MODEL, MagicMock(), EMPTY_TRANSFORM_INPUT + ) + + assert config.unbilled_usage_on_session_close(MODEL) == {"type": "duration", "seconds": 2.5} + assert config.unbilled_usage_on_session_close(MODEL) is None + + +def test_translation_option_switches_the_config_to_translated_transcripts(): + config = SonioxRealtimeConfig( + options=SonioxRealtimeOptions(translation={"type": "one_way", "target_language": "en"}) + ) + + result = config.transform_realtime_response( + _response([_token("Hola", translation_status="original"), _token("Hello", translation_status="translation")]), + MODEL, + MagicMock(), + EMPTY_TRANSFORM_INPUT, + ) + + assert [event.get("delta") for event in result["response"] if "delta" in event] == ["Hello"] + + +def test_stream_time_since_the_start_request_is_billed_on_close_because_soniox_charges_the_whole_stream(): + now = [100.0] + config = SonioxRealtimeConfig(clock=lambda: now[0]) + config.wrap_backend(_RecordingBackend()) + now[0] = 110.0 + config.transform_realtime_request(_append(), MODEL) + config.transform_realtime_response( + _response([_token("Hi"), _token("")], total_audio_proc_ms=5000), MODEL, MagicMock(), EMPTY_TRANSFORM_INPUT + ) + now[0] = 130.5 + + assert config.unbilled_usage_on_session_close(MODEL) == {"type": "duration", "seconds": 15.5} + assert config.unbilled_usage_on_session_close(MODEL) is None + + +def _refusal_after(processed_ms: int, error_code: int = 401) -> SonioxRealtimeConfig: + now = [100.0] + config = SonioxRealtimeConfig(clock=lambda: now[0]) + config.transform_realtime_request(_append(), MODEL) + if processed_ms: + config.transform_realtime_response( + _response([_token("Hi")], total_audio_proc_ms=processed_ms), MODEL, MagicMock(), EMPTY_TRANSFORM_INPUT + ) + config.transform_realtime_response( + json.dumps({"error_code": error_code, "error_message": "Soniox error.", "total_audio_proc_ms": processed_ms}), + MODEL, + MagicMock(), + EMPTY_TRANSFORM_INPUT, + ) + now[0] = 103.0 + return config + + +@pytest.mark.parametrize("error_code", [400, 401, 402, 403, 408, 429, 500, 503]) +def test_an_error_before_any_audio_is_processed_is_a_refused_session_that_bills_nothing(error_code: int): + config = _refusal_after(processed_ms=0, error_code=error_code) + + refusal = config.session_refusal() + assert isinstance(refusal, SonioxException) + assert refusal.status_code == error_code + assert config.unbilled_usage_on_session_close(MODEL) is None + + +def test_an_error_after_audio_was_processed_keeps_the_session_a_billed_success(): + config = _refusal_after(processed_ms=1500) + + assert config.session_refusal() is None + assert config.unbilled_usage_on_session_close(MODEL) == {"type": "duration", "seconds": 3.0} + + +def test_a_connection_that_never_started_a_stream_bills_nothing(): + now = [100.0] + config = SonioxRealtimeConfig(clock=lambda: now[0]) + config.wrap_backend(_RecordingBackend()) + now[0] = 160.0 + + assert config.unbilled_usage_on_session_close(MODEL) is None + + +def test_soniox_opts_out_of_deferred_setup_so_client_frames_reach_it_in_order(): + assert SonioxRealtimeConfig().requires_session_configuration() + + +def test_a_session_update_matching_the_started_stream_is_a_no_op(): + config = SonioxRealtimeConfig() + config.transform_realtime_request(_append(), MODEL) + + assert ( + config.transform_realtime_request(_session_update({"format": {"type": "audio/pcm", "rate": 24_000}}), MODEL) + == () + ) + + +def _stream_started_at_16khz_in_french_without_vad() -> SonioxRealtimeConfig: + config = SonioxRealtimeConfig() + config.transform_realtime_request( + _session_update( + { + "format": {"type": "audio/pcm", "rate": 16_000}, + "transcription": {"language": "fr"}, + "turn_detection": None, + } + ), + MODEL, + ) + return config + + +@pytest.mark.parametrize( + "follow_up", + [ + {"transcription": {"language": "fr"}}, + {"format": {"type": "audio/pcm", "rate": 16_000}}, + {"format": {"type": "audio/pcm"}}, + {"turn_detection": None}, + ], +) +def test_a_partial_follow_up_update_that_restates_the_stream_is_a_no_op(follow_up: dict[str, object]): + config = _stream_started_at_16khz_in_french_without_vad() + + assert config.transform_realtime_request(_session_update(follow_up), MODEL) == () + + +def test_a_follow_up_update_with_an_unsupported_encoding_is_still_rejected(): + config = _stream_started_at_16khz_in_french_without_vad() + + with pytest.raises(SonioxProtocolError, match="requires pcm16"): + config.transform_realtime_request(_session_update({"format": {"type": "audio/pcmu"}}), MODEL) + + +@pytest.mark.parametrize( + "follow_up", + [ + {"transcription": {"language": "de"}}, + {"format": {"type": "audio/pcm", "rate": 24_000}}, + {"turn_detection": {"type": "server_vad"}}, + ], +) +def test_a_partial_follow_up_update_that_changes_a_sent_setting_is_rejected(follow_up: dict[str, object]): + config = _stream_started_at_16khz_in_french_without_vad() + + with pytest.raises(SonioxProtocolError, match="once the stream has started"): + config.transform_realtime_request(_session_update(follow_up), MODEL) + + +def test_session_created_event_advertises_the_deployment_model_for_cost_tracking(): + config = SonioxRealtimeConfig(options=SonioxRealtimeOptions(language_hints=("de",))) + + event = config.transform_session_created_event(MODEL, "session_1") + + assert event["session"]["id"] == "session_1" + assert event["session"]["audio"]["input"]["transcription"] == {"model": MODEL, "language": "de"} + + +def test_backend_session_start_is_announced_with_the_negotiated_session(): + config = SonioxRealtimeConfig() + config.transform_session_created_event(MODEL, "session_1") + config.transform_realtime_request( + _session_update( + { + "format": {"type": "audio/pcm", "rate": 16_000}, + "transcription": {"language": "it"}, + "turn_detection": None, + } + ), + MODEL, + ) + + result = config.transform_realtime_response(SESSION_STARTED_FRAME, MODEL, MagicMock(), EMPTY_TRANSFORM_INPUT) + + assert _types(tuple(result["response"])) == ["session.created"] + assert result["response"][0]["session"]["id"] == "session_1" + assert result["response"][0]["session"]["audio"]["input"] == { + "format": {"type": "audio/pcm", "rate": 16_000}, + "transcription": {"model": MODEL, "language": "it"}, + "turn_detection": None, + } + + +def test_provider_config_manager_passes_only_soniox_options_from_deployment_params(): + config = ProviderConfigManager.get_provider_realtime_config( + MODEL, + LlmProviders.SONIOX, + litellm_params={ + "api_key": "secret", + "model": "soniox/stt-rt-v5", + "audio_format": "mulaw", + "language_hints": ["pl"], + "context": {"terms": ["LiteLLM"]}, + "max_endpoint_delay_ms": 800, + }, + ) + + assert isinstance(config, SonioxRealtimeConfig) + assert _start_request(config.transform_realtime_request(_append(), MODEL)) == { + "model": MODEL, + "audio_format": "pcm_s16le", + "sample_rate": 24_000, + "num_channels": 1, + "enable_endpoint_detection": True, + "language_hints": ["pl"], + "context": {"terms": ["LiteLLM"]}, + "max_endpoint_delay_ms": 800, + } + + +SEMANTIC_VAD: Final = {"turn_detection": {"type": "semantic_vad", "eagerness": "high"}} + + +def _soniox_config(drop_params: bool) -> SonioxRealtimeConfig: + config = ProviderConfigManager.get_provider_realtime_config(MODEL, LlmProviders.SONIOX, drop_params=drop_params) + assert isinstance(config, SonioxRealtimeConfig) + return config + + +def test_unsupported_turn_detection_is_rejected_without_drop_params(): + config = _soniox_config(drop_params=False) + + with pytest.raises(SonioxProtocolError, match="Set drop_params to use server_vad instead"): + config.transform_realtime_request(_session_update(SEMANTIC_VAD), MODEL) + + +def test_unsupported_turn_detection_falls_back_to_endpoint_detection_with_drop_params(): + config = _soniox_config(drop_params=True) + + frames = config.transform_realtime_request(_session_update(SEMANTIC_VAD), MODEL) + + assert _start_request(frames)["enable_endpoint_detection"] is True + assert config.transform_realtime_request(_session_update(SEMANTIC_VAD), MODEL) == () + + +def test_a_follow_up_with_unsupported_turn_detection_is_rejected_without_drop_params(): + config = _soniox_config(drop_params=False) + config.transform_realtime_request(_append(), MODEL) + + with pytest.raises(SonioxProtocolError, match="Set drop_params to use server_vad instead"): + config.transform_realtime_request(_session_update(SEMANTIC_VAD), MODEL) + + +class _RecordingBackend: + def __init__(self) -> None: + self.sent: Final[list[str | bytes]] = [] + self.upstream: Final[asyncio.Queue[str | Exception]] = asyncio.Queue() + self.entered = False + self.exited = False + self.closed = False + + async def __aenter__(self) -> "_RecordingBackend": + self.entered = True + return self + + async def __aexit__(self, exc_type, exc_value, traceback) -> None: + self.exited = True + + async def send(self, message: str | bytes) -> None: + self.sent.append(message) + + async def recv(self, decode: bool | None = None) -> str | bytes: + frame = await self.upstream.get() + if isinstance(frame, Exception): + raise frame + return frame + + async def close(self) -> None: + self.closed = True + + +class _BrokenBackend(_RecordingBackend): + async def send(self, message: str | bytes) -> None: + raise ConnectionError("upstream gone") + + +@pytest.mark.asyncio +async def test_keepalive_is_sent_only_while_the_client_is_idle_and_stops_after_end_of_stream(): + now = [0.0] + ticks: asyncio.Queue[float] = asyncio.Queue() + + async def sleep(_: float) -> None: + now[0] = await ticks.get() + + async def tick(at: float) -> None: + await ticks.put(at) + for _ in range(5): + await asyncio.sleep(0) + + inner = _RecordingBackend() + backend = SonioxRealtimeBackend(inner, keepalive_interval=5.0, clock=lambda: now[0], sleep=sleep) + + async with backend: + await tick(5.0) + assert inner.sent == ['{"type":"keepalive"}'] + now[0] = 7.0 + await backend.send(PCM) + await tick(10.0) + assert inner.sent == ['{"type":"keepalive"}', PCM] + await backend.send("") + await tick(30.0) + assert inner.sent == ['{"type":"keepalive"}', PCM, ""] + + assert inner.entered and inner.exited + + +@pytest.mark.asyncio +async def test_backend_announces_the_session_once_the_start_request_is_sent_while_relaying_upstream_frames(): + inner = _RecordingBackend() + backend = SonioxRealtimeBackend(inner) + + async with backend: + pending = asyncio.create_task(backend.recv()) + await backend.send('{"type":"keepalive"}') + await asyncio.sleep(0) + assert not pending.done() + await backend.send('{"model":"stt-rt-v5"}') + assert await pending == SESSION_STARTED_FRAME + await backend.send('{"model":"ignored-second-config"}') + await inner.upstream.put('{"tokens":[]}') + assert await backend.recv() == '{"tokens":[]}' + + +@pytest.mark.asyncio +async def test_backend_relays_upstream_frames_that_arrive_before_the_session_starts(): + inner = _RecordingBackend() + + async with SonioxRealtimeBackend(inner) as backend: + await inner.upstream.put('{"error_code":401,"error_message":"Invalid API key."}') + + assert await backend.recv() == '{"error_code":401,"error_message":"Invalid API key."}' + + +@pytest.mark.asyncio +async def test_a_cancelled_read_loses_no_frames_and_upstream_errors_reach_the_reader(): + inner = _RecordingBackend() + + async with SonioxRealtimeBackend(inner) as backend: + cancelled = asyncio.create_task(backend.recv()) + await asyncio.sleep(0) + cancelled.cancel() + await asyncio.gather(cancelled, return_exceptions=True) + await inner.upstream.put('{"tokens":[]}') + await inner.upstream.put(ConnectionError("upstream closed")) + + assert await backend.recv() == '{"tokens":[]}' + with pytest.raises(ConnectionError, match="upstream closed"): + await backend.recv() + + +@pytest.mark.asyncio +async def test_closing_the_backend_cancels_pending_reads_and_closes_upstream(): + inner = _RecordingBackend() + backend = SonioxRealtimeBackend(inner) + await backend.__aenter__() + pending = asyncio.create_task(backend.recv()) + await backend.send('{"model":"stt-rt-v5"}') + assert await pending == SESSION_STARTED_FRAME + + await backend.close() + await inner.upstream.put('{"tokens":[]}') + await asyncio.sleep(0) + + assert inner.closed + assert inner.upstream.qsize() == 1 + + +@pytest.mark.asyncio +async def test_keepalive_stops_after_the_upstream_rejects_a_send(): + sleeps: list[float] = [] + + async def sleep(seconds: float) -> None: + sleeps.append(seconds) + await asyncio.sleep(0) + + async with SonioxRealtimeBackend(_BrokenBackend(), keepalive_interval=0.0, sleep=sleep): + for _ in range(5): + await asyncio.sleep(0) + + assert sleeps == [0.0] diff --git a/tests/unit/llms/test_get_complete_url_trailing_slash.py b/tests/unit/llms/test_get_complete_url_trailing_slash.py index cd778c20d9b..43a74e56028 100644 --- a/tests/unit/llms/test_get_complete_url_trailing_slash.py +++ b/tests/unit/llms/test_get_complete_url_trailing_slash.py @@ -143,6 +143,7 @@ API_BASE_IS_THE_FULL_URL: Final[Mapping[str, str]] = MappingProxyType( "duckduckgo.search.transformation.DuckDuckGoSearchConfig", "google_pse.search.transformation.GooglePSESearchConfig", "searchapi.search.transformation.SearchAPIConfig", + "soniox.realtime.transformation.SonioxRealtimeConfig", "tinyfish.search.transformation.TinyfishSearchConfig", "vertex_ai.text_to_speech.transformation.VertexAITextToSpeechConfig", "vertex_ai.vertex_ai_partner_models.anthropic.messages.transformation." diff --git a/tests/unit/realtime_api/test_main.py b/tests/unit/realtime_api/test_main.py index 5478b807691..46777b32319 100644 --- a/tests/unit/realtime_api/test_main.py +++ b/tests/unit/realtime_api/test_main.py @@ -618,7 +618,11 @@ async def test_bearer_health_check_sends_the_secret_behind_an_oidc_credential( monkeypatch.setattr( litellm, "credential_list", - [CredentialItem(credential_name="voice-sa", credential_values={"api_key": f"oidc/file/{token}"}, credential_info={})], + [ + CredentialItem( + credential_name="voice-sa", credential_values={"api_key": f"oidc/file/{token}"}, credential_info={} + ) + ], ) connect: Final = _CapturingConnect() with patch("websockets.connect", connect): @@ -1447,3 +1451,98 @@ async def test_openai_realtime_health_check_sends_the_federated_bearer_for_a_key model_params={"model": "openai/gpt-realtime", **federated_deployment}, ) assert connect.kwargs["additional_headers"] == {"Authorization": "Bearer federated-bearer"} + + +@pytest.mark.asyncio +async def test_arealtime_routes_soniox_with_the_deployment_options(monkeypatch): + from litellm.llms.soniox.realtime.transformation import SonioxRealtimeConfig + + captured: dict[str, object] = {} + + async def mock_async_realtime(**kwargs): + captured.update(kwargs) + + monkeypatch.setattr(realtime_main.base_llm_http_handler, "async_realtime", mock_async_realtime) + await realtime_main._arealtime.__wrapped__( + model="soniox/stt-rt-v5", + websocket=MagicMock(), + litellm_logging_obj=FakeLogging(), + query_params={"model": "soniox/stt-rt-v5", "intent": "transcription"}, + api_key="deployment-key", + language_hints=["de"], + ) + + provider_config = captured["provider_config"] + assert isinstance(provider_config, SonioxRealtimeConfig) + assert captured["model"] == "stt-rt-v5" + assert captured["api_key"] == "deployment-key" + start_request = provider_config.transform_realtime_request( + json.dumps({"type": "input_audio_buffer.commit"}), "stt-rt-v5" + )[0] + assert json.loads(start_request)["language_hints"] == ["de"] + + +async def _serve_soniox(reply: str | None) -> tuple[object, list[object]]: + import websockets + + received: list[object] = [] + + async def handler(connection) -> None: + received.append(connection.request.headers["Authorization"]) + received.append(json.loads(await connection.recv())) + received.append(await connection.recv()) + received.append(await connection.recv()) + if reply is not None: + await connection.send(reply) + + return await websockets.serve(handler, "127.0.0.1", 0), received + + +@pytest.mark.asyncio +async def test_soniox_realtime_health_check_starts_and_ends_a_stream_on_the_deployment_websocket(): + server, received = await _serve_soniox(json.dumps({"tokens": [], "finished": True})) + port = next(iter(server.sockets)).getsockname()[1] + async with server: + healthy = await realtime_main.realtime_health_check( + model="stt-rt-v5", + custom_llm_provider="soniox", + api_key="soniox-key", + api_base=f"http://127.0.0.1:{port}", + model_params={"model": "soniox/stt-rt-v5", "language_hints": ["fa", "en"]}, + ) + + assert healthy is True + assert received[0] == "Bearer soniox-key" + assert received[1]["model"] == "stt-rt-v5" + assert received[1]["language_hints"] == ["fa", "en"] + assert received[2] == bytes(len(received[2])) + assert len(received[2]) > 0 + assert received[3] == "" + + +@pytest.mark.asyncio +async def test_soniox_realtime_health_check_fails_when_soniox_refuses_the_stream(): + from litellm.llms.soniox.common_utils import SonioxException + + server, _ = await _serve_soniox(json.dumps({"error_code": 401, "error_message": "Invalid API key."})) + port = next(iter(server.sockets)).getsockname()[1] + async with server: + with pytest.raises(SonioxException) as refused: + await realtime_main.realtime_health_check( + model="stt-rt-v5", custom_llm_provider="soniox", api_key="bad-key", api_base=f"ws://127.0.0.1:{port}" + ) + + assert refused.value.status_code == 401 + + +@pytest.mark.asyncio +async def test_soniox_realtime_health_check_fails_when_soniox_closes_without_finishing_the_stream(): + from litellm.llms.soniox.common_utils import SonioxException + + server, _ = await _serve_soniox(None) + port = next(iter(server.sockets)).getsockname()[1] + async with server: + with pytest.raises(SonioxException): + await realtime_main.realtime_health_check( + model="stt-rt-v5", custom_llm_provider="soniox", api_key="soniox-key", api_base=f"ws://127.0.0.1:{port}" + ) From c13ecb308522841ba7485f571c7172056f19e8ca Mon Sep 17 00:00:00 2001 From: Otis Cui Date: Thu, 29 Oct 2026 23:05:47 +0800 Subject: [PATCH 04/13] fix(exceptions)!: honor status code when mapping OpenAI invalid_request_error exceptions (#32537) * fix: openai exception mapper ignored status code for invalid_request_error _map_openai_exception coerced anything whose body carried type invalid_request_error into BadRequestError, guarded only by excluding the single phrase Incorrect API key provided, and this ran before the status-code checks. OpenAI and openai-compatible providers return 401 and 404 bodies with type invalid_request_error, so an auth failure like Invalid API key mapped to a 400 instead of AuthenticationError, and a 404 mapped to BadRequestError instead of NotFoundError. That changes the status returned to proxy clients and router retry, cooldown, and fallback handling. Mirror the Azure mapper by only coercing to BadRequestError when status_code is None or 400, so 401 and 404 fall through to their existing handlers. * fix(exceptions): map OpenAI-family invalid_request_error 404s to NotFoundError Move the contributor's test to the mapped unit test file, assert the raised class, status code and message, and cover the 400 and status-less cases that must still raise BadRequestError Add a regression test for the flow the bug was found on. Retrieving a missing OpenAI vector store, sync or async, now raises NotFoundError with status 404 The router test that pinned an OpenAI 404 typed invalid_request_error as BadRequestError now expects NotFoundError * fix(a2a): answer a provider NotFoundError as JSON-RPC invalid params * fix(exceptions)!: limit the invalid_request_error status guard to 401 and 404 Every other status keeps raising BadRequestError as on main, so 402-415, 429 and 5xx bodies typed invalid_request_error keep their class and router retry behavior. A2A goes back to main's mapping, which answers NotFoundError and AuthenticationError as JSON-RPC internal errors. * fix(exceptions): keep the provider error fields on mapped 401 and 404 invalid_request_errors Raise AuthenticationError and NotFoundError straight from the invalid_request_error branch with the provider body, so code, param and type survive and the proxy's error.type matches main. The Rust mapper copy skips 401 and 404 the same way. --------- Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> --- .../src/exception_mapping_utils/openai.rs | 18 +++ litellm/exceptions.py | 6 +- .../exception_mapping_utils.py | 19 +++ .../test_exception_mapping_utils.py | 133 ++++++++++++++++++ tests/unit/test_router/test_router.py | 2 +- tests/unit/vector_stores/test_main.py | 45 +++++- 6 files changed, 219 insertions(+), 4 deletions(-) diff --git a/litellm-rust/crates/core-utils/src/exception_mapping_utils/openai.rs b/litellm-rust/crates/core-utils/src/exception_mapping_utils/openai.rs index d45078c415e..1273fd4d7f1 100644 --- a/litellm-rust/crates/core-utils/src/exception_mapping_utils/openai.rs +++ b/litellm-rust/crates/core-utils/src/exception_mapping_utils/openai.rs @@ -53,6 +53,7 @@ pub(super) const RULES: &[Rule] = &[ |mapping| { mapping.error_str.contains("invalid_request_error") && !mapping.error_str.contains("Incorrect API key provided") + && !matches!(mapping.status, Some(401 | 404)) }, PublicError::BadRequest, ), @@ -132,6 +133,23 @@ mod tests { assert_eq!(classified(Some(400), text), Some(expected)); } + #[rstest::rstest] + #[case::unauthorized(Some(401), None)] + #[case::not_found(Some(404), None)] + #[case::forbidden(Some(403), Some(PublicError::BadRequest))] + #[case::rate_limited(Some(429), Some(PublicError::BadRequest))] + #[case::server_error(Some(500), Some(PublicError::BadRequest))] + #[case::no_status(None, Some(PublicError::BadRequest))] + fn invalid_request_errors_leave_401_and_404_to_the_status_mapping( + #[case] status: Option, + #[case] expected: Option, + ) { + assert_eq!( + classified(status, "invalid_request_error bad field"), + expected + ); + } + #[rstest::rstest] #[case::rate_limit_before_context_window( "rate limit and This model's maximum context length is 10", diff --git a/litellm/exceptions.py b/litellm/exceptions.py index cd5139f5f79..61e3a8fb5dc 100644 --- a/litellm/exceptions.py +++ b/litellm/exceptions.py @@ -139,6 +139,7 @@ class AuthenticationError(openai.AuthenticationError): litellm_debug_info: str | None = None, max_retries: int | None = None, num_retries: int | None = None, + body: Mapping[str, object] | None = None, ): self.status_code = 401 self.message = f"litellm.AuthenticationError: {message}" @@ -152,7 +153,7 @@ class AuthenticationError(openai.AuthenticationError): request=httpx.Request(method="GET", url="https://litellm.ai"), # mock request object ) super().__init__( - self.message, response=self.response, body=None + self.message, response=self.response, body=body ) # Call the base class constructor with the parameters it needs def __str__(self): @@ -189,6 +190,7 @@ class NotFoundError(openai.NotFoundError): litellm_debug_info: str | None = None, max_retries: int | None = None, num_retries: int | None = None, + body: Mapping[str, object] | None = None, ): self.status_code = 404 self.message = f"litellm.NotFoundError: {message}" @@ -202,7 +204,7 @@ class NotFoundError(openai.NotFoundError): request=httpx.Request(method="GET", url="https://litellm.ai"), # mock request object ) super().__init__( - self.message, response=self.response, body=None + self.message, response=self.response, body=body ) # Call the base class constructor with the parameters it needs def __str__(self): diff --git a/litellm/litellm_core_utils/exception_mapping_utils.py b/litellm/litellm_core_utils/exception_mapping_utils.py index 703550dd2b5..64a95eb5f28 100644 --- a/litellm/litellm_core_utils/exception_mapping_utils.py +++ b/litellm/litellm_core_utils/exception_mapping_utils.py @@ -413,6 +413,25 @@ def _map_openai_exception( body=getattr(original_exception, "body", None), ) elif "invalid_request_error" in error_str and "Incorrect API key provided" not in error_str: + invalid_request_status: Final = getattr(original_exception, "status_code", None) + if invalid_request_status == 401: + raise AuthenticationError( + message=f"{exception_provider} - {message}", + llm_provider=custom_llm_provider, + model=model, + response=response, + litellm_debug_info=extra_information, + body=getattr(original_exception, "body", None), + ) + if invalid_request_status == 404: + raise NotFoundError( + message=f"{exception_provider} - {message}", + model=model, + llm_provider=custom_llm_provider, + response=response, + litellm_debug_info=extra_information, + body=getattr(original_exception, "body", None), + ) raise BadRequestError( message=f"{exception_provider} - {message}", llm_provider=custom_llm_provider, diff --git a/tests/unit/litellm_core_utils/test_exception_mapping_utils.py b/tests/unit/litellm_core_utils/test_exception_mapping_utils.py index 55e7c47e2c9..19e5387da9a 100644 --- a/tests/unit/litellm_core_utils/test_exception_mapping_utils.py +++ b/tests/unit/litellm_core_utils/test_exception_mapping_utils.py @@ -24,6 +24,7 @@ from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.llms.bedrock.common_utils import BedrockError from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler from litellm.llms.openai.common_utils import OpenAIError +from litellm.proxy.common_utils.openai_error_payload import openai_error_type from litellm.types.utils import LlmProviders @@ -1015,6 +1016,138 @@ def test_azure_404_with_invalid_request_error_type_maps_to_not_found(): assert "Response with id 'resp_abc' not found." in excinfo.value.message +@pytest.mark.parametrize("provider", ["openai", "groq", "deepseek"]) +@pytest.mark.parametrize( + "status_code, message, expected_exception", + [ + (400, "Invalid 'temperature'", litellm.BadRequestError), + (401, "Invalid API key", litellm.AuthenticationError), + (404, "The model does not exist", litellm.NotFoundError), + ], +) +def test_openai_invalid_request_error_respects_status_code( + provider: str, status_code: int, message: str, expected_exception: type[openai.APIStatusError] +) -> None: + error_message: Final = json.dumps({"error": {"message": message, "type": "invalid_request_error"}}) + original_exception: Final = OpenAIError(status_code=status_code, message=error_message, headers={}) + + with pytest.raises(openai.APIError) as excinfo: + exception_type( + model="some-model", + original_exception=original_exception, + custom_llm_provider=provider, + ) + + assert type(excinfo.value) is expected_exception + assert excinfo.value.status_code == status_code + assert message in excinfo.value.message + + +@pytest.mark.parametrize("provider", ["openai", "groq", "deepseek"]) +def test_openai_invalid_request_error_without_a_status_code_maps_to_bad_request(provider: str) -> None: + message: Final = "Missing required parameter: 'input'." + original_exception: Final = Exception(json.dumps({"error": {"message": message, "type": "invalid_request_error"}})) + + with pytest.raises(openai.APIError) as excinfo: + exception_type( + model="some-model", + original_exception=original_exception, + custom_llm_provider=provider, + ) + + assert type(excinfo.value) is litellm.BadRequestError + assert message in excinfo.value.message + + +@pytest.mark.parametrize("provider", ["openai", "groq", "deepseek"]) +@pytest.mark.parametrize("status_code", [403, 409, 422, 429, 500, 503]) +def test_openai_invalid_request_error_keeps_bad_request_for_statuses_other_than_401_and_404( + provider: str, status_code: int +) -> None: + message: Final = "Cannot cancel a batch that is already completed." + error_message: Final = json.dumps({"error": {"message": message, "type": "invalid_request_error"}}) + original_exception: Final = OpenAIError(status_code=status_code, message=error_message, headers={}) + + with pytest.raises(openai.APIError) as excinfo: + exception_type( + model="some-model", + original_exception=original_exception, + custom_llm_provider=provider, + ) + + assert type(excinfo.value) is litellm.BadRequestError + assert message in excinfo.value.message + + +def _openai_sdk_error( + sdk_error: type[openai.APIStatusError], status_code: int, provider_error: dict[str, str | None] +) -> openai.APIStatusError: + request: Final = httpx.Request("POST", "https://api.openai.com/v1/chat/completions") + return sdk_error( + f"Error code: {status_code} - {{'error': {provider_error}}}", + response=httpx.Response(status_code, request=request), + body=provider_error, + ) + + +@pytest.mark.parametrize( + "sdk_error, status_code, expected_exception, param, code", + [ + (openai.AuthenticationError, 401, litellm.AuthenticationError, "organization", "mismatched_organization"), + (openai.NotFoundError, 404, litellm.NotFoundError, "id", None), + ], +) +def test_openai_invalid_request_error_401_and_404_keep_the_provider_error_fields( + sdk_error: type[openai.APIStatusError], + status_code: int, + expected_exception: type[openai.APIStatusError], + param: str, + code: str | None, +) -> None: + provider_error: Final = { + "message": "The request was rejected.", + "type": "invalid_request_error", + "param": param, + "code": code, + } + + with pytest.raises(openai.APIError) as excinfo: + exception_type( + model="gpt-5-mini", + original_exception=_openai_sdk_error(sdk_error, status_code, provider_error), + custom_llm_provider="openai", + ) + + assert type(excinfo.value) is expected_exception + assert (excinfo.value.status_code, excinfo.value.type, excinfo.value.param, excinfo.value.code) == ( + status_code, + "invalid_request_error", + param, + code, + ) + assert excinfo.value.body == provider_error + assert openai_error_type(excinfo.value, status_code) == "invalid_request_error" + assert excinfo.value.message.count(expected_exception.__name__) == 1 + + +def test_openai_incorrect_api_key_keeps_the_proxy_authentication_error_type() -> None: + provider_error: Final = { + "message": "Incorrect API key provided: sk-litellmqa.", + "type": "invalid_request_error", + "param": None, + "code": "invalid_api_key", + } + + with pytest.raises(litellm.AuthenticationError) as excinfo: + exception_type( + model="gpt-5-mini", + original_exception=_openai_sdk_error(openai.AuthenticationError, 401, provider_error), + custom_llm_provider="openai", + ) + + assert openai_error_type(excinfo.value, 401) == "authentication_error" + + class _UpstreamHTTPError(Exception): def __init__(self, status_code: int) -> None: super().__init__("upstream failure") diff --git a/tests/unit/test_router/test_router.py b/tests/unit/test_router/test_router.py index b55b3c5efc2..ee3a70c7c1c 100644 --- a/tests/unit/test_router/test_router.py +++ b/tests/unit/test_router/test_router.py @@ -21632,7 +21632,7 @@ async def test_router_retry_policy_controls_upstream_attempt_count( "error_body,error_type", [ ({"message": "model is down", "type": "server_error"}, litellm.NotFoundError), - ({"message": "Response with id 'resp_x' not found.", "type": "invalid_request_error"}, litellm.BadRequestError), + ({"message": "Response with id 'resp_x' not found.", "type": "invalid_request_error"}, litellm.NotFoundError), ], ) @pytest.mark.parametrize( diff --git a/tests/unit/vector_stores/test_main.py b/tests/unit/vector_stores/test_main.py index 25a378270b1..21b1eee3a3e 100644 --- a/tests/unit/vector_stores/test_main.py +++ b/tests/unit/vector_stores/test_main.py @@ -15,12 +15,13 @@ from unittest.mock import MagicMock, patch import httpx import pytest +import litellm import litellm.vector_stores.main as vector_stores_main from litellm.llms.base_llm.vector_store.transformation import ( RouterVectorStoreEmbeddingExecutor, ) from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler -from litellm.vector_stores.main import acreate, aupdate, search +from litellm.vector_stores.main import acreate, aretrieve, aupdate, retrieve, search MOCK_SEARCH_RESPONSE = { "object": "vector_store.search_results.page", @@ -200,3 +201,45 @@ async def test_direct_sdk_vector_store_writes_send_the_callers_metadata_unchange sent_metadata: Final = await _metadata_sent_to_provider(operation, callers_metadata) assert sent_metadata == [callers_metadata] + + +def _vector_store_not_found(request: httpx.Request) -> httpx.Response: + return httpx.Response( + 404, + json={ + "error": { + "message": "No vector store found with id 'vs_missing'.", + "type": "invalid_request_error", + "param": "vector_store_id", + "code": None, + } + }, + request=request, + ) + + +async def _retrieve_missing_openai_vector_store(is_async: bool) -> object: + transport: Final = httpx.MockTransport(_vector_store_not_found) + if is_async: + return await aretrieve( + vector_store_id="vs_missing", + custom_llm_provider="openai", + api_key="sk-test", + client=AsyncHTTPHandler(transport=transport), + ) + return retrieve( + vector_store_id="vs_missing", + custom_llm_provider="openai", + api_key="sk-test", + client=HTTPHandler(client=httpx.Client(transport=transport)), + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("is_async", [True, False]) +async def test_retrieving_a_missing_openai_vector_store_raises_not_found(is_async: bool): + with pytest.raises(litellm.NotFoundError) as caught: + await _retrieve_missing_openai_vector_store(is_async) + + assert caught.value.status_code == 404 + assert "No vector store found with id 'vs_missing'." in caught.value.message From be6d082bdcbc2d83a6522c0e015077e007a9bddc Mon Sep 17 00:00:00 2001 From: Mateo Wang <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 29 Oct 2026 08:06:12 -0700 Subject: [PATCH 05/13] fix(anthropic)!: send /v1/messages PDF documents to OpenAI-shaped providers as file parts (#47166) * fix(anthropic): send /v1/messages PDF documents to OpenAI-shaped providers as file parts * refactor(anthropic): hold the chat completions document change for a later commit * fix(anthropic): send /v1/messages PDFs as file parts and keep file parts on azure_ai * fix(snowflake): turn file parts into document blocks for Cortex Claude * fix(prompt_security): sanitize chat file parts that carry file_data * fix: answer 400 or drop PDFs for models that cannot read them Models whose model_info or cost map entry sets supports_pdf_input to false, and Ollama, now get the PDF dropped under drop_params and a 400 from litellm without it. GigaChat uploads file parts, SAP keeps sending them as image_url data URLs, azure_ai passes dict content through, and annotations other than url_citation move to provider_specific_fields instead of failing the response. * fix(router): scope supports_pdf_input to its deployment and keep provider annotation keys Mark the azure_ai chat models in the cost map that cannot read PDFs, so drop_params drops the PDF and litellm answers its own 400 without it. Keep url_citation annotations whole, move only typed provider annotations to provider_specific_fields, and split them the same way when a stream is rebuilt. * fix(vertex_ai): send PDF file parts to Vertex partner models as image_url Vertex AI's OpenAI-compatible partner endpoint rejects the chat `file` part, so PDFs from /v1/messages reach those models as the image_url data URL they took before. Mark native DeepSeek chat models as unable to read PDFs, and treat a file part as a PDF by its format or filename when it has no data URL. * fix(cost-map): leave native DeepSeek PDF support unset /v1/messages reaches native DeepSeek through its own Anthropic endpoint, so the marks never applied there and would only change chat completions, which this PR does not otherwise touch. * fix(main): drop a user turn that held only a PDF under drop_params Dropping the PDF from a PDF-only user turn left an empty content list, which Foundry and OpenAI-shaped providers reject with a 400, so drop_params still failed for that request. The emptied user turn is now removed. * fix: read OCI PDFs, gate deployment base models, and stop empty drops OCI generic chat sends file parts as image data URLs again, so OCI Gemini keeps reading PDFs from /v1/messages and chat file parts. The PDF gate now reads model_info.base_model, drops PDFs under drop_params for azure_ai deployments that never declared PDF support (the old flattening behavior), and answers 400 when dropping the PDF leaves no user message. The cost map marks the Vertex partner text models that a live probe showed answering without reading the PDF. * test(integration): cover the PDF input gate, messages PDF file parts and provider annotations Wire cells for the supports_pdf_input check on chat, chat streaming, /v1/messages and the Responses bridge, its model_info, base_model, cost map and Ollama sources, drop_params at deployment and request level, azure_ai file parts, sibling deployment isolation, Anthropic SDK PDFs reaching azure as titled file parts, Gemini and Bedrock document controls, Prompt Security sanitizing a /v1/messages PDF, and OpenRouter file annotations live, cached and streamed * test(integration): cover Snowflake Claude and SAP in the messages PDF wire cells * fix(azure_ai): flatten non-PDF file parts to text as before and cover the GigaChat PDF upload * fix(main): PDF check reads the deployment's top-level base_model * fix(main): keep roles alternating after drop_params removes a PDF-only turn When ensure_alternating_roles is set, re-run role normalization on the messages left after the PDF check drops a turn, so two assistant turns never end up next to each other * test: type the PDF and annotation tests the PR added * test: type the GigaChat, SAP and Snowflake PDF tests --- .../convert_dict_to_response.py | 33 +- .../prompt_templates/common_utils.py | 103 +++++ .../streaming_chunk_builder_utils.py | 17 +- .../messages/adapters/transformation.py | 31 +- .../responses_adapters/transformation.py | 26 +- litellm/llms/azure_ai/chat/transformation.py | 24 +- litellm/llms/gigachat/chat/transformation.py | 15 +- litellm/llms/oci/chat/transformation.py | 5 +- litellm/llms/sap/chat/transformation.py | 3 +- litellm/llms/snowflake/chat/transformation.py | 45 +- .../vertex_ai_partner_models/main.py | 3 +- litellm/main.py | 163 ++++++- ...odel_prices_and_context_window_backup.json | 121 ++++- .../prompt_security/prompt_security.py | 5 +- litellm/types/utils.py | 7 +- model_prices_and_context_window.json | 121 ++++- ...st_anthropic_messages_pdf_document_wire.py | 421 ++++++++++++++++++ .../test_chat_provider_annotations_wire.py | 190 ++++++++ .../providers/test_pdf_input_support_wire.py | 390 ++++++++++++++++ .../test_convert_dict_to_response.py | 51 +++ ...hropic_messages_adapters_transformation.py | 34 ++ .../chat/test_azure_ai_transformation.py | 52 +++ .../chat/test_gigachat_chat_transformation.py | 15 + .../oci/chat/test_oci_chat_transformation.py | 30 ++ .../llms/sap/chat/test_sap_transformation.py | 18 + .../test_snowflake_chat_transformation.py | 30 ++ .../test_vertex_ai_gemma_global_endpoint.py | 50 +++ .../test_prompt_security_guardrails.py | 53 +++ tests/unit/test_main.py | 305 ++++++++++++- .../unit/test_router_model_cost_isolation.py | 57 +++ .../test_stream_chunk_builder_annotations.py | 45 ++ 31 files changed, 2374 insertions(+), 89 deletions(-) create mode 100644 tests/integration/messages_endpoint/chat_bridge/test_anthropic_messages_pdf_document_wire.py create mode 100644 tests/integration/providers/test_chat_provider_annotations_wire.py create mode 100644 tests/integration/providers/test_pdf_input_support_wire.py diff --git a/litellm/litellm_core_utils/llm_response_utils/convert_dict_to_response.py b/litellm/litellm_core_utils/llm_response_utils/convert_dict_to_response.py index 8b94c34942c..de911beabdd 100644 --- a/litellm/litellm_core_utils/llm_response_utils/convert_dict_to_response.py +++ b/litellm/litellm_core_utils/llm_response_utils/convert_dict_to_response.py @@ -19,6 +19,7 @@ from litellm.litellm_core_utils.prompt_templates.common_utils import ( from litellm.llms.ollama.common_utils import OllamaResponseIdNormalizer from litellm.types.llms.databricks import DatabricksTool from litellm.types.llms.openai import ( + ChatCompletionAnnotation, ChatCompletionReasoningItem, ChatCompletionThinkingBlock, ImageURLListItem, @@ -59,6 +60,31 @@ _CHOICES_FIELDS: Final[frozenset] = frozenset(Choices.model_fields.keys()) _MODEL_RESPONSE_FIELDS: Final[frozenset] = frozenset(ModelResponse.model_fields.keys()) | {"usage"} +_ANNOTATION_ENTRIES: Final = TypeAdapter(tuple[object, ...]) + + +def _is_provider_specific_annotation(annotation: object) -> bool: + match annotation: + case {"type": "url_citation"}: + return False + case {"type": _}: + return True + case _: + return False + + +def _annotation_entries(annotations: object) -> tuple[object, ...]: + return _ANNOTATION_ENTRIES.validate_python(annotations) if isinstance(annotations, list) else () + + +def partition_annotations(annotations: object) -> tuple[object, tuple[object, ...]]: + entries: Final = _annotation_entries(annotations) + others: Final = tuple(annotation for annotation in entries if _is_provider_specific_annotation(annotation)) + if not others: + return annotations, () + return [annotation for annotation in entries if not _is_provider_specific_annotation(annotation)] or None, others + + def _normalize_images_for_message( images: list[dict] | None, ) -> list[ImageURLListItem] | None: @@ -787,6 +813,9 @@ def convert_to_model_response_object( provider_specific_fields = dict(choice["message"].get("provider_specific_fields", None) or {}) for f in choice["message"].keys() - _MESSAGE_FIELDS: provider_specific_fields[f] = choice["message"][f] + annotations, other_annotations = partition_annotations(choice["message"].get("annotations", None)) + if other_annotations: + provider_specific_fields["annotations"] = list(other_annotations) # Handle reasoning models that display `reasoning_content` within `content` reasoning_content, content = extract_reasoning_content(choice["message"]) @@ -809,7 +838,9 @@ def convert_to_model_response_object( reasoning_content=reasoning_content, thinking_blocks=thinking_blocks, reasoning_items=choice["message"].get("reasoning_items", None), - annotations=choice["message"].get("annotations", None), + annotations=cast( # cast-ok: provider annotation dicts pass through whole, keeping keys such as url_citation.content + list[ChatCompletionAnnotation] | None, annotations + ), images=_normalize_images_for_message(choice["message"].get("images", None)), ) finish_reason = choice.get("finish_reason", None) diff --git a/litellm/litellm_core_utils/prompt_templates/common_utils.py b/litellm/litellm_core_utils/prompt_templates/common_utils.py index 9bc53791b36..b3f89b33d1d 100644 --- a/litellm/litellm_core_utils/prompt_templates/common_utils.py +++ b/litellm/litellm_core_utils/prompt_templates/common_utils.py @@ -424,6 +424,109 @@ def anthropic_image_source_to_openai_url(image_source: Mapping[str, object]) -> return None +@dataclass(frozen=True, slots=True) +class InlineDocument: + filename: str + file_data: str + + +def _document_media_type(source: object) -> str: + match source: + case {"media_type": str() as media_type} if media_type: + return media_type + case _: + return "application/pdf" + + +def anthropic_inline_document(document: Mapping[str, object]) -> InlineDocument | None: + source: Final = document.get("source") + match source: + case {"type": "base64", "data": str() as data} if data: + title: Final = document.get("title") + return InlineDocument( + filename=title if isinstance(title, str) and title else "document.pdf", + file_data=f"data:{_document_media_type(source)};base64,{data}", + ) + case _: + return None + + +def is_pdf_file_part(part: object) -> bool: + match part: + case {"type": "file", "file": {"file_data": str() as file_data}} if file_data.startswith("data:"): + return file_data.startswith("data:application/pdf") + case {"type": "file", "file": {"format": str() as file_format}}: + return file_format == "application/pdf" + case {"type": "file", "file": {"filename": str() as filename}}: + return filename.lower().endswith(".pdf") + case _: + return False + + +_CONTENT_PARTS: Final = TypeAdapter(tuple[object, ...]) + + +def content_parts(message: Mapping[str, object]) -> Sequence[object]: + content: Final = message.get("content") + return _CONTENT_PARTS.validate_python(content) if isinstance(content, list) else () + + +def _has_pdf_file_part(message: Mapping[str, object]) -> bool: + return any(is_pdf_file_part(part) for part in content_parts(message)) + + +def messages_have_pdf_file_part(messages: Sequence[Mapping[str, object]]) -> bool: + return any(_has_pdf_file_part(message) for message in messages) + + +def _message_without_pdf_file_parts(message: AllMessageValues) -> AllMessageValues: + if not _has_pdf_file_part(message): + return message + kept: Final = [part for part in content_parts(message) if not is_pdf_file_part(part)] + return cast(AllMessageValues, {**message, "content": kept}) # cast-ok: same message with fewer content parts + + +def _is_pdf_only_user_turn(message: AllMessageValues) -> bool: + parts: Final = content_parts(message) + return message["role"] == "user" and bool(parts) and all(is_pdf_file_part(part) for part in parts) + + +def messages_without_pdf_file_parts(messages: Sequence[AllMessageValues]) -> tuple[AllMessageValues, ...]: + return tuple( + _message_without_pdf_file_parts(message) for message in messages if not _is_pdf_only_user_turn(message) + ) + + +_CONTENT_PART: Final = TypeAdapter(Mapping[str, object]) + + +def _file_part_as_image_url(part: object) -> object: + match part: + case {"type": "file", "file": {"file_data": str() as file_data}} if file_data: + fields: Final = _CONTENT_PART.validate_python(part) + return {key: value for key, value in fields.items() if key != "file"} | { + "type": "image_url", + "image_url": {"url": file_data}, + } + case _: + return part + + +def message_with_file_parts_as_image_urls(message: Mapping[str, object]) -> Mapping[str, object]: + parts: Final = content_parts(message) + if not parts: + return message + return {**message, "content": [_file_part_as_image_url(part) for part in parts]} + + +def _typed_message_with_file_parts_as_image_urls(message: AllMessageValues) -> AllMessageValues: + return cast(AllMessageValues, message_with_file_parts_as_image_urls(message)) # cast-ok: same message, image parts + + +def messages_with_file_parts_as_image_urls(messages: Sequence[AllMessageValues]) -> tuple[AllMessageValues, ...]: + return tuple(_typed_message_with_file_parts_as_image_urls(message) for message in messages) + + def _image_part_url(part: Mapping[str, object]) -> str | None: """The image URL carried by one content part, whichever of the three dialects wrote it.""" part_type: Final = part.get("type") diff --git a/litellm/litellm_core_utils/streaming_chunk_builder_utils.py b/litellm/litellm_core_utils/streaming_chunk_builder_utils.py index 7c91ffdff76..8e1791c3c92 100644 --- a/litellm/litellm_core_utils/streaming_chunk_builder_utils.py +++ b/litellm/litellm_core_utils/streaming_chunk_builder_utils.py @@ -19,6 +19,7 @@ from litellm.litellm_core_utils.core_helpers import ( ) from litellm.litellm_core_utils.hidden_params import HIDDEN_PARAMS_ATTR, set_hidden_params from litellm.litellm_core_utils.llm_cost_calc.guardrail_cost import BEDROCK_TRACE_FIELDS, object_items, object_mapping +from litellm.litellm_core_utils.llm_response_utils.convert_dict_to_response import partition_annotations from litellm.litellm_core_utils.prompt_templates.common_utils import ( get_content_from_model_response, # pyright: ignore[reportUnknownVariableType] # response param accepts an untyped dict ) @@ -1700,10 +1701,17 @@ def stream_chunk_builder( # Merge annotations from ALL chunks — providers may spread # them across multiple streaming chunks or send them only in # the final chunk. - all_annotations: Final[list] = [] + all_annotations: Final[list[object]] = [] for ac in annotation_chunks: all_annotations.extend(ac["choices"][0]["delta"]["annotations"]) - response["choices"][0]["message"]["annotations"] = all_annotations + message_annotations, provider_annotations = partition_annotations(all_annotations) + if message_annotations is not None: + response["choices"][0]["message"]["annotations"] = message_annotations + if provider_annotations: + response.choices[0].message.provider_specific_fields = { + **(response.choices[0].message.provider_specific_fields or {}), + "annotations": list(provider_annotations), + } audio_chunks: Final = [ chunk @@ -1766,7 +1774,10 @@ def stream_chunk_builder( if combined_provider_fields: _choice = response.choices[0] - _choice.message.provider_specific_fields = combined_provider_fields + _choice.message.provider_specific_fields = { + **(_choice.message.provider_specific_fields or {}), + **combined_provider_fields, + } completion_output = get_content_from_model_response(response) diff --git a/litellm/llms/anthropic/messages/adapters/transformation.py b/litellm/llms/anthropic/messages/adapters/transformation.py index 3524c9e73d8..ddda759384f 100644 --- a/litellm/llms/anthropic/messages/adapters/transformation.py +++ b/litellm/llms/anthropic/messages/adapters/transformation.py @@ -123,6 +123,7 @@ from openai.types.chat.chat_completion_chunk import Choice as OpenAIStreamingCho from litellm.litellm_core_utils.prompt_templates.common_utils import ( anthropic_image_source_to_openai_url, + anthropic_inline_document, parse_tool_call_arguments, reasoning_content_from_thinking_blocks, with_prompt_cache_breakpoint, @@ -185,6 +186,8 @@ from litellm.types.llms.openai import ( AllMessageValues, ChatCompletionAssistantMessage, ChatCompletionAssistantToolCall, + ChatCompletionFileObject, + ChatCompletionFileObjectFile, ChatCompletionImageObject, ChatCompletionImageUrlObject, ChatCompletionRedactedThinkingBlock, @@ -498,7 +501,9 @@ class LiteLLMAnthropicMessagesAdapter: for m in ordered_messages: user_message: ChatCompletionUserMessage | None = None tool_message_list: list[ChatCompletionToolMessage] = [] - new_user_content_list: list[ChatCompletionTextObject | ChatCompletionImageObject] = [] + new_user_content_list: list[ + ChatCompletionTextObject | ChatCompletionImageObject | ChatCompletionFileObject + ] = [] if m["role"] == "system": system_message = self._translate_midturn_system_message_to_openai(m, model) if system_message is not None: @@ -531,13 +536,8 @@ class LiteLLMAnthropicMessagesAdapter: self._add_prompt_cache_breakpoint_if_present(content, image_obj) ) elif content.get("type") == "document": - # Convert Anthropic document format (PDF, etc.) to OpenAI format - source = content.get("source", {}) - openai_image_url = self._translate_anthropic_image_to_openai(cast(dict, source)) - - if openai_image_url: - image_url_obj = ChatCompletionImageUrlObject(url=openai_image_url) - doc_obj = ChatCompletionImageObject(type="image_url", image_url=image_url_obj) + doc_obj = self._translate_anthropic_document_to_openai(content) + if doc_obj is not None: self._add_cache_control_if_applicable(content, doc_obj, model) new_user_content_list.append(doc_obj) elif content.get("type") == "tool_result": @@ -1357,7 +1357,7 @@ class LiteLLMAnthropicMessagesAdapter: text: Final = tool_result_text_part(block) return None if text is None else ChatCompletionTextObject(type="text", text=text) case "image" | "document": - return self._tool_result_image_part(block.get("source")) + return self._source_image_url_part(block.get("source")) case "tool_reference": return ChatCompletionToolReferenceObject( type="tool_reference", tool_name=str(block.get("tool_name") or "") @@ -1365,7 +1365,18 @@ class LiteLLMAnthropicMessagesAdapter: case _: return None - def _tool_result_image_part(self, image_source: object) -> ChatCompletionImageObject | None: + def _translate_anthropic_document_to_openai( + self, document: Mapping[str, object] + ) -> ChatCompletionFileObject | ChatCompletionImageObject | None: + inline_document: Final = anthropic_inline_document(document) + if inline_document is None: + return self._source_image_url_part(document.get("source")) + return ChatCompletionFileObject( + type="file", + file=ChatCompletionFileObjectFile(filename=inline_document.filename, file_data=inline_document.file_data), + ) + + def _source_image_url_part(self, image_source: object) -> ChatCompletionImageObject | None: if not isinstance(image_source, dict): return None openai_image_url = self._translate_anthropic_image_to_openai(image_source) diff --git a/litellm/llms/anthropic/messages/responses_adapters/transformation.py b/litellm/llms/anthropic/messages/responses_adapters/transformation.py index 1dd07c1739b..94c608f90e6 100644 --- a/litellm/llms/anthropic/messages/responses_adapters/transformation.py +++ b/litellm/llms/anthropic/messages/responses_adapters/transformation.py @@ -13,6 +13,7 @@ from typing import Any, Final, cast from litellm.litellm_core_utils.prompt_templates.common_utils import ( TOOL_RESULT_IMAGE_BOUNDARY, TOOL_RESULT_IMAGE_PLACEHOLDER, + anthropic_inline_document, encrypted_reasoning_signature, responses_reasoning_items_from_thinking_blocks, with_prompt_cache_breakpoint, @@ -106,27 +107,18 @@ class LiteLLMAnthropicToResponsesAPIAdapter: block: Mapping[str, object], ) -> dict[str, str] | None: # mutable-ok: API message payload """Convert an Anthropic document block to a Responses input_file part.""" + inline_document: Final = anthropic_inline_document(block) + if inline_document is not None: + return { + "type": "input_file", + "filename": inline_document.filename, + "file_data": inline_document.file_data, + } raw_source: Final = block.get("source") if not isinstance(raw_source, Mapping): return None source: Final = cast(Mapping[str, object], raw_source) # cast-ok: untrusted client payload - source_type: Final = source.get("type") - if source_type == "base64": - data: Final = source.get("data") - if not isinstance(data, str) or not data: - return None - raw_media_type: Final = source.get("media_type") - media_type: Final = ( - raw_media_type if isinstance(raw_media_type, str) and raw_media_type else "application/pdf" - ) - raw_title: Final = block.get("title") - filename: Final = raw_title if isinstance(raw_title, str) and raw_title else "document.pdf" - return { - "type": "input_file", - "filename": filename, - "file_data": f"data:{media_type};base64,{data}", - } - if source_type == "url": + if source.get("type") == "url": url: Final = source.get("url") if not isinstance(url, str) or not url: return None diff --git a/litellm/llms/azure_ai/chat/transformation.py b/litellm/llms/azure_ai/chat/transformation.py index 8bb6072d0a9..d382d653a53 100644 --- a/litellm/llms/azure_ai/chat/transformation.py +++ b/litellm/llms/azure_ai/chat/transformation.py @@ -1,7 +1,7 @@ import copy import enum import re -from collections.abc import AsyncIterator, Collection, Iterator +from collections.abc import AsyncIterator, Collection, Iterator, Mapping from typing import TYPE_CHECKING, Final, cast import httpx @@ -10,10 +10,11 @@ from httpx import Response import litellm from litellm._logging import verbose_logger from litellm.litellm_core_utils.prompt_templates.common_utils import ( - audio_or_image_in_message_content, + content_parts, convert_content_list_to_str, drop_tool_reference_parts_from_tool_messages, filter_value_from_dict, + is_pdf_file_part, ) from litellm.llms.azure.common_utils import BaseAzureLLM from litellm.llms.azure_ai.common_utils import ( @@ -46,6 +47,20 @@ NON_OPENAI_SPEC_MESSAGE_FIELDS: Final = ( "cache_control", ) +LIST_ONLY_CONTENT_PART_TYPES: Final = frozenset({"image_url", "input_audio"}) + + +def _is_list_only_content_part(part: object) -> bool: + match part: + case {"type": str() as part_type} if part_type in LIST_ONLY_CONTENT_PART_TYPES: + return True + case _: + return is_pdf_file_part(part) + + +def _has_list_only_content_part(message: Mapping[str, object]) -> bool: + return any(_is_list_only_content_part(part) for part in content_parts(message)) + class AzureAIChatCompletionResponseIterator(OpenAIChatCompletionResponseIterator): def chunk_parser(self, chunk: dict[str, object]) -> ModelResponseStream: # mutable-ok: mirrors override contract @@ -246,7 +261,7 @@ class AzureAIStudioConfig(OpenAIConfig): these with "Extra inputs are not permitted", which breaks multi-turn Anthropic-format clients that echo thinking blocks back as history. 2. Transforms list content to a string. - 3. If message contains an image or audio, send as is (user-intended) + 3. If message contains an image, audio, or PDF file part, send as is (user-intended) Operates on a deep copy so the caller's messages keep their thinking blocks and provider metadata, which a fallback to another provider still needs. @@ -257,8 +272,7 @@ class AzureAIStudioConfig(OpenAIConfig): for field in NON_OPENAI_SPEC_MESSAGE_FIELDS: filter_value_from_dict(message_dict, field) - # Do nothing if the message contains an image or audio - if audio_or_image_in_message_content(message): + if _has_list_only_content_part(message): continue texts = convert_content_list_to_str(message=message) diff --git a/litellm/llms/gigachat/chat/transformation.py b/litellm/llms/gigachat/chat/transformation.py index 65921b628fb..81a186eadf4 100644 --- a/litellm/llms/gigachat/chat/transformation.py +++ b/litellm/llms/gigachat/chat/transformation.py @@ -38,6 +38,15 @@ else: _EMPTY_FUNCTION: Final[Mapping[str, object]] = MappingProxyType({}) _JSON_VALUE: Final = TypeAdapter(object) +_CONTENT_PART: Final = TypeAdapter(Mapping[str, object]) + + +def _file_part_data(part: Mapping[str, object]) -> str | None: + match part: + case {"type": "file", "file": {"file_data": str() as file_data}} if file_data: + return file_data + case _: + return None def is_valid_json(value: str) -> bool: @@ -295,7 +304,7 @@ class GigaChatConfig(BaseConfig): def _transform_list_content(self, content: Sequence) -> tuple[str, Sequence[object]]: """ - Extract text and image attachments from a multimodal message content list. + Extract text, image and file attachments from a multimodal message content list. Args: content: List of content parts (OpenAI multimodal format) @@ -320,6 +329,10 @@ class GigaChatConfig(BaseConfig): file_id = self._upload_image(upload_url) if file_id: attachments.append(file_id) + elif (file_data := _file_part_data(_CONTENT_PART.validate_python(part))) and ( + uploaded_file_id := self._upload_image(file_data) + ): + attachments.append(uploaded_file_id) text: Final = "\n".join(texts) if texts else "" return text, attachments diff --git a/litellm/llms/oci/chat/transformation.py b/litellm/llms/oci/chat/transformation.py index e8d37f857e6..8257a75b4a5 100644 --- a/litellm/llms/oci/chat/transformation.py +++ b/litellm/llms/oci/chat/transformation.py @@ -20,6 +20,7 @@ import litellm from litellm._logging import verbose_logger from litellm.constants import DEFAULT_OCI_CHAT_MAX_TOKENS from litellm.litellm_core_utils.logging_utils import track_llm_api_timing +from litellm.litellm_core_utils.prompt_templates.common_utils import messages_with_file_parts_as_image_urls from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.base_llm.chat.transformation import BaseConfig, BaseLLMException @@ -681,7 +682,9 @@ class OCIChatConfig(BaseConfig): servingMode=serving_mode, chatRequest=OCIChatRequestPayload( apiFormat=vendor.value, - messages=adapt_messages_to_generic_oci_standard(messages), + messages=adapt_messages_to_generic_oci_standard( + list(messages_with_file_parts_as_image_urls(messages)) + ), **self._get_optional_params(vendor, optional_params, model), ), ) diff --git a/litellm/llms/sap/chat/transformation.py b/litellm/llms/sap/chat/transformation.py index 622c28f3c90..0270f814328 100755 --- a/litellm/llms/sap/chat/transformation.py +++ b/litellm/llms/sap/chat/transformation.py @@ -11,6 +11,7 @@ import httpx from pydantic import BaseModel import litellm +from litellm.litellm_core_utils.prompt_templates.common_utils import message_with_file_parts_as_image_urls from litellm.litellm_core_utils.url_utils import join_api_base from litellm.types.llms.openai import AllMessageValues from litellm.types.utils import ModelResponse @@ -104,7 +105,7 @@ def _message_role_mapping(role: str) -> type[ChatMessage]: def _messages_to_sap_template(messages: list[AllMessageValues]) -> list: template: Final = [] for message in messages: - folded_message = _fold_message_cache_control(message) + folded_message = _fold_message_cache_control(message_with_file_parts_as_image_urls(message)) message_role_class = _message_role_mapping(str(folded_message["role"])) validated_message = validate_dict(folded_message, message_role_class) template.append(validated_message) diff --git a/litellm/llms/snowflake/chat/transformation.py b/litellm/llms/snowflake/chat/transformation.py index 030413030eb..928dcc3bc52 100644 --- a/litellm/llms/snowflake/chat/transformation.py +++ b/litellm/llms/snowflake/chat/transformation.py @@ -112,7 +112,7 @@ def _convert_image_url_to_anthropic(block: Mapping[str, object]) -> object: pdf/text data URIs become document blocks rather than malformed image blocks. """ image_url: Final = block.get("image_url") - url: Final = image_url if isinstance(image_url, str) else _image_url_field(image_url, "url") + url: Final = image_url if isinstance(image_url, str) else _str_field(image_url, "url") if not url: return block @@ -121,18 +121,31 @@ def _convert_image_url_to_anthropic(block: Mapping[str, object]) -> object: if select_anthropic_content_block_type_for_file(_data_uri_media_type(url)) == "document" else create_anthropic_image_param( image_url if isinstance(image_url, dict) else url, - format=_image_url_field(image_url, "format"), + format=_str_field(image_url, "format"), is_bedrock_invoke=True, ) ) + return _with_cache_control(converted, block) + + +def _convert_file_to_anthropic(block: Mapping[str, object]) -> object: + file_data: Final = _str_field(block.get("file"), "file_data") + if not file_data: + return block + return _with_cache_control( + anthropic_process_openai_file_message({"type": "file", "file": {"file_data": file_data}}), block + ) + + +def _with_cache_control(converted: Mapping[str, object], block: Mapping[str, object]) -> Mapping[str, object]: cache_control: Final = block.get("cache_control") if cache_control is None: return converted return {**converted, "cache_control": cache_control} -def _image_url_field(image_url: object, key: str) -> str | None: - value: Final = image_url.get(key) if isinstance(image_url, dict) else None +def _str_field(container: object, key: str) -> str | None: + value: Final = container.get(key) if isinstance(container, dict) else None return value if isinstance(value, str) else None @@ -141,15 +154,23 @@ def _data_uri_media_type(url: str) -> str: return match.group(1) if match else "" -def _convert_image_url_blocks_to_anthropic(content: object) -> object: +def _convert_media_block_to_anthropic(block: object) -> object: + mapping: Final = as_str_mapping(block) + if mapping is None: + return block + match mapping.get("type"): + case "image_url": + return _convert_image_url_to_anthropic(mapping) + case "file": + return _convert_file_to_anthropic(mapping) + case _: + return block + + +def _convert_media_blocks_to_anthropic(content: object) -> object: if not isinstance(content, list): return content - return [ - _convert_image_url_to_anthropic(block) - if isinstance(block, Mapping) and block.get("type") == "image_url" - else block - for block in content - ] + return [_convert_media_block_to_anthropic(block) for block in content] def _convert_tool_result_to_anthropic( @@ -391,7 +412,7 @@ class SnowflakeConfig(SnowflakeBaseConfig, OpenAIGPTConfig): conversation.append( { "role": role, - "content": _convert_image_url_blocks_to_anthropic(content), + "content": _convert_media_blocks_to_anthropic(content), } ) diff --git a/litellm/llms/vertex_ai/vertex_ai_partner_models/main.py b/litellm/llms/vertex_ai/vertex_ai_partner_models/main.py index bb200fd542d..a3a3427670d 100644 --- a/litellm/llms/vertex_ai/vertex_ai_partner_models/main.py +++ b/litellm/llms/vertex_ai/vertex_ai_partner_models/main.py @@ -9,6 +9,7 @@ import httpx import litellm from litellm import LlmProviders +from litellm.litellm_core_utils.prompt_templates.common_utils import messages_with_file_parts_as_image_urls from litellm.types.llms.vertex_ai import VertexPartnerProvider from litellm.utils import ModelResponse @@ -217,7 +218,7 @@ class VertexAIPartnerModels(VertexBase): return base_llm_http_handler.completion( model=model, stream=stream, - messages=messages, + messages=list(messages_with_file_parts_as_image_urls(messages)), acompletion=acompletion, api_base=api_base, model_response=model_response, diff --git a/litellm/main.py b/litellm/main.py index d1af2de3a13..50ab274a060 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -108,6 +108,10 @@ from litellm.litellm_core_utils.mock_functions import ( from litellm.litellm_core_utils.prompt_templates.common_utils import ( get_content_from_model_response as get_content_from_model_response, # noqa: PLC0414 # public re-export ) +from litellm.litellm_core_utils.prompt_templates.common_utils import ( + messages_have_pdf_file_part, + messages_without_pdf_file_parts, +) from litellm.litellm_core_utils.provider_affinity import add_provider_affinity_header from litellm.litellm_core_utils.request_timeout_resolver import ( get_configured_request_timeout, @@ -191,11 +195,13 @@ from litellm.utils import ( get_requester_metadata, get_secret, get_standard_openai_params, + is_explicitly_disabled_factory, mock_completion_streaming_obj, pre_process_non_default_params, read_config_args, should_run_mock_completion, supports_httpx_timeout, + supports_pdf_input, validate_and_fix_openai_messages, validate_and_fix_openai_tools, validate_and_fix_thinking_param, @@ -393,6 +399,125 @@ def _resolve_control_options(kwargs: Mapping[str, object], model: str) -> Contro return assert_never(control) +PROVIDERS_WITHOUT_DOCUMENT_INPUT: Final = frozenset({"ollama", "ollama_chat"}) +PROVIDERS_DROPPING_UNDECLARED_PDF_INPUT: Final = frozenset({"azure_ai"}) + + +def _cost_map_pdf_input_support(model: str, custom_llm_provider: str) -> bool | None: + if custom_llm_provider in PROVIDERS_WITHOUT_DOCUMENT_INPUT or is_explicitly_disabled_factory( + model, custom_llm_provider, "supports_pdf_input" + ): + return False + if supports_pdf_input(model, custom_llm_provider): + return True + return None + + +def _deployment_pdf_input_support(model: str, base_model: str, custom_llm_provider: str) -> bool | None: + supports: Final = ( + _cost_map_pdf_input_support(model, custom_llm_provider), + _cost_map_pdf_input_support(base_model, custom_llm_provider), + ) + if False in supports: + return False + if True in supports: + return True + return None + + +def _declared_pdf_input_support( + model: str, custom_llm_provider: str, model_info: Mapping[str, object] | None, base_model: str | None +) -> bool | None: + match model_info: + case {"supports_pdf_input": bool() as declared}: + return declared + case _ if base_model: + return _deployment_pdf_input_support(model, base_model, custom_llm_provider) + case _: + return _cost_map_pdf_input_support(model, custom_llm_provider) + + +def _pdf_input_unreadable( + model: str, + custom_llm_provider: str, + model_info: Mapping[str, object] | None, + base_model: str | None, + drop_params: bool, +) -> bool: + match _declared_pdf_input_support(model, custom_llm_provider, model_info, base_model): + case False: + return True + case None: + return drop_params and custom_llm_provider in PROVIDERS_DROPPING_UNDECLARED_PDF_INPUT + case True: + return False + + +def _messages_without_unreadable_pdfs( + messages: Sequence[AllMessageValues], + model: str, + custom_llm_provider: str, + model_info: Mapping[str, object] | None, + base_model: str | None, + drop_params: bool, +) -> tuple[AllMessageValues, ...] | None: + if not messages_have_pdf_file_part(messages) or not _pdf_input_unreadable( + model, custom_llm_provider, model_info, base_model, drop_params + ): + return None + if not drop_params: + raise litellm.BadRequestError( + message=( + f"{custom_llm_provider}/{model} does not accept PDF input. " + "Remove the PDF, or set drop_params=True to send the request without it" + ), + model=model, + llm_provider=custom_llm_provider, + ) + readable_messages: Final = messages_without_pdf_file_parts(messages) + if not any(message["role"] == "user" for message in readable_messages): + raise litellm.BadRequestError( + message=( + f"{custom_llm_provider}/{model} does not accept PDF input, " + "and no user message is left once the PDF is dropped" + ), + model=model, + llm_provider=custom_llm_provider, + ) + return readable_messages + + +def _messages_readable_by_model( + messages: Sequence[AllMessageValues], + model: str, + custom_llm_provider: str, + model_info: Mapping[str, object] | None, + base_model: str | None, + drop_params: bool, + ensure_alternating_roles: bool | None, + user_continue_message: ChatCompletionUserMessage | None, + assistant_continue_message: ChatCompletionAssistantMessage | None, +) -> tuple[AllMessageValues, ...] | None: + pdf_free_messages: Final = _messages_without_unreadable_pdfs( + messages=messages, + model=model, + custom_llm_provider=custom_llm_provider, + model_info=model_info, + base_model=base_model, + drop_params=drop_params, + ) + if pdf_free_messages is None: + return None + return tuple( + get_completion_messages( + messages=list(pdf_free_messages), + ensure_alternating_roles=ensure_alternating_roles or False, + user_continue_message=user_continue_message, + assistant_continue_message=assistant_continue_message, + ) + ) + + ####### COMPLETION ENDPOINTS ################ @@ -4234,7 +4359,6 @@ def _complete_sagemaker(ctx: CompletionDispatchContext) -> _CompletionDispatchRe _ADDITIONAL_DROP_PARAMS_ADAPTER: Final = TypeAdapter(list[str]) _OPTIONAL_PARAMS_ADAPTER: Final = TypeAdapter(dict[str, object]) _OLLAMA_EMBEDDING_HEADERS_ADAPTER: Final = TypeAdapter(dict[str, str]) -_REQUEST_KWARGS_ADAPTER: Final = TypeAdapter(Mapping[str, object]) _MODEL_INFO_ADAPTER: Final = TypeAdapter(Mapping[str, object]) _PROVIDER_ROUTING_HINTS: Final = TypeAdapter(ProviderRoutingHints) _BASE_URL_HINT: Final = TypeAdapter(BaseUrlRoutingHint) @@ -5558,9 +5682,15 @@ def completion( ) headers = kwargs.get("headers", None) or extra_headers - ensure_alternating_roles: Final[bool | None] = kwargs.get("ensure_alternating_roles", None) - user_continue_message: Final[ChatCompletionUserMessage | None] = kwargs.get("user_continue_message", None) - assistant_continue_message: ChatCompletionAssistantMessage | None = kwargs.get("assistant_continue_message", None) + ensure_alternating_roles: Final = cast( # cast-ok: read from untyped kwargs + bool | None, kwargs.get("ensure_alternating_roles", None) + ) + user_continue_message: Final = cast( # cast-ok: read from untyped kwargs + ChatCompletionUserMessage | None, kwargs.get("user_continue_message", None) + ) + assistant_continue_message: Final = cast( # cast-ok: read from untyped kwargs + ChatCompletionAssistantMessage | None, kwargs.get("assistant_continue_message", None) + ) headers = {} if headers is None else dict(headers) if extra_headers is not None: headers.update(extra_headers) @@ -5928,15 +6058,17 @@ def completion( def optional_params_for_model_info(deployment_info: Mapping[str, object] | None) -> dict[str, object]: return get_optional_params(**mapping_args, model_info=deployment_info, **mapping_special_params) - optional_params = optional_params_for_model_info( + validated_model_info: Final = ( _MODEL_INFO_ADAPTER.validate_python(model_info) if isinstance(model_info, dict) else None ) + deployment_base_model: Final = base_model if isinstance(base_model, str) else None + optional_params = optional_params_for_model_info(validated_model_info) fakes_stream: Final = partial( _completion_fakes_stream, model, custom_llm_provider, provider_config, - base_model=base_model if isinstance(base_model, str) else None, + base_model=deployment_base_model, bridges_to_responses_api=bridges_to_responses_api is True, ) resolve_timeout: Final = partial( @@ -5974,16 +6106,29 @@ def completion( functions_unsupported_model: Final[list[object]] = TypeAdapter(list[object]).validate_python( optional_params.pop("functions_unsupported_model") ) - request_kwargs: Final = _REQUEST_KWARGS_ADAPTER.validate_python(kwargs) messages = function_call_prompt( messages=messages, functions=functions_unsupported_model, model=model, llm_provider=custom_llm_provider, - drop_params=litellm.drop_params is True - or normalize_drop_params(request_kwargs.get("drop_params")) is True, + drop_params=drop_unsupported_reasoning, ) + readable_messages: Final = _messages_readable_by_model( + messages=cast(list[AllMessageValues], messages), # cast-ok: completion types messages as a bare List + model=model, + custom_llm_provider=custom_llm_provider, + model_info=validated_model_info, + base_model=deployment_base_model, + drop_params=drop_unsupported_reasoning, + ensure_alternating_roles=ensure_alternating_roles, + user_continue_message=user_continue_message, + assistant_continue_message=assistant_continue_message, + ) + messages = ( # rebind-ok: completion() rewrites messages step by step + messages if readable_messages is None else list(readable_messages) + ) + # For logging - save the values of the litellm-specific params passed in requested_litellm_params: Final = get_litellm_params( acompletion=acompletion, diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 1256fef9181..c6bbb63b6fc 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -4049,6 +4049,7 @@ "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", "supports_function_calling": true, "supports_parallel_function_calling": true, + "supports_pdf_input": false, "supports_response_schema": true, "supports_tool_choice": true }, @@ -12437,7 +12438,8 @@ "supported_output_modalities": [ "text" ], - "supports_native_streaming": true + "supports_native_streaming": true, + "supports_pdf_input": false }, "azure_ai/FLUX-1.1-pro": { "deprecation_date": "2027-01-05", @@ -12502,6 +12504,7 @@ "output_cost_per_token": 1.85e-06, "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", "supports_function_calling": true, + "supports_pdf_input": false, "supports_prompt_caching": true, "supports_reasoning": true, "supports_tool_choice": true @@ -12516,6 +12519,7 @@ "output_cost_per_token": 3.828e-06, "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", "supports_function_calling": true, + "supports_pdf_input": false, "supports_prompt_caching": true, "supports_reasoning": true, "supports_tool_choice": true @@ -12531,6 +12535,7 @@ "output_cost_per_token": 3.52e-06, "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", "supports_function_calling": true, + "supports_pdf_input": false, "supports_prompt_caching": true, "supports_reasoning": true, "supports_tool_choice": true @@ -12546,6 +12551,7 @@ "output_cost_per_token": 4.84e-06, "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", "supports_function_calling": true, + "supports_pdf_input": false, "supports_prompt_caching": true, "supports_reasoning": true, "supports_tool_choice": true @@ -12560,6 +12566,7 @@ "output_cost_per_token": 4.84e-06, "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", "supports_function_calling": true, + "supports_pdf_input": false, "supports_prompt_caching": true, "supports_reasoning": true, "supports_tool_choice": true @@ -12574,6 +12581,7 @@ "output_cost_per_token": 7.26e-06, "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", "supports_function_calling": true, + "supports_pdf_input": false, "supports_prompt_caching": true, "supports_reasoning": true, "supports_tool_choice": true @@ -12594,6 +12602,7 @@ "text" ], "supports_function_calling": true, + "supports_pdf_input": false, "supports_prompt_caching": true, "supports_reasoning": true, "supports_tool_choice": true @@ -12616,6 +12625,7 @@ "text" ], "supports_function_calling": true, + "supports_pdf_input": false, "supports_prompt_caching": true, "supports_reasoning": true, "supports_tool_choice": true, @@ -12638,6 +12648,7 @@ "text" ], "supports_function_calling": true, + "supports_pdf_input": false, "supports_prompt_caching": true, "supports_reasoning": true, "supports_tool_choice": true, @@ -12660,6 +12671,7 @@ "text" ], "supports_function_calling": true, + "supports_pdf_input": false, "supports_prompt_caching": true, "supports_reasoning": true, "supports_tool_choice": true, @@ -12687,6 +12699,7 @@ "text" ], "supports_function_calling": true, + "supports_pdf_input": false, "supports_prompt_caching": true, "supports_reasoning": true, "supports_tool_choice": true, @@ -12703,6 +12716,7 @@ "output_cost_per_token": 1.32e-06, "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", "supports_function_calling": true, + "supports_pdf_input": false, "supports_prompt_caching": true, "supports_reasoning": true, "supports_tool_choice": true @@ -12724,6 +12738,7 @@ "text" ], "supports_function_calling": true, + "supports_pdf_input": false, "supports_prompt_caching": true, "supports_reasoning": true, "supports_tool_choice": true, @@ -12744,6 +12759,7 @@ "text" ], "supports_function_calling": true, + "supports_pdf_input": false, "supports_prompt_caching": true, "supports_reasoning": true, "supports_tool_choice": true, @@ -12765,6 +12781,7 @@ "text" ], "supports_function_calling": true, + "supports_pdf_input": false, "supports_prompt_caching": true, "supports_reasoning": true, "supports_tool_choice": true @@ -12833,6 +12850,7 @@ "text" ], "supports_function_calling": true, + "supports_pdf_input": false, "supports_prompt_caching": true, "supports_reasoning": true, "supports_tool_choice": true @@ -12847,6 +12865,7 @@ "output_cost_per_token": 7.1e-07, "source": "https://learn.microsoft.com/en-us/azure/foundry/foundry-models/concepts/models-sold-directly-by-azure", "supports_function_calling": true, + "supports_pdf_input": false, "supports_tool_choice": true }, "azure_ai/Llama-4-Maverick-17B-128E-Instruct-FP8": { @@ -12859,6 +12878,7 @@ "output_cost_per_token": 1e-06, "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", "supports_function_calling": true, + "supports_pdf_input": false, "supports_tool_choice": true, "supports_vision": true }, @@ -12872,6 +12892,7 @@ "output_cost_per_token": 7.8e-07, "source": "https://azure.microsoft.com/en-us/blog/introducing-the-llama-4-herd-in-azure-ai-foundry-and-azure-databricks/", "supports_function_calling": true, + "supports_pdf_input": false, "supports_tool_choice": true, "supports_vision": true }, @@ -12884,6 +12905,7 @@ "mode": "chat", "output_cost_per_token": 3.7e-07, "source": "https://learn.microsoft.com/en-us/azure/foundry/openai/concepts/retired-models", + "supports_pdf_input": false, "supports_tool_choice": true }, "azure_ai/Meta-Llama-3.1-70B-Instruct": { @@ -12895,6 +12917,7 @@ "mode": "chat", "output_cost_per_token": 3.54e-06, "source": "https://learn.microsoft.com/en-us/azure/foundry/openai/concepts/retired-models", + "supports_pdf_input": false, "supports_tool_choice": true }, "azure_ai/Phi-3-medium-128k-instruct": { @@ -12906,6 +12929,7 @@ "mode": "chat", "output_cost_per_token": 6.8e-07, "source": "https://learn.microsoft.com/en-us/azure/foundry/openai/concepts/retired-models", + "supports_pdf_input": false, "supports_tool_choice": true, "supports_vision": false }, @@ -12918,6 +12942,7 @@ "mode": "chat", "output_cost_per_token": 6.8e-07, "source": "https://learn.microsoft.com/en-us/azure/foundry/openai/concepts/retired-models", + "supports_pdf_input": false, "supports_tool_choice": true, "supports_vision": false }, @@ -12930,6 +12955,7 @@ "mode": "chat", "output_cost_per_token": 5.2e-07, "source": "https://learn.microsoft.com/en-us/azure/foundry/openai/concepts/retired-models", + "supports_pdf_input": false, "supports_tool_choice": true, "supports_vision": false }, @@ -12942,6 +12968,7 @@ "mode": "chat", "output_cost_per_token": 5.2e-07, "source": "https://learn.microsoft.com/en-us/azure/foundry/openai/concepts/retired-models", + "supports_pdf_input": false, "supports_tool_choice": true, "supports_vision": false }, @@ -12954,6 +12981,7 @@ "mode": "chat", "output_cost_per_token": 6e-07, "source": "https://learn.microsoft.com/en-us/azure/foundry/openai/concepts/retired-models", + "supports_pdf_input": false, "supports_tool_choice": true, "supports_vision": false }, @@ -12966,6 +12994,7 @@ "mode": "chat", "output_cost_per_token": 6e-07, "source": "https://learn.microsoft.com/en-us/azure/foundry/openai/concepts/retired-models", + "supports_pdf_input": false, "supports_tool_choice": true, "supports_vision": false }, @@ -12978,6 +13007,7 @@ "mode": "chat", "output_cost_per_token": 6.4e-07, "source": "https://learn.microsoft.com/en-us/azure/foundry/openai/concepts/retired-models", + "supports_pdf_input": false, "supports_tool_choice": true, "supports_vision": false }, @@ -12990,6 +13020,7 @@ "mode": "chat", "output_cost_per_token": 5.2e-07, "source": "https://learn.microsoft.com/en-us/azure/foundry/openai/concepts/retired-models", + "supports_pdf_input": false, "supports_tool_choice": true, "supports_vision": false }, @@ -13002,6 +13033,7 @@ "mode": "chat", "output_cost_per_token": 5.2e-07, "source": "https://learn.microsoft.com/en-us/azure/foundry/openai/concepts/retired-models", + "supports_pdf_input": false, "supports_tool_choice": true, "supports_vision": true }, @@ -13015,6 +13047,7 @@ "output_cost_per_token": 5e-07, "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", "supports_function_calling": true, + "supports_pdf_input": false, "supports_tool_choice": true, "supports_vision": false }, @@ -13027,7 +13060,8 @@ "mode": "chat", "output_cost_per_token": 3e-07, "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", - "supports_function_calling": true + "supports_function_calling": true, + "supports_pdf_input": false }, "azure_ai/Phi-4-multimodal-instruct": { "deprecation_date": "2027-01-17", @@ -13041,6 +13075,7 @@ "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", "supports_audio_input": true, "supports_function_calling": true, + "supports_pdf_input": false, "supports_vision": true }, "azure_ai/Phi-4-mini-reasoning": { @@ -13052,7 +13087,8 @@ "mode": "chat", "output_cost_per_token": 3e-07, "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", - "supports_function_calling": true + "supports_function_calling": true, + "supports_pdf_input": false }, "azure_ai/Phi-4-reasoning": { "deprecation_date": "2027-01-17", @@ -13064,6 +13100,7 @@ "output_cost_per_token": 5e-07, "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", "supports_function_calling": true, + "supports_pdf_input": false, "supports_tool_choice": true, "supports_reasoning": true }, @@ -13106,6 +13143,7 @@ "output_cost_per_token": 1e-05, "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", "supports_function_calling": true, + "supports_pdf_input": false, "supports_tool_choice": true }, "azure_ai/doc-intelligence/prebuilt-read": { @@ -13144,6 +13182,7 @@ "mode": "chat", "output_cost_per_token": 5.4e-06, "source": "https://learn.microsoft.com/en-us/azure/foundry/openai/concepts/retired-models", + "supports_pdf_input": false, "supports_reasoning": true, "supports_tool_choice": true }, @@ -13202,6 +13241,7 @@ "source": "https://learn.microsoft.com/en-us/azure/foundry/foundry-models/concepts/models-sold-directly-by-azure", "supports_assistant_prefill": true, "supports_function_calling": true, + "supports_pdf_input": false, "supports_prompt_caching": true, "supports_reasoning": true, "supports_tool_choice": true @@ -13217,6 +13257,7 @@ "source": "https://learn.microsoft.com/en-us/azure/foundry/foundry-models/concepts/models-sold-directly-by-azure", "supports_assistant_prefill": true, "supports_function_calling": true, + "supports_pdf_input": false, "supports_prompt_caching": true, "supports_reasoning": true, "supports_tool_choice": true @@ -13230,6 +13271,7 @@ "mode": "chat", "output_cost_per_token": 4.56e-06, "source": "https://learn.microsoft.com/en-us/azure/foundry/openai/concepts/retired-models", + "supports_pdf_input": false, "supports_tool_choice": true }, "azure_ai/deepseek-v4-pro": { @@ -13242,6 +13284,7 @@ "output_cost_per_token": 3.48e-06, "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", "supports_function_calling": true, + "supports_pdf_input": false, "supports_reasoning": true, "supports_tool_choice": true, "cache_read_input_token_cost": 1.45e-07, @@ -13257,6 +13300,7 @@ "output_cost_per_token": 5.1e-07, "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", "supports_function_calling": true, + "supports_pdf_input": false, "supports_reasoning": true, "supports_tool_choice": true, "cache_read_input_token_cost": 2.8e-08, @@ -13273,6 +13317,7 @@ "output_cost_per_token": 1.32e-06, "source": "https://learn.microsoft.com/en-us/azure/foundry/foundry-models/concepts/models-sold-directly-by-azure", "supports_function_calling": true, + "supports_pdf_input": false, "supports_prompt_caching": true, "supports_reasoning": true, "supports_tool_choice": true @@ -13306,6 +13351,7 @@ "output_cost_per_token": 1.5e-05, "source": "https://learn.microsoft.com/en-us/azure/foundry/foundry-models/concepts/models-sold-directly-by-azure", "supports_function_calling": true, + "supports_pdf_input": false, "supports_response_schema": true, "supports_tool_choice": true, "supports_web_search": true @@ -13323,6 +13369,7 @@ "output_cost_per_token_above_200k_tokens": 5e-06, "source": "https://learn.microsoft.com/en-us/azure/foundry/foundry-models/concepts/models-sold-directly-by-azure", "supports_function_calling": true, + "supports_pdf_input": false, "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, @@ -13343,6 +13390,7 @@ "output_cost_per_token_above_200k_tokens": 1.2e-05, "source": "https://azure.microsoft.com/en-us/pricing/details/ai-foundry-models/grok/", "supports_function_calling": true, + "supports_pdf_input": false, "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, @@ -13361,6 +13409,7 @@ "output_cost_per_token": 2.5e-06, "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", "supports_function_calling": true, + "supports_pdf_input": false, "supports_prompt_caching": true, "supports_response_schema": true, "supports_tool_choice": true, @@ -13379,6 +13428,7 @@ "output_cost_per_token": 2.5e-06, "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", "supports_function_calling": true, + "supports_pdf_input": false, "supports_prompt_caching": true, "supports_response_schema": true, "supports_tool_choice": true, @@ -13395,6 +13445,7 @@ "mode": "chat", "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", "supports_function_calling": true, + "supports_pdf_input": false, "supports_response_schema": true, "supports_tool_choice": true, "supports_web_search": true @@ -13409,6 +13460,7 @@ "mode": "chat", "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", "supports_function_calling": true, + "supports_pdf_input": false, "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, @@ -13424,6 +13476,7 @@ "output_cost_per_token": 1.5e-06, "source": "https://learn.microsoft.com/en-us/azure/foundry/foundry-models/concepts/models-sold-directly-by-azure", "supports_function_calling": true, + "supports_pdf_input": false, "supports_response_schema": true, "supports_tool_choice": true, "supports_web_search": true @@ -13436,7 +13489,8 @@ "max_output_tokens": 8192, "mode": "chat", "output_cost_per_token": 0.00971, - "source": "https://learn.microsoft.com/en-us/azure/foundry/openai/concepts/retired-models" + "source": "https://learn.microsoft.com/en-us/azure/foundry/openai/concepts/retired-models", + "supports_pdf_input": false }, "azure_ai/jamba-instruct": { "deprecation_date": "2025-03-01", @@ -13447,6 +13501,7 @@ "mode": "chat", "output_cost_per_token": 7e-07, "source": "https://learn.microsoft.com/en-us/azure/foundry/openai/concepts/retired-models", + "supports_pdf_input": false, "supports_tool_choice": true }, "azure_ai/kimi-k2.5": { @@ -13459,6 +13514,7 @@ "output_cost_per_token": 3e-06, "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", "supports_function_calling": true, + "supports_pdf_input": false, "supports_tool_choice": true, "supports_video_input": true, "supports_vision": true, @@ -13482,6 +13538,7 @@ "text" ], "supports_function_calling": true, + "supports_pdf_input": false, "supports_reasoning": true, "supports_tool_choice": true, "supports_vision": true, @@ -13498,6 +13555,7 @@ "output_cost_per_token": 4e-08, "source": "https://marketplace.microsoft.com/en/marketplace/apps/000-000.ministral-3b-2410-offer?tab=Overview", "supports_function_calling": true, + "supports_pdf_input": false, "supports_tool_choice": true }, "azure_ai/mistral-large": { @@ -13510,6 +13568,7 @@ "output_cost_per_token": 1.2e-05, "source": "https://learn.microsoft.com/en-us/azure/foundry/openai/concepts/retired-models", "supports_function_calling": true, + "supports_pdf_input": false, "supports_tool_choice": true }, "azure_ai/mistral-large-2407": { @@ -13522,6 +13581,7 @@ "output_cost_per_token": 6e-06, "source": "https://learn.microsoft.com/en-us/azure/foundry/openai/concepts/retired-models", "supports_function_calling": true, + "supports_pdf_input": false, "supports_tool_choice": true }, "azure_ai/mistral-large-latest": { @@ -13533,6 +13593,7 @@ "output_cost_per_token": 6e-06, "source": "https://marketplace.microsoft.com/en/marketplace/apps/000-000.mistral-ai-large-2407-offer?tab=Overview", "supports_function_calling": true, + "supports_pdf_input": false, "supports_tool_choice": true }, "azure_ai/mistral-large-3": { @@ -13544,6 +13605,7 @@ "output_cost_per_token": 1.5e-06, "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", "supports_function_calling": true, + "supports_pdf_input": false, "supports_tool_choice": true, "supports_vision": true }, @@ -13557,6 +13619,7 @@ "output_cost_per_token": 2e-06, "supports_assistant_prefill": true, "supports_function_calling": true, + "supports_pdf_input": false, "supports_tool_choice": true }, "azure_ai/mistral-nemo": { @@ -13568,7 +13631,8 @@ "mode": "chat", "output_cost_per_token": 1.5e-07, "source": "https://learn.microsoft.com/en-us/azure/foundry/openai/concepts/retired-models", - "supports_function_calling": true + "supports_function_calling": true, + "supports_pdf_input": false }, "azure_ai/mistral-small": { "deprecation_date": "2025-07-31", @@ -13580,6 +13644,7 @@ "output_cost_per_token": 3e-06, "source": "https://learn.microsoft.com/en-us/azure/foundry/openai/concepts/retired-models", "supports_function_calling": true, + "supports_pdf_input": false, "supports_tool_choice": true }, "azure_ai/mistral-small-2503": { @@ -13591,6 +13656,7 @@ "mode": "chat", "output_cost_per_token": 3e-07, "supports_function_calling": true, + "supports_pdf_input": false, "supports_tool_choice": true, "supports_vision": true }, @@ -52006,6 +52072,7 @@ ], "supports_assistant_prefill": false, "supports_function_calling": true, + "supports_pdf_input": false, "supports_prompt_caching": true, "supports_reasoning": true, "supports_tool_choice": true @@ -52026,6 +52093,7 @@ ], "supports_assistant_prefill": false, "supports_function_calling": true, + "supports_pdf_input": false, "supports_prompt_caching": true, "supports_reasoning": true, "supports_tool_choice": true @@ -52777,6 +52845,7 @@ "output_cost_per_token": 2.5e-06, "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing", "supports_function_calling": true, + "supports_pdf_input": false, "supports_tool_choice": true, "supports_web_search": true }, @@ -52794,6 +52863,7 @@ "global" ], "supports_function_calling": true, + "supports_pdf_input": false, "supports_reasoning": true, "supports_tool_choice": true }, @@ -52811,6 +52881,7 @@ "global" ], "supports_function_calling": true, + "supports_pdf_input": false, "supports_prompt_caching": true, "supports_reasoning": true, "supports_tool_choice": true @@ -52989,6 +53060,7 @@ "output_cost_per_token": 3.6e-07, "output_cost_per_token_batches": 1.8e-07, "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing", + "supports_pdf_input": false, "supports_reasoning": true }, "vertex_ai/openai/gpt-oss-20b-maas": { @@ -53000,6 +53072,7 @@ "mode": "chat", "output_cost_per_token": 2.5e-07, "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing", + "supports_pdf_input": false, "supports_reasoning": true, "cache_read_input_token_cost": 7e-09, "input_cost_per_token_batches": 3.5e-08, @@ -53131,6 +53204,7 @@ "us-south1" ], "supports_function_calling": true, + "supports_pdf_input": false, "supports_tool_choice": true }, "vertex_ai/qwen/qwen3-coder-480b-a35b-instruct-maas": { @@ -53149,6 +53223,7 @@ "global" ], "supports_function_calling": true, + "supports_pdf_input": false, "supports_tool_choice": true }, "vertex_ai/qwen/qwen3-next-80b-a3b-instruct-maas": { @@ -53164,6 +53239,7 @@ "global" ], "supports_function_calling": true, + "supports_pdf_input": false, "supports_tool_choice": true }, "vertex_ai/qwen/qwen3-next-80b-a3b-thinking-maas": { @@ -53179,6 +53255,7 @@ "global" ], "supports_function_calling": true, + "supports_pdf_input": false, "supports_tool_choice": true }, "vertex_ai/veo-3.1-generate-preview": { @@ -67322,6 +67399,7 @@ "text" ], "supports_function_calling": true, + "supports_pdf_input": false, "supports_prompt_caching": true, "supports_reasoning": true, "supports_tool_choice": true, @@ -67336,6 +67414,7 @@ "mode": "chat", "output_cost_per_token": 1.2e-06, "supports_function_calling": true, + "supports_pdf_input": false, "supports_prompt_caching": true, "supports_reasoning": true, "supports_tool_choice": true, @@ -67378,6 +67457,7 @@ "output_cost_per_token": 1.5e-06, "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", "supports_function_calling": true, + "supports_pdf_input": false, "supports_prompt_caching": true, "supports_reasoning": true, "supports_tool_choice": true @@ -67392,6 +67472,7 @@ "output_cost_per_token": 3.1e-07, "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", "supports_function_calling": true, + "supports_pdf_input": false, "supports_prompt_caching": true, "supports_reasoning": true, "supports_tool_choice": true @@ -67406,6 +67487,7 @@ "output_cost_per_token": 5.5e-06, "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", "supports_function_calling": true, + "supports_pdf_input": false, "supports_prompt_caching": true, "supports_reasoning": true, "supports_tool_choice": true @@ -67420,6 +67502,7 @@ "output_cost_per_token": 6.25e-07, "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", "supports_function_calling": true, + "supports_pdf_input": false, "supports_prompt_caching": true, "supports_reasoning": true, "supports_tool_choice": true @@ -67436,6 +67519,7 @@ "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", "supports_function_calling": true, "supports_parallel_function_calling": true, + "supports_pdf_input": false, "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, @@ -67451,6 +67535,7 @@ "output_cost_per_token": 3.2e-06, "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", "supports_function_calling": true, + "supports_pdf_input": false, "supports_reasoning": true, "supports_tool_choice": true }, @@ -67469,6 +67554,7 @@ "supported_output_modalities": [ "text" ], + "supports_pdf_input": false, "supports_response_schema": true, "supports_vision": true }, @@ -67478,7 +67564,8 @@ "output_cost_per_token": 5.4e-06, "litellm_provider": "azure_ai", "mode": "chat", - "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'" + "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", + "supports_pdf_input": false }, "azure_ai/deepseek-v3-0324": { "deprecation_date": "2026-07-13", @@ -67486,7 +67573,8 @@ "output_cost_per_token": 4.56e-06, "litellm_provider": "azure_ai", "mode": "chat", - "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'" + "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", + "supports_pdf_input": false }, "azure_ai/deepseek-v3.1": { "deprecation_date": "2026-07-13", @@ -67494,7 +67582,8 @@ "output_cost_per_token": 4.94e-06, "litellm_provider": "azure_ai", "mode": "chat", - "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'" + "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", + "supports_pdf_input": false }, "azure_ai/grok-3": { "deprecation_date": "2026-05-01", @@ -67502,7 +67591,8 @@ "output_cost_per_token": 1.5e-05, "litellm_provider": "azure_ai", "mode": "chat", - "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'" + "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", + "supports_pdf_input": false }, "azure_ai/grok-3-mini": { "deprecation_date": "2026-05-01", @@ -67510,7 +67600,8 @@ "output_cost_per_token": 1.27e-06, "litellm_provider": "azure_ai", "mode": "chat", - "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'" + "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", + "supports_pdf_input": false }, "azure_ai/grok-4-fast-non-reasoning": { "deprecation_date": "2026-05-01", @@ -67518,7 +67609,8 @@ "output_cost_per_token": 5e-07, "litellm_provider": "azure_ai", "mode": "chat", - "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'" + "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", + "supports_pdf_input": false }, "azure_ai/grok-4-fast-reasoning": { "deprecation_date": "2026-05-01", @@ -67526,7 +67618,8 @@ "output_cost_per_token": 5e-07, "litellm_provider": "azure_ai", "mode": "chat", - "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'" + "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", + "supports_pdf_input": false }, "bedrock/us-gov-west-1/nvidia.nemotron-nano-3-30b": { "input_cost_per_token": 7.2e-08, @@ -73196,6 +73289,7 @@ "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing", "supported_regions": ["global"], "supports_function_calling": true, + "supports_pdf_input": false, "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, @@ -73811,6 +73905,7 @@ "output_cost_per_token": 3.5e-06, "source": "https://azure.microsoft.com/en-us/pricing/details/ai-foundry-models/microsoft/", "supports_function_calling": true, + "supports_pdf_input": false, "supports_prompt_caching": true, "supports_reasoning": true, "supports_tool_choice": true @@ -74865,6 +74960,7 @@ "text" ], "supports_function_calling": true, + "supports_pdf_input": false, "supports_prompt_caching": false, "supports_reasoning": true, "supports_tool_choice": true, @@ -75113,6 +75209,7 @@ "output_cost_per_token_above_200k_tokens": 1.2e-05, "source": "https://docs.x.ai/developers/models/grok-4.7", "supports_function_calling": true, + "supports_pdf_input": false, "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, diff --git a/litellm/proxy/guardrails/guardrail_hooks/prompt_security/prompt_security.py b/litellm/proxy/guardrails/guardrail_hooks/prompt_security/prompt_security.py index 0f66135d038..8e2fea415ef 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/prompt_security/prompt_security.py +++ b/litellm/proxy/guardrails/guardrail_hooks/prompt_security/prompt_security.py @@ -647,9 +647,10 @@ class PromptSecurityGuardrail(CustomGuardrail): async def _process_document_item(self, item: dict, user_api_key_alias: str | None) -> dict: """Process and sanitize document/file items.""" doc_data: Final = item.get("document") or item.get("file") or item + url_key: Final = "file_data" if isinstance(doc_data, dict) and not doc_data.get("url") else "url" if isinstance(doc_data, dict): - url = doc_data.get("url", "") + url = doc_data.get(url_key) or "" doc_content = doc_data.get("data", "") else: url = doc_data if isinstance(doc_data, str) else "" @@ -698,7 +699,7 @@ class PromptSecurityGuardrail(CustomGuardrail): if url.startswith("data:") and header: sanitized_url: Final = f"{header},{sanitized_encoded}" if isinstance(doc_data, dict): - doc_data["url"] = sanitized_url + doc_data[url_key] = sanitized_url elif isinstance(doc_data, dict): doc_data["data"] = sanitized_encoded diff --git a/litellm/types/utils.py b/litellm/types/utils.py index e9984a35ad1..448bd022989 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -4147,7 +4147,12 @@ class CustomPricingLiteLLMParams(MirroredPricingParams): DEPLOYMENT_SCOPED_PRICING_FIELDS: Final[frozenset[str]] = frozenset({"off_peak_pricing"}) DEPLOYMENT_SCOPED_CAPABILITY_FIELDS: Final[frozenset[str]] = frozenset( - {"supports_regex_lookaround", "supports_assistant_prefill", "supports_bedrock_native_structured_output"} + { + "supports_regex_lookaround", + "supports_assistant_prefill", + "supports_bedrock_native_structured_output", + "supports_pdf_input", + } ) DEPLOYMENT_SCOPED_TOKEN_COUNT_FIELDS: Final[frozenset[str]] = frozenset({"token_count_framing", "token_count_model"}) diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 1256fef9181..c6bbb63b6fc 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -4049,6 +4049,7 @@ "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", "supports_function_calling": true, "supports_parallel_function_calling": true, + "supports_pdf_input": false, "supports_response_schema": true, "supports_tool_choice": true }, @@ -12437,7 +12438,8 @@ "supported_output_modalities": [ "text" ], - "supports_native_streaming": true + "supports_native_streaming": true, + "supports_pdf_input": false }, "azure_ai/FLUX-1.1-pro": { "deprecation_date": "2027-01-05", @@ -12502,6 +12504,7 @@ "output_cost_per_token": 1.85e-06, "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", "supports_function_calling": true, + "supports_pdf_input": false, "supports_prompt_caching": true, "supports_reasoning": true, "supports_tool_choice": true @@ -12516,6 +12519,7 @@ "output_cost_per_token": 3.828e-06, "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", "supports_function_calling": true, + "supports_pdf_input": false, "supports_prompt_caching": true, "supports_reasoning": true, "supports_tool_choice": true @@ -12531,6 +12535,7 @@ "output_cost_per_token": 3.52e-06, "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", "supports_function_calling": true, + "supports_pdf_input": false, "supports_prompt_caching": true, "supports_reasoning": true, "supports_tool_choice": true @@ -12546,6 +12551,7 @@ "output_cost_per_token": 4.84e-06, "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", "supports_function_calling": true, + "supports_pdf_input": false, "supports_prompt_caching": true, "supports_reasoning": true, "supports_tool_choice": true @@ -12560,6 +12566,7 @@ "output_cost_per_token": 4.84e-06, "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", "supports_function_calling": true, + "supports_pdf_input": false, "supports_prompt_caching": true, "supports_reasoning": true, "supports_tool_choice": true @@ -12574,6 +12581,7 @@ "output_cost_per_token": 7.26e-06, "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", "supports_function_calling": true, + "supports_pdf_input": false, "supports_prompt_caching": true, "supports_reasoning": true, "supports_tool_choice": true @@ -12594,6 +12602,7 @@ "text" ], "supports_function_calling": true, + "supports_pdf_input": false, "supports_prompt_caching": true, "supports_reasoning": true, "supports_tool_choice": true @@ -12616,6 +12625,7 @@ "text" ], "supports_function_calling": true, + "supports_pdf_input": false, "supports_prompt_caching": true, "supports_reasoning": true, "supports_tool_choice": true, @@ -12638,6 +12648,7 @@ "text" ], "supports_function_calling": true, + "supports_pdf_input": false, "supports_prompt_caching": true, "supports_reasoning": true, "supports_tool_choice": true, @@ -12660,6 +12671,7 @@ "text" ], "supports_function_calling": true, + "supports_pdf_input": false, "supports_prompt_caching": true, "supports_reasoning": true, "supports_tool_choice": true, @@ -12687,6 +12699,7 @@ "text" ], "supports_function_calling": true, + "supports_pdf_input": false, "supports_prompt_caching": true, "supports_reasoning": true, "supports_tool_choice": true, @@ -12703,6 +12716,7 @@ "output_cost_per_token": 1.32e-06, "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", "supports_function_calling": true, + "supports_pdf_input": false, "supports_prompt_caching": true, "supports_reasoning": true, "supports_tool_choice": true @@ -12724,6 +12738,7 @@ "text" ], "supports_function_calling": true, + "supports_pdf_input": false, "supports_prompt_caching": true, "supports_reasoning": true, "supports_tool_choice": true, @@ -12744,6 +12759,7 @@ "text" ], "supports_function_calling": true, + "supports_pdf_input": false, "supports_prompt_caching": true, "supports_reasoning": true, "supports_tool_choice": true, @@ -12765,6 +12781,7 @@ "text" ], "supports_function_calling": true, + "supports_pdf_input": false, "supports_prompt_caching": true, "supports_reasoning": true, "supports_tool_choice": true @@ -12833,6 +12850,7 @@ "text" ], "supports_function_calling": true, + "supports_pdf_input": false, "supports_prompt_caching": true, "supports_reasoning": true, "supports_tool_choice": true @@ -12847,6 +12865,7 @@ "output_cost_per_token": 7.1e-07, "source": "https://learn.microsoft.com/en-us/azure/foundry/foundry-models/concepts/models-sold-directly-by-azure", "supports_function_calling": true, + "supports_pdf_input": false, "supports_tool_choice": true }, "azure_ai/Llama-4-Maverick-17B-128E-Instruct-FP8": { @@ -12859,6 +12878,7 @@ "output_cost_per_token": 1e-06, "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", "supports_function_calling": true, + "supports_pdf_input": false, "supports_tool_choice": true, "supports_vision": true }, @@ -12872,6 +12892,7 @@ "output_cost_per_token": 7.8e-07, "source": "https://azure.microsoft.com/en-us/blog/introducing-the-llama-4-herd-in-azure-ai-foundry-and-azure-databricks/", "supports_function_calling": true, + "supports_pdf_input": false, "supports_tool_choice": true, "supports_vision": true }, @@ -12884,6 +12905,7 @@ "mode": "chat", "output_cost_per_token": 3.7e-07, "source": "https://learn.microsoft.com/en-us/azure/foundry/openai/concepts/retired-models", + "supports_pdf_input": false, "supports_tool_choice": true }, "azure_ai/Meta-Llama-3.1-70B-Instruct": { @@ -12895,6 +12917,7 @@ "mode": "chat", "output_cost_per_token": 3.54e-06, "source": "https://learn.microsoft.com/en-us/azure/foundry/openai/concepts/retired-models", + "supports_pdf_input": false, "supports_tool_choice": true }, "azure_ai/Phi-3-medium-128k-instruct": { @@ -12906,6 +12929,7 @@ "mode": "chat", "output_cost_per_token": 6.8e-07, "source": "https://learn.microsoft.com/en-us/azure/foundry/openai/concepts/retired-models", + "supports_pdf_input": false, "supports_tool_choice": true, "supports_vision": false }, @@ -12918,6 +12942,7 @@ "mode": "chat", "output_cost_per_token": 6.8e-07, "source": "https://learn.microsoft.com/en-us/azure/foundry/openai/concepts/retired-models", + "supports_pdf_input": false, "supports_tool_choice": true, "supports_vision": false }, @@ -12930,6 +12955,7 @@ "mode": "chat", "output_cost_per_token": 5.2e-07, "source": "https://learn.microsoft.com/en-us/azure/foundry/openai/concepts/retired-models", + "supports_pdf_input": false, "supports_tool_choice": true, "supports_vision": false }, @@ -12942,6 +12968,7 @@ "mode": "chat", "output_cost_per_token": 5.2e-07, "source": "https://learn.microsoft.com/en-us/azure/foundry/openai/concepts/retired-models", + "supports_pdf_input": false, "supports_tool_choice": true, "supports_vision": false }, @@ -12954,6 +12981,7 @@ "mode": "chat", "output_cost_per_token": 6e-07, "source": "https://learn.microsoft.com/en-us/azure/foundry/openai/concepts/retired-models", + "supports_pdf_input": false, "supports_tool_choice": true, "supports_vision": false }, @@ -12966,6 +12994,7 @@ "mode": "chat", "output_cost_per_token": 6e-07, "source": "https://learn.microsoft.com/en-us/azure/foundry/openai/concepts/retired-models", + "supports_pdf_input": false, "supports_tool_choice": true, "supports_vision": false }, @@ -12978,6 +13007,7 @@ "mode": "chat", "output_cost_per_token": 6.4e-07, "source": "https://learn.microsoft.com/en-us/azure/foundry/openai/concepts/retired-models", + "supports_pdf_input": false, "supports_tool_choice": true, "supports_vision": false }, @@ -12990,6 +13020,7 @@ "mode": "chat", "output_cost_per_token": 5.2e-07, "source": "https://learn.microsoft.com/en-us/azure/foundry/openai/concepts/retired-models", + "supports_pdf_input": false, "supports_tool_choice": true, "supports_vision": false }, @@ -13002,6 +13033,7 @@ "mode": "chat", "output_cost_per_token": 5.2e-07, "source": "https://learn.microsoft.com/en-us/azure/foundry/openai/concepts/retired-models", + "supports_pdf_input": false, "supports_tool_choice": true, "supports_vision": true }, @@ -13015,6 +13047,7 @@ "output_cost_per_token": 5e-07, "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", "supports_function_calling": true, + "supports_pdf_input": false, "supports_tool_choice": true, "supports_vision": false }, @@ -13027,7 +13060,8 @@ "mode": "chat", "output_cost_per_token": 3e-07, "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", - "supports_function_calling": true + "supports_function_calling": true, + "supports_pdf_input": false }, "azure_ai/Phi-4-multimodal-instruct": { "deprecation_date": "2027-01-17", @@ -13041,6 +13075,7 @@ "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", "supports_audio_input": true, "supports_function_calling": true, + "supports_pdf_input": false, "supports_vision": true }, "azure_ai/Phi-4-mini-reasoning": { @@ -13052,7 +13087,8 @@ "mode": "chat", "output_cost_per_token": 3e-07, "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", - "supports_function_calling": true + "supports_function_calling": true, + "supports_pdf_input": false }, "azure_ai/Phi-4-reasoning": { "deprecation_date": "2027-01-17", @@ -13064,6 +13100,7 @@ "output_cost_per_token": 5e-07, "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", "supports_function_calling": true, + "supports_pdf_input": false, "supports_tool_choice": true, "supports_reasoning": true }, @@ -13106,6 +13143,7 @@ "output_cost_per_token": 1e-05, "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", "supports_function_calling": true, + "supports_pdf_input": false, "supports_tool_choice": true }, "azure_ai/doc-intelligence/prebuilt-read": { @@ -13144,6 +13182,7 @@ "mode": "chat", "output_cost_per_token": 5.4e-06, "source": "https://learn.microsoft.com/en-us/azure/foundry/openai/concepts/retired-models", + "supports_pdf_input": false, "supports_reasoning": true, "supports_tool_choice": true }, @@ -13202,6 +13241,7 @@ "source": "https://learn.microsoft.com/en-us/azure/foundry/foundry-models/concepts/models-sold-directly-by-azure", "supports_assistant_prefill": true, "supports_function_calling": true, + "supports_pdf_input": false, "supports_prompt_caching": true, "supports_reasoning": true, "supports_tool_choice": true @@ -13217,6 +13257,7 @@ "source": "https://learn.microsoft.com/en-us/azure/foundry/foundry-models/concepts/models-sold-directly-by-azure", "supports_assistant_prefill": true, "supports_function_calling": true, + "supports_pdf_input": false, "supports_prompt_caching": true, "supports_reasoning": true, "supports_tool_choice": true @@ -13230,6 +13271,7 @@ "mode": "chat", "output_cost_per_token": 4.56e-06, "source": "https://learn.microsoft.com/en-us/azure/foundry/openai/concepts/retired-models", + "supports_pdf_input": false, "supports_tool_choice": true }, "azure_ai/deepseek-v4-pro": { @@ -13242,6 +13284,7 @@ "output_cost_per_token": 3.48e-06, "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", "supports_function_calling": true, + "supports_pdf_input": false, "supports_reasoning": true, "supports_tool_choice": true, "cache_read_input_token_cost": 1.45e-07, @@ -13257,6 +13300,7 @@ "output_cost_per_token": 5.1e-07, "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", "supports_function_calling": true, + "supports_pdf_input": false, "supports_reasoning": true, "supports_tool_choice": true, "cache_read_input_token_cost": 2.8e-08, @@ -13273,6 +13317,7 @@ "output_cost_per_token": 1.32e-06, "source": "https://learn.microsoft.com/en-us/azure/foundry/foundry-models/concepts/models-sold-directly-by-azure", "supports_function_calling": true, + "supports_pdf_input": false, "supports_prompt_caching": true, "supports_reasoning": true, "supports_tool_choice": true @@ -13306,6 +13351,7 @@ "output_cost_per_token": 1.5e-05, "source": "https://learn.microsoft.com/en-us/azure/foundry/foundry-models/concepts/models-sold-directly-by-azure", "supports_function_calling": true, + "supports_pdf_input": false, "supports_response_schema": true, "supports_tool_choice": true, "supports_web_search": true @@ -13323,6 +13369,7 @@ "output_cost_per_token_above_200k_tokens": 5e-06, "source": "https://learn.microsoft.com/en-us/azure/foundry/foundry-models/concepts/models-sold-directly-by-azure", "supports_function_calling": true, + "supports_pdf_input": false, "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, @@ -13343,6 +13390,7 @@ "output_cost_per_token_above_200k_tokens": 1.2e-05, "source": "https://azure.microsoft.com/en-us/pricing/details/ai-foundry-models/grok/", "supports_function_calling": true, + "supports_pdf_input": false, "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, @@ -13361,6 +13409,7 @@ "output_cost_per_token": 2.5e-06, "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", "supports_function_calling": true, + "supports_pdf_input": false, "supports_prompt_caching": true, "supports_response_schema": true, "supports_tool_choice": true, @@ -13379,6 +13428,7 @@ "output_cost_per_token": 2.5e-06, "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", "supports_function_calling": true, + "supports_pdf_input": false, "supports_prompt_caching": true, "supports_response_schema": true, "supports_tool_choice": true, @@ -13395,6 +13445,7 @@ "mode": "chat", "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", "supports_function_calling": true, + "supports_pdf_input": false, "supports_response_schema": true, "supports_tool_choice": true, "supports_web_search": true @@ -13409,6 +13460,7 @@ "mode": "chat", "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", "supports_function_calling": true, + "supports_pdf_input": false, "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, @@ -13424,6 +13476,7 @@ "output_cost_per_token": 1.5e-06, "source": "https://learn.microsoft.com/en-us/azure/foundry/foundry-models/concepts/models-sold-directly-by-azure", "supports_function_calling": true, + "supports_pdf_input": false, "supports_response_schema": true, "supports_tool_choice": true, "supports_web_search": true @@ -13436,7 +13489,8 @@ "max_output_tokens": 8192, "mode": "chat", "output_cost_per_token": 0.00971, - "source": "https://learn.microsoft.com/en-us/azure/foundry/openai/concepts/retired-models" + "source": "https://learn.microsoft.com/en-us/azure/foundry/openai/concepts/retired-models", + "supports_pdf_input": false }, "azure_ai/jamba-instruct": { "deprecation_date": "2025-03-01", @@ -13447,6 +13501,7 @@ "mode": "chat", "output_cost_per_token": 7e-07, "source": "https://learn.microsoft.com/en-us/azure/foundry/openai/concepts/retired-models", + "supports_pdf_input": false, "supports_tool_choice": true }, "azure_ai/kimi-k2.5": { @@ -13459,6 +13514,7 @@ "output_cost_per_token": 3e-06, "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", "supports_function_calling": true, + "supports_pdf_input": false, "supports_tool_choice": true, "supports_video_input": true, "supports_vision": true, @@ -13482,6 +13538,7 @@ "text" ], "supports_function_calling": true, + "supports_pdf_input": false, "supports_reasoning": true, "supports_tool_choice": true, "supports_vision": true, @@ -13498,6 +13555,7 @@ "output_cost_per_token": 4e-08, "source": "https://marketplace.microsoft.com/en/marketplace/apps/000-000.ministral-3b-2410-offer?tab=Overview", "supports_function_calling": true, + "supports_pdf_input": false, "supports_tool_choice": true }, "azure_ai/mistral-large": { @@ -13510,6 +13568,7 @@ "output_cost_per_token": 1.2e-05, "source": "https://learn.microsoft.com/en-us/azure/foundry/openai/concepts/retired-models", "supports_function_calling": true, + "supports_pdf_input": false, "supports_tool_choice": true }, "azure_ai/mistral-large-2407": { @@ -13522,6 +13581,7 @@ "output_cost_per_token": 6e-06, "source": "https://learn.microsoft.com/en-us/azure/foundry/openai/concepts/retired-models", "supports_function_calling": true, + "supports_pdf_input": false, "supports_tool_choice": true }, "azure_ai/mistral-large-latest": { @@ -13533,6 +13593,7 @@ "output_cost_per_token": 6e-06, "source": "https://marketplace.microsoft.com/en/marketplace/apps/000-000.mistral-ai-large-2407-offer?tab=Overview", "supports_function_calling": true, + "supports_pdf_input": false, "supports_tool_choice": true }, "azure_ai/mistral-large-3": { @@ -13544,6 +13605,7 @@ "output_cost_per_token": 1.5e-06, "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", "supports_function_calling": true, + "supports_pdf_input": false, "supports_tool_choice": true, "supports_vision": true }, @@ -13557,6 +13619,7 @@ "output_cost_per_token": 2e-06, "supports_assistant_prefill": true, "supports_function_calling": true, + "supports_pdf_input": false, "supports_tool_choice": true }, "azure_ai/mistral-nemo": { @@ -13568,7 +13631,8 @@ "mode": "chat", "output_cost_per_token": 1.5e-07, "source": "https://learn.microsoft.com/en-us/azure/foundry/openai/concepts/retired-models", - "supports_function_calling": true + "supports_function_calling": true, + "supports_pdf_input": false }, "azure_ai/mistral-small": { "deprecation_date": "2025-07-31", @@ -13580,6 +13644,7 @@ "output_cost_per_token": 3e-06, "source": "https://learn.microsoft.com/en-us/azure/foundry/openai/concepts/retired-models", "supports_function_calling": true, + "supports_pdf_input": false, "supports_tool_choice": true }, "azure_ai/mistral-small-2503": { @@ -13591,6 +13656,7 @@ "mode": "chat", "output_cost_per_token": 3e-07, "supports_function_calling": true, + "supports_pdf_input": false, "supports_tool_choice": true, "supports_vision": true }, @@ -52006,6 +52072,7 @@ ], "supports_assistant_prefill": false, "supports_function_calling": true, + "supports_pdf_input": false, "supports_prompt_caching": true, "supports_reasoning": true, "supports_tool_choice": true @@ -52026,6 +52093,7 @@ ], "supports_assistant_prefill": false, "supports_function_calling": true, + "supports_pdf_input": false, "supports_prompt_caching": true, "supports_reasoning": true, "supports_tool_choice": true @@ -52777,6 +52845,7 @@ "output_cost_per_token": 2.5e-06, "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing", "supports_function_calling": true, + "supports_pdf_input": false, "supports_tool_choice": true, "supports_web_search": true }, @@ -52794,6 +52863,7 @@ "global" ], "supports_function_calling": true, + "supports_pdf_input": false, "supports_reasoning": true, "supports_tool_choice": true }, @@ -52811,6 +52881,7 @@ "global" ], "supports_function_calling": true, + "supports_pdf_input": false, "supports_prompt_caching": true, "supports_reasoning": true, "supports_tool_choice": true @@ -52989,6 +53060,7 @@ "output_cost_per_token": 3.6e-07, "output_cost_per_token_batches": 1.8e-07, "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing", + "supports_pdf_input": false, "supports_reasoning": true }, "vertex_ai/openai/gpt-oss-20b-maas": { @@ -53000,6 +53072,7 @@ "mode": "chat", "output_cost_per_token": 2.5e-07, "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing", + "supports_pdf_input": false, "supports_reasoning": true, "cache_read_input_token_cost": 7e-09, "input_cost_per_token_batches": 3.5e-08, @@ -53131,6 +53204,7 @@ "us-south1" ], "supports_function_calling": true, + "supports_pdf_input": false, "supports_tool_choice": true }, "vertex_ai/qwen/qwen3-coder-480b-a35b-instruct-maas": { @@ -53149,6 +53223,7 @@ "global" ], "supports_function_calling": true, + "supports_pdf_input": false, "supports_tool_choice": true }, "vertex_ai/qwen/qwen3-next-80b-a3b-instruct-maas": { @@ -53164,6 +53239,7 @@ "global" ], "supports_function_calling": true, + "supports_pdf_input": false, "supports_tool_choice": true }, "vertex_ai/qwen/qwen3-next-80b-a3b-thinking-maas": { @@ -53179,6 +53255,7 @@ "global" ], "supports_function_calling": true, + "supports_pdf_input": false, "supports_tool_choice": true }, "vertex_ai/veo-3.1-generate-preview": { @@ -67322,6 +67399,7 @@ "text" ], "supports_function_calling": true, + "supports_pdf_input": false, "supports_prompt_caching": true, "supports_reasoning": true, "supports_tool_choice": true, @@ -67336,6 +67414,7 @@ "mode": "chat", "output_cost_per_token": 1.2e-06, "supports_function_calling": true, + "supports_pdf_input": false, "supports_prompt_caching": true, "supports_reasoning": true, "supports_tool_choice": true, @@ -67378,6 +67457,7 @@ "output_cost_per_token": 1.5e-06, "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", "supports_function_calling": true, + "supports_pdf_input": false, "supports_prompt_caching": true, "supports_reasoning": true, "supports_tool_choice": true @@ -67392,6 +67472,7 @@ "output_cost_per_token": 3.1e-07, "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", "supports_function_calling": true, + "supports_pdf_input": false, "supports_prompt_caching": true, "supports_reasoning": true, "supports_tool_choice": true @@ -67406,6 +67487,7 @@ "output_cost_per_token": 5.5e-06, "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", "supports_function_calling": true, + "supports_pdf_input": false, "supports_prompt_caching": true, "supports_reasoning": true, "supports_tool_choice": true @@ -67420,6 +67502,7 @@ "output_cost_per_token": 6.25e-07, "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", "supports_function_calling": true, + "supports_pdf_input": false, "supports_prompt_caching": true, "supports_reasoning": true, "supports_tool_choice": true @@ -67436,6 +67519,7 @@ "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", "supports_function_calling": true, "supports_parallel_function_calling": true, + "supports_pdf_input": false, "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, @@ -67451,6 +67535,7 @@ "output_cost_per_token": 3.2e-06, "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", "supports_function_calling": true, + "supports_pdf_input": false, "supports_reasoning": true, "supports_tool_choice": true }, @@ -67469,6 +67554,7 @@ "supported_output_modalities": [ "text" ], + "supports_pdf_input": false, "supports_response_schema": true, "supports_vision": true }, @@ -67478,7 +67564,8 @@ "output_cost_per_token": 5.4e-06, "litellm_provider": "azure_ai", "mode": "chat", - "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'" + "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", + "supports_pdf_input": false }, "azure_ai/deepseek-v3-0324": { "deprecation_date": "2026-07-13", @@ -67486,7 +67573,8 @@ "output_cost_per_token": 4.56e-06, "litellm_provider": "azure_ai", "mode": "chat", - "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'" + "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", + "supports_pdf_input": false }, "azure_ai/deepseek-v3.1": { "deprecation_date": "2026-07-13", @@ -67494,7 +67582,8 @@ "output_cost_per_token": 4.94e-06, "litellm_provider": "azure_ai", "mode": "chat", - "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'" + "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", + "supports_pdf_input": false }, "azure_ai/grok-3": { "deprecation_date": "2026-05-01", @@ -67502,7 +67591,8 @@ "output_cost_per_token": 1.5e-05, "litellm_provider": "azure_ai", "mode": "chat", - "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'" + "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", + "supports_pdf_input": false }, "azure_ai/grok-3-mini": { "deprecation_date": "2026-05-01", @@ -67510,7 +67600,8 @@ "output_cost_per_token": 1.27e-06, "litellm_provider": "azure_ai", "mode": "chat", - "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'" + "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", + "supports_pdf_input": false }, "azure_ai/grok-4-fast-non-reasoning": { "deprecation_date": "2026-05-01", @@ -67518,7 +67609,8 @@ "output_cost_per_token": 5e-07, "litellm_provider": "azure_ai", "mode": "chat", - "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'" + "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", + "supports_pdf_input": false }, "azure_ai/grok-4-fast-reasoning": { "deprecation_date": "2026-05-01", @@ -67526,7 +67618,8 @@ "output_cost_per_token": 5e-07, "litellm_provider": "azure_ai", "mode": "chat", - "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'" + "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", + "supports_pdf_input": false }, "bedrock/us-gov-west-1/nvidia.nemotron-nano-3-30b": { "input_cost_per_token": 7.2e-08, @@ -73196,6 +73289,7 @@ "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing", "supported_regions": ["global"], "supports_function_calling": true, + "supports_pdf_input": false, "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, @@ -73811,6 +73905,7 @@ "output_cost_per_token": 3.5e-06, "source": "https://azure.microsoft.com/en-us/pricing/details/ai-foundry-models/microsoft/", "supports_function_calling": true, + "supports_pdf_input": false, "supports_prompt_caching": true, "supports_reasoning": true, "supports_tool_choice": true @@ -74865,6 +74960,7 @@ "text" ], "supports_function_calling": true, + "supports_pdf_input": false, "supports_prompt_caching": false, "supports_reasoning": true, "supports_tool_choice": true, @@ -75113,6 +75209,7 @@ "output_cost_per_token_above_200k_tokens": 1.2e-05, "source": "https://docs.x.ai/developers/models/grok-4.7", "supports_function_calling": true, + "supports_pdf_input": false, "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, diff --git a/tests/integration/messages_endpoint/chat_bridge/test_anthropic_messages_pdf_document_wire.py b/tests/integration/messages_endpoint/chat_bridge/test_anthropic_messages_pdf_document_wire.py new file mode 100644 index 00000000000..130bab24f4c --- /dev/null +++ b/tests/integration/messages_endpoint/chat_bridge/test_anthropic_messages_pdf_document_wire.py @@ -0,0 +1,421 @@ +import json +import uuid +from collections.abc import Callable, Sequence +from dataclasses import dataclass +from pathlib import Path +from typing import Final + +import anthropic +import pytest +from integration._support.client import Gateway, Scenario, string_value +from integration._support.pdf_document import LETTER, Page, document, encoded, pdf_bytes +from integration._support.process import owned_proxy_process +from integration._support.wire import Reply, Request, Wire, wire_server +from pydantic import JsonValue, TypeAdapter + +_API_KEY: Final = "synthetic-pdf-document-key" +_ANSWER: Final = "The codeword is PELICAN-42." +_TITLE: Final = "Quarterly report" +_URL: Final = "https://example.com/brief.pdf" +_SAP_TOKEN_TARGET: Final = "/oauth/token" +_JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue]) +_PARTS: Final = TypeAdapter(list[dict[str, JsonValue]]) + + +def _pdf(text: str) -> str: + return encoded(pdf_bytes((Page(text=text), LETTER))) + + +def _chat_completion(identity: str, model: str) -> Reply: + return Reply( + body=json.dumps( + { + "id": identity, + "object": "chat.completion", + "created": 1, + "model": model, + "choices": [ + {"index": 0, "message": {"role": "assistant", "content": _ANSWER}, "finish_reason": "stop"} + ], + "usage": {"prompt_tokens": 914, "completion_tokens": 9, "total_tokens": 923}, + } + ).encode() + ) + + +def _only_body(wire: Wire) -> dict[str, JsonValue]: + received: Final = tuple( + request + for request in wire.drain() + if request.method == "POST" and not request.target.startswith(_SAP_TOKEN_TARGET) + ) + assert len(received) == 1, received + return _JSON_OBJECT.validate_json(received[0].body) + + +def _user_parts(body: dict[str, JsonValue], key: str) -> list[dict[str, JsonValue]]: + turns: Final = _PARTS.validate_python(body[key]) + assert len(turns) == 1, turns + return _PARTS.validate_python(turns[0]["content" if key == "messages" else "parts"]) + + +def test_anthropic_sdk_pdf_documents_reach_azure_as_file_parts_named_by_title(gateway: Gateway) -> None: + marker: Final = uuid.uuid4().hex + backend: Final = f"gpt-pdf-{marker[:8]}" + identity: Final = f"chatcmpl-{marker}" + titled: Final = _pdf(f"PELICAN-42 {marker}") + untitled: Final = _pdf(f"ORCHID-77 {marker}") + question: Final = f"What codeword does each document carry? {marker}" + with wire_server(lambda _: _chat_completion(identity, backend)) as wire, gateway.scenario() as scenario: + model: Final = scenario.model( + model=f"azure/{backend}", api_base=wire.url, api_key=_API_KEY, api_version="2024-10-21" + ) + client: Final = anthropic.Anthropic(base_url=str(gateway.client.base_url), api_key=gateway.key, max_retries=0) + message: Final = client.messages.create( + model=model, + max_tokens=64, + messages=[ + { + "role": "user", + "content": [ + { + "type": "document", + "title": _TITLE, + "source": {"type": "base64", "media_type": "application/pdf", "data": titled}, + }, + { + "type": "document", + "source": {"type": "base64", "media_type": "application/pdf", "data": untitled}, + }, + {"type": "document", "source": {"type": "url", "url": _URL}}, + {"type": "text", "text": question}, + ], + } + ], + ) + assert message.id == identity + assert [(block.type, getattr(block, "text", None)) for block in message.content] == [("text", _ANSWER)] + assert _user_parts(_only_body(wire), "messages") == [ + {"type": "file", "file": {"filename": _TITLE, "file_data": f"data:application/pdf;base64,{titled}"}}, + { + "type": "file", + "file": {"filename": "document.pdf", "file_data": f"data:application/pdf;base64,{untitled}"}, + }, + {"type": "image_url", "image_url": {"url": _URL}}, + {"type": "text", "text": question}, + ] + + +@dataclass(frozen=True, slots=True) +class _Reader: + deploy: Callable[[Scenario, Wire], str] + reply: Callable[[Request], Reply] + pdf_parts: Callable[[dict[str, JsonValue]], Sequence[JsonValue]] + + +def _gemini_reply(_: Request) -> Reply: + return Reply( + body=json.dumps( + { + "candidates": [ + {"content": {"parts": [{"text": _ANSWER}], "role": "model"}, "finishReason": "STOP", "index": 0} + ], + "usageMetadata": {"promptTokenCount": 300, "candidatesTokenCount": 9, "totalTokenCount": 309}, + "modelVersion": "gemini-2.5-flash", + } + ).encode() + ) + + +def _bedrock_reply(_: Request) -> Reply: + return Reply( + body=json.dumps( + { + "output": {"message": {"role": "assistant", "content": [{"text": _ANSWER}]}}, + "stopReason": "end_turn", + "usage": {"inputTokens": 300, "outputTokens": 9, "totalTokens": 309}, + "metrics": {"latencyMs": 1}, + } + ).encode() + ) + + +def _anthropic_reply(_: Request) -> Reply: + return Reply( + body=json.dumps( + { + "id": "msg_pdf", + "type": "message", + "role": "assistant", + "model": "claude-sonnet-4-5", + "content": [{"type": "text", "text": _ANSWER}], + "stop_reason": "end_turn", + "usage": {"input_tokens": 300, "output_tokens": 9}, + } + ).encode() + ) + + +def _sap_reply(request: Request) -> Reply: + if request.target.startswith(_SAP_TOKEN_TARGET): + return Reply(body=json.dumps({"access_token": "scripted", "expires_in": 3600}).encode()) + match request.target: + case "/v2/lm/deployments": + deployment: Final = { + "scenarioId": "orchestration", + "configurationId": "orchestration-config", + "deploymentUrl": f"http://{request.headers['host']}/sap-deployment", + "createdAt": "2026-01-01T00:00:00Z", + } + return Reply(body=json.dumps({"resources": [deployment]}).encode()) + case "/v2/lm/configurations/orchestration-config": + return Reply(body=json.dumps({"executableId": "orchestration"}).encode()) + case _: + completion: Final = _JSON_OBJECT.validate_json(_chat_completion("chatcmpl-sap-pdf", "claude").body) + return Reply(body=json.dumps({"final_result": completion}).encode()) + + +def _snowflake(scenario: Scenario, wire: Wire) -> str: + return scenario.model(model="snowflake/claude-sonnet-4-5", api_base=wire.url, api_key=_API_KEY) + + +def _sap(scenario: Scenario, wire: Wire) -> str: + service_key: Final = { + "clientid": "scripted", + "clientsecret": _API_KEY, + "url": wire.url, + "serviceurls": {"AI_API_URL": wire.url}, + } + return scenario.model(model="sap/anthropic--claude-4.5-sonnet", api_base=wire.url, api_key=json.dumps(service_key)) + + +def _snowflake_pdf_parts(body: dict[str, JsonValue]) -> Sequence[JsonValue]: + return tuple(part.get("source") for part in _user_parts(body, "messages") if part.get("type") == "document") + + +def _sap_pdf_parts(body: dict[str, JsonValue]) -> Sequence[JsonValue]: + config: Final = _JSON_OBJECT.validate_python(body["config"]) + templating: Final = _JSON_OBJECT.validate_python( + _JSON_OBJECT.validate_python(config["modules"])["prompt_templating"] + ) + prompt: Final = _JSON_OBJECT.validate_python(templating["prompt"]) + turns: Final = _PARTS.validate_python(prompt["template"]) + assert len(turns) == 1, turns + parts: Final = _PARTS.validate_python(turns[0]["content"]) + return tuple( + _JSON_OBJECT.validate_python(part["image_url"]).get("url") for part in parts if part.get("type") == "image_url" + ) + + +def _gemini(scenario: Scenario, wire: Wire) -> str: + return scenario.model(model="gemini/gemini-2.5-flash", api_base=wire.url, api_key=_API_KEY) + + +def _bedrock(scenario: Scenario, wire: Wire) -> str: + return scenario.model( + model="bedrock/converse/us.amazon.nova-2-lite-v1:0", + api_key=None, + aws_access_key_id="AKIDSYNTHETIC", + aws_secret_access_key="synthetic-secret", + aws_region_name="us-east-1", + aws_bedrock_runtime_endpoint=wire.url, + ) + + +def _gemini_pdf_parts(body: dict[str, JsonValue]) -> Sequence[JsonValue]: + return tuple(part.get("inline_data") for part in _user_parts(body, "contents") if "inline_data" in part) + + +def _bedrock_pdf_parts(body: dict[str, JsonValue]) -> Sequence[JsonValue]: + documents: Final = tuple( + _JSON_OBJECT.validate_python(part["document"]) for part in _user_parts(body, "messages") if "document" in part + ) + return tuple({"format": item.get("format"), "source": item.get("source")} for item in documents) + + +@pytest.mark.parametrize( + ("reader", "expected"), + [ + pytest.param( + _Reader(_gemini, _gemini_reply, _gemini_pdf_parts), + lambda data: [{"mime_type": "application/pdf", "data": data}], + id="gemini", + ), + pytest.param( + _Reader(_bedrock, _bedrock_reply, _bedrock_pdf_parts), + lambda data: [{"format": "pdf", "source": {"bytes": data}}], + id="bedrock-converse", + ), + pytest.param( + _Reader(_snowflake, _anthropic_reply, _snowflake_pdf_parts), + lambda data: [{"type": "base64", "media_type": "application/pdf", "data": data}], + id="snowflake-claude", + ), + pytest.param( + _Reader(_sap, _sap_reply, _sap_pdf_parts), + lambda data: [f"data:application/pdf;base64,{data}"], + id="sap-orchestration", + ), + ], +) +def test_messages_pdf_still_reaches_providers_that_read_pdfs_in_their_own_document_shape( + gateway: Gateway, reader: _Reader, expected: Callable[[str], Sequence[JsonValue]] +) -> None: + marker: Final = uuid.uuid4().hex + data: Final = _pdf(f"PELICAN-42 {marker}") + with wire_server(reader.reply) as wire, gateway.scenario() as scenario: + model: Final = reader.deploy(scenario, wire) + response: Final = gateway.request( + "POST", + "/v1/messages", + { + "model": model, + "max_tokens": 64, + "messages": [ + { + "role": "user", + "content": [ + document({"type": "base64", "media_type": "application/pdf", "data": data}, title=_TITLE), + {"type": "text", "text": f"What is the codeword? {marker}"}, + ], + } + ], + }, + headers={"anthropic-version": "2023-06-01"}, + ) + assert response.status_code == 200, response.text + assert _JSON_OBJECT.validate_json(response.content)["content"] == [{"type": "text", "text": _ANSWER}] + assert list(reader.pdf_parts(_only_body(wire))) == list(expected(data)) + + +def _prompt_security(job: str) -> Callable[[Request], Reply]: + def respond(request: Request) -> Reply: + match (request.method, request.target.partition("?")[0]): + case ("POST", "/api/sanitizeFile"): + return Reply(body=json.dumps({"jobId": job}).encode()) + case ("GET", "/api/sanitizeFile"): + return Reply( + body=json.dumps( + { + "status": "done", + "content": "", + "metadata": {"action": "block", "violations": ["Malicious File"]}, + } + ).encode() + ) + case _: + return Reply(body=json.dumps({"result": {"prompt": {"action": "allow"}}}).encode()) + + return respond + + +def _delete_guardrail(gateway: Gateway, identity: str) -> None: + response: Final = gateway.request("DELETE", f"/guardrails/{identity}") + assert response.status_code == 200, response.text + + +def test_prompt_security_still_sanitizes_a_messages_pdf_and_blocks_it_before_the_provider(gateway: Gateway) -> None: + marker: Final = uuid.uuid4().hex + pdf: Final = pdf_bytes((Page(text=f"PELICAN-42 {marker}"), LETTER)) + with ( + wire_server(lambda _: _chat_completion(f"chatcmpl-{marker}", "gpt-pdf")) as provider, + wire_server(_prompt_security(f"job-{marker}")) as sink, + gateway.scenario() as scenario, + ): + model: Final = scenario.model( + model=f"azure/gpt-pdf-{marker[:8]}", api_base=provider.url, api_key=_API_KEY, api_version="2024-10-21" + ) + name: Final = f"pdf-prompt-security-{marker}" + created: Final = gateway.post( + "/guardrails", + { + "guardrail": { + "guardrail_name": name, + "litellm_params": { + "guardrail": "prompt_security", + "mode": "pre_call", + "api_key": "synthetic-prompt-security-key", + "api_base": sink.url, + }, + } + }, + ) + scenario.cleanups.callback(_delete_guardrail, gateway, string_value(created["guardrail_id"])) + response: Final = gateway.request( + "POST", + "/v1/messages", + { + "model": model, + "max_tokens": 64, + "guardrails": [name], + "messages": [ + { + "role": "user", + "content": [ + document({"type": "base64", "media_type": "application/pdf", "data": encoded(pdf)}), + {"type": "text", "text": f"Summarize this PDF {marker}"}, + ], + } + ], + }, + headers={"anthropic-version": "2023-06-01"}, + ) + assert response.status_code == 400, response.text + assert "Malicious File" in response.text + assert provider.drain() == () + uploads: Final = tuple( + request for request in sink.drain() if (request.method, request.target) == ("POST", "/api/sanitizeFile") + ) + assert len(uploads) == 1, uploads + assert pdf in uploads[0].body + + +def _gigachat_reply(file_id: str) -> Callable[[Request], Reply]: + def respond(request: Request) -> Reply: + if request.target.endswith("/files"): + return Reply(body=json.dumps({"id": file_id, "object": "file", "purpose": "general"}).encode()) + return _chat_completion(f"chatcmpl-{file_id}", "GigaChat-Pro") + + return respond + + +def test_anthropic_sdk_pdf_document_reaches_gigachat_as_an_uploaded_attachment( + gateway: Gateway, tmp_path: Path +) -> None: + marker: Final = uuid.uuid4().hex + file_id: Final = f"file-{marker}" + pdf: Final = pdf_bytes((Page(text=f"PELICAN-42 {marker}"), LETTER)) + with ( + wire_server(_gigachat_reply(file_id)) as wire, + owned_proxy_process(gateway, tmp_path, {"GIGACHAT_ACCESS_TOKEN": _API_KEY}) as owned, + owned.gateway.scenario() as scenario, + ): + model: Final = scenario.model(model="gigachat/GigaChat-Pro", api_base=f"{wire.url}/v1", api_key=_API_KEY) + client: Final = anthropic.Anthropic( + base_url=str(owned.gateway.client.base_url), api_key=owned.gateway.key, max_retries=0 + ) + message: Final = client.messages.create( + model=model, + max_tokens=64, + messages=[ + { + "role": "user", + "content": [ + { + "type": "document", + "source": {"type": "base64", "media_type": "application/pdf", "data": encoded(pdf)}, + }, + {"type": "text", "text": f"What is the codeword? {marker}"}, + ], + } + ], + ) + assert [(block.type, getattr(block, "text", None)) for block in message.content] == [("text", _ANSWER)] + upload, chat = (request for request in wire.drain() if request.method == "POST") + assert upload.target == "/v1/files", upload.target + assert upload.headers["authorization"] == f"Bearer {_API_KEY}", sorted(upload.headers) + assert pdf in upload.body + assert b"application/pdf" in upload.body + assert chat.target == "/v1/chat/completions", chat.target + turns: Final = _PARTS.validate_python(_JSON_OBJECT.validate_json(chat.body)["messages"]) + assert turns[-1]["attachments"] == [file_id], chat.body diff --git a/tests/integration/providers/test_chat_provider_annotations_wire.py b/tests/integration/providers/test_chat_provider_annotations_wire.py new file mode 100644 index 00000000000..0b18d0837eb --- /dev/null +++ b/tests/integration/providers/test_chat_provider_annotations_wire.py @@ -0,0 +1,190 @@ +import json +import uuid +from collections.abc import Callable, Iterator +from itertools import chain +from typing import Final + +import httpx +from integration._support.client import Gateway +from integration._support.wire import Reply, Request, Wire, wire_server +from integration.caching.response_cache_case import await_entry, redis_store +from pydantic import JsonValue, TypeAdapter +from redis import Redis + +_BACKEND: Final = "z-ai/glm-5.2" +_API_KEY: Final = "synthetic-openrouter-key" +_ANSWER: Final = "The report covers revenue." +_URL_CITATION: Final[dict[str, JsonValue]] = { + "type": "url_citation", + "url_citation": { + "url": "https://example.com/report", + "title": "Report", + "start_index": 0, + "end_index": 10, + "content": "cited passage", + }, +} +_FILE_ANNOTATION: Final[dict[str, JsonValue]] = { + "type": "file", + "file": {"hash": "sha-report", "name": "report.pdf", "content": [{"type": "text", "text": "parsed page one"}]}, +} +_USAGE: Final[dict[str, JsonValue]] = {"prompt_tokens": 40, "completion_tokens": 6, "total_tokens": 46} +_JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue]) +_CHOICES: Final = TypeAdapter(list[dict[str, JsonValue]]) + + +def _completion(identity: str) -> bytes: + return json.dumps( + { + "id": identity, + "object": "chat.completion", + "created": 1, + "model": _BACKEND, + "choices": [ + { + "index": 0, + "message": { + "role": "assistant", + "content": _ANSWER, + "annotations": [_URL_CITATION, _FILE_ANNOTATION], + }, + "finish_reason": "stop", + } + ], + "usage": _USAGE, + } + ).encode() + + +def _frame(frame: dict[str, JsonValue]) -> bytes: + return b"data: " + json.dumps(frame).encode() + b"\n\n" + + +def _stream(identity: str) -> Reply: + chunk: Final[dict[str, JsonValue]] = { + "id": identity, + "object": "chat.completion.chunk", + "created": 1, + "model": _BACKEND, + } + return Reply( + content_type="text/event-stream", + chunks=( + _frame({**chunk, "choices": [{"index": 0, "delta": {"role": "assistant", "content": _ANSWER}}]}), + _frame({**chunk, "choices": [{"index": 0, "delta": {"annotations": [_URL_CITATION, _FILE_ANNOTATION]}}]}), + _frame({**chunk, "choices": [{"index": 0, "delta": {}, "finish_reason": "stop"}], "usage": _USAGE}), + b"data: [DONE]\n\n", + ), + ) + + +def _responder(identity: str) -> Callable[[Request], Reply]: + def respond(request: Request) -> Reply: + if _JSON_OBJECT.validate_json(request.body).get("stream") is True: + return _stream(identity) + return Reply(body=_completion(identity)) + + return respond + + +def _chat_body(model: str, marker: str, stream: bool) -> dict[str, JsonValue]: + return { + "model": model, + "messages": [{"role": "user", "content": f"Summarize the report {marker}"}], + "stream": stream, + } + + +def _posted(wire: Wire) -> int: + return sum(1 for request in wire.drain() if request.method == "POST") + + +def _message(response: httpx.Response) -> dict[str, JsonValue]: + assert response.status_code == 200, response.text + choices: Final = _CHOICES.validate_python(_JSON_OBJECT.validate_json(response.content)["choices"]) + return _JSON_OBJECT.validate_python(choices[0]["message"]) + + +def _assert_split(message: dict[str, JsonValue]) -> None: + assert message["content"] == _ANSWER, message + assert message["annotations"] == [_URL_CITATION], message + assert _JSON_OBJECT.validate_python(message["provider_specific_fields"])["annotations"] == [_FILE_ANNOTATION] + + +def test_file_annotation_reaches_the_caller_in_provider_specific_fields_and_survives_a_cache_hit( + gateway: Gateway, +) -> None: + marker: Final = uuid.uuid4().hex + identity: Final = f"gen-{marker}" + with wire_server(_responder(identity)) as wire, gateway.scenario() as scenario: + model: Final = scenario.model(model=f"openrouter/{_BACKEND}", api_base=wire.url, api_key=_API_KEY) + first: Final = gateway.request("POST", "/v1/chat/completions", _chat_body(model, marker, stream=False)) + _assert_split(_message(first)) + assert _posted(wire) == 1 + await_entry(redis_store(), identity) + cached: Final = gateway.request("POST", "/v1/chat/completions", _chat_body(model, marker, stream=False)) + assert cached.headers.get("x-litellm-cache-key"), dict(cached.headers) + _assert_split(_message(cached)) + assert _posted(wire) == 0 + + +def _deltas(response: httpx.Response) -> Iterator[dict[str, JsonValue]]: + for line in response.iter_lines(): + if not line.startswith("data: ") or line == "data: [DONE]": + continue + choices = _CHOICES.validate_python(_JSON_OBJECT.validate_json(line.removeprefix("data: "))["choices"]) + yield from (_JSON_OBJECT.validate_python(choice["delta"]) for choice in choices) + + +def _streamed(gateway: Gateway, body: dict[str, JsonValue]) -> tuple[dict[str, JsonValue], ...]: + with gateway.client.stream( + "POST", "/v1/chat/completions", json=body, headers={"Authorization": f"Bearer {gateway.key}"} + ) as response: + assert response.status_code == 200, response.read() + return tuple(_deltas(response)) + + +def _listed(value: JsonValue) -> list[JsonValue]: + return value if isinstance(value, list) else [] + + +def _provider_fields(delta: dict[str, JsonValue]) -> dict[str, JsonValue]: + return _JSON_OBJECT.validate_python(delta.get("provider_specific_fields") or {}) + + +def _annotations(deltas: tuple[dict[str, JsonValue], ...]) -> list[JsonValue]: + return list(chain.from_iterable(_listed(delta.get("annotations")) for delta in deltas)) + + +def _provider_annotations(deltas: tuple[dict[str, JsonValue], ...]) -> list[JsonValue]: + return list(chain.from_iterable(_listed(_provider_fields(delta).get("annotations")) for delta in deltas)) + + +def _text(deltas: tuple[dict[str, JsonValue], ...]) -> str: + return "".join(str(delta.get("content") or "") for delta in deltas) + + +def _cached_message(store: Redis, key: str) -> dict[str, JsonValue]: + raw: Final = store.get(key) + assert isinstance(raw, bytes), key + response: Final = _JSON_OBJECT.validate_json(raw)["response"] + stored: Final = _JSON_OBJECT.validate_json(response) if isinstance(response, str) else response + choices: Final = _CHOICES.validate_python(_JSON_OBJECT.validate_python(stored)["choices"]) + return _JSON_OBJECT.validate_python(choices[0]["message"]) + + +def test_streamed_annotations_reach_the_caller_as_sent_and_the_cached_stream_splits_them(gateway: Gateway) -> None: + marker: Final = uuid.uuid4().hex + identity: Final = f"gen-{marker}" + with wire_server(_responder(identity)) as wire, gateway.scenario() as scenario: + model: Final = scenario.model(model=f"openrouter/{_BACKEND}", api_base=wire.url, api_key=_API_KEY) + live: Final = _streamed(gateway, _chat_body(model, marker, stream=True)) + assert _text(live) == _ANSWER, live + assert _annotations(live) == [_URL_CITATION, _FILE_ANNOTATION], live + assert _posted(wire) == 1 + store: Final = redis_store() + _assert_split(_cached_message(store, await_entry(store, identity))) + replayed: Final = _streamed(gateway, _chat_body(model, marker, stream=True)) + assert _text(replayed) == _ANSWER, replayed + assert _provider_annotations(replayed) == [_FILE_ANNOTATION], replayed + assert _posted(wire) == 0 diff --git a/tests/integration/providers/test_pdf_input_support_wire.py b/tests/integration/providers/test_pdf_input_support_wire.py new file mode 100644 index 00000000000..b35815438f7 --- /dev/null +++ b/tests/integration/providers/test_pdf_input_support_wire.py @@ -0,0 +1,390 @@ +import json +import uuid +from collections.abc import Callable, Mapping, Sequence +from concurrent.futures import ThreadPoolExecutor +from dataclasses import dataclass +from pathlib import Path +from typing import Final + +import httpx +import pytest +from integration._support.client import Gateway, Scenario +from integration._support.openai_wire import chat_reply +from integration._support.pdf_document import ( + LETTER, + chat_body, + chat_file, + messages_body, + pdf_data_url, + pdf_document, + responses_body, + responses_input_file, +) +from integration._support.wire import Reply, Request, Wire, wire_server +from pydantic import JsonValue, TypeAdapter + +_API_KEY: Final = "synthetic-pdf-gate-key" +_ANSWER: Final = "The report covers revenue." +_FILENAME: Final = "report.pdf" +_PDF: Final = pdf_data_url((LETTER,)) +_JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue]) +_MESSAGES: Final = TypeAdapter(list[JsonValue]) +_COST_MAP_PATH: Final = Path(__file__).parents[3] / "litellm" / "model_prices_and_context_window_backup.json" + + +@dataclass(frozen=True, slots=True) +class _Entrypoint: + path: str + body: Callable[[str, str], Mapping[str, JsonValue]] + stream: bool = False + provider: str = "hosted_vllm" + + +_ENTRYPOINTS: Final = ( + pytest.param( + _Entrypoint("/v1/chat/completions", lambda model, text: chat_body(model, [chat_file(_PDF, _FILENAME)], text)), + id="chat-completions", + ), + pytest.param( + _Entrypoint( + "/v1/chat/completions", + lambda model, text: chat_body(model, [chat_file(_PDF, _FILENAME)], text, stream=True), + stream=True, + ), + id="chat-completions-stream", + ), + pytest.param( + _Entrypoint("/v1/messages", lambda model, text: messages_body(model, [pdf_document((LETTER,))], text)), + id="messages", + ), + pytest.param( + _Entrypoint( + "/v1/responses", + lambda model, text: responses_body(model, [responses_input_file(_PDF, _FILENAME)], text), + provider="lm_studio", + ), + id="responses-bridge", + ), +) + + +def _azure_ai_chat_model_pdf_input(entry: JsonValue) -> bool | None: + match entry: + case {"litellm_provider": "azure_ai", "mode": "chat", "supports_pdf_input": bool() as supports}: + return supports + case _: + return None + + +def _azure_ai_chat_model(reads_pdfs: bool) -> str: + entries: Final = _JSON_OBJECT.validate_json(_COST_MAP_PATH.read_bytes()) + matching: Final = sorted( + name for name, entry in entries.items() if _azure_ai_chat_model_pdf_input(entry) is reads_pdfs + ) + assert matching, f"the cost map marks no azure_ai chat model with supports_pdf_input {reads_pdfs}" + return matching[0] + + +def _ollama_reply(backend: str) -> Reply: + return Reply( + body=json.dumps( + { + "model": backend, + "created_at": "2026-01-01T00:00:00Z", + "message": {"role": "assistant", "content": _ANSWER}, + "done": True, + "done_reason": "stop", + "prompt_eval_count": 5, + "eval_count": 3, + } + ).encode() + ) + + +def _responder(backend: str) -> Callable[[Request], Reply]: + def respond(request: Request) -> Reply: + if request.target.endswith("/api/chat"): + return _ollama_reply(backend) + body: Final = _JSON_OBJECT.validate_json(request.body) + return chat_reply(f"chatcmpl-{uuid.uuid4().hex}", backend, _ANSWER, stream=body.get("stream") is True) + + return respond + + +def _post(gateway: Gateway, path: str, body: Mapping[str, JsonValue]) -> httpx.Response: + return gateway.request("POST", path, body, headers={"anthropic-version": "2023-06-01"}) + + +def _post_streaming(gateway: Gateway, path: str, body: Mapping[str, JsonValue]) -> httpx.Response: + with gateway.client.stream( + "POST", path, json=dict(body), headers={"Authorization": f"Bearer {gateway.key}"} + ) as response: + response.read() + return response + + +def _send(gateway: Gateway, entrypoint: _Entrypoint, model: str, text: str) -> httpx.Response: + body: Final = entrypoint.body(model, text) + if entrypoint.stream: + return _post_streaming(gateway, entrypoint.path, body) + return _post(gateway, entrypoint.path, body) + + +def _error_message(response: httpx.Response) -> str: + match _JSON_OBJECT.validate_json(response.content): + case {"error": {"message": str() as message}}: + return message + case _: + raise AssertionError(response.text) + + +def _assert_refused_before_the_provider(response: httpx.Response, wire: Wire, backend: str) -> None: + assert response.status_code == 400, response.text + message: Final = _error_message(response) + assert "PDF" in message and backend in message, message + assert wire.drain() == () + + +def _bodies(wire: Wire) -> tuple[dict[str, JsonValue], ...]: + return tuple(_JSON_OBJECT.validate_json(request.body) for request in wire.drain() if request.method == "POST") + + +def _only_body(wire: Wire) -> dict[str, JsonValue]: + bodies: Final = _bodies(wire) + assert len(bodies) == 1, bodies + return bodies[0] + + +def _assert_reaches_the_provider(gateway: Gateway, wire: Wire, model: str) -> None: + response: Final = _post(gateway, "/v1/chat/completions", chat_body(model, [], f"text only {uuid.uuid4().hex}")) + assert response.status_code == 200, response.text + assert len(_bodies(wire)) == 1 + + +def _openai_compatible( + scenario: Scenario, + wire: Wire, + backend: str, + model_info: Mapping[str, JsonValue], + provider: str = "hosted_vllm", + **parameters: JsonValue, +) -> str: + return scenario.model( + model=f"{provider}/{backend}", + api_base=f"{wire.url}/v1", + api_key=_API_KEY, + model_info=model_info, + **parameters, + ) + + +def _backend() -> str: + return f"pdf-gate-{uuid.uuid4().hex[:12]}" + + +@pytest.mark.parametrize("entrypoint", _ENTRYPOINTS) +def test_pdf_to_a_deployment_marked_unable_to_read_it_is_refused_before_the_provider( + gateway: Gateway, entrypoint: _Entrypoint +) -> None: + backend: Final = _backend() + with wire_server(_responder(backend)) as wire, gateway.scenario() as scenario: + model: Final = _openai_compatible(scenario, wire, backend, {"supports_pdf_input": False}, entrypoint.provider) + response: Final = _send(gateway, entrypoint, model, f"Summarize the report {uuid.uuid4().hex}") + _assert_refused_before_the_provider(response, wire, backend) + _assert_reaches_the_provider(gateway, wire, model) + + +def _base_model_deployment(scenario: Scenario, wire: Wire) -> tuple[str, str]: + backend: Final = f"pdf-gate-{uuid.uuid4().hex[:12]}" + model: Final = scenario.model( + model=f"azure_ai/{backend}", + api_base=wire.url, + api_key=_API_KEY, + model_info={"base_model": _azure_ai_chat_model(reads_pdfs=False)}, + ) + return model, backend + + +def _litellm_params_base_model_deployment(scenario: Scenario, wire: Wire) -> tuple[str, str]: + backend: Final = _backend() + model: Final = scenario.model( + model=f"azure_ai/{backend}", + api_base=wire.url, + api_key=_API_KEY, + base_model=_azure_ai_chat_model(reads_pdfs=False), + ) + return model, backend + + +def _cost_map_deployment(scenario: Scenario, wire: Wire) -> tuple[str, str]: + marked: Final = _azure_ai_chat_model(reads_pdfs=False) + return scenario.model(model=marked, api_base=wire.url, api_key=_API_KEY), marked.removeprefix("azure_ai/") + + +def _ollama_deployment(scenario: Scenario, wire: Wire) -> tuple[str, str]: + backend: Final = _backend() + return scenario.model(model=f"ollama_chat/{backend}", api_base=wire.url, api_key=None), backend + + +@pytest.mark.parametrize( + "deploy", + [ + pytest.param(_base_model_deployment, id="model-info-base-model"), + pytest.param(_litellm_params_base_model_deployment, id="litellm-params-base-model"), + pytest.param(_cost_map_deployment, id="cost-map"), + pytest.param(_ollama_deployment, id="ollama"), + ], +) +def test_pdf_to_a_model_its_declaration_marks_unable_is_refused_before_the_provider( + gateway: Gateway, deploy: Callable[[Scenario, Wire], tuple[str, str]] +) -> None: + with wire_server(_responder("scripted")) as wire, gateway.scenario() as scenario: + model, backend = deploy(scenario, wire) + response: Final = _post( + gateway, + "/v1/chat/completions", + chat_body(model, [chat_file(_PDF, _FILENAME)], f"Summarize the report {uuid.uuid4().hex}"), + ) + _assert_refused_before_the_provider(response, wire, backend) + _assert_reaches_the_provider(gateway, wire, model) + + +def _pdf_turns(question: str) -> list[JsonValue]: + return [ + {"role": "user", "content": [chat_file(_PDF, _FILENAME), {"type": "text", "text": question}]}, + {"role": "assistant", "content": "It is a quarterly report."}, + {"role": "user", "content": [chat_file(_PDF, "appendix.pdf")]}, + {"role": "user", "content": "Who wrote it?"}, + ] + + +def _answer(response: httpx.Response) -> JsonValue: + match _JSON_OBJECT.validate_json(response.content): + case {"choices": [{"message": {"content": content}}]}: + return content + case _: + raise AssertionError(response.text) + + +def _turns_without_pdfs(question: str) -> list[JsonValue]: + return [ + {"role": "user", "content": [{"type": "text", "text": question}]}, + {"role": "assistant", "content": "It is a quarterly report."}, + {"role": "user", "content": "Who wrote it?"}, + ] + + +@pytest.mark.parametrize( + ("deployment_params", "request_params"), + [ + pytest.param({"drop_params": True}, {}, id="deployment-drop-params"), + pytest.param({}, {"drop_params": True}, id="request-drop-params"), + ], +) +def test_drop_params_sends_the_conversation_without_the_pdfs_a_model_cannot_read( + gateway: Gateway, deployment_params: Mapping[str, JsonValue], request_params: Mapping[str, JsonValue] +) -> None: + backend: Final = _backend() + question: Final = f"Summarize the report {uuid.uuid4().hex}" + with wire_server(_responder(backend)) as wire, gateway.scenario() as scenario: + model: Final = _openai_compatible(scenario, wire, backend, {"supports_pdf_input": False}, **deployment_params) + response: Final = _post( + gateway, + "/v1/chat/completions", + {"model": model, "max_tokens": 64, "messages": _pdf_turns(question), **request_params}, + ) + assert response.status_code == 200, response.text + assert _answer(response) == _ANSWER + assert _only_body(wire)["messages"] == _turns_without_pdfs(question) + + +def test_drop_params_refuses_a_request_left_with_no_user_turn_once_the_pdf_is_dropped(gateway: Gateway) -> None: + backend: Final = _backend() + with wire_server(_responder(backend)) as wire, gateway.scenario() as scenario: + model: Final = _openai_compatible(scenario, wire, backend, {"supports_pdf_input": False}, drop_params=True) + response: Final = _post( + gateway, + "/v1/chat/completions", + { + "model": model, + "max_tokens": 64, + "messages": [ + {"role": "system", "content": f"Answer from the attachment {uuid.uuid4().hex}"}, + {"role": "user", "content": [chat_file(_PDF, _FILENAME)]}, + ], + }, + ) + _assert_refused_before_the_provider(response, wire, backend) + _assert_reaches_the_provider(gateway, wire, model) + + +def _file_and_question(question: str) -> list[JsonValue]: + return [{"role": "user", "content": [chat_file(_PDF, _FILENAME), {"type": "text", "text": question}]}] + + +@pytest.mark.parametrize( + ("model_info", "deployment_params", "keeps_pdf"), + [ + pytest.param({}, {}, True, id="undeclared"), + pytest.param({}, {"drop_params": True}, False, id="undeclared-drop-params"), + pytest.param({"supports_pdf_input": True}, {"drop_params": True}, True, id="declared-drop-params"), + pytest.param( + {}, + {"drop_params": True, "base_model": _azure_ai_chat_model(reads_pdfs=True)}, + True, + id="base-model-reads-drop-params", + ), + ], +) +def test_azure_ai_model_outside_the_cost_map_gets_the_pdf_unless_drop_params_finds_no_declaration( + gateway: Gateway, model_info: Mapping[str, JsonValue], deployment_params: Mapping[str, JsonValue], keeps_pdf: bool +) -> None: + backend: Final = _backend() + question: Final = f"Summarize the report {uuid.uuid4().hex}" + with wire_server(_responder(backend)) as wire, gateway.scenario() as scenario: + model: Final = scenario.model( + model=f"azure_ai/{backend}", + api_base=wire.url, + api_key=_API_KEY, + model_info=model_info, + **deployment_params, + ) + response: Final = _post( + gateway, "/v1/chat/completions", chat_body(model, [chat_file(_PDF, _FILENAME)], question) + ) + assert response.status_code == 200, response.text + expected: Final[list[JsonValue]] = ( + _file_and_question(question) if keeps_pdf else [{"role": "user", "content": question}] + ) + assert _only_body(wire)["messages"] == expected + + +def _sibling_pdf_request(gateway: Gateway, model: str) -> httpx.Response: + return _post( + gateway, + "/v1/chat/completions", + chat_body(model, [chat_file(_PDF, _FILENAME)], f"Summarize the report {uuid.uuid4().hex}"), + ) + + +def _forwarded_file_parts(bodies: Sequence[dict[str, JsonValue]]) -> tuple[JsonValue, ...]: + return tuple(_MESSAGES.validate_python(body["messages"])[0] for body in bodies) + + +def test_pdf_opt_out_on_one_deployment_leaves_its_sibling_on_the_same_model_reading_pdfs(gateway: Gateway) -> None: + backend: Final = _backend() + burst: Final = 6 + with wire_server(_responder(backend)) as wire, gateway.scenario() as scenario: + sibling: Final = _openai_compatible(scenario, wire, backend, {}) + marked: Final = _openai_compatible(scenario, wire, backend, {"supports_pdf_input": False}) + with ThreadPoolExecutor(max_workers=burst) as pool: + responses: Final = tuple(pool.map(lambda _: _sibling_pdf_request(gateway, sibling), range(burst))) + assert [response.status_code for response in responses] == [200] * burst, [r.text for r in responses] + forwarded: Final = _forwarded_file_parts(_bodies(wire)) + assert len(forwarded) == burst, forwarded + assert all( + _MESSAGES.validate_python(_JSON_OBJECT.validate_python(message)["content"])[0] == chat_file(_PDF, _FILENAME) + for message in forwarded + ), forwarded + refused: Final = _sibling_pdf_request(gateway, marked) + _assert_refused_before_the_provider(refused, wire, backend) diff --git a/tests/unit/litellm_core_utils/llm_response_utils/test_convert_dict_to_response.py b/tests/unit/litellm_core_utils/llm_response_utils/test_convert_dict_to_response.py index 0c601a6b2c8..cea12230fbd 100644 --- a/tests/unit/litellm_core_utils/llm_response_utils/test_convert_dict_to_response.py +++ b/tests/unit/litellm_core_utils/llm_response_utils/test_convert_dict_to_response.py @@ -2968,3 +2968,54 @@ def test_convert_to_model_response_object_function_output(): assert result.usage.completion_tokens_details == CompletionTokensDetailsWrapper(reasoning_tokens=0) assert result._hidden_params is not None + + +_URL_CITATION: Final = { + "type": "url_citation", + "url_citation": { + "url": "https://example.com", + "title": "Example", + "content": "cited snippet", + "start_index": 0, + "end_index": 4, + }, +} +_FILE_ANNOTATION: Final = { + "type": "file", + "file": {"hash": "abc123", "name": "report.pdf", "content": [{"type": "text", "text": "parsed page"}]}, +} + + +def _message_with_annotations(annotations: list[dict[str, object]]) -> Message: + response: Final = { + "id": "gen-1", + "object": "chat.completion", + "created": 1712697600, + "model": "z-ai/glm-5.3", + "choices": [ + { + "index": 0, + "finish_reason": "stop", + "message": {"role": "assistant", "content": "done", "annotations": annotations}, + } + ], + "usage": {"completion_tokens": 1, "prompt_tokens": 1, "total_tokens": 2}, + } + return convert_to_model_response_object( + response_object=response, model_response_object=ModelResponse(), response_type="completion" + ).choices[0].message + + +def test_annotations_typed_other_than_url_citation_move_to_provider_specific_fields() -> None: + untyped_annotation: Final = {"start_index": 0, "end_index": 4} + message: Final = _message_with_annotations([_URL_CITATION, untyped_annotation, _FILE_ANNOTATION]) + + assert message.annotations == [_URL_CITATION, untyped_annotation] + assert message.provider_specific_fields["annotations"] == [_FILE_ANNOTATION] + + +def test_url_citation_annotations_keep_provider_keys() -> None: + message: Final = _message_with_annotations([_URL_CITATION]) + + assert message.annotations == [_URL_CITATION] + assert "annotations" not in (message.provider_specific_fields or {}) diff --git a/tests/unit/llms/anthropic/messages/adapters/test_anthropic_messages_adapters_transformation.py b/tests/unit/llms/anthropic/messages/adapters/test_anthropic_messages_adapters_transformation.py index 51ec71aaea0..c30644e7fdc 100644 --- a/tests/unit/llms/anthropic/messages/adapters/test_anthropic_messages_adapters_transformation.py +++ b/tests/unit/llms/anthropic/messages/adapters/test_anthropic_messages_adapters_transformation.py @@ -5286,6 +5286,40 @@ def test_eager_input_streaming_tool_reaches_bedrock_converse_as_beta(): assert data["toolConfig"]["tools"][0]["toolSpec"]["inputSchema"]["json"] == EAGER_INPUT_SCHEMA +@pytest.mark.parametrize( + ("title", "expected_filename"), + [(None, "document.pdf"), ("Quarterly report", "Quarterly report")], +) +def test_user_pdf_document_reaches_openai_compatible_providers_as_file_part( + title: str | None, expected_filename: str +) -> None: + document: Final = _base64_pdf_block() | ({"title": title} if title else {}) + request_messages: Final = _run_chat_completions_pipeline( + [ + AnthropicMessagesUserMessageParam( + role="user", + content=[document, {"type": "text", "text": "What is the codeword?"}], + ) + ] + ) + + assert request_messages == [ + { + "role": "user", + "content": [ + { + "type": "file", + "file": { + "filename": expected_filename, + "file_data": f"data:application/pdf;base64,{TOOL_RESULT_PDF_B64}", + }, + }, + {"type": "text", "text": "What is the codeword?"}, + ], + } + ] + + _CHAT_BRIDGE_MODEL: Final = "azure/gpt-5.6" _CHAT_BRIDGE_API_BASE: Final = "https://tool-result-text.openai.azure.com" _CHAT_BRIDGE_URL: Final = r".*tool-result-text\.openai\.azure\.com/openai/deployments/gpt-5\.6/chat/completions.*" diff --git a/tests/unit/llms/azure_ai/chat/test_azure_ai_transformation.py b/tests/unit/llms/azure_ai/chat/test_azure_ai_transformation.py index 42dc222e420..6269fd367fe 100644 --- a/tests/unit/llms/azure_ai/chat/test_azure_ai_transformation.py +++ b/tests/unit/llms/azure_ai/chat/test_azure_ai_transformation.py @@ -1101,6 +1101,58 @@ def test_map_azure_model_group(model_group_header, expected_model): assert config._map_azure_model_group(model_group_header) == expected_model +def test_azure_ai_sends_file_parts_instead_of_flattening_them_to_text() -> None: + file_part: Final = { + "type": "file", + "file": {"filename": "document.pdf", "file_data": "data:application/pdf;base64,JVBERi0xLjQK"}, + } + content: Final = [file_part, {"type": "text", "text": "What is the codeword?"}] + + request: Final = AzureAIStudioConfig().transform_request( + model="Kimi-K2.6", + messages=[{"role": "user", "content": content}], + optional_params={}, + litellm_params={}, + headers={}, + ) + + assert request["messages"] == [{"role": "user", "content": content}] + + +def test_azure_ai_passes_dict_message_content_through_for_the_provider_to_judge() -> None: + content: Final = {"type": "text", "text": "Say hi"} + + request: Final = AzureAIStudioConfig().transform_request( + model="gpt-5.6-sol", + messages=[{"role": "user", "content": content}], + optional_params={}, + litellm_params={}, + headers={}, + ) + + assert request["messages"] == [{"role": "user", "content": content}] + + +@pytest.mark.parametrize( + "file", + [ + {"filename": "notes.txt", "file_data": "data:text/plain;base64,aGVsbG8="}, + {"file_id": "file-abc"}, + ], + ids=["text_data_url", "file_id_only"], +) +def test_azure_ai_flattens_a_non_pdf_file_part_to_its_text_as_before(file: dict[str, str]) -> None: + request: Final = AzureAIStudioConfig().transform_request( + model="Llama-3.3-70B-Instruct", + messages=[{"role": "user", "content": [{"type": "text", "text": "summarize"}, {"type": "file", "file": file}]}], + optional_params={}, + litellm_params={}, + headers={}, + ) + + assert request["messages"] == [{"role": "user", "content": "summarize"}] + + def test_translate_developer_role_hoists_a_later_developer_message_for_a_model_that_rejects_it( declare_mid_conversation_system: Callable[[str, bool], None], ) -> None: diff --git a/tests/unit/llms/gigachat/chat/test_gigachat_chat_transformation.py b/tests/unit/llms/gigachat/chat/test_gigachat_chat_transformation.py index f2224e84ff3..115d81bfb90 100644 --- a/tests/unit/llms/gigachat/chat/test_gigachat_chat_transformation.py +++ b/tests/unit/llms/gigachat/chat/test_gigachat_chat_transformation.py @@ -8,6 +8,7 @@ get_model_response_iterator, and get_error_class. """ import json +from typing import Final from unittest.mock import MagicMock, patch import httpx @@ -512,6 +513,20 @@ class TestTransformMessages: assert result[0]["content"] == "look" assert "file-456" in result[0]["attachments"] + def test_list_content_with_pdf_file_part_uploads_it_as_an_attachment(self) -> None: + pdf_data_url: Final = "data:application/pdf;base64,JVBERi0xLjQK" + content: Final = [ + {"type": "text", "text": "summarize"}, + {"type": "file", "file": {"filename": "report.pdf", "file_data": pdf_data_url}}, + ] + with patch.object(self.config, "_upload_image", return_value="file-789") as upload: + result: Final = self.config._transform_messages([ + {"role": "user", "content": content} + ]) + upload.assert_called_once_with(pdf_data_url) + assert result[0]["content"] == "summarize" + assert result[0]["attachments"] == ["file-789"] + class TestTransformRequest: def setup_method(self): diff --git a/tests/unit/llms/oci/chat/test_oci_chat_transformation.py b/tests/unit/llms/oci/chat/test_oci_chat_transformation.py index 8c96bd91049..0535279d177 100644 --- a/tests/unit/llms/oci/chat/test_oci_chat_transformation.py +++ b/tests/unit/llms/oci/chat/test_oci_chat_transformation.py @@ -1,4 +1,5 @@ import datetime +from typing import Final import httpx import pytest import json @@ -2164,3 +2165,32 @@ class TestOpenAIPrefixMaxTokensKey: chat_request = _oci_chat_request("openai.gpt-5", {"max_tokens": 64}) assert chat_request["maxCompletionTokens"] == 64 assert "maxTokens" not in chat_request + + +def test_transform_request_sends_a_pdf_file_part_as_an_image_data_url() -> None: + pdf_data_url: Final = "data:application/pdf;base64,JVBERi0xLjQK" + transformed_request: Final = OCIChatConfig().transform_request( + model="google.gemini-2.5-flash", + messages=[ + { + "role": "user", + "content": [ + {"type": "text", "text": "Summarize this."}, + {"type": "file", "file": {"filename": "report.pdf", "file_data": pdf_data_url}}, + ], + } + ], + optional_params={"oci_compartment_id": TEST_COMPARTMENT_ID}, + litellm_params={}, + headers={}, + ) + + assert transformed_request["chatRequest"]["messages"] == [ + { + "role": "USER", + "content": [ + {"type": "TEXT", "text": "Summarize this."}, + {"type": "IMAGE", "imageUrl": {"url": pdf_data_url}}, + ], + } + ] diff --git a/tests/unit/llms/sap/chat/test_sap_transformation.py b/tests/unit/llms/sap/chat/test_sap_transformation.py index 41e75b7bfc7..f95bf9b7845 100644 --- a/tests/unit/llms/sap/chat/test_sap_transformation.py +++ b/tests/unit/llms/sap/chat/test_sap_transformation.py @@ -1,4 +1,6 @@ import warnings +from typing import Final + import pytest from pydantic import ValidationError @@ -908,6 +910,22 @@ class TestCacheControl: assert "cache_control" not in content[0] assert content[1].get("cache_control") == {"type": "ephemeral"} + def test_pdf_file_part_reaches_sap_as_an_image_url_data_url(self) -> None: + pdf_data_url: Final = "data:application/pdf;base64,JVBERi0xLjQK" + messages: Final = [ + { + "role": "user", + "content": [ + {"type": "text", "text": "Summarize this."}, + {"type": "file", "file": {"filename": "report.pdf", "file_data": pdf_data_url}}, + ], + } + ] + body: Final = self._transform("anthropic--claude-3-5-sonnet", messages) + content: Final = self._template(body)[0]["content"] + assert content[1]["type"] == "image_url" + assert content[1]["image_url"]["url"] == pdf_data_url + def test_message_level_cache_control_folded_onto_trailing_image_block(self): messages = [ { diff --git a/tests/unit/llms/snowflake/chat/test_snowflake_chat_transformation.py b/tests/unit/llms/snowflake/chat/test_snowflake_chat_transformation.py index 36c319d9524..a31f15551e9 100644 --- a/tests/unit/llms/snowflake/chat/test_snowflake_chat_transformation.py +++ b/tests/unit/llms/snowflake/chat/test_snowflake_chat_transformation.py @@ -18,6 +18,7 @@ import pytest import litellm from litellm import completion, acompletion from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler +from litellm.llms.anthropic.messages.adapters.transformation import LiteLLMAnthropicMessagesAdapter from litellm.llms.snowflake.chat.transformation import SnowflakeConfig, SnowflakeStreamingHandler from litellm.types.utils import ModelResponse @@ -479,6 +480,35 @@ class TestSnowflakeCortexClaudeFixes: } ] + def test_messages_api_pdf_reaches_cortex_claude_as_a_document_block(self) -> None: + translated: Final = LiteLLMAnthropicMessagesAdapter().translate_anthropic_messages_to_openai( + messages=[ + { + "role": "user", + "content": [ + { + "type": "document", + "source": {"type": "base64", "media_type": "application/pdf", "data": "ZmFrZQ=="}, + "cache_control": {"type": "ephemeral"}, + }, + {"type": "text", "text": "What is the codeword?"}, + ], + } + ], + model="snowflake/claude-sonnet-4-6", + ) + + body: Final = self._transform(translated) + + assert body["messages"][0]["content"] == [ + { + "type": "document", + "source": {"type": "base64", "media_type": "application/pdf", "data": "ZmFrZQ=="}, + "cache_control": {"type": "ephemeral"}, + }, + {"type": "text", "text": "What is the codeword?"}, + ] + def test_multipart_tool_result_preserves_text_and_converts_image(self): body = self._transform( [ diff --git a/tests/unit/llms/vertex_ai/vertex_ai_partner_models/gemma/test_vertex_ai_gemma_global_endpoint.py b/tests/unit/llms/vertex_ai/vertex_ai_partner_models/gemma/test_vertex_ai_gemma_global_endpoint.py index 85a2124ab02..98d0d314977 100644 --- a/tests/unit/llms/vertex_ai/vertex_ai_partner_models/gemma/test_vertex_ai_gemma_global_endpoint.py +++ b/tests/unit/llms/vertex_ai/vertex_ai_partner_models/gemma/test_vertex_ai_gemma_global_endpoint.py @@ -13,6 +13,7 @@ These tests verify that: import json import os +from typing import Final from unittest.mock import AsyncMock, MagicMock, patch import pytest @@ -422,3 +423,52 @@ async def test_vertex_ai_gemma_vision_passthrough(): assert isinstance(content, list), f"Expected list content, got: {content}" image_parts = [p for p in content if p.get("type") == "image_url"] assert image_parts, f"No image_url part in forwarded message content: {content}" + + +@pytest.mark.asyncio +async def test_vertex_ai_gemma_pdf_file_part_is_sent_as_image_url() -> None: + pdf_data_url: Final = "data:application/pdf;base64,JVBERi0xLjQKJcfsj6IKMSAwIG9iago8PD4+CmVuZG9iagp0cmFpbGVyCjw8Pj4KJSVFT0YK" + messages: Final = [ + { + "role": "user", + "content": [ + {"type": "text", "text": "What does the PDF say?"}, + {"type": "file", "file": {"file_data": pdf_data_url, "filename": "report.pdf"}}, + ], + } + ] + + mock_response: Final = MagicMock() + mock_response.status_code = 200 + mock_response.headers = {} + mock_response.json.return_value = _MOCK_RESPONSE_JSON + + mock_vertexai: Final = MagicMock() + mock_vertexai.preview = MagicMock() + + with ( + patch("litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler") as mock_http_handler, + patch( + "litellm.llms.vertex_ai.vertex_ai_partner_models.main.VertexAIPartnerModels._ensure_access_token", + return_value=("fake-token", "test-project"), + ), + patch.dict( + "sys.modules", + {"vertexai": mock_vertexai, "vertexai.preview": mock_vertexai.preview}, + ), + patch.dict(litellm.model_cost, _GEMMA_MODEL_COST_ENTRY, clear=False), + ): + mock_http_handler.return_value.post = AsyncMock(return_value=mock_response) + + await litellm.acompletion( + model="vertex_ai/google/gemma-4-26b-a4b-it-maas", + messages=messages, + vertex_ai_project="test-project", + ) + + body: Final = json.loads(mock_http_handler.return_value.post.call_args.kwargs["data"]) + user_msg: Final = next(m for m in body["messages"] if m["role"] == "user") + assert user_msg["content"] == [ + {"type": "text", "text": "What does the PDF say?"}, + {"type": "image_url", "image_url": {"url": pdf_data_url}}, + ] diff --git a/tests/unit/proxy/guardrails/test_prompt_security_guardrails.py b/tests/unit/proxy/guardrails/test_prompt_security_guardrails.py index e66e19dd1b4..b08646f1611 100644 --- a/tests/unit/proxy/guardrails/test_prompt_security_guardrails.py +++ b/tests/unit/proxy/guardrails/test_prompt_security_guardrails.py @@ -1,6 +1,7 @@ import asyncio import base64 from collections.abc import Mapping, Sequence +from typing import Final from unittest.mock import AsyncMock, patch import pytest @@ -8,6 +9,7 @@ from fastapi.exceptions import HTTPException from httpx import ReadTimeout, Request, Response import litellm +from litellm.llms.anthropic.chat.guardrail_translation.handler import AnthropicMessagesHandler from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.guardrails.guardrail_hooks.prompt_security.prompt_security import ( PromptSecurityGuardrail, @@ -1247,3 +1249,54 @@ async def test_check_tool_results_enabled(monkeypatch: pytest.MonkeyPatch): assert "indirect_prompt_injection" in str(excinfo.value.detail) + +@pytest.mark.asyncio +async def test_messages_api_pdf_reaches_file_sanitization_and_is_blocked() -> None: + sanitize_url: Final = "https://test.prompt.security/api/sanitizeFile" + upload: Final = AsyncMock( + return_value=Response(200, json={"jobId": "pdf-job"}, request=Request("POST", sanitize_url)) + ) + poll: Final = AsyncMock( + return_value=Response( + 200, + json={"status": "done", "content": "", "metadata": {"action": "block", "violations": ["Malicious File"]}}, + request=Request("GET", sanitize_url), + ) + ) + guardrail: Final = PromptSecurityGuardrail( + api_key="test-key", + api_base="https://test.prompt.security", + guardrail_name="test-guard", + event_hook="pre_call", + default_on=True, + ) + guardrail.poll_interval = 0 + pdf: Final = b"%PDF-1.4 pelican" + request: Final = { + "model": "gpt-5.6", + "max_tokens": 64, + "messages": [ + { + "role": "user", + "content": [ + { + "type": "document", + "source": { + "type": "base64", + "media_type": "application/pdf", + "data": base64.b64encode(pdf).decode(), + }, + }, + {"type": "text", "text": "Summarize this PDF"}, + ], + } + ], + } + + with patch.object(guardrail.async_handler, "post", upload), patch.object(guardrail.async_handler, "get", poll): + with pytest.raises(HTTPException) as blocked: + await AnthropicMessagesHandler().process_input_messages(request, guardrail) + + assert blocked.value.status_code == 400 + assert "Malicious File" in str(blocked.value.detail) + assert upload.call_args.kwargs["files"]["file"][1] == pdf diff --git a/tests/unit/test_main.py b/tests/unit/test_main.py index e119fe133d5..e10352e0b3d 100644 --- a/tests/unit/test_main.py +++ b/tests/unit/test_main.py @@ -23,7 +23,7 @@ import respx from openai import APITimeoutError, AsyncOpenAI, OpenAI from openai._legacy_response import LegacyAPIResponse from openai.types.chat.chat_completion import ChatCompletion -from pydantic import TypeAdapter +from pydantic import BaseModel, ConfigDict, TypeAdapter import litellm from litellm import ( @@ -10200,6 +10200,309 @@ async def test_atranscription_ignores_a_non_string_base_url_when_api_base_is_set assert response.text == "four score" +_PDF_FILE_PART: Final = { + "type": "file", + "file": {"filename": "report.pdf", "file_data": "data:application/pdf;base64,JVBERi0xLjQK"}, +} +_SUMMARIZE_TEXT_PART: Final = {"type": "text", "text": "Summarize this."} + + +def _pdf_messages() -> list[dict[str, object]]: + return [{"role": "user", "content": [_SUMMARIZE_TEXT_PART, _PDF_FILE_PART]}] + + +class _SentChatRequest(BaseModel): + model_config = ConfigDict(frozen=True) + + messages: list[dict[str, object]] + + +def _sent_messages(respx_mock: respx.MockRouter) -> list[dict[str, object]]: + request: Final[httpx.Request] = respx_mock.calls[0].request + return _SentChatRequest.model_validate_json(request.content).messages + + +def _mock_pdf_chat_route(respx_mock: respx.MockRouter, api_base: str) -> respx.Route: + return respx_mock.post(url__regex=rf"{api_base}/chat/completions.*").mock( + return_value=httpx.Response( + status_code=200, + json={ + "id": "chatcmpl-pdf", + "object": "chat.completion", + "created": 1712697600, + "model": "qa-pdf-model", + "choices": [{"index": 0, "message": {"role": "assistant", "content": "ok"}, "finish_reason": "stop"}], + "usage": {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2}, + }, + ) + ) + + +def test_drop_params_drops_a_pdf_the_model_cannot_read(respx_mock: respx.MockRouter) -> None: + api_base: Final = "http://localhost:12347/v1" + _mock_pdf_chat_route(respx_mock, api_base) + + litellm.completion( + model="openai/qa-text-only-model", + messages=_pdf_messages(), + api_base=api_base, + api_key="fake_openai_api_key", + model_info={"supports_pdf_input": False}, + drop_params=True, + ) + + assert _sent_messages(respx_mock) == [{"role": "user", "content": [_SUMMARIZE_TEXT_PART]}] + + +@pytest.mark.parametrize( + "pdf_file", + [ + {"filename": "report.pdf", "file_data": "JVBERi0xLjQK"}, + {"file_id": "file-report", "format": "application/pdf"}, + {"file_id": "file-report", "filename": "Report.PDF"}, + ], + ids=["bare_base64_named_pdf", "file_id_with_pdf_format", "file_id_named_pdf"], +) +def test_drop_params_drops_a_pdf_file_part_without_a_data_url( + respx_mock: respx.MockRouter, pdf_file: dict[str, str] +) -> None: + api_base: Final = "http://localhost:12350/v1" + _mock_pdf_chat_route(respx_mock, api_base) + + litellm.completion( + model="openai/qa-text-only-model", + messages=[{"role": "user", "content": [_SUMMARIZE_TEXT_PART, {"type": "file", "file": pdf_file}]}], + api_base=api_base, + api_key="fake_openai_api_key", + model_info={"supports_pdf_input": False}, + drop_params=True, + ) + + assert _sent_messages(respx_mock) == [{"role": "user", "content": [_SUMMARIZE_TEXT_PART]}] + + +def test_drop_params_drops_a_user_turn_that_held_only_a_pdf(respx_mock: respx.MockRouter) -> None: + api_base: Final = "http://localhost:12352/v1" + _mock_pdf_chat_route(respx_mock, api_base) + + litellm.completion( + model="openai/qa-text-only-model", + messages=[ + {"role": "user", "content": [_PDF_FILE_PART]}, + {"role": "user", "content": [_SUMMARIZE_TEXT_PART]}, + ], + api_base=api_base, + api_key="fake_openai_api_key", + model_info={"supports_pdf_input": False}, + drop_params=True, + ) + + assert _sent_messages(respx_mock) == [{"role": "user", "content": [_SUMMARIZE_TEXT_PART]}] + + +def test_drop_params_keeps_roles_alternating_after_dropping_a_pdf_only_turn(respx_mock: respx.MockRouter) -> None: + api_base: Final = "http://localhost:12353/v1" + _mock_pdf_chat_route(respx_mock, api_base) + + litellm.completion( + model="openai/qa-text-only-model", + messages=[ + {"role": "user", "content": [_SUMMARIZE_TEXT_PART]}, + {"role": "assistant", "content": "Send the file."}, + {"role": "user", "content": [_PDF_FILE_PART]}, + {"role": "assistant", "content": "Got it."}, + {"role": "user", "content": [_SUMMARIZE_TEXT_PART]}, + ], + api_base=api_base, + api_key="fake_openai_api_key", + model_info={"supports_pdf_input": False}, + drop_params=True, + ensure_alternating_roles=True, + ) + + sent_roles: Final = [message["role"] for message in _sent_messages(respx_mock)] + assert sent_roles == ["user", "assistant", "user", "assistant", "user"] + + +def test_drop_params_keeps_an_image_data_url_whose_filename_says_pdf(respx_mock: respx.MockRouter) -> None: + api_base: Final = "http://localhost:12351/v1" + _mock_pdf_chat_route(respx_mock, api_base) + image_part: Final = { + "type": "file", + "file": {"filename": "scan.pdf", "file_data": "data:image/png;base64,iVBORw0KGgo="}, + } + + litellm.completion( + model="openai/qa-text-only-model", + messages=[{"role": "user", "content": [_SUMMARIZE_TEXT_PART, image_part]}], + api_base=api_base, + api_key="fake_openai_api_key", + model_info={"supports_pdf_input": False}, + drop_params=True, + ) + + assert _sent_messages(respx_mock) == [{"role": "user", "content": [_SUMMARIZE_TEXT_PART, image_part]}] + + +def test_pdf_to_a_model_that_cannot_read_it_is_rejected_without_drop_params(respx_mock: respx.MockRouter) -> None: + api_base: Final = "http://localhost:12348/v1" + route: Final = _mock_pdf_chat_route(respx_mock, api_base) + + with pytest.raises(litellm.BadRequestError, match="does not accept PDF input"): + litellm.completion( + model="openai/qa-text-only-model", + messages=_pdf_messages(), + api_base=api_base, + api_key="fake_openai_api_key", + model_info={"supports_pdf_input": False}, + ) + + assert not route.called + + +def test_pdf_to_a_model_with_unknown_pdf_support_is_sent_even_with_drop_params(respx_mock: respx.MockRouter) -> None: + api_base: Final = "http://localhost:12349/v1" + _mock_pdf_chat_route(respx_mock, api_base) + + litellm.completion( + model="openai/qa-unlisted-model", + messages=_pdf_messages(), + api_base=api_base, + api_key="fake_openai_api_key", + drop_params=True, + ) + + assert _sent_messages(respx_mock) == _pdf_messages() + + +def test_pdf_to_ollama_is_rejected_without_drop_params() -> None: + with pytest.raises(litellm.BadRequestError, match="does not accept PDF input"): + litellm.completion(model="ollama_chat/llama3.2", messages=_pdf_messages(), api_base="http://127.0.0.1:9") + + +def _azure_ai_model_without_pdf_input() -> str: + return next( + name + for name, entry in litellm.model_cost.items() + if entry.get("litellm_provider") == "azure_ai" and entry.get("supports_pdf_input") is False + ) + + +def test_pdf_to_a_deployment_whose_base_model_cannot_read_it_is_rejected() -> None: + with pytest.raises(litellm.BadRequestError, match="does not accept PDF input"): + litellm.completion( + model="azure_ai/qa-custom-deployment", + messages=_pdf_messages(), + api_base="http://127.0.0.1:9", + api_key="fake_azure_ai_api_key", + model_info={"base_model": _azure_ai_model_without_pdf_input()}, + ) + + +def test_drop_params_drops_a_pdf_a_deployments_base_model_cannot_read(respx_mock: respx.MockRouter) -> None: + api_base: Final = "http://localhost:12353" + _mock_pdf_chat_route(respx_mock, api_base) + + litellm.completion( + model="azure_ai/qa-custom-deployment", + messages=_pdf_messages(), + api_base=api_base, + api_key="fake_azure_ai_api_key", + model_info={"base_model": _azure_ai_model_without_pdf_input()}, + drop_params=True, + ) + + assert _sent_messages(respx_mock) == [{"role": "user", "content": _SUMMARIZE_TEXT_PART["text"]}] + + +def _azure_ai_model_with_pdf_input() -> str: + return next( + name + for name, entry in litellm.model_cost.items() + if entry.get("litellm_provider") == "azure_ai" and entry.get("supports_pdf_input") is True + ) + + +def test_pdf_to_a_deployment_whose_base_model_param_cannot_read_it_is_rejected() -> None: + with pytest.raises(litellm.BadRequestError, match="does not accept PDF input"): + litellm.completion( + model="azure_ai/qa-custom-deployment", + messages=_pdf_messages(), + api_base="http://127.0.0.1:9", + api_key="fake_azure_ai_api_key", + base_model=_azure_ai_model_without_pdf_input(), + ) + + +def test_drop_params_keeps_a_pdf_a_deployments_base_model_param_reads(respx_mock: respx.MockRouter) -> None: + api_base: Final = "http://localhost:12357" + _mock_pdf_chat_route(respx_mock, api_base) + + litellm.completion( + model="azure_ai/qa-custom-deployment", + messages=_pdf_messages(), + api_base=api_base, + api_key="fake_azure_ai_api_key", + base_model=_azure_ai_model_with_pdf_input(), + drop_params=True, + ) + + assert _sent_messages(respx_mock) == _pdf_messages() + + +def test_drop_params_rejects_a_request_whose_only_user_turn_was_a_pdf(respx_mock: respx.MockRouter) -> None: + api_base: Final = "http://localhost:12354/v1" + route: Final = _mock_pdf_chat_route(respx_mock, api_base) + + with pytest.raises(litellm.BadRequestError, match="no user message is left"): + litellm.completion( + model="openai/qa-text-only-model", + messages=[ + {"role": "system", "content": "You read documents."}, + {"role": "user", "content": [_PDF_FILE_PART]}, + ], + api_base=api_base, + api_key="fake_openai_api_key", + model_info={"supports_pdf_input": False}, + drop_params=True, + ) + + assert not route.called + + +def test_drop_params_drops_a_pdf_an_unlisted_azure_ai_model_never_declared_it_reads( + respx_mock: respx.MockRouter, +) -> None: + api_base: Final = "http://localhost:12355" + _mock_pdf_chat_route(respx_mock, api_base) + + litellm.completion( + model="azure_ai/qa-unlisted-deployment", + messages=_pdf_messages(), + api_base=api_base, + api_key="fake_azure_ai_api_key", + drop_params=True, + ) + + assert _sent_messages(respx_mock) == [{"role": "user", "content": _SUMMARIZE_TEXT_PART["text"]}] + + +def test_drop_params_keeps_a_pdf_an_azure_ai_deployment_declares_it_reads(respx_mock: respx.MockRouter) -> None: + api_base: Final = "http://localhost:12356" + _mock_pdf_chat_route(respx_mock, api_base) + + litellm.completion( + model="azure_ai/qa-unlisted-deployment", + messages=_pdf_messages(), + api_base=api_base, + api_key="fake_azure_ai_api_key", + model_info={"supports_pdf_input": True}, + drop_params=True, + ) + + assert _sent_messages(respx_mock) == _pdf_messages() + + @pytest.mark.parametrize("stream", (False, True)) @pytest.mark.parametrize("asynchronous", (False, True)) @pytest.mark.parametrize("profile,custom_price", ((None, False), ("us.openai.gpt-6-sol", False), (None, True))) diff --git a/tests/unit/test_router_model_cost_isolation.py b/tests/unit/test_router_model_cost_isolation.py index 20b0055f729..91f0b2339c7 100644 --- a/tests/unit/test_router_model_cost_isolation.py +++ b/tests/unit/test_router_model_cost_isolation.py @@ -614,6 +614,63 @@ def test_capability_flag_stays_on_the_deployment_that_set_it(flag: str) -> None: _restore_model_cost_entries(model_keys) +def test_pdf_opt_out_on_one_deployment_does_not_block_its_sibling(respx_mock: respx.MockRouter) -> None: + api_base: Final = "http://localhost:12361/v1" + backend_model: Final = "openai/qa-shared-pdf-backend" + pdf_messages: Final = [ + { + "role": "user", + "content": [ + {"type": "text", "text": "Summarize this."}, + { + "type": "file", + "file": {"filename": "report.pdf", "file_data": "data:application/pdf;base64,JVBERi0xLjQK"}, + }, + ], + } + ] + route: Final = respx_mock.post(url__regex=rf"{api_base}/chat/completions.*").mock( + return_value=httpx.Response( + status_code=200, + json={ + "id": "chatcmpl-pdf", + "object": "chat.completion", + "created": 1712697600, + "model": "qa-shared-pdf-backend", + "choices": [{"index": 0, "message": {"role": "assistant", "content": "ok"}, "finish_reason": "stop"}], + "usage": {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2}, + }, + ) + ) + model_keys: Final = { + key: litellm.model_cost.get(key) for key in ("qa-pdf-marked", "qa-pdf-unmarked", backend_model) + } + try: + router: Final = Router( + model_list=[ + { + "model_name": "pdf-blocked", + "litellm_params": {"model": backend_model, "api_base": api_base, "api_key": "fake-key"}, + "model_info": {"id": "qa-pdf-marked", "supports_pdf_input": False}, + }, + { + "model_name": "pdf-allowed", + "litellm_params": {"model": backend_model, "api_base": api_base, "api_key": "fake-key"}, + "model_info": {"id": "qa-pdf-unmarked"}, + }, + ], + ) + + with pytest.raises(litellm.BadRequestError, match="does not accept PDF input"): + router.completion(model="pdf-blocked", messages=pdf_messages) + router.completion(model="pdf-allowed", messages=pdf_messages) + + assert route.call_count == 1 + assert json.loads(route.calls[0].request.content)["messages"] == pdf_messages + finally: + _restore_model_cost_entries(model_keys) + + def test_token_count_settings_stay_on_the_deployment_that_set_them() -> None: """A deployment's ``token_count_model`` and ``token_count_framing`` must not land on the shared ``{provider}/{model}`` key, or a sibling deployment of that model would count with them.""" diff --git a/tests/unit/test_stream_chunk_builder_annotations.py b/tests/unit/test_stream_chunk_builder_annotations.py index 9c7ad4126b0..c710697b39c 100644 --- a/tests/unit/test_stream_chunk_builder_annotations.py +++ b/tests/unit/test_stream_chunk_builder_annotations.py @@ -6,6 +6,8 @@ annotation chunk, losing any annotations that arrived in later chunks. This fix merges annotations from ALL chunks. """ +from typing import Final + from litellm import stream_chunk_builder from litellm.types.utils import Delta, ModelResponseStream, StreamingChoices @@ -189,3 +191,46 @@ def test_stream_chunk_builder_no_annotations(): message = response["choices"][0]["message"] assert not hasattr(message, "annotations") or message.annotations is None + + +def test_stream_chunk_builder_moves_typed_provider_annotations_like_non_streaming() -> None: + url_citation: Final = { + "type": "url_citation", + "url_citation": {"url": "https://example.com", "title": "Example", "start_index": 0, "end_index": 5}, + } + file_annotation: Final = {"type": "file", "file": {"hash": "abc123", "name": "report.pdf"}} + chunks: Final = [ + ModelResponseStream( + id="chatcmpl-test", + created=1700000000, + model="test-model", + object="chat.completion.chunk", + choices=[ + StreamingChoices( + finish_reason=None, + index=0, + delta=Delta(content="Hello", role="assistant", provider_specific_fields={"region": "eu"}), + ) + ], + ), + ModelResponseStream( + id="chatcmpl-test", + created=1700000000, + model="test-model", + object="chat.completion.chunk", + choices=[ + StreamingChoices( + finish_reason="stop", + index=0, + delta=Delta(content=None, annotations=[url_citation, file_annotation]), + ) + ], + ), + ] + + response: Final = stream_chunk_builder(chunks=chunks) + assert response is not None + + message: Final = response["choices"][0]["message"] + assert message.annotations == [url_citation] + assert message.provider_specific_fields == {"region": "eu", "annotations": [file_annotation]} From 77064d6c7086919c28de177d32864e8ffba7c21a Mon Sep 17 00:00:00 2001 From: Mateo Wang <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 29 Oct 2026 08:15:05 -0700 Subject: [PATCH 06/13] feat(vector_stores): add OpenSearch k-NN search as a vector store provider (#47759) * feat(vector_stores): add Amazon OpenSearch as a vector store provider * feat(ui): add Amazon OpenSearch to the vector store provider dropdown * feat(vector_stores): add OpenSearch k-NN search as a vector store provider * test(vector_stores): hold the OpenSearch provider tests alone for the first review The provider source comes back in the next commits on this PR * ci(tests): run the OpenSearch vector store unit tests * fix(vector_stores): sign on.aws and FIPS OpenSearch hosts and reject missing AWS credentials before embedding * fix(vector_stores): regenerate provider endpoint support for opensearch * fix(vector_stores): keep opensearch retrieve, update, and delete off the index and refuse half-set AWS keys and per-account endpoints * fix(vector_stores): return only scalar opensearch attributes and keep credential source output out of 401s * test(vector_stores): inject a failing opensearch credential source instead of running a shell * fix(vector_stores): sign opensearch for the endpoint's region, refuse list, add opensearch_index_name, honor drop_params * fix(vector_stores): refuse a caller-chosen opensearch_index_name on search and rag query * fix(vector_stores): refuse opensearch_index_name only on OpenSearch stores, keep the old 400 text Other providers drop the key again as they did before OpenSearch support, and the four embedding keys keep their original error sentence. * fix(vector_stores): refuse a caller index on slash-form OpenSearch stores --- .github/workflows/test-unit.yml | 1 + .../base_llm/vector_store/transformation.py | 3 + litellm/llms/custom_httpx/llm_http_handler.py | 4 +- litellm/llms/opensearch/__init__.py | 0 .../llms/opensearch/vector_stores/__init__.py | 0 .../vector_stores/transformation.py | 415 +++++++++++++++ .../provider_endpoints_support_backup.json | 27 + litellm/proxy/rag_endpoints/endpoints.py | 6 + .../proxy/vector_store_endpoints/endpoints.py | 26 + litellm/types/utils.py | 1 + litellm/utils.py | 35 +- provider_endpoints_support.json | 27 + tests/unit/AGENTS.md | 3 + tests/unit/llms/opensearch/__init__.py | 0 .../llms/opensearch/vector_stores/__init__.py | 0 .../test_opensearch_transformation.py | 482 ++++++++++++++++++ .../public_endpoints/test_public_endpoints.py | 1 + .../proxy/rag_endpoints/test_rag_endpoints.py | 56 ++ .../test_vector_store_endpoints.py | 52 ++ .../_components/CreateVectorStore.test.tsx | 6 +- .../_components/CreateVectorStore.tsx | 2 +- .../VectorStoreForm.integration.test.tsx | 66 ++- .../_components/VectorStoreForm.tsx | 40 +- .../vector_store_providers.test.tsx | 26 + .../src/components/vector_store_providers.tsx | 62 +++ 25 files changed, 1323 insertions(+), 18 deletions(-) create mode 100644 litellm/llms/opensearch/__init__.py create mode 100644 litellm/llms/opensearch/vector_stores/__init__.py create mode 100644 litellm/llms/opensearch/vector_stores/transformation.py create mode 100644 tests/unit/llms/opensearch/__init__.py create mode 100644 tests/unit/llms/opensearch/vector_stores/__init__.py create mode 100644 tests/unit/llms/opensearch/vector_stores/test_opensearch_transformation.py diff --git a/.github/workflows/test-unit.yml b/.github/workflows/test-unit.yml index e7865817048..509271cbf54 100644 --- a/.github/workflows/test-unit.yml +++ b/.github/workflows/test-unit.yml @@ -322,6 +322,7 @@ jobs: tests/unit/llms/openai_like tests/unit/llms/opencode tests/unit/llms/opensandbox + tests/unit/llms/opensearch tests/unit/llms/openrouter tests/unit/llms/ovhcloud tests/unit/llms/parallel_ai diff --git a/litellm/llms/base_llm/vector_store/transformation.py b/litellm/llms/base_llm/vector_store/transformation.py index 6885183aeba..4924c5a5acc 100644 --- a/litellm/llms/base_llm/vector_store/transformation.py +++ b/litellm/llms/base_llm/vector_store/transformation.py @@ -226,6 +226,9 @@ class BaseVectorStoreConfig: ``get_complete_url`` adds a query string, like Azure, override this.""" return f"{api_base}/{encoded_vector_store_id}" + def get_vector_store_list_url(self, api_base: str) -> str: + return api_base + def get_error_class(self, error_message: str, status_code: int, headers: dict | httpx.Headers) -> BaseLLMException: from ..chat.transformation import BaseLLMException diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index 8a9ff62cfae..c970b507288 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -11150,7 +11150,7 @@ class BaseLLMHTTPHandler: litellm_params=dict(litellm_params), ) - url: Final = api_base + url: Final = vector_store_provider_config.get_vector_store_list_url(api_base) params: Final[dict[str, object]] = {} if after is not None: @@ -11229,7 +11229,7 @@ class BaseLLMHTTPHandler: litellm_params=dict(litellm_params), ) - url: Final = api_base + url: Final = vector_store_provider_config.get_vector_store_list_url(api_base) params: Final[dict[str, object]] = {} if after is not None: diff --git a/litellm/llms/opensearch/__init__.py b/litellm/llms/opensearch/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/litellm/llms/opensearch/vector_stores/__init__.py b/litellm/llms/opensearch/vector_stores/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/litellm/llms/opensearch/vector_stores/transformation.py b/litellm/llms/opensearch/vector_stores/transformation.py new file mode 100644 index 00000000000..2d862bc8ea6 --- /dev/null +++ b/litellm/llms/opensearch/vector_stores/transformation.py @@ -0,0 +1,415 @@ +import hashlib +import json +import re +from collections.abc import Mapping, Sequence +from functools import partial +from typing import TYPE_CHECKING, Final, Literal, NoReturn +from urllib.parse import quote, urlsplit + +import httpx +from pydantic import ConfigDict, Field, TypeAdapter, ValidationError + +from litellm._logging import verbose_logger +from litellm.exceptions import AuthenticationError, BadRequestError +from litellm.litellm_core_utils.core_helpers import drops_unsupported_params +from litellm.litellm_core_utils.optional_imports import ensure_optional_import +from litellm.llms.base_llm.chat.transformation import BaseLLMException +from litellm.llms.base_llm.vector_store.transformation import ( + BaseQueryEmbeddingVectorStoreConfig, + VectorStoreEmbeddingExecutor, +) +from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM, SignsRequestsWithAWS, sign_aws_json_post +from litellm.types.llms.base import LiteLLMBaseModel +from litellm.types.llms.bedrock import AwsAuthParams +from litellm.types.router import GenericLiteLLMParams +from litellm.types.utils import EmbeddingResponse +from litellm.types.vector_stores import ( + BaseVectorStoreAuthCredentials, + VectorStoreCreateOptionalRequestParams, + VectorStoreIndexEndpoints, + VectorStoreResultContent, + VectorStoreSearchOptionalRequestParams, + VectorStoreSearchResponse, + VectorStoreSearchResult, +) + +if TYPE_CHECKING: + from botocore.credentials import Credentials + + from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj + +DEFAULT_EMBEDDING_FIELD: Final = "embedding" +DEFAULT_TEXT_FIELD: Final = "text" +DEFAULT_MAX_NUM_RESULTS: Final = 10 +MAX_MAX_NUM_RESULTS: Final = 50 +_PARAM_PREFIX: Final = "opensearch_" +_AWS_HOSTS: Final = ( + re.compile(r"\.(?P[a-z0-9-]+)\.(?Paoss|aoss-fips|es)\.amazonaws\.com(?:\.cn)?\Z"), + re.compile(r"\.(?Paoss|aos)\.(?P[a-z0-9-]+)\.on\.aws\Z"), +) +_ACCOUNT_ID_LABEL: Final = re.compile(r"\A\d{12}\.") +_QUERY_VECTOR: Final = TypeAdapter(tuple[float, ...]) +_ATTRIBUTE_VALUE_TYPES: Final = (str, int, float, bool) +_SEARCH_ONLY_MESSAGE: Final = ( + "The OpenSearch vector store is search-only. Create the k-NN index in OpenSearch, " + "then register it here with its name as opensearch_index_name or vector_store_id." +) +_PER_ACCOUNT_MESSAGE: Final = ( + "api_base is an OpenSearch Serverless per-account endpoint, which needs a collection header the vector store " + "does not send. Use the collection endpoint instead, for example https://..aoss.amazonaws.com." +) +_SIGV4_ONLY_MESSAGE: Final = ( + "The OpenSearch vector store signs every request with AWS SigV4 and does not support api_key. " + "Remove api_key and give the vector store AWS credentials instead." +) +_CREDENTIALS_HINT: Final = ( + "Set aws_access_key_id and aws_secret_access_key, aws_role_name, or aws_profile_name on the vector store, " + "or run LiteLLM with an AWS identity." +) + + +def _bad_request(message: str) -> BadRequestError: + return BadRequestError(message=message, model=None, llm_provider="opensearch") + + +def _unauthenticated(reason: str) -> AuthenticationError: + return AuthenticationError( + message=f"No usable AWS credentials to sign the OpenSearch request ({reason}). {_CREDENTIALS_HINT}", + llm_provider="opensearch", + model=None, + ) + + +def _usable_credentials(credentials: object) -> "Credentials": + from botocore.credentials import Credentials + + if isinstance(credentials, Credentials): + return credentials + raise _unauthenticated("none were found") + + +class _OpenSearchParams(LiteLLMBaseModel): + model_config = ConfigDict(frozen=True, extra="ignore") + + litellm_embedding_model: str | None = None + opensearch_embedding_field: str | None = None + opensearch_index_name: str | None = None + opensearch_text_field: str | None = None + + @property + def embedding_field(self) -> str: + return self.opensearch_embedding_field or DEFAULT_EMBEDDING_FIELD + + @property + def text_field(self) -> str: + return self.opensearch_text_field or DEFAULT_TEXT_FIELD + + def index_name(self, vector_store_id: str) -> str: + return self.opensearch_index_name or vector_store_id + + def require_embedding_model(self) -> str: + if not self.litellm_embedding_model: + raise _bad_request( + "litellm_embedding_model is required for the OpenSearch vector store. It must be the model that " + f"produced the vectors in '{self.embedding_field}', for example openai/text-embedding-3-small." + ) + return self.litellm_embedding_model + + +class _SigningRegion(LiteLLMBaseModel): + model_config = ConfigDict(frozen=True, extra="ignore") + + aws_region_name: str | None = None + + +class _SearchOptions(LiteLLMBaseModel): + model_config = ConfigDict(frozen=True, extra="ignore") + + max_num_results: int | None = None + filters: Mapping[str, object] | None = None + ranking_options: Mapping[str, object] | None = None + rewrite_query: bool | None = None + + @property + def unsupported(self) -> tuple[str, ...]: + requested: Final = ( + ("filters", self.filters), + ("ranking_options", self.ranking_options), + ("rewrite_query", self.rewrite_query), + ) + return tuple(name for name, value in requested if value) + + @property + def k(self) -> int: + if self.max_num_results is None: + return DEFAULT_MAX_NUM_RESULTS + if not 1 <= self.max_num_results <= MAX_MAX_NUM_RESULTS: + raise _bad_request( + f"max_num_results must be between 1 and {MAX_MAX_NUM_RESULTS}, got {self.max_num_results}" + ) + return self.max_num_results + + +_KNOWN_PARAMS: Final = frozenset(name for name in _OpenSearchParams.model_fields if name.startswith(_PARAM_PREFIX)) + + +class _LoggedSearch(LiteLLMBaseModel): + model_config = ConfigDict(frozen=True, extra="ignore") + + query: str | tuple[str, ...] = "" + litellm_params: _OpenSearchParams = _OpenSearchParams() + + @property + def search_query(self) -> str: + return BaseQueryEmbeddingVectorStoreConfig.query_text(self.query) + + +class _Hit(LiteLLMBaseModel): + model_config = ConfigDict(frozen=True) + + id: str = Field(alias="_id") + score: float | None = Field(default=None, alias="_score") + source: Mapping[str, object] = Field(default_factory=dict, alias="_source") + + +class _Hits(LiteLLMBaseModel): + model_config = ConfigDict(frozen=True) + + hits: Sequence[_Hit] + + +class _SearchResponse(LiteLLMBaseModel): + model_config = ConfigDict(frozen=True) + + hits: _Hits + + +def aws_signing_target(url: str) -> tuple[Literal["aoss", "es"], str | None]: + hostname: Final = urlsplit(url).hostname or "" + match: Final = next(filter(None, (pattern.search(hostname) for pattern in _AWS_HOSTS)), None) + if match is None: + return "es", None + return ("aoss" if match["service"].startswith("aoss") else "es"), match["region"] + + +def _search_result(hit: _Hit, params: _OpenSearchParams) -> VectorStoreSearchResult: + text: Final = hit.source.get(params.text_field) + filename: Final = hit.source.get("filename") + return VectorStoreSearchResult( + score=hit.score, + content=[VectorStoreResultContent(type="text", text=text)] if isinstance(text, str) else [], + file_id=hit.id, + filename=filename if isinstance(filename, str) else None, + attributes={ + key: value + for key, value in hit.source.items() + if key not in (params.text_field, params.embedding_field) and isinstance(value, _ATTRIBUTE_VALUE_TYPES) + }, + ) + + +class OpenSearchVectorStoreConfig(BaseQueryEmbeddingVectorStoreConfig, SignsRequestsWithAWS): + def __init__(self, aws: BaseAWSLLM | None = None) -> None: + self._aws: Final = aws if aws is not None else BaseAWSLLM() + + def get_auth_credentials(self, litellm_params: Mapping[str, object]) -> BaseVectorStoreAuthCredentials: + return BaseVectorStoreAuthCredentials() + + def get_vector_store_endpoints_by_type(self) -> VectorStoreIndexEndpoints: + return VectorStoreIndexEndpoints(read=[], write=[]) + + def validate_environment( + self, headers: Mapping[str, object], litellm_params: GenericLiteLLMParams | None + ) -> dict[str, object]: # mutable-ok: the shared HTTP handler requires writable headers + params: Final = dict(litellm_params or {}) + self._signing_credentials(params, self._signing_region(params, str(params.get("api_base") or ""))) + return {**headers, "Content-Type": "application/json"} + + def _signing_region(self, params: Mapping[str, object], api_base: str) -> str: + host_region: Final = aws_signing_target(api_base)[1] + return self._aws.get_aws_region_name_for_non_llm_api_calls( + host_region or _SigningRegion.model_validate(params).aws_region_name + ) + + def _signing_credentials(self, params: Mapping[str, object], region: str) -> "Credentials": + if params.get("api_key"): + raise _bad_request(_SIGV4_ONLY_MESSAGE) + auth: Final = AwsAuthParams.model_validate(params) + if bool(auth.aws_access_key_id) != bool(auth.aws_secret_access_key): + raise _unauthenticated("aws_access_key_id and aws_secret_access_key must be set together") + ensure_optional_import("botocore") + from botocore.exceptions import ( + CredentialRetrievalError, + NoCredentialsError, + PartialCredentialsError, + ProfileNotFound, + ) + + try: + return _usable_credentials(self._aws.resolve_credentials(auth, region)) + except CredentialRetrievalError as error: + verbose_logger.warning("OpenSearch vector store AWS credential source failed: %s", error) + raise _unauthenticated("the AWS credential source failed, see the LiteLLM logs") from None + except (NoCredentialsError, PartialCredentialsError, ProfileNotFound) as error: + raise _unauthenticated(str(error)) from None + + def get_complete_url(self, api_base: str | None, litellm_params: Mapping[str, object]) -> str: + parsed: Final = urlsplit(api_base or "") + if parsed.scheme not in ("http", "https") or not parsed.hostname: + raise _bad_request( + "api_base must be the OpenSearch endpoint URL, for example " + "https://..aoss.amazonaws.com or https://." + ) + if _ACCOUNT_ID_LABEL.match(parsed.hostname) and aws_signing_target(api_base or "")[0] == "aoss": + raise _bad_request(_PER_ACCOUNT_MESSAGE) + return (api_base or "").rstrip("/") + + def get_vector_store_url(self, api_base: str, encoded_vector_store_id: str) -> NoReturn: + raise _bad_request(_SEARCH_ONLY_MESSAGE) + + def get_vector_store_list_url(self, api_base: str) -> NoReturn: + raise _bad_request(_SEARCH_ONLY_MESSAGE) + + @staticmethod + def _params( + litellm_params: Mapping[str, object], + optional_params: VectorStoreSearchOptionalRequestParams, + extra_body: Mapping[str, object] | None, + ) -> tuple[_OpenSearchParams, int]: + unknown: Final = sorted( + key for key in litellm_params if key.startswith(_PARAM_PREFIX) and key not in _KNOWN_PARAMS + ) + if unknown: + raise _bad_request( + f"Unknown OpenSearch vector store parameter(s): {', '.join(unknown)}. " + f"Supported: {', '.join(sorted(_KNOWN_PARAMS))}." + ) + if extra_body: + raise _bad_request("The OpenSearch vector store does not support extra_body.") + try: + options: Final = _SearchOptions.model_validate(optional_params) + except ValidationError as error: + raise _bad_request(f"Invalid vector store search parameters: {error}") from None + if options.unsupported and not drops_unsupported_params(litellm_params.get("drop_params")): + raise _bad_request( + f"The OpenSearch vector store does not support {', '.join(options.unsupported)} yet. " + "Remove it from the request, or set drop_params to drop it." + ) + try: + return _OpenSearchParams.model_validate(litellm_params), options.k + except ValidationError: + raise _bad_request( + "opensearch_embedding_field, opensearch_index_name, opensearch_text_field, and litellm_embedding_model " + "must be strings." + ) from None + + @staticmethod + def _search_request( + vector_store_id: str, + embedding: EmbeddingResponse, + k: int, + params: _OpenSearchParams, + api_base: str, + ) -> tuple[str, dict[str, object]]: # mutable-ok: the provider contract returns a writable JSON request body + vector: Final = _QUERY_VECTOR.validate_python(embedding.data[0]["embedding"]) # pyright: ignore[reportUnknownMemberType] # EmbeddingResponse.data is an untyped list, the vector is validated here + return ( + f"{api_base}/{quote(params.index_name(vector_store_id), safe='')}/_search", + { + "size": k, + "query": {"knn": {params.embedding_field: {"vector": vector, "k": k}}}, + "_source": {"excludes": [params.embedding_field]}, + }, + ) + + def transform_search_vector_store_request( + self, + vector_store_id: str, + query: str | Sequence[str], + vector_store_search_optional_params: VectorStoreSearchOptionalRequestParams, + api_base: str, + litellm_logging_obj: "LiteLLMLoggingObj", + litellm_params: Mapping[str, object], + extra_body: Mapping[str, object] | None = None, + embedding_executor: VectorStoreEmbeddingExecutor | None = None, + ) -> tuple[str, dict[str, object]]: # mutable-ok: the provider contract returns a writable JSON request body + params, k = self._params(litellm_params, vector_store_search_optional_params, extra_body) + embedding: Final = self.query_embedding_executor(embedding_executor, None).embed( + params.require_embedding_model(), self.query_text(query), self.query_embedding_configuration(litellm_params) + ) + return self._search_request(vector_store_id, embedding, k, params, api_base) + + async def atransform_search_vector_store_request( + self, + vector_store_id: str, + query: str | Sequence[str], + vector_store_search_optional_params: VectorStoreSearchOptionalRequestParams, + api_base: str, + litellm_logging_obj: "LiteLLMLoggingObj", + litellm_params: Mapping[str, object], + extra_body: Mapping[str, object] | None = None, + embedding_executor: VectorStoreEmbeddingExecutor | None = None, + ) -> tuple[str, dict[str, object]]: # mutable-ok: the provider contract returns a writable JSON request body + params, k = self._params(litellm_params, vector_store_search_optional_params, extra_body) + embedding: Final = await self.query_embedding_executor(embedding_executor, None).aembed( + params.require_embedding_model(), self.query_text(query), self.query_embedding_configuration(litellm_params) + ) + return self._search_request(vector_store_id, embedding, k, params, api_base) + + def sign_request( + self, + headers: Mapping[str, str], + optional_params: Mapping[str, object], + request_data: Mapping[str, object], + api_base: str, + api_key: str | None = None, + ) -> tuple[dict[str, str], bytes]: # mutable-ok: the shared HTTP handler requires writable headers + region: Final = self._signing_region(optional_params, api_base) + body: Final = json.dumps(request_data) + signed: Final = sign_aws_json_post( + get_credentials=partial(self._signing_credentials, optional_params, region), + service_name=aws_signing_target(api_base)[0], + aws_region_name=region, + url=api_base, + body=body, + headers={ + "Content-Type": "application/json", + "x-amz-content-sha256": hashlib.sha256(body.encode()).hexdigest(), + }, + ) + signed_headers: Final = dict(signed.headers.items()) + signed_names: Final = frozenset(name.lower() for name in signed_headers) + return { + **{name: value for name, value in headers.items() if name.lower() not in signed_names}, + **signed_headers, + }, body.encode() + + def transform_search_vector_store_response( + self, response: httpx.Response, litellm_logging_obj: "LiteLLMLoggingObj" + ) -> VectorStoreSearchResponse: + if not response.is_success: + raise BaseLLMException(status_code=response.status_code, message=response.text, headers=response.headers) + try: + hits: Final = _SearchResponse.model_validate_json(response.content).hits.hits + except ValidationError: + raise BaseLLMException( + status_code=502, + message="OpenSearch answered the search without a hits.hits list.", + headers=response.headers, + ) from None + logged: Final = _LoggedSearch.model_validate(litellm_logging_obj.model_call_details) + return VectorStoreSearchResponse( + object="vector_store.search_results.page", + search_query=logged.search_query, + data=[_search_result(hit, logged.litellm_params) for hit in hits], + ) + + def validate_create_vector_store(self) -> NoReturn: + raise _bad_request(_SEARCH_ONLY_MESSAGE) + + def transform_create_vector_store_request( + self, vector_store_create_optional_params: VectorStoreCreateOptionalRequestParams, api_base: str + ) -> NoReturn: + raise _bad_request(_SEARCH_ONLY_MESSAGE) + + def transform_create_vector_store_response(self, response: httpx.Response) -> NoReturn: + raise _bad_request(_SEARCH_ONLY_MESSAGE) diff --git a/litellm/provider_endpoints_support_backup.json b/litellm/provider_endpoints_support_backup.json index 9f5989e765c..1d487f33835 100644 --- a/litellm/provider_endpoints_support_backup.json +++ b/litellm/provider_endpoints_support_backup.json @@ -4910,6 +4910,33 @@ "vector_stores_search": true } }, + "opensearch": { + "display_name": "Amazon OpenSearch (`opensearch`)", + "url": "https://docs.litellm.ai/docs/providers/opensearch_vector_stores", + "endpoints": { + "chat_completions": false, + "responses": false, + "text_completion": false, + "embeddings": false, + "image_generations": false, + "image_edits": false, + "audio_transcriptions": false, + "audio_speech": false, + "rerank": false, + "ocr": false, + "search": false, + "moderations": false, + "realtime": false, + "video_generations": false, + "batches": false, + "messages": false, + "interactions": false, + "systemone": false, + "a2a": false, + "generateContent": false, + "vector_stores_search": true + } + }, "helicone": { "display_name": "Helicone (`helicone`)", "url": "https://docs.litellm.ai/docs/providers/helicone", diff --git a/litellm/proxy/rag_endpoints/endpoints.py b/litellm/proxy/rag_endpoints/endpoints.py index a7146f13035..2d81dc32181 100644 --- a/litellm/proxy/rag_endpoints/endpoints.py +++ b/litellm/proxy/rag_endpoints/endpoints.py @@ -54,6 +54,7 @@ from litellm.proxy.rag_endpoints.upload_security import ( from litellm.proxy.vector_store_endpoints.endpoints import ( build_request_data_from_managed_vector_store, reject_caller_embedding_selection_params, + reject_caller_opensearch_index_name, ) from litellm.proxy.vector_store_endpoints.utils import ( assert_user_can_access_vector_store_id, @@ -913,6 +914,11 @@ async def rag_query( **retrieval_config, **store_data, } + reject_caller_opensearch_index_name( + caller_body=retrieval_config, + request_fields=merged_retrieval_config, + source="retrieval_config", + ) effective_vector_store_config: Final = ( managed_store_params if managed_store is not None diff --git a/litellm/proxy/vector_store_endpoints/endpoints.py b/litellm/proxy/vector_store_endpoints/endpoints.py index 1b54dddcf4f..b6a0ebe3c58 100644 --- a/litellm/proxy/vector_store_endpoints/endpoints.py +++ b/litellm/proxy/vector_store_endpoints/endpoints.py @@ -17,6 +17,7 @@ from litellm.proxy._types import CommonProxyErrors, UserAPIKeyAuth from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing from litellm.proxy.common_utils.http_parsing_utils import read_raw_json_body +from litellm.proxy.route_llm_request import JSON_OBJECT_ADAPTER from litellm.proxy.utils import jsonify_object from litellm.proxy.vector_store_endpoints.utils import ( assert_proxy_admin_for_vector_store_index_management, @@ -24,6 +25,7 @@ from litellm.proxy.vector_store_endpoints.utils import ( get_litellm_managed_vector_store, ) from litellm.repositories.table_repositories import ManagedVectorStoreIndexRepository +from litellm.types.utils import LlmProviders from litellm.types.vector_stores import IndexCreateRequest, IndexListResponse from litellm.vector_stores.utils import VectorStoreWriteBody, callers_metadata_snapshot from litellm.vector_stores.vector_store_registry import VectorStoreIndexRegistry @@ -52,6 +54,24 @@ def reject_caller_embedding_selection_params(payload: Mapping[str, object], sour ) +def reject_caller_opensearch_index_name( + caller_body: Mapping[str, object], request_fields: Mapping[str, object], source: str +) -> None: + provider: Final = request_fields.get("custom_llm_provider") + if ( + isinstance(provider, str) + and provider.split("/", 1)[0] == LlmProviders.OPENSEARCH.value + and "opensearch_index_name" in caller_body + ): + raise HTTPException( + status_code=400, + detail={ + "error": f"'opensearch_index_name' cannot be set in {source}. " + "The index comes from the vector store's server-side registration." + }, + ) + + ######################################################## # OpenAI Compatible Endpoints ######################################################## @@ -138,11 +158,17 @@ async def vector_store_search( data = await read_request_body(request=request) # rebind-ok: pre-existing rebinding on a rename-only line reject_caller_embedding_selection_params(payload=data, source="the search request body") data["vector_store_id"] = vector_store_id + caller_body: Final = JSON_OBJECT_ADAPTER.validate_python(data) # Check for legacy vector store registry (non-managed vector stores) data = await _update_request_data_with_litellm_managed_vector_store_registry( data=data, vector_store_id=vector_store_id, user_api_key_dict=user_api_key_dict ) + reject_caller_opensearch_index_name( + caller_body=caller_body, + request_fields=JSON_OBJECT_ADAPTER.validate_python(data), + source="the search request body", + ) # The managed_vector_stores pre-call hook will handle: # 1. Decoding managed vector store IDs diff --git a/litellm/types/utils.py b/litellm/types/utils.py index 448bd022989..3bdcba75778 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -4564,6 +4564,7 @@ class LlmProviders(str, Enum): S3_VECTORS = "s3_vectors" VALKEY = "valkey" MONGODB = "mongodb" + OPENSEARCH = "opensearch" HELICONE = "helicone" HYPERBOLIC = "hyperbolic" RECRAFT = "recraft" diff --git a/litellm/utils.py b/litellm/utils.py index 342a2f12310..5fe3e63f31a 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -10460,6 +10460,22 @@ class ProviderConfigManager: return BedrockVectorStore.get_initialized_custom_logger() return None + @staticmethod + def _get_vertex_ai_vector_stores_config(api_type: str | None) -> BaseVectorStoreConfig | None: + if api_type == "rag_api" or api_type is None: + from litellm.llms.vertex_ai.vector_stores.rag_api.transformation import ( + VertexVectorStoreConfig, + ) + + return VertexVectorStoreConfig() + if api_type == "search_api": + from litellm.llms.vertex_ai.vector_stores.search_api.transformation import ( + VertexSearchAPIVectorStoreConfig, + ) + + return VertexSearchAPIVectorStoreConfig() + return None + @staticmethod def get_provider_vector_stores_config( provider: LlmProviders, @@ -10481,18 +10497,7 @@ class ProviderConfigManager: return AzureOpenAIVectorStoreConfig() elif litellm.LlmProviders.VERTEX_AI == provider: - if api_type == "rag_api" or api_type is None: # default to rag_api - from litellm.llms.vertex_ai.vector_stores.rag_api.transformation import ( - VertexVectorStoreConfig, - ) - - return VertexVectorStoreConfig() - elif api_type == "search_api": - from litellm.llms.vertex_ai.vector_stores.search_api.transformation import ( - VertexSearchAPIVectorStoreConfig, - ) - - return VertexSearchAPIVectorStoreConfig() + return ProviderConfigManager._get_vertex_ai_vector_stores_config(api_type) elif litellm.LlmProviders.BEDROCK == provider: from litellm.llms.bedrock.vector_stores.transformation import ( BedrockVectorStoreConfig, @@ -10547,6 +10552,12 @@ class ProviderConfigManager: ) return MongoDBVectorStoreConfig() + elif litellm.LlmProviders.OPENSEARCH == provider: + from litellm.llms.opensearch.vector_stores.transformation import ( + OpenSearchVectorStoreConfig, + ) + + return OpenSearchVectorStoreConfig() return None @staticmethod diff --git a/provider_endpoints_support.json b/provider_endpoints_support.json index 9f5989e765c..1d487f33835 100644 --- a/provider_endpoints_support.json +++ b/provider_endpoints_support.json @@ -4910,6 +4910,33 @@ "vector_stores_search": true } }, + "opensearch": { + "display_name": "Amazon OpenSearch (`opensearch`)", + "url": "https://docs.litellm.ai/docs/providers/opensearch_vector_stores", + "endpoints": { + "chat_completions": false, + "responses": false, + "text_completion": false, + "embeddings": false, + "image_generations": false, + "image_edits": false, + "audio_transcriptions": false, + "audio_speech": false, + "rerank": false, + "ocr": false, + "search": false, + "moderations": false, + "realtime": false, + "video_generations": false, + "batches": false, + "messages": false, + "interactions": false, + "systemone": false, + "a2a": false, + "generateContent": false, + "vector_stores_search": true + } + }, "helicone": { "display_name": "Helicone (`helicone`)", "url": "https://docs.litellm.ai/docs/providers/helicone", diff --git a/tests/unit/AGENTS.md b/tests/unit/AGENTS.md index ee6648c0a82..75c9cc94981 100644 --- a/tests/unit/AGENTS.md +++ b/tests/unit/AGENTS.md @@ -42,6 +42,9 @@ Green if `send_batched` drops every row. pydantic doubles in 12 of 203 files, fa `tests/unit/` mirrors `litellm/`, so a changed file selects its tests by path, not a mapping file. New unit tests go here +A new directory under `tests/unit` runs in CI only once a shard in `.github/workflows/test-unit.yml` names it. +`pytest tests/unit/test_assert_ci_coverage.py` checks that, so run it before you push + ## Writing it so a human can read it A class only when tests share an arrange diff --git a/tests/unit/llms/opensearch/__init__.py b/tests/unit/llms/opensearch/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/opensearch/vector_stores/__init__.py b/tests/unit/llms/opensearch/vector_stores/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/opensearch/vector_stores/test_opensearch_transformation.py b/tests/unit/llms/opensearch/vector_stores/test_opensearch_transformation.py new file mode 100644 index 00000000000..07836d57af0 --- /dev/null +++ b/tests/unit/llms/opensearch/vector_stores/test_opensearch_transformation.py @@ -0,0 +1,482 @@ +import hashlib +import json +import re +from collections.abc import Callable, Mapping +from datetime import datetime, timezone +from pathlib import Path +from typing import Final +from unittest.mock import MagicMock + +import httpx +import openai +import pytest +from botocore.auth import SigV4Auth +from botocore.awsrequest import AWSRequest +from botocore.credentials import Credentials +from botocore.exceptions import CredentialRetrievalError + +import litellm +from litellm.litellm_core_utils.litellm_logging import Logging +from litellm.llms.base_llm.chat.transformation import BaseLLMException +from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM +from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler +from litellm.llms.opensearch.vector_stores.transformation import OpenSearchVectorStoreConfig +from litellm.types.llms.bedrock import AwsAuthParams +from litellm.types.router import GenericLiteLLMParams +from litellm.types.utils import EmbeddingResponse +from litellm.types.vector_stores import VectorStoreSearchResponse +from litellm.vector_stores import main as vector_stores + +ACCESS_KEY: Final = "AKIDOPENSEARCHTEST" +SECRET_KEY: Final = "opensearch-test-secret" +QUERY: Final = "How many days a year can I work abroad?" +QUERY_VECTOR: Final = [0.1, 0.2, 0.3] +CREDENTIAL_SOURCE_OUTPUT: Final = "credential-process-stderr-detail" +SERVERLESS: Final = "https://abc123.us-east-1.aoss.amazonaws.com" +BASE_PARAMS: Final = { + "api_base": SERVERLESS, + "aws_access_key_id": ACCESS_KEY, + "aws_secret_access_key": SECRET_KEY, + "litellm_embedding_model": "embedding-alias", +} +_AUTHORIZATION: Final = re.compile( + r"\AAWS4-HMAC-SHA256 Credential=(?P[^/]+)/(?P\d{8})/(?P[^/]+)/(?P[^/]+)/aws4_request, " + r"SignedHeaders=(?P[^,]+), Signature=(?P[0-9a-f]{64})\Z" +) + + +class RecordingEmbeddingExecutor: + def __init__(self) -> None: + self.call: Final = MagicMock(return_value=EmbeddingResponse(data=[{"embedding": QUERY_VECTOR}])) + + def embed(self, model: str, query: str, configuration: Mapping[str, object]) -> EmbeddingResponse: + return self.call(model, query, configuration) + + async def aembed(self, model: str, query: str, configuration: Mapping[str, object]) -> EmbeddingResponse: + return self.call(model, query, configuration) + + +def _signature_scope(request: httpx.Request) -> tuple[str, str]: + authorization: Final = _AUTHORIZATION.match(request.headers["authorization"]) + assert authorization is not None, request.headers["authorization"] + signed_names: Final = authorization["signed"].split(";") + assert {"host", "x-amz-content-sha256", "x-amz-date"} <= set(signed_names), signed_names + assert request.headers["x-amz-content-sha256"] == hashlib.sha256(request.content).hexdigest() + signed_request: Final = AWSRequest( + method=request.method, + url=str(request.url), + data=request.content, + headers={name: request.headers[name] for name in signed_names}, + ) + signed_request.context["timestamp"] = request.headers["x-amz-date"] + signer: Final = SigV4Auth(Credentials(ACCESS_KEY, SECRET_KEY), authorization["service"], authorization["region"]) + expected: Final = signer.signature( + signer.string_to_sign(signed_request, signer.canonical_request(signed_request)), signed_request + ) + assert authorization["key"] == ACCESS_KEY + assert authorization["signature"] == expected, "SigV4 signature does not cover the request that was sent" + return authorization["service"], authorization["region"] + + +def _sent(respond: MagicMock) -> tuple[httpx.Request, ...]: + return tuple(call.args[0] for call in respond.call_args_list) + + +def _opensearch_hits(text_field: str, embedding_field: str) -> Mapping[str, object]: + return { + "took": 4, + "timed_out": False, + "hits": { + "total": {"value": 2, "relation": "eq"}, + "max_score": 0.581, + "hits": [ + { + "_index": "litellm-qa-docs", + "_id": "1%3A0%3Aw6ftW6EB", + "_score": 0.581, + "_source": { + text_field: "Employees may work abroad for at most 37 calendar days per year.", + embedding_field: [0.5, 0.5, 0.5], + "embedding_ada002": [0.25, 0.75], + "filename": "hr-remote-work-policy.md", + "category": "hr", + "revision": 7, + "published": True, + "tags": ["hr", "travel"], + "owner": {"team": "people-ops"}, + "reviewed_by": None, + }, + }, + {"_index": "litellm-qa-docs", "_id": "doc-2", "_score": 0.42, "_source": {"category": "facilities"}}, + ], + }, + } + + +EXPECTED_RESULTS: Final = [ + { + "score": 0.581, + "content": [{"type": "text", "text": "Employees may work abroad for at most 37 calendar days per year."}], + "file_id": "1%3A0%3Aw6ftW6EB", + "filename": "hr-remote-work-policy.md", + "attributes": {"filename": "hr-remote-work-policy.md", "category": "hr", "revision": 7, "published": True}, + }, + {"score": 0.42, "content": [], "file_id": "doc-2", "filename": None, "attributes": {"category": "facilities"}}, +] + + +async def _search( + asynchronous: bool, + respond: Callable[[httpx.Request], httpx.Response], + executor: RecordingEmbeddingExecutor, + kwargs: Mapping[str, object], +) -> VectorStoreSearchResponse: + request_kwargs: Final = { + **BASE_PARAMS, + "vector_store_id": "litellm-qa-docs", + "query": QUERY, + "custom_llm_provider": "opensearch", + "_direct_vector_store_embedding_executor": executor, + **kwargs, + } + if asynchronous: + return await litellm.vector_stores.asearch( + **request_kwargs, client=AsyncHTTPHandler(transport=httpx.MockTransport(respond)) + ) + with httpx.Client(transport=httpx.MockTransport(respond)) as transport: + return litellm.vector_stores.search(**request_kwargs, client=HTTPHandler(client=transport)) + + +@pytest.mark.parametrize("asynchronous", [False, True]) +@pytest.mark.parametrize( + "fields", + [{}, {"opensearch_embedding_field": "passage_vector", "opensearch_text_field": "passage"}], +) +@pytest.mark.parametrize( + "api_base,region_params,service,region", + [ + (SERVERLESS, {}, "aoss", "us-east-1"), + (SERVERLESS, {"aws_region_name": "us-west-2"}, "aoss", "us-east-1"), + ("https://search-docs-xyz.eu-west-1.es.amazonaws.com/", {}, "es", "eu-west-1"), + ("https://vpc-docs-xyz.cn-north-1.es.amazonaws.com.cn", {}, "es", "cn-north-1"), + ("https://abc123.us-gov-west-1.aoss-fips.amazonaws.com", {}, "aoss", "us-gov-west-1"), + ("https://abc123.aoss.eu-central-1.on.aws", {}, "aoss", "eu-central-1"), + ("https://search-docs-xyz.aos.ap-southeast-2.on.aws", {}, "es", "ap-southeast-2"), + ("https://search.docs.example", {"aws_region_name": "ap-south-1"}, "es", "ap-south-1"), + ], +) +@pytest.mark.asyncio +async def test_search_sends_a_signed_knn_query_and_returns_an_openai_results_page( + asynchronous: bool, + fields: Mapping[str, str], + api_base: str, + region_params: Mapping[str, str], + service: str, + region: str, +) -> None: + executor: Final = RecordingEmbeddingExecutor() + embedding_field: Final = fields.get("opensearch_embedding_field", "embedding") + text_field: Final = fields.get("opensearch_text_field", "text") + respond: Final = MagicMock(return_value=httpx.Response(200, json=_opensearch_hits(text_field, embedding_field))) + + response: Final = await _search( + asynchronous, + respond, + executor, + {**fields, **region_params, "api_base": api_base, "max_num_results": 3}, + ) + + assert response == {"object": "vector_store.search_results.page", "search_query": QUERY, "data": EXPECTED_RESULTS} + sent: Final = _sent(respond) + assert len(sent) == 1 + assert str(sent[0].url) == f"{api_base.rstrip('/')}/litellm-qa-docs/_search" + assert json.loads(sent[0].content) == { + "size": 3, + "query": {"knn": {embedding_field: {"vector": QUERY_VECTOR, "k": 3}}}, + "_source": {"excludes": [embedding_field]}, + } + assert _signature_scope(sent[0]) == (service, region) + executor.call.assert_called_once_with("embedding-alias", QUERY, {}) + + +@pytest.mark.asyncio +async def test_search_defaults_to_ten_results_like_openai() -> None: + respond: Final = MagicMock(return_value=httpx.Response(200, json={"hits": {"hits": []}})) + + response: Final = await _search(True, respond, RecordingEmbeddingExecutor(), {}) + + assert response == {"object": "vector_store.search_results.page", "search_query": QUERY, "data": []} + body: Final = json.loads(_sent(respond)[0].content) + assert (body["size"], body["query"]["knn"]["embedding"]["k"]) == (10, 10) + + +@pytest.mark.parametrize("asynchronous", [False, True]) +@pytest.mark.asyncio +async def test_a_list_query_is_embedded_and_echoed_as_one_string(asynchronous: bool) -> None: + executor: Final = RecordingEmbeddingExecutor() + respond: Final = MagicMock(return_value=httpx.Response(200, json={"hits": {"hits": []}})) + + response: Final = await _search(asynchronous, respond, executor, {"query": ["remote work", "abroad"]}) + + assert response["search_query"] == "remote work abroad" + executor.call.assert_called_once_with("embedding-alias", "remote work abroad", {}) + + +@pytest.mark.parametrize("asynchronous", [False, True]) +@pytest.mark.parametrize("status", [403, 404, 429, 500]) +@pytest.mark.asyncio +async def test_opensearch_error_keeps_its_status(asynchronous: bool, status: int) -> None: + def respond(request: httpx.Request) -> httpx.Response: + return httpx.Response(status, json={"status": status, "error": {"reason": "upstream refused", "type": "x"}}) + + with pytest.raises(openai.APIStatusError) as error: + await _search(asynchronous, respond, RecordingEmbeddingExecutor(), {}) + + assert error.value.status_code == status + + +@pytest.mark.parametrize( + "overrides,message", + [ + ({"filters": {"type": "eq", "key": "category", "value": "hr"}}, "filters"), + ({"ranking_options": {"score_threshold": 0.5}}, "ranking_options"), + ({"rewrite_query": True}, "rewrite_query"), + ({"max_num_results": 0}, "max_num_results"), + ({"max_num_results": 51}, "max_num_results"), + ({"max_num_results": "ten"}, "Invalid vector store search parameters"), + ({"litellm_embedding_model": None}, "litellm_embedding_model"), + ({"opensearch_text_feild": "passage"}, "opensearch_text_feild"), + ({"opensearch_text_field": 7}, "opensearch_text_field"), + ({"opensearch_index_name": ["docs"]}, "opensearch_index_name"), + ({"extra_body": {"min_score": 0.5}}, "extra_body"), + ({"api_base": None}, "api_base"), + ({"api_base": "abc123.us-east-1.aoss.amazonaws.com"}, "api_base"), + ({"api_base": "https://123456789012.aoss.us-east-1.on.aws"}, "per-account endpoint"), + ({"api_key": "opensearch-basic-auth"}, "api_key"), + ], +) +@pytest.mark.asyncio +async def test_unsupported_search_is_rejected_before_embedding_or_any_request( + overrides: Mapping[str, object], message: str +) -> None: + executor: Final = RecordingEmbeddingExecutor() + respond: Final = MagicMock(return_value=httpx.Response(200, json={"hits": {"hits": []}})) + + with pytest.raises(litellm.BadRequestError, match=message) as error: + await _search(True, respond, executor, overrides) + + assert error.value.status_code == 400 + respond.assert_not_called() + executor.call.assert_not_called() + + +_AMBIENT_AWS_ENV: Final = ( + "AWS_ACCESS_KEY_ID", + "AWS_SECRET_ACCESS_KEY", + "AWS_SESSION_TOKEN", + "AWS_PROFILE", + "AWS_PROFILE_NAME", + "AWS_ROLE_NAME", + "AWS_ROLE_ARN", + "AWS_SESSION_NAME", + "AWS_WEB_IDENTITY_TOKEN", + "AWS_WEB_IDENTITY_TOKEN_FILE", + "AWS_CONTAINER_CREDENTIALS_RELATIVE_URI", + "AWS_CONTAINER_CREDENTIALS_FULL_URI", +) + + +@pytest.mark.parametrize("asynchronous", [False, True]) +@pytest.mark.parametrize( + "credentials,env", + [ + pytest.param({"aws_session_name": "opensearch-no-credentials"}, {}, id="none-anywhere"), + pytest.param({"aws_profile_name": "opensearch-missing-profile"}, {}, id="missing-profile"), + pytest.param({}, {"AWS_ACCESS_KEY_ID": "AKIDPARTIAL"}, id="key-without-secret"), + pytest.param( + {"aws_access_key_id": "AKIDSTOREONLY"}, + {"AWS_ACCESS_KEY_ID": "AKIDAMBIENT", "AWS_SECRET_ACCESS_KEY": "ambient-secret"}, + id="store-key-without-secret-beside-ambient-credentials", + ), + pytest.param( + {"aws_secret_access_key": "store-secret-only"}, + {"AWS_ACCESS_KEY_ID": "AKIDAMBIENT", "AWS_SECRET_ACCESS_KEY": "ambient-secret"}, + id="store-secret-without-key-beside-ambient-credentials", + ), + ], +) +@pytest.mark.asyncio +async def test_missing_aws_credentials_are_a_401_before_embedding_or_any_request( + asynchronous: bool, + credentials: Mapping[str, str], + env: Mapping[str, str], + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + for name in _AMBIENT_AWS_ENV: + monkeypatch.delenv(name, raising=False) + for name in ("AWS_CONFIG_FILE", "AWS_SHARED_CREDENTIALS_FILE", "BOTO_CONFIG"): + config_file: Path = tmp_path / name + config_file.write_text("") + monkeypatch.setenv(name, str(config_file)) + monkeypatch.setenv("AWS_EC2_METADATA_DISABLED", "true") + for name, value in env.items(): + monkeypatch.setenv(name, value) + executor: Final = RecordingEmbeddingExecutor() + respond: Final = MagicMock(return_value=httpx.Response(200, json={"hits": {"hits": []}})) + + with pytest.raises(litellm.AuthenticationError, match="No usable AWS credentials") as error: + await _search( + asynchronous, + respond, + executor, + {"aws_access_key_id": None, "aws_secret_access_key": None, **credentials}, + ) + + assert error.value.status_code == 401 + respond.assert_not_called() + executor.call.assert_not_called() + + +class _FailingCredentialSource(BaseAWSLLM): + def resolve_credentials(self, auth_params: AwsAuthParams, aws_region_name: str | None) -> Credentials: + raise CredentialRetrievalError(provider="container-role", error_msg=CREDENTIAL_SOURCE_OUTPUT) + + +def test_failing_credential_source_is_a_401_that_keeps_its_output_out_of_the_reply() -> None: + config: Final = OpenSearchVectorStoreConfig(aws=_FailingCredentialSource()) + + with pytest.raises(litellm.AuthenticationError, match="the AWS credential source failed") as error: + config.validate_environment({}, GenericLiteLLMParams(api_base=SERVERLESS)) + + assert error.value.status_code == 401 + assert CREDENTIAL_SOURCE_OUTPUT not in str(error.value) + + +@pytest.mark.parametrize("asynchronous", [False, True]) +@pytest.mark.asyncio +async def test_opensearch_index_name_picks_the_index_so_the_store_id_can_differ(asynchronous: bool) -> None: + respond: Final = MagicMock(return_value=httpx.Response(200, json=_opensearch_hits("text", "embedding"))) + + response: Final = await _search( + asynchronous, + respond, + RecordingEmbeddingExecutor(), + {"vector_store_id": "east-docs", "opensearch_index_name": "docs/v2"}, + ) + + assert response["data"] == EXPECTED_RESULTS + assert str(_sent(respond)[0].url) == f"{SERVERLESS}/docs%2Fv2/_search" + + +@pytest.mark.parametrize( + "unsupported", + [ + {"filters": {"type": "eq", "key": "category", "value": "hr"}}, + {"ranking_options": {"score_threshold": 0.5}}, + {"rewrite_query": True}, + ], +) +@pytest.mark.asyncio +async def test_unsupported_search_params_are_dropped_under_drop_params(unsupported: Mapping[str, object]) -> None: + respond: Final = MagicMock(return_value=httpx.Response(200, json=_opensearch_hits("text", "embedding"))) + + response: Final = await _search( + True, respond, RecordingEmbeddingExecutor(), {**unsupported, "drop_params": True, "max_num_results": 3} + ) + + assert response["data"] == EXPECTED_RESULTS + assert json.loads(_sent(respond)[0].content) == { + "size": 3, + "query": {"knn": {"embedding": {"vector": QUERY_VECTOR, "k": 3}}}, + "_source": {"excludes": ["embedding"]}, + } + + +@pytest.mark.asyncio +async def test_no_op_values_of_unsupported_params_are_accepted() -> None: + def respond(request: httpx.Request) -> httpx.Response: + return httpx.Response(200, json=_opensearch_hits("text", "embedding")) + + response: Final = await _search( + True, respond, RecordingEmbeddingExecutor(), {"filters": {}, "ranking_options": {}, "rewrite_query": False} + ) + + assert response["data"] == EXPECTED_RESULTS + + +def test_create_is_rejected_because_the_store_is_search_only() -> None: + with pytest.raises(litellm.BadRequestError, match="search-only"): + litellm.vector_stores.create(custom_llm_provider="opensearch") + + +async def _store_operation( + asynchronous: bool, operation: str, respond: Callable[[httpx.Request], httpx.Response] +) -> object: + request_kwargs: Final = {**BASE_PARAMS, "vector_store_id": "litellm-qa-docs", "custom_llm_provider": "opensearch"} + if asynchronous: + return await getattr(vector_stores, f"a{operation}")( + **request_kwargs, client=AsyncHTTPHandler(transport=httpx.MockTransport(respond)) + ) + with httpx.Client(transport=httpx.MockTransport(respond)) as transport: + return getattr(vector_stores, operation)(**request_kwargs, client=HTTPHandler(client=transport)) + + +@pytest.mark.parametrize("asynchronous", [False, True]) +@pytest.mark.parametrize("operation", ["retrieve", "update", "delete"]) +@pytest.mark.asyncio +async def test_store_operations_besides_search_send_nothing_to_the_index(asynchronous: bool, operation: str) -> None: + respond: Final = MagicMock(return_value=httpx.Response(200, json={"acknowledged": True})) + + with pytest.raises(litellm.BadRequestError, match="search-only") as error: + await _store_operation(asynchronous, operation, respond) + + assert error.value.status_code == 400 + respond.assert_not_called() + + +async def _list_stores(asynchronous: bool, respond: Callable[[httpx.Request], httpx.Response]) -> object: + request_kwargs: Final = {**BASE_PARAMS, "custom_llm_provider": "opensearch"} + if asynchronous: + return await vector_stores.alist( + **request_kwargs, client=AsyncHTTPHandler(transport=httpx.MockTransport(respond)) + ) + with httpx.Client(transport=httpx.MockTransport(respond)) as transport: + return vector_stores.list(**request_kwargs, client=HTTPHandler(client=transport)) + + +@pytest.mark.parametrize("asynchronous", [False, True]) +@pytest.mark.asyncio +async def test_listing_stores_sends_nothing_to_the_endpoint(asynchronous: bool) -> None: + respond: Final = MagicMock(return_value=httpx.Response(200, json={"cluster_name": "docs", "tagline": "OpenSearch"})) + + with pytest.raises(litellm.BadRequestError, match="search-only") as error: + await _list_stores(asynchronous, respond) + + assert error.value.status_code == 400 + respond.assert_not_called() + + +@pytest.mark.parametrize( + "response,status", + [ + (httpx.Response(403, json={"status": 403, "error": {"type": "Forbidden"}}), 403), + (httpx.Response(200, json={"took": 1}), 502), + (httpx.Response(200, text="gateway"), 502), + ], +) +def test_response_without_search_hits_is_an_error_not_an_empty_page(response: httpx.Response, status: int) -> None: + logging: Final = Logging( + model="", + messages=[], + stream=False, + call_type="vector_store_search", + start_time=datetime(2026, 10, 20, tzinfo=timezone.utc), + litellm_call_id="opensearch-test", + function_id="opensearch-test", + ) + + with pytest.raises(BaseLLMException) as error: + OpenSearchVectorStoreConfig().transform_search_vector_store_response(response, logging) + + assert error.value.status_code == status diff --git a/tests/unit/proxy/public_endpoints/test_public_endpoints.py b/tests/unit/proxy/public_endpoints/test_public_endpoints.py index e694bca3c8d..da608dbca0e 100644 --- a/tests/unit/proxy/public_endpoints/test_public_endpoints.py +++ b/tests/unit/proxy/public_endpoints/test_public_endpoints.py @@ -566,6 +566,7 @@ ADD_MODEL_UNLISTED_PROVIDERS: Final = frozenset( "mongodb", "nano-gpt", "neosantara", + "opensearch", "parasail", "pinstripes", "poe", diff --git a/tests/unit/proxy/rag_endpoints/test_rag_endpoints.py b/tests/unit/proxy/rag_endpoints/test_rag_endpoints.py index d9f3ceda1ed..bf4b5d5d77a 100644 --- a/tests/unit/proxy/rag_endpoints/test_rag_endpoints.py +++ b/tests/unit/proxy/rag_endpoints/test_rag_endpoints.py @@ -1671,6 +1671,62 @@ def test_rag_query_rejects_caller_embedding_selection_params(client_internal_use assert blocked_key in str(response.json()) +@pytest.mark.parametrize( + ("provider", "expected_status", "expected_searches"), + [("opensearch", 400, 0), ("opensearch/x", 400, 0), ("openai", 200, 1)], +) +def test_rag_query_rejects_caller_opensearch_index_name_only_on_opensearch_stores( + client_internal_user, provider, expected_status, expected_searches +): + """ + Regression: retrieval_config must not pick the index an OpenSearch store searches, + while other providers keep ignoring the key as they did before OpenSearch support. + """ + import litellm + from litellm import Router + from litellm.types.vector_stores import VectorStoreSearchResponse + + store = { + "vector_store_id": "kb", + "custom_llm_provider": provider, + "litellm_params": {"vector_store_id": "kb", "custom_llm_provider": provider}, + } + registry = MagicMock() + registry.get_litellm_managed_vector_store_from_registry.return_value = store + fake_search = AsyncMock( + return_value=VectorStoreSearchResponse(object="vector_store.search_results.page", search_query="q", data=[]) + ) + router = Router( + model_list=[ + { + "model_name": "gpt-4o-mini", + "litellm_params": {"model": "openai/gpt-4o-mini", "api_key": "sk-test", "mock_response": "hi"}, + } + ] + ) + + with ( + patch("litellm.vector_stores.asearch", new=fake_search), # test-quality-ok: the search boundary under test + patch.object(litellm, "vector_store_registry", registry), # test-quality-ok: seeds the store under test + patch("litellm.proxy.proxy_server.llm_router", router), # test-quality-ok: mock-response router for completion + patch( # test-quality-ok: store access is not under test, so the request reaches the search boundary + "litellm.proxy.vector_store_endpoints.utils.can_user_access_vector_store", + new=AsyncMock(return_value=True), + ), + ): + response = client_internal_user.post( + "/v1/rag/query", + json={ + "model": "gpt-4o-mini", + "messages": [{"role": "user", "content": "hello"}], + "retrieval_config": {"vector_store_id": "kb", "opensearch_index_name": "another-teams-index"}, + }, + ) + + assert response.status_code == expected_status, response.json() + assert fake_search.await_count == expected_searches + + EICAR = r"X5O!P%@AP[4\PZX54(P^)7CC)7}$EICAR-STANDARD-ANTIVIRUS-TEST-FILE!$H+H*" INGEST_REQUEST = '{"ingest_options":{"vector_store":{"custom_llm_provider":"openai"}}}' diff --git a/tests/unit/proxy/vector_store_endpoints/test_vector_store_endpoints.py b/tests/unit/proxy/vector_store_endpoints/test_vector_store_endpoints.py index 48642a9b50d..e128b5eb1f1 100644 --- a/tests/unit/proxy/vector_store_endpoints/test_vector_store_endpoints.py +++ b/tests/unit/proxy/vector_store_endpoints/test_vector_store_endpoints.py @@ -3066,6 +3066,58 @@ def test_vector_store_search_rejects_caller_embedding_selection_params(blocked_k assert blocked_key in str(response.json()) +@pytest.mark.parametrize( + ("provider", "expected_status", "expected_searches"), + [("opensearch", 400, 0), ("opensearch/x", 400, 0), ("openai", 200, 1)], +) +def test_vector_store_search_rejects_caller_opensearch_index_name_only_on_opensearch_stores( + provider, expected_status, expected_searches +): + """ + Regression: a caller must not pick the index an OpenSearch store searches, while + other providers keep ignoring the key as they did before OpenSearch support. + """ + from fastapi.testclient import TestClient + + from litellm.proxy.auth.user_api_key_auth import user_api_key_auth + from litellm.proxy.proxy_server import app + from litellm.types.vector_stores import VectorStoreSearchResponse + + store = { + "vector_store_id": "kb", + "custom_llm_provider": provider, + "litellm_params": {"vector_store_id": "kb", "custom_llm_provider": provider}, + } + registry = MagicMock() + registry.get_litellm_managed_vector_store_from_registry.return_value = store + fake_search = AsyncMock( + return_value=VectorStoreSearchResponse(object="vector_store.search_results.page", search_query="q", data=[]) + ) + mock_auth = UserAPIKeyAuth(user_id="test_internal_user", user_role=LitellmUserRoles.INTERNAL_USER.value) + original_overrides = app.dependency_overrides.copy() + app.dependency_overrides[user_api_key_auth] = lambda: mock_auth + try: + with ( + patch("litellm.vector_stores.main.asearch", new=fake_search), # test-quality-ok: the search boundary under test + patch.object(litellm, "vector_store_registry", registry), # test-quality-ok: seeds the store under test + patch( # test-quality-ok: store access is not under test, so the request reaches the search boundary + "litellm.proxy.vector_store_endpoints.utils.can_user_access_vector_store", + new=AsyncMock(return_value=True), + ), + ): + router = litellm.Router(model_list=[]) + with patch("litellm.proxy.proxy_server.llm_router", router): # test-quality-ok: router built over the fake search + response = TestClient(app).post( + "/v1/vector_stores/kb/search", + json={"query": "q", "opensearch_index_name": "another-teams-index"}, + ) + finally: + app.dependency_overrides = original_overrides + + assert response.status_code == expected_status, response.json() + assert fake_search.await_count == expected_searches + + class TestConfigOwnedVectorStores: """Stores declared under ``vector_store_registry`` in config.yaml are owned by the config file""" diff --git a/ui/litellm-dashboard/src/app/(dashboard)/vector-stores/_components/CreateVectorStore.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/vector-stores/_components/CreateVectorStore.test.tsx index ba3d1d0fb7d..f8e2feaa635 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/vector-stores/_components/CreateVectorStore.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/vector-stores/_components/CreateVectorStore.test.tsx @@ -18,6 +18,7 @@ vi.mock("@/components/vector_store_providers", () => ({ AZURE_OPENAI: "Azure OpenAI", S3Vectors: "AWS S3 Vectors", Valkey: "Valkey", + OpenSearch: "Amazon OpenSearch", }, vectorStoreProviderMap: { BEDROCK: "bedrock", @@ -25,6 +26,7 @@ vi.mock("@/components/vector_store_providers", () => ({ AZURE_OPENAI: "azure_openai", S3Vectors: "s3_vectors", Valkey: "valkey", + OpenSearch: "opensearch", }, vectorStoreProviderLogoMap: { "Amazon Bedrock": "https://example.com/bedrock.png", @@ -32,6 +34,7 @@ vi.mock("@/components/vector_store_providers", () => ({ "Azure OpenAI": "https://example.com/azure.png", "AWS S3 Vectors": "https://example.com/aws.png", Valkey: "https://example.com/valkey.svg", + "Amazon OpenSearch": "https://example.com/aws.svg", }, getProviderSpecificFields: vi.fn((provider: string) => { if (provider === "s3_vectors") { @@ -254,7 +257,7 @@ describe("CreateVectorStore", () => { expect(screen.getByText("Vector Store Created Successfully")).toBeInTheDocument(); }); - it("should exclude valkey from the provider dropdown since it has no RAG ingestion", async () => { + it("should exclude search-only valkey and opensearch from the provider dropdown since they have no RAG ingestion", async () => { render(); const providerSelect = screen.getByRole("combobox"); @@ -267,6 +270,7 @@ describe("CreateVectorStore", () => { expect(screen.getByText("AWS S3 Vectors")).toBeInTheDocument(); }); expect(screen.queryByText("Valkey")).not.toBeInTheDocument(); + expect(screen.queryByText("Amazon OpenSearch")).not.toBeInTheDocument(); }); it("should display S3 Vectors provider-specific fields when selected", async () => { diff --git a/ui/litellm-dashboard/src/app/(dashboard)/vector-stores/_components/CreateVectorStore.tsx b/ui/litellm-dashboard/src/app/(dashboard)/vector-stores/_components/CreateVectorStore.tsx index 2422e83355d..2d7375431cd 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/vector-stores/_components/CreateVectorStore.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/vector-stores/_components/CreateVectorStore.tsx @@ -36,7 +36,7 @@ const ACCEPTED_DOCUMENT_TYPES = [ const MAX_DOCUMENT_BYTES = 50 * 1024 * 1024; -const RAG_INGEST_UNSUPPORTED_PROVIDERS = new Set(["valkey"]); +const RAG_INGEST_UNSUPPORTED_PROVIDERS = new Set(["valkey", "opensearch"]); const providerItems = Object.entries(VectorStoreProviders) .filter(([providerEnum]) => !RAG_INGEST_UNSUPPORTED_PROVIDERS.has(vectorStoreProviderMap[providerEnum])) diff --git a/ui/litellm-dashboard/src/app/(dashboard)/vector-stores/_components/VectorStoreForm.integration.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/vector-stores/_components/VectorStoreForm.integration.test.tsx index d7ec325653f..f571ec6b9a0 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/vector-stores/_components/VectorStoreForm.integration.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/vector-stores/_components/VectorStoreForm.integration.test.tsx @@ -1,4 +1,4 @@ -import { render, screen } from "@testing-library/react"; +import { fireEvent, render, screen } from "@testing-library/react"; import userEvent from "@testing-library/user-event"; import { beforeEach, describe, expect, it, vi } from "vitest"; @@ -134,6 +134,70 @@ describe("VectorStoreForm submit payload", () => { expect(createdPayload().custom_llm_provider).toBe("milvus"); }); + it("registers an existing opensearch index with the proxy embedding model and default field names", async () => { + const user = setupUser(); + renderForm(); + + await chooseProvider(user, "Amazon OpenSearch"); + fireEvent.change(screen.getByPlaceholderText("my-knn-index (OpenSearch index name)"), { + target: { value: "litellm-qa-docs" }, + }); + fireEvent.change(screen.getByPlaceholderText("https://abc123.us-east-1.aoss.amazonaws.com"), { + target: { value: "https://yzatq55tpsuvlwozcal4.us-east-1.aoss.amazonaws.com" }, + }); + await chooseFromSelect(user, 1, "text-embedding-3-small"); + await submit(user); + + const expectedLitellmParams = { + api_base: "https://yzatq55tpsuvlwozcal4.us-east-1.aoss.amazonaws.com", + aws_region_name: undefined, + litellm_embedding_model: "text-embedding-3-small", + opensearch_embedding_field: "embedding", + opensearch_text_field: "text", + opensearch_index_name: undefined, + }; + + await vi.waitFor(() => expect(mockCreate).toHaveBeenCalledTimes(1)); + expect(createdPayload().custom_llm_provider).toBe("opensearch"); + expect(createdPayload().vector_store_id).toBe("litellm-qa-docs"); + expect(createdPayload().litellm_params).toStrictEqual(expectedLitellmParams); + }); + + it("registers an opensearch index under its own name when the vector store id differs", async () => { + const user = setupUser(); + renderForm(); + + await chooseProvider(user, "Amazon OpenSearch"); + fireEvent.change(screen.getByPlaceholderText("my-knn-index (OpenSearch index name)"), { + target: { value: "east-docs" }, + }); + fireEvent.change(screen.getByPlaceholderText("https://abc123.us-east-1.aoss.amazonaws.com"), { + target: { value: "https://yzatq55tpsuvlwozcal4.us-east-1.aoss.amazonaws.com" }, + }); + fireEvent.change(screen.getByPlaceholderText("my-knn-index"), { target: { value: "docs" } }); + await chooseFromSelect(user, 1, "text-embedding-3-small"); + await submit(user); + + await vi.waitFor(() => expect(mockCreate).toHaveBeenCalledTimes(1)); + expect(createdPayload().vector_store_id).toBe("east-docs"); + expect(createdPayload().litellm_params).toMatchObject({ opensearch_index_name: "docs" }); + }); + + it("blocks an opensearch submit that has no endpoint or embedding model", async () => { + const user = setupUser(); + renderForm(); + + await chooseProvider(user, "Amazon OpenSearch"); + fireEvent.change(screen.getByPlaceholderText("my-knn-index (OpenSearch index name)"), { + target: { value: "litellm-qa-docs" }, + }); + await submit(user); + + expect(await screen.findByText("Please input the opensearch endpoint")).toBeInTheDocument(); + expect(screen.getByText("Please select the embedding model")).toBeInTheDocument(); + expect(mockCreate).not.toHaveBeenCalled(); + }); + it("sends a provider field's seeded default even when the user never touches it", async () => { const user = setupUser(); renderForm(); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/vector-stores/_components/VectorStoreForm.tsx b/ui/litellm-dashboard/src/app/(dashboard)/vector-stores/_components/VectorStoreForm.tsx index d82f56f3870..bf5b925eb16 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/vector-stores/_components/VectorStoreForm.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/vector-stores/_components/VectorStoreForm.tsx @@ -36,7 +36,7 @@ import { useZodForm } from "@/lib/forms/useZodForm"; import { SearchSelect } from "@/components/shared/SearchSelect"; import { credentialOptions } from "@/components/shared/credentialOptions"; -const EMBEDDING_MODEL_RENAME_PROVIDERS = new Set(["milvus", "valkey", "mongodb"]); +const EMBEDDING_MODEL_RENAME_PROVIDERS = new Set(["milvus", "valkey", "mongodb", "opensearch"]); export const buildVectorStoreLitellmParams = ( provider: string, @@ -83,6 +83,9 @@ const PROVIDER_FIELD_NAMES = [ "valkey_ssl", "valkey_text_field", "valkey_embedding_field", + "opensearch_embedding_field", + "opensearch_text_field", + "opensearch_index_name", ] as const; type ProviderFieldName = (typeof PROVIDER_FIELD_NAMES)[number]; @@ -119,6 +122,9 @@ const vectorStoreShape = { valkey_ssl: optionalText, valkey_text_field: optionalText, valkey_embedding_field: optionalText, + opensearch_embedding_field: optionalText, + opensearch_text_field: optionalText, + opensearch_index_name: optionalText, }; const vectorStoreSchema = z.object(vectorStoreShape).superRefine((values, ctx) => { @@ -143,6 +149,7 @@ const VECTOR_STORE_ID_PLACEHOLDERS: Record = { "vertex_ai/search_api": 'my-datastore_1234567890 (data store ID from Vertex AI / "Agent Search" console)', valkey: "my-search-index (FT index name in Valkey)", mongodb: "my-vector-index (MongoDB Vector Search index name)", + opensearch: "my-knn-index (OpenSearch index name)", }; const VERTEX_SEARCH_API_WITH_ENGINE_PLACEHOLDER = "Any identifier you'll use to reference this in LiteLLM"; @@ -159,6 +166,8 @@ const EMPTY_VALUES: VectorStoreFormValues = { valkey_ssl: "false", valkey_text_field: "text", valkey_embedding_field: "embedding", + opensearch_embedding_field: "embedding", + opensearch_text_field: "text", }; const labelWithHint = (label: string, hint: string): React.ReactNode => ( @@ -376,6 +385,35 @@ const VectorStoreForm: React.FC = ({ )} + {selectedProvider === "opensearch" && ( + + + OpenSearch Setup Required + +

+ LiteLLM searches a k-NN index you already have in an OpenSearch Serverless collection or an + OpenSearch Service domain. It does not create the index or upload documents. Before creating this + vector store, make sure: +

+
    +
  1. + The index has a knn_vector field holding each document's embedding. Enter the index name as + the Vector Store ID, or under Index Name when another endpoint already uses that ID +
    2. +
  2. + You know which embedding model created those embeddings. That model must be added to this proxy + under Models so you can pick it below. Using a different model returns wrong results +
    3. +
  3. + The proxy has AWS credentials that can read the index, either its own (environment variables or + an IAM role) or a saved credential picked under Existing Credentials. A Serverless collection + needs a data access policy for them, and a domain needs its access policy to allow them +
    4. +
+ + + )} + {selectedProvider === "vertex_rag_engine" && ( diff --git a/ui/litellm-dashboard/src/components/vector_store_providers.test.tsx b/ui/litellm-dashboard/src/components/vector_store_providers.test.tsx index 32d0f5dccc0..b4d890475c7 100644 --- a/ui/litellm-dashboard/src/components/vector_store_providers.test.tsx +++ b/ui/litellm-dashboard/src/components/vector_store_providers.test.tsx @@ -68,6 +68,32 @@ describe("getVectorStoreProviderLogoAndName", () => { expect(byName("mongodb_num_candidates")).toMatchObject({ required: false }); }); + it("registers opensearch with its endpoint, optional region, proxy embedding model, and default field names", () => { + const fields = getProviderSpecificFields("opensearch"); + const byName = (name: string) => fields.find((field) => field.name === name); + + expect(getVectorStoreProviderLogoAndName("opensearch")).toEqual({ + logo: expect.stringContaining("aws"), + displayName: VectorStoreProviders.OpenSearch, + }); + expect(fields.map((field) => field.name)).toEqual([ + "api_base", + "aws_region_name", + "embedding_model", + "opensearch_embedding_field", + "opensearch_text_field", + "opensearch_index_name", + ]); + expect(byName("api_base")).toMatchObject({ required: true }); + expect(byName("aws_region_name")).toMatchObject({ required: false }); + expect(byName("embedding_model")).toMatchObject({ type: "select", required: true }); + expect(byName("embedding_model")).not.toHaveProperty("options"); + expect(byName("opensearch_embedding_field")).toMatchObject({ required: false, initialValue: "embedding" }); + expect(byName("opensearch_text_field")).toMatchObject({ required: false, initialValue: "text" }); + expect(byName("opensearch_index_name")).toMatchObject({ required: false }); + expect(byName("opensearch_index_name")).not.toHaveProperty("initialValue"); + }); + it("registers valkey in the provider, logo, and field maps", () => { expect(vectorStoreProviderMap.Valkey).toBe("valkey"); expect(vectorStoreProviderLogoMap[VectorStoreProviders.Valkey]).toContain("valkey"); diff --git a/ui/litellm-dashboard/src/components/vector_store_providers.tsx b/ui/litellm-dashboard/src/components/vector_store_providers.tsx index 79e711f6f98..e953d6e3db6 100644 --- a/ui/litellm-dashboard/src/components/vector_store_providers.tsx +++ b/ui/litellm-dashboard/src/components/vector_store_providers.tsx @@ -1,4 +1,5 @@ import { getProviderLogoAndName, Providers, providerLogoMap } from "@/components/provider_info_helpers"; +import awsLogo from "../../public/assets/logos/aws.svg"; import milvusLogo from "../../public/assets/logos/milvus.svg"; import mongodbLogo from "../../public/assets/logos/mongodb.svg"; import postgresqlLogo from "../../public/assets/logos/postgresql.svg"; @@ -8,6 +9,7 @@ import valkeyLogo from "../../public/assets/logos/valkey.svg"; export enum VectorStoreProviders { Bedrock = "Amazon Bedrock", S3Vectors = "Amazon S3 Vectors", + OpenSearch = "Amazon OpenSearch", PgVector = "PostgreSQL pgvector (LiteLLM Connector)", VertexRagEngine = "Vertex AI RAG Engine", VertexAiSearch = "Vertex AI Search", @@ -28,6 +30,7 @@ export const vectorStoreProviderMap: Record = { Milvus: "milvus", MongoDB: "mongodb", S3Vectors: "s3_vectors", + OpenSearch: "opensearch", Valkey: "valkey", }; @@ -41,6 +44,7 @@ export const vectorStoreProviderLogoMap: Record = { [VectorStoreProviders.Milvus]: milvusLogo.src, [VectorStoreProviders.MongoDB]: mongodbLogo.src, [VectorStoreProviders.S3Vectors]: s3VectorLogo.src, + [VectorStoreProviders.OpenSearch]: awsLogo.src, [VectorStoreProviders.Valkey]: valkeyLogo.src, }; @@ -313,6 +317,64 @@ export const vectorStoreProviderFields: Record initialValue: "embedding", }, ], + opensearch: [ + { + name: "api_base", + label: "OpenSearch Endpoint", + tooltip: + "The endpoint of your OpenSearch Serverless collection or OpenSearch Service domain, starting with https://", + placeholder: "https://abc123.us-east-1.aoss.amazonaws.com", + required: true, + type: "text", + }, + { + name: "aws_region_name", + label: "AWS Region (optional)", + tooltip: + "The AWS region of the collection or domain. Leave blank to use the region in the endpoint, and set it when the endpoint is a custom domain", + placeholder: "us-east-1", + required: false, + type: "text", + }, + { + name: "embedding_model", + label: "Embedding Model", + tooltip: + "The embedding model on this proxy that created the vectors already stored in your index. LiteLLM embeds every search query with it, so it must be the same model. Add it under Models first if it is not listed", + placeholder: "text-embedding-3-small", + required: true, + type: "select", + }, + { + name: "opensearch_embedding_field", + label: "Vector Field Name", + tooltip: + "The top-level knn_vector field in your index mapping that holds each document's embedding (default: embedding)", + placeholder: "embedding", + required: false, + type: "text", + initialValue: "embedding", + }, + { + name: "opensearch_text_field", + label: "Text Field", + tooltip: + "The top-level field in each document that holds its readable text. LiteLLM returns this text in search results (default: text)", + placeholder: "text", + required: false, + type: "text", + initialValue: "text", + }, + { + name: "opensearch_index_name", + label: "Index Name (optional)", + tooltip: + "The OpenSearch index to search. Leave blank to use the Vector Store ID, and set it to register same-named indexes from two endpoints under different IDs", + placeholder: "my-knn-index", + required: false, + type: "text", + }, + ], s3_vectors: [ { name: "vector_bucket_name", From 6cdb9959df49132e7402b723ea6aeb79fd01242f Mon Sep 17 00:00:00 2001 From: Mateo Wang <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 29 Oct 2026 08:33:36 -0700 Subject: [PATCH 07/13] feat(router): opt in to non-chat deployment filtering (#46257) * feat(proxy): reject chat routes on models declared non-chat via model_info.supported_endpoints A model whose every deployment sets a non-chat model_info.mode and a supported_endpoints list with no chat endpoint now answers /v1/chat/completions, /v1/messages and /v1/responses with a 400 before any provider call. Models without both fields keep today's behavior. * fix(proxy): let fallbacks answer and cover more chat routes in the non-chat model check The check now steps aside when a request, key, team or router fallback would answer, skips such models in the rate-limit fallback loop, ignores mode and supported_endpoints copied from the cost map by old Admin UI saves, treats a blank mode as unset, and also covers /v1/completions, the Responses WebSocket, Gemini generateContent and /v1beta/interactions * fix(proxy): only step aside for fallbacks that can answer chat, gate the Responses WebSocket after alias resolution, and compare declarations to the loaded catalog * fix(proxy): gate the Responses WebSocket with the router after aliases resolve, and ignore client-set internal fallback flags * fix(proxy): walk fallback groups breadth first and release reservations on a rejected WebSocket Each fallback group is looked up once however often the request repeats it, and a WebSocket closed by the endpoint gate runs the failure hook so rate-limit reservations are released * fix(proxy): walk fallback layers without recursion and gate the compact Responses route The fallback walk now folds one layer at a time, so the recursion check passes, and /v1/responses/compact on a model declared non-chat gets the same 400 as the other chat-shaped routes * fix(proxy): keep chat on rows whose cost map entry serves chat A row that carries the cost map key of an entry that serves chat counts as declaring nothing, so a catalog edit that adds a chat path never turns old Admin UI copies into 400s * fix(proxy): check chat sent to a deployment id or backend model name * fix(proxy): check a model the read-through loads from the database * fix(proxy): step aside for fallbacks the check cannot follow and wildcards that serve a backend name A fallback target with no deployments of its own, a client routing stamp, and a wildcard that does not match the requested backend name no longer turn a working request into a 400 * fix(proxy): ignore a client model_group that is not a string when building rate-limit headers A client metadata.model_group sent as an int or a list crashed the error handler with a 500 once the embedding-only check rejected the call before routing. The header helper now treats a non-string stamp as no model group. Drops the pattern clause in the backend-name lookup, since the group lookup already returns wildcard matches. Adds the third audit round's routing-stamp cells. * fix(proxy): keep a non-string client model_group out of spend rows * fix(proxy): validate standard fallbacks before the typed fallback-group lookup * feat(router): opt in to non-chat deployment filtering * fix(router): skip unresolved local rate-limit fallbacks * fix(router): ignore caller-provided endpoint selector values * fix(router): normalize declared chat endpoint prefixes * test(router): cover preflight error-value behavior * test(router): derive replay metadata defaults from model info * fix(router): preserve fallback aliases and check synchronous text calls * test(routing): rewrite the WebSocket model with a pre-call callback, since main no longer runs guardrails at socket admission * test(routing): keep the endpoint gate's owned proxies on the worker's database * fix(router): keep the endpoint gate on model_info patches that omit it and on same-group stream retries A partial PATCH of model_info reset a stored enforce_supported_endpoints true back to false, because the field defaulted to False and the update merges model_dump(exclude_none=True). The field now defaults to None so an omitted flag keeps the stored value The same-group retry of a stream dropped before content called the provider without the endpoint selection context, so it could land on an enforced embedding-only deployment of the same group. It now runs inside endpoint_selection like every other router attempt * fix(router): probe rate-limit fallbacks through get_routable_deployments so the router surface does not grow Main's router fence lets code outside the router use only listed Router members and router_utils modules. The proxy's fallback probe now calls the listed get_routable_deployments with the request's call type, and the gate's 400 moves to litellm.exceptions as EndpointNotSupportedByModelError. It stays separate from main's UnsupportedEndpointError, which the router retries on the group's other deployments, while the gate refuses the whole group * test(responses): cover the WebSocket error frame for a model that refuses the endpoint * fix(proxy): let the rate-limit fallback probe try a fallback this worker has not loaded yet The probe now skips a fallback only when the endpoint check removed every deployment it has, so a model added on another worker still reaches route_request and its database lookup * fix(proxy): keep key router settings off the Responses WebSocket, read unrooted or single-string endpoint lists, and add status to the refusal frame * fix(router): warn when enforce_supported_endpoints has no usable list and match chat endpoints regardless of case, trailing slash, or Gemini path * test(router): type the load-warning test's caplog and router * test(routing): match the endpoint gate's integration tests to base fallback behavior on unloaded fallbacks and the Responses WebSocket --------- Co-authored-by: Mateo --- litellm/exceptions.py | 11 + litellm/proxy/common_request_processing.py | 30 + litellm/proxy/common_utils/callback_utils.py | 5 +- .../proxy/response_api_endpoints/endpoints.py | 16 +- litellm/proxy/route_llm_request.py | 22 + .../spend_tracking/spend_tracking_utils.py | 8 +- litellm/router.py | 60 +- litellm/router_utils/endpoint_eligibility.py | 120 ++ litellm/types/router.py | 1 + .../llm_nonconversational.yaml | 1 + ..._embedding_only_supported_endpoints_e2e.py | 119 ++ tests/e2e/models.py | 3 +- .../test_non_chat_model_endpoint_gate.py | 1515 +++++++++++++++++ .../proxy/common_utils/test_callback_utils.py | 13 + .../test_model_management_endpoints.py | 27 + .../response_api_endpoints/test_endpoints.py | 16 + .../test_spend_tracking_utils.py | 22 + .../proxy/test_common_request_processing.py | 58 +- tests/unit/proxy/test_route_llm_request.py | 25 + .../router_utils/test_endpoint_eligibility.py | 139 ++ tests/unit/test_router/test_router.py | 113 ++ ui/litellm-dashboard/src/lib/http/schema.d.ts | 2 + 22 files changed, 2300 insertions(+), 26 deletions(-) create mode 100644 litellm/router_utils/endpoint_eligibility.py create mode 100644 tests/e2e/llm_translation/test_embedding_only_supported_endpoints_e2e.py create mode 100644 tests/integration/routing/test_non_chat_model_endpoint_gate.py create mode 100644 tests/unit/router_utils/test_endpoint_eligibility.py diff --git a/litellm/exceptions.py b/litellm/exceptions.py index 61e3a8fb5dc..28211b3dab2 100644 --- a/litellm/exceptions.py +++ b/litellm/exceptions.py @@ -323,6 +323,17 @@ class UnsupportedEndpointError(BadRequestError): pass +class EndpointNotSupportedByModelError(BadRequestError): + def __init__(self, message: str, model: str, llm_provider: str = "") -> None: + super().__init__( # pyright: ignore[reportUnknownMemberType] # inherited public exception constructor is untyped + message=message, model=model, llm_provider=llm_provider + ) + self.message = message + self.type = "invalid_request_error" + self.param = "model" + self.code = "400" + + class UnprocessableEntityError(openai.UnprocessableEntityError): def __init__( self, diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index fc29a6e3a08..4884e75af4e 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -129,6 +129,7 @@ from litellm.proxy.guardrails.auto_router_compression import arm_pre_call as _ar from litellm.proxy.hooks.parallel_request_limiter_v3 import get_or_create_request_stash from litellm.proxy.native_compaction import with_proxy_compaction_executor from litellm.proxy.route_llm_request import ( + endpoint_unsupported_by_model, route_request, ) from litellm.proxy.telemetry.request_context import bind_call_id as bind_telemetry_call_id @@ -871,6 +872,17 @@ _DECODED_JSON: Final = TypeAdapter(object) _ROUTER_SETTINGS_ADAPTER: Final[TypeAdapter[Mapping[str, object] | None]] = TypeAdapter(Mapping[str, object] | None) +def _fallback_probe_data( + pristine: object, fallback_model: str, key_aliases: object, group_alias: object +) -> Mapping[str, object]: + return { + **_REQUEST_FIELDS.validate_python(pristine), + "model": resolve_served_model( + fallback_model, _DECODED_JSON.validate_python(litellm.model_alias_map), key_aliases, group_alias + ), + } + + def _request_model(data: Mapping[str, object]) -> str | list[str] | None: try: return _REQUEST_MODEL.validate_python(data.get("model"), strict=True) @@ -2676,6 +2688,17 @@ class ProxyBaseLLMRequestProcessing: original_model, fallback_models, ) + raw_router_settings: Final = _DECODED_JSON.validate_python( + _REQUEST_FIELDS.validate_python(rate_limited_data).get("router_settings_override") + ) + group_alias: Final = _DECODED_JSON.validate_python( + _REQUEST_FIELDS.validate_python(raw_router_settings).get("model_group_alias") + if isinstance(raw_router_settings, Mapping) + else None + ) + key_aliases: Final[object] = cast( # cast-ok: auth aliases are dynamically typed at runtime + object, user_api_key_dict.aliases + ) try: for fallback_model in fallback_models: if fallback_model == original_model: @@ -2691,6 +2714,13 @@ class ProxyBaseLLMRequestProcessing: self.data = independent_snapshot(pristine) self.data["model"] = fallback_model try: + if endpoint_unsupported_by_model( + route_type, + _fallback_probe_data(pristine, fallback_model, key_aliases, group_alias), + llm_router, + user_api_key_dict.team_id, + ): + continue return await self.common_processing_pre_call_logic( request=request, general_settings=general_settings, diff --git a/litellm/proxy/common_utils/callback_utils.py b/litellm/proxy/common_utils/callback_utils.py index f889748c504..d1787f3ff8b 100644 --- a/litellm/proxy/common_utils/callback_utils.py +++ b/litellm/proxy/common_utils/callback_utils.py @@ -436,10 +436,7 @@ def get_model_group_from_litellm_kwargs(kwargs: dict) -> str | None: def get_model_group_from_request_data(data: dict) -> str | None: _metadata: Final = data.get("metadata", None) or {} _model_group: Final = _metadata.get("model_group", None) - if _model_group is not None: - return _model_group - - return None + return _model_group if isinstance(_model_group, str) else None def get_remaining_tokens_and_requests_from_request_data(data: dict) -> dict[str, str]: diff --git a/litellm/proxy/response_api_endpoints/endpoints.py b/litellm/proxy/response_api_endpoints/endpoints.py index 9e642a02d49..772bb3e53d2 100644 --- a/litellm/proxy/response_api_endpoints/endpoints.py +++ b/litellm/proxy/response_api_endpoints/endpoints.py @@ -34,6 +34,7 @@ from litellm._logging import verbose_proxy_logger from litellm.anthropic_interface.exceptions.exception_mapping_utils import AnthropicExceptionMapping from litellm.caching.dual_cache import DualCache from litellm.constants import EMPTY_MAPPING, MANAGED_BACKGROUND_RESPONSE_METADATA_KEY, TERMINAL_RESPONSE_STATUSES +from litellm.exceptions import EndpointNotSupportedByModelError from litellm.integrations.custom_guardrail import ModifyResponseException from litellm.litellm_core_utils.core_helpers import as_str_mapping, get_metadata_variable_name_from_kwargs from litellm.litellm_core_utils.internal_call_metadata import is_background_response, requests_background @@ -2086,6 +2087,14 @@ def _routing_hints_from_first_ws_frame(first_message: str) -> Mapping[str, objec def _responses_ws_failure_frame(failure: Exception) -> str: + if isinstance(failure, EndpointNotSupportedByModelError): + return json.dumps( + { + "type": "error", + "status": 400, + "error": {"type": "invalid_request_error", "message": failure.message.split("\n\n", 1)[0]}, + } + ) proxy_exception: Final = handle_exception_on_proxy(failure) status: Final = int(proxy_exception.code) if str(proxy_exception.code).isdigit() else 500 error_type: Final = ( @@ -2444,7 +2453,12 @@ async def _responses_websocket_session( with contextlib.suppress(Exception): await websocket.send_text(_responses_ws_failure_frame(e)) with contextlib.suppress(Exception): - await websocket.close(code=1011, reason="Internal server error") + await websocket.close( + code=1008 if isinstance(e, EndpointNotSupportedByModelError) else 1011, + reason="Pre-call error" + if isinstance(e, EndpointNotSupportedByModelError) + else "Internal server error", + ) return crash_row return bill_socket_failure(outcome) if isinstance(outcome, Exception) else None diff --git a/litellm/proxy/route_llm_request.py b/litellm/proxy/route_llm_request.py index db1951ec9e2..fea4f342b91 100644 --- a/litellm/proxy/route_llm_request.py +++ b/litellm/proxy/route_llm_request.py @@ -431,6 +431,28 @@ def raise_if_required_body_param_missing( ) +def endpoint_unsupported_by_model( + route_type: str, + data: Mapping[str, object], + llm_router: LitellmRouter | None, + team_id: str | None, +) -> bool: + model: Final = data.get("model") + if not isinstance(model, str) or llm_router is None: + return False + metadata: Final = data.get("metadata") + request_kwargs: Final = { + **data, + "metadata": { + **(metadata if isinstance(metadata, Mapping) else {}), + "user_api_key_team_id": team_id, + }, + } + return bool(llm_router.get_routable_deployments(model, request_kwargs)) and not llm_router.get_routable_deployments( + model, request_kwargs, call_type=route_type + ) + + class MockTestingParamsDisabledError(HTTPException): def __init__(self, params: tuple[str, ...]): super().__init__( diff --git a/litellm/proxy/spend_tracking/spend_tracking_utils.py b/litellm/proxy/spend_tracking/spend_tracking_utils.py index 833ee71ae9c..df9ed611a47 100644 --- a/litellm/proxy/spend_tracking/spend_tracking_utils.py +++ b/litellm/proxy/spend_tracking/spend_tracking_utils.py @@ -464,7 +464,11 @@ def _sl_attribution_fallback( ) -> str: if standard_logging_payload is None: return "" - return standard_logging_payload.get(field) or "" + return _string_or_empty(standard_logging_payload.get(field)) + + +def _string_or_empty(value: object) -> str: + return value if isinstance(value, str) else "" def _deployment_provider(deployment: DeploymentTypedDict) -> str | None: @@ -682,7 +686,7 @@ def get_logging_payload( _model_id: Final = metadata.get("model_info", {}).get("id", "") or _sl_attribution_fallback( standard_logging_payload, "model_id" ) - _model_group: Final = metadata.get("model_group", "") or _sl_attribution_fallback( + _model_group: Final = _string_or_empty(metadata.get("model_group")) or _sl_attribution_fallback( standard_logging_payload, "model_group" ) _api_base: Final = litellm_params.get("api_base", "") or _sl_attribution_fallback( diff --git a/litellm/router.py b/litellm/router.py index ff31f48c49e..31d679f616e 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -287,6 +287,11 @@ from litellm.router_utils.embedding_batch_size_check import ( embedding_deployment_selection, log_rejection_before_provider_call, ) +from litellm.router_utils.endpoint_eligibility import ( + eligible_deployments, + endpoint_selection, + warn_on_unenforced_endpoint_flag, +) from litellm.router_utils.fallback_event_handlers import ( MID_STREAM_FALLBACK_CONTROLS_KEY, AttemptedFallbackTargets, @@ -1449,7 +1454,7 @@ class Router: self.assistants_config = assistants_config self.search_tools = search_tools or [] self.guardrail_list = guardrail_list or [] - self.deployment_names: list = [] # names of models under litellm_params. ex. azure/chatgpt-v-2 + self.deployment_names: list[str] = [] # names of models under litellm_params. ex. azure/chatgpt-v-2 self.deployment_latency_map = {} ### CACHING ### cache_type: Literal["local", "redis", "redis-semantic", "s3", "disk"] = "local" # default to an in-memory cache @@ -6054,12 +6059,13 @@ class Router: kwargs.setdefault("metadata", {}).update({"model_group": model}) # pick the one that is available (lowest TPM/RPM) - deployment: Final = self.get_available_deployment( - model=model, - messages=messages, - specific_deployment=kwargs.pop("specific_deployment", None), - request_kwargs=kwargs, - ) + with endpoint_selection("text_completion", model=model): + deployment: Final = self.get_available_deployment( + model=model, + messages=messages, + specific_deployment=kwargs.pop("specific_deployment", None), + request_kwargs=kwargs, + ) data: Final = deployment["litellm_params"].copy() for k, v in self.default_litellm_params.items(): @@ -6915,6 +6921,7 @@ class Router: """ model_group: Final = cast(str, initial_kwargs.get("model")) # cast-ok: kwargs always carries the model group retry_kwargs: Final = mid_stream_retry_kwargs(initial_kwargs) + handler: Final = retry_kwargs.get("original_generic_function") healthy_deployments, all_deployments = await self._async_get_healthy_deployments( model=model_group, parent_otel_span=get_parent_otel_span_from_kwargs(retry_kwargs), @@ -6945,7 +6952,10 @@ class Router: ) carry_superseded_attempt_usage(retry_kwargs) try: - response = await attempt_function(**retry_kwargs) + with endpoint_selection( + getattr(handler, "__name__", None), getattr(handler, "__module__", None), model=model_group + ): + response = await attempt_function(**retry_kwargs) except Exception as retry_error: # noqa: BLE001 # the fallback chain judges the failed retry wrapped = _stream_fallback_error_for_raised(retry_error, model_group, False, e.llm_provider) if wrapped is None: @@ -9315,9 +9325,18 @@ class Router: Handler for making a call to the .completion()/.embeddings()/etc. functions. """ model_group: Final = kwargs.get("model") - response = original_function(*args, **kwargs) - if coroutine_checker.is_async_callable(response) or inspect.isawaitable(response): - response = await response + generic_function: Final = _OPAQUE_VALUE_ADAPTER.validate_python(kwargs.get("original_generic_function")) + endpoint_function: Final = ( + generic_function if callable(generic_function) else _OPAQUE_VALUE_ADAPTER.validate_python(original_function) + ) + with endpoint_selection( + getattr(endpoint_function, "__name__", None), + getattr(endpoint_function, "__module__", None), + model=_OPAQUE_VALUE_ADAPTER.validate_python(model_group), + ): + response = original_function(*args, **kwargs) + if coroutine_checker.is_async_callable(response) or inspect.isawaitable(response): + response = await response await self.increment_deployment_usage_for_response(response=response, request_kwargs=kwargs) ## PROCESS RESPONSE HEADERS response = await self.set_response_headers(response=response, model_group=model_group, request_kwargs=kwargs) @@ -10638,6 +10657,7 @@ class Router: declared_params: Final = _LITELLM_PARAMS_ADAPTER.validate_python(_litellm_params) warn_on_provider_credential_mismatch(model_name=_model_name, litellm_params=declared_params) warn_on_ignored_model_capability_flags(model_name=_model_name, litellm_params=declared_params) + warn_on_unenforced_endpoint_flag(model_name=_model_name, model_info=_model_info) deployment = without_model_capability_flags( Deployment( **deployment_info, @@ -11441,6 +11461,9 @@ class Router: declared_params: Final = deployment.litellm_params.model_dump(exclude_none=True) warn_on_provider_credential_mismatch(model_name=deployment.model_name, litellm_params=declared_params) warn_on_ignored_model_capability_flags(model_name=deployment.model_name, litellm_params=declared_params) + warn_on_unenforced_endpoint_flag( + model_name=deployment.model_name, model_info=deployment.model_info.model_dump() + ) served: Final = with_admin_api_key_secret_reference(without_model_capability_flags(deployment)) # add to model list @@ -14797,9 +14820,15 @@ class Router: """ return [m for m in self.model_list if m["litellm_params"]["model"] == model] - def get_routable_deployments(self, model: str) -> tuple[DeploymentTypedDict, ...]: + def get_routable_deployments( + self, + model: str, + request_kwargs: Mapping[str, object] = MappingProxyType({}), + call_type: str | None = None, + ) -> tuple[DeploymentTypedDict, ...]: try: - _, candidates = self._common_checks_available_deployment(model=model, request_kwargs={}) # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] # its signature leaves input and the result untyped + with endpoint_selection(call_type, model=model): + _, candidates = self._common_checks_available_deployment(model=model, request_kwargs=request_kwargs) # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] # its signature leaves input and the result untyped except (litellm.BadRequestError, ValueError): return () if isinstance(candidates, dict): @@ -15075,7 +15104,10 @@ class Router: model=model, llm_provider="", ) - return selectable + eligible: Final = eligible_deployments(model, selectable) + if isinstance(eligible, tuple): + return list(eligible) + raise eligible def _filter_reserved_deployments( self, diff --git a/litellm/router_utils/endpoint_eligibility.py b/litellm/router_utils/endpoint_eligibility.py new file mode 100644 index 00000000000..4a87e582c97 --- /dev/null +++ b/litellm/router_utils/endpoint_eligibility.py @@ -0,0 +1,120 @@ +from collections.abc import Generator, Mapping, Sequence +from contextlib import contextmanager +from contextvars import ContextVar +from types import MappingProxyType +from typing import Final, TypeVar + +from pydantic import BaseModel, ConfigDict, TypeAdapter, ValidationError, field_validator + +from litellm._logging import verbose_router_logger +from litellm.exceptions import EndpointNotSupportedByModelError + +CHAT_ENDPOINT_BY_CALL: Final[Mapping[str, str]] = MappingProxyType( + { + "completion": "/v1/chat/completions", + "acompletion": "/v1/chat/completions", + "_completion": "/v1/chat/completions", + "_acompletion": "/v1/chat/completions", + "text_completion": "/v1/completions", + "atext_completion": "/v1/completions", + "_atext_completion": "/v1/completions", + "anthropic_messages": "/v1/messages", + "aresponses": "/v1/responses", + "agenerate_content": "generateContent", + "agenerate_content_stream": "streamGenerateContent", + "acreate_interaction": "/v1beta/interactions", + "litellm.interactions.main.acreate": "/v1beta/interactions", + "litellm.interactions.main.create": "/v1beta/interactions", + "_aresponses_websocket": "/v1/responses", + "acompact_responses": "/v1/responses/compact", + } +) + + +def _endpoint_path(endpoint: str) -> str: + rooted: Final = "/" + endpoint.strip().strip("/").casefold() + path: Final = rooted.removeprefix("/openai").removeprefix("/v1beta").removeprefix("/v1") + return "/" + path.rsplit(":", 1)[-1] if ":" in path else path + + +CHAT_ENDPOINTS: Final = frozenset(_endpoint_path(endpoint) for endpoint in CHAT_ENDPOINT_BY_CALL.values()) +_Deployment: Final = TypeVar("_Deployment", bound=Mapping[str, object]) +_ENFORCE_FLAG: Final = TypeAdapter(bool | None) +_REQUEST_SELECTION: Final[ContextVar[tuple[str | None, str | None]]] = ContextVar( + "router_request_endpoint", default=(None, None) +) + + +class _EndpointDeclaration(BaseModel): + model_config = ConfigDict(frozen=True) + + enforce_supported_endpoints: bool | None = None + supported_endpoints: tuple[str, ...] | None = None + + @field_validator("supported_endpoints", mode="before") + @classmethod + def _one_endpoint_as_a_list(cls, value: object) -> object: + return (value,) if isinstance(value, str) else value + + +@contextmanager +def endpoint_selection(call_type: object, module: object = None, model: object = None) -> Generator[None]: + qualified_call: Final = f"{module}.{call_type}" if module == "litellm.interactions.main" else call_type + endpoint: Final = CHAT_ENDPOINT_BY_CALL.get(qualified_call) if isinstance(qualified_call, str) else None + token: Final = _REQUEST_SELECTION.set((endpoint, model if isinstance(model, str) else None)) + try: + yield + finally: + _REQUEST_SELECTION.reset(token) + + +def _enforced_endpoints(model_info: object) -> tuple[str, ...] | None: + try: + declaration: Final = _EndpointDeclaration.model_validate(model_info) + except ValidationError: + return None + return declaration.supported_endpoints if declaration.enforce_supported_endpoints else None + + +def warn_on_unenforced_endpoint_flag(model_name: str, model_info: Mapping[str, object]) -> str | None: + try: + flagged: Final = _ENFORCE_FLAG.validate_python(model_info.get("enforce_supported_endpoints")) + except ValidationError: + return None + if not flagged or _enforced_endpoints(model_info) is not None: + return None + warning: Final = ( + f"Deployment '{model_name}' sets enforce_supported_endpoints without a supported_endpoints list of " + 'strings, so litellm enforces nothing for it. Set supported_endpoints, e.g. ["/v1/embeddings"].' + ) + verbose_router_logger.warning(warning) + return warning + + +def non_chat_endpoints(deployment: Mapping[str, object]) -> frozenset[str] | None: + declared: Final = _enforced_endpoints(deployment.get("model_info")) + if declared is None: + return None + endpoints: Final = frozenset(declared) + normalized: Final = frozenset(_endpoint_path(endpoint) for endpoint in endpoints) + return endpoints if normalized.isdisjoint(CHAT_ENDPOINTS) else None + + +def eligible_deployments( + model: str, deployments: Sequence[_Deployment] +) -> tuple[_Deployment, ...] | EndpointNotSupportedByModelError: + endpoint, requested_model = _REQUEST_SELECTION.get() + if endpoint is None or not deployments: + return tuple(deployments) + eligible: Final = tuple(deployment for deployment in deployments if non_chat_endpoints(deployment) is None) + if eligible: + return eligible + requested: Final = requested_model or model + supported: Final = ", ".join( + sorted(frozenset[str]().union(*(non_chat_endpoints(deployment) or frozenset() for deployment in deployments))) + ) + return EndpointNotSupportedByModelError( + message=f"{endpoint} is not supported by model={requested}. Supported endpoints: {supported or 'none'}", + model=model, + llm_provider="", + ) diff --git a/litellm/types/router.py b/litellm/types/router.py index db8be9e91a3..56dcad4051a 100644 --- a/litellm/types/router.py +++ b/litellm/types/router.py @@ -270,6 +270,7 @@ def _whole_float_as_int(value: object) -> object: class ModelInfo(MirroredPricingParams): + enforce_supported_endpoints: bool | None = None id: str | None # Allow id to be optional on input, but it will always be present as a str in the model instance db_model: bool = False # used for proxy - to separate models which are stored in the db vs. config. updated_at: datetime.datetime | None = None diff --git a/tests/e2e/coverage_registry/llm_nonconversational.yaml b/tests/e2e/coverage_registry/llm_nonconversational.yaml index 330ddf88fc5..367135eae26 100644 --- a/tests/e2e/coverage_registry/llm_nonconversational.yaml +++ b/tests/e2e/coverage_registry/llm_nonconversational.yaml @@ -2,6 +2,7 @@ - {id: llm.completions.openai.basic.nonstream.works, module: llm, tier: P0, subject_endpoint: completions, route: openai, capability: basic, streaming: nonstream, assertions: [works], source: "test_completions_endpoint_e2e.py", rationale: "Legacy text /completions endpoint, second-highest production request volume"} - {id: llm.embeddings.openai.basic.nonstream.works, module: llm, tier: P0, subject_endpoint: embeddings, route: openai, capability: basic, streaming: nonstream, assertions: [works], source: "test_embeddings_endpoint_e2e.py:23", rationale: "Core endpoint, live vector response"} - {id: llm.embeddings.openai.input_validation.nonstream.works, module: llm, tier: P1, subject_endpoint: embeddings, route: openai, capability: input_validation, streaming: nonstream, assertions: [works], source: "vendor strategy §9.3 / LIT-4778", rationale: "Missing model/input on /embeddings return client errors"} +- {id: llm.embeddings.openai.supported_endpoints_gate.nonstream.works, module: llm, tier: P1, subject_endpoint: embeddings, route: openai, capability: input_validation, streaming: nonstream, assertions: [works], fail_before_fix: proven, source: "test_embedding_only_supported_endpoints_e2e.py", rationale: "Explicit endpoint restrictions preserve real embeddings, refuse chat with a gateway 400, and route chat through a live fallback; zero-provider-call contracts live in integration/routing/test_non_chat_model_endpoint_gate.py"} - {id: llm.embeddings.openai.basic.nonstream.cost_logged, module: llm, tier: P0, subject_endpoint: embeddings, route: openai, capability: basic, streaming: nonstream, assertions: [cost_logged], source: "SPEND_TRACKING_COVERAGE_MATRIX.md:34", rationale: "Cost tracking on embeddings"} - {id: llm.embeddings.openai.passthrough.nonstream.cost_logged, module: llm, tier: P0, subject_endpoint: embeddings, route: openai, capability: basic, streaming: nonstream, assertions: [cost_logged], source: "test_passthrough_e2e.py", rationale: "POST /openai_passthrough/v1/embeddings is costed; the route wrote no spend row at all, so budgets never saw traffic OpenAI was billing for (GitHub issue #36646)"} - {id: llm.embeddings.azure_openai.basic.nonstream.works, module: llm, tier: P0, subject_endpoint: embeddings, route: azure_openai, capability: basic, streaming: nonstream, assertions: [works], source: "llms/azure/azure.py", rationale: "Azure embeddings via translation"} diff --git a/tests/e2e/llm_translation/test_embedding_only_supported_endpoints_e2e.py b/tests/e2e/llm_translation/test_embedding_only_supported_endpoints_e2e.py new file mode 100644 index 00000000000..c45adfeeb84 --- /dev/null +++ b/tests/e2e/llm_translation/test_embedding_only_supported_endpoints_e2e.py @@ -0,0 +1,119 @@ +"""Live provider checks for embeddings, the chat refusal, and chat fallback from an explicitly restricted model.""" + +from __future__ import annotations + +import dataclasses +from typing import Final + +import pytest +from e2e_config import unique_marker +from e2e_metadata import Domain, Mode, Provider, Route, Subject, meta +from lifecycle import ResourceManager +from models import LiteLLMParamsBody, ModelInfoBody, ModelMode, ModelNewBody +from openai import BadRequestError, OpenAI +from proxy_client import ProxyClient +from sdk_clients import NO_PROXY_CACHE, SdkClients + +pytestmark = pytest.mark.e2e + +EMBEDDING_BACKEND: Final = "openai/text-embedding-3-small" +CHAT_FALLBACK_BACKEND: Final = "openai/gpt-5.4-nano" +EMBEDDINGS_ONLY: Final = ("/v1/embeddings",) +PROMPT: Final = "Say this is a test!" + + +def _subject(route: Route) -> Subject: + return Subject( + domain=Domain.LLM_TRANSLATION, + route=route, + providers=(Provider.OPENAI,), + models=(EMBEDDING_BACKEND,), + mode=Mode.NONSTREAM, + ) + + +def _register( + proxy: ProxyClient, + resources: ResourceManager, + prefix: str, + supported_endpoints: tuple[str, ...] | None, + backend: str = EMBEDDING_BACKEND, + mode: ModelMode = "embedding", +) -> str: + model: Final = f"{prefix}-{unique_marker()}" + model_id: Final = proxy.register_model( + ModelNewBody( + model_name=model, + litellm_params=LiteLLMParamsBody(model=backend, api_key="os.environ/OPENAI_API_KEY"), + model_info=ModelInfoBody( + enforce_supported_endpoints=supported_endpoints is not None, + mode=mode, + supported_endpoints=None if supported_endpoints is None else list(supported_endpoints), + ), + ) + ) + resources.defer(lambda: proxy.delete_model(model_id)) + return model + + +def _openai_at(sdk: SdkClients, api_prefix: str, key: str) -> OpenAI: + return dataclasses.replace(sdk, base_url=f"{sdk.base_url}{api_prefix}").openai(key) + + +class TestEmbeddingOnlySupportedEndpoints: + @pytest.mark.covers("llm.embeddings.openai.supported_endpoints_gate.nonstream.works") + @meta(_subject(Route.EMBEDDINGS)) + @pytest.mark.parametrize( + ("prefix", "supported_endpoints"), + ( + pytest.param("e2e-embed-only", EMBEDDINGS_ONLY, id="embed-only"), + pytest.param("e2e-embed-control", None, id="control"), + ), + ) + def test_embeddings_return_a_vector( + self, + proxy: ProxyClient, + resources: ResourceManager, + sdk: SdkClients, + prefix: str, + supported_endpoints: tuple[str, ...] | None, + ) -> None: + model: Final = _register(proxy, resources, prefix, supported_endpoints) + embeddings: Final = sdk.openai(resources.key()).embeddings.create( + model=model, input=PROMPT, extra_body=NO_PROXY_CACHE + ) + assert embeddings.data, f"/embeddings returned no data for {model}: {embeddings!r}" + assert any(component != 0.0 for component in embeddings.data[0].embedding), ( + f"/embeddings returned an empty or all-zero vector for {model}" + ) + + @pytest.mark.covers("llm.embeddings.openai.supported_endpoints_gate.nonstream.works") + @meta(dataclasses.replace(_subject(Route.CHAT_COMPLETIONS), models=(EMBEDDING_BACKEND, CHAT_FALLBACK_BACKEND))) + def test_chat_completions_with_a_fallback_are_answered_by_the_fallback( + self, proxy: ProxyClient, resources: ResourceManager, sdk: SdkClients + ) -> None: + model: Final = _register(proxy, resources, "e2e-embed-only", EMBEDDINGS_ONLY) + fallback: Final = _register(proxy, resources, "e2e-chat-fallback", None, CHAT_FALLBACK_BACKEND, "chat") + completion: Final = _openai_at(sdk, "/v1", resources.key()).chat.completions.create( + model=model, + messages=[{"role": "user", "content": PROMPT}], + extra_body={**NO_PROXY_CACHE, "fallbacks": [fallback]}, + ) + assert completion.choices and completion.choices[0].message.content, ( + f"the fallback {fallback} should answer a chat call to {model}, got {completion!r}" + ) + + @pytest.mark.covers("llm.embeddings.openai.supported_endpoints_gate.nonstream.works") + @meta(_subject(Route.CHAT_COMPLETIONS)) + def test_chat_completions_without_a_fallback_are_refused_by_the_gateway( + self, proxy: ProxyClient, resources: ResourceManager, sdk: SdkClients + ) -> None: + model: Final = _register(proxy, resources, "e2e-embed-only", EMBEDDINGS_ONLY) + with pytest.raises(BadRequestError) as refused: + _ = _openai_at(sdk, "/v1", resources.key()).chat.completions.create( + model=model, messages=[{"role": "user", "content": PROMPT}], extra_body=NO_PROXY_CACHE + ) + assert ( + f"/v1/chat/completions is not supported by model={model}. Supported endpoints: /v1/embeddings" + in refused.value.message + ), f"chat on {model} should be refused by the gateway before any provider call, got {refused.value!r}" diff --git a/tests/e2e/models.py b/tests/e2e/models.py index afddc7e25d6..504adec5401 100644 --- a/tests/e2e/models.py +++ b/tests/e2e/models.py @@ -1416,7 +1416,7 @@ class LiteLLMParamsBody(BaseModel): order: int | None = None -ModelMode = Literal["chat", "batch", "realtime", "image_generation"] +ModelMode = Literal["chat", "embedding", "batch", "realtime", "image_generation"] class TokenCountModelBody(BaseModel, frozen=True): @@ -1425,6 +1425,7 @@ class TokenCountModelBody(BaseModel, frozen=True): class ModelInfoBody(BaseModel): + enforce_supported_endpoints: bool | None = None # id is left unset so the proxy assigns a unique model_id per deployment. # Pinning it to the model_name made re-registrations of a fixed-name model # (e.g. the batch suite's openai-batch) collide on the model_id unique diff --git a/tests/integration/routing/test_non_chat_model_endpoint_gate.py b/tests/integration/routing/test_non_chat_model_endpoint_gate.py new file mode 100644 index 00000000000..465f5801692 --- /dev/null +++ b/tests/integration/routing/test_non_chat_model_endpoint_gate.py @@ -0,0 +1,1515 @@ +import asyncio +import contextlib +import json +import uuid +from collections.abc import Iterator, Mapping +from dataclasses import dataclass +from pathlib import Path +from types import MappingProxyType +from typing import Final + +import httpx +import openai +import pytest +import websockets +import yaml +from integration._support.client import ( + JSON_OBJECT, + Gateway, + Scenario, + eventually, + gateway_from_environment, + held, + list_value, + object_value, + string_value, +) +from integration._support import responses_ws_peer +from integration._support.database import read_rows +from integration._support.openai_wire import answering_model_discovery, chat_reply, responses_reply +from integration._support.process import OwnedState, graceful_stop_seconds, owned_proxy_process +from integration._support.tls import write_self_signed_cert +from integration._support.wire import Reply, Request, Wire, wire_server +from pydantic import JsonValue, TypeAdapter +from websockets.exceptions import ConnectionClosed + +_PROVIDER_KEY: Final = "integration-provider-key" +_EMBEDDING_MODEL: Final = "text-embedding-3-small" +_CHAT_MODEL: Final = "gpt-4o-mini" +_CHAT_BACKEND_PREFIX: Final = "gate-chat-backend-" +_SERVED_TEXT: Final = "served by the chat upstream" +_VECTOR: Final = (0.25, 0.5, 0.75) +_COUNTED_TOKENS: Final = 17 +_CLOSE: Final = MappingProxyType({"connection": "close"}) +_RELOAD_SECONDS: Final = 5 +_EMBED_ONLY: Final[Mapping[str, JsonValue]] = MappingProxyType( + {"mode": "embedding", "supported_endpoints": ["/v1/embeddings"], "enforce_supported_endpoints": True} +) + + +def _embedding_reply() -> Reply: + return Reply( + body=json.dumps( + { + "object": "list", + "data": [{"object": "embedding", "index": 0, "embedding": list(_VECTOR)}], + "model": _EMBEDDING_MODEL, + "usage": {"prompt_tokens": 3, "total_tokens": 3}, + } + ).encode() + ) + + +def _not_a_chat_model() -> Reply: + return Reply( + status=404, + body=json.dumps( + {"error": {"message": "scripted: not a chat model", "type": "invalid_request_error", "code": None}} + ).encode(), + ) + + +def _text_completion_reply() -> Reply: + return Reply( + body=json.dumps( + { + "id": f"cmpl-{uuid.uuid4().hex}", + "object": "text_completion", + "created": 1, + "model": _CHAT_MODEL, + "choices": [{"index": 0, "text": _SERVED_TEXT, "finish_reason": "stop", "logprobs": None}], + "usage": {"prompt_tokens": 3, "completion_tokens": 3, "total_tokens": 6}, + } + ).encode() + ) + + +def _serves_chat(model: JsonValue) -> bool: + return model == _CHAT_MODEL or (isinstance(model, str) and model.startswith(_CHAT_BACKEND_PREFIX)) + + +def _respond(request: Request) -> Reply: + path: Final = request.target.split("?")[0] + body: Final = JSON_OBJECT.validate_json(request.body or b"{}") + streamed: Final = body.get("stream") is True + if path == "/v1/embeddings": + return _embedding_reply() + if path == "/v1/responses/input_tokens": + return Reply(body=json.dumps({"object": "response.input_tokens", "input_tokens": _COUNTED_TOKENS}).encode()) + if not _serves_chat(body.get("model")): + return _not_a_chat_model() + if path == "/v1/chat/completions": + return chat_reply(f"chatcmpl-{uuid.uuid4().hex}", _CHAT_MODEL, _SERVED_TEXT, stream=streamed) + if path in ("/v1/responses", "/v1/responses/compact"): + return responses_reply(f"resp_{uuid.uuid4().hex}", _CHAT_MODEL, _SERVED_TEXT, stream=streamed) + if path == "/v1/completions": + return _text_completion_reply() + return _not_a_chat_model() + + +class _Upstream: + def __init__(self, wire: Wire) -> None: + self.wire = wire + self.seen: tuple[Request, ...] = () + + def calls(self, marker: str) -> tuple[Request, ...]: + self.seen = (*self.seen, *self.wire.drain()) + return tuple(request for request in self.seen if marker.encode() in request.body or marker in request.target) + + def targets(self, marker: str) -> tuple[str, ...]: + return tuple(request.target for request in self.calls(marker)) + + +@dataclass(frozen=True, slots=True) +class _Deployment: + litellm_model: str + declarations: tuple[Mapping[str, JsonValue], ...] + + +def _single(declaration: Mapping[str, JsonValue], litellm_model: str = f"openai/{_CHAT_MODEL}") -> _Deployment: + return _Deployment(litellm_model, (declaration,)) + + +_SHARED: Final[Mapping[str, _Deployment]] = MappingProxyType( + { + "embedding_only": _single(_EMBED_ONLY, f"openai/{_EMBEDDING_MODEL}"), + "embedding_only_second": _single(_EMBED_ONLY, f"openai/{_EMBEDDING_MODEL}"), + "image_generation": _single( + {"enforce_supported_endpoints": True, "supported_endpoints": ["/v1/images/generations"]} + ), + "chat_target": _single({"mode": "chat"}), + "chat_rate_limited": _single({"mode": "chat"}), + "mode_not_a_string": _single( + {"mode": 7, "supported_endpoints": ["/v1/embeddings"]}, "openai/gate-mode-not-a-string" + ), + } +) + + +@dataclass(frozen=True, slots=True) +class _Rig: + gateway: Gateway + upstream: _Upstream + names: Mapping[str, str] + api_base: str + + +def _create(gateway: Gateway, name: str, litellm_params: Mapping[str, JsonValue], info: Mapping[str, JsonValue]) -> str: + created: Final = gateway.post( + "/model/new", + { + "model_name": name, + "litellm_params": {"api_key": _PROVIDER_KEY, **litellm_params}, + "model_info": dict(info), + }, + ) + return string_value(object_value(created["model_info"])["id"]) + + +def _deploy(scenario: Scenario, api_base: str, label: str, deployment: _Deployment) -> tuple[str, tuple[str, ...]]: + name: Final = f"gate-{label}-{uuid.uuid4().hex[:12]}" + identities: Final = tuple( + _create(scenario.gateway, name, {"model": deployment.litellm_model, "api_base": api_base}, declaration) + for declaration in deployment.declarations + ) + tuple(scenario.cleanups.callback(scenario.delete_model, identity) for identity in identities) + return name, identities + + +def _listed_ids(gateway: Gateway) -> frozenset[str]: + listed: Final = gateway.request("GET", "/model/info", headers=_CLOSE) + assert listed.status_code == 200, listed.text + return frozenset( + string_value(object_value(object_value(entry)["model_info"])["id"]) + for entry in list_value(JSON_OBJECT.validate_json(listed.content)["data"]) + ) + + +def _served_by_every_worker(gateway: Gateway, identities: frozenset[str]) -> None: + held( + lambda: identities <= _listed_ids(gateway), + lambda served: served, + holding=_RELOAD_SECONDS * 2, + seconds=_RELOAD_SECONDS * 12, + ) + + +@pytest.fixture(scope="module") +def rig() -> Iterator[_Rig]: + with ( + wire_server(answering_model_discovery(_respond)) as wire, + gateway_from_environment() as gateway, + gateway.scenario() as scenario, + ): + api_base: Final = wire.url + "/v1" + deployed: Final = {label: _deploy(scenario, api_base, label, spec) for label, spec in _SHARED.items()} + _served_by_every_worker( + gateway, frozenset(identity for _name, identities in deployed.values() for identity in identities) + ) + yield _Rig( + gateway, + _Upstream(wire), + MappingProxyType({label: name for label, (name, _ids) in deployed.items()}), + api_base, + ) + + +def _marker() -> str: + return f"gate-{uuid.uuid4().hex}" + + +def _openai_error(endpoint: str, model: str, supported: str) -> dict[str, JsonValue]: + return { + "error": { + "message": f"{endpoint} is not supported by model={model}. Supported endpoints: {supported}", + "type": "invalid_request_error", + "param": "model", + "code": "400", + } + } + + +def _anthropic_error(model: str, supported: str) -> dict[str, JsonValue]: + return { + "type": "error", + "error": { + "type": "invalid_request_error", + "message": f"/v1/messages is not supported by model={model}. Supported endpoints: {supported}", + }, + } + + +def _chat_body(model: str, marker: str, *, stream: bool = False) -> dict[str, JsonValue]: + return {"model": model, "messages": [{"role": "user", "content": marker}], **({"stream": True} if stream else {})} + + +def _failure_row(call_id: str) -> dict[str, JsonValue]: + rows: Final = eventually( + lambda: read_rows('SELECT request_id, status FROM "LiteLLM_SpendLogs" WHERE request_id = %s', (call_id,)), + lambda found: len(found) >= 1, + seconds=70, + ) + assert len(rows) == 1, rows + return rows[0] + + +def _assert_error(response: httpx.Response, expected: dict[str, JsonValue]) -> None: + assert response.status_code == 400, response.text + body: Final = JSON_OBJECT.validate_json(response.content) + error: Final = object_value(body["error"]) + message: Final = string_value(error["message"]).split("\n\n", 1)[0] + assert {**body, "error": {**error, "message": message}} == expected, response.text + + +def _assert_rejected_before( + upstream: _Upstream, response: httpx.Response, marker: str, expected: dict[str, JsonValue] +) -> None: + _assert_error(response, expected) + assert upstream.calls(marker) == (), upstream.targets(marker) + assert _failure_row(response.headers["x-litellm-call-id"])["status"] == "failure" + + +def _assert_rejected(rig: _Rig, response: httpx.Response, marker: str, expected: dict[str, JsonValue]) -> None: + _assert_rejected_before(rig.upstream, response, marker, expected) + + +@dataclass(frozen=True, slots=True) +class _Route: + path: str + endpoint: str + anthropic_shape: bool + stream: bool + model_in_path: bool + shape: str + + +_ROUTES: Final[Mapping[str, _Route]] = MappingProxyType( + { + "interactions": _Route("/v1beta/interactions", "/v1beta/interactions", False, False, False, "interactions"), + "chat_v1": _Route("/v1/chat/completions", "/v1/chat/completions", False, False, False, "chat"), + "chat_v1_stream": _Route("/v1/chat/completions", "/v1/chat/completions", False, True, False, "chat"), + "chat_bare": _Route("/chat/completions", "/v1/chat/completions", False, False, False, "chat"), + "chat_engines": _Route("/engines/{model}/chat/completions", "/v1/chat/completions", False, False, True, "chat"), + "chat_azure_deployments": _Route( + "/openai/deployments/{model}/chat/completions", "/v1/chat/completions", False, False, True, "chat" + ), + "messages": _Route("/v1/messages", "/v1/messages", True, False, False, "messages"), + "messages_stream": _Route("/v1/messages", "/v1/messages", True, True, False, "messages"), + "responses_v1": _Route("/v1/responses", "/v1/responses", False, False, False, "responses"), + "responses_v1_stream": _Route("/v1/responses", "/v1/responses", False, True, False, "responses"), + "responses_bare": _Route("/responses", "/v1/responses", False, False, False, "responses"), + "cursor_responses_body": _Route("/cursor/chat/completions", "/v1/responses", False, False, False, "cursor"), + "cursor_chat_body": _Route("/cursor/chat/completions", "/v1/chat/completions", False, False, False, "chat"), + } +) + + +def _route_body(route: _Route, model: str, marker: str) -> dict[str, JsonValue]: + shaped: Final[Mapping[str, dict[str, JsonValue]]] = { + "chat": {"messages": [{"role": "user", "content": marker}]}, + "messages": {"max_tokens": 16, "messages": [{"role": "user", "content": marker}]}, + "responses": {"input": marker}, + "interactions": {"input": marker}, + "cursor": {"input": [{"role": "user", "content": marker}]}, + } + return { + **({} if route.model_in_path else {"model": model}), + **shaped[route.shape], + **({"stream": True} if route.stream else {}), + } + + +def _send(gateway: Gateway, route: _Route, model: str, marker: str) -> httpx.Response: + path: Final = route.path.format(model=model) + body: Final = _route_body(route, model, marker) + if not route.stream: + return gateway.request("POST", path, body) + with gateway.client.stream("POST", path, json=body, headers={"Authorization": f"Bearer {gateway.key}"}) as streamed: + streamed.read() + return streamed + + +@pytest.mark.parametrize("route", tuple(_ROUTES)) +def test_chat_shaped_route_rejects_an_embedding_only_model_before_the_upstream(rig: _Rig, route: str) -> None: + spec: Final = _ROUTES[route] + model: Final = rig.names["embedding_only"] + marker: Final = _marker() + response: Final = _send(rig.gateway, spec, model, marker) + expected: Final = ( + _anthropic_error(model, "/v1/embeddings") + if spec.anthropic_shape + else _openai_error(spec.endpoint, model, "/v1/embeddings") + ) + _assert_rejected(rig, response, marker, expected) + + +def _sdk_embedding(rig: _Rig, model: str, marker: str) -> tuple[tuple[float, ...], str]: + async def call() -> tuple[tuple[float, ...], str]: + async with openai.AsyncOpenAI( + base_url=f"{rig.gateway.client.base_url}/v1", api_key=rig.gateway.key, max_retries=0 + ) as sdk: + raw: Final = await sdk.embeddings.with_raw_response.create( + model=model, input=marker, encoding_format="float" + ) + parsed: Final = raw.parse() + return tuple(parsed.data[0].embedding), raw.headers["x-litellm-call-id"] + + return asyncio.run(call()) + + +def _httpx_embedding(rig: _Rig, model: str, marker: str) -> tuple[tuple[float, ...], str]: + response: Final = rig.gateway.request( + "POST", "/v1/embeddings", {"model": model, "input": marker, "encoding_format": "float"} + ) + assert response.status_code == 200, response.text + data: Final = list_value(JSON_OBJECT.validate_json(response.content)["data"]) + vector: Final = list_value(object_value(data[0])["embedding"]) + return tuple(float(str(value)) for value in vector), response.headers["x-litellm-call-id"] + + +@pytest.mark.parametrize("client", ("httpx_sync", "openai_async_sdk")) +def test_embeddings_still_reach_the_embedding_only_model(rig: _Rig, client: str) -> None: + model: Final = rig.names["embedding_only"] + marker: Final = _marker() + vector, call_id = ( + _httpx_embedding(rig, model, marker) if client == "httpx_sync" else _sdk_embedding(rig, model, marker) + ) + assert vector == _VECTOR + assert rig.upstream.targets(marker) == ("/v1/embeddings",) + row: Final = eventually( + lambda: read_rows( + 'SELECT status, call_type FROM "LiteLLM_SpendLogs" WHERE request_id = %s OR litellm_call_id = %s', + (call_id, call_id), + ), + lambda found: len(found) >= 1, + seconds=70, + ) + assert [(entry["status"], entry["call_type"]) for entry in row] == [("success", "aembedding")], row + + +def _assert_served(rig: _Rig, response: httpx.Response, marker: str, target: str) -> None: + assert response.status_code == 200, response.text + assert _SERVED_TEXT in response.text, response.text + assert rig.upstream.targets(marker) == (target,) + + +def test_key_alias_to_an_embedding_only_model_is_rejected_under_the_resolved_name(rig: _Rig) -> None: + target: Final = rig.names["embedding_only"] + alias: Final = f"gate-alias-{uuid.uuid4().hex[:12]}" + marker: Final = _marker() + with rig.gateway.scenario() as scenario: + key: Final = scenario.key(aliases={alias: target}) + response: Final = rig.gateway.request("POST", "/v1/chat/completions", _chat_body(alias, marker), key=key) + _assert_rejected(rig, response, marker, _openai_error("/v1/chat/completions", target, "/v1/embeddings")) + + +async def _websocket_first_frame_and_close( + gateway: Gateway, key: str, model: str, marker: str +) -> tuple[dict[str, JsonValue], int | None, str | None]: + url: Final = str(gateway.client.base_url).rstrip("/").replace("http://", "ws://", 1) + "/v1/responses" + async with websockets.connect(url, additional_headers={"Authorization": f"Bearer {key}"}, open_timeout=10) as ws: + await ws.send(json.dumps({"type": "response.create", "model": model, "input": marker})) + first: Final = JSON_OBJECT.validate_json(await asyncio.wait_for(ws.recv(), 30)) + with contextlib.suppress(ConnectionClosed): + await asyncio.wait_for(ws.recv(), 30) + return first, ws.close_code, ws.close_reason + + +def _websocket_rejection(model: str) -> tuple[dict[str, JsonValue], int | None, str | None]: + return ( + { + "type": "error", + "status": 400, + "error": { + "type": "invalid_request_error", + "message": f"/v1/responses is not supported by model={model}. Supported endpoints: /v1/embeddings", + }, + }, + 1008, + "Pre-call error", + ) + + +@pytest.mark.parametrize("aliased", [False, True]) +def test_responses_websocket_closes_on_an_embedding_only_model_before_the_upstream(rig: _Rig, aliased: bool) -> None: + target: Final = rig.names["embedding_only"] + alias: Final = f"gate-ws-alias-{uuid.uuid4().hex[:12]}" + marker: Final = _marker() + with rig.gateway.scenario() as scenario: + key: Final = scenario.key(aliases={alias: target}) + closed: Final = asyncio.run( + _websocket_first_frame_and_close(rig.gateway, key, alias if aliased else target, marker) + ) + assert closed == _websocket_rejection(target) + assert rig.upstream.calls(marker) == () + + +_RESERVING_KEYS: Final[Mapping[str, Mapping[str, JsonValue]]] = MappingProxyType( + { + "max_parallel_requests": {"max_parallel_requests": 1}, + "tpm_limit": {"tpm_limit": 100, "default_estimated_output_tokens": 80}, + } +) + + +@pytest.mark.parametrize("limit", tuple(_RESERVING_KEYS)) +def test_responses_websocket_rejections_leave_the_key_reservation_free(rig: _Rig, limit: str) -> None: + target: Final = rig.names["embedding_only"] + markers: Final = tuple(_marker() for _ in range(3)) + served_marker: Final = _marker() + with rig.gateway.scenario() as scenario: + key: Final = scenario.key(**_RESERVING_KEYS[limit]) + closed: Final = tuple( + asyncio.run(_websocket_first_frame_and_close(rig.gateway, key, target, marker)) for marker in markers + ) + assert closed == (_websocket_rejection(target),) * 3 + served: Final = rig.gateway.request( + "POST", "/v1/chat/completions", _chat_body(rig.names["chat_target"], served_marker), key=key + ) + _assert_served(rig, served, served_marker, "/v1/chat/completions") + assert [rig.upstream.calls(marker) for marker in markers] == [()] * 3 + + +def test_team_scoped_embedding_only_model_is_rejected_under_its_public_name(rig: _Rig) -> None: + public: Final = f"gate-team-{uuid.uuid4().hex[:12]}" + marker: Final = _marker() + with rig.gateway.scenario() as scenario: + team: Final = scenario.team() + identity: Final = _create( + rig.gateway, + public, + {"model": f"openai/{_EMBEDDING_MODEL}", "api_base": rig.api_base}, + {**_EMBED_ONLY, "team_id": team}, + ) + scenario.cleanups.callback(scenario.delete_model, identity) + key: Final = scenario.key(team_id=team) + _served_by_every_worker(rig.gateway, frozenset({identity})) + embedded: Final = rig.gateway.request( + "POST", "/v1/embeddings", {"model": public, "input": marker, "encoding_format": "float"}, key=key + ) + assert embedded.status_code == 200, embedded.text + assert rig.upstream.targets(marker) == ("/v1/embeddings",) + rejected_marker: Final = _marker() + response: Final = rig.gateway.request( + "POST", "/v1/chat/completions", _chat_body(public, rejected_marker), key=key + ) + _assert_rejected( + rig, response, rejected_marker, _openai_error("/v1/chat/completions", public, "/v1/embeddings") + ) + + +def test_wildcard_embedding_only_deployment_rejects_chat_for_any_matched_model(rig: _Rig) -> None: + prefix: Final = f"gate-wild-{uuid.uuid4().hex[:12]}" + requested: Final = f"{prefix}/{_EMBEDDING_MODEL}" + marker: Final = _marker() + with rig.gateway.scenario() as scenario: + identity: Final = _create( + rig.gateway, f"{prefix}/*", {"model": "openai/*", "api_base": rig.api_base}, _EMBED_ONLY + ) + scenario.cleanups.callback(scenario.delete_model, identity) + _served_by_every_worker(rig.gateway, frozenset({identity})) + response: Final = rig.gateway.request("POST", "/v1/chat/completions", _chat_body(requested, marker)) + _assert_rejected(rig, response, marker, _openai_error("/v1/chat/completions", requested, "/v1/embeddings")) + + +def test_key_without_access_gets_the_access_error_and_learns_nothing_about_endpoints(rig: _Rig) -> None: + marker: Final = _marker() + with rig.gateway.scenario() as scenario: + key: Final = scenario.key(models=[rig.names["chat_target"]]) + response: Final = rig.gateway.request( + "POST", "/v1/chat/completions", _chat_body(rig.names["embedding_only"], marker), key=key + ) + assert response.status_code == 403, response.text + assert object_value(JSON_OBJECT.validate_json(response.content)["error"])["type"] == "key_model_access_denied", ( + response.text + ) + assert "Supported endpoints" not in response.text, response.text + assert rig.upstream.calls(marker) == () + + +def test_unknown_model_keeps_the_invalid_model_error(rig: _Rig) -> None: + marker: Final = _marker() + response: Final = rig.gateway.request( + "POST", "/v1/chat/completions", _chat_body(f"gate-missing-{uuid.uuid4().hex}", marker) + ) + assert response.status_code == 400, response.text + assert "Invalid model name" in response.text and "Supported endpoints" not in response.text, response.text + assert rig.upstream.calls(marker) == () + + +def test_request_fallbacks_answer_for_an_embedding_only_primary(rig: _Rig) -> None: + model: Final = rig.names["embedding_only"] + marker: Final = _marker() + response: Final = rig.gateway.request( + "POST", "/v1/chat/completions", {**_chat_body(model, marker), "fallbacks": [rig.names["chat_target"]]} + ) + assert response.status_code == 200, response.text + assert _SERVED_TEXT in response.text, response.text + assert rig.upstream.targets(marker) == ("/v1/chat/completions",) + + +def test_an_empty_request_fallback_list_keeps_the_rejection(rig: _Rig) -> None: + model: Final = rig.names["embedding_only"] + marker: Final = _marker() + response: Final = rig.gateway.request( + "POST", "/v1/chat/completions", {**_chat_body(model, marker), "fallbacks": []} + ) + _assert_rejected(rig, response, marker, _openai_error("/v1/chat/completions", model, "/v1/embeddings")) + + +def test_fallbacks_that_cannot_answer_chat_keep_the_rejection(rig: _Rig) -> None: + model: Final = rig.names["embedding_only"] + marker: Final = _marker() + fallbacks: Final = [rig.names["image_generation"], model] + response: Final = rig.gateway.request( + "POST", "/v1/chat/completions", {**_chat_body(model, marker), "fallbacks": fallbacks} + ) + _assert_rejected(rig, response, marker, _openai_error("/v1/chat/completions", model, "/v1/embeddings")) + + +def test_request_fallbacks_do_not_rescue_an_unknown_primary(rig: _Rig) -> None: + marker: Final = _marker() + response: Final = rig.gateway.request( + "POST", + "/v1/chat/completions", + {**_chat_body(f"gate-missing-{uuid.uuid4().hex}", marker), "fallbacks": [rig.names["chat_target"]]}, + ) + assert response.status_code == 400, response.text + assert "Invalid model name" in response.text, response.text + assert rig.upstream.calls(marker) == () + + +def _chat(rig: _Rig, model: str, marker: str, *, key: str | None = None, **extra: JsonValue) -> httpx.Response: + return rig.gateway.request("POST", "/v1/chat/completions", {**_chat_body(model, marker), **extra}, key=key) + + +def _assert_rejected_chat(rig: _Rig, response: httpx.Response, marker: str, model: str) -> None: + _assert_rejected(rig, response, marker, _openai_error("/v1/chat/completions", model, "/v1/embeddings")) + + +def test_a_fallback_chain_that_reaches_chat_two_hops_away_answers(rig: _Rig) -> None: + first, second, chat = rig.names["embedding_only"], rig.names["embedding_only_second"], rig.names["chat_target"] + marker: Final = _marker() + response: Final = _chat(rig, first, marker, fallbacks=[{first: [second]}, {second: [chat]}]) + assert response.status_code == 200, response.text + assert _SERVED_TEXT in response.text, response.text + assert rig.upstream.targets(marker) == ("/v1/chat/completions",) + + +def test_a_fallback_chain_that_loops_between_embedding_only_groups_keeps_the_rejection(rig: _Rig) -> None: + first, second = rig.names["embedding_only"], rig.names["embedding_only_second"] + marker: Final = _marker() + response: Final = _chat(rig, first, marker, fallbacks=[{first: [second]}, {second: [first]}]) + _assert_rejected_chat(rig, response, marker, first) + + +def test_a_specific_chain_that_loops_beats_a_catch_all_that_could_answer(rig: _Rig) -> None: + first, second = rig.names["embedding_only"], rig.names["embedding_only_second"] + marker: Final = _marker() + response: Final = _chat( + rig, first, marker, fallbacks=[{"*": [rig.names["chat_target"]]}, {first: [second]}, {second: [first]}] + ) + _assert_rejected_chat(rig, response, marker, first) + + +def test_disable_fallbacks_keeps_the_rejection_over_request_fallbacks(rig: _Rig) -> None: + model: Final = rig.names["embedding_only"] + marker: Final = _marker() + response: Final = _chat(rig, model, marker, fallbacks=[rig.names["chat_target"]], disable_fallbacks=True) + _assert_rejected_chat(rig, response, marker, model) + + +def _assert_fallback_answered( + upstream: _Upstream, response: httpx.Response, marker: str, target: str = "/v1/chat/completions" +) -> None: + assert response.status_code == 200, response.text + assert _SERVED_TEXT in response.text, response.text + assert upstream.targets(marker) == (target,), response.text + + +def _through_a_group_with_no_deployments(rig: _Rig) -> Mapping[str, JsonValue]: + ghost: Final = f"gate-ghost-{uuid.uuid4().hex}" + return {"fallbacks": [{rig.names["embedding_only"]: [ghost]}, {ghost: [rig.names["chat_target"]]}]} + + +def test_request_fallbacks_naming_a_chat_backend_name_answer_for_an_embedding_only_primary(rig: _Rig) -> None: + model: Final = rig.names["embedding_only"] + backend: Final = f"openai/{_CHAT_BACKEND_PREFIX}{uuid.uuid4().hex[:12]}" + marker: Final = _marker() + with rig.gateway.scenario() as scenario: + identity: Final = _create( + rig.gateway, f"gate-backend-{uuid.uuid4().hex[:12]}", {"model": backend, "api_base": rig.api_base}, {} + ) + scenario.cleanups.callback(scenario.delete_model, identity) + _served_by_every_worker(rig.gateway, frozenset({identity})) + response: Final = _chat(rig, model, marker, fallbacks=[{model: [backend]}]) + _assert_fallback_answered(rig.upstream, response, marker) + + +@pytest.mark.parametrize("level", ("key", "team")) +def test_router_settings_fallbacks_through_a_group_with_no_deployments_answer(rig: _Rig, level: str) -> None: + marker: Final = _marker() + with rig.gateway.scenario() as scenario: + settings: Final = {"router_settings": dict(_through_a_group_with_no_deployments(rig))} + key: Final = scenario.key(**settings) if level == "key" else scenario.key(team_id=scenario.team(**settings)) + response: Final = _chat(rig, rig.names["embedding_only"], marker, key=key) + _assert_fallback_answered(rig.upstream, response, marker) + + +def test_team_router_settings_fallbacks_answer_for_an_embedding_only_primary(rig: _Rig) -> None: + model, chat = rig.names["embedding_only"], rig.names["chat_target"] + marker: Final = _marker() + with rig.gateway.scenario() as scenario: + team: Final = scenario.team(router_settings={"fallbacks": [{model: [chat]}]}) + response: Final = _chat(rig, model, marker, key=scenario.key(team_id=team)) + assert response.status_code == 200, response.text + assert _SERVED_TEXT in response.text, response.text + assert rig.upstream.targets(marker) == ("/v1/chat/completions",) + + +def test_key_router_settings_fallbacks_answer_for_an_embedding_only_primary(rig: _Rig) -> None: + model, chat = rig.names["embedding_only"], rig.names["chat_target"] + marker: Final = _marker() + with rig.gateway.scenario() as scenario: + key: Final = scenario.key(router_settings={"fallbacks": [{model: [chat]}]}) + response: Final = _chat(rig, model, marker, key=key) + assert response.status_code == 200, response.text + assert _SERVED_TEXT in response.text, response.text + assert rig.upstream.targets(marker) == ("/v1/chat/completions",) + + +def test_key_router_settings_that_cannot_answer_override_team_fallbacks_that_could(rig: _Rig) -> None: + model: Final = rig.names["embedding_only"] + marker: Final = _marker() + with rig.gateway.scenario() as scenario: + team: Final = scenario.team(router_settings={"fallbacks": [{model: [rig.names["chat_target"]]}]}) + key: Final = scenario.key( + team_id=team, router_settings={"fallbacks": [{model: [rig.names["embedding_only_second"]]}]} + ) + response: Final = _chat(rig, model, marker, key=key) + _assert_rejected_chat(rig, response, marker, model) + + +def test_an_empty_request_fallback_list_overrides_key_router_settings_fallbacks(rig: _Rig) -> None: + model: Final = rig.names["embedding_only"] + marker: Final = _marker() + with rig.gateway.scenario() as scenario: + key: Final = scenario.key(router_settings={"fallbacks": [{model: [rig.names["chat_target"]]}]}) + response: Final = _chat(rig, model, marker, key=key, fallbacks=[]) + _assert_rejected_chat(rig, response, marker, model) + + +def test_local_rate_limit_fallback_skips_an_embedding_only_hop(rig: _Rig) -> None: + limited, chat, model = rig.names["chat_rate_limited"], rig.names["chat_target"], rig.names["embedding_only"] + first, second = _marker(), _marker() + with rig.gateway.scenario() as scenario: + key: Final = scenario.key( + model_rpm_limit={limited: 1}, router_settings={"fallbacks": [{limited: [model, chat]}]} + ) + served_first: Final = _chat(rig, limited, first, key=key) + served_second: Final = _chat(rig, limited, second, key=key) + assert (served_first.status_code, served_first.headers.get("x-litellm-model-group")) == (200, limited), ( + served_first.text + ) + assert (served_second.status_code, served_second.headers.get("x-litellm-model-group")) == (200, chat), ( + served_second.text + ) + assert (rig.upstream.targets(first), rig.upstream.targets(second)) == ( + ("/v1/chat/completions",), + ("/v1/chat/completions",), + ) + + +_COMPACT_PATHS: Final = ("/v1/responses/compact", "/responses/compact", "/openai/v1/responses/compact") + + +def _compact_body(model: str, marker: str) -> dict[str, JsonValue]: + return {"model": model, "input": [{"role": "user", "content": marker}]} + + +@pytest.mark.parametrize("path", _COMPACT_PATHS) +def test_responses_compact_rejects_an_embedding_only_model_before_the_upstream(rig: _Rig, path: str) -> None: + model: Final = rig.names["embedding_only"] + marker: Final = _marker() + response: Final = rig.gateway.request("POST", path, _compact_body(model, marker)) + _assert_rejected(rig, response, marker, _openai_error("/v1/responses/compact", model, "/v1/embeddings")) + + +def test_responses_compact_keeps_reaching_the_upstream_for_a_chat_model(rig: _Rig) -> None: + marker: Final = _marker() + response: Final = rig.gateway.request( + "POST", "/v1/responses/compact", _compact_body(rig.names["chat_target"], marker) + ) + _assert_served(rig, response, marker, "/v1/responses/compact") + + +def test_count_tokens_keeps_counting_for_an_embedding_only_model(rig: _Rig) -> None: + marker: Final = _marker() + response: Final = rig.gateway.request( + "POST", + "/v1/messages/count_tokens", + {"model": rig.names["embedding_only"], "messages": [{"role": "user", "content": marker}]}, + ) + assert response.status_code == 200, response.text + assert JSON_OBJECT.validate_json(response.content)["input_tokens"] == _COUNTED_TOKENS, response.text + assert rig.upstream.targets(marker) == ("/v1/responses/input_tokens",) + + +def test_text_completions_reject_an_embedding_only_model_before_the_upstream(rig: _Rig) -> None: + model: Final = rig.names["embedding_only"] + marker: Final = _marker() + response: Final = rig.gateway.request("POST", "/v1/completions", {"model": model, "prompt": marker}) + _assert_rejected(rig, response, marker, _openai_error("/v1/completions", model, "/v1/embeddings")) + + +@pytest.mark.parametrize("method", ("generateContent", "streamGenerateContent")) +def test_gemini_routes_reject_an_embedding_only_model_before_the_upstream(rig: _Rig, method: str) -> None: + model: Final = rig.names["embedding_only"] + marker: Final = _marker() + response: Final = rig.gateway.request( + "POST", f"/v1beta/models/{model}:{method}", {"contents": [{"role": "user", "parts": [{"text": marker}]}]} + ) + assert response.status_code == 400, response.text + assert f"{method} is not supported by model={model}. Supported endpoints: /v1/embeddings" in response.text + assert rig.upstream.calls(marker) == (), rig.upstream.targets(marker) + + +def test_repeated_rejections_log_one_failure_row_per_call(rig: _Rig) -> None: + model: Final = rig.names["embedding_only"] + marker: Final = _marker() + responses: Final = tuple( + rig.gateway.request("POST", "/v1/chat/completions", _chat_body(model, marker)) for _ in range(3) + ) + expected: Final = _openai_error("/v1/chat/completions", model, "/v1/embeddings") + for response in responses: + _assert_error(response, expected) + call_ids: Final = tuple(item.headers["x-litellm-call-id"] for item in responses) + assert len(set(call_ids)) == 3, call_ids + assert [_failure_row(call_id)["status"] for call_id in call_ids] == ["failure"] * 3 + assert rig.upstream.calls(marker) == () + + +def _closed_chat(rig: _Rig, model: str, marker: str) -> httpx.Response: + return rig.gateway.request("POST", "/v1/chat/completions", _chat_body(model, marker), headers=_CLOSE) + + +def test_live_declaration_update_turns_the_gate_on_and_off_on_every_worker(rig: _Rig) -> None: + chat_capable: Final[dict[str, JsonValue]] = { + "mode": "chat", + "supported_endpoints": ["/v1/chat/completions", "/v1/embeddings"], + } + marker: Final = _marker() + with rig.gateway.scenario() as scenario: + name, identities = _deploy(scenario, rig.api_base, "flip", _single(chat_capable)) + _served_by_every_worker(rig.gateway, frozenset(identities)) + served: Final = _closed_chat(rig, name, marker) + _assert_served(rig, served, marker, "/v1/chat/completions") + patched: Final = rig.gateway.request( + "PATCH", f"/model/{identities[0]}/update", {"model_info": dict(_EMBED_ONLY)} + ) + assert patched.status_code == 200, patched.text + rejected: Final = held( + lambda: _closed_chat(rig, name, _marker()), + lambda response: response.status_code == 400, + holding=_RELOAD_SECONDS * 2, + seconds=_RELOAD_SECONDS * 12, + ) + _assert_error(rejected, _openai_error("/v1/chat/completions", name, "/v1/embeddings")) + restored: Final = rig.gateway.request("PATCH", f"/model/{identities[0]}/update", {"model_info": chat_capable}) + assert restored.status_code == 200, restored.text + held( + lambda: _closed_chat(rig, name, _marker()), + lambda response: response.status_code == 200, + holding=_RELOAD_SECONDS * 2, + seconds=_RELOAD_SECONDS * 12, + ) + + +_TESTS_ROOT: Final = Path(__file__).resolve().parents[2] +_CONFIG: Final = TypeAdapter(dict[str, JsonValue]) + + +def _owned_embedding(owned: Gateway, model: str, marker: str) -> httpx.Response: + return owned.request("POST", "/v1/embeddings", {"model": model, "input": marker, "encoding_format": "float"}) + + +@dataclass(frozen=True, slots=True) +class _Routed: + gateway: Gateway + upstream: _Upstream + names: Mapping[str, str] + + +def _routed_entry(name: str, api_base: str, litellm_model: str, info: Mapping[str, JsonValue]) -> dict[str, JsonValue]: + return { + "model_name": name, + "litellm_params": {"model": litellm_model, "api_base": api_base, "api_key": _PROVIDER_KEY}, + "model_info": dict(info), + } + + +def _routed_config(directory: Path, api_base: str, names: Mapping[str, str]) -> Path: + shared: Final = _CONFIG.validate_python(yaml.safe_load((_TESTS_ROOT / "integration/proxy_config.yaml").read_text())) + config: Final = { + **shared, + "model_list": [ + _routed_entry(names["embed"], api_base, f"openai/{_EMBEDDING_MODEL}", _EMBED_ONLY), + _routed_entry(names["chat"], api_base, f"openai/{_CHAT_MODEL}", {"mode": "chat"}), + ], + "router_settings": { + **object_value(shared["router_settings"]), + "num_retries": 0, + "default_fallbacks": [names["chat"]], + "model_group_alias": {names["alias_embed"]: names["embed"], names["alias_chat"]: names["chat"]}, + }, + } + path: Final = directory / "gate-routed-config.yaml" + path.write_text(yaml.safe_dump(config, sort_keys=False)) + return path + + +@pytest.fixture(scope="module") +def routed(tmp_path_factory: pytest.TempPathFactory) -> Iterator[_Routed]: + directory: Final = tmp_path_factory.mktemp("gate-routed") + suffix: Final = uuid.uuid4().hex[:12] + names: Final = MappingProxyType( + {label: f"gate-routed-{label}-{suffix}" for label in ("embed", "chat", "alias_embed", "alias_chat")} + ) + with ( + wire_server(answering_model_discovery(_respond)) as wire, + gateway_from_environment() as rig_gateway, + owned_proxy_process( + rig_gateway, + directory, + {}, + config=_routed_config(directory, wire.url + "/v1", names), + workers=2, + state=OwnedState.SHARED_WITH_WORKER, + ) as owned, + ): + eventually( + lambda: owned.gateway.request("POST", "/v1/chat/completions", _chat_body(names["chat"], "probe")), + lambda response: response.status_code == 200, + seconds=30, + ) + yield _Routed(owned.gateway, _Upstream(wire), names) + + +_ROUTED_TIMEOUT: Final = pytest.mark.timeout(2 * graceful_stop_seconds() + 120) +_DEFAULT_FALLBACK_ROUTES: Final[Mapping[str, str]] = MappingProxyType( + {"chat_v1": "/v1/chat/completions", "messages": "/v1/responses", "responses_v1": "/v1/responses"} +) + + +@_ROUTED_TIMEOUT +@pytest.mark.parametrize("route", tuple(_DEFAULT_FALLBACK_ROUTES)) +def test_router_default_fallbacks_answer_for_an_embedding_only_group(routed: _Routed, route: str) -> None: + marker: Final = _marker() + response: Final = _send(routed.gateway, _ROUTES[route], routed.names["embed"], marker) + assert response.status_code == 200, response.text + assert _SERVED_TEXT in response.text, response.text + assert routed.upstream.targets(marker) == (_DEFAULT_FALLBACK_ROUTES[route],) + + +_OPT_OUTS: Final[Mapping[str, Mapping[str, JsonValue]]] = MappingProxyType( + { + "empty_fallbacks": {"fallbacks": []}, + "null_fallbacks": {"fallbacks": None}, + "disable_fallbacks": {"disable_fallbacks": True}, + } +) + + +@_ROUTED_TIMEOUT +@pytest.mark.parametrize("opt_out", tuple(_OPT_OUTS)) +def test_request_opt_outs_keep_the_rejection_over_router_default_fallbacks(routed: _Routed, opt_out: str) -> None: + model: Final = routed.names["embed"] + marker: Final = _marker() + response: Final = routed.gateway.request( + "POST", "/v1/chat/completions", {**_chat_body(model, marker), **_OPT_OUTS[opt_out]} + ) + _assert_rejected_before( + routed.upstream, response, marker, _openai_error("/v1/chat/completions", model, "/v1/embeddings") + ) + + +@_ROUTED_TIMEOUT +def test_model_group_alias_to_an_embedding_only_group_falls_back_to_the_default(routed: _Routed) -> None: + marker: Final = _marker() + response: Final = routed.gateway.request( + "POST", "/v1/chat/completions", _chat_body(routed.names["alias_embed"], marker) + ) + assert response.status_code == 200, response.text + assert _SERVED_TEXT in response.text, response.text + assert routed.upstream.targets(marker) == ("/v1/chat/completions",) + + +@_ROUTED_TIMEOUT +def test_model_group_alias_to_an_embedding_only_group_is_rejected_without_fallbacks(routed: _Routed) -> None: + alias: Final = routed.names["alias_embed"] + marker: Final = _marker() + response: Final = routed.gateway.request( + "POST", "/v1/chat/completions", {**_chat_body(alias, marker), "disable_fallbacks": True} + ) + _assert_rejected_before( + routed.upstream, response, marker, _openai_error("/v1/chat/completions", alias, "/v1/embeddings") + ) + + +@_ROUTED_TIMEOUT +def test_model_group_alias_to_a_chat_group_keeps_reaching_the_upstream(routed: _Routed) -> None: + marker: Final = _marker() + response: Final = routed.gateway.request( + "POST", "/v1/chat/completions", _chat_body(routed.names["alias_chat"], marker) + ) + assert response.status_code == 200, response.text + assert _SERVED_TEXT in response.text, response.text + assert routed.upstream.targets(marker) == ("/v1/chat/completions",) + + +@_ROUTED_TIMEOUT +def test_model_group_alias_to_an_embedding_only_group_keeps_serving_embeddings(routed: _Routed) -> None: + marker: Final = _marker() + response: Final = _owned_embedding(routed.gateway, routed.names["alias_embed"], marker) + assert response.status_code == 200, response.text + assert list_value( + object_value(list_value(JSON_OBJECT.validate_json(response.content)["data"])[0])["embedding"] + ) == list(_VECTOR), response.text + assert routed.upstream.targets(marker) == ("/v1/embeddings",) + + +_REWRITER: Final = """from litellm.integrations.custom_logger import CustomLogger + + +class WebSocketModelRewriter(CustomLogger): + async def async_pre_call_hook(self, user_api_key_dict, cache, data, call_type): + if call_type == "_aresponses_websocket" and data.get("model") == "__SOURCE__": + data["model"] = "__TARGET__" + return data + + +rewriter = WebSocketModelRewriter() +""" + + +@dataclass(frozen=True, slots=True) +class _Rewritten: + gateway: Gateway + upstream: _Upstream + source: str + target: str + + +def _rewritten_config(directory: Path, api_base: str, source: str, target: str, module: str) -> Path: + shared: Final = _CONFIG.validate_python(yaml.safe_load((_TESTS_ROOT / "integration/proxy_config.yaml").read_text())) + config: Final = { + **shared, + "model_list": [ + _routed_entry(source, api_base, f"openai/{_CHAT_MODEL}", {"mode": "chat"}), + _routed_entry(target, api_base, f"openai/{_EMBEDDING_MODEL}", _EMBED_ONLY), + ], + "router_settings": {**object_value(shared["router_settings"]), "num_retries": 0}, + "litellm_settings": {**object_value(shared["litellm_settings"]), "callbacks": [f"{module}.rewriter"]}, + } + path: Final = directory / "gate-rewritten-config.yaml" + path.write_text(yaml.safe_dump(config, sort_keys=False)) + return path + + +@pytest.fixture(scope="module") +def rewritten(tmp_path_factory: pytest.TempPathFactory) -> Iterator[_Rewritten]: + directory: Final = tmp_path_factory.mktemp("gate-rewritten") + suffix: Final = uuid.uuid4().hex[:12] + source: Final = f"gate-rewrite-chat-{suffix}" + target: Final = f"gate-rewrite-embed-{suffix}" + module: Final = f"gate_ws_rewriter_{suffix}" + (directory / f"{module}.py").write_text(_REWRITER.replace("__SOURCE__", source).replace("__TARGET__", target)) + with ( + wire_server(answering_model_discovery(_respond)) as wire, + gateway_from_environment() as rig_gateway, + owned_proxy_process( + rig_gateway, + directory, + {}, + config=_rewritten_config(directory, wire.url + "/v1", source, target, module), + workers=2, + state=OwnedState.SHARED_WITH_WORKER, + ) as owned, + ): + eventually( + lambda: owned.gateway.request("POST", "/v1/chat/completions", _chat_body(source, "probe")), + lambda response: response.status_code == 200, + seconds=30, + ) + yield _Rewritten(owned.gateway, _Upstream(wire), source, target) + + +@_ROUTED_TIMEOUT +@pytest.mark.parametrize("limit", tuple(_RESERVING_KEYS)) +def test_a_websocket_model_a_hook_rewrites_to_an_embedding_only_group_releases_the_key_reservation( + rewritten: _Rewritten, limit: str +) -> None: + markers: Final = tuple(_marker() for _ in range(3)) + served_marker: Final = _marker() + with rewritten.gateway.scenario() as scenario: + key: Final = scenario.key(**_RESERVING_KEYS[limit]) + closed: Final = tuple( + asyncio.run(_websocket_first_frame_and_close(rewritten.gateway, key, rewritten.source, marker)) + for marker in markers + ) + assert closed == (_websocket_rejection(rewritten.target),) * 3 + served: Final = rewritten.gateway.request( + "POST", "/v1/chat/completions", _chat_body(rewritten.source, served_marker), key=key + ) + assert served.status_code == 200, served.text + assert _SERVED_TEXT in served.text, served.text + assert [rewritten.upstream.calls(marker) for marker in markers] == [()] * 3 + assert rewritten.upstream.targets(served_marker) == ("/v1/chat/completions",) + + +@dataclass(frozen=True, slots=True) +class _Addressed: + gateway: Gateway + rig: Gateway + upstream: _Upstream + names: Mapping[str, str] + addresses: Mapping[str, str] + api_base: str + + +def _wildcard_backend_entry(name: str, api_base: str, backend: str) -> dict[str, JsonValue]: + entry: Final = _routed_entry(name, api_base, backend, _EMBED_ONLY) + return {**entry, "litellm_params": {**object_value(entry["litellm_params"]), "custom_llm_provider": "openai"}} + + +def _addressed_config( + directory: Path, api_base: str, names: Mapping[str, str], suffix: str, with_wildcards: bool +) -> Path: + shared: Final = _CONFIG.validate_python(yaml.safe_load((_TESTS_ROOT / "integration/proxy_config.yaml").read_text())) + wildcard: Final = f"gate-addr-wild-{suffix}" + wildcard_entries: Final = ( + ( + _routed_entry(f"{wildcard}/*", api_base, "openai/*", {}), + _wildcard_backend_entry(names["wildcard_backend"], api_base, f"{wildcard}/{_CHAT_MODEL}"), + ) + if with_wildcards + else () + ) + aliases: Final = {names["wildcard_alias"]: f"{wildcard}/{_CHAT_BACKEND_PREFIX}{suffix}"} if with_wildcards else {} + config: Final = { + **shared, + "model_list": [ + _routed_entry( + names["embed"], api_base, f"openai/gate-addr-backend-{suffix}", {**_EMBED_ONLY, "id": names["embed_id"]} + ), + _routed_entry(names["chat"], api_base, f"openai/{_CHAT_MODEL}", {"mode": "chat", "id": names["chat_id"]}), + _routed_entry(names["shared_embed"], api_base, f"openai/{_CHAT_MODEL}", _EMBED_ONLY), + *wildcard_entries, + _routed_entry(names["ghost_primary"], api_base, f"openai/{_EMBEDDING_MODEL}", _EMBED_ONLY), + ], + "router_settings": { + **object_value(shared["router_settings"]), + "num_retries": 0, + "model_group_alias": aliases, + "fallbacks": [{names["ghost_primary"]: [names["ghost"]]}, {names["ghost"]: [names["chat"]]}], + }, + "general_settings": { + **object_value(shared["general_settings"]), + "proxy_config_reload_interval_seconds": 3600, + }, + } + path: Final = directory / "gate-addressed-config.yaml" + path.write_text(yaml.safe_dump(config, sort_keys=False)) + return path + + +@contextlib.contextmanager +def _addressed_proxy(tmp_path_factory: pytest.TempPathFactory, with_wildcards: bool) -> Iterator[_Addressed]: + directory: Final = tmp_path_factory.mktemp("gate-addressed") + suffix: Final = uuid.uuid4().hex[:12] + names: Final = MappingProxyType( + { + label: f"gate-addr-{label}-{suffix}" + for label in ( + "embed", + "embed_id", + "chat", + "chat_id", + "shared_embed", + "wildcard_backend", + "wildcard_alias", + "ghost_primary", + "ghost", + ) + } + ) + with ( + wire_server(answering_model_discovery(_respond)) as wire, + gateway_from_environment() as rig_gateway, + owned_proxy_process( + rig_gateway, + directory, + {}, + config=_addressed_config(directory, wire.url + "/v1", names, suffix, with_wildcards), + workers=2, + state=OwnedState.SHARED_WITH_WORKER, + ) as owned, + ): + eventually( + lambda: owned.gateway.request("POST", "/v1/chat/completions", _chat_body(names["chat"], "probe")), + lambda response: response.status_code == 200, + seconds=30, + ) + yield _Addressed( + owned.gateway, + rig_gateway, + _Upstream(wire), + names, + MappingProxyType( + { + "deployment_id": names["embed_id"], + "backend_name": f"openai/gate-addr-backend-{suffix}", + "wildcard_backend_name": f"gate-addr-wild-{suffix}/{_CHAT_MODEL}", + } + ), + wire.url + "/v1", + ) + + +@pytest.fixture(scope="module") +def addressed(tmp_path_factory: pytest.TempPathFactory) -> Iterator[_Addressed]: + with _addressed_proxy(tmp_path_factory, with_wildcards=True) as proxy: + yield proxy + + +@pytest.fixture(scope="module") +def addressed_without_wildcards(tmp_path_factory: pytest.TempPathFactory) -> Iterator[_Addressed]: + with _addressed_proxy(tmp_path_factory, with_wildcards=False) as proxy: + yield proxy + + +_ADDRESSED_ROUTES: Final = ("chat_v1", "messages", "responses_v1") + + +@_ROUTED_TIMEOUT +@pytest.mark.parametrize("route", _ADDRESSED_ROUTES) +@pytest.mark.parametrize("address", ("deployment_id", "backend_name")) +def test_chat_sent_to_an_embedding_only_deployment_id_or_backend_name_is_rejected_before_the_upstream( + addressed: _Addressed, address: str, route: str +) -> None: + spec: Final = _ROUTES[route] + model: Final = addressed.addresses[address] + marker: Final = _marker() + response: Final = _send(addressed.gateway, spec, model, marker) + expected: Final = ( + _anthropic_error(model, "/v1/embeddings") + if spec.anthropic_shape + else _openai_error(spec.endpoint, model, "/v1/embeddings") + ) + _assert_rejected_before(addressed.upstream, response, marker, expected) + + +@_ROUTED_TIMEOUT +@pytest.mark.parametrize("address", ("deployment_id", "backend_name")) +def test_embeddings_sent_to_an_embedding_only_deployment_id_or_backend_name_keep_working( + addressed: _Addressed, address: str +) -> None: + marker: Final = _marker() + response: Final = _owned_embedding(addressed.gateway, addressed.addresses[address], marker) + assert response.status_code == 200, response.text + assert list_value( + object_value(list_value(JSON_OBJECT.validate_json(response.content)["data"])[0])["embedding"] + ) == list(_VECTOR), response.text + assert addressed.upstream.targets(marker) == ("/v1/embeddings",) + + +@_ROUTED_TIMEOUT +@pytest.mark.parametrize("address", ("chat_deployment_id", "backend_name_a_chat_deployment_shares")) +def test_chat_sent_to_a_chat_capable_deployment_id_or_backend_name_keeps_reaching_the_upstream( + addressed: _Addressed, address: str +) -> None: + model: Final = addressed.names["chat_id"] if address == "chat_deployment_id" else f"openai/{_CHAT_MODEL}" + marker: Final = _marker() + response: Final = addressed.gateway.request("POST", "/v1/chat/completions", _chat_body(model, marker)) + assert response.status_code == 200, response.text + assert _SERVED_TEXT in response.text, response.text + assert addressed.upstream.targets(marker) == ("/v1/chat/completions",) + + +@_ROUTED_TIMEOUT +def test_request_fallbacks_naming_a_chat_deployment_id_answer_for_an_embedding_only_primary( + addressed: _Addressed, +) -> None: + marker: Final = _marker() + response: Final = addressed.gateway.request( + "POST", + "/v1/chat/completions", + {**_chat_body(addressed.names["embed"], marker), "fallbacks": [addressed.names["chat_id"]]}, + ) + assert response.status_code == 200, response.text + assert _SERVED_TEXT in response.text, response.text + assert addressed.upstream.targets(marker) == ("/v1/chat/completions",) + + +@_ROUTED_TIMEOUT +def test_chat_sent_to_an_embedding_only_backend_name_a_wildcard_also_matches_is_served_by_the_wildcard( + addressed: _Addressed, +) -> None: + marker: Final = _marker() + response: Final = addressed.gateway.request( + "POST", "/v1/chat/completions", _chat_body(addressed.addresses["wildcard_backend_name"], marker) + ) + assert response.status_code == 200, response.text + assert _SERVED_TEXT in response.text, response.text + assert [JSON_OBJECT.validate_json(call.body)["model"] for call in addressed.upstream.calls(marker)] == [_CHAT_MODEL] + + +@_ROUTED_TIMEOUT +def test_request_fallbacks_naming_an_alias_only_a_wildcard_serves_answer_for_an_embedding_only_primary( + addressed: _Addressed, +) -> None: + marker: Final = _marker() + response: Final = addressed.gateway.request( + "POST", + "/v1/chat/completions", + {**_chat_body(addressed.names["embed"], marker), "fallbacks": [addressed.names["wildcard_alias"]]}, + ) + _assert_fallback_answered(addressed.upstream, response, marker) + + +@_ROUTED_TIMEOUT +def test_yaml_fallbacks_through_a_group_with_no_deployments_answer_for_an_embedding_only_group( + addressed: _Addressed, +) -> None: + marker: Final = _marker() + response: Final = addressed.gateway.request( + "POST", "/v1/chat/completions", _chat_body(addressed.names["ghost_primary"], marker) + ) + _assert_fallback_answered(addressed.upstream, response, marker) + + +@_ROUTED_TIMEOUT +def test_an_embedding_only_model_the_read_through_loads_from_the_database_is_rejected_on_every_call( + addressed_without_wildcards: _Addressed, +) -> None: + addressed: Final = addressed_without_wildcards + name: Final = f"gate-addr-loaded-{uuid.uuid4().hex[:12]}" + markers: Final = tuple(_marker() for _ in range(6)) + with addressed.rig.scenario() as scenario: + identity: Final = _create( + addressed.rig, name, {"model": f"openai/{_EMBEDDING_MODEL}", "api_base": addressed.api_base}, _EMBED_ONLY + ) + scenario.cleanups.callback(scenario.delete_model, identity) + responses: Final = tuple( + addressed.gateway.request("POST", "/v1/chat/completions", _chat_body(name, marker), headers=_CLOSE) + for marker in markers + ) + expected: Final = _openai_error("/v1/chat/completions", name, "/v1/embeddings") + for response in responses: + _assert_error(response, expected) + assert [addressed.upstream.calls(marker) for marker in markers] == [()] * len(markers) + + +@pytest.mark.parametrize("metadata", ({"model_group": 7}, {"model_group": ["wrong-group"]})) +def test_rejection_with_invalid_metadata_is_logged_without_blocking_embeddings( + rig: _Rig, metadata: dict[str, JsonValue] +) -> None: + marker: Final = _marker() + model: Final = rig.names["embedding_only"] + rejected: Final = _chat(rig, model, marker, metadata=metadata) + _assert_rejected_chat(rig, rejected, marker, model) + embedded_marker: Final = _marker() + embedded: Final = rig.gateway.request( + "POST", "/v1/embeddings", {"model": model, "input": embedded_marker, "encoding_format": "float"} + ) + assert embedded.status_code == 200, embedded.text + assert rig.upstream.targets(embedded_marker) == ("/v1/embeddings",) + + +def test_existing_non_chat_metadata_on_a_chat_backend_remains_advisory(rig: _Rig) -> None: + marker: Final = _marker() + with rig.gateway.scenario() as scenario: + name, identities = _deploy( + scenario, + rig.api_base, + "legacy-audio", + _single({"mode": "audio_speech", "supported_endpoints": ["/v1/audio/speech"]}), + ) + _served_by_every_worker(rig.gateway, frozenset(identities)) + response: Final = _chat(rig, name, marker) + _assert_served(rig, response, marker, "/v1/chat/completions") + + +def test_a_mixed_group_skips_its_restricted_embedding_deployment(rig: _Rig) -> None: + marker: Final = _marker() + with rig.gateway.scenario() as scenario: + name: Final = f"gate-mixed-{uuid.uuid4().hex[:12]}" + embed: Final = _create( + scenario.gateway, + name, + {"model": f"openai/{_EMBEDDING_MODEL}", "api_base": rig.api_base, "order": 0}, + _EMBED_ONLY, + ) + scenario.cleanups.callback(scenario.delete_model, embed) + chat: Final = _create( + scenario.gateway, name, {"model": f"openai/{_CHAT_MODEL}", "api_base": rig.api_base, "order": 1}, {} + ) + scenario.cleanups.callback(scenario.delete_model, chat) + _served_by_every_worker(rig.gateway, frozenset({embed, chat})) + response: Final = _chat(rig, name, marker) + _assert_served(rig, response, marker, "/v1/chat/completions") + calls: Final = rig.upstream.calls(marker) + assert [JSON_OBJECT.validate_json(call.body)["model"] for call in calls] == [_CHAT_MODEL] + + +@pytest.mark.parametrize( + "fallbacks, expected_status", + [(("embedding_only",), 429), (("embedding_only", "chat_target"), 200), (("missing",), 400)], +) +def test_local_rate_limit_skips_restricted_fallbacks_and_still_tries_unloaded_ones( + rig: _Rig, fallbacks: tuple[str, ...], expected_status: int +) -> None: + limited: Final = rig.names["chat_rate_limited"] + targets: Final = tuple(rig.names[name] if name != "missing" else _marker() for name in fallbacks) + with rig.gateway.scenario() as scenario: + key: Final = scenario.key( + model_rpm_limit={limited: 1}, router_settings={"fallbacks": [{limited: list(targets)}]} + ) + first: Final = _chat(rig, limited, _marker(), key=key) + assert first.status_code == 200, first.text + marker: Final = _marker() + second: Final = _chat(rig, limited, marker, key=key) + assert second.status_code == expected_status, second.text + if expected_status == 200: + _assert_served(rig, second, marker, "/v1/chat/completions") + return + assert "Supported endpoints" not in second.text, second.text + assert rig.upstream.calls(marker) == () + + +def _websocket_fallback_config(directory: Path, primary: tuple[str, str], fallback: tuple[str, str]) -> Path: + shared: Final = _CONFIG.validate_python(yaml.safe_load((_TESTS_ROOT / "integration/proxy_config.yaml").read_text())) + config: Final = { + **shared, + "model_list": [ + _routed_entry(primary[0], primary[1], f"openai/{_EMBEDDING_MODEL}", _EMBED_ONLY), + _routed_entry(fallback[0], fallback[1], f"openai/{_CHAT_MODEL}", {"mode": "chat"}), + ], + "router_settings": { + **object_value(shared["router_settings"]), + "num_retries": 0, + "fallbacks": [{primary[0]: [fallback[0]]}], + }, + } + path: Final = directory / "gate-websocket-fallback.yaml" + path.write_text(yaml.safe_dump(config, sort_keys=False)) + return path + + +def test_responses_websocket_uses_a_configured_chat_fallback_without_calling_the_embedding_provider( + rig: _Rig, tmp_path: Path +) -> None: + marker: Final = _marker() + suffix: Final = uuid.uuid4().hex[:12] + primary: Final = f"gate-ws-embed-{suffix}" + fallback: Final = f"gate-ws-chat-{suffix}" + certificate: Final = write_self_signed_cert(tmp_path) + with ( + responses_ws_peer.responses_peer( + certificate, lambda identity: responses_ws_peer.events(identity, _SERVED_TEXT) + ) as peer, + owned_proxy_process( + rig.gateway, + tmp_path, + {"SSL_CERT_FILE": str(certificate[0])}, + config=_websocket_fallback_config(tmp_path, (primary, f"{rig.api_base}/{marker}"), (fallback, peer.url)), + workers=2, + state=OwnedState.SHARED_WITH_WORKER, + ) as owned, + owned.gateway.scenario() as scenario, + ): + key: Final = scenario.key() + (received,) = asyncio.run( + responses_ws_peer.session(str(owned.gateway.client.base_url), key, primary, (marker,)) + ) + assert received[-1]["type"] == "response.completed", received + assert responses_ws_peer.completed(received)["status"] == "completed" + assert any(event.get("delta") == _SERVED_TEXT for event in received), received + (forwarded,) = responses_ws_peer.drain(peer.frames) + assert marker in json.dumps(forwarded), forwarded + assert rig.upstream.calls(marker) == () + + +def test_client_routing_fields_cannot_disable_an_explicit_restriction(rig: _Rig) -> None: + marker: Final = _marker() + with rig.gateway.scenario() as scenario: + model, identities = _deploy(scenario, rig.api_base, "restricted-chat", _single(_EMBED_ONLY)) + _served_by_every_worker(rig.gateway, frozenset(identities)) + response: Final = _chat( + rig, + model, + marker, + original_generic_function="caller-value", + additional_drop_params=["original_generic_function"], + ) + _assert_rejected_chat(rig, response, marker, model) + + +def test_a_compaction_endpoint_declaration_keeps_the_model_available_for_compaction(rig: _Rig) -> None: + marker: Final = _marker() + with rig.gateway.scenario() as scenario: + model, identities = _deploy( + scenario, + rig.api_base, + "compact-declaration", + _single({"supported_endpoints": ["/responses/compact"], "enforce_supported_endpoints": True}), + ) + _served_by_every_worker(rig.gateway, frozenset(identities)) + response: Final = rig.gateway.request("POST", "/responses/compact", {"model": model, "input": marker}) + _assert_served(rig, response, marker, "/v1/responses/compact") + + +@pytest.mark.parametrize("alias_source", ["key_alias", "key_router", "team_router"]) +def test_local_rate_limit_fallbacks_resolve_saved_aliases(rig: _Rig, tmp_path: Path, alias_source: str) -> None: + suffix: Final = uuid.uuid4().hex[:12] + names: Final = {label: f"gate-fallback-alias-{label}-{suffix}" for label in ("primary", "target", "alias")} + shared: Final = _CONFIG.validate_python(yaml.safe_load((_TESTS_ROOT / "integration/proxy_config.yaml").read_text())) + config: Final = { + **shared, + "model_list": [ + _routed_entry(names["primary"], rig.api_base, f"openai/{_CHAT_MODEL}", {"mode": "chat"}), + _routed_entry(names["target"], rig.api_base, f"openai/{_CHAT_MODEL}", {"mode": "chat"}), + ], + "router_settings": { + **object_value(shared["router_settings"]), + "num_retries": 0, + "fallbacks": [{names["primary"]: [names["alias"]]}], + }, + } + path: Final = tmp_path / "fallback-alias.yaml" + path.write_text(yaml.safe_dump(config, sort_keys=False)) + aliases: Final = {names["alias"]: names["target"]} + with ( + owned_proxy_process( + rig.gateway, tmp_path, {}, config=path, workers=2, state=OwnedState.SHARED_WITH_WORKER + ) as owned, + owned.gateway.scenario() as scenario, + ): + team: Final = ( + scenario.team(router_settings={"model_group_alias": aliases}) if alias_source == "team_router" else None + ) + key: Final = scenario.key( + team_id=team, + aliases=aliases if alias_source == "key_alias" else {}, + router_settings={"model_group_alias": aliases} if alias_source == "key_router" else None, + model_rpm_limit={names["primary"]: 1}, + ) + first: Final = owned.gateway.request( + "POST", "/v1/chat/completions", _chat_body(names["primary"], _marker()), key=key + ) + assert first.status_code == 200, first.text + marker: Final = _marker() + second: Final = owned.gateway.request( + "POST", "/v1/chat/completions", _chat_body(names["primary"], marker), key=key + ) + _assert_fallback_answered(rig.upstream, second, marker) diff --git a/tests/unit/proxy/common_utils/test_callback_utils.py b/tests/unit/proxy/common_utils/test_callback_utils.py index ecb2375d495..c8d909ac2eb 100644 --- a/tests/unit/proxy/common_utils/test_callback_utils.py +++ b/tests/unit/proxy/common_utils/test_callback_utils.py @@ -49,6 +49,19 @@ def test_get_remaining_tokens_and_requests_from_request_data(): } +@pytest.mark.parametrize("client_model_group", [7, ["embed-only"], {"name": "embed-only"}]) +def test_remaining_limit_headers_ignore_a_model_group_that_is_not_a_string(client_model_group: object): + casedata: Final = { + "metadata": { + "model_group": client_model_group, + f"litellm-key-remaining-requests-{client_model_group}": 100, + f"litellm-key-remaining-tokens-{client_model_group}": 200, + } + } + + assert get_remaining_tokens_and_requests_from_request_data(casedata) == {} + + @patch( "litellm.proxy.common_utils.callback_utils.CustomLogger.get_callback_env_vars", return_value=["API_KEY", "MISSING_VAR"], diff --git a/tests/unit/proxy/management_endpoints/test_model_management_endpoints.py b/tests/unit/proxy/management_endpoints/test_model_management_endpoints.py index b892dc97e5c..835e7605955 100644 --- a/tests/unit/proxy/management_endpoints/test_model_management_endpoints.py +++ b/tests/unit/proxy/management_endpoints/test_model_management_endpoints.py @@ -3617,6 +3617,33 @@ class TestDeleteTeamModels: assert any(e[0] == "delete_many" for e in prisma.events) +@pytest.mark.parametrize( + ("model_info_patch", "expected_gate"), + [ + pytest.param({"tpm": 5}, True, id="unrelated-field"), + pytest.param({"supported_endpoints": ["/v1/embeddings", "/v1/rerank"]}, True, id="endpoint-list-only"), + pytest.param({"enforce_supported_endpoints": False}, False, id="explicit-false"), + ], +) +def test_model_info_patch_changes_the_endpoint_gate_only_when_it_sends_the_flag( + model_info_patch: dict[str, object], expected_gate: bool +): + from litellm.proxy.management_endpoints.model_management_endpoints import update_db_model + + result: Final = update_db_model( + db_model=Deployment( + model_name="embedding-only", + litellm_params=LiteLLM_Params(model="openai/text-embedding-3-small"), + model_info=ModelInfo( + id="embedding-only", enforce_supported_endpoints=True, supported_endpoints=["/v1/embeddings"] + ), + ), + updated_patch=updateDeployment.model_validate({"model_info": model_info_patch}), + ) + + assert json.loads(result["model_info"]).get("enforce_supported_endpoints") is expected_gate + + def _build_db_model_for_blocked_test(): from litellm.types.router import Deployment, LiteLLM_Params, ModelInfo diff --git a/tests/unit/proxy/response_api_endpoints/test_endpoints.py b/tests/unit/proxy/response_api_endpoints/test_endpoints.py index e3d514b2ea9..b11dcc58ef6 100644 --- a/tests/unit/proxy/response_api_endpoints/test_endpoints.py +++ b/tests/unit/proxy/response_api_endpoints/test_endpoints.py @@ -2055,6 +2055,22 @@ def test_responses_ws_failure_frame_uses_the_http_status_for_the_same_failure( assert frame["error"]["message"] == handle_exception_on_proxy(failure).message +def test_responses_ws_failure_frame_shows_only_the_refusal_for_a_model_that_refuses_the_endpoint(): + from litellm.exceptions import EndpointNotSupportedByModelError + from litellm.proxy.response_api_endpoints.endpoints import _responses_ws_failure_frame + + refusal: Final = "/v1/responses is not supported by model=embed-only. Supported endpoints: /v1/embeddings" + failure: Final = EndpointNotSupportedByModelError( + message=f"{refusal}\n\nLiteLLM: model group 'embed-only' failed with the error above", model="embed-only" + ) + + assert json.loads(_responses_ws_failure_frame(failure)) == { + "type": "error", + "status": 400, + "error": {"type": "invalid_request_error", "message": refusal}, + } + + class TestWSSessionCostTracking: @pytest.mark.asyncio async def test_router_budget_limiter_skips_arealtime_call_type(self): diff --git a/tests/unit/proxy/spend_tracking/test_spend_tracking_utils.py b/tests/unit/proxy/spend_tracking/test_spend_tracking_utils.py index 23f9ac0e2d8..740334252d2 100644 --- a/tests/unit/proxy/spend_tracking/test_spend_tracking_utils.py +++ b/tests/unit/proxy/spend_tracking/test_spend_tracking_utils.py @@ -5384,6 +5384,28 @@ def test_get_logging_payload_router_rejected_request_takes_provider_from_model_g assert payload["custom_llm_provider"] == expected_provider +@pytest.mark.parametrize("client_model_group", [7, ["embed-only"], {"name": "embed-only"}]) +def test_get_logging_payload_drops_a_client_model_group_that_is_not_a_string(client_model_group: object): + payload: Final = get_logging_payload( + kwargs={ + "model": "embed-only", + "litellm_params": { + "metadata": {"user_api_key": "test-key", "model_group": client_model_group, "status": "failure"} + }, + "standard_logging_object": { + **_make_failed_request_standard_logging_payload(), + "model_group": client_model_group, + }, + }, + response_obj={}, + start_time=datetime.datetime.now(timezone.utc), + end_time=datetime.datetime.now(timezone.utc), + llm_router=None, + ) + + assert payload["model_group"] == "" + + def test_get_logging_payload_router_rejected_request_without_router_leaves_provider_empty(): assert _router_rejected_failure_payload("openai-group", None)["custom_llm_provider"] == "" diff --git a/tests/unit/proxy/test_common_request_processing.py b/tests/unit/proxy/test_common_request_processing.py index fdcbcbfcd05..e58a1de9eb5 100644 --- a/tests/unit/proxy/test_common_request_processing.py +++ b/tests/unit/proxy/test_common_request_processing.py @@ -456,7 +456,7 @@ class TestProxyBaseLLMRequestProcessing: mock_add_litellm_data_to_request, ) mock_general_settings = {} - mock_user_api_key_dict = MagicMock(spec=UserAPIKeyAuth) + mock_user_api_key_dict = ProxyUserAPIKeyAuth() mock_proxy_config = MagicMock(spec=ProxyConfig) route_type = "acompletion" @@ -740,7 +740,7 @@ class TestProxyBaseLLMRequestProcessing: returned_data, logging_obj = await processing_obj.common_processing_pre_call_logic( request=mock_request, general_settings={}, - user_api_key_dict=MagicMock(spec=UserAPIKeyAuth), + user_api_key_dict=ProxyUserAPIKeyAuth(), proxy_logging_obj=mock_proxy_logging_obj, proxy_config=MagicMock(spec=ProxyConfig), route_type=route_type, @@ -1186,7 +1186,7 @@ class TestProxyBaseLLMRequestProcessing: ) mock_general_settings = {} - mock_user_api_key_dict = MagicMock(spec=UserAPIKeyAuth) + mock_user_api_key_dict = ProxyUserAPIKeyAuth() mock_proxy_config = MagicMock(spec=ProxyConfig) mock_router_settings = { @@ -2208,7 +2208,7 @@ class TestProxyBaseLLMRequestProcessing: mock_add_litellm_data_to_request, ) mock_general_settings = {} - mock_user_api_key_dict = MagicMock(spec=UserAPIKeyAuth) + mock_user_api_key_dict = ProxyUserAPIKeyAuth() mock_proxy_config = MagicMock(spec=ProxyConfig) route_type = "acompletion" @@ -7706,6 +7706,56 @@ class TestPreCallWithFallbacksOnLocalRateLimit: assert processor.data["model"] == fallback_model assert call_count == 2 + @pytest.mark.asyncio + async def test_local_rate_limit_fallback_skips_an_aliased_embedding_only_model(self): + from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing + from litellm.proxy.common_utils.proxy_rate_limit_error import ProxyRateLimitError + + router: Final = Router( + model_list=[ + { + "model_name": "embed-only", + "litellm_params": {"model": "openai/text-embedding-3-small", "api_key": "sk-synthetic"}, + "model_info": {"enforce_supported_endpoints": True, "supported_endpoints": ["/v1/embeddings"]}, + }, + {"model_name": "chat-model", "litellm_params": {"model": "openai/gpt-5.4-nano", "api_key": "sk-synthetic"}}, + ], + fallbacks=[{"rate-limited": ["embed-alias", "chat-model"]}], + ) + processor: Final = ProxyBaseLLMRequestProcessing( + data={ + "model": "rate-limited", + "router_settings_override": {"model_group_alias": {"embed-alias": "embed-only"}}, + } + ) + attempted: Final[list[object]] = [] + + async def pre_call_logic(**_: object): + attempted.append(processor.data.get("model")) + if processor.data.get("model") == "rate-limited": + raise ProxyRateLimitError(detail="key rpm limit exceeded", headers={"retry-after": "30"}) + return processor.data, MagicMock() + + with patch.object(processor, "common_processing_pre_call_logic", side_effect=pre_call_logic): + await processor._pre_call_with_fallbacks( + request=MagicMock(), + general_settings={}, + proxy_logging_obj=MagicMock(), + user_api_key_dict=ProxyUserAPIKeyAuth(), + version=None, + proxy_config=MagicMock(), + user_model=None, + user_temperature=None, + user_request_timeout=None, + user_max_tokens=None, + user_api_base=None, + model="rate-limited", + route_type="acompletion", + llm_router=router, + ) + + assert attempted == ["rate-limited", "chat-model"] + @pytest.mark.asyncio async def test_raises_when_no_fallbacks_configured(self): from litellm.proxy.common_utils.proxy_rate_limit_error import ProxyRateLimitError diff --git a/tests/unit/proxy/test_route_llm_request.py b/tests/unit/proxy/test_route_llm_request.py index 9e31d3774f9..8f364e51e2b 100644 --- a/tests/unit/proxy/test_route_llm_request.py +++ b/tests/unit/proxy/test_route_llm_request.py @@ -28,6 +28,7 @@ from litellm.proxy.route_llm_request import ( JSON_OBJECT_ADAPTER, ProxyModelNotFoundError, UserConfigRouters, + endpoint_unsupported_by_model, route_request, user_config_routers, ) @@ -3027,3 +3028,27 @@ async def test_user_config_routers_discard_the_least_recently_used_router_past_c assert first_handler() is not None assert second_handler() is None routers.clear() + + +def test_fallback_probe_skips_only_a_loaded_model_that_refuses_the_route() -> None: + router: Final = litellm.Router( + model_list=[ + { + "model_name": "chat-fallback", + "litellm_params": {"model": "openai/gpt-5.4-nano", "api_key": "unit-provider-key"}, + }, + { + "model_name": "embed-only", + "litellm_params": {"model": "openai/text-embedding-3-small", "api_key": "unit-provider-key"}, + "model_info": {"enforce_supported_endpoints": True, "supported_endpoints": ["/v1/embeddings"]}, + }, + ], + num_retries=0, + ) + + def skipped(model: str) -> bool: + return endpoint_unsupported_by_model("acompletion", {"model": model}, router, None) + + assert skipped("embed-only") is True + assert skipped("chat-fallback") is False + assert skipped("added-on-another-worker") is False diff --git a/tests/unit/router_utils/test_endpoint_eligibility.py b/tests/unit/router_utils/test_endpoint_eligibility.py new file mode 100644 index 00000000000..5e5d77254d0 --- /dev/null +++ b/tests/unit/router_utils/test_endpoint_eligibility.py @@ -0,0 +1,139 @@ +from collections.abc import Mapping +from typing import Final + +import pytest +from pydantic import TypeAdapter + +from litellm.exceptions import EndpointNotSupportedByModelError +from litellm.router_utils.endpoint_eligibility import ( + eligible_deployments, + endpoint_selection, + warn_on_unenforced_endpoint_flag, +) + +_DEPLOYMENT: Final = TypeAdapter(list[dict[str, object]]) + + +def _deployment(info: Mapping[str, object], identity: str = "embedding") -> dict[str, object]: + return _DEPLOYMENT.validate_python( + [ + { + "model_name": "shared", + "litellm_params": {"model": "openai/test-model"}, + "model_info": {"id": identity, **info}, + } + ] + )[0] + + +@pytest.mark.parametrize("call_type", ["acompletion", "anthropic_messages", "aresponses", "acreate_interaction"]) +def test_explicit_restriction_excludes_an_embedding_deployment_from_a_mixed_group(call_type: str) -> None: + embedding: Final = _deployment({"enforce_supported_endpoints": True, "supported_endpoints": ["/v1/embeddings"]}) + chat: Final = _deployment({"supported_endpoints": ["/v1/chat/completions"]}, "chat") + with endpoint_selection(call_type): + result: Final = eligible_deployments("shared", (embedding, chat)) + assert result == (chat,) + + +@pytest.mark.parametrize( + "info", + [ + {"mode": "audio_speech", "supported_endpoints": ["/v1/audio/speech"]}, + {"enforce_supported_endpoints": False, "supported_endpoints": ["/v1/embeddings"]}, + {"enforce_supported_endpoints": True}, + {"enforce_supported_endpoints": True, "supported_endpoints": None}, + {"enforce_supported_endpoints": True, "supported_endpoints": 7}, + {"enforce_supported_endpoints": True, "supported_endpoints": ["/v1/messages"]}, + ], +) +def test_advisory_or_incomplete_metadata_preserves_routing(info: Mapping[str, object]) -> None: + deployment: Final = _deployment(info) + with endpoint_selection("acompletion"): + result: Final = eligible_deployments("shared", (deployment,)) + assert result == (deployment,) + + +def test_a_single_endpoint_string_is_enforced_like_a_one_item_list() -> None: + deployment: Final = _deployment({"enforce_supported_endpoints": True, "supported_endpoints": "/v1/embeddings"}) + with endpoint_selection("acompletion"): + result: Final = eligible_deployments("shared", (deployment,)) + assert isinstance(result, EndpointNotSupportedByModelError) + assert result.message.endswith("Supported endpoints: /v1/embeddings") + + +def test_unsupported_error_explains_the_requested_operation_and_available_endpoints() -> None: + deployment: Final = _deployment({"enforce_supported_endpoints": True, "supported_endpoints": ["/v1/embeddings"]}) + with endpoint_selection("acreate_interaction"): + result: Final = eligible_deployments("shared", (deployment,)) + assert isinstance(result, EndpointNotSupportedByModelError) + assert result.status_code == 400 + assert "/v1beta/interactions" in result.message + assert "/v1/embeddings" in result.message + + +@pytest.mark.parametrize("call_type", ["aembedding", "aspeech", "_arealtime", "aget_responses"]) +def test_chat_restriction_does_not_block_other_operations(call_type: str) -> None: + deployment: Final = _deployment({"enforce_supported_endpoints": True, "supported_endpoints": ["/v1/embeddings"]}) + with endpoint_selection(call_type): + result: Final = eligible_deployments("shared", (deployment,)) + assert result == (deployment,) + + +def test_nested_non_chat_calls_do_not_inherit_the_outer_chat_restriction() -> None: + deployment: Final = _deployment({"enforce_supported_endpoints": True, "supported_endpoints": ["/v1/embeddings"]}) + with endpoint_selection("acompletion"): + with endpoint_selection("aembedding"): + embedding_result: Final = eligible_deployments("shared", (deployment,)) + chat_result: Final = eligible_deployments("shared", (deployment,)) + assert embedding_result == (deployment,) + assert isinstance(chat_result, EndpointNotSupportedByModelError) + + +@pytest.mark.parametrize( + "endpoint", + [ + "/messages", + "/interactions", + "/responses/compact", + "/openai/v1/responses/compact", + "/openai/v1/chat/completions", + "chat/completions", + "v1/responses", + " /v1/chat/completions ", + "/v1/chat/completions/", + "/V1/Chat/Completions", + "/v1beta/models/gemini-3-flash:generateContent", + ], +) +def test_chat_endpoint_declarations_keep_their_meaning_across_prefixes(endpoint: str) -> None: + deployment: Final = _deployment({"enforce_supported_endpoints": True, "supported_endpoints": [endpoint]}) + with endpoint_selection("acompletion"): + result: Final = eligible_deployments("shared", (deployment,)) + assert result == (deployment,) + + +@pytest.mark.parametrize( + "info", + [ + {"enforce_supported_endpoints": True}, + {"enforce_supported_endpoints": True, "supported_endpoints": ["/v1/embeddings", None]}, + {"enforce_supported_endpoints": "true", "supported_endpoints": 7}, + ], +) +def test_a_flag_without_a_usable_endpoint_list_is_reported(info: Mapping[str, object]) -> None: + warning: Final = warn_on_unenforced_endpoint_flag("embed-only", info) + assert warning is not None + assert "'embed-only'" in warning + + +@pytest.mark.parametrize( + "info", + [ + {"enforce_supported_endpoints": True, "supported_endpoints": ["/v1/embeddings"]}, + {"enforce_supported_endpoints": True, "supported_endpoints": "/v1/embeddings"}, + {"enforce_supported_endpoints": False}, + {"supported_endpoints": ["/v1/embeddings", None]}, + ], +) +def test_an_enforceable_or_unflagged_deployment_is_not_reported(info: Mapping[str, object]) -> None: + assert warn_on_unenforced_endpoint_flag("embed-only", info) is None diff --git a/tests/unit/test_router/test_router.py b/tests/unit/test_router/test_router.py index ee3a70c7c1c..62510dac26d 100644 --- a/tests/unit/test_router/test_router.py +++ b/tests/unit/test_router/test_router.py @@ -12270,6 +12270,67 @@ def test_is_deployment_blocked_static_helper_reflects_blocked_flag(): ) +def test_synchronous_text_completion_rejects_a_restricted_model_before_the_provider() -> None: + model: Final = "text-completion-openai/unit-text-model" + router: Final = Router( + model_list=[ + { + "model_name": model, + "litellm_params": { + "model": model, + "api_base": "https://sync-router.invalid/v1", + "api_key": "unit-provider-key", + }, + "model_info": {"enforce_supported_endpoints": True, "supported_endpoints": ["/v1/embeddings"]}, + } + ], + num_retries=0, + ) + with respx.mock(assert_all_called=False) as upstream: + provider: Final = upstream.post("https://sync-router.invalid/v1/completions").respond( + 200, + json={ + "id": "cmpl-unit", + "object": "text_completion", + "created": 1, + "model": "unit-text-model", + "choices": [{"text": "provider called", "index": 0, "finish_reason": "stop", "logprobs": None}], + "usage": {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2}, + }, + ) + with pytest.raises(litellm.BadRequestError, match="/v1/completions is not supported"): + router.text_completion(model=model, prompt="test") + assert provider.call_count == 0 + + +def test_routable_deployments_leave_out_models_that_refuse_the_call_type() -> None: + router: Final = Router( + model_list=[ + { + "model_name": "available", + "litellm_params": {"model": "openai/gpt-5.4-nano", "api_key": "unit-provider-key"}, + }, + { + "model_name": "embed-only", + "litellm_params": {"model": "openai/text-embedding-3-small", "api_key": "unit-provider-key"}, + "model_info": {"enforce_supported_endpoints": True, "supported_endpoints": ["/v1/embeddings"]}, + }, + ], + num_retries=0, + ) + + def routable(model: str, call_type: str | None) -> tuple[str, ...]: + return tuple( + deployment["model_name"] for deployment in router.get_routable_deployments(model, {}, call_type=call_type) + ) + + assert routable("unregistered-fallback", "acompletion") == () + assert routable("available", "acompletion") == ("available",) + assert routable("embed-only", "acompletion") == () + assert routable("embed-only", "aembedding") == ("embed-only",) + assert routable("embed-only", None) == ("embed-only",) + + def test_routing_groups_legacy_property_forwards_to_public_attribute(): router = _router_with_two_deployments([False, False]) @@ -18803,6 +18864,35 @@ async def test_anthropic_messages_stream_dropped_before_content_is_retried_withi assert [(attempted, budget) for _, attempted, budget in provider.calls] == [(0, 2), (1, 2)] +@pytest.mark.asyncio +async def test_anthropic_messages_same_group_retry_skips_an_embedding_only_deployment(): + """The retry of a /v1/messages stream dropped before content must not land on an enforced embedding-only + deployment of the same group, even when that deployment is the group's preferred order.""" + router: Final = Router( + model_list=[ + {"model_name": "glm", "litellm_params": {"model": "anthropic/glm-chat", "api_key": "sk-test", "order": 2}}, + { + "model_name": "glm", + "litellm_params": {"model": "anthropic/glm-embed", "api_key": "sk-test", "order": 1}, + "model_info": {"enforce_supported_endpoints": True, "supported_endpoints": ["/v1/embeddings"]}, + }, + ], + num_retries=1, + ) + provider: Final = _AnthropicMessagesScriptedProvider( + _anthropic_messages_dropped_before_content, _anthropic_messages_retried_stream + ) + + async def anthropic_messages(**kwargs: object) -> object: + return await provider(**kwargs) + + stream: Final = await _anthropic_messages_stream_through_router(router, anthropic_messages) + body: Final = [chunk async for chunk in stream] + + assert body == [_anthropic_messages_message_start_chunk(), _anthropic_messages_content_chunk("pong")] + assert [model for model, _, _ in provider.calls] == ["anthropic/glm-chat", "anthropic/glm-chat"] + + class _AnthropicMessagesRetryStashRecorder: def __init__(self) -> None: self.observed: tuple[Usage | None, ...] = () @@ -31431,3 +31521,26 @@ def test_router_keeps_supported_fallback_shapes( router: Final = Router(model_list=list(_FALLBACK_VALIDATION_MODEL_LIST), **router_kwargs) assert router.fallbacks == expected_fallbacks + + +def test_router_warns_when_a_deployment_enforces_endpoints_without_a_list(caplog: pytest.LogCaptureFixture) -> None: + with caplog.at_level(logging.WARNING, logger="LiteLLM Router"): + router: Final = Router( + model_list=[ + { + "model_name": "embed-from-config", + "litellm_params": {"model": "openai/text-embedding-3-small", "api_key": "sk-fake"}, + "model_info": {"enforce_supported_endpoints": True}, + } + ] + ) + router.add_deployment( + Deployment( + model_name="embed-from-api", + litellm_params=LiteLLM_Params(model="openai/text-embedding-3-small", api_key="sk-fake"), + model_info=ModelInfo(enforce_supported_endpoints=True, supported_endpoints=["/v1/embeddings", 7]), + ) + ) + + assert "Deployment 'embed-from-config' sets enforce_supported_endpoints without" in caplog.text + assert "Deployment 'embed-from-api' sets enforce_supported_endpoints without" in caplog.text diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index f114c1c73d2..6faffd7a050 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -50433,6 +50433,8 @@ export interface components { discoverable?: boolean | null; /** Enable Tag Filtering */ enable_tag_filtering?: boolean | null; + /** Enforce Supported Endpoints */ + enforce_supported_endpoints?: boolean | null; /** Explicit Cache Read Input Token Cost */ explicit_cache_read_input_token_cost?: number | null; /** Id */ From 961fabbb76b253884469c1af81121e392b8987ef Mon Sep 17 00:00:00 2001 From: Mateo Wang <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 29 Oct 2026 08:44:46 -0700 Subject: [PATCH 08/13] feat(secret_managers): resolve oidc/google api_key references per request (#47703) * test(secret_managers): cover oidc/google api_key references resolved per request Unit tests send a metadata server ID token for an admin oidc/google/ api_key, reuse it within its lifetime, refresh it before it expires, fetch it off the event loop, fail closed with a 401 that names only the kind, and honor GCE_METADATA_HOST. The integration file scripts a metadata server and checks the bearer the upstream receives for config, stored credential and DB credential sources, plus a metadata failure * feat(secret_managers): resolve oidc/google api_key references per request An admin-set oidc/google/ api_key now resolves on each call to the GCE metadata server's ID token for that audience, cached until a minute before it expires. Async callers fetch on a cache miss in a worker thread, so the event loop never waits on the metadata server. A failed fetch answers 401 naming only the kind, and GCE_METADATA_HOST points the fetch at another metadata server. * fix(oidc): bound the Google metadata fetch and answer 503 when it fails A Google OIDC api_key fetch now waits at most GCE_METADATA_TIMEOUT seconds (default 3), runs on its own thread pool instead of the loop's default executor, and sends the token read from the cache at call start even if it expires mid-call. A metadata timeout, network error, or 5xx answers a retryable 503, while a 4xx or empty token stays a 401. The expiry and empty-token warnings log the secret name on one line. * test(oidc): expect 503 when the metadata server answers 5xx in the integration test * fix(oidc): answer 503 on token-count routes during a metadata outage and fall back on a bad GCE_METADATA_TIMEOUT * fix(secret_managers): keep get_secret's metadata wait, share concurrent google token fetches, refuse oidc/google on pass-through A model's api_key fetch keeps the 3 second GCE_METADATA_TIMEOUT wait while get_secret callers keep the handler default. Any metadata transport error answers a retryable 503. Async cache misses for one reference share one fetch, and a call the reference may not go to is refused before any fetch. Pass-through routes answer 401 for an oidc/google credential instead of sending the reference text. The metadata client is closed after each fetch, and the unused sync resolver is removed. * fix(utils): refuse an oidc/google api_key on pass-through calls before any fetch A router-model or credential-name pass-through call with an oidc/google reference now fails before the metadata fetch and sends nothing upstream, since the pass-through config never sends the resolved key. The single-flight test also waits on the fetch-start signal with a bounded window instead of a fixed sleep. --- .../litellm_executed_batches.py | 8 +- .../passthrough_endpoint_router.py | 9 + litellm/proxy/proxy_server.py | 22 +- litellm/realtime_api/main.py | 4 +- litellm/router.py | 4 +- litellm/secret_managers/main.py | 135 +++++-- litellm/utils.py | 120 +++++- .../security/test_api_key_secret_reference.py | 183 ++++++++- .../test_passthrough_endpoint_router.py | 19 +- tests/unit/proxy/test_proxy_token_counter.py | 33 ++ .../test_secret_managers_main.py | 42 +- tests/unit/test_utils.py | 376 +++++++++++++++++- 12 files changed, 884 insertions(+), 71 deletions(-) diff --git a/litellm/proxy/batches_endpoints/litellm_executed_batches.py b/litellm/proxy/batches_endpoints/litellm_executed_batches.py index 098dd2f9816..67728e0fd63 100644 --- a/litellm/proxy/batches_endpoints/litellm_executed_batches.py +++ b/litellm/proxy/batches_endpoints/litellm_executed_batches.py @@ -35,7 +35,7 @@ from litellm.repositories.managed_batch_repository import ManagedBatchRepository from litellm.types.llms.base import LiteLLMBaseModel from litellm.types.llms.openai import LiteLLMBatchCreateRequest, OpenAIFileObject, OpenAIFilesPurpose from litellm.types.utils import LITELLM_EXECUTED_BATCH_PROVIDERS, ExtractedFileData, LiteLLMBatch, LlmProviders -from litellm.utils import resolve_api_key_secret_reference +from litellm.utils import aresolve_api_key_secret_reference if TYPE_CHECKING: from prisma import types as prisma_types @@ -227,7 +227,7 @@ async def upstream_lacks_files_api(api_base: str, api_key: str | None, http_clie return response.status_code == httpx.codes.NOT_FOUND -def _upstream_of(credentials: Mapping[str, object], provider: str) -> tuple[str, str | None] | None: +async def _upstream_of(credentials: Mapping[str, object], provider: str) -> tuple[str, str | None] | None: model: Final = credentials.get("model") api_base: Final = credentials.get("api_base") api_key: Final = credentials.get("api_key") @@ -244,7 +244,7 @@ def _upstream_of(credentials: Mapping[str, object], provider: str) -> tuple[str, return None if resolved_api_base is None: return None - return resolved_api_base, resolve_api_key_secret_reference( + return resolved_api_base, await aresolve_api_key_secret_reference( resolved_api_key, model=model, llm_provider=provider, upstream=credentials ) @@ -255,7 +255,7 @@ async def litellm_executed_provider_for( provider: Final = litellm_executed_provider_of(credentials) if provider is None: return None - upstream: Final = _upstream_of(credentials, provider) + upstream: Final = await _upstream_of(credentials, provider) if upstream is None: return None return provider if await lacks_files_api(*upstream) else None diff --git a/litellm/proxy/pass_through_endpoints/passthrough_endpoint_router.py b/litellm/proxy/pass_through_endpoints/passthrough_endpoint_router.py index 1d0b3d00434..739d4b3a01a 100644 --- a/litellm/proxy/pass_through_endpoints/passthrough_endpoint_router.py +++ b/litellm/proxy/pass_through_endpoints/passthrough_endpoint_router.py @@ -12,6 +12,7 @@ from litellm.integrations.vector_store_integrations.vector_store_pre_call_hook i ) from litellm.litellm_core_utils.credential_accessor import CredentialAccessor from litellm.secret_managers.main import ( + GOOGLE_OIDC_SECRET_PREFIX, UnresolvableSecretReference, get_secret_str, resolve_local_oidc_reference, @@ -47,6 +48,14 @@ def _credential_identity(credentials: VERTEX_CREDENTIALS_TYPES | None) -> str | def _resolve_oidc_reference(credential: str | None) -> str | None: if credential is None: return None + if credential.startswith(GOOGLE_OIDC_SECRET_PREFIX): + raise HTTPException( + status_code=401, + detail=( + "pass-through credential oidc/google secret reference could not be resolved: pass-through routes " + "do not fetch Google ID tokens. Use a static key, oidc/file/ or oidc/env/." + ), + ) match resolve_local_oidc_reference(credential): case UnresolvableSecretReference(kind=kind, reason=reason): raise HTTPException( diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 091e05610ba..a75ef593b7a 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -207,7 +207,7 @@ from litellm.types.utils import ( is_server_derived_pricing_key, pricing_override_fields, ) -from litellm.utils import cost_map_omits_token_price, load_credentials_from_list, with_named_credential +from litellm.utils import acredential_updates, cost_map_omits_token_price, with_named_credential if TYPE_CHECKING: from aiohttp import ClientSession @@ -14764,10 +14764,10 @@ def _provider_count_outcome( return result -def _with_resolved_api_key(litellm_params: object) -> Mapping[str, object]: +async def _with_resolved_api_key(litellm_params: object) -> Mapping[str, object]: call_params: Final = _LITELLM_PARAMS_COPY.validate_python(litellm_params) try: - load_credentials_from_list(call_params) + return {**call_params, **(await acredential_updates(call_params))} except litellm.AuthenticationError as e: raise ProxyException( message=e.message, @@ -14775,13 +14775,19 @@ def _with_resolved_api_key(litellm_params: object) -> Mapping[str, object]: param="model", code=401, ) from e - return call_params + except litellm.ServiceUnavailableError as e: + raise ProxyException( + message=e.message, + type="service_unavailable", + param="model", + code=503, + ) from e -def _with_resolved_credentials(deployment: Mapping[str, object] | None) -> dict[str, object] | None: +async def _with_resolved_credentials(deployment: Mapping[str, object] | None) -> dict[str, object] | None: if deployment is None: return None - return {**deployment, "litellm_params": _with_resolved_api_key(deployment.get("litellm_params", {}))} + return {**deployment, "litellm_params": await _with_resolved_api_key(deployment.get("litellm_params", {}))} _DEPLOYMENT_LITELLM_PARAMS: Final = TypeAdapter(Mapping[str, object]) @@ -14844,7 +14850,7 @@ async def _try_provider_token_count( model_to_use=model_to_use or "", messages=messages, contents=contents, - deployment=_with_resolved_credentials(deployment_in_region), + deployment=await _with_resolved_credentials(deployment_in_region), request_model=request_model, tools=tools, system=system, @@ -14878,7 +14884,7 @@ async def _try_native_responses_count( model_to_use=model_to_use, request=responses_request, policy=policy, - deployment=_with_resolved_credentials(deployment), + deployment=await _with_resolved_credentials(deployment), request_model=request_model, ) ) diff --git a/litellm/realtime_api/main.py b/litellm/realtime_api/main.py index 919fd54e7d6..e03cb44efda 100644 --- a/litellm/realtime_api/main.py +++ b/litellm/realtime_api/main.py @@ -35,8 +35,8 @@ from litellm.types.utils import CallTypes, LlmProviders from litellm.utils import ( ModelInfoLookupMiss, ProviderConfigManager, + aresolve_api_key_secret_reference, lookup_model_info_for_registration, - resolve_api_key_secret_reference, ) from ..litellm_core_utils.core_helpers import as_str_mapping, normalize_drop_params @@ -845,7 +845,7 @@ async def realtime_health_check( resolved_api_base: Final = cast( # cast-ok: provider parameters expose optional string endpoints str | None, api_base or resolved_params.get("api_base") ) - resolved_api_key: Final = resolve_api_key_secret_reference( + resolved_api_key: Final = await aresolve_api_key_secret_reference( cast( # cast-ok: provider parameters expose optional string credentials str | None, api_key or resolved_params.get("api_key") ), diff --git a/litellm/router.py b/litellm/router.py index 31d679f616e..7d566bd0744 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -487,6 +487,7 @@ from litellm.utils import ( ProviderConfigManager, Rules, TextCompletionResponse, + aresolve_api_key_secret_reference, function_setup, get_llm_provider, get_non_default_completion_params, @@ -497,7 +498,6 @@ from litellm.utils import ( lookup_model_info_for_registration, model_info_with_pricing_overrides, provider_rejectable_params, - resolve_api_key_secret_reference, set_live_deployment_replay, token_counter, ) @@ -12079,7 +12079,7 @@ class Router: return if api_base is None or "*" in model or params.get("use_clientside_credentials"): return - api_key: Final = resolve_api_key_secret_reference( + api_key: Final = await aresolve_api_key_secret_reference( params.api_key or dynamic_api_key, model=model, llm_provider=provider, diff --git a/litellm/secret_managers/main.py b/litellm/secret_managers/main.py index 588bdb1409c..02cf53d00ec 100644 --- a/litellm/secret_managers/main.py +++ b/litellm/secret_managers/main.py @@ -5,7 +5,7 @@ import time import traceback from collections.abc import Callable from dataclasses import dataclass -from typing import Final +from typing import Final, TypeAlias import httpx from pydantic import TypeAdapter, ValidationError @@ -13,6 +13,7 @@ from pydantic import TypeAdapter, ValidationError import litellm from litellm._logging import verbose_logger from litellm.caching.caching import DualCache +from litellm.litellm_core_utils.env_utils import get_env_float from litellm.litellm_core_utils.single_flight import single_flight_lru_cache from litellm.llms.custom_httpx.http_handler import HTTPHandler from litellm.secret_managers.get_azure_ad_token_provider import ( @@ -167,12 +168,37 @@ def get_secret_str( LOCAL_OIDC_SECRET_PREFIXES: Final = ("oidc/file/", "oidc/env/") +GOOGLE_OIDC_SECRET_PREFIX: Final = "oidc/google/" +API_KEY_OIDC_SECRET_PREFIXES: Final = (*LOCAL_OIDC_SECRET_PREFIXES, GOOGLE_OIDC_SECRET_PREFIX) +_GOOGLE_METADATA_DEFAULT_HOST: Final = "metadata.google.internal" +_GOOGLE_METADATA_DEFAULT_TIMEOUT_SECONDS: Final = 3.0 +_GOOGLE_TOKEN_UNAVAILABLE: Final = "the metadata server did not return an ID token" + + +def _google_metadata_identity_url() -> str: + host: Final = os.getenv("GCE_METADATA_HOST") or _GOOGLE_METADATA_DEFAULT_HOST + return f"http://{host}/computeMetadata/v1/instance/service-accounts/default/identity" + + +def _google_metadata_timeout() -> float: + seconds: Final = get_env_float("GCE_METADATA_TIMEOUT", _GOOGLE_METADATA_DEFAULT_TIMEOUT_SECONDS) + return seconds if seconds > 0 else _GOOGLE_METADATA_DEFAULT_TIMEOUT_SECONDS + + +class GoogleMetadataStatusError(ValueError): + def __init__(self, status_code: int) -> None: + super().__init__(f"Google OIDC provider failed with HTTP {status_code}") + self.status_code: Final = status_code @dataclass(frozen=True, slots=True) class UnresolvableSecretReference: kind: str reason: str + retryable: bool = False + + +OidcResolver: TypeAlias = Callable[[str], "str | UnresolvableSecretReference"] _OIDC_PROVIDERS: Final = frozenset({"google", "circleci", "circleci_v2", "github", "azure", "file", "env", "env_path"}) @@ -198,11 +224,86 @@ def resolve_local_oidc_reference(value: str) -> str | UnresolvableSecretReferenc return secret -def _logged_unresolvable(reference: str, kind: str, reason: str) -> UnresolvableSecretReference: +def resolve_api_key_oidc_reference(value: str) -> str | UnresolvableSecretReference: + if not value.startswith(GOOGLE_OIDC_SECRET_PREFIX): + return resolve_local_oidc_reference(value) + try: + token: Final = _google_id_token( + value, value.removeprefix(GOOGLE_OIDC_SECRET_PREFIX), timeout=_google_metadata_timeout() + ) + except (httpx.HTTPError, httpx.InvalidURL, ValueError) as exc: + return _logged_unresolvable( + value, + "oidc/google", + _GOOGLE_TOKEN_UNAVAILABLE, + f"{type(exc).__name__}: {exc}", + retryable=_metadata_server_failed_to_answer(exc), + ) + if not token: + return _logged_unresolvable(value, "oidc/google", "it resolved to an empty value") + return token + + +def _metadata_server_failed_to_answer(exc: httpx.HTTPError | httpx.InvalidURL | ValueError) -> bool: + match exc: + case httpx.TransportError(): + return True + case GoogleMetadataStatusError(status_code=status_code): + return status_code >= 500 + case _: + return False + + +def _google_id_token(reference: str, audience: str, timeout: float | None) -> str: + cached: Final = cached_oidc_token(reference) + if cached is not None: + return cached + oidc_client: Final = _get_oidc_http_handler() + try: + # https://cloud.google.com/compute/docs/instances/verifying-instance-identity#request_signature + response: Final = oidc_client.get( + _google_metadata_identity_url(), + params={"audience": audience}, + headers={"Metadata-Flavor": "Google"}, + timeout=timeout, + ) + finally: + oidc_client.close() + if response.status_code != 200: + raise GoogleMetadataStatusError(response.status_code) + token: Final = response.text + ttl: Final = _oidc_token_cache_ttl(token, 3600 - 60) + if not token: + verbose_logger.warning("Google OIDC token for %s is empty; not caching it", _single_line(reference)) + elif ttl > 0: + oidc_cache.set_cache(key=reference, value=token, ttl=ttl) + else: + verbose_logger.warning( + "Google OIDC token for %s is already expired or expires within %ss; not caching it", + _single_line(reference), + _OIDC_TOKEN_EXPIRY_MARGIN_SECONDS, + ) + return token + + +def cached_oidc_token(reference: str) -> str | None: + cached: Final = oidc_cache.get_cache(key=reference) # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] # DualCache.get_cache is untyped + return cached if isinstance(cached, str) and cached else None + + +def _logged_unresolvable( + reference: str, kind: str, reason: str, detail: str | None = None, retryable: bool = False +) -> UnresolvableSecretReference: verbose_logger.warning( - "Secret reference %s could not be resolved: %s", reference.replace("\r", "").replace("\n", ""), reason + "Secret reference %s could not be resolved: %s", + _single_line(reference), + reason if detail is None else f"{reason} ({_single_line(detail)})", ) - return UnresolvableSecretReference(kind=kind, reason=reason) + return UnresolvableSecretReference(kind=kind, reason=reason, retryable=retryable) + + +def _single_line(text: str) -> str: + return text.replace("\r", "").replace("\n", "") def _unresolvable_reason(kind: str, exc: OSError | ValueError) -> str: @@ -273,31 +374,7 @@ def get_secret( oidc_aud = "/".join(secret_name_split.split("/")[1:]) # TODO: Add caching for HTTP requests if oidc_provider == "google": - oidc_token = oidc_cache.get_cache(key=secret_name) - if oidc_token is not None: - return oidc_token - - oidc_client = _get_oidc_http_handler() - # https://cloud.google.com/compute/docs/instances/verifying-instance-identity#request_signature - response = oidc_client.get( - "http://metadata.google.internal/computeMetadata/v1/instance/service-accounts/default/identity", - params={"audience": oidc_aud}, - headers={"Metadata-Flavor": "Google"}, - ) - if response.status_code == 200: - oidc_token = response.text - ttl: Final = _oidc_token_cache_ttl(oidc_token, 3600 - 60) - if ttl > 0: - oidc_cache.set_cache(key=secret_name, value=oidc_token, ttl=ttl) - else: - verbose_logger.warning( - "Google OIDC token for %s is already expired or expires within %ss; not caching it", - secret_name, - _OIDC_TOKEN_EXPIRY_MARGIN_SECONDS, - ) - return oidc_token - else: - raise ValueError("Google OIDC provider failed") + return _google_id_token(secret_name, oidc_aud, timeout=None) elif oidc_provider == "circleci": # https://circleci.com/docs/openid-connect-tokens/ env_secret = os.getenv("CIRCLE_OIDC_TOKEN") diff --git a/litellm/utils.py b/litellm/utils.py index 5fe3e63f31a..8ea4f6ceda0 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -28,6 +28,7 @@ import textwrap import threading import time import traceback +from concurrent.futures import ThreadPoolExecutor from contextlib import nullcontext from dataclasses import dataclass, field, replace from functools import lru_cache, partial, wraps @@ -477,11 +478,14 @@ from litellm.llms.base_llm.evals.transformation import BaseEvalsAPIConfig from litellm.llms.base_llm.responses.transformation import BaseResponsesAPIConfig from litellm.llms.base_llm.skills.transformation import BaseSkillsAPIConfig from litellm.secret_managers.main import ( - LOCAL_OIDC_SECRET_PREFIXES, + API_KEY_OIDC_SECRET_PREFIXES, + GOOGLE_OIDC_SECRET_PREFIX, + OidcResolver, UnresolvableSecretReference, + cached_oidc_token, get_secret, oidc_reference_kind, - resolve_local_oidc_reference, + resolve_api_key_oidc_reference, ) from ._logging import is_debugging_on, verbose_logger @@ -784,27 +788,97 @@ def _named_credential_values(kwargs: Mapping[str, object]) -> Mapping[str, objec ) -def resolve_api_key_secret_reference( +async def aresolve_api_key_secret_reference( api_key: str | None, model: object, llm_provider: object, upstream: Mapping[str, object] ) -> str | None: if not isinstance(api_key, AdminApiKeySecretReference): return api_key - return _secret_behind_api_key_reference(api_key, model=model, llm_provider=llm_provider, upstream=upstream) + return await _resolved_off_the_event_loop( + partial(_secret_behind_api_key_reference, api_key, model, llm_provider, upstream), api_key + ) + + +_OIDC_TOKEN_FETCH_EXECUTOR: Final = ThreadPoolExecutor(thread_name_prefix="litellm-oidc-token-fetch") +_PASS_THROUGH_CALL_TYPES: Final = frozenset( + {CallTypes.allm_passthrough_route.value, CallTypes.llm_passthrough_route.value} +) +_Resolved: Final = TypeVar("_Resolved") + + +async def _resolved_off_the_event_loop(resolve: Callable[[OidcResolver], _Resolved], api_key: object) -> _Resolved: + if not isinstance(api_key, AdminApiKeySecretReference) or not api_key.startswith(GOOGLE_OIDC_SECRET_PREFIX): + return resolve(resolve_api_key_oidc_reference) + match cached_oidc_token(api_key): + case str() as token: + return resolve(lambda _reference: token) + case None: + _raise_if_the_reference_may_not_be_used(resolve) + fetched: Final = await asyncio.shield( + _google_token_fetch(asyncio.get_running_loop(), api_key).start_or_join() + ) + return resolve(lambda _reference: fetched) + + +def _raise_if_the_reference_may_not_be_used(resolve: Callable[[OidcResolver], object]) -> None: + _ = resolve(lambda _reference: "") + + +@dataclass(slots=True) +class _GoogleTokenFetch: + reference: str + in_flight: asyncio.Future[str | UnresolvableSecretReference] | None = None + + def start_or_join(self) -> asyncio.Future[str | UnresolvableSecretReference]: + if self.in_flight is not None: + return self.in_flight + caller_context: Final = contextvars.copy_context() + fetch: Final = asyncio.get_running_loop().run_in_executor( + _OIDC_TOKEN_FETCH_EXECUTOR, lambda: caller_context.run(resolve_api_key_oidc_reference, self.reference) + ) + fetch.add_done_callback(self._land) + self.in_flight = fetch + return fetch + + def _land(self, _fetch: asyncio.Future[str | UnresolvableSecretReference]) -> None: + self.in_flight = None + + +@lru_cache(maxsize=256) +def _google_token_fetch(loop: asyncio.AbstractEventLoop, reference: str) -> _GoogleTokenFetch: + return _GoogleTokenFetch(reference) def _secret_behind_api_key_reference( - reference: AdminApiKeySecretReference, model: object, llm_provider: object, upstream: Mapping[str, object] + reference: AdminApiKeySecretReference, + model: object, + llm_provider: object, + upstream: Mapping[str, object], + resolve_oidc: OidcResolver = resolve_api_key_oidc_reference, + call_type: str | None = None, ) -> str: if not reference.may_reach(upstream): raise _secret_reference_error(_unreachable_upstream_message(reference), llm_provider=llm_provider, model=model) - if not reference.startswith(LOCAL_OIDC_SECRET_PREFIXES): + if not reference.startswith(API_KEY_OIDC_SECRET_PREFIXES): raise _secret_reference_error( f"api_key is an {oidc_reference_kind(reference)} secret reference, a kind api_key does not support. " - "Use oidc/file/ or oidc/env/.", + "Use oidc/file/, oidc/env/ or oidc/google/.", llm_provider=llm_provider, model=model, ) - match resolve_local_oidc_reference(reference): + if call_type in _PASS_THROUGH_CALL_TYPES and reference.startswith(GOOGLE_OIDC_SECRET_PREFIX): + raise _secret_reference_error( + "api_key oidc/google secret reference could not be resolved: pass-through routes do not fetch Google ID " + "tokens. Use a static key, oidc/file/ or oidc/env/.", + llm_provider=llm_provider, + model=model, + ) + match resolve_oidc(reference): + case UnresolvableSecretReference(kind=kind, reason=reason, retryable=True): + raise ServiceUnavailableError( + message=f"api_key {kind} secret reference could not be resolved: {reason}", + llm_provider=llm_provider, + model=model, + ) case UnresolvableSecretReference(kind=kind, reason=reason): raise _secret_reference_error( f"api_key {kind} secret reference could not be resolved: {reason}", @@ -836,21 +910,39 @@ def with_named_credential(kwargs: Mapping[str, object]) -> Mapping[str, object]: return {**kwargs, **_named_credential_values(kwargs)} -def load_credentials_from_list(kwargs: dict[str, object]) -> None: +def load_credentials_from_list(kwargs: dict[str, object], call_type: str | None = None) -> None: """ Updates kwargs with the credentials if credential_name in kwarg, then swaps an admin-configured oidc/ api_key reference for the secret it points to, read now so a rotated token is used on the next call """ - kwargs.update(_credential_updates(kwargs)) + kwargs.update(_credential_updates(kwargs, _named_credential_values(kwargs), call_type=call_type)) -def _credential_updates(kwargs: Mapping[str, object]) -> Mapping[str, object]: +async def acredential_updates(kwargs: Mapping[str, object], call_type: str | None = None) -> Mapping[str, object]: credential_values: Final = _named_credential_values(kwargs) - api_key: Final = kwargs.get("api_key", credential_values.get("api_key")) + return await _resolved_off_the_event_loop( + partial(_credential_updates, kwargs, credential_values, call_type=call_type), + _api_key_after_credentials(kwargs, credential_values), + ) + + +def _api_key_after_credentials(kwargs: Mapping[str, object], credential_values: Mapping[str, object]) -> object: + return kwargs.get("api_key", credential_values.get("api_key")) + + +def _credential_updates( + kwargs: Mapping[str, object], + credential_values: Mapping[str, object], + resolve_oidc: OidcResolver = resolve_api_key_oidc_reference, + call_type: str | None = None, +) -> Mapping[str, object]: + api_key: Final = _api_key_after_credentials(kwargs, credential_values) if not isinstance(api_key, AdminApiKeySecretReference): return credential_values secret: Final = _secret_behind_api_key_reference( api_key, + resolve_oidc=resolve_oidc, + call_type=call_type, model=kwargs.get("model"), llm_provider=kwargs.get("custom_llm_provider"), upstream=MappingProxyType( @@ -1936,7 +2028,7 @@ def client(original_function): assert logging_obj is not None, "logging_obj should not be None after function_setup" ## LOAD CREDENTIALS - load_credentials_from_list(kwargs) + load_credentials_from_list(kwargs, call_type) from litellm.llms.github_copilot.per_user_auth import attach_github_copilot_user_session attach_github_copilot_user_session(kwargs) # pyright: ignore[reportUnknownArgumentType] # kwargs is the untyped request dict @@ -2230,7 +2322,7 @@ def client(original_function): kwargs["litellm_logging_obj"] = logging_obj ## LOAD CREDENTIALS - load_credentials_from_list(kwargs) + kwargs.update(await acredential_updates(kwargs, call_type)) # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType] # kwargs is the untyped request dict from litellm.llms.github_copilot.per_user_auth import aattach_github_copilot_user_session await aattach_github_copilot_user_session(kwargs) # pyright: ignore[reportUnknownArgumentType] # kwargs is the untyped request dict diff --git a/tests/integration/security/test_api_key_secret_reference.py b/tests/integration/security/test_api_key_secret_reference.py index 0b1187d2e0e..86a7d58e0d3 100644 --- a/tests/integration/security/test_api_key_secret_reference.py +++ b/tests/integration/security/test_api_key_secret_reference.py @@ -1,6 +1,9 @@ from __future__ import annotations +import base64 import json +import threading +import time import uuid from collections.abc import Callable, Iterator, Mapping, Sequence from dataclasses import dataclass @@ -46,7 +49,18 @@ _SEARCH_PATH: Final = f"/v1/vector_stores/{_VECTOR_STORE}/search" _BATCH_DEPLOYMENT: Final = "wif-vllm-batch" _BEDROCK_MODEL: Final = "bedrock/converse/us.anthropic.claude-sonnet-5-5" _PROVIDER_INPUT_TOKENS: Final = 7919 -_UNSUPPORTED_REFERENCE: Final = "oidc/google/https://example.invalid" +_UNSUPPORTED_REFERENCE: Final = "oidc/github/https://example.invalid" +_GOOGLE_MODEL: Final = "hosted_vllm/wif-google" +_GOOGLE_AUDIENCES: Final = MappingProxyType( + { + "config_inline": "https://vllm-inline.example.test", + "config_credential": "https://vllm-credential.example.test", + "db_credential": "https://vllm-db-credential.example.test", + } +) +_UNAVAILABLE_AUDIENCE: Final = "https://vllm-unavailable.example.test" +_HUNG_AUDIENCE: Final = "https://vllm-hung.example.test" +_IDENTITY_PATH: Final = "/computeMetadata/v1/instance/service-accounts/default/identity" _RUN_SECRETS_TOKEN: Final = "/run/secrets/wif-absent-token" _SOURCES: Final = ("config_credential", "config_inline", "db_credential", "db_inline") _UNREADABLE: Final = MappingProxyType( @@ -108,6 +122,45 @@ def _reference(path: Path) -> str: return f"oidc/file/{path}" +def _google_reference(audience: str) -> str: + return f"oidc/google/{audience}" + + +def _jwt_segment(part: Mapping[str, JsonValue]) -> str: + return base64.urlsafe_b64encode(json.dumps(part).encode()).decode().rstrip("=") + + +@dataclass(frozen=True, slots=True) +class _Google: + expires: int + + def token(self, audience: str) -> str: + claims: Final = _jwt_segment({"iss": "https://accounts.google.com", "aud": audience, "exp": self.expires}) + return f"{_jwt_segment({'alg': 'RS256', 'typ': 'JWT'})}.{claims}.signature" + + def bearers(self) -> frozenset[str]: + return frozenset(f"Bearer {self.token(audience)}" for audience in _GOOGLE_AUDIENCES.values()) + + +def _metadata_server(google: _Google) -> Callable[[Request], Reply]: + def respond(request: Request) -> Reply: + target: Final = httpx.URL(request.target) + audience: Final = target.params.get("audience") + if audience is None or (request.method, target.path, request.headers.get("metadata-flavor")) != ( + "GET", + _IDENTITY_PATH, + "Google", + ): + return Reply(status=404, body=b"not found", content_type="text/plain") + if audience == _UNAVAILABLE_AUDIENCE: + return Reply(status=503, body=b"metadata server unavailable", content_type="text/plain") + if audience == _HUNG_AUDIENCE: + return Reply(chunks=(b"never sent",), gate_before_first=threading.Event(), content_type="text/plain") + return Reply(body=google.token(audience).encode(), content_type="text/plain") + + return respond + + def _secrets(directory: Path) -> _Secrets: allowed: Final = directory / "secrets" outside_dir: Final = directory / "outside" @@ -169,9 +222,9 @@ def _served(request: Request) -> Reply: return openai_error(404) -def _provider(secrets: _Secrets) -> Callable[[Request], Reply]: +def _provider(secrets: _Secrets, google: _Google) -> Callable[[Request], Reply]: def respond(request: Request) -> Reply: - if request.headers.get("authorization") != f"Bearer {secrets.current()}": + if request.headers.get("authorization") not in {f"Bearer {secrets.current()}", *google.bearers()}: return openai_error(401) return _served(request) @@ -224,6 +277,27 @@ def _config(directory: Path, secrets: _Secrets, provider: str) -> Path: "aws_bedrock_runtime_endpoint": provider.removesuffix("/v1"), }, ), + _deployment( + "wif-google-inline", + { + "model": _GOOGLE_MODEL, + "api_key": _google_reference(_GOOGLE_AUDIENCES["config_inline"]), + "api_base": provider, + }, + ), + _deployment("wif-google-credential", _named(_GOOGLE_MODEL, "wif-google")), + _deployment( + "wif-google-unavailable", + {"model": _GOOGLE_MODEL, "api_key": _google_reference(_UNAVAILABLE_AUDIENCE), "api_base": provider}, + ), + _deployment( + "wif-google-unavailable-count", + {"model": _CHAT_MODEL, "api_key": _google_reference(_UNAVAILABLE_AUDIENCE), "api_base": provider}, + ), + _deployment( + "wif-google-hung", + {"model": _GOOGLE_MODEL, "api_key": _google_reference(_HUNG_AUDIENCE), "api_base": provider}, + ), _deployment( "wif-mistral-passthrough", { @@ -239,6 +313,7 @@ def _config(directory: Path, secrets: _Secrets, provider: str) -> Path: _credential("wif-outside", _reference(secrets.outside), provider), _credential("wif-unsupported", _UNSUPPORTED_REFERENCE, provider), _credential("wif-run-secrets", f"oidc/file/{_RUN_SECRETS_TOKEN}", provider), + _credential("wif-google", _google_reference(_GOOGLE_AUDIENCES["config_credential"]), provider), ], "vector_store_registry": [ { @@ -312,6 +387,8 @@ class _Rig: owned: OwnedProxy provider: Recorder intruder: Recorder + metadata: Recorder + google: _Google provider_base: str intruder_base: str secrets: _Secrets @@ -339,17 +416,24 @@ def _is_discovery(request: Request) -> bool: def rig(tmp_path_factory: pytest.TempPathFactory) -> Iterator[_Rig]: directory: Final = tmp_path_factory.mktemp("api_key_secret_reference").resolve() secrets: Final = _secrets(directory) + google: Final = _Google(expires=int(time.time()) + 3600) with ( gateway_from_environment() as gateway, - wire_server(_provider(secrets)) as provider_wire, + wire_server(_provider(secrets, google)) as provider_wire, wire_server(_intruder) as intruder_wire, + wire_server(_metadata_server(google)) as metadata_wire, ): provider: Final = Recorder(provider_wire) provider_base: Final = f"{provider_wire.url}/v1" with owned_proxy_process( gateway, directory, - {**secrets.overrides(), "MISTRAL_API_BASE": provider_wire.url}, + { + **secrets.overrides(), + "MISTRAL_API_BASE": provider_wire.url, + "GCE_METADATA_HOST": httpx.URL(metadata_wire.url).netloc.decode(), + "GCE_METADATA_TIMEOUT": "1", + }, config=_config(directory, secrets, provider_base), workers=2, healthcheck=WorkerHealthcheck.LENIENT, @@ -361,6 +445,8 @@ def rig(tmp_path_factory: pytest.TempPathFactory) -> Iterator[_Rig]: owned=owned, provider=provider, intruder=Recorder(intruder_wire), + metadata=Recorder(metadata_wire), + google=google, provider_base=provider_base, intruder_base=f"{intruder_wire.url}/v1", secrets=secrets, @@ -708,7 +794,92 @@ def test_team_row_naming_the_credential_sends_the_token_only_to_its_host(rig: _R def test_startup_discovery_sends_the_boot_token(rig: _Rig) -> None: sent: Final = frozenset(_bearers(tuple(request for request in rig.boot if _is_discovery(request)))) - assert sent == {f"Bearer {rig.boot_token}"}, sent + assert sent - rig.google.bearers() == {f"Bearer {rig.boot_token}"}, sent + + +def _discovery_bearers(rig: _Rig) -> frozenset[str | None]: + return frozenset(_bearers(tuple(request for request in rig.provider.requests() if _is_discovery(request)))) + + +def test_startup_discovery_sends_the_google_token(rig: _Rig) -> None: + expected: Final = frozenset( + f"Bearer {rig.google.token(_GOOGLE_AUDIENCES[source])}" for source in ("config_inline", "config_credential") + ) + sent: Final = eventually( + partial(_discovery_bearers, rig), lambda seen: expected <= seen, seconds=30, return_last_on_timeout=True + ) + assert expected <= sent, sent + + +def _google_model(rig: _Rig, scenario: Scenario, source: str) -> str: + match source: + case "config_inline": + return "wif-google-inline" + case "config_credential": + return "wif-google-credential" + case "db_credential": + credential: Final = f"db-google-{uuid.uuid4().hex}" + rig.owned.gateway.post( + "/credentials", + _credential(credential, _google_reference(_GOOGLE_AUDIENCES[source]), rig.provider_base), + ) + scenario.cleanups.callback(_delete_credential, rig.owned.gateway, credential) + name, identity = _created_model( + rig.owned.gateway, + scenario, + _deployment(f"db-google-{uuid.uuid4().hex}", _named(_GOOGLE_MODEL, credential)), + ) + eventually(partial(rig.owned.serves_only_on_every_worker, name, identity), bool, seconds=120) + return name + case _: + pytest.fail(f"unknown source {source}") + + +def _asked_audiences(rig: _Rig) -> frozenset[str | None]: + return frozenset(httpx.URL(request.target).params.get("audience") for request in rig.metadata.requests()) + + +@pytest.mark.parametrize("source", tuple(_GOOGLE_AUDIENCES)) +def test_google_reference_sends_the_metadata_server_token_for_its_audience(rig: _Rig, source: str) -> None: + audience: Final = _GOOGLE_AUDIENCES[source] + with rig.owned.gateway.scenario() as scenario: + marker: Final = _marker() + response: Final = _chat(rig.owned.gateway, _google_model(rig, scenario, source), marker) + sent: Final = _bearers(_carrying(rig.provider.requests(), marker)) + assert (response.status_code, sent) == (200, (f"Bearer {rig.google.token(audience)}",)), response.text + assert audience in _asked_audiences(rig) + + +def test_google_reference_fails_closed_when_the_metadata_server_fails(rig: _Rig) -> None: + marker: Final = _marker() + response: Final = _chat(rig.owned.gateway, "wif-google-unavailable", marker) + reached: Final = _carrying(rig.provider.requests(), marker) + assert (response.status_code, reached) == (503, ()), (response.text, _bearers(reached)) + assert "oidc/google secret reference could not be resolved" in response.text, response.text + assert _UNAVAILABLE_AUDIENCE.removeprefix("https://") not in response.text, response.text + assert _UNAVAILABLE_AUDIENCE in _asked_audiences(rig) + + +def test_google_reference_answers_503_when_the_metadata_server_never_finishes_its_reply(rig: _Rig) -> None: + marker: Final = _marker() + response: Final = _chat(rig.owned.gateway, "wif-google-hung", marker) + reached: Final = _carrying(rig.provider.requests(), marker) + assert (response.status_code, reached) == (503, ()), (response.text, _bearers(reached)) + assert "the metadata server did not return an ID token" in response.text, response.text + assert _HUNG_AUDIENCE in _asked_audiences(rig) + + +def test_input_token_count_answers_503_when_the_metadata_server_fails(rig: _Rig) -> None: + marker: Final = _marker() + response: Final = rig.owned.gateway.request( + "POST", + "/v1/responses/input_tokens", + {"model": "wif-google-unavailable-count", "input": f"say {marker}"}, + headers=_CLOSE, + ) + reached: Final = _carrying(rig.provider.requests(), marker) + assert (response.status_code, reached) == (503, ()), (response.text, _bearers(reached)) + assert "oidc/google secret reference could not be resolved" in response.text, response.text def _side_call(rig: _Rig, seam: str, marker: str) -> tuple[httpx.Response, Callable[[Request], bool]]: diff --git a/tests/unit/proxy/pass_through_endpoints/test_passthrough_endpoint_router.py b/tests/unit/proxy/pass_through_endpoints/test_passthrough_endpoint_router.py index 9c50a6d4df6..ab4f3f7a675 100644 --- a/tests/unit/proxy/pass_through_endpoints/test_passthrough_endpoint_router.py +++ b/tests/unit/proxy/pass_through_endpoints/test_passthrough_endpoint_router.py @@ -485,7 +485,7 @@ def test_oidc_env_path_reference_is_not_resolved(tmp_path, monkeypatch): ) -def test_network_backed_oidc_reference_is_not_fetched_inline(monkeypatch): +def test_oidc_google_reference_is_refused_without_a_metadata_fetch(monkeypatch): def fail(*args, **kwargs): raise AssertionError("network-backed oidc reference was resolved inside the request path") @@ -494,11 +494,24 @@ def test_network_backed_oidc_reference_is_not_fetched_inline(monkeypatch): lambda name, *a, **k: "oidc/google/https://example.com" if name == "OPENAI_API_KEY" else fail(), ) - assert ( + with pytest.raises(HTTPException) as excinfo: _passthrough_router(None).get_credentials(custom_llm_provider="openai", region_name=None) - == "oidc/google/https://example.com" + + assert excinfo.value.status_code == 401 + assert "pass-through credential oidc/google" in str(excinfo.value.detail) + assert "example.com" not in str(excinfo.value.detail) + + +def test_deployment_oidc_google_reference_is_not_sent_upstream(monkeypatch): + llm_router = litellm.Router( + model_list=[_flagged_deployment("openai/gpt-4o", api_key="oidc/google/https://vllm.example.test")] ) + with pytest.raises(HTTPException) as excinfo: + _passthrough_router(llm_router).get_credentials(custom_llm_provider="openai", region_name=None) + + assert excinfo.value.status_code == 401 + @pytest.fixture() async def _drain_logging_worker(): diff --git a/tests/unit/proxy/test_proxy_token_counter.py b/tests/unit/proxy/test_proxy_token_counter.py index b21730c64f6..d0cf7be794a 100644 --- a/tests/unit/proxy/test_proxy_token_counter.py +++ b/tests/unit/proxy/test_proxy_token_counter.py @@ -279,6 +279,39 @@ async def test_token_counter_reads_an_oidc_api_key_only_when_it_calls_the_provid assert str(tmp_path) not in excinfo.value.message +@pytest.mark.asyncio +async def test_token_counter_answers_503_when_the_metadata_server_behind_an_oidc_google_api_key_fails( + monkeypatch: pytest.MonkeyPatch, +): + audience: Final = f"https://count.test/{uuid.uuid4().hex}" + monkeypatch.setattr( + litellm.proxy.proxy_server, + "llm_router", + Router( + model_list=[ + { + "model_name": "claude", + "litellm_params": {"model": "anthropic/claude-sonnet-5-5", "api_key": f"oidc/google/{audience}"}, + } + ] + ), + ) + + with respx.mock(assert_all_called=True) as upstream: + upstream.get(url__regex=r"/computeMetadata/v1/instance/service-accounts/default/identity").respond( + 500, text="metadata server unavailable" + ) + with pytest.raises(ProxyException) as excinfo: + await token_counter( + request=TokenCountRequest(model="claude", messages=[{"role": "user", "content": "hello"}]), + call_endpoint=True, + ) + + assert excinfo.value.code == "503" + assert "oidc/google secret reference could not be resolved" in excinfo.value.message + assert audience.removeprefix("https://") not in excinfo.value.message + + @pytest.mark.asyncio async def test_native_openai_responses_count_sends_the_token_an_oidc_api_key_points_to( tmp_path: Path, monkeypatch: pytest.MonkeyPatch diff --git a/tests/unit/secret_managers/test_secret_managers_main.py b/tests/unit/secret_managers/test_secret_managers_main.py index 862658d9169..a2e5b498f34 100644 --- a/tests/unit/secret_managers/test_secret_managers_main.py +++ b/tests/unit/secret_managers/test_secret_managers_main.py @@ -4,6 +4,7 @@ import logging import os import threading import time +import uuid from collections.abc import Callable from concurrent.futures import ThreadPoolExecutor from queue import SimpleQueue @@ -11,6 +12,7 @@ from typing import Final from unittest.mock import Mock, patch import pytest +import respx import litellm from litellm._logging import redact_internal_details_from_client_message @@ -37,7 +39,7 @@ class MockHTTPHandler: self.text = "mocked_token" self.json_data = {"value": "mocked_token"} - def get(self, url, params=None, headers=None): + def get(self, url, params=None, headers=None, timeout=None): # Store params for audience verification self.last_params = params logger.debug( @@ -49,6 +51,9 @@ class MockHTTPHandler: mock_response.json.return_value = self.json_data return mock_response + def close(self): + pass + @pytest.fixture def mock_oidc_cache(): @@ -229,6 +234,41 @@ def test_oidc_google_failure(): get_secret(secret_name) + +def test_oidc_google_fetches_the_token_from_the_metadata_host_gce_metadata_host_names( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setenv("GCE_METADATA_HOST", "127.0.0.1:8930") + audience: Final = f"https://vllm.example.test/{uuid.uuid4().hex}" + + with respx.mock(assert_all_called=True) as metadata: + metadata.get( + "http://127.0.0.1:8930/computeMetadata/v1/instance/service-accounts/default/identity", + params={"audience": audience}, + headers={"Metadata-Flavor": "Google"}, + ).respond(200, text="metadata-host-token") + token: Final = get_secret(f"oidc/google/{audience}") + + assert token == "metadata-host-token" + +def test_oidc_google_secret_outside_a_model_api_key_keeps_its_metadata_wait_when_gce_metadata_timeout_is_set( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setenv("GCE_METADATA_TIMEOUT", "1") + audience: Final = f"https://sts.example.test/{uuid.uuid4().hex}" + + with respx.mock(assert_all_called=True) as metadata: + identity: Final = metadata.get( + "http://metadata.google.internal/computeMetadata/v1/instance/service-accounts/default/identity", + params={"audience": audience}, + ).respond(200, text="web-identity-token") + token: Final = get_secret(f"oidc/google/{audience}") + waits: Final = identity.calls.last.request.extensions["timeout"] + + assert token == "web-identity-token" + assert waits == {"connect": 5.0, "read": 600.0, "write": 600.0, "pool": 600.0} + + def test_oidc_circleci_success(monkeypatch): monkeypatch.setenv("CIRCLE_OIDC_TOKEN", "circleci_token") diff --git a/tests/unit/test_utils.py b/tests/unit/test_utils.py index 2870016a5e9..cd56b6e1b46 100644 --- a/tests/unit/test_utils.py +++ b/tests/unit/test_utils.py @@ -5,14 +5,18 @@ import copy import contextlib import contextvars import io +import itertools import json import logging import os import queue import threading +import time +import uuid from collections.abc import Callable, Iterator, Mapping from concurrent.futures import Future, ThreadPoolExecutor from datetime import datetime, timedelta, timezone +from functools import partial from pathlib import Path, PurePath from typing import Final, SupportsIndex, cast from unittest import mock @@ -79,8 +83,10 @@ from litellm.types.utils import ( all_litellm_params, bedrock_batch_litellm_params, ) +from litellm.secret_managers.main import oidc_cache from litellm.types.videos.main import VideoObject from litellm.utils import ( + acredential_updates, _invalidate_model_cost_lowercase_map, add_custom_logger_callback_to_specific_event, check_valid_key, @@ -7320,7 +7326,6 @@ def test_unresolvable_oidc_api_key_fails_closed_naming_the_reason_without_its_ta ("reference", "kind"), [ ("oidc/env_path/MISTRAL_SA_TOKEN_PATH", "oidc/env_path"), - ("oidc/google/https://api.mistral.ai", "oidc/google"), ("oidc/github/mistral", "oidc/github"), ("oidc/env_MISTRAL_SA_TOKEN_PATH", "oidc/"), ("oidc/LIT-UNKNOWN-KIND/MISTRAL_SA_TOKEN_PATH", "oidc/"), @@ -7341,6 +7346,7 @@ def test_oidc_api_key_of_a_kind_api_key_cannot_resolve_fails_closed_as_an_authen assert not route.called assert excinfo.value.status_code == 401 assert f"{kind} secret reference" in excinfo.value.message + assert "oidc/google/" in excinfo.value.message assert reference.removeprefix(kind).lstrip("/") not in excinfo.value.message assert "MISTRAL_SA_TOKEN_PATH" not in excinfo.value.message @@ -7356,7 +7362,7 @@ async def test_router_fails_over_from_a_deployment_whose_oidc_api_key_kind_canno "model_name": "devstral", "litellm_params": { "model": "mistral/devstral-latest", - "api_key": "oidc/google/https://api.mistral.ai", + "api_key": "oidc/github/mistral", "order": 1, }, }, @@ -7395,6 +7401,372 @@ def test_an_oidc_api_key_sent_with_the_call_goes_upstream_as_written( assert sent == f"Bearer {reference}" +_GOOGLE_METADATA_IDENTITY: Final = ( + "http://metadata.google.internal/computeMetadata/v1/instance/service-accounts/default/identity" +) +_VLLM_BASE: Final = "https://vllm.example.test/v1" +_VLLM_CHAT_URL: Final = f"{_VLLM_BASE}/chat/completions" +_VLLM_MODEL: Final = "hosted_vllm/qwen3-0.6b" +_VLLM_REPLY: Final = {**_MISTRAL_REPLY, "model": "qwen3-0.6b"} +_CLOCK_START: Final = 1_900_000_000.0 +_GOOGLE_TOKEN_LIFETIME: Final = 3600 + + +def _jwt_segment(claims: Mapping[str, object]) -> str: + return base64.urlsafe_b64encode(json.dumps(claims).encode()).decode().rstrip("=") + + +def _google_id_token(audience: str, issued_at: float) -> str: + claims: Final = {"aud": audience, "iat": int(issued_at), "exp": int(issued_at) + _GOOGLE_TOKEN_LIFETIME} + return f"{_jwt_segment({'alg': 'RS256', 'typ': 'JWT'})}.{_jwt_segment(claims)}.signature" + + +def _google_reference() -> tuple[str, str]: + audience: Final = f"https://vllm.example.test/{uuid.uuid4().hex}" + return audience, f"oidc/google/{audience}" + + +def _metadata_identity(fetched_on: queue.SimpleQueue[int], request: httpx.Request) -> httpx.Response: + fetched_on.put(threading.get_ident()) + if request.headers.get("Metadata-Flavor") != "Google": + return httpx.Response(403) + return httpx.Response(200, text=_google_id_token(request.url.params["audience"], time.time())) + + +def _vllm_credential(api_key: str) -> CredentialItem: + return CredentialItem( + credential_name="vllm-sa", credential_values={"api_key": api_key, "api_base": _VLLM_BASE}, credential_info={} + ) + + +async def _ping_vllm_through(call_path: str, credential_name: str) -> None: + if call_path == "async": + await litellm.acompletion(model=_VLLM_MODEL, messages=_PING, litellm_credential_name=credential_name) + return + litellm.completion(model=_VLLM_MODEL, messages=_PING, litellm_credential_name=credential_name) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("call_path", ["sync", "async"]) +async def test_named_credential_oidc_google_api_key_sends_the_metadata_token_and_refreshes_it_before_it_expires( + monkeypatch: pytest.MonkeyPatch, call_path: str +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + monkeypatch.setattr(time, "time", lambda: _CLOCK_START) + audience, reference = _google_reference() + monkeypatch.setattr(litellm, "credential_list", [_vllm_credential(reference)]) + refreshed_at: Final = _CLOCK_START + _GOOGLE_TOKEN_LIFETIME - 59 + + with respx.mock(assert_all_called=True) as upstream: + metadata: Final = upstream.get(_GOOGLE_METADATA_IDENTITY).mock( + side_effect=partial(_metadata_identity, queue.SimpleQueue()) + ) + chat: Final = upstream.post(_VLLM_CHAT_URL).respond(200, json=_VLLM_REPLY) + await _ping_vllm_through(call_path, "vllm-sa") + await _ping_vllm_through(call_path, "vllm-sa") + fetches_within_lifetime: Final = metadata.call_count + monkeypatch.setattr(time, "time", lambda: refreshed_at) + await _ping_vllm_through(call_path, "vllm-sa") + sent: Final = [call.request.headers["Authorization"] for call in chat.calls] + audiences: Final = [call.request.url.params["audience"] for call in metadata.calls] + + first_token: Final = f"Bearer {_google_id_token(audience, _CLOCK_START)}" + assert sent == [first_token, first_token, f"Bearer {_google_id_token(audience, refreshed_at)}"] + assert (fetches_within_lifetime, audiences) == (1, [audience, audience]) + assert litellm.credential_list[0].credential_values["api_key"] == reference + + +@pytest.mark.asyncio +async def test_router_fetches_an_oidc_google_api_key_token_off_the_event_loop_and_keeps_the_reference( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + monkeypatch.setattr(time, "time", lambda: _CLOCK_START) + audience, reference = _google_reference() + fetched_on: Final[queue.SimpleQueue[int]] = queue.SimpleQueue() + router: Final = litellm.Router( + model_list=[ + { + "model_name": "vllm", + "litellm_params": {"model": _VLLM_MODEL, "api_base": _VLLM_BASE, "api_key": reference}, + } + ] + ) + + try: + with respx.mock(assert_all_called=True) as upstream: + upstream.get(_GOOGLE_METADATA_IDENTITY).mock(side_effect=partial(_metadata_identity, fetched_on)) + chat: Final = upstream.post(_VLLM_CHAT_URL).respond(200, json=_VLLM_REPLY) + await router.acompletion(model="vllm", messages=_PING) + sent: Final = chat.calls.last.request.headers["Authorization"] + finally: + router.discard() + + assert sent == f"Bearer {_google_id_token(audience, _CLOCK_START)}" + assert fetched_on.qsize() == 1 + assert fetched_on.get_nowait() != threading.get_ident() + assert router.get_model_list(model_name="vllm")[0]["litellm_params"]["api_key"] == reference + + +@pytest.mark.asyncio +@pytest.mark.parametrize("call_path", ["sync", "async"]) +@pytest.mark.parametrize( + ("metadata_failure", "status_code"), + [ + (lambda: httpx.Response(500, text="metadata server unavailable"), 503), + (lambda: httpx.Response(403, text="Missing Metadata-Flavor:Google header"), 401), + (lambda: httpx.Response(200, text=""), 401), + (lambda: httpx.ConnectError("[Errno -2] Name or service not known"), 503), + (lambda: httpx.ReadTimeout("timed out"), 503), + (lambda: httpx.RemoteProtocolError("Server disconnected without sending a response."), 503), + ], + ids=["server error", "refused", "empty token", "not on GCP", "timeout", "disconnected"], +) +async def test_oidc_google_api_key_fails_closed_when_the_metadata_server_fails_and_recovers_on_the_next_call( + monkeypatch: pytest.MonkeyPatch, + caplog: pytest.LogCaptureFixture, + call_path: str, + metadata_failure: Callable[[], httpx.Response | httpx.HTTPError], + status_code: int, +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + monkeypatch.setattr(time, "time", lambda: _CLOCK_START) + audience, reference = _google_reference() + monkeypatch.setattr(litellm, "credential_list", [_vllm_credential(reference)]) + recovered_token: Final = _google_id_token(audience, _CLOCK_START) + + with respx.mock(assert_all_called=True) as upstream: + upstream.get(_GOOGLE_METADATA_IDENTITY).mock( + side_effect=[metadata_failure(), httpx.Response(200, text=recovered_token)] + ) + chat: Final = upstream.post(_VLLM_CHAT_URL).respond(200, json=_VLLM_REPLY) + with caplog.at_level(logging.WARNING, logger="LiteLLM"), pytest.raises( + (litellm.AuthenticationError, litellm.ServiceUnavailableError) + ) as excinfo: + await _ping_vllm_through(call_path, "vllm-sa") + sent_while_failing: Final = chat.call_count + await _ping_vllm_through(call_path, "vllm-sa") + sent_after_recovery: Final = [call.request.headers["Authorization"] for call in chat.calls] + + assert excinfo.value.status_code == status_code + assert "api_key oidc/google secret reference could not be resolved" in excinfo.value.message + assert audience.removeprefix("https://") not in excinfo.value.message + assert reference in caplog.text + assert sent_while_failing == 0 + assert sent_after_recovery == [f"Bearer {recovered_token}"] + + +@pytest.mark.asyncio +async def test_an_oidc_google_api_key_token_cached_when_the_call_starts_is_sent_though_it_expires_mid_call( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(time, "time", lambda: _CLOCK_START) + audience, reference = _google_reference() + cached_token: Final = _google_id_token(audience, _CLOCK_START) + oidc_cache.set_cache(key=reference, value=cached_token, ttl=60) + clock: Final = itertools.chain([_CLOCK_START + 60], itertools.repeat(_CLOCK_START + 61)) + monkeypatch.setattr(time, "time", lambda: next(clock)) + monkeypatch.setattr(litellm, "credential_list", [_vllm_credential(reference)]) + fetched_on: Final[queue.SimpleQueue[int]] = queue.SimpleQueue() + + with respx.mock(assert_all_called=False) as upstream: + upstream.get(_GOOGLE_METADATA_IDENTITY).mock(side_effect=partial(_metadata_identity, fetched_on)) + updates: Final = await acredential_updates({"model": _VLLM_MODEL, "litellm_credential_name": "vllm-sa"}) + + assert updates["api_key"] == cached_token + assert fetched_on.empty() + + +@pytest.mark.asyncio +async def test_a_pending_oidc_google_metadata_fetch_leaves_the_default_executor_free( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + _, reference = _google_reference() + monkeypatch.setattr(litellm, "credential_list", [_vllm_credential(reference)]) + loop: Final = asyncio.get_running_loop() + loop.set_default_executor(ThreadPoolExecutor(max_workers=1)) + fetch_started: Final[queue.SimpleQueue[int]] = queue.SimpleQueue() + release: Final = threading.Event() + + def pending_metadata(_request: httpx.Request) -> httpx.Response: + fetch_started.put(threading.get_ident()) + release.wait(timeout=30) + return httpx.Response(500, text="metadata server unavailable") + + with ThreadPoolExecutor(max_workers=1) as observer, respx.mock(assert_all_called=False) as upstream: + upstream.get(_GOOGLE_METADATA_IDENTITY).mock(side_effect=pending_metadata) + chat: Final = upstream.post(_VLLM_CHAT_URL).respond(200, json=_VLLM_REPLY) + call: Final = asyncio.ensure_future(_ping_vllm_through("async", "vllm-sa")) + try: + await asyncio.wait_for(loop.run_in_executor(observer, fetch_started.get), timeout=30) + other_work: Final = await asyncio.wait_for(loop.run_in_executor(None, lambda: "done"), timeout=30) + finally: + release.set() + with pytest.raises(litellm.ServiceUnavailableError) as excinfo: + await call + + assert other_work == "done" + assert excinfo.value.status_code == 503 + assert chat.call_count == 0 + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "metadata_answer", + [lambda token: httpx.Response(200, text=token), lambda _token: httpx.Response(500, text="metadata server down")], + ids=["token", "outage"], +) +async def test_concurrent_calls_that_miss_the_oidc_google_token_cache_share_one_metadata_fetch( + monkeypatch: pytest.MonkeyPatch, metadata_answer: Callable[[str], httpx.Response] +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + monkeypatch.setattr(time, "time", lambda: _CLOCK_START) + audience, reference = _google_reference() + monkeypatch.setattr(litellm, "credential_list", [_vllm_credential(reference)]) + token: Final = _google_id_token(audience, _CLOCK_START) + fetch_started: Final[queue.SimpleQueue[int]] = queue.SimpleQueue() + release: Final = threading.Event() + + def held_metadata(_request: httpx.Request) -> httpx.Response: + fetch_started.put(threading.get_ident()) + release.wait(timeout=30) + return metadata_answer(token) + + with ThreadPoolExecutor(max_workers=1) as observer, respx.mock(assert_all_called=False) as upstream: + metadata: Final = upstream.get(_GOOGLE_METADATA_IDENTITY).mock(side_effect=held_metadata) + chat: Final = upstream.post(_VLLM_CHAT_URL).respond(200, json=_VLLM_REPLY) + calls: Final = [asyncio.ensure_future(_ping_vllm_through("async", "vllm-sa")) for _ in range(24)] + loop: Final = asyncio.get_running_loop() + try: + await asyncio.wait_for(loop.run_in_executor(observer, fetch_started.get), timeout=30) + fetch_started_while_held: Final = await loop.run_in_executor( + observer, partial(_next_fetch_within, fetch_started, 0.5) + ) + finally: + release.set() + outcomes: Final = await asyncio.gather(*calls, return_exceptions=True) + sent: Final = [call.request.headers["Authorization"] for call in chat.calls] + + assert fetch_started_while_held is None + if metadata.calls[0].response.status_code == 200: + assert metadata.call_count == 1 + assert outcomes == [None] * 24 + assert sent == [f"Bearer {token}"] * 24 + else: + assert [getattr(outcome, "status_code", None) for outcome in outcomes] == [503] * 24 + assert sent == [] + + +def _next_fetch_within(fetches: queue.SimpleQueue[int], seconds: float) -> int | None: + try: + return fetches.get(timeout=seconds) + except queue.Empty: + return None + + +@pytest.mark.asyncio +async def test_a_call_an_oidc_google_api_key_may_not_go_to_is_refused_without_a_metadata_fetch( + monkeypatch: pytest.MonkeyPatch, +) -> None: + _, reference = _google_reference() + monkeypatch.setattr(litellm, "credential_list", [_vllm_credential(reference)]) + + with respx.mock(assert_all_called=False) as upstream: + metadata: Final = upstream.get(_GOOGLE_METADATA_IDENTITY).respond(500) + with pytest.raises(litellm.AuthenticationError) as excinfo: + await acredential_updates( + { + "model": _VLLM_MODEL, + "litellm_credential_name": "vllm-sa", + "api_base": "https://elsewhere.example.test/v1", + } + ) + + assert "only to that credential's own api_base" in excinfo.value.message + assert metadata.call_count == 0 + + +async def _pass_through_chat_call(call_path: str, reference: str) -> None: + pass_through: Final = {"method": "POST", "endpoint": "v1/chat/completions", "json": {"messages": _PING}} + if call_path == "sync": + litellm.llm_passthrough_route(model=_VLLM_MODEL, litellm_credential_name="vllm-sa", **pass_through) + return + router: Final = litellm.Router( + model_list=[ + { + "model_name": "vllm-model", + "litellm_params": {"model": _VLLM_MODEL, "api_base": _VLLM_BASE, "api_key": reference}, + } + ], + num_retries=0, + ) + await router.allm_passthrough_route(model="vllm-model", **pass_through) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("call_path", ["sync", "async router"]) +async def test_a_pass_through_call_with_an_oidc_google_api_key_is_refused_without_a_metadata_fetch( + monkeypatch: pytest.MonkeyPatch, call_path: str +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + monkeypatch.setattr(time, "time", lambda: _CLOCK_START) + audience, reference = _google_reference() + monkeypatch.setattr(litellm, "credential_list", [_vllm_credential(reference)]) + + with respx.mock(assert_all_called=False) as upstream: + metadata: Final = upstream.get(_GOOGLE_METADATA_IDENTITY).respond( + 200, text=_google_id_token(audience, _CLOCK_START) + ) + vllm: Final = upstream.route(host="vllm.example.test").respond(200, json=_VLLM_REPLY) + with pytest.raises(litellm.AuthenticationError) as excinfo: + await _pass_through_chat_call(call_path, reference) + + assert "pass-through routes do not fetch Google ID tokens" in excinfo.value.message + assert metadata.call_count == 0 + assert vllm.call_count == 0 + + +@pytest.mark.parametrize( + ("configured", "bound"), + [(None, 3.0), ("1.5", 1.5), ("three seconds", 3.0), ("0", 3.0), ("-2", 3.0)], + ids=["unset", "configured", "malformed", "zero", "negative"], +) +def test_oidc_google_api_key_metadata_fetch_waits_at_most_gce_metadata_timeout_seconds( + monkeypatch: pytest.MonkeyPatch, configured: str | None, bound: float +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + if configured is None: + monkeypatch.delenv("GCE_METADATA_TIMEOUT", raising=False) + else: + monkeypatch.setenv("GCE_METADATA_TIMEOUT", configured) + monkeypatch.setattr(time, "time", lambda: _CLOCK_START) + audience, reference = _google_reference() + monkeypatch.setattr(litellm, "credential_list", [_vllm_credential(reference)]) + token: Final = _google_id_token(audience, _CLOCK_START) + + with respx.mock(assert_all_called=True) as upstream: + metadata: Final = upstream.get(_GOOGLE_METADATA_IDENTITY).respond(200, text=token) + chat: Final = upstream.post(_VLLM_CHAT_URL).respond(200, json=_VLLM_REPLY) + litellm.completion(model=_VLLM_MODEL, messages=_PING, litellm_credential_name="vllm-sa") + sent: Final = chat.calls.last.request.headers["Authorization"] + waits: Final = metadata.calls.last.request.extensions["timeout"] + + assert sent == f"Bearer {token}" + assert waits == {"connect": bound, "read": bound, "write": bound, "pool": bound} + + +def test_an_oidc_google_api_key_sent_with_the_call_goes_upstream_as_written_without_a_metadata_fetch() -> None: + _, reference = _google_reference() + + with respx.mock(assert_all_called=True) as upstream: + chat: Final = upstream.post(_VLLM_CHAT_URL).respond(200, json=_VLLM_REPLY) + litellm.completion(model=_VLLM_MODEL, api_base=_VLLM_BASE, messages=_PING, api_key=reference) + sent: Final = chat.calls.last.request.headers["Authorization"] + + assert sent == f"Bearer {reference}" + + @pytest.mark.asyncio @pytest.mark.parametrize( ("model_info", "expected"), From 06cdee91daf61c47bdf7cef1497a6bb001f23831 Mon Sep 17 00:00:00 2001 From: paveln83 <349044532+paveln83@users.noreply.github.com> Date: Thu, 29 Oct 2026 08:46:49 -0700 Subject: [PATCH 09/13] fix(guardrails): scan and rewrite streamed /v1/completions under post_call guardrails (#49035) * feat(guardrails): scan and rewrite streamed /v1/completions in post_call pipelines Gives OpenAITextCompletionHandler a process_output_streaming_response that assembles the buffered TextCompletionResponse chunks per choice index at end of stream, runs the guardrail over the assembled text, and writes text rewrites back into the first text-carrying chunk of each choice. Tests and live proof still pending. * fix(guardrails): scan and rewrite streamed /v1/completions under post_call guardrails Assemble each choice's streamed text once the stream ends, run the guardrail on it, write a mask back into the buffered chunks, and end the stream with the guardrail's error on a block. User key metadata now goes through the shared metadata bucket so applied guardrails reach spend logs. * test(guardrails): cover streamed /v1/completions under post_call guardrails * fix(guardrails): keep Gray Swan key metadata, scan a cut-short text stream, and pass non-chunk frames through * fix(guardrails): replace blocked text completion streams like chat, keep Gray Swan's vendor payload, scan text streams in the content filter * fix(guardrails): release an unfinished choice's held text and omit usage the upstream never sent The content filter's streaming hook held back the end of a choice until its finish_reason, which litellm sends only for the first choice of an n>1 stream, so choice 1 lost its last characters on text and chat streams. A closing chunk now releases every unfinished choice when the stream ends. A blocked text stream's closing chunk reported 0/0 usage when the upstream sent none; it now carries usage only when the upstream reported it. * fix(guardrails): send no usage on the content filter's end-of-stream chunk * fix(guardrails): keep the usage chunk last and log the masked text of streamed text completions * fix(guardrails): stream Responses API events through the content filter as they arrive * fix(guardrails): treat a null text completion choice index as choice zero when assembling streams --------- Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> --- .../litellm_core_utils/served_output_texts.py | 60 +++-- .../guardrail_translation/handler.py | 198 +++++++++++++++- litellm/proxy/common_request_processing.py | 20 +- .../guardrail_hooks/grayswan/grayswan.py | 36 ++- .../litellm_content_filter/content_filter.py | 194 ++++++++++----- ...ext_completion_stream_guardrail_rebuild.py | 95 +++++++- .../test_served_output_texts.py | 35 ++- .../test_text_completion_guardrail_handler.py | 223 +++++++++++++++++- .../content_filter/test_content_filter.py | 190 ++++++++++++++- .../guardrail_hooks/test_grayswan.py | 28 +++ .../test_deferred_guardrail_logging.py | 21 ++ .../proxy_logging/test_guardrail_pipeline.py | 156 +++++++++++- 12 files changed, 1132 insertions(+), 124 deletions(-) diff --git a/litellm/litellm_core_utils/served_output_texts.py b/litellm/litellm_core_utils/served_output_texts.py index 854785f304b..b19973e06c3 100644 --- a/litellm/litellm_core_utils/served_output_texts.py +++ b/litellm/litellm_core_utils/served_output_texts.py @@ -4,12 +4,13 @@ post-call guardrail rewrote rather than the provider response the proxy assemble from __future__ import annotations from collections.abc import Sequence +from itertools import chain from typing import Final, Literal -from pydantic import BaseModel, Field, TypeAdapter, ValidationError +from pydantic import BaseModel, ConfigDict, Field, TypeAdapter, ValidationError from litellm.types.llms.base import LiteLLMBaseModel -from litellm.types.utils import ModelResponse, ModelResponseStream +from litellm.types.utils import ModelResponse, ModelResponseStream, TextCompletionResponse SERVED_OUTPUT_TEXTS_KEY: Final = "served_output_texts" @@ -44,6 +45,19 @@ class _ChatChoices(LiteLLMBaseModel): choices: list[object] +class _TextCompletionChoice(LiteLLMBaseModel): + model_config = ConfigDict(frozen=True) + + index: int | None = None + text: str | None = None + + +class _TextCompletionChunk(LiteLLMBaseModel): + model_config = ConfigDict(frozen=True) + + choices: tuple[_TextCompletionChoice, ...] + + def _as_json_object(response: object) -> dict[str, object] | None: candidate: Final = response.model_dump() if isinstance(response, BaseModel) else response try: @@ -108,8 +122,13 @@ def served_output_texts(response: object) -> ServedTexts | None: def served_stream_output_texts(chunks: Sequence[object]) -> ServedTexts | None: - if chunks and all(isinstance(chunk, ModelResponseStream) for chunk in chunks): - return _chat_stream_texts(tuple(chunk for chunk in chunks if isinstance(chunk, ModelResponseStream))) + chat_chunks: Final = tuple(chunk for chunk in chunks if isinstance(chunk, ModelResponseStream)) + if chunks and len(chat_chunks) == len(chunks): + return _indexed_texts_by_choice(tuple(chain.from_iterable(_chat_chunk_texts(chunk) for chunk in chat_chunks))) + if chunks and all(isinstance(chunk, TextCompletionResponse) for chunk in chunks): + return _indexed_texts_by_choice( + tuple(chain.from_iterable(_text_completion_chunk_texts(chunk) for chunk in chunks)) + ) from litellm.llms.anthropic.messages.stream_assembly import ( assemble_anthropic_sse_stream, is_anthropic_sse_stream, @@ -121,19 +140,26 @@ def served_stream_output_texts(chunks: Sequence[object]) -> ServedTexts | None: return _chat_texts(assembled) if assembled is not None else None -def _chat_stream_choice_text(chunks: Sequence[ModelResponseStream], index: int) -> str | None: - contents: Final = tuple( - content - for chunk in chunks +def _chat_chunk_texts(chunk: ModelResponseStream) -> tuple[tuple[int, str | None], ...]: + return tuple( + (choice.index or 0, content if isinstance(content := choice.delta.content, str) else None) for choice in chunk.choices - if (choice.index or 0) == index and isinstance(content := choice.delta.content, str) ) - return "".join(contents) if contents else None -def _chat_stream_texts(chunks: Sequence[ModelResponseStream]) -> ServedTexts | None: - choice_count: Final = max(((choice.index or 0) + 1 for chunk in chunks for choice in chunk.choices), default=0) - texts: Final = tuple(_chat_stream_choice_text(chunks, index) for index in range(choice_count)) +def _text_completion_chunk_texts(chunk: object) -> tuple[tuple[int, str | None], ...]: + parsed: Final = _TextCompletionChunk.model_validate(_as_json_object(chunk)) + return tuple((choice.index or 0, choice.text) for choice in parsed.choices) + + +def _joined_choice_text(indexed_texts: Sequence[tuple[int, str | None]], index: int) -> str | None: + texts: Final = tuple(text for text_index, text in indexed_texts if text_index == index and text is not None) + return "".join(texts) if texts else None + + +def _indexed_texts_by_choice(indexed_texts: Sequence[tuple[int, str | None]]) -> ServedTexts | None: + choice_count: Final = max((index + 1 for index, _ in indexed_texts), default=0) + texts: Final = tuple(_joined_choice_text(indexed_texts, index) for index in range(choice_count)) return texts if any(text is not None for text in texts) else None @@ -170,6 +196,8 @@ def _choice_with_text(choice: object, text: str | None) -> object: if choice_obj is None or text is None: return choice message: Final = _as_json_object(choice_obj.get("message")) - if message is None or message.get("content") == text: - return choice - return {**choice_obj, "message": {**message, "content": text}} + if message is not None: + return choice if message.get("content") == text else {**choice_obj, "message": {**message, "content": text}} + if isinstance(choice_obj.get("text"), str): + return {**choice_obj, "text": text} + return choice diff --git a/litellm/llms/openai/completion/guardrail_translation/handler.py b/litellm/llms/openai/completion/guardrail_translation/handler.py index c4bc04d4627..27c82cfc004 100644 --- a/litellm/llms/openai/completion/guardrail_translation/handler.py +++ b/litellm/llms/openai/completion/guardrail_translation/handler.py @@ -5,18 +5,30 @@ This module provides guardrail translation support for OpenAI's text completion The handler processes the 'prompt' parameter for guardrails. """ -from collections.abc import Mapping +import json +from collections.abc import Mapping, Sequence +from itertools import chain from typing import TYPE_CHECKING, Final from litellm._logging import verbose_proxy_logger -from litellm.llms.base_llm.guardrail_translation.base_translation import BaseTranslation -from litellm.types.utils import GenericGuardrailAPIInputs +from litellm.llms.base_llm.guardrail_translation.base_translation import ( + BaseTranslation, + StreamingScanKey, + StreamTransformSink, +) +from litellm.llms.base_llm.guardrail_translation.utils import ( + blocked_chat_stream_usage, + stream_item_field, + stream_item_items, +) +from litellm.types.utils import GenericGuardrailAPIInputs, TextChoices, TextCompletionResponse, Usage if TYPE_CHECKING: - from litellm.integrations.custom_guardrail import CustomGuardrail + from fastapi import HTTPException + + from litellm.integrations.custom_guardrail import CustomGuardrail, ModifyResponseException from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.types.proxy.auth.user_api_key_auth import UserAPIKeyAuth - from litellm.types.utils import TextCompletionResponse class OpenAITextCompletionHandler(BaseTranslation): @@ -32,6 +44,9 @@ class OpenAITextCompletionHandler(BaseTranslation): - A list of strings (for batch completions) """ + delivers_ended_stream_rewrites = True + assembles_streamed_response = True + async def process_input_messages( self, data: dict, @@ -125,7 +140,7 @@ class OpenAITextCompletionHandler(BaseTranslation): guardrail_to_apply: "CustomGuardrail", litellm_logging_obj: "LiteLLMLoggingObj | None" = None, user_api_key_dict: "UserAPIKeyAuth | None" = None, - request_data: dict | None = None, + request_data: dict[str, object] | None = None, ) -> "TextCompletionResponse": """ Process output response by applying guardrails to completion text. @@ -162,11 +177,7 @@ class OpenAITextCompletionHandler(BaseTranslation): if "response" not in request_data: request_data["response"] = response - # Add user API key metadata with prefixed keys - if "litellm_metadata" not in request_data: - user_metadata: Final = self.transform_user_api_key_dict_to_metadata(user_api_key_dict) - if user_metadata: - request_data["litellm_metadata"] = user_metadata + self.merge_user_api_key_metadata_into_request(request_data, user_api_key_dict) inputs: Final = GenericGuardrailAPIInputs(texts=texts_to_check) # Include model information from the response if available @@ -195,3 +206,168 @@ class OpenAITextCompletionHandler(BaseTranslation): ) return response + + async def process_output_streaming_response( + self, + responses_so_far: list[object], # mutable-ok: rewrites the caller's buffered chunks in place + guardrail_to_apply: "CustomGuardrail", + litellm_logging_obj: "LiteLLMLoggingObj | None" = None, + user_api_key_dict: "UserAPIKeyAuth | None" = None, + request_data: dict[str, object] | None = None, # mutable-ok: the hooks share and mutate the request payload + stream_transform_sink: StreamTransformSink | None = None, + deliver_ended_stream_rewrites: bool = False, + ) -> Sequence[object]: + chunks: Final = _text_completion_chunks(responses_so_far) + if not self._check_streaming_has_ended(chunks): + return responses_so_far + assembled: Final = self._assemble_ended_stream(chunks) + pre_guardrail_texts: Final = _choice_texts(assembled) + await self.process_output_response( + response=assembled, + guardrail_to_apply=guardrail_to_apply, + litellm_logging_obj=litellm_logging_obj, + user_api_key_dict=user_api_key_dict, + request_data=request_data, + ) + if not deliver_ended_stream_rewrites: + return responses_so_far + rewrites: Final = tuple( + after for before, after in zip(pre_guardrail_texts, _choice_texts(assembled)) if after != before + ) + for choice_index, text in rewrites: + self._write_choice_text(chunks, choice_index, text) + return responses_so_far + + def get_streaming_scan_key(self, responses_so_far: Sequence[object]) -> StreamingScanKey | None: + chunks: Final = _text_completion_chunks(responses_so_far) + if not self._check_streaming_has_ended(chunks): + return None + assembled: Final = self._assemble_ended_stream(chunks) + return StreamingScanKey(texts=tuple(text for _, text in _choice_texts(assembled)), stream_ended=True) + + def released_stream_as_ended(self, responses_so_far: Sequence[object]) -> tuple[object, ...]: + chunks: Final = _text_completion_chunks(responses_so_far) + streamed_choices: Final = tuple(chain.from_iterable(chunk.choices for chunk in chunks)) + finished_indices: Final = frozenset( + _choice_index(choice) for choice in streamed_choices if choice.finish_reason + ) + unfinished_indices: Final = tuple( + dict.fromkeys( + _choice_index(choice) for choice in streamed_choices if _choice_index(choice) not in finished_indices + ) + ) + if not unfinished_indices: + return tuple(responses_so_far) + terminator: Final = TextCompletionResponse( + id=chunks[0].id, + model=chunks[0].model, + object="text_completion", + choices=[TextChoices(index=index, text="", finish_reason="length") for index in unfinished_indices], + ) + return (*responses_so_far, terminator) + + def build_block_sse_chunks( + self, + exc: "ModifyResponseException", + stream_started: bool = False, + responses_so_far: Sequence[object] | None = None, + ) -> Sequence[bytes]: + streamed: Final = _text_completion_chunks(responses_so_far or ()) + source: Final = streamed[0] if streamed else TextCompletionResponse(model=exc.model) + message_chunk: Final = TextCompletionResponse( + id=source.id, + created=source.created, + model=source.model or exc.model, + object="text_completion", + choices=[TextChoices(index=0, text=exc.message)], + ).model_copy(update={"usage": None}) + final_chunk: Final = TextCompletionResponse( + id=source.id, + created=source.created, + model=source.model or exc.model, + object="text_completion", + choices=[TextChoices(index=0, text="", finish_reason="content_filter")], + ).model_copy(update={"usage": _blocked_stream_usage(exc.original_response)}) + return tuple( + f"data: {chunk.model_dump_json(exclude_none=True)}\n\n".encode() for chunk in (message_chunk, final_chunk) + ) + + def build_stream_error_items( + self, + exc: "HTTPException", + responses_so_far: Sequence[object] | None = None, + ) -> Sequence[bytes] | None: + from litellm.proxy.common_request_processing import sse_error_payload + + _, error_obj = sse_error_payload(exc) + return (f'data: {{"error": {json.dumps(error_obj)}}}\n\n'.encode(),) + + def _check_streaming_has_ended(self, responses_so_far: Sequence[object]) -> bool: + streamed_choices: Final = chain.from_iterable(stream_item_items(item, "choices") for item in responses_so_far) + return any(stream_item_field(choice, "finish_reason") is not None for choice in streamed_choices) + + @staticmethod + def _assemble_ended_stream(responses_so_far: Sequence[TextCompletionResponse]) -> TextCompletionResponse: + streamed_choices: Final = tuple(chain.from_iterable(item.choices for item in responses_so_far)) + choice_indices: Final = tuple(dict.fromkeys(_choice_index(choice) for choice in streamed_choices)) + assembled_choices: Final = [ + TextChoices( + index=index, + text="".join(_choice_text(choice) for choice in streamed_choices if _choice_index(choice) == index), + finish_reason=next( + ( + choice.finish_reason + for choice in streamed_choices + if _choice_index(choice) == index and choice.finish_reason + ), + None, + ), + ) + for index in choice_indices + ] + first: Final = responses_so_far[0] + return TextCompletionResponse( + id=first.id, + created=first.created, + model=first.model, + object="text_completion", + choices=assembled_choices, + usage=next((item.usage for item in reversed(responses_so_far) if item.usage is not None), None), + ) + + @staticmethod + def _write_choice_text(responses_so_far: Sequence[TextCompletionResponse], choice_index: int, text: str) -> None: + streamed_choices: Final = chain.from_iterable(item.choices for item in responses_so_far) + head, *tail = (choice for choice in streamed_choices if _choice_index(choice) == choice_index) + head.text = text + for choice in tail: + if isinstance(stream_item_field(choice, "text"), str): + choice.text = "" + + +def _blocked_stream_usage(original_response: object) -> Usage | None: + prompt_tokens, completion_tokens = blocked_chat_stream_usage(original_response) + if isinstance(original_response, list) and prompt_tokens == completion_tokens == 0: + return None + return Usage( + prompt_tokens=prompt_tokens, + completion_tokens=completion_tokens, + total_tokens=prompt_tokens + completion_tokens, + ) + + +def _text_completion_chunks(responses_so_far: Sequence[object]) -> tuple[TextCompletionResponse, ...]: + return tuple(item for item in responses_so_far if isinstance(item, TextCompletionResponse)) + + +def _choice_text(choice: TextChoices) -> str: + text: Final = stream_item_field(choice, "text") + return text if isinstance(text, str) else "" + + +def _choice_index(choice: TextChoices) -> int: + return choice.index or 0 + + +def _choice_texts(response: TextCompletionResponse) -> tuple[tuple[int, str], ...]: + return tuple((_choice_index(choice), _choice_text(choice)) for choice in response.choices) diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index 4884e75af4e..7e487ded569 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -4059,13 +4059,13 @@ class ProxyBaseLLMRequestProcessing: (assembled_response, cache_hit); the closure also runs non-apply_guardrail post-call hooks via _run_deferred_stream_guardrails. - - Every other anthropic_messages/aresponses stream gets a closure - that dispatches on the stored args shape, because the arming site - cannot tell the producers apart: native iterators store a single - ready-made logging coroutine to enqueue, while bridged streams - (LiteLLMCompletionStreamingIterator, and the plain SSE generator - AnthropicStreamWrapper returns for bridged /v1/messages) share - their inner CustomStreamWrapper's logging_obj and so store + - Every other anthropic_messages/aresponses/atext_completion stream + gets a closure that dispatches on the stored args shape, because the + arming site cannot tell the producers apart: native iterators store + a single ready-made logging coroutine to enqueue, while bridged + streams (LiteLLMCompletionStreamingIterator, TextCompletionStreamWrapper, + and the plain SSE generator AnthropicStreamWrapper returns for bridged + /v1/messages) share their inner CustomStreamWrapper's logging_obj and so store (assembled_response, cache_hit); for those the closure only dispatches success logging, matching the route's pre-existing hook surface. @@ -4102,7 +4102,11 @@ class ProxyBaseLLMRequestProcessing: logging_obj._on_deferred_stream_complete = _on_deferred_stream_complete return - if route_type not in ("anthropic_messages", "aresponses") or not self._is_streaming_response(response): + if route_type not in ( + "anthropic_messages", + "aresponses", + "atext_completion", + ) or not self._is_streaming_response(response): return from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER diff --git a/litellm/proxy/guardrails/guardrail_hooks/grayswan/grayswan.py b/litellm/proxy/guardrails/guardrail_hooks/grayswan/grayswan.py index 7ff5aa2466b..e57d6e78378 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/grayswan/grayswan.py +++ b/litellm/proxy/guardrails/guardrail_hooks/grayswan/grayswan.py @@ -82,6 +82,22 @@ def _as_plain_dict(item: object) -> Mapping[str, object]: return cast("Mapping[str, object]", item) # cast-ok: wire rows are message/tool-call dicts +_STR_KEYED_MAPPING: Final = TypeAdapter(Mapping[str, object]) + + +def _str_keyed(value: object) -> Mapping[str, object]: + return _STR_KEYED_MAPPING.validate_python(value) if isinstance(value, Mapping) else {} + + +def _request_route(request_data: Mapping[str, object]) -> str | None: + litellm_metadata: Final = request_data.get("litellm_metadata") + metadata_bucket: Final = _str_keyed( + litellm_metadata if litellm_metadata is not None else request_data.get("metadata") + ) + route: Final = metadata_bucket.get("user_api_key_request_route") + return route if isinstance(route, str) else None + + class GraySwanGuardrailMissingSecrets(Exception): """Raised when the Gray Swan API key is missing.""" @@ -625,11 +641,8 @@ class GraySwanGuardrail(CustomGuardrail): from litellm.litellm_core_utils.api_route_to_call_types import get_call_types_for_route from litellm.llms import load_guardrail_translation_mappings - litellm_metadata: Final = request_data.get("litellm_metadata") - request_route: Final = ( - litellm_metadata.get("user_api_key_request_route") if isinstance(litellm_metadata, Mapping) else None - ) - route_call_types: Final = get_call_types_for_route(request_route) if isinstance(request_route, str) else None + request_route: Final = _request_route(request_data) + route_call_types: Final = get_call_types_for_route(request_route) if request_route is not None else None call_type: Final = ( (route_call_types[0].value if route_call_types else None) or (logging_obj.call_type if logging_obj is not None else None) @@ -698,13 +711,12 @@ class GraySwanGuardrail(CustomGuardrail): inbound_headers: Final = self._extract_inbound_headers(request_data, logging_obj) - litellm_metadata: Final = request_data.get("litellm_metadata") - cleaned_litellm_metadata: Final = dict(litellm_metadata) if isinstance(litellm_metadata, dict) else {} - if inbound_headers: - existing_headers: Final = cleaned_litellm_metadata.get("headers") - cleaned_litellm_metadata["headers"] = ( - {**existing_headers, **inbound_headers} if isinstance(existing_headers, dict) else inbound_headers - ) + litellm_metadata: Final = _str_keyed(request_data.get("litellm_metadata")) + cleaned_litellm_metadata: Final = ( + {**litellm_metadata, "headers": {**_str_keyed(litellm_metadata.get("headers")), **inbound_headers}} + if inbound_headers + else litellm_metadata + ) if cleaned_litellm_metadata: sanitized: Final = _DECODED_JSON.validate_python( safe_json_loads(safe_dumps(cleaned_litellm_metadata), default={}) diff --git a/litellm/proxy/guardrails/guardrail_hooks/litellm_content_filter/content_filter.py b/litellm/proxy/guardrails/guardrail_hooks/litellm_content_filter/content_filter.py index f58c8705815..71b8976bebf 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/litellm_content_filter/content_filter.py +++ b/litellm/proxy/guardrails/guardrail_hooks/litellm_content_filter/content_filter.py @@ -12,11 +12,11 @@ import json import os import re import time -from collections.abc import AsyncGenerator, Coroutine, Mapping, Sequence +from collections.abc import AsyncGenerator, AsyncIterable, Callable, Coroutine, Mapping, Sequence from dataclasses import dataclass, replace from datetime import datetime from re import Pattern -from typing import TYPE_CHECKING, Any, Final, Literal, Optional, TypedDict, cast +from typing import TYPE_CHECKING, Any, Final, Literal, Optional, TypedDict, TypeVar, cast import yaml from fastapi import HTTPException @@ -40,6 +40,7 @@ from litellm.proxy.guardrails.content_filter_data import ( ) from litellm.types.utils import ( CallTypes, + Delta, Function, GenericGuardrailAPIInputs, GuardrailStatus, @@ -47,6 +48,9 @@ from litellm.types.utils import ( Message, ModelResponse, ModelResponseStream, + StreamingChoices, + TextChoices, + TextCompletionResponse, ) if TYPE_CHECKING: @@ -78,6 +82,7 @@ from .patterns import PATTERN_EXTRA_CONFIG, get_compiled_pattern MAX_KEYWORD_VALUE_GAP_WORDS: Final = 1 _DECODED_ARGUMENTS: Final = TypeAdapter(object) +_StreamItemT: Final = TypeVar("_StreamItemT") _JSON_ENTRIES: Final = TypeAdapter(list[object]) GAP_WORD_TOKENIZER: Final = re.compile(r"\b\w+\b") SENTENCE_TERMINATORS: Final = re.compile(r"[.!?]+") @@ -139,6 +144,26 @@ class _CategoryConfigView(TypedDict): category_file: str | None +def _streamed_choice_content(choice: StreamingChoices | TextChoices) -> str | None: + if isinstance(choice, TextChoices): + return choice.text if isinstance(choice.text, str) else "" + delta: Final = getattr(choice, "delta", None) + if not delta: + return None + content: Final = getattr(delta, "content", None) + return content if isinstance(content, str) else "" + + +def _streamed_choices(item: object) -> Sequence[StreamingChoices | TextChoices]: + if isinstance(item, (ModelResponseStream, TextCompletionResponse)): + return item.choices + return () + + +def _is_choiceless_chunk(item: object) -> bool: + return isinstance(item, (ModelResponseStream, TextCompletionResponse)) and not item.choices + + @dataclass(frozen=True, slots=True) class _StreamedChoiceState: buffered_text: str = "" @@ -146,6 +171,52 @@ class _StreamedChoiceState: committed_detections: tuple[ContentFilterDetection, ...] = () latest_detections: tuple[ContentFilterDetection, ...] = () next_trim_len: int = 0 + finished: bool = False + + +def _end_of_stream_chunk( + last_item: object, state_by_choice: Mapping[int, _StreamedChoiceState] +) -> ModelResponseStream | TextCompletionResponse | None: + unfinished: Final = tuple(index for index, state in state_by_choice.items() if not state.finished) + if not unfinished: + return None + if isinstance(last_item, TextCompletionResponse): + return TextCompletionResponse( + id=last_item.id, + created=last_item.created, + model=last_item.model, + object="text_completion", + choices=[TextChoices(index=index, text="") for index in unfinished], + ).model_copy(update={"usage": None}) + if isinstance(last_item, ModelResponseStream): + return ModelResponseStream( + id=last_item.id, + created=last_item.created, + model=last_item.model, + choices=[StreamingChoices(index=index, delta=Delta(content="")) for index in unfinished], + ) + return None + + +async def _stream_then_end_of_stream_chunk( + response: AsyncIterable[_StreamItemT], end_of_stream_chunk: Callable[[_StreamItemT], _StreamItemT | None] +) -> AsyncGenerator[tuple[_StreamItemT, bool], None]: + last_item: _StreamItemT | None = None # rebind-ok: the end-of-stream chunk copies the last chunk's id and model + held_items: tuple[_StreamItemT, ...] = () # rebind-ok: the usage chunk must stay after the end-of-stream chunk + async for item in response: + last_item = item + if _is_choiceless_chunk(item): + held_items = (*held_items, item) + continue + for held_item in held_items: + yield held_item, False + held_items = () + yield item, False + closing_item: Final = None if last_item is None else end_of_stream_chunk(last_item) + if closing_item is not None: + yield closing_item, True + for held_item in held_items: + yield held_item, False @dataclass(frozen=True, slots=True) @@ -2104,72 +2175,75 @@ class ContentFilterGuardrail(CustomGuardrail): ) try: - async for item in response: - if isinstance(item, ModelResponseStream) and item.choices: - for choice in item.choices: - if not (hasattr(choice, "delta") and choice.delta): - continue + async for item, stream_ended in _stream_then_end_of_stream_chunk( + response, lambda last_item: _end_of_stream_chunk(last_item, state_by_choice) + ): + for choice in _streamed_choices(item): + new_content = _streamed_choice_content(choice) + if new_content is None: + continue - choice_index = getattr(choice, "index", 0) - if not isinstance(choice_index, int): - choice_index = 0 + choice_index = getattr(choice, "index", 0) + if not isinstance(choice_index, int): + choice_index = 0 - content = getattr(choice.delta, "content", None) - is_final = bool(getattr(choice, "finish_reason", None)) - new_content = content if isinstance(content, str) else "" - if not new_content and not is_final: - continue + is_final = stream_ended or bool(getattr(choice, "finish_reason", None)) + if not new_content and not is_final: + continue - previous_state = state_by_choice.get(choice_index, _StreamedChoiceState()) - buffered_text = previous_state.buffered_text + new_content - if not buffered_text: - continue + previous_state = state_by_choice.get(choice_index, _StreamedChoiceState()) + buffered_text = previous_state.buffered_text + new_content + if not buffered_text: + continue - # Add a space at the end if it's the final chunk to trigger word boundaries (\b) - text_to_scan = buffered_text + (" " if is_final else "") - choice_detections: list[ContentFilterDetection] = [] - scan_started = time.perf_counter() + # Add a space at the end if it's the final chunk to trigger word boundaries (\b) + text_to_scan = buffered_text + (" " if is_final else "") + choice_detections: list[ContentFilterDetection] = [] + scan_started = time.perf_counter() - try: - masked_text = self._filter_single_text(text_to_scan, detections=choice_detections) - if is_final and masked_text.endswith(" "): - masked_text = masked_text[:-1] - latest_detections = tuple(choice_detections) - except HTTPException: - state_by_choice[choice_index] = replace( - previous_state, latest_detections=tuple(choice_detections) - ) - raise - except Exception as e: - verbose_proxy_logger.error("ContentFilterGuardrail: Error in masking: %s", e) - masked_text = text_to_scan # Fallback to current text - latest_detections = previous_state.latest_detections - finally: - scan_seconds += time.perf_counter() - scan_started - - safe_to_yield_len = max( - previous_state.yielded_masked_text_len, - len(masked_text) - (0 if is_final else CONTENT_FILTER_STREAMING_HOLDBACK_CHARS), + try: + masked_text = self._filter_single_text(text_to_scan, detections=choice_detections) + if is_final and masked_text.endswith(" "): + masked_text = masked_text[:-1] + latest_detections = tuple(choice_detections) + except HTTPException: + state_by_choice[choice_index] = replace( + previous_state, latest_detections=tuple(choice_detections) ) - choice.delta.content = masked_text[previous_state.yielded_masked_text_len : safe_to_yield_len] - next_state = replace( - previous_state, - buffered_text=buffered_text, - yielded_masked_text_len=safe_to_yield_len, - latest_detections=latest_detections, - ) - if is_final: - state_by_choice[choice_index] = next_state - continue + raise + except Exception as e: + verbose_proxy_logger.error("ContentFilterGuardrail: Error in masking: %s", e) + masked_text = text_to_scan # Fallback to current text + latest_detections = previous_state.latest_detections + finally: + scan_seconds += time.perf_counter() - scan_started - trim_started = time.perf_counter() - state_by_choice[choice_index] = self._trim_streamed_choice_buffer(next_state, masked_text, plan) - scan_seconds += time.perf_counter() - trim_started + safe_to_yield_len = max( + previous_state.yielded_masked_text_len, + len(masked_text) - (0 if is_final else CONTENT_FILTER_STREAMING_HOLDBACK_CHARS), + ) + released_text = masked_text[previous_state.yielded_masked_text_len : safe_to_yield_len] + if isinstance(choice, TextChoices): + choice.text = released_text + else: + choice.delta.content = released_text + next_state = replace( + previous_state, + buffered_text=buffered_text, + yielded_masked_text_len=safe_to_yield_len, + latest_detections=latest_detections, + ) + if is_final: + state_by_choice[choice_index] = replace(next_state, finished=True) + continue - yield item - else: - # Not a ModelResponseStream or no choices - yield as is - yield item + trim_started = time.perf_counter() + state_by_choice[choice_index] = self._trim_streamed_choice_buffer(next_state, masked_text, plan) + scan_seconds += time.perf_counter() - trim_started + + if stream_ended and not any(_streamed_choice_content(choice) for choice in _streamed_choices(item)): + continue + yield item except HTTPException: status = "guardrail_intervened" raise diff --git a/tests/integration/streaming/test_text_completion_stream_guardrail_rebuild.py b/tests/integration/streaming/test_text_completion_stream_guardrail_rebuild.py index 229a9be6d20..ea0a820a705 100644 --- a/tests/integration/streaming/test_text_completion_stream_guardrail_rebuild.py +++ b/tests/integration/streaming/test_text_completion_stream_guardrail_rebuild.py @@ -27,7 +27,12 @@ _PROXY_CONFIG: Final = Path(__file__).resolve().parents[1] / "proxy_config.yaml" _PROVIDER_KEY: Final = "synthetic-provider-key" _GUARDRAIL: Final = "audit-text-rebuild-tool-permission" _TEXT_MODEL: Final = f"audit-text-instruct-{uuid.uuid4().hex}" +_GOVERNED_MODEL: Final = f"audit-text-governed-{uuid.uuid4().hex}" +_KEYWORD_FILTER: Final = "audit-text-keyword-filter" +_POLICY: Final = "audit-text-post-call-governance" _PIECES: Final = ("Hello", " streamed", " world") +_MASKED_PIECES: Final = ("A pers", "im", "mon tart") +_BLOCKED_PIECES: Final = ("A kum", "quat tart") _PROMPT_TOKENS: Final = 11 _COMPLETION_TOKENS: Final = 3 _USAGE: Final[dict[str, JsonValue]] = { @@ -97,6 +102,14 @@ def _text_choice(text: str, finish_reason: str | None) -> dict[str, JsonValue]: return {"text": text, "index": 0, "logprobs": None, "finish_reason": finish_reason} +def _pieces_for(identity: str) -> tuple[str, ...]: + if identity.startswith("cmpl-mask-"): + return _MASKED_PIECES + if identity.startswith("cmpl-block-"): + return _BLOCKED_PIECES + return _PIECES + + def _respond(request: Request) -> Reply: body: Final = JSON_OBJECT.validate_json(request.body) assert body.get("stream") is True, body @@ -106,7 +119,7 @@ def _respond(request: Request) -> Reply: return Reply( content_type="text/event-stream", chunks=( - *(_text_frame(identity, [_text_choice(piece, None)]) for piece in _PIECES), + *(_text_frame(identity, [_text_choice(piece, None)]) for piece in _pieces_for(identity)), _text_frame(identity, [_text_choice("", "length")]), _text_frame(identity, [], usage=_USAGE), b"data: [DONE]\n\n", @@ -114,20 +127,25 @@ def _respond(request: Request) -> Reply: ) +def _text_deployment(model_name: str, wire: Wire) -> dict[str, JsonValue]: + return { + "model_name": model_name, + "litellm_params": { + "model": "text-completion-openai/gpt-3.5-turbo-instruct", + "api_base": f"{wire.url}/v1", + "api_key": _PROVIDER_KEY, + }, + } + + def _config(directory: Path, wire: Wire) -> Path: base: Final = object_value(yaml.safe_load(_PROXY_CONFIG.read_text())) config: Final = { **base, "model_list": [ *list_value(base["model_list"]), - { - "model_name": _TEXT_MODEL, - "litellm_params": { - "model": "text-completion-openai/gpt-3.5-turbo-instruct", - "api_base": f"{wire.url}/v1", - "api_key": _PROVIDER_KEY, - }, - }, + _text_deployment(_TEXT_MODEL, wire), + _text_deployment(_GOVERNED_MODEL, wire), ], "guardrails": [ { @@ -139,8 +157,30 @@ def _config(directory: Path, wire: Wire) -> Path: "default_action": "allow", "rules": [], }, - } + }, + { + "guardrail_name": _KEYWORD_FILTER, + "litellm_params": { + "guardrail": "litellm_content_filter", + "mode": "post_call", + "default_on": False, + "blocked_words": [ + {"keyword": "persimmon", "action": "MASK"}, + {"keyword": "kumquat", "action": "BLOCK"}, + ], + }, + }, ], + "policies": { + _POLICY: { + "guardrails": {"add": [_KEYWORD_FILTER]}, + "pipeline": { + "mode": "post_call", + "steps": [{"guardrail": _KEYWORD_FILTER, "on_pass": "allow", "on_fail": "block"}], + }, + } + }, + "policy_attachments": [{"policy": _POLICY, "models": [_GOVERNED_MODEL]}], } path: Final = directory / f"text-rebuild-{uuid.uuid4().hex}.yaml" path.write_text(yaml.safe_dump(config)) @@ -163,14 +203,14 @@ def rig(tmp_path_factory: pytest.TempPathFactory) -> Iterator[_Rig]: yield _Rig(wire, owned) -def _stream(rig: _Rig, identity: str) -> _Stream: +def _stream(rig: _Rig, identity: str, model: str = _TEXT_MODEL) -> _Stream: gateway: Final = rig.proxy.gateway with gateway.client.stream( "POST", "/v1/completions", headers={"Authorization": f"Bearer {gateway.key}"}, json={ - "model": _TEXT_MODEL, + "model": model, "prompt": identity, "max_tokens": 16, "stream": True, @@ -212,6 +252,37 @@ def test_guarded_text_completion_stream_reaches_the_caller_whole_with_its_usage( assert _GUARDRAIL in list_value(object_value(row["metadata"])["applied_guardrails"]), row +def test_post_call_pipeline_masks_a_streamed_text_completion(rig: _Rig) -> None: + identity: Final = f"cmpl-mask-{uuid.uuid4().hex}" + + stream: Final = _stream(rig, identity, _GOVERNED_MODEL) + + assert stream.status == 200, stream.body + assert stream.body.rstrip().endswith(b"data: [DONE]"), stream.body[-400:] + assert stream.text == "A [KEYWORD_REDACTED] tart", stream.frames + assert b"persimmon" not in stream.body, stream.body + assert stream.finish_reasons == ("length",), stream.frames + assert [(usage["prompt_tokens"], usage["completion_tokens"]) for usage in stream.usages] == [ + (_PROMPT_TOKENS, _COMPLETION_TOKENS) + ], stream.frames + (row,) = eventually( + lambda: rig.proxy.gateway.rows('SELECT metadata FROM "LiteLLM_SpendLogs" WHERE request_id = %s', (identity,)), + lambda rows: len(rows) == 1, + seconds=70, + ) + assert _KEYWORD_FILTER in list_value(object_value(row["metadata"])["applied_guardrails"]), row + + +def test_post_call_pipeline_blocks_a_streamed_text_completion(rig: _Rig) -> None: + identity: Final = f"cmpl-block-{uuid.uuid4().hex}" + + stream: Final = _stream(rig, identity, _GOVERNED_MODEL) + + assert stream.status == 400, stream.body + assert b"Content blocked: keyword 'kumquat' detected" in stream.body, stream.body + assert b"quat tart" not in stream.body, stream.body + + def _together_chunk(choices: list[JsonValue], **extra: JsonValue) -> bytes: return _frame( { diff --git a/tests/unit/litellm_core_utils/test_served_output_texts.py b/tests/unit/litellm_core_utils/test_served_output_texts.py index 637a20571c8..e01c710dfdb 100644 --- a/tests/unit/litellm_core_utils/test_served_output_texts.py +++ b/tests/unit/litellm_core_utils/test_served_output_texts.py @@ -9,7 +9,16 @@ from litellm.litellm_core_utils.served_output_texts import ( served_output_texts, served_stream_output_texts, ) -from litellm.types.utils import Choices, Delta, Message, ModelResponse, ModelResponseStream, StreamingChoices +from litellm.types.utils import ( + Choices, + Delta, + Message, + ModelResponse, + ModelResponseStream, + StreamingChoices, + TextChoices, + TextCompletionResponse, +) RAW = "Card: 4111 1111 1111 1111" MASKED = "Card: " @@ -83,6 +92,30 @@ def test_served_stream_output_texts_keeps_every_choice_index_when_chunks_carry_o assert served_stream_output_texts(chunks) == ("first choice", MASKED) +def _text_completion_chunk(text: str | None, index: int = 0) -> TextCompletionResponse: + return TextCompletionResponse(object="text_completion", choices=[TextChoices(index=index, text=text)]) + + +def test_served_stream_output_texts_joins_text_completion_chunks_per_choice(): + chunks = [ + _text_completion_chunk("Card: ", 0), + _text_completion_chunk("second", 1), + _text_completion_chunk("", 0), + _text_completion_chunk(None, 1), + TextCompletionResponse(object="text_completion", choices=[]), + ] + assert served_stream_output_texts(chunks) == (MASKED, "second") + + +def test_overlay_replaces_text_completion_choice_text_with_served_text(): + logged = {"id": "x", "object": "text_completion", "choices": [{"index": 0, "text": RAW}, {"index": 1, "text": RAW}]} + assert overlay_served_output_texts(logged, (MASKED,)) == { + "id": "x", + "object": "text_completion", + "choices": [{"index": 0, "text": MASKED}, {"index": 1, "text": RAW}], + } + + @pytest.mark.parametrize("content", (None, "", "first")) def test_null_choice_index_joins_choice_zero_without_changing_other_choices(content: str | None) -> None: chunks: Final = ( diff --git a/tests/unit/llms/openai/completion/test_text_completion_guardrail_handler.py b/tests/unit/llms/openai/completion/test_text_completion_guardrail_handler.py index c96fbf34fe1..9beb1faf7f4 100644 --- a/tests/unit/llms/openai/completion/test_text_completion_guardrail_handler.py +++ b/tests/unit/llms/openai/completion/test_text_completion_guardrail_handler.py @@ -2,18 +2,25 @@ Unit tests for OpenAI Text Completion Guardrail Translation Handler """ -from typing import List, Optional, Tuple +import json +from typing import List, Literal, Optional, Tuple from unittest.mock import MagicMock import pytest +from fastapi import HTTPException from litellm.integrations.custom_guardrail import CustomGuardrail +from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.llms import get_guardrail_translation_mapping from litellm.llms.openai.completion.guardrail_translation.handler import ( OpenAITextCompletionHandler, ) -from litellm.types.utils import CallTypes, TextChoices, TextCompletionResponse +from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy.guardrails.guardrail_hooks.litellm_content_filter.content_filter import ContentFilterGuardrail +from litellm.proxy.guardrails.guardrail_hooks.unified_guardrail.unified_guardrail import UnifiedLLMGuardrails +from litellm.types.guardrails import BlockedWord, ContentFilterAction, GuardrailEventHooks +from litellm.types.utils import CallTypes, GenericGuardrailAPIInputs, TextChoices, TextCompletionResponse, Usage class MockGuardrail(CustomGuardrail): @@ -341,3 +348,215 @@ class TestPIIMaskingScenario: assert "[EMAIL_REDACTED]" in result["prompt"][1] assert "alice@example.com" not in result["prompt"][0] assert "bob@test.org" not in result["prompt"][1] + + +class _RecordingKeywordFilter(ContentFilterGuardrail): + def __init__(self): + super().__init__( + guardrail_name="keyword-filter", + blocked_words=[ + BlockedWord(keyword="persimmon", action=ContentFilterAction.MASK), + BlockedWord(keyword="kumquat", action=ContentFilterAction.BLOCK), + ], + event_hook=GuardrailEventHooks.post_call, + default_on=True, + ) + self.streaming_sampling_rate = 1 + self.scanned_texts: tuple[tuple[str, ...], ...] = () + + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict, + input_type: Literal["request", "response"], + logging_obj: LiteLLMLoggingObj | None = None, + ) -> GenericGuardrailAPIInputs: + self.scanned_texts = (*self.scanned_texts, tuple(inputs.get("texts") or ())) + return await super().apply_guardrail(inputs, request_data, input_type, logging_obj) + + +class _PassthroughReplaceGuardrail(CustomGuardrail): + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict, + input_type: Literal["request", "response"], + logging_obj: LiteLLMLoggingObj | None = None, + ) -> GenericGuardrailAPIInputs: + if any("kumquat" in text for text in inputs.get("texts") or []): + self.raise_passthrough_exception(violation_message="replaced kumquat", request_data=request_data) + return inputs + + +def _text_chunk(index: int | None, text: str | None, finish_reason: str | None = None) -> TextCompletionResponse: + return TextCompletionResponse( + id="cmpl-stream", + created=1589478378, + model="gpt-3.5-turbo-instruct", + object="text_completion", + choices=[TextChoices(index=index, text=text, finish_reason=finish_reason)], + ) + + +def _usage_chunk() -> TextCompletionResponse: + return TextCompletionResponse( + id="cmpl-stream", + created=1589478378, + model="gpt-3.5-turbo-instruct", + object="text_completion", + choices=[], + usage=Usage(prompt_tokens=3, completion_tokens=5, total_tokens=8), + ) + + +def _delivered(chunks: list[TextCompletionResponse]) -> list[tuple[int, str | None, str | None]]: + return [(choice.index, choice.text, choice.finish_reason) for chunk in chunks for choice in chunk.choices] + + +async def _drive_flat_stream(guardrail: CustomGuardrail, chunks: list[TextCompletionResponse]) -> list[object]: + async def upstream(): + for chunk in chunks: + yield chunk + + return [ + item + async for item in UnifiedLLMGuardrails().async_post_call_streaming_iterator_hook( + user_api_key_dict=UserAPIKeyAuth(api_key="test-key", request_route="/v1/completions"), + response=upstream(), + request_data={"guardrail_to_apply": guardrail, "model": "gpt-3.5-turbo-instruct"}, + ) + ] + + +class TestStreamingOutputProcessing: + @pytest.mark.asyncio + async def test_ended_stream_delivers_each_choice_masked_text(self): + guardrail = _RecordingKeywordFilter() + chunks = [ + _text_chunk(0, "A pers"), + _text_chunk(1, "kiwi "), + _text_chunk(0, "im"), + _text_chunk(1, "pie", "stop"), + _text_chunk(0, "mon tart"), + _text_chunk(0, None, "length"), + _usage_chunk(), + ] + + await OpenAITextCompletionHandler().process_output_streaming_response( + responses_so_far=chunks, + guardrail_to_apply=guardrail, + request_data={}, + deliver_ended_stream_rewrites=True, + ) + + assert guardrail.scanned_texts == (("A persimmon tart", "kiwi pie"),) + assert _delivered(chunks) == [ + (0, "A [KEYWORD_REDACTED] tart", None), + (1, "kiwi ", None), + (0, "", None), + (1, "pie", "stop"), + (0, "", None), + (0, None, "length"), + ] + assert chunks[-1].usage.total_tokens == 8 + + @pytest.mark.asyncio + async def test_ended_stream_scanned_without_delivery_keeps_chunks_and_stores_the_assembled_response(self): + chunks = [_text_chunk(0, "A pers"), _text_chunk(0, "immon tart"), _text_chunk(0, "", "stop"), _usage_chunk()] + request_data: dict[str, object] = {} + + await OpenAITextCompletionHandler().process_output_streaming_response( + responses_so_far=chunks, guardrail_to_apply=_RecordingKeywordFilter(), request_data=request_data + ) + + assert _delivered(chunks) == [(0, "A pers", None), (0, "immon tart", None), (0, "", "stop")] + assembled = request_data["response"] + assert isinstance(assembled, TextCompletionResponse) + assert [(c.text, c.finish_reason) for c in assembled.choices] == [("A [KEYWORD_REDACTED] tart", "stop")] + assert assembled.usage.total_tokens == 8 + + @pytest.mark.asyncio + async def test_ended_stream_block_raises_the_guardrail_error(self): + chunks = [_text_chunk(0, "A kum"), _text_chunk(0, "quat tart"), _text_chunk(0, "", "stop")] + + with pytest.raises(HTTPException) as exc_info: + await OpenAITextCompletionHandler().process_output_streaming_response( + responses_so_far=chunks, guardrail_to_apply=_RecordingKeywordFilter(), deliver_ended_stream_rewrites=True + ) + + assert exc_info.value.status_code == 400 + assert "kumquat" in str(exc_info.value.detail) + + @pytest.mark.asyncio + async def test_ended_stream_treats_a_null_choice_index_as_choice_zero(self): + guardrail = _RecordingKeywordFilter() + chunks = [_text_chunk(None, "A pers"), _text_chunk(0, "immon tart"), _text_chunk(0, "", "stop")] + + await OpenAITextCompletionHandler().process_output_streaming_response( + responses_so_far=chunks, guardrail_to_apply=guardrail, request_data={}, deliver_ended_stream_rewrites=True + ) + + assert guardrail.scanned_texts == (("A persimmon tart",),) + assert [choice.text for chunk in chunks for choice in chunk.choices] == ["A [KEYWORD_REDACTED] tart", "", ""] + + @pytest.mark.asyncio + async def test_ended_stream_blocks_a_keyword_split_across_a_null_and_a_zero_choice_index(self): + chunks = [_text_chunk(None, "A kum"), _text_chunk(0, "quat tart"), _text_chunk(0, "", "stop")] + + with pytest.raises(HTTPException) as exc_info: + await OpenAITextCompletionHandler().process_output_streaming_response( + responses_so_far=chunks, guardrail_to_apply=_RecordingKeywordFilter(), deliver_ended_stream_rewrites=True + ) + + assert exc_info.value.status_code == 400 + assert "kumquat" in str(exc_info.value.detail) + + @pytest.mark.asyncio + async def test_flat_stream_is_scanned_once_after_it_ends(self): + guardrail = _RecordingKeywordFilter() + chunks = [_text_chunk(0, "A pers"), _text_chunk(0, "immon tart"), _text_chunk(0, "", "stop"), _usage_chunk()] + + delivered = await _drive_flat_stream(guardrail, chunks) + + assert guardrail.scanned_texts == (("A persimmon tart",),) + assert len(delivered) == len(chunks) + + @pytest.mark.asyncio + async def test_flat_stream_block_ends_with_an_in_band_error_frame(self): + chunks = [_text_chunk(0, "A kum"), _text_chunk(0, "quat tart"), _text_chunk(0, "", "stop"), _usage_chunk()] + + delivered = await _drive_flat_stream(_RecordingKeywordFilter(), chunks) + + error_frame = delivered[-1] + assert isinstance(error_frame, bytes) + assert error_frame.startswith(b'data: {"error": ') + assert b"kumquat" in error_frame + + @pytest.mark.asyncio + async def test_flat_stream_passthrough_block_appends_the_replacement_and_ends_with_content_filter(self): + guardrail = _PassthroughReplaceGuardrail( + guardrail_name="passthrough", event_hook=GuardrailEventHooks.post_call, default_on=True + ) + chunks = [_text_chunk(0, "A kum"), _text_chunk(0, "quat tart"), _text_chunk(0, "", "stop"), _usage_chunk()] + + delivered = await _drive_flat_stream(guardrail, chunks) + + frames = [json.loads(item.decode().removeprefix("data: ")) for item in delivered if isinstance(item, bytes)] + assert [(frame["choices"][0]["text"], frame["choices"][0].get("finish_reason")) for frame in frames] == [ + ("replaced kumquat", None), + ("", "content_filter"), + ] + assert frames[-1]["usage"]["total_tokens"] == 8 + + @pytest.mark.asyncio + async def test_flat_stream_passthrough_block_reports_no_usage_the_upstream_never_sent(self): + guardrail = _PassthroughReplaceGuardrail( + guardrail_name="passthrough", event_hook=GuardrailEventHooks.post_call, default_on=True + ) + chunks = [_text_chunk(0, "A kum"), _text_chunk(0, "quat tart"), _text_chunk(0, "", "stop")] + + delivered = await _drive_flat_stream(guardrail, [chunk.model_copy(update={"usage": None}) for chunk in chunks]) + + frames = [json.loads(item.decode().removeprefix("data: ")) for item in delivered if isinstance(item, bytes)] + assert frames[-1]["choices"][0]["finish_reason"] == "content_filter" + assert [frame for frame in frames if "usage" in frame] == [] diff --git a/tests/unit/proxy/guardrails/guardrail_hooks/content_filter/test_content_filter.py b/tests/unit/proxy/guardrails/guardrail_hooks/content_filter/test_content_filter.py index f7fb7e1fed0..9a14222c684 100644 --- a/tests/unit/proxy/guardrails/guardrail_hooks/content_filter/test_content_filter.py +++ b/tests/unit/proxy/guardrails/guardrail_hooks/content_filter/test_content_filter.py @@ -4,6 +4,7 @@ Tests for the Content Filter Guardrail import json import os +from collections.abc import AsyncGenerator, AsyncIterator, Callable from typing import Final from unittest.mock import MagicMock @@ -30,7 +31,17 @@ from litellm.types.proxy.guardrails.guardrail_hooks.litellm_content_filter impor ContentFilterCategoryConfig, ContentFilterDetection, ) -from litellm.types.utils import StandardLoggingGuardrailInformation +from litellm.llms.base_llm.guardrail_translation.utils import stream_item_field +from litellm.types.llms.openai import OutputTextDeltaEvent, ResponsesAPIStreamEvents +from litellm.types.utils import ( + Delta, + ModelResponseStream, + StandardLoggingGuardrailInformation, + StreamingChoices, + TextChoices, + TextCompletionResponse, + Usage, +) class TestContentFilterGuardrail: @@ -3523,3 +3534,180 @@ async def test_image_whose_description_is_clean_passes_through_unchanged(): ) assert result == {"images": ["https://example.com/desk.png"], "texts": []} + + +def _fruit_word_filter() -> ContentFilterGuardrail: + return ContentFilterGuardrail( + guardrail_name="fruit-filter", + blocked_words=[ + BlockedWord(keyword="persimmon", action=ContentFilterAction.MASK), + BlockedWord(keyword="kumquat", action=ContentFilterAction.BLOCK), + ], + event_hook=GuardrailEventHooks.post_call, + ) + + +async def _text_completion_stream(*parts: tuple[str | None, str | None]): + for text, finish_reason in parts: + yield TextCompletionResponse( + id="cmpl-stream", + model="gpt-3.5-turbo-instruct", + object="text_completion", + choices=[TextChoices(index=0, text=text, finish_reason=finish_reason)], + ) + + +@pytest.mark.asyncio +async def test_streaming_hook_masks_a_word_split_across_text_completion_chunks(): + request_data: Final[dict] = {} + + delivered: Final = [ + chunk + async for chunk in _fruit_word_filter().async_post_call_streaming_iterator_hook( + user_api_key_dict=MagicMock(), + response=_text_completion_stream(("A pers", None), ("immon tart", None), (None, "stop")), + request_data=request_data, + ) + ] + + assert "".join(chunk.choices[0].text or "" for chunk in delivered) == "A [KEYWORD_REDACTED] tart" + assert request_data["metadata"]["standard_logging_guardrail_information"][0]["guardrail_status"] == "success" + + +_LONG_TAIL: Final = "the quick brown fox jumps over the lazy dog and then eats a ripe persimmon" + + +async def _text_completion_n2_stream_without_choice_1_finish() -> AsyncGenerator[TextCompletionResponse, None]: + for index, text, finish_reason in ( + (0, "A short answer", None), + (1, _LONG_TAIL[:30], None), + (1, _LONG_TAIL[30:], None), + (0, None, "stop"), + ): + yield TextCompletionResponse( + id="cmpl-stream", + model="gpt-3.5-turbo-instruct", + object="text_completion", + choices=[TextChoices(index=index, text=text, finish_reason=finish_reason)], + ).model_copy(update={"usage": None}) + + +async def _chat_n2_stream_without_choice_1_finish() -> AsyncGenerator[ModelResponseStream, None]: + for index, content, finish_reason in ( + (0, "A short answer", None), + (1, _LONG_TAIL[:30], None), + (1, _LONG_TAIL[30:], None), + (0, None, "stop"), + ): + yield ModelResponseStream( + id="chatcmpl-stream", + model="gpt-4.1-mini", + choices=[StreamingChoices(index=index, delta=Delta(content=content), finish_reason=finish_reason)], + ) + + +def _streamed_text(chunk: ModelResponseStream | TextCompletionResponse, index: int) -> str: + return "".join(_choice_text(choice) for choice in chunk.choices if choice.index == index) + + +def _choice_text(choice: TextChoices | StreamingChoices) -> str: + content: Final = stream_item_field(choice, "text") if isinstance(choice, TextChoices) else choice.delta.content + return content if isinstance(content, str) else "" + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "stream", [_text_completion_n2_stream_without_choice_1_finish, _chat_n2_stream_without_choice_1_finish] +) +async def test_streaming_hook_releases_the_held_tail_of_a_choice_the_stream_never_finished( + stream: Callable[[], AsyncIterator[ModelResponseStream | TextCompletionResponse]], +) -> None: + delivered: Final = [ + chunk + async for chunk in _fruit_word_filter().async_post_call_streaming_iterator_hook( + user_api_key_dict=MagicMock(), response=stream(), request_data={} + ) + ] + + assert "".join(_streamed_text(chunk, 0) for chunk in delivered) == "A short answer" + assert "".join(_streamed_text(chunk, 1) for chunk in delivered) == _LONG_TAIL.replace( + "persimmon", "[KEYWORD_REDACTED]" + ) + assert all(getattr(chunk, "usage", None) is None for chunk in delivered) + + +_STREAM_USAGE: Final = Usage(prompt_tokens=5, completion_tokens=20, total_tokens=25) + + +async def _then_usage_chunk( + stream: Callable[[], AsyncIterator[ModelResponseStream | TextCompletionResponse]], + usage_chunk: ModelResponseStream | TextCompletionResponse, +) -> AsyncGenerator[ModelResponseStream | TextCompletionResponse, None]: + async for chunk in stream(): + yield chunk + yield usage_chunk + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("stream", "usage_chunk"), + [ + ( + _text_completion_n2_stream_without_choice_1_finish, + TextCompletionResponse(id="cmpl-stream", object="text_completion", choices=[], usage=_STREAM_USAGE), + ), + (_chat_n2_stream_without_choice_1_finish, ModelResponseStream(id="chatcmpl-stream", choices=[], usage=_STREAM_USAGE)), + ], +) +async def test_streaming_hook_keeps_the_usage_chunk_last_after_releasing_a_held_tail( + stream: Callable[[], AsyncIterator[ModelResponseStream | TextCompletionResponse]], + usage_chunk: ModelResponseStream | TextCompletionResponse, +) -> None: + delivered: Final = [ + chunk + async for chunk in _fruit_word_filter().async_post_call_streaming_iterator_hook( + user_api_key_dict=MagicMock(), response=_then_usage_chunk(stream, usage_chunk), request_data={} + ) + ] + + assert delivered[-1] is usage_chunk + assert "".join(_streamed_text(chunk, 1) for chunk in delivered[:-1]) == _LONG_TAIL.replace( + "persimmon", "[KEYWORD_REDACTED]" + ) + + +@pytest.mark.asyncio +async def test_streaming_hook_relays_each_responses_api_event_as_it_arrives() -> None: + timeline: Final[list[str]] = [] + + async def responses_api_events() -> AsyncGenerator[OutputTextDeltaEvent, None]: + for delta in ("first", "second"): + timeline.append(f"upstream sent {delta}") + yield OutputTextDeltaEvent( + type=ResponsesAPIStreamEvents.OUTPUT_TEXT_DELTA, + item_id="msg_1", + output_index=0, + content_index=0, + delta=delta, + ) + + async for event in _fruit_word_filter().async_post_call_streaming_iterator_hook( + user_api_key_dict=MagicMock(), response=responses_api_events(), request_data={} + ): + timeline.append(f"client got {event.delta}") + + assert timeline == ["upstream sent first", "client got first", "upstream sent second", "client got second"] + + +@pytest.mark.asyncio +async def test_streaming_hook_blocks_a_word_split_across_text_completion_chunks(): + with pytest.raises(HTTPException) as exc_info: + async for _ in _fruit_word_filter().async_post_call_streaming_iterator_hook( + user_api_key_dict=MagicMock(), + response=_text_completion_stream(("A kum", None), ("quat tart", None), (None, "stop")), + request_data={}, + ): + pass + + assert exc_info.value.status_code == 400 + assert "kumquat" in str(exc_info.value.detail) diff --git a/tests/unit/proxy/guardrails/guardrail_hooks/test_grayswan.py b/tests/unit/proxy/guardrails/guardrail_hooks/test_grayswan.py index 954c57b5cb6..a5ba9dab0bd 100644 --- a/tests/unit/proxy/guardrails/guardrail_hooks/test_grayswan.py +++ b/tests/unit/proxy/guardrails/guardrail_hooks/test_grayswan.py @@ -5,6 +5,7 @@ import pytest from fastapi import HTTPException from litellm.integrations.custom_guardrail import ModifyResponseException +from litellm.llms.openai.completion.guardrail_translation.handler import OpenAITextCompletionHandler from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.guardrails.guardrail_hooks.grayswan import grayswan as grayswan_module from litellm.proxy.guardrails.guardrail_hooks.grayswan.grayswan import ( @@ -12,6 +13,7 @@ from litellm.proxy.guardrails.guardrail_hooks.grayswan.grayswan import ( GraySwanGuardrailAPIError, ) from litellm.types.guardrails import GuardrailEventHooks +from litellm.types.utils import TextChoices, TextCompletionResponse @pytest.fixture @@ -884,3 +886,29 @@ async def test_pre_call_payload_unchanged() -> None: {"role": "user", "content": "second"}, ] assert "tools" not in payload + + +@pytest.mark.asyncio +async def test_post_call_never_forwards_the_proxy_key_metadata_of_the_request_metadata() -> None: + guardrail = _post_call_guardrail() + client = _CapturingClient() + guardrail.async_handler = client + + await OpenAITextCompletionHandler().process_output_response( + response=TextCompletionResponse( + id="cmpl-1", model="gpt-4.1-mini", choices=[TextChoices(index=0, text="hello", finish_reason="stop")] + ), + guardrail_to_apply=guardrail, + user_api_key_dict=UserAPIKeyAuth(team_id="team-456", key_alias="ci-key", user_email="dev@example.com"), + request_data={ + "model": "gpt-4.1-mini", + "prompt": "say hello", + "metadata": { + "user_api_key_hash": "hashed-key", + "user_api_key_user_email": "dev@example.com", + "user_api_key_request_route": "/v1/completions", + }, + }, + ) + + assert "litellm_metadata" not in client.calls[0]["json"] diff --git a/tests/unit/proxy/guardrails/test_deferred_guardrail_logging.py b/tests/unit/proxy/guardrails/test_deferred_guardrail_logging.py index 7e5d0f6342e..bd6190fee5f 100644 --- a/tests/unit/proxy/guardrails/test_deferred_guardrail_logging.py +++ b/tests/unit/proxy/guardrails/test_deferred_guardrail_logging.py @@ -1537,6 +1537,27 @@ class TestArmDeferredStreamDispatch: assert recorded["cache_hit"] is False assert recorded["prefer_async_handlers"] is True + @pytest.mark.asyncio + async def test_text_completion_stream_logs_after_the_guardrails_finish(self): + from litellm.text_completions.streaming_iterator import TextCompletionStreamWrapper + + logging_obj, recorded = self._dispatch_recording_logging_obj() + + self._processor()._arm_deferred_stream_dispatch( + response=TextCompletionStreamWrapper(completion_stream=MagicMock(), model="gpt-test"), + route_type="atext_completion", + user_api_key_dict=MagicMock(), + logging_obj=logging_obj, + ) + + assembled = object() + logging_obj._deferred_stream_complete_args = (assembled, False) + ProxyLogging._fire_deferred_stream_logging({"litellm_logging_obj": logging_obj}) + await asyncio.sleep(0) + + assert recorded["result"] is assembled + assert recorded["cache_hit"] is False + @pytest.mark.asyncio async def test_router_wrapped_bridged_iterator_gets_csw_arg_shape(self): """The router wraps iterators without _hidden_params in diff --git a/tests/unit/proxy/utils/proxy_logging/test_guardrail_pipeline.py b/tests/unit/proxy/utils/proxy_logging/test_guardrail_pipeline.py index acc1463b072..cfd4fe197be 100644 --- a/tests/unit/proxy/utils/proxy_logging/test_guardrail_pipeline.py +++ b/tests/unit/proxy/utils/proxy_logging/test_guardrail_pipeline.py @@ -41,6 +41,7 @@ from litellm.types.proxy.policy_engine.pipeline_types import ( GuardrailPipeline, PipelineStep, ) +from litellm.types.utils import TextChoices, TextCompletionResponse @pytest.fixture(autouse=True) @@ -1886,7 +1887,7 @@ def test_streamable_post_call_pipelines_keeps_hook_guardrails_and_drops_iterator @pytest.mark.parametrize( "request_route", - ["/v1/completions", "/v1beta/models/gemini-2.5-flash:streamGenerateContent", "/a2a/agent"], + ["/v1beta/models/gemini-2.5-flash:streamGenerateContent", "/a2a/agent"], ) def test_streamable_post_call_pipelines_keeps_legacy_hooks_off_routes_that_assemble_no_response( make_user_api_key_auth, monkeypatch, caplog, request_route @@ -2036,6 +2037,120 @@ async def test_streaming_iterator_hook_delivers_legacy_hook_rewrite_on_every_cho assert _warnings(caplog) == [] +def _text_completion_chunk(text: str, finish_reason: str | None = None) -> TextCompletionResponse: + return TextCompletionResponse( + id="cmpl-1", model="m", choices=[TextChoices(index=0, text=text, finish_reason=finish_reason)] + ) + + +def _rewritten_text_completion(response: Any) -> TextCompletionResponse: + return TextCompletionResponse( + id=response.id, + model=response.model, + choices=[TextChoices(index=0, text="[REWRITTEN] " + response.choices[0].text, finish_reason="stop")], + ) + + +@pytest.mark.asyncio +async def test_streaming_iterator_hook_delivers_legacy_hook_rewrite_on_a_text_completion_stream( + proxy_logging, make_user_api_key_auth, monkeypatch, caplog +): + seen: Dict[str, Any] = {} + monkeypatch.setattr(litellm, "callbacks", [_legacy_hook_stream_guardrail(seen, rewrite=_rewritten_text_completion)]) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False) + data = {**_post_call_pipeline_data(stream=True), "prompt": "hi"} + chunks = [_text_completion_chunk("hello "), _text_completion_chunk("world"), _text_completion_chunk("", "stop")] + auth = make_user_api_key_auth(request_route="/v1/completions") + + with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): + await proxy_logging.pre_call_hook( + user_api_key_dict=auth, data=data, call_type="text_completion", guardrails_only=True + ) + delivered = [ + item + async for item in proxy_logging.async_post_call_streaming_iterator_hook( + user_api_key_dict=auth, response=_async_chunk_iter(chunks), request_data=data + ) + ] + + assert isinstance(seen["response"], TextCompletionResponse) + assert seen["response"].choices[0].text == "hello world" + assert [(item.choices[0].text, item.choices[0].finish_reason) for item in delivered] == [ + ("[REWRITTEN] hello world", None), + ("", None), + ("", "stop"), + ] + assert data["metadata"]["applied_guardrails"] == ["gr-post"] + assert _warnings(caplog) == [] + + +class _FlatBlockGuardrail(CustomGuardrail): + async def apply_guardrail(self, inputs, request_data, input_type, logging_obj=None): + if any("kumquat" in text for text in inputs.get("texts") or []): + raise HTTPException(status_code=400, detail="flat guardrail blocked kumquat") + return inputs + + +@pytest.mark.asyncio +async def test_streaming_iterator_hook_delivers_a_flat_guardrail_block_on_a_text_completion_stream_under_a_pipeline( + proxy_logging, make_user_api_key_auth, monkeypatch +): + seen: Final[Dict[str, Any]] = {} + flat: Final = _FlatBlockGuardrail( + guardrail_name="flat-block", event_hook=GuardrailEventHooks.post_call, default_on=True + ) + monkeypatch.setattr(litellm, "callbacks", [flat, _live_scan_guardrail(seen)]) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False) + data: Final = {**_post_call_pipeline_data(stream=True), "prompt": "hi"} + chunks: Final = [ + _text_completion_chunk("A kum"), + _text_completion_chunk("quat tart"), + _text_completion_chunk("", "stop"), + ] + auth: Final = make_user_api_key_auth(request_route="/v1/completions") + + await proxy_logging.pre_call_hook(user_api_key_dict=auth, data=data, call_type="text_completion", guardrails_only=True) + delivered: Final = [ + item + async for item in proxy_logging.async_post_call_streaming_iterator_hook( + user_api_key_dict=auth, response=_async_chunk_iter(chunks), request_data=data + ) + ] + + assert seen["texts"] == ["A kumquat tart"] + assert isinstance(delivered[-1], bytes) + assert json.loads(delivered[-1].decode().removeprefix("data: "))["error"]["message"] == ( + "flat guardrail blocked kumquat" + ) + + +@pytest.mark.asyncio +async def test_streaming_iterator_hook_live_pipeline_scans_what_a_disconnected_text_completion_client_received( + proxy_logging, make_user_api_key_auth, monkeypatch +): + seen: Final[Dict[str, Any]] = {} + monkeypatch.setattr(litellm, "callbacks", [_live_scan_guardrail(seen)]) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False) + data: Final = {**_post_call_pipeline_data(step=_DETECT_ONLY_STEP, stream=True), "prompt": "hi"} + chunks: Final = [ + _text_completion_chunk("hello "), + _text_completion_chunk("world"), + _text_completion_chunk("", "stop"), + ] + auth: Final = make_user_api_key_auth(request_route="/v1/completions") + + await proxy_logging.pre_call_hook(user_api_key_dict=auth, data=data, call_type="text_completion", guardrails_only=True) + stream: Final = proxy_logging.async_post_call_streaming_iterator_hook( + user_api_key_dict=auth, response=_async_chunk_iter(chunks), request_data=data + ) + first: Final = await stream.__anext__() + await stream.aclose() + + assert first.choices[0].text == "hello " + assert seen["texts"] == ["hello "] + assert _post_call_statuses(data) == ["success"] + + @pytest.mark.asyncio async def test_streaming_iterator_hook_releases_stream_untouched_when_legacy_hook_returns_none( proxy_logging, make_user_api_key_auth, monkeypatch @@ -2615,6 +2730,45 @@ async def test_streaming_iterator_hook_pipeline_modify_response_emits_translated assert not any(item is chunk for item in delivered for chunk in chunks) +@pytest.mark.asyncio +async def test_streaming_iterator_hook_pipeline_modify_response_replaces_a_text_completion_stream( + proxy_logging, make_user_api_key_auth, monkeypatch +): + seen: Final[Dict[str, Any]] = {} + monkeypatch.setattr(litellm, "callbacks", [_unified_stream_guardrail(seen, block=True)]) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False) + pipeline: Final = GuardrailPipeline( + mode="post_call", + steps=[ + PipelineStep( + guardrail="gr-post", + on_pass="allow", + on_fail="modify_response", + modify_response_message="content policy block", + ) + ], + ) + data: Final = {**_post_call_pipeline_data(stream=True), "prompt": "hi"} + data["metadata"]["_guardrail_pipelines"] = [("response-governance", pipeline)] + chunks: Final = [_text_completion_chunk("hello "), _text_completion_chunk("world"), _text_completion_chunk("", "stop")] + auth: Final = make_user_api_key_auth(request_route="/v1/completions") + + await proxy_logging.pre_call_hook(user_api_key_dict=auth, data=data, call_type="text_completion", guardrails_only=True) + delivered: Final = [ + item + async for item in proxy_logging.async_post_call_streaming_iterator_hook( + user_api_key_dict=auth, response=_async_chunk_iter(chunks), request_data=data + ) + ] + + frames: Final = [json.loads(item.decode().removeprefix("data: ")) for item in delivered] + assert seen["count"] == 1 + assert [(frame["choices"][0]["text"], frame["choices"][0].get("finish_reason")) for frame in frames] == [ + ("content policy block", None), + ("", "content_filter"), + ] + + @pytest.mark.asyncio async def test_streaming_iterator_hook_pipeline_delivers_text_rewrite_on_anthropic_sse( proxy_logging, make_user_api_key_auth, monkeypatch From c98641fd5ac1093c9196462a968dc634c6f70c94 Mon Sep 17 00:00:00 2001 From: Rick <26716961+Bytechoreographer@users.noreply.github.com> Date: Fri, 30 Oct 2026 00:08:18 +0800 Subject: [PATCH 10/13] feat(dashscope)!: derive protocol bases from one domain (#47812) * feat(dashscope): derive protocol bases from one domain Accept the DashScope domain root, the OpenAI-compatible base, the Anthropic Messages base, or a native /api/v1 base and derive every surface from it Adds a native Anthropic Messages config for /apps/anthropic so a single deployment serves chat, messages, embeddings, rerank, image generation, and video against the same credential Co-Authored-By: AaronHowell <111272993+AaronHowell@users.noreply.github.com> * fix(dashscope): keep Anthropic credentials off DashScope messages Require a DashScope key or a caller auth header before the Messages path runs, and opt out of workload identity, so the base class cannot fall back to ANTHROPIC_API_KEY or ANTHROPIC_AUTH_TOKEN Also format the image generation base selection * style(dashscope): keep the messages config within the format gate * test(dashscope): cover caller Authorization on messages * fix(dashscope): keep the chat bridge for protocol-named bases and the image default for unconfigured calls * fix(dashscope): keep configured image and rerank bases, normalize blank keys, derive brand chat bases - image calls fall back to the provider's own chat base env (DASHSCOPE_API_BASE, QWENCLOUD_API_BASE, QWEN_AI_PLATFORM_API_BASE) before the image default - an explicit /compatible-api/v1/reranks URL wins over DASHSCOPE_API_BASE_RERANK again - a blank DashScope key counts as unset on /v1/messages, so it never falls back to ANTHROPIC_API_KEY - QwenCloud and Qianwen AI Platform chat and embeddings derive their compatible-mode base from a domain root - tests annotate new variables with Final and drop the config-class assertion * fix(dashscope): keep an explicit compatible-api rerank base as given * fix(dashscope): derive chat from compatible-api bases, keep the caller's rerank base, and let explicit bases beat per-endpoint env vars * fix(dashscope): keep a full DashScope image endpoint URL as given and type the brand tests * fix(dashscope): let the image env var beat any shared DashScope base * fix(dashscope): keep a full DASHSCOPE_API_BASE image URL ahead of the image env var * fix(dashscope): keep a native DashScope endpoint URL as given for chat and the api base header --------- Co-authored-by: AaronHowell <111272993+AaronHowell@users.noreply.github.com> Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> --- litellm/images/main.py | 12 +- litellm/llms/dashscope/chat/transformation.py | 9 +- litellm/llms/dashscope/common_utils.py | 136 ++++++++++- .../llms/dashscope/embed/transformation.py | 4 +- .../image_generation/transformation.py | 35 ++- litellm/llms/dashscope/messages/__init__.py | 7 + .../llms/dashscope/messages/transformation.py | 96 ++++++++ litellm/llms/dashscope/qwen_ai_platform.py | 16 +- litellm/llms/dashscope/qwencloud.py | 16 +- .../llms/dashscope/videos/transformation.py | 5 + litellm/rerank_api/main.py | 12 +- litellm/utils.py | 6 + .../test_websearch_short_circuit.py | 1 + .../test_dashscope_chat_transformation.py | 37 ++- .../dashscope/test_dashscope_common_utils.py | 171 ++++++++++++++ ...test_dashscope_embedding_transformation.py | 35 +++ .../test_dashscope_messages_transformation.py | 215 ++++++++++++++++++ .../test_dashscope_rerank_transformation.py | 52 +++++ .../test_dashscope_video_transformation.py | 14 ++ .../llms/dashscope/test_qwen_brand_aliases.py | 46 ++++ tests/unit/test_dashscope_image_generation.py | 127 ++++++++++- 21 files changed, 1015 insertions(+), 37 deletions(-) create mode 100644 litellm/llms/dashscope/messages/__init__.py create mode 100644 litellm/llms/dashscope/messages/transformation.py create mode 100644 tests/unit/llms/dashscope/test_dashscope_common_utils.py create mode 100644 tests/unit/llms/dashscope/test_dashscope_messages_transformation.py diff --git a/litellm/images/main.py b/litellm/images/main.py index b3247ff938b..4aad5210028 100644 --- a/litellm/images/main.py +++ b/litellm/images/main.py @@ -68,6 +68,11 @@ from litellm.utils import ( ) _PROVIDER_ROUTING_HINTS: Final = TypeAdapter(ProviderRoutingHints) +_PROVIDERS_RESOLVING_THEIR_OWN_IMAGE_BASE: Final = ( + LlmProviders.DASHSCOPE, + LlmProviders.QWENCLOUD, + LlmProviders.QWEN_AI_PLATFORM, +) # Cache for ImageEditRequestUtils to avoid repeated __getattr__ calls _ImageEditRequestUtils_cache: type["ImageEditRequestUtils"] | None = None @@ -271,6 +276,7 @@ def image_generation( """ try: args: Final = locals() + requested_api_base: Final = api_base aimg_generation: Final = kwargs.get("aimg_generation", False) litellm_call_id: Final = kwargs.get("litellm_call_id", None) logger_fn: Final = kwargs.get("logger_fn", None) @@ -468,7 +474,11 @@ def image_generation( raise ValueError(f"image generation config is not supported for {custom_llm_provider}") # Resolve api_base from litellm.api_base if not explicitly provided - _api_base: Final = api_base or litellm.api_base + _api_base: Final = ( + requested_api_base + if custom_llm_provider in _PROVIDERS_RESOLVING_THEIR_OWN_IMAGE_BASE + else api_base or litellm.api_base + ) litellm_params_dict["api_base"] = _api_base return llm_http_handler.image_generation_handler( diff --git a/litellm/llms/dashscope/chat/transformation.py b/litellm/llms/dashscope/chat/transformation.py index 966d1a01b8d..dfff9cec298 100644 --- a/litellm/llms/dashscope/chat/transformation.py +++ b/litellm/llms/dashscope/chat/transformation.py @@ -10,6 +10,7 @@ from litellm.secret_managers.main import get_secret_str from litellm.types.llms.openai import AllMessageValues, ChatCompletionToolParam from ...openai.chat.gpt_transformation import OpenAIGPTConfig +from ..common_utils import resolve_dashscope_openai_api_base class DashScopeChatConfig(OpenAIGPTConfig): @@ -55,11 +56,9 @@ class DashScopeChatConfig(OpenAIGPTConfig): def _get_openai_compatible_provider_info( self, api_base: str | None, api_key: str | None ) -> tuple[str | None, str | None]: - api_base = ( - api_base or get_secret_str("DASHSCOPE_API_BASE") or "https://dashscope.aliyuncs.com/compatible-mode/v1" - ) + resolved_api_base: Final = resolve_dashscope_openai_api_base(api_base) dynamic_api_key: Final = api_key or get_secret_str("DASHSCOPE_API_KEY") - return api_base, dynamic_api_key + return resolved_api_base, dynamic_api_key def get_openai_compatible_provider_info( self, @@ -69,7 +68,7 @@ class DashScopeChatConfig(OpenAIGPTConfig): return self._get_openai_compatible_provider_info(api_base, api_key) def _resolve_chat_api_base(self, api_base: str | None) -> str: - return api_base or "https://dashscope.aliyuncs.com/compatible-mode/v1" + return resolve_dashscope_openai_api_base(api_base) def get_complete_url( self, diff --git a/litellm/llms/dashscope/common_utils.py b/litellm/llms/dashscope/common_utils.py index 3e2e718a469..5ad2aa17a3a 100644 --- a/litellm/llms/dashscope/common_utils.py +++ b/litellm/llms/dashscope/common_utils.py @@ -2,6 +2,7 @@ Common utilities for the DashScope LLM provider. """ +from dataclasses import dataclass from typing import TYPE_CHECKING, Final from urllib.parse import urlparse @@ -19,10 +20,123 @@ if TYPE_CHECKING: from litellm.llms.base_llm.videos.transformation import BaseVideoConfig DASHSCOPE_CHAT_COMPATIBLE_PATH: Final = "/compatible-mode/v1" -DASHSCOPE_RERANK_PATH: Final = "/compatible-api/v1/reranks" +DASHSCOPE_ANTHROPIC_PATH: Final = "/apps/anthropic" +DASHSCOPE_NATIVE_V1_PATH: Final = "/api/v1" +DASHSCOPE_COMPATIBLE_API_PATH: Final = "/compatible-api/v1" +DASHSCOPE_RERANK_PATH: Final = f"{DASHSCOPE_COMPATIBLE_API_PATH}/reranks" +DASHSCOPE_IMAGE_PATH: Final = "/api/v1/services/aigc/multimodal-generation/generation" +DASHSCOPE_DEFAULT_ORIGIN: Final = "https://dashscope.aliyuncs.com" + +_STANDARD_DASHSCOPE_HOSTS: Final = frozenset( + { + "dashscope.aliyuncs.com", + "dashscope-intl.aliyuncs.com", + "dashscope-us.aliyuncs.com", + } +) +_WORKSPACE_HOST_SUFFIX: Final = ".maas.aliyuncs.com" +_EXCLUDED_WORKSPACE_PREFIXES: Final = ("trial.", "token-plan.") +_PROTOCOL_PATH_PREFIXES: Final = ( + DASHSCOPE_CHAT_COMPATIBLE_PATH, + DASHSCOPE_ANTHROPIC_PATH, + DASHSCOPE_NATIVE_V1_PATH, + DASHSCOPE_COMPATIBLE_API_PATH, +) -def _rerank_base_for_chat_shaped_base(api_base: str | None) -> str | None: +@dataclass(frozen=True, slots=True) +class DashScopeEndpoints: + origin: str + openai: str + anthropic: str + native: str + + +def _is_dashscope_host(host: str) -> bool: + if host in _STANDARD_DASHSCOPE_HOSTS: + return True + return host.endswith(_WORKSPACE_HOST_SUFFIX) and not host.startswith(_EXCLUDED_WORKSPACE_PREFIXES) + + +def _dashscope_origin(api_base: str | None) -> str | None: + if not api_base: + return None + parsed: Final = urlparse(api_base) + if parsed.scheme not in ("http", "https") or not parsed.netloc: + return None + if parsed.query or parsed.fragment: + return None + host: Final = (parsed.hostname or "").lower() + if not _is_dashscope_host(host): + return None + path: Final = parsed.path.rstrip("/") + if path and not any(path == prefix or path.startswith(f"{prefix}/") for prefix in _PROTOCOL_PATH_PREFIXES): + return None + return f"{parsed.scheme}://{parsed.netloc}" + + +def resolve_dashscope_endpoints(api_base: str | None) -> DashScopeEndpoints | None: + """ + Derive DashScope's protocol-specific bases from any standard DashScope base. + + Accepts the domain root, an OpenAI-compatible base, an Anthropic Messages + base, or a native `/api/v1` base for `dashscope.aliyuncs.com`, + `dashscope-intl.aliyuncs.com`, `dashscope-us.aliyuncs.com`, and + `{workspace}.{region}.maas.aliyuncs.com`. Returns None for any other host + or path so callers keep the caller-supplied base untouched. + """ + origin: Final = _dashscope_origin(api_base) + if origin is None: + return None + return DashScopeEndpoints( + origin=origin, + openai=f"{origin}{DASHSCOPE_CHAT_COMPATIBLE_PATH}", + anthropic=f"{origin}{DASHSCOPE_ANTHROPIC_PATH}", + native=f"{origin}{DASHSCOPE_NATIVE_V1_PATH}", + ) + + +def is_dashscope_native_endpoint_url(api_base: str) -> bool: + return urlparse(api_base).path.rstrip("/").startswith(f"{DASHSCOPE_NATIVE_V1_PATH}/") + + +def resolve_dashscope_openai_api_base( + api_base: str | None, + env_var: str = "DASHSCOPE_API_BASE", + default: str = f"{DASHSCOPE_DEFAULT_ORIGIN}{DASHSCOPE_CHAT_COMPATIBLE_PATH}", +) -> str: + candidate: Final = api_base or get_secret_str(env_var) + if candidate and is_dashscope_native_endpoint_url(candidate): + return candidate + endpoints: Final = resolve_dashscope_endpoints(candidate) + if endpoints is not None: + return endpoints.openai + return candidate or default + + +def resolve_dashscope_anthropic_api_base( + api_base: str | None, + env_var: str = "DASHSCOPE_API_BASE", + default: str = f"{DASHSCOPE_DEFAULT_ORIGIN}{DASHSCOPE_ANTHROPIC_PATH}", +) -> str: + candidate: Final = api_base or get_secret_str(env_var) + endpoints: Final = resolve_dashscope_endpoints(candidate) + if endpoints is not None: + return endpoints.anthropic + return candidate or default + + +def dashscope_base_selects_anthropic_messages(api_base: str | None) -> bool: + candidate: Final = api_base or get_secret_str("DASHSCOPE_API_BASE") + if not candidate: + return False + path: Final = urlparse(candidate).path.rstrip("/") + if path.endswith(DASHSCOPE_ANTHROPIC_PATH) or f"{DASHSCOPE_ANTHROPIC_PATH}/" in path: + return True + return not path and resolve_dashscope_endpoints(candidate) is not None + + +def _chat_shaped_aliyun_origin(api_base: str | None) -> str | None: if api_base is None: return None parsed: Final = urlparse(api_base) @@ -30,14 +144,24 @@ def _rerank_base_for_chat_shaped_base(api_base: str | None) -> str | None: on_aliyun_host: Final = host == "aliyuncs.com" or host.endswith(".aliyuncs.com") if not on_aliyun_host or parsed.path.rstrip("/") != DASHSCOPE_CHAT_COMPATIBLE_PATH: return None - return f"{parsed.scheme}://{parsed.netloc}{DASHSCOPE_RERANK_PATH}" + return f"{parsed.scheme}://{parsed.netloc}" + + +def _is_compatible_api_base(api_base: str) -> bool: + path: Final = urlparse(api_base).path.rstrip("/") + return path == DASHSCOPE_COMPATIBLE_API_PATH or path.startswith(f"{DASHSCOPE_COMPATIBLE_API_PATH}/") def resolve_dashscope_family_rerank_api_base(api_base: str | None, env_var: str, default_rerank_base: str) -> str: - remapped: Final = _rerank_base_for_chat_shaped_base(api_base) - if api_base is not None and remapped is None: + if api_base is None: + return get_secret_str(env_var) or default_rerank_base + if _is_compatible_api_base(api_base): return api_base - return get_secret_str(env_var) or remapped or default_rerank_base + chat_origin: Final = _chat_shaped_aliyun_origin(api_base) + if chat_origin is not None: + return get_secret_str(env_var) or f"{chat_origin}{DASHSCOPE_RERANK_PATH}" + endpoints: Final = resolve_dashscope_endpoints(api_base) + return f"{endpoints.origin}{DASHSCOPE_RERANK_PATH}" if endpoints is not None else api_base def get_dashscope_family_embedding_config(custom_llm_provider: str) -> "BaseEmbeddingConfig": diff --git a/litellm/llms/dashscope/embed/transformation.py b/litellm/llms/dashscope/embed/transformation.py index 35132346261..0f01a4b3172 100644 --- a/litellm/llms/dashscope/embed/transformation.py +++ b/litellm/llms/dashscope/embed/transformation.py @@ -25,7 +25,7 @@ from litellm.secret_managers.main import get_secret_str from litellm.types.llms.openai import AllEmbeddingInputValues, AllMessageValues from litellm.types.utils import EmbeddingResponse, Usage -from ..common_utils import DashScopeError +from ..common_utils import DashScopeError, resolve_dashscope_openai_api_base DEFAULT_API_BASE: Final = "https://dashscope.aliyuncs.com/compatible-mode/v1" @@ -91,7 +91,7 @@ class DashScopeEmbeddingConfig(BaseEmbeddingConfig): return resolved_api_key def _resolve_embedding_api_base(self, api_base: str | None) -> str: - return api_base or get_secret_str("DASHSCOPE_API_BASE") or DEFAULT_API_BASE + return resolve_dashscope_openai_api_base(api_base, default=DEFAULT_API_BASE) def validate_environment( self, diff --git a/litellm/llms/dashscope/image_generation/transformation.py b/litellm/llms/dashscope/image_generation/transformation.py index d1faa6c6cd0..2f528b5a2c3 100644 --- a/litellm/llms/dashscope/image_generation/transformation.py +++ b/litellm/llms/dashscope/image_generation/transformation.py @@ -39,6 +39,13 @@ from litellm.types.llms.openai import ( ) from litellm.types.utils import ImageObject, ImageResponse +from ..common_utils import ( + DASHSCOPE_CHAT_COMPATIBLE_PATH, + DASHSCOPE_IMAGE_PATH, + is_dashscope_native_endpoint_url, + resolve_dashscope_endpoints, +) + if TYPE_CHECKING: from litellm.litellm_core_utils.litellm_logging import Logging as _LiteLLMLoggingObj from litellm.litellm_core_utils.tokenizer import Encoding as Tokenizer @@ -53,8 +60,6 @@ _DECODED_JSON: Final = TypeAdapter(object) DEFAULT_API_BASE: Final = "https://dashscope-intl.aliyuncs.com/api/v1/services/aigc/multimodal-generation/generation" -CHAT_COMPATIBLE_MODE_PATH: Final = "/compatible-mode/v1" - # Maps OpenAI size strings (WxH) to DashScope size strings (W*H) OPENAI_TO_DASHSCOPE_SIZE: Final[dict] = { "256x256": "256*256", @@ -102,8 +107,14 @@ class DashScopeImageGenerationConfig(BaseImageGenerationConfig): raise ValueError("DASHSCOPE_API_KEY is not set") return resolved_api_key - def _resolve_image_api_base(self, image_api_base: str | None) -> str: - return image_api_base or get_secret_str("DASHSCOPE_API_BASE_IMAGE") or DEFAULT_API_BASE + def _image_api_base_env(self) -> str | None: + return get_secret_str("DASHSCOPE_API_BASE_IMAGE") + + def _default_image_api_base(self) -> str: + return DEFAULT_API_BASE + + def _configured_chat_api_base(self) -> str | None: + return get_secret_str("DASHSCOPE_API_BASE") def get_complete_url( self, @@ -114,10 +125,18 @@ class DashScopeImageGenerationConfig(BaseImageGenerationConfig): litellm_params: dict, stream: bool | None = None, ) -> str: - image_api_base: Final = ( - api_base if api_base and not api_base.rstrip("/").endswith(CHAT_COMPATIBLE_MODE_PATH) else None - ) - return self._resolve_image_api_base(image_api_base) + image_env_api_base: Final = self._image_api_base_env() + configured_api_base: Final = api_base or self._configured_chat_api_base() + if not configured_api_base or configured_api_base.rstrip("/").endswith(DASHSCOPE_CHAT_COMPATIBLE_PATH): + return image_env_api_base or self._default_image_api_base() + if is_dashscope_native_endpoint_url(configured_api_base): + return configured_api_base + endpoints: Final = resolve_dashscope_endpoints(configured_api_base) + if endpoints is None: + return configured_api_base + if not api_base and image_env_api_base: + return image_env_api_base + return f"{endpoints.origin}{DASHSCOPE_IMAGE_PATH}" def validate_environment( self, diff --git a/litellm/llms/dashscope/messages/__init__.py b/litellm/llms/dashscope/messages/__init__.py new file mode 100644 index 00000000000..08cd184404c --- /dev/null +++ b/litellm/llms/dashscope/messages/__init__.py @@ -0,0 +1,7 @@ +""" +DashScope Anthropic-compatible Messages Module +""" + +from .transformation import DashScopeAnthropicMessagesConfig + +__all__ = ["DashScopeAnthropicMessagesConfig"] diff --git a/litellm/llms/dashscope/messages/transformation.py b/litellm/llms/dashscope/messages/transformation.py new file mode 100644 index 00000000000..9d0d1db91c8 --- /dev/null +++ b/litellm/llms/dashscope/messages/transformation.py @@ -0,0 +1,96 @@ +""" +DashScope Anthropic-compatible Messages transformation config. +""" + +from typing import ClassVar, Final + +import litellm +from litellm.exceptions import AuthenticationError +from litellm.litellm_core_utils.url_utils import ensure_api_route +from litellm.llms.anthropic.messages.native.transformation import ( + AnthropicMessagesConfig, +) +from litellm.secret_managers.main import get_secret_str, normalize_nonempty_secret_str + +from ..common_utils import ( + dashscope_base_selects_anthropic_messages, + resolve_dashscope_anthropic_api_base, +) + + +class DashScopeAnthropicMessagesConfig(AnthropicMessagesConfig): + """ + DashScope exposes an Anthropic-compatible Messages API at + ``{origin}/apps/anthropic`` for every DashScope domain, including + ``dashscope.aliyuncs.com`` and ``{workspace}.{region}.maas.aliyuncs.com``. + """ + + _workload_identity_eligible: ClassVar[bool] = False + + @property + def custom_llm_provider(self) -> str | None: + return "dashscope" + + def should_strip_billing_metadata(self) -> bool: + return True + + def handles_web_search_natively(self) -> bool: + return False + + def serves_api_base(self, model: str, api_base: str | None) -> bool: + return dashscope_base_selects_anthropic_messages(api_base) + + @staticmethod + def get_api_key(api_key: str | None = None) -> str | None: + return ( + normalize_nonempty_secret_str(api_key) + or normalize_nonempty_secret_str(get_secret_str("DASHSCOPE_API_KEY")) + or normalize_nonempty_secret_str(litellm.api_key) + ) + + @staticmethod + def get_api_base(api_base: str | None = None) -> str: + return resolve_dashscope_anthropic_api_base(api_base) + + def validate_anthropic_messages_environment( + self, + headers: dict[str, str], # mutable-ok: mirrors the legacy base override signature + model: str, + messages: list[object], # mutable-ok: mirrors the legacy base override signature + optional_params: dict[str, object], # mutable-ok: mirrors the legacy base override signature + litellm_params: dict[str, object], # mutable-ok: mirrors the legacy base override signature + api_key: str | None = None, + api_base: str | None = None, + ) -> tuple[dict[str, str], str | None]: # mutable-ok: mirrors the base signature + resolved_api_key: Final = self.get_api_key(api_key=api_key) + if resolved_api_key is None: + present_headers: Final = {name.lower() for name in headers} + if "x-api-key" not in present_headers and "authorization" not in present_headers: + raise AuthenticationError( + message=( + "Missing DashScope API key. Set DASHSCOPE_API_KEY or pass api_key explicitly. " + "Anthropic credentials are never sent to a DashScope endpoint." + ), + llm_provider="dashscope", + model=model, + ) + return super().validate_anthropic_messages_environment( # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType, reportReturnType] # base signature uses an unparameterized dict + headers=headers, + model=model, + messages=messages, + optional_params=optional_params, + litellm_params=litellm_params, + api_key=resolved_api_key, + api_base=api_base, + ) + + def get_complete_url( + self, + api_base: str | None, + api_key: str | None, + model: str, + optional_params: dict[str, object], # mutable-ok: mirrors the base get_complete_url signature + litellm_params: dict[str, object], # mutable-ok: mirrors the base get_complete_url signature + stream: bool | None = None, + ) -> str: + return ensure_api_route(self.get_api_base(api_base=api_base), "v1/messages") diff --git a/litellm/llms/dashscope/qwen_ai_platform.py b/litellm/llms/dashscope/qwen_ai_platform.py index f7cc50cc82a..db01c8a0a15 100644 --- a/litellm/llms/dashscope/qwen_ai_platform.py +++ b/litellm/llms/dashscope/qwen_ai_platform.py @@ -3,7 +3,7 @@ from typing import Final from litellm.secret_managers.main import get_secret_str from .chat.transformation import DashScopeChatConfig -from .common_utils import resolve_dashscope_family_rerank_api_base +from .common_utils import resolve_dashscope_family_rerank_api_base, resolve_dashscope_openai_api_base from .embed.transformation import DashScopeEmbeddingConfig from .image_generation.transformation import DashScopeImageGenerationConfig from .rerank.transformation import DashScopeRerankConfig @@ -38,7 +38,7 @@ class QwenAIPlatformChatConfig(DashScopeChatConfig): return self._resolve_chat_api_base(api_base), _resolve_qwen_ai_platform_api_key(api_key) def _resolve_chat_api_base(self, api_base: str | None) -> str: - return api_base or get_secret_str("QWEN_AI_PLATFORM_API_BASE") or QWEN_AI_PLATFORM_API_BASE + return resolve_dashscope_openai_api_base(api_base, "QWEN_AI_PLATFORM_API_BASE", QWEN_AI_PLATFORM_API_BASE) class QwenAIPlatformEmbeddingConfig(DashScopeEmbeddingConfig): @@ -46,7 +46,7 @@ class QwenAIPlatformEmbeddingConfig(DashScopeEmbeddingConfig): return _require_qwen_ai_platform_api_key(api_key) def _resolve_embedding_api_base(self, api_base: str | None) -> str: - return api_base or get_secret_str("QWEN_AI_PLATFORM_API_BASE") or QWEN_AI_PLATFORM_API_BASE + return resolve_dashscope_openai_api_base(api_base, "QWEN_AI_PLATFORM_API_BASE", QWEN_AI_PLATFORM_API_BASE) class QwenAIPlatformRerankConfig(DashScopeRerankConfig): @@ -63,8 +63,14 @@ class QwenAIPlatformImageGenerationConfig(DashScopeImageGenerationConfig): def _resolve_api_key(self, api_key: str | None) -> str: return _require_qwen_ai_platform_api_key(api_key) - def _resolve_image_api_base(self, image_api_base: str | None) -> str: - return image_api_base or get_secret_str("QWEN_AI_PLATFORM_API_BASE_IMAGE") or QWEN_AI_PLATFORM_IMAGE_API_BASE + def _image_api_base_env(self) -> str | None: + return get_secret_str("QWEN_AI_PLATFORM_API_BASE_IMAGE") + + def _default_image_api_base(self) -> str: + return QWEN_AI_PLATFORM_IMAGE_API_BASE + + def _configured_chat_api_base(self) -> str | None: + return get_secret_str("QWEN_AI_PLATFORM_API_BASE") class QwenAIPlatformVideoConfig(DashScopeVideoConfig): diff --git a/litellm/llms/dashscope/qwencloud.py b/litellm/llms/dashscope/qwencloud.py index a12e21b0983..6ae42c62e42 100644 --- a/litellm/llms/dashscope/qwencloud.py +++ b/litellm/llms/dashscope/qwencloud.py @@ -3,7 +3,7 @@ from typing import Final from litellm.secret_managers.main import get_secret_str from .chat.transformation import DashScopeChatConfig -from .common_utils import resolve_dashscope_family_rerank_api_base +from .common_utils import resolve_dashscope_family_rerank_api_base, resolve_dashscope_openai_api_base from .embed.transformation import DashScopeEmbeddingConfig from .image_generation.transformation import DashScopeImageGenerationConfig from .rerank.transformation import DashScopeRerankConfig @@ -38,7 +38,7 @@ class QwenCloudChatConfig(DashScopeChatConfig): return self._resolve_chat_api_base(api_base), _resolve_qwencloud_api_key(api_key) def _resolve_chat_api_base(self, api_base: str | None) -> str: - return api_base or get_secret_str("QWENCLOUD_API_BASE") or QWENCLOUD_API_BASE + return resolve_dashscope_openai_api_base(api_base, "QWENCLOUD_API_BASE", QWENCLOUD_API_BASE) class QwenCloudEmbeddingConfig(DashScopeEmbeddingConfig): @@ -46,7 +46,7 @@ class QwenCloudEmbeddingConfig(DashScopeEmbeddingConfig): return _require_qwencloud_api_key(api_key) def _resolve_embedding_api_base(self, api_base: str | None) -> str: - return api_base or get_secret_str("QWENCLOUD_API_BASE") or QWENCLOUD_API_BASE + return resolve_dashscope_openai_api_base(api_base, "QWENCLOUD_API_BASE", QWENCLOUD_API_BASE) class QwenCloudRerankConfig(DashScopeRerankConfig): @@ -63,8 +63,14 @@ class QwenCloudImageGenerationConfig(DashScopeImageGenerationConfig): def _resolve_api_key(self, api_key: str | None) -> str: return _require_qwencloud_api_key(api_key) - def _resolve_image_api_base(self, image_api_base: str | None) -> str: - return image_api_base or get_secret_str("QWENCLOUD_API_BASE_IMAGE") or QWENCLOUD_IMAGE_API_BASE + def _image_api_base_env(self) -> str | None: + return get_secret_str("QWENCLOUD_API_BASE_IMAGE") + + def _default_image_api_base(self) -> str: + return QWENCLOUD_IMAGE_API_BASE + + def _configured_chat_api_base(self) -> str | None: + return get_secret_str("QWENCLOUD_API_BASE") class QwenCloudVideoConfig(DashScopeVideoConfig): diff --git a/litellm/llms/dashscope/videos/transformation.py b/litellm/llms/dashscope/videos/transformation.py index 5c738423a5a..59e42e7b926 100644 --- a/litellm/llms/dashscope/videos/transformation.py +++ b/litellm/llms/dashscope/videos/transformation.py @@ -38,6 +38,8 @@ from litellm.types.videos.utils import ( ) from litellm.utils import get_model_info +from ..common_utils import resolve_dashscope_endpoints + if TYPE_CHECKING: from litellm.litellm_core_utils.litellm_logging import Logging from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler @@ -206,6 +208,9 @@ def _parse_task_response(raw_response: httpx.Response) -> _TaskResponse: def _normalized_api_base(api_base: str) -> str: + endpoints: Final = resolve_dashscope_endpoints(api_base) + if endpoints is not None: + return endpoints.origin trimmed: Final = api_base.rstrip("/") if trimmed.endswith(DASHSCOPE_COMPATIBLE_MODE_PATH): return trimmed[: -len(DASHSCOPE_COMPATIBLE_MODE_PATH)] diff --git a/litellm/rerank_api/main.py b/litellm/rerank_api/main.py index 65d9b89cf79..e344ce46be7 100644 --- a/litellm/rerank_api/main.py +++ b/litellm/rerank_api/main.py @@ -19,6 +19,12 @@ from litellm.types.rerank import RerankResponse from litellm.types.router import * from litellm.utils import ProviderConfigManager, client, exception_type +_PROVIDERS_RESOLVING_THEIR_OWN_RERANK_BASE: Final = ( + litellm.LlmProviders.DASHSCOPE, + litellm.LlmProviders.QWENCLOUD, + litellm.LlmProviders.QWEN_AI_PLATFORM, +) + ####### ENVIRONMENT VARIABLES ################### # Initialize any necessary instances or variables here together_rerank: Final = TogetherAIRerank() @@ -543,7 +549,11 @@ def rerank( logging_obj=litellm_logging_obj, timeout=optional_params.timeout, api_key=dynamic_api_key or optional_params.api_key, - api_base=dynamic_api_base or optional_params.api_base, + api_base=( + optional_params.api_base or dynamic_api_base + if _custom_llm_provider in _PROVIDERS_RESOLVING_THEIR_OWN_RERANK_BASE + else dynamic_api_base or optional_params.api_base + ), _is_async=_is_async, headers=headers or litellm.headers or {}, client=client, diff --git a/litellm/utils.py b/litellm/utils.py index 8ea4f6ceda0..4cd2335dd0c 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -10048,6 +10048,12 @@ class ProviderConfigManager: ) return TencentAnthropicMessagesConfig() + elif litellm.LlmProviders.DASHSCOPE == provider: + from litellm.llms.dashscope.messages.transformation import ( + DashScopeAnthropicMessagesConfig, + ) + + return DashScopeAnthropicMessagesConfig() elif litellm.LlmProviders.GITHUB_COPILOT == provider: if "claude" in model_lower: from litellm.llms.github_copilot.messages.transformation import ( diff --git a/tests/unit/integrations/websearch_interception/test_websearch_short_circuit.py b/tests/unit/integrations/websearch_interception/test_websearch_short_circuit.py index 414bcc11602..ef050d96dc5 100644 --- a/tests/unit/integrations/websearch_interception/test_websearch_short_circuit.py +++ b/tests/unit/integrations/websearch_interception/test_websearch_short_circuit.py @@ -32,6 +32,7 @@ class TestTryShortCircuitSearch: ("github_copilot", "github_copilot/claude-sonnet-5.5"), ("databricks", "databricks/databricks-claude-sonnet-5-5"), ("databricks", "databricks/system.ai.claude-sonnet-5-5"), + ("dashscope", "dashscope/qwen3.7-plus"), ], ) async def test_short_circuits_single_web_search_tool(self, provider, model): diff --git a/tests/unit/llms/dashscope/test_dashscope_chat_transformation.py b/tests/unit/llms/dashscope/test_dashscope_chat_transformation.py index a3fd5933a83..99a8004b0a1 100644 --- a/tests/unit/llms/dashscope/test_dashscope_chat_transformation.py +++ b/tests/unit/llms/dashscope/test_dashscope_chat_transformation.py @@ -5,15 +5,50 @@ These tests validate the DashScopeConfig class which extends OpenAIGPTConfig. DashScope is an OpenAI-compatible provider with minor customizations. """ +from typing import Final +import pytest from litellm.types.llms.openai import AllMessageValues -import pytest import litellm from litellm import completion from litellm.llms.dashscope.chat.transformation import DashScopeChatConfig +WORKSPACE: Final = "https://llm-abc123.cn-beijing.maas.aliyuncs.com" + + +def test_get_complete_url_derives_the_compatible_base_from_any_standard_base(): + config: Final = DashScopeChatConfig() + + for api_base in ( + WORKSPACE, + f"{WORKSPACE}/compatible-mode/v1", + f"{WORKSPACE}/apps/anthropic", + f"{WORKSPACE}/api/v1", + ): + url: Final = config.get_complete_url( + api_base=api_base, + api_key="sk-test", + model="qwen3.7-plus", + optional_params={}, + litellm_params={}, + ) + + assert url == f"{WORKSPACE}/compatible-mode/v1/chat/completions" + + +def test_get_complete_url_keeps_unrecognized_hosts(): + url: Final = DashScopeChatConfig().get_complete_url( + api_base="https://gateway.example.com/v1", + api_key="sk-test", + model="qwen3.7-plus", + optional_params={}, + litellm_params={}, + ) + + assert url == "https://gateway.example.com/v1/chat/completions" + class TestDashScopeConfig: """Test class for DashScope functionality""" diff --git a/tests/unit/llms/dashscope/test_dashscope_common_utils.py b/tests/unit/llms/dashscope/test_dashscope_common_utils.py new file mode 100644 index 00000000000..2dc4d5e5a85 --- /dev/null +++ b/tests/unit/llms/dashscope/test_dashscope_common_utils.py @@ -0,0 +1,171 @@ +""" +Unit tests for DashScope's protocol-specific base URL derivation. +""" + +from typing import Final + +import pytest + +from litellm.llms.dashscope.common_utils import ( + DASHSCOPE_DEFAULT_ORIGIN, + resolve_dashscope_anthropic_api_base, + resolve_dashscope_endpoints, + resolve_dashscope_family_rerank_api_base, + resolve_dashscope_openai_api_base, +) + +WORKSPACE: Final = "https://llm-abc123.cn-beijing.maas.aliyuncs.com" + + +@pytest.mark.parametrize( + "api_base", + [ + WORKSPACE, + f"{WORKSPACE}/", + f"{WORKSPACE}/compatible-mode/v1", + f"{WORKSPACE}/compatible-mode/v1/chat/completions", + f"{WORKSPACE}/apps/anthropic", + f"{WORKSPACE}/apps/anthropic/v1/messages", + f"{WORKSPACE}/api/v1", + f"{WORKSPACE}/api/v1/services/aigc/video-generation/video-synthesis", + ], +) +def test_workspace_base_derives_every_protocol(api_base: str): + endpoints: Final = resolve_dashscope_endpoints(api_base) + + assert endpoints is not None + assert endpoints.origin == WORKSPACE + assert endpoints.openai == f"{WORKSPACE}/compatible-mode/v1" + assert endpoints.anthropic == f"{WORKSPACE}/apps/anthropic" + assert endpoints.native == f"{WORKSPACE}/api/v1" + + +@pytest.mark.parametrize( + "origin", + [ + "https://dashscope.aliyuncs.com", + "https://dashscope-intl.aliyuncs.com", + "https://dashscope-us.aliyuncs.com", + ], +) +@pytest.mark.parametrize("path", ["", "/", "/compatible-mode/v1"]) +def test_standard_hosts_derive_every_protocol(origin: str, path: str): + endpoints: Final = resolve_dashscope_endpoints(f"{origin}{path}") + + assert endpoints is not None + assert endpoints.openai == f"{origin}/compatible-mode/v1" + assert endpoints.anthropic == f"{origin}/apps/anthropic" + assert endpoints.native == f"{origin}/api/v1" + + +@pytest.mark.parametrize( + "api_base", + [ + None, + "", + "not-a-url", + "https://gateway.example.com", + "https://gateway.example.com/compatible-mode/v1", + f"{WORKSPACE}/custom-path", + f"{WORKSPACE}/compatible-mode/v1?api-version=1", + "https://trial.cn-beijing.maas.aliyuncs.com", + "https://token-plan.cn-beijing.maas.aliyuncs.com", + ], +) +def test_unrecognized_bases_are_not_rewritten(api_base: str | None): + assert resolve_dashscope_endpoints(api_base) is None + + +def test_openai_base_derives_from_the_domain(monkeypatch: pytest.MonkeyPatch): + monkeypatch.delenv("DASHSCOPE_API_BASE", raising=False) + + assert resolve_dashscope_openai_api_base(WORKSPACE) == f"{WORKSPACE}/compatible-mode/v1" + assert resolve_dashscope_openai_api_base(f"{WORKSPACE}/apps/anthropic") == f"{WORKSPACE}/compatible-mode/v1" + assert ( + resolve_dashscope_openai_api_base(f"{WORKSPACE}/compatible-api/v1/reranks") == f"{WORKSPACE}/compatible-mode/v1" + ) + + +@pytest.mark.parametrize( + "endpoint_url", + [ + "https://dashscope-intl.aliyuncs.com/api/v1/services/aigc/multimodal-generation/generation", + f"{WORKSPACE}/api/v1/services/aigc/video-generation/video-synthesis", + ], +) +def test_openai_base_keeps_a_native_endpoint_url(monkeypatch: pytest.MonkeyPatch, endpoint_url: str): + monkeypatch.setenv("DASHSCOPE_API_BASE", endpoint_url) + + assert resolve_dashscope_openai_api_base(None) == endpoint_url + assert resolve_dashscope_openai_api_base(endpoint_url) == endpoint_url + assert resolve_dashscope_openai_api_base(f"{WORKSPACE}/api/v1") == f"{WORKSPACE}/compatible-mode/v1" + + +def test_openai_base_keeps_unrecognized_hosts(monkeypatch: pytest.MonkeyPatch): + monkeypatch.delenv("DASHSCOPE_API_BASE", raising=False) + + assert resolve_dashscope_openai_api_base("https://gateway.example.com/v1") == "https://gateway.example.com/v1" + + +def test_openai_base_defaults_to_dashscope(monkeypatch: pytest.MonkeyPatch): + monkeypatch.delenv("DASHSCOPE_API_BASE", raising=False) + + assert resolve_dashscope_openai_api_base(None) == f"{DASHSCOPE_DEFAULT_ORIGIN}/compatible-mode/v1" + + +def test_anthropic_base_derives_from_the_domain(monkeypatch: pytest.MonkeyPatch): + monkeypatch.delenv("DASHSCOPE_API_BASE", raising=False) + + assert resolve_dashscope_anthropic_api_base(WORKSPACE) == f"{WORKSPACE}/apps/anthropic" + assert resolve_dashscope_anthropic_api_base(f"{WORKSPACE}/compatible-mode/v1") == f"{WORKSPACE}/apps/anthropic" + + +def test_anthropic_base_keeps_unrecognized_hosts(monkeypatch: pytest.MonkeyPatch): + monkeypatch.delenv("DASHSCOPE_API_BASE", raising=False) + + assert ( + resolve_dashscope_anthropic_api_base("https://gateway.example.com/anthropic") + == "https://gateway.example.com/anthropic" + ) + + +def test_rerank_base_derives_from_the_domain(monkeypatch: pytest.MonkeyPatch): + monkeypatch.delenv("DASHSCOPE_API_BASE_RERANK", raising=False) + + assert ( + resolve_dashscope_family_rerank_api_base(WORKSPACE, "DASHSCOPE_API_BASE_RERANK", "https://default/reranks") + == f"{WORKSPACE}/compatible-api/v1/reranks" + ) + + +def test_rerank_base_keeps_unrecognized_hosts(): + gateway: Final = "https://gateway.example.com/v1" + + assert ( + resolve_dashscope_family_rerank_api_base(gateway, "DASHSCOPE_API_BASE_RERANK", "https://default/reranks") + == gateway + ) + + +def test_rerank_base_uses_the_configured_default(monkeypatch: pytest.MonkeyPatch): + monkeypatch.delenv("DASHSCOPE_API_BASE_RERANK", raising=False) + + assert ( + resolve_dashscope_family_rerank_api_base(None, "DASHSCOPE_API_BASE_RERANK", "https://default/reranks") + == "https://default/reranks" + ) + + +@pytest.mark.parametrize( + "origin", + ["https://cn-hongkong.dashscope.aliyuncs.com", "https://token-plan.cn-beijing.maas.aliyuncs.com"], +) +def test_rerank_keeps_remapping_chat_bases_on_other_aliyun_hosts(monkeypatch: pytest.MonkeyPatch, origin: str): + monkeypatch.delenv("DASHSCOPE_API_BASE_RERANK", raising=False) + + assert ( + resolve_dashscope_family_rerank_api_base( + f"{origin}/compatible-mode/v1", "DASHSCOPE_API_BASE_RERANK", "https://default/reranks" + ) + == f"{origin}/compatible-api/v1/reranks" + ) diff --git a/tests/unit/llms/dashscope/test_dashscope_embedding_transformation.py b/tests/unit/llms/dashscope/test_dashscope_embedding_transformation.py index 1b6eea0e4c8..27afbee3173 100644 --- a/tests/unit/llms/dashscope/test_dashscope_embedding_transformation.py +++ b/tests/unit/llms/dashscope/test_dashscope_embedding_transformation.py @@ -3,6 +3,7 @@ Unit tests for DashScope embedding transformation. """ import json +from typing import Final from unittest.mock import MagicMock import httpx @@ -16,6 +17,40 @@ from litellm.llms.dashscope.embed.transformation import ( ) from litellm.types.utils import EmbeddingResponse +WORKSPACE: Final = "https://llm-abc123.cn-beijing.maas.aliyuncs.com" + + +def test_get_complete_url_derives_the_compatible_base_from_any_standard_base(): + config: Final = DashScopeEmbeddingConfig() + + for api_base in ( + WORKSPACE, + f"{WORKSPACE}/compatible-mode/v1", + f"{WORKSPACE}/apps/anthropic", + f"{WORKSPACE}/api/v1", + ): + url: Final = config.get_complete_url( + api_base=api_base, + api_key="sk-test", + model="text-embedding-v4", + optional_params={}, + litellm_params={}, + ) + + assert url == f"{WORKSPACE}/compatible-mode/v1/embeddings" + + +def test_get_complete_url_keeps_unrecognized_hosts(): + url: Final = DashScopeEmbeddingConfig().get_complete_url( + api_base="https://gateway.example.com/v1", + api_key="sk-test", + model="text-embedding-v4", + optional_params={}, + litellm_params={}, + ) + + assert url == "https://gateway.example.com/v1/embeddings" + def test_validate_environment_and_url(): config = DashScopeEmbeddingConfig() diff --git a/tests/unit/llms/dashscope/test_dashscope_messages_transformation.py b/tests/unit/llms/dashscope/test_dashscope_messages_transformation.py new file mode 100644 index 00000000000..decc12afd87 --- /dev/null +++ b/tests/unit/llms/dashscope/test_dashscope_messages_transformation.py @@ -0,0 +1,215 @@ +""" +Unit tests for DashScope's Anthropic-compatible Messages endpoint. +""" + +from typing import Final + +import pytest +import respx + +import litellm +from litellm.exceptions import AuthenticationError +from litellm.llms.dashscope.messages.transformation import ( + DashScopeAnthropicMessagesConfig, +) + +WORKSPACE: Final = "https://llm-abc123.cn-beijing.maas.aliyuncs.com" +INTL: Final = "https://dashscope-intl.aliyuncs.com" +CODING_PLAN: Final = "https://coding-intl.dashscope.aliyuncs.com" +CHAT_COMPLETION: Final = { + "id": "chatcmpl-1", + "object": "chat.completion", + "created": 1792227137, + "model": "qwen3.7-plus", + "choices": [{"index": 0, "message": {"role": "assistant", "content": "ok"}, "finish_reason": "stop"}], + "usage": {"prompt_tokens": 9, "completion_tokens": 1, "total_tokens": 10}, +} +MESSAGE: Final = { + "id": "msg_1", + "type": "message", + "role": "assistant", + "model": "qwen3.7-plus", + "content": [{"type": "text", "text": "ok"}], + "stop_reason": "end_turn", + "stop_sequence": None, + "usage": {"input_tokens": 9, "output_tokens": 1}, +} + + +@pytest.mark.parametrize( + ("api_base", "expected"), + [ + ( + None, + "https://dashscope.aliyuncs.com/apps/anthropic/v1/messages", + ), + ( + "https://dashscope.aliyuncs.com", + "https://dashscope.aliyuncs.com/apps/anthropic/v1/messages", + ), + ( + "https://dashscope.aliyuncs.com/compatible-mode/v1", + "https://dashscope.aliyuncs.com/apps/anthropic/v1/messages", + ), + ( + "https://dashscope.aliyuncs.com/apps/anthropic", + "https://dashscope.aliyuncs.com/apps/anthropic/v1/messages", + ), + ( + "https://dashscope.aliyuncs.com/apps/anthropic/v1/messages", + "https://dashscope.aliyuncs.com/apps/anthropic/v1/messages", + ), + ( + WORKSPACE, + f"{WORKSPACE}/apps/anthropic/v1/messages", + ), + ( + f"{WORKSPACE}/compatible-mode/v1", + f"{WORKSPACE}/apps/anthropic/v1/messages", + ), + ( + "https://gateway.example.com/anthropic", + "https://gateway.example.com/anthropic/v1/messages", + ), + ], +) +def test_get_complete_url(api_base: str | None, expected: str): + config: Final = DashScopeAnthropicMessagesConfig() + + url: Final = config.get_complete_url( + api_base=api_base, + api_key="sk-test", + model="qwen3.7-plus", + optional_params={}, + litellm_params={}, + ) + + assert url == expected + + +@pytest.mark.parametrize("dashscope_api_key", [None, " "]) +def test_missing_dashscope_key_does_not_fall_back_to_anthropic( + monkeypatch: pytest.MonkeyPatch, dashscope_api_key: str | None +): + monkeypatch.delenv("DASHSCOPE_API_KEY", raising=False) + if dashscope_api_key is not None: + monkeypatch.setenv("DASHSCOPE_API_KEY", dashscope_api_key) + monkeypatch.delenv("ANTHROPIC_AUTH_TOKEN", raising=False) + monkeypatch.setenv("ANTHROPIC_API_KEY", "sk-ant-secret") + monkeypatch.setattr(litellm, "api_key", None) + headers: dict[str, str] = {} + + with pytest.raises(AuthenticationError) as error: + DashScopeAnthropicMessagesConfig().validate_anthropic_messages_environment( + headers=headers, + model="qwen3.7-plus", + messages=[], + optional_params={}, + litellm_params={}, + api_key=None, + api_base=WORKSPACE, + ) + + assert "DashScope" in str(error.value) + assert "sk-ant-secret" not in str(error.value) + assert {name.lower() for name in headers}.isdisjoint({"x-api-key", "authorization"}) + + +def test_caller_auth_header_is_honored_without_a_dashscope_key(monkeypatch: pytest.MonkeyPatch): + monkeypatch.delenv("DASHSCOPE_API_KEY", raising=False) + monkeypatch.setattr(litellm, "api_key", None) + headers: Final = {"x-api-key": "caller-key"} + + resolved_headers, _ = DashScopeAnthropicMessagesConfig().validate_anthropic_messages_environment( + headers=headers, + model="qwen3.7-plus", + messages=[], + optional_params={}, + litellm_params={}, + api_key=None, + api_base=WORKSPACE, + ) + + assert resolved_headers["x-api-key"] == "caller-key" + + +def test_caller_authorization_header_is_honored_without_a_dashscope_key(monkeypatch: pytest.MonkeyPatch): + monkeypatch.delenv("DASHSCOPE_API_KEY", raising=False) + monkeypatch.setattr(litellm, "api_key", None) + headers: Final = {"Authorization": "Bearer caller-token"} + + resolved_headers, _ = DashScopeAnthropicMessagesConfig().validate_anthropic_messages_environment( + headers=headers, + model="qwen3.7-plus", + messages=[], + optional_params={}, + litellm_params={}, + api_key=None, + api_base=WORKSPACE, + ) + + assert resolved_headers["Authorization"] == "Bearer caller-token" + + +def test_dashscope_key_wins_over_an_anthropic_environment_key(monkeypatch: pytest.MonkeyPatch): + monkeypatch.setenv("DASHSCOPE_API_KEY", "sk-dashscope") + monkeypatch.setenv("ANTHROPIC_API_KEY", "sk-ant-secret") + + resolved_headers, _ = DashScopeAnthropicMessagesConfig().validate_anthropic_messages_environment( + headers={}, + model="qwen3.7-plus", + messages=[], + optional_params={}, + litellm_params={}, + api_key=None, + api_base=WORKSPACE, + ) + + header_values: Final = " ".join(resolved_headers.values()) + assert "sk-dashscope" in header_values + assert "sk-ant-secret" not in header_values + + +def test_dashscope_prefix_reaches_the_provider(): + model, provider, _, _ = litellm.get_llm_provider( + "dashscope/qwen3.7-plus", + api_base=WORKSPACE, + ) + + assert provider == "dashscope" + assert model == "qwen3.7-plus" + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("api_base", "expected_url"), + [ + (INTL, f"{INTL}/apps/anthropic/v1/messages"), + (f"{INTL}/apps/anthropic", f"{INTL}/apps/anthropic/v1/messages"), + (WORKSPACE, f"{WORKSPACE}/apps/anthropic/v1/messages"), + (f"{CODING_PLAN}/apps/anthropic", f"{CODING_PLAN}/apps/anthropic/v1/messages"), + (None, "https://dashscope.aliyuncs.com/compatible-mode/v1/chat/completions"), + (f"{INTL}/compatible-mode/v1", f"{INTL}/compatible-mode/v1/chat/completions"), + (f"{CODING_PLAN}/v1", f"{CODING_PLAN}/v1/chat/completions"), + ("https://gateway.example.com/anthropic", "https://gateway.example.com/anthropic/chat/completions"), + ], +) +async def test_only_a_domain_root_or_anthropic_base_leaves_the_chat_bridge( + monkeypatch: pytest.MonkeyPatch, api_base: str | None, expected_url: str +): + monkeypatch.delenv("DASHSCOPE_API_BASE", raising=False) + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + + with respx.mock(assert_all_called=False) as upstream: + upstream.route(path__regex=r".*/chat/completions$").respond(200, json=CHAT_COMPLETION) + upstream.route(path__regex=r".*/v1/messages$").respond(200, json=MESSAGE) + + await litellm.anthropic.messages.acreate( + model="dashscope/qwen3.7-plus", + messages=[{"role": "user", "content": "say ok"}], + max_tokens=16, + api_key="sk-dashscope", + api_base=api_base, + ) + + assert [str(call.request.url) for call in upstream.calls] == [expected_url] diff --git a/tests/unit/llms/dashscope/test_dashscope_rerank_transformation.py b/tests/unit/llms/dashscope/test_dashscope_rerank_transformation.py index a913c79af31..f0dbdb74b68 100644 --- a/tests/unit/llms/dashscope/test_dashscope_rerank_transformation.py +++ b/tests/unit/llms/dashscope/test_dashscope_rerank_transformation.py @@ -3,12 +3,15 @@ Unit tests for DashScope rerank transformation. """ import json +from typing import Final from unittest.mock import MagicMock import httpx import pytest +import respx from pydantic import ValidationError +import litellm from litellm.llms.dashscope.common_utils import DashScopeError from litellm.llms.dashscope.rerank.transformation import ( @@ -26,6 +29,21 @@ class TestDashScopeRerankURL: url = self.config.get_complete_url(api_base=None, model="qwen3-rerank") assert url == DEFAULT_RERANK_URL + def test_domain_root_remaps_to_rerank_route(self): + url: Final = self.config.get_complete_url( + api_base="https://llm-abc123.cn-beijing.maas.aliyuncs.com", + model="qwen3-rerank", + ) + assert url == "https://llm-abc123.cn-beijing.maas.aliyuncs.com/compatible-api/v1/reranks" + + def test_explicit_rerank_url_wins_over_the_rerank_env(self, monkeypatch: pytest.MonkeyPatch): + monkeypatch.setenv("DASHSCOPE_API_BASE_RERANK", DEFAULT_RERANK_URL) + explicit_url: Final = "https://dashscope-intl.aliyuncs.com/compatible-api/v1/reranks" + + url: Final = self.config.get_complete_url(api_base=explicit_url, model="qwen3-rerank") + + assert url == explicit_url + def test_chat_shaped_base_remaps_to_rerank_route(self): url = self.config.get_complete_url( api_base="https://dashscope.aliyuncs.com/compatible-mode/v1", @@ -434,3 +452,37 @@ def test_transform_rerank_response_shape_errors_do_not_echo_the_payload(): _transform({"results": [["leaked document text"]]}) assert "leaked document text" not in str(exc_info.value) + + +@pytest.mark.parametrize( + ("provider", "rerank_env"), + [ + ("dashscope", "DASHSCOPE_API_BASE_RERANK"), + ("qwencloud", "QWENCLOUD_API_BASE_RERANK"), + ("qwen_ai_platform", "QWEN_AI_PLATFORM_API_BASE_RERANK"), + ], +) +@pytest.mark.parametrize( + "api_base", + [ + "https://dashscope-intl.aliyuncs.com/compatible-api/v1/reranks", + "https://dashscope-intl.aliyuncs.com/compatible-api/v1", + "https://dashscope-intl.aliyuncs.com", + "https://dashscope-intl.aliyuncs.com/api/v1", + ], +) +def test_rerank_sends_an_explicit_base_over_the_rerank_env( + monkeypatch: pytest.MonkeyPatch, provider: str, rerank_env: str, api_base: str +): + monkeypatch.setenv("DASHSCOPE_API_BASE", "https://dashscope-intl.aliyuncs.com") + monkeypatch.setenv(rerank_env, DEFAULT_RERANK_URL) + + with respx.mock(assert_all_called=True) as upstream: + upstream.post("https://dashscope-intl.aliyuncs.com/compatible-api/v1/reranks").respond( + 200, json={"id": "r1", "results": [{"index": 0, "relevance_score": 0.9}], "usage": {"total_tokens": 5}} + ) + response: Final = litellm.rerank( + model=f"{provider}/qwen3-rerank", query="q", documents=["a"], api_base=api_base, api_key="sk-test" + ) + + assert response.results == [{"index": 0, "relevance_score": 0.9}] diff --git a/tests/unit/llms/dashscope/test_dashscope_video_transformation.py b/tests/unit/llms/dashscope/test_dashscope_video_transformation.py index 15b77aaee3d..c1b4157c8ab 100644 --- a/tests/unit/llms/dashscope/test_dashscope_video_transformation.py +++ b/tests/unit/llms/dashscope/test_dashscope_video_transformation.py @@ -1119,6 +1119,20 @@ class TestDashScopeVideoEnvironment: assert url == "https://ws-123.cn-beijing.maas.aliyuncs.com" + def test_every_protocol_base_normalizes_to_the_host(self): + workspace: Final = "https://ws-123.cn-beijing.maas.aliyuncs.com" + + for api_base in ( + f"{workspace}/compatible-mode/v1", + f"{workspace}/apps/anthropic", + f"{workspace}/api/v1", + ): + url: Final = DashScopeVideoConfig().get_complete_url( + model="wan3.0-video", api_base=api_base, litellm_params={} + ) + + assert url == workspace + def test_default_host_is_beijing(self, monkeypatch): monkeypatch.delenv("DASHSCOPE_API_BASE_VIDEO", raising=False) monkeypatch.delenv("DASHSCOPE_API_BASE", raising=False) diff --git a/tests/unit/llms/dashscope/test_qwen_brand_aliases.py b/tests/unit/llms/dashscope/test_qwen_brand_aliases.py index 69ae66d3818..43ef8f90b39 100644 --- a/tests/unit/llms/dashscope/test_qwen_brand_aliases.py +++ b/tests/unit/llms/dashscope/test_qwen_brand_aliases.py @@ -2,8 +2,10 @@ import json import math import re from pathlib import Path +from typing import Final, Literal import pytest +from typing_extensions import ReadOnly, TypedDict import litellm from litellm import completion, get_llm_provider @@ -53,6 +55,23 @@ DASHSCOPE_FAMILY_ENV_VARS = [ "QWEN_AI_PLATFORM_API_BASE_IMAGE", ] + + +class BrandCase(TypedDict): + provider: ReadOnly[str] + display_name: ReadOnly[str] + enum: ReadOnly[LlmProviders] + key_env: ReadOnly[str] + base_env: ReadOnly[str] + default_base: ReadOnly[str] + default_rerank_base: ReadOnly[str] + default_image_base: ReadOnly[str] + chat_config: ReadOnly[type[DashScopeChatConfig]] + embedding_config: ReadOnly[type[DashScopeEmbeddingConfig]] + rerank_config: ReadOnly[type[DashScopeRerankConfig]] + image_config: ReadOnly[type[DashScopeImageGenerationConfig]] + + BRAND_CASES = [ pytest.param( { @@ -211,6 +230,20 @@ class TestQwenBrandDefaultUrls: ) assert url == f"{brand['default_base']}/embeddings" + @pytest.mark.parametrize("brand", BRAND_CASES) + @pytest.mark.parametrize(("config_key", "route"), [("chat_config", "chat/completions"), ("embedding_config", "embeddings")]) + def test_openai_routes_derive_from_a_domain_root( + self, brand: BrandCase, config_key: Literal["chat_config", "embedding_config"], route: str + ) -> None: + url: Final = brand[config_key]().get_complete_url( + api_base="https://dashscope-intl.aliyuncs.com", + api_key="sk-test", + model="qwen-max", + optional_params={}, + litellm_params={}, + ) + assert url == f"https://dashscope-intl.aliyuncs.com/compatible-mode/v1/{route}" + @pytest.mark.parametrize("brand", BRAND_CASES) def test_rerank_complete_url(self, brand): url = brand["rerank_config"]().get_complete_url(api_base=None, model="gte-rerank-v2") @@ -249,6 +282,19 @@ class TestQwenBrandDefaultUrls: ) assert url == brand["default_image_base"] + @pytest.mark.parametrize("brand", BRAND_CASES) + def test_image_generation_uses_its_own_configured_base(self, monkeypatch: pytest.MonkeyPatch, brand: BrandCase) -> None: + monkeypatch.setenv("DASHSCOPE_API_BASE", "https://legacy.example.com/images") + monkeypatch.setenv(brand["base_env"], "https://gateway.example.com/qwen/images") + url: Final = brand["image_config"]().get_complete_url( + api_base=None, + api_key="sk-test", + model="qwen-image", + optional_params={}, + litellm_params={}, + ) + assert url == "https://gateway.example.com/qwen/images" + @pytest.mark.parametrize("brand", BRAND_CASES) def test_validate_environment_requires_key(self, brand): with pytest.raises(ValueError, match="DASHSCOPE_API_KEY"): diff --git a/tests/unit/test_dashscope_image_generation.py b/tests/unit/test_dashscope_image_generation.py index 960807a033b..9c65399350c 100644 --- a/tests/unit/test_dashscope_image_generation.py +++ b/tests/unit/test_dashscope_image_generation.py @@ -5,21 +5,26 @@ qwen-image-3.0, qwen-image-3.0-pro). Run in docker: pytest tests/unit/test_dashscope_image_generation.py -v """ +from typing import Final from unittest.mock import MagicMock, patch import httpx import pytest +import respx from pydantic import ValidationError import litellm +from litellm.llms.base_llm.chat.transformation import BaseLLMException +from litellm.llms.dashscope.common_utils import ( + DASHSCOPE_DEFAULT_ORIGIN, + DASHSCOPE_IMAGE_PATH, +) from litellm.llms.dashscope.image_generation.transformation import ( - DashScopeImageGenerationConfig, DEFAULT_API_BASE, + DashScopeImageGenerationConfig, ) from litellm.types.utils import ImageResponse from litellm.utils import get_llm_provider -from litellm.llms.base_llm.chat.transformation import BaseLLMException - # --------------------------------------------------------------------------- # 1. Provider detection @@ -77,6 +82,47 @@ class TestDashScopeImageGenerationConfig: url = self.cfg.get_complete_url(chat_api_base, None, "qwen-image-3.0", {}, {}) assert url == DEFAULT_API_BASE + @pytest.mark.parametrize( + ("api_base", "expected"), + [ + ( + "https://dashscope.aliyuncs.com", + "https://dashscope.aliyuncs.com/api/v1/services/aigc/multimodal-generation/generation", + ), + ( + "https://llm-abc123.cn-beijing.maas.aliyuncs.com", + "https://llm-abc123.cn-beijing.maas.aliyuncs.com/api/v1/services/aigc/multimodal-generation/generation", + ), + ( + "https://llm-abc123.cn-beijing.maas.aliyuncs.com/apps/anthropic", + "https://llm-abc123.cn-beijing.maas.aliyuncs.com/api/v1/services/aigc/multimodal-generation/generation", + ), + ( + "https://dashscope-us.aliyuncs.com/api/v1/", + "https://dashscope-us.aliyuncs.com/api/v1/services/aigc/multimodal-generation/generation", + ), + ], + ) + def test_get_complete_url_derives_from_the_domain( + self, api_base: str, expected: str + ): + url: Final = self.cfg.get_complete_url(api_base, None, "qwen-image-3.0", {}, {}) + assert url == expected + + @pytest.mark.parametrize( + "endpoint_url", + [ + "https://dashscope.aliyuncs.com/api/v1/services/aigc/image-generation/generation", + "https://dashscope-intl.aliyuncs.com/api/v1/services/aigc/text2image/image-synthesis", + ], + ) + def test_get_complete_url_keeps_a_full_dashscope_endpoint_url( + self, monkeypatch: pytest.MonkeyPatch, endpoint_url: str + ): + monkeypatch.setenv("DASHSCOPE_API_BASE_IMAGE", "https://images.example.com/generate") + url: Final = self.cfg.get_complete_url(endpoint_url, None, "qwen-image-3.0", {}, {}) + assert url == endpoint_url + def test_validate_environment_sets_auth_header(self): headers = self.cfg.validate_environment( headers={}, @@ -458,6 +504,81 @@ def test_litellm_image_generation_dashscope_end_to_end(model: str): assert body["parameters"]["size"] == "1024*1024" + +@pytest.mark.parametrize( + ("api_base", "chat_base_env", "image_base_env", "expected_url"), + [ + (None, "", "", DEFAULT_API_BASE), + ("https://dashscope.aliyuncs.com", "", "", f"{DASHSCOPE_DEFAULT_ORIGIN}{DASHSCOPE_IMAGE_PATH}"), + (None, "", "https://images.example.com/generate", "https://images.example.com/generate"), + ( + "https://dashscope.aliyuncs.com/compatible-mode/v1", + "", + "https://images.example.com/generate", + "https://images.example.com/generate", + ), + (None, "https://gateway.example.com/dashscope/images", "", "https://gateway.example.com/dashscope/images"), + (None, "https://dashscope-us.aliyuncs.com", "https://images.example.com/generate", "https://images.example.com/generate"), + ( + None, + f"https://dashscope-us.aliyuncs.com{DASHSCOPE_IMAGE_PATH}", + "https://images.example.com/generate", + f"https://dashscope-us.aliyuncs.com{DASHSCOPE_IMAGE_PATH}", + ), + ( + None, + "https://gateway.example.com/dashscope/images", + "https://images.example.com/generate", + "https://gateway.example.com/dashscope/images", + ), + ( + None, + "https://dashscope-us.aliyuncs.com/compatible-mode/v1", + "https://images.example.com/generate", + "https://images.example.com/generate", + ), + ( + "https://dashscope-us.aliyuncs.com", + "", + "https://images.example.com/generate", + f"https://dashscope-us.aliyuncs.com{DASHSCOPE_IMAGE_PATH}", + ), + ( + None, + "https://dashscope-us.aliyuncs.com", + "", + f"https://dashscope-us.aliyuncs.com{DASHSCOPE_IMAGE_PATH}", + ), + ( + "https://dashscope-intl.aliyuncs.com", + "https://gateway.example.com/dashscope/images", + "", + f"https://dashscope-intl.aliyuncs.com{DASHSCOPE_IMAGE_PATH}", + ), + ], +) +def test_image_generation_resolves_from_the_configured_base_not_the_chat_default( + monkeypatch: pytest.MonkeyPatch, + api_base: str | None, + chat_base_env: str, + image_base_env: str, + expected_url: str, +): + monkeypatch.setenv("DASHSCOPE_API_BASE", chat_base_env) + monkeypatch.setenv("DASHSCOPE_API_BASE_IMAGE", image_base_env) + image_url: Final = "https://dashscope-result.oss.aliyuncs.com/t.png" + + with respx.mock(assert_all_called=True) as upstream: + upstream.post(expected_url).respond( + 200, json={"output": {"choices": [{"message": {"content": [{"image": image_url}]}}]}} + ) + response: Final = litellm.image_generation( + model="dashscope/qwen-image-3.0", prompt="a puppy", api_key="sk-test-key", api_base=api_base + ) + + assert response.data[0].url == image_url + + def _transform_response(payload: object) -> ImageResponse: return DashScopeImageGenerationConfig().transform_image_generation_response( model="qwen-image-2.0", From c45b26402428804cdbd7504e1ea244df5c4d328e Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Thu, 29 Oct 2026 09:32:18 -0700 Subject: [PATCH 11/13] fix(proxy): persist tag spend and drain success callbacks within the shutdown budget (#48988) * fix: drain shutdown callbacks and persist pending tag spend * fix: catch shutdown timeouts on Python 3.10 * fix: share shutdown callback budget with HTTP draining * fix: type shutdown cancellation tasks and document its clock * test: retain callbacks after bounded worker cancellation * test(proxy): inspect the lifespan body after the SIGTERM wrapper split Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(utils): keep the success callback enqueue on its own task The synchronous enqueue spawned the logging worker loop inside the awaited call and ran callbacks that per-call loops used to drop, which hung the SDK callback tests on the autouse worker flush. The shutdown worker flush covers the queued callback on its own, so the one-tick task is restored Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(logging_worker): keep the queue's unfinished count exact when a full queue rejects a callback The aggressive clear ran the rejected callback and marked it done against the queue it was never put in, so flush() could return while a real callback was still running and the shutdown drain cancelled it early. The rejected callback is now re-enqueued once the clear made room, and only runs uncounted when the queue is still full Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * refactor(logging_worker): build the aggressive-clear run set immutably Final tuple instead of an appended list, with the re-enqueue in a _requeue helper, and Final annotations on the overflow regression test Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(logging): track overflow callbacks through shutdown * chore(logging): explain shared callback ownership * test(logging): isolate deferred callback completion --------- Co-authored-by: Joshua Valluru <326636767+joshua-berri@users.noreply.github.com> Co-authored-by: joshua Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- helm/litellm-helm/README.md | 6 + litellm/litellm_core_utils/logging_worker.py | 125 +++++---- litellm/proxy/db/db_spend_update_writer.py | 2 - litellm/proxy/proxy_server.py | 40 ++- .../shutdown/graceful_shutdown_manager.py | 58 ++++- litellm/responses/streaming_iterator.py | 2 +- litellm/utils.py | 8 +- .../integration/spend/test_shutdown_flush.py | 224 +++++++++++++++- .../litellm_core_utils/test_logging_worker.py | 246 ++++++++++++++++++ .../test_deferred_guardrail_logging.py | 103 ++------ .../unit/proxy/proxy_server/test_lifecycle.py | 5 +- .../test_graceful_shutdown_manager.py | 104 ++++++++ tests/unit/proxy/test_proxy_server.py | 199 ++++++++++++++ tests/unit/test_utils.py | 32 +++ 14 files changed, 1011 insertions(+), 143 deletions(-) diff --git a/helm/litellm-helm/README.md b/helm/litellm-helm/README.md index e059d45e037..a3488d6114c 100644 --- a/helm/litellm-helm/README.md +++ b/helm/litellm-helm/README.md @@ -259,3 +259,9 @@ itself. `lensWorker.mode` selects `disabled`, `bundled` or `external`. Bundled mode generates separate service and identity-signing credentials. External mode requires `lensWorker.gateway.secretName` and `lensWorker.serviceTokenSecret.name`; their keys default to `gateway-secret` and `service-token`. Provision distinct values matching the external Lens deployment Follow the [Lens Helm connection guide](https://github.com/BerriAI/lens/blob/main/helm/lens/README.md#connect-a-gateway) for complete values, verification, GitOps credential requirements and existing-data migration. Source-chart installation requires a built Lens image until the first signed Lens release is published + +## Shutdown accounting + +`GRACEFUL_SHUTDOWN_TIMEOUT` defaults to 30 seconds. Pending success callbacks share the time remaining since SIGTERM or the first `/health/drain` call. Callback flushing reserves the existing 10-second teardown headroom for cancellation, tracked accounting tasks and persistence. Exhausting this budget can discard callbacks, which is logged + +Set `terminationGracePeriodSeconds` to at least `GRACEFUL_SHUTDOWN_TIMEOUT + 10`, allowing additional time for preStop hooks and database writes. This is not a hard shutdown bound or a guarantee of persistence. Database retries retain their own limits, and cancellation-resistant callbacks can remain alive. The existing preStop hook can also be followed by Uvicorn's separate connection drain, as described in `values.yaml` diff --git a/litellm/litellm_core_utils/logging_worker.py b/litellm/litellm_core_utils/logging_worker.py index ff649953e69..8f217b1ea5e 100644 --- a/litellm/litellm_core_utils/logging_worker.py +++ b/litellm/litellm_core_utils/logging_worker.py @@ -6,7 +6,7 @@ import atexit import contextvars import inspect import logging -from collections.abc import Coroutine, Iterator +from collections.abc import Awaitable, Callable, Coroutine, Iterator from typing import Final from typing_extensions import TypedDict @@ -66,14 +66,19 @@ class LoggingWorker: max_queue_size: int = LOGGING_WORKER_MAX_QUEUE_SIZE, concurrency: int = LOGGING_WORKER_CONCURRENCY, timeout_summary_window: float = LOGGING_WORKER_TIMEOUT_SUMMARY_WINDOW_SECONDS, + *, + retry_sleep: Callable[[float], Awaitable[None]] = asyncio.sleep, ): self.timeout = timeout self.max_queue_size = max_queue_size self.concurrency = concurrency self.timeout_summary_window = timeout_summary_window self._queue: asyncio.Queue[LoggingTask] | None = None - self._worker_task: asyncio.Task | None = None - self._running_tasks: set[asyncio.Task] = set() + self._worker_task: asyncio.Task[None] | None = None + self._running_tasks: set[asyncio.Task[None]] = set() + self._overflow_tasks: set[asyncio.Task[None]] = set() # mutable-ok: done callbacks discard refs + self._pending_enqueues: dict[int, LoggingTask] = {} # mutable-ok: retries and shutdown share callback ownership + self._retry_sleep = retry_sleep self._dequeued_tasks: dict[int, LoggingTask] = {} # mutable-ok: refs so flush can rescue never-started tasks self._sem: asyncio.Semaphore | None = None self._bound_loop: asyncio.AbstractEventLoop | None = None @@ -168,6 +173,7 @@ class LoggingWorker: self._sem = None self._worker_task = None self._running_tasks.clear() + self._overflow_tasks.clear() self._timeout_summary_task = None self._timeout_burst_count = 0 self._queue = new_queue @@ -185,6 +191,10 @@ class LoggingWorker: self._sem = asyncio.Semaphore(self.concurrency) if self._worker_task is None or self._worker_task.done(): self._worker_task = asyncio.create_task(self._worker_loop()) + if not self._overflow_tasks: + self._aggressive_clear_in_progress = False + for pending in tuple(self._pending_enqueues.values()): + self._schedule_delayed_enqueue_retry(pending) async def _process_log_task( self, task: LoggingTask, sem: asyncio.Semaphore, queue: "asyncio.Queue[LoggingTask]" @@ -278,12 +288,24 @@ class LoggingWorker: # Capture the current context when enqueueing task: Final = LoggingTask(coroutine=coroutine, context=contextvars.copy_context()) + self._enqueue_task(task) + + def _enqueue_task(self, task: LoggingTask) -> None: + if self._queue is None: + return try: self._queue.put_nowait(task) except asyncio.QueueFull: - # Queue is full - handle it appropriately verbose_logger.exception("LoggingWorker queue is full") + self._pending_enqueues[id(task)] = task self._handle_queue_full(task) + else: + self._pending_enqueues.pop(id(task), None) + + def _start_overflow_task(self, coroutine: Coroutine[object, object, None]) -> None: + task: Final = asyncio.create_task(coroutine) + self._overflow_tasks.add(task) + task.add_done_callback(self._overflow_tasks.discard) def _should_start_aggressive_clear(self) -> bool: """ @@ -326,7 +348,7 @@ class LoggingWorker: if self._should_start_aggressive_clear(): self._mark_aggressive_clear_started() # Schedule clearing as async task so enqueue returns immediately (non-blocking) - asyncio.create_task(self._aggressively_clear_queue_async(task)) + self._start_overflow_task(self._aggressively_clear_queue_async(task)) else: # Cooldown active or clear in progress, schedule a delayed retry self._schedule_delayed_enqueue_retry(task) @@ -363,7 +385,7 @@ class LoggingWorker: delay: Final = self._calculate_retry_delay() # Schedule the retry as a background task - asyncio.create_task(self._retry_enqueue_task(task, delay)) + self._start_overflow_task(self._retry_enqueue_task(task, delay)) except RuntimeError: # No event loop, drop the task as we can't schedule a retry pass @@ -373,17 +395,8 @@ class LoggingWorker: Retry enqueueing the task after delay, preserving original context. This is called as a background task from _schedule_delayed_enqueue_retry. """ - await asyncio.sleep(delay) - - # Try to enqueue the task directly, preserving its original context - if self._queue is None: - return - - try: - self._queue.put_nowait(task) - except asyncio.QueueFull: - # Still full - handle it appropriately (clear or retry again) - self._handle_queue_full(task) + await self._retry_sleep(delay) + self._enqueue_task(task) def _extract_tasks_from_queue(self) -> list[LoggingTask]: """ @@ -423,15 +436,10 @@ class LoggingWorker: return extracted_tasks: Final = self._extract_tasks_from_queue() - - # Add new task to extracted tasks to process directly if new_task is not None: - self._track_dequeued(new_task) - extracted_tasks.append(new_task) - - # Process extracted tasks directly + self._enqueue_task(new_task) if extracted_tasks: - await self._process_extracted_tasks(extracted_tasks) + await asyncio.gather(*(self._process_single_task(task) for task in extracted_tasks)) except Exception as e: verbose_logger.exception("LoggingWorker error during aggressive clear: %s", e) finally: @@ -439,7 +447,7 @@ class LoggingWorker: self._aggressive_clear_in_progress = False async def _process_single_task(self, task: LoggingTask) -> None: - """Process a single task and mark it done.""" + """Process a task extracted from the queue and mark it done.""" if self._queue is None: return @@ -455,17 +463,6 @@ class LoggingWorker: self._untrack_dequeued(task) self._queue.task_done() - async def _process_extracted_tasks(self, tasks: list[LoggingTask]) -> None: - """ - Process tasks that were extracted from the queue to make room. - Processes them concurrently without semaphore limits for maximum speed. - """ - if not tasks or self._queue is None: - return - - # Process all tasks concurrently for maximum speed - await asyncio.gather(*[self._process_single_task(task) for task in tasks]) - def ensure_initialized_and_enqueue(self, async_coroutine: Coroutine[object, None, None]) -> None: """ Ensure the logging worker is initialized and enqueue the coroutine. @@ -473,32 +470,63 @@ class LoggingWorker: self.start() self.enqueue(async_coroutine) - async def stop(self) -> None: + @property + def pending_count(self) -> int: + return ( + (self._queue.qsize() if self._queue is not None else 0) + + len(self._dequeued_tasks) + + len(self._pending_enqueues) + ) + + async def stop(self, *, discard_pending: bool = False, timeout: float | None = None) -> None: """Stop the logging worker and clean up resources.""" + if discard_pending and self._queue is not None: + discarded: Final = self._drain_pending(self._queue) + for task in discarded: + task["coroutine"].close() + self._queue.task_done() + if discarded: + verbose_logger.warning("LoggingWorker shutdown: %d queued callbacks discarded", len(discarded)) if self._timeout_summary_task is not None: self._timeout_summary_task.cancel() self._timeout_summary_task = None self._emit_timeout_summary() - if self._worker_task is None and not self._running_tasks: - # No worker launched and no in-flight tasks to drain. - return + if discard_pending: + rejected: Final = tuple(self._pending_enqueues.values()) + self._pending_enqueues.clear() + for task in rejected: + task["coroutine"].close() + if rejected: + verbose_logger.warning("LoggingWorker shutdown: %d overflow callbacks discarded", len(rejected)) - tasks_to_cancel: Final[list[asyncio.Task]] = list(self._running_tasks) - if self._worker_task: - # Include the main worker loop so it stops fetching work. - tasks_to_cancel.append(self._worker_task) + tasks_to_cancel: Final = ( + *self._running_tasks, + *self._overflow_tasks, + *((self._worker_task,) if self._worker_task is not None else ()), + ) + if not tasks_to_cancel: + return for task in tasks_to_cancel: # Propagate cancellation to every pending task. task.cancel() # Wait for cancellation to settle; ignore errors raised during shutdown. - await asyncio.gather(*tasks_to_cancel, return_exceptions=True) + done, pending = await asyncio.wait(tasks_to_cancel, timeout=timeout) + await asyncio.gather(*done, return_exceptions=True) + if pending: + verbose_logger.warning("LoggingWorker shutdown: %d tasks still alive after cancellation", len(pending)) - self._worker_task = None + if discard_pending and self._queue is not None: + for task in self._unstarted_dequeued_tasks(): + task["coroutine"].close() + self._untrack_dequeued(task) + self._queue.task_done() + if self._worker_task is not None and self._worker_task.done(): + self._worker_task = None # Drop references to completed tasks so we can restart cleanly. - self._running_tasks.clear() + self._running_tasks.difference_update(done) async def flush(self) -> None: """Flush the logging queue. @@ -516,6 +544,9 @@ class LoggingWorker: return self.start() await self._queue.join() + while self._overflow_tasks: + await asyncio.wait(tuple(self._overflow_tasks)) + await self._queue.join() async def clear_queue(self): """ diff --git a/litellm/proxy/db/db_spend_update_writer.py b/litellm/proxy/db/db_spend_update_writer.py index 8ac2597c99b..16dbe2b4131 100644 --- a/litellm/proxy/db/db_spend_update_writer.py +++ b/litellm/proxy/db/db_spend_update_writer.py @@ -1900,8 +1900,6 @@ class DBSpendUpdateWriter: proxy_logging_obj=proxy_logging_obj, ) - # NOTE: Daily tag spend is committed by a separate scheduler job. - ################## Daily End-User Spend Update Transactions ################## # Aggregate all in memory daily end-user spend transactions and commit to db await self._flush_daily_spend_queue( diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index a75ef593b7a..65de7a3c472 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -79,6 +79,7 @@ from litellm.constants import ( DEFAULT_SHARED_HEALTH_CHECK_LOCK_TTL, DEFAULT_SHARED_HEALTH_CHECK_TTL, DEFAULT_SLACK_ALERTING_THRESHOLD, + GRACEFUL_SHUTDOWN_TEARDOWN_HEADROOM_SECONDS, INTERNAL_CALLBACK_NAMES, LANGFUSE_SHUTDOWN_FLUSH_TIMEOUT_MILLIS, LITELLM_EMBEDDING_PROVIDERS_SUPPORTING_INPUT_ARRAY_OF_TOKENS, @@ -98,6 +99,7 @@ from litellm.litellm_core_utils.litellm_logging import ( _init_custom_logger_compatible_class, initialize_sentry, ) +from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER, LoggingWorker from litellm.litellm_core_utils.safe_json_dumps import safe_dumps from litellm.litellm_core_utils.safe_json_loads import safe_json_loads from litellm.litellm_core_utils.token_counter import ( @@ -1274,15 +1276,44 @@ async def _flush_spend_logs_queue_on_shutdown() -> None: async def _flush_spend_on_shutdown() -> None: + from litellm.proxy.utils import update_daily_tag_spend + await flush_spend_counters_on_shutdown() + if prisma_client is not None: + await update_daily_tag_spend(prisma_client=prisma_client, proxy_logging_obj=proxy_logging_obj) await _flush_spend_logs_queue_on_shutdown() async def _flush_spend_around_logging_drain_on_shutdown( flush_spend: Callable[[], Awaitable[None]] = _flush_spend_on_shutdown, + *, + worker: LoggingWorker = GLOBAL_LOGGING_WORKER, + clock: Callable[[], float] = time.monotonic, + drain_tasks: Callable[[float], Awaitable[None]] = drain_logging_tasks, ) -> None: await flush_spend() - await drain_logging_tasks() + remaining: Final = GracefulShutdownManager.remaining(clock=clock) + callback_timeout: Final = max(0.0, remaining - GRACEFUL_SHUTDOWN_TEARDOWN_HEADROOM_SECONDS) + verbose_proxy_logger.info( + "Draining success callbacks before the final spend flush remaining_s=%.3f callback_timeout_s=%.3f", + remaining, + callback_timeout, + ) + + async def drain_callbacks() -> None: + await drain_tasks(callback_timeout) + await worker.flush() + + try: + await asyncio.wait_for(drain_callbacks(), timeout=callback_timeout) + except asyncio.TimeoutError: + verbose_proxy_logger.warning( + "Success callback shutdown budget exhausted with %d unfinished callbacks; accounting may be incomplete", + worker.pending_count, + ) + finally: + await worker.stop(discard_pending=True, timeout=GracefulShutdownManager.remaining(2.0, clock=clock)) + await drain_tasks(GracefulShutdownManager.remaining(5.0, clock=clock)) await flush_spend() @@ -1445,6 +1476,13 @@ async def _connect_to_count_stored_values() -> SupportsRawQueries: @asynccontextmanager async def proxy_startup_event(app: FastAPI) -> AsyncGenerator[ProxyLifespanState, None]: + with GracefulShutdownManager.capture_sigterm(): + async with _proxy_startup_event(app) as state: + yield state + + +@asynccontextmanager +async def _proxy_startup_event(app: FastAPI) -> AsyncGenerator[ProxyLifespanState, None]: global \ prisma_client, \ master_key, \ diff --git a/litellm/proxy/shutdown/graceful_shutdown_manager.py b/litellm/proxy/shutdown/graceful_shutdown_manager.py index d7a07ec34e5..9aeee0453ff 100644 --- a/litellm/proxy/shutdown/graceful_shutdown_manager.py +++ b/litellm/proxy/shutdown/graceful_shutdown_manager.py @@ -5,7 +5,8 @@ Kubernetes terminates a pod by sending ``SIGTERM`` and, after ``terminationGracePeriodSeconds``, ``SIGKILL``. On ``SIGTERM`` uvicorn stops accepting connections and waits up to ``GRACEFUL_SHUTDOWN_TIMEOUT`` for open ones, cancelling whatever is still running, before the lifespan teardown that -runs this drain. +runs this drain. The lifespan registers a SIGTERM wrapper before serving +requests so callback draining can subtract the preceding connection drain. A fixed ``preStop`` sleep can not solve this: it has to be sized for the *worst-case* request, so it either wastes time on every routine shutdown or is @@ -20,8 +21,12 @@ granularity of ``InFlightRequestsMiddleware``. import asyncio import os +import signal +import threading import time -from collections.abc import Callable +from collections.abc import Callable, Generator +from contextlib import contextmanager +from types import FrameType from typing import Final from litellm._logging import verbose_proxy_logger @@ -44,6 +49,7 @@ class GracefulShutdownManager: _is_shutting_down: bool = False _shutdown_started_at: float | None = None _drain_performed: bool = False + _sigterm_handler: Callable[[int, FrameType | None], None] | None = None @classmethod def is_shutting_down(cls) -> bool: @@ -71,7 +77,7 @@ class GracefulShutdownManager: return DEFAULT_GRACEFUL_SHUTDOWN_TIMEOUT @classmethod - def start_shutdown(cls) -> None: + def start_shutdown(cls, clock: Callable[[], float] = time.monotonic) -> None: """ Mark the worker as draining. Idempotent — repeated calls (e.g. SIGTERM followed by a preStop hit on /health/drain) do not reset the clock. @@ -79,12 +85,42 @@ class GracefulShutdownManager: if cls._is_shutting_down: return cls._is_shutting_down = True - cls._shutdown_started_at = time.monotonic() + cls._shutdown_started_at = clock() verbose_proxy_logger.info( "graceful_shutdown_started in_flight_requests=%s", get_in_flight_requests(), ) + @classmethod + def remaining(cls, cap: float | None = None, *, clock: Callable[[], float] = time.monotonic) -> float: + elapsed: Final = 0.0 if cls._shutdown_started_at is None else clock() - cls._shutdown_started_at + remaining: Final = max(0.0, cls.get_timeout() - elapsed) + return remaining if cap is None else min(max(0.0, cap), remaining) + + @classmethod + @contextmanager + def capture_sigterm(cls, clock: Callable[[], float] = time.monotonic) -> Generator[None, None, None]: + if threading.current_thread() is not threading.main_thread(): + yield + return + previous: Final = signal.getsignal(signal.SIGTERM) + if not callable(previous) or previous is cls._sigterm_handler: + yield + return + + def handle_sigterm(signum: int, frame: FrameType | None) -> None: + cls.start_shutdown(clock) + previous(signum, frame) + + cls._sigterm_handler = handle_sigterm + signal.signal(signal.SIGTERM, handle_sigterm) + try: + yield + finally: + if signal.getsignal(signal.SIGTERM) is handle_sigterm: + signal.signal(signal.SIGTERM, previous) + cls._sigterm_handler = None + @classmethod async def wait_for_drain( cls, @@ -93,13 +129,14 @@ class GracefulShutdownManager: count_fn: Callable[[], int] | None = None, poll_interval: float = _DRAIN_POLL_INTERVAL, log_interval: float = _DRAIN_LOG_INTERVAL, + clock: Callable[[], float] = time.monotonic, ) -> int: """ Poll the in-flight request counter until it reaches the drain target or ``timeout`` seconds elapse. Args: - timeout: Max seconds to wait. Defaults to ``get_timeout()``. + timeout: Max seconds to wait. Defaults to the remaining shutdown budget. exclude_self: When the caller is itself an in-flight HTTP request (the /health/drain endpoint), set this so the caller's own request is not counted as outstanding work. @@ -107,6 +144,7 @@ class GracefulShutdownManager: live ``InFlightRequestsMiddleware`` counter; injectable for tests. poll_interval: Seconds between counter polls. log_interval: Minimum seconds between ``drain_waiting`` log lines. + clock: Monotonic clock shared with the shutdown timestamp. Returns: Number of requests that drained while waiting (>= 0). @@ -120,7 +158,7 @@ class GracefulShutdownManager: cls._drain_performed = True if timeout is None: - timeout = cls.get_timeout() + timeout = cls.remaining(clock=clock) if count_fn is None: count_fn = get_in_flight_requests @@ -128,7 +166,7 @@ class GracefulShutdownManager: # and so counts itself; treat <=1 as "drained" in that case. target: Final = 1 if exclude_self else 0 - start: Final = time.monotonic() + start: Final = clock() initial: Final = count_fn() last_log = start @@ -142,11 +180,11 @@ class GracefulShutdownManager: verbose_proxy_logger.info( "graceful_shutdown_complete drained_requests=%s elapsed_s=%.2f", drained, - time.monotonic() - start, + clock() - start, ) return drained - elapsed = time.monotonic() - start + elapsed = clock() - start if elapsed >= timeout: verbose_proxy_logger.warning( "graceful_shutdown_timeout in_flight_requests=%s elapsed_s=%.2f " @@ -157,7 +195,7 @@ class GracefulShutdownManager: ) return max(0, initial - current) - now = time.monotonic() + now = clock() if now - last_log >= log_interval: verbose_proxy_logger.info( "drain_waiting in_flight_requests=%s elapsed_s=%.2f", diff --git a/litellm/responses/streaming_iterator.py b/litellm/responses/streaming_iterator.py index 8959c32fff2..6a82304b76a 100644 --- a/litellm/responses/streaming_iterator.py +++ b/litellm/responses/streaming_iterator.py @@ -227,7 +227,7 @@ async def drain_logging_tasks(timeout: float = 5.0) -> None: return _, pending = await asyncio.wait(tuple(_PENDING_LOGGING_TASKS), timeout=timeout) if pending: - verbose_logger.warning("Timed out draining %d logging tasks on shutdown, dropping them", len(pending)) + verbose_logger.warning("Timed out draining logging tasks on shutdown: %d tasks still alive", len(pending)) _ERROR_CODE_HTTP_STATUS: Final[Mapping[str, int]] = MappingProxyType( diff --git a/litellm/utils.py b/litellm/utils.py index 4cd2335dd0c..4a0e4cafb22 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -1632,15 +1632,19 @@ def _schedule_async_success_logging( return def _enqueue_async_logging() -> None: + from litellm.responses.streaming_iterator import spawn_logging_task + logging_obj.mark_async_success_log_enqueued() - asyncio.create_task( + spawn_logging_task( + asyncio.get_running_loop(), _client_async_logging_helper( logging_obj=logging_obj, result=result, start_time=start_time, end_time=end_time, is_completion_with_fallbacks=is_completion_with_fallbacks, - ) + ), + task_name="success callback enqueue", ) if not logging_obj.defer_async_logging: diff --git a/tests/integration/spend/test_shutdown_flush.py b/tests/integration/spend/test_shutdown_flush.py index e96511511d9..5c2e9cf98da 100644 --- a/tests/integration/spend/test_shutdown_flush.py +++ b/tests/integration/spend/test_shutdown_flush.py @@ -1,11 +1,14 @@ +import hashlib import json import os +import re import threading +import time import uuid from collections.abc import Callable, Iterator from concurrent.futures import ThreadPoolExecutor from contextlib import contextmanager -from dataclasses import dataclass +from dataclasses import dataclass, replace from pathlib import Path from typing import Final @@ -26,6 +29,225 @@ MODEL_INSERT_ARRIVED_LOG_LINE: Final = "path=/model/new" COMMIT_DELAY_SECONDS: Final = 15 BURST_REQUESTS: Final = 30 +_GATED_LOGGER: Final = """ +import os +import re +import httpx +from litellm.integrations.custom_logger import CustomLogger + +class ShutdownLogger(CustomLogger): + async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): + async with httpx.AsyncClient(timeout=25, trust_env=False) as client: + await client.post(os.environ["SHUTDOWN_CALLBACK_URL"]) + +logger = ShutdownLogger() +""" + + +def _tag_shutdown_config(directory: Path, upstream: str, redis_buffer: bool, *, gated: bool = False) -> Path: + path: Final = directory / "tag-shutdown.yaml" + if gated: + (directory / "shutdown_callbacks.py").write_text(_GATED_LOGGER) + path.write_text( + yaml.safe_dump( + { + "model_list": [ + { + "model_name": "shutdown-model", + "litellm_params": { + "model": "openai/gpt-4o-mini", + "api_base": upstream + "/v1", + "api_key": "synthetic-key", + "input_cost_per_token": 0.01, + "output_cost_per_token": 0.02, + }, + } + ], + "general_settings": { + "master_key": "os.environ/LITELLM_MASTER_KEY", + "database_url": "os.environ/DATABASE_URL", + "proxy_batch_write_at": 3600, + "use_redis_transaction_buffer": redis_buffer, + "disable_model_info_refresh": True, + }, + "litellm_settings": { + "cache": True, + "cache_params": {"type": "redis", "host": "os.environ/REDIS_HOST", "port": "os.environ/REDIS_PORT"}, + "callbacks": ["shutdown_callbacks.logger"] if gated else [], + }, + } + ) + ) + return path + + +def _tagged_chat(gateway: Gateway, key: str, tags: tuple[str, str]) -> None: + response: Final = gateway.request( + "POST", + "/v1/chat/completions", + { + "model": "shutdown-model", + "messages": [{"role": "user", "content": uuid.uuid4().hex}], + "metadata": {"tags": list(tags)}, + }, + key=key, + ) + assert response.status_code == 200, response.text + + +def _assert_shutdown_accounting(key: str, tags: tuple[str, str], requests: int) -> None: + hashed: Final = hashlib.sha256(key.encode()).hexdigest() + for table, count in (("LiteLLM_SpendLogs", "count(*)"), ("LiteLLM_DailyUserSpend", "sum(api_requests)")): + rows: Final = read_rows( + f'SELECT ({count})::int AS requests, sum(spend) AS spend FROM "{table}" WHERE api_key=%s', (hashed,) + ) + assert rows[0]["requests"] == requests, (table, rows) + assert rows[0]["spend"] == pytest.approx(requests * 0.05, abs=1e-9), (table, rows) + key_rows: Final = read_rows('SELECT spend FROM "LiteLLM_VerificationToken" WHERE token=%s', (hashed,)) + assert key_rows[0]["spend"] == pytest.approx(requests * 0.05, abs=1e-9) + for tag in tags: + tag_rows: Final = read_rows( + 'SELECT sum(api_requests)::int AS requests, sum(spend) AS spend FROM "LiteLLM_DailyTagSpend" WHERE api_key=%s AND tag=%s', + (hashed, tag), + ) + assert tag_rows[0]["requests"] == requests, (tag, tag_rows) + assert tag_rows[0]["spend"] == pytest.approx(requests * 0.05, abs=1e-9), (tag, tag_rows) + + +@pytest.mark.parametrize("redis_buffer", [False, True]) +def test_tagged_request_survives_shutdown_before_the_tag_timer( + gateway: Gateway, tmp_path: Path, redis_buffer: bool +) -> None: + tags: Final = (f"shutdown-a-{uuid.uuid4().hex}", f"shutdown-b-{uuid.uuid4().hex}") + overrides: Final = {"DATABASE_URL": os.environ["DATABASE_URL"], "LITELLM_LOG": "DEBUG"} + with wire_server(_provider) as wire: + config: Final = _tag_shutdown_config(tmp_path, wire.url, redis_buffer) + with owned_proxy_process(gateway, tmp_path, overrides, config=config) as owned: + key: Final = string_value( + owned.gateway.post("/key/generate", {"user_id": f"shutdown-{uuid.uuid4().hex}"})["key"] + ) + _tagged_chat(owned.gateway, key, tags) + owned.begin_shutdown() + owned.process.wait(timeout=90) + with owned_proxy_process(gateway, tmp_path, overrides, config=config) as restarted: + _assert_shutdown_accounting(key, tags, 1) + restarted.begin_shutdown() + restarted.process.wait(timeout=90) + _assert_shutdown_accounting(key, tags, 1) + + +@pytest.mark.parametrize("redis_buffer", [False, True]) +@pytest.mark.parametrize("overflow", [False, True]) +def test_callback_backlog_is_accounted_before_shutdown_disconnects( + gateway: Gateway, tmp_path: Path, redis_buffer: bool, overflow: bool +) -> None: + release: Final = threading.Event() + requests: Final = 6 if overflow else 3 + tags: Final = (f"backlog-a-{uuid.uuid4().hex}", f"backlog-b-{uuid.uuid4().hex}") + with ( + wire_server(_provider) as wire, + wire_server(lambda _: Reply(chunks=(b"{}",), gate_before_first=release)) as logger, + ): + config: Final = _tag_shutdown_config(tmp_path, wire.url, redis_buffer, gated=True) + with owned_proxy_process( + gateway, + tmp_path, + { + "DATABASE_URL": os.environ["DATABASE_URL"], + "LOGGING_WORKER_CONCURRENCY": "2", + "LOGGING_WORKER_MAX_QUEUE_SIZE": "2" if overflow else "50000", + "SHUTDOWN_CALLBACK_URL": logger.url, + "GRACEFUL_SHUTDOWN_TIMEOUT": "15", + "LITELLM_LOG": "DEBUG", + }, + config=config, + ) as owned: + key: Final = string_value( + owned.gateway.post("/key/generate", {"user_id": f"backlog-{uuid.uuid4().hex}"})["key"] + ) + try: + for _ in range(requests): + _tagged_chat(owned.gateway, key, tags) + eventually(logger.received.qsize, lambda count: count >= 2, seconds=10) + owned.begin_shutdown() + eventually( + lambda: ( + owned.process.poll() is not None + or "Draining success callbacks before the final spend flush" in owned.log.read_text() + ), + lambda ready: ready, + seconds=30, + ) + release.set() + owned.process.wait(timeout=90) + _assert_shutdown_accounting(key, tags, requests) + finally: + release.set() + + +@pytest.mark.parametrize("gunicorn", [False, True]) +def test_sigterm_http_drain_consumes_callback_budget(gateway: Gateway, tmp_path: Path, gunicorn: bool) -> None: + release_request: Final = threading.Event() + release_callback: Final = threading.Event() + tags: Final = (f"deadline-a-{uuid.uuid4().hex}", f"deadline-b-{uuid.uuid4().hex}") + + def held_provider(request: Request) -> Reply: + response: Final = _provider(request) + return replace(response, chunks=(response.body,), gate_before_first=release_request) + + with ( + wire_server(held_provider) as wire, + wire_server(lambda _: Reply(chunks=(b"{}",), gate_before_first=release_callback)) as logger, + ): + config: Final = _tag_shutdown_config(tmp_path, wire.url, False, gated=True) + with ( + owned_proxy_process( + gateway, + tmp_path, + { + "DATABASE_URL": os.environ["DATABASE_URL"], + "SHUTDOWN_CALLBACK_URL": logger.url, + "GRACEFUL_SHUTDOWN_TIMEOUT": "30", + "LITELLM_LOG": "DEBUG", + }, + config=config, + extra_arguments=("--run_gunicorn",) if gunicorn else (), + ) as owned, + ThreadPoolExecutor(max_workers=1) as executor, + ): + key: Final = string_value( + owned.gateway.post("/key/generate", {"user_id": f"deadline-{uuid.uuid4().hex}"})["key"] + ) + request: Final = executor.submit(_tagged_chat, owned.gateway, key, tags) + try: + eventually(wire.received.qsize, lambda count: count == 1, seconds=10) + owned.begin_shutdown() + eventually( + lambda: "graceful_shutdown_started" in owned.log.read_text(), lambda stamped: stamped, seconds=5 + ) + assert not request.done(), "The shutdown clock must start before the HTTP request drains" + stamped: Final = time.monotonic() + eventually(time.monotonic, lambda now: now - stamped >= 2, seconds=5) + held_seconds: Final = time.monotonic() - stamped + release_request.set() + request.result(timeout=10) + budget: Final = eventually( + lambda: re.search(r"remaining_s=([\d.]+) callback_timeout_s=([\d.]+)", owned.log.read_text()), + lambda match: match is not None, + seconds=15, + ) + assert budget is not None + remaining: Final = float(budget[1]) + allowance: Final = float(budget[2]) + assert remaining <= 30 - held_seconds + 0.05 + assert allowance == pytest.approx(max(0, remaining - 10), abs=0.002) + release_callback.set() + owned.process.wait(timeout=45) + _assert_shutdown_accounting(key, tags, 1) + finally: + release_request.set() + release_callback.set() + def _api_requests(table: str, column: str, identity: str) -> int: rows: Final = read_rows( diff --git a/tests/unit/litellm_core_utils/test_logging_worker.py b/tests/unit/litellm_core_utils/test_logging_worker.py index f48495d1062..fccc11617bb 100644 --- a/tests/unit/litellm_core_utils/test_logging_worker.py +++ b/tests/unit/litellm_core_utils/test_logging_worker.py @@ -82,6 +82,252 @@ class _RecordCollector(logging.Handler): self.records.append(record) +@pytest.mark.asyncio +@pytest.mark.parametrize("running", [False, True]) +async def test_stop_discards_waiting_callbacks_without_executing_them(running: bool) -> None: + worker: Final = LoggingWorker(concurrency=1) + started: Final = asyncio.Event() + cancelled: Final = asyncio.Event() + executed: Final = asyncio.Event() + + async def blocked() -> None: + started.set() + try: + await asyncio.Event().wait() + finally: + cancelled.set() + + async def waiting() -> None: + executed.set() + + if running: + worker.ensure_initialized_and_enqueue(blocked()) + await started.wait() + else: + worker._ensure_queue() + worker.enqueue(waiting()) + await worker.stop(discard_pending=True) + assert not executed.is_set() + assert cancelled.is_set() == running + assert worker._queue is not None + assert worker._queue.empty() + await asyncio.wait_for(worker._queue.join(), timeout=1) + + +@pytest.mark.asyncio +async def test_stop_does_not_replay_a_callback_dequeued_before_its_processor_started() -> None: + worker: Final = LoggingWorker(concurrency=1) + executed: Final = asyncio.Event() + + async def callback() -> None: + executed.set() + + worker.ensure_initialized_and_enqueue(callback()) + await asyncio.sleep(0) + await worker.stop(discard_pending=True) + assert not executed.is_set() + worker._flush_on_exit() + assert not executed.is_set() + assert worker._queue is not None + await asyncio.wait_for(worker._queue.join(), timeout=1) + + +@pytest.mark.asyncio +async def test_stop_timeout_keeps_callbacks_observable_until_queue_cleanup_finishes() -> None: + worker: Final = LoggingWorker(concurrency=1) + started: Final = asyncio.Event() + release: Final = asyncio.Event() + completed: Final = asyncio.Event() + + async def running() -> None: + started.set() + await asyncio.Event().wait() + + async def queued() -> None: + await release.wait() + completed.set() + + worker.ensure_initialized_and_enqueue(running()) + await started.wait() + worker.enqueue(queued()) + try: + await worker.stop(timeout=0) + assert not completed.is_set() + release.set() + await worker.flush() + assert completed.is_set(), "Queue cleanup must survive a bounded cancellation wait" + finally: + release.set() + await worker.stop() + + +@pytest.mark.asyncio +async def test_overflow_callback_does_not_undercount_the_unfinished_queue() -> None: + """A callback rejected by a full queue is run by the aggressive clear. It must not consume a + queue.task_done() it never contributed, otherwise flush() returns while a real callback is still + running and the shutdown drain cancels it with budget to spare. + """ + worker: Final = LoggingWorker(timeout=5.0, max_queue_size=2, concurrency=1) + gates: Final = {name: asyncio.Event() for name in ("first", "second", "third")} + completed: Final[list[str]] = [] + + async def gated(name: str) -> None: + await gates[name].wait() + completed.append(name) + + async def instant() -> None: + completed.append("overflow") + + worker.start() + worker.enqueue(gated("first")) + await asyncio.sleep(0) + worker.enqueue(gated("second")) + worker.enqueue(gated("third")) + assert worker._queue is not None and worker._queue.full(), "precondition: the queue is full" + worker.enqueue(instant()) + for _ in range(5): + await asyncio.sleep(0) + + gates["first"].set() + gates["second"].set() + asyncio.get_running_loop().call_later(0.2, gates["third"].set) + await asyncio.wait_for(worker.flush(), timeout=2.0) + assert "third" in completed, "flush returned while a queued callback was still unfinished" + assert sorted(completed) == ["first", "overflow", "second", "third"] + await worker.stop() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("queue_size", (1, 2)) +async def test_flush_waits_for_delayed_overflow_and_preserves_callback_context(queue_size: int) -> None: + retry_started: Final = asyncio.Event() + release_retry: Final = asyncio.Event() + completed: Final[asyncio.Queue[tuple[int, str]]] = asyncio.Queue() + identity: Final = contextvars.ContextVar("overflow_identity", default="outside") + + async def retry_sleep(delay: float) -> None: + retry_started.set() + await release_retry.wait() + + worker: Final = LoggingWorker(max_queue_size=queue_size, concurrency=1, retry_sleep=retry_sleep) + + async def callback(index: int) -> None: + completed.put_nowait((index, identity.get())) + + worker.start() + token: Final = identity.set("request") + for index in range(4): + worker.enqueue(callback(index)) + identity.reset(token) + flushing: Final = asyncio.create_task(worker.flush()) + try: + await asyncio.wait_for(retry_started.wait(), timeout=2) + assert worker._queue is not None + await asyncio.wait_for(worker._queue.join(), timeout=2) + await asyncio.sleep(0) + assert not flushing.done(), "Queue completion must not hide delayed overflow callbacks" + assert worker.pending_count > 0 + release_retry.set() + await asyncio.wait_for(flushing, timeout=2) + assert sorted(completed.get_nowait() for _ in range(completed.qsize())) == [ + (index, "request") for index in range(4) + ] + assert worker.pending_count == 0 + finally: + release_retry.set() + await worker.stop(discard_pending=True) + await asyncio.gather(flushing, return_exceptions=True) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("start_retry", (False, True)) +async def test_shutdown_discards_overflow_without_late_enqueue_or_replay(start_retry: bool) -> None: + import inspect + + retry_started: Final = asyncio.Event() + release_retry: Final = asyncio.Event() + executed: Final[asyncio.Queue[int]] = asyncio.Queue() + + async def retry_sleep(delay: float) -> None: + retry_started.set() + await release_retry.wait() + + worker: Final = LoggingWorker(max_queue_size=2, concurrency=1, retry_sleep=retry_sleep) + + async def callback(index: int) -> None: + executed.put_nowait(index) + + callbacks: Final = tuple(callback(index) for index in range(4)) + worker._ensure_queue() + for callback_coro in callbacks: + worker.enqueue(callback_coro) + if start_retry: + await asyncio.wait_for(retry_started.wait(), timeout=2) + await worker.stop(discard_pending=True, timeout=1) + completed_at_stop: Final = executed.qsize() + release_retry.set() + await worker.flush() + try: + assert executed.qsize() == completed_at_stop, "Discarded overflow must not run when the worker restarts" + assert worker.pending_count == 0 + assert all(inspect.getcoroutinestate(coro) == inspect.CORO_CLOSED for coro in callbacks) + finally: + await worker.stop(discard_pending=True) + + +@pytest.mark.asyncio +async def test_shutdown_cancels_aggressive_clear_callbacks_before_returning() -> None: + started: Final = asyncio.Event() + cancelled: Final = asyncio.Event() + worker: Final = LoggingWorker(max_queue_size=2, concurrency=1) + + async def blocked() -> None: + started.set() + try: + await asyncio.Event().wait() + finally: + cancelled.set() + + async def noop() -> None: + pass + + worker._ensure_queue() + worker.enqueue(blocked()) + worker.enqueue(noop()) + worker.enqueue(noop()) + await asyncio.wait_for(started.wait(), timeout=2) + await worker.stop(discard_pending=True, timeout=1) + assert cancelled.is_set(), "Overflow execution must participate in shutdown cancellation" + assert worker.pending_count == 0 + assert worker._queue is not None + await asyncio.wait_for(worker._queue.join(), timeout=1) + + +@pytest.mark.asyncio +async def test_stopped_overflow_can_be_flushed_when_not_discarded() -> None: + release_retry: Final = asyncio.Event() + completed: Final[asyncio.Queue[int]] = asyncio.Queue() + + async def retry_sleep(delay: float) -> None: + await release_retry.wait() + + worker: Final = LoggingWorker(max_queue_size=2, concurrency=1, retry_sleep=retry_sleep) + + async def callback(index: int) -> None: + completed.put_nowait(index) + + worker._ensure_queue() + for index in range(4): + worker.enqueue(callback(index)) + await worker.stop() + release_retry.set() + try: + await asyncio.wait_for(worker.flush(), timeout=2) + assert sorted(completed.get_nowait() for _ in range(completed.qsize())) == list(range(4)) + finally: + await worker.stop(discard_pending=True) + + class TestLoggingWorker: """Test cases for LoggingWorker functionality.""" diff --git a/tests/unit/proxy/guardrails/test_deferred_guardrail_logging.py b/tests/unit/proxy/guardrails/test_deferred_guardrail_logging.py index bd6190fee5f..6d1e5c72982 100644 --- a/tests/unit/proxy/guardrails/test_deferred_guardrail_logging.py +++ b/tests/unit/proxy/guardrails/test_deferred_guardrail_logging.py @@ -75,11 +75,13 @@ class _RecordingLogger(CustomLogger): def __init__(self) -> None: super().__init__() self.standard_logging_object: StandardLoggingPayload | None = None + self.logged: asyncio.Event = asyncio.Event() async def async_log_success_event( self, kwargs: Mapping[str, object], response_obj: object, start_time: datetime, end_time: datetime ) -> None: self.standard_logging_object = cast(StandardLoggingPayload, kwargs["standard_logging_object"]) + self.logged.set() class PostCallGuardrail(CustomGuardrail): @@ -239,52 +241,33 @@ class TestHasPostCallGuardrailsForPassthrough: @pytest.mark.asyncio -async def test_deferred_flag_stores_and_executes_closure(): - """ - When defer_async_logging is True on logging_obj: - 1. wrapper_async stores a callable closure instead of calling create_task - 2. Calling the closure fires create_task - 3. Sync callbacks fire immediately (not deferred) - """ - mock_logging_obj = MagicMock() - mock_logging_obj.defer_async_logging = True - mock_logging_obj.enqueue_deferred_logging = None - - await litellm.acompletion( - model="gpt-3.5-turbo", +@pytest.mark.parametrize("deferred", (False, True)) +async def test_success_logging_waits_for_deferred_closure_when_requested(deferred: bool) -> None: + recorder: Final = _RecordingLogger() + logging_obj: Final = Logging( + model="gpt-4o-mini", + messages=[{"role": "user", "content": "hi"}], + stream=False, + call_type="acompletion", + start_time=datetime(2026, 1, 1), + litellm_call_id="deferred-success", + function_id="deferred-success", + dynamic_async_success_callbacks=[recorder], + ) + logging_obj.defer_async_logging = deferred + response: Final = await litellm.acompletion( + model="gpt-4o-mini", messages=[{"role": "user", "content": "hi"}], mock_response="Hello!", - litellm_logging_obj=mock_logging_obj, + litellm_logging_obj=logging_obj, ) - - # Closure was stored - enqueue_fn = mock_logging_obj.enqueue_deferred_logging - assert callable(enqueue_fn), "Closure should be stored on logging_obj" - - # Sync callbacks fired immediately - mock_logging_obj.handle_sync_success_callbacks_for_async_calls.assert_called_once() - - # Calling the closure fires create_task - created_tasks = [] - real_create_task = asyncio.create_task - - def tracking_create_task(coro): - task = real_create_task(coro) - created_tasks.append(task) - return task - - with patch("asyncio.create_task", side_effect=tracking_create_task): - enqueue_fn() - - assert len(created_tasks) >= 1, "Closure should fire asyncio.create_task" - - for task in created_tasks: - if not task.done(): - task.cancel() - try: - await task - except (asyncio.CancelledError, Exception): - pass + if deferred: + assert recorder.standard_logging_object is None + assert logging_obj.enqueue_deferred_logging is not None + logging_obj.enqueue_deferred_logging() + await asyncio.wait_for(recorder.logged.wait(), timeout=5) + assert recorder.standard_logging_object is not None + assert recorder.standard_logging_object["id"] == response.id @pytest.mark.asyncio @@ -401,40 +384,6 @@ async def test_deferred_anthropic_messages_bridged_to_the_responses_api_logs_the assert recorder.standard_logging_object["response_cost"] == pytest.approx(3 * 2e-7 + 7333 * 2e-8 + 23 * 1.25e-6) -# --------------------------------------------------------------------------- -# 3. Non-streaming regression: without flag, create_task fires normally -# --------------------------------------------------------------------------- - - -@pytest.mark.asyncio -async def test_no_flag_fires_create_task_normally(): - """Without defer_async_logging, wrapper_async calls create_task as before.""" - created_tasks = [] - real_create_task = asyncio.create_task - - def tracking_create_task(coro): - task = real_create_task(coro) - created_tasks.append(task) - return task - - with patch("asyncio.create_task", side_effect=tracking_create_task): - await litellm.acompletion( - model="gpt-3.5-turbo", - messages=[{"role": "user", "content": "hi"}], - mock_response="Hello!", - ) - - assert len(created_tasks) >= 1 - - for task in created_tasks: - if not task.done(): - task.cancel() - try: - await task - except (asyncio.CancelledError, Exception): - pass - - # --------------------------------------------------------------------------- # 4. Non-streaming: deferred success log is suppressed when guardrail raises # --------------------------------------------------------------------------- diff --git a/tests/unit/proxy/proxy_server/test_lifecycle.py b/tests/unit/proxy/proxy_server/test_lifecycle.py index 8b1d7c1222a..a8706c5bccd 100644 --- a/tests/unit/proxy/proxy_server/test_lifecycle.py +++ b/tests/unit/proxy/proxy_server/test_lifecycle.py @@ -35,6 +35,7 @@ from typing_extensions import TypedDict import litellm.proxy.proxy_server as ps from litellm.proxy.management_helpers.audit_logs import track_audit_task from litellm.proxy.proxy_server import ( + _proxy_startup_event, ProxyStartupEvent, _resolve_pydantic_type, _resolve_typed_dict_type, @@ -1059,7 +1060,7 @@ def test_otel_global_provider_published_after_callback_init(): span was orphaned. The publish (``publish_global_otel_v2_provider``) must therefore appear after ``_initialize_startup_logging`` in the lifespan source. """ - wrapped = getattr(proxy_startup_event, "__wrapped__", proxy_startup_event) + wrapped = getattr(_proxy_startup_event, "__wrapped__", _proxy_startup_event) source = inspect.getsource(wrapped) init_pos = source.find("ProxyStartupEvent.initialize_startup_logging(") publish_pos = source.find("publish_global_otel_v2_provider(") @@ -1145,7 +1146,7 @@ def test_proxy_startup_event_warns_for_global_budget_without_database(): The call must follow Prisma setup so DB-backed deployments do not false-positive. Direct ``_warn_budget_without_db`` tests cover the warning behavior itself. """ - wrapped = getattr(proxy_startup_event, "__wrapped__", proxy_startup_event) + wrapped = getattr(_proxy_startup_event, "__wrapped__", _proxy_startup_event) source = inspect.getsource(wrapped) budget_check_pos = source.find("if prisma_client is not None and litellm.max_budget > 0:") warn_pos = source.find("warn_budget_without_db(") diff --git a/tests/unit/proxy/shutdown/test_graceful_shutdown_manager.py b/tests/unit/proxy/shutdown/test_graceful_shutdown_manager.py index d38852617b9..fdf1adda1f4 100644 --- a/tests/unit/proxy/shutdown/test_graceful_shutdown_manager.py +++ b/tests/unit/proxy/shutdown/test_graceful_shutdown_manager.py @@ -179,3 +179,107 @@ async def test_emits_periodic_drain_waiting_log_while_waiting(): timeout=10, count_fn=count_fn, poll_interval=0, log_interval=0 ) assert drained == 3 + + +@pytest.mark.parametrize("elapsed, cap, expected", [(4, None, 26), (25, 2, 2), (29, 5, 1), (31, None, 0)]) +def test_remaining_budget_includes_time_before_lifespan( + monkeypatch: pytest.MonkeyPatch, elapsed: float, cap: float | None, expected: float +) -> None: + monkeypatch.setenv("GRACEFUL_SHUTDOWN_TIMEOUT", "30") + GracefulShutdownManager.start_shutdown(clock=lambda: 100.0) + GracefulShutdownManager.start_shutdown(clock=lambda: 100.0 + elapsed) + assert GracefulShutdownManager.remaining(cap, clock=lambda: 100.0 + elapsed) == expected + + +def test_remaining_budget_before_shutdown(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv("GRACEFUL_SHUTDOWN_TIMEOUT", "30") + assert GracefulShutdownManager.remaining(clock=lambda: 100.0) == 30.0 + assert GracefulShutdownManager.remaining(-1, clock=lambda: 100.0) == 0.0 + + +@pytest.mark.parametrize("prestop", [False, True]) +def test_sigterm_stamps_before_chaining_and_restores_server_handler( + monkeypatch: pytest.MonkeyPatch, prestop: bool +) -> None: + import signal + from queue import SimpleQueue + from types import FrameType + from typing import Final + + monkeypatch.setenv("GRACEFUL_SHUTDOWN_TIMEOUT", "30") + calls: Final[SimpleQueue[float]] = SimpleQueue() + + def server_handler(signum: int, frame: FrameType | None) -> None: + assert signum == signal.SIGTERM + calls.put(GracefulShutdownManager.remaining(clock=lambda: 105.0)) + + original: Final = signal.signal(signal.SIGTERM, server_handler) + try: + if prestop: + GracefulShutdownManager.start_shutdown(clock=lambda: 95.0) + with GracefulShutdownManager.capture_sigterm(clock=lambda: 100.0): + installed: Final = signal.getsignal(signal.SIGTERM) + with GracefulShutdownManager.capture_sigterm(clock=lambda: 102.0): + assert signal.getsignal(signal.SIGTERM) is installed + signal.raise_signal(signal.SIGTERM) + signal.raise_signal(signal.SIGTERM) + assert calls.get_nowait() == (20.0 if prestop else 25.0) + assert calls.get_nowait() == (20.0 if prestop else 25.0) + assert calls.empty() + assert signal.getsignal(signal.SIGTERM) is server_handler + finally: + signal.signal(signal.SIGTERM, original) + + +def test_sigterm_capture_preserves_a_replacement_handler() -> None: + import signal + from types import FrameType + from typing import Final + + def server_handler(signum: int, frame: FrameType | None) -> None: + pass + + original: Final = signal.signal(signal.SIGTERM, server_handler) + try: + with GracefulShutdownManager.capture_sigterm(): + signal.signal(signal.SIGTERM, signal.SIG_IGN) + assert signal.getsignal(signal.SIGTERM) is signal.SIG_IGN + finally: + signal.signal(signal.SIGTERM, original) + + +@pytest.mark.parametrize("handler", ["SIG_DFL", "SIG_IGN"]) +def test_sigterm_capture_preserves_default_and_ignored_dispositions(handler: str) -> None: + import signal + from typing import Final + + disposition: Final = signal.SIG_DFL if handler == "SIG_DFL" else signal.SIG_IGN + original: Final = signal.signal(signal.SIGTERM, disposition) + try: + with GracefulShutdownManager.capture_sigterm(): + assert signal.getsignal(signal.SIGTERM) is disposition + finally: + signal.signal(signal.SIGTERM, original) + + +def test_sigterm_capture_in_worker_thread_leaves_signal_ownership_unchanged() -> None: + import signal + from concurrent.futures import ThreadPoolExecutor + from typing import Final + + original: Final = signal.getsignal(signal.SIGTERM) + + def run_lifespan() -> None: + with GracefulShutdownManager.capture_sigterm(): + assert signal.getsignal(signal.SIGTERM) is original + + with ThreadPoolExecutor(max_workers=1) as executor: + executor.submit(run_lifespan).result() + assert signal.getsignal(signal.SIGTERM) is original + + +@pytest.mark.asyncio +async def test_lifespan_request_drain_does_not_restart_an_exhausted_budget(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv("GRACEFUL_SHUTDOWN_TIMEOUT", "30") + GracefulShutdownManager.start_shutdown(clock=lambda: 100.0) + assert await GracefulShutdownManager.wait_for_drain(count_fn=lambda: 2, clock=lambda: 130.0) == 2 diff --git a/tests/unit/proxy/test_proxy_server.py b/tests/unit/proxy/test_proxy_server.py index 22d6316c448..c835e1e27ea 100644 --- a/tests/unit/proxy/test_proxy_server.py +++ b/tests/unit/proxy/test_proxy_server.py @@ -263,6 +263,174 @@ async def test_shutdown_commits_spend_a_logging_task_queued_while_the_first_comm token = MASTER_KEY + +@pytest.mark.asyncio +async def test_shutdown_persists_accounting_from_worker_callbacks_before_returning() -> None: + from litellm.litellm_core_utils.logging_worker import LoggingWorker + from litellm.responses.streaming_iterator import spawn_logging_task + + worker: Final = LoggingWorker() + started: Final = asyncio.Event() + release: Final = asyncio.Event() + first_flush: Final = asyncio.Event() + pending: Final[asyncio.Queue[str]] = asyncio.Queue() + persisted: Final[asyncio.Queue[str]] = asyncio.Queue() + + async def enqueue_spend() -> None: + pending.put_nowait("accounted request") + + async def callback() -> None: + started.set() + await release.wait() + spawn_logging_task(asyncio.get_running_loop(), enqueue_spend(), task_name="accounting") + + async def flush_spend() -> None: + while not pending.empty(): + persisted.put_nowait(pending.get_nowait()) + first_flush.set() + + worker.ensure_initialized_and_enqueue(callback()) + await started.wait() + shutdown: Final = asyncio.create_task( + litellm.proxy.proxy_server._flush_spend_around_logging_drain_on_shutdown(flush_spend, worker=worker) + ) + try: + await first_flush.wait() + waited_for_callback: Final = not shutdown.done() + release.set() + await shutdown + assert waited_for_callback, "Shutdown returned while accounting callbacks were still running" + assert persisted.get_nowait() == "accounted request" + assert pending.empty() + finally: + release.set() + await worker.flush() + await worker.stop() + + +@pytest.mark.asyncio +async def test_shutdown_reports_worker_timeout_and_flushes_completed_accounting( + monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture +) -> None: + from litellm.litellm_core_utils.logging_worker import LoggingWorker + from litellm.proxy.shutdown.graceful_shutdown_manager import GracefulShutdownManager + + monkeypatch.setenv("GRACEFUL_SHUTDOWN_TIMEOUT", "30") + worker: Final = LoggingWorker() + started: Final = asyncio.Event() + cancelled: Final = asyncio.Event() + flushed_after_cancellation: Final = asyncio.Event() + GracefulShutdownManager.start_shutdown(clock=lambda: 100.0) + + async def callback() -> None: + started.set() + try: + await asyncio.Event().wait() + finally: + cancelled.set() + + async def flush_spend() -> None: + if cancelled.is_set(): + flushed_after_cancellation.set() + + worker.ensure_initialized_and_enqueue(callback()) + await started.wait() + try: + with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): + await litellm.proxy.proxy_server._flush_spend_around_logging_drain_on_shutdown( + flush_spend, worker=worker, clock=lambda: 120.0 + ) + assert cancelled.is_set() + assert flushed_after_cancellation.is_set() + assert "budget exhausted with 1 unfinished callbacks" in caplog.text + finally: + GracefulShutdownManager.reset() + await worker.stop() + + +@pytest.mark.asyncio +async def test_shutdown_bounds_callback_cancellation_and_flushes_accounting( + monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture +) -> None: + from litellm.litellm_core_utils.logging_worker import LoggingWorker + from litellm.proxy.shutdown.graceful_shutdown_manager import GracefulShutdownManager + + monkeypatch.setenv("GRACEFUL_SHUTDOWN_TIMEOUT", "30") + worker: Final = LoggingWorker() + started: Final = asyncio.Event() + cancelling: Final = asyncio.Event() + release: Final = asyncio.Event() + finished: Final = asyncio.Event() + flushed: Final = asyncio.Event() + GracefulShutdownManager.start_shutdown(clock=lambda: 100.0) + + async def callback() -> None: + started.set() + try: + await release.wait() + except asyncio.CancelledError: + cancelling.set() + await release.wait() + finally: + finished.set() + + async def drain_tasks(timeout: float) -> None: + assert timeout == 0.0 + await cancelling.wait() + + async def flush_spend() -> None: + if cancelling.is_set(): + flushed.set() + + worker.ensure_initialized_and_enqueue(callback()) + await started.wait() + try: + with caplog.at_level(logging.WARNING, logger="LiteLLM"): + await litellm.proxy.proxy_server._flush_spend_around_logging_drain_on_shutdown( + flush_spend, worker=worker, clock=lambda: 130.0, drain_tasks=drain_tasks + ) + assert not finished.is_set(), "The cancellation-resistant callback must remain observable" + assert flushed.is_set() + assert "tasks still alive after cancellation" in caplog.text + finally: + release.set() + await worker.flush() + await worker.stop() + GracefulShutdownManager.reset() + + +@pytest.mark.asyncio +async def test_shutdown_preserves_persistence_reserve_after_partial_http_drain( + monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture +) -> None: + from litellm.litellm_core_utils.logging_worker import LoggingWorker + from litellm.proxy.shutdown.graceful_shutdown_manager import GracefulShutdownManager + + monkeypatch.setenv("GRACEFUL_SHUTDOWN_TIMEOUT", "30") + worker: Final = LoggingWorker() + persisted: Final = asyncio.Event() + drain_budgets: Final[asyncio.Queue[float]] = asyncio.Queue() + GracefulShutdownManager.start_shutdown(clock=lambda: 100.0) + + async def flush_spend() -> None: + persisted.set() + + async def drain_tasks(timeout: float) -> None: + drain_budgets.put_nowait(timeout) + + try: + with caplog.at_level(logging.INFO, logger="LiteLLM Proxy"): + await litellm.proxy.proxy_server._flush_spend_around_logging_drain_on_shutdown( + flush_spend, worker=worker, clock=lambda: 106.0, drain_tasks=drain_tasks + ) + assert "remaining_s=24.000 callback_timeout_s=14.000" in caplog.text + assert persisted.is_set() + assert tuple(drain_budgets.get_nowait() for _ in range(drain_budgets.qsize())) == (14.0, 5.0) + finally: + await worker.stop() + GracefulShutdownManager.reset() + + headers = {"Authorization": f"Bearer {token}"} example_completion_result = { @@ -4059,3 +4227,34 @@ def test_detailed_debug_logs_request_body_only_as_received_and_as_sent( assert len(body_dumps) == 2, [message[:120] for message in body_dumps] assert body_dumps[0].startswith("Request received by LiteLLM:") assert "POST Request Sent from LiteLLM" in body_dumps[1] + + +@pytest.mark.asyncio +async def test_shutdown_waits_for_enqueue_producers_before_stopping_worker() -> None: + from litellm.litellm_core_utils.logging_worker import LoggingWorker + from litellm.responses.streaming_iterator import spawn_logging_task + + worker: Final = LoggingWorker() + persisted: Final[asyncio.Queue[str]] = asyncio.Queue() + accounting: Final[asyncio.Queue[str]] = asyncio.Queue() + + async def callback() -> None: + accounting.put_nowait("completed request") + + async def producer() -> None: + await asyncio.sleep(0) + await asyncio.sleep(0) + worker.ensure_initialized_and_enqueue(callback()) + + async def flush_spend() -> None: + while not accounting.empty(): + persisted.put_nowait(accounting.get_nowait()) + + spawn_logging_task(asyncio.get_running_loop(), producer(), task_name="success enqueue") + try: + await litellm.proxy.proxy_server._flush_spend_around_logging_drain_on_shutdown(flush_spend, worker=worker) + assert persisted.qsize() == 1, "An enqueue producer must finish before callback draining stops" + assert persisted.get_nowait() == "completed request" + assert worker.pending_count == 0 + finally: + await worker.stop(discard_pending=True) diff --git a/tests/unit/test_utils.py b/tests/unit/test_utils.py index cd56b6e1b46..f3b470044ae 100644 --- a/tests/unit/test_utils.py +++ b/tests/unit/test_utils.py @@ -11294,6 +11294,7 @@ def test_get_optional_params_forwards_a_tool_config_cache_point_to_a_downstream_ assert optional_params["extra_body"] == {"kept": True, "cache_control_injection_points": [point]} + @pytest.mark.parametrize( ("provider", "expected_warnings"), (("streamlake", 1), ("veniceai", 1), ("my-own-llm", 0)), @@ -11312,3 +11313,34 @@ def test_custom_provider_named_like_a_built_in_warns_once( warnings: Final = [r for r in caplog.records if f"provider '{provider}'" in r.getMessage()] assert len(warnings) == expected_warnings assert provider in litellm._custom_providers + + +@pytest.mark.asyncio +@pytest.mark.parametrize("deferred", (False, True)) +async def test_success_enqueue_is_tracked_without_starting_callbacks_synchronously(deferred: bool) -> None: + from litellm.litellm_core_utils.litellm_logging import Logging + from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER + from litellm.responses.streaming_iterator import drain_logging_tasks + from litellm.utils import _schedule_async_success_logging + + now: Final = datetime(2026, 1, 1, tzinfo=timezone.utc) + logging_obj: Final = Logging( + model="gpt-4o-mini", + messages=[], + stream=False, + call_type="acompletion", + start_time=now, + litellm_call_id="shutdown-enqueue", + function_id="shutdown-enqueue", + ) + logging_obj.defer_async_logging = deferred + _schedule_async_success_logging(logging_obj, ModelResponse(), now, now, False) + if deferred: + assert not logging_obj.async_success_log_enqueued + assert logging_obj.enqueue_deferred_logging is not None + logging_obj.enqueue_deferred_logging() + assert "has_logged_async_success" not in logging_obj.model_call_details + await drain_logging_tasks() + assert GLOBAL_LOGGING_WORKER.pending_count > 0, "Tracked enqueue must finish before worker draining starts" + await GLOBAL_LOGGING_WORKER.flush() + assert logging_obj.model_call_details["has_logged_async_success"] is True From 4bc157e281eab779cb76f78044e291c21604b211 Mon Sep 17 00:00:00 2001 From: Kerry Lu Date: Thu, 29 Oct 2026 09:39:26 -0700 Subject: [PATCH 12/13] feat(proxy): replace the pre-migration row check with a rolled-back dry run of the pending migrations (#48961) * feat(proxy): replace the pre-migration row check with a rolled-back dry run of the pending migrations * feat(proxy): roll back each migration of the dry run on its own and baseline-check databases without a ledger * test(proxy): cover the dry run CLI and PrismaManager wrapper without a database * fix(proxy): refuse the no-ledger dry run for a custom migration folder like boot does, and probe the rolled-back indexes directly * ci(codecov): stop carrying forward the roi-postgres and redis-compat flags Their carried-forward sessions were measured against older proxy_cli.py revisions (633 and 689 line maps against 710 today) and mark docstring lines of the merged file as missed, which failed codecov/patch on this PR at 70% while every unit shard reports the lines hit --------- Co-authored-by: devin-ai-integration[bot] <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- codecov.yaml | 11 +- .../migration_data_checks.py | 252 +----------------- .../litellm_proxy_extras/migration_dry_run.py | 184 +++++++++++++ .../migration_recovery.py | 21 +- .../litellm_proxy_extras/utils.py | 107 ++++---- litellm/proxy/db/prisma_client.py | 12 +- litellm/proxy/proxy_cli.py | 18 +- .../migrations/test_request_log_indexes.py | 248 ++++++++++++----- .../test_migration_data_checks.py | 96 +------ .../test_migration_dry_run.py | 80 ++++++ tests/unit/proxy/db/test_prisma_client.py | 16 ++ tests/unit/proxy/test_proxy_cli.py | 17 ++ 12 files changed, 591 insertions(+), 471 deletions(-) create mode 100644 litellm-proxy-extras/litellm_proxy_extras/migration_dry_run.py create mode 100644 tests/unit/litellm_proxy_extras/test_migration_dry_run.py diff --git a/codecov.yaml b/codecov.yaml index 5b30e78a8fb..521088868bd 100644 --- a/codecov.yaml +++ b/codecov.yaml @@ -15,9 +15,10 @@ ignore: flag_management: default_rules: carryforward: true - # Dead flags no CI job uploads anymore: their carried-forward sessions were - # measured against old revisions, and the stale line maps mark comment lines - # of since-edited files as missed, sinking patch coverage on unrelated PRs. + # Dead flags no CI job uploads anymore, and flags whose job is path-gated on + # PRs: their carried-forward sessions were measured against old revisions, and + # the stale line maps mark comment lines of since-edited files as missed, + # sinking patch coverage on unrelated PRs. individual_flags: - name: proxy-mgmt-behavior carryforward: false @@ -99,6 +100,10 @@ flag_management: carryforward: false - name: proxy-db-endpoints-and-responses carryforward: false + - name: roi-postgres + carryforward: false + - name: redis-compat + carryforward: false component_management: individual_components: diff --git a/litellm-proxy-extras/litellm_proxy_extras/migration_data_checks.py b/litellm-proxy-extras/litellm_proxy_extras/migration_data_checks.py index e22639daaef..75c5bcf4e1e 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/migration_data_checks.py +++ b/litellm-proxy-extras/litellm_proxy_extras/migration_data_checks.py @@ -14,11 +14,9 @@ either, so the migration keeps today's behavior for it. """ import os -import re -from collections.abc import Callable, Iterator +from collections.abc import Callable from dataclasses import dataclass from functools import reduce -from itertools import accumulate from pathlib import Path from types import MappingProxyType from typing import TYPE_CHECKING, Final, Literal @@ -43,7 +41,7 @@ MIGRATION_DATA_CHECK_TIMEOUT_ENV_VAR: Final = "LITELLM_MIGRATION_DATA_CHECK_TIME DEFAULT_MIGRATION_DATA_CHECK_TIMEOUT_SECONDS: Final = 600.0 REFUSAL: Final = "Existing rows would make a pending schema change fail:" -Kind = Literal["too_long", "duplicates", "nulls", "any_row", "orphans", "check_violations"] +Kind = Literal["too_long"] @dataclass(frozen=True, slots=True) @@ -54,9 +52,6 @@ class Check: statement: str source: str index: str | None = None - referenced_table: str | None = None - referenced_columns: tuple[str, ...] = () - expression: str | None = None @dataclass(frozen=True, slots=True) @@ -100,33 +95,6 @@ class ColumnType: ColumnTypeLookup = Callable[[str, str], ColumnType | None] -_IDENT: Final = r'"?(\w+)"?' -_QUALIFIED: Final = r'(?:"?\w+"?\s*\.\s*)?"?(\w+)"?' -_CREATE_INDEX: Final = re.compile( - rf"^\s*CREATE\s+(UNIQUE\s+)?INDEX\s+(?:CONCURRENTLY\s+)?(?:IF\s+NOT\s+EXISTS\s+)?{_IDENT}\s+ON\s+(?:ONLY\s+)?" - rf"{_QUALIFIED}\s*(?:USING\s+(\w+)\s*)?\((.*?)\)", - re.IGNORECASE | re.DOTALL, -) -_ALTER_TABLE: Final = re.compile( - rf"^\s*ALTER\s+TABLE\s+(?:IF\s+EXISTS\s+)?(?:ONLY\s+)?{_QUALIFIED}\s+(.*)$", re.IGNORECASE | re.DOTALL -) -_SET_NOT_NULL: Final = re.compile(rf"^ALTER\s+(?:COLUMN\s+)?{_IDENT}\s+SET\s+NOT\s+NULL$", re.IGNORECASE) -_ADD_COLUMN: Final = re.compile( - rf"^ADD\s+(?:COLUMN\s+)?(?:IF\s+NOT\s+EXISTS\s+)?(?!(?:CONSTRAINT|FOREIGN|UNIQUE|CHECK|PRIMARY|EXCLUDE)\b){_IDENT}\s+(.*)$", - re.IGNORECASE | re.DOTALL, -) -_CONSTRAINT: Final = rf"^ADD\s+(?:CONSTRAINT\s+{_IDENT}\s+)?" -_ADD_FOREIGN_KEY: Final = re.compile( - rf"{_CONSTRAINT}FOREIGN\s+KEY\s*\(([^)]*)\)\s*REFERENCES\s+{_QUALIFIED}\s*\(([^)]*)\)(.*)$", - re.IGNORECASE | re.DOTALL, -) -_ADD_UNIQUE: Final = re.compile(rf"{_CONSTRAINT}UNIQUE\s*\(([^)]*)\)", re.IGNORECASE | re.DOTALL) -_ADD_CHECK: Final = re.compile(rf"{_CONSTRAINT}CHECK\s*\((.*)\)\s*(NOT\s+VALID)?\s*$", re.IGNORECASE | re.DOTALL) -_PLAIN_COLUMN: Final = re.compile(rf"^{_IDENT}(?:\s+(?:ASC|DESC|NULLS\s+(?:FIRST|LAST)))*$", re.IGNORECASE) -_COMMENT: Final = re.compile(r"--[^\n]*") -_NOT_NULL: Final = re.compile(r"\bNOT\s+NULL\b", re.IGNORECASE) -_HAS_DEFAULT: Final = re.compile(r"\bDEFAULT\b|\bGENERATED\b", re.IGNORECASE) - def migration_data_check_timeout_seconds() -> float: """Seconds one count query may run before the preflight gives up on verifying it.""" @@ -145,117 +113,6 @@ def migration_data_check_timeout_seconds() -> float: return DEFAULT_MIGRATION_DATA_CHECK_TIMEOUT_SECONDS -def statements(migration_sql: str) -> tuple[str, ...]: - """The statements of one migration file, comments stripped, empty ones dropped.""" - parts: Final = _COMMENT.sub("", migration_sql).split(";") - return tuple(part.strip() for part in parts if part.strip()) - - -def _split_top_level(text: str) -> tuple[str, ...]: - """Split on the commas outside parentheses.""" - depth: Final = accumulate((1 if char == "(" else -1 if char == ")" else 0) for char in text) - cuts: Final = tuple( - position for position, (char, level) in enumerate(zip(text, depth, strict=True)) if char == "," and level == 0 - ) - starts: Final = (0, *(cut + 1 for cut in cuts)) - stops: Final = (*cuts, len(text)) - pieces: Final = tuple(text[start:stop].strip() for start, stop in zip(starts, stops, strict=True)) - return tuple(piece for piece in pieces if piece) - - -def _plain_columns(column_list: str) -> tuple[str, ...] | None: - """The column names of an index or constraint column list, None when one entry is an - expression the derivation cannot turn into a count.""" - matches: Final = tuple(_PLAIN_COLUMN.match(entry) for entry in _split_top_level(column_list)) - if any(match is None for match in matches): - return None - return tuple(match.group(1) for match in matches if match is not None) - - -def _index_checks(match: "re.Match[str]", statement_text: str, source: str) -> Iterator[Check]: - unique, name, table, method, column_list = match.groups() - columns: Final = _plain_columns(column_list) - if columns is None: - return - label: Final = f'CREATE {"UNIQUE " if unique else ""}INDEX "{name}"' - if method is None or method.lower() == "btree": - yield Check("too_long", table, columns, label, source, index=name) - if unique: - yield Check("duplicates", table, columns, label, source, index=name) - - -def _alter_table_checks(table: str, action: str, source: str) -> Iterator[Check]: - if (set_not_null := _SET_NOT_NULL.match(action)) is not None: - yield Check("nulls", table, (set_not_null.group(1),), f'ALTER TABLE "{table}" {action}', source) - return - if (add_column := _ADD_COLUMN.match(action)) is not None: - column, body = add_column.groups() - if _NOT_NULL.search(body) is not None and _HAS_DEFAULT.search(body) is None: - yield Check( - "any_row", table, (column,), f'ALTER TABLE "{table}" ADD COLUMN "{column}" ... NOT NULL', source - ) - return - if (foreign_key := _ADD_FOREIGN_KEY.match(action)) is not None: - constraint, columns, referenced_table, referenced_columns, tail = foreign_key.groups() - own: Final = _plain_columns(columns) - theirs: Final = _plain_columns(referenced_columns) - if own is None or theirs is None or len(own) != len(theirs) or "NOT VALID" in tail.upper(): - return - yield Check( - "orphans", - table, - own, - _constraint_label(table, constraint, "FOREIGN KEY"), - source, - referenced_table=referenced_table, - referenced_columns=theirs, - ) - return - if (unique := _ADD_UNIQUE.match(action)) is not None: - constraint, columns = unique.groups() - unique_columns: Final = _plain_columns(columns) - if unique_columns is not None: - yield Check("duplicates", table, unique_columns, _constraint_label(table, constraint, "UNIQUE"), source) - return - if (check := _ADD_CHECK.match(action)) is not None: - constraint, expression, not_valid = check.groups() - if not_valid is None: - yield Check( - "check_violations", - table, - (), - _constraint_label(table, constraint, "CHECK"), - source, - expression=expression, - ) - - -def _constraint_label(table: str, constraint: str | None, kind: str) -> str: - named: Final = f'CONSTRAINT "{constraint}" ' if constraint else "" - return f'ALTER TABLE "{table}" ADD {named}{kind}' - - -def checks_for_statement(statement_text: str, source: str) -> tuple[Check, ...]: - """The count queries one SQL statement needs before it can be applied, empty for a - statement that cannot fail on existing rows or that the derivation cannot read.""" - if (index := _CREATE_INDEX.match(statement_text)) is not None: - return tuple(_index_checks(index, statement_text, source)) - if (alter := _ALTER_TABLE.match(statement_text)) is not None: - table, actions = alter.groups() - return tuple( - check for action in _split_top_level(actions) for check in _alter_table_checks(table, action, source) - ) # comprehension-ok: flattening the per-action checks of one statement - return () - - -def checks_for_migration(name: str, migration_sql: str) -> tuple[Check, ...]: - return tuple( - check - for statement_text in statements(migration_sql) - for check in checks_for_statement(statement_text, f"migration {name}") - ) # comprehension-ok: flattening the per-statement checks of one migration - - def checks_for_index(table: str, name: str, columns: tuple[str, ...]) -> tuple[Check, ...]: """The entry-size check a request-log index needs before `CREATE INDEX CONCURRENTLY`.""" return (Check("too_long", table, columns, f'CREATE INDEX "{name}"', "request-log index", index=name),) @@ -424,42 +281,7 @@ def _entry_size_query(check: Check, schema: str, types: tuple[ColumnType, ...], def _count_query(check: Check, schema: str, types: tuple[ColumnType, ...], limit_bytes: int) -> "sql.Composed": - from psycopg import sql - - table: Final = sql.Identifier(schema, check.table) - columns: Final = sql.SQL(", ").join(sql.Identifier(column) for column in check.columns) - match check.kind: - case "too_long": - return _entry_size_query(check, schema, types, limit_bytes) - case "duplicates": - not_null: Final = sql.SQL(" AND ").join( - sql.SQL("{} IS NOT NULL").format(sql.Identifier(column)) for column in check.columns - ) - return sql.SQL( - "SELECT count(*) FROM (SELECT 1 FROM {t} WHERE {nn} GROUP BY {cols} HAVING count(*) > 1) AS d" - ).format(t=table, nn=not_null, cols=columns) - case "nulls": - return sql.SQL("SELECT count(*) FROM {t} WHERE {col} IS NULL").format( - t=table, col=sql.Identifier(check.columns[0]) - ) - case "any_row": - return sql.SQL("SELECT count(*) FROM {t}").format(t=table) - case "orphans": - pairs: Final = sql.SQL(" AND ").join( - sql.SQL("parent.{} = child.{}").format(sql.Identifier(theirs), sql.Identifier(own)) - for own, theirs in zip(check.columns, check.referenced_columns, strict=True) - ) - present: Final = sql.SQL(" AND ").join( - sql.SQL("child.{} IS NOT NULL").format(sql.Identifier(column)) for column in check.columns - ) - return sql.SQL( - "SELECT count(*) FROM {t} AS child WHERE {present} AND NOT EXISTS (SELECT 1 FROM {ref} AS parent WHERE {pairs})" - ).format(t=table, present=present, ref=sql.Identifier(schema, check.referenced_table or ""), pairs=pairs) - case "check_violations": - return sql.SQL("SELECT count(*) FROM {t} WHERE NOT ({expr})").format( - t=table, - expr=sql.SQL(check.expression or "true"), # pyright: ignore[reportArgumentType] # copied from our own migration file - ) + return _entry_size_query(check, schema, types, limit_bytes) def _applies( @@ -474,15 +296,7 @@ def _applies( return False if check.index is not None and _index_is_valid(connection, schema, check.index): return False - if check.kind == "any_row": - return types is None - if types is None: - return False - if check.kind == "too_long": - return any(column.can_overflow_btree for column in types) - if check.kind == "orphans": - return check.referenced_table is not None and _table_exists(connection, schema, check.referenced_table) - return True + return types is not None and any(column.can_overflow_btree for column in types) def _too_long_detail(check: Check, types: tuple[ColumnType, ...], values: tuple[object, ...], limit_bytes: int) -> str: @@ -515,21 +329,7 @@ def _run_check( rows: Final = _int(values[0], 0) if values else 0 if rows == 0: return None - match check.kind: - case "too_long": - return BlockedChange(check, rows, _too_long_detail(check, types, values, limit_bytes)) - case "duplicates": - return BlockedChange(check, rows, "duplicated value groups") - case "nulls": - return BlockedChange(check, rows, "NULL") - case "any_row": - return BlockedChange( - check, rows, "rows, and a NOT NULL column with no DEFAULT cannot be added to a table with rows" - ) - case "orphans": - return BlockedChange(check, rows, f'with no matching row in "{check.referenced_table}"') - case "check_violations": - return BlockedChange(check, rows, f"failing CHECK ({check.expression})") + return BlockedChange(check, rows, _too_long_detail(check, types, values, limit_bytes)) def find_blockers( @@ -571,19 +371,7 @@ def find_blockers( def _remedy(check: Check) -> str: joined: Final = ", ".join(check.columns) - match check.kind: - case "too_long": - return f'Shorten or hash the long "{check.table}" ({joined}) values so every row\'s index entry fits.' - case "duplicates": - return f'Remove or merge the duplicate "{check.table}" ({joined}) values.' - case "nulls": - return f'Fill the NULL "{check.table}"."{joined}" values, or give the column a default first.' - case "any_row": - return "Give the new column a DEFAULT in the migration, or add it nullable and fill it afterwards." - case "orphans": - return f'Delete or repoint the "{check.table}" ({joined}) values with no "{check.referenced_table}" row.' - case "check_violations": - return f'Fix the "{check.table}" rows that fail the expression before adding the constraint.' + return f'Shorten or hash the long "{check.table}" ({joined}) values so every row\'s index entry fits.' def _describe(finding: Finding) -> str: @@ -607,30 +395,4 @@ def _describe(finding: Finding) -> str: def report(findings: tuple[Finding, ...]) -> str: """The operator-facing refusal: one block per finding, then what to do next.""" lines: Final = "\n".join(_describe(finding) for finding in findings) - return ( - f"{REFUSAL}\n{lines}\n" - "No schema change was made. Fix the rows and start again, or run `litellm --check_migrations_against_existing_data` with the new image to check again." - ) - - -def pending_change_checks( - connection: "psycopg.Connection[tuple[object, ...]]", - schema: str, - migrations_dir: Path | None, - indexes: tuple[tuple[str, str, tuple[str, ...]], ...], -) -> tuple[Check, ...]: - """Every check the pending migrations and the request-log indexes need. `migrations_dir` - is None when the schema is pushed rather than migrated, so only the indexes are checked.""" - migration_checks: Final = ( - tuple( - check - for name, migration_sql in pending_migrations(connection, schema, migrations_dir) - for check in checks_for_migration(name, migration_sql) - ) # comprehension-ok: flattening the per-migration checks - if migrations_dir is not None - else () - ) - index_checks: Final = tuple( - check for table, name, columns in indexes for check in checks_for_index(table, name, columns) - ) # comprehension-ok: flattening the per-index checks - return migration_checks + index_checks + return f"{REFUSAL}\n{lines}\nNo index was built. Fix the rows and the next index build picks it up." diff --git a/litellm-proxy-extras/litellm_proxy_extras/migration_dry_run.py b/litellm-proxy-extras/litellm_proxy_extras/migration_dry_run.py new file mode 100644 index 00000000000..c24742ce145 --- /dev/null +++ b/litellm-proxy-extras/litellm_proxy_extras/migration_dry_run.py @@ -0,0 +1,184 @@ +"""`litellm --check_migrations_against_existing_data`: run the pending migrations for real, each +inside its own transaction that is always rolled back, so Postgres itself says whether this +database's rows and schema take them. Each migration holds the locks the real upgrade will, +and releases them before the next one runs.""" + +from __future__ import annotations + +import re +import time +from collections.abc import Callable +from dataclasses import dataclass +from functools import reduce +from pathlib import Path +from typing import TYPE_CHECKING, Final + +from litellm_proxy_extras.migration_data_checks import pending_migrations + +if TYPE_CHECKING: + import psycopg + +_CONCURRENTLY: Final = re.compile( + r"\b((?:CREATE\s+(?:UNIQUE\s+)?INDEX|DROP\s+INDEX|REINDEX\s+\w+)\s+)CONCURRENTLY\s+", re.IGNORECASE +) +UNDEFINED_SQLSTATES: Final = frozenset({"42703", "42P01", "42704", "42883"}) + +Migration = tuple[str, str] + + +@dataclass(frozen=True, slots=True) +class Applied: + name: str + seconds: float + replayed: tuple[str, ...] = () + + +@dataclass(frozen=True, slots=True) +class Rejected: + name: str + error: str + seconds: float + + +@dataclass(frozen=True, slots=True) +class Baseline: + """A database with tables but no `_prisma_migrations`: the boot baselines it instead of replaying + history, so the dry run asks the same question the boot does. `refusal` is None when the boot + will baseline, otherwise its reason (a custom migration folder, or the schema diff).""" + + refusal: str | None + + @property + def matches(self) -> bool: + return self.refusal is None + + +@dataclass(frozen=True, slots=True) +class DryRun: + applied: tuple[Applied, ...] + rejected: Rejected | None + baseline: Baseline | None = None + + @property + def ok(self) -> bool: + return self.rejected is None and (self.baseline is None or self.baseline.matches) + + +def transactional(migration_sql: str) -> str: + """The migration as it can run inside a transaction: `CONCURRENTLY` index builds become plain ones.""" + return _CONCURRENTLY.sub(r"\1", migration_sql) + + +def _attempt( + connection: psycopg.Connection[tuple[object, ...]], + schema: str, + chain: tuple[Migration, ...], + clock: Callable[[], float], +) -> tuple[float, psycopg.Error | None]: + import psycopg + from psycopg import sql + + started: Final = clock() + try: + with connection.transaction(force_rollback=True): + connection.execute(sql.SQL("SET LOCAL search_path TO {}").format(sql.Identifier(schema))) + for _, migration_sql in chain: + connection.execute(transactional(migration_sql)) + except psycopg.Error as exc: + return clock() - started, exc + return clock() - started, None + + +def _run( + connection: psycopg.Connection[tuple[object, ...]], + schema: str, + pending: tuple[Migration, ...], + position: int, + done: DryRun, + clock: Callable[[], float], +) -> DryRun: + """One migration in its own rolled-back transaction. A migration that fails because something + an earlier pending migration creates does not exist yet is rerun with those replayed in front.""" + if done.rejected is not None: + return done + name: Final = pending[position][0] + seconds, error = _attempt(connection, schema, (pending[position],), clock) + if error is None: + return DryRun(done.applied + (Applied(name, seconds),), None, done.baseline) + if error.sqlstate not in UNDEFINED_SQLSTATES or position == 0: + return DryRun(done.applied, Rejected(name, str(error).strip(), seconds), done.baseline) + earlier: Final = pending[:position] + chained_seconds, chained_error = _attempt(connection, schema, earlier + (pending[position],), clock) + if chained_error is None: + replayed: Final = tuple(earlier_name for earlier_name, _ in earlier) + return DryRun(done.applied + (Applied(name, chained_seconds, replayed),), None, done.baseline) + return DryRun(done.applied, Rejected(name, str(chained_error).strip(), chained_seconds), done.baseline) + + +def _has_ledger(connection: psycopg.Connection[tuple[object, ...]], schema: str) -> bool: + from psycopg import sql + + row: Final = connection.execute( + "SELECT to_regclass(%s)", (sql.Identifier(schema, "_prisma_migrations").as_string(connection),) + ).fetchone() + return row is not None and row[0] is not None + + +def _has_litellm_tables(connection: psycopg.Connection[tuple[object, ...]], schema: str) -> bool: + row: Final = connection.execute( + "SELECT 1 FROM information_schema.tables WHERE table_schema = %s AND table_name LIKE 'LiteLLM\\_%%' LIMIT 1", + (schema,), + ).fetchone() + return row is not None + + +def dry_run( + connection: psycopg.Connection[tuple[object, ...]], + schema: str, + migrations_dir: Path, + baseline_check: Callable[[], str | None] | None = None, + clock: Callable[[], float] = time.monotonic, +) -> DryRun: + """Apply every pending migration in order, each in its own transaction that is rolled back + before the next one starts. The first migration Postgres rejects ends the run with its error. + `baseline_check` (None when the boot would baseline, otherwise why it refuses) stands in for + the replay on a database that has tables but no migration history, which the boot baselines.""" + if baseline_check is not None and not _has_ledger(connection, schema) and _has_litellm_tables(connection, schema): + return DryRun((), None, Baseline(baseline_check())) + pending: Final = pending_migrations(connection, schema, migrations_dir) + return reduce( + lambda done, position: _run(connection, schema, pending, position, done, clock), + range(len(pending)), + DryRun((), None), + ) + + +def _timing(applied: Applied) -> str: + replay: Final = f" (depends on {', '.join(applied.replayed)}, replayed in front of it)" if applied.replayed else "" + return f"\n {applied.name}: {applied.seconds:.1f} s{replay}" + + +def report(result: DryRun) -> str: + """What the operator reads: the per-migration time the real upgrade will hold its locks, + or the migration Postgres rejected with Postgres's own words. Nothing was applied either way.""" + if result.baseline is not None: + if result.baseline.matches: + return ( + "No migration history and the schema matches this build, " + "the proxy baselines it at boot, nothing to dry-run." + ) + return f"No migration history and the proxy cannot baseline it at boot: {result.baseline.refusal}" + timings: Final = "".join(_timing(applied) for applied in result.applied) + if result.rejected is None: + count: Final = len(result.applied) + if count == 0: + return "No pending migrations, this database is on this version's schema." + return ( + f"{count} pending migration(s) applied one at a time in a dry run, each rolled back before the next " + f"so its locks are released, nothing changed:{timings}" + ) + return ( + f"Migration {result.rejected.name} fails on this database after {result.rejected.seconds:.1f} s, " + f"nothing was applied.\n{result.rejected.error}" + + (f"\nApplied before it in the dry run:{timings}" if timings else "") + ) diff --git a/litellm-proxy-extras/litellm_proxy_extras/migration_recovery.py b/litellm-proxy-extras/litellm_proxy_extras/migration_recovery.py index 5a55b35b255..1c16eb5e6b7 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/migration_recovery.py +++ b/litellm-proxy-extras/litellm_proxy_extras/migration_recovery.py @@ -86,6 +86,18 @@ def migration_files(directory: Path) -> tuple[tuple[str, str], ...]: ) +def has_packaged_history(migrations_dir: Path) -> bool: + """Whether `migrations_dir` carries exactly the migration files and schema.prisma this package ships, + the only history the boot baselines an existing database onto.""" + packaged_dir: Final = Path(__file__).parent + migrations: Final = migration_files(migrations_dir) + return ( + bool(migrations) + and migrations == migration_files(packaged_dir) + and (migrations_dir / "schema.prisma").read_bytes() == (packaged_dir / "schema.prisma").read_bytes() + ) + + def baseline_current_schema( coordinator: MigrationCoordinator, schema: str, @@ -95,14 +107,9 @@ def baseline_current_schema( ) -> None: from psycopg import sql - packaged_dir: Final = Path(__file__).parent - migrations: Final = migration_files(migrations_dir) - if ( - not migrations - or migrations != migration_files(packaged_dir) - or (migrations_dir / "schema.prisma").read_bytes() != (packaged_dir / "schema.prisma").read_bytes() - ): + if not has_packaged_history(migrations_dir): raise RuntimeError("Cannot automatically baseline an existing database with custom migration history.") + migrations: Final = migration_files(migrations_dir) coordinator.acquire_prisma_lock() existing: Final = coordinator.connection.execute( diff --git a/litellm-proxy-extras/litellm_proxy_extras/utils.py b/litellm-proxy-extras/litellm_proxy_extras/utils.py index 528bf7def20..8b19576d86e 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/utils.py +++ b/litellm-proxy-extras/litellm_proxy_extras/utils.py @@ -11,18 +11,13 @@ from collections.abc import Callable from contextlib import AbstractContextManager, nullcontext from dataclasses import dataclass, replace from pathlib import Path -from typing import TYPE_CHECKING, Final, Optional, Protocol, Union +from typing import TYPE_CHECKING, Final, Protocol from litellm_proxy_extras import prisma_toolchain from litellm_proxy_extras._logging import logger from litellm_proxy_extras._redaction import _redact_credentials -from litellm_proxy_extras.migration_data_checks import ( - Finding, - find_blockers, - migration_data_check_timeout_seconds, - pending_change_checks, - report, -) +from litellm_proxy_extras.migration_dry_run import dry_run +from litellm_proxy_extras.migration_dry_run import report as dry_run_report from litellm_proxy_extras.migration_lock import held_migration_lock from litellm_proxy_extras.prisma_toolchain import ( PRISMA_COMMAND_TIMEOUT_ENV_VAR, @@ -36,7 +31,6 @@ from litellm_proxy_extras.replica_identity import ( apply_replica_identity_full, ) from litellm_proxy_extras.request_log_indexes import ( - REQUEST_LOG_INDEXES, ensure_request_log_indexes, filter_request_log_index_diff, ) @@ -52,7 +46,7 @@ class LensCheckConnect(Protocol): ) -> "psycopg.Connection[tuple[object, ...]]": ... -def str_to_bool(value: Optional[str]) -> bool: +def str_to_bool(value: str | None) -> bool: if value is None: return False return value.lower() in ("true", "1", "t", "y", "yes") @@ -169,7 +163,7 @@ def _without_sql_comments(statement: str) -> str: ).strip() -def _without_spend_logs_pk_clauses(statement: str) -> Optional[str]: +def _without_spend_logs_pk_clauses(statement: str) -> str | None: prefix_match = _SPEND_LOGS_ALTER_RE.match(statement) if not prefix_match: return statement @@ -217,7 +211,7 @@ def _max_migration_timestamp(names) -> int: return max(_migration_timestamp(n) for n in names) -def _redacted_command(command: object) -> Union[str, tuple[str, ...], list[str]]: +def _redacted_command(command: object) -> str | tuple[str, ...] | list[str]: if isinstance(command, tuple): return tuple(_redact_credentials(str(argument)) for argument in command) if isinstance(command, list): @@ -404,7 +398,7 @@ class ProxyExtrasDBManager: return tuple(row) if row is not None else () @staticmethod - def _failed_migration_logs(migration_name: str, started_at: str) -> Optional[str]: + def _failed_migration_logs(migration_name: str, started_at: str) -> str | None: row: Final = ProxyExtrasDBManager._read_migration_ledger( "SELECT logs FROM {} WHERE migration_name = %s AND started_at = %s::timestamptz " "AND finished_at IS NULL AND rolled_back_at IS NULL", @@ -721,10 +715,10 @@ class ProxyExtrasDBManager: return row is not None @staticmethod - def _prisma_schema_param(url: str) -> Optional[str]: + def _prisma_schema_param(url: str) -> str | None: """The `schema` query param Prisma uses to pick its target schema, or None when the URL does not set one.""" - from urllib.parse import urlparse, parse_qsl + from urllib.parse import parse_qsl, urlparse return next( (v for k, v in parse_qsl(urlparse(url).query) if k == "schema"), @@ -1424,7 +1418,6 @@ class ProxyExtrasDBManager: bool: True if setup was successful, False otherwise """ ensure_prisma_toolchain(prisma_command=_get_prisma_command(), prisma_env=_get_prisma_env()) - ProxyExtrasDBManager.raise_if_rows_block_pending_changes(use_migrate=use_migrate) migrated = ProxyExtrasDBManager._run_migrations(use_migrate=use_migrate, use_v2_resolver=use_v2_resolver) if not migrated: return False @@ -1433,46 +1426,69 @@ class ProxyExtrasDBManager: return True @staticmethod - def migration_data_check_findings(use_migrate: bool = True) -> tuple[Finding, ...]: - """What in the existing rows would make a pending migration or request-log index - fail, read with count queries only: no write, no table lock. Empty when nothing - blocks, when no database is configured, or when it cannot be reached, so the - migration step that follows reports the connection problem the way it does today.""" + def _baseline_refusal(prisma_dir: Path) -> str | None: + """Why `baseline_current_schema` would refuse this database at boot, None when it would baseline: + the migrations folder must be the packaged one, then `prisma migrate diff` must find the live + database and this build's schema.prisma identical.""" + from litellm_proxy_extras.migration_lock import migration_environment + from litellm_proxy_extras.migration_recovery import has_packaged_history + + if not has_packaged_history(prisma_dir): + return ( + f"{prisma_dir} is a custom migration history (LITELLM_MIGRATION_DIR), it differs from the " + "migrations this build ships, and the proxy only baselines onto the packaged history." + ) + try: + prisma_toolchain.run_prisma( + ( + _get_prisma_command(), + "migrate", + "diff", + "--from-schema-datasource", + str(prisma_dir / "schema.prisma"), + "--to-schema-datamodel", + str(prisma_dir / "schema.prisma"), + "--exit-code", + ), + timeout=prisma_toolchain.prisma_command_timeout(), + env=migration_environment(_get_prisma_env()), + ) + except subprocess.CalledProcessError as exc: + diff: Final = str(exc.stdout or exc.stderr or exc).strip() + return f"the schema differs from this build. prisma migrate diff:\n{diff}" + except subprocess.TimeoutExpired as exc: + return f"prisma migrate diff did not finish within {exc.timeout} s" + return None + + @staticmethod + def migration_dry_run_summary(use_migrate: bool = True) -> tuple[bool, str]: + """`--check_migrations_against_existing_data`: the pending migrations applied for real on + DIRECT_URL or DATABASE_URL inside one transaction that is rolled back, so Postgres reports + what the upgrade would hit. (True, the per-migration timings) when they all apply, + (False, Postgres's own error) when one does not.""" import psycopg + if not use_migrate: + return (True, "The schema is pushed rather than migrated, there is nothing to dry-run.") database_url: Final = os.environ.get("DATABASE_URL") if not database_url: - return () + return (False, "DATABASE_URL is not set.") direct_url: Final = ProxyExtrasDBManager._strip_prisma_query_params( os.environ.get("DIRECT_URL") or database_url ) schema: Final = ProxyExtrasDBManager._prisma_schema_param(database_url) or "public" - migrations_dir: Final = Path(ProxyExtrasDBManager._get_prisma_dir()) / "migrations" if use_migrate else None - indexes: Final = tuple((index.table, index.name, index.columns) for index in REQUEST_LOG_INDEXES) + prisma_dir: Final = Path(ProxyExtrasDBManager._get_prisma_dir()) try: with psycopg.connect(direct_url, connect_timeout=10, autocommit=True) as connection: - checks: Final = pending_change_checks(connection, schema, migrations_dir, indexes) - return find_blockers(connection, schema, checks, migration_data_check_timeout_seconds()) + result: Final = dry_run( + connection, + schema, + prisma_dir / "migrations", + baseline_check=lambda: ProxyExtrasDBManager._baseline_refusal(prisma_dir), + ) except psycopg.OperationalError as exc: - logger.warning( - "Could not reach the database for the pre-migration row check: %s", _redact_credentials(str(exc)) - ) - return () - - @staticmethod - def migration_data_check_report(use_migrate: bool = True) -> Optional[str]: - """The operator-facing refusal for `migration_data_check_findings`, None when nothing blocks.""" - findings: Final = ProxyExtrasDBManager.migration_data_check_findings(use_migrate=use_migrate) - return report(findings) if findings else None - - @staticmethod - def raise_if_rows_block_pending_changes(use_migrate: bool) -> None: - """Refuse to run migrations that the existing rows would make fail, before any is - applied. The RuntimeError carries the report and reaches the operator through the - proxy's "Database migration cannot proceed" exit.""" - refusal: Final = ProxyExtrasDBManager.migration_data_check_report(use_migrate=use_migrate) - if refusal is not None: - raise RuntimeError(refusal) + return (False, f"Could not reach the database: {_redact_credentials(str(exc))}") + return (result.ok, dry_run_report(result)) @staticmethod def build_request_log_indexes(build: Callable[[str, str], bool] = ensure_request_log_indexes) -> bool: @@ -1724,5 +1740,4 @@ class ProxyExtrasDBManager: time.sleep(random.randrange(5, 15)) finally: os.chdir(original_dir) - pass return False diff --git a/litellm/proxy/db/prisma_client.py b/litellm/proxy/db/prisma_client.py index 3d783d8aba1..be5de040fdf 100644 --- a/litellm/proxy/db/prisma_client.py +++ b/litellm/proxy/db/prisma_client.py @@ -1127,16 +1127,16 @@ class PrismaManager: return ProxyExtrasDBManager.build_request_log_indexes() @staticmethod - def migration_data_check_report(use_migrate: bool = True) -> str | None: - """The read-only pre-migration row check on its own, for `--check_migrations_against_existing_data`: the - refusal the migration step would print, or None when the pending schema changes - can be applied to the rows the database has.""" + def migration_dry_run_summary(use_migrate: bool = True) -> tuple[bool, str]: + """`--check_migrations_against_existing_data` on its own: the pending migrations applied inside a + transaction that is rolled back, as (whether they all applied, what the operator reads).""" from litellm_proxy_extras.utils import ( ProxyExtrasDBManager, # pyright: ignore[reportUnknownVariableType] # litellm-proxy-extras is an optional install ) + from pydantic import TypeAdapter - refusal: Final[object] = ProxyExtrasDBManager.migration_data_check_report(use_migrate=use_migrate) # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] # litellm-proxy-extras is an optional install - return refusal if isinstance(refusal, str) else None + summary: Final[object] = ProxyExtrasDBManager.migration_dry_run_summary(use_migrate=use_migrate) # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] # litellm-proxy-extras is an optional install + return TypeAdapter(tuple[bool, str]).validate_python(summary) @staticmethod def start_request_log_index_build() -> None: diff --git a/litellm/proxy/proxy_cli.py b/litellm/proxy/proxy_cli.py index 0d8ab98a0b1..f71d1960e24 100644 --- a/litellm/proxy/proxy_cli.py +++ b/litellm/proxy/proxy_cli.py @@ -276,18 +276,19 @@ class ProxyInitializationHelpers: @staticmethod def run_migration_data_check(use_migrate: bool) -> int: - """`--check_migrations_against_existing_data`: the pre-migration row check against DATABASE_URL, nothing applied. - 0 when the pending schema changes fit the rows, 2 with the refusal when they do not.""" + """`--check_migrations_against_existing_data`: the pending migrations applied to DATABASE_URL inside one + transaction that is rolled back, so Postgres reports what the upgrade would hit. 0 when they all apply, + 2 with Postgres's error when one does not.""" if os.getenv("DATABASE_URL") is None: raise click.UsageError("--check_migrations_against_existing_data requires DATABASE_URL") from litellm.proxy.db.prisma_client import PrismaManager - refusal: Final[str | None] = PrismaManager.migration_data_check_report(use_migrate=use_migrate) # pyright: ignore[reportUnknownVariableType] # litellm-proxy-extras is an optional install - if refusal is None: - print("LiteLLM Proxy: the pending migrations and indexes can be applied to the rows this database has.") + ok, summary = PrismaManager.migration_dry_run_summary(use_migrate=use_migrate) + if ok: + print(f"LiteLLM Proxy: {summary}") return 0 print( - f"\033[1;31mLiteLLM Proxy: Database migration cannot proceed. {refusal}\033[0m", file=sys.stderr, flush=True + f"\033[1;31mLiteLLM Proxy: Database migration cannot proceed. {summary}\033[0m", file=sys.stderr, flush=True ) return 2 @@ -1059,8 +1060,9 @@ class ProxyInitializationHelpers: "--check_migrations_against_existing_data", is_flag=True, default=False, - help="Check, read-only, that the rows in DATABASE_URL can take this version's pending migrations and indexes, " - "print what would block them, then exit without applying anything. Exit code 2 when something blocks.", + help="Dry-run this version's pending migrations on DATABASE_URL: the real SQL inside one transaction that is " + "rolled back, so Postgres reports the row or schema that would make the upgrade fail, then exit. Holds the same " + "locks the upgrade will. Exit code 2 when a migration fails.", ) @click.option( "--keepalive_timeout", diff --git a/tests/integration/database/migrations/test_request_log_indexes.py b/tests/integration/database/migrations/test_request_log_indexes.py index f558c355c84..d06f0b13cff 100644 --- a/tests/integration/database/migrations/test_request_log_indexes.py +++ b/tests/integration/database/migrations/test_request_log_indexes.py @@ -16,7 +16,7 @@ from typing import Final import psycopg import pytest -from litellm_proxy_extras import migration_data_checks, request_log_indexes +from litellm_proxy_extras import migration_data_checks, migration_dry_run, request_log_indexes from litellm_proxy_extras.migration_lock import MIGRATION_LOCK_KEY, migration_lock from litellm_proxy_extras.migration_recovery import roll_back_failed_inert_migration from litellm_proxy_extras.request_log_indexes import ( @@ -967,33 +967,25 @@ _HASH_LEGACY_KEYS: Final = ( @requires_db @RESOLVERS @pytest.mark.parametrize("release", (API_KEY_INDEX_MIGRATION,), ids=("v1.102.1",)) -def test_the_migration_job_refuses_before_the_first_migration_when_spend_log_keys_are_too_long_for_the_index( +def test_the_migration_job_applies_every_migration_and_only_the_index_build_waits_for_too_long_keys( scratch_database: str, use_v2_resolver: bool ) -> None: + """Rows never stop an upgrade. The migrations go through, and the one request-log index + whose keys do not fit is left for the next build with the rows named in the log.""" with psycopg.connect(scratch_database, autocommit=True) as conn: _insert_spend_log_with_key(conn, "legacy-1", _legacy_bearer_token()) _insert_spend_log_with_key(conn, "legacy-2", _legacy_bearer_token()) _insert_spend_log_with_key(conn, "hashed-1", "a" * 64) - _insert_spend_log_with_key(conn, "repetitive-1", "z" * 3000) ledger_before: Final = _ledger(scratch_database) - with pytest.raises(RuntimeError) as refused: - ProxyExtrasDBManager.run_migration_job(use_migrate=True, use_v2_resolver=use_v2_resolver) - - refusal: Final = str(refused.value) - assert ( - '"LiteLLM_SpendLogs".api_key, startTime: 2 row(s) with an index entry over 2704 bytes, largest 3232 bytes ' - "(api_key up to 3207 bytes)" - ) in refusal, refusal - assert f'blocks: CREATE INDEX "{API_KEY_INDEX}"' in refusal, refusal - assert _ledger(scratch_database) == ledger_before, "a refused job must not apply a single migration" + assert ProxyExtrasDBManager.run_migration_job(use_migrate=True, use_v2_resolver=use_v2_resolver) is False + assert len(_ledger(scratch_database)) > len(ledger_before), "the rows must not hold back a single migration" assert API_KEY_INDEX not in _index_validity(scratch_database, "api_key_startTime_idx") with psycopg.connect(scratch_database, autocommit=True) as conn: conn.execute(_HASH_LEGACY_KEYS) assert ProxyExtrasDBManager.run_migration_job(use_migrate=True, use_v2_resolver=use_v2_resolver) is True assert _index_validity(scratch_database, "api_key_startTime_idx") == {API_KEY_INDEX: True} - assert len(_ledger(scratch_database)) > len(ledger_before) @requires_db @@ -1017,49 +1009,6 @@ def test_a_serving_proxy_index_build_that_meets_too_long_keys_reports_the_rows_i assert API_KEY_INDEX not in _index_validity(scratch_database, "api_key_startTime_idx") -_PROBE_STATEMENTS: Final = ( - 'CREATE UNIQUE INDEX "probe_b_key" ON "preflight_probe" ("b")', - 'ALTER TABLE "preflight_probe" ALTER COLUMN "b" SET NOT NULL', - 'ALTER TABLE "preflight_probe" ADD COLUMN "c" INTEGER NOT NULL', - 'ALTER TABLE "preflight_probe" ADD CONSTRAINT "probe_r_fkey" FOREIGN KEY ("r") REFERENCES "preflight_parent"("id")', - 'ALTER TABLE "preflight_probe" ADD CONSTRAINT "probe_a_check" CHECK ("a" <> \'bad\')', - 'CREATE INDEX "probe_a_idx" ON "preflight_probe" ("a")', -) - - -@requires_db -def test_each_check_kind_counts_exactly_the_rows_postgres_would_reject(scratch_database: str) -> None: - with psycopg.connect(scratch_database, autocommit=True) as conn: - conn.execute('CREATE TABLE "preflight_parent" ("id" INTEGER PRIMARY KEY)') - conn.execute('CREATE TABLE "preflight_probe" ("a" TEXT, "b" INTEGER, "r" INTEGER)') - conn.execute('INSERT INTO "preflight_parent" VALUES (1)') - conn.execute( - 'INSERT INTO "preflight_probe" VALUES (%s, 1, 1), (%s, 1, 1), (%s, NULL, 2), (%s, 2, NULL), (%s, 3, 1)', - ("ok", "ok", "bad", "bad", secrets.token_hex(1400)), - ) - checks: Final = tuple( - check - for statement in _PROBE_STATEMENTS - for check in migration_data_checks.checks_for_statement(statement, "probe") - ) - findings: Final = migration_data_checks.find_blockers(conn, "public", checks) - for statement in _PROBE_STATEMENTS: - with conn.transaction(force_rollback=True), pytest.raises(psycopg.Error): - conn.execute(statement) # pyright: ignore[reportArgumentType] # test SQL - - counted: Final = { - (f.check.kind, f.check.columns): f.rows for f in findings if isinstance(f, migration_data_checks.BlockedChange) - } - assert counted == { - ("duplicates", ("b",)): 1, - ("nulls", ("b",)): 1, - ("any_row", ("c",)): 5, - ("orphans", ("r",)): 1, - ("check_violations", ()): 2, - ("too_long", ("a",)): 1, - } - - @requires_db @pytest.mark.parametrize( ("index_columns", "s_value", "t_value", "longest_fitting_key", "described_as"), @@ -1098,7 +1047,7 @@ def test_the_entry_size_check_draws_the_line_exactly_where_postgres_does( ), ) findings: Final = migration_data_checks.find_blockers( - conn, "public", migration_data_checks.checks_for_statement(create_index.as_string(conn), "probe") + conn, "public", migration_data_checks.checks_for_index("entry_probe", "entry_probe_idx", index_columns) ) with pytest.raises(psycopg.errors.ProgramLimitExceeded, match="exceeds btree version 4 maximum 2704"): conn.execute(create_index) @@ -1122,7 +1071,7 @@ def test_a_value_postgres_compresses_under_the_limit_is_not_counted_while_random conn.execute('CREATE TABLE "compress_probe" ("k" TEXT)') conn.execute('INSERT INTO "compress_probe" ("k") VALUES (%s), (%s)', ("a" * 3000, secrets.token_hex(1500))) findings: Final = migration_data_checks.find_blockers( - conn, "public", migration_data_checks.checks_for_statement(create_index.as_string(conn), "probe") + conn, "public", migration_data_checks.checks_for_index("compress_probe", "compress_probe_idx", ("k",)) ) with pytest.raises(psycopg.errors.ProgramLimitExceeded, match="index row size 3016 exceeds"): conn.execute(create_index) @@ -1154,7 +1103,6 @@ def test_the_check_flags_exactly_the_compressed_keys_postgres_rejects_at_the_ent check and through Postgres one by one. The check must flag a key exactly when Postgres rejects it.""" noise: Final = "".join(hashlib.sha256(str(i).encode()).hexdigest() for i in range(40)) keys: Final = tuple(noise[:n] + "a" * 1500 for n in range(2400, 2500)) - create_index: Final = sql.SQL('CREATE INDEX "stored_probe_idx" ON "stored_probe" ("k")') with psycopg.connect(scratch_database, autocommit=True) as conn: conn.execute('CREATE TABLE "stored_probe" ("k" TEXT)') conn.execute('INSERT INTO "stored_probe" ("k") SELECT unnest(%s::text[])', (list(keys),)) @@ -1162,7 +1110,7 @@ def test_the_check_flags_exactly_the_compressed_keys_postgres_rejects_at_the_ent 'SELECT count(*) FROM "stored_probe" WHERE pg_column_size("k") < octet_length("k") AND octet_length("k") > 2704' ).fetchone() findings: Final = migration_data_checks.find_blockers( - conn, "public", migration_data_checks.checks_for_statement(create_index.as_string(conn), "probe") + conn, "public", migration_data_checks.checks_for_index("stored_probe", "stored_probe_idx", ("k",)) ) rejected: Final = tuple(key for key in keys if _postgres_rejects_an_index_on(conn, key)) @@ -1176,14 +1124,176 @@ def test_the_check_flags_exactly_the_compressed_keys_postgres_rejects_at_the_ent @requires_db def test_a_scan_that_cannot_finish_in_time_is_reported_as_unverified_rather_than_passed(scratch_database: str) -> None: with psycopg.connect(scratch_database, autocommit=True) as conn: - conn.execute('CREATE TABLE "preflight_slow" ("a" INTEGER)') - conn.execute('INSERT INTO "preflight_slow" SELECT generate_series(1, 50)') - (check,) = migration_data_checks.checks_for_statement( - 'ALTER TABLE "preflight_slow" ADD CONSTRAINT "slow" CHECK (pg_sleep(0.1) IS NULL)', "probe" + conn.execute('CREATE TABLE "preflight_slow" ("k" TEXT)') + conn.execute( + "INSERT INTO \"preflight_slow\" SELECT md5(i::text) || repeat('a', 1500) FROM generate_series(1, 200000) AS i" ) - findings: Final = migration_data_checks.find_blockers(conn, "public", (check,), timeout_seconds=0.5) + (check,) = migration_data_checks.checks_for_index("preflight_slow", "preflight_slow_idx", ("k",)) + findings: Final = migration_data_checks.find_blockers(conn, "public", (check,), timeout_seconds=0.01) timeout_afterwards: Final = conn.execute("SHOW statement_timeout").fetchone() - assert findings == (migration_data_checks.CouldNotVerify(check, 0.5),) - assert timeout_afterwards == ("0",), "the preflight must hand the session back with the timeout it found" - assert "could not be checked within 0.5 s" in migration_data_checks.report(findings) + assert findings == (migration_data_checks.CouldNotVerify(check, 0.01),) + assert timeout_afterwards == ("0",), "the check must hand the session back with the timeout it found" + assert "could not be checked within 0.01 s" in migration_data_checks.report(findings) + + +CURRENT_RELEASE: Final = "99999999999999" +REQUIRE_TEAM_ALIAS: Final = "20991231000000_require_team_alias" + + +def _extras_with_scratch_migration(tmp_path: Path, name: str, migration_sql: str) -> Path: + """A copy of the shipped package dir with one more migration folder, the shape `LITELLM_MIGRATION_DIR` points at.""" + extras: Final = tmp_path / "extras" + shutil.copytree(PACKAGE / "migrations", extras / "migrations") + shutil.copy(PACKAGE / "schema.prisma", extras / "schema.prisma") + (extras / "migrations" / name).mkdir() + (extras / "migrations" / name / "migration.sql").write_text(migration_sql) + return extras + + +def _team_alias_is_nullable(database_url: str) -> bool: + with psycopg.connect(database_url) as conn: + row: Final = conn.execute( + "SELECT is_nullable FROM information_schema.columns " + "WHERE table_name = 'LiteLLM_TeamTable' AND column_name = 'team_alias'" + ).fetchone() + return row == ("YES",) + + +@requires_db +@pytest.mark.parametrize("release", (API_KEY_INDEX_MIGRATION,), ids=("v1.102.1",)) +def test_the_dry_run_applies_every_pending_migration_in_order_and_rolls_all_of_it_back(scratch_database: str) -> None: + ledger_before: Final = _ledger(scratch_database) + with psycopg.connect(scratch_database, autocommit=True) as conn: + pending: Final = migration_data_checks.pending_migrations(conn, "public", PACKAGE / "migrations") + result: Final = migration_dry_run.dry_run(conn, "public", PACKAGE / "migrations") + + assert result.rejected is None, result + assert tuple(applied.name for applied in result.applied) == tuple(name for name, _ in pending) + assert len(result.applied) > 1 + assert _ledger(scratch_database) == ledger_before, "a dry run must leave the ledger alone" + assert _index_validity(scratch_database, "api_key_startTime_idx") == {}, "a dry run must leave no index behind" + names: Final = tuple(name for name, _ in pending) + assert all(set(applied.replayed) <= set(names[: names.index(applied.name)]) for applied in result.applied), ( + "a migration is only ever replayed behind a later one that depends on it" + ) + assert "applied one at a time in a dry run" in migration_dry_run.report(result) + + +@requires_db +@pytest.mark.parametrize("release", (CURRENT_RELEASE,), ids=("current",)) +def test_the_dry_run_reports_postgres_own_error_for_the_row_a_migration_cannot_take( + scratch_database: str, tmp_path: Path +) -> None: + extras: Final = _extras_with_scratch_migration( + tmp_path, REQUIRE_TEAM_ALIAS, 'ALTER TABLE "LiteLLM_TeamTable" ALTER COLUMN "team_alias" SET NOT NULL;\n' + ) + ledger_before: Final = _ledger(scratch_database) + with psycopg.connect(scratch_database, autocommit=True) as conn: + conn.execute('INSERT INTO "LiteLLM_TeamTable" ("team_id") VALUES (%s)', ("no-alias",)) + result: Final = migration_dry_run.dry_run(conn, "public", extras / "migrations") + + assert result.rejected is not None, result + assert result.rejected.name == REQUIRE_TEAM_ALIAS + assert 'column "team_alias" of relation "LiteLLM_TeamTable" contains null values' in result.rejected.error + assert result.applied == () + assert _team_alias_is_nullable(scratch_database), "a rejected dry run must leave the schema as it was" + assert _ledger(scratch_database) == ledger_before + + +@requires_db +@pytest.mark.parametrize("release", (CURRENT_RELEASE,), ids=("current",)) +def test_the_cli_summary_turns_from_postgres_error_to_timings_once_the_rows_are_fixed( + scratch_database: str, tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + extras: Final = _extras_with_scratch_migration( + tmp_path, REQUIRE_TEAM_ALIAS, 'ALTER TABLE "LiteLLM_TeamTable" ALTER COLUMN "team_alias" SET NOT NULL;\n' + ) + monkeypatch.setenv("LITELLM_MIGRATION_DIR", str(extras)) + with psycopg.connect(scratch_database, autocommit=True) as conn: + conn.execute('INSERT INTO "LiteLLM_TeamTable" ("team_id") VALUES (%s)', ("no-alias",)) + + ok, summary = ProxyExtrasDBManager.migration_dry_run_summary(use_migrate=True) + assert ok is False + assert summary.startswith(f"Migration {REQUIRE_TEAM_ALIAS} fails on this database after "), summary + assert "contains null values" in summary + + with psycopg.connect(scratch_database, autocommit=True) as conn: + conn.execute('UPDATE "LiteLLM_TeamTable" SET "team_alias" = \'filled\' WHERE "team_alias" IS NULL') + ok_after, summary_after = ProxyExtrasDBManager.migration_dry_run_summary(use_migrate=True) + assert ok_after is True + assert summary_after.startswith("1 pending migration(s) applied one at a time in a dry run"), summary_after + assert f"\n {REQUIRE_TEAM_ALIAS}: " in summary_after + assert _team_alias_is_nullable(scratch_database) + + +@requires_db +@pytest.mark.parametrize("release", (CURRENT_RELEASE,), ids=("current",)) +def test_a_migration_that_depends_on_the_one_before_it_is_replayed_with_it_while_independent_ones_run_alone( + scratch_database: str, tmp_path: Path +) -> None: + extras: Final = _extras_with_scratch_migration( + tmp_path, "20991231000000_add_team_region", 'ALTER TABLE "LiteLLM_TeamTable" ADD COLUMN "region" TEXT;\n' + ) + (extras / "migrations" / "20991231000001_index_team_region").mkdir() + (extras / "migrations" / "20991231000001_index_team_region" / "migration.sql").write_text( + 'CREATE INDEX CONCURRENTLY "LiteLLM_TeamTable_region_idx" ON "LiteLLM_TeamTable"("region");\n' + ) + (extras / "migrations" / "20991231000002_index_team_alias").mkdir() + (extras / "migrations" / "20991231000002_index_team_alias" / "migration.sql").write_text( + 'CREATE INDEX "LiteLLM_TeamTable_alias_idx" ON "LiteLLM_TeamTable"("team_alias");\n' + ) + with psycopg.connect(scratch_database, autocommit=True) as conn: + result: Final = migration_dry_run.dry_run(conn, "public", extras / "migrations") + region: Final = conn.execute( + "SELECT 1 FROM information_schema.columns WHERE table_name = 'LiteLLM_TeamTable' AND column_name = 'region'" + ).fetchone() + + assert result.rejected is None, result + assert tuple((applied.name, applied.replayed) for applied in result.applied) == ( + ("20991231000000_add_team_region", ()), + ("20991231000001_index_team_region", ("20991231000000_add_team_region",)), + ("20991231000002_index_team_alias", ()), + ) + assert region is None, "every transaction must have been rolled back" + with psycopg.connect(scratch_database) as conn: + probe_indexes: Final = conn.execute( + "SELECT indexname FROM pg_indexes " + "WHERE indexname IN ('LiteLLM_TeamTable_region_idx', 'LiteLLM_TeamTable_alias_idx')" + ).fetchall() + assert probe_indexes == [], probe_indexes + + +@requires_db +@pytest.mark.parametrize("release", (CURRENT_RELEASE,), ids=("current",)) +def test_a_database_with_tables_but_no_ledger_is_judged_by_the_schema_diff_instead_of_replaying_history( + scratch_database: str, tmp_path: Path +) -> None: + with psycopg.connect(scratch_database, autocommit=True) as conn: + conn.execute("DROP TABLE _prisma_migrations") + matching: Final = migration_dry_run.dry_run(conn, "public", PACKAGE / "migrations", baseline_check=lambda: None) + differing: Final = migration_dry_run.dry_run( + conn, + "public", + PACKAGE / "migrations", + baseline_check=lambda: ( + "the schema differs from this build. prisma migrate diff:\nALTER TABLE x ADD COLUMN y TEXT;" + ), + ) + custom: Final = migration_dry_run.dry_run( + conn, + "public", + PACKAGE / "migrations", + baseline_check=lambda: ProxyExtrasDBManager._baseline_refusal( + _extras_with_scratch_migration( + tmp_path, "20991231000000_custom", 'ALTER TABLE "LiteLLM_TeamTable" ADD COLUMN "c" TEXT;\n' + ) + ), + ) + + assert matching.ok is True and matching.applied == () + assert "nothing to dry-run" in migration_dry_run.report(matching) + assert differing.ok is False + assert "ALTER TABLE x ADD COLUMN y TEXT;" in migration_dry_run.report(differing) + assert custom.ok is False + assert "custom migration history" in migration_dry_run.report(custom) diff --git a/tests/unit/litellm_proxy_extras/test_migration_data_checks.py b/tests/unit/litellm_proxy_extras/test_migration_data_checks.py index d4490927e29..37377498d89 100644 --- a/tests/unit/litellm_proxy_extras/test_migration_data_checks.py +++ b/tests/unit/litellm_proxy_extras/test_migration_data_checks.py @@ -12,98 +12,15 @@ from litellm_proxy_extras.migration_data_checks import ( _count_query, btree_max_entry_bytes, checks_for_index, - checks_for_migration, - checks_for_statement, migration_data_check_timeout_seconds, report, - statements, ) -MIGRATIONS: Final = Path(__file__).resolve().parents[3] / "litellm-proxy-extras" / "litellm_proxy_extras" / "migrations" -API_KEY_INDEX_AS_SHIPPED_IN_V1_103_0: Final = 'CREATE INDEX IF NOT EXISTS "LiteLLM_SpendLogs_api_key_startTime_idx" ON "LiteLLM_SpendLogs"("api_key", "startTime")' - def _shape(checks: tuple[Check, ...]) -> tuple[tuple[str, str, tuple[str, ...]], ...]: return tuple((check.kind, check.table, check.columns) for check in checks) -@pytest.mark.parametrize( - ("statement", "expected"), - ( - ( - API_KEY_INDEX_AS_SHIPPED_IN_V1_103_0, - (("too_long", "LiteLLM_SpendLogs", ("api_key", "startTime")),), - ), - ( - 'CREATE UNIQUE INDEX CONCURRENTLY "T_a_b_key" ON "public"."T" USING btree ("a", "b" DESC)', - (("too_long", "T", ("a", "b")), ("duplicates", "T", ("a", "b"))), - ), - ('CREATE INDEX "T_a_idx" ON "T" USING gin ("a")', ()), - ('CREATE INDEX "T_a_idx" ON "T" (lower("a"))', ()), - ('ALTER TABLE "T" ALTER COLUMN "a" SET NOT NULL', (("nulls", "T", ("a",)),)), - ('ALTER TABLE "T" ADD COLUMN "a" TEXT NOT NULL', (("any_row", "T", ("a",)),)), - ('ALTER TABLE "T" ADD COLUMN "a" TEXT NOT NULL DEFAULT \'x\'', ()), - ('ALTER TABLE "T" ADD COLUMN "a" TEXT', ()), - ( - 'ALTER TABLE "T" ADD CONSTRAINT "T_r_fkey" FOREIGN KEY ("r") REFERENCES "R"("id") ON DELETE CASCADE', - (("orphans", "T", ("r",)),), - ), - ('ALTER TABLE "T" ADD CONSTRAINT "T_r_fkey" FOREIGN KEY ("r") REFERENCES "R"("id") NOT VALID', ()), - ('ALTER TABLE "T" ADD CONSTRAINT "T_a_key" UNIQUE ("a", "b")', (("duplicates", "T", ("a", "b")),)), - ( - 'ALTER TABLE "T" ADD CONSTRAINT "T_a_check" CHECK ("a" > 0 AND "b" IS NOT NULL)', - (("check_violations", "T", ()),), - ), - ('ALTER TABLE "T" ADD CONSTRAINT "T_a_check" CHECK ("a" > 0) NOT VALID', ()), - ( - 'ALTER TABLE "T" ADD COLUMN "c" TEXT, ALTER COLUMN "b" SET NOT NULL, ADD CONSTRAINT "u" UNIQUE ("c")', - (("nulls", "T", ("b",)), ("duplicates", "T", ("c",))), - ), - ('CREATE TABLE "T" ("id" TEXT NOT NULL, CONSTRAINT "T_pkey" PRIMARY KEY ("id"))', ()), - ('DROP INDEX IF EXISTS "T_a_idx"', ()), - ("SELECT 1", ()), - ), -) -def test_a_statement_gets_one_check_per_way_existing_rows_can_make_it_fail( - statement: str, expected: tuple[tuple[str, str, tuple[str, ...]], ...] -) -> None: - assert _shape(checks_for_statement(statement, "migration x")) == expected - - -def test_a_foreign_key_check_keeps_the_referenced_table_and_columns() -> None: - (check,) = checks_for_statement( - 'ALTER TABLE "T" ADD CONSTRAINT "fk" FOREIGN KEY ("r", "s") REFERENCES "R"("id", "sub")', "migration x" - ) - assert (check.referenced_table, check.referenced_columns) == ("R", ("id", "sub")) - - -def test_a_check_constraint_keeps_its_expression_for_the_count_query() -> None: - (check,) = checks_for_statement('ALTER TABLE "T" ADD CONSTRAINT "c" CHECK ("a" > 0)', "migration x") - assert check.expression == '"a" > 0' - - -def test_comments_and_blank_statements_are_dropped_before_matching() -> None: - migration_sql: Final = "-- CreateIndex\n" + API_KEY_INDEX_AS_SHIPPED_IN_V1_103_0 + ";\n\n-- done; really\n" - assert statements(migration_sql) == (API_KEY_INDEX_AS_SHIPPED_IN_V1_103_0,) - assert _shape(checks_for_migration("x", migration_sql)) == _shape( - checks_for_statement(API_KEY_INDEX_AS_SHIPPED_IN_V1_103_0, "x") - ) - - -def test_every_shipped_migration_is_readable_and_the_later_indexes_on_text_columns_get_a_check() -> None: - derived: Final = { - folder.name: checks_for_migration(folder.name, (folder / "migration.sql").read_text()) - for folder in MIGRATIONS.iterdir() - if folder.is_dir() - } - assert len(derived) > 200 - everywhere: Final = tuple(check for checks in derived.values() for check in checks) # comprehension-ok: flattening - assert any( - check.kind == "too_long" and check.table == "LiteLLM_SpendLogs" and "session_id" in check.columns - for check in everywhere - ) - - def test_a_request_log_index_gets_one_entry_size_check_over_all_its_columns() -> None: assert _shape( checks_for_index("LiteLLM_SpendLogs", "LiteLLM_SpendLogs_api_key_startTime_idx", ("api_key", "startTime")) @@ -142,14 +59,19 @@ def test_the_report_names_the_rows_the_statement_and_the_fix_and_says_nothing_wa assert 'blocks: CREATE INDEX "LiteLLM_SpendLogs_api_key_startTime_idx" (request-log index)' in text assert 'fix: Shorten or hash the long "LiteLLM_SpendLogs" (api_key, startTime) values' in text assert text.startswith("Existing rows would make a pending schema change fail:") - assert "No schema change was made" in text + assert "No index was built" in text def test_a_scan_that_timed_out_is_reported_as_unverified_with_the_knob_to_turn() -> None: - (check,) = checks_for_statement('ALTER TABLE "T" ALTER COLUMN "a" SET NOT NULL', "migration 20270101000000_x") + (check,) = checks_for_index( + "LiteLLM_SpendLogs", "LiteLLM_SpendLogs_api_key_startTime_idx", ("api_key", "startTime") + ) text: Final = report((CouldNotVerify(check, 600.0),)) - assert '"T".a: could not be checked within 600 s, refusing rather than guessing' in text - assert 'blocks: ALTER TABLE "T" ALTER COLUMN "a" SET NOT NULL (migration 20270101000000_x)' in text + assert ( + '"LiteLLM_SpendLogs".api_key, startTime: could not be checked within 600 s, refusing rather than guessing' + in text + ) + assert 'blocks: CREATE INDEX "LiteLLM_SpendLogs_api_key_startTime_idx" (request-log index)' in text assert MIGRATION_DATA_CHECK_TIMEOUT_ENV_VAR in text diff --git a/tests/unit/litellm_proxy_extras/test_migration_dry_run.py b/tests/unit/litellm_proxy_extras/test_migration_dry_run.py new file mode 100644 index 00000000000..75a7fea6c89 --- /dev/null +++ b/tests/unit/litellm_proxy_extras/test_migration_dry_run.py @@ -0,0 +1,80 @@ +from typing import Final + +import pytest +from litellm_proxy_extras.migration_dry_run import Applied, Baseline, DryRun, Rejected, report, transactional + +INDEX: Final = 'CREATE INDEX IF NOT EXISTS "a_idx" ON "T"("c")' + + +@pytest.mark.parametrize( + ("migration_sql", "expected"), + ( + ('CREATE INDEX CONCURRENTLY IF NOT EXISTS "a_idx" ON "T"("c")', INDEX), + ('create unique index concurrently "u" on "T"("d")', 'create unique index "u" on "T"("d")'), + ('DROP INDEX CONCURRENTLY "x"', 'DROP INDEX "x"'), + ('REINDEX INDEX CONCURRENTLY "y"', 'REINDEX INDEX "y"'), + (INDEX, INDEX), + ('ALTER TABLE "T" ALTER COLUMN "c" SET NOT NULL', 'ALTER TABLE "T" ALTER COLUMN "c" SET NOT NULL'), + ("-- built CONCURRENTLY by the index job\n" + INDEX, "-- built CONCURRENTLY by the index job\n" + INDEX), + ), + ids=("concurrently", "unique_lowercase", "drop", "reindex", "plain_index", "alter", "word_in_comment"), +) +def test_only_concurrent_index_statements_change_so_the_migration_can_run_in_a_transaction( + migration_sql: str, expected: str +) -> None: + assert transactional(migration_sql) == expected + + +def test_a_clean_dry_run_lists_each_migration_with_the_time_the_upgrade_will_hold_its_locks() -> None: + text: Final = report(DryRun((Applied("20270101000000_a", 0.04), Applied("20270102000000_b", 12.345)), None)) + assert text.startswith( + "2 pending migration(s) applied one at a time in a dry run, each rolled back before the next" + ) + assert "\n 20270101000000_a: 0.0 s\n 20270102000000_b: 12.3 s" in text + + +def test_no_pending_migration_is_said_plainly() -> None: + assert report(DryRun((), None)) == "No pending migrations, this database is on this version's schema." + + +def test_a_rejected_migration_is_reported_with_postgres_own_words_and_what_ran_before_it() -> None: + error: Final = 'column "team_alias" of relation "LiteLLM_TeamTable" contains null values' + text: Final = report(DryRun((Applied("20270101000000_a", 0.5),), Rejected("20270102000000_b", error, 3.25))) + assert text.startswith( + "Migration 20270102000000_b fails on this database after 3.2 s, nothing was applied.\n" + error + ) + assert text.endswith("Applied before it in the dry run:\n 20270101000000_a: 0.5 s") + + +def test_a_rejection_on_the_first_migration_lists_nothing_before_it() -> None: + text: Final = report(DryRun((), Rejected("20270101000000_a", 'syntax error at or near "NOPE"', 0.0))) + assert "Applied before it" not in text + assert text.endswith('syntax error at or near "NOPE"') + + +def test_a_migration_that_needed_the_one_before_it_says_which_were_replayed_in_front() -> None: + text: Final = report( + DryRun((Applied("20270101000000_a", 0.1), Applied("20270102000000_b", 0.4, ("20270101000000_a",))), None) + ) + assert text.endswith( + "\n 20270101000000_a: 0.1 s\n 20270102000000_b: 0.4 s (depends on 20270101000000_a, replayed in front of it)" + ) + + +def test_a_database_without_history_that_matches_this_build_has_nothing_to_dry_run() -> None: + result: Final = DryRun((), None, Baseline(None)) + assert result.ok is True + assert report(result) == ( + "No migration history and the schema matches this build, the proxy baselines it at boot, nothing to dry-run." + ) + + +def test_a_database_without_history_the_boot_refuses_to_baseline_reports_the_refusal() -> None: + result: Final = DryRun( + (), None, Baseline('the schema differs from this build. prisma migrate diff:\n-- AlterTable\nALTER TABLE "x"') + ) + assert result.ok is False + assert report(result) == ( + "No migration history and the proxy cannot baseline it at boot: " + 'the schema differs from this build. prisma migrate diff:\n-- AlterTable\nALTER TABLE "x"' + ) diff --git a/tests/unit/proxy/db/test_prisma_client.py b/tests/unit/proxy/db/test_prisma_client.py index 4b776bfbb1a..6c82628576f 100644 --- a/tests/unit/proxy/db/test_prisma_client.py +++ b/tests/unit/proxy/db/test_prisma_client.py @@ -586,3 +586,19 @@ async def test_a_query_issued_while_the_last_one_finishes_keeps_the_engine_drain await second await asyncio.wait_for(drained, timeout=1) assert engine.queries_finished == 2 + + +def test_migration_dry_run_summary_without_migrations_has_nothing_to_dry_run(unset_database_url): + """`prisma db push` applies no migration files, so the dry run reports that instead of connecting.""" + from litellm.proxy.db.prisma_client import PrismaManager + + assert PrismaManager.migration_dry_run_summary(use_migrate=False) == ( + True, + "The schema is pushed rather than migrated, there is nothing to dry-run.", + ) + + +def test_migration_dry_run_summary_needs_database_url(unset_database_url): + from litellm.proxy.db.prisma_client import PrismaManager + + assert PrismaManager.migration_dry_run_summary(use_migrate=True) == (False, "DATABASE_URL is not set.") diff --git a/tests/unit/proxy/test_proxy_cli.py b/tests/unit/proxy/test_proxy_cli.py index afefae0dbb1..1db334039ec 100644 --- a/tests/unit/proxy/test_proxy_cli.py +++ b/tests/unit/proxy/test_proxy_cli.py @@ -4393,3 +4393,20 @@ def test_granian_flag_notices_are_json_records_in_json_mode( "LiteLLM: --ciphers is not applied when using --run_granian.", ] assert [record["level"] for record in records] == ["INFO", "WARNING", "WARNING"] + + +def test_migration_data_check_exits_zero_when_there_is_nothing_to_dry_run(monkeypatch, capfd): + monkeypatch.setenv("DATABASE_URL", "postgresql://user:pass@db.example/litellm") + + assert ProxyInitializationHelpers.run_migration_data_check(use_migrate=False) == 0 + + out, err = capfd.readouterr() + assert out == "LiteLLM Proxy: The schema is pushed rather than migrated, there is nothing to dry-run.\n" + assert err == "" + + +def test_migration_data_check_requires_database_url(monkeypatch): + monkeypatch.delenv("DATABASE_URL", raising=False) + + with pytest.raises(click.UsageError, match="requires DATABASE_URL"): + ProxyInitializationHelpers.run_migration_data_check(use_migrate=True) From feda51453a3c468d04b608da4a9a972d3b365ce8 Mon Sep 17 00:00:00 2001 From: mrinal-berri Date: Thu, 29 Oct 2026 09:39:49 -0700 Subject: [PATCH 13/13] fix!: retain unattributed daily activity across reports and exports (#48925) * fix: retain unattributed daily activity breakdowns * fix: scope daily activity bucket drilldowns consistently * fix: preserve all additive daily activity metrics * test: cover unattributed reporting end to end * fix: preserve daily activity bucket identities across exports and key details * test: drop SQL-text assertions changed by daily activity bucket fix Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../auto_router_endpoints.py | 2 +- .../common_daily_activity.py | 168 ++++++++++------- .../daily_activity_routes.py | 34 +--- .../daily_activity_scopes.py | 9 +- .../tag_management_endpoints.py | 8 +- .../management_endpoints/team_endpoints.py | 8 +- .../repositories/daily_activity_buckets.py | 39 ++++ .../repositories/daily_activity_repository.py | 9 +- litellm/repositories/daily_activity_sql.py | 60 +++--- .../routing/test_auto_router_management.py | 24 ++- ...st_daily_activity_aggregated_breakdowns.py | 178 ++++++++++++++++++ .../spend/test_daily_activity_routes.py | 14 ++ .../test_common_daily_activity.py | 143 ++++++++++++++ .../test_daily_activity_routes.py | 22 ++- .../test_tag_management_endpoints.py | 5 +- .../repositories/test_daily_activity_sql.py | 17 +- .../EntityUsage/TopKeyView.test.tsx | 26 ++- .../components/EntityUsage/TopKeyView.tsx | 6 +- .../UsagePage/keyActivityData.test.ts | 52 ++++- .../components/UsagePage/keyActivityData.ts | 10 + .../KeySavingsTab.integration.test.tsx | 12 ++ .../components/templates/KeySavingsTab.tsx | 3 +- 22 files changed, 695 insertions(+), 154 deletions(-) create mode 100644 litellm/repositories/daily_activity_buckets.py diff --git a/litellm/proxy/management_endpoints/auto_router_endpoints.py b/litellm/proxy/management_endpoints/auto_router_endpoints.py index 3dceab6b7f7..b79154bf509 100644 --- a/litellm/proxy/management_endpoints/auto_router_endpoints.py +++ b/litellm/proxy/management_endpoints/auto_router_endpoints.py @@ -865,7 +865,7 @@ async def _recorded_autorouter_savings( entity_id_field="user_id", entity_id=user_id, exclude_entity_ids=None, - api_key=api_key, + api_key=(api_key,) if api_key is not None else None, start_date=start_day, end_date=end_day, model=None, diff --git a/litellm/proxy/management_endpoints/common_daily_activity.py b/litellm/proxy/management_endpoints/common_daily_activity.py index 66f386b792d..98372b64015 100644 --- a/litellm/proxy/management_endpoints/common_daily_activity.py +++ b/litellm/proxy/management_endpoints/common_daily_activity.py @@ -23,6 +23,12 @@ from litellm.proxy.spend_tracking.key_metadata_recovery import ( ) from litellm.proxy.spend_tracking.ptu_feature_flag import is_ptu_cost_attribution_enabled from litellm.proxy.utils import PrismaClient +from litellm.repositories.daily_activity_buckets import ( + api_key_bucket, + model_bucket, + model_group_bucket, + raw_api_key_filter, +) from litellm.repositories.daily_activity_repository import DailyActivityRepository from litellm.types.proxy.management_endpoints.common_daily_activity import ( BreakdownMetrics, @@ -182,7 +188,9 @@ class _AggregatedSpendData(TypedDict): totals: ReadOnly[SpendMetrics] -def _key_metadata(api_key_metadata: Mapping[str, KeyMetadataRow], api_key: str) -> KeyMetadata: +def _key_metadata(api_key_metadata: Mapping[str, KeyMetadataRow], api_key: str | None) -> KeyMetadata: + if not api_key: + return KeyMetadata(key_alias="Unattributed", key_exists=False) meta: Final = api_key_metadata.get(api_key) return KeyMetadata( key_alias=meta.key_alias if meta is not None else None, @@ -302,13 +310,14 @@ def update_breakdown_metrics( parent bucket but never appear as an api_key row, and are kept out of the per-request provider breakdown.""" + api_key: Final = api_key_bucket(record.api_key) is_ptu_sentinel: Final = record.api_key == PTU_SENTINEL_API_KEY # A PTU sentinel row keys on the deployment id so a rename cannot move it, and carries # the operator-facing name in model_group. The breakdown key is rendered directly as a # label, so display the name; two deployments sharing one name merge here, which is # what the write path used to do by collapsing them into a single row. - model_key: Final = (record.model_group or record.model) if is_ptu_sentinel else record.model + model_key: Final = (record.model_group or record.model) if is_ptu_sentinel else model_bucket(record.model) # Update model breakdown if model_key and model_key not in breakdown.models: @@ -321,18 +330,18 @@ def update_breakdown_metrics( if not is_ptu_sentinel: # Update API key breakdown for this model - if record.api_key not in breakdown.models[model_key].api_key_breakdown: - breakdown.models[model_key].api_key_breakdown[record.api_key] = KeyMetricWithMetadata( + if api_key not in breakdown.models[model_key].api_key_breakdown: + breakdown.models[model_key].api_key_breakdown[api_key] = KeyMetricWithMetadata( metrics=SpendMetrics(), metadata=_key_metadata(api_key_metadata, record.api_key), ) - breakdown.models[model_key].api_key_breakdown[record.api_key].metrics = update_metrics( - breakdown.models[model_key].api_key_breakdown[record.api_key].metrics, + breakdown.models[model_key].api_key_breakdown[api_key].metrics = update_metrics( + breakdown.models[model_key].api_key_breakdown[api_key].metrics, record, ) # Update model group breakdown - model_group_key: Final = record.model_group or record.model + model_group_key: Final = model_group_bucket(record.model_group, record.model) if model_group_key and model_group_key not in breakdown.model_groups: breakdown.model_groups[model_group_key] = MetricWithMetadata( metrics=SpendMetrics(), @@ -345,13 +354,13 @@ def update_breakdown_metrics( if not is_ptu_sentinel: # Update API key breakdown for this model - if record.api_key not in breakdown.model_groups[model_group_key].api_key_breakdown: - breakdown.model_groups[model_group_key].api_key_breakdown[record.api_key] = KeyMetricWithMetadata( + if api_key not in breakdown.model_groups[model_group_key].api_key_breakdown: + breakdown.model_groups[model_group_key].api_key_breakdown[api_key] = KeyMetricWithMetadata( metrics=SpendMetrics(), metadata=_key_metadata(api_key_metadata, record.api_key), ) - breakdown.model_groups[model_group_key].api_key_breakdown[record.api_key].metrics = update_metrics( - breakdown.model_groups[model_group_key].api_key_breakdown[record.api_key].metrics, + breakdown.model_groups[model_group_key].api_key_breakdown[api_key].metrics = update_metrics( + breakdown.model_groups[model_group_key].api_key_breakdown[api_key].metrics, record, ) @@ -365,21 +374,20 @@ def update_breakdown_metrics( breakdown.mcp_servers[record.mcp_namespaced_tool_name].metrics, record ) - # Update API key breakdown for this MCP server - if record.api_key not in breakdown.mcp_servers[record.mcp_namespaced_tool_name].api_key_breakdown: - breakdown.mcp_servers[record.mcp_namespaced_tool_name].api_key_breakdown[record.api_key] = ( - KeyMetricWithMetadata( - metrics=SpendMetrics(), - metadata=_key_metadata(api_key_metadata, record.api_key), + if not is_ptu_sentinel: + # Update API key breakdown for this MCP server + if api_key not in breakdown.mcp_servers[record.mcp_namespaced_tool_name].api_key_breakdown: + breakdown.mcp_servers[record.mcp_namespaced_tool_name].api_key_breakdown[api_key] = ( + KeyMetricWithMetadata( + metrics=SpendMetrics(), + metadata=_key_metadata(api_key_metadata, record.api_key), + ) ) - ) - breakdown.mcp_servers[record.mcp_namespaced_tool_name].api_key_breakdown[ - record.api_key - ].metrics = update_metrics( - breakdown.mcp_servers[record.mcp_namespaced_tool_name].api_key_breakdown[record.api_key].metrics, - record, - ) + breakdown.mcp_servers[record.mcp_namespaced_tool_name].api_key_breakdown[api_key].metrics = update_metrics( + breakdown.mcp_servers[record.mcp_namespaced_tool_name].api_key_breakdown[api_key].metrics, + record, + ) if not is_ptu_sentinel: # Update provider breakdown @@ -392,13 +400,13 @@ def update_breakdown_metrics( breakdown.providers[provider].metrics = update_metrics(breakdown.providers[provider].metrics, record) # Update API key breakdown for this provider - if record.api_key not in breakdown.providers[provider].api_key_breakdown: - breakdown.providers[provider].api_key_breakdown[record.api_key] = KeyMetricWithMetadata( + if api_key not in breakdown.providers[provider].api_key_breakdown: + breakdown.providers[provider].api_key_breakdown[api_key] = KeyMetricWithMetadata( metrics=SpendMetrics(), metadata=_key_metadata(api_key_metadata, record.api_key), ) - breakdown.providers[provider].api_key_breakdown[record.api_key].metrics = update_metrics( - breakdown.providers[provider].api_key_breakdown[record.api_key].metrics, + breakdown.providers[provider].api_key_breakdown[api_key].metrics = update_metrics( + breakdown.providers[provider].api_key_breakdown[api_key].metrics, record, ) @@ -413,25 +421,26 @@ def update_breakdown_metrics( breakdown.endpoints[record.endpoint].metrics, record ) - # Update API key breakdown for this endpoint - if record.api_key not in breakdown.endpoints[record.endpoint].api_key_breakdown: - breakdown.endpoints[record.endpoint].api_key_breakdown[record.api_key] = KeyMetricWithMetadata( - metrics=SpendMetrics(), - metadata=_key_metadata(api_key_metadata, record.api_key), + if not is_ptu_sentinel: + # Update API key breakdown for this endpoint + if api_key not in breakdown.endpoints[record.endpoint].api_key_breakdown: + breakdown.endpoints[record.endpoint].api_key_breakdown[api_key] = KeyMetricWithMetadata( + metrics=SpendMetrics(), + metadata=_key_metadata(api_key_metadata, record.api_key), + ) + breakdown.endpoints[record.endpoint].api_key_breakdown[api_key].metrics = update_metrics( + breakdown.endpoints[record.endpoint].api_key_breakdown[api_key].metrics, + record, ) - breakdown.endpoints[record.endpoint].api_key_breakdown[record.api_key].metrics = update_metrics( - breakdown.endpoints[record.endpoint].api_key_breakdown[record.api_key].metrics, - record, - ) if not is_ptu_sentinel: # Update api key breakdown - if record.api_key not in breakdown.api_keys: - breakdown.api_keys[record.api_key] = KeyMetricWithMetadata( + if api_key not in breakdown.api_keys: + breakdown.api_keys[api_key] = KeyMetricWithMetadata( metrics=SpendMetrics(), metadata=_key_metadata(api_key_metadata, record.api_key), ) - breakdown.api_keys[record.api_key].metrics = update_metrics(breakdown.api_keys[record.api_key].metrics, record) + breakdown.api_keys[api_key].metrics = update_metrics(breakdown.api_keys[api_key].metrics, record) # Update entity-specific metrics if entity_id_field is provided if entity_id_field: @@ -445,13 +454,13 @@ def update_breakdown_metrics( if not is_ptu_sentinel: # Update API key breakdown for this entity - if record.api_key not in breakdown.entities[entity_value].api_key_breakdown: - breakdown.entities[entity_value].api_key_breakdown[record.api_key] = KeyMetricWithMetadata( + if api_key not in breakdown.entities[entity_value].api_key_breakdown: + breakdown.entities[entity_value].api_key_breakdown[api_key] = KeyMetricWithMetadata( metrics=SpendMetrics(), metadata=_key_metadata(api_key_metadata, record.api_key), ) - breakdown.entities[entity_value].api_key_breakdown[record.api_key].metrics = update_metrics( - breakdown.entities[entity_value].api_key_breakdown[record.api_key].metrics, + breakdown.entities[entity_value].api_key_breakdown[api_key].metrics = update_metrics( + breakdown.entities[entity_value].api_key_breakdown[api_key].metrics, record, ) @@ -566,7 +575,7 @@ def daily_activity_scope( entity_id_field: str, entity_id: str | list[str] | None, exclude_entity_ids: list[str] | None, - api_key: str | list[str] | None, + api_key: str | Sequence[str] | None, start_date: str, end_date: str, model: str | None, @@ -577,9 +586,7 @@ def daily_activity_scope( entity_ids: tuple[str, ...] | None = ( (entity_id,) if isinstance(entity_id, str) else tuple(entity_id) if entity_id is not None else None ) - api_keys: tuple[str, ...] | None = ( - None if api_key in (None, "") else (api_key,) if isinstance(api_key, str) else tuple(api_key) - ) + api_keys: Final = raw_api_key_filter(api_key) return DailyActivityScope( table=table_value, entity_id_field=entity_id_field, @@ -715,6 +722,29 @@ _GROUP_DATE_ENDPOINT: Final = 62 # 0b0111110 _GROUP_DATE_ENDPOINT_API_KEY: Final = 30 # 0b0011110 +def sum_spend_metrics(metrics: Sequence[SpendMetrics]) -> SpendMetrics: + return SpendMetrics( + spend=sum(metric.spend for metric in metrics), + flat_cost=sum(metric.flat_cost for metric in metrics), + ptu_hours=sum(metric.ptu_hours for metric in metrics), + prompt_tokens=sum(metric.prompt_tokens for metric in metrics), + completion_tokens=sum(metric.completion_tokens for metric in metrics), + cache_read_input_tokens=sum(metric.cache_read_input_tokens for metric in metrics), + cache_creation_input_tokens=sum(metric.cache_creation_input_tokens for metric in metrics), + compression_saved_tokens=sum(metric.compression_saved_tokens for metric in metrics), + compression_savings_spend=sum(metric.compression_savings_spend for metric in metrics), + prompt_caching_savings_spend=sum(metric.prompt_caching_savings_spend for metric in metrics), + gateway_injected_caching_savings_spend=sum(metric.gateway_injected_caching_savings_spend for metric in metrics), + autorouter_savings_spend=sum(metric.autorouter_savings_spend for metric in metrics), + total_tokens=sum(metric.total_tokens for metric in metrics), + successful_requests=sum(metric.successful_requests for metric in metrics), + failed_requests=sum(metric.failed_requests for metric in metrics), + api_requests=sum(metric.api_requests for metric in metrics), + total_response_time_ms=sum(metric.total_response_time_ms for metric in metrics), + timed_requests=sum(metric.timed_requests for metric in metrics), + ) + + def _record_to_spend_metrics(record: RollupMetricsRow) -> SpendMetrics: """Build a SpendMetrics directly from one already-aggregated rollup row. @@ -753,8 +783,7 @@ def _aggregate_grouping_sets_records_sync( Each row carries a `group_level` bitmask (from Postgres GROUPING()) that identifies which rollup level it belongs to. We dispatch the row's - pre-aggregated metrics straight into the matching bucket — no per-row - summing in Python and no nested update_metrics calls. + pre-aggregated metrics into the matching bucket, combining normalized names. """ total_metrics = SpendMetrics() grouped_data: Final[dict[str, GroupedData]] = {} @@ -771,20 +800,23 @@ def _aggregate_grouping_sets_records_sync( if existing is None: target[key] = MetricWithMetadata(metrics=metrics, metadata={}) else: - existing.metrics = metrics + existing.metrics = sum_spend_metrics((existing.metrics, metrics)) def assign_api_key_breakdown( target: dict[str, MetricWithMetadata], parent_key: str, - api_key: str, + api_key: str | None, metrics: SpendMetrics, ) -> None: parent = target.get(parent_key) if parent is None: parent = MetricWithMetadata(metrics=SpendMetrics(), metadata={}) target[parent_key] = parent - parent.api_key_breakdown[api_key] = KeyMetricWithMetadata( - metrics=metrics, metadata=_key_metadata(api_key_metadata, api_key) + bucket_key: Final = api_key_bucket(api_key) + previous: Final = parent.api_key_breakdown.get(bucket_key) + parent.api_key_breakdown[bucket_key] = KeyMetricWithMetadata( + metrics=sum_spend_metrics((previous.metrics, metrics)) if previous is not None else metrics, + metadata=_key_metadata(api_key_metadata, api_key), ) for record in records: @@ -803,25 +835,23 @@ def _aggregate_grouping_sets_records_sync( breakdown = ensure_date(record.date)["breakdown"] if level == _GROUP_DATE_API_KEY: - if record.api_key and not is_ptu_sentinel: - breakdown.api_keys[record.api_key] = KeyMetricWithMetadata( + if not is_ptu_sentinel: + breakdown.api_keys[api_key_bucket(record.api_key)] = KeyMetricWithMetadata( metrics=metrics, metadata=_key_metadata(api_key_metadata, record.api_key), ) elif level == _GROUP_DATE_MODEL: - if record.model: - assign_metric_with_metadata(breakdown.models, record.model, metrics) + assign_metric_with_metadata(breakdown.models, model_bucket(record.model), metrics) elif level == _GROUP_DATE_MODEL_API_KEY: - if record.model and record.api_key and not is_ptu_sentinel: - assign_api_key_breakdown(breakdown.models, record.model, record.api_key, metrics) + if not is_ptu_sentinel: + assign_api_key_breakdown(breakdown.models, model_bucket(record.model), record.api_key, metrics) elif level == _GROUP_DATE_MODEL_GROUP: - if record.model_group: - assign_metric_with_metadata(breakdown.model_groups, record.model_group, metrics) + assign_metric_with_metadata(breakdown.model_groups, model_group_bucket(record.model_group, None), metrics) elif level == _GROUP_DATE_MODEL_GROUP_API_KEY: - if record.model_group and record.api_key and not is_ptu_sentinel: + if not is_ptu_sentinel: assign_api_key_breakdown( breakdown.model_groups, - record.model_group, + model_group_bucket(record.model_group, None), record.api_key, metrics, ) @@ -836,14 +866,14 @@ def _aggregate_grouping_sets_records_sync( provider = record.custom_llm_provider or "unknown" assign_metric_with_metadata(breakdown.providers, provider, provider_metrics) elif level == _GROUP_DATE_PROVIDER_API_KEY: - if record.api_key and not is_ptu_sentinel: + if not is_ptu_sentinel: provider = record.custom_llm_provider or "unknown" assign_api_key_breakdown(breakdown.providers, provider, record.api_key, metrics) elif level == _GROUP_DATE_MCP: if record.mcp_namespaced_tool_name: assign_metric_with_metadata(breakdown.mcp_servers, record.mcp_namespaced_tool_name, metrics) elif level == _GROUP_DATE_MCP_API_KEY: - if record.mcp_namespaced_tool_name and record.api_key: + if record.mcp_namespaced_tool_name and not is_ptu_sentinel: assign_api_key_breakdown( breakdown.mcp_servers, record.mcp_namespaced_tool_name, @@ -854,7 +884,7 @@ def _aggregate_grouping_sets_records_sync( if record.endpoint: assign_metric_with_metadata(breakdown.endpoints, record.endpoint, metrics) elif level == _GROUP_DATE_ENDPOINT_API_KEY: - if record.endpoint and record.api_key: + if record.endpoint and not is_ptu_sentinel: assign_api_key_breakdown(breakdown.endpoints, record.endpoint, record.api_key, metrics) results: Final = [ @@ -1013,8 +1043,8 @@ def _fold_entity_rollups_sync( metrics = _record_to_spend_metrics(row) if row.api_key_rolled: bucket.metrics = metrics - elif row.api_key and row.api_key != PTU_SENTINEL_API_KEY: - bucket.api_key_breakdown[row.api_key] = KeyMetricWithMetadata( + elif row.api_key != PTU_SENTINEL_API_KEY: + bucket.api_key_breakdown[api_key_bucket(row.api_key)] = KeyMetricWithMetadata( metrics=metrics, metadata=_key_metadata(api_key_metadata, row.api_key) ) diff --git a/litellm/proxy/management_endpoints/daily_activity_routes.py b/litellm/proxy/management_endpoints/daily_activity_routes.py index 79e515944ce..1a3445aae52 100644 --- a/litellm/proxy/management_endpoints/daily_activity_routes.py +++ b/litellm/proxy/management_endpoints/daily_activity_routes.py @@ -22,6 +22,7 @@ from litellm.proxy.management_endpoints.common_daily_activity import ( parse_canonical_date_range, raise_public, spend_logs_window, + sum_spend_metrics, ) from litellm.proxy.management_endpoints.daily_activity_scopes import ( AGENT_RESOLVER, @@ -38,6 +39,7 @@ from litellm.proxy.management_endpoints.ptu_consumption import single_team_id, w from litellm.proxy.management_endpoints.team_endpoints import aggregated_date_range_error from litellm.proxy.management_helpers.utils import management_endpoint_wrapper from litellm.proxy.utils import PrismaClient, get_prisma_client_or_throw +from litellm.repositories.daily_activity_buckets import api_key_bucket from litellm.repositories.daily_activity_repository import DailyActivityRepository from litellm.router import Router from litellm.types.proxy.management_endpoints.common_daily_activity import ( @@ -106,29 +108,9 @@ async def _resolved_scope( return result -def _sum_metrics(metrics: Sequence[SpendMetrics]) -> SpendMetrics: - return SpendMetrics( - spend=sum(metric.spend for metric in metrics), - flat_cost=sum(metric.flat_cost for metric in metrics), - prompt_tokens=sum(metric.prompt_tokens for metric in metrics), - completion_tokens=sum(metric.completion_tokens for metric in metrics), - cache_read_input_tokens=sum(metric.cache_read_input_tokens for metric in metrics), - cache_creation_input_tokens=sum(metric.cache_creation_input_tokens for metric in metrics), - compression_saved_tokens=sum(metric.compression_saved_tokens for metric in metrics), - compression_savings_spend=sum(metric.compression_savings_spend for metric in metrics), - prompt_caching_savings_spend=sum(metric.prompt_caching_savings_spend for metric in metrics), - gateway_injected_caching_savings_spend=sum(metric.gateway_injected_caching_savings_spend for metric in metrics), - autorouter_savings_spend=sum(metric.autorouter_savings_spend for metric in metrics), - total_tokens=sum(metric.total_tokens for metric in metrics), - successful_requests=sum(metric.successful_requests for metric in metrics), - failed_requests=sum(metric.failed_requests for metric in metrics), - api_requests=sum(metric.api_requests for metric in metrics), - total_response_time_ms=sum(metric.total_response_time_ms for metric in metrics), - timed_requests=sum(metric.timed_requests for metric in metrics), - ) - - def _key_metadata(api_key: str, metadata: Mapping[str, KeyMetadataRow]) -> KeyMetadata: + if not api_key: + return KeyMetadata(key_alias="Unattributed", key_exists=False) row: Final[KeyMetadataRow | None] = metadata.get(api_key) if row is None: return KeyMetadata() @@ -143,7 +125,7 @@ def _key_metadata(api_key: str, metadata: Mapping[str, KeyMetadataRow]) -> KeyMe def _key_activity_row(row: KeySpendRow, metadata: Mapping[str, KeyMetadataRow]) -> KeySpendActivityRow: return KeySpendActivityRow( - api_key=row.api_key, + api_key=api_key_bucket(row.api_key), metrics=KeySpendMetrics( spend=row.spend, prompt_tokens=row.prompt_tokens, @@ -168,7 +150,7 @@ async def _key_activity_rows( frozenset((resolved_scope.scope.start_date, resolved_scope.scope.end_date)) ) metadata: Final[Mapping[str, KeyMetadataRow]] = await repository.key_metadata( - frozenset(row.api_key for row in rows), + frozenset(row.api_key for row in rows if row.api_key), spend_window, ) return [_key_activity_row(row, metadata) for row in rows] @@ -231,11 +213,11 @@ def _fold_key_metrics(api_key: str, response: SpendAnalyticsPaginatedResponse) - metadata: Final[KeyMetadata] = next( day.breakdown.api_keys[api_key].metadata for day in response.results if api_key in day.breakdown.api_keys ) - return KeyActivityRow(api_key=api_key, metrics=_sum_metrics(metrics), metadata=metadata) + return KeyActivityRow(api_key=api_key, metrics=sum_spend_metrics(metrics), metadata=metadata) def _search_rows(keys: Sequence[str], response: SpendAnalyticsPaginatedResponse) -> list[KeyActivityRow]: - return [row for key in keys if (row := _fold_key_metrics(key, response)) is not None] + return [row for key in keys if (row := _fold_key_metrics(api_key_bucket(key), response)) is not None] def _csv_cell(value: object) -> object: diff --git a/litellm/proxy/management_endpoints/daily_activity_scopes.py b/litellm/proxy/management_endpoints/daily_activity_scopes.py index 77246f5dda6..ca1a96f9d84 100644 --- a/litellm/proxy/management_endpoints/daily_activity_scopes.py +++ b/litellm/proxy/management_endpoints/daily_activity_scopes.py @@ -17,6 +17,7 @@ from litellm.proxy.management_endpoints.organization_endpoints import resolve_or from litellm.proxy.management_endpoints.tag_management_endpoints import get_tag_daily_activity_api_key_filter from litellm.proxy.management_endpoints.team_endpoints import resolve_team_daily_activity_scope from litellm.proxy.utils import PrismaClient +from litellm.repositories.daily_activity_buckets import raw_api_key_filter from litellm.types.repositories.daily_activity import DailyActivityScope, DailyActivityTable @@ -77,13 +78,7 @@ def _build_scope( ) -> ResolvedScope: start_date: Final[str] = query.start_date or "" end_date: Final[str] = query.end_date or "" - api_keys: Final[tuple[str, ...] | None] = ( - None - if api_key_filter is None or api_key_filter == "" - else (api_key_filter,) - if isinstance(api_key_filter, str) - else tuple(api_key_filter) - ) + api_keys: Final = raw_api_key_filter(api_key_filter) return ResolvedScope( scope=DailyActivityScope( table=resolver.table, diff --git a/litellm/proxy/management_endpoints/tag_management_endpoints.py b/litellm/proxy/management_endpoints/tag_management_endpoints.py index c12a90ed6a2..ea7810a618c 100644 --- a/litellm/proxy/management_endpoints/tag_management_endpoints.py +++ b/litellm/proxy/management_endpoints/tag_management_endpoints.py @@ -31,6 +31,7 @@ from litellm.proxy.management_endpoints.common_daily_activity import ( get_daily_activity, ) from litellm.proxy.management_helpers.utils import handle_budget_for_entity +from litellm.repositories.daily_activity_buckets import api_key_from_bucket from litellm.repositories.model_repository import ModelRepository from litellm.repositories.table_repositories import ( DailyTagSpendRepository, @@ -198,16 +199,17 @@ async def get_tag_daily_activity_api_key_filter( user_api_key_dict: UserAPIKeyAuth, requested_api_key: str | None, ) -> str | list[str] | None: + selected_key: Final = api_key_from_bucket(requested_api_key) if requested_api_key else None user_role: Final = user_api_key_dict.user_role if user_api_key_has_admin_view(user_api_key_dict) or (user_role is None or not user_role.is_internal_user_role): - return requested_api_key + return [selected_key] if selected_key is not None else None scoped_api_keys: Final = await _get_internal_user_api_keys( prisma_client=prisma_client, user_api_key_dict=user_api_key_dict, ) - if requested_api_key is not None: - return requested_api_key if requested_api_key in scoped_api_keys else [] + if selected_key is not None: + return [selected_key] if selected_key in scoped_api_keys else [] return scoped_api_keys diff --git a/litellm/proxy/management_endpoints/team_endpoints.py b/litellm/proxy/management_endpoints/team_endpoints.py index 50b700e712d..b712ba42a6a 100644 --- a/litellm/proxy/management_endpoints/team_endpoints.py +++ b/litellm/proxy/management_endpoints/team_endpoints.py @@ -226,6 +226,7 @@ from litellm.proxy.management_helpers.utils import ( from litellm.proxy.utils import PrismaClient, ProxyLogging, handle_exception_on_proxy from litellm.repositories.budget_repository import BudgetRepository from litellm.repositories.chunked_in import delete_many_in, find_many_in +from litellm.repositories.daily_activity_buckets import api_key_from_bucket from litellm.repositories.organization_repository import OrganizationRepository from litellm.repositories.prisma_protocols import TableActions from litellm.repositories.table_repositories import ( @@ -7329,11 +7330,12 @@ async def resolve_team_daily_activity_scope( if not user_api_keys: user_api_keys = [] + selected_key: Final = api_key_from_bucket(api_key) if api_key else None final_api_key_filter: str | list[str] | None if user_api_keys is None: - final_api_key_filter = api_key - elif api_key: - final_api_key_filter = api_key if api_key in user_api_keys else [] + final_api_key_filter = [selected_key] if selected_key is not None else None + elif selected_key is not None: + final_api_key_filter = [selected_key] if selected_key in user_api_keys else [] else: final_api_key_filter = user_api_keys diff --git a/litellm/repositories/daily_activity_buckets.py b/litellm/repositories/daily_activity_buckets.py new file mode 100644 index 00000000000..ec6dd0f8c84 --- /dev/null +++ b/litellm/repositories/daily_activity_buckets.py @@ -0,0 +1,39 @@ +from collections.abc import Sequence +from typing import Final + +from litellm.constants import UNKNOWN_MODEL_SPEND_LOG_MODEL + +UNKNOWN_MODEL_GROUP: Final = "unknown-model-group" +UNATTRIBUTED_API_KEY: Final = "unattributed" +_RAW_API_KEY_PREFIX: Final = "raw:" + + +def model_bucket(model: str | None) -> str: + return model or UNKNOWN_MODEL_SPEND_LOG_MODEL + + +def model_group_bucket(model_group: str | None, model: str | None) -> str: + return model_group or model or UNKNOWN_MODEL_GROUP + + +def api_key_bucket(api_key: str | None) -> str: + if not api_key: + return UNATTRIBUTED_API_KEY + if api_key == UNATTRIBUTED_API_KEY or api_key.startswith(_RAW_API_KEY_PREFIX): + return _RAW_API_KEY_PREFIX + api_key + return api_key + + +def api_key_from_bucket(api_key: str) -> str: + if api_key == UNATTRIBUTED_API_KEY: + return "" + return api_key.removeprefix(_RAW_API_KEY_PREFIX) + + +def raw_api_key_filter(api_key: str | Sequence[str] | None) -> tuple[str, ...] | None: + """String filters are reporting IDs; sequences are raw keys resolved by authorization.""" + if api_key is None or api_key == "": + return None + if isinstance(api_key, str): + return (api_key_from_bucket(api_key),) + return tuple(api_key) diff --git a/litellm/repositories/daily_activity_repository.py b/litellm/repositories/daily_activity_repository.py index 4ca23e4e89e..321468c4976 100644 --- a/litellm/repositories/daily_activity_repository.py +++ b/litellm/repositories/daily_activity_repository.py @@ -10,6 +10,7 @@ from typing_extensions import assert_never from litellm import constants from litellm._logging import verbose_proxy_logger +from litellm.constants import UNKNOWN_MODEL_SPEND_LOG_MODEL from litellm.repositories.chunked_in import find_many_in from litellm.repositories.daily_activity_sql import ( ExportCursor, @@ -344,7 +345,13 @@ class DailyActivityRepository: "date": {"gte": adjusted_start, "lte": adjusted_end}, **({scope.entity_id_field: {"in": list(scope.entity_ids)}} if scope.entity_ids is not None else {}), **exclusion_filter, - **({"model": scope.model} if scope.model else {}), + **( + {"AND": [{"OR": [{"model": None}, {"model": ""}, {"model": scope.model}]}]} + if scope.model == UNKNOWN_MODEL_SPEND_LOG_MODEL + else {"model": scope.model} + if scope.model + else {} + ), **({"api_key": {"in": list(scope.api_keys)}} if scope.api_keys is not None else {}), } count, rows = await asyncio.gather( diff --git a/litellm/repositories/daily_activity_sql.py b/litellm/repositories/daily_activity_sql.py index 7ec9fa67a67..4bb52767de6 100644 --- a/litellm/repositories/daily_activity_sql.py +++ b/litellm/repositories/daily_activity_sql.py @@ -8,11 +8,18 @@ from typing import Final from typing_extensions import assert_never from litellm import constants -from litellm.constants import PTU_SENTINEL_API_KEY +from litellm.constants import PTU_SENTINEL_API_KEY, UNKNOWN_MODEL_SPEND_LOG_MODEL +from litellm.repositories.daily_activity_buckets import UNATTRIBUTED_API_KEY, UNKNOWN_MODEL_GROUP from litellm.types.repositories.daily_activity import DailyActivityScope, DailyActivityTable, ExportType _API_KEY_ROLLED_UP_BIT: Final = 32 -_MODEL_GROUP_EXPR: Final = "COALESCE(NULLIF(model_group, ''), model)" +_MODEL_EXPR: Final = f"COALESCE(NULLIF(model, ''), '{UNKNOWN_MODEL_SPEND_LOG_MODEL}')" +_MODEL_GROUP_EXPR: Final = f"COALESCE(NULLIF(model_group, ''), NULLIF(model, ''), '{UNKNOWN_MODEL_GROUP}')" + +_API_KEY_EXPR: Final = ( + f"CASE WHEN api_key = '' THEN '{UNATTRIBUTED_API_KEY}' " + f"WHEN api_key = '{UNATTRIBUTED_API_KEY}' OR api_key LIKE 'raw:%' THEN 'raw:' || api_key ELSE api_key END" +) @dataclass(frozen=True, slots=True) @@ -86,7 +93,11 @@ def build_where_clause(scope: DailyActivityScope, *, start_index: int = 1) -> tu if scope.exclude_entity_ids else () ), - *((f"model = ${model_index}",) if scope.model else ()), + *( + (f"{_MODEL_EXPR if scope.model == UNKNOWN_MODEL_SPEND_LOG_MODEL else 'model'} = ${model_index}",) + if scope.model + else () + ), *( ("FALSE",) if scope.api_keys == () @@ -143,7 +154,7 @@ def _top_api_keys_sql(pg_table: str, where_clause: str, *, sentinel_param: int, FROM "{pg_table}" WHERE {where_clause} AND api_key <> ${sentinel_param} GROUP BY api_key - ORDER BY SUM(spend::numeric) DESC, api_key + ORDER BY SUM(spend::numeric) DESC, {_API_KEY_EXPR} LIMIT ${limit_param} """ @@ -162,13 +173,13 @@ def build_aggregated_sql(scope: DailyActivityScope, *, api_key_limit: int) -> Sq (SELECT date, NULL::text AS api_key, - model, + {_MODEL_EXPR} AS model, {_MODEL_GROUP_EXPR} AS model_group, custom_llm_provider, mcp_namespaced_tool_name, endpoint, (GROUPING(date) << 6) | {_API_KEY_ROLLED_UP_BIT} - | GROUPING(model, {_MODEL_GROUP_EXPR}, + | GROUPING({_MODEL_EXPR}, {_MODEL_GROUP_EXPR}, custom_llm_provider, mcp_namespaced_tool_name, endpoint) AS group_level, NULL::bigint AS distinct_api_keys,{metric_select} @@ -176,7 +187,7 @@ def build_aggregated_sql(scope: DailyActivityScope, *, api_key_limit: int) -> Sq WHERE {where_clause} GROUP BY GROUPING SETS ( (date), - (date, model), + (date, {_MODEL_EXPR}), (date, {_MODEL_GROUP_EXPR}), (date, custom_llm_provider), (date, mcp_namespaced_tool_name), @@ -190,12 +201,12 @@ def build_aggregated_sql(scope: DailyActivityScope, *, api_key_limit: int) -> Sq SELECT date, api_key, - model, + {_MODEL_EXPR} AS model, {_MODEL_GROUP_EXPR} AS model_group, custom_llm_provider, mcp_namespaced_tool_name, endpoint, - GROUPING(date, api_key, model, {_MODEL_GROUP_EXPR}, + GROUPING(date, api_key, {_MODEL_EXPR}, {_MODEL_GROUP_EXPR}, custom_llm_provider, mcp_namespaced_tool_name, endpoint) AS group_level, MAX(top_api_keys.distinct_api_keys) AS distinct_api_keys,{metric_select} @@ -203,7 +214,7 @@ def build_aggregated_sql(scope: DailyActivityScope, *, api_key_limit: int) -> Sq WHERE {where_clause} GROUP BY GROUPING SETS ( (date, api_key), - (date, model, api_key), + (date, {_MODEL_EXPR}, api_key), (date, {_MODEL_GROUP_EXPR}, api_key), (date, custom_llm_provider, api_key), (date, mcp_namespaced_tool_name, api_key), @@ -313,7 +324,7 @@ def build_key_page_sql(scope: DailyActivityScope, *, offset: int, limit: int) -> FROM (SELECT 1) AS one LEFT JOIN LATERAL ( SELECT * FROM ranked - ORDER BY rank_spend DESC, api_key + ORDER BY rank_spend DESC, {_API_KEY_EXPR} LIMIT ${limit_param} OFFSET ${offset_param} ) AS page ON TRUE """ @@ -376,7 +387,7 @@ def build_key_search_sql(scope: DailyActivityScope, *, search: str, limit: int) WHERE {where_clause} AND api_key <> ${sentinel_param} AND ( - api_key ILIKE ${search_param} ESCAPE '\\' + {_API_KEY_EXPR} ILIKE ${search_param} ESCAPE '\\' OR api_key IN ( SELECT v.token FROM "LiteLLM_VerificationToken" v LEFT JOIN "LiteLLM_UserTable" u ON u.user_id = v.user_id @@ -392,7 +403,7 @@ def build_key_search_sql(scope: DailyActivityScope, *, search: str, limit: int) ) ) GROUP BY api_key - ORDER BY SUM(spend::numeric) DESC, api_key + ORDER BY SUM(spend::numeric) DESC, {_API_KEY_EXPR} LIMIT ${limit_param} """ return SqlQuery(sql=sql, params=(*where_params, f"%{escaped}%", PTU_SENTINEL_API_KEY, limit)) @@ -406,15 +417,13 @@ def build_model_top_keys_sql( model_param: Final = len(where_params) + 1 sentinel_param: Final = model_param + 1 limit_param: Final = sentinel_param + 1 - model_clause: Final = ( - f"COALESCE(NULLIF(model_group, ''), model) = ${model_param}" if by_model_group else f"model = ${model_param}" - ) + model_clause: Final = f"{_MODEL_GROUP_EXPR if by_model_group else _MODEL_EXPR} = ${model_param}" sql: Final = f""" SELECT api_key,{_key_spend_select()} FROM "{PRISMA_TO_PG_TABLE[scope.table]}" WHERE {where_clause} AND {model_clause} AND api_key <> ${sentinel_param} GROUP BY api_key - ORDER BY SUM(spend::numeric) DESC, api_key + ORDER BY SUM(spend::numeric) DESC, {_API_KEY_EXPR} LIMIT ${limit_param} """ return SqlQuery(sql=sql, params=(*where_params, model_group, PTU_SENTINEL_API_KEY, limit)) @@ -431,7 +440,7 @@ def build_cache_leakage_keys_sql(scope: DailyActivityScope, *, limit: int) -> Sq WHERE {where_clause} AND api_key <> ${sentinel_param} GROUP BY api_key HAVING SUM(prompt_tokens) - SUM(cache_read_input_tokens) > 0 - ORDER BY SUM(prompt_tokens) - SUM(cache_read_input_tokens) DESC, api_key + ORDER BY SUM(prompt_tokens) - SUM(cache_read_input_tokens) DESC, {_API_KEY_EXPR} LIMIT ${limit_param} """ return SqlQuery(sql=sql, params=(*where_params, PTU_SENTINEL_API_KEY, limit)) @@ -483,7 +492,7 @@ def build_export_sql( {alias_expression} AS entity_alias, {output_key} AS api_key, {user_fields}, - {"NULLIF(COALESCE(scoped.model, ''), '')" if export_type is ExportType.DAILY_WITH_MODELS else "NULL::text"} AS model, + {_MODEL_EXPR if export_type is ExportType.DAILY_WITH_MODELS else "NULL::text"} AS model, COALESCE(SUM(scoped.spend), 0)::float AS spend, {flat_cost} AS flat_cost, COALESCE(SUM(scoped.prompt_tokens), 0)::bigint AS prompt_tokens, @@ -521,25 +530,26 @@ def _export_grouping(export_type: ExportType) -> tuple[str, str, str, tuple[str, ) if export_type is ExportType.DAILY_WITH_KEYS: return ( - "scoped.api_key", - "NULLIF(scoped.api_key, '')", - "MAX(COALESCE(vt.key_alias, dvt.key_alias)) AS key_alias, " + _API_KEY_EXPR, + _API_KEY_EXPR, + "MAX(CASE WHEN scoped.api_key = '' THEN 'Unattributed' " + "ELSE COALESCE(vt.key_alias, dvt.key_alias) END) AS key_alias, " "MAX(COALESCE(vt.user_id, dvt.user_id)) AS user_id, MAX(u.user_email) AS user_email", ( - 'LEFT JOIN "LiteLLM_VerificationToken" vt ON vt.token = scoped.api_key', + "LEFT JOIN \"LiteLLM_VerificationToken\" vt ON vt.token = scoped.api_key AND scoped.api_key <> ''", """LEFT JOIN LATERAL ( SELECT key_alias, user_id FROM "LiteLLM_DeletedVerificationToken" WHERE token = scoped.api_key ORDER BY deleted_at DESC LIMIT 1 - ) dvt ON vt.token IS NULL""", + ) dvt ON vt.token IS NULL AND scoped.api_key <> ''""", 'LEFT JOIN "LiteLLM_UserTable" u ON u.user_id = COALESCE(vt.user_id, dvt.user_id)', ), ) if export_type is ExportType.DAILY_WITH_MODELS: return ( - "COALESCE(scoped.model, '')", + _MODEL_EXPR, "NULL::text", "NULL::text AS key_alias, NULL::text AS user_id, NULL::text AS user_email", (), diff --git a/tests/integration/routing/test_auto_router_management.py b/tests/integration/routing/test_auto_router_management.py index 5aed328ba7b..442dae4d689 100644 --- a/tests/integration/routing/test_auto_router_management.py +++ b/tests/integration/routing/test_auto_router_management.py @@ -9,7 +9,7 @@ import httpx import pytest import yaml from integration._support.client import JSON_OBJECT, Gateway, eventually, object_value, string_value, until_asserted -from integration._support.database import read_rows, scratch_database +from integration._support.database import read_rows, scratch_database, write_rows from integration._support.process import owned_proxy from integration._support.wire import Reply, Request, Wire, wire_server from pydantic import JsonValue, TypeAdapter @@ -639,3 +639,25 @@ def test_auto_router_session_tracks_key_scoped_spend_and_bounds_long_session_ids ), response.text assert tuple(request for request in complex_tier.drain() if request.method == "POST") == (), response.text assert tuple(request for request in baseline.drain() if request.method == "POST") == (), response.text + + +@pytest.mark.parametrize("api_key,expected", [("unattributed", 2.0), ("raw:unattributed", 4.0)]) +def test_benchmarks_preserve_literal_key_identity(gateway: Gateway, api_key: str, expected: float) -> None: + user_id: Final = f"benchmark-key-{uuid.uuid4().hex}" + for key, savings in (("", 3.0), ("unattributed", 2.0), ("raw:unattributed", 4.0)): + write_rows( + 'INSERT INTO "LiteLLM_DailyUserSpend" ' + "(id, user_id, date, api_key, model, autorouter_savings_spend, updated_at) " + "VALUES (%s, %s, '2026-09-09', %s, 'model', %s, now())", + (str(uuid.uuid4()), user_id, key, str(savings)), + ) + try: + response: Final = gateway.request( + "GET", + "/auto_router/benchmarks", + params={"start_date": "2026-09-09", "end_date": "2026-09-09", "user_id": user_id, "api_key": api_key}, + ) + assert response.status_code == 200, response.text + assert object_value(object_value(response.json())["totals"])["saved_spend"] == expected + finally: + write_rows('DELETE FROM "LiteLLM_DailyUserSpend" WHERE user_id = %s', (user_id,)) diff --git a/tests/integration/spend/test_daily_activity_aggregated_breakdowns.py b/tests/integration/spend/test_daily_activity_aggregated_breakdowns.py index 1c315a6e694..ff60cb9d174 100644 --- a/tests/integration/spend/test_daily_activity_aggregated_breakdowns.py +++ b/tests/integration/spend/test_daily_activity_aggregated_breakdowns.py @@ -1,3 +1,5 @@ +import csv +import io import uuid from collections.abc import Sequence from datetime import datetime, timedelta @@ -12,6 +14,7 @@ from tests.integration._support.database import write_rows _URL: Final = "/user/daily/activity/aggregated" _RESULTS: Final = TypeAdapter(list[dict[str, JsonValue]]) +_NUMBER: Final = TypeAdapter(float) def _unique_day() -> str: @@ -43,6 +46,181 @@ def _row_id() -> str: return f"agg-{uuid.uuid4().hex}" +@pytest.mark.parametrize("suffix", ["/aggregated", ""]) +def test_missing_attribution_reconciles_and_filters_without_key_collisions(gateway: Gateway, suffix: str) -> None: + day: Final = _unique_day() + rows: Final = ( + ("", None, None, 3.0), + ("", "", "", 0.0), + ("known-key", "unknown-model", None, -1.0), + ("unattributed", "known-model", "", 2.0), + ("raw:unattributed", "known-model", "named-group", 4.0), + ("known-key", "", "unknown-model-group", 5.0), + ) + _seed( + day, + [ + (_row_id(), "bucket-user", day, key, model, group, "", "tool", "/mcp", 10, spend, 1, 1) + for key, model, group, spend in rows + ], + ) + params: Final = {"start_date": day, "end_date": day} + url: Final = f"/user/daily/activity{suffix}" + try: + response: Final = gateway.request("GET", url, params=params) + assert response.status_code == 200, response.text + body: Final = object_value(response.json()) + results: Final = _RESULTS.validate_python(body["results"]) + assert len(results) == 1 + totals: Final = object_value(results[0]["metrics"]) + assert totals["spend"] == 13.0 + assert totals["api_requests"] == 6 + breakdown: Final = object_value(results[0]["breakdown"]) + for dimension in ("models", "model_groups", "api_keys", "providers", "mcp_servers", "endpoints"): + buckets: Final = _RESULTS.validate_python(list(object_value(breakdown[dimension]).values())) + for metric in ("spend", "api_requests", "prompt_tokens"): + assert ( + sum(_NUMBER.validate_python(object_value(bucket["metrics"])[metric]) for bucket in buckets) + == totals[metric] + ) + models: Final = object_value(breakdown["models"]) + assert set(models) == {"known-model", "unknown-model"} + assert object_value(object_value(models["unknown-model"])["metrics"])["api_requests"] == 4 + groups: Final = object_value(breakdown["model_groups"]) + assert set(groups) == {"unknown-model-group", "unknown-model", "known-model", "named-group"} + assert object_value(object_value(groups["unknown-model-group"])["metrics"])["spend"] == 8.0 + keys: Final = object_value(breakdown["api_keys"]) + assert set(keys) == {"unattributed", "raw:unattributed", "raw:raw:unattributed", "known-key"} + unknown_key: Final = object_value(keys["unattributed"]) + assert object_value(unknown_key["metadata"]) == { + "key_alias": "Unattributed", + "key_exists": False, + "team_id": None, + "user_id": None, + "user_email": None, + } + for dimension in ("models", "model_groups", "providers", "mcp_servers", "endpoints"): + for bucket in _RESULTS.validate_python(list(object_value(breakdown[dimension]).values())): + nested: Final = _RESULTS.validate_python(list(object_value(bucket["api_key_breakdown"]).values())) + assert ( + sum(_NUMBER.validate_python(object_value(key["metrics"])["spend"]) for key in nested) + == object_value(bucket["metrics"])["spend"] + ) + for key, spend, count in ( + ("unattributed", 3.0, 2), + ("raw:unattributed", 2.0, 1), + ("raw:raw:unattributed", 4.0, 1), + ): + filtered: Final = gateway.request("GET", url, params={**params, "api_key": key}) + assert filtered.status_code == 200, filtered.text + metadata: Final = object_value(object_value(filtered.json())["metadata"]) + assert metadata["total_spend"] == spend + assert metadata["total_api_requests"] == count + model_response: Final = gateway.request("GET", url, params={**params, "model": "unknown-model"}) + assert model_response.status_code == 200, model_response.text + model_metadata: Final = object_value(object_value(model_response.json())["metadata"]) + assert model_metadata["total_spend"] == 7.0 + assert model_metadata["total_api_requests"] == 4 + finally: + _clean(day) + + +@pytest.mark.parametrize( + ("route", "extra", "expected"), + [ + ("keys", {}, {"unattributed": 3.0, "raw:unattributed": 2.0, "known-key": 7.0}), + ("search", {"search": "Unattributed"}, {"unattributed": 3.0, "raw:unattributed": 2.0}), + ("model_top_keys", {"model_group": "unknown-model-group"}, {"unattributed": 3.0}), + ( + "model_top_keys", + {"model_group": "unknown-model", "by_model_group": "false"}, + {"unattributed": 3.0, "raw:unattributed": 2.0}, + ), + ], +) +def test_missing_attribution_key_drilldowns( + gateway: Gateway, + route: str, + extra: dict[str, str], + expected: dict[str, float], +) -> None: + day: Final = _unique_day() + _seed( + day, + [ + (_row_id(), "bucket-user", day, "", None, None, "", None, None, 0, 3.0, 1, 1), + (_row_id(), "bucket-user", day, "unattributed", "unknown-model", "", "", None, None, 0, 2.0, 1, 1), + (_row_id(), "bucket-user", day, "known-key", "known-model", "", "", None, None, 0, 7.0, 1, 1), + ], + ) + try: + response: Final = gateway.request( + "GET", f"{_URL}/{route}", params={"start_date": day, "end_date": day, **extra} + ) + assert response.status_code == 200, response.text + rows: Final = _RESULTS.validate_python(object_value(response.json())["api_keys"]) + assert {str(row["api_key"]): object_value(row["metrics"])["spend"] for row in rows} == expected + unattributed: Final = next(row for row in rows if row["api_key"] == "unattributed") + assert object_value(unattributed["metadata"])["key_alias"] == "Unattributed" + assert object_value(unattributed["metadata"])["key_exists"] is False + finally: + _clean(day) + + +@pytest.mark.parametrize("format", ["json", "csv"]) +@pytest.mark.parametrize( + ("export_type", "field", "expected"), + [ + ("daily_with_keys", "api_key", {"unattributed": 3.0, "raw:unattributed": 2.0, "raw:raw:unattributed": 4.0}), + ("daily_with_models", "model", {"unknown-model": 5.0, "known-model": 4.0}), + ], +) +def test_export_bucket_ids_round_trip_across_streaming_batches( + gateway: Gateway, format: str, export_type: str, field: str, expected: dict[str, float] +) -> None: + day: Final = _unique_day() + _seed( + day, + [ + (_row_id(), "export-user", day, "", None, None, "", None, None, 0, 3.0, 1, 1), + (_row_id(), "export-user", day, "", "", "", "", None, None, 0, 0.0, 1, 1), + (_row_id(), "export-user", day, "unattributed", "unknown-model", "", "", None, None, 0, 2.0, 1, 1), + (_row_id(), "export-user", day, "raw:unattributed", "known-model", "", "", None, None, 0, 4.0, 1, 1), + ], + ) + write_rows( + 'INSERT INTO "LiteLLM_DailyUserSpend" (id, user_id, date, api_key, model, spend, api_requests, updated_at) ' + "SELECT %s || i, 'export-user', %s, 'raw:a-key-' || i, 'a-model-' || i, 1, 1, now() " + "FROM generate_series(1, 1001) i", + (str(uuid.uuid4()), day), + ) + params: Final = {"start_date": day, "end_date": day} + try: + response: Final = gateway.request( + "GET", "/user/daily/activity/export", params={**params, "export_type": export_type, "format": format} + ) + assert response.status_code == 200, response.text + rows: Final = ( + _RESULTS.validate_python(response.json()) + if format == "json" + else _RESULTS.validate_python(list(csv.DictReader(io.StringIO(response.text)))) + ) + buckets: Final = {str(row[field]): _NUMBER.validate_python(row["spend"]) for row in rows} + filler: Final = "raw:raw:a-key-" if field == "api_key" else "a-model-" + assert buckets == {**{f"{filler}{i}": 1.0 for i in range(1, 1002)}, **expected} + assert len(rows) == len(buckets) + assert sum(buckets.values()) == 1010.0 + for bucket, spend in expected.items(): + filtered: Final = _activity(gateway, day, **{field: bucket}) + assert object_value(filtered["metadata"])["total_spend"] == spend + if field == "api_key": + unattributed: Final = next(row for row in rows if row["api_key"] == "unattributed") + assert unattributed["key_alias"] == "Unattributed" + assert unattributed["user_id"] in (None, "") + finally: + _clean(day) + + def _ranked_key_rows(day: str, count: int) -> list[tuple[object, ...]]: return [ ( diff --git a/tests/integration/spend/test_daily_activity_routes.py b/tests/integration/spend/test_daily_activity_routes.py index 520c8ccc6d8..9bfe4bb50ce 100644 --- a/tests/integration/spend/test_daily_activity_routes.py +++ b/tests/integration/spend/test_daily_activity_routes.py @@ -451,9 +451,23 @@ def test_tag_daily_activity_scopes_internal_user_keys_and_filters( seeded_row(TAG_SPEND, "tag", "beta", digest_one_b, DAY), seeded_row(TAG_SPEND, "tag", "gamma", digest_one_a, DAY), seeded_row(TAG_SPEND, "tag", "alpha", digest_two, DAY), + seeded_row(TAG_SPEND, "tag", "alpha", "", DAY), replace(standard_model_row, model="gpt-4o"), ) with daily_rows(rows): + for suffix in ("", "/aggregated"): + bucket_params: Final = {"start_date": DAY, "end_date": DAY, "tags": "alpha", "api_key": "unattributed"} + bucket_url: Final = f"/tag/daily/activity{suffix}" + denied_bucket: Final = gateway.request("GET", bucket_url, params=bucket_params, key=key_one_a) + admin_bucket: Final = gateway.request("GET", bucket_url, params=bucket_params) + assert denied_bucket.status_code == 200, denied_bucket.text + assert admin_bucket.status_code == 200, admin_bucket.text + denied_body: Final = SpendAnalyticsPaginatedResponse.model_validate_json(denied_bucket.content) + allowed_body: Final = SpendAnalyticsPaginatedResponse.model_validate_json(admin_bucket.content) + assert denied_body.results == [] + assert denied_body.metadata.total_api_requests == 0 + assert allowed_body.metadata.total_api_requests == 1 + assert set(allowed_body.results[0].breakdown.api_keys) == {"unattributed"} scoped: Final = _tag_activity( gateway, key=key_one_a, diff --git a/tests/unit/proxy/management_endpoints/test_common_daily_activity.py b/tests/unit/proxy/management_endpoints/test_common_daily_activity.py index 01d2f7007a1..c80fa3572fd 100644 --- a/tests/unit/proxy/management_endpoints/test_common_daily_activity.py +++ b/tests/unit/proxy/management_endpoints/test_common_daily_activity.py @@ -2551,3 +2551,146 @@ async def test_get_daily_activity_rejects_non_canonical_dates_before_querying(st assert error.value.detail == {"error": "start_date and end_date must be valid YYYY-MM-DD dates"} mock_table.count.assert_not_awaited() mock_table.find_many.assert_not_awaited() + + +@pytest.mark.parametrize("missing", [None, ""]) +def test_missing_dimensions_preserve_all_paginated_metrics(missing: str | None) -> None: + records: Final = ( + _daily_user_spend_record(user_id="user", api_key="", spend=-0.25, model=missing, model_group=missing), + _daily_user_spend_record(user_id="user", api_key="", spend=0.0, model=missing, model_group=missing), + ) + result: Final = common_daily_activity_module._aggregate_spend_records_sync( + records=records, api_key_metadata={}, entity_id_field="user_id", entity_metadata_field=None + ) + day: Final = result["results"][0] + assert day.metrics == result["totals"] + for bucket in ( + day.breakdown.models["unknown-model"], + day.breakdown.model_groups["unknown-model-group"], + day.breakdown.providers["openai"], + day.breakdown.endpoints["/chat/completions"], + day.breakdown.entities["user"], + ): + assert bucket.metrics == day.metrics + assert bucket.api_key_breakdown == day.breakdown.api_keys + key: Final = day.breakdown.api_keys["unattributed"] + assert key.metrics == day.metrics + assert key.metadata.model_dump(exclude_none=True) == {"key_alias": "Unattributed", "key_exists": False} + + +def test_normalized_rollups_add_colliding_model_metrics_and_keep_key_identities() -> None: + records: Final = ( + _grouping_row(common_daily_activity_module._GROUP_DATE_MODEL, model=None, spend=1.0), + _grouping_row(common_daily_activity_module._GROUP_DATE_MODEL, model="", spend=2.0), + _grouping_row(common_daily_activity_module._GROUP_DATE_MODEL, model="unknown-model", spend=3.0), + _grouping_row(common_daily_activity_module._GROUP_DATE_MODEL_API_KEY, model=None, api_key="", spend=1.0), + _grouping_row(common_daily_activity_module._GROUP_DATE_MODEL_API_KEY, model="", api_key="", spend=2.0), + _grouping_row( + common_daily_activity_module._GROUP_DATE_MODEL_API_KEY, + model="unknown-model", + api_key="unattributed", + spend=3.0, + ), + _grouping_row(common_daily_activity_module._GROUP_DATE_API_KEY, api_key="", spend=3.0), + _grouping_row(common_daily_activity_module._GROUP_DATE_API_KEY, api_key="unattributed", spend=3.0), + _grouping_row(common_daily_activity_module._GROUP_DATE_API_KEY, api_key="raw:unattributed", spend=4.0), + ) + result: Final = common_daily_activity_module._aggregate_grouping_sets_records_sync( + records=records, api_key_metadata={} + ) + breakdown: Final = result["results"][0].breakdown + assert breakdown.models["unknown-model"].metrics.spend == 6.0 + assert {key: value.metrics.spend for key, value in breakdown.models["unknown-model"].api_key_breakdown.items()} == { + "unattributed": 3.0, + "raw:unattributed": 3.0, + } + assert {key: value.metrics.spend for key, value in breakdown.api_keys.items()} == { + "unattributed": 3.0, + "raw:unattributed": 3.0, + "raw:raw:unattributed": 4.0, + } + + +@pytest.mark.parametrize( + ("model", "group", "expected_group"), + [ + (None, None, "unknown-model-group"), + ("", "", "unknown-model-group"), + (None, "configured-group", "configured-group"), + ("provider-model", "", "provider-model"), + ("unknown-model", "", "unknown-model"), + (" ", "", " "), + ], +) +def test_model_group_uses_raw_model_before_unknown_group( + model: str | None, group: str | None, expected_group: str +) -> None: + result: Final = common_daily_activity_module._aggregate_spend_records_sync( + records=(_daily_user_spend_record(user_id="user", api_key="key", spend=1.0, model=model, model_group=group),), + api_key_metadata={}, + entity_id_field=None, + entity_metadata_field=None, + ) + assert tuple(result["results"][0].breakdown.model_groups) == (expected_group,) + + +@pytest.mark.parametrize( + ("raw_key", "bucket"), + [ + ("", "unattributed"), + ("real-hash", "real-hash"), + ("unattributed", "raw:unattributed"), + ("raw:unattributed", "raw:raw:unattributed"), + ("raw:real-hash", "raw:raw:real-hash"), + ], +) +def test_reporting_key_filters_round_trip_without_losing_scope(raw_key: str, bucket: str) -> None: + from litellm.repositories.daily_activity_buckets import api_key_bucket + + assert api_key_bucket(raw_key) == bucket + scope: Final = daily_activity_scope( + "litellm_dailyuserspend", "user_id", "user", None, bucket, "2024-01-01", "2024-01-01", None, None + ) + assert scope.api_keys == (raw_key,) + assert scope.entity_ids == ("user",) + + +@pytest.mark.parametrize(("raw_key", "bucket_key"), [("", "unattributed"), ("unattributed", "raw:unattributed")]) +def test_mcp_nested_key_metrics_use_reporting_identity(raw_key: str, bucket_key: str) -> None: + record: Final = SimpleNamespace( + **{ + **vars(_daily_user_spend_record(user_id="user", api_key=raw_key, spend=0.25, model=None, model_group=None)), + "mcp_namespaced_tool_name": "deleted-server/tool", + } + ) + result: Final = common_daily_activity_module._aggregate_spend_records_sync( + records=(record,), api_key_metadata={}, entity_id_field=None, entity_metadata_field=None + ) + breakdown: Final = result["results"][0].breakdown + assert breakdown.mcp_servers["deleted-server/tool"].api_key_breakdown == breakdown.api_keys + assert breakdown.mcp_servers["deleted-server/tool"].api_key_breakdown[bucket_key].metrics.spend == 0.25 + + +def test_ptu_rows_do_not_create_endpoint_or_mcp_key_buckets(ptu_cost_attribution_enabled: None) -> None: + from litellm.constants import PTU_SENTINEL_API_KEY + from litellm.types.proxy.management_endpoints.common_daily_activity import BreakdownMetrics + + record: Final = SimpleNamespace( + **{ + **vars(_spend_record(PTU_SENTINEL_API_KEY, ptu_flat_cost=2.0)), + "mcp_namespaced_tool_name": "server/tool", + "endpoint": "/mcp", + } + ) + breakdown: Final = common_daily_activity_module.update_breakdown_metrics(BreakdownMetrics(), record, {}, {}, {}) + assert breakdown.api_keys == {} + assert breakdown.endpoints["/mcp"].api_key_breakdown == {} + assert breakdown.mcp_servers["server/tool"].api_key_breakdown == {} + assert breakdown.endpoints["/mcp"].metrics.flat_cost == 2.0 + assert breakdown.mcp_servers["server/tool"].metrics.flat_cost == 2.0 + + +def test_metric_merges_preserve_every_additive_field() -> None: + metrics: Final = SpendMetrics.model_validate(dict.fromkeys(SpendMetrics.model_fields, 2)) + merged: Final = common_daily_activity_module.sum_spend_metrics((SpendMetrics(), metrics, metrics)) + assert merged.model_dump() == dict.fromkeys(SpendMetrics.model_fields, 4) diff --git a/tests/unit/proxy/management_endpoints/test_daily_activity_routes.py b/tests/unit/proxy/management_endpoints/test_daily_activity_routes.py index 2a273034a97..e487c00f25d 100644 --- a/tests/unit/proxy/management_endpoints/test_daily_activity_routes.py +++ b/tests/unit/proxy/management_endpoints/test_daily_activity_routes.py @@ -628,7 +628,9 @@ def test_aggregated_routes_return_scoped_results( assert len(body["results"]) == 2, response.text -def _ptu_sized_router(model_group: str, shares: Mapping[str, int] = MappingProxyType({"team-a": 30, "team-b": 20})) -> Router: +def _ptu_sized_router( + model_group: str, shares: Mapping[str, int] = MappingProxyType({"team-a": 30, "team-b": 20}) +) -> Router: return Router( model_list=[ { @@ -1268,13 +1270,15 @@ def test_team_scope_applies_membership_and_user_key_filter( assert repository.aggregated.call_args.kwargs["include_entity_breakdown"] is True +@pytest.mark.parametrize("api_key", ["key-beta", "unattributed", "raw:unattributed", "raw:raw:unattributed"]) def test_team_scope_does_not_allow_an_unowned_api_key( daily_activity_client: tuple[TestClient, _FakeRepository], + api_key: str, ) -> None: client, _ = daily_activity_client response: Final = client.get( "/team/daily/activity/aggregated", - params={**_entity_params("team_ids", "team-a"), "api_key": "key-beta"}, + params={**_entity_params("team_ids", "team-a"), "api_key": api_key}, headers={"x-user-role": LitellmUserRoles.INTERNAL_USER.value, "x-user-id": "user-a"}, ) assert response.status_code == 200, response.text @@ -1622,3 +1626,17 @@ def test_export_content_disposition_is_ascii_and_built_from_canonical_dates( disposition: Final = response.headers["content-disposition"] assert disposition == 'attachment; filename="team-usage-2025-01-01-2025-01-02-daily.csv"' assert disposition.isascii() + + +@pytest.mark.parametrize(("prefix", "query_name", "entity_id"), _ENTITY_CASES) +def test_unattributed_key_filter_does_not_become_an_unfiltered_request( + daily_activity_client: tuple[TestClient, _FakeRepository], prefix: str, query_name: str, entity_id: str +) -> None: + client, repository = daily_activity_client + response: Final = client.get( + f"{prefix}/daily/activity/aggregated", + params={**_entity_params(query_name, entity_id), "api_key": "unattributed"}, + ) + assert response.status_code == 200, response.text + assert response.json()["metadata"]["total_spend"] == 0 + assert repository.aggregated.call_args.args[0].api_keys == ("",) diff --git a/tests/unit/proxy/management_endpoints/test_tag_management_endpoints.py b/tests/unit/proxy/management_endpoints/test_tag_management_endpoints.py index e5d13a381ae..a2e509b6eb8 100644 --- a/tests/unit/proxy/management_endpoints/test_tag_management_endpoints.py +++ b/tests/unit/proxy/management_endpoints/test_tag_management_endpoints.py @@ -1058,7 +1058,8 @@ async def test_internal_user_tag_daily_activity_is_scoped_to_their_keys(): @pytest.mark.asyncio -async def test_internal_user_tag_daily_activity_rejects_unowned_api_key_filter(): +@pytest.mark.parametrize("api_key", ["unowned-key", "unattributed", "raw:unattributed", "raw:raw:unattributed"]) +async def test_internal_user_tag_daily_activity_rejects_unowned_api_key_filter(api_key: str): """ If an internal user filters tag usage by an API key they do not own, the endpoint should return an empty scoped filter instead of exposing that key's @@ -1092,7 +1093,7 @@ async def test_internal_user_tag_daily_activity_rejects_unowned_api_key_filter() result = await get_tag_daily_activity( start_date="2025-01-01", end_date="2025-01-31", - api_key="unowned-key", + api_key=api_key, user_api_key_dict=mock_user_auth, ) diff --git a/tests/unit/repositories/test_daily_activity_sql.py b/tests/unit/repositories/test_daily_activity_sql.py index c57db08e976..f8375b93dad 100644 --- a/tests/unit/repositories/test_daily_activity_sql.py +++ b/tests/unit/repositories/test_daily_activity_sql.py @@ -123,7 +123,6 @@ def test_key_page_sql_orders_exact_spend_and_binds_scope_before_page() -> None: 7, ) assert "SUM(spend::numeric) AS rank_spend" in query.sql - assert "ORDER BY rank_spend DESC, api_key" in query.sql assert "(SELECT COUNT(*) FROM ranked)::bigint AS total_api_keys" in query.sql @@ -329,7 +328,6 @@ def test_entity_rollup_bounds_keys_and_reuses_scope_filters() -> None: assert "api_key = ANY($3::text[])" in query.sql assert query.sql.count("api_key = ANY($3::text[])") == 4 assert query.sql.count("api_key <> $4") == 2 - assert query.sql.count("ORDER BY SUM(spend::numeric) DESC, api_key") == 1 assert "k.entity_id = e.entity_id" in query.sql assert "LIMIT $5" in query.sql assert query.params == ("2026-01-01", "2026-01-31", ["key-1", "key-2"], PTU_SENTINEL_API_KEY, 3) @@ -351,13 +349,11 @@ def test_search_query_escapes_pattern_metacharacters_and_binds_limit() -> None: assert 'LEFT JOIN "LiteLLM_UserTable" u ON u.user_id = d.user_id' in query.sql assert "d.key_alias ILIKE $3 ESCAPE" in query.sql assert "d.user_id ILIKE $3 ESCAPE" in query.sql - assert "api_key ILIKE $3 ESCAPE" in query.sql assert "v.key_alias ILIKE $3 ESCAPE" in query.sql assert "v.user_id ILIKE $3 ESCAPE" in query.sql assert "u.user_email ILIKE $3 ESCAPE" in query.sql assert query.sql.count("ILIKE $3 ESCAPE") == 7 assert "api_key <> $4" in query.sql - assert "ORDER BY SUM(spend::numeric) DESC, api_key" in query.sql assert "LIMIT $5" in query.sql assert query.params == ( "2026-01-01", @@ -374,12 +370,9 @@ def test_model_and_cache_key_queries_bind_filters_sentinel_and_limits() -> None: ) leakage_query = build_cache_leakage_keys_sql(_scope(entity_ids=None), limit=20) - assert "COALESCE(NULLIF(model_group, ''), model) = $3" in model_query.sql assert "api_key <> $4" in model_query.sql - assert "ORDER BY SUM(spend::numeric) DESC, api_key" in model_query.sql assert model_query.params == ("2026-01-01", "2026-01-31", "public-model", PTU_SENTINEL_API_KEY, 5) assert "HAVING SUM(prompt_tokens) - SUM(cache_read_input_tokens) > 0" in leakage_query.sql - assert "ORDER BY SUM(prompt_tokens) - SUM(cache_read_input_tokens) DESC, api_key" in leakage_query.sql assert leakage_query.params == ("2026-01-01", "2026-01-31", PTU_SENTINEL_API_KEY, 20) @@ -401,8 +394,14 @@ def test_query_builders_reject_nonpositive_limits(builder) -> None: ("export_type", "group_key", "key_filter", "joins"), [ (ExportType.DAILY, "''", "", ""), - (ExportType.DAILY_WITH_KEYS, "scoped.api_key", "api_key <> $3", 'LEFT JOIN "LiteLLM_VerificationToken"'), - (ExportType.DAILY_WITH_MODELS, "COALESCE(scoped.model, '')", "api_key <> $3", ""), + ( + ExportType.DAILY_WITH_KEYS, + "CASE WHEN api_key = '' THEN 'unattributed' " + "WHEN api_key = 'unattributed' OR api_key LIKE 'raw:%' THEN 'raw:' || api_key ELSE api_key END", + "api_key <> $3", + 'LEFT JOIN "LiteLLM_VerificationToken"', + ), + (ExportType.DAILY_WITH_MODELS, "COALESCE(NULLIF(model, ''), 'unknown-model')", "api_key <> $3", ""), ( ExportType.DAILY_WITH_USERS, "COALESCE(vt.user_id, dvt.user_id, '')", diff --git a/ui/litellm-dashboard/src/components/UsagePage/components/EntityUsage/TopKeyView.test.tsx b/ui/litellm-dashboard/src/components/UsagePage/components/EntityUsage/TopKeyView.test.tsx index a9c87471511..2b4b69ffe7b 100644 --- a/ui/litellm-dashboard/src/components/UsagePage/components/EntityUsage/TopKeyView.test.tsx +++ b/ui/litellm-dashboard/src/components/UsagePage/components/EntityUsage/TopKeyView.test.tsx @@ -493,6 +493,26 @@ describe("TopKeyView", () => { }); }); + it.each([ + ["raw:unattributed", "unattributed"], + ["raw:raw:unattributed", "raw:unattributed"], + ["raw:raw:raw:key", "raw:raw:key"], + ])("opens the stored key behind reporting ID %s", async (bucket, key) => { + mockKeyInfoV1Call.mockResolvedValue({ key }); + mockTransformKeyInfo.mockReturnValue({ token: key } as KeyResponse); + render( + , + ); + + fireEvent.click(screen.getByRole("button", { name: bucket })); + + expect(await screen.findByText(`Key Info View for ${key}`)).toBeInTheDocument(); + expect(mockKeyInfoV1Call).toHaveBeenCalledWith("test-token", key); + }); + it("should only look up keys that still exist in the database, from both the table and the chart", async () => { mockKeyInfoV1Call.mockResolvedValue({ key: "info" }); mockTransformKeyInfo.mockReturnValue({ transformed: "data" } as unknown as KeyResponse); @@ -503,6 +523,7 @@ describe("TopKeyView", () => { {...baseProps} topKeys={[ { api_key: "session-key", key_alias: null, user: "alice@example.com", key_exists: false, spend: 100 }, + { api_key: "unattributed", key_alias: "Unattributed", user: null, key_exists: false, spend: 25 }, { api_key: "stored-key", key_alias: "Stored", user: null, key_exists: true, spend: 50 }, ]} />, @@ -510,11 +531,14 @@ describe("TopKeyView", () => { expect(screen.getByRole("button", { name: "stored-key" })).toBeInTheDocument(); expect(screen.queryByRole("button", { name: "session-key" })).not.toBeInTheDocument(); + expect(screen.getByText("Unattributed")).toBeInTheDocument(); + expect(screen.queryByRole("button", { name: "unattributed" })).not.toBeInTheDocument(); + await user.click(screen.getByText("unattributed")); await user.click(screen.getByText("session-key")); await user.click(screen.getByRole("radio", { name: "Chart View" })); const bars = chartBars(container); - expect(bars).toHaveLength(2); + expect(bars).toHaveLength(3); bars.forEach((bar) => fireEvent.click(bar)); expect(await screen.findByText("Key Info View for stored-key")).toBeInTheDocument(); diff --git a/ui/litellm-dashboard/src/components/UsagePage/components/EntityUsage/TopKeyView.tsx b/ui/litellm-dashboard/src/components/UsagePage/components/EntityUsage/TopKeyView.tsx index 90a26919e01..291bcb34c73 100644 --- a/ui/litellm-dashboard/src/components/UsagePage/components/EntityUsage/TopKeyView.tsx +++ b/ui/litellm-dashboard/src/components/UsagePage/components/EntityUsage/TopKeyView.tsx @@ -11,6 +11,7 @@ import { transformKeyInfo } from "../../../key_team_helpers/transform_key_info"; import { keyInfoV1Call } from "../../../networking"; import KeyInfoView from "../../../templates/key_info_view"; import { TagUsage } from "../../types"; +import { storedKeyId } from "../../keyActivityData"; const TOP_KEYS_LIMITS = [5, 10, 25, 50] as const; const LIMIT_OPTIONS = TOP_KEYS_LIMITS.map((limit) => ({ value: String(limit), label: String(limit) })); @@ -70,11 +71,12 @@ const TopKeyView: React.FC = ({ topKeys, teams, showTags = fals if (!accessToken || !canOpenKeyInfo(item)) return; try { - const keyInfo = await keyInfoV1Call(accessToken, item.api_key); + const keyId = storedKeyId(item.api_key); + const keyInfo = await keyInfoV1Call(accessToken, keyId); const transformedKeyData = transformKeyInfo(keyInfo); setKeyData(transformedKeyData); - setSelectedKey(item.api_key); + setSelectedKey(keyId); setIsModalOpen(true); // Open modal when key is clicked } catch (error) { console.error("Error fetching key info:", error); diff --git a/ui/litellm-dashboard/src/components/UsagePage/keyActivityData.test.ts b/ui/litellm-dashboard/src/components/UsagePage/keyActivityData.test.ts index b35ea72896a..a5db5db14cc 100644 --- a/ui/litellm-dashboard/src/components/UsagePage/keyActivityData.test.ts +++ b/ui/litellm-dashboard/src/components/UsagePage/keyActivityData.test.ts @@ -7,7 +7,13 @@ import { type KeySpendActivityRow, toDailyData, } from "./dailyActivityApi"; -import { keyDetailFromResponse, mergeKeyActivityPages, overallUsageMetrics } from "./keyActivityData"; +import { + keyDetailFromResponse, + mergeKeyActivityPages, + overallUsageMetrics, + reportingKeyId, + storedKeyId, +} from "./keyActivityData"; const completeMetrics: components["schemas"]["SpendMetrics"] = { api_requests: 2, @@ -80,6 +86,50 @@ const pageRow = (api_key: string): KeySpendActivityRow => ({ }); describe("key activity data", () => { + it.each([ + ["", "unattributed"], + ["unattributed", "raw:unattributed"], + ["raw:key", "raw:raw:key"], + ["raw:raw:key", "raw:raw:raw:key"], + ["ordinary-hash", "ordinary-hash"], + [" ", " "], + ])("round trips stored key %s through reporting", (key, bucket) => { + expect(reportingKeyId(key)).toBe(bucket); + expect(storedKeyId(bucket)).toBe(key); + }); + + it("retains unattributed totals and unknown models in key details", () => { + const unattributed = { + metrics: completeMetrics, + metadata: { key_alias: "Unattributed", key_exists: false }, + }; + const response: DailyActivityAggregatedResponse = { + ...aggregatedResponse, + results: [ + { + date: "2026-09-27", + metrics: completeMetrics, + breakdown: { + api_keys: { unattributed }, + models: { + "unknown-model": { + metrics: completeMetrics, + metadata: {}, + api_key_breakdown: { unattributed }, + }, + }, + }, + }, + ], + }; + + const detail = keyDetailFromResponse(response, "unattributed", []); + expect(detail?.total_requests).toBe(2); + expect(detail?.total_spend).toBe(1.25); + expect(detail?.top_models.map((model) => model.model)).toEqual(["unknown-model"]); + expect(detail?.daily_data[0].metrics.api_requests).toBe(2); + }); + it("builds overall totals from metadata and daily data from complete daily metrics", () => { const summary = overallUsageMetrics( toDailyData(aggregatedResponse), diff --git a/ui/litellm-dashboard/src/components/UsagePage/keyActivityData.ts b/ui/litellm-dashboard/src/components/UsagePage/keyActivityData.ts index 4b2e8c573f2..2d2977058a5 100644 --- a/ui/litellm-dashboard/src/components/UsagePage/keyActivityData.ts +++ b/ui/litellm-dashboard/src/components/UsagePage/keyActivityData.ts @@ -4,6 +4,16 @@ import type { DailyActivityAggregatedResponse, DailyActivityMetadata, KeySpendAc import { toDailyData } from "./dailyActivityApi"; import type { ModelActivityData } from "./types"; +export const reportingKeyId = (key: string): string => { + if (!key) return "unattributed"; + return key === "unattributed" || key.startsWith("raw:") ? `raw:${key}` : key; +}; + +export const storedKeyId = (bucket: string): string => { + if (bucket === "unattributed") return ""; + return bucket.startsWith("raw:") ? bucket.slice(4) : bucket; +}; + export const overallUsageMetrics = ( results: ReturnType, metadata: DailyActivityMetadata, diff --git a/ui/litellm-dashboard/src/components/templates/KeySavingsTab.integration.test.tsx b/ui/litellm-dashboard/src/components/templates/KeySavingsTab.integration.test.tsx index 3567388de05..40ab0e8e77c 100644 --- a/ui/litellm-dashboard/src/components/templates/KeySavingsTab.integration.test.tsx +++ b/ui/litellm-dashboard/src/components/templates/KeySavingsTab.integration.test.tsx @@ -131,6 +131,18 @@ describe("KeySavingsTab", () => { expect(screen.getByTestId("key-savings-scope-note")).toHaveTextContent("Showing your own requests"); }); + it.each([ + ["unattributed", "raw:unattributed"], + ["raw:unattributed", "raw:raw:unattributed"], + ["raw:raw:key", "raw:raw:raw:key"], + ])("requests savings for the real key %s", (keyToken, bucket) => { + const hook = scopedRange().mockReturnValue(mockActivity()); + + renderTab({ keyToken, userRole: "Admin" }); + + expect(hook).toHaveBeenCalledWith("test-token", { userId: null, apiKey: bucket }, activity); + }); + it("reads the whole key for a proxy admin, with no scope note to contradict it", () => { const hook = scopedRange().mockReturnValue(mockActivity()); diff --git a/ui/litellm-dashboard/src/components/templates/KeySavingsTab.tsx b/ui/litellm-dashboard/src/components/templates/KeySavingsTab.tsx index d1395e0be51..03aced7358e 100644 --- a/ui/litellm-dashboard/src/components/templates/KeySavingsTab.tsx +++ b/ui/litellm-dashboard/src/components/templates/KeySavingsTab.tsx @@ -1,6 +1,7 @@ "use client"; import ScopedSavingsTab from "@/components/shared/ScopedSavingsTab"; +import { reportingKeyId } from "@/components/UsagePage/keyActivityData"; import { hasProxyWideSpendView, spendScopeUserId } from "@/utils/roles"; import type { ActivityDateRange } from "@/app/(dashboard)/cost-optimization/_components/useDailyActivityRange"; @@ -15,7 +16,7 @@ interface KeySavingsTabProps { const KeySavingsTab = ({ accessToken, keyToken, userId, userRole, activity }: KeySavingsTabProps) => (