From 5e23db8e03d770d1dee78a5d7b3e4903619cb95e Mon Sep 17 00:00:00 2001 From: yujonglee Date: Fri, 11 Sep 2026 16:22:55 -0700 Subject: [PATCH] feat(ocr): add Azure Mistral adapter with native authentication (#40502) * feat(ocr): move Azure credential resolution to Rust * fix(auth): keep shared primitives warning-free * fix(auth): preserve missing key provider errors * fix(auth): enforce Azure input provenance * fix(ocr): preserve proxy credential provenance --- .github/scripts/verify_linux_native_wheel.py | 4 +- litellm-rust/Cargo.lock | 605 ++++++++++++++- litellm-rust/Cargo.toml | 5 + .../crates/ai-gateway/src/io/realtime.rs | 6 +- .../crates/ai-gateway/src/io/responses_ws.rs | 6 +- litellm-rust/crates/core/Cargo.toml | 6 + .../crates/core/src/auth/credential.rs | 168 +++++ litellm-rust/crates/core/src/auth/error.rs | 120 +++ litellm-rust/crates/core/src/auth/http.rs | 86 +++ litellm-rust/crates/core/src/auth/mod.rs | 55 ++ litellm-rust/crates/core/src/auth/policy.rs | 114 +++ litellm-rust/crates/core/src/auth/secret.rs | 41 + litellm-rust/crates/core/src/auth/token.rs | 47 ++ litellm-rust/crates/core/src/error.rs | 19 +- litellm-rust/crates/core/src/lib.rs | 2 + .../core/src/ocr/adapters/azure_mistral.rs | 129 +++- litellm-rust/crates/core/src/ocr/types.rs | 10 + litellm-rust/crates/core/src/ocr/wire.rs | 15 + .../anthropic/messages/transformation.rs | 9 +- .../auth/credential_provider_cache.rs | 43 ++ .../core/src/providers/azure_ai/auth/mod.rs | 7 + .../src/providers/azure_ai/auth/native.rs | 702 ++++++++++++++++++ .../src/providers/azure_ai/auth/resolve.rs | 683 +++++++++++++++++ .../core/src/providers/azure_ai/auth/types.rs | 195 +++++ .../azure_ai/messages/transformation.rs | 16 +- .../crates/core/src/providers/azure_ai/mod.rs | 1 + .../crates/core/tests/azure_ai_ocr.rs | 22 + litellm-rust/crates/core/tests/ocr.rs | 2 + litellm-rust/crates/core/tests/ocr/support.rs | 1 + .../python-bridge/src/routes/definition.rs | 2 +- .../crates/python-bridge/src/routes/ocr.rs | 9 + litellm/ocr/main.py | 454 ++++++----- litellm/proxy/litellm_pre_call_utils.py | 1 + litellm/rust_bridge/ocr.py | 22 +- ...cr_azure_document_intelligence_api_base.py | 24 +- .../ocr/test_ocr_native_format.py | 41 +- tests/test_litellm/ocr/test_rust_bridge.py | 431 +++++++++-- .../proxy/test_litellm_pre_call_utils.py | 3 + .../rust_bridge/native_route_wheel_test.py | 2 +- 39 files changed, 3767 insertions(+), 341 deletions(-) create mode 100644 litellm-rust/crates/core/src/auth/credential.rs create mode 100644 litellm-rust/crates/core/src/auth/error.rs create mode 100644 litellm-rust/crates/core/src/auth/http.rs create mode 100644 litellm-rust/crates/core/src/auth/mod.rs create mode 100644 litellm-rust/crates/core/src/auth/policy.rs create mode 100644 litellm-rust/crates/core/src/auth/secret.rs create mode 100644 litellm-rust/crates/core/src/auth/token.rs create mode 100644 litellm-rust/crates/core/src/providers/azure_ai/auth/credential_provider_cache.rs create mode 100644 litellm-rust/crates/core/src/providers/azure_ai/auth/mod.rs create mode 100644 litellm-rust/crates/core/src/providers/azure_ai/auth/native.rs create mode 100644 litellm-rust/crates/core/src/providers/azure_ai/auth/resolve.rs create mode 100644 litellm-rust/crates/core/src/providers/azure_ai/auth/types.rs diff --git a/.github/scripts/verify_linux_native_wheel.py b/.github/scripts/verify_linux_native_wheel.py index 899e2a211c0..4fb8f068eb0 100644 --- a/.github/scripts/verify_linux_native_wheel.py +++ b/.github/scripts/verify_linux_native_wheel.py @@ -205,7 +205,7 @@ def main( native_module: Final = load_native_module(native_path) native_module_loads: Final = native_module is not None panic_test_hook_absent: Final = native_module is not None and not hasattr(native_module, "_panic_for_test") - native_size_limit: Final = 20_000_000 + native_size_limit: Final = 25_000_000 native_size_within_limit: Final = native_member.file_size <= native_size_limit validations: Final = ( (f"Python tag is {EXPECTED_PYTHON_TAG}", python_tag == EXPECTED_PYTHON_TAG), @@ -222,7 +222,7 @@ def main( ("Python extension entry point is present", extension_entry_point_present), ("Native module loads", native_module_loads), ("Production module omits the panic test hook", panic_test_hook_absent), - ("Native extension does not exceed 20 MB", native_size_within_limit), + ("Native extension does not exceed 25 MB", native_size_within_limit), ("Wheel contents are valid", not unexpected_members), ) diff --git a/litellm-rust/Cargo.lock b/litellm-rust/Cargo.lock index a7c514270da..9c0a8cb7fe7 100644 --- a/litellm-rust/Cargo.lock +++ b/litellm-rust/Cargo.lock @@ -2,6 +2,12 @@ # It is not intended for manual editing. version = 4 +[[package]] +name = "adler2" +version = "2.0.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "320119579fcad9c21884f5c4861d16174d0e06250625266f50fe6898340abefa" + [[package]] name = "ahash" version = "0.8.12" @@ -55,6 +61,29 @@ dependencies = [ "rustversion", ] +[[package]] +name = "async-compression" +version = "0.4.46" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4f10dafd0c8d2e51ae9a748805777613ed0bbe17bf586b76c8311f45c020a32f" +dependencies = [ + "compression-codecs", + "compression-core", + "pin-project-lite", + "tokio", +] + +[[package]] +name = "async-lock" +version = "3.4.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "290f7f2596bd5b78a9fec8088ccd89180d7f9f55b94b0576823bbbdc72ee8311" +dependencies = [ + "event-listener", + "event-listener-strategy", + "pin-project-lite", +] + [[package]] name = "async-trait" version = "0.1.91" @@ -482,6 +511,58 @@ dependencies = [ "tracing", ] +[[package]] +name = "azure_core" +version = "1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4e41cbd819986ba41904c207d8ffc4106f8f8352a548d773e9554906379bb2fb" +dependencies = [ + "async-lock", + "async-trait", + "azure_core_macros", + "bytes", + "futures", + "pin-project", + "rustc_version", + "serde", + "serde_json", + "tokio", + "tracing", + "typespec", + "typespec_client_core", +] + +[[package]] +name = "azure_core_macros" +version = "1.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b9b52dba6a345f3ad2d42ff8d0d63df9d0994cfa29657bf18ffdbf149f78a4f5" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.119", + "tracing", +] + +[[package]] +name = "azure_identity" +version = "1.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "32edf96b356ca7c51d7590c4925cc36efc3947a5da4468e8e0b25c56ecbb3de5" +dependencies = [ + "async-lock", + "async-trait", + "azure_core", + "futures", + "pin-project", + "serde", + "serde_json", + "time", + "tokio", + "tracing", + "url", +] + [[package]] name = "base64" version = "0.13.1" @@ -494,6 +575,12 @@ version = "0.22.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "72b3254f16251a8381aa12e40e3c4d2f0199f8c6508fbecb9d91f575e0fbb8c6" +[[package]] +name = "base64" +version = "0.23.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ac07cdecf99051d9a5238b80f35af32cdeba5b336e55d957b318b50137e18da5" + [[package]] name = "base64-simd" version = "0.8.0" @@ -673,6 +760,16 @@ version = "0.5.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "0c9ea0ac24bc397ab3c98583a3c9ba74fa56b09a4449bbe172b9b1ddb016027a" +[[package]] +name = "combine" +version = "4.6.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cfc320937d09e6de266b31b9afb480f197d7a861be86be7cb2ea7e5d1bfffc5e" +dependencies = [ + "bytes", + "memchr", +] + [[package]] name = "compact_str" version = "0.9.1" @@ -688,6 +785,23 @@ dependencies = [ "static_assertions", ] +[[package]] +name = "compression-codecs" +version = "0.4.41" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "58a6d0db8759036a783bc7c3f7a07f8cef3bf9470eb1db3bc86e8bcd1c5d0fe8" +dependencies = [ + "compression-core", + "flate2", + "memchr", +] + +[[package]] +name = "compression-core" +version = "0.4.33" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6e8ccc4ea9f6acc32d102c0f6d471d11d913ad15f20c04de743374861fa1d414" + [[package]] name = "const-oid" version = "0.10.2" @@ -728,6 +842,15 @@ dependencies = [ "libc", ] +[[package]] +name = "crc32fast" +version = "1.5.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8498c871161e1742aaa9d52551b2d6ebdd4c3d45a3be423e3728f33b955be550" +dependencies = [ + "cfg-if", +] + [[package]] name = "criterion" version = "0.8.2" @@ -763,6 +886,15 @@ dependencies = [ "itertools 0.13.0", ] +[[package]] +name = "crossbeam-channel" +version = "0.5.17" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "98b0cc327b5bc766e7fda9c9260cc0fa81b43a8e240440422dff70788e3f9ef1" +dependencies = [ + "crossbeam-utils", +] + [[package]] name = "crossbeam-deque" version = "0.8.7" @@ -889,6 +1021,9 @@ name = "deranged" version = "0.5.8" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "7cd812cc2bc1d69d4764bd80df88b4317eaef9e773c75226407d9bc0876b211c" +dependencies = [ + "serde_core", +] [[package]] name = "derive_builder" @@ -960,6 +1095,12 @@ version = "1.0.5" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "92773504d58c093f6de2459af4af33faa518c13451eb8f2b5698ed3d36e7c813" +[[package]] +name = "dyn-clone" +version = "1.0.20" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d0881ea181b1df73ff77ffaaf9c7544ecc11e82fba9b5f27b262a3c73a332555" + [[package]] name = "either" version = "1.16.0" @@ -972,12 +1113,42 @@ version = "1.0.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "877a4ace8713b0bcf2a4e7eec82529c029f1d0619886d18145fea96c3ffe5c0f" +[[package]] +name = "errno" +version = "0.3.14" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "39cab71617ae0d63f51a36d69f866391735b51691dbda63cf6f96d042b63efeb" +dependencies = [ + "libc", + "windows-sys 0.61.2", +] + [[package]] name = "esaxx-rs" version = "0.1.10" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "d817e038c30374a4bcb22f94d0a8a0e216958d4c3dcde369b1439fec4bdda6e6" +[[package]] +name = "event-listener" +version = "5.4.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5a23add41df1562121a9393cb065eab5146a1242410f23a644851e90cfd669d2" +dependencies = [ + "parking", + "pin-project-lite", +] + +[[package]] +name = "event-listener-strategy" +version = "0.5.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8be9f3dfaaffdae2972880079a491a1a8bb7cbed0b8dd7a347f668b4150a3b93" +dependencies = [ + "event-listener", + "pin-project-lite", +] + [[package]] name = "fastrand" version = "2.5.0" @@ -990,6 +1161,17 @@ version = "0.1.9" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "5baebc0774151f905a1a2cc41989300b1e6fbb29aff0ceffa1064fdd3088d582" +[[package]] +name = "flate2" +version = "1.1.10" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6e634e2e0ebac1ee034020da1ca582e17ffe4e0f5e985823721e168928136dcb" +dependencies = [ + "crc32fast", + "miniz_oxide", + "zlib-rs", +] + [[package]] name = "fnv" version = "1.0.7" @@ -1011,6 +1193,21 @@ version = "1.3.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "42703706b716c37f96a77aea830392ad231f44c9e9a67872fa5548707e11b11c" +[[package]] +name = "futures" +version = "0.3.33" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a88cf1f829d945f548cf8fec32c61b1f202b6d93b45848602fc02af4b12ad218" +dependencies = [ + "futures-channel", + "futures-core", + "futures-executor", + "futures-io", + "futures-sink", + "futures-task", + "futures-util", +] + [[package]] name = "futures-channel" version = "0.3.33" @@ -1027,6 +1224,17 @@ version = "0.3.33" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "2cd50c473c80f6d7c3670a752354b8e569b1a7cbfdc0419ec88e5edad85e0dc7" +[[package]] +name = "futures-executor" +version = "0.3.33" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6754879cc9f2c66f88c6e5c35344bb0bdb0708b0352b1201815667c7eabc7458" +dependencies = [ + "futures-core", + "futures-task", + "futures-util", +] + [[package]] name = "futures-io" version = "0.3.33" @@ -1068,6 +1276,7 @@ version = "0.3.33" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "a77a90a256fce34da66415271e30f94ee91c57b04b8a2c042d9cf3220179deaa" dependencies = [ + "futures-channel", "futures-core", "futures-io", "futures-macro", @@ -1537,6 +1746,55 @@ version = "1.0.18" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "8f42a60cbdf9a97f5d2305f08a87dc4e09308d1276d28c869c684d7777685682" +[[package]] +name = "jni" +version = "0.22.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5efd9a482cf3a427f00d6b35f14332adc7902ce91efb778580e180ff90fa3498" +dependencies = [ + "cfg-if", + "combine", + "jni-macros", + "jni-sys", + "log", + "simd_cesu8", + "thiserror 2.0.19", + "walkdir", + "windows-link", +] + +[[package]] +name = "jni-macros" +version = "0.22.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a00109accc170f0bdb141fed3e393c565b6f5e072365c3bd58f5b062591560a3" +dependencies = [ + "proc-macro2", + "quote", + "rustc_version", + "simd_cesu8", + "syn 2.0.119", +] + +[[package]] +name = "jni-sys" +version = "0.4.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c6377a88cb3910bee9b0fa88d4f42e1d2da8e79915598f65fb0c7ee14c878af2" +dependencies = [ + "jni-sys-macros", +] + +[[package]] +name = "jni-sys-macros" +version = "0.4.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "38c0b942f458fe50cdac086d2f946512305e5631e720728f2a61aabcd47a6264" +dependencies = [ + "quote", + "syn 2.0.119", +] + [[package]] name = "jobserver" version = "0.1.35" @@ -1580,7 +1838,7 @@ dependencies = [ "futures-util", "litellm-config", "litellm-core", - "reqwest", + "reqwest 0.12.28", "rustls 0.23.42", "rustls-native-certs", "serde", @@ -1613,20 +1871,26 @@ dependencies = [ "aws-sigv4", "aws-smithy-runtime-api", "aws-types", + "azure_core", + "azure_identity", "base64 0.22.1", "data-url", + "moka", "rand 0.8.7", - "reqwest", + "reqwest 0.12.28", "rstest", "serde", "serde_json", "serde_path_to_error", "sha2 0.10.9", + "strum", + "subtle", "thiserror 2.0.19", "tokio", "tracing", "tracing-subscriber", "url", + "veil", ] [[package]] @@ -1681,6 +1945,15 @@ version = "0.8.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "92daf443525c4cce67b150400bc2316076100ce0b3686209eb8cf3c31612e6f0" +[[package]] +name = "lock_api" +version = "0.4.14" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "224399e74b87b5f3557511d98dff8b14089b3dadafcab6bb93eab67d3aace965" +dependencies = [ + "scopeguard", +] + [[package]] name = "log" version = "0.4.33" @@ -1743,6 +2016,16 @@ version = "0.2.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "68354c5c6bd36d73ff3feceb05efa59b6acb7626617f4962be322a825e61f79a" +[[package]] +name = "miniz_oxide" +version = "0.9.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b63fbc4a50860e98e7b2aa7804ded1db5cbc3aff9193adaff57a6931bf7c4b4c" +dependencies = [ + "adler2", + "simd-adler32", +] + [[package]] name = "mio" version = "1.2.2" @@ -1754,6 +2037,26 @@ dependencies = [ "windows-sys 0.61.2", ] +[[package]] +name = "moka" +version = "0.12.16" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4293f18e7567a1caf3c584855554377025c65e0aa445344d04171f5ad63d19b9" +dependencies = [ + "async-lock", + "crossbeam-channel", + "crossbeam-epoch", + "crossbeam-utils", + "equivalent", + "event-listener", + "futures-util", + "parking_lot", + "portable-atomic", + "smallvec", + "tagptr", + "uuid", +] + [[package]] name = "monostate" version = "0.1.18" @@ -1866,6 +2169,35 @@ dependencies = [ "winapi", ] +[[package]] +name = "parking" +version = "2.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f38d5652c16fde515bb1ecef450ab0f6a219d619a7274976324d5e377f7dceba" + +[[package]] +name = "parking_lot" +version = "0.12.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "93857453250e3077bd71ff98b6a65ea6621a19bb0f559a85248955ac12c45a1a" +dependencies = [ + "lock_api", + "parking_lot_core", +] + +[[package]] +name = "parking_lot_core" +version = "0.9.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2621685985a2ebf1c516881c026032ac7deafcda1a2c9b7850dc81e3dfcb64c1" +dependencies = [ + "cfg-if", + "libc", + "redox_syscall", + "smallvec", + "windows-link", +] + [[package]] name = "paste" version = "1.0.15" @@ -1884,6 +2216,26 @@ version = "2.3.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9b4f627cb1b25917193a259e49bdad08f671f8d9708acfd5fe0a8c1455d87220" +[[package]] +name = "pin-project" +version = "1.1.13" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2466b2336ed02bcdca6b294417127b90ec92038d1d5c4fbeac971a922e0e0924" +dependencies = [ + "pin-project-internal", +] + +[[package]] +name = "pin-project-internal" +version = "1.1.13" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c96395f0a926bc13b1c17622aaddda1ecb55d49c8f1bf9777e4d877800a43f8b" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.119", +] + [[package]] name = "pin-project-lite" version = "0.2.17" @@ -2085,6 +2437,7 @@ version = "0.11.16" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "2f4bfc015262b9df63c8845072ce59068853ff5872180c2ce2f13038b970e560" dependencies = [ + "aws-lc-rs", "bytes", "getrandom 0.4.3", "lru-slab", @@ -2252,6 +2605,15 @@ dependencies = [ "crossbeam-utils", ] +[[package]] +name = "redox_syscall" +version = "0.5.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ed2bf2547551a7053d6fdfafda3f938979645c44812fbfcda098faae3f1a362d" +dependencies = [ + "bitflags", +] + [[package]] name = "regex" version = "1.13.1" @@ -2332,11 +2694,49 @@ dependencies = [ "url", "wasm-bindgen", "wasm-bindgen-futures", - "wasm-streams", + "wasm-streams 0.4.2", "web-sys", "webpki-roots", ] +[[package]] +name = "reqwest" +version = "0.13.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "16a1cfa75cc186dd73d5818e510e042e40927bccc9c236b061cea97e1eb08029" +dependencies = [ + "base64 0.23.1", + "bytes", + "futures-core", + "futures-util", + "http 1.4.2", + "http-body 1.1.0", + "http-body-util", + "hyper 1.10.1", + "hyper-rustls 0.27.9", + "hyper-util", + "js-sys", + "log", + "percent-encoding", + "pin-project-lite", + "quinn", + "rustls 0.23.42", + "rustls-pki-types", + "rustls-platform-verifier", + "sync_wrapper", + "tokio", + "tokio-rustls 0.26.4", + "tokio-util", + "tower", + "tower-http", + "tower-service", + "url", + "wasm-bindgen", + "wasm-bindgen-futures", + "wasm-streams 0.5.0", + "web-sys", +] + [[package]] name = "ring" version = "0.17.14" @@ -2444,6 +2844,33 @@ dependencies = [ "zeroize", ] +[[package]] +name = "rustls-platform-verifier" +version = "0.7.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "26d1e2536ce4f35f4846aa13bff16bd0ff40157cdb14cc056c7b14ba41233ba0" +dependencies = [ + "core-foundation", + "core-foundation-sys", + "jni", + "log", + "once_cell", + "rustls 0.23.42", + "rustls-native-certs", + "rustls-platform-verifier-android", + "rustls-webpki 0.103.13", + "security-framework", + "security-framework-sys", + "webpki-root-certs", + "windows-sys 0.61.2", +] + +[[package]] +name = "rustls-platform-verifier-android" +version = "0.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f87165f0995f63a9fbeea62b64d10b4d9d8e78ec6d7d51fb2125fda7bb36788f" + [[package]] name = "rustls-webpki" version = "0.101.7" @@ -2496,6 +2923,12 @@ dependencies = [ "windows-sys 0.61.2", ] +[[package]] +name = "scopeguard" +version = "1.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "94143f37725109f92c262ed2cf5e59bce7498c01bcc1502d7b9afe439a4e9f49" + [[package]] name = "sct" version = "0.7.1" @@ -2649,6 +3082,38 @@ version = "2.0.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "f8fadd59c855ef2080decdef8ff161eb6661b86933c9d82e5ba29dc602a55aba" +[[package]] +name = "signal-hook-registry" +version = "1.4.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c4db69cba1110affc0e9f7bcd48bbf87b3f4fc7c61fc9155afd4c469eb3d6c1b" +dependencies = [ + "errno", + "libc", +] + +[[package]] +name = "simd-adler32" +version = "0.3.10" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3a219298ac11a56ea9a6d2120044824d6f01aeb034955e7af7bc16858527deea" + +[[package]] +name = "simd_cesu8" +version = "1.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "11031e251abf8611c80f460e19dbdeb54a66db918e49c65a7065b46ac7aec520" +dependencies = [ + "rustc_version", + "simdutf8", +] + +[[package]] +name = "simdutf8" +version = "0.1.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e3a9fe34e3e7a50316060351f37187a3f546bce95496156754b601a5fa71b76e" + [[package]] name = "slab" version = "0.4.12" @@ -2711,6 +3176,27 @@ version = "0.11.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "7da8b5736845d9f2fcb837ea5d9e2628564b3b043a70948a3f0b778838c5fb4f" +[[package]] +name = "strum" +version = "0.28.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9628de9b8791db39ceda2b119bbe13134770b56c138ec1d3af810d045c04f9bd" +dependencies = [ + "strum_macros", +] + +[[package]] +name = "strum_macros" +version = "0.28.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ab85eea0270ee17587ed4156089e10b9e6880ee688791d45a905f5b1ca36f664" +dependencies = [ + "heck", + "proc-macro2", + "quote", + "syn 2.0.119", +] + [[package]] name = "subtle" version = "2.6.1" @@ -2759,6 +3245,12 @@ dependencies = [ "syn 2.0.119", ] +[[package]] +name = "tagptr" +version = "0.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7b2093cf4c8eb1e67749a6762251bc9cd836b6fc171623bd0a9d324d37af2417" + [[package]] name = "target-lexicon" version = "0.13.5" @@ -2922,6 +3414,7 @@ dependencies = [ "libc", "mio", "pin-project-lite", + "signal-hook-registry", "socket2 0.6.5", "tokio-macros", "windows-sys 0.61.2", @@ -3039,12 +3532,17 @@ version = "0.6.11" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "4cfcf7e2740e6fc6d4d688b4ef00650406bb94adf4731e43c096c3a19fe40840" dependencies = [ + "async-compression", "bitflags", "bytes", + "futures-core", "futures-util", "http 1.4.2", "http-body 1.1.0", + "http-body-util", "pin-project-lite", + "tokio", + "tokio-util", "tower", "tower-layer", "tower-service", @@ -3138,6 +3636,57 @@ version = "1.20.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "b6f5e870be6c3b371b77fe0ee0bafb859fa4964b4404c27de1d380043c4dda20" +[[package]] +name = "typespec" +version = "1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "753a2fe021e407d4fc9ee6f4f0a33403cc306d5c54c4e4ebe1b8cbde0ca052b9" +dependencies = [ + "base64 0.22.1", + "bytes", + "futures", + "serde", + "serde_json", + "url", +] + +[[package]] +name = "typespec_client_core" +version = "1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0373af0f9d4f580b3a1a9d9639cedaabe015ed262b35bfbe13941bfb14fe1ea6" +dependencies = [ + "async-trait", + "base64 0.22.1", + "bytes", + "dyn-clone", + "futures", + "pin-project", + "rand 0.10.2", + "reqwest 0.13.5", + "serde", + "serde_json", + "time", + "tokio", + "tracing", + "typespec", + "typespec_macros", + "url", + "uuid", +] + +[[package]] +name = "typespec_macros" +version = "1.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2c608f4427943f8adb211abc95c87672b1b98847152783507d54e3246e502f60" +dependencies = [ + "proc-macro2", + "quote", + "rustc_version", + "syn 2.0.119", +] + [[package]] name = "unicase" version = "2.9.0" @@ -3213,10 +3762,32 @@ version = "1.24.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "bf3923a6f5c4c6382e0b653c4117f48d631ea17f38ed86e2a828e6f7412f5239" dependencies = [ + "getrandom 0.4.3", "js-sys", "wasm-bindgen", ] +[[package]] +name = "veil" +version = "0.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c7352f0bbf3ab98911b0c0277065094c1b1ec79bbc85fa3b7d16bf1859c3d96f" +dependencies = [ + "once_cell", + "veil-macros", +] + +[[package]] +name = "veil-macros" +version = "0.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "47a3f4f06d904eb789b935253752ba6bcc1dfa61349f8d5341c66abe070b44e5" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.119", +] + [[package]] name = "version_check" version = "0.9.5" @@ -3331,6 +3902,19 @@ dependencies = [ "web-sys", ] +[[package]] +name = "wasm-streams" +version = "0.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9d1ec4f6517c9e11ae630e200b2b65d193279042e28edd4a2cda233e46670bbb" +dependencies = [ + "futures-util", + "js-sys", + "wasm-bindgen", + "wasm-bindgen-futures", + "web-sys", +] + [[package]] name = "web-sys" version = "0.3.103" @@ -3351,6 +3935,15 @@ dependencies = [ "wasm-bindgen", ] +[[package]] +name = "webpki-root-certs" +version = "1.0.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b96554aa2acc8ccdb7e1c9a58a7a68dd5d13bccc69cd124cb09406db612a1c9b" +dependencies = [ + "rustls-pki-types", +] + [[package]] name = "webpki-roots" version = "1.0.9" @@ -3609,6 +4202,12 @@ dependencies = [ "syn 2.0.119", ] +[[package]] +name = "zlib-rs" +version = "0.6.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "34b31d188d9d685a4f9c7b46d6e36631b07058d2cfe190267adce54dc230bf12" + [[package]] name = "zmij" version = "1.0.23" diff --git a/litellm-rust/Cargo.toml b/litellm-rust/Cargo.toml index f3e54e5b2aa..0f30ac5cf7d 100644 --- a/litellm-rust/Cargo.toml +++ b/litellm-rust/Cargo.toml @@ -41,8 +41,13 @@ tokio = { version = "1", features = ["rt-multi-thread", "macros", "time", "net"] tokio-tungstenite = { version = "0.24", default-features = false, features = ["connect", "rustls-tls-native-roots"] } futures-util = { version = "0.3", default-features = false, features = ["sink", "std"] } base64 = "0.22" +azure_core = "1.0.0" +azure_identity = { version = "1.0.0", features = ["tokio"] } +moka = { version = "0.12.16", features = ["future"] } +strum = { version = "0.28.0", features = ["derive"] } url = "2.5.8" criterion = "0.8.2" +veil = "0.3.0" [profile.release] opt-level = 3 diff --git a/litellm-rust/crates/ai-gateway/src/io/realtime.rs b/litellm-rust/crates/ai-gateway/src/io/realtime.rs index 207c31dffa0..1aa31adcc38 100644 --- a/litellm-rust/crates/ai-gateway/src/io/realtime.rs +++ b/litellm-rust/crates/ai-gateway/src/io/realtime.rs @@ -15,6 +15,8 @@ use std::time::Duration; use futures_util::stream::{SplitSink, SplitStream}; use futures_util::{Sink, SinkExt, Stream, StreamExt}; +use litellm_core::AuthError; +use litellm_core::auth::error::MissingCredential; use litellm_core::error::Error; use litellm_core::realtime::transformation::RealtimeProviderConfig; use litellm_core::realtime::types::RealtimeEvent; @@ -32,8 +34,6 @@ use crate::io::tls::connect_upstream; /// Environment variable holding the OpenAI API key (last-resort fallback). const OPENAI_API_KEY_ENV: &str = "OPENAI_API_KEY"; -const MISSING_KEY_MESSAGE: &str = "Missing OpenAI API Key - a realtime call is being made but no key was passed via params or the OPENAI_API_KEY environment variable"; - /// Default **idle** timeout: if neither side sends a frame for this long, the /// session is reaped. It resets on any activity, so it does not cap a healthy /// (continuously streaming) session — it only frees a stalled one (e.g. a @@ -59,7 +59,7 @@ pub(crate) fn resolve_api_key(api_key: Option<&str>) -> Result { .ok() .filter(|key| !key.trim().is_empty()) }) - .ok_or_else(|| Error::Auth(MISSING_KEY_MESSAGE.to_string())) + .ok_or_else(|| Error::from(AuthError::from(MissingCredential::OpenAiRealtimeApiKey))) } /// Open the upstream WebSocket to OpenAI for `(model, api_key, api_base)`. diff --git a/litellm-rust/crates/ai-gateway/src/io/responses_ws.rs b/litellm-rust/crates/ai-gateway/src/io/responses_ws.rs index 9df3d0c6cc5..7f3b6b0650f 100644 --- a/litellm-rust/crates/ai-gateway/src/io/responses_ws.rs +++ b/litellm-rust/crates/ai-gateway/src/io/responses_ws.rs @@ -4,7 +4,9 @@ use std::time::Duration; use futures_util::stream::{SplitSink, SplitStream}; use futures_util::{Sink, SinkExt, Stream, StreamExt}; +use litellm_core::AuthError; use litellm_core::Error; +use litellm_core::auth::error::MissingCredential; use litellm_core::providers::openai::responses::transformation::OPENAI_RESPONSES_WS_CONFIG; use litellm_core::responses::types::ResponsesWsEvent; use litellm_core::responses::websocket::ResponsesWebSocketProviderConfig; @@ -23,8 +25,6 @@ use crate::constants::{ }; const OPENAI_API_KEY_ENV: &str = "OPENAI_API_KEY"; -const MISSING_KEY_MESSAGE: &str = "Missing OpenAI API Key - a Responses WebSocket call is being made but no key was passed via params or the OPENAI_API_KEY environment variable"; - pub type ResponsesUpstreamWs = WebSocketStream>; type UpstreamTx = SplitSink; type UpstreamRx = SplitStream; @@ -120,7 +120,7 @@ pub(crate) fn resolve_api_key(api_key: Option<&str>) -> Result { .ok() .filter(|value| !value.trim().is_empty()) }) - .ok_or_else(|| Error::Auth(MISSING_KEY_MESSAGE.to_string())) + .ok_or_else(|| Error::from(AuthError::from(MissingCredential::OpenAiResponsesApiKey))) } async fn dial_upstream( diff --git a/litellm-rust/crates/core/Cargo.toml b/litellm-rust/crates/core/Cargo.toml index 6d0a2fa775b..dc4a1acea16 100644 --- a/litellm-rust/crates/core/Cargo.toml +++ b/litellm-rust/crates/core/Cargo.toml @@ -12,18 +12,24 @@ path = "tests/workspace_crate_allowlist.rs" [dependencies] base64.workspace = true +azure_core.workspace = true +azure_identity.workspace = true data-url = "0.3.2" +moka.workspace = true rand.workspace = true reqwest.workspace = true serde.workspace = true serde_json.workspace = true serde_path_to_error = "0.1" +strum.workspace = true +subtle.workspace = true tokio.workspace = true thiserror.workspace = true tracing.workspace = true tracing-subscriber = { workspace = true, optional = true } sha2.workspace = true url.workspace = true +veil.workspace = true aws-config = { version = "1.9.0", default-features = false, features = ["rustls", "rt-tokio"], optional = true } aws-credential-types = { version = "1.3.0", features = ["hardcoded-credentials"], optional = true } aws-sdk-sts = { version = "1.108.0", default-features = false, features = ["rustls", "rt-tokio"], optional = true } diff --git a/litellm-rust/crates/core/src/auth/credential.rs b/litellm-rust/crates/core/src/auth/credential.rs new file mode 100644 index 00000000000..b5235b6780c --- /dev/null +++ b/litellm-rust/crates/core/src/auth/credential.rs @@ -0,0 +1,168 @@ +use std::future::Future; +use std::path::PathBuf; +use std::pin::Pin; +use std::sync::Arc; + +use veil::Redact; + +use crate::AuthError; + +use super::{ResolvedCredential, SecretValue, TokenProviderHandle}; + +#[derive(Clone, Debug, PartialEq, Eq)] +pub enum CredentialFileRef { + Path(PathBuf), + EnvironmentVariable(String), +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub enum CredentialRef { + Explicit(SecretValue), + Env(String), + File(CredentialFileRef), + Request(String), + Host(String), + None, +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub enum CredentialLookup { + Found(SecretValue), + Missing, + Declined, +} + +pub type CredentialLookupFuture<'a> = + Pin> + Send + 'a>>; + +pub trait CredentialResolver: std::fmt::Debug + Send + Sync { + fn resolve<'a>(&'a self, reference: &'a CredentialRef) -> CredentialLookupFuture<'a>; +} + +#[derive(Clone, Redact)] +pub struct CredentialResolverHandle(#[redact(with = "[REDACTED]")] Arc); + +impl CredentialResolverHandle { + pub fn new(resolver: Arc) -> Self { + Self(resolver) + } + + pub async fn resolve(&self, reference: &CredentialRef) -> Result { + self.0.resolve(reference).await + } +} + +#[derive(Clone, Debug)] +pub enum CredentialPlan { + Static(CredentialRef), + Caller(TokenProviderHandle), + None, +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub enum CredentialPlanResolution { + Resolved(ResolvedCredential), + Unavailable, +} + +impl CredentialPlan { + pub async fn resolve( + &self, + resolver: &CredentialResolverHandle, + ) -> Result { + match self { + Self::Static(CredentialRef::Explicit(secret)) => Ok( + CredentialPlanResolution::Resolved(ResolvedCredential::Static(secret.clone())), + ), + Self::Static(CredentialRef::None) | Self::None => { + Ok(CredentialPlanResolution::Unavailable) + } + Self::Static(reference) => match resolver.resolve(reference).await? { + CredentialLookup::Found(secret) => Ok(CredentialPlanResolution::Resolved( + ResolvedCredential::Static(secret), + )), + CredentialLookup::Missing | CredentialLookup::Declined => { + Ok(CredentialPlanResolution::Unavailable) + } + }, + Self::Caller(caller) => { + let credential = caller.acquire().await?; + if credential.secret().expose().is_empty() { + return Err(AuthError::EmptyCallerCredential); + } + Ok(CredentialPlanResolution::Resolved(credential)) + } + } + } +} + +#[cfg(test)] +mod tests { + use std::sync::Arc; + + use super::{ + CredentialLookup, CredentialLookupFuture, CredentialPlan, CredentialPlanResolution, + CredentialRef, CredentialResolver, CredentialResolverHandle, + }; + use crate::AuthError; + use crate::auth::SecretValue; + + #[derive(Debug)] + struct HostResolver; + + impl CredentialResolver for HostResolver { + fn resolve<'a>(&'a self, reference: &'a CredentialRef) -> CredentialLookupFuture<'a> { + Box::pin(async move { + Ok(match reference { + CredentialRef::Host(name) if name == "rotating-token" => { + CredentialLookup::Found(SecretValue::new("resolved")) + } + _ => CredentialLookup::Declined, + }) + }) + } + } + + #[tokio::test] + async fn static_host_reference_resolves_at_acquisition_time() { + let resolver = CredentialResolverHandle::new(Arc::new(HostResolver)); + let plan = CredentialPlan::Static(CredentialRef::Host("rotating-token".to_string())); + + let resolved = plan.resolve(&resolver).await.unwrap(); + + assert!(matches!(resolved, CredentialPlanResolution::Resolved(_))); + } + + #[tokio::test] + async fn declined_reference_is_available_for_pre_acquisition_fallback() { + let resolver = CredentialResolverHandle::new(Arc::new(HostResolver)); + let plan = CredentialPlan::Static(CredentialRef::Request("api-key".to_string())); + + assert_eq!( + plan.resolve(&resolver).await.unwrap(), + CredentialPlanResolution::Unavailable + ); + } + + #[derive(Debug)] + struct FailingResolver; + + impl CredentialResolver for FailingResolver { + fn resolve<'a>(&'a self, _reference: &'a CredentialRef) -> CredentialLookupFuture<'a> { + Box::pin(async { Err(AuthError::UnresolvedOidcReference) }) + } + } + + #[tokio::test] + async fn acquisition_failure_is_terminal() { + let resolver = CredentialResolverHandle::new(Arc::new(FailingResolver)); + let plan = CredentialPlan::Static(CredentialRef::Host("token".to_string())); + + let error = plan + .resolve(&resolver) + .await + .expect_err("acquisition errors cannot become fallback"); + + assert_eq!(error, AuthError::UnresolvedOidcReference); + } +} diff --git a/litellm-rust/crates/core/src/auth/error.rs b/litellm-rust/crates/core/src/auth/error.rs new file mode 100644 index 00000000000..ddd4a6d016e --- /dev/null +++ b/litellm-rust/crates/core/src/auth/error.rs @@ -0,0 +1,120 @@ +use thiserror::Error; + +#[derive(Clone, Debug, Error, PartialEq, Eq)] +pub enum AuthError { + #[error("invalid authentication configuration: {0}")] + Configuration(#[from] AuthConfigurationError), + #[error("credential acquisition failed: {0}")] + AzureTokenAcquisition(String), + #[error("credential acquisition failed: {}", .0.iter().map(ToString::to_string).collect::>().join("; "))] + CredentialChain(Vec), + #[error("credential caller failed: credential caller returned an empty credential")] + EmptyCallerCredential, + #[error("credential caller failed: Azure AD token provider returned an empty token")] + EmptyAzureToken, + #[error("credential acquisition failed: Azure OIDC reference did not resolve to a value")] + UnresolvedOidcReference, + #[error( + "Missing {provider} API Key - A call is being made to {provider} but no key is set either in the environment variables or via params" + )] + MissingApiKey { provider: &'static str }, + #[error( + "Missing {provider} API Base - Set {environment_variable} environment variable or pass api_base parameter" + )] + MissingApiBase { + provider: &'static str, + environment_variable: &'static str, + }, + #[error("{0}")] + MissingCredential(#[from] MissingCredential), + #[error("{0}")] + Aws(#[from] AwsAuthError), + #[error("invalid authentication header")] + InvalidHeader, +} + +#[derive(Clone, Debug, Error, PartialEq, Eq)] +pub enum AuthConfigurationError { + #[error("credential header already exists")] + ExistingCredentialHeader, + #[error("credential plan is not allowed by the provider auth policy")] + DisallowedCredentialPlan, + #[error("credential cannot be empty")] + EmptyCredential, + #[error("invalid Azure credential selector")] + InvalidAzureSelector, + #[error("ClientSecretCredential requires tenant_id, client_id, and client_secret")] + MissingClientSecretFields, + #[error("WorkloadIdentityCredential requires tenant_id")] + MissingWorkloadTenant, + #[error("WorkloadIdentityCredential requires client_id")] + MissingWorkloadClient, + #[error("WorkloadIdentityCredential requires azure_federated_token_file")] + MissingWorkloadTokenFile, + #[error("credential reference requires a host credential resolver")] + MissingHostResolver, + #[error("caller credential plan requires provider-specific inputs")] + MissingCallerInputs, + #[error("credential header {0} already exists")] + DuplicateHeader(&'static str), + #[error("{0} must be a string or null")] + InvalidFieldType(String), + #[error("unsupported OIDC reference")] + UnsupportedOidcReference, + #[error("{0} cannot be empty")] + EmptyReference(String), + #[error("Azure credential initialization failed: {0}")] + AzureCredentialInitialization(String), + #[error("Azure authority must be an HTTPS origin without credentials, query, or fragment")] + InvalidAzureAuthority, + #[error("request-controlled Azure auth inputs cannot be combined with host credentials")] + MixedAzureCredentialSources, + #[error("request-controlled Azure credential references are not allowed")] + RequestAzureCredentialReference, + #[error("host credentials cannot be sent to a request-controlled Azure endpoint")] + RequestAzureCredentialDestination, +} + +#[derive(Clone, Debug, Error, PartialEq, Eq)] +pub enum MissingCredential { + #[error( + "Missing Anthropic API Key - Set `api_key` or the ANTHROPIC_API_KEY environment variable" + )] + AnthropicApiKey, + #[error("Missing Azure API Key - Set `api_key` or the AZURE_API_KEY environment variable")] + AzureApiKey, + #[error( + "Missing Azure API Base - Set `api_base` or the AZURE_API_BASE environment variable. Expected format: https://.services.ai.azure.com/anthropic" + )] + AzureApiBase, + #[error( + "Missing OpenAI API Key - a realtime call is being made but no key was passed via params or the OPENAI_API_KEY environment variable" + )] + OpenAiRealtimeApiKey, + #[error( + "Missing OpenAI API Key - a Responses WebSocket call is being made but no key was passed via params or the OPENAI_API_KEY environment variable" + )] + OpenAiResponsesApiKey, +} + +#[derive(Clone, Debug, Error, PartialEq, Eq)] +pub enum AwsAuthError { + #[error("AWS profile credentials failed: {0}")] + Profile(String), + #[error("AWS default credentials failed: {0}")] + DefaultChain(String), + #[error("AWS role credentials failed: {0}")] + AssumeRole(String), + #[error("AWS web identity credentials failed: {0}")] + WebIdentity(String), + #[error("AWS web identity expiration was invalid: {0}")] + WebIdentityExpiration(String), + #[error("AWS signing parameters failed: {0}")] + SigningParameters(String), + #[error("AWS signable request failed: {0}")] + SignableRequest(String), + #[error("AWS request signing failed: {0}")] + Signing(String), + #[error("AWS web identity response had no credentials")] + MissingWebIdentityCredentials, +} diff --git a/litellm-rust/crates/core/src/auth/http.rs b/litellm-rust/crates/core/src/auth/http.rs new file mode 100644 index 00000000000..83931311550 --- /dev/null +++ b/litellm-rust/crates/core/src/auth/http.rs @@ -0,0 +1,86 @@ +use crate::AuthError; +use crate::auth::error::AuthConfigurationError; + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum CredentialPlacement { + Bearer, + Header(&'static str), +} + +impl CredentialPlacement { + pub fn header_name(self) -> &'static str { + match self { + Self::Bearer => "Authorization", + Self::Header(name) => name, + } + } +} + +pub(crate) fn apply_credential( + headers: Vec<(String, String)>, + credential: &str, + placement: CredentialPlacement, +) -> Result, AuthError> { + if credential.trim().is_empty() { + return Err(AuthError::Configuration( + AuthConfigurationError::EmptyCredential, + )); + } + if headers + .iter() + .any(|(name, _)| name.eq_ignore_ascii_case(placement.header_name())) + { + return Err(AuthError::Configuration( + AuthConfigurationError::DuplicateHeader(placement.header_name()), + )); + } + let value = match placement { + CredentialPlacement::Bearer => format!("Bearer {credential}"), + CredentialPlacement::Header(_) => credential.to_string(), + }; + Ok( + std::iter::once((placement.header_name().to_string(), value)) + .chain(headers) + .collect(), + ) +} + +/// How the upstream call is authenticated. API-key strategies are resolved in +/// `prepare`; SigV4 needs the serialized body, so the handler signs it. +#[derive(Clone, Debug, PartialEq, Eq)] +pub enum RequestAuth { + Header { name: &'static str, value: String }, + Bearer { token: String }, + AwsSigV4 { region: String }, +} + +#[cfg(test)] +mod tests { + use super::{CredentialPlacement, apply_credential}; + + #[test] + fn bearer_uses_authorization_header() { + let headers = apply_credential(Vec::new(), "key", CredentialPlacement::Bearer) + .expect("credential applies"); + + assert_eq!( + headers, + vec![("Authorization".to_string(), "Bearer key".to_string())] + ); + } + + #[test] + fn named_header_rejects_existing_value() { + let error = apply_credential( + vec![( + "ocp-apim-subscription-key".to_string(), + "caller-key".to_string(), + )], + "configured-key", + CredentialPlacement::Header("Ocp-Apim-Subscription-Key"), + ) + .expect_err("provider policy must handle existing credentials"); + + assert!(error.to_string().contains("already exists")); + } +} diff --git a/litellm-rust/crates/core/src/auth/mod.rs b/litellm-rust/crates/core/src/auth/mod.rs new file mode 100644 index 00000000000..2ca2f3c3016 --- /dev/null +++ b/litellm-rust/crates/core/src/auth/mod.rs @@ -0,0 +1,55 @@ +mod credential; +pub mod error; +pub use error::AuthError; +pub(crate) mod http; +mod policy; +mod secret; +mod token; + +use serde::{Deserialize, Serialize}; + +#[derive(Clone, Copy, Debug, Default, Deserialize, Eq, Hash, PartialEq, Serialize)] +#[serde(rename_all = "snake_case")] +pub enum InputSource { + Request, + #[default] + Deployment, + Environment, +} + +#[derive(Clone, Debug, Default, Eq, PartialEq)] +pub struct Sourced { + value: T, + source: InputSource, +} + +impl Sourced { + pub fn new(value: T, source: InputSource) -> Self { + Self { value, source } + } + + pub fn value(&self) -> &T { + &self.value + } + + pub fn source(&self) -> InputSource { + self.source + } + + pub fn into_value(self) -> T { + self.value + } + + pub fn map(self, map: impl FnOnce(T) -> U) -> Sourced { + Sourced::new(map(self.value), self.source) + } +} + +pub use credential::{ + CredentialFileRef, CredentialLookup, CredentialLookupFuture, CredentialPlan, + CredentialPlanResolution, CredentialRef, CredentialResolver, CredentialResolverHandle, +}; +pub use http::{CredentialPlacement, RequestAuth}; +pub use policy::{CredentialPlanKind, CredentialRule, ExistingHeaderBehavior, ProviderAuthPolicy}; +pub use secret::SecretValue; +pub use token::{ResolvedCredential, TokenFuture, TokenProvider, TokenProviderHandle}; diff --git a/litellm-rust/crates/core/src/auth/policy.rs b/litellm-rust/crates/core/src/auth/policy.rs new file mode 100644 index 00000000000..b796dedf0d8 --- /dev/null +++ b/litellm-rust/crates/core/src/auth/policy.rs @@ -0,0 +1,114 @@ +use crate::AuthError; +use crate::auth::error::AuthConfigurationError; + +use super::http::apply_credential; +use super::{CredentialPlacement, ResolvedCredential}; + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum CredentialPlanKind { + Static, + Entra, + Caller, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub struct CredentialRule { + pub kind: CredentialPlanKind, + pub placement: CredentialPlacement, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum ExistingHeaderBehavior { + Preserve, + Reject, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub struct ProviderAuthPolicy { + pub rules: &'static [CredentialRule], + pub accepted_existing_headers: &'static [&'static str], + pub existing_header_behavior: ExistingHeaderBehavior, + pub scope: Option<&'static str>, + pub audience: Option<&'static str>, +} + +impl ProviderAuthPolicy { + pub fn has_existing_credential(&self, headers: &[(String, String)]) -> bool { + headers.iter().any(|(name, _)| { + self.accepted_existing_headers + .iter() + .any(|accepted| name.eq_ignore_ascii_case(accepted)) + }) + } + + pub fn apply( + &self, + headers: Vec<(String, String)>, + kind: CredentialPlanKind, + credential: &ResolvedCredential, + ) -> Result, AuthError> { + if self.has_existing_credential(&headers) { + return match self.existing_header_behavior { + ExistingHeaderBehavior::Preserve => Ok(headers), + ExistingHeaderBehavior::Reject => Err(AuthError::Configuration( + AuthConfigurationError::ExistingCredentialHeader, + )), + }; + } + let rule = + self.rules + .iter() + .find(|rule| rule.kind == kind) + .ok_or(AuthError::Configuration( + AuthConfigurationError::DisallowedCredentialPlan, + ))?; + apply_credential(headers, credential.secret().expose(), rule.placement) + } +} + +#[cfg(test)] +mod tests { + use super::{CredentialPlanKind, CredentialRule, ExistingHeaderBehavior, ProviderAuthPolicy}; + use crate::auth::{CredentialPlacement, ResolvedCredential, SecretValue}; + + const RULES: &[CredentialRule] = &[CredentialRule { + kind: CredentialPlanKind::Static, + placement: CredentialPlacement::Header("x-api-key"), + }]; + const POLICY: ProviderAuthPolicy = ProviderAuthPolicy { + rules: RULES, + accepted_existing_headers: &["x-api-key"], + existing_header_behavior: ExistingHeaderBehavior::Preserve, + scope: None, + audience: None, + }; + + #[test] + fn rules_define_allowed_plans_and_credential_placement() { + let headers = POLICY + .apply( + Vec::new(), + CredentialPlanKind::Static, + &ResolvedCredential::Static(SecretValue::new("secret")), + ) + .unwrap(); + + assert_eq!( + headers, + vec![("x-api-key".to_string(), "secret".to_string())] + ); + } + + #[test] + fn unsupported_plan_is_rejected() { + let error = POLICY + .apply( + Vec::new(), + CredentialPlanKind::Entra, + &ResolvedCredential::Static(SecretValue::new("secret")), + ) + .unwrap_err(); + + assert!(error.to_string().contains("not allowed")); + } +} diff --git a/litellm-rust/crates/core/src/auth/secret.rs b/litellm-rust/crates/core/src/auth/secret.rs new file mode 100644 index 00000000000..3ecb0a835ee --- /dev/null +++ b/litellm-rust/crates/core/src/auth/secret.rs @@ -0,0 +1,41 @@ +use veil::Redact; + +#[derive(Redact, Clone)] +pub struct SecretValue(#[redact(with = "[REDACTED]")] String); + +impl SecretValue { + pub fn new(value: impl Into) -> Self { + Self(value.into()) + } + + pub fn expose(&self) -> &str { + &self.0 + } +} + +impl PartialEq for SecretValue { + fn eq(&self, other: &Self) -> bool { + subtle::ConstantTimeEq::ct_eq(self.0.as_bytes(), other.0.as_bytes()).into() + } +} + +impl Eq for SecretValue {} + +#[cfg(test)] +mod tests { + use super::SecretValue; + + #[test] + fn debug_redacts_plaintext() { + let debug = format!("{:?}", SecretValue::new("credential-value")); + + assert!(!debug.contains("credential-value")); + assert!(debug.contains("REDACTED")); + } + + #[test] + fn equality_compares_plaintext_values() { + assert_eq!(SecretValue::new("same"), SecretValue::new("same")); + assert_ne!(SecretValue::new("same"), SecretValue::new("different")); + } +} diff --git a/litellm-rust/crates/core/src/auth/token.rs b/litellm-rust/crates/core/src/auth/token.rs new file mode 100644 index 00000000000..cfc6b8f0d6b --- /dev/null +++ b/litellm-rust/crates/core/src/auth/token.rs @@ -0,0 +1,47 @@ +use std::future::Future; +use std::pin::Pin; +use std::sync::Arc; +use std::time::SystemTime; + +use veil::Redact; + +use crate::AuthError; + +use super::secret::SecretValue; + +#[derive(Clone, Debug, PartialEq, Eq)] +pub enum ResolvedCredential { + Static(SecretValue), + AccessToken { + token: SecretValue, + expires_on: Option, + }, +} + +impl ResolvedCredential { + pub fn secret(&self) -> &SecretValue { + match self { + Self::Static(secret) | Self::AccessToken { token: secret, .. } => secret, + } + } +} + +pub type TokenFuture<'a> = + Pin> + Send + 'a>>; + +pub trait TokenProvider: std::fmt::Debug + Send + Sync { + fn acquire(&self) -> TokenFuture<'_>; +} + +#[derive(Clone, Redact)] +pub struct TokenProviderHandle(#[redact(with = "[REDACTED]")] Arc); + +impl TokenProviderHandle { + pub fn new(caller: Arc) -> Self { + Self(caller) + } + + pub async fn acquire(&self) -> Result { + self.0.acquire().await + } +} diff --git a/litellm-rust/crates/core/src/error.rs b/litellm-rust/crates/core/src/error.rs index eefa7d606d8..b171c0c4274 100644 --- a/litellm-rust/crates/core/src/error.rs +++ b/litellm-rust/crates/core/src/error.rs @@ -22,7 +22,7 @@ pub enum Error { )] MissingApiKey { provider: &'static str }, #[error( - "Missing Azure AI credentials - set AZURE_AI_API_KEY or provide an Authorization header" + "invalid authentication configuration: Missing Azure AI credentials - set AZURE_AI_API_KEY or configure Entra ID" )] MissingAzureAiCredentials, #[error("Missing Azure AI credentials - set AZURE_AI_API_KEY or provide azure_ad_token")] @@ -121,6 +121,15 @@ impl From for Error { } } +impl From for Error { + fn from(error: crate::AuthError) -> Self { + match error { + crate::AuthError::MissingApiKey { provider } => Self::MissingApiKey { provider }, + error => Self::Auth(error.to_string()), + } + } +} + pub fn json_type_name(value: &serde_json::Value) -> &'static str { match value { serde_json::Value::Null => "null", @@ -136,6 +145,14 @@ pub fn json_type_name(value: &serde_json::Value) -> &'static str { mod transport_tests { use super::*; + #[test] + fn missing_auth_key_preserves_provider_in_public_error() { + assert_eq!( + Error::from(crate::AuthError::MissingApiKey { provider: "Vertex" }), + Error::MissingApiKey { provider: "Vertex" } + ); + } + #[tokio::test] async fn transport_errors_remove_urls_and_keep_dispatch_context() { let error = reqwest::Client::builder() diff --git a/litellm-rust/crates/core/src/lib.rs b/litellm-rust/crates/core/src/lib.rs index 7e81b292441..0b3573deab2 100644 --- a/litellm-rust/crates/core/src/lib.rs +++ b/litellm-rust/crates/core/src/lib.rs @@ -1,4 +1,5 @@ pub mod audio_transcription; +pub mod auth; pub mod caching; pub mod call_lifecycle; pub mod chat_completions; @@ -17,4 +18,5 @@ pub mod router; pub mod routing_utils; mod url_utils; +pub use auth::AuthError; pub use error::Error; diff --git a/litellm-rust/crates/core/src/ocr/adapters/azure_mistral.rs b/litellm-rust/crates/core/src/ocr/adapters/azure_mistral.rs index 468c883a1dd..0e52bc61249 100644 --- a/litellm-rust/crates/core/src/ocr/adapters/azure_mistral.rs +++ b/litellm-rust/crates/core/src/ocr/adapters/azure_mistral.rs @@ -1,5 +1,9 @@ +use std::sync::OnceLock; + use super::OcrAdapter; use crate::Error; +use crate::auth::error::AuthConfigurationError; +use crate::auth::{InputSource, Sourced}; use crate::constants::AZURE_AI_OCR_PATH; use crate::ocr::OcrClient; use crate::ocr::codecs::mistral::{self, MistralOcrParams, MistralOcrResponse}; @@ -10,6 +14,7 @@ use crate::ocr::prepare::{ }; use crate::ocr::registry::OcrProvider; use crate::ocr::types::{LiteLLMOcrRequest, LiteLLMOcrResponse, OcrConnection}; +use crate::providers::azure_ai::auth::{AzureAuthInputs, AzureAuthService}; use crate::url_utils::ApiUrl; const AZURE_AI_API_KEY_ENV: &str = "AZURE_AI_API_KEY"; @@ -31,7 +36,12 @@ impl OcrAdapter for AzureMistralAdapter { known: params, extra_params: _extra_params, } = _prepare_ocr_request::(request)?; - let headers = authenticate(&request.connection, &credential_env)?; + let config = AzureAuthInputs::from_sourced_optional_params( + &request.optional_params, + &request.input_sources, + ) + .map_err(Error::from)?; + let headers = validate_environment(&request.connection, &config, &credential_env).await?; let url = get_complete_url(request.connection.api_base.as_deref(), &credential_env)?; let document = inline_remote_document( client.document_fetcher(), @@ -76,21 +86,61 @@ fn get_complete_url( }) } -fn authenticate( +async fn validate_environment( connection: &OcrConnection, - env_lookup: &dyn Fn(&str) -> Option, + config: &AzureAuthInputs, + env_lookup: &(dyn Fn(&str) -> Option + Sync), ) -> Result, OcrError> { if crate::http_utils::has_header(&connection.extra_headers, "authorization") { + validate_destination(connection, connection.extra_headers_source)?; return Ok(connection.extra_headers.clone()); } let key = nonblank(connection.api_key.clone()) - .or_else(|| nonblank(env_lookup(AZURE_AI_API_KEY_ENV))) + .map(|value| Sourced::new(value, connection.api_key_source)) + .or_else(|| { + nonblank(env_lookup(AZURE_AI_API_KEY_ENV)) + .map(|value| Sourced::new(value, InputSource::Environment)) + }); + if let Some(key) = key { + validate_destination(connection, key.source())?; + return Ok(bearer_headers(connection, key.value())); + } + static SERVICE: OnceLock = OnceLock::new(); + let key = SERVICE + .get_or_init(AzureAuthService::default) + .get_azure_ad_token(config, env_lookup) + .await + .map_err(Error::from)? + .map(|credential| { + let source = credential.source(); + let value = credential.value().secret().expose().to_string(); + Sourced::new(value, source) + }) .ok_or(Error::MissingAzureAiCredentials)?; - Ok( - std::iter::once(("Authorization".into(), format!("Bearer {key}"))) - .chain(connection.extra_headers.clone()) - .collect(), - ) + validate_destination(connection, key.source())?; + Ok(bearer_headers(connection, key.value())) +} + +fn validate_destination( + connection: &OcrConnection, + credential_source: InputSource, +) -> Result<(), OcrError> { + if connection.api_base.is_some() + && connection.api_base_source == InputSource::Request + && credential_source != InputSource::Request + { + return Err(Error::from(crate::AuthError::Configuration( + AuthConfigurationError::RequestAzureCredentialDestination, + )) + .into()); + } + Ok(()) +} + +fn bearer_headers(connection: &OcrConnection, key: &str) -> Vec<(String, String)> { + std::iter::once(("Authorization".into(), format!("Bearer {key}"))) + .chain(connection.extra_headers.clone()) + .collect() } fn nonblank(value: Option) -> Option { @@ -119,27 +169,76 @@ mod tests { ); } - #[test] - fn supplied_authorization_precedes_keys() { + #[tokio::test] + async fn supplied_authorization_precedes_keys() { let connection = OcrConnection { api_key: Some("request-key".into()), extra_headers: vec![("authorization".into(), "Bearer prepared".into())], ..Default::default() }; assert_eq!( - authenticate(&connection, &|_| Some("environment-key".into())).unwrap(), + validate_environment(&connection, &Default::default(), &|_| Some( + "environment-key".into() + )) + .await + .unwrap(), connection.extra_headers ); } - #[test] - fn request_key_precedes_environment_key() { + #[tokio::test] + async fn request_key_precedes_environment_key() { let connection = OcrConnection { api_key: Some("request-key".into()), ..Default::default() }; assert_eq!( - authenticate(&connection, &|_| Some("environment-key".into())).unwrap()[0], + validate_environment(&connection, &Default::default(), &|_| Some( + "environment-key".into() + )) + .await + .unwrap()[0], + ("Authorization".into(), "Bearer request-key".into()) + ); + } + + #[tokio::test] + async fn request_endpoint_cannot_receive_environment_key() { + let connection = OcrConnection { + api_base: Some("https://request.example".into()), + api_base_source: InputSource::Request, + ..Default::default() + }; + + let error = validate_environment(&connection, &Default::default(), &|name| { + (name == AZURE_AI_API_KEY_ENV).then(|| "environment-key".into()) + }) + .await + .unwrap_err(); + + assert!( + error + .to_string() + .contains("request-controlled Azure endpoint") + ); + } + + #[tokio::test] + async fn request_endpoint_accepts_request_owned_key() { + let connection = OcrConnection { + api_key: Some("request-key".into()), + api_key_source: InputSource::Request, + api_base: Some("https://request.example".into()), + api_base_source: InputSource::Request, + ..Default::default() + }; + + let headers = validate_environment(&connection, &Default::default(), &|_| None) + .await + .unwrap(); + + assert_eq!( + headers[0], ("Authorization".into(), "Bearer request-key".into()) ); } diff --git a/litellm-rust/crates/core/src/ocr/types.rs b/litellm-rust/crates/core/src/ocr/types.rs index eeda94738a0..dec474876fb 100644 --- a/litellm-rust/crates/core/src/ocr/types.rs +++ b/litellm-rust/crates/core/src/ocr/types.rs @@ -1,3 +1,4 @@ +use std::collections::BTreeMap; use std::sync::Arc; use std::time::Duration; @@ -7,6 +8,7 @@ use serde_json::{Map, Value}; use super::hooks::{NoopOcrHooks, OcrHooks}; use super::registry::{OcrAdapterKind, resolve_wire_adapter}; use crate::Error; +use crate::auth::InputSource; use crate::constants::OCR_HTTP_TIMEOUT_SECS; #[derive(Clone, Debug, PartialEq, Serialize, Deserialize)] @@ -65,8 +67,11 @@ pub enum OcrResponseFormat { #[derive(Clone)] pub struct OcrConnection { pub api_key: Option, + pub api_key_source: InputSource, pub api_base: Option, + pub api_base_source: InputSource, pub extra_headers: Vec<(String, String)>, + pub extra_headers_source: InputSource, pub timeout: Duration, pub max_download_bytes: u64, } @@ -75,8 +80,11 @@ impl Default for OcrConnection { fn default() -> Self { Self { api_key: None, + api_key_source: InputSource::Deployment, api_base: None, + api_base_source: InputSource::Deployment, extra_headers: Vec::new(), + extra_headers_source: InputSource::Deployment, timeout: Duration::from_secs(OCR_HTTP_TIMEOUT_SECS), max_download_bytes: crate::constants::OCR_DOWNLOAD_MAX_BYTES, } @@ -90,6 +98,7 @@ pub struct LiteLLMOcrRequest { pub hooks: Arc, pub litellm_call_id: Option, pub optional_params: Map, + pub input_sources: BTreeMap, pub(crate) adapter: OcrAdapterKind, } @@ -109,6 +118,7 @@ impl LiteLLMOcrRequest { hooks: Arc::new(NoopOcrHooks), litellm_call_id: None, optional_params, + input_sources: BTreeMap::new(), adapter: adapter_kind, }) } diff --git a/litellm-rust/crates/core/src/ocr/wire.rs b/litellm-rust/crates/core/src/ocr/wire.rs index db8d91905c4..d0fe32378b9 100644 --- a/litellm-rust/crates/core/src/ocr/wire.rs +++ b/litellm-rust/crates/core/src/ocr/wire.rs @@ -1,10 +1,12 @@ use crate::ocr::error::OcrRequestError; use crate::ocr::error::OcrResponseError; +use std::collections::BTreeMap; use std::time::Duration; use super::hooks::{OcrDuringCallRequest, OcrPreCallRequest}; use super::types::{LiteLLMOcrRequest, OcrConnection, OcrDocument}; use crate::Error; +use crate::auth::InputSource; use serde::{ Deserialize, de::{DeserializeOwned, IntoDeserializer}, @@ -28,6 +30,8 @@ pub struct OcrWireRequest { pub extra_headers: Option>, #[serde(default)] pub optional_params: Map, + #[serde(default)] + pub input_sources: BTreeMap, pub timeout_seconds: Option, } @@ -36,6 +40,9 @@ pub fn is_supported_request(model: &str, custom_llm_provider: Option<&str>) -> b } pub fn decode_request(wire: OcrWireRequest) -> Result { + let api_key_source = source_for(&wire.input_sources, "api_key"); + let api_base_source = source_for(&wire.input_sources, "api_base"); + let extra_headers_source = source_for(&wire.input_sources, "extra_headers"); let document = decode_request_value(wire.document, "document")?; let headers = wire .extra_headers @@ -67,17 +74,25 @@ pub fn decode_request(wire: OcrWireRequest) -> Result )?; let connection = OcrConnection { api_key: nonblank(wire.api_key), + api_key_source, api_base: nonblank(wire.api_base), + api_base_source, extra_headers: headers, + extra_headers_source, timeout: timeout.unwrap_or(defaults.timeout), max_download_bytes: defaults.max_download_bytes, }; Ok(LiteLLMOcrRequest { connection, + input_sources: wire.input_sources, ..request }) } +fn source_for(sources: &BTreeMap, name: &str) -> InputSource { + sources.get(name).copied().unwrap_or_default() +} + fn nonblank(value: Option) -> Option { value .map(|s| s.trim().to_string()) diff --git a/litellm-rust/crates/core/src/providers/anthropic/messages/transformation.rs b/litellm-rust/crates/core/src/providers/anthropic/messages/transformation.rs index f31b961e78a..3ed00b7cc5f 100644 --- a/litellm-rust/crates/core/src/providers/anthropic/messages/transformation.rs +++ b/litellm-rust/crates/core/src/providers/anthropic/messages/transformation.rs @@ -1,3 +1,4 @@ +use crate::auth::error::MissingCredential; use crate::error::Error; use crate::messages::transformation::{AnthropicMessagesProviderConfig, MessagesAuthStrategy}; @@ -21,13 +22,7 @@ pub fn resolve_anthropic_api_key( non_empty(api_key) .map(str::to_string) .or_else(|| env_lookup(ANTHROPIC_API_KEY_ENV).filter(|value| !value.trim().is_empty())) - .ok_or_else(|| { - Error::Auth( - "Missing Anthropic API Key - Set `api_key` or the ANTHROPIC_API_KEY \ - environment variable" - .to_string(), - ) - }) + .ok_or_else(|| Error::from(crate::AuthError::from(MissingCredential::AnthropicApiKey))) } pub fn complete_anthropic_url( diff --git a/litellm-rust/crates/core/src/providers/azure_ai/auth/credential_provider_cache.rs b/litellm-rust/crates/core/src/providers/azure_ai/auth/credential_provider_cache.rs new file mode 100644 index 00000000000..297e4cc6502 --- /dev/null +++ b/litellm-rust/crates/core/src/providers/azure_ai/auth/credential_provider_cache.rs @@ -0,0 +1,43 @@ +use std::future::Future; +use std::sync::Arc; + +use azure_core::credentials::TokenCredential; +use moka::future::Cache; + +use crate::AuthError; + +#[derive(Clone, Debug, PartialEq, Eq, Hash)] +pub(crate) struct AzureCredentialProviderCacheKey { + pub(crate) mechanism: &'static str, + pub(crate) authority: String, + pub(crate) tenant_id: String, + pub(crate) client_id: String, + pub(crate) scope: String, + pub(crate) secret_identity: String, +} + +pub(crate) struct AzureCredentialProviderCache { + entries: Cache>, +} + +impl AzureCredentialProviderCache { + pub(crate) fn new(capacity: u64) -> Self { + Self { + entries: Cache::builder().max_capacity(capacity).build(), + } + } + + pub(crate) async fn get_or_create( + &self, + key: AzureCredentialProviderCacheKey, + create: F, + ) -> Result, AuthError> + where + F: Future, AuthError>>, + { + self.entries + .try_get_with(key, create) + .await + .map_err(|error| (*error).clone()) + } +} diff --git a/litellm-rust/crates/core/src/providers/azure_ai/auth/mod.rs b/litellm-rust/crates/core/src/providers/azure_ai/auth/mod.rs new file mode 100644 index 00000000000..33d007c1945 --- /dev/null +++ b/litellm-rust/crates/core/src/providers/azure_ai/auth/mod.rs @@ -0,0 +1,7 @@ +mod credential_provider_cache; +mod native; +mod resolve; +mod types; + +pub(crate) use resolve::AzureAuthService; +pub(crate) use types::AzureAuthInputs; diff --git a/litellm-rust/crates/core/src/providers/azure_ai/auth/native.rs b/litellm-rust/crates/core/src/providers/azure_ai/auth/native.rs new file mode 100644 index 00000000000..b8f19818d16 --- /dev/null +++ b/litellm-rust/crates/core/src/providers/azure_ai/auth/native.rs @@ -0,0 +1,702 @@ +use crate::auth::error::AuthConfigurationError; +use std::sync::Arc; +use std::time::{Duration, UNIX_EPOCH}; + +use azure_core::cloud::{CloudConfiguration, CustomConfiguration}; +use azure_core::credentials::{Secret, TokenCredential}; +use azure_core::http::ClientOptions; +use azure_identity::{ + ClientAssertion, ClientAssertionCredential, ClientAssertionCredentialOptions, + ClientSecretCredential, ClientSecretCredentialOptions, DeveloperToolsCredential, + ManagedIdentityCredential, ManagedIdentityCredentialOptions, UserAssignedId, + WorkloadIdentityCredential, WorkloadIdentityCredentialOptions, +}; +use sha2::{Digest, Sha256}; + +use crate::AuthError; +use crate::auth::{InputSource, ResolvedCredential, SecretValue, Sourced}; + +use super::credential_provider_cache::{ + AzureCredentialProviderCache, AzureCredentialProviderCacheKey, +}; + +#[derive(Clone, Debug)] +pub(crate) enum NativeAzureRequest { + ClientSecret { + tenant_id: Sourced, + client_id: Sourced, + client_secret: Sourced, + scope: Sourced, + authority: Option>, + }, + ClientAssertion { + tenant_id: Sourced, + client_id: Sourced, + assertion: Sourced, + assertion_identity: String, + scope: Sourced, + authority: Option>, + }, + WorkloadIdentity { + tenant_id: Sourced, + client_id: Sourced, + token_file_path: Sourced, + scope: Sourced, + authority: Option>, + }, + ManagedIdentity { + client_id: Option>, + scope: Sourced, + selection_source: InputSource, + }, + DeveloperTools { + scope: Sourced, + selection_source: InputSource, + }, +} + +#[derive(Clone, Debug)] +pub(crate) struct ValidatedAzureRequest { + request: NativeAzureRequest, + credential_source: InputSource, +} + +impl ValidatedAzureRequest { + pub(crate) fn new(request: NativeAzureRequest) -> Result { + validate_authority(&request)?; + let credential_source = validate_sources(&request)?; + Ok(Self { + request, + credential_source, + }) + } + + pub(crate) fn credential_source(&self) -> InputSource { + self.credential_source + } + + #[cfg(test)] + pub(super) fn kind(&self) -> &'static str { + match self.request { + NativeAzureRequest::ClientSecret { .. } => "client-secret", + NativeAzureRequest::ClientAssertion { .. } => "client-assertion", + NativeAzureRequest::WorkloadIdentity { .. } => "workload-identity", + NativeAzureRequest::ManagedIdentity { .. } => "managed-identity", + NativeAzureRequest::DeveloperTools { .. } => "developer-tools", + } + } +} + +pub(crate) struct NativeAzureTokenAcquirer { + cache: AzureCredentialProviderCache, + transport: Option, +} + +impl Default for NativeAzureTokenAcquirer { + fn default() -> Self { + Self::new(64) + } +} + +impl NativeAzureTokenAcquirer { + pub(crate) fn new(cache_capacity: u64) -> Self { + Self { + cache: AzureCredentialProviderCache::new(cache_capacity), + transport: None, + } + } + + #[cfg(test)] + pub(super) fn with_transport( + cache_capacity: u64, + transport: azure_core::http::Transport, + ) -> Self { + Self { + cache: AzureCredentialProviderCache::new(cache_capacity), + transport: Some(transport), + } + } + + pub(crate) async fn acquire( + &self, + request: ValidatedAzureRequest, + ) -> Result { + let scope = request.request.scope().to_string(); + let key = request.request.cache_key(); + let transport = self.transport.clone(); + let credential = self + .cache + .get_or_create( + key, + async move { build_credential(request.request, transport) }, + ) + .await?; + let token = credential + .get_token(&[scope.as_str()], None) + .await + .map_err(|error| AuthError::AzureTokenAcquisition(error.to_string()))?; + let expires_on = u64::try_from(token.expires_on.unix_timestamp()) + .ok() + .map(|seconds| UNIX_EPOCH + Duration::from_secs(seconds)); + + Ok(ResolvedCredential::AccessToken { + token: SecretValue::new(token.token.secret()), + expires_on, + }) + } +} + +impl NativeAzureRequest { + fn scope(&self) -> &str { + match self { + Self::ClientSecret { scope, .. } + | Self::ClientAssertion { scope, .. } + | Self::WorkloadIdentity { scope, .. } + | Self::ManagedIdentity { scope, .. } + | Self::DeveloperTools { scope, .. } => scope.value(), + } + } + + fn cache_key(&self) -> AzureCredentialProviderCacheKey { + match self { + Self::ClientSecret { + tenant_id, + client_id, + client_secret, + scope, + authority, + } => AzureCredentialProviderCacheKey { + mechanism: "client-secret", + authority: authority + .as_ref() + .map(|value| value.value().clone()) + .unwrap_or_default(), + tenant_id: tenant_id.value().clone(), + client_id: client_id.value().clone(), + scope: scope.value().clone(), + secret_identity: secret_digest(client_secret.value().expose()), + }, + Self::ClientAssertion { + tenant_id, + client_id, + assertion, + assertion_identity, + scope, + authority, + } => AzureCredentialProviderCacheKey { + mechanism: "client-assertion", + authority: authority + .as_ref() + .map(|value| value.value().clone()) + .unwrap_or_default(), + tenant_id: tenant_id.value().clone(), + client_id: client_id.value().clone(), + scope: scope.value().clone(), + secret_identity: format!( + "{assertion_identity}:{}", + secret_digest(assertion.value().expose()) + ), + }, + Self::WorkloadIdentity { + tenant_id, + client_id, + token_file_path, + scope, + authority, + } => AzureCredentialProviderCacheKey { + mechanism: "workload-identity", + authority: authority + .as_ref() + .map(|value| value.value().clone()) + .unwrap_or_default(), + tenant_id: tenant_id.value().clone(), + client_id: client_id.value().clone(), + scope: scope.value().clone(), + secret_identity: token_file_path.value().clone(), + }, + Self::ManagedIdentity { + client_id, scope, .. + } => AzureCredentialProviderCacheKey { + mechanism: "managed-identity", + authority: String::new(), + tenant_id: String::new(), + client_id: client_id + .as_ref() + .map(|value| value.value().clone()) + .unwrap_or_default(), + scope: scope.value().clone(), + secret_identity: String::new(), + }, + Self::DeveloperTools { scope, .. } => AzureCredentialProviderCacheKey { + mechanism: "developer-tools", + authority: String::new(), + tenant_id: String::new(), + client_id: String::new(), + scope: scope.value().clone(), + secret_identity: String::new(), + }, + } + } +} + +fn validate_authority(request: &NativeAzureRequest) -> Result<(), AuthError> { + let authority = match request { + NativeAzureRequest::ClientSecret { authority, .. } + | NativeAzureRequest::ClientAssertion { authority, .. } + | NativeAzureRequest::WorkloadIdentity { authority, .. } => authority.as_ref(), + NativeAzureRequest::ManagedIdentity { .. } | NativeAzureRequest::DeveloperTools { .. } => { + None + } + }; + let Some(authority) = authority else { + return Ok(()); + }; + let url = url::Url::parse(authority.value()) + .map_err(|_| AuthError::Configuration(AuthConfigurationError::InvalidAzureAuthority))?; + if url.scheme() != "https" + || url.host_str().is_none() + || !url.username().is_empty() + || url.password().is_some() + || url.query().is_some() + || url.fragment().is_some() + || !matches!(url.path(), "" | "/") + { + return Err(AuthError::Configuration( + AuthConfigurationError::InvalidAzureAuthority, + )); + } + Ok(()) +} + +fn validate_sources(request: &NativeAzureRequest) -> Result { + match request { + NativeAzureRequest::ClientSecret { + tenant_id, + client_id, + client_secret, + scope, + authority, + } => { + let identity_sources = [ + tenant_id.source(), + client_id.source(), + client_secret.source(), + ]; + let request_identity = identity_sources.contains(&InputSource::Request); + if request_identity + && !identity_sources + .iter() + .all(|source| *source == InputSource::Request) + { + return mixed_sources(); + } + if !request_identity && is_request_controlled(scope, authority.as_ref()) { + return mixed_sources(); + } + Ok(if request_identity { + InputSource::Request + } else { + trusted_source(&identity_sources) + }) + } + NativeAzureRequest::ClientAssertion { + tenant_id, + client_id, + assertion, + scope, + authority, + .. + } => trusted_only(&[ + tenant_id.source(), + client_id.source(), + assertion.source(), + scope.source(), + authority + .as_ref() + .map(Sourced::source) + .unwrap_or(InputSource::Environment), + ]), + NativeAzureRequest::WorkloadIdentity { + tenant_id, + client_id, + token_file_path, + scope, + authority, + } => trusted_only(&[ + tenant_id.source(), + client_id.source(), + token_file_path.source(), + scope.source(), + authority + .as_ref() + .map(Sourced::source) + .unwrap_or(InputSource::Environment), + ]), + NativeAzureRequest::ManagedIdentity { + client_id, + scope, + selection_source, + } => trusted_only(&[ + client_id + .as_ref() + .map(Sourced::source) + .unwrap_or(InputSource::Environment), + scope.source(), + *selection_source, + ]), + NativeAzureRequest::DeveloperTools { + scope, + selection_source, + } => trusted_only(&[scope.source(), *selection_source]), + } +} + +fn is_request_controlled(value: &Sourced, optional: Option<&Sourced>) -> bool { + value.source() == InputSource::Request + || optional.is_some_and(|value| value.source() == InputSource::Request) +} + +fn trusted_only(sources: &[InputSource]) -> Result { + if sources.contains(&InputSource::Request) { + return mixed_sources(); + } + Ok(trusted_source(sources)) +} + +fn trusted_source(sources: &[InputSource]) -> InputSource { + if sources.contains(&InputSource::Deployment) { + InputSource::Deployment + } else { + InputSource::Environment + } +} + +fn mixed_sources() -> Result { + Err(AuthError::Configuration( + AuthConfigurationError::MixedAzureCredentialSources, + )) +} + +fn build_credential( + request: NativeAzureRequest, + transport: Option, +) -> Result, AuthError> { + match request { + NativeAzureRequest::ClientSecret { + tenant_id, + client_id, + client_secret, + authority, + .. + } => ClientSecretCredential::new( + tenant_id.value(), + client_id.into_value(), + Secret::new(client_secret.value().expose().to_string()), + Some(ClientSecretCredentialOptions { + client_options: client_options(authority.map(Sourced::into_value), transport), + }), + ) + .map(|credential| credential as Arc), + NativeAzureRequest::ClientAssertion { + tenant_id, + client_id, + assertion, + authority, + .. + } => ClientAssertionCredential::new( + tenant_id.into_value(), + client_id.into_value(), + StaticAssertion(assertion.into_value()), + Some(ClientAssertionCredentialOptions { + client_options: client_options(authority.map(Sourced::into_value), transport), + }), + ) + .map(|credential| credential as Arc), + NativeAzureRequest::WorkloadIdentity { + tenant_id, + client_id, + token_file_path, + authority, + .. + } => WorkloadIdentityCredential::new(Some(WorkloadIdentityCredentialOptions { + credential_options: azure_identity::ClientAssertionCredentialOptions { + client_options: client_options(authority.map(Sourced::into_value), transport), + }, + client_id: Some(client_id.into_value()), + tenant_id: Some(tenant_id.into_value()), + token_file_path: Some(token_file_path.into_value().into()), + })) + .map(|credential| credential as Arc), + NativeAzureRequest::ManagedIdentity { client_id, .. } => { + ManagedIdentityCredential::new(Some(ManagedIdentityCredentialOptions { + user_assigned_id: client_id + .map(Sourced::into_value) + .map(UserAssignedId::ClientId), + client_options: client_options(None, transport), + })) + .map(|credential| credential as Arc) + } + NativeAzureRequest::DeveloperTools { .. } => DeveloperToolsCredential::new(None) + .map(|credential| credential as Arc), + } + .map_err(|error| { + AuthError::Configuration(AuthConfigurationError::AzureCredentialInitialization( + error.to_string(), + )) + }) +} + +fn client_options( + authority: Option, + transport: Option, +) -> ClientOptions { + let cloud = authority.map(|authority_host| { + let mut custom = CustomConfiguration::default(); + custom.authority_host = authority_host; + Arc::new(CloudConfiguration::from(custom)) + }); + ClientOptions { + cloud, + transport, + ..Default::default() + } +} + +fn secret_digest(secret: &str) -> String { + format!("{:x}", Sha256::digest(secret.as_bytes())) +} + +#[derive(Debug)] +struct StaticAssertion(SecretValue); + +impl ClientAssertion for StaticAssertion { + fn secret<'life0, 'life1, 'async_trait>( + &'life0 self, + _options: Option>, + ) -> std::pin::Pin< + Box> + Send + 'async_trait>, + > + where + 'life0: 'async_trait, + 'life1: 'async_trait, + Self: 'async_trait, + { + Box::pin(async move { Ok(self.0.expose().to_string()) }) + } +} + +#[cfg(test)] +mod tests { + use std::sync::{Arc, Mutex}; + + use azure_core::http::headers::Headers; + use azure_core::http::{AsyncRawResponse, HttpClient, Request, StatusCode, Transport}; + use azure_core::{Bytes, Result}; + + use super::{NativeAzureRequest, NativeAzureTokenAcquirer, ValidatedAzureRequest}; + use crate::auth::{InputSource, SecretValue, Sourced}; + + fn deployment(value: T) -> Sourced { + Sourced::new(value, InputSource::Deployment) + } + + fn sourced_client_secret( + credential_source: InputSource, + authority_source: InputSource, + authority: &str, + ) -> NativeAzureRequest { + NativeAzureRequest::ClientSecret { + tenant_id: Sourced::new("tenant".to_string(), credential_source), + client_id: Sourced::new("client".to_string(), credential_source), + client_secret: Sourced::new(SecretValue::new("secret"), credential_source), + scope: Sourced::new("scope".to_string(), InputSource::Environment), + authority: Some(Sourced::new(authority.to_string(), authority_source)), + } + } + + fn client_secret_request( + tenant: &str, + client: &str, + secret: &str, + scope: &str, + authority: &str, + ) -> ValidatedAzureRequest { + ValidatedAzureRequest::new(NativeAzureRequest::ClientSecret { + tenant_id: deployment(tenant.to_string()), + client_id: deployment(client.to_string()), + client_secret: deployment(SecretValue::new(secret)), + scope: deployment(scope.to_string()), + authority: Some(deployment(authority.to_string())), + }) + .unwrap() + } + + #[derive(Debug, Default)] + struct RecordingTokenClient { + requests: Mutex>, + } + + impl HttpClient for RecordingTokenClient { + fn execute_request<'life0, 'life1, 'async_trait>( + &'life0 self, + request: &'life1 Request, + ) -> std::pin::Pin< + Box> + Send + 'async_trait>, + > + where + 'life0: 'async_trait, + 'life1: 'async_trait, + Self: 'async_trait, + { + Box::pin(async move { + let body = Bytes::from(request.body()); + self.requests.lock().unwrap().push(( + request.url().to_string(), + String::from_utf8(body.to_vec()).unwrap(), + )); + Ok(AsyncRawResponse::from_bytes( + StatusCode::Ok, + Headers::new(), + r#"{"token_type":"Bearer","expires_in":3600,"ext_expires_in":3600,"access_token":"native-token"}"#, + )) + }) + } + } + + #[tokio::test] + async fn client_secret_uses_sdk_protocol_and_reuses_cached_credential() { + let transport = Arc::new(RecordingTokenClient::default()); + let acquirer = + NativeAzureTokenAcquirer::with_transport(4, Transport::new(transport.clone())); + let request = client_secret_request( + "tenant", + "client", + "secret", + "https://service.test/.default", + "https://login.test", + ); + + let first = acquirer.acquire(request.clone()).await.unwrap(); + let second = acquirer.acquire(request).await.unwrap(); + + assert_eq!(first.secret().expose(), "native-token"); + assert_eq!(second.secret().expose(), "native-token"); + let requests = transport.requests.lock().unwrap(); + assert_eq!(requests.len(), 1); + assert_eq!(requests[0].0, "https://login.test/tenant/oauth2/v2.0/token"); + assert!(requests[0].1.contains("client_id=client")); + assert!(requests[0].1.contains("client_secret=secret")); + assert!( + requests[0] + .1 + .contains("scope=https%3A%2F%2Fservice.test%2F.default") + ); + } + + #[tokio::test] + async fn credential_provider_cache_isolates_every_client_secret_identity_field() { + let transport = Arc::new(RecordingTokenClient::default()); + let acquirer = + NativeAzureTokenAcquirer::with_transport(16, Transport::new(transport.clone())); + let request = client_secret_request; + let base = request("tenant", "client", "secret", "scope", "https://login.test"); + let variants = [ + base.clone(), + request( + "other-tenant", + "client", + "secret", + "scope", + "https://login.test", + ), + request( + "tenant", + "other-client", + "secret", + "scope", + "https://login.test", + ), + request( + "tenant", + "client", + "other-secret", + "scope", + "https://login.test", + ), + request( + "tenant", + "client", + "secret", + "other-scope", + "https://login.test", + ), + request( + "tenant", + "client", + "secret", + "scope", + "https://other-login.test", + ), + ]; + + acquirer.acquire(base.clone()).await.unwrap(); + acquirer.acquire(base).await.unwrap(); + for request in variants.into_iter().skip(1) { + acquirer.acquire(request).await.unwrap(); + } + + assert_eq!(transport.requests.lock().unwrap().len(), 6); + } + + #[test] + fn request_authority_requires_request_owned_client_secret_identity() { + let error = ValidatedAzureRequest::new(sourced_client_secret( + InputSource::Deployment, + InputSource::Request, + "https://login.example", + )) + .unwrap_err(); + + assert!(matches!( + error, + crate::AuthError::Configuration( + crate::auth::error::AuthConfigurationError::MixedAzureCredentialSources + ) + )); + } + + #[test] + fn request_owned_client_secret_identity_can_select_custom_authority() { + let request = ValidatedAzureRequest::new(sourced_client_secret( + InputSource::Request, + InputSource::Request, + "https://login.example", + )) + .unwrap(); + + assert_eq!(request.credential_source(), InputSource::Request); + } + + #[test] + fn authority_is_restricted_to_an_https_origin() { + for authority in [ + "http://login.example", + "https://user@login.example", + "https://login.example/tenant", + "https://login.example?target=other", + ] { + let error = ValidatedAzureRequest::new(sourced_client_secret( + InputSource::Deployment, + InputSource::Deployment, + authority, + )) + .unwrap_err(); + assert!(matches!( + error, + crate::AuthError::Configuration( + crate::auth::error::AuthConfigurationError::InvalidAzureAuthority + ) + )); + } + } +} diff --git a/litellm-rust/crates/core/src/providers/azure_ai/auth/resolve.rs b/litellm-rust/crates/core/src/providers/azure_ai/auth/resolve.rs new file mode 100644 index 00000000000..025dd4f8740 --- /dev/null +++ b/litellm-rust/crates/core/src/providers/azure_ai/auth/resolve.rs @@ -0,0 +1,683 @@ +use crate::AuthError; +use crate::auth::error::AuthConfigurationError; +use crate::auth::{ + CredentialFileRef, CredentialLookup, CredentialRef, InputSource, ResolvedCredential, + SecretValue, Sourced, TokenProviderHandle, +}; +use std::future::Future; +use std::pin::Pin; +use std::sync::Arc; + +use super::native::{NativeAzureRequest, NativeAzureTokenAcquirer, ValidatedAzureRequest}; +use super::types::{AzureAuthInputs, AzureCredentialType, ConfigValue, DEFAULT_AZURE_SCOPE}; + +const AZURE_AD_TOKEN_ENV: &str = "AZURE_AD_TOKEN"; +const AZURE_TENANT_ID_ENV: &str = "AZURE_TENANT_ID"; +const AZURE_CLIENT_ID_ENV: &str = "AZURE_CLIENT_ID"; +const AZURE_CLIENT_SECRET_ENV: &str = "AZURE_CLIENT_SECRET"; +const AZURE_SCOPE_ENV: &str = "AZURE_SCOPE"; +const AZURE_AUTHORITY_HOST_ENV: &str = "AZURE_AUTHORITY_HOST"; +const AZURE_CREDENTIAL_ENV: &str = "AZURE_CREDENTIAL"; +const AZURE_FEDERATED_TOKEN_FILE_ENV: &str = "AZURE_FEDERATED_TOKEN_FILE"; + +#[derive(Clone, Debug)] +pub(crate) enum AzureCredentialPlan { + Supplied(Sourced), + Caller(TokenProviderHandle), + Oidc { + reference: Sourced, + tenant_id: Sourced, + client_id: Sourced, + scope: Sourced, + authority: Option>, + }, + Native(ValidatedAzureRequest), + Chain(Vec), + Missing, +} + +/// Rust counterpart to Python's `get_azure_ad_token`, not `BaseAzureLLM`. +pub(crate) struct AzureAuthService { + native: Arc, +} + +trait AzureTokenAcquirer: Send + Sync { + fn acquire( + &self, + request: ValidatedAzureRequest, + ) -> Pin> + Send + '_>>; +} + +impl AzureTokenAcquirer for NativeAzureTokenAcquirer { + fn acquire( + &self, + request: ValidatedAzureRequest, + ) -> Pin> + Send + '_>> { + Box::pin(NativeAzureTokenAcquirer::acquire(self, request)) + } +} + +impl Default for AzureAuthService { + fn default() -> Self { + Self { + native: Arc::new(NativeAzureTokenAcquirer::default()), + } + } +} + +impl AzureAuthService { + #[cfg(test)] + fn with_acquirer(native: Arc) -> Self { + Self { native } + } + + pub(crate) async fn get_azure_ad_token( + &self, + inputs: &AzureAuthInputs, + env_lookup: &(dyn Fn(&str) -> Option + Sync), + ) -> Result>, AuthError> { + match select_auth_plan(inputs, env_lookup)? { + AzureCredentialPlan::Supplied(credential) => Ok(Some(credential)), + AzureCredentialPlan::Caller(caller) => { + let credential = caller.acquire().await?; + if credential.secret().expose().is_empty() { + return Err(AuthError::EmptyAzureToken); + } + Ok(Some(Sourced::new(credential, InputSource::Deployment))) + } + AzureCredentialPlan::Oidc { + reference, + tenant_id, + client_id, + scope, + authority, + } => { + let assertion = resolve_reference(inputs, env_lookup, reference.value()) + .await? + .ok_or(AuthError::UnresolvedOidcReference)?; + let request = ValidatedAzureRequest::new(NativeAzureRequest::ClientAssertion { + tenant_id, + client_id, + assertion: Sourced::new(assertion, reference.source()), + assertion_identity: format!("{:?}", reference.value()), + scope, + authority, + })?; + let source = request.credential_source(); + self.native + .acquire(request) + .await + .map(|credential| Sourced::new(credential, source)) + .map(Some) + } + AzureCredentialPlan::Native(request) => { + let source = request.credential_source(); + self.native + .acquire(request) + .await + .map(|credential| Some(Sourced::new(credential, source))) + } + AzureCredentialPlan::Chain(requests) => { + let mut failures = Vec::new(); + for request in requests { + let source = request.credential_source(); + match self.native.acquire(request).await { + Ok(credential) => return Ok(Some(Sourced::new(credential, source))), + Err(error) => failures.push(error), + } + } + Err(AuthError::CredentialChain(failures)) + } + AzureCredentialPlan::Missing => Ok(None), + } + } +} + +pub(crate) fn select_auth_plan( + inputs: &AzureAuthInputs, + env_lookup: &dyn Fn(&str) -> Option, +) -> Result { + let token = configured_secret(&inputs.azure_ad_token, AZURE_AD_TOKEN_ENV, env_lookup); + let tenant_id = configured_string(&inputs.tenant_id, AZURE_TENANT_ID_ENV, env_lookup); + let client_id = configured_string(&inputs.client_id, AZURE_CLIENT_ID_ENV, env_lookup); + let client_secret = + configured_secret(&inputs.client_secret, AZURE_CLIENT_SECRET_ENV, env_lookup); + let scope = configured_string(&inputs.azure_scope, AZURE_SCOPE_ENV, env_lookup) + .unwrap_or_else(|| Sourced::new(DEFAULT_AZURE_SCOPE.to_string(), InputSource::Environment)); + let authority = configured_string( + &inputs.azure_authority_host, + AZURE_AUTHORITY_HOST_ENV, + env_lookup, + ); + let selector = configured_string(&inputs.azure_credential, AZURE_CREDENTIAL_ENV, env_lookup) + .map(|value| { + value + .value() + .parse::() + .map(|selector| Sourced::new(selector, value.source())) + }) + .transpose() + .map_err(|_| AuthError::Configuration(AuthConfigurationError::InvalidAzureSelector))?; + let federated_token_file = configured_string( + &inputs.federated_token_file, + AZURE_FEDERATED_TOKEN_FILE_ENV, + env_lookup, + ); + + if inputs.azure_ad_token_provider.is_none() + && let (Some(tenant_id), Some(client_id), Some(client_secret)) = + (tenant_id.clone(), client_id.clone(), client_secret) + { + return Ok(AzureCredentialPlan::Native(ValidatedAzureRequest::new( + NativeAzureRequest::ClientSecret { + tenant_id, + client_id, + client_secret, + scope, + authority, + }, + )?)); + } + + if let (Some(reference), Some(tenant_id), Some(client_id)) = ( + oidc_reference(&token)?, + tenant_id.clone(), + client_id.clone(), + ) { + return Ok(AzureCredentialPlan::Oidc { + reference, + tenant_id, + client_id, + scope, + authority, + }); + } + + if let Some(caller) = &inputs.azure_ad_token_provider { + return Ok(AzureCredentialPlan::Caller(caller.clone())); + } + + if let Some(token) = token { + return Ok(AzureCredentialPlan::Supplied(token.map(|token| { + ResolvedCredential::AccessToken { + token, + expires_on: None, + } + }))); + } + + if !*inputs.enable_azure_ad_token_refresh.value() && selector.is_none() { + return Ok(AzureCredentialPlan::Missing); + } + + select_native_plan( + selector, + tenant_id, + client_id, + federated_token_file, + scope, + authority, + inputs.enable_azure_ad_token_refresh.source(), + ) +} + +fn select_native_plan( + selector: Option>, + tenant_id: Option>, + client_id: Option>, + federated_token_file: Option>, + scope: Sourced, + authority: Option>, + refresh_source: InputSource, +) -> Result { + let selected = selector.unwrap_or_else(|| { + Sourced::new( + { + if federated_token_file.is_some() { + AzureCredentialType::DefaultAzureCredential + } else if client_id.is_some() { + AzureCredentialType::ManagedIdentityCredential + } else { + AzureCredentialType::DefaultAzureCredential + } + }, + refresh_source, + ) + }); + let selection_source = selected.source(); + + match selected.into_value() { + AzureCredentialType::ClientSecretCredential => Err(AuthError::Configuration( + AuthConfigurationError::MissingClientSecretFields, + )), + AzureCredentialType::WorkloadIdentityCredential => { + Ok(AzureCredentialPlan::Native(ValidatedAzureRequest::new( + workload_request(tenant_id, client_id, federated_token_file, scope, authority)?, + )?)) + } + AzureCredentialType::ManagedIdentityCredential => Ok(AzureCredentialPlan::Native( + ValidatedAzureRequest::new(NativeAzureRequest::ManagedIdentity { + client_id, + scope, + selection_source, + })?, + )), + AzureCredentialType::DefaultAzureCredential => { + let workload = match (tenant_id, client_id.clone(), federated_token_file) { + (Some(tenant_id), Some(client_id), Some(token_file_path)) => { + Some(NativeAzureRequest::WorkloadIdentity { + tenant_id, + client_id, + token_file_path, + scope: scope.clone(), + authority, + }) + } + _ => None, + }; + Ok(AzureCredentialPlan::Chain( + workload + .into_iter() + .chain(std::iter::once(NativeAzureRequest::ManagedIdentity { + client_id, + scope: scope.clone(), + selection_source, + })) + .chain(std::iter::once(NativeAzureRequest::DeveloperTools { + scope, + selection_source, + })) + .map(ValidatedAzureRequest::new) + .collect::, _>>()?, + )) + } + AzureCredentialType::DeploymentIdentityCredential => { + let workload = match (tenant_id, client_id.clone(), federated_token_file) { + (Some(tenant_id), Some(client_id), Some(token_file_path)) => { + Some(NativeAzureRequest::WorkloadIdentity { + tenant_id, + client_id, + token_file_path, + scope: scope.clone(), + authority, + }) + } + _ => None, + }; + let user_assigned = client_id.map(|client_id| NativeAzureRequest::ManagedIdentity { + client_id: Some(client_id), + scope: scope.clone(), + selection_source, + }); + Ok(AzureCredentialPlan::Chain( + workload + .into_iter() + .chain(user_assigned) + .chain(std::iter::once(NativeAzureRequest::ManagedIdentity { + client_id: None, + scope, + selection_source, + })) + .map(ValidatedAzureRequest::new) + .collect::, _>>()?, + )) + } + } +} + +fn workload_request( + tenant_id: Option>, + client_id: Option>, + token_file_path: Option>, + scope: Sourced, + authority: Option>, +) -> Result { + Ok(NativeAzureRequest::WorkloadIdentity { + tenant_id: tenant_id.ok_or(AuthError::Configuration( + AuthConfigurationError::MissingWorkloadTenant, + ))?, + client_id: client_id.ok_or(AuthError::Configuration( + AuthConfigurationError::MissingWorkloadClient, + ))?, + token_file_path: token_file_path.ok_or(AuthError::Configuration( + AuthConfigurationError::MissingWorkloadTokenFile, + ))?, + scope, + authority, + }) +} + +fn configured_string( + configured: &ConfigValue, + environment_name: &str, + env_lookup: &dyn Fn(&str) -> Option, +) -> Option> { + configured + .as_value() + .filter(|value| !value.value().is_empty()) + .cloned() + .or_else(|| { + env_lookup(environment_name) + .filter(|value| !value.is_empty()) + .map(|value| Sourced::new(value, InputSource::Environment)) + }) +} + +fn configured_secret( + configured: &ConfigValue, + environment_name: &str, + env_lookup: &dyn Fn(&str) -> Option, +) -> Option> { + configured + .as_value() + .filter(|value| !value.value().expose().is_empty()) + .cloned() + .or_else(|| { + env_lookup(environment_name) + .filter(|value| !value.is_empty()) + .map(|value| Sourced::new(SecretValue::new(value), InputSource::Environment)) + }) +} + +async fn resolve_reference( + inputs: &AzureAuthInputs, + env_lookup: &(dyn Fn(&str) -> Option + Sync), + reference: &CredentialRef, +) -> Result, AuthError> { + let lookup = match reference { + CredentialRef::Explicit(secret) => return Ok(Some(secret.clone())), + CredentialRef::Env(name) => env_lookup(name) + .filter(|value| !value.is_empty()) + .map(SecretValue::new) + .map_or(CredentialLookup::Missing, CredentialLookup::Found), + CredentialRef::None => return Ok(None), + CredentialRef::File(_) | CredentialRef::Request(_) | CredentialRef::Host(_) => { + let resolver = inputs + .credential_resolver + .as_ref() + .ok_or(AuthError::Configuration( + AuthConfigurationError::MissingHostResolver, + ))?; + resolver.resolve(reference).await? + } + }; + Ok(match lookup { + CredentialLookup::Found(secret) => Some(secret), + CredentialLookup::Missing | CredentialLookup::Declined => None, + }) +} + +fn oidc_reference( + token: &Option>, +) -> Result>, AuthError> { + let Some(token) = token.as_ref() else { + return Ok(None); + }; + let value = token.value().expose(); + if token.source() == InputSource::Request && value.starts_with("oidc/") { + return Err(AuthError::Configuration( + AuthConfigurationError::RequestAzureCredentialReference, + )); + } + if let Some(name) = value.strip_prefix("oidc/env/") { + return non_empty_reference(name, "OIDC environment reference") + .map(CredentialRef::Env) + .map(|reference| Sourced::new(reference, token.source())) + .map(Some); + } + if let Some(name) = value.strip_prefix("oidc/env_path/") { + return non_empty_reference(name, "OIDC environment path reference") + .map(|name| CredentialRef::File(CredentialFileRef::EnvironmentVariable(name))) + .map(|reference| Sourced::new(reference, token.source())) + .map(Some); + } + if let Some(path) = value.strip_prefix("oidc/file/") { + let path = non_empty_reference(path, "OIDC file reference")?; + return Ok(Some(Sourced::new( + CredentialRef::File(CredentialFileRef::Path(path.into())), + token.source(), + ))); + } + if value.starts_with("oidc/") { + return Err(AuthError::Configuration( + AuthConfigurationError::UnsupportedOidcReference, + )); + } + Ok(None) +} + +fn non_empty_reference(value: &str, kind: &str) -> Result { + if value.is_empty() { + return Err(AuthError::Configuration( + AuthConfigurationError::EmptyReference(kind.to_string()), + )); + } + Ok(value.to_string()) +} + +#[cfg(test)] +mod tests { + use std::future::Future; + use std::sync::{Arc, Mutex}; + + use serde_json::json; + + use super::{ + AzureAuthService, AzureCredentialPlan, AzureTokenAcquirer, oidc_reference, + resolve_reference, select_auth_plan, + }; + use crate::AuthError; + use crate::auth::ResolvedCredential; + use crate::auth::{ + CredentialFileRef, CredentialLookup, CredentialLookupFuture, CredentialRef, + CredentialResolver, CredentialResolverHandle, InputSource, SecretValue, Sourced, + }; + use crate::providers::azure_ai::auth::native::ValidatedAzureRequest; + use crate::providers::azure_ai::auth::types::AzureAuthInputs; + + #[derive(Debug)] + struct FileResolver; + + struct ChainAcquirer { + requests: Mutex>, + succeed_on: Option<&'static str>, + } + + impl AzureTokenAcquirer for ChainAcquirer { + fn acquire( + &self, + request: ValidatedAzureRequest, + ) -> std::pin::Pin< + Box> + Send + '_>, + > { + let kind = request.kind(); + self.requests.lock().unwrap().push(kind); + Box::pin(async move { + if self.succeed_on == Some(kind) { + Ok(ResolvedCredential::AccessToken { + token: SecretValue::new("chain-token"), + expires_on: None, + }) + } else { + Err(AuthError::AzureTokenAcquisition(format!("{kind} failed"))) + } + }) + } + } + + impl CredentialResolver for FileResolver { + fn resolve<'a>(&'a self, reference: &'a CredentialRef) -> CredentialLookupFuture<'a> { + Box::pin(async move { + Ok(match reference { + CredentialRef::File(CredentialFileRef::Path(path)) + if path == std::path::Path::new("/run/secrets/assertion") => + { + CredentialLookup::Found(SecretValue::new("rotated-assertion")) + } + _ => CredentialLookup::Declined, + }) + }) + } + } + + #[test] + fn null_and_empty_values_fall_back_to_environment() { + let params = json!({"tenant_id": null, "client_id": "", "client_secret": null}); + let inputs = AzureAuthInputs::from_optional_params(params.as_object().unwrap()).unwrap(); + let plan = select_auth_plan(&inputs, &|name| match name { + "AZURE_TENANT_ID" => Some("tenant".to_string()), + "AZURE_CLIENT_ID" => Some("client".to_string()), + "AZURE_CLIENT_SECRET" => Some("secret".to_string()), + _ => None, + }) + .unwrap(); + + assert!(matches!(plan, AzureCredentialPlan::Native(_))); + } + + #[test] + fn supplied_token_does_not_require_refresh() { + let params = json!({"azure_ad_token": "token"}); + let inputs = AzureAuthInputs::from_optional_params(params.as_object().unwrap()).unwrap(); + + assert!(matches!( + select_auth_plan(&inputs, &|_| None).unwrap(), + AzureCredentialPlan::Supplied(_) + )); + } + + #[test] + fn oidc_reference_is_deferred() { + let params = json!({ + "azure_ad_token": "oidc/env/ASSERTION", + "tenant_id": "tenant", + "client_id": "client" + }); + let inputs = AzureAuthInputs::from_optional_params(params.as_object().unwrap()).unwrap(); + + assert!(matches!( + select_auth_plan(&inputs, &|_| None).unwrap(), + AzureCredentialPlan::Oidc { + reference, + .. + } if reference.value() == &CredentialRef::Env("ASSERTION".to_string()) + )); + } + + #[test] + fn oidc_file_location_is_typed_before_resolution() { + assert_eq!( + oidc_reference(&Some(Sourced::new( + SecretValue::new("oidc/file//run/secrets/assertion"), + InputSource::Deployment, + ))) + .unwrap() + .map(Sourced::into_value), + Some(CredentialRef::File(CredentialFileRef::Path( + "/run/secrets/assertion".into() + ))) + ); + } + + #[test] + fn unsupported_oidc_reference_is_rejected_during_plan_creation() { + let error = oidc_reference(&Some(Sourced::new( + SecretValue::new("oidc/vault/assertion"), + InputSource::Deployment, + ))) + .expect_err("unsupported backend must fail validation"); + + assert!(error.to_string().contains("unsupported OIDC reference")); + } + + #[test] + fn request_oidc_reference_is_rejected_before_lookup() { + let params = json!({ + "azure_ad_token": "oidc/env/ASSERTION", + "tenant_id": "tenant", + "client_id": "client" + }); + let sources = std::collections::BTreeMap::from([ + ("azure_ad_token".to_string(), InputSource::Request), + ("tenant_id".to_string(), InputSource::Request), + ("client_id".to_string(), InputSource::Request), + ]); + let inputs = + AzureAuthInputs::from_sourced_optional_params(params.as_object().unwrap(), &sources) + .unwrap(); + + let error = select_auth_plan(&inputs, &|name| { + assert_ne!(name, "ASSERTION"); + None + }) + .unwrap_err(); + + assert!(matches!( + error, + AuthError::Configuration( + crate::auth::error::AuthConfigurationError::RequestAzureCredentialReference + ) + )); + } + + #[tokio::test] + async fn host_resolver_owns_file_access() { + let inputs = AzureAuthInputs { + credential_resolver: Some(CredentialResolverHandle::new(Arc::new(FileResolver))), + ..AzureAuthInputs::default() + }; + let reference = + CredentialRef::File(CredentialFileRef::Path("/run/secrets/assertion".into())); + + let resolved = resolve_reference(&inputs, &|_| None, &reference) + .await + .unwrap(); + + assert_eq!(resolved, Some(SecretValue::new("rotated-assertion"))); + } + + #[tokio::test] + async fn default_chain_uses_declared_order_and_stops_after_success() { + let acquirer = Arc::new(ChainAcquirer { + requests: Mutex::new(Vec::new()), + succeed_on: Some("developer-tools"), + }); + let service = AzureAuthService::with_acquirer(acquirer.clone()); + let inputs = AzureAuthInputs { + enable_azure_ad_token_refresh: Sourced::new(true, InputSource::Deployment), + ..Default::default() + }; + + let credential = service + .get_azure_ad_token(&inputs, &|_| None) + .await + .unwrap() + .unwrap(); + + assert_eq!(credential.value().secret().expose(), "chain-token"); + assert_eq!( + *acquirer.requests.lock().unwrap(), + ["managed-identity", "developer-tools"] + ); + } + + #[tokio::test] + async fn chain_reports_each_acquisition_failure() { + let acquirer = Arc::new(ChainAcquirer { + requests: Mutex::new(Vec::new()), + succeed_on: None, + }); + let service = AzureAuthService::with_acquirer(acquirer); + let inputs = AzureAuthInputs { + enable_azure_ad_token_refresh: Sourced::new(true, InputSource::Deployment), + ..Default::default() + }; + + let error = service + .get_azure_ad_token(&inputs, &|_| None) + .await + .unwrap_err(); + + assert!(matches!(error, AuthError::CredentialChain(errors) if errors.len() == 2)); + } +} diff --git a/litellm-rust/crates/core/src/providers/azure_ai/auth/types.rs b/litellm-rust/crates/core/src/providers/azure_ai/auth/types.rs new file mode 100644 index 00000000000..f15d526d945 --- /dev/null +++ b/litellm-rust/crates/core/src/providers/azure_ai/auth/types.rs @@ -0,0 +1,195 @@ +use crate::auth::error::AuthConfigurationError; +use serde_json::{Map, Value}; +use std::collections::BTreeMap; +use strum::EnumString; + +use crate::AuthError; +use crate::auth::{ + CredentialResolverHandle, InputSource, SecretValue, Sourced, TokenProviderHandle, +}; + +pub const DEFAULT_AZURE_SCOPE: &str = "https://cognitiveservices.azure.com/.default"; + +#[derive(Clone, Debug, Default, PartialEq, Eq)] +pub enum ConfigValue { + #[default] + Absent, + ExplicitNone(InputSource), + Value(Sourced), +} + +impl ConfigValue { + pub fn as_value(&self) -> Option<&Sourced> { + match self { + Self::Value(value) => Some(value), + Self::Absent | Self::ExplicitNone(_) => None, + } + } +} + +#[derive(Clone, Copy, Debug, EnumString, PartialEq, Eq, Hash)] +#[allow(clippy::enum_variant_names)] +pub enum AzureCredentialType { + ClientSecretCredential, + ManagedIdentityCredential, + DefaultAzureCredential, + DeploymentIdentityCredential, + WorkloadIdentityCredential, +} + +#[derive(Clone, Debug, Default)] +pub struct AzureAuthInputs { + pub azure_ad_token: ConfigValue, + pub azure_ad_token_provider: Option, + pub credential_resolver: Option, + pub tenant_id: ConfigValue, + pub client_id: ConfigValue, + pub client_secret: ConfigValue, + pub azure_scope: ConfigValue, + pub azure_authority_host: ConfigValue, + pub azure_credential: ConfigValue, + pub federated_token_file: ConfigValue, + pub enable_azure_ad_token_refresh: Sourced, +} + +impl AzureAuthInputs { + #[cfg(test)] + pub fn from_optional_params(params: &Map) -> Result { + Self::from_sourced_optional_params(params, &BTreeMap::new()) + } + + pub fn from_sourced_optional_params( + params: &Map, + sources: &BTreeMap, + ) -> Result { + Ok(Self { + azure_ad_token: secret_config(params, sources, "azure_ad_token")?, + azure_ad_token_provider: None, + credential_resolver: None, + tenant_id: string_config(params, sources, "tenant_id")?, + client_id: string_config(params, sources, "client_id")?, + client_secret: secret_config(params, sources, "client_secret")?, + azure_scope: string_config(params, sources, "azure_scope")?, + azure_authority_host: string_config(params, sources, "azure_authority_host")?, + azure_credential: string_config(params, sources, "azure_credential")?, + federated_token_file: string_config(params, sources, "azure_federated_token_file")?, + enable_azure_ad_token_refresh: Sourced::new( + params + .get("enable_azure_ad_token_refresh") + .and_then(Value::as_bool) + .unwrap_or(false), + source_for(sources, "enable_azure_ad_token_refresh"), + ), + }) + } +} + +fn string_config( + params: &Map, + sources: &BTreeMap, + name: &str, +) -> Result, AuthError> { + let source = source_for(sources, name); + match params.get(name) { + None => Ok(ConfigValue::Absent), + Some(Value::Null) => Ok(ConfigValue::ExplicitNone(source)), + Some(Value::String(value)) => Ok(ConfigValue::Value(Sourced::new(value.clone(), source))), + Some(_) => Err(AuthError::Configuration( + AuthConfigurationError::InvalidFieldType(name.to_string()), + )), + } +} + +fn secret_config( + params: &Map, + sources: &BTreeMap, + name: &str, +) -> Result, AuthError> { + Ok(match string_config(params, sources, name)? { + ConfigValue::Absent => ConfigValue::Absent, + ConfigValue::ExplicitNone(source) => ConfigValue::ExplicitNone(source), + ConfigValue::Value(value) => ConfigValue::Value(value.map(SecretValue::new)), + }) +} + +fn source_for(sources: &BTreeMap, name: &str) -> InputSource { + sources.get(name).copied().unwrap_or_default() +} + +#[cfg(test)] +mod tests { + use serde_json::json; + + use std::collections::BTreeMap; + + use super::{AzureAuthInputs, AzureCredentialType, ConfigValue}; + use crate::auth::{InputSource, Sourced}; + + #[test] + fn selector_parsing_is_exact() { + assert_eq!( + "ClientSecretCredential".parse::(), + Ok(AzureCredentialType::ClientSecretCredential) + ); + assert!( + "clientsecretcredential" + .parse::() + .is_err() + ); + } + + #[test] + fn defaults_preserve_absence() { + let inputs = AzureAuthInputs::default(); + + assert_eq!(inputs.tenant_id, ConfigValue::Absent); + assert_eq!(inputs.azure_ad_token, ConfigValue::Absent); + } + + #[test] + fn parsing_distinguishes_null_empty_and_absent() { + let params = json!({"tenant_id": null, "client_id": ""}); + let inputs = AzureAuthInputs::from_optional_params(params.as_object().unwrap()).unwrap(); + + assert_eq!( + inputs.tenant_id, + ConfigValue::ExplicitNone(InputSource::Deployment) + ); + assert_eq!( + inputs.client_id, + ConfigValue::Value(Sourced::new(String::new(), InputSource::Deployment)) + ); + assert_eq!(inputs.client_secret, ConfigValue::Absent); + } + + #[test] + fn parsing_preserves_trusted_input_sources() { + let params = json!({"tenant_id": "tenant", "client_secret": null}); + let sources = BTreeMap::from([ + ("tenant_id".to_string(), InputSource::Request), + ("client_secret".to_string(), InputSource::Request), + ]); + let inputs = + AzureAuthInputs::from_sourced_optional_params(params.as_object().unwrap(), &sources) + .unwrap(); + + assert_eq!( + inputs.tenant_id, + ConfigValue::Value(Sourced::new("tenant".to_string(), InputSource::Request)) + ); + assert_eq!( + inputs.client_secret, + ConfigValue::ExplicitNone(InputSource::Request) + ); + } + + #[test] + fn debug_does_not_expose_secrets() { + let params = json!({"azure_ad_token": "token-value", "client_secret": "secret-value"}); + let inputs = AzureAuthInputs::from_optional_params(params.as_object().unwrap()).unwrap(); + let debug = format!("{inputs:?}"); + + assert!(!debug.contains("token-value")); + assert!(!debug.contains("secret-value")); + } +} diff --git a/litellm-rust/crates/core/src/providers/azure_ai/messages/transformation.rs b/litellm-rust/crates/core/src/providers/azure_ai/messages/transformation.rs index b8ca10461fb..585b34f393f 100644 --- a/litellm-rust/crates/core/src/providers/azure_ai/messages/transformation.rs +++ b/litellm-rust/crates/core/src/providers/azure_ai/messages/transformation.rs @@ -1,3 +1,4 @@ +use crate::auth::error::MissingCredential; use crate::error::Error; use crate::messages::transformation::{AnthropicMessagesProviderConfig, MessagesAuthStrategy}; use crate::messages::types::{ @@ -32,12 +33,7 @@ pub fn resolve_azure_api_key( non_empty(api_key) .map(str::to_string) .or_else(|| env_lookup(AZURE_API_KEY_ENV).filter(|value| !value.trim().is_empty())) - .ok_or_else(|| { - Error::Auth( - "Missing Azure API Key - Set `api_key` or the AZURE_API_KEY environment variable" - .to_string(), - ) - }) + .ok_or_else(|| Error::from(crate::AuthError::from(MissingCredential::AzureApiKey))) } pub fn complete_azure_anthropic_url( @@ -47,13 +43,7 @@ pub fn complete_azure_anthropic_url( let api_base = non_empty(api_base) .map(str::to_string) .or_else(|| env_lookup(AZURE_API_BASE_ENV).filter(|value| !value.trim().is_empty())) - .ok_or_else(|| { - Error::Auth( - "Missing Azure API Base - Set `api_base` or the AZURE_API_BASE environment variable. \ - Expected format: https://.services.ai.azure.com/anthropic" - .to_string(), - ) - })?; + .ok_or_else(|| Error::from(crate::AuthError::from(MissingCredential::AzureApiBase)))?; let api_base = api_base.trim_end_matches('/'); diff --git a/litellm-rust/crates/core/src/providers/azure_ai/mod.rs b/litellm-rust/crates/core/src/providers/azure_ai/mod.rs index 5d13fa93e00..f2d5b679aee 100644 --- a/litellm-rust/crates/core/src/providers/azure_ai/mod.rs +++ b/litellm-rust/crates/core/src/providers/azure_ai/mod.rs @@ -1,2 +1,3 @@ +pub(crate) mod auth; pub mod messages; pub mod ocr; diff --git a/litellm-rust/crates/core/tests/azure_ai_ocr.rs b/litellm-rust/crates/core/tests/azure_ai_ocr.rs index 3d7fb2e54a3..d7d532cfef1 100644 --- a/litellm-rust/crates/core/tests/azure_ai_ocr.rs +++ b/litellm-rust/crates/core/tests/azure_ai_ocr.rs @@ -45,6 +45,28 @@ async fn facade_executes_azure_mistral_with_prepared_auth() { ); } +#[tokio::test] +async fn facade_acquires_supplied_entra_token_for_final_request() { + let (base, seen, server) = mock_server(vec![MockResponse::json(json!({"pages":[]}))]).await; + let mut request = wire_request( + "azure_ai/model", + &base, + json!({"azure_ad_token":"rust-owned-token"}), + ); + request.connection.api_key = None; + + perform_ocr(request).await.unwrap(); + server.await.unwrap(); + + let requests = seen.lock().unwrap(); + assert_eq!(requests.len(), 1); + assert!( + requests[0] + .to_ascii_lowercase() + .contains("authorization: bearer rust-owned-token\r\n") + ); +} + struct ReplaceBodyDocument; impl OcrHooks for ReplaceBodyDocument { diff --git a/litellm-rust/crates/core/tests/ocr.rs b/litellm-rust/crates/core/tests/ocr.rs index d1828dfb816..cecd8869741 100644 --- a/litellm-rust/crates/core/tests/ocr.rs +++ b/litellm-rust/crates/core/tests/ocr.rs @@ -21,6 +21,7 @@ fn request_boundary_selects_mistral_and_rejects_unknown_providers() { .as_object() .unwrap() .clone(), + input_sources: Default::default(), timeout_seconds: None, }; assert!(decode_request(request).is_ok()); @@ -33,6 +34,7 @@ fn request_boundary_selects_mistral_and_rejects_unknown_providers() { custom_llm_provider: Some("unknown".into()), extra_headers: None, optional_params: serde_json::Map::new(), + input_sources: Default::default(), timeout_seconds: None, }) .is_err() diff --git a/litellm-rust/crates/core/tests/ocr/support.rs b/litellm-rust/crates/core/tests/ocr/support.rs index 45047a0e62a..a2e67dffc7d 100644 --- a/litellm-rust/crates/core/tests/ocr/support.rs +++ b/litellm-rust/crates/core/tests/ocr/support.rs @@ -30,6 +30,7 @@ pub(crate) fn wire_request(model: &str, base: &str, options: Value) -> LiteLLMOc custom_llm_provider: None, extra_headers: None, optional_params: options.as_object().unwrap().clone(), + input_sources: Default::default(), timeout_seconds: Some(2.0), }) .unwrap() diff --git a/litellm-rust/crates/python-bridge/src/routes/definition.rs b/litellm-rust/crates/python-bridge/src/routes/definition.rs index bc51647cbad..97313651011 100644 --- a/litellm-rust/crates/python-bridge/src/routes/definition.rs +++ b/litellm-rust/crates/python-bridge/src/routes/definition.rs @@ -225,7 +225,7 @@ mod tests { ( "ocr", "aocr", - "(model, document, api_key=None, api_base=None, custom_llm_provider=None, extra_headers=None, optional_params=None, timeout_seconds=None)", + "(model, document, api_key=None, api_base=None, custom_llm_provider=None, extra_headers=None, optional_params=None, input_sources=None, timeout_seconds=None)", ), ( "transcription", diff --git a/litellm-rust/crates/python-bridge/src/routes/ocr.rs b/litellm-rust/crates/python-bridge/src/routes/ocr.rs index 2e6900f784f..50095e3ebf2 100644 --- a/litellm-rust/crates/python-bridge/src/routes/ocr.rs +++ b/litellm-rust/crates/python-bridge/src/routes/ocr.rs @@ -22,6 +22,12 @@ fn prepare_ocr( timeout_seconds: inputs.timeout_seconds, })?; let optional_params = object_or_empty("optional_params", inputs.optional_params)?; + let input_sources = inputs + .input_sources + .map(serde_json::from_value) + .transpose() + .map_err(|error| pyo3::exceptions::PyValueError::new_err(error.to_string()))? + .unwrap_or_default(); Ok(async move { let RouteOptions { @@ -41,6 +47,7 @@ fn prepare_ocr( custom_llm_provider, extra_headers, optional_params, + input_sources, timeout_seconds: timeout.map(|value| value.as_secs_f64()), })?; return litellm_core::ocr::ocr(request) @@ -82,6 +89,8 @@ bridge_route! { extra_headers: Option, #[pyo3(from_py_with = litellm_python_interop::from_py)] optional_params: Option, + #[pyo3(from_py_with = litellm_python_interop::from_py)] + input_sources: Option, timeout_seconds: Option, }, prepare = prepare_ocr, diff --git a/litellm/ocr/main.py b/litellm/ocr/main.py index df3f9d2096b..56bfd98895d 100644 --- a/litellm/ocr/main.py +++ b/litellm/ocr/main.py @@ -10,6 +10,7 @@ import re from collections.abc import Callable, Coroutine, Mapping from dataclasses import dataclass from io import IOBase +from types import MappingProxyType from typing import Any, Final, cast import httpx @@ -19,6 +20,7 @@ from litellm._logging import verbose_logger from litellm.constants import request_timeout from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.llms.azure_ai.ocr.common_utils import ( + is_azure_cohere_parse_model, is_azure_document_intelligence_model, ) from litellm.llms.base_llm.ocr.transformation import ( @@ -52,21 +54,32 @@ class _PreparedOCRRequest: litellm_params: dict[str, object] effective_timeout: float | httpx.Timeout litellm_logging_obj: LiteLLMLoggingObj + caller_supplied_api_key: bool = True + caller_supplied_api_base: bool = True -@dataclass -class _PreparedRustOCRCall: - api_key: str | None - api_base: str | None - headers: dict[str, object] - optional_params: dict[str, object] - - -_RUST_OCR_PROVIDERS: Final = { - "mistral", - "azure_ai", - "vertex_ai", -} +_RUST_OCR_PROVIDERS: Final = frozenset({"mistral", "azure_ai", "vertex_ai"}) +_RUST_OCR_CONFIG_FIELDS: Final = frozenset( + { + "azure_ad_token", + "tenant_id", + "client_id", + "client_secret", + "azure_scope", + "azure_authority_host", + "azure_credential", + "azure_federated_token_file", + "vertex_credentials", + "vertex_ai_credentials", + "vertex_project", + "vertex_ai_project", + "vertex_location", + "vertex_ai_location", + } +) +_RUST_OCR_SECRET_FIELDS: Final = frozenset( + {"azure_ad_token", "client_secret", "azure_federated_token_file", "vertex_credentials", "vertex_ai_credentials"} +) def _prepare_ocr_request( @@ -94,6 +107,7 @@ def _prepare_ocr_request( if doc_type not in ["document_url", "image_url"]: raise ValueError(f"Invalid document type: {doc_type}. Must be 'document_url', 'image_url', or 'file'") + caller_supplied_api_key: Final = api_key is not None caller_supplied_api_base: Final = api_base is not None ( @@ -187,182 +201,256 @@ def _prepare_ocr_request( litellm_params=dict(litellm_params), effective_timeout=effective_timeout, litellm_logging_obj=litellm_logging_obj, + caller_supplied_api_key=caller_supplied_api_key, + caller_supplied_api_base=caller_supplied_api_base, ) -def _rust_ocr_supported(prepared_request: _PreparedOCRRequest) -> bool: - if prepared_request.optional_params.get(OCR_REQUEST_FORMAT_PARAM) == "native": - return False - if not prepared_request.provider_config.supports_rust_bridge(): - return False - return prepared_request.custom_llm_provider in _RUST_OCR_PROVIDERS - - -def _rust_bridge_optional_params( - prepared_request: _PreparedOCRRequest, - resolve_secret: Callable[[str], str | None], -) -> dict[str, object]: - optional_params: Final = dict(prepared_request.optional_params) - if prepared_request.custom_llm_provider == "vertex_ai": - vertex_project: Final = ( - prepared_request.litellm_params.get("vertex_project") - or prepared_request.litellm_params.get("vertex_ai_project") - or litellm.vertex_project - or resolve_secret("VERTEXAI_PROJECT") - ) - vertex_location: Final = ( - prepared_request.litellm_params.get("vertex_location") - or prepared_request.litellm_params.get("vertex_ai_location") - or litellm.vertex_location - or resolve_secret("VERTEXAI_LOCATION") - or resolve_secret("VERTEX_LOCATION") - ) - if vertex_project is not None: - optional_params["vertex_project"] = vertex_project - if vertex_location is not None: - optional_params["vertex_location"] = vertex_location - return optional_params - - -def _rust_bridge_api_base( - prepared_request: _PreparedOCRRequest, - resolve_secret: Callable[[str], str | None], -) -> str | None: - if prepared_request.api_base is not None: - return prepared_request.api_base - if prepared_request.custom_llm_provider == "azure_ai": - if is_azure_document_intelligence_model(prepared_request.model): - return resolve_secret("AZURE_DOCUMENT_INTELLIGENCE_ENDPOINT") - return resolve_secret("AZURE_AI_API_BASE") +def _rust_ocr_provider(request: rust_ocr_bridge.LiteLLMOcrRequest) -> str | None: + if request.custom_llm_provider is not None: + return request.custom_llm_provider + prefix: Final = request.model.partition("/")[0] + if prefix in _RUST_OCR_PROVIDERS: + return prefix + if request.model.startswith("mistral-ocr"): + return "mistral" return None -def _prepare_rust_ocr_call( - prepared_request: _PreparedOCRRequest, - resolve_api_key: Callable[[str], str | None], -) -> _PreparedRustOCRCall: - provider_config: Final = prepared_request.provider_config - api_key_env_var: Final = provider_config.get_api_key_env_var() - resolved_api_key: Final = prepared_request.api_key or ( - resolve_api_key(api_key_env_var) if api_key_env_var is not None else None +def _rust_ocr_supported(request: rust_ocr_bridge.LiteLLMOcrRequest) -> bool: + provider: Final = _rust_ocr_provider(request) + if provider not in _RUST_OCR_PROVIDERS or request.kwargs.get(OCR_REQUEST_FORMAT_PARAM) == "native": + return False + if provider == "azure_ai": + return ( + not is_azure_cohere_parse_model(request.model) + and not callable(request.kwargs.get("azure_ad_token_provider")) + and request.kwargs.get("azure_username") is None + and request.kwargs.get("azure_password") is None + ) + return True + + +def _rust_bridge_optional_params( + request: rust_ocr_bridge.LiteLLMOcrRequest, + resolve_secret: Callable[[str], str | None], +) -> Mapping[str, object]: + optional_params: Final = MappingProxyType( + { + name: value + for name, value in request.kwargs.items() + if (name not in GenericLiteLLMParams.model_fields or name in _RUST_OCR_CONFIG_FIELDS) + and name not in {"litellm_logging_obj", "aocr", "litellm_call_id", "proxy_server_request"} + } ) - resolved_headers: Final = provider_config.validate_environment( - headers=prepared_request.extra_headers or {}, - model=prepared_request.model, - api_key=resolved_api_key, - api_base=prepared_request.api_base, - litellm_params=prepared_request.litellm_params, + provider: Final = _rust_ocr_provider(request) + if provider == "azure_ai" and litellm.enable_azure_ad_token_refresh is True: + return MappingProxyType({**optional_params, "enable_azure_ad_token_refresh": True}) + if provider != "vertex_ai": + return optional_params + project: Final = ( + request.kwargs.get("vertex_project") + or request.kwargs.get("vertex_ai_project") + or litellm.vertex_project + or resolve_secret("VERTEXAI_PROJECT") ) - resolved_complete_url: Final = provider_config.get_complete_url( - api_base=prepared_request.api_base, - model=prepared_request.model, - optional_params=prepared_request.optional_params, - litellm_params=prepared_request.litellm_params, + location: Final = ( + request.kwargs.get("vertex_location") + or request.kwargs.get("vertex_ai_location") + or litellm.vertex_location + or resolve_secret("VERTEXAI_LOCATION") + or resolve_secret("VERTEX_LOCATION") ) - rust_api_base: Final = _rust_bridge_api_base(prepared_request, resolve_api_key) - rust_optional_params: Final = _rust_bridge_optional_params(prepared_request, resolve_api_key) - prepared_request.litellm_logging_obj.pre_call( + credentials: Final = ( + request.kwargs.get("vertex_credentials") + or request.kwargs.get("vertex_ai_credentials") + or resolve_secret("VERTEXAI_CREDENTIALS") + ) + vertex_params: Final = MappingProxyType( + { + name: value + for name, value in ( + ("vertex_project", project), + ("vertex_location", location), + ("vertex_credentials", credentials), + ) + if value is not None + } + ) + return MappingProxyType({**optional_params, **vertex_params}) + + +def _rust_bridge_input_sources( + request: rust_ocr_bridge.LiteLLMOcrRequest, + optional_params: Mapping[str, object], +) -> Mapping[str, str]: + proxy_request: Final = request.kwargs.get("proxy_server_request") + if not isinstance(proxy_request, Mapping): + return MappingProxyType({}) + proxy_request_mapping: Final = cast( # cast-ok: runtime Mapping check loses generic key and value types + Mapping[object, object], proxy_request + ) + body_value: Final = proxy_request_mapping.get("body") + if not isinstance(body_value, Mapping): + return MappingProxyType({}) + body: Final = cast( # cast-ok: runtime Mapping check loses generic key and value types + Mapping[object, object], body_value + ) + credential_fields_value: Final = proxy_request_mapping.get("credential_fields", ()) + credential_fields: Final = ( + frozenset(name for name in credential_fields_value if isinstance(name, str)) + if isinstance(credential_fields_value, (list, tuple, set, frozenset)) + else frozenset() + ) + names: Final = frozenset(optional_params) | frozenset({"api_key", "api_base", "extra_headers"}) + request_sources: Final = MappingProxyType( + {name: "request" for name in names if name in body or name in credential_fields} + ) + if litellm.enable_azure_ad_token_refresh is True and "enable_azure_ad_token_refresh" in optional_params: + return MappingProxyType({**request_sources, "enable_azure_ad_token_refresh": "deployment"}) + return request_sources + + +def _marshal_rust_ocr_request( + request: rust_ocr_bridge.LiteLLMOcrRequest, + resolve_secret: Callable[[str], str | None], +) -> rust_ocr_bridge.LiteLLMOcrRequest: + if not isinstance(request.document, dict): + raise TypeError(f"document must be a dict with 'type' and URL/file field, got {type(request.document)}") + document: Final = ( + convert_file_document_to_url_document(request.document) + if request.document.get("type") == "file" + else request.document + ) + provider: Final = _rust_ocr_provider(request) + api_key: Final = request.api_key or resolve_secret("MISTRAL_API_KEY") if provider == "mistral" else request.api_key + optional_params: Final = _rust_bridge_optional_params(request, resolve_secret) + input_sources: Final = _rust_bridge_input_sources(request, optional_params) + logged_optional_params: Final = MappingProxyType( + {name: "****" if name in _RUST_OCR_SECRET_FIELDS else value for name, value in optional_params.items()} + ) + logged_kwargs: Final = MappingProxyType( + { + name: "****" if name in _RUST_OCR_SECRET_FIELDS else value + for name, value in request.kwargs.items() + if name != "proxy_server_request" + } + ) + logging_obj: Final = cast( # cast-ok: bridge kwargs carry the prepared logging object + LiteLLMLoggingObj, request.kwargs["litellm_logging_obj"] + ) + logging_obj.update_from_kwargs( + kwargs=dict(logged_kwargs), # mutable-ok: logging API requires an owned dict + model=request.model, + optional_params=dict(logged_optional_params), # mutable-ok: logging API requires an owned dict + litellm_params={ + "litellm_call_id": request.kwargs.get("litellm_call_id"), + "api_base": request.api_base, + }, # mutable-ok: legacy logging requires a concrete params dict + custom_llm_provider=provider, + ) + logging_obj.pre_call( input="OCR document processing", - api_key=resolved_api_key, - additional_args={ + api_key=api_key, + additional_args={ # mutable-ok: pre_call mutates the additional_args dict "complete_input_dict": { - "model": prepared_request.model, - "document": prepared_request.document, - **rust_optional_params, - }, - "api_base": resolved_complete_url, - "headers": resolved_headers, + "model": request.model, + "document": document, + **logged_optional_params, + }, # mutable-ok: callbacks consume a JSON-serializable request dict + "api_base": request.api_base or "", + "headers": request.extra_headers or {}, # mutable-ok: logging callbacks consume a concrete headers dict }, ) - return _PreparedRustOCRCall( - api_key=resolved_api_key, - api_base=rust_api_base, - headers=cast(dict[str, object], resolved_headers), - optional_params=rust_optional_params, + return rust_ocr_bridge.LiteLLMOcrRequest( + model=request.model, + document=document, + api_key=api_key, + api_base=request.api_base, + timeout=request.timeout if request.timeout is not None else request_timeout, + custom_llm_provider=request.custom_llm_provider, + extra_headers=request.extra_headers, + kwargs=optional_params, + input_sources=input_sources, ) def _map_rust_ocr_error( error: Exception, - prepared_request: _PreparedOCRRequest, + request: rust_ocr_bridge.LiteLLMOcrRequest, exception_types: tuple[type[BaseException], type[BaseException]] | None, ) -> Exception: - if exception_types is None: + if exception_types is None or not isinstance(error, exception_types[1]): return error - _, upstream_error = exception_types - if not isinstance(error, upstream_error): + provider: Final = _rust_ocr_provider(request) + if provider is None: return error - error_args: Final = cast( # cast-ok: BaseException.args is typed with Any in the standard library stubs + provider_config: Final = ProviderConfigManager.get_provider_ocr_config( + model=request.model.removeprefix(f"{provider}/"), provider=litellm.LlmProviders(provider) + ) + if provider_config is None: + return error + error_args: Final = cast( # cast-ok: Python exceptions expose positional args as a tuple tuple[object, ...], error.args ) - status_value: Final = error_args[0] if error_args else 0 - message_value: Final = error_args[1] if len(error_args) > 1 else str(error) - status: Final = status_value if isinstance(status_value, int) else 0 - message: Final = message_value if isinstance(message_value, str) else str(message_value) - error_factory: Final = cast( # cast-ok: the legacy provider interface leaves callable parameters untyped - Callable[..., Exception], prepared_request.provider_config.get_error_class + status: Final = error_args[0] if error_args and isinstance(error_args[0], int) else 500 + message: Final = str(error_args[1]) if len(error_args) > 1 else str(error) + error_factory: Final = cast( # cast-ok: provider configs expose heterogeneous exception factories + Callable[..., Exception], provider_config.get_error_class ) return error_factory( - error_message=message, - status_code=status or 500, - headers={}, # mutable-ok: provider error factories require a concrete header dict - ) + error_message=message, status_code=status or 500, headers={} + ) # mutable-ok: provider error factories require a concrete headers dict def _run_rust_ocr( - prepared_request: _PreparedOCRRequest, + request: rust_ocr_bridge.LiteLLMOcrRequest, resolve_api_key: Callable[[str], str | None], ) -> OCRResponse | None: if rust_ocr_bridge.load_rust_ocr() is None: return None - prepared: Final = _prepare_rust_ocr_call( - prepared_request=prepared_request, - resolve_api_key=resolve_api_key, - ) + marshalled: Final = _marshal_rust_ocr_request(request, resolve_api_key) + input_sources: Final = marshalled.input_sources try: - rust_response: Final = rust_ocr_bridge.ocr( - model=prepared_request.model, - document=prepared_request.document, - api_key=prepared.api_key, - api_base=prepared.api_base, - custom_llm_provider=prepared_request.custom_llm_provider, - extra_headers=prepared.headers, - optional_params=prepared.optional_params, - timeout=prepared_request.effective_timeout, + response: Final = rust_ocr_bridge.ocr( + model=marshalled.model, + document=dict(marshalled.document), # mutable-ok: PyO3 OCR binding requires a concrete dict + api_key=marshalled.api_key, + api_base=marshalled.api_base, + custom_llm_provider=marshalled.custom_llm_provider, + extra_headers=marshalled.extra_headers, + optional_params=dict(marshalled.kwargs), # mutable-ok: PyO3 OCR binding requires a concrete dict + input_sources=input_sources, + timeout=marshalled.timeout, ) except Exception as error: - raise _map_rust_ocr_error(error, prepared_request, native_exception_types()) from error - if rust_response is None: - return None - return OCRResponse.model_validate(rust_response) + raise _map_rust_ocr_error(error, request, native_exception_types()) from error + return OCRResponse.model_validate(response) if response is not None else None async def _run_rust_aocr( - prepared_request: _PreparedOCRRequest, + request: rust_ocr_bridge.LiteLLMOcrRequest, resolve_api_key: Callable[[str], str | None], ) -> OCRResponse | None: if rust_ocr_bridge.load_rust_aocr() is None: return None - prepared: Final = _prepare_rust_ocr_call( - prepared_request=prepared_request, - resolve_api_key=resolve_api_key, - ) + marshalled: Final = _marshal_rust_ocr_request(request, resolve_api_key) + input_sources: Final = marshalled.input_sources try: - rust_response: Final = await rust_ocr_bridge.aocr( - model=prepared_request.model, - document=prepared_request.document, - api_key=prepared.api_key, - api_base=prepared.api_base, - custom_llm_provider=prepared_request.custom_llm_provider, - extra_headers=prepared.headers, - optional_params=prepared.optional_params, - timeout=prepared_request.effective_timeout, + response: Final = await rust_ocr_bridge.aocr( + model=marshalled.model, + document=dict(marshalled.document), # mutable-ok: PyO3 OCR binding requires a concrete dict + api_key=marshalled.api_key, + api_base=marshalled.api_base, + custom_llm_provider=marshalled.custom_llm_provider, + extra_headers=marshalled.extra_headers, + optional_params=dict(marshalled.kwargs), # mutable-ok: PyO3 OCR binding requires a concrete dict + input_sources=input_sources, + timeout=marshalled.timeout, ) except Exception as error: - raise _map_rust_ocr_error(error, prepared_request, native_exception_types()) from error - if rust_response is None: - return None - return OCRResponse.model_validate(rust_response) + raise _map_rust_ocr_error(error, request, native_exception_types()) from error + return OCRResponse.model_validate(response) if response is not None else None @client @@ -444,7 +532,29 @@ async def aocr( "extra_headers": extra_headers, "kwargs": kwargs, } + request: Final = rust_ocr_bridge.LiteLLMOcrRequest( + model=model, + document=document, + api_key=api_key, + api_base=api_base, + timeout=timeout, + custom_llm_provider=custom_llm_provider, + extra_headers=extra_headers, + kwargs=kwargs, + ) try: + if rust_enabled() and _rust_ocr_supported(request): + from litellm.secret_managers.main import get_secret_str + + rust_response: Final = await _run_rust_aocr( + request=request, + resolve_api_key=get_secret_str, + ) + if rust_response is None: + verbose_logger.debug("Async Rust OCR bridge unavailable; falling back to Python path") + else: + return rust_response + prepared: Final = _prepare_ocr_request( model=model, document=document, @@ -459,18 +569,6 @@ async def aocr( custom_llm_provider = prepared.custom_llm_provider completion_kwargs.update({"model": model, "custom_llm_provider": custom_llm_provider}) - if _rust_ocr_supported(prepared) and rust_enabled(): - from litellm.secret_managers.main import get_secret_str - - rust_response: Final = await _run_rust_aocr( - prepared_request=prepared, - resolve_api_key=get_secret_str, - ) - if rust_response is None: - verbose_logger.debug("Async Rust OCR bridge unavailable; falling back to Python path") - else: - return rust_response - response = base_llm_http_handler.ocr( model=prepared.model, document=prepared.document, @@ -494,9 +592,11 @@ async def aocr( return response except Exception as e: + error_provider: Final = custom_llm_provider or _rust_ocr_provider(request) + error_model: Final = model.removeprefix(f"{error_provider}/") if error_provider else model raise litellm.exception_type( - model=model, - custom_llm_provider=custom_llm_provider, + model=error_model, + custom_llm_provider=error_provider, original_exception=e, completion_kwargs=completion_kwargs, extra_kwargs=kwargs, @@ -714,9 +814,31 @@ def ocr( "extra_headers": extra_headers, "kwargs": kwargs, } + request: Final = rust_ocr_bridge.LiteLLMOcrRequest( + model=model, + document=document, + api_key=api_key, + api_base=api_base, + timeout=timeout, + custom_llm_provider=custom_llm_provider, + extra_headers=extra_headers, + kwargs=kwargs, + ) try: _is_async: Final = kwargs.pop("aocr", False) is True completion_kwargs["aocr"] = _is_async + if rust_enabled() and _rust_ocr_supported(request): + from litellm.secret_managers.main import get_secret_str + + rust_response: Final = _run_rust_ocr( + request=request, + resolve_api_key=get_secret_str, + ) + if rust_response is None: + verbose_logger.debug("Rust OCR bridge unavailable; falling back to Python path") + else: + return rust_response + prepared: Final = _prepare_ocr_request( model=model, document=document, @@ -731,18 +853,6 @@ def ocr( custom_llm_provider = prepared.custom_llm_provider completion_kwargs.update({"model": model, "custom_llm_provider": custom_llm_provider}) - if _rust_ocr_supported(prepared) and rust_enabled(): - from litellm.secret_managers.main import get_secret_str - - rust_response: Final = _run_rust_ocr( - prepared_request=prepared, - resolve_api_key=get_secret_str, - ) - if rust_response is None: - verbose_logger.debug("Rust OCR bridge unavailable; falling back to Python path") - else: - return rust_response - response: Final = base_llm_http_handler.ocr( model=prepared.model, document=prepared.document, @@ -760,9 +870,11 @@ def ocr( return response except Exception as e: + error_provider: Final = custom_llm_provider or _rust_ocr_provider(request) + error_model: Final = model.removeprefix(f"{error_provider}/") if error_provider else model raise litellm.exception_type( - model=model, - custom_llm_provider=custom_llm_provider, + model=error_model, + custom_llm_provider=error_provider, original_exception=e, completion_kwargs=completion_kwargs, extra_kwargs=kwargs, diff --git a/litellm/proxy/litellm_pre_call_utils.py b/litellm/proxy/litellm_pre_call_utils.py index c2805d00c2e..8f7b515c22a 100644 --- a/litellm/proxy/litellm_pre_call_utils.py +++ b/litellm/proxy/litellm_pre_call_utils.py @@ -2016,6 +2016,7 @@ async def add_litellm_data_to_request( "method": request.method, "headers": _logging_safe_headers, "body": None, # filled in post-strip; see below + "credential_fields": tuple(sorted(name for name in _TRANSPORT_ONLY_CREDENTIAL_KEYS if name in data)), "arrival_time": arrival_time, # Track when request arrived at proxy } diff --git a/litellm/rust_bridge/ocr.py b/litellm/rust_bridge/ocr.py index b7fdb5a98ef..db959e76f7c 100644 --- a/litellm/rust_bridge/ocr.py +++ b/litellm/rust_bridge/ocr.py @@ -2,7 +2,8 @@ from __future__ import annotations -from collections.abc import Awaitable +from collections.abc import Awaitable, Mapping +from dataclasses import dataclass from typing import Final, Protocol, cast # noqa: TID251 # native extension exposes dynamically typed callables import httpx @@ -11,6 +12,19 @@ from litellm.rust_bridge.bindings import NativeBinding from litellm.rust_bridge.timeouts import timeout_to_seconds as _timeout_to_seconds +@dataclass(frozen=True, slots=True) +class LiteLLMOcrRequest: + model: str + document: Mapping[str, object] + api_key: str | None + api_base: str | None + timeout: float | httpx.Timeout | None + custom_llm_provider: str | None + extra_headers: dict[str, object] | None + kwargs: Mapping[str, object] + input_sources: Mapping[str, str] | None = None + + class RustOcr(Protocol): def __call__( self, @@ -21,6 +35,7 @@ class RustOcr(Protocol): custom_llm_provider: str | None, extra_headers: dict[str, object] | None, optional_params: dict[str, object], + input_sources: dict[str, str], timeout_seconds: float | None, ) -> dict[str, object]: raise NotImplementedError @@ -36,6 +51,7 @@ class RustAocr(Protocol): custom_llm_provider: str | None, extra_headers: dict[str, object] | None, optional_params: dict[str, object], + input_sources: dict[str, str], timeout_seconds: float | None, ) -> Awaitable[dict[str, object]]: raise NotImplementedError @@ -71,6 +87,7 @@ def ocr( extra_headers: dict[str, object] | None, optional_params: dict[str, object], timeout: float | httpx.Timeout | None, + input_sources: Mapping[str, str] | None = None, ) -> dict[str, object] | None: rust_ocr: Final = load_rust_ocr() if rust_ocr is None: @@ -83,6 +100,7 @@ def ocr( custom_llm_provider=custom_llm_provider, extra_headers=extra_headers, optional_params=optional_params, + input_sources=dict(input_sources or {}), # mutable-ok: native boundary requires a concrete dict timeout_seconds=_timeout_to_seconds(timeout), ) @@ -97,6 +115,7 @@ async def aocr( extra_headers: dict[str, object] | None, optional_params: dict[str, object], timeout: float | httpx.Timeout | None, + input_sources: Mapping[str, str] | None = None, ) -> dict[str, object] | None: rust_aocr: Final = load_rust_aocr() if rust_aocr is None: @@ -109,5 +128,6 @@ async def aocr( custom_llm_provider=custom_llm_provider, extra_headers=extra_headers, optional_params=optional_params, + input_sources=dict(input_sources or {}), # mutable-ok: native boundary requires a concrete dict timeout_seconds=_timeout_to_seconds(timeout), ) diff --git a/tests/test_litellm/ocr/test_ocr_azure_document_intelligence_api_base.py b/tests/test_litellm/ocr/test_ocr_azure_document_intelligence_api_base.py index 0c8b1cc2836..460aff3e8d1 100644 --- a/tests/test_litellm/ocr/test_ocr_azure_document_intelligence_api_base.py +++ b/tests/test_litellm/ocr/test_ocr_azure_document_intelligence_api_base.py @@ -1,20 +1,18 @@ """ -Regression tests for Azure Document Intelligence api_base resolution in OCR. +Regression tests for Azure Document Intelligence api_base ownership in OCR. `azure_ai` exposes two OCR services on one provider; the `doc-intelligence` -sub-route must resolve to `AZURE_DOCUMENT_INTELLIGENCE_ENDPOINT`, not to the -generic `AZURE_AI_API_BASE` fallback that `get_llm_provider` injects. These tests -pin that routing and guard the backwards-compatibility contract that an explicitly -supplied api_base is always honoured. +sub-route must defer environment resolution to Rust, not accept the generic +`AZURE_AI_API_BASE` fallback that `get_llm_provider` injects. An explicitly +supplied api_base is still always honoured. """ from litellm.llms.azure_ai.ocr.common_utils import ( is_azure_document_intelligence_model, ) -from litellm.ocr.main import _prepare_ocr_request, _rust_bridge_api_base +from litellm.ocr.main import _prepare_ocr_request _DOC = {"type": "document_url", "document_url": "https://example.com/doc.pdf"} -_DOC_INTELLIGENCE_ENDPOINT = "https://di.cognitiveservices.azure.com" _AZURE_AI_API_BASE = "https://generic-azure-ai.example.com" @@ -23,13 +21,6 @@ class _FakeLogging: return None -def _resolve_secret(name: str) -> str | None: - return { - "AZURE_DOCUMENT_INTELLIGENCE_ENDPOINT": _DOC_INTELLIGENCE_ENDPOINT, - "AZURE_AI_API_BASE": _AZURE_AI_API_BASE, - }.get(name) - - def _prepare(model: str, api_base: str | None): return _prepare_ocr_request( model=model, @@ -56,15 +47,13 @@ class TestIsAzureDocumentIntelligenceModel: class TestDocIntelligenceApiBaseResolution: def test_generic_azure_ai_base_does_not_hijack_doc_intelligence(self, monkeypatch): - """Without an explicit api_base, the AZURE_AI_API_BASE fallback must not - overwrite the endpoint, so it resolves to the Document Intelligence one.""" + """The generic Azure base must not overwrite Rust-owned DI resolution.""" monkeypatch.setenv("AZURE_AI_API_BASE", _AZURE_AI_API_BASE) monkeypatch.delenv("AZURE_DOCUMENT_INTELLIGENCE_ENDPOINT", raising=False) prepared = _prepare("azure_ai/doc-intelligence/prebuilt-layout", None) assert prepared.api_base is None - assert _rust_bridge_api_base(prepared, _resolve_secret) == _DOC_INTELLIGENCE_ENDPOINT def test_explicit_api_base_is_honoured_for_doc_intelligence(self, monkeypatch): """A caller-supplied api_base must always win, even for doc-intelligence.""" @@ -74,7 +63,6 @@ class TestDocIntelligenceApiBaseResolution: prepared = _prepare("azure_ai/doc-intelligence/prebuilt-layout", custom) assert prepared.api_base == custom - assert _rust_bridge_api_base(prepared, _resolve_secret) == custom def test_generic_azure_ai_base_still_applies_to_mistral_ocr(self, monkeypatch): """Non doc-intelligence azure_ai models keep using AZURE_AI_API_BASE.""" diff --git a/tests/test_litellm/ocr/test_ocr_native_format.py b/tests/test_litellm/ocr/test_ocr_native_format.py index 249fbda713e..5f69708fe91 100644 --- a/tests/test_litellm/ocr/test_ocr_native_format.py +++ b/tests/test_litellm/ocr/test_ocr_native_format.py @@ -4,49 +4,42 @@ providers that don't support a native response must reject it, and the Rust bridge (which only returns the normalized shape) must not serve native requests. """ -import dataclasses -from unittest.mock import MagicMock - import pytest import litellm -from litellm.llms.azure_ai.ocr.cohere_parse_transformation import AzureAICohereParseConfig -from litellm.llms.cohere.ocr.transformation import CohereParseConfig -from litellm.ocr.main import _PreparedOCRRequest, _rust_ocr_supported +from litellm.ocr.main import _rust_ocr_supported +from litellm.rust_bridge.ocr import LiteLLMOcrRequest DOCUMENT = {"type": "document_url", "document_url": "https://example.com/doc.pdf"} -def _prepared(optional_params: dict[str, object]) -> _PreparedOCRRequest: - return _PreparedOCRRequest( - model="doc-intelligence/prebuilt-layout", - document=dict(DOCUMENT), +def _request( + optional_params: dict[str, object], model: str = "azure_ai/doc-intelligence/prebuilt-layout" +) -> LiteLLMOcrRequest: + return LiteLLMOcrRequest( + model=model, + document=DOCUMENT, api_key="fake-key", - api_base="https://example.cognitiveservices.azure.com", - custom_llm_provider="azure_ai", + api_base=None, + custom_llm_provider=None, extra_headers=None, - provider_config=MagicMock(), - optional_params=optional_params, - litellm_params={}, - effective_timeout=60.0, - litellm_logging_obj=MagicMock(), + timeout=60.0, + kwargs=optional_params, ) @pytest.mark.parametrize("optional_params", [{}, {"req_format": "litellm"}]) def test_rust_ocr_serves_default_format(optional_params): - assert _rust_ocr_supported(_prepared(optional_params)) is True + assert _rust_ocr_supported(_request(optional_params)) is True def test_rust_ocr_skipped_for_native_format(): - assert _rust_ocr_supported(_prepared({"req_format": "native"})) is False + assert _rust_ocr_supported(_request({"req_format": "native"})) is False -@pytest.mark.parametrize("provider_config", [CohereParseConfig(), AzureAICohereParseConfig()]) -def test_rust_ocr_skipped_for_configs_without_bridge_support(provider_config): - prepared = dataclasses.replace(_prepared({}), provider_config=provider_config) - - assert _rust_ocr_supported(prepared) is False +@pytest.mark.parametrize("model", ["cohere/cohere-parse", "azure_ai/cohere-parse"]) +def test_rust_ocr_skipped_for_unsupported_models(model): + assert _rust_ocr_supported(_request({}, model)) is False @pytest.mark.asyncio diff --git a/tests/test_litellm/ocr/test_rust_bridge.py b/tests/test_litellm/ocr/test_rust_bridge.py index c34833221cc..3441bd4de34 100644 --- a/tests/test_litellm/ocr/test_rust_bridge.py +++ b/tests/test_litellm/ocr/test_rust_bridge.py @@ -3,7 +3,6 @@ import builtins import importlib import types -from typing import Any import httpx import pytest @@ -59,6 +58,7 @@ class RecordingBridge: custom_llm_provider: str | None, extra_headers: dict[str, object] | None, optional_params: dict[str, object], + input_sources: dict[str, str], timeout_seconds: float | None, ) -> dict[str, object]: self.calls.append( @@ -70,6 +70,7 @@ class RecordingBridge: "custom_llm_provider": custom_llm_provider, "extra_headers": extra_headers, "optional_params": optional_params, + "input_sources": input_sources, "timeout_seconds": timeout_seconds, } ) @@ -91,6 +92,7 @@ class RecordingAsyncBridge: custom_llm_provider: str | None, extra_headers: dict[str, object] | None, optional_params: dict[str, object], + input_sources: dict[str, str], timeout_seconds: float | None, ) -> dict[str, object]: self.calls.append( @@ -102,6 +104,7 @@ class RecordingAsyncBridge: "custom_llm_provider": custom_llm_provider, "extra_headers": extra_headers, "optional_params": optional_params, + "input_sources": input_sources, "timeout_seconds": timeout_seconds, } ) @@ -118,6 +121,7 @@ class RaisingBridge: custom_llm_provider: str | None, extra_headers: dict[str, object] | None, optional_params: dict[str, object], + input_sources: dict[str, str], timeout_seconds: float | None, ) -> dict[str, object]: raise RuntimeError("bridge failed") @@ -133,6 +137,7 @@ class RaisingAsyncBridge: custom_llm_provider: str | None, extra_headers: dict[str, object] | None, optional_params: dict[str, object], + input_sources: dict[str, str], timeout_seconds: float | None, ) -> dict[str, object]: raise RuntimeError("bridge failed") @@ -144,6 +149,9 @@ class RecordingLogging: def __init__(self) -> None: self.pre_call_kwargs: dict[str, object] | None = None + def update_from_kwargs(self, **kwargs: object) -> None: + self.update_kwargs = kwargs + def pre_call( self, *, @@ -158,66 +166,32 @@ class RecordingLogging: } -class FakeOCRConfig: - """A stand-in ``BaseOCRConfig`` that echoes the request it would build.""" - - def __init__(self, api_key_env_var: str = "MISTRAL_API_KEY") -> None: - self.api_key_env_var = api_key_env_var - - def get_api_key_env_var(self) -> str: - return self.api_key_env_var - - def validate_environment( - self, - *, - headers: dict[str, object], - model: str, - api_key: str | None, - api_base: str | None, - litellm_params: dict[str, object], - ) -> dict[str, object]: - return {"Authorization": f"Bearer {api_key}", **headers} - - def get_complete_url( - self, - *, - api_base: str | None, - model: str, - optional_params: dict[str, object], - litellm_params: dict[str, object], - ) -> str: - return f"{api_base or 'https://api.mistral.ai/v1'}/ocr" - - def get_error_class(self, error_message: str, status_code: int, headers: dict[str, str]) -> BaseLLMException: - return BaseLLMException(status_code=status_code, message=error_message, headers=headers) - - -def build_prepared_request( +def build_request( *, logging_obj: RecordingLogging | None = None, - provider_config: FakeOCRConfig | None = None, model: str = "mistral-ocr-latest", document: dict[str, object] = DOCUMENT, api_key: str | None = "sk-test", api_base: str | None = None, - custom_llm_provider: str = "mistral", + custom_llm_provider: str | None = "mistral", extra_headers: dict[str, object] | None = None, optional_params: dict[str, object] | None = None, litellm_params: dict[str, object] | None = None, timeout: float | httpx.Timeout | None = 12.5, -) -> Any: - return ocr_main._PreparedOCRRequest( +) -> rust_bridge.LiteLLMOcrRequest: + return rust_bridge.LiteLLMOcrRequest( model=model, document=document, api_key=api_key, api_base=api_base, custom_llm_provider=custom_llm_provider, extra_headers=extra_headers, - provider_config=provider_config or FakeOCRConfig(), - optional_params=optional_params or {}, - litellm_params=litellm_params or {}, - effective_timeout=timeout, - litellm_logging_obj=logging_obj or RecordingLogging(), + timeout=timeout, + kwargs={ + **(optional_params or {}), + **(litellm_params or {}), + "litellm_logging_obj": logging_obj or RecordingLogging(), + }, ) @@ -425,6 +399,7 @@ def test_bridge_wrapper_forwards_prepared_args_and_wraps_response(): "x-trace-id": "trace-1", }, "optional_params": {"include_image_base64": True, "pages": [0]}, + "input_sources": {}, "timeout_seconds": 12.5, } @@ -456,6 +431,7 @@ async def test_bridge_wrapper_forwards_prepared_async_args_and_wraps_response(): "custom_llm_provider": "vertex_ai", "extra_headers": None, "optional_params": {"vertex_project": "project-1"}, + "input_sources": {}, "timeout_seconds": 42.0, } @@ -467,7 +443,7 @@ def test_run_rust_ocr_prepares_request_and_wraps_response(): rust_bridge._OCR.override(bridge) response = ocr_main._run_rust_ocr( - prepared_request=build_prepared_request( + request=build_request( logging_obj=logging_obj, api_base="https://proxy.internal", extra_headers={"x-trace-id": "trace-1"}, @@ -486,10 +462,10 @@ def test_run_rust_ocr_prepares_request_and_wraps_response(): "api_base": "https://proxy.internal", "custom_llm_provider": "mistral", "extra_headers": { - "Authorization": "Bearer sk-test", "x-trace-id": "trace-1", }, "optional_params": {"include_image_base64": True}, + "input_sources": {}, "timeout_seconds": 12.5, } @@ -499,7 +475,7 @@ def test_rust_upstream_error_uses_ocr_provider_error_mapping(): mapped = ocr_main._map_rust_ocr_error( error, - build_prepared_request(), + build_request(), (RuntimeError, RustUpstreamError), ) @@ -514,7 +490,7 @@ def test_run_rust_ocr_resolves_key_via_secret_manager_when_missing(): rust_bridge._OCR.override(bridge) ocr_main._run_rust_ocr( - prepared_request=build_prepared_request(api_key=None, timeout=None), + request=build_request(api_key=None, timeout=None), resolve_api_key=lambda name: "sk-from-vault" if name == "MISTRAL_API_KEY" else None, ) @@ -530,7 +506,7 @@ def test_run_rust_ocr_prefers_explicit_key_over_resolver(): raise AssertionError(f"resolver should not be called for {name}") ocr_main._run_rust_ocr( - prepared_request=build_prepared_request( + request=build_request( api_key="sk-explicit", timeout=None, ), @@ -540,7 +516,7 @@ def test_run_rust_ocr_prefers_explicit_key_over_resolver(): assert bridge.calls[0]["api_key"] == "sk-explicit" -def test_run_rust_ocr_uses_provider_api_key_env_var(): +def test_run_rust_ocr_uses_mistral_secret_manager_without_provider_config(): bridge = RecordingBridge() resolver_calls = [] litellm.rust(True) @@ -551,16 +527,15 @@ def test_run_rust_ocr_uses_provider_api_key_env_var(): return "sk-provider-env" ocr_main._run_rust_ocr( - prepared_request=build_prepared_request( - provider_config=FakeOCRConfig(api_key_env_var="PROVIDER_OCR_API_KEY"), - model="provider-ocr-model", + request=build_request( + model="mistral-ocr-latest", api_key=None, timeout=None, ), resolve_api_key=_resolver, ) - assert resolver_calls == ["PROVIDER_OCR_API_KEY"] + assert resolver_calls == ["MISTRAL_API_KEY"] assert bridge.calls[0]["api_key"] == "sk-provider-env" @@ -570,7 +545,7 @@ def test_prepare_rust_ocr_call_forwards_vertex_routing_metadata(): rust_bridge._OCR.override(bridge) ocr_main._run_rust_ocr( - prepared_request=build_prepared_request( + request=build_request( custom_llm_provider="vertex_ai", model="mistral-ocr-maas", litellm_params={ @@ -588,6 +563,7 @@ def test_prepare_rust_ocr_call_forwards_vertex_routing_metadata(): "include_image_base64": True, "vertex_project": "project-1", "vertex_location": "us-central1", + "vertex_credentials": "redacted", } @@ -603,7 +579,7 @@ def test_prepare_rust_ocr_call_resolves_vertex_routing_metadata_from_secret_mana }.get(name) ocr_main._run_rust_ocr( - prepared_request=build_prepared_request( + request=build_request( custom_llm_provider="vertex_ai", model="mistral-ocr-maas", timeout=None, @@ -615,42 +591,189 @@ def test_prepare_rust_ocr_call_resolves_vertex_routing_metadata_from_secret_mana assert bridge.calls[0]["optional_params"]["vertex_location"] == "us-east5" -def test_prepare_rust_ocr_call_resolves_azure_ai_api_base_from_secret_manager(): +def test_prepare_rust_ocr_call_defers_azure_environment_resolution_to_rust(): bridge = RecordingBridge() litellm.rust(True) rust_bridge._OCR.override(bridge) ocr_main._run_rust_ocr( - prepared_request=build_prepared_request( + request=build_request( custom_llm_provider="azure_ai", model="pixtral-12b-2409", + api_key=None, api_base=None, timeout=None, ), - resolve_api_key=lambda name: "https://azure.example.com" if name == "AZURE_AI_API_BASE" else None, + resolve_api_key=lambda name: pytest.fail(f"Python resolved Azure secret {name}"), ) - assert bridge.calls[0]["api_base"] == "https://azure.example.com" + assert bridge.calls[0]["api_base"] is None + assert bridge.calls[0]["api_key"] is None + assert bridge.calls[0]["extra_headers"] is None -def test_prepare_rust_ocr_call_resolves_document_intelligence_endpoint(): +def test_prepare_rust_ocr_call_defers_document_intelligence_environment_to_rust(): bridge = RecordingBridge() litellm.rust(True) rust_bridge._OCR.override(bridge) ocr_main._run_rust_ocr( - prepared_request=build_prepared_request( + request=build_request( custom_llm_provider="azure_ai", model="doc-intelligence/prebuilt-layout", api_base=None, timeout=None, ), - resolve_api_key=lambda name: ( - "https://document-intelligence.example.com" if name == "AZURE_DOCUMENT_INTELLIGENCE_ENDPOINT" else None - ), + resolve_api_key=lambda name: pytest.fail(f"Python resolved Azure secret {name}"), ) - assert bridge.calls[0]["api_base"] == "https://document-intelligence.example.com" + assert bridge.calls[0]["api_base"] is None + + +def test_prepare_rust_ocr_call_forwards_raw_azure_auth_inputs(): + bridge = RecordingBridge() + litellm.rust(True) + rust_bridge._OCR.override(bridge) + + ocr_main._run_rust_ocr( + request=build_request( + custom_llm_provider="azure_ai", + model="pixtral-12b-2409", + api_key=None, + api_base="https://azure.example.com", + extra_headers={"x-trace-id": "trace-1"}, + litellm_params={ + "azure_ad_token": "entra-token", + "tenant_id": "tenant", + "client_id": "client", + "client_secret": "secret", + "azure_scope": "scope", + "azure_authority_host": "https://login.example.com", + "azure_credential": "ClientSecretCredential", + "azure_federated_token_file": "/token", + }, + timeout=None, + ), + resolve_api_key=lambda name: pytest.fail(f"Python resolved Azure secret {name}"), + ) + + call = bridge.calls[0] + assert call["api_key"] is None + assert call["api_base"] == "https://azure.example.com" + assert call["extra_headers"] == {"x-trace-id": "trace-1"} + assert call["optional_params"] == { + "azure_ad_token": "entra-token", + "tenant_id": "tenant", + "client_id": "client", + "client_secret": "secret", + "azure_scope": "scope", + "azure_authority_host": "https://login.example.com", + "azure_credential": "ClientSecretCredential", + "azure_federated_token_file": "/token", + } + assert call["input_sources"] == {} + + +def test_prepare_rust_ocr_call_preserves_proxy_input_sources(): + bridge = RecordingBridge() + litellm.rust(True) + rust_bridge._OCR.override(bridge) + request_values = { + "tenant_id": "tenant", + "client_id": "client", + "client_secret": "secret", + "azure_authority_host": "https://login.example.com", + "api_base": "https://azure.example.com", + } + + ocr_main._run_rust_ocr( + request=build_request( + custom_llm_provider="azure_ai", + model="pixtral-12b-2409", + api_key="request-key", + api_base="https://azure.example.com", + litellm_params={ + "tenant_id": "tenant", + "client_id": "client", + "client_secret": "secret", + "azure_authority_host": "https://login.example.com", + "proxy_server_request": {"body": request_values, "credential_fields": ("api_key",)}, + }, + ), + resolve_api_key=lambda _name: None, + ) + + assert bridge.calls[0]["input_sources"] == { + **{name: "request" for name in request_values}, + "api_key": "request", + } + + +def test_rust_ocr_logging_redacts_azure_credentials(): + bridge = RecordingBridge() + logging_obj = RecordingLogging() + litellm.rust(True) + rust_bridge._OCR.override(bridge) + + ocr_main._run_rust_ocr( + request=build_request( + logging_obj=logging_obj, + custom_llm_provider="azure_ai", + model="pixtral-12b-2409", + api_key=None, + litellm_params={"azure_ad_token": "token", "client_secret": "secret"}, + ), + resolve_api_key=lambda _name: None, + ) + + assert logging_obj.update_kwargs["optional_params"] == { + "azure_ad_token": "****", + "client_secret": "****", + } + assert logging_obj.pre_call_kwargs is not None + additional_args = logging_obj.pre_call_kwargs["additional_args"] + assert isinstance(additional_args, dict) + complete_input = additional_args["complete_input_dict"] + assert isinstance(complete_input, dict) + assert complete_input["azure_ad_token"] == "****" + assert complete_input["client_secret"] == "****" + + +def test_rust_eligibility_rejects_python_only_azure_auth_modes(): + for params in ( + {"azure_ad_token_provider": lambda: "token"}, + {"azure_username": "user"}, + {"azure_password": "password"}, + ): + assert not ocr_main._rust_ocr_supported( + build_request( + custom_llm_provider="azure_ai", + model="pixtral-12b-2409", + litellm_params=params, + ) + ) + + +def test_prepare_rust_ocr_call_forwards_global_azure_refresh(monkeypatch: pytest.MonkeyPatch): + bridge = RecordingBridge() + litellm.rust(True) + rust_bridge._OCR.override(bridge) + monkeypatch.setattr(litellm, "enable_azure_ad_token_refresh", True) + + ocr_main._run_rust_ocr( + request=build_request( + custom_llm_provider="azure_ai", + model="pixtral-12b-2409", + api_key=None, + api_base="https://azure.example.com", + litellm_params={"proxy_server_request": {"body": {"enable_azure_ad_token_refresh": True}}}, + timeout=None, + ), + resolve_api_key=lambda _name: None, + ) + + assert bridge.calls[0]["optional_params"] == {"enable_azure_ad_token_refresh": True} + assert bridge.calls[0]["input_sources"] == {"enable_azure_ad_token_refresh": "deployment"} def test_run_rust_ocr_runs_pre_call_logging(): @@ -660,7 +783,7 @@ def test_run_rust_ocr_runs_pre_call_logging(): rust_bridge._OCR.override(bridge) ocr_main._run_rust_ocr( - prepared_request=build_prepared_request( + request=build_request( logging_obj=logging_obj, api_base="https://api.mistral.ai/v1", extra_headers={"x-trace-id": "trace-1"}, @@ -676,9 +799,8 @@ def test_run_rust_ocr_runs_pre_call_logging(): complete_input = additional_args["complete_input_dict"] assert complete_input["document"] == DOCUMENT assert complete_input["include_image_base64"] is True - assert additional_args["api_base"] == "https://api.mistral.ai/v1/ocr" + assert additional_args["api_base"] == "https://api.mistral.ai/v1" assert additional_args["headers"] == { - "Authorization": "Bearer sk-test", "x-trace-id": "trace-1", } @@ -696,12 +818,11 @@ def test_ocr_routes_to_rust_when_enabled(fake_bridge): assert response.pages[0].markdown == "hello world" assert len(fake_bridge.calls) == 1 call = fake_bridge.calls[0] - assert call["model"] == "mistral-ocr-latest" + assert call["model"] == MODEL assert call["document"] == DOCUMENT assert call["api_key"] == "sk-test" - assert call["custom_llm_provider"] == "mistral" + assert call["custom_llm_provider"] is None assert call["extra_headers"] == { - "Authorization": "Bearer sk-test", "x-trace-id": "trace-1", } assert call["optional_params"].get("include_image_base64") is True @@ -717,8 +838,29 @@ def test_ocr_routes_azure_ai_to_rust_when_enabled(fake_bridge): assert isinstance(response, OCRResponse) assert len(fake_bridge.calls) == 1 - assert fake_bridge.calls[0]["model"] == "pixtral-12b-2409" - assert fake_bridge.calls[0]["custom_llm_provider"] == "azure_ai" + assert fake_bridge.calls[0]["model"] == "azure_ai/pixtral-12b-2409" + assert fake_bridge.calls[0]["custom_llm_provider"] is None + assert fake_bridge.calls[0]["extra_headers"] is None + + +def test_ocr_routes_azure_entra_inputs_to_rust_without_python_auth(fake_bridge): + response = litellm.ocr( + model="azure_ai/pixtral-12b-2409", + document=DOCUMENT, + api_base="https://example.services.ai.azure.com", + azure_ad_token="entra-token", + tenant_id="tenant", + client_id="client", + ) + + assert isinstance(response, OCRResponse) + assert fake_bridge.calls[0]["api_key"] is None + assert fake_bridge.calls[0]["extra_headers"] is None + assert fake_bridge.calls[0]["optional_params"] == { + "azure_ad_token": "entra-token", + "tenant_id": "tenant", + "client_id": "client", + } def test_ocr_rust_path_converts_file_document_before_bridge(fake_bridge): @@ -768,12 +910,11 @@ async def test_aocr_routes_to_async_rust_when_enabled(fake_async_bridge): assert response.pages[0].markdown == "hello world" assert len(fake_async_bridge.calls) == 1 call = fake_async_bridge.calls[0] - assert call["model"] == "mistral-ocr-latest" + assert call["model"] == MODEL assert call["document"] == DOCUMENT assert call["api_key"] == "sk-test" - assert call["custom_llm_provider"] == "mistral" + assert call["custom_llm_provider"] is None assert call["extra_headers"] == { - "Authorization": "Bearer sk-test", "x-trace-id": "trace-1", } assert call["optional_params"].get("include_image_base64") is True @@ -864,3 +1005,137 @@ def test_ocr_provider_configs_expose_api_key_env_vars(): assert AzureDocumentIntelligenceOCRConfig().get_api_key_env_var() == "AZURE_DOCUMENT_INTELLIGENCE_API_KEY" assert VertexAIOCRConfig().get_api_key_env_var() == "VERTEX_AI_API_KEY" assert VertexAIDeepSeekOCRConfig().get_api_key_env_var() == "VERTEX_AI_API_KEY" + + +@pytest.mark.parametrize("asynchronous", [False, True]) +@pytest.mark.asyncio +async def test_rust_receives_unmapped_azure_options(asynchronous, fake_bridge, fake_async_bridge): + from typing import Final + + arguments: Final = { + "model": "azure_ai/doc-intelligence/prebuilt-layout", + "document": DOCUMENT, + "api_key": "test-key", + "pages": [0, 2], + "features": ["languages", "style"], + "provider_extension": {"enabled": True}, + } + if asynchronous: + await litellm.aocr(**arguments) + else: + litellm.ocr(**arguments) + call: Final = (fake_async_bridge if asynchronous else fake_bridge).calls[0] + assert call["model"] == arguments["model"] + assert call["custom_llm_provider"] is None + assert call["extra_headers"] is None + assert call["optional_params"] == { + "pages": [0, 2], + "features": ["languages", "style"], + "provider_extension": {"enabled": True}, + } + + +@pytest.mark.parametrize("enabled", [False, True]) +@pytest.mark.asyncio +async def test_python_fallback_maps_original_options_once(enabled, monkeypatch): + from io import BytesIO + from typing import Final + + class PythonHandler: + def __init__(self): + self.calls = [] + + def ocr(self, **kwargs): + self.calls.append(kwargs) + return OCRResponse(pages=[], model=kwargs["model"]) + + handler: Final = PythonHandler() + monkeypatch.setattr(ocr_main, "base_llm_http_handler", handler) + litellm.rust(enabled) + rust_bridge._OCR.override(None) + rust_bridge._AOCR.override(None) + for asynchronous in (False, True): + file: Final = BytesIO(b"test document") + arguments: Final = { + "model": "azure_ai/doc-intelligence/prebuilt-layout", + "document": {"type": "file", "file": file}, + "api_key": "test-key", + "pages": [0, 2], + } + if asynchronous: + await litellm.aocr(**arguments) + else: + litellm.ocr(**arguments) + assert handler.calls[-1]["optional_params"]["pages"] == "1,3" + assert handler.calls[-1]["document"]["document_url"].endswith("dGVzdCBkb2N1bWVudA==") + assert len(handler.calls) == 2 + + +@pytest.mark.parametrize("asynchronous", [False, True]) +@pytest.mark.parametrize("model", ["mistral/mistral-ocr-latest", "azure_ai/doc-intelligence/prebuilt-read"]) +@pytest.mark.asyncio +async def test_native_public_ocr_matches_python(model, asynchronous): + import json + from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer + from threading import Thread + from typing import Final + from urllib.parse import parse_qsl, urlsplit + + native: Final = rust_bridge_loader.get_native_bridge() + if native is None: + pytest.skip("requires the compiled Rust extension") + calls: Final = [] + + class Handler(BaseHTTPRequestHandler): + def do_POST(self): + body: Final = json.loads(self.rfile.read(int(self.headers["Content-Length"]))) + target: Final = urlsplit(self.path) + calls.append( + ( + target.path, + parse_qsl(target.query), + self.headers.get("Authorization"), + self.headers.get("Ocp-Apim-Subscription-Key"), + body, + ) + ) + payload: Final = ( + {"status": "succeeded", "analyzeResult": {"pages": []}} + if "doc-intelligence" in model + else {"pages": [{"index": 0, "markdown": "hello"}]} + ) + encoded: Final = json.dumps(payload).encode() + self.send_response(200) + self.send_header("Content-Type", "application/json") + self.send_header("Content-Length", str(len(encoded))) + self.end_headers() + self.wfile.write(encoded) + + def log_message(self, *_args): + pass + + server: Final = ThreadingHTTPServer(("127.0.0.1", 0), Handler) + thread: Final = Thread(target=server.serve_forever, daemon=True) + thread.start() + responses: Final = [] + try: + for enabled in (False, True): + litellm.rust(enabled) + arguments: Final = { + "model": model, + "document": {"type": "document_url", "document_url": "data:application/pdf;base64,YWJj"}, + "api_key": "test-key", + "api_base": f"http://127.0.0.1:{server.server_port}", + "pages": [0, 2], + "timeout": 3.0, + } + response: Final = await litellm.aocr(**arguments) if asynchronous else litellm.ocr(**arguments) + responses.append(response.model_dump()) + assert len(calls) == 2 + assert calls[0] == calls[1] + for key in ("model", "pages", "object"): + assert responses[0][key] == responses[1][key] + finally: + server.shutdown() + server.server_close() + thread.join(timeout=3) diff --git a/tests/test_litellm/proxy/test_litellm_pre_call_utils.py b/tests/test_litellm/proxy/test_litellm_pre_call_utils.py index ec9025a5220..37b983d709a 100644 --- a/tests/test_litellm/proxy/test_litellm_pre_call_utils.py +++ b/tests/test_litellm/proxy/test_litellm_pre_call_utils.py @@ -769,6 +769,7 @@ async def test_add_litellm_data_to_request_body_snapshot_excludes_proxy_server_r data = { "model": "gpt-3.5-turbo", "messages": [{"role": "user", "content": "hello"}], + "api_key": "request-key", } user_api_key_dict = UserAPIKeyAuth( @@ -796,6 +797,8 @@ async def test_add_litellm_data_to_request_body_snapshot_excludes_proxy_server_r assert "proxy_server_request" not in snapshot_body, ( "proxy_server_request must be excluded from its own body snapshot to prevent the body from self-referencing" ) + assert "api_key" not in snapshot_body + assert updated["proxy_server_request"]["credential_fields"] == ("api_key",) def test_refresh_proxy_server_request_body_snapshot_picks_up_guardrail_masking(): diff --git a/tests/test_litellm/rust_bridge/native_route_wheel_test.py b/tests/test_litellm/rust_bridge/native_route_wheel_test.py index 5b2af3cf02e..3b70043fada 100644 --- a/tests/test_litellm/rust_bridge/native_route_wheel_test.py +++ b/tests/test_litellm/rust_bridge/native_route_wheel_test.py @@ -194,10 +194,10 @@ def azure_ocr_kwargs(api_base: str) -> dict[str, object]: "api_base": api_base, "custom_llm_provider": "azure_ai", "extra_headers": { - "Authorization": "Bearer prepared-azure-token", "x-test-outcome": "success", "x-test-route": "azure_ocr", }, + "optional_params": {"azure_ad_token": "prepared-azure-token"}, }