From 98f05d52d7741cd76243187c5975e641735b20fc Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Fri, 4 Sep 2026 10:45:13 -0700 Subject: [PATCH] feat(ocr): move provider auth into Rust --- basedpyright-code-budget.json | 4 +- litellm-rust/Cargo.lock | 929 +++++++++++++++++- .../crates/ai-gateway/src/ocr/hooks.rs | 13 +- litellm-rust/crates/core/CLAUDE.md | 6 + litellm-rust/crates/core/Cargo.toml | 3 +- litellm-rust/crates/core/src/constants.rs | 6 + .../crates/core/src/messages/prepare.rs | 11 +- .../core/src/messages/transformation.rs | 20 +- .../crates/core/src/ocr/common_utils.rs | 11 +- litellm-rust/crates/core/src/ocr/handler.rs | 15 +- litellm-rust/crates/core/src/ocr/tests.rs | 134 ++- .../crates/core/src/ocr/transformation.rs | 69 +- .../anthropic/messages/transformation.rs | 7 +- .../core/src/providers/azure_ai/auth.rs | 540 ++++++++++ .../azure_ai/messages/transformation.rs | 5 +- .../crates/core/src/providers/azure_ai/mod.rs | 1 + .../providers/azure_ai/ocr/transformation.rs | 167 ++-- .../providers/mistral/ocr/transformation.rs | 2 +- .../core/src/providers/reducto/ocr/tests.rs | 21 +- .../providers/reducto/ocr/transformation.rs | 4 +- .../core/src/providers/vertex_ai/auth.rs | 290 ++++++ .../core/src/providers/vertex_ai/mod.rs | 1 + .../providers/vertex_ai/ocr/transformation.rs | 43 +- .../document_intelligence/transformation.py | 6 +- litellm/ocr/main.py | 66 +- ...ocument_intelligence_ocr_transformation.py | 27 +- .../llms/test_polling_url_origin_match.py | 29 +- tests/test_litellm/ocr/test_rust_bridge.py | 79 +- type-discipline-budget.json | 2 +- 29 files changed, 2282 insertions(+), 229 deletions(-) create mode 100644 litellm-rust/crates/core/src/providers/azure_ai/auth.rs create mode 100644 litellm-rust/crates/core/src/providers/vertex_ai/auth.rs diff --git a/basedpyright-code-budget.json b/basedpyright-code-budget.json index 1b0650c70d8..3b48c1c33d2 100644 --- a/basedpyright-code-budget.json +++ b/basedpyright-code-budget.json @@ -57,7 +57,7 @@ "limit": 5601 }, "reportMissingTypeArgument": { - "limit": 15285 + "limit": 15284 }, "reportMissingTypeStubs": { "limit": 40 @@ -108,7 +108,7 @@ "limit": 38309 }, "reportUnknownParameterType": { - "limit": 19622 + "limit": 19621 }, "reportUnknownVariableType": { "limit": 29846 diff --git a/litellm-rust/Cargo.lock b/litellm-rust/Cargo.lock index 676e9669c81..196ba8e81e3 100644 --- a/litellm-rust/Cargo.lock +++ b/litellm-rust/Cargo.lock @@ -2,6 +2,12 @@ # It is not intended for manual editing. version = 4 +[[package]] +name = "adler2" +version = "2.0.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "320119579fcad9c21884f5c4861d16174d0e06250625266f50fe6898340abefa" + [[package]] name = "aho-corasick" version = "1.1.5" @@ -20,6 +26,15 @@ dependencies = [ "cc", ] +[[package]] +name = "android_system_properties" +version = "0.1.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ae221649c9976a6f6c56ae1facf410f3ddb33cc661c4b7b61020a912d4237fbc" +dependencies = [ + "libc", +] + [[package]] name = "anes" version = "0.1.6" @@ -41,6 +56,29 @@ dependencies = [ "rustversion", ] +[[package]] +name = "async-compression" +version = "0.4.44" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "515a1f282e33d55983c499d7e9e87082e81cbc32974825bf9032f928392d5844" +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" @@ -468,6 +506,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.22.1" @@ -484,6 +574,12 @@ dependencies = [ "vsimd", ] +[[package]] +name = "bitflags" +version = "1.3.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bef38d45163c2f1dde094a7dfd33ccf595c92905c8f8f4fdc18d06fb1037718a" + [[package]] name = "bitflags" version = "2.13.1" @@ -508,6 +604,15 @@ dependencies = [ "hybrid-array", ] +[[package]] +name = "bs58" +version = "0.5.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bf88ba1141d185c399bee5288d850d63b8369520c1eafc32a0430b5b6c287bf4" +dependencies = [ + "tinyvec", +] + [[package]] name = "bumpalo" version = "3.20.3" @@ -525,6 +630,9 @@ name = "bytes" version = "1.12.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "fc652a48c352aef3ea3aed32080501cf3ef6ed5da78602a020c991775b0aff04" +dependencies = [ + "serde", +] [[package]] name = "bytes-utils" @@ -577,6 +685,18 @@ dependencies = [ "rand_core 0.10.1", ] +[[package]] +name = "chrono" +version = "0.4.45" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1aa79e62e7697b8e29b513a68abacf485adcd1fe8284a4316c5ae868e6633327" +dependencies = [ + "iana-time-zone", + "num-traits", + "serde", + "windows-link", +] + [[package]] name = "ciborium" version = "0.2.2" @@ -644,6 +764,33 @@ 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 = "compression-codecs" +version = "0.4.39" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2fe67f2944eef52fc7b106b8c9450d243a88701a0c065f7f57235e76abaed7df" +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" @@ -684,6 +831,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" @@ -778,17 +934,85 @@ dependencies = [ "cmov", ] +[[package]] +name = "darling" +version = "0.23.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "25ae13da2f202d56bd7f91c25fba009e7717a1e4a1cc98a76d844b65ae912e9d" +dependencies = [ + "darling_core", + "darling_macro", +] + +[[package]] +name = "darling_core" +version = "0.23.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9865a50f7c335f53564bb694ef660825eb8610e0a53d3e11bf1b0d3df31e03b0" +dependencies = [ + "ident_case", + "proc-macro2", + "quote", + "strsim", + "syn 2.0.119", +] + +[[package]] +name = "darling_macro" +version = "0.23.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ac3984ec7bd6cfa798e62b4a642426a5be0e68f9401cfc2a01e3fa9ea2fcdb8d" +dependencies = [ + "darling_core", + "quote", + "syn 2.0.119", +] + [[package]] name = "data-encoding" version = "2.11.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "a4ae5f15dda3c708c0ade84bfee31ccab44a3da4f88015ed22f63732abe300c8" +[[package]] +name = "defmt" +version = "1.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e2953bfe4f93bbd20cc71198842756f77d161884c99ebbabc41d80231ded88d1" +dependencies = [ + "bitflags 1.3.2", + "defmt-macros", +] + +[[package]] +name = "defmt-macros" +version = "1.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bad9c72e7ca2137e0dc3813245a0d282fd6daad32fd800af018306a9169b5fe8" +dependencies = [ + "defmt-parser", + "proc-macro2", + "quote", + "syn 2.0.119", +] + +[[package]] +name = "defmt-parser" +version = "1.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "10d60334b3b2e7c9d91ef8150abfb6fa4c1c39ebbcf4a81c2e346aad939fee3e" +dependencies = [ + "thiserror 2.0.19", +] + [[package]] name = "deranged" version = "0.5.8" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "7cd812cc2bc1d69d4764bd80df88b4317eaef9e773c75226407d9bc0876b211c" +dependencies = [ + "serde_core", +] [[package]] name = "digest" @@ -829,6 +1053,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" @@ -841,6 +1071,36 @@ 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 = "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" @@ -853,6 +1113,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" @@ -874,6 +1145,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" @@ -890,6 +1176,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" @@ -931,6 +1228,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", @@ -984,6 +1282,82 @@ version = "0.3.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "e4eba85ea1d0a966a983acd07deee566e67395d2d96b6fb39e62b5a833f1eb0b" +[[package]] +name = "google-cloud-auth" +version = "1.15.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f54aab44c16b8463ae11b165a87c3d484780231f157bb1ed65843d591beb5abd" +dependencies = [ + "async-trait", + "base64", + "bytes", + "chrono", + "google-cloud-gax", + "hex", + "hmac", + "http 1.4.2", + "reqwest 0.13.4", + "rustc_version", + "rustls 0.23.42", + "rustls-pki-types", + "serde", + "serde_json", + "sha2 0.11.0", + "thiserror 2.0.19", + "time", + "tokio", + "url", +] + +[[package]] +name = "google-cloud-gax" +version = "1.14.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c5615cff28ee59cfe52fbb4c11b8b1e77f650296e2ea4f4c2b7757ac6b19e752" +dependencies = [ + "bytes", + "futures", + "google-cloud-rpc", + "google-cloud-wkt", + "http 1.4.2", + "pin-project", + "rand 0.10.2", + "serde", + "serde_json", + "thiserror 2.0.19", + "tokio", + "tokio-stream", +] + +[[package]] +name = "google-cloud-rpc" +version = "1.6.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e2162c08a89118130979ba261080e960e44cdcb2d6e2ab8ca9b1da245285d353" +dependencies = [ + "bytes", + "google-cloud-wkt", + "serde", + "serde_json", + "serde_with", +] + +[[package]] +name = "google-cloud-wkt" +version = "1.7.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7fccf98cfd5481a5f5a285181ab0c62123d7d47cd2bb7299448440649349e4e7" +dependencies = [ + "base64", + "bytes", + "serde", + "serde_json", + "serde_with", + "thiserror 2.0.19", + "time", + "url", +] + [[package]] name = "h2" version = "0.3.27" @@ -996,7 +1370,7 @@ dependencies = [ "futures-sink", "futures-util", "http 0.2.12", - "indexmap", + "indexmap 2.14.0", "slab", "tokio", "tokio-util", @@ -1015,7 +1389,7 @@ dependencies = [ "futures-core", "futures-sink", "http 1.4.2", - "indexmap", + "indexmap 2.14.0", "slab", "tokio", "tokio-util", @@ -1033,6 +1407,12 @@ dependencies = [ "zerocopy", ] +[[package]] +name = "hashbrown" +version = "0.12.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8a9ee70c43aaf417c914396645a0fa852624801b24ebb7ae78fe8272889ac888" + [[package]] name = "hashbrown" version = "0.17.1" @@ -1237,6 +1617,30 @@ dependencies = [ "tracing", ] +[[package]] +name = "iana-time-zone" +version = "0.1.65" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e31bc9ad994ba00e440a8aa5c9ef0ec67d5cb5e5cb0cc7f8b744a35b389cc470" +dependencies = [ + "android_system_properties", + "core-foundation-sys", + "iana-time-zone-haiku", + "js-sys", + "log", + "wasm-bindgen", + "windows-core", +] + +[[package]] +name = "iana-time-zone-haiku" +version = "0.1.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f31827a206f56af32e590ba56d5d2d085f558508192593743f16b2306495269f" +dependencies = [ + "cc", +] + [[package]] name = "icu_collections" version = "2.2.0" @@ -1319,6 +1723,12 @@ dependencies = [ "zerovec", ] +[[package]] +name = "ident_case" +version = "1.0.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b9e0384b61958566e926dc50660321d12159025e767c18e043daf26b70104c39" + [[package]] name = "idna" version = "1.1.0" @@ -1340,6 +1750,17 @@ dependencies = [ "icu_properties", ] +[[package]] +name = "indexmap" +version = "1.9.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bd070e393353796e801d209ad339e89596eb4c8d430d18ede6a1cced8fafbd99" +dependencies = [ + "autocfg", + "hashbrown 0.12.3", + "serde", +] + [[package]] name = "indexmap" version = "2.14.0" @@ -1347,7 +1768,9 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "d466e9454f08e4a911e14806c24e16fba1b4c121d1ea474396f396069cf949d9" dependencies = [ "equivalent", - "hashbrown", + "hashbrown 0.17.1", + "serde", + "serde_core", ] [[package]] @@ -1371,6 +1794,108 @@ version = "1.0.18" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "8f42a60cbdf9a97f5d2305f08a87dc4e09308d1276d28c869c684d7777685682" +[[package]] +name = "jiff" +version = "0.2.35" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "668b7183bd07af9a4885f5c35b0cc5c83c4607a913c16b7e17291832910d2dcc" +dependencies = [ + "defmt", + "jiff-core", + "jiff-static", + "jiff-tzdb-platform", + "log", + "portable-atomic", + "portable-atomic-util", + "serde_core", + "windows-link", +] + +[[package]] +name = "jiff-core" +version = "0.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7feca88439efe53da3754500c1851dedf3cb36c524dd5cf8225cc0794de95d09" +dependencies = [ + "defmt", +] + +[[package]] +name = "jiff-static" +version = "0.2.35" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3a69dcb3a21cfb32ce1cd056169337ca284af0766dd766e7878819b251a49204" +dependencies = [ + "jiff-core", + "proc-macro2", + "quote", + "syn 2.0.119", +] + +[[package]] +name = "jiff-tzdb" +version = "0.1.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "142bd39932ad231f10513df9ab62661fead8719872150b7ad02a2df79f4e141e" + +[[package]] +name = "jiff-tzdb-platform" +version = "0.1.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "875a5a69ac2bab1a891711cf5eccbec1ce0341ea805560dcd90b7a2e925132e8" +dependencies = [ + "jiff-tzdb", +] + +[[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" @@ -1414,7 +1939,7 @@ dependencies = [ "futures-util", "litellm-config", "litellm-core", - "reqwest", + "reqwest 0.12.28", "serde", "serde_json", "sha2 0.10.9", @@ -1446,9 +1971,12 @@ dependencies = [ "aws-sigv4", "aws-smithy-runtime-api", "aws-types", + "azure_core", + "azure_identity", "base64", + "google-cloud-auth", "rand 0.8.7", - "reqwest", + "reqwest 0.12.28", "rstest", "serde", "serde_json", @@ -1471,7 +1999,7 @@ dependencies = [ "litellm-python-interop", "pyo3", "pyo3-async-runtimes", - "reqwest", + "reqwest 0.12.28", "serde", "serde_json", "tokio", @@ -1536,6 +2064,16 @@ dependencies = [ "unicase", ] +[[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" @@ -1605,12 +2143,38 @@ dependencies = [ "winapi", ] +[[package]] +name = "parking" +version = "2.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f38d5652c16fde515bb1ecef450ab0f6a219d619a7274976324d5e377f7dceba" + [[package]] name = "percent-encoding" 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" @@ -1663,6 +2227,15 @@ version = "1.14.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "3d20d5497ef88037a52ff98267d066e7f11fcc5e99bbfbd58a42336193aacec3" +[[package]] +name = "portable-atomic-util" +version = "0.2.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "10ab3eb7f3becc3a1cbc4f2c6f20267996cfc1a6467a873763411b136a122715" +dependencies = [ + "portable-atomic", +] + [[package]] name = "potential_utf" version = "0.1.5" @@ -1812,6 +2385,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", @@ -1933,6 +2507,26 @@ dependencies = [ "crossbeam-utils", ] +[[package]] +name = "ref-cast" +version = "1.0.27" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7e440fb4e4b4147295338efb76001ab9e4efc0e5839df2c47fc5ac2381d365c3" +dependencies = [ + "ref-cast-impl", +] + +[[package]] +name = "ref-cast-impl" +version = "1.0.27" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "92ecd8964f8453721699a1ed72037b0db49ce2f5a5138486ee89bed6f67cdf3a" +dependencies = [ + "proc-macro2", + "quote", + "syn 3.0.0", +] + [[package]] name = "regex" version = "1.13.1" @@ -2013,11 +2607,52 @@ dependencies = [ "url", "wasm-bindgen", "wasm-bindgen-futures", - "wasm-streams", + "wasm-streams 0.4.2", "web-sys", "webpki-roots", ] +[[package]] +name = "reqwest" +version = "0.13.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "219c5811de6525e5416c7d5d53bb656d3afdbc6c5af816e0802bcfa42dbdc1c3" +dependencies = [ + "base64", + "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", + "serde", + "serde_json", + "serde_urlencoded", + "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" @@ -2095,6 +2730,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "3c54fcab019b409d04215d3a17cb438fd7fbf192ee61461f20f4fe18704bc138" dependencies = [ "aws-lc-rs", + "log", "once_cell", "ring", "rustls-pki-types", @@ -2125,6 +2761,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" @@ -2177,6 +2840,30 @@ dependencies = [ "windows-sys 0.61.2", ] +[[package]] +name = "schemars" +version = "0.9.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4cd191f9397d57d581cddd31014772520aa448f65ef991055d7f61582c65165f" +dependencies = [ + "dyn-clone", + "ref-cast", + "serde", + "serde_json", +] + +[[package]] +name = "schemars" +version = "1.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "687274d293b6cdc6e73e0fee520bf2049650090d7164f87672d212a3c530cf4a" +dependencies = [ + "dyn-clone", + "ref-cast", + "serde", + "serde_json", +] + [[package]] name = "sct" version = "0.7.1" @@ -2193,7 +2880,7 @@ version = "3.7.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "b7f4bc775c73d9a02cde8bf7b2ec4c9d12743edf609006c7facc23998404cd1d" dependencies = [ - "bitflags", + "bitflags 2.13.1", "core-foundation", "core-foundation-sys", "libc", @@ -2282,6 +2969,39 @@ dependencies = [ "serde", ] +[[package]] +name = "serde_with" +version = "3.22.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ee78f1fbe43ac4a0e47aadb3dbd357b69eb0d3793e948624cd03dd2750ab1c0a" +dependencies = [ + "base64", + "bs58", + "chrono", + "hex", + "indexmap 1.9.3", + "indexmap 2.14.0", + "jiff", + "schemars 0.9.0", + "schemars 1.2.2", + "serde_core", + "serde_json", + "serde_with_macros", + "time", +] + +[[package]] +name = "serde_with_macros" +version = "3.22.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8705578779c2b6bd90d84d66eb2e206b708b1a4d7b9f17641b293545bf1c7e46" +dependencies = [ + "darling", + "proc-macro2", + "quote", + "syn 2.0.119", +] + [[package]] name = "sha1" version = "0.10.7" @@ -2330,6 +3050,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" @@ -2368,6 +3120,12 @@ version = "1.2.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6ce2be8dc25455e1f91df71bfa12ad37d7af1092ae736f3a6cd0e37bc7810596" +[[package]] +name = "strsim" +version = "0.11.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7da8b5736845d9f2fcb837ea5d9e2628564b3b043a70948a3f0b778838c5fb4f" + [[package]] name = "subtle" version = "2.6.1" @@ -2546,6 +3304,7 @@ dependencies = [ "libc", "mio", "pin-project-lite", + "signal-hook-registry", "socket2 0.6.5", "tokio-macros", "windows-sys 0.61.2", @@ -2582,6 +3341,18 @@ dependencies = [ "tokio", ] +[[package]] +name = "tokio-stream" +version = "0.1.19" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a3d06f0b082ba57c26b79407372e57cf2a1e28124f78e9479fe80322cf53420b" +dependencies = [ + "futures-core", + "pin-project-lite", + "tokio", + "tokio-util", +] + [[package]] name = "tokio-tungstenite" version = "0.24.0" @@ -2626,7 +3397,7 @@ version = "0.25.13+spec-1.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6975367e4d2ef766d86af01ffad14b622fecc8d4357a998fbc4deb6e9bacaf9b" dependencies = [ - "indexmap", + "indexmap 2.14.0", "toml_datetime", "toml_parser", "winnow", @@ -2663,12 +3434,17 @@ version = "0.6.11" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "4cfcf7e2740e6fc6d4d688b4ef00650406bb94adf4731e43c096c3a19fe40840" dependencies = [ - "bitflags", + "async-compression", + "bitflags 2.13.1", "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", @@ -2762,6 +3538,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", + "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", + "bytes", + "dyn-clone", + "futures", + "pin-project", + "rand 0.10.2", + "reqwest 0.13.4", + "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" @@ -2816,6 +3643,7 @@ 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", ] @@ -2925,6 +3753,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" @@ -2945,6 +3786,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" @@ -2985,12 +3835,65 @@ version = "0.4.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "712e227841d057c1ee1cd2fb22fa7e5a5461ae8e48fa2ca79ec42cfc1931183f" +[[package]] +name = "windows-core" +version = "0.62.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b8e83a14d34d0623b51dce9581199302a221863196a1dde71a7663a4c2be9deb" +dependencies = [ + "windows-implement", + "windows-interface", + "windows-link", + "windows-result", + "windows-strings", +] + +[[package]] +name = "windows-implement" +version = "0.60.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "053e2e040ab57b9dc951b72c264860db7eb3b0200ba345b4e4c3b14f67855ddf" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.119", +] + +[[package]] +name = "windows-interface" +version = "0.59.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3f316c4a2570ba26bbec722032c4099d8c8bc095efccdc15688708623367e358" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.119", +] + [[package]] name = "windows-link" version = "0.2.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "f0805222e57f7521d6a62e36fa9163bc891acd422f971defe97d64e70d0a4fe5" +[[package]] +name = "windows-result" +version = "0.4.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7781fa89eaf60850ac3d2da7af8e5242a5ea78d1a11c49bf2910bb5a73853eb5" +dependencies = [ + "windows-link", +] + +[[package]] +name = "windows-strings" +version = "0.5.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7837d08f69c77cf6b07689544538e017c1bfcf57e34b4c0ff58e6c2cd3b37091" +dependencies = [ + "windows-link", +] + [[package]] name = "windows-sys" version = "0.52.0" @@ -3197,6 +4100,12 @@ dependencies = [ "syn 2.0.119", ] +[[package]] +name = "zlib-rs" +version = "0.6.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "34b31d188d9d685a4f9c7b46d6e36631b07058d2cfe190267adce54dc230bf12" + [[package]] name = "zmij" version = "1.0.23" diff --git a/litellm-rust/crates/ai-gateway/src/ocr/hooks.rs b/litellm-rust/crates/ai-gateway/src/ocr/hooks.rs index a05799b6120..00fdf79cd26 100644 --- a/litellm-rust/crates/ai-gateway/src/ocr/hooks.rs +++ b/litellm-rust/crates/ai-gateway/src/ocr/hooks.rs @@ -88,11 +88,14 @@ impl OcrLifecycleHooks { ) -> Result { let config = request.config?; let env_lookup = |key: &str| std::env::var(key).ok(); - let upstream_headers = config.validate_environment( - string_headers(request.extra_headers)?, - request.api_key.as_deref(), - &env_lookup, - )?; + let (upstream_headers, _) = config + .validate_environment( + string_headers(request.extra_headers)?, + request.api_key.as_deref(), + &request.optional_params, + &env_lookup, + ) + .await?; let url = config.complete_url( request.api_base.as_deref(), &request.model, diff --git a/litellm-rust/crates/core/CLAUDE.md b/litellm-rust/crates/core/CLAUDE.md index a973e1311d5..ca9871b571e 100644 --- a/litellm-rust/crates/core/CLAUDE.md +++ b/litellm-rust/crates/core/CLAUDE.md @@ -30,6 +30,12 @@ Not allowed: - Provider-specific branching that belongs in `providers`. - Panics for user/provider-controlled input. +Provider authentication is the narrow exception to the filesystem and global +state rules. Like Bedrock credential discovery, ADC and managed identity may +read SDK-defined credential files or contact metadata services. Credential and +token caches may be process-wide, must be secret-safe, and may contain auth +state only. + ## Typed Contracts (core rule) Trait and function boundaries MUST be strongly typed. No stringly-typed JSON diff --git a/litellm-rust/crates/core/Cargo.toml b/litellm-rust/crates/core/Cargo.toml index 56e16dbe3ca..465e15ea16b 100644 --- a/litellm-rust/crates/core/Cargo.toml +++ b/litellm-rust/crates/core/Cargo.toml @@ -8,12 +8,13 @@ repository.workspace = true [dependencies] async-trait = "0.1" base64.workspace = true +google-cloud-auth = { version = "=1.15.0", default-features = false, features = ["default-rustls-provider"] } rand.workspace = true reqwest.workspace = true serde.workspace = true serde_json.workspace = true thiserror.workspace = true -tokio = { workspace = true, features = ["net", "time"] } +tokio = { workspace = true, features = ["net", "sync", "time"] } tracing.workspace = true tokio = { workspace = true, features = ["sync", "rt"] } tracing-subscriber = { workspace = true, optional = true } diff --git a/litellm-rust/crates/core/src/constants.rs b/litellm-rust/crates/core/src/constants.rs index e97fbc75313..1333ddbf849 100644 --- a/litellm-rust/crates/core/src/constants.rs +++ b/litellm-rust/crates/core/src/constants.rs @@ -34,6 +34,12 @@ pub(crate) const AUDIO_TRANSCRIPTION_TIMEOUT_SECS: u64 = 600; pub(crate) const OCR_DOCUMENT_INTELLIGENCE_POLL_TIMEOUT_SECS: u64 = 120; pub(crate) const OCR_MAX_SAFE_FETCH_REDIRECTS: usize = 10; +pub(crate) const AZURE_AUTH_DEFAULT_SCOPE: &str = "https://cognitiveservices.azure.com/.default"; +pub(crate) const AZURE_AUTH_TOKEN_REFRESH_WINDOW_SECS: u64 = 120; +pub(crate) const VERTEX_AUTH_CLOUD_PLATFORM_SCOPE: &str = + "https://www.googleapis.com/auth/cloud-platform"; +pub(crate) const VERTEX_AUTH_SDK_TOKEN_CACHE_TTL_SECS: u64 = 30; +pub(crate) const VERTEX_AUTH_TOKEN_REFRESH_WINDOW_SECS: u64 = 5; /// `object` field every non-streaming chat completion response carries. pub const CHAT_COMPLETION_OBJECT: &str = "chat.completion"; diff --git a/litellm-rust/crates/core/src/messages/prepare.rs b/litellm-rust/crates/core/src/messages/prepare.rs index ec83d03f535..f216276fe2e 100644 --- a/litellm-rust/crates/core/src/messages/prepare.rs +++ b/litellm-rust/crates/core/src/messages/prepare.rs @@ -1,8 +1,9 @@ +use crate::auth::ResolvedAuth; use crate::error::Error; use crate::routing_utils::provider::{CustomLlmProvider, get_custom_llm_provider}; use super::common_utils::{has_bearer_auth, has_header, messages_provider_config, string_headers}; -use super::transformation::{AnthropicMessagesProviderConfig, MessagesAuthStrategy}; +use super::transformation::AnthropicMessagesProviderConfig; use super::types::{MessagesRequest, ProviderMessagesRequest}; use serde_json::{Map, Value}; @@ -70,13 +71,7 @@ fn validate_environment( || (config.accepts_bearer_auth() && has_bearer_auth(&headers)); if !already_authorized { let api_key = config.resolve_api_key(api_key, env_lookup)?; - let auth_header = match auth_strategy { - MessagesAuthStrategy::Bearer => { - ("authorization".to_string(), format!("Bearer {api_key}")) - } - MessagesAuthStrategy::Header(name) => (name.to_string(), api_key), - }; - headers.push(auth_header); + headers.push(ResolvedAuth::from_credential(auth_strategy, api_key).credential_header()); } for (name, value) in config.default_headers() { diff --git a/litellm-rust/crates/core/src/messages/transformation.rs b/litellm-rust/crates/core/src/messages/transformation.rs index a5904c085a0..395580a72fc 100644 --- a/litellm-rust/crates/core/src/messages/transformation.rs +++ b/litellm-rust/crates/core/src/messages/transformation.rs @@ -1,20 +1,6 @@ use super::types::{AnthropicMessagesRequest, AnthropicMessagesResponse}; use crate::Error; - -#[derive(Clone, Copy, Debug, PartialEq, Eq)] -pub enum MessagesAuthStrategy { - Bearer, - Header(&'static str), -} - -impl MessagesAuthStrategy { - pub fn header_name(self) -> &'static str { - match self { - Self::Bearer => "authorization", - Self::Header(header_name) => header_name, - } - } -} +use crate::auth::AuthHeaderKind; pub trait AnthropicMessagesProviderConfig: Sync { fn complete_url( @@ -30,8 +16,8 @@ pub trait AnthropicMessagesProviderConfig: Sync { env_lookup: &dyn Fn(&str) -> Option, ) -> Result; - fn auth_strategy(&self) -> MessagesAuthStrategy { - MessagesAuthStrategy::Header("x-api-key") + fn auth_strategy(&self) -> AuthHeaderKind { + AuthHeaderKind::Header("x-api-key") } fn accepts_bearer_auth(&self) -> bool { diff --git a/litellm-rust/crates/core/src/ocr/common_utils.rs b/litellm-rust/crates/core/src/ocr/common_utils.rs index 3fa4c0438aa..ea3452752fe 100644 --- a/litellm-rust/crates/core/src/ocr/common_utils.rs +++ b/litellm-rust/crates/core/src/ocr/common_utils.rs @@ -6,6 +6,7 @@ use base64::engine::general_purpose::STANDARD as BASE64_STANDARD; use reqwest::Url; use serde_json::{Map, Value}; +use crate::auth::ResolvedAuth; use crate::constants::{OCR_DOCUMENT_INTELLIGENCE_POLL_TIMEOUT_SECS, OCR_MAX_SAFE_FETCH_REDIRECTS}; use crate::error::Error; use crate::http_utils::truncate_error_body; @@ -322,7 +323,7 @@ pub(super) async fn poll_document_intelligence( client: &reqwest::Client, operation_url: &str, original_url: &str, - headers: &[(String, String)], + auth: &ResolvedAuth, timeout: Option, ) -> Result { if !same_origin(operation_url, original_url) { @@ -343,12 +344,8 @@ pub(super) async fn poll_document_intelligence( ))); } - let mut request_builder = client.get(operation_url); - for (key, value) in headers { - if key.eq_ignore_ascii_case("ocp-apim-subscription-key") { - request_builder = request_builder.header(key, value); - } - } + let (name, value) = auth.credential_header(); + let request_builder = client.get(operation_url).header(name, value); let response = request_builder .send() .await diff --git a/litellm-rust/crates/core/src/ocr/handler.rs b/litellm-rust/crates/core/src/ocr/handler.rs index ad2f9a3b785..53aefeac515 100644 --- a/litellm-rust/crates/core/src/ocr/handler.rs +++ b/litellm-rust/crates/core/src/ocr/handler.rs @@ -44,11 +44,14 @@ pub(crate) async fn execute_ocr_provider_call( max_document_download_bytes, } = request; let env_lookup = |key: &str| std::env::var(key).ok(); - let upstream_headers = config.validate_environment( - string_headers(extra_headers)?, - api_key.as_deref(), - &env_lookup, - )?; + let (upstream_headers, resolved_auth) = config + .validate_environment( + string_headers(extra_headers)?, + api_key.as_deref(), + &url_params, + &env_lookup, + ) + .await?; let url = config.complete_url(api_base.as_deref(), &model, &url_params, &env_lookup)?; let document = if config.requires_data_uri_document() { convert_document_url_to_data_uri(client, document, max_document_download_bytes).await? @@ -97,7 +100,7 @@ pub(crate) async fn execute_ocr_provider_call( ) })?; let response_json = - poll_document_intelligence(client, &operation_url, &url, &upstream_headers, timeout) + poll_document_intelligence(client, &operation_url, &url, &resolved_auth, timeout) .await?; return public_response(config, &model, &optional_params, response_json); } diff --git a/litellm-rust/crates/core/src/ocr/tests.rs b/litellm-rust/crates/core/src/ocr/tests.rs index e5bed878190..fced0329bfe 100644 --- a/litellm-rust/crates/core/src/ocr/tests.rs +++ b/litellm-rust/crates/core/src/ocr/tests.rs @@ -50,8 +50,21 @@ async fn read_http_request(socket: &mut TcpStream) -> String { } async fn write_response(socket: &mut TcpStream, status: &str, body: &str) { + write_response_with_headers(socket, status, &[], body).await; +} + +async fn write_response_with_headers( + socket: &mut TcpStream, + status: &str, + headers: &[(&str, &str)], + body: &str, +) { + let extra_headers = headers + .iter() + .map(|(name, value)| format!("{name}: {value}\r\n")) + .collect::(); let response = format!( - "HTTP/1.1 {status}\r\ncontent-type: application/json\r\ncontent-length: {}\r\nconnection: close\r\n\r\n{}", + "HTTP/1.1 {status}\r\n{extra_headers}content-type: application/json\r\ncontent-length: {}\r\nconnection: close\r\n\r\n{}", body.len(), body ); @@ -61,6 +74,54 @@ async fn write_response(socket: &mut TcpStream, status: &str, body: &str) { .expect("writes response"); } +async fn document_intelligence_round_trip( + api_key: Option<&str>, + azure_ad_token: Option<&str>, +) -> (String, String) { + let listener = TcpListener::bind("127.0.0.1:0") + .await + .expect("test listener binds"); + let address = listener.local_addr().expect("listener has local address"); + let operation_url = format!("http://{address}/operations/1"); + let server = tokio::spawn(async move { + let (mut post_socket, _) = listener.accept().await.expect("accepts POST"); + let post = read_http_request(&mut post_socket).await; + write_response_with_headers( + &mut post_socket, + "202 Accepted", + &[("operation-location", &operation_url), ("retry-after", "0")], + "", + ) + .await; + + let (mut poll_socket, _) = listener.accept().await.expect("accepts poll"); + let poll = read_http_request(&mut poll_socket).await; + write_response( + &mut poll_socket, + "200 OK", + r#"{"status":"succeeded","analyzeResult":{"content":"hello","pages":[{"pageNumber":1,"lines":[{"content":"hello"}]}]}}"#, + ) + .await; + (post, poll) + }); + let api_base = format!("http://{address}"); + let mut request = base_request("azure_ai/doc-intelligence/prebuilt-layout", &api_base); + request.api_key = api_key; + request.extra_headers = Some(Map::from_iter([( + "x-request-only".to_string(), + json!("must-not-poll"), + )])); + request.optional_params = azure_ad_token + .map(|token| Map::from_iter([("azure_ad_token".to_string(), json!(token))])) + .unwrap_or_default(); + + let client = test_client(); + ocr(&client, request) + .await + .expect("Document Intelligence OCR succeeds"); + server.await.expect("server task completes") +} + fn base_request<'a>(model: &'a str, api_base: &'a str) -> OcrRequest<'a> { OcrRequest { model, @@ -168,6 +229,77 @@ async fn vertex_routing_params_are_not_forwarded_in_provider_body() { assert!(!provider_request.contains("vertex_location")); } +#[tokio::test] +async fn document_intelligence_subscription_key_is_reused_for_polling_only() { + let (post, poll) = document_intelligence_round_trip(Some("subscription-key"), None).await; + let post = post.to_ascii_lowercase(); + let poll = poll.to_ascii_lowercase(); + assert!(post.contains("ocp-apim-subscription-key: subscription-key")); + assert!(poll.contains("ocp-apim-subscription-key: subscription-key")); + assert!(!post.contains("authorization:")); + assert!(!poll.contains("authorization:")); + assert!(post.contains("x-request-only: must-not-poll")); + assert!(!poll.contains("x-request-only:")); +} + +#[tokio::test] +async fn document_intelligence_entra_token_is_reused_for_polling_only() { + let (post, poll) = document_intelligence_round_trip(None, Some("entra-token")).await; + let post = post.to_ascii_lowercase(); + let poll = poll.to_ascii_lowercase(); + assert!(post.contains("authorization: bearer entra-token")); + assert!(poll.contains("authorization: bearer entra-token")); + assert!(!post.contains("ocp-apim-subscription-key:")); + assert!(!poll.contains("ocp-apim-subscription-key:")); + assert!(post.contains("x-request-only: must-not-poll")); + assert!(!poll.contains("x-request-only:")); +} + +#[tokio::test] +async fn document_intelligence_rejects_cross_origin_polling_before_sending_credentials() { + let provider = TcpListener::bind("127.0.0.1:0") + .await + .expect("provider listener binds"); + let provider_address = provider.local_addr().expect("provider address"); + let attacker = TcpListener::bind("127.0.0.1:0") + .await + .expect("attacker listener binds"); + let attacker_address = attacker.local_addr().expect("attacker address"); + let operation_url = format!("http://{attacker_address}/steal"); + let server = tokio::spawn(async move { + let (mut socket, _) = provider.accept().await.expect("accepts POST"); + let request = read_http_request(&mut socket).await; + write_response_with_headers( + &mut socket, + "202 Accepted", + &[("operation-location", &operation_url)], + "", + ) + .await; + request + }); + let api_base = format!("http://{provider_address}"); + let client = test_client(); + let error = ocr( + &client, + base_request("azure_ai/doc-intelligence/prebuilt-layout", &api_base), + ) + .await + .expect_err("cross-origin operation URL must fail"); + assert!(error.to_string().contains("cross-origin")); + let initial = server.await.expect("server task completes"); + assert!( + initial + .to_ascii_lowercase() + .contains("ocp-apim-subscription-key: sk-test") + ); + assert!( + tokio::time::timeout(std::time::Duration::from_millis(50), attacker.accept()) + .await + .is_err() + ); +} + #[tokio::test] async fn provider_http_error_preserves_status_and_bounds_body() { let listener = TcpListener::bind("127.0.0.1:0") diff --git a/litellm-rust/crates/core/src/ocr/transformation.rs b/litellm-rust/crates/core/src/ocr/transformation.rs index 62299faf9ed..13979a563d1 100644 --- a/litellm-rust/crates/core/src/ocr/transformation.rs +++ b/litellm-rust/crates/core/src/ocr/transformation.rs @@ -1,29 +1,18 @@ +use async_trait::async_trait; + use crate::Error; +use crate::auth::{AuthHeaderKind, Environment, ResolvedAuth}; use serde_json::{Map, Value}; use super::types::{OcrRequestData, OcrResponseData}; -#[derive(Clone, Copy, Debug, PartialEq, Eq)] -pub enum OcrAuthStrategy { - Bearer, - Header(&'static str), -} - -impl OcrAuthStrategy { - pub fn header_name(self) -> &'static str { - match self { - Self::Bearer => "authorization", - Self::Header(header_name) => header_name, - } - } -} - #[derive(Clone, Copy, Debug, PartialEq, Eq)] pub enum OcrResponseHandling { Json, AzureDocumentIntelligencePoll, } +#[async_trait] pub trait OcrProviderConfig: Sync { fn supported_ocr_params(&self) -> &'static [&'static str]; @@ -71,30 +60,50 @@ pub trait OcrProviderConfig: Sync { fn resolve_api_key( &self, api_key: Option<&str>, - env_lookup: &dyn Fn(&str) -> Option, + env_lookup: &(dyn Fn(&str) -> Option + Sync), ) -> Result; + fn forwarded_auth(&self, headers: &[(String, String)]) -> Option { + let kind = self.auth_header_kind(); + headers + .iter() + .find(|(name, _)| name.eq_ignore_ascii_case(kind.header_name())) + .and_then(|(_, value)| ResolvedAuth::from_header(kind, value)) + } + + async fn resolve_auth( + &self, + api_key: Option<&str>, + _optional_params: &Map, + environment: &dyn Environment, + ) -> Result { + let env_lookup = |name: &str| environment.get(name); + Ok(ResolvedAuth::from_credential( + self.auth_header_kind(), + self.resolve_api_key(api_key, &env_lookup)?, + )) + } + #[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)] - fn validate_environment( + async fn validate_environment( &self, headers: Vec<(String, String)>, api_key: Option<&str>, - env_lookup: &dyn Fn(&str) -> Option, - ) -> Result, Error> { - let strategy = self.auth_strategy(); - if crate::http_utils::has_header(&headers, strategy.header_name()) { - return Ok(headers); - } - let api_key = self.resolve_api_key(api_key, env_lookup)?; - let auth_header = match strategy { - OcrAuthStrategy::Bearer => ("Authorization".to_string(), format!("Bearer {api_key}")), - OcrAuthStrategy::Header(name) => (name.to_string(), api_key), + optional_params: &Map, + environment: &dyn Environment, + ) -> Result<(Vec<(String, String)>, ResolvedAuth), Error> { + let auth = match self.forwarded_auth(&headers) { + Some(auth) => auth, + None => { + self.resolve_auth(api_key, optional_params, environment) + .await? + } }; - Ok(std::iter::once(auth_header).chain(headers).collect()) + Ok((auth.apply_preserving_existing(headers), auth)) } - fn auth_strategy(&self) -> OcrAuthStrategy { - OcrAuthStrategy::Bearer + fn auth_header_kind(&self) -> AuthHeaderKind { + AuthHeaderKind::Bearer } fn requires_data_uri_document(&self) -> bool { diff --git a/litellm-rust/crates/core/src/providers/anthropic/messages/transformation.rs b/litellm-rust/crates/core/src/providers/anthropic/messages/transformation.rs index f31b961e78a..5049b514a4f 100644 --- a/litellm-rust/crates/core/src/providers/anthropic/messages/transformation.rs +++ b/litellm-rust/crates/core/src/providers/anthropic/messages/transformation.rs @@ -1,5 +1,6 @@ +use crate::auth::AuthHeaderKind; use crate::error::Error; -use crate::messages::transformation::{AnthropicMessagesProviderConfig, MessagesAuthStrategy}; +use crate::messages::transformation::AnthropicMessagesProviderConfig; const ANTHROPIC_API_KEY_ENV: &str = "ANTHROPIC_API_KEY"; const ANTHROPIC_API_BASE_ENV: &str = "ANTHROPIC_API_BASE"; @@ -65,8 +66,8 @@ impl AnthropicMessagesProviderConfig for AnthropicMessagesConfig { resolve_anthropic_api_key(api_key, env_lookup) } - fn auth_strategy(&self) -> MessagesAuthStrategy { - MessagesAuthStrategy::Header("x-api-key") + fn auth_strategy(&self) -> AuthHeaderKind { + AuthHeaderKind::Header("x-api-key") } } diff --git a/litellm-rust/crates/core/src/providers/azure_ai/auth.rs b/litellm-rust/crates/core/src/providers/azure_ai/auth.rs new file mode 100644 index 00000000000..d8e037bf556 --- /dev/null +++ b/litellm-rust/crates/core/src/providers/azure_ai/auth.rs @@ -0,0 +1,540 @@ +use std::path::PathBuf; +use std::sync::{Arc, OnceLock}; +use std::time::{Duration, Instant, SystemTime}; + +use async_trait::async_trait; +use azure_core::credentials::{Secret, TokenCredential}; +use azure_core::http::ClientMethodOptions; +use azure_identity::{ + ClientAssertion, ClientAssertionCredential, ClientSecretCredential, DeveloperToolsCredential, + ManagedIdentityCredential, ManagedIdentityCredentialOptions, UserAssignedId, + WorkloadIdentityCredential, WorkloadIdentityCredentialOptions, +}; +use serde_json::{Map, Value}; +use sha2::{Digest, Sha256}; + +use crate::Error; +use crate::auth::{ + AuthHeaderKind, Environment, ExpiringToken, ResolvedAuth, SecretString, TokenCache, + TokenProvider, +}; +use crate::constants::{AZURE_AUTH_DEFAULT_SCOPE, AZURE_AUTH_TOKEN_REFRESH_WINDOW_SECS}; + +static TOKENS: OnceLock = OnceLock::new(); + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum AzureCredentialKind { + ApiKey, + StaticToken, + ClientSecret, + ClientAssertion, + WorkloadIdentity, + ManagedIdentity, + Default, +} + +pub struct AzureAuthInput<'a> { + pub api_key: Option<&'a str>, + pub api_key_env: &'static str, + pub api_key_kind: AuthHeaderKind, + pub params: &'a Map, + pub environment: &'a dyn Environment, +} + +enum AzureTokenCredential { + ClientSecret(Arc), + ClientAssertion(Arc>), + WorkloadIdentity(Arc), + ManagedIdentity(Arc), + Default { + workload_identity: Option>, + managed_identity: Arc, + developer_tools: Option>, + }, +} + +impl AzureTokenCredential { + async fn token(&self, scope: &str) -> azure_core::Result { + match self { + Self::ClientSecret(value) => value.get_token(&[scope], None).await, + Self::ClientAssertion(value) => value.get_token(&[scope], None).await, + Self::WorkloadIdentity(value) => value.get_token(&[scope], None).await, + Self::ManagedIdentity(value) => value.get_token(&[scope], None).await, + Self::Default { + workload_identity, + managed_identity, + developer_tools, + } => { + if let Some(value) = workload_identity + && let Ok(token) = value.get_token(&[scope], None).await + { + return Ok(token); + } + match managed_identity.get_token(&[scope], None).await { + Ok(token) => Ok(token), + Err(managed_identity_error) => match developer_tools { + Some(value) => value.get_token(&[scope], None).await, + None => Err(managed_identity_error), + }, + } + } + } + } +} + +#[derive(Clone)] +struct StaticClientAssertion(SecretString); + +impl std::fmt::Debug for StaticClientAssertion { + fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + formatter.write_str("[redacted]") + } +} + +#[async_trait] +impl ClientAssertion for StaticClientAssertion { + async fn secret( + &self, + _options: Option>, + ) -> azure_core::Result { + Ok(self.0.expose().to_string()) + } +} + +struct AzureTokenProvider { + credential: AzureTokenCredential, + scope: String, +} + +#[async_trait] +impl TokenProvider for AzureTokenProvider { + async fn token(&self) -> Result { + let token = self + .credential + .token(&self.scope) + .await + .map_err(|_| Error::Auth("Azure credential refresh failed".to_string()))?; + let seconds = token + .expires_on + .unix_timestamp() + .saturating_sub( + SystemTime::now() + .duration_since(SystemTime::UNIX_EPOCH) + .unwrap_or_default() + .as_secs() as i64, + ) + .max(0) as u64; + Ok(ExpiringToken { + token: SecretString::new(token.token.secret()), + expires_at: Instant::now() + Duration::from_secs(seconds), + }) + } +} + +fn non_empty(value: Option<&str>) -> Option<&str> { + value.map(str::trim).filter(|value| !value.is_empty()) +} + +fn param_or_env( + params: &Map, + param: &str, + env: &str, + environment: &dyn Environment, +) -> Option { + params + .get(param) + .and_then(Value::as_str) + .and_then(|value| non_empty(Some(value))) + .map(str::to_string) + .or_else(|| { + environment + .get(env) + .filter(|value| !value.trim().is_empty()) + }) +} + +fn refresh_enabled(params: &Map) -> bool { + params + .get("azure_ad_token_refresh") + .and_then(Value::as_bool) + .unwrap_or(false) +} + +pub fn classify_azure_credential(input: &AzureAuthInput<'_>) -> Result { + if non_empty(input.api_key).is_some() + || input + .environment + .get(input.api_key_env) + .is_some_and(|value| !value.trim().is_empty()) + { + return Ok(AzureCredentialKind::ApiKey); + } + let azure_ad_token = param_or_env( + input.params, + "azure_ad_token", + "AZURE_AD_TOKEN", + input.environment, + ); + if azure_ad_token + .as_deref() + .is_some_and(|token| !token.starts_with("oidc/")) + { + return Ok(AzureCredentialKind::StaticToken); + } + let tenant = param_or_env( + input.params, + "tenant_id", + "AZURE_TENANT_ID", + input.environment, + ); + let client = param_or_env( + input.params, + "client_id", + "AZURE_CLIENT_ID", + input.environment, + ); + let secret = param_or_env( + input.params, + "client_secret", + "AZURE_CLIENT_SECRET", + input.environment, + ); + if tenant.is_some() && client.is_some() && secret.is_some() { + return Ok(AzureCredentialKind::ClientSecret); + } + if azure_ad_token + .as_deref() + .is_some_and(|token| token.starts_with("oidc/")) + { + if tenant.is_some() + && client.is_some() + && param_or_env( + input.params, + "azure_ad_client_assertion", + "AZURE_AD_CLIENT_ASSERTION", + input.environment, + ) + .is_some() + { + return Ok(AzureCredentialKind::ClientAssertion); + } + return Err(Error::Auth( + "Azure OIDC authentication requires tenant_id, client_id, and a resolved client assertion" + .to_string(), + )); + } + if !refresh_enabled(input.params) { + return Err(Error::Auth( + "Missing Azure credentials; set an API key, azure_ad_token, service principal, or enable Azure token refresh" + .to_string(), + )); + } + let credential = input.environment.get("AZURE_CREDENTIAL"); + if credential.as_deref() == Some("WorkloadIdentityCredential") { + return Ok(AzureCredentialKind::WorkloadIdentity); + } + if credential.as_deref() == Some("ManagedIdentityCredential") + || input.environment.get("IDENTITY_ENDPOINT").is_some() + || input.environment.get("MSI_ENDPOINT").is_some() + { + return Ok(AzureCredentialKind::ManagedIdentity); + } + Ok(AzureCredentialKind::Default) +} + +fn cache_key(kind: AzureCredentialKind, values: &[&str]) -> String { + let mut hasher = Sha256::new(); + hasher.update(format!("{kind:?}")); + values.iter().for_each(|value| { + hasher.update(value.len().to_le_bytes()); + hasher.update(value.as_bytes()); + }); + hasher + .finalize() + .iter() + .map(|byte| format!("{byte:02x}")) + .collect() +} + +fn token_credential( + kind: AzureCredentialKind, + params: &Map, + environment: &dyn Environment, +) -> Result<(String, AzureTokenCredential), Error> { + let scope = param_or_env(params, "azure_scope", "AZURE_SCOPE", environment) + .unwrap_or_else(|| AZURE_AUTH_DEFAULT_SCOPE.to_string()); + let tenant = param_or_env(params, "tenant_id", "AZURE_TENANT_ID", environment); + let client = param_or_env(params, "client_id", "AZURE_CLIENT_ID", environment); + let token_file = environment + .get("AZURE_FEDERATED_TOKEN_FILE") + .filter(|value| !value.trim().is_empty()); + let result = match kind { + AzureCredentialKind::ClientSecret => { + let tenant = + tenant.ok_or_else(|| Error::Auth("Azure tenant ID is missing".to_string()))?; + let client = + client.ok_or_else(|| Error::Auth("Azure client ID is missing".to_string()))?; + let secret = param_or_env(params, "client_secret", "AZURE_CLIENT_SECRET", environment) + .ok_or_else(|| Error::Auth("Azure client secret is missing".to_string()))?; + let key = cache_key(kind, &[&tenant, &client, &secret, &scope]); + let credential = + ClientSecretCredential::new(&tenant, client, Secret::new(secret), None) + .map_err(|_| Error::Auth("Azure service principal is invalid".to_string()))?; + (key, AzureTokenCredential::ClientSecret(credential)) + } + AzureCredentialKind::ClientAssertion => { + let tenant = + tenant.ok_or_else(|| Error::Auth("Azure tenant ID is missing".to_string()))?; + let client = + client.ok_or_else(|| Error::Auth("Azure client ID is missing".to_string()))?; + let assertion = param_or_env( + params, + "azure_ad_client_assertion", + "AZURE_AD_CLIENT_ASSERTION", + environment, + ) + .ok_or_else(|| Error::Auth("Azure client assertion is missing".to_string()))?; + let key = cache_key(kind, &[&tenant, &client, &assertion, &scope]); + let credential = ClientAssertionCredential::new( + tenant, + client, + StaticClientAssertion(SecretString::new(assertion)), + None, + ) + .map_err(|_| Error::Auth("Azure client assertion is invalid".to_string()))?; + (key, AzureTokenCredential::ClientAssertion(credential)) + } + AzureCredentialKind::WorkloadIdentity => { + let credential = + WorkloadIdentityCredential::new(Some(WorkloadIdentityCredentialOptions { + tenant_id: tenant.clone(), + client_id: client.clone(), + token_file_path: token_file.clone().map(PathBuf::from), + ..Default::default() + })) + .map_err(|_| Error::Auth("Azure workload identity is invalid".to_string()))?; + ( + cache_key( + kind, + &[ + tenant.as_deref().unwrap_or(""), + client.as_deref().unwrap_or(""), + token_file.as_deref().unwrap_or(""), + &scope, + ], + ), + AzureTokenCredential::WorkloadIdentity(credential), + ) + } + AzureCredentialKind::ManagedIdentity => { + let options = ManagedIdentityCredentialOptions { + user_assigned_id: client.clone().map(UserAssignedId::ClientId), + ..Default::default() + }; + let credential = ManagedIdentityCredential::new(Some(options)) + .map_err(|_| Error::Auth("Azure managed identity is unavailable".to_string()))?; + ( + cache_key(kind, &[client.as_deref().unwrap_or("system"), &scope]), + AzureTokenCredential::ManagedIdentity(credential), + ) + } + AzureCredentialKind::Default => { + let workload_identity = match (&tenant, &client, &token_file) { + (Some(tenant_id), Some(client_id), Some(token_file_path)) => { + WorkloadIdentityCredential::new(Some(WorkloadIdentityCredentialOptions { + tenant_id: Some(tenant_id.clone()), + client_id: Some(client_id.clone()), + token_file_path: Some(PathBuf::from(token_file_path)), + ..Default::default() + })) + .ok() + } + _ => None, + }; + let options = ManagedIdentityCredentialOptions { + user_assigned_id: client.clone().map(UserAssignedId::ClientId), + ..Default::default() + }; + let managed_identity = ManagedIdentityCredential::new(Some(options)).map_err(|_| { + Error::Auth("Azure default credentials are unavailable".to_string()) + })?; + let developer_tools = DeveloperToolsCredential::new(None).ok(); + ( + cache_key( + kind, + &[ + tenant.as_deref().unwrap_or(""), + client.as_deref().unwrap_or("system"), + token_file.as_deref().unwrap_or(""), + &scope, + ], + ), + AzureTokenCredential::Default { + workload_identity, + managed_identity, + developer_tools, + }, + ) + } + AzureCredentialKind::ApiKey | AzureCredentialKind::StaticToken => { + return Err(Error::Auth( + "Azure credential classification failed".to_string(), + )); + } + }; + Ok((format!("{}:{scope}", result.0), result.1)) +} + +pub async fn resolve_azure_auth(input: AzureAuthInput<'_>) -> Result { + let kind = classify_azure_credential(&input)?; + match kind { + AzureCredentialKind::ApiKey => { + let key = non_empty(input.api_key) + .map(str::to_string) + .or_else(|| input.environment.get(input.api_key_env)) + .ok_or_else(|| Error::Auth("Azure API key is missing".to_string()))?; + Ok(ResolvedAuth::from_credential(input.api_key_kind, key)) + } + AzureCredentialKind::StaticToken => { + let token = param_or_env( + input.params, + "azure_ad_token", + "AZURE_AD_TOKEN", + input.environment, + ) + .ok_or_else(|| Error::Auth("Azure access token is missing".to_string()))?; + Ok(ResolvedAuth::Bearer(SecretString::new(token))) + } + kind => { + let scope = param_or_env( + input.params, + "azure_scope", + "AZURE_SCOPE", + input.environment, + ) + .unwrap_or_else(|| AZURE_AUTH_DEFAULT_SCOPE.to_string()); + let (key, credential) = token_credential(kind, input.params, input.environment)?; + let provider = AzureTokenProvider { credential, scope }; + let token = TOKENS + .get_or_init(|| { + TokenCache::new(Duration::from_secs(AZURE_AUTH_TOKEN_REFRESH_WINDOW_SECS)) + }) + .token(key, &provider) + .await?; + Ok(ResolvedAuth::Bearer(token)) + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + + fn empty_environment(_: &str) -> Option { + None + } + + fn input<'a>( + api_key: Option<&'a str>, + params: &'a Map, + env: &'a dyn Environment, + ) -> AzureAuthInput<'a> { + AzureAuthInput { + api_key, + api_key_env: "AZURE_DOCUMENT_INTELLIGENCE_API_KEY", + api_key_kind: AuthHeaderKind::Header("Ocp-Apim-Subscription-Key"), + params, + environment: env, + } + } + + #[test] + fn api_key_precedes_static_token_and_service_principal() { + let params = Map::from_iter([ + ( + "azure_ad_token".to_string(), + Value::String("token".to_string()), + ), + ("tenant_id".to_string(), Value::String("tenant".to_string())), + ("client_id".to_string(), Value::String("client".to_string())), + ( + "client_secret".to_string(), + Value::String("secret".to_string()), + ), + ]); + assert_eq!( + classify_azure_credential(&input(Some("key"), ¶ms, &empty_environment)).unwrap(), + AzureCredentialKind::ApiKey + ); + } + + #[test] + fn static_token_precedes_service_principal_and_refresh_requires_opt_in() { + let token_params = Map::from_iter([ + ( + "azure_ad_token".to_string(), + Value::String("token".to_string()), + ), + ("tenant_id".to_string(), Value::String("tenant".to_string())), + ("client_id".to_string(), Value::String("client".to_string())), + ( + "client_secret".to_string(), + Value::String("secret".to_string()), + ), + ]); + assert_eq!( + classify_azure_credential(&input(None, &token_params, &empty_environment)).unwrap(), + AzureCredentialKind::StaticToken + ); + let params = Map::new(); + assert!(classify_azure_credential(&input(None, ¶ms, &empty_environment)).is_err()); + } + + #[test] + fn oidc_reference_requires_a_resolved_assertion() { + let params = Map::from_iter([ + ( + "azure_ad_token".to_string(), + Value::String("oidc/token-reference".to_string()), + ), + ("tenant_id".to_string(), Value::String("tenant".to_string())), + ("client_id".to_string(), Value::String("client".to_string())), + ]); + assert!(classify_azure_credential(&input(None, ¶ms, &empty_environment)).is_err()); + + let mut resolved = params; + resolved.insert( + "azure_ad_client_assertion".to_string(), + Value::String("assertion".to_string()), + ); + assert_eq!( + classify_azure_credential(&input(None, &resolved, &empty_environment)).unwrap(), + AzureCredentialKind::ClientAssertion + ); + } + + #[test] + fn refresh_uses_managed_identity_signals_or_default_chain() { + let params = Map::from_iter([("azure_ad_token_refresh".to_string(), Value::Bool(true))]); + let managed_environment = |name: &str| { + (name == "IDENTITY_ENDPOINT").then(|| "https://identity.example".to_string()) + }; + + assert_eq!( + classify_azure_credential(&input(None, ¶ms, &managed_environment)).unwrap(), + AzureCredentialKind::ManagedIdentity + ); + assert_eq!( + classify_azure_credential(&input(None, ¶ms, &empty_environment)).unwrap(), + AzureCredentialKind::Default + ); + } + + #[test] + fn cache_keys_preserve_credential_field_boundaries() { + assert_ne!( + cache_key(AzureCredentialKind::ClientSecret, &["ab", "c"]), + cache_key(AzureCredentialKind::ClientSecret, &["a", "bc"]) + ); + } +} diff --git a/litellm-rust/crates/core/src/providers/azure_ai/messages/transformation.rs b/litellm-rust/crates/core/src/providers/azure_ai/messages/transformation.rs index b8ca10461fb..88507590506 100644 --- a/litellm-rust/crates/core/src/providers/azure_ai/messages/transformation.rs +++ b/litellm-rust/crates/core/src/providers/azure_ai/messages/transformation.rs @@ -1,5 +1,6 @@ +use crate::auth::AuthHeaderKind; use crate::error::Error; -use crate::messages::transformation::{AnthropicMessagesProviderConfig, MessagesAuthStrategy}; +use crate::messages::transformation::AnthropicMessagesProviderConfig; use crate::messages::types::{ AnthropicMessage, AnthropicMessagesRequest, AnthropicMessagesResponse, ContentBlock, MessageContent, SystemPrompt, @@ -160,7 +161,7 @@ impl AnthropicMessagesProviderConfig for AzureAnthropicMessagesConfig { resolve_azure_api_key(api_key, env_lookup) } - fn auth_strategy(&self) -> MessagesAuthStrategy { + fn auth_strategy(&self) -> AuthHeaderKind { self.anthropic.auth_strategy() } diff --git a/litellm-rust/crates/core/src/providers/azure_ai/mod.rs b/litellm-rust/crates/core/src/providers/azure_ai/mod.rs index 5d13fa93e00..4b2a66a3b0a 100644 --- a/litellm-rust/crates/core/src/providers/azure_ai/mod.rs +++ b/litellm-rust/crates/core/src/providers/azure_ai/mod.rs @@ -1,2 +1,3 @@ +pub mod auth; pub mod messages; pub mod ocr; diff --git a/litellm-rust/crates/core/src/providers/azure_ai/ocr/transformation.rs b/litellm-rust/crates/core/src/providers/azure_ai/ocr/transformation.rs index 3cecbbbab19..62a41be9e45 100644 --- a/litellm-rust/crates/core/src/providers/azure_ai/ocr/transformation.rs +++ b/litellm-rust/crates/core/src/providers/azure_ai/ocr/transformation.rs @@ -1,8 +1,12 @@ use std::collections::BTreeSet; +use async_trait::async_trait; + +use crate::auth::{AuthHeaderKind, Environment, ResolvedAuth}; use crate::error::{Error, json_type_name}; -use crate::ocr::transformation::{OcrAuthStrategy, OcrProviderConfig, OcrResponseHandling}; +use crate::ocr::transformation::{OcrProviderConfig, OcrResponseHandling}; use crate::ocr::types::{OcrRequestData, OcrResponseData}; +use crate::providers::azure_ai::auth::{AzureAuthInput, resolve_azure_auth}; use serde_json::{Map, Value, json}; use crate::providers::mistral::ocr::transformation::MISTRAL_OCR_CONFIG; @@ -99,68 +103,6 @@ pub fn resolve_document_intelligence_endpoint( ) } -fn prepend_auth_header( - headers: Vec<(String, String)>, - name: &str, - value: String, -) -> Vec<(String, String)> { - std::iter::once((name.to_string(), value)) - .chain(headers) - .collect() -} - -pub fn validate_azure_ai_environment( - headers: Vec<(String, String)>, - api_key: Option<&str>, - azure_ad_token: Option<&str>, - env_lookup: &dyn Fn(&str) -> Option, -) -> Result, Error> { - if crate::http_utils::has_header(&headers, "Authorization") - || crate::http_utils::has_header(&headers, "Api-Key") - { - return Ok(headers); - } - if let Ok(api_key) = resolve_azure_ai_api_key(api_key, env_lookup) { - return Ok(prepend_auth_header(headers, "Api-Key", api_key)); - } - non_empty(azure_ad_token) - .map(|token| prepend_auth_header(headers, "Authorization", format!("Bearer {token}"))) - .ok_or_else(|| { - Error::Auth( - "Missing Azure AI credentials - set AZURE_AI_API_KEY or provide azure_ad_token" - .to_string(), - ) - }) -} - -pub fn validate_document_intelligence_environment( - headers: Vec<(String, String)>, - api_key: Option<&str>, - azure_ad_token: Option<&str>, - env_lookup: &dyn Fn(&str) -> Option, -) -> Result, Error> { - if crate::http_utils::has_header(&headers, "Authorization") - || crate::http_utils::has_header(&headers, "Ocp-Apim-Subscription-Key") - { - return Ok(headers); - } - if let Ok(api_key) = resolve_document_intelligence_api_key(api_key, env_lookup) { - return Ok(prepend_auth_header( - headers, - "Ocp-Apim-Subscription-Key", - api_key, - )); - } - non_empty(azure_ad_token) - .map(|token| prepend_auth_header(headers, "Authorization", format!("Bearer {token}"))) - .ok_or_else(|| { - Error::Auth( - "Missing Azure Document Intelligence credentials - set AZURE_DOCUMENT_INTELLIGENCE_API_KEY or provide azure_ad_token" - .to_string(), - ) - }) -} - fn encode_model_id(model: &str) -> Result { let model_id = model.rsplit('/').next().unwrap_or(model); if matches!(model_id, "." | "..") { @@ -502,6 +444,7 @@ fn transform_document_intelligence_response( }) } +#[async_trait] impl OcrProviderConfig for AzureAiOcrConfig { fn supported_ocr_params(&self) -> &'static [&'static str] { MISTRAL_OCR_CONFIG.supported_ocr_params() @@ -539,16 +482,33 @@ impl OcrProviderConfig for AzureAiOcrConfig { fn resolve_api_key( &self, api_key: Option<&str>, - env_lookup: &dyn Fn(&str) -> Option, + env_lookup: &(dyn Fn(&str) -> Option + Sync), ) -> Result { resolve_azure_ai_api_key(api_key, env_lookup) } + async fn resolve_auth( + &self, + api_key: Option<&str>, + optional_params: &Map, + environment: &dyn Environment, + ) -> Result { + resolve_azure_auth(AzureAuthInput { + api_key, + api_key_env: AZURE_AI_API_KEY_ENV, + api_key_kind: AuthHeaderKind::Bearer, + params: optional_params, + environment, + }) + .await + } + fn requires_data_uri_document(&self) -> bool { true } } +#[async_trait] impl OcrProviderConfig for AzureDocumentIntelligenceOcrConfig { #[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)] fn supported_ocr_params(&self) -> &'static [&'static str] { @@ -636,13 +596,43 @@ impl OcrProviderConfig for AzureDocumentIntelligenceOcrConfig { fn resolve_api_key( &self, api_key: Option<&str>, - env_lookup: &dyn Fn(&str) -> Option, + env_lookup: &(dyn Fn(&str) -> Option + Sync), ) -> Result { resolve_document_intelligence_api_key(api_key, env_lookup) } - fn auth_strategy(&self) -> OcrAuthStrategy { - OcrAuthStrategy::Header("Ocp-Apim-Subscription-Key") + fn forwarded_auth(&self, headers: &[(String, String)]) -> Option { + [ + AuthHeaderKind::Header("Ocp-Apim-Subscription-Key"), + AuthHeaderKind::Bearer, + ] + .iter() + .find_map(|kind| { + headers + .iter() + .find(|(name, _)| name.eq_ignore_ascii_case(kind.header_name())) + .and_then(|(_, value)| ResolvedAuth::from_header(*kind, value)) + }) + } + + async fn resolve_auth( + &self, + api_key: Option<&str>, + optional_params: &Map, + environment: &dyn Environment, + ) -> Result { + resolve_azure_auth(AzureAuthInput { + api_key, + api_key_env: AZURE_DOCUMENT_INTELLIGENCE_API_KEY_ENV, + api_key_kind: AuthHeaderKind::Header("Ocp-Apim-Subscription-Key"), + params: optional_params, + environment, + }) + .await + } + + fn auth_header_kind(&self) -> AuthHeaderKind { + AuthHeaderKind::Header("Ocp-Apim-Subscription-Key") } fn response_handling(&self) -> OcrResponseHandling { @@ -669,6 +659,10 @@ mod tests { .map(|(_, value)| value.as_str()) } + fn empty_environment(_: &str) -> Option { + None + } + #[fixture] fn native_operation() -> Value { json!({ @@ -1103,11 +1097,12 @@ mod tests { assert!(!url.contains("req_format")); } - #[test] - fn document_intelligence_validate_environment_uses_subscription_key() { - let headers = - validate_document_intelligence_environment(Vec::new(), Some("my-key"), None, &|_| None) - .expect("api key authenticates"); + #[tokio::test] + async fn document_intelligence_validate_environment_uses_subscription_key() { + let (headers, _) = AZURE_DOCUMENT_INTELLIGENCE_OCR_CONFIG + .validate_environment(Vec::new(), Some("my-key"), &Map::new(), &empty_environment) + .await + .expect("api key authenticates"); assert_eq!( header_value(&headers, "Ocp-Apim-Subscription-Key"), @@ -1115,15 +1110,13 @@ mod tests { ); } - #[test] - fn document_intelligence_validate_environment_falls_back_to_entra_token() { - let headers = validate_document_intelligence_environment( - Vec::new(), - None, - Some("entra-token"), - &|_| None, - ) - .expect("Entra token authenticates"); + #[tokio::test] + async fn document_intelligence_validate_environment_falls_back_to_entra_token() { + let params = Map::from_iter([("azure_ad_token".to_string(), json!("entra-token"))]); + let (headers, _) = AZURE_DOCUMENT_INTELLIGENCE_OCR_CONFIG + .validate_environment(Vec::new(), None, ¶ms, &empty_environment) + .await + .expect("Entra token authenticates"); assert_eq!( header_value(&headers, "Authorization"), @@ -1373,11 +1366,13 @@ mod tests { assert_eq!(resolved, "https://generic-azure-ai.example.com"); } - #[test] - fn azure_ai_ocr_authenticates_with_entra_token() { - let headers = - validate_azure_ai_environment(Vec::new(), None, Some("entra-token"), &|_| None) - .expect("Entra token authenticates"); + #[tokio::test] + async fn azure_ai_ocr_authenticates_with_entra_token() { + let params = Map::from_iter([("azure_ad_token".to_string(), json!("entra-token"))]); + let (headers, _) = AZURE_AI_OCR_CONFIG + .validate_environment(Vec::new(), None, ¶ms, &empty_environment) + .await + .expect("Entra token authenticates"); assert_eq!( header_value(&headers, "Authorization"), diff --git a/litellm-rust/crates/core/src/providers/mistral/ocr/transformation.rs b/litellm-rust/crates/core/src/providers/mistral/ocr/transformation.rs index 11e8fe7db18..0db2e6c7084 100644 --- a/litellm-rust/crates/core/src/providers/mistral/ocr/transformation.rs +++ b/litellm-rust/crates/core/src/providers/mistral/ocr/transformation.rs @@ -153,7 +153,7 @@ impl OcrProviderConfig for MistralOcrConfig { fn resolve_api_key( &self, api_key: Option<&str>, - env_lookup: &dyn Fn(&str) -> Option, + env_lookup: &(dyn Fn(&str) -> Option + Sync), ) -> Result { resolve_api_key(api_key, env_lookup) } diff --git a/litellm-rust/crates/core/src/providers/reducto/ocr/tests.rs b/litellm-rust/crates/core/src/providers/reducto/ocr/tests.rs index 2b66d058b5d..e51d74f945d 100644 --- a/litellm-rust/crates/core/src/providers/reducto/ocr/tests.rs +++ b/litellm-rust/crates/core/src/providers/reducto/ocr/tests.rs @@ -1,9 +1,13 @@ use rstest::{fixture, rstest}; -use serde_json::{Value, json}; +use serde_json::{Map, Value, json}; use super::transformation::*; use crate::ocr::transformation::OcrProviderConfig; +fn reducto_environment(_: &str) -> Option { + Some("env-reducto-key".to_string()) +} + #[fixture] fn parse_response() -> Value { json!({ @@ -185,15 +189,20 @@ fn test_parse_v3_rejects_plain_http_urls(#[case] source: &str) { } #[rstest] -fn test_parse_v3_uses_programmatic_api_key_over_env() { +#[tokio::test] +async fn test_parse_v3_uses_programmatic_api_key_over_env() { let key = resolve_api_key(Some("passed-key"), &|_| Some("env-reducto-key".to_string())) .expect("explicit key should resolve"); assert_eq!(key, "passed-key"); - let headers = REDUCTO_PARSE_V3_CONFIG - .validate_environment(Vec::new(), Some("passed-key"), &|_| { - Some("env-reducto-key".to_string()) - }) + let (headers, _) = REDUCTO_PARSE_V3_CONFIG + .validate_environment( + Vec::new(), + Some("passed-key"), + &Map::new(), + &reducto_environment, + ) + .await .expect("headers should validate"); assert_eq!( headers, diff --git a/litellm-rust/crates/core/src/providers/reducto/ocr/transformation.rs b/litellm-rust/crates/core/src/providers/reducto/ocr/transformation.rs index b8507541e18..9658055fc49 100644 --- a/litellm-rust/crates/core/src/providers/reducto/ocr/transformation.rs +++ b/litellm-rust/crates/core/src/providers/reducto/ocr/transformation.rs @@ -358,7 +358,7 @@ impl OcrProviderConfig for ReductoParseV3Config { fn resolve_api_key( &self, api_key: Option<&str>, - env_lookup: &dyn Fn(&str) -> Option, + env_lookup: &(dyn Fn(&str) -> Option + Sync), ) -> Result { resolve_api_key(api_key, env_lookup) } @@ -400,7 +400,7 @@ impl OcrProviderConfig for ReductoParseLegacyConfig { fn resolve_api_key( &self, api_key: Option<&str>, - env_lookup: &dyn Fn(&str) -> Option, + env_lookup: &(dyn Fn(&str) -> Option + Sync), ) -> Result { resolve_api_key(api_key, env_lookup) } diff --git a/litellm-rust/crates/core/src/providers/vertex_ai/auth.rs b/litellm-rust/crates/core/src/providers/vertex_ai/auth.rs new file mode 100644 index 00000000000..ea010f18afb --- /dev/null +++ b/litellm-rust/crates/core/src/providers/vertex_ai/auth.rs @@ -0,0 +1,290 @@ +use std::collections::HashMap; +use std::sync::OnceLock; +use std::time::{Duration, Instant}; + +use async_trait::async_trait; +use google_cloud_auth::credentials::impersonated; +use google_cloud_auth::credentials::{ + AccessTokenCredentials, external_account, service_account, user_account, +}; +use serde_json::{Map, Value}; +use sha2::{Digest, Sha256}; +use tokio::sync::Mutex; + +use crate::Error; +use crate::auth::{ + Environment, ExpiringToken, ResolvedAuth, SecretString, TokenCache, TokenProvider, +}; +use crate::constants::{ + VERTEX_AUTH_CLOUD_PLATFORM_SCOPE, VERTEX_AUTH_SDK_TOKEN_CACHE_TTL_SECS, + VERTEX_AUTH_TOKEN_REFRESH_WINDOW_SECS, +}; + +static CREDENTIALS: OnceLock>> = OnceLock::new(); +static TOKENS: OnceLock = OnceLock::new(); + +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct VertexAuth { + pub auth: ResolvedAuth, + pub project_id: String, +} + +struct GoogleTokenProvider { + credentials: AccessTokenCredentials, +} + +#[async_trait] +impl TokenProvider for GoogleTokenProvider { + async fn token(&self) -> Result { + let token = self + .credentials + .access_token() + .await + .map_err(|_| Error::Auth("Vertex AI credential refresh failed".to_string()))?; + Ok(ExpiringToken { + token: SecretString::new(token.token), + expires_at: Instant::now() + Duration::from_secs(VERTEX_AUTH_SDK_TOKEN_CACHE_TTL_SECS), + }) + } +} + +fn non_empty(value: Option<&str>) -> Option<&str> { + value.map(str::trim).filter(|value| !value.is_empty()) +} + +fn value_string<'a>(params: &'a Map, names: &[&str]) -> Option<&'a str> { + names + .iter() + .find_map(|name| params.get(*name).and_then(Value::as_str)) + .and_then(|value| non_empty(Some(value))) +} + +fn fingerprint(value: &str) -> String { + Sha256::digest(value.as_bytes()) + .iter() + .map(|byte| format!("{byte:02x}")) + .collect() +} + +fn credentials_json(value: &Value) -> Result, Error> { + match value { + Value::Null => Ok(None), + Value::Object(_) => Ok(Some(value.clone())), + Value::String(raw) if raw.trim_start().starts_with('{') => serde_json::from_str(raw) + .map(Some) + .map_err(|_| Error::Auth("Vertex AI credentials JSON is invalid".to_string())), + Value::String(path) => std::fs::read_to_string(path) + .map_err(|_| Error::Auth("Vertex AI credentials file could not be read".to_string())) + .and_then(|raw| { + serde_json::from_str(&raw).map(Some).map_err(|_| { + Error::Auth("Vertex AI credentials file contains invalid JSON".to_string()) + }) + }), + _ => Err(Error::Auth( + "Vertex AI credentials must be a JSON object, JSON string, or file path".to_string(), + )), + } +} + +fn build_credentials(json: Option) -> Result { + let built = match json { + None => google_cloud_auth::credentials::Builder::default() + .with_scopes([VERTEX_AUTH_CLOUD_PLATFORM_SCOPE]) + .build_access_token_credentials(), + Some(value) => match value.get("type").and_then(Value::as_str) { + Some("service_account") => service_account::Builder::new(value) + .with_access_specifier(service_account::AccessSpecifier::from_scopes([ + VERTEX_AUTH_CLOUD_PLATFORM_SCOPE, + ])) + .build_access_token_credentials(), + Some("authorized_user") => user_account::Builder::new(value) + .with_scopes([VERTEX_AUTH_CLOUD_PLATFORM_SCOPE]) + .build_access_token_credentials(), + Some("external_account") => external_account::Builder::new(value) + .with_scopes([VERTEX_AUTH_CLOUD_PLATFORM_SCOPE]) + .build_access_token_credentials(), + Some("impersonated_service_account") => impersonated::Builder::new(value) + .with_scopes([VERTEX_AUTH_CLOUD_PLATFORM_SCOPE]) + .build_access_token_credentials(), + _ => { + return Err(Error::Auth( + "Vertex AI credential type is unsupported".to_string(), + )); + } + }, + }; + built.map_err(|_| Error::Auth("Vertex AI credentials could not be loaded".to_string())) +} + +fn configured_project( + explicit: Option<&str>, + env_lookup: &dyn Fn(&str) -> Option, +) -> Option { + non_empty(explicit).map(str::to_string).or_else(|| { + ["VERTEXAI_PROJECT", "GOOGLE_CLOUD_PROJECT", "GCLOUD_PROJECT"] + .iter() + .find_map(|name| env_lookup(name)) + .filter(|value| !value.trim().is_empty()) + }) +} + +fn inferred_project( + explicit: Option<&str>, + json: Option<&Value>, + env_lookup: &dyn Fn(&str) -> Option, +) -> Result { + configured_project(explicit, env_lookup) + .or_else(|| { + json.and_then(|value| { + value + .get("project_id") + .or_else(|| value.get("quota_project_id")) + .and_then(Value::as_str) + .and_then(|value| non_empty(Some(value))) + .map(str::to_string) + }) + }) + .ok_or_else(|| { + Error::InvalidRequest( + "Could not resolve Vertex AI project; pass vertex_project or set VERTEXAI_PROJECT" + .to_string(), + ) + }) +} + +pub fn resolve_vertex_project( + params: &Map, + env_lookup: &dyn Fn(&str) -> Option, +) -> Result { + let explicit = value_string(params, &["vertex_project", "vertex_ai_project"]); + if let Some(project) = configured_project(explicit, env_lookup) { + return Ok(project); + } + let json = params + .get("vertex_credentials") + .map(credentials_json) + .transpose()? + .flatten(); + inferred_project(None, json.as_ref(), env_lookup) +} + +async fn cached_credentials( + key: String, + json: Option, +) -> Result { + let credentials = CREDENTIALS.get_or_init(|| Mutex::new(HashMap::new())); + let mut entries = credentials.lock().await; + if let Some(credentials) = entries.get(&key) { + return Ok(credentials.clone()); + } + let loaded = build_credentials(json)?; + entries.insert(key, loaded.clone()); + Ok(loaded) +} + +pub async fn resolve_vertex_auth( + api_key: Option<&str>, + params: &Map, + environment: &dyn Environment, +) -> Result { + let explicit_project = value_string(params, &["vertex_project", "vertex_ai_project"]); + let env_lookup = |name: &str| environment.get(name); + let configured_project = configured_project(explicit_project, &env_lookup); + let explicit_token = non_empty(api_key) + .map(str::to_string) + .or_else(|| { + environment + .get("VERTEX_AI_API_KEY") + .filter(|value| !value.trim().is_empty()) + }) + .or_else(|| { + environment + .get("VERTEXAI_API_KEY") + .filter(|value| !value.trim().is_empty()) + }); + let credential_value = params.get("vertex_credentials"); + let json = match (configured_project.as_ref(), explicit_token.as_ref()) { + (Some(_), Some(_)) => None, + _ => credential_value + .map(credentials_json) + .transpose()? + .flatten(), + }; + let project_id = configured_project + .map(Ok) + .unwrap_or_else(|| inferred_project(None, json.as_ref(), &env_lookup))?; + let token = match explicit_token { + Some(token) => SecretString::new(token), + None => { + let identity = credential_value + .map(Value::to_string) + .unwrap_or_else(|| "application-default".to_string()); + let key = fingerprint(&identity); + let provider = GoogleTokenProvider { + credentials: cached_credentials(key.clone(), json).await?, + }; + TOKENS + .get_or_init(|| { + TokenCache::new(Duration::from_secs(VERTEX_AUTH_TOKEN_REFRESH_WINDOW_SECS)) + }) + .token(key, &provider) + .await? + } + }; + Ok(VertexAuth { + auth: ResolvedAuth::Bearer(token), + project_id, + }) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[tokio::test] + async fn explicit_token_and_project_precede_environment_without_loading_adc() { + let params = Map::from_iter([ + ( + "vertex_project".to_string(), + Value::String("param-project".to_string()), + ), + ( + "vertex_credentials".to_string(), + Value::String("missing-credentials.json".to_string()), + ), + ]); + let env = |name: &str| match name { + "VERTEXAI_PROJECT" => Some("env-project".to_string()), + "VERTEX_AI_API_KEY" => Some("env-token".to_string()), + _ => None, + }; + let resolved = resolve_vertex_auth(Some("explicit-token"), ¶ms, &env) + .await + .expect("explicit credentials resolve"); + assert_eq!(resolved.project_id, "param-project"); + assert_eq!( + resolved.auth.credential_header(), + ( + "Authorization".to_string(), + "Bearer explicit-token".to_string() + ) + ); + } + + #[test] + fn credential_json_infers_service_account_and_authorized_user_projects() { + let env = |_: &str| None; + let service = + serde_json::json!({"type": "service_account", "project_id": "service-project"}); + let user = + serde_json::json!({"type": "authorized_user", "quota_project_id": "quota-project"}); + assert_eq!( + inferred_project(None, Some(&service), &env).unwrap(), + "service-project" + ); + assert_eq!( + inferred_project(None, Some(&user), &env).unwrap(), + "quota-project" + ); + } +} diff --git a/litellm-rust/crates/core/src/providers/vertex_ai/mod.rs b/litellm-rust/crates/core/src/providers/vertex_ai/mod.rs index 3621ff6a2fd..47c3f6b6c66 100644 --- a/litellm-rust/crates/core/src/providers/vertex_ai/mod.rs +++ b/litellm-rust/crates/core/src/providers/vertex_ai/mod.rs @@ -1 +1,2 @@ +pub mod auth; pub mod ocr; diff --git a/litellm-rust/crates/core/src/providers/vertex_ai/ocr/transformation.rs b/litellm-rust/crates/core/src/providers/vertex_ai/ocr/transformation.rs index c324de8cb45..1a73a9a594d 100644 --- a/litellm-rust/crates/core/src/providers/vertex_ai/ocr/transformation.rs +++ b/litellm-rust/crates/core/src/providers/vertex_ai/ocr/transformation.rs @@ -1,15 +1,18 @@ +use async_trait::async_trait; + +use crate::auth::{Environment, ResolvedAuth}; use crate::error::{Error, json_type_name}; use crate::ocr::transformation::OcrProviderConfig; use crate::ocr::types::{OcrRequestData, OcrResponseData}; use serde_json::{Map, Value, json}; use crate::providers::mistral::ocr::transformation::MISTRAL_OCR_CONFIG; +use crate::providers::vertex_ai::auth::{resolve_vertex_auth, resolve_vertex_project}; const VERTEX_DEFAULT_LOCATION: &str = "us-central1"; const VERTEX_DEFAULT_DEEPSEEK_API_BASE: &str = "https://aiplatform.googleapis.com"; const VERTEX_AI_API_KEY_ENV: &str = "VERTEX_AI_API_KEY"; const VERTEXAI_API_KEY_ENV: &str = "VERTEXAI_API_KEY"; -const VERTEXAI_PROJECT_ENV: &str = "VERTEXAI_PROJECT"; const VERTEXAI_LOCATION_ENV: &str = "VERTEXAI_LOCATION"; const VERTEX_LOCATION_ENV: &str = "VERTEX_LOCATION"; @@ -62,15 +65,7 @@ fn vertex_project( params: &Map, env_lookup: &dyn Fn(&str) -> Option, ) -> Result { - string_param(params, &["vertex_project", "vertex_ai_project"]) - .map(str::to_string) - .or_else(|| env_lookup(VERTEXAI_PROJECT_ENV).filter(|value| !value.trim().is_empty())) - .ok_or_else(|| { - Error::InvalidRequest( - "Missing vertex_project - Set VERTEXAI_PROJECT environment variable or pass vertex_project parameter" - .to_string(), - ) - }) + resolve_vertex_project(params, env_lookup) } fn vertex_location( @@ -207,6 +202,7 @@ fn ocr_data_from_content(content: Value, usage: Option, model: &str) -> V } } +#[async_trait] impl OcrProviderConfig for VertexAiOcrConfig { fn supported_ocr_params(&self) -> &'static [&'static str] { MISTRAL_OCR_CONFIG.supported_ocr_params() @@ -244,16 +240,28 @@ impl OcrProviderConfig for VertexAiOcrConfig { fn resolve_api_key( &self, api_key: Option<&str>, - env_lookup: &dyn Fn(&str) -> Option, + env_lookup: &(dyn Fn(&str) -> Option + Sync), ) -> Result { resolve_vertex_api_key(api_key, env_lookup) } + async fn resolve_auth( + &self, + api_key: Option<&str>, + optional_params: &Map, + environment: &dyn Environment, + ) -> Result { + Ok(resolve_vertex_auth(api_key, optional_params, environment) + .await? + .auth) + } + fn requires_data_uri_document(&self) -> bool { true } } +#[async_trait] impl OcrProviderConfig for VertexAiDeepSeekOcrConfig { #[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)] fn supported_ocr_params(&self) -> &'static [&'static str] { @@ -368,10 +376,21 @@ impl OcrProviderConfig for VertexAiDeepSeekOcrConfig { fn resolve_api_key( &self, api_key: Option<&str>, - env_lookup: &dyn Fn(&str) -> Option, + env_lookup: &(dyn Fn(&str) -> Option + Sync), ) -> Result { resolve_vertex_api_key(api_key, env_lookup) } + + async fn resolve_auth( + &self, + api_key: Option<&str>, + optional_params: &Map, + environment: &dyn Environment, + ) -> Result { + Ok(resolve_vertex_auth(api_key, optional_params, environment) + .await? + .auth) + } } #[cfg(test)] diff --git a/litellm/llms/azure_ai/ocr/document_intelligence/transformation.py b/litellm/llms/azure_ai/ocr/document_intelligence/transformation.py index f5126f81006..2632c02c67a 100644 --- a/litellm/llms/azure_ai/ocr/document_intelligence/transformation.py +++ b/litellm/llms/azure_ai/ocr/document_intelligence/transformation.py @@ -618,7 +618,11 @@ class AzureDocumentIntelligenceOCRConfig(BaseOCRConfig): except SSRFError as ssrf_err: raise ValueError(f"Azure Document Intelligence: rejected polling URL ({ssrf_err})") - poll_headers = {"Ocp-Apim-Subscription-Key": raw_response.request.headers.get("Ocp-Apim-Subscription-Key", "")} + poll_headers: Final = { # mutable-ok: httpx requires a concrete headers dict + name: value + for name in ("Ocp-Apim-Subscription-Key", "Authorization") + if (value := raw_response.request.headers.get(name)) + } return operation_url, poll_headers @staticmethod diff --git a/litellm/ocr/main.py b/litellm/ocr/main.py index 777bb504e5c..6882842d493 100644 --- a/litellm/ocr/main.py +++ b/litellm/ocr/main.py @@ -196,7 +196,7 @@ 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) + optional_params: Final = prepared_request.optional_params if prepared_request.custom_llm_provider == "vertex_ai": vertex_project: Final = ( prepared_request.litellm_params.get("vertex_project") @@ -211,11 +211,45 @@ def _rust_bridge_optional_params( 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 + vertex_credentials: Final = prepared_request.litellm_params.get("vertex_credentials") + vertex_auth_params: Final = { # mutable-ok: the Rust bridge requires a concrete dict + name: value + for name, value in ( + ("vertex_project", vertex_project), + ("vertex_location", vertex_location), + ("vertex_credentials", vertex_credentials), + ) + if value is not None + } + return { # mutable-ok: the Rust bridge requires a concrete dict + **optional_params, + **vertex_auth_params, + } + if prepared_request.custom_llm_provider == "azure_ai": + azure_ad_token: Final = prepared_request.litellm_params.get("azure_ad_token") + azure_ad_client_assertion: Final = ( + resolve_secret(azure_ad_token) + if isinstance(azure_ad_token, str) and azure_ad_token.startswith("oidc/") + else None + ) + azure_auth_params: Final = { # mutable-ok: the Rust bridge requires a concrete dict + name: value + for name, value in ( + ("azure_ad_token", azure_ad_token), + ("azure_ad_client_assertion", azure_ad_client_assertion), + ("tenant_id", prepared_request.litellm_params.get("tenant_id")), + ("client_id", prepared_request.litellm_params.get("client_id")), + ("client_secret", prepared_request.litellm_params.get("client_secret")), + ("azure_scope", prepared_request.litellm_params.get("azure_scope")), + ) + if value is not None + } + return { # mutable-ok: the Rust bridge requires a concrete dict + **optional_params, + **azure_auth_params, + "azure_ad_token_refresh": litellm.enable_azure_ad_token_refresh is True, + } + return dict(optional_params) # mutable-ok: the Rust bridge requires a concrete dict def _rust_bridge_api_base( @@ -237,15 +271,21 @@ def _prepare_rust_ocr_call( ) -> rust_ocr_bridge.RustOCRRequest: provider_config: Final = prepared_request.provider_config api_key_env_var: Final = provider_config.get_api_key_env_var() + rust_owns_auth: Final = prepared_request.custom_llm_provider in ("vertex_ai", "azure_ai") 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 + resolve_api_key(api_key_env_var) if api_key_env_var is not None and not rust_owns_auth else None ) - 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, + resolved_headers: Final = ( + dict(prepared_request.extra_headers or {}) # mutable-ok: the Rust bridge requires a concrete dict + if rust_owns_auth + else provider_config.validate_environment( + headers=prepared_request.extra_headers + or {}, # mutable-ok: provider validation requires mutable request headers + model=prepared_request.model, + api_key=resolved_api_key, + api_base=prepared_request.api_base, + litellm_params=prepared_request.litellm_params, + ) ) resolved_complete_url: Final = provider_config.get_complete_url( api_base=prepared_request.api_base, diff --git a/tests/test_litellm/llms/azure_ai/test_azure_document_intelligence_ocr_transformation.py b/tests/test_litellm/llms/azure_ai/test_azure_document_intelligence_ocr_transformation.py index ef4c78553f1..ca15bd5ec46 100644 --- a/tests/test_litellm/llms/azure_ai/test_azure_document_intelligence_ocr_transformation.py +++ b/tests/test_litellm/llms/azure_ai/test_azure_document_intelligence_ocr_transformation.py @@ -1,10 +1,10 @@ +from typing import Final from unittest.mock import MagicMock import httpx import pytest from litellm.exceptions import UnsupportedParamsError - from litellm.llms.azure_ai.ocr.document_intelligence.transformation import ( AzureDocumentIntelligenceOCRConfig, ) @@ -96,6 +96,31 @@ def _completed_response(payload: dict) -> httpx.Response: ) +@pytest.mark.parametrize( + "credential_header", + [ + {"Ocp-Apim-Subscription-Key": "subscription-key"}, + {"Authorization": "Bearer entra-token"}, + ], +) +def test_polling_target_reuses_only_same_origin_credential(credential_header: dict[str, str]) -> None: + operation_url: Final = "https://example.cognitiveservices.azure.com/analyzeResults/xyz" + initial_response: Final = httpx.Response( + status_code=202, + headers={"Operation-Location": operation_url}, + request=httpx.Request( + "POST", + "https://example.cognitiveservices.azure.com/documentintelligence/documentModels/layout:analyze", + headers={**credential_header, "x-trace-id": "request-only"}, + ), + ) + + target, poll_headers = AzureDocumentIntelligenceOCRConfig()._get_polling_target(initial_response) + + assert target == operation_url + assert poll_headers == credential_header + + def _assert_native_fields_preserved(serialized: dict) -> None: assert serialized["content"] == "Invoice\nInvoice No: INV-12345\nTotal: $100.00" assert serialized["tables"] == AZURE_TABLES diff --git a/tests/test_litellm/llms/test_polling_url_origin_match.py b/tests/test_litellm/llms/test_polling_url_origin_match.py index f1f910bc731..9df61c9fc31 100644 --- a/tests/test_litellm/llms/test_polling_url_origin_match.py +++ b/tests/test_litellm/llms/test_polling_url_origin_match.py @@ -12,10 +12,8 @@ share an origin with the original request URL. from unittest.mock import MagicMock, patch -import httpx import pytest - # Azure DALL-E sync + async paths route through ``assert_same_origin`` # the same way as the cases below. The helper itself is unit-tested in # ``tests/test_litellm/litellm_core_utils/test_url_utils.py``; the @@ -43,17 +41,22 @@ def test_azure_di_sync_rejects_cross_origin_polling(): ) raw_response.request.headers = {"Ocp-Apim-Subscription-Key": "leak-me"} - with pytest.raises(ValueError, match="rejected polling URL"): - config.transform_ocr_response( - model="azure-doc-intel", - raw_response=raw_response, - logging_obj=MagicMock(), - request_data={}, - optional_params={}, - litellm_params={}, - encoding=None, - response={}, - ) + with patch( # test-quality-ok: verifies URL validation runs before the internal client factory + "litellm.llms.custom_httpx.http_handler._get_httpx_client" + ) as get_client: + with pytest.raises(ValueError, match="rejected polling URL"): + config.transform_ocr_response( + model="azure-doc-intel", + raw_response=raw_response, + logging_obj=MagicMock(), + request_data={}, + optional_params={}, + litellm_params={}, + encoding=None, + response={}, + ) + + get_client.assert_not_called() # ── Black Forest Labs polling ───────────────────────────────────────────────── diff --git a/tests/test_litellm/ocr/test_rust_bridge.py b/tests/test_litellm/ocr/test_rust_bridge.py index 14b5294985e..4b212d80f00 100644 --- a/tests/test_litellm/ocr/test_rust_bridge.py +++ b/tests/test_litellm/ocr/test_rust_bridge.py @@ -295,6 +295,19 @@ def test_explicit_false_overrides_process_enable(): assert ocr_main._rust_ocr_enabled(build_prepared_request(litellm_params={"rust": False})) is False +@pytest.mark.parametrize( + "litellm_params", + [ + {"azure_ad_token_provider": lambda: "token"}, + {"azure_username": "user", "azure_password": "password"}, + ], +) +def test_unsupported_azure_credentials_stay_on_python(litellm_params: dict[str, object]) -> None: + prepared: Final = build_prepared_request(custom_llm_provider="azure_ai", litellm_params=litellm_params) + + assert ocr_main._rust_ocr_supported(prepared) is False + + def test_load_rust_ocr_returns_injected_impl(): bridge = RecordingBridge() litellm.rust(True) @@ -662,6 +675,7 @@ def test_prepare_rust_ocr_call_forwards_vertex_routing_metadata(): "vertex_location": "us-central1", "vertex_credentials": "redacted", }, + extra_headers={"x-trace-id": "trace-1"}, optional_params={"include_image_base64": True}, timeout=None, ), @@ -672,7 +686,70 @@ 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", } + assert bridge.calls[0]["api_key"] == "sk-test" + assert bridge.calls[0]["extra_headers"] == {"x-trace-id": "trace-1"} + + +@pytest.mark.parametrize( + "litellm_params,resolved_assertion,expected_auth_params", + [ + ( + {"azure_ad_token": "entra-token", "azure_scope": "scope/.default"}, + None, + {"azure_ad_token": "entra-token", "azure_scope": "scope/.default"}, + ), + ( + {"tenant_id": "tenant", "client_id": "client", "client_secret": "secret"}, + None, + {"tenant_id": "tenant", "client_id": "client", "client_secret": "secret"}, + ), + ( + {"azure_ad_token": "oidc/token-reference", "tenant_id": "tenant", "client_id": "client"}, + "assertion", + { + "azure_ad_token": "oidc/token-reference", + "azure_ad_client_assertion": "assertion", + "tenant_id": "tenant", + "client_id": "client", + }, + ), + ], +) +def test_prepare_rust_ocr_call_forwards_azure_auth_configuration( + litellm_params: dict[str, object], + resolved_assertion: str | None, + expected_auth_params: dict[str, object], +) -> None: + bridge: Final = RecordingBridge() + logging_obj: Final = RecordingLogging() + litellm.rust(True) + rust_bridge.set_rust_ocr(ocr=bridge) + + ocr_main._run_rust_ocr( + prepared_request=build_prepared_request( + logging_obj=logging_obj, + custom_llm_provider="azure_ai", + model="doc-intelligence/prebuilt-layout", + api_key=None, + extra_headers={"x-trace-id": "trace-1"}, + litellm_params=litellm_params, + timeout=None, + ), + resolve_api_key=lambda name: resolved_assertion if name == "oidc/token-reference" else None, + ) + + assert bridge.calls[0]["api_key"] is None + assert bridge.calls[0]["extra_headers"] == {"x-trace-id": "trace-1"} + assert bridge.calls[0]["optional_params"] == { + **expected_auth_params, + "azure_ad_token_refresh": False, + } + assert logging_obj.pre_call_kwargs is not None + logged_input: Final = logging_obj.pre_call_kwargs["additional_args"] + assert "document" not in str(logged_input) + assert all(name not in str(logged_input) for name in expected_auth_params) def test_prepare_rust_ocr_call_resolves_vertex_routing_metadata_from_secret_manager(): @@ -758,7 +835,7 @@ def test_run_rust_ocr_runs_pre_call_logging(): assert logging_obj.pre_call_kwargs["input"] == "OCR document processing" additional_args = logging_obj.pre_call_kwargs["additional_args"] complete_input = additional_args["complete_input_dict"] - assert complete_input["document"] == DOCUMENT + assert "document" not in complete_input assert complete_input["include_image_base64"] is True assert additional_args["api_base"] == "https://api.mistral.ai/v1/ocr" assert additional_args["headers"] == { diff --git a/type-discipline-budget.json b/type-discipline-budget.json index 2f30972675c..d21f2bf12b1 100644 --- a/type-discipline-budget.json +++ b/type-discipline-budget.json @@ -3,7 +3,7 @@ "limit": 22292 }, "LIT002": { - "limit": 26749 + "limit": 26744 }, "LIT003": { "limit": 261