diff --git a/litellm-rust/Cargo.lock b/litellm-rust/Cargo.lock index 422dfb20065..bb957cc50da 100644 --- a/litellm-rust/Cargo.lock +++ b/litellm-rust/Cargo.lock @@ -2,6 +2,15 @@ # It is not intended for manual editing. version = 4 +[[package]] +name = "android_system_properties" +version = "0.1.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "819e7219dbd41043ac279b19830f2efc897156490d7fd6ea916720117ee66311" +dependencies = [ + "libc", +] + [[package]] name = "async-trait" version = "0.1.89" @@ -104,6 +113,15 @@ dependencies = [ "generic-array", ] +[[package]] +name = "block-buffer" +version = "0.12.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d2f6c7dbe95a6ed67ad9f18e57daf93a2f034c524b99fd2b76d18fdfeb6660aa" +dependencies = [ + "hybrid-array", +] + [[package]] name = "bumpalo" version = "3.20.3" @@ -121,6 +139,9 @@ name = "bytes" version = "1.12.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "8ae3f5d315924270530207e2a68396c3cc547f6dca3fbdca317cfb1a51edb593" +dependencies = [ + "serde", +] [[package]] name = "cc" @@ -144,6 +165,45 @@ version = "0.2.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "613afe47fcd5fac7ccf1db93babcb082c5994d996f20b8b159f2ad1658eb5724" +[[package]] +name = "chacha20" +version = "0.10.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d524456ba66e72eb8b115ff89e01e497f8e6d11d78b70b1aa13c0fbd97540a81" +dependencies = [ + "cfg-if", + "cpufeatures 0.3.0", + "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 = "cmov" +version = "0.5.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0c9ea0ac24bc397ab3c98583a3c9ba74fa56b09a4449bbe172b9b1ddb016027a" + +[[package]] +name = "combine" +version = "4.6.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ba5a308b75df32fe02788e748662718f03fde005016435c444eea572398219fd" +dependencies = [ + "bytes", + "memchr", +] + [[package]] name = "core-foundation" version = "0.10.1" @@ -169,6 +229,15 @@ dependencies = [ "libc", ] +[[package]] +name = "cpufeatures" +version = "0.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8b2a41393f66f16b0823bb79094d54ac5fbd34ab292ddafb9a0456ac9f87d201" +dependencies = [ + "libc", +] + [[package]] name = "crypto-common" version = "0.1.7" @@ -179,20 +248,94 @@ dependencies = [ "typenum", ] +[[package]] +name = "crypto-common" +version = "0.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ce6e4c961d6cd6c9a86db418387425e8bdeaf05b3c8bc1411e6dca4c252f1453" +dependencies = [ + "hybrid-array", +] + +[[package]] +name = "ctutils" +version = "0.4.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7d5515a3834141de9eafb9717ad39eea8247b5674e6066c404e8c4b365d2a29e" +dependencies = [ + "cmov", +] + +[[package]] +name = "darling" +version = "0.21.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9cdf337090841a411e2a7f3deb9187445851f91b309c0c0a29e05f74a00a48c0" +dependencies = [ + "darling_core", + "darling_macro", +] + +[[package]] +name = "darling_core" +version = "0.21.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1247195ecd7e3c85f83c8d2a366e4210d588e802133e1e355180a9870b517ea4" +dependencies = [ + "fnv", + "ident_case", + "proc-macro2", + "quote", + "strsim", + "syn", +] + +[[package]] +name = "darling_macro" +version = "0.21.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d38308df82d1080de0afee5d069fa14b0326a88c14f15c5ccda35b4a6c414c81" +dependencies = [ + "darling_core", + "quote", + "syn", +] + [[package]] name = "data-encoding" version = "2.11.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "a4ae5f15dda3c708c0ade84bfee31ccab44a3da4f88015ed22f63732abe300c8" +[[package]] +name = "deranged" +version = "0.5.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7cd812cc2bc1d69d4764bd80df88b4317eaef9e773c75226407d9bc0876b211c" +dependencies = [ + "powerfmt", + "serde_core", +] + [[package]] name = "digest" version = "0.10.7" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9ed9a281f7bc9b7576e61468ba615a66a5c8cfdff42420a70aa82701a3b1e292" dependencies = [ - "block-buffer", - "crypto-common", + "block-buffer 0.10.4", + "crypto-common 0.1.7", +] + +[[package]] +name = "digest" +version = "0.11.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f1dd6dbb5841937940781866fa1281a1ff7bd3bf827091440879f9994983d5c2" +dependencies = [ + "block-buffer 0.12.1", + "crypto-common 0.2.2", + "ctutils", ] [[package]] @@ -206,12 +349,28 @@ dependencies = [ "syn", ] +[[package]] +name = "dyn-clone" +version = "1.0.20" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d0881ea181b1df73ff77ffaaf9c7544ecc11e82fba9b5f27b262a3c73a332555" + [[package]] name = "equivalent" 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.52.0", +] + [[package]] name = "find-msvc-tools" version = "0.1.9" @@ -353,11 +512,99 @@ dependencies = [ "cfg-if", "js-sys", "libc", - "r-efi", + "r-efi 5.3.0", "wasip2", "wasm-bindgen", ] +[[package]] +name = "getrandom" +version = "0.4.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "300e883d756b2e4ec94e02791f39b04b522276138852cfc41d9fb7e904106099" +dependencies = [ + "cfg-if", + "libc", + "r-efi 6.0.0", + "rand_core 0.10.1", +] + +[[package]] +name = "google-cloud-auth" +version = "1.9.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "54a26c047222f874ea87177368ad07c65a9f66534ad3a3f9401f1322c802ccac" +dependencies = [ + "async-trait", + "base64", + "bytes", + "chrono", + "google-cloud-gax", + "hex", + "hmac", + "http", + "reqwest 0.13.4", + "rustc_version", + "rustls", + "rustls-pki-types", + "serde", + "serde_json", + "sha2 0.11.0", + "thiserror 2.0.18", + "time", + "tokio", + "url", +] + +[[package]] +name = "google-cloud-gax" +version = "1.9.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7dc387965cc2efc28d73896d6707125815c16792c23c33a0c67794f3d6e31cc8" +dependencies = [ + "base64", + "bytes", + "futures", + "google-cloud-rpc", + "google-cloud-wkt", + "http", + "pin-project", + "rand 0.10.2", + "serde", + "serde_json", + "thiserror 2.0.18", + "tokio", +] + +[[package]] +name = "google-cloud-rpc" +version = "1.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7e3b123ea17ff20539fbdf145e6213e0464cc0a30b0d078a68bf90405ef17fb7" +dependencies = [ + "bytes", + "google-cloud-wkt", + "serde", + "serde_json", + "serde_with", +] + +[[package]] +name = "google-cloud-wkt" +version = "1.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0b30ccefdb9276269bb0336afe207c5e0ba1a544a5eb0034763af051f2b9eb63" +dependencies = [ + "base64", + "bytes", + "serde", + "serde_json", + "serde_with", + "thiserror 2.0.18", + "time", + "url", +] + [[package]] name = "h2" version = "0.4.15" @@ -370,13 +617,19 @@ dependencies = [ "futures-core", "futures-sink", "http", - "indexmap", + "indexmap 2.14.0", "slab", "tokio", "tokio-util", "tracing", ] +[[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" @@ -389,6 +642,21 @@ version = "0.5.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "2304e00983f87ffb38b55b444b5e3b60a884b5d30c0fca7d82fe33449bbe55ea" +[[package]] +name = "hex" +version = "0.4.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7f24254aa9a54b5c858eaee2f5bccdb46aaf0e486a595ed5fd8f86ba55232a70" + +[[package]] +name = "hmac" +version = "0.13.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6303bc9732ae41b04cb554b844a762b4115a61bfaa81e3e83050991eeb56863f" +dependencies = [ + "digest 0.11.3", +] + [[package]] name = "http" version = "1.4.2" @@ -434,6 +702,15 @@ version = "1.0.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "df3b46402a9d5adb4c86a0cf463f42e19994e3ee891101b1841f30a545cb49a9" +[[package]] +name = "hybrid-array" +version = "0.4.13" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "818356c5132c1fede50f837ca96afbe78ff42413047f4abb886217845e1b6c8c" +dependencies = [ + "typenum", +] + [[package]] name = "hyper" version = "1.10.1" @@ -495,6 +772,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" @@ -577,6 +878,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" @@ -598,6 +905,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" @@ -605,7 +923,9 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "d466e9454f08e4a911e14806c24e16fba1b4c121d1ea474396f396069cf949d9" dependencies = [ "equivalent", - "hashbrown", + "hashbrown 0.17.1", + "serde", + "serde_core", ] [[package]] @@ -629,6 +949,55 @@ version = "1.0.18" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "8f42a60cbdf9a97f5d2305f08a87dc4e09308d1276d28c869c684d7777685682" +[[package]] +name = "jni" +version = "0.22.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5efd9a482cf3a427f00d6b35f14332adc7902ce91efb778580e180ff90fa3498" +dependencies = [ + "cfg-if", + "combine", + "jni-macros", + "jni-sys", + "log", + "simd_cesu8", + "thiserror 2.0.18", + "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", +] + +[[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", +] + [[package]] name = "js-sys" version = "0.3.103" @@ -654,12 +1023,14 @@ dependencies = [ "base64", "futures-channel", "futures-util", + "google-cloud-auth", "litellm-core", "pyo3", - "reqwest", + "reqwest 0.12.28", + "rustls", "serde", "serde_json", - "sha2", + "sha2 0.10.9", "subtle", "tokio", "tokio-tungstenite", @@ -753,6 +1124,21 @@ dependencies = [ "windows-sys 0.61.2", ] +[[package]] +name = "num-conv" +version = "0.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "51d515d32fb182ee37cda2ccdcb92950d6a3c2893aa280e540671c2cd0f3b1d9" + +[[package]] +name = "num-traits" +version = "0.2.19" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "071dfc062690e90b734c0b2273ce72ad0ffa95f0c74596bc250dcfd960262841" +dependencies = [ + "autocfg", +] + [[package]] name = "once_cell" version = "1.21.4" @@ -771,6 +1157,26 @@ version = "2.3.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9b4f627cb1b25917193a259e49bdad08f671f8d9708acfd5fe0a8c1455d87220" +[[package]] +name = "pin-project" +version = "1.1.13" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2466b2336ed02bcdca6b294417127b90ec92038d1d5c4fbeac971a922e0e0924" +dependencies = [ + "pin-project-internal", +] + +[[package]] +name = "pin-project-internal" +version = "1.1.13" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c96395f0a926bc13b1c17622aaddda1ecb55d49c8f1bf9777e4d877800a43f8b" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "pin-project-lite" version = "0.2.17" @@ -792,6 +1198,12 @@ dependencies = [ "zerovec", ] +[[package]] +name = "powerfmt" +version = "0.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "439ee305def115ba05938db6eb1644ff94165c5ab5e9420d1c1bcedbba909391" + [[package]] name = "ppv-lite86" version = "0.2.21" @@ -956,6 +1368,12 @@ version = "5.3.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "69cdb34c158ceb288df11e18b4bd39de994f6657d83847bdffdbd7f346754b0f" +[[package]] +name = "r-efi" +version = "6.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f8dcc9c7d52a811697d2151c701e0d08956f92b0e24136cf4cf27b57a6a0d9bf" + [[package]] name = "rand" version = "0.8.6" @@ -977,6 +1395,17 @@ dependencies = [ "rand_core 0.9.5", ] +[[package]] +name = "rand" +version = "0.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c7f5fa3a058cd35567ef9bfa5e75732bee0f9e4c55fa90477bef2dfcdbc4be80" +dependencies = [ + "chacha20", + "getrandom 0.4.3", + "rand_core 0.10.1", +] + [[package]] name = "rand_chacha" version = "0.3.1" @@ -1015,6 +1444,32 @@ dependencies = [ "getrandom 0.3.4", ] +[[package]] +name = "rand_core" +version = "0.10.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "63b8176103e19a2643978565ca18b50549f6101881c443590420e4dc998a3c69" + +[[package]] +name = "ref-cast" +version = "1.0.25" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f354300ae66f76f1c85c5f84693f0ce81d747e2c3f21a45fef496d89c960bf7d" +dependencies = [ + "ref-cast-impl", +] + +[[package]] +name = "ref-cast-impl" +version = "1.0.25" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b7186006dcb21920990093f30e3dea63b7d6e977bf1256be20c3563a5db070da" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "reqwest" version = "0.12.28" @@ -1059,6 +1514,43 @@ dependencies = [ "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", + "http", + "http-body", + "http-body-util", + "hyper", + "hyper-rustls", + "hyper-util", + "js-sys", + "log", + "percent-encoding", + "pin-project-lite", + "rustls", + "rustls-pki-types", + "rustls-platform-verifier", + "serde", + "serde_json", + "serde_urlencoded", + "sync_wrapper", + "tokio", + "tokio-rustls", + "tower", + "tower-http", + "tower-service", + "url", + "wasm-bindgen", + "wasm-bindgen-futures", + "web-sys", +] + [[package]] name = "ring" version = "0.17.14" @@ -1079,12 +1571,22 @@ version = "2.1.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "94300abf3f1ae2e2b8ffb7b58043de3d399c73fa6f4b73826402a5c457614dbe" +[[package]] +name = "rustc_version" +version = "0.4.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cfcb3a22ef46e85b45de6ee7e79d063319ebb6594faafcf1c225ea92ab6e9b92" +dependencies = [ + "semver", +] + [[package]] name = "rustls" version = "0.23.41" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6b92b125634d9b795e7beca796cc790df15a7fb38323bf3196fda83292d06b1f" dependencies = [ + "log", "once_cell", "ring", "rustls-pki-types", @@ -1115,6 +1617,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", + "rustls-native-certs", + "rustls-platform-verifier-android", + "rustls-webpki", + "security-framework", + "security-framework-sys", + "webpki-root-certs", + "windows-sys 0.52.0", +] + +[[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.103.13" @@ -1138,6 +1667,15 @@ version = "1.0.23" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9774ba4a74de5f7b1c1451ed6cd5285a32eddb5cccb8cc655a4e50009e06477f" +[[package]] +name = "same-file" +version = "1.0.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "93fc1dc3aaa9bfed95e02e6eadabb4baf7e3078b0bd1b4d7b6b0b68378900502" +dependencies = [ + "winapi-util", +] + [[package]] name = "schannel" version = "0.1.29" @@ -1147,6 +1685,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.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a2b42f36aa1cd011945615b92222f6bf73c599a102a300334cd7f8dbeec726cc" +dependencies = [ + "dyn-clone", + "ref-cast", + "serde", + "serde_json", +] + [[package]] name = "security-framework" version = "3.7.0" @@ -1170,6 +1732,12 @@ dependencies = [ "libc", ] +[[package]] +name = "semver" +version = "1.0.28" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8a7852d02fc848982e0c167ef163aaff9cd91dc640ba85e263cb1ce46fae51cd" + [[package]] name = "serde" version = "1.0.228" @@ -1236,6 +1804,37 @@ dependencies = [ "serde", ] +[[package]] +name = "serde_with" +version = "3.17.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "381b283ce7bc6b476d903296fb59d0d36633652b633b27f64db4fb46dcbfc3b9" +dependencies = [ + "base64", + "chrono", + "hex", + "indexmap 1.9.3", + "indexmap 2.14.0", + "schemars 0.9.0", + "schemars 1.2.1", + "serde_core", + "serde_json", + "serde_with_macros", + "time", +] + +[[package]] +name = "serde_with_macros" +version = "3.17.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a6d4e30573c8cb306ed6ab1dca8423eec9a463ea0e155f45399455e0368b27e0" +dependencies = [ + "darling", + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "sha1" version = "0.10.6" @@ -1243,8 +1842,8 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "e3bf829a2d51ab4a5ddf1352d8470c140cadc8301b2ae1789db023f01cedd6ba" dependencies = [ "cfg-if", - "cpufeatures", - "digest", + "cpufeatures 0.2.17", + "digest 0.10.7", ] [[package]] @@ -1254,8 +1853,19 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "a7507d819769d01a365ab707794a4084392c824f54a7a6a7862f8c3d0892b283" dependencies = [ "cfg-if", - "cpufeatures", - "digest", + "cpufeatures 0.2.17", + "digest 0.10.7", +] + +[[package]] +name = "sha2" +version = "0.11.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "446ba717509524cb3f22f17ecc096f10f4822d76ab5c0b9822c5f9c284e825f4" +dependencies = [ + "cfg-if", + "cpufeatures 0.3.0", + "digest 0.11.3", ] [[package]] @@ -1264,6 +1874,32 @@ 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_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" @@ -1283,7 +1919,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "52d1cfed4120b4d927bf7c0f86d2087a4a7d6027c906d9f9d525a80573b9be51" dependencies = [ "libc", - "windows-sys 0.61.2", + "windows-sys 0.60.2", ] [[package]] @@ -1292,6 +1928,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" @@ -1375,6 +2017,37 @@ dependencies = [ "syn", ] +[[package]] +name = "time" +version = "0.3.45" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9e442fc33d7fdb45aa9bfeb312c095964abdf596f7567261062b2a7107aaabd" +dependencies = [ + "deranged", + "itoa", + "num-conv", + "powerfmt", + "serde_core", + "time-core", + "time-macros", +] + +[[package]] +name = "time-core" +version = "0.1.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8b36ee98fd31ec7426d599183e8fe26932a8dc1fb76ddb6214d05493377d34ca" + +[[package]] +name = "time-macros" +version = "0.2.25" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "71e552d1249bf61ac2a52db88179fd0673def1e1ad8243a00d9ec9ed71fee3dd" +dependencies = [ + "num-conv", + "time-core", +] + [[package]] name = "tinystr" version = "0.8.3" @@ -1410,6 +2083,7 @@ dependencies = [ "libc", "mio", "pin-project-lite", + "signal-hook-registry", "socket2", "tokio-macros", "windows-sys 0.61.2", @@ -1617,6 +2291,16 @@ version = "0.9.5" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "0b928f33d975fc6ad9f86c8f283853ad26bdd5b10b7f1542aa2fa15e2289105a" +[[package]] +name = "walkdir" +version = "2.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "29790946404f91d9c5d06f9874efddea1dc06c5efe94541a7d6863108e3a5e4b" +dependencies = [ + "same-file", + "winapi-util", +] + [[package]] name = "want" version = "0.3.1" @@ -1729,6 +2413,15 @@ dependencies = [ "wasm-bindgen", ] +[[package]] +name = "webpki-root-certs" +version = "1.0.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0d46a5a140e6f7afeccd8eae97eff335163939eac8b929834875168b29b3d267" +dependencies = [ + "rustls-pki-types", +] + [[package]] name = "webpki-roots" version = "1.0.8" @@ -1738,12 +2431,74 @@ dependencies = [ "rustls-pki-types", ] +[[package]] +name = "winapi-util" +version = "0.1.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c2a7b1c03c876122aa43f3020e6c3c3ee5c05081c9a00739faf7503aeba10d22" +dependencies = [ + "windows-sys 0.52.0", +] + +[[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", +] + +[[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", +] + [[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" diff --git a/litellm-rust/Cargo.toml b/litellm-rust/Cargo.toml index 00af23b4c00..fa933856ba7 100644 --- a/litellm-rust/Cargo.toml +++ b/litellm-rust/Cargo.toml @@ -4,12 +4,13 @@ members = [ "crates/ai-gateway", "crates/python-bridge", ] -resolver = "2" +resolver = "3" [workspace.package] edition = "2021" license = "MIT" repository = "https://github.com/BerriAI/litellm" +rust-version = "1.86" [workspace.dependencies] litellm-core = { path = "crates/core" } @@ -21,6 +22,8 @@ rand = "0.8" reqwest = { version = "0.12", default-features = false, features = ["blocking", "json", "multipart", "rustls-tls", "http2", "stream"] } serde = { version = "1.0", features = ["derive"] } serde_json = "1.0" +google-cloud-auth = { version = "=1.9.0", default-features = false } +rustls = { version = "=0.23.41", default-features = false, features = ["ring"] } sha2 = "0.10" subtle = "2" thiserror = "2.0" diff --git a/litellm-rust/crates/ai-gateway/Cargo.toml b/litellm-rust/crates/ai-gateway/Cargo.toml index 2f414159158..760f43b6829 100644 --- a/litellm-rust/crates/ai-gateway/Cargo.toml +++ b/litellm-rust/crates/ai-gateway/Cargo.toml @@ -4,6 +4,7 @@ version = "0.1.0" edition.workspace = true license.workspace = true repository.workspace = true +rust-version.workspace = true [lib] name = "litellm_ai_gateway" @@ -24,17 +25,17 @@ tokio-tungstenite.workspace = true futures-util.workspace = true serde_json.workspace = true base64.workspace = true +google-cloud-auth.workspace = true +rustls.workspace = true +sha2.workspace = true axum = { workspace = true, features = ["ws"], optional = true } serde = { workspace = true, optional = true } subtle = { workspace = true, optional = true } -# sha2 hashes the master key into user_api_key_hash (matches the proxy's -# SHA-256 hash_token) so the plaintext credential never enters a log payload. -sha2 = { workspace = true, optional = true } pyo3 = { workspace = true, features = ["auto-initialize"], optional = true } [features] default = [] -server = ["dep:axum", "dep:subtle", "dep:serde", "dep:sha2"] +server = ["dep:axum", "dep:subtle", "dep:serde"] # Build the gateway's config from the proxy YAML via an embedded Python # interpreter (links libpython; requires `litellm` importable at runtime). python-config = ["dep:pyo3"] diff --git a/litellm-rust/crates/ai-gateway/src/constants.rs b/litellm-rust/crates/ai-gateway/src/constants.rs index 10689943bff..879a057cf39 100644 --- a/litellm-rust/crates/ai-gateway/src/constants.rs +++ b/litellm-rust/crates/ai-gateway/src/constants.rs @@ -34,4 +34,10 @@ pub(crate) const DEFAULT_FLUSH_INTERVAL_MS: u64 = 500; #[cfg(feature = "server")] pub(crate) const DEFAULT_PROVIDER: &str = "openai"; +pub(crate) const CLOUD_PLATFORM_SCOPE: &str = "https://www.googleapis.com/auth/cloud-platform"; + +pub(crate) const VERTEXAI_CREDENTIALS_ENV: &str = "VERTEXAI_CREDENTIALS"; + +pub(crate) const VERTEX_CREDENTIALS_CACHE_CAPACITY: usize = 64; + pub(crate) const ENV_REFERENCE_PREFIX: &str = "os.environ/"; diff --git a/litellm-rust/crates/ai-gateway/src/io/mod.rs b/litellm-rust/crates/ai-gateway/src/io/mod.rs index 3b566027646..090c3eddd9f 100644 --- a/litellm-rust/crates/ai-gateway/src/io/mod.rs +++ b/litellm-rust/crates/ai-gateway/src/io/mod.rs @@ -1,3 +1,4 @@ pub mod ocr; pub mod realtime; pub mod realtime_pool; +pub mod vertex_ai; diff --git a/litellm-rust/crates/ai-gateway/src/io/ocr.rs b/litellm-rust/crates/ai-gateway/src/io/ocr.rs index 0d42125bdbd..efd823f6f17 100644 --- a/litellm-rust/crates/ai-gateway/src/io/ocr.rs +++ b/litellm-rust/crates/ai-gateway/src/io/ocr.rs @@ -9,18 +9,23 @@ use std::time::Duration; use litellm_core::error::CoreError; use litellm_core::ocr::transformation::{ - OcrAuthStrategy, OcrDocumentPreparation, OcrResponseHandling, + OcrAuth, OcrAuthStrategy, OcrDocumentPreparation, OcrResponseHandling, +}; +use litellm_core::providers::vertex_ai::ocr::transformation::{ + classify_vertex_bearer, validate_vertex_authorization_headers, VertexTokenSource, }; use litellm_core::CoreResult; use serde_json::{Map, Value}; use crate::config::resolve_env_reference; +use crate::io::vertex_ai::VertexAiBase; mod common_utils; use common_utils::{ - classify_reqwest_error, convert_document_url_to_data_uri, has_header, ocr_provider_config, - poll_document_intelligence, string_headers, truncate_error_body, upload_reducto_document, + classify_reqwest_error, convert_document_url_to_data_uri, has_header, header_values, + ocr_provider_config, poll_document_intelligence, string_headers, truncate_error_body, + upload_reducto_document, }; /// OCR over large documents can take a while; bound it generously rather than @@ -95,9 +100,31 @@ async fn ocr_with_env( let headers = string_headers(request.extra_headers)?; let auth_strategy = config.auth_strategy(); - let api_key = (!has_header(&headers, auth_strategy.header_name())) - .then(|| config.resolve_api_key(api_key.as_deref(), env_lookup)) - .transpose()?; + let api_key = if has_header(&headers, auth_strategy.header_name()) { + if config.ocr_auth() == OcrAuth::VertexOauth { + validate_vertex_authorization_headers(&header_values( + &headers, + auth_strategy.header_name(), + ))?; + } + None + } else { + Some(match config.ocr_auth() { + OcrAuth::ProviderKey => config.resolve_api_key(api_key.as_deref(), env_lookup)?, + OcrAuth::VertexOauth => match classify_vertex_bearer(api_key.as_deref(), env_lookup)? { + VertexTokenSource::Explicit(token) => token, + VertexTokenSource::Mint => { + let credentials = VertexAiBase::resolve_credential_source( + &request.optional_params, + env_lookup, + ); + VertexAiBase::shared() + .get_access_token(credentials.as_deref()) + .await? + } + }, + }) + }; let url = config.complete_url( api_base.as_deref(), model, diff --git a/litellm-rust/crates/ai-gateway/src/io/ocr/common_utils.rs b/litellm-rust/crates/ai-gateway/src/io/ocr/common_utils.rs index e965209a0dd..babee4afbd2 100644 --- a/litellm-rust/crates/ai-gateway/src/io/ocr/common_utils.rs +++ b/litellm-rust/crates/ai-gateway/src/io/ocr/common_utils.rs @@ -94,6 +94,14 @@ pub(super) fn has_header(headers: &[(String, String)], name: &str) -> bool { .any(|(key, _)| key.eq_ignore_ascii_case(name)) } +pub(super) fn header_values<'a>(headers: &'a [(String, String)], name: &str) -> Vec<&'a str> { + headers + .iter() + .filter(|(key, _)| key.eq_ignore_ascii_case(name)) + .map(|(_, value)| value.as_str()) + .collect() +} + fn document_url_field(document: &Value) -> CoreResult> { let Some(object) = document.as_object() else { return Ok(None); diff --git a/litellm-rust/crates/ai-gateway/src/io/ocr/tests.rs b/litellm-rust/crates/ai-gateway/src/io/ocr/tests.rs index 901e20275bc..9bff72bb2c0 100644 --- a/litellm-rust/crates/ai-gateway/src/io/ocr/tests.rs +++ b/litellm-rust/crates/ai-gateway/src/io/ocr/tests.rs @@ -571,3 +571,216 @@ async fn ocr_maps_unregistered_provider_to_invalid_provider() { assert_eq!(err.public_status_code(), Some(400)); assert_eq!(err.public_message(), "Invalid OCR request"); } + +#[tokio::test] +async fn vertex_ocr_sends_explicit_oauth_bearer_to_raw_predict_endpoint() { + let listener = TcpListener::bind("127.0.0.1:0") + .await + .expect("test listener binds"); + let addr = listener.local_addr().expect("listener has local addr"); + + let server = tokio::spawn(async move { + let (mut socket, _) = listener.accept().await.expect("accepts one request"); + let request = read_http_headers(&mut socket).await; + let response_body = r#"{"pages":[{"index":0,"markdown":"ok"}],"model":"mistral-ocr-maas","usage_info":{"pages_processed":1}}"#; + let response = format!( + "HTTP/1.1 200 OK\r\ncontent-type: application/json\r\ncontent-length: {}\r\nconnection: close\r\n\r\n{}", + response_body.len(), + response_body + ); + socket + .write_all(response.as_bytes()) + .await + .expect("writes response"); + request + }); + + let mut optional_params = Map::new(); + optional_params.insert("vertex_project".to_string(), Value::String("proj-1".into())); + optional_params.insert( + "vertex_location".to_string(), + Value::String("global".into()), + ); + + let response = ocr(OcrRequest { + model: "mistral-ocr-maas", + document: json!({ + "type": "image_url", + "image_url": "data:image/png;base64,abc" + }), + api_key: Some("ya29.explicit-oauth-token"), + api_base: Some(&format!("http://{addr}")), + custom_llm_provider: "vertex_ai", + extra_headers: None, + optional_params, + timeout: Some(Duration::from_secs(5)), + }) + .await + .expect("vertex ocr request succeeds"); + + assert_eq!(response["pages"][0]["markdown"], "ok"); + + let request = server.await.expect("server task completes"); + let request_line = request.lines().next().unwrap_or_default(); + assert!( + request_line + .contains("/v1/projects/proj-1/locations/global/publishers/mistralai/models/mistral-ocr-maas:rawPredict"), + "{request}" + ); + let authorization_count = request + .lines() + .filter(|line| line.to_ascii_lowercase().starts_with("authorization:")) + .count(); + assert_eq!(authorization_count, 1, "{request}"); + assert!( + request.contains("authorization: Bearer ya29.explicit-oauth-token") + || request.contains("Authorization: Bearer ya29.explicit-oauth-token"), + "{request}" + ); +} + +#[tokio::test] +async fn vertex_ocr_rejects_google_api_key_shaped_token_before_calling_upstream() { + let mut optional_params = Map::new(); + optional_params.insert("vertex_project".to_string(), Value::String("proj-1".into())); + optional_params.insert( + "vertex_location".to_string(), + Value::String("global".into()), + ); + + let err = ocr(OcrRequest { + model: "mistral-ocr-maas", + document: json!({ + "type": "image_url", + "image_url": "data:image/png;base64,abc" + }), + api_key: Some("AIzaSyExampleApiKeyValue000000000000000"), + api_base: Some("http://192.0.2.1:9"), + custom_llm_provider: "vertex_ai", + extra_headers: None, + optional_params, + timeout: Some(Duration::from_secs(5)), + }) + .await + .expect_err("google api key is rejected"); + + match err { + CoreError::Auth(message) => assert!(message.contains("OAuth"), "{message}"), + other => panic!("expected auth error, got {other:?}"), + } +} + +#[tokio::test] +async fn vertex_ocr_rejects_google_api_key_in_caller_authorization_header_before_upstream() { + let mut optional_params = Map::new(); + optional_params.insert("vertex_project".to_string(), Value::String("proj-1".into())); + optional_params.insert( + "vertex_location".to_string(), + Value::String("global".into()), + ); + + let mut headers = Map::new(); + headers.insert( + "Authorization".to_string(), + Value::String("Bearer AIzaSyExampleApiKeyValue000000000000000".to_string()), + ); + + let err = ocr(OcrRequest { + model: "mistral-ocr-maas", + document: json!({ + "type": "image_url", + "image_url": "data:image/png;base64,abc" + }), + api_key: None, + api_base: Some("http://192.0.2.1:9"), + custom_llm_provider: "vertex_ai", + extra_headers: Some(headers), + optional_params, + timeout: Some(Duration::from_secs(5)), + }) + .await + .expect_err("google api key in authorization header is rejected"); + + match err { + CoreError::Auth(message) => assert!(message.contains("OAuth"), "{message}"), + other => panic!("expected auth error, got {other:?}"), + } +} + +#[tokio::test] +async fn vertex_ocr_rejects_malformed_authorization_scheme_before_upstream() { + let mut optional_params = Map::new(); + optional_params.insert("vertex_project".to_string(), Value::String("proj-1".into())); + optional_params.insert( + "vertex_location".to_string(), + Value::String("global".into()), + ); + + let mut headers = Map::new(); + headers.insert( + "Authorization".to_string(), + Value::String("Basic dXNlcjpwYXNzd29yZA==".to_string()), + ); + + let err = ocr(OcrRequest { + model: "mistral-ocr-maas", + document: json!({ + "type": "image_url", + "image_url": "data:image/png;base64,abc" + }), + api_key: None, + api_base: Some("http://192.0.2.1:9"), + custom_llm_provider: "vertex_ai", + extra_headers: Some(headers), + optional_params, + timeout: Some(Duration::from_secs(5)), + }) + .await + .expect_err("malformed authorization scheme is rejected"); + + match err { + CoreError::Auth(message) => assert!(message.contains("Bearer"), "{message}"), + other => panic!("expected auth error, got {other:?}"), + } +} + +#[tokio::test] +async fn vertex_ocr_rejects_duplicate_authorization_headers_before_upstream() { + let mut optional_params = Map::new(); + optional_params.insert("vertex_project".to_string(), Value::String("proj-1".into())); + optional_params.insert( + "vertex_location".to_string(), + Value::String("global".into()), + ); + + let mut headers = Map::new(); + headers.insert( + "Authorization".to_string(), + Value::String("Bearer ya29.first-token".to_string()), + ); + headers.insert( + "authorization".to_string(), + Value::String("Bearer ya29.second-token".to_string()), + ); + + let err = ocr(OcrRequest { + model: "mistral-ocr-maas", + document: json!({ + "type": "image_url", + "image_url": "data:image/png;base64,abc" + }), + api_key: None, + api_base: Some("http://192.0.2.1:9"), + custom_llm_provider: "vertex_ai", + extra_headers: Some(headers), + optional_params, + timeout: Some(Duration::from_secs(5)), + }) + .await + .expect_err("duplicate authorization headers are rejected"); + + match err { + CoreError::Auth(message) => assert!(message.contains("exactly one"), "{message}"), + other => panic!("expected auth error, got {other:?}"), + } +} diff --git a/litellm-rust/crates/ai-gateway/src/io/vertex_ai/mod.rs b/litellm-rust/crates/ai-gateway/src/io/vertex_ai/mod.rs new file mode 100644 index 00000000000..7813827cbab --- /dev/null +++ b/litellm-rust/crates/ai-gateway/src/io/vertex_ai/mod.rs @@ -0,0 +1,3 @@ +mod vertex_ai_base; + +pub use vertex_ai_base::VertexAiBase; diff --git a/litellm-rust/crates/ai-gateway/src/io/vertex_ai/vertex_ai_base.rs b/litellm-rust/crates/ai-gateway/src/io/vertex_ai/vertex_ai_base.rs new file mode 100644 index 00000000000..e78aae7e13f --- /dev/null +++ b/litellm-rust/crates/ai-gateway/src/io/vertex_ai/vertex_ai_base.rs @@ -0,0 +1,556 @@ +//! Shared Vertex AI host-layer authentication. +//! +//! Rust counterpart of Python's `litellm/llms/vertex_ai/vertex_llm_base.py` +//! (`VertexBase`) scoped to what current routes need: resolve a credential +//! source (inline JSON, file path, `VERTEXAI_CREDENTIALS`, or ADC), build +//! Google credentials with the cloud-platform scope, cache them by content, +//! and mint OAuth access tokens. Any Vertex route (OCR today, others later) +//! should use this instead of owning Google auth. + +use std::sync::{Once, OnceLock}; + +use google_cloud_auth::credentials::service_account::AccessSpecifier; +use google_cloud_auth::credentials::{ + external_account, impersonated, service_account, user_account, AccessTokenCredentials, + Builder as AdcBuilder, +}; +use litellm_core::cache::in_memory::InMemoryCache; +use litellm_core::error::CoreError; +use litellm_core::CoreResult; +use serde_json::{Map, Value}; +use sha2::{Digest, Sha256}; +use tokio::sync::Mutex; + +use crate::config::resolve_env_reference; +use crate::constants::{ + CLOUD_PLATFORM_SCOPE, VERTEXAI_CREDENTIALS_ENV, VERTEX_CREDENTIALS_CACHE_CAPACITY, +}; + +type CacheKey = [u8; 32]; + +fn ensure_crypto_provider() { + static INSTALL: Once = Once::new(); + INSTALL.call_once(|| { + let _ = rustls::crypto::ring::default_provider().install_default(); + }); +} + +fn adc_cache_key() -> CacheKey { + let mut hasher = Sha256::new(); + hasher.update(b"adc"); + hasher.finalize().into() +} + +fn content_cache_key(contents: &str) -> CacheKey { + let mut hasher = Sha256::new(); + hasher.update(b"inline:"); + hasher.update(contents.trim().as_bytes()); + hasher.finalize().into() +} + +pub struct VertexAiBase { + cache: Mutex>, +} + +impl VertexAiBase { + pub fn new() -> Self { + Self { + cache: Mutex::new(InMemoryCache::new(VERTEX_CREDENTIALS_CACHE_CAPACITY)), + } + } + + pub fn shared() -> &'static VertexAiBase { + static SHARED: OnceLock = OnceLock::new(); + SHARED.get_or_init(VertexAiBase::new) + } + + /// Mirrors Python `VertexBase.safe_get_vertex_ai_credentials`: request + /// params (`vertex_credentials`, then `vertex_ai_credentials`) take + /// precedence over the `VERTEXAI_CREDENTIALS` environment variable. + pub fn resolve_credential_source( + optional_params: &Map, + env_lookup: &(dyn Fn(&str) -> Option + Sync), + ) -> Option { + Self::credential_source_param(optional_params) + .and_then(|source| resolve_env_reference(Some(&source), env_lookup)) + .or_else(|| { + env_lookup(VERTEXAI_CREDENTIALS_ENV) + .map(|value| value.trim().to_string()) + .filter(|value| !value.is_empty()) + }) + } + + fn credential_source_param(optional_params: &Map) -> Option { + ["vertex_credentials", "vertex_ai_credentials"] + .iter() + .find_map(|key| optional_params.get(*key)) + .and_then(|value| match value { + Value::String(raw) => { + let trimmed = raw.trim(); + (!trimmed.is_empty()).then(|| trimmed.to_string()) + } + Value::Object(_) => Some(value.to_string()), + _ => None, + }) + } + + /// Mirrors Python `VertexBase.get_access_token`: load credentials from the + /// given source (or ADC when absent), cache them by content, and return an + /// OAuth access token. The Google auth crate refreshes expired tokens + /// internally, matching Python's cached-credential refresh behavior. + pub async fn get_access_token(&self, credentials: Option<&str>) -> CoreResult { + ensure_crypto_provider(); + let token = self + .resolve_credentials(credentials) + .await? + .access_token() + .await + .map_err(|_| CoreError::Auth("Failed to obtain Vertex access token".to_string()))?; + Ok(token.token) + } + + async fn resolve_credentials( + &self, + credentials: Option<&str>, + ) -> CoreResult { + let (key, built) = match credentials { + None => (adc_cache_key(), None), + Some(raw) => { + let contents = load_credentials_contents(raw).await?; + (content_cache_key(&contents), Some(contents)) + } + }; + if let Some(existing) = self.cache.lock().await.get(&key) { + return Ok(existing); + } + let credentials = match built { + None => build_adc_credentials()?, + Some(contents) => build_from_json(parse_credentials_json(&contents)?)?, + }; + Ok(self.cache.lock().await.get_or_insert(key, credentials)) + } + + #[cfg(test)] + async fn cached_credential_count(&self) -> usize { + self.cache.lock().await.len() + } +} + +impl Default for VertexAiBase { + fn default() -> Self { + Self::new() + } +} + +fn build_adc_credentials() -> CoreResult { + AdcBuilder::default() + .with_scopes([CLOUD_PLATFORM_SCOPE]) + .build_access_token_credentials() + .map_err(|_| CoreError::Auth("Failed to load Vertex ADC credentials".to_string())) +} + +async fn load_credentials_contents(raw: &str) -> CoreResult { + let trimmed = raw.trim(); + if trimmed.starts_with('{') { + Ok(trimmed.to_string()) + } else { + tokio::fs::read_to_string(trimmed) + .await + .map_err(|_| CoreError::Auth("Failed to read Vertex credentials file".to_string())) + } +} + +fn parse_credentials_json(contents: &str) -> CoreResult { + serde_json::from_str(contents.trim()) + .map_err(|_| CoreError::Auth("Vertex credentials are not valid JSON".to_string())) +} + +fn build_from_json(json: Value) -> CoreResult { + let scopes = [CLOUD_PLATFORM_SCOPE]; + let credentials = match json.get("type").and_then(Value::as_str) { + Some("service_account") => service_account::Builder::new(json) + .with_access_specifier(AccessSpecifier::from_scopes(scopes)) + .build_access_token_credentials(), + Some("authorized_user") => user_account::Builder::new(json) + .with_scopes(scopes) + .build_access_token_credentials(), + Some("external_account") => external_account::Builder::new(json) + .with_scopes(scopes) + .build_access_token_credentials(), + Some("impersonated_service_account") => impersonated::Builder::new(json) + .with_scopes(scopes) + .build_access_token_credentials(), + Some(_) => { + return Err(CoreError::Auth( + "Unsupported Vertex credential type".to_string(), + )) + } + None => { + return Err(CoreError::Auth( + "Vertex credentials JSON is missing the required `type` field".to_string(), + )) + } + }; + credentials.map_err(|_| CoreError::Auth("Failed to load Vertex credentials".to_string())) +} + +#[cfg(test)] +mod tests { + use super::*; + use serde_json::json; + use std::sync::Arc; + + fn inline_service_account_json() -> String { + json!({ + "type": "service_account", + "project_id": "proj-1", + "private_key_id": "key-id", + "private_key": "-----BEGIN PRIVATE KEY-----\nnot-a-real-key\n-----END PRIVATE KEY-----\n", + "client_email": "sa@proj-1.iam.gserviceaccount.com" + }) + .to_string() + } + + #[test] + fn content_cache_key_distinguishes_adc_from_inline_and_matches_on_repeat() { + let adc = adc_cache_key(); + let inline = content_cache_key("{\"type\":\"service_account\"}"); + assert_ne!(adc, inline); + assert_eq!( + inline, + content_cache_key(" {\"type\":\"service_account\"} ") + ); + } + + #[test] + fn content_cache_key_differs_for_distinct_credential_sources() { + let adc = adc_cache_key(); + let first = content_cache_key("{\"type\":\"service_account\",\"client_email\":\"a\"}"); + let second = content_cache_key("{\"type\":\"service_account\",\"client_email\":\"b\"}"); + assert_ne!(first, second); + assert_ne!(adc, first); + assert_ne!(adc, second); + } + + #[tokio::test] + async fn file_backed_credentials_key_tracks_content_not_path() { + let path = std::env::temp_dir().join(format!( + "vertex-cred-{}-{:?}.json", + std::process::id(), + std::thread::current().id() + )); + tokio::fs::write( + &path, + b"{\"type\":\"service_account\",\"client_email\":\"old\"}", + ) + .await + .expect("writes first credential file"); + let first = load_credentials_contents(path.to_str().expect("utf-8 path")) + .await + .expect("reads first credential file"); + + tokio::fs::write( + &path, + b"{\"type\":\"service_account\",\"client_email\":\"new\"}", + ) + .await + .expect("rotates credential file"); + let second = load_credentials_contents(path.to_str().expect("utf-8 path")) + .await + .expect("reads rotated credential file"); + tokio::fs::remove_file(&path).await.ok(); + + assert_ne!( + content_cache_key(&first), + content_cache_key(&second), + "rotating a credential file at the same path must not reuse the old cache entry" + ); + } + + #[tokio::test] + async fn resolve_credentials_caches_inline_service_account_per_instance() { + let base = VertexAiBase::new(); + let inline = inline_service_account_json(); + base.resolve_credentials(Some(&inline)) + .await + .expect("first resolve builds credentials"); + base.resolve_credentials(Some(&inline)) + .await + .expect("second resolve reuses cache"); + assert_eq!(base.cached_credential_count().await, 1); + } + + #[tokio::test] + async fn resolve_credentials_reads_service_account_from_file() { + let path = std::env::temp_dir().join(format!( + "vertex-file-cred-{}-{:?}.json", + std::process::id(), + std::thread::current().id() + )); + tokio::fs::write(&path, inline_service_account_json()) + .await + .expect("writes credential file"); + let base = VertexAiBase::new(); + base.resolve_credentials(Some(path.to_str().expect("utf-8 path"))) + .await + .expect("file-backed service account resolves"); + tokio::fs::remove_file(&path).await.ok(); + assert_eq!(base.cached_credential_count().await, 1); + } + + #[tokio::test] + async fn resolve_credentials_single_flight_under_concurrency() { + let base = Arc::new(VertexAiBase::new()); + let inline = Arc::new(inline_service_account_json()); + let tasks: Vec<_> = (0..16) + .map(|_| { + let base = Arc::clone(&base); + let inline = Arc::clone(&inline); + tokio::spawn(async move { + base.resolve_credentials(Some(inline.as_str())) + .await + .expect("concurrent resolve succeeds"); + }) + }) + .collect(); + for task in tasks { + task.await.expect("task completes"); + } + assert_eq!(base.cached_credential_count().await, 1); + } + + #[tokio::test] + async fn build_from_json_builds_service_account_credentials() { + build_from_json(json!({ + "type": "service_account", + "project_id": "proj-1", + "private_key_id": "key-id", + "private_key": "-----BEGIN PRIVATE KEY-----\nnot-a-real-key\n-----END PRIVATE KEY-----\n", + "client_email": "sa@proj-1.iam.gserviceaccount.com" + })) + .expect("service account dispatches and builds"); + } + + #[tokio::test] + async fn build_from_json_builds_authorized_user_credentials() { + build_from_json(json!({ + "type": "authorized_user", + "client_id": "client-id.apps.googleusercontent.com", + "client_secret": "client-secret", + "refresh_token": "refresh-token" + })) + .expect("authorized user dispatches and builds"); + } + + #[tokio::test] + async fn build_from_json_builds_impersonated_service_account_credentials() { + build_from_json(json!({ + "type": "impersonated_service_account", + "service_account_impersonation_url": "https://iamcredentials.googleapis.com/v1/projects/-/serviceAccounts/target@proj-1.iam.gserviceaccount.com:generateAccessToken", + "source_credentials": { + "type": "authorized_user", + "client_id": "client-id.apps.googleusercontent.com", + "client_secret": "client-secret", + "refresh_token": "refresh-token" + } + })) + .expect("impersonated service account dispatches and builds"); + } + + #[tokio::test] + async fn build_from_json_builds_external_account_for_all_standard_source_mechanisms() { + let base = |source: Value| { + json!({ + "type": "external_account", + "audience": "//iam.googleapis.com/projects/1/locations/global/workloadIdentityPools/p/providers/pr", + "subject_token_type": "urn:ietf:params:oauth:token-type:jwt", + "token_url": "https://sts.googleapis.com/v1/token", + "credential_source": source + }) + }; + + build_from_json(base(json!({"file": "/var/run/secrets/token"}))) + .expect("file-sourced external account builds"); + build_from_json(base(json!({ + "url": "https://169.254.169.254/token", + "headers": {"Metadata": "true"}, + "format": {"type": "json", "subject_token_field_name": "access_token"} + }))) + .expect("url-sourced external account builds"); + build_from_json(base(json!({ + "executable": {"command": "/usr/bin/token-helper", "timeout_millis": 5000} + }))) + .expect("executable-sourced external account builds"); + build_from_json(base(json!({ + "environment_id": "aws1", + "region_url": "http://169.254.169.254/latest/meta-data/placement/availability-zone", + "regional_cred_verification_url": "https://sts.{region}.amazonaws.com?Action=GetCallerIdentity&Version=2011-06-15" + }))) + .expect("aws-sourced external account builds"); + } + + #[test] + fn build_from_json_rejects_unknown_credential_type() { + let err = + build_from_json(json!({"type": "totally_made_up"})).expect_err("unknown type rejected"); + assert!(matches!(err, CoreError::Auth(_)), "{err:?}"); + } + + #[test] + fn build_from_json_requires_type_field() { + let err = build_from_json(json!({"client_email": "x"})).expect_err("missing type rejected"); + assert!(matches!(err, CoreError::Auth(_)), "{err:?}"); + } + + #[test] + fn build_from_json_unknown_type_error_omits_attacker_controlled_type() { + let err = build_from_json(json!({"type": "attacker-controlled-type-string"})) + .expect_err("unknown type rejected"); + match err { + CoreError::Auth(message) => { + assert!( + !message.contains("attacker-controlled-type-string"), + "{message}" + ); + } + other => panic!("expected auth error, got {other:?}"), + } + } + + #[test] + fn parse_credentials_json_reports_invalid_json_without_echoing_contents() { + let err = parse_credentials_json("{not-valid-json-secret-value") + .expect_err("invalid json rejected"); + match err { + CoreError::Auth(message) => { + assert!( + !message.contains("not-valid-json-secret-value"), + "{message}" + ); + } + other => panic!("expected auth error, got {other:?}"), + } + } + + #[tokio::test] + async fn load_credentials_contents_missing_file_error_omits_attacker_controlled_path() { + let err = load_credentials_contents("/attacker/controlled/secret-credentials-path.json") + .await + .expect_err("missing file rejected"); + match err { + CoreError::Auth(message) => { + assert!(!message.contains("attacker"), "{message}"); + assert!(!message.contains("secret-credentials-path"), "{message}"); + } + other => panic!("expected auth error, got {other:?}"), + } + } + + #[test] + fn resolve_credential_source_reads_string_object_and_treats_blank_as_absent() { + let mut inline = Map::new(); + inline.insert( + "vertex_credentials".to_string(), + Value::String(" /path/to/sa.json ".into()), + ); + assert_eq!( + VertexAiBase::resolve_credential_source(&inline, &|_| None).as_deref(), + Some("/path/to/sa.json") + ); + + let mut object = Map::new(); + object.insert( + "vertex_credentials".to_string(), + json!({"type": "service_account"}), + ); + assert_eq!( + VertexAiBase::resolve_credential_source(&object, &|_| None).as_deref(), + Some("{\"type\":\"service_account\"}") + ); + + let mut blank = Map::new(); + blank.insert( + "vertex_credentials".to_string(), + Value::String(" ".into()), + ); + assert_eq!( + VertexAiBase::resolve_credential_source(&blank, &|_| None), + None + ); + + assert_eq!( + VertexAiBase::resolve_credential_source(&Map::new(), &|_| None), + None + ); + } + + #[test] + fn resolve_credential_source_prefers_optional_param_then_env() { + let mut params = Map::new(); + params.insert( + "vertex_credentials".to_string(), + Value::String("/from/param.json".into()), + ); + let env = + |key: &str| (key == VERTEXAI_CREDENTIALS_ENV).then(|| "/from/env.json".to_string()); + assert_eq!( + VertexAiBase::resolve_credential_source(¶ms, &env).as_deref(), + Some("/from/param.json") + ); + assert_eq!( + VertexAiBase::resolve_credential_source(&Map::new(), &env).as_deref(), + Some("/from/env.json") + ); + + let blank_env = |key: &str| (key == VERTEXAI_CREDENTIALS_ENV).then(|| " ".to_string()); + assert_eq!( + VertexAiBase::resolve_credential_source(&Map::new(), &blank_env), + None + ); + + assert_eq!( + VertexAiBase::resolve_credential_source(&Map::new(), &|_| None), + None + ); + } + + #[test] + fn resolve_credential_source_falls_back_through_alias_param() { + let mut params = Map::new(); + params.insert( + "vertex_ai_credentials".to_string(), + Value::String("/from/alias.json".into()), + ); + assert_eq!( + VertexAiBase::resolve_credential_source(¶ms, &|_| None).as_deref(), + Some("/from/alias.json") + ); + } + + #[test] + fn resolve_credential_source_resolves_exact_environment_reference() { + let params = Map::from_iter([( + "vertex_credentials".to_string(), + Value::String("os.environ/CUSTOM_VERTEX_CREDENTIALS".into()), + )]); + let env = |key: &str| { + (key == "CUSTOM_VERTEX_CREDENTIALS").then(|| "{\"type\":\"authorized_user\"}".into()) + }; + assert_eq!( + VertexAiBase::resolve_credential_source(¶ms, &env).as_deref(), + Some("{\"type\":\"authorized_user\"}") + ); + } + + #[test] + fn resolve_credential_source_treats_unresolved_reference_as_absent() { + let params = Map::from_iter([( + "vertex_credentials".to_string(), + Value::String("os.environ/MISSING_VERTEX_CREDENTIALS".into()), + )]); + assert_eq!( + VertexAiBase::resolve_credential_source(¶ms, &|_| None), + None + ); + } +} diff --git a/litellm-rust/crates/core/Cargo.toml b/litellm-rust/crates/core/Cargo.toml index 1881bcfa602..2b2bcbd67a3 100644 --- a/litellm-rust/crates/core/Cargo.toml +++ b/litellm-rust/crates/core/Cargo.toml @@ -4,6 +4,7 @@ version = "0.1.0" edition.workspace = true license.workspace = true repository.workspace = true +rust-version.workspace = true [dependencies] rand.workspace = true diff --git a/litellm-rust/crates/core/src/cache/in_memory.rs b/litellm-rust/crates/core/src/cache/in_memory.rs new file mode 100644 index 00000000000..48a1d9d2688 --- /dev/null +++ b/litellm-rust/crates/core/src/cache/in_memory.rs @@ -0,0 +1,89 @@ +use std::collections::{HashMap, VecDeque}; +use std::hash::Hash; + +pub struct InMemoryCache { + capacity: usize, + entries: HashMap, + order: VecDeque, +} + +impl InMemoryCache +where + K: Clone + Eq + Hash, + V: Clone, +{ + pub fn new(capacity: usize) -> Self { + Self { + capacity: capacity.max(1), + entries: HashMap::new(), + order: VecDeque::new(), + } + } + + pub fn get(&self, key: &K) -> Option { + self.entries.get(key).cloned() + } + + pub fn get_or_insert(&mut self, key: K, value: V) -> V { + if let Some(existing) = self.entries.get(&key) { + return existing.clone(); + } + while self.entries.len() >= self.capacity { + let Some(evicted) = self.order.pop_front() else { + break; + }; + self.entries.remove(&evicted); + } + self.order.push_back(key.clone()); + self.entries.insert(key, value.clone()); + value + } + + pub fn len(&self) -> usize { + self.entries.len() + } + + pub fn is_empty(&self) -> bool { + self.entries.is_empty() + } +} + +#[cfg(test)] +mod tests { + use super::InMemoryCache; + + #[test] + fn evicts_oldest_entry_at_capacity() { + let mut cache = InMemoryCache::new(2); + cache.get_or_insert(1, "one"); + cache.get_or_insert(2, "two"); + cache.get_or_insert(3, "three"); + + assert_eq!(cache.get(&1), None); + assert_eq!(cache.get(&2), Some("two")); + assert_eq!(cache.get(&3), Some("three")); + assert_eq!(cache.len(), 2); + } + + #[test] + fn reuses_existing_entry_without_replacement() { + let mut cache = InMemoryCache::new(2); + + assert!(cache.is_empty()); + assert_eq!(cache.get_or_insert(1, "first"), "first"); + assert_eq!(cache.get_or_insert(1, "replacement"), "first"); + assert_eq!(cache.get(&1), Some("first")); + assert_eq!(cache.len(), 1); + } + + #[test] + fn zero_capacity_still_retains_one_entry() { + let mut cache = InMemoryCache::new(0); + cache.get_or_insert(1, "one"); + cache.get_or_insert(2, "two"); + + assert_eq!(cache.get(&1), None); + assert_eq!(cache.get(&2), Some("two")); + assert_eq!(cache.len(), 1); + } +} diff --git a/litellm-rust/crates/core/src/cache/mod.rs b/litellm-rust/crates/core/src/cache/mod.rs new file mode 100644 index 00000000000..fa2e12699f5 --- /dev/null +++ b/litellm-rust/crates/core/src/cache/mod.rs @@ -0,0 +1 @@ +pub mod in_memory; diff --git a/litellm-rust/crates/core/src/constants.rs b/litellm-rust/crates/core/src/constants.rs new file mode 100644 index 00000000000..b6b8c080f76 --- /dev/null +++ b/litellm-rust/crates/core/src/constants.rs @@ -0,0 +1,4 @@ +pub(crate) const VERTEX_GLOBAL_LOCATION: &str = "global"; +pub(crate) const VERTEX_GLOBAL_API_BASE: &str = "https://aiplatform.googleapis.com"; +pub(crate) const GOOGLE_API_KEY_PREFIX: &str = "AIza"; +pub(crate) const BEARER_SCHEME: &str = "Bearer"; diff --git a/litellm-rust/crates/core/src/lib.rs b/litellm-rust/crates/core/src/lib.rs index 2ac479cc725..92dc4f4007d 100644 --- a/litellm-rust/crates/core/src/lib.rs +++ b/litellm-rust/crates/core/src/lib.rs @@ -1,3 +1,5 @@ +pub mod cache; +pub(crate) mod constants; pub mod error; pub mod ocr; pub mod providers; diff --git a/litellm-rust/crates/core/src/ocr/transformation.rs b/litellm-rust/crates/core/src/ocr/transformation.rs index 87fb658a3bd..bd726c1d87d 100644 --- a/litellm-rust/crates/core/src/ocr/transformation.rs +++ b/litellm-rust/crates/core/src/ocr/transformation.rs @@ -27,6 +27,12 @@ pub enum OcrResponseHandling { AzureDocumentIntelligencePoll, } +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum OcrAuth { + ProviderKey, + VertexOauth, +} + #[derive(Clone, Copy, Debug, PartialEq, Eq)] pub enum OcrDocumentPreparation { None, @@ -72,9 +78,17 @@ pub trait OcrProviderConfig: Sync { fn resolve_api_key( &self, - api_key: Option<&str>, - env_lookup: &dyn Fn(&str) -> Option, - ) -> CoreResult; + _api_key: Option<&str>, + _env_lookup: &dyn Fn(&str) -> Option, + ) -> CoreResult { + Err(crate::error::CoreError::Auth( + "provider does not use direct api-key auth".to_string(), + )) + } + + fn ocr_auth(&self) -> OcrAuth { + OcrAuth::ProviderKey + } fn auth_strategy(&self) -> OcrAuthStrategy { OcrAuthStrategy::Bearer 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 7a5b2c6094f..b2e84f25b4a 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,12 +1,15 @@ +use crate::constants::{ + BEARER_SCHEME, GOOGLE_API_KEY_PREFIX, VERTEX_GLOBAL_API_BASE, VERTEX_GLOBAL_LOCATION, +}; use crate::error::{json_type_name, CoreError, CoreResult}; -use crate::ocr::transformation::OcrProviderConfig; +use crate::ocr::transformation::{OcrAuth, OcrProviderConfig}; use crate::ocr::types::{OcrRequestData, OcrResponseData}; use serde_json::{json, Map, Value}; use crate::providers::mistral::ocr::transformation::MISTRAL_OCR_CONFIG; const VERTEX_DEFAULT_LOCATION: &str = "us-central1"; -const VERTEX_DEFAULT_DEEPSEEK_API_BASE: &str = "https://aiplatform.googleapis.com"; +const VERTEX_DEFAULT_DEEPSEEK_API_BASE: &str = VERTEX_GLOBAL_API_BASE; 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"; @@ -40,22 +43,76 @@ pub fn is_deepseek_model(model: &str) -> bool { model.to_ascii_lowercase().contains("deepseek") } -pub fn resolve_vertex_api_key( +#[derive(Clone, Debug, PartialEq, Eq)] +pub enum VertexTokenSource { + Explicit(String), + Mint, +} + +fn is_google_api_key(token: &str) -> bool { + token.starts_with(GOOGLE_API_KEY_PREFIX) +} + +fn google_api_key_not_oauth_error() -> CoreError { + CoreError::Auth( + "Received a Google API key (AIza...) for Vertex AI, which is not an OAuth access token. \ + Provide service-account credentials/ADC or an OAuth access token instead" + .to_string(), + ) +} + +pub fn classify_vertex_bearer( api_key: Option<&str>, env_lookup: &dyn Fn(&str) -> Option, -) -> CoreResult { - api_key +) -> CoreResult { + let token = api_key .map(str::trim) .filter(|key| !key.is_empty()) .map(str::to_string) .or_else(|| env_lookup(VERTEX_AI_API_KEY_ENV).filter(|key| !key.trim().is_empty())) - .or_else(|| env_lookup(VERTEXAI_API_KEY_ENV).filter(|key| !key.trim().is_empty())) - .ok_or_else(|| { - CoreError::Auth( - "Missing Vertex AI access token - pass api_key or provide Authorization via extra_headers" - .to_string(), - ) - }) + .or_else(|| env_lookup(VERTEXAI_API_KEY_ENV).filter(|key| !key.trim().is_empty())); + + match token { + Some(token) if is_google_api_key(token.trim()) => Err(google_api_key_not_oauth_error()), + Some(token) => Ok(VertexTokenSource::Explicit(token.trim().to_string())), + None => Ok(VertexTokenSource::Mint), + } +} + +fn malformed_vertex_authorization_error() -> CoreError { + CoreError::Auth( + "Vertex AI requires exactly one `Authorization: Bearer ` header. \ + Provide a valid OAuth Bearer token, or omit the header to mint one from credentials/ADC" + .to_string(), + ) +} + +pub fn validate_vertex_authorization_headers(values: &[&str]) -> CoreResult<()> { + match values { + [] => Ok(()), + [single] => validate_vertex_authorization_value(single), + _ => Err(malformed_vertex_authorization_error()), + } +} + +fn validate_vertex_authorization_value(header_value: &str) -> CoreResult<()> { + let mut parts = header_value.split_whitespace(); + let scheme = parts + .next() + .ok_or_else(malformed_vertex_authorization_error)?; + let token = parts + .next() + .ok_or_else(malformed_vertex_authorization_error)?; + if parts.next().is_some() { + return Err(malformed_vertex_authorization_error()); + } + if !scheme.eq_ignore_ascii_case(BEARER_SCHEME) { + return Err(malformed_vertex_authorization_error()); + } + if is_google_api_key(token) { + return Err(google_api_key_not_oauth_error()); + } + Ok(()) } fn vertex_project( @@ -84,12 +141,22 @@ fn vertex_location( .unwrap_or_else(|| VERTEX_DEFAULT_LOCATION.to_string()) } +fn vertex_base_url(location: &str) -> String { + match location { + VERTEX_GLOBAL_LOCATION => VERTEX_GLOBAL_API_BASE.to_string(), + location if !location.contains('-') => { + format!("https://aiplatform.{location}.rep.googleapis.com") + } + location => format!("https://{location}-aiplatform.googleapis.com"), + } +} + fn vertex_mistral_api_base(api_base: Option<&str>, location: &str) -> String { api_base .map(str::trim) .filter(|value| !value.is_empty()) .map(str::to_string) - .unwrap_or_else(|| format!("https://{location}-aiplatform.googleapis.com")) + .unwrap_or_else(|| vertex_base_url(location)) .trim_end_matches('/') .to_string() } @@ -241,12 +308,8 @@ impl OcrProviderConfig for VertexAiOcrConfig { complete_vertex_mistral_url(api_base, model, optional_params, env_lookup) } - fn resolve_api_key( - &self, - api_key: Option<&str>, - env_lookup: &dyn Fn(&str) -> Option, - ) -> CoreResult { - resolve_vertex_api_key(api_key, env_lookup) + fn ocr_auth(&self) -> OcrAuth { + OcrAuth::VertexOauth } fn requires_data_uri_document(&self) -> bool { @@ -346,12 +409,8 @@ impl OcrProviderConfig for VertexAiDeepSeekOcrConfig { complete_vertex_deepseek_url(api_base, optional_params, env_lookup) } - fn resolve_api_key( - &self, - api_key: Option<&str>, - env_lookup: &dyn Fn(&str) -> Option, - ) -> CoreResult { - resolve_vertex_api_key(api_key, env_lookup) + fn ocr_auth(&self) -> OcrAuth { + OcrAuth::VertexOauth } } @@ -375,6 +434,171 @@ mod tests { ); } + #[test] + fn vertex_mistral_url_uses_global_host_without_region_prefix() { + let params = Map::from_iter([ + ("vertex_project".to_string(), json!("proj-1")), + ("vertex_location".to_string(), json!("global")), + ]); + + let url = complete_vertex_mistral_url(None, "mistral-ocr-maas", ¶ms, &|_| None) + .expect("url builds"); + + assert_eq!( + url, + "https://aiplatform.googleapis.com/v1/projects/proj-1/locations/global/publishers/mistralai/models/mistral-ocr-maas:rawPredict" + ); + } + + #[test] + fn vertex_mistral_url_uses_residency_host_for_single_token_location() { + let params = Map::from_iter([ + ("vertex_project".to_string(), json!("proj-1")), + ("vertex_location".to_string(), json!("eu")), + ]); + + let url = complete_vertex_mistral_url(None, "mistral-ocr-maas", ¶ms, &|_| None) + .expect("url builds"); + + assert_eq!( + url, + "https://aiplatform.eu.rep.googleapis.com/v1/projects/proj-1/locations/eu/publishers/mistralai/models/mistral-ocr-maas:rawPredict" + ); + } + + #[test] + fn vertex_mistral_url_prefers_explicit_api_base() { + let params = Map::from_iter([ + ("vertex_project".to_string(), json!("proj-1")), + ("vertex_location".to_string(), json!("global")), + ]); + + let url = complete_vertex_mistral_url( + Some("https://custom.example.com/"), + "mistral-ocr-maas", + ¶ms, + &|_| None, + ) + .expect("url builds"); + + assert_eq!( + url, + "https://custom.example.com/v1/projects/proj-1/locations/global/publishers/mistralai/models/mistral-ocr-maas:rawPredict" + ); + } + + #[test] + fn classify_vertex_bearer_uses_explicit_oauth_token() { + let source = classify_vertex_bearer(Some(" ya29.oauth-token "), &|_| None) + .expect("token classifies"); + assert_eq!( + source, + VertexTokenSource::Explicit("ya29.oauth-token".to_string()) + ); + } + + #[test] + fn classify_vertex_bearer_reads_oauth_token_from_env() { + let source = classify_vertex_bearer(None, &|key| { + (key == VERTEX_AI_API_KEY_ENV).then(|| "ya29.from-env".to_string()) + }) + .expect("token classifies"); + assert_eq!( + source, + VertexTokenSource::Explicit("ya29.from-env".to_string()) + ); + } + + #[test] + fn classify_vertex_bearer_mints_when_no_token_supplied() { + let source = classify_vertex_bearer(None, &|_| None).expect("token classifies"); + assert_eq!(source, VertexTokenSource::Mint); + } + + #[test] + fn classify_vertex_bearer_rejects_google_api_key() { + let err = classify_vertex_bearer(Some("AIzaSyExampleApiKeyValue"), &|_| None) + .expect_err("google api key is rejected"); + assert!(matches!(err, CoreError::Auth(_)), "{err:?}"); + } + + #[test] + fn classify_vertex_bearer_rejects_google_api_key_from_env() { + let err = classify_vertex_bearer(None, &|key| { + (key == VERTEXAI_API_KEY_ENV).then(|| "AIzaSyExampleApiKeyValue".to_string()) + }) + .expect_err("google api key from env is rejected"); + assert!(matches!(err, CoreError::Auth(_)), "{err:?}"); + } + + #[test] + fn validate_vertex_authorization_headers_rejects_api_key_bearer() { + for header in [ + "Bearer AIzaSyExampleApiKeyValue", + " bearer AIzaSyExampleApiKeyValue ", + "BEARER AIzaSyExampleApiKeyValue", + ] { + let err = validate_vertex_authorization_headers(&[header]) + .expect_err(&format!("api key bearer rejected: {header:?}")); + assert!(matches!(err, CoreError::Auth(_)), "{header:?} -> {err:?}"); + } + } + + #[test] + fn validate_vertex_authorization_headers_rejects_malformed_values() { + for header in [ + "", + " ", + "AIzaSyExampleApiKeyValue", + "ya29.raw-token-without-scheme", + "Basic dXNlcjpwYXNz", + "Token ya29.some-token", + "Bearer2 ya29.token", + "Bearer", + "Bearer ", + "Bearer ya29.token extra-part", + "Bearer ya29.token AIzaExtra", + ] { + let err = validate_vertex_authorization_headers(&[header]) + .expect_err(&format!("expected rejection for {header:?}")); + assert!(matches!(err, CoreError::Auth(_)), "{header:?} -> {err:?}"); + } + } + + #[test] + fn validate_vertex_authorization_headers_rejects_duplicate_headers() { + let err = validate_vertex_authorization_headers(&[ + "Bearer ya29.first-token", + "Bearer ya29.second-token", + ]) + .expect_err("duplicate authorization headers rejected"); + assert!(matches!(err, CoreError::Auth(_)), "{err:?}"); + } + + #[test] + fn validate_vertex_authorization_headers_allows_single_oauth_bearer() { + for header in [ + "Bearer ya29.real-oauth-token", + "bearer ya29.real-oauth-token", + "BEARER ya29.real-oauth-token", + " Bearer ya29.real-oauth-token ", + ] { + validate_vertex_authorization_headers(&[header]) + .unwrap_or_else(|err| panic!("oauth bearer allowed: {header:?} -> {err:?}")); + } + validate_vertex_authorization_headers(&[]) + .expect("no authorization header defers to minting"); + } + + #[test] + fn vertex_configs_use_google_oauth() { + assert_eq!(VERTEX_AI_OCR_CONFIG.ocr_auth(), OcrAuth::VertexOauth); + assert_eq!( + VERTEX_AI_DEEPSEEK_OCR_CONFIG.ocr_auth(), + OcrAuth::VertexOauth + ); + } + #[test] fn vertex_mistral_reuses_mistral_body_transform() { let body = VERTEX_AI_OCR_CONFIG diff --git a/litellm-rust/crates/python-bridge/Cargo.toml b/litellm-rust/crates/python-bridge/Cargo.toml index 83e163c38f1..b81f63f112a 100644 --- a/litellm-rust/crates/python-bridge/Cargo.toml +++ b/litellm-rust/crates/python-bridge/Cargo.toml @@ -4,6 +4,7 @@ version = "0.1.0" edition.workspace = true license.workspace = true repository.workspace = true +rust-version.workspace = true [lib] name = "_native"