From f2e340cf2bfce78928fd6377d8afd2f9cb2e6b84 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Sat, 18 Jul 2026 19:12:00 -0700 Subject: [PATCH] feat(rust): port BaseAWSLLM auth (credential resolution + SigV4) to litellm-core as a base provider (#33888) * feat(rust): add feature-gated Bedrock AWS auth Co-Authored-By: Ishaan Jaffer <155045088+ishaan-berri@users.noreply.github.com> * refactor(rust): move Bedrock auth into core Co-Authored-By: Ishaan Jaffer <155045088+ishaan-berri@users.noreply.github.com> * fix(rust): fall through caller identity lookup errors Co-Authored-By: Ishaan Jaffer <155045088+ishaan-berri@users.noreply.github.com> * test(rust): add live Bedrock proof and CI coverage Co-Authored-By: Ishaan Jaffer <155045088+ishaan-berri@users.noreply.github.com> * refactor(rust): share in-memory cache with Bedrock auth Co-Authored-By: Ishaan Jaffer <155045088+ishaan-berri@users.noreply.github.com> * fix(rust): preserve web identity credential expiry 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> --- .github/workflows/test-rust.yml | 6 + litellm-rust/Cargo.lock | 922 +++++++++++++++++- litellm-rust/crates/core/Cargo.toml | 19 + .../core/src/caching/in_memory_cache.rs | 258 +++++ litellm-rust/crates/core/src/caching/mod.rs | 1 + litellm-rust/crates/core/src/lib.rs | 1 + .../core/src/providers/bedrock/aws_base.rs | 724 ++++++++++++++ .../core/src/providers/bedrock/constants.rs | 14 + .../crates/core/src/providers/bedrock/mod.rs | 6 + litellm-rust/crates/core/src/providers/mod.rs | 2 + .../llms/bedrock/test_base_aws_llm.py | 25 + 11 files changed, 1930 insertions(+), 48 deletions(-) create mode 100644 litellm-rust/crates/core/src/caching/in_memory_cache.rs create mode 100644 litellm-rust/crates/core/src/caching/mod.rs create mode 100644 litellm-rust/crates/core/src/providers/bedrock/aws_base.rs create mode 100644 litellm-rust/crates/core/src/providers/bedrock/constants.rs create mode 100644 litellm-rust/crates/core/src/providers/bedrock/mod.rs diff --git a/.github/workflows/test-rust.yml b/.github/workflows/test-rust.yml index 13e1dc4ad5e..21e1bcb90c6 100644 --- a/.github/workflows/test-rust.yml +++ b/.github/workflows/test-rust.yml @@ -61,5 +61,11 @@ jobs: - name: Run Clippy run: cargo clippy --workspace --all-targets --locked -- -D warnings + - name: Run Clippy with Bedrock auth + run: cargo clippy -p litellm-core --all-targets --features bedrock-auth --locked -- -D warnings + - name: Run Rust tests run: cargo test --workspace --locked + + - name: Run core tests with Bedrock auth + run: cargo test -p litellm-core --features bedrock-auth --locked diff --git a/litellm-rust/Cargo.lock b/litellm-rust/Cargo.lock index 7daff1dfc91..402d16715a3 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 = "arc-swap" +version = "1.9.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c049c0be4daef0b145cb3555416b3b8ef5b7888a38aea1a3a155801fe7b0810b" +dependencies = [ + "rustversion", +] + [[package]] name = "async-trait" version = "0.1.89" @@ -19,6 +28,358 @@ version = "1.1.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "1505bd5d3d116872e7271a6d4e16d81d0c8570876c8de68093a09ac269d8aac0" +[[package]] +name = "autocfg" +version = "1.5.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f2032f911046de80f0a198e0901378627c33f59ea0ac00e363d481118bd70a53" + +[[package]] +name = "aws-config" +version = "1.9.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "47712fde1909402600ccfbb26e47d482d2e58bb9e9e603d9f17e67cc435a6319" +dependencies = [ + "aws-credential-types", + "aws-runtime", + "aws-sdk-sts", + "aws-smithy-async", + "aws-smithy-http", + "aws-smithy-json", + "aws-smithy-runtime", + "aws-smithy-runtime-api", + "aws-smithy-schema", + "aws-smithy-types", + "aws-types", + "bytes", + "fastrand", + "http 1.4.2", + "time", + "tokio", + "tracing", + "url", +] + +[[package]] +name = "aws-credential-types" +version = "1.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e93964ffdaf57857f544be3666a5f57570bb699e934700f11b49708f61bb556e" +dependencies = [ + "aws-smithy-async", + "aws-smithy-runtime-api", + "aws-smithy-types", + "zeroize", +] + +[[package]] +name = "aws-lc-rs" +version = "1.17.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "00bdb5da18dac48ca2cc7cd4a98e533e8635a58e2361d13a1a4ee3888e0d72f1" +dependencies = [ + "aws-lc-sys", + "zeroize", +] + +[[package]] +name = "aws-lc-sys" +version = "0.43.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "43103168cc76fe62678a375e722fc9cb3a0146159ac5828bc4f0dfd755c2224c" +dependencies = [ + "cc", + "cmake", + "dunce", + "fs_extra", + "pkg-config", +] + +[[package]] +name = "aws-runtime" +version = "1.8.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7816e98ee912159f45d307e5ee6bfea4a335a55aee15f7f3e32f81a6f3000f1d" +dependencies = [ + "aws-credential-types", + "aws-sigv4", + "aws-smithy-async", + "aws-smithy-http", + "aws-smithy-runtime", + "aws-smithy-runtime-api", + "aws-smithy-types", + "aws-types", + "bytes", + "bytes-utils", + "fastrand", + "http 1.4.2", + "http-body 1.0.1", + "percent-encoding", + "pin-project-lite", + "tracing", + "uuid", +] + +[[package]] +name = "aws-sdk-sts" +version = "1.108.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3c72b08911d8128dd360fe1b22a9fec0fa8b552dde8ec828dcf20ef5ec974e9f" +dependencies = [ + "arc-swap", + "aws-credential-types", + "aws-runtime", + "aws-smithy-async", + "aws-smithy-http", + "aws-smithy-json", + "aws-smithy-observability", + "aws-smithy-query", + "aws-smithy-runtime", + "aws-smithy-runtime-api", + "aws-smithy-schema", + "aws-smithy-types", + "aws-smithy-xml", + "aws-types", + "fastrand", + "http 0.2.12", + "http 1.4.2", + "regex-lite", + "tracing", +] + +[[package]] +name = "aws-sigv4" +version = "1.5.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "723c2234ad7511ceef63eab016b7ba6ff7c55590fefb96fa8467af014a07309f" +dependencies = [ + "aws-credential-types", + "aws-smithy-http", + "aws-smithy-runtime-api", + "aws-smithy-types", + "bytes", + "form_urlencoded", + "hex", + "hmac", + "http 0.2.12", + "http 1.4.2", + "percent-encoding", + "sha2 0.11.0", + "time", + "tracing", +] + +[[package]] +name = "aws-smithy-async" +version = "1.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f02e407fb3b54891734224b9ffac8a71fdd35f542500fa1af95754a6b2beb316" +dependencies = [ + "futures-util", + "pin-project-lite", + "tokio", +] + +[[package]] +name = "aws-smithy-http" +version = "0.64.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "37843d9add67c3aff5856f409c6dc315d3cdff60f9c0cb5b670dab1e9920306d" +dependencies = [ + "aws-smithy-runtime-api", + "aws-smithy-types", + "bytes", + "bytes-utils", + "futures-core", + "futures-util", + "http 1.4.2", + "http-body 1.0.1", + "http-body-util", + "percent-encoding", + "pin-project-lite", + "pin-utils", + "tracing", +] + +[[package]] +name = "aws-smithy-http-client" +version = "1.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "635d23afda0a6ab48d666c4d447c4873e8d1e83518a2be2093122397e50b838e" +dependencies = [ + "aws-smithy-async", + "aws-smithy-runtime-api", + "aws-smithy-types", + "h2 0.3.27", + "h2 0.4.15", + "http 0.2.12", + "http 1.4.2", + "http-body 0.4.6", + "hyper 0.14.32", + "hyper 1.10.1", + "hyper-rustls 0.24.2", + "hyper-rustls 0.27.9", + "hyper-util", + "pin-project-lite", + "rustls 0.21.12", + "rustls 0.23.41", + "rustls-native-certs", + "rustls-pki-types", + "tokio", + "tokio-rustls 0.26.4", + "tower", + "tracing", +] + +[[package]] +name = "aws-smithy-json" +version = "0.63.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3dc65a121adb4b33729919fcfa14fa36fb33c1555a8f06bb0e2188dbfdc1d9ef" +dependencies = [ + "aws-smithy-runtime-api", + "aws-smithy-schema", + "aws-smithy-types", +] + +[[package]] +name = "aws-smithy-observability" +version = "0.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8e86338c869539a581bf161247762a6e87f92c5c075060057b5ed6d06632ed0c" +dependencies = [ + "aws-smithy-runtime-api", +] + +[[package]] +name = "aws-smithy-query" +version = "0.61.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "dd22a6ba36e3f113cb8d5b3d1fe0ed31c76ee608ef63322d753bb8d2c9479e77" +dependencies = [ + "aws-smithy-types", + "urlencoding", +] + +[[package]] +name = "aws-smithy-runtime" +version = "1.12.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bea94a9ff8464016338c851e24b472d7131c388c88898a502e781815b2ee6045" +dependencies = [ + "aws-smithy-async", + "aws-smithy-http", + "aws-smithy-http-client", + "aws-smithy-observability", + "aws-smithy-runtime-api", + "aws-smithy-schema", + "aws-smithy-types", + "bytes", + "fastrand", + "http 0.2.12", + "http 1.4.2", + "http-body 0.4.6", + "http-body 1.0.1", + "http-body-util", + "pin-project-lite", + "pin-utils", + "tokio", + "tracing", +] + +[[package]] +name = "aws-smithy-runtime-api" +version = "1.13.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "22ed1ebe6e0a95ea84570225f5a8208dec4b8f77e61a9b0d6f51773fcb4612f0" +dependencies = [ + "aws-smithy-async", + "aws-smithy-runtime-api-macros", + "aws-smithy-types", + "bytes", + "http 0.2.12", + "http 1.4.2", + "pin-project-lite", + "tokio", + "tracing", + "zeroize", +] + +[[package]] +name = "aws-smithy-runtime-api-macros" +version = "1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "221eaa237ddf1ca79b60d1372aad77e47f9c0ea5b3ce5099da8c61d027dc77b3" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "aws-smithy-schema" +version = "0.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7d56e0a4e53127a632224e43633b0fe045fa9e1e3cfc68b9830f1115e103f910" +dependencies = [ + "aws-smithy-runtime-api", + "aws-smithy-types", + "http 1.4.2", +] + +[[package]] +name = "aws-smithy-types" +version = "1.6.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d6dc683efb34b9e755675b37fedbe0103141e5b6df7bdc9eb6967756a8c167d8" +dependencies = [ + "base64-simd", + "bytes", + "bytes-utils", + "futures-core", + "http 0.2.12", + "http 1.4.2", + "http-body 0.4.6", + "http-body 1.0.1", + "http-body-util", + "itoa", + "num-integer", + "pin-project-lite", + "pin-utils", + "ryu", + "serde", + "time", + "tokio", + "tokio-util", +] + +[[package]] +name = "aws-smithy-xml" +version = "0.61.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ea3f68eec3607f02acd24067969ce2abc6ba16aa7d5ce59ca450ed2fb5f78957" +dependencies = [ + "aws-smithy-runtime-api", + "aws-smithy-schema", + "aws-smithy-types", + "xmlparser", +] + +[[package]] +name = "aws-types" +version = "1.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e957a6c6dbce82b7a91f44231c09273159703769f447cbe85e854dfe9cf67f86" +dependencies = [ + "aws-credential-types", + "aws-smithy-async", + "aws-smithy-runtime-api", + "aws-smithy-schema", + "aws-smithy-types", + "rustc_version", + "tracing", +] + [[package]] name = "axum" version = "0.7.9" @@ -30,10 +391,10 @@ dependencies = [ "base64", "bytes", "futures-util", - "http", - "http-body", + "http 1.4.2", + "http-body 1.0.1", "http-body-util", - "hyper", + "hyper 1.10.1", "hyper-util", "itoa", "matchit", @@ -65,8 +426,8 @@ dependencies = [ "async-trait", "bytes", "futures-util", - "http", - "http-body", + "http 1.4.2", + "http-body 1.0.1", "http-body-util", "mime", "pin-project-lite", @@ -83,6 +444,16 @@ version = "0.22.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "72b3254f16251a8381aa12e40e3c4d2f0199f8c6508fbecb9d91f575e0fbb8c6" +[[package]] +name = "base64-simd" +version = "0.8.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "339abbe78e73178762e23bea9dfd08e697eb3f3301cd4be981c0f78ba5859195" +dependencies = [ + "outref", + "vsimd", +] + [[package]] name = "bitflags" version = "2.13.0" @@ -98,6 +469,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" @@ -116,6 +496,16 @@ version = "1.12.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "8ae3f5d315924270530207e2a68396c3cc547f6dca3fbdca317cfb1a51edb593" +[[package]] +name = "bytes-utils" +version = "0.1.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7dafe3a8757b027e2be6e4e5601ed563c55989fcf1546e933c66c8eb3a058d35" +dependencies = [ + "bytes", + "either", +] + [[package]] name = "cc" version = "1.2.65" @@ -123,6 +513,8 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "e228eec9be7c17ccb640b59b36a5cd805ea2a564a4c5e162c2f659fea30d3b96" dependencies = [ "find-msvc-tools", + "jobserver", + "libc", "shlex", ] @@ -138,6 +530,27 @@ version = "0.2.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "613afe47fcd5fac7ccf1db93babcb082c5994d996f20b8b159f2ad1658eb5724" +[[package]] +name = "cmake" +version = "0.1.58" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c0f78a02292a74a88ac736019ab962ece0bc380e3f977bf72e376c5d78ff0678" +dependencies = [ + "cc", +] + +[[package]] +name = "cmov" +version = "0.5.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0c9ea0ac24bc397ab3c98583a3c9ba74fa56b09a4449bbe172b9b1ddb016027a" + +[[package]] +name = "const-oid" +version = "0.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a6ef517f0926dd24a1582492c791b6a4818a4d94e789a334894aa15b0d12f55c" + [[package]] name = "core-foundation" version = "0.10.1" @@ -163,6 +576,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" @@ -173,20 +595,56 @@ 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 = "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" + [[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", + "const-oid", + "crypto-common 0.2.2", + "ctutils", ] [[package]] @@ -200,12 +658,30 @@ dependencies = [ "syn", ] +[[package]] +name = "dunce" +version = "1.0.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "92773504d58c093f6de2459af4af33faa518c13451eb8f2b5698ed3d36e7c813" + +[[package]] +name = "either" +version = "1.16.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "91622ff5e7162018101f2fea40d6ebf4a78bbe5a49736a2020649edf9693679e" + [[package]] name = "equivalent" version = "1.0.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "877a4ace8713b0bcf2a4e7eec82529c029f1d0619886d18145fea96c3ffe5c0f" +[[package]] +name = "fastrand" +version = "2.4.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9f1f227452a390804cdb637b74a86990f2a7d7ba4b7d5693aac9b4dd6defd8d6" + [[package]] name = "find-msvc-tools" version = "0.1.9" @@ -227,6 +703,12 @@ dependencies = [ "percent-encoding", ] +[[package]] +name = "fs_extra" +version = "1.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "42703706b716c37f96a77aea830392ad231f44c9e9a67872fa5548707e11b11c" + [[package]] name = "futures-channel" version = "0.3.32" @@ -320,11 +802,41 @@ 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", +] + +[[package]] +name = "h2" +version = "0.3.27" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0beca50380b1fc32983fc1cb4587bfa4bb9e78fc259aad4a0032d2080309222d" +dependencies = [ + "bytes", + "fnv", + "futures-core", + "futures-sink", + "futures-util", + "http 0.2.12", + "indexmap", + "slab", + "tokio", + "tokio-util", + "tracing", +] + [[package]] name = "h2" version = "0.4.15" @@ -336,7 +848,7 @@ dependencies = [ "fnv", "futures-core", "futures-sink", - "http", + "http 1.4.2", "indexmap", "slab", "tokio", @@ -356,6 +868,32 @@ 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 = "0.2.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "601cbb57e577e2f5ef5be8e7b83f0f63994f25aa94d673e54a92d5c516d101f1" +dependencies = [ + "bytes", + "fnv", + "itoa", +] + [[package]] name = "http" version = "1.4.2" @@ -366,6 +904,17 @@ dependencies = [ "itoa", ] +[[package]] +name = "http-body" +version = "0.4.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7ceab25649e9960c0311ea418d17bee82c0dcec1bd053b5f9a66e265a693bed2" +dependencies = [ + "bytes", + "http 0.2.12", + "pin-project-lite", +] + [[package]] name = "http-body" version = "1.0.1" @@ -373,7 +922,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "1efedce1fb8e6913f23e0c92de8e62cd5b772a67e7b3946df930a62566c93184" dependencies = [ "bytes", - "http", + "http 1.4.2", ] [[package]] @@ -384,8 +933,8 @@ checksum = "b021d93e26becf5dc7e1b75b1bed1fd93124b374ceb73f43d4d4eafec896a64a" dependencies = [ "bytes", "futures-core", - "http", - "http-body", + "http 1.4.2", + "http-body 1.0.1", "pin-project-lite", ] @@ -401,6 +950,39 @@ 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 = "0.14.32" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "41dfc780fdec9373c01bae43289ea34c972e40ee3c9f6b3c8801a35f35586ce7" +dependencies = [ + "bytes", + "futures-channel", + "futures-core", + "futures-util", + "h2 0.3.27", + "http 0.2.12", + "http-body 0.4.6", + "httparse", + "httpdate", + "itoa", + "pin-project-lite", + "socket2 0.5.10", + "tokio", + "tower-service", + "tracing", + "want", +] + [[package]] name = "hyper" version = "1.10.1" @@ -411,9 +993,9 @@ dependencies = [ "bytes", "futures-channel", "futures-core", - "h2", - "http", - "http-body", + "h2 0.4.15", + "http 1.4.2", + "http-body 1.0.1", "httparse", "httpdate", "itoa", @@ -423,18 +1005,34 @@ dependencies = [ "want", ] +[[package]] +name = "hyper-rustls" +version = "0.24.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ec3efd23720e2049821a693cbc7e65ea87c72f1c58ff2f9522ff332b1491e590" +dependencies = [ + "futures-util", + "http 0.2.12", + "hyper 0.14.32", + "log", + "rustls 0.21.12", + "tokio", + "tokio-rustls 0.24.1", +] + [[package]] name = "hyper-rustls" version = "0.27.9" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "33ca68d021ef39cf6463ab54c1d0f5daf03377b70561305bb89a8f83aab66e0f" dependencies = [ - "http", - "hyper", + "http 1.4.2", + "hyper 1.10.1", "hyper-util", - "rustls", + "rustls 0.23.41", + "rustls-native-certs", "tokio", - "tokio-rustls", + "tokio-rustls 0.26.4", "tower-service", "webpki-roots", ] @@ -449,14 +1047,14 @@ dependencies = [ "bytes", "futures-channel", "futures-util", - "http", - "http-body", - "hyper", + "http 1.4.2", + "http-body 1.0.1", + "hyper 1.10.1", "ipnet", "libc", "percent-encoding", "pin-project-lite", - "socket2", + "socket2 0.6.4", "tokio", "tower-service", "tracing", @@ -587,6 +1185,16 @@ version = "1.0.18" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "8f42a60cbdf9a97f5d2305f08a87dc4e09308d1276d28c869c684d7777685682" +[[package]] +name = "jobserver" +version = "0.1.35" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1c00acbd29eabad4a2392fa0e921c874934dbbf4194312ad20f04a0ed67a3cb3" +dependencies = [ + "getrandom 0.4.3", + "libc", +] + [[package]] name = "js-sys" version = "0.3.103" @@ -617,7 +1225,7 @@ dependencies = [ "reqwest", "serde", "serde_json", - "sha2", + "sha2 0.10.9", "subtle", "tokio", "tokio-tungstenite", @@ -628,9 +1236,17 @@ dependencies = [ name = "litellm-core" version = "0.1.0" dependencies = [ + "aws-config", + "aws-credential-types", + "aws-sdk-sts", + "aws-sigv4", + "aws-smithy-runtime-api", + "aws-types", "rand 0.8.6", + "reqwest", "serde", "serde_json", + "sha2 0.10.9", "thiserror 2.0.18", "tokio", ] @@ -694,6 +1310,30 @@ dependencies = [ "windows-sys 0.61.2", ] +[[package]] +name = "num-conv" +version = "0.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "521739c6d2bac4aa25192232afe6841231376b2b26d4d9fae5ecf8ca5772e441" + +[[package]] +name = "num-integer" +version = "0.1.46" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7969661fd2958a5cb096e56c8e1ad0444ac2bbcd0061bd28660485a44879858f" +dependencies = [ + "num-traits", +] + +[[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" @@ -706,6 +1346,12 @@ version = "0.2.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "7c87def4c32ab89d880effc9e097653c8da5d6ef28e6b539d313baaacfbafcbe" +[[package]] +name = "outref" +version = "0.5.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1a80800c0488c3a21695ea981a54918fbb37abf04f4d0720c453632255e2ff0e" + [[package]] name = "percent-encoding" version = "2.3.2" @@ -718,6 +1364,18 @@ version = "0.2.17" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "a89322df9ebe1c1578d689c92318e070967d1042b512afbe49518723f4e6d5cd" +[[package]] +name = "pin-utils" +version = "0.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8b870d8c151b6f2fb93e84a13146138f05d02ed11c7e7c54f8826aaaf7c9f184" + +[[package]] +name = "pkg-config" +version = "0.3.33" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "19f132c84eca552bf34cab8ec81f1c1dcc229b811638f9d283dceabe58c5569e" + [[package]] name = "portable-atomic" version = "1.13.1" @@ -733,6 +1391,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" @@ -834,8 +1498,8 @@ dependencies = [ "quinn-proto", "quinn-udp", "rustc-hash", - "rustls", - "socket2", + "rustls 0.23.41", + "socket2 0.6.4", "thiserror 2.0.18", "tokio", "tracing", @@ -854,7 +1518,7 @@ dependencies = [ "rand 0.9.4", "ring", "rustc-hash", - "rustls", + "rustls 0.23.41", "rustls-pki-types", "slab", "thiserror 2.0.18", @@ -872,7 +1536,7 @@ dependencies = [ "cfg_aliases", "libc", "once_cell", - "socket2", + "socket2 0.6.4", "tracing", "windows-sys 0.60.2", ] @@ -892,6 +1556,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" @@ -951,6 +1621,12 @@ dependencies = [ "getrandom 0.3.4", ] +[[package]] +name = "regex-lite" +version = "0.1.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cab834c73d247e67f4fae452806d17d3c7501756d98c8808d7c9c7aa7d18f973" + [[package]] name = "reqwest" version = "0.12.28" @@ -962,26 +1638,26 @@ dependencies = [ "futures-channel", "futures-core", "futures-util", - "h2", - "http", - "http-body", + "h2 0.4.15", + "http 1.4.2", + "http-body 1.0.1", "http-body-util", - "hyper", - "hyper-rustls", + "hyper 1.10.1", + "hyper-rustls 0.27.9", "hyper-util", "js-sys", "log", "percent-encoding", "pin-project-lite", "quinn", - "rustls", + "rustls 0.23.41", "rustls-pki-types", "serde", "serde_json", "serde_urlencoded", "sync_wrapper", "tokio", - "tokio-rustls", + "tokio-rustls 0.26.4", "tokio-util", "tower", "tower-http", @@ -1014,16 +1690,38 @@ 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.21.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3f56a14d1f48b391359b22f731fd4bd7e43c97f3c50eee276f3aa09c94784d3e" +dependencies = [ + "log", + "ring", + "rustls-webpki 0.101.7", + "sct", +] + [[package]] name = "rustls" version = "0.23.41" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6b92b125634d9b795e7beca796cc790df15a7fb38323bf3196fda83292d06b1f" dependencies = [ + "aws-lc-rs", "once_cell", "ring", "rustls-pki-types", - "rustls-webpki", + "rustls-webpki 0.103.13", "subtle", "zeroize", ] @@ -1050,12 +1748,23 @@ dependencies = [ "zeroize", ] +[[package]] +name = "rustls-webpki" +version = "0.101.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8b6275d1ee7a1cd780b64aca7726599a1dbc893b1e64144529e55c3c2f745765" +dependencies = [ + "ring", + "untrusted", +] + [[package]] name = "rustls-webpki" version = "0.103.13" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "61c429a8649f110dddef65e2a5ad240f747e85f7758a6bccc7e5777bd33f756e" dependencies = [ + "aws-lc-rs", "ring", "rustls-pki-types", "untrusted", @@ -1082,6 +1791,16 @@ dependencies = [ "windows-sys 0.61.2", ] +[[package]] +name = "sct" +version = "0.7.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "da046153aa2352493d6cb7da4b6e5c0c057d8a1d0a9aa8560baffdd945acd414" +dependencies = [ + "ring", + "untrusted", +] + [[package]] name = "security-framework" version = "3.7.0" @@ -1105,6 +1824,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" @@ -1178,8 +1903,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]] @@ -1189,8 +1914,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]] @@ -1211,6 +1947,16 @@ version = "1.15.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "8ed6a63f02c8539c91a8685a86f4099661ba3da017932f6ebbea6de3f0fa7c90" +[[package]] +name = "socket2" +version = "0.5.10" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e22376abed350d73dd1cd119b57ffccad95b4e585a7cda43e286245ce23c0678" +dependencies = [ + "libc", + "windows-sys 0.52.0", +] + [[package]] name = "socket2" version = "0.6.4" @@ -1310,6 +2056,36 @@ dependencies = [ "syn", ] +[[package]] +name = "time" +version = "0.3.53" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "18dfaaeddcb932337b5e7866ee7d0ce9b76d2fd092997146f187ec09b4558a50" +dependencies = [ + "deranged", + "num-conv", + "powerfmt", + "serde_core", + "time-core", + "time-macros", +] + +[[package]] +name = "time-core" +version = "0.1.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9e1c906769ad99c88eaa54e728060edef082f8e358ff32030cb7c7d315e81109" + +[[package]] +name = "time-macros" +version = "0.2.31" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c431b87111666e491a90baa837f914fb45cd5dc3c268591b0220ff5057f2085f" +dependencies = [ + "num-conv", + "time-core", +] + [[package]] name = "tinystr" version = "0.8.3" @@ -1345,7 +2121,7 @@ dependencies = [ "libc", "mio", "pin-project-lite", - "socket2", + "socket2 0.6.4", "tokio-macros", "windows-sys 0.61.2", ] @@ -1361,13 +2137,23 @@ dependencies = [ "syn", ] +[[package]] +name = "tokio-rustls" +version = "0.24.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c28327cf380ac148141087fbfb9de9d7bd4e84ab5d2c28fbc911d753de8a7081" +dependencies = [ + "rustls 0.21.12", + "tokio", +] + [[package]] name = "tokio-rustls" version = "0.26.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "1729aa945f29d91ba541258c8df89027d5792d85a8841fb65e8bf0f4ede4ef61" dependencies = [ - "rustls", + "rustls 0.23.41", "tokio", ] @@ -1379,11 +2165,11 @@ checksum = "edc5f74e248dc973e0dbb7b74c7e0d6fcc301c694ff50049504004ef4d0cdcd9" dependencies = [ "futures-util", "log", - "rustls", + "rustls 0.23.41", "rustls-native-certs", "rustls-pki-types", "tokio", - "tokio-rustls", + "tokio-rustls 0.26.4", "tungstenite", ] @@ -1425,8 +2211,8 @@ dependencies = [ "bitflags", "bytes", "futures-util", - "http", - "http-body", + "http 1.4.2", + "http-body 1.0.1", "pin-project-lite", "tower", "tower-layer", @@ -1454,9 +2240,21 @@ checksum = "63e71662fa4b2a2c3a26f570f037eb95bb1f85397f3cd8076caed2f026a6d100" dependencies = [ "log", "pin-project-lite", + "tracing-attributes", "tracing-core", ] +[[package]] +name = "tracing-attributes" +version = "0.1.31" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7490cfa5ec963746568740651ac6781f701c9c5ea257c58e057f3ba8cf69e8da" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "tracing-core" version = "0.1.36" @@ -1481,11 +2279,11 @@ dependencies = [ "byteorder", "bytes", "data-encoding", - "http", + "http 1.4.2", "httparse", "log", "rand 0.8.6", - "rustls", + "rustls 0.23.41", "rustls-pki-types", "sha1", "thiserror 1.0.69", @@ -1522,6 +2320,12 @@ dependencies = [ "serde", ] +[[package]] +name = "urlencoding" +version = "2.1.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "daf8dba3b7eb870caf1ddeed7bc9d2a049f3cfdfae7cb521b087cc33ae4c49da" + [[package]] name = "utf-8" version = "0.7.6" @@ -1534,12 +2338,28 @@ version = "1.0.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "b6c140620e7ffbb22c2dee59cafe6084a59b5ffc27a8859a5f0d494b5d52b6be" +[[package]] +name = "uuid" +version = "1.24.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bf3923a6f5c4c6382e0b653c4117f48d631ea17f38ed86e2a828e6f7412f5239" +dependencies = [ + "js-sys", + "wasm-bindgen", +] + [[package]] name = "version_check" version = "0.9.5" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "0b928f33d975fc6ad9f86c8f283853ad26bdd5b10b7f1542aa2fa15e2289105a" +[[package]] +name = "vsimd" +version = "0.8.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5c3082ca00d5a5ef149bb8b555a72ae84c9c59f7250f013ac822ac2e49b19c64" + [[package]] name = "want" version = "0.3.1" @@ -1835,6 +2655,12 @@ version = "0.6.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "1ffae5123b2d3fc086436f8834ae3ab053a283cfac8fe0a0b8eaae044768a4c4" +[[package]] +name = "xmlparser" +version = "0.13.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "66fee0b777b0f5ac1c69bb06d361268faafa61cd4682ae064a171c16c433e9e4" + [[package]] name = "yoke" version = "0.8.3" diff --git a/litellm-rust/crates/core/Cargo.toml b/litellm-rust/crates/core/Cargo.toml index 9bd4634cc2a..65c6db7412c 100644 --- a/litellm-rust/crates/core/Cargo.toml +++ b/litellm-rust/crates/core/Cargo.toml @@ -10,6 +10,25 @@ rand.workspace = true serde.workspace = true serde_json.workspace = true thiserror.workspace = true +sha2.workspace = true +aws-config = { version = "1.9.0", default-features = false, features = ["rustls", "rt-tokio"], optional = true } +aws-credential-types = { version = "1.3.0", features = ["hardcoded-credentials"], optional = true } +aws-sdk-sts = { version = "1.108.0", default-features = false, features = ["rustls", "rt-tokio"], optional = true } +aws-sigv4 = { version = "1.5.1", optional = true } +aws-types = { version = "1.4.0", optional = true } +aws-smithy-runtime-api = { version = "1.13.0", optional = true } + +[features] +default = [] +bedrock-auth = [ + "dep:aws-config", + "dep:aws-credential-types", + "dep:aws-sdk-sts", + "dep:aws-sigv4", + "dep:aws-types", + "dep:aws-smithy-runtime-api", +] [dev-dependencies] +reqwest.workspace = true tokio = { workspace = true, features = ["macros", "rt-multi-thread"] } diff --git a/litellm-rust/crates/core/src/caching/in_memory_cache.rs b/litellm-rust/crates/core/src/caching/in_memory_cache.rs new file mode 100644 index 00000000000..0ceeedb8b71 --- /dev/null +++ b/litellm-rust/crates/core/src/caching/in_memory_cache.rs @@ -0,0 +1,258 @@ +use std::cmp::Reverse; +use std::collections::{BinaryHeap, HashMap}; +use std::time::{Duration, SystemTime, UNIX_EPOCH}; + +const DEFAULT_MAX_SIZE_IN_MEMORY: usize = 200; +const DEFAULT_TTL: Duration = Duration::from_secs(600); + +pub struct InMemoryCache { + pub cache_dict: HashMap, + pub ttl_dict: HashMap, + pub expiration_heap: BinaryHeap>, + pub max_size_in_memory: usize, + pub default_ttl: Duration, + now: Box Duration + Send + Sync>, +} + +impl Default for InMemoryCache { + fn default() -> Self { + Self::new(None, None) + } +} + +impl InMemoryCache { + pub fn new(max_size_in_memory: Option, default_ttl: Option) -> Self { + Self::with_clock(max_size_in_memory, default_ttl, || { + SystemTime::now() + .duration_since(UNIX_EPOCH) + .unwrap_or_default() + }) + } + + pub fn with_clock( + max_size_in_memory: Option, + default_ttl: Option, + now: impl Fn() -> Duration + Send + Sync + 'static, + ) -> Self { + Self { + cache_dict: HashMap::new(), + ttl_dict: HashMap::new(), + expiration_heap: BinaryHeap::new(), + max_size_in_memory: max_size_in_memory.unwrap_or(DEFAULT_MAX_SIZE_IN_MEMORY), + default_ttl: default_ttl.unwrap_or(DEFAULT_TTL), + now: Box::new(now), + } + } + + pub fn evict_cache(&mut self) { + if self.max_size_in_memory == 0 { + return; + } + + let current_time = (self.now)(); + while let Some(Reverse((expiration_time, key))) = self.expiration_heap.peek().cloned() { + if self.ttl_dict.get(&key).copied() != Some(expiration_time) { + self.expiration_heap.pop(); + } else if expiration_time <= current_time { + self.expiration_heap.pop(); + self.remove_key(&key); + } else { + break; + } + } + + while self.cache_dict.len() >= self.max_size_in_memory { + let Some(Reverse((expiration_time, key))) = self.expiration_heap.pop() else { + break; + }; + if self.ttl_dict.get(&key).copied() == Some(expiration_time) { + self.remove_key(&key); + } + } + } + + pub fn allow_ttl_override(&self, key: &str) -> bool { + match self.ttl_dict.get(key).copied() { + None => true, + Some(expiration_time) => expiration_time < (self.now)(), + } + } + + pub fn set_cache(&mut self, key: impl Into, value: V, ttl: Option) { + if self.max_size_in_memory == 0 { + return; + } + + self.evict_cache(); + let key = key.into(); + self.cache_dict.insert(key.clone(), value); + if self.allow_ttl_override(&key) { + let expiration_time = (self.now)() + ttl.unwrap_or(self.default_ttl); + self.ttl_dict.insert(key.clone(), expiration_time); + self.expiration_heap.push(Reverse((expiration_time, key))); + } + } + + // Generic values intentionally omit Python's per-item size check. + pub fn get_cache(&mut self, key: &str) -> Option { + if self.cache_dict.contains_key(key) { + if self.is_key_expired(key) { + self.remove_key(key); + return None; + } + return self.cache_dict.get(key).cloned(); + } + None + } + + pub fn get_ttl(&self, key: &str) -> Option { + self.ttl_dict.get(key).copied() + } + + pub fn delete_cache(&mut self, key: &str) { + self.remove_key(key); + } + + pub fn flush_cache(&mut self) { + self.cache_dict.clear(); + self.ttl_dict.clear(); + self.expiration_heap.clear(); + } + + fn is_key_expired(&self, key: &str) -> bool { + self.ttl_dict + .get(key) + .is_some_and(|expiration_time| *expiration_time < (self.now)()) + } + + fn remove_key(&mut self, key: &str) { + self.cache_dict.remove(key); + self.ttl_dict.remove(key); + } +} + +#[cfg(test)] +mod tests { + use std::sync::{ + atomic::{AtomicU64, Ordering}, + Arc, + }; + + use super::InMemoryCache; + use std::time::Duration; + + fn cache(now: Arc, max_size: usize, default_ttl: Duration) -> InMemoryCache { + InMemoryCache::with_clock(Some(max_size), Some(default_ttl), move || { + Duration::from_secs(now.load(Ordering::Relaxed)) + }) + } + + #[test] + fn ttl_expiry_is_deterministic() { + let now = Arc::new(AtomicU64::new(100)); + let mut cache = cache(now.clone(), 10, Duration::from_secs(60)); + cache.set_cache("key", "value".to_string(), None); + assert_eq!(cache.get_cache("key"), Some("value".to_string())); + now.store(161, Ordering::Relaxed); + assert_eq!(cache.get_cache("key"), None); + assert_eq!(cache.get_ttl("key"), None); + } + + #[test] + fn default_and_per_set_ttl_are_applied() { + let now = Arc::new(AtomicU64::new(100)); + let mut cache = cache(now.clone(), 10, Duration::from_secs(60)); + cache.set_cache("default", "value".to_string(), None); + cache.set_cache("custom", "value".to_string(), Some(Duration::from_secs(20))); + assert_eq!(cache.get_ttl("default"), Some(Duration::from_secs(160))); + assert_eq!(cache.get_ttl("custom"), Some(Duration::from_secs(120))); + } + + #[test] + fn unexpired_entries_do_not_allow_ttl_override() { + let now = Arc::new(AtomicU64::new(100)); + let mut cache = cache(now.clone(), 10, Duration::from_secs(60)); + cache.set_cache("key", "first".to_string(), Some(Duration::from_secs(20))); + cache.set_cache("key", "second".to_string(), Some(Duration::from_secs(80))); + assert_eq!(cache.get_cache("key"), Some("second".to_string())); + assert_eq!(cache.get_ttl("key"), Some(Duration::from_secs(120))); + now.store(121, Ordering::Relaxed); + cache.set_cache("key", "third".to_string(), Some(Duration::from_secs(80))); + assert_eq!(cache.get_ttl("key"), Some(Duration::from_secs(201))); + } + + #[test] + fn max_size_evicts_earliest_expiration() { + let now = Arc::new(AtomicU64::new(100)); + let mut cache = cache(now, 2, Duration::from_secs(60)); + cache.set_cache("early", "value".to_string(), Some(Duration::from_secs(10))); + cache.set_cache("late", "value".to_string(), Some(Duration::from_secs(20))); + cache.set_cache("new", "value".to_string(), Some(Duration::from_secs(30))); + assert_eq!(cache.get_cache("early"), None); + assert!(cache.get_cache("late").is_some()); + assert!(cache.get_cache("new").is_some()); + } + + #[test] + fn expired_entries_are_evicted_before_live_entries() { + let now = Arc::new(AtomicU64::new(100)); + let mut cache = cache(now.clone(), 3, Duration::from_secs(60)); + cache.set_cache( + "expired-one", + "value".to_string(), + Some(Duration::from_secs(10)), + ); + cache.set_cache( + "expired-two", + "value".to_string(), + Some(Duration::from_secs(20)), + ); + cache.set_cache("live", "value".to_string(), Some(Duration::from_secs(100))); + now.store(121, Ordering::Relaxed); + cache.set_cache("new", "value".to_string(), Some(Duration::from_secs(100))); + assert_eq!(cache.get_cache("expired-one"), None); + assert_eq!(cache.get_cache("expired-two"), None); + assert!(cache.get_cache("live").is_some()); + assert!(cache.get_cache("new").is_some()); + } + + #[test] + fn stale_heap_entries_are_skipped() { + let now = Arc::new(AtomicU64::new(100)); + let mut cache = cache(now, 1, Duration::from_secs(60)); + cache.set_cache( + "removed", + "value".to_string(), + Some(Duration::from_secs(10)), + ); + cache.delete_cache("removed"); + cache.set_cache("kept", "value".to_string(), Some(Duration::from_secs(20))); + cache.set_cache("new", "value".to_string(), Some(Duration::from_secs(30))); + assert_eq!(cache.get_cache("removed"), None); + assert_eq!(cache.get_cache("kept"), None); + assert!(cache.get_cache("new").is_some()); + } + + #[test] + fn delete_and_flush_remove_values_and_ttls() { + let now = Arc::new(AtomicU64::new(100)); + let mut cache = cache(now, 10, Duration::from_secs(60)); + cache.set_cache("one", "value".to_string(), None); + cache.set_cache("two", "value".to_string(), None); + cache.delete_cache("one"); + assert_eq!(cache.get_cache("one"), None); + cache.flush_cache(); + assert!(cache.cache_dict.is_empty()); + assert!(cache.ttl_dict.is_empty()); + assert!(cache.expiration_heap.is_empty()); + } + + #[test] + fn zero_max_size_does_not_cache() { + let now = Arc::new(AtomicU64::new(100)); + let mut cache = cache(now, 0, Duration::from_secs(60)); + cache.set_cache("key", "value".to_string(), None); + assert_eq!(cache.get_cache("key"), None); + assert!(cache.cache_dict.is_empty()); + } +} diff --git a/litellm-rust/crates/core/src/caching/mod.rs b/litellm-rust/crates/core/src/caching/mod.rs new file mode 100644 index 00000000000..5fb8a0e5174 --- /dev/null +++ b/litellm-rust/crates/core/src/caching/mod.rs @@ -0,0 +1 @@ +pub mod in_memory_cache; diff --git a/litellm-rust/crates/core/src/lib.rs b/litellm-rust/crates/core/src/lib.rs index 117596f53d5..3989fb441bc 100644 --- a/litellm-rust/crates/core/src/lib.rs +++ b/litellm-rust/crates/core/src/lib.rs @@ -1,3 +1,4 @@ +pub mod caching; pub mod call_lifecycle; pub mod constants; pub mod error; diff --git a/litellm-rust/crates/core/src/providers/bedrock/aws_base.rs b/litellm-rust/crates/core/src/providers/bedrock/aws_base.rs new file mode 100644 index 00000000000..82d1e8fdf91 --- /dev/null +++ b/litellm-rust/crates/core/src/providers/bedrock/aws_base.rs @@ -0,0 +1,724 @@ +use std::collections::BTreeMap; +use std::sync::{Mutex, OnceLock}; +use std::time::Duration; +use std::time::{SystemTime, UNIX_EPOCH}; + +use crate::caching::in_memory_cache::InMemoryCache; +use crate::error::{CoreError, CoreResult}; +use aws_credential_types::provider::ProvideCredentials; +use aws_credential_types::Credentials; +use aws_sigv4::http_request::{ + sign, SignableBody, SignableRequest, SigningParams, SigningSettings, +}; +use aws_sigv4::sign::v4; +use aws_smithy_runtime_api::client::identity::Identity; +use sha2::{Digest, Sha256}; + +use super::constants::{ + AWS_ACCESS_KEY_ID, AWS_EXTERNAL_ID, AWS_PROFILE_NAME, AWS_REGION_NAME, AWS_ROLE_ARN, + AWS_ROLE_NAME, AWS_SECRET_ACCESS_KEY, AWS_SESSION_NAME, AWS_SESSION_TOKEN, AWS_STS_ENDPOINT, + AWS_WEB_IDENTITY_TOKEN, AWS_WEB_IDENTITY_TOKEN_FILE, BEDROCK_SERVICE, + DEFAULT_SESSION_NAME_PREFIX, +}; + +const STATIC_CREDENTIALS_TTL: Duration = Duration::from_secs(3600 - 60); +const AMBIENT_CREDENTIALS_TTL: Duration = Duration::from_secs(600); + +static IAM_CREDENTIALS_CACHE: OnceLock>> = OnceLock::new(); + +fn credential_cache_ttl(flow: &AwsAuthFlow) -> Option { + match flow { + AwsAuthFlow::StaticKeys { .. } => Some(STATIC_CREDENTIALS_TTL), + AwsAuthFlow::DefaultChain => Some(AMBIENT_CREDENTIALS_TTL), + AwsAuthFlow::WebIdentity { .. } + | AwsAuthFlow::AssumeRole { .. } + | AwsAuthFlow::Profile { .. } + | AwsAuthFlow::SessionToken { .. } => None, + } +} + +#[derive(Clone, Debug, Default, PartialEq, Eq)] +pub struct AwsAuthConfig { + pub access_key_id: Option, + pub secret_access_key: Option, + pub session_token: Option, + pub region_name: Option, + pub session_name: Option, + pub profile_name: Option, + pub role_name: Option, + pub web_identity_token: Option, + pub sts_endpoint: Option, + pub external_id: Option, +} + +impl AwsAuthConfig { + fn with_environment(self, env_lookup: &dyn Fn(&str) -> Option) -> Self { + Self { + access_key_id: self.access_key_id.or_else(|| env_lookup(AWS_ACCESS_KEY_ID)), + secret_access_key: self + .secret_access_key + .or_else(|| env_lookup(AWS_SECRET_ACCESS_KEY)), + session_token: self.session_token.or_else(|| env_lookup(AWS_SESSION_TOKEN)), + region_name: self.region_name.or_else(|| env_lookup(AWS_REGION_NAME)), + session_name: self.session_name.or_else(|| env_lookup(AWS_SESSION_NAME)), + profile_name: self.profile_name.or_else(|| env_lookup(AWS_PROFILE_NAME)), + role_name: self.role_name.or_else(|| env_lookup(AWS_ROLE_NAME)), + web_identity_token: self + .web_identity_token + .or_else(|| env_lookup(AWS_WEB_IDENTITY_TOKEN)), + sts_endpoint: self.sts_endpoint.or_else(|| env_lookup(AWS_STS_ENDPOINT)), + external_id: self.external_id.or_else(|| env_lookup(AWS_EXTERNAL_ID)), + } + } +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub enum AwsAuthFlow { + WebIdentity { + token: String, + role: String, + session_name: String, + }, + AssumeRole { + role: String, + session_name: Option, + }, + Profile { + name: String, + }, + SessionToken { + access_key_id: String, + secret_access_key: String, + session_token: String, + }, + StaticKeys { + access_key_id: String, + secret_access_key: String, + region_name: String, + }, + DefaultChain, +} + +fn cache_key(config: &AwsAuthConfig, flow: &AwsAuthFlow) -> String { + let mut hasher = Sha256::new(); + hasher.update(format!("{config:?}:{flow:?}")); + format!("{:x}", hasher.finalize()) +} + +fn get_cached_credentials(key: &str) -> Option { + let cache = IAM_CREDENTIALS_CACHE.get_or_init(|| Mutex::new(InMemoryCache::default())); + let mut entries = cache.lock().ok()?; + entries.get_cache(key) +} + +fn set_cached_credentials(key: String, credentials: Credentials, ttl: Duration) { + let cache = IAM_CREDENTIALS_CACHE.get_or_init(|| Mutex::new(InMemoryCache::default())); + if let Ok(mut entries) = cache.lock() { + entries.set_cache(key, credentials, Some(ttl)); + } +} + +fn role_identity(arn: &str) -> Option<(&str, &str, &str)> { + let mut parts = arn.splitn(6, ':'); + let ("arn", partition, _, _, account, resource) = ( + parts.next()?, + parts.next()?, + parts.next()?, + parts.next()?, + parts.next()?, + parts.next()?, + ) else { + return None; + }; + let role = if let Some(role) = resource.strip_prefix("role/") { + role.rsplit('/').next()? + } else { + resource.strip_prefix("assumed-role/")?.split('/').next()? + }; + Some((partition, account, role)) +} + +fn same_role_arns(target: &str, caller: &str) -> bool { + role_identity(target) == role_identity(caller) +} + +pub fn classify_auth( + config: AwsAuthConfig, + env_lookup: &dyn Fn(&str) -> Option, +) -> AwsAuthFlow { + let config = config.with_environment(env_lookup); + if let (Some(token), Some(role), Some(session_name)) = ( + config.web_identity_token.clone(), + config.role_name.clone(), + config.session_name.clone(), + ) { + return AwsAuthFlow::WebIdentity { + token, + role, + session_name, + }; + } + if let Some(role) = config.role_name.clone() { + return AwsAuthFlow::AssumeRole { + role, + session_name: config.session_name.clone(), + }; + } + if let Some(name) = config.profile_name { + return AwsAuthFlow::Profile { name }; + } + if let (Some(access_key_id), Some(secret_access_key), Some(session_token)) = ( + config.access_key_id.clone(), + config.secret_access_key.clone(), + config.session_token, + ) { + return AwsAuthFlow::SessionToken { + access_key_id, + secret_access_key, + session_token, + }; + } + if let (Some(access_key_id), Some(secret_access_key), Some(region_name)) = ( + config.access_key_id, + config.secret_access_key, + config.region_name, + ) { + return AwsAuthFlow::StaticKeys { + access_key_id, + secret_access_key, + region_name, + }; + } + AwsAuthFlow::DefaultChain +} + +pub async fn resolve_credentials( + config: AwsAuthConfig, + env_lookup: &dyn Fn(&str) -> Option, +) -> CoreResult { + let resolved = config.clone().with_environment(env_lookup); + let flow = classify_auth(config, env_lookup); + match flow { + AwsAuthFlow::SessionToken { + access_key_id, + secret_access_key, + session_token, + } => Ok(Credentials::new( + access_key_id, + secret_access_key, + Some(session_token), + None, + "litellm-static-session", + )), + AwsAuthFlow::StaticKeys { + access_key_id, + secret_access_key, + region_name, + } => { + let flow = AwsAuthFlow::StaticKeys { + access_key_id: access_key_id.clone(), + secret_access_key: secret_access_key.clone(), + region_name, + }; + let key = cache_key(&resolved, &flow); + if let Some(credentials) = get_cached_credentials(&key) { + return Ok(credentials); + } + let credentials = Credentials::new( + access_key_id, + secret_access_key, + None, + None, + "litellm-static", + ); + set_cached_credentials( + key, + credentials.clone(), + credential_cache_ttl(&flow).unwrap_or(STATIC_CREDENTIALS_TTL), + ); + Ok(credentials) + } + AwsAuthFlow::Profile { name } => { + let provider = aws_config::profile::ProfileFileCredentialsProvider::builder() + .profile_name(name) + .build(); + provider.provide_credentials().await.map_err(|error| { + CoreError::Auth(format!("AWS profile credentials failed: {error}")) + }) + } + AwsAuthFlow::AssumeRole { role, session_name } => { + if is_already_running_as_role(&role, &resolved).await? { + let ambient_flow = AwsAuthFlow::DefaultChain; + let key = cache_key(&resolved, &ambient_flow); + if let Some(credentials) = get_cached_credentials(&key) { + return Ok(credentials); + } + let provider = + aws_config::default_provider::credentials::DefaultCredentialsChain::builder() + .build() + .await; + let credentials = provider.provide_credentials().await.map_err(|error| { + CoreError::Auth(format!("AWS default credentials failed: {error}")) + })?; + set_cached_credentials( + key, + credentials.clone(), + credential_cache_ttl(&ambient_flow).unwrap_or(AMBIENT_CREDENTIALS_TTL), + ); + return Ok(credentials); + } + let mut loader = aws_config::defaults(aws_config::BehaviorVersion::latest()); + if let Some(region) = resolved.region_name.clone() { + loader = loader.region(aws_types::region::Region::new(region)); + } + if let Some(endpoint) = resolved.sts_endpoint.clone() { + loader = loader.endpoint_url(endpoint); + } + if let (Some(access_key_id), Some(secret_access_key)) = + (resolved.access_key_id, resolved.secret_access_key) + { + loader = loader.credentials_provider(Credentials::new( + access_key_id, + secret_access_key, + resolved.session_token, + None, + "litellm-role-source", + )); + } + let sdk_config = loader.load().await; + let builder = aws_config::sts::AssumeRoleProvider::builder(role); + let builder = match session_name { + Some(name) => builder.session_name(name), + None => builder.session_name(default_session_name()), + }; + let builder = match resolved.external_id { + Some(id) => builder.external_id(id), + None => builder, + }; + let provider = builder.configure(&sdk_config).build().await; + provider + .provide_credentials() + .await + .map_err(|error| CoreError::Auth(format!("AWS role credentials failed: {error}"))) + } + AwsAuthFlow::WebIdentity { + token, + role, + session_name, + } => { + let mut loader = aws_config::defaults(aws_config::BehaviorVersion::latest()); + if let Some(region) = resolved.region_name { + loader = loader.region(aws_types::region::Region::new(region)); + } + if let Some(endpoint) = resolved.sts_endpoint { + loader = loader.endpoint_url(endpoint); + } + let sdk_config = loader.load().await; + let client = aws_sdk_sts::Client::new(&sdk_config); + let response = client + .assume_role_with_web_identity() + .role_arn(role) + .role_session_name(session_name) + .web_identity_token(token) + .send() + .await + .map_err(|error| { + CoreError::Auth(format!("AWS web identity credentials failed: {error}")) + })?; + let credentials = response.credentials().ok_or_else(|| { + CoreError::Auth("AWS web identity response had no credentials".to_string()) + })?; + let expiration = SystemTime::try_from(*credentials.expiration()).map_err(|error| { + CoreError::Auth(format!("AWS web identity expiration was invalid: {error}")) + })?; + Ok(Credentials::new( + credentials.access_key_id(), + credentials.secret_access_key(), + Some(credentials.session_token().to_string()), + Some(expiration), + "litellm-web-identity", + )) + } + AwsAuthFlow::DefaultChain => { + let key = cache_key(&resolved, &AwsAuthFlow::DefaultChain); + if let Some(credentials) = get_cached_credentials(&key) { + return Ok(credentials); + } + let provider = + aws_config::default_provider::credentials::DefaultCredentialsChain::builder() + .build() + .await; + let credentials = provider.provide_credentials().await.map_err(|error| { + CoreError::Auth(format!("AWS default credentials failed: {error}")) + })?; + set_cached_credentials( + key, + credentials.clone(), + credential_cache_ttl(&AwsAuthFlow::DefaultChain).unwrap_or(AMBIENT_CREDENTIALS_TTL), + ); + Ok(credentials) + } + } +} + +async fn is_already_running_as_role(role: &str, config: &AwsAuthConfig) -> CoreResult { + if role_identity(role).is_none() { + return Ok(false); + } + if let (Ok(current_role), Ok(token_file)) = ( + std::env::var(AWS_ROLE_ARN), + std::env::var(AWS_WEB_IDENTITY_TOKEN_FILE), + ) { + if !token_file.is_empty() { + return Ok(same_role_arns(role, ¤t_role)); + } + } + let mut loader = aws_config::defaults(aws_config::BehaviorVersion::latest()); + if let Some(region) = config.region_name.clone() { + loader = loader.region(aws_types::region::Region::new(region)); + } + if let Some(endpoint) = config.sts_endpoint.clone() { + loader = loader.endpoint_url(endpoint); + } + let sdk_config = loader.load().await; + let response = match aws_sdk_sts::Client::new(&sdk_config) + .get_caller_identity() + .send() + .await + { + Ok(response) => response, + Err(_) => return Ok(false), + }; + Ok(response + .arn() + .is_some_and(|caller| same_role_arns(role, caller))) +} + +fn default_session_name() -> String { + let seconds = SystemTime::now() + .duration_since(UNIX_EPOCH) + .map_or(0, |duration| duration.as_secs()); + format!("{DEFAULT_SESSION_NAME_PREFIX}-{seconds}") +} + +pub fn sign_bedrock_post( + url: &str, + body: &[u8], + headers: &BTreeMap, + region: &str, + credentials: &Credentials, + signing_time: SystemTime, +) -> CoreResult> { + let identity: Identity = credentials.clone().into(); + let params = v4::SigningParams::builder() + .identity(&identity) + .region(region) + .name(BEDROCK_SERVICE) + .time(signing_time) + .settings(SigningSettings::default()) + .build() + .map(SigningParams::from) + .map_err(|error| CoreError::Auth(format!("AWS signing parameters failed: {error}")))?; + let header_refs = headers + .iter() + .map(|(name, value)| (name.as_str(), value.as_str())); + let request = SignableRequest::new("POST", url, header_refs, SignableBody::Bytes(body)) + .map_err(|error| CoreError::Auth(format!("AWS signable request failed: {error}")))?; + let (instructions, _) = sign(request, ¶ms) + .map_err(|error| CoreError::Auth(format!("AWS request signing failed: {error}")))? + .into_parts(); + Ok(instructions + .headers() + .map(|(name, value)| { + let normalized_name = match name { + "authorization" => "Authorization", + "x-amz-date" => "X-Amz-Date", + "x-amz-security-token" => "X-Amz-Security-Token", + _ => name, + }; + (normalized_name.to_string(), value.to_string()) + }) + .collect()) +} + +#[cfg(test)] +mod tests { + use super::*; + + fn no_env(_: &str) -> Option { + None + } + + fn parity_inputs() -> (String, Vec, BTreeMap) { + ( + "https://bedrock-runtime.us-east-1.amazonaws.com/model/amazon.titan-text-express-v1/invoke" + .to_string(), + br#"{"input":"hello"}"#.to_vec(), + BTreeMap::from([("Content-Type".to_string(), "application/json".to_string())]), + ) + } + + #[test] + fn classification_preserves_python_precedence() { + let config = AwsAuthConfig { + access_key_id: Some("ak".into()), + secret_access_key: Some("sk".into()), + session_token: Some("token".into()), + region_name: Some("us-east-1".into()), + session_name: Some("session".into()), + profile_name: Some("profile".into()), + role_name: Some("role".into()), + web_identity_token: Some("oidc".into()), + ..Default::default() + }; + assert!(matches!( + classify_auth(config, &no_env), + AwsAuthFlow::WebIdentity { .. } + )); + } + + #[test] + fn classification_covers_fallthroughs() { + let env = |key: &str| match key { + AWS_PROFILE_NAME => Some("profile".into()), + _ => None, + }; + assert!(matches!( + classify_auth(AwsAuthConfig::default(), &env), + AwsAuthFlow::Profile { .. } + )); + assert!(matches!( + classify_auth( + AwsAuthConfig { + access_key_id: Some("ak".into()), + secret_access_key: Some("sk".into()), + session_token: Some("token".into()), + ..Default::default() + }, + &no_env + ), + AwsAuthFlow::SessionToken { .. } + )); + assert!(matches!( + classify_auth( + AwsAuthConfig { + access_key_id: Some("ak".into()), + secret_access_key: Some("sk".into()), + region_name: Some("us-east-1".into()), + ..Default::default() + }, + &no_env + ), + AwsAuthFlow::StaticKeys { .. } + )); + assert_eq!( + classify_auth(AwsAuthConfig::default(), &no_env), + AwsAuthFlow::DefaultChain + ); + } + + #[tokio::test] + async fn static_credentials_do_not_use_network() { + let credentials = resolve_credentials( + AwsAuthConfig { + access_key_id: Some("ak".into()), + secret_access_key: Some("sk".into()), + region_name: Some("us-east-1".into()), + ..Default::default() + }, + &no_env, + ) + .await + .expect("static credentials"); + assert_eq!(credentials.access_key_id(), "ak"); + assert_eq!(credentials.session_token(), None); + } + + #[test] + fn cache_policy_matches_python_flows() { + assert_eq!( + credential_cache_ttl(&AwsAuthFlow::StaticKeys { + access_key_id: "ak".into(), + secret_access_key: "sk".into(), + region_name: "us-east-1".into(), + }), + Some(STATIC_CREDENTIALS_TTL) + ); + assert_eq!( + credential_cache_ttl(&AwsAuthFlow::DefaultChain), + Some(AMBIENT_CREDENTIALS_TTL) + ); + assert_eq!( + credential_cache_ttl(&AwsAuthFlow::SessionToken { + access_key_id: "ak".into(), + secret_access_key: "sk".into(), + session_token: "token".into(), + }), + None + ); + assert_eq!( + credential_cache_ttl(&AwsAuthFlow::Profile { + name: "profile".into() + }), + None + ); + assert_eq!( + credential_cache_ttl(&AwsAuthFlow::AssumeRole { + role: "arn:aws:iam::123456789012:role/demo".into(), + session_name: None, + }), + None + ); + assert_eq!( + credential_cache_ttl(&AwsAuthFlow::WebIdentity { + token: "token".into(), + role: "arn:aws:iam::123456789012:role/demo".into(), + session_name: "session".into(), + }), + None + ); + } + + #[test] + fn cache_round_trip_preserves_credentials() { + let key = format!("cache-test-{}", std::process::id()); + let credentials = Credentials::new("cache-ak", "cache-sk", None, None, "test"); + set_cached_credentials(key.clone(), credentials.clone(), STATIC_CREDENTIALS_TTL); + assert_eq!( + get_cached_credentials(&key).map(|value| value.access_key_id().to_string()), + Some("cache-ak".to_string()) + ); + } + + #[test] + fn same_role_comparison_matches_partition_account_and_role() { + assert!(same_role_arns( + "arn:aws:iam::123456789012:role/path/demo", + "arn:aws:sts::123456789012:assumed-role/demo/session" + )); + assert!(!same_role_arns( + "arn:aws:iam::123456789012:role/demo", + "arn:aws:iam::999999999999:role/demo" + )); + assert!(!same_role_arns( + "arn:aws:iam::123456789012:role/demo", + "arn:aws-cn:iam::123456789012:role/demo" + )); + assert!(!same_role_arns( + "arn:aws:iam::123456789012:user/demo", + "arn:aws:iam::123456789012:role/demo" + )); + } + + #[test] + fn signing_matches_botocore_golden_vector() { + let (url, body, headers) = parity_inputs(); + let credentials = Credentials::new( + "AKIDEXAMPLE", + "wJalrXUtnFEMI/K7MDENG+bPxRfiCYEXAMPLEKEY", + Some("session-token".to_string()), + None, + "test", + ); + let signed = sign_bedrock_post( + &url, + &body, + &headers, + "us-east-1", + &credentials, + UNIX_EPOCH + std::time::Duration::from_secs(1_704_164_645), + ) + .expect("golden signature"); + assert_eq!( + signed.get("X-Amz-Date").map(String::as_str), + Some("20240102T030405Z") + ); + assert_eq!( + signed.get("X-Amz-Security-Token").map(String::as_str), + Some("session-token") + ); + assert_eq!( + signed.get("Authorization").map(String::as_str), + Some("AWS4-HMAC-SHA256 Credential=AKIDEXAMPLE/20240102/us-east-1/bedrock/aws4_request, SignedHeaders=content-type;host;x-amz-date;x-amz-security-token, Signature=55c027ef47527d3ad63f1735f9d099efdbc99f296ff914bd94e727e24ec0e464") + ); + } + + #[test] + fn signing_without_session_token_omits_security_header() { + let (url, body, headers) = parity_inputs(); + let credentials = Credentials::new( + "AKIDEXAMPLE", + "wJalrXUtnFEMI/K7MDENG+bPxRfiCYEXAMPLEKEY", + None, + None, + "test", + ); + let signed = sign_bedrock_post( + &url, + &body, + &headers, + "us-east-1", + &credentials, + UNIX_EPOCH + std::time::Duration::from_secs(1_704_164_645), + ) + .expect("signature"); + assert!(!signed.contains_key("X-Amz-Security-Token")); + } + + #[ignore] + #[tokio::test] + async fn live_bedrock_invoke_model_returns_200() -> Result<(), Box> { + let access_key_id = std::env::var("AWS_BEDROCK_TEST_ACCESS_KEY_ID")?; + let secret_access_key = std::env::var("AWS_BEDROCK_TEST_SECRET_ACCESS_KEY")?; + let body = br#"{"anthropic_version":"bedrock-2023-05-31","max_tokens":1,"messages":[{"role":"user","content":[{"type":"text","text":"ping"}]}]}"#.to_vec(); + let headers = + BTreeMap::from([("Content-Type".to_string(), "application/json".to_string())]); + let credentials = resolve_credentials( + AwsAuthConfig { + access_key_id: Some(access_key_id), + secret_access_key: Some(secret_access_key), + region_name: Some("us-west-2".to_string()), + ..Default::default() + }, + &no_env, + ) + .await?; + let client = reqwest::Client::new(); + let mut failures = Vec::new(); + + for region in ["us-west-2", "us-east-1"] { + let url = format!( + "https://bedrock-runtime.{region}.amazonaws.com/model/us.anthropic.claude-opus-4-8/invoke" + ); + let signed_headers = sign_bedrock_post( + &url, + &body, + &headers, + region, + &credentials, + SystemTime::now(), + )?; + let mut request = client.post(&url).body(body.clone()); + for (name, value) in &headers { + request = request.header(name, value); + } + for (name, value) in signed_headers { + request = request.header(name, value); + } + let response = request.send().await?; + let status = response.status(); + let response_body = response.text().await?; + let snippet: String = response_body.chars().take(240).collect(); + println!("region={region} status={status} response={snippet}"); + if status == reqwest::StatusCode::OK { + return Ok(()); + } + failures.push(format!("{region}: {status} {snippet}")); + } + + panic!( + "no Bedrock region returned HTTP 200: {}", + failures.join("; ") + ); + } +} diff --git a/litellm-rust/crates/core/src/providers/bedrock/constants.rs b/litellm-rust/crates/core/src/providers/bedrock/constants.rs new file mode 100644 index 00000000000..a08ae9de146 --- /dev/null +++ b/litellm-rust/crates/core/src/providers/bedrock/constants.rs @@ -0,0 +1,14 @@ +pub const AWS_ACCESS_KEY_ID: &str = "AWS_ACCESS_KEY_ID"; +pub const AWS_SECRET_ACCESS_KEY: &str = "AWS_SECRET_ACCESS_KEY"; +pub const AWS_SESSION_TOKEN: &str = "AWS_SESSION_TOKEN"; +pub const AWS_REGION_NAME: &str = "AWS_REGION_NAME"; +pub const AWS_SESSION_NAME: &str = "AWS_SESSION_NAME"; +pub const AWS_PROFILE_NAME: &str = "AWS_PROFILE_NAME"; +pub const AWS_ROLE_NAME: &str = "AWS_ROLE_NAME"; +pub const AWS_WEB_IDENTITY_TOKEN: &str = "AWS_WEB_IDENTITY_TOKEN"; +pub const AWS_ROLE_ARN: &str = "AWS_ROLE_ARN"; +pub const AWS_WEB_IDENTITY_TOKEN_FILE: &str = "AWS_WEB_IDENTITY_TOKEN_FILE"; +pub const AWS_STS_ENDPOINT: &str = "AWS_STS_ENDPOINT"; +pub const AWS_EXTERNAL_ID: &str = "AWS_EXTERNAL_ID"; +pub const BEDROCK_SERVICE: &str = "bedrock"; +pub const DEFAULT_SESSION_NAME_PREFIX: &str = "litellm-session"; diff --git a/litellm-rust/crates/core/src/providers/bedrock/mod.rs b/litellm-rust/crates/core/src/providers/bedrock/mod.rs new file mode 100644 index 00000000000..8027260a78b --- /dev/null +++ b/litellm-rust/crates/core/src/providers/bedrock/mod.rs @@ -0,0 +1,6 @@ +//! User-directed exception: this base provider owns AWS auth I/O for parity +//! with Python's `BaseAWSLLM`; the broader core purity guidance is reconciled +//! separately. + +pub mod aws_base; +mod constants; diff --git a/litellm-rust/crates/core/src/providers/mod.rs b/litellm-rust/crates/core/src/providers/mod.rs index dc9dc515e7d..805600d6dbe 100644 --- a/litellm-rust/crates/core/src/providers/mod.rs +++ b/litellm-rust/crates/core/src/providers/mod.rs @@ -1,5 +1,7 @@ pub mod anthropic; pub mod azure_ai; +#[cfg(feature = "bedrock-auth")] +pub mod bedrock; pub mod mistral; pub mod openai; pub mod vertex_ai; diff --git a/tests/test_litellm/llms/bedrock/test_base_aws_llm.py b/tests/test_litellm/llms/bedrock/test_base_aws_llm.py index 470448251c9..aaf523eacd5 100644 --- a/tests/test_litellm/llms/bedrock/test_base_aws_llm.py +++ b/tests/test_litellm/llms/bedrock/test_base_aws_llm.py @@ -15,6 +15,7 @@ from typing import Any, Dict, Optional from unittest.mock import MagicMock, patch from botocore.awsrequest import AWSPreparedRequest, AWSRequest +from botocore.auth import SigV4Auth from botocore.credentials import Credentials import litellm @@ -768,6 +769,30 @@ def test_get_request_headers_with_sigv4(): assert result == mock_request.prepare.return_value +def test_sigv4_matches_rust_golden_vector(): + request = AWSRequest( + method="POST", + url="https://bedrock-runtime.us-east-1.amazonaws.com/model/amazon.titan-text-express-v1/invoke", + data=b'{"input":"hello"}', + headers={"Content-Type": "application/json"}, + ) + credentials = Credentials( + "AKIDEXAMPLE", + "wJalrXUtnFEMI/K7MDENG+bPxRfiCYEXAMPLEKEY", + "session-token", + ) + with patch("botocore.auth.get_current_datetime", return_value=datetime(2024, 1, 2, 3, 4, 5)): + SigV4Auth(credentials, "bedrock", "us-east-1").add_auth(request) + assert request.headers["X-Amz-Date"] == "20240102T030405Z" + assert request.headers["X-Amz-Security-Token"] == "session-token" + assert ( + request.headers["Authorization"] + == "AWS4-HMAC-SHA256 Credential=AKIDEXAMPLE/20240102/us-east-1/bedrock/aws4_request, " + "SignedHeaders=content-type;host;x-amz-date;x-amz-security-token, " + "Signature=55c027ef47527d3ad63f1735f9d099efdbc99f296ff914bd94e727e24ec0e464" + ) + + def test_get_request_headers_with_api_key_bearer_token(): """ Test that get_request_headers uses the api_key parameter as a bearer token when provided