mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-07 08:26:10 +00:00
perf(rust): share media buffers and stream replayable JSON
This commit is contained in:
Yujong Lee
parent
e058aa68c4
commit
77e85554a7
60 changed files with 3050 additions and 452 deletions
20
litellm-rust/Cargo.lock
generated
20
litellm-rust/Cargo.lock
generated
|
|
@ -536,6 +536,15 @@ dependencies = [
|
|||
"either",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "bytestring"
|
||||
version = "1.5.1"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "86566c496f2f47d9b8147a4c8b02ffdb69c919fe0c2b2e7195d22cbba0e635c9"
|
||||
dependencies = [
|
||||
"bytes",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "cast"
|
||||
version = "0.3.0"
|
||||
|
|
@ -1435,6 +1444,12 @@ dependencies = [
|
|||
"aws-sigv4",
|
||||
"aws-smithy-runtime-api",
|
||||
"aws-types",
|
||||
"base64",
|
||||
"bytes",
|
||||
"bytestring",
|
||||
"futures-util",
|
||||
"h2 0.4.15",
|
||||
"http 1.4.2",
|
||||
"rand 0.8.7",
|
||||
"reqwest",
|
||||
"serde",
|
||||
|
|
@ -1449,6 +1464,7 @@ dependencies = [
|
|||
name = "litellm-python-bridge"
|
||||
version = "0.1.0"
|
||||
dependencies = [
|
||||
"bytes",
|
||||
"criterion",
|
||||
"futures-util",
|
||||
"litellm-ai-gateway",
|
||||
|
|
@ -1456,8 +1472,10 @@ dependencies = [
|
|||
"litellm-python-interop",
|
||||
"pyo3",
|
||||
"pyo3-async-runtimes",
|
||||
"reqwest",
|
||||
"serde",
|
||||
"serde_json",
|
||||
"sha2 0.10.9",
|
||||
"tokio",
|
||||
"tokio-tungstenite",
|
||||
"tracing",
|
||||
|
|
@ -1468,6 +1486,7 @@ dependencies = [
|
|||
name = "litellm-python-interop"
|
||||
version = "0.1.0"
|
||||
dependencies = [
|
||||
"bytes",
|
||||
"pyo3",
|
||||
"pythonize",
|
||||
"rstest",
|
||||
|
|
@ -1686,6 +1705,7 @@ version = "0.29.2"
|
|||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "4688ddedf473e32662b9b067670129a8afb8c18e351482c70d62ba4a88171e8b"
|
||||
dependencies = [
|
||||
"bytes",
|
||||
"libc",
|
||||
"once_cell",
|
||||
"portable-atomic",
|
||||
|
|
|
|||
|
|
@ -14,13 +14,16 @@ license = "MIT"
|
|||
repository = "https://github.com/BerriAI/litellm"
|
||||
|
||||
[workspace.dependencies]
|
||||
bytes = "1.10"
|
||||
bytestring = "1.5.1"
|
||||
h2 = "0.4"
|
||||
tracing = "0.1"
|
||||
tracing-subscriber = { version = "0.3", default-features = false, features = ["registry", "std"] }
|
||||
litellm-core = { path = "crates/core" }
|
||||
litellm-ai-gateway = { path = "crates/ai-gateway", default-features = false }
|
||||
litellm-python-interop = { path = "crates/python-interop" }
|
||||
axum = "0.7"
|
||||
pyo3 = "0.29.2"
|
||||
pyo3 = { version = "0.29.2", features = ["bytes"] }
|
||||
pyo3-async-runtimes = { version = "0.29.0", features = ["tokio-runtime"] }
|
||||
pythonize = "0.29.0"
|
||||
rand = "0.8"
|
||||
|
|
|
|||
|
|
@ -1,4 +1,4 @@
|
|||
use std::io::{Read, Write};
|
||||
use std::io::{BufRead, BufReader, Read, Write};
|
||||
use std::net::TcpListener;
|
||||
use std::thread;
|
||||
|
||||
|
|
@ -12,11 +12,25 @@ async fn bedrock_request_is_signed_and_contains_audio() {
|
|||
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);
|
||||
let mut reader = BufReader::new(&mut stream);
|
||||
let mut headers = String::new();
|
||||
loop {
|
||||
let mut line = String::new();
|
||||
reader.read_line(&mut line).expect("headers");
|
||||
if line == "\r\n" {
|
||||
break;
|
||||
}
|
||||
headers.push_str(&line);
|
||||
}
|
||||
let length = headers
|
||||
.lines()
|
||||
.find_map(|line| line.strip_prefix("content-length: "))
|
||||
.expect("content length")
|
||||
.parse::<usize>()
|
||||
.unwrap();
|
||||
let mut body = vec![0; length];
|
||||
reader.read_exact(&mut body).expect("complete request body");
|
||||
let request = headers + &String::from_utf8(body).unwrap();
|
||||
assert!(request.contains("POST /model/mistral.voxtral-mini-3b-2507/converse"));
|
||||
assert!(request.contains("authorization: AWS4-HMAC-SHA256"));
|
||||
assert!(request.contains("x-amz-date:"));
|
||||
|
|
|
|||
|
|
@ -1,14 +0,0 @@
|
|||
use std::sync::OnceLock;
|
||||
use std::time::Duration;
|
||||
|
||||
const HTTP_CLIENT_TIMEOUT_SECS: u64 = 600;
|
||||
|
||||
pub(crate) fn http_client() -> &'static reqwest::Client {
|
||||
static CLIENT: OnceLock<reqwest::Client> = OnceLock::new();
|
||||
CLIENT.get_or_init(|| {
|
||||
reqwest::Client::builder()
|
||||
.timeout(Duration::from_secs(HTTP_CLIENT_TIMEOUT_SECS))
|
||||
.build()
|
||||
.expect("failed to build reqwest client")
|
||||
})
|
||||
}
|
||||
|
|
@ -14,7 +14,6 @@
|
|||
//! for the load-time config reader.
|
||||
|
||||
pub mod audio_transcription;
|
||||
mod client;
|
||||
pub mod io;
|
||||
pub mod ocr;
|
||||
|
||||
|
|
|
|||
|
|
@ -1,5 +1,5 @@
|
|||
use litellm_core::http_utils::body::JsonPayload;
|
||||
use std::net::IpAddr;
|
||||
use std::time::{Duration, Instant};
|
||||
|
||||
use base64::Engine;
|
||||
use base64::engine::general_purpose::STANDARD as BASE64_STANDARD;
|
||||
|
|
@ -17,10 +17,7 @@ use litellm_core::providers::vertex_ai::ocr::transformation::{
|
|||
VERTEX_AI_DEEPSEEK_OCR_CONFIG, VERTEX_AI_OCR_CONFIG,
|
||||
};
|
||||
|
||||
use crate::client::http_client;
|
||||
|
||||
const ERROR_BODY_MAX_CHARS: usize = 256;
|
||||
const AZURE_DOCUMENT_INTELLIGENCE_POLL_TIMEOUT_SECS: u64 = 120;
|
||||
const DEFAULT_MAX_IMAGE_URL_DOWNLOAD_SIZE_MB: f64 = 50.0;
|
||||
const MAX_SAFE_FETCH_REDIRECTS: usize = 10;
|
||||
|
||||
|
|
@ -74,11 +71,11 @@ pub(super) fn string_headers(
|
|||
.collect()
|
||||
}
|
||||
|
||||
fn document_url_field(document: &Value) -> Result<Option<(&str, &str)>, Error> {
|
||||
fn document_url_field(document: &JsonPayload) -> Result<Option<(&str, &str)>, Error> {
|
||||
let Some(object) = document.as_object() else {
|
||||
return Ok(None);
|
||||
};
|
||||
let Some(doc_type) = object.get("type").and_then(Value::as_str) else {
|
||||
let Some(doc_type) = object.get("type").and_then(JsonPayload::as_str) else {
|
||||
return Ok(None);
|
||||
};
|
||||
let field = match doc_type {
|
||||
|
|
@ -86,7 +83,7 @@ fn document_url_field(document: &Value) -> Result<Option<(&str, &str)>, Error> {
|
|||
"image_url" => "image_url",
|
||||
_ => return Ok(None),
|
||||
};
|
||||
let Some(url) = object.get(field).and_then(Value::as_str) else {
|
||||
let Some(url) = object.get(field).and_then(JsonPayload::as_str) else {
|
||||
return Ok(None);
|
||||
};
|
||||
Ok(Some((field, url)))
|
||||
|
|
@ -249,7 +246,10 @@ async fn read_response_with_limit(
|
|||
Ok(bytes)
|
||||
}
|
||||
|
||||
pub(super) async fn convert_document_url_to_data_uri(document: Value) -> Result<Value, Error> {
|
||||
pub(super) async fn convert_document_url_to_data_uri<D: Into<JsonPayload>>(
|
||||
document: D,
|
||||
) -> Result<JsonPayload, Error> {
|
||||
let document = document.into();
|
||||
let Some((field, url)) = document_url_field(&document)? else {
|
||||
return Ok(document);
|
||||
};
|
||||
|
|
@ -285,109 +285,8 @@ pub(super) async fn convert_document_url_to_data_uri(document: Value) -> Result<
|
|||
.as_object()
|
||||
.cloned()
|
||||
.ok_or_else(|| Error::InvalidRequest("OCR document must be an object".to_string()))?;
|
||||
transformed.insert(field.to_string(), Value::String(data_uri));
|
||||
Ok(Value::Object(transformed))
|
||||
}
|
||||
|
||||
fn same_origin(left: &str, right: &str) -> bool {
|
||||
let Ok(left) = reqwest::Url::parse(left) else {
|
||||
return false;
|
||||
};
|
||||
let Ok(right) = reqwest::Url::parse(right) else {
|
||||
return false;
|
||||
};
|
||||
left.scheme() == right.scheme()
|
||||
&& left.host_str() == right.host_str()
|
||||
&& left.port_or_known_default() == right.port_or_known_default()
|
||||
}
|
||||
|
||||
fn retry_after_secs(response: &reqwest::Response) -> u64 {
|
||||
response
|
||||
.headers()
|
||||
.get(reqwest::header::RETRY_AFTER)
|
||||
.and_then(|value| value.to_str().ok())
|
||||
.and_then(|value| value.parse::<u64>().ok())
|
||||
.unwrap_or(2)
|
||||
}
|
||||
|
||||
fn operation_status(response_json: &Value) -> Result<&str, Error> {
|
||||
let status = response_json
|
||||
.get("status")
|
||||
.and_then(Value::as_str)
|
||||
.ok_or(Error::MissingField("status"))?;
|
||||
match status {
|
||||
"succeeded" => Ok("succeeded"),
|
||||
"running" | "notStarted" => Ok("running"),
|
||||
"failed" => {
|
||||
let message = response_json
|
||||
.get("error")
|
||||
.and_then(|error| error.get("message"))
|
||||
.and_then(Value::as_str)
|
||||
.unwrap_or("Unknown error");
|
||||
Err(Error::InvalidResponse(format!(
|
||||
"Azure Document Intelligence analysis failed: {message}"
|
||||
)))
|
||||
}
|
||||
other => Err(Error::InvalidResponse(format!(
|
||||
"Unknown operation status: {other}"
|
||||
))),
|
||||
}
|
||||
}
|
||||
|
||||
pub(super) async fn poll_document_intelligence(
|
||||
operation_url: &str,
|
||||
original_url: &str,
|
||||
headers: &[(String, String)],
|
||||
timeout: Option<Duration>,
|
||||
) -> Result<Value, Error> {
|
||||
if !same_origin(operation_url, original_url) {
|
||||
return Err(Error::InvalidResponse(
|
||||
"Azure Document Intelligence: rejected cross-origin polling URL".to_string(),
|
||||
));
|
||||
}
|
||||
|
||||
let start = Instant::now();
|
||||
let timeout = timeout.unwrap_or(Duration::from_secs(
|
||||
AZURE_DOCUMENT_INTELLIGENCE_POLL_TIMEOUT_SECS,
|
||||
));
|
||||
loop {
|
||||
if start.elapsed() > timeout {
|
||||
return Err(Error::Network(format!(
|
||||
"Azure Document Intelligence operation polling timed out after {} seconds",
|
||||
timeout.as_secs()
|
||||
)));
|
||||
}
|
||||
|
||||
let mut request_builder = http_client().get(operation_url);
|
||||
for (key, value) in headers {
|
||||
if key.eq_ignore_ascii_case("ocp-apim-subscription-key") {
|
||||
request_builder = request_builder.header(key, value);
|
||||
}
|
||||
}
|
||||
let response = request_builder
|
||||
.send()
|
||||
.await
|
||||
.map_err(|err| Error::Network(err.to_string()))?;
|
||||
let retry_after = retry_after_secs(&response);
|
||||
let status = response.status();
|
||||
let text = response
|
||||
.text()
|
||||
.await
|
||||
.map_err(|err| Error::Network(err.to_string()))?;
|
||||
if !status.is_success() {
|
||||
return Err(Error::Http {
|
||||
status: status.as_u16(),
|
||||
body: truncate_error_body(&text),
|
||||
});
|
||||
}
|
||||
let response_json: Value = serde_json::from_str(&text).map_err(|err| {
|
||||
Error::InvalidResponse(format!("invalid Azure DI poll response JSON: {err}"))
|
||||
})?;
|
||||
if operation_status(&response_json)? == "succeeded" {
|
||||
return Ok(response_json);
|
||||
}
|
||||
tokio::time::sleep(Duration::from_secs(retry_after)).await;
|
||||
}
|
||||
transformed.insert(field.to_string(), JsonPayload::from(data_uri));
|
||||
Ok(JsonPayload::Object(transformed))
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
|
|
|
|||
|
|
@ -1,12 +1,8 @@
|
|||
use litellm_core::error::Error;
|
||||
use litellm_core::http_utils::http_request;
|
||||
use litellm_core::ocr::transformation::OcrResponseHandling;
|
||||
use serde_json::Value;
|
||||
|
||||
use super::common_utils::{poll_document_intelligence, truncate_error_body};
|
||||
use super::hooks::OcrLifecycleHooks;
|
||||
use super::types::PreparedOcrRequest;
|
||||
use crate::client::http_client;
|
||||
|
||||
#[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)]
|
||||
pub(crate) async fn execute_ocr_provider_call(
|
||||
|
|
@ -14,63 +10,5 @@ pub(crate) async fn execute_ocr_provider_call(
|
|||
hooks: &OcrLifecycleHooks,
|
||||
) -> Result<Value, Error> {
|
||||
let request = hooks.prepare_provider_request(request).await?;
|
||||
let mut request_builder = http_client().post(&request.url).json(&request.body);
|
||||
for (key, value) in &request.upstream_headers {
|
||||
request_builder = request_builder.header(key, value);
|
||||
}
|
||||
if let Some(duration) = request.timeout {
|
||||
request_builder = request_builder.timeout(duration);
|
||||
}
|
||||
|
||||
let response = http_request(request_builder)
|
||||
.await
|
||||
.map_err(|err| Error::Network(err.to_string()))?;
|
||||
|
||||
let status = response.status();
|
||||
if request.config.response_handling() == OcrResponseHandling::AzureDocumentIntelligencePoll
|
||||
&& status.as_u16() == 202
|
||||
{
|
||||
let operation_url = response
|
||||
.headers()
|
||||
.get("operation-location")
|
||||
.and_then(|value| value.to_str().ok())
|
||||
.map(str::to_string)
|
||||
.ok_or_else(|| {
|
||||
Error::InvalidResponse(
|
||||
"Azure Document Intelligence returned 202 but no Operation-Location header found"
|
||||
.to_string(),
|
||||
)
|
||||
})?;
|
||||
let response_json = poll_document_intelligence(
|
||||
&operation_url,
|
||||
&request.url,
|
||||
&request.upstream_headers,
|
||||
request.timeout,
|
||||
)
|
||||
.await?;
|
||||
return Ok(request
|
||||
.config
|
||||
.transform_ocr_response(&request.model, response_json)?
|
||||
.into_json());
|
||||
}
|
||||
|
||||
let text = response
|
||||
.text()
|
||||
.await
|
||||
.map_err(|err| Error::Network(err.to_string()))?;
|
||||
|
||||
if !status.is_success() {
|
||||
return Err(Error::Http {
|
||||
status: status.as_u16(),
|
||||
body: truncate_error_body(&text),
|
||||
});
|
||||
}
|
||||
|
||||
let response_json: Value = serde_json::from_str(&text)
|
||||
.map_err(|err| Error::InvalidResponse(format!("invalid OCR response JSON: {err}")))?;
|
||||
|
||||
Ok(request
|
||||
.config
|
||||
.transform_ocr_response(&request.model, response_json)?
|
||||
.into_json())
|
||||
litellm_core::ocr::execute_ocr_provider_call(request).await
|
||||
}
|
||||
|
|
|
|||
|
|
@ -1,5 +1,6 @@
|
|||
use litellm_core::call_lifecycle::{CallLifecycleContext, CallLifecycleHooks, CallLifecycleTiming};
|
||||
use litellm_core::error::Error;
|
||||
use litellm_core::http_utils::body::JsonPayload;
|
||||
use serde_json::{Map, Value, json};
|
||||
use std::future::Future;
|
||||
use std::pin::Pin;
|
||||
|
|
@ -64,7 +65,7 @@ impl OcrLifecycleHooks {
|
|||
Err(_) => optional_params,
|
||||
};
|
||||
Ok(PreparedOcrRequest {
|
||||
document,
|
||||
document: document.into(),
|
||||
optional_params,
|
||||
..request
|
||||
})
|
||||
|
|
@ -95,7 +96,7 @@ impl OcrLifecycleHooks {
|
|||
request.document
|
||||
};
|
||||
let body = config
|
||||
.transform_ocr_request(&request.model, document, request.optional_params)?
|
||||
.transform_ocr_payload(&request.model, document, request.optional_params)?
|
||||
.data;
|
||||
let body = self
|
||||
.run_during_call_guardrails(&model, &custom_llm_provider, &url, body)
|
||||
|
|
@ -115,8 +116,8 @@ impl OcrLifecycleHooks {
|
|||
model: &str,
|
||||
custom_llm_provider: &str,
|
||||
url: &str,
|
||||
body: Value,
|
||||
) -> Result<Value, Error> {
|
||||
body: JsonPayload,
|
||||
) -> Result<JsonPayload, Error> {
|
||||
if self.guardrail_runner.is_empty() {
|
||||
return Ok(body);
|
||||
}
|
||||
|
|
@ -133,7 +134,7 @@ impl OcrLifecycleHooks {
|
|||
.run_during_call(&context, guardrail_request)
|
||||
.await
|
||||
.map_err(guardrail_error_to_core_error)?;
|
||||
parse_ocr_during_call_guardrail_request(guardrail_request)
|
||||
parse_ocr_during_call_guardrail_request(guardrail_request).map(JsonPayload::from)
|
||||
}
|
||||
|
||||
fn standard_logging_payload(
|
||||
|
|
|
|||
|
|
@ -1,5 +1,6 @@
|
|||
use litellm_core::Error;
|
||||
use litellm_core::call_lifecycle::CallLifecycle;
|
||||
use litellm_core::http_utils::body::JsonPayload;
|
||||
use serde_json::Value;
|
||||
|
||||
mod common_utils;
|
||||
|
|
@ -14,7 +15,7 @@ use handler::execute_ocr_provider_call;
|
|||
use prepare::{PreparedOcrCall, prepare_ocr_call};
|
||||
|
||||
#[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)]
|
||||
pub async fn ocr(request: OcrRequest<'_>) -> Result<Value, Error> {
|
||||
pub async fn ocr<D: Into<JsonPayload>>(request: OcrRequest<'_, D>) -> Result<Value, Error> {
|
||||
let PreparedOcrCall { request, hooks } = prepare_ocr_call(request);
|
||||
CallLifecycle::default()
|
||||
.run_request(request, &hooks, |request| {
|
||||
|
|
|
|||
|
|
@ -1,3 +1,4 @@
|
|||
use litellm_core::http_utils::body::JsonPayload;
|
||||
use std::sync::atomic::{AtomicU64, Ordering};
|
||||
use std::time::{SystemTime, UNIX_EPOCH};
|
||||
|
||||
|
|
@ -15,7 +16,9 @@ pub(crate) struct PreparedOcrCall {
|
|||
}
|
||||
|
||||
#[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)]
|
||||
pub(crate) fn prepare_ocr_call(request: OcrRequest<'_>) -> PreparedOcrCall {
|
||||
pub(crate) fn prepare_ocr_call<D: Into<JsonPayload>>(
|
||||
request: OcrRequest<'_, D>,
|
||||
) -> PreparedOcrCall {
|
||||
let call_id = request
|
||||
.litellm_call_id
|
||||
.map(str::to_string)
|
||||
|
|
@ -49,7 +52,7 @@ pub(crate) fn prepare_ocr_call(request: OcrRequest<'_>) -> PreparedOcrCall {
|
|||
model,
|
||||
custom_llm_provider,
|
||||
litellm_call_id: call_id,
|
||||
document: request.document,
|
||||
document: request.document.into(),
|
||||
api_key: request.api_key.map(str::to_string),
|
||||
api_base: request.api_base.map(str::to_string),
|
||||
extra_headers: request.extra_headers,
|
||||
|
|
|
|||
|
|
@ -1,3 +1,5 @@
|
|||
use litellm_core::http_utils::body::JsonPayload;
|
||||
pub use litellm_core::ocr::types::ProviderOcrRequest;
|
||||
use std::sync::Arc;
|
||||
use std::time::Duration;
|
||||
|
||||
|
|
@ -9,9 +11,9 @@ use crate::integrations::custom_guardrail::CustomGuardrail;
|
|||
use crate::integrations::custom_logger::CustomLogger;
|
||||
use crate::integrations::types::RequestMetadata;
|
||||
|
||||
pub struct OcrRequest<'a> {
|
||||
pub struct OcrRequest<'a, D = Value> {
|
||||
pub model: &'a str,
|
||||
pub document: Value,
|
||||
pub document: D,
|
||||
pub api_key: Option<&'a str>,
|
||||
pub api_base: Option<&'a str>,
|
||||
pub custom_llm_provider: Option<&'a str>,
|
||||
|
|
@ -29,7 +31,7 @@ pub(crate) struct PreparedOcrRequest {
|
|||
pub(crate) model: String,
|
||||
pub(crate) custom_llm_provider: String,
|
||||
pub(crate) litellm_call_id: String,
|
||||
pub(crate) document: Value,
|
||||
pub(crate) document: JsonPayload,
|
||||
pub(crate) api_key: Option<String>,
|
||||
pub(crate) api_base: Option<String>,
|
||||
pub(crate) extra_headers: Option<Map<String, Value>>,
|
||||
|
|
@ -47,12 +49,3 @@ impl CallLifecycleRequest for PreparedOcrRequest {
|
|||
)
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) struct ProviderOcrRequest {
|
||||
pub(crate) model: String,
|
||||
pub(crate) config: &'static dyn OcrProviderConfig,
|
||||
pub(crate) url: String,
|
||||
pub(crate) body: Value,
|
||||
pub(crate) upstream_headers: Vec<(String, String)>,
|
||||
pub(crate) timeout: Option<Duration>,
|
||||
}
|
||||
|
|
|
|||
|
|
@ -6,6 +6,12 @@ license.workspace = true
|
|||
repository.workspace = true
|
||||
|
||||
[dependencies]
|
||||
base64.workspace = true
|
||||
futures-util.workspace = true
|
||||
h2.workspace = true
|
||||
tokio.workspace = true
|
||||
bytes.workspace = true
|
||||
bytestring.workspace = true
|
||||
rand.workspace = true
|
||||
reqwest.workspace = true
|
||||
serde.workspace = true
|
||||
|
|
@ -32,4 +38,5 @@ bedrock-auth = [
|
|||
]
|
||||
|
||||
[dev-dependencies]
|
||||
tokio = { workspace = true, features = ["macros", "rt-multi-thread"] }
|
||||
http = "1"
|
||||
tokio = { workspace = true, features = ["macros", "rt-multi-thread", "io-util"] }
|
||||
|
|
|
|||
|
|
@ -1,7 +1,11 @@
|
|||
use serde_json::Value;
|
||||
|
||||
use crate::constants::AUDIO_TRANSCRIPTION_TIMEOUT_SECS;
|
||||
use crate::error::Error;
|
||||
use crate::http_utils::{http_request, truncate_error_body};
|
||||
use crate::http_utils::body::PreparedJsonBody;
|
||||
use crate::http_utils::replay::{BodySigner, replay_client, send_json};
|
||||
use crate::http_utils::truncate_error_body;
|
||||
use std::time::Duration;
|
||||
|
||||
use super::client::http_client;
|
||||
use super::types::ProviderAudioTranscriptionRequest;
|
||||
|
|
@ -10,19 +14,23 @@ use super::types::ProviderAudioTranscriptionRequest;
|
|||
pub async fn execute_audio_transcription_provider_call(
|
||||
request: ProviderAudioTranscriptionRequest,
|
||||
) -> Result<Value, Error> {
|
||||
let body = serde_json::to_vec(&request.body)
|
||||
.map_err(|error| Error::InvalidRequest(format!("invalid audio request body: {error}")))?;
|
||||
let headers = signed_headers(&request, &body).await?;
|
||||
let mut request_builder = http_client().post(&request.url).body(body);
|
||||
for (key, value) in headers {
|
||||
request_builder = request_builder.header(key, value);
|
||||
}
|
||||
if let Some(duration) = request.timeout {
|
||||
request_builder = request_builder.timeout(duration);
|
||||
}
|
||||
let response = http_request(request_builder)
|
||||
.await
|
||||
.map_err(|error| Error::Network(error.to_string()))?;
|
||||
let signer = request_signer(&request).await?;
|
||||
let body = PreparedJsonBody::new(request.body.clone())?;
|
||||
let response = send_json(
|
||||
if body.is_streamed() {
|
||||
replay_client()?
|
||||
} else {
|
||||
http_client()
|
||||
},
|
||||
&request.url,
|
||||
&body,
|
||||
&request.upstream_headers,
|
||||
request
|
||||
.timeout
|
||||
.unwrap_or(Duration::from_secs(AUDIO_TRANSCRIPTION_TIMEOUT_SECS)),
|
||||
signer.as_deref(),
|
||||
)
|
||||
.await?;
|
||||
let status = response.status();
|
||||
let text = response
|
||||
.text()
|
||||
|
|
@ -43,19 +51,17 @@ pub async fn execute_audio_transcription_provider_call(
|
|||
}
|
||||
|
||||
#[cfg(feature = "bedrock-auth")]
|
||||
async fn signed_headers(
|
||||
async fn request_signer(
|
||||
request: &ProviderAudioTranscriptionRequest,
|
||||
body: &[u8],
|
||||
) -> Result<Vec<(String, String)>, Error> {
|
||||
) -> Result<Option<Box<BodySigner<'static>>>, Error> {
|
||||
use crate::audio_transcription::transformation::AudioTranscriptionAuth;
|
||||
use crate::providers::bedrock::audio_transcription::aws_auth_config;
|
||||
use crate::providers::bedrock::aws_base::{resolve_credentials, sign_bedrock_digest};
|
||||
use std::collections::BTreeMap;
|
||||
use std::time::SystemTime;
|
||||
|
||||
use crate::audio_transcription::transformation::AudioTranscriptionAuth;
|
||||
use crate::providers::bedrock::audio_transcription::aws_auth_config;
|
||||
use crate::providers::bedrock::aws_base::{resolve_credentials, sign_bedrock_post};
|
||||
|
||||
let AudioTranscriptionAuth::AwsSigV4 { region, .. } = &request.auth else {
|
||||
return Ok(request.upstream_headers.clone());
|
||||
return Ok(None);
|
||||
};
|
||||
let env_lookup = |key: &str| std::env::var(key).ok();
|
||||
let credentials = resolve_credentials(
|
||||
|
|
@ -63,29 +69,56 @@ async fn signed_headers(
|
|||
&env_lookup,
|
||||
)
|
||||
.await?;
|
||||
let unsigned: BTreeMap<String, String> = request.upstream_headers.iter().cloned().collect();
|
||||
let signature = sign_bedrock_post(
|
||||
&request.url,
|
||||
body,
|
||||
&unsigned,
|
||||
region,
|
||||
&credentials,
|
||||
SystemTime::now(),
|
||||
)?;
|
||||
Ok(unsigned.into_iter().chain(signature).collect())
|
||||
let region = region.clone();
|
||||
Ok(Some(Box::new(move |url, digest, headers| {
|
||||
let unsigned = headers
|
||||
.iter()
|
||||
.filter(|(name, _)| {
|
||||
!matches!(
|
||||
name.as_str(),
|
||||
"authorization" | "x-amz-date" | "x-amz-security-token" | "host"
|
||||
)
|
||||
})
|
||||
.map(|(name, value)| {
|
||||
Ok((
|
||||
name.to_string(),
|
||||
value
|
||||
.to_str()
|
||||
.map_err(|_| Error::InvalidRequest("invalid signing header".into()))?
|
||||
.to_owned(),
|
||||
))
|
||||
})
|
||||
.collect::<Result<BTreeMap<_, _>, Error>>()?;
|
||||
let signed = sign_bedrock_digest(
|
||||
url.as_str(),
|
||||
digest,
|
||||
&unsigned,
|
||||
®ion,
|
||||
&credentials,
|
||||
SystemTime::now(),
|
||||
)?;
|
||||
let mut result = headers.clone();
|
||||
for (name, value) in signed {
|
||||
result.insert(
|
||||
reqwest::header::HeaderName::from_bytes(name.as_bytes())
|
||||
.map_err(|_| Error::Auth("invalid signing header".into()))?,
|
||||
reqwest::header::HeaderValue::from_str(&value)
|
||||
.map_err(|_| Error::Auth("invalid signing value".into()))?,
|
||||
);
|
||||
}
|
||||
Ok(result)
|
||||
})))
|
||||
}
|
||||
|
||||
#[cfg(not(feature = "bedrock-auth"))]
|
||||
async fn signed_headers(
|
||||
async fn request_signer(
|
||||
request: &ProviderAudioTranscriptionRequest,
|
||||
_body: &[u8],
|
||||
) -> Result<Vec<(String, String)>, Error> {
|
||||
) -> Result<Option<Box<BodySigner<'static>>>, Error> {
|
||||
use crate::audio_transcription::transformation::AudioTranscriptionAuth;
|
||||
|
||||
match request.auth {
|
||||
AudioTranscriptionAuth::AwsSigV4 { .. } => Err(Error::Unsupported(
|
||||
"AWS SigV4 requires the bedrock-auth feature",
|
||||
)),
|
||||
AudioTranscriptionAuth::Bearer => Ok(request.upstream_headers.clone()),
|
||||
AudioTranscriptionAuth::Bearer => Ok(None),
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -1,4 +1,5 @@
|
|||
use crate::Error;
|
||||
use crate::http_utils::body::JsonPayload;
|
||||
mod client;
|
||||
mod handler;
|
||||
mod prepare;
|
||||
|
|
@ -12,7 +13,9 @@ pub use prepare::prepare_audio_transcription_provider_call;
|
|||
pub use types::{AudioTranscriptionRequest, ProviderAudioTranscriptionRequest};
|
||||
|
||||
#[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)]
|
||||
pub async fn audio_transcription(request: AudioTranscriptionRequest<'_>) -> Result<Value, Error> {
|
||||
pub async fn audio_transcription<A: Into<JsonPayload>>(
|
||||
request: AudioTranscriptionRequest<'_, A>,
|
||||
) -> Result<Value, Error> {
|
||||
execute_audio_transcription_provider_call(prepare_audio_transcription_provider_call(request)?)
|
||||
.await
|
||||
}
|
||||
|
|
|
|||
|
|
@ -1,4 +1,5 @@
|
|||
use crate::error::Error;
|
||||
use crate::http_utils::body::JsonPayload;
|
||||
use crate::http_utils::{has_header, string_headers};
|
||||
#[cfg(feature = "bedrock-auth")]
|
||||
use crate::providers::bedrock::audio_transcription::BEDROCK_AUDIO_TRANSCRIPTION_CONFIG;
|
||||
|
|
@ -18,8 +19,8 @@ fn provider_config(provider: &str) -> Option<&'static dyn AudioTranscriptionProv
|
|||
}
|
||||
|
||||
#[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)]
|
||||
pub fn prepare_audio_transcription_provider_call(
|
||||
request: AudioTranscriptionRequest<'_>,
|
||||
pub fn prepare_audio_transcription_provider_call<A: Into<JsonPayload>>(
|
||||
request: AudioTranscriptionRequest<'_, A>,
|
||||
) -> Result<ProviderAudioTranscriptionRequest, Error> {
|
||||
let provider_info = get_custom_llm_provider(request.model, request.custom_llm_provider)
|
||||
.or_else(|| {
|
||||
|
|
@ -58,7 +59,7 @@ pub fn prepare_audio_transcription_provider_call(
|
|||
)?;
|
||||
let filtered_params = config.map_transcription_params(&request.optional_params);
|
||||
let transformed =
|
||||
config.transform_transcription_request(&model, request.audio, filtered_params)?;
|
||||
config.transform_transcription_payload(&model, request.audio.into(), filtered_params)?;
|
||||
Ok(ProviderAudioTranscriptionRequest {
|
||||
model,
|
||||
custom_llm_provider: provider_info.custom_llm_provider.to_string(),
|
||||
|
|
|
|||
|
|
@ -1,4 +1,4 @@
|
|||
use std::io::{Read, Write};
|
||||
use std::io::{BufRead, BufReader, Read, Write};
|
||||
use std::net::TcpListener;
|
||||
use std::thread;
|
||||
|
||||
|
|
@ -13,11 +13,25 @@ async fn bedrock_request_is_signed_and_contains_audio() {
|
|||
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);
|
||||
let mut reader = BufReader::new(&mut stream);
|
||||
let mut headers = String::new();
|
||||
loop {
|
||||
let mut line = String::new();
|
||||
reader.read_line(&mut line).expect("headers");
|
||||
if line == "\r\n" {
|
||||
break;
|
||||
}
|
||||
headers.push_str(&line);
|
||||
}
|
||||
let length = headers
|
||||
.lines()
|
||||
.find_map(|line| line.strip_prefix("content-length: "))
|
||||
.expect("content length")
|
||||
.parse::<usize>()
|
||||
.unwrap();
|
||||
let mut body = vec![0; length];
|
||||
reader.read_exact(&mut body).expect("complete request body");
|
||||
let request = headers + &String::from_utf8(body).unwrap();
|
||||
assert!(request.contains("POST /model/mistral.voxtral-mini-3b-2507/converse"));
|
||||
assert!(request.contains("authorization: AWS4-HMAC-SHA256"));
|
||||
assert!(request.contains("x-amz-date:"));
|
||||
|
|
|
|||
|
|
@ -1,4 +1,5 @@
|
|||
use crate::Error;
|
||||
use crate::http_utils::body::JsonPayload;
|
||||
use serde_json::{Map, Value};
|
||||
|
||||
use super::types::{AudioTranscriptionRequestData, AudioTranscriptionResponseData};
|
||||
|
|
@ -32,6 +33,15 @@ pub trait AudioTranscriptionProviderConfig: Sync {
|
|||
model: &str,
|
||||
audio: Value,
|
||||
optional_params: Map<String, Value>,
|
||||
) -> Result<AudioTranscriptionRequestData, Error> {
|
||||
self.transform_transcription_payload(model, audio.into(), optional_params)
|
||||
}
|
||||
|
||||
fn transform_transcription_payload(
|
||||
&self,
|
||||
model: &str,
|
||||
audio: JsonPayload,
|
||||
optional_params: Map<String, Value>,
|
||||
) -> Result<AudioTranscriptionRequestData, Error>;
|
||||
|
||||
fn transform_transcription_response(
|
||||
|
|
|
|||
|
|
@ -1,3 +1,4 @@
|
|||
use crate::http_utils::body::JsonPayload;
|
||||
use std::time::Duration;
|
||||
|
||||
use serde::{Deserialize, Serialize};
|
||||
|
|
@ -5,9 +6,9 @@ use serde_json::{Map, Value};
|
|||
|
||||
use super::transformation::{AudioTranscriptionAuth, AudioTranscriptionProviderConfig};
|
||||
|
||||
pub struct AudioTranscriptionRequest<'a> {
|
||||
pub struct AudioTranscriptionRequest<'a, A = Value> {
|
||||
pub model: &'a str,
|
||||
pub audio: Value,
|
||||
pub audio: A,
|
||||
pub api_key: Option<&'a str>,
|
||||
pub api_base: Option<&'a str>,
|
||||
pub custom_llm_provider: Option<&'a str>,
|
||||
|
|
@ -22,7 +23,7 @@ pub struct ProviderAudioTranscriptionRequest {
|
|||
pub(super) custom_llm_provider: String,
|
||||
pub(super) config: &'static dyn AudioTranscriptionProviderConfig,
|
||||
pub(super) url: String,
|
||||
pub(super) body: Value,
|
||||
pub(super) body: JsonPayload,
|
||||
pub(super) upstream_headers: Vec<(String, String)>,
|
||||
pub(super) auth: AudioTranscriptionAuth,
|
||||
#[cfg(feature = "bedrock-auth")]
|
||||
|
|
@ -43,18 +44,21 @@ impl ProviderAudioTranscriptionRequest {
|
|||
&self.url
|
||||
}
|
||||
|
||||
pub fn body(&self) -> &Value {
|
||||
pub fn body(&self) -> &JsonPayload {
|
||||
&self.body
|
||||
}
|
||||
|
||||
pub fn with_body(self, body: Value) -> Self {
|
||||
Self { body, ..self }
|
||||
Self {
|
||||
body: body.into(),
|
||||
..self
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
|
||||
pub struct AudioTranscriptionRequestData {
|
||||
pub body: Value,
|
||||
pub body: JsonPayload,
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
|
||||
|
|
|
|||
|
|
@ -41,3 +41,12 @@ pub const CHAT_COMPLETION_OBJECT: &str = "chat.completion";
|
|||
/// `litellm/litellm_core_utils/prompt_templates/factory.py`.
|
||||
pub const EMPTY_TEXT_PLACEHOLDER: &str =
|
||||
"[System: Empty message content sanitised to satisfy protocol]";
|
||||
|
||||
pub const JSON_BODY_CHUNK_BYTES: usize = 64 * 1024;
|
||||
pub const JSON_BODY_MAX_REDIRECTS: usize = 10;
|
||||
pub const JSON_BODY_PROTOCOL_RETRIES: usize = 2;
|
||||
|
||||
pub const JSON_PAYLOAD_MAX_DEPTH: usize = 128;
|
||||
|
||||
pub const OCR_TIMEOUT_SECS: u64 = 600;
|
||||
pub const AZURE_DOCUMENT_INTELLIGENCE_POLL_TIMEOUT_SECS: u64 = 120;
|
||||
|
|
|
|||
|
|
@ -1,5 +1,8 @@
|
|||
//! Header and upstream-body helpers shared by every route module.
|
||||
|
||||
pub mod body;
|
||||
pub mod replay;
|
||||
|
||||
use serde_json::{Map, Value};
|
||||
|
||||
use crate::constants::UPSTREAM_ERROR_BODY_MAX_CHARS;
|
||||
|
|
|
|||
219
litellm-rust/crates/core/src/http_utils/body/mod.rs
Normal file
219
litellm-rust/crates/core/src/http_utils/body/mod.rs
Normal file
|
|
@ -0,0 +1,219 @@
|
|||
use std::sync::Arc;
|
||||
|
||||
use base64::Engine;
|
||||
use bytes::Bytes;
|
||||
use sha2::{Digest, Sha256};
|
||||
|
||||
use crate::Error;
|
||||
use crate::constants::JSON_BODY_CHUNK_BYTES;
|
||||
mod value;
|
||||
pub use value::{JsonPayload, SharedText};
|
||||
|
||||
#[derive(Clone)]
|
||||
enum Part {
|
||||
Bytes(Bytes),
|
||||
Quoted(SharedText),
|
||||
Base64(Bytes),
|
||||
}
|
||||
|
||||
#[derive(Clone)]
|
||||
pub struct PreparedJsonBody {
|
||||
parts: Arc<[Part]>,
|
||||
length: u64,
|
||||
streamed: bool,
|
||||
}
|
||||
|
||||
impl PreparedJsonBody {
|
||||
pub fn new(payload: JsonPayload) -> Result<Self, Error> {
|
||||
if !payload.contains_media() {
|
||||
let bytes = serde_json::to_vec(&payload)
|
||||
.map_err(|_| Error::InvalidRequest("invalid JSON request body".into()))?;
|
||||
return Ok(Self::buffered(bytes.into()));
|
||||
}
|
||||
Self::streamed(payload)
|
||||
}
|
||||
|
||||
pub fn streamed(payload: JsonPayload) -> Result<Self, Error> {
|
||||
let mut parts = Vec::new();
|
||||
append_payload(payload, &mut parts)?;
|
||||
let mut body = Self {
|
||||
parts: parts.into(),
|
||||
length: 0,
|
||||
streamed: true,
|
||||
};
|
||||
body.length = body.chunks().try_fold(0_u64, |size, chunk| {
|
||||
size.checked_add(chunk.len() as u64).ok_or_else(size_error)
|
||||
})?;
|
||||
Ok(body)
|
||||
}
|
||||
|
||||
pub fn buffered(bytes: Bytes) -> Self {
|
||||
Self {
|
||||
length: bytes.len() as u64,
|
||||
parts: Arc::from([Part::Bytes(bytes)]),
|
||||
streamed: false,
|
||||
}
|
||||
}
|
||||
|
||||
pub fn content_length(&self) -> u64 {
|
||||
self.length
|
||||
}
|
||||
pub fn is_streamed(&self) -> bool {
|
||||
self.streamed
|
||||
}
|
||||
|
||||
pub fn chunks(&self) -> impl Iterator<Item = Bytes> + Send + 'static {
|
||||
BodyChunks {
|
||||
parts: self.parts.clone(),
|
||||
part: 0,
|
||||
offset: 0,
|
||||
opened: false,
|
||||
}
|
||||
}
|
||||
|
||||
pub fn sha256(&self) -> String {
|
||||
let mut digest = Sha256::new();
|
||||
for chunk in self.chunks() {
|
||||
digest.update(chunk);
|
||||
}
|
||||
format!("{:x}", digest.finalize())
|
||||
}
|
||||
|
||||
pub(super) fn buffered_bytes(&self) -> Option<Bytes> {
|
||||
match self.parts.first() {
|
||||
Some(Part::Bytes(bytes)) if !self.streamed => Some(bytes.clone()),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn size_error() -> Error {
|
||||
Error::InvalidRequest("JSON request body is too large".into())
|
||||
}
|
||||
|
||||
fn append_payload(payload: JsonPayload, parts: &mut Vec<Part>) -> Result<(), Error> {
|
||||
match payload {
|
||||
JsonPayload::String(text) => parts.push(Part::Quoted(text)),
|
||||
JsonPayload::Base64(bytes) => {
|
||||
parts.push(Part::Bytes(Bytes::from_static(b"\"")));
|
||||
parts.push(Part::Base64(bytes));
|
||||
parts.push(Part::Bytes(Bytes::from_static(b"\"")));
|
||||
}
|
||||
JsonPayload::Array(items) => {
|
||||
parts.push(Part::Bytes(Bytes::from_static(b"[")));
|
||||
for (index, item) in items.into_iter().enumerate() {
|
||||
if index > 0 {
|
||||
parts.push(Part::Bytes(Bytes::from_static(b",")));
|
||||
}
|
||||
append_payload(item, parts)?;
|
||||
}
|
||||
parts.push(Part::Bytes(Bytes::from_static(b"]")));
|
||||
}
|
||||
JsonPayload::Object(fields) => {
|
||||
parts.push(Part::Bytes(Bytes::from_static(b"{")));
|
||||
for (index, (key, value)) in fields.into_iter().enumerate() {
|
||||
if index > 0 {
|
||||
parts.push(Part::Bytes(Bytes::from_static(b",")));
|
||||
}
|
||||
parts.push(Part::Quoted(key.into()));
|
||||
parts.push(Part::Bytes(Bytes::from_static(b":")));
|
||||
append_payload(value, parts)?;
|
||||
}
|
||||
parts.push(Part::Bytes(Bytes::from_static(b"}")));
|
||||
}
|
||||
scalar => parts.push(Part::Bytes(
|
||||
serde_json::to_vec(&scalar)
|
||||
.map_err(|_| Error::InvalidRequest("invalid JSON value".into()))?
|
||||
.into(),
|
||||
)),
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn needs_escape(byte: u8) -> bool {
|
||||
byte < 32 || byte == b'"' || byte == b'\\'
|
||||
}
|
||||
|
||||
struct BodyChunks {
|
||||
parts: Arc<[Part]>,
|
||||
part: usize,
|
||||
offset: usize,
|
||||
opened: bool,
|
||||
}
|
||||
|
||||
impl Iterator for BodyChunks {
|
||||
type Item = Bytes;
|
||||
|
||||
fn next(&mut self) -> Option<Bytes> {
|
||||
loop {
|
||||
let part = self.parts.get(self.part)?;
|
||||
match part {
|
||||
Part::Bytes(bytes) if self.offset < bytes.len() => {
|
||||
let end = self
|
||||
.offset
|
||||
.saturating_add(JSON_BODY_CHUNK_BYTES)
|
||||
.min(bytes.len());
|
||||
let chunk = bytes.slice(self.offset..end);
|
||||
self.offset = end;
|
||||
return Some(chunk);
|
||||
}
|
||||
Part::Base64(bytes) if self.offset < bytes.len() => {
|
||||
let end = self
|
||||
.offset
|
||||
.saturating_add(JSON_BODY_CHUNK_BYTES / 4 * 3)
|
||||
.min(bytes.len());
|
||||
let chunk =
|
||||
base64::engine::general_purpose::STANDARD.encode(&bytes[self.offset..end]);
|
||||
self.offset = end;
|
||||
return Some(chunk.into());
|
||||
}
|
||||
Part::Quoted(_) if !self.opened => {
|
||||
self.opened = true;
|
||||
return Some(Bytes::from_static(b"\""));
|
||||
}
|
||||
Part::Quoted(text) if self.offset < text.bytes().len() => {
|
||||
let bytes = text.bytes();
|
||||
let limit = self
|
||||
.offset
|
||||
.saturating_add(JSON_BODY_CHUNK_BYTES)
|
||||
.min(bytes.len());
|
||||
let raw_length = bytes[self.offset..limit]
|
||||
.iter()
|
||||
.position(|byte| needs_escape(*byte))
|
||||
.unwrap_or(limit - self.offset);
|
||||
if raw_length > 0 {
|
||||
let start = self.offset;
|
||||
self.offset += raw_length;
|
||||
return Some(bytes.slice(start..self.offset));
|
||||
}
|
||||
let start = self.offset;
|
||||
let end = (start + (JSON_BODY_CHUNK_BYTES - 2) / 6).min(bytes.len());
|
||||
let count = bytes[start..end]
|
||||
.iter()
|
||||
.take_while(|byte| needs_escape(**byte))
|
||||
.count();
|
||||
self.offset += count;
|
||||
let quoted = serde_json::to_vec(&text.as_str()[start..self.offset])
|
||||
.unwrap_or_else(|_| {
|
||||
unreachable!("serializing a string into Vec cannot fail")
|
||||
});
|
||||
let output = Bytes::from(quoted);
|
||||
return Some(output.slice(1..output.len() - 1));
|
||||
}
|
||||
Part::Quoted(_) => {
|
||||
self.part += 1;
|
||||
self.offset = 0;
|
||||
self.opened = false;
|
||||
return Some(Bytes::from_static(b"\""));
|
||||
}
|
||||
_ => {}
|
||||
}
|
||||
self.part += 1;
|
||||
self.offset = 0;
|
||||
self.opened = false;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests;
|
||||
138
litellm-rust/crates/core/src/http_utils/body/tests.rs
Normal file
138
litellm-rust/crates/core/src/http_utils/body/tests.rs
Normal file
|
|
@ -0,0 +1,138 @@
|
|||
use std::collections::BTreeMap;
|
||||
use std::sync::Arc;
|
||||
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn streams_match_serde_including_escaping_and_base64_boundaries() {
|
||||
let text = format!(
|
||||
"{}{}é💙\\\"\n",
|
||||
"x".repeat(JSON_BODY_CHUNK_BYTES - 1),
|
||||
(0_u8..32).map(char::from).collect::<String>()
|
||||
);
|
||||
for length in [0, 1, 2, 3, 49151, 49152, 49153, 131072] {
|
||||
let payload = JsonPayload::Object(BTreeMap::from([
|
||||
("image".into(), JsonPayload::String(text.clone().into())),
|
||||
(
|
||||
"audio".into(),
|
||||
JsonPayload::Base64(vec![123; length].into()),
|
||||
),
|
||||
(
|
||||
"nested".into(),
|
||||
serde_json::json!([null, true, 2.5, {"key": "value"}]).into(),
|
||||
),
|
||||
]));
|
||||
let expected = serde_json::to_vec(&payload).unwrap();
|
||||
let body = PreparedJsonBody::streamed(payload).unwrap();
|
||||
let actual: Vec<u8> = body.chunks().flat_map(|bytes| bytes.to_vec()).collect();
|
||||
assert_eq!(actual, expected);
|
||||
assert_eq!(body.content_length(), expected.len() as u64);
|
||||
assert_eq!(body.sha256(), format!("{:x}", Sha256::digest(&expected)));
|
||||
assert!(
|
||||
body.chunks()
|
||||
.all(|chunk| chunk.len() <= JSON_BODY_CHUNK_BYTES)
|
||||
);
|
||||
assert_eq!(body.chunks().flatten().collect::<Vec<_>>(), expected);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn shared_payload_slices_and_replays_retain_the_owner_without_copying() {
|
||||
struct Owner(Arc<Vec<u8>>);
|
||||
impl AsRef<[u8]> for Owner {
|
||||
fn as_ref(&self) -> &[u8] {
|
||||
&self.0
|
||||
}
|
||||
}
|
||||
let owner = Arc::new(vec![b'A'; 1024 * 1024]);
|
||||
let pointer = owner.as_ptr();
|
||||
let bytes = Bytes::from_owner(Owner(owner.clone()));
|
||||
let text = SharedText::new(bytes).unwrap();
|
||||
let slice = text.slice(1..text.bytes().len()).unwrap();
|
||||
assert_eq!(slice.bytes().as_ptr(), pointer.wrapping_add(1));
|
||||
let body = PreparedJsonBody::streamed(JsonPayload::String(text)).unwrap();
|
||||
drop(slice);
|
||||
assert_eq!(Arc::strong_count(&owner), 2);
|
||||
for _ in 0..2 {
|
||||
let mut chunks = body.chunks();
|
||||
assert_eq!(chunks.next().unwrap(), "\"");
|
||||
assert_eq!(chunks.next().unwrap().as_ptr(), pointer);
|
||||
}
|
||||
drop(body);
|
||||
assert_eq!(Arc::strong_count(&owner), 1);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn non_media_remains_buffered_and_data_uris_stream() {
|
||||
assert!(
|
||||
!PreparedJsonBody::new(serde_json::json!({"text":"hello"}).into())
|
||||
.unwrap()
|
||||
.is_streamed()
|
||||
);
|
||||
assert!(
|
||||
PreparedJsonBody::new(
|
||||
serde_json::json!({"document_url":"data:application/pdf;base64,AA=="}).into()
|
||||
)
|
||||
.unwrap()
|
||||
.is_streamed()
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn serde_escaping_stays_bounded_for_long_control_runs_and_unicode_splits() {
|
||||
for text in [
|
||||
"\u{0001}".repeat(JSON_BODY_CHUNK_BYTES * 3),
|
||||
format!("{}é💙\n\\\"", "a".repeat(JSON_BODY_CHUNK_BYTES - 1)),
|
||||
] {
|
||||
let expected = serde_json::to_vec(&text).unwrap();
|
||||
let body = PreparedJsonBody::streamed(JsonPayload::String(text.into())).unwrap();
|
||||
assert_eq!(body.content_length(), expected.len() as u64);
|
||||
assert!(
|
||||
body.chunks()
|
||||
.all(|chunk| chunk.len() <= JSON_BODY_CHUNK_BYTES)
|
||||
);
|
||||
assert_eq!(body.chunks().flatten().collect::<Vec<_>>(), expected);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn typed_messages_keep_nested_media_in_outgoing_slices() {
|
||||
use crate::messages::transformation::AnthropicMessagesProviderConfig;
|
||||
use crate::messages::types::{AnthropicMessagesRequest, MessageContent};
|
||||
use crate::providers::azure_ai::messages::transformation::AZURE_ANTHROPIC_MESSAGES_CONFIG;
|
||||
let bytes = Bytes::from(vec![b'A'; JSON_BODY_CHUNK_BYTES * 3]);
|
||||
let mut request: AnthropicMessagesRequest = serde_json::from_value(serde_json::json!({"model":"model","messages":[{"role":"system","content":[{"type":"tool_result","content":[]}]}]})).unwrap();
|
||||
let MessageContent::Blocks(blocks) = &mut request.messages[0].content else {
|
||||
panic!("blocks")
|
||||
};
|
||||
blocks[0].extra.insert(
|
||||
"content".into(),
|
||||
JsonPayload::Array(vec![JsonPayload::object([
|
||||
("type", "image".into()),
|
||||
(
|
||||
"source",
|
||||
JsonPayload::object([
|
||||
("type", "base64".into()),
|
||||
(
|
||||
"data",
|
||||
JsonPayload::String(SharedText::new(bytes.clone()).unwrap()),
|
||||
),
|
||||
]),
|
||||
),
|
||||
])]),
|
||||
);
|
||||
let request = AZURE_ANTHROPIC_MESSAGES_CONFIG
|
||||
.transform_request(request)
|
||||
.unwrap();
|
||||
let expected = serde_json::to_vec(&request).unwrap();
|
||||
let body = PreparedJsonBody::new(request.into_payload().unwrap()).unwrap();
|
||||
let emitted = body.chunks().flatten().collect::<Vec<_>>();
|
||||
assert_eq!(
|
||||
serde_json::from_slice::<serde_json::Value>(&emitted).unwrap(),
|
||||
serde_json::from_slice::<serde_json::Value>(&expected).unwrap()
|
||||
);
|
||||
assert!(
|
||||
body.chunks()
|
||||
.any(|chunk| chunk.as_ptr() == bytes.as_ptr() && chunk.len() == JSON_BODY_CHUNK_BYTES)
|
||||
);
|
||||
}
|
||||
261
litellm-rust/crates/core/src/http_utils/body/value.rs
Normal file
261
litellm-rust/crates/core/src/http_utils/body/value.rs
Normal file
|
|
@ -0,0 +1,261 @@
|
|||
use std::collections::BTreeMap;
|
||||
use std::fmt;
|
||||
use std::ops::{Index, Range};
|
||||
|
||||
use base64::Engine;
|
||||
use bytes::Bytes;
|
||||
use bytestring::ByteString;
|
||||
use serde::{Deserialize, Deserializer, Serialize, Serializer};
|
||||
use serde_json::Value;
|
||||
|
||||
use crate::Error;
|
||||
|
||||
#[derive(Clone, PartialEq, Eq)]
|
||||
pub struct SharedText(ByteString);
|
||||
|
||||
impl SharedText {
|
||||
pub fn new(bytes: Bytes) -> Result<Self, std::str::Utf8Error> {
|
||||
ByteString::try_from(bytes).map(Self)
|
||||
}
|
||||
|
||||
pub fn as_str(&self) -> &str {
|
||||
&self.0
|
||||
}
|
||||
|
||||
pub fn bytes(&self) -> &Bytes {
|
||||
self.0.as_bytes()
|
||||
}
|
||||
|
||||
pub fn slice(&self, range: Range<usize>) -> Result<Self, Error> {
|
||||
if self.as_str().get(range.clone()).is_none() {
|
||||
return Err(Error::InvalidRequest("invalid shared text range".into()));
|
||||
}
|
||||
Ok(Self(self.0.slice_ref(&self.0[range])))
|
||||
}
|
||||
}
|
||||
|
||||
impl From<String> for SharedText {
|
||||
fn from(value: String) -> Self {
|
||||
Self(ByteString::from(value))
|
||||
}
|
||||
}
|
||||
|
||||
impl fmt::Debug for SharedText {
|
||||
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
|
||||
f.debug_struct("SharedText")
|
||||
.field("len", &self.0.len())
|
||||
.finish()
|
||||
}
|
||||
}
|
||||
|
||||
impl Serialize for SharedText {
|
||||
fn serialize<S: Serializer>(&self, serializer: S) -> Result<S::Ok, S::Error> {
|
||||
serializer.serialize_str(self.as_str())
|
||||
}
|
||||
}
|
||||
|
||||
impl<'de> Deserialize<'de> for SharedText {
|
||||
fn deserialize<D: Deserializer<'de>>(deserializer: D) -> Result<Self, D::Error> {
|
||||
String::deserialize(deserializer).map(Self::from)
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Clone, PartialEq)]
|
||||
pub enum JsonPayload {
|
||||
Null,
|
||||
Bool(bool),
|
||||
Number(serde_json::Number),
|
||||
String(SharedText),
|
||||
Base64(Bytes),
|
||||
Array(Vec<Self>),
|
||||
Object(BTreeMap<String, Self>),
|
||||
}
|
||||
|
||||
impl JsonPayload {
|
||||
pub fn object<const N: usize>(fields: [(&str, Self); N]) -> Self {
|
||||
Self::Object(
|
||||
fields
|
||||
.into_iter()
|
||||
.map(|(key, value)| (key.to_owned(), value))
|
||||
.collect(),
|
||||
)
|
||||
}
|
||||
|
||||
pub fn type_name(&self) -> &'static str {
|
||||
match self {
|
||||
Self::Null => "null",
|
||||
Self::Bool(_) => "bool",
|
||||
Self::Number(_) => "number",
|
||||
Self::String(_) => "string",
|
||||
Self::Base64(_) => "bytes",
|
||||
Self::Array(_) => "array",
|
||||
Self::Object(_) => "object",
|
||||
}
|
||||
}
|
||||
|
||||
pub fn as_str(&self) -> Option<&str> {
|
||||
self.as_text().map(SharedText::as_str)
|
||||
}
|
||||
|
||||
pub fn as_text(&self) -> Option<&SharedText> {
|
||||
match self {
|
||||
Self::String(value) => Some(value),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
|
||||
pub fn as_object(&self) -> Option<&BTreeMap<String, Self>> {
|
||||
match self {
|
||||
Self::Object(value) => Some(value),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
|
||||
pub fn as_array(&self) -> Option<&Vec<Self>> {
|
||||
match self {
|
||||
Self::Array(value) => Some(value),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
|
||||
pub fn is_object(&self) -> bool {
|
||||
self.as_object().is_some()
|
||||
}
|
||||
|
||||
pub fn get(&self, key: &str) -> Option<&Self> {
|
||||
self.as_object()?.get(key)
|
||||
}
|
||||
|
||||
pub fn into_object(self) -> Result<BTreeMap<String, Self>, Error> {
|
||||
match self {
|
||||
Self::Object(value) => Ok(value),
|
||||
other => Err(Error::InvalidType {
|
||||
expected: "object",
|
||||
actual: other.type_name(),
|
||||
}),
|
||||
}
|
||||
}
|
||||
|
||||
pub fn contains_media(&self) -> bool {
|
||||
match self {
|
||||
Self::Base64(_) => true,
|
||||
Self::String(value) => value.as_str().starts_with("data:"),
|
||||
Self::Array(items) => items.iter().any(Self::contains_media),
|
||||
Self::Object(fields) => {
|
||||
fields.get("type").and_then(Self::as_str) == Some("base64")
|
||||
|| fields.contains_key("audio")
|
||||
|| fields.contains_key("base64Source")
|
||||
|| fields.values().any(Self::contains_media)
|
||||
}
|
||||
_ => false,
|
||||
}
|
||||
}
|
||||
|
||||
pub fn materialize(&self) -> Value {
|
||||
match self {
|
||||
Self::Null => Value::Null,
|
||||
Self::Bool(value) => Value::Bool(*value),
|
||||
Self::Number(value) => Value::Number(value.clone()),
|
||||
Self::String(value) => Value::String(value.as_str().to_owned()),
|
||||
Self::Base64(value) => {
|
||||
Value::String(base64::engine::general_purpose::STANDARD.encode(value))
|
||||
}
|
||||
Self::Array(items) => Value::Array(items.iter().map(Self::materialize).collect()),
|
||||
Self::Object(fields) => Value::Object(
|
||||
fields
|
||||
.iter()
|
||||
.map(|(key, value)| (key.clone(), value.materialize()))
|
||||
.collect(),
|
||||
),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl fmt::Debug for JsonPayload {
|
||||
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
|
||||
f.debug_struct("JsonPayload")
|
||||
.field("kind", &self.type_name())
|
||||
.finish()
|
||||
}
|
||||
}
|
||||
|
||||
impl From<Value> for JsonPayload {
|
||||
fn from(value: Value) -> Self {
|
||||
match value {
|
||||
Value::Null => Self::Null,
|
||||
Value::Bool(value) => Self::Bool(value),
|
||||
Value::Number(value) => Self::Number(value),
|
||||
Value::String(value) => Self::String(value.into()),
|
||||
Value::Array(items) => Self::Array(items.into_iter().map(Self::from).collect()),
|
||||
Value::Object(fields) => Self::Object(
|
||||
fields
|
||||
.into_iter()
|
||||
.map(|(key, value)| (key, value.into()))
|
||||
.collect(),
|
||||
),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl From<String> for JsonPayload {
|
||||
fn from(value: String) -> Self {
|
||||
Self::String(value.into())
|
||||
}
|
||||
}
|
||||
|
||||
impl From<&str> for JsonPayload {
|
||||
fn from(value: &str) -> Self {
|
||||
value.to_owned().into()
|
||||
}
|
||||
}
|
||||
|
||||
impl Serialize for JsonPayload {
|
||||
fn serialize<S: Serializer>(&self, serializer: S) -> Result<S::Ok, S::Error> {
|
||||
match self {
|
||||
Self::Null => serializer.serialize_unit(),
|
||||
Self::Bool(value) => serializer.serialize_bool(*value),
|
||||
Self::Number(value) => value.serialize(serializer),
|
||||
Self::String(value) => value.serialize(serializer),
|
||||
Self::Base64(value) => serializer.collect_str(&base64::display::Base64Display::new(
|
||||
value,
|
||||
&base64::engine::general_purpose::STANDARD,
|
||||
)),
|
||||
Self::Array(items) => items.serialize(serializer),
|
||||
Self::Object(fields) => fields.serialize(serializer),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl<'de> Deserialize<'de> for JsonPayload {
|
||||
fn deserialize<D: Deserializer<'de>>(deserializer: D) -> Result<Self, D::Error> {
|
||||
Value::deserialize(deserializer).map(Self::from)
|
||||
}
|
||||
}
|
||||
|
||||
impl Index<&str> for JsonPayload {
|
||||
type Output = Self;
|
||||
fn index(&self, key: &str) -> &Self {
|
||||
self.get(key).unwrap_or(&Self::Null)
|
||||
}
|
||||
}
|
||||
|
||||
impl Index<usize> for JsonPayload {
|
||||
type Output = Self;
|
||||
fn index(&self, index: usize) -> &Self {
|
||||
self.as_array()
|
||||
.and_then(|items| items.get(index))
|
||||
.unwrap_or(&Self::Null)
|
||||
}
|
||||
}
|
||||
|
||||
impl PartialEq<Value> for JsonPayload {
|
||||
fn eq(&self, other: &Value) -> bool {
|
||||
self.materialize() == *other
|
||||
}
|
||||
}
|
||||
|
||||
impl PartialEq<&str> for JsonPayload {
|
||||
fn eq(&self, other: &&str) -> bool {
|
||||
self.as_str() == Some(*other)
|
||||
}
|
||||
}
|
||||
225
litellm-rust/crates/core/src/http_utils/replay.rs
Normal file
225
litellm-rust/crates/core/src/http_utils/replay.rs
Normal file
|
|
@ -0,0 +1,225 @@
|
|||
use std::error::Error as _;
|
||||
use std::sync::OnceLock;
|
||||
use std::time::{Duration, Instant};
|
||||
|
||||
use reqwest::header::{CONTENT_LENGTH, CONTENT_TYPE, HeaderMap, HeaderName, HeaderValue};
|
||||
use reqwest::{Client, Method, Response, StatusCode, Url};
|
||||
|
||||
use crate::Error;
|
||||
use crate::constants::{
|
||||
JSON_BODY_MAX_REDIRECTS, JSON_BODY_PROTOCOL_RETRIES, MESSAGES_CONNECT_TIMEOUT_SECS,
|
||||
};
|
||||
|
||||
use super::body::PreparedJsonBody;
|
||||
|
||||
pub type BodySigner<'a> =
|
||||
dyn Fn(&Url, &str, &HeaderMap) -> Result<HeaderMap, Error> + Send + Sync + 'a;
|
||||
|
||||
pub fn replay_client() -> Result<&'static Client, Error> {
|
||||
static CLIENT: OnceLock<Result<Client, String>> = OnceLock::new();
|
||||
CLIENT
|
||||
.get_or_init(|| {
|
||||
Client::builder()
|
||||
.connect_timeout(Duration::from_secs(MESSAGES_CONNECT_TIMEOUT_SECS))
|
||||
.redirect(reqwest::redirect::Policy::none())
|
||||
.retry(reqwest::retry::never())
|
||||
.build()
|
||||
.map_err(|error| error.to_string())
|
||||
})
|
||||
.as_ref()
|
||||
.map_err(|error| Error::Network(error.clone()))
|
||||
}
|
||||
|
||||
pub async fn send_json(
|
||||
client: &Client,
|
||||
url: &str,
|
||||
body: &PreparedJsonBody,
|
||||
headers: &[(String, String)],
|
||||
timeout: Duration,
|
||||
signer: Option<&BodySigner<'_>>,
|
||||
) -> Result<Response, Error> {
|
||||
let started = Instant::now();
|
||||
let initial_url =
|
||||
Url::parse(url).map_err(|_| Error::InvalidRequest("invalid provider URL".into()))?;
|
||||
let digest = signer.map(|_| body.sha256());
|
||||
let headers = headers
|
||||
.iter()
|
||||
.map(|(key, value)| {
|
||||
let key = HeaderName::from_bytes(key.as_bytes())
|
||||
.map_err(|_| Error::InvalidRequest("invalid request header name".into()))?;
|
||||
let value = HeaderValue::from_str(value)
|
||||
.map_err(|_| Error::InvalidRequest("invalid request header value".into()))?;
|
||||
Ok((key, value))
|
||||
})
|
||||
.collect::<Result<HeaderMap, Error>>()?;
|
||||
if !body.is_streamed() {
|
||||
let headers = signed_headers(signer, &initial_url, digest.as_deref(), &headers)?;
|
||||
return client
|
||||
.post(initial_url)
|
||||
.headers(headers)
|
||||
.body(request_body(body))
|
||||
.timeout(remaining(started, timeout)?)
|
||||
.send()
|
||||
.await
|
||||
.map_err(network_error);
|
||||
}
|
||||
let mut url = initial_url;
|
||||
let mut headers = headers;
|
||||
let mut signing = signer;
|
||||
let mut method = Method::POST;
|
||||
let mut redirects = 0;
|
||||
let mut retries = 0;
|
||||
loop {
|
||||
let mut attempt_headers = headers.clone();
|
||||
let request = if method == Method::POST {
|
||||
attempt_headers.remove(reqwest::header::TRANSFER_ENCODING);
|
||||
attempt_headers.insert(CONTENT_LENGTH, HeaderValue::from(body.content_length()));
|
||||
attempt_headers
|
||||
.entry(CONTENT_TYPE)
|
||||
.or_insert(HeaderValue::from_static("application/json"));
|
||||
client
|
||||
.request(method.clone(), url.clone())
|
||||
.body(request_body(body))
|
||||
} else {
|
||||
client.request(method.clone(), url.clone())
|
||||
};
|
||||
let attempt_headers = if method == Method::POST {
|
||||
signed_headers(signing, &url, digest.as_deref(), &attempt_headers)?
|
||||
} else {
|
||||
attempt_headers
|
||||
};
|
||||
let response = request
|
||||
.headers(attempt_headers)
|
||||
.timeout(remaining(started, timeout)?)
|
||||
.send()
|
||||
.await;
|
||||
let response = match response {
|
||||
Ok(response) => response,
|
||||
Err(error)
|
||||
if retries < JSON_BODY_PROTOCOL_RETRIES && retryable_protocol_error(&error) =>
|
||||
{
|
||||
retries += 1;
|
||||
continue;
|
||||
}
|
||||
Err(error) => return Err(network_error(error)),
|
||||
};
|
||||
let status = response.status();
|
||||
if !matches!(status.as_u16(), 301 | 302 | 303 | 307 | 308) {
|
||||
return Ok(response);
|
||||
}
|
||||
let Some(location) = response
|
||||
.headers()
|
||||
.get(reqwest::header::LOCATION)
|
||||
.and_then(|value| value.to_str().ok())
|
||||
else {
|
||||
return Ok(response);
|
||||
};
|
||||
let Ok(next_url) = url.join(location) else {
|
||||
return Ok(response);
|
||||
};
|
||||
if !matches!(next_url.scheme(), "http" | "https") {
|
||||
return Err(Error::Network("unsupported redirect scheme".into()));
|
||||
}
|
||||
if redirects >= JSON_BODY_MAX_REDIRECTS {
|
||||
return Err(Error::Network("too many provider redirects".into()));
|
||||
}
|
||||
redirects += 1;
|
||||
if matches!(
|
||||
status,
|
||||
StatusCode::MOVED_PERMANENTLY | StatusCode::FOUND | StatusCode::SEE_OTHER
|
||||
) {
|
||||
method = Method::GET;
|
||||
for name in [
|
||||
"content-type",
|
||||
"content-length",
|
||||
"content-encoding",
|
||||
"transfer-encoding",
|
||||
] {
|
||||
headers.remove(name);
|
||||
}
|
||||
}
|
||||
if url.origin() != next_url.origin() {
|
||||
for name in [
|
||||
"authorization",
|
||||
"proxy-authorization",
|
||||
"cookie",
|
||||
"www-authenticate",
|
||||
"x-api-key",
|
||||
"api-key",
|
||||
"x-amz-security-token",
|
||||
"x-amz-date",
|
||||
"x-amz-content-sha256",
|
||||
"ocp-apim-subscription-key",
|
||||
] {
|
||||
headers.remove(name);
|
||||
}
|
||||
signing = None;
|
||||
}
|
||||
headers.remove(reqwest::header::HOST);
|
||||
if !(url.scheme() == "https" && next_url.scheme() == "http") {
|
||||
let mut referer = url.clone();
|
||||
let _ = referer.set_username("");
|
||||
let _ = referer.set_password(None);
|
||||
referer.set_fragment(None);
|
||||
if let Ok(value) = HeaderValue::from_str(referer.as_str()) {
|
||||
headers.insert(reqwest::header::REFERER, value);
|
||||
}
|
||||
} else {
|
||||
headers.remove(reqwest::header::REFERER);
|
||||
}
|
||||
url = next_url;
|
||||
}
|
||||
}
|
||||
|
||||
fn signed_headers(
|
||||
signer: Option<&BodySigner<'_>>,
|
||||
url: &Url,
|
||||
digest: Option<&str>,
|
||||
headers: &HeaderMap,
|
||||
) -> Result<HeaderMap, Error> {
|
||||
match (signer, digest) {
|
||||
(Some(signer), Some(digest)) => signer(url, digest, headers),
|
||||
_ => Ok(headers.clone()),
|
||||
}
|
||||
}
|
||||
|
||||
fn remaining(started: Instant, timeout: Duration) -> Result<Duration, Error> {
|
||||
timeout
|
||||
.checked_sub(started.elapsed())
|
||||
.filter(|duration| !duration.is_zero())
|
||||
.ok_or_else(|| Error::Network("request timed out".into()))
|
||||
}
|
||||
|
||||
fn network_error(error: reqwest::Error) -> Error {
|
||||
if error.is_connect() || error.is_builder() {
|
||||
Error::Connect(error.to_string())
|
||||
} else {
|
||||
Error::Network(error.to_string())
|
||||
}
|
||||
}
|
||||
|
||||
fn retryable_protocol_error(error: &reqwest::Error) -> bool {
|
||||
let mut source = error.source();
|
||||
while let Some(error) = source {
|
||||
if let Some(error) = error.downcast_ref::<h2::Error>() {
|
||||
return error.is_remote()
|
||||
&& ((error.is_go_away() && error.reason() == Some(h2::Reason::NO_ERROR))
|
||||
|| (error.is_reset() && error.reason() == Some(h2::Reason::REFUSED_STREAM)));
|
||||
}
|
||||
source = error.source();
|
||||
}
|
||||
false
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
#[path = "replay_tests.rs"]
|
||||
mod tests;
|
||||
|
||||
fn request_body(body: &PreparedJsonBody) -> reqwest::Body {
|
||||
if let Some(bytes) = body.buffered_bytes() {
|
||||
return bytes.into();
|
||||
}
|
||||
reqwest::Body::wrap_stream(futures_util::stream::iter(
|
||||
body.chunks().map(Ok::<_, std::io::Error>),
|
||||
))
|
||||
}
|
||||
390
litellm-rust/crates/core/src/http_utils/replay_tests.rs
Normal file
390
litellm-rust/crates/core/src/http_utils/replay_tests.rs
Normal file
|
|
@ -0,0 +1,390 @@
|
|||
use super::*;
|
||||
use bytes::Bytes;
|
||||
use sha2::{Digest, Sha256};
|
||||
use std::sync::atomic::{AtomicUsize, Ordering};
|
||||
use tokio::io::{AsyncBufReadExt, AsyncReadExt, AsyncWriteExt, BufReader};
|
||||
use tokio::net::TcpListener;
|
||||
|
||||
use crate::http_utils::body::JsonPayload;
|
||||
|
||||
struct Received {
|
||||
method: String,
|
||||
headers: String,
|
||||
length: usize,
|
||||
digest: String,
|
||||
}
|
||||
|
||||
async fn listener() -> (TcpListener, String) {
|
||||
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
|
||||
let url = format!("http://{}", listener.local_addr().unwrap());
|
||||
(listener, url)
|
||||
}
|
||||
|
||||
async fn serve(listener: TcpListener, replies: Vec<(u16, String, Duration)>) -> Vec<Received> {
|
||||
let mut received = Vec::new();
|
||||
for (status, location, delay) in replies {
|
||||
let (socket, _) = listener.accept().await.unwrap();
|
||||
let mut socket = BufReader::new(socket);
|
||||
let mut line = String::new();
|
||||
socket.read_line(&mut line).await.unwrap();
|
||||
let method = line.split_whitespace().next().unwrap().to_owned();
|
||||
let mut headers = String::new();
|
||||
loop {
|
||||
line.clear();
|
||||
socket.read_line(&mut line).await.unwrap();
|
||||
if line == "\r\n" {
|
||||
break;
|
||||
}
|
||||
headers.push_str(&line.to_ascii_lowercase());
|
||||
}
|
||||
let length: usize = headers
|
||||
.lines()
|
||||
.find_map(|line| line.strip_prefix("content-length: "))
|
||||
.map(|length| length.parse().unwrap())
|
||||
.unwrap_or(0);
|
||||
let mut left = length;
|
||||
let mut digest = Sha256::new();
|
||||
let mut buffer = [0_u8; 8192];
|
||||
while left > 0 {
|
||||
let read = socket.read(&mut buffer[..left.min(8192)]).await.unwrap();
|
||||
assert_ne!(read, 0);
|
||||
digest.update(&buffer[..read]);
|
||||
left -= read;
|
||||
tokio::task::yield_now().await;
|
||||
}
|
||||
received.push(Received {
|
||||
method,
|
||||
headers,
|
||||
length,
|
||||
digest: format!("{:x}", digest.finalize()),
|
||||
});
|
||||
tokio::time::sleep(delay).await;
|
||||
if status == 0 {
|
||||
continue;
|
||||
}
|
||||
let location = if location.is_empty() {
|
||||
String::new()
|
||||
} else {
|
||||
format!("Location: {location}\r\n")
|
||||
};
|
||||
let response = format!(
|
||||
"HTTP/1.1 {status} Test\r\n{location}Content-Length: 2\r\nConnection: close\r\n\r\n{{}}"
|
||||
);
|
||||
let _ = socket.write_all(response.as_bytes()).await;
|
||||
}
|
||||
received
|
||||
}
|
||||
|
||||
fn body() -> PreparedJsonBody {
|
||||
PreparedJsonBody::streamed(JsonPayload::object([(
|
||||
"audio",
|
||||
JsonPayload::Base64(Bytes::from(vec![7; 1024 * 1024 + 1])),
|
||||
)]))
|
||||
.unwrap()
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn redirects_replay_identical_bodies_resign_same_origin_and_strip_cross_origin_secrets() {
|
||||
let (first, url) = listener().await;
|
||||
let (second, next_url) = listener().await;
|
||||
let first = tokio::spawn(serve(
|
||||
first,
|
||||
vec![
|
||||
(307, "/again".into(), Duration::ZERO),
|
||||
(308, next_url, Duration::ZERO),
|
||||
],
|
||||
));
|
||||
let second = tokio::spawn(serve(second, vec![(200, String::new(), Duration::ZERO)]));
|
||||
let signed = AtomicUsize::new(0);
|
||||
let body = body();
|
||||
let signer = |url: &Url, digest: &str, headers: &HeaderMap| {
|
||||
assert_eq!(digest, body.sha256());
|
||||
assert_eq!(headers[CONTENT_LENGTH], body.content_length().to_string());
|
||||
signed.fetch_add(1, Ordering::SeqCst);
|
||||
let mut headers = headers.clone();
|
||||
headers.insert("authorization", HeaderValue::from_str(url.path()).unwrap());
|
||||
headers.insert("x-amz-security-token", HeaderValue::from_static("secret"));
|
||||
Ok(headers)
|
||||
};
|
||||
let response = send_json(
|
||||
&Client::new(),
|
||||
&url,
|
||||
&body,
|
||||
&[
|
||||
("x-api-key".into(), "secret".into()),
|
||||
("cookie".into(), "secret".into()),
|
||||
],
|
||||
Duration::from_secs(10),
|
||||
Some(&signer),
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(response.status(), 200);
|
||||
assert_eq!(signed.load(Ordering::SeqCst), 2);
|
||||
let requests = first.await.unwrap();
|
||||
assert!(requests[0].headers.contains("authorization: /\r\n"));
|
||||
assert!(requests[1].headers.contains("authorization: /again\r\n"));
|
||||
let cross_origin = second.await.unwrap();
|
||||
assert!(!cross_origin[0].headers.contains("secret"));
|
||||
assert!(!cross_origin[0].headers.contains("authorization:"));
|
||||
for request in requests.into_iter().chain(cross_origin) {
|
||||
assert_eq!(request.method, "POST");
|
||||
assert_eq!(request.length as u64, body.content_length());
|
||||
assert_eq!(request.digest, body.sha256());
|
||||
assert!(!request.headers.contains("transfer-encoding:"));
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn redirects_that_change_post_to_get_drop_the_body() {
|
||||
for status in [301, 302, 303] {
|
||||
let (listener, url) = listener().await;
|
||||
let server = tokio::spawn(serve(
|
||||
listener,
|
||||
vec![
|
||||
(status, "/next".into(), Duration::ZERO),
|
||||
(200, String::new(), Duration::ZERO),
|
||||
],
|
||||
));
|
||||
send_json(
|
||||
&Client::new(),
|
||||
&url,
|
||||
&body(),
|
||||
&[],
|
||||
Duration::from_secs(10),
|
||||
None,
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
let requests = server.await.unwrap();
|
||||
assert_eq!(requests[1].method, "GET");
|
||||
assert_eq!(requests[1].length, 0);
|
||||
assert!(!requests[1].headers.contains("content-type:"));
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn redirects_share_a_deadline_and_have_a_limit() {
|
||||
let (listener, url) = listener().await;
|
||||
let server = tokio::spawn(serve(
|
||||
listener,
|
||||
vec![
|
||||
(307, "/next".into(), Duration::from_millis(70)),
|
||||
(200, String::new(), Duration::from_millis(70)),
|
||||
],
|
||||
));
|
||||
let body = PreparedJsonBody::streamed("data:test".into()).unwrap();
|
||||
let start = Instant::now();
|
||||
assert!(
|
||||
send_json(
|
||||
&Client::new(),
|
||||
&url,
|
||||
&body,
|
||||
&[],
|
||||
Duration::from_millis(110),
|
||||
None
|
||||
)
|
||||
.await
|
||||
.is_err()
|
||||
);
|
||||
assert!(start.elapsed() < Duration::from_millis(190));
|
||||
assert_eq!(server.await.unwrap().len(), 2);
|
||||
let (listener, url) = self::listener().await;
|
||||
let server = tokio::spawn(serve(
|
||||
listener,
|
||||
vec![(307, "/next".into(), Duration::ZERO); JSON_BODY_MAX_REDIRECTS + 1],
|
||||
));
|
||||
let error = send_json(
|
||||
&Client::new(),
|
||||
&url,
|
||||
&body,
|
||||
&[],
|
||||
Duration::from_secs(10),
|
||||
None,
|
||||
)
|
||||
.await
|
||||
.unwrap_err();
|
||||
assert!(error.to_string().contains("too many"));
|
||||
assert_eq!(server.await.unwrap().len(), JSON_BODY_MAX_REDIRECTS + 1);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn provider_errors_and_disconnects_are_not_retried() {
|
||||
for status in [429, 500, 0] {
|
||||
let (listener, url) = listener().await;
|
||||
let server = tokio::spawn(serve(
|
||||
listener,
|
||||
vec![(status, String::new(), Duration::ZERO)],
|
||||
));
|
||||
let result = send_json(
|
||||
&Client::new(),
|
||||
&url,
|
||||
&body(),
|
||||
&[],
|
||||
Duration::from_secs(10),
|
||||
None,
|
||||
)
|
||||
.await;
|
||||
if status == 0 {
|
||||
assert!(result.is_err());
|
||||
} else {
|
||||
assert_eq!(result.unwrap().status().as_u16(), status);
|
||||
}
|
||||
assert_eq!(server.await.unwrap().len(), 1);
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn http2_refused_streams_are_replayed_at_most_twice() {
|
||||
for failures in [2, 3] {
|
||||
let (listener, url) = listener().await;
|
||||
let attempts = std::sync::Arc::new(AtomicUsize::new(0));
|
||||
let observed = attempts.clone();
|
||||
let server = tokio::spawn(async move {
|
||||
loop {
|
||||
let (socket, _) = listener.accept().await.unwrap();
|
||||
let observed = observed.clone();
|
||||
tokio::spawn(async move {
|
||||
let mut connection = h2::server::handshake(socket).await.unwrap();
|
||||
while let Some(request) = connection.accept().await {
|
||||
let Ok((_, mut response)) = request else {
|
||||
break;
|
||||
};
|
||||
let attempt = observed.fetch_add(1, Ordering::SeqCst);
|
||||
if attempt < failures {
|
||||
response.send_reset(h2::Reason::REFUSED_STREAM);
|
||||
} else {
|
||||
response
|
||||
.send_response(http::Response::new(()), true)
|
||||
.unwrap();
|
||||
}
|
||||
}
|
||||
});
|
||||
}
|
||||
});
|
||||
let client = Client::builder()
|
||||
.http2_prior_knowledge()
|
||||
.redirect(reqwest::redirect::Policy::none())
|
||||
.retry(reqwest::retry::never())
|
||||
.build()
|
||||
.unwrap();
|
||||
let result = send_json(&client, &url, &body(), &[], Duration::from_secs(10), None).await;
|
||||
if failures == 2 {
|
||||
assert_eq!(result.unwrap().status(), 200);
|
||||
} else {
|
||||
assert!(result.is_err());
|
||||
}
|
||||
assert_eq!(attempts.load(Ordering::SeqCst), 3);
|
||||
server.abort();
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn cancellation_releases_media_while_a_slow_consumer_stalls_the_upload() {
|
||||
use std::sync::Arc;
|
||||
struct Owner(Arc<Vec<u8>>);
|
||||
impl AsRef<[u8]> for Owner {
|
||||
fn as_ref(&self) -> &[u8] {
|
||||
&self.0
|
||||
}
|
||||
}
|
||||
let owner = Arc::new(vec![b'A'; 16 * 1024 * 1024]);
|
||||
let body = PreparedJsonBody::streamed(JsonPayload::object([(
|
||||
"audio",
|
||||
JsonPayload::Base64(Bytes::from_owner(Owner(owner.clone()))),
|
||||
)]))
|
||||
.unwrap();
|
||||
let (listener, url) = listener().await;
|
||||
let (started, ready) = tokio::sync::oneshot::channel();
|
||||
let (close, closed) = tokio::sync::oneshot::channel();
|
||||
let server = tokio::spawn(async move {
|
||||
let (socket, _) = listener.accept().await.unwrap();
|
||||
let mut socket = BufReader::new(socket);
|
||||
loop {
|
||||
let mut line = String::new();
|
||||
socket.read_line(&mut line).await.unwrap();
|
||||
if line == "\r\n" {
|
||||
break;
|
||||
}
|
||||
}
|
||||
started.send(()).unwrap();
|
||||
let _ = closed.await;
|
||||
});
|
||||
let client = replay_client().unwrap().clone();
|
||||
let upload = tokio::spawn(async move {
|
||||
send_json(&client, &url, &body, &[], Duration::from_secs(30), None).await
|
||||
});
|
||||
ready.await.unwrap();
|
||||
assert_eq!(Arc::strong_count(&owner), 2);
|
||||
assert!(!upload.is_finished());
|
||||
upload.abort();
|
||||
assert!(upload.await.unwrap_err().is_cancelled());
|
||||
tokio::time::timeout(Duration::from_secs(2), async {
|
||||
while Arc::strong_count(&owner) > 1 {
|
||||
tokio::task::yield_now().await;
|
||||
}
|
||||
})
|
||||
.await
|
||||
.unwrap();
|
||||
close.send(()).unwrap();
|
||||
server.await.unwrap();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn media_upload_uses_the_supplied_clients_configuration() {
|
||||
let (listener, url) = listener().await;
|
||||
let server = tokio::spawn(serve(listener, vec![(200, String::new(), Duration::ZERO)]));
|
||||
let mut headers = HeaderMap::new();
|
||||
headers.insert("x-client", HeaderValue::from_static("supplied"));
|
||||
let client = Client::builder()
|
||||
.default_headers(headers)
|
||||
.redirect(reqwest::redirect::Policy::none())
|
||||
.retry(reqwest::retry::never())
|
||||
.build()
|
||||
.unwrap();
|
||||
let response = send_json(&client, &url, &body(), &[], Duration::from_secs(5), None)
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(response.status(), 200);
|
||||
assert!(
|
||||
server.await.unwrap()[0]
|
||||
.headers
|
||||
.contains("x-client: supplied")
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn http2_unprocessed_goaway_is_replayed() {
|
||||
let (listener, url) = listener().await;
|
||||
let server = tokio::spawn(async move {
|
||||
for attempt in 0..3 {
|
||||
let (socket, _) = listener.accept().await.unwrap();
|
||||
let mut connection = h2::server::handshake(socket).await.unwrap();
|
||||
if attempt < 2 {
|
||||
connection.abrupt_shutdown(h2::Reason::NO_ERROR);
|
||||
} else {
|
||||
let (_, mut response) = connection.accept().await.unwrap().unwrap();
|
||||
response
|
||||
.send_response(http::Response::new(()), true)
|
||||
.unwrap();
|
||||
}
|
||||
let _ = std::future::poll_fn(|cx| connection.poll_closed(cx)).await;
|
||||
}
|
||||
});
|
||||
let client = Client::builder()
|
||||
.http2_prior_knowledge()
|
||||
.redirect(reqwest::redirect::Policy::none())
|
||||
.retry(reqwest::retry::never())
|
||||
.build()
|
||||
.unwrap();
|
||||
let response = send_json(&client, &url, &body(), &[], Duration::from_secs(5), None)
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(response.status(), 200);
|
||||
drop(response);
|
||||
drop(client);
|
||||
tokio::time::timeout(Duration::from_secs(5), server)
|
||||
.await
|
||||
.unwrap()
|
||||
.unwrap();
|
||||
}
|
||||
|
|
@ -1,6 +1,10 @@
|
|||
use super::types::IntoMessagesRequest;
|
||||
use crate::constants::ANTHROPIC_MESSAGES_PROVIDER;
|
||||
use crate::constants::MESSAGES_TIMEOUT_SECS;
|
||||
use crate::error::Error;
|
||||
use crate::http_utils::http_request;
|
||||
use crate::http_utils::body::PreparedJsonBody;
|
||||
use crate::http_utils::replay::{replay_client, send_json};
|
||||
use std::time::Duration;
|
||||
|
||||
use super::client::http_client;
|
||||
use super::common_utils::truncate_error_body;
|
||||
|
|
@ -8,21 +12,26 @@ use super::prepare::prepare_provider_request;
|
|||
use super::types::{AnthropicMessagesResponse, MessagesRequest};
|
||||
|
||||
#[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)]
|
||||
pub(super) async fn execute_messages_provider_call(
|
||||
request: MessagesRequest<'_>,
|
||||
pub(super) async fn execute_messages_provider_call<B: IntoMessagesRequest>(
|
||||
request: MessagesRequest<'_, B>,
|
||||
) -> Result<AnthropicMessagesResponse, Error> {
|
||||
let request = prepare_provider_request(request)?;
|
||||
let mut request_builder = http_client().post(&request.url).json(&request.body);
|
||||
for (key, value) in &request.upstream_headers {
|
||||
request_builder = request_builder.header(key, value);
|
||||
}
|
||||
if let Some(duration) = request.timeout {
|
||||
request_builder = request_builder.timeout(duration);
|
||||
}
|
||||
|
||||
let response = http_request(request_builder)
|
||||
.await
|
||||
.map_err(|err| Error::Network(err.to_string()))?;
|
||||
let body = PreparedJsonBody::new(request.body.into_payload()?)?;
|
||||
let response = send_json(
|
||||
if body.is_streamed() {
|
||||
replay_client()?
|
||||
} else {
|
||||
http_client()
|
||||
},
|
||||
&request.url,
|
||||
&body,
|
||||
&request.upstream_headers,
|
||||
request
|
||||
.timeout
|
||||
.unwrap_or(Duration::from_secs(MESSAGES_TIMEOUT_SECS)),
|
||||
None,
|
||||
)
|
||||
.await?;
|
||||
|
||||
let status = response.status();
|
||||
let text = response
|
||||
|
|
@ -42,8 +51,8 @@ pub(super) async fn execute_messages_provider_call(
|
|||
request.config.transform_response(&request.model, response)
|
||||
}
|
||||
|
||||
pub(super) async fn execute_messages_provider_stream(
|
||||
request: MessagesRequest<'_>,
|
||||
pub(super) async fn execute_messages_provider_stream<B: IntoMessagesRequest>(
|
||||
request: MessagesRequest<'_, B>,
|
||||
) -> Result<reqwest::Response, Error> {
|
||||
let request = prepare_provider_request(request)?;
|
||||
if request.provider != ANTHROPIC_MESSAGES_PROVIDER {
|
||||
|
|
@ -52,17 +61,23 @@ pub(super) async fn execute_messages_provider_stream(
|
|||
));
|
||||
}
|
||||
|
||||
let mut request_builder = http_client().post(&request.url).json(&request.body);
|
||||
for (key, value) in &request.upstream_headers {
|
||||
request_builder = request_builder.header(key, value);
|
||||
}
|
||||
if let Some(duration) = request.timeout {
|
||||
request_builder = request_builder.timeout(duration);
|
||||
}
|
||||
let body = PreparedJsonBody::new(request.body.into_payload()?)?;
|
||||
let response = send_json(
|
||||
if body.is_streamed() {
|
||||
replay_client()?
|
||||
} else {
|
||||
http_client()
|
||||
},
|
||||
&request.url,
|
||||
&body,
|
||||
&request.upstream_headers,
|
||||
request
|
||||
.timeout
|
||||
.unwrap_or(Duration::from_secs(MESSAGES_TIMEOUT_SECS)),
|
||||
None,
|
||||
)
|
||||
.await?;
|
||||
|
||||
let response = http_request(request_builder)
|
||||
.await
|
||||
.map_err(|err| Error::Network(err.to_string()))?;
|
||||
let status = response.status();
|
||||
if !status.is_success() {
|
||||
let text = response
|
||||
|
|
|
|||
|
|
@ -11,7 +11,9 @@ use crate::Error;
|
|||
mod client;
|
||||
mod common_utils;
|
||||
mod handler;
|
||||
mod payload;
|
||||
mod prepare;
|
||||
use types::IntoMessagesRequest;
|
||||
pub mod transformation;
|
||||
pub mod types;
|
||||
|
||||
|
|
@ -19,11 +21,15 @@ use handler::{execute_messages_provider_call, execute_messages_provider_stream};
|
|||
use types::{AnthropicMessagesResponse, MessagesRequest};
|
||||
|
||||
#[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)]
|
||||
pub async fn messages(request: MessagesRequest<'_>) -> Result<AnthropicMessagesResponse, Error> {
|
||||
pub async fn messages<B: IntoMessagesRequest>(
|
||||
request: MessagesRequest<'_, B>,
|
||||
) -> Result<AnthropicMessagesResponse, Error> {
|
||||
execute_messages_provider_call(request).await
|
||||
}
|
||||
|
||||
pub async fn messages_stream(request: MessagesRequest<'_>) -> Result<reqwest::Response, Error> {
|
||||
pub async fn messages_stream<B: IntoMessagesRequest>(
|
||||
request: MessagesRequest<'_, B>,
|
||||
) -> Result<reqwest::Response, Error> {
|
||||
execute_messages_provider_stream(request).await
|
||||
}
|
||||
|
||||
|
|
|
|||
109
litellm-rust/crates/core/src/messages/payload.rs
Normal file
109
litellm-rust/crates/core/src/messages/payload.rs
Normal file
|
|
@ -0,0 +1,109 @@
|
|||
use super::types::{
|
||||
AnthropicMessage, AnthropicMessagesRequest, CacheControl, ContentBlock, MessageContent,
|
||||
SystemPrompt,
|
||||
};
|
||||
use crate::Error;
|
||||
use crate::http_utils::body::JsonPayload;
|
||||
|
||||
fn metadata(value: impl serde::Serialize) -> Result<JsonPayload, Error> {
|
||||
serde_json::to_value(value)
|
||||
.map(JsonPayload::from)
|
||||
.map_err(|_| Error::InvalidRequest("invalid messages metadata".into()))
|
||||
}
|
||||
|
||||
impl AnthropicMessagesRequest {
|
||||
pub(crate) fn into_payload(self) -> Result<JsonPayload, Error> {
|
||||
let mut fields = self.extra;
|
||||
fields.insert("model".into(), self.model.into());
|
||||
fields.insert(
|
||||
"messages".into(),
|
||||
JsonPayload::Array(
|
||||
self.messages
|
||||
.into_iter()
|
||||
.map(AnthropicMessage::into_payload)
|
||||
.collect::<Result<_, _>>()?,
|
||||
),
|
||||
);
|
||||
if let Some(system) = self.system {
|
||||
let content = match system {
|
||||
SystemPrompt::Text(text) => MessageContent::Text(text),
|
||||
SystemPrompt::Blocks(blocks) => MessageContent::Blocks(blocks),
|
||||
};
|
||||
fields.insert("system".into(), content.into_payload()?);
|
||||
}
|
||||
macro_rules! optional {
|
||||
($($field:ident => $convert:expr),* $(,)?) => { $(
|
||||
if let Some(value) = self.$field { fields.insert(stringify!($field).into(), $convert(value)?); }
|
||||
)* };
|
||||
}
|
||||
optional! {
|
||||
max_tokens => metadata,
|
||||
stop_sequences => metadata,
|
||||
stream => metadata,
|
||||
temperature => metadata,
|
||||
top_p => metadata,
|
||||
top_k => metadata,
|
||||
service_tier => metadata,
|
||||
speed => metadata,
|
||||
inference_geo => metadata,
|
||||
metadata => Ok::<_, Error>,
|
||||
tool_choice => Ok::<_, Error>,
|
||||
thinking => Ok::<_, Error>,
|
||||
container => Ok::<_, Error>,
|
||||
context_management => Ok::<_, Error>,
|
||||
output_format => Ok::<_, Error>,
|
||||
output_config => Ok::<_, Error>,
|
||||
tools => |value| Ok::<_, Error>(JsonPayload::Array(value)),
|
||||
mcp_servers => |value| Ok::<_, Error>(JsonPayload::Array(value)),
|
||||
}
|
||||
Ok(JsonPayload::Object(fields))
|
||||
}
|
||||
}
|
||||
|
||||
impl AnthropicMessage {
|
||||
fn into_payload(self) -> Result<JsonPayload, Error> {
|
||||
let mut fields = self.extra;
|
||||
fields.insert("role".into(), self.role.into());
|
||||
fields.insert("content".into(), self.content.into_payload()?);
|
||||
Ok(JsonPayload::Object(fields))
|
||||
}
|
||||
}
|
||||
|
||||
impl MessageContent {
|
||||
fn into_payload(self) -> Result<JsonPayload, Error> {
|
||||
match self {
|
||||
Self::Text(text) => Ok(JsonPayload::String(text)),
|
||||
Self::Blocks(blocks) => blocks
|
||||
.into_iter()
|
||||
.map(ContentBlock::into_payload)
|
||||
.collect::<Result<_, _>>()
|
||||
.map(JsonPayload::Array),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl ContentBlock {
|
||||
fn into_payload(self) -> Result<JsonPayload, Error> {
|
||||
let mut fields = self.extra;
|
||||
if let Some(cache) = self.cache_control {
|
||||
fields.insert("cache_control".into(), cache.into_payload()?);
|
||||
}
|
||||
Ok(JsonPayload::Object(fields))
|
||||
}
|
||||
}
|
||||
|
||||
impl CacheControl {
|
||||
fn into_payload(self) -> Result<JsonPayload, Error> {
|
||||
let mut fields = self.extra;
|
||||
for (key, value) in [
|
||||
("type", self.cache_type),
|
||||
("ttl", self.ttl),
|
||||
("scope", self.scope),
|
||||
] {
|
||||
if let Some(value) = value {
|
||||
fields.insert(key.into(), value.into());
|
||||
}
|
||||
}
|
||||
Ok(JsonPayload::Object(fields))
|
||||
}
|
||||
}
|
||||
|
|
@ -1,3 +1,4 @@
|
|||
use super::types::IntoMessagesRequest;
|
||||
use crate::error::Error;
|
||||
use crate::routing_utils::provider::{CustomLlmProvider, get_custom_llm_provider};
|
||||
|
||||
|
|
@ -6,8 +7,8 @@ use super::transformation::{AnthropicMessagesProviderConfig, MessagesAuthStrateg
|
|||
use super::types::{MessagesRequest, ProviderMessagesRequest};
|
||||
use serde_json::{Map, Value};
|
||||
|
||||
pub(super) fn prepare_provider_request(
|
||||
request: MessagesRequest<'_>,
|
||||
pub(super) fn prepare_provider_request<B: IntoMessagesRequest>(
|
||||
request: MessagesRequest<'_, B>,
|
||||
) -> Result<ProviderMessagesRequest, Error> {
|
||||
let provider_info = get_custom_llm_provider(request.model, request.custom_llm_provider)
|
||||
.or_else(|| {
|
||||
|
|
@ -33,15 +34,8 @@ pub(super) fn prepare_provider_request(
|
|||
let headers =
|
||||
validate_environment(config, request.extra_headers, request.api_key, &env_lookup)?;
|
||||
|
||||
let typed_request = serde_json::from_value(request.body).map_err(|err| {
|
||||
Error::InvalidRequest(format!("invalid Anthropic messages request: {err}"))
|
||||
})?;
|
||||
let transformed = config.transform_request(typed_request)?;
|
||||
let body = serde_json::to_value(transformed).map_err(|err| {
|
||||
Error::InvalidRequest(format!(
|
||||
"failed to serialize Anthropic messages request: {err}"
|
||||
))
|
||||
})?;
|
||||
let typed_request = request.body.into_messages_request()?;
|
||||
let body = config.transform_request(typed_request)?;
|
||||
|
||||
let url = config.complete_url(request.api_base, &model, &env_lookup)?;
|
||||
|
||||
|
|
|
|||
|
|
@ -1,3 +1,5 @@
|
|||
use crate::http_utils::body::{JsonPayload, SharedText};
|
||||
use std::collections::BTreeMap;
|
||||
use std::time::Duration;
|
||||
|
||||
use serde::{Deserialize, Serialize};
|
||||
|
|
@ -5,9 +7,9 @@ use serde_json::{Map, Value};
|
|||
|
||||
use super::transformation::AnthropicMessagesProviderConfig;
|
||||
|
||||
pub struct MessagesRequest<'a> {
|
||||
pub struct MessagesRequest<'a, B = Value> {
|
||||
pub model: &'a str,
|
||||
pub body: Value,
|
||||
pub body: B,
|
||||
pub api_key: Option<&'a str>,
|
||||
pub api_base: Option<&'a str>,
|
||||
pub custom_llm_provider: Option<&'a str>,
|
||||
|
|
@ -20,7 +22,7 @@ pub(super) struct ProviderMessagesRequest {
|
|||
pub(super) model: String,
|
||||
pub(super) config: &'static dyn AnthropicMessagesProviderConfig,
|
||||
pub(super) url: String,
|
||||
pub(super) body: Value,
|
||||
pub(super) body: AnthropicMessagesRequest,
|
||||
pub(super) upstream_headers: Vec<(String, String)>,
|
||||
pub(super) timeout: Option<Duration>,
|
||||
}
|
||||
|
|
@ -28,14 +30,14 @@ pub(super) struct ProviderMessagesRequest {
|
|||
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
|
||||
#[serde(untagged)]
|
||||
pub enum SystemPrompt {
|
||||
Text(String),
|
||||
Text(SharedText),
|
||||
Blocks(Vec<ContentBlock>),
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
|
||||
#[serde(untagged)]
|
||||
pub enum MessageContent {
|
||||
Text(String),
|
||||
Text(SharedText),
|
||||
Blocks(Vec<ContentBlock>),
|
||||
}
|
||||
|
||||
|
|
@ -44,7 +46,7 @@ pub struct ContentBlock {
|
|||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub cache_control: Option<CacheControl>,
|
||||
#[serde(flatten)]
|
||||
pub extra: Map<String, Value>,
|
||||
pub extra: BTreeMap<String, JsonPayload>,
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, Default, PartialEq, Serialize, Deserialize)]
|
||||
|
|
@ -56,7 +58,7 @@ pub struct CacheControl {
|
|||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub scope: Option<String>,
|
||||
#[serde(flatten)]
|
||||
pub extra: Map<String, Value>,
|
||||
pub extra: BTreeMap<String, JsonPayload>,
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
|
||||
|
|
@ -64,7 +66,7 @@ pub struct AnthropicMessage {
|
|||
pub role: String,
|
||||
pub content: MessageContent,
|
||||
#[serde(flatten)]
|
||||
pub extra: Map<String, Value>,
|
||||
pub extra: BTreeMap<String, JsonPayload>,
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
|
||||
|
|
@ -76,7 +78,7 @@ pub struct AnthropicMessagesRequest {
|
|||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub system: Option<SystemPrompt>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub metadata: Option<Value>,
|
||||
pub metadata: Option<JsonPayload>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub stop_sequences: Option<Vec<String>>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
|
|
@ -88,29 +90,29 @@ pub struct AnthropicMessagesRequest {
|
|||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub top_k: Option<i64>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub tools: Option<Vec<Value>>,
|
||||
pub tools: Option<Vec<JsonPayload>>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub tool_choice: Option<Value>,
|
||||
pub tool_choice: Option<JsonPayload>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub thinking: Option<Value>,
|
||||
pub thinking: Option<JsonPayload>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub service_tier: Option<String>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub container: Option<Value>,
|
||||
pub container: Option<JsonPayload>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub mcp_servers: Option<Vec<Value>>,
|
||||
pub mcp_servers: Option<Vec<JsonPayload>>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub context_management: Option<Value>,
|
||||
pub context_management: Option<JsonPayload>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub output_format: Option<Value>,
|
||||
pub output_format: Option<JsonPayload>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub output_config: Option<Value>,
|
||||
pub output_config: Option<JsonPayload>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub speed: Option<String>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub inference_geo: Option<String>,
|
||||
#[serde(flatten)]
|
||||
pub extra: Map<String, Value>,
|
||||
pub extra: BTreeMap<String, JsonPayload>,
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
|
||||
|
|
@ -132,3 +134,21 @@ pub struct AnthropicMessagesResponse {
|
|||
#[serde(flatten)]
|
||||
pub extra: Map<String, Value>,
|
||||
}
|
||||
|
||||
pub trait IntoMessagesRequest {
|
||||
fn into_messages_request(self) -> Result<AnthropicMessagesRequest, crate::Error>;
|
||||
}
|
||||
|
||||
impl IntoMessagesRequest for AnthropicMessagesRequest {
|
||||
fn into_messages_request(self) -> Result<AnthropicMessagesRequest, crate::Error> {
|
||||
Ok(self)
|
||||
}
|
||||
}
|
||||
|
||||
impl IntoMessagesRequest for Value {
|
||||
fn into_messages_request(self) -> Result<AnthropicMessagesRequest, crate::Error> {
|
||||
serde_json::from_value(self).map_err(|error| {
|
||||
crate::Error::InvalidRequest(format!("invalid Anthropic messages request: {error}"))
|
||||
})
|
||||
}
|
||||
}
|
||||
|
|
|
|||
19
litellm-rust/crates/core/src/ocr/client.rs
Normal file
19
litellm-rust/crates/core/src/ocr/client.rs
Normal file
|
|
@ -0,0 +1,19 @@
|
|||
use std::sync::OnceLock;
|
||||
use std::time::Duration;
|
||||
|
||||
use crate::Error;
|
||||
use crate::constants::{MESSAGES_CONNECT_TIMEOUT_SECS, OCR_TIMEOUT_SECS};
|
||||
|
||||
pub(super) fn http_client() -> Result<&'static reqwest::Client, Error> {
|
||||
static CLIENT: OnceLock<Result<reqwest::Client, String>> = OnceLock::new();
|
||||
CLIENT
|
||||
.get_or_init(|| {
|
||||
reqwest::Client::builder()
|
||||
.connect_timeout(Duration::from_secs(MESSAGES_CONNECT_TIMEOUT_SECS))
|
||||
.timeout(Duration::from_secs(OCR_TIMEOUT_SECS))
|
||||
.build()
|
||||
.map_err(|error| error.to_string())
|
||||
})
|
||||
.as_ref()
|
||||
.map_err(|error| Error::Network(error.clone()))
|
||||
}
|
||||
187
litellm-rust/crates/core/src/ocr/handler.rs
Normal file
187
litellm-rust/crates/core/src/ocr/handler.rs
Normal file
|
|
@ -0,0 +1,187 @@
|
|||
use super::client::http_client;
|
||||
use super::transformation::OcrResponseHandling;
|
||||
use super::types::ProviderOcrRequest;
|
||||
use crate::Error;
|
||||
use crate::constants::{AZURE_DOCUMENT_INTELLIGENCE_POLL_TIMEOUT_SECS, OCR_TIMEOUT_SECS};
|
||||
use crate::http_utils::body::PreparedJsonBody;
|
||||
use crate::http_utils::replay::{replay_client, send_json};
|
||||
use crate::http_utils::truncate_error_body;
|
||||
use serde_json::Value;
|
||||
use std::time::{Duration, Instant};
|
||||
|
||||
pub async fn execute_ocr_provider_call(request: ProviderOcrRequest) -> Result<Value, Error> {
|
||||
let started = Instant::now();
|
||||
let body = PreparedJsonBody::new(request.body)?;
|
||||
let response = send_json(
|
||||
if body.is_streamed() {
|
||||
replay_client()?
|
||||
} else {
|
||||
http_client()?
|
||||
},
|
||||
&request.url,
|
||||
&body,
|
||||
&request.upstream_headers,
|
||||
request
|
||||
.timeout
|
||||
.unwrap_or(Duration::from_secs(OCR_TIMEOUT_SECS)),
|
||||
None,
|
||||
)
|
||||
.await?;
|
||||
let status = response.status();
|
||||
if request.config.response_handling() == OcrResponseHandling::AzureDocumentIntelligencePoll
|
||||
&& status.as_u16() == 202
|
||||
{
|
||||
let operation_url = response
|
||||
.headers()
|
||||
.get("operation-location")
|
||||
.and_then(|value| value.to_str().ok())
|
||||
.map(str::to_string)
|
||||
.ok_or_else(|| {
|
||||
Error::InvalidResponse(
|
||||
"Azure Document Intelligence returned 202 but no Operation-Location header found"
|
||||
.to_string(),
|
||||
)
|
||||
})?;
|
||||
let response_json = poll_document_intelligence(
|
||||
&operation_url,
|
||||
&request.url,
|
||||
&request.upstream_headers,
|
||||
request
|
||||
.timeout
|
||||
.map(|timeout| timeout.saturating_sub(started.elapsed())),
|
||||
)
|
||||
.await?;
|
||||
return Ok(request
|
||||
.config
|
||||
.transform_ocr_response(&request.model, response_json)?
|
||||
.into_json());
|
||||
}
|
||||
|
||||
let text = response
|
||||
.text()
|
||||
.await
|
||||
.map_err(|err| Error::Network(err.to_string()))?;
|
||||
|
||||
if !status.is_success() {
|
||||
return Err(Error::Http {
|
||||
status: status.as_u16(),
|
||||
body: truncate_error_body(&text),
|
||||
});
|
||||
}
|
||||
|
||||
let response_json: Value = serde_json::from_str(&text)
|
||||
.map_err(|err| Error::InvalidResponse(format!("invalid OCR response JSON: {err}")))?;
|
||||
|
||||
Ok(request
|
||||
.config
|
||||
.transform_ocr_response(&request.model, response_json)?
|
||||
.into_json())
|
||||
}
|
||||
|
||||
fn same_origin(left: &str, right: &str) -> bool {
|
||||
let Ok(left) = reqwest::Url::parse(left) else {
|
||||
return false;
|
||||
};
|
||||
let Ok(right) = reqwest::Url::parse(right) else {
|
||||
return false;
|
||||
};
|
||||
left.scheme() == right.scheme()
|
||||
&& left.host_str() == right.host_str()
|
||||
&& left.port_or_known_default() == right.port_or_known_default()
|
||||
}
|
||||
|
||||
fn retry_after_secs(response: &reqwest::Response) -> u64 {
|
||||
response
|
||||
.headers()
|
||||
.get(reqwest::header::RETRY_AFTER)
|
||||
.and_then(|value| value.to_str().ok())
|
||||
.and_then(|value| value.parse::<u64>().ok())
|
||||
.unwrap_or(2)
|
||||
}
|
||||
|
||||
fn operation_status(response_json: &Value) -> Result<&str, Error> {
|
||||
let status = response_json
|
||||
.get("status")
|
||||
.and_then(Value::as_str)
|
||||
.ok_or(Error::MissingField("status"))?;
|
||||
match status {
|
||||
"succeeded" => Ok("succeeded"),
|
||||
"running" | "notStarted" => Ok("running"),
|
||||
"failed" => {
|
||||
let message = response_json
|
||||
.get("error")
|
||||
.and_then(|error| error.get("message"))
|
||||
.and_then(Value::as_str)
|
||||
.unwrap_or("Unknown error");
|
||||
Err(Error::InvalidResponse(format!(
|
||||
"Azure Document Intelligence analysis failed: {message}"
|
||||
)))
|
||||
}
|
||||
other => Err(Error::InvalidResponse(format!(
|
||||
"Unknown operation status: {other}"
|
||||
))),
|
||||
}
|
||||
}
|
||||
|
||||
async fn poll_document_intelligence(
|
||||
operation_url: &str,
|
||||
original_url: &str,
|
||||
headers: &[(String, String)],
|
||||
timeout: Option<Duration>,
|
||||
) -> Result<Value, Error> {
|
||||
if !same_origin(operation_url, original_url) {
|
||||
return Err(Error::InvalidResponse(
|
||||
"Azure Document Intelligence: rejected cross-origin polling URL".to_string(),
|
||||
));
|
||||
}
|
||||
|
||||
let start = Instant::now();
|
||||
let timeout = timeout.unwrap_or(Duration::from_secs(
|
||||
AZURE_DOCUMENT_INTELLIGENCE_POLL_TIMEOUT_SECS,
|
||||
));
|
||||
loop {
|
||||
if start.elapsed() > timeout {
|
||||
return Err(Error::Network(format!(
|
||||
"Azure Document Intelligence operation polling timed out after {} seconds",
|
||||
timeout.as_secs()
|
||||
)));
|
||||
}
|
||||
|
||||
let mut request_builder = http_client()?.get(operation_url);
|
||||
for (key, value) in headers {
|
||||
if key.eq_ignore_ascii_case("ocp-apim-subscription-key") {
|
||||
request_builder = request_builder.header(key, value);
|
||||
}
|
||||
}
|
||||
let remaining = timeout
|
||||
.checked_sub(start.elapsed())
|
||||
.ok_or_else(|| Error::Network("OCR polling timed out".into()))?;
|
||||
let response = request_builder
|
||||
.timeout(remaining)
|
||||
.send()
|
||||
.await
|
||||
.map_err(|err| Error::Network(err.to_string()))?;
|
||||
let retry_after = retry_after_secs(&response);
|
||||
let status = response.status();
|
||||
let text = response
|
||||
.text()
|
||||
.await
|
||||
.map_err(|err| Error::Network(err.to_string()))?;
|
||||
if !status.is_success() {
|
||||
return Err(Error::Http {
|
||||
status: status.as_u16(),
|
||||
body: truncate_error_body(&text),
|
||||
});
|
||||
}
|
||||
let response_json: Value = serde_json::from_str(&text).map_err(|err| {
|
||||
Error::InvalidResponse(format!("invalid Azure DI poll response JSON: {err}"))
|
||||
})?;
|
||||
if operation_status(&response_json)? == "succeeded" {
|
||||
return Ok(response_json);
|
||||
}
|
||||
tokio::time::sleep(
|
||||
Duration::from_secs(retry_after).min(timeout.saturating_sub(start.elapsed())),
|
||||
)
|
||||
.await;
|
||||
}
|
||||
}
|
||||
|
|
@ -1,2 +1,6 @@
|
|||
pub mod transformation;
|
||||
pub mod types;
|
||||
|
||||
mod client;
|
||||
mod handler;
|
||||
pub use handler::execute_ocr_provider_call;
|
||||
|
|
|
|||
|
|
@ -1,4 +1,5 @@
|
|||
use crate::Error;
|
||||
use crate::http_utils::body::JsonPayload;
|
||||
use serde_json::{Map, Value};
|
||||
|
||||
use super::types::{OcrRequestData, OcrResponseData};
|
||||
|
|
@ -43,6 +44,15 @@ pub trait OcrProviderConfig: Sync {
|
|||
model: &str,
|
||||
document: Value,
|
||||
optional_params: Map<String, Value>,
|
||||
) -> Result<OcrRequestData, Error> {
|
||||
self.transform_ocr_payload(model, document.into(), optional_params)
|
||||
}
|
||||
|
||||
fn transform_ocr_payload(
|
||||
&self,
|
||||
model: &str,
|
||||
document: JsonPayload,
|
||||
optional_params: Map<String, Value>,
|
||||
) -> Result<OcrRequestData, Error>;
|
||||
|
||||
fn transform_ocr_response(
|
||||
|
|
|
|||
|
|
@ -1,9 +1,10 @@
|
|||
use crate::http_utils::body::JsonPayload;
|
||||
use serde::{Deserialize, Serialize};
|
||||
use serde_json::Value;
|
||||
|
||||
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
|
||||
pub struct OcrRequestData {
|
||||
pub data: Value,
|
||||
pub data: JsonPayload,
|
||||
pub files: Option<Value>,
|
||||
}
|
||||
|
||||
|
|
@ -27,3 +28,12 @@ impl OcrResponseData {
|
|||
})
|
||||
}
|
||||
}
|
||||
|
||||
pub struct ProviderOcrRequest {
|
||||
pub model: String,
|
||||
pub config: &'static dyn super::transformation::OcrProviderConfig,
|
||||
pub url: String,
|
||||
pub body: JsonPayload,
|
||||
pub upstream_headers: Vec<(String, String)>,
|
||||
pub timeout: Option<std::time::Duration>,
|
||||
}
|
||||
|
|
|
|||
|
|
@ -7,7 +7,6 @@ use crate::messages::types::{
|
|||
use crate::providers::anthropic::messages::transformation::{
|
||||
ANTHROPIC_MESSAGES_CONFIG, AnthropicMessagesConfig, non_empty,
|
||||
};
|
||||
use serde_json::{Map, Value};
|
||||
|
||||
const AZURE_API_KEY_ENV: &str = "AZURE_API_KEY";
|
||||
const AZURE_API_BASE_ENV: &str = "AZURE_API_BASE";
|
||||
|
|
@ -86,13 +85,16 @@ fn strip_scope_from_message(message: &mut AnthropicMessage) {
|
|||
}
|
||||
}
|
||||
|
||||
fn text_content_block(text: String) -> ContentBlock {
|
||||
let extra = Map::from_iter([
|
||||
fn text_content_block(text: crate::http_utils::body::SharedText) -> ContentBlock {
|
||||
let extra = std::collections::BTreeMap::from_iter([
|
||||
(
|
||||
"type".to_string(),
|
||||
Value::String(TEXT_BLOCK_TYPE.to_string()),
|
||||
crate::http_utils::body::JsonPayload::from(TEXT_BLOCK_TYPE),
|
||||
),
|
||||
(
|
||||
"text".to_string(),
|
||||
crate::http_utils::body::JsonPayload::String(text),
|
||||
),
|
||||
("text".to_string(), Value::String(text)),
|
||||
]);
|
||||
ContentBlock {
|
||||
cache_control: None,
|
||||
|
|
|
|||
|
|
@ -1,3 +1,4 @@
|
|||
use crate::http_utils::body::JsonPayload;
|
||||
use std::collections::BTreeSet;
|
||||
|
||||
use crate::error::{Error, json_type_name};
|
||||
|
|
@ -216,14 +217,16 @@ pub fn complete_document_intelligence_url(
|
|||
Ok(url)
|
||||
}
|
||||
|
||||
fn document_url_from_mistral_document(document: &Value) -> Result<&str, Error> {
|
||||
fn document_url_from_mistral_document(
|
||||
document: &JsonPayload,
|
||||
) -> Result<&crate::http_utils::body::SharedText, Error> {
|
||||
let object = document.as_object().ok_or_else(|| Error::InvalidType {
|
||||
expected: "object",
|
||||
actual: json_type_name(document),
|
||||
actual: document.type_name(),
|
||||
})?;
|
||||
let doc_type = object
|
||||
.get("type")
|
||||
.and_then(Value::as_str)
|
||||
.and_then(JsonPayload::as_str)
|
||||
.ok_or(Error::MissingField("document.type"))?;
|
||||
let field_name = match doc_type {
|
||||
"document_url" => "document_url",
|
||||
|
|
@ -236,18 +239,11 @@ fn document_url_from_mistral_document(document: &Value) -> Result<&str, Error> {
|
|||
};
|
||||
object
|
||||
.get(field_name)
|
||||
.and_then(Value::as_str)
|
||||
.filter(|value| !value.is_empty())
|
||||
.and_then(JsonPayload::as_text)
|
||||
.filter(|value| !value.as_str().is_empty())
|
||||
.ok_or(Error::MissingField(field_name))
|
||||
}
|
||||
|
||||
fn extract_base64_from_data_uri(data_uri: &str) -> &str {
|
||||
data_uri
|
||||
.split_once(',')
|
||||
.map(|(_, data)| data)
|
||||
.unwrap_or(data_uri)
|
||||
}
|
||||
|
||||
fn page_markdown(page: &Map<String, Value>) -> String {
|
||||
page.get("lines")
|
||||
.and_then(Value::as_array)
|
||||
|
|
@ -285,13 +281,13 @@ impl OcrProviderConfig for AzureAiOcrConfig {
|
|||
MISTRAL_OCR_CONFIG.supported_ocr_params()
|
||||
}
|
||||
|
||||
fn transform_ocr_request(
|
||||
fn transform_ocr_payload(
|
||||
&self,
|
||||
model: &str,
|
||||
document: Value,
|
||||
document: JsonPayload,
|
||||
optional_params: Map<String, Value>,
|
||||
) -> Result<OcrRequestData, Error> {
|
||||
MISTRAL_OCR_CONFIG.transform_ocr_request(model, document, optional_params)
|
||||
MISTRAL_OCR_CONFIG.transform_ocr_payload(model, document, optional_params)
|
||||
}
|
||||
|
||||
fn transform_ocr_response(
|
||||
|
|
@ -330,27 +326,24 @@ impl OcrProviderConfig for AzureDocumentIntelligenceOcrConfig {
|
|||
AZURE_DOCUMENT_INTELLIGENCE_SUPPORTED_OCR_PARAMS
|
||||
}
|
||||
|
||||
fn transform_ocr_request(
|
||||
fn transform_ocr_payload(
|
||||
&self,
|
||||
_model: &str,
|
||||
document: Value,
|
||||
document: JsonPayload,
|
||||
_optional_params: Map<String, Value>,
|
||||
) -> Result<OcrRequestData, Error> {
|
||||
let document_url = document_url_from_mistral_document(&document)?;
|
||||
let mut data = Map::new();
|
||||
if document_url.starts_with("data:") {
|
||||
data.insert(
|
||||
"base64Source".to_string(),
|
||||
Value::String(extract_base64_from_data_uri(document_url).to_string()),
|
||||
);
|
||||
let (field, value) = if document_url.as_str().starts_with("data:") {
|
||||
let start = document_url.as_str().find(',').map_or(0, |index| index + 1);
|
||||
(
|
||||
"base64Source",
|
||||
document_url.slice(start..document_url.bytes().len())?,
|
||||
)
|
||||
} else {
|
||||
data.insert(
|
||||
"urlSource".to_string(),
|
||||
Value::String(document_url.to_string()),
|
||||
);
|
||||
}
|
||||
("urlSource", document_url.clone())
|
||||
};
|
||||
Ok(OcrRequestData {
|
||||
data: Value::Object(data),
|
||||
data: JsonPayload::object([(field, JsonPayload::String(value))]),
|
||||
files: None,
|
||||
})
|
||||
}
|
||||
|
|
@ -439,6 +432,28 @@ impl OcrProviderConfig for AzureDocumentIntelligenceOcrConfig {
|
|||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn document_intelligence_slices_the_shared_data_uri() {
|
||||
use crate::http_utils::body::SharedText;
|
||||
use bytes::Bytes;
|
||||
let source = Bytes::from(String::from("data:application/pdf;base64,AQIDBA=="));
|
||||
let offset = source.iter().position(|byte| *byte == b',').unwrap() + 1;
|
||||
let document = JsonPayload::object([
|
||||
("type", "document_url".into()),
|
||||
(
|
||||
"document_url",
|
||||
JsonPayload::String(SharedText::new(source.clone()).unwrap()),
|
||||
),
|
||||
]);
|
||||
let body = AZURE_DOCUMENT_INTELLIGENCE_OCR_CONFIG
|
||||
.transform_ocr_payload("prebuilt-layout", document, Map::new())
|
||||
.unwrap()
|
||||
.data;
|
||||
let shared = body["base64Source"].as_text().unwrap();
|
||||
assert_eq!(shared.as_str(), "AQIDBA==");
|
||||
assert_eq!(shared.bytes().as_ptr(), source[offset..].as_ptr());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn azure_ai_reuses_mistral_body_transform() {
|
||||
let body = AZURE_AI_OCR_CONFIG
|
||||
|
|
@ -451,7 +466,7 @@ mod tests {
|
|||
.data;
|
||||
|
||||
assert_eq!(body["model"], "pixtral-12b-2409");
|
||||
assert_eq!(body["include_image_base64"], true);
|
||||
assert_eq!(body["include_image_base64"], json!(true));
|
||||
assert_eq!(
|
||||
body["document"]["document_url"],
|
||||
"data:application/pdf;base64,abc"
|
||||
|
|
|
|||
|
|
@ -1,3 +1,4 @@
|
|||
use crate::http_utils::body::JsonPayload;
|
||||
use serde_json::{Map, Value, json};
|
||||
|
||||
use crate::audio_transcription::transformation::{
|
||||
|
|
@ -6,7 +7,7 @@ use crate::audio_transcription::transformation::{
|
|||
use crate::audio_transcription::types::{
|
||||
AudioTranscriptionRequestData, AudioTranscriptionResponseData,
|
||||
};
|
||||
use crate::error::{Error, json_type_name};
|
||||
use crate::error::Error;
|
||||
|
||||
pub use super::aws_base::{aws_auth_config, bedrock_model_id_and_region, resolve_bedrock_region};
|
||||
use super::constants::{BEDROCK_RUNTIME_ENDPOINT_TEMPLATE, BEDROCK_SERVICE};
|
||||
|
|
@ -18,24 +19,24 @@ pub static BEDROCK_AUDIO_TRANSCRIPTION_CONFIG: BedrockAudioTranscriptionConfig =
|
|||
|
||||
pub struct BedrockAudioTranscriptionConfig;
|
||||
|
||||
fn audio_fields(audio: Value) -> Result<(String, String), Error> {
|
||||
let object = audio.as_object().ok_or_else(|| Error::InvalidType {
|
||||
expected: "object",
|
||||
actual: json_type_name(&audio),
|
||||
})?;
|
||||
fn audio_fields(audio: JsonPayload) -> Result<(JsonPayload, String), Error> {
|
||||
let mut object = audio.into_object()?;
|
||||
let data = object
|
||||
.get("data")
|
||||
.and_then(Value::as_str)
|
||||
.filter(|value| !value.is_empty())
|
||||
.remove("data")
|
||||
.filter(|value| match value {
|
||||
JsonPayload::String(text) => !text.as_str().is_empty(),
|
||||
JsonPayload::Base64(bytes) => !bytes.is_empty(),
|
||||
_ => false,
|
||||
})
|
||||
.ok_or(Error::MissingField("audio.data"))?;
|
||||
let format = object
|
||||
.get("format")
|
||||
.and_then(Value::as_str)
|
||||
.and_then(JsonPayload::as_str)
|
||||
.filter(|value| matches!(*value, "wav" | "mp3" | "flac" | "ogg"))
|
||||
.ok_or_else(|| {
|
||||
Error::InvalidRequest("audio.format must be wav, mp3, flac, or ogg".to_string())
|
||||
Error::InvalidRequest("audio.format must be wav, mp3, flac, or ogg".into())
|
||||
})?;
|
||||
Ok((data.to_string(), format.to_string()))
|
||||
Ok((data, format.to_owned()))
|
||||
}
|
||||
|
||||
fn optional_string<'a>(params: &'a Map<String, Value>, key: &str) -> Option<&'a str> {
|
||||
|
|
@ -52,10 +53,10 @@ impl AudioTranscriptionProviderConfig for BedrockAudioTranscriptionConfig {
|
|||
}
|
||||
|
||||
#[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)]
|
||||
fn transform_transcription_request(
|
||||
fn transform_transcription_payload(
|
||||
&self,
|
||||
_model: &str,
|
||||
audio: Value,
|
||||
audio: JsonPayload,
|
||||
optional_params: Map<String, Value>,
|
||||
) -> Result<AudioTranscriptionRequestData, Error> {
|
||||
let (data, format) = audio_fields(audio)?;
|
||||
|
|
@ -71,17 +72,32 @@ impl AudioTranscriptionProviderConfig for BedrockAudioTranscriptionConfig {
|
|||
inference_config.insert("temperature".to_string(), temperature.clone());
|
||||
}
|
||||
Ok(AudioTranscriptionRequestData {
|
||||
body: json!({
|
||||
"messages": [{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{"audio": {"format": format, "source": {"bytes": data}}},
|
||||
{"text": instruction}
|
||||
]
|
||||
}],
|
||||
"system": [{"text": "You are a transcription assistant."}],
|
||||
"inferenceConfig": inference_config,
|
||||
}),
|
||||
body: JsonPayload::object([
|
||||
(
|
||||
"messages",
|
||||
JsonPayload::Array(vec![JsonPayload::object([
|
||||
("role", "user".into()),
|
||||
(
|
||||
"content",
|
||||
JsonPayload::Array(vec![
|
||||
JsonPayload::object([(
|
||||
"audio",
|
||||
JsonPayload::object([
|
||||
("format", format.into()),
|
||||
("source", JsonPayload::object([("bytes", data)])),
|
||||
]),
|
||||
)]),
|
||||
JsonPayload::object([("text", instruction.into())]),
|
||||
]),
|
||||
),
|
||||
])]),
|
||||
),
|
||||
(
|
||||
"system",
|
||||
json!([{"text": "You are a transcription assistant."}]).into(),
|
||||
),
|
||||
("inferenceConfig", Value::Object(inference_config).into()),
|
||||
]),
|
||||
})
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -439,6 +439,42 @@ pub fn sign_bedrock_post(
|
|||
region: &str,
|
||||
credentials: &Credentials,
|
||||
signing_time: SystemTime,
|
||||
) -> Result<BTreeMap<String, String>, Error> {
|
||||
sign_bedrock_body(
|
||||
url,
|
||||
SignableBody::Bytes(body),
|
||||
headers,
|
||||
region,
|
||||
credentials,
|
||||
signing_time,
|
||||
)
|
||||
}
|
||||
|
||||
pub fn sign_bedrock_digest(
|
||||
url: &str,
|
||||
digest: &str,
|
||||
headers: &BTreeMap<String, String>,
|
||||
region: &str,
|
||||
credentials: &Credentials,
|
||||
signing_time: SystemTime,
|
||||
) -> Result<BTreeMap<String, String>, Error> {
|
||||
sign_bedrock_body(
|
||||
url,
|
||||
SignableBody::Precomputed(digest.to_owned()),
|
||||
headers,
|
||||
region,
|
||||
credentials,
|
||||
signing_time,
|
||||
)
|
||||
}
|
||||
|
||||
fn sign_bedrock_body(
|
||||
url: &str,
|
||||
body: SignableBody<'_>,
|
||||
headers: &BTreeMap<String, String>,
|
||||
region: &str,
|
||||
credentials: &Credentials,
|
||||
signing_time: SystemTime,
|
||||
) -> Result<BTreeMap<String, String>, Error> {
|
||||
let identity: Identity = credentials.clone().into();
|
||||
let params = v4::SigningParams::builder()
|
||||
|
|
@ -453,7 +489,7 @@ pub fn sign_bedrock_post(
|
|||
let header_refs = headers
|
||||
.iter()
|
||||
.map(|(name, value)| (name.as_str(), value.as_str()));
|
||||
let request = SignableRequest::new("POST", url, header_refs, SignableBody::Bytes(body))
|
||||
let request = SignableRequest::new("POST", url, header_refs, body)
|
||||
.map_err(|error| Error::Auth(format!("AWS signable request failed: {error}")))?;
|
||||
let (instructions, _) = sign(request, ¶ms)
|
||||
.map_err(|error| Error::Auth(format!("AWS request signing failed: {error}")))?
|
||||
|
|
@ -849,6 +885,39 @@ mod tests {
|
|||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn streamed_raw_and_encoded_audio_sign_like_buffered_json() {
|
||||
use crate::http_utils::body::JsonPayload;
|
||||
use crate::http_utils::body::PreparedJsonBody;
|
||||
use bytes::Bytes;
|
||||
let credentials = Credentials::new("access", "secret", Some("token".into()), None, "test");
|
||||
let headers = BTreeMap::from([("Content-Type".into(), "application/json".into())]);
|
||||
let url = "https://bedrock-runtime.us-east-1.amazonaws.com/model/test/converse";
|
||||
for payload in [
|
||||
JsonPayload::Base64(Bytes::from(vec![5; 131_073])),
|
||||
"AQI=".into(),
|
||||
] {
|
||||
let payload = JsonPayload::object([("audio", payload)]);
|
||||
let buffered = serde_json::to_vec(&payload).unwrap();
|
||||
let streamed = PreparedJsonBody::streamed(payload).unwrap();
|
||||
assert_eq!(streamed.content_length(), buffered.len() as u64);
|
||||
let time = UNIX_EPOCH + std::time::Duration::from_secs(1_704_164_645);
|
||||
assert_eq!(
|
||||
sign_bedrock_digest(
|
||||
url,
|
||||
&streamed.sha256(),
|
||||
&headers,
|
||||
"us-east-1",
|
||||
&credentials,
|
||||
time
|
||||
)
|
||||
.unwrap(),
|
||||
sign_bedrock_post(url, &buffered, &headers, "us-east-1", &credentials, time)
|
||||
.unwrap()
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn signing_without_session_token_omits_security_header() {
|
||||
let (url, body, headers) = parity_inputs();
|
||||
|
|
|
|||
|
|
@ -1,7 +1,9 @@
|
|||
use crate::error::{Error, json_type_name};
|
||||
use crate::http_utils::body::JsonPayload;
|
||||
use crate::ocr::transformation::OcrProviderConfig;
|
||||
use crate::ocr::types::{OcrRequestData, OcrResponseData};
|
||||
use serde_json::{Map, Value};
|
||||
use std::collections::BTreeMap;
|
||||
|
||||
const SUPPORTED_OCR_PARAMS: &[&str] = &[
|
||||
"pages",
|
||||
|
|
@ -76,28 +78,28 @@ impl OcrProviderConfig for MistralOcrConfig {
|
|||
}
|
||||
|
||||
#[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)]
|
||||
fn transform_ocr_request(
|
||||
fn transform_ocr_payload(
|
||||
&self,
|
||||
model: &str,
|
||||
document: Value,
|
||||
document: JsonPayload,
|
||||
optional_params: Map<String, Value>,
|
||||
) -> Result<OcrRequestData, Error> {
|
||||
if !document.is_object() {
|
||||
return Err(Error::InvalidType {
|
||||
expected: "object",
|
||||
actual: json_type_name(&document),
|
||||
actual: document.type_name(),
|
||||
});
|
||||
}
|
||||
|
||||
let mut data = Map::new();
|
||||
data.insert("model".to_string(), Value::String(model.to_string()));
|
||||
let mut data = BTreeMap::new();
|
||||
data.insert("model".to_string(), JsonPayload::from(model));
|
||||
data.insert("document".to_string(), document);
|
||||
for (param, value) in optional_params {
|
||||
data.insert(param, value);
|
||||
data.insert(param, value.into());
|
||||
}
|
||||
|
||||
Ok(OcrRequestData {
|
||||
data: Value::Object(data),
|
||||
data: JsonPayload::Object(data),
|
||||
files: None,
|
||||
})
|
||||
}
|
||||
|
|
@ -301,9 +303,12 @@ mod tests {
|
|||
json!({param: value}).as_object().unwrap().clone(),
|
||||
)
|
||||
.expect("request should transform");
|
||||
assert_eq!(result.data.get(param), Some(&value));
|
||||
assert_eq!(result.data.get("model"), Some(&json!("mistral-ocr-latest")));
|
||||
assert_eq!(result.data.get("document"), Some(&document));
|
||||
assert_eq!(result.data.materialize().get(param), Some(&value));
|
||||
assert_eq!(
|
||||
result.data.materialize().get("model"),
|
||||
Some(&json!("mistral-ocr-latest"))
|
||||
);
|
||||
assert_eq!(result.data.materialize().get("document"), Some(&document));
|
||||
assert_eq!(result.files, None);
|
||||
}
|
||||
}
|
||||
|
|
@ -324,12 +329,21 @@ mod tests {
|
|||
.clone();
|
||||
let result = transform_ocr_request("mistral-ocr-latest", document, optional_params)
|
||||
.expect("request should transform");
|
||||
assert_eq!(result.data.get("table_format"), Some(&json!("html")));
|
||||
assert_eq!(
|
||||
result.data.get("confidence_scores_granularity"),
|
||||
result.data.materialize().get("table_format"),
|
||||
Some(&json!("html"))
|
||||
);
|
||||
assert_eq!(
|
||||
result
|
||||
.data
|
||||
.materialize()
|
||||
.get("confidence_scores_granularity"),
|
||||
Some(&json!("page"))
|
||||
);
|
||||
assert_eq!(result.data.get("extract_header"), Some(&json!(true)));
|
||||
assert_eq!(
|
||||
result.data.materialize().get("extract_header"),
|
||||
Some(&json!(true))
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
|
|
|
|||
|
|
@ -1,7 +1,9 @@
|
|||
use crate::error::{Error, json_type_name};
|
||||
use crate::http_utils::body::JsonPayload;
|
||||
use crate::ocr::transformation::OcrProviderConfig;
|
||||
use crate::ocr::types::{OcrRequestData, OcrResponseData};
|
||||
use serde_json::{Map, Value, json};
|
||||
use std::collections::BTreeMap;
|
||||
|
||||
use crate::providers::mistral::ocr::transformation::MISTRAL_OCR_CONFIG;
|
||||
|
||||
|
|
@ -125,14 +127,14 @@ pub fn complete_vertex_deepseek_url(
|
|||
))
|
||||
}
|
||||
|
||||
fn document_content_item(document: &Value) -> Result<Value, Error> {
|
||||
fn document_content_item(document: &JsonPayload) -> Result<JsonPayload, Error> {
|
||||
let object = document.as_object().ok_or_else(|| Error::InvalidType {
|
||||
expected: "object",
|
||||
actual: json_type_name(document),
|
||||
actual: document.type_name(),
|
||||
})?;
|
||||
let doc_type = object
|
||||
.get("type")
|
||||
.and_then(Value::as_str)
|
||||
.and_then(JsonPayload::as_str)
|
||||
.ok_or(Error::MissingField("document.type"))?;
|
||||
let url_field = match doc_type {
|
||||
"image_url" => "image_url",
|
||||
|
|
@ -143,16 +145,22 @@ fn document_content_item(document: &Value) -> Result<Value, Error> {
|
|||
)));
|
||||
}
|
||||
};
|
||||
let url = object
|
||||
let _url = object
|
||||
.get(url_field)
|
||||
.and_then(Value::as_str)
|
||||
.and_then(JsonPayload::as_str)
|
||||
.filter(|value| !value.is_empty())
|
||||
.ok_or(Error::MissingField(url_field))?;
|
||||
|
||||
Ok(json!({
|
||||
"type": "image_url",
|
||||
"image_url": url,
|
||||
}))
|
||||
Ok(JsonPayload::object([
|
||||
("type", "image_url".into()),
|
||||
(
|
||||
"image_url",
|
||||
object
|
||||
.get(url_field)
|
||||
.ok_or(Error::MissingField(url_field))?
|
||||
.clone(),
|
||||
),
|
||||
]))
|
||||
}
|
||||
|
||||
fn deepseek_model_name(model: &str) -> String {
|
||||
|
|
@ -212,13 +220,13 @@ impl OcrProviderConfig for VertexAiOcrConfig {
|
|||
MISTRAL_OCR_CONFIG.supported_ocr_params()
|
||||
}
|
||||
|
||||
fn transform_ocr_request(
|
||||
fn transform_ocr_payload(
|
||||
&self,
|
||||
model: &str,
|
||||
document: Value,
|
||||
document: JsonPayload,
|
||||
optional_params: Map<String, Value>,
|
||||
) -> Result<OcrRequestData, Error> {
|
||||
MISTRAL_OCR_CONFIG.transform_ocr_request(model, document, optional_params)
|
||||
MISTRAL_OCR_CONFIG.transform_ocr_payload(model, document, optional_params)
|
||||
}
|
||||
|
||||
fn transform_ocr_response(
|
||||
|
|
@ -257,28 +265,35 @@ impl OcrProviderConfig for VertexAiDeepSeekOcrConfig {
|
|||
DEEPSEEK_SUPPORTED_OCR_PARAMS
|
||||
}
|
||||
|
||||
fn transform_ocr_request(
|
||||
fn transform_ocr_payload(
|
||||
&self,
|
||||
model: &str,
|
||||
document: Value,
|
||||
document: JsonPayload,
|
||||
optional_params: Map<String, Value>,
|
||||
) -> Result<OcrRequestData, Error> {
|
||||
let mut data = Map::new();
|
||||
data.insert(
|
||||
"model".to_string(),
|
||||
Value::String(deepseek_model_name(model)),
|
||||
);
|
||||
data.insert(
|
||||
"messages".to_string(),
|
||||
json!([{"role": "user", "content": [document_content_item(&document)?]}]),
|
||||
);
|
||||
let mut data = BTreeMap::from([
|
||||
(
|
||||
"model".into(),
|
||||
JsonPayload::from(deepseek_model_name(model)),
|
||||
),
|
||||
(
|
||||
"messages".into(),
|
||||
JsonPayload::Array(vec![JsonPayload::object([
|
||||
("role", "user".into()),
|
||||
(
|
||||
"content",
|
||||
JsonPayload::Array(vec![document_content_item(&document)?]),
|
||||
),
|
||||
])]),
|
||||
),
|
||||
]);
|
||||
for (key, value) in optional_params {
|
||||
if DEEPSEEK_SUPPORTED_OCR_PARAMS.contains(&key.as_str()) {
|
||||
data.insert(key, value);
|
||||
data.insert(key, value.into());
|
||||
}
|
||||
}
|
||||
Ok(OcrRequestData {
|
||||
data: Value::Object(data),
|
||||
data: JsonPayload::Object(data),
|
||||
files: None,
|
||||
})
|
||||
}
|
||||
|
|
@ -404,7 +419,7 @@ mod tests {
|
|||
.data;
|
||||
|
||||
assert_eq!(body["model"], "deepseek-ai/deepseek-ocr-maas");
|
||||
assert_eq!(body["temperature"], 0.1);
|
||||
assert_eq!(body["temperature"], json!(0.1));
|
||||
assert_eq!(
|
||||
body["messages"][0]["content"][0],
|
||||
json!({"type": "image_url", "image_url": "gs://bucket/doc.pdf"})
|
||||
|
|
|
|||
|
|
@ -16,6 +16,7 @@ extension-module = ["pyo3/extension-module"]
|
|||
panic-test = []
|
||||
|
||||
[dependencies]
|
||||
bytes.workspace = true
|
||||
futures-util.workspace = true
|
||||
tracing.workspace = true
|
||||
tracing-subscriber.workspace = true
|
||||
|
|
@ -29,6 +30,8 @@ serde_json.workspace = true
|
|||
tokio.workspace = true
|
||||
|
||||
[dev-dependencies]
|
||||
reqwest.workspace = true
|
||||
sha2.workspace = true
|
||||
criterion = "0.8.2"
|
||||
tokio-tungstenite.workspace = true
|
||||
|
||||
|
|
|
|||
55
litellm-rust/crates/python-bridge/benches/media/README.md
Normal file
55
litellm-rust/crates/python-bridge/benches/media/README.md
Normal file
|
|
@ -0,0 +1,55 @@
|
|||
# Media request memory benchmark
|
||||
|
||||
The refactor keeps existing encoded media in shared buffers through Python extraction, provider transformation and outgoing chunks. Raw audio is encoded in bounded chunks. The benchmark includes request preparation, AWS signing and transmission to a local HTTP sink that hashes each incoming chunk without collecting the body
|
||||
|
||||
## Method
|
||||
|
||||
Run on macOS arm64 with Rust 1.98.0 and CPython 3.11, release profile, on 2026-09-02. Each scenario uses distinct Python inputs of 1, 16 or 64 MiB at concurrency 1 or 16. Inputs stay alive until all requests finish. Encoded sizes describe existing ASCII base64 strings; raw sizes describe unencoded bytes and produce about 4/3 as many wire bytes
|
||||
|
||||
`python` uses Python base64, JSON, botocore signing and concurrent `http.client` calls. It constructs the same Bedrock transcription body as Rust, without the rest of the Python SDK. `buffered` extracts through the existing Python-to-Serde boundary, runs the Rust provider transform and allocates the full outgoing JSON body. `current` is the pre-refactor shared-body implementation saved before this pass. `refactor` uses ByteString, the smaller body module and the explicit transport client
|
||||
|
||||
Rust cases call the actual bridge media extractor and Bedrock transcription transform. They use fixed test credentials and signing time. Python uses botocore with the same test credentials and the current signing time. All approaches hash the payload once for signing, then the sink independently hashes received bytes. All four produce identical body lengths and hashes for each scenario
|
||||
|
||||
Timing, RSS and allocation measurements run in separate fresh processes. CPU and throughput are medians of three sequential timing runs, interleaved by approach, after builds and test processes finished. CPU includes both the client and local sink threads. Throughput includes extraction, transforms, preparation, signing and sending; it is not remote-provider throughput. Per-stage times are wall-clock medians
|
||||
|
||||
RSS includes the interpreter, retained inputs, Rust runtime, HTTP client and sink. `input_rss_mib` is the process high-water RSS after input setup, not live RSS. The Rust allocator counter measures cumulative requested bytes, including reallocations, in the client and sink after setup. Python allocations are measured separately with tracemalloc's peak live traced bytes; these two allocation columns are different quantities and must not be added. Native allocation totals are not a peak-memory metric: bounded encoding can allocate many successive chunks while keeping only a few live
|
||||
|
||||
No TLS, remote download, guardrail materialization, retries or response payload optimization is included in these timings. Local transport tests separately cover redirects, retries, disconnects and stalled-consumer cancellation. This benchmark is a focused body pipeline comparison, not a complete LiteLLM deployment memory profile
|
||||
|
||||
## Results
|
||||
|
||||
At 64 MiB per input and concurrency 16, the retained inputs alone total 1024 MiB
|
||||
|
||||
| Input | Approach | Peak RSS, MiB | CPU seconds | Wire MiB/s |
|
||||
|---|---|---:|---:|---:|
|
||||
| encoded | python | 2150.9 | 6.14 | 350.2 |
|
||||
| encoded | buffered | 3111.5 | 6.18 | 283.4 |
|
||||
| encoded | current | 1062.8 | 8.34 | 225.0 |
|
||||
| encoded | refactor | 1063.3 | 7.44 | 247.9 |
|
||||
| raw | python | 3879.5 | 9.40 | 267.9 |
|
||||
| raw | buffered | 3879.7 | 9.64 | 237.9 |
|
||||
| raw | current | 1077.1 | 8.96 | 312.6 |
|
||||
| raw | refactor | 1081.0 | 9.46 | 279.1 |
|
||||
|
||||
The full matrix, allocation measurements, stage timings and body hashes are in [results.csv](results.csv)
|
||||
|
||||
The refactor counts its emitted chunks to establish Content-Length. Raw media is therefore encoded during length calculation, signing and transmission, one more pass than the pre-refactor implementation. Existing encoded media uses shared slices for all three passes. Encoding and escaping chunks remain at most 64 KiB; their live memory does not scale with the payload
|
||||
|
||||
## Reproduction
|
||||
|
||||
The existing serialization benchmark also supports one-shot media measurements. The embedded Python interpreter needs botocore and typing_extensions. Set `PYO3_PYTHON` at build time and, when using a virtual environment, make its site-packages available through `PYTHONPATH`
|
||||
|
||||
```bash
|
||||
cd litellm-rust
|
||||
cargo bench -p litellm-python-bridge --bench serialization -- --media refactor 64 16 raw memory
|
||||
cargo bench -p litellm-python-bridge --bench serialization -- --media refactor 64 16 raw allocation
|
||||
cargo bench -p litellm-python-bridge --bench serialization -- --media refactor 64 16 raw timing
|
||||
```
|
||||
|
||||
Approaches are `python`, `buffered`, or a shared-body label such as `refactor`. Input modes are `encoded` and `raw`. Repeat for 1, 16 and 64 MiB, and concurrency 1 and 16. For timing comparisons, build first, identify the executable from Cargo's `compiler-artifact` output, and invoke it directly in fresh processes. Record the executable hash when comparing different work-in-progress versions
|
||||
|
||||
The pre-refactor executable was captured from the uncommitted implementation based on staging `e058aa68c4`, before the ByteString and typed-extraction changes. Its source snapshot remained separate while the branch was refactored in place. Executable SHA-256 values used for this report:
|
||||
|
||||
Pre-refactor Rust: `4874abe92aeb9e1ed77bef6668c11f4ca545458c2bc189a18404ad61cd804327`
|
||||
|
||||
Refactor and Python: `c551eaef47e7a7eceecfed8b4ab76a42d84b8392342be01e8a22f25cce658287`
|
||||
266
litellm-rust/crates/python-bridge/benches/media/mod.rs
Normal file
266
litellm-rust/crates/python-bridge/benches/media/mod.rs
Normal file
|
|
@ -0,0 +1,266 @@
|
|||
use std::alloc::{GlobalAlloc, Layout, System};
|
||||
use std::collections::BTreeMap;
|
||||
use std::ffi::CString;
|
||||
use std::io::{BufRead, BufReader, Read, Write};
|
||||
use std::net::TcpListener;
|
||||
use std::sync::atomic::{AtomicBool, AtomicU64, Ordering};
|
||||
use std::time::{Duration, Instant, UNIX_EPOCH};
|
||||
|
||||
use litellm_core::audio_transcription::transformation::AudioTranscriptionProviderConfig;
|
||||
use litellm_core::http_utils::body::PreparedJsonBody;
|
||||
use litellm_core::http_utils::replay::send_json;
|
||||
use litellm_core::providers::bedrock::audio_transcription::BEDROCK_AUDIO_TRANSCRIPTION_CONFIG;
|
||||
use litellm_core::providers::bedrock::aws_base::{
|
||||
AwsAuthConfig, resolve_credentials, sign_bedrock_digest,
|
||||
};
|
||||
use litellm_python_interop::from_py;
|
||||
use pyo3::prelude::*;
|
||||
use pyo3::types::PyModule;
|
||||
use serde_json::{Value, json};
|
||||
use sha2::{Digest, Sha256};
|
||||
|
||||
#[path = "../../src/payload.rs"]
|
||||
mod payload;
|
||||
|
||||
struct Allocator;
|
||||
static TRACK: AtomicBool = AtomicBool::new(false);
|
||||
static ALLOCATED: AtomicU64 = AtomicU64::new(0);
|
||||
|
||||
#[global_allocator]
|
||||
static ALLOCATOR: Allocator = Allocator;
|
||||
|
||||
unsafe impl GlobalAlloc for Allocator {
|
||||
unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
|
||||
if TRACK.load(Ordering::Relaxed) {
|
||||
ALLOCATED.fetch_add(layout.size() as u64, Ordering::Relaxed);
|
||||
}
|
||||
unsafe { System.alloc(layout) }
|
||||
}
|
||||
unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
|
||||
unsafe { System.dealloc(ptr, layout) }
|
||||
}
|
||||
unsafe fn realloc(&self, ptr: *mut u8, layout: Layout, size: usize) -> *mut u8 {
|
||||
if TRACK.load(Ordering::Relaxed) {
|
||||
ALLOCATED.fetch_add(size as u64, Ordering::Relaxed);
|
||||
}
|
||||
unsafe { System.realloc(ptr, layout, size) }
|
||||
}
|
||||
}
|
||||
|
||||
fn sink(concurrency: usize) -> (String, std::thread::JoinHandle<Vec<(usize, String)>>) {
|
||||
let listener = TcpListener::bind("127.0.0.1:0").unwrap();
|
||||
let url = format!(
|
||||
"http://{}/model/benchmark/converse",
|
||||
listener.local_addr().unwrap()
|
||||
);
|
||||
let thread = std::thread::spawn(move || {
|
||||
let workers: Vec<_> = (0..concurrency)
|
||||
.map(|_| {
|
||||
let (socket, _) = listener.accept().unwrap();
|
||||
std::thread::spawn(move || {
|
||||
let mut reader = BufReader::new(socket);
|
||||
let mut length = 0;
|
||||
loop {
|
||||
let mut line = String::new();
|
||||
assert_ne!(reader.read_line(&mut line).unwrap(), 0);
|
||||
if line == "\r\n" {
|
||||
break;
|
||||
}
|
||||
if let Some(value) = line.to_lowercase().strip_prefix("content-length:") {
|
||||
length = value.trim().parse::<usize>().unwrap();
|
||||
}
|
||||
}
|
||||
let mut digest = Sha256::new();
|
||||
let mut remaining = length;
|
||||
let mut chunk = [0; 65536];
|
||||
while remaining > 0 {
|
||||
let count = remaining.min(chunk.len());
|
||||
reader.read_exact(&mut chunk[..count]).unwrap();
|
||||
digest.update(&chunk[..count]);
|
||||
remaining -= count;
|
||||
}
|
||||
let digest = format!("{:x}", digest.finalize());
|
||||
write!(
|
||||
reader.get_mut(),
|
||||
"HTTP/1.1 200 OK\r\nContent-Length: {}\r\nConnection: close\r\n\r\n{}",
|
||||
digest.len(),
|
||||
digest
|
||||
)
|
||||
.unwrap();
|
||||
(length, digest)
|
||||
})
|
||||
})
|
||||
.collect();
|
||||
workers
|
||||
.into_iter()
|
||||
.map(|worker| worker.join().unwrap())
|
||||
.collect()
|
||||
});
|
||||
(url, thread)
|
||||
}
|
||||
|
||||
pub fn run() {
|
||||
let args: Vec<_> = std::env::args().collect();
|
||||
let approach = &args[2];
|
||||
let mib: usize = args[3].parse().unwrap();
|
||||
let concurrency: usize = args[4].parse().unwrap();
|
||||
let encoding = &args[5];
|
||||
let measurement = &args[6];
|
||||
Python::initialize();
|
||||
let runtime = tokio::runtime::Builder::new_multi_thread()
|
||||
.worker_threads(4)
|
||||
.enable_all()
|
||||
.build()
|
||||
.unwrap();
|
||||
let client = reqwest::Client::builder()
|
||||
.no_proxy()
|
||||
.redirect(reqwest::redirect::Policy::none())
|
||||
.retry(reqwest::retry::never())
|
||||
.build()
|
||||
.unwrap();
|
||||
let (url, server) = sink(concurrency);
|
||||
let credentials = runtime
|
||||
.block_on(resolve_credentials(
|
||||
AwsAuthConfig {
|
||||
region_name: Some("us-east-1".into()),
|
||||
access_key_id: Some("benchmark".into()),
|
||||
secret_access_key: Some("benchmark".into()),
|
||||
..Default::default()
|
||||
},
|
||||
&|_| None,
|
||||
))
|
||||
.unwrap();
|
||||
let module = Python::attach(|py| {
|
||||
PyModule::from_code(
|
||||
py,
|
||||
&CString::new(include_str!("pipeline.py")).unwrap(),
|
||||
c"pipeline.py",
|
||||
c"pipeline",
|
||||
)
|
||||
.unwrap()
|
||||
.unbind()
|
||||
});
|
||||
let inputs = Python::attach(|py| {
|
||||
module
|
||||
.bind(py)
|
||||
.call_method1("inputs", (mib * 1024 * 1024, concurrency, encoding))
|
||||
.unwrap()
|
||||
.unbind()
|
||||
});
|
||||
let before: (f64, u64) = Python::attach(|py| {
|
||||
module
|
||||
.bind(py)
|
||||
.call_method1("start", (measurement,))
|
||||
.unwrap()
|
||||
.extract()
|
||||
.unwrap()
|
||||
});
|
||||
TRACK.store(measurement == "allocation", Ordering::Relaxed);
|
||||
let started = Instant::now();
|
||||
let stages = if approach == "python" {
|
||||
Python::attach(|py| {
|
||||
module
|
||||
.bind(py)
|
||||
.call_method1("run", (inputs.bind(py), &url))
|
||||
.unwrap()
|
||||
.extract::<Vec<f64>>()
|
||||
.unwrap()
|
||||
})
|
||||
} else {
|
||||
let extracted = Python::attach(|py| {
|
||||
inputs
|
||||
.bind(py)
|
||||
.try_iter()
|
||||
.unwrap()
|
||||
.map(|input| {
|
||||
let input = input.unwrap();
|
||||
if approach == "buffered" {
|
||||
let encoded = module.bind(py).call_method1("encode", (&input,)).unwrap();
|
||||
from_py::<Value>(&encoded).unwrap().into()
|
||||
} else if encoding == "raw" {
|
||||
payload::audio_payload_from_py(&input).unwrap()
|
||||
} else {
|
||||
payload::payload_from_py(&input).unwrap()
|
||||
}
|
||||
})
|
||||
.collect::<Vec<_>>()
|
||||
});
|
||||
let extraction = started.elapsed().as_secs_f64();
|
||||
let transformed = extracted
|
||||
.into_iter()
|
||||
.map(|input| {
|
||||
BEDROCK_AUDIO_TRANSCRIPTION_CONFIG
|
||||
.transform_transcription_payload("benchmark", input, Default::default())
|
||||
.unwrap()
|
||||
.body
|
||||
})
|
||||
.collect::<Vec<_>>();
|
||||
let transform = started.elapsed().as_secs_f64();
|
||||
let bodies = transformed
|
||||
.into_iter()
|
||||
.map(|body| {
|
||||
if approach == "buffered" {
|
||||
PreparedJsonBody::buffered(serde_json::to_vec(&body).unwrap().into())
|
||||
} else {
|
||||
PreparedJsonBody::new(body).unwrap()
|
||||
}
|
||||
})
|
||||
.collect::<Vec<_>>();
|
||||
let preparation = started.elapsed().as_secs_f64();
|
||||
let signed = bodies
|
||||
.iter()
|
||||
.map(|body| {
|
||||
let digest = body.sha256();
|
||||
let headers = sign_bedrock_digest(
|
||||
&url,
|
||||
&digest,
|
||||
&BTreeMap::new(),
|
||||
"us-east-1",
|
||||
&credentials,
|
||||
UNIX_EPOCH + Duration::from_secs(1_700_000_000),
|
||||
)
|
||||
.unwrap();
|
||||
(digest, headers.into_iter().collect::<Vec<_>>())
|
||||
})
|
||||
.collect::<Vec<_>>();
|
||||
let signing = started.elapsed().as_secs_f64();
|
||||
runtime.block_on(async {
|
||||
let responses = futures_util::future::join_all(bodies.iter().zip(&signed).map(
|
||||
|(body, (digest, headers))| async {
|
||||
let response =
|
||||
send_json(&client, &url, body, headers, Duration::from_secs(120), None)
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(response.text().await.unwrap(), *digest);
|
||||
},
|
||||
))
|
||||
.await;
|
||||
std::hint::black_box(responses);
|
||||
});
|
||||
vec![
|
||||
extraction,
|
||||
transform - extraction,
|
||||
preparation - transform,
|
||||
signing - preparation,
|
||||
started.elapsed().as_secs_f64() - signing,
|
||||
]
|
||||
};
|
||||
let elapsed = started.elapsed().as_secs_f64();
|
||||
TRACK.store(false, Ordering::Relaxed);
|
||||
let allocated = ALLOCATED.load(Ordering::Relaxed);
|
||||
let after: (f64, u64, u64) = Python::attach(|py| {
|
||||
module
|
||||
.bind(py)
|
||||
.call_method1("finish", (measurement,))
|
||||
.unwrap()
|
||||
.extract()
|
||||
.unwrap()
|
||||
});
|
||||
let received = server.join().unwrap();
|
||||
let wire_bytes: usize = received.iter().map(|(length, _)| length).sum();
|
||||
assert!(received.iter().all(|(_, hash)| hash == &received[0].1));
|
||||
println!(
|
||||
"{}",
|
||||
json!({"approach":approach,"mib":mib,"concurrency":concurrency,"encoding":encoding,"measurement":measurement,"seconds":elapsed,"cpu_seconds":after.0-before.0,"peak_rss_bytes":after.1,"input_rss_bytes":before.1,"rust_allocated_bytes":allocated,"python_peak_traced_bytes":after.2,"wire_mib_per_second":wire_bytes as f64/1048576.0/elapsed,"stage_seconds":stages,"sha256":received[0].1})
|
||||
);
|
||||
}
|
||||
117
litellm-rust/crates/python-bridge/benches/media/pipeline.py
Normal file
117
litellm-rust/crates/python-bridge/benches/media/pipeline.py
Normal file
|
|
@ -0,0 +1,117 @@
|
|||
import base64
|
||||
import concurrent.futures
|
||||
import hashlib
|
||||
import http.client
|
||||
import json
|
||||
import resource
|
||||
import sys
|
||||
import time
|
||||
import tracemalloc
|
||||
from typing import Final, TypeAlias, TypedDict
|
||||
from urllib.parse import urlsplit
|
||||
|
||||
from botocore.auth import SigV4Auth
|
||||
from botocore.awsrequest import AWSRequest
|
||||
from botocore.credentials import Credentials
|
||||
from typing_extensions import ReadOnly
|
||||
|
||||
Json: TypeAlias = bool | int | float | str | list["Json"] | dict[str, "Json"] | None
|
||||
|
||||
|
||||
class AudioInput(TypedDict):
|
||||
data: ReadOnly[str | bytes]
|
||||
format: ReadOnly[str]
|
||||
|
||||
|
||||
class EncodedAudio(TypedDict):
|
||||
data: ReadOnly[str]
|
||||
format: ReadOnly[str]
|
||||
|
||||
|
||||
def inputs(size: int, concurrency: int, encoding: str) -> tuple[AudioInput, ...]:
|
||||
return tuple(
|
||||
{"data": b"a" * size if encoding == "raw" else "A" * size, "format": "wav"} for _ in range(concurrency)
|
||||
)
|
||||
|
||||
|
||||
def encode(audio: AudioInput) -> EncodedAudio:
|
||||
data: Final = audio["data"]
|
||||
return {
|
||||
"data": base64.b64encode(data).decode("ascii") if isinstance(data, bytes) else data,
|
||||
"format": audio["format"],
|
||||
}
|
||||
|
||||
|
||||
def transform(audio: EncodedAudio) -> dict[str, Json]:
|
||||
return {
|
||||
"messages": [
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{"audio": {"format": audio["format"], "source": {"bytes": audio["data"]}}},
|
||||
{"text": "Transcribe the audio. Respond with only the transcript."},
|
||||
],
|
||||
}
|
||||
],
|
||||
"system": [{"text": "You are a transcription assistant."}],
|
||||
"inferenceConfig": {"maxTokens": 4096},
|
||||
}
|
||||
|
||||
|
||||
def stats() -> tuple[float, int]:
|
||||
usage: Final = resource.getrusage(resource.RUSAGE_SELF)
|
||||
return time.process_time(), usage.ru_maxrss * (1 if sys.platform == "darwin" else 1024)
|
||||
|
||||
|
||||
def start(measurement: str) -> tuple[float, int]:
|
||||
if measurement == "allocation":
|
||||
tracemalloc.start()
|
||||
return stats()
|
||||
|
||||
|
||||
def finish(measurement: str) -> tuple[float, int, int]:
|
||||
return (*stats(), tracemalloc.get_traced_memory()[1] if measurement == "allocation" else 0)
|
||||
|
||||
|
||||
def send(url: str, body: bytes, headers: dict[str, str], digest: str) -> None:
|
||||
target: Final = urlsplit(url)
|
||||
assert target.hostname is not None
|
||||
connection: Final = http.client.HTTPConnection(target.hostname, target.port, timeout=120)
|
||||
connection.request("POST", target.path, body=body, headers=headers)
|
||||
response: Final = connection.getresponse()
|
||||
assert response.status == 200
|
||||
assert response.read().decode() == digest
|
||||
connection.close()
|
||||
|
||||
|
||||
def run(audio: tuple[AudioInput, ...], url: str) -> list[float]:
|
||||
started: Final = time.perf_counter()
|
||||
encoded: Final = tuple(encode(item) for item in audio)
|
||||
extraction: Final = time.perf_counter()
|
||||
transformed: Final = tuple(transform(item) for item in encoded)
|
||||
transformation: Final = time.perf_counter()
|
||||
bodies: Final = tuple(
|
||||
json.dumps(item, separators=(",", ":"), ensure_ascii=False, sort_keys=True).encode() for item in transformed
|
||||
)
|
||||
preparation: Final = time.perf_counter()
|
||||
digests: Final = tuple(hashlib.sha256(body).hexdigest() for body in bodies)
|
||||
requests: Final = tuple(
|
||||
AWSRequest(method="POST", url=url, data=body, headers={"X-Amz-Content-SHA256": digest})
|
||||
for body, digest in zip(bodies, digests)
|
||||
)
|
||||
for request in requests:
|
||||
SigV4Auth(Credentials("benchmark", "benchmark"), "bedrock", "us-east-1").add_auth(request)
|
||||
signing: Final = time.perf_counter()
|
||||
with concurrent.futures.ThreadPoolExecutor(max_workers=len(bodies)) as executor:
|
||||
tuple(
|
||||
executor.map(
|
||||
lambda args: send(url, *args), zip(bodies, (dict(request.headers) for request in requests), digests)
|
||||
)
|
||||
)
|
||||
return [
|
||||
extraction - started,
|
||||
transformation - extraction,
|
||||
preparation - transformation,
|
||||
signing - preparation,
|
||||
time.perf_counter() - signing,
|
||||
]
|
||||
49
litellm-rust/crates/python-bridge/benches/media/results.csv
Normal file
49
litellm-rust/crates/python-bridge/benches/media/results.csv
Normal file
|
|
@ -0,0 +1,49 @@
|
|||
approach,encoding,input_mib,concurrency,input_rss_mib,peak_rss_mib,rust_allocated_mib,python_peak_traced_mib,cpu_seconds,wall_seconds,wire_mib_per_second,extraction_seconds,transform_seconds,preparation_seconds,signing_seconds,sending_seconds,body_sha256
|
||||
python,encoded,1,1,36.015625,38.5,0.008942,2.00469,0.007635,0.007596,131.680602,1e-06,1e-06,0.002508,0.00082,0.004382,f8307139b09453fd8af8ec4551aff42459dc60d0ba84090a5e223e1421e037b1
|
||||
buffered,encoded,1,1,35.515625,39.8125,4.070372,0.000585,0.006976,0.006722,148.81207,0.000109,1e-05,0.000588,0.002739,0.0033,f8307139b09453fd8af8ec4551aff42459dc60d0ba84090a5e223e1421e037b1
|
||||
current,encoded,1,1,35.6875,36.921875,0.10866,0.000532,0.008435,0.007915,126.376088,3.8e-05,1.3e-05,0.000858,0.003445,0.003544,f8307139b09453fd8af8ec4551aff42459dc60d0ba84090a5e223e1421e037b1
|
||||
refactor,encoded,1,1,36.109375,37.265625,0.08128,0.000418,0.008424,0.007869,127.107453,5.8e-05,1.3e-05,0.000739,0.003435,0.003662,f8307139b09453fd8af8ec4551aff42459dc60d0ba84090a5e223e1421e037b1
|
||||
python,encoded,1,16,51.578125,72.578125,0.144142,17.035325,0.100369,0.052824,302.962791,7e-06,9e-06,0.034603,0.008962,0.009126,f8307139b09453fd8af8ec4551aff42459dc60d0ba84090a5e223e1421e037b1
|
||||
buffered,encoded,1,16,51.0625,88.09375,65.147511,0.001463,0.103402,0.059328,269.750457,0.001444,4.4e-05,0.008051,0.043155,0.006447,f8307139b09453fd8af8ec4551aff42459dc60d0ba84090a5e223e1421e037b1
|
||||
current,encoded,1,16,50.828125,54.4375,1.318544,0.000624,0.124145,0.072433,220.947769,0.000556,4.9e-05,0.013589,0.052113,0.006186,f8307139b09453fd8af8ec4551aff42459dc60d0ba84090a5e223e1421e037b1
|
||||
refactor,encoded,1,16,51.4375,54.953125,1.321156,0.000624,0.11704,0.065675,243.682147,0.000874,4.9e-05,0.008303,0.050536,0.005896,f8307139b09453fd8af8ec4551aff42459dc60d0ba84090a5e223e1421e037b1
|
||||
python,encoded,16,1,51.09375,83.578125,0.008944,32.00469,0.089575,0.0881,181.614275,2e-06,2e-06,0.033938,0.008023,0.045926,0715e88725a53c5ee6fd0e1913e6593394af313e30063c64c1d193c938b754e2
|
||||
buffered,encoded,16,1,50.765625,84.109375,64.07048,0.000615,0.097848,0.095846,166.936747,0.001316,9e-06,0.007207,0.042597,0.044757,0715e88725a53c5ee6fd0e1913e6593394af313e30063c64c1d193c938b754e2
|
||||
current,encoded,16,1,50.53125,51.796875,0.108662,0.000563,0.120787,0.110464,144.846236,0.000435,1.3e-05,0.013541,0.051401,0.044892,0715e88725a53c5ee6fd0e1913e6593394af313e30063c64c1d193c938b754e2
|
||||
refactor,encoded,16,1,51.03125,52.21875,0.081282,0.000448,0.113516,0.104137,153.645635,0.000746,1.3e-05,0.007921,0.050322,0.045007,0715e88725a53c5ee6fd0e1913e6593394af313e30063c64c1d193c938b754e2
|
||||
python,encoded,16,16,292.125,567.1875,0.144173,272.035325,1.374041,0.728761,351.286222,1e-05,1.1e-05,0.549075,0.124654,0.050559,0715e88725a53c5ee6fd0e1913e6593394af313e30063c64c1d193c938b754e2
|
||||
buffered,encoded,16,16,291.1875,807.46875,1025.131947,0.001524,1.542114,0.90781,282.001427,0.031269,5.5e-05,0.120981,0.707183,0.050253,0715e88725a53c5ee6fd0e1913e6593394af313e30063c64c1d193c938b754e2
|
||||
current,encoded,16,16,291.046875,294.75,1.349794,0.000685,1.912532,1.152532,222.123063,0.009276,5.7e-05,0.221244,0.836345,0.066117,0715e88725a53c5ee6fd0e1913e6593394af313e30063c64c1d193c938b754e2
|
||||
refactor,encoded,16,16,291.5625,295.21875,1.305561,0.000685,1.779405,1.027048,249.261802,0.014938,5.3e-05,0.128839,0.82514,0.057373,0715e88725a53c5ee6fd0e1913e6593394af313e30063c64c1d193c938b754e2
|
||||
python,encoded,64,1,99.296875,227.765625,0.008944,128.00469,0.362755,0.360889,177.340738,3e-06,2e-06,0.146347,0.031805,0.180794,8f48a63c8caaf885f5477dc99e0841c3c671777812b6ad8aa6a45687f26f9449
|
||||
buffered,encoded,64,1,98.71875,227.96875,256.070374,0.000566,0.399823,0.393474,162.654323,0.005253,9e-06,0.031351,0.175565,0.18033,8f48a63c8caaf885f5477dc99e0841c3c671777812b6ad8aa6a45687f26f9449
|
||||
current,encoded,64,1,98.546875,99.84375,0.108662,0.000566,0.484495,0.44534,143.711044,0.001701,1.5e-05,0.055285,0.208477,0.178448,8f48a63c8caaf885f5477dc99e0841c3c671777812b6ad8aa6a45687f26f9449
|
||||
refactor,encoded,64,1,99.0625,100.234375,0.081282,0.000479,0.464134,0.427744,149.622678,0.003932,1.4e-05,0.035907,0.20703,0.180438,8f48a63c8caaf885f5477dc99e0841c3c671777812b6ad8aa6a45687f26f9449
|
||||
python,encoded,64,16,1059.8125,2150.875,0.144173,1088.035325,6.143728,2.92395,350.212522,1.3e-05,1.3e-05,2.188775,0.509833,0.234596,8f48a63c8caaf885f5477dc99e0841c3c671777812b6ad8aa6a45687f26f9449
|
||||
buffered,encoded,64,16,1059.234375,3111.53125,4097.116352,0.000736,6.180851,3.613478,283.384591,0.108224,4.9e-05,0.506384,2.804387,0.177961,8f48a63c8caaf885f5477dc99e0841c3c671777812b6ad8aa6a45687f26f9449
|
||||
current,encoded,64,16,1059.21875,1062.78125,1.318575,0.000689,8.338055,4.551881,224.962775,0.042777,6.5e-05,0.888808,3.344306,0.275328,8f48a63c8caaf885f5477dc99e0841c3c671777812b6ad8aa6a45687f26f9449
|
||||
refactor,encoded,64,16,1059.5625,1063.265625,1.321186,0.000689,7.441413,4.130159,247.933269,0.066971,5.7e-05,0.53399,3.280206,0.232016,8f48a63c8caaf885f5477dc99e0841c3c671777812b6ad8aa6a45687f26f9449
|
||||
python,raw,1,1,36.09375,42.140625,0.008942,4.004745,0.010809,0.010715,124.464623,0.001153,2e-06,0.003358,0.00099,0.005278,b93702f8e244a497f0558f231ed9932e11d638c059a68fae04c65d4cdc12ede1
|
||||
buffered,raw,1,1,35.734375,42.390625,5.403822,2.666848,0.010026,0.009779,136.371429,0.001274,9e-06,0.000649,0.003622,0.004201,b93702f8e244a497f0558f231ed9932e11d638c059a68fae04c65d4cdc12ede1
|
||||
current,raw,1,1,35.90625,37.453125,2.775385,0.000532,0.008849,0.008395,158.859134,1.3e-05,1.3e-05,4e-06,0.00396,0.004419,b93702f8e244a497f0558f231ed9932e11d638c059a68fae04c65d4cdc12ede1
|
||||
refactor,raw,1,1,36.09375,37.515625,4.081371,0.000418,0.009241,0.008737,152.631848,1.1e-05,1.5e-05,0.000365,0.003993,0.004366,b93702f8e244a497f0558f231ed9932e11d638c059a68fae04c65d4cdc12ede1
|
||||
python,raw,1,16,51.390625,99.90625,0.144142,44.036157,0.144836,0.082271,259.353123,0.01641,1.3e-05,0.045342,0.01115,0.009353,b93702f8e244a497f0558f231ed9932e11d638c059a68fae04c65d4cdc12ede1
|
||||
buffered,raw,1,16,51.109375,100.6875,86.465351,2.667868,0.147116,0.090379,236.086682,0.016809,4.1e-05,0.010212,0.057063,0.006142,b93702f8e244a497f0558f231ed9932e11d638c059a68fae04c65d4cdc12ede1
|
||||
current,raw,1,16,50.875,60.421875,43.986605,0.000654,0.132194,0.069414,307.391271,1.8e-05,4.6e-05,2.5e-05,0.062784,0.006526,b93702f8e244a497f0558f231ed9932e11d638c059a68fae04c65d4cdc12ede1
|
||||
refactor,raw,1,16,51.453125,61.03125,65.353718,0.000624,0.137993,0.075325,283.26719,1.8e-05,4.4e-05,0.005736,0.062697,0.006735,b93702f8e244a497f0558f231ed9932e11d638c059a68fae04c65d4cdc12ede1
|
||||
python,raw,16,1,51.359375,115.84375,0.008944,64.004745,0.143382,0.140933,151.373531,0.017633,3e-06,0.051982,0.0106,0.060808,ac6672aaeabdb47defc79154d4677581266dcc943692458b2ab7b91dbe505e0c
|
||||
buffered,raw,16,1,50.625,94.625,85.403717,42.666848,0.14594,0.142832,149.360885,0.018153,1.2e-05,0.0084,0.056755,0.059511,ac6672aaeabdb47defc79154d4677581266dcc943692458b2ab7b91dbe505e0c
|
||||
current,raw,16,1,50.6875,52.171875,42.775417,0.000505,0.129723,0.122208,174.567696,1.3e-05,1.3e-05,5e-06,0.062537,0.059637,ac6672aaeabdb47defc79154d4677581266dcc943692458b2ab7b91dbe505e0c
|
||||
refactor,raw,16,1,51.125,52.484375,64.081373,0.000448,0.135518,0.127975,166.701123,1.2e-05,1.3e-05,0.005697,0.062554,0.059747,ac6672aaeabdb47defc79154d4677581266dcc943692458b2ab7b91dbe505e0c
|
||||
python,raw,16,16,291.796875,999.703125,0.144173,704.036157,2.077596,1.226351,278.335767,0.269383,2.2e-05,0.72474,0.164234,0.067562,ac6672aaeabdb47defc79154d4677581266dcc943692458b2ab7b91dbe505e0c
|
||||
buffered,raw,16,16,291.03125,999.859375,1366.465473,42.667868,2.265842,1.422768,239.910695,0.275638,5e-05,0.136665,0.943933,0.063732,ac6672aaeabdb47defc79154d4677581266dcc943692458b2ab7b91dbe505e0c
|
||||
current,raw,16,16,291.140625,308.703125,684.002291,0.000685,1.999672,1.075116,317.488772,2.8e-05,4.9e-05,2.5e-05,1.005901,0.069095,ac6672aaeabdb47defc79154d4677581266dcc943692458b2ab7b91dbe505e0c
|
||||
refactor,raw,16,16,291.5625,307.46875,1025.307026,0.000685,2.094675,1.166656,292.577455,3e-05,4.7e-05,0.092605,1.004331,0.069312,ac6672aaeabdb47defc79154d4677581266dcc943692458b2ab7b91dbe505e0c
|
||||
python,raw,64,1,99.546875,356.046875,0.008944,256.004745,0.570477,0.562132,151.80346,0.071589,5e-06,0.210723,0.041346,0.239411,d20093ebb9fec29c958e0c49b2ed0741d6a11d526c028ff3ebc6a98db2c25041
|
||||
buffered,raw,64,1,98.734375,270.703125,341.403717,170.666848,0.587426,0.575344,148.317433,0.073881,1e-05,0.033467,0.228159,0.239493,d20093ebb9fec29c958e0c49b2ed0741d6a11d526c028ff3ebc6a98db2c25041
|
||||
current,raw,64,1,98.71875,100.75,170.775417,0.000624,0.522024,0.489979,174.157649,1.4e-05,1.4e-05,5e-06,0.250626,0.239619,d20093ebb9fec29c958e0c49b2ed0741d6a11d526c028ff3ebc6a98db2c25041
|
||||
refactor,raw,64,1,99.125,101.15625,256.081373,0.000479,0.542276,0.511243,166.913858,1.4e-05,1.3e-05,0.022774,0.250823,0.237421,d20093ebb9fec29c958e0c49b2ed0741d6a11d526c028ff3ebc6a98db2c25041
|
||||
python,raw,64,16,1059.5625,3879.453125,0.144173,2816.036157,9.39656,5.096071,267.919599,1.086042,2.3e-05,3.029026,0.670795,0.303795,d20093ebb9fec29c958e0c49b2ed0741d6a11d526c028ff3ebc6a98db2c25041
|
||||
buffered,raw,64,16,1059.0,3879.671875,5462.465473,170.667868,9.639721,5.738253,237.936046,1.120327,5.6e-05,0.555259,3.785224,0.264696,d20093ebb9fec29c958e0c49b2ed0741d6a11d526c028ff3ebc6a98db2c25041
|
||||
current,raw,64,16,1059.171875,1077.09375,2731.986666,0.000689,8.961106,4.367673,312.600592,3.2e-05,6.5e-05,2.7e-05,4.064023,0.315124,d20093ebb9fec29c958e0c49b2ed0741d6a11d526c028ff3ebc6a98db2c25041
|
||||
refactor,raw,64,16,1059.515625,1080.984375,4097.291401,0.000631,9.460124,4.892549,279.064576,3.1e-05,4.9e-05,0.383175,4.199578,0.301434,d20093ebb9fec29c958e0c49b2ed0741d6a11d526c028ff3ebc6a98db2c25041
|
||||
|
|
|
@ -1,7 +1,7 @@
|
|||
use std::hint::black_box;
|
||||
use std::time::Duration;
|
||||
|
||||
use criterion::{BenchmarkId, Criterion, criterion_group, criterion_main};
|
||||
use criterion::{BenchmarkId, Criterion, criterion_group};
|
||||
use litellm_python_interop::{from_py, to_py};
|
||||
use pyo3::prelude::*;
|
||||
use pyo3::types::PyDict;
|
||||
|
|
@ -99,4 +99,13 @@ criterion_group! {
|
|||
.measurement_time(Duration::from_secs(4));
|
||||
targets = bridge_serialization
|
||||
}
|
||||
criterion_main!(benches);
|
||||
mod media;
|
||||
|
||||
fn main() {
|
||||
if std::env::args().nth(1).as_deref() == Some("--media") {
|
||||
media::run();
|
||||
} else {
|
||||
benches();
|
||||
Criterion::default().configure_from_args().final_summary();
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -185,3 +185,6 @@ asyncio.run(asyncio.wait_for(exercise(), timeout=5))
|
|||
.expect("server task should not panic");
|
||||
}
|
||||
}
|
||||
|
||||
mod messages;
|
||||
mod payload;
|
||||
|
|
|
|||
236
litellm-rust/crates/python-bridge/src/messages.rs
Normal file
236
litellm-rust/crates/python-bridge/src/messages.rs
Normal file
|
|
@ -0,0 +1,236 @@
|
|||
use crate::payload::payload_from_py;
|
||||
use litellm_core::http_utils::body::{JsonPayload, SharedText};
|
||||
use litellm_core::messages::types::{
|
||||
AnthropicMessage, AnthropicMessagesRequest, CacheControl, ContentBlock, MessageContent,
|
||||
SystemPrompt,
|
||||
};
|
||||
use litellm_python_interop::{from_py, text_bytes_from_py};
|
||||
use pyo3::exceptions::PyValueError;
|
||||
use pyo3::prelude::*;
|
||||
use pyo3::types::{PyDict, PyList, PyString, PyTuple};
|
||||
use serde::de::DeserializeOwned;
|
||||
use std::collections::BTreeMap;
|
||||
|
||||
fn invalid(message: impl std::fmt::Display) -> PyErr {
|
||||
PyValueError::new_err(format!("invalid Anthropic messages request: {message}"))
|
||||
}
|
||||
|
||||
fn required<'py>(dict: &Bound<'py, PyDict>, name: &str) -> PyResult<Bound<'py, PyAny>> {
|
||||
dict.get_item(name)?
|
||||
.ok_or_else(|| invalid(format!("missing field `{name}`")))
|
||||
}
|
||||
|
||||
fn optional<T: DeserializeOwned>(dict: &Bound<'_, PyDict>, name: &str) -> PyResult<Option<T>> {
|
||||
dict.get_item(name)?
|
||||
.filter(|value| !value.is_none())
|
||||
.map(|value| from_py(&value))
|
||||
.transpose()
|
||||
}
|
||||
|
||||
fn finite_number(dict: &Bound<'_, PyDict>, name: &str) -> PyResult<Option<f64>> {
|
||||
optional::<f64>(dict, name).map(|number| number.filter(|value| value.is_finite()))
|
||||
}
|
||||
|
||||
fn shared(dict: &Bound<'_, PyDict>, name: &str) -> PyResult<Option<JsonPayload>> {
|
||||
dict.get_item(name)?
|
||||
.filter(|value| !value.is_none())
|
||||
.map(|value| payload_from_py(&value))
|
||||
.transpose()
|
||||
}
|
||||
|
||||
fn shared_array(dict: &Bound<'_, PyDict>, name: &str) -> PyResult<Option<Vec<JsonPayload>>> {
|
||||
dict.get_item(name)?
|
||||
.filter(|value| !value.is_none())
|
||||
.map(|value| sequence(&value, payload_from_py))
|
||||
.transpose()
|
||||
}
|
||||
|
||||
fn extras(dict: &Bound<'_, PyDict>, known: &[&str]) -> PyResult<BTreeMap<String, JsonPayload>> {
|
||||
dict.iter()
|
||||
.filter_map(|(key, value)| match key.extract::<String>() {
|
||||
Ok(key) if known.contains(&key.as_str()) => None,
|
||||
Ok(key) => Some(payload_from_py(&value).map(|value| (key, value))),
|
||||
Err(error) => Some(Err(error)),
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
|
||||
fn sequence<T>(
|
||||
value: &Bound<'_, PyAny>,
|
||||
extract: impl Fn(&Bound<'_, PyAny>) -> PyResult<T>,
|
||||
) -> PyResult<Vec<T>> {
|
||||
if !value.is_instance_of::<PyList>() && !value.is_instance_of::<PyTuple>() {
|
||||
return Err(invalid("expected an array"));
|
||||
}
|
||||
value.try_iter()?.map(|value| extract(&value?)).collect()
|
||||
}
|
||||
|
||||
fn content(value: &Bound<'_, PyAny>) -> PyResult<MessageContent> {
|
||||
if value.is_instance_of::<PyString>() {
|
||||
return SharedText::new(text_bytes_from_py(value)?)
|
||||
.map(MessageContent::Text)
|
||||
.map_err(invalid);
|
||||
}
|
||||
sequence(value, block).map(MessageContent::Blocks)
|
||||
}
|
||||
|
||||
fn block(value: &Bound<'_, PyAny>) -> PyResult<ContentBlock> {
|
||||
let dict = value
|
||||
.cast::<PyDict>()
|
||||
.map_err(|_| invalid("content block must be an object"))?;
|
||||
let cache_control = dict
|
||||
.get_item("cache_control")?
|
||||
.filter(|value| !value.is_none())
|
||||
.map(|value| {
|
||||
let cache = value
|
||||
.cast::<PyDict>()
|
||||
.map_err(|_| invalid("cache_control must be an object"))?;
|
||||
Ok::<_, PyErr>(CacheControl {
|
||||
cache_type: optional(cache, "type")?,
|
||||
ttl: optional(cache, "ttl")?,
|
||||
scope: optional(cache, "scope")?,
|
||||
extra: extras(cache, &["type", "ttl", "scope"])?,
|
||||
})
|
||||
})
|
||||
.transpose()?;
|
||||
Ok(ContentBlock {
|
||||
cache_control,
|
||||
extra: extras(dict, &["cache_control"])?,
|
||||
})
|
||||
}
|
||||
|
||||
fn message(value: &Bound<'_, PyAny>) -> PyResult<AnthropicMessage> {
|
||||
let dict = value
|
||||
.cast::<PyDict>()
|
||||
.map_err(|_| invalid("message must be an object"))?;
|
||||
Ok(AnthropicMessage {
|
||||
role: from_py(&required(dict, "role")?)?,
|
||||
content: content(&required(dict, "content")?)?,
|
||||
extra: extras(dict, &["role", "content"])?,
|
||||
})
|
||||
}
|
||||
|
||||
pub(crate) fn messages_from_py(value: &Bound<'_, PyAny>) -> PyResult<AnthropicMessagesRequest> {
|
||||
let dict = value
|
||||
.cast::<PyDict>()
|
||||
.map_err(|_| PyValueError::new_err("body must be a dict"))?;
|
||||
let system = dict
|
||||
.get_item("system")?
|
||||
.filter(|value| !value.is_none())
|
||||
.map(|value| {
|
||||
content(&value).map(|value| match value {
|
||||
MessageContent::Text(text) => SystemPrompt::Text(text),
|
||||
MessageContent::Blocks(blocks) => SystemPrompt::Blocks(blocks),
|
||||
})
|
||||
})
|
||||
.transpose()?;
|
||||
Ok(AnthropicMessagesRequest {
|
||||
model: from_py(&required(dict, "model")?)?,
|
||||
messages: sequence(&required(dict, "messages")?, message)?,
|
||||
system,
|
||||
max_tokens: optional(dict, "max_tokens")?,
|
||||
stop_sequences: optional(dict, "stop_sequences")?,
|
||||
stream: optional(dict, "stream")?,
|
||||
temperature: finite_number(dict, "temperature")?,
|
||||
top_p: finite_number(dict, "top_p")?,
|
||||
top_k: optional(dict, "top_k")?,
|
||||
service_tier: optional(dict, "service_tier")?,
|
||||
speed: optional(dict, "speed")?,
|
||||
inference_geo: optional(dict, "inference_geo")?,
|
||||
metadata: shared(dict, "metadata")?,
|
||||
tool_choice: shared(dict, "tool_choice")?,
|
||||
thinking: shared(dict, "thinking")?,
|
||||
container: shared(dict, "container")?,
|
||||
context_management: shared(dict, "context_management")?,
|
||||
output_format: shared(dict, "output_format")?,
|
||||
output_config: shared(dict, "output_config")?,
|
||||
tools: shared_array(dict, "tools")?,
|
||||
mcp_servers: shared_array(dict, "mcp_servers")?,
|
||||
extra: extras(
|
||||
dict,
|
||||
&[
|
||||
"model",
|
||||
"messages",
|
||||
"system",
|
||||
"max_tokens",
|
||||
"stop_sequences",
|
||||
"stream",
|
||||
"temperature",
|
||||
"top_p",
|
||||
"top_k",
|
||||
"service_tier",
|
||||
"speed",
|
||||
"inference_geo",
|
||||
"metadata",
|
||||
"tool_choice",
|
||||
"thinking",
|
||||
"container",
|
||||
"context_management",
|
||||
"output_format",
|
||||
"output_config",
|
||||
"tools",
|
||||
"mcp_servers",
|
||||
],
|
||||
)?,
|
||||
})
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use litellm_core::messages::transformation::AnthropicMessagesProviderConfig;
|
||||
use litellm_core::providers::azure_ai::messages::transformation::AZURE_ANTHROPIC_MESSAGES_CONFIG;
|
||||
|
||||
#[test]
|
||||
fn typed_extraction_retains_nested_media_and_extension_owners() {
|
||||
Python::initialize();
|
||||
let (request, pointer) = Python::attach(|py| {
|
||||
let globals = PyDict::new(py);
|
||||
py.run(c"media = 'A' * (1024 * 1024)\nbody = {'model': 'model', 'messages': [{'role': 'system', 'content': [{'type': 'tool_result', 'content': [{'type': 'image', 'source': {'type': 'base64', 'data': media}}]}]}], 'extension': {'data': media}}", Some(&globals), None).unwrap();
|
||||
let media = globals.get_item("media").unwrap().unwrap();
|
||||
let pointer = media.cast::<PyString>().unwrap().to_str().unwrap().as_ptr() as usize;
|
||||
let body = globals.get_item("body").unwrap().unwrap();
|
||||
(messages_from_py(&body).unwrap(), pointer)
|
||||
});
|
||||
let request = std::thread::spawn(move || {
|
||||
AZURE_ANTHROPIC_MESSAGES_CONFIG
|
||||
.transform_request(request)
|
||||
.unwrap()
|
||||
})
|
||||
.join()
|
||||
.unwrap();
|
||||
let SystemPrompt::Blocks(blocks) = request.system.unwrap() else {
|
||||
panic!("blocks")
|
||||
};
|
||||
assert_eq!(
|
||||
blocks[0].extra["content"][0]["source"]["data"]
|
||||
.as_text()
|
||||
.unwrap()
|
||||
.bytes()
|
||||
.as_ptr() as usize,
|
||||
pointer
|
||||
);
|
||||
assert_eq!(
|
||||
request.extra["extension"]["data"]
|
||||
.as_text()
|
||||
.unwrap()
|
||||
.bytes()
|
||||
.as_ptr() as usize,
|
||||
pointer
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn direct_extraction_matches_serde_for_known_fields_and_extensions() {
|
||||
Python::initialize();
|
||||
Python::attach(|py| {
|
||||
let globals = PyDict::new(py);
|
||||
py.run(c"body = {'model': 'model', 'messages': [{'role': 'user', 'content': [{'type': 'text', 'text': 'héllo', 'cache_control': {'type': 'ephemeral', 'scope': 'global', 'unknown': 4}}], 'extension': [None, True]}], 'system': 'system', 'max_tokens': 3, 'tools': [{'name': 'tool'}], 'metadata': {'a': 1}, 'thinking': None, 'temperature': float('nan'), 'top_p': float('inf'), 'unknown': {'data': 'AAAA'}}", Some(&globals), None).unwrap();
|
||||
let value = globals.get_item("body").unwrap().unwrap();
|
||||
let direct = messages_from_py(&value).unwrap();
|
||||
let reference: AnthropicMessagesRequest =
|
||||
serde_json::from_value(from_py::<serde_json::Value>(&value).unwrap()).unwrap();
|
||||
assert_eq!(direct, reference);
|
||||
});
|
||||
}
|
||||
}
|
||||
51
litellm-rust/crates/python-bridge/src/payload.rs
Normal file
51
litellm-rust/crates/python-bridge/src/payload.rs
Normal file
|
|
@ -0,0 +1,51 @@
|
|||
use litellm_core::http_utils::body::{JsonPayload, SharedText};
|
||||
use litellm_python_interop::{bytes_from_py, from_py, text_bytes_from_py};
|
||||
use pyo3::exceptions::PyValueError;
|
||||
use pyo3::prelude::*;
|
||||
use pyo3::types::{PyByteArray, PyBytes, PyDict, PyList, PyString, PyTuple};
|
||||
|
||||
pub(crate) fn payload_from_py(value: &Bound<'_, PyAny>) -> PyResult<JsonPayload> {
|
||||
extract(value, 0, false)
|
||||
}
|
||||
|
||||
pub(crate) fn audio_payload_from_py(value: &Bound<'_, PyAny>) -> PyResult<JsonPayload> {
|
||||
extract(value, 0, true)
|
||||
}
|
||||
|
||||
fn extract(value: &Bound<'_, PyAny>, depth: usize, audio: bool) -> PyResult<JsonPayload> {
|
||||
if depth > litellm_core::constants::JSON_PAYLOAD_MAX_DEPTH {
|
||||
return Err(PyValueError::new_err(
|
||||
"request nesting exceeds the JSON depth limit",
|
||||
));
|
||||
}
|
||||
if value.is_instance_of::<PyString>() {
|
||||
return SharedText::new(text_bytes_from_py(value)?)
|
||||
.map(JsonPayload::String)
|
||||
.map_err(|_| PyValueError::new_err("invalid UTF-8 string"));
|
||||
}
|
||||
if audio
|
||||
&& depth == 1
|
||||
&& (value.is_instance_of::<PyBytes>() || value.is_instance_of::<PyByteArray>())
|
||||
{
|
||||
return Ok(JsonPayload::Base64(bytes_from_py(value)?));
|
||||
}
|
||||
if let Ok(dict) = value.cast::<PyDict>() {
|
||||
return dict
|
||||
.iter()
|
||||
.map(|(key, value)| {
|
||||
let key = key.extract::<String>()?;
|
||||
let value = extract(&value, depth + 1, audio && depth == 0 && key == "data")?;
|
||||
Ok((key, value))
|
||||
})
|
||||
.collect::<PyResult<_>>()
|
||||
.map(JsonPayload::Object);
|
||||
}
|
||||
if value.is_instance_of::<PyList>() || value.is_instance_of::<PyTuple>() {
|
||||
return value
|
||||
.try_iter()?
|
||||
.map(|item| extract(&item?, depth + 1, false))
|
||||
.collect::<PyResult<_>>()
|
||||
.map(JsonPayload::Array);
|
||||
}
|
||||
from_py::<serde_json::Value>(value).map(JsonPayload::from)
|
||||
}
|
||||
|
|
@ -1,4 +1,5 @@
|
|||
use litellm_core::Error;
|
||||
use litellm_core::http_utils::body::JsonPayload;
|
||||
use std::future::Future;
|
||||
|
||||
use litellm_core::audio_transcription::{
|
||||
|
|
@ -53,8 +54,8 @@ bridge_route! {
|
|||
inputs = AudioTranscriptionInputs,
|
||||
required = {
|
||||
model: String,
|
||||
#[pyo3(from_py_with = litellm_python_interop::from_py)]
|
||||
audio: Value,
|
||||
#[pyo3(from_py_with = crate::payload::audio_payload_from_py)]
|
||||
audio: JsonPayload,
|
||||
},
|
||||
optional = {
|
||||
api_key: Option<String>,
|
||||
|
|
|
|||
|
|
@ -1,17 +1,19 @@
|
|||
use litellm_core::Error;
|
||||
use litellm_core::messages::messages as run_messages;
|
||||
use litellm_core::messages::types::{AnthropicMessagesResponse, MessagesRequest};
|
||||
use litellm_core::messages::types::{
|
||||
AnthropicMessagesRequest, AnthropicMessagesResponse, MessagesRequest,
|
||||
};
|
||||
use pyo3::prelude::*;
|
||||
use serde_json::Value;
|
||||
use std::future::Future;
|
||||
|
||||
use crate::errors::core_error_to_pyerr;
|
||||
use crate::marshal::{RouteOptions, RouteOptionsInputs, required_value};
|
||||
use crate::marshal::{RouteOptions, RouteOptionsInputs};
|
||||
|
||||
fn prepare_messages(
|
||||
inputs: MessagesInputs,
|
||||
) -> PyResult<impl Future<Output = Result<AnthropicMessagesResponse, Error>> + Send + 'static> {
|
||||
let body = required_value("body", inputs.body, Value::is_object, "dict")?;
|
||||
let body = inputs.body;
|
||||
let options = RouteOptions::from_python(RouteOptionsInputs {
|
||||
model: inputs.model,
|
||||
api_key: inputs.api_key,
|
||||
|
|
@ -49,8 +51,8 @@ bridge_route! {
|
|||
inputs = MessagesInputs,
|
||||
required = {
|
||||
model: String,
|
||||
#[pyo3(from_py_with = litellm_python_interop::from_py)]
|
||||
body: Value,
|
||||
#[pyo3(from_py_with = crate::messages::messages_from_py)]
|
||||
body: AnthropicMessagesRequest,
|
||||
},
|
||||
optional = {
|
||||
api_key: Option<String>,
|
||||
|
|
|
|||
|
|
@ -1,4 +1,5 @@
|
|||
use litellm_core::Error;
|
||||
use litellm_core::http_utils::body::JsonPayload;
|
||||
use std::future::Future;
|
||||
|
||||
use litellm_ai_gateway::io::ocr::{OcrRequest, ocr as run_ocr};
|
||||
|
|
@ -55,8 +56,8 @@ bridge_route! {
|
|||
inputs = OcrInputs,
|
||||
required = {
|
||||
model: String,
|
||||
#[pyo3(from_py_with = litellm_python_interop::from_py)]
|
||||
document: Value,
|
||||
#[pyo3(from_py_with = crate::payload::payload_from_py)]
|
||||
document: JsonPayload,
|
||||
},
|
||||
optional = {
|
||||
api_key: Option<String>,
|
||||
|
|
|
|||
|
|
@ -6,6 +6,7 @@ license.workspace = true
|
|||
repository.workspace = true
|
||||
|
||||
[dependencies]
|
||||
bytes.workspace = true
|
||||
pyo3.workspace = true
|
||||
pythonize.workspace = true
|
||||
serde.workspace = true
|
||||
|
|
|
|||
11
litellm-rust/crates/python-interop/src/bytes.rs
Normal file
11
litellm-rust/crates/python-interop/src/bytes.rs
Normal file
|
|
@ -0,0 +1,11 @@
|
|||
use bytes::Bytes;
|
||||
use pyo3::prelude::*;
|
||||
use pyo3::pybacked::PyBackedStr;
|
||||
|
||||
pub fn bytes_from_py(value: &Bound<'_, PyAny>) -> PyResult<Bytes> {
|
||||
Ok(value.extract()?)
|
||||
}
|
||||
|
||||
pub fn text_bytes_from_py(value: &Bound<'_, PyAny>) -> PyResult<Bytes> {
|
||||
Ok(Bytes::from_owner(value.extract::<PyBackedStr>()?))
|
||||
}
|
||||
|
|
@ -3,3 +3,6 @@ mod marshal;
|
|||
|
||||
pub use gil::{release_count, release_gil};
|
||||
pub use marshal::{Pythonized, from_py, panic_to_pyerr, to_py};
|
||||
|
||||
mod bytes;
|
||||
pub use bytes::{bytes_from_py, text_bytes_from_py};
|
||||
|
|
|
|||
|
|
@ -42,3 +42,80 @@ fn release_gil_runs_work_and_records_it(#[from(initialized_python)] python: &Ini
|
|||
assert_eq!(result, 42);
|
||||
assert_eq!(release_count(), before + 1);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
fn immutable_payloads_share_storage_and_mutable_payloads_take_snapshots(
|
||||
#[from(initialized_python)] python: &InitializedPython,
|
||||
) {
|
||||
use litellm_python_interop::{bytes_from_py, text_bytes_from_py};
|
||||
use pyo3::prelude::*;
|
||||
use pyo3::types::{PyByteArray, PyBytes, PyBytesMethods, PyString, PyStringMethods};
|
||||
|
||||
let (raw, text) = python.attach(|py| {
|
||||
let source = PyBytes::new(py, b"large immutable media");
|
||||
let references = refcount(&source);
|
||||
let raw = bytes_from_py(source.as_any()).unwrap();
|
||||
assert_eq!(raw.as_ptr(), source.as_bytes().as_ptr());
|
||||
assert_eq!(refcount(&source), references + 1);
|
||||
let replay = raw.clone().slice(6..15);
|
||||
assert_eq!(replay.as_ptr(), source.as_bytes()[6..].as_ptr());
|
||||
drop(replay);
|
||||
assert_eq!(refcount(&source), references + 1);
|
||||
let text = PyString::new(py, "QUJDREVGR0g=");
|
||||
let shared_text = text_bytes_from_py(text.as_any()).unwrap();
|
||||
assert_eq!(shared_text.as_ptr(), text.to_str().unwrap().as_ptr());
|
||||
let mutable = PyByteArray::new(py, b"original");
|
||||
let snapshot = bytes_from_py(mutable.as_any()).unwrap();
|
||||
mutable.set_item(0, b'X').unwrap();
|
||||
assert_eq!(snapshot.as_ref(), b"original");
|
||||
(raw, shared_text)
|
||||
});
|
||||
std::thread::spawn(move || {
|
||||
assert_eq!(raw.as_ref(), b"large immutable media");
|
||||
assert_eq!(text.slice(4..).as_ref(), b"REVGR0g=");
|
||||
})
|
||||
.join()
|
||||
.unwrap();
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
fn owners_are_released_on_success_error_and_cancelled_suspension(
|
||||
#[from(initialized_python)] python: &InitializedPython,
|
||||
) {
|
||||
use litellm_python_interop::bytes_from_py;
|
||||
use pyo3::types::PyBytes;
|
||||
use std::future::Future;
|
||||
use std::task::{Context, Poll, Waker};
|
||||
|
||||
let source = python.attach(|py| PyBytes::new(py, b"owner lifetime").unbind());
|
||||
let baseline = python.attach(|py| refcount(source.bind(py)));
|
||||
for fail in [false, true] {
|
||||
let bytes = python.attach(|py| bytes_from_py(source.bind(py).as_any()).unwrap());
|
||||
let result = std::thread::spawn(move || {
|
||||
assert_eq!(bytes.as_ref(), b"owner lifetime");
|
||||
if fail { Err(()) } else { Ok(()) }
|
||||
})
|
||||
.join()
|
||||
.unwrap();
|
||||
assert_eq!(result.is_err(), fail);
|
||||
python.attach(|py| assert_eq!(refcount(source.bind(py)), baseline));
|
||||
}
|
||||
let bytes = python.attach(|py| bytes_from_py(source.bind(py).as_any()).unwrap());
|
||||
let mut future = Box::pin(async move {
|
||||
std::future::pending::<()>().await;
|
||||
assert_eq!(bytes.as_ref(), b"owner lifetime");
|
||||
});
|
||||
assert_eq!(
|
||||
future
|
||||
.as_mut()
|
||||
.poll(&mut Context::from_waker(Waker::noop())),
|
||||
Poll::Pending
|
||||
);
|
||||
python.attach(|py| assert_eq!(refcount(source.bind(py)), baseline + 1));
|
||||
drop(future);
|
||||
python.attach(|py| assert_eq!(refcount(source.bind(py)), baseline));
|
||||
}
|
||||
|
||||
fn refcount<T>(value: &pyo3::Bound<'_, T>) -> isize {
|
||||
unsafe { pyo3::ffi::Py_REFCNT(value.as_ptr()) }
|
||||
}
|
||||
|
|
|
|||
|
|
@ -1,4 +1,3 @@
|
|||
import base64
|
||||
from typing import Final
|
||||
|
||||
import httpx
|
||||
|
|
@ -26,7 +25,7 @@ class BedrockAudioTranscriptionRustDispatch:
|
|||
if audio_format not in {"wav", "mp3", "flac", "ogg"}:
|
||||
raise ValueError(f"Unsupported Bedrock audio format for file {processed_audio.filename!r}")
|
||||
return {
|
||||
"data": base64.b64encode(processed_audio.file_content).decode("ascii"),
|
||||
"data": processed_audio.file_content,
|
||||
"format": audio_format,
|
||||
"filename": processed_audio.filename,
|
||||
}
|
||||
|
|
|
|||
|
|
@ -16,6 +16,9 @@ from pathlib import Path
|
|||
from socket import socket as Socket
|
||||
from typing import Final
|
||||
|
||||
MEDIA: Final = "AQID" * (256 * 1024)
|
||||
DOCUMENT: Final = "data:application/pdf;base64," + MEDIA
|
||||
|
||||
REQUEST_STARTED: Final = threading.Event()
|
||||
REQUEST_CANCELLED: Final = threading.Event()
|
||||
|
||||
|
|
@ -83,7 +86,7 @@ def assert_native_request(
|
|||
assert path == "/v1/ocr"
|
||||
assert headers.get("authorization") == "Bearer sk-native"
|
||||
assert body["model"] == "mistral-ocr-latest"
|
||||
assert body["document"]["document_url"] == "https://example.com/document.pdf"
|
||||
assert body["document"]["document_url"] == DOCUMENT
|
||||
assert body["include_image_base64"] is True
|
||||
return
|
||||
if route == "transcription":
|
||||
|
|
@ -98,7 +101,7 @@ def assert_native_request(
|
|||
assert body["model"] == "claude-sonnet-4-5"
|
||||
if route == "messages":
|
||||
assert body["max_tokens"] == 16
|
||||
assert body["messages"][0]["content"] == "hello-from-messages"
|
||||
assert body["messages"][0]["content"][0]["content"][0]["source"]["data"] == MEDIA
|
||||
return
|
||||
assert body["max_tokens"] == 17
|
||||
assert body["messages"][0]["content"] == [{"type": "text", "text": "hello-from-chat"}]
|
||||
|
|
@ -132,7 +135,7 @@ def route_kwargs(route: str, api_base: str, outcome: str) -> dict[str, object]:
|
|||
if route == "ocr":
|
||||
return common | {
|
||||
"model": "mistral-ocr-latest",
|
||||
"document": {"type": "document_url", "document_url": "https://example.com/document.pdf"},
|
||||
"document": {"type": "document_url", "document_url": DOCUMENT},
|
||||
"api_key": "sk-native",
|
||||
"custom_llm_provider": "mistral",
|
||||
"optional_params": {"include_image_base64": True},
|
||||
|
|
@ -140,7 +143,7 @@ def route_kwargs(route: str, api_base: str, outcome: str) -> dict[str, object]:
|
|||
if route == "transcription":
|
||||
return common | {
|
||||
"model": "mistral.voxtral-mini-3b-2507",
|
||||
"audio": {"data": "AQI=", "format": "wav", "filename": "audio.wav"},
|
||||
"audio": {"data": b"\x01\x02", "format": "wav", "filename": "audio.wav"},
|
||||
"custom_llm_provider": "bedrock",
|
||||
"optional_params": {
|
||||
"aws_access_key_id": "native-access-key",
|
||||
|
|
@ -155,7 +158,23 @@ def route_kwargs(route: str, api_base: str, outcome: str) -> dict[str, object]:
|
|||
"body": {
|
||||
"model": "claude-sonnet-4-5",
|
||||
"max_tokens": 16,
|
||||
"messages": [{"role": "user", "content": "hello-from-messages"}],
|
||||
"messages": [
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{
|
||||
"type": "tool_result",
|
||||
"tool_use_id": "tool_1",
|
||||
"content": [
|
||||
{
|
||||
"type": "image",
|
||||
"source": {"type": "base64", "media_type": "image/png", "data": MEDIA},
|
||||
}
|
||||
],
|
||||
}
|
||||
],
|
||||
}
|
||||
],
|
||||
},
|
||||
"api_key": "sk-native",
|
||||
"custom_llm_provider": "anthropic",
|
||||
|
|
@ -237,12 +256,7 @@ async def exercise_async(native: object, api_base: str) -> None:
|
|||
|
||||
async def exercise_async_concurrency(native: object, api_base: str) -> None:
|
||||
responses: Final = await asyncio.wait_for(
|
||||
asyncio.gather(
|
||||
*(
|
||||
native.amessages(**route_kwargs("messages", api_base, "success"))
|
||||
for _ in range(32)
|
||||
)
|
||||
),
|
||||
asyncio.gather(*(native.amessages(**route_kwargs("messages", api_base, "success")) for _ in range(32))),
|
||||
timeout=15,
|
||||
)
|
||||
for response in responses:
|
||||
|
|
@ -252,6 +266,17 @@ async def exercise_async_concurrency(native: object, api_base: str) -> None:
|
|||
def exercise_routes(native_path: Path, api_base: str) -> object:
|
||||
native: Final = load_native(native_path)
|
||||
exercise_sync(native, api_base)
|
||||
assert_success(
|
||||
"transcription",
|
||||
native.transcription(
|
||||
**(
|
||||
route_kwargs("transcription", api_base, "success")
|
||||
| {
|
||||
"audio": {"data": "AQI=", "format": "wav", "filename": "audio.wav"},
|
||||
}
|
||||
)
|
||||
),
|
||||
)
|
||||
asyncio.run(exercise_async(native, api_base))
|
||||
asyncio.run(exercise_async_concurrency(native, api_base))
|
||||
return native
|
||||
|
|
|
|||
|
|
@ -149,3 +149,13 @@ async def test_bedrock_atranscription_uses_rust_only_path() -> None:
|
|||
rust_bridge.configure_rust_transcription(transcription=None, atranscription=None)
|
||||
|
||||
assert response.text == "rust"
|
||||
|
||||
|
||||
def test_dispatch_shares_raw_audio_with_native_bridge() -> None:
|
||||
from typing import Final
|
||||
|
||||
raw: Final = b"audio" * 1024
|
||||
payload: Final = BedrockAudioTranscriptionRustDispatch._audio_payload(("audio.wav", raw, "audio/wav"))
|
||||
assert payload["data"] is raw
|
||||
assert payload["format"] == "wav"
|
||||
assert payload["filename"] == "audio.wav"
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue