Merge branch 'litellm_internal_staging' into litellm_model_armor_logging_only_streaming
Some checks are pending
ai-gateway image / ai-gateway release image (push) Waiting to run
LiteLLM Rust / rust-lint (push) Waiting to run
LiteLLM Rust / rust-test (push) Waiting to run

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
yucheng 2026-09-11 23:48:35 +00:00
commit 7a083353a8
255 changed files with 24259 additions and 7564 deletions

View file

@ -106,11 +106,6 @@ dockerfiles:
and lint workflows already exercise that output, so building the image adds no signal about it
paths:
- ui/Dockerfile
- reason: >-
The Rust gateway ships as its own chart and package with a separate release pipeline, so its
image is not part of this repo's Python image set
paths:
- litellm-rust/crates/ai-gateway/Dockerfile
- reason: >-
An example image under cookbook/ that is documentation rather than a shipped artifact
paths:

View file

@ -205,7 +205,7 @@ def main(
native_module: Final = load_native_module(native_path)
native_module_loads: Final = native_module is not None
panic_test_hook_absent: Final = native_module is not None and not hasattr(native_module, "_panic_for_test")
native_size_limit: Final = 20_000_000
native_size_limit: Final = 25_000_000
native_size_within_limit: Final = native_member.file_size <= native_size_limit
validations: Final = (
(f"Python tag is {EXPECTED_PYTHON_TAG}", python_tag == EXPECTED_PYTHON_TAG),
@ -222,7 +222,7 @@ def main(
("Python extension entry point is present", extension_entry_point_present),
("Native module loads", native_module_loads),
("Production module omits the panic test hook", panic_test_hook_absent),
("Native extension does not exceed 20 MB", native_size_within_limit),
("Native extension does not exceed 25 MB", native_size_within_limit),
("Wheel contents are valid", not unexpected_members),
)

73
.github/workflows/ai-gateway-image.yml vendored Normal file
View file

@ -0,0 +1,73 @@
name: ai-gateway image
on:
push:
paths:
- "litellm-rust/**"
- "litellm/**"
- "enterprise/**"
- "litellm-proxy-extras/**"
- "pyproject.toml"
- "rust-toolchain.toml"
- ".github/workflows/ai-gateway-image.yml"
pull_request:
branches:
- main
- litellm_internal_staging
- litellm_oss_staging
- "litellm_**"
paths:
- "litellm-rust/**"
- "litellm/**"
- "enterprise/**"
- "litellm-proxy-extras/**"
- "pyproject.toml"
- "rust-toolchain.toml"
- ".github/workflows/ai-gateway-image.yml"
permissions:
contents: read
concurrency:
group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.ref }}
cancel-in-progress: true
jobs:
ai-gateway-image:
name: ai-gateway release image
runs-on: ubuntu-latest
timeout-minutes: 60
permissions:
contents: read
steps:
- name: Checkout repository
uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0
with:
persist-credentials: false
- name: Build the release image
run: docker build -f litellm-rust/crates/ai-gateway/Dockerfile -t litellm-ai-gateway:${{ github.sha }} .
- name: Start the gateway and wait for readiness
env:
IMAGE: litellm-ai-gateway:${{ github.sha }}
run: |
docker run -d --name ai-gateway -p 4001:4001 \
-e LITELLM_MASTER_KEY=sk-ci-not-a-real-key \
-e OPENAI_API_KEY=sk-ci-not-a-real-key \
"$IMAGE"
for _ in $(seq 1 60); do
if curl -fsS http://127.0.0.1:4001/health/readiness; then
echo "gateway is serving readiness"
exit 0
fi
sleep 2
done
echo "gateway never became ready" >&2
docker logs ai-gateway >&2
exit 1
- name: Assert the gateway loaded the baked config
run: |
docker logs ai-gateway 2>&1 | tee gateway.log
grep 'via python config reader' gateway.log
- name: Stop the gateway
if: always()
run: docker rm -f ai-gateway || true

View file

@ -262,6 +262,8 @@ curl -X POST 'http://0.0.0.0:4000/v1/chat/completions' \
}
```
For MCP OAuth, an upstream may advertise dynamic client registration but refuse requests with HTTP 401 or 403. If the provider requires a pre-registered OAuth app, configure its `credentials.client_id` and, when required, `credentials.client_secret` on the MCP server. This skips dynamic registration in the gateway sign-in flow. The provider must approve the app for MCP access; reaching its authorization page does not establish that login or tool calls will succeed
[**Docs: MCP Gateway**](https://docs.litellm.ai/docs/mcp)
</details>

View file

@ -0,0 +1 @@
ALTER TABLE "LiteLLM_AutoRouterSession" ADD COLUMN IF NOT EXISTS "baseline_models" JSONB NOT NULL DEFAULT '{}';

View file

@ -1514,6 +1514,7 @@ model LiteLLM_AutoRouterSession {
classifier_cost Float @default(0)
classifier_cost_recorded_turns Int @default(0)
tier_turns Json @default("{}")
baseline_models Json @default("{}")
@@id([api_key, session_id, router_name])
@@index([last_turn_at], map: "idx_autorouter_session_last_turn")

750
litellm-rust/Cargo.lock generated

File diff suppressed because it is too large Load diff

View file

@ -41,8 +41,14 @@ tokio = { version = "1", features = ["rt-multi-thread", "macros", "time", "net"]
tokio-tungstenite = { version = "0.24", default-features = false, features = ["connect", "rustls-tls-native-roots"] }
futures-util = { version = "0.3", default-features = false, features = ["sink", "std"] }
base64 = "0.22"
gcp_auth = "0.12.7"
azure_core = "1.0.0"
azure_identity = { version = "1.0.0", features = ["tokio"] }
moka = { version = "0.12.16", features = ["future"] }
strum = { version = "0.28.0", features = ["derive"] }
url = "2.5.8"
criterion = "0.8.2"
veil = "0.3.0"
[profile.release]
opt-level = 3

View file

@ -14,15 +14,20 @@
# ---- Chef -------------------------------------------------------------------
# cargo-chef caches the dependency build so only the gateway crate recompiles on
# a source-only change. python3-dev is present in every rust stage because the
# `python-config` feature links libpython via pyo3 (even in the cook step).
FROM rust:1.90-slim-bookworm AS chef
# `python-config` feature links libpython via pyo3 (even in the cook step), and
# python3-pip builds the litellm wheel in the builder stage.
FROM rust:1.98-slim-bookworm AS chef
ENV PYO3_PYTHON=python3.11
# rustup reads rust-toolchain.toml from any parent of the working directory, so
# copying it in is what keeps every cargo call below on the repo's pinned
# channel rather than on whatever the base image happens to ship.
COPY rust-toolchain.toml /build/rust-toolchain.toml
WORKDIR /build/litellm-rust
RUN apt-get update \
&& apt-get install -y --no-install-recommends \
python3 python3-dev pkg-config libssl-dev clang \
python3 python3-dev python3-pip pkg-config libssl-dev clang \
&& rm -rf /var/lib/apt/lists/* \
&& cargo install cargo-chef --locked --version 0.1.77
WORKDIR /build/litellm-rust
# ---- Planner ----------------------------------------------------------------
# Produce the dependency recipe from the rust workspace manifests + Cargo.lock.
@ -43,6 +48,19 @@ RUN cargo chef cook --locked --release \
COPY litellm-rust/ .
RUN cargo build --locked --release -p litellm-ai-gateway --bin litellm-ai-gateway --features server,python-config
# The root pyproject builds with maturin against litellm-rust/crates/python-bridge,
# so the wheel is built here, next to the crate sources and the cargo toolchain,
# and the runtime stage installs the artifact instead of compiling anything.
# litellm[proxy] pins litellm-enterprise and litellm-proxy-extras to the versions
# in this repo, and those hit PyPI hours after every version bump merges, so both
# wheels are built from the repo too instead of being resolved from PyPI.
COPY pyproject.toml README.md LICENSE /build/
COPY litellm/ /build/litellm/
COPY enterprise/ /build/enterprise/
COPY litellm-proxy-extras/ /build/litellm-proxy-extras/
RUN pip3 wheel --no-cache-dir --no-deps --wheel-dir /build/dist \
/build /build/enterprise /build/litellm-proxy-extras
# ---- Runtime ----------------------------------------------------------------
# python:3.11-slim-bookworm ships libpython3.11, matching the builder's PyO3
# 3.11 ABI so the embedded interpreter links and imports cleanly.
@ -56,11 +74,16 @@ RUN apt-get update \
WORKDIR /app
# Install litellm (with proxy extras) FROM THIS REPO'S SOURCE so
# `import litellm.proxy.read_model_list` works — it is not on PyPI yet. Copy the
# package + packaging metadata, then pip install the proxy extra.
COPY pyproject.toml README.md LICENSE ./
COPY litellm/ ./litellm/
RUN pip install --no-cache-dir ".[proxy]"
# `import litellm.proxy.read_model_list` works — it is not on PyPI yet. The two
# sibling wheels come from the builder as well, so the pins in litellm[proxy]
# resolve against them and never wait on a PyPI publish.
COPY --from=builder /build/dist/*.whl /tmp/wheels/
RUN wheel="$(ls /tmp/wheels/litellm-*.whl)" \
&& pip install --no-cache-dir \
/tmp/wheels/litellm_enterprise-*.whl \
/tmp/wheels/litellm_proxy_extras-*.whl \
"${wheel}[proxy]" \
&& rm -rf /tmp/wheels
# The compiled gateway binary (pure-Rust realtime hot path; Python is load-time
# only).

View file

@ -9,19 +9,28 @@
# Strategy: ignore everything, then re-include only what the build needs:
# - litellm/ (pip install . needs the full package + proxy reader)
# - litellm-rust/ (the rust workspace; Cargo.lock + crate sources)
# - pyproject.toml / README.md / LICENSE (packaging metadata for pip install)
# - enterprise/ (litellm/proxy/enterprise symlinks into it; maturin walks it)
# - litellm-proxy-extras/ (built into a wheel alongside enterprise/ for litellm[proxy])
# - pyproject.toml / README.md / LICENSE (packaging metadata for the wheel build)
# - rust-toolchain.toml (the pinned channel every cargo call in the build uses)
*
# --- re-include the build inputs ---
!litellm/
!litellm-rust/
!enterprise/
!litellm-proxy-extras/
!pyproject.toml
!rust-toolchain.toml
!README.md
!LICENSE
# --- prune heavy / irrelevant subpaths back out of the re-included trees ---
# Rust build artifacts (huge; regenerated in the builder).
**/target/
# Committed python distribution artifacts; the wheel build does not read them.
enterprise/dist/
litellm-proxy-extras/dist/
# Python caches and compiled bytecode.
**/__pycache__/
**/*.pyc

View file

@ -269,7 +269,11 @@ fn guardrail_error_to_core_error(error: GuardrailError) -> Error {
fn core_error_kind(error: &Error) -> &'static str {
match error {
Error::Auth(_) | Error::MissingApiKey { .. } => "AuthError",
Error::Auth(_)
| Error::MissingApiKey { .. }
| Error::MissingAzureAiCredentials
| Error::MissingAzureDocumentIntelligenceCredentials
| Error::MissingReductoApiKey => "AuthError",
Error::InvalidProvider(_) => "InvalidProvider",
Error::InvalidRequest(_) => "InvalidRequest",
Error::InvalidType { .. } => "InvalidType",

View file

@ -15,6 +15,8 @@ use std::time::Duration;
use futures_util::stream::{SplitSink, SplitStream};
use futures_util::{Sink, SinkExt, Stream, StreamExt};
use litellm_core::AuthError;
use litellm_core::auth::error::MissingCredential;
use litellm_core::error::Error;
use litellm_core::realtime::transformation::RealtimeProviderConfig;
use litellm_core::realtime::types::RealtimeEvent;
@ -32,8 +34,6 @@ use crate::io::tls::connect_upstream;
/// Environment variable holding the OpenAI API key (last-resort fallback).
const OPENAI_API_KEY_ENV: &str = "OPENAI_API_KEY";
const MISSING_KEY_MESSAGE: &str = "Missing OpenAI API Key - a realtime call is being made but no key was passed via params or the OPENAI_API_KEY environment variable";
/// Default **idle** timeout: if neither side sends a frame for this long, the
/// session is reaped. It resets on any activity, so it does not cap a healthy
/// (continuously streaming) session — it only frees a stalled one (e.g. a
@ -59,7 +59,7 @@ pub(crate) fn resolve_api_key(api_key: Option<&str>) -> Result<String, Error> {
.ok()
.filter(|key| !key.trim().is_empty())
})
.ok_or_else(|| Error::Auth(MISSING_KEY_MESSAGE.to_string()))
.ok_or_else(|| Error::from(AuthError::from(MissingCredential::OpenAiRealtimeApiKey)))
}
/// Open the upstream WebSocket to OpenAI for `(model, api_key, api_base)`.

View file

@ -4,7 +4,9 @@ use std::time::Duration;
use futures_util::stream::{SplitSink, SplitStream};
use futures_util::{Sink, SinkExt, Stream, StreamExt};
use litellm_core::AuthError;
use litellm_core::Error;
use litellm_core::auth::error::MissingCredential;
use litellm_core::providers::openai::responses::transformation::OPENAI_RESPONSES_WS_CONFIG;
use litellm_core::responses::types::ResponsesWsEvent;
use litellm_core::responses::websocket::ResponsesWebSocketProviderConfig;
@ -23,8 +25,6 @@ use crate::constants::{
};
const OPENAI_API_KEY_ENV: &str = "OPENAI_API_KEY";
const MISSING_KEY_MESSAGE: &str = "Missing OpenAI API Key - a Responses WebSocket call is being made but no key was passed via params or the OPENAI_API_KEY environment variable";
pub type ResponsesUpstreamWs = WebSocketStream<MaybeTlsStream<TcpStream>>;
type UpstreamTx = SplitSink<ResponsesUpstreamWs, Message>;
type UpstreamRx = SplitStream<ResponsesUpstreamWs>;
@ -120,7 +120,7 @@ pub(crate) fn resolve_api_key(api_key: Option<&str>) -> Result<String, Error> {
.ok()
.filter(|value| !value.trim().is_empty())
})
.ok_or_else(|| Error::Auth(MISSING_KEY_MESSAGE.to_string()))
.ok_or_else(|| Error::from(AuthError::from(MissingCredential::OpenAiResponsesApiKey)))
}
async fn dial_upstream(

View file

@ -1,529 +0,0 @@
use std::net::IpAddr;
use std::time::{Duration, Instant};
use base64::Engine;
use base64::engine::general_purpose::STANDARD as BASE64_STANDARD;
use litellm_core::error::Error;
use litellm_core::ocr::transformation::OcrProviderConfig;
use reqwest::Url;
use serde_json::{Map, Value};
use litellm_core::providers::azure_ai::ocr::transformation::{
AZURE_AI_OCR_CONFIG, AZURE_DOCUMENT_INTELLIGENCE_OCR_CONFIG,
};
use litellm_core::providers::mistral::ocr::transformation::MISTRAL_OCR_CONFIG;
use litellm_core::providers::reducto::ocr::transformation as reducto;
use litellm_core::providers::vertex_ai::ocr::transformation as vertex_ai;
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;
pub(super) fn truncate_error_body(body: &str) -> String {
if body.chars().count() <= ERROR_BODY_MAX_CHARS {
return body.to_string();
}
let truncated: String = body.chars().take(ERROR_BODY_MAX_CHARS).collect();
format!("{truncated}... (truncated)")
}
#[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)]
pub(super) fn ocr_provider_config(
provider: &str,
model: &str,
) -> Option<&'static dyn OcrProviderConfig> {
match provider {
"mistral" => Some(&MISTRAL_OCR_CONFIG),
"reducto" => reducto::config_for_model(model),
"azure_ai" if is_azure_document_intelligence_model(model) => {
Some(&AZURE_DOCUMENT_INTELLIGENCE_OCR_CONFIG)
}
"azure_ai" => Some(&AZURE_AI_OCR_CONFIG),
"vertex_ai" if vertex_ai::is_deepseek_model(model) => Some(&VERTEX_AI_DEEPSEEK_OCR_CONFIG),
"vertex_ai" => Some(&VERTEX_AI_OCR_CONFIG),
_ => None,
}
}
fn is_azure_document_intelligence_model(model: &str) -> bool {
let model = model.to_ascii_lowercase();
model.contains("doc-intelligence") || model.contains("documentintelligence")
}
pub(super) fn string_headers(
extra_headers: Option<Map<String, Value>>,
) -> Result<Vec<(String, String)>, Error> {
extra_headers
.unwrap_or_default()
.into_iter()
.map(|(key, value)| {
value
.as_str()
.map(|value| (key.clone(), value.to_string()))
.ok_or_else(|| {
Error::InvalidRequest(format!(
"OCR extra_headers.{key} must be a string, got {}",
litellm_core::error::json_type_name(&value)
))
})
})
.collect()
}
fn document_url_field(document: &Value) -> 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 {
return Ok(None);
};
let field = match doc_type {
"document_url" => "document_url",
"image_url" => "image_url",
_ => return Ok(None),
};
let Some(url) = object.get(field).and_then(Value::as_str) else {
return Ok(None);
};
Ok(Some((field, url)))
}
fn is_url_requiring_fetch(url: &str) -> bool {
!url.starts_with("data:") && (url.starts_with("http://") || url.starts_with("https://"))
}
fn max_document_download_bytes() -> u64 {
let max_size_mb = std::env::var("MAX_IMAGE_URL_DOWNLOAD_SIZE_MB")
.ok()
.and_then(|value| value.parse::<f64>().ok())
.unwrap_or(DEFAULT_MAX_IMAGE_URL_DOWNLOAD_SIZE_MB);
(max_size_mb.max(0.0) * 1024.0 * 1024.0) as u64
}
fn is_blocked_ip(ip: IpAddr) -> bool {
match ip {
IpAddr::V4(ip) => {
ip.is_private()
|| ip.is_loopback()
|| ip.is_link_local()
|| ip.is_broadcast()
|| ip.is_multicast()
|| ip.is_unspecified()
}
IpAddr::V6(ip) => {
let first_segment = ip.segments()[0];
let is_unique_local = (first_segment & 0xfe00) == 0xfc00;
let is_link_local = (first_segment & 0xffc0) == 0xfe80;
ip.is_loopback()
|| ip.is_unspecified()
|| ip.is_multicast()
|| is_unique_local
|| is_link_local
|| ip
.to_ipv4_mapped()
.or_else(|| ip.to_ipv4())
.map(|v4| is_blocked_ip(IpAddr::V4(v4)))
.unwrap_or(false)
}
}
}
fn blocked_url_error(url: &Url) -> Error {
Error::InvalidRequest(format!(
"OCR document URL rejected by SSRF protection: {url}"
))
}
async fn validate_safe_fetch_url(url: &Url) -> Result<(), Error> {
if !matches!(url.scheme(), "http" | "https") {
return Err(blocked_url_error(url));
}
let host = url.host_str().ok_or_else(|| blocked_url_error(url))?;
if let Ok(ip) = host.parse::<IpAddr>() {
if is_blocked_ip(ip) {
return Err(blocked_url_error(url));
}
return Ok(());
}
let port = url
.port_or_known_default()
.ok_or_else(|| blocked_url_error(url))?;
let addresses = tokio::net::lookup_host((host, port))
.await
.map_err(|err| Error::Network(err.to_string()))?;
let mut saw_address = false;
for address in addresses {
saw_address = true;
if is_blocked_ip(address.ip()) {
return Err(blocked_url_error(url));
}
}
if !saw_address {
return Err(blocked_url_error(url));
}
Ok(())
}
fn redirect_location(response: &reqwest::Response, url: &Url) -> Result<Url, Error> {
let location = response
.headers()
.get(reqwest::header::LOCATION)
.and_then(|value| value.to_str().ok())
.ok_or_else(|| {
Error::InvalidResponse("OCR document redirect missing Location header".to_string())
})?;
url.join(location)
.map_err(|err| Error::InvalidResponse(format!("invalid OCR document redirect: {err}")))
}
async fn safe_get_document_url(url: &str) -> Result<(Url, reqwest::Response), Error> {
let client = reqwest::Client::builder()
.redirect(reqwest::redirect::Policy::none())
.build()
.map_err(|err| Error::Network(err.to_string()))?;
let mut current_url = Url::parse(url)
.map_err(|err| Error::InvalidRequest(format!("invalid OCR document URL: {err}")))?;
for _ in 0..MAX_SAFE_FETCH_REDIRECTS {
validate_safe_fetch_url(&current_url).await?;
let response = client
.get(current_url.clone())
.send()
.await
.map_err(|err| Error::Network(err.to_string()))?;
if !response.status().is_redirection() {
return Ok((current_url, response));
}
current_url = redirect_location(&response, &current_url)?;
}
Err(Error::InvalidRequest(
"Too many redirects while fetching OCR document URL".to_string(),
))
}
fn enforce_download_size(content_length: u64, max_bytes: u64, url: &Url) -> Result<(), Error> {
if max_bytes == 0 {
return Err(Error::InvalidRequest(format!(
"OCR document URL download is disabled (MAX_IMAGE_URL_DOWNLOAD_SIZE_MB=0). url={url}"
)));
}
if content_length > max_bytes {
let size_mb = content_length as f64 / (1024.0 * 1024.0);
let max_size_mb = max_bytes as f64 / (1024.0 * 1024.0);
return Err(Error::InvalidRequest(format!(
"OCR document size ({size_mb:.2}MB) exceeds maximum allowed size ({max_size_mb:.2}MB). url={url}"
)));
}
Ok(())
}
async fn read_response_with_limit(
mut response: reqwest::Response,
url: &Url,
) -> Result<Vec<u8>, Error> {
let max_bytes = max_document_download_bytes();
if let Some(content_length) = response.content_length() {
enforce_download_size(content_length, max_bytes, url)?;
} else {
enforce_download_size(0, max_bytes, url)?;
}
let mut bytes = Vec::new();
let mut bytes_downloaded: u64 = 0;
while let Some(chunk) = response
.chunk()
.await
.map_err(|err| Error::Network(err.to_string()))?
{
bytes_downloaded += chunk.len() as u64;
enforce_download_size(bytes_downloaded, max_bytes, url)?;
bytes.extend_from_slice(&chunk);
}
Ok(bytes)
}
pub(super) async fn convert_document_url_to_data_uri(document: Value) -> Result<Value, Error> {
let Some((field, url)) = document_url_field(&document)? else {
return Ok(document);
};
if !is_url_requiring_fetch(url) {
return Ok(document);
}
let (final_url, response) = safe_get_document_url(url).await?;
let status = response.status();
if !status.is_success() {
let body = response.text().await.unwrap_or_default();
return Err(Error::Http {
status: status.as_u16(),
body: truncate_error_body(&body),
});
}
let content_type = response
.headers()
.get(reqwest::header::CONTENT_TYPE)
.and_then(|value| value.to_str().ok())
.and_then(|value| value.split(';').next())
.map(str::trim)
.filter(|value| !value.is_empty())
.unwrap_or("application/octet-stream")
.to_string();
let bytes = read_response_with_limit(response, &final_url).await?;
let data_uri = format!(
"data:{content_type};base64,{}",
BASE64_STANDARD.encode(bytes)
);
let mut transformed = document
.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}"
))),
}
}
#[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)]
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;
}
}
#[cfg(test)]
mod tests {
use litellm_core::ocr::transformation::OcrResponseHandling;
use serde_json::json;
use super::*;
#[test]
fn blocks_private_and_metadata_ips() {
assert!(is_blocked_ip("127.0.0.1".parse().unwrap()));
assert!(is_blocked_ip("10.0.0.1".parse().unwrap()));
assert!(is_blocked_ip("169.254.169.254".parse().unwrap()));
assert!(is_blocked_ip("::1".parse().unwrap()));
assert!(is_blocked_ip("fd00::1".parse().unwrap()));
assert!(is_blocked_ip("fe80::1".parse().unwrap()));
assert!(is_blocked_ip("::ffff:169.254.169.254".parse().unwrap()));
assert!(is_blocked_ip("::ffff:10.0.0.1".parse().unwrap()));
assert!(!is_blocked_ip("8.8.8.8".parse().unwrap()));
assert!(!is_blocked_ip("::ffff:8.8.8.8".parse().unwrap()));
}
#[tokio::test]
async fn convert_document_url_rejects_loopback_fetch() {
let error = convert_document_url_to_data_uri(json!({
"type": "image_url",
"image_url": "http://127.0.0.1/image.png"
}))
.await
.unwrap_err();
assert!(matches!(
error,
Error::InvalidRequest(message)
if message.contains("SSRF protection")
));
}
#[tokio::test]
async fn convert_document_url_leaves_data_uri_untouched() {
let document = json!({
"type": "image_url",
"image_url": "data:image/png;base64,abcd"
});
let transformed = convert_document_url_to_data_uri(document.clone())
.await
.unwrap();
assert_eq!(transformed, document);
}
#[test]
fn truncate_error_body_passes_short_strings_through() {
let body = "Unauthorized";
assert_eq!(truncate_error_body(body), "Unauthorized");
}
#[test]
fn truncate_error_body_caps_long_payloads() {
let body = "x".repeat(306);
let truncated = truncate_error_body(&body);
assert!(truncated.ends_with("... (truncated)"));
let prefix_chars = truncated
.strip_suffix("... (truncated)")
.expect("truncated marker present")
.chars()
.count();
assert_eq!(prefix_chars, 256);
}
#[test]
fn truncate_error_body_does_not_split_multibyte_chars() {
let body = "é".repeat(266);
let truncated = truncate_error_body(&body);
assert!(truncated.is_char_boundary(truncated.len()));
}
#[test]
fn ocr_dispatch_supports_migrated_providers() {
assert!(ocr_provider_config("mistral", "mistral-ocr-latest").is_some());
assert!(
ocr_provider_config("azure_ai", "pixtral-12b-2409")
.expect("azure ai config resolves")
.requires_data_uri_document()
);
assert_eq!(
ocr_provider_config("azure_ai", "doc-intelligence/prebuilt-read")
.expect("document intelligence config resolves")
.response_handling(),
OcrResponseHandling::AzureDocumentIntelligencePoll
);
assert!(
ocr_provider_config("vertex_ai", "deepseek-ocr-maas")
.expect("vertex deepseek config resolves")
.supported_ocr_params()
.contains(&"temperature")
);
assert!(ocr_provider_config("openai", "gpt-4o").is_none());
}
#[test]
fn string_headers_accepts_string_values() {
let headers = json!({
"x-trace-id": "trace-1"
})
.as_object()
.unwrap()
.clone();
assert_eq!(
string_headers(Some(headers)).expect("string headers accepted"),
vec![("x-trace-id".to_string(), "trace-1".to_string())]
);
}
#[test]
fn string_headers_rejects_non_string_values() {
let headers = json!({
"x-retry-count": 3
})
.as_object()
.unwrap()
.clone();
let err = string_headers(Some(headers)).expect_err("non-string header rejected");
assert_eq!(
err,
Error::InvalidRequest(
"OCR extra_headers.x-retry-count must be a string, got number".to_string()
)
);
}
}

View file

@ -1,84 +0,0 @@
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(
request: PreparedOcrRequest,
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_with_params(
&request.model,
response_json,
&request.optional_params,
)?
.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_with_params(
&request.model,
response_json,
&request.optional_params,
)?
.into_json())
}

View file

@ -1,401 +0,0 @@
use litellm_core::call_lifecycle::{CallLifecycleContext, CallLifecycleHooks, CallLifecycleTiming};
use litellm_core::error::Error;
use litellm_core::providers::reducto::ocr::transformation::{
build_upload_request, extract_document_source, extract_upload_file_id,
};
use serde_json::{Map, Value, json};
use std::future::Future;
use std::pin::Pin;
use super::common_utils::{convert_document_url_to_data_uri, string_headers, truncate_error_body};
use super::types::{PreparedOcrRequest, ProviderOcrRequest};
use crate::client::http_client;
use crate::integrations::custom_guardrail::{
CustomGuardrailRunner, GuardrailContext, GuardrailError, GuardrailRequest,
};
use crate::integrations::custom_logger::{
CallType, CallbackTiming, CallbackValue, CustomLoggerRunner, LoggingError, ModelCallDetails,
};
use crate::integrations::types::{
RequestMetadata, StandardLoggingMetadata, StandardLoggingPayload,
};
pub(crate) struct OcrLifecycleHooks {
logger_runner: CustomLoggerRunner,
guardrail_runner: CustomGuardrailRunner,
request_metadata: RequestMetadata,
}
type OcrFuture<'a, T> = Pin<Box<dyn Future<Output = Result<T, Error>> + Send + 'a>>;
type OcrLogFuture<'a> = Pin<Box<dyn Future<Output = ()> + Send + 'a>>;
impl OcrLifecycleHooks {
pub(crate) fn new(
logger_runner: CustomLoggerRunner,
guardrail_runner: CustomGuardrailRunner,
request_metadata: RequestMetadata,
) -> Self {
Self {
logger_runner,
guardrail_runner,
request_metadata,
}
}
async fn run_pre_call_guardrails(
&self,
request: PreparedOcrRequest,
) -> Result<PreparedOcrRequest, Error> {
if self.guardrail_runner.is_empty() {
return Ok(request);
}
let context = guardrail_context(&self.request_metadata);
let guardrail_request = GuardrailRequest::new(json!({
"model": request.model,
"custom_llm_provider": request.custom_llm_provider,
"document": request.document,
"optional_params": request.optional_params,
}));
let (guardrail_request, _) = self
.guardrail_runner
.run_pre_call(&context, guardrail_request)
.await
.map_err(guardrail_error_to_core_error)?;
let (document, optional_params) = parse_ocr_pre_call_guardrail_request(guardrail_request)?;
let optional_params = match &request.config {
Ok(config) => config.map_ocr_params(&optional_params),
Err(_) => optional_params,
};
Ok(PreparedOcrRequest {
document,
optional_params,
..request
})
}
pub(crate) async fn prepare_provider_request(
&self,
request: PreparedOcrRequest,
) -> Result<ProviderOcrRequest, Error> {
let config = request.config?;
let env_lookup = |key: &str| std::env::var(key).ok();
let upstream_headers = config.validate_environment(
string_headers(request.extra_headers)?,
request.api_key.as_deref(),
&env_lookup,
)?;
let url = config.complete_url(
request.api_base.as_deref(),
&request.model,
&request.optional_params,
&env_lookup,
)?;
let model = request.model.clone();
let custom_llm_provider = request.custom_llm_provider.clone();
let is_reducto = custom_llm_provider == "reducto";
let document = if is_reducto {
let guarded_document = self
.run_during_call_guardrails(&model, &custom_llm_provider, &url, request.document)
.await?;
upload_reducto_document(
&guarded_document,
request.api_base.as_deref(),
request.timeout,
&upstream_headers,
)
.await?
} else if config.requires_data_uri_document() {
convert_document_url_to_data_uri(request.document).await?
} else {
request.document
};
let optional_params = request.optional_params;
let body = config
.transform_ocr_request(&request.model, document, optional_params.clone())?
.data;
let body = if is_reducto {
body
} else {
self.run_during_call_guardrails(&model, &custom_llm_provider, &url, body)
.await?
};
Ok(ProviderOcrRequest {
model,
config,
url,
body,
optional_params,
upstream_headers,
timeout: request.timeout,
})
}
async fn run_during_call_guardrails(
&self,
model: &str,
custom_llm_provider: &str,
url: &str,
body: Value,
) -> Result<Value, Error> {
if self.guardrail_runner.is_empty() {
return Ok(body);
}
let context = guardrail_context(&self.request_metadata);
let guardrail_request = GuardrailRequest::new(json!({
"model": model,
"custom_llm_provider": custom_llm_provider,
"url": url,
"body": body,
}));
let (guardrail_request, _) = self
.guardrail_runner
.run_during_call(&context, guardrail_request)
.await
.map_err(guardrail_error_to_core_error)?;
parse_ocr_during_call_guardrail_request(guardrail_request)
}
fn standard_logging_payload(
&self,
context: &CallLifecycleContext,
timing: &CallLifecycleTiming,
) -> StandardLoggingPayload {
StandardLoggingPayload {
id: context.litellm_call_id.clone(),
litellm_call_id: context.litellm_call_id.clone(),
call_type: context.call_type.clone(),
model: context.model.clone(),
custom_llm_provider: context.custom_llm_provider.clone(),
response_cost: 0.0,
prompt_tokens: 0,
completion_tokens: 0,
total_tokens: 0,
start_time: timing.start_time,
end_time: timing.end_time,
stream: false,
metadata: StandardLoggingMetadata {
user_api_key_hash: self.request_metadata.user_api_key_hash.clone(),
user_api_key_user_id: self.request_metadata.user_api_key_user_id.clone(),
user_api_key_team_id: self.request_metadata.user_api_key_team_id.clone(),
..Default::default()
},
messages: None,
}
}
}
async fn upload_reducto_document(
document: &Value,
api_base: Option<&str>,
timeout: Option<std::time::Duration>,
upstream_headers: &[(String, String)],
) -> Result<Value, Error> {
let source = extract_document_source(document)?;
let Some(authorization) = upstream_headers
.iter()
.find(|(name, _)| name.eq_ignore_ascii_case("authorization"))
.map(|(_, value)| value.as_str())
else {
return Err(Error::Auth(
"Reducto upload requires an Authorization header".to_string(),
));
};
let Some(upload) = build_upload_request(source, authorization, api_base) else {
return Ok(document.clone());
};
let part = reqwest::multipart::Part::bytes(upload.bytes)
.file_name(upload.file_name)
.mime_str(&upload.mime_type)
.map_err(|error| Error::InvalidRequest(error.to_string()))?;
let form = reqwest::multipart::Form::new().part("file", part);
let mut request_builder = http_client().post(upload.url).multipart(form);
for (name, value) in upstream_headers {
if !name.eq_ignore_ascii_case("content-type")
&& !name.eq_ignore_ascii_case("content-length")
{
request_builder = request_builder.header(name, value);
}
}
if let Some(timeout) = timeout {
request_builder = request_builder.timeout(timeout);
}
let response = request_builder
.send()
.await
.map_err(|error| Error::Network(error.to_string()))?;
let status = response.status();
let body = response
.text()
.await
.map_err(|error| Error::Network(error.to_string()))?;
if !status.is_success() {
return Err(Error::Http {
status: status.as_u16(),
body: truncate_error_body(&body),
});
}
let response_json: Value = serde_json::from_str(&body).map_err(|error| {
Error::InvalidResponse(format!("invalid Reducto upload response JSON: {error}"))
})?;
let file_id = extract_upload_file_id(&response_json)?;
Ok(json!({"type": "document_url", "document_url": file_id}))
}
impl CallLifecycleHooks<PreparedOcrRequest, PreparedOcrRequest, Value> for OcrLifecycleHooks {
type PreCallFuture<'a> = OcrFuture<'a, PreparedOcrRequest>;
type DuringCallFuture<'a> = OcrFuture<'a, PreparedOcrRequest>;
type SuccessFuture<'a> = OcrLogFuture<'a>;
type FailureFuture<'a> = OcrLogFuture<'a>;
fn async_pre_call_hook<'a>(
&'a self,
_context: &'a CallLifecycleContext,
request: PreparedOcrRequest,
) -> Self::PreCallFuture<'a> {
Box::pin(async move { self.run_pre_call_guardrails(request).await })
}
fn async_during_call_hook<'a>(
&'a self,
_context: &'a CallLifecycleContext,
request: PreparedOcrRequest,
) -> Self::DuringCallFuture<'a> {
Box::pin(async move { Ok(request) })
}
#[tracing::instrument(
name = "success_callback",
target = "litellm::function_trace",
level = "trace",
skip_all
)]
fn async_log_success_event<'a>(
&'a self,
context: &'a CallLifecycleContext,
response: &'a Value,
timing: &'a CallLifecycleTiming,
) -> Self::SuccessFuture<'a> {
Box::pin(async move {
if self.logger_runner.is_empty() {
return;
}
let response_obj = CallbackValue::new("ocr", response.clone());
self.logger_runner
.async_log_success_event(
&ModelCallDetails::from_standard_logging_payload(
self.standard_logging_payload(context, timing),
),
&response_obj,
CallbackTiming::new(timing.start_time, timing.end_time),
)
.await;
})
}
#[tracing::instrument(
name = "failure_callback",
target = "litellm::function_trace",
level = "trace",
skip_all
)]
fn async_log_failure_event<'a>(
&'a self,
context: &'a CallLifecycleContext,
error: &'a Error,
timing: &'a CallLifecycleTiming,
) -> Self::FailureFuture<'a> {
Box::pin(async move {
if self.logger_runner.is_empty() {
return;
}
let logging_error = LoggingError {
message: error.to_string(),
kind: core_error_kind(error).to_string(),
};
let response_obj = CallbackValue::new(
"error",
json!({
"message": logging_error.message,
"kind": logging_error.kind,
}),
);
self.logger_runner
.async_log_failure_event(
&ModelCallDetails::from_standard_logging_payload(
self.standard_logging_payload(context, timing),
)
.with_failure_error(logging_error),
Some(&response_obj),
CallbackTiming::new(timing.start_time, timing.end_time),
)
.await;
})
}
}
fn guardrail_context(metadata: &RequestMetadata) -> GuardrailContext {
GuardrailContext {
call_type: CallType::Ocr,
selected_guardrails: Vec::new(),
metadata: std::collections::HashMap::new(),
user_api_key_hash: metadata.user_api_key_hash.clone(),
user_api_key_user_id: metadata.user_api_key_user_id.clone(),
user_api_key_team_id: metadata.user_api_key_team_id.clone(),
trace_parent: None,
}
}
fn parse_ocr_pre_call_guardrail_request(
request: GuardrailRequest,
) -> Result<(Value, Map<String, Value>), Error> {
let Value::Object(mut data) = request.data else {
return Err(Error::InvalidRequest(
"OCR pre_call guardrail must return an object".to_string(),
));
};
let document = data.remove("document").ok_or_else(|| {
Error::InvalidRequest("OCR pre_call guardrail removed document".to_string())
})?;
let optional_params = match data.remove("optional_params") {
Some(Value::Object(params)) => params,
Some(_) => {
return Err(Error::InvalidRequest(
"OCR pre_call guardrail optional_params must be an object".to_string(),
));
}
None => Map::new(),
};
Ok((document, optional_params))
}
fn parse_ocr_during_call_guardrail_request(request: GuardrailRequest) -> Result<Value, Error> {
let Value::Object(mut data) = request.data else {
return Err(Error::InvalidRequest(
"OCR during_call guardrail must return an object".to_string(),
));
};
data.remove("body")
.ok_or_else(|| Error::InvalidRequest("OCR during_call guardrail removed body".to_string()))
}
fn guardrail_error_to_core_error(error: GuardrailError) -> Error {
Error::InvalidRequest(format!("{}: {}", error.kind, error.message))
}
fn core_error_kind(error: &Error) -> &'static str {
match error {
Error::Auth(_) | Error::MissingApiKey { .. } => "AuthError",
Error::InvalidProvider(_) => "InvalidProvider",
Error::InvalidRequest(_) => "InvalidRequest",
Error::InvalidType { .. } => "InvalidType",
Error::MissingField(_) => "MissingField",
Error::Http { .. } => "HttpError",
Error::InvalidResponse(_) => "InvalidResponse",
Error::Network(_) => "NetworkError",
Error::Connect(_) => "ConnectError",
Error::Routing(_) => "RoutingError",
Error::Unsupported(_) => "UnsupportedRequest",
}
}

View file

@ -1,174 +1,127 @@
use litellm_core::Error;
use litellm_core::call_lifecycle::CallLifecycle;
use litellm_core::ocr::{
OcrClient,
wire::{OcrWireRequest, decode_request},
};
use serde_json::Value;
mod common_utils;
mod handler;
mod hooks;
mod prepare;
mod types;
pub use types::OcrRequest;
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> {
let PreparedOcrCall { request, hooks } = prepare_ocr_call(request);
CallLifecycle::default()
.run_request(request, &hooks, |request| {
execute_ocr_provider_call(request, &hooks)
})
core_ocr(request).await
}
async fn core_ocr(request: OcrRequest<'_>) -> Result<Value, Error> {
validate_host_hooks(&request)?;
let client = OcrClient::new(crate::client::http_client().clone())?;
let core_request = decode_request(OcrWireRequest {
model: request.model.to_string(),
document: request.document,
api_key: request.api_key.map(str::to_string),
api_base: request.api_base.map(str::to_string),
custom_llm_provider: request.custom_llm_provider.map(str::to_string),
extra_headers: request.extra_headers,
optional_params: request.optional_params,
input_sources: Default::default(),
timeout_seconds: request.timeout.map(|timeout| timeout.as_secs_f64()),
})?;
client
.perform(core_request)
.await
.map(|response| response.into_json())
}
fn validate_host_hooks(request: &OcrRequest<'_>) -> Result<(), Error> {
if !request.guardrails.is_empty() {
return Err(Error::Unsupported(
"OCR host guardrails are not wired to the core path",
));
}
if !request.callbacks.is_empty() {
return Err(Error::Unsupported(
"OCR host callbacks are not wired to the core path",
));
}
Ok(())
}
#[cfg(test)]
mod tests {
use std::sync::Arc;
use litellm_core::ocr::wire::is_supported_request;
use serde_json::{Map, json};
use tokio::io::{AsyncReadExt, AsyncWriteExt};
use tokio::net::{TcpListener, TcpStream};
use super::{OcrRequest, ocr};
use crate::integrations::types::RequestMetadata;
use super::{OcrRequest, validate_host_hooks};
use crate::integrations::custom_guardrail::{CustomGuardrail, GuardrailEventHook};
use crate::integrations::custom_logger::CustomLogger;
async fn read_http_request(socket: &mut TcpStream) -> String {
let mut request = Vec::new();
let mut buffer = [0_u8; 1024];
let header_end = loop {
let n = socket.read(&mut buffer).await.expect("reads request");
if n == 0 {
break request.len();
}
request.extend_from_slice(&buffer[..n]);
if let Some(position) = request.windows(4).position(|window| window == b"\r\n\r\n") {
break position + 4;
}
};
let headers = String::from_utf8_lossy(&request[..header_end]);
let content_length = headers
.lines()
.find_map(|line| {
let (name, value) = line.split_once(':')?;
name.eq_ignore_ascii_case("content-length")
.then(|| value.trim().parse::<usize>().ok())
.flatten()
})
.unwrap_or(0);
while request.len().saturating_sub(header_end) < content_length {
let n = socket.read(&mut buffer).await.expect("reads body");
if n == 0 {
break;
}
request.extend_from_slice(&buffer[..n]);
struct TestGuardrail;
impl CustomGuardrail for TestGuardrail {
fn guardrail_name(&self) -> &str {
"test"
}
fn supported_event_hooks(&self) -> &[GuardrailEventHook] {
&[]
}
String::from_utf8(request).expect("request is utf8")
}
fn base_ocr_request(model: &str) -> OcrRequest<'_> {
struct TestLogger;
impl CustomLogger for TestLogger {}
fn request() -> OcrRequest<'static> {
OcrRequest {
model,
document: json!({
"type": "document_url",
"document_url": "https://example.com/doc.pdf"
}),
api_key: Some("sk-test"),
model: "model",
document: json!({"type":"image_url","image_url":"data:image/png;base64,YQ=="}),
api_key: None,
api_base: None,
custom_llm_provider: None,
custom_llm_provider: Some("mistral"),
extra_headers: None,
optional_params: Map::new(),
timeout: None,
callbacks: Vec::new(),
guardrails: Vec::new(),
request_metadata: RequestMetadata::default(),
request_metadata: Default::default(),
litellm_call_id: None,
}
}
#[tokio::test]
async fn reducto_file_upload_then_parse_maps_response() {
let listener = TcpListener::bind("127.0.0.1:0")
.await
.expect("test listener binds");
let address = listener.local_addr().expect("listener has local address");
let server = tokio::spawn(async move {
let (mut upload_socket, _) = listener.accept().await.expect("accepts upload request");
let upload_request = read_http_request(&mut upload_socket).await;
let upload_body = r#"{"file_id":"reducto://uploaded.pdf"}"#;
let upload_response = format!(
"HTTP/1.1 200 OK\r\ncontent-type: application/json\r\ncontent-length: {}\r\nconnection: close\r\n\r\n{}",
upload_body.len(),
upload_body
);
upload_socket
.write_all(upload_response.as_bytes())
.await
.expect("writes upload response");
#[test]
fn core_activation_includes_migrated_providers() {
assert!(is_supported_request("model", Some("mistral")));
assert!(is_supported_request("pixtral-12b", Some("azure_ai")));
assert!(is_supported_request(
"doc-intelligence/prebuilt-layout",
Some("azure_ai")
));
assert!(is_supported_request("parse-v3", Some("reducto")));
assert!(is_supported_request("mistral-ocr", Some("vertex_ai")));
assert!(is_supported_request("deepseek-ocr", Some("vertex_ai")));
}
let (mut parse_socket, _) = listener.accept().await.expect("accepts parse request");
let parse_request = read_http_request(&mut parse_socket).await;
let parse_body = r#"{"job_id":"job_123","usage":{"num_pages":3,"credits":3},"result":{"chunks":[{"content":"Page 1 block A","blocks":[{"content":"Page 1 block A","bbox":{"page":1},"kind":"text"}]},{"content":"Page 2 block A","blocks":[{"content":"Page 2 block A","bbox":{"page":2},"kind":"table"}]},{"content":"Page 1 block B","blocks":[{"content":"Page 1 block B","bbox":{"page":1},"kind":"text"}]},{"content":"Page 3 block A","blocks":[{"content":"Page 3 block A","bbox":{"page":3},"kind":"figure"}]}]}}"#;
let parse_response = format!(
"HTTP/1.1 200 OK\r\ncontent-type: application/json\r\ncontent-length: {}\r\nconnection: close\r\n\r\n{}",
parse_body.len(),
parse_body
);
parse_socket
.write_all(parse_response.as_bytes())
.await
.expect("writes parse response");
(upload_request, parse_request)
});
let api_base = format!("http://{address}");
let mut request = base_ocr_request("reducto/parse-v3");
request.api_base = Some(&api_base);
request.api_key = None;
request.extra_headers = Some(Map::from_iter([
("Authorization".to_string(), json!("Bearer test-key")),
("x-trace-id".to_string(), json!("trace-1")),
]));
request.document = json!({
"type": "document_url",
"document_url": "data:application/pdf;base64,JVBERi0xLjQ="
});
request.optional_params = Map::from_iter([
(
"formatting".to_string(),
json!({"table_output_format": "html"}),
),
("retrieval".to_string(), json!({"chunk_mode": "section"})),
("settings".to_string(), json!({"ocr_system": "standard"})),
]);
#[test]
fn core_path_rejects_unwired_guardrails() {
let request = OcrRequest {
guardrails: vec![Arc::new(TestGuardrail)],
..request()
};
let error = validate_host_hooks(&request).unwrap_err();
assert!(error.to_string().contains("guardrails are not wired"));
}
let response = ocr(request).await.expect("Reducto OCR succeeds");
assert_eq!(response["pages"].as_array().map(Vec::len), Some(3));
assert_eq!(
response["pages"][0]["markdown"],
"Page 1 block A\n\nPage 1 block B"
);
assert_eq!(response["pages"][1]["markdown"], "Page 2 block A");
assert_eq!(response["pages"][2]["markdown"], "Page 3 block A");
assert_eq!(response["usage_info"]["pages_processed"], 3);
assert_eq!(response["usage_info"]["credits"], 3);
assert_eq!(response["provider_native_response"]["job_id"], "job_123");
let (upload_request, parse_request) = server.await.expect("server task completes");
assert!(
upload_request
.to_ascii_lowercase()
.contains("authorization: bearer test-key")
);
assert!(upload_request.contains("application/pdf"));
assert!(upload_request.contains("%PDF-1.4"));
assert!(upload_request.contains("x-trace-id: trace-1"));
assert!(
parse_request
.to_ascii_lowercase()
.contains("authorization: bearer test-key")
);
assert!(parse_request.contains(r#""input":"reducto://uploaded.pdf""#));
assert!(parse_request.contains(r#""table_output_format":"html""#));
assert!(parse_request.contains(r#""chunk_mode":"section""#));
assert!(parse_request.contains(r#""ocr_system":"standard""#));
#[test]
fn core_path_rejects_unwired_callbacks() {
let request = OcrRequest {
callbacks: vec![Arc::new(TestLogger)],
..request()
};
let error = validate_host_hooks(&request).unwrap_err();
assert!(error.to_string().contains("callbacks are not wired"));
}
}

View file

@ -1,163 +0,0 @@
use std::sync::atomic::{AtomicU64, Ordering};
use std::time::{SystemTime, UNIX_EPOCH};
use litellm_core::routing_utils::provider::{CustomLlmProvider, get_custom_llm_provider};
use serde_json::{Map, Value};
use super::common_utils::ocr_provider_config;
use super::hooks::OcrLifecycleHooks;
use super::types::{OcrRequest, PreparedOcrRequest};
use crate::integrations::custom_guardrail::CustomGuardrailRunner;
use crate::integrations::custom_logger::CustomLoggerRunner;
pub(crate) struct PreparedOcrCall {
pub(crate) request: PreparedOcrRequest,
pub(crate) hooks: OcrLifecycleHooks,
}
#[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)]
pub(crate) fn prepare_ocr_call(request: OcrRequest<'_>) -> PreparedOcrCall {
let call_id = request
.litellm_call_id
.map(str::to_string)
.unwrap_or_else(new_ocr_call_id);
let provider_info = get_custom_llm_provider(request.model, request.custom_llm_provider)
.unwrap_or(CustomLlmProvider {
model: request.model,
custom_llm_provider: "mistral",
});
let model = provider_info.model.to_string();
let custom_llm_provider = provider_info.custom_llm_provider.to_string();
let config = ocr_provider_config(&custom_llm_provider, &model)
.ok_or_else(|| litellm_core::Error::InvalidProvider(custom_llm_provider.clone()))
.and_then(|config| {
validate_request_format(config, &request.optional_params, &custom_llm_provider)?;
Ok(config)
});
let optional_params = match &config {
Ok(config) => {
let supported = config.supported_ocr_params();
let mut mapped = config.map_ocr_params(
&request
.optional_params
.iter()
.filter(|(name, _)| supported.contains(&name.as_str()))
.map(|(name, value)| (name.clone(), value.clone()))
.collect(),
);
for name in [
"vertex_project",
"vertex_ai_project",
"vertex_location",
"vertex_ai_location",
] {
if let Some(value) = request.optional_params.get(name) {
mapped.insert(name.to_string(), value.clone());
}
}
mapped
}
Err(_) => request.optional_params,
};
PreparedOcrCall {
request: PreparedOcrRequest {
config,
model,
custom_llm_provider,
litellm_call_id: call_id,
document: request.document,
api_key: request.api_key.map(str::to_string),
api_base: request.api_base.map(str::to_string),
extra_headers: request.extra_headers,
optional_params,
timeout: request.timeout,
},
hooks: OcrLifecycleHooks::new(
CustomLoggerRunner::new(request.callbacks),
CustomGuardrailRunner::new(request.guardrails),
request.request_metadata,
),
}
}
fn validate_request_format(
config: &'static dyn litellm_core::ocr::transformation::OcrProviderConfig,
optional_params: &Map<String, Value>,
provider: &str,
) -> Result<(), litellm_core::Error> {
let Some(format) = optional_params.get("req_format") else {
return Ok(());
};
match format.as_str() {
Some("litellm") => Ok(()),
Some("native") if config.supported_ocr_params().contains(&"req_format") => Ok(()),
Some("native") => Err(litellm_core::Error::InvalidRequest(format!(
"`req_format=native` is not supported for provider {provider}"
))),
_ => Err(litellm_core::Error::InvalidRequest(format!(
"Invalid `req_format`: {format}. Expected `litellm` or `native`"
))),
}
}
fn new_ocr_call_id() -> String {
static COUNTER: AtomicU64 = AtomicU64::new(1);
let sequence = COUNTER.fetch_add(1, Ordering::Relaxed);
let timestamp = SystemTime::now()
.duration_since(UNIX_EPOCH)
.map(|duration| duration.as_nanos())
.unwrap_or(0);
format!("ocr-{timestamp}-{sequence}")
}
#[cfg(test)]
mod tests {
use litellm_core::error::Error;
use serde_json::{Map, json};
use super::{OcrRequest, prepare_ocr_call};
use crate::integrations::types::RequestMetadata;
fn base_ocr_request(model: &str) -> OcrRequest<'_> {
OcrRequest {
model,
document: json!({
"type": "document_url",
"document_url": "https://example.com/doc.pdf"
}),
api_key: Some("sk-test"),
api_base: None,
custom_llm_provider: None,
extra_headers: None,
optional_params: Map::new(),
timeout: None,
callbacks: Vec::new(),
guardrails: Vec::new(),
request_metadata: RequestMetadata::default(),
litellm_call_id: None,
}
}
fn request_with_format(format: &str) -> OcrRequest<'_> {
let mut request = base_ocr_request("mistral/mistral-ocr-latest");
request.optional_params = Map::from_iter([("req_format".to_string(), json!(format))]);
request
}
#[test]
fn native_format_rejected_for_provider_without_support_as_bad_request() {
let prepared = prepare_ocr_call(request_with_format("native"));
assert!(
matches!(prepared.request.config, Err(Error::InvalidRequest(message)) if message.contains("not supported for provider"))
);
}
#[test]
fn unknown_format_rejected_for_provider_without_support_as_bad_request() {
let prepared = prepare_ocr_call(request_with_format("raw"));
assert!(
matches!(prepared.request.config, Err(Error::InvalidRequest(message)) if message.contains("Invalid `req_format`"))
);
}
}

View file

@ -1,8 +1,6 @@
use std::sync::Arc;
use std::time::Duration;
use litellm_core::call_lifecycle::{CallLifecycleContext, CallLifecycleRequest};
use litellm_core::ocr::transformation::OcrProviderConfig;
use serde_json::{Map, Value};
use crate::integrations::custom_guardrail::CustomGuardrail;
@ -23,37 +21,3 @@ pub struct OcrRequest<'a> {
pub request_metadata: RequestMetadata,
pub litellm_call_id: Option<&'a str>,
}
pub(crate) struct PreparedOcrRequest {
pub(crate) config: Result<&'static dyn OcrProviderConfig, litellm_core::Error>,
pub(crate) model: String,
pub(crate) custom_llm_provider: String,
pub(crate) litellm_call_id: String,
pub(crate) document: Value,
pub(crate) api_key: Option<String>,
pub(crate) api_base: Option<String>,
pub(crate) extra_headers: Option<Map<String, Value>>,
pub(crate) optional_params: Map<String, Value>,
pub(crate) timeout: Option<Duration>,
}
impl CallLifecycleRequest for PreparedOcrRequest {
fn lifecycle_context(&self) -> CallLifecycleContext {
CallLifecycleContext::new(
"ocr",
self.model.clone(),
self.custom_llm_provider.clone(),
self.litellm_call_id.clone(),
)
}
}
pub(crate) struct ProviderOcrRequest {
pub(crate) model: String,
pub(crate) config: &'static dyn OcrProviderConfig,
pub(crate) url: String,
pub(crate) body: Value,
pub(crate) optional_params: Map<String, Value>,
pub(crate) upstream_headers: Vec<(String, String)>,
pub(crate) timeout: Option<Duration>,
}

View file

@ -105,7 +105,11 @@ impl IntoResponse for MessagesRouteError {
StatusCode::NOT_FOUND,
"no messages deployment is configured for this model".to_string(),
),
Error::Auth(_) => (
Error::Auth(_)
| Error::MissingApiKey { .. }
| Error::MissingAzureAiCredentials
| Error::MissingAzureDocumentIntelligenceCredentials
| Error::MissingReductoApiKey => (
StatusCode::BAD_GATEWAY,
"messages provider authentication failed".to_string(),
),
@ -114,8 +118,7 @@ impl IntoResponse for MessagesRouteError {
| Error::Connect(_)
| Error::InvalidResponse(_)
| Error::InvalidType { .. }
| Error::MissingField(_)
| Error::MissingApiKey { .. } => (
| Error::MissingField(_) => (
StatusCode::BAD_GATEWAY,
"messages provider request failed".to_string(),
),

View file

@ -1,641 +0,0 @@
use std::sync::{Arc, Mutex};
use std::time::Duration;
use litellm_ai_gateway::integrations::custom_guardrail::{
CustomGuardrail, GuardrailContext, GuardrailDecision, GuardrailError, GuardrailEventHook,
GuardrailFuture, GuardrailRequest,
};
use litellm_ai_gateway::integrations::custom_logger::{
CallbackTiming, CallbackValue, CustomLogger, LogFuture, ModelCallDetails,
};
use litellm_ai_gateway::integrations::types::RequestMetadata;
use litellm_ai_gateway::ocr::{OcrRequest, ocr};
use litellm_core::error::Error;
#[cfg(feature = "trace-parity")]
use litellm_core::observability::FunctionTrace;
use serde_json::{Map, Value, json};
use tokio::io::{AsyncReadExt, AsyncWriteExt};
use tokio::net::{TcpListener, TcpStream};
#[cfg(feature = "trace-parity")]
use tracing::instrument::WithSubscriber;
async fn read_http_headers(socket: &mut TcpStream) -> String {
let mut request = Vec::new();
let mut buffer = [0_u8; 1024];
loop {
let n = socket.read(&mut buffer).await.expect("reads request");
if n == 0 {
break;
}
request.extend_from_slice(&buffer[..n]);
if request.windows(4).any(|window| window == b"\r\n\r\n") {
break;
}
}
String::from_utf8(request).expect("request is utf8")
}
async fn read_http_request(socket: &mut TcpStream) -> String {
let mut request = Vec::new();
let mut buffer = [0_u8; 1024];
let header_end = loop {
let n = socket.read(&mut buffer).await.expect("reads request");
if n == 0 {
break request.len();
}
request.extend_from_slice(&buffer[..n]);
if let Some(position) = request.windows(4).position(|window| window == b"\r\n\r\n") {
break position + 4;
}
};
let headers = String::from_utf8_lossy(&request[..header_end]);
let content_length = headers
.lines()
.find_map(|line| {
let (name, value) = line.split_once(':')?;
name.eq_ignore_ascii_case("content-length")
.then(|| value.trim().parse::<usize>().ok())
.flatten()
})
.unwrap_or(0);
while request.len().saturating_sub(header_end) < content_length {
let n = socket.read(&mut buffer).await.expect("reads body");
if n == 0 {
break;
}
request.extend_from_slice(&buffer[..n]);
}
String::from_utf8(request).expect("request is utf8")
}
#[derive(Clone, Debug, PartialEq)]
struct RecordedLogEvent {
hook: &'static str,
model: String,
call_type: String,
user_id: Option<String>,
response_object: Option<String>,
error_kind: Option<String>,
}
#[derive(Default)]
struct RecordingOcrLogger {
events: Mutex<Vec<RecordedLogEvent>>,
}
impl RecordingOcrLogger {
fn events(&self) -> Vec<RecordedLogEvent> {
self.events.lock().unwrap().clone()
}
}
impl CustomLogger for RecordingOcrLogger {
fn async_log_success_event<'a>(
&'a self,
model_call_details: &'a ModelCallDetails,
response_obj: &'a CallbackValue,
_timing: CallbackTiming,
) -> LogFuture<'a> {
Box::pin(async move {
self.events.lock().unwrap().push(RecordedLogEvent {
hook: "async_log_success_event",
model: model_call_details.model.clone(),
call_type: model_call_details.call_type.to_string(),
user_id: model_call_details.metadata.user_api_key_user_id.clone(),
response_object: Some(response_obj.object.clone()),
error_kind: None,
});
Ok(())
})
}
fn async_log_failure_event<'a>(
&'a self,
model_call_details: &'a ModelCallDetails,
response_obj: Option<&'a CallbackValue>,
_timing: CallbackTiming,
) -> LogFuture<'a> {
Box::pin(async move {
self.events.lock().unwrap().push(RecordedLogEvent {
hook: "async_log_failure_event",
model: model_call_details.model.clone(),
call_type: model_call_details.call_type.to_string(),
user_id: model_call_details.metadata.user_api_key_user_id.clone(),
response_object: response_obj.map(|value| value.object.clone()),
error_kind: model_call_details
.failure_error
.as_ref()
.map(|error| error.kind.clone()),
});
Ok(())
})
}
}
struct RecordingOcrGuardrail {
hooks: Vec<GuardrailEventHook>,
events: Mutex<Vec<&'static str>>,
block_pre_call: bool,
block_during_call: bool,
}
impl RecordingOcrGuardrail {
fn new(hooks: Vec<GuardrailEventHook>) -> Self {
Self {
hooks,
events: Mutex::new(Vec::new()),
block_pre_call: false,
block_during_call: false,
}
}
fn blocking_pre_call() -> Self {
Self {
hooks: vec![GuardrailEventHook::PreCall],
events: Mutex::new(Vec::new()),
block_pre_call: true,
block_during_call: false,
}
}
fn blocking_during_call() -> Self {
Self {
hooks: vec![GuardrailEventHook::DuringCall],
events: Mutex::new(Vec::new()),
block_pre_call: false,
block_during_call: true,
}
}
fn events(&self) -> Vec<&'static str> {
self.events.lock().unwrap().clone()
}
}
impl CustomGuardrail for RecordingOcrGuardrail {
fn guardrail_name(&self) -> &str {
"recording-ocr-guardrail"
}
fn supported_event_hooks(&self) -> &[GuardrailEventHook] {
&self.hooks
}
fn async_pre_call_hook<'a>(
&'a self,
_context: &'a GuardrailContext,
mut request: GuardrailRequest,
) -> GuardrailFuture<'a> {
Box::pin(async move {
self.events.lock().unwrap().push("async_pre_call_hook");
if self.block_pre_call {
return Ok(GuardrailDecision::Block(GuardrailError::blocked(
"blocked before provider",
)));
}
request.data["document"]["guarded_pre"] = json!(true);
Ok(GuardrailDecision::Mask(request))
})
}
fn async_moderation_hook<'a>(
&'a self,
_context: &'a GuardrailContext,
mut request: GuardrailRequest,
) -> GuardrailFuture<'a> {
Box::pin(async move {
self.events.lock().unwrap().push("async_moderation_hook");
if self.block_during_call {
return Ok(GuardrailDecision::Block(GuardrailError::blocked(
"blocked before provider",
)));
}
request.data["body"]["guarded_during"] = json!(true);
Ok(GuardrailDecision::Mask(request))
})
}
}
fn base_ocr_request(model: &str) -> OcrRequest<'_> {
OcrRequest {
model,
document: json!({
"type": "document_url",
"document_url": "https://example.com/doc.pdf"
}),
api_key: Some("sk-test"),
api_base: None,
custom_llm_provider: None,
extra_headers: None,
optional_params: Map::new(),
timeout: None,
callbacks: Vec::new(),
guardrails: Vec::new(),
request_metadata: RequestMetadata::default(),
litellm_call_id: None,
}
}
#[tokio::test]
async fn reducto_during_call_guardrail_blocks_before_upload() {
let listener = TcpListener::bind("127.0.0.1:0")
.await
.expect("test listener binds");
let address = listener.local_addr().expect("listener has local address");
let api_base = format!("http://{address}");
let guardrail = Arc::new(RecordingOcrGuardrail::blocking_during_call());
let mut request = base_ocr_request("reducto/parse-v3");
request.api_base = Some(&api_base);
request.document = json!({
"type": "document_url",
"document_url": "data:application/pdf;base64,JVBERi0xLjQ="
});
request.guardrails = vec![guardrail.clone()];
let error = ocr(request).await.expect_err("guardrail blocks upload");
assert!(matches!(error, Error::InvalidRequest(_)));
assert_eq!(guardrail.events(), vec!["async_moderation_hook"]);
let accepted = tokio::time::timeout(Duration::from_millis(100), listener.accept()).await;
assert!(accepted.is_err(), "upload socket should not be touched");
}
#[tokio::test]
async fn reducto_upload_error_body_is_truncated() {
let listener = TcpListener::bind("127.0.0.1:0")
.await
.expect("test listener binds");
let address = listener.local_addr().expect("listener has local address");
let server = tokio::spawn(async move {
let (mut socket, _) = listener.accept().await.expect("accepts upload request");
let _request = read_http_request(&mut socket).await;
let body = "x".repeat(300);
let response = format!(
"HTTP/1.1 500 Internal Server Error\r\ncontent-length: {}\r\nconnection: close\r\n\r\n{}",
body.len(),
body
);
socket
.write_all(response.as_bytes())
.await
.expect("writes upload response");
});
let api_base = format!("http://{address}");
let mut request = base_ocr_request("reducto/parse-v3");
request.api_base = Some(&api_base);
request.document = json!({
"type": "document_url",
"document_url": "data:application/pdf;base64,JVBERi0xLjQ="
});
let error = ocr(request).await.expect_err("upload should fail");
assert!(
matches!(error, Error::Http { status: 500, body } if body.chars().count() < 300 && body.ends_with("... (truncated)"))
);
server.await.expect("server task completes");
}
#[tokio::test]
async fn ocr_lifecycle_runs_pre_during_and_success_hooks() {
let listener = TcpListener::bind("127.0.0.1:0")
.await
.expect("test listener binds");
let addr = listener.local_addr().expect("listener has local addr");
let server = tokio::spawn(async move {
let (mut socket, _) = listener.accept().await.expect("accepts one request");
let request = read_http_request(&mut socket).await;
let response_body = r#"{"pages":[{"index":0,"markdown":"ok"}],"model":"mistral-ocr-latest","usage_info":{"pages_processed":1}}"#;
let response = format!(
"HTTP/1.1 200 OK\r\ncontent-type: application/json\r\ncontent-length: {}\r\nconnection: close\r\n\r\n{}",
response_body.len(),
response_body
);
socket
.write_all(response.as_bytes())
.await
.expect("writes response");
request
});
let logger = Arc::new(RecordingOcrLogger::default());
let guardrail = Arc::new(RecordingOcrGuardrail::new(vec![
GuardrailEventHook::PreCall,
GuardrailEventHook::DuringCall,
]));
#[cfg(feature = "trace-parity")]
let trace = FunctionTrace::default();
let api_base = format!("http://{addr}");
let call = ocr(OcrRequest {
model: "mistral-ocr-latest",
document: json!({
"type": "document_url",
"document_url": "https://example.com/doc.pdf"
}),
api_key: Some("sk-test"),
api_base: Some(&api_base),
custom_llm_provider: Some("mistral"),
extra_headers: None,
optional_params: Map::new(),
timeout: Some(Duration::from_secs(5)),
callbacks: vec![logger.clone()],
guardrails: vec![guardrail.clone()],
request_metadata: RequestMetadata {
user_api_key_user_id: Some("user-1".to_string()),
..Default::default()
},
litellm_call_id: Some("ocr-call-1"),
});
#[cfg(feature = "trace-parity")]
let call = call.with_subscriber(trace.dispatcher());
let response = call.await.expect("ocr request succeeds");
assert_eq!(response["pages"][0]["markdown"], "ok");
assert_eq!(
guardrail.events(),
vec!["async_pre_call_hook", "async_moderation_hook"]
);
assert_eq!(
logger.events(),
vec![RecordedLogEvent {
hook: "async_log_success_event",
model: "mistral-ocr-latest".to_string(),
call_type: "ocr".to_string(),
user_id: Some("user-1".to_string()),
response_object: Some("ocr".to_string()),
error_kind: None,
}]
);
#[cfg(feature = "trace-parity")]
assert_eq!(
trace
.events()
.iter()
.filter(|event| event.function.ends_with("_callback"))
.map(|event| event.function)
.collect::<Vec<_>>(),
vec!["success_callback"]
);
let request = server.await.expect("server task completes");
assert!(request.contains(r#""guarded_pre":true"#), "{request}");
assert!(request.contains(r#""guarded_during":true"#), "{request}");
}
#[tokio::test]
async fn ocr_lifecycle_runs_failure_hook_on_provider_error() {
let listener = TcpListener::bind("127.0.0.1:0")
.await
.expect("test listener binds");
let addr = listener.local_addr().expect("listener has local addr");
let server = tokio::spawn(async move {
let (mut socket, _) = listener.accept().await.expect("accepts one request");
let _request = read_http_request(&mut socket).await;
let response_body = "provider failed";
let response = format!(
"HTTP/1.1 500 Internal Server Error\r\ncontent-length: {}\r\nconnection: close\r\n\r\n{}",
response_body.len(),
response_body
);
socket
.write_all(response.as_bytes())
.await
.expect("writes response");
});
let logger = Arc::new(RecordingOcrLogger::default());
#[cfg(feature = "trace-parity")]
let trace = FunctionTrace::default();
let api_base = format!("http://{addr}");
let call = ocr(OcrRequest {
model: "mistral-ocr-latest",
document: json!({
"type": "document_url",
"document_url": "https://example.com/doc.pdf"
}),
api_key: Some("sk-test"),
api_base: Some(&api_base),
custom_llm_provider: Some("mistral"),
extra_headers: None,
optional_params: Map::new(),
timeout: Some(Duration::from_secs(5)),
callbacks: vec![logger.clone()],
guardrails: Vec::new(),
request_metadata: RequestMetadata::default(),
litellm_call_id: Some("ocr-call-2"),
});
#[cfg(feature = "trace-parity")]
let call = call.with_subscriber(trace.dispatcher());
let err = call.await.expect_err("provider error propagates");
assert!(matches!(err, Error::Http { status: 500, .. }));
server.await.expect("server task completes");
assert_eq!(
logger.events(),
vec![RecordedLogEvent {
hook: "async_log_failure_event",
model: "mistral-ocr-latest".to_string(),
call_type: "ocr".to_string(),
user_id: None,
response_object: Some("error".to_string()),
error_kind: Some("HttpError".to_string()),
}]
);
#[cfg(feature = "trace-parity")]
assert_eq!(
trace
.events()
.iter()
.filter(|event| event.function.ends_with("_callback"))
.map(|event| event.function)
.collect::<Vec<_>>(),
vec!["failure_callback"]
);
}
#[tokio::test]
async fn ocr_lifecycle_pre_call_block_skips_provider_socket() {
let listener = TcpListener::bind("127.0.0.1:0")
.await
.expect("test listener binds");
let addr = listener.local_addr().expect("listener has local addr");
let logger = Arc::new(RecordingOcrLogger::default());
let guardrail = Arc::new(RecordingOcrGuardrail::blocking_pre_call());
let err = ocr(OcrRequest {
model: "mistral-ocr-latest",
document: json!({
"type": "document_url",
"document_url": "https://example.com/doc.pdf"
}),
api_key: Some("sk-test"),
api_base: Some(&format!("http://{addr}")),
custom_llm_provider: Some("mistral"),
extra_headers: None,
optional_params: Map::new(),
timeout: Some(Duration::from_millis(100)),
callbacks: vec![logger.clone()],
guardrails: vec![guardrail.clone()],
request_metadata: RequestMetadata::default(),
litellm_call_id: Some("ocr-call-3"),
})
.await
.expect_err("guardrail blocks request");
assert!(matches!(err, Error::InvalidRequest(_)));
assert_eq!(guardrail.events(), vec!["async_pre_call_hook"]);
assert_eq!(
logger.events(),
vec![RecordedLogEvent {
hook: "async_log_failure_event",
model: "mistral-ocr-latest".to_string(),
call_type: "ocr".to_string(),
user_id: None,
response_object: Some("error".to_string()),
error_kind: Some("InvalidRequest".to_string()),
}]
);
let accepted = tokio::time::timeout(Duration::from_millis(100), listener.accept()).await;
assert!(accepted.is_err(), "provider socket should not be touched");
}
#[tokio::test]
async fn ocr_does_not_duplicate_authorization_header_when_header_is_supplied() {
let listener = TcpListener::bind("127.0.0.1:0")
.await
.expect("test listener binds");
let addr = listener.local_addr().expect("listener has local addr");
let server = tokio::spawn(async move {
let (mut socket, _) = listener.accept().await.expect("accepts one request");
let request = read_http_headers(&mut socket).await;
let response_body = r#"{"pages":[{"index":0,"markdown":"ok"}],"model":"mistral-ocr-latest","usage_info":{"pages_processed":1}}"#;
let response = format!(
"HTTP/1.1 200 OK\r\ncontent-type: application/json\r\ncontent-length: {}\r\nconnection: close\r\n\r\n{}",
response_body.len(),
response_body
);
socket
.write_all(response.as_bytes())
.await
.expect("writes response");
request
});
let mut headers = Map::new();
headers.insert(
"Authorization".to_string(),
Value::String("Bearer sk-from-python".to_string()),
);
headers.insert(
"x-trace-id".to_string(),
Value::String("trace-1".to_string()),
);
let response = ocr(OcrRequest {
model: "mistral-ocr-latest",
document: json!({
"type": "document_url",
"document_url": "https://example.com/doc.pdf"
}),
api_key: Some("sk-for-rust-fallback"),
api_base: Some(&format!("http://{addr}")),
custom_llm_provider: Some("mistral"),
extra_headers: Some(headers),
optional_params: Map::new(),
timeout: Some(Duration::from_secs(5)),
callbacks: Vec::new(),
guardrails: Vec::new(),
request_metadata: RequestMetadata::default(),
litellm_call_id: None,
})
.await
.expect("ocr request succeeds");
assert_eq!(response["pages"][0]["markdown"], "ok");
let request = server.await.expect("server task completes");
let authorization_count = request
.lines()
.filter(|line| line.to_ascii_lowercase().starts_with("authorization:"))
.count();
assert_eq!(authorization_count, 1, "{request}");
assert!(
request.contains("authorization: Bearer sk-from-python")
|| request.contains("Authorization: Bearer sk-from-python"),
"{request}"
);
}
#[tokio::test]
async fn document_intelligence_poll_uses_resolved_subscription_key() {
let listener = TcpListener::bind("127.0.0.1:0")
.await
.expect("test listener binds");
let addr = listener.local_addr().expect("listener has local addr");
let operation_url = format!("http://{addr}/operations/1");
let server = tokio::spawn(async move {
let (mut post_socket, _) = listener.accept().await.expect("accepts post request");
let post_request = read_http_headers(&mut post_socket).await;
let post_response = format!(
"HTTP/1.1 202 Accepted\r\noperation-location: {operation_url}\r\ncontent-length: 0\r\nconnection: close\r\n\r\n"
);
post_socket
.write_all(post_response.as_bytes())
.await
.expect("writes post response");
let (mut poll_socket, _) = listener.accept().await.expect("accepts poll request");
let poll_request = read_http_headers(&mut poll_socket).await;
let response_body = r#"{"status":"succeeded","analyzeResult":{"pages":[{"pageNumber":1,"lines":[{"content":"ok"}]}]}}"#;
let poll_response = format!(
"HTTP/1.1 200 OK\r\ncontent-type: application/json\r\ncontent-length: {}\r\nconnection: close\r\n\r\n{}",
response_body.len(),
response_body
);
poll_socket
.write_all(poll_response.as_bytes())
.await
.expect("writes poll response");
(post_request, poll_request)
});
let response = ocr(OcrRequest {
model: "doc-intelligence/prebuilt-read",
document: json!({
"type": "document_url",
"document_url": "https://example.com/doc.pdf"
}),
api_key: Some("di-key"),
api_base: Some(&format!("http://{addr}")),
custom_llm_provider: Some("azure_ai"),
extra_headers: None,
optional_params: Map::new(),
timeout: Some(Duration::from_secs(5)),
callbacks: Vec::new(),
guardrails: Vec::new(),
request_metadata: RequestMetadata::default(),
litellm_call_id: None,
})
.await
.expect("document intelligence request succeeds");
assert_eq!(response["pages"][0]["markdown"], "ok");
let (post_request, poll_request) = server.await.expect("server task completes");
assert!(
post_request
.to_ascii_lowercase()
.contains("ocp-apim-subscription-key: di-key"),
"{post_request}"
);
assert!(
poll_request
.to_ascii_lowercase()
.contains("ocp-apim-subscription-key: di-key"),
"{poll_request}"
);
}

View file

@ -12,17 +12,25 @@ path = "tests/workspace_crate_allowlist.rs"
[dependencies]
base64.workspace = true
azure_core.workspace = true
azure_identity.workspace = true
data-url = "0.3.2"
gcp_auth.workspace = true
moka.workspace = true
rand.workspace = true
reqwest.workspace = true
serde.workspace = true
serde_json.workspace = true
serde_path_to_error = "0.1"
strum.workspace = true
subtle.workspace = true
tokio.workspace = true
thiserror.workspace = true
tracing.workspace = true
tracing-subscriber = { workspace = true, optional = true }
sha2.workspace = true
url.workspace = true
veil.workspace = true
aws-config = { version = "1.9.0", default-features = false, features = ["rustls", "rt-tokio"], optional = true }
aws-credential-types = { version = "1.3.0", features = ["hardcoded-credentials"], optional = true }
aws-sdk-sts = { version = "1.108.0", default-features = false, features = ["rustls", "rt-tokio"], optional = true }
@ -44,5 +52,4 @@ observability = ["dep:tracing-subscriber"]
[dev-dependencies]
rstest.workspace = true
tokio = { workspace = true, features = ["macros", "rt-multi-thread"] }
tracing-subscriber.workspace = true

View file

@ -0,0 +1,168 @@
use std::future::Future;
use std::path::PathBuf;
use std::pin::Pin;
use std::sync::Arc;
use veil::Redact;
use crate::AuthError;
use super::{ResolvedCredential, SecretValue, TokenProviderHandle};
#[derive(Clone, Debug, PartialEq, Eq)]
pub enum CredentialFileRef {
Path(PathBuf),
EnvironmentVariable(String),
}
#[derive(Clone, Debug, PartialEq, Eq)]
pub enum CredentialRef {
Explicit(SecretValue),
Env(String),
File(CredentialFileRef),
Request(String),
Host(String),
None,
}
#[derive(Clone, Debug, PartialEq, Eq)]
pub enum CredentialLookup {
Found(SecretValue),
Missing,
Declined,
}
pub type CredentialLookupFuture<'a> =
Pin<Box<dyn Future<Output = Result<CredentialLookup, AuthError>> + Send + 'a>>;
pub trait CredentialResolver: std::fmt::Debug + Send + Sync {
fn resolve<'a>(&'a self, reference: &'a CredentialRef) -> CredentialLookupFuture<'a>;
}
#[derive(Clone, Redact)]
pub struct CredentialResolverHandle(#[redact(with = "[REDACTED]")] Arc<dyn CredentialResolver>);
impl CredentialResolverHandle {
pub fn new(resolver: Arc<dyn CredentialResolver>) -> Self {
Self(resolver)
}
pub async fn resolve(&self, reference: &CredentialRef) -> Result<CredentialLookup, AuthError> {
self.0.resolve(reference).await
}
}
#[derive(Clone, Debug)]
pub enum CredentialPlan {
Static(CredentialRef),
Caller(TokenProviderHandle),
None,
}
#[derive(Clone, Debug, PartialEq, Eq)]
pub enum CredentialPlanResolution {
Resolved(ResolvedCredential),
Unavailable,
}
impl CredentialPlan {
pub async fn resolve(
&self,
resolver: &CredentialResolverHandle,
) -> Result<CredentialPlanResolution, AuthError> {
match self {
Self::Static(CredentialRef::Explicit(secret)) => Ok(
CredentialPlanResolution::Resolved(ResolvedCredential::Static(secret.clone())),
),
Self::Static(CredentialRef::None) | Self::None => {
Ok(CredentialPlanResolution::Unavailable)
}
Self::Static(reference) => match resolver.resolve(reference).await? {
CredentialLookup::Found(secret) => Ok(CredentialPlanResolution::Resolved(
ResolvedCredential::Static(secret),
)),
CredentialLookup::Missing | CredentialLookup::Declined => {
Ok(CredentialPlanResolution::Unavailable)
}
},
Self::Caller(caller) => {
let credential = caller.acquire().await?;
if credential.secret().expose().is_empty() {
return Err(AuthError::EmptyCallerCredential);
}
Ok(CredentialPlanResolution::Resolved(credential))
}
}
}
}
#[cfg(test)]
mod tests {
use std::sync::Arc;
use super::{
CredentialLookup, CredentialLookupFuture, CredentialPlan, CredentialPlanResolution,
CredentialRef, CredentialResolver, CredentialResolverHandle,
};
use crate::AuthError;
use crate::auth::SecretValue;
#[derive(Debug)]
struct HostResolver;
impl CredentialResolver for HostResolver {
fn resolve<'a>(&'a self, reference: &'a CredentialRef) -> CredentialLookupFuture<'a> {
Box::pin(async move {
Ok(match reference {
CredentialRef::Host(name) if name == "rotating-token" => {
CredentialLookup::Found(SecretValue::new("resolved"))
}
_ => CredentialLookup::Declined,
})
})
}
}
#[tokio::test]
async fn static_host_reference_resolves_at_acquisition_time() {
let resolver = CredentialResolverHandle::new(Arc::new(HostResolver));
let plan = CredentialPlan::Static(CredentialRef::Host("rotating-token".to_string()));
let resolved = plan.resolve(&resolver).await.unwrap();
assert!(matches!(resolved, CredentialPlanResolution::Resolved(_)));
}
#[tokio::test]
async fn declined_reference_is_available_for_pre_acquisition_fallback() {
let resolver = CredentialResolverHandle::new(Arc::new(HostResolver));
let plan = CredentialPlan::Static(CredentialRef::Request("api-key".to_string()));
assert_eq!(
plan.resolve(&resolver).await.unwrap(),
CredentialPlanResolution::Unavailable
);
}
#[derive(Debug)]
struct FailingResolver;
impl CredentialResolver for FailingResolver {
fn resolve<'a>(&'a self, _reference: &'a CredentialRef) -> CredentialLookupFuture<'a> {
Box::pin(async { Err(AuthError::UnresolvedOidcReference) })
}
}
#[tokio::test]
async fn acquisition_failure_is_terminal() {
let resolver = CredentialResolverHandle::new(Arc::new(FailingResolver));
let plan = CredentialPlan::Static(CredentialRef::Host("token".to_string()));
let error = plan
.resolve(&resolver)
.await
.expect_err("acquisition errors cannot become fallback");
assert_eq!(error, AuthError::UnresolvedOidcReference);
}
}

View file

@ -0,0 +1,128 @@
use thiserror::Error;
#[derive(Clone, Debug, Error, PartialEq, Eq)]
pub enum AuthError {
#[error("invalid authentication configuration: {0}")]
Configuration(#[from] AuthConfigurationError),
#[error("credential acquisition failed: {0}")]
AzureTokenAcquisition(String),
#[error("credential acquisition failed: Vertex AI credentials: {0}")]
VertexTokenAcquisition(String),
#[error("credential acquisition failed: {}", .0.iter().map(ToString::to_string).collect::<Vec<_>>().join("; "))]
CredentialChain(Vec<AuthError>),
#[error("credential caller failed: credential caller returned an empty credential")]
EmptyCallerCredential,
#[error("credential caller failed: Azure AD token provider returned an empty token")]
EmptyAzureToken,
#[error("credential acquisition failed: Azure OIDC reference did not resolve to a value")]
UnresolvedOidcReference,
#[error(
"Missing {provider} API Key - A call is being made to {provider} but no key is set either in the environment variables or via params"
)]
MissingApiKey { provider: &'static str },
#[error(
"Missing {provider} API Base - Set {environment_variable} environment variable or pass api_base parameter"
)]
MissingApiBase {
provider: &'static str,
environment_variable: &'static str,
},
#[error("{0}")]
MissingCredential(#[from] MissingCredential),
#[error("{0}")]
Aws(#[from] AwsAuthError),
#[error("invalid authentication header")]
InvalidHeader,
}
#[derive(Clone, Debug, Error, PartialEq, Eq)]
pub enum AuthConfigurationError {
#[error("credential header already exists")]
ExistingCredentialHeader,
#[error("credential plan is not allowed by the provider auth policy")]
DisallowedCredentialPlan,
#[error("credential cannot be empty")]
EmptyCredential,
#[error("invalid Azure credential selector")]
InvalidAzureSelector,
#[error("ClientSecretCredential requires tenant_id, client_id, and client_secret")]
MissingClientSecretFields,
#[error("WorkloadIdentityCredential requires tenant_id")]
MissingWorkloadTenant,
#[error("WorkloadIdentityCredential requires client_id")]
MissingWorkloadClient,
#[error("WorkloadIdentityCredential requires azure_federated_token_file")]
MissingWorkloadTokenFile,
#[error("credential reference requires a host credential resolver")]
MissingHostResolver,
#[error("caller credential plan requires provider-specific inputs")]
MissingCallerInputs,
#[error("credential header {0} already exists")]
DuplicateHeader(&'static str),
#[error("{0} must be a string or null")]
InvalidFieldType(String),
#[error("unsupported OIDC reference")]
UnsupportedOidcReference,
#[error("{0} cannot be empty")]
EmptyReference(String),
#[error("Azure credential initialization failed: {0}")]
AzureCredentialInitialization(String),
#[error("Azure authority must be an HTTPS origin without credentials, query, or fragment")]
InvalidAzureAuthority,
#[error("request-controlled Azure auth inputs cannot be combined with host credentials")]
MixedAzureCredentialSources,
#[error("request-controlled Azure credential references are not allowed")]
RequestAzureCredentialReference,
#[error("host credentials cannot be sent to a request-controlled Azure endpoint")]
RequestAzureCredentialDestination,
#[error("credentials cannot be sent to a request-controlled Vertex AI endpoint")]
RequestVertexCredentialDestination,
#[error(
"request-controlled Vertex credentials must use the canonical Google OAuth token endpoint"
)]
RequestVertexTokenEndpoint,
}
#[derive(Clone, Debug, Error, PartialEq, Eq)]
pub enum MissingCredential {
#[error(
"Missing Anthropic API Key - Set `api_key` or the ANTHROPIC_API_KEY environment variable"
)]
AnthropicApiKey,
#[error("Missing Azure API Key - Set `api_key` or the AZURE_API_KEY environment variable")]
AzureApiKey,
#[error(
"Missing Azure API Base - Set `api_base` or the AZURE_API_BASE environment variable. Expected format: https://<resource-name>.services.ai.azure.com/anthropic"
)]
AzureApiBase,
#[error(
"Missing OpenAI API Key - a realtime call is being made but no key was passed via params or the OPENAI_API_KEY environment variable"
)]
OpenAiRealtimeApiKey,
#[error(
"Missing OpenAI API Key - a Responses WebSocket call is being made but no key was passed via params or the OPENAI_API_KEY environment variable"
)]
OpenAiResponsesApiKey,
}
#[derive(Clone, Debug, Error, PartialEq, Eq)]
pub enum AwsAuthError {
#[error("AWS profile credentials failed: {0}")]
Profile(String),
#[error("AWS default credentials failed: {0}")]
DefaultChain(String),
#[error("AWS role credentials failed: {0}")]
AssumeRole(String),
#[error("AWS web identity credentials failed: {0}")]
WebIdentity(String),
#[error("AWS web identity expiration was invalid: {0}")]
WebIdentityExpiration(String),
#[error("AWS signing parameters failed: {0}")]
SigningParameters(String),
#[error("AWS signable request failed: {0}")]
SignableRequest(String),
#[error("AWS request signing failed: {0}")]
Signing(String),
#[error("AWS web identity response had no credentials")]
MissingWebIdentityCredentials,
}

View file

@ -0,0 +1,86 @@
use crate::AuthError;
use crate::auth::error::AuthConfigurationError;
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum CredentialPlacement {
Bearer,
Header(&'static str),
}
impl CredentialPlacement {
pub fn header_name(self) -> &'static str {
match self {
Self::Bearer => "Authorization",
Self::Header(name) => name,
}
}
}
pub(crate) fn apply_credential(
headers: Vec<(String, String)>,
credential: &str,
placement: CredentialPlacement,
) -> Result<Vec<(String, String)>, AuthError> {
if credential.trim().is_empty() {
return Err(AuthError::Configuration(
AuthConfigurationError::EmptyCredential,
));
}
if headers
.iter()
.any(|(name, _)| name.eq_ignore_ascii_case(placement.header_name()))
{
return Err(AuthError::Configuration(
AuthConfigurationError::DuplicateHeader(placement.header_name()),
));
}
let value = match placement {
CredentialPlacement::Bearer => format!("Bearer {credential}"),
CredentialPlacement::Header(_) => credential.to_string(),
};
Ok(
std::iter::once((placement.header_name().to_string(), value))
.chain(headers)
.collect(),
)
}
/// How the upstream call is authenticated. API-key strategies are resolved in
/// `prepare`; SigV4 needs the serialized body, so the handler signs it.
#[derive(Clone, Debug, PartialEq, Eq)]
pub enum RequestAuth {
Header { name: &'static str, value: String },
Bearer { token: String },
AwsSigV4 { region: String },
}
#[cfg(test)]
mod tests {
use super::{CredentialPlacement, apply_credential};
#[test]
fn bearer_uses_authorization_header() {
let headers = apply_credential(Vec::new(), "key", CredentialPlacement::Bearer)
.expect("credential applies");
assert_eq!(
headers,
vec![("Authorization".to_string(), "Bearer key".to_string())]
);
}
#[test]
fn named_header_rejects_existing_value() {
let error = apply_credential(
vec![(
"ocp-apim-subscription-key".to_string(),
"caller-key".to_string(),
)],
"configured-key",
CredentialPlacement::Header("Ocp-Apim-Subscription-Key"),
)
.expect_err("provider policy must handle existing credentials");
assert!(error.to_string().contains("already exists"));
}
}

View file

@ -0,0 +1,56 @@
mod credential;
pub mod error;
pub(crate) mod vertex;
pub use error::AuthError;
pub(crate) mod http;
mod policy;
mod secret;
mod token;
use serde::{Deserialize, Serialize};
#[derive(Clone, Copy, Debug, Default, Deserialize, Eq, Hash, PartialEq, Serialize)]
#[serde(rename_all = "snake_case")]
pub enum InputSource {
Request,
#[default]
Deployment,
Environment,
}
#[derive(Clone, Debug, Default, Eq, PartialEq)]
pub struct Sourced<T> {
value: T,
source: InputSource,
}
impl<T> Sourced<T> {
pub fn new(value: T, source: InputSource) -> Self {
Self { value, source }
}
pub fn value(&self) -> &T {
&self.value
}
pub fn source(&self) -> InputSource {
self.source
}
pub fn into_value(self) -> T {
self.value
}
pub fn map<U>(self, map: impl FnOnce(T) -> U) -> Sourced<U> {
Sourced::new(map(self.value), self.source)
}
}
pub use credential::{
CredentialFileRef, CredentialLookup, CredentialLookupFuture, CredentialPlan,
CredentialPlanResolution, CredentialRef, CredentialResolver, CredentialResolverHandle,
};
pub use http::{CredentialPlacement, RequestAuth};
pub use policy::{CredentialPlanKind, CredentialRule, ExistingHeaderBehavior, ProviderAuthPolicy};
pub use secret::SecretValue;
pub use token::{ResolvedCredential, TokenFuture, TokenProvider, TokenProviderHandle};

View file

@ -0,0 +1,114 @@
use crate::AuthError;
use crate::auth::error::AuthConfigurationError;
use super::http::apply_credential;
use super::{CredentialPlacement, ResolvedCredential};
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum CredentialPlanKind {
Static,
Entra,
Caller,
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct CredentialRule {
pub kind: CredentialPlanKind,
pub placement: CredentialPlacement,
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum ExistingHeaderBehavior {
Preserve,
Reject,
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct ProviderAuthPolicy {
pub rules: &'static [CredentialRule],
pub accepted_existing_headers: &'static [&'static str],
pub existing_header_behavior: ExistingHeaderBehavior,
pub scope: Option<&'static str>,
pub audience: Option<&'static str>,
}
impl ProviderAuthPolicy {
pub fn has_existing_credential(&self, headers: &[(String, String)]) -> bool {
headers.iter().any(|(name, _)| {
self.accepted_existing_headers
.iter()
.any(|accepted| name.eq_ignore_ascii_case(accepted))
})
}
pub fn apply(
&self,
headers: Vec<(String, String)>,
kind: CredentialPlanKind,
credential: &ResolvedCredential,
) -> Result<Vec<(String, String)>, AuthError> {
if self.has_existing_credential(&headers) {
return match self.existing_header_behavior {
ExistingHeaderBehavior::Preserve => Ok(headers),
ExistingHeaderBehavior::Reject => Err(AuthError::Configuration(
AuthConfigurationError::ExistingCredentialHeader,
)),
};
}
let rule =
self.rules
.iter()
.find(|rule| rule.kind == kind)
.ok_or(AuthError::Configuration(
AuthConfigurationError::DisallowedCredentialPlan,
))?;
apply_credential(headers, credential.secret().expose(), rule.placement)
}
}
#[cfg(test)]
mod tests {
use super::{CredentialPlanKind, CredentialRule, ExistingHeaderBehavior, ProviderAuthPolicy};
use crate::auth::{CredentialPlacement, ResolvedCredential, SecretValue};
const RULES: &[CredentialRule] = &[CredentialRule {
kind: CredentialPlanKind::Static,
placement: CredentialPlacement::Header("x-api-key"),
}];
const POLICY: ProviderAuthPolicy = ProviderAuthPolicy {
rules: RULES,
accepted_existing_headers: &["x-api-key"],
existing_header_behavior: ExistingHeaderBehavior::Preserve,
scope: None,
audience: None,
};
#[test]
fn rules_define_allowed_plans_and_credential_placement() {
let headers = POLICY
.apply(
Vec::new(),
CredentialPlanKind::Static,
&ResolvedCredential::Static(SecretValue::new("secret")),
)
.unwrap();
assert_eq!(
headers,
vec![("x-api-key".to_string(), "secret".to_string())]
);
}
#[test]
fn unsupported_plan_is_rejected() {
let error = POLICY
.apply(
Vec::new(),
CredentialPlanKind::Entra,
&ResolvedCredential::Static(SecretValue::new("secret")),
)
.unwrap_err();
assert!(error.to_string().contains("not allowed"));
}
}

View file

@ -0,0 +1,41 @@
use veil::Redact;
#[derive(Redact, Clone)]
pub struct SecretValue(#[redact(with = "[REDACTED]")] String);
impl SecretValue {
pub fn new(value: impl Into<String>) -> Self {
Self(value.into())
}
pub fn expose(&self) -> &str {
&self.0
}
}
impl PartialEq for SecretValue {
fn eq(&self, other: &Self) -> bool {
subtle::ConstantTimeEq::ct_eq(self.0.as_bytes(), other.0.as_bytes()).into()
}
}
impl Eq for SecretValue {}
#[cfg(test)]
mod tests {
use super::SecretValue;
#[test]
fn debug_redacts_plaintext() {
let debug = format!("{:?}", SecretValue::new("credential-value"));
assert!(!debug.contains("credential-value"));
assert!(debug.contains("REDACTED"));
}
#[test]
fn equality_compares_plaintext_values() {
assert_eq!(SecretValue::new("same"), SecretValue::new("same"));
assert_ne!(SecretValue::new("same"), SecretValue::new("different"));
}
}

View file

@ -0,0 +1,47 @@
use std::future::Future;
use std::pin::Pin;
use std::sync::Arc;
use std::time::SystemTime;
use veil::Redact;
use crate::AuthError;
use super::secret::SecretValue;
#[derive(Clone, Debug, PartialEq, Eq)]
pub enum ResolvedCredential {
Static(SecretValue),
AccessToken {
token: SecretValue,
expires_on: Option<SystemTime>,
},
}
impl ResolvedCredential {
pub fn secret(&self) -> &SecretValue {
match self {
Self::Static(secret) | Self::AccessToken { token: secret, .. } => secret,
}
}
}
pub type TokenFuture<'a> =
Pin<Box<dyn Future<Output = Result<ResolvedCredential, AuthError>> + Send + 'a>>;
pub trait TokenProvider: std::fmt::Debug + Send + Sync {
fn acquire(&self) -> TokenFuture<'_>;
}
#[derive(Clone, Redact)]
pub struct TokenProviderHandle(#[redact(with = "[REDACTED]")] Arc<dyn TokenProvider>);
impl TokenProviderHandle {
pub fn new(caller: Arc<dyn TokenProvider>) -> Self {
Self(caller)
}
pub async fn acquire(&self) -> Result<ResolvedCredential, AuthError> {
self.0.acquire().await
}
}

View file

@ -0,0 +1,592 @@
use std::collections::BTreeMap;
use std::future::Future;
use std::path::Path;
use std::pin::Pin;
use std::sync::Arc;
use gcp_auth::{CustomServiceAccount, TokenProvider};
use moka::future::Cache;
use serde_json::{Map, Value};
use sha2::{Digest, Sha256};
use crate::auth::error::AuthConfigurationError;
use crate::auth::http::apply_credential;
use crate::auth::{AuthError, CredentialPlacement, InputSource, SecretValue, Sourced};
const CLOUD_PLATFORM_SCOPE: &str = "https://www.googleapis.com/auth/cloud-platform";
const GOOGLE_OAUTH_TOKEN_ENDPOINT: &str = "https://oauth2.googleapis.com/token";
const GOOGLE_APPLICATION_CREDENTIALS_ENV: &str = "GOOGLE_APPLICATION_CREDENTIALS";
const VERTEX_AI_API_KEY_ENV: &str = "VERTEX_AI_API_KEY";
const VERTEXAI_API_KEY_ENV: &str = "VERTEXAI_API_KEY";
const VERTEXAI_CREDENTIALS_ENV: &str = "VERTEXAI_CREDENTIALS";
const VERTEXAI_PROJECT_ENV: &str = "VERTEXAI_PROJECT";
const VERTEXAI_LOCATION_ENV: &str = "VERTEXAI_LOCATION";
const VERTEX_LOCATION_ENV: &str = "VERTEX_LOCATION";
#[derive(Clone, Debug, Default)]
pub(crate) struct VertexConfig {
credentials: Option<Sourced<SecretValue>>,
project_id: Option<String>,
location: Option<String>,
}
impl VertexConfig {
pub(crate) fn from_sourced_optional_params(
params: &Map<String, Value>,
sources: &BTreeMap<String, InputSource>,
) -> Result<Self, AuthError> {
Ok(Self {
credentials: optional_credentials(
params,
sources,
&["vertex_credentials", "vertex_ai_credentials"],
)?,
project_id: optional_string(params, &["vertex_project", "vertex_ai_project"])?,
location: optional_string(params, &["vertex_location", "vertex_ai_location"])?,
})
}
pub(crate) fn project_id(&self) -> Option<&str> {
self.project_id.as_deref()
}
pub(crate) fn location(&self) -> Option<&str> {
self.location.as_deref()
}
}
pub(crate) struct VertexEnvironment {
pub headers: Vec<(String, String)>,
pub project_id: String,
}
struct VertexAccessToken {
token: String,
project_id: String,
}
pub(crate) fn get_vertex_ai_project(
config: &VertexConfig,
env_lookup: &dyn Fn(&str) -> Option<String>,
) -> Option<String> {
config
.project_id()
.map(str::to_string)
.or_else(|| non_empty_env(env_lookup, VERTEXAI_PROJECT_ENV))
}
pub(crate) fn get_vertex_ai_location(
config: &VertexConfig,
env_lookup: &dyn Fn(&str) -> Option<String>,
) -> Option<String> {
config
.location()
.map(str::to_string)
.or_else(|| non_empty_env(env_lookup, VERTEXAI_LOCATION_ENV))
.or_else(|| non_empty_env(env_lookup, VERTEX_LOCATION_ENV))
}
#[derive(Clone)]
pub(crate) struct VertexAuth {
providers: Cache<CredentialCacheKey, Arc<dyn VertexTokenSource>>,
loader: Arc<dyn VertexProviderLoader>,
}
impl Default for VertexAuth {
fn default() -> Self {
Self::new(Arc::new(GcpProviderLoader))
}
}
impl VertexAuth {
fn new(loader: Arc<dyn VertexProviderLoader>) -> Self {
Self {
providers: Cache::builder().max_capacity(64).build(),
loader,
}
}
#[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)]
pub(crate) async fn validate_environment(
&self,
headers: Vec<(String, String)>,
api_key: Option<&str>,
config: &VertexConfig,
env_lookup: &(dyn Fn(&str) -> Option<String> + Sync),
) -> Result<VertexEnvironment, AuthError> {
let has_authorization = headers
.iter()
.any(|(name, _)| name.eq_ignore_ascii_case("Authorization"));
let static_token = api_key
.map(str::trim)
.filter(|value| !value.is_empty())
.map(str::to_string)
.or_else(|| non_empty_env(env_lookup, VERTEX_AI_API_KEY_ENV))
.or_else(|| non_empty_env(env_lookup, VERTEXAI_API_KEY_ENV));
let project_id = get_vertex_ai_project(config, env_lookup);
if !has_authorization && static_token.is_none() {
let access = self.get_access_token(config, env_lookup).await?;
return Ok(VertexEnvironment {
headers: apply_credential(headers, &access.token, CredentialPlacement::Bearer)?,
project_id: project_id.unwrap_or(access.project_id),
});
}
let project_id = match project_id {
Some(project_id) => project_id,
None => {
self.load_provider(config, env_lookup)
.await?
.project_id()
.await?
}
};
let headers = if has_authorization {
headers
} else {
apply_credential(
headers,
static_token.as_deref().expect("static token was checked"),
CredentialPlacement::Bearer,
)?
};
Ok(VertexEnvironment {
headers,
project_id,
})
}
async fn get_access_token(
&self,
config: &VertexConfig,
env_lookup: &(dyn Fn(&str) -> Option<String> + Sync),
) -> Result<VertexAccessToken, AuthError> {
let provider = self.load_provider(config, env_lookup).await?;
let (token, project_id) = tokio::try_join!(provider.token(), provider.project_id())?;
Ok(VertexAccessToken { token, project_id })
}
async fn load_provider(
&self,
config: &VertexConfig,
env_lookup: &(dyn Fn(&str) -> Option<String> + Sync),
) -> Result<Arc<dyn VertexTokenSource>, AuthError> {
let source = credential_source(config, env_lookup);
let key = source.cache_key();
self.providers
.try_get_with(key, self.loader.load(source))
.await
.map_err(|error| (*error).clone())
}
}
trait VertexTokenSource: Send + Sync {
fn project_id(&self) -> VertexAuthFuture<'_, String>;
fn token(&self) -> VertexAuthFuture<'_, String>;
}
trait VertexProviderLoader: Send + Sync {
fn load(&self, source: CredentialSource) -> VertexAuthFuture<'_, Arc<dyn VertexTokenSource>>;
}
type VertexAuthFuture<'a, T> = Pin<Box<dyn Future<Output = Result<T, AuthError>> + Send + 'a>>;
struct GcpTokenSource(Arc<dyn TokenProvider>);
impl VertexTokenSource for GcpTokenSource {
fn project_id(&self) -> VertexAuthFuture<'_, String> {
Box::pin(async move {
self.0
.project_id()
.await
.map(|project| project.to_string())
.map_err(auth_acquisition_error)
})
}
fn token(&self) -> VertexAuthFuture<'_, String> {
Box::pin(async move {
self.0
.token(&[CLOUD_PLATFORM_SCOPE])
.await
.map(|token| token.as_str().to_string())
.map_err(auth_acquisition_error)
})
}
}
struct GcpProviderLoader;
impl VertexProviderLoader for GcpProviderLoader {
fn load(&self, source: CredentialSource) -> VertexAuthFuture<'_, Arc<dyn VertexTokenSource>> {
Box::pin(async move {
let provider: Arc<dyn TokenProvider> = match source {
CredentialSource::Inline(configured) => Arc::new(
CustomServiceAccount::from_json(validate_request_credentials(
configured.expose(),
)?)
.map_err(auth_acquisition_error)?,
),
CredentialSource::Trusted(configured) => {
let configured = configured.expose();
let service_account = if Path::new(configured).is_file() {
CustomServiceAccount::from_file(configured)
} else {
CustomServiceAccount::from_json(configured)
}
.map_err(auth_acquisition_error)?;
Arc::new(service_account)
}
CredentialSource::ApplicationCredentials(path) => {
Arc::new(CustomServiceAccount::from_file(path).map_err(auth_acquisition_error)?)
}
CredentialSource::Adc => {
gcp_auth::provider().await.map_err(auth_acquisition_error)?
}
};
Ok(Arc::new(GcpTokenSource(provider)) as Arc<dyn VertexTokenSource>)
})
}
}
fn validate_request_credentials(configured: &str) -> Result<&str, AuthError> {
let token_uri = serde_json::from_str::<Value>(configured)
.ok()
.and_then(|credentials| {
credentials
.get("token_uri")
.and_then(Value::as_str)
.map(str::to_string)
});
if token_uri.as_deref() != Some(GOOGLE_OAUTH_TOKEN_ENDPOINT) {
return Err(AuthConfigurationError::RequestVertexTokenEndpoint.into());
}
Ok(configured)
}
#[derive(Clone, Debug)]
enum CredentialSource {
Inline(SecretValue),
Trusted(SecretValue),
ApplicationCredentials(String),
Adc,
}
impl CredentialSource {
fn cache_key(&self) -> CredentialCacheKey {
match self {
Self::Inline(configured) => {
CredentialCacheKey::Inline(Sha256::digest(configured.expose()).into())
}
Self::Trusted(configured) => {
CredentialCacheKey::Trusted(Sha256::digest(configured.expose()).into())
}
Self::ApplicationCredentials(path) => {
CredentialCacheKey::ApplicationCredentials(path.clone())
}
Self::Adc => CredentialCacheKey::Adc,
}
}
}
#[derive(Clone, Debug, Hash, PartialEq, Eq)]
enum CredentialCacheKey {
Inline([u8; 32]),
Trusted([u8; 32]),
ApplicationCredentials(String),
Adc,
}
fn credential_source(
config: &VertexConfig,
env_lookup: &dyn Fn(&str) -> Option<String>,
) -> CredentialSource {
if let Some(configured) = config.credentials.clone() {
return match configured.source() {
InputSource::Request => CredentialSource::Inline(configured.into_value()),
InputSource::Deployment | InputSource::Environment => {
CredentialSource::Trusted(configured.into_value())
}
};
}
if let Some(configured) = non_empty_env(env_lookup, VERTEXAI_CREDENTIALS_ENV) {
return CredentialSource::Trusted(SecretValue::new(configured));
}
non_empty_env(env_lookup, GOOGLE_APPLICATION_CREDENTIALS_ENV)
.map(CredentialSource::ApplicationCredentials)
.unwrap_or(CredentialSource::Adc)
}
fn optional_credentials(
params: &Map<String, Value>,
sources: &BTreeMap<String, InputSource>,
names: &[&str],
) -> Result<Option<Sourced<SecretValue>>, AuthError> {
for name in names {
let source = source_for(sources, name);
match params.get(*name) {
None | Some(Value::Null) => continue,
Some(Value::String(value)) if value.trim().is_empty() => continue,
Some(Value::String(value)) => {
return Ok(Some(Sourced::new(SecretValue::new(value), source)));
}
Some(Value::Object(value)) if value.is_empty() => continue,
Some(Value::Object(value)) => {
return serde_json::to_string(value)
.map(SecretValue::new)
.map(|value| Sourced::new(value, source))
.map(Some)
.map_err(|error| {
AuthError::Configuration(AuthConfigurationError::InvalidFieldType(format!(
"{}: {error}",
names[0]
)))
});
}
Some(_) => {
return Err(AuthError::Configuration(
AuthConfigurationError::InvalidFieldType(names[0].to_string()),
));
}
}
}
Ok(None)
}
fn source_for(sources: &BTreeMap<String, InputSource>, name: &str) -> InputSource {
sources.get(name).copied().unwrap_or_default()
}
fn optional_string(
params: &Map<String, Value>,
names: &[&str],
) -> Result<Option<String>, AuthError> {
for name in names {
match params.get(*name) {
None | Some(Value::Null) => continue,
Some(Value::String(value)) if value.trim().is_empty() => continue,
Some(Value::String(value)) => return Ok(Some(value.clone())),
Some(_) => {
return Err(AuthError::Configuration(
AuthConfigurationError::InvalidFieldType(names[0].to_string()),
));
}
}
}
Ok(None)
}
fn non_empty_env(env_lookup: &dyn Fn(&str) -> Option<String>, name: &str) -> Option<String> {
env_lookup(name)
.map(|value| value.trim().to_string())
.filter(|value| !value.is_empty())
}
fn auth_acquisition_error(error: gcp_auth::Error) -> AuthError {
AuthError::VertexTokenAcquisition(error.to_string())
}
#[cfg(test)]
mod tests {
use std::sync::atomic::{AtomicUsize, Ordering};
use serde_json::json;
use super::*;
struct FakeProvider {
calls: Arc<AtomicUsize>,
}
impl VertexTokenSource for FakeProvider {
fn project_id(&self) -> VertexAuthFuture<'_, String> {
self.calls.fetch_add(1, Ordering::SeqCst);
Box::pin(async { Ok("adc-project".into()) })
}
fn token(&self) -> VertexAuthFuture<'_, String> {
self.calls.fetch_add(1, Ordering::SeqCst);
Box::pin(async { Ok("adc-token".into()) })
}
}
struct FakeLoader {
loads: Arc<AtomicUsize>,
provider: Arc<dyn VertexTokenSource>,
}
impl VertexProviderLoader for FakeLoader {
fn load(
&self,
_source: CredentialSource,
) -> VertexAuthFuture<'_, Arc<dyn VertexTokenSource>> {
let loads = self.loads.clone();
let provider = self.provider.clone();
Box::pin(async move {
loads.fetch_add(1, Ordering::SeqCst);
Ok(provider)
})
}
}
fn config(value: Value) -> VertexConfig {
VertexConfig::from_sourced_optional_params(value.as_object().unwrap(), &BTreeMap::new())
.unwrap()
}
fn auth(calls: Arc<AtomicUsize>, loads: Arc<AtomicUsize>) -> VertexAuth {
let provider: Arc<dyn VertexTokenSource> = Arc::new(FakeProvider { calls });
VertexAuth::new(Arc::new(FakeLoader { loads, provider }))
}
#[test]
fn config_is_typed_and_secrets_are_redacted() {
let config = config(json!({
"vertex_credentials":{"private_key":"secret-key"},
"vertex_project":"project-1",
"vertex_location":"europe-west4"
}));
assert_eq!(config.project_id(), Some("project-1"));
assert_eq!(config.location(), Some("europe-west4"));
assert!(!format!("{config:?}").contains("secret-key"));
assert!(
VertexConfig::from_sourced_optional_params(
json!({"vertex_credentials":true}).as_object().unwrap(),
&BTreeMap::new()
)
.is_err()
);
}
#[test]
fn empty_primary_values_fall_back_to_python_aliases() {
let config = config(json!({
"vertex_credentials": null,
"vertex_ai_credentials": "alias-credentials",
"vertex_project": " ",
"vertex_ai_project": "alias-project",
"vertex_location": null,
"vertex_ai_location": "alias-location"
}));
assert_eq!(
config.credentials.as_ref().unwrap().value().expose(),
"alias-credentials"
);
assert_eq!(config.project_id(), Some("alias-project"));
assert_eq!(config.location(), Some("alias-location"));
}
#[test]
fn project_and_location_prefer_input_then_environment() {
let configured =
config(json!({"vertex_project":"input-project","vertex_location":"input-location"}));
let env = |name: &str| Some(format!("env-{name}"));
assert_eq!(
get_vertex_ai_project(&configured, &env).as_deref(),
Some("input-project")
);
assert_eq!(
get_vertex_ai_location(&configured, &env).as_deref(),
Some("input-location")
);
let empty = VertexConfig::default();
assert_eq!(
get_vertex_ai_project(&empty, &|_| Some("env-project".into())).as_deref(),
Some("env-project")
);
assert_eq!(
get_vertex_ai_location(&empty, &|name| (name == VERTEX_LOCATION_ENV)
.then(|| "fallback-location".into()))
.as_deref(),
Some("fallback-location")
);
}
#[test]
fn credential_discovery_prefers_input_then_environment_then_adc() {
let params = json!({"vertex_credentials":"input-json"});
let sources = BTreeMap::from([("vertex_credentials".to_string(), InputSource::Request)]);
let configured =
VertexConfig::from_sourced_optional_params(params.as_object().unwrap(), &sources)
.unwrap();
assert!(
matches!(credential_source(&configured, &|_| Some("environment-value".into())), CredentialSource::Inline(value) if value.expose() == "input-json")
);
let empty = VertexConfig::default();
assert!(
matches!(credential_source(&empty, &|name| (name == VERTEXAI_CREDENTIALS_ENV).then(|| "environment-json".into())), CredentialSource::Trusted(value) if value.expose() == "environment-json")
);
assert!(
matches!(credential_source(&empty, &|name| (name == GOOGLE_APPLICATION_CREDENTIALS_ENV).then(|| "adc.json".into())), CredentialSource::ApplicationCredentials(path) if path == "adc.json")
);
assert!(matches!(
credential_source(&empty, &|_| None),
CredentialSource::Adc
));
assert_ne!(
CredentialSource::Inline(SecretValue::new("same-value")).cache_key(),
CredentialSource::Trusted(SecretValue::new("same-value")).cache_key()
);
}
#[test]
fn request_credentials_require_canonical_token_endpoint() {
assert!(
validate_request_credentials(r#"{"token_uri":"https://oauth2.googleapis.com/token"}"#)
.is_ok()
);
assert!(matches!(
validate_request_credentials(r#"{"token_uri":"http://127.0.0.1/token"}"#),
Err(AuthError::Configuration(
AuthConfigurationError::RequestVertexTokenEndpoint
))
));
assert!(matches!(
validate_request_credentials("{}"),
Err(AuthError::Configuration(
AuthConfigurationError::RequestVertexTokenEndpoint
))
));
}
#[tokio::test]
async fn explicit_token_and_header_do_not_acquire_adc() {
let loads = Arc::new(AtomicUsize::new(0));
let auth = auth(Arc::new(AtomicUsize::new(0)), loads.clone());
let configured = config(json!({"vertex_project":"project-1"}));
let explicit = auth
.validate_environment(Vec::new(), Some("access-token"), &configured, &|_| None)
.await
.unwrap();
assert_eq!(explicit.headers[0].1, "Bearer access-token");
let existing = auth
.validate_environment(
vec![("authorization".into(), "Bearer existing".into())],
None,
&configured,
&|_| None,
)
.await
.unwrap();
assert_eq!(existing.headers[0].1, "Bearer existing");
assert_eq!(loads.load(Ordering::SeqCst), 0);
}
#[tokio::test]
async fn provider_is_reused_across_authentication_calls() {
let calls = Arc::new(AtomicUsize::new(0));
let loads = Arc::new(AtomicUsize::new(0));
let auth = auth(calls.clone(), loads.clone());
for _ in 0..2 {
let environment = auth
.validate_environment(Vec::new(), None, &VertexConfig::default(), &|_| None)
.await
.unwrap();
assert_eq!(environment.project_id, "adc-project");
assert_eq!(environment.headers[0].1, "Bearer adc-token");
}
assert_eq!(loads.load(Ordering::SeqCst), 1);
assert_eq!(calls.load(Ordering::SeqCst), 4);
}
}

View file

@ -43,6 +43,23 @@ pub const EMPTY_TEXT_PLACEHOLDER: &str =
"[System: Empty message content sanitised to satisfy protocol]";
pub const FUNCTION_TRACE_TARGET: &str = "litellm::function_trace";
pub(crate) const MEDIA_CONNECT_TIMEOUT_SECS: u64 = 10;
pub(crate) const OCR_HTTP_TIMEOUT_SECS: u64 = 600;
pub(crate) const OCR_CONNECT_TIMEOUT_SECS: u64 = 10;
pub(crate) const OCR_INLINE_MAX_BYTES: usize = 50 * 1024 * 1024;
pub(crate) const OCR_DOWNLOAD_MAX_BYTES: u64 = 50 * 1024 * 1024;
pub(crate) const OCR_MAX_FETCH_REDIRECTS: usize = 10;
pub(crate) const OCR_POLL_TIMEOUT_SECS: u64 = 120;
pub(crate) const OCR_POLL_RETRY_SECS: u64 = 2;
pub(crate) const AZURE_DI_API_VERSION: &str = "2024-11-30";
pub(crate) const AZURE_DI_SUBSCRIPTION_HEADER: &str = "Ocp-Apim-Subscription-Key";
pub(crate) const AZURE_DI_DEFAULT_DPI: i64 = 96;
pub(crate) const AZURE_DI_DEFAULT_WIDTH: f64 = 8.5;
pub(crate) const AZURE_DI_DEFAULT_HEIGHT: f64 = 11.0;
pub(crate) const REDUCTO_API_BASE: &str = "https://platform.reducto.ai";
pub(crate) const REDUCTO_API_KEY_ENV: &str = "REDUCTO_API_KEY";
pub(crate) const REDUCTO_ID_PREFIX: &str = "reducto://";
pub(crate) const AZURE_AI_OCR_PATH: &str = "/providers/mistral/azure/ocr";
pub(crate) const MISTRAL_OCR_API_BASE: &str = "https://api.mistral.ai/v1";

View file

@ -21,6 +21,18 @@ pub enum Error {
"Missing {provider} API Key - A call is being made to {provider} but no key is set either in the environment variables or via params"
)]
MissingApiKey { provider: &'static str },
#[error(
"invalid authentication configuration: Missing Azure AI credentials - set AZURE_AI_API_KEY or configure Entra ID"
)]
MissingAzureAiCredentials,
#[error(
"invalid authentication configuration: Missing Azure Document Intelligence credentials - set AZURE_DOCUMENT_INTELLIGENCE_API_KEY or configure Entra ID"
)]
MissingAzureDocumentIntelligenceCredentials,
#[error(
"Missing REDUCTO_API_KEY - set it in the environment or pass api_key to litellm.ocr()/litellm.aocr()"
)]
MissingReductoApiKey,
#[error("upstream request failed with status {status}: {body}")]
Http { status: u16, body: String },
#[error("upstream network error: {0}")]
@ -40,6 +52,28 @@ pub enum Error {
Unsupported(&'static str),
}
#[derive(Debug, ThisError)]
pub(crate) enum MediaError {
#[error("media URL rejected by network policy")]
BlockedUrl,
#[error("media download is disabled")]
DownloadDisabled,
#[error("media download exceeds the maximum size")]
DownloadTooLarge,
#[error("too many redirects while fetching media")]
TooManyRedirects,
#[error("media redirect is missing a Location header")]
MissingRedirectLocation,
#[error("invalid media redirect")]
InvalidRedirect,
#[error("media download failed with status {0}")]
Http(u16),
#[error("media download timed out")]
Timeout,
#[error("{0}")]
Transport(#[from] TransportError),
}
#[derive(Clone, Debug, ThisError, PartialEq, Eq)]
pub enum TransportError {
#[error("upstream request failed with status {status}: {body}")]
@ -93,6 +127,15 @@ impl From<TransportError> for Error {
}
}
impl From<crate::AuthError> for Error {
fn from(error: crate::AuthError) -> Self {
match error {
crate::AuthError::MissingApiKey { provider } => Self::MissingApiKey { provider },
error => Self::Auth(error.to_string()),
}
}
}
pub fn json_type_name(value: &serde_json::Value) -> &'static str {
match value {
serde_json::Value::Null => "null",
@ -108,6 +151,14 @@ pub fn json_type_name(value: &serde_json::Value) -> &'static str {
mod transport_tests {
use super::*;
#[test]
fn missing_auth_key_preserves_provider_in_public_error() {
assert_eq!(
Error::from(crate::AuthError::MissingApiKey { provider: "Vertex" }),
Error::MissingApiKey { provider: "Vertex" }
);
}
#[tokio::test]
async fn transport_errors_remove_urls_and_keep_dispatch_context() {
let error = reqwest::Client::builder()

View file

@ -1,10 +1,12 @@
pub mod audio_transcription;
pub mod auth;
pub mod caching;
pub mod call_lifecycle;
pub mod chat_completions;
pub mod constants;
pub mod error;
pub mod http_utils;
mod media;
pub mod messages;
#[cfg(any(feature = "observability", test))]
pub mod observability;
@ -16,4 +18,5 @@ pub mod router;
pub mod routing_utils;
mod url_utils;
pub use auth::AuthError;
pub use error::Error;

View file

@ -0,0 +1,528 @@
use std::future::Future;
use std::io;
use std::net::{IpAddr, SocketAddr};
use std::pin::Pin;
use std::sync::Arc;
use std::time::Duration;
use reqwest::Url;
use reqwest::dns::{Addrs, Name, Resolve, Resolving};
use crate::constants::MEDIA_CONNECT_TIMEOUT_SECS;
use crate::error::{MediaError, TransportError};
#[derive(Clone)]
pub(crate) struct MediaFetcher {
client: reqwest::Client,
address_resolver: Arc<dyn AddressResolver>,
allow_private_network: bool,
}
type AddressResolution<'a> = Pin<Box<dyn Future<Output = io::Result<Vec<SocketAddr>>> + Send + 'a>>;
trait AddressResolver: Send + Sync {
fn resolve<'a>(&'a self, host: &'a str, port: u16) -> AddressResolution<'a>;
}
#[derive(Clone, Copy)]
pub(crate) struct DownloadPolicy {
pub(crate) timeout: Duration,
pub(crate) max_bytes: u64,
pub(crate) max_redirects: usize,
}
#[derive(Debug)]
pub(crate) struct DownloadedMedia {
pub(crate) bytes: Vec<u8>,
pub(crate) content_type: String,
}
impl MediaFetcher {
pub(crate) fn new() -> Result<Self, reqwest::Error> {
Self::with_resolvers(Arc::new(PublicDnsResolver), Arc::new(SystemAddressResolver))
}
fn with_resolvers<R>(
transport_resolver: Arc<R>,
address_resolver: Arc<dyn AddressResolver>,
) -> Result<Self, reqwest::Error>
where
R: Resolve + 'static,
{
let client = reqwest::Client::builder()
.connect_timeout(Duration::from_secs(MEDIA_CONNECT_TIMEOUT_SECS))
.redirect(reqwest::redirect::Policy::none())
.no_proxy()
.dns_resolver(transport_resolver)
.build()?;
Ok(Self {
client,
address_resolver,
allow_private_network: false,
})
}
#[cfg(test)]
pub(crate) fn for_test(client: reqwest::Client) -> Self {
Self {
client,
address_resolver: Arc::new(AllowPrivateResolver),
allow_private_network: true,
}
}
pub(crate) async fn fetch(
&self,
url: Url,
policy: DownloadPolicy,
) -> Result<DownloadedMedia, MediaError> {
if policy.max_bytes == 0 {
return Err(MediaError::DownloadDisabled);
}
tokio::time::timeout(policy.timeout, self.fetch_before_deadline(url, policy))
.await
.map_err(|_| MediaError::Timeout)?
}
async fn fetch_before_deadline(
&self,
mut url: Url,
policy: DownloadPolicy,
) -> Result<DownloadedMedia, MediaError> {
let mut redirects_followed = 0;
loop {
self.validate_url(&url).await?;
let mut response = self
.client
.get(url.clone())
.send()
.await
.map_err(TransportError::from)?;
if response.status().is_redirection() {
if redirects_followed == policy.max_redirects {
return Err(MediaError::TooManyRedirects);
}
let location = response
.headers()
.get(reqwest::header::LOCATION)
.and_then(|value| value.to_str().ok())
.ok_or(MediaError::MissingRedirectLocation)?;
url = url
.join(location)
.map_err(|_| MediaError::InvalidRedirect)?;
redirects_followed += 1;
continue;
}
if !response.status().is_success() {
return Err(MediaError::Http(response.status().as_u16()));
}
enforce_download_size(response.content_length().unwrap_or(0), policy.max_bytes)?;
let content_type = response
.headers()
.get(reqwest::header::CONTENT_TYPE)
.and_then(|value| value.to_str().ok())
.and_then(|value| value.split(';').next())
.map(str::trim)
.filter(|value| !value.is_empty())
.unwrap_or("application/octet-stream")
.to_string();
let mut bytes = Vec::new();
while let Some(chunk) = response.chunk().await.map_err(TransportError::from)? {
enforce_download_size(bytes.len() as u64 + chunk.len() as u64, policy.max_bytes)?;
bytes.extend_from_slice(&chunk);
}
return Ok(DownloadedMedia {
bytes,
content_type,
});
}
}
async fn validate_url(&self, url: &Url) -> Result<(), MediaError> {
if !matches!(url.scheme(), "http" | "https")
|| !url.username().is_empty()
|| url.password().is_some()
{
return Err(MediaError::BlockedUrl);
}
let host = url.host_str().ok_or(MediaError::BlockedUrl)?;
if self.allow_private_network {
return Ok(());
}
if let Ok(ip) = host.parse::<IpAddr>() {
return (!is_blocked_ip(ip))
.then_some(())
.ok_or(MediaError::BlockedUrl);
}
let port = url.port_or_known_default().ok_or(MediaError::BlockedUrl)?;
let addresses = self
.address_resolver
.resolve(host, port)
.await
.map_err(|error| TransportError::Network(error.to_string()))?;
validate_addresses(&addresses)
}
}
fn enforce_download_size(length: u64, max_bytes: u64) -> Result<(), MediaError> {
if length > max_bytes {
return Err(MediaError::DownloadTooLarge);
}
Ok(())
}
fn validate_addresses(addresses: &[SocketAddr]) -> Result<(), MediaError> {
if addresses.is_empty() || addresses.iter().any(|address| is_blocked_ip(address.ip())) {
return Err(MediaError::BlockedUrl);
}
Ok(())
}
fn is_blocked_ip(ip: IpAddr) -> bool {
match ip {
IpAddr::V4(ip) => {
let [first, second, third, _] = ip.octets();
first == 0
|| first == 10
|| first == 127
|| (first == 100 && (64..=127).contains(&second))
|| (first == 169 && second == 254)
|| (first == 172 && (16..=31).contains(&second))
|| (first == 192 && second == 0 && (third == 0 || third == 2))
|| (first == 192 && second == 168)
|| (first == 192 && second == 88 && third == 99)
|| (first == 198 && (second == 18 || second == 19))
|| (first == 198 && second == 51 && third == 100)
|| (first == 203 && second == 0 && third == 113)
|| first >= 224
}
IpAddr::V6(ip) => {
let segments = ip.segments();
ip.is_loopback()
|| ip.is_unspecified()
|| ip.is_multicast()
|| (segments[0] & 0xfe00) == 0xfc00
|| (segments[0] & 0xffc0) == 0xfe80
|| (segments[0] & 0xffc0) == 0xfec0
|| (segments[0] == 0x2001 && segments[1] == 0x0db8)
|| ip
.to_ipv4_mapped()
.or_else(|| ip.to_ipv4())
.map(|ipv4| is_blocked_ip(IpAddr::V4(ipv4)))
.unwrap_or(false)
}
}
}
#[derive(Default)]
struct PublicDnsResolver;
struct SystemAddressResolver;
impl AddressResolver for SystemAddressResolver {
fn resolve<'a>(&'a self, host: &'a str, port: u16) -> AddressResolution<'a> {
Box::pin(async move {
Ok(tokio::net::lookup_host((host, port))
.await?
.collect::<Vec<_>>())
})
}
}
#[cfg(test)]
struct AllowPrivateResolver;
#[cfg(test)]
impl AddressResolver for AllowPrivateResolver {
fn resolve<'a>(&'a self, _host: &'a str, port: u16) -> AddressResolution<'a> {
Box::pin(async move { Ok(vec![SocketAddr::from(([8, 8, 8, 8], port))]) })
}
}
impl Resolve for PublicDnsResolver {
fn resolve(&self, name: Name) -> Resolving {
let host = name.as_str().to_string();
Box::pin(async move {
let addresses = tokio::net::lookup_host((host.as_str(), 0))
.await
.map_err(|error| Box::new(error) as Box<dyn std::error::Error + Send + Sync>)?
.collect::<Vec<_>>();
validate_addresses(&addresses).map_err(|_| {
Box::new(io::Error::other("destination rejected by network policy"))
as Box<dyn std::error::Error + Send + Sync>
})?;
Ok(Box::new(addresses.into_iter()) as Addrs)
})
}
}
#[cfg(test)]
mod tests {
use super::*;
use std::collections::HashSet;
use tokio::io::{AsyncReadExt, AsyncWriteExt};
use tokio::net::TcpListener;
async fn serve(response: &'static [u8]) -> (Url, tokio::task::JoinHandle<()>) {
let listener = TcpListener::bind("127.0.0.1:0")
.await
.expect("test listener binds");
let address = listener.local_addr().expect("listener has address");
let task = tokio::spawn(async move {
let (mut socket, _) = listener.accept().await.expect("accepts request");
let mut request = [0_u8; 1024];
let bytes_read = socket.read(&mut request).await.expect("reads request");
assert!(bytes_read > 0);
socket.write_all(response).await.expect("writes response");
});
(
Url::parse(&format!("http://{address}/document")).expect("valid test URL"),
task,
)
}
async fn serve_named(
host: &str,
responses: Vec<&'static [u8]>,
) -> (Url, tokio::task::JoinHandle<Vec<String>>, SocketAddr) {
let listener = TcpListener::bind("127.0.0.1:0")
.await
.expect("test listener binds");
let address = listener.local_addr().expect("listener has address");
let task = tokio::spawn(async move {
let mut requests = Vec::with_capacity(responses.len());
for response in responses {
let (mut socket, _) = listener.accept().await.expect("accepts request");
let mut request = [0_u8; 4096];
let bytes_read = socket.read(&mut request).await.expect("reads request");
requests.push(String::from_utf8_lossy(&request[..bytes_read]).into_owned());
socket.write_all(response).await.expect("writes response");
}
requests
});
(
Url::parse(&format!("http://{host}:{}/document", address.port()))
.expect("valid test URL"),
task,
address,
)
}
struct LoopbackDnsResolver(SocketAddr);
impl Resolve for LoopbackDnsResolver {
fn resolve(&self, _name: Name) -> Resolving {
let address = self.0;
Box::pin(async move { Ok(Box::new(vec![address].into_iter()) as Addrs) })
}
}
struct TestAddressResolver {
blocked_hosts: HashSet<&'static str>,
}
impl AddressResolver for TestAddressResolver {
fn resolve<'a>(&'a self, host: &'a str, port: u16) -> AddressResolution<'a> {
let blocked = self.blocked_hosts.contains(host);
Box::pin(async move {
let ip = if blocked {
IpAddr::from([127, 0, 0, 1])
} else {
IpAddr::from([8, 8, 8, 8])
};
Ok(vec![SocketAddr::new(ip, port)])
})
}
}
fn policy_checked_fetcher(
address: SocketAddr,
blocked_hosts: HashSet<&'static str>,
) -> MediaFetcher {
MediaFetcher::with_resolvers(
Arc::new(LoopbackDnsResolver(address)),
Arc::new(TestAddressResolver { blocked_hosts }),
)
.expect("test fetcher builds")
}
fn policy(max_bytes: u64, max_redirects: usize) -> DownloadPolicy {
DownloadPolicy {
timeout: Duration::from_secs(1),
max_bytes,
max_redirects,
}
}
#[test]
fn blocks_non_public_addresses() {
for address in [
"0.0.0.1",
"10.0.0.1",
"100.64.0.1",
"127.0.0.1",
"169.254.1.1",
"172.16.0.1",
"192.168.0.1",
"198.18.0.1",
"198.51.100.1",
"203.0.113.1",
"224.0.0.1",
"::1",
"fc00::1",
"fe80::1",
"2001:db8::1",
"::ffff:127.0.0.1",
] {
assert!(is_blocked_ip(address.parse().expect("valid test address")));
}
assert!(!is_blocked_ip(
"8.8.8.8".parse().expect("valid public address")
));
}
#[tokio::test]
async fn fetches_exact_limit_and_normalizes_content_type() {
let (url, server) = serve(
b"HTTP/1.1 200 OK\r\nContent-Type: application/pdf; charset=binary\r\nContent-Length: 3\r\nConnection: close\r\n\r\nabc",
)
.await;
let client = reqwest::Client::builder()
.redirect(reqwest::redirect::Policy::none())
.build()
.expect("test client builds");
let media = MediaFetcher::for_test(client)
.fetch(url, policy(3, 0))
.await
.expect("download succeeds at exact limit");
server.await.expect("server completes");
assert_eq!(media.bytes, b"abc");
assert_eq!(media.content_type, "application/pdf");
}
#[tokio::test]
async fn rejects_declared_oversize_body() {
let (url, server) = serve(
b"HTTP/1.1 200 OK\r\nContent-Type: application/pdf\r\nContent-Length: 3\r\nConnection: close\r\n\r\nabc",
)
.await;
let client = reqwest::Client::builder()
.redirect(reqwest::redirect::Policy::none())
.build()
.expect("test client builds");
let error = MediaFetcher::for_test(client)
.fetch(url, policy(2, 0))
.await
.expect_err("oversize body is rejected");
server.await.expect("server completes");
assert!(matches!(error, MediaError::DownloadTooLarge));
}
#[tokio::test]
async fn rejects_streamed_oversize_body() {
let (url, server) = serve(
b"HTTP/1.1 200 OK\r\nTransfer-Encoding: chunked\r\nConnection: close\r\n\r\n2\r\nab\r\n2\r\ncd\r\n0\r\n\r\n",
)
.await;
let client = reqwest::Client::builder()
.redirect(reqwest::redirect::Policy::none())
.build()
.expect("test client builds");
let error = MediaFetcher::for_test(client)
.fetch(url, policy(3, 0))
.await
.expect_err("stream crossing limit is rejected");
server.await.expect("server completes");
assert!(matches!(error, MediaError::DownloadTooLarge));
}
#[tokio::test]
async fn follows_allowed_redirects_and_revalidates_each_destination() {
let (url, server, address) = serve_named(
"public.test",
vec![
b"HTTP/1.1 302 Found\r\nLocation: /final\r\nContent-Length: 0\r\n\r\n",
b"HTTP/1.1 200 OK\r\nContent-Type: text/plain\r\nContent-Length: 2\r\nConnection: close\r\n\r\nok",
],
)
.await;
let media = policy_checked_fetcher(address, HashSet::new())
.fetch(url, policy(2, 1))
.await
.expect("redirected fetch succeeds");
let requests = server.await.expect("server completes");
assert_eq!(requests.len(), 2);
assert!(requests[1].starts_with("GET /final "));
assert_eq!(media.bytes, b"ok");
}
#[tokio::test]
async fn blocks_redirected_private_destination_before_second_request() {
let (url, server, address) = serve_named(
"public.test",
vec![b"HTTP/1.1 302 Found\r\nLocation: http://blocked.test/document\r\nContent-Length: 0\r\n\r\n"],
)
.await;
let error = policy_checked_fetcher(address, HashSet::from(["blocked.test"]))
.fetch(url, policy(10, 1))
.await
.expect_err("private redirect is rejected");
let requests = server.await.expect("server completes");
assert_eq!(requests.len(), 1);
assert!(matches!(error, MediaError::BlockedUrl));
}
#[tokio::test]
async fn enforces_total_timeout() {
let listener = TcpListener::bind("127.0.0.1:0")
.await
.expect("test listener binds");
let address = listener.local_addr().expect("listener has address");
let server = tokio::spawn(async move {
let (_socket, _) = listener.accept().await.expect("accepts request");
tokio::time::sleep(Duration::from_millis(100)).await;
});
let url = Url::parse(&format!("http://public.test:{}/document", address.port()))
.expect("valid test URL");
let error = policy_checked_fetcher(address, HashSet::new())
.fetch(
url,
DownloadPolicy {
timeout: Duration::from_millis(20),
max_bytes: 10,
max_redirects: 0,
},
)
.await
.expect_err("fetch times out");
server.await.expect("server completes");
assert!(matches!(error, MediaError::Timeout));
}
#[tokio::test]
async fn document_client_does_not_send_ambient_credentials() {
let (url, server, address) = serve_named(
"public.test",
vec![b"HTTP/1.1 200 OK\r\nContent-Length: 2\r\nConnection: close\r\n\r\nok"],
)
.await;
policy_checked_fetcher(address, HashSet::new())
.fetch(url, policy(2, 0))
.await
.expect("fetch succeeds");
let requests = server.await.expect("server completes");
assert!(!requests[0].to_ascii_lowercase().contains("authorization:"));
assert!(!requests[0].to_ascii_lowercase().contains("api-key:"));
}
#[tokio::test]
async fn rejects_url_credentials_before_network_access() {
let fetcher = MediaFetcher::new().expect("media fetcher builds");
let url =
Url::parse("https://user:password@8.8.8.8/document").expect("credentialed URL parses");
assert!(matches!(
fetcher.validate_url(&url).await,
Err(MediaError::BlockedUrl)
));
}
}

View file

@ -0,0 +1,214 @@
use super::super::OcrAdapter;
use crate::Error;
use crate::auth::{InputSource, Sourced};
use crate::constants::{AZURE_DI_API_VERSION, AZURE_DI_SUBSCRIPTION_HEADER};
use crate::ocr::OcrClient;
use crate::ocr::codecs::document_intelligence::{
self, AzureDocumentIntelligenceOperation, DocumentIntelligenceParams,
};
use crate::ocr::error::{OcrError, OcrRequestError, OcrResponseError};
use crate::ocr::prepare::{credential_env, transform_request_body};
use crate::ocr::registry::OcrProvider;
use crate::ocr::types::{LiteLLMOcrRequest, LiteLLMOcrResponse, OcrConnection, OcrResponseFormat};
use crate::ocr::wire::DecodedOcrResponse;
use crate::providers::azure_ai::auth::AzureAuthInputs;
use crate::url_utils::ApiUrl;
mod polling;
const AZURE_DI_API_KEY_ENV: &str = "AZURE_DOCUMENT_INTELLIGENCE_API_KEY";
const AZURE_DI_ENDPOINT_ENV: &str = "AZURE_DOCUMENT_INTELLIGENCE_ENDPOINT";
#[derive(Clone, Debug)]
pub(crate) struct AzureDocumentIntelligenceAdapter;
impl OcrAdapter for AzureDocumentIntelligenceAdapter {
type ProviderResponse = AzureDocumentIntelligenceOperation;
const PROVIDER: OcrProvider = OcrProvider::AzureAi;
async fn prepare_request(
&self,
request: &LiteLLMOcrRequest,
client: &OcrClient,
) -> Result<reqwest::Request, OcrError> {
let params = map_ocr_params(request)?;
let config = AzureAuthInputs::from_sourced_optional_params(
&request.optional_params,
&request.input_sources,
)
.map_err(Error::from)?;
let headers = validate_environment(&request.connection, &config, &credential_env).await?;
let endpoint = nonblank(request.connection.api_base.clone())
.or_else(|| nonblank(credential_env(AZURE_DI_ENDPOINT_ENV)))
.ok_or_else(|| Error::Auth("Missing Azure Document Intelligence API Base - Set AZURE_DOCUMENT_INTELLIGENCE_ENDPOINT or pass api_base".into()))?;
let url = get_complete_url(&endpoint, &request.model, &params)?;
let body = document_intelligence::transform_ocr_request(request.document.clone())?;
transform_request_body(client, request, &url, &headers, body, |_| Ok(())).await
}
fn transform_ocr_response(
&self,
request: &LiteLLMOcrRequest,
response: Self::ProviderResponse,
) -> Result<LiteLLMOcrResponse, OcrResponseError> {
document_intelligence::transform_ocr_response(&request.model, response)
}
async fn read_response(
&self,
client: &OcrClient,
response: reqwest::Response,
url: &str,
headers: &[(String, String)],
request: &LiteLLMOcrRequest,
) -> Result<DecodedOcrResponse<Self::ProviderResponse>, OcrError> {
polling::read_operation_response(
client.polling_http(),
response,
url,
headers,
&request.connection,
request.response_format()? == OcrResponseFormat::Native,
)
.await
}
}
#[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)]
fn map_ocr_params(
request: &LiteLLMOcrRequest,
) -> Result<DocumentIntelligenceParams, OcrRequestError> {
let params = document_intelligence::decode_input_params(
request.optional_params.clone(),
"optional_params",
)?;
let crate::ocr::prepare::ParsedProviderParams {
known: params,
extra_params: _extra_params,
} = params;
document_intelligence::map_ocr_params(params)
}
fn get_complete_url(
endpoint: &str,
model: &str,
params: &DocumentIntelligenceParams,
) -> Result<String, OcrError> {
let model = format!("{}:analyze", model_id(model)?);
ApiUrl::parse(endpoint)
.and_then(|url| url.complete_path(&["documentintelligence", "documentModels", &model]))
.map(|url| {
url.append_query_pairs(
[("api-version", AZURE_DI_API_VERSION)]
.into_iter()
.chain(params.pages.iter().map(|pages| ("pages", pages.as_str())))
.chain(
params
.features
.iter()
.map(|features| ("features", features.as_str())),
),
)
.into_string()
})
.map_err(|_| OcrRequestError::RequestField {
path: "api_base".into(),
})
.map_err(OcrError::from)
}
async fn validate_environment(
connection: &OcrConnection,
config: &AzureAuthInputs,
env_lookup: &(dyn Fn(&str) -> Option<String> + Sync),
) -> Result<Vec<(String, String)>, OcrError> {
if crate::http_utils::has_header(&connection.extra_headers, "authorization")
|| crate::http_utils::has_header(&connection.extra_headers, AZURE_DI_SUBSCRIPTION_HEADER)
{
super::validate_destination(connection, connection.extra_headers_source)?;
return Ok(connection.extra_headers.clone());
}
let key = nonblank(connection.api_key.clone())
.map(|value| Sourced::new(value, connection.api_key_source))
.or_else(|| {
nonblank(env_lookup(AZURE_DI_API_KEY_ENV))
.map(|value| Sourced::new(value, InputSource::Environment))
});
if let Some(key) = key {
super::validate_destination(connection, key.source())?;
return Ok(
std::iter::once((AZURE_DI_SUBSCRIPTION_HEADER.into(), key.into_value()))
.chain(connection.extra_headers.clone())
.collect(),
);
}
let token = super::resolve_entra(config, env_lookup)
.await?
.ok_or(Error::MissingAzureDocumentIntelligenceCredentials)?;
super::validate_destination(connection, token.source())?;
Ok(
std::iter::once(("Authorization".into(), format!("Bearer {}", token.value())))
.chain(connection.extra_headers.clone())
.collect(),
)
}
fn model_id(model: &str) -> Result<&str, OcrRequestError> {
let model = model.rsplit('/').next().unwrap_or(model);
if matches!(model, "." | "..") {
return Err(OcrRequestError::DotModel);
}
Ok(model)
}
fn nonblank(value: Option<String>) -> Option<String> {
value
.map(|value| value.trim().to_string())
.filter(|value| !value.is_empty())
}
#[cfg(test)]
mod tests {
use super::*;
#[tokio::test]
async fn request_endpoint_cannot_receive_environment_key() {
let connection = OcrConnection {
api_base: Some("https://request.example".into()),
api_base_source: InputSource::Request,
..Default::default()
};
let error = validate_environment(&connection, &Default::default(), &|name| {
(name == AZURE_DI_API_KEY_ENV).then(|| "environment-key".into())
})
.await
.unwrap_err();
assert!(
error
.to_string()
.contains("request-controlled Azure endpoint")
);
}
#[tokio::test]
async fn request_endpoint_accepts_request_owned_key() {
let connection = OcrConnection {
api_key: Some("request-key".into()),
api_key_source: InputSource::Request,
api_base: Some("https://request.example".into()),
api_base_source: InputSource::Request,
..Default::default()
};
let headers = validate_environment(&connection, &Default::default(), &|_| None)
.await
.unwrap();
assert_eq!(
headers[0],
(AZURE_DI_SUBSCRIPTION_HEADER.into(), "request-key".into())
);
}
}

View file

@ -0,0 +1,100 @@
use std::time::Duration;
use reqwest::Url;
use tokio::time::Instant;
use crate::constants::{AZURE_DI_SUBSCRIPTION_HEADER, OCR_POLL_RETRY_SECS};
use crate::ocr::client::read_json_response;
use crate::ocr::codecs::document_intelligence::{
AzureDocumentIntelligenceOperation, OperationStatus,
};
use crate::ocr::error::{OcrError, OcrPollingError, OcrResponseError};
use crate::ocr::types::OcrConnection;
use crate::ocr::wire::DecodedOcrResponse;
pub(super) async fn read_operation_response(
http_client: &reqwest::Client,
response: reqwest::Response,
original_url: &str,
headers: &[(String, String)],
connection: &OcrConnection,
native: bool,
) -> Result<DecodedOcrResponse<AzureDocumentIntelligenceOperation>, OcrError> {
if response.status() != reqwest::StatusCode::ACCEPTED {
return read_json_response(response, native).await;
}
let location = response
.headers()
.get("operation-location")
.and_then(|value| value.to_str().ok())
.ok_or(OcrPollingError::PollLocation)?;
let original = Url::parse(original_url).map_err(|_| OcrPollingError::PollOrigin)?;
let operation = Url::parse(location).map_err(|_| OcrPollingError::PollOrigin)?;
if original.origin() != operation.origin()
|| !operation.username().is_empty()
|| operation.password().is_some()
{
return Err(OcrPollingError::PollOrigin.into());
}
poll_operation(http_client, operation, headers, connection, native).await
}
async fn poll_operation(
http_client: &reqwest::Client,
url: Url,
headers: &[(String, String)],
connection: &OcrConnection,
native: bool,
) -> Result<DecodedOcrResponse<AzureDocumentIntelligenceOperation>, OcrError> {
let deadline = Instant::now()
.checked_add(connection.poll_timeout)
.ok_or(OcrPollingError::PollTimeout)?;
loop {
let remaining = deadline
.checked_duration_since(Instant::now())
.filter(|remaining| !remaining.is_zero())
.ok_or(OcrPollingError::PollTimeout)?;
let builder = http_client
.get(url.clone())
.timeout(remaining.min(connection.timeout));
let builder = crate::http_utils::with_headers(
builder,
headers,
crate::http_utils::HeaderPolicy::Only(&[AZURE_DI_SUBSCRIPTION_HEADER, "authorization"]),
);
let response = tokio::time::timeout_at(deadline, crate::http_utils::http_request(builder))
.await
.map_err(|_| OcrPollingError::PollTimeout)?
.map_err(crate::error::TransportError::from)?;
let retry = response
.headers()
.get(reqwest::header::RETRY_AFTER)
.and_then(|value| value.to_str().ok())
.and_then(|value| value.parse::<u64>().ok())
.unwrap_or(OCR_POLL_RETRY_SECS)
.max(1);
let decoded = tokio::time::timeout_at(
deadline,
read_json_response::<AzureDocumentIntelligenceOperation>(response, native),
)
.await
.map_err(|_| OcrPollingError::PollTimeout)??;
match &decoded.data.status {
Some(OperationStatus::Succeeded) => return Ok(decoded),
Some(OperationStatus::Running | OperationStatus::NotStarted) => {
tokio::time::timeout_at(deadline, tokio::time::sleep(Duration::from_secs(retry)))
.await
.map_err(|_| OcrPollingError::PollTimeout)?;
}
status => {
return Err(OcrResponseError::OperationStatus(
status
.as_ref()
.map(ToString::to_string)
.unwrap_or_else(|| "None".into()),
)
.into());
}
}
}
}

View file

@ -0,0 +1,217 @@
use super::super::OcrAdapter;
use crate::Error;
use crate::auth::{InputSource, Sourced};
use crate::constants::AZURE_AI_OCR_PATH;
use crate::ocr::OcrClient;
use crate::ocr::codecs::mistral::{self, MistralOcrParams, MistralOcrResponse};
use crate::ocr::document::{inline_remote_document, validate_inline_document};
use crate::ocr::error::{OcrError, OcrRequestError, OcrResponseError};
use crate::ocr::prepare::{
_prepare_ocr_request, ParsedProviderParams, credential_env, transform_request_body,
};
use crate::ocr::registry::OcrProvider;
use crate::ocr::types::{LiteLLMOcrRequest, LiteLLMOcrResponse, OcrConnection};
use crate::providers::azure_ai::auth::AzureAuthInputs;
use crate::url_utils::ApiUrl;
const AZURE_AI_API_KEY_ENV: &str = "AZURE_AI_API_KEY";
const AZURE_AI_API_BASE_ENV: &str = "AZURE_AI_API_BASE";
#[derive(Clone, Debug)]
pub(crate) struct AzureMistralAdapter;
impl OcrAdapter for AzureMistralAdapter {
type ProviderResponse = MistralOcrResponse;
const PROVIDER: OcrProvider = OcrProvider::AzureAi;
async fn prepare_request(
&self,
request: &LiteLLMOcrRequest,
client: &OcrClient,
) -> Result<reqwest::Request, OcrError> {
let ParsedProviderParams {
known: params,
extra_params: _extra_params,
} = _prepare_ocr_request::<MistralOcrParams>(request)?;
let config = AzureAuthInputs::from_sourced_optional_params(
&request.optional_params,
&request.input_sources,
)
.map_err(Error::from)?;
let headers = validate_environment(&request.connection, &config, &credential_env).await?;
let url = get_complete_url(request.connection.api_base.as_deref(), &credential_env)?;
let document = inline_remote_document(
client.document_fetcher(),
request.document.clone(),
&request.connection,
)
.await?;
let body = mistral::transform_ocr_request(&request.model, document, &params)?;
transform_request_body(client, request, &url, &headers, body, |body| {
validate_inline_document(&body.document)
})
.await
}
fn transform_ocr_response(
&self,
request: &LiteLLMOcrRequest,
response: Self::ProviderResponse,
) -> Result<LiteLLMOcrResponse, OcrResponseError> {
mistral::transform_ocr_response(&request.model, response)
}
}
fn get_complete_url(
api_base: Option<&str>,
env_lookup: &dyn Fn(&str) -> Option<String>,
) -> Result<String, OcrError> {
let base = nonblank(api_base.map(str::to_string))
.or_else(|| nonblank(env_lookup(AZURE_AI_API_BASE_ENV)))
.ok_or_else(|| Error::Auth(
"Missing Azure AI API Base - Set AZURE_AI_API_BASE environment variable or pass api_base parameter".into(),
))?;
let path: Vec<&str> = AZURE_AI_OCR_PATH.trim_matches('/').split('/').collect();
ApiUrl::parse(&base)
.and_then(|url| url.complete_path(&path))
.map(|url| url.into_string())
.map_err(|_| {
OcrRequestError::RequestField {
path: "api_base".into(),
}
.into()
})
}
async fn validate_environment(
connection: &OcrConnection,
config: &AzureAuthInputs,
env_lookup: &(dyn Fn(&str) -> Option<String> + Sync),
) -> Result<Vec<(String, String)>, OcrError> {
if crate::http_utils::has_header(&connection.extra_headers, "authorization") {
super::validate_destination(connection, connection.extra_headers_source)?;
return Ok(connection.extra_headers.clone());
}
let key = nonblank(connection.api_key.clone())
.map(|value| Sourced::new(value, connection.api_key_source))
.or_else(|| {
nonblank(env_lookup(AZURE_AI_API_KEY_ENV))
.map(|value| Sourced::new(value, InputSource::Environment))
});
if let Some(key) = key {
super::validate_destination(connection, key.source())?;
return Ok(bearer_headers(connection, key.value()));
}
let key = super::resolve_entra(config, env_lookup)
.await?
.ok_or(Error::MissingAzureAiCredentials)?;
super::validate_destination(connection, key.source())?;
Ok(bearer_headers(connection, key.value()))
}
fn bearer_headers(connection: &OcrConnection, key: &str) -> Vec<(String, String)> {
std::iter::once(("Authorization".into(), format!("Bearer {key}")))
.chain(connection.extra_headers.clone())
.collect()
}
fn nonblank(value: Option<String>) -> Option<String> {
value
.map(|value| value.trim().to_string())
.filter(|value| !value.is_empty())
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn completes_azure_path_and_preserves_query() {
assert_eq!(
get_complete_url(Some("https://example.com/?tenant=a"), &|_| None).unwrap(),
"https://example.com/providers/mistral/azure/ocr?tenant=a"
);
assert_eq!(
get_complete_url(
Some("https://example.com/providers/mistral/azure/ocr"),
&|_| None
)
.unwrap(),
"https://example.com/providers/mistral/azure/ocr"
);
}
#[tokio::test]
async fn supplied_authorization_precedes_keys() {
let connection = OcrConnection {
api_key: Some("request-key".into()),
extra_headers: vec![("authorization".into(), "Bearer prepared".into())],
..Default::default()
};
assert_eq!(
validate_environment(&connection, &Default::default(), &|_| {
Some("environment-key".into())
})
.await
.unwrap(),
connection.extra_headers
);
}
#[tokio::test]
async fn request_key_precedes_environment_key() {
let connection = OcrConnection {
api_key: Some("request-key".into()),
..Default::default()
};
assert_eq!(
validate_environment(&connection, &Default::default(), &|_| {
Some("environment-key".into())
})
.await
.unwrap()[0],
("Authorization".into(), "Bearer request-key".into())
);
}
#[tokio::test]
async fn request_endpoint_cannot_receive_environment_key() {
let connection = OcrConnection {
api_base: Some("https://request.example".into()),
api_base_source: InputSource::Request,
..Default::default()
};
let error = validate_environment(&connection, &Default::default(), &|name| {
(name == AZURE_AI_API_KEY_ENV).then(|| "environment-key".into())
})
.await
.unwrap_err();
assert!(
error
.to_string()
.contains("request-controlled Azure endpoint")
);
}
#[tokio::test]
async fn request_endpoint_accepts_request_owned_key() {
let connection = OcrConnection {
api_key: Some("request-key".into()),
api_key_source: InputSource::Request,
api_base: Some("https://request.example".into()),
api_base_source: InputSource::Request,
..Default::default()
};
let headers = validate_environment(&connection, &Default::default(), &|_| None)
.await
.unwrap();
assert_eq!(
headers[0],
("Authorization".into(), "Bearer request-key".into())
);
}
}

View file

@ -0,0 +1,49 @@
mod document_intelligence;
mod mistral;
use std::sync::OnceLock;
use crate::Error;
use crate::auth::error::AuthConfigurationError;
use crate::auth::{InputSource, Sourced};
use crate::ocr::error::OcrError;
use crate::ocr::types::OcrConnection;
use crate::providers::azure_ai::auth::{AzureAuthInputs, AzureAuthService};
pub(crate) use document_intelligence::AzureDocumentIntelligenceAdapter;
pub(crate) use mistral::AzureMistralAdapter;
async fn resolve_entra(
config: &AzureAuthInputs,
env_lookup: &(dyn Fn(&str) -> Option<String> + Sync),
) -> Result<Option<Sourced<String>>, Error> {
static SERVICE: OnceLock<AzureAuthService> = OnceLock::new();
SERVICE
.get_or_init(AzureAuthService::default)
.get_azure_ad_token(config, env_lookup)
.await
.map(|credential| {
credential.map(|credential| {
let source = credential.source();
let value = credential.value().secret().expose().to_string();
Sourced::new(value, source)
})
})
.map_err(Error::from)
}
fn validate_destination(
connection: &OcrConnection,
credential_source: InputSource,
) -> Result<(), OcrError> {
if connection.api_base.is_some()
&& connection.api_base_source == InputSource::Request
&& credential_source != InputSource::Request
{
return Err(Error::from(crate::AuthError::Configuration(
AuthConfigurationError::RequestAzureCredentialDestination,
))
.into());
}
Ok(())
}

View file

@ -8,9 +8,15 @@ use super::registry::OcrProvider;
use super::types::{LiteLLMOcrRequest, LiteLLMOcrResponse, OcrResponseFormat};
use super::wire::DecodedOcrResponse;
mod azure;
mod mistral;
mod reducto;
mod vertex;
pub(crate) use azure::{AzureDocumentIntelligenceAdapter, AzureMistralAdapter};
pub(crate) use mistral::MistralAdapter;
pub(crate) use reducto::{ReductoLegacyAdapter, ReductoV3Adapter};
pub(crate) use vertex::{VertexDeepSeekAdapter, VertexMistralAdapter};
/// Converts a complete LiteLLM OCR call to provider HTTP and normalizes its response.
pub(crate) trait OcrAdapter: Send + Sync + Sized + 'static {
@ -62,6 +68,12 @@ macro_rules! for_each_ocr_adapter {
($callback:ident) => {
$callback! {
Mistral, $crate::ocr::adapters::MistralAdapter, $crate::ocr::adapters::MistralAdapter, Mistral;
AzureMistral, $crate::ocr::adapters::AzureMistralAdapter, $crate::ocr::adapters::AzureMistralAdapter, AzureAi;
AzureDocumentIntelligence, $crate::ocr::adapters::AzureDocumentIntelligenceAdapter, $crate::ocr::adapters::AzureDocumentIntelligenceAdapter, AzureAi;
ReductoLegacy, $crate::ocr::adapters::ReductoLegacyAdapter, $crate::ocr::adapters::ReductoLegacyAdapter, Reducto;
ReductoV3, $crate::ocr::adapters::ReductoV3Adapter, $crate::ocr::adapters::ReductoV3Adapter, Reducto;
VertexMistral, $crate::ocr::adapters::VertexMistralAdapter, $crate::ocr::adapters::VertexMistralAdapter, VertexAi;
VertexDeepSeek, $crate::ocr::adapters::VertexDeepSeekAdapter, $crate::ocr::adapters::VertexDeepSeekAdapter, VertexAi;
}
};
}

View file

@ -0,0 +1,45 @@
use super::super::OcrAdapter;
use crate::ocr::OcrClient;
use crate::ocr::codecs::reducto::{self, ReductoLegacyParams, ReductoResponse};
use crate::ocr::error::{OcrError, OcrResponseError};
use crate::ocr::prepare::{
_prepare_ocr_request, ParsedProviderParams, build_http_request, credential_env,
guardrail_document, merge_extra_params,
};
use crate::ocr::registry::OcrProvider;
use crate::ocr::types::{LiteLLMOcrRequest, LiteLLMOcrResponse};
#[derive(Clone, Debug)]
pub(crate) struct ReductoLegacyAdapter;
impl OcrAdapter for ReductoLegacyAdapter {
type ProviderResponse = ReductoResponse;
const PROVIDER: OcrProvider = OcrProvider::Reducto;
async fn prepare_request(
&self,
request: &LiteLLMOcrRequest,
client: &OcrClient,
) -> Result<reqwest::Request, OcrError> {
let ParsedProviderParams {
known: params,
extra_params,
} = _prepare_ocr_request::<ReductoLegacyParams>(request)?;
let headers = super::validate_environment(&request.connection, &credential_env)?;
let url = super::get_complete_url(request.connection.api_base.as_deref(), "parse")?;
let document = guardrail_document(request, &url).await?;
let document =
super::prepare_document(client, document, &request.connection, &headers).await?;
let body = reducto::transform_legacy_ocr_request(&request.model, document, &params)?;
let body = merge_extra_params(&body, extra_params)?;
build_http_request(client, request, &url, &headers, &body)
}
fn transform_ocr_response(
&self,
request: &LiteLLMOcrRequest,
response: Self::ProviderResponse,
) -> Result<LiteLLMOcrResponse, OcrResponseError> {
reducto::transform_ocr_response(&request.model, response)
}
}

View file

@ -0,0 +1,148 @@
mod legacy;
mod v3;
use crate::Error;
use crate::constants::{REDUCTO_API_BASE, REDUCTO_API_KEY_ENV, REDUCTO_ID_PREFIX};
use crate::ocr::document::InlineDocument;
use crate::ocr::error::{OcrError, OcrRequestError, OcrResponseError};
use crate::ocr::types::{OcrConnection, OcrDocument};
use crate::url_utils::ApiUrl;
pub(crate) use legacy::ReductoLegacyAdapter;
pub(crate) use v3::ReductoV3Adapter;
pub(super) fn get_complete_url(api_base: Option<&str>, path: &str) -> Result<String, OcrError> {
let base = api_base
.map(str::trim)
.filter(|base| !base.is_empty())
.unwrap_or(REDUCTO_API_BASE);
ApiUrl::parse(base)
.and_then(|url| url.complete_path(&[path]))
.map(|url| url.into_string())
.map_err(|_| {
OcrRequestError::RequestField {
path: "api_base".into(),
}
.into()
})
}
pub(super) fn validate_environment(
connection: &OcrConnection,
env_lookup: &(dyn Fn(&str) -> Option<String> + Sync),
) -> Result<Vec<(String, String)>, OcrError> {
if crate::http_utils::has_header(&connection.extra_headers, "authorization") {
return Ok(connection.extra_headers.clone());
}
let api_key = connection
.api_key
.as_deref()
.map(str::trim)
.filter(|key| !key.is_empty())
.map(str::to_string)
.or_else(|| {
env_lookup(REDUCTO_API_KEY_ENV)
.map(|key| key.trim().to_string())
.filter(|key| !key.is_empty())
})
.ok_or(Error::MissingReductoApiKey)?;
Ok(
std::iter::once(("Authorization".into(), format!("Bearer {api_key}")))
.chain(connection.extra_headers.clone())
.collect(),
)
}
pub(super) async fn prepare_document(
client: &crate::ocr::OcrClient,
document: OcrDocument,
connection: &OcrConnection,
headers: &[(String, String)],
) -> Result<OcrDocument, OcrError> {
if document.source().starts_with(REDUCTO_ID_PREFIX) {
if document.source()[REDUCTO_ID_PREFIX.len()..]
.trim()
.is_empty()
{
return Err(OcrRequestError::RequestField {
path: "document file id".into(),
}
.into());
}
return Ok(document);
}
let inline = InlineDocument::parse(document.source())?.ok_or(OcrRequestError::ReductoSource)?;
let mime = inline.mime_type().to_string();
let bytes = inline.decode(crate::constants::OCR_INLINE_MAX_BYTES)?;
let part = reqwest::multipart::Part::bytes(bytes)
.file_name("document")
.mime_str(&mime)
.map_err(|_| OcrRequestError::InvalidDataUri)?;
let builder = client
.provider_http()
.post(get_complete_url(connection.api_base.as_deref(), "upload")?)
.multipart(reqwest::multipart::Form::new().part("file", part))
.timeout(connection.timeout);
let builder = crate::http_utils::with_headers(
builder,
headers,
crate::http_utils::HeaderPolicy::Except(&["content-type", "content-length"]),
);
let response = crate::http_utils::http_request(builder)
.await
.map_err(crate::error::TransportError::from)?;
let uploaded = crate::ocr::client::read_json_response::<
crate::ocr::codecs::reducto::ReductoUploadResponse,
>(response, false)
.await?
.data;
let file_id = uploaded
.file_id
.as_deref()
.map(str::trim)
.filter(|id| !id.is_empty());
let Some(file_id) = file_id else {
return Err(OcrResponseError::ResponseField {
path: "file_id".into(),
}
.into());
};
Ok(document.with_source(file_id.to_string()))
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn explicit_key_precedes_environment_key() {
let connection = OcrConnection {
api_key: Some("passed-key".into()),
..Default::default()
};
let headers = validate_environment(&connection, &|_| Some("env-key".into())).unwrap();
assert_eq!(headers[0].1, "Bearer passed-key");
}
#[test]
fn blank_explicit_key_uses_environment_key() {
let connection = OcrConnection {
api_key: Some(" ".into()),
..Default::default()
};
let headers = validate_environment(&connection, &|_| Some(" env-key ".into())).unwrap();
assert_eq!(headers[0].1, "Bearer env-key");
}
#[test]
fn existing_authorization_skips_key_lookup() {
let connection = OcrConnection {
extra_headers: vec![("authorization".into(), "Bearer existing".into())],
..Default::default()
};
assert_eq!(
validate_environment(&connection, &|_| None).unwrap(),
connection.extra_headers
);
}
}

View file

@ -0,0 +1,45 @@
use super::super::OcrAdapter;
use crate::ocr::OcrClient;
use crate::ocr::codecs::reducto::{self, ReductoResponse, ReductoV3Params};
use crate::ocr::error::{OcrError, OcrResponseError};
use crate::ocr::prepare::{
_prepare_ocr_request, ParsedProviderParams, build_http_request, credential_env,
guardrail_document, merge_extra_params,
};
use crate::ocr::registry::OcrProvider;
use crate::ocr::types::{LiteLLMOcrRequest, LiteLLMOcrResponse};
#[derive(Clone, Debug)]
pub(crate) struct ReductoV3Adapter;
impl OcrAdapter for ReductoV3Adapter {
type ProviderResponse = ReductoResponse;
const PROVIDER: OcrProvider = OcrProvider::Reducto;
async fn prepare_request(
&self,
request: &LiteLLMOcrRequest,
client: &OcrClient,
) -> Result<reqwest::Request, OcrError> {
let ParsedProviderParams {
known: params,
extra_params,
} = _prepare_ocr_request::<ReductoV3Params>(request)?;
let headers = super::validate_environment(&request.connection, &credential_env)?;
let url = super::get_complete_url(request.connection.api_base.as_deref(), "parse")?;
let document = guardrail_document(request, &url).await?;
let document =
super::prepare_document(client, document, &request.connection, &headers).await?;
let body = reducto::transform_v3_ocr_request(&request.model, document, &params)?;
let body = merge_extra_params(&body, extra_params)?;
build_http_request(client, request, &url, &headers, &body)
}
fn transform_ocr_response(
&self,
request: &LiteLLMOcrRequest,
response: Self::ProviderResponse,
) -> Result<LiteLLMOcrResponse, OcrResponseError> {
reducto::transform_ocr_response(&request.model, response)
}
}

View file

@ -0,0 +1,134 @@
use super::super::OcrAdapter;
use super::validate_destination;
use crate::Error;
use crate::auth::vertex::{self, VertexConfig};
use crate::ocr::OcrClient;
use crate::ocr::codecs::deepseek::{self, DeepSeekOcrParams, DeepSeekOcrResponse};
use crate::ocr::error::{OcrError, OcrRequestError, OcrResponseError};
use crate::ocr::prepare::{
_prepare_ocr_request, ParsedProviderParams, credential_env, transform_request_body,
};
use crate::ocr::registry::OcrProvider;
use crate::ocr::types::{LiteLLMOcrRequest, LiteLLMOcrResponse};
use crate::url_utils::ApiUrl;
const DEFAULT_API_BASE: &str = "https://aiplatform.googleapis.com";
const MODEL_NAMESPACE: &str = "deepseek-ai";
const DEFAULT_LOCATION: &str = "us-central1";
#[derive(Clone, Debug)]
pub(crate) struct VertexDeepSeekAdapter;
impl OcrAdapter for VertexDeepSeekAdapter {
type ProviderResponse = DeepSeekOcrResponse;
const PROVIDER: OcrProvider = OcrProvider::VertexAi;
async fn prepare_request(
&self,
request: &LiteLLMOcrRequest,
client: &OcrClient,
) -> Result<reqwest::Request, OcrError> {
validate_destination(&request.connection)?;
let ParsedProviderParams {
known: params,
extra_params: _extra_params,
} = _prepare_ocr_request::<DeepSeekOcrParams>(request)?;
let config = VertexConfig::from_sourced_optional_params(
&request.optional_params,
&request.input_sources,
)
.map_err(Error::from)?;
let authentication = client
.vertex_auth()
.validate_environment(
request.connection.extra_headers.clone(),
request.connection.api_key.as_deref(),
&config,
&credential_env,
)
.await
.map_err(Error::from)?;
let location = vertex::get_vertex_ai_location(&config, &credential_env)
.unwrap_or_else(|| DEFAULT_LOCATION.to_string());
let url = get_complete_url(
request.connection.api_base.as_deref(),
&authentication.project_id,
&location,
)?;
let document = request.document.clone();
let body =
deepseek::transform_ocr_request(&provider_model(&request.model), document, &params)?;
transform_request_body(client, request, &url, &authentication.headers, body, |_| {
Ok(())
})
.await
}
fn transform_ocr_response(
&self,
request: &LiteLLMOcrRequest,
response: Self::ProviderResponse,
) -> Result<LiteLLMOcrResponse, OcrResponseError> {
deepseek::transform_ocr_response(&request.model, response)
}
}
fn provider_model(model: &str) -> String {
if model.starts_with(&format!("{MODEL_NAMESPACE}/")) {
model.to_string()
} else {
format!("{MODEL_NAMESPACE}/{model}")
}
}
fn get_complete_url(
api_base: Option<&str>,
project: &str,
location: &str,
) -> Result<String, OcrError> {
let base = api_base
.map(str::trim)
.filter(|base| !base.is_empty())
.unwrap_or(DEFAULT_API_BASE);
ApiUrl::parse(base)
.and_then(|url| {
url.complete_path(&[
"v1",
"projects",
project,
"locations",
location,
"endpoints",
"openapi",
"chat",
"completions",
])
})
.map(|url| url.into_string())
.map_err(|_| {
OcrRequestError::RequestField {
path: "api_base".into(),
}
.into()
})
}
#[cfg(test)]
mod tests {
use super::{get_complete_url, provider_model};
#[test]
fn adapter_owns_model_namespace_and_endpoint() {
assert_eq!(
provider_model("deepseek-ocr-maas"),
"deepseek-ai/deepseek-ocr-maas"
);
assert_eq!(
provider_model("deepseek-ai/deepseek-ocr-maas"),
"deepseek-ai/deepseek-ocr-maas"
);
assert_eq!(
get_complete_url(None, "proj-1", "europe-west4").unwrap(),
"https://aiplatform.googleapis.com/v1/projects/proj-1/locations/europe-west4/endpoints/openapi/chat/completions"
);
}
}

View file

@ -0,0 +1,154 @@
use super::super::OcrAdapter;
use super::validate_destination;
use crate::Error;
use crate::auth::vertex::{self, VertexConfig};
use crate::ocr::OcrClient;
use crate::ocr::codecs::mistral::{self, MistralOcrParams, MistralOcrResponse};
use crate::ocr::document::{inline_remote_document, validate_inline_document};
use crate::ocr::error::{OcrError, OcrRequestError, OcrResponseError};
use crate::ocr::prepare::{
_prepare_ocr_request, ParsedProviderParams, credential_env, transform_request_body,
};
use crate::ocr::registry::OcrProvider;
use crate::ocr::types::{LiteLLMOcrRequest, LiteLLMOcrResponse};
use crate::url_utils::ApiUrl;
const DEFAULT_LOCATION: &str = "us-central1";
#[derive(Clone, Debug)]
pub(crate) struct VertexMistralAdapter;
impl OcrAdapter for VertexMistralAdapter {
type ProviderResponse = MistralOcrResponse;
const PROVIDER: OcrProvider = OcrProvider::VertexAi;
async fn prepare_request(
&self,
request: &LiteLLMOcrRequest,
client: &OcrClient,
) -> Result<reqwest::Request, OcrError> {
validate_destination(&request.connection)?;
let ParsedProviderParams {
known: params,
extra_params: _extra_params,
} = _prepare_ocr_request::<MistralOcrParams>(request)?;
let config = VertexConfig::from_sourced_optional_params(
&request.optional_params,
&request.input_sources,
)
.map_err(Error::from)?;
let authentication = client
.vertex_auth()
.validate_environment(
request.connection.extra_headers.clone(),
request.connection.api_key.as_deref(),
&config,
&credential_env,
)
.await
.map_err(Error::from)?;
let location = vertex::get_vertex_ai_location(&config, &credential_env)
.unwrap_or_else(|| DEFAULT_LOCATION.to_string());
let url = get_complete_url(
request.connection.api_base.as_deref(),
&authentication.project_id,
&location,
&request.model,
)?;
let document = inline_remote_document(
client.document_fetcher(),
request.document.clone(),
&request.connection,
)
.await?;
let body = mistral::transform_ocr_request(&request.model, document, &params)?;
transform_request_body(
client,
request,
&url,
&authentication.headers,
body,
|body| validate_inline_document(&body.document),
)
.await
}
fn transform_ocr_response(
&self,
request: &LiteLLMOcrRequest,
response: Self::ProviderResponse,
) -> Result<LiteLLMOcrResponse, OcrResponseError> {
mistral::transform_ocr_response(&request.model, response)
}
}
fn get_complete_url(
api_base: Option<&str>,
project: &str,
location: &str,
model: &str,
) -> Result<String, OcrError> {
validate_location(location)?;
let default_base = format!("https://{location}-aiplatform.googleapis.com");
let base = api_base
.map(str::trim)
.filter(|base| !base.is_empty())
.unwrap_or(&default_base);
let prediction = format!("{model}:rawPredict");
ApiUrl::parse(base)
.and_then(|url| {
url.complete_path(&[
"v1",
"projects",
project,
"locations",
location,
"publishers",
"mistralai",
"models",
&prediction,
])
})
.map(|url| url.into_string())
.map_err(|_| {
OcrRequestError::RequestField {
path: "api_base".into(),
}
.into()
})
}
fn validate_location(location: &str) -> Result<(), OcrError> {
let valid = !location.is_empty()
&& location
.bytes()
.all(|value| value.is_ascii_lowercase() || value.is_ascii_digit() || value == b'-')
&& location
.as_bytes()
.first()
.is_some_and(u8::is_ascii_alphanumeric)
&& location
.as_bytes()
.last()
.is_some_and(u8::is_ascii_alphanumeric);
if valid {
return Ok(());
}
Err(OcrRequestError::RequestField {
path: "vertex_location".into(),
}
.into())
}
#[cfg(test)]
mod tests {
use super::get_complete_url;
#[test]
fn endpoint_uses_location_project_and_model() {
assert_eq!(
get_complete_url(None, "proj-1", "europe-west4", "mistral-ocr-maas").unwrap(),
"https://europe-west4-aiplatform.googleapis.com/v1/projects/proj-1/locations/europe-west4/publishers/mistralai/models/mistral-ocr-maas:rawPredict"
);
assert!(get_complete_url(None, "proj-1", "attacker.example/path", "model").is_err());
}
}

View file

@ -0,0 +1,21 @@
mod deepseek;
mod mistral;
use crate::Error;
use crate::auth::InputSource;
use crate::auth::error::AuthConfigurationError;
use crate::ocr::error::OcrError;
use crate::ocr::types::OcrConnection;
pub(crate) use deepseek::VertexDeepSeekAdapter;
pub(crate) use mistral::VertexMistralAdapter;
fn validate_destination(connection: &OcrConnection) -> Result<(), OcrError> {
if connection.api_base.is_some() && connection.api_base_source == InputSource::Request {
return Err(Error::from(crate::AuthError::Configuration(
AuthConfigurationError::RequestVertexCredentialDestination,
))
.into());
}
Ok(())
}

View file

@ -8,17 +8,28 @@ use super::handler::perform_ocr_request;
use super::types::{LiteLLMOcrRequest, LiteLLMOcrResponse};
use super::wire::{DecodedOcrResponse, decode_response};
use crate::Error;
use crate::auth::vertex::VertexAuth;
use crate::constants::OCR_CONNECT_TIMEOUT_SECS;
use crate::error::TransportError;
use crate::media::MediaFetcher;
#[derive(Clone)]
pub struct OcrClient {
provider_http: reqwest::Client,
polling_http: reqwest::Client,
document_fetcher: MediaFetcher,
vertex_auth: VertexAuth,
}
impl OcrClient {
pub fn new(provider_http: reqwest::Client) -> Result<Self, TransportError> {
Ok(Self { provider_http })
let document_fetcher = MediaFetcher::new().map_err(TransportError::from)?;
Ok(Self {
provider_http,
polling_http: no_redirect_http()?,
document_fetcher,
vertex_auth: VertexAuth::default(),
})
}
#[tracing::instrument(
@ -35,10 +46,35 @@ impl OcrClient {
&self.provider_http
}
#[cfg(test)]
pub(crate) fn for_test(provider_http: reqwest::Client) -> Self {
Self { provider_http }
pub(crate) fn polling_http(&self) -> &reqwest::Client {
&self.polling_http
}
pub(crate) fn document_fetcher(&self) -> &MediaFetcher {
&self.document_fetcher
}
pub(crate) fn vertex_auth(&self) -> &VertexAuth {
&self.vertex_auth
}
#[cfg(test)]
pub(crate) fn for_test(provider_http: reqwest::Client, document_http: reqwest::Client) -> Self {
Self {
provider_http,
polling_http: no_redirect_http().expect("test polling client builds"),
document_fetcher: MediaFetcher::for_test(document_http),
vertex_auth: VertexAuth::default(),
}
}
}
fn no_redirect_http() -> Result<reqwest::Client, TransportError> {
reqwest::Client::builder()
.connect_timeout(Duration::from_secs(OCR_CONNECT_TIMEOUT_SECS))
.redirect(reqwest::redirect::Policy::none())
.build()
.map_err(TransportError::from)
}
pub async fn ocr(request: LiteLLMOcrRequest) -> Result<LiteLLMOcrResponse, Error> {

View file

@ -0,0 +1,5 @@
mod transformation;
mod types;
pub(crate) use transformation::{transform_ocr_request, transform_ocr_response};
pub(crate) use types::{DeepSeekOcrParams, DeepSeekOcrResponse};

View file

@ -0,0 +1,98 @@
use serde::de::IntoDeserializer;
use serde_json::{Value, json};
use super::types::*;
use crate::ocr::error::{OcrRequestError, OcrResponseError};
use crate::ocr::types::{LiteLLMOcrResponse, OcrDocument};
#[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)]
pub(crate) fn transform_ocr_request(
provider_model: &str,
document: OcrDocument,
params: &DeepSeekOcrParams,
) -> Result<DeepSeekOcrRequest, OcrRequestError> {
if document.source().is_empty() {
return Err(OcrRequestError::MissingField("document URL"));
}
Ok(DeepSeekOcrRequest {
model: provider_model.to_string(),
messages: vec![DeepSeekOcrMessage {
role: UserRole::User,
content: vec![document],
}],
params: params.clone(),
})
}
pub(crate) fn transform_ocr_response(
model: &str,
response: DeepSeekOcrResponse,
) -> Result<LiteLLMOcrResponse, OcrResponseError> {
let content = response
.choices
.into_iter()
.next()
.and_then(|choice| choice.message.content)
.ok_or(OcrResponseError::EmptyContent)?;
let decoded = decode_content(content)?;
let pages = match decoded.result.pages {
Some(pages) if !pages.is_empty() => pages
.into_iter()
.map(|page| serde_json::to_value(page).expect("DeepSeek page serializes"))
.collect(),
_ => vec![json!({
"index":0,
"markdown":decoded.fallback_markdown,
"images":null
})],
};
Ok(LiteLLMOcrResponse {
pages,
model: decoded.result.model.unwrap_or_else(|| model.to_string()),
document_annotation: decoded.result.document_annotation,
usage_info: decoded.result.usage_info.or(response.usage),
object: "ocr".into(),
extra_fields: decoded.result.extra_fields,
provider_native_response: None,
})
}
struct DecodedContent {
result: DeepSeekOcrResult,
fallback_markdown: String,
}
fn decode_content(content: DeepSeekContent) -> Result<DecodedContent, OcrResponseError> {
let (result, fallback_markdown) = match content {
DeepSeekContent::Text(text) if text.is_empty() => {
return Err(OcrResponseError::EmptyContent);
}
DeepSeekContent::Text(text) => (decode_json_content(&text)?, text),
DeepSeekContent::Object(object) => {
let fallback =
serde_json::to_string(&object).map_err(|_| OcrResponseError::ResponseField {
path: "choices[0].message.content".into(),
})?;
(Some(object), fallback)
}
};
Ok(DecodedContent {
result: result.unwrap_or_default(),
fallback_markdown,
})
}
fn decode_json_content(text: &str) -> Result<Option<DeepSeekOcrResult>, OcrResponseError> {
if !text.trim_start().starts_with('{') {
return Ok(None);
}
let value = match serde_json::from_str::<Value>(text) {
Ok(value) => value,
Err(_) => return Ok(None),
};
serde_path_to_error::deserialize(value.into_deserializer())
.map(Some)
.map_err(|error| OcrResponseError::ResponseField {
path: format!("choices[0].message.content.{}", error.path()),
})
}

View file

@ -0,0 +1,95 @@
use serde::{Deserialize, Serialize};
use serde_json::{Map, Value};
#[derive(Clone, Debug, Default, Serialize, Deserialize)]
pub(crate) struct DeepSeekOcrParams {
#[serde(skip_serializing_if = "Option::is_none")]
pub stream: Option<bool>,
#[serde(skip_serializing_if = "Option::is_none")]
pub temperature: Option<f64>,
#[serde(skip_serializing_if = "Option::is_none")]
pub max_tokens: Option<i64>,
#[serde(skip_serializing_if = "Option::is_none")]
pub top_p: Option<f64>,
#[serde(skip_serializing_if = "Option::is_none")]
pub n: Option<i64>,
#[serde(skip_serializing_if = "Option::is_none")]
pub stop: Option<StopSequences>,
}
#[derive(Clone, Debug, Serialize, Deserialize)]
#[serde(untagged)]
pub(crate) enum StopSequences {
One(String),
Many(Vec<String>),
}
#[derive(Clone, Debug, Serialize, Deserialize)]
pub(crate) struct DeepSeekOcrRequest {
pub model: String,
pub messages: Vec<DeepSeekOcrMessage>,
#[serde(flatten)]
pub params: DeepSeekOcrParams,
}
#[derive(Clone, Debug, Serialize, Deserialize)]
pub(crate) struct DeepSeekOcrMessage {
pub role: UserRole,
pub content: Vec<crate::ocr::types::OcrDocument>,
}
#[derive(Clone, Debug, Serialize, Deserialize)]
#[serde(rename_all = "lowercase")]
pub(crate) enum UserRole {
User,
}
#[derive(Clone, Debug, Deserialize)]
pub(crate) struct DeepSeekOcrResponse {
#[serde(default)]
pub choices: Vec<DeepSeekChoice>,
pub usage: Option<Value>,
}
#[derive(Clone, Debug, Deserialize)]
pub(crate) struct DeepSeekChoice {
pub message: DeepSeekResponseMessage,
}
#[derive(Clone, Debug, Deserialize)]
pub(crate) struct DeepSeekResponseMessage {
pub content: Option<DeepSeekContent>,
}
#[derive(Clone, Debug, Deserialize)]
#[serde(untagged)]
pub(crate) enum DeepSeekContent {
Text(String),
Object(DeepSeekOcrResult),
}
#[derive(Clone, Debug, Default, Serialize, Deserialize)]
pub(crate) struct DeepSeekOcrResult {
#[serde(skip_serializing_if = "Option::is_none")]
pub pages: Option<Vec<DeepSeekPage>>,
#[serde(skip_serializing_if = "Option::is_none")]
pub model: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
pub usage_info: Option<Value>,
#[serde(skip_serializing_if = "Option::is_none")]
pub document_annotation: Option<Value>,
#[serde(flatten)]
pub extra_fields: Map<String, Value>,
}
#[derive(Clone, Debug, Serialize, Deserialize)]
pub(crate) struct DeepSeekPage {
#[serde(default)]
pub index: i64,
#[serde(default)]
pub markdown: String,
pub images: Option<Value>,
pub dimensions: Option<Value>,
#[serde(flatten)]
pub extra_fields: Map<String, Value>,
}

View file

@ -0,0 +1,9 @@
mod params;
mod transformation;
mod types;
pub(crate) use params::{decode_input_params, map_ocr_params};
pub(crate) use transformation::{transform_ocr_request, transform_ocr_response};
pub(crate) use types::{
AzureDocumentIntelligenceOperation, DocumentIntelligenceParams, OperationStatus,
};

View file

@ -0,0 +1,195 @@
use std::collections::BTreeSet;
use serde_json::{Map, Value};
use super::types::{
DocumentIntelligenceInputParams, DocumentIntelligenceParams, FeaturesInput, PagesInput,
};
use crate::ocr::error::OcrRequestError;
use crate::ocr::prepare::ParsedProviderParams;
pub(crate) fn decode_input_params(
params: Map<String, Value>,
prefix: &str,
) -> Result<ParsedProviderParams<DocumentIntelligenceInputParams>, OcrRequestError> {
if let Some(Value::Array(pages)) = params.get("pages") {
if pages.iter().any(Value::is_boolean) {
return Err(OcrRequestError::Pages("boolean page index".into()));
}
if pages
.iter()
.any(|page| page.is_number() && page.as_i64().is_none())
{
return Err(OcrRequestError::Pages("page index is out of range".into()));
}
if !pages.iter().all(Value::is_i64) && !pages.iter().all(Value::is_string) {
return Err(OcrRequestError::Pages("mixed page element types".into()));
}
}
crate::ocr::wire::decode_request_value(Value::Object(params), prefix)
}
pub(crate) fn map_ocr_params(
params: DocumentIntelligenceInputParams,
) -> Result<DocumentIntelligenceParams, OcrRequestError> {
Ok(DocumentIntelligenceParams {
pages: params.pages.map(normalize_pages).transpose()?.flatten(),
features: params
.features
.map(normalize_features)
.transpose()?
.flatten(),
})
}
fn normalize_pages(pages: PagesInput) -> Result<Option<String>, OcrRequestError> {
let normalized = match pages {
PagesInput::ZeroBasedIndices(indices) => {
if indices.is_empty() {
return Ok(None);
}
indices
.into_iter()
.map(|page| {
if page < 0 {
return Err(OcrRequestError::Pages("negative page index".into()));
}
page.checked_add(1)
.ok_or_else(|| OcrRequestError::Pages("page index is out of range".into()))
})
.collect::<Result<BTreeSet<_>, _>>()?
.into_iter()
.map(|page| page.to_string())
.collect::<Vec<_>>()
.join(",")
}
PagesInput::NativeTokens(tokens) => {
if tokens.is_empty() {
return Ok(None);
}
tokens
.iter()
.map(|token| token.trim())
.collect::<Vec<_>>()
.join(",")
}
PagesInput::NativeRange(range) => range
.split(',')
.map(str::trim)
.collect::<Vec<_>>()
.join(","),
};
if !normalized.split(',').all(valid_page_token) {
return Err(OcrRequestError::Pages("invalid native page range".into()));
}
Ok(Some(normalized))
}
fn valid_page_token(token: &str) -> bool {
let mut parts = token.split('-');
let start = parts.next().unwrap_or_default();
if start.is_empty() || !start.chars().all(|character| character.is_ascii_digit()) {
return false;
}
match parts.next() {
None => true,
Some(end) => {
!end.is_empty()
&& end.chars().all(|character| character.is_ascii_digit())
&& parts.next().is_none()
}
}
}
fn normalize_features(features: FeaturesInput) -> Result<Option<String>, OcrRequestError> {
let tokens = match features {
FeaturesInput::Names(names) => names,
FeaturesInput::CommaSeparated(names) => names.split(',').map(str::to_string).collect(),
};
if tokens.is_empty() {
return Ok(None);
}
let normalized = tokens.iter().map(|token| token.trim()).collect::<Vec<_>>();
if !normalized.iter().all(|token| {
let Some((first, rest)) = token.as_bytes().split_first() else {
return false;
};
first.is_ascii_alphabetic() && rest.iter().all(u8::is_ascii_alphanumeric)
}) {
return Err(OcrRequestError::Features);
}
Ok(Some(normalized.join(",")))
}
#[cfg(test)]
mod tests {
use rstest::rstest;
use serde_json::{Value, json};
use super::*;
fn map(value: Value) -> Result<DocumentIntelligenceParams, OcrRequestError> {
let fields = value.as_object().unwrap().clone();
map_ocr_params(decode_input_params(fields, "optional_params")?.known)
}
#[test]
fn input_params_retain_unknown_fields() {
let parsed = decode_input_params(
json!({
"pages": [0],
"future_ocr_option": true,
"extra_body": {"provider_option": "value"}
})
.as_object()
.unwrap()
.clone(),
"optional_params",
)
.unwrap();
assert_eq!(
parsed.known.pages,
Some(PagesInput::ZeroBasedIndices(vec![0]))
);
assert_eq!(parsed.extra_params["future_ocr_option"], true);
assert_eq!(
parsed.extra_params["extra_body"],
json!({"provider_option": "value"})
);
assert_eq!(
serde_json::to_value(map_ocr_params(parsed.known).unwrap()).unwrap(),
json!({"pages": "1", "features": null})
);
}
#[rstest]
#[case(json!(["keyValuePairs"]), "keyValuePairs")]
#[case(json!(["keyValuePairs", "languages"]), "keyValuePairs,languages")]
#[case(json!("keyValuePairs"), "keyValuePairs")]
#[case(json!("keyValuePairs,languages"), "keyValuePairs,languages")]
#[case(json!("keyValuePairs, languages"), "keyValuePairs,languages")]
fn feature_mapping_matches_python(#[case] input: Value, #[case] expected: &str) {
assert_eq!(
map(json!({"features": input})).unwrap().features.as_deref(),
Some(expected)
);
}
#[rstest]
#[case(json!("keyValuePairs&pages=9"))]
#[case(json!("key value pairs"))]
#[case(json!(""))]
#[case(json!([1, 2]))]
#[case(json!([["keyValuePairs"]]))]
#[case(json!({"feature":"keyValuePairs"}))]
#[case(json!(5))]
fn invalid_feature_mapping_matches_python(#[case] input: Value) {
assert!(map(json!({"features": input})).is_err());
}
#[test]
fn empty_feature_list_is_omitted() {
assert_eq!(map(json!({"features": []})).unwrap().features, None);
}
}

View file

@ -0,0 +1,111 @@
use base64::{Engine, engine::general_purpose::STANDARD};
use serde_json::{Map, Value, json};
use super::types::*;
use crate::constants::{AZURE_DI_DEFAULT_DPI, AZURE_DI_DEFAULT_HEIGHT, AZURE_DI_DEFAULT_WIDTH};
use crate::ocr::document::InlineDocument;
use crate::ocr::error::{OcrRequestError, OcrResponseError};
use crate::ocr::types::{LiteLLMOcrResponse, OcrDocument};
#[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)]
pub(crate) fn transform_ocr_request(
document: OcrDocument,
) -> Result<DocumentIntelligenceRequest, OcrRequestError> {
let source = document.source();
if source.is_empty() {
return Err(OcrRequestError::MissingField("document URL"));
}
Ok(if let Some(document) = InlineDocument::parse(source)? {
DocumentIntelligenceRequest::Base64Source(
STANDARD.encode(document.decode(crate::constants::OCR_INLINE_MAX_BYTES)?),
)
} else {
DocumentIntelligenceRequest::UrlSource(source.to_string())
})
}
pub(crate) fn transform_ocr_response(
model: &str,
response: AzureDocumentIntelligenceOperation,
) -> Result<LiteLLMOcrResponse, OcrResponseError> {
if response.status != Some(OperationStatus::Succeeded) {
return Err(OcrResponseError::OperationStatus(
response
.status
.map(|status| status.to_string())
.unwrap_or_else(|| "None".into()),
));
}
let result = response.analyze_result.unwrap_or_default();
let pages = result
.pages
.into_iter()
.map(normalize_page)
.collect::<Result<Vec<_>, _>>()?;
let pages_processed = pages.len();
let mut extra_fields = Map::new();
extra_fields.insert("content".into(), option_value(result.content));
extra_fields.insert("tables".into(), option_value(result.tables));
extra_fields.insert(
"key_value_pairs".into(),
option_value(result.key_value_pairs),
);
Ok(LiteLLMOcrResponse {
pages,
model: model.into(),
document_annotation: None,
usage_info: Some(json!({"pages_processed":pages_processed})),
object: "ocr".into(),
extra_fields,
provider_native_response: None,
})
}
fn normalize_page(page: AzureDocumentIntelligencePage) -> Result<Value, OcrResponseError> {
let index = page
.page_number
.unwrap_or(1)
.checked_sub(1)
.ok_or(OcrResponseError::NumericRange("page.pageNumber"))?;
let scale = if page.unit.as_deref().unwrap_or("inch") == "inch" {
AZURE_DI_DEFAULT_DPI as f64
} else {
1.0
};
let width = pixel_dimension(
page.width.unwrap_or(AZURE_DI_DEFAULT_WIDTH),
scale,
"page.width",
)?;
let height = pixel_dimension(
page.height.unwrap_or(AZURE_DI_DEFAULT_HEIGHT),
scale,
"page.height",
)?;
let markdown = page
.lines
.iter()
.map(|line| line.content.as_deref().unwrap_or_default())
.collect::<Vec<_>>()
.join("\n");
Ok(json!({
"index":index,
"markdown":markdown,
"images":null,
"dimensions":{"width":width,"height":height,"dpi":AZURE_DI_DEFAULT_DPI}
}))
}
fn pixel_dimension(value: f64, scale: f64, field: &'static str) -> Result<i64, OcrResponseError> {
let value = value * scale;
if !value.is_finite() || value < i64::MIN as f64 || value > i64::MAX as f64 {
return Err(OcrResponseError::NumericRange(field));
}
Ok(value.trunc() as i64)
}
fn option_value<T: serde::Serialize>(value: Option<T>) -> Value {
value
.and_then(|value| serde_json::to_value(value).ok())
.unwrap_or(Value::Null)
}

View file

@ -0,0 +1,138 @@
use serde::{Deserialize, Deserializer, Serialize};
use serde_json::{Map, Value};
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
#[serde(untagged)]
pub(crate) enum PagesInput {
ZeroBasedIndices(Vec<i64>),
NativeTokens(Vec<String>),
NativeRange(String),
}
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
#[serde(untagged)]
pub(crate) enum FeaturesInput {
Names(Vec<String>),
CommaSeparated(String),
}
#[derive(Clone, Debug, Default, PartialEq, Serialize, Deserialize)]
pub(crate) struct DocumentIntelligenceInputParams {
pub pages: Option<PagesInput>,
pub features: Option<FeaturesInput>,
}
#[derive(Clone, Debug, PartialEq, Serialize)]
pub(crate) struct DocumentIntelligenceParams {
pub pages: Option<String>,
pub features: Option<String>,
}
#[derive(Clone, Debug, Serialize, Deserialize)]
pub(crate) enum DocumentIntelligenceRequest {
#[serde(rename = "urlSource")]
UrlSource(String),
#[serde(rename = "base64Source")]
Base64Source(String),
}
#[derive(Clone, Debug, PartialEq)]
pub(crate) enum OperationStatus {
Succeeded,
Running,
NotStarted,
Failed,
Unknown(String),
}
impl<'de> Deserialize<'de> for OperationStatus {
fn deserialize<D: Deserializer<'de>>(deserializer: D) -> Result<Self, D::Error> {
Ok(match String::deserialize(deserializer)?.as_str() {
"succeeded" => Self::Succeeded,
"running" => Self::Running,
"notStarted" => Self::NotStarted,
"failed" => Self::Failed,
value => Self::Unknown(value.to_string()),
})
}
}
impl std::fmt::Display for OperationStatus {
fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
formatter.write_str(match self {
Self::Succeeded => "succeeded",
Self::Running => "running",
Self::NotStarted => "notStarted",
Self::Failed => "failed",
Self::Unknown(value) => value,
})
}
}
#[derive(Clone, Debug, Deserialize)]
pub(crate) struct AzureDocumentIntelligenceOperation {
pub status: Option<OperationStatus>,
#[serde(rename = "analyzeResult")]
pub analyze_result: Option<AzureDocumentIntelligenceAnalyzeResult>,
}
#[derive(Clone, Debug, Default, Deserialize)]
pub(crate) struct AzureDocumentIntelligenceAnalyzeResult {
pub content: Option<String>,
#[serde(default)]
pub pages: Vec<AzureDocumentIntelligencePage>,
pub tables: Option<Vec<Map<String, Value>>>,
#[serde(rename = "keyValuePairs")]
pub key_value_pairs: Option<Vec<Map<String, Value>>>,
}
#[derive(Clone, Debug, Deserialize)]
pub(crate) struct AzureDocumentIntelligencePage {
#[serde(rename = "pageNumber", default, deserialize_with = "optional_i64")]
pub page_number: Option<i64>,
#[serde(default, deserialize_with = "optional_f64")]
pub width: Option<f64>,
#[serde(default, deserialize_with = "optional_f64")]
pub height: Option<f64>,
pub unit: Option<String>,
#[serde(default)]
pub lines: Vec<AzureDocumentIntelligenceLine>,
}
#[derive(Clone, Debug, Deserialize)]
pub(crate) struct AzureDocumentIntelligenceLine {
pub content: Option<String>,
}
fn optional_i64<'de, D: Deserializer<'de>>(deserializer: D) -> Result<Option<i64>, D::Error> {
match Option::<Value>::deserialize(deserializer)? {
None | Some(Value::Null) => Ok(None),
Some(Value::Number(number)) => number
.as_i64()
.map(Some)
.ok_or_else(|| serde::de::Error::custom("expected an integer")),
Some(Value::String(value)) => value
.parse::<i64>()
.map(Some)
.map_err(|_| serde::de::Error::custom("expected an integer")),
Some(_) => Err(serde::de::Error::custom("expected an integer")),
}
}
fn optional_f64<'de, D: Deserializer<'de>>(deserializer: D) -> Result<Option<f64>, D::Error> {
match Option::<Value>::deserialize(deserializer)? {
None | Some(Value::Null) => Ok(None),
Some(Value::Number(number)) => number
.as_f64()
.filter(|value| value.is_finite())
.map(Some)
.ok_or_else(|| serde::de::Error::custom("expected a finite number")),
Some(Value::String(value)) => value
.parse::<f64>()
.ok()
.filter(|value| value.is_finite())
.map(Some)
.ok_or_else(|| serde::de::Error::custom("expected a finite number")),
Some(_) => Err(serde::de::Error::custom("expected a number")),
}
}

View file

@ -36,6 +36,101 @@ mod tests {
use rstest::rstest;
use serde_json::{Value, json};
fn mapped_params(value: Value) -> Value {
serde_json::to_value(serde_json::from_value::<MistralOcrParams>(value).unwrap()).unwrap()
}
fn document() -> OcrDocument {
serde_json::from_value(
json!({"type":"document_url","document_url":"https://example.com/a.pdf"}),
)
.unwrap()
}
#[rstest]
fn extract_header_is_a_supported_ocr_param() {
assert_eq!(
mapped_params(json!({"extract_header":true}))["extract_header"],
true
);
}
#[rstest]
fn extract_footer_is_a_supported_ocr_param() {
assert_eq!(
mapped_params(json!({"extract_footer":false}))["extract_footer"],
false
);
}
#[rstest]
fn existing_ocr_params_remain_supported() {
let mapped = mapped_params(json!({
"pages":[0,2],
"include_image_base64":true,
"image_limit":2,
"image_min_size":100,
"bbox_annotation_format":{"type":"json_schema"},
"document_annotation_format":{"type":"json_schema"}
}));
assert_eq!(mapped["pages"], json!([0, 2]));
assert_eq!(mapped["include_image_base64"], true);
assert_eq!(mapped["image_limit"], 2);
assert_eq!(mapped["image_min_size"], 100);
assert_eq!(mapped["bbox_annotation_format"]["type"], "json_schema");
assert_eq!(mapped["document_annotation_format"]["type"], "json_schema");
}
#[rstest]
fn map_ocr_params_forwards_extract_header() {
assert_eq!(
mapped_params(json!({"extract_header":true}))["extract_header"],
true
);
}
#[rstest]
fn map_ocr_params_forwards_extract_footer() {
assert_eq!(
mapped_params(json!({"extract_footer":true}))["extract_footer"],
true
);
}
#[rstest]
fn map_ocr_params_forwards_extract_header_and_footer() {
let mapped = mapped_params(json!({"extract_header":true,"extract_footer":false}));
assert_eq!(mapped["extract_header"], true);
assert_eq!(mapped["extract_footer"], false);
}
#[rstest]
fn map_ocr_params_drops_unknown_params() {
let mapped = mapped_params(json!({"extract_header":true,"unsupported_param":"value"}));
assert_eq!(mapped["extract_header"], true);
assert!(mapped.get("unsupported_param").is_none());
}
#[rstest]
#[case("table_format", json!("html"))]
#[case("confidence_scores_granularity", json!("word"))]
#[case("document_annotation_prompt", json!("extract"))]
#[case("include_blocks", json!(true))]
#[case("id", json!("req-123"))]
fn new_ocr_params_are_supported(#[case] name: &str, #[case] value: Value) {
assert_eq!(mapped_params(json!({name:value.clone()}))[name], value);
}
#[rstest]
#[case("table_format", json!("html"))]
#[case("confidence_scores_granularity", json!("word"))]
#[case("document_annotation_prompt", json!("extract"))]
#[case("include_blocks", json!(true))]
#[case("id", json!("req-123"))]
fn map_ocr_params_forwards_new_ocr_params(#[case] name: &str, #[case] value: Value) {
assert_eq!(mapped_params(json!({name:value.clone()}))[name], value);
}
#[rstest]
#[case("pages", json!([0, 2]))]
#[case("include_image_base64", json!(true))]
@ -53,50 +148,81 @@ mod tests {
fn request_mapping_matches_python(#[case] name: &str, #[case] value: Value) {
let params: MistralOcrParams =
serde_json::from_value(json!({name: value.clone()})).unwrap();
let document: OcrDocument = serde_json::from_value(
json!({"type":"document_url","document_url":"https://example.com/a.pdf"}),
)
.unwrap();
let result =
serde_json::to_value(transform_ocr_request("model", document, &params).unwrap())
serde_json::to_value(transform_ocr_request("model", document(), &params).unwrap())
.unwrap();
assert_eq!(result["model"], "model");
assert_eq!(result[name], value);
}
#[test]
fn request_mapping_filters_unknown_fields() {
let params: MistralOcrParams = serde_json::from_value(json!({"unknown": true})).unwrap();
let document: OcrDocument = serde_json::from_value(
json!({"type":"document_url","document_url":"https://example.com/a.pdf"}),
#[rstest]
#[case("table_format", json!("html"))]
#[case("confidence_scores_granularity", json!("word"))]
#[case("document_annotation_prompt", json!("extract"))]
#[case("id", json!("req-123"))]
#[case("extract_header", json!(true))]
#[case("include_blocks", json!(true))]
#[case("pages", json!([0,1]))]
fn transform_ocr_request_includes_each_optional_param(
#[case] name: &str,
#[case] value: Value,
) {
let params: MistralOcrParams = serde_json::from_value(json!({name:value.clone()})).unwrap();
let result = serde_json::to_value(
transform_ocr_request("mistral-ocr-latest", document(), &params).unwrap(),
)
.unwrap();
let result =
serde_json::to_value(transform_ocr_request("model", document, &params).unwrap())
.unwrap();
assert!(result.get("unknown").is_none());
assert_eq!(result[name], value);
assert_eq!(result["model"], "mistral-ocr-latest");
}
#[test]
fn response_preserves_provider_fields() {
#[rstest]
fn transform_ocr_request_includes_multiple_new_params() {
let params: MistralOcrParams = serde_json::from_value(json!({
"table_format":"html",
"confidence_scores_granularity":"page",
"extract_header":true
}))
.unwrap();
let result = serde_json::to_value(
transform_ocr_request("mistral-ocr-latest", document(), &params).unwrap(),
)
.unwrap();
assert_eq!(result["table_format"], "html");
assert_eq!(result["confidence_scores_granularity"], "page");
assert_eq!(result["extract_header"], true);
}
#[rstest]
fn transform_ocr_response_preserves_blocks_and_confidence_scores() {
let response: MistralOcrResponse = serde_json::from_value(json!({
"pages":[{"index":0,"markdown":"hello","header":"head","confidence_scores":{"mean":0.99}}],
"pages":[{"index":0,"markdown":"hello","blocks":[{"type":"title"}],"confidence_scores":{"mean":0.99}}],
"model":"returned-model",
"usage_info":{"pages_processed":1,"future_counter":5},
"future_response_field":"kept"
"usage_info":{"pages_processed":1}
}))
.unwrap();
let result = transform_ocr_response("model", response)
.unwrap()
.into_json();
assert_eq!(result["pages"][0]["header"], "head");
assert_eq!(result["usage_info"]["future_counter"], 5);
assert_eq!(result["future_response_field"], "kept");
assert_eq!(result["model"], "returned-model");
assert_eq!(result["pages"][0]["blocks"][0]["type"], "title");
assert_eq!(result["pages"][0]["confidence_scores"]["mean"], 0.99);
}
#[test]
fn response_rejects_null_pages() {
assert!(serde_json::from_value::<MistralOcrResponse>(json!({"pages":null})).is_err());
#[rstest]
fn transform_ocr_response_preserves_ocr4_page_fields() {
let page = json!({
"index":0,
"markdown":"table page",
"tables":[{"rows":2,"cols":3}],
"hyperlinks":["https://example.com"],
"header":"header",
"footer":"footer"
});
let response: MistralOcrResponse =
serde_json::from_value(json!({"pages":[page.clone()]})).unwrap();
let result = transform_ocr_response("model", response)
.unwrap()
.into_json();
assert_eq!(result["pages"][0], page);
}
}

View file

@ -1 +1,4 @@
pub(crate) mod deepseek;
pub(crate) mod document_intelligence;
pub(crate) mod mistral;
pub(crate) mod reducto;

View file

@ -0,0 +1,9 @@
mod transformation;
mod types;
pub(crate) use transformation::{
transform_legacy_ocr_request, transform_ocr_response, transform_v3_ocr_request,
};
pub(crate) use types::{
ReductoLegacyParams, ReductoResponse, ReductoUploadResponse, ReductoV3Params,
};

View file

@ -0,0 +1,115 @@
use std::collections::BTreeMap;
use serde_json::{Value, json};
use super::types::*;
use crate::ocr::error::{OcrRequestError, OcrResponseError};
use crate::ocr::types::{LiteLLMOcrResponse, OcrDocument};
#[tracing::instrument(
name = "transform_ocr_request",
target = "litellm::function_trace",
level = "trace",
skip_all
)]
pub(crate) fn transform_v3_ocr_request(
_model: &str,
document: OcrDocument,
params: &ReductoV3Params,
) -> Result<ReductoV3Request, OcrRequestError> {
Ok(ReductoV3Request {
input: document.source().to_string(),
params: params.clone(),
})
}
#[tracing::instrument(
name = "transform_ocr_request",
target = "litellm::function_trace",
level = "trace",
skip_all
)]
pub(crate) fn transform_legacy_ocr_request(
_model: &str,
document: OcrDocument,
params: &ReductoLegacyParams,
) -> Result<ReductoLegacyRequest, OcrRequestError> {
Ok(ReductoLegacyRequest {
document_url: document.source().to_string(),
options: params.enhance.as_ref().map(|_| params.clone()),
})
}
pub(crate) fn transform_ocr_response(
model: &str,
response: ReductoResponse,
) -> Result<LiteLLMOcrResponse, OcrResponseError> {
let result = match response.result {
Some(result) => result.unwrap_or_default(),
None => ReductoResult {
chunks: response.chunks,
},
};
let usage = response.usage.unwrap_or_default();
Ok(LiteLLMOcrResponse {
pages: build_pages(result.chunks.unwrap_or_default()),
model: model.to_string(),
document_annotation: None,
usage_info: Some(json!({
"pages_processed": usage.num_pages,
"credits": usage.credits,
})),
object: "ocr".to_string(),
extra_fields: serde_json::Map::new(),
provider_native_response: None,
})
}
fn build_pages(chunks: Vec<ReductoChunk>) -> Vec<Value> {
let blocks_by_page = chunks
.iter()
.flat_map(|chunk| chunk.blocks.iter().flatten())
.filter_map(|block| block.bbox.as_ref()?.page.map(|page| (page, block)))
.fold(
BTreeMap::<i64, Vec<&ReductoBlock>>::new(),
|mut pages, (page, block)| {
pages.entry(page).or_default().push(block);
pages
},
);
if blocks_by_page.is_empty() {
let markdown = join_content(chunks.iter().map(|chunk| chunk.content.as_deref()));
return if markdown.is_empty() {
Vec::new()
} else {
vec![page(0, markdown, None)]
};
}
blocks_by_page
.into_iter()
.map(|(index, blocks)| {
let markdown = join_content(blocks.iter().map(|block| block.content.as_deref()));
page(
index.saturating_sub(1).max(0),
markdown,
Some(json!(blocks)),
)
})
.collect()
}
fn join_content<'a>(content: impl Iterator<Item = Option<&'a str>>) -> String {
content
.flatten()
.filter(|text| !text.is_empty())
.collect::<Vec<_>>()
.join("\n\n")
}
fn page(index: i64, markdown: String, blocks: Option<Value>) -> Value {
let mut result = json!({"index":index,"markdown":markdown,"images":null});
if let (Value::Object(fields), Some(blocks)) = (&mut result, blocks) {
fields.insert("blocks".into(), blocks);
}
result
}

View file

@ -0,0 +1,128 @@
use serde::{Deserialize, Deserializer, Serialize};
use serde_json::{Map, Value};
#[derive(Clone, Debug, Default, Serialize, Deserialize)]
pub(crate) struct ReductoV3Params {
#[serde(skip_serializing_if = "Option::is_none")]
pub formatting: Option<Map<String, Value>>,
#[serde(skip_serializing_if = "Option::is_none")]
pub retrieval: Option<Map<String, Value>>,
#[serde(skip_serializing_if = "Option::is_none")]
pub settings: Option<Map<String, Value>>,
}
#[derive(Clone, Debug, Default, Serialize, Deserialize)]
pub(crate) struct ReductoLegacyParams {
#[serde(skip_serializing_if = "Option::is_none")]
pub enhance: Option<Map<String, Value>>,
}
#[derive(Clone, Debug, Serialize, Deserialize)]
pub(crate) struct ReductoV3Request {
pub input: String,
#[serde(flatten)]
pub params: ReductoV3Params,
}
#[derive(Clone, Debug, Serialize, Deserialize)]
pub(crate) struct ReductoLegacyRequest {
pub document_url: String,
#[serde(skip_serializing_if = "Option::is_none")]
pub options: Option<ReductoLegacyParams>,
}
#[derive(Deserialize)]
pub(crate) struct ReductoUploadResponse {
pub file_id: Option<String>,
}
#[derive(Clone, Debug, Deserialize)]
pub(crate) struct ReductoResponse {
#[serde(default, deserialize_with = "present_nullable")]
pub result: Option<Option<ReductoResult>>,
pub usage: Option<ReductoUsage>,
#[serde(default)]
pub chunks: Option<Vec<ReductoChunk>>,
}
fn present_nullable<'de, D: Deserializer<'de>, T: Deserialize<'de>>(
deserializer: D,
) -> Result<Option<Option<T>>, D::Error> {
Option::<T>::deserialize(deserializer).map(Some)
}
#[derive(Clone, Debug, Default, Deserialize)]
pub(crate) struct ReductoResult {
pub chunks: Option<Vec<ReductoChunk>>,
}
#[derive(Clone, Debug, Default, Deserialize)]
pub(crate) struct ReductoUsage {
#[serde(default, deserialize_with = "optional_i64")]
pub num_pages: Option<i64>,
#[serde(default, deserialize_with = "optional_f64")]
pub credits: Option<f64>,
}
#[derive(Clone, Debug, Deserialize)]
pub(crate) struct ReductoChunk {
pub content: Option<String>,
pub blocks: Option<Vec<ReductoBlock>>,
}
#[derive(Clone, Debug, Serialize, Deserialize)]
pub(crate) struct ReductoBlock {
#[serde(skip_serializing_if = "Option::is_none")]
pub content: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
pub bbox: Option<ReductoBoundingBox>,
#[serde(flatten)]
pub extra_fields: Map<String, Value>,
}
#[derive(Clone, Debug, Serialize, Deserialize)]
pub(crate) struct ReductoBoundingBox {
#[serde(default, deserialize_with = "optional_i64")]
pub page: Option<i64>,
#[serde(flatten)]
pub extra_fields: Map<String, Value>,
}
fn optional_i64<'de, D: Deserializer<'de>>(deserializer: D) -> Result<Option<i64>, D::Error> {
match Option::<Value>::deserialize(deserializer)? {
None | Some(Value::Null) => Ok(None),
Some(Value::Number(number)) => number
.as_i64()
.or_else(|| number.as_f64().and_then(checked_truncated_i64))
.map(Some)
.ok_or_else(|| serde::de::Error::custom("expected an integer")),
Some(Value::String(value)) => value
.trim()
.parse::<i64>()
.map(Some)
.map_err(|_| serde::de::Error::custom("expected an integer")),
Some(Value::Bool(value)) => Ok(Some(i64::from(value))),
Some(_) => Ok(None),
}
}
fn optional_f64<'de, D: Deserializer<'de>>(deserializer: D) -> Result<Option<f64>, D::Error> {
match Option::<Value>::deserialize(deserializer)? {
None | Some(Value::Null) => Ok(None),
Some(Value::Number(number)) => number
.as_f64()
.map(Some)
.ok_or_else(|| serde::de::Error::custom("expected a number")),
Some(Value::String(value)) => value
.trim()
.parse::<f64>()
.map(Some)
.map_err(|_| serde::de::Error::custom("expected a number")),
Some(_) => Ok(None),
}
}
fn checked_truncated_i64(value: f64) -> Option<i64> {
(value.is_finite() && value >= i64::MIN as f64 && value <= i64::MAX as f64)
.then(|| value.trunc() as i64)
}

View file

@ -0,0 +1,212 @@
use base64::{Engine, engine::general_purpose::STANDARD};
use data_url::mime::Mime;
use data_url::{DataUrl, DataUrlError, forgiving_base64::DecodeError};
use reqwest::Url;
use super::error::{OcrError, OcrRequestError, OcrResponseError};
use super::types::{OcrConnection, OcrDocument};
use crate::constants::OCR_MAX_FETCH_REDIRECTS;
use crate::error::{MediaError, TransportError};
use crate::media::{DownloadPolicy, MediaFetcher};
pub(crate) struct InlineDocument<'a>(DataUrl<'a>);
impl<'a> InlineDocument<'a> {
pub(crate) fn parse(source: &'a str) -> Result<Option<Self>, OcrRequestError> {
match DataUrl::process(source) {
Ok(url) => Ok(Some(Self(url))),
Err(DataUrlError::NotADataUrl) => Ok(None),
Err(DataUrlError::NoComma) => Err(OcrRequestError::InvalidDataUri),
}
}
pub(crate) fn mime_type(&self) -> &Mime {
self.0.mime_type()
}
pub(crate) fn decode(&self, max_bytes: usize) -> Result<Vec<u8>, OcrRequestError> {
let mut body = Vec::new();
self.0
.decode(|bytes| {
if bytes.len() > max_bytes.saturating_sub(body.len()) {
return Err(OcrRequestError::InlineDocumentTooLarge);
}
body.extend_from_slice(bytes);
Ok(())
})
.map_err(|error| match error {
DecodeError::InvalidBase64(_) => OcrRequestError::InvalidDataUri,
DecodeError::WriteError(error) => error,
})?;
Ok(body)
}
}
pub(crate) fn validate_inline_document(document: &OcrDocument) -> Result<(), OcrRequestError> {
let inline =
InlineDocument::parse(document.source())?.ok_or(OcrRequestError::InvalidDataUri)?;
inline.decode(crate::constants::OCR_INLINE_MAX_BYTES)?;
Ok(())
}
pub(crate) async fn inline_remote_document(
fetcher: &MediaFetcher,
document: OcrDocument,
connection: &OcrConnection,
) -> Result<OcrDocument, OcrError> {
let source = document.source();
if !source.starts_with("http://") && !source.starts_with("https://") {
validate_inline_document(&document)?;
return Ok(document);
}
let url = Url::parse(source).map_err(|_| OcrRequestError::RequestField {
path: "document URL".into(),
})?;
let downloaded = fetcher
.fetch(
url,
DownloadPolicy {
timeout: connection.timeout,
max_bytes: connection.max_download_bytes,
max_redirects: OCR_MAX_FETCH_REDIRECTS,
},
)
.await
.map_err(map_media_error)?;
let result = document.with_source(format!(
"data:{};base64,{}",
downloaded.content_type,
STANDARD.encode(downloaded.bytes)
));
validate_inline_document(&result)?;
Ok(result)
}
fn map_media_error(error: MediaError) -> OcrError {
match error {
MediaError::BlockedUrl => OcrRequestError::BlockedDocumentUrl.into(),
MediaError::DownloadDisabled => OcrRequestError::DownloadDisabled.into(),
MediaError::DownloadTooLarge => OcrRequestError::DownloadTooLarge.into(),
MediaError::TooManyRedirects => OcrRequestError::TooManyRedirects.into(),
MediaError::MissingRedirectLocation => OcrResponseError::MissingRedirectLocation.into(),
MediaError::InvalidRedirect => OcrResponseError::InvalidRedirect.into(),
MediaError::Http(status) => TransportError::Http {
status,
body: "OCR document download failed".into(),
}
.into(),
MediaError::Timeout => {
TransportError::Network("OCR document download timed out".into()).into()
}
MediaError::Transport(error) => error.into(),
}
}
#[cfg(test)]
mod tests {
use super::*;
use serde_json::Map;
fn document(source: &str) -> OcrDocument {
OcrDocument::DocumentUrl {
document_url: source.into(),
extra_fields: Map::new(),
}
}
#[test]
fn decodes_data_urls_and_limits_decoded_size() {
for (source, expected) in [
("data:application/pdf;base64,YWJj", b"abc".as_slice()),
("DATA:application/pdf;BASE64,YWI", b"ab".as_slice()),
("data:,a%20b%00%FF", b"a b\0\xff".as_slice()),
] {
let inline = InlineDocument::parse(source).unwrap().unwrap();
assert_eq!(inline.decode(expected.len()).unwrap(), expected);
assert_eq!(
inline.decode(expected.len() - 1),
Err(OcrRequestError::InlineDocumentTooLarge)
);
}
}
#[test]
fn preserves_mime_parameters_and_standard_default() {
let inline = InlineDocument::parse("data:application/pdf;version=1.7;base64,YQ==")
.unwrap()
.unwrap();
assert!(inline.mime_type().matches("application", "pdf"));
assert_eq!(inline.mime_type().get_parameter("version"), Some("1.7"));
let default = InlineDocument::parse("data:,a").unwrap().unwrap();
assert!(default.mime_type().matches("text", "plain"));
assert_eq!(
default.mime_type().get_parameter("charset"),
Some("US-ASCII")
);
}
#[test]
fn rejects_invalid_inline_documents() {
for source in [
"https://example.com/document.pdf",
"data:application/pdf;base64",
"data:application/pdf;base64,INVALID!",
] {
assert!(validate_inline_document(&document(source)).is_err());
}
}
#[tokio::test]
async fn remote_conversion_preserves_kind_and_isolates_provider_credentials() {
use tokio::io::{AsyncReadExt, AsyncWriteExt};
use tokio::net::TcpListener;
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
let address = listener.local_addr().unwrap();
let server = tokio::spawn(async move {
let (mut socket, _) = listener.accept().await.unwrap();
let mut request = vec![0_u8; 2048];
let count = socket.read(&mut request).await.unwrap();
socket
.write_all(b"HTTP/1.1 200 OK\r\nContent-Type: image/png; charset=binary\r\nContent-Length: 3\r\nConnection: close\r\n\r\nabc")
.await
.unwrap();
String::from_utf8_lossy(&request[..count]).into_owned()
});
let mut provider_headers = reqwest::header::HeaderMap::new();
provider_headers.insert(
reqwest::header::AUTHORIZATION,
reqwest::header::HeaderValue::from_static("Bearer provider-secret"),
);
let provider_http = reqwest::Client::builder()
.default_headers(provider_headers)
.build()
.unwrap();
let document_http = reqwest::Client::builder()
.redirect(reqwest::redirect::Policy::none())
.build()
.unwrap();
let client = super::super::OcrClient::for_test(provider_http, document_http);
let converted = inline_remote_document(
client.document_fetcher(),
OcrDocument::ImageUrl {
image_url: format!("http://{address}/image"),
extra_fields: Map::from_iter([("detail".into(), serde_json::json!("high"))]),
},
&OcrConnection::default(),
)
.await
.unwrap();
let request = server.await.unwrap();
assert_eq!(
converted,
OcrDocument::ImageUrl {
image_url: "data:image/png;base64,YWJj".into(),
extra_fields: Map::from_iter([("detail".into(), serde_json::json!("high"))]),
}
);
assert!(!request.to_ascii_lowercase().contains("authorization"));
assert!(!request.contains("provider-secret"));
}
}

View file

@ -10,12 +10,52 @@ pub enum OcrRequestError {
RequestField { path: String },
#[error("missing required field: {0}")]
MissingField(&'static str),
#[error("invalid OCR document data URI")]
InvalidDataUri,
#[error("Reducto requires a reducto:// id or a data URI")]
ReductoSource,
#[error("inline OCR document exceeds the size limit")]
InlineDocumentTooLarge,
#[error("OCR document URL is blocked by network policy")]
BlockedDocumentUrl,
#[error("OCR document downloads are disabled")]
DownloadDisabled,
#[error("OCR document download exceeds the size limit")]
DownloadTooLarge,
#[error("OCR document download exceeded the redirect limit")]
TooManyRedirects,
#[error("invalid OCR pages: {0}")]
Pages(String),
#[error("invalid OCR features")]
Features,
#[error("OCR model cannot be a dot segment")]
DotModel,
}
#[derive(Debug, Clone, PartialEq, Eq, Error)]
pub enum OcrResponseError {
#[error("invalid OCR response field: {path}")]
ResponseField { path: String },
#[error("OCR response is missing non-empty content")]
EmptyContent,
#[error("OCR document redirect is missing a location")]
MissingRedirectLocation,
#[error("OCR document redirect location is invalid")]
InvalidRedirect,
#[error("OCR operation ended with status {0}")]
OperationStatus(String),
#[error("OCR response numeric value is out of range: {0}")]
NumericRange(&'static str),
}
#[derive(Debug, Clone, PartialEq, Eq, Error)]
pub enum OcrPollingError {
#[error("OCR accepted response is missing a valid operation-location")]
PollLocation,
#[error("OCR operation-location must use the submission origin without credentials")]
PollOrigin,
#[error("OCR polling timed out")]
PollTimeout,
}
#[derive(Debug, Error)]
@ -27,6 +67,8 @@ pub enum OcrError {
#[error("{0}")]
Transport(#[from] TransportError),
#[error("{0}")]
Polling(#[from] OcrPollingError),
#[error("{0}")]
Public(#[from] crate::Error),
}
@ -36,6 +78,7 @@ impl From<OcrError> for crate::Error {
OcrError::Request(error) => error.into(),
OcrError::Response(error) => error.into(),
OcrError::Transport(error) => error.into(),
OcrError::Polling(error) => crate::Error::InvalidResponse(error.to_string()),
OcrError::Public(error) => error,
}
}

View file

@ -1,21 +1,39 @@
mod adapters;
pub mod client;
mod codecs;
mod document;
pub mod error;
mod handler;
pub mod hooks;
mod prepare;
mod registry;
pub mod transformation;
pub mod types;
pub mod wire;
pub use client::{OcrClient, ocr};
pub use types::{LiteLLMOcrRequest, LiteLLMOcrResponse, OcrConnection, OcrDocument};
#[cfg(test)]
#[path = "../../tests/azure_ai_ocr.rs"]
mod azure_ai_tests;
#[cfg(test)]
#[path = "../../tests/azure_document_intelligence_ocr.rs"]
mod azure_document_intelligence_tests;
#[cfg(test)]
#[path = "../../tests/deepseek_ocr.rs"]
mod deepseek_tests;
#[cfg(test)]
#[path = "../../tests/reducto_ocr.rs"]
mod reducto_tests;
#[cfg(test)]
#[path = "../../tests/ocr/support.rs"]
pub(crate) mod test_support;
#[cfg(test)]
#[path = "../../tests/ocr.rs"]
pub(crate) mod tests;
#[cfg(test)]
#[path = "../../tests/vertex_ai_deepseek_ocr.rs"]
mod vertex_ai_deepseek_tests;
#[cfg(test)]
#[path = "../../tests/vertex_ai_ocr.rs"]
mod vertex_ai_tests;

View file

@ -4,7 +4,7 @@ use serde_json::{Map, Value};
use super::OcrClient;
use super::error::{OcrError, OcrRequestError};
use super::hooks::OcrDuringCallRequest;
use super::types::LiteLLMOcrRequest;
use super::types::{LiteLLMOcrRequest, OcrDocument};
#[derive(Debug, Deserialize)]
pub(crate) struct ParsedProviderParams<T> {
@ -24,6 +24,39 @@ pub(crate) fn _prepare_ocr_request<T: DeserializeOwned>(
)
}
pub(crate) fn merge_extra_params<B: Serialize>(
body: &B,
extra_params: Map<String, Value>,
) -> Result<Value, OcrRequestError> {
let Value::Object(fields) =
serde_json::to_value(body).map_err(|_| OcrRequestError::RequestField {
path: "body".into(),
})?
else {
return Err(OcrRequestError::RequestField {
path: "body".into(),
});
};
let extra_body = extra_params
.get("extra_body")
.and_then(Value::as_object)
.cloned()
.unwrap_or_default()
.into_iter()
.collect::<Map<String, Value>>();
Ok(Value::Object(
fields
.into_iter()
.chain(
extra_params
.into_iter()
.filter(|(name, _)| name != "extra_body"),
)
.chain(extra_body)
.collect(),
))
}
pub(crate) async fn transform_request_body<B>(
client: &OcrClient,
request: &LiteLLMOcrRequest,
@ -77,6 +110,29 @@ pub(crate) fn build_http_request<B: Serialize>(
.map_err(OcrError::from)
}
pub(crate) async fn guardrail_document(
request: &LiteLLMOcrRequest,
url: &str,
) -> Result<OcrDocument, OcrError> {
if !request.hooks.has_guardrails() {
return Ok(request.document.clone());
}
let changed = request
.hooks
.during_call(OcrDuringCallRequest {
model: request.model.clone(),
custom_llm_provider: request.adapter.provider().as_str().into(),
url: url.into(),
body: serde_json::to_value(&request.document).map_err(|_| {
OcrRequestError::RequestField {
path: "document".into(),
}
})?,
})
.await?;
super::wire::decode_request_value(changed.body, "guardrail.document").map_err(OcrError::from)
}
#[derive(Serialize)]
struct OcrWireBody<B> {
#[serde(flatten)]
@ -107,7 +163,6 @@ impl<B: Serialize + DeserializeOwned> OcrWireBody<B> {
pub(crate) fn credential_env(name: &str) -> Option<String> {
std::env::var(name).ok()
}
#[cfg(test)]
mod tests {
use serde_json::json;

View file

@ -24,12 +24,18 @@ super::adapters::for_each_ocr_adapter!(define_adapter_types);
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub(crate) enum OcrProvider {
Mistral,
AzureAi,
Reducto,
VertexAi,
}
impl OcrProvider {
pub(crate) const fn as_str(self) -> &'static str {
match self {
Self::Mistral => "mistral",
Self::AzureAi => "azure_ai",
Self::Reducto => "reducto",
Self::VertexAi => "vertex_ai",
}
}
}
@ -45,9 +51,78 @@ pub(crate) fn resolve_wire_adapter(
});
let typed_provider = match provider.custom_llm_provider {
"mistral" => OcrProvider::Mistral,
"azure_ai" => OcrProvider::AzureAi,
"reducto" => OcrProvider::Reducto,
"vertex_ai" => OcrProvider::VertexAi,
value => return Err(Error::InvalidProvider(value.to_string())),
};
match typed_provider {
OcrProvider::Mistral => Ok((provider.model.to_string(), OcrAdapterKind::Mistral)),
let adapter = match typed_provider {
OcrProvider::Mistral => OcrAdapterKind::Mistral,
OcrProvider::AzureAi if is_document_intelligence_model(provider.model) => {
OcrAdapterKind::AzureDocumentIntelligence
}
OcrProvider::AzureAi => OcrAdapterKind::AzureMistral,
OcrProvider::Reducto if provider.model.eq_ignore_ascii_case("parse-legacy") => {
OcrAdapterKind::ReductoLegacy
}
OcrProvider::Reducto if provider.model.eq_ignore_ascii_case("parse-v3") => {
OcrAdapterKind::ReductoV3
}
OcrProvider::Reducto => {
return Err(Error::InvalidRequest(format!(
"unsupported Reducto OCR model: {}",
provider.model
)));
}
OcrProvider::VertexAi if provider.model.to_ascii_lowercase().contains("deepseek") => {
OcrAdapterKind::VertexDeepSeek
}
OcrProvider::VertexAi => OcrAdapterKind::VertexMistral,
};
Ok((provider.model.to_string(), adapter))
}
fn is_document_intelligence_model(model: &str) -> bool {
let model = model.to_ascii_lowercase();
model.contains("doc-intelligence") || model.contains("documentintelligence")
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn provider_models_are_preserved_without_a_local_allowlist() {
let cases = [
("mistral/future-ocr-model", OcrAdapterKind::Mistral),
("azure_ai/future-ocr-model", OcrAdapterKind::AzureMistral),
];
for (qualified_model, expected_adapter) in cases {
let expected_model = qualified_model.split_once('/').unwrap().1;
let (model, adapter) = resolve_wire_adapter(qualified_model, None).unwrap();
assert_eq!(model, expected_model);
assert_eq!(adapter, expected_adapter);
}
}
#[test]
fn unknown_reducto_models_are_rejected() {
assert!(matches!(
resolve_wire_adapter("reducto/future-parse-model", None),
Err(Error::InvalidRequest(_))
));
}
#[test]
fn known_protocol_models_still_select_specialized_adapters() {
let (model, adapter) = resolve_wire_adapter("reducto/parse-legacy", None).unwrap();
assert_eq!(model, "parse-legacy");
assert_eq!(adapter, OcrAdapterKind::ReductoLegacy);
let (model, adapter) =
resolve_wire_adapter("azure_ai/doc-intelligence/prebuilt-layout", None).unwrap();
assert_eq!(model, "doc-intelligence/prebuilt-layout");
assert_eq!(adapter, OcrAdapterKind::AzureDocumentIntelligence);
}
}

View file

@ -1,107 +0,0 @@
use crate::Error;
use serde_json::{Map, Value};
use super::types::{LiteLLMOcrResponse, OcrRequestData};
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum OcrAuthStrategy {
Bearer,
Header(&'static str),
}
impl OcrAuthStrategy {
pub fn header_name(self) -> &'static str {
match self {
Self::Bearer => "authorization",
Self::Header(header_name) => header_name,
}
}
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum OcrResponseHandling {
Json,
AzureDocumentIntelligencePoll,
}
pub trait OcrProviderConfig: Sync {
fn supported_ocr_params(&self) -> &'static [&'static str];
#[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)]
fn map_ocr_params(&self, non_default_params: &Map<String, Value>) -> Map<String, Value> {
let mut mapped_params = Map::new();
for (param, value) in non_default_params {
if self.supported_ocr_params().contains(&param.as_str()) {
mapped_params.insert(param.clone(), value.clone());
}
}
mapped_params
}
fn transform_ocr_request(
&self,
model: &str,
document: Value,
optional_params: Map<String, Value>,
) -> Result<OcrRequestData, Error>;
fn transform_ocr_response(
&self,
model: &str,
response_json: Value,
) -> Result<LiteLLMOcrResponse, Error>;
fn transform_ocr_response_with_params(
&self,
model: &str,
response_json: Value,
_optional_params: &Map<String, Value>,
) -> Result<LiteLLMOcrResponse, Error> {
self.transform_ocr_response(model, response_json)
}
fn complete_url(
&self,
api_base: Option<&str>,
model: &str,
optional_params: &Map<String, Value>,
env_lookup: &dyn Fn(&str) -> Option<String>,
) -> Result<String, Error>;
fn resolve_api_key(
&self,
api_key: Option<&str>,
env_lookup: &dyn Fn(&str) -> Option<String>,
) -> Result<String, Error>;
#[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)]
fn validate_environment(
&self,
headers: Vec<(String, String)>,
api_key: Option<&str>,
env_lookup: &dyn Fn(&str) -> Option<String>,
) -> Result<Vec<(String, String)>, Error> {
let strategy = self.auth_strategy();
if crate::http_utils::has_header(&headers, strategy.header_name()) {
return Ok(headers);
}
let api_key = self.resolve_api_key(api_key, env_lookup)?;
let auth_header = match strategy {
OcrAuthStrategy::Bearer => ("Authorization".to_string(), format!("Bearer {api_key}")),
OcrAuthStrategy::Header(name) => (name.to_string(), api_key),
};
Ok(std::iter::once(auth_header).chain(headers).collect())
}
fn auth_strategy(&self) -> OcrAuthStrategy {
OcrAuthStrategy::Bearer
}
fn requires_data_uri_document(&self) -> bool {
false
}
fn response_handling(&self) -> OcrResponseHandling {
OcrResponseHandling::Json
}
}

View file

@ -1,3 +1,4 @@
use std::collections::BTreeMap;
use std::sync::Arc;
use std::time::Duration;
@ -7,14 +8,9 @@ use serde_json::{Map, Value};
use super::hooks::{NoopOcrHooks, OcrHooks};
use super::registry::{OcrAdapterKind, resolve_wire_adapter};
use crate::Error;
use crate::auth::InputSource;
use crate::constants::OCR_HTTP_TIMEOUT_SECS;
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
pub struct OcrRequestData {
pub data: Value,
pub files: Option<Value>,
}
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
#[serde(tag = "type")]
pub enum OcrDocument {
@ -32,6 +28,28 @@ pub enum OcrDocument {
},
}
impl OcrDocument {
pub(crate) fn source(&self) -> &str {
match self {
Self::DocumentUrl { document_url, .. } => document_url,
Self::ImageUrl { image_url, .. } => image_url,
}
}
pub(crate) fn with_source(self, source: String) -> Self {
match self {
Self::DocumentUrl { extra_fields, .. } => Self::DocumentUrl {
document_url: source,
extra_fields,
},
Self::ImageUrl { extra_fields, .. } => Self::ImageUrl {
image_url: source,
extra_fields,
},
}
}
}
#[derive(Clone, Copy, Debug, Default, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "lowercase")]
pub enum OcrResponseFormat {
@ -43,18 +61,28 @@ pub enum OcrResponseFormat {
#[derive(Clone)]
pub struct OcrConnection {
pub api_key: Option<String>,
pub api_key_source: InputSource,
pub api_base: Option<String>,
pub api_base_source: InputSource,
pub extra_headers: Vec<(String, String)>,
pub extra_headers_source: InputSource,
pub timeout: Duration,
pub max_download_bytes: u64,
pub poll_timeout: Duration,
}
impl Default for OcrConnection {
fn default() -> Self {
Self {
api_key: None,
api_key_source: InputSource::Deployment,
api_base: None,
api_base_source: InputSource::Deployment,
extra_headers: Vec::new(),
extra_headers_source: InputSource::Deployment,
timeout: Duration::from_secs(OCR_HTTP_TIMEOUT_SECS),
max_download_bytes: crate::constants::OCR_DOWNLOAD_MAX_BYTES,
poll_timeout: Duration::from_secs(crate::constants::OCR_POLL_TIMEOUT_SECS),
}
}
}
@ -66,6 +94,7 @@ pub struct LiteLLMOcrRequest {
pub hooks: Arc<dyn OcrHooks>,
pub litellm_call_id: Option<String>,
pub optional_params: Map<String, Value>,
pub input_sources: BTreeMap<String, InputSource>,
pub(crate) adapter: OcrAdapterKind,
}
@ -85,6 +114,7 @@ impl LiteLLMOcrRequest {
hooks: Arc::new(NoopOcrHooks),
litellm_call_id: None,
optional_params,
input_sources: BTreeMap::new(),
adapter: adapter_kind,
})
}

View file

@ -1,10 +1,12 @@
use crate::ocr::error::OcrRequestError;
use crate::ocr::error::OcrResponseError;
use std::collections::BTreeMap;
use std::time::Duration;
use super::hooks::{OcrDuringCallRequest, OcrPreCallRequest};
use super::types::{LiteLLMOcrRequest, OcrConnection, OcrDocument};
use crate::Error;
use crate::auth::InputSource;
use serde::{
Deserialize,
de::{DeserializeOwned, IntoDeserializer},
@ -28,6 +30,8 @@ pub struct OcrWireRequest {
pub extra_headers: Option<Map<String, Value>>,
#[serde(default)]
pub optional_params: Map<String, Value>,
#[serde(default)]
pub input_sources: BTreeMap<String, InputSource>,
pub timeout_seconds: Option<f64>,
}
@ -36,6 +40,9 @@ pub fn is_supported_request(model: &str, custom_llm_provider: Option<&str>) -> b
}
pub fn decode_request(wire: OcrWireRequest) -> Result<LiteLLMOcrRequest, Error> {
let api_key_source = source_for(&wire.input_sources, "api_key");
let api_base_source = source_for(&wire.input_sources, "api_base");
let extra_headers_source = source_for(&wire.input_sources, "extra_headers");
let document = decode_request_value(wire.document, "document")?;
let headers = wire
.extra_headers
@ -67,16 +74,26 @@ pub fn decode_request(wire: OcrWireRequest) -> Result<LiteLLMOcrRequest, Error>
)?;
let connection = OcrConnection {
api_key: nonblank(wire.api_key),
api_key_source,
api_base: nonblank(wire.api_base),
api_base_source,
extra_headers: headers,
extra_headers_source,
timeout: timeout.unwrap_or(defaults.timeout),
max_download_bytes: defaults.max_download_bytes,
poll_timeout: defaults.poll_timeout,
};
Ok(LiteLLMOcrRequest {
connection,
input_sources: wire.input_sources,
..request
})
}
fn source_for(sources: &BTreeMap<String, InputSource>, name: &str) -> InputSource {
sources.get(name).copied().unwrap_or_default()
}
fn nonblank(value: Option<String>) -> Option<String> {
value
.map(|s| s.trim().to_string())

View file

@ -1,3 +1,4 @@
use crate::auth::error::MissingCredential;
use crate::error::Error;
use crate::messages::transformation::{AnthropicMessagesProviderConfig, MessagesAuthStrategy};
@ -21,13 +22,7 @@ pub fn resolve_anthropic_api_key(
non_empty(api_key)
.map(str::to_string)
.or_else(|| env_lookup(ANTHROPIC_API_KEY_ENV).filter(|value| !value.trim().is_empty()))
.ok_or_else(|| {
Error::Auth(
"Missing Anthropic API Key - Set `api_key` or the ANTHROPIC_API_KEY \
environment variable"
.to_string(),
)
})
.ok_or_else(|| Error::from(crate::AuthError::from(MissingCredential::AnthropicApiKey)))
}
pub fn complete_anthropic_url(

View file

@ -0,0 +1,43 @@
use std::future::Future;
use std::sync::Arc;
use azure_core::credentials::TokenCredential;
use moka::future::Cache;
use crate::AuthError;
#[derive(Clone, Debug, PartialEq, Eq, Hash)]
pub(crate) struct AzureCredentialProviderCacheKey {
pub(crate) mechanism: &'static str,
pub(crate) authority: String,
pub(crate) tenant_id: String,
pub(crate) client_id: String,
pub(crate) scope: String,
pub(crate) secret_identity: String,
}
pub(crate) struct AzureCredentialProviderCache {
entries: Cache<AzureCredentialProviderCacheKey, Arc<dyn TokenCredential>>,
}
impl AzureCredentialProviderCache {
pub(crate) fn new(capacity: u64) -> Self {
Self {
entries: Cache::builder().max_capacity(capacity).build(),
}
}
pub(crate) async fn get_or_create<F>(
&self,
key: AzureCredentialProviderCacheKey,
create: F,
) -> Result<Arc<dyn TokenCredential>, AuthError>
where
F: Future<Output = Result<Arc<dyn TokenCredential>, AuthError>>,
{
self.entries
.try_get_with(key, create)
.await
.map_err(|error| (*error).clone())
}
}

View file

@ -0,0 +1,7 @@
mod credential_provider_cache;
mod native;
mod resolve;
mod types;
pub(crate) use resolve::AzureAuthService;
pub(crate) use types::AzureAuthInputs;

View file

@ -0,0 +1,702 @@
use crate::auth::error::AuthConfigurationError;
use std::sync::Arc;
use std::time::{Duration, UNIX_EPOCH};
use azure_core::cloud::{CloudConfiguration, CustomConfiguration};
use azure_core::credentials::{Secret, TokenCredential};
use azure_core::http::ClientOptions;
use azure_identity::{
ClientAssertion, ClientAssertionCredential, ClientAssertionCredentialOptions,
ClientSecretCredential, ClientSecretCredentialOptions, DeveloperToolsCredential,
ManagedIdentityCredential, ManagedIdentityCredentialOptions, UserAssignedId,
WorkloadIdentityCredential, WorkloadIdentityCredentialOptions,
};
use sha2::{Digest, Sha256};
use crate::AuthError;
use crate::auth::{InputSource, ResolvedCredential, SecretValue, Sourced};
use super::credential_provider_cache::{
AzureCredentialProviderCache, AzureCredentialProviderCacheKey,
};
#[derive(Clone, Debug)]
pub(crate) enum NativeAzureRequest {
ClientSecret {
tenant_id: Sourced<String>,
client_id: Sourced<String>,
client_secret: Sourced<SecretValue>,
scope: Sourced<String>,
authority: Option<Sourced<String>>,
},
ClientAssertion {
tenant_id: Sourced<String>,
client_id: Sourced<String>,
assertion: Sourced<SecretValue>,
assertion_identity: String,
scope: Sourced<String>,
authority: Option<Sourced<String>>,
},
WorkloadIdentity {
tenant_id: Sourced<String>,
client_id: Sourced<String>,
token_file_path: Sourced<String>,
scope: Sourced<String>,
authority: Option<Sourced<String>>,
},
ManagedIdentity {
client_id: Option<Sourced<String>>,
scope: Sourced<String>,
selection_source: InputSource,
},
DeveloperTools {
scope: Sourced<String>,
selection_source: InputSource,
},
}
#[derive(Clone, Debug)]
pub(crate) struct ValidatedAzureRequest {
request: NativeAzureRequest,
credential_source: InputSource,
}
impl ValidatedAzureRequest {
pub(crate) fn new(request: NativeAzureRequest) -> Result<Self, AuthError> {
validate_authority(&request)?;
let credential_source = validate_sources(&request)?;
Ok(Self {
request,
credential_source,
})
}
pub(crate) fn credential_source(&self) -> InputSource {
self.credential_source
}
#[cfg(test)]
pub(super) fn kind(&self) -> &'static str {
match self.request {
NativeAzureRequest::ClientSecret { .. } => "client-secret",
NativeAzureRequest::ClientAssertion { .. } => "client-assertion",
NativeAzureRequest::WorkloadIdentity { .. } => "workload-identity",
NativeAzureRequest::ManagedIdentity { .. } => "managed-identity",
NativeAzureRequest::DeveloperTools { .. } => "developer-tools",
}
}
}
pub(crate) struct NativeAzureTokenAcquirer {
cache: AzureCredentialProviderCache,
transport: Option<azure_core::http::Transport>,
}
impl Default for NativeAzureTokenAcquirer {
fn default() -> Self {
Self::new(64)
}
}
impl NativeAzureTokenAcquirer {
pub(crate) fn new(cache_capacity: u64) -> Self {
Self {
cache: AzureCredentialProviderCache::new(cache_capacity),
transport: None,
}
}
#[cfg(test)]
pub(super) fn with_transport(
cache_capacity: u64,
transport: azure_core::http::Transport,
) -> Self {
Self {
cache: AzureCredentialProviderCache::new(cache_capacity),
transport: Some(transport),
}
}
pub(crate) async fn acquire(
&self,
request: ValidatedAzureRequest,
) -> Result<ResolvedCredential, AuthError> {
let scope = request.request.scope().to_string();
let key = request.request.cache_key();
let transport = self.transport.clone();
let credential = self
.cache
.get_or_create(
key,
async move { build_credential(request.request, transport) },
)
.await?;
let token = credential
.get_token(&[scope.as_str()], None)
.await
.map_err(|error| AuthError::AzureTokenAcquisition(error.to_string()))?;
let expires_on = u64::try_from(token.expires_on.unix_timestamp())
.ok()
.map(|seconds| UNIX_EPOCH + Duration::from_secs(seconds));
Ok(ResolvedCredential::AccessToken {
token: SecretValue::new(token.token.secret()),
expires_on,
})
}
}
impl NativeAzureRequest {
fn scope(&self) -> &str {
match self {
Self::ClientSecret { scope, .. }
| Self::ClientAssertion { scope, .. }
| Self::WorkloadIdentity { scope, .. }
| Self::ManagedIdentity { scope, .. }
| Self::DeveloperTools { scope, .. } => scope.value(),
}
}
fn cache_key(&self) -> AzureCredentialProviderCacheKey {
match self {
Self::ClientSecret {
tenant_id,
client_id,
client_secret,
scope,
authority,
} => AzureCredentialProviderCacheKey {
mechanism: "client-secret",
authority: authority
.as_ref()
.map(|value| value.value().clone())
.unwrap_or_default(),
tenant_id: tenant_id.value().clone(),
client_id: client_id.value().clone(),
scope: scope.value().clone(),
secret_identity: secret_digest(client_secret.value().expose()),
},
Self::ClientAssertion {
tenant_id,
client_id,
assertion,
assertion_identity,
scope,
authority,
} => AzureCredentialProviderCacheKey {
mechanism: "client-assertion",
authority: authority
.as_ref()
.map(|value| value.value().clone())
.unwrap_or_default(),
tenant_id: tenant_id.value().clone(),
client_id: client_id.value().clone(),
scope: scope.value().clone(),
secret_identity: format!(
"{assertion_identity}:{}",
secret_digest(assertion.value().expose())
),
},
Self::WorkloadIdentity {
tenant_id,
client_id,
token_file_path,
scope,
authority,
} => AzureCredentialProviderCacheKey {
mechanism: "workload-identity",
authority: authority
.as_ref()
.map(|value| value.value().clone())
.unwrap_or_default(),
tenant_id: tenant_id.value().clone(),
client_id: client_id.value().clone(),
scope: scope.value().clone(),
secret_identity: token_file_path.value().clone(),
},
Self::ManagedIdentity {
client_id, scope, ..
} => AzureCredentialProviderCacheKey {
mechanism: "managed-identity",
authority: String::new(),
tenant_id: String::new(),
client_id: client_id
.as_ref()
.map(|value| value.value().clone())
.unwrap_or_default(),
scope: scope.value().clone(),
secret_identity: String::new(),
},
Self::DeveloperTools { scope, .. } => AzureCredentialProviderCacheKey {
mechanism: "developer-tools",
authority: String::new(),
tenant_id: String::new(),
client_id: String::new(),
scope: scope.value().clone(),
secret_identity: String::new(),
},
}
}
}
fn validate_authority(request: &NativeAzureRequest) -> Result<(), AuthError> {
let authority = match request {
NativeAzureRequest::ClientSecret { authority, .. }
| NativeAzureRequest::ClientAssertion { authority, .. }
| NativeAzureRequest::WorkloadIdentity { authority, .. } => authority.as_ref(),
NativeAzureRequest::ManagedIdentity { .. } | NativeAzureRequest::DeveloperTools { .. } => {
None
}
};
let Some(authority) = authority else {
return Ok(());
};
let url = url::Url::parse(authority.value())
.map_err(|_| AuthError::Configuration(AuthConfigurationError::InvalidAzureAuthority))?;
if url.scheme() != "https"
|| url.host_str().is_none()
|| !url.username().is_empty()
|| url.password().is_some()
|| url.query().is_some()
|| url.fragment().is_some()
|| !matches!(url.path(), "" | "/")
{
return Err(AuthError::Configuration(
AuthConfigurationError::InvalidAzureAuthority,
));
}
Ok(())
}
fn validate_sources(request: &NativeAzureRequest) -> Result<InputSource, AuthError> {
match request {
NativeAzureRequest::ClientSecret {
tenant_id,
client_id,
client_secret,
scope,
authority,
} => {
let identity_sources = [
tenant_id.source(),
client_id.source(),
client_secret.source(),
];
let request_identity = identity_sources.contains(&InputSource::Request);
if request_identity
&& !identity_sources
.iter()
.all(|source| *source == InputSource::Request)
{
return mixed_sources();
}
if !request_identity && is_request_controlled(scope, authority.as_ref()) {
return mixed_sources();
}
Ok(if request_identity {
InputSource::Request
} else {
trusted_source(&identity_sources)
})
}
NativeAzureRequest::ClientAssertion {
tenant_id,
client_id,
assertion,
scope,
authority,
..
} => trusted_only(&[
tenant_id.source(),
client_id.source(),
assertion.source(),
scope.source(),
authority
.as_ref()
.map(Sourced::source)
.unwrap_or(InputSource::Environment),
]),
NativeAzureRequest::WorkloadIdentity {
tenant_id,
client_id,
token_file_path,
scope,
authority,
} => trusted_only(&[
tenant_id.source(),
client_id.source(),
token_file_path.source(),
scope.source(),
authority
.as_ref()
.map(Sourced::source)
.unwrap_or(InputSource::Environment),
]),
NativeAzureRequest::ManagedIdentity {
client_id,
scope,
selection_source,
} => trusted_only(&[
client_id
.as_ref()
.map(Sourced::source)
.unwrap_or(InputSource::Environment),
scope.source(),
*selection_source,
]),
NativeAzureRequest::DeveloperTools {
scope,
selection_source,
} => trusted_only(&[scope.source(), *selection_source]),
}
}
fn is_request_controlled<T>(value: &Sourced<T>, optional: Option<&Sourced<String>>) -> bool {
value.source() == InputSource::Request
|| optional.is_some_and(|value| value.source() == InputSource::Request)
}
fn trusted_only(sources: &[InputSource]) -> Result<InputSource, AuthError> {
if sources.contains(&InputSource::Request) {
return mixed_sources();
}
Ok(trusted_source(sources))
}
fn trusted_source(sources: &[InputSource]) -> InputSource {
if sources.contains(&InputSource::Deployment) {
InputSource::Deployment
} else {
InputSource::Environment
}
}
fn mixed_sources<T>() -> Result<T, AuthError> {
Err(AuthError::Configuration(
AuthConfigurationError::MixedAzureCredentialSources,
))
}
fn build_credential(
request: NativeAzureRequest,
transport: Option<azure_core::http::Transport>,
) -> Result<Arc<dyn TokenCredential>, AuthError> {
match request {
NativeAzureRequest::ClientSecret {
tenant_id,
client_id,
client_secret,
authority,
..
} => ClientSecretCredential::new(
tenant_id.value(),
client_id.into_value(),
Secret::new(client_secret.value().expose().to_string()),
Some(ClientSecretCredentialOptions {
client_options: client_options(authority.map(Sourced::into_value), transport),
}),
)
.map(|credential| credential as Arc<dyn TokenCredential>),
NativeAzureRequest::ClientAssertion {
tenant_id,
client_id,
assertion,
authority,
..
} => ClientAssertionCredential::new(
tenant_id.into_value(),
client_id.into_value(),
StaticAssertion(assertion.into_value()),
Some(ClientAssertionCredentialOptions {
client_options: client_options(authority.map(Sourced::into_value), transport),
}),
)
.map(|credential| credential as Arc<dyn TokenCredential>),
NativeAzureRequest::WorkloadIdentity {
tenant_id,
client_id,
token_file_path,
authority,
..
} => WorkloadIdentityCredential::new(Some(WorkloadIdentityCredentialOptions {
credential_options: azure_identity::ClientAssertionCredentialOptions {
client_options: client_options(authority.map(Sourced::into_value), transport),
},
client_id: Some(client_id.into_value()),
tenant_id: Some(tenant_id.into_value()),
token_file_path: Some(token_file_path.into_value().into()),
}))
.map(|credential| credential as Arc<dyn TokenCredential>),
NativeAzureRequest::ManagedIdentity { client_id, .. } => {
ManagedIdentityCredential::new(Some(ManagedIdentityCredentialOptions {
user_assigned_id: client_id
.map(Sourced::into_value)
.map(UserAssignedId::ClientId),
client_options: client_options(None, transport),
}))
.map(|credential| credential as Arc<dyn TokenCredential>)
}
NativeAzureRequest::DeveloperTools { .. } => DeveloperToolsCredential::new(None)
.map(|credential| credential as Arc<dyn TokenCredential>),
}
.map_err(|error| {
AuthError::Configuration(AuthConfigurationError::AzureCredentialInitialization(
error.to_string(),
))
})
}
fn client_options(
authority: Option<String>,
transport: Option<azure_core::http::Transport>,
) -> ClientOptions {
let cloud = authority.map(|authority_host| {
let mut custom = CustomConfiguration::default();
custom.authority_host = authority_host;
Arc::new(CloudConfiguration::from(custom))
});
ClientOptions {
cloud,
transport,
..Default::default()
}
}
fn secret_digest(secret: &str) -> String {
format!("{:x}", Sha256::digest(secret.as_bytes()))
}
#[derive(Debug)]
struct StaticAssertion(SecretValue);
impl ClientAssertion for StaticAssertion {
fn secret<'life0, 'life1, 'async_trait>(
&'life0 self,
_options: Option<azure_core::http::ClientMethodOptions<'life1>>,
) -> std::pin::Pin<
Box<dyn std::future::Future<Output = azure_core::Result<String>> + Send + 'async_trait>,
>
where
'life0: 'async_trait,
'life1: 'async_trait,
Self: 'async_trait,
{
Box::pin(async move { Ok(self.0.expose().to_string()) })
}
}
#[cfg(test)]
mod tests {
use std::sync::{Arc, Mutex};
use azure_core::http::headers::Headers;
use azure_core::http::{AsyncRawResponse, HttpClient, Request, StatusCode, Transport};
use azure_core::{Bytes, Result};
use super::{NativeAzureRequest, NativeAzureTokenAcquirer, ValidatedAzureRequest};
use crate::auth::{InputSource, SecretValue, Sourced};
fn deployment<T>(value: T) -> Sourced<T> {
Sourced::new(value, InputSource::Deployment)
}
fn sourced_client_secret(
credential_source: InputSource,
authority_source: InputSource,
authority: &str,
) -> NativeAzureRequest {
NativeAzureRequest::ClientSecret {
tenant_id: Sourced::new("tenant".to_string(), credential_source),
client_id: Sourced::new("client".to_string(), credential_source),
client_secret: Sourced::new(SecretValue::new("secret"), credential_source),
scope: Sourced::new("scope".to_string(), InputSource::Environment),
authority: Some(Sourced::new(authority.to_string(), authority_source)),
}
}
fn client_secret_request(
tenant: &str,
client: &str,
secret: &str,
scope: &str,
authority: &str,
) -> ValidatedAzureRequest {
ValidatedAzureRequest::new(NativeAzureRequest::ClientSecret {
tenant_id: deployment(tenant.to_string()),
client_id: deployment(client.to_string()),
client_secret: deployment(SecretValue::new(secret)),
scope: deployment(scope.to_string()),
authority: Some(deployment(authority.to_string())),
})
.unwrap()
}
#[derive(Debug, Default)]
struct RecordingTokenClient {
requests: Mutex<Vec<(String, String)>>,
}
impl HttpClient for RecordingTokenClient {
fn execute_request<'life0, 'life1, 'async_trait>(
&'life0 self,
request: &'life1 Request,
) -> std::pin::Pin<
Box<dyn std::future::Future<Output = Result<AsyncRawResponse>> + Send + 'async_trait>,
>
where
'life0: 'async_trait,
'life1: 'async_trait,
Self: 'async_trait,
{
Box::pin(async move {
let body = Bytes::from(request.body());
self.requests.lock().unwrap().push((
request.url().to_string(),
String::from_utf8(body.to_vec()).unwrap(),
));
Ok(AsyncRawResponse::from_bytes(
StatusCode::Ok,
Headers::new(),
r#"{"token_type":"Bearer","expires_in":3600,"ext_expires_in":3600,"access_token":"native-token"}"#,
))
})
}
}
#[tokio::test]
async fn client_secret_uses_sdk_protocol_and_reuses_cached_credential() {
let transport = Arc::new(RecordingTokenClient::default());
let acquirer =
NativeAzureTokenAcquirer::with_transport(4, Transport::new(transport.clone()));
let request = client_secret_request(
"tenant",
"client",
"secret",
"https://service.test/.default",
"https://login.test",
);
let first = acquirer.acquire(request.clone()).await.unwrap();
let second = acquirer.acquire(request).await.unwrap();
assert_eq!(first.secret().expose(), "native-token");
assert_eq!(second.secret().expose(), "native-token");
let requests = transport.requests.lock().unwrap();
assert_eq!(requests.len(), 1);
assert_eq!(requests[0].0, "https://login.test/tenant/oauth2/v2.0/token");
assert!(requests[0].1.contains("client_id=client"));
assert!(requests[0].1.contains("client_secret=secret"));
assert!(
requests[0]
.1
.contains("scope=https%3A%2F%2Fservice.test%2F.default")
);
}
#[tokio::test]
async fn credential_provider_cache_isolates_every_client_secret_identity_field() {
let transport = Arc::new(RecordingTokenClient::default());
let acquirer =
NativeAzureTokenAcquirer::with_transport(16, Transport::new(transport.clone()));
let request = client_secret_request;
let base = request("tenant", "client", "secret", "scope", "https://login.test");
let variants = [
base.clone(),
request(
"other-tenant",
"client",
"secret",
"scope",
"https://login.test",
),
request(
"tenant",
"other-client",
"secret",
"scope",
"https://login.test",
),
request(
"tenant",
"client",
"other-secret",
"scope",
"https://login.test",
),
request(
"tenant",
"client",
"secret",
"other-scope",
"https://login.test",
),
request(
"tenant",
"client",
"secret",
"scope",
"https://other-login.test",
),
];
acquirer.acquire(base.clone()).await.unwrap();
acquirer.acquire(base).await.unwrap();
for request in variants.into_iter().skip(1) {
acquirer.acquire(request).await.unwrap();
}
assert_eq!(transport.requests.lock().unwrap().len(), 6);
}
#[test]
fn request_authority_requires_request_owned_client_secret_identity() {
let error = ValidatedAzureRequest::new(sourced_client_secret(
InputSource::Deployment,
InputSource::Request,
"https://login.example",
))
.unwrap_err();
assert!(matches!(
error,
crate::AuthError::Configuration(
crate::auth::error::AuthConfigurationError::MixedAzureCredentialSources
)
));
}
#[test]
fn request_owned_client_secret_identity_can_select_custom_authority() {
let request = ValidatedAzureRequest::new(sourced_client_secret(
InputSource::Request,
InputSource::Request,
"https://login.example",
))
.unwrap();
assert_eq!(request.credential_source(), InputSource::Request);
}
#[test]
fn authority_is_restricted_to_an_https_origin() {
for authority in [
"http://login.example",
"https://user@login.example",
"https://login.example/tenant",
"https://login.example?target=other",
] {
let error = ValidatedAzureRequest::new(sourced_client_secret(
InputSource::Deployment,
InputSource::Deployment,
authority,
))
.unwrap_err();
assert!(matches!(
error,
crate::AuthError::Configuration(
crate::auth::error::AuthConfigurationError::InvalidAzureAuthority
)
));
}
}
}

View file

@ -0,0 +1,683 @@
use crate::AuthError;
use crate::auth::error::AuthConfigurationError;
use crate::auth::{
CredentialFileRef, CredentialLookup, CredentialRef, InputSource, ResolvedCredential,
SecretValue, Sourced, TokenProviderHandle,
};
use std::future::Future;
use std::pin::Pin;
use std::sync::Arc;
use super::native::{NativeAzureRequest, NativeAzureTokenAcquirer, ValidatedAzureRequest};
use super::types::{AzureAuthInputs, AzureCredentialType, ConfigValue, DEFAULT_AZURE_SCOPE};
const AZURE_AD_TOKEN_ENV: &str = "AZURE_AD_TOKEN";
const AZURE_TENANT_ID_ENV: &str = "AZURE_TENANT_ID";
const AZURE_CLIENT_ID_ENV: &str = "AZURE_CLIENT_ID";
const AZURE_CLIENT_SECRET_ENV: &str = "AZURE_CLIENT_SECRET";
const AZURE_SCOPE_ENV: &str = "AZURE_SCOPE";
const AZURE_AUTHORITY_HOST_ENV: &str = "AZURE_AUTHORITY_HOST";
const AZURE_CREDENTIAL_ENV: &str = "AZURE_CREDENTIAL";
const AZURE_FEDERATED_TOKEN_FILE_ENV: &str = "AZURE_FEDERATED_TOKEN_FILE";
#[derive(Clone, Debug)]
pub(crate) enum AzureCredentialPlan {
Supplied(Sourced<ResolvedCredential>),
Caller(TokenProviderHandle),
Oidc {
reference: Sourced<CredentialRef>,
tenant_id: Sourced<String>,
client_id: Sourced<String>,
scope: Sourced<String>,
authority: Option<Sourced<String>>,
},
Native(ValidatedAzureRequest),
Chain(Vec<ValidatedAzureRequest>),
Missing,
}
/// Rust counterpart to Python's `get_azure_ad_token`, not `BaseAzureLLM`.
pub(crate) struct AzureAuthService {
native: Arc<dyn AzureTokenAcquirer>,
}
trait AzureTokenAcquirer: Send + Sync {
fn acquire(
&self,
request: ValidatedAzureRequest,
) -> Pin<Box<dyn Future<Output = Result<ResolvedCredential, AuthError>> + Send + '_>>;
}
impl AzureTokenAcquirer for NativeAzureTokenAcquirer {
fn acquire(
&self,
request: ValidatedAzureRequest,
) -> Pin<Box<dyn Future<Output = Result<ResolvedCredential, AuthError>> + Send + '_>> {
Box::pin(NativeAzureTokenAcquirer::acquire(self, request))
}
}
impl Default for AzureAuthService {
fn default() -> Self {
Self {
native: Arc::new(NativeAzureTokenAcquirer::default()),
}
}
}
impl AzureAuthService {
#[cfg(test)]
fn with_acquirer(native: Arc<dyn AzureTokenAcquirer>) -> Self {
Self { native }
}
pub(crate) async fn get_azure_ad_token(
&self,
inputs: &AzureAuthInputs,
env_lookup: &(dyn Fn(&str) -> Option<String> + Sync),
) -> Result<Option<Sourced<ResolvedCredential>>, AuthError> {
match select_auth_plan(inputs, env_lookup)? {
AzureCredentialPlan::Supplied(credential) => Ok(Some(credential)),
AzureCredentialPlan::Caller(caller) => {
let credential = caller.acquire().await?;
if credential.secret().expose().is_empty() {
return Err(AuthError::EmptyAzureToken);
}
Ok(Some(Sourced::new(credential, InputSource::Deployment)))
}
AzureCredentialPlan::Oidc {
reference,
tenant_id,
client_id,
scope,
authority,
} => {
let assertion = resolve_reference(inputs, env_lookup, reference.value())
.await?
.ok_or(AuthError::UnresolvedOidcReference)?;
let request = ValidatedAzureRequest::new(NativeAzureRequest::ClientAssertion {
tenant_id,
client_id,
assertion: Sourced::new(assertion, reference.source()),
assertion_identity: format!("{:?}", reference.value()),
scope,
authority,
})?;
let source = request.credential_source();
self.native
.acquire(request)
.await
.map(|credential| Sourced::new(credential, source))
.map(Some)
}
AzureCredentialPlan::Native(request) => {
let source = request.credential_source();
self.native
.acquire(request)
.await
.map(|credential| Some(Sourced::new(credential, source)))
}
AzureCredentialPlan::Chain(requests) => {
let mut failures = Vec::new();
for request in requests {
let source = request.credential_source();
match self.native.acquire(request).await {
Ok(credential) => return Ok(Some(Sourced::new(credential, source))),
Err(error) => failures.push(error),
}
}
Err(AuthError::CredentialChain(failures))
}
AzureCredentialPlan::Missing => Ok(None),
}
}
}
pub(crate) fn select_auth_plan(
inputs: &AzureAuthInputs,
env_lookup: &dyn Fn(&str) -> Option<String>,
) -> Result<AzureCredentialPlan, AuthError> {
let token = configured_secret(&inputs.azure_ad_token, AZURE_AD_TOKEN_ENV, env_lookup);
let tenant_id = configured_string(&inputs.tenant_id, AZURE_TENANT_ID_ENV, env_lookup);
let client_id = configured_string(&inputs.client_id, AZURE_CLIENT_ID_ENV, env_lookup);
let client_secret =
configured_secret(&inputs.client_secret, AZURE_CLIENT_SECRET_ENV, env_lookup);
let scope = configured_string(&inputs.azure_scope, AZURE_SCOPE_ENV, env_lookup)
.unwrap_or_else(|| Sourced::new(DEFAULT_AZURE_SCOPE.to_string(), InputSource::Environment));
let authority = configured_string(
&inputs.azure_authority_host,
AZURE_AUTHORITY_HOST_ENV,
env_lookup,
);
let selector = configured_string(&inputs.azure_credential, AZURE_CREDENTIAL_ENV, env_lookup)
.map(|value| {
value
.value()
.parse::<AzureCredentialType>()
.map(|selector| Sourced::new(selector, value.source()))
})
.transpose()
.map_err(|_| AuthError::Configuration(AuthConfigurationError::InvalidAzureSelector))?;
let federated_token_file = configured_string(
&inputs.federated_token_file,
AZURE_FEDERATED_TOKEN_FILE_ENV,
env_lookup,
);
if inputs.azure_ad_token_provider.is_none()
&& let (Some(tenant_id), Some(client_id), Some(client_secret)) =
(tenant_id.clone(), client_id.clone(), client_secret)
{
return Ok(AzureCredentialPlan::Native(ValidatedAzureRequest::new(
NativeAzureRequest::ClientSecret {
tenant_id,
client_id,
client_secret,
scope,
authority,
},
)?));
}
if let (Some(reference), Some(tenant_id), Some(client_id)) = (
oidc_reference(&token)?,
tenant_id.clone(),
client_id.clone(),
) {
return Ok(AzureCredentialPlan::Oidc {
reference,
tenant_id,
client_id,
scope,
authority,
});
}
if let Some(caller) = &inputs.azure_ad_token_provider {
return Ok(AzureCredentialPlan::Caller(caller.clone()));
}
if let Some(token) = token {
return Ok(AzureCredentialPlan::Supplied(token.map(|token| {
ResolvedCredential::AccessToken {
token,
expires_on: None,
}
})));
}
if !*inputs.enable_azure_ad_token_refresh.value() && selector.is_none() {
return Ok(AzureCredentialPlan::Missing);
}
select_native_plan(
selector,
tenant_id,
client_id,
federated_token_file,
scope,
authority,
inputs.enable_azure_ad_token_refresh.source(),
)
}
fn select_native_plan(
selector: Option<Sourced<AzureCredentialType>>,
tenant_id: Option<Sourced<String>>,
client_id: Option<Sourced<String>>,
federated_token_file: Option<Sourced<String>>,
scope: Sourced<String>,
authority: Option<Sourced<String>>,
refresh_source: InputSource,
) -> Result<AzureCredentialPlan, AuthError> {
let selected = selector.unwrap_or_else(|| {
Sourced::new(
{
if federated_token_file.is_some() {
AzureCredentialType::DefaultAzureCredential
} else if client_id.is_some() {
AzureCredentialType::ManagedIdentityCredential
} else {
AzureCredentialType::DefaultAzureCredential
}
},
refresh_source,
)
});
let selection_source = selected.source();
match selected.into_value() {
AzureCredentialType::ClientSecretCredential => Err(AuthError::Configuration(
AuthConfigurationError::MissingClientSecretFields,
)),
AzureCredentialType::WorkloadIdentityCredential => {
Ok(AzureCredentialPlan::Native(ValidatedAzureRequest::new(
workload_request(tenant_id, client_id, federated_token_file, scope, authority)?,
)?))
}
AzureCredentialType::ManagedIdentityCredential => Ok(AzureCredentialPlan::Native(
ValidatedAzureRequest::new(NativeAzureRequest::ManagedIdentity {
client_id,
scope,
selection_source,
})?,
)),
AzureCredentialType::DefaultAzureCredential => {
let workload = match (tenant_id, client_id.clone(), federated_token_file) {
(Some(tenant_id), Some(client_id), Some(token_file_path)) => {
Some(NativeAzureRequest::WorkloadIdentity {
tenant_id,
client_id,
token_file_path,
scope: scope.clone(),
authority,
})
}
_ => None,
};
Ok(AzureCredentialPlan::Chain(
workload
.into_iter()
.chain(std::iter::once(NativeAzureRequest::ManagedIdentity {
client_id,
scope: scope.clone(),
selection_source,
}))
.chain(std::iter::once(NativeAzureRequest::DeveloperTools {
scope,
selection_source,
}))
.map(ValidatedAzureRequest::new)
.collect::<Result<Vec<_>, _>>()?,
))
}
AzureCredentialType::DeploymentIdentityCredential => {
let workload = match (tenant_id, client_id.clone(), federated_token_file) {
(Some(tenant_id), Some(client_id), Some(token_file_path)) => {
Some(NativeAzureRequest::WorkloadIdentity {
tenant_id,
client_id,
token_file_path,
scope: scope.clone(),
authority,
})
}
_ => None,
};
let user_assigned = client_id.map(|client_id| NativeAzureRequest::ManagedIdentity {
client_id: Some(client_id),
scope: scope.clone(),
selection_source,
});
Ok(AzureCredentialPlan::Chain(
workload
.into_iter()
.chain(user_assigned)
.chain(std::iter::once(NativeAzureRequest::ManagedIdentity {
client_id: None,
scope,
selection_source,
}))
.map(ValidatedAzureRequest::new)
.collect::<Result<Vec<_>, _>>()?,
))
}
}
}
fn workload_request(
tenant_id: Option<Sourced<String>>,
client_id: Option<Sourced<String>>,
token_file_path: Option<Sourced<String>>,
scope: Sourced<String>,
authority: Option<Sourced<String>>,
) -> Result<NativeAzureRequest, AuthError> {
Ok(NativeAzureRequest::WorkloadIdentity {
tenant_id: tenant_id.ok_or(AuthError::Configuration(
AuthConfigurationError::MissingWorkloadTenant,
))?,
client_id: client_id.ok_or(AuthError::Configuration(
AuthConfigurationError::MissingWorkloadClient,
))?,
token_file_path: token_file_path.ok_or(AuthError::Configuration(
AuthConfigurationError::MissingWorkloadTokenFile,
))?,
scope,
authority,
})
}
fn configured_string(
configured: &ConfigValue<String>,
environment_name: &str,
env_lookup: &dyn Fn(&str) -> Option<String>,
) -> Option<Sourced<String>> {
configured
.as_value()
.filter(|value| !value.value().is_empty())
.cloned()
.or_else(|| {
env_lookup(environment_name)
.filter(|value| !value.is_empty())
.map(|value| Sourced::new(value, InputSource::Environment))
})
}
fn configured_secret(
configured: &ConfigValue<SecretValue>,
environment_name: &str,
env_lookup: &dyn Fn(&str) -> Option<String>,
) -> Option<Sourced<SecretValue>> {
configured
.as_value()
.filter(|value| !value.value().expose().is_empty())
.cloned()
.or_else(|| {
env_lookup(environment_name)
.filter(|value| !value.is_empty())
.map(|value| Sourced::new(SecretValue::new(value), InputSource::Environment))
})
}
async fn resolve_reference(
inputs: &AzureAuthInputs,
env_lookup: &(dyn Fn(&str) -> Option<String> + Sync),
reference: &CredentialRef,
) -> Result<Option<SecretValue>, AuthError> {
let lookup = match reference {
CredentialRef::Explicit(secret) => return Ok(Some(secret.clone())),
CredentialRef::Env(name) => env_lookup(name)
.filter(|value| !value.is_empty())
.map(SecretValue::new)
.map_or(CredentialLookup::Missing, CredentialLookup::Found),
CredentialRef::None => return Ok(None),
CredentialRef::File(_) | CredentialRef::Request(_) | CredentialRef::Host(_) => {
let resolver = inputs
.credential_resolver
.as_ref()
.ok_or(AuthError::Configuration(
AuthConfigurationError::MissingHostResolver,
))?;
resolver.resolve(reference).await?
}
};
Ok(match lookup {
CredentialLookup::Found(secret) => Some(secret),
CredentialLookup::Missing | CredentialLookup::Declined => None,
})
}
fn oidc_reference(
token: &Option<Sourced<SecretValue>>,
) -> Result<Option<Sourced<CredentialRef>>, AuthError> {
let Some(token) = token.as_ref() else {
return Ok(None);
};
let value = token.value().expose();
if token.source() == InputSource::Request && value.starts_with("oidc/") {
return Err(AuthError::Configuration(
AuthConfigurationError::RequestAzureCredentialReference,
));
}
if let Some(name) = value.strip_prefix("oidc/env/") {
return non_empty_reference(name, "OIDC environment reference")
.map(CredentialRef::Env)
.map(|reference| Sourced::new(reference, token.source()))
.map(Some);
}
if let Some(name) = value.strip_prefix("oidc/env_path/") {
return non_empty_reference(name, "OIDC environment path reference")
.map(|name| CredentialRef::File(CredentialFileRef::EnvironmentVariable(name)))
.map(|reference| Sourced::new(reference, token.source()))
.map(Some);
}
if let Some(path) = value.strip_prefix("oidc/file/") {
let path = non_empty_reference(path, "OIDC file reference")?;
return Ok(Some(Sourced::new(
CredentialRef::File(CredentialFileRef::Path(path.into())),
token.source(),
)));
}
if value.starts_with("oidc/") {
return Err(AuthError::Configuration(
AuthConfigurationError::UnsupportedOidcReference,
));
}
Ok(None)
}
fn non_empty_reference(value: &str, kind: &str) -> Result<String, AuthError> {
if value.is_empty() {
return Err(AuthError::Configuration(
AuthConfigurationError::EmptyReference(kind.to_string()),
));
}
Ok(value.to_string())
}
#[cfg(test)]
mod tests {
use std::future::Future;
use std::sync::{Arc, Mutex};
use serde_json::json;
use super::{
AzureAuthService, AzureCredentialPlan, AzureTokenAcquirer, oidc_reference,
resolve_reference, select_auth_plan,
};
use crate::AuthError;
use crate::auth::ResolvedCredential;
use crate::auth::{
CredentialFileRef, CredentialLookup, CredentialLookupFuture, CredentialRef,
CredentialResolver, CredentialResolverHandle, InputSource, SecretValue, Sourced,
};
use crate::providers::azure_ai::auth::native::ValidatedAzureRequest;
use crate::providers::azure_ai::auth::types::AzureAuthInputs;
#[derive(Debug)]
struct FileResolver;
struct ChainAcquirer {
requests: Mutex<Vec<&'static str>>,
succeed_on: Option<&'static str>,
}
impl AzureTokenAcquirer for ChainAcquirer {
fn acquire(
&self,
request: ValidatedAzureRequest,
) -> std::pin::Pin<
Box<dyn Future<Output = Result<ResolvedCredential, AuthError>> + Send + '_>,
> {
let kind = request.kind();
self.requests.lock().unwrap().push(kind);
Box::pin(async move {
if self.succeed_on == Some(kind) {
Ok(ResolvedCredential::AccessToken {
token: SecretValue::new("chain-token"),
expires_on: None,
})
} else {
Err(AuthError::AzureTokenAcquisition(format!("{kind} failed")))
}
})
}
}
impl CredentialResolver for FileResolver {
fn resolve<'a>(&'a self, reference: &'a CredentialRef) -> CredentialLookupFuture<'a> {
Box::pin(async move {
Ok(match reference {
CredentialRef::File(CredentialFileRef::Path(path))
if path == std::path::Path::new("/run/secrets/assertion") =>
{
CredentialLookup::Found(SecretValue::new("rotated-assertion"))
}
_ => CredentialLookup::Declined,
})
})
}
}
#[test]
fn null_and_empty_values_fall_back_to_environment() {
let params = json!({"tenant_id": null, "client_id": "", "client_secret": null});
let inputs = AzureAuthInputs::from_optional_params(params.as_object().unwrap()).unwrap();
let plan = select_auth_plan(&inputs, &|name| match name {
"AZURE_TENANT_ID" => Some("tenant".to_string()),
"AZURE_CLIENT_ID" => Some("client".to_string()),
"AZURE_CLIENT_SECRET" => Some("secret".to_string()),
_ => None,
})
.unwrap();
assert!(matches!(plan, AzureCredentialPlan::Native(_)));
}
#[test]
fn supplied_token_does_not_require_refresh() {
let params = json!({"azure_ad_token": "token"});
let inputs = AzureAuthInputs::from_optional_params(params.as_object().unwrap()).unwrap();
assert!(matches!(
select_auth_plan(&inputs, &|_| None).unwrap(),
AzureCredentialPlan::Supplied(_)
));
}
#[test]
fn oidc_reference_is_deferred() {
let params = json!({
"azure_ad_token": "oidc/env/ASSERTION",
"tenant_id": "tenant",
"client_id": "client"
});
let inputs = AzureAuthInputs::from_optional_params(params.as_object().unwrap()).unwrap();
assert!(matches!(
select_auth_plan(&inputs, &|_| None).unwrap(),
AzureCredentialPlan::Oidc {
reference,
..
} if reference.value() == &CredentialRef::Env("ASSERTION".to_string())
));
}
#[test]
fn oidc_file_location_is_typed_before_resolution() {
assert_eq!(
oidc_reference(&Some(Sourced::new(
SecretValue::new("oidc/file//run/secrets/assertion"),
InputSource::Deployment,
)))
.unwrap()
.map(Sourced::into_value),
Some(CredentialRef::File(CredentialFileRef::Path(
"/run/secrets/assertion".into()
)))
);
}
#[test]
fn unsupported_oidc_reference_is_rejected_during_plan_creation() {
let error = oidc_reference(&Some(Sourced::new(
SecretValue::new("oidc/vault/assertion"),
InputSource::Deployment,
)))
.expect_err("unsupported backend must fail validation");
assert!(error.to_string().contains("unsupported OIDC reference"));
}
#[test]
fn request_oidc_reference_is_rejected_before_lookup() {
let params = json!({
"azure_ad_token": "oidc/env/ASSERTION",
"tenant_id": "tenant",
"client_id": "client"
});
let sources = std::collections::BTreeMap::from([
("azure_ad_token".to_string(), InputSource::Request),
("tenant_id".to_string(), InputSource::Request),
("client_id".to_string(), InputSource::Request),
]);
let inputs =
AzureAuthInputs::from_sourced_optional_params(params.as_object().unwrap(), &sources)
.unwrap();
let error = select_auth_plan(&inputs, &|name| {
assert_ne!(name, "ASSERTION");
None
})
.unwrap_err();
assert!(matches!(
error,
AuthError::Configuration(
crate::auth::error::AuthConfigurationError::RequestAzureCredentialReference
)
));
}
#[tokio::test]
async fn host_resolver_owns_file_access() {
let inputs = AzureAuthInputs {
credential_resolver: Some(CredentialResolverHandle::new(Arc::new(FileResolver))),
..AzureAuthInputs::default()
};
let reference =
CredentialRef::File(CredentialFileRef::Path("/run/secrets/assertion".into()));
let resolved = resolve_reference(&inputs, &|_| None, &reference)
.await
.unwrap();
assert_eq!(resolved, Some(SecretValue::new("rotated-assertion")));
}
#[tokio::test]
async fn default_chain_uses_declared_order_and_stops_after_success() {
let acquirer = Arc::new(ChainAcquirer {
requests: Mutex::new(Vec::new()),
succeed_on: Some("developer-tools"),
});
let service = AzureAuthService::with_acquirer(acquirer.clone());
let inputs = AzureAuthInputs {
enable_azure_ad_token_refresh: Sourced::new(true, InputSource::Deployment),
..Default::default()
};
let credential = service
.get_azure_ad_token(&inputs, &|_| None)
.await
.unwrap()
.unwrap();
assert_eq!(credential.value().secret().expose(), "chain-token");
assert_eq!(
*acquirer.requests.lock().unwrap(),
["managed-identity", "developer-tools"]
);
}
#[tokio::test]
async fn chain_reports_each_acquisition_failure() {
let acquirer = Arc::new(ChainAcquirer {
requests: Mutex::new(Vec::new()),
succeed_on: None,
});
let service = AzureAuthService::with_acquirer(acquirer);
let inputs = AzureAuthInputs {
enable_azure_ad_token_refresh: Sourced::new(true, InputSource::Deployment),
..Default::default()
};
let error = service
.get_azure_ad_token(&inputs, &|_| None)
.await
.unwrap_err();
assert!(matches!(error, AuthError::CredentialChain(errors) if errors.len() == 2));
}
}

View file

@ -0,0 +1,195 @@
use crate::auth::error::AuthConfigurationError;
use serde_json::{Map, Value};
use std::collections::BTreeMap;
use strum::EnumString;
use crate::AuthError;
use crate::auth::{
CredentialResolverHandle, InputSource, SecretValue, Sourced, TokenProviderHandle,
};
pub const DEFAULT_AZURE_SCOPE: &str = "https://cognitiveservices.azure.com/.default";
#[derive(Clone, Debug, Default, PartialEq, Eq)]
pub enum ConfigValue<T> {
#[default]
Absent,
ExplicitNone(InputSource),
Value(Sourced<T>),
}
impl<T> ConfigValue<T> {
pub fn as_value(&self) -> Option<&Sourced<T>> {
match self {
Self::Value(value) => Some(value),
Self::Absent | Self::ExplicitNone(_) => None,
}
}
}
#[derive(Clone, Copy, Debug, EnumString, PartialEq, Eq, Hash)]
#[allow(clippy::enum_variant_names)]
pub enum AzureCredentialType {
ClientSecretCredential,
ManagedIdentityCredential,
DefaultAzureCredential,
DeploymentIdentityCredential,
WorkloadIdentityCredential,
}
#[derive(Clone, Debug, Default)]
pub struct AzureAuthInputs {
pub azure_ad_token: ConfigValue<SecretValue>,
pub azure_ad_token_provider: Option<TokenProviderHandle>,
pub credential_resolver: Option<CredentialResolverHandle>,
pub tenant_id: ConfigValue<String>,
pub client_id: ConfigValue<String>,
pub client_secret: ConfigValue<SecretValue>,
pub azure_scope: ConfigValue<String>,
pub azure_authority_host: ConfigValue<String>,
pub azure_credential: ConfigValue<String>,
pub federated_token_file: ConfigValue<String>,
pub enable_azure_ad_token_refresh: Sourced<bool>,
}
impl AzureAuthInputs {
#[cfg(test)]
pub fn from_optional_params(params: &Map<String, Value>) -> Result<Self, AuthError> {
Self::from_sourced_optional_params(params, &BTreeMap::new())
}
pub fn from_sourced_optional_params(
params: &Map<String, Value>,
sources: &BTreeMap<String, InputSource>,
) -> Result<Self, AuthError> {
Ok(Self {
azure_ad_token: secret_config(params, sources, "azure_ad_token")?,
azure_ad_token_provider: None,
credential_resolver: None,
tenant_id: string_config(params, sources, "tenant_id")?,
client_id: string_config(params, sources, "client_id")?,
client_secret: secret_config(params, sources, "client_secret")?,
azure_scope: string_config(params, sources, "azure_scope")?,
azure_authority_host: string_config(params, sources, "azure_authority_host")?,
azure_credential: string_config(params, sources, "azure_credential")?,
federated_token_file: string_config(params, sources, "azure_federated_token_file")?,
enable_azure_ad_token_refresh: Sourced::new(
params
.get("enable_azure_ad_token_refresh")
.and_then(Value::as_bool)
.unwrap_or(false),
source_for(sources, "enable_azure_ad_token_refresh"),
),
})
}
}
fn string_config(
params: &Map<String, Value>,
sources: &BTreeMap<String, InputSource>,
name: &str,
) -> Result<ConfigValue<String>, AuthError> {
let source = source_for(sources, name);
match params.get(name) {
None => Ok(ConfigValue::Absent),
Some(Value::Null) => Ok(ConfigValue::ExplicitNone(source)),
Some(Value::String(value)) => Ok(ConfigValue::Value(Sourced::new(value.clone(), source))),
Some(_) => Err(AuthError::Configuration(
AuthConfigurationError::InvalidFieldType(name.to_string()),
)),
}
}
fn secret_config(
params: &Map<String, Value>,
sources: &BTreeMap<String, InputSource>,
name: &str,
) -> Result<ConfigValue<SecretValue>, AuthError> {
Ok(match string_config(params, sources, name)? {
ConfigValue::Absent => ConfigValue::Absent,
ConfigValue::ExplicitNone(source) => ConfigValue::ExplicitNone(source),
ConfigValue::Value(value) => ConfigValue::Value(value.map(SecretValue::new)),
})
}
fn source_for(sources: &BTreeMap<String, InputSource>, name: &str) -> InputSource {
sources.get(name).copied().unwrap_or_default()
}
#[cfg(test)]
mod tests {
use serde_json::json;
use std::collections::BTreeMap;
use super::{AzureAuthInputs, AzureCredentialType, ConfigValue};
use crate::auth::{InputSource, Sourced};
#[test]
fn selector_parsing_is_exact() {
assert_eq!(
"ClientSecretCredential".parse::<AzureCredentialType>(),
Ok(AzureCredentialType::ClientSecretCredential)
);
assert!(
"clientsecretcredential"
.parse::<AzureCredentialType>()
.is_err()
);
}
#[test]
fn defaults_preserve_absence() {
let inputs = AzureAuthInputs::default();
assert_eq!(inputs.tenant_id, ConfigValue::Absent);
assert_eq!(inputs.azure_ad_token, ConfigValue::Absent);
}
#[test]
fn parsing_distinguishes_null_empty_and_absent() {
let params = json!({"tenant_id": null, "client_id": ""});
let inputs = AzureAuthInputs::from_optional_params(params.as_object().unwrap()).unwrap();
assert_eq!(
inputs.tenant_id,
ConfigValue::ExplicitNone(InputSource::Deployment)
);
assert_eq!(
inputs.client_id,
ConfigValue::Value(Sourced::new(String::new(), InputSource::Deployment))
);
assert_eq!(inputs.client_secret, ConfigValue::Absent);
}
#[test]
fn parsing_preserves_trusted_input_sources() {
let params = json!({"tenant_id": "tenant", "client_secret": null});
let sources = BTreeMap::from([
("tenant_id".to_string(), InputSource::Request),
("client_secret".to_string(), InputSource::Request),
]);
let inputs =
AzureAuthInputs::from_sourced_optional_params(params.as_object().unwrap(), &sources)
.unwrap();
assert_eq!(
inputs.tenant_id,
ConfigValue::Value(Sourced::new("tenant".to_string(), InputSource::Request))
);
assert_eq!(
inputs.client_secret,
ConfigValue::ExplicitNone(InputSource::Request)
);
}
#[test]
fn debug_does_not_expose_secrets() {
let params = json!({"azure_ad_token": "token-value", "client_secret": "secret-value"});
let inputs = AzureAuthInputs::from_optional_params(params.as_object().unwrap()).unwrap();
let debug = format!("{inputs:?}");
assert!(!debug.contains("token-value"));
assert!(!debug.contains("secret-value"));
}
}

View file

@ -1,3 +1,4 @@
use crate::auth::error::MissingCredential;
use crate::error::Error;
use crate::messages::transformation::{AnthropicMessagesProviderConfig, MessagesAuthStrategy};
use crate::messages::types::{
@ -32,12 +33,7 @@ pub fn resolve_azure_api_key(
non_empty(api_key)
.map(str::to_string)
.or_else(|| env_lookup(AZURE_API_KEY_ENV).filter(|value| !value.trim().is_empty()))
.ok_or_else(|| {
Error::Auth(
"Missing Azure API Key - Set `api_key` or the AZURE_API_KEY environment variable"
.to_string(),
)
})
.ok_or_else(|| Error::from(crate::AuthError::from(MissingCredential::AzureApiKey)))
}
pub fn complete_azure_anthropic_url(
@ -47,13 +43,7 @@ pub fn complete_azure_anthropic_url(
let api_base = non_empty(api_base)
.map(str::to_string)
.or_else(|| env_lookup(AZURE_API_BASE_ENV).filter(|value| !value.trim().is_empty()))
.ok_or_else(|| {
Error::Auth(
"Missing Azure API Base - Set `api_base` or the AZURE_API_BASE environment variable. \
Expected format: https://<resource-name>.services.ai.azure.com/anthropic"
.to_string(),
)
})?;
.ok_or_else(|| Error::from(crate::AuthError::from(MissingCredential::AzureApiBase)))?;
let api_base = api_base.trim_end_matches('/');

View file

@ -1,2 +1,2 @@
pub(crate) mod auth;
pub mod messages;
pub mod ocr;

View file

@ -1 +0,0 @@
pub mod transformation;

View file

@ -1 +0,0 @@
pub mod ocr;

View file

@ -1 +0,0 @@
pub mod transformation;

View file

@ -1,436 +0,0 @@
use crate::error::{Error, json_type_name};
use crate::ocr::transformation::OcrProviderConfig;
use crate::ocr::types::{LiteLLMOcrResponse, OcrRequestData};
use serde_json::{Map, Value};
const SUPPORTED_OCR_PARAMS: &[&str] = &[
"pages",
"include_image_base64",
"image_limit",
"image_min_size",
"bbox_annotation_format",
"document_annotation_format",
"document_annotation_prompt",
"extract_header",
"extract_footer",
"table_format",
"confidence_scores_granularity",
"include_blocks",
"id",
];
/// Default Mistral API base, used when the caller does not override `api_base`.
pub const MISTRAL_DEFAULT_API_BASE: &str = "https://api.mistral.ai/v1";
/// Environment variable holding the Mistral API key.
pub const MISTRAL_API_KEY_ENV: &str = "MISTRAL_API_KEY";
/// Error message raised when no Mistral API key can be resolved.
pub const MISSING_KEY_MESSAGE: &str = "Missing Mistral API Key - A call is being made to Mistral but no key is set either in the environment variables or via params";
/// Build the complete OCR endpoint URL, de-duplicating a trailing `/v1`.
///
/// Blank/whitespace `api_base` is treated as absent (guard at resolution time).
pub fn complete_url(api_base: Option<&str>) -> String {
let base = api_base
.map(str::trim)
.filter(|base| !base.is_empty())
.unwrap_or(MISTRAL_DEFAULT_API_BASE)
.trim_end_matches('/');
if base.ends_with("/v1") {
format!("{base}/ocr")
} else {
format!("{base}/v1/ocr")
}
}
/// Resolve the Mistral API key from the explicit param or the environment.
///
/// Blank/whitespace values are treated as absent. Returns `Error::Auth`
/// when no usable key is available.
///
/// Note: the env fallback only reads the process environment. Secret-manager
/// backends (AWS/Azure/GCP/Vault) are resolved on the Python side and passed in
/// via `api_key`; this fallback is a last resort for direct/standalone use.
pub fn resolve_api_key(
api_key: Option<&str>,
env_lookup: &dyn Fn(&str) -> Option<String>,
) -> Result<String, Error> {
api_key
.map(str::trim)
.filter(|key| !key.is_empty())
.map(str::to_string)
.or_else(|| env_lookup(MISTRAL_API_KEY_ENV).filter(|key| !key.trim().is_empty()))
.ok_or_else(|| Error::Auth(MISSING_KEY_MESSAGE.to_string()))
}
pub struct MistralOcrConfig;
pub const MISTRAL_OCR_CONFIG: MistralOcrConfig = MistralOcrConfig;
impl OcrProviderConfig for MistralOcrConfig {
#[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)]
fn supported_ocr_params(&self) -> &'static [&'static str] {
SUPPORTED_OCR_PARAMS
}
#[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)]
fn transform_ocr_request(
&self,
model: &str,
document: Value,
optional_params: Map<String, Value>,
) -> Result<OcrRequestData, Error> {
if !document.is_object() {
return Err(Error::InvalidType {
expected: "object",
actual: json_type_name(&document),
});
}
let mut data = Map::new();
data.insert("model".to_string(), Value::String(model.to_string()));
data.insert("document".to_string(), document);
for (param, value) in optional_params {
data.insert(param, value);
}
Ok(OcrRequestData {
data: Value::Object(data),
files: None,
})
}
#[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)]
fn transform_ocr_response(
&self,
model: &str,
response_json: Value,
) -> Result<LiteLLMOcrResponse, Error> {
let response_object = response_json
.as_object()
.ok_or_else(|| Error::InvalidType {
expected: "object",
actual: json_type_name(&response_json),
})?;
let pages = response_object
.get("pages")
.and_then(Value::as_array)
.cloned()
.unwrap_or_default();
let model = response_object
.get("model")
.and_then(Value::as_str)
.unwrap_or(model)
.to_string();
let document_annotation = response_object.get("document_annotation").cloned();
let usage_info = response_object.get("usage_info").cloned();
Ok(LiteLLMOcrResponse {
pages,
model,
document_annotation,
usage_info,
object: "ocr".to_string(),
extra_fields: Map::new(),
provider_native_response: None,
})
}
#[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)]
fn complete_url(
&self,
api_base: Option<&str>,
_model: &str,
_optional_params: &Map<String, Value>,
_env_lookup: &dyn Fn(&str) -> Option<String>,
) -> Result<String, Error> {
Ok(complete_url(api_base))
}
fn resolve_api_key(
&self,
api_key: Option<&str>,
env_lookup: &dyn Fn(&str) -> Option<String>,
) -> Result<String, Error> {
resolve_api_key(api_key, env_lookup)
}
}
#[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)]
pub fn supported_ocr_params() -> &'static [&'static str] {
MISTRAL_OCR_CONFIG.supported_ocr_params()
}
pub fn map_ocr_params(non_default_params: &Map<String, Value>) -> Map<String, Value> {
MISTRAL_OCR_CONFIG.map_ocr_params(non_default_params)
}
#[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)]
pub fn transform_ocr_request(
model: &str,
document: Value,
optional_params: Map<String, Value>,
) -> Result<OcrRequestData, Error> {
MISTRAL_OCR_CONFIG.transform_ocr_request(model, document, optional_params)
}
#[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)]
pub fn transform_ocr_response(
model: &str,
response_json: Value,
) -> Result<LiteLLMOcrResponse, Error> {
MISTRAL_OCR_CONFIG.transform_ocr_response(model, response_json)
}
#[cfg(test)]
mod tests {
use super::*;
use serde_json::json;
#[test]
fn extract_header_is_a_supported_ocr_param() {
assert!(supported_ocr_params().contains(&"extract_header"));
}
#[test]
fn extract_footer_is_a_supported_ocr_param() {
assert!(supported_ocr_params().contains(&"extract_footer"));
}
#[test]
fn existing_ocr_params_remain_supported() {
for param in [
"pages",
"include_image_base64",
"image_limit",
"image_min_size",
"bbox_annotation_format",
"document_annotation_format",
] {
assert!(supported_ocr_params().contains(&param));
}
}
#[test]
fn map_ocr_params_forwards_extract_header() {
let params = json!({"extract_header": true});
assert_eq!(
map_ocr_params(params.as_object().unwrap()),
params.as_object().unwrap().clone()
);
}
#[test]
fn map_ocr_params_forwards_extract_footer() {
let params = json!({"extract_footer": true});
assert_eq!(
map_ocr_params(params.as_object().unwrap()),
params.as_object().unwrap().clone()
);
}
#[test]
fn map_ocr_params_forwards_extract_header_and_footer() {
let params = json!({"extract_header": true, "extract_footer": false});
assert_eq!(
map_ocr_params(params.as_object().unwrap()),
params.as_object().unwrap().clone()
);
}
#[test]
fn map_ocr_params_drops_unknown_params() {
let params = json!({"extract_header": true, "unsupported_param": "value"});
let mapped = map_ocr_params(params.as_object().unwrap());
assert_eq!(mapped.get("extract_header"), Some(&json!(true)));
assert!(!mapped.contains_key("unsupported_param"));
}
#[test]
fn new_ocr_params_are_supported() {
for param in [
"table_format",
"confidence_scores_granularity",
"document_annotation_prompt",
"include_blocks",
"id",
] {
assert!(supported_ocr_params().contains(&param));
}
}
#[test]
fn map_ocr_params_forwards_new_ocr_params() {
for (param, value) in [
("table_format", json!("html")),
("confidence_scores_granularity", json!("word")),
(
"document_annotation_prompt",
json!("Extract all invoice line items"),
),
("include_blocks", json!(true)),
("id", json!("req-123")),
] {
let params = json!({param: value});
assert_eq!(
map_ocr_params(params.as_object().unwrap()),
params.as_object().unwrap().clone()
);
}
}
#[test]
fn transform_ocr_request_includes_each_optional_param() {
let document = json!({
"type": "document_url",
"document_url": "https://example.com/doc.pdf"
});
for (param, value) in [
("table_format", json!("html")),
("confidence_scores_granularity", json!("word")),
(
"document_annotation_prompt",
json!("Extract all invoice line items"),
),
("id", json!("req-123")),
("extract_header", json!(true)),
("include_blocks", json!(true)),
("pages", json!([0, 1])),
] {
let result = transform_ocr_request(
"mistral-ocr-latest",
document.clone(),
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.files, None);
}
}
#[test]
fn transform_ocr_request_includes_multiple_new_params() {
let document = json!({
"type": "document_url",
"document_url": "https://example.com/doc.pdf"
});
let optional_params = json!({
"table_format": "html",
"confidence_scores_granularity": "page",
"extract_header": true
})
.as_object()
.unwrap()
.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"),
Some(&json!("page"))
);
assert_eq!(result.data.get("extract_header"), Some(&json!(true)));
}
#[test]
fn transform_ocr_response_preserves_blocks_and_confidence_scores() {
let blocks = json!([{"type": "title", "content": "Invoice"}]);
let confidence_scores = json!({"page": 0.98});
let response = json!({
"pages": [{"index": 0, "markdown": "# Invoice", "blocks": blocks, "confidence_scores": confidence_scores}],
"model": "mistral-ocr-4-0",
"usage_info": {"pages_processed": 1}
});
let result =
transform_ocr_response("mistral-ocr-4-0", response).expect("response should transform");
assert_eq!(result.pages[0].get("blocks"), Some(&blocks));
assert_eq!(
result.pages[0].get("confidence_scores"),
Some(&confidence_scores)
);
}
#[test]
fn transform_ocr_response_preserves_ocr4_page_fields() {
let response = json!({
"pages": [{"index": 0, "markdown": "table page", "tables": [{"rows": 2, "cols": 3}], "hyperlinks": ["https://example.com"], "header": "Acme Corp", "footer": "Page 1"}],
"model": "mistral-ocr-4-0",
"usage_info": {"pages_processed": 1}
});
let result = transform_ocr_response("mistral-ocr-4-0", response.clone())
.expect("response should transform");
assert_eq!(result.pages[0], response["pages"][0]);
}
#[test]
fn transform_ocr_request_rejects_non_object_document() {
let err = transform_ocr_request("mistral-ocr-latest", json!("bad"), Map::new())
.expect_err("string document should be rejected");
assert_eq!(
err,
Error::InvalidType {
expected: "object",
actual: "string",
}
);
}
#[test]
fn transform_ocr_response_normalizes_mistral_json() {
let response = json!({
"pages": [{"index": 0, "markdown": "hello"}],
"model": "mistral-ocr-2505-completion",
"document_annotation": null,
"usage_info": {"pages_processed": 1}
});
let result = transform_ocr_response("mistral-ocr-latest", response)
.expect("response should transform");
assert_eq!(result.pages, vec![json!({"index": 0, "markdown": "hello"})]);
assert_eq!(result.model, "mistral-ocr-2505-completion");
assert_eq!(result.document_annotation, Some(Value::Null));
assert_eq!(result.usage_info, Some(json!({"pages_processed": 1})));
assert_eq!(result.object, "ocr");
}
#[test]
fn complete_url_defaults_and_dedupes_v1() {
assert_eq!(complete_url(None), "https://api.mistral.ai/v1/ocr");
assert_eq!(complete_url(Some(" ")), "https://api.mistral.ai/v1/ocr");
assert_eq!(
complete_url(Some("https://proxy.internal")),
"https://proxy.internal/v1/ocr"
);
assert_eq!(
complete_url(Some("https://proxy.internal/v1/")),
"https://proxy.internal/v1/ocr"
);
}
#[test]
fn resolve_api_key_prefers_param_then_env() {
let no_env = |_: &str| None;
assert_eq!(
resolve_api_key(Some("sk-param"), &no_env).unwrap(),
"sk-param"
);
let with_env = |key: &str| (key == MISTRAL_API_KEY_ENV).then(|| "sk-env".to_string());
assert_eq!(resolve_api_key(None, &with_env).unwrap(), "sk-env");
// Blank param falls through to the environment.
assert_eq!(resolve_api_key(Some(" "), &with_env).unwrap(), "sk-env");
}
#[test]
fn resolve_api_key_errors_when_absent() {
let err = resolve_api_key(None, &|_| None).expect_err("missing key should error");
assert_eq!(err, Error::Auth(MISSING_KEY_MESSAGE.to_string()));
}
}

View file

@ -2,7 +2,4 @@ pub mod anthropic;
pub mod azure_ai;
#[cfg(feature = "bedrock-auth")]
pub mod bedrock;
pub mod mistral;
pub mod openai;
pub mod reducto;
pub mod vertex_ai;

View file

@ -1 +0,0 @@
pub mod ocr;

View file

@ -1,4 +0,0 @@
pub mod transformation;
#[cfg(test)]
mod tests;

View file

@ -1,202 +0,0 @@
use rstest::{fixture, rstest};
use serde_json::{Value, json};
use super::transformation::*;
use crate::ocr::transformation::OcrProviderConfig;
#[fixture]
fn parse_response() -> Value {
json!({
"job_id": "job_123",
"usage": {"num_pages": 3, "credits": 3},
"result": {
"chunks": [
{
"content": "Page 1 block A",
"blocks": [{
"content": "Page 1 block A",
"bbox": {"page": 1},
"kind": "text",
}],
},
{
"content": "Page 2 block A",
"blocks": [{
"content": "Page 2 block A",
"bbox": {"page": 2},
"kind": "table",
}],
},
{
"content": "Page 1 block B",
"blocks": [{
"content": "Page 1 block B",
"bbox": {"page": 1},
"kind": "text",
}],
},
{
"content": "Page 3 block A",
"blocks": [{
"content": "Page 3 block A",
"bbox": {"page": 3},
"kind": "figure",
}],
},
],
},
})
}
#[rstest]
fn test_parse_v3_file_upload_and_response_mapping(parse_response: Value) {
let source = classify_document_source("data:application/pdf;base64,JVBERi0xLjQ=")
.expect("PDF data URI should be valid");
let upload = build_upload_request(
source,
"Bearer test-key",
Some("https://platform.reducto.ai"),
)
.expect("data URI should require upload");
assert_eq!(upload.url, "https://platform.reducto.ai/upload");
assert_eq!(upload.authorization, "Bearer test-key");
assert_eq!(upload.file_name, "document");
assert_eq!(upload.mime_type, "application/pdf");
assert_eq!(upload.bytes, b"%PDF-1.4");
let optional_params = json!({
"formatting": {"table_output_format": "html"},
"retrieval": {"chunk_mode": "section"},
"settings": {"ocr_system": "standard"},
})
.as_object()
.expect("params should be an object")
.clone();
let request = build_parse_v3_request("reducto://uploaded.pdf", optional_params);
assert_eq!(
request.data,
json!({
"input": "reducto://uploaded.pdf",
"formatting": {"table_output_format": "html"},
"retrieval": {"chunk_mode": "section"},
"settings": {"ocr_system": "standard"},
})
);
let transformed = transform_reducto_response("parse-v3", parse_response.clone())
.expect("response should transform");
assert_eq!(
transformed.usage_info,
Some(json!({"pages_processed": 3, "credits": 3}))
);
assert_eq!(transformed.pages.len(), 3);
assert_eq!(
transformed.pages[0],
json!({
"index": 0,
"markdown": "Page 1 block A\n\nPage 1 block B",
"blocks": [
{"content": "Page 1 block A", "bbox": {"page": 1}, "kind": "text"},
{"content": "Page 1 block B", "bbox": {"page": 1}, "kind": "text"},
],
})
);
assert_eq!(transformed.pages[1]["markdown"], "Page 2 block A");
assert_eq!(transformed.pages[2]["markdown"], "Page 3 block A");
assert_eq!(transformed.provider_native_response, Some(parse_response));
}
#[rstest]
fn test_parse_v3_reducto_id_passthrough_skips_upload(parse_response: Value) {
let document = json!({
"type": "document_url",
"document_url": "reducto://already-uploaded.pdf",
});
let source = extract_document_source(&document).expect("Reducto ID should be valid");
assert!(build_upload_request(source.clone(), "Bearer test-key", None).is_none());
assert_eq!(
source,
ReductoDocumentSource::FileId("reducto://already-uploaded.pdf".to_string())
);
let request = REDUCTO_PARSE_V3_CONFIG
.transform_ocr_request(
"parse-v3",
document,
json!({"retrieval": {"chunk_mode": "section"}})
.as_object()
.expect("params should be object")
.clone(),
)
.expect("direct ID should transform");
assert_eq!(request.data["input"], "reducto://already-uploaded.pdf");
assert_eq!(request.data["retrieval"]["chunk_mode"], "section");
let response = REDUCTO_PARSE_V3_CONFIG
.transform_ocr_response("parse-v3", parse_response)
.expect("response should transform");
assert!(
response.pages[0]["markdown"]
.as_str()
.expect("markdown should be string")
.starts_with("Page 1 block A")
);
}
#[rstest]
fn test_parse_legacy_wraps_enhance_under_options() {
let request = build_parse_legacy_request(
"reducto://legacy.pdf",
json!({"enhance": {"agentic": [{"type": "table"}]}})
.as_object()
.expect("params should be object"),
);
assert_eq!(
request.data,
json!({
"document_url": "reducto://legacy.pdf",
"options": {"enhance": {"agentic": [{"type": "table"}]}},
})
);
}
#[rstest]
fn test_parse_v3_image_data_uri_upload_uses_image_mime() {
let source = classify_document_source("data:image/png;base64,iVBORw0KGgo=")
.expect("PNG data URI should be valid");
let upload = build_upload_request(
source,
"Bearer programmatic-key",
Some("https://custom.reducto.test/"),
)
.expect("data URI should require upload");
assert_eq!(upload.url, "https://custom.reducto.test/upload");
assert_eq!(upload.authorization, "Bearer programmatic-key");
assert_eq!(upload.mime_type, "image/png");
assert_eq!(upload.bytes, b"\x89PNG\r\n\x1a\n");
}
#[rstest]
#[case::http("http://example.com/document.pdf")]
#[case::https("https://example.com/document.pdf")]
fn test_parse_v3_rejects_plain_http_urls(#[case] source: &str) {
let error = classify_document_source(source).expect_err("plain URL should be rejected");
assert!(error.to_string().contains("upload the file first"));
}
#[rstest]
fn test_parse_v3_uses_programmatic_api_key_over_env() {
let key = resolve_api_key(Some("passed-key"), &|_| Some("env-reducto-key".to_string()))
.expect("explicit key should resolve");
assert_eq!(key, "passed-key");
let headers = REDUCTO_PARSE_V3_CONFIG
.validate_environment(Vec::new(), Some("passed-key"), &|_| {
Some("env-reducto-key".to_string())
})
.expect("headers should validate");
assert_eq!(
headers,
vec![("Authorization".to_string(), "Bearer passed-key".to_string())]
);
}

View file

@ -1,407 +0,0 @@
use std::collections::BTreeMap;
use base64::Engine;
use base64::engine::general_purpose::STANDARD as BASE64_STANDARD;
use serde_json::{Map, Value, json};
use crate::error::{Error, json_type_name};
use crate::ocr::transformation::OcrProviderConfig;
use crate::ocr::types::{LiteLLMOcrResponse, OcrRequestData};
pub const REDUCTO_API_BASE: &str = "https://platform.reducto.ai";
pub const REDUCTO_API_KEY_ENV: &str = "REDUCTO_API_KEY";
pub const REDUCTO_ID_PREFIX: &str = "reducto://";
const PARSE_V3_SUPPORTED_OCR_PARAMS: &[&str] = &["formatting", "retrieval", "settings"];
const PARSE_LEGACY_SUPPORTED_OCR_PARAMS: &[&str] = &["enhance"];
const MISSING_KEY_MESSAGE: &str = "Missing REDUCTO_API_KEY - set it in the environment or pass api_key to litellm.ocr()/litellm.aocr()";
const DATA_URI_UPLOAD_REQUIRED: &str =
"Reducto data URI upload must complete before OCR request transformation";
#[derive(Clone, Debug, PartialEq, Eq)]
pub enum ReductoDocumentSource {
FileId(String),
Upload { bytes: Vec<u8>, mime_type: String },
}
#[derive(Clone, PartialEq, Eq)]
pub struct ReductoUploadRequest {
pub url: String,
pub authorization: String,
pub file_name: &'static str,
pub bytes: Vec<u8>,
pub mime_type: String,
}
pub struct ReductoParseV3Config;
pub struct ReductoParseLegacyConfig;
pub const REDUCTO_PARSE_V3_CONFIG: ReductoParseV3Config = ReductoParseV3Config;
pub const REDUCTO_PARSE_LEGACY_CONFIG: ReductoParseLegacyConfig = ReductoParseLegacyConfig;
pub fn config_for_model(model: &str) -> Option<&'static dyn OcrProviderConfig> {
match model {
"parse-v3" => Some(&REDUCTO_PARSE_V3_CONFIG),
"parse-legacy" => Some(&REDUCTO_PARSE_LEGACY_CONFIG),
_ => None,
}
}
pub fn normalize_api_base(api_base: Option<&str>) -> String {
api_base
.map(str::trim)
.filter(|base| !base.is_empty())
.unwrap_or(REDUCTO_API_BASE)
.trim_end_matches('/')
.to_string()
}
pub fn parse_url(api_base: Option<&str>) -> String {
format!("{}/parse", normalize_api_base(api_base))
}
pub fn upload_url(api_base: Option<&str>) -> String {
format!("{}/upload", normalize_api_base(api_base))
}
pub fn resolve_api_key(
api_key: Option<&str>,
env_lookup: &dyn Fn(&str) -> Option<String>,
) -> Result<String, Error> {
api_key
.map(str::trim)
.filter(|key| !key.is_empty())
.map(str::to_string)
.or_else(|| {
env_lookup(REDUCTO_API_KEY_ENV)
.map(|key| key.trim().to_string())
.filter(|key| !key.is_empty())
})
.ok_or_else(|| Error::Auth(MISSING_KEY_MESSAGE.to_string()))
}
pub fn extract_document_source(document: &Value) -> Result<ReductoDocumentSource, Error> {
let document = document.as_object().ok_or_else(|| Error::InvalidType {
expected: "object",
actual: json_type_name(document),
})?;
let source = document
.get("document_url")
.and_then(Value::as_str)
.filter(|source| !source.is_empty())
.or_else(|| document.get("image_url").and_then(Value::as_str))
.ok_or_else(|| {
Error::InvalidRequest(
"Reducto expected OCR preprocessing to produce document_url or image_url"
.to_string(),
)
})?;
classify_document_source(source)
}
pub fn classify_document_source(source: &str) -> Result<ReductoDocumentSource, Error> {
if source.starts_with(REDUCTO_ID_PREFIX) {
return Ok(ReductoDocumentSource::FileId(source.to_string()));
}
if source.starts_with("http://") || source.starts_with("https://") {
return Err(Error::InvalidRequest(
"Reducto requires type='file' (auto-uploaded) or a reducto:// id. Plain http(s) URLs are not supported; upload the file first."
.to_string(),
));
}
if !source.starts_with("data:") {
return Err(Error::InvalidRequest(
"Reducto requires a reducto:// id or a base64 data URI after OCR preprocessing."
.to_string(),
));
}
let (header, encoded) = source
.split_once(',')
.ok_or_else(|| Error::InvalidRequest("Invalid Reducto data URI provided.".to_string()))?;
if !header.split(';').any(|part| part == "base64") {
return Err(Error::InvalidRequest(
"Reducto only supports base64-encoded data URIs.".to_string(),
));
}
let mime_type = header
.strip_prefix("data:")
.and_then(|header| header.split(';').next())
.filter(|mime| !mime.is_empty())
.unwrap_or("application/octet-stream")
.to_string();
let bytes = BASE64_STANDARD.decode(encoded).map_err(|_| {
Error::InvalidRequest("Invalid Reducto base64 payload provided.".to_string())
})?;
Ok(ReductoDocumentSource::Upload { bytes, mime_type })
}
pub fn build_upload_request(
source: ReductoDocumentSource,
authorization: &str,
api_base: Option<&str>,
) -> Option<ReductoUploadRequest> {
let ReductoDocumentSource::Upload { bytes, mime_type } = source else {
return None;
};
Some(ReductoUploadRequest {
url: upload_url(api_base),
authorization: authorization.to_string(),
file_name: "document",
bytes,
mime_type,
})
}
pub fn extract_upload_file_id(response_json: &Value) -> Result<&str, Error> {
response_json
.as_object()
.and_then(|response| response.get("file_id"))
.and_then(Value::as_str)
.filter(|file_id| !file_id.is_empty())
.ok_or_else(|| {
Error::InvalidResponse(format!(
"Reducto /upload returned 200 without a file_id; got payload={response_json}"
))
})
}
pub fn build_parse_v3_request(
file_id: &str,
optional_params: Map<String, Value>,
) -> OcrRequestData {
let data = std::iter::once(("input".to_string(), Value::String(file_id.to_string())))
.chain(optional_params)
.collect();
OcrRequestData {
data: Value::Object(data),
files: None,
}
}
pub fn build_parse_legacy_request(
file_id: &str,
optional_params: &Map<String, Value>,
) -> OcrRequestData {
let options = optional_params
.get("enhance")
.filter(|enhance| !enhance.is_null())
.map(|enhance| json!({"options": {"enhance": enhance}}));
let data = match options {
Some(Value::Object(options)) => std::iter::once((
"document_url".to_string(),
Value::String(file_id.to_string()),
))
.chain(options)
.collect(),
_ => Map::from_iter([(
"document_url".to_string(),
Value::String(file_id.to_string()),
)]),
};
OcrRequestData {
data: Value::Object(data),
files: None,
}
}
fn source_file_id(document: &Value) -> Result<String, Error> {
match extract_document_source(document)? {
ReductoDocumentSource::FileId(file_id) => Ok(file_id),
ReductoDocumentSource::Upload { .. } => Err(Error::Unsupported(DATA_URI_UPLOAD_REQUIRED)),
}
}
fn page_number(block: &Map<String, Value>) -> Option<i64> {
let page = block.get("bbox")?.as_object()?.get("page")?;
page.as_i64()
.or_else(|| page.as_u64().and_then(|page| i64::try_from(page).ok()))
.or_else(|| page.as_str().and_then(|page| page.parse().ok()))
}
fn chunks(result: &Map<String, Value>) -> &[Value] {
result
.get("chunks")
.and_then(Value::as_array)
.map(Vec::as_slice)
.unwrap_or_default()
}
fn build_pages(result: &Map<String, Value>) -> Vec<Value> {
let blocks_by_page = chunks(result)
.iter()
.filter_map(Value::as_object)
.filter_map(|chunk| chunk.get("blocks").and_then(Value::as_array))
.flatten()
.filter_map(|block| block.as_object().map(|object| (block, object)))
.filter_map(|(block, object)| page_number(object).map(|page| (page, block.clone())))
.fold(
BTreeMap::<i64, Vec<Value>>::new(),
|mut pages, (page, block)| {
pages.entry(page).or_default().push(block);
pages
},
);
if blocks_by_page.is_empty() {
let markdown = chunks(result)
.iter()
.filter_map(Value::as_object)
.filter_map(|chunk| chunk.get("content").and_then(Value::as_str))
.filter(|content| !content.is_empty())
.collect::<Vec<_>>()
.join("\n\n");
return if markdown.is_empty() {
Vec::new()
} else {
vec![json!({"index": 0, "markdown": markdown})]
};
}
blocks_by_page
.into_iter()
.map(|(page, blocks)| {
let markdown = blocks
.iter()
.filter_map(Value::as_object)
.filter_map(|block| block.get("content").and_then(Value::as_str))
.filter(|content| !content.is_empty())
.collect::<Vec<_>>()
.join("\n\n");
json!({
"index": page.saturating_sub(1).max(0),
"markdown": markdown,
"blocks": blocks,
})
})
.collect()
}
pub fn transform_reducto_response(
model: &str,
response_json: Value,
) -> Result<LiteLLMOcrResponse, Error> {
let response = response_json
.as_object()
.ok_or_else(|| Error::InvalidType {
expected: "object",
actual: json_type_name(&response_json),
})?;
let empty_result = Map::new();
let result = match response.get("result") {
Some(Value::Object(result)) => result,
Some(Value::Null) => &empty_result,
Some(_) => {
return Err(Error::InvalidResponse(
"Reducto result must be an object".to_string(),
));
}
None => response,
};
let usage = response
.get("usage")
.and_then(Value::as_object)
.cloned()
.unwrap_or_default();
let usage_info = Some(json!({
"pages_processed": usage.get("num_pages").cloned().unwrap_or(Value::Null),
"credits": usage.get("credits").cloned().unwrap_or(Value::Null),
}));
Ok(LiteLLMOcrResponse {
pages: build_pages(result),
model: model.to_string(),
document_annotation: None,
usage_info,
object: "ocr".to_string(),
extra_fields: Map::new(),
provider_native_response: Some(response_json),
})
}
impl OcrProviderConfig for ReductoParseV3Config {
fn supported_ocr_params(&self) -> &'static [&'static str] {
PARSE_V3_SUPPORTED_OCR_PARAMS
}
fn transform_ocr_request(
&self,
_model: &str,
document: Value,
optional_params: Map<String, Value>,
) -> Result<OcrRequestData, Error> {
let file_id = source_file_id(&document)?;
Ok(build_parse_v3_request(&file_id, optional_params))
}
fn transform_ocr_response(
&self,
model: &str,
response_json: Value,
) -> Result<LiteLLMOcrResponse, Error> {
transform_reducto_response(model, response_json)
}
fn complete_url(
&self,
api_base: Option<&str>,
_model: &str,
_optional_params: &Map<String, Value>,
_env_lookup: &dyn Fn(&str) -> Option<String>,
) -> Result<String, Error> {
Ok(parse_url(api_base))
}
fn resolve_api_key(
&self,
api_key: Option<&str>,
env_lookup: &dyn Fn(&str) -> Option<String>,
) -> Result<String, Error> {
resolve_api_key(api_key, env_lookup)
}
}
impl OcrProviderConfig for ReductoParseLegacyConfig {
fn supported_ocr_params(&self) -> &'static [&'static str] {
PARSE_LEGACY_SUPPORTED_OCR_PARAMS
}
fn transform_ocr_request(
&self,
_model: &str,
document: Value,
optional_params: Map<String, Value>,
) -> Result<OcrRequestData, Error> {
let file_id = source_file_id(&document)?;
Ok(build_parse_legacy_request(&file_id, &optional_params))
}
fn transform_ocr_response(
&self,
model: &str,
response_json: Value,
) -> Result<LiteLLMOcrResponse, Error> {
transform_reducto_response(model, response_json)
}
fn complete_url(
&self,
api_base: Option<&str>,
_model: &str,
_optional_params: &Map<String, Value>,
_env_lookup: &dyn Fn(&str) -> Option<String>,
) -> Result<String, Error> {
Ok(parse_url(api_base))
}
fn resolve_api_key(
&self,
api_key: Option<&str>,
env_lookup: &dyn Fn(&str) -> Option<String>,
) -> Result<String, Error> {
resolve_api_key(api_key, env_lookup)
}
}

View file

@ -1 +0,0 @@
pub mod ocr;

View file

@ -1 +0,0 @@
pub mod transformation;

View file

@ -1,467 +0,0 @@
use crate::error::{Error, json_type_name};
use crate::ocr::transformation::OcrProviderConfig;
use crate::ocr::types::{LiteLLMOcrResponse, OcrRequestData};
use serde_json::{Map, Value, json};
use crate::providers::mistral::ocr::transformation::MISTRAL_OCR_CONFIG;
const VERTEX_DEFAULT_LOCATION: &str = "us-central1";
const VERTEX_DEFAULT_DEEPSEEK_API_BASE: &str = "https://aiplatform.googleapis.com";
const VERTEX_AI_API_KEY_ENV: &str = "VERTEX_AI_API_KEY";
const VERTEXAI_API_KEY_ENV: &str = "VERTEXAI_API_KEY";
const VERTEXAI_PROJECT_ENV: &str = "VERTEXAI_PROJECT";
const VERTEXAI_LOCATION_ENV: &str = "VERTEXAI_LOCATION";
const VERTEX_LOCATION_ENV: &str = "VERTEX_LOCATION";
#[rustfmt::skip]
const DEEPSEEK_SUPPORTED_OCR_PARAMS: &[&str] = &[
"stream",
"temperature",
"max_tokens",
"top_p",
"n",
"stop",
];
pub struct VertexAiOcrConfig;
pub struct VertexAiDeepSeekOcrConfig;
pub const VERTEX_AI_OCR_CONFIG: VertexAiOcrConfig = VertexAiOcrConfig;
pub const VERTEX_AI_DEEPSEEK_OCR_CONFIG: VertexAiDeepSeekOcrConfig = VertexAiDeepSeekOcrConfig;
fn string_param<'a>(params: &'a Map<String, Value>, keys: &[&str]) -> Option<&'a str> {
keys.iter()
.find_map(|key| params.get(*key).and_then(Value::as_str))
.map(str::trim)
.filter(|value| !value.is_empty())
}
pub fn is_deepseek_model(model: &str) -> bool {
model.to_ascii_lowercase().contains("deepseek")
}
pub fn resolve_vertex_api_key(
api_key: Option<&str>,
env_lookup: &dyn Fn(&str) -> Option<String>,
) -> Result<String, Error> {
api_key
.map(str::trim)
.filter(|key| !key.is_empty())
.map(str::to_string)
.or_else(|| env_lookup(VERTEX_AI_API_KEY_ENV).filter(|key| !key.trim().is_empty()))
.or_else(|| env_lookup(VERTEXAI_API_KEY_ENV).filter(|key| !key.trim().is_empty()))
.ok_or_else(|| {
Error::Auth(
"Missing Vertex AI access token - pass api_key or provide Authorization via extra_headers"
.to_string(),
)
})
}
fn vertex_project(
params: &Map<String, Value>,
env_lookup: &dyn Fn(&str) -> Option<String>,
) -> Result<String, Error> {
string_param(params, &["vertex_project", "vertex_ai_project"])
.map(str::to_string)
.or_else(|| env_lookup(VERTEXAI_PROJECT_ENV).filter(|value| !value.trim().is_empty()))
.ok_or_else(|| {
Error::InvalidRequest(
"Missing vertex_project - Set VERTEXAI_PROJECT environment variable or pass vertex_project parameter"
.to_string(),
)
})
}
fn vertex_location(
params: &Map<String, Value>,
env_lookup: &dyn Fn(&str) -> Option<String>,
) -> String {
string_param(params, &["vertex_location", "vertex_ai_location"])
.map(str::to_string)
.or_else(|| env_lookup(VERTEXAI_LOCATION_ENV).filter(|value| !value.trim().is_empty()))
.or_else(|| env_lookup(VERTEX_LOCATION_ENV).filter(|value| !value.trim().is_empty()))
.unwrap_or_else(|| VERTEX_DEFAULT_LOCATION.to_string())
}
fn vertex_mistral_api_base(api_base: Option<&str>, location: &str) -> String {
api_base
.map(str::trim)
.filter(|value| !value.is_empty())
.map(str::to_string)
.unwrap_or_else(|| format!("https://{location}-aiplatform.googleapis.com"))
.trim_end_matches('/')
.to_string()
}
pub fn complete_vertex_mistral_url(
api_base: Option<&str>,
model: &str,
optional_params: &Map<String, Value>,
env_lookup: &dyn Fn(&str) -> Option<String>,
) -> Result<String, Error> {
let project = vertex_project(optional_params, env_lookup)?;
let location = vertex_location(optional_params, env_lookup);
let base = vertex_mistral_api_base(api_base, &location);
Ok(format!(
"{base}/v1/projects/{project}/locations/{location}/publishers/mistralai/models/{model}:rawPredict"
))
}
pub fn complete_vertex_deepseek_url(
api_base: Option<&str>,
optional_params: &Map<String, Value>,
env_lookup: &dyn Fn(&str) -> Option<String>,
) -> Result<String, Error> {
let project = vertex_project(optional_params, env_lookup)?;
let location = vertex_location(optional_params, env_lookup);
let base = api_base
.map(str::trim)
.filter(|value| !value.is_empty())
.unwrap_or(VERTEX_DEFAULT_DEEPSEEK_API_BASE)
.trim_end_matches('/');
Ok(format!(
"{base}/v1/projects/{project}/locations/{location}/endpoints/openapi/chat/completions"
))
}
fn document_content_item(document: &Value) -> Result<Value, Error> {
let object = document.as_object().ok_or_else(|| Error::InvalidType {
expected: "object",
actual: json_type_name(document),
})?;
let doc_type = object
.get("type")
.and_then(Value::as_str)
.ok_or(Error::MissingField("document.type"))?;
let url_field = match doc_type {
"image_url" => "image_url",
"document_url" => "document_url",
other => {
return Err(Error::InvalidRequest(format!(
"Unsupported document type: {other}. Expected 'image_url' or 'document_url'"
)));
}
};
let url = object
.get(url_field)
.and_then(Value::as_str)
.filter(|value| !value.is_empty())
.ok_or(Error::MissingField(url_field))?;
Ok(json!({
"type": "image_url",
"image_url": url,
}))
}
fn deepseek_model_name(model: &str) -> String {
if model.starts_with("deepseek-ai/") {
model.to_string()
} else {
format!("deepseek-ai/{model}")
}
}
fn first_choice_content(response: &Value) -> Result<Value, Error> {
response
.get("choices")
.and_then(Value::as_array)
.and_then(|choices| choices.first())
.and_then(|choice| choice.get("message"))
.and_then(|message| message.get("content"))
.cloned()
.filter(|content| match content {
Value::String(value) => !value.is_empty(),
Value::Object(_) => true,
_ => false,
})
.ok_or_else(|| Error::InvalidResponse("No content in DeepSeek OCR response".to_string()))
}
fn ocr_data_from_content(content: Value, usage: Option<Value>, model: &str) -> Value {
match content {
Value::String(content) => {
if content.trim_start().starts_with('{') {
serde_json::from_str(&content).unwrap_or_else(|_| {
json!({
"pages": [{"index": 0, "markdown": content}],
"model": model,
"usage_info": usage.unwrap_or_else(|| json!({})),
})
})
} else {
json!({
"pages": [{"index": 0, "markdown": content}],
"model": model,
"usage_info": usage.unwrap_or_else(|| json!({})),
})
}
}
Value::Object(_) => content,
other => json!({
"pages": [{"index": 0, "markdown": other.to_string()}],
"model": model,
"usage_info": usage.unwrap_or_else(|| json!({})),
}),
}
}
impl OcrProviderConfig for VertexAiOcrConfig {
fn supported_ocr_params(&self) -> &'static [&'static str] {
MISTRAL_OCR_CONFIG.supported_ocr_params()
}
#[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)]
fn transform_ocr_request(
&self,
model: &str,
document: Value,
optional_params: Map<String, Value>,
) -> Result<OcrRequestData, Error> {
MISTRAL_OCR_CONFIG.transform_ocr_request(model, document, optional_params)
}
fn transform_ocr_response(
&self,
model: &str,
response_json: Value,
) -> Result<LiteLLMOcrResponse, Error> {
MISTRAL_OCR_CONFIG.transform_ocr_response(model, response_json)
}
#[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)]
fn complete_url(
&self,
api_base: Option<&str>,
model: &str,
optional_params: &Map<String, Value>,
env_lookup: &dyn Fn(&str) -> Option<String>,
) -> Result<String, Error> {
complete_vertex_mistral_url(api_base, model, optional_params, env_lookup)
}
fn resolve_api_key(
&self,
api_key: Option<&str>,
env_lookup: &dyn Fn(&str) -> Option<String>,
) -> Result<String, Error> {
resolve_vertex_api_key(api_key, env_lookup)
}
fn requires_data_uri_document(&self) -> bool {
true
}
}
impl OcrProviderConfig for VertexAiDeepSeekOcrConfig {
#[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)]
fn supported_ocr_params(&self) -> &'static [&'static str] {
DEEPSEEK_SUPPORTED_OCR_PARAMS
}
#[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)]
fn map_ocr_params(&self, non_default_params: &Map<String, Value>) -> Map<String, Value> {
non_default_params
.iter()
.filter(|(name, _)| DEEPSEEK_SUPPORTED_OCR_PARAMS.contains(&name.as_str()))
.map(|(name, value)| (name.clone(), value.clone()))
.collect()
}
#[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)]
fn transform_ocr_request(
&self,
model: &str,
document: Value,
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)?]}]),
);
for (key, value) in optional_params {
if DEEPSEEK_SUPPORTED_OCR_PARAMS.contains(&key.as_str()) {
data.insert(key, value);
}
}
Ok(OcrRequestData {
data: Value::Object(data),
files: None,
})
}
#[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)]
fn transform_ocr_response(
&self,
model: &str,
response_json: Value,
) -> Result<LiteLLMOcrResponse, Error> {
let response = response_json
.as_object()
.ok_or_else(|| Error::InvalidType {
expected: "object",
actual: json_type_name(&response_json),
})?;
let usage = response.get("usage").cloned();
let content = first_choice_content(&response_json)?;
let mut ocr_data = ocr_data_from_content(content.clone(), usage.clone(), model);
if !ocr_data.get("pages").is_some_and(Value::is_array) {
ocr_data = json!({
"pages": [{
"index": 0,
"markdown": match content {
Value::String(value) => value,
other => other.to_string(),
}
}],
"model": ocr_data.get("model").and_then(Value::as_str).unwrap_or(model),
"usage_info": ocr_data.get("usage_info").cloned().or(usage).unwrap_or_else(|| json!({})),
});
}
let object = ocr_data.as_object().ok_or_else(|| Error::InvalidType {
expected: "object",
actual: json_type_name(&ocr_data),
})?;
let pages = object
.get("pages")
.and_then(Value::as_array)
.cloned()
.unwrap_or_default();
let usage_info = object
.get("usage_info")
.cloned()
.or_else(|| response.get("usage").cloned());
Ok(LiteLLMOcrResponse {
pages,
model: object
.get("model")
.and_then(Value::as_str)
.unwrap_or(model)
.to_string(),
document_annotation: object.get("document_annotation").cloned(),
usage_info,
object: "ocr".to_string(),
extra_fields: Map::new(),
provider_native_response: None,
})
}
#[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)]
fn complete_url(
&self,
api_base: Option<&str>,
_model: &str,
optional_params: &Map<String, Value>,
env_lookup: &dyn Fn(&str) -> Option<String>,
) -> Result<String, Error> {
complete_vertex_deepseek_url(api_base, optional_params, env_lookup)
}
fn resolve_api_key(
&self,
api_key: Option<&str>,
env_lookup: &dyn Fn(&str) -> Option<String>,
) -> Result<String, Error> {
resolve_vertex_api_key(api_key, env_lookup)
}
}
#[cfg(test)]
mod tests {
use super::*;
use rstest::rstest;
#[test]
fn vertex_mistral_url_uses_project_location_and_model() {
let params = Map::from_iter([
("vertex_project".to_string(), json!("proj-1")),
("vertex_location".to_string(), json!("europe-west4")),
]);
let url = complete_vertex_mistral_url(None, "mistral-ocr-maas", &params, &|_| None)
.expect("url builds");
assert_eq!(
url,
"https://europe-west4-aiplatform.googleapis.com/v1/projects/proj-1/locations/europe-west4/publishers/mistralai/models/mistral-ocr-maas:rawPredict"
);
}
#[test]
fn vertex_mistral_reuses_mistral_body_transform() {
let body = VERTEX_AI_OCR_CONFIG
.transform_ocr_request(
"mistral-ocr-maas",
json!({"type": "image_url", "image_url": "data:image/png;base64,abc"}),
Map::new(),
)
.expect("request transforms")
.data;
assert_eq!(body["model"], "mistral-ocr-maas");
assert_eq!(body["document"]["image_url"], "data:image/png;base64,abc");
}
#[test]
fn vertex_deepseek_request_uses_ocr_endpoint_shape() {
let body = VERTEX_AI_DEEPSEEK_OCR_CONFIG
.transform_ocr_request(
"deepseek-ocr-maas",
json!({"type": "document_url", "document_url": "gs://bucket/doc.pdf"}),
Map::from_iter([("temperature".to_string(), json!(0.1))]),
)
.expect("request transforms")
.data;
assert_eq!(body["model"], "deepseek-ai/deepseek-ocr-maas");
assert_eq!(body["temperature"], 0.1);
assert_eq!(
body["messages"][0]["content"][0],
json!({"type": "image_url", "image_url": "gs://bucket/doc.pdf"})
);
}
#[rstest]
#[case::bare_model("deepseek-ocr-maas")]
#[case::namespaced_model("deepseek-ai/deepseek-ocr-maas")]
fn vertex_deepseek_request_uses_single_provider_namespace(#[case] model: &str) {
let body = VERTEX_AI_DEEPSEEK_OCR_CONFIG
.transform_ocr_request(
model,
json!({"type": "image_url", "image_url": "data:image/png;base64,AA=="}),
Map::new(),
)
.expect("request transforms")
.data;
assert_eq!(body["model"], "deepseek-ai/deepseek-ocr-maas");
}
#[test]
fn vertex_deepseek_response_wraps_markdown_content() {
let response = VERTEX_AI_DEEPSEEK_OCR_CONFIG
.transform_ocr_response(
"deepseek-ocr-maas",
json!({
"choices": [{"message": {"content": "# OCR text"}}],
"usage": {"prompt_tokens": 1}
}),
)
.expect("response transforms");
assert_eq!(
response.pages,
vec![json!({"index": 0, "markdown": "# OCR text"})]
);
assert_eq!(response.model, "deepseek-ocr-maas");
assert_eq!(response.usage_info, Some(json!({"prompt_tokens": 1})));
}
}

View file

@ -60,6 +60,14 @@ impl ApiUrl<Base> {
}
impl ApiUrl<Complete> {
pub(crate) fn append_query_pairs<'a>(
mut self,
pairs: impl IntoIterator<Item = (&'a str, &'a str)>,
) -> Self {
self.url.query_pairs_mut().extend_pairs(pairs);
self
}
pub(crate) fn into_string(self) -> String {
self.url.into()
}
@ -92,4 +100,19 @@ mod tests {
.expect("url builds");
assert_eq!(actual, "https://example.test/v1/ocr?tenant=a");
}
#[test]
fn appended_query_pairs_are_encoded() {
let actual = ApiUrl::parse("https://example.test")
.and_then(|url| url.complete_path(&["analyze"]))
.map(|url| {
url.append_query_pairs([("model", "name with spaces")])
.into_string()
})
.expect("url builds");
assert_eq!(
actual,
"https://example.test/analyze?model=name+with+spaces"
);
}
}

View file

@ -0,0 +1,97 @@
use std::sync::Arc;
use serde_json::{Value, json};
use super::hooks::{OcrDuringCallRequest, OcrHookFuture, OcrHooks};
use super::test_support::{MockResponse, mock_server, perform_ocr, wire_request};
#[tokio::test]
async fn facade_executes_azure_mistral_with_prepared_auth() {
let (base, seen, server) = mock_server(vec![MockResponse::json(json!({
"pages":[{"index":0,"markdown":"hello"}],
"usage_info":{"pages_processed":1}
}))])
.await;
let mut request = wire_request(
"azure_ai/model",
&base,
json!({"include_image_base64":true}),
);
request.connection.api_key = None;
request.connection.extra_headers = vec![(
"Authorization".into(),
"Bearer python-prepared-token".into(),
)];
let result = perform_ocr(request).await.unwrap();
server.await.unwrap();
assert_eq!(result.pages[0]["markdown"], "hello");
let requests = seen.lock().unwrap();
assert_eq!(requests.len(), 1);
assert!(requests[0].starts_with("POST /providers/mistral/azure/ocr "));
assert!(
requests[0]
.to_ascii_lowercase()
.contains("authorization: bearer python-prepared-token\r\n")
);
let body: Value = serde_json::from_str(requests[0].split_once("\r\n\r\n").unwrap().1).unwrap();
assert_eq!(
body,
json!({
"model":"model",
"document":{"type":"document_url","document_url":"data:application/pdf;base64,YWJj"},
"include_image_base64":true
})
);
}
#[tokio::test]
async fn facade_acquires_supplied_entra_token_for_final_request() {
let (base, seen, server) = mock_server(vec![MockResponse::json(json!({"pages":[]}))]).await;
let mut request = wire_request(
"azure_ai/model",
&base,
json!({"azure_ad_token":"rust-owned-token"}),
);
request.connection.api_key = None;
perform_ocr(request).await.unwrap();
server.await.unwrap();
let requests = seen.lock().unwrap();
assert_eq!(requests.len(), 1);
assert!(
requests[0]
.to_ascii_lowercase()
.contains("authorization: bearer rust-owned-token\r\n")
);
}
struct ReplaceBodyDocument;
impl OcrHooks for ReplaceBodyDocument {
fn has_guardrails(&self) -> bool {
true
}
fn during_call(
&self,
mut request: OcrDuringCallRequest,
) -> OcrHookFuture<'_, OcrDuringCallRequest> {
Box::pin(async move {
request.body["document"] = json!({
"type":"document_url",
"document_url":"https://example.com/not-inline.pdf"
});
Ok(request)
})
}
}
#[tokio::test]
async fn rejects_non_inline_body_after_guardrails() {
let mut request = wire_request("azure_ai/model", "http://127.0.0.1:1", json!({}));
request.hooks = Arc::new(ReplaceBodyDocument);
let error = perform_ocr(request).await.unwrap_err();
assert!(error.to_string().contains("data URI"));
}

View file

@ -0,0 +1,395 @@
use serde_json::{Value, json};
use super::test_support::{MockResponse, mock_server, perform_ocr, wire_request};
use super::wire::{OcrWireRequest, decode_request};
fn query_value(url: &str, key: &str) -> Option<String> {
url::Url::parse(url)
.unwrap()
.query_pairs()
.find_map(|(name, value)| (name == key).then(|| value.into_owned()))
}
#[tokio::test]
async fn facade_maps_pages_features_and_url_document() {
let (base, seen, server) = mock_server(vec![MockResponse::json(json!({
"status":"succeeded",
"analyzeResult":{"pages":[]}
}))])
.await;
let mut request = wire_request(
"azure_ai/doc-intelligence/prebuilt-read",
&base,
json!({"pages":[2,0,0,1],"features":["keyValuePairs","languages"]}),
);
request.document = serde_json::from_value(json!({
"type":"document_url",
"document_url":"https://example.com/document.pdf"
}))
.unwrap();
perform_ocr(request).await.unwrap();
server.await.unwrap();
let request = &seen.lock().unwrap()[0];
let target = request.split_whitespace().nth(1).unwrap();
let url = format!("{base}{target}");
assert_eq!(query_value(&url, "pages").as_deref(), Some("1,2,3"));
assert_eq!(
query_value(&url, "features").as_deref(),
Some("keyValuePairs,languages")
);
let body: Value = serde_json::from_str(request.split_once("\r\n\r\n").unwrap().1).unwrap();
assert_eq!(
body,
json!({"urlSource":"https://example.com/document.pdf"})
);
}
#[tokio::test]
async fn rejects_invalid_pages_features_and_format() {
for options in [
json!({"pages":[true]}),
json!({"pages":[1,"2"]}),
json!({"pages":[-1]}),
json!({"pages":"1&&features=bad"}),
json!({"features":"languages&pages=1"}),
json!({"req_format":"azure"}),
] {
let result = decode_request(OcrWireRequest {
model: "azure_ai/doc-intelligence/prebuilt-read".into(),
document: json!({"type":"document_url","document_url":"https://example.com/a.pdf"}),
api_key: Some("key".into()),
api_base: Some("http://127.0.0.1:1".into()),
custom_llm_provider: None,
extra_headers: None,
optional_params: options.as_object().unwrap().clone(),
input_sources: Default::default(),
timeout_seconds: None,
});
let rejected = match result {
Ok(request) => perform_ocr(request).await.is_err(),
Err(_) => true,
};
assert!(rejected, "accepted {options}");
}
}
#[tokio::test]
async fn inline_document_decodes_to_base64_source() {
let (base, seen, server) = mock_server(vec![MockResponse::json(json!({
"status":"succeeded"
}))])
.await;
let request = wire_request("azure_ai/doc-intelligence/prebuilt-read", &base, json!({}));
perform_ocr(request).await.unwrap();
server.await.unwrap();
let request = &seen.lock().unwrap()[0];
let body: Value = serde_json::from_str(request.split_once("\r\n\r\n").unwrap().1).unwrap();
assert_eq!(body, json!({"base64Source":"YWJj"}));
}
#[tokio::test]
async fn immediate_response_normalizes_pages_and_preserves_native() {
let operation = json!({
"status":"succeeded",
"operationExtension":42,
"analyzeResult":{
"content":"A\n\nB",
"tables":[{"cells":[]}],
"keyValuePairs":[{"key":{"content":"A"}}],
"pages":[{
"pageNumber":"2",
"width":"8.5",
"height":11,
"unit":"inch",
"lines":[{"content":"A"},{"content":null},{"content":"B"}]
}]
}
});
let (base, _, server) = mock_server(vec![MockResponse::json(operation.clone())]).await;
let result = perform_ocr(wire_request(
"azure_ai/doc-intelligence/prebuilt-read",
&base,
json!({"req_format":"native"}),
))
.await
.unwrap();
server.await.unwrap();
assert_eq!(result.pages[0]["index"], 1);
assert_eq!(result.pages[0]["markdown"], "A\n\nB");
assert_eq!(
result.pages[0]["dimensions"],
json!({"width":816,"height":1056,"dpi":96})
);
assert_eq!(result.usage_info, Some(json!({"pages_processed":1})));
assert_eq!(result.provider_native_response, Some(operation));
}
#[tokio::test]
async fn accepted_response_polls_to_success_with_only_credentials() {
let operation = json!({"status":"succeeded","analyzeResult":{"pages":[]}});
let (base, seen, server) = mock_server(vec![
MockResponse {
status: 202,
headers: vec![("Operation-Location", "{base}/operation".into())],
body: json!({}),
},
MockResponse {
status: 200,
headers: vec![("Retry-After", "0".into())],
body: json!({"status":"running"}),
},
MockResponse::json(operation.clone()),
])
.await;
let mut request = wire_request(
"azure_ai/doc-intelligence/prebuilt-read",
&base,
json!({"req_format":"native"}),
);
request
.connection
.extra_headers
.push(("X-Trace".into(), "initial-only".into()));
let result = perform_ocr(request).await.unwrap();
server.await.unwrap();
assert_eq!(result.provider_native_response, Some(operation));
let requests = seen.lock().unwrap();
assert_eq!(requests.len(), 3);
assert!(requests[0].to_ascii_lowercase().contains("x-trace:"));
for poll in &requests[1..] {
assert!(!poll.to_ascii_lowercase().contains("x-trace:"));
assert!(
poll.to_ascii_lowercase()
.contains("ocp-apim-subscription-key: test-key")
);
}
}
#[tokio::test]
async fn polling_forwards_bearer_credentials() {
let (base, seen, server) = mock_server(vec![
MockResponse {
status: 202,
headers: vec![("Operation-Location", "{base}/operation".into())],
body: json!({}),
},
MockResponse::json(json!({"status":"succeeded"})),
])
.await;
let mut request = wire_request("azure_ai/doc-intelligence/prebuilt-read", &base, json!({}));
request.connection.api_key = None;
request.connection.extra_headers = vec![("Authorization".into(), "Bearer token".into())];
perform_ocr(request).await.unwrap();
server.await.unwrap();
let requests = seen.lock().unwrap();
assert!(
requests[1]
.to_ascii_lowercase()
.contains("authorization: bearer token")
);
}
#[tokio::test]
async fn polling_does_not_follow_redirects() {
let (base, seen, server) = mock_server(vec![
MockResponse {
status: 202,
headers: vec![("Operation-Location", "{base}/operation".into())],
body: json!({}),
},
MockResponse {
status: 302,
headers: vec![("Location", "{base}/redirected".into())],
body: json!({}),
},
MockResponse::json(json!({"status":"succeeded"})),
])
.await;
let error = perform_ocr(wire_request(
"azure_ai/doc-intelligence/prebuilt-read",
&base,
json!({}),
))
.await
.unwrap_err();
assert!(error.to_string().contains("status 302"), "{error}");
assert_eq!(seen.lock().unwrap().len(), 2);
server.abort();
}
#[tokio::test]
async fn polling_rejects_terminal_failure() {
let (base, _, server) = mock_server(vec![
MockResponse {
status: 202,
headers: vec![("Operation-Location", "{base}/operation".into())],
body: json!({}),
},
MockResponse::json(json!({"status":"failed"})),
])
.await;
let error = perform_ocr(wire_request(
"azure_ai/doc-intelligence/prebuilt-read",
&base,
json!({}),
))
.await
.unwrap_err();
server.await.unwrap();
assert!(error.to_string().contains("status failed"));
}
#[tokio::test]
async fn malformed_provider_pages_report_response_paths() {
for (analysis, path) in [
(json!({"pages":null}), "pages"),
(json!({"pages":[null]}), "pages[0]"),
(json!({"pages":[{"lines":null}]}), "lines"),
(json!({"pages":[{"width":"bad"}]}), "width"),
] {
let (base, _, server) = mock_server(vec![MockResponse::json(json!({
"status":"succeeded",
"analyzeResult":analysis
}))])
.await;
let error = perform_ocr(wire_request(
"azure_ai/doc-intelligence/prebuilt-read",
&base,
json!({}),
))
.await
.unwrap_err();
server.await.unwrap();
assert!(error.to_string().contains(path), "{error}");
}
}
#[tokio::test]
async fn rejects_missing_invalid_and_cross_origin_operation_locations() {
for headers in [
Vec::new(),
vec![("Operation-Location", "/relative".into())],
vec![("Operation-Location", "http://example.com/operation".into())],
vec![(
"Operation-Location",
"http://user:password@127.0.0.1/operation".into(),
)],
] {
let (base, _, server) = mock_server(vec![MockResponse {
status: 202,
headers,
body: json!({}),
}])
.await;
let error = perform_ocr(wire_request(
"azure_ai/doc-intelligence/prebuilt-read",
&base,
json!({}),
))
.await
.unwrap_err();
server.await.unwrap();
assert!(error.to_string().contains("operation-location"));
}
}
#[tokio::test]
async fn polling_deadline_bounds_retry_delay() {
let (base, _, server) = mock_server(vec![
MockResponse {
status: 202,
headers: vec![("Operation-Location", "{base}/operation".into())],
body: json!({}),
},
MockResponse {
status: 200,
headers: vec![("Retry-After", "9999".into())],
body: json!({"status":"notStarted"}),
},
])
.await;
let mut request = wire_request("azure_ai/doc-intelligence/prebuilt-read", &base, json!({}));
request.connection.poll_timeout = std::time::Duration::from_millis(100);
let error = tokio::time::timeout(std::time::Duration::from_secs(1), perform_ocr(request))
.await
.unwrap()
.unwrap_err();
server.await.unwrap();
assert!(error.to_string().contains("timed out"));
}
#[tokio::test]
async fn model_id_is_encoded_and_dot_segments_are_rejected() {
let (base, seen, server) = mock_server(vec![MockResponse::json(json!({
"status":"succeeded"
}))])
.await;
perform_ocr(wire_request(
"azure_ai/doc-intelligence/a ?#é",
&base,
json!({}),
))
.await
.unwrap();
server.await.unwrap();
assert!(seen.lock().unwrap()[0].contains("a%20%3F%23%C3%A9:analyze"));
for model in [
"azure_ai/doc-intelligence/.",
"azure_ai/doc-intelligence/..",
] {
let error = perform_ocr(wire_request(model, "http://127.0.0.1:1", json!({})))
.await
.unwrap_err();
assert!(error.to_string().contains("dot segment"));
}
}
#[tokio::test]
async fn pre_call_guardrail_receives_caller_pages_before_mapping() {
use crate::ocr::hooks::{OcrHookFuture, OcrHooks, OcrPreCallRequest};
use std::sync::Arc;
struct RewritePages;
impl OcrHooks for RewritePages {
fn has_guardrails(&self) -> bool {
true
}
fn pre_call(&self, request: OcrPreCallRequest) -> OcrHookFuture<'_, OcrPreCallRequest> {
Box::pin(async move {
assert_eq!(request.optional_params["pages"], json!([0, 2]));
Ok(OcrPreCallRequest {
optional_params: json!({"pages": [1]}),
..request
})
})
}
}
let (base, seen, server) =
mock_server(vec![MockResponse::json(json!({"status": "succeeded"}))]).await;
let request = wire_request(
"azure_ai/doc-intelligence/prebuilt-read",
&base,
json!({"pages": [0, 2]}),
)
.with_host_hooks(Arc::new(RewritePages), None);
perform_ocr(request).await.unwrap();
server.await.unwrap();
let requests = seen.lock().unwrap();
let target = requests[0].split_whitespace().nth(1).unwrap();
assert_eq!(
query_value(&format!("{base}{target}"), "pages").as_deref(),
Some("2")
);
assert_eq!(requests.len(), 1);
}

View file

@ -0,0 +1,95 @@
use rstest::rstest;
use serde_json::{Value, json};
use crate::ocr::codecs::deepseek::{
DeepSeekOcrParams, DeepSeekOcrResponse, transform_ocr_request, transform_ocr_response,
};
use crate::ocr::types::OcrDocument;
fn document() -> OcrDocument {
serde_json::from_value(json!({"type":"image_url","image_url":"gs://bucket/a.png"})).unwrap()
}
#[rstest]
#[case("stream", json!(true))]
#[case("temperature", json!(0.1))]
#[case("max_tokens", json!(1024))]
#[case("top_p", json!(0.9))]
#[case("n", json!(2))]
#[case("stop", json!("done"))]
#[case("stop", json!(["done", "stop"]))]
fn request_mapping_matches_python(#[case] name: &str, #[case] value: Value) {
let params: DeepSeekOcrParams =
serde_json::from_value(json!({name: value.clone(), "ignored": true})).unwrap();
let result = serde_json::to_value(
transform_ocr_request("deepseek-ai/deepseek-ocr-maas", document(), &params).unwrap(),
)
.unwrap();
assert_eq!(result["model"], "deepseek-ai/deepseek-ocr-maas");
assert_eq!(
result["messages"][0]["content"][0],
json!({"type":"image_url","image_url":"gs://bucket/a.png"})
);
assert_eq!(result[name], value);
assert!(result.get("ignored").is_none());
}
#[rstest]
#[case(json!("# hello"), "# hello")]
#[case(json!("{broken"), "{broken")]
#[case(json!(" {\"pages\":[]} "), " {\"pages\":[]} ")]
#[case(json!({"pages":[]}), "{\"pages\":[]}")]
#[case(json!({}), "{}")]
#[case(json!("[]"), "[]")]
#[case(json!("{\"pages\":[{\"markdown\":\"json text\"}]}"), "json text")]
#[case(json!({"pages":[{"markdown":"object"}]}), "object")]
fn response_codec_handles_text_json_and_objects(#[case] content: Value, #[case] expected: &str) {
let response: DeepSeekOcrResponse = serde_json::from_value(
json!({"choices":[{"message":{"content":content}}],"usage":{"prompt_tokens":1}}),
)
.unwrap();
let result = transform_ocr_response("model", response)
.unwrap()
.into_json();
assert_eq!(result["pages"][0]["markdown"], expected);
assert_eq!(result["pages"][0]["index"], 0);
assert_eq!(result["usage_info"]["prompt_tokens"], 1);
}
#[test]
fn structured_result_maps_pages_usage_model_and_annotation() {
let response: DeepSeekOcrResponse = serde_json::from_value(json!({
"choices":[{"message":{"content":{
"pages":[{"index":2,"markdown":"page","images":[{"id":"one"}],"dimensions":{"width":10}}],
"model":"provider-model",
"usage_info":{"pages_processed":1},
"document_annotation":{"language":"en"},
"future":"kept"
}}}]
}))
.unwrap();
let result = transform_ocr_response("requested", response)
.unwrap()
.into_json();
assert_eq!(result["pages"][0]["index"], 2);
assert_eq!(result["pages"][0]["images"][0]["id"], "one");
assert_eq!(result["model"], "provider-model");
assert_eq!(result["usage_info"]["pages_processed"], 1);
assert_eq!(result["document_annotation"]["language"], "en");
assert_eq!(result["future"], "kept");
}
#[test]
fn response_codec_rejects_missing_empty_and_malformed_content() {
for value in [
json!({"choices":[]}),
json!({"choices":[{"message":{"content":""}}]}),
json!({"choices":[{"message":{"content":"{\"pages\":[{\"markdown\":42}]}"}}]}),
json!({"choices":[{"message":{"content":{"pages":[{"markdown":42}]}}}]}),
] {
let result = serde_json::from_value::<DeepSeekOcrResponse>(value)
.map_err(|_| ())
.and_then(|response| transform_ocr_response("model", response).map_err(|_| ()));
assert!(result.is_err());
}
}

View file

@ -21,6 +21,7 @@ fn request_boundary_selects_mistral_and_rejects_unknown_providers() {
.as_object()
.unwrap()
.clone(),
input_sources: Default::default(),
timeout_seconds: None,
};
assert!(decode_request(request).is_ok());
@ -33,6 +34,7 @@ fn request_boundary_selects_mistral_and_rejects_unknown_providers() {
custom_llm_provider: Some("unknown".into()),
extra_headers: None,
optional_params: serde_json::Map::new(),
input_sources: Default::default(),
timeout_seconds: None,
})
.is_err()

View file

@ -8,7 +8,11 @@ use crate::ocr::wire::{OcrWireRequest, decode_request};
use crate::ocr::{LiteLLMOcrRequest, LiteLLMOcrResponse, OcrClient};
pub(crate) fn ocr_client() -> OcrClient {
OcrClient::for_test(reqwest::Client::new())
let document_http = reqwest::Client::builder()
.redirect(reqwest::redirect::Policy::none())
.build()
.expect("test document client builds");
OcrClient::for_test(reqwest::Client::new(), document_http)
}
pub(crate) async fn perform_ocr(
@ -26,6 +30,7 @@ pub(crate) fn wire_request(model: &str, base: &str, options: Value) -> LiteLLMOc
custom_llm_provider: None,
extra_headers: None,
optional_params: options.as_object().unwrap().clone(),
input_sources: Default::default(),
timeout_seconds: Some(2.0),
})
.unwrap()

View file

@ -0,0 +1,232 @@
use std::sync::Arc;
use rstest::rstest;
use serde_json::{Value, json};
use super::hooks::{OcrDuringCallRequest, OcrHookFuture, OcrHooks};
use super::test_support::{MockResponse, mock_server, perform_ocr, wire_request};
fn request_body(request: &str) -> Value {
serde_json::from_str(request.split_once("\r\n\r\n").unwrap().1).unwrap()
}
#[rstest]
#[case(
"reducto/parse-v3",
json!({
"formatting":{"table_output_format":"html"},
"retrieval":{"chunk_mode":"section"},
"settings":{"ocr_system":"standard"},
"future_ocr_option":true,
"extra_body":{"provider_option":"value"}
}),
"reducto://already.pdf",
json!({
"input":"reducto://already.pdf",
"formatting":{"table_output_format":"html"},
"retrieval":{"chunk_mode":"section"},
"settings":{"ocr_system":"standard"},
"future_ocr_option":true,
"provider_option":"value"
})
)]
#[case(
"reducto/parse-legacy",
json!({
"enhance":{"agentic":[{"type":"table"}]},
"future_ocr_option":true,
"extra_body":{"provider_option":"value"}
}),
"reducto://legacy.pdf",
json!({
"document_url":"reducto://legacy.pdf",
"options":{"enhance":{"agentic":[{"type":"table"}]}},
"future_ocr_option":true,
"provider_option":"value"
})
)]
#[tokio::test]
async fn request_mapping_matches_python(
#[case] model: &str,
#[case] options: Value,
#[case] source: &str,
#[case] expected: Value,
) {
let (base, seen, server) = mock_server(vec![MockResponse::json(json!({
"result":{"chunks":[]}
}))])
.await;
let mut request = wire_request(model, &base, options);
request.document = request.document.with_source(source.into());
perform_ocr(request).await.unwrap();
server.await.unwrap();
let requests = seen.lock().unwrap();
assert_eq!(requests.len(), 1);
assert!(requests[0].starts_with("POST /parse "));
assert_eq!(request_body(&requests[0]), expected);
}
#[rstest]
#[case("parse-v3")]
#[case("parse-legacy")]
#[tokio::test]
async fn data_uri_upload_preserves_multipart_headers(#[case] model: &str) {
let (base, seen, server) = mock_server(vec![
MockResponse::json(json!({"file_id":"reducto://uploaded.pdf"})),
MockResponse::json(json!({"result":{"chunks":[{"content":"hello"}]}})),
])
.await;
let mut request = wire_request(&format!("reducto/{model}"), &base, json!({}));
request.connection.extra_headers = vec![
("Content-Type".into(), "application/json".into()),
("X-Trace".into(), "upload-test".into()),
];
let response = perform_ocr(request).await.unwrap();
server.await.unwrap();
assert_eq!(response.pages[0]["markdown"], "hello");
let requests = seen.lock().unwrap();
assert_eq!(requests.len(), 2);
assert!(requests[0].starts_with("POST /upload "));
assert!(
requests[0]
.to_ascii_lowercase()
.contains("content-type: multipart/form-data; boundary=")
);
assert!(requests[0].contains("x-trace: upload-test"));
assert!(requests[0].contains("application/pdf"));
assert!(requests[0].contains("abc"));
assert!(requests[1].starts_with("POST /parse "));
}
#[rstest]
#[case(json!({"file_id":""}))]
#[case(json!({}))]
#[case(json!({"file_id":null}))]
#[tokio::test]
async fn invalid_upload_ids_stop_before_parse(#[case] response: Value) {
let (base, seen, server) = mock_server(vec![MockResponse::json(response)]).await;
let error = perform_ocr(wire_request("reducto/parse-v3", &base, json!({})))
.await
.unwrap_err();
server.await.unwrap();
assert!(error.to_string().contains("file_id"));
assert_eq!(seen.lock().unwrap().len(), 1);
}
#[tokio::test]
async fn upload_failure_stops_before_parse() {
let (base, seen, server) = mock_server(vec![MockResponse {
status: 503,
headers: vec![],
body: json!({"error":"unavailable"}),
}])
.await;
assert!(
perform_ocr(wire_request("reducto/parse-v3", &base, json!({})))
.await
.is_err()
);
server.await.unwrap();
assert_eq!(seen.lock().unwrap().len(), 1);
}
#[rstest]
#[case("https://example.com/a.pdf")]
#[case("reducto://")]
#[case("data:application/pdf;base64")]
#[case("data:application/pdf;base64,INVALID!")]
#[tokio::test]
async fn rejects_invalid_document_sources_before_network(#[case] source: &str) {
let mut request = wire_request("reducto/parse-v3", "http://127.0.0.1:1", json!({}));
request.document = request.document.with_source(source.into());
assert!(perform_ocr(request).await.is_err());
}
#[test]
fn response_normalization_groups_blocks_and_distinguishes_null_result() {
use crate::ocr::codecs::reducto::{ReductoResponse, transform_ocr_response};
let raw = json!({"usage":{"num_pages":"2","credits":"3"},"result":{"chunks":[
{"blocks":[{"content":"B","bbox":{"page":2},"kind":"table"}]},
{"blocks":[{"content":"A","bbox":{"page":1},"kind":"text"},{"content":"C","bbox":{"page":1}}]}
]}});
let response: ReductoResponse = serde_json::from_value(raw).unwrap();
let normalized = transform_ocr_response("parse-v3", response)
.unwrap()
.into_json();
assert_eq!(normalized["pages"][0]["markdown"], "A\n\nC");
assert_eq!(normalized["pages"][1]["markdown"], "B");
assert_eq!(normalized["pages"][1]["blocks"][0]["kind"], "table");
assert_eq!(normalized["usage_info"]["pages_processed"], 2);
assert_eq!(normalized["usage_info"]["credits"], 3.0);
let missing: ReductoResponse =
serde_json::from_value(json!({"chunks":[{"content":"text"}]})).unwrap();
let missing = transform_ocr_response("parse-v3", missing).unwrap();
assert_eq!(missing.pages[0]["markdown"], "text");
let null: ReductoResponse = serde_json::from_value(
json!({"result":null,"chunks":[{"content":"ignored"}],"usage":null}),
)
.unwrap();
let null = transform_ocr_response("parse-v3", null).unwrap();
assert!(null.pages.is_empty());
}
#[tokio::test]
async fn facade_omits_native_response_by_default_and_preserves_auth_priority() {
let raw = json!({"job_id":"job-1","result":{"chunks":[]}});
let (base, seen, server) = mock_server(vec![MockResponse::json(raw)]).await;
let mut request = wire_request("reducto/parse-v3", &base, json!({}));
request.document = request.document.with_source("reducto://ready.pdf".into());
request.connection.extra_headers = vec![("authorization".into(), "Bearer existing".into())];
let response = perform_ocr(request).await.unwrap();
server.await.unwrap();
assert_eq!(response.provider_native_response, None);
assert!(
seen.lock().unwrap()[0]
.to_ascii_lowercase()
.contains("authorization: bearer existing")
);
}
struct RewriteDocument;
impl OcrHooks for RewriteDocument {
fn has_guardrails(&self) -> bool {
true
}
fn during_call(
&self,
request: OcrDuringCallRequest,
) -> OcrHookFuture<'_, OcrDuringCallRequest> {
Box::pin(async move {
assert_eq!(
request.body["document_url"],
"data:application/pdf;base64,YWJj"
);
Ok(OcrDuringCallRequest {
body: json!({"type":"document_url","document_url":"reducto://guarded.pdf"}),
..request
})
})
}
}
#[tokio::test]
async fn guardrail_rewrites_document_before_upload() {
let (base, seen, server) =
mock_server(vec![MockResponse::json(json!({"result":{"chunks":[]}}))]).await;
let mut request = wire_request("reducto/parse-v3", &base, json!({}));
request.hooks = Arc::new(RewriteDocument);
perform_ocr(request).await.unwrap();
server.await.unwrap();
let requests = seen.lock().unwrap();
assert_eq!(requests.len(), 1);
assert!(requests[0].starts_with("POST /parse "));
assert!(requests[0].contains("reducto://guarded.pdf"));
}

View file

@ -0,0 +1,83 @@
use serde_json::{Value, json};
use super::test_support::{MockResponse, mock_server, perform_ocr, wire_request};
use crate::auth::InputSource;
fn request_body(request: &str) -> Value {
serde_json::from_str(request.split_once("\r\n\r\n").unwrap().1).unwrap()
}
#[tokio::test]
async fn facade_executes_vertex_deepseek_at_the_openai_endpoint() {
let (base, seen, server) = mock_server(vec![MockResponse::json(json!({
"choices":[{"message":{"content":"recognized"}}],
"usage":{"prompt_tokens":1}
}))])
.await;
let mut request = wire_request(
"vertex_ai/deepseek-ocr-maas",
&base,
json!({
"vertex_project":"project-1",
"vertex_location":"europe-west4",
"temperature":0.1,
"future_ocr_option":true,
"extra_body":{"provider_option":"value"}
}),
);
request.document = request
.document
.with_source("gs://bucket/document.pdf".into());
let response = perform_ocr(request).await.unwrap();
server.await.unwrap();
assert_eq!(response.pages[0]["markdown"], "recognized");
assert_eq!(response.usage_info.unwrap()["prompt_tokens"], 1);
let requests = seen.lock().unwrap();
assert!(requests[0].starts_with(
"POST /v1/projects/project-1/locations/europe-west4/endpoints/openapi/chat/completions "
));
assert!(
requests[0]
.to_ascii_lowercase()
.contains("authorization: bearer test-key")
);
let body = request_body(&requests[0]);
assert_eq!(body["model"], "deepseek-ai/deepseek-ocr-maas");
assert_eq!(body["temperature"], 0.1);
assert!(body.get("future_ocr_option").is_none());
assert!(body.get("extra_body").is_none());
assert_eq!(
body["messages"][0]["content"][0],
json!({"type":"document_url","document_url":"gs://bucket/document.pdf"})
);
}
#[test]
fn host_registration_selects_deepseek_without_affecting_mistral() {
assert!(crate::ocr::wire::is_supported_request(
"deepseek-ocr-maas",
Some("vertex_ai")
));
assert!(crate::ocr::wire::is_supported_request(
"mistral-ocr-maas",
Some("vertex_ai")
));
}
#[tokio::test]
async fn request_controlled_api_base_is_rejected_before_vertex_auth() {
let mut request = wire_request(
"vertex_ai/deepseek-ocr-maas",
"https://caller.example",
json!({"vertex_project":"project-1"}),
);
request.connection.api_base_source = InputSource::Request;
let error = perform_ocr(request).await.unwrap_err();
assert!(
error
.to_string()
.contains("request-controlled Vertex AI endpoint")
);
}

View file

@ -0,0 +1,161 @@
use serde_json::{Value, json};
use super::test_support::{MockResponse, mock_server, perform_ocr, wire_request};
use crate::auth::InputSource;
fn request_body(request: &str) -> Value {
serde_json::from_str(request.split_once("\r\n\r\n").unwrap().1).unwrap()
}
#[tokio::test]
async fn facade_executes_vertex_mistral_with_resolved_project_and_location() {
let (base, seen, server) = mock_server(vec![MockResponse::json(json!({
"pages":[{"index":0,"markdown":"hello"}],
"usage_info":{"pages_processed":1}
}))])
.await;
let request = wire_request(
"vertex_ai/mistral-ocr-maas",
&base,
json!({
"vertex_project":"project-1",
"vertex_location":"europe-west4",
"extract_footer":true
}),
);
let response = perform_ocr(request).await.unwrap();
server.await.unwrap();
assert_eq!(response.pages[0]["markdown"], "hello");
let requests = seen.lock().unwrap();
assert_eq!(requests.len(), 1);
assert!(requests[0].starts_with(
"POST /v1/projects/project-1/locations/europe-west4/publishers/mistralai/models/mistral-ocr-maas:rawPredict "
));
assert!(
requests[0]
.to_ascii_lowercase()
.contains("authorization: bearer test-key")
);
assert_eq!(
request_body(&requests[0]),
json!({
"model":"mistral-ocr-maas",
"document":{"type":"document_url","document_url":"data:application/pdf;base64,YWJj"},
"extract_footer":true
})
);
}
#[tokio::test]
async fn supplied_authorization_is_forwarded_without_a_static_token() {
let (base, seen, server) = mock_server(vec![MockResponse::json(json!({"pages":[]}))]).await;
let mut request = wire_request(
"vertex_ai/model",
&base,
json!({"vertex_project":"project-1"}),
);
request.connection.api_key = None;
request.connection.extra_headers = vec![("authorization".into(), "Bearer supplied".into())];
perform_ocr(request).await.unwrap();
server.await.unwrap();
assert!(
seen.lock().unwrap()[0]
.to_ascii_lowercase()
.contains("authorization: bearer supplied")
);
}
#[tokio::test]
async fn invalid_credentials_fail_before_provider_http() {
let request = wire_request(
"vertex_ai/model",
"http://127.0.0.1:1",
json!({"vertex_credentials": true}),
);
let error = perform_ocr(request).await.unwrap_err();
assert!(error.to_string().contains("vertex_credentials"));
}
#[tokio::test]
async fn request_controlled_api_base_is_rejected_before_vertex_auth() {
let mut request = wire_request(
"vertex_ai/mistral-ocr-maas",
"https://caller.example",
json!({"vertex_project":"project-1"}),
);
request.connection.api_base_source = InputSource::Request;
let error = perform_ocr(request).await.unwrap_err();
assert!(
error
.to_string()
.contains("request-controlled Vertex AI endpoint")
);
}
#[tokio::test]
async fn adapters_build_complete_requests_and_share_mistral_normalization() {
use std::time::Duration;
use crate::ocr::adapters::{MistralAdapter, OcrAdapter, VertexMistralAdapter};
use crate::ocr::test_support::ocr_client;
let client = ocr_client();
let options = json!({
"pages": [0, 2],
"include_image_base64": true,
"vertex_project": "project-1",
"vertex_location": "us-central1",
"unknown": "ignored"
});
let direct = wire_request(
"mistral/mistral-ocr-maas",
"https://mistral.test",
options.clone(),
);
let vertex = wire_request("vertex_ai/mistral-ocr-maas", "https://vertex.test", options);
let direct_http = MistralAdapter
.prepare_request(&direct, &client)
.await
.unwrap();
let vertex_http = VertexMistralAdapter
.prepare_request(&vertex, &client)
.await
.unwrap();
assert_eq!(direct_http.url().as_str(), "https://mistral.test/v1/ocr");
assert_eq!(
vertex_http.url().as_str(),
"https://vertex.test/v1/projects/project-1/locations/us-central1/publishers/mistralai/models/mistral-ocr-maas:rawPredict"
);
for http in [&direct_http, &vertex_http] {
assert_eq!(http.method(), reqwest::Method::POST);
assert_eq!(http.headers()["authorization"], "Bearer test-key");
assert_eq!(http.headers()["content-type"], "application/json");
assert_eq!(http.timeout(), Some(&Duration::from_secs(2)));
let body: Value = serde_json::from_slice(http.body().unwrap().as_bytes().unwrap()).unwrap();
assert_eq!(
body,
json!({
"model": "mistral-ocr-maas",
"document": {"type": "document_url", "document_url": "data:application/pdf;base64,YWJj"},
"pages": [0, 2],
"include_image_base64": true
})
);
}
let payload = json!({"pages": [{"index": 0, "markdown": "hello"}], "extra": "preserved"});
let direct_response = MistralAdapter
.transform_ocr_response(&direct, serde_json::from_value(payload.clone()).unwrap())
.unwrap()
.into_json();
let vertex_response = VertexMistralAdapter
.transform_ocr_response(&vertex, serde_json::from_value(payload).unwrap())
.unwrap()
.into_json();
assert_eq!(direct_response, vertex_response);
assert_eq!(direct_response["model"], "mistral-ocr-maas");
assert_eq!(direct_response["object"], "ocr");
assert_eq!(direct_response["extra"], "preserved");
}

View file

@ -42,6 +42,9 @@ pub(crate) fn chat_completions_error_to_pyerr(err: Error) -> PyErr {
| Error::InvalidType { .. }
| Error::MissingField(_)
| Error::MissingApiKey { .. }
| Error::MissingAzureAiCredentials
| Error::MissingAzureDocumentIntelligenceCredentials
| Error::MissingReductoApiKey
| Error::Routing(_)
// Nothing reached the provider, so serving it on Python cannot double
// bill and is the only way the caller gets an answer at all.

View file

@ -225,7 +225,7 @@ mod tests {
(
"ocr",
"aocr",
"(model, document, api_key=None, api_base=None, custom_llm_provider=None, extra_headers=None, optional_params=None, timeout_seconds=None)",
"(model, document, api_key=None, api_base=None, custom_llm_provider=None, extra_headers=None, optional_params=None, input_sources=None, timeout_seconds=None)",
),
(
"transcription",

View file

@ -2,6 +2,7 @@ use litellm_core::Error;
use std::future::Future;
use litellm_ai_gateway::io::ocr::{OcrRequest, ocr as run_ocr};
use litellm_core::ocr::wire::{OcrWireRequest, decode_request, is_supported_request};
use pyo3::prelude::*;
use serde_json::Value;
@ -21,6 +22,12 @@ fn prepare_ocr(
timeout_seconds: inputs.timeout_seconds,
})?;
let optional_params = object_or_empty("optional_params", inputs.optional_params)?;
let input_sources = inputs
.input_sources
.map(serde_json::from_value)
.transpose()
.map_err(|error| pyo3::exceptions::PyValueError::new_err(error.to_string()))?
.unwrap_or_default();
Ok(async move {
let RouteOptions {
@ -31,6 +38,22 @@ fn prepare_ocr(
extra_headers,
timeout,
} = options;
if is_supported_request(&model, custom_llm_provider.as_deref()) {
let request = decode_request(OcrWireRequest {
model,
document,
api_key,
api_base,
custom_llm_provider,
extra_headers,
optional_params,
input_sources,
timeout_seconds: timeout.map(|value| value.as_secs_f64()),
})?;
return litellm_core::ocr::ocr(request)
.await
.map(|response| response.into_json());
}
run_ocr(OcrRequest {
model: &model,
document,
@ -66,8 +89,29 @@ bridge_route! {
extra_headers: Option<serde_json::Value>,
#[pyo3(from_py_with = litellm_python_interop::from_py)]
optional_params: Option<serde_json::Value>,
#[pyo3(from_py_with = litellm_python_interop::from_py)]
input_sources: Option<serde_json::Value>,
timeout_seconds: Option<f64>,
},
prepare = prepare_ocr,
errors = ocr_error_to_pyerr,
}
#[cfg(test)]
mod tests {
use litellm_core::ocr::wire::is_supported_request;
#[test]
fn native_activation_includes_migrated_providers() {
assert!(is_supported_request("model", Some("mistral")));
assert!(is_supported_request("pixtral-12b", Some("azure_ai")));
assert!(is_supported_request(
"documentintelligence/prebuilt-read",
Some("azure_ai")
));
assert!(is_supported_request("parse-v3", Some("reducto")));
assert!(is_supported_request("parse-legacy", Some("reducto")));
assert!(is_supported_request("mistral-ocr", Some("vertex_ai")));
assert!(is_supported_request("deepseek-ocr", Some("vertex_ai")));
}
}

View file

@ -1370,6 +1370,7 @@ from .exceptions import (
InvalidRequestError,
BadRequestError,
ImageFetchError,
VectorStoreSearchError,
NotFoundError,
PermissionDeniedError,
RateLimitError,
@ -1471,9 +1472,11 @@ from .vector_stores.vector_store_registry import (
VectorStoreRegistry,
VectorStoreIndexRegistry,
)
from .types.vector_stores import VectorStoreSearchFailureMode
vector_store_registry: Optional[VectorStoreRegistry] = None
vector_store_index_registry: Optional[VectorStoreIndexRegistry] = None
vector_store_search_failure_mode: VectorStoreSearchFailureMode = "annotate"
### RAG ###
from . import rag

Some files were not shown because too many files have changed in this diff Show more