From 10413796c651aa4369a9db88c244dbb217ce6849 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Fri, 25 Sep 2026 16:18:36 +0000 Subject: [PATCH] test(rust): reorganize core crate tests and split cache and OCR suites (#43177) * test(rust): group cache tests under cache/ and fold test_ocr.py into ocr/ The two failure cases in test_ocr.py duplicated the upstream-500 and timeout rows of PUBLIC_FAILURES, so only the file-input encoding case moves to ocr/test_requests.py Co-Authored-By: Claude Opus 5.5 * test(rust): split the response cache suite into one file per backend test_response_cache.py grew to 2400 lines. Each backend now has its own file, shared fixtures live in cache/conftest.py and shared helpers in support/cache.py. The helpers alias the private native test handles once, dropping the per-call reportPrivateUsage hits Co-Authored-By: Claude Opus 5.5 * split tokenizer test * test(core): consolidate route integration tests under tests/ with rstest and wiremock Moves the public-API OCR route tests out of src/ocr/route.rs and document.rs into tests/ocr/, split per provider plus lifecycle, machine, and document tests, merging the duplicated pairs. Messages, audio transcription, and chat completions share one wiremock-based upstream and recording secret source in tests/support, and gain table-driven cases for auth, routing, upstream errors, streaming, and declines. Tests of litellm-llms items move to that crate. Co-Authored-By: Claude Opus 5.5 * test(messages): keep the stream relay test independent of the stream head contents The stream head carries no headers on main, so the relay test asserts the open-then-deliver order and the relayed body instead of header hand-off. Co-Authored-By: Claude Sonnet 5 --------- Co-authored-by: Yujong Lee Co-authored-by: Claude Opus 5.5 --- litellm-rust/Cargo.lock | 1 + litellm-rust/crates/core/Cargo.toml | 1 + .../core/src/chat_completions/handler.rs | 26 + .../core/src/chat_completions/prepare.rs | 244 -- .../crates/core/src/messages/common_utils.rs | 145 +- litellm-rust/crates/core/src/ocr/document.rs | 157 - litellm-rust/crates/core/src/ocr/mod.rs | 209 - litellm-rust/crates/core/src/ocr/prepare.rs | 205 +- litellm-rust/crates/core/src/ocr/route.rs | 3521 ----------------- .../crates/core/tests/audio_transcription.rs | 268 +- .../crates/core/tests/chat_completions.rs | 320 ++ litellm-rust/crates/core/tests/messages.rs | 471 --- .../crates/core/tests/messages/main.rs | 94 + .../crates/core/tests/messages/request.rs | 269 ++ .../crates/core/tests/messages/response.rs | 136 + .../crates/core/tests/messages/secrets.rs | 94 + .../crates/core/tests/messages/stream.rs | 163 + .../crates/core/tests/ocr/aws_textract.rs | 173 + .../crates/core/tests/ocr/azure_ai.rs | 270 ++ .../tests/ocr/azure_document_intelligence.rs | 441 +++ litellm-rust/crates/core/tests/ocr/cohere.rs | 42 + .../crates/core/tests/ocr/documents.rs | 182 + .../crates/core/tests/ocr/lifecycle.rs | 269 ++ litellm-rust/crates/core/tests/ocr/machine.rs | 284 ++ litellm-rust/crates/core/tests/ocr/main.rs | 125 + litellm-rust/crates/core/tests/ocr/mistral.rs | 248 ++ litellm-rust/crates/core/tests/ocr/reducto.rs | 321 ++ .../crates/core/tests/ocr/vertex_ai.rs | 184 + litellm-rust/crates/core/tests/support/mod.rs | 155 + .../llms/src/reducto/ocr/transformation.rs | 44 + litellm-rust/crates/llms/tests/ocr_handler.rs | 79 + tests/test_litellm_rust/AGENTS.md | 1 + tests/test_litellm_rust/cache/__init__.py | 1 + tests/test_litellm_rust/cache/conftest.py | 30 + .../cache/test_azure_blob.py | 173 + tests/test_litellm_rust/cache/test_disk.py | 117 + tests/test_litellm_rust/cache/test_facade.py | 397 ++ tests/test_litellm_rust/cache/test_gcs.py | 242 ++ .../cache/test_qdrant_semantic.py | 286 ++ tests/test_litellm_rust/cache/test_redis.py | 228 ++ .../cache/test_redis_semantic.py | 606 +++ tests/test_litellm_rust/cache/test_rollout.py | 264 ++ tests/test_litellm_rust/cache/test_s3.py | 187 + .../test_valkey_semantic.py} | 0 tests/test_litellm_rust/ocr/test_requests.py | 15 + tests/test_litellm_rust/support/cache.py | 40 + tests/test_litellm_rust/test_cache.py | 2397 ----------- tests/test_litellm_rust/test_ocr.py | 134 - tests/test_litellm_rust/tokenizer/__init__.py | 0 .../test_fast_count.py} | 61 - .../tokenizer/test_huggingface.py | 24 + .../tokenizer/test_tiktoken.py | 53 + 52 files changed, 7015 insertions(+), 7382 deletions(-) create mode 100644 litellm-rust/crates/core/tests/chat_completions.rs delete mode 100644 litellm-rust/crates/core/tests/messages.rs create mode 100644 litellm-rust/crates/core/tests/messages/main.rs create mode 100644 litellm-rust/crates/core/tests/messages/request.rs create mode 100644 litellm-rust/crates/core/tests/messages/response.rs create mode 100644 litellm-rust/crates/core/tests/messages/secrets.rs create mode 100644 litellm-rust/crates/core/tests/messages/stream.rs create mode 100644 litellm-rust/crates/core/tests/ocr/aws_textract.rs create mode 100644 litellm-rust/crates/core/tests/ocr/azure_ai.rs create mode 100644 litellm-rust/crates/core/tests/ocr/azure_document_intelligence.rs create mode 100644 litellm-rust/crates/core/tests/ocr/cohere.rs create mode 100644 litellm-rust/crates/core/tests/ocr/documents.rs create mode 100644 litellm-rust/crates/core/tests/ocr/lifecycle.rs create mode 100644 litellm-rust/crates/core/tests/ocr/machine.rs create mode 100644 litellm-rust/crates/core/tests/ocr/main.rs create mode 100644 litellm-rust/crates/core/tests/ocr/mistral.rs create mode 100644 litellm-rust/crates/core/tests/ocr/reducto.rs create mode 100644 litellm-rust/crates/core/tests/ocr/vertex_ai.rs create mode 100644 litellm-rust/crates/core/tests/support/mod.rs create mode 100644 litellm-rust/crates/llms/tests/ocr_handler.rs create mode 100644 tests/test_litellm_rust/AGENTS.md create mode 100644 tests/test_litellm_rust/cache/__init__.py create mode 100644 tests/test_litellm_rust/cache/conftest.py create mode 100644 tests/test_litellm_rust/cache/test_azure_blob.py create mode 100644 tests/test_litellm_rust/cache/test_disk.py create mode 100644 tests/test_litellm_rust/cache/test_facade.py create mode 100644 tests/test_litellm_rust/cache/test_gcs.py create mode 100644 tests/test_litellm_rust/cache/test_qdrant_semantic.py create mode 100644 tests/test_litellm_rust/cache/test_redis.py create mode 100644 tests/test_litellm_rust/cache/test_redis_semantic.py create mode 100644 tests/test_litellm_rust/cache/test_rollout.py create mode 100644 tests/test_litellm_rust/cache/test_s3.py rename tests/test_litellm_rust/{test_valkey_semantic_cache_native.py => cache/test_valkey_semantic.py} (100%) create mode 100644 tests/test_litellm_rust/support/cache.py delete mode 100644 tests/test_litellm_rust/test_cache.py delete mode 100644 tests/test_litellm_rust/test_ocr.py create mode 100644 tests/test_litellm_rust/tokenizer/__init__.py rename tests/test_litellm_rust/{test_tokenizer.py => tokenizer/test_fast_count.py} (51%) create mode 100644 tests/test_litellm_rust/tokenizer/test_huggingface.py create mode 100644 tests/test_litellm_rust/tokenizer/test_tiktoken.py diff --git a/litellm-rust/Cargo.lock b/litellm-rust/Cargo.lock index 5de32f5b622..2d96efa6077 100644 --- a/litellm-rust/Cargo.lock +++ b/litellm-rust/Cargo.lock @@ -3092,6 +3092,7 @@ dependencies = [ "tokio-tungstenite", "url", "veil", + "wiremock", ] [[package]] diff --git a/litellm-rust/crates/core/Cargo.toml b/litellm-rust/crates/core/Cargo.toml index d7096cdd774..12410c187e2 100644 --- a/litellm-rust/crates/core/Cargo.toml +++ b/litellm-rust/crates/core/Cargo.toml @@ -40,3 +40,4 @@ litellm-auth-gcp.workspace = true litellm-llms = { workspace = true, features = ["test-support"] } rstest.workspace = true rstest_reuse.workspace = true +wiremock = "0.6.5" diff --git a/litellm-rust/crates/core/src/chat_completions/handler.rs b/litellm-rust/crates/core/src/chat_completions/handler.rs index de926c715d5..2391ab83a60 100644 --- a/litellm-rust/crates/core/src/chat_completions/handler.rs +++ b/litellm-rust/crates/core/src/chat_completions/handler.rs @@ -85,3 +85,29 @@ pub(super) async fn outbound_request( other => other, }) } + +#[cfg(test)] +mod tests { + use super::{Error, as_response_error}; + + #[test] + fn response_errors_collapse_to_one_variant_that_can_only_mean_already_sent() { + for original in [ + Error::MissingField("usage"), + Error::Unsupported("non-text response content block"), + Error::InvalidRequest("whatever".to_string()), + Error::Auth(litellm_auth::Error::InvalidHeader), + ] { + let label = format!("{original:?}"); + assert!( + matches!(as_response_error(original), Error::InvalidResponse(_)), + "{label} must not stay retryable once the provider has answered" + ); + } + let upstream = Error::Transport(litellm_http::transport::Error::Http { + status: 500, + body: "boom".to_string(), + }); + assert_eq!(as_response_error(upstream.clone()), upstream); + } +} diff --git a/litellm-rust/crates/core/src/chat_completions/prepare.rs b/litellm-rust/crates/core/src/chat_completions/prepare.rs index afea46221f5..b6425773964 100644 --- a/litellm-rust/crates/core/src/chat_completions/prepare.rs +++ b/litellm-rust/crates/core/src/chat_completions/prepare.rs @@ -736,248 +736,4 @@ mod tests { .unwrap_or_else(|error| panic!("prepare declined {messages}: {error}")); } } - - mod round_trip { - use tokio::{ - io::{AsyncReadExt, AsyncWriteExt}, - net::{TcpListener, TcpStream}, - }; - - use super::*; - use crate::chat_completions::chat_completions; - - async fn read_http_request(socket: &mut TcpStream) -> String { - let mut request = Vec::new(); - let mut buffer = [0_u8; 1024]; - let header_end = loop { - let n = socket.read(&mut buffer).await.expect("reads request"); - if n == 0 { - break request.len(); - } - request.extend_from_slice(&buffer[..n]); - if let Some(position) = request.windows(4).position(|window| window == b"\r\n\r\n") - { - break position + 4; - } - }; - let headers = String::from_utf8_lossy(&request[..header_end]); - let content_length = headers - .lines() - .find_map(|line| { - let (name, value) = line.split_once(':')?; - name.eq_ignore_ascii_case("content-length") - .then(|| value.trim().parse::().ok()) - .flatten() - }) - .unwrap_or(0); - while request.len().saturating_sub(header_end) < content_length { - let n = socket.read(&mut buffer).await.expect("reads body"); - if n == 0 { - break; - } - request.extend_from_slice(&buffer[..n]); - } - String::from_utf8(request).expect("request is utf8") - } - - fn http_response(status: &str, body: &str) -> String { - format!( - "HTTP/1.1 {status}\r\ncontent-type: application/json\r\ncontent-length: {}\r\nconnection: close\r\n\r\n{}", - body.len(), - body - ) - } - - /// Serve one request from a stub upstream and hand back what it received. - async fn serve_once( - status: &'static str, - body: &'static str, - ) -> (String, tokio::task::JoinHandle) { - let listener = TcpListener::bind("127.0.0.1:0").await.expect("binds"); - let port = listener.local_addr().expect("addr").port(); - let handle = tokio::spawn(async move { - let (mut socket, _) = listener.accept().await.expect("accepts"); - let received = read_http_request(&mut socket).await; - socket - .write_all(http_response(status, body).as_bytes()) - .await - .expect("writes response"); - socket.flush().await.expect("flushes"); - received - }); - (format!("http://127.0.0.1:{port}/v1/messages"), handle) - } - - fn call(api_base: &str, messages: Value, params: Value) -> ChatCompletionsRequest<'_> { - ChatCompletionsRequest { - model: "anthropic/claude-sonnet-4-5", - messages, - optional_params: match params { - Value::Object(map) => map, - other => panic!("params must be an object, got {other}"), - }, - api_key: Some("sk-test"), - api_base: Some(api_base), - custom_llm_provider: None, - extra_headers: None, - timeout: Some(std::time::Duration::from_secs(10)), - } - } - - const GOOD_BODY: &str = r#"{"id":"msg_1","type":"message","role":"assistant","model":"claude-sonnet-4-5-20260101","content":[{"type":"text","text":"hello"}],"stop_reason":"end_turn","stop_sequence":null,"usage":{"input_tokens":11,"output_tokens":4}}"#; - - #[tokio::test] - async fn round_trip_sends_the_translated_body_and_normalizes_the_response() { - let (api_base, handle) = serve_once("200 OK", GOOD_BODY).await; - let response = chat_completions(call( - &api_base, - json!([ - {"role": "system", "content": "be terse"}, - {"role": "user", "content": "hi"} - ]), - json!({"max_tokens": 16}), - )) - .await - .expect("call succeeds"); - - let received = handle.await.expect("server task"); - let sent: Value = serde_json::from_str( - received - .split_once("\r\n\r\n") - .expect("request has a body") - .1, - ) - .expect("body is json"); - assert_eq!( - sent["messages"], - json!([{"role": "user", "content": [{"type": "text", "text": "hi"}]}]) - ); - assert_eq!( - sent["system"], - json!([{"type": "text", "text": "be terse"}]) - ); - assert_eq!(sent["max_tokens"], json!(16)); - assert!(received.to_lowercase().contains("x-api-key: sk-test")); - - assert_eq!( - response.choices[0].message.content.as_deref(), - Some("hello") - ); - assert_eq!(response.usage.total_tokens, 15); - } - - #[tokio::test] - async fn a_response_it_cannot_normalize_is_reported_as_already_sent() { - // The provider was called and billed, so the host must not retry this - // on its own path. `MissingField` here would read as a pre-send - // decline and be retried; `InvalidResponse` cannot. - const NO_USAGE: &str = - r#"{"model":"m","content":[{"type":"text","text":"hi"}],"stop_reason":"end_turn"}"#; - let (api_base, handle) = serve_once("200 OK", NO_USAGE).await; - let err = chat_completions(call( - &api_base, - json!([{"role": "user", "content": "hi"}]), - json!({"max_tokens": 16}), - )) - .await - .expect_err("response cannot be normalized"); - handle.await.expect("server task"); - assert!( - matches!(err, Error::InvalidResponse(_)), - "expected a post-send error, got {err:?}" - ); - } - - #[tokio::test] - async fn a_tool_use_block_in_the_response_is_also_reported_as_already_sent() { - const TOOL_USE: &str = r#"{"model":"m","content":[{"type":"tool_use","id":"t","name":"f","input":{}}],"stop_reason":"tool_use","usage":{"input_tokens":1,"output_tokens":1}}"#; - let (api_base, handle) = serve_once("200 OK", TOOL_USE).await; - let err = chat_completions(call( - &api_base, - json!([{"role": "user", "content": "hi"}]), - json!({"max_tokens": 16}), - )) - .await - .expect_err("response cannot be normalized"); - handle.await.expect("server task"); - assert!( - matches!(err, Error::InvalidResponse(_)), - "expected a post-send error, got {err:?}" - ); - } - - #[tokio::test] - async fn an_upstream_error_status_keeps_its_code() { - let (api_base, handle) = - serve_once("429 Too Many Requests", r#"{"error":"slow down"}"#).await; - let err = chat_completions(call( - &api_base, - json!([{"role": "user", "content": "hi"}]), - json!({"max_tokens": 16}), - )) - .await - .expect_err("upstream rejects"); - handle.await.expect("server task"); - assert!( - matches!( - err, - Error::Transport(litellm_http::transport::Error::Http { status: 429, .. }) - ), - "expected a 429, got {err:?}" - ); - } - - #[tokio::test] - async fn a_connection_that_is_never_established_declines_instead_of_failing() { - // Nothing was sent, so nothing was billed and the host can still serve - // the request. Classing this with the post-send failures would turn a - // recoverable fallback into a user-facing error on exactly the - // deployments whose transport is configured only on the Python client. - let port = { - let listener = TcpListener::bind("127.0.0.1:0").await.expect("binds"); - listener.local_addr().expect("has an address").port() - // Dropped here, so the port is closed and the connect is refused. - }; - let err = chat_completions(call( - &format!("http://127.0.0.1:{port}/v1/messages"), - json!([{"role": "user", "content": "hi"}]), - json!({"max_tokens": 16}), - )) - .await - .expect_err("nothing is listening"); - assert!( - matches!( - err, - Error::Transport(litellm_http::transport::Error::Connect(_)) - ), - "expected a pre-send connect failure, got {err:?}" - ); - } - - #[test] - fn response_errors_collapse_to_one_variant_that_can_only_mean_already_sent() { - use crate::chat_completions::handler::as_response_error; - - for original in [ - Error::MissingField("usage"), - Error::Unsupported("non-text response content block"), - Error::InvalidRequest("whatever".to_string()), - Error::Auth(litellm_auth::Error::InvalidHeader), - ] { - let label = format!("{original:?}"); - assert!( - matches!(as_response_error(original), Error::InvalidResponse(_)), - "{label} must not stay retryable once the provider has answered" - ); - } - // An upstream status is already unambiguous, so it survives intact. - assert!(matches!( - as_response_error(Error::Transport(litellm_http::transport::Error::Http { - status: 500, - body: "boom".to_string() - })), - Error::Transport(litellm_http::transport::Error::Http { status: 500, .. }) - )); - } - } } diff --git a/litellm-rust/crates/core/src/messages/common_utils.rs b/litellm-rust/crates/core/src/messages/common_utils.rs index 4327754ed05..d27b79bdc04 100644 --- a/litellm-rust/crates/core/src/messages/common_utils.rs +++ b/litellm-rust/crates/core/src/messages/common_utils.rs @@ -29,151 +29,10 @@ pub(super) fn string_headers( #[cfg(test)] mod tests { - use std::{sync::Arc, time::Duration}; - - use futures_util::future::BoxFuture; - use litellm_secrets::{SecretValue, source::SecretSource}; - use serde_json::{Value, json}; - use tokio::{ - io::{AsyncReadExt, AsyncWriteExt}, - net::{TcpListener, TcpStream}, - }; + use serde_json::json; use super::{messages_provider_config, string_headers, truncate_error_body}; - use crate::messages::{ - Error, - route::{LocalMessagesHost, MessagesCall, MessagesOutput, messages_machine}, - types::MessagesShaping, - }; - - struct RecordingSecrets { - values: Vec<(&'static str, String)>, - requested: std::sync::Mutex>, - } - - impl SecretSource for RecordingSecrets { - fn get_secret_str<'a>( - &'a self, - name: &'a str, - ) -> BoxFuture<'a, Result, litellm_secrets::Error>> { - Box::pin(async move { - self.requested.lock().unwrap().push(name.to_string()); - Ok(self - .values - .iter() - .find(|(key, _)| *key == name) - .map(|(_, value)| SecretValue::new(value.clone()))) - }) - } - } - - fn secrets_call() -> MessagesCall { - let Value::Object(body) = json!({ - "model": "claude-sonnet-4-5", - "max_tokens": 16, - "messages": [{"role": "user", "content": "hi"}] - }) else { - unreachable!("literal object") - }; - MessagesCall { - model: "claude-sonnet-4-5".into(), - body, - api_key: None, - api_base: None, - custom_llm_provider: Some("anthropic".into()), - extra_headers: None, - provider_specific_header: None, - timeout: Some(Duration::from_secs(5)), - shaping: MessagesShaping::default(), - } - } - - async fn read_http_request(socket: &mut TcpStream) -> String { - let mut request = Vec::new(); - let mut buffer = [0_u8; 1024]; - let header_end = loop { - let n = socket.read(&mut buffer).await.expect("reads request"); - if n == 0 { - break request.len(); - } - request.extend_from_slice(&buffer[..n]); - if let Some(position) = request.windows(4).position(|window| window == b"\r\n\r\n") { - break position + 4; - } - }; - let headers = String::from_utf8_lossy(&request[..header_end]); - let content_length = headers - .lines() - .find_map(|line| { - let (name, value) = line.split_once(':')?; - name.eq_ignore_ascii_case("content-length") - .then(|| value.trim().parse::().ok()) - .flatten() - }) - .unwrap_or(0); - while request.len().saturating_sub(header_end) < content_length { - let n = socket.read(&mut buffer).await.expect("reads body"); - if n == 0 { - break; - } - request.extend_from_slice(&buffer[..n]); - } - String::from_utf8(request).expect("request is utf8") - } - - #[tokio::test] - async fn route_reads_the_provider_credential_and_base_from_the_secret_source() { - let listener = TcpListener::bind("127.0.0.1:0").await.expect("binds"); - let addr = listener.local_addr().expect("addr"); - let server = tokio::spawn(async move { - let (mut socket, _) = listener.accept().await.expect("accepts request"); - let request = read_http_request(&mut socket).await; - let response_body = r#"{"id":"msg_1","type":"message","role":"assistant","content":[],"model":"claude-sonnet-4-5","stop_reason":"end_turn","usage":{"input_tokens":1,"output_tokens":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 secrets = Arc::new(RecordingSecrets { - values: vec![ - ("ANTHROPIC_API_KEY", "sk-from-manager".to_string()), - ("ANTHROPIC_BASE_URL", format!("http://{addr}")), - ], - requested: std::sync::Mutex::new(Vec::new()), - }); - - let output = litellm_host::run::run( - messages_machine(secrets.clone()), - &LocalMessagesHost::new(secrets_call()), - ) - .await - .expect("messages request succeeds"); - - assert!(matches!(output, MessagesOutput::Message(_))); - let request = server.await.expect("server task completes"); - assert!( - request - .to_ascii_lowercase() - .contains("x-api-key: sk-from-manager"), - "{request}" - ); - let requested = secrets.requested.lock().unwrap().clone(); - assert_eq!( - requested, - messages_provider_config("anthropic") - .unwrap() - .secret_names() - .iter() - .map(ToString::to_string) - .collect::>() - ); - } + use crate::messages::Error; #[test] fn provider_config_resolves_anthropic_and_azure_ai() { diff --git a/litellm-rust/crates/core/src/ocr/document.rs b/litellm-rust/crates/core/src/ocr/document.rs index b78c09298de..c33ee053422 100644 --- a/litellm-rust/crates/core/src/ocr/document.rs +++ b/litellm-rust/crates/core/src/ocr/document.rs @@ -244,160 +244,3 @@ mod tests { } } } - -#[cfg(test)] -mod document_tests { - use litellm_host::event::WireRequest; - use litellm_llms::base_llm::ocr::error::Error; - use rstest::rstest; - use serde_json::{Value, json}; - - use crate::ocr::route::LocalOcrHost; - use crate::ocr::test_support::{ - MockResponse, SERVED_DOCUMENT, document_server, mock_server, perform_ocr_with, - request_body, wire_request_with_document, - }; - - #[derive(Clone, Copy, Debug)] - enum Route { - Mistral, - AzureAi, - VertexMistral, - AzureCohereParse, - Cohere, - } - - impl Route { - fn model(self) -> &'static str { - match self { - Self::Mistral => "mistral/model", - Self::AzureAi => "azure_ai/model", - Self::VertexMistral => "vertex_ai/mistral-ocr-maas", - Self::AzureCohereParse => "azure_ai/cohere-parse", - Self::Cohere => "cohere/model", - } - } - - fn document_type(self) -> &'static str { - match self { - Self::Mistral | Self::AzureAi | Self::VertexMistral => "document_url", - Self::AzureCohereParse | Self::Cohere => "image_url", - } - } - - fn options(self) -> Value { - match self { - Self::Mistral | Self::AzureAi => json!({"pages": [0]}), - Self::VertexMistral => json!({"pages": [0], "vertex_project": "project-1"}), - Self::AzureCohereParse | Self::Cohere => json!({"output_format": "markdown"}), - } - } - } - - /// What the host does to the wire request in `before_send`. - #[derive(Clone, Copy, Debug)] - enum Host { - Detached, - ReplacesDocument, - } - - const REPLACED_DOCUMENT: &str = "data:image/png;base64,cmVwbGFjZWQ="; - - impl Host { - fn before_send(self, wire: WireRequest) -> WireRequest { - let Value::Object(fields) = wire.body else { - return wire; - }; - let body = fields - .into_iter() - .map(|(name, value)| match self { - Self::Detached => (name, value), - Self::ReplacesDocument if name == "document" => { - let document_type = value["type"].clone(); - let key = document_type.as_str().unwrap_or_default().to_string(); - (name, json!({"type": document_type, key: REPLACED_DOCUMENT})) - } - Self::ReplacesDocument => (name, value), - }) - .collect(); - WireRequest { - body: Value::Object(body), - ..wire - } - } - } - - struct Sent { - result: Result<(), Error>, - provider_body: Option, - } - - async fn send(route: Route, host: Host, document_base: &str) -> Sent { - let (base, seen, provider) = - mock_server(vec![MockResponse::json(json!({"pages": []}))]).await; - let document_type = route.document_type(); - let document = - json!({"type": document_type, document_type: format!("{document_base}/scan.png")}); - let request = wire_request_with_document(route.model(), &base, document, route.options()); - let local = - LocalOcrHost::new(request).with_before_send(move |wire, _| Ok(host.before_send(wire))); - let result = perform_ocr_with(local).await.map(|_| ()); - match result { - Ok(()) => provider.await.unwrap(), - Err(_) => provider.abort(), - } - let provider_body = seen - .lock() - .unwrap() - .first() - .map(|request| request_body(request)); - Sent { - result, - provider_body, - } - } - - fn served_document_uri() -> String { - use base64::Engine; - format!( - "data:image/png;base64,{}", - base64::engine::general_purpose::STANDARD.encode(SERVED_DOCUMENT) - ) - } - - #[rstest] - #[case::azure_ai(Route::AzureAi)] - #[case::vertex_mistral(Route::VertexMistral)] - #[case::azure_cohere_parse(Route::AzureCohereParse)] - #[tokio::test] - async fn inlining_routes_send_the_downloaded_document(#[case] route: Route) { - let (document_base, _documents) = document_server().await; - let sent = send(route, Host::Detached, &document_base).await; - sent.result.unwrap(); - assert_eq!( - sent.provider_body.unwrap()["document"][route.document_type()], - json!(served_document_uri()) - ); - } - - #[rstest] - #[tokio::test] - async fn document_replaced_by_the_host_reaches_the_provider( - #[values( - Route::Mistral, - Route::AzureAi, - Route::VertexMistral, - Route::AzureCohereParse, - Route::Cohere - )] - route: Route, - ) { - let (document_base, _documents) = document_server().await; - let sent = send(route, Host::ReplacesDocument, &document_base).await; - sent.result.unwrap(); - assert_eq!( - sent.provider_body.unwrap()["document"][route.document_type()], - json!(REPLACED_DOCUMENT) - ); - } -} diff --git a/litellm-rust/crates/core/src/ocr/mod.rs b/litellm-rust/crates/core/src/ocr/mod.rs index 270a402c9fa..2a0d20f69c9 100644 --- a/litellm-rust/crates/core/src/ocr/mod.rs +++ b/litellm-rust/crates/core/src/ocr/mod.rs @@ -7,212 +7,3 @@ pub mod provider_config; pub mod route; pub mod types; pub mod wire; - -#[cfg(test)] -pub(crate) mod test_support { - use std::sync::{Arc, Mutex}; - - use futures_util::future::BoxFuture; - use litellm_host::event::WireRequest; - use litellm_llms::base_llm::ocr::{ - error::Error, - handler::{CallHooks, OcrClient}, - transformation::LiteLLMOcrResponse, - }; - use serde_json::{Value, json}; - use tokio::{ - io::{AsyncReadExt, AsyncWriteExt}, - net::TcpListener, - }; - - use crate::ocr::{ - route::{LocalOcrHost, ocr_machine}, - types::LiteLLMOcrRequest, - wire::{OcrWireRequest, decode_request}, - }; - - /// Stands in for a host with no hooks registered: the wire request goes out unchanged - /// and response events go nowhere. - pub(crate) struct NoHooks; - - impl CallHooks for NoHooks { - fn before_send(&self, wire: WireRequest) -> BoxFuture<'_, Result> { - Box::pin(async move { Ok(wire) }) - } - - fn response_received<'a>(&'a self, _body: &'a [u8]) -> BoxFuture<'a, Result<(), Error>> { - Box::pin(async { Ok(()) }) - } - } - - pub(crate) fn ocr_client() -> OcrClient { - let document_http = reqwest::Client::builder() - .redirect(reqwest::redirect::Policy::none()) - .build() - .expect("test document client builds"); - OcrClient::for_test(reqwest::Client::new(), document_http) - } - - pub(crate) async fn perform_ocr( - request: LiteLLMOcrRequest, - ) -> Result { - crate::ocr::client::perform(&ocr_client(), request).await - } - - pub(crate) async fn perform_ocr_with(host: LocalOcrHost) -> Result { - litellm_host::run::run(ocr_machine(ocr_client()), &host).await - } - - pub(crate) fn wire_request(model: &str, base: &str, options: Value) -> LiteLLMOcrRequest { - wire_request_with_document( - model, - base, - json!({"type":"document_url","document_url":"data:application/pdf;base64,YWJj"}), - options, - ) - } - - pub(crate) fn wire_request_with_document( - model: &str, - base: &str, - document: Value, - options: Value, - ) -> LiteLLMOcrRequest { - decode_request(OcrWireRequest { - model: model.into(), - document, - api_key: Some(litellm_auth::SecretValue::new("test-key")), - api_base: Some(base.into()), - custom_llm_provider: None, - extra_headers: None, - optional_params: options.as_object().unwrap().clone(), - input_sources: Default::default(), - timeout_seconds: Some(2.0), - }) - .unwrap() - } - - pub(crate) fn resolved_request( - request: LiteLLMOcrRequest, - ) -> crate::ocr::types::ResolvedOcrRequest { - request - .map_document(crate::ocr::document::prepare_document) - .unwrap() - } - - pub(crate) fn with_source(request: LiteLLMOcrRequest, source: &str) -> LiteLLMOcrRequest { - let request = resolved_request(request); - let document = request.document.clone().with_source(source.into()); - request.with_document(document.into()) - } - - pub(crate) fn request_body(request: &str) -> Value { - serde_json::from_str(request.split_once("\r\n\r\n").unwrap().1).unwrap() - } - - pub(crate) const SERVED_DOCUMENT: &[u8] = b"\x89PNG served document"; - - /// Serves [`SERVED_DOCUMENT`] as `image/png` to every connection until aborted. - pub(crate) async fn document_server() -> (String, tokio::task::JoinHandle<()>) { - let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); - let base = format!("http://{}", listener.local_addr().unwrap()); - let task = tokio::spawn(async move { - loop { - let (mut socket, _) = listener.accept().await.unwrap(); - let mut buffer = [0u8; 4096]; - let _ = socket.read(&mut buffer).await.unwrap(); - let head = format!( - "HTTP/1.1 200 OK\r\nContent-Type: image/png\r\nContent-Length: {}\r\nConnection: close\r\n\r\n", - SERVED_DOCUMENT.len() - ); - socket.write_all(head.as_bytes()).await.unwrap(); - socket.write_all(SERVED_DOCUMENT).await.unwrap(); - } - }); - (base, task) - } - - pub(crate) struct MockResponse { - pub status: u16, - pub headers: Vec<(&'static str, String)>, - pub body: Value, - } - - impl MockResponse { - pub fn json(body: Value) -> Self { - Self { - status: 200, - headers: vec![], - body, - } - } - } - - pub(crate) async fn mock_server( - responses: Vec, - ) -> (String, Arc>>, tokio::task::JoinHandle<()>) { - let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); - let base = format!("http://{}", listener.local_addr().unwrap()); - let requests = Arc::new(Mutex::new(Vec::new())); - let seen = requests.clone(); - let server_base = base.clone(); - let task = tokio::spawn(async move { - for response in responses { - let (mut socket, _) = listener.accept().await.unwrap(); - let mut bytes = Vec::new(); - let mut buffer = [0u8; 4096]; - let header_end = loop { - let n = socket.read(&mut buffer).await.unwrap(); - assert!(n > 0); - bytes.extend_from_slice(&buffer[..n]); - if let Some(index) = bytes.windows(4).position(|s| s == b"\r\n\r\n") { - break index + 4; - } - }; - let length = String::from_utf8_lossy(&bytes[..header_end]) - .lines() - .find_map(|line| { - let (name, value) = line.split_once(':')?; - name.eq_ignore_ascii_case("content-length") - .then(|| value.trim().parse::().unwrap()) - }) - .unwrap_or(0); - while bytes.len() < header_end + length { - let n = socket.read(&mut buffer).await.unwrap(); - assert!(n > 0); - bytes.extend_from_slice(&buffer[..n]); - } - seen.lock() - .unwrap() - .push(String::from_utf8_lossy(&bytes).into_owned()); - let body = serde_json::to_vec(&response.body).unwrap(); - let headers = response - .headers - .into_iter() - .map(|(name, value)| { - format!("{name}: {}\r\n", value.replace("{base}", &server_base)) - }) - .collect::(); - let head = format!( - "HTTP/1.1 {} OK\r\nContent-Type: application/json\r\nContent-Length: {}\r\nConnection: close\r\n{}\r\n", - response.status, - body.len(), - headers - ); - socket.write_all(head.as_bytes()).await.unwrap(); - socket.write_all(&body).await.unwrap(); - } - }); - (base, requests, task) - } - - pub(crate) fn header<'a>(request: &'a str, name: &str) -> Option<&'a str> { - request - .lines() - .take_while(|line| !line.is_empty()) - .find_map(|line| { - let (key, value) = line.split_once(':')?; - key.eq_ignore_ascii_case(name).then(|| value.trim()) - }) - } -} diff --git a/litellm-rust/crates/core/src/ocr/prepare.rs b/litellm-rust/crates/core/src/ocr/prepare.rs index 37c0f18f659..18961ec96fa 100644 --- a/litellm-rust/crates/core/src/ocr/prepare.rs +++ b/litellm-rust/crates/core/src/ocr/prepare.rs @@ -70,20 +70,203 @@ pub(crate) fn prepare_request( } } -#[cfg(test)] -pub(crate) fn prepare_request_for_test(request: ResolvedOcrRequest) -> PreparedOcrRequest { - prepare_request( - request, - true, - &OcrClient::for_test(reqwest::Client::new(), reqwest::Client::new()), - std::sync::Arc::new(litellm_core_utils::settings::ProcessEnvironment), - ) -} - #[cfg(test)] mod tests { + use std::time::Duration; + + use futures_util::future::BoxFuture; use litellm_core_utils::call_arguments::{CallArguments, compose_body, parse_options}; - use serde_json::json; + use litellm_host::event::WireRequest; + use litellm_llms::{ + base_llm::ocr::{ + error::Error, + handler::{CallHooks, OcrClient}, + transformation::{BaseOcrConfig, OcrResponseFormat}, + }, + cohere::ocr::transformation::CohereParseConfig, + mistral::ocr::transformation::MistralOcrConfig, + vertex_ai::ocr::transformation::VertexAiOcrConfig, + }; + use serde_json::{Value, json}; + + use super::*; + use crate::ocr::{ + document::prepare_document, + types::LiteLLMOcrRequest, + wire::{OcrWireRequest, decode_request}, + }; + + /// Stands in for a host with no hooks registered. + struct NoHooks; + + impl CallHooks for NoHooks { + fn before_send(&self, wire: WireRequest) -> BoxFuture<'_, Result> { + Box::pin(async move { Ok(wire) }) + } + + fn response_received<'a>(&'a self, _body: &'a [u8]) -> BoxFuture<'a, Result<(), Error>> { + Box::pin(async { Ok(()) }) + } + } + + fn client() -> OcrClient { + OcrClient::for_test(reqwest::Client::new(), reqwest::Client::new()) + } + + fn request(model: &str, base: &str, document: Value, options: Value) -> LiteLLMOcrRequest { + decode_request(OcrWireRequest { + model: model.into(), + document, + api_key: Some(litellm_auth::SecretValue::new("test-key")), + api_base: Some(base.into()), + custom_llm_provider: None, + extra_headers: None, + optional_params: options.as_object().unwrap().clone(), + input_sources: Default::default(), + timeout_seconds: Some(2.0), + }) + .unwrap() + } + + fn prepared(request: LiteLLMOcrRequest) -> PreparedOcrRequest { + prepare_request( + request.map_document(prepare_document).unwrap(), + true, + &client(), + std::sync::Arc::new(litellm_core_utils::settings::ProcessEnvironment), + ) + } + + fn image(url: &str) -> Value { + json!({"type": "image_url", "image_url": url}) + } + + #[tokio::test] + async fn cohere_body_keeps_native_document_fields_and_untyped_overrides() { + let request = request( + "cohere/parse", + "https://example.com", + image("https://example.com/original.png"), + json!({ + "output_format": "markdown", "timeout": 30, + "extra_body": { + "output_format": {"future": true}, + "document": {"type": "image_url", "image_url": "https://example.com/a.png", + "provider_options": {"nested": [false, 0, null]}} + } + }), + ); + + let http = CohereParseConfig + .prepare_request(&prepared(request), &client(), &NoHooks) + .await + .unwrap(); + + let body: Value = serde_json::from_slice(http.body()).unwrap(); + assert_eq!( + body, + json!({ + "model": "parse", "output_format": {"future": true}, + "document": {"type": "image_url", "image_url": "https://example.com/a.png", + "provider_options": {"nested": [false, 0, null]}} + }) + ); + } + + #[tokio::test] + async fn explicit_null_options_use_defaults_before_http() { + let request = request( + "cohere/parse", + "https://example.com", + image("https://example.com/a.png"), + json!({"output_format": null, "req_format": null}), + ); + assert_eq!( + request.response_format().unwrap(), + OcrResponseFormat::Litellm + ); + + let http = CohereParseConfig + .prepare_request(&prepared(request), &client(), &NoHooks) + .await + .unwrap(); + + let body: Value = serde_json::from_slice(http.body()).unwrap(); + assert_eq!(body["output_format"], "markdown"); + assert!(body.get("req_format").is_none()); + } + + #[tokio::test] + async fn direct_and_vertex_mistral_build_the_same_request_and_share_normalization() { + let options = json!({ + "pages": [0, 2], + "include_image_base64": true, + "vertex_project": "project-1", + "vertex_location": "us-central1", + "unknown": "preserved" + }); + let document = + json!({"type": "document_url", "document_url": "data:application/pdf;base64,YWJj"}); + let direct = prepared(request( + "mistral/mistral-ocr-maas", + "https://mistral.test", + document.clone(), + options.clone(), + )); + let vertex = prepared(request( + "vertex_ai/mistral-ocr-maas", + "https://vertex.test", + document, + options, + )); + + let direct_http = MistralOcrConfig + .prepare_request(&direct, &client(), &NoHooks) + .await + .unwrap(); + let vertex_http = VertexAiOcrConfig + .prepare_request(&vertex, &client(), &NoHooks) + .await + .unwrap(); + + assert_eq!(direct_http.url(), "https://mistral.test/v1/ocr"); + assert_eq!( + vertex_http.url(), + "https://vertex.test/v1/projects/project-1/locations/us-central1/publishers/mistralai/models/mistral-ocr-maas:rawPredict" + ); + for http in [&direct_http, &vertex_http] { + assert_eq!(http.header("authorization").unwrap(), "Bearer test-key"); + assert_eq!(http.header("content-type").unwrap(), "application/json"); + assert_eq!(http.timeout(), Some(Duration::from_secs(2))); + let body: Value = serde_json::from_slice(http.body()).unwrap(); + assert_eq!( + body, + json!({ + "model": "mistral-ocr-maas", + "document": {"type": "document_url", "document_url": "data:application/pdf;base64,YWJj"}, + "pages": [0, 2], + "include_image_base64": true, + "unknown": "preserved" + }) + ); + } + let payload = serde_json::to_vec( + &json!({"pages": [{"index": 0, "markdown": "hello"}], "extra": "preserved"}), + ) + .unwrap(); + let direct_response = MistralOcrConfig + .transform_ocr_response(&direct.model, &payload, OcrResponseFormat::Litellm) + .unwrap() + .into_json(); + let vertex_response = VertexAiOcrConfig + .transform_ocr_response(&vertex.model, &payload, OcrResponseFormat::Litellm) + .unwrap() + .into_json(); + assert_eq!(direct_response, vertex_response); + assert_eq!(direct_response["model"], "mistral-ocr-maas"); + assert_eq!(direct_response["object"], "ocr"); + assert_eq!(direct_response["extra"], "preserved"); + } #[derive(serde::Deserialize)] struct KnownParams { diff --git a/litellm-rust/crates/core/src/ocr/route.rs b/litellm-rust/crates/core/src/ocr/route.rs index e3e57bd2d77..adb704a15d1 100644 --- a/litellm-rust/crates/core/src/ocr/route.rs +++ b/litellm-rust/crates/core/src/ocr/route.rs @@ -161,3524 +161,3 @@ impl litellm_host::host::Host for LocalOcrHost { Ok(()) } } - -#[cfg(test)] -mod aws_textract_tests { - use std::{collections::BTreeMap, time::SystemTime}; - - use litellm_auth_aws::{Credentials, aws_signature_headers, sign_post}; - use litellm_llms::base_llm::ocr::error::Error; - use serde_json::{Value, json}; - use time::{PrimitiveDateTime, format_description}; - - use crate::ocr::{ - route::LocalOcrHost, - test_support::{ - MockResponse, header, mock_server, perform_ocr_with, request_body, - wire_request_with_document, - }, - types::LiteLLMOcrRequest, - }; - - const ACCESS_KEY_ID: &str = "AKIDEXAMPLE"; - const SECRET_ACCESS_KEY: &str = "wJalrXUtnFEMI/K7MDENG+bPxRfiCYEXAMPLEKEY"; - - fn textract_request(base: &str) -> LiteLLMOcrRequest { - textract_request_for("aws_textract/detect-document-text", base) - } - - fn textract_request_for(model: &str, base: &str) -> LiteLLMOcrRequest { - wire_request_with_document( - model, - &format!("{base}/"), - json!({"type": "image_url", "image_url": "data:image/png;base64,b3JpZ2luYWw="}), - json!({ - "aws_access_key_id": ACCESS_KEY_ID, - "aws_secret_access_key": SECRET_ACCESS_KEY, - "aws_region_name": "eu-west-1" - }), - ) - } - - fn textract_response() -> MockResponse { - MockResponse::json(json!({ - "DocumentMetadata": {"Pages": 1}, - "Blocks": [{"BlockType": "PAGE"}, {"BlockType": "LINE", "Text": "Invoice 12345"}] - })) - } - - /// Recomputes SigV4 over the bytes the server received, at the time the client claimed. - fn expected_authorization(url: &str, raw_request: &str) -> String { - let format = - format_description::parse_borrowed::<2>("[year][month][day]T[hour][minute][second]Z") - .unwrap(); - let signed_at: SystemTime = - PrimitiveDateTime::parse(header(raw_request, "x-amz-date").unwrap(), &format) - .unwrap() - .assume_utc() - .into(); - let headers: BTreeMap = ["content-type", "x-amz-target"] - .into_iter() - .map(|name| { - ( - name.to_string(), - header(raw_request, name).unwrap().to_string(), - ) - }) - .collect(); - let body = raw_request.split_once("\r\n\r\n").unwrap().1; - sign_post( - url, - body.as_bytes(), - &aws_signature_headers(&headers), - "eu-west-1", - "textract", - &Credentials::new(ACCESS_KEY_ID, SECRET_ACCESS_KEY, None, None, "test"), - signed_at, - ) - .unwrap()["Authorization"] - .clone() - } - - #[tokio::test] - async fn the_request_is_signed_for_textract_and_lines_become_the_page() { - let (base, seen, server) = mock_server(vec![textract_response()]).await; - - let response = perform_ocr_with(LocalOcrHost::new(textract_request(&base))) - .await - .unwrap(); - server.await.unwrap(); - - let raw = seen.lock().unwrap()[0].clone(); - assert_eq!( - header(&raw, "x-amz-target"), - Some("Textract.DetectDocumentText") - ); - assert_eq!( - header(&raw, "content-type"), - Some("application/x-amz-json-1.1") - ); - assert_eq!( - request_body(&raw), - json!({"Document": {"Bytes": "b3JpZ2luYWw="}}) - ); - assert_eq!( - header(&raw, "authorization"), - Some(expected_authorization(&format!("{base}/"), &raw).as_str()) - ); - assert_eq!(response.pages[0].markdown, "Invoice 12345"); - assert_eq!(response.usage_info.unwrap().pages_processed, Some(1)); - } - - #[tokio::test] - async fn a_body_rewritten_by_before_send_is_what_gets_signed_and_sent() { - let (base, seen, server) = mock_server(vec![textract_response()]).await; - let host = LocalOcrHost::new(textract_request(&base)).with_before_send(|mut wire, _| { - assert!( - !wire - .headers - .iter() - .any(|(name, _)| name.eq_ignore_ascii_case("authorization")), - "the hook ran after signing" - ); - wire.body["Document"]["Bytes"] = Value::from("cmVkYWN0ZWQ="); - Ok(wire) - }); - - perform_ocr_with(host).await.unwrap(); - server.await.unwrap(); - - let raw = seen.lock().unwrap()[0].clone(); - assert_eq!( - request_body(&raw), - json!({"Document": {"Bytes": "cmVkYWN0ZWQ="}}) - ); - assert_eq!( - header(&raw, "authorization"), - Some(expected_authorization(&format!("{base}/"), &raw).as_str()) - ); - } - - #[tokio::test] - async fn a_multi_page_rejection_reaches_the_caller_with_the_single_page_limit() { - let (base, _, server) = mock_server(vec![MockResponse { - status: 400, - headers: vec![], - body: json!({ - "__type": "UnsupportedDocumentException", - "Message": "Request has unsupported document format" - }), - }]) - .await; - - let error = perform_ocr_with(LocalOcrHost::new(textract_request(&base))) - .await - .unwrap_err(); - server.await.unwrap(); - - let Error::Provider { status, body, .. } = error else { - panic!("expected a provider error, got {error:?}"); - }; - assert_eq!(status, 400); - assert!( - body.contains("multi-page documents are not supported"), - "{body}" - ); - } - - #[tokio::test] - async fn analyze_document_asks_for_layout_and_tables_and_returns_markdown() { - let (base, seen, server) = mock_server(vec![MockResponse::json(json!({ - "DocumentMetadata": {"Pages": 1}, - "Blocks": [ - {"Id": "l1", "BlockType": "LINE", "Text": "Quarterly Report"}, - {"Id": "t", "BlockType": "LAYOUT_TITLE", - "Relationships": [{"Type": "CHILD", "Ids": ["l1"]}]} - ] - }))]) - .await; - let request = textract_request_for("aws_textract/analyze-document", &base); - - let response = perform_ocr_with(LocalOcrHost::new(request)).await.unwrap(); - server.await.unwrap(); - - let raw = seen.lock().unwrap()[0].clone(); - assert_eq!( - header(&raw, "x-amz-target"), - Some("Textract.AnalyzeDocument") - ); - assert_eq!( - request_body(&raw)["FeatureTypes"], - json!(["LAYOUT", "TABLES"]) - ); - assert_eq!( - header(&raw, "authorization"), - Some(expected_authorization(&format!("{base}/"), &raw).as_str()) - ); - assert_eq!(response.pages[0].markdown, "# Quarterly Report"); - } -} - -#[cfg(test)] -mod azure_ai_tests { - use litellm_llms::base_llm::ocr::error::Error; - use serde_json::{Value, json}; - - use crate::ocr::route::LocalOcrHost; - use crate::ocr::test_support::{ - MockResponse, mock_server, perform_ocr, perform_ocr_with, wire_request, - }; - - #[tokio::test] - async fn facade_executes_azure_mistral_with_prepared_auth() { - let (base, seen, server) = mock_server(vec![MockResponse::json(json!({ - "pages":[{"index":0,"markdown":"hello"}], - "usage_info":{"pages_processed":1} - }))]) - .await; - let mut request = wire_request( - "azure_ai/model", - &base, - json!({"include_image_base64":true}), - ); - request.credentials.api_key = None; - request.transport.extra_headers = vec![( - "Authorization".into(), - "Bearer python-prepared-token".into(), - )]; - - let result = perform_ocr(request).await.unwrap(); - server.await.unwrap(); - assert_eq!(result.pages[0].markdown, "hello"); - let requests = seen.lock().unwrap(); - assert_eq!(requests.len(), 1); - assert!(requests[0].starts_with("POST /providers/mistral/azure/ocr ")); - assert!( - requests[0] - .to_ascii_lowercase() - .contains("authorization: bearer python-prepared-token\r\n") - ); - let body: Value = - serde_json::from_str(requests[0].split_once("\r\n\r\n").unwrap().1).unwrap(); - assert_eq!( - body, - json!({ - "model":"model", - "document":{"type":"document_url","document_url":"data:application/pdf;base64,YWJj"}, - "include_image_base64":true - }) - ); - } - - #[tokio::test] - async fn facade_acquires_supplied_entra_token_for_final_request() { - let (base, seen, server) = mock_server(vec![MockResponse::json(json!({"pages":[]}))]).await; - let mut request = wire_request( - "azure_ai/model", - &base, - json!({"azure_ad_token":"rust-owned-token"}), - ); - request.credentials.api_key = None; - - perform_ocr(request).await.unwrap(); - server.await.unwrap(); - - let requests = seen.lock().unwrap(); - assert_eq!(requests.len(), 1); - assert!( - requests[0] - .to_ascii_lowercase() - .contains("authorization: bearer rust-owned-token\r\n") - ); - } - - #[tokio::test] - async fn rejects_non_inline_body_after_guardrails() { - let request = wire_request("azure_ai/model", "http://127.0.0.1:1", json!({})); - let host = LocalOcrHost::new(request).with_before_send(|mut wire, _| { - wire.body["document"] = json!({ - "type":"document_url", - "document_url":"https://example.com/not-inline.pdf" - }); - Ok(wire) - }); - let error = perform_ocr_with(host).await.unwrap_err(); - assert!(error.to_string().contains("data URI")); - } - - mod transformation { - use std::sync::{ - Arc, - atomic::{AtomicUsize, Ordering}, - }; - - use litellm_auth::{ - ResolvedCredential, SecretValue, TokenFuture, TokenProvider, TokenProviderHandle, - }; - use rstest::rstest; - use serde_json::json; - - use super::*; - use crate::ocr::{ - test_support::{MockResponse, header, mock_server, perform_ocr}, - types::LiteLLMOcrRequest, - wire::decode_request, - }; - - #[derive(Debug)] - struct CountingToken { - token: fn(usize) -> String, - calls: AtomicUsize, - } - - impl CountingToken { - fn new(token: fn(usize) -> String) -> Arc { - Arc::new(Self { - token, - calls: AtomicUsize::new(0), - }) - } - - fn calls(&self) -> usize { - self.calls.load(Ordering::SeqCst) - } - } - - impl TokenProvider for CountingToken { - fn acquire(&self) -> TokenFuture<'_> { - let call = self.calls.fetch_add(1, Ordering::SeqCst) + 1; - let token = SecretValue::new((self.token)(call)); - Box::pin(async move { - Ok(ResolvedCredential::AccessToken { - token, - expires_on: None, - }) - }) - } - } - - fn numbered_token(call: usize) -> String { - format!("callback-{call}") - } - - fn azure_request( - provider: &Arc, - api_base: Option<&str>, - api_key: Option<&str>, - extra_headers: Value, - optional_params: Value, - ) -> LiteLLMOcrRequest { - let wire = serde_json::from_value(json!({ - "model": "azure_ai/mistral-ocr-latest", - "document": {"type":"document_url","document_url":"data:application/pdf;base64,YWJj"}, - "api_key": api_key, - "api_base": api_base, - "custom_llm_provider": null, - "extra_headers": extra_headers, - "optional_params": optional_params, - "timeout_seconds": 2.0 - })) - .unwrap(); - LiteLLMOcrRequest { - azure_ad_token_provider: Some(TokenProviderHandle::new(provider.clone())), - ..decode_request(wire).unwrap() - } - } - - fn ocr_page() -> MockResponse { - MockResponse::json(json!({"pages":[{"index":0,"markdown":"hello"}]})) - } - - #[tokio::test] - async fn token_provider_result_is_the_bearer_and_is_acquired_for_each_request() { - let provider = CountingToken::new(numbered_token); - let (base, seen, server) = mock_server(vec![ocr_page(), ocr_page()]).await; - - for _ in 0..2 { - perform_ocr(azure_request( - &provider, - Some(&base), - None, - Value::Null, - json!({}), - )) - .await - .unwrap(); - } - server.await.unwrap(); - - assert_eq!(provider.calls(), 2); - let requests = seen.lock().unwrap(); - assert_eq!( - requests - .iter() - .map(|request| header(request, "authorization")) - .collect::>(), - [Some("Bearer callback-1"), Some("Bearer callback-2")] - ); - } - - #[rstest] - #[case::api_key_skips_provider(Some("resource-key"), Value::Null, json!({}), "Bearer resource-key", 0)] - #[case::provider_beats_static_token( - None, - Value::Null, - json!({"azure_ad_token":"static-token"}), - "Bearer callback-1", - 1 - )] - #[case::header_wins_on_the_wire_but_provider_still_runs( - None, - json!({"Authorization":"Bearer override"}), - json!({}), - "Bearer override", - 1 - )] - #[tokio::test] - async fn credential_precedence( - #[case] api_key: Option<&str>, - #[case] extra_headers: Value, - #[case] optional_params: Value, - #[case] expected_authorization: &str, - #[case] expected_calls: usize, - ) { - let provider = CountingToken::new(numbered_token); - let (base, seen, server) = mock_server(vec![ocr_page()]).await; - - perform_ocr(azure_request( - &provider, - Some(&base), - api_key, - extra_headers, - optional_params, - )) - .await - .unwrap(); - server.await.unwrap(); - - assert_eq!(provider.calls(), expected_calls); - let requests = seen.lock().unwrap(); - assert_eq!(requests.len(), 1); - assert_eq!( - header(&requests[0], "authorization"), - Some(expected_authorization) - ); - } - - #[rstest] - #[case::missing_api_base( - false, - json!({}), - numbered_token, - |error: &Error| matches!(error, Error::Auth(litellm_auth::Error::MissingApiBase { - provider: "Azure AI", - environment_variable: "AZURE_AI_API_BASE", - })), - 0 - )] - #[case::unsupported_oidc_reference( - true, - json!({"azure_ad_token":"oidc/assertion","client_id":"client","tenant_id":"tenant"}), - numbered_token, - |error: &Error| matches!(error, Error::Auth(litellm_auth::Error::UnsupportedOidcReference)), - 0 - )] - #[case::empty_provider_token_ignores_static_token( - true, - json!({"azure_ad_token":"static-token"}), - |_| String::new(), - |error: &Error| matches!(error, Error::MissingAzureAiCredentials), - 1 - )] - #[tokio::test] - async fn credential_failures_send_no_provider_request( - #[case] with_api_base: bool, - #[case] optional_params: Value, - #[case] token: fn(usize) -> String, - #[case] expected: fn(&Error) -> bool, - #[case] expected_calls: usize, - ) { - let provider = CountingToken::new(token); - let (base, seen, server) = mock_server(vec![ocr_page()]).await; - - let error = perform_ocr(azure_request( - &provider, - with_api_base.then_some(base.as_str()), - None, - Value::Null, - optional_params, - )) - .await - .unwrap_err(); - server.abort(); - - assert!(expected(&error), "unexpected error: {error:?}"); - assert_eq!(provider.calls(), expected_calls); - assert!(seen.lock().unwrap().is_empty()); - } - } -} - -#[cfg(test)] -mod azure_document_intelligence_tests { - use litellm_host::event::{CallEvent, MachineEvent}; - use litellm_llms::base_llm::ocr::{error::Error, settings::OcrSettings}; - use rstest::rstest; - use serde_json::{Value, json}; - - use crate::ocr::route::LocalOcrHost; - use crate::ocr::{ - test_support::{ - MockResponse, mock_server, ocr_client, perform_ocr, perform_ocr_with, wire_request, - }, - wire::{OcrWireRequest, decode_request}, - }; - - fn query_value(url: &str, key: &str) -> Option { - url::Url::parse(url) - .unwrap() - .query_pairs() - .find_map(|(name, value)| (name == key).then(|| value.into_owned())) - } - - #[tokio::test] - async fn facade_maps_pages_features_and_url_document() { - let (base, seen, server) = mock_server(vec![MockResponse::json(json!({ - "status":"succeeded", - "analyzeResult":{"pages":[]} - }))]) - .await; - let mut request = wire_request( - "azure_ai/doc-intelligence/prebuilt-read", - &base, - json!({"pages":[2,0,0,1],"features":["keyValuePairs","languages"]}), - ); - request.document = serde_json::from_value::< - litellm_llms::base_llm::ocr::transformation::OcrDocument, - >(json!({ - "type":"document_url", - "document_url":"https://example.com/document.pdf" - })) - .unwrap() - .into(); - - perform_ocr(request).await.unwrap(); - server.await.unwrap(); - let request = &seen.lock().unwrap()[0]; - let target = request.split_whitespace().nth(1).unwrap(); - let url = format!("{base}{target}"); - assert_eq!(query_value(&url, "pages").as_deref(), Some("1,2,3")); - assert_eq!( - query_value(&url, "features").as_deref(), - Some("keyValuePairs,languages") - ); - let body: Value = serde_json::from_str(request.split_once("\r\n\r\n").unwrap().1).unwrap(); - assert_eq!( - body, - json!({"urlSource":"https://example.com/document.pdf"}) - ); - } - - #[rstest] - #[case(json!({"pages":[true]}), Error::Pages("expected only integers or only strings".into()))] - #[case(json!({"pages":[1,"2"]}), Error::Pages("expected only integers or only strings".into()))] - #[case(json!({"pages":[-1]}), Error::Pages("negative page index".into()))] - #[case(json!({"pages":"1&&features=bad"}), Error::Pages("invalid native page range".into()))] - #[case(json!({"features":"languages&pages=1"}), Error::Features)] - #[case(json!({"req_format":"azure"}), Error::RequestFormat)] - #[tokio::test] - async fn rejects_invalid_pages_features_and_format( - #[case] options: Value, - #[case] expected: Error, - ) { - let (base, seen, server) = mock_server(vec![MockResponse::json(json!({}))]).await; - let result = decode_request(OcrWireRequest { - model: "azure_ai/doc-intelligence/prebuilt-read".into(), - document: json!({"type":"document_url","document_url":"https://example.com/a.pdf"}), - api_key: Some(litellm_auth::SecretValue::new("key")), - api_base: Some(base), - custom_llm_provider: None, - extra_headers: None, - optional_params: options.as_object().unwrap().clone(), - input_sources: Default::default(), - timeout_seconds: Some(2.0), - }); - let result = match result { - Ok(request) => perform_ocr(request).await, - Err(error) => Err(error), - }; - server.abort(); - let _ = server.await; - assert!( - seen.lock().unwrap().is_empty(), - "sent invalid options: {options}" - ); - let error = result.unwrap_err(); - assert_eq!( - std::mem::discriminant(&error), - std::mem::discriminant(&expected) - ); - assert_eq!(error.http_status_code(), Some(400)); - assert_eq!(error.to_string(), expected.to_string()); - } - - #[rstest] - #[case(json!({}))] - #[case(json!({"req_format":"litellm"}))] - #[tokio::test] - async fn missing_native_fields_keep_page_text_without_retaining_raw_response( - #[case] options: Value, - ) { - let operation = json!({ - "status":"succeeded", - "analyzeResult":{"pages":[{"pageNumber":1,"lines":[{"content":"hello"}]}]} - }); - let (base, seen, server) = mock_server(vec![MockResponse::json(operation)]).await; - let response = perform_ocr(wire_request( - "azure_ai/doc-intelligence/prebuilt-read", - &base, - options, - )) - .await - .unwrap(); - server.await.unwrap(); - - assert_eq!(response.pages.len(), 1); - assert_eq!(response.pages[0].index, 0); - assert_eq!(response.pages[0].markdown, "hello"); - assert_eq!(response.provider_native_response, None); - let serialized = response.into_json(); - assert_eq!(serialized.get("content"), Some(&Value::Null)); - assert_eq!(serialized.get("tables"), Some(&Value::Null)); - assert_eq!(serialized.get("keyValuePairs"), Some(&Value::Null)); - let requests = seen.lock().unwrap(); - assert_eq!(requests.len(), 1); - let target = requests[0].split_whitespace().nth(1).unwrap(); - let url = format!("{base}{target}"); - for field in ["pages", "features", "req_format"] { - assert_eq!(query_value(&url, field), None); - } - let body: Value = - serde_json::from_str(requests[0].split_once("\r\n\r\n").unwrap().1).unwrap(); - assert_eq!(body, json!({"base64Source":"YWJj"})); - } - - #[tokio::test] - async fn inline_document_decodes_to_base64_source() { - let (base, seen, server) = mock_server(vec![MockResponse::json(json!({ - "status":"succeeded" - }))]) - .await; - let request = wire_request("azure_ai/doc-intelligence/prebuilt-read", &base, json!({})); - - perform_ocr(request).await.unwrap(); - server.await.unwrap(); - let request = &seen.lock().unwrap()[0]; - let body: Value = serde_json::from_str(request.split_once("\r\n\r\n").unwrap().1).unwrap(); - assert_eq!(body, json!({"base64Source":"YWJj"})); - } - - #[tokio::test] - async fn immediate_response_normalizes_pages_and_preserves_native() { - let operation = json!({ - "status":"succeeded", - "operationExtension":42, - "analyzeResult":{ - "content":"A\n\nB", - "tables":[{"cells":[]}], - "keyValuePairs":[{"key":{"content":"A"}}], - "pages":[{ - "pageNumber":"2", - "width":"8.5", - "height":11, - "unit":"inch", - "lines":[{"content":"A"},{"content":null},{"content":"B"}] - }] - } - }); - let (base, _, server) = mock_server(vec![MockResponse::json(operation.clone())]).await; - let result = perform_ocr(wire_request( - "azure_ai/doc-intelligence/prebuilt-read", - &base, - json!({"req_format":"native"}), - )) - .await - .unwrap(); - server.await.unwrap(); - - assert_eq!(result.pages[0].index, 1); - assert_eq!(result.pages[0].markdown, "A\n\nB"); - assert_eq!( - serde_json::to_value(&result.pages[0].dimensions).unwrap(), - json!({"width":816,"height":1056,"dpi":96}) - ); - assert_eq!(result.usage_info.as_ref().unwrap().pages_processed, Some(1)); - let serialized = result.clone().into_json(); - assert_eq!(serialized["content"], "A\n\nB"); - assert_eq!(serialized["tables"], json!([{"cells":[]}])); - assert_eq!( - serialized["keyValuePairs"], - json!([{"key":{"content":"A"}}]) - ); - assert!(serialized.get("key_value_pairs").is_none()); - assert_eq!( - result.provider_native_response.map(Value::Object), - Some(operation) - ); - } - - #[tokio::test] - async fn client_settings_choose_the_api_version_and_the_inch_to_pixel_dpi() { - let (base, seen, server) = mock_server(vec![MockResponse::json(json!({ - "status":"succeeded", - "analyzeResult":{"pages":[{"pageNumber":1,"width":8.5,"height":11,"unit":"inch"}]} - }))]) - .await; - let client = ocr_client().with_settings(OcrSettings { - document_intelligence_api_version: "2099-01-01".into(), - document_intelligence_dpi: 72, - ..OcrSettings::default() - }); - - let result = crate::ocr::client::perform( - &client, - wire_request("azure_ai/doc-intelligence/prebuilt-read", &base, json!({})), - ) - .await - .unwrap(); - server.await.unwrap(); - - let target = seen.lock().unwrap()[0] - .split_whitespace() - .nth(1) - .unwrap() - .to_string(); - assert_eq!( - query_value(&format!("{base}{target}"), "api-version").as_deref(), - Some("2099-01-01") - ); - assert_eq!( - serde_json::to_value(&result.pages[0].dimensions).unwrap(), - json!({"width":612,"height":792,"dpi":72}) - ); - } - - #[tokio::test] - async fn accepted_response_polls_to_success_with_only_credentials() { - let operation = json!({"status":"succeeded","analyzeResult":{"pages":[]}}); - let (base, seen, server) = mock_server(vec![ - MockResponse { - status: 202, - headers: vec![("Operation-Location", "{base}/operation".into())], - body: json!({}), - }, - MockResponse { - status: 200, - headers: vec![("Retry-After", "0".into())], - body: json!({"status":"running"}), - }, - MockResponse::json(operation.clone()), - ]) - .await; - let mut request = wire_request( - "azure_ai/doc-intelligence/prebuilt-read", - &base, - json!({"req_format":"native"}), - ); - request - .transport - .extra_headers - .push(("X-Trace".into(), "initial-only".into())); - - let result = perform_ocr(request).await.unwrap(); - server.await.unwrap(); - assert_eq!( - result.provider_native_response.map(Value::Object), - Some(operation) - ); - let requests = seen.lock().unwrap(); - assert_eq!(requests.len(), 3); - assert!(requests[0].to_ascii_lowercase().contains("x-trace:")); - for poll in &requests[1..] { - assert!(!poll.to_ascii_lowercase().contains("x-trace:")); - assert!( - poll.to_ascii_lowercase() - .contains("ocp-apim-subscription-key: test-key") - ); - } - } - - #[tokio::test] - async fn accepted_response_emits_response_received_before_polling() { - let (base, seen, server) = mock_server(vec![ - MockResponse { - status: 202, - headers: vec![("Operation-Location", "{base}/operation".into())], - body: json!({"submitted": true}), - }, - MockResponse::json(json!({"status":"succeeded"})), - ]) - .await; - let request_count = seen.clone(); - let host = LocalOcrHost::new(wire_request( - "azure_ai/doc-intelligence/prebuilt-read", - &base, - json!({}), - )) - .with_observer(move |event| { - let CallEvent::Machine(MachineEvent::ResponseReceived { raw }) = event else { - return; - }; - match request_count.lock().unwrap().len() { - 1 => assert_eq!(raw.body, r#"{"submitted":true}"#), - 2 => assert!(raw.body.contains("succeeded")), - count => panic!("unexpected callback after {count} requests"), - } - }); - - perform_ocr_with(host).await.unwrap(); - server.await.unwrap(); - assert_eq!(seen.lock().unwrap().len(), 2); - } - - #[tokio::test] - async fn polling_forwards_bearer_credentials() { - let (base, seen, server) = mock_server(vec![ - MockResponse { - status: 202, - headers: vec![("Operation-Location", "{base}/operation".into())], - body: json!({}), - }, - MockResponse::json(json!({"status":"succeeded"})), - ]) - .await; - let mut request = wire_request("azure_ai/doc-intelligence/prebuilt-read", &base, json!({})); - request.credentials.api_key = None; - request.transport.extra_headers = vec![("Authorization".into(), "Bearer token".into())]; - - perform_ocr(request).await.unwrap(); - server.await.unwrap(); - let requests = seen.lock().unwrap(); - assert!( - requests[1] - .to_ascii_lowercase() - .contains("authorization: bearer token") - ); - } - - #[tokio::test] - async fn polling_does_not_follow_redirects() { - let (base, seen, server) = mock_server(vec![ - MockResponse { - status: 202, - headers: vec![("Operation-Location", "{base}/operation".into())], - body: json!({}), - }, - MockResponse { - status: 302, - headers: vec![("Location", "{base}/redirected".into())], - body: json!({}), - }, - MockResponse::json(json!({"status":"succeeded"})), - ]) - .await; - - let error = perform_ocr(wire_request( - "azure_ai/doc-intelligence/prebuilt-read", - &base, - json!({}), - )) - .await - .unwrap_err(); - - assert!(error.to_string().contains("status 302"), "{error}"); - assert_eq!(seen.lock().unwrap().len(), 2); - server.abort(); - } - - #[tokio::test] - async fn polling_rejects_terminal_failure() { - let (base, _, server) = mock_server(vec![ - MockResponse { - status: 202, - headers: vec![("Operation-Location", "{base}/operation".into())], - body: json!({}), - }, - MockResponse::json(json!({"status":"failed"})), - ]) - .await; - - let error = perform_ocr(wire_request( - "azure_ai/doc-intelligence/prebuilt-read", - &base, - json!({}), - )) - .await - .unwrap_err(); - server.await.unwrap(); - assert!(error.to_string().contains("status failed")); - } - - #[tokio::test] - async fn malformed_provider_pages_report_response_paths() { - for (analysis, path) in [ - (json!({"pages":null}), "pages"), - (json!({"pages":[null]}), "pages[0]"), - (json!({"pages":[{"lines":null}]}), "lines"), - (json!({"pages":[{"width":"bad"}]}), "width"), - ] { - let (base, _, server) = mock_server(vec![MockResponse::json(json!({ - "status":"succeeded", - "analyzeResult":analysis - }))]) - .await; - let error = perform_ocr(wire_request( - "azure_ai/doc-intelligence/prebuilt-read", - &base, - json!({}), - )) - .await - .unwrap_err(); - server.await.unwrap(); - assert!(error.to_string().contains(path), "{error}"); - } - } - - #[tokio::test] - async fn rejects_missing_invalid_and_cross_origin_operation_locations() { - for headers in [ - Vec::new(), - vec![("Operation-Location", "/relative".into())], - vec![("Operation-Location", "http://example.com/operation".into())], - vec![( - "Operation-Location", - "http://user:password@127.0.0.1/operation".into(), - )], - ] { - let (base, _, server) = mock_server(vec![MockResponse { - status: 202, - headers, - body: json!({}), - }]) - .await; - let error = perform_ocr(wire_request( - "azure_ai/doc-intelligence/prebuilt-read", - &base, - json!({}), - )) - .await - .unwrap_err(); - server.await.unwrap(); - assert!(error.to_string().contains("operation-location")); - } - } - - #[tokio::test] - async fn polling_deadline_bounds_retry_delay() { - let (base, _, server) = mock_server(vec![ - MockResponse { - status: 202, - headers: vec![("Operation-Location", "{base}/operation".into())], - body: json!({}), - }, - MockResponse { - status: 200, - headers: vec![("Retry-After", "9999".into())], - body: json!({"status":"notStarted"}), - }, - ]) - .await; - let request = wire_request("azure_ai/doc-intelligence/prebuilt-read", &base, json!({})); - let client = ocr_client().with_settings(OcrSettings { - poll_timeout: std::time::Duration::from_millis(100), - ..OcrSettings::default() - }); - - let error = tokio::time::timeout( - std::time::Duration::from_secs(1), - crate::ocr::client::perform(&client, request), - ) - .await - .unwrap() - .unwrap_err(); - server.await.unwrap(); - assert!(error.to_string().contains("timed out")); - } - - #[tokio::test] - async fn model_id_is_encoded_and_dot_segments_are_rejected() { - let (base, seen, server) = mock_server(vec![MockResponse::json(json!({ - "status":"succeeded" - }))]) - .await; - perform_ocr(wire_request( - "azure_ai/doc-intelligence/a ?#é", - &base, - json!({}), - )) - .await - .unwrap(); - server.await.unwrap(); - assert!(seen.lock().unwrap()[0].contains("a%20%3F%23%C3%A9:analyze")); - - for model in [ - "azure_ai/doc-intelligence/.", - "azure_ai/doc-intelligence/..", - ] { - let error = perform_ocr(wire_request(model, "http://127.0.0.1:1", json!({}))) - .await - .unwrap_err(); - assert!(error.to_string().contains("dot segment")); - } - } - - mod transformation { - use std::sync::{Arc, Mutex}; - - use litellm_host::event::{CallEvent, MachineEvent}; - use litellm_llms::base_llm::ocr::transformation::OcrDocument; - use serde_json::{Value, json}; - - use super::*; - use crate::ocr::{ - route::LocalOcrHost, - test_support::{ - MockResponse, mock_server, perform_ocr, perform_ocr_with, wire_request, - }, - }; - - #[tokio::test] - async fn facade_maps_pages_features_and_url_document() { - let (base, seen, server) = mock_server(vec![MockResponse::json(json!({ - "status":"succeeded", - "analyzeResult":{"pages":[]} - }))]) - .await; - let mut request = wire_request( - "azure_ai/doc-intelligence/prebuilt-read", - &base, - json!({"pages":[2,0,0,1],"features":["keyValuePairs","languages"], "future_option": {"nested":null}, "extra_body":{"provider_option":false}}), - ); - request.document = serde_json::from_value::(json!({ - "type":"document_url", - "document_url":"https://example.com/document.pdf" - })) - .unwrap() - .into(); - - perform_ocr(request).await.unwrap(); - server.await.unwrap(); - let request = &seen.lock().unwrap()[0]; - let target = request.split_whitespace().nth(1).unwrap(); - let url = format!("{base}{target}"); - assert_eq!(query_value(&url, "pages").as_deref(), Some("1,2,3")); - assert_eq!( - query_value(&url, "features").as_deref(), - Some("keyValuePairs,languages") - ); - let body: Value = - serde_json::from_str(request.split_once("\r\n\r\n").unwrap().1).unwrap(); - assert_eq!( - body, - json!({"urlSource":"https://example.com/document.pdf", "future_option":{"nested":null}, "provider_option":false}) - ); - } - - #[tokio::test] - async fn rejects_invalid_pages_features_and_format() { - for options in [ - json!({"pages":[true]}), - json!({"pages":[1,"2"]}), - json!({"pages":[-1]}), - json!({"pages":"1&&features=bad"}), - json!({"features":"languages&pages=1"}), - json!({"req_format":"azure"}), - ] { - let request = wire_request( - "azure_ai/doc-intelligence/prebuilt-read", - "http://127.0.0.1:1", - options.clone(), - ); - let rejected = perform_ocr(request).await.is_err(); - assert!(rejected, "accepted {options}"); - } - } - - #[tokio::test] - async fn immediate_response_normalizes_pages_and_preserves_native() { - let operation = json!({ - "status":"succeeded", - "operationExtension":42, - "analyzeResult":{ - "content":"A\n\nB", - "tables":[{"cells":[]}], - "keyValuePairs":[{"key":{"content":"A"}}], - "pages":[{ - "pageNumber":"2", - "width":"8.5", - "height":11, - "unit":"inch", - "lines":[{"content":"A"},{"content":null},{"content":"B"}] - }] - } - }); - let (base, _, server) = mock_server(vec![MockResponse::json(operation.clone())]).await; - let result = perform_ocr(wire_request( - "azure_ai/doc-intelligence/prebuilt-read", - &base, - json!({"req_format":"native"}), - )) - .await - .unwrap(); - server.await.unwrap(); - - assert_eq!(result.pages[0].index, 1); - assert_eq!(result.pages[0].markdown, "A\n\nB"); - assert_eq!( - serde_json::to_value(&result.pages[0].dimensions).unwrap(), - json!({"width":816,"height":1056,"dpi":96}) - ); - assert_eq!(result.usage_info.as_ref().unwrap().pages_processed, Some(1)); - let serialized = result.clone().into_json(); - assert_eq!(serialized["content"], "A\n\nB"); - assert_eq!(serialized["tables"], json!([{"cells":[]}])); - assert_eq!( - serialized["keyValuePairs"], - json!([{"key":{"content":"A"}}]) - ); - assert!(serialized.get("key_value_pairs").is_none()); - assert_eq!( - result.provider_native_response.as_ref(), - operation.as_object() - ); - } - - #[tokio::test] - async fn accepted_response_polls_to_success_with_only_credentials() { - let operation = json!({"status":"succeeded","analyzeResult":{"pages":[]}}); - let (base, seen, server) = mock_server(vec![ - MockResponse { - status: 202, - headers: vec![("Operation-Location", "{base}/operation".into())], - body: json!({}), - }, - MockResponse { - status: 200, - headers: vec![("Retry-After", "0".into())], - body: json!({"status":"running"}), - }, - MockResponse::json(operation.clone()), - ]) - .await; - let mut request = wire_request( - "azure_ai/doc-intelligence/prebuilt-read", - &base, - json!({"req_format":"native"}), - ); - request - .transport - .extra_headers - .push(("X-Trace".into(), "initial-only".into())); - - let result = perform_ocr(request).await.unwrap(); - server.await.unwrap(); - assert_eq!( - result.provider_native_response.as_ref(), - operation.as_object() - ); - let requests = seen.lock().unwrap(); - assert_eq!(requests.len(), 3); - assert!(requests[0].to_ascii_lowercase().contains("x-trace:")); - for poll in &requests[1..] { - assert!(!poll.to_ascii_lowercase().contains("x-trace:")); - assert!( - poll.to_ascii_lowercase() - .contains("ocp-apim-subscription-key: test-key") - ); - } - } - - #[tokio::test] - async fn accepted_response_emits_response_received_for_submission_and_completed_poll() { - let (base, seen, server) = mock_server(vec![ - MockResponse { - status: 202, - headers: vec![("Operation-Location", "{base}/operation".into())], - body: json!({"submitted": true}), - }, - MockResponse::json(json!({"status":"succeeded"})), - ]) - .await; - let responses_received = Arc::new(Mutex::new(Vec::new())); - let request_count = seen.clone(); - let observed = responses_received.clone(); - let host = LocalOcrHost::new(wire_request( - "azure_ai/doc-intelligence/prebuilt-read", - &base, - json!({}), - )) - .with_observer(move |event| { - if let CallEvent::Machine(MachineEvent::ResponseReceived { raw }) = event { - observed - .lock() - .unwrap() - .push((request_count.lock().unwrap().len(), raw.body.clone())); - } - }); - - perform_ocr_with(host).await.unwrap(); - server.await.unwrap(); - assert_eq!(seen.lock().unwrap().len(), 2); - assert_eq!( - *responses_received.lock().unwrap(), - [ - (1, r#"{"submitted":true}"#.to_string()), - (2, r#"{"status":"succeeded"}"#.to_string()), - ] - ); - } - } -} - -#[cfg(test)] -mod cohere_tests { - mod transformation { - use litellm_llms::{ - base_llm::ocr::{ - error::Error, - transformation::{BaseOcrConfig, OcrDocument, OcrResponseFormat}, - }, - cohere::ocr::transformation::*, - }; - use rstest::rstest; - use serde_json::{Value, json}; - - #[tokio::test] - async fn composed_body_preserves_native_document_fields_and_untyped_overrides() { - let request = crate::ocr::test_support::wire_request( - "cohere/parse", - "https://example.com", - json!({ - "output_format":"markdown", "timeout":30, - "extra_body":{ - "output_format": {"future":true}, - "document":{"type":"image_url","image_url":"https://example.com/a.png", - "provider_options":{"nested":[false,0,null]}} - } - }), - ); - let request = request.with_document( - serde_json::from_value(json!({ - "type":"image_url","image_url":"https://example.com/original.png" - })) - .unwrap(), - ); - let request = crate::ocr::prepare::prepare_request_for_test(request); - let http = CohereParseConfig - .prepare_request( - &request, - &crate::ocr::test_support::ocr_client(), - &crate::ocr::test_support::NoHooks, - ) - .await - .unwrap(); - let body: Value = serde_json::from_slice(http.body()).unwrap(); - assert_eq!( - body, - json!({ - "model":"parse", "output_format":{"future":true}, - "document":{"type":"image_url","image_url":"https://example.com/a.png", - "provider_options":{"nested":[false,0,null]}} - }) - ); - } - - #[tokio::test] - async fn explicit_null_options_use_defaults_before_http() { - let request = crate::ocr::test_support::wire_request( - "cohere/parse", - "https://example.com", - json!({"output_format":null,"req_format":null}), - ); - let request = request.with_document( - serde_json::from_value( - json!({"type":"image_url","image_url":"https://example.com/a.png"}), - ) - .unwrap(), - ); - assert_eq!( - request.response_format().unwrap(), - OcrResponseFormat::Litellm - ); - let request = crate::ocr::prepare::prepare_request_for_test(request); - let http = CohereParseConfig - .prepare_request( - &request, - &crate::ocr::test_support::ocr_client(), - &crate::ocr::test_support::NoHooks, - ) - .await - .unwrap(); - let body: Value = serde_json::from_slice(http.body()).unwrap(); - assert_eq!(body["output_format"], "markdown"); - assert!(body.get("req_format").is_none()); - } - - #[rstest] - #[case::cohere("cohere/parse-v5.0", "POST /v2/parse ")] - #[case::azure_ai("azure_ai/Cohere-parse-v5.0", "POST /providers/cohere/v2/parse ")] - #[tokio::test] - async fn route_sends_image_to_its_parse_endpoint_with_the_bearer_key( - #[case] model: &str, - #[case] request_line: &str, - ) { - use crate::ocr::test_support::{MockResponse, header, mock_server, perform_ocr}; - - let (base, seen, server) = - mock_server(vec![MockResponse::json(json!({"pages":[]}))]).await; - let request = crate::ocr::test_support::wire_request(model, &base, json!({})) - .with_document( - serde_json::from_value::( - json!({"type":"image_url","image_url":"data:image/png;base64,YWJj"}), - ) - .unwrap() - .into(), - ); - - perform_ocr(request).await.unwrap(); - server.await.unwrap(); - - let requests = seen.lock().unwrap(); - assert_eq!(requests.len(), 1); - assert!(requests[0].starts_with(request_line), "{}", requests[0]); - assert_eq!( - header(&requests[0], "authorization"), - Some("Bearer test-key") - ); - } - - #[rstest] - #[tokio::test] - async fn route_rejects_non_image_document_without_a_request( - #[values("cohere/parse-v5.0", "azure_ai/Cohere-parse-v5.0")] model: &str, - ) { - use crate::ocr::test_support::{MockResponse, mock_server, perform_ocr}; - - let (base, seen, server) = - mock_server(vec![MockResponse::json(json!({"pages":[]}))]).await; - - let error = perform_ocr(crate::ocr::test_support::wire_request( - model, - &base, - json!({}), - )) - .await - .unwrap_err(); - server.abort(); - - assert!(matches!(error, Error::CohereImageOnly), "{error:?}"); - assert!(seen.lock().unwrap().is_empty()); - } - } -} - -#[cfg(test)] -mod deepseek_tests { - use litellm_llms::{ - base_llm::ocr::transformation::{BaseOcrConfig, OcrDocument}, - vertex_ai::ocr::deepseek_transformation::{ - DeepSeekOcrParams, DeepSeekOcrResponse, VertexAIDeepSeekOCRConfig, - normalize_response as transform_ocr_response, - }, - }; - use rstest::rstest; - use serde_json::{Value, json}; - - fn document() -> OcrDocument { - serde_json::from_value(json!({"type":"image_url","image_url":"gs://bucket/a.png"})).unwrap() - } - - #[rstest] - #[case("stream", json!(true))] - #[case("temperature", json!(0.1))] - #[case("max_tokens", json!(1024))] - #[case("top_p", json!(0.9))] - #[case("n", json!(2))] - #[case("stop", json!("done"))] - #[case("stop", json!(["done", "stop"]))] - fn request_mapping_matches_python(#[case] name: &str, #[case] value: Value) { - let params: DeepSeekOcrParams = - serde_json::from_value(json!({name: value.clone(), "ignored": true})).unwrap(); - let result = serde_json::to_value( - VertexAIDeepSeekOCRConfig - .transform_ocr_request("deepseek-ai/deepseek-ocr-maas", document(), ¶ms, &[]) - .unwrap(), - ) - .unwrap(); - assert_eq!(result["model"], "deepseek-ai/deepseek-ocr-maas"); - assert_eq!( - result["messages"][0]["content"][0], - json!({"type":"image_url","image_url":"gs://bucket/a.png"}) - ); - assert_eq!(result[name], value); - assert!(result.get("ignored").is_none()); - } - - #[rstest] - #[case(json!({"type":"image_url","image_url":"data:image/png;base64,AA=="}))] - #[case(json!({"type":"document_url","document_url":"data:application/pdf;base64,AA=="}))] - fn request_maps_both_document_types_to_image_content(#[case] document: Value) { - let source = document - .get("image_url") - .or_else(|| document.get("document_url")) - .unwrap() - .clone(); - let request = VertexAIDeepSeekOCRConfig - .transform_ocr_request( - "deepseek-ai/deepseek-ocr-maas", - serde_json::from_value(document).unwrap(), - &DeepSeekOcrParams::default(), - &[], - ) - .unwrap(); - let result = serde_json::to_value(request).unwrap(); - assert_eq!( - result["messages"][0]["content"][0], - json!({"type":"image_url","image_url":source}) - ); - } - - #[rstest] - #[case(json!("# hello"), "# hello")] - #[case(json!("{broken"), "{broken")] - #[case(json!(" {\"pages\":[]} "), " {\"pages\":[]} ")] - #[case(json!({"pages":[]}), "")] - #[case(json!("[]"), "[]")] - #[case(json!("{\"pages\":[{\"markdown\":\"json text\"}]}"), "json text")] - #[case(json!({"pages":[{"markdown":"object"}]}), "object")] - fn response_codec_handles_text_json_and_objects( - #[case] content: Value, - #[case] expected: &str, - ) { - let structured = content - .as_object() - .is_some_and(|object| object.contains_key("pages")) - || content - .as_str() - .is_some_and(|text| text.contains("\"pages\"")); - let response: DeepSeekOcrResponse = serde_json::from_value( - json!({"choices":[{"message":{"content":content}}],"usage":{"prompt_tokens":1}}), - ) - .unwrap(); - let result = transform_ocr_response("model", response) - .unwrap() - .into_json(); - assert_eq!(result["pages"][0]["markdown"], expected); - assert_eq!(result["pages"][0]["index"], 0); - if structured { - assert!(result["usage_info"].is_null()); - } else { - assert_eq!(result["usage_info"]["prompt_tokens"], 1); - } - } - - #[test] - fn structured_result_maps_pages_usage_model_and_annotation() { - let response: DeepSeekOcrResponse = serde_json::from_value(json!({ - "choices":[{"message":{"content":{ - "pages":[{"index":2,"markdown":"page","images":[{"id":"one"}],"dimensions":{"width":10}}], - "model":"provider-model", - "usage_info":{"pages_processed":1}, - "document_annotation":{"language":"en"}, - "future":"kept" - }}}] - })) - .unwrap(); - let result = transform_ocr_response("requested", response) - .unwrap() - .into_json(); - assert_eq!(result["pages"][0]["index"], 2); - assert_eq!(result["pages"][0]["images"][0]["id"], "one"); - assert_eq!(result["model"], "provider-model"); - assert_eq!(result["usage_info"]["pages_processed"], 1); - assert_eq!(result["document_annotation"]["language"], "en"); - assert_eq!(result["future"], "kept"); - } - - #[test] - fn response_codec_rejects_missing_empty_and_malformed_content() { - for value in [ - json!({"choices":[{"message":{"content":{}}}]}), - json!({"choices":[]}), - json!({"choices":[{"message":{"content":""}}]}), - json!({"choices":[{"message":{"content":"{\"pages\":[{\"markdown\":42}]}"}}]}), - json!({"choices":[{"message":{"content":{"pages":[{"markdown":42}]}}}]}), - ] { - let result = serde_json::from_value::(value) - .map_err(|_| ()) - .and_then(|response| transform_ocr_response("model", response).map_err(|_| ())); - assert!(result.is_err()); - } - } -} - -#[cfg(test)] -mod reducto_tests { - use litellm_host::event::{CallEvent, MachineEvent, WireRequest}; - use litellm_llms::base_llm::ocr::{error::Error, transformation::OcrDocument}; - use rstest::rstest; - use serde_json::{Value, json}; - - use crate::ocr::route::LocalOcrHost; - use crate::ocr::test_support::{ - MockResponse, mock_server, perform_ocr, perform_ocr_with, wire_request, - }; - - fn request_body(request: &str) -> Value { - serde_json::from_str(request.split_once("\r\n\r\n").unwrap().1).unwrap() - } - - #[rstest] - #[case( - "reducto/parse-v3", - json!({ - "formatting":{"table_output_format":"html"}, - "retrieval":{"chunk_mode":"section"}, - "settings":{"ocr_system":"standard"}, - "future_ocr_option":true, - "extra_body":{"provider_option":"value"} - }), - "reducto://already.pdf", - json!({ - "input":"reducto://already.pdf", - "formatting":{"table_output_format":"html"}, - "retrieval":{"chunk_mode":"section"}, - "settings":{"ocr_system":"standard"}, - "future_ocr_option":true, - "provider_option":"value" - }) - )] - #[case( - "reducto/parse-legacy", - json!({ - "enhance":{"agentic":[{"type":"table"}]}, - "future_ocr_option":true, - "extra_body":{"provider_option":"value"} - }), - "reducto://legacy.pdf", - json!({ - "document_url":"reducto://legacy.pdf", - "options":{"enhance":{"agentic":[{"type":"table"}]}}, - "future_ocr_option":true, - "provider_option":"value" - }) - )] - #[tokio::test] - async fn request_mapping_matches_python( - #[case] model: &str, - #[case] options: Value, - #[case] source: &str, - #[case] expected: Value, - ) { - let (base, seen, server) = mock_server(vec![MockResponse::json(json!({ - "result":{"chunks":[]} - }))]) - .await; - let request = - crate::ocr::test_support::with_source(wire_request(model, &base, options), source); - - perform_ocr(request).await.unwrap(); - server.await.unwrap(); - let requests = seen.lock().unwrap(); - assert_eq!(requests.len(), 1); - assert!(requests[0].starts_with("POST /parse ")); - assert_eq!(request_body(&requests[0]), expected); - } - - #[rstest] - #[case("parse-v3")] - #[case("parse-legacy")] - #[tokio::test] - async fn data_uri_upload_preserves_multipart_headers( - #[case] model: &str, - #[values("application/pdf", "image/png")] mime_type: &str, - ) { - let (base, seen, server) = mock_server(vec![ - MockResponse::json(json!({"file_id":"reducto://uploaded.pdf"})), - MockResponse::json(json!({"result":{"chunks":[{"content":"hello"}]}})), - ]) - .await; - let document = if mime_type.starts_with("image/") { - json!({"type":"image_url","image_url":format!("data:{mime_type};base64,YWJj")}) - } else { - json!({"type":"document_url","document_url":format!("data:{mime_type};base64,YWJj")}) - }; - let mut request = crate::ocr::types::LiteLLMOcrRequest { - document: serde_json::from_value::(document) - .unwrap() - .into(), - ..wire_request(&format!("reducto/{model}"), &base, json!({})) - }; - request.transport.extra_headers = vec![ - ("Content-Type".into(), "application/json".into()), - ("X-Trace".into(), "upload-test".into()), - ]; - - let response = perform_ocr(request).await.unwrap(); - server.await.unwrap(); - assert_eq!(response.pages[0].markdown, "hello"); - let requests = seen.lock().unwrap(); - assert_eq!(requests.len(), 2); - assert!(requests[0].starts_with("POST /upload ")); - assert!( - requests[0] - .to_ascii_lowercase() - .contains("content-type: multipart/form-data; boundary=") - ); - assert!(requests[0].contains("x-trace: upload-test")); - let multipart = requests[0].split_once("\r\n\r\n").unwrap().1; - assert!(multipart.contains(&format!("Content-Type: {mime_type}\r\n"))); - assert!(multipart.contains("\r\n\r\nabc\r\n--")); - assert!(requests[1].starts_with("POST /parse ")); - let source_field = if model == "parse-legacy" { - "document_url" - } else { - "input" - }; - assert_eq!( - request_body(&requests[1]), - json!({source_field:"reducto://uploaded.pdf"}) - ); - for request in requests.iter() { - assert!( - request - .to_ascii_lowercase() - .contains("authorization: bearer test-key\r\n") - ); - } - } - - #[tokio::test] - async fn response_received_stays_after_reducto_upload_and_parse() { - let (base, seen, server) = mock_server(vec![ - MockResponse::json(json!({"file_id":"reducto://uploaded.pdf"})), - MockResponse::json(json!({"result":{"chunks":[]}})), - ]) - .await; - let request_count = seen.clone(); - let host = LocalOcrHost::new(wire_request("reducto/parse-v3", &base, json!({}))) - .with_observer(move |event| { - if let CallEvent::Machine(MachineEvent::ResponseReceived { raw }) = event { - assert_eq!(request_count.lock().unwrap().len(), 2); - assert_eq!(raw.body, r#"{"result":{"chunks":[]}}"#); - } - }); - - perform_ocr_with(host).await.unwrap(); - server.await.unwrap(); - assert_eq!(seen.lock().unwrap().len(), 2); - } - - #[rstest] - #[case(json!({"file_id":""}))] - #[case(json!({}))] - #[case(json!({"file_id":null}))] - #[tokio::test] - async fn invalid_upload_ids_stop_before_parse(#[case] response: Value) { - let (base, seen, server) = mock_server(vec![MockResponse::json(response)]).await; - let error = perform_ocr(wire_request("reducto/parse-v3", &base, json!({}))) - .await - .unwrap_err(); - server.await.unwrap(); - assert!(error.to_string().contains("file_id")); - assert_eq!(seen.lock().unwrap().len(), 1); - } - - #[tokio::test] - async fn upload_failure_stops_before_parse() { - let (base, seen, server) = mock_server(vec![MockResponse { - status: 503, - headers: vec![], - body: json!({"error":"unavailable"}), - }]) - .await; - assert!( - perform_ocr(wire_request("reducto/parse-v3", &base, json!({}))) - .await - .is_err() - ); - server.await.unwrap(); - assert_eq!(seen.lock().unwrap().len(), 1); - } - - #[rstest] - #[case("https://example.com/a.pdf", Error::ReductoSource)] - #[case("reducto://", Error::RequestField { path: "document file id".into() })] - #[case("data:application/pdf;base64", Error::InvalidDataUri)] - #[case("data:application/pdf;base64,INVALID!", Error::InvalidDataUri)] - #[tokio::test] - async fn rejects_invalid_document_sources_before_network( - #[case] source: &str, - #[case] expected: Error, - ) { - let (base, seen, server) = mock_server(vec![MockResponse::json(json!({}))]).await; - let request = crate::ocr::test_support::with_source( - wire_request("reducto/parse-v3", &base, json!({})), - source, - ); - let result = perform_ocr(request).await; - server.abort(); - let _ = server.await; - assert!( - seen.lock().unwrap().is_empty(), - "sent invalid source: {source}" - ); - let error = result.unwrap_err(); - assert_eq!( - std::mem::discriminant(&error), - std::mem::discriminant(&expected) - ); - assert_eq!(error.http_status_code(), Some(400)); - assert_eq!(error.to_string(), expected.to_string()); - } - - #[test] - fn response_normalization_groups_blocks_and_distinguishes_null_result() { - use litellm_llms::reducto::ocr::transformation::{ - ReductoResponse, normalize_response as transform_ocr_response, - }; - - let raw = json!({"usage":{"num_pages":"2","credits":"3"},"result":{"type":"full","chunks":[ - {"blocks":[{ - "type":"Table", - "content":"B", - "bbox":{"left":0.1,"top":0.2,"width":0.8,"height":0.3,"page":2,"original_page":4}, - "confidence":"high", - "granular_confidence":{"parse_confidence":0.95,"extract_confidence":null}, - "image_url":null - }]}, - {"blocks":[{"content":"A","bbox":{"page":1},"type":"Text"},{"content":"C","bbox":{"page":1}}]} - ]}}); - let response: ReductoResponse = serde_json::from_value(raw).unwrap(); - let normalized = transform_ocr_response("parse-v3", response) - .unwrap() - .into_json(); - assert_eq!(normalized["pages"][0]["markdown"], "A\n\nC"); - assert_eq!(normalized["pages"][1]["markdown"], "B"); - assert_eq!(normalized["pages"][1]["blocks"][0]["type"], "Table"); - assert_eq!( - normalized["pages"][1]["blocks"][0]["bbox"], - json!({"left":0.1,"top":0.2,"width":0.8,"height":0.3,"page":2,"original_page":4}) - ); - assert_eq!(normalized["pages"][1]["blocks"][0]["confidence"], "high"); - assert_eq!( - normalized["pages"][1]["blocks"][0]["granular_confidence"]["parse_confidence"], - 0.95 - ); - assert!(normalized["pages"][1]["blocks"][0]["image_url"].is_null()); - assert_eq!(normalized["usage_info"]["pages_processed"], 2); - assert_eq!(normalized["usage_info"]["credits"], 3.0); - - let missing: ReductoResponse = - serde_json::from_value(json!({"chunks":[{"content":"text"}]})).unwrap(); - let missing = transform_ocr_response("parse-v3", missing).unwrap(); - assert_eq!(missing.pages[0].markdown, "text"); - let null: ReductoResponse = serde_json::from_value( - json!({"result":null,"chunks":[{"content":"ignored"}],"usage":null}), - ) - .unwrap(); - let null = transform_ocr_response("parse-v3", null).unwrap(); - assert!(null.pages.is_empty()); - } - - #[tokio::test] - async fn facade_omits_native_response_by_default_and_preserves_auth_priority() { - let raw = json!({"job_id":"job-1","result":{"chunks":[]}}); - let (base, seen, server) = mock_server(vec![MockResponse::json(raw)]).await; - let mut request = crate::ocr::test_support::with_source( - wire_request("reducto/parse-v3", &base, json!({})), - "reducto://ready.pdf", - ); - request.transport.extra_headers = vec![("authorization".into(), "Bearer existing".into())]; - - let response = perform_ocr(request).await.unwrap(); - server.await.unwrap(); - assert_eq!(response.provider_native_response, None); - assert!( - seen.lock().unwrap()[0] - .to_ascii_lowercase() - .contains("authorization: bearer existing") - ); - } - - #[tokio::test] - async fn native_format_retains_the_provider_response() { - let raw = json!({ - "result":{"chunks":[{"content":"native OCR response"}]}, - "usage":{"num_pages":1} - }); - let (base, _, server) = mock_server(vec![MockResponse::json(raw.clone())]).await; - let request = crate::ocr::test_support::with_source( - wire_request("reducto/parse-v3", &base, json!({"req_format":"native"})), - "reducto://ready.pdf", - ); - - let response = perform_ocr(request).await.unwrap(); - server.await.unwrap(); - - assert_eq!(response.pages[0].markdown, "native OCR response"); - assert_eq!(response.provider_native_response.as_ref(), raw.as_object()); - } - - #[tokio::test] - async fn unknown_model_reaches_parse_and_keeps_its_name() { - let (base, seen, server) = mock_server(vec![MockResponse::json(json!({ - "result":{"chunks":[{"content":"future model response"}]} - }))]) - .await; - let request = crate::ocr::test_support::with_source( - wire_request("reducto/future-parse-model", &base, json!({})), - "reducto://ready.pdf", - ); - - let response = perform_ocr(request).await.unwrap(); - server.await.unwrap(); - - assert_eq!(response.model, "future-parse-model"); - assert_eq!(response.pages[0].markdown, "future model response"); - let requests = seen.lock().unwrap(); - assert!(requests[0].starts_with("POST /parse ")); - assert_eq!( - request_body(&requests[0]), - json!({"input":"reducto://ready.pdf"}) - ); - } - - #[tokio::test] - async fn guardrail_rewrites_document_before_upload() { - let (base, seen, server) = - mock_server(vec![MockResponse::json(json!({"result":{"chunks":[]}}))]).await; - let host = LocalOcrHost::new(wire_request("reducto/parse-v3", &base, json!({}))) - .with_before_send(|wire, _| { - assert_eq!( - wire.body["document_url"], - "data:application/pdf;base64,YWJj" - ); - Ok(WireRequest { - body: json!({"type":"document_url","document_url":"reducto://guarded.pdf"}), - ..wire - }) - }); - - perform_ocr_with(host).await.unwrap(); - server.await.unwrap(); - let requests = seen.lock().unwrap(); - assert_eq!(requests.len(), 1); - assert!(requests[0].starts_with("POST /parse ")); - assert!(requests[0].contains("reducto://guarded.pdf")); - } - - mod transformation { - use litellm_host::event::{CallEvent, MachineEvent, WireRequest}; - use litellm_llms::{ - base_llm::ocr::transformation::{BaseOcrConfig, OcrConnection, OcrRequestContext}, - reducto::ocr::transformation::*, - }; - use rstest::rstest; - - use super::*; - use crate::ocr::{ - route::LocalOcrHost, - test_support::{ - MockResponse, mock_server, perform_ocr, perform_ocr_with, wire_request, - }, - }; - - #[tokio::test] - async fn v3_options_preserve_explicit_null() { - let overrides = - serde_json::from_value(json!({"formatting":null,"settings":{},"unknown":true})) - .unwrap(); - let params = ReductoParseV3Config - .map_ocr_params(&overrides, "parse-v3") - .unwrap(); - let client = crate::ocr::test_support::ocr_client(); - let connection = OcrConnection::default(); - let document = serde_json::from_value( - json!({"type":"document_url","document_url":"reducto://ready.pdf"}), - ) - .unwrap(); - let body = ReductoParseV3Config - .async_transform_ocr_request( - "parse-v3", - document, - ¶ms, - &[], - OcrRequestContext { - client: &client, - connection: &connection, - }, - ) - .await - .unwrap(); - assert_eq!( - serde_json::to_value(body).unwrap(), - json!({ - "input":"reducto://ready.pdf", "formatting":null, "settings":{} - }) - ); - let absent = ReductoParseV3Config - .map_ocr_params( - &litellm_core_utils::call_arguments::CallArguments::default(), - "parse-v3", - ) - .unwrap(); - assert_eq!(serde_json::to_value(absent).unwrap(), json!({})); - } - - #[rstest] - #[case( - "reducto/parse-v3", - json!({ - "formatting":{"table_output_format":"html"}, - "retrieval":{"chunk_mode":"section"}, - "settings":{"ocr_system":"standard"}, - "future_ocr_option":true, - "extra_body":{"provider_option":"value"} - }), - "reducto://already.pdf", - json!({ - "input":"reducto://already.pdf", - "formatting":{"table_output_format":"html"}, - "retrieval":{"chunk_mode":"section"}, - "settings":{"ocr_system":"standard"}, - "future_ocr_option":true, - "provider_option":"value" - }) - )] - #[case( - "reducto/parse-legacy", - json!({ - "enhance":{"agentic":[{"type":"table"}]}, - "future_ocr_option":true, - "extra_body":{"provider_option":"value"} - }), - "reducto://legacy.pdf", - json!({ - "document_url":"reducto://legacy.pdf", - "options":{"enhance":{"agentic":[{"type":"table"}]}}, - "future_ocr_option":true, - "provider_option":"value" - }) - )] - #[tokio::test] - async fn request_mapping_matches_python( - #[case] model: &str, - #[case] options: Value, - #[case] source: &str, - #[case] expected: Value, - ) { - let (base, seen, server) = mock_server(vec![MockResponse::json(json!({ - "result":{"chunks":[]} - }))]) - .await; - let request = - crate::ocr::test_support::with_source(wire_request(model, &base, options), source); - - perform_ocr(request).await.unwrap(); - server.await.unwrap(); - let requests = seen.lock().unwrap(); - assert_eq!(requests.len(), 1); - assert!(requests[0].starts_with("POST /parse ")); - assert_eq!(request_body(&requests[0]), expected); - } - - #[rstest] - #[case("parse-v3")] - #[case("parse-legacy")] - #[tokio::test] - async fn data_uri_upload_preserves_multipart_headers(#[case] model: &str) { - let (base, seen, server) = mock_server(vec![ - MockResponse::json(json!({"file_id":"reducto://uploaded.pdf"})), - MockResponse::json(json!({"result":{"chunks":[{"content":"hello"}]}})), - ]) - .await; - let mut request = wire_request(&format!("reducto/{model}"), &base, json!({})); - request.transport.extra_headers = vec![ - ("Content-Type".into(), "application/json".into()), - ("X-Trace".into(), "upload-test".into()), - ]; - - let response = perform_ocr(request).await.unwrap(); - server.await.unwrap(); - assert_eq!(response.pages[0].markdown, "hello"); - let requests = seen.lock().unwrap(); - assert_eq!(requests.len(), 2); - assert!(requests[0].starts_with("POST /upload ")); - assert!( - requests[0] - .to_ascii_lowercase() - .contains("content-type: multipart/form-data; boundary=") - ); - assert!(requests[0].contains("x-trace: upload-test")); - assert!(requests[0].contains("application/pdf")); - assert!(requests[0].contains("abc")); - assert!(requests[1].starts_with("POST /parse ")); - } - - #[tokio::test] - async fn response_received_stays_after_reducto_upload_and_parse() { - let (base, seen, server) = mock_server(vec![ - MockResponse::json(json!({"file_id":"reducto://uploaded.pdf"})), - MockResponse::json(json!({"result":{"chunks":[]}})), - ]) - .await; - let request_count = seen.clone(); - let host = LocalOcrHost::new(wire_request("reducto/parse-v3", &base, json!({}))) - .with_observer(move |event| { - if let CallEvent::Machine(MachineEvent::ResponseReceived { raw }) = event { - assert_eq!(request_count.lock().unwrap().len(), 2); - assert_eq!(raw.body, r#"{"result":{"chunks":[]}}"#); - } - }); - - perform_ocr_with(host).await.unwrap(); - server.await.unwrap(); - assert_eq!(seen.lock().unwrap().len(), 2); - } - - #[rstest] - #[case("https://example.com/a.pdf")] - #[case("reducto://")] - #[case("data:application/pdf;base64")] - #[case("data:application/pdf;base64,INVALID!")] - #[tokio::test] - async fn rejects_invalid_document_sources_before_network(#[case] source: &str) { - let request = crate::ocr::test_support::with_source( - wire_request("reducto/parse-v3", "http://127.0.0.1:1", json!({})), - source, - ); - assert!(perform_ocr(request).await.is_err()); - } - - #[tokio::test] - async fn facade_omits_native_response_by_default_and_preserves_auth_priority() { - let raw = json!({"job_id":"job-1","result":{"chunks":[]}}); - let (base, seen, server) = mock_server(vec![MockResponse::json(raw)]).await; - let mut request = crate::ocr::test_support::with_source( - wire_request("reducto/parse-v3", &base, json!({})), - "reducto://ready.pdf", - ); - request.transport.extra_headers = - vec![("authorization".into(), "Bearer existing".into())]; - - let response = perform_ocr(request).await.unwrap(); - server.await.unwrap(); - assert_eq!(response.provider_native_response, None); - assert!( - seen.lock().unwrap()[0] - .to_ascii_lowercase() - .contains("authorization: bearer existing") - ); - } - - #[rstest] - #[case("reducto/parse-v3")] - #[case("reducto/parse-legacy")] - #[tokio::test] - async fn guardrail_headers_reach_upload_and_parse(#[case] model: &str) { - let (base, seen, server) = mock_server(vec![ - MockResponse::json(json!({"file_id":"reducto://uploaded.pdf"})), - MockResponse::json(json!({"result":{"chunks":[]}})), - ]) - .await; - let mut request = wire_request(model, &base, json!({})); - request.transport.extra_headers = - vec![("authorization".into(), "Bearer original".into())]; - let host = LocalOcrHost::new(request).with_before_send(|wire, _| { - Ok(WireRequest { - headers: vec![("authorization".into(), "Bearer guarded".into())], - ..wire - }) - }); - - perform_ocr_with(host).await.unwrap(); - server.await.unwrap(); - let requests = seen.lock().unwrap(); - assert_eq!(requests.len(), 2); - assert!(requests[0].starts_with("POST /upload ")); - assert!(requests[1].starts_with("POST /parse ")); - for request in requests.iter() { - assert!(request.contains("authorization: Bearer guarded")); - assert!(!request.contains("Bearer original")); - } - } - } -} - -#[cfg(test)] -mod vertex_ai_tests { - use litellm_auth::InputSource; - use litellm_llms::base_llm::ocr::{settings::OcrSettings, transformation::OcrResponseFormat}; - use serde_json::{Value, json}; - - use crate::ocr::test_support::{ - MockResponse, mock_server, ocr_client, perform_ocr, wire_request, - }; - - fn request_body(request: &str) -> Value { - serde_json::from_str(request.split_once("\r\n\r\n").unwrap().1).unwrap() - } - - #[tokio::test] - async fn facade_executes_vertex_mistral_with_resolved_project_and_location() { - let (base, seen, server) = mock_server(vec![MockResponse::json(json!({ - "pages":[{"index":0,"markdown":"hello"}], - "usage_info":{"pages_processed":1} - }))]) - .await; - let request = wire_request( - "vertex_ai/mistral-ocr-maas", - &base, - json!({ - "vertex_project":"project-1", - "vertex_location":"europe-west4", - "extract_footer":true - }), - ); - - let response = perform_ocr(request).await.unwrap(); - server.await.unwrap(); - assert_eq!(response.pages[0].markdown, "hello"); - let requests = seen.lock().unwrap(); - assert_eq!(requests.len(), 1); - assert!(requests[0].starts_with( - "POST /v1/projects/project-1/locations/europe-west4/publishers/mistralai/models/mistral-ocr-maas:rawPredict " - )); - assert!( - requests[0] - .to_ascii_lowercase() - .contains("authorization: bearer test-key") - ); - assert_eq!( - request_body(&requests[0]), - json!({ - "model":"mistral-ocr-maas", - "document":{"type":"document_url","document_url":"data:application/pdf;base64,YWJj"}, - "extract_footer":true - }) - ); - } - - #[tokio::test] - async fn configured_project_and_location_apply_when_the_call_sets_neither() { - let (base, seen, server) = mock_server(vec![MockResponse::json(json!({"pages":[]}))]).await; - let client = ocr_client().with_settings(OcrSettings { - vertex_project: Some("configured-project".into()), - vertex_location: Some("europe-west4".into()), - ..OcrSettings::default() - }); - - crate::ocr::client::perform( - &client, - wire_request("vertex_ai/mistral-ocr-maas", &base, json!({})), - ) - .await - .unwrap(); - server.await.unwrap(); - assert!(seen.lock().unwrap()[0].starts_with( - "POST /v1/projects/configured-project/locations/europe-west4/publishers/mistralai/models/mistral-ocr-maas:rawPredict " - )); - } - - #[tokio::test] - async fn supplied_authorization_is_forwarded_without_a_static_token() { - let (base, seen, server) = mock_server(vec![MockResponse::json(json!({"pages":[]}))]).await; - let mut request = wire_request( - "vertex_ai/model", - &base, - json!({"vertex_project":"project-1"}), - ); - request.credentials.api_key = None; - request.transport.extra_headers = vec![("authorization".into(), "Bearer supplied".into())]; - - perform_ocr(request).await.unwrap(); - server.await.unwrap(); - assert!( - seen.lock().unwrap()[0] - .to_ascii_lowercase() - .contains("authorization: bearer supplied") - ); - } - - #[tokio::test] - async fn invalid_credentials_fail_before_provider_http() { - let request = wire_request( - "vertex_ai/model", - "http://127.0.0.1:1", - json!({"vertex_credentials": true}), - ); - let error = perform_ocr(request).await.unwrap_err(); - assert!(error.to_string().contains("vertex_credentials")); - } - - #[tokio::test] - async fn request_controlled_api_base_is_rejected_before_vertex_auth() { - let mut request = wire_request( - "vertex_ai/mistral-ocr-maas", - "https://caller.example", - json!({"vertex_project":"project-1"}), - ); - request.credentials.api_base = Some(litellm_auth::Sourced::new( - "https://caller.example".into(), - InputSource::Request, - )); - - let error = perform_ocr(request).await.unwrap_err(); - assert!( - error - .to_string() - .contains("request-controlled Vertex AI endpoint") - ); - } - - #[tokio::test] - async fn adapters_build_complete_requests_and_share_mistral_normalization() { - use std::time::Duration; - - use litellm_llms::{ - base_llm::ocr::transformation::BaseOcrConfig, - mistral::ocr::transformation::MistralOcrConfig, - vertex_ai::ocr::transformation::VertexAiOcrConfig, - }; - - use crate::ocr::test_support::ocr_client; - - let client = ocr_client(); - let options = json!({ - "pages": [0, 2], - "include_image_base64": true, - "vertex_project": "project-1", - "vertex_location": "us-central1", - "unknown": "ignored" - }); - let direct = wire_request( - "mistral/mistral-ocr-maas", - "https://mistral.test", - options.clone(), - ); - let vertex = wire_request("vertex_ai/mistral-ocr-maas", "https://vertex.test", options); - let direct = crate::ocr::prepare::prepare_request_for_test( - crate::ocr::test_support::resolved_request(direct), - ); - let vertex = crate::ocr::prepare::prepare_request_for_test( - crate::ocr::test_support::resolved_request(vertex), - ); - let direct_http = MistralOcrConfig - .prepare_request(&direct, &client, &crate::ocr::test_support::NoHooks) - .await - .unwrap(); - let vertex_http = VertexAiOcrConfig - .prepare_request(&vertex, &client, &crate::ocr::test_support::NoHooks) - .await - .unwrap(); - assert_eq!(direct_http.url(), "https://mistral.test/v1/ocr"); - assert_eq!( - vertex_http.url(), - "https://vertex.test/v1/projects/project-1/locations/us-central1/publishers/mistralai/models/mistral-ocr-maas:rawPredict" - ); - for http in [&direct_http, &vertex_http] { - assert_eq!(http.header("authorization").unwrap(), "Bearer test-key"); - assert_eq!(http.header("content-type").unwrap(), "application/json"); - assert_eq!(http.timeout(), Some(Duration::from_secs(2))); - let body: Value = serde_json::from_slice(http.body()).unwrap(); - assert_eq!( - body, - json!({ - "model": "mistral-ocr-maas", - "document": {"type": "document_url", "document_url": "data:application/pdf;base64,YWJj"}, - "pages": [0, 2], - "include_image_base64": true, - "unknown": "ignored" - }) - ); - } - let payload = json!({"pages": [{"index": 0, "markdown": "hello"}], "extra": "preserved"}); - let raw = serde_json::to_vec(&payload).unwrap(); - let direct_response = MistralOcrConfig - .transform_ocr_response(&direct.model, &raw, OcrResponseFormat::Litellm) - .unwrap() - .into_json(); - let vertex_response = VertexAiOcrConfig - .transform_ocr_response(&vertex.model, &raw, OcrResponseFormat::Litellm) - .unwrap() - .into_json(); - assert_eq!(direct_response, vertex_response); - assert_eq!(direct_response["model"], "mistral-ocr-maas"); - assert_eq!(direct_response["object"], "ocr"); - assert_eq!(direct_response["extra"], "preserved"); - } - - mod transformation { - - use rstest::rstest; - use serde_json::{Value, json}; - - use crate::ocr::test_support::wire_request; - - #[rstest] - #[case::mistral(false)] - #[case::vertex(true)] - #[tokio::test] - async fn configs_build_complete_requests_and_share_mistral_normalization( - #[case] use_vertex: bool, - ) { - use std::time::Duration; - - use litellm_llms::{ - base_llm::ocr::transformation::BaseOcrConfig, - mistral::ocr::transformation::MistralOcrConfig, - vertex_ai::ocr::transformation::VertexAiOcrConfig, - }; - - use crate::ocr::test_support::ocr_client; - - let client = ocr_client(); - let options = json!({ - "pages": [0, 2], - "include_image_base64": true, - "vertex_project": "project-1", - "vertex_location": "us-central1", - "unknown": "preserved" - }); - let direct = wire_request( - "mistral/mistral-ocr-maas", - "https://mistral.test", - options.clone(), - ); - let vertex = wire_request("vertex_ai/mistral-ocr-maas", "https://vertex.test", options); - let direct = crate::ocr::prepare::prepare_request_for_test( - crate::ocr::test_support::resolved_request(direct), - ); - let vertex = crate::ocr::prepare::prepare_request_for_test( - crate::ocr::test_support::resolved_request(vertex), - ); - let direct_http = MistralOcrConfig - .prepare_request(&direct, &client, &crate::ocr::test_support::NoHooks) - .await - .unwrap(); - let vertex_http = VertexAiOcrConfig - .prepare_request(&vertex, &client, &crate::ocr::test_support::NoHooks) - .await - .unwrap(); - assert_eq!(direct_http.url(), "https://mistral.test/v1/ocr"); - assert_eq!( - vertex_http.url(), - "https://vertex.test/v1/projects/project-1/locations/us-central1/publishers/mistralai/models/mistral-ocr-maas:rawPredict" - ); - let http = if use_vertex { - &vertex_http - } else { - &direct_http - }; - assert_eq!(http.header("authorization").unwrap(), "Bearer test-key"); - assert_eq!(http.header("content-type").unwrap(), "application/json"); - assert_eq!(http.timeout(), Some(Duration::from_secs(2))); - let body: Value = serde_json::from_slice(http.body()).unwrap(); - assert_eq!( - body, - json!({ - "model": "mistral-ocr-maas", - "document": {"type": "document_url", "document_url": "data:application/pdf;base64,YWJj"}, - "pages": [0, 2], - "include_image_base64": true, - "unknown": "preserved" - }) - ); - let payload = serde_json::to_vec( - &json!({"pages": [{"index": 0, "markdown": "hello"}], "extra": "preserved"}), - ) - .unwrap(); - let direct_response = MistralOcrConfig - .transform_ocr_response(&direct.model, &payload, Default::default()) - .unwrap() - .into_json(); - let vertex_response = VertexAiOcrConfig - .transform_ocr_response(&vertex.model, &payload, Default::default()) - .unwrap() - .into_json(); - assert_eq!(direct_response, vertex_response); - assert_eq!(direct_response["model"], "mistral-ocr-maas"); - assert_eq!(direct_response["object"], "ocr"); - assert_eq!(direct_response["extra"], "preserved"); - } - } -} - -#[cfg(test)] -mod vertex_ai_deepseek_tests { - use litellm_auth::InputSource; - use serde_json::{Value, json}; - - use crate::ocr::test_support::{MockResponse, mock_server, perform_ocr, wire_request}; - - fn request_body(request: &str) -> Value { - serde_json::from_str(request.split_once("\r\n\r\n").unwrap().1).unwrap() - } - - #[tokio::test] - async fn facade_executes_vertex_deepseek_at_the_openai_endpoint() { - let (base, seen, server) = mock_server(vec![MockResponse::json(json!({ - "choices":[{"message":{"content":"recognized"}}], - "usage":{"prompt_tokens":1} - }))]) - .await; - let request = wire_request( - "vertex_ai/deepseek-ocr-maas", - &base, - json!({ - "vertex_project":"project-1", - "vertex_location":"europe-west4", - "temperature":0.1, - "future_ocr_option":true, - "extra_body":{"provider_option":"value"} - }), - ); - let request = crate::ocr::test_support::with_source(request, "gs://bucket/document.pdf"); - - let response = perform_ocr(request).await.unwrap(); - server.await.unwrap(); - assert_eq!(response.pages[0].markdown, "recognized"); - assert_eq!( - response.usage_info.unwrap().extra_fields["prompt_tokens"], - 1 - ); - let requests = seen.lock().unwrap(); - assert!(requests[0].starts_with( - "POST /v1/projects/project-1/locations/europe-west4/endpoints/openapi/chat/completions " - )); - assert!( - requests[0] - .to_ascii_lowercase() - .contains("authorization: bearer test-key") - ); - let body = request_body(&requests[0]); - assert_eq!(body["model"], "deepseek-ai/deepseek-ocr-maas"); - assert_eq!(body["temperature"], 0.1); - assert_eq!(body["future_ocr_option"], true); - assert!(body.get("extra_body").is_none()); - assert_eq!( - body["messages"][0]["content"][0], - json!({"type":"image_url","image_url":"gs://bucket/document.pdf"}) - ); - } - - #[test] - fn host_registration_selects_deepseek_without_affecting_mistral() { - assert!(crate::ocr::arguments::is_supported_request( - "deepseek-ocr-maas", - Some("vertex_ai") - )); - assert!(crate::ocr::arguments::is_supported_request( - "mistral-ocr-maas", - Some("vertex_ai") - )); - } - - #[tokio::test] - async fn request_controlled_api_base_is_rejected_before_vertex_auth() { - let mut request = wire_request( - "vertex_ai/deepseek-ocr-maas", - "https://caller.example", - json!({"vertex_project":"project-1"}), - ); - request.credentials.api_base = Some(litellm_auth::Sourced::new( - "https://caller.example".into(), - InputSource::Request, - )); - - let error = perform_ocr(request).await.unwrap_err(); - assert!( - error - .to_string() - .contains("request-controlled Vertex AI endpoint") - ); - } - - mod deepseek_transformation { - use serde_json::json; - - use super::*; - use crate::ocr::test_support::{MockResponse, mock_server, perform_ocr, wire_request}; - - #[tokio::test] - async fn facade_executes_vertex_deepseek_at_the_openai_endpoint() { - let (base, seen, server) = mock_server(vec![MockResponse::json(json!({ - "choices":[{"message":{"content":"recognized"}}], - "usage":{"prompt_tokens":1} - }))]) - .await; - let request = wire_request( - "vertex_ai/deepseek-ocr-maas", - &base, - json!({ - "vertex_project":"project-1", - "vertex_location":"europe-west4", - "temperature":0.1, - "future_ocr_option":true, - "extra_body":{"provider_option":"value"} - }), - ); - let request = - crate::ocr::test_support::with_source(request, "gs://bucket/document.pdf"); - - let response = perform_ocr(request).await.unwrap(); - server.await.unwrap(); - assert_eq!(response.pages[0].markdown, "recognized"); - assert_eq!( - response.usage_info.unwrap().extra_fields["prompt_tokens"], - 1 - ); - let requests = seen.lock().unwrap(); - assert!(requests[0].starts_with( - "POST /v1/projects/project-1/locations/europe-west4/endpoints/openapi/chat/completions " - )); - assert!( - requests[0] - .to_ascii_lowercase() - .contains("authorization: bearer test-key") - ); - let body = request_body(&requests[0]); - assert_eq!(body["model"], "deepseek-ai/deepseek-ocr-maas"); - assert_eq!(body["temperature"], 0.1); - assert_eq!(body["future_ocr_option"], true); - assert_eq!(body["provider_option"], "value"); - assert!(body.get("vertex_project").is_none()); - assert!(body.get("extra_body").is_none()); - assert_eq!( - body["messages"][0]["content"][0], - json!({"type":"image_url","image_url":"gs://bucket/document.pdf"}) - ); - } - } -} - -#[cfg(test)] -pub(crate) mod tests { - use std::sync::{Arc, Mutex}; - - use futures_util::future::BoxFuture; - use litellm_auth_gcp::VertexAuth; - use litellm_host::{ - event::{CallEvent, MachineEvent, WireRequest}, - host::{Host, HostOp}, - machine::{HostFailure, Machine, MachineStep}, - }; - use litellm_http::{ - HttpClientPool, HttpSettings, Resolution, - media::{PublicDnsResolver, UrlPolicy}, - }; - use litellm_llms::base_llm::ocr::{ - error::Error as OcrError, - handler::OcrClient, - settings::OcrSettings, - transformation::{ - BaseOcrConfig, LiteLLMOcrResponse, OCR_RESPONSE_MAX_BYTES, OcrTransportConfig, - }, - }; - use litellm_secrets::source::SecretSource; - use rstest::rstest; - use serde_json::{Value, json}; - - use crate::ocr::route::{LocalOcrHost, OcrOp, OcrProjection, ocr_machine}; - use crate::ocr::{ - test_support::{ - MockResponse, mock_server, ocr_client, perform_ocr, perform_ocr_with, wire_request, - }, - wire::{OcrWireRequest, decode_request}, - }; - - struct RecordingSecretSource { - names: Arc>>, - values: &'static [(&'static str, &'static str)], - api_base: String, - } - - impl SecretSource for RecordingSecretSource { - fn get_secret_str<'a>( - &'a self, - name: &'a str, - ) -> BoxFuture<'a, Result, litellm_secrets::Error>> - { - self.names.lock().unwrap().push(name.to_owned()); - Box::pin(async move { - Ok(match name { - "MISTRAL_AZURE_API_BASE" => Some(self.api_base.clone()), - _ => self - .values - .iter() - .find(|(key, _)| *key == name) - .map(|(_, value)| value.to_string()), - } - .map(litellm_secrets::SecretValue::new)) - }) - } - } - - #[rstest] - #[case::mistral("mistral/model", json!({}))] - #[case::vertex("vertex_ai/mistral-ocr-latest", json!({"vertex_project":"test-project", "vertex_location":"us-central1"}))] - #[tokio::test] - async fn ocr_contract_upstream_error_preserves_status_body_and_headers( - #[case] model: &str, - #[case] options: Value, - ) { - let payload = json!({"message": format!("{} END-OF-PROVIDER-BODY", "x".repeat(4096))}); - let expected_body = serde_json::to_string(&payload).unwrap(); - let (base, seen, server) = mock_server(vec![MockResponse { - status: 422, - headers: vec![ - ("Retry-After", "17".into()), - ("X-Request-ID", "request-123".into()), - ("X-Future-Header", "retained".into()), - ], - body: payload, - }]) - .await; - let error = perform_ocr(wire_request(model, &base, options)) - .await - .unwrap_err(); - server.await.unwrap(); - assert_eq!(seen.lock().unwrap().len(), 1); - let OcrError::Provider { - status, - body, - headers, - } = error - else { - panic!("expected provider error, got {error:?}"); - }; - assert_eq!(status, 422); - for (name, value) in [ - ("retry-after", "17"), - ("x-request-id", "request-123"), - ("x-future-header", "retained"), - ] { - assert!( - headers - .iter() - .any(|(key, actual)| key.eq_ignore_ascii_case(name) && actual == value) - ); - } - assert_eq!( - body.len(), - expected_body.len(), - "provider error body was truncated" - ); - assert_eq!(body, expected_body); - } - - #[test] - fn request_boundary_selects_mistral_and_rejects_unknown_providers() { - let request = OcrWireRequest { - model: "mistral/model".into(), - document: json!({"type":"document_url","document_url":"https://example.com/doc.pdf"}), - api_key: Some(litellm_auth::SecretValue::new("key")), - api_base: None, - custom_llm_provider: None, - extra_headers: None, - optional_params: json!({"extract_header":true,"unknown":42}) - .as_object() - .unwrap() - .clone(), - input_sources: Default::default(), - timeout_seconds: None, - }; - assert!(decode_request(request).is_ok()); - assert!( - decode_request(OcrWireRequest { - model: "model".into(), - document: json!({"type":"document_url","document_url":"https://example.com/doc.pdf"}), - api_key: Some(litellm_auth::SecretValue::new("key")), - api_base: None, - custom_llm_provider: Some("unknown".into()), - extra_headers: None, - optional_params: serde_json::Map::new(), - input_sources: Default::default(), - timeout_seconds: None, - }) - .is_err() - ); - } - - #[tokio::test] - async fn facade_executes_direct_mistral_once() { - let (base, seen, server) = mock_server(vec![MockResponse::json(json!({ - "pages":[{"index":0,"markdown":"hello","custom":"preserved"}], - "usage_info":{"pages_processed":1} - }))]) - .await; - let result = perform_ocr(wire_request( - "mistral/model", - &base, - json!({"pages":"0,2-4","extract_header":true,"unknown":"ignored"}), - )) - .await - .unwrap(); - server.await.unwrap(); - assert_eq!(result.pages[0].markdown, "hello"); - assert_eq!(result.pages[0].extra_fields["custom"], "preserved"); - let requests = seen.lock().unwrap(); - assert_eq!(requests.len(), 1); - assert!(requests[0].starts_with("POST /v1/ocr ")); - assert!( - requests[0] - .to_ascii_lowercase() - .contains("authorization: bearer test-key\r\n") - ); - let body: Value = - serde_json::from_str(requests[0].split_once("\r\n\r\n").unwrap().1).unwrap(); - assert_eq!( - body, - json!({ - "model":"model", - "document":{"type":"document_url","document_url":"data:application/pdf;base64,YWJj"}, - "pages":"0,2-4", - "extract_header":true, - "unknown":"ignored" - }) - ); - } - - #[tokio::test] - async fn facade_retains_native_response_when_requested() { - let provider_response = json!({ - "pages":[{"index":0,"markdown":"hello"}], - "usage_info":{"pages_processed":1}, - "provider_only":"preserved" - }); - let (base, _, server) = - mock_server(vec![MockResponse::json(provider_response.clone())]).await; - let response = perform_ocr(wire_request( - "mistral/model", - &base, - json!({"req_format":"native"}), - )) - .await - .unwrap(); - - server.await.unwrap(); - assert_eq!( - response.provider_native_response.map(Value::Object), - Some(provider_response) - ); - } - - #[rstest] - #[case::plain_key(&[("MISTRAL_API_KEY", "plain")], "plain")] - #[case::azure_key_wins(&[("MISTRAL_AZURE_API_KEY", "azure"), ("MISTRAL_API_KEY", "plain")], "azure")] - #[case::empty_azure_key_falls_through(&[("MISTRAL_AZURE_API_KEY", ""), ("MISTRAL_API_KEY", "plain")], "plain")] - #[tokio::test] - async fn mistral_env_fallbacks_follow_python_through_the_injected_secret_source( - #[case] secrets: &'static [(&'static str, &'static str)], - #[case] expected_key: &str, - ) { - let (base, seen, server) = mock_server(vec![MockResponse::json(json!({"pages":[]}))]).await; - let names = Arc::new(Mutex::new(Vec::new())); - let client = ocr_client().with_secrets(Arc::new(RecordingSecretSource { - names: names.clone(), - values: secrets, - api_base: base.clone(), - })); - let request = decode_request(OcrWireRequest { - model: "mistral/model".into(), - document: json!({"type":"document_url","document_url":"data:application/pdf;base64,YWJj"}), - api_key: None, - api_base: None, - custom_llm_provider: None, - extra_headers: None, - optional_params: Default::default(), - input_sources: Default::default(), - timeout_seconds: Some(2.0), - }) - .unwrap(); - - crate::ocr::client::perform(&client, request).await.unwrap(); - server.await.unwrap(); - assert_eq!( - *names.lock().unwrap(), - litellm_llms::mistral::ocr::transformation::MistralOcrConfig.secret_names() - ); - assert!(seen.lock().unwrap()[0].contains(&format!("authorization: Bearer {expected_key}"))); - } - - #[tokio::test] - async fn mistral_ocr_resolves_provider_secrets_before_transformation() { - let (base, seen, server) = mock_server(vec![MockResponse::json(json!({"pages":[]}))]).await; - let names = Arc::new(Mutex::new(Vec::new())); - let client = ocr_client().with_secrets(Arc::new(RecordingSecretSource { - names: names.clone(), - values: &[("MISTRAL_API_KEY", "source-key")], - api_base: base.clone(), - })); - let request = decode_request(OcrWireRequest { - model: "mistral/mistral-ocr-latest".into(), - document: json!({ - "type":"document_url", - "document_url":"data:application/pdf;base64,YWJj" - }), - api_key: None, - api_base: None, - custom_llm_provider: None, - extra_headers: None, - optional_params: Default::default(), - input_sources: Default::default(), - timeout_seconds: Some(2.0), - }) - .unwrap(); - - crate::ocr::client::perform(&client, request).await.unwrap(); - server.await.unwrap(); - assert_eq!( - *names.lock().unwrap(), - litellm_llms::mistral::ocr::transformation::MistralOcrConfig.secret_names() - ); - assert!(seen.lock().unwrap()[0].contains("authorization: Bearer source-key")); - } - - #[tokio::test] - async fn ocr_client_uses_the_injected_http_pool_configuration() { - let (base, seen, server) = mock_server(vec![MockResponse::json(json!({"pages":[]}))]).await; - let settings = HttpSettings { - user_agent: Some("host-owned/1".into()), - ..HttpSettings::default() - }; - let client = OcrClient::new( - &HttpClientPool::new(Arc::new(PublicDnsResolver)), - &Resolution::from(&settings).config, - UrlPolicy::default(), - VertexAuth::default(), - OcrSettings::default(), - Arc::new(litellm_secrets::source::EnvironmentSecrets::default()), - ) - .unwrap(); - crate::ocr::client::perform(&client, wire_request("mistral/model", &base, json!({}))) - .await - .unwrap(); - server.await.unwrap(); - assert!(seen.lock().unwrap()[0].contains("user-agent: host-owned/1")); - } - - fn event_name(event: &CallEvent) -> &'static str { - match event { - CallEvent::Started { .. } => "started", - CallEvent::Machine(MachineEvent::ResponseReceived { .. }) => "response", - CallEvent::Succeeded { .. } => "success", - CallEvent::Failed { .. } => "failure", - } - } - - fn recording_host( - request: crate::ocr::types::LiteLLMOcrRequest, - events: Arc>>, - block: bool, - ) -> LocalOcrHost { - let before_send_events = events.clone(); - LocalOcrHost::new(request) - .with_before_send(move |wire, _| { - before_send_events.lock().unwrap().push("before_send"); - if block { - return Err(OcrError::InvalidRequest("blocked".into())); - } - Ok(wire) - }) - .with_observer(move |event| events.lock().unwrap().push(event_name(event))) - } - - #[tokio::test] - async fn lifecycle_sends_headers_returned_by_the_before_send_operation() { - let (base, seen, server) = mock_server(vec![MockResponse::json(json!({"pages":[]}))]).await; - let host = LocalOcrHost::new(wire_request("mistral/model", &base, json!({}))) - .with_before_send(|mut wire, _| { - wire.headers - .push(("x-core-callback".into(), "edited".into())); - Ok(wire) - }); - - perform_ocr_with(host).await.unwrap(); - server.await.unwrap(); - - assert!(seen.lock().unwrap()[0].contains("x-core-callback: edited")); - } - - #[tokio::test] - async fn before_send_context_names_the_route_and_its_secrets() { - let (base, _, server) = mock_server(vec![MockResponse::json(json!({"pages":[]}))]).await; - let observed = Arc::new(Mutex::new(None)); - let captured = observed.clone(); - let host = LocalOcrHost::new(wire_request( - "mistral/model", - &base, - json!({"pages": [0], "req_format": "native"}), - )) - .with_before_send(move |wire, context| { - *captured.lock().unwrap() = Some((wire.clone(), context.clone())); - Ok(wire) - }); - perform_ocr_with(host).await.unwrap(); - server.await.unwrap(); - let (wire, context) = observed.lock().unwrap().take().unwrap(); - assert_eq!(context.custom_llm_provider, "mistral"); - assert_eq!(context.model, "model"); - assert_eq!(wire.body["pages"], json!([0])); - assert!(context.secret_fields.is_empty()); - assert_eq!(context.optional_params["req_format"], "native"); - - let (base, _, server) = mock_server(vec![MockResponse::json(json!({"pages":[]}))]).await; - let observed = Arc::new(Mutex::new(None)); - let captured = observed.clone(); - let request = wire_request( - "azure_ai/model", - &base, - json!({"client_secret": "shh", "tenant_id": "t"}), - ); - let request = request.with_document(crate::ocr::types::OcrDocumentInput::Bytes { - bytes: b"abc".as_slice().into(), - file_name: None, - mime_type: Some("application/pdf".into()), - }); - let host = LocalOcrHost::new(request).with_before_send(move |wire, context| { - *captured.lock().unwrap() = Some(context.clone()); - Ok(wire) - }); - perform_ocr_with(host).await.unwrap(); - server.await.unwrap(); - let context = observed.lock().unwrap().take().unwrap(); - assert_eq!(context.secret_fields, ["client_secret"]); - } - - #[tokio::test] - async fn lifecycle_orders_hooks_and_emits_one_success() { - let (base, seen, server) = mock_server(vec![MockResponse::json(json!({"pages":[]}))]).await; - let events = Arc::new(Mutex::new(Vec::new())); - let host = recording_host( - wire_request("mistral/model", &base, json!({})), - events.clone(), - false, - ); - perform_ocr_with(host).await.unwrap(); - server.await.unwrap(); - assert_eq!( - *events.lock().unwrap(), - ["started", "before_send", "response", "success"] - ); - assert_eq!(seen.lock().unwrap().len(), 1); - } - - #[tokio::test] - async fn lifecycle_blocking_prevents_execution_and_emits_one_failure() { - let events = Arc::new(Mutex::new(Vec::new())); - let host = recording_host( - wire_request("mistral/model", "http://127.0.0.1:1", json!({})), - events.clone(), - true, - ); - let error = perform_ocr_with(host).await.unwrap_err(); - assert!(matches!(error, OcrError::InvalidRequest(message) if message == "blocked")); - assert_eq!( - *events.lock().unwrap(), - ["started", "before_send", "failure"] - ); - } - - #[tokio::test] - async fn upstream_failure_emits_one_terminal_failure() { - let (base, seen, server) = mock_server(vec![MockResponse { - status: 500, - headers: vec![], - body: json!({"error":"failed"}), - }]) - .await; - let events = Arc::new(Mutex::new(Vec::new())); - let host = recording_host( - wire_request("mistral/model", &base, json!({})), - events.clone(), - false, - ); - assert!(perform_ocr_with(host).await.is_err()); - server.await.unwrap(); - assert_eq!( - *events.lock().unwrap(), - ["started", "before_send", "failure"] - ); - assert_eq!(seen.lock().unwrap().len(), 1); - } - - /// Drives the machine by hand, answering every op through `host` except `before_send`, - /// which `intercept` answers so a test can fail or cancel exactly there. - async fn drive_until( - client: OcrClient, - host: &LocalOcrHost, - mut intercept: impl FnMut(WireRequest) -> Result>, - ) -> ( - Result, - Vec<&'static str>, - crate::ocr::route::OcrMachine, - ) { - let mut machine = ocr_machine(client); - let mut ops = Vec::new(); - let outcome = loop { - let op = match machine.resume().await { - Ok(MachineStep::Host(op)) => op, - Ok(MachineStep::Complete(response)) => break Ok(response), - Err(error) => break Err(error), - }; - let answer = match op { - HostOp::Project(reply) => { - ops.push("Project"); - host.project() - .await - .map(|projection| reply.send(projection)) - .map_err(HostFailure::Error) - } - HostOp::Custom(op) => { - ops.push(match op { - OcrOp::AcquireAzureAdToken(_) => "AcquireAzureAdToken", - }); - host.custom_op(op).await.map_err(HostFailure::Error) - } - HostOp::BeforeSend { wire, reply, .. } => { - ops.push("BeforeSend"); - intercept(*wire).map(|wire| reply.send(wire)) - } - HostOp::Emit(event, reply) => { - let event = CallEvent::Machine(event); - ops.push(event_name(&event)); - host.emit(&event) - .await - .map(|()| reply.send(())) - .map_err(HostFailure::Error) - } - }; - if let Err(failure) = answer { - break machine.interrupt(failure).await; - } - }; - (outcome, ops, machine) - } - - #[tokio::test] - async fn failed_before_send_does_not_replay_or_reach_transport() { - let host = LocalOcrHost::new(wire_request( - "mistral/model", - "http://127.0.0.1:1", - json!({}), - )); - let (outcome, ops, mut machine) = drive_until(ocr_client(), &host, |_| { - Err(HostFailure::Error(OcrError::InvalidRequest( - "before_send failed".into(), - ))) - }) - .await; - assert!( - matches!(outcome, Err(OcrError::InvalidRequest(message)) if message == "before_send failed") - ); - assert_eq!(ops, ["Project", "BeforeSend"]); - assert!(machine.resume().await.is_err()); - } - - #[tokio::test] - async fn invalid_provider_response_emits_response_received_before_normalization_failure() { - let (base, seen, server) = - mock_server(vec![MockResponse::json(json!({"pages":"invalid"}))]).await; - let responses_received = Arc::new(Mutex::new(Vec::new())); - let observed = responses_received.clone(); - let host = LocalOcrHost::new(wire_request("mistral/model", &base, json!({}))) - .with_observer(move |event| { - if let CallEvent::Machine(MachineEvent::ResponseReceived { raw }) = event { - observed.lock().unwrap().push(raw.body.clone()); - } - }); - let error = perform_ocr_with(host).await.unwrap_err(); - server.await.unwrap(); - assert!(matches!(error, OcrError::ResponseField { .. })); - assert_eq!(seen.lock().unwrap().len(), 1); - assert_eq!( - *responses_received.lock().unwrap(), - [r#"{"pages":"invalid"}"#] - ); - } - - #[tokio::test] - async fn direct_native_host_drives_the_same_state_machine() { - let (base, seen, server) = mock_server(vec![MockResponse::json(json!({ - "pages":[{"index":0,"markdown":"native"}] - }))]) - .await; - let host = LocalOcrHost::new(wire_request("mistral/model", &base, json!({}))); - let (outcome, ops, mut machine) = drive_until(ocr_client(), &host, Ok).await; - server.await.unwrap(); - assert_eq!(outcome.unwrap().pages[0].markdown, "native"); - assert_eq!(seen.lock().unwrap().len(), 1); - assert_eq!(ops, ["Project", "BeforeSend", "response"]); - assert!(matches!( - machine.resume().await, - Err(OcrError::InvalidRequest(_)) - )); - } - - #[tokio::test] - async fn empty_byte_documents_fail_before_the_provider_is_called() { - let (base, seen, _server) = mock_server(vec![]).await; - let request = wire_request("mistral/model", &base, json!({})).with_document( - crate::ocr::types::OcrDocumentInput::Bytes { - bytes: Default::default(), - file_name: None, - mime_type: None, - }, - ); - let response = perform_ocr_with(LocalOcrHost::new(request)).await; - assert!(matches!(response.unwrap_err(), OcrError::EmptyFile)); - assert!(seen.lock().unwrap().is_empty()); - } - - #[tokio::test] - async fn path_documents_are_read_by_core_without_a_host_operation() { - let (base, seen, server) = mock_server(vec![MockResponse::json(json!({ - "pages":[{"index":0,"markdown":"path"}] - }))]) - .await; - let dir = std::env::temp_dir().join(format!("litellm-ocr-{}", rand::random::())); - std::fs::create_dir_all(&dir).unwrap(); - let path = dir.join("scan.png"); - std::fs::write(&path, b"abc").unwrap(); - let request = wire_request("mistral/model", &base, json!({})).with_document( - crate::ocr::types::OcrDocumentInput::Path { - path: path.clone(), - mime_type: None, - }, - ); - let (response, ops, _) = drive_until(ocr_client(), &LocalOcrHost::new(request), Ok).await; - server.await.unwrap(); - std::fs::remove_dir_all(&dir).unwrap(); - assert_eq!(response.unwrap().pages[0].markdown, "path"); - assert_eq!(ops, ["Project", "BeforeSend", "response"]); - assert!(seen.lock().unwrap()[0].contains("data:image/png;base64,YWJj")); - - let (base, seen, _server) = mock_server(vec![]).await; - let request = wire_request("mistral/model", &base, json!({})).with_document( - crate::ocr::types::OcrDocumentInput::Path { - path: path.clone(), - mime_type: None, - }, - ); - let response = perform_ocr_with(LocalOcrHost::new(request)).await; - assert!(matches!( - response.unwrap_err(), - OcrError::FileRead { path: failed, source } if failed == path && source.kind() == std::io::ErrorKind::NotFound - )); - assert!(seen.lock().unwrap().is_empty()); - } - - #[tokio::test] - async fn cancellation_at_before_send_prevents_execution_and_further_resumption() { - let host = LocalOcrHost::new(wire_request( - "mistral/model", - "http://127.0.0.1:1", - json!({}), - )); - let (outcome, ops, mut machine) = drive_until(ocr_client(), &host, |_| { - Err(HostFailure::Cancelled(OcrError::InvalidRequest( - "cancelled".into(), - ))) - }) - .await; - assert!( - matches!(outcome, Err(OcrError::InvalidRequest(message)) if message == "cancelled") - ); - assert_eq!(ops, ["Project", "BeforeSend"]); - assert!(machine.resume().await.is_err()); - } - - #[tokio::test] - async fn resuming_before_answering_preserves_pending_operation() { - let request = wire_request("mistral/model", "http://127.0.0.1:1", json!({})); - let mut machine = ocr_machine(ocr_client()); - let Ok(MachineStep::Host(HostOp::Project(reply))) = machine.resume().await else { - panic!("expected the projection op first"); - }; - assert!(machine.resume().await.is_err()); - reply.send(OcrProjection { - request, - caller_token: false, - }); - assert!(matches!( - machine.resume().await, - Ok(MachineStep::Host(HostOp::BeforeSend { .. })) - )); - } - - async fn read_bounded_response( - response: Vec, - limit: usize, - ) -> Result { - use tokio::io::{AsyncReadExt, AsyncWriteExt}; - - let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); - let address = listener.local_addr().unwrap(); - let server = tokio::spawn(async move { - let (mut socket, _) = listener.accept().await.unwrap(); - let mut request = [0; 4096]; - assert!(socket.read(&mut request).await.unwrap() > 0); - socket.write_all(&response).await.unwrap(); - std::future::pending::<()>().await; - }); - let response = reqwest::Client::new() - .get(format!("http://{address}")) - .send() - .await - .unwrap(); - let result = tokio::time::timeout( - std::time::Duration::from_secs(2), - litellm_llms::base_llm::ocr::handler::read_response_bytes(response, limit), - ) - .await; - server.abort(); - let _ = server.await; - result.expect("bounded reads must finish without waiting for the rest of an oversized body") - } - - #[tokio::test] - async fn response_limit_accepts_exact_size_and_rejects_declared_and_chunked_overflow() { - use litellm_llms::base_llm::ocr::error::Error; - - for response in [ - "HTTP/1.1 200 OK\r\nContent-Length: 8\r\n\r\nabcdefgh", - "HTTP/1.1 200 OK\r\nTransfer-Encoding: chunked\r\n\r\n4\r\nabcd\r\n4\r\nefgh\r\n0\r\n\r\n", - ] { - assert_eq!( - read_bounded_response(response.as_bytes().to_vec(), 8) - .await - .unwrap(), - "abcdefgh" - ); - } - for response in [ - "HTTP/1.1 200 OK\r\nContent-Length: 9\r\n\r\n", - "HTTP/1.1 200 OK\r\nTransfer-Encoding: chunked\r\n\r\n4\r\nabcd\r\n5\r\nefghi\r\n", - ] { - assert!(matches!( - read_bounded_response(response.as_bytes().to_vec(), 8).await, - Err(Error::TooLarge { limit: 8 }) - )); - } - } - - #[rstest] - #[case::declared("Content-Length: 1000000")] - #[case::chunked("Transfer-Encoding: chunked")] - #[tokio::test] - async fn oversized_error_retains_http_status_and_bounded_diagnostics_without_draining( - #[case] headers: &str, - ) { - let prefix = "x".repeat(4096); - let body = if headers.starts_with("Transfer") { - format!("{:x}\r\n{prefix}\r\n", prefix.len()) - } else { - prefix.clone() - }; - let response = format!("HTTP/1.1 429 Too Many Requests\r\n{headers}\r\n\r\n{body}"); - let error = read_bounded_response(response.into_bytes(), prefix.len()) - .await - .unwrap_err(); - match error { - OcrError::Transport(litellm_http::transport::Error::Http { status, body }) => { - assert_eq!(status, 429); - assert_eq!(body, prefix); - } - error => panic!("unexpected error: {error}"), - } - } - - #[test] - fn response_limit_is_validated_and_not_forwarded_to_the_provider() { - let request = wire_request( - "mistral/model", - "http://localhost", - json!({"max_response_bytes": 123}), - ); - assert_eq!(request.transport.max_response_bytes, 123); - assert!(!request.optional_params.contains_key("max_response_bytes")); - for value in [ - json!(0), - json!(-1), - json!(true), - json!("123"), - json!(1.5), - json!(OCR_RESPONSE_MAX_BYTES + 1), - Value::Null, - ] { - let wire = serde_json::from_value(json!({ - "model": "mistral/model", "document": {"type": "document_url", "document_url": "data:application/pdf;base64,YWJj"}, - "optional_params": {"max_response_bytes": value} - })).unwrap(); - let Err(error) = decode_request(wire) else { - panic!("invalid response limit accepted") - }; - assert!(error.to_string().contains("max_response_bytes")); - } - } - - #[derive(Debug)] - struct PendingToken { - entered: Arc, - dropped: Arc, - } - - struct TokenFutureDrop(Arc); - - impl Drop for TokenFutureDrop { - fn drop(&mut self) { - self.0.store(true, std::sync::atomic::Ordering::SeqCst); - } - } - - impl litellm_auth::TokenProvider for PendingToken { - fn acquire(&self) -> litellm_auth::TokenFuture<'_> { - Box::pin(async move { - let _guard = TokenFutureDrop(self.dropped.clone()); - self.entered.notify_one(); - std::future::pending().await - }) - } - } - - #[tokio::test] - async fn interrupt_drops_provider_captures_before_returning() { - use std::sync::atomic::{AtomicBool, Ordering}; - - let entered = Arc::new(tokio::sync::Notify::new()); - let dropped = Arc::new(AtomicBool::new(false)); - let request = wire_request("azure_ai/mistral-ocr", "https://example.invalid", json!({})); - let request = crate::ocr::types::LiteLLMOcrRequest { - transport: OcrTransportConfig { - extra_headers: vec![("authorization".into(), "Bearer test-key".into())], - ..request.transport - }, - azure_ad_token_provider: Some(litellm_auth::TokenProviderHandle::new(Arc::new( - PendingToken { - entered: entered.clone(), - dropped: dropped.clone(), - }, - ))), - ..request - }; - let host = LocalOcrHost::new(request); - let mut machine = ocr_machine(ocr_client()); - tokio::time::timeout(std::time::Duration::from_secs(2), async { - loop { - tokio::select! { - _ = entered.notified() => break, - step = machine.resume() => { - match step.unwrap() { - MachineStep::Host(HostOp::Project(reply)) => reply.send(host.project().await.unwrap()), - MachineStep::Host(HostOp::Custom(op)) => host.custom_op(op).await.unwrap(), - MachineStep::Host(HostOp::BeforeSend { wire, reply, .. }) => reply.send(*wire), - MachineStep::Host(HostOp::Emit(_, reply)) => reply.send(()), - MachineStep::Complete(_) => panic!("pending provider completed"), - } - } - } - } - }) - .await - .unwrap(); - assert!(!dropped.load(Ordering::SeqCst)); - let selected = OcrError::InvalidRequest("cancelled".into()); - let acknowledgement = machine.interrupt(HostFailure::Cancelled(selected.clone())); - assert!( - dropped.load(Ordering::SeqCst), - "interrupt returned while provider captures were still alive" - ); - assert!( - matches!(acknowledgement.await, Err(OcrError::InvalidRequest(message)) if message == "cancelled") - ); - } - - struct CallerTokenHost { - request: Mutex>, - trace: Mutex>, - } - - impl Host for CallerTokenHost { - async fn project(&self) -> Result { - self.trace.lock().unwrap().push("project".into()); - Ok(OcrProjection { - request: self.request.lock().unwrap().take().unwrap(), - caller_token: true, - }) - } - - async fn custom_op(&self, op: OcrOp) -> Result<(), OcrError> { - match op { - OcrOp::AcquireAzureAdToken(reply) => { - self.trace.lock().unwrap().push("token".into()); - reply.send(litellm_auth::ResolvedCredential::Static( - litellm_auth::SecretValue::new("caller-token"), - )); - Ok(()) - } - } - } - - async fn before_send( - &self, - wire: WireRequest, - _: &litellm_host::event::RequestContext, - ) -> Result { - let is_authorization = |name: &str| name.eq_ignore_ascii_case("authorization"); - let authorization = wire - .headers - .iter() - .find(|(name, _)| is_authorization(name)) - .map(|(_, value)| value.clone()) - .unwrap_or_default(); - self.trace - .lock() - .unwrap() - .push(format!("before_send:{authorization}")); - let headers = wire - .headers - .into_iter() - .map(|(name, value)| match is_authorization(&name) { - true => (name, "Bearer edited".to_string()), - false => (name, value), - }) - .collect(); - Ok(WireRequest { headers, ..wire }) - } - } - - #[tokio::test] - async fn the_callers_azure_token_is_acquired_before_before_send_which_can_still_replace_it() { - let (base, seen, server) = mock_server(vec![MockResponse::json(json!({"pages":[]}))]).await; - let mut request = wire_request("azure_ai/model", &base, json!({})); - request.credentials.api_key = None; - let host = CallerTokenHost { - request: Mutex::new(Some(request)), - trace: Mutex::new(Vec::new()), - }; - - litellm_host::run::run(ocr_machine(ocr_client()), &host) - .await - .unwrap(); - server.await.unwrap(); - - assert_eq!( - *host.trace.lock().unwrap(), - ["project", "token", "before_send:Bearer caller-token"] - ); - assert!( - seen.lock().unwrap()[0] - .to_ascii_lowercase() - .contains("authorization: bearer edited\r\n") - ); - } - - #[tokio::test] - async fn interrupting_an_in_flight_provider_request_closes_its_connection() { - use tokio::io::AsyncReadExt; - - let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); - let base = format!("http://{}", listener.local_addr().unwrap()); - let received = Arc::new(tokio::sync::Notify::new()); - let server_received = received.clone(); - let server = tokio::spawn(async move { - let (mut socket, _) = listener.accept().await.unwrap(); - let mut request = Vec::new(); - let mut buffer = [0u8; 4096]; - while !request.windows(4).any(|window| window == b"\r\n\r\n") { - let read = socket.read(&mut buffer).await.unwrap(); - request.extend_from_slice(&buffer[..read]); - } - server_received.notify_one(); - loop { - if socket.read(&mut buffer).await.unwrap() == 0 { - break; - } - } - }); - let host = LocalOcrHost::new(wire_request("mistral/model", &base, json!({}))); - let mut machine = ocr_machine(ocr_client()); - tokio::time::timeout(std::time::Duration::from_secs(2), async { - loop { - tokio::select! { - _ = received.notified() => break, - step = machine.resume() => { - match step.unwrap() { - MachineStep::Host(HostOp::Project(reply)) => reply.send(host.project().await.unwrap()), - MachineStep::Host(HostOp::Custom(op)) => host.custom_op(op).await.unwrap(), - MachineStep::Host(HostOp::BeforeSend { wire, reply, .. }) => reply.send(*wire), - MachineStep::Host(HostOp::Emit(_, reply)) => reply.send(()), - MachineStep::Complete(_) => panic!("the stalled provider completed"), - } - } - } - } - }) - .await - .unwrap(); - - let cancelled = OcrError::InvalidRequest("cancelled".into()); - assert!( - machine - .interrupt(HostFailure::Cancelled(cancelled)) - .await - .is_err() - ); - tokio::time::timeout(std::time::Duration::from_secs(1), server) - .await - .expect("the provider connection stayed open after the interrupt") - .unwrap(); - } -} diff --git a/litellm-rust/crates/core/tests/audio_transcription.rs b/litellm-rust/crates/core/tests/audio_transcription.rs index aa5aef0149b..196f085a6c3 100644 --- a/litellm-rust/crates/core/tests/audio_transcription.rs +++ b/litellm-rust/crates/core/tests/audio_transcription.rs @@ -1,50 +1,250 @@ -use std::{ - io::{Read, Write}, - net::TcpListener, - thread, +use litellm_core::audio_transcription::{ + Error, audio_transcription, types::AudioTranscriptionRequest, }; +use rstest::{fixture, rstest}; +use serde_json::{Map, Value, json}; +use wiremock::ResponseTemplate; -use litellm_core::audio_transcription::{audio_transcription, types::AudioTranscriptionRequest}; -use serde_json::{Map, json}; +mod support; +use support::*; -#[tokio::test] -async fn bedrock_request_is_signed_and_contains_audio() { - let listener = TcpListener::bind("127.0.0.1:0").expect("listener"); - let address = listener.local_addr().expect("address"); - let server = thread::spawn(move || { - let (mut stream, _) = listener.accept().expect("connection"); - let mut request = Vec::new(); - let mut buffer = [0_u8; 16_384]; - let count = stream.read(&mut buffer).expect("request"); - request.extend_from_slice(&buffer[..count]); - let request = String::from_utf8_lossy(&request); - assert!(request.contains("POST /model/mistral.voxtral-mini-3b-2507/converse")); - assert!(request.contains("authorization: AWS4-HMAC-SHA256")); - assert!(request.contains("x-amz-date:")); - assert!(request.contains("\"bytes\":\"AQI=\"")); - assert!(request.contains("Transcribe the audio. Respond with only the transcript.")); - let response = b"HTTP/1.1 200 OK\r\nContent-Type: application/json\r\nContent-Length: 53\r\nConnection: close\r\n\r\n{\"output\":{\"message\":{\"content\":[{\"text\":\"hello\"}]}}}"; - stream.write_all(response).expect("response"); - }); +const MODEL: &str = "mistral.voxtral-mini-3b-2507"; - let optional_params = Map::from_iter([ +fn transcript_response(text: &str) -> ResponseTemplate { + json_response(json!({"output": {"message": {"content": [{"text": text}]}}})) +} + +fn aws_params(region: &str) -> Map { + Map::from_iter([ ("aws_access_key_id".to_string(), json!("access-key")), ("aws_secret_access_key".to_string(), json!("secret-key")), - ("aws_region_name".to_string(), json!("us-east-1")), - ]); - let api_base = format!("http://{address}"); - let response = audio_transcription(AudioTranscriptionRequest { - model: "mistral.voxtral-mini-3b-2507", + ("aws_region_name".to_string(), json!(region)), + ]) +} + +#[fixture] +fn request() -> AudioTranscriptionRequest<'static> { + AudioTranscriptionRequest { + model: MODEL, audio: json!({"data": "AQI=", "format": "wav", "filename": "audio.wav"}), api_key: None, - api_base: Some(&api_base), + api_base: None, custom_llm_provider: Some("bedrock"), extra_headers: None, - optional_params, + optional_params: aws_params("us-east-1"), timeout: None, + } +} + +#[rstest] +#[case::us_east_1("us-east-1")] +#[case::eu_west_1("eu-west-1")] +#[tokio::test] +async fn bedrock_converse_request_is_signed_for_the_requested_region( + request: AudioTranscriptionRequest<'static>, + #[case] region: &str, +) { + let upstream = upstream([transcript_response("hello")]).await; + let base = upstream.uri(); + + let response = audio_transcription(AudioTranscriptionRequest { + api_base: Some(&base), + optional_params: aws_params(region), + ..request }) .await .expect("transcription"); + assert_eq!(response, json!({"text": "hello"})); - server.join().expect("server"); + let sent = only_request(&upstream).await; + assert_eq!(sent.method.as_str(), "POST"); + assert_eq!(sent.url.path(), format!("/model/{MODEL}/converse")); + let authorization = sent.header("authorization").expect("request is signed"); + assert!( + authorization.starts_with("AWS4-HMAC-SHA256 Credential=access-key/"), + "{authorization}" + ); + assert!( + authorization.contains(&format!("/{region}/bedrock/aws4_request")), + "{authorization}" + ); + assert!(sent.header("x-amz-date").is_some()); + assert!(!sent.body_text().contains("secret-key")); +} + +#[rstest] +#[tokio::test] +async fn the_provider_can_come_from_the_model_prefix(request: AudioTranscriptionRequest<'static>) { + let upstream = upstream([transcript_response("hello")]).await; + let base = upstream.uri(); + let model = format!("bedrock/{MODEL}"); + + audio_transcription(AudioTranscriptionRequest { + model: &model, + custom_llm_provider: None, + api_base: Some(&base), + ..request + }) + .await + .expect("transcription"); + + assert_eq!( + only_request(&upstream).await.url.path(), + format!("/model/{MODEL}/converse") + ); +} + +#[rstest] +#[tokio::test] +async fn audio_and_transcription_params_reach_the_converse_body( + request: AudioTranscriptionRequest<'static>, + #[values("wav", "mp3", "flac", "ogg")] format: &str, +) { + let upstream = upstream([transcript_response("hello")]).await; + let base = upstream.uri(); + let optional_params = aws_params("us-east-1") + .into_iter() + .chain([ + ("language".to_string(), json!("fr")), + ("temperature".to_string(), json!(0.2)), + ]) + .collect(); + + audio_transcription(AudioTranscriptionRequest { + audio: json!({"data": "AQI=", "format": format}), + api_base: Some(&base), + optional_params, + ..request + }) + .await + .expect("transcription"); + + let body = only_request(&upstream).await.json(); + let content = &body["messages"][0]["content"]; + assert_eq!( + content[0], + json!({"audio": {"format": format, "source": {"bytes": "AQI="}}}) + ); + let instruction = content[1]["text"].as_str().expect("instruction text"); + assert!(instruction.contains("fr"), "{instruction}"); + assert_eq!(body["inferenceConfig"]["temperature"], 0.2); +} + +#[rstest] +#[case::unknown_format(json!({"data": "AQI=", "format": "aac"}))] +#[case::missing_data(json!({"format": "wav"}))] +#[case::not_an_object(json!("AQI="))] +#[tokio::test] +async fn invalid_audio_is_rejected_before_sending( + request: AudioTranscriptionRequest<'static>, + #[case] audio: Value, +) { + let upstream = upstream([transcript_response("hello")]).await; + let base = upstream.uri(); + + let error = audio_transcription(AudioTranscriptionRequest { + audio, + api_base: Some(&base), + ..request + }) + .await + .expect_err("invalid audio is rejected"); + + assert!( + matches!( + error, + Error::InvalidRequest(_) | Error::MissingField(_) | Error::InvalidType { .. } + ), + "{error:?}" + ); + assert!(received(&upstream).await.is_empty()); +} + +#[rstest] +#[case::unknown_provider(MODEL, Some("openai"), "openai")] +#[case::unresolvable_model( + "no-such-model", + None, + "unable to resolve custom_llm_provider for audio transcription request" +)] +#[tokio::test] +async fn unsupported_providers_are_rejected_before_sending( + request: AudioTranscriptionRequest<'static>, + #[case] model: &'static str, + #[case] provider: Option<&'static str>, + #[case] reported: &str, +) { + let error = audio_transcription(AudioTranscriptionRequest { + model, + custom_llm_provider: provider, + api_base: Some(UNREACHABLE_BASE), + ..request + }) + .await + .expect_err("unsupported provider errors"); + + assert_eq!(error, Error::InvalidProvider(reported.into())); +} + +#[rstest] +#[tokio::test] +async fn a_non_string_extra_header_is_rejected(request: AudioTranscriptionRequest<'static>) { + let error = audio_transcription(AudioTranscriptionRequest { + extra_headers: Some(Map::from_iter([("x-count".to_string(), json!(3))])), + api_base: Some(UNREACHABLE_BASE), + ..request + }) + .await + .expect_err("a non-string header is rejected"); + + assert!(matches!(error, Error::Headers(_)), "{error:?}"); +} + +#[rstest] +#[case::throttled(429)] +#[case::server_error(500)] +#[tokio::test] +async fn an_upstream_error_keeps_its_status_and_body( + request: AudioTranscriptionRequest<'static>, + #[case] status: u16, +) { + let upstream = + upstream([ResponseTemplate::new(status).set_body_string("upstream said no")]).await; + let base = upstream.uri(); + + let error = audio_transcription(AudioTranscriptionRequest { + api_base: Some(&base), + ..request + }) + .await + .expect_err("upstream error propagates"); + + assert_eq!( + error, + Error::Transport(litellm_http::transport::Error::Http { + status, + body: "upstream said no".into() + }) + ); +} + +#[rstest] +#[case::not_json(ResponseTemplate::new(200).set_body_string("not json"))] +#[case::no_output(json_response(json!({"unexpected": true})))] +#[tokio::test] +async fn an_unreadable_success_body_is_an_invalid_response( + request: AudioTranscriptionRequest<'static>, + #[case] response: ResponseTemplate, +) { + let upstream = upstream([response]).await; + let base = upstream.uri(); + + let error = audio_transcription(AudioTranscriptionRequest { + api_base: Some(&base), + ..request + }) + .await + .expect_err("an unreadable body fails"); + + assert!(matches!(error, Error::InvalidResponse(_)), "{error:?}"); } diff --git a/litellm-rust/crates/core/tests/chat_completions.rs b/litellm-rust/crates/core/tests/chat_completions.rs new file mode 100644 index 00000000000..ae96509fe2e --- /dev/null +++ b/litellm-rust/crates/core/tests/chat_completions.rs @@ -0,0 +1,320 @@ +use std::time::Duration; + +use litellm_core::chat_completions::{ + Error, chat_completions, chat_completions_decline_reason, types::ChatCompletionsRequest, +}; +use litellm_http::transport::Error as TransportError; +use rstest::{fixture, rstest}; +use serde_json::{Map, Value, json}; +use wiremock::ResponseTemplate; + +mod support; +use support::*; + +const ANTHROPIC_MESSAGE: &str = r#"{"id":"msg_1","type":"message","role":"assistant","model":"claude-sonnet-4-5-20260101","content":[{"type":"text","text":"hello"}],"stop_reason":"end_turn","stop_sequence":null,"usage":{"input_tokens":11,"output_tokens":4}}"#; + +fn object(value: Value) -> Map { + let Value::Object(map) = value else { + panic!("expected a json object, got {value}"); + }; + map +} + +fn anthropic_response(body: &str) -> ResponseTemplate { + ResponseTemplate::new(200).set_body_raw(body, "application/json") +} + +fn hi() -> Value { + json!([{"role": "user", "content": "hi"}]) +} + +#[fixture] +fn request() -> ChatCompletionsRequest<'static> { + ChatCompletionsRequest { + model: "anthropic/claude-sonnet-4-5", + messages: hi(), + optional_params: object(json!({"max_tokens": 16})), + api_key: Some("sk-test"), + api_base: None, + custom_llm_provider: None, + extra_headers: None, + timeout: Some(Duration::from_secs(10)), + } +} + +#[rstest] +#[tokio::test] +async fn anthropic_round_trip_translates_the_conversation_and_normalizes_the_response( + request: ChatCompletionsRequest<'static>, +) { + let upstream = upstream([anthropic_response(ANTHROPIC_MESSAGE)]).await; + let base = upstream.uri(); + + let response = chat_completions(ChatCompletionsRequest { + messages: json!([ + {"role": "system", "content": "be terse"}, + {"role": "user", "content": "hi"} + ]), + api_base: Some(&base), + ..request + }) + .await + .expect("call succeeds"); + + let sent = only_request(&upstream).await; + assert_eq!(sent.url.path(), "/v1/messages"); + assert_eq!(sent.header_values("x-api-key"), ["sk-test"]); + let body = sent.json(); + assert_eq!(body["model"], "claude-sonnet-4-5"); + assert_eq!( + body["messages"], + json!([{"role": "user", "content": [{"type": "text", "text": "hi"}]}]) + ); + assert_eq!( + body["system"], + json!([{"type": "text", "text": "be terse"}]) + ); + assert_eq!(body["max_tokens"], 16); + assert_eq!( + response.choices[0].message.content.as_deref(), + Some("hello") + ); + assert_eq!(response.usage.total_tokens, 15); +} + +#[rstest] +#[tokio::test] +async fn the_deployment_key_replaces_a_caller_supplied_x_api_key( + request: ChatCompletionsRequest<'static>, +) { + let upstream = upstream([anthropic_response(ANTHROPIC_MESSAGE)]).await; + let base = upstream.uri(); + + chat_completions(ChatCompletionsRequest { + api_base: Some(&base), + extra_headers: Some(object( + json!({"x-api-key": "caller-key", "x-trace": "kept"}), + )), + ..request + }) + .await + .expect("call succeeds"); + + let sent = only_request(&upstream).await; + assert_eq!(sent.header_values("x-api-key"), ["sk-test"]); + assert_eq!(sent.header("x-trace"), Some("kept")); +} + +#[rstest] +#[tokio::test] +async fn bedrock_round_trip_is_signed_and_normalized(request: ChatCompletionsRequest<'static>) { + let upstream = upstream([json_response(json!({ + "output": {"message": {"role": "assistant", "content": [{"text": "hello"}]}}, + "stopReason": "end_turn", + "usage": {"inputTokens": 11, "outputTokens": 4, "totalTokens": 15} + }))]) + .await; + let base = upstream.uri(); + + let response = chat_completions(ChatCompletionsRequest { + model: "bedrock/anthropic.claude-sonnet-4-5", + optional_params: object(json!({ + "aws_access_key_id": "access-key", + "aws_secret_access_key": "secret-key", + "aws_region_name": "eu-west-1" + })), + api_key: None, + api_base: Some(&base), + ..request + }) + .await + .expect("call succeeds"); + + let sent = only_request(&upstream).await; + assert_eq!( + sent.url.path(), + "/model/anthropic.claude-sonnet-4-5/converse" + ); + let authorization = sent.header("authorization").expect("request is signed"); + assert!( + authorization.contains("/eu-west-1/bedrock/aws4_request"), + "{authorization}" + ); + assert_eq!( + sent.json()["messages"], + json!([{"role": "user", "content": [{"text": "hi"}]}]) + ); + assert_eq!( + response.choices[0].message.content.as_deref(), + Some("hello") + ); + assert_eq!(response.usage.total_tokens, 15); +} + +/// The provider already answered and billed these, so the host must not retry them on +/// its own path: they surface as `InvalidResponse`, never as a pre-send decline. +#[rstest] +#[case::missing_usage( + r#"{"model":"m","content":[{"type":"text","text":"hi"}],"stop_reason":"end_turn"}"# +)] +#[case::tool_use_block(r#"{"model":"m","content":[{"type":"tool_use","id":"t","name":"f","input":{}}],"stop_reason":"tool_use","usage":{"input_tokens":1,"output_tokens":1}}"#)] +#[case::not_json("not json")] +#[tokio::test] +async fn a_response_it_cannot_normalize_is_reported_as_already_sent( + request: ChatCompletionsRequest<'static>, + #[case] body: &str, +) { + let upstream = upstream([anthropic_response(body)]).await; + let base = upstream.uri(); + + let error = chat_completions(ChatCompletionsRequest { + api_base: Some(&base), + ..request + }) + .await + .expect_err("response cannot be normalized"); + + assert!(matches!(error, Error::InvalidResponse(_)), "{error:?}"); +} + +#[rstest] +#[case::rate_limited(429)] +#[case::server_error(500)] +#[tokio::test] +async fn an_upstream_error_status_keeps_its_code_and_body( + request: ChatCompletionsRequest<'static>, + #[case] status: u16, +) { + let upstream = upstream([ResponseTemplate::new(status).set_body_string("slow down")]).await; + let base = upstream.uri(); + + let error = chat_completions(ChatCompletionsRequest { + api_base: Some(&base), + ..request + }) + .await + .expect_err("upstream rejects"); + + assert_eq!( + error, + Error::Transport(TransportError::Http { + status, + body: "slow down".into() + }) + ); +} + +/// Nothing was sent, so nothing was billed and the host can still serve the request. +#[rstest] +#[tokio::test] +async fn a_connection_that_is_never_established_declines_instead_of_failing( + request: ChatCompletionsRequest<'static>, +) { + let error = chat_completions(ChatCompletionsRequest { + api_base: Some(UNREACHABLE_BASE), + ..request + }) + .await + .expect_err("nothing is listening"); + + assert!( + matches!(error, Error::Transport(TransportError::Connect(_))), + "{error:?}" + ); +} + +#[rstest] +#[tokio::test] +async fn a_timeout_after_sending_is_not_a_pre_send_decline( + request: ChatCompletionsRequest<'static>, +) { + let upstream = + upstream([anthropic_response(ANTHROPIC_MESSAGE).set_delay(Duration::from_secs(5))]).await; + let base = upstream.uri(); + + let error = chat_completions(ChatCompletionsRequest { + api_base: Some(&base), + timeout: Some(Duration::from_millis(100)), + ..request + }) + .await + .expect_err("the call times out"); + + assert!( + matches!(error, Error::Transport(TransportError::Network(_))), + "{error:?}" + ); +} + +#[rstest] +#[case::accepted("anthropic/claude-sonnet-4-5", None, hi(), json!({"max_tokens": 16}), None)] +#[case::accepted_bedrock("bedrock/anthropic.claude-sonnet-4-5", None, hi(), json!({}), None)] +#[case::unknown_provider( + "gpt-4o", + Some("openai"), + hi(), + json!({}), + Some("provider is not on the rust chat completions path") +)] +#[case::unreadable_messages( + "anthropic/claude-sonnet-4-5", + None, + json!("hi"), + json!({}), + Some("unreadable message list") +)] +#[case::empty_messages("anthropic/claude-sonnet-4-5", None, json!([]), json!({}), Some("empty message list"))] +#[case::streaming( + "anthropic/claude-sonnet-4-5", + None, + hi(), + json!({"stream": true}), + Some("streaming") +)] +#[case::unrecognized_param( + "anthropic/claude-sonnet-4-5", + None, + hi(), + json!({"not_a_param": 1}), + Some("unrecognized request parameter") +)] +#[case::opens_on_assistant_turn( + "anthropic/claude-sonnet-4-5", + None, + json!([{"role": "assistant", "content": "hi"}]), + json!({}), + Some("conversation does not open on a user turn") +)] +fn decline_reason_names_why_the_core_would_not_serve_the_request( + #[case] model: &str, + #[case] provider: Option<&str>, + #[case] messages: Value, + #[case] params: Value, + #[case] reason: Option<&str>, +) { + assert_eq!( + chat_completions_decline_reason(model, provider, messages, &object(params)), + reason + ); +} + +/// A request the decline check accepts must not be declined by the call itself. +#[rstest] +#[tokio::test] +async fn a_declined_request_fails_the_call_before_sending( + request: ChatCompletionsRequest<'static>, +) { + let upstream = upstream([anthropic_response(ANTHROPIC_MESSAGE)]).await; + let base = upstream.uri(); + + let error = chat_completions(ChatCompletionsRequest { + optional_params: object(json!({"stream": true})), + api_base: Some(&base), + ..request + }) + .await + .expect_err("streaming is declined"); + + assert_eq!(error, Error::Unsupported("streaming")); + assert!(received(&upstream).await.is_empty()); +} diff --git a/litellm-rust/crates/core/tests/messages.rs b/litellm-rust/crates/core/tests/messages.rs deleted file mode 100644 index 18af8a7d619..00000000000 --- a/litellm-rust/crates/core/tests/messages.rs +++ /dev/null @@ -1,471 +0,0 @@ -use std::{sync::Arc, time::Duration}; - -use futures_util::future::BoxFuture; -use litellm_core::messages::{ - Error, messages, - route::{LocalMessagesHost, MessagesCall, messages_machine}, - types::{MessagesRequest, MessagesShaping}, -}; -use litellm_secrets::{SecretValue, source::SecretSource}; -use serde_json::{Map, Value, json}; -use tokio::{ - io::{AsyncReadExt, AsyncWriteExt}, - net::{TcpListener, TcpStream}, -}; - -struct RecordingSecrets { - values: Vec<(&'static str, String)>, - fails: bool, - requested: std::sync::Mutex>, -} - -impl RecordingSecrets { - fn new(values: Vec<(&'static str, String)>, fails: bool) -> Self { - Self { - values, - fails, - requested: std::sync::Mutex::new(Vec::new()), - } - } -} - -impl SecretSource for RecordingSecrets { - fn get_secret_str<'a>( - &'a self, - name: &'a str, - ) -> BoxFuture<'a, Result, litellm_secrets::Error>> { - Box::pin(async move { - self.requested.lock().unwrap().push(name.to_string()); - if self.fails { - return Err(litellm_secrets::Error::ManagedSecretMissing); - } - Ok(self - .values - .iter() - .find(|(key, _)| *key == name) - .map(|(_, value)| SecretValue::new(value.clone()))) - }) - } -} - -fn secrets_call() -> MessagesCall { - let Value::Object(body) = json!({ - "model": "claude-sonnet-4-5", - "max_tokens": 16, - "messages": [{"role": "user", "content": "hi"}] - }) else { - unreachable!("literal object") - }; - MessagesCall { - model: "claude-sonnet-4-5".into(), - body, - api_key: None, - api_base: None, - custom_llm_provider: Some("anthropic".into()), - extra_headers: None, - provider_specific_header: None, - timeout: Some(Duration::from_secs(5)), - shaping: MessagesShaping::default(), - } -} - -#[tokio::test] -async fn route_surfaces_a_secret_manager_failure_before_the_call() { - let Err(error) = litellm_host::run::run( - messages_machine(Arc::new(RecordingSecrets::new(Vec::new(), true))), - &LocalMessagesHost::new(secrets_call()), - ) - .await - else { - panic!("a secret manager failure fails the call"); - }; - assert!( - matches!(&error, Error::Secret(source) if matches!(source.source_error(), litellm_secrets::Error::ManagedSecretMissing)), - "{error:?}" - ); -} - -async fn read_http_request(socket: &mut TcpStream) -> String { - let mut request = Vec::new(); - let mut buffer = [0_u8; 1024]; - let header_end = loop { - let n = socket.read(&mut buffer).await.expect("reads request"); - if n == 0 { - break request.len(); - } - request.extend_from_slice(&buffer[..n]); - if let Some(position) = request.windows(4).position(|window| window == b"\r\n\r\n") { - break position + 4; - } - }; - let headers = String::from_utf8_lossy(&request[..header_end]); - let content_length = headers - .lines() - .find_map(|line| { - let (name, value) = line.split_once(':')?; - name.eq_ignore_ascii_case("content-length") - .then(|| value.trim().parse::().ok()) - .flatten() - }) - .unwrap_or(0); - while request.len().saturating_sub(header_end) < content_length { - let n = socket.read(&mut buffer).await.expect("reads body"); - if n == 0 { - break; - } - request.extend_from_slice(&buffer[..n]); - } - String::from_utf8(request).expect("request is utf8") -} - -fn write_response(body: &str) -> String { - format!( - "HTTP/1.1 200 OK\r\ncontent-type: application/json\r\ncontent-length: {}\r\nconnection: close\r\n\r\n{}", - body.len(), - body - ) -} - -#[tokio::test] -async fn messages_round_trip_builds_azure_request_and_passes_response_through() { - let listener = TcpListener::bind("127.0.0.1:0").await.expect("binds"); - let addr = listener.local_addr().expect("addr"); - - let server = tokio::spawn(async move { - let (mut socket, _) = listener.accept().await.expect("accepts request"); - let request = read_http_request(&mut socket).await; - let response_body = r#"{"id":"msg_1","type":"message","role":"assistant","content":[{"type":"text","text":"hi"}],"model":"claude-sonnet-4-5","stop_reason":"end_turn","usage":{"input_tokens":1,"output_tokens":2}}"#; - socket - .write_all(write_response(response_body).as_bytes()) - .await - .expect("writes response"); - request - }); - - let response = messages(MessagesRequest { - model: "claude-sonnet-4-5", - body: json!({ - "model": "claude-sonnet-4-5", - "max_tokens": 1024, - "messages": [{ - "role": "user", - "content": [{ - "type": "text", - "text": "hi", - "cache_control": {"type": "ephemeral", "scope": "global"} - }] - }] - }), - api_key: Some("sk-azure"), - api_base: Some(&format!("http://{addr}")), - custom_llm_provider: Some("azure_ai"), - extra_headers: None, - provider_specific_header: None, - timeout: Some(Duration::from_secs(5)), - shaping: MessagesShaping::default(), - }) - .await - .expect("messages request succeeds"); - - assert_eq!(response.content[0]["text"], "hi"); - assert_eq!(response.stop_reason.as_deref(), Some("end_turn")); - - let request = server.await.expect("server task completes"); - let (head, body) = request.split_once("\r\n\r\n").expect("has body"); - assert!(head.starts_with("POST /anthropic/v1/messages "), "{head}"); - let head_lower = head.to_ascii_lowercase(); - assert!(head_lower.contains("x-api-key: sk-azure"), "{head}"); - assert!( - head_lower.contains("anthropic-version: 2023-06-01"), - "{head}" - ); - assert!( - head_lower.contains("content-type: application/json"), - "{head}" - ); - - let sent_body: Value = serde_json::from_str(body).expect("body is json"); - assert_eq!( - sent_body["messages"][0]["content"][0]["cache_control"], - json!({"type": "ephemeral"}) - ); -} - -#[tokio::test] -async fn messages_round_trip_builds_native_anthropic_request() { - let listener = TcpListener::bind("127.0.0.1:0").await.expect("binds"); - let addr = listener.local_addr().expect("addr"); - - let server = tokio::spawn(async move { - let (mut socket, _) = listener.accept().await.expect("accepts request"); - let request = read_http_request(&mut socket).await; - let response_body = r#"{"id":"msg_1","type":"message","role":"assistant","content":[{"type":"text","text":"hi"}],"model":"claude-sonnet-4-5","stop_reason":"end_turn","usage":{"input_tokens":1,"output_tokens":2}}"#; - socket - .write_all(write_response(response_body).as_bytes()) - .await - .expect("writes response"); - request - }); - - let response = messages(MessagesRequest { - model: "claude-sonnet-4-5", - body: json!({ - "model": "claude-sonnet-4-5", - "max_tokens": 1024, - "messages": [{"role": "user", "content": "hi"}] - }), - api_key: Some("sk-ant"), - api_base: Some(&format!("http://{addr}")), - custom_llm_provider: Some("anthropic"), - extra_headers: None, - provider_specific_header: None, - timeout: Some(Duration::from_secs(5)), - shaping: MessagesShaping::default(), - }) - .await - .expect("messages request succeeds"); - - assert_eq!(response.content[0]["text"], "hi"); - assert_eq!(response.stop_reason.as_deref(), Some("end_turn")); - - let request = server.await.expect("server task completes"); - let (head, _) = request.split_once("\r\n\r\n").expect("has body"); - assert!(head.starts_with("POST /v1/messages "), "{head}"); - let head_lower = head.to_ascii_lowercase(); - assert!(head_lower.contains("x-api-key: sk-ant"), "{head}"); - assert!( - head_lower.contains("anthropic-version: 2023-06-01"), - "{head}" - ); -} - -#[tokio::test] -async fn messages_does_not_duplicate_auth_when_x_api_key_supplied() { - let listener = TcpListener::bind("127.0.0.1:0").await.expect("binds"); - let addr = listener.local_addr().expect("addr"); - - let server = tokio::spawn(async move { - let (mut socket, _) = listener.accept().await.expect("accepts request"); - let request = read_http_request(&mut socket).await; - let response_body = - r#"{"id":"msg_2","type":"message","role":"assistant","content":[],"model":"m"}"#; - socket - .write_all(write_response(response_body).as_bytes()) - .await - .expect("writes response"); - request - }); - - let mut headers = Map::new(); - headers.insert( - "x-api-key".to_string(), - Value::String("from-python".to_string()), - ); - headers.insert( - "anthropic-beta".to_string(), - Value::String("token-efficient-tools-2025-02-19".to_string()), - ); - - messages(MessagesRequest { - model: "claude-sonnet-4-5", - body: json!({"model": "claude-sonnet-4-5", "max_tokens": 8, "messages": []}), - api_key: Some("rust-fallback-key"), - api_base: Some(&format!("http://{addr}")), - custom_llm_provider: Some("azure_ai"), - extra_headers: Some(headers), - provider_specific_header: None, - timeout: Some(Duration::from_secs(5)), - shaping: MessagesShaping::default(), - }) - .await - .expect("messages request succeeds"); - - let request = server.await.expect("server task completes"); - let head = request - .split_once("\r\n\r\n") - .expect("has body") - .0 - .to_ascii_lowercase(); - let api_key_count = head - .lines() - .filter(|line| line.starts_with("x-api-key:")) - .count(); - assert_eq!(api_key_count, 1, "{head}"); - assert!(head.contains("x-api-key: from-python"), "{head}"); - assert!( - head.contains("anthropic-beta: token-efficient-tools-2025-02-19"), - "{head}" - ); - assert!(!head.contains("rust-fallback-key"), "{head}"); -} - -#[tokio::test] -async fn messages_forwards_entra_id_bearer_without_requiring_api_key() { - let listener = TcpListener::bind("127.0.0.1:0").await.expect("binds"); - let addr = listener.local_addr().expect("addr"); - - let server = tokio::spawn(async move { - let (mut socket, _) = listener.accept().await.expect("accepts request"); - let request = read_http_request(&mut socket).await; - let response_body = - r#"{"id":"msg_3","type":"message","role":"assistant","content":[],"model":"m"}"#; - socket - .write_all(write_response(response_body).as_bytes()) - .await - .expect("writes response"); - request - }); - - let mut headers = Map::new(); - headers.insert( - "Authorization".to_string(), - Value::String("Bearer entra-token".to_string()), - ); - - messages(MessagesRequest { - model: "claude-sonnet-4-5", - body: json!({"model": "claude-sonnet-4-5", "max_tokens": 8, "messages": []}), - api_key: None, - api_base: Some(&format!("http://{addr}")), - custom_llm_provider: Some("azure_ai"), - extra_headers: Some(headers), - provider_specific_header: None, - timeout: Some(Duration::from_secs(5)), - shaping: MessagesShaping::default(), - }) - .await - .expect("entra id request succeeds without api key"); - - let request = server.await.expect("server task completes"); - let head = request - .split_once("\r\n\r\n") - .expect("has body") - .0 - .to_ascii_lowercase(); - assert!(head.contains("authorization: bearer entra-token"), "{head}"); - assert!(!head.contains("x-api-key"), "{head}"); -} - -#[tokio::test] -async fn messages_requires_auth_when_no_key_and_no_header() { - let err = messages(MessagesRequest { - model: "claude-sonnet-4-5", - body: json!({"model": "claude-sonnet-4-5", "max_tokens": 8, "messages": []}), - api_key: None, - api_base: Some("http://127.0.0.1:1"), - custom_llm_provider: Some("azure_ai"), - extra_headers: None, - provider_specific_header: None, - timeout: Some(Duration::from_millis(50)), - shaping: MessagesShaping::default(), - }) - .await - .expect_err("missing auth errors"); - - assert!(matches!(err, Error::Auth(_))); -} - -#[tokio::test] -async fn messages_ignores_malformed_authorization_and_uses_api_key() { - let listener = TcpListener::bind("127.0.0.1:0").await.expect("binds"); - let addr = listener.local_addr().expect("addr"); - - let server = tokio::spawn(async move { - let (mut socket, _) = listener.accept().await.expect("accepts request"); - let request = read_http_request(&mut socket).await; - let response_body = - r#"{"id":"msg_4","type":"message","role":"assistant","content":[],"model":"m"}"#; - socket - .write_all(write_response(response_body).as_bytes()) - .await - .expect("writes response"); - request - }); - - let mut headers = Map::new(); - headers.insert( - "Authorization".to_string(), - Value::String("Bearer ".to_string()), - ); - - messages(MessagesRequest { - model: "claude-sonnet-4-5", - body: json!({"model": "claude-sonnet-4-5", "max_tokens": 8, "messages": []}), - api_key: Some("sk-azure"), - api_base: Some(&format!("http://{addr}")), - custom_llm_provider: Some("azure_ai"), - extra_headers: Some(headers), - provider_specific_header: None, - timeout: Some(Duration::from_secs(5)), - shaping: MessagesShaping::default(), - }) - .await - .expect("falls back to api key"); - - let request = server.await.expect("server task completes"); - let head = request - .split_once("\r\n\r\n") - .expect("has body") - .0 - .to_ascii_lowercase(); - assert!(head.contains("x-api-key: sk-azure"), "{head}"); -} - -#[tokio::test] -async fn messages_maps_provider_error_status_to_http_error() { - let listener = TcpListener::bind("127.0.0.1:0").await.expect("binds"); - let addr = listener.local_addr().expect("addr"); - - tokio::spawn(async move { - let (mut socket, _) = listener.accept().await.expect("accepts request"); - let _ = read_http_request(&mut socket).await; - let body = "unauthorized"; - let response = format!( - "HTTP/1.1 401 Unauthorized\r\ncontent-length: {}\r\nconnection: close\r\n\r\n{}", - body.len(), - body - ); - socket - .write_all(response.as_bytes()) - .await - .expect("writes response"); - }); - - let err = messages(MessagesRequest { - model: "claude-sonnet-4-5", - body: json!({"model": "claude-sonnet-4-5", "max_tokens": 8, "messages": []}), - api_key: Some("sk-azure"), - api_base: Some(&format!("http://{addr}")), - custom_llm_provider: Some("azure_ai"), - extra_headers: None, - provider_specific_header: None, - timeout: Some(Duration::from_secs(5)), - shaping: MessagesShaping::default(), - }) - .await - .expect_err("provider error propagates"); - - assert!(matches!( - err, - Error::Transport(litellm_http::transport::Error::Http { status: 401, .. }) - )); -} - -#[tokio::test] -async fn messages_rejects_unsupported_provider() { - let err = messages(MessagesRequest { - model: "claude-3-5-sonnet", - body: json!({"model": "claude-3-5-sonnet", "max_tokens": 8, "messages": []}), - api_key: Some("sk"), - api_base: Some("http://127.0.0.1:1"), - custom_llm_provider: Some("openai"), - extra_headers: None, - provider_specific_header: None, - timeout: Some(Duration::from_millis(50)), - shaping: MessagesShaping::default(), - }) - .await - .expect_err("unsupported provider errors"); - - assert!(matches!(err, Error::InvalidProvider(provider) if provider == "openai")); -} diff --git a/litellm-rust/crates/core/tests/messages/main.rs b/litellm-rust/crates/core/tests/messages/main.rs new file mode 100644 index 00000000000..4e549bae309 --- /dev/null +++ b/litellm-rust/crates/core/tests/messages/main.rs @@ -0,0 +1,94 @@ +use std::{sync::Arc, time::Duration}; + +use litellm_core::messages::{ + Error, + route::{LocalMessagesHost, MessagesCall, MessagesOutput, messages_machine}, + types::MessagesShaping, +}; +use litellm_types::llms::anthropic_messages::anthropic_response::AnthropicMessagesResponse; +use rstest::fixture; +use serde_json::{Map, Value, json}; +use wiremock::ResponseTemplate; + +#[path = "../support/mod.rs"] +mod support; +use support::*; + +mod request; +mod response; +mod secrets; +mod stream; + +const MODEL: &str = "claude-sonnet-4-5"; + +fn object(value: Value) -> Map { + let Value::Object(map) = value else { + panic!("expected a json object, got {value}"); + }; + map +} + +fn message_body() -> Value { + json!({ + "id": "msg_1", + "type": "message", + "role": "assistant", + "content": [{"type": "text", "text": "hi"}], + "model": MODEL, + "stop_reason": "end_turn", + "usage": {"input_tokens": 1, "output_tokens": 2} + }) +} + +fn message_response() -> ResponseTemplate { + json_response(message_body()) +} + +/// A non-streaming call with nothing that would authenticate or route it, so each test +/// states the provider, credentials, and base it depends on. +#[fixture] +fn call() -> MessagesCall { + MessagesCall { + model: MODEL.into(), + body: object(json!({ + "model": MODEL, + "max_tokens": 16, + "messages": [{"role": "user", "content": "hi"}] + })), + api_key: None, + api_base: None, + custom_llm_provider: Some("anthropic".into()), + extra_headers: None, + provider_specific_header: None, + timeout: Some(Duration::from_secs(5)), + shaping: MessagesShaping::default(), + } +} + +fn headers<'a>(pairs: impl IntoIterator) -> Option> { + Some( + pairs + .into_iter() + .map(|(name, value)| (name.to_string(), Value::from(value))) + .collect(), + ) +} + +async fn run_with( + secrets: Arc, + call: MessagesCall, +) -> Result { + litellm_host::run::run(messages_machine(secrets), &LocalMessagesHost::new(call)).await +} + +/// Runs the route with a secret source that knows nothing, so no environment leaks in. +async fn run(call: MessagesCall) -> Result { + run_with(Arc::new(RecordingSecrets::empty()), call).await +} + +async fn run_message(call: MessagesCall) -> AnthropicMessagesResponse { + match run(call).await.expect("messages call succeeds") { + MessagesOutput::Message(message) => *message, + MessagesOutput::Streamed => panic!("a non-streaming call returned a stream"), + } +} diff --git a/litellm-rust/crates/core/tests/messages/request.rs b/litellm-rust/crates/core/tests/messages/request.rs new file mode 100644 index 00000000000..9353324d370 --- /dev/null +++ b/litellm-rust/crates/core/tests/messages/request.rs @@ -0,0 +1,269 @@ +use litellm_types::utils::{ProviderSpecificHeader, ProviderSpecificHeaders}; +use rstest::rstest; + +use super::*; + +#[rstest] +#[case::anthropic_key("anthropic", Some("sk-ant"), &[], ("x-api-key", "sk-ant"), &["authorization"])] +#[case::azure_key("azure_ai", Some("sk-azure"), &[], ("x-api-key", "sk-azure"), &["authorization"])] +#[case::caller_x_api_key_wins( + "azure_ai", + Some("rust-fallback-key"), + &[("x-api-key", "from-python")], + ("x-api-key", "from-python"), + &["authorization"] +)] +#[case::entra_bearer_without_key( + "azure_ai", + None, + &[("Authorization", "Bearer entra-token")], + ("authorization", "Bearer entra-token"), + &["x-api-key"] +)] +#[case::empty_bearer_falls_back_to_key( + "azure_ai", + Some("sk-azure"), + &[("Authorization", "Bearer ")], + ("x-api-key", "sk-azure"), + &[] +)] +#[case::anthropic_forwards_caller_authorization( + "anthropic", + Some("sk-ant"), + &[("Authorization", "Bearer caller")], + ("authorization", "Bearer caller"), + &["x-api-key"] +)] +#[case::anthropic_oauth_key_becomes_bearer( + "anthropic", + Some("sk-ant-oat01-token"), + &[], + ("authorization", "Bearer sk-ant-oat01-token"), + &["x-api-key"] +)] +#[tokio::test] +async fn credentials_become_exactly_one_auth_header( + call: MessagesCall, + #[case] provider: &str, + #[case] api_key: Option<&str>, + #[case] extra_headers: &[(&str, &str)], + #[case] expected: (&str, &str), + #[case] absent: &[&str], +) { + let upstream = upstream([message_response()]).await; + + run_message(MessagesCall { + custom_llm_provider: Some(provider.into()), + api_key: api_key.map(Into::into), + api_base: Some(upstream.uri()), + extra_headers: headers(extra_headers.iter().copied()), + ..call + }) + .await; + + let request = only_request(&upstream).await; + let (name, value) = expected; + assert_eq!(request.header_values(name), [value]); + for name in absent { + assert_eq!(request.header(name), None, "{name} must not be sent"); + } +} + +#[rstest] +#[case::anthropic("anthropic")] +#[case::azure_ai("azure_ai")] +#[tokio::test] +async fn a_call_without_credentials_fails_before_sending( + call: MessagesCall, + #[case] provider: &str, +) { + let upstream = upstream([message_response()]).await; + + let error = run(MessagesCall { + custom_llm_provider: Some(provider.into()), + api_base: Some(upstream.uri()), + ..call + }) + .await + .err() + .expect("a call without credentials fails"); + + assert!( + matches!( + error, + Error::Auth(litellm_auth::Error::MissingApiKey { .. }) + ), + "{error:?}" + ); + assert!(received(&upstream).await.is_empty()); +} + +#[rstest] +#[case::anthropic(MODEL, Some("anthropic"), "", "/v1/messages")] +#[case::anthropic_base_with_trailing_slash(MODEL, Some("anthropic"), "/", "/v1/messages")] +#[case::anthropic_base_with_the_messages_path( + MODEL, + Some("anthropic"), + "/v1/messages", + "/v1/messages" +)] +#[case::azure_ai(MODEL, Some("azure_ai"), "", "/anthropic/v1/messages")] +#[case::provider_from_model_prefix("anthropic/claude-sonnet-4-5", None, "", "/v1/messages")] +#[tokio::test] +async fn each_provider_posts_to_its_messages_endpoint( + call: MessagesCall, + #[case] model: &str, + #[case] provider: Option<&str>, + #[case] base_suffix: &str, + #[case] path: &str, +) { + let upstream = upstream([message_response()]).await; + + run_message(MessagesCall { + model: model.into(), + custom_llm_provider: provider.map(Into::into), + api_key: Some("sk".into()), + api_base: Some(format!("{}{base_suffix}", upstream.uri())), + ..call + }) + .await; + + let request = only_request(&upstream).await; + assert_eq!(request.method.as_str(), "POST"); + assert_eq!(request.url.path(), path); + assert_eq!(request.json()["model"], MODEL); + assert_eq!(request.header("anthropic-version"), Some("2023-06-01")); + assert_eq!(request.header("content-type"), Some("application/json")); +} + +#[rstest] +#[case::unknown_provider(MODEL, Some("openai"), "openai")] +#[case::unresolvable_model( + "no-such-model", + None, + "unable to resolve custom_llm_provider for messages request" +)] +#[tokio::test] +async fn unsupported_providers_are_rejected_before_sending( + call: MessagesCall, + #[case] model: &str, + #[case] provider: Option<&str>, + #[case] reported: &str, +) { + let error = run(MessagesCall { + model: model.into(), + custom_llm_provider: provider.map(Into::into), + api_key: Some("sk".into()), + api_base: Some(UNREACHABLE_BASE.into()), + ..call + }) + .await + .err() + .expect("unsupported provider errors"); + + assert_eq!(error, Error::InvalidProvider(reported.into())); +} + +#[rstest] +#[tokio::test] +async fn caller_headers_and_provider_scoped_headers_are_forwarded(call: MessagesCall) { + let upstream = upstream([message_response()]).await; + let scoped = |provider: &str, value: &str| ProviderSpecificHeader { + custom_llm_provider: provider.into(), + extra_headers: object(json!({"x-scoped": value})), + }; + + run_message(MessagesCall { + api_key: Some("sk".into()), + api_base: Some(upstream.uri()), + extra_headers: headers([("anthropic-beta", "token-efficient-tools-2025-02-19")]), + provider_specific_header: Some(ProviderSpecificHeaders::Many(vec![ + scoped("bedrock", "other-provider"), + scoped("azure_ai, anthropic", "this-provider"), + ])), + ..call + }) + .await; + + let request = only_request(&upstream).await; + assert_eq!( + request.header("anthropic-beta"), + Some("token-efficient-tools-2025-02-19") + ); + assert_eq!(request.header_values("x-scoped"), ["this-provider"]); +} + +#[rstest] +#[tokio::test] +async fn azure_strips_the_cache_control_scope_anthropic_rejects(call: MessagesCall) { + let upstream = upstream([message_response()]).await; + + run_message(MessagesCall { + custom_llm_provider: Some("azure_ai".into()), + api_key: Some("sk-azure".into()), + api_base: Some(upstream.uri()), + body: object(json!({ + "model": MODEL, + "max_tokens": 16, + "messages": [{ + "role": "user", + "content": [{ + "type": "text", + "text": "hi", + "cache_control": {"type": "ephemeral", "scope": "global"} + }] + }] + })), + ..call + }) + .await; + + assert_eq!( + only_request(&upstream).await.json()["messages"][0]["content"][0]["cache_control"], + json!({"type": "ephemeral"}) + ); +} + +#[rstest] +#[tokio::test] +async fn additional_drop_params_remove_fields_before_sending(call: MessagesCall) { + let upstream = upstream([message_response()]).await; + let mut body = call.body.clone(); + body.insert("temperature".into(), json!(0.5)); + body.insert("top_k".into(), json!(3)); + + run_message(MessagesCall { + api_key: Some("sk".into()), + api_base: Some(upstream.uri()), + body, + shaping: MessagesShaping { + additional_drop_params: vec!["temperature".into()], + ..MessagesShaping::default() + }, + ..call + }) + .await; + + let sent = only_request(&upstream).await.json(); + assert_eq!(sent.get("temperature"), None); + assert_eq!(sent["top_k"], 3); +} + +#[rstest] +#[case::anthropic_streams(MODEL, Some("anthropic"), true, true)] +#[case::anthropic_prefix_streams("anthropic/claude-sonnet-4-5", None, true, true)] +#[case::azure_without_stream(MODEL, Some("azure_ai"), false, true)] +#[case::azure_stream(MODEL, Some("azure_ai"), true, false)] +#[case::other_provider(MODEL, Some("openai"), false, false)] +#[case::unresolvable_model("no-such-model", None, false, false)] +fn supports_matches_what_the_route_can_serve( + #[case] model: &str, + #[case] provider: Option<&str>, + #[case] stream: bool, + #[case] supported: bool, +) { + assert_eq!( + litellm_core::messages::route::supports(model, provider, stream), + supported + ); +} diff --git a/litellm-rust/crates/core/tests/messages/response.rs b/litellm-rust/crates/core/tests/messages/response.rs new file mode 100644 index 00000000000..38a18c415ba --- /dev/null +++ b/litellm-rust/crates/core/tests/messages/response.rs @@ -0,0 +1,136 @@ +use litellm_core::messages::{messages, types::MessagesRequest}; +use litellm_http::transport::Error as TransportError; +use rstest::rstest; + +use super::*; + +#[rstest] +#[tokio::test] +async fn the_provider_message_is_returned(call: MessagesCall) { + let upstream = upstream([message_response()]).await; + + let message = run_message(MessagesCall { + api_key: Some("sk".into()), + api_base: Some(upstream.uri()), + ..call + }) + .await; + + assert_eq!(message.id, "msg_1"); + assert_eq!(message.content, [json!({"type": "text", "text": "hi"})]); + assert_eq!(message.stop_reason.as_deref(), Some("end_turn")); +} + +#[rstest] +#[case::bad_request(400)] +#[case::unauthorized(401)] +#[case::rate_limited(429)] +#[case::server_error(500)] +#[case::overloaded(529)] +#[tokio::test] +async fn an_upstream_error_keeps_its_status_and_body(call: MessagesCall, #[case] status: u16) { + let upstream = + upstream([ResponseTemplate::new(status).set_body_string("upstream said no")]).await; + + let error = run(MessagesCall { + api_key: Some("sk".into()), + api_base: Some(upstream.uri()), + ..call + }) + .await + .err() + .expect("upstream error propagates"); + + assert_eq!( + error, + Error::Transport(TransportError::Http { + status, + body: "upstream said no".into() + }) + ); +} + +#[rstest] +#[case::not_json(ResponseTemplate::new(200).set_body_string("not json"))] +#[case::not_a_message(json_response(json!({"unexpected": true})))] +#[tokio::test] +async fn an_unreadable_success_body_is_an_invalid_response( + call: MessagesCall, + #[case] response: ResponseTemplate, +) { + let upstream = upstream([response]).await; + + let error = run(MessagesCall { + api_key: Some("sk".into()), + api_base: Some(upstream.uri()), + ..call + }) + .await + .err() + .expect("an unreadable body fails"); + + assert!(error.is_response(), "{error:?}"); +} + +#[rstest] +#[tokio::test] +async fn a_provider_slower_than_the_timeout_fails_the_call(call: MessagesCall) { + let upstream = upstream([message_response().set_delay(Duration::from_secs(5))]).await; + + let error = run(MessagesCall { + api_key: Some("sk".into()), + api_base: Some(upstream.uri()), + timeout: Some(Duration::from_millis(100)), + ..call + }) + .await + .err() + .expect("the call times out"); + + assert!(matches!(error, Error::Transport(_)), "{error:?}"); +} + +fn facade_request(body: Value, api_base: &str) -> MessagesRequest<'_> { + MessagesRequest { + model: MODEL, + body, + api_key: Some("sk-ant"), + api_base: Some(api_base), + custom_llm_provider: Some("anthropic"), + extra_headers: None, + provider_specific_header: None, + timeout: Some(Duration::from_secs(5)), + shaping: MessagesShaping::default(), + } +} + +#[tokio::test] +async fn the_facade_runs_the_route_in_process() { + let upstream = upstream([message_response()]).await; + let base = upstream.uri(); + + let message = messages(facade_request( + json!({"model": MODEL, "max_tokens": 16, "messages": [{"role": "user", "content": "hi"}]}), + &base, + )) + .await + .expect("messages request succeeds"); + + assert_eq!(message.id, "msg_1"); + assert_eq!( + only_request(&upstream).await.header("x-api-key"), + Some("sk-ant") + ); +} + +#[tokio::test] +async fn the_facade_rejects_a_body_that_is_not_an_object() { + let error = messages(facade_request(json!([]), UNREACHABLE_BASE)) + .await + .expect_err("a non-object body is rejected"); + + assert_eq!( + error, + Error::InvalidRequest("messages body must be an object".into()) + ); +} diff --git a/litellm-rust/crates/core/tests/messages/secrets.rs b/litellm-rust/crates/core/tests/messages/secrets.rs new file mode 100644 index 00000000000..419b6d6c753 --- /dev/null +++ b/litellm-rust/crates/core/tests/messages/secrets.rs @@ -0,0 +1,94 @@ +use litellm_llms::{ + anthropic::experimental_pass_through::messages::transformation::ANTHROPIC_MESSAGES_CONFIG, + azure_ai::anthropic::messages_transformation::AZURE_ANTHROPIC_MESSAGES_CONFIG, + base_llm::anthropic_messages::transformation::BaseAnthropicMessagesConfig, +}; +use rstest::rstest; + +use super::*; + +#[rstest] +#[case::anthropic("anthropic", &ANTHROPIC_MESSAGES_CONFIG, "ANTHROPIC_API_KEY", "ANTHROPIC_BASE_URL", "/v1/messages")] +#[case::azure_ai("azure_ai", &AZURE_ANTHROPIC_MESSAGES_CONFIG, "AZURE_API_KEY", "AZURE_API_BASE", "/anthropic/v1/messages")] +#[tokio::test] +async fn the_credential_and_base_come_from_the_secret_source( + call: MessagesCall, + #[case] provider: &str, + #[case] config: &dyn BaseAnthropicMessagesConfig, + #[case] key_name: &str, + #[case] base_name: &str, + #[case] path: &str, +) { + let upstream = upstream([message_response()]).await; + let base = upstream.uri(); + let secrets = Arc::new(RecordingSecrets::new([ + (key_name, "sk-from-manager"), + (base_name, base.as_str()), + ])); + + let output = run_with( + secrets.clone(), + MessagesCall { + custom_llm_provider: Some(provider.into()), + ..call + }, + ) + .await + .expect("messages call succeeds"); + + assert!(matches!(output, MessagesOutput::Message(_))); + let request = only_request(&upstream).await; + assert_eq!(request.url.path(), path); + assert_eq!(request.header("x-api-key"), Some("sk-from-manager")); + assert_eq!(secrets.requested(), config.secret_names()); +} + +#[rstest] +#[tokio::test] +async fn call_arguments_win_over_the_secret_source(call: MessagesCall) { + let upstream = upstream([message_response()]).await; + let secrets = Arc::new(RecordingSecrets::new([ + ("ANTHROPIC_API_KEY", "sk-from-manager"), + ("ANTHROPIC_BASE_URL", UNREACHABLE_BASE), + ])); + + run_with( + secrets, + MessagesCall { + api_key: Some("sk-from-call".into()), + api_base: Some(upstream.uri()), + ..call + }, + ) + .await + .expect("messages call succeeds"); + + assert_eq!( + only_request(&upstream).await.header("x-api-key"), + Some("sk-from-call") + ); +} + +#[rstest] +#[tokio::test] +async fn a_secret_manager_failure_fails_the_call_before_sending(call: MessagesCall) { + let upstream = upstream([message_response()]).await; + + let error = run_with( + Arc::new(RecordingSecrets::failing()), + MessagesCall { + api_key: Some("sk".into()), + api_base: Some(upstream.uri()), + ..call + }, + ) + .await + .err() + .expect("a secret manager failure fails the call"); + + assert!( + matches!(&error, Error::Secret(source) if matches!(source.source_error(), litellm_secrets::Error::ManagedSecretMissing)), + "{error:?}" + ); + assert!(received(&upstream).await.is_empty()); +} diff --git a/litellm-rust/crates/core/tests/messages/stream.rs b/litellm-rust/crates/core/tests/messages/stream.rs new file mode 100644 index 00000000000..ea23a9e8e38 --- /dev/null +++ b/litellm-rust/crates/core/tests/messages/stream.rs @@ -0,0 +1,163 @@ +use std::{convert::Infallible, sync::Mutex}; + +use bytes::Bytes; +use litellm_core::messages::route::Messages; +use litellm_host::host::{Demand, Host}; +use rstest::rstest; + +use super::*; + +const SSE_BODY: &str = "event: message_start\ndata: {\"type\":\"message_start\"}\n\nevent: message_stop\ndata: {\"type\":\"message_stop\"}\n\n"; + +enum Seen { + Open, + Deliver(Bytes), +} + +/// Projects like `LocalMessagesHost`, records every stream op in the order the route +/// performs it, and detaches after `detach_after` ops. +struct RecordingStreamHost { + call: LocalMessagesHost, + detach_after: usize, + seen: Mutex>, +} + +impl RecordingStreamHost { + fn new(call: MessagesCall, detach_after: usize) -> Self { + Self { + call: LocalMessagesHost::new(call), + detach_after, + seen: Mutex::new(Vec::new()), + } + } + + fn record(&self, op: Seen) -> Demand { + let mut seen = self.seen.lock().unwrap(); + seen.push(op); + match seen.len() < self.detach_after { + true => Demand::More, + false => Demand::Detached, + } + } +} + +impl Host for RecordingStreamHost { + async fn project(&self) -> Result { + self.call.project().await + } + + async fn custom_op(&self, op: Infallible) -> Result<(), Error> { + match op {} + } + + async fn open(&self, (): ()) -> Result { + Ok(self.record(Seen::Open)) + } + + async fn deliver(&self, chunk: Bytes) -> Result { + Ok(self.record(Seen::Deliver(chunk))) + } +} + +fn streaming(call: MessagesCall, api_base: String) -> MessagesCall { + let mut body = call.body.clone(); + body.insert("stream".into(), json!(true)); + MessagesCall { + api_key: Some("sk-ant".into()), + api_base: Some(api_base), + body, + ..call + } +} + +fn sse_response() -> ResponseTemplate { + ResponseTemplate::new(200).set_body_raw(SSE_BODY, "text/event-stream") +} + +async fn stream_through(host: &RecordingStreamHost) -> Result { + litellm_host::run::run(messages_machine(Arc::new(RecordingSecrets::empty())), host).await +} + +#[rstest] +#[tokio::test] +async fn the_stream_opens_once_before_relaying_the_upstream_body(call: MessagesCall) { + let upstream = upstream([sse_response()]).await; + let host = RecordingStreamHost::new(streaming(call, upstream.uri()), usize::MAX); + + let outcome = stream_through(&host).await.expect("streamed call succeeds"); + + assert!(matches!(outcome, MessagesOutput::Streamed)); + let seen = host.seen.into_inner().unwrap(); + let [Seen::Open, chunks @ ..] = seen.as_slice() else { + panic!("the stream opens before any chunk is delivered"); + }; + let delivered: Vec = chunks + .iter() + .flat_map(|step| match step { + Seen::Deliver(chunk) => chunk.to_vec(), + Seen::Open => panic!("the stream opens exactly once"), + }) + .collect(); + assert_eq!(delivered, SSE_BODY.as_bytes()); +} + +#[rstest] +#[case::at_open(1)] +#[case::after_the_first_chunk(2)] +#[tokio::test] +async fn a_detached_caller_receives_nothing_more(call: MessagesCall, #[case] detach_after: usize) { + let upstream = upstream([sse_response()]).await; + let host = RecordingStreamHost::new(streaming(call, upstream.uri()), detach_after); + + let outcome = stream_through(&host) + .await + .expect("a detached stream still completes"); + + assert!(matches!(outcome, MessagesOutput::Streamed)); + assert_eq!(host.seen.into_inner().unwrap().len(), detach_after); +} + +#[rstest] +#[tokio::test] +async fn an_upstream_error_fails_the_call_without_opening_the_stream(call: MessagesCall) { + let upstream = upstream([ResponseTemplate::new(429).set_body_string("slow down")]).await; + let host = RecordingStreamHost::new(streaming(call, upstream.uri()), usize::MAX); + + let error = stream_through(&host) + .await + .err() + .expect("upstream error propagates"); + + assert!( + matches!( + error, + Error::Transport(litellm_http::transport::Error::Http { status: 429, .. }) + ), + "{error:?}" + ); + assert!(host.seen.into_inner().unwrap().is_empty()); +} + +#[rstest] +#[tokio::test] +async fn streaming_is_refused_for_providers_that_cannot_stream(call: MessagesCall) { + let upstream = upstream([sse_response()]).await; + let host = RecordingStreamHost::new( + MessagesCall { + custom_llm_provider: Some("azure_ai".into()), + ..streaming(call, upstream.uri()) + }, + usize::MAX, + ); + + let error = stream_through(&host) + .await + .err() + .expect("azure streaming is refused"); + + assert_eq!( + error, + Error::Unsupported("streaming messages for this provider") + ); + assert!(received(&upstream).await.is_empty()); +} diff --git a/litellm-rust/crates/core/tests/ocr/aws_textract.rs b/litellm-rust/crates/core/tests/ocr/aws_textract.rs new file mode 100644 index 00000000000..790e16a95ec --- /dev/null +++ b/litellm-rust/crates/core/tests/ocr/aws_textract.rs @@ -0,0 +1,173 @@ +use std::{collections::BTreeMap, time::SystemTime}; + +use litellm_auth_aws::{Credentials, aws_signature_headers, sign_post}; +use rstest::rstest; +use time::{PrimitiveDateTime, format_description}; +use wiremock::Request; + +use super::*; + +const ACCESS_KEY_ID: &str = "AKIDEXAMPLE"; +const SECRET_ACCESS_KEY: &str = "wJalrXUtnFEMI/K7MDENG+bPxRfiCYEXAMPLEKEY"; +const DETECT: &str = "aws_textract/detect-document-text"; +const ANALYZE: &str = "aws_textract/analyze-document"; + +fn textract_request(model: &str, base: &str) -> LiteLLMOcrRequest { + ocr_request_with_document( + model, + &format!("{base}/"), + json!({"type": "image_url", "image_url": "data:image/png;base64,b3JpZ2luYWw="}), + json!({ + "aws_access_key_id": ACCESS_KEY_ID, + "aws_secret_access_key": SECRET_ACCESS_KEY, + "aws_region_name": "eu-west-1" + }), + ) +} + +fn textract_response() -> ResponseTemplate { + json_response(json!({ + "DocumentMetadata": {"Pages": 1}, + "Blocks": [{"BlockType": "PAGE"}, {"BlockType": "LINE", "Text": "Invoice 12345"}] + })) +} + +/// Recomputes SigV4 over the request the upstream received, at the time the client claimed. +fn expected_authorization(url: &str, sent: &Request) -> String { + let format = + format_description::parse_borrowed::<2>("[year][month][day]T[hour][minute][second]Z") + .unwrap(); + let signed_at: SystemTime = + PrimitiveDateTime::parse(sent.header("x-amz-date").unwrap(), &format) + .unwrap() + .assume_utc() + .into(); + let headers: BTreeMap = ["content-type", "x-amz-target"] + .into_iter() + .map(|name| (name.to_string(), sent.header(name).unwrap().to_string())) + .collect(); + sign_post( + url, + &sent.body, + &aws_signature_headers(&headers), + "eu-west-1", + "textract", + &Credentials::new(ACCESS_KEY_ID, SECRET_ACCESS_KEY, None, None, "test"), + signed_at, + ) + .unwrap()["Authorization"] + .clone() +} + +/// The recorded URL names wiremock's host, not the address the client signed for. +fn assert_signed(upstream: &MockServer, sent: &Request) { + let url = format!("{}/", upstream.uri()); + assert_eq!( + sent.header("authorization"), + Some(expected_authorization(&url, sent).as_str()) + ); +} + +#[tokio::test] +async fn detect_document_text_is_signed_and_lines_become_the_page() { + let upstream = upstream([textract_response()]).await; + + let response = perform_with(LocalOcrHost::new(textract_request(DETECT, &upstream.uri()))) + .await + .unwrap(); + + let sent = only_request(&upstream).await; + assert_eq!( + sent.header("x-amz-target"), + Some("Textract.DetectDocumentText") + ); + assert_eq!( + sent.header("content-type"), + Some("application/x-amz-json-1.1") + ); + assert_eq!(sent.json(), json!({"Document": {"Bytes": "b3JpZ2luYWw="}})); + assert_signed(&upstream, &sent); + assert_eq!(response.pages[0].markdown, "Invoice 12345"); + assert_eq!(response.usage_info.unwrap().pages_processed, Some(1)); +} + +#[tokio::test] +async fn a_body_rewritten_by_before_send_is_what_gets_signed_and_sent() { + let upstream = upstream([textract_response()]).await; + let host = LocalOcrHost::new(textract_request(DETECT, &upstream.uri())).with_before_send( + |mut wire, _| { + assert!( + !wire + .headers + .iter() + .any(|(name, _)| name.eq_ignore_ascii_case("authorization")), + "the hook ran after signing" + ); + wire.body["Document"]["Bytes"] = Value::from("cmVkYWN0ZWQ="); + Ok(wire) + }, + ); + + perform_with(host).await.unwrap(); + + let sent = only_request(&upstream).await; + assert_eq!(sent.json(), json!({"Document": {"Bytes": "cmVkYWN0ZWQ="}})); + assert_signed(&upstream, &sent); +} + +#[tokio::test] +async fn analyze_document_asks_for_layout_and_tables_and_returns_markdown() { + let upstream = upstream([json_response(json!({ + "DocumentMetadata": {"Pages": 1}, + "Blocks": [ + {"Id": "l1", "BlockType": "LINE", "Text": "Quarterly Report"}, + {"Id": "t", "BlockType": "LAYOUT_TITLE", + "Relationships": [{"Type": "CHILD", "Ids": ["l1"]}]} + ] + }))]) + .await; + + let response = perform_with(LocalOcrHost::new(textract_request( + ANALYZE, + &upstream.uri(), + ))) + .await + .unwrap(); + + let sent = only_request(&upstream).await; + assert_eq!( + sent.header("x-amz-target"), + Some("Textract.AnalyzeDocument") + ); + assert_eq!(sent.json()["FeatureTypes"], json!(["LAYOUT", "TABLES"])); + assert_signed(&upstream, &sent); + assert_eq!(response.pages[0].markdown, "# Quarterly Report"); +} + +#[rstest] +#[case::detect(DETECT)] +#[case::analyze(ANALYZE)] +#[tokio::test] +async fn a_multi_page_rejection_reaches_the_caller_with_the_single_page_limit(#[case] model: &str) { + let upstream = upstream([status_response( + 400, + json!({ + "__type": "UnsupportedDocumentException", + "Message": "Request has unsupported document format" + }), + )]) + .await; + + let error = perform_with(LocalOcrHost::new(textract_request(model, &upstream.uri()))) + .await + .unwrap_err(); + + let Error::Provider { status, body, .. } = error else { + panic!("expected a provider error, got {error:?}"); + }; + assert_eq!(status, 400); + assert!( + body.contains("multi-page documents are not supported"), + "{body}" + ); +} diff --git a/litellm-rust/crates/core/tests/ocr/azure_ai.rs b/litellm-rust/crates/core/tests/ocr/azure_ai.rs new file mode 100644 index 00000000000..0d6ba024c15 --- /dev/null +++ b/litellm-rust/crates/core/tests/ocr/azure_ai.rs @@ -0,0 +1,270 @@ +use std::sync::{ + Arc, + atomic::{AtomicUsize, Ordering}, +}; + +use litellm_auth::{ + ResolvedCredential, SecretValue, TokenFuture, TokenProvider, TokenProviderHandle, +}; +use rstest::rstest; + +use super::*; + +#[derive(Debug)] +struct CountingToken { + token: fn(usize) -> String, + calls: AtomicUsize, +} + +impl CountingToken { + fn new(token: fn(usize) -> String) -> Arc { + Arc::new(Self { + token, + calls: AtomicUsize::new(0), + }) + } + + fn calls(&self) -> usize { + self.calls.load(Ordering::SeqCst) + } +} + +impl TokenProvider for CountingToken { + fn acquire(&self) -> TokenFuture<'_> { + let call = self.calls.fetch_add(1, Ordering::SeqCst) + 1; + let token = SecretValue::new((self.token)(call)); + Box::pin(async move { + Ok(ResolvedCredential::AccessToken { + token, + expires_on: None, + }) + }) + } +} + +fn numbered_token(call: usize) -> String { + format!("callback-{call}") +} + +fn azure_request( + provider: &Arc, + api_base: Option<&str>, + api_key: Option<&str>, + extra_headers: Value, + optional_params: Value, +) -> LiteLLMOcrRequest { + let wire = serde_json::from_value(json!({ + "model": "azure_ai/mistral-ocr-latest", + "document": {"type": "document_url", "document_url": INLINE_PDF}, + "api_key": api_key, + "api_base": api_base, + "custom_llm_provider": null, + "extra_headers": extra_headers, + "optional_params": optional_params, + "timeout_seconds": 2.0 + })) + .unwrap(); + let mut request = decode_request(wire).unwrap(); + request.azure_ad_token_provider = Some(TokenProviderHandle::new(provider.clone())); + request +} + +fn ocr_page() -> ResponseTemplate { + json_response(json!({"pages": [{"index": 0, "markdown": "hello"}]})) +} + +#[tokio::test] +async fn mistral_on_azure_sends_the_prepared_bearer_and_the_mistral_body() { + let upstream = upstream([json_response(json!({ + "pages": [{"index": 0, "markdown": "hello"}], + "usage_info": {"pages_processed": 1} + }))]) + .await; + let request = with_headers( + without_api_key(ocr_request( + "azure_ai/model", + &upstream.uri(), + json!({"include_image_base64": true}), + )), + &[("Authorization", "Bearer python-prepared-token")], + ); + + let result = perform(request).await.unwrap(); + + assert_eq!(result.pages[0].markdown, "hello"); + let sent = only_request(&upstream).await; + assert_eq!(sent.url.path(), "/providers/mistral/azure/ocr"); + assert_eq!( + sent.header("authorization"), + Some("Bearer python-prepared-token") + ); + assert_eq!( + sent.json(), + json!({ + "model": "model", + "document": {"type": "document_url", "document_url": INLINE_PDF}, + "include_image_base64": true + }) + ); +} + +#[tokio::test] +async fn a_static_entra_token_becomes_the_bearer() { + let upstream = upstream([pages_response()]).await; + let request = without_api_key(ocr_request( + "azure_ai/model", + &upstream.uri(), + json!({"azure_ad_token": "rust-owned-token"}), + )); + + perform(request).await.unwrap(); + + assert_eq!( + only_request(&upstream).await.header("authorization"), + Some("Bearer rust-owned-token") + ); +} + +#[tokio::test] +async fn a_guardrail_that_swaps_in_a_remote_document_is_rejected() { + let host = LocalOcrHost::new(ocr_request("azure_ai/model", UNREACHABLE_BASE, json!({}))) + .with_before_send(|mut wire, _| { + wire.body["document"] = json!({ + "type": "document_url", + "document_url": "https://example.com/not-inline.pdf" + }); + Ok(wire) + }); + + let error = perform_with(host).await.unwrap_err(); + + assert!(error.to_string().contains("data URI"), "{error}"); +} + +#[tokio::test] +async fn the_token_provider_is_the_bearer_and_is_acquired_for_each_request() { + let provider = CountingToken::new(numbered_token); + let upstream = upstream([ocr_page(), ocr_page()]).await; + let base = upstream.uri(); + + for _ in 0..2 { + perform(azure_request( + &provider, + Some(&base), + None, + Value::Null, + json!({}), + )) + .await + .unwrap(); + } + + assert_eq!(provider.calls(), 2); + let authorizations: Vec = received(&upstream) + .await + .iter() + .map(|request| { + request + .header("authorization") + .unwrap_or_default() + .to_string() + }) + .collect(); + assert_eq!(authorizations, ["Bearer callback-1", "Bearer callback-2"]); +} + +#[rstest] +#[case::api_key_skips_provider(Some("resource-key"), Value::Null, json!({}), "Bearer resource-key", 0)] +#[case::provider_beats_static_token( + None, + Value::Null, + json!({"azure_ad_token": "static-token"}), + "Bearer callback-1", + 1 +)] +#[case::header_wins_on_the_wire_but_provider_still_runs( + None, + json!({"Authorization": "Bearer override"}), + json!({}), + "Bearer override", + 1 +)] +#[tokio::test] +async fn credential_precedence( + #[case] api_key: Option<&str>, + #[case] extra_headers: Value, + #[case] optional_params: Value, + #[case] expected_authorization: &str, + #[case] expected_calls: usize, +) { + let provider = CountingToken::new(numbered_token); + let upstream = upstream([ocr_page()]).await; + + perform(azure_request( + &provider, + Some(&upstream.uri()), + api_key, + extra_headers, + optional_params, + )) + .await + .unwrap(); + + assert_eq!(provider.calls(), expected_calls); + assert_eq!( + only_request(&upstream).await.header_values("authorization"), + [expected_authorization] + ); +} + +#[rstest] +#[case::missing_api_base( + false, + json!({}), + numbered_token, + |error: &Error| matches!(error, Error::Auth(litellm_auth::Error::MissingApiBase { + provider: "Azure AI", + environment_variable: "AZURE_AI_API_BASE", + })), + 0 +)] +#[case::unsupported_oidc_reference( + true, + json!({"azure_ad_token": "oidc/assertion", "client_id": "client", "tenant_id": "tenant"}), + numbered_token, + |error: &Error| matches!(error, Error::Auth(litellm_auth::Error::UnsupportedOidcReference)), + 0 +)] +#[case::empty_provider_token_ignores_static_token( + true, + json!({"azure_ad_token": "static-token"}), + |_| String::new(), + |error: &Error| matches!(error, Error::MissingAzureAiCredentials), + 1 +)] +#[tokio::test] +async fn credential_failures_send_no_provider_request( + #[case] with_api_base: bool, + #[case] optional_params: Value, + #[case] token: fn(usize) -> String, + #[case] expected: fn(&Error) -> bool, + #[case] expected_calls: usize, +) { + let provider = CountingToken::new(token); + let upstream = upstream([ocr_page()]).await; + let base = upstream.uri(); + + let error = perform(azure_request( + &provider, + with_api_base.then_some(base.as_str()), + None, + Value::Null, + optional_params, + )) + .await + .unwrap_err(); + + assert!(expected(&error), "unexpected error: {error:?}"); + assert_eq!(provider.calls(), expected_calls); + assert!(received(&upstream).await.is_empty()); +} diff --git a/litellm-rust/crates/core/tests/ocr/azure_document_intelligence.rs b/litellm-rust/crates/core/tests/ocr/azure_document_intelligence.rs new file mode 100644 index 00000000000..1921176d158 --- /dev/null +++ b/litellm-rust/crates/core/tests/ocr/azure_document_intelligence.rs @@ -0,0 +1,441 @@ +use std::{ + sync::{Arc, Mutex}, + time::Duration, +}; + +use litellm_host::event::{CallEvent, MachineEvent}; +use litellm_llms::base_llm::ocr::settings::OcrSettings; +use rstest::rstest; + +use super::*; + +const MODEL: &str = "azure_ai/doc-intelligence/prebuilt-read"; + +fn read_request(base: &str, options: Value) -> LiteLLMOcrRequest { + ocr_request(MODEL, base, options) +} + +#[tokio::test] +async fn pages_features_and_extra_options_map_to_the_analyze_call() { + let upstream = upstream([json_response(json!({ + "status": "succeeded", + "analyzeResult": {"pages": []} + }))]) + .await; + let request = read_request( + &upstream.uri(), + json!({ + "pages": [2, 0, 0, 1], + "features": ["keyValuePairs", "languages"], + "future_option": {"nested": null}, + "extra_body": {"provider_option": false} + }), + ) + .with_document( + document( + json!({"type": "document_url", "document_url": "https://example.com/document.pdf"}), + ) + .into(), + ); + + perform(request).await.unwrap(); + + let sent = only_request(&upstream).await; + assert!( + sent.url.path().ends_with("/prebuilt-read:analyze"), + "{}", + sent.url + ); + assert_eq!(sent.query("pages").as_deref(), Some("1,2,3")); + assert_eq!( + sent.query("features").as_deref(), + Some("keyValuePairs,languages") + ); + assert_eq!( + sent.json(), + json!({ + "urlSource": "https://example.com/document.pdf", + "future_option": {"nested": null}, + "provider_option": false + }) + ); +} + +#[rstest] +#[case(json!({"pages": [true]}), Error::Pages("expected only integers or only strings".into()))] +#[case(json!({"pages": [1, "2"]}), Error::Pages("expected only integers or only strings".into()))] +#[case(json!({"pages": [-1]}), Error::Pages("negative page index".into()))] +#[case(json!({"pages": "1&&features=bad"}), Error::Pages("invalid native page range".into()))] +#[case(json!({"features": "languages&pages=1"}), Error::Features)] +#[case(json!({"req_format": "azure"}), Error::RequestFormat)] +#[tokio::test] +async fn invalid_pages_features_and_format_are_rejected_before_sending( + #[case] options: Value, + #[case] expected: Error, +) { + let upstream = upstream([json_response(json!({}))]).await; + + let result = match decode_request(wire( + MODEL, + &upstream.uri(), + json!({"type": "document_url", "document_url": "https://example.com/a.pdf"}), + options.clone(), + )) { + Ok(request) => perform(request).await, + Err(error) => Err(error), + }; + + assert!( + received(&upstream).await.is_empty(), + "sent invalid options: {options}" + ); + let error = result.unwrap_err(); + assert_eq!( + std::mem::discriminant(&error), + std::mem::discriminant(&expected) + ); + assert_eq!(error.http_status_code(), Some(400)); + assert_eq!(error.to_string(), expected.to_string()); +} + +#[rstest] +#[case::no_options(json!({}))] +#[case::litellm_format(json!({"req_format": "litellm"}))] +#[tokio::test] +async fn an_inline_document_is_sent_as_base64_and_only_page_text_is_kept(#[case] options: Value) { + let upstream = upstream([json_response(json!({ + "status": "succeeded", + "analyzeResult": {"pages": [{"pageNumber": 1, "lines": [{"content": "hello"}]}]} + }))]) + .await; + + let response = perform(read_request(&upstream.uri(), options)) + .await + .unwrap(); + + assert_eq!(response.pages.len(), 1); + assert_eq!(response.pages[0].index, 0); + assert_eq!(response.pages[0].markdown, "hello"); + assert_eq!(response.provider_native_response, None); + let serialized = response.into_json(); + for field in ["content", "tables", "keyValuePairs"] { + assert_eq!(serialized.get(field), Some(&Value::Null), "{field}"); + } + let sent = only_request(&upstream).await; + for field in ["pages", "features", "req_format"] { + assert_eq!(sent.query(field), None, "{field}"); + } + assert_eq!(sent.json(), json!({"base64Source": "YWJj"})); +} + +#[tokio::test] +async fn native_format_normalizes_pages_and_keeps_the_provider_response() { + let operation = json!({ + "status": "succeeded", + "operationExtension": 42, + "analyzeResult": { + "content": "A\n\nB", + "tables": [{"cells": []}], + "keyValuePairs": [{"key": {"content": "A"}}], + "pages": [{ + "pageNumber": "2", + "width": "8.5", + "height": 11, + "unit": "inch", + "lines": [{"content": "A"}, {"content": null}, {"content": "B"}] + }] + } + }); + let upstream = upstream([json_response(operation.clone())]).await; + + let result = perform(read_request( + &upstream.uri(), + json!({"req_format": "native"}), + )) + .await + .unwrap(); + + assert_eq!(result.pages[0].index, 1); + assert_eq!(result.pages[0].markdown, "A\n\nB"); + assert_eq!( + serde_json::to_value(&result.pages[0].dimensions).unwrap(), + json!({"width": 816, "height": 1056, "dpi": 96}) + ); + assert_eq!(result.usage_info.as_ref().unwrap().pages_processed, Some(1)); + let serialized = result.clone().into_json(); + assert_eq!(serialized["content"], "A\n\nB"); + assert_eq!(serialized["tables"], json!([{"cells": []}])); + assert_eq!( + serialized["keyValuePairs"], + json!([{"key": {"content": "A"}}]) + ); + assert!(serialized.get("key_value_pairs").is_none()); + assert_eq!( + result.provider_native_response.map(Value::Object), + Some(operation) + ); +} + +#[tokio::test] +async fn client_settings_choose_the_api_version_and_the_inch_to_pixel_dpi() { + let upstream = upstream([json_response(json!({ + "status": "succeeded", + "analyzeResult": {"pages": [{"pageNumber": 1, "width": 8.5, "height": 11, "unit": "inch"}]} + }))]) + .await; + let client = ocr_client().with_settings(OcrSettings { + document_intelligence_api_version: "2099-01-01".into(), + document_intelligence_dpi: 72, + ..OcrSettings::default() + }); + + let result = + litellm_core::ocr::client::perform(&client, read_request(&upstream.uri(), json!({}))) + .await + .unwrap(); + + assert_eq!( + only_request(&upstream) + .await + .query("api-version") + .as_deref(), + Some("2099-01-01") + ); + assert_eq!( + serde_json::to_value(&result.pages[0].dimensions).unwrap(), + json!({"width": 612, "height": 792, "dpi": 72}) + ); +} + +#[tokio::test] +async fn an_accepted_response_polls_to_success_with_only_credentials() { + let operation = json!({"status": "succeeded", "analyzeResult": {"pages": []}}); + let upstream = MockServer::start().await; + respond_in_order( + &upstream, + [ + accepted(&upstream, json!({})), + json_response(json!({"status": "running"})).insert_header("Retry-After", "0"), + json_response(operation.clone()), + ], + ) + .await; + let request = with_headers( + read_request(&upstream.uri(), json!({"req_format": "native"})), + &[("X-Trace", "initial-only")], + ); + + let result = perform(request).await.unwrap(); + + assert_eq!( + result.provider_native_response.map(Value::Object), + Some(operation) + ); + let requests = received(&upstream).await; + assert_eq!(requests.len(), 3); + assert_eq!(requests[0].header("x-trace"), Some("initial-only")); + for poll in &requests[1..] { + assert_eq!(poll.method.as_str(), "GET"); + assert_eq!(poll.url.path(), "/operation"); + assert_eq!(poll.header("x-trace"), None); + assert_eq!(poll.header("ocp-apim-subscription-key"), Some("test-key")); + } +} + +#[tokio::test] +async fn polling_forwards_bearer_credentials() { + let upstream = MockServer::start().await; + respond_in_order( + &upstream, + [ + accepted(&upstream, json!({})), + json_response(json!({"status": "succeeded"})), + ], + ) + .await; + let request = with_headers( + without_api_key(read_request(&upstream.uri(), json!({}))), + &[("Authorization", "Bearer token")], + ); + + perform(request).await.unwrap(); + + assert_eq!( + received(&upstream).await[1].header("authorization"), + Some("Bearer token") + ); +} + +#[tokio::test] +async fn response_received_fires_for_the_submission_and_the_completed_poll() { + let upstream = MockServer::start().await; + respond_in_order( + &upstream, + [ + accepted(&upstream, json!({"submitted": true})), + json_response(json!({"status": "succeeded"})), + ], + ) + .await; + let observed = Arc::new(Mutex::new(Vec::new())); + let recorder = observed.clone(); + let host = + LocalOcrHost::new(read_request(&upstream.uri(), json!({}))).with_observer(move |event| { + if let CallEvent::Machine(MachineEvent::ResponseReceived { raw }) = event { + recorder.lock().unwrap().push(raw.body.clone()); + } + }); + + perform_with(host).await.unwrap(); + + assert_eq!(received(&upstream).await.len(), 2); + assert_eq!( + *observed.lock().unwrap(), + [r#"{"submitted":true}"#, r#"{"status":"succeeded"}"#] + ); +} + +#[tokio::test] +async fn polling_does_not_follow_redirects() { + let upstream = MockServer::start().await; + respond_in_order( + &upstream, + [ + accepted(&upstream, json!({})), + ResponseTemplate::new(302) + .insert_header("Location", format!("{}/redirected", upstream.uri())), + json_response(json!({"status": "succeeded"})), + ], + ) + .await; + + let error = perform(read_request(&upstream.uri(), json!({}))) + .await + .unwrap_err(); + + assert!(error.to_string().contains("status 302"), "{error}"); + assert_eq!(received(&upstream).await.len(), 2); +} + +#[tokio::test] +async fn a_failed_operation_is_an_error() { + let upstream = MockServer::start().await; + respond_in_order( + &upstream, + [ + accepted(&upstream, json!({})), + json_response(json!({"status": "failed"})), + ], + ) + .await; + + let error = perform(read_request(&upstream.uri(), json!({}))) + .await + .unwrap_err(); + + assert!(error.to_string().contains("status failed"), "{error}"); +} + +#[tokio::test] +async fn the_polling_deadline_bounds_the_retry_delay() { + let upstream = MockServer::start().await; + respond_in_order( + &upstream, + [ + accepted(&upstream, json!({})), + json_response(json!({"status": "notStarted"})).insert_header("Retry-After", "9999"), + ], + ) + .await; + let client = ocr_client().with_settings(OcrSettings { + poll_timeout: Duration::from_millis(100), + ..OcrSettings::default() + }); + + let error = tokio::time::timeout( + Duration::from_secs(1), + litellm_core::ocr::client::perform(&client, read_request(&upstream.uri(), json!({}))), + ) + .await + .expect("the deadline cuts the retry delay short") + .unwrap_err(); + + assert!(error.to_string().contains("timed out"), "{error}"); +} + +#[rstest] +#[case::null_pages(json!({"pages": null}), "pages")] +#[case::null_page(json!({"pages": [null]}), "pages[0]")] +#[case::null_lines(json!({"pages": [{"lines": null}]}), "lines")] +#[case::bad_width(json!({"pages": [{"width": "bad"}]}), "width")] +#[tokio::test] +async fn malformed_provider_pages_report_the_response_path( + #[case] analysis: Value, + #[case] path: &str, +) { + let upstream = upstream([json_response(json!({ + "status": "succeeded", + "analyzeResult": analysis + }))]) + .await; + + let error = perform(read_request(&upstream.uri(), json!({}))) + .await + .unwrap_err(); + + assert!(error.to_string().contains(path), "{error}"); +} + +#[rstest] +#[case::missing(None)] +#[case::relative(Some("/relative"))] +#[case::cross_origin(Some("http://example.com/operation"))] +#[case::with_userinfo(Some("http://user:password@127.0.0.1/operation"))] +#[tokio::test] +async fn an_unusable_operation_location_is_rejected(#[case] location: Option<&str>) { + let response = location + .into_iter() + .fold(ResponseTemplate::new(202), |response, location| { + response.insert_header("Operation-Location", location) + }); + let upstream = upstream([response]).await; + + let error = perform(read_request(&upstream.uri(), json!({}))) + .await + .unwrap_err(); + + assert!(error.to_string().contains("operation-location"), "{error}"); + assert_eq!(received(&upstream).await.len(), 1); +} + +#[tokio::test] +async fn the_model_id_is_percent_encoded() { + let upstream = upstream([json_response(json!({"status": "succeeded"}))]).await; + + perform(ocr_request( + "azure_ai/doc-intelligence/a ?#é", + &upstream.uri(), + json!({}), + )) + .await + .unwrap(); + + let sent = only_request(&upstream).await; + assert!( + sent.url.path().ends_with("/a%20%3F%23%C3%A9:analyze"), + "{}", + sent.url + ); +} + +#[rstest] +#[case::dot("azure_ai/doc-intelligence/.")] +#[case::dot_dot("azure_ai/doc-intelligence/..")] +#[tokio::test] +async fn dot_segment_model_ids_are_rejected(#[case] model: &str) { + let error = perform(ocr_request(model, UNREACHABLE_BASE, json!({}))) + .await + .unwrap_err(); + + assert!(error.to_string().contains("dot segment"), "{error}"); +} diff --git a/litellm-rust/crates/core/tests/ocr/cohere.rs b/litellm-rust/crates/core/tests/ocr/cohere.rs new file mode 100644 index 00000000000..007aa49a2fd --- /dev/null +++ b/litellm-rust/crates/core/tests/ocr/cohere.rs @@ -0,0 +1,42 @@ +use rstest::rstest; + +use super::*; + +#[rstest] +#[case::cohere("cohere/parse-v5.0", "/v2/parse")] +#[case::azure_ai("azure_ai/Cohere-parse-v5.0", "/providers/cohere/v2/parse")] +#[tokio::test] +async fn an_image_goes_to_the_parse_endpoint_with_the_bearer_key( + #[case] model: &str, + #[case] path: &str, +) { + let upstream = upstream([pages_response()]).await; + let request = ocr_request_with_document( + model, + &upstream.uri(), + json!({"type": "image_url", "image_url": "data:image/png;base64,YWJj"}), + json!({}), + ); + + perform(request).await.unwrap(); + + let sent = only_request(&upstream).await; + assert_eq!(sent.method.as_str(), "POST"); + assert_eq!(sent.url.path(), path); + assert_eq!(sent.header("authorization"), Some("Bearer test-key")); +} + +#[rstest] +#[tokio::test] +async fn a_non_image_document_is_rejected_before_sending( + #[values("cohere/parse-v5.0", "azure_ai/Cohere-parse-v5.0")] model: &str, +) { + let upstream = upstream([pages_response()]).await; + + let error = perform(ocr_request(model, &upstream.uri(), json!({}))) + .await + .unwrap_err(); + + assert!(matches!(error, Error::CohereImageOnly), "{error:?}"); + assert!(received(&upstream).await.is_empty()); +} diff --git a/litellm-rust/crates/core/tests/ocr/documents.rs b/litellm-rust/crates/core/tests/ocr/documents.rs new file mode 100644 index 00000000000..e29ff3e9ee9 --- /dev/null +++ b/litellm-rust/crates/core/tests/ocr/documents.rs @@ -0,0 +1,182 @@ +use base64::Engine; +use litellm_core::ocr::types::OcrDocumentInput; +use litellm_host::event::WireRequest; +use rstest::rstest; +use wiremock::{Mock, matchers::any}; + +use super::*; + +const SERVED_DOCUMENT: &[u8] = b"\x89PNG served document"; +const REPLACED_DOCUMENT: &str = "data:image/png;base64,cmVwbGFjZWQ="; + +#[derive(Clone, Copy, Debug)] +enum Route { + Mistral, + AzureAi, + VertexMistral, + AzureCohereParse, + Cohere, +} + +impl Route { + fn model(self) -> &'static str { + match self { + Self::Mistral => "mistral/model", + Self::AzureAi => "azure_ai/model", + Self::VertexMistral => "vertex_ai/mistral-ocr-maas", + Self::AzureCohereParse => "azure_ai/cohere-parse", + Self::Cohere => "cohere/model", + } + } + + fn document_type(self) -> &'static str { + match self { + Self::Mistral | Self::AzureAi | Self::VertexMistral => "document_url", + Self::AzureCohereParse | Self::Cohere => "image_url", + } + } + + fn options(self) -> Value { + match self { + Self::Mistral | Self::AzureAi => json!({"pages": [0]}), + Self::VertexMistral => json!({"pages": [0], "vertex_project": "project-1"}), + Self::AzureCohereParse | Self::Cohere => json!({"output_format": "markdown"}), + } + } +} + +/// What the host does to the wire request in `before_send`. +#[derive(Clone, Copy, Debug)] +enum Guardrail { + Detached, + ReplacesDocument, +} + +impl Guardrail { + fn before_send(self, wire: WireRequest) -> WireRequest { + let Value::Object(fields) = wire.body else { + return wire; + }; + let body = fields + .into_iter() + .map(|(name, value)| match self { + Self::ReplacesDocument if name == "document" => { + let document_type = value["type"].clone(); + let key = document_type.as_str().unwrap_or_default().to_string(); + (name, json!({"type": document_type, key: REPLACED_DOCUMENT})) + } + Self::Detached | Self::ReplacesDocument => (name, value), + }) + .collect(); + WireRequest { + body: Value::Object(body), + ..wire + } + } +} + +/// Serves [`SERVED_DOCUMENT`] as `image/png` to every request. +async fn document_server() -> MockServer { + let server = MockServer::start().await; + Mock::given(any()) + .respond_with(ResponseTemplate::new(200).set_body_raw(SERVED_DOCUMENT, "image/png")) + .mount(&server) + .await; + server +} + +/// Sends a remote document through `route` and returns the document the provider saw. +async fn provider_document(route: Route, guardrail: Guardrail) -> Value { + let documents = document_server().await; + let upstream = upstream([pages_response()]).await; + let document_type = route.document_type(); + let request = ocr_request_with_document( + route.model(), + &upstream.uri(), + json!({"type": document_type, document_type: format!("{}/scan.png", documents.uri())}), + route.options(), + ); + let host = + LocalOcrHost::new(request).with_before_send(move |wire, _| Ok(guardrail.before_send(wire))); + + perform_with(host).await.unwrap(); + + only_request(&upstream).await.json()["document"][document_type].clone() +} + +#[rstest] +#[case::azure_ai(Route::AzureAi)] +#[case::vertex_mistral(Route::VertexMistral)] +#[case::azure_cohere_parse(Route::AzureCohereParse)] +#[tokio::test] +async fn inlining_routes_send_the_downloaded_document(#[case] route: Route) { + let expected = format!( + "data:image/png;base64,{}", + base64::engine::general_purpose::STANDARD.encode(SERVED_DOCUMENT) + ); + + assert_eq!( + provider_document(route, Guardrail::Detached).await, + expected + ); +} + +#[rstest] +#[tokio::test] +async fn a_document_replaced_by_the_host_reaches_the_provider( + #[values( + Route::Mistral, + Route::AzureAi, + Route::VertexMistral, + Route::AzureCohereParse, + Route::Cohere + )] + route: Route, +) { + assert_eq!( + provider_document(route, Guardrail::ReplacesDocument).await, + REPLACED_DOCUMENT + ); +} + +#[tokio::test] +async fn an_empty_byte_document_fails_before_sending() { + let upstream = upstream([pages_response()]).await; + let request = ocr_request("mistral/model", &upstream.uri(), json!({})).with_document( + OcrDocumentInput::Bytes { + bytes: Default::default(), + file_name: None, + mime_type: None, + }, + ); + + let error = perform(request).await.unwrap_err(); + + assert!(matches!(error, Error::EmptyFile), "{error:?}"); + assert!(received(&upstream).await.is_empty()); +} + +#[tokio::test] +async fn a_missing_path_document_fails_before_sending() { + let upstream = upstream([pages_response()]).await; + let path = + std::env::temp_dir().join(format!("litellm-ocr-missing-{}.png", rand::random::())); + let request = ocr_request("mistral/model", &upstream.uri(), json!({})).with_document( + OcrDocumentInput::Path { + path: path.clone(), + mime_type: None, + }, + ); + + let error = perform(request).await.unwrap_err(); + + assert!( + matches!( + &error, + Error::FileRead { path: failed, source } + if *failed == path && source.kind() == std::io::ErrorKind::NotFound + ), + "{error:?}" + ); + assert!(received(&upstream).await.is_empty()); +} diff --git a/litellm-rust/crates/core/tests/ocr/lifecycle.rs b/litellm-rust/crates/core/tests/ocr/lifecycle.rs new file mode 100644 index 00000000000..65e64cce79b --- /dev/null +++ b/litellm-rust/crates/core/tests/ocr/lifecycle.rs @@ -0,0 +1,269 @@ +use std::sync::{Arc, Mutex}; + +use litellm_core::ocr::{ + route::{Ocr, OcrOp, OcrProjection, ocr_machine}, + types::OcrDocumentInput, +}; +use litellm_host::{ + event::{CallEvent, MachineEvent, RequestContext, WireRequest}, + host::Host, +}; +use rstest::rstest; + +use super::*; + +pub(crate) fn event_name(event: &CallEvent) -> &'static str { + match event { + CallEvent::Started { .. } => "started", + CallEvent::Machine(MachineEvent::ResponseReceived { .. }) => "response", + CallEvent::Succeeded { .. } => "success", + CallEvent::Failed { .. } => "failure", + } +} + +fn recording_host( + request: LiteLLMOcrRequest, + events: Arc>>, + block: bool, +) -> LocalOcrHost { + let before_send_events = events.clone(); + LocalOcrHost::new(request) + .with_before_send(move |wire, _| { + before_send_events.lock().unwrap().push("before_send"); + match block { + true => Err(Error::InvalidRequest("blocked".into())), + false => Ok(wire), + } + }) + .with_observer(move |event| events.lock().unwrap().push(event_name(event))) +} + +#[tokio::test] +async fn hooks_run_in_order_and_one_success_is_emitted() { + let upstream = upstream([pages_response()]).await; + let events = Arc::new(Mutex::new(Vec::new())); + + perform_with(recording_host( + ocr_request("mistral/model", &upstream.uri(), json!({})), + events.clone(), + false, + )) + .await + .unwrap(); + + assert_eq!( + *events.lock().unwrap(), + ["started", "before_send", "response", "success"] + ); + assert_eq!(received(&upstream).await.len(), 1); +} + +#[tokio::test] +async fn a_blocking_before_send_prevents_the_call_and_emits_one_failure() { + let upstream = upstream([pages_response()]).await; + let events = Arc::new(Mutex::new(Vec::new())); + + let error = perform_with(recording_host( + ocr_request("mistral/model", &upstream.uri(), json!({})), + events.clone(), + true, + )) + .await + .unwrap_err(); + + assert!( + matches!(&error, Error::InvalidRequest(message) if message == "blocked"), + "{error:?}" + ); + assert_eq!( + *events.lock().unwrap(), + ["started", "before_send", "failure"] + ); + assert!(received(&upstream).await.is_empty()); +} + +#[tokio::test] +async fn an_upstream_failure_emits_one_terminal_failure() { + let upstream = upstream([status_response(500, json!({"error": "failed"}))]).await; + let events = Arc::new(Mutex::new(Vec::new())); + + let result = perform_with(recording_host( + ocr_request("mistral/model", &upstream.uri(), json!({})), + events.clone(), + false, + )) + .await; + + assert!(result.is_err()); + assert_eq!( + *events.lock().unwrap(), + ["started", "before_send", "failure"] + ); + assert_eq!(received(&upstream).await.len(), 1); +} + +#[tokio::test] +async fn an_invalid_provider_response_is_observed_before_normalization_fails() { + let upstream = upstream([json_response(json!({"pages": "invalid"}))]).await; + let observed = Arc::new(Mutex::new(Vec::new())); + let recorder = observed.clone(); + let host = LocalOcrHost::new(ocr_request("mistral/model", &upstream.uri(), json!({}))) + .with_observer(move |event| { + if let CallEvent::Machine(MachineEvent::ResponseReceived { raw }) = event { + recorder.lock().unwrap().push(raw.body.clone()); + } + }); + + let error = perform_with(host).await.unwrap_err(); + + assert!(matches!(error, Error::ResponseField { .. }), "{error:?}"); + assert_eq!(*observed.lock().unwrap(), [r#"{"pages":"invalid"}"#]); +} + +#[tokio::test] +async fn headers_returned_by_before_send_are_sent() { + let upstream = upstream([pages_response()]).await; + let host = LocalOcrHost::new(ocr_request("mistral/model", &upstream.uri(), json!({}))) + .with_before_send(|mut wire, _| { + wire.headers + .push(("x-core-callback".into(), "edited".into())); + Ok(wire) + }); + + perform_with(host).await.unwrap(); + + assert_eq!( + only_request(&upstream).await.header("x-core-callback"), + Some("edited") + ); +} + +async fn before_send_context(request: LiteLLMOcrRequest) -> (WireRequest, RequestContext) { + let observed = Arc::new(Mutex::new(None)); + let captured = observed.clone(); + let host = LocalOcrHost::new(request).with_before_send(move |wire, context| { + *captured.lock().unwrap() = Some((wire.clone(), context.clone())); + Ok(wire) + }); + perform_with(host).await.unwrap(); + let context = observed.lock().unwrap().take(); + context.expect("before_send ran") +} + +#[tokio::test] +async fn before_send_sees_the_route_its_params_and_the_body() { + let upstream = upstream([pages_response()]).await; + + let (wire, context) = before_send_context(ocr_request( + "mistral/model", + &upstream.uri(), + json!({"pages": [0], "req_format": "native"}), + )) + .await; + + assert_eq!(context.custom_llm_provider, "mistral"); + assert_eq!(context.model, "model"); + assert_eq!(context.optional_params["req_format"], "native"); + assert!(context.secret_fields.is_empty()); + assert_eq!(wire.body["pages"], json!([0])); +} + +#[rstest] +#[case::client_secret(json!({"client_secret": "shh", "tenant_id": "t"}), &["client_secret"])] +#[case::no_secrets(json!({"tenant_id": "t"}), &[])] +#[tokio::test] +async fn before_send_names_the_secret_params(#[case] options: Value, #[case] secrets: &[&str]) { + let upstream = upstream([pages_response()]).await; + let request = ocr_request("azure_ai/model", &upstream.uri(), options).with_document( + OcrDocumentInput::Bytes { + bytes: b"abc".as_slice().into(), + file_name: None, + mime_type: Some("application/pdf".into()), + }, + ); + + let (_, context) = before_send_context(request).await; + + assert_eq!(context.secret_fields, secrets); +} + +/// Hands the route a caller-owned Azure token and rewrites the bearer in `before_send`. +struct CallerTokenHost { + request: Mutex>, + trace: Mutex>, +} + +impl Host for CallerTokenHost { + async fn project(&self) -> Result { + self.trace.lock().unwrap().push("project".into()); + Ok(OcrProjection { + request: self.request.lock().unwrap().take().unwrap(), + caller_token: true, + }) + } + + async fn custom_op(&self, op: OcrOp) -> Result<(), Error> { + match op { + OcrOp::AcquireAzureAdToken(reply) => { + self.trace.lock().unwrap().push("token".into()); + reply.send(litellm_auth::ResolvedCredential::Static( + litellm_auth::SecretValue::new("caller-token"), + )); + Ok(()) + } + } + } + + async fn before_send( + &self, + wire: WireRequest, + _: &RequestContext, + ) -> Result { + let is_authorization = |name: &str| name.eq_ignore_ascii_case("authorization"); + let authorization = wire + .headers + .iter() + .find(|(name, _)| is_authorization(name)) + .map(|(_, value)| value.clone()) + .unwrap_or_default(); + self.trace + .lock() + .unwrap() + .push(format!("before_send:{authorization}")); + let headers = wire + .headers + .into_iter() + .map(|(name, value)| match is_authorization(&name) { + true => (name, "Bearer edited".to_string()), + false => (name, value), + }) + .collect(); + Ok(WireRequest { headers, ..wire }) + } +} + +#[tokio::test] +async fn the_callers_azure_token_is_acquired_before_before_send_which_can_still_replace_it() { + let upstream = upstream([pages_response()]).await; + let host = CallerTokenHost { + request: Mutex::new(Some(without_api_key(ocr_request( + "azure_ai/model", + &upstream.uri(), + json!({}), + )))), + trace: Mutex::new(Vec::new()), + }; + + litellm_host::run::run(ocr_machine(ocr_client()), &host) + .await + .unwrap(); + + assert_eq!( + *host.trace.lock().unwrap(), + ["project", "token", "before_send:Bearer caller-token"] + ); + assert_eq!( + only_request(&upstream).await.header_values("authorization"), + ["Bearer edited"] + ); +} diff --git a/litellm-rust/crates/core/tests/ocr/machine.rs b/litellm-rust/crates/core/tests/ocr/machine.rs new file mode 100644 index 00000000000..073ce67e74b --- /dev/null +++ b/litellm-rust/crates/core/tests/ocr/machine.rs @@ -0,0 +1,284 @@ +use std::{ + sync::{ + Arc, + atomic::{AtomicBool, Ordering}, + }, + time::Duration, +}; + +use litellm_core::ocr::{ + route::{OcrMachine, OcrOp, OcrProjection}, + types::OcrDocumentInput, +}; +use litellm_host::{ + event::{CallEvent, WireRequest}, + host::{Host, HostOp}, + machine::{HostFailure, Machine, MachineStep}, +}; +use litellm_llms::base_llm::ocr::transformation::OcrTransportConfig; +use rstest::rstest; +use tokio::{io::AsyncReadExt, net::TcpListener, sync::Notify}; + +use super::{lifecycle::event_name, *}; + +/// Drives the machine by hand, answering every op through `host` except `before_send`, +/// which `intercept` answers so a test can fail or cancel exactly there. +async fn drive_until( + host: &LocalOcrHost, + mut intercept: impl FnMut(WireRequest) -> Result>, +) -> ( + Result, + Vec<&'static str>, + OcrMachine, +) { + let mut machine = ocr_machine(ocr_client()); + let mut ops = Vec::new(); + let outcome = loop { + let op = match machine.resume().await { + Ok(MachineStep::Host(op)) => op, + Ok(MachineStep::Complete(response)) => break Ok(response), + Err(error) => break Err(error), + }; + let answer = match op { + HostOp::Project(reply) => { + ops.push("Project"); + host.project() + .await + .map(|projection| reply.send(projection)) + .map_err(HostFailure::Error) + } + HostOp::Custom(op) => { + ops.push(match op { + OcrOp::AcquireAzureAdToken(_) => "AcquireAzureAdToken", + }); + host.custom_op(op).await.map_err(HostFailure::Error) + } + HostOp::BeforeSend { wire, reply, .. } => { + ops.push("BeforeSend"); + intercept(*wire).map(|wire| reply.send(wire)) + } + HostOp::Emit(event, reply) => { + let event = CallEvent::Machine(event); + ops.push(event_name(&event)); + host.emit(&event) + .await + .map(|()| reply.send(())) + .map_err(HostFailure::Error) + } + }; + if let Err(failure) = answer { + break machine.interrupt(failure).await; + } + }; + (outcome, ops, machine) +} + +/// Answers every op until `stop` fires, leaving the machine suspended mid-call. +async fn drive_until_notified(machine: &mut OcrMachine, host: &LocalOcrHost, stop: &Notify) { + tokio::time::timeout(Duration::from_secs(2), async { + loop { + tokio::select! { + _ = stop.notified() => break, + step = machine.resume() => { + match step.unwrap() { + MachineStep::Host(HostOp::Project(reply)) => reply.send(host.project().await.unwrap()), + MachineStep::Host(HostOp::Custom(op)) => host.custom_op(op).await.unwrap(), + MachineStep::Host(HostOp::BeforeSend { wire, reply, .. }) => reply.send(*wire), + MachineStep::Host(HostOp::Emit(_, reply)) => reply.send(()), + MachineStep::Complete(_) => panic!("the stalled call completed"), + } + } + } + } + }) + .await + .expect("the call reached the stall point"); +} + +#[tokio::test] +async fn a_hand_driven_machine_performs_the_same_call() { + let upstream = upstream([json_response(json!({ + "pages": [{"index": 0, "markdown": "native"}] + }))]) + .await; + let host = LocalOcrHost::new(ocr_request("mistral/model", &upstream.uri(), json!({}))); + + let (outcome, ops, mut machine) = drive_until(&host, Ok).await; + + assert_eq!(outcome.unwrap().pages[0].markdown, "native"); + assert_eq!(received(&upstream).await.len(), 1); + assert_eq!(ops, ["Project", "BeforeSend", "response"]); + assert!(matches!( + machine.resume().await, + Err(Error::InvalidRequest(_)) + )); +} + +#[tokio::test] +async fn a_path_document_is_read_by_core_without_a_host_operation() { + let upstream = upstream([json_response(json!({ + "pages": [{"index": 0, "markdown": "path"}] + }))]) + .await; + let dir = std::env::temp_dir().join(format!("litellm-ocr-{}", rand::random::())); + std::fs::create_dir_all(&dir).unwrap(); + let path = dir.join("scan.png"); + std::fs::write(&path, b"abc").unwrap(); + let request = ocr_request("mistral/model", &upstream.uri(), json!({})).with_document( + OcrDocumentInput::Path { + path, + mime_type: None, + }, + ); + + let (response, ops, _) = drive_until(&LocalOcrHost::new(request), Ok).await; + std::fs::remove_dir_all(&dir).unwrap(); + + assert_eq!(response.unwrap().pages[0].markdown, "path"); + assert_eq!(ops, ["Project", "BeforeSend", "response"]); + assert_eq!( + only_request(&upstream).await.json()["document"]["image_url"], + "data:image/png;base64,YWJj" + ); +} + +#[rstest] +#[case::failed(HostFailure::Error(Error::InvalidRequest("before_send failed".into())), "before_send failed")] +#[case::cancelled(HostFailure::Cancelled(Error::InvalidRequest("cancelled".into())), "cancelled")] +#[tokio::test] +async fn a_before_send_failure_ends_the_call_without_reaching_transport( + #[case] failure: HostFailure, + #[case] message: &str, +) { + let upstream = upstream([pages_response()]).await; + let host = LocalOcrHost::new(ocr_request("mistral/model", &upstream.uri(), json!({}))); + let failure = Arc::new(std::sync::Mutex::new(Some(failure))); + + let (outcome, ops, mut machine) = drive_until(&host, |_| { + Err(failure + .lock() + .unwrap() + .take() + .expect("before_send is asked once")) + }) + .await; + + assert!( + matches!(&outcome, Err(Error::InvalidRequest(actual)) if actual == message), + "{outcome:?}" + ); + assert_eq!(ops, ["Project", "BeforeSend"]); + assert!(machine.resume().await.is_err()); + assert!(received(&upstream).await.is_empty()); +} + +#[tokio::test] +async fn resuming_before_answering_keeps_the_pending_operation() { + let request = ocr_request("mistral/model", UNREACHABLE_BASE, json!({})); + let mut machine = ocr_machine(ocr_client()); + let Ok(MachineStep::Host(HostOp::Project(reply))) = machine.resume().await else { + panic!("expected the projection op first"); + }; + + assert!(machine.resume().await.is_err()); + reply.send(OcrProjection { + request, + caller_token: false, + }); + assert!(matches!( + machine.resume().await, + Ok(MachineStep::Host(HostOp::BeforeSend { .. })) + )); +} + +#[derive(Debug)] +struct PendingToken { + entered: Arc, + dropped: Arc, +} + +struct TokenFutureDrop(Arc); + +impl Drop for TokenFutureDrop { + fn drop(&mut self) { + self.0.store(true, Ordering::SeqCst); + } +} + +impl litellm_auth::TokenProvider for PendingToken { + fn acquire(&self) -> litellm_auth::TokenFuture<'_> { + Box::pin(async move { + let _guard = TokenFutureDrop(self.dropped.clone()); + self.entered.notify_one(); + std::future::pending().await + }) + } +} + +#[tokio::test] +async fn interrupt_drops_provider_captures_before_returning() { + let entered = Arc::new(Notify::new()); + let dropped = Arc::new(AtomicBool::new(false)); + let mut request = ocr_request("azure_ai/mistral-ocr", "https://example.invalid", json!({})); + request.transport = OcrTransportConfig { + extra_headers: vec![("authorization".into(), "Bearer test-key".into())], + ..request.transport + }; + request.azure_ad_token_provider = Some(litellm_auth::TokenProviderHandle::new(Arc::new( + PendingToken { + entered: entered.clone(), + dropped: dropped.clone(), + }, + ))); + let host = LocalOcrHost::new(request); + let mut machine = ocr_machine(ocr_client()); + + drive_until_notified(&mut machine, &host, &entered).await; + assert!(!dropped.load(Ordering::SeqCst)); + let acknowledgement = machine.interrupt(HostFailure::Cancelled(Error::InvalidRequest( + "cancelled".into(), + ))); + + assert!( + dropped.load(Ordering::SeqCst), + "interrupt returned while provider captures were still alive" + ); + assert!( + matches!(acknowledgement.await, Err(Error::InvalidRequest(message)) if message == "cancelled") + ); +} + +#[tokio::test] +async fn interrupting_an_in_flight_provider_request_closes_its_connection() { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let base = format!("http://{}", listener.local_addr().unwrap()); + let received = Arc::new(Notify::new()); + let server_received = received.clone(); + let server = tokio::spawn(async move { + let (mut socket, _) = listener.accept().await.unwrap(); + let mut request = Vec::new(); + let mut buffer = [0u8; 4096]; + while !request.windows(4).any(|window| window == b"\r\n\r\n") { + let read = socket.read(&mut buffer).await.unwrap(); + request.extend_from_slice(&buffer[..read]); + } + server_received.notify_one(); + while socket.read(&mut buffer).await.unwrap() != 0 {} + }); + let host = LocalOcrHost::new(ocr_request("mistral/model", &base, json!({}))); + let mut machine = ocr_machine(ocr_client()); + + drive_until_notified(&mut machine, &host, &received).await; + let cancelled = Error::InvalidRequest("cancelled".into()); + + assert!( + machine + .interrupt(HostFailure::Cancelled(cancelled)) + .await + .is_err() + ); + tokio::time::timeout(Duration::from_secs(1), server) + .await + .expect("the provider connection stayed open after the interrupt") + .unwrap(); +} diff --git a/litellm-rust/crates/core/tests/ocr/main.rs b/litellm-rust/crates/core/tests/ocr/main.rs new file mode 100644 index 00000000000..1a915389b20 --- /dev/null +++ b/litellm-rust/crates/core/tests/ocr/main.rs @@ -0,0 +1,125 @@ +use litellm_core::ocr::{ + document::prepare_document, + route::{LocalOcrHost, ocr_machine}, + types::LiteLLMOcrRequest, + wire::{OcrWireRequest, decode_request}, +}; +use litellm_llms::base_llm::ocr::{ + error::Error, + handler::OcrClient, + transformation::{LiteLLMOcrResponse, OcrDocument}, +}; +use serde_json::{Map, Value, json}; +use wiremock::{MockServer, ResponseTemplate}; + +#[path = "../support/mod.rs"] +mod support; +use support::*; + +mod aws_textract; +mod azure_ai; +mod azure_document_intelligence; +mod cohere; +mod documents; +mod lifecycle; +mod machine; +mod mistral; +mod reducto; +mod vertex_ai; + +const INLINE_PDF: &str = "data:application/pdf;base64,YWJj"; + +fn object(value: Value) -> Map { + let Value::Object(map) = value else { + panic!("expected a json object, got {value}"); + }; + map +} + +fn ocr_client() -> OcrClient { + let document_http = reqwest::Client::builder() + .redirect(reqwest::redirect::Policy::none()) + .build() + .expect("test document client builds"); + OcrClient::for_test(reqwest::Client::new(), document_http) +} + +async fn perform(request: LiteLLMOcrRequest) -> Result { + litellm_core::ocr::client::perform(&ocr_client(), request).await +} + +async fn perform_with(host: LocalOcrHost) -> Result { + litellm_host::run::run(ocr_machine(ocr_client()), &host).await +} + +fn wire(model: &str, base: &str, document: Value, options: Value) -> OcrWireRequest { + OcrWireRequest { + model: model.into(), + document, + api_key: Some(litellm_auth::SecretValue::new("test-key")), + api_base: Some(base.into()), + custom_llm_provider: None, + extra_headers: None, + optional_params: object(options), + input_sources: Default::default(), + timeout_seconds: Some(2.0), + } +} + +/// A request for an inline PDF, authenticated with `test-key`. +fn ocr_request(model: &str, base: &str, options: Value) -> LiteLLMOcrRequest { + ocr_request_with_document( + model, + base, + json!({"type": "document_url", "document_url": INLINE_PDF}), + options, + ) +} + +fn ocr_request_with_document( + model: &str, + base: &str, + document: Value, + options: Value, +) -> LiteLLMOcrRequest { + decode_request(wire(model, base, document, options)).expect("request decodes") +} + +fn document(value: Value) -> OcrDocument { + serde_json::from_value(value).expect("document parses") +} + +/// Points the request's resolved document at `source`, keeping its type. +fn with_source(request: LiteLLMOcrRequest, source: &str) -> LiteLLMOcrRequest { + let resolved = request + .map_document(prepare_document) + .expect("document resolves"); + let document = resolved.document.clone().with_source(source.into()); + resolved.with_document(document.into()) +} + +fn with_headers(request: LiteLLMOcrRequest, headers: &[(&str, &str)]) -> LiteLLMOcrRequest { + let mut request = request; + request.transport.extra_headers = headers + .iter() + .map(|(name, value)| (name.to_string(), value.to_string())) + .collect(); + request +} + +fn without_api_key(request: LiteLLMOcrRequest) -> LiteLLMOcrRequest { + let mut request = request; + request.credentials.api_key = None; + request +} + +fn pages_response() -> ResponseTemplate { + json_response(json!({"pages": []})) +} + +/// An Azure Document Intelligence 202 whose operation lives on `server`. +fn accepted(server: &MockServer, body: Value) -> ResponseTemplate { + ResponseTemplate::new(202) + .insert_header("Operation-Location", format!("{}/operation", server.uri())) + .set_body_json(body) +} diff --git a/litellm-rust/crates/core/tests/ocr/mistral.rs b/litellm-rust/crates/core/tests/ocr/mistral.rs new file mode 100644 index 00000000000..f80e564b03f --- /dev/null +++ b/litellm-rust/crates/core/tests/ocr/mistral.rs @@ -0,0 +1,248 @@ +use std::sync::Arc; + +use litellm_auth_gcp::VertexAuth; +use litellm_http::{ + HttpClientPool, HttpSettings, Resolution, + media::{PublicDnsResolver, UrlPolicy}, +}; +use litellm_llms::{ + base_llm::ocr::{ + settings::OcrSettings, + transformation::{BaseOcrConfig, OCR_RESPONSE_MAX_BYTES}, + }, + mistral::ocr::transformation::MistralOcrConfig, +}; +use rstest::rstest; + +use super::*; + +#[tokio::test] +async fn direct_mistral_sends_one_request_with_every_option() { + let upstream = upstream([json_response(json!({ + "pages": [{"index": 0, "markdown": "hello", "custom": "preserved"}], + "usage_info": {"pages_processed": 1} + }))]) + .await; + + let result = perform(ocr_request( + "mistral/model", + &upstream.uri(), + json!({"pages": "0,2-4", "extract_header": true, "unknown": "ignored"}), + )) + .await + .unwrap(); + + assert_eq!(result.pages[0].markdown, "hello"); + assert_eq!(result.pages[0].extra_fields["custom"], "preserved"); + let sent = only_request(&upstream).await; + assert_eq!(sent.url.path(), "/v1/ocr"); + assert_eq!(sent.header("authorization"), Some("Bearer test-key")); + assert_eq!( + sent.json(), + json!({ + "model": "model", + "document": {"type": "document_url", "document_url": INLINE_PDF}, + "pages": "0,2-4", + "extract_header": true, + "unknown": "ignored" + }) + ); +} + +#[rstest] +#[case::litellm_format(json!({}), false)] +#[case::native_format(json!({"req_format": "native"}), true)] +#[tokio::test] +async fn the_native_response_is_kept_only_when_requested( + #[case] options: Value, + #[case] kept: bool, +) { + let provider_response = json!({ + "pages": [{"index": 0, "markdown": "hello"}], + "usage_info": {"pages_processed": 1}, + "provider_only": "preserved" + }); + let upstream = upstream([json_response(provider_response.clone())]).await; + + let response = perform(ocr_request("mistral/model", &upstream.uri(), options)) + .await + .unwrap(); + + assert_eq!( + response.provider_native_response.map(Value::Object), + kept.then_some(provider_response) + ); +} + +#[rstest] +#[case::mistral("mistral/model", json!({}))] +#[case::vertex( + "vertex_ai/mistral-ocr-latest", + json!({"vertex_project": "test-project", "vertex_location": "us-central1"}) +)] +#[tokio::test] +async fn an_upstream_error_keeps_its_status_whole_body_and_headers( + #[case] model: &str, + #[case] options: Value, +) { + let payload = json!({"message": format!("{} END-OF-PROVIDER-BODY", "x".repeat(4096))}); + let expected_body = serde_json::to_string(&payload).unwrap(); + let upstream = upstream([status_response(422, payload) + .insert_header("Retry-After", "17") + .insert_header("X-Request-ID", "request-123") + .insert_header("X-Future-Header", "retained")]) + .await; + + let error = perform(ocr_request(model, &upstream.uri(), options)) + .await + .unwrap_err(); + + assert_eq!(received(&upstream).await.len(), 1); + let Error::Provider { + status, + body, + headers, + } = error + else { + panic!("expected provider error, got {error:?}"); + }; + assert_eq!(status, 422); + for (name, value) in [ + ("retry-after", "17"), + ("x-request-id", "request-123"), + ("x-future-header", "retained"), + ] { + assert!( + headers + .iter() + .any(|(key, actual)| key.eq_ignore_ascii_case(name) && actual == value), + "{name} missing from {headers:?}" + ); + } + assert_eq!(body, expected_body); +} + +#[rstest] +#[case::mistral_prefix("mistral/model", None, true)] +#[case::unknown_provider("model", Some("unknown"), false)] +fn decoding_accepts_known_providers_and_rejects_unknown_ones( + #[case] model: &str, + #[case] provider: Option<&str>, + #[case] accepted: bool, +) { + let request = OcrWireRequest { + custom_llm_provider: provider.map(Into::into), + ..wire( + model, + "https://example.com", + json!({"type": "document_url", "document_url": "https://example.com/doc.pdf"}), + json!({"extract_header": true, "unknown": 42}), + ) + }; + + assert_eq!(decode_request(request).is_ok(), accepted); +} + +#[rstest] +#[case::plain_key(&[("MISTRAL_API_KEY", "plain")], "plain")] +#[case::azure_key_wins(&[("MISTRAL_AZURE_API_KEY", "azure"), ("MISTRAL_API_KEY", "plain")], "azure")] +#[case::empty_azure_key_falls_through(&[("MISTRAL_AZURE_API_KEY", ""), ("MISTRAL_API_KEY", "plain")], "plain")] +#[tokio::test] +async fn missing_credentials_come_from_the_injected_secret_source( + #[case] secrets: &[(&str, &str)], + #[case] expected_key: &str, +) { + let upstream = upstream([pages_response()]).await; + let base = upstream.uri(); + let source = Arc::new(RecordingSecrets::new( + secrets + .iter() + .copied() + .chain([("MISTRAL_AZURE_API_BASE", base.as_str())]), + )); + let client = ocr_client().with_secrets(source.clone()); + let request = decode_request(OcrWireRequest { + api_key: None, + api_base: None, + ..wire( + "mistral/model", + &base, + json!({"type": "document_url", "document_url": INLINE_PDF}), + json!({}), + ) + }) + .unwrap(); + + litellm_core::ocr::client::perform(&client, request) + .await + .unwrap(); + + assert_eq!(source.requested(), MistralOcrConfig.secret_names()); + assert_eq!( + only_request(&upstream).await.header("authorization"), + Some(format!("Bearer {expected_key}").as_str()) + ); +} + +#[tokio::test] +async fn the_client_uses_the_injected_http_pool_configuration() { + let upstream = upstream([pages_response()]).await; + let settings = HttpSettings { + user_agent: Some("host-owned/1".into()), + ..HttpSettings::default() + }; + let client = OcrClient::new( + &HttpClientPool::new(Arc::new(PublicDnsResolver)), + &Resolution::from(&settings).config, + UrlPolicy::default(), + VertexAuth::default(), + OcrSettings::default(), + Arc::new(litellm_secrets::source::EnvironmentSecrets::default()), + ) + .unwrap(); + + litellm_core::ocr::client::perform( + &client, + ocr_request("mistral/model", &upstream.uri(), json!({})), + ) + .await + .unwrap(); + + assert_eq!( + only_request(&upstream).await.header("user-agent"), + Some("host-owned/1") + ); +} + +#[test] +fn a_valid_response_limit_is_consumed_and_not_forwarded() { + let request = ocr_request( + "mistral/model", + UNREACHABLE_BASE, + json!({"max_response_bytes": 123}), + ); + + assert_eq!(request.transport.max_response_bytes, 123); + assert!(!request.optional_params.contains_key("max_response_bytes")); +} + +#[rstest] +#[case::zero(json!(0))] +#[case::negative(json!(-1))] +#[case::boolean(json!(true))] +#[case::string(json!("123"))] +#[case::fraction(json!(1.5))] +#[case::above_the_cap(json!(OCR_RESPONSE_MAX_BYTES + 1))] +#[case::null(Value::Null)] +fn an_invalid_response_limit_is_rejected(#[case] limit: Value) { + let Err(error) = decode_request(wire( + "mistral/model", + UNREACHABLE_BASE, + json!({"type": "document_url", "document_url": INLINE_PDF}), + json!({"max_response_bytes": limit}), + )) else { + panic!("invalid response limit {limit} accepted"); + }; + + assert!(error.to_string().contains("max_response_bytes"), "{error}"); +} diff --git a/litellm-rust/crates/core/tests/ocr/reducto.rs b/litellm-rust/crates/core/tests/ocr/reducto.rs new file mode 100644 index 00000000000..8ccab27e58d --- /dev/null +++ b/litellm-rust/crates/core/tests/ocr/reducto.rs @@ -0,0 +1,321 @@ +use std::sync::{Arc, Mutex}; + +use litellm_host::event::{CallEvent, MachineEvent, WireRequest}; +use rstest::rstest; + +use super::*; + +fn upload_response() -> ResponseTemplate { + json_response(json!({"file_id": "reducto://uploaded.pdf"})) +} + +fn chunks_response(chunks: Value) -> ResponseTemplate { + json_response(json!({"result": {"chunks": chunks}})) +} + +fn source_field(model: &str) -> &'static str { + match model.ends_with("parse-legacy") { + true => "document_url", + false => "input", + } +} + +#[rstest] +#[case::v3( + "reducto/parse-v3", + json!({ + "formatting": {"table_output_format": "html"}, + "retrieval": {"chunk_mode": "section"}, + "settings": {"ocr_system": "standard"}, + "future_ocr_option": true, + "extra_body": {"provider_option": "value"} + }), + "reducto://already.pdf", + json!({ + "input": "reducto://already.pdf", + "formatting": {"table_output_format": "html"}, + "retrieval": {"chunk_mode": "section"}, + "settings": {"ocr_system": "standard"}, + "future_ocr_option": true, + "provider_option": "value" + }) +)] +#[case::legacy( + "reducto/parse-legacy", + json!({ + "enhance": {"agentic": [{"type": "table"}]}, + "future_ocr_option": true, + "extra_body": {"provider_option": "value"} + }), + "reducto://legacy.pdf", + json!({ + "document_url": "reducto://legacy.pdf", + "options": {"enhance": {"agentic": [{"type": "table"}]}}, + "future_ocr_option": true, + "provider_option": "value" + }) +)] +#[tokio::test] +async fn an_uploaded_document_is_parsed_with_mapped_options( + #[case] model: &str, + #[case] options: Value, + #[case] source: &str, + #[case] expected: Value, +) { + let upstream = upstream([chunks_response(json!([]))]).await; + + perform(with_source( + ocr_request(model, &upstream.uri(), options), + source, + )) + .await + .unwrap(); + + let sent = only_request(&upstream).await; + assert_eq!(sent.url.path(), "/parse"); + assert_eq!(sent.json(), expected); +} + +#[rstest] +#[tokio::test] +async fn an_inline_document_is_uploaded_as_multipart_then_parsed( + #[values("parse-v3", "parse-legacy")] model: &str, + #[values("application/pdf", "image/png")] mime_type: &str, +) { + let upstream = upstream([ + upload_response(), + chunks_response(json!([{"content": "hello"}])), + ]) + .await; + let data_uri = format!("data:{mime_type};base64,YWJj"); + let document = match mime_type.starts_with("image/") { + true => json!({"type": "image_url", "image_url": data_uri}), + false => json!({"type": "document_url", "document_url": data_uri}), + }; + let request = with_headers( + ocr_request_with_document( + &format!("reducto/{model}"), + &upstream.uri(), + document, + json!({}), + ), + &[ + ("Content-Type", "application/json"), + ("X-Trace", "upload-test"), + ], + ); + + let response = perform(request).await.unwrap(); + + assert_eq!(response.pages[0].markdown, "hello"); + let requests = received(&upstream).await; + let [upload, parse] = requests.as_slice() else { + panic!( + "expected an upload and a parse, got {} requests", + requests.len() + ); + }; + assert_eq!(upload.url.path(), "/upload"); + assert!( + upload + .header("content-type") + .is_some_and(|value| value.starts_with("multipart/form-data; boundary=")), + "{:?}", + upload.header("content-type") + ); + assert_eq!(upload.header("x-trace"), Some("upload-test")); + let multipart = upload.body_text(); + assert!( + multipart.contains(&format!("Content-Type: {mime_type}\r\n")), + "{multipart}" + ); + assert!(multipart.contains("\r\n\r\nabc\r\n--"), "{multipart}"); + assert_eq!(parse.url.path(), "/parse"); + assert_eq!( + parse.json(), + json!({source_field(model): "reducto://uploaded.pdf"}) + ); + for request in &requests { + assert_eq!(request.header("authorization"), Some("Bearer test-key")); + } +} + +#[tokio::test] +async fn response_received_fires_once_for_the_parse_response() { + let upstream = upstream([upload_response(), chunks_response(json!([]))]).await; + let observed = Arc::new(Mutex::new(Vec::new())); + let recorder = observed.clone(); + let host = LocalOcrHost::new(ocr_request("reducto/parse-v3", &upstream.uri(), json!({}))) + .with_observer(move |event| { + if let CallEvent::Machine(MachineEvent::ResponseReceived { raw }) = event { + recorder.lock().unwrap().push(raw.body.clone()); + } + }); + + perform_with(host).await.unwrap(); + + assert_eq!(received(&upstream).await.len(), 2); + assert_eq!(*observed.lock().unwrap(), [r#"{"result":{"chunks":[]}}"#]); +} + +#[rstest] +#[case::empty_id(json_response(json!({"file_id": ""})))] +#[case::missing_id(json_response(json!({})))] +#[case::null_id(json_response(json!({"file_id": null})))] +#[case::upload_failure(status_response(503, json!({"error": "unavailable"})))] +#[tokio::test] +async fn a_failed_upload_stops_before_parse(#[case] upload: ResponseTemplate) { + let upstream = upstream([upload]).await; + + let result = perform(ocr_request("reducto/parse-v3", &upstream.uri(), json!({}))).await; + + assert!(result.is_err()); + assert_eq!(received(&upstream).await.len(), 1); +} + +#[rstest] +#[case::remote_url("https://example.com/a.pdf", Error::ReductoSource)] +#[case::empty_file_id("reducto://", Error::RequestField { path: "document file id".into() })] +#[case::data_uri_without_payload("data:application/pdf;base64", Error::InvalidDataUri)] +#[case::invalid_base64("data:application/pdf;base64,INVALID!", Error::InvalidDataUri)] +#[tokio::test] +async fn invalid_document_sources_are_rejected_before_sending( + #[case] source: &str, + #[case] expected: Error, +) { + let upstream = upstream([json_response(json!({}))]).await; + + let result = perform(with_source( + ocr_request("reducto/parse-v3", &upstream.uri(), json!({})), + source, + )) + .await; + + assert!( + received(&upstream).await.is_empty(), + "sent invalid source: {source}" + ); + let error = result.unwrap_err(); + assert_eq!( + std::mem::discriminant(&error), + std::mem::discriminant(&expected) + ); + assert_eq!(error.http_status_code(), Some(400)); + assert_eq!(error.to_string(), expected.to_string()); +} + +#[tokio::test] +async fn a_forwarded_authorization_wins_and_the_native_response_is_omitted_by_default() { + let upstream = upstream([json_response( + json!({"job_id": "job-1", "result": {"chunks": []}}), + )]) + .await; + let request = with_headers( + with_source( + ocr_request("reducto/parse-v3", &upstream.uri(), json!({})), + "reducto://ready.pdf", + ), + &[("authorization", "Bearer existing")], + ); + + let response = perform(request).await.unwrap(); + + assert_eq!(response.provider_native_response, None); + assert_eq!( + only_request(&upstream).await.header_values("authorization"), + ["Bearer existing"] + ); +} + +#[tokio::test] +async fn native_format_retains_the_provider_response() { + let raw = json!({ + "result": {"chunks": [{"content": "native OCR response"}]}, + "usage": {"num_pages": 1} + }); + let upstream = upstream([json_response(raw.clone())]).await; + + let response = perform(with_source( + ocr_request( + "reducto/parse-v3", + &upstream.uri(), + json!({"req_format": "native"}), + ), + "reducto://ready.pdf", + )) + .await + .unwrap(); + + assert_eq!(response.pages[0].markdown, "native OCR response"); + assert_eq!( + response.provider_native_response.map(Value::Object), + Some(raw) + ); +} + +#[tokio::test] +async fn an_unknown_model_reaches_parse_and_keeps_its_name() { + let upstream = upstream([chunks_response( + json!([{"content": "future model response"}]), + )]) + .await; + + let response = perform(with_source( + ocr_request("reducto/future-parse-model", &upstream.uri(), json!({})), + "reducto://ready.pdf", + )) + .await + .unwrap(); + + assert_eq!(response.model, "future-parse-model"); + assert_eq!(response.pages[0].markdown, "future model response"); + let sent = only_request(&upstream).await; + assert_eq!(sent.url.path(), "/parse"); + assert_eq!(sent.json(), json!({"input": "reducto://ready.pdf"})); +} + +#[tokio::test] +async fn a_guardrail_can_replace_the_document_before_upload() { + let upstream = upstream([chunks_response(json!([]))]).await; + let host = LocalOcrHost::new(ocr_request("reducto/parse-v3", &upstream.uri(), json!({}))) + .with_before_send(|wire, _| { + assert_eq!(wire.body["document_url"], INLINE_PDF); + Ok(WireRequest { + body: json!({"type": "document_url", "document_url": "reducto://guarded.pdf"}), + ..wire + }) + }); + + perform_with(host).await.unwrap(); + + let sent = only_request(&upstream).await; + assert_eq!(sent.url.path(), "/parse"); + assert_eq!(sent.json(), json!({"input": "reducto://guarded.pdf"})); +} + +#[rstest] +#[tokio::test] +async fn guardrail_headers_reach_both_upload_and_parse( + #[values("reducto/parse-v3", "reducto/parse-legacy")] model: &str, +) { + let upstream = upstream([upload_response(), chunks_response(json!([]))]).await; + let request = with_headers( + ocr_request(model, &upstream.uri(), json!({})), + &[("authorization", "Bearer original")], + ); + let host = LocalOcrHost::new(request).with_before_send(|wire, _| { + Ok(WireRequest { + headers: vec![("authorization".into(), "Bearer guarded".into())], + ..wire + }) + }); + + perform_with(host).await.unwrap(); + + let requests = received(&upstream).await; + let paths: Vec<&str> = requests.iter().map(|request| request.url.path()).collect(); + assert_eq!(paths, ["/upload", "/parse"]); + for request in &requests { + assert_eq!(request.header_values("authorization"), ["Bearer guarded"]); + } +} diff --git a/litellm-rust/crates/core/tests/ocr/vertex_ai.rs b/litellm-rust/crates/core/tests/ocr/vertex_ai.rs new file mode 100644 index 00000000000..f0b2488e494 --- /dev/null +++ b/litellm-rust/crates/core/tests/ocr/vertex_ai.rs @@ -0,0 +1,184 @@ +use litellm_auth::{InputSource, Sourced}; +use litellm_core::ocr::arguments::is_supported_request; +use litellm_llms::base_llm::ocr::settings::OcrSettings; +use rstest::rstest; + +use super::*; + +#[tokio::test] +async fn mistral_is_served_at_the_resolved_project_and_location() { + let upstream = upstream([json_response(json!({ + "pages": [{"index": 0, "markdown": "hello"}], + "usage_info": {"pages_processed": 1} + }))]) + .await; + + let response = perform(ocr_request( + "vertex_ai/mistral-ocr-maas", + &upstream.uri(), + json!({ + "vertex_project": "project-1", + "vertex_location": "europe-west4", + "extract_footer": true + }), + )) + .await + .unwrap(); + + assert_eq!(response.pages[0].markdown, "hello"); + let sent = only_request(&upstream).await; + assert_eq!( + sent.url.path(), + "/v1/projects/project-1/locations/europe-west4/publishers/mistralai/models/mistral-ocr-maas:rawPredict" + ); + assert_eq!(sent.header("authorization"), Some("Bearer test-key")); + assert_eq!( + sent.json(), + json!({ + "model": "mistral-ocr-maas", + "document": {"type": "document_url", "document_url": INLINE_PDF}, + "extract_footer": true + }) + ); +} + +#[tokio::test] +async fn configured_project_and_location_apply_when_the_call_sets_neither() { + let upstream = upstream([pages_response()]).await; + let client = ocr_client().with_settings(OcrSettings { + vertex_project: Some("configured-project".into()), + vertex_location: Some("europe-west4".into()), + ..OcrSettings::default() + }); + + litellm_core::ocr::client::perform( + &client, + ocr_request("vertex_ai/mistral-ocr-maas", &upstream.uri(), json!({})), + ) + .await + .unwrap(); + + assert_eq!( + only_request(&upstream).await.url.path(), + "/v1/projects/configured-project/locations/europe-west4/publishers/mistralai/models/mistral-ocr-maas:rawPredict" + ); +} + +#[tokio::test] +async fn a_supplied_authorization_is_forwarded_without_a_static_token() { + let upstream = upstream([pages_response()]).await; + let request = with_headers( + without_api_key(ocr_request( + "vertex_ai/model", + &upstream.uri(), + json!({"vertex_project": "project-1"}), + )), + &[("authorization", "Bearer supplied")], + ); + + perform(request).await.unwrap(); + + assert_eq!( + only_request(&upstream).await.header_values("authorization"), + ["Bearer supplied"] + ); +} + +#[tokio::test] +async fn invalid_credentials_fail_before_sending() { + let error = perform(ocr_request( + "vertex_ai/model", + UNREACHABLE_BASE, + json!({"vertex_credentials": true}), + )) + .await + .unwrap_err(); + + assert!(error.to_string().contains("vertex_credentials"), "{error}"); +} + +#[rstest] +#[tokio::test] +async fn a_request_controlled_api_base_is_rejected_before_vertex_auth( + #[values("vertex_ai/mistral-ocr-maas", "vertex_ai/deepseek-ocr-maas")] model: &str, +) { + let mut request = ocr_request( + model, + "https://caller.example", + json!({"vertex_project": "project-1"}), + ); + request.credentials.api_base = Some(Sourced::new( + "https://caller.example".into(), + InputSource::Request, + )); + + let error = perform(request).await.unwrap_err(); + + assert!( + error + .to_string() + .contains("request-controlled Vertex AI endpoint"), + "{error}" + ); +} + +#[tokio::test] +async fn deepseek_is_served_at_the_openai_compatible_endpoint() { + let upstream = upstream([json_response(json!({ + "choices": [{"message": {"content": "recognized"}}], + "usage": {"prompt_tokens": 1} + }))]) + .await; + let request = with_source( + ocr_request( + "vertex_ai/deepseek-ocr-maas", + &upstream.uri(), + json!({ + "vertex_project": "project-1", + "vertex_location": "europe-west4", + "temperature": 0.1, + "future_ocr_option": true, + "extra_body": {"provider_option": "value"} + }), + ), + "gs://bucket/document.pdf", + ); + + let response = perform(request).await.unwrap(); + + assert_eq!(response.pages[0].markdown, "recognized"); + assert_eq!( + response.usage_info.unwrap().extra_fields["prompt_tokens"], + 1 + ); + let sent = only_request(&upstream).await; + assert_eq!( + sent.url.path(), + "/v1/projects/project-1/locations/europe-west4/endpoints/openapi/chat/completions" + ); + assert_eq!(sent.header("authorization"), Some("Bearer test-key")); + let body = sent.json(); + assert_eq!(body["model"], "deepseek-ai/deepseek-ocr-maas"); + assert_eq!(body["temperature"], 0.1); + assert_eq!(body["future_ocr_option"], true); + assert_eq!(body["provider_option"], "value"); + assert!(body.get("vertex_project").is_none()); + assert!(body.get("extra_body").is_none()); + assert_eq!( + body["messages"][0]["content"][0], + json!({"type": "image_url", "image_url": "gs://bucket/document.pdf"}) + ); +} + +#[rstest] +#[case::deepseek("deepseek-ocr-maas", Some("vertex_ai"), true)] +#[case::mistral("mistral-ocr-maas", Some("vertex_ai"), true)] +#[case::prefixed("vertex_ai/mistral-ocr-maas", None, true)] +#[case::unknown_provider("model", Some("unknown"), false)] +fn supported_requests_follow_the_registered_configs( + #[case] model: &str, + #[case] provider: Option<&str>, + #[case] supported: bool, +) { + assert_eq!(is_supported_request(model, provider), supported); +} diff --git a/litellm-rust/crates/core/tests/support/mod.rs b/litellm-rust/crates/core/tests/support/mod.rs new file mode 100644 index 00000000000..4d2fe0232d0 --- /dev/null +++ b/litellm-rust/crates/core/tests/support/mod.rs @@ -0,0 +1,155 @@ +//! Shared fixtures for route integration tests: a scripted upstream and a recording +//! secret source. + +#![allow(dead_code)] // each test binary compiles this module on its own and uses a different subset + +use std::sync::Mutex; + +use futures_util::future::BoxFuture; +use litellm_secrets::{SecretValue, source::SecretSource}; +use serde_json::Value; +use wiremock::{Mock, MockServer, Request, ResponseTemplate, matchers::any}; + +/// A port nothing listens on, for calls that must fail before any request is sent. +pub const UNREACHABLE_BASE: &str = "http://127.0.0.1:1"; + +/// Starts an upstream that answers its n-th request with the n-th response and 404s after. +pub async fn upstream(responses: impl IntoIterator) -> MockServer { + let server = MockServer::start().await; + respond_in_order(&server, responses).await; + server +} + +/// Scripts responses on a started server, for responses that need its address. +pub async fn respond_in_order( + server: &MockServer, + responses: impl IntoIterator, +) { + for response in responses { + Mock::given(any()) + .respond_with(response) + .up_to_n_times(1) + .mount(server) + .await; + } +} + +pub async fn received(server: &MockServer) -> Vec { + server + .received_requests() + .await + .expect("request recording is on") +} + +pub async fn only_request(server: &MockServer) -> Request { + let [request] = <[Request; 1]>::try_from(received(server).await) + .unwrap_or_else(|requests| panic!("expected one request, got {}", requests.len())); + request +} + +pub fn json_response(body: Value) -> ResponseTemplate { + ResponseTemplate::new(200).set_body_json(body) +} + +pub fn status_response(status: u16, body: Value) -> ResponseTemplate { + ResponseTemplate::new(status).set_body_json(body) +} + +pub trait ReceivedRequest { + fn header(&self, name: &str) -> Option<&str>; + fn header_values(&self, name: &str) -> Vec<&str>; + fn json(&self) -> Value; + fn body_text(&self) -> String; + /// The path and query, as the request line carried them. + fn target(&self) -> String; + fn query(&self, name: &str) -> Option; +} + +impl ReceivedRequest for Request { + fn header(&self, name: &str) -> Option<&str> { + self.headers.get(name).and_then(|value| value.to_str().ok()) + } + + fn header_values(&self, name: &str) -> Vec<&str> { + self.headers + .get_all(name) + .iter() + .filter_map(|value| value.to_str().ok()) + .collect() + } + + fn json(&self) -> Value { + serde_json::from_slice(&self.body).expect("request body is json") + } + + fn body_text(&self) -> String { + String::from_utf8_lossy(&self.body).into_owned() + } + + fn target(&self) -> String { + match self.url.query() { + Some(query) => format!("{}?{query}", self.url.path()), + None => self.url.path().to_string(), + } + } + + fn query(&self, name: &str) -> Option { + self.url + .query_pairs() + .find_map(|(key, value)| (key == name).then(|| value.into_owned())) + } +} + +/// A secret source that answers from a fixed table and records every name it was asked for. +pub struct RecordingSecrets { + values: Vec<(String, String)>, + fails: bool, + requested: Mutex>, +} + +impl RecordingSecrets { + pub fn new<'a>(values: impl IntoIterator) -> Self { + Self { + values: values + .into_iter() + .map(|(name, value)| (name.to_string(), value.to_string())) + .collect(), + fails: false, + requested: Mutex::new(Vec::new()), + } + } + + pub fn empty() -> Self { + Self::new([]) + } + + pub fn failing() -> Self { + Self { + fails: true, + ..Self::empty() + } + } + + pub fn requested(&self) -> Vec { + self.requested.lock().unwrap().clone() + } +} + +impl SecretSource for RecordingSecrets { + fn get_secret_str<'a>( + &'a self, + name: &'a str, + ) -> BoxFuture<'a, Result, litellm_secrets::Error>> { + Box::pin(async move { + self.requested.lock().unwrap().push(name.to_string()); + if self.fails { + return Err(litellm_secrets::Error::ManagedSecretMissing); + } + Ok(self + .values + .iter() + .find(|(key, _)| key == name) + .map(|(_, value)| SecretValue::new(value.clone()))) + }) + } +} diff --git a/litellm-rust/crates/llms/src/reducto/ocr/transformation.rs b/litellm-rust/crates/llms/src/reducto/ocr/transformation.rs index f00259984ba..c3377536545 100644 --- a/litellm-rust/crates/llms/src/reducto/ocr/transformation.rs +++ b/litellm-rust/crates/llms/src/reducto/ocr/transformation.rs @@ -554,6 +554,50 @@ async fn upload_bytes_async( mod tests { use super::*; + #[tokio::test] + async fn v3_body_keeps_explicit_null_options_and_drops_unknown_ones() { + use crate::base_llm::ocr::{handler::OcrClient, transformation::OcrRequestContext}; + + let overrides = + serde_json::from_value(json!({"formatting":null,"settings":{},"unknown":true})) + .unwrap(); + let params = ReductoParseV3Config + .map_ocr_params(&overrides, "parse-v3") + .unwrap(); + let client = OcrClient::for_test(reqwest::Client::new(), reqwest::Client::new()); + let connection = OcrConnection::default(); + let document = serde_json::from_value( + json!({"type":"document_url","document_url":"reducto://ready.pdf"}), + ) + .unwrap(); + + let body = ReductoParseV3Config + .async_transform_ocr_request( + "parse-v3", + document, + ¶ms, + &[], + OcrRequestContext { + client: &client, + connection: &connection, + }, + ) + .await + .unwrap(); + + assert_eq!( + serde_json::to_value(body).unwrap(), + json!({"input":"reducto://ready.pdf", "formatting":null, "settings":{}}) + ); + let absent = ReductoParseV3Config + .map_ocr_params( + &litellm_core_utils::call_arguments::CallArguments::default(), + "parse-v3", + ) + .unwrap(); + assert_eq!(serde_json::to_value(absent).unwrap(), json!({})); + } + #[test] fn options_preserve_null_and_select_the_provider_fields() { let overrides = serde_json::from_value(json!({ diff --git a/litellm-rust/crates/llms/tests/ocr_handler.rs b/litellm-rust/crates/llms/tests/ocr_handler.rs new file mode 100644 index 00000000000..6e46e6f76d4 --- /dev/null +++ b/litellm-rust/crates/llms/tests/ocr_handler.rs @@ -0,0 +1,79 @@ +use std::time::Duration; + +use litellm_llms::base_llm::ocr::{error::Error, handler::read_response_bytes}; +use rstest::rstest; +use tokio::{ + io::{AsyncReadExt, AsyncWriteExt}, + net::TcpListener, +}; + +/// Answers one request with raw `response` bytes and then holds the connection open, so a +/// read that waits for the rest of an oversized body hangs instead of passing. +async fn read_bounded(response: String, limit: usize) -> Result { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let address = listener.local_addr().unwrap(); + let server = tokio::spawn(async move { + let (mut socket, _) = listener.accept().await.unwrap(); + let mut request = [0; 4096]; + assert!(socket.read(&mut request).await.unwrap() > 0); + socket.write_all(response.as_bytes()).await.unwrap(); + std::future::pending::<()>().await; + }); + let response = reqwest::Client::new() + .get(format!("http://{address}")) + .send() + .await + .unwrap(); + let result = + tokio::time::timeout(Duration::from_secs(2), read_response_bytes(response, limit)).await; + server.abort(); + result.expect("bounded reads must finish without waiting for the rest of an oversized body") +} + +#[rstest] +#[case::declared("HTTP/1.1 200 OK\r\nContent-Length: 8\r\n\r\nabcdefgh")] +#[case::chunked( + "HTTP/1.1 200 OK\r\nTransfer-Encoding: chunked\r\n\r\n4\r\nabcd\r\n4\r\nefgh\r\n0\r\n\r\n" +)] +#[tokio::test] +async fn a_body_of_exactly_the_limit_is_read(#[case] response: &str) { + assert_eq!(read_bounded(response.into(), 8).await.unwrap(), "abcdefgh"); +} + +#[rstest] +#[case::declared("HTTP/1.1 200 OK\r\nContent-Length: 9\r\n\r\n")] +#[case::chunked("HTTP/1.1 200 OK\r\nTransfer-Encoding: chunked\r\n\r\n4\r\nabcd\r\n5\r\nefghi\r\n")] +#[tokio::test] +async fn a_body_over_the_limit_is_rejected(#[case] response: &str) { + assert!(matches!( + read_bounded(response.into(), 8).await, + Err(Error::TooLarge { limit: 8 }) + )); +} + +#[rstest] +#[case::declared("Content-Length: 1000000")] +#[case::chunked("Transfer-Encoding: chunked")] +#[tokio::test] +async fn an_oversized_error_keeps_its_status_and_a_bounded_body_without_draining( + #[case] headers: &str, +) { + let prefix = "x".repeat(4096); + let body = match headers.starts_with("Transfer") { + true => format!("{:x}\r\n{prefix}\r\n", prefix.len()), + false => prefix.clone(), + }; + + let error = read_bounded( + format!("HTTP/1.1 429 Too Many Requests\r\n{headers}\r\n\r\n{body}"), + prefix.len(), + ) + .await + .unwrap_err(); + + let Error::Transport(litellm_http::transport::Error::Http { status, body }) = error else { + panic!("unexpected error: {error}"); + }; + assert_eq!(status, 429); + assert_eq!(body, prefix); +} diff --git a/tests/test_litellm_rust/AGENTS.md b/tests/test_litellm_rust/AGENTS.md new file mode 100644 index 00000000000..d65ffd613aa --- /dev/null +++ b/tests/test_litellm_rust/AGENTS.md @@ -0,0 +1 @@ +This directory holds only the tests that cannot be written in the Rust code diff --git a/tests/test_litellm_rust/cache/__init__.py b/tests/test_litellm_rust/cache/__init__.py new file mode 100644 index 00000000000..8b137891791 --- /dev/null +++ b/tests/test_litellm_rust/cache/__init__.py @@ -0,0 +1 @@ + diff --git a/tests/test_litellm_rust/cache/conftest.py b/tests/test_litellm_rust/cache/conftest.py new file mode 100644 index 00000000000..07ccd0fde2b --- /dev/null +++ b/tests/test_litellm_rust/cache/conftest.py @@ -0,0 +1,30 @@ +import threading +from collections.abc import Generator +from typing import Final + +import fakeredis +import pytest + +from tests.test_litellm_rust.support.s3_stub import S3Stub + + +@pytest.fixture +def redis_url() -> Generator[str]: + server: Final = fakeredis.TcpFakeServer(("127.0.0.1", 0), server_type="redis") + worker: Final = threading.Thread(target=server.serve_forever, daemon=True) + worker.start() + try: + yield f"redis://127.0.0.1:{server.server_address[1]}" + finally: + server.shutdown() + server.server_close() + worker.join(timeout=5) + + +@pytest.fixture +def s3_stub() -> Generator[S3Stub]: + stub: Final = S3Stub() + try: + yield stub + finally: + stub.close() diff --git a/tests/test_litellm_rust/cache/test_azure_blob.py b/tests/test_litellm_rust/cache/test_azure_blob.py new file mode 100644 index 00000000000..bbbab22baca --- /dev/null +++ b/tests/test_litellm_rust/cache/test_azure_blob.py @@ -0,0 +1,173 @@ +import asyncio +import json +import os +import time +import uuid +from collections.abc import Generator +from types import SimpleNamespace +from typing import Final, cast + +import pytest +from azure.storage.blob import ContainerClient + +from litellm.caching.azure_blob_cache import AzureBlobCache +from litellm.caching.caching import Cache +from litellm.rust_bridge import _native +from litellm.types.caching import LiteLLMCacheType +from tests.test_litellm_rust.support.cache import ( + CacheLookup, + CacheTestHandle, + CacheTestResolver, + assert_native_runtime, + completion_kwargs, + request, + require_rust, +) +from tests.test_litellm_rust.support.isolation import rebound + +pytestmark: Final = pytest.mark.requires_rust_extension + + +@pytest.fixture +def azure_blob_facade() -> Generator[Cache]: + account_url: Final = os.environ.get("AZURE_BLOB_CACHE_ACCOUNT_URL") + if account_url is None: + pytest.skip( + "live Azure Blob parity needs AZURE_BLOB_CACHE_ACCOUNT_URL plus DefaultAzureCredential inputs in the environment" + ) + facade: Final = Cache( + type=LiteLLMCacheType.AZURE_BLOB, + azure_account_url=account_url, + azure_blob_container=f"litellm-parity-{uuid.uuid4().hex[:12]}", + ) + backend: Final = facade.cache + assert isinstance(backend, AzureBlobCache) + try: + yield facade + finally: + backend.container_client.delete_container() + asyncio.run(backend.disconnect()) + + +def azure_blob_handle(facade: Cache) -> _native._CacheTestHandle: + backend: Final = facade.cache + assert isinstance(backend, AzureBlobCache) + return CacheTestHandle.azure_blob( + backend.container_client.url.removesuffix(f"/{backend.container_client.container_name}"), + backend.container_client.container_name, + ) + + +def test_azure_blob_facade_serves_natively_and_python_reads_the_same_blobs(azure_blob_facade: Cache) -> None: + backend: Final = azure_blob_facade.cache + assert isinstance(backend, AzureBlobCache) + handle: Final = azure_blob_handle(azure_blob_facade) + assert handle.backend == "azure-blob" + account_url: Final = backend.container_client.url.removesuffix(f"/{backend.container_client.container_name}") + with pytest.raises(TypeError, match="containers must match"): + CacheTestHandle.azure_blob(account_url, f"{backend.container_client.container_name}-other")._bind_facade( + azure_blob_facade + ) + handle._bind_facade(azure_blob_facade) + resolver: Final = CacheTestResolver(SimpleNamespace(cache=azure_blob_facade)) + native: Final = resolver.resolve() + assert native.kind == "native" + + response: Final = { + "choices": [{"text": "caf\u00e9 \u2603"}], + "usage": {"total_tokens": 3}, + "flag": True, + "empty": None, + } + native.store({**request("sync"), "ttl_seconds": 0.001}, response) + native.store(request("sync"), {"choices": [{"text": "second"}]}) + time.sleep(0.01) + stored: Final = json.loads(backend.container_client.download_blob("sync").readall()) + assert stored["response"] == response + assert isinstance(stored["timestamp"], float) + assert native.lookup(request("sync")) == response + assert cast(CacheLookup, azure_blob_facade).get_cache(cache_key="sync") == response + + backend.set_cache("python", {"timestamp": time.time(), "response": response}) + backend.set_cache("legacy", "bare legacy value") + backend.container_client.upload_blob("invalid", b"{not json", overwrite=True) + assert native.lookup(request("python")) == response + assert native.lookup(request("legacy")) == cast(CacheLookup, azure_blob_facade).get_cache(cache_key="legacy") + assert native.lookup_batch([request("python"), request("missing"), request("invalid"), request("sync")]) == { + "values": [response, None, None, response], + "missing_indices": [1, 2], + } + + with rebound(azure_blob_facade, "ttl", 12): + assert resolver.resolve().kind == "python_callback" + with rebound(backend, "container_client", ContainerClient.from_container_url(backend.container_client.url)): + assert resolver.resolve().kind == "python_callback" + + def custom_get(*_args: object, **_kwargs: object) -> None: + return None + + with rebound(backend, "get_cache", custom_get): + assert resolver.resolve().kind == "python_callback" + assert resolver.resolve().kind == "python_callback" + assert cast(CacheLookup, azure_blob_facade).get_cache(cache_key="sync") == response + + class CustomBlobCache(AzureBlobCache): + pass + + with rebound(azure_blob_facade, "cache", CustomBlobCache(account_url, backend.container_client.container_name)): + assert resolver.resolve().kind == "python_callback" + with pytest.raises(TypeError): + azure_blob_handle(azure_blob_facade)._bind_facade(azure_blob_facade) + + +async def test_azure_blob_native_async_writes_overwrite_batch_and_flush_like_python(azure_blob_facade: Cache) -> None: + backend: Final = azure_blob_facade.cache + assert isinstance(backend, AzureBlobCache) + azure_blob_handle(azure_blob_facade)._bind_facade(azure_blob_facade) + binding: Final = CacheTestResolver(SimpleNamespace(cache=azure_blob_facade)).resolve() + assert binding.kind == "native" + ping: Final = cast(dict[str, object], await binding.ping()) + assert ping["status"] == "success", ping + + await binding.async_store(request("async"), {"value": 1}) + await binding.async_store({**request("async"), "ttl_seconds": 0.001}, {"value": 2}) + time.sleep(0.01) + assert await binding.async_lookup(request("async")) == {"value": 2} + assert await backend.async_get_cache("async") == json.loads( + backend.container_client.download_blob("async").readall() + ) + assert cast(CacheLookup, azure_blob_facade).get_cache(cache_key="async") == {"value": 2} + + await binding.async_store_batch([request("first"), request("second")], [{"value": 3}, {"value": 4}]) + assert await binding.async_lookup_batch([request("second"), request("missing"), request("first")]) == { + "values": [{"value": 4}, None, {"value": 3}], + "missing_indices": [1], + } + await binding.async_flush() + assert [blob.name for blob in backend.container_client.list_blobs()] == [] + assert await binding.async_lookup(request("async")) is None + + +async def test_azure_blob_rust_required_rule_activates_natively(monkeypatch: pytest.MonkeyPatch) -> None: + account_url: Final = os.environ.get("AZURE_BLOB_CACHE_ACCOUNT_URL") + if account_url is None: + pytest.skip( + "live Azure Blob parity needs AZURE_BLOB_CACHE_ACCOUNT_URL plus DefaultAzureCredential inputs in the environment" + ) + require_rust(monkeypatch, LiteLLMCacheType.AZURE_BLOB) + facade: Final = Cache( + type=LiteLLMCacheType.AZURE_BLOB, + azure_account_url=account_url, + azure_blob_container=f"litellm-parity-{uuid.uuid4().hex[:12]}", + ) + backend: Final = facade.cache + assert isinstance(backend, AzureBlobCache) + try: + assert_native_runtime(facade) + kwargs: Final = completion_kwargs("azure") + await facade.async_add_cache({"answer": "azure"}, **kwargs) + assert await facade.async_get_cache(**kwargs) == {"answer": "azure"} + assert backend.get_cache(facade.get_cache_key(**kwargs))["response"] == {"answer": "azure"} + finally: + backend.container_client.delete_container() + await backend.disconnect() diff --git a/tests/test_litellm_rust/cache/test_disk.py b/tests/test_litellm_rust/cache/test_disk.py new file mode 100644 index 00000000000..4f2907e6a09 --- /dev/null +++ b/tests/test_litellm_rust/cache/test_disk.py @@ -0,0 +1,117 @@ +import asyncio +import json +import time +from pathlib import Path +from types import SimpleNamespace +from typing import Final + +import diskcache +import pytest + +from litellm.caching.caching import Cache +from litellm.caching.disk_cache import DiskCache +from litellm.types.caching import LiteLLMCacheType +from tests.test_litellm_rust.support.cache import CacheTestHandle, CacheTestResolver, request +from tests.test_litellm_rust.support.isolation import rebound + +pytestmark: Final = pytest.mark.requires_rust_extension + + +async def test_disk_reads_python_entries_and_python_reads_native_entries(tmp_path: Path) -> None: + disk_cache: Final = DiskCache(disk_cache_dir=str(tmp_path)) + response: Final = {"choices": [{"text": "cached"}], "usage": {"total_tokens": 3}} + disk_cache.disk_cache.set( + "sync", + {"timestamp": time.time(), "response": json.dumps(response)}, + ) + disk_cache.disk_cache.set("async", json.dumps({"timestamp": time.time(), "response": response})) + disk_cache.disk_cache.set("raw", json.dumps(response)) + disk_cache.disk_cache.set("invalid", "not a cache entry") + disk_cache.disk_cache.set( + "large", + {"timestamp": time.time(), "response": {"text": "x" * 70_000}}, + ) + binding: Final = CacheTestResolver(SimpleNamespace(cache=CacheTestHandle.disk(str(tmp_path)))).resolve() + + assert binding.lookup(request("sync")) == response + assert await binding.async_lookup(request("async")) == response + assert binding.lookup(request("raw")) == response + assert await binding.async_lookup(request("invalid")) is None + assert binding.lookup(request("large")) == {"text": "x" * 70_000} + + await binding.async_store({**request("native"), "ttl_seconds": 12.0}, response) + stored_response: Final = disk_cache.get_cache("native") + assert isinstance(stored_response, dict) + assert stored_response["response"] == response + stored, expire_time = disk_cache.disk_cache.get("native", expire_time=True) + assert stored is not None + assert time.time() < expire_time <= time.time() + 12.0 + await binding.async_store(request("no-ttl"), response) + _, no_expiry = disk_cache.disk_cache.get("no-ttl", expire_time=True) + assert no_expiry is None + + +async def test_disk_entries_survive_a_fresh_handle_and_expire_on_time(tmp_path: Path) -> None: + first: Final = CacheTestResolver(SimpleNamespace(cache=CacheTestHandle.disk(str(tmp_path)))).resolve() + await first.async_store(request("persistent"), {"value": "persistent"}) + await first.async_store({**request("expiring"), "ttl_seconds": 0.3}, {"value": "expiring"}) + fresh: Final = CacheTestResolver(SimpleNamespace(cache=CacheTestHandle.disk(str(tmp_path)))).resolve() + assert fresh.lookup(request("persistent")) == {"value": "persistent"} + assert fresh.lookup(request("expiring")) == {"value": "expiring"} + await asyncio.sleep(0.4) + assert fresh.lookup(request("expiring")) is None + assert fresh.lookup(request("persistent")) == {"value": "persistent"} + + +def test_disk_facade_registers_and_store_changes_fall_back(tmp_path: Path) -> None: + facade: Final = Cache(type=LiteLLMCacheType.DISK, disk_cache_dir=str(tmp_path)) + with pytest.raises(TypeError, match="directories must match"): + CacheTestHandle.disk(str(tmp_path / "other"))._bind_facade(facade) + handle: Final = CacheTestHandle.disk(str(tmp_path)) + handle._bind_facade(facade) + resolver: Final = CacheTestResolver(SimpleNamespace(cache=facade)) + binding: Final = resolver.resolve() + assert binding.kind == "native" + binding.store(request("native"), {"value": "native"}) + assert facade.get_cache(cache_key="native") == {"value": "native"} + + with rebound(facade.cache, "disk_cache", diskcache.Cache(str(tmp_path))): + assert resolver.resolve().kind == "python_callback" + assert resolver.resolve().kind == "native" + + class CustomDiskCache(DiskCache): + pass + + with rebound(facade, "cache", CustomDiskCache(disk_cache_dir=str(tmp_path))): + assert resolver.resolve().kind == "python_callback" + + class CustomStore(diskcache.Cache): + pass + + custom_facade: Final = Cache(type=LiteLLMCacheType.DISK, disk_cache_dir=str(tmp_path)) + custom_facade.cache.disk_cache = CustomStore(str(tmp_path)) + with pytest.raises(TypeError, match="built-in diskcache store"): + CacheTestHandle.disk(str(tmp_path))._bind_facade(custom_facade) + + +async def test_disk_native_batch_lookup_and_store_report_partial_hits(tmp_path: Path) -> None: + binding: Final = CacheTestResolver(SimpleNamespace(cache=CacheTestHandle.disk(str(tmp_path)))).resolve() + requests: Final = [request("hit"), request("miss"), request("disabled")] + requests[2]["controls"] = { + "supported_call_type": True, + "configured": True, + "native_backend": True, + "default_on": True, + "caching": False, + "no_cache": False, + "no_store": False, + "use_cache": False, + } + await binding.async_store_batch(requests, [{"value": 1}, {"value": 2}, {"value": 3}]) + + partial: Final = await binding.async_lookup_batch(requests) + + assert partial == { + "values": [{"value": 1}, {"value": 2}, None], + "missing_indices": [2], + } diff --git a/tests/test_litellm_rust/cache/test_facade.py b/tests/test_litellm_rust/cache/test_facade.py new file mode 100644 index 00000000000..d99ea4e2baa --- /dev/null +++ b/tests/test_litellm_rust/cache/test_facade.py @@ -0,0 +1,397 @@ +import asyncio +import contextvars +import gc +import weakref +from types import SimpleNamespace +from typing import Final, cast + +import pytest + +import litellm +from litellm.caching.caching import Cache, disable_cache, enable_cache, update_cache +from litellm.caching.in_memory_cache import InMemoryCache +from litellm.rust_bridge import _native +from litellm.rust_bridge.catalog import CacheRule, Route, RouteRule, SecretManagerRule +from litellm.rust_bridge.configuration import Rollout +from litellm.rust_bridge.response_cache import ResponseCacheRuntime, resolve_response_cache +from litellm.types.caching import LiteLLMCacheType +from tests.test_litellm_rust.support.cache import CacheLookup, CacheTestHandle, CacheTestResolver, request +from tests.test_litellm_rust.support.isolation import rebound + +pytestmark: Final = pytest.mark.requires_rust_extension + + +def test_existing_constructor_and_global_are_unchanged() -> None: + facade: Final = Cache(type=LiteLLMCacheType.LOCAL) + assert type(facade.cache) is InMemoryCache + assert "_native_cache_handle" not in vars(facade) + assert resolve_response_cache(facade) is None + with rebound(litellm, "cache", facade): + resolver: Final = CacheTestResolver(litellm) + assert resolver.resolve().kind == "python_callback" + resolver.resolve().store(None, {"answer": 7}, callback_kwargs={"cache_key": "key"}) + assert cast(CacheLookup, facade).get_cache(cache_key="key") == {"answer": 7} + + +async def test_catalog_constructs_native_runtime_from_public_cache_configuration() -> None: + rules: Final = ( + RouteRule(Route.OCR, Rollout.PYTHON_ONLY), + SecretManagerRule(Rollout.PYTHON_ONLY, systems=frozenset({"local"})), + CacheRule(Rollout.RUST_REQUIRED, backends=frozenset({"local"})), + ) + facade: Final = Cache(type=LiteLLMCacheType.LOCAL) + runtime: Final = resolve_response_cache(facade, rules) + assert isinstance(runtime, ResponseCacheRuntime) + assert runtime.kind == "native" + + sync_request: Final = runtime.request(facade, {"cache_key": "sync"}) + assert sync_request is not None + runtime.store(sync_request, {"answer": 1}) + assert runtime.lookup(sync_request) == {"answer": 1} + assert facade.cache.get_cache("sync") is None + + async_request: Final = runtime.request(facade, {"cache_key": "async"}) + assert async_request is not None + await runtime.async_store(async_request, {"answer": 2}) + assert await runtime.async_lookup(async_request) == {"answer": 2} + assert await facade.cache.async_get_cache("async") is None + + requests: Final = (sync_request, async_request) + expected: Final = { + "values": [{"answer": 1}, {"answer": 2}], + "missing_indices": [], + } + assert runtime.lookup_batch(requests) == expected + assert await runtime.async_lookup_batch(requests) == expected + + await runtime.async_flush() + assert runtime.lookup(sync_request) is None + assert await runtime.async_lookup(async_request) is None + + +async def test_inference_resolver_uses_the_configured_native_cache_directly() -> None: + rules: Final = ( + RouteRule(Route.OCR, Rollout.PYTHON_ONLY), + SecretManagerRule(Rollout.PYTHON_ONLY, systems=frozenset({"local"})), + CacheRule(Rollout.RUST_REQUIRED, backends=frozenset({"local"})), + ) + facade: Final = Cache(type=LiteLLMCacheType.LOCAL) + runtime: Final = resolve_response_cache(facade, rules) + assert isinstance(runtime, ResponseCacheRuntime) + facade._native_cache = runtime + + selected: Final = _native._CacheResolver(SimpleNamespace(cache=facade)).resolve() + assert selected.kind == "native" + request: Final = runtime.request(facade, {"cache_key": "inference-native"}) + assert request is not None + await selected.async_store(request, {"answer": 42}) + assert await selected.async_lookup(request) == {"answer": 42} + assert await runtime.async_lookup(request) == {"answer": 42} + assert facade.cache.get_cache("inference-native") is None + + facade._native_cache = None + fallback: Final = _native._CacheResolver(SimpleNamespace(cache=facade)).resolve() + assert fallback.kind == "python_callback" + await fallback.async_store(None, {"answer": 7}, callback_kwargs={"cache_key": "inference-python"}) + assert facade.get_cache(cache_key="inference-python") == {"answer": 7} + assert facade.cache.get_cache("inference-python") is not None + + +async def test_inference_resolver_declines_a_native_runtime_whose_facade_changed() -> None: + rules: Final = ( + RouteRule(Route.OCR, Rollout.PYTHON_ONLY), + SecretManagerRule(Rollout.PYTHON_ONLY, systems=frozenset({"local"})), + CacheRule(Rollout.RUST_REQUIRED, backends=frozenset({"local"})), + ) + facade: Final = Cache(type=LiteLLMCacheType.LOCAL) + runtime: Final = resolve_response_cache(facade, rules) + assert isinstance(runtime, ResponseCacheRuntime) + facade._native_cache = runtime + stale_request: Final = runtime.request(facade, {"cache_key": "stale-only"}) + assert stale_request is not None + await runtime.async_store(stale_request, {"answer": "stale"}) + + replacement: Final = InMemoryCache() + facade.cache = replacement + with pytest.raises(_native.RustBridgeDeclined): + _native._CacheResolver(SimpleNamespace(cache=facade)).resolve() + assert await runtime.async_lookup(stale_request) == {"answer": "stale"} + assert replacement.get_cache("stale-only") is None + assert replacement.get_cache("swapped-backend") is None + + +def test_existing_global_lifecycle_remains_the_resolver_source_of_truth() -> None: + resolver: Final = CacheTestResolver(litellm) + + enable_cache(type=LiteLLMCacheType.LOCAL, ttl=30) + enabled: Final = litellm.cache + assert isinstance(enabled, Cache) + assert enabled.ttl == 30 + assert resolver.resolve().kind == "python_callback" + + enable_cache(type=LiteLLMCacheType.LOCAL, ttl=60) + assert litellm.cache is enabled + + update_cache(type=LiteLLMCacheType.LOCAL, ttl=60) + updated: Final = litellm.cache + assert isinstance(updated, Cache) + assert updated is not enabled + assert updated.ttl == 60 + + disable_cache() + assert litellm.cache is None + assert resolver.resolve().kind == "disabled" + + +async def test_native_bindings_survive_replacement_and_capture_writes_before_dispatch() -> None: + namespace: Final = SimpleNamespace(cache=CacheTestHandle.memory()) + resolver: Final = CacheTestResolver(namespace) + selected: Final = resolver.resolve() + assert selected.kind == "native" + selected.store(request(), {"answer": 1}) + assert await selected.async_lookup(request()) == {"answer": 1} + with rebound(namespace, "cache", CacheTestHandle.memory()): + replacement: Final = resolver.resolve() + await selected.async_store(request(), {"answer": 2}) + assert replacement.lookup(request()) is None + assert selected.lookup(request()) == {"answer": 2} + with rebound(namespace, "cache", None): + disabled: Final = resolver.resolve() + assert disabled.kind == "disabled" + assert disabled.lookup(None) is None + await disabled.async_store(None, object()) + assert await disabled.async_lookup(None) is None + assert selected.lookup(request()) == {"answer": 2} + + +async def test_python_callback_preserves_identity_caller_task_context_and_errors() -> None: + context: Final = contextvars.ContextVar("cache_context", default="caller") + caller: Final = asyncio.current_task() + sentinel: Final = object() + failure: Final = RuntimeError("callback failed") + + class CustomCache: + async def async_get_cache(self, *, marker: object) -> object: + assert marker is sentinel + assert asyncio.current_task() is caller + context.set("callback") + return marker + + async def async_add_cache(self, response: object, *, marker: object) -> None: + assert response is sentinel + assert marker is sentinel + raise failure + + namespace: Final = SimpleNamespace(cache=CustomCache()) + binding: Final = CacheTestResolver(namespace).resolve() + assert binding.kind == "python_callback" + assert await binding.async_lookup(None, callback_kwargs={"marker": sentinel}) is sentinel + assert context.get() == "callback" + with pytest.raises(RuntimeError) as caught: + await binding.async_store(None, sentinel, callback_kwargs={"marker": sentinel}) + assert caught.value is failure + + +async def test_callback_cancellation_stays_in_the_callers_task() -> None: + entered: Final = asyncio.Event() + finished: Final = asyncio.Event() + + class CustomCache: + async def async_get_cache(self) -> None: + entered.set() + try: + await asyncio.Future() + finally: + finished.set() + + binding: Final = CacheTestResolver(SimpleNamespace(cache=CustomCache())).resolve() + + async def lookup() -> object: + return await binding.async_lookup(None, callback_kwargs={}) + + task: Final = asyncio.create_task(lookup()) + await entered.wait() + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + assert finished.is_set() + + +def test_registered_facade_uses_native_and_instance_overrides_fall_back() -> None: + facade: Final = Cache(type=LiteLLMCacheType.LOCAL) + handle: Final = CacheTestHandle.memory() + handle._bind_facade(facade) + resolver: Final = CacheTestResolver(SimpleNamespace(cache=facade)) + native: Final = resolver.resolve() + assert native.kind == "native" + native.store(request(), {"source": "native"}) + assert native.lookup(request()) == {"source": "native"} + assert cast(CacheLookup, facade).get_cache(cache_key="key") is None + sentinel: Final = object() + + def outer_override(**_kwargs: object) -> object: + return sentinel + + def backend_override(*_args: object, **_kwargs: object) -> dict[str, str]: + return {"source": "override"} + + with rebound(facade, "get_cache", outer_override): + fallback: Final = resolver.resolve() + assert fallback.kind == "python_callback" + assert fallback.lookup(None, callback_kwargs={"cache_key": "key"}) is sentinel + assert resolver.resolve().kind == "python_callback" + delattr(facade, "get_cache") + assert resolver.resolve().kind == "native" + with rebound(facade.cache, "get_cache", backend_override): + backend_fallback: Final = resolver.resolve() + assert backend_fallback.kind == "python_callback" + assert backend_fallback.lookup(None, callback_kwargs={"cache_key": "key"}) == {"source": "override"} + + +def test_facade_subclasses_backend_replacement_and_configuration_changes_are_not_bypassed() -> None: + class CustomCache(Cache): + pass + + handle: Final = CacheTestHandle.memory() + with pytest.raises(TypeError): + handle._bind_facade(CustomCache(type=LiteLLMCacheType.LOCAL)) + facade: Final = Cache(type=LiteLLMCacheType.LOCAL) + handle._bind_facade(facade) + resolver: Final = CacheTestResolver(SimpleNamespace(cache=facade)) + with rebound(facade, "cache", InMemoryCache()): + assert resolver.resolve().kind == "python_callback" + with rebound(facade, "ttl", 12): + assert resolver.resolve().kind == "python_callback" + with rebound(facade, "semantic_cache_scope", "end_user"): + assert resolver.resolve().kind == "python_callback" + + def custom_key(**_kwargs: object) -> str: + return "custom" + + with rebound(facade, "get_cache_key", custom_key): + assert resolver.resolve().kind == "python_callback" + assert resolver.resolve().kind == "python_callback" + delattr(facade, "get_cache_key") + assert resolver.resolve().kind == "native" + + +def test_resolver_and_callback_cycles_can_be_collected() -> None: + class CustomCache: + pass + + def cyclic_reference() -> weakref.ReferenceType[CustomCache]: + callback: Final = CustomCache() + namespace: Final = SimpleNamespace(cache=callback) + binding: Final = CacheTestResolver(namespace).resolve() + setattr(callback, "binding", binding) + return weakref.ref(callback) + + reference: Final = cyclic_reference() + gc.collect() + assert reference() is None + + +def test_invalid_duration_and_request_shape_fail_before_storage() -> None: + binding: Final = CacheTestResolver(SimpleNamespace(cache=CacheTestHandle.memory())).resolve() + for seconds in (-1.0, float("nan"), float("inf")): + with pytest.raises(ValueError, match="cache durations must be finite and nonnegative"): + binding.store({**request(), "ttl_seconds": seconds}, {"answer": 1}) + assert binding.lookup(request()) is None + with pytest.raises(ValueError, match="cache durations must be finite and nonnegative"): + CacheTestHandle.memory(ttl_seconds=-1) + + +async def test_memory_size_policy_is_applied_by_the_native_host() -> None: + handle: Final = CacheTestHandle.memory(capacity=2, max_entry_bytes=128) + binding: Final = CacheTestResolver(SimpleNamespace(cache=handle)).resolve() + small: Final = {"answer": "ok"} + binding.store(request("small"), small) + assert await binding.async_lookup(request("small")) == small + await binding.async_store(request("large"), {"answer": "x" * 256}) + assert binding.lookup(request("large")) is None + assert binding.lookup(request("small")) == small + disabled: Final = CacheTestResolver(SimpleNamespace(cache=CacheTestHandle.memory(capacity=0))).resolve() + await disabled.async_store(request(), small) + assert await disabled.async_lookup(request()) is None + + +async def test_native_batch_lookup_and_store_report_partial_hits() -> None: + binding: Final = CacheTestResolver(SimpleNamespace(cache=CacheTestHandle.memory())).resolve() + requests: Final = [request("hit"), request("miss"), request("disabled")] + requests[2]["controls"] = { + "supported_call_type": True, + "configured": True, + "native_backend": True, + "default_on": True, + "caching": False, + "no_cache": False, + "no_store": False, + "use_cache": False, + } + await binding.async_store_batch(requests, [{"value": 1}, {"value": 2}, {"value": 3}]) + + partial: Final = await binding.async_lookup_batch(requests) + + assert partial == { + "values": [{"value": 1}, {"value": 2}, None], + "missing_indices": [2], + } + + +async def test_python_batch_callbacks_use_the_builtin_cache_api() -> None: + result: Final = object() + marker: Final = object() + + class CustomCache(Cache): + def get_cache(self, dynamic_cache_object: object = None, **kwargs: object) -> object: + return ("sync", kwargs) + + async def async_get_cache(self, dynamic_cache_object: object = None, **kwargs: object) -> object: + return ("async", kwargs) + + async def async_add_cache_pipeline( + self, result: object, dynamic_cache_object: object = None, **kwargs: object + ) -> object: + return result, kwargs + + binding: Final = CacheTestResolver(SimpleNamespace(cache=CustomCache(type=LiteLLMCacheType.LOCAL))).resolve() + assert binding.kind == "python_callback" + requests: Final = [request("first"), request("second")] + kwargs: Final = [{"cache_key": "first"}, {"cache_key": "second"}] + + assert binding.lookup_batch(requests, callback_kwargs=kwargs) == [("sync", kwargs[0]), ("sync", kwargs[1])] + assert await binding.async_lookup_batch(requests, callback_kwargs=kwargs) == [ + ("async", kwargs[0]), + ("async", kwargs[1]), + ] + with pytest.raises(ValueError, match="equal lengths"): + binding.lookup_batch(requests, callback_kwargs=kwargs[:1]) + with pytest.raises(TypeError, match="callback_result"): + await binding.async_store_batch(requests, [1, 2], callback_kwargs={"marker": marker}) + stored: Final = cast( + tuple[object, dict[str, object]], + await binding.async_store_batch(requests, [1, 2], callback_result=result, callback_kwargs={"marker": marker}), + ) + assert stored[0] is result + assert stored[1] == {"marker": marker} + + +async def test_unmodified_builtin_cache_callbacks_can_ping_and_flush() -> None: + async def ping() -> str: + return "pong" + + cache: Final = Cache(type=LiteLLMCacheType.LOCAL) + cache.cache.set_cache("key", "value") + binding: Final = CacheTestResolver(SimpleNamespace(cache=cache)).resolve() + assert binding.kind == "python_callback" + + setattr(cache.cache, "ping", ping) + assert await binding.ping() == "pong" + await binding.async_flush() + assert cache.cache.get_cache("key") is None + + +def test_facade_registration_rejects_mismatched_capacity() -> None: + facade: Final = Cache(type=LiteLLMCacheType.LOCAL) + with pytest.raises(TypeError, match="capacities must match"): + CacheTestHandle.memory(capacity=7)._bind_facade(facade) diff --git a/tests/test_litellm_rust/cache/test_gcs.py b/tests/test_litellm_rust/cache/test_gcs.py new file mode 100644 index 00000000000..bfc9ebbb4d7 --- /dev/null +++ b/tests/test_litellm_rust/cache/test_gcs.py @@ -0,0 +1,242 @@ +import json +import time +from collections.abc import Generator +from types import SimpleNamespace +from typing import Final, cast + +import pytest + +from litellm.caching.caching import Cache +from litellm.caching.gcs_cache import GCSCache +from litellm.types.caching import LiteLLMCacheType +from tests.test_litellm_rust.support.cache import CacheLookup, CacheTestHandle, CacheTestResolver, request +from tests.test_litellm_rust.support.fake_gcs import FakeGcs +from tests.test_litellm_rust.support.isolation import rebound + +pytestmark: Final = pytest.mark.requires_rust_extension + + +@pytest.fixture +def fake_gcs() -> Generator[FakeGcs]: + server: Final = FakeGcs() + try: + yield server + finally: + server.close() + + +async def test_gcs_reads_python_entries_and_writes_python_compatible_objects( + fake_gcs: FakeGcs, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.delenv("GCS_PATH_SERVICE_ACCOUNT", raising=False) + monkeypatch.delenv("GCS_BUCKET_NAME", raising=False) + response: Final = {"choices": [{"text": "cached"}], "usage": {"total_tokens": 3}, "flag": True, "empty": None} + fake_gcs.put( + "bucket", + "cache/sync", + json.dumps({"timestamp": time.time(), "response": json.dumps(response)}).encode(), + ) + fake_gcs.put("bucket", "cache/async", json.dumps({"timestamp": time.time(), "response": response}).encode()) + fake_gcs.put("bucket", "cache/raw", json.dumps(response).encode()) + fake_gcs.put("bucket", "cache/invalid", b"not a cache entry") + binding: Final = CacheTestResolver( + SimpleNamespace( + cache=CacheTestHandle.gcs( + "bucket", + gcs_path="cache", + endpoint=fake_gcs.url, + token=fake_gcs.token, + ) + ) + ).resolve() + + assert binding.lookup(request("sync")) == response + assert await binding.async_lookup(request("async")) == response + assert binding.lookup(request("raw")) == response + assert await binding.async_lookup(request("invalid")) is None + assert binding.lookup(request("missing")) is None + + await binding.async_store({**request("native"), "ttl_seconds": 12.0}, response) + stored: Final = fake_gcs.objects[("bucket", "cache/native")] + stored_value: Final = cast(dict[str, object], json.loads(stored)) + assert stored_value["response"] == response + assert isinstance(stored_value["timestamp"], float) + upload: Final = next(item for item in fake_gcs.requests if item.method == "POST") + assert upload.path == "/upload/storage/v1/b/bucket/o" + assert upload.query == "uploadType=media&name=cache%2Fnative" + assert upload.headers["Authorization"] == f"Bearer {fake_gcs.token}" + assert upload.headers["Content-Type"] == "application/json" + upload_text: Final = f"{upload.path}?{upload.query}{upload.headers}" + assert "ttl" not in upload_text.lower() + assert "expiry" not in upload_text.lower() + download: Final = next(item for item in fake_gcs.requests if item.path.endswith("/cache%2Fsync")) + assert download.path == "/storage/v1/b/bucket/o/cache%2Fsync" + assert download.query == "alt=media" + + binding.store(request("sync2"), response) + assert binding.lookup(request("sync2")) == response + assert GCSCache(bucket_name="bucket", gcs_path="cache").key_prefix == "cache/" + assert GCSCache(bucket_name="bucket", gcs_path="cache/").key_prefix == "cache/" + assert GCSCache(bucket_name="bucket").key_prefix == "" + + +async def test_gcs_batch_lookup_preserves_order_and_treats_malformed_entries_as_misses(fake_gcs: FakeGcs) -> None: + fake_gcs.put("bucket", "cache/hit", json.dumps({"timestamp": time.time(), "response": {"value": 1}}).encode()) + fake_gcs.put("bucket", "cache/invalid", b"not a cache entry") + binding: Final = CacheTestResolver( + SimpleNamespace( + cache=CacheTestHandle.gcs( + "bucket", + gcs_path="cache", + endpoint=fake_gcs.url, + token=fake_gcs.token, + ) + ) + ).resolve() + requests: Final = [request("hit"), request("missing"), request("invalid")] + expected: Final = {"values": [{"value": 1}, None, None], "missing_indices": [1, 2]} + + assert await binding.async_lookup_batch(requests) == expected + assert binding.lookup_batch(requests) == expected + await binding.async_store_batch([request("first"), request("second")], [{"value": 1}, {"value": 2}]) + assert ("bucket", "cache/first") in fake_gcs.objects + assert ("bucket", "cache/second") in fake_gcs.objects + + +async def test_gcs_facade_binds_only_exact_matching_configuration( + fake_gcs: FakeGcs, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.delenv("GCS_PATH_SERVICE_ACCOUNT", raising=False) + monkeypatch.delenv("GCS_BUCKET_NAME", raising=False) + monkeypatch.setenv("GOOGLE_APPLICATION_CREDENTIALS", "/nonexistent") + facade: Final = Cache(type=LiteLLMCacheType.GCS, gcs_bucket_name="bucket", gcs_path="cache/") + assert type(facade.cache) is GCSCache + + mismatched_bucket: Final = CacheTestHandle.gcs( + "other", + gcs_path="cache", + endpoint=fake_gcs.url, + token=fake_gcs.token, + ) + with pytest.raises(TypeError, match="buckets must match"): + mismatched_bucket._bind_facade(facade) + mismatched_prefix: Final = CacheTestHandle.gcs( + "bucket", + gcs_path="x", + endpoint=fake_gcs.url, + token=fake_gcs.token, + ) + with pytest.raises(TypeError, match="key prefixes must match"): + mismatched_prefix._bind_facade(facade) + mismatched_credentials: Final = CacheTestHandle.gcs( + "bucket", + gcs_path="cache", + path_service_account="sa.json", + endpoint=fake_gcs.url, + token=fake_gcs.token, + ) + with pytest.raises(TypeError, match="credentials must match"): + mismatched_credentials._bind_facade(facade) + with pytest.raises(TypeError, match="types must match"): + CacheTestHandle.memory()._bind_facade(facade) + + matching: Final = CacheTestHandle.gcs( + "bucket", + gcs_path="cache", + endpoint=fake_gcs.url, + token=fake_gcs.token, + ) + matching._bind_facade(facade) + resolver: Final = CacheTestResolver(SimpleNamespace(cache=facade)) + binding: Final = resolver.resolve() + assert binding.kind == "native" + await binding.async_store(request("native"), {"value": "native"}) + assert await binding.async_lookup(request("native")) == {"value": "native"} + assert cast(CacheLookup, facade).get_cache(cache_key="native") is None + + with rebound(facade.cache, "bucket_name", "other"): + assert resolver.resolve().kind == "python_callback" + with rebound(facade.cache, "key_prefix", "x/"): + assert resolver.resolve().kind == "python_callback" + with rebound(facade.cache, "path_service_account", "sa.json"): + assert resolver.resolve().kind == "python_callback" + + def no_get_cache(*args: object, **kwargs: object) -> None: + return None + + with rebound(facade.cache, "get_cache", no_get_cache): + assert resolver.resolve().kind == "python_callback" + with rebound(facade, "ttl", 12): + assert resolver.resolve().kind == "python_callback" + + class CustomGcs(GCSCache): + pass + + with rebound(facade, "cache", CustomGcs(bucket_name="bucket", gcs_path="cache/")): + assert resolver.resolve().kind == "python_callback" + custom_facade: Final = Cache(type=LiteLLMCacheType.GCS, gcs_bucket_name="bucket", gcs_path="cache/") + with rebound(custom_facade, "cache", CustomGcs(bucket_name="bucket", gcs_path="cache/")): + with pytest.raises(TypeError, match="types must match"): + matching._bind_facade(custom_facade) + + missing_bucket: Final = Cache(type=LiteLLMCacheType.GCS) + with pytest.raises(TypeError, match="requires a configured bucket name"): + matching._bind_facade(missing_bucket) + + +async def test_gcs_flush_is_a_no_op_and_ping_is_not_implemented( + fake_gcs: FakeGcs, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.delenv("GCS_PATH_SERVICE_ACCOUNT", raising=False) + monkeypatch.delenv("GCS_BUCKET_NAME", raising=False) + binding: Final = CacheTestResolver( + SimpleNamespace( + cache=CacheTestHandle.gcs( + "bucket", + gcs_path="cache", + endpoint=fake_gcs.url, + token=fake_gcs.token, + ) + ) + ).resolve() + await binding.async_store(request("key"), {"value": "stored"}) + await binding.async_flush() + assert ("bucket", "cache/key") in fake_gcs.objects + assert await binding.async_lookup(request("key")) == {"value": "stored"} + with pytest.raises(NotImplementedError): + await binding.ping() + + facade: Final = Cache(type=LiteLLMCacheType.GCS, gcs_bucket_name="bucket", gcs_path="cache/") + with pytest.raises(AttributeError): + await facade.ping() + assert cast(CacheLookup, facade.cache).flush_cache() is None + + +async def test_gcs_unauthorized_and_server_errors_surface_as_runtime_errors(fake_gcs: FakeGcs) -> None: + wrong_token: Final = CacheTestResolver( + SimpleNamespace( + cache=CacheTestHandle.gcs( + "bucket", + gcs_path="cache", + endpoint=fake_gcs.url, + token="wrong-token", + ) + ) + ).resolve() + with pytest.raises(RuntimeError): + wrong_token.lookup(request("missing")) + assert not fake_gcs.objects + + binding: Final = CacheTestResolver( + SimpleNamespace( + cache=CacheTestHandle.gcs( + "bucket", + gcs_path="cache", + endpoint=fake_gcs.url, + token=fake_gcs.token, + ) + ) + ).resolve() + with pytest.raises(RuntimeError): + binding.lookup(request("server-error")) + assert binding.lookup(request("missing")) is None diff --git a/tests/test_litellm_rust/cache/test_qdrant_semantic.py b/tests/test_litellm_rust/cache/test_qdrant_semantic.py new file mode 100644 index 00000000000..160089c9002 --- /dev/null +++ b/tests/test_litellm_rust/cache/test_qdrant_semantic.py @@ -0,0 +1,286 @@ +import hashlib +import http.server +import json +import math +import os +import threading +import time +from collections.abc import Generator +from types import SimpleNamespace +from typing import Final +from uuid import uuid4 + +import pytest + +from litellm.caching.caching import Cache +from litellm.types.caching import LiteLLMCacheType +from tests.test_litellm_rust.support.cache import ( + CacheTestHandle, + CacheTestResolver, + assert_native_runtime, + request, + require_rust, +) + +pytestmark: Final = pytest.mark.requires_rust_extension + + +def qdrant_request( + key: str, + messages: list[dict[str, object]], + **kwargs: object, +) -> dict[str, object]: + return {**request(key), "messages": messages, **kwargs} + + +def embedding_vector(text: str) -> list[float]: + raw: Final = hashlib.sha256(text.encode()).digest()[:8] + values: Final = [byte / 127.5 - 1 for byte in raw] + norm: Final = math.sqrt(sum(value * value for value in values)) + return [value / norm for value in values] + + +@pytest.fixture +def qdrant_url() -> str: + value: Final[str | None] = os.environ.get("QDRANT_URL") + if not value: + pytest.skip("QDRANT_URL is required for Qdrant semantic cache tests") + return value.rstrip("/") + + +@pytest.fixture +def fake_embedding_endpoint(monkeypatch: pytest.MonkeyPatch) -> Generator[str]: + class EmbeddingHandler(http.server.BaseHTTPRequestHandler): + def do_POST(self) -> None: + length: Final = int(self.headers["Content-Length"]) + body: Final = json.loads(self.rfile.read(length)) + text: Final = body["input"] + response: Final = { + "object": "list", + "data": [ + { + "object": "embedding", + "index": 0, + "embedding": embedding_vector(text), + } + ], + "model": body["model"], + "usage": {"prompt_tokens": 1, "total_tokens": 1}, + } + encoded: Final = json.dumps(response).encode() + self.send_response(200) + self.send_header("Content-Type", "application/json") + self.send_header("Content-Length", str(len(encoded))) + self.end_headers() + self.wfile.write(encoded) + + def log_message(self, *_args: object) -> None: + return + + server: Final = http.server.ThreadingHTTPServer(("127.0.0.1", 0), EmbeddingHandler) + worker: Final = threading.Thread(target=server.serve_forever, daemon=True) + worker.start() + monkeypatch.setenv("OPENAI_API_BASE", f"http://127.0.0.1:{server.server_address[1]}") + monkeypatch.setenv("OPENAI_API_KEY", "sk-test") + try: + yield f"http://127.0.0.1:{server.server_address[1]}" + finally: + server.shutdown() + server.server_close() + worker.join(timeout=5) + + +def qdrant_facade(qdrant_url: str, collection_name: str) -> Cache: + return Cache( + type=LiteLLMCacheType.QDRANT_SEMANTIC, + qdrant_api_base=qdrant_url, + qdrant_collection_name=collection_name, + similarity_threshold=0.99, + qdrant_semantic_cache_embedding_model="text-embedding-3-small", + qdrant_semantic_cache_vector_size=8, + ) + + +def test_qdrant_semantic_facade_binds_native_and_shares_entries(qdrant_url: str, fake_embedding_endpoint: str) -> None: + del fake_embedding_endpoint + messages: Final = [{"role": "user", "content": "shared prompt"}] + collection: Final = f"cache_{uuid4().hex}" + facade: Final = qdrant_facade(qdrant_url, collection) + facade.cache.set_cache( + "python-key", + {"timestamp": time.time(), "response": json.dumps({"id": "py"})}, + messages=messages, + ) + handle: Final = CacheTestHandle.qdrant_semantic( + qdrant_url, + collection_name=collection, + similarity_threshold=0.99, + vector_size=8, + ) + handle._bind_facade(facade) + binding: Final = CacheTestResolver(SimpleNamespace(cache=facade)).resolve() + assert binding.kind == "native" + assert binding.lookup(qdrant_request("python-key", messages)) == {"id": "py"} + binding.store(qdrant_request("native-key", messages), {"id": "native"}) + python_value: Final = facade.cache.get_cache("native-key", messages=messages) + assert isinstance(python_value, dict) + assert python_value["response"] == {"id": "native"} + unrelated: Final = [{"role": "user", "content": "unrelated prompt"}] + assert binding.lookup(qdrant_request("native-key", unrelated)) is None + assert facade.cache.get_cache("native-key", messages=unrelated) is None + assert binding.lookup(qdrant_request("different-key", messages)) is None + assert facade.cache.get_cache("different-key", messages=messages) is None + + +async def test_qdrant_semantic_async_parity(qdrant_url: str, fake_embedding_endpoint: str) -> None: + del fake_embedding_endpoint + messages: Final = [{"role": "user", "content": "async prompt"}] + collection: Final = f"cache_{uuid4().hex}" + facade: Final = qdrant_facade(qdrant_url, collection) + handle: Final = CacheTestHandle.qdrant_semantic( + qdrant_url, + collection_name=collection, + similarity_threshold=0.99, + vector_size=8, + ) + handle._bind_facade(facade) + binding: Final = CacheTestResolver(SimpleNamespace(cache=facade)).resolve() + await facade.cache.async_set_cache( + "python-key", + {"timestamp": time.time(), "response": json.dumps({"id": "py"})}, + messages=messages, + ) + assert await binding.async_lookup(qdrant_request("python-key", messages)) == {"id": "py"} + await binding.async_store(qdrant_request("native-key", messages), {"id": "native"}) + python_value: Final = await facade.cache.async_get_cache("native-key", messages=messages) + assert isinstance(python_value, dict) + assert python_value["response"] == {"id": "native"} + + +async def test_qdrant_semantic_async_store_batch_shares_entries(qdrant_url: str, fake_embedding_endpoint: str) -> None: + del fake_embedding_endpoint + collection: Final = f"cache_{uuid4().hex}" + facade: Final = qdrant_facade(qdrant_url, collection) + handle: Final = CacheTestHandle.qdrant_semantic( + qdrant_url, + collection_name=collection, + similarity_threshold=0.99, + vector_size=8, + ) + handle._bind_facade(facade) + binding: Final = CacheTestResolver(SimpleNamespace(cache=facade)).resolve() + entries: Final = [ + qdrant_request("batch-one", [{"role": "user", "content": "first batch prompt"}]), + qdrant_request("batch-two", [{"role": "user", "content": "second batch prompt"}]), + ] + await binding.async_store_batch(entries, [{"id": "one"}, {"id": "two"}]) + + assert binding.lookup(entries[0]) == {"id": "one"} + assert binding.lookup(entries[1]) == {"id": "two"} + assert (await facade.cache.async_get_cache("batch-one", messages=entries[0]["messages"]))["response"] == { + "id": "one" + } + assert (await facade.cache.async_get_cache("batch-two", messages=entries[1]["messages"]))["response"] == { + "id": "two" + } + + +async def test_qdrant_semantic_malformed_entries_and_unsupported_operations( + qdrant_url: str, fake_embedding_endpoint: str +) -> None: + del fake_embedding_endpoint + messages: Final = [{"role": "user", "content": "malformed prompt"}] + collection: Final = f"cache_{uuid4().hex}" + facade: Final = qdrant_facade(qdrant_url, collection) + handle: Final = CacheTestHandle.qdrant_semantic( + qdrant_url, + collection_name=collection, + similarity_threshold=0.99, + vector_size=8, + ) + handle._bind_facade(facade) + binding: Final = CacheTestResolver(SimpleNamespace(cache=facade)).resolve() + key: Final = "malformed-key" + response: Final = { + "points": [ + { + "id": str(uuid4()), + "vector": embedding_vector("malformed prompt"), + "payload": { + "litellm_cache_key": key, + "text": "malformed prompt", + "response": "not json", + }, + } + ] + } + facade.cache.sync_client.put( + url=f"{qdrant_url}/collections/{collection}/points", + headers=facade.cache.headers, + json=response, + ) + assert binding.lookup(qdrant_request(key, messages)) is None + with pytest.raises(RuntimeError, match="operation is not supported"): + binding.lookup_batch([qdrant_request(key, messages)]) + with pytest.raises(RuntimeError, match="operation is not supported"): + await binding.async_flush() + with pytest.raises(RuntimeError, match="operation is not supported"): + await binding.ping() + + +def test_qdrant_semantic_ignores_request_expiry(qdrant_url: str, fake_embedding_endpoint: str) -> None: + del fake_embedding_endpoint + messages: Final = [{"role": "user", "content": "persistent prompt"}] + collection: Final = f"cache_{uuid4().hex}" + facade: Final = qdrant_facade(qdrant_url, collection) + handle: Final = CacheTestHandle.qdrant_semantic( + qdrant_url, + collection_name=collection, + similarity_threshold=0.99, + vector_size=8, + ) + handle._bind_facade(facade) + binding: Final = CacheTestResolver(SimpleNamespace(cache=facade)).resolve() + binding.store(qdrant_request("persistent-key", messages, ttl_seconds=1.0), {"id": "persistent"}) + time.sleep(1.2) + assert binding.lookup(qdrant_request("persistent-key", messages)) == {"id": "persistent"} + python_value: Final = facade.cache.get_cache("persistent-key", messages=messages) + assert isinstance(python_value, dict) + assert python_value["response"] == {"id": "persistent"} + + +def test_qdrant_semantic_mutation_and_projection_fallback(qdrant_url: str, fake_embedding_endpoint: str) -> None: + del fake_embedding_endpoint + collection: Final = f"cache_{uuid4().hex}" + facade: Final = qdrant_facade(qdrant_url, collection) + handle: Final = CacheTestHandle.qdrant_semantic( + qdrant_url, + collection_name=collection, + similarity_threshold=0.99, + vector_size=8, + ) + handle._bind_facade(facade) + facade.cache.qdrant_api_key = "rotated" + assert CacheTestResolver(SimpleNamespace(cache=facade)).resolve().kind == "python_callback" + facade.cache.similarity_threshold = 0.5 + assert CacheTestResolver(SimpleNamespace(cache=facade)).resolve().kind == "python_callback" + unsupported: Final = qdrant_facade(qdrant_url, f"cache_{uuid4().hex}") + unsupported.cache.embedding_max_input_tokens = 100 + with pytest.raises(TypeError, match="requires Python"): + handle._bind_facade(unsupported) + unsupported.cache.embedding_max_input_tokens = None + unsupported.cache.qdrant_api_base = "http://127.0.0.1:7777" + with pytest.raises(TypeError, match="gRPC"): + handle._bind_facade(unsupported) + + +def test_qdrant_semantic_rust_required_rule_activates_natively( + qdrant_url: str, fake_embedding_endpoint: str, monkeypatch: pytest.MonkeyPatch +) -> None: + del fake_embedding_endpoint + require_rust(monkeypatch, LiteLLMCacheType.QDRANT_SEMANTIC) + facade: Final = qdrant_facade(qdrant_url, f"cache_{uuid4().hex}") + assert_native_runtime(facade) + kwargs: Final = {"model": "gpt-4o", "messages": [{"role": "user", "content": "qdrant activation"}]} + facade.add_cache({"answer": "qdrant"}, **kwargs) + assert facade.get_cache(**kwargs) == {"answer": "qdrant"} diff --git a/tests/test_litellm_rust/cache/test_redis.py b/tests/test_litellm_rust/cache/test_redis.py new file mode 100644 index 00000000000..dd88145ef21 --- /dev/null +++ b/tests/test_litellm_rust/cache/test_redis.py @@ -0,0 +1,228 @@ +import json +import os +import time +from types import SimpleNamespace +from typing import Final +from urllib.parse import urlparse + +import pytest +import redis + +import litellm +from litellm.caching.caching import Cache +from litellm.caching.redis_cluster_cache import RedisClusterCache +from litellm.rust_bridge import catalog +from litellm.rust_bridge.catalog import CacheRule +from litellm.rust_bridge.configuration import Rollout +from litellm.types.caching import LiteLLMCacheType +from tests.test_litellm_rust.support.cache import ( + CacheTestHandle, + CacheTestResolver, + assert_native_runtime, + completion_kwargs, + request, + require_rust, +) +from tests.test_litellm_rust.support.isolation import rebound + +pytestmark: Final = pytest.mark.requires_rust_extension + + +@pytest.fixture +def cluster_nodes() -> tuple[tuple[str, int], ...]: + configured: Final = os.environ.get("LITELLM_TEST_REDIS_CLUSTER_NODES") + if not configured: + pytest.skip("LITELLM_TEST_REDIS_CLUSTER_NODES is not set") + return tuple((host, int(port)) for host, _, port in (node.partition(":") for node in configured.split(","))) + + +async def test_redis_reads_python_sync_and_async_entries_and_writes_without_hidden_prefix(redis_url: str) -> None: + client: Final = redis.Redis.from_url(redis_url) + namespace: Final = SimpleNamespace(cache=CacheTestHandle.redis(redis_url, namespace="team")) + binding: Final = CacheTestResolver(namespace).resolve() + response: Final = {"choices": [{"text": "cached"}], "usage": {"total_tokens": 3}, "flag": True, "empty": None} + envelope: Final = {"timestamp": time.time(), "response": json.dumps(response)} + client.set("team:sync", str(envelope)) + client.set("team:async", json.dumps({"timestamp": time.time(), "response": response})) + client.set("team:raw", json.dumps(response)) + client.set("team:invalid", "not a cache entry") + assert binding.lookup(request("sync")) == response + assert await binding.async_lookup(request("team:async")) == response + assert binding.lookup(request("raw")) == response + assert await binding.async_lookup(request("invalid")) is None + await binding.async_store({**request("native"), "ttl_seconds": 12.0}, response) + stored: Final = client.get("team:native") + assert isinstance(stored, bytes) + assert json.loads(stored)["response"] == response + assert 0 < client.ttl("team:native") <= 12 + assert client.get("litellm-cache:team:native") is None + assert client.get("team:team:async") is None + client.close() + + +async def test_redis_facade_buffers_native_async_writes(redis_url: str) -> None: + parsed: Final = urlparse(redis_url) + with rebound(litellm, "default_redis_ttl", 60): + facade: Final = Cache( + type=LiteLLMCacheType.REDIS, + host=parsed.hostname, + port=str(parsed.port), + redis_flush_size=2, + ) + with pytest.raises(TypeError, match="default TTLs must match"): + CacheTestHandle.redis(redis_url, ttl_seconds=61)._bind_facade(facade) + with pytest.raises(TypeError, match="namespaces must match"): + CacheTestHandle.redis(redis_url, namespace="other")._bind_facade(facade) + CacheTestHandle.redis(redis_url, ttl_seconds=60)._bind_facade(facade) + binding: Final = CacheTestResolver(SimpleNamespace(cache=facade)).resolve() + client: Final = redis.Redis.from_url(redis_url) + + with rebound(facade.cache, "redis_kwargs", {**facade.cache.redis_kwargs, "ssl": True}): + assert CacheTestResolver(SimpleNamespace(cache=facade)).resolve().kind == "python_callback" + + pool: Final = facade.cache.redis_client.connection_pool + with rebound(pool, "connection_kwargs", {**pool.connection_kwargs, "db": 1}): + assert CacheTestResolver(SimpleNamespace(cache=facade)).resolve().kind == "python_callback" + + await binding.async_store(request("first"), {"value": 1}) + assert client.get("first") is None + await binding.async_store(request("second"), {"value": 2}) + + assert client.get("first") is not None + assert client.get("second") is not None + await facade.cache.disconnect() + client.close() + + +async def test_redis_cluster_facade_serves_multi_slot_batches_and_scoped_flush_natively( + cluster_nodes: tuple[tuple[str, int], ...], +) -> None: + startup_nodes: Final = [{"host": host, "port": port} for host, port in cluster_nodes] + url: Final = f"redis://{cluster_nodes[0][0]}:{cluster_nodes[0][1]}" + with rebound(litellm, "default_redis_ttl", 60): + facade: Final = Cache(type=LiteLLMCacheType.REDIS, redis_startup_nodes=startup_nodes, namespace="parity") + assert type(facade.cache) is RedisClusterCache + with pytest.raises(TypeError, match="types must match"): + CacheTestHandle.redis(url, namespace="parity")._bind_facade(facade) + CacheTestHandle.redis(url, namespace="parity", startup_nodes=list(cluster_nodes))._bind_facade(facade) + resolver: Final = CacheTestResolver(SimpleNamespace(cache=facade)) + assert resolver.resolve().kind == "native" + + manager: Final = facade.cache.redis_client.nodes_manager + with rebound(manager, "connection_kwargs", {**manager.connection_kwargs, "db": 1}): + assert resolver.resolve().kind == "python_callback" + with rebound(facade.cache, "redis_kwargs", {**facade.cache.redis_kwargs, "startup_nodes": startup_nodes[:1]}): + assert resolver.resolve().kind == "python_callback" + binding: Final = resolver.resolve() + assert binding.kind == "native" + + client: Final = redis.RedisCluster(startup_nodes=[redis.cluster.ClusterNode(*node) for node in cluster_nodes]) + keys: Final = tuple(f"slot-{index}" for index in range(12)) + slots: Final = {client.keyslot(f"parity:{key}") for key in keys} + assert len(slots) > 1, slots + requests: Final = [request(key) for key in keys] + values: Final = [{"index": index} for index in range(len(keys))] + await binding.async_store_batch(requests, values) + client.set("parity:slot-3", "not a cache entry") + client.set("parity:slot-7", json.dumps({"timestamp": time.time(), "response": {"index": 7, "python": True}})) + + batch: Final = await binding.async_lookup_batch(requests) + assert batch == { + "values": [ + None if index == 3 else {"index": 7, "python": True} if index == 7 else value + for index, value in enumerate(values) + ], + "missing_indices": [3], + } + assert facade.cache.get_cache("parity:slot-0")["response"] == {"index": 0} + assert (await facade.cache.async_get_cache("parity:slot-11"))["response"] == {"index": 11} + assert facade.cache.redis_client.mget_nonatomic([f"parity:{key}" for key in keys[:2]]) == [ + client.get("parity:slot-0"), + client.get("parity:slot-1"), + ] + + await binding.async_store({**request("pinned"), "ttl_seconds": 12.0}, {"pinned": True}) + assert 0 < client.ttl("parity:pinned") <= 12 + client.set("unscoped", "stays") + + await binding.async_flush() + + remaining: Final = tuple( + sorted(key for node in client.get_primaries() for key in client.keys("parity:*", target_nodes=node)) + ) + assert remaining == (), remaining + assert client.get("unscoped") == b"stays" + client.delete("unscoped") + client.close() + facade.cache.redis_client.close() + + +def redis_facade(redis_url: str, **settings: object) -> Cache: + parsed: Final = urlparse(redis_url) + return Cache(type=LiteLLMCacheType.REDIS, host=parsed.hostname, port=str(parsed.port), **settings) + + +@pytest.mark.parametrize( + ("settings", "message"), + [ + pytest.param({"max_connections": 10}, "max_connections requires Python", id="pool-size"), + pytest.param({"socket_timeout": 1.0}, "socket_timeout and socket_connect_timeout", id="socket-timeout"), + pytest.param( + {"socket_connect_timeout": 1.0}, "socket_timeout and socket_connect_timeout", id="connect-timeout" + ), + pytest.param({"socket_keepalive": True}, "does not support socket_keepalive", id="keepalive"), + pytest.param({"health_check_interval": 5}, "does not support health_check_interval", id="health-check"), + pytest.param({"client_name": "litellm"}, "does not support client_name", id="client-name"), + pytest.param({"ssl": True}, "ssl_check_hostname=false require Python", id="tls-default-hostname-check"), + pytest.param({"ssl": True, "ssl_cert_reqs": "none"}, "ssl_cert_reqs=none", id="tls-without-verification"), + pytest.param( + {"ssl": True, "ssl_check_hostname": True, "ssl_ca_certs": "/ca.pem"}, + "does not support ssl_ca_certs", + id="tls-custom-ca", + ), + pytest.param( + {"ssl": True, "ssl_check_hostname": True, "ssl_certfile": "/client.pem", "ssl_keyfile": "/client.key"}, + "does not support ssl_ca_certs, ssl_ca_data, ssl_certfile or ssl_keyfile", + id="tls-client-certificate", + ), + ], +) +def test_redis_settings_the_native_client_cannot_honor_decline( + redis_url: str, monkeypatch: pytest.MonkeyPatch, settings: dict[str, object], message: str +) -> None: + require_rust(monkeypatch, LiteLLMCacheType.REDIS) + with pytest.raises(RuntimeError, match=f"declined the cache: native Redis.*{message}"): + redis_facade(redis_url, **settings) + + +def test_redis_verified_tls_activates_natively(redis_url: str, monkeypatch: pytest.MonkeyPatch) -> None: + require_rust(monkeypatch, LiteLLMCacheType.REDIS) + assert_native_runtime(redis_facade(redis_url, ssl=True, ssl_check_hostname=True)) + + +async def test_redis_flush_size_buffers_native_facade_writes(redis_url: str, monkeypatch: pytest.MonkeyPatch) -> None: + require_rust(monkeypatch, LiteLLMCacheType.REDIS) + facade: Final = redis_facade(redis_url, redis_flush_size=2, namespace="team") + assert_native_runtime(facade) + client: Final = redis.Redis.from_url(redis_url) + first: Final = completion_kwargs("first") + await facade.async_add_cache({"value": 1}, **first) + first_key: Final = facade.get_cache_key(**first) + assert first_key.startswith("team:") + assert client.get(first_key) is None + second: Final = completion_kwargs("second") + await facade.async_add_cache({"value": 2}, **second) + assert client.get(first_key) is not None + assert client.get(facade.get_cache_key(**second)) is not None + client.close() + + +def test_rust_with_fallback_keeps_python_when_the_native_client_declines( + redis_url: str, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr( + catalog, + "RULES", + (CacheRule(Rollout.RUST_OPT_OUT, backends=frozenset({LiteLLMCacheType.REDIS})),), + ) + assert redis_facade(redis_url, socket_timeout=1.0)._native_cache is None # pyright: ignore[reportPrivateUsage] # the activation under test has no public accessor diff --git a/tests/test_litellm_rust/cache/test_redis_semantic.py b/tests/test_litellm_rust/cache/test_redis_semantic.py new file mode 100644 index 00000000000..279330d9060 --- /dev/null +++ b/tests/test_litellm_rust/cache/test_redis_semantic.py @@ -0,0 +1,606 @@ +import asyncio +import contextvars +import hashlib +import json +import math +import os +from collections.abc import Callable, Generator +from contextlib import ExitStack +from types import SimpleNamespace +from typing import Final, cast +from uuid import uuid4 + +import pytest +import redis + +import litellm +from litellm.caching.caching import Cache +from litellm.caching.redis_semantic_cache import RedisSemanticCache +from litellm.types.caching import LiteLLMCacheType +from litellm.types.llms.custom_llm import CustomLLMItem +from litellm.types.utils import EmbeddingResponse +from tests.test_litellm_rust.support.cache import ( + CacheTestHandle, + CacheTestResolver, + assert_native_runtime, + request, + require_rust, +) +from tests.test_litellm_rust.support.isolation import rebound + +pytestmark: Final = pytest.mark.requires_rust_extension + + +PARAPHRASE_MARKER: Final = " (paraphrase)" + + +SEMANTIC_EMBEDDING_MODEL: Final = "semantic-test/deterministic" + + +SEMANTIC_INDEX_PREFIX: Final = "litellm_test_semantic_" + + +SEMANTIC_CONTEXT: Final = contextvars.ContextVar("semantic_test_context", default="unset") + + +def _normalized(vector: list[float]) -> list[float]: + norm: Final = math.sqrt(sum(component * component for component in vector)) + return [component / norm for component in vector] + + +def _base_embedding(prompt: str) -> list[float]: + digest: Final = hashlib.sha256(prompt.encode("utf-8")).digest() + return _normalized([float(digest[index] + 1) for index in range(8)]) + + +def _semantic_embedding(prompt: str) -> list[float]: + if PARAPHRASE_MARKER not in prompt: + return _base_embedding(prompt) + base: Final = _base_embedding(prompt.replace(PARAPHRASE_MARKER, "").strip()) + pivot: Final = min(range(8), key=lambda index: abs(base[index])) + direction: Final = _normalized( + [(1.0 - base[pivot] * base[pivot]) if index == pivot else -base[index] * base[pivot] for index in range(8)] + ) + # Rotating an orthogonal unit direction by 0.329 produces ~0.05 cosine distance + return _normalized([base[index] + 0.329 * direction[index] for index in range(8)]) + + +class DeterministicEmbedding(litellm.CustomLLM): + def __init__(self) -> None: + self.calls: list[dict[str, object]] = [] + self.async_calls: list[dict[str, object]] = [] + self.entered = asyncio.Event() + self.gate: asyncio.Event | None = None + + def _respond( + self, + model: str, + input: object, + model_response: EmbeddingResponse, + ) -> EmbeddingResponse: + texts: Final = cast(list[object], input if isinstance(input, list) else [input]) + self.calls.append({"model": model, "input": texts}) + model_response.model = model + model_response.data = [ + {"object": "embedding", "index": index, "embedding": _semantic_embedding(str(text))} + for index, text in enumerate(texts) + ] + return model_response + + def embedding( + self, + model: str, + input: list[object], + model_response: EmbeddingResponse, + print_verbose: Callable[..., object], + logging_obj: object, + optional_params: dict[str, object], + api_key: object = None, + api_base: object = None, + timeout: object = None, + litellm_params: object = None, + ) -> EmbeddingResponse: + return self._respond(model, input, model_response) + + async def aembedding( + self, + model: str, + input: list[object], + model_response: EmbeddingResponse, + print_verbose: Callable[..., object], + logging_obj: object, + optional_params: dict[str, object], + api_key: object = None, + api_base: object = None, + timeout: object = None, + litellm_params: object = None, + ) -> EmbeddingResponse: + texts: Final = cast(list[object], input if isinstance(input, list) else [input]) + self.async_calls.append( + { + "model": model, + "input": texts, + "task": asyncio.current_task(), + "context": SEMANTIC_CONTEXT.get(), + } + ) + SEMANTIC_CONTEXT.set("written-in-aembedding") + self.entered.set() + if self.gate is not None: + await self.gate.wait() + return self._respond(model, input, model_response) + + +@pytest.fixture +def semantic_embedding() -> Generator[DeterministicEmbedding]: + handler: Final = DeterministicEmbedding() + with ExitStack() as stack: + stack.enter_context( + rebound( + litellm, + "custom_provider_map", + [ + *litellm.custom_provider_map, + cast( + CustomLLMItem, + {"provider": "semantic-test", "custom_handler": handler}, + ), + ], + ) + ) + stack.enter_context( + rebound( + litellm, + "_custom_providers", # pyright: ignore[reportPrivateUsage] # no public provider-registration hook + [*litellm._custom_providers, "semantic-test"], # pyright: ignore[reportPrivateUsage] # no public provider-registration hook + ) + ) + stack.enter_context(rebound(litellm, "provider_list", [*litellm.provider_list, "semantic-test"])) + yield handler + + +@pytest.fixture +def redis_stack() -> Generator[tuple[str, str]]: + url: Final = os.environ.get("LITELLM_REDIS_STACK_URL") + if url is None: + pytest.skip("LITELLM_REDIS_STACK_URL is not set") + index: Final = f"{SEMANTIC_INDEX_PREFIX}{uuid4().hex}" + yield url, index + client: Final = redis.Redis.from_url(url) + try: + client.execute_command("FT.DROPINDEX", index, "DD") # pyright: ignore[reportUnknownMemberType] # redis-py leaves execute_command partially unknown + except redis.RedisError: + pass + client.close() + + +def semantic_request(key: str, prompt: str, **extra: object) -> dict[str, object]: + return { + "key": {"preset": key}, + "messages": [{"role": "user", "content": prompt}], + **extra, + } + + +def semantic_messages(prompt: str) -> list[dict[str, object]]: + return [{"role": "user", "content": prompt}] + + +def semantic_entry_id(prompt: str, tag: str) -> str: + return hashlib.sha256(f"{prompt}litellm_cache_key{tag}".encode()).hexdigest() + + +def semantic_facade(url: str, index: str, *, similarity_threshold: float = 0.8) -> Cache: + facade: Final = Cache( + type=LiteLLMCacheType.REDIS_SEMANTIC, + redis_url=url, + similarity_threshold=similarity_threshold, + redis_semantic_cache_embedding_model=SEMANTIC_EMBEDDING_MODEL, + redis_semantic_cache_index_name=index, + ) + CacheTestHandle.redis_semantic(facade.cache)._bind_facade(facade) + return facade + + +def test_redis_semantic_constructor_identity_and_provenance( + redis_stack: tuple[str, str], semantic_embedding: DeterministicEmbedding +) -> None: + url, index = redis_stack + facade: Final = semantic_facade(url, index) + backend: Final = cast(RedisSemanticCache, facade.cache) + assert backend.__class__.__module__ == "litellm.caching.redis_semantic_cache" + assert type(backend) is RedisSemanticCache + assert backend._redis_url == url # pyright: ignore[reportPrivateUsage] # provenance check needs the projected config + assert backend._index_name == index # pyright: ignore[reportPrivateUsage] # provenance check needs the projected config + assert backend.similarity_threshold == 0.8 + assert backend.embedding_model == SEMANTIC_EMBEDDING_MODEL + handle: Final = cast(object, getattr(facade, "_native_cache_handle")) + assert isinstance(handle, CacheTestHandle) + assert handle.backend == "redis_semantic" + binding: Final = CacheTestResolver(SimpleNamespace(cache=facade)).resolve() + assert binding.kind == "native" + + +def test_redis_semantic_native_and_python_sync_entries_share_one_layout( + redis_stack: tuple[str, str], semantic_embedding: DeterministicEmbedding +) -> None: + url, index = redis_stack + facade: Final = semantic_facade(url, index) + binding: Final = CacheTestResolver(SimpleNamespace(cache=facade)).resolve() + client: Final = redis.Redis.from_url(url) + response: Final = {"choices": [{"text": "paris"}], "usage": {"total_tokens": 2}} + + binding.store(semantic_request("geo", "what is the capital of france"), response) + + native_hash_key: Final = f"{index}:{semantic_entry_id('what is the capital of france', 'geo')}" + stored: Final = client.hgetall(native_hash_key) + assert set(stored) == { + b"entry_id", + b"prompt", + b"response", + b"prompt_vector", + b"inserted_at", + b"updated_at", + b"litellm_cache_key", + }, stored + assert stored[b"entry_id"].decode() == native_hash_key.split(":", 1)[1] + assert stored[b"prompt"] == b"what is the capital of france" + assert stored[b"litellm_cache_key"] == b"geo" + assert len(stored[b"prompt_vector"]) == 32 + decoded: Final = cast(dict[str, object], json.loads(stored[b"response"])) + assert decoded["response"] == response + assert ( + cast(RedisSemanticCache, facade.cache).get_cache( # pyright: ignore[reportUnknownMemberType] # **kwargs stays unknown on the backend class + "geo", messages=semantic_messages("what is the capital of france") + ) + == decoded + ) + assert semantic_embedding.calls == [ + {"model": "deterministic", "input": ["what is the capital of france"]}, + {"model": "deterministic", "input": ["what is the capital of france"]}, + {"model": "deterministic", "input": ["dimension test"]}, + ] + + cast(RedisSemanticCache, facade.cache).set_cache( # pyright: ignore[reportUnknownMemberType] # **kwargs stays unknown on the backend class + "math", + json.dumps({"timestamp": 1700000000.0, "response": {"answer": 42}}), + messages=semantic_messages("what is 6 times 7"), + ) + python_hash_key: Final = f"{index}:{semantic_entry_id('what is 6 times 7', 'math')}" + assert json.loads(cast(bytes, client.hget(python_hash_key, "response"))) == { + "timestamp": 1700000000.0, + "response": {"answer": 42}, + } + assert binding.lookup(semantic_request("math", "what is 6 times 7")) == {"answer": 42} + client.close() + + +async def test_redis_semantic_async_paths_and_store_batch_share_one_layout( + redis_stack: tuple[str, str], semantic_embedding: DeterministicEmbedding +) -> None: + url, index = redis_stack + facade: Final = semantic_facade(url, index) + binding: Final = CacheTestResolver(SimpleNamespace(cache=facade)).resolve() + client: Final = redis.Redis.from_url(url) + + await binding.async_store(semantic_request("async", "name a primary color"), {"answer": "blue"}) + hash_key: Final = f"{index}:{semantic_entry_id('name a primary color', 'async')}" + decoded: Final = cast(dict[str, object], json.loads(cast(bytes, client.hget(hash_key, "response")))) + python_read: Final = await cast(RedisSemanticCache, facade.cache).async_get_cache( # pyright: ignore[reportUnknownMemberType] # **kwargs stays unknown on the backend class + "async", messages=semantic_messages("name a primary color") + ) + assert python_read == decoded + + await binding.async_store_batch( + [ + semantic_request("batch-one", "first batch prompt"), + semantic_request("batch-two", "second batch prompt"), + ], + [{"answer": 1}, {"answer": 2}], + ) + expected: Final = { + key: json.loads(cast(bytes, client.hget(f"{index}:{semantic_entry_id(prompt, key)}", "response"))) + for key, prompt in ( + ("batch-one", "first batch prompt"), + ("batch-two", "second batch prompt"), + ) + } + for key, prompt in ( + ("batch-one", "first batch prompt"), + ("batch-two", "second batch prompt"), + ): + assert ( + cast(RedisSemanticCache, facade.cache).get_cache( # pyright: ignore[reportUnknownMemberType] # **kwargs stays unknown on the backend class + key, messages=semantic_messages(prompt) + ) + == expected[key] + ), key + + cast(RedisSemanticCache, facade.cache).set_cache( # pyright: ignore[reportUnknownMemberType] # **kwargs stays unknown on the backend class + "async-python", + json.dumps({"timestamp": 1700000000.0, "response": {"answer": "python"}}), + messages=semantic_messages("python written prompt"), + ) + assert await binding.async_lookup(semantic_request("async-python", "python written prompt")) == {"answer": "python"} + client.close() + + +async def test_native_semantic_async_embedding_runs_inline_in_the_callers_task( + redis_stack: tuple[str, str], semantic_embedding: DeterministicEmbedding +) -> None: + url, index = redis_stack + facade: Final = semantic_facade(url, index) + binding: Final = CacheTestResolver(SimpleNamespace(cache=facade)).resolve() + assert binding.kind == "native" + caller: Final = asyncio.current_task() + SEMANTIC_CONTEXT.set("caller-sentinel") + response: Final = {"choices": [{"text": "paris"}]} + + await binding.async_store(semantic_request("inline", "what is the capital of france"), response) + assert ( + await binding.async_lookup(semantic_request("inline", f"what is the capital of france{PARAPHRASE_MARKER}")) + == response + ) + assert await binding.async_lookup(semantic_request("inline", "python written prompt")) is None + assert SEMANTIC_CONTEXT.get() == "written-in-aembedding" + assert semantic_embedding.async_calls == [ + { + "model": "deterministic", + "input": ["what is the capital of france"], + "task": caller, + "context": "caller-sentinel", + }, + { + "model": "deterministic", + "input": [f"what is the capital of france{PARAPHRASE_MARKER}"], + "task": caller, + "context": "written-in-aembedding", + }, + { + "model": "deterministic", + "input": ["python written prompt"], + "task": caller, + "context": "written-in-aembedding", + }, + ], semantic_embedding.async_calls + + +async def test_native_semantic_cancellation_during_embedding_skips_the_backend( + redis_stack: tuple[str, str], semantic_embedding: DeterministicEmbedding +) -> None: + url, index = redis_stack + facade: Final = semantic_facade(url, index) + binding: Final = CacheTestResolver(SimpleNamespace(cache=facade)).resolve() + assert binding.kind == "native" + semantic_embedding.gate = asyncio.Event() + + async def lookup() -> object: + return await binding.async_lookup(semantic_request("cancel", "cancelled prompt")) + + task: Final = asyncio.create_task(lookup()) + await semantic_embedding.entered.wait() + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + semantic_embedding.gate.set() + + assert len(semantic_embedding.async_calls) == 1 + assert ( + await cast(RedisSemanticCache, facade.cache).async_get_cache( # pyright: ignore[reportUnknownMemberType] # **kwargs stays unknown on the backend class + "cancel", messages=semantic_messages("cancelled prompt") + ) + is None + ) + + +def test_redis_semantic_similarity_tag_and_threshold_boundaries( + redis_stack: tuple[str, str], semantic_embedding: DeterministicEmbedding +) -> None: + url, index = redis_stack + facade: Final = semantic_facade(url, index) + binding: Final = CacheTestResolver(SimpleNamespace(cache=facade)).resolve() + + binding.store(semantic_request("sim", "tell me a joke"), {"answer": "haha"}) + paraphrase: Final = f"tell me a joke{PARAPHRASE_MARKER}" + assert binding.lookup(semantic_request("sim", paraphrase)) == {"answer": "haha"} + assert binding.lookup(semantic_request("sim", "an unrelated question about spreadsheets")) is None + assert binding.lookup(semantic_request("other-key", "tell me a joke")) is None + + strict: Final = semantic_facade(url, index, similarity_threshold=0.99) + strict_binding: Final = CacheTestResolver(SimpleNamespace(cache=strict)).resolve() + assert strict_binding.lookup(semantic_request("sim", paraphrase)) is None + assert strict_binding.lookup(semantic_request("sim", "tell me a joke")) == {"answer": "haha"} + + +def test_redis_semantic_ttl_is_written_only_when_requested( + redis_stack: tuple[str, str], semantic_embedding: DeterministicEmbedding +) -> None: + url, index = redis_stack + facade: Final = semantic_facade(url, index) + binding: Final = CacheTestResolver(SimpleNamespace(cache=facade)).resolve() + client: Final = redis.Redis.from_url(url) + + binding.store({**semantic_request("ttl", "ttl prompt"), "ttl_seconds": 12.0}, {"answer": 1}) + expiring: Final = f"{index}:{semantic_entry_id('ttl prompt', 'ttl')}" + assert 0 < client.ttl(expiring) <= 12 + + binding.store(semantic_request("ttl-none", "untimed prompt"), {"answer": 2}) + persistent: Final = f"{index}:{semantic_entry_id('untimed prompt', 'ttl-none')}" + assert client.ttl(persistent) == -1 + + binding.store( + {**semantic_request("ttl-fraction", "fractional prompt"), "ttl_seconds": 1.5}, + {"answer": 3}, + ) + fractional: Final = f"{index}:{semantic_entry_id('fractional prompt', 'ttl-fraction')}" + assert client.ttl(fractional) == 2 + client.close() + + +def test_redis_semantic_malformed_response_is_a_miss_for_both_readers( + redis_stack: tuple[str, str], semantic_embedding: DeterministicEmbedding +) -> None: + url, index = redis_stack + facade: Final = semantic_facade(url, index) + binding: Final = CacheTestResolver(SimpleNamespace(cache=facade)).resolve() + client: Final = redis.Redis.from_url(url) + + binding.store(semantic_request("bad", "corrupt me"), {"answer": 1}) + hash_key: Final = f"{index}:{semantic_entry_id('corrupt me', 'bad')}" + client.hset(hash_key, "response", b"{not json") + assert binding.lookup(semantic_request("bad", "corrupt me")) is None + assert ( + cast(RedisSemanticCache, facade.cache).get_cache( # pyright: ignore[reportUnknownMemberType] # **kwargs stays unknown on the backend class + "bad", messages=semantic_messages("corrupt me") + ) + is None + ) + client.close() + + +async def test_redis_semantic_unsupported_operations_raise_not_implemented( + redis_stack: tuple[str, str], semantic_embedding: DeterministicEmbedding +) -> None: + url, index = redis_stack + facade: Final = semantic_facade(url, index) + binding: Final = CacheTestResolver(SimpleNamespace(cache=facade)).resolve() + + with pytest.raises(NotImplementedError): + binding.lookup_batch([semantic_request("batch", "prompt one")]) + with pytest.raises(NotImplementedError): + await binding.async_lookup_batch([semantic_request("batch", "prompt one")]) + with pytest.raises(NotImplementedError): + await binding.async_flush() + with pytest.raises(NotImplementedError): + await binding.ping() + + +def test_redis_semantic_requests_without_prompt_are_noops( + redis_stack: tuple[str, str], semantic_embedding: DeterministicEmbedding +) -> None: + url, index = redis_stack + facade: Final = semantic_facade(url, index) + binding: Final = CacheTestResolver(SimpleNamespace(cache=facade)).resolve() + client: Final = redis.Redis.from_url(url) + + binding.store(request("plain"), {"answer": 1}) + assert binding.lookup(request("plain")) is None + assert semantic_embedding.calls == [] + assert client.keys(f"{index}:*") == [] + client.close() + + +def test_redis_semantic_scope_overrides_the_tag_and_isolates_entries( + redis_stack: tuple[str, str], semantic_embedding: DeterministicEmbedding +) -> None: + url, index = redis_stack + facade: Final = semantic_facade(url, index) + binding: Final = CacheTestResolver(SimpleNamespace(cache=facade)).resolve() + client: Final = redis.Redis.from_url(url) + + scoped: Final = {**semantic_request("scoped", "scoped prompt"), "scope": "team-a"} + binding.store(scoped, {"answer": "kept"}) + hash_key: Final = f"{index}:{semantic_entry_id('scoped prompt', 'team-a')}" + assert client.hget(hash_key, "litellm_cache_key") == b"team-a" + assert binding.lookup(scoped) == {"answer": "kept"} + assert binding.lookup(semantic_request("scoped", "scoped prompt")) is None + assert binding.lookup({**scoped, "scope": "team-b"}) is None + client.close() + + +def test_redis_semantic_configuration_drift_falls_back_to_python( + redis_stack: tuple[str, str], + semantic_embedding: DeterministicEmbedding, + monkeypatch: pytest.MonkeyPatch, +) -> None: + url, index = redis_stack + facade: Final = semantic_facade(url, index) + resolver: Final = CacheTestResolver(SimpleNamespace(cache=facade)) + assert resolver.resolve().kind == "native" + + with rebound(facade.cache, "similarity_threshold", 0.5): + assert resolver.resolve().kind == "python_callback" + with rebound(facade, "semantic_cache_scope", "end_user"): + assert resolver.resolve().kind == "python_callback" + with rebound(facade.cache, "embedding_model", "other-model"): + assert resolver.resolve().kind == "python_callback" + with rebound(facade.cache, "_index_name", "other-index"): + assert resolver.resolve().kind == "python_callback" + with rebound(facade.cache, "CACHE_KEY_FIELD_NAME", "other-field"): + assert resolver.resolve().kind == "python_callback" + + def patched_embedding(self: object, prompt: str, metadata: object = None) -> list[float]: + return _semantic_embedding(prompt) + + monkeypatch.setattr(RedisSemanticCache, "_get_embedding", patched_embedding) + assert resolver.resolve().kind == "python_callback" + + +def test_redis_semantic_handle_rejects_wrong_backends( + redis_stack: tuple[str, str], semantic_embedding: DeterministicEmbedding +) -> None: + url, index = redis_stack + + class CustomSemanticCache(RedisSemanticCache): + pass + + with pytest.raises(TypeError, match="built-in RedisSemanticCache"): + CacheTestHandle.redis_semantic(object()) + with pytest.raises(TypeError, match="built-in RedisSemanticCache"): + CacheTestHandle.redis_semantic( + CustomSemanticCache( + redis_url=url, + similarity_threshold=0.8, + embedding_model=SEMANTIC_EMBEDDING_MODEL, + index_name=f"{index}_subclass", + ) + ) + + facade: Final = semantic_facade(url, index) + with pytest.raises(TypeError, match="backend types must match"): + CacheTestHandle.redis(url)._bind_facade(facade) + + subclassed_facade: Final = Cache( + type=LiteLLMCacheType.REDIS_SEMANTIC, + redis_url=url, + similarity_threshold=0.8, + redis_semantic_cache_embedding_model=SEMANTIC_EMBEDDING_MODEL, + redis_semantic_cache_index_name=index, + ) + subclassed_facade.cache = CustomSemanticCache( # pyright: ignore[reportAttributeAccessIssue] # facade backend slot is not declared + redis_url=url, + similarity_threshold=0.8, + embedding_model=SEMANTIC_EMBEDDING_MODEL, + index_name=index, + ) + with pytest.raises(TypeError): + CacheTestHandle.redis_semantic(subclassed_facade.cache)._bind_facade(subclassed_facade) + + replacement_facade: Final = Cache( + type=LiteLLMCacheType.REDIS_SEMANTIC, + redis_url=url, + similarity_threshold=0.8, + redis_semantic_cache_embedding_model=SEMANTIC_EMBEDDING_MODEL, + redis_semantic_cache_index_name=index, + ) + with pytest.raises(TypeError, match="must be the native embedder"): + CacheTestHandle.redis_semantic(facade.cache)._bind_facade(replacement_facade) + + +async def test_redis_semantic_rust_required_rule_activates_natively( + redis_stack: tuple[str, str], semantic_embedding: DeterministicEmbedding, monkeypatch: pytest.MonkeyPatch +) -> None: + del semantic_embedding + url, index = redis_stack + require_rust(monkeypatch, LiteLLMCacheType.REDIS_SEMANTIC) + facade: Final = Cache( + type=LiteLLMCacheType.REDIS_SEMANTIC, + redis_url=url, + similarity_threshold=0.8, + redis_semantic_cache_embedding_model=SEMANTIC_EMBEDDING_MODEL, + redis_semantic_cache_index_name=index, + ) + assert_native_runtime(facade) + kwargs: Final = {"model": "gpt-4o", "messages": semantic_messages("name a primary color")} + await facade.async_add_cache({"answer": "blue"}, **kwargs) + assert await facade.async_get_cache(**kwargs) == {"answer": "blue"} diff --git a/tests/test_litellm_rust/cache/test_rollout.py b/tests/test_litellm_rust/cache/test_rollout.py new file mode 100644 index 00000000000..7f33e31599f --- /dev/null +++ b/tests/test_litellm_rust/cache/test_rollout.py @@ -0,0 +1,264 @@ +import asyncio +from collections.abc import Callable +from pathlib import Path +from types import SimpleNamespace +from typing import Final, TypeAlias, cast +from urllib.parse import urlparse +from uuid import uuid4 + +import pytest + +from litellm.caching.caching import Cache +from litellm.rust_bridge.response_cache import NativeResponseCacheRuntime, ResponseCacheRuntime, resolve_response_cache +from litellm.types.caching import LiteLLMCacheType +from litellm.types.utils import EmbeddingResponse +from tests.test_litellm_rust.support.cache import assert_native_runtime, completion_kwargs, require_rust +from tests.test_litellm_rust.support.s3_stub import S3Stub + +pytestmark: Final = pytest.mark.requires_rust_extension + + +CacheFactory: TypeAlias = Callable[[], Cache] + + +@pytest.fixture +def cache_factory(request: pytest.FixtureRequest, tmp_path: Path) -> CacheFactory: + backend: Final = cast(LiteLLMCacheType, request.param) + match backend: + case LiteLLMCacheType.LOCAL: + return lambda: Cache(type=backend) + case LiteLLMCacheType.DISK: + return lambda: Cache(type=backend, disk_cache_dir=str(tmp_path)) + case LiteLLMCacheType.REDIS: + parsed: Final = urlparse(cast(str, request.getfixturevalue("redis_url"))) + return lambda: Cache(type=backend, host=parsed.hostname, port=str(parsed.port)) + case LiteLLMCacheType.S3: + stub: Final = cast(S3Stub, request.getfixturevalue("s3_stub")) + return lambda: Cache( + type=backend, + s3_bucket_name="cache-bucket", + s3_region_name="us-east-1", + s3_endpoint_url=stub.url, + s3_aws_access_key_id="key", + s3_aws_secret_access_key="secret", + s3_path="team", + ) + case LiteLLMCacheType.GCS: + return lambda: Cache(type=backend, gcs_bucket_name="bucket", gcs_path="cache/") + case LiteLLMCacheType.REDIS_SEMANTIC: + return lambda: Cache( + type=backend, + redis_url="redis://127.0.0.1:6379", + similarity_threshold=0.8, + redis_semantic_cache_embedding_model="text-embedding-3-small", + ) + case LiteLLMCacheType.VALKEY_SEMANTIC: + return lambda: Cache(type=backend, redis_url="redis://127.0.0.1:6390/0", similarity_threshold=0.8) + case _: + raise AssertionError(f"no local factory for {backend}") + + +ROUND_TRIP_BACKENDS: Final = ( + LiteLLMCacheType.LOCAL, + LiteLLMCacheType.DISK, + LiteLLMCacheType.REDIS, + LiteLLMCacheType.S3, +) + + +SHARED_STORE_BACKENDS: Final = (LiteLLMCacheType.DISK, LiteLLMCacheType.REDIS, LiteLLMCacheType.S3) + + +@pytest.mark.parametrize("backend", list(LiteLLMCacheType)) +def test_shipped_rules_keep_every_backend_on_python(backend: LiteLLMCacheType) -> None: + assert resolve_response_cache(cast(Cache, SimpleNamespace(type=backend))) is None + + +@pytest.mark.parametrize( + "cache_factory", + [ + LiteLLMCacheType.LOCAL, + LiteLLMCacheType.DISK, + LiteLLMCacheType.REDIS, + LiteLLMCacheType.S3, + LiteLLMCacheType.GCS, + LiteLLMCacheType.REDIS_SEMANTIC, + LiteLLMCacheType.VALKEY_SEMANTIC, + ], + indirect=True, +) +def test_shipped_rules_construct_python_backed_facades(cache_factory: CacheFactory) -> None: + assert cache_factory()._native_cache is None # pyright: ignore[reportPrivateUsage] # the activation under test has no public accessor + + +@pytest.mark.parametrize( + "cache_factory", + [ + LiteLLMCacheType.LOCAL, + LiteLLMCacheType.DISK, + LiteLLMCacheType.REDIS, + LiteLLMCacheType.S3, + LiteLLMCacheType.GCS, + LiteLLMCacheType.REDIS_SEMANTIC, + LiteLLMCacheType.VALKEY_SEMANTIC, + ], + indirect=True, +) +def test_rust_required_rule_activates_the_native_backend( + cache_factory: CacheFactory, monkeypatch: pytest.MonkeyPatch, request: pytest.FixtureRequest +) -> None: + require_rust(monkeypatch, cast(LiteLLMCacheType, request.node.callspec.params["cache_factory"])) + assert_native_runtime(cache_factory()) + + +@pytest.mark.parametrize("cache_factory", ROUND_TRIP_BACKENDS, indirect=True) +async def test_facade_storage_calls_round_trip_through_the_native_backend( + cache_factory: CacheFactory, monkeypatch: pytest.MonkeyPatch, request: pytest.FixtureRequest +) -> None: + require_rust(monkeypatch, cast(LiteLLMCacheType, request.node.callspec.params["cache_factory"])) + facade: Final = cache_factory() + assert_native_runtime(facade) + + sync_kwargs: Final = completion_kwargs("sync") + facade.add_cache({"answer": 1}, **sync_kwargs) + assert facade.get_cache(**sync_kwargs) == {"answer": 1} + + async_kwargs: Final = completion_kwargs("async") + await facade.async_add_cache({"answer": 2}, **async_kwargs) + assert await facade.async_get_cache(**async_kwargs) == {"answer": 2} + assert facade.get_cache(**completion_kwargs("absent")) is None + + +async def test_memory_facade_writes_bypass_the_python_backend(monkeypatch: pytest.MonkeyPatch) -> None: + require_rust(monkeypatch, LiteLLMCacheType.LOCAL) + facade: Final = Cache(type=LiteLLMCacheType.LOCAL) + assert_native_runtime(facade) + kwargs: Final = completion_kwargs("memory") + facade.add_cache({"answer": 1}, **kwargs) + assert facade.cache.get_cache(facade.get_cache_key(**kwargs)) is None + assert facade.get_cache(**kwargs) == {"answer": 1} + + +@pytest.mark.parametrize("cache_factory", SHARED_STORE_BACKENDS, indirect=True) +async def test_native_and_python_facades_share_one_wire_format( + cache_factory: CacheFactory, monkeypatch: pytest.MonkeyPatch, request: pytest.FixtureRequest +) -> None: + python_facade: Final = cache_factory() + assert python_facade._native_cache is None # pyright: ignore[reportPrivateUsage] # the activation under test has no public accessor + require_rust(monkeypatch, cast(LiteLLMCacheType, request.node.callspec.params["cache_factory"])) + native_facade: Final = cache_factory() + assert_native_runtime(native_facade) + + native_written: Final = completion_kwargs("native") + native_facade.add_cache({"writer": "native"}, **native_written) + assert python_facade.get_cache(**native_written) == {"writer": "native"} + + python_written: Final = completion_kwargs("python") + python_facade.add_cache({"writer": "python"}, **python_written) + assert native_facade.get_cache(**python_written) == {"writer": "python"} + + async_native: Final = completion_kwargs("async-native") + await native_facade.async_add_cache({"writer": "async-native"}, **async_native) + assert await python_facade.async_get_cache(**async_native) == {"writer": "async-native"} + + async_python: Final = completion_kwargs("async-python") + await python_facade.async_add_cache({"writer": "async-python"}, **async_python) + assert await native_facade.async_get_cache(**async_python) == {"writer": "async-python"} + + +@pytest.mark.parametrize("cache_factory", ROUND_TRIP_BACKENDS, indirect=True) +async def test_embedding_pipeline_stores_one_native_entry_per_input( + cache_factory: CacheFactory, monkeypatch: pytest.MonkeyPatch, request: pytest.FixtureRequest +) -> None: + require_rust(monkeypatch, cast(LiteLLMCacheType, request.node.callspec.params["cache_factory"])) + facade: Final = cache_factory() + assert_native_runtime(facade) + inputs: Final = [f"alpha {uuid4().hex}", f"beta {uuid4().hex}"] + result: Final = EmbeddingResponse( + model="text-embedding-3-small", + data=[ + {"object": "embedding", "index": 0, "embedding": [0.1, 0.2]}, + {"object": "embedding", "index": 1, "embedding": [0.3, 0.4]}, + ], + ) + await facade.async_add_cache_pipeline(result, model="text-embedding-3-small", input=inputs) + + keys: Final = [facade.get_cache_key(model="text-embedding-3-small", input=text) for text in inputs] + assert len(set(keys)) == len(inputs) + for text, expected in zip(inputs, ([0.1, 0.2], [0.3, 0.4]), strict=True): + cached = await facade.async_get_cache(model="text-embedding-3-small", input=text) + assert isinstance(cached, dict) + assert cached["embedding"] == expected + assert await facade.async_get_cache(model="text-embedding-3-small", input=inputs) is None + + +@pytest.mark.parametrize( + ("backend", "settings", "message"), + [ + pytest.param( + LiteLLMCacheType.VALKEY_SEMANTIC, + {"redis_url": "rediss://127.0.0.1:6390/0", "similarity_threshold": 0.8}, + "native Valkey semantic cache does not support TLS connections", + id="valkey-tls", + ), + pytest.param( + LiteLLMCacheType.VALKEY_SEMANTIC, + {"redis_url": "redis://127.0.0.1:6390/0?socket_timeout=1", "similarity_threshold": 0.8}, + "native Redis uses fixed socket timeouts; socket_timeout and socket_connect_timeout require Python", + id="valkey-socket-timeout", + ), + pytest.param( + LiteLLMCacheType.REDIS_SEMANTIC, + {"redis_url": "rediss://127.0.0.1:6380", "similarity_threshold": 0.8}, + "native Redis semantic cache does not support TLS or query options in redis_url", + id="redis-semantic-tls", + ), + pytest.param( + LiteLLMCacheType.REDIS_SEMANTIC, + {"redis_url": "redis://127.0.0.1:6379?socket_timeout=1", "similarity_threshold": 0.8}, + "native Redis semantic cache does not support TLS or query options in redis_url", + id="redis-semantic-query", + ), + ], +) +def test_semantic_settings_the_native_client_cannot_honor_decline( + monkeypatch: pytest.MonkeyPatch, backend: LiteLLMCacheType, settings: dict[str, object], message: str +) -> None: + require_rust(monkeypatch, backend) + with pytest.raises(RuntimeError, match=f"declined the cache: {message}"): + Cache(type=backend, **settings) + + +class _SemanticHit: + """A native semantic runtime that answers every lookup with one cached response.""" + + kind: Final = "native" + + def lookup_semantic(self, request: object) -> tuple[object, float | None]: + return {"answer": 42}, 0.97 + + async def async_lookup_semantic(self, request: object) -> tuple[object, float | None]: + return {"answer": 42}, 0.97 + + +@pytest.mark.parametrize("semantic_type", [LiteLLMCacheType.QDRANT_SEMANTIC, LiteLLMCacheType.REDIS_SEMANTIC]) +@pytest.mark.parametrize("use_async", [False, True], ids=["sync", "async"]) +def test_native_semantic_hit_stamps_similarity_on_request_metadata( + semantic_type: LiteLLMCacheType, use_async: bool +) -> None: + """Python semantic backends write `metadata["semantic-similarity"]` on every lookup, and the + facade copies it to the caller's metadata; the native path must report it the same way.""" + facade: Final = Cache() + facade.type = semantic_type + facade._native_cache = ResponseCacheRuntime(cast(NativeResponseCacheRuntime, _SemanticHit())) # pyright: ignore[reportPrivateUsage] # the native path under test has no public setter + metadata: Final[dict[str, object]] = {} + kwargs: Final = { + "cache_key": "semantic-key", + "messages": [{"role": "user", "content": "hello"}], + "metadata": metadata, + } + + result: Final = asyncio.run(facade.async_get_cache(**kwargs)) if use_async else facade.get_cache(**kwargs) + + assert result == {"answer": 42} + assert metadata["semantic-similarity"] == 0.97 diff --git a/tests/test_litellm_rust/cache/test_s3.py b/tests/test_litellm_rust/cache/test_s3.py new file mode 100644 index 00000000000..044bfc39f8d --- /dev/null +++ b/tests/test_litellm_rust/cache/test_s3.py @@ -0,0 +1,187 @@ +import json +import time +from datetime import datetime +from types import SimpleNamespace +from typing import Final, cast +from unittest.mock import Mock + +import boto3 +import botocore.config +import pytest + +from litellm.caching.caching import Cache +from litellm.caching.s3_cache import S3Cache +from litellm.types.caching import LiteLLMCacheType +from tests.test_litellm_rust.support.cache import CacheTestHandle, CacheTestResolver, request +from tests.test_litellm_rust.support.isolation import rebound +from tests.test_litellm_rust.support.s3_stub import S3Stub + +pytestmark: Final = pytest.mark.requires_rust_extension + + +def python_s3(url: str) -> S3Cache: + return S3Cache( + s3_bucket_name="cache-bucket", + s3_region_name="us-east-1", + s3_endpoint_url=url, + s3_aws_access_key_id="key", + s3_aws_secret_access_key="secret", + s3_path="team", + ) + + +async def test_s3_reads_python_entries_and_writes_with_python_metadata(s3_stub: S3Stub) -> None: + python_cache: Final = python_s3(s3_stub.url) + response: Final = {"choices": [{"text": "cached"}], "usage": {"total_tokens": 3}} + python_cache.set_cache("sync:key", {"timestamp": time.time(), "response": response}, ttl=90) + python_cache.set_cache("plain", {"timestamp": time.time(), "response": response}) + s3_stub.put_object("team/malformed", b"not a cache entry") + s3_stub.put_object( + "team/expired", + json.dumps({"timestamp": time.time(), "response": response}).encode(), + {"expires": "Thu, 01 Jan 1970 00:00:00 GMT"}, + ) + binding: Final = CacheTestResolver( + SimpleNamespace( + cache=CacheTestHandle.s3( + "cache-bucket", + region="us-east-1", + endpoint_url=s3_stub.url, + key_prefix="team/", + access_key_id="key", + secret_access_key="secret", + ) + ) + ).resolve() + + assert binding.lookup(request("sync:key")) == response + assert await binding.async_lookup(request("plain")) == response + assert binding.lookup(request("malformed")) is None + assert binding.lookup(request("expired")) is None + assert binding.lookup(request("absent")) is None + + binding.store({**request("native:key"), "ttl_seconds": 90.0}, response) + await binding.async_store(request("no_ttl"), response) + stored: Final = s3_stub.objects["team/native/key"] + assert stored.headers["content-type"] == "application/json" + assert stored.headers["content-language"] == "en" + assert stored.headers["content-disposition"] == 'inline; filename="team/native/key.json"' + assert stored.headers["cache-control"] == "immutable, max-age=90, s-maxage=90" + expires: Final = cast(datetime, s3_stub.expires("team/native/key")) + remaining: Final = (expires - datetime.now(expires.tzinfo)).total_seconds() + assert 60 < remaining <= 91 + no_ttl: Final = s3_stub.objects["team/no_ttl"] + assert no_ttl.headers["cache-control"] == "immutable, max-age=31536000, s-maxage=31536000" + assert "expires" not in no_ttl.headers + assert python_cache.get_cache("native:key")["response"] == response + + partial: Final = await binding.async_lookup_batch([request("native:key"), request("absent"), request("malformed")]) + assert partial == {"values": [response, None, None], "missing_indices": [1, 2]} + + +def test_s3_facade_binds_only_exact_configuration_and_falls_back_on_mutation(s3_stub: S3Stub) -> None: + facade: Final = Cache( + type=LiteLLMCacheType.S3, + s3_bucket_name="cache-bucket", + s3_region_name="us-east-1", + s3_endpoint_url=s3_stub.url, + s3_aws_access_key_id="key", + s3_aws_secret_access_key="secret", + s3_path="team", + ) + handle: Final = CacheTestHandle.s3( + "cache-bucket", + region="us-east-1", + endpoint_url=s3_stub.url, + key_prefix="team/", + access_key_id="key", + secret_access_key="secret", + ) + with pytest.raises(TypeError, match="buckets must match"): + CacheTestHandle.s3("other", region="us-east-1", endpoint_url=s3_stub.url)._bind_facade(facade) + with pytest.raises(TypeError, match="key prefixes must match"): + CacheTestHandle.s3( + "cache-bucket", region="us-east-1", endpoint_url=s3_stub.url, key_prefix="other/" + )._bind_facade(facade) + handle._bind_facade(facade) + resolver: Final = CacheTestResolver(SimpleNamespace(cache=facade)) + binding: Final = resolver.resolve() + assert binding.kind == "native" + + handler: Final = Mock() + facade.cache.s3_client.meta.events.register("before-call.s3.*", handler) + binding.store(request("native"), {"answer": 1}) + assert binding.lookup(request("native")) == {"answer": 1} + assert handler.call_count == 0 + assert "team/native" in s3_stub.objects + + with rebound(facade.cache, "bucket_name", "other"): + assert resolver.resolve().kind == "python_callback" + other_client: Final = boto3.client( + "s3", + region_name="us-east-1", + endpoint_url=s3_stub.url, + aws_access_key_id="key", + aws_secret_access_key="secret", + ) + with rebound(facade.cache, "s3_client", other_client): + assert resolver.resolve().kind == "python_callback" + + class CustomS3Cache(S3Cache): + pass + + subclassed: Final = Cache( + type=LiteLLMCacheType.S3, + s3_bucket_name="cache-bucket", + s3_region_name="us-east-1", + s3_endpoint_url=s3_stub.url, + s3_aws_access_key_id="key", + s3_aws_secret_access_key="secret", + s3_path="team", + ) + subclassed.cache = CustomS3Cache( + s3_bucket_name="cache-bucket", + s3_region_name="us-east-1", + s3_endpoint_url=s3_stub.url, + s3_aws_access_key_id="key", + s3_aws_secret_access_key="secret", + s3_path="team", + ) + with pytest.raises(TypeError): + handle._bind_facade(subclassed) + assert CacheTestResolver(SimpleNamespace(cache=subclassed)).resolve().kind == "python_callback" + + +def test_s3_facade_rejects_configurations_that_require_python(s3_stub: S3Stub) -> None: + handle: Final = CacheTestHandle.s3( + "cache-bucket", + region="us-east-1", + endpoint_url=s3_stub.url, + key_prefix="team/", + access_key_id="key", + secret_access_key="secret", + ) + unverified: Final = Cache( + type=LiteLLMCacheType.S3, + s3_bucket_name="cache-bucket", + s3_region_name="us-east-1", + s3_endpoint_url="https://s3.example.test", + s3_aws_access_key_id="key", + s3_aws_secret_access_key="secret", + s3_path="team", + s3_verify=False, + ) + with pytest.raises(TypeError, match="requires Python"): + handle._bind_facade(unverified) + proxied: Final = Cache( + type=LiteLLMCacheType.S3, + s3_bucket_name="cache-bucket", + s3_region_name="us-east-1", + s3_endpoint_url=s3_stub.url, + s3_aws_access_key_id="key", + s3_aws_secret_access_key="secret", + s3_path="team", + s3_config=botocore.config.Config(proxies={"https": "http://proxy.test"}), + ) + with pytest.raises(TypeError, match="requires Python"): + handle._bind_facade(proxied) diff --git a/tests/test_litellm_rust/test_valkey_semantic_cache_native.py b/tests/test_litellm_rust/cache/test_valkey_semantic.py similarity index 100% rename from tests/test_litellm_rust/test_valkey_semantic_cache_native.py rename to tests/test_litellm_rust/cache/test_valkey_semantic.py diff --git a/tests/test_litellm_rust/ocr/test_requests.py b/tests/test_litellm_rust/ocr/test_requests.py index 0d3b8ba472d..d09e60784fa 100644 --- a/tests/test_litellm_rust/ocr/test_requests.py +++ b/tests/test_litellm_rust/ocr/test_requests.py @@ -244,6 +244,21 @@ def test_native_ocr_maps_provider_400_with_public_provider_details(ocr_server: R assert "invalid OCR request" in str(caught.value) +def test_native_ocr_encodes_python_file_input_and_drops_unknown_arguments(ocr_server: RecordingServer) -> None: + response: Final = call_native_ocr( + ocr_server, + document={"type": "file", "file": BytesIO(b"abc"), "mime_type": "image/png"}, + opaque_extension=object(), + ) + + assert response.pages[0].markdown == "native OCR response" + assert_native_request(ocr_server) + assert ocr_server.requests[0].body == { + "model": "mistral-ocr-latest", + "document": {"type": "image_url", "image_url": "data:image/png;base64,YWJj"}, + } + + class TokenAbort(BaseException): pass diff --git a/tests/test_litellm_rust/support/cache.py b/tests/test_litellm_rust/support/cache.py new file mode 100644 index 00000000000..41eb4d25257 --- /dev/null +++ b/tests/test_litellm_rust/support/cache.py @@ -0,0 +1,40 @@ +from typing import Final, Protocol +from uuid import uuid4 + +import pytest + +from litellm.caching.caching import Cache +from litellm.rust_bridge import _native, catalog +from litellm.rust_bridge.catalog import CacheRule +from litellm.rust_bridge.configuration import Rollout +from litellm.rust_bridge.response_cache import ResponseCacheRuntime +from litellm.types.caching import LiteLLMCacheType + +CacheTestHandle: Final = _native._CacheTestHandle # pyright: ignore[reportPrivateUsage] # test-only handle has no public module name + + +CacheTestResolver: Final = _native._CacheTestResolver # pyright: ignore[reportPrivateUsage] # test-only resolver has no public module name + + +class CacheLookup(Protocol): + def get_cache(self, **kwargs: object) -> object: ... + def flush_cache(self) -> object: ... + + +def request(key: str = "key") -> dict[str, object]: + return {"key": {"preset": key}} + + +def require_rust(monkeypatch: pytest.MonkeyPatch, backend: LiteLLMCacheType) -> None: + monkeypatch.setattr(catalog, "RULES", (CacheRule(Rollout.RUST_REQUIRED, backends=frozenset({backend})),)) + + +def assert_native_runtime(facade: Cache) -> ResponseCacheRuntime: + runtime: Final = facade._native_cache # pyright: ignore[reportPrivateUsage] # the activation under test has no public accessor + assert isinstance(runtime, ResponseCacheRuntime) + assert runtime.kind == "native" + return runtime + + +def completion_kwargs(label: str) -> dict[str, object]: + return {"model": "gpt-4o", "messages": [{"role": "user", "content": f"{label} {uuid4().hex}"}]} diff --git a/tests/test_litellm_rust/test_cache.py b/tests/test_litellm_rust/test_cache.py deleted file mode 100644 index 96b3674fde3..00000000000 --- a/tests/test_litellm_rust/test_cache.py +++ /dev/null @@ -1,2397 +0,0 @@ -import asyncio -import contextvars -import gc -import hashlib -import http.server -import json -import math -import os -import threading -import time -import uuid -import weakref -from collections.abc import Callable, Generator -from contextlib import ExitStack -from datetime import datetime -from pathlib import Path -from types import SimpleNamespace -from typing import Final, Protocol, TypeAlias, cast -from unittest.mock import Mock -from urllib.parse import urlparse -from uuid import uuid4 - -import boto3 -import botocore.config -import diskcache -import fakeredis -import pytest -import redis -from azure.storage.blob import ContainerClient - -import litellm -from litellm.caching.azure_blob_cache import AzureBlobCache -from litellm.caching.caching import Cache, disable_cache, enable_cache, update_cache -from litellm.caching.disk_cache import DiskCache -from litellm.caching.gcs_cache import GCSCache -from litellm.caching.in_memory_cache import InMemoryCache -from litellm.caching.redis_cluster_cache import RedisClusterCache -from litellm.caching.redis_semantic_cache import RedisSemanticCache -from litellm.caching.s3_cache import S3Cache -from litellm.rust_bridge import _native, catalog -from litellm.rust_bridge.catalog import CacheRule, Route, RouteRule, SecretManagerRule -from litellm.rust_bridge.configuration import Rollout -from litellm.rust_bridge.response_cache import NativeResponseCacheRuntime, ResponseCacheRuntime, resolve_response_cache -from litellm.types.caching import LiteLLMCacheType -from litellm.types.llms.custom_llm import CustomLLMItem -from litellm.types.utils import EmbeddingResponse -from tests.test_litellm_rust.support.fake_gcs import FakeGcs -from tests.test_litellm_rust.support.isolation import rebound -from tests.test_litellm_rust.support.s3_stub import S3Stub - -_CacheTestHandle: Final = _native._CacheTestHandle # pyright: ignore[reportPrivateUsage] # test-only handle has no public module name -_CacheTestResolver: Final = _native._CacheTestResolver # pyright: ignore[reportPrivateUsage] # test-only resolver has no public module name - -pytestmark: Final = pytest.mark.requires_rust_extension - - -class CacheLookup(Protocol): - def get_cache(self, **kwargs: object) -> object: ... - def flush_cache(self) -> object: ... - - -def request(key: str = "key") -> dict[str, object]: - return {"key": {"preset": key}} - - -def qdrant_request( - key: str, - messages: list[dict[str, object]], - **kwargs: object, -) -> dict[str, object]: - return {**request(key), "messages": messages, **kwargs} - - -def embedding_vector(text: str) -> list[float]: - raw: Final = hashlib.sha256(text.encode()).digest()[:8] - values: Final = [byte / 127.5 - 1 for byte in raw] - norm: Final = math.sqrt(sum(value * value for value in values)) - return [value / norm for value in values] - - -@pytest.fixture -def qdrant_url() -> str: - value: Final[str | None] = os.environ.get("QDRANT_URL") - if not value: - pytest.skip("QDRANT_URL is required for Qdrant semantic cache tests") - return value.rstrip("/") - - -@pytest.fixture -def fake_embedding_endpoint(monkeypatch: pytest.MonkeyPatch) -> Generator[str]: - class EmbeddingHandler(http.server.BaseHTTPRequestHandler): - def do_POST(self) -> None: - length: Final = int(self.headers["Content-Length"]) - body: Final = json.loads(self.rfile.read(length)) - text: Final = body["input"] - response: Final = { - "object": "list", - "data": [ - { - "object": "embedding", - "index": 0, - "embedding": embedding_vector(text), - } - ], - "model": body["model"], - "usage": {"prompt_tokens": 1, "total_tokens": 1}, - } - encoded: Final = json.dumps(response).encode() - self.send_response(200) - self.send_header("Content-Type", "application/json") - self.send_header("Content-Length", str(len(encoded))) - self.end_headers() - self.wfile.write(encoded) - - def log_message(self, *_args: object) -> None: - return - - server: Final = http.server.ThreadingHTTPServer(("127.0.0.1", 0), EmbeddingHandler) - worker: Final = threading.Thread(target=server.serve_forever, daemon=True) - worker.start() - monkeypatch.setenv("OPENAI_API_BASE", f"http://127.0.0.1:{server.server_address[1]}") - monkeypatch.setenv("OPENAI_API_KEY", "sk-test") - try: - yield f"http://127.0.0.1:{server.server_address[1]}" - finally: - server.shutdown() - server.server_close() - worker.join(timeout=5) - - -@pytest.fixture -def redis_url() -> Generator[str]: - server: Final = fakeredis.TcpFakeServer(("127.0.0.1", 0), server_type="redis") - worker: Final = threading.Thread(target=server.serve_forever, daemon=True) - worker.start() - try: - yield f"redis://127.0.0.1:{server.server_address[1]}" - finally: - server.shutdown() - server.server_close() - worker.join(timeout=5) - - -@pytest.fixture -def fake_gcs() -> Generator[FakeGcs]: - server: Final = FakeGcs() - try: - yield server - finally: - server.close() - - -@pytest.fixture -def azure_blob_facade() -> Generator[Cache]: - account_url: Final = os.environ.get("AZURE_BLOB_CACHE_ACCOUNT_URL") - if account_url is None: - pytest.skip( - "live Azure Blob parity needs AZURE_BLOB_CACHE_ACCOUNT_URL plus DefaultAzureCredential inputs in the environment" - ) - facade: Final = Cache( - type=LiteLLMCacheType.AZURE_BLOB, - azure_account_url=account_url, - azure_blob_container=f"litellm-parity-{uuid.uuid4().hex[:12]}", - ) - backend: Final = facade.cache - assert isinstance(backend, AzureBlobCache) - try: - yield facade - finally: - backend.container_client.delete_container() - asyncio.run(backend.disconnect()) - - -def azure_blob_handle(facade: Cache) -> _native._CacheTestHandle: - backend: Final = facade.cache - assert isinstance(backend, AzureBlobCache) - return _native._CacheTestHandle.azure_blob( - backend.container_client.url.removesuffix(f"/{backend.container_client.container_name}"), - backend.container_client.container_name, - ) - - -@pytest.fixture -def cluster_nodes() -> tuple[tuple[str, int], ...]: - configured: Final = os.environ.get("LITELLM_TEST_REDIS_CLUSTER_NODES") - if not configured: - pytest.skip("LITELLM_TEST_REDIS_CLUSTER_NODES is not set") - return tuple((host, int(port)) for host, _, port in (node.partition(":") for node in configured.split(","))) - - -def test_existing_constructor_and_global_are_unchanged() -> None: - facade: Final = Cache(type=LiteLLMCacheType.LOCAL) - assert type(facade.cache) is InMemoryCache - assert "_native_cache_handle" not in vars(facade) - assert resolve_response_cache(facade) is None - with rebound(litellm, "cache", facade): - resolver: Final = _CacheTestResolver(litellm) - assert resolver.resolve().kind == "python_callback" - resolver.resolve().store(None, {"answer": 7}, callback_kwargs={"cache_key": "key"}) - assert cast(CacheLookup, facade).get_cache(cache_key="key") == {"answer": 7} - - -async def test_catalog_constructs_native_runtime_from_public_cache_configuration() -> None: - rules: Final = ( - RouteRule(Route.OCR, Rollout.PYTHON_ONLY), - SecretManagerRule(Rollout.PYTHON_ONLY, systems=frozenset({"local"})), - CacheRule(Rollout.RUST_REQUIRED, backends=frozenset({"local"})), - ) - facade: Final = Cache(type=LiteLLMCacheType.LOCAL) - runtime: Final = resolve_response_cache(facade, rules) - assert isinstance(runtime, ResponseCacheRuntime) - assert runtime.kind == "native" - - sync_request: Final = runtime.request(facade, {"cache_key": "sync"}) - assert sync_request is not None - runtime.store(sync_request, {"answer": 1}) - assert runtime.lookup(sync_request) == {"answer": 1} - assert facade.cache.get_cache("sync") is None - - async_request: Final = runtime.request(facade, {"cache_key": "async"}) - assert async_request is not None - await runtime.async_store(async_request, {"answer": 2}) - assert await runtime.async_lookup(async_request) == {"answer": 2} - assert await facade.cache.async_get_cache("async") is None - - requests: Final = (sync_request, async_request) - expected: Final = { - "values": [{"answer": 1}, {"answer": 2}], - "missing_indices": [], - } - assert runtime.lookup_batch(requests) == expected - assert await runtime.async_lookup_batch(requests) == expected - - await runtime.async_flush() - assert runtime.lookup(sync_request) is None - assert await runtime.async_lookup(async_request) is None - - -async def test_inference_resolver_uses_the_configured_native_cache_directly() -> None: - rules: Final = ( - RouteRule(Route.OCR, Rollout.PYTHON_ONLY), - SecretManagerRule(Rollout.PYTHON_ONLY, systems=frozenset({"local"})), - CacheRule(Rollout.RUST_REQUIRED, backends=frozenset({"local"})), - ) - facade: Final = Cache(type=LiteLLMCacheType.LOCAL) - runtime: Final = resolve_response_cache(facade, rules) - assert isinstance(runtime, ResponseCacheRuntime) - facade._native_cache = runtime - - selected: Final = _native._CacheResolver(SimpleNamespace(cache=facade)).resolve() - assert selected.kind == "native" - request: Final = runtime.request(facade, {"cache_key": "inference-native"}) - assert request is not None - await selected.async_store(request, {"answer": 42}) - assert await selected.async_lookup(request) == {"answer": 42} - assert await runtime.async_lookup(request) == {"answer": 42} - assert facade.cache.get_cache("inference-native") is None - - facade._native_cache = None - fallback: Final = _native._CacheResolver(SimpleNamespace(cache=facade)).resolve() - assert fallback.kind == "python_callback" - await fallback.async_store(None, {"answer": 7}, callback_kwargs={"cache_key": "inference-python"}) - assert facade.get_cache(cache_key="inference-python") == {"answer": 7} - assert facade.cache.get_cache("inference-python") is not None - - -async def test_inference_resolver_declines_a_native_runtime_whose_facade_changed() -> None: - rules: Final = ( - RouteRule(Route.OCR, Rollout.PYTHON_ONLY), - SecretManagerRule(Rollout.PYTHON_ONLY, systems=frozenset({"local"})), - CacheRule(Rollout.RUST_REQUIRED, backends=frozenset({"local"})), - ) - facade: Final = Cache(type=LiteLLMCacheType.LOCAL) - runtime: Final = resolve_response_cache(facade, rules) - assert isinstance(runtime, ResponseCacheRuntime) - facade._native_cache = runtime - stale_request: Final = runtime.request(facade, {"cache_key": "stale-only"}) - assert stale_request is not None - await runtime.async_store(stale_request, {"answer": "stale"}) - - replacement: Final = InMemoryCache() - facade.cache = replacement - with pytest.raises(_native.RustBridgeDeclined): - _native._CacheResolver(SimpleNamespace(cache=facade)).resolve() - assert await runtime.async_lookup(stale_request) == {"answer": "stale"} - assert replacement.get_cache("stale-only") is None - assert replacement.get_cache("swapped-backend") is None - - -def test_existing_global_lifecycle_remains_the_resolver_source_of_truth() -> None: - resolver: Final = _CacheTestResolver(litellm) - - enable_cache(type=LiteLLMCacheType.LOCAL, ttl=30) - enabled: Final = litellm.cache - assert isinstance(enabled, Cache) - assert enabled.ttl == 30 - assert resolver.resolve().kind == "python_callback" - - enable_cache(type=LiteLLMCacheType.LOCAL, ttl=60) - assert litellm.cache is enabled - - update_cache(type=LiteLLMCacheType.LOCAL, ttl=60) - updated: Final = litellm.cache - assert isinstance(updated, Cache) - assert updated is not enabled - assert updated.ttl == 60 - - disable_cache() - assert litellm.cache is None - assert resolver.resolve().kind == "disabled" - - -async def test_native_bindings_survive_replacement_and_capture_writes_before_dispatch() -> None: - namespace: Final = SimpleNamespace(cache=_CacheTestHandle.memory()) - resolver: Final = _CacheTestResolver(namespace) - selected: Final = resolver.resolve() - assert selected.kind == "native" - selected.store(request(), {"answer": 1}) - assert await selected.async_lookup(request()) == {"answer": 1} - with rebound(namespace, "cache", _CacheTestHandle.memory()): - replacement: Final = resolver.resolve() - await selected.async_store(request(), {"answer": 2}) - assert replacement.lookup(request()) is None - assert selected.lookup(request()) == {"answer": 2} - with rebound(namespace, "cache", None): - disabled: Final = resolver.resolve() - assert disabled.kind == "disabled" - assert disabled.lookup(None) is None - await disabled.async_store(None, object()) - assert await disabled.async_lookup(None) is None - assert selected.lookup(request()) == {"answer": 2} - - -async def test_python_callback_preserves_identity_caller_task_context_and_errors() -> None: - context: Final = contextvars.ContextVar("cache_context", default="caller") - caller: Final = asyncio.current_task() - sentinel: Final = object() - failure: Final = RuntimeError("callback failed") - - class CustomCache: - async def async_get_cache(self, *, marker: object) -> object: - assert marker is sentinel - assert asyncio.current_task() is caller - context.set("callback") - return marker - - async def async_add_cache(self, response: object, *, marker: object) -> None: - assert response is sentinel - assert marker is sentinel - raise failure - - namespace: Final = SimpleNamespace(cache=CustomCache()) - binding: Final = _CacheTestResolver(namespace).resolve() - assert binding.kind == "python_callback" - assert await binding.async_lookup(None, callback_kwargs={"marker": sentinel}) is sentinel - assert context.get() == "callback" - with pytest.raises(RuntimeError) as caught: - await binding.async_store(None, sentinel, callback_kwargs={"marker": sentinel}) - assert caught.value is failure - - -async def test_callback_cancellation_stays_in_the_callers_task() -> None: - entered: Final = asyncio.Event() - finished: Final = asyncio.Event() - - class CustomCache: - async def async_get_cache(self) -> None: - entered.set() - try: - await asyncio.Future() - finally: - finished.set() - - binding: Final = _CacheTestResolver(SimpleNamespace(cache=CustomCache())).resolve() - - async def lookup() -> object: - return await binding.async_lookup(None, callback_kwargs={}) - - task: Final = asyncio.create_task(lookup()) - await entered.wait() - task.cancel() - with pytest.raises(asyncio.CancelledError): - await task - assert finished.is_set() - - -def test_registered_facade_uses_native_and_instance_overrides_fall_back() -> None: - facade: Final = Cache(type=LiteLLMCacheType.LOCAL) - handle: Final = _CacheTestHandle.memory() - handle._bind_facade(facade) - resolver: Final = _CacheTestResolver(SimpleNamespace(cache=facade)) - native: Final = resolver.resolve() - assert native.kind == "native" - native.store(request(), {"source": "native"}) - assert native.lookup(request()) == {"source": "native"} - assert cast(CacheLookup, facade).get_cache(cache_key="key") is None - sentinel: Final = object() - - def outer_override(**_kwargs: object) -> object: - return sentinel - - def backend_override(*_args: object, **_kwargs: object) -> dict[str, str]: - return {"source": "override"} - - with rebound(facade, "get_cache", outer_override): - fallback: Final = resolver.resolve() - assert fallback.kind == "python_callback" - assert fallback.lookup(None, callback_kwargs={"cache_key": "key"}) is sentinel - assert resolver.resolve().kind == "python_callback" - delattr(facade, "get_cache") - assert resolver.resolve().kind == "native" - with rebound(facade.cache, "get_cache", backend_override): - backend_fallback: Final = resolver.resolve() - assert backend_fallback.kind == "python_callback" - assert backend_fallback.lookup(None, callback_kwargs={"cache_key": "key"}) == {"source": "override"} - - -def test_facade_subclasses_backend_replacement_and_configuration_changes_are_not_bypassed() -> None: - class CustomCache(Cache): - pass - - handle: Final = _CacheTestHandle.memory() - with pytest.raises(TypeError): - handle._bind_facade(CustomCache(type=LiteLLMCacheType.LOCAL)) - facade: Final = Cache(type=LiteLLMCacheType.LOCAL) - handle._bind_facade(facade) - resolver: Final = _CacheTestResolver(SimpleNamespace(cache=facade)) - with rebound(facade, "cache", InMemoryCache()): - assert resolver.resolve().kind == "python_callback" - with rebound(facade, "ttl", 12): - assert resolver.resolve().kind == "python_callback" - with rebound(facade, "semantic_cache_scope", "end_user"): - assert resolver.resolve().kind == "python_callback" - - def custom_key(**_kwargs: object) -> str: - return "custom" - - with rebound(facade, "get_cache_key", custom_key): - assert resolver.resolve().kind == "python_callback" - assert resolver.resolve().kind == "python_callback" - delattr(facade, "get_cache_key") - assert resolver.resolve().kind == "native" - - -def test_resolver_and_callback_cycles_can_be_collected() -> None: - class CustomCache: - pass - - def cyclic_reference() -> weakref.ReferenceType[CustomCache]: - callback: Final = CustomCache() - namespace: Final = SimpleNamespace(cache=callback) - binding: Final = _CacheTestResolver(namespace).resolve() - setattr(callback, "binding", binding) - return weakref.ref(callback) - - reference: Final = cyclic_reference() - gc.collect() - assert reference() is None - - -async def test_redis_reads_python_sync_and_async_entries_and_writes_without_hidden_prefix(redis_url: str) -> None: - client: Final = redis.Redis.from_url(redis_url) - namespace: Final = SimpleNamespace(cache=_CacheTestHandle.redis(redis_url, namespace="team")) - binding: Final = _CacheTestResolver(namespace).resolve() - response: Final = {"choices": [{"text": "cached"}], "usage": {"total_tokens": 3}, "flag": True, "empty": None} - envelope: Final = {"timestamp": time.time(), "response": json.dumps(response)} - client.set("team:sync", str(envelope)) - client.set("team:async", json.dumps({"timestamp": time.time(), "response": response})) - client.set("team:raw", json.dumps(response)) - client.set("team:invalid", "not a cache entry") - assert binding.lookup(request("sync")) == response - assert await binding.async_lookup(request("team:async")) == response - assert binding.lookup(request("raw")) == response - assert await binding.async_lookup(request("invalid")) is None - await binding.async_store({**request("native"), "ttl_seconds": 12.0}, response) - stored: Final = client.get("team:native") - assert isinstance(stored, bytes) - assert json.loads(stored)["response"] == response - assert 0 < client.ttl("team:native") <= 12 - assert client.get("litellm-cache:team:native") is None - assert client.get("team:team:async") is None - client.close() - - -def test_invalid_duration_and_request_shape_fail_before_storage() -> None: - binding: Final = _CacheTestResolver(SimpleNamespace(cache=_CacheTestHandle.memory())).resolve() - for seconds in (-1.0, float("nan"), float("inf")): - with pytest.raises(ValueError, match="cache durations must be finite and nonnegative"): - binding.store({**request(), "ttl_seconds": seconds}, {"answer": 1}) - assert binding.lookup(request()) is None - with pytest.raises(ValueError, match="cache durations must be finite and nonnegative"): - _CacheTestHandle.memory(ttl_seconds=-1) - - -async def test_memory_size_policy_is_applied_by_the_native_host() -> None: - handle: Final = _CacheTestHandle.memory(capacity=2, max_entry_bytes=128) - binding: Final = _CacheTestResolver(SimpleNamespace(cache=handle)).resolve() - small: Final = {"answer": "ok"} - binding.store(request("small"), small) - assert await binding.async_lookup(request("small")) == small - await binding.async_store(request("large"), {"answer": "x" * 256}) - assert binding.lookup(request("large")) is None - assert binding.lookup(request("small")) == small - disabled: Final = _CacheTestResolver(SimpleNamespace(cache=_CacheTestHandle.memory(capacity=0))).resolve() - await disabled.async_store(request(), small) - assert await disabled.async_lookup(request()) is None - - -async def test_native_batch_lookup_and_store_report_partial_hits() -> None: - binding: Final = _CacheTestResolver(SimpleNamespace(cache=_CacheTestHandle.memory())).resolve() - requests: Final = [request("hit"), request("miss"), request("disabled")] - requests[2]["controls"] = { - "supported_call_type": True, - "configured": True, - "native_backend": True, - "default_on": True, - "caching": False, - "no_cache": False, - "no_store": False, - "use_cache": False, - } - await binding.async_store_batch(requests, [{"value": 1}, {"value": 2}, {"value": 3}]) - - partial: Final = await binding.async_lookup_batch(requests) - - assert partial == { - "values": [{"value": 1}, {"value": 2}, None], - "missing_indices": [2], - } - - -async def test_python_batch_callbacks_use_the_builtin_cache_api() -> None: - result: Final = object() - marker: Final = object() - - class CustomCache(Cache): - def get_cache(self, dynamic_cache_object: object = None, **kwargs: object) -> object: - return ("sync", kwargs) - - async def async_get_cache(self, dynamic_cache_object: object = None, **kwargs: object) -> object: - return ("async", kwargs) - - async def async_add_cache_pipeline( - self, result: object, dynamic_cache_object: object = None, **kwargs: object - ) -> object: - return result, kwargs - - binding: Final = _CacheTestResolver(SimpleNamespace(cache=CustomCache(type=LiteLLMCacheType.LOCAL))).resolve() - assert binding.kind == "python_callback" - requests: Final = [request("first"), request("second")] - kwargs: Final = [{"cache_key": "first"}, {"cache_key": "second"}] - - assert binding.lookup_batch(requests, callback_kwargs=kwargs) == [("sync", kwargs[0]), ("sync", kwargs[1])] - assert await binding.async_lookup_batch(requests, callback_kwargs=kwargs) == [ - ("async", kwargs[0]), - ("async", kwargs[1]), - ] - with pytest.raises(ValueError, match="equal lengths"): - binding.lookup_batch(requests, callback_kwargs=kwargs[:1]) - with pytest.raises(TypeError, match="callback_result"): - await binding.async_store_batch(requests, [1, 2], callback_kwargs={"marker": marker}) - stored: Final = cast( - tuple[object, dict[str, object]], - await binding.async_store_batch(requests, [1, 2], callback_result=result, callback_kwargs={"marker": marker}), - ) - assert stored[0] is result - assert stored[1] == {"marker": marker} - - -async def test_unmodified_builtin_cache_callbacks_can_ping_and_flush() -> None: - async def ping() -> str: - return "pong" - - cache: Final = Cache(type=LiteLLMCacheType.LOCAL) - cache.cache.set_cache("key", "value") - binding: Final = _CacheTestResolver(SimpleNamespace(cache=cache)).resolve() - assert binding.kind == "python_callback" - - setattr(cache.cache, "ping", ping) - assert await binding.ping() == "pong" - await binding.async_flush() - assert cache.cache.get_cache("key") is None - - -def test_facade_registration_rejects_mismatched_capacity() -> None: - facade: Final = Cache(type=LiteLLMCacheType.LOCAL) - with pytest.raises(TypeError, match="capacities must match"): - _CacheTestHandle.memory(capacity=7)._bind_facade(facade) - - -def test_azure_blob_facade_serves_natively_and_python_reads_the_same_blobs(azure_blob_facade: Cache) -> None: - backend: Final = azure_blob_facade.cache - assert isinstance(backend, AzureBlobCache) - handle: Final = azure_blob_handle(azure_blob_facade) - assert handle.backend == "azure-blob" - account_url: Final = backend.container_client.url.removesuffix(f"/{backend.container_client.container_name}") - with pytest.raises(TypeError, match="containers must match"): - _native._CacheTestHandle.azure_blob( - account_url, f"{backend.container_client.container_name}-other" - )._bind_facade(azure_blob_facade) - handle._bind_facade(azure_blob_facade) - resolver: Final = _native._CacheTestResolver(SimpleNamespace(cache=azure_blob_facade)) - native: Final = resolver.resolve() - assert native.kind == "native" - - response: Final = { - "choices": [{"text": "caf\u00e9 \u2603"}], - "usage": {"total_tokens": 3}, - "flag": True, - "empty": None, - } - native.store({**request("sync"), "ttl_seconds": 0.001}, response) - native.store(request("sync"), {"choices": [{"text": "second"}]}) - time.sleep(0.01) - stored: Final = json.loads(backend.container_client.download_blob("sync").readall()) - assert stored["response"] == response - assert isinstance(stored["timestamp"], float) - assert native.lookup(request("sync")) == response - assert cast(CacheLookup, azure_blob_facade).get_cache(cache_key="sync") == response - - backend.set_cache("python", {"timestamp": time.time(), "response": response}) - backend.set_cache("legacy", "bare legacy value") - backend.container_client.upload_blob("invalid", b"{not json", overwrite=True) - assert native.lookup(request("python")) == response - assert native.lookup(request("legacy")) == cast(CacheLookup, azure_blob_facade).get_cache(cache_key="legacy") - assert native.lookup_batch([request("python"), request("missing"), request("invalid"), request("sync")]) == { - "values": [response, None, None, response], - "missing_indices": [1, 2], - } - - with rebound(azure_blob_facade, "ttl", 12): - assert resolver.resolve().kind == "python_callback" - with rebound(backend, "container_client", ContainerClient.from_container_url(backend.container_client.url)): - assert resolver.resolve().kind == "python_callback" - - def custom_get(*_args: object, **_kwargs: object) -> None: - return None - - with rebound(backend, "get_cache", custom_get): - assert resolver.resolve().kind == "python_callback" - assert resolver.resolve().kind == "python_callback" - assert cast(CacheLookup, azure_blob_facade).get_cache(cache_key="sync") == response - - class CustomBlobCache(AzureBlobCache): - pass - - with rebound(azure_blob_facade, "cache", CustomBlobCache(account_url, backend.container_client.container_name)): - assert resolver.resolve().kind == "python_callback" - with pytest.raises(TypeError): - azure_blob_handle(azure_blob_facade)._bind_facade(azure_blob_facade) - - -async def test_azure_blob_native_async_writes_overwrite_batch_and_flush_like_python(azure_blob_facade: Cache) -> None: - backend: Final = azure_blob_facade.cache - assert isinstance(backend, AzureBlobCache) - azure_blob_handle(azure_blob_facade)._bind_facade(azure_blob_facade) - binding: Final = _native._CacheTestResolver(SimpleNamespace(cache=azure_blob_facade)).resolve() - assert binding.kind == "native" - ping: Final = cast(dict[str, object], await binding.ping()) - assert ping["status"] == "success", ping - - await binding.async_store(request("async"), {"value": 1}) - await binding.async_store({**request("async"), "ttl_seconds": 0.001}, {"value": 2}) - time.sleep(0.01) - assert await binding.async_lookup(request("async")) == {"value": 2} - assert await backend.async_get_cache("async") == json.loads( - backend.container_client.download_blob("async").readall() - ) - assert cast(CacheLookup, azure_blob_facade).get_cache(cache_key="async") == {"value": 2} - - await binding.async_store_batch([request("first"), request("second")], [{"value": 3}, {"value": 4}]) - assert await binding.async_lookup_batch([request("second"), request("missing"), request("first")]) == { - "values": [{"value": 4}, None, {"value": 3}], - "missing_indices": [1], - } - await binding.async_flush() - assert [blob.name for blob in backend.container_client.list_blobs()] == [] - assert await binding.async_lookup(request("async")) is None - - -async def test_redis_facade_buffers_native_async_writes(redis_url: str) -> None: - parsed: Final = urlparse(redis_url) - with rebound(litellm, "default_redis_ttl", 60): - facade: Final = Cache( - type=LiteLLMCacheType.REDIS, - host=parsed.hostname, - port=str(parsed.port), - redis_flush_size=2, - ) - with pytest.raises(TypeError, match="default TTLs must match"): - _CacheTestHandle.redis(redis_url, ttl_seconds=61)._bind_facade(facade) - with pytest.raises(TypeError, match="namespaces must match"): - _CacheTestHandle.redis(redis_url, namespace="other")._bind_facade(facade) - _CacheTestHandle.redis(redis_url, ttl_seconds=60)._bind_facade(facade) - binding: Final = _CacheTestResolver(SimpleNamespace(cache=facade)).resolve() - client: Final = redis.Redis.from_url(redis_url) - - with rebound(facade.cache, "redis_kwargs", {**facade.cache.redis_kwargs, "ssl": True}): - assert _CacheTestResolver(SimpleNamespace(cache=facade)).resolve().kind == "python_callback" - - pool: Final = facade.cache.redis_client.connection_pool - with rebound(pool, "connection_kwargs", {**pool.connection_kwargs, "db": 1}): - assert _CacheTestResolver(SimpleNamespace(cache=facade)).resolve().kind == "python_callback" - - await binding.async_store(request("first"), {"value": 1}) - assert client.get("first") is None - await binding.async_store(request("second"), {"value": 2}) - - assert client.get("first") is not None - assert client.get("second") is not None - await facade.cache.disconnect() - client.close() - - -async def test_disk_reads_python_entries_and_python_reads_native_entries(tmp_path: Path) -> None: - disk_cache: Final = DiskCache(disk_cache_dir=str(tmp_path)) - response: Final = {"choices": [{"text": "cached"}], "usage": {"total_tokens": 3}} - disk_cache.disk_cache.set( - "sync", - {"timestamp": time.time(), "response": json.dumps(response)}, - ) - disk_cache.disk_cache.set("async", json.dumps({"timestamp": time.time(), "response": response})) - disk_cache.disk_cache.set("raw", json.dumps(response)) - disk_cache.disk_cache.set("invalid", "not a cache entry") - disk_cache.disk_cache.set( - "large", - {"timestamp": time.time(), "response": {"text": "x" * 70_000}}, - ) - binding: Final = _native._CacheTestResolver( - SimpleNamespace(cache=_native._CacheTestHandle.disk(str(tmp_path))) - ).resolve() - - assert binding.lookup(request("sync")) == response - assert await binding.async_lookup(request("async")) == response - assert binding.lookup(request("raw")) == response - assert await binding.async_lookup(request("invalid")) is None - assert binding.lookup(request("large")) == {"text": "x" * 70_000} - - await binding.async_store({**request("native"), "ttl_seconds": 12.0}, response) - stored_response: Final = disk_cache.get_cache("native") - assert isinstance(stored_response, dict) - assert stored_response["response"] == response - stored, expire_time = disk_cache.disk_cache.get("native", expire_time=True) - assert stored is not None - assert time.time() < expire_time <= time.time() + 12.0 - await binding.async_store(request("no-ttl"), response) - _, no_expiry = disk_cache.disk_cache.get("no-ttl", expire_time=True) - assert no_expiry is None - - -async def test_disk_entries_survive_a_fresh_handle_and_expire_on_time(tmp_path: Path) -> None: - first: Final = _native._CacheTestResolver( - SimpleNamespace(cache=_native._CacheTestHandle.disk(str(tmp_path))) - ).resolve() - await first.async_store(request("persistent"), {"value": "persistent"}) - await first.async_store({**request("expiring"), "ttl_seconds": 0.3}, {"value": "expiring"}) - fresh: Final = _native._CacheTestResolver( - SimpleNamespace(cache=_native._CacheTestHandle.disk(str(tmp_path))) - ).resolve() - assert fresh.lookup(request("persistent")) == {"value": "persistent"} - assert fresh.lookup(request("expiring")) == {"value": "expiring"} - await asyncio.sleep(0.4) - assert fresh.lookup(request("expiring")) is None - assert fresh.lookup(request("persistent")) == {"value": "persistent"} - - -def test_disk_facade_registers_and_store_changes_fall_back(tmp_path: Path) -> None: - facade: Final = Cache(type=LiteLLMCacheType.DISK, disk_cache_dir=str(tmp_path)) - with pytest.raises(TypeError, match="directories must match"): - _native._CacheTestHandle.disk(str(tmp_path / "other"))._bind_facade(facade) - handle: Final = _native._CacheTestHandle.disk(str(tmp_path)) - handle._bind_facade(facade) - resolver: Final = _native._CacheTestResolver(SimpleNamespace(cache=facade)) - binding: Final = resolver.resolve() - assert binding.kind == "native" - binding.store(request("native"), {"value": "native"}) - assert facade.get_cache(cache_key="native") == {"value": "native"} - - with rebound(facade.cache, "disk_cache", diskcache.Cache(str(tmp_path))): - assert resolver.resolve().kind == "python_callback" - assert resolver.resolve().kind == "native" - - class CustomDiskCache(DiskCache): - pass - - with rebound(facade, "cache", CustomDiskCache(disk_cache_dir=str(tmp_path))): - assert resolver.resolve().kind == "python_callback" - - class CustomStore(diskcache.Cache): - pass - - custom_facade: Final = Cache(type=LiteLLMCacheType.DISK, disk_cache_dir=str(tmp_path)) - custom_facade.cache.disk_cache = CustomStore(str(tmp_path)) - with pytest.raises(TypeError, match="built-in diskcache store"): - _native._CacheTestHandle.disk(str(tmp_path))._bind_facade(custom_facade) - - -async def test_disk_native_batch_lookup_and_store_report_partial_hits(tmp_path: Path) -> None: - binding: Final = _native._CacheTestResolver( - SimpleNamespace(cache=_native._CacheTestHandle.disk(str(tmp_path))) - ).resolve() - requests: Final = [request("hit"), request("miss"), request("disabled")] - requests[2]["controls"] = { - "supported_call_type": True, - "configured": True, - "native_backend": True, - "default_on": True, - "caching": False, - "no_cache": False, - "no_store": False, - "use_cache": False, - } - await binding.async_store_batch(requests, [{"value": 1}, {"value": 2}, {"value": 3}]) - - partial: Final = await binding.async_lookup_batch(requests) - - assert partial == { - "values": [{"value": 1}, {"value": 2}, None], - "missing_indices": [2], - } - - -@pytest.fixture -def s3_stub() -> Generator[S3Stub]: - stub: Final = S3Stub() - try: - yield stub - finally: - stub.close() - - -def python_s3(url: str) -> S3Cache: - return S3Cache( - s3_bucket_name="cache-bucket", - s3_region_name="us-east-1", - s3_endpoint_url=url, - s3_aws_access_key_id="key", - s3_aws_secret_access_key="secret", - s3_path="team", - ) - - -async def test_s3_reads_python_entries_and_writes_with_python_metadata(s3_stub: S3Stub) -> None: - python_cache: Final = python_s3(s3_stub.url) - response: Final = {"choices": [{"text": "cached"}], "usage": {"total_tokens": 3}} - python_cache.set_cache("sync:key", {"timestamp": time.time(), "response": response}, ttl=90) - python_cache.set_cache("plain", {"timestamp": time.time(), "response": response}) - s3_stub.put_object("team/malformed", b"not a cache entry") - s3_stub.put_object( - "team/expired", - json.dumps({"timestamp": time.time(), "response": response}).encode(), - {"expires": "Thu, 01 Jan 1970 00:00:00 GMT"}, - ) - binding: Final = _native._CacheTestResolver( - SimpleNamespace( - cache=_native._CacheTestHandle.s3( - "cache-bucket", - region="us-east-1", - endpoint_url=s3_stub.url, - key_prefix="team/", - access_key_id="key", - secret_access_key="secret", - ) - ) - ).resolve() - - assert binding.lookup(request("sync:key")) == response - assert await binding.async_lookup(request("plain")) == response - assert binding.lookup(request("malformed")) is None - assert binding.lookup(request("expired")) is None - assert binding.lookup(request("absent")) is None - - binding.store({**request("native:key"), "ttl_seconds": 90.0}, response) - await binding.async_store(request("no_ttl"), response) - stored: Final = s3_stub.objects["team/native/key"] - assert stored.headers["content-type"] == "application/json" - assert stored.headers["content-language"] == "en" - assert stored.headers["content-disposition"] == 'inline; filename="team/native/key.json"' - assert stored.headers["cache-control"] == "immutable, max-age=90, s-maxage=90" - expires: Final = cast(datetime, s3_stub.expires("team/native/key")) - remaining: Final = (expires - datetime.now(expires.tzinfo)).total_seconds() - assert 60 < remaining <= 91 - no_ttl: Final = s3_stub.objects["team/no_ttl"] - assert no_ttl.headers["cache-control"] == "immutable, max-age=31536000, s-maxage=31536000" - assert "expires" not in no_ttl.headers - assert python_cache.get_cache("native:key")["response"] == response - - partial: Final = await binding.async_lookup_batch([request("native:key"), request("absent"), request("malformed")]) - assert partial == {"values": [response, None, None], "missing_indices": [1, 2]} - - -def test_s3_facade_binds_only_exact_configuration_and_falls_back_on_mutation(s3_stub: S3Stub) -> None: - facade: Final = Cache( - type=LiteLLMCacheType.S3, - s3_bucket_name="cache-bucket", - s3_region_name="us-east-1", - s3_endpoint_url=s3_stub.url, - s3_aws_access_key_id="key", - s3_aws_secret_access_key="secret", - s3_path="team", - ) - handle: Final = _native._CacheTestHandle.s3( - "cache-bucket", - region="us-east-1", - endpoint_url=s3_stub.url, - key_prefix="team/", - access_key_id="key", - secret_access_key="secret", - ) - with pytest.raises(TypeError, match="buckets must match"): - _native._CacheTestHandle.s3("other", region="us-east-1", endpoint_url=s3_stub.url)._bind_facade(facade) - with pytest.raises(TypeError, match="key prefixes must match"): - _native._CacheTestHandle.s3( - "cache-bucket", region="us-east-1", endpoint_url=s3_stub.url, key_prefix="other/" - )._bind_facade(facade) - handle._bind_facade(facade) - resolver: Final = _native._CacheTestResolver(SimpleNamespace(cache=facade)) - binding: Final = resolver.resolve() - assert binding.kind == "native" - - handler: Final = Mock() - facade.cache.s3_client.meta.events.register("before-call.s3.*", handler) - binding.store(request("native"), {"answer": 1}) - assert binding.lookup(request("native")) == {"answer": 1} - assert handler.call_count == 0 - assert "team/native" in s3_stub.objects - - with rebound(facade.cache, "bucket_name", "other"): - assert resolver.resolve().kind == "python_callback" - other_client: Final = boto3.client( - "s3", - region_name="us-east-1", - endpoint_url=s3_stub.url, - aws_access_key_id="key", - aws_secret_access_key="secret", - ) - with rebound(facade.cache, "s3_client", other_client): - assert resolver.resolve().kind == "python_callback" - - class CustomS3Cache(S3Cache): - pass - - subclassed: Final = Cache( - type=LiteLLMCacheType.S3, - s3_bucket_name="cache-bucket", - s3_region_name="us-east-1", - s3_endpoint_url=s3_stub.url, - s3_aws_access_key_id="key", - s3_aws_secret_access_key="secret", - s3_path="team", - ) - subclassed.cache = CustomS3Cache( - s3_bucket_name="cache-bucket", - s3_region_name="us-east-1", - s3_endpoint_url=s3_stub.url, - s3_aws_access_key_id="key", - s3_aws_secret_access_key="secret", - s3_path="team", - ) - with pytest.raises(TypeError): - handle._bind_facade(subclassed) - assert _native._CacheTestResolver(SimpleNamespace(cache=subclassed)).resolve().kind == "python_callback" - - -def test_s3_facade_rejects_configurations_that_require_python(s3_stub: S3Stub) -> None: - handle: Final = _native._CacheTestHandle.s3( - "cache-bucket", - region="us-east-1", - endpoint_url=s3_stub.url, - key_prefix="team/", - access_key_id="key", - secret_access_key="secret", - ) - unverified: Final = Cache( - type=LiteLLMCacheType.S3, - s3_bucket_name="cache-bucket", - s3_region_name="us-east-1", - s3_endpoint_url="https://s3.example.test", - s3_aws_access_key_id="key", - s3_aws_secret_access_key="secret", - s3_path="team", - s3_verify=False, - ) - with pytest.raises(TypeError, match="requires Python"): - handle._bind_facade(unverified) - proxied: Final = Cache( - type=LiteLLMCacheType.S3, - s3_bucket_name="cache-bucket", - s3_region_name="us-east-1", - s3_endpoint_url=s3_stub.url, - s3_aws_access_key_id="key", - s3_aws_secret_access_key="secret", - s3_path="team", - s3_config=botocore.config.Config(proxies={"https": "http://proxy.test"}), - ) - with pytest.raises(TypeError, match="requires Python"): - handle._bind_facade(proxied) - - -async def test_gcs_reads_python_entries_and_writes_python_compatible_objects( - fake_gcs: FakeGcs, monkeypatch: pytest.MonkeyPatch -) -> None: - monkeypatch.delenv("GCS_PATH_SERVICE_ACCOUNT", raising=False) - monkeypatch.delenv("GCS_BUCKET_NAME", raising=False) - response: Final = {"choices": [{"text": "cached"}], "usage": {"total_tokens": 3}, "flag": True, "empty": None} - fake_gcs.put( - "bucket", - "cache/sync", - json.dumps({"timestamp": time.time(), "response": json.dumps(response)}).encode(), - ) - fake_gcs.put("bucket", "cache/async", json.dumps({"timestamp": time.time(), "response": response}).encode()) - fake_gcs.put("bucket", "cache/raw", json.dumps(response).encode()) - fake_gcs.put("bucket", "cache/invalid", b"not a cache entry") - binding: Final = _native._CacheTestResolver( - SimpleNamespace( - cache=_native._CacheTestHandle.gcs( - "bucket", - gcs_path="cache", - endpoint=fake_gcs.url, - token=fake_gcs.token, - ) - ) - ).resolve() - - assert binding.lookup(request("sync")) == response - assert await binding.async_lookup(request("async")) == response - assert binding.lookup(request("raw")) == response - assert await binding.async_lookup(request("invalid")) is None - assert binding.lookup(request("missing")) is None - - await binding.async_store({**request("native"), "ttl_seconds": 12.0}, response) - stored: Final = fake_gcs.objects[("bucket", "cache/native")] - stored_value: Final = cast(dict[str, object], json.loads(stored)) - assert stored_value["response"] == response - assert isinstance(stored_value["timestamp"], float) - upload: Final = next(item for item in fake_gcs.requests if item.method == "POST") - assert upload.path == "/upload/storage/v1/b/bucket/o" - assert upload.query == "uploadType=media&name=cache%2Fnative" - assert upload.headers["Authorization"] == f"Bearer {fake_gcs.token}" - assert upload.headers["Content-Type"] == "application/json" - upload_text: Final = f"{upload.path}?{upload.query}{upload.headers}" - assert "ttl" not in upload_text.lower() - assert "expiry" not in upload_text.lower() - download: Final = next(item for item in fake_gcs.requests if item.path.endswith("/cache%2Fsync")) - assert download.path == "/storage/v1/b/bucket/o/cache%2Fsync" - assert download.query == "alt=media" - - binding.store(request("sync2"), response) - assert binding.lookup(request("sync2")) == response - assert GCSCache(bucket_name="bucket", gcs_path="cache").key_prefix == "cache/" - assert GCSCache(bucket_name="bucket", gcs_path="cache/").key_prefix == "cache/" - assert GCSCache(bucket_name="bucket").key_prefix == "" - - -async def test_gcs_batch_lookup_preserves_order_and_treats_malformed_entries_as_misses(fake_gcs: FakeGcs) -> None: - fake_gcs.put("bucket", "cache/hit", json.dumps({"timestamp": time.time(), "response": {"value": 1}}).encode()) - fake_gcs.put("bucket", "cache/invalid", b"not a cache entry") - binding: Final = _native._CacheTestResolver( - SimpleNamespace( - cache=_native._CacheTestHandle.gcs( - "bucket", - gcs_path="cache", - endpoint=fake_gcs.url, - token=fake_gcs.token, - ) - ) - ).resolve() - requests: Final = [request("hit"), request("missing"), request("invalid")] - expected: Final = {"values": [{"value": 1}, None, None], "missing_indices": [1, 2]} - - assert await binding.async_lookup_batch(requests) == expected - assert binding.lookup_batch(requests) == expected - await binding.async_store_batch([request("first"), request("second")], [{"value": 1}, {"value": 2}]) - assert ("bucket", "cache/first") in fake_gcs.objects - assert ("bucket", "cache/second") in fake_gcs.objects - - -async def test_gcs_facade_binds_only_exact_matching_configuration( - fake_gcs: FakeGcs, monkeypatch: pytest.MonkeyPatch -) -> None: - monkeypatch.delenv("GCS_PATH_SERVICE_ACCOUNT", raising=False) - monkeypatch.delenv("GCS_BUCKET_NAME", raising=False) - monkeypatch.setenv("GOOGLE_APPLICATION_CREDENTIALS", "/nonexistent") - facade: Final = Cache(type=LiteLLMCacheType.GCS, gcs_bucket_name="bucket", gcs_path="cache/") - assert type(facade.cache) is GCSCache - - mismatched_bucket: Final = _native._CacheTestHandle.gcs( - "other", - gcs_path="cache", - endpoint=fake_gcs.url, - token=fake_gcs.token, - ) - with pytest.raises(TypeError, match="buckets must match"): - mismatched_bucket._bind_facade(facade) - mismatched_prefix: Final = _native._CacheTestHandle.gcs( - "bucket", - gcs_path="x", - endpoint=fake_gcs.url, - token=fake_gcs.token, - ) - with pytest.raises(TypeError, match="key prefixes must match"): - mismatched_prefix._bind_facade(facade) - mismatched_credentials: Final = _native._CacheTestHandle.gcs( - "bucket", - gcs_path="cache", - path_service_account="sa.json", - endpoint=fake_gcs.url, - token=fake_gcs.token, - ) - with pytest.raises(TypeError, match="credentials must match"): - mismatched_credentials._bind_facade(facade) - with pytest.raises(TypeError, match="types must match"): - _native._CacheTestHandle.memory()._bind_facade(facade) - - matching: Final = _native._CacheTestHandle.gcs( - "bucket", - gcs_path="cache", - endpoint=fake_gcs.url, - token=fake_gcs.token, - ) - matching._bind_facade(facade) - resolver: Final = _native._CacheTestResolver(SimpleNamespace(cache=facade)) - binding: Final = resolver.resolve() - assert binding.kind == "native" - await binding.async_store(request("native"), {"value": "native"}) - assert await binding.async_lookup(request("native")) == {"value": "native"} - assert cast(CacheLookup, facade).get_cache(cache_key="native") is None - - with rebound(facade.cache, "bucket_name", "other"): - assert resolver.resolve().kind == "python_callback" - with rebound(facade.cache, "key_prefix", "x/"): - assert resolver.resolve().kind == "python_callback" - with rebound(facade.cache, "path_service_account", "sa.json"): - assert resolver.resolve().kind == "python_callback" - - def no_get_cache(*args: object, **kwargs: object) -> None: - return None - - with rebound(facade.cache, "get_cache", no_get_cache): - assert resolver.resolve().kind == "python_callback" - with rebound(facade, "ttl", 12): - assert resolver.resolve().kind == "python_callback" - - class CustomGcs(GCSCache): - pass - - with rebound(facade, "cache", CustomGcs(bucket_name="bucket", gcs_path="cache/")): - assert resolver.resolve().kind == "python_callback" - custom_facade: Final = Cache(type=LiteLLMCacheType.GCS, gcs_bucket_name="bucket", gcs_path="cache/") - with rebound(custom_facade, "cache", CustomGcs(bucket_name="bucket", gcs_path="cache/")): - with pytest.raises(TypeError, match="types must match"): - matching._bind_facade(custom_facade) - - missing_bucket: Final = Cache(type=LiteLLMCacheType.GCS) - with pytest.raises(TypeError, match="requires a configured bucket name"): - matching._bind_facade(missing_bucket) - - -async def test_gcs_flush_is_a_no_op_and_ping_is_not_implemented( - fake_gcs: FakeGcs, monkeypatch: pytest.MonkeyPatch -) -> None: - monkeypatch.delenv("GCS_PATH_SERVICE_ACCOUNT", raising=False) - monkeypatch.delenv("GCS_BUCKET_NAME", raising=False) - binding: Final = _native._CacheTestResolver( - SimpleNamespace( - cache=_native._CacheTestHandle.gcs( - "bucket", - gcs_path="cache", - endpoint=fake_gcs.url, - token=fake_gcs.token, - ) - ) - ).resolve() - await binding.async_store(request("key"), {"value": "stored"}) - await binding.async_flush() - assert ("bucket", "cache/key") in fake_gcs.objects - assert await binding.async_lookup(request("key")) == {"value": "stored"} - with pytest.raises(NotImplementedError): - await binding.ping() - - facade: Final = Cache(type=LiteLLMCacheType.GCS, gcs_bucket_name="bucket", gcs_path="cache/") - with pytest.raises(AttributeError): - await facade.ping() - assert cast(CacheLookup, facade.cache).flush_cache() is None - - -async def test_gcs_unauthorized_and_server_errors_surface_as_runtime_errors(fake_gcs: FakeGcs) -> None: - wrong_token: Final = _native._CacheTestResolver( - SimpleNamespace( - cache=_native._CacheTestHandle.gcs( - "bucket", - gcs_path="cache", - endpoint=fake_gcs.url, - token="wrong-token", - ) - ) - ).resolve() - with pytest.raises(RuntimeError): - wrong_token.lookup(request("missing")) - assert not fake_gcs.objects - - binding: Final = _native._CacheTestResolver( - SimpleNamespace( - cache=_native._CacheTestHandle.gcs( - "bucket", - gcs_path="cache", - endpoint=fake_gcs.url, - token=fake_gcs.token, - ) - ) - ).resolve() - with pytest.raises(RuntimeError): - binding.lookup(request("server-error")) - assert binding.lookup(request("missing")) is None - - -async def test_redis_cluster_facade_serves_multi_slot_batches_and_scoped_flush_natively( - cluster_nodes: tuple[tuple[str, int], ...], -) -> None: - startup_nodes: Final = [{"host": host, "port": port} for host, port in cluster_nodes] - url: Final = f"redis://{cluster_nodes[0][0]}:{cluster_nodes[0][1]}" - with rebound(litellm, "default_redis_ttl", 60): - facade: Final = Cache(type=LiteLLMCacheType.REDIS, redis_startup_nodes=startup_nodes, namespace="parity") - assert type(facade.cache) is RedisClusterCache - with pytest.raises(TypeError, match="types must match"): - _native._CacheTestHandle.redis(url, namespace="parity")._bind_facade(facade) - _native._CacheTestHandle.redis(url, namespace="parity", startup_nodes=list(cluster_nodes))._bind_facade(facade) - resolver: Final = _native._CacheTestResolver(SimpleNamespace(cache=facade)) - assert resolver.resolve().kind == "native" - - manager: Final = facade.cache.redis_client.nodes_manager - with rebound(manager, "connection_kwargs", {**manager.connection_kwargs, "db": 1}): - assert resolver.resolve().kind == "python_callback" - with rebound(facade.cache, "redis_kwargs", {**facade.cache.redis_kwargs, "startup_nodes": startup_nodes[:1]}): - assert resolver.resolve().kind == "python_callback" - binding: Final = resolver.resolve() - assert binding.kind == "native" - - client: Final = redis.RedisCluster(startup_nodes=[redis.cluster.ClusterNode(*node) for node in cluster_nodes]) - keys: Final = tuple(f"slot-{index}" for index in range(12)) - slots: Final = {client.keyslot(f"parity:{key}") for key in keys} - assert len(slots) > 1, slots - requests: Final = [request(key) for key in keys] - values: Final = [{"index": index} for index in range(len(keys))] - await binding.async_store_batch(requests, values) - client.set("parity:slot-3", "not a cache entry") - client.set("parity:slot-7", json.dumps({"timestamp": time.time(), "response": {"index": 7, "python": True}})) - - batch: Final = await binding.async_lookup_batch(requests) - assert batch == { - "values": [ - None if index == 3 else {"index": 7, "python": True} if index == 7 else value - for index, value in enumerate(values) - ], - "missing_indices": [3], - } - assert facade.cache.get_cache("parity:slot-0")["response"] == {"index": 0} - assert (await facade.cache.async_get_cache("parity:slot-11"))["response"] == {"index": 11} - assert facade.cache.redis_client.mget_nonatomic([f"parity:{key}" for key in keys[:2]]) == [ - client.get("parity:slot-0"), - client.get("parity:slot-1"), - ] - - await binding.async_store({**request("pinned"), "ttl_seconds": 12.0}, {"pinned": True}) - assert 0 < client.ttl("parity:pinned") <= 12 - client.set("unscoped", "stays") - - await binding.async_flush() - - remaining: Final = tuple( - sorted(key for node in client.get_primaries() for key in client.keys("parity:*", target_nodes=node)) - ) - assert remaining == (), remaining - assert client.get("unscoped") == b"stays" - client.delete("unscoped") - client.close() - facade.cache.redis_client.close() - - -PARAPHRASE_MARKER: Final = " (paraphrase)" -SEMANTIC_EMBEDDING_MODEL: Final = "semantic-test/deterministic" -SEMANTIC_INDEX_PREFIX: Final = "litellm_test_semantic_" -SEMANTIC_CONTEXT: Final = contextvars.ContextVar("semantic_test_context", default="unset") - - -def _normalized(vector: list[float]) -> list[float]: - norm: Final = math.sqrt(sum(component * component for component in vector)) - return [component / norm for component in vector] - - -def _base_embedding(prompt: str) -> list[float]: - digest: Final = hashlib.sha256(prompt.encode("utf-8")).digest() - return _normalized([float(digest[index] + 1) for index in range(8)]) - - -def _semantic_embedding(prompt: str) -> list[float]: - if PARAPHRASE_MARKER not in prompt: - return _base_embedding(prompt) - base: Final = _base_embedding(prompt.replace(PARAPHRASE_MARKER, "").strip()) - pivot: Final = min(range(8), key=lambda index: abs(base[index])) - direction: Final = _normalized( - [(1.0 - base[pivot] * base[pivot]) if index == pivot else -base[index] * base[pivot] for index in range(8)] - ) - # Rotating an orthogonal unit direction by 0.329 produces ~0.05 cosine distance - return _normalized([base[index] + 0.329 * direction[index] for index in range(8)]) - - -class DeterministicEmbedding(litellm.CustomLLM): - def __init__(self) -> None: - self.calls: list[dict[str, object]] = [] - self.async_calls: list[dict[str, object]] = [] - self.entered = asyncio.Event() - self.gate: asyncio.Event | None = None - - def _respond( - self, - model: str, - input: object, - model_response: EmbeddingResponse, - ) -> EmbeddingResponse: - texts: Final = cast(list[object], input if isinstance(input, list) else [input]) - self.calls.append({"model": model, "input": texts}) - model_response.model = model - model_response.data = [ - {"object": "embedding", "index": index, "embedding": _semantic_embedding(str(text))} - for index, text in enumerate(texts) - ] - return model_response - - def embedding( - self, - model: str, - input: list[object], - model_response: EmbeddingResponse, - print_verbose: Callable[..., object], - logging_obj: object, - optional_params: dict[str, object], - api_key: object = None, - api_base: object = None, - timeout: object = None, - litellm_params: object = None, - ) -> EmbeddingResponse: - return self._respond(model, input, model_response) - - async def aembedding( - self, - model: str, - input: list[object], - model_response: EmbeddingResponse, - print_verbose: Callable[..., object], - logging_obj: object, - optional_params: dict[str, object], - api_key: object = None, - api_base: object = None, - timeout: object = None, - litellm_params: object = None, - ) -> EmbeddingResponse: - texts: Final = cast(list[object], input if isinstance(input, list) else [input]) - self.async_calls.append( - { - "model": model, - "input": texts, - "task": asyncio.current_task(), - "context": SEMANTIC_CONTEXT.get(), - } - ) - SEMANTIC_CONTEXT.set("written-in-aembedding") - self.entered.set() - if self.gate is not None: - await self.gate.wait() - return self._respond(model, input, model_response) - - -@pytest.fixture -def semantic_embedding() -> Generator[DeterministicEmbedding]: - handler: Final = DeterministicEmbedding() - with ExitStack() as stack: - stack.enter_context( - rebound( - litellm, - "custom_provider_map", - [ - *litellm.custom_provider_map, - cast( - CustomLLMItem, - {"provider": "semantic-test", "custom_handler": handler}, - ), - ], - ) - ) - stack.enter_context( - rebound( - litellm, - "_custom_providers", # pyright: ignore[reportPrivateUsage] # no public provider-registration hook - [*litellm._custom_providers, "semantic-test"], # pyright: ignore[reportPrivateUsage] # no public provider-registration hook - ) - ) - stack.enter_context(rebound(litellm, "provider_list", [*litellm.provider_list, "semantic-test"])) - yield handler - - -@pytest.fixture -def redis_stack() -> Generator[tuple[str, str]]: - url: Final = os.environ.get("LITELLM_REDIS_STACK_URL") - if url is None: - pytest.skip("LITELLM_REDIS_STACK_URL is not set") - index: Final = f"{SEMANTIC_INDEX_PREFIX}{uuid4().hex}" - yield url, index - client: Final = redis.Redis.from_url(url) - try: - client.execute_command("FT.DROPINDEX", index, "DD") # pyright: ignore[reportUnknownMemberType] # redis-py leaves execute_command partially unknown - except redis.RedisError: - pass - client.close() - - -def semantic_request(key: str, prompt: str, **extra: object) -> dict[str, object]: - return { - "key": {"preset": key}, - "messages": [{"role": "user", "content": prompt}], - **extra, - } - - -def semantic_messages(prompt: str) -> list[dict[str, object]]: - return [{"role": "user", "content": prompt}] - - -def semantic_entry_id(prompt: str, tag: str) -> str: - return hashlib.sha256(f"{prompt}litellm_cache_key{tag}".encode()).hexdigest() - - -def semantic_facade(url: str, index: str, *, similarity_threshold: float = 0.8) -> Cache: - facade: Final = Cache( - type=LiteLLMCacheType.REDIS_SEMANTIC, - redis_url=url, - similarity_threshold=similarity_threshold, - redis_semantic_cache_embedding_model=SEMANTIC_EMBEDDING_MODEL, - redis_semantic_cache_index_name=index, - ) - _CacheTestHandle.redis_semantic(facade.cache)._bind_facade(facade) - return facade - - -def test_redis_semantic_constructor_identity_and_provenance( - redis_stack: tuple[str, str], semantic_embedding: DeterministicEmbedding -) -> None: - url, index = redis_stack - facade: Final = semantic_facade(url, index) - backend: Final = cast(RedisSemanticCache, facade.cache) - assert backend.__class__.__module__ == "litellm.caching.redis_semantic_cache" - assert type(backend) is RedisSemanticCache - assert backend._redis_url == url # pyright: ignore[reportPrivateUsage] # provenance check needs the projected config - assert backend._index_name == index # pyright: ignore[reportPrivateUsage] # provenance check needs the projected config - assert backend.similarity_threshold == 0.8 - assert backend.embedding_model == SEMANTIC_EMBEDDING_MODEL - handle: Final = cast(object, getattr(facade, "_native_cache_handle")) - assert isinstance(handle, _CacheTestHandle) - assert handle.backend == "redis_semantic" - binding: Final = _CacheTestResolver(SimpleNamespace(cache=facade)).resolve() - assert binding.kind == "native" - - -def test_redis_semantic_native_and_python_sync_entries_share_one_layout( - redis_stack: tuple[str, str], semantic_embedding: DeterministicEmbedding -) -> None: - url, index = redis_stack - facade: Final = semantic_facade(url, index) - binding: Final = _CacheTestResolver(SimpleNamespace(cache=facade)).resolve() - client: Final = redis.Redis.from_url(url) - response: Final = {"choices": [{"text": "paris"}], "usage": {"total_tokens": 2}} - - binding.store(semantic_request("geo", "what is the capital of france"), response) - - native_hash_key: Final = f"{index}:{semantic_entry_id('what is the capital of france', 'geo')}" - stored: Final = client.hgetall(native_hash_key) - assert set(stored) == { - b"entry_id", - b"prompt", - b"response", - b"prompt_vector", - b"inserted_at", - b"updated_at", - b"litellm_cache_key", - }, stored - assert stored[b"entry_id"].decode() == native_hash_key.split(":", 1)[1] - assert stored[b"prompt"] == b"what is the capital of france" - assert stored[b"litellm_cache_key"] == b"geo" - assert len(stored[b"prompt_vector"]) == 32 - decoded: Final = cast(dict[str, object], json.loads(stored[b"response"])) - assert decoded["response"] == response - assert ( - cast(RedisSemanticCache, facade.cache).get_cache( # pyright: ignore[reportUnknownMemberType] # **kwargs stays unknown on the backend class - "geo", messages=semantic_messages("what is the capital of france") - ) - == decoded - ) - assert semantic_embedding.calls == [ - {"model": "deterministic", "input": ["what is the capital of france"]}, - {"model": "deterministic", "input": ["what is the capital of france"]}, - {"model": "deterministic", "input": ["dimension test"]}, - ] - - cast(RedisSemanticCache, facade.cache).set_cache( # pyright: ignore[reportUnknownMemberType] # **kwargs stays unknown on the backend class - "math", - json.dumps({"timestamp": 1700000000.0, "response": {"answer": 42}}), - messages=semantic_messages("what is 6 times 7"), - ) - python_hash_key: Final = f"{index}:{semantic_entry_id('what is 6 times 7', 'math')}" - assert json.loads(cast(bytes, client.hget(python_hash_key, "response"))) == { - "timestamp": 1700000000.0, - "response": {"answer": 42}, - } - assert binding.lookup(semantic_request("math", "what is 6 times 7")) == {"answer": 42} - client.close() - - -async def test_redis_semantic_async_paths_and_store_batch_share_one_layout( - redis_stack: tuple[str, str], semantic_embedding: DeterministicEmbedding -) -> None: - url, index = redis_stack - facade: Final = semantic_facade(url, index) - binding: Final = _CacheTestResolver(SimpleNamespace(cache=facade)).resolve() - client: Final = redis.Redis.from_url(url) - - await binding.async_store(semantic_request("async", "name a primary color"), {"answer": "blue"}) - hash_key: Final = f"{index}:{semantic_entry_id('name a primary color', 'async')}" - decoded: Final = cast(dict[str, object], json.loads(cast(bytes, client.hget(hash_key, "response")))) - python_read: Final = await cast(RedisSemanticCache, facade.cache).async_get_cache( # pyright: ignore[reportUnknownMemberType] # **kwargs stays unknown on the backend class - "async", messages=semantic_messages("name a primary color") - ) - assert python_read == decoded - - await binding.async_store_batch( - [ - semantic_request("batch-one", "first batch prompt"), - semantic_request("batch-two", "second batch prompt"), - ], - [{"answer": 1}, {"answer": 2}], - ) - expected: Final = { - key: json.loads(cast(bytes, client.hget(f"{index}:{semantic_entry_id(prompt, key)}", "response"))) - for key, prompt in ( - ("batch-one", "first batch prompt"), - ("batch-two", "second batch prompt"), - ) - } - for key, prompt in ( - ("batch-one", "first batch prompt"), - ("batch-two", "second batch prompt"), - ): - assert ( - cast(RedisSemanticCache, facade.cache).get_cache( # pyright: ignore[reportUnknownMemberType] # **kwargs stays unknown on the backend class - key, messages=semantic_messages(prompt) - ) - == expected[key] - ), key - - cast(RedisSemanticCache, facade.cache).set_cache( # pyright: ignore[reportUnknownMemberType] # **kwargs stays unknown on the backend class - "async-python", - json.dumps({"timestamp": 1700000000.0, "response": {"answer": "python"}}), - messages=semantic_messages("python written prompt"), - ) - assert await binding.async_lookup(semantic_request("async-python", "python written prompt")) == {"answer": "python"} - client.close() - - -async def test_native_semantic_async_embedding_runs_inline_in_the_callers_task( - redis_stack: tuple[str, str], semantic_embedding: DeterministicEmbedding -) -> None: - url, index = redis_stack - facade: Final = semantic_facade(url, index) - binding: Final = _CacheTestResolver(SimpleNamespace(cache=facade)).resolve() - assert binding.kind == "native" - caller: Final = asyncio.current_task() - SEMANTIC_CONTEXT.set("caller-sentinel") - response: Final = {"choices": [{"text": "paris"}]} - - await binding.async_store(semantic_request("inline", "what is the capital of france"), response) - assert ( - await binding.async_lookup(semantic_request("inline", f"what is the capital of france{PARAPHRASE_MARKER}")) - == response - ) - assert await binding.async_lookup(semantic_request("inline", "python written prompt")) is None - assert SEMANTIC_CONTEXT.get() == "written-in-aembedding" - assert semantic_embedding.async_calls == [ - { - "model": "deterministic", - "input": ["what is the capital of france"], - "task": caller, - "context": "caller-sentinel", - }, - { - "model": "deterministic", - "input": [f"what is the capital of france{PARAPHRASE_MARKER}"], - "task": caller, - "context": "written-in-aembedding", - }, - { - "model": "deterministic", - "input": ["python written prompt"], - "task": caller, - "context": "written-in-aembedding", - }, - ], semantic_embedding.async_calls - - -async def test_native_semantic_cancellation_during_embedding_skips_the_backend( - redis_stack: tuple[str, str], semantic_embedding: DeterministicEmbedding -) -> None: - url, index = redis_stack - facade: Final = semantic_facade(url, index) - binding: Final = _CacheTestResolver(SimpleNamespace(cache=facade)).resolve() - assert binding.kind == "native" - semantic_embedding.gate = asyncio.Event() - - async def lookup() -> object: - return await binding.async_lookup(semantic_request("cancel", "cancelled prompt")) - - task: Final = asyncio.create_task(lookup()) - await semantic_embedding.entered.wait() - task.cancel() - with pytest.raises(asyncio.CancelledError): - await task - semantic_embedding.gate.set() - - assert len(semantic_embedding.async_calls) == 1 - assert ( - await cast(RedisSemanticCache, facade.cache).async_get_cache( # pyright: ignore[reportUnknownMemberType] # **kwargs stays unknown on the backend class - "cancel", messages=semantic_messages("cancelled prompt") - ) - is None - ) - - -def test_redis_semantic_similarity_tag_and_threshold_boundaries( - redis_stack: tuple[str, str], semantic_embedding: DeterministicEmbedding -) -> None: - url, index = redis_stack - facade: Final = semantic_facade(url, index) - binding: Final = _CacheTestResolver(SimpleNamespace(cache=facade)).resolve() - - binding.store(semantic_request("sim", "tell me a joke"), {"answer": "haha"}) - paraphrase: Final = f"tell me a joke{PARAPHRASE_MARKER}" - assert binding.lookup(semantic_request("sim", paraphrase)) == {"answer": "haha"} - assert binding.lookup(semantic_request("sim", "an unrelated question about spreadsheets")) is None - assert binding.lookup(semantic_request("other-key", "tell me a joke")) is None - - strict: Final = semantic_facade(url, index, similarity_threshold=0.99) - strict_binding: Final = _CacheTestResolver(SimpleNamespace(cache=strict)).resolve() - assert strict_binding.lookup(semantic_request("sim", paraphrase)) is None - assert strict_binding.lookup(semantic_request("sim", "tell me a joke")) == {"answer": "haha"} - - -def test_redis_semantic_ttl_is_written_only_when_requested( - redis_stack: tuple[str, str], semantic_embedding: DeterministicEmbedding -) -> None: - url, index = redis_stack - facade: Final = semantic_facade(url, index) - binding: Final = _CacheTestResolver(SimpleNamespace(cache=facade)).resolve() - client: Final = redis.Redis.from_url(url) - - binding.store({**semantic_request("ttl", "ttl prompt"), "ttl_seconds": 12.0}, {"answer": 1}) - expiring: Final = f"{index}:{semantic_entry_id('ttl prompt', 'ttl')}" - assert 0 < client.ttl(expiring) <= 12 - - binding.store(semantic_request("ttl-none", "untimed prompt"), {"answer": 2}) - persistent: Final = f"{index}:{semantic_entry_id('untimed prompt', 'ttl-none')}" - assert client.ttl(persistent) == -1 - - binding.store( - {**semantic_request("ttl-fraction", "fractional prompt"), "ttl_seconds": 1.5}, - {"answer": 3}, - ) - fractional: Final = f"{index}:{semantic_entry_id('fractional prompt', 'ttl-fraction')}" - assert client.ttl(fractional) == 2 - client.close() - - -def test_redis_semantic_malformed_response_is_a_miss_for_both_readers( - redis_stack: tuple[str, str], semantic_embedding: DeterministicEmbedding -) -> None: - url, index = redis_stack - facade: Final = semantic_facade(url, index) - binding: Final = _CacheTestResolver(SimpleNamespace(cache=facade)).resolve() - client: Final = redis.Redis.from_url(url) - - binding.store(semantic_request("bad", "corrupt me"), {"answer": 1}) - hash_key: Final = f"{index}:{semantic_entry_id('corrupt me', 'bad')}" - client.hset(hash_key, "response", b"{not json") - assert binding.lookup(semantic_request("bad", "corrupt me")) is None - assert ( - cast(RedisSemanticCache, facade.cache).get_cache( # pyright: ignore[reportUnknownMemberType] # **kwargs stays unknown on the backend class - "bad", messages=semantic_messages("corrupt me") - ) - is None - ) - client.close() - - -async def test_redis_semantic_unsupported_operations_raise_not_implemented( - redis_stack: tuple[str, str], semantic_embedding: DeterministicEmbedding -) -> None: - url, index = redis_stack - facade: Final = semantic_facade(url, index) - binding: Final = _CacheTestResolver(SimpleNamespace(cache=facade)).resolve() - - with pytest.raises(NotImplementedError): - binding.lookup_batch([semantic_request("batch", "prompt one")]) - with pytest.raises(NotImplementedError): - await binding.async_lookup_batch([semantic_request("batch", "prompt one")]) - with pytest.raises(NotImplementedError): - await binding.async_flush() - with pytest.raises(NotImplementedError): - await binding.ping() - - -def test_redis_semantic_requests_without_prompt_are_noops( - redis_stack: tuple[str, str], semantic_embedding: DeterministicEmbedding -) -> None: - url, index = redis_stack - facade: Final = semantic_facade(url, index) - binding: Final = _CacheTestResolver(SimpleNamespace(cache=facade)).resolve() - client: Final = redis.Redis.from_url(url) - - binding.store(request("plain"), {"answer": 1}) - assert binding.lookup(request("plain")) is None - assert semantic_embedding.calls == [] - assert client.keys(f"{index}:*") == [] - client.close() - - -def test_redis_semantic_scope_overrides_the_tag_and_isolates_entries( - redis_stack: tuple[str, str], semantic_embedding: DeterministicEmbedding -) -> None: - url, index = redis_stack - facade: Final = semantic_facade(url, index) - binding: Final = _CacheTestResolver(SimpleNamespace(cache=facade)).resolve() - client: Final = redis.Redis.from_url(url) - - scoped: Final = {**semantic_request("scoped", "scoped prompt"), "scope": "team-a"} - binding.store(scoped, {"answer": "kept"}) - hash_key: Final = f"{index}:{semantic_entry_id('scoped prompt', 'team-a')}" - assert client.hget(hash_key, "litellm_cache_key") == b"team-a" - assert binding.lookup(scoped) == {"answer": "kept"} - assert binding.lookup(semantic_request("scoped", "scoped prompt")) is None - assert binding.lookup({**scoped, "scope": "team-b"}) is None - client.close() - - -def test_redis_semantic_configuration_drift_falls_back_to_python( - redis_stack: tuple[str, str], - semantic_embedding: DeterministicEmbedding, - monkeypatch: pytest.MonkeyPatch, -) -> None: - url, index = redis_stack - facade: Final = semantic_facade(url, index) - resolver: Final = _CacheTestResolver(SimpleNamespace(cache=facade)) - assert resolver.resolve().kind == "native" - - with rebound(facade.cache, "similarity_threshold", 0.5): - assert resolver.resolve().kind == "python_callback" - with rebound(facade, "semantic_cache_scope", "end_user"): - assert resolver.resolve().kind == "python_callback" - with rebound(facade.cache, "embedding_model", "other-model"): - assert resolver.resolve().kind == "python_callback" - with rebound(facade.cache, "_index_name", "other-index"): - assert resolver.resolve().kind == "python_callback" - with rebound(facade.cache, "CACHE_KEY_FIELD_NAME", "other-field"): - assert resolver.resolve().kind == "python_callback" - - def patched_embedding(self: object, prompt: str, metadata: object = None) -> list[float]: - return _semantic_embedding(prompt) - - monkeypatch.setattr(RedisSemanticCache, "_get_embedding", patched_embedding) - assert resolver.resolve().kind == "python_callback" - - -def test_redis_semantic_handle_rejects_wrong_backends( - redis_stack: tuple[str, str], semantic_embedding: DeterministicEmbedding -) -> None: - url, index = redis_stack - - class CustomSemanticCache(RedisSemanticCache): - pass - - with pytest.raises(TypeError, match="built-in RedisSemanticCache"): - _CacheTestHandle.redis_semantic(object()) - with pytest.raises(TypeError, match="built-in RedisSemanticCache"): - _CacheTestHandle.redis_semantic( - CustomSemanticCache( - redis_url=url, - similarity_threshold=0.8, - embedding_model=SEMANTIC_EMBEDDING_MODEL, - index_name=f"{index}_subclass", - ) - ) - - facade: Final = semantic_facade(url, index) - with pytest.raises(TypeError, match="backend types must match"): - _CacheTestHandle.redis(url)._bind_facade(facade) - - subclassed_facade: Final = Cache( - type=LiteLLMCacheType.REDIS_SEMANTIC, - redis_url=url, - similarity_threshold=0.8, - redis_semantic_cache_embedding_model=SEMANTIC_EMBEDDING_MODEL, - redis_semantic_cache_index_name=index, - ) - subclassed_facade.cache = CustomSemanticCache( # pyright: ignore[reportAttributeAccessIssue] # facade backend slot is not declared - redis_url=url, - similarity_threshold=0.8, - embedding_model=SEMANTIC_EMBEDDING_MODEL, - index_name=index, - ) - with pytest.raises(TypeError): - _CacheTestHandle.redis_semantic(subclassed_facade.cache)._bind_facade(subclassed_facade) - - replacement_facade: Final = Cache( - type=LiteLLMCacheType.REDIS_SEMANTIC, - redis_url=url, - similarity_threshold=0.8, - redis_semantic_cache_embedding_model=SEMANTIC_EMBEDDING_MODEL, - redis_semantic_cache_index_name=index, - ) - with pytest.raises(TypeError, match="must be the native embedder"): - _CacheTestHandle.redis_semantic(facade.cache)._bind_facade(replacement_facade) - - -def qdrant_facade(qdrant_url: str, collection_name: str) -> Cache: - return Cache( - type=LiteLLMCacheType.QDRANT_SEMANTIC, - qdrant_api_base=qdrant_url, - qdrant_collection_name=collection_name, - similarity_threshold=0.99, - qdrant_semantic_cache_embedding_model="text-embedding-3-small", - qdrant_semantic_cache_vector_size=8, - ) - - -def test_qdrant_semantic_facade_binds_native_and_shares_entries(qdrant_url: str, fake_embedding_endpoint: str) -> None: - del fake_embedding_endpoint - messages: Final = [{"role": "user", "content": "shared prompt"}] - collection: Final = f"cache_{uuid4().hex}" - facade: Final = qdrant_facade(qdrant_url, collection) - facade.cache.set_cache( - "python-key", - {"timestamp": time.time(), "response": json.dumps({"id": "py"})}, - messages=messages, - ) - handle: Final = _native._CacheTestHandle.qdrant_semantic( - qdrant_url, - collection_name=collection, - similarity_threshold=0.99, - vector_size=8, - ) - handle._bind_facade(facade) - binding: Final = _native._CacheTestResolver(SimpleNamespace(cache=facade)).resolve() - assert binding.kind == "native" - assert binding.lookup(qdrant_request("python-key", messages)) == {"id": "py"} - binding.store(qdrant_request("native-key", messages), {"id": "native"}) - python_value: Final = facade.cache.get_cache("native-key", messages=messages) - assert isinstance(python_value, dict) - assert python_value["response"] == {"id": "native"} - unrelated: Final = [{"role": "user", "content": "unrelated prompt"}] - assert binding.lookup(qdrant_request("native-key", unrelated)) is None - assert facade.cache.get_cache("native-key", messages=unrelated) is None - assert binding.lookup(qdrant_request("different-key", messages)) is None - assert facade.cache.get_cache("different-key", messages=messages) is None - - -async def test_qdrant_semantic_async_parity(qdrant_url: str, fake_embedding_endpoint: str) -> None: - del fake_embedding_endpoint - messages: Final = [{"role": "user", "content": "async prompt"}] - collection: Final = f"cache_{uuid4().hex}" - facade: Final = qdrant_facade(qdrant_url, collection) - handle: Final = _native._CacheTestHandle.qdrant_semantic( - qdrant_url, - collection_name=collection, - similarity_threshold=0.99, - vector_size=8, - ) - handle._bind_facade(facade) - binding: Final = _native._CacheTestResolver(SimpleNamespace(cache=facade)).resolve() - await facade.cache.async_set_cache( - "python-key", - {"timestamp": time.time(), "response": json.dumps({"id": "py"})}, - messages=messages, - ) - assert await binding.async_lookup(qdrant_request("python-key", messages)) == {"id": "py"} - await binding.async_store(qdrant_request("native-key", messages), {"id": "native"}) - python_value: Final = await facade.cache.async_get_cache("native-key", messages=messages) - assert isinstance(python_value, dict) - assert python_value["response"] == {"id": "native"} - - -async def test_qdrant_semantic_async_store_batch_shares_entries(qdrant_url: str, fake_embedding_endpoint: str) -> None: - del fake_embedding_endpoint - collection: Final = f"cache_{uuid4().hex}" - facade: Final = qdrant_facade(qdrant_url, collection) - handle: Final = _native._CacheTestHandle.qdrant_semantic( - qdrant_url, - collection_name=collection, - similarity_threshold=0.99, - vector_size=8, - ) - handle._bind_facade(facade) - binding: Final = _native._CacheTestResolver(SimpleNamespace(cache=facade)).resolve() - entries: Final = [ - qdrant_request("batch-one", [{"role": "user", "content": "first batch prompt"}]), - qdrant_request("batch-two", [{"role": "user", "content": "second batch prompt"}]), - ] - await binding.async_store_batch(entries, [{"id": "one"}, {"id": "two"}]) - - assert binding.lookup(entries[0]) == {"id": "one"} - assert binding.lookup(entries[1]) == {"id": "two"} - assert (await facade.cache.async_get_cache("batch-one", messages=entries[0]["messages"]))["response"] == { - "id": "one" - } - assert (await facade.cache.async_get_cache("batch-two", messages=entries[1]["messages"]))["response"] == { - "id": "two" - } - - -async def test_qdrant_semantic_malformed_entries_and_unsupported_operations( - qdrant_url: str, fake_embedding_endpoint: str -) -> None: - del fake_embedding_endpoint - messages: Final = [{"role": "user", "content": "malformed prompt"}] - collection: Final = f"cache_{uuid4().hex}" - facade: Final = qdrant_facade(qdrant_url, collection) - handle: Final = _native._CacheTestHandle.qdrant_semantic( - qdrant_url, - collection_name=collection, - similarity_threshold=0.99, - vector_size=8, - ) - handle._bind_facade(facade) - binding: Final = _native._CacheTestResolver(SimpleNamespace(cache=facade)).resolve() - key: Final = "malformed-key" - response: Final = { - "points": [ - { - "id": str(uuid4()), - "vector": embedding_vector("malformed prompt"), - "payload": { - "litellm_cache_key": key, - "text": "malformed prompt", - "response": "not json", - }, - } - ] - } - facade.cache.sync_client.put( - url=f"{qdrant_url}/collections/{collection}/points", - headers=facade.cache.headers, - json=response, - ) - assert binding.lookup(qdrant_request(key, messages)) is None - with pytest.raises(RuntimeError, match="operation is not supported"): - binding.lookup_batch([qdrant_request(key, messages)]) - with pytest.raises(RuntimeError, match="operation is not supported"): - await binding.async_flush() - with pytest.raises(RuntimeError, match="operation is not supported"): - await binding.ping() - - -def test_qdrant_semantic_ignores_request_expiry(qdrant_url: str, fake_embedding_endpoint: str) -> None: - del fake_embedding_endpoint - messages: Final = [{"role": "user", "content": "persistent prompt"}] - collection: Final = f"cache_{uuid4().hex}" - facade: Final = qdrant_facade(qdrant_url, collection) - handle: Final = _native._CacheTestHandle.qdrant_semantic( - qdrant_url, - collection_name=collection, - similarity_threshold=0.99, - vector_size=8, - ) - handle._bind_facade(facade) - binding: Final = _native._CacheTestResolver(SimpleNamespace(cache=facade)).resolve() - binding.store(qdrant_request("persistent-key", messages, ttl_seconds=1.0), {"id": "persistent"}) - time.sleep(1.2) - assert binding.lookup(qdrant_request("persistent-key", messages)) == {"id": "persistent"} - python_value: Final = facade.cache.get_cache("persistent-key", messages=messages) - assert isinstance(python_value, dict) - assert python_value["response"] == {"id": "persistent"} - - -def test_qdrant_semantic_mutation_and_projection_fallback(qdrant_url: str, fake_embedding_endpoint: str) -> None: - del fake_embedding_endpoint - collection: Final = f"cache_{uuid4().hex}" - facade: Final = qdrant_facade(qdrant_url, collection) - handle: Final = _native._CacheTestHandle.qdrant_semantic( - qdrant_url, - collection_name=collection, - similarity_threshold=0.99, - vector_size=8, - ) - handle._bind_facade(facade) - facade.cache.qdrant_api_key = "rotated" - assert _native._CacheTestResolver(SimpleNamespace(cache=facade)).resolve().kind == "python_callback" - facade.cache.similarity_threshold = 0.5 - assert _native._CacheTestResolver(SimpleNamespace(cache=facade)).resolve().kind == "python_callback" - unsupported: Final = qdrant_facade(qdrant_url, f"cache_{uuid4().hex}") - unsupported.cache.embedding_max_input_tokens = 100 - with pytest.raises(TypeError, match="requires Python"): - handle._bind_facade(unsupported) - unsupported.cache.embedding_max_input_tokens = None - unsupported.cache.qdrant_api_base = "http://127.0.0.1:7777" - with pytest.raises(TypeError, match="gRPC"): - handle._bind_facade(unsupported) - - -CacheFactory: TypeAlias = Callable[[], Cache] - - -def require_rust(monkeypatch: pytest.MonkeyPatch, backend: LiteLLMCacheType) -> None: - monkeypatch.setattr(catalog, "RULES", (CacheRule(Rollout.RUST_REQUIRED, backends=frozenset({backend})),)) - - -def native_runtime(facade: Cache) -> ResponseCacheRuntime: - runtime: Final = facade._native_cache # pyright: ignore[reportPrivateUsage] # the activation under test has no public accessor - assert isinstance(runtime, ResponseCacheRuntime) - assert runtime.kind == "native" - return runtime - - -@pytest.fixture -def cache_factory(request: pytest.FixtureRequest, tmp_path: Path) -> CacheFactory: - backend: Final = cast(LiteLLMCacheType, request.param) - match backend: - case LiteLLMCacheType.LOCAL: - return lambda: Cache(type=backend) - case LiteLLMCacheType.DISK: - return lambda: Cache(type=backend, disk_cache_dir=str(tmp_path)) - case LiteLLMCacheType.REDIS: - parsed: Final = urlparse(cast(str, request.getfixturevalue("redis_url"))) - return lambda: Cache(type=backend, host=parsed.hostname, port=str(parsed.port)) - case LiteLLMCacheType.S3: - stub: Final = cast(S3Stub, request.getfixturevalue("s3_stub")) - return lambda: Cache( - type=backend, - s3_bucket_name="cache-bucket", - s3_region_name="us-east-1", - s3_endpoint_url=stub.url, - s3_aws_access_key_id="key", - s3_aws_secret_access_key="secret", - s3_path="team", - ) - case LiteLLMCacheType.GCS: - return lambda: Cache(type=backend, gcs_bucket_name="bucket", gcs_path="cache/") - case LiteLLMCacheType.REDIS_SEMANTIC: - return lambda: Cache( - type=backend, - redis_url="redis://127.0.0.1:6379", - similarity_threshold=0.8, - redis_semantic_cache_embedding_model="text-embedding-3-small", - ) - case LiteLLMCacheType.VALKEY_SEMANTIC: - return lambda: Cache(type=backend, redis_url="redis://127.0.0.1:6390/0", similarity_threshold=0.8) - case _: - raise AssertionError(f"no local factory for {backend}") - - -ROUND_TRIP_BACKENDS: Final = ( - LiteLLMCacheType.LOCAL, - LiteLLMCacheType.DISK, - LiteLLMCacheType.REDIS, - LiteLLMCacheType.S3, -) -SHARED_STORE_BACKENDS: Final = (LiteLLMCacheType.DISK, LiteLLMCacheType.REDIS, LiteLLMCacheType.S3) - - -def completion_kwargs(label: str) -> dict[str, object]: - return {"model": "gpt-4o", "messages": [{"role": "user", "content": f"{label} {uuid4().hex}"}]} - - -@pytest.mark.parametrize("backend", list(LiteLLMCacheType)) -def test_shipped_rules_keep_every_backend_on_python(backend: LiteLLMCacheType) -> None: - assert resolve_response_cache(cast(Cache, SimpleNamespace(type=backend))) is None - - -@pytest.mark.parametrize( - "cache_factory", - [ - LiteLLMCacheType.LOCAL, - LiteLLMCacheType.DISK, - LiteLLMCacheType.REDIS, - LiteLLMCacheType.S3, - LiteLLMCacheType.GCS, - LiteLLMCacheType.REDIS_SEMANTIC, - LiteLLMCacheType.VALKEY_SEMANTIC, - ], - indirect=True, -) -def test_shipped_rules_construct_python_backed_facades(cache_factory: CacheFactory) -> None: - assert cache_factory()._native_cache is None # pyright: ignore[reportPrivateUsage] # the activation under test has no public accessor - - -@pytest.mark.parametrize( - "cache_factory", - [ - LiteLLMCacheType.LOCAL, - LiteLLMCacheType.DISK, - LiteLLMCacheType.REDIS, - LiteLLMCacheType.S3, - LiteLLMCacheType.GCS, - LiteLLMCacheType.REDIS_SEMANTIC, - LiteLLMCacheType.VALKEY_SEMANTIC, - ], - indirect=True, -) -def test_rust_required_rule_activates_the_native_backend( - cache_factory: CacheFactory, monkeypatch: pytest.MonkeyPatch, request: pytest.FixtureRequest -) -> None: - require_rust(monkeypatch, cast(LiteLLMCacheType, request.node.callspec.params["cache_factory"])) - native_runtime(cache_factory()) - - -@pytest.mark.parametrize("cache_factory", ROUND_TRIP_BACKENDS, indirect=True) -async def test_facade_storage_calls_round_trip_through_the_native_backend( - cache_factory: CacheFactory, monkeypatch: pytest.MonkeyPatch, request: pytest.FixtureRequest -) -> None: - require_rust(monkeypatch, cast(LiteLLMCacheType, request.node.callspec.params["cache_factory"])) - facade: Final = cache_factory() - native_runtime(facade) - - sync_kwargs: Final = completion_kwargs("sync") - facade.add_cache({"answer": 1}, **sync_kwargs) - assert facade.get_cache(**sync_kwargs) == {"answer": 1} - - async_kwargs: Final = completion_kwargs("async") - await facade.async_add_cache({"answer": 2}, **async_kwargs) - assert await facade.async_get_cache(**async_kwargs) == {"answer": 2} - assert facade.get_cache(**completion_kwargs("absent")) is None - - -async def test_memory_facade_writes_bypass_the_python_backend(monkeypatch: pytest.MonkeyPatch) -> None: - require_rust(monkeypatch, LiteLLMCacheType.LOCAL) - facade: Final = Cache(type=LiteLLMCacheType.LOCAL) - native_runtime(facade) - kwargs: Final = completion_kwargs("memory") - facade.add_cache({"answer": 1}, **kwargs) - assert facade.cache.get_cache(facade.get_cache_key(**kwargs)) is None - assert facade.get_cache(**kwargs) == {"answer": 1} - - -@pytest.mark.parametrize("cache_factory", SHARED_STORE_BACKENDS, indirect=True) -async def test_native_and_python_facades_share_one_wire_format( - cache_factory: CacheFactory, monkeypatch: pytest.MonkeyPatch, request: pytest.FixtureRequest -) -> None: - python_facade: Final = cache_factory() - assert python_facade._native_cache is None # pyright: ignore[reportPrivateUsage] # the activation under test has no public accessor - require_rust(monkeypatch, cast(LiteLLMCacheType, request.node.callspec.params["cache_factory"])) - native_facade: Final = cache_factory() - native_runtime(native_facade) - - native_written: Final = completion_kwargs("native") - native_facade.add_cache({"writer": "native"}, **native_written) - assert python_facade.get_cache(**native_written) == {"writer": "native"} - - python_written: Final = completion_kwargs("python") - python_facade.add_cache({"writer": "python"}, **python_written) - assert native_facade.get_cache(**python_written) == {"writer": "python"} - - async_native: Final = completion_kwargs("async-native") - await native_facade.async_add_cache({"writer": "async-native"}, **async_native) - assert await python_facade.async_get_cache(**async_native) == {"writer": "async-native"} - - async_python: Final = completion_kwargs("async-python") - await python_facade.async_add_cache({"writer": "async-python"}, **async_python) - assert await native_facade.async_get_cache(**async_python) == {"writer": "async-python"} - - -@pytest.mark.parametrize("cache_factory", ROUND_TRIP_BACKENDS, indirect=True) -async def test_embedding_pipeline_stores_one_native_entry_per_input( - cache_factory: CacheFactory, monkeypatch: pytest.MonkeyPatch, request: pytest.FixtureRequest -) -> None: - require_rust(monkeypatch, cast(LiteLLMCacheType, request.node.callspec.params["cache_factory"])) - facade: Final = cache_factory() - native_runtime(facade) - inputs: Final = [f"alpha {uuid4().hex}", f"beta {uuid4().hex}"] - result: Final = EmbeddingResponse( - model="text-embedding-3-small", - data=[ - {"object": "embedding", "index": 0, "embedding": [0.1, 0.2]}, - {"object": "embedding", "index": 1, "embedding": [0.3, 0.4]}, - ], - ) - await facade.async_add_cache_pipeline(result, model="text-embedding-3-small", input=inputs) - - keys: Final = [facade.get_cache_key(model="text-embedding-3-small", input=text) for text in inputs] - assert len(set(keys)) == len(inputs) - for text, expected in zip(inputs, ([0.1, 0.2], [0.3, 0.4]), strict=True): - cached = await facade.async_get_cache(model="text-embedding-3-small", input=text) - assert isinstance(cached, dict) - assert cached["embedding"] == expected - assert await facade.async_get_cache(model="text-embedding-3-small", input=inputs) is None - - -def redis_facade(redis_url: str, **settings: object) -> Cache: - parsed: Final = urlparse(redis_url) - return Cache(type=LiteLLMCacheType.REDIS, host=parsed.hostname, port=str(parsed.port), **settings) - - -@pytest.mark.parametrize( - ("settings", "message"), - [ - pytest.param({"max_connections": 10}, "max_connections requires Python", id="pool-size"), - pytest.param({"socket_timeout": 1.0}, "socket_timeout and socket_connect_timeout", id="socket-timeout"), - pytest.param( - {"socket_connect_timeout": 1.0}, "socket_timeout and socket_connect_timeout", id="connect-timeout" - ), - pytest.param({"socket_keepalive": True}, "does not support socket_keepalive", id="keepalive"), - pytest.param({"health_check_interval": 5}, "does not support health_check_interval", id="health-check"), - pytest.param({"client_name": "litellm"}, "does not support client_name", id="client-name"), - pytest.param({"ssl": True}, "ssl_check_hostname=false require Python", id="tls-default-hostname-check"), - pytest.param({"ssl": True, "ssl_cert_reqs": "none"}, "ssl_cert_reqs=none", id="tls-without-verification"), - pytest.param( - {"ssl": True, "ssl_check_hostname": True, "ssl_ca_certs": "/ca.pem"}, - "does not support ssl_ca_certs", - id="tls-custom-ca", - ), - pytest.param( - {"ssl": True, "ssl_check_hostname": True, "ssl_certfile": "/client.pem", "ssl_keyfile": "/client.key"}, - "does not support ssl_ca_certs, ssl_ca_data, ssl_certfile or ssl_keyfile", - id="tls-client-certificate", - ), - ], -) -def test_redis_settings_the_native_client_cannot_honor_decline( - redis_url: str, monkeypatch: pytest.MonkeyPatch, settings: dict[str, object], message: str -) -> None: - require_rust(monkeypatch, LiteLLMCacheType.REDIS) - with pytest.raises(RuntimeError, match=f"declined the cache: native Redis.*{message}"): - redis_facade(redis_url, **settings) - - -def test_redis_verified_tls_activates_natively(redis_url: str, monkeypatch: pytest.MonkeyPatch) -> None: - require_rust(monkeypatch, LiteLLMCacheType.REDIS) - native_runtime(redis_facade(redis_url, ssl=True, ssl_check_hostname=True)) - - -async def test_redis_flush_size_buffers_native_facade_writes(redis_url: str, monkeypatch: pytest.MonkeyPatch) -> None: - require_rust(monkeypatch, LiteLLMCacheType.REDIS) - facade: Final = redis_facade(redis_url, redis_flush_size=2, namespace="team") - native_runtime(facade) - client: Final = redis.Redis.from_url(redis_url) - first: Final = completion_kwargs("first") - await facade.async_add_cache({"value": 1}, **first) - first_key: Final = facade.get_cache_key(**first) - assert first_key.startswith("team:") - assert client.get(first_key) is None - second: Final = completion_kwargs("second") - await facade.async_add_cache({"value": 2}, **second) - assert client.get(first_key) is not None - assert client.get(facade.get_cache_key(**second)) is not None - client.close() - - -@pytest.mark.parametrize( - ("backend", "settings", "message"), - [ - pytest.param( - LiteLLMCacheType.VALKEY_SEMANTIC, - {"redis_url": "rediss://127.0.0.1:6390/0", "similarity_threshold": 0.8}, - "native Valkey semantic cache does not support TLS connections", - id="valkey-tls", - ), - pytest.param( - LiteLLMCacheType.VALKEY_SEMANTIC, - {"redis_url": "redis://127.0.0.1:6390/0?socket_timeout=1", "similarity_threshold": 0.8}, - "native Redis uses fixed socket timeouts; socket_timeout and socket_connect_timeout require Python", - id="valkey-socket-timeout", - ), - pytest.param( - LiteLLMCacheType.REDIS_SEMANTIC, - {"redis_url": "rediss://127.0.0.1:6380", "similarity_threshold": 0.8}, - "native Redis semantic cache does not support TLS or query options in redis_url", - id="redis-semantic-tls", - ), - pytest.param( - LiteLLMCacheType.REDIS_SEMANTIC, - {"redis_url": "redis://127.0.0.1:6379?socket_timeout=1", "similarity_threshold": 0.8}, - "native Redis semantic cache does not support TLS or query options in redis_url", - id="redis-semantic-query", - ), - ], -) -def test_semantic_settings_the_native_client_cannot_honor_decline( - monkeypatch: pytest.MonkeyPatch, backend: LiteLLMCacheType, settings: dict[str, object], message: str -) -> None: - require_rust(monkeypatch, backend) - with pytest.raises(RuntimeError, match=f"declined the cache: {message}"): - Cache(type=backend, **settings) - - -def test_rust_with_fallback_keeps_python_when_the_native_client_declines( - redis_url: str, monkeypatch: pytest.MonkeyPatch -) -> None: - monkeypatch.setattr( - catalog, - "RULES", - (CacheRule(Rollout.RUST_OPT_OUT, backends=frozenset({LiteLLMCacheType.REDIS})),), - ) - assert redis_facade(redis_url, socket_timeout=1.0)._native_cache is None # pyright: ignore[reportPrivateUsage] # the activation under test has no public accessor - - -def test_qdrant_semantic_rust_required_rule_activates_natively( - qdrant_url: str, fake_embedding_endpoint: str, monkeypatch: pytest.MonkeyPatch -) -> None: - del fake_embedding_endpoint - require_rust(monkeypatch, LiteLLMCacheType.QDRANT_SEMANTIC) - facade: Final = qdrant_facade(qdrant_url, f"cache_{uuid4().hex}") - native_runtime(facade) - kwargs: Final = {"model": "gpt-4o", "messages": [{"role": "user", "content": "qdrant activation"}]} - facade.add_cache({"answer": "qdrant"}, **kwargs) - assert facade.get_cache(**kwargs) == {"answer": "qdrant"} - - -async def test_redis_semantic_rust_required_rule_activates_natively( - redis_stack: tuple[str, str], semantic_embedding: DeterministicEmbedding, monkeypatch: pytest.MonkeyPatch -) -> None: - del semantic_embedding - url, index = redis_stack - require_rust(monkeypatch, LiteLLMCacheType.REDIS_SEMANTIC) - facade: Final = Cache( - type=LiteLLMCacheType.REDIS_SEMANTIC, - redis_url=url, - similarity_threshold=0.8, - redis_semantic_cache_embedding_model=SEMANTIC_EMBEDDING_MODEL, - redis_semantic_cache_index_name=index, - ) - native_runtime(facade) - kwargs: Final = {"model": "gpt-4o", "messages": semantic_messages("name a primary color")} - await facade.async_add_cache({"answer": "blue"}, **kwargs) - assert await facade.async_get_cache(**kwargs) == {"answer": "blue"} - - -async def test_azure_blob_rust_required_rule_activates_natively(monkeypatch: pytest.MonkeyPatch) -> None: - account_url: Final = os.environ.get("AZURE_BLOB_CACHE_ACCOUNT_URL") - if account_url is None: - pytest.skip( - "live Azure Blob parity needs AZURE_BLOB_CACHE_ACCOUNT_URL plus DefaultAzureCredential inputs in the environment" - ) - require_rust(monkeypatch, LiteLLMCacheType.AZURE_BLOB) - facade: Final = Cache( - type=LiteLLMCacheType.AZURE_BLOB, - azure_account_url=account_url, - azure_blob_container=f"litellm-parity-{uuid.uuid4().hex[:12]}", - ) - backend: Final = facade.cache - assert isinstance(backend, AzureBlobCache) - try: - native_runtime(facade) - kwargs: Final = completion_kwargs("azure") - await facade.async_add_cache({"answer": "azure"}, **kwargs) - assert await facade.async_get_cache(**kwargs) == {"answer": "azure"} - assert backend.get_cache(facade.get_cache_key(**kwargs))["response"] == {"answer": "azure"} - finally: - backend.container_client.delete_container() - await backend.disconnect() - - -class _SemanticHit: - """A native semantic runtime that answers every lookup with one cached response.""" - - kind: Final = "native" - - def lookup_semantic(self, request: object) -> tuple[object, float | None]: - return {"answer": 42}, 0.97 - - async def async_lookup_semantic(self, request: object) -> tuple[object, float | None]: - return {"answer": 42}, 0.97 - - -@pytest.mark.parametrize("semantic_type", [LiteLLMCacheType.QDRANT_SEMANTIC, LiteLLMCacheType.REDIS_SEMANTIC]) -@pytest.mark.parametrize("use_async", [False, True], ids=["sync", "async"]) -def test_native_semantic_hit_stamps_similarity_on_request_metadata( - semantic_type: LiteLLMCacheType, use_async: bool -) -> None: - """Python semantic backends write `metadata["semantic-similarity"]` on every lookup, and the - facade copies it to the caller's metadata; the native path must report it the same way.""" - facade: Final = Cache() - facade.type = semantic_type - facade._native_cache = ResponseCacheRuntime(cast(NativeResponseCacheRuntime, _SemanticHit())) # pyright: ignore[reportPrivateUsage] # the native path under test has no public setter - metadata: Final[dict[str, object]] = {} - kwargs: Final = { - "cache_key": "semantic-key", - "messages": [{"role": "user", "content": "hello"}], - "metadata": metadata, - } - - result: Final = asyncio.run(facade.async_get_cache(**kwargs)) if use_async else facade.get_cache(**kwargs) - - assert result == {"answer": 42} - assert metadata["semantic-similarity"] == 0.97 diff --git a/tests/test_litellm_rust/test_ocr.py b/tests/test_litellm_rust/test_ocr.py deleted file mode 100644 index 2fbf9817a53..00000000000 --- a/tests/test_litellm_rust/test_ocr.py +++ /dev/null @@ -1,134 +0,0 @@ -import json -import threading -from collections.abc import Generator -from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer -from io import BytesIO -from typing import Final - -import pytest - -import litellm - -pytestmark = pytest.mark.requires_rust_extension - - -@pytest.fixture -def ocr_server() -> Generator[tuple[ThreadingHTTPServer, list[dict[str, object]]]]: - requests: Final[list[dict[str, object]]] = [] - - class Handler(BaseHTTPRequestHandler): - def do_POST(self) -> None: - requests.append( - { - "headers": {name.lower(): value for name, value in self.headers.items()}, - "body": json.loads(self.rfile.read(int(self.headers["Content-Length"]))), - } - ) - if self.headers.get("x-test-stall") == "true": - self.connection.settimeout(2) - try: - self.rfile.read(1) - except TimeoutError: - pass - return - if self.headers.get("User-Agent", "").startswith("python-httpx"): - self.send_response(418) - self.end_headers() - return - status = int(self.headers.get("x-test-status", "200")) - if status != 200: - body = b'{"error":"provider unavailable"}' - self.send_response(status) - self.send_header("Content-Length", str(len(body))) - self.end_headers() - self.wfile.write(body) - return - response: Final = json.dumps( - { - "pages": [{"index": 0, "markdown": "native OCR response", "images": [], "dimensions": None}], - "model": "mistral-ocr-latest", - "usage_info": {"pages_processed": 1, "doc_size_bytes": 3}, - } - ).encode() - self.send_response(200) - self.send_header("Content-Type", "application/json") - self.send_header("Content-Length", str(len(response))) - self.end_headers() - self.wfile.write(response) - - def log_message(self, format: str, *args: object) -> None: - pass - - server: Final = ThreadingHTTPServer(("127.0.0.1", 0), Handler) - thread: Final = threading.Thread(target=lambda: server.serve_forever(poll_interval=0.01), daemon=True) - thread.start() - try: - yield server, requests - finally: - server.shutdown() - server.server_close() - thread.join() - - -def test_native_lifecycle_core_encodes_python_file_input(ocr_server): - server, requests = ocr_server - litellm.rust(True) - response = litellm.ocr( - model="mistral/mistral-ocr-latest", - document={"type": "file", "file": BytesIO(b"abc"), "mime_type": "image/png"}, - api_key="test-key", - api_base=f"http://127.0.0.1:{server.server_port}", - opaque_extension=object(), - ) - assert response.pages[0].markdown == "native OCR response" - assert requests[0]["body"]["document"] == {"type": "image_url", "image_url": "data:image/png;base64,YWJj"} - assert "opaque_extension" not in requests[0]["body"] - - -@pytest.mark.parametrize("asynchronous", [False, True]) -@pytest.mark.asyncio -async def test_native_ocr_failures_do_not_retry_on_python(ocr_server, asynchronous): - server, requests = ocr_server - arguments = { - "model": "mistral-ocr-latest", - "custom_llm_provider": "mistral", - "document": {"type": "document_url", "document_url": "data:application/pdf;base64,YWJj"}, - "api_key": "test-key", - "api_base": f"http://127.0.0.1:{server.server_port}", - "extra_headers": {"x-test-status": "503"}, - "num_retries": 0, - } - litellm.rust(True) - with pytest.raises(litellm.ServiceUnavailableError) as caught: - await litellm.aocr(**arguments) if asynchronous else litellm.ocr(**arguments) - assert caught.value.status_code == 503 - assert len(requests) == 1 - assert not requests[0]["headers"].get("user-agent", "").startswith("python-httpx") - - -@pytest.mark.parametrize("asynchronous", [False, True]) -@pytest.mark.asyncio -async def test_native_ocr_enforces_request_deadline_without_fallback(ocr_server, asynchronous): - import asyncio - import time - - server, requests = ocr_server - litellm.rust(True) - arguments = { - "model": "mistral/mistral-ocr-latest", - "document": {"type": "document_url", "document_url": "data:application/pdf;base64,YWJj"}, - "api_key": "test-key", - "api_base": f"http://127.0.0.1:{server.server_port}", - "extra_headers": {"x-test-stall": "true"}, - "timeout": 0.1, - "num_retries": 0, - } - started = time.monotonic() - with pytest.raises(litellm.Timeout): - await asyncio.wait_for( - litellm.aocr(**arguments) if asynchronous else asyncio.to_thread(litellm.ocr, **arguments), - timeout=3, - ) - assert 0.09 <= time.monotonic() - started < 3 - assert len(requests) == 1 - assert not requests[0]["headers"].get("user-agent", "").startswith("python-httpx") diff --git a/tests/test_litellm_rust/tokenizer/__init__.py b/tests/test_litellm_rust/tokenizer/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/test_litellm_rust/test_tokenizer.py b/tests/test_litellm_rust/tokenizer/test_fast_count.py similarity index 51% rename from tests/test_litellm_rust/test_tokenizer.py rename to tests/test_litellm_rust/tokenizer/test_fast_count.py index 98d5259b652..2902b79dca8 100644 --- a/tests/test_litellm_rust/test_tokenizer.py +++ b/tests/test_litellm_rust/tokenizer/test_fast_count.py @@ -12,67 +12,6 @@ from tests.test_litellm.litellm_core_utils.test_decode_special_tokens import TOK pytestmark = pytest.mark.requires_rust_extension -def test_tiktoken_codec_round_trips_and_counts() -> None: - tokenizer: Final = _native.Tokenizer.from_tiktoken("cl100k_base") - encoded: Final = tokenizer.encode("hello world") - - assert tokenizer.name == "cl100k_base" - assert tokenizer.count("hello world") == len(encoded) - assert tokenizer.decode(encoded) == "hello world" - - -def test_huggingface_codec_skips_special_tokens() -> None: - tokenizer: Final = _native.Tokenizer.from_json(claude_json_str) - encoded: Final = tokenizer.encode("hello") - - assert "" in tokenizer.decode(encoded, skip_special_tokens=False) - assert tokenizer.decode(encoded, skip_special_tokens=True) == "hello" - - -def test_tiktoken_codec_keeps_the_requested_encoding_name() -> None: - assert _native.Tokenizer.from_tiktoken("gpt2").name == "gpt2" - assert _native.Tokenizer.from_tiktoken("r50k_base").name == "r50k_base" - assert _native.Tokenizer.from_tiktoken("gpt2").encode("hi") == _native.Tokenizer.from_tiktoken("r50k_base").encode( - "hi" - ) - - -def test_tiktoken_codec_exposes_its_vocabulary() -> None: - reference: Final = tiktoken.get_encoding("cl100k_base") - tokenizer: Final = _native.Tokenizer.from_tiktoken("cl100k_base") - - assert tokenizer.special_tokens() == reference._special_tokens - assert tokenizer.max_token_value() == reference.max_token_value - assert tokenizer.token_byte_values() == reference.token_byte_values() - assert tokenizer.encode_single_token(b"hello") == reference.encode_single_token("hello") - assert tokenizer.is_special_token(reference.eot_token) and not tokenizer.is_special_token(0) - with pytest.raises(KeyError): - tokenizer.encode_single_token(b"<|not-a-token|>") - - -def test_huggingface_codec_rejects_tiktoken_only_calls() -> None: - tokenizer: Final = _native.Tokenizer.from_json(claude_json_str) - with pytest.raises(ValueError, match="requires a tiktoken encoding"): - tokenizer.token_byte_values() - with pytest.raises(ValueError, match="requires a Hugging Face tokenizer"): - _native.Tokenizer.from_tiktoken("cl100k_base").get_vocab() - - -def test_unknown_tiktoken_encoding_raises_value_error() -> None: - with pytest.raises(ValueError, match="unsupported tokenizer"): - _native.Tokenizer.from_tiktoken("unknown-encoding") - - -def test_tiktoken_codec_decodes_truncated_unicode_like_python() -> None: - reference: Final = tiktoken.get_encoding("cl100k_base") - tokenizer: Final = _native.Tokenizer.from_tiktoken(reference.name) - encoded: Final = reference.encode("🙂漢字") - - assert tuple(tokenizer.decode(encoded[:end]) for end in range(1, len(encoded) + 1)) == tuple( - reference.decode(encoded[:end]) for end in range(1, len(encoded) + 1) - ) - - FAST_TEXTS: Final = ( "", "hello world <|endoftext|>", diff --git a/tests/test_litellm_rust/tokenizer/test_huggingface.py b/tests/test_litellm_rust/tokenizer/test_huggingface.py new file mode 100644 index 00000000000..05c5989c676 --- /dev/null +++ b/tests/test_litellm_rust/tokenizer/test_huggingface.py @@ -0,0 +1,24 @@ +from typing import Final + +import pytest + +from litellm.rust_bridge import _native +from litellm.utils import claude_json_str + +pytestmark = pytest.mark.requires_rust_extension + + +def test_huggingface_codec_skips_special_tokens() -> None: + tokenizer: Final = _native.Tokenizer.from_json(claude_json_str) + encoded: Final = tokenizer.encode("hello") + + assert "" in tokenizer.decode(encoded, skip_special_tokens=False) + assert tokenizer.decode(encoded, skip_special_tokens=True) == "hello" + + +def test_huggingface_codec_rejects_tiktoken_only_calls() -> None: + tokenizer: Final = _native.Tokenizer.from_json(claude_json_str) + with pytest.raises(ValueError, match="requires a tiktoken encoding"): + tokenizer.token_byte_values() + with pytest.raises(ValueError, match="requires a Hugging Face tokenizer"): + _native.Tokenizer.from_tiktoken("cl100k_base").get_vocab() diff --git a/tests/test_litellm_rust/tokenizer/test_tiktoken.py b/tests/test_litellm_rust/tokenizer/test_tiktoken.py new file mode 100644 index 00000000000..c204d7eaf2f --- /dev/null +++ b/tests/test_litellm_rust/tokenizer/test_tiktoken.py @@ -0,0 +1,53 @@ +from typing import Final + +import pytest +import tiktoken + +from litellm.rust_bridge import _native + +pytestmark = pytest.mark.requires_rust_extension + + +def test_tiktoken_codec_round_trips_and_counts() -> None: + tokenizer: Final = _native.Tokenizer.from_tiktoken("cl100k_base") + encoded: Final = tokenizer.encode("hello world") + + assert tokenizer.name == "cl100k_base" + assert tokenizer.count("hello world") == len(encoded) + assert tokenizer.decode(encoded) == "hello world" + + +def test_tiktoken_codec_keeps_the_requested_encoding_name() -> None: + assert _native.Tokenizer.from_tiktoken("gpt2").name == "gpt2" + assert _native.Tokenizer.from_tiktoken("r50k_base").name == "r50k_base" + assert _native.Tokenizer.from_tiktoken("gpt2").encode("hi") == _native.Tokenizer.from_tiktoken("r50k_base").encode( + "hi" + ) + + +def test_tiktoken_codec_exposes_its_vocabulary() -> None: + reference: Final = tiktoken.get_encoding("cl100k_base") + tokenizer: Final = _native.Tokenizer.from_tiktoken("cl100k_base") + + assert tokenizer.special_tokens() == reference._special_tokens + assert tokenizer.max_token_value() == reference.max_token_value + assert tokenizer.token_byte_values() == reference.token_byte_values() + assert tokenizer.encode_single_token(b"hello") == reference.encode_single_token("hello") + assert tokenizer.is_special_token(reference.eot_token) and not tokenizer.is_special_token(0) + with pytest.raises(KeyError): + tokenizer.encode_single_token(b"<|not-a-token|>") + + +def test_unknown_tiktoken_encoding_raises_value_error() -> None: + with pytest.raises(ValueError, match="unsupported tokenizer"): + _native.Tokenizer.from_tiktoken("unknown-encoding") + + +def test_tiktoken_codec_decodes_truncated_unicode_like_python() -> None: + reference: Final = tiktoken.get_encoding("cl100k_base") + tokenizer: Final = _native.Tokenizer.from_tiktoken(reference.name) + encoded: Final = reference.encode("🙂漢字") + + assert tuple(tokenizer.decode(encoded[:end]) for end in range(1, len(encoded) + 1)) == tuple( + reference.decode(encoded[:end]) for end in range(1, len(encoded) + 1) + )