diff --git a/litellm-rust/crates/ai-gateway/src/audio_transcription/hooks.rs b/litellm-rust/crates/ai-gateway/src/audio_transcription/hooks.rs index 17f5591d1fc..6f48f38c9f6 100644 --- a/litellm-rust/crates/ai-gateway/src/audio_transcription/hooks.rs +++ b/litellm-rust/crates/ai-gateway/src/audio_transcription/hooks.rs @@ -272,7 +272,6 @@ fn core_error_kind(error: &Error) -> &'static str { Error::Auth(_) | Error::MissingApiKey { .. } | Error::MissingAzureAiCredentials - | Error::MissingAzureAiCredentialsOrAdToken | Error::MissingAzureDocumentIntelligenceCredentials | Error::MissingReductoApiKey => "AuthError", Error::InvalidProvider(_) => "InvalidProvider", diff --git a/litellm-rust/crates/ai-gateway/src/ocr/common_utils.rs b/litellm-rust/crates/ai-gateway/src/ocr/common_utils.rs deleted file mode 100644 index 8305cc80a1d..00000000000 --- a/litellm-rust/crates/ai-gateway/src/ocr/common_utils.rs +++ /dev/null @@ -1,525 +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::vertex_ai::ocr::transformation as vertex_ai; -use litellm_core::providers::vertex_ai::ocr::transformation::VERTEX_AI_DEEPSEEK_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), - "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" => None, - _ => 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>, -) -> Result, 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, 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::().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::() { - 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 { - 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(¤t_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, ¤t_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, 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 { - 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::().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, -) -> Result { - 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() - ) - ); - } -} diff --git a/litellm-rust/crates/ai-gateway/src/ocr/handler.rs b/litellm-rust/crates/ai-gateway/src/ocr/handler.rs deleted file mode 100644 index 6c6e12724cd..00000000000 --- a/litellm-rust/crates/ai-gateway/src/ocr/handler.rs +++ /dev/null @@ -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 { - 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()) -} diff --git a/litellm-rust/crates/ai-gateway/src/ocr/hooks.rs b/litellm-rust/crates/ai-gateway/src/ocr/hooks.rs deleted file mode 100644 index 3d8246af3f5..00000000000 --- a/litellm-rust/crates/ai-gateway/src/ocr/hooks.rs +++ /dev/null @@ -1,330 +0,0 @@ -use litellm_core::call_lifecycle::{CallLifecycleContext, CallLifecycleHooks, CallLifecycleTiming}; -use litellm_core::error::Error; -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}; -use super::types::{PreparedOcrRequest, ProviderOcrRequest}; -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> + Send + 'a>>; -type OcrLogFuture<'a> = Pin + 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 { - 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 { - 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 document = 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 = 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 { - 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, - } - } -} - -impl CallLifecycleHooks 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), 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 { - 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 { .. } - | Error::MissingAzureAiCredentials - | Error::MissingAzureAiCredentialsOrAdToken - | Error::MissingAzureDocumentIntelligenceCredentials - | Error::MissingReductoApiKey => "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", - } -} diff --git a/litellm-rust/crates/ai-gateway/src/ocr/mod.rs b/litellm-rust/crates/ai-gateway/src/ocr/mod.rs index 116c0f4a5e9..fb63a02f7ad 100644 --- a/litellm-rust/crates/ai-gateway/src/ocr/mod.rs +++ b/litellm-rust/crates/ai-gateway/src/ocr/mod.rs @@ -1,31 +1,96 @@ 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 { - 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 { + 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 super::{OcrRequest, validate_host_hooks}; + use crate::integrations::custom_guardrail::{CustomGuardrail, GuardrailEventHook}; + use crate::integrations::custom_logger::CustomLogger; + + struct TestGuardrail; + + impl CustomGuardrail for TestGuardrail { + fn guardrail_name(&self) -> &str { + "test" + } + + fn supported_event_hooks(&self) -> &[GuardrailEventHook] { + &[] + } + } + + struct TestLogger; + + impl CustomLogger for TestLogger {} + + fn request() -> OcrRequest<'static> { + OcrRequest { + model: "model", + document: json!({"type":"image_url","image_url":"data:image/png;base64,YQ=="}), + api_key: None, + api_base: None, + custom_llm_provider: Some("mistral"), + extra_headers: None, + optional_params: Map::new(), + timeout: None, + callbacks: Vec::new(), + guardrails: Vec::new(), + request_metadata: Default::default(), + litellm_call_id: None, + } + } #[test] fn core_activation_includes_migrated_providers() { @@ -37,6 +102,26 @@ mod tests { )); 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"))); + assert!(is_supported_request("deepseek-ocr", Some("vertex_ai"))); + } + + #[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")); + } + + #[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")); } } diff --git a/litellm-rust/crates/ai-gateway/src/ocr/prepare.rs b/litellm-rust/crates/ai-gateway/src/ocr/prepare.rs deleted file mode 100644 index fa9ca1a193e..00000000000 --- a/litellm-rust/crates/ai-gateway/src/ocr/prepare.rs +++ /dev/null @@ -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, - 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`")) - ); - } -} diff --git a/litellm-rust/crates/ai-gateway/src/ocr/types.rs b/litellm-rust/crates/ai-gateway/src/ocr/types.rs index 75a8e61ddbf..e96d2df1adb 100644 --- a/litellm-rust/crates/ai-gateway/src/ocr/types.rs +++ b/litellm-rust/crates/ai-gateway/src/ocr/types.rs @@ -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, - pub(crate) api_base: Option, - pub(crate) extra_headers: Option>, - pub(crate) optional_params: Map, - pub(crate) timeout: Option, -} - -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, - pub(crate) upstream_headers: Vec<(String, String)>, - pub(crate) timeout: Option, -} diff --git a/litellm-rust/crates/ai-gateway/src/routes/messages/mod.rs b/litellm-rust/crates/ai-gateway/src/routes/messages/mod.rs index 9707e9f2611..39465e28e84 100644 --- a/litellm-rust/crates/ai-gateway/src/routes/messages/mod.rs +++ b/litellm-rust/crates/ai-gateway/src/routes/messages/mod.rs @@ -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,12 +118,7 @@ impl IntoResponse for MessagesRouteError { | Error::Connect(_) | Error::InvalidResponse(_) | Error::InvalidType { .. } - | Error::MissingField(_) - | Error::MissingApiKey { .. } - | Error::MissingAzureAiCredentials - | Error::MissingAzureAiCredentialsOrAdToken - | Error::MissingAzureDocumentIntelligenceCredentials - | Error::MissingReductoApiKey => ( + | Error::MissingField(_) => ( StatusCode::BAD_GATEWAY, "messages provider request failed".to_string(), ), diff --git a/litellm-rust/crates/ai-gateway/tests/ocr_lifecycle.rs b/litellm-rust/crates/ai-gateway/tests/ocr_lifecycle.rs deleted file mode 100644 index 2fbd25d986f..00000000000 --- a/litellm-rust/crates/ai-gateway/tests/ocr_lifecycle.rs +++ /dev/null @@ -1,603 +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::().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, - response_object: Option, - error_kind: Option, -} - -#[derive(Default)] -struct RecordingOcrLogger { - events: Mutex>, -} - -impl RecordingOcrLogger { - fn events(&self) -> Vec { - 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, - events: Mutex>, - block_pre_call: bool, - block_during_call: bool, -} - -impl RecordingOcrGuardrail { - fn new(hooks: Vec) -> 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 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)) - }) - } -} - -#[tokio::test] -async fn azure_mistral_uses_prepared_authorization_through_gateway() { - let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); - let api_base = format!("http://{}", listener.local_addr().unwrap()); - let server = tokio::spawn(async move { - let (mut socket, _) = listener.accept().await.unwrap(); - let request = read_http_request(&mut socket).await; - let body = br#"{"pages":[]}"#; - socket - .write_all( - format!( - "HTTP/1.1 200 OK\r\ncontent-type: application/json\r\ncontent-length: {}\r\nconnection: close\r\n\r\n", - body.len() - ) - .as_bytes(), - ) - .await - .unwrap(); - socket.write_all(body).await.unwrap(); - request - }); - let request = OcrRequest { - model: "mistral-ocr-2505", - document: json!({ - "type":"document_url", - "document_url":"data:application/pdf;base64,YWJj" - }), - api_key: None, - api_base: Some(&api_base), - custom_llm_provider: Some("azure_ai"), - extra_headers: Some(Map::from_iter([( - "Authorization".into(), - json!("Bearer python-prepared-token"), - )])), - optional_params: Map::new(), - timeout: None, - callbacks: Vec::new(), - guardrails: Vec::new(), - request_metadata: RequestMetadata::default(), - litellm_call_id: None, - }; - - ocr(request).await.unwrap(); - let sent = server.await.unwrap(); - assert!(sent.starts_with("POST /providers/mistral/azure/ocr ")); - assert!( - sent.to_ascii_lowercase() - .contains("authorization: bearer python-prepared-token\r\n") - ); -} - -#[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!["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!["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}" - ); -} diff --git a/litellm-rust/crates/core/src/error.rs b/litellm-rust/crates/core/src/error.rs index 2a4cbad96c0..fa4a9d36e03 100644 --- a/litellm-rust/crates/core/src/error.rs +++ b/litellm-rust/crates/core/src/error.rs @@ -25,8 +25,6 @@ pub enum Error { "invalid authentication configuration: Missing Azure AI credentials - set AZURE_AI_API_KEY or configure Entra ID" )] MissingAzureAiCredentials, - #[error("Missing Azure AI credentials - set AZURE_AI_API_KEY or provide azure_ad_token")] - MissingAzureAiCredentialsOrAdToken, #[error( "invalid authentication configuration: Missing Azure Document Intelligence credentials - set AZURE_DOCUMENT_INTELLIGENCE_API_KEY or configure Entra ID" )] diff --git a/litellm-rust/crates/core/src/ocr/adapters/mod.rs b/litellm-rust/crates/core/src/ocr/adapters/mod.rs index a96a7fcdf38..9171d11836c 100644 --- a/litellm-rust/crates/core/src/ocr/adapters/mod.rs +++ b/litellm-rust/crates/core/src/ocr/adapters/mod.rs @@ -16,7 +16,7 @@ mod vertex; pub(crate) use azure::{AzureDocumentIntelligenceAdapter, AzureMistralAdapter}; pub(crate) use mistral::MistralAdapter; pub(crate) use reducto::{ReductoLegacyAdapter, ReductoV3Adapter}; -pub(crate) use vertex::VertexMistralAdapter; +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 { @@ -73,6 +73,7 @@ macro_rules! for_each_ocr_adapter { 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; } }; } diff --git a/litellm-rust/crates/core/src/ocr/adapters/vertex/deepseek.rs b/litellm-rust/crates/core/src/ocr/adapters/vertex/deepseek.rs new file mode 100644 index 00000000000..ef188f8b9ac --- /dev/null +++ b/litellm-rust/crates/core/src/ocr/adapters/vertex/deepseek.rs @@ -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 { + validate_destination(&request.connection)?; + let ParsedProviderParams { + known: params, + extra_params: _extra_params, + } = _prepare_ocr_request::(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, ¶ms)?; + transform_request_body(client, request, &url, &authentication.headers, body, |_| { + Ok(()) + }) + .await + } + + fn transform_ocr_response( + &self, + request: &LiteLLMOcrRequest, + response: Self::ProviderResponse, + ) -> Result { + 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 { + 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" + ); + } +} diff --git a/litellm-rust/crates/core/src/ocr/adapters/vertex/mod.rs b/litellm-rust/crates/core/src/ocr/adapters/vertex/mod.rs index ce6f884b41d..270c41e647d 100644 --- a/litellm-rust/crates/core/src/ocr/adapters/vertex/mod.rs +++ b/litellm-rust/crates/core/src/ocr/adapters/vertex/mod.rs @@ -1,3 +1,4 @@ +mod deepseek; mod mistral; use crate::Error; @@ -6,6 +7,7 @@ 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> { diff --git a/litellm-rust/crates/core/src/ocr/codecs/deepseek/mod.rs b/litellm-rust/crates/core/src/ocr/codecs/deepseek/mod.rs new file mode 100644 index 00000000000..682b3addde7 --- /dev/null +++ b/litellm-rust/crates/core/src/ocr/codecs/deepseek/mod.rs @@ -0,0 +1,5 @@ +mod transformation; +mod types; + +pub(crate) use transformation::{transform_ocr_request, transform_ocr_response}; +pub(crate) use types::{DeepSeekOcrParams, DeepSeekOcrResponse}; diff --git a/litellm-rust/crates/core/src/ocr/codecs/deepseek/transformation.rs b/litellm-rust/crates/core/src/ocr/codecs/deepseek/transformation.rs new file mode 100644 index 00000000000..98cfc0db78d --- /dev/null +++ b/litellm-rust/crates/core/src/ocr/codecs/deepseek/transformation.rs @@ -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 { + 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 { + 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 { + 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, OcrResponseError> { + if !text.trim_start().starts_with('{') { + return Ok(None); + } + let value = match serde_json::from_str::(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()), + }) +} diff --git a/litellm-rust/crates/core/src/ocr/codecs/deepseek/types.rs b/litellm-rust/crates/core/src/ocr/codecs/deepseek/types.rs new file mode 100644 index 00000000000..0ce2d9913f7 --- /dev/null +++ b/litellm-rust/crates/core/src/ocr/codecs/deepseek/types.rs @@ -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, + #[serde(skip_serializing_if = "Option::is_none")] + pub temperature: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub max_tokens: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub top_p: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub n: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub stop: Option, +} + +#[derive(Clone, Debug, Serialize, Deserialize)] +#[serde(untagged)] +pub(crate) enum StopSequences { + One(String), + Many(Vec), +} + +#[derive(Clone, Debug, Serialize, Deserialize)] +pub(crate) struct DeepSeekOcrRequest { + pub model: String, + pub messages: Vec, + #[serde(flatten)] + pub params: DeepSeekOcrParams, +} + +#[derive(Clone, Debug, Serialize, Deserialize)] +pub(crate) struct DeepSeekOcrMessage { + pub role: UserRole, + pub content: Vec, +} + +#[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, + pub usage: Option, +} + +#[derive(Clone, Debug, Deserialize)] +pub(crate) struct DeepSeekChoice { + pub message: DeepSeekResponseMessage, +} + +#[derive(Clone, Debug, Deserialize)] +pub(crate) struct DeepSeekResponseMessage { + pub content: Option, +} + +#[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>, + #[serde(skip_serializing_if = "Option::is_none")] + pub model: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub usage_info: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub document_annotation: Option, + #[serde(flatten)] + pub extra_fields: Map, +} + +#[derive(Clone, Debug, Serialize, Deserialize)] +pub(crate) struct DeepSeekPage { + #[serde(default)] + pub index: i64, + #[serde(default)] + pub markdown: String, + pub images: Option, + pub dimensions: Option, + #[serde(flatten)] + pub extra_fields: Map, +} diff --git a/litellm-rust/crates/core/src/ocr/codecs/mistral/transformation.rs b/litellm-rust/crates/core/src/ocr/codecs/mistral/transformation.rs index cd0a1dc6b17..5bd7e555a1e 100644 --- a/litellm-rust/crates/core/src/ocr/codecs/mistral/transformation.rs +++ b/litellm-rust/crates/core/src/ocr/codecs/mistral/transformation.rs @@ -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::(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, ¶ms).unwrap()) + serde_json::to_value(transform_ocr_request("model", document(), ¶ms).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(), ¶ms).unwrap(), ) .unwrap(); - let result = - serde_json::to_value(transform_ocr_request("model", document, ¶ms).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(), ¶ms).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::(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); } } diff --git a/litellm-rust/crates/core/src/ocr/codecs/mod.rs b/litellm-rust/crates/core/src/ocr/codecs/mod.rs index 79dcd150f5f..7c752749901 100644 --- a/litellm-rust/crates/core/src/ocr/codecs/mod.rs +++ b/litellm-rust/crates/core/src/ocr/codecs/mod.rs @@ -1,3 +1,4 @@ +pub(crate) mod deepseek; pub(crate) mod document_intelligence; pub(crate) mod mistral; pub(crate) mod reducto; diff --git a/litellm-rust/crates/core/src/ocr/error.rs b/litellm-rust/crates/core/src/ocr/error.rs index f42ac2ceb18..522d059ec48 100644 --- a/litellm-rust/crates/core/src/ocr/error.rs +++ b/litellm-rust/crates/core/src/ocr/error.rs @@ -36,6 +36,8 @@ pub enum OcrRequestError { 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")] diff --git a/litellm-rust/crates/core/src/ocr/mod.rs b/litellm-rust/crates/core/src/ocr/mod.rs index 69b2958483a..1e975c3f521 100644 --- a/litellm-rust/crates/core/src/ocr/mod.rs +++ b/litellm-rust/crates/core/src/ocr/mod.rs @@ -7,7 +7,6 @@ mod handler; pub mod hooks; mod prepare; mod registry; -pub mod transformation; pub mod types; pub mod wire; @@ -21,6 +20,9 @@ mod azure_ai_tests; #[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)] @@ -30,5 +32,8 @@ pub(crate) mod test_support; #[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; diff --git a/litellm-rust/crates/core/src/ocr/prepare.rs b/litellm-rust/crates/core/src/ocr/prepare.rs index 363f963a66c..bf6f924088c 100644 --- a/litellm-rust/crates/core/src/ocr/prepare.rs +++ b/litellm-rust/crates/core/src/ocr/prepare.rs @@ -163,7 +163,6 @@ impl OcrWireBody { pub(crate) fn credential_env(name: &str) -> Option { std::env::var(name).ok() } - #[cfg(test)] mod tests { use serde_json::json; diff --git a/litellm-rust/crates/core/src/ocr/registry.rs b/litellm-rust/crates/core/src/ocr/registry.rs index 1097d102a20..1b20a91143b 100644 --- a/litellm-rust/crates/core/src/ocr/registry.rs +++ b/litellm-rust/crates/core/src/ocr/registry.rs @@ -75,7 +75,7 @@ pub(crate) fn resolve_wire_adapter( ))); } OcrProvider::VertexAi if provider.model.to_ascii_lowercase().contains("deepseek") => { - return Err(Error::Unsupported("Vertex DeepSeek OCR")); + OcrAdapterKind::VertexDeepSeek } OcrProvider::VertexAi => OcrAdapterKind::VertexMistral, }; diff --git a/litellm-rust/crates/core/src/ocr/transformation.rs b/litellm-rust/crates/core/src/ocr/transformation.rs deleted file mode 100644 index ac4f10bf15b..00000000000 --- a/litellm-rust/crates/core/src/ocr/transformation.rs +++ /dev/null @@ -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) -> Map { - let mut mapped_params = Map::new(); - for (param, value) in non_default_params { - if self.supported_ocr_params().contains(¶m.as_str()) { - mapped_params.insert(param.clone(), value.clone()); - } - } - mapped_params - } - - fn transform_ocr_request( - &self, - model: &str, - document: Value, - optional_params: Map, - ) -> Result; - - fn transform_ocr_response( - &self, - model: &str, - response_json: Value, - ) -> Result; - - fn transform_ocr_response_with_params( - &self, - model: &str, - response_json: Value, - _optional_params: &Map, - ) -> Result { - self.transform_ocr_response(model, response_json) - } - - fn complete_url( - &self, - api_base: Option<&str>, - model: &str, - optional_params: &Map, - env_lookup: &dyn Fn(&str) -> Option, - ) -> Result; - - fn resolve_api_key( - &self, - api_key: Option<&str>, - env_lookup: &dyn Fn(&str) -> Option, - ) -> Result; - - #[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, - ) -> Result, 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 - } -} diff --git a/litellm-rust/crates/core/src/ocr/types.rs b/litellm-rust/crates/core/src/ocr/types.rs index 0e92b0b6868..06519f86c91 100644 --- a/litellm-rust/crates/core/src/ocr/types.rs +++ b/litellm-rust/crates/core/src/ocr/types.rs @@ -11,12 +11,6 @@ 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, -} - #[derive(Clone, Debug, PartialEq, Serialize, Deserialize)] #[serde(tag = "type")] pub enum OcrDocument { diff --git a/litellm-rust/crates/core/src/providers/azure_ai/mod.rs b/litellm-rust/crates/core/src/providers/azure_ai/mod.rs index f2d5b679aee..4f41d1d6abb 100644 --- a/litellm-rust/crates/core/src/providers/azure_ai/mod.rs +++ b/litellm-rust/crates/core/src/providers/azure_ai/mod.rs @@ -1,3 +1,2 @@ pub(crate) mod auth; pub mod messages; -pub mod ocr; diff --git a/litellm-rust/crates/core/src/providers/azure_ai/ocr/mod.rs b/litellm-rust/crates/core/src/providers/azure_ai/ocr/mod.rs deleted file mode 100644 index f239b6921fa..00000000000 --- a/litellm-rust/crates/core/src/providers/azure_ai/ocr/mod.rs +++ /dev/null @@ -1 +0,0 @@ -pub mod transformation; diff --git a/litellm-rust/crates/core/src/providers/azure_ai/ocr/transformation.rs b/litellm-rust/crates/core/src/providers/azure_ai/ocr/transformation.rs deleted file mode 100644 index 2dbec8e2187..00000000000 --- a/litellm-rust/crates/core/src/providers/azure_ai/ocr/transformation.rs +++ /dev/null @@ -1,1376 +0,0 @@ -use std::collections::BTreeSet; - -use crate::error::{Error, json_type_name}; -use crate::ocr::transformation::{OcrAuthStrategy, OcrProviderConfig, OcrResponseHandling}; -use crate::ocr::types::{LiteLLMOcrResponse, OcrRequestData}; -use serde_json::{Map, Value, json}; - -use crate::providers::mistral::ocr::transformation::MISTRAL_OCR_CONFIG; - -const AZURE_AI_API_KEY_ENV: &str = "AZURE_AI_API_KEY"; -const AZURE_AI_API_BASE_ENV: &str = "AZURE_AI_API_BASE"; -const AZURE_DOCUMENT_INTELLIGENCE_API_KEY_ENV: &str = "AZURE_DOCUMENT_INTELLIGENCE_API_KEY"; -const AZURE_DOCUMENT_INTELLIGENCE_ENDPOINT_ENV: &str = "AZURE_DOCUMENT_INTELLIGENCE_ENDPOINT"; -const AZURE_DOCUMENT_INTELLIGENCE_API_VERSION: &str = "2024-11-30"; -const AZURE_DOCUMENT_INTELLIGENCE_DEFAULT_DPI: i64 = 96; - -const AZURE_DOCUMENT_INTELLIGENCE_SUPPORTED_OCR_PARAMS: &[&str] = - &["pages", "features", "req_format"]; - -pub struct AzureAiOcrConfig; -pub struct AzureDocumentIntelligenceOcrConfig; - -pub const AZURE_AI_OCR_CONFIG: AzureAiOcrConfig = AzureAiOcrConfig; -pub const AZURE_DOCUMENT_INTELLIGENCE_OCR_CONFIG: AzureDocumentIntelligenceOcrConfig = - AzureDocumentIntelligenceOcrConfig; - -fn non_empty(value: Option<&str>) -> Option<&str> { - value.map(str::trim).filter(|value| !value.is_empty()) -} - -fn resolve_value( - explicit: Option<&str>, - env_name: &str, - env_lookup: &dyn Fn(&str) -> Option, - missing_message: &str, -) -> Result { - non_empty(explicit) - .map(str::to_string) - .or_else(|| env_lookup(env_name).filter(|value| !value.trim().is_empty())) - .ok_or_else(|| Error::Auth(missing_message.to_string())) -} - -pub fn resolve_azure_ai_api_key( - api_key: Option<&str>, - env_lookup: &dyn Fn(&str) -> Option, -) -> Result { - resolve_value( - api_key, - AZURE_AI_API_KEY_ENV, - env_lookup, - "Missing Azure AI API Key - A call is being made to Azure AI but no key is set either in the environment variables or via params", - ) -} - -pub fn resolve_azure_ai_api_base( - api_base: Option<&str>, - env_lookup: &dyn Fn(&str) -> Option, -) -> Result { - resolve_value( - api_base, - AZURE_AI_API_BASE_ENV, - env_lookup, - "Missing Azure AI API Base - Set AZURE_AI_API_BASE environment variable or pass api_base parameter", - ) -} - -pub fn complete_azure_ai_url( - api_base: Option<&str>, - env_lookup: &dyn Fn(&str) -> Option, -) -> Result { - let base = resolve_azure_ai_api_base(api_base, env_lookup)?; - Ok(format!( - "{}/providers/mistral/azure/ocr", - base.trim_end_matches('/') - )) -} - -pub fn resolve_document_intelligence_api_key( - api_key: Option<&str>, - env_lookup: &dyn Fn(&str) -> Option, -) -> Result { - resolve_value( - api_key, - AZURE_DOCUMENT_INTELLIGENCE_API_KEY_ENV, - env_lookup, - "Missing Azure Document Intelligence API Key - Set AZURE_DOCUMENT_INTELLIGENCE_API_KEY environment variable or pass api_key parameter", - ) -} - -pub fn resolve_document_intelligence_endpoint( - api_base: Option<&str>, - env_lookup: &dyn Fn(&str) -> Option, -) -> Result { - resolve_value( - api_base, - AZURE_DOCUMENT_INTELLIGENCE_ENDPOINT_ENV, - env_lookup, - "Missing Azure Document Intelligence Endpoint - Set AZURE_DOCUMENT_INTELLIGENCE_ENDPOINT environment variable or pass api_base parameter", - ) -} - -fn prepend_auth_header( - headers: Vec<(String, String)>, - name: &str, - value: String, -) -> Vec<(String, String)> { - std::iter::once((name.to_string(), value)) - .chain(headers) - .collect() -} - -pub fn validate_azure_ai_environment( - headers: Vec<(String, String)>, - api_key: Option<&str>, - azure_ad_token: Option<&str>, - env_lookup: &dyn Fn(&str) -> Option, -) -> Result, Error> { - if crate::http_utils::has_header(&headers, "Authorization") - || crate::http_utils::has_header(&headers, "Api-Key") - { - return Ok(headers); - } - if let Ok(api_key) = resolve_azure_ai_api_key(api_key, env_lookup) { - return Ok(prepend_auth_header(headers, "Api-Key", api_key)); - } - non_empty(azure_ad_token) - .map(|token| prepend_auth_header(headers, "Authorization", format!("Bearer {token}"))) - .ok_or(Error::MissingAzureAiCredentialsOrAdToken) -} - -pub fn validate_document_intelligence_environment( - headers: Vec<(String, String)>, - api_key: Option<&str>, - azure_ad_token: Option<&str>, - env_lookup: &dyn Fn(&str) -> Option, -) -> Result, Error> { - if crate::http_utils::has_header(&headers, "Authorization") - || crate::http_utils::has_header(&headers, "Ocp-Apim-Subscription-Key") - { - return Ok(headers); - } - if let Ok(api_key) = resolve_document_intelligence_api_key(api_key, env_lookup) { - return Ok(prepend_auth_header( - headers, - "Ocp-Apim-Subscription-Key", - api_key, - )); - } - non_empty(azure_ad_token) - .map(|token| prepend_auth_header(headers, "Authorization", format!("Bearer {token}"))) - .ok_or_else(|| { - Error::Auth( - "Missing Azure Document Intelligence credentials - set AZURE_DOCUMENT_INTELLIGENCE_API_KEY or provide azure_ad_token" - .to_string(), - ) - }) -} - -fn encode_model_id(model: &str) -> Result { - let model_id = model.rsplit('/').next().unwrap_or(model); - if matches!(model_id, "." | "..") { - return Err(Error::InvalidRequest( - "model_id cannot be a dot path segment".to_string(), - )); - } - Ok(model_id - .bytes() - .flat_map(|byte| match byte { - b'A'..=b'Z' | b'a'..=b'z' | b'0'..=b'9' | b'-' | b'_' | b'.' | b'~' => { - vec![byte as char] - } - _ => format!("%{byte:02X}").chars().collect(), - }) - .collect()) -} - -fn pages_token_is_valid(token: &str) -> bool { - let mut parts = token.split('-'); - let Some(start) = parts.next() else { - return false; - }; - if start.is_empty() || !start.chars().all(|ch| ch.is_ascii_digit()) { - return false; - } - match parts.next() { - None => true, - Some(end) => { - !end.is_empty() && end.chars().all(|ch| ch.is_ascii_digit()) && parts.next().is_none() - } - } -} - -fn normalize_pages_param(pages: &Value) -> Result, Error> { - match pages { - Value::String(value) => { - let normalized = value - .split(',') - .map(str::trim) - .collect::>() - .join(","); - if normalized.split(',').all(pages_token_is_valid) { - Ok(Some(normalized)) - } else { - Err(Error::InvalidRequest(format!( - "Invalid `pages` string for Azure Document Intelligence: {value:?}. Expected format like '1-3,5,7-9'." - ))) - } - } - Value::Array(values) => { - if values.is_empty() { - return Ok(None); - } - if values.iter().any(Value::is_boolean) { - return Err(Error::InvalidRequest( - "`pages` must be integers, not booleans".to_string(), - )); - } - if values.iter().all(Value::is_i64) { - let mut pages = BTreeSet::new(); - for value in values { - let page = value.as_i64().expect("checked is_i64"); - if page < 0 { - return Err(Error::InvalidRequest( - "`pages` integers must be >= 0 (Mistral 0-based indices)".to_string(), - )); - } - pages.insert(page + 1); - } - return Ok(Some( - pages - .into_iter() - .map(|page| page.to_string()) - .collect::>() - .join(","), - )); - } - if values.iter().all(Value::is_string) { - let normalized = values - .iter() - .filter_map(Value::as_str) - .map(str::trim) - .collect::>() - .join(","); - if normalized.split(',').all(pages_token_is_valid) { - return Ok(Some(normalized)); - } - return Err(Error::InvalidRequest(format!( - "Invalid `pages` list for Azure Document Intelligence: {values:?}. Expected tokens like '1' or '3-5'." - ))); - } - Err(Error::InvalidRequest( - "`pages` must be a list[int] (0-based, Mistral-style) or a string like '1-3,5,7-9'." - .to_string(), - )) - } - _ => Err(Error::InvalidRequest( - "`pages` must be a list[int] (0-based, Mistral-style) or a string like '1-3,5,7-9'." - .to_string(), - )), - } -} - -fn feature_token_is_valid(token: &str) -> bool { - let Some((first, rest)) = token.as_bytes().split_first() else { - return false; - }; - first.is_ascii_alphabetic() && rest.iter().all(u8::is_ascii_alphanumeric) -} - -fn invalid_features_error(features: &Value) -> Error { - Error::InvalidRequest(format!( - "Invalid `features` for Azure Document Intelligence: {features:?}. Expected a list of feature names or a comma-separated string like 'keyValuePairs' or 'keyValuePairs,languages'." - )) -} - -fn normalize_features_param(features: &Value) -> Result, Error> { - let normalized = match features { - Value::String(value) => value - .split(',') - .map(str::trim) - .collect::>() - .join(","), - Value::Array(values) if values.is_empty() => return Ok(None), - Value::Array(values) => values - .iter() - .map(Value::as_str) - .collect::>>() - .ok_or_else(|| invalid_features_error(features))? - .into_iter() - .map(str::trim) - .collect::>() - .join(","), - _ => return Err(invalid_features_error(features)), - }; - - if normalized.split(',').all(feature_token_is_valid) { - Ok(Some(normalized)) - } else { - Err(invalid_features_error(features)) - } -} - -fn normalize_req_format(req_format: &Value) -> Result { - match req_format.as_str() { - Some(value @ ("native" | "litellm")) => Ok(value.to_string()), - _ => Err(Error::InvalidRequest(format!( - "Invalid `req_format` for Azure Document Intelligence: {req_format:?}. Expected 'native' or 'litellm'." - ))), - } -} - -pub fn map_document_intelligence_ocr_params( - non_default_params: &Map, -) -> Result, Error> { - let mut mapped = Map::new(); - if let Some(pages) = non_default_params.get("pages") - && let Some(normalized) = normalize_pages_param(pages)? - { - mapped.insert("pages".to_string(), Value::String(normalized)); - } - if let Some(features) = non_default_params.get("features") - && let Some(normalized) = normalize_features_param(features)? - { - mapped.insert("features".to_string(), Value::String(normalized)); - } - if let Some(req_format) = non_default_params.get("req_format") { - mapped.insert( - "req_format".to_string(), - Value::String(normalize_req_format(req_format)?), - ); - } - Ok(mapped) -} - -pub fn complete_document_intelligence_url( - api_base: Option<&str>, - model: &str, - optional_params: &Map, - env_lookup: &dyn Fn(&str) -> Option, -) -> Result { - let endpoint = resolve_document_intelligence_endpoint(api_base, env_lookup)?; - let mut url = format!( - "{}/documentintelligence/documentModels/{}:analyze?api-version={}", - endpoint.trim_end_matches('/'), - encode_model_id(model)?, - AZURE_DOCUMENT_INTELLIGENCE_API_VERSION - ); - - if let Some(pages) = optional_params.get("pages") - && let Some(normalized) = normalize_pages_param(pages)? - { - url.push_str("&pages="); - url.push_str(&normalized); - } - - if let Some(features) = optional_params.get("features") - && let Some(normalized) = normalize_features_param(features)? - { - url.push_str("&features="); - url.push_str(&normalized); - } - - if let Some(req_format) = optional_params.get("req_format") { - normalize_req_format(req_format)?; - } - - Ok(url) -} - -fn document_url_from_mistral_document(document: &Value) -> Result<&str, 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 field_name = match doc_type { - "document_url" => "document_url", - "image_url" => "image_url", - other => { - return Err(Error::InvalidRequest(format!( - "Invalid document type: {other}. Must be 'document_url' or 'image_url'" - ))); - } - }; - object - .get(field_name) - .and_then(Value::as_str) - .filter(|value| !value.is_empty()) - .ok_or(Error::MissingField(field_name)) -} - -fn extract_base64_from_data_uri(data_uri: &str) -> &str { - data_uri - .split_once(',') - .map(|(_, data)| data) - .unwrap_or(data_uri) -} - -fn page_markdown(page: &Map) -> String { - page.get("lines") - .and_then(Value::as_array) - .map(|lines| { - lines - .iter() - .filter_map(|line| line.get("content").and_then(Value::as_str)) - .collect::>() - .join("\n") - }) - .unwrap_or_default() -} - -fn page_dimensions(page: &Map) -> Value { - let width = page.get("width").and_then(Value::as_f64).unwrap_or(8.5); - let height = page.get("height").and_then(Value::as_f64).unwrap_or(11.0); - let unit = page.get("unit").and_then(Value::as_str).unwrap_or("inch"); - let (width, height) = if unit == "inch" { - ( - (width * AZURE_DOCUMENT_INTELLIGENCE_DEFAULT_DPI as f64) as i64, - (height * AZURE_DOCUMENT_INTELLIGENCE_DEFAULT_DPI as f64) as i64, - ) - } else { - (width as i64, height as i64) - }; - json!({ - "width": width, - "height": height, - "dpi": AZURE_DOCUMENT_INTELLIGENCE_DEFAULT_DPI, - }) -} - -fn transform_document_intelligence_response( - model: &str, - response_json: Value, - preserve_native_response: bool, -) -> Result { - let response = response_json - .as_object() - .ok_or_else(|| Error::InvalidType { - expected: "object", - actual: json_type_name(&response_json), - })?; - let status = response - .get("status") - .and_then(Value::as_str) - .ok_or(Error::MissingField("status"))?; - if status != "succeeded" { - return Err(Error::InvalidResponse(format!( - "Azure Document Intelligence analysis failed with status: {status}" - ))); - } - - let analyze_result = response.get("analyzeResult").and_then(Value::as_object); - let azure_pages = analyze_result - .and_then(|result| result.get("pages")) - .and_then(Value::as_array) - .cloned() - .unwrap_or_default(); - let pages = azure_pages - .iter() - .filter_map(Value::as_object) - .map(|page| { - let page_number = page.get("pageNumber").and_then(Value::as_i64).unwrap_or(1); - json!({ - "index": page_number - 1, - "markdown": page_markdown(page), - "dimensions": page_dimensions(page), - }) - }) - .collect::>(); - let extra_fields = ["content", "tables", "keyValuePairs"] - .into_iter() - .map(|field| { - ( - field.to_string(), - analyze_result - .and_then(|result| result.get(field)) - .cloned() - .unwrap_or(Value::Null), - ) - }) - .collect(); - - Ok(LiteLLMOcrResponse { - usage_info: Some(json!({ - "pages_processed": pages.len(), - "doc_size_bytes": null, - })), - pages, - model: model.to_string(), - document_annotation: None, - object: "ocr".to_string(), - extra_fields, - provider_native_response: preserve_native_response.then_some(response_json), - }) -} - -impl OcrProviderConfig for AzureAiOcrConfig { - 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, - ) -> Result { - MISTRAL_OCR_CONFIG.transform_ocr_request(model, document, optional_params) - } - - fn transform_ocr_response( - &self, - model: &str, - response_json: Value, - ) -> Result { - 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, - env_lookup: &dyn Fn(&str) -> Option, - ) -> Result { - complete_azure_ai_url(api_base, env_lookup) - } - - fn resolve_api_key( - &self, - api_key: Option<&str>, - env_lookup: &dyn Fn(&str) -> Option, - ) -> Result { - resolve_azure_ai_api_key(api_key, env_lookup) - } - - fn requires_data_uri_document(&self) -> bool { - true - } -} - -impl OcrProviderConfig for AzureDocumentIntelligenceOcrConfig { - #[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)] - fn supported_ocr_params(&self) -> &'static [&'static str] { - AZURE_DOCUMENT_INTELLIGENCE_SUPPORTED_OCR_PARAMS - } - - #[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)] - fn map_ocr_params(&self, non_default_params: &Map) -> Map { - map_document_intelligence_ocr_params(non_default_params).unwrap_or_else(|_| { - non_default_params - .iter() - .filter(|(name, _)| { - AZURE_DOCUMENT_INTELLIGENCE_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, - ) -> Result { - let document_url = document_url_from_mistral_document(&document)?; - let mut data = Map::new(); - if document_url.starts_with("data:") { - data.insert( - "base64Source".to_string(), - Value::String(extract_base64_from_data_uri(document_url).to_string()), - ); - } else { - data.insert( - "urlSource".to_string(), - Value::String(document_url.to_string()), - ); - } - 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 { - transform_document_intelligence_response(model, response_json, false) - } - - fn transform_ocr_response_with_params( - &self, - model: &str, - response_json: Value, - optional_params: &Map, - ) -> Result { - transform_document_intelligence_response( - model, - response_json, - optional_params.get("req_format").and_then(Value::as_str) == Some("native"), - ) - } - - #[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)] - fn complete_url( - &self, - api_base: Option<&str>, - model: &str, - optional_params: &Map, - env_lookup: &dyn Fn(&str) -> Option, - ) -> Result { - complete_document_intelligence_url(api_base, model, optional_params, env_lookup) - } - - fn resolve_api_key( - &self, - api_key: Option<&str>, - env_lookup: &dyn Fn(&str) -> Option, - ) -> Result { - resolve_document_intelligence_api_key(api_key, env_lookup) - } - - fn auth_strategy(&self) -> OcrAuthStrategy { - OcrAuthStrategy::Header("Ocp-Apim-Subscription-Key") - } - - fn response_handling(&self) -> OcrResponseHandling { - OcrResponseHandling::AzureDocumentIntelligencePoll - } -} - -#[cfg(test)] -mod tests { - use super::*; - use rstest::{fixture, rstest}; - - const ENDPOINT: &str = "https://example.cognitiveservices.azure.com"; - - #[fixture] - fn document_intelligence_config() -> AzureDocumentIntelligenceOcrConfig { - AzureDocumentIntelligenceOcrConfig - } - - fn header_value<'a>(headers: &'a [(String, String)], name: &str) -> Option<&'a str> { - headers - .iter() - .find(|(header_name, _)| header_name.eq_ignore_ascii_case(name)) - .map(|(_, value)| value.as_str()) - } - - #[fixture] - fn native_operation() -> Value { - json!({ - "status": "succeeded", - "createdDateTime": "2026-07-02T00:00:00Z", - "lastUpdatedDateTime": "2026-07-02T00:00:05Z", - "analyzeResult": { - "content": "Invoice\nInvoice No: INV-12345\nTotal: $100.00", - "pages": [{ - "pageNumber": 1, - "width": 8.5, - "height": 11, - "unit": "inch", - "angle": 0.13, - "lines": [ - {"content": "Invoice"}, - {"content": "Invoice No: INV-12345"}, - {"content": "Total: $100.00"} - ], - "words": [{"content": "Invoice", "confidence": 0.994}] - }], - "tables": [ - { - "rowCount": 2, - "columnCount": 2, - "cells": [ - {"kind": "columnHeader", "rowIndex": 0, "columnIndex": 0, "content": "Item"}, - {"kind": "columnHeader", "rowIndex": 0, "columnIndex": 1, "content": "Price"}, - {"rowIndex": 1, "columnIndex": 0, "content": "Widget"}, - {"rowIndex": 1, "columnIndex": 1, "content": "$100.00"} - ] - }, - { - "rowCount": 1, - "columnCount": 1, - "cells": [{"rowIndex": 0, "columnIndex": 0, "content": "Totals"}] - } - ], - "keyValuePairs": [ - { - "key": {"content": "Invoice No"}, - "value": {"content": "INV-12345"}, - "confidence": 0.98 - }, - { - "key": {"content": "Total"}, - "value": {"content": "$100.00"}, - "confidence": 0.95 - } - ], - "paragraphs": [{"content": "Invoice"}] - } - }) - } - - fn assert_native_fields_preserved(response: &LiteLLMOcrResponse, operation: &Value) { - let analyze_result = &operation["analyzeResult"]; - - assert_eq!(response.extra_fields["content"], analyze_result["content"]); - assert_eq!(response.extra_fields["tables"], analyze_result["tables"]); - assert_eq!( - response.extra_fields["keyValuePairs"], - analyze_result["keyValuePairs"] - ); - assert_eq!(response.object, "ocr"); - assert_eq!( - response.usage_info, - Some(json!({"pages_processed": 1, "doc_size_bytes": null})) - ); - assert_eq!(response.pages[0]["index"], 0); - assert_eq!( - response.pages[0]["markdown"], - "Invoice\nInvoice No: INV-12345\nTotal: $100.00" - ); - assert_eq!( - response.pages[0]["dimensions"], - json!({"width": 816, "height": 1056, "dpi": 96}) - ); - } - - #[test] - fn azure_ai_reuses_mistral_body_transform() { - let body = AZURE_AI_OCR_CONFIG - .transform_ocr_request( - "pixtral-12b-2409", - json!({"type": "document_url", "document_url": "data:application/pdf;base64,abc"}), - serde_json::Map::from_iter([("include_image_base64".to_string(), json!(true))]), - ) - .expect("request transforms") - .data; - - assert_eq!(body["model"], "pixtral-12b-2409"); - assert_eq!(body["include_image_base64"], true); - assert_eq!( - body["document"]["document_url"], - "data:application/pdf;base64,abc" - ); - } - - #[test] - fn document_intelligence_url_normalizes_zero_based_pages() { - let params = serde_json::Map::from_iter([("pages".to_string(), json!([2, 0, 2]))]); - let url = complete_document_intelligence_url( - Some("https://example.cognitiveservices.azure.com/"), - "azure_ai/doc-intelligence/prebuilt-layout", - ¶ms, - &|_| None, - ) - .expect("url builds"); - - assert_eq!( - url, - "https://example.cognitiveservices.azure.com/documentintelligence/documentModels/prebuilt-layout:analyze?api-version=2024-11-30&pages=1,3" - ); - } - - #[test] - fn document_intelligence_url_normalizes_features() { - let params = serde_json::Map::from_iter([( - "features".to_string(), - json!("keyValuePairs, languages"), - )]); - let url = complete_document_intelligence_url( - Some("https://example.cognitiveservices.azure.com"), - "prebuilt-layout", - ¶ms, - &|_| None, - ) - .expect("url builds"); - - assert_eq!( - url, - "https://example.cognitiveservices.azure.com/documentintelligence/documentModels/prebuilt-layout:analyze?api-version=2024-11-30&features=keyValuePairs,languages" - ); - } - - #[test] - fn document_intelligence_url_combines_pages_and_feature_list() { - let params = serde_json::Map::from_iter([ - ("pages".to_string(), json!([0, 1, 2])), - ( - "features".to_string(), - json!([" keyValuePairs ", "languages"]), - ), - ]); - let url = complete_document_intelligence_url( - Some("https://example.cognitiveservices.azure.com"), - "prebuilt-layout", - ¶ms, - &|_| None, - ) - .expect("url builds"); - - assert_eq!( - url, - "https://example.cognitiveservices.azure.com/documentintelligence/documentModels/prebuilt-layout:analyze?api-version=2024-11-30&pages=1,2,3&features=keyValuePairs,languages" - ); - } - - #[test] - fn document_intelligence_url_omits_empty_feature_list() { - let params = serde_json::Map::from_iter([("features".to_string(), json!([]))]); - assert!( - map_document_intelligence_ocr_params(¶ms) - .expect("empty features map") - .is_empty() - ); - let url = complete_document_intelligence_url( - Some("https://example.cognitiveservices.azure.com"), - "prebuilt-layout", - ¶ms, - &|_| None, - ) - .expect("url builds"); - - assert_eq!( - url, - "https://example.cognitiveservices.azure.com/documentintelligence/documentModels/prebuilt-layout:analyze?api-version=2024-11-30" - ); - } - - #[rstest] - #[case::query_injection(json!("keyValuePairs&pages=9"))] - #[case::spaces(json!("key value pairs"))] - #[case::empty_string(json!(""))] - #[case::integer_list(json!([1, 2]))] - #[case::nested_list(json!([["keyValuePairs"]]))] - #[case::object(json!({"feature": "keyValuePairs"}))] - #[case::number(json!(5))] - fn document_intelligence_mapping_rejects_invalid_features(#[case] features: Value) { - let params = serde_json::Map::from_iter([("features".to_string(), features)]); - let error = - map_document_intelligence_ocr_params(¶ms).expect_err("invalid features must fail"); - - assert!(matches!( - error, - Error::InvalidRequest(message) if message.contains("Invalid `features`") - )); - } - - #[rstest] - #[case::single_list(json!(["keyValuePairs"]), "keyValuePairs")] - #[case::multiple_list( - json!(["keyValuePairs", "languages"]), - "keyValuePairs,languages" - )] - #[case::single_string(json!("keyValuePairs"), "keyValuePairs")] - #[case::comma_separated(json!("keyValuePairs,languages"), "keyValuePairs,languages")] - #[case::spaces(json!("keyValuePairs, languages"), "keyValuePairs,languages")] - fn document_intelligence_maps_features(#[case] features: Value, #[case] expected: &str) { - let params = Map::from_iter([ - ("features".to_string(), features), - ("unsupported".to_string(), json!(true)), - ]); - - assert_eq!( - AZURE_DOCUMENT_INTELLIGENCE_OCR_CONFIG.map_ocr_params(¶ms), - Map::from_iter([("features".to_string(), json!(expected))]) - ); - } - - #[test] - fn document_intelligence_request_uses_base64_source_for_data_uri() { - let body = AZURE_DOCUMENT_INTELLIGENCE_OCR_CONFIG - .transform_ocr_request( - "prebuilt-read", - json!({"type": "document_url", "document_url": "data:application/pdf;base64,abc123"}), - Map::new(), - ) - .expect("request transforms") - .data; - - assert_eq!(body, json!({"base64Source": "abc123"})); - } - - #[rstest] - fn document_intelligence_response_normalizes_pages(native_operation: Value) { - let response = AZURE_DOCUMENT_INTELLIGENCE_OCR_CONFIG - .transform_ocr_response("prebuilt-layout", native_operation.clone()) - .expect("response transforms"); - - assert_native_fields_preserved(&response, &native_operation); - } - - #[test] - fn azure_document_intelligence_model_id_is_encoded() { - let url = complete_document_intelligence_url( - Some(ENDPOINT), - "prebuilt-layout?x=1#frag", - &Map::new(), - &|_| None, - ) - .expect("url builds"); - - assert_eq!( - url, - "https://example.cognitiveservices.azure.com/documentintelligence/documentModels/prebuilt-layout%3Fx%3D1%23frag:analyze?api-version=2024-11-30" - ); - } - - #[test] - fn azure_document_intelligence_dot_segment_model_id_is_rejected() { - let error = complete_document_intelligence_url( - Some(ENDPOINT), - "azure_ai/doc-intelligence/..", - &Map::new(), - &|_| None, - ) - .expect_err("dot segment must fail"); - - assert_eq!( - error, - Error::InvalidRequest("model_id cannot be a dot path segment".to_string()) - ); - } - - #[rstest] - fn document_intelligence_async_response_preserves_normalized_fields(native_operation: Value) { - let response = AZURE_DOCUMENT_INTELLIGENCE_OCR_CONFIG - .transform_ocr_response( - "azure_ai/doc-intelligence/prebuilt-layout", - native_operation.clone(), - ) - .expect("response transforms"); - - assert_native_fields_preserved(&response, &native_operation); - } - - #[test] - fn document_intelligence_response_tolerates_missing_native_fields() { - let response = AZURE_DOCUMENT_INTELLIGENCE_OCR_CONFIG - .transform_ocr_response( - "azure_ai/doc-intelligence/prebuilt-read", - json!({ - "status": "succeeded", - "analyzeResult": { - "pages": [{ - "pageNumber": 1, - "width": 8.5, - "height": 11, - "unit": "inch", - "lines": [{"content": "hello"}] - }] - } - }), - ) - .expect("missing optional fields are allowed"); - - assert_eq!(response.pages[0]["markdown"], "hello"); - assert_eq!(response.extra_fields["content"], Value::Null); - assert_eq!(response.extra_fields["tables"], Value::Null); - assert_eq!(response.extra_fields["keyValuePairs"], Value::Null); - } - - #[test] - fn document_intelligence_non_succeeded_status_is_rejected() { - let error = AZURE_DOCUMENT_INTELLIGENCE_OCR_CONFIG - .transform_ocr_response( - "azure_ai/doc-intelligence/prebuilt-layout", - json!({"status": "failed"}), - ) - .expect_err("failed status must fail"); - - assert_eq!( - error, - Error::InvalidResponse( - "Azure Document Intelligence analysis failed with status: failed".to_string() - ) - ); - } - - #[test] - fn document_intelligence_supported_params_include_features() { - assert_eq!( - AZURE_DOCUMENT_INTELLIGENCE_OCR_CONFIG.supported_ocr_params(), - &["pages", "features", "req_format"] - ); - } - - #[rstest] - fn document_intelligence_native_format_carries_raw_operation(native_operation: Value) { - let response = AZURE_DOCUMENT_INTELLIGENCE_OCR_CONFIG - .transform_ocr_response_with_params( - "azure_ai/doc-intelligence/prebuilt-layout", - native_operation.clone(), - &Map::from_iter([("req_format".to_string(), json!("native"))]), - ) - .expect("native response transforms"); - - assert_eq!( - response.provider_native_response, - Some(native_operation.clone()) - ); - assert_native_fields_preserved(&response, &native_operation); - } - - #[rstest] - fn document_intelligence_async_native_format_carries_raw_operation(native_operation: Value) { - let response = AZURE_DOCUMENT_INTELLIGENCE_OCR_CONFIG - .transform_ocr_response_with_params( - "azure_ai/doc-intelligence/prebuilt-layout", - native_operation.clone(), - &Map::from_iter([("req_format".to_string(), json!("native"))]), - ) - .expect("native response transforms"); - - assert_eq!( - response.provider_native_response, - Some(native_operation.clone()) - ); - assert_native_fields_preserved(&response, &native_operation); - } - - #[rstest] - #[case::default(Map::new())] - #[case::litellm(Map::from_iter([("req_format".to_string(), json!("litellm"))]))] - fn document_intelligence_default_format_omits_raw_operation( - #[case] optional_params: Map, - native_operation: Value, - ) { - let response = AZURE_DOCUMENT_INTELLIGENCE_OCR_CONFIG - .transform_ocr_response_with_params( - "azure_ai/doc-intelligence/prebuilt-layout", - native_operation.clone(), - &optional_params, - ) - .expect("response transforms"); - - assert_eq!(response.provider_native_response, None); - assert_native_fields_preserved(&response, &native_operation); - } - - #[rstest] - #[case::native("native")] - #[case::litellm("litellm")] - fn document_intelligence_maps_req_format(#[case] req_format: &str) { - let mapped = map_document_intelligence_ocr_params(&Map::from_iter([( - "req_format".to_string(), - json!(req_format), - )])) - .expect("req_format maps"); - - assert_eq!( - mapped, - Map::from_iter([("req_format".to_string(), json!(req_format))]) - ); - } - - #[test] - fn document_intelligence_rejects_unknown_req_format() { - let error = map_document_intelligence_ocr_params(&Map::from_iter([( - "req_format".to_string(), - json!("azure"), - )])) - .expect_err("unknown req_format must fail"); - - assert!( - matches!(error, Error::InvalidRequest(message) if message.contains("Invalid `req_format`")) - ); - } - - #[test] - fn document_intelligence_url_omits_req_format() { - let url = complete_document_intelligence_url( - Some(ENDPOINT), - "prebuilt-layout", - &Map::from_iter([("req_format".to_string(), json!("native"))]), - &|_| None, - ) - .expect("url builds"); - - assert!(!url.contains("req_format")); - } - - #[test] - fn document_intelligence_validate_environment_uses_subscription_key() { - let headers = - validate_document_intelligence_environment(Vec::new(), Some("my-key"), None, &|_| None) - .expect("api key authenticates"); - - assert_eq!( - header_value(&headers, "Ocp-Apim-Subscription-Key"), - Some("my-key") - ); - } - - #[test] - fn document_intelligence_validate_environment_falls_back_to_entra_token() { - let headers = validate_document_intelligence_environment( - Vec::new(), - None, - Some("entra-token"), - &|_| None, - ) - .expect("Entra token authenticates"); - - assert_eq!( - header_value(&headers, "Authorization"), - Some("Bearer entra-token") - ); - assert_eq!(header_value(&headers, "Ocp-Apim-Subscription-Key"), None); - } - - #[test] - fn document_intelligence_supported_params_include_pages_features_and_req_format() { - assert_eq!( - AZURE_DOCUMENT_INTELLIGENCE_OCR_CONFIG.supported_ocr_params(), - &["pages", "features", "req_format"] - ); - } - - #[test] - fn document_intelligence_maps_zero_based_page_list() { - let mapped = map_document_intelligence_ocr_params(&Map::from_iter([( - "pages".to_string(), - json!([0, 1, 2]), - )])) - .expect("pages map"); - - assert_eq!( - mapped, - Map::from_iter([("pages".to_string(), json!("1,2,3"))]) - ); - } - - #[test] - fn document_intelligence_page_mapping_dedupes_and_sorts() { - let mapped = map_document_intelligence_ocr_params(&Map::from_iter([( - "pages".to_string(), - json!([2, 0, 0, 1]), - )])) - .expect("pages map"); - - assert_eq!(mapped["pages"], "1,2,3"); - } - - #[test] - fn document_intelligence_page_mapping_omits_empty_list() { - let mapped = map_document_intelligence_ocr_params(&Map::from_iter([( - "pages".to_string(), - json!([]), - )])) - .expect("empty pages map"); - - assert!(mapped.is_empty()); - } - - #[test] - fn document_intelligence_page_mapping_accepts_native_range() { - let mapped = map_document_intelligence_ocr_params(&Map::from_iter([( - "pages".to_string(), - json!("3-9"), - )])) - .expect("range maps"); - - assert_eq!(mapped["pages"], "3-9"); - } - - #[test] - fn document_intelligence_page_mapping_strips_spaces() { - let mapped = map_document_intelligence_ocr_params(&Map::from_iter([( - "pages".to_string(), - json!("1-3, 5"), - )])) - .expect("range maps"); - - assert_eq!(mapped["pages"], "1-3,5"); - } - - #[test] - fn document_intelligence_page_mapping_accepts_string_tokens() { - let mapped = map_document_intelligence_ocr_params(&Map::from_iter([( - "pages".to_string(), - json!(["1", "3-5"]), - )])) - .expect("tokens map"); - - assert_eq!(mapped["pages"], "1,3-5"); - } - - #[test] - fn document_intelligence_page_mapping_rejects_invalid_string() { - let error = map_document_intelligence_ocr_params(&Map::from_iter([( - "pages".to_string(), - json!("a,b"), - )])) - .expect_err("invalid pages must fail"); - - assert!( - matches!(error, Error::InvalidRequest(message) if message.contains("Invalid `pages` string")) - ); - } - - #[test] - fn document_intelligence_page_mapping_rejects_negative_index() { - let error = map_document_intelligence_ocr_params(&Map::from_iter([( - "pages".to_string(), - json!([-1]), - )])) - .expect_err("negative pages must fail"); - - assert!( - matches!(error, Error::InvalidRequest(message) if message.contains("must be >= 0")) - ); - } - - #[test] - fn document_intelligence_page_mapping_rejects_bool_list() { - let error = map_document_intelligence_ocr_params(&Map::from_iter([( - "pages".to_string(), - json!([true, false]), - )])) - .expect_err("boolean pages must fail"); - - assert!( - matches!(error, Error::InvalidRequest(message) if message.contains("integers, not booleans")) - ); - } - - #[test] - fn document_intelligence_page_mapping_rejects_unsupported_type() { - let error = map_document_intelligence_ocr_params(&Map::from_iter([( - "pages".to_string(), - json!(5), - )])) - .expect_err("unsupported pages must fail"); - - assert!( - matches!(error, Error::InvalidRequest(message) if message.contains("Mistral-style")) - ); - } - - #[test] - fn document_intelligence_url_appends_pages_query() { - let url = complete_document_intelligence_url( - Some("https://example.cognitiveservices.azure.com/"), - "azure_ai/doc-intelligence/prebuilt-layout", - &Map::from_iter([("pages".to_string(), json!("1-3,5"))]), - &|_| None, - ) - .expect("url builds"); - - assert!(url.contains("api-version=2024-11-30")); - assert!(url.contains("pages=1-3,5")); - assert!(url.contains("/documentintelligence/documentModels/prebuilt-layout:analyze")); - } - - #[test] - fn document_intelligence_url_has_no_pages_when_params_are_empty() { - let url = complete_document_intelligence_url( - Some(ENDPOINT), - "prebuilt-layout", - &Map::new(), - &|_| None, - ) - .expect("url builds"); - - assert!(!url.contains("pages=")); - } - - #[rstest] - fn document_intelligence_request_keeps_pages_out_of_body( - document_intelligence_config: AzureDocumentIntelligenceOcrConfig, - ) { - let request = document_intelligence_config - .transform_ocr_request( - "prebuilt-layout", - json!({"type": "document_url", "document_url": "https://example.com/x.pdf"}), - Map::from_iter([("pages".to_string(), json!("1,2,3"))]), - ) - .expect("request transforms"); - - assert_eq!( - request.data, - json!({"urlSource": "https://example.com/x.pdf"}) - ); - } - - #[test] - fn document_intelligence_mistral_pages_flow_to_query_only() { - let mapped = map_document_intelligence_ocr_params(&Map::from_iter([( - "pages".to_string(), - json!([2, 3, 4, 5, 6, 7, 8]), - )])) - .expect("pages map"); - let url = - complete_document_intelligence_url(Some(ENDPOINT), "prebuilt-layout", &mapped, &|_| { - None - }) - .expect("url builds"); - let request = AZURE_DOCUMENT_INTELLIGENCE_OCR_CONFIG - .transform_ocr_request( - "prebuilt-layout", - json!({"type": "document_url", "document_url": "https://example.com/x.pdf"}), - mapped, - ) - .expect("request transforms"); - - assert!(url.contains("pages=3,4,5,6,7,8,9")); - assert_eq!( - request.data, - json!({"urlSource": "https://example.com/x.pdf"}) - ); - } - - #[test] - fn document_intelligence_endpoint_ignores_generic_azure_ai_base() { - let resolved = resolve_document_intelligence_endpoint(None, &|name| match name { - AZURE_DOCUMENT_INTELLIGENCE_ENDPOINT_ENV => Some(ENDPOINT.to_string()), - AZURE_AI_API_BASE_ENV => Some("https://generic.example.com".to_string()), - _ => None, - }) - .expect("endpoint resolves"); - - assert_eq!(resolved, ENDPOINT); - } - - #[test] - fn document_intelligence_endpoint_honors_explicit_api_base() { - let resolved = resolve_document_intelligence_endpoint( - Some("https://my-di.cognitiveservices.azure.com"), - &|name| match name { - AZURE_DOCUMENT_INTELLIGENCE_ENDPOINT_ENV => Some(ENDPOINT.to_string()), - AZURE_AI_API_BASE_ENV => Some("https://generic.example.com".to_string()), - _ => None, - }, - ) - .expect("endpoint resolves"); - - assert_eq!(resolved, "https://my-di.cognitiveservices.azure.com"); - } - - #[test] - fn azure_ai_mistral_ocr_uses_generic_api_base() { - let resolved = resolve_azure_ai_api_base(None, &|name| match name { - AZURE_AI_API_BASE_ENV => Some("https://generic-azure-ai.example.com".to_string()), - AZURE_DOCUMENT_INTELLIGENCE_ENDPOINT_ENV => Some(ENDPOINT.to_string()), - _ => None, - }) - .expect("api base resolves"); - - assert_eq!(resolved, "https://generic-azure-ai.example.com"); - } - - #[test] - fn azure_ai_ocr_authenticates_with_entra_token() { - let headers = - validate_azure_ai_environment(Vec::new(), None, Some("entra-token"), &|_| None) - .expect("Entra token authenticates"); - - assert_eq!( - header_value(&headers, "Authorization"), - Some("Bearer entra-token") - ); - } -} diff --git a/litellm-rust/crates/core/src/providers/mistral/mod.rs b/litellm-rust/crates/core/src/providers/mistral/mod.rs deleted file mode 100644 index 3621ff6a2fd..00000000000 --- a/litellm-rust/crates/core/src/providers/mistral/mod.rs +++ /dev/null @@ -1 +0,0 @@ -pub mod ocr; diff --git a/litellm-rust/crates/core/src/providers/mistral/ocr/mod.rs b/litellm-rust/crates/core/src/providers/mistral/ocr/mod.rs deleted file mode 100644 index f239b6921fa..00000000000 --- a/litellm-rust/crates/core/src/providers/mistral/ocr/mod.rs +++ /dev/null @@ -1 +0,0 @@ -pub mod transformation; diff --git a/litellm-rust/crates/core/src/providers/mistral/ocr/transformation.rs b/litellm-rust/crates/core/src/providers/mistral/ocr/transformation.rs deleted file mode 100644 index 044fc587c22..00000000000 --- a/litellm-rust/crates/core/src/providers/mistral/ocr/transformation.rs +++ /dev/null @@ -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, -) -> Result { - 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, - ) -> Result { - 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 { - 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, - _env_lookup: &dyn Fn(&str) -> Option, - ) -> Result { - Ok(complete_url(api_base)) - } - - fn resolve_api_key( - &self, - api_key: Option<&str>, - env_lookup: &dyn Fn(&str) -> Option, - ) -> Result { - 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) -> Map { - 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, -) -> Result { - 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 { - 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(¶m)); - } - } - - #[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(¶m)); - } - } - - #[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())); - } -} diff --git a/litellm-rust/crates/core/src/providers/mod.rs b/litellm-rust/crates/core/src/providers/mod.rs index 805600d6dbe..1aeb75063d6 100644 --- a/litellm-rust/crates/core/src/providers/mod.rs +++ b/litellm-rust/crates/core/src/providers/mod.rs @@ -2,6 +2,4 @@ pub mod anthropic; pub mod azure_ai; #[cfg(feature = "bedrock-auth")] pub mod bedrock; -pub mod mistral; pub mod openai; -pub mod vertex_ai; diff --git a/litellm-rust/crates/core/src/providers/vertex_ai/mod.rs b/litellm-rust/crates/core/src/providers/vertex_ai/mod.rs deleted file mode 100644 index 3621ff6a2fd..00000000000 --- a/litellm-rust/crates/core/src/providers/vertex_ai/mod.rs +++ /dev/null @@ -1 +0,0 @@ -pub mod ocr; diff --git a/litellm-rust/crates/core/src/providers/vertex_ai/ocr/mod.rs b/litellm-rust/crates/core/src/providers/vertex_ai/ocr/mod.rs deleted file mode 100644 index f239b6921fa..00000000000 --- a/litellm-rust/crates/core/src/providers/vertex_ai/ocr/mod.rs +++ /dev/null @@ -1 +0,0 @@ -pub mod transformation; diff --git a/litellm-rust/crates/core/src/providers/vertex_ai/ocr/transformation.rs b/litellm-rust/crates/core/src/providers/vertex_ai/ocr/transformation.rs deleted file mode 100644 index c2d810822ae..00000000000 --- a/litellm-rust/crates/core/src/providers/vertex_ai/ocr/transformation.rs +++ /dev/null @@ -1,361 +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}; - -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 VertexAiDeepSeekOcrConfig; - -pub const VERTEX_AI_DEEPSEEK_OCR_CONFIG: VertexAiDeepSeekOcrConfig = VertexAiDeepSeekOcrConfig; - -fn string_param<'a>(params: &'a Map, 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, -) -> Result { - 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, - env_lookup: &dyn Fn(&str) -> Option, -) -> Result { - 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, - env_lookup: &dyn Fn(&str) -> Option, -) -> 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()) -} - -pub fn complete_vertex_deepseek_url( - api_base: Option<&str>, - optional_params: &Map, - env_lookup: &dyn Fn(&str) -> Option, -) -> Result { - 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 { - 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 { - 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, 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 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) -> Map { - 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, - ) -> Result { - 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 { - 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, - env_lookup: &dyn Fn(&str) -> Option, - ) -> Result { - 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, - ) -> Result { - resolve_vertex_api_key(api_key, env_lookup) - } -} - -#[cfg(test)] -mod tests { - use super::*; - use rstest::rstest; - - #[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}))); - } -} diff --git a/litellm-rust/crates/core/tests/deepseek_ocr.rs b/litellm-rust/crates/core/tests/deepseek_ocr.rs new file mode 100644 index 00000000000..875fc9e3dc6 --- /dev/null +++ b/litellm-rust/crates/core/tests/deepseek_ocr.rs @@ -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(), ¶ms).unwrap(), + ) + .unwrap(); + assert_eq!(result["model"], "deepseek-ai/deepseek-ocr-maas"); + assert_eq!( + result["messages"][0]["content"][0], + json!({"type":"image_url","image_url":"gs://bucket/a.png"}) + ); + assert_eq!(result[name], value); + assert!(result.get("ignored").is_none()); +} + +#[rstest] +#[case(json!("# 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::(value) + .map_err(|_| ()) + .and_then(|response| transform_ocr_response("model", response).map_err(|_| ())); + assert!(result.is_err()); + } +} diff --git a/litellm-rust/crates/core/tests/vertex_ai_deepseek_ocr.rs b/litellm-rust/crates/core/tests/vertex_ai_deepseek_ocr.rs new file mode 100644 index 00000000000..6d3061d8f5d --- /dev/null +++ b/litellm-rust/crates/core/tests/vertex_ai_deepseek_ocr.rs @@ -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") + ); +} diff --git a/litellm-rust/crates/core/tests/vertex_ai_ocr.rs b/litellm-rust/crates/core/tests/vertex_ai_ocr.rs index 2358552c742..96a19dd62b4 100644 --- a/litellm-rust/crates/core/tests/vertex_ai_ocr.rs +++ b/litellm-rust/crates/core/tests/vertex_ai_ocr.rs @@ -2,7 +2,6 @@ use serde_json::{Value, json}; use super::test_support::{MockResponse, mock_server, perform_ocr, wire_request}; use crate::auth::InputSource; -use crate::ocr::wire::{OcrWireRequest, decode_request}; fn request_body(request: &str) -> Value { serde_json::from_str(request.split_once("\r\n\r\n").unwrap().1).unwrap() @@ -81,30 +80,82 @@ async fn invalid_credentials_fail_before_provider_http() { #[tokio::test] async fn request_controlled_api_base_is_rejected_before_vertex_auth() { - let request = decode_request(OcrWireRequest { - model: "vertex_ai/model".into(), - document: json!({"type":"document_url","document_url":"data:application/pdf;base64,YWJj"}), - api_key: Some("test-key".into()), - api_base: Some("https://attacker.example".into()), - custom_llm_provider: None, - extra_headers: None, - optional_params: json!({"vertex_project":"project-1"}) - .as_object() - .unwrap() - .clone(), - input_sources: std::collections::BTreeMap::from([( - "api_base".to_string(), - InputSource::Request, - )]), - timeout_seconds: Some(2.0), - }) - .unwrap(); + 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"); +} diff --git a/litellm-rust/crates/python-bridge/src/errors.rs b/litellm-rust/crates/python-bridge/src/errors.rs index 864f30db6a9..e1f458ea0bc 100644 --- a/litellm-rust/crates/python-bridge/src/errors.rs +++ b/litellm-rust/crates/python-bridge/src/errors.rs @@ -43,7 +43,6 @@ pub(crate) fn chat_completions_error_to_pyerr(err: Error) -> PyErr { | Error::MissingField(_) | Error::MissingApiKey { .. } | Error::MissingAzureAiCredentials - | Error::MissingAzureAiCredentialsOrAdToken | Error::MissingAzureDocumentIntelligenceCredentials | Error::MissingReductoApiKey | Error::Routing(_) diff --git a/litellm-rust/crates/python-bridge/src/routes/ocr.rs b/litellm-rust/crates/python-bridge/src/routes/ocr.rs index 0aab11e3cfc..c5def64c2f1 100644 --- a/litellm-rust/crates/python-bridge/src/routes/ocr.rs +++ b/litellm-rust/crates/python-bridge/src/routes/ocr.rs @@ -112,6 +112,6 @@ mod tests { 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"))); + assert!(is_supported_request("deepseek-ocr", Some("vertex_ai"))); } }