From dc585235bef4bfd4f0f57e8c688dec2cf2a9350e Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Thu, 16 Jul 2026 21:37:28 -0700 Subject: [PATCH] feat(vertex): add shared VertexAiBase Rust host authentication (#33602) * feat(vertex): add shared Rust Google authentication Mint and refresh Vertex OAuth access tokens in the ai-gateway host layer via the official google-cloud-auth crate, preserving the inline service-account JSON, ADC and GOOGLE_APPLICATION_CREDENTIALS contract with library-managed caching/refresh and no hand-rolled signing. Core stays auth/IO-free: it only classifies the bearer source and rejects Google AIza API keys rather than sending them as OAuth. Also fix Vertex Mistral rawPredict endpoint construction so the global location targets aiplatform.googleapis.com. Co-Authored-By: Ishaan Jaffer <155045088+ishaan-berri@users.noreply.github.com> * refactor(vertex): address review; SHA-256 cache key and env credentials Move newly introduced Vertex constants into crate-level constants.rs, drop the doc/inline comments added by the auth PR, and replace the DefaultHasher u64 credential-cache key with a collision-resistant SHA-256 digest (sha2 is now a non-optional ai-gateway dependency so the non-server host path can use it). When no explicit vertex_credentials optional param is supplied, the host now reads VERTEXAI_CREDENTIALS from the environment before falling back to standard ADC/GOOGLE_APPLICATION_CREDENTIALS, without exposing credential content. Adds tests covering env-based inline credential selection and that distinct credential sources cannot collide on one cache entry. Co-Authored-By: Ishaan Jaffer <155045088+ishaan-berri@users.noreply.github.com> * build(vertex): pin google-cloud-auth to =1.13.0 for dependency-age policy Co-Authored-By: Ishaan Jaffer <155045088+ishaan-berri@users.noreply.github.com> * fix(vertex): bound credential cache, reject AIza in auth header, data-minimize auth errors Co-Authored-By: Ishaan Jaffer <155045088+ishaan-berri@users.noreply.github.com> * build(vertex): pin google-cloud-auth to =1.9.0 and adopt MSRV-aware resolver for Rust 1.86 Co-Authored-By: Ishaan Jaffer <155045088+ishaan-berri@users.noreply.github.com> * fix(vertex): require well-formed Bearer scheme on caller Authorization header Co-Authored-By: Ishaan Jaffer <155045088+ishaan-berri@users.noreply.github.com> * fix(vertex): require exactly one well-formed Bearer authorization header Co-Authored-By: Ishaan Jaffer <155045088+ishaan-berri@users.noreply.github.com> * fix(vertex): key credential cache by content and pin rustls exactly Co-Authored-By: Ishaan Jaffer <155045088+ishaan-berri@users.noreply.github.com> * chore(bridge): remove accidentally committed native extension binary Co-Authored-By: Ishaan Jaffer <155045088+ishaan-berri@users.noreply.github.com> * refactor(vertex): split base auth and shared cache Co-Authored-By: Ishaan Jaffer <155045088+ishaan-berri@users.noreply.github.com> --------- Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Co-authored-by: Ishaan Jaffer <155045088+ishaan-berri@users.noreply.github.com> --- litellm-rust/Cargo.lock | 779 +++++++++++++++++- litellm-rust/Cargo.toml | 5 +- litellm-rust/crates/ai-gateway/Cargo.toml | 9 +- .../crates/ai-gateway/src/constants.rs | 6 + litellm-rust/crates/ai-gateway/src/io/mod.rs | 1 + litellm-rust/crates/ai-gateway/src/io/ocr.rs | 39 +- .../ai-gateway/src/io/ocr/common_utils.rs | 8 + .../crates/ai-gateway/src/io/ocr/tests.rs | 213 +++++ .../crates/ai-gateway/src/io/vertex_ai/mod.rs | 3 + .../src/io/vertex_ai/vertex_ai_base.rs | 556 +++++++++++++ litellm-rust/crates/core/Cargo.toml | 1 + .../crates/core/src/cache/in_memory.rs | 89 ++ litellm-rust/crates/core/src/cache/mod.rs | 1 + litellm-rust/crates/core/src/constants.rs | 4 + litellm-rust/crates/core/src/lib.rs | 2 + .../crates/core/src/ocr/transformation.rs | 20 +- .../providers/vertex_ai/ocr/transformation.rs | 274 +++++- litellm-rust/crates/python-bridge/Cargo.toml | 1 + 18 files changed, 1960 insertions(+), 51 deletions(-) create mode 100644 litellm-rust/crates/ai-gateway/src/io/vertex_ai/mod.rs create mode 100644 litellm-rust/crates/ai-gateway/src/io/vertex_ai/vertex_ai_base.rs create mode 100644 litellm-rust/crates/core/src/cache/in_memory.rs create mode 100644 litellm-rust/crates/core/src/cache/mod.rs create mode 100644 litellm-rust/crates/core/src/constants.rs 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"