scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-26 01:12:21 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 99

test(rust): reorganize core crate tests and split cache and OCR suites (#43177)

Browse source 
* 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 <noreply@anthropic.com>

* 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 <noreply@anthropic.com>

* 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 <noreply@anthropic.com>

* 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 <noreply@anthropic.com>

---------

Co-authored-by: Yujong Lee <yujong@berri.ai>
Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>
This commit is contained in:
devin-ai-integration[bot] 2026-09-25 16:18:36 +00:00 • committed by GitHub
parent c289d5d6fb
commit 10413796c6
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
52 changed files with 7015 additions and 7382 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

1
litellm-rust/Cargo.lock generated
View file

@ -3092,6 +3092,7 @@ dependencies = [
"tokio-tungstenite",
"url",
"veil",
"wiremock",
]
[[package]]

1
litellm-rust/crates/core/Cargo.toml
View file

@ -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"

26
litellm-rust/crates/core/src/chat_completions/handler.rs
View file

@ -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);
}
}

244
litellm-rust/crates/core/src/chat_completions/prepare.rs
View file

@ -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::<usize>().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<String>) {
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, .. })
));
}
}
}

145
litellm-rust/crates/core/src/messages/common_utils.rs
View file

@ -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<Vec<String>>,
}
impl SecretSource for RecordingSecrets {
fn get_secret_str<'a>(
&'a self,
name: &'a str,
) -> BoxFuture<'a, Result<Option<SecretValue>, 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::<usize>().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::<Vec<_>>()
);
}
use crate::messages::Error;
#[test]
fn provider_config_resolves_anthropic_and_azure_ai() {

157
litellm-rust/crates/core/src/ocr/document.rs
View file

@ -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<Value>,
}
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)
);
}
}

209
litellm-rust/crates/core/src/ocr/mod.rs
View file

@ -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<Error> for NoHooks {
fn before_send(&self, wire: WireRequest) -> BoxFuture<'_, Result<WireRequest, Error>> {
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<LiteLLMOcrResponse, Error> {
crate::ocr::client::perform(&ocr_client(), request).await
}
pub(crate) async fn perform_ocr_with(host: LocalOcrHost) -> Result<LiteLLMOcrResponse, Error> {
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<MockResponse>,
) -> (String, Arc<Mutex<Vec<String>>>, 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::<usize>().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::<String>();
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())
})
}
}

205
litellm-rust/crates/core/src/ocr/prepare.rs
View file

@ -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<Error> for NoHooks {
fn before_send(&self, wire: WireRequest) -> BoxFuture<'_, Result<WireRequest, Error>> {
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 {

3521
litellm-rust/crates/core/src/ocr/route.rs
View file

File diff suppressed because it is too large Load diff

268
litellm-rust/crates/core/tests/audio_transcription.rs
View file

@ -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<String, Value> {
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:?}");
}

320
litellm-rust/crates/core/tests/chat_completions.rs Normal file
View file

@ -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<String, Value> {
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());
}

471
litellm-rust/crates/core/tests/messages.rs
View file

@ -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<Vec<String>>,
}
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<Option<SecretValue>, 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::<usize>().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"));
}

94
litellm-rust/crates/core/tests/messages/main.rs Normal file
View file

@ -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<String, Value> {
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<Item = (&'a str, &'a str)>) -> Option<Map<String, Value>> {
Some(
pairs
.into_iter()
.map(|(name, value)| (name.to_string(), Value::from(value)))
.collect(),
)
}
async fn run_with(
secrets: Arc<RecordingSecrets>,
call: MessagesCall,
) -> Result<MessagesOutput, Error> {
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<MessagesOutput, Error> {
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"),
}
}

269
litellm-rust/crates/core/tests/messages/request.rs Normal file
View file

@ -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
);
}

136
litellm-rust/crates/core/tests/messages/response.rs Normal file
View file

@ -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())
);
}

94
litellm-rust/crates/core/tests/messages/secrets.rs Normal file
View file

@ -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());
}

163
litellm-rust/crates/core/tests/messages/stream.rs Normal file
View file

@ -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<Vec<Seen>>,
}
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<Messages> for RecordingStreamHost {
async fn project(&self) -> Result<MessagesCall, Error> {
self.call.project().await
}
async fn custom_op(&self, op: Infallible) -> Result<(), Error> {
match op {}
}
async fn open(&self, (): ()) -> Result<Demand, Error> {
Ok(self.record(Seen::Open))
}
async fn deliver(&self, chunk: Bytes) -> Result<Demand, Error> {
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<MessagesOutput, Error> {
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<u8> = 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());
}

173
litellm-rust/crates/core/tests/ocr/aws_textract.rs Normal file
View file

@ -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<String, String> = ["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}"
);
}

270
litellm-rust/crates/core/tests/ocr/azure_ai.rs Normal file
View file

@ -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<Self> {
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<CountingToken>,
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<String> = 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());
}

441
litellm-rust/crates/core/tests/ocr/azure_document_intelligence.rs Normal file
View file

@ -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}");
}

42
litellm-rust/crates/core/tests/ocr/cohere.rs Normal file
View file

@ -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());
}

182
litellm-rust/crates/core/tests/ocr/documents.rs Normal file
View file

@ -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::<u64>()));
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());
}

269
litellm-rust/crates/core/tests/ocr/lifecycle.rs Normal file
View file

@ -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<Mutex<Vec<&'static str>>>,
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<Option<LiteLLMOcrRequest>>,
trace: Mutex<Vec<String>>,
}
impl Host<Ocr> for CallerTokenHost {
async fn project(&self) -> Result<OcrProjection, Error> {
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<WireRequest, Error> {
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"]
);
}

284
litellm-rust/crates/core/tests/ocr/machine.rs Normal file
View file

@ -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<WireRequest, HostFailure<Error>>,
) -> (
Result<LiteLLMOcrResponse, Error>,
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::<u64>()));
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<Error>,
#[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<Notify>,
dropped: Arc<AtomicBool>,
}
struct TokenFutureDrop(Arc<AtomicBool>);
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();
}

125
litellm-rust/crates/core/tests/ocr/main.rs Normal file
View file

@ -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<String, Value> {
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<LiteLLMOcrResponse, Error> {
litellm_core::ocr::client::perform(&ocr_client(), request).await
}
async fn perform_with(host: LocalOcrHost) -> Result<LiteLLMOcrResponse, Error> {
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)
}

248
litellm-rust/crates/core/tests/ocr/mistral.rs Normal file
View file

@ -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}");
}

321
litellm-rust/crates/core/tests/ocr/reducto.rs Normal file
View file

@ -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"]);
}
}

184
litellm-rust/crates/core/tests/ocr/vertex_ai.rs Normal file
View file

@ -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);
}

155
litellm-rust/crates/core/tests/support/mod.rs Normal file
View file

@ -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<Item = ResponseTemplate>) -> 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<Item = ResponseTemplate>,
) {
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<Request> {
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<String>;
}
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<String> {
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<Vec<String>>,
}
impl RecordingSecrets {
pub fn new<'a>(values: impl IntoIterator<Item = (&'a str, &'a str)>) -> 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<String> {
self.requested.lock().unwrap().clone()
}
}
impl SecretSource for RecordingSecrets {
fn get_secret_str<'a>(
&'a self,
name: &'a str,
) -> BoxFuture<'a, Result<Option<SecretValue>, 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())))
})
}
}

44
litellm-rust/crates/llms/src/reducto/ocr/transformation.rs
View file

@ -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,
&params,
&[],
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!({

79
litellm-rust/crates/llms/tests/ocr_handler.rs Normal file
View file

@ -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<bytes::Bytes, Error> {
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);
}

1
tests/test_litellm_rust/AGENTS.md Normal file
View file

@ -0,0 +1 @@
This directory holds only the tests that cannot be written in the Rust code

1
tests/test_litellm_rust/cache/__init__.py vendored Normal file
View file

@ -0,0 +1 @@

30
tests/test_litellm_rust/cache/conftest.py vendored Normal file
View file

@ -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()

173
tests/test_litellm_rust/cache/test_azure_blob.py vendored Normal file
View file

@ -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()

117
tests/test_litellm_rust/cache/test_disk.py vendored Normal file
View file

@ -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],
}

397
tests/test_litellm_rust/cache/test_facade.py vendored Normal file
View file

@ -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)

242
tests/test_litellm_rust/cache/test_gcs.py vendored Normal file
View file

@ -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

286
tests/test_litellm_rust/cache/test_qdrant_semantic.py vendored Normal file
View file

@ -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"}

228
tests/test_litellm_rust/cache/test_redis.py vendored Normal file
View file

@ -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

606
tests/test_litellm_rust/cache/test_redis_semantic.py vendored Normal file
View file

@ -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"}

264
tests/test_litellm_rust/cache/test_rollout.py vendored Normal file
View file

@ -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

187
tests/test_litellm_rust/cache/test_s3.py vendored Normal file
View file

@ -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)

0
tests/test_litellm_rust/test_valkey_semantic_cache_native.py → tests/test_litellm_rust/cache/test_valkey_semantic.py vendored
View file

15
tests/test_litellm_rust/ocr/test_requests.py
View file

@ -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

40
tests/test_litellm_rust/support/cache.py Normal file
View file

@ -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}"}]}

2397
tests/test_litellm_rust/test_cache.py
View file

File diff suppressed because it is too large Load diff

134
tests/test_litellm_rust/test_ocr.py
View file

@ -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")

0
tests/test_litellm_rust/tokenizer/__init__.py Normal file
View file

61
tests/test_litellm_rust/test_tokenizer.py → tests/test_litellm_rust/tokenizer/test_fast_count.py
View file

@ -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("<SOS>hello<EOT>")
assert "<SOS>" 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|>",

24
tests/test_litellm_rust/tokenizer/test_huggingface.py Normal file
View file

@ -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("<SOS>hello<EOT>")
assert "<SOS>" 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()

53
tests/test_litellm_rust/tokenizer/test_tiktoken.py Normal file
View file

@ -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)
)