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

feat(ocr): add Azure Mistral adapter with native authentication (#40502)

Browse source 
* 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
This commit is contained in:
yujonglee 2026-09-11 16:22:55 -07:00 • committed by GitHub
parent d70e64d973
commit 5e23db8e03
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
39 changed files with 3767 additions and 341 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

4
.github/scripts/verify_linux_native_wheel.py vendored
View file

@ -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),
)

605
litellm-rust/Cargo.lock generated
View file

@ -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"

5
litellm-rust/Cargo.toml
View file

@ -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

6
litellm-rust/crates/ai-gateway/src/io/realtime.rs
View file

@ -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<String, Error> {
.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)`.

6
litellm-rust/crates/ai-gateway/src/io/responses_ws.rs
View file

@ -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<MaybeTlsStream<TcpStream>>;
type UpstreamTx = SplitSink<ResponsesUpstreamWs, Message>;
type UpstreamRx = SplitStream<ResponsesUpstreamWs>;
@ -120,7 +120,7 @@ pub(crate) fn resolve_api_key(api_key: Option<&str>) -> Result<String, Error> {
.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(

6
litellm-rust/crates/core/Cargo.toml
View file

@ -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 }

168
litellm-rust/crates/core/src/auth/credential.rs Normal file
View file

@ -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<Box<dyn Future<Output = Result<CredentialLookup, AuthError>> + 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<dyn CredentialResolver>);
impl CredentialResolverHandle {
pub fn new(resolver: Arc<dyn CredentialResolver>) -> Self {
Self(resolver)
}
pub async fn resolve(&self, reference: &CredentialRef) -> Result<CredentialLookup, AuthError> {
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<CredentialPlanResolution, AuthError> {
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);
}
}

120
litellm-rust/crates/core/src/auth/error.rs Normal file
View file

@ -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::<Vec<_>>().join("; "))]
CredentialChain(Vec<AuthError>),
#[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://<resource-name>.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,
}

86
litellm-rust/crates/core/src/auth/http.rs Normal file
View file

@ -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<Vec<(String, String)>, 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"));
}
}

55
litellm-rust/crates/core/src/auth/mod.rs Normal file
View file

@ -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<T> {
value: T,
source: InputSource,
}
impl<T> Sourced<T> {
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<U>(self, map: impl FnOnce(T) -> U) -> Sourced<U> {
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};

114
litellm-rust/crates/core/src/auth/policy.rs Normal file
View file

@ -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<Vec<(String, String)>, 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"));
}
}

41
litellm-rust/crates/core/src/auth/secret.rs Normal file
View file

@ -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<String>) -> 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"));
}
}

47
litellm-rust/crates/core/src/auth/token.rs Normal file
View file

@ -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<SystemTime>,
},
}
impl ResolvedCredential {
pub fn secret(&self) -> &SecretValue {
match self {
Self::Static(secret) | Self::AccessToken { token: secret, .. } => secret,
}
}
}
pub type TokenFuture<'a> =
Pin<Box<dyn Future<Output = Result<ResolvedCredential, AuthError>> + Send + 'a>>;
pub trait TokenProvider: std::fmt::Debug + Send + Sync {
fn acquire(&self) -> TokenFuture<'_>;
}
#[derive(Clone, Redact)]
pub struct TokenProviderHandle(#[redact(with = "[REDACTED]")] Arc<dyn TokenProvider>);
impl TokenProviderHandle {
pub fn new(caller: Arc<dyn TokenProvider>) -> Self {
Self(caller)
}
pub async fn acquire(&self) -> Result<ResolvedCredential, AuthError> {
self.0.acquire().await
}
}

19
litellm-rust/crates/core/src/error.rs
View file

@ -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<TransportError> for Error {
}
}
impl From<crate::AuthError> 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()

2
litellm-rust/crates/core/src/lib.rs
View file

@ -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;

129
litellm-rust/crates/core/src/ocr/adapters/azure_mistral.rs
View file

@ -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::<MistralOcrParams>(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<String>,
config: &AzureAuthInputs,
env_lookup: &(dyn Fn(&str) -> Option<String> + Sync),
) -> Result<Vec<(String, String)>, 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<AzureAuthService> = 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<String>) -> Option<String> {
@ -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())
);
}

10
litellm-rust/crates/core/src/ocr/types.rs
View file

@ -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<String>,
pub api_key_source: InputSource,
pub api_base: Option<String>,
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<dyn OcrHooks>,
pub litellm_call_id: Option<String>,
pub optional_params: Map<String, Value>,
pub input_sources: BTreeMap<String, InputSource>,
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,
})
}

15
litellm-rust/crates/core/src/ocr/wire.rs
View file

@ -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<Map<String, Value>>,
#[serde(default)]
pub optional_params: Map<String, Value>,
#[serde(default)]
pub input_sources: BTreeMap<String, InputSource>,
pub timeout_seconds: Option<f64>,
}
@ -36,6 +40,9 @@ pub fn is_supported_request(model: &str, custom_llm_provider: Option<&str>) -> b
}
pub fn decode_request(wire: OcrWireRequest) -> Result<LiteLLMOcrRequest, Error> {
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<LiteLLMOcrRequest, Error>
)?;
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<String, InputSource>, name: &str) -> InputSource {
sources.get(name).copied().unwrap_or_default()
}
fn nonblank(value: Option<String>) -> Option<String> {
value
.map(|s| s.trim().to_string())

9
litellm-rust/crates/core/src/providers/anthropic/messages/transformation.rs
View file

@ -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(

43
litellm-rust/crates/core/src/providers/azure_ai/auth/credential_provider_cache.rs Normal file
View file

@ -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<AzureCredentialProviderCacheKey, Arc<dyn TokenCredential>>,
}
impl AzureCredentialProviderCache {
pub(crate) fn new(capacity: u64) -> Self {
Self {
entries: Cache::builder().max_capacity(capacity).build(),
}
}
pub(crate) async fn get_or_create<F>(
&self,
key: AzureCredentialProviderCacheKey,
create: F,
) -> Result<Arc<dyn TokenCredential>, AuthError>
where
F: Future<Output = Result<Arc<dyn TokenCredential>, AuthError>>,
{
self.entries
.try_get_with(key, create)
.await
.map_err(|error| (*error).clone())
}
}

7
litellm-rust/crates/core/src/providers/azure_ai/auth/mod.rs Normal file
View file

@ -0,0 +1,7 @@
mod credential_provider_cache;
mod native;
mod resolve;
mod types;
pub(crate) use resolve::AzureAuthService;
pub(crate) use types::AzureAuthInputs;

702
litellm-rust/crates/core/src/providers/azure_ai/auth/native.rs Normal file
View file

@ -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<String>,
client_id: Sourced<String>,
client_secret: Sourced<SecretValue>,
scope: Sourced<String>,
authority: Option<Sourced<String>>,
},
ClientAssertion {
tenant_id: Sourced<String>,
client_id: Sourced<String>,
assertion: Sourced<SecretValue>,
assertion_identity: String,
scope: Sourced<String>,
authority: Option<Sourced<String>>,
},
WorkloadIdentity {
tenant_id: Sourced<String>,
client_id: Sourced<String>,
token_file_path: Sourced<String>,
scope: Sourced<String>,
authority: Option<Sourced<String>>,
},
ManagedIdentity {
client_id: Option<Sourced<String>>,
scope: Sourced<String>,
selection_source: InputSource,
},
DeveloperTools {
scope: Sourced<String>,
selection_source: InputSource,
},
}
#[derive(Clone, Debug)]
pub(crate) struct ValidatedAzureRequest {
request: NativeAzureRequest,
credential_source: InputSource,
}
impl ValidatedAzureRequest {
pub(crate) fn new(request: NativeAzureRequest) -> Result<Self, AuthError> {
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<azure_core::http::Transport>,
}
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<ResolvedCredential, AuthError> {
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<InputSource, AuthError> {
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<T>(value: &Sourced<T>, optional: Option<&Sourced<String>>) -> bool {
value.source() == InputSource::Request
|| optional.is_some_and(|value| value.source() == InputSource::Request)
}
fn trusted_only(sources: &[InputSource]) -> Result<InputSource, AuthError> {
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<T>() -> Result<T, AuthError> {
Err(AuthError::Configuration(
AuthConfigurationError::MixedAzureCredentialSources,
))
}
fn build_credential(
request: NativeAzureRequest,
transport: Option<azure_core::http::Transport>,
) -> Result<Arc<dyn TokenCredential>, 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<dyn TokenCredential>),
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<dyn TokenCredential>),
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<dyn TokenCredential>),
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<dyn TokenCredential>)
}
NativeAzureRequest::DeveloperTools { .. } => DeveloperToolsCredential::new(None)
.map(|credential| credential as Arc<dyn TokenCredential>),
}
.map_err(|error| {
AuthError::Configuration(AuthConfigurationError::AzureCredentialInitialization(
error.to_string(),
))
})
}
fn client_options(
authority: Option<String>,
transport: Option<azure_core::http::Transport>,
) -> 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<azure_core::http::ClientMethodOptions<'life1>>,
) -> std::pin::Pin<
Box<dyn std::future::Future<Output = azure_core::Result<String>> + 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<T>(value: T) -> Sourced<T> {
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<Vec<(String, String)>>,
}
impl HttpClient for RecordingTokenClient {
fn execute_request<'life0, 'life1, 'async_trait>(
&'life0 self,
request: &'life1 Request,
) -> std::pin::Pin<
Box<dyn std::future::Future<Output = Result<AsyncRawResponse>> + 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
)
));
}
}
}

683
litellm-rust/crates/core/src/providers/azure_ai/auth/resolve.rs Normal file
View file

@ -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<ResolvedCredential>),
Caller(TokenProviderHandle),
Oidc {
reference: Sourced<CredentialRef>,
tenant_id: Sourced<String>,
client_id: Sourced<String>,
scope: Sourced<String>,
authority: Option<Sourced<String>>,
},
Native(ValidatedAzureRequest),
Chain(Vec<ValidatedAzureRequest>),
Missing,
}
/// Rust counterpart to Python's `get_azure_ad_token`, not `BaseAzureLLM`.
pub(crate) struct AzureAuthService {
native: Arc<dyn AzureTokenAcquirer>,
}
trait AzureTokenAcquirer: Send + Sync {
fn acquire(
&self,
request: ValidatedAzureRequest,
) -> Pin<Box<dyn Future<Output = Result<ResolvedCredential, AuthError>> + Send + '_>>;
}
impl AzureTokenAcquirer for NativeAzureTokenAcquirer {
fn acquire(
&self,
request: ValidatedAzureRequest,
) -> Pin<Box<dyn Future<Output = Result<ResolvedCredential, AuthError>> + 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<dyn AzureTokenAcquirer>) -> Self {
Self { native }
}
pub(crate) async fn get_azure_ad_token(
&self,
inputs: &AzureAuthInputs,
env_lookup: &(dyn Fn(&str) -> Option<String> + Sync),
) -> Result<Option<Sourced<ResolvedCredential>>, 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<String>,
) -> Result<AzureCredentialPlan, AuthError> {
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::<AzureCredentialType>()
.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<Sourced<AzureCredentialType>>,
tenant_id: Option<Sourced<String>>,
client_id: Option<Sourced<String>>,
federated_token_file: Option<Sourced<String>>,
scope: Sourced<String>,
authority: Option<Sourced<String>>,
refresh_source: InputSource,
) -> Result<AzureCredentialPlan, AuthError> {
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::<Result<Vec<_>, _>>()?,
))
}
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::<Result<Vec<_>, _>>()?,
))
}
}
}
fn workload_request(
tenant_id: Option<Sourced<String>>,
client_id: Option<Sourced<String>>,
token_file_path: Option<Sourced<String>>,
scope: Sourced<String>,
authority: Option<Sourced<String>>,
) -> Result<NativeAzureRequest, AuthError> {
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<String>,
environment_name: &str,
env_lookup: &dyn Fn(&str) -> Option<String>,
) -> Option<Sourced<String>> {
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<SecretValue>,
environment_name: &str,
env_lookup: &dyn Fn(&str) -> Option<String>,
) -> Option<Sourced<SecretValue>> {
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<String> + Sync),
reference: &CredentialRef,
) -> Result<Option<SecretValue>, 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<Sourced<SecretValue>>,
) -> Result<Option<Sourced<CredentialRef>>, 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<String, AuthError> {
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<Vec<&'static str>>,
succeed_on: Option<&'static str>,
}
impl AzureTokenAcquirer for ChainAcquirer {
fn acquire(
&self,
request: ValidatedAzureRequest,
) -> std::pin::Pin<
Box<dyn Future<Output = Result<ResolvedCredential, AuthError>> + 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));
}
}

195
litellm-rust/crates/core/src/providers/azure_ai/auth/types.rs Normal file
View file

@ -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<T> {
#[default]
Absent,
ExplicitNone(InputSource),
Value(Sourced<T>),
}
impl<T> ConfigValue<T> {
pub fn as_value(&self) -> Option<&Sourced<T>> {
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<SecretValue>,
pub azure_ad_token_provider: Option<TokenProviderHandle>,
pub credential_resolver: Option<CredentialResolverHandle>,
pub tenant_id: ConfigValue<String>,
pub client_id: ConfigValue<String>,
pub client_secret: ConfigValue<SecretValue>,
pub azure_scope: ConfigValue<String>,
pub azure_authority_host: ConfigValue<String>,
pub azure_credential: ConfigValue<String>,
pub federated_token_file: ConfigValue<String>,
pub enable_azure_ad_token_refresh: Sourced<bool>,
}
impl AzureAuthInputs {
#[cfg(test)]
pub fn from_optional_params(params: &Map<String, Value>) -> Result<Self, AuthError> {
Self::from_sourced_optional_params(params, &BTreeMap::new())
}
pub fn from_sourced_optional_params(
params: &Map<String, Value>,
sources: &BTreeMap<String, InputSource>,
) -> Result<Self, AuthError> {
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<String, Value>,
sources: &BTreeMap<String, InputSource>,
name: &str,
) -> Result<ConfigValue<String>, 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<String, Value>,
sources: &BTreeMap<String, InputSource>,
name: &str,
) -> Result<ConfigValue<SecretValue>, 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<String, InputSource>, 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::<AzureCredentialType>(),
Ok(AzureCredentialType::ClientSecretCredential)
);
assert!(
"clientsecretcredential"
.parse::<AzureCredentialType>()
.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"));
}
}

16
litellm-rust/crates/core/src/providers/azure_ai/messages/transformation.rs
View file

@ -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://<resource-name>.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('/');

1
litellm-rust/crates/core/src/providers/azure_ai/mod.rs
View file

@ -1,2 +1,3 @@
pub(crate) mod auth;
pub mod messages;
pub mod ocr;

22
litellm-rust/crates/core/tests/azure_ai_ocr.rs
View file

@ -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 {

2
litellm-rust/crates/core/tests/ocr.rs
View file

@ -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()

1
litellm-rust/crates/core/tests/ocr/support.rs
View file

@ -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()

2
litellm-rust/crates/python-bridge/src/routes/definition.rs
View file

@ -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",

9
litellm-rust/crates/python-bridge/src/routes/ocr.rs
View file

@ -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<serde_json::Value>,
#[pyo3(from_py_with = litellm_python_interop::from_py)]
optional_params: Option<serde_json::Value>,
#[pyo3(from_py_with = litellm_python_interop::from_py)]
input_sources: Option<serde_json::Value>,
timeout_seconds: Option<f64>,
},
prepare = prepare_ocr,

454
litellm/ocr/main.py
View file

@ -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,

1
litellm/proxy/litellm_pre_call_utils.py
View file

@ -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
}

22
litellm/rust_bridge/ocr.py
View file

@ -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),
)

24
tests/test_litellm/ocr/test_ocr_azure_document_intelligence_api_base.py
View file

@ -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."""

41
tests/test_litellm/ocr/test_ocr_native_format.py
View file

@ -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

431
tests/test_litellm/ocr/test_rust_bridge.py
View file

@ -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)

3
tests/test_litellm/proxy/test_litellm_pre_call_utils.py
View file

@ -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():

2
tests/test_litellm/rust_bridge/native_route_wheel_test.py
View file

@ -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"},
}