diff --git a/basedpyright-code-budget.json b/basedpyright-code-budget.json index 962d1266fd7..ef88ae574fb 100644 --- a/basedpyright-code-budget.json +++ b/basedpyright-code-budget.json @@ -1,6 +1,6 @@ { "reportAny": { - "limit": 17271 + "limit": 17270 }, "reportArgumentType": { "limit": 2539 diff --git a/enterprise/litellm_enterprise/proxy/hooks/managed_files.py b/enterprise/litellm_enterprise/proxy/hooks/managed_files.py index cf2cee9b6ef..570b306d6df 100644 --- a/enterprise/litellm_enterprise/proxy/hooks/managed_files.py +++ b/enterprise/litellm_enterprise/proxy/hooks/managed_files.py @@ -36,6 +36,7 @@ from litellm.llms.base_llm.managed_resources.isolation import ( build_list_page, build_owner_filter, can_access_resource, + resolve_resource_owner_id, ) from litellm.proxy._types import ( CallTypes, @@ -222,7 +223,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): file_object=file_object, model_mappings=model_mappings, flat_model_file_ids=list(model_mappings.values()), - created_by=user_api_key_dict.user_id, + created_by=resolve_resource_owner_id(user_api_key_dict), team_id=user_api_key_dict.team_id, updated_by=user_api_key_dict.user_id, ) @@ -238,7 +239,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): "unified_file_id": file_id, "model_mappings": json.dumps(model_mappings), "flat_model_file_ids": list(model_mappings.values()), - "created_by": user_api_key_dict.user_id, + "created_by": resolve_resource_owner_id(user_api_key_dict), "team_id": user_api_key_dict.team_id, "updated_by": user_api_key_dict.user_id, } @@ -342,7 +343,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): "file_object": file_object.model_dump_json(), "model_object_id": model_object_id, "file_purpose": file_purpose, - "created_by": user_api_key_dict.user_id, + "created_by": resolve_resource_owner_id(user_api_key_dict), "team_id": user_api_key_dict.team_id, "updated_by": user_api_key_dict.user_id, "status": file_object.status, diff --git a/enterprise/pyproject.toml b/enterprise/pyproject.toml index 3653aba67ef..cac98b69793 100644 --- a/enterprise/pyproject.toml +++ b/enterprise/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "litellm-enterprise" -version = "0.1.61" +version = "0.1.62" description = "Package for LiteLLM Enterprise features" readme = "README.md" requires-python = ">=3.9" @@ -26,7 +26,7 @@ required-version = ">=0.10.9" module-root = "" [tool.commitizen] -version = "0.1.61" +version = "0.1.62" version_files = [ "pyproject.toml:^version", "../pyproject.toml:litellm-enterprise==", diff --git a/litellm-proxy-extras/pyproject.toml b/litellm-proxy-extras/pyproject.toml index 0ef5cd1e856..d5741d479bf 100644 --- a/litellm-proxy-extras/pyproject.toml +++ b/litellm-proxy-extras/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "litellm-proxy-extras" -version = "0.4.90" +version = "0.4.91" description = "Additional files for the LiteLLM Proxy. Reduces the size of the main litellm package." readme = "README.md" requires-python = ">=3.9" @@ -26,7 +26,7 @@ required-version = ">=0.10.9" module-root = "" [tool.commitizen] -version = "0.4.90" +version = "0.4.91" version_files = [ "pyproject.toml:^version", "../pyproject.toml:litellm-proxy-extras==", diff --git a/litellm-rust/Cargo.lock b/litellm-rust/Cargo.lock index ce28f737334..4388e561026 100644 --- a/litellm-rust/Cargo.lock +++ b/litellm-rust/Cargo.lock @@ -2,6 +2,36 @@ # It is not intended for manual editing. version = 4 +[[package]] +name = "aho-corasick" +version = "1.1.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c982642fa9e8606056828ee9a8505737230110bb1099153c79efe865c59d12ba" +dependencies = [ + "memchr", +] + +[[package]] +name = "alloca" +version = "0.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e5a7d05ea6aea7e9e64d25b9156ba2fee3fdd659e34e41063cd2fc7cd020d7f4" +dependencies = [ + "cc", +] + +[[package]] +name = "anes" +version = "0.1.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4b46cbb362ab8752921c97e041f5e366ee6297bd428a31275b9fcf1e380f7299" + +[[package]] +name = "anstyle" +version = "1.0.14" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "940b3a0ca603d1eade50a4846a2afffd5ef57a9feac2c0e2ec2e14f9ead76000" + [[package]] name = "arc-swap" version = "1.9.2" @@ -506,6 +536,12 @@ dependencies = [ "either", ] +[[package]] +name = "cast" +version = "0.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "37b2a672a2cb129a2e41c10b1224bb368f9f37a2b16b612598138befd7b37eb5" + [[package]] name = "cc" version = "1.3.0" @@ -541,6 +577,58 @@ dependencies = [ "rand_core 0.10.1", ] +[[package]] +name = "ciborium" +version = "0.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "42e69ffd6f0917f5c029256a24d0161db17cea3997d185db0d35926308770f0e" +dependencies = [ + "ciborium-io", + "ciborium-ll", + "serde", +] + +[[package]] +name = "ciborium-io" +version = "0.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "05afea1e0a06c9be33d539b876f1ce3692f4afea2cb41f740e7743225ed1c757" + +[[package]] +name = "ciborium-ll" +version = "0.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "57663b653d948a338bfb3eeba9bb2fd5fcfaecb9e199e87e1eda4d9e8b240fd9" +dependencies = [ + "ciborium-io", + "half", +] + +[[package]] +name = "clap" +version = "4.6.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "473c7e07f409a8d772161724aa8db6a765a2532a70f9667eeb7b49d3d02fbdca" +dependencies = [ + "clap_builder", +] + +[[package]] +name = "clap_builder" +version = "4.6.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7b48fea5a88e9ae728a2dcbedbfc0e730f7d60da42e1cb049a83c9fb8b789889" +dependencies = [ + "anstyle", + "clap_lex", +] + +[[package]] +name = "clap_lex" +version = "1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8d4a3bb8b1e0c1050499d1815f5ab16d04f0959b233085fb31653fbfc9d98f9" + [[package]] name = "cmake" version = "0.1.58" @@ -596,6 +684,72 @@ dependencies = [ "libc", ] +[[package]] +name = "criterion" +version = "0.8.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "950046b2aa2492f9a536f5f4f9a3de7b9e2476e575e05bd6c333371add4d98f3" +dependencies = [ + "alloca", + "anes", + "cast", + "ciborium", + "clap", + "criterion-plot", + "itertools", + "num-traits", + "oorandom", + "page_size", + "plotters", + "rayon", + "regex", + "serde", + "serde_json", + "tinytemplate", + "walkdir", +] + +[[package]] +name = "criterion-plot" +version = "0.8.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d8d80a2f4f5b554395e47b5d8305bc3d27813bacb73493eb1001e8f76dae29ea" +dependencies = [ + "cast", + "itertools", +] + +[[package]] +name = "crossbeam-deque" +version = "0.8.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5181e0de7b61eb03a81e347d6dd8797bae9da5146707b51077e2d71a54ec0ceb" +dependencies = [ + "crossbeam-epoch", + "crossbeam-utils", +] + +[[package]] +name = "crossbeam-epoch" +version = "0.9.20" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2d6914041f254d6e9176c01941b21115dcfb7089e55135a35411081bd106ef3f" +dependencies = [ + "crossbeam-utils", +] + +[[package]] +name = "crossbeam-utils" +version = "0.8.22" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "61803da095bee82a81bb1a452ecc25d3b2f1416d1897eb86430c6159ef717c17" + +[[package]] +name = "crunchy" +version = "0.2.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "460fbee9c2c2f33933d720630a6a0bac33ba7053db5344fac858d4b8952d77d5" + [[package]] name = "crypto-common" version = "0.1.7" @@ -856,6 +1010,17 @@ dependencies = [ "tracing", ] +[[package]] +name = "half" +version = "2.7.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" +dependencies = [ + "cfg-if", + "crunchy", + "zerocopy", +] + [[package]] name = "hashbrown" version = "0.17.1" @@ -1179,6 +1344,15 @@ version = "2.12.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "d98f6fed1fde3f8c21bc40a1abb88dd75e67924f9cffc3ef95607bad8017f8e2" +[[package]] +name = "itertools" +version = "0.13.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "413ee7dfc52ee1a4949ceeb7dbc8a33f2d6c088194d9f922fb8318faf1f01186" +dependencies = [ + "either", +] + [[package]] name = "itoa" version = "1.0.18" @@ -1255,10 +1429,13 @@ dependencies = [ name = "litellm-python-bridge" version = "0.1.0" dependencies = [ + "criterion", "litellm-ai-gateway", "litellm-core", "pyo3", "pyo3-async-runtimes", + "pythonize", + "serde", "serde_json", "tokio", ] @@ -1340,6 +1517,12 @@ version = "1.21.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9f7c3e4beb33f85d45ae3e3a1792185706c8e16d043238c593331cc7cd313b50" +[[package]] +name = "oorandom" +version = "11.1.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d6790f58c7ff633d8771f42965289203411a5e5c68388703c06e14f24770b41e" + [[package]] name = "openssl-probe" version = "0.2.1" @@ -1352,6 +1535,16 @@ version = "0.5.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "1a80800c0488c3a21695ea981a54918fbb37abf04f4d0720c453632255e2ff0e" +[[package]] +name = "page_size" +version = "0.6.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "30d5b2194ed13191c1999ae0704b7839fb18384fa22e49b57eeaa97d79ce40da" +dependencies = [ + "libc", + "winapi", +] + [[package]] name = "percent-encoding" version = "2.3.2" @@ -1376,6 +1569,34 @@ version = "0.3.33" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "19f132c84eca552bf34cab8ec81f1c1dcc229b811638f9d283dceabe58c5569e" +[[package]] +name = "plotters" +version = "0.3.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5aeb6f403d7a4911efb1e33402027fc44f29b5bf6def3effcc22d7bb75f2b747" +dependencies = [ + "num-traits", + "plotters-backend", + "plotters-svg", + "wasm-bindgen", + "web-sys", +] + +[[package]] +name = "plotters-backend" +version = "0.3.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "df42e13c12958a16b3f7f4386b9ab1f3e7933914ecea48da7139435263a4172a" + +[[package]] +name = "plotters-svg" +version = "0.3.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "51bae2ac328883f7acdfea3d66a7c35751187f870bc81f94563733a154d7a670" +dependencies = [ + "plotters-backend", +] + [[package]] name = "portable-atomic" version = "1.14.0" @@ -1486,6 +1707,16 @@ dependencies = [ "syn 2.0.119", ] +[[package]] +name = "pythonize" +version = "0.29.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6ec376e1216e0c929a74964ce2020012a1a39f32d80e78aa688721219ea7fb89" +dependencies = [ + "pyo3", + "serde", +] + [[package]] name = "quinn" version = "0.11.11" @@ -1613,12 +1844,61 @@ dependencies = [ "rand_core 0.10.1", ] +[[package]] +name = "rayon" +version = "1.12.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fb39b166781f92d482534ef4b4b1b2568f42613b53e5b6c160e24cfbfa30926d" +dependencies = [ + "either", + "rayon-core", +] + +[[package]] +name = "rayon-core" +version = "1.13.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "22e18b0f0062d30d4230b2e85ff77fdfe4326feb054b9783a3460d8435c8ab91" +dependencies = [ + "crossbeam-deque", + "crossbeam-utils", +] + +[[package]] +name = "regex" +version = "1.13.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f020237b6c8eed93db2e2cb53c00c60a8e1bc73da7d073199a1180401450218d" +dependencies = [ + "aho-corasick", + "memchr", + "regex-automata", + "regex-syntax", +] + +[[package]] +name = "regex-automata" +version = "0.4.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ad8553b9b26413251cbf30e620595c7a41b3887f03da04579c0e6b0d6a06b4b2" +dependencies = [ + "aho-corasick", + "memchr", + "regex-syntax", +] + [[package]] name = "regex-lite" version = "0.1.9" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "cab834c73d247e67f4fae452806d17d3c7501756d98c8808d7c9c7aa7d18f973" +[[package]] +name = "regex-syntax" +version = "0.8.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d6f6ff9a378485b298a5286656da665ba74413d36db0979633275d2e708145d4" + [[package]] name = "reqwest" version = "0.12.28" @@ -1774,6 +2054,15 @@ version = "1.0.23" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9774ba4a74de5f7b1c1451ed6cd5285a32eddb5cccb8cc655a4e50009e06477f" +[[package]] +name = "same-file" +version = "1.0.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "93fc1dc3aaa9bfed95e02e6eadabb4baf7e3078b0bd1b4d7b6b0b68378900502" +dependencies = [ + "winapi-util", +] + [[package]] name = "schannel" version = "0.1.29" @@ -2099,6 +2388,16 @@ dependencies = [ "zerovec", ] +[[package]] +name = "tinytemplate" +version = "1.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "be4d6b5f19ff7664e8c98d03e2139cb510db9b0a60b55f8e8709b689d939b6bc" +dependencies = [ + "serde", + "serde_json", +] + [[package]] name = "tinyvec" version = "1.12.0" @@ -2363,6 +2662,16 @@ version = "0.8.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "5c3082ca00d5a5ef149bb8b555a72ae84c9c59f7250f013ac822ac2e49b19c64" +[[package]] +name = "walkdir" +version = "2.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "29790946404f91d9c5d06f9874efddea1dc06c5efe94541a7d6863108e3a5e4b" +dependencies = [ + "same-file", + "winapi-util", +] + [[package]] name = "want" version = "0.3.1" @@ -2475,6 +2784,37 @@ dependencies = [ "rustls-pki-types", ] +[[package]] +name = "winapi" +version = "0.3.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5c839a674fcd7a98952e593242ea400abe93992746761e38641405d28b00f419" +dependencies = [ + "winapi-i686-pc-windows-gnu", + "winapi-x86_64-pc-windows-gnu", +] + +[[package]] +name = "winapi-i686-pc-windows-gnu" +version = "0.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ac3b87c63620426dd9b991e5ce0329eff545bccbbb34f3be09ff6fb6ab51b7b6" + +[[package]] +name = "winapi-util" +version = "0.1.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c2a7b1c03c876122aa43f3020e6c3c3ee5c05081c9a00739faf7503aeba10d22" +dependencies = [ + "windows-sys 0.61.2", +] + +[[package]] +name = "winapi-x86_64-pc-windows-gnu" +version = "0.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "712e227841d057c1ee1cd2fb22fa7e5a5461ae8e48fa2ca79ec42cfc1931183f" + [[package]] name = "windows-link" version = "0.2.1" diff --git a/litellm-rust/Cargo.toml b/litellm-rust/Cargo.toml index 6d63be05d00..481ea3f8f66 100644 --- a/litellm-rust/Cargo.toml +++ b/litellm-rust/Cargo.toml @@ -18,6 +18,7 @@ litellm-ai-gateway = { path = "crates/ai-gateway", default-features = false } axum = "0.7" pyo3 = "0.29.0" pyo3-async-runtimes = { version = "0.29.0", features = ["tokio-runtime"] } +pythonize = "0.29.0" rand = "0.8" reqwest = { version = "0.12", default-features = false, features = ["blocking", "json", "rustls-tls", "http2", "stream"] } serde = { version = "1.0", features = ["derive"] } diff --git a/litellm-rust/crates/python-bridge/Cargo.toml b/litellm-rust/crates/python-bridge/Cargo.toml index 20a9ba789ce..0c4a753f762 100644 --- a/litellm-rust/crates/python-bridge/Cargo.toml +++ b/litellm-rust/crates/python-bridge/Cargo.toml @@ -9,10 +9,23 @@ repository.workspace = true name = "_native" crate-type = ["cdylib"] +[features] +default = ["extension-module"] +extension-module = ["pyo3/extension-module"] + [dependencies] litellm-core = { workspace = true, features = ["bedrock-auth"] } litellm-ai-gateway = { workspace = true, default-features = false } -pyo3 = { workspace = true, features = ["extension-module"] } +pyo3.workspace = true pyo3-async-runtimes.workspace = true +pythonize.workspace = true +serde.workspace = true serde_json.workspace = true tokio.workspace = true + +[dev-dependencies] +criterion = "0.8.2" + +[[bench]] +name = "serialization" +harness = false diff --git a/litellm-rust/crates/python-bridge/benches/serialization.rs b/litellm-rust/crates/python-bridge/benches/serialization.rs new file mode 100644 index 00000000000..8a90cf667d0 --- /dev/null +++ b/litellm-rust/crates/python-bridge/benches/serialization.rs @@ -0,0 +1,103 @@ +use std::hint::black_box; +use std::time::Duration; + +use criterion::{BenchmarkId, Criterion, criterion_group, criterion_main}; +use pyo3::prelude::*; +use pyo3::types::PyDict; +use serde_json::{Value, json}; + +const PAYLOAD_SIZES: &[(&str, usize)] = &[ + ("1_KiB", 1024), + ("64_KiB", 64 * 1024), + ("1_MiB", 1024 * 1024), + ("4_MiB", 4 * 1024 * 1024), + ("16_MiB", 16 * 1024 * 1024), +]; + +fn former_json_roundtrip_from_py(py: Python<'_>, value: &Bound<'_, PyAny>) -> Value { + let json = py.import("json").expect("Python json module should import"); + let encoded: String = json + .call_method1("dumps", (value,)) + .expect("payload should serialize") + .extract() + .expect("json.dumps should return a string"); + serde_json::from_str(&encoded).expect("serialized JSON should parse") +} + +fn pythonize_from_py(value: &Bound<'_, PyAny>) -> Value { + pythonize::depythonize(value).expect("payload should depythonize") +} + +fn former_json_roundtrip_to_py(py: Python<'_>, value: &Value) -> Py { + let json = py.import("json").expect("Python json module should import"); + let encoded = serde_json::to_string(value).expect("response should serialize"); + json.call_method1("loads", (encoded,)) + .expect("serialized response should parse in Python") + .unbind() +} + +fn pythonize_to_py(py: Python<'_>, value: &Value) -> Py { + pythonize::pythonize(py, value) + .expect("response should pythonize") + .unbind() +} + +fn serialization(c: &mut Criterion) { + Python::initialize(); + Python::attach(|py| { + for &(label, payload_bytes) in PAYLOAD_SIZES { + let data_uri = format!("data:image/png;base64,{}", "A".repeat(payload_bytes)); + let document = PyDict::new(py); + document + .set_item("type", "image_url") + .expect("document type should be set"); + document + .set_item("image_url", &data_uri) + .expect("document URL should be set"); + let response = json!({ + "pages": [{ + "index": 0, + "markdown": "OCR text", + "images": [{"image_base64": data_uri}], + }], + "model": "mistral-ocr-latest", + "document_annotation": null, + "usage_info": {"pages_processed": 1}, + "object": "ocr", + }); + + c.bench_with_input( + BenchmarkId::new("python_to_rust_json", label), + &document, + |b, document| { + b.iter(|| former_json_roundtrip_from_py(py, black_box(document.as_any()))) + }, + ); + c.bench_with_input( + BenchmarkId::new("python_to_rust_pythonize", label), + &document, + |b, document| b.iter(|| pythonize_from_py(black_box(document.as_any()))), + ); + c.bench_with_input( + BenchmarkId::new("rust_to_python_json", label), + &response, + |b, response| b.iter(|| former_json_roundtrip_to_py(py, black_box(response))), + ); + c.bench_with_input( + BenchmarkId::new("rust_to_python_pythonize", label), + &response, + |b, response| b.iter(|| pythonize_to_py(py, black_box(response))), + ); + } + }); +} + +criterion_group! { + name = benches; + config = Criterion::default() + .sample_size(20) + .warm_up_time(Duration::from_secs(1)) + .measurement_time(Duration::from_secs(4)); + targets = serialization +} +criterion_main!(benches); diff --git a/litellm-rust/crates/python-bridge/src/lib.rs b/litellm-rust/crates/python-bridge/src/lib.rs index c6f81cf6916..f9e75f45f75 100644 --- a/litellm-rust/crates/python-bridge/src/lib.rs +++ b/litellm-rust/crates/python-bridge/src/lib.rs @@ -19,6 +19,9 @@ use pyo3::types::{PyAny, PyDict}; use serde_json::{Map, Value}; mod gil; +mod marshal; + +use marshal::{from_py, to_py}; pyo3::create_exception!( _native, @@ -41,35 +44,18 @@ type MarshaledOcrInputs = ( Option, ); -fn py_to_json(py: Python<'_>, value: &Bound<'_, PyAny>) -> PyResult { - let json = py.import("json")?; - let encoded: String = json.call_method1("dumps", (value,))?.extract()?; - serde_json::from_str(&encoded).map_err(|err| PyValueError::new_err(err.to_string())) -} - -fn json_to_py(py: Python<'_>, value: Value) -> PyResult> { - let json = py.import("json")?; - let encoded = - serde_json::to_string(&value).map_err(|err| PyValueError::new_err(err.to_string()))?; - Ok(json.call_method1("loads", (encoded,))?.unbind()) -} - fn messages_response_to_py( py: Python<'_>, response: AnthropicMessagesResponse, ) -> PyResult> { - let value = - serde_json::to_value(response).map_err(|err| PyValueError::new_err(err.to_string()))?; - json_to_py(py, value) + to_py(py, &response) } fn chat_completions_response_to_py( py: Python<'_>, response: ChatCompletionsResponse, ) -> PyResult> { - let value = - serde_json::to_value(response).map_err(|err| PyValueError::new_err(err.to_string()))?; - json_to_py(py, value) + to_py(py, &response) } fn core_error_to_pyerr(err: CoreError) -> PyErr { @@ -116,7 +102,7 @@ fn optional_object_to_map( value: Option>, ) -> PyResult> { match value { - Some(value) => match py_to_json(py, value.bind(py))? { + Some(value) => match from_py(value.bind(py))? { Value::Object(map) => Ok(map), _ => Err(PyValueError::new_err(format!("{name} must be a dict"))), }, @@ -139,7 +125,7 @@ fn marshal_headers( headers: Option>, ) -> PyResult> { let value = match headers { - Some(headers) => py_to_json(py, headers.bind(py))?, + Some(headers) => from_py(headers.bind(py))?, None => Value::Object(Map::new()), }; let Value::Object(headers) = value else { @@ -211,7 +197,7 @@ fn marshal_inputs( optional_params: Option>, timeout_seconds: Option, ) -> PyResult { - let document = py_to_json(py, document.bind(py))?; + let document = from_py(document.bind(py))?; let extra_headers = match extra_headers { Some(headers) => Some(optional_object_to_map(py, "extra_headers", Some(headers))?), None => None, @@ -262,7 +248,7 @@ fn ocr( }); match result { - Ok(value) => json_to_py(py, value), + Ok(value) => to_py(py, &value), Err(err) => Err(core_error_to_pyerr(err)), } } @@ -307,7 +293,7 @@ fn aocr( .await .map_err(core_error_to_pyerr)?; - Python::attach(|py| json_to_py(py, value)) + Python::attach(|py| to_py(py, &value)) }) } @@ -325,7 +311,7 @@ fn transcription( optional_params: Option>, timeout_seconds: Option, ) -> PyResult> { - let audio = py_to_json(py, audio.bind(py))?; + let audio = from_py(audio.bind(py))?; let extra_headers = match extra_headers { Some(headers) => Some(optional_object_to_map(py, "extra_headers", Some(headers))?), None => None, @@ -351,7 +337,7 @@ fn transcription( )) }); match result { - Ok(value) => json_to_py(py, value), + Ok(value) => to_py(py, &value), Err(err) => Err(core_error_to_pyerr(err)), } } @@ -370,7 +356,7 @@ fn atranscription( optional_params: Option>, timeout_seconds: Option, ) -> PyResult> { - let audio = py_to_json(py, audio.bind(py))?; + let audio = from_py(audio.bind(py))?; let extra_headers = match extra_headers { Some(headers) => Some(optional_object_to_map(py, "extra_headers", Some(headers))?), None => None, @@ -394,7 +380,7 @@ fn atranscription( }) .await .map_err(core_error_to_pyerr)?; - Python::attach(|py| json_to_py(py, value)) + Python::attach(|py| to_py(py, &value)) }) } @@ -406,7 +392,7 @@ fn marshal_messages_inputs( extra_headers: Option>, timeout_seconds: Option, ) -> PyResult { - let body = py_to_json(py, body.bind(py))?; + let body: Value = from_py(body.bind(py))?; if !body.is_object() { return Err(PyValueError::new_err("body must be a dict")); } @@ -498,7 +484,7 @@ fn marshal_chat_completions_inputs( extra_headers: Option>, timeout_seconds: Option, ) -> PyResult { - let messages = py_to_json(py, messages.bind(py))?; + let messages: Value = from_py(messages.bind(py))?; if !messages.is_array() { return Err(PyValueError::new_err("messages must be a list")); } @@ -527,7 +513,7 @@ fn chat_completions_decline( optional_params: Option>, custom_llm_provider: Option, ) -> PyResult> { - let messages = py_to_json(py, messages.bind(py))?; + let messages = from_py(messages.bind(py))?; let optional_params = optional_object_to_map(py, "optional_params", optional_params)?; Ok(chat_completions_decline_reason( &model, diff --git a/litellm-rust/crates/python-bridge/src/marshal.rs b/litellm-rust/crates/python-bridge/src/marshal.rs new file mode 100644 index 00000000000..c3d0638427c --- /dev/null +++ b/litellm-rust/crates/python-bridge/src/marshal.rs @@ -0,0 +1,20 @@ +use pyo3::exceptions::PyValueError; +use pyo3::prelude::*; +use serde::Serialize; +use serde::de::DeserializeOwned; + +pub fn from_py(value: &Bound<'_, PyAny>) -> PyResult +where + T: DeserializeOwned, +{ + pythonize::depythonize(value).map_err(|error| PyValueError::new_err(error.to_string())) +} + +pub fn to_py(py: Python<'_>, value: &T) -> PyResult> +where + T: Serialize + ?Sized, +{ + pythonize::pythonize(py, value) + .map(Bound::unbind) + .map_err(|error| PyValueError::new_err(error.to_string())) +} diff --git a/litellm-rust/crates/python-bridge/tests/marshal_boundary.rs b/litellm-rust/crates/python-bridge/tests/marshal_boundary.rs new file mode 100644 index 00000000000..6a6ede22e85 --- /dev/null +++ b/litellm-rust/crates/python-bridge/tests/marshal_boundary.rs @@ -0,0 +1,52 @@ +use std::fs; +use std::path::{Path, PathBuf}; + +const DISALLOWED_OUTSIDE_MARSHAL: &[&str] = &[ + "py.import(\"json\")", + "pythonize::", + "serde_json::to_string", + "serde_json::from_str", +]; + +fn source_root() -> PathBuf { + Path::new(env!("CARGO_MANIFEST_DIR")).join("src") +} + +fn rust_sources(directory: &Path) -> Vec { + fs::read_dir(directory) + .expect("bridge source directory should be readable") + .map(|entry| { + entry + .expect("bridge source entry should be readable") + .path() + }) + .flat_map(|path| { + if path.is_dir() { + rust_sources(&path) + } else if path.extension().is_some_and(|extension| extension == "rs") { + vec![path] + } else { + Vec::new() + } + }) + .collect() +} + +#[test] +fn serialization_is_centralized_in_marshal_module() { + let root = source_root(); + + for path in rust_sources(&root) { + if path == root.join("marshal.rs") { + continue; + } + let source = fs::read_to_string(&path).expect("bridge source should be readable"); + for disallowed in DISALLOWED_OUTSIDE_MARSHAL { + assert!( + !source.contains(disallowed), + "{} bypasses the typed marshal module with `{disallowed}`", + path.display() + ); + } + } +} diff --git a/litellm/llms/base_llm/managed_resources/base_managed_resource.py b/litellm/llms/base_llm/managed_resources/base_managed_resource.py index 2a59eddf88a..cced330d873 100644 --- a/litellm/llms/base_llm/managed_resources/base_managed_resource.py +++ b/litellm/llms/base_llm/managed_resources/base_managed_resource.py @@ -12,6 +12,7 @@ from litellm.llms.base_llm.managed_resources.isolation import ( build_list_page, build_owner_filter, can_access_resource, + resolve_resource_owner_id, ) from litellm.proxy._types import UserAPIKeyAuth from litellm.types.utils import SpecialEnums @@ -157,7 +158,7 @@ class BaseManagedResource(ABC, Generic[ResourceObjectType]): "resource_object": resource_object, "model_mappings": model_mappings, "flat_model_resource_ids": list(model_mappings.values()), - "created_by": user_api_key_dict.user_id, + "created_by": resolve_resource_owner_id(user_api_key_dict), "team_id": user_api_key_dict.team_id, "updated_by": user_api_key_dict.user_id, } @@ -179,7 +180,7 @@ class BaseManagedResource(ABC, Generic[ResourceObjectType]): "unified_resource_id": unified_resource_id, "model_mappings": json.dumps(model_mappings), "flat_model_resource_ids": list(model_mappings.values()), - "created_by": user_api_key_dict.user_id, + "created_by": resolve_resource_owner_id(user_api_key_dict), "team_id": user_api_key_dict.team_id, "updated_by": user_api_key_dict.user_id, } diff --git a/litellm/llms/base_llm/managed_resources/isolation.py b/litellm/llms/base_llm/managed_resources/isolation.py index e1b204214d7..6a71e8e9223 100644 --- a/litellm/llms/base_llm/managed_resources/isolation.py +++ b/litellm/llms/base_llm/managed_resources/isolation.py @@ -3,10 +3,11 @@ Tenant-isolation helpers for managed file/batch/vector-store resources. Returns a Prisma filter and an ownership check that scope managed resources to the caller's identity: proxy admins see everything, user-keyed callers -see records they created, and service-account keys (no user_id) fall back -to the resource's owning team. Callers with no admin role and no -identifying ids are denied so an empty user_id can never select an -unscoped query. +see records they created, service-account keys (no user_id) fall back to +the resource's owning team, and keys with neither a user_id nor a team_id +fall back to their own hashed token so they can still reach the resources +they created. Callers with no admin role and no identifying ids at all +are denied so an empty user_id can never select an unscoped query. """ from typing import Any, Final @@ -19,6 +20,32 @@ from litellm.proxy._types import ( ) +def resolve_resource_owner_id( + user_api_key_dict: UserAPIKeyAuth, +) -> str | None: + """Return the identity to stamp on (and match against) a managed + resource's ``created_by``. + + A key with neither a user_id nor a team_id would otherwise stamp + ``created_by=None`` and be locked out of its own resources, so it owns + them under its hashed token instead, using the ``key:`` scope prefix + already used by ``proxy/common_utils/resource_ownership.py``. ``None`` + means the caller has no usable identity of its own and must fall back + to team scoping, or be denied. + """ + if user_api_key_dict.user_id is not None: + return user_api_key_dict.user_id + + if user_api_key_dict.team_id is not None: + return None + + token: Final = user_api_key_dict.token or user_api_key_dict.api_key + if token: + return f"key:{token}" + + return None + + def build_list_page(items: list[Any], has_more: bool = False) -> dict[str, Any]: """Build the OpenAI-style paginated list response shape used by managed file/batch/vector-store listings. ``first_id`` and ``last_id`` are @@ -39,7 +66,8 @@ def build_owner_filter( to records the caller is allowed to see. - ``{}`` means no scoping (proxy admins). - - ``{"created_by": }`` for user-keyed callers. + - ``{"created_by": }`` for user-keyed callers, and for keys + with no user_id and no team_id (owner id is their hashed token). - ``{"team_id": }`` for service-account callers that have a team but no user_id. - ``{"OR": [...]}`` when the caller has both — listing must include @@ -62,12 +90,13 @@ def build_owner_filter( ] } - if user_id is not None: - return {"created_by": user_id} - if team_id is not None: return {"team_id": team_id} + owner_id: Final = resolve_resource_owner_id(user_api_key_dict) + if owner_id is not None: + return {"created_by": owner_id} + return None @@ -86,8 +115,8 @@ def can_access_resource( if _user_has_admin_view(user_api_key_dict): return True - user_id: Final = user_api_key_dict.user_id - if user_id is not None and created_by is not None and created_by == user_id: + owner_id: Final = resolve_resource_owner_id(user_api_key_dict) + if owner_id is not None and created_by is not None and created_by == owner_id: return True team_id: Final = user_api_key_dict.team_id diff --git a/litellm/main.py b/litellm/main.py index cafa1e4718f..c341db08155 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -8612,9 +8612,9 @@ def _joined_streamed_citations(streamed_citations: "tuple[object, ...]") -> "lis def _stream_builder_model_map_cost(response: ModelResponse) -> float | None: - model_name: Final = getattr(response, "model", None) + model_name: Final = response.model usage: Final = getattr(response, "usage", None) - if not isinstance(model_name, str) or not model_name or not isinstance(usage, Usage): + if not model_name or not isinstance(usage, Usage): return None try: prompt_cost, completion_tokens_cost = litellm.cost_per_token(model=model_name, usage_object=usage) diff --git a/litellm/proxy/guardrails/guardrail_hooks/lakera_ai_v2.py b/litellm/proxy/guardrails/guardrail_hooks/lakera_ai_v2.py index 7791adeb41e..bcaffa8e91c 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/lakera_ai_v2.py +++ b/litellm/proxy/guardrails/guardrail_hooks/lakera_ai_v2.py @@ -301,12 +301,12 @@ class LakeraAIGuardrail(CustomGuardrail): explicit sync below a hot reload that changes mode would pass validation but keep dispatching on the stale event_hook. """ - new_event_hook: Final = getattr(litellm_params, "mode", None) or self.event_hook - prospective_payload: Final = getattr(litellm_params, "payload", None) - prospective_breakdown: Final = getattr(litellm_params, "breakdown", None) + new_event_hook: Final = litellm_params.mode or self.event_hook + prospective_payload: Final = litellm_params.payload + prospective_breakdown: Final = litellm_params.breakdown self._validate_advisory_config( - on_flagged=getattr(litellm_params, "on_flagged", None) or self.on_flagged, - advisory_system_message=getattr(litellm_params, "advisory_system_message", None), + on_flagged=litellm_params.on_flagged or self.on_flagged, + advisory_system_message=litellm_params.advisory_system_message, payload=self.payload if prospective_payload is None else prospective_payload, breakdown=self.breakdown if prospective_breakdown is None else prospective_breakdown, ) diff --git a/litellm/proxy/guardrails/guardrail_hooks/qualifire/qualifire.py b/litellm/proxy/guardrails/guardrail_hooks/qualifire/qualifire.py index daeb91eb2bd..f834426d619 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/qualifire/qualifire.py +++ b/litellm/proxy/guardrails/guardrail_hooks/qualifire/qualifire.py @@ -121,7 +121,7 @@ class QualifireGuardrail(CustomGuardrail): the live instance untouched instead of raising after it's already been corrupted. Mirrors LakeraAIGuardrail's own override of this same method. """ - prospective_on_flagged: Final = getattr(litellm_params, "on_flagged", None) or self.on_flagged + prospective_on_flagged: Final = litellm_params.on_flagged or self.on_flagged self._validate_on_flagged(prospective_on_flagged) super().update_in_memory_litellm_params(litellm_params=litellm_params) diff --git a/litellm/proxy/guardrails/guardrail_registry.py b/litellm/proxy/guardrails/guardrail_registry.py index 90d5f6f4970..dc13c09dd38 100644 --- a/litellm/proxy/guardrails/guardrail_registry.py +++ b/litellm/proxy/guardrails/guardrail_registry.py @@ -413,14 +413,15 @@ class GuardrailRegistry: raise Exception(f"Error getting guardrail from DB: {e}") -def _apply_configured_bool_override(instance: CustomGuardrail, litellm_params: LitellmParams, param_name: str) -> None: - """Override ``instance.`` only when ``litellm_params`` explicitly - sets it, preserving whatever default the guardrail's own constructor chose +def _apply_configured_bool_overrides(instance: CustomGuardrail, litellm_params: LitellmParams) -> None: + """Override the parallel/raw-scan flags only when ``litellm_params`` explicitly + sets them, preserving whatever default the guardrail's own constructor chose otherwise (its constructor default may be True, so blindly copying an absent/None config value would silently clobber it back to False).""" - configured: Final = getattr(litellm_params, param_name, None) - if configured is not None: - setattr(instance, param_name, bool(configured)) + if litellm_params.run_in_parallel is not None: + instance.run_in_parallel = bool(litellm_params.run_in_parallel) + if litellm_params.scan_raw_request is not None: + instance.scan_raw_request = bool(litellm_params.scan_raw_request) class InMemoryGuardrailHandler: @@ -544,8 +545,7 @@ class InMemoryGuardrailHandler: "skip_tool_message_in_guardrail are enabled together, which excludes every message from " "scanning, so no request content would ever be scanned. Remove one of the two." ) - for override_param in ("run_in_parallel", "scan_raw_request"): - _apply_configured_bool_override(custom_guardrail_callback, litellm_params, override_param) + _apply_configured_bool_overrides(custom_guardrail_callback, litellm_params) parsed_guardrail: Final = Guardrail( guardrail_id=guardrail.get("guardrail_id"), @@ -803,7 +803,6 @@ class InMemoryGuardrailHandler: previous_guardrail: Final = self.IN_MEMORY_GUARDRAILS.get(guardrail_id) previous_source: Final = self._sources.get(guardrail_id, source) - # Remove from memory if exists (also removes from callbacks) if guardrail_id in self.IN_MEMORY_GUARDRAILS: self.delete_in_memory_guardrail(guardrail_id) diff --git a/litellm/proxy/management_endpoints/policy_endpoints/ai_policy_suggester.py b/litellm/proxy/management_endpoints/policy_endpoints/ai_policy_suggester.py index 4bc53678c23..1096954536a 100644 --- a/litellm/proxy/management_endpoints/policy_endpoints/ai_policy_suggester.py +++ b/litellm/proxy/management_endpoints/policy_endpoints/ai_policy_suggester.py @@ -9,6 +9,7 @@ from typing import Final import litellm from litellm._logging import verbose_proxy_logger from litellm.constants import DEFAULT_COMPETITOR_DISCOVERY_MODEL +from litellm.proxy._types import ProxyErrorTypes, ProxyException SUGGEST_TOOL: Final = { "type": "function", @@ -60,6 +61,18 @@ class AiPolicySuggester: system_prompt: Final = self._build_system_prompt(templates) user_prompt: Final = self._build_user_prompt(attack_examples, description) model = model or DEFAULT_COMPETITOR_DISCOVERY_MODEL + custom_llm_provider: Final = model.split("/", 1)[0] if "/" in model else None + supported_params: Final = litellm.get_supported_openai_params( + model=model, + custom_llm_provider=custom_llm_provider, + ) + if supported_params is not None and "tools" not in supported_params: + raise ProxyException( + message=(f"AI policy suggestion requires tool calling; model '{model}' does not support it"), + type=ProxyErrorTypes.validation_error.value, + param="model", + code=400, + ) try: response: Final = await litellm.acompletion( @@ -74,6 +87,7 @@ class AiPolicySuggester: "function": {"name": "select_policy_templates"}, }, temperature=0.2, + drop_params=True, ) tool_calls: Final = response.choices[0].message.tool_calls diff --git a/litellm/proxy/pass_through_endpoints/managed_id_rewriter.py b/litellm/proxy/pass_through_endpoints/managed_id_rewriter.py index 23cfef6576c..567d8375737 100644 --- a/litellm/proxy/pass_through_endpoints/managed_id_rewriter.py +++ b/litellm/proxy/pass_through_endpoints/managed_id_rewriter.py @@ -49,6 +49,7 @@ from litellm._logging import verbose_proxy_logger from litellm.llms.base_llm.managed_resources.isolation import ( build_owner_filter, can_access_resource, + resolve_resource_owner_id, ) from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.batches_endpoints.common_utils import validate_batch_list_limit @@ -686,7 +687,7 @@ async def _mint_or_reuse_object( "file_object": json.dumps(body_snapshot), "model_object_id": namespaced_model_object_id, "file_purpose": file_purpose, - "created_by": user_api_key_dict.user_id, + "created_by": resolve_resource_owner_id(user_api_key_dict), "team_id": user_api_key_dict.team_id, "updated_by": user_api_key_dict.user_id, }, diff --git a/litellm/proxy/policy_engine/pipeline_executor.py b/litellm/proxy/policy_engine/pipeline_executor.py index a5619821197..50cb813c6fa 100644 --- a/litellm/proxy/policy_engine/pipeline_executor.py +++ b/litellm/proxy/policy_engine/pipeline_executor.py @@ -185,7 +185,7 @@ class PipelineExecutor: # snapshot instead of `data` (which earlier pass_data steps in # this same pipeline may have already rewritten), same reason # the normal sequential/parallel guardrail loops do this. - scans_raw_request: Final = getattr(callback, "scan_raw_request", False) + scans_raw_request: Final = callback.scan_raw_request hook_input: Final[dict] = ( # mutable-ok: same request-payload shape as data independent_snapshot(raw_request_snapshot) if scans_raw_request and raw_request_snapshot is not None diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 597acdf661f..3a70750528f 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -16,7 +16,7 @@ import threading import time import traceback import warnings -from collections.abc import AsyncGenerator, AsyncIterator, Callable, Mapping, MutableMapping, Sequence +from collections.abc import AsyncGenerator, AsyncIterator, Callable, Collection, Mapping, MutableMapping, Sequence from datetime import datetime, timedelta, timezone from types import MappingProxyType, UnionType from typing import ( @@ -12792,6 +12792,7 @@ async def _fetch_db_models_for_search( size: int, sort_by: str | None, is_byok_outside_caller_teams: Callable[[dict[str, JsonValue]], bool], + model_name: str | None = None, ) -> tuple[list[dict[str, Any]], int]: """ Run the bounded DB query that backs `/v2/model/info?search=`. Returns @@ -12808,7 +12809,9 @@ async def _fetch_db_models_for_search( filter for `team_public_model_name` instead and keep the DB cost bounded by `search`. """ - db_where_condition: Final[dict[str, Any]] = {"model_name": {"contains": search_lower, "mode": "insensitive"}} + db_where_condition: Final[dict[str, Any]] = { + "model_name": {"contains": search_lower, "mode": "insensitive"} if model_name is None else model_name + } if db_model_ids_in_router: db_where_condition["model_id"] = {"not": {"in": list(db_model_ids_in_router)}} @@ -12855,6 +12858,7 @@ async def _apply_search_filter_to_models( page: int = 1, size: int = 50, sort_by: str | None = None, + model_name: str | None = None, ) -> tuple[list[dict[str, Any]], int | None]: """ Apply search filter to models, querying database for additional matching models. @@ -12875,6 +12879,11 @@ async def _apply_search_filter_to_models( sort_by: Sort field. When set, results must be sorted across the full match set, so the DB fetch is capped at ``_SORTED_SEARCH_DB_FETCH_CAP`` instead of one page. + model_name: Exact ``model_name`` the caller already narrowed + ``all_models`` to (``?model=``). The DB query matches it + exactly instead of the substring, and is skipped when the + substring cannot occur in it, otherwise rows from other model + groups leak into the result and the count. Returns: Tuple of (filtered_models, total_count). total_count is None if not searching. @@ -12932,7 +12941,8 @@ async def _apply_search_filter_to_models( # Query database for additional models with search term db_models: list[dict[str, Any]] = [] - if prisma_client is not None: + exact_name_can_match: Final = model_name is None or search_lower in model_name.lower() + if prisma_client is not None and exact_name_can_match: try: db_models, db_models_total_count = await _fetch_db_models_for_search( prisma_client=prisma_client, @@ -12944,6 +12954,7 @@ async def _apply_search_filter_to_models( size=size, sort_by=sort_by, is_byok_outside_caller_teams=_is_byok_outside_caller_teams, + model_name=model_name, ) search_total_count = router_models_count + db_models_total_count except Exception as e: @@ -13497,7 +13508,7 @@ async def model_info_v2( all_models += [user_model] if model is not None: - all_models = [m for m in all_models if m["model_name"] == model] + all_models = [m for m in all_models if _deployment_matches_allowed_model_names(m, frozenset((model,)))] # Apply search filter if provided all_models, search_total_count = await _apply_search_filter_to_models( @@ -13509,6 +13520,7 @@ async def model_info_v2( page=page, size=size, sort_by=sortBy, + model_name=model, ) if user_models_only: @@ -14023,7 +14035,7 @@ async def model_metrics_exceptions( return {"data": response, "exception_types": list(exception_types)} -def _deployment_matches_allowed_model_names(model: dict[str, JsonValue], allowed_model_names: set[str]) -> bool: +def _deployment_matches_allowed_model_names(model: dict[str, JsonValue], allowed_model_names: Collection[str]) -> bool: """Match a router deployment against allowed public model names. Team-scoped rows store an internal routing key in ``model_name``; callers diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 1d16fa63607..eab56c31c39 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -1416,7 +1416,7 @@ class ProxyLogging: mutation is discarded and a warning is logged so the misconfiguration is visible instead of silently forwarding unredacted content. """ - scans_raw_request: Final = getattr(callback, "scan_raw_request", False) + scans_raw_request: Final = callback.scan_raw_request should_use_raw_snapshot: Final = scans_raw_request and raw_request_snapshot is not None input_data: Final = ( # mutable-ok: same request-payload shape as data independent_snapshot(raw_request_snapshot) if should_use_raw_snapshot else data @@ -1453,7 +1453,7 @@ class ProxyLogging: "scan_raw_request is for block-only guardrails and this mutation is being " "discarded. Remove scan_raw_request from this guardrail's config if it needs " "to mask/rewrite content.", - getattr(callback, "guardrail_name", None) or callback.__class__.__name__, + callback.guardrail_name or callback.__class__.__name__, ) if scans_raw_request: if result is not None: @@ -1778,7 +1778,7 @@ class ProxyLogging: # guarantee must hold even under litellm.safe_memory_mode, which # otherwise makes deep copies return the original object. needs_raw_request_snapshot: Final = any( - isinstance(cb, CustomGuardrail) and getattr(cb, "scan_raw_request", False) + isinstance(cb, CustomGuardrail) and cb.scan_raw_request for cb in ProxyLogging._callback_capabilities().resolved_callbacks ) raw_request_snapshot: Final[dict | None] = ( # mutable-ok: same request-payload shape as data @@ -1938,7 +1938,7 @@ class ProxyLogging: """ def _input_for(callback: CustomGuardrail) -> dict: # mutable-ok: same request-payload shape as data - if not getattr(callback, "scan_raw_request", False) or raw_request_snapshot is None: + if not callback.scan_raw_request or raw_request_snapshot is None: return data return independent_snapshot(raw_request_snapshot) @@ -1962,11 +1962,7 @@ class ProxyLogging: # deployment-level guardrail sharing this name would see no marker # via _pre_call_hook_already_ran and re-run it a second time on # live kwargs. - if ( - getattr(callback, "scan_raw_request", False) - and not isinstance(result, BaseException) - and result is not None - ): + if callback.scan_raw_request and not isinstance(result, BaseException) and result is not None: callback.mark_pre_call_hook_ran(data) raised: Final = tuple(result for result in results if isinstance(result, BaseException)) blocking: Final = next((exc for exc in raised if not _exception_changes_request_flow(exc)), None) diff --git a/pyproject.toml b/pyproject.toml index a0db4d49467..34c1fec1c11 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -67,8 +67,8 @@ proxy = [ "azure-identity>=1.25.2,<2.0", "azure-storage-blob>=12.28.0,<13.0", "mcp>=1.28.1,<2.0", - "litellm-proxy-extras==0.4.90", - "litellm-enterprise==0.1.61", + "litellm-proxy-extras==0.4.91", + "litellm-enterprise==0.1.62", "RestrictedPython>=8.5,<9.0", "rich>=13.9.4,<14.0", "InquirerPy>=0.3.4,<1.0", diff --git a/tests/test_litellm/enterprise/proxy/test_managed_files_access_check.py b/tests/test_litellm/enterprise/proxy/test_managed_files_access_check.py index c75c8099ea1..ad46798b788 100644 --- a/tests/test_litellm/enterprise/proxy/test_managed_files_access_check.py +++ b/tests/test_litellm/enterprise/proxy/test_managed_files_access_check.py @@ -10,11 +10,14 @@ with deployment credentials, bypassing the managed files access-control hooks. import base64 import pytest +from types import SimpleNamespace from unittest.mock import AsyncMock, MagicMock, patch from fastapi import HTTPException -from litellm.proxy._types import UserAPIKeyAuth +from litellm.caching.dual_cache import DualCache +from litellm.proxy._types import CallTypes, UserAPIKeyAuth +from litellm.types.utils import LiteLLMBatch def _make_user_api_key_dict(user_id: str) -> UserAPIKeyAuth: @@ -161,6 +164,108 @@ async def test_service_account_blocked_from_other_team_file(): assert exc_info.value.status_code == 403 +# --- Keyless key must not be locked out of the batch it created --- + + +def _make_unified_batch_id() -> str: + raw = "litellm_proxy;model_id:my-model-id;llm_batch_id:batch_raw_123" + return base64.urlsafe_b64encode(raw.encode()).decode().rstrip("=") + + +def _make_managed_files_instance_with_object_store(): + """Managed-files hook backed by an in-memory stand-in for the managed + object table, so create and retrieve exercise the same stored row.""" + from litellm_enterprise.proxy.hooks.managed_files import ( + _PROXY_LiteLLMManagedFiles, + ) + + store = {} + + async def upsert(where, data): + store[where["unified_object_id"]] = SimpleNamespace(**data["create"]) + + async def find_first(where): + return store.get(where["unified_object_id"]) + + mock_prisma = MagicMock() + mock_prisma.db.litellm_managedobjecttable.upsert = AsyncMock(side_effect=upsert) + mock_prisma.db.litellm_managedobjecttable.find_first = AsyncMock( + side_effect=find_first + ) + + return ( + _PROXY_LiteLLMManagedFiles( + internal_usage_cache=DualCache(), + prisma_client=mock_prisma, + ), + store, + ) + + +async def _store_batch(managed_files, unified_batch_id: str, creator: UserAPIKeyAuth): + await managed_files.store_unified_object_id( + unified_object_id=unified_batch_id, + file_object=LiteLLMBatch( + id="batch_raw_123", + completion_window="24h", + created_at=0, + endpoint="/v1/chat/completions", + input_file_id="file-1", + object="batch", + status="validating", + ), + litellm_parent_otel_span=None, + model_object_id="batch_raw_123", + file_purpose="batch", + user_api_key_dict=creator, + ) + + +@pytest.mark.asyncio +async def test_keyless_key_can_retrieve_the_batch_it_created(): + """Regression: a key with no user_id and no team_id (what `/key/generate` + by a proxy admin and service-account keys produce) stamped + `created_by=None` and was then denied its own managed batch with + "User None does not have access".""" + unified_batch_id = _make_unified_batch_id() + managed_files, store = _make_managed_files_instance_with_object_store() + keyless = UserAPIKeyAuth(api_key="sk-keyless", parent_otel_span=None) + + await _store_batch(managed_files, unified_batch_id, keyless) + assert store[unified_batch_id].created_by == f"key:{keyless.token}" + + data = {"batch_id": unified_batch_id} + await managed_files.async_pre_call_hook( + user_api_key_dict=keyless, + cache=DualCache(), + data=data, + call_type=CallTypes.aretrieve_batch.value, + ) + assert data["batch_id"] == "batch_raw_123" + assert data["model"] == "my-model-id" + + +@pytest.mark.asyncio +async def test_other_keyless_key_still_denied_the_batch(): + unified_batch_id = _make_unified_batch_id() + managed_files, _ = _make_managed_files_instance_with_object_store() + + await _store_batch( + managed_files, + unified_batch_id, + UserAPIKeyAuth(api_key="sk-creator", parent_otel_span=None), + ) + + with pytest.raises(HTTPException) as exc_info: + await managed_files.async_pre_call_hook( + user_api_key_dict=UserAPIKeyAuth(api_key="sk-other", parent_otel_span=None), + cache=DualCache(), + data={"batch_id": unified_batch_id}, + call_type=CallTypes.aretrieve_batch.value, + ) + assert exc_info.value.status_code == 403 + + # --- Option C fix test: check_batch_cost bypasses managed files hook --- diff --git a/tests/test_litellm/enterprise/proxy/test_managed_files_hook.py b/tests/test_litellm/enterprise/proxy/test_managed_files_hook.py index eddfc4fbd34..f3ad8a8592e 100644 --- a/tests/test_litellm/enterprise/proxy/test_managed_files_hook.py +++ b/tests/test_litellm/enterprise/proxy/test_managed_files_hook.py @@ -527,13 +527,33 @@ async def test_afile_list_orders_newest_first_and_breaks_ties_on_the_cursor_colu @pytest.mark.asyncio -async def test_afile_list_denies_a_caller_without_a_user_or_team(): +async def test_afile_list_scopes_a_keyless_key_to_its_own_hashed_token(): + caller = UserAPIKeyAuth(api_key="sk-test", parent_otel_span=None) + managed_files, table = _make_managed_files_over_rows( + [ + _make_managed_file_row("unified-mine", created_by=f"key:{caller.token}"), + _make_managed_file_row("unified-theirs", created_by="other-user"), + ] + ) + + response = await managed_files.afile_list( + purpose=None, + litellm_parent_otel_span=None, + user_api_key_dict=caller, + ) + + assert [file.id for file in response.data] == ["unified-mine"] + assert table.find_many_calls[0]["where"] == {"created_by": f"key:{caller.token}"} + + +@pytest.mark.asyncio +async def test_afile_list_denies_a_caller_with_no_identity_at_all(): managed_files, table = _make_managed_files_over_rows([_make_managed_file_row("unified-mine")]) response = await managed_files.afile_list( purpose=None, litellm_parent_otel_span=None, - user_api_key_dict=UserAPIKeyAuth(api_key="sk-test", parent_otel_span=None), + user_api_key_dict=UserAPIKeyAuth(parent_otel_span=None), ) assert response.data == [] diff --git a/tests/test_litellm/llms/base_llm/test_managed_resource_isolation.py b/tests/test_litellm/llms/base_llm/test_managed_resource_isolation.py index b5fcd9d8219..1746926c689 100644 --- a/tests/test_litellm/llms/base_llm/test_managed_resource_isolation.py +++ b/tests/test_litellm/llms/base_llm/test_managed_resource_isolation.py @@ -7,6 +7,7 @@ import pytest from litellm.llms.base_llm.managed_resources.isolation import ( build_owner_filter, can_access_resource, + resolve_resource_owner_id, ) from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth @@ -154,3 +155,46 @@ def test_access_identity_less_caller_always_denied(created_by, resource_team_id) ) is False ) + + +# --------------------------------------------------------------------------- +# keyless keys (no user_id, no team_id) own their resources by hashed token +# --------------------------------------------------------------------------- + + +def test_owner_id_prefers_user_id_then_falls_back_to_token(): + assert resolve_resource_owner_id(UserAPIKeyAuth(user_id="alice")) == "alice" + assert resolve_resource_owner_id(UserAPIKeyAuth(team_id="team-eng")) is None + assert resolve_resource_owner_id(UserAPIKeyAuth()) is None + + keyless = UserAPIKeyAuth(api_key="sk-keyless") + assert resolve_resource_owner_id(keyless) == f"key:{keyless.token}" + + +def test_keyless_key_can_access_its_own_resource(): + """Regression for the self-lockout: a key generated by a proxy admin (or a + service-account key) has no user_id and no team_id, so it used to stamp + `created_by=None` and then be denied its own batches and files.""" + keyless = UserAPIKeyAuth(api_key="sk-keyless") + owner_id = resolve_resource_owner_id(keyless) + + assert build_owner_filter(keyless) == {"created_by": owner_id} + assert ( + can_access_resource(keyless, created_by=owner_id, resource_team_id=None) is True + ) + + +def test_keyless_key_denied_another_keyless_keys_resource(): + """The #27004 isolation invariant: two distinct keyless keys must not see + each other's resources.""" + creator = UserAPIKeyAuth(api_key="sk-creator") + other = UserAPIKeyAuth(api_key="sk-other") + + assert ( + can_access_resource( + other, + created_by=resolve_resource_owner_id(creator), + resource_team_id=None, + ) + is False + ) diff --git a/tests/test_litellm/proxy/management_endpoints/policy_endpoints/test_ai_policy_suggester.py b/tests/test_litellm/proxy/management_endpoints/policy_endpoints/test_ai_policy_suggester.py index e3893a66094..93dc429168f 100644 --- a/tests/test_litellm/proxy/management_endpoints/policy_endpoints/test_ai_policy_suggester.py +++ b/tests/test_litellm/proxy/management_endpoints/policy_endpoints/test_ai_policy_suggester.py @@ -7,6 +7,9 @@ from unittest.mock import AsyncMock, MagicMock, patch import pytest +import litellm + +from litellm.proxy._types import ProxyException from litellm.proxy.management_endpoints.policy_endpoints.ai_policy_suggester import ( SUGGEST_TOOL, AiPolicySuggester, @@ -234,6 +237,7 @@ class TestAiPolicySuggester: call_kwargs = mock_acompletion.call_args.kwargs assert call_kwargs["model"] == "gpt-4o-mini" assert call_kwargs["temperature"] == 0.2 + assert call_kwargs["drop_params"] is True assert len(call_kwargs["tools"]) == 1 assert call_kwargs["tools"][0]["function"]["name"] == "select_policy_templates" assert ( @@ -242,3 +246,76 @@ class TestAiPolicySuggester: assert len(call_kwargs["messages"]) == 2 assert call_kwargs["messages"][0]["role"] == "system" assert call_kwargs["messages"][1]["role"] == "user" + + +class TestSuggesterRejectsModelsWithoutToolCalling: + @pytest.mark.asyncio + async def test_a_tools_less_model_is_rejected(self, local_model_cost_map): + with pytest.raises(ProxyException) as exc: + await AiPolicySuggester().suggest( + templates=SAMPLE_TEMPLATES, + attack_examples=["Ignore all previous instructions"], + description="Block prompt injection attempts", + model="perplexity/sonar", + ) + + assert int(exc.value.code) == 400 + assert exc.value.param == "model" + assert "tool calling" in exc.value.message + + def test_a_model_without_forced_tool_choice_support_remains_eligible(self, local_model_cost_map): + supported_params = litellm.get_supported_openai_params( + model="amazon.nova-pro-v1:0", + custom_llm_provider="bedrock", + ) + + assert supported_params is not None + assert "tools" in supported_params + assert "tool_choice" not in supported_params + + +class TestSuggesterToleratesAModelThatRefusesItsSamplingParams: + """The model is operator-supplied, so it can be a reasoning model whose only accepted + temperature is 1. This call pins temperature=0.2 for tool-selection determinism, which such + a model rejects outright: without drop_params litellm raises UnsupportedParamsError and the + whole suggestion fails rather than degrading. Every other internal LLM call in the proxy + already opts in through judge_acompletion; this one was the exception. + """ + + @pytest.mark.asyncio + async def test_a_reasoning_model_gets_past_param_mapping(self, monkeypatch, local_model_cost_map): + """Drives the real entry point with no patching and no network. Which exception escapes is + the discriminator: param mapping runs before any credential check, so UnsupportedParamsError + means the call died on the pinned temperature, while AuthenticationError means it survived + that and got as far as needing a key. Asserting the latter is what the caller observes. + """ + monkeypatch.delenv("OPENAI_API_KEY", raising=False) + + with pytest.raises(litellm.AuthenticationError): + await AiPolicySuggester().suggest( + templates=SAMPLE_TEMPLATES, + attack_examples=["My SSN is 123-45-6789"], + description="", + model="gpt-5.6-terra", + ) + + def test_the_pinned_temperature_is_what_such_a_model_refuses(self, local_model_cost_map): + """The other half of the discriminator above: the same temperature this call pins is + exactly what the model rejects, and drop_params is what removes it.""" + from litellm.utils import get_optional_params + + optional_params = get_optional_params( + model="gpt-5.6-terra", + custom_llm_provider="openai", + temperature=0.2, + tools=[SUGGEST_TOOL], + tool_choice={"type": "function", "function": {"name": "select_policy_templates"}}, + drop_params=True, + ) + + assert "temperature" not in optional_params + assert optional_params["tools"] == [SUGGEST_TOOL] + assert optional_params["tool_choice"] == { + "type": "function", + "function": {"name": "select_policy_templates"}, + } diff --git a/tests/test_litellm/proxy/proxy_server/test_team_model_name_translation.py b/tests/test_litellm/proxy/proxy_server/test_team_model_name_translation.py index 2f4018b55ab..038d061350f 100644 --- a/tests/test_litellm/proxy/proxy_server/test_team_model_name_translation.py +++ b/tests/test_litellm/proxy/proxy_server/test_team_model_name_translation.py @@ -154,6 +154,59 @@ async def test_model_info_v2_translates_team_model_name(monkeypatch): assert "model_name_team-abc-123_4a6b8" not in names +@pytest.mark.asyncio +async def test_model_info_v2_exact_model_filter_matches_team_public_name(monkeypatch): + """`/v2/model/info?model=` must keep the team-scoped row whose + `model_name` is the internal routing key: the dashboard links team model + chips with the public name, and the exact filter ran before translation.""" + global_row = { + "model_name": "gpt-4o", + "litellm_params": {"model": "gpt-4o"}, + "model_info": {"id": "normal-id-1", "db_model": False}, + } + router = MagicMock() + router.model_list = [_team_row(), global_row] + + monkeypatch.setattr(ps, "llm_router", router) + monkeypatch.setattr(ps, "user_model", None) + monkeypatch.setattr(ps, "prisma_client", MagicMock()) + monkeypatch.setattr(ps.proxy_config, "get_config", AsyncMock(return_value={})) + monkeypatch.setattr( + ps, + "_apply_search_filter_to_models", + AsyncMock(side_effect=lambda all_models, **kw: (all_models, len(all_models))), + ) + monkeypatch.setattr( + ps, "_enrich_model_info_with_litellm_data", lambda model, **kw: model + ) + import litellm.proxy.agent_endpoints.model_list_helpers as mlh + + monkeypatch.setattr( + mlh, + "append_agents_to_model_info", + AsyncMock(side_effect=lambda models, **kw: models), + ) + + admin = UserAPIKeyAuth(user_id="u", user_role=LitellmUserRoles.PROXY_ADMIN) + resp = await ps.model_info_v2( + user_api_key_dict=admin, + model="team-claude-sonnet", + user_models_only=False, + include_team_models=False, + debug=False, + page=1, + size=50, + search=None, + modelId=None, + teamId=None, + sortBy=None, + sortOrder="asc", + ) + + assert [m["model_name"] for m in resp["data"]] == ["team-claude-sonnet"] + assert resp["total_count"] == 1 + + @pytest.mark.asyncio async def test_model_info_v1_list_path_translates_team_model_name(monkeypatch): """/v1/model/info list path (no litellm_model_id) must include team-scoped diff --git a/tests/test_litellm/proxy/test_proxy_server.py b/tests/test_litellm/proxy/test_proxy_server.py index 8e470cc663b..949088ea3ba 100644 --- a/tests/test_litellm/proxy/test_proxy_server.py +++ b/tests/test_litellm/proxy/test_proxy_server.py @@ -2126,6 +2126,53 @@ async def test_apply_search_filter_bounds_db_fetch_by_page_and_cap(): assert take < 10_000, "sorted search must cap below the full match set" +@pytest.mark.asyncio +async def test_apply_search_filter_honours_exact_model_name_in_db_query(): + """ + `/v2/model/info?model=&search=`: the router list is already + narrowed to the exact group, so the DB count and fetch must be too, or + other groups' rows leak into the page and inflate total_count. + """ + from litellm.proxy.proxy_server import _apply_search_filter_to_models + + prisma_client = MagicMock() + prisma_client.db.litellm_proxymodeltable.count = AsyncMock(return_value=0) + prisma_client.db.litellm_proxymodeltable.find_many = AsyncMock(return_value=[]) + proxy_config = MagicMock() + proxy_config.decrypt_model_list_from_db = lambda rows: [] + + await _apply_search_filter_to_models( + all_models=[], + search="sonnet", + prisma_client=prisma_client, + proxy_config=proxy_config, + model_name="anthropic-sonnet-5", + ) + where = prisma_client.db.litellm_proxymodeltable.count.call_args.kwargs["where"] + assert where["model_name"] == "anthropic-sonnet-5" + assert prisma_client.db.litellm_proxymodeltable.find_many.call_args.kwargs["where"] == where + + prisma_client.db.litellm_proxymodeltable.count.reset_mock() + _, total_count = await _apply_search_filter_to_models( + all_models=[], + search="opus", + prisma_client=prisma_client, + proxy_config=proxy_config, + model_name="anthropic-sonnet-5", + ) + prisma_client.db.litellm_proxymodeltable.count.assert_not_called() + assert total_count == 0 + + await _apply_search_filter_to_models( + all_models=[], + search="sonnet", + prisma_client=prisma_client, + proxy_config=proxy_config, + ) + where = prisma_client.db.litellm_proxymodeltable.count.call_args.kwargs["where"] + assert where["model_name"] == {"contains": "sonnet", "mode": "insensitive"} + + @pytest.mark.asyncio async def test_filter_models_by_team_id_excludes_viewer_direct_access(): """ diff --git a/tests/test_litellm/proxy/utils/proxy_logging/test_pre_call_hook.py b/tests/test_litellm/proxy/utils/proxy_logging/test_pre_call_hook.py index 0971ce09d79..9d2a27ce9d3 100644 --- a/tests/test_litellm/proxy/utils/proxy_logging/test_pre_call_hook.py +++ b/tests/test_litellm/proxy/utils/proxy_logging/test_pre_call_hook.py @@ -9,11 +9,14 @@ import pytest from fastapi import HTTPException import litellm +from litellm.caching.caching import DualCache from litellm.exceptions import RejectedRequestError from litellm.integrations.custom_guardrail import CustomGuardrail from litellm.integrations.custom_logger import CustomLogger +from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.utils import ProxyLogging from litellm.types.guardrails import GuardrailEventHooks +from litellm.types.utils import CallTypesLiteral def _load(module: str, name: str): @@ -473,7 +476,13 @@ class _RedactingGuardrail(CustomGuardrail): kwargs.setdefault("event_hook", GuardrailEventHooks.pre_call) super().__init__(guardrail_name="redactor", **kwargs) - async def async_pre_call_hook(self, user_api_key_dict, cache, data, call_type): # type: ignore[override] + async def async_pre_call_hook( + self, + user_api_key_dict: UserAPIKeyAuth, + cache: DualCache, + data: dict, + call_type: CallTypesLiteral, + ) -> dict | None: for msg in data.get("messages", []): if "SECRET" in msg.get("content", ""): msg["content"] = msg["content"].replace("SECRET", "[REDACTED]") @@ -488,7 +497,13 @@ class _BlockOnSecretGuardrail(CustomGuardrail): kwargs.setdefault("event_hook", GuardrailEventHooks.pre_call) super().__init__(guardrail_name="blocker", **kwargs) - async def async_pre_call_hook(self, user_api_key_dict, cache, data, call_type): # type: ignore[override] + async def async_pre_call_hook( + self, + user_api_key_dict: UserAPIKeyAuth, + cache: DualCache, + data: dict, + call_type: CallTypesLiteral, + ) -> dict | None: if any("SECRET" in msg.get("content", "") for msg in data.get("messages", [])): raise HTTPException(status_code=400, detail="blocked: SECRET detected") return None @@ -560,7 +575,13 @@ async def test_scan_raw_request_guardrail_does_not_undo_later_masking( separate marker (PII_TOKEN) that only the redactor reacts to.""" class _PiiRedactor(_RedactingGuardrail): - async def async_pre_call_hook(self, user_api_key_dict, cache, data, call_type): # type: ignore[override] + async def async_pre_call_hook( + self, + user_api_key_dict: UserAPIKeyAuth, + cache: DualCache, + data: dict, + call_type: CallTypesLiteral, + ) -> dict | None: for msg in data.get("messages", []): if "PII_TOKEN" in msg.get("content", ""): msg["content"] = msg["content"].replace("PII_TOKEN", "[REDACTED]") @@ -692,7 +713,13 @@ async def test_scan_raw_request_warns_when_guardrail_mutation_discarded( super().__init__(**kwargs) self.scan_raw_request = True - async def async_pre_call_hook(self, user_api_key_dict, cache, data, call_type): # type: ignore[override] + async def async_pre_call_hook( + self, + user_api_key_dict: UserAPIKeyAuth, + cache: DualCache, + data: dict, + call_type: CallTypesLiteral, + ) -> dict | None: for msg in data.get("messages", []): msg["content"] = msg["content"].replace("SECRET", "[REDACTED]") return data diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/models/useModels.test.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/models/useModels.test.ts index 411e8402e11..7231c126a63 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/hooks/models/useModels.test.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/models/useModels.test.ts @@ -118,6 +118,7 @@ describe("useModelsInfo", () => { // exclude_auto_routers defaults off: only the Models + Endpoints table opts in, so // every other consumer of this hook keeps seeing auto-routers. false, + undefined, ); expect(modelInfoCall).toHaveBeenCalledTimes(1); }); @@ -145,6 +146,7 @@ describe("useModelsInfo", () => { // exclude_auto_routers defaults off: only the Models + Endpoints table opts in, so // every other consumer of this hook keeps seeing auto-routers. false, + undefined, ); }); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/models/useModels.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/models/useModels.ts index a5fbc433ea3..a9f7c54698a 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/hooks/models/useModels.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/models/useModels.ts @@ -38,6 +38,7 @@ export const useModelsInfo = ( sortBy?: string, sortOrder?: string, excludeAutoRouters: boolean = false, + modelName?: string, ) => { const { accessToken, userId, userRole } = useAuthorized(); return useQuery({ @@ -48,6 +49,7 @@ export const useModelsInfo = ( page, size, ...(search && { search }), + ...(modelName && { modelName }), ...(modelId && { modelId }), ...(teamId && { teamId }), ...(sortBy && { sortBy }), @@ -70,6 +72,7 @@ export const useModelsInfo = ( sortBy, sortOrder, excludeAutoRouters, + modelName, ), enabled: Boolean(accessToken && userId && userRole), }); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/AllModelsTab.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/AllModelsTab.test.tsx index 4d0b1c466a4..65faa85e29e 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/AllModelsTab.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/AllModelsTab.test.tsx @@ -33,6 +33,7 @@ interface ModelsInfoArgs { teamId?: string; sortBy?: string; sortOrder?: string; + modelName?: string; } const modelsInfoCalls: ModelsInfoArgs[] = []; @@ -47,12 +48,14 @@ type UseModelsInfoArgs = [ teamId?: string, sortBy?: string, sortOrder?: string, + excludeAutoRouters?: boolean, + modelName?: string, ]; vi.mock("../../hooks/models/useModels", () => ({ useModelsInfo: (...args: UseModelsInfoArgs) => { - const [page, size, search, , teamId, sortBy, sortOrder] = args; - const call: ModelsInfoArgs = { page, size, search, teamId, sortBy, sortOrder }; + const [page, size, search, , teamId, sortBy, sortOrder, , modelName] = args; + const call: ModelsInfoArgs = { page, size, search, teamId, sortBy, sortOrder, modelName }; modelsInfoCalls.push(call); return { ...modelsInfoResult, refetch: mockRefetch }; }, @@ -260,6 +263,28 @@ describe("AllModelsTab", () => { expect(within(table).queryByText("gpt-4")).not.toBeInTheDocument(); }); + it("asks the server for the exact selected model group so deployments beyond the first page are found", () => { + render(); + + expect(lastModelsInfoCall().modelName).toBe("claude-opus"); + expect(lastModelsInfoCall().search).toBeUndefined(); + }); + + it.each(["all", "wildcard"])("sends no exact model name for the %s pseudo group", (group) => { + render(); + + expect(lastModelsInfoCall().modelName).toBeUndefined(); + }); + + it("keeps the exact model group alongside a typed search", async () => { + render(); + + fireEvent.change(screen.getByPlaceholderText("Search model names…"), { target: { value: "opus" } }); + + await waitFor(() => expect(lastModelsInfoCall().search).toBe("opus")); + expect(lastModelsInfoCall().modelName).toBe("claude-opus"); + }); + it("resets search, filters, team and sorting from the drawer reset button", async () => { const user = userEvent.setup(); render(); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/AllModelsTab.tsx b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/AllModelsTab.tsx index 1a9d33a50bc..be2cf22d71a 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/AllModelsTab.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/AllModelsTab.tsx @@ -81,6 +81,11 @@ const AllModelsTab = ({ }, [modelNameSearch, debouncedUpdateSearch]); const teamIdForQuery = selectedTeamValue === PERSONAL_TEAM_VALUE ? undefined : selectedTeamValue; + const isConcreteModelGroup = + Boolean(selectedModelGroup) && + selectedModelGroup !== ALL_MODEL_GROUPS_VALUE && + selectedModelGroup !== WILDCARD_MODEL_GROUP_VALUE; + const modelNameForQuery = isConcreteModelGroup ? selectedModelGroup ?? undefined : undefined; const sortBy = useMemo(() => { if (sorting.length === 0) return undefined; @@ -108,6 +113,7 @@ const AllModelsTab = ({ // Auto-routers are routing constructs, not deployments; the sibling Auto-Routers tab // lists and manages them. Excluded server-side so total_count stays honest. true, + modelNameForQuery, ); const isLoading = isLoadingModelsInfo || isLoadingModelCostMap; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/detailNavigation.test.ts b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/detailNavigation.test.ts index 292b27618bd..c4ea206022c 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/detailNavigation.test.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/detailNavigation.test.ts @@ -1,7 +1,7 @@ import { act, renderHook, waitFor } from "@testing-library/react"; import { withNuqsTestingAdapter, type UrlUpdateEvent } from "nuqs/adapters/testing"; import { describe, expect, it, vi } from "vitest"; -import { useModelDetailRouting } from "./detailNavigation"; +import { useModelDetailRouting, useModelGroupFilterRouting } from "./detailNavigation"; describe("useModelDetailRouting", () => { it("openModel sets ?model= with a history push", async () => { @@ -54,3 +54,29 @@ describe("useModelDetailRouting", () => { expect(result.current.teamId).toBeNull(); }); }); + +describe("useModelGroupFilterRouting", () => { + it("reads the selected group from ?model_group=", () => { + const { result } = renderHook(() => useModelGroupFilterRouting(), { + wrapper: withNuqsTestingAdapter({ searchParams: "?model_group=gpt-4.1" }), + }); + expect(result.current.modelGroup).toBe("gpt-4.1"); + }); + + it("writes the selected group to ?model_group= and clears it on null", async () => { + const onUrlUpdate = vi.fn<(event: UrlUpdateEvent) => void>(); + const { result } = renderHook(() => useModelGroupFilterRouting(), { + wrapper: withNuqsTestingAdapter({ onUrlUpdate }), + }); + await act(async () => { + result.current.setModelGroup("claude-sonnet-5"); + }); + await waitFor(() => expect(onUrlUpdate).toHaveBeenCalled()); + expect(onUrlUpdate.mock.calls.at(-1)?.[0].searchParams.get("model_group")).toBe("claude-sonnet-5"); + + await act(async () => { + result.current.setModelGroup(null); + }); + await waitFor(() => expect(onUrlUpdate.mock.calls.at(-1)?.[0].searchParams.has("model_group")).toBe(false)); + }); +}); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/detailNavigation.ts b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/detailNavigation.ts index 2cfad341d25..e83a81a53cf 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/detailNavigation.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/detailNavigation.ts @@ -1,4 +1,4 @@ -import { parseAsString, useQueryStates } from "nuqs"; +import { parseAsString, useQueryState, useQueryStates } from "nuqs"; import { useCallback } from "react"; export interface ModelDetailRouting { @@ -41,3 +41,21 @@ export function useModelDetailRouting(): ModelDetailRouting { close, }; } + +export interface ModelGroupFilterRouting { + modelGroup: string | null; + setModelGroup: (modelGroup: string | null) => void; +} + +export function useModelGroupFilterRouting(): ModelGroupFilterRouting { + const [modelGroup, setParam] = useQueryState("model_group", parseAsString); + + const setModelGroup = useCallback( + (next: string | null) => { + void setParam(next); + }, + [setParam], + ); + + return { modelGroup, setModelGroup }; +} diff --git a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/panels/AllModelsPanel.tsx b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/panels/AllModelsPanel.tsx index 9d40ea32185..552a4f57b24 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/panels/AllModelsPanel.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/panels/AllModelsPanel.tsx @@ -1,19 +1,22 @@ "use client"; -import { useState } from "react"; import AllModelsTab from "@/app/(dashboard)/models-and-endpoints/components/AllModelsTab"; +import { ALL_MODEL_GROUPS_VALUE } from "@/app/(dashboard)/models-and-endpoints/components/AllModelsTable"; import { useModelDashboardData } from "@/app/(dashboard)/models-and-endpoints/useModelDashboardData"; -import { useModelDetailRouting } from "@/app/(dashboard)/models-and-endpoints/detailNavigation"; +import { + useModelDetailRouting, + useModelGroupFilterRouting, +} from "@/app/(dashboard)/models-and-endpoints/detailNavigation"; export default function AllModelsPanel() { - const [selectedModelGroup, setSelectedModelGroup] = useState(null); + const { modelGroup, setModelGroup } = useModelGroupFilterRouting(); const { availableModelGroups, availableModelAccessGroups } = useModelDashboardData(); const { openModel, openTeam } = useModelDetailRouting(); return ( setModelGroup(group === ALL_MODEL_GROUPS_VALUE ? null : group)} availableModelGroups={availableModelGroups} availableModelAccessGroups={availableModelAccessGroups} setSelectedModelId={openModel} diff --git a/ui/litellm-dashboard/src/components/CodeBlock.tsx b/ui/litellm-dashboard/src/components/CodeBlock.tsx index 88ef9acf08e..a39d7417302 100644 --- a/ui/litellm-dashboard/src/components/CodeBlock.tsx +++ b/ui/litellm-dashboard/src/components/CodeBlock.tsx @@ -20,10 +20,10 @@ const CodeBlock = ({ code, language }: CodeBlockProps) => { }; return ( -
+
@@ -1738,9 +1748,9 @@ const TeamInfoView: React.FC = ({

Default Member Models

{info.default_team_member_models.map((model, index) => ( - + {model} - + ))}
diff --git a/ui/litellm-dashboard/src/components/templates/key_info_view.test.tsx b/ui/litellm-dashboard/src/components/templates/key_info_view.test.tsx index 0d41d199b22..7817d4e7cea 100644 --- a/ui/litellm-dashboard/src/components/templates/key_info_view.test.tsx +++ b/ui/litellm-dashboard/src/components/templates/key_info_view.test.tsx @@ -561,6 +561,32 @@ describe("KeyInfoView", () => { ); }); + it("links each model chip to the models page filtered to that model group", async () => { + const keyData = { ...MOCK_KEY_DATA, models: ["gpt-4.1", "anthropic/*"] }; + renderWithProviders( + {}} keyId="test-key-id" onKeyDataUpdate={() => {}} teams={[]} />, + ); + + expect(await screen.findByRole("link", { name: "gpt-4.1" })).toHaveAttribute( + "href", + expect.stringContaining("/models-and-endpoints?model_group=gpt-4.1"), + ); + expect(screen.getByRole("link", { name: "anthropic/*" })).toHaveAttribute( + "href", + expect.stringContaining("/models-and-endpoints?model_group=anthropic%2F*"), + ); + }); + + it("keeps the all-proxy-models grant chip non-clickable", async () => { + const keyData = { ...MOCK_KEY_DATA, models: ["all-proxy-models"] }; + renderWithProviders( + {}} keyId="test-key-id" onKeyDataUpdate={() => {}} teams={[]} />, + ); + + expect((await screen.findAllByText("all-proxy-models")).length).toBeGreaterThan(0); + expect(screen.queryByRole("link", { name: "all-proxy-models" })).not.toBeInTheDocument(); + }); + it("renders no team link when the key has no team", async () => { renderWithProviders( {currentKeyData.models && currentKeyData.models.length > 0 ? ( currentKeyData.models.map((model, index) => ( - + {model} - + )) ) : (

No models specified

@@ -996,9 +997,9 @@ export default function KeyInfoView({
{currentKeyData.models && currentKeyData.models.length > 0 ? ( currentKeyData.models.map((model, index) => ( - + {model} - + )) ) : (

No models specified

diff --git a/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/JsonViewer.test.tsx b/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/JsonViewer.test.tsx index 055e400167b..a0ccb7f5323 100644 --- a/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/JsonViewer.test.tsx +++ b/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/JsonViewer.test.tsx @@ -1,17 +1,26 @@ import { render, screen } from "@testing-library/react"; +import { ThemeProvider } from "next-themes"; +import { darkStyles, defaultStyles } from "react-json-view-lite"; import { describe, expect, it } from "vitest"; import { JsonViewer } from "./JsonViewer"; +const renderWithTheme = (theme: "light" | "dark", data: unknown) => + render( + + + , + ); + describe("JsonViewer", () => { it("should render a placeholder and no tree when the log entry carries no payload", () => { - render(); + renderWithTheme("light", null); expect(screen.getByText("No data")).toBeInTheDocument(); expect(screen.queryByRole("tree")).not.toBeInTheDocument(); }); it("should render the payload as a tree exposing its keys", () => { - render(); + renderWithTheme("light", { model: "claude-opus-4-5", stream: true }); expect(screen.getByRole("tree")).toBeInTheDocument(); expect(screen.getByText(/model/)).toBeInTheDocument(); @@ -20,9 +29,26 @@ describe("JsonViewer", () => { }); it("should treat an empty payload as data rather than showing the placeholder", () => { - render(); + renderWithTheme("light", {}); expect(screen.getByRole("tree")).toBeInTheDocument(); expect(screen.queryByText("No data")).not.toBeInTheDocument(); }); + + it("should style the tree with the light palette when the dashboard theme is light", () => { + renderWithTheme("light", { model: "claude-opus-4-5" }); + + expect(screen.getByRole("tree")).toHaveClass(...defaultStyles.container.split(" ")); + }); + + it("should style the tree with the dark palette when the dashboard theme is dark", () => { + renderWithTheme("dark", { model: "claude-opus-4-5" }); + + const tree = screen.getByRole("tree"); + expect(tree).toHaveClass(...darkStyles.container.split(" ")); + defaultStyles.container + .split(" ") + .filter((className) => !darkStyles.container.split(" ").includes(className)) + .forEach((lightOnlyClassName) => expect(tree).not.toHaveClass(lightOnlyClassName)); + }); }); diff --git a/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/JsonViewer.tsx b/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/JsonViewer.tsx index 86b897f526f..980573a65cc 100644 --- a/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/JsonViewer.tsx +++ b/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/JsonViewer.tsx @@ -1,4 +1,5 @@ -import { JsonView, defaultStyles } from "react-json-view-lite"; +import { useTheme } from "next-themes"; +import { JsonView, darkStyles, defaultStyles } from "react-json-view-lite"; import "react-json-view-lite/dist/index.css"; import { JSON_MAX_HEIGHT, SPACING_LARGE } from "./constants"; @@ -12,6 +13,8 @@ interface JsonViewerProps { * Uses an interactive tree component for easy navigation. */ export function JsonViewer({ data }: JsonViewerProps) { + const { resolvedTheme } = useTheme(); + if (!data) return No data; return ( @@ -24,8 +27,8 @@ export function JsonViewer({ data }: JsonViewerProps) { borderRadius: 4, }} > -
- +
+
); diff --git a/ui/litellm-dashboard/src/utils/entityLinks.test.ts b/ui/litellm-dashboard/src/utils/entityLinks.test.ts new file mode 100644 index 00000000000..47161a903ed --- /dev/null +++ b/ui/litellm-dashboard/src/utils/entityLinks.test.ts @@ -0,0 +1,19 @@ +import { describe, expect, it, vi } from "vitest"; + +vi.mock("@/components/networking", () => ({ serverRootPath: "" })); + +import { modelGroupHref } from "./entityLinks"; + +describe("modelGroupHref", () => { + it("targets the models page filtered to the encoded model group", () => { + expect(modelGroupHref("gpt-4.1")).toMatch(/\/models-and-endpoints\?model_group=gpt-4\.1$/); + expect(modelGroupHref("openai/*")).toMatch(/\?model_group=openai%2F\*$/); + }); + + it.each(["all-proxy-models", "all-team-models", "no-default-models"])( + "returns no href for the %s grant sentinel", + (sentinel) => { + expect(modelGroupHref(sentinel)).toBeUndefined(); + }, + ); +}); diff --git a/ui/litellm-dashboard/src/utils/entityLinks.ts b/ui/litellm-dashboard/src/utils/entityLinks.ts index 675ac8d0554..ad257ec7969 100644 --- a/ui/litellm-dashboard/src/utils/entityLinks.ts +++ b/ui/litellm-dashboard/src/utils/entityLinks.ts @@ -1,5 +1,11 @@ import { migratedHref } from "@/utils/migratedPages"; +const MODEL_GRANT_SENTINELS: ReadonlySet = new Set([ + "all-proxy-models", + "all-team-models", + "no-default-models", +]); + export function teamDetailHref(teamId: string): string { return `${migratedHref("teams")}?team=${encodeURIComponent(teamId)}`; } @@ -15,3 +21,8 @@ export function userDetailHref(userId: string): string { export function orgDetailHref(orgId: string): string { return `${migratedHref("organizations")}?org=${encodeURIComponent(orgId)}`; } + +export function modelGroupHref(modelGroup: string): string | undefined { + if (MODEL_GRANT_SENTINELS.has(modelGroup)) return undefined; + return `${migratedHref("models-and-endpoints")}?model_group=${encodeURIComponent(modelGroup)}`; +} diff --git a/uv.lock b/uv.lock index 019c5a70f9e..8ef72116466 100644 --- a/uv.lock +++ b/uv.lock @@ -10,7 +10,7 @@ resolution-markers = [ ] [options] -exclude-newer = "2026-08-25T23:16:47.126855Z" +exclude-newer = "2026-08-26T18:33:25.773031Z" exclude-newer-span = "P3D" [manifest] @@ -4665,12 +4665,12 @@ proxy-dev = [ [[package]] name = "litellm-enterprise" -version = "0.1.61" +version = "0.1.62" source = { editable = "enterprise" } [[package]] name = "litellm-proxy-extras" -version = "0.4.90" +version = "0.4.91" source = { editable = "litellm-proxy-extras" } [[package]]