diff --git a/litellm-rust/Cargo.lock b/litellm-rust/Cargo.lock index 7f794fede17..9e1bfb0786c 100644 --- a/litellm-rust/Cargo.lock +++ b/litellm-rust/Cargo.lock @@ -2027,6 +2027,7 @@ dependencies = [ "serde_path_to_error", "serde_with", "sha2 0.10.9", + "strum", "subtle", "thiserror 2.0.19", "tokio", diff --git a/litellm-rust/crates/core/Cargo.toml b/litellm-rust/crates/core/Cargo.toml index d9adcd3e31b..8991cf294f5 100644 --- a/litellm-rust/crates/core/Cargo.toml +++ b/litellm-rust/crates/core/Cargo.toml @@ -4,7 +4,6 @@ version = "0.1.0" edition.workspace = true license.workspace = true repository.workspace = true -autotests = false [dependencies] bytes.workspace = true @@ -25,6 +24,7 @@ serde.workspace = true serde_json = { workspace = true, features = ["preserve_order"] } serde_with.workspace = true serde_path_to_error = "0.1" +strum.workspace = true subtle.workspace = true tokio = { workspace = true, features = ["sync"] } tokio-tungstenite.workspace = true @@ -32,6 +32,7 @@ thiserror.workspace = true sha2.workspace = true url.workspace = true veil.workspace = true + [features] default = [] observability = ["dep:tracing-subscriber"] diff --git a/litellm-rust/crates/core/src/call_lifecycle/host.rs b/litellm-rust/crates/core/src/call_lifecycle/host.rs index 97eb9c4c650..a62b61f1b34 100644 --- a/litellm-rust/crates/core/src/call_lifecycle/host.rs +++ b/litellm-rust/crates/core/src/call_lifecycle/host.rs @@ -120,3 +120,126 @@ impl HostLifecycle { }; } } + +#[cfg(test)] +mod tests { + use super::{HostFailure, HostLifecycle, HostPhase}; + + fn run( + fail_at: Option, + asynchronous: bool, + ) -> (Vec, Vec) { + let mut lifecycle = HostLifecycle::new(asynchronous); + let mut events = Vec::new(); + let mut failures = Vec::new(); + + while lifecycle.phase() != HostPhase::Complete { + let phase = lifecycle.phase(); + events.push(phase); + let result = if Some(phase) == fail_at { + Err(HostFailure::Error(crate::ocr::Error::InvalidRequest( + "selected failure".into(), + ))) + } else { + Ok(()) + }; + if let Some(error) = lifecycle.accept(result) { + failures.push(error); + } + } + (events, failures) + } + + #[test] + fn public_outcome_is_finalized_before_a_single_terminal_dispatch() { + for asynchronous in [false, true] { + let (events, failures) = run(None, asynchronous); + assert!(failures.is_empty()); + assert_eq!( + &events[events.len() - 2..], + &[HostPhase::Finalize, HostPhase::Success] + ); + assert_eq!( + events + .iter() + .filter(|phase| **phase == HostPhase::Execute) + .count(), + 1 + ); + assert_eq!( + events.contains(&HostPhase::DeploymentPostCall), + asynchronous + ); + } + } + + #[test] + fn only_provider_and_response_construction_failures_use_provider_mapping() { + for phase in [ + HostPhase::Setup, + HostPhase::DeploymentPreCall, + HostPhase::Prepare, + HostPhase::Execute, + HostPhase::ConstructResponse, + HostPhase::DeploymentPostCall, + HostPhase::Finalize, + ] { + let (events, failures) = run(Some(phase), true); + assert_eq!(failures.len(), 1); + assert!(!events.contains(&HostPhase::Success)); + let mapped = matches!(phase, HostPhase::Execute | HostPhase::ConstructResponse); + assert_eq!(events.contains(&HostPhase::MapFailure), mapped); + assert_eq!(events.contains(&HostPhase::DeploymentFailure), mapped); + assert_eq!( + &events[events.len() - 2..], + &[HostPhase::Failure, HostPhase::AsyncFailure] + ); + assert!( + events + .iter() + .filter(|phase| **phase == HostPhase::Execute) + .count() + <= 1 + ); + } + } + + #[test] + fn failure_handler_errors_do_not_replace_selected_failure_or_suppress_async_dispatch() { + let mut lifecycle = HostLifecycle::new(true); + while lifecycle.phase() != HostPhase::Execute { + lifecycle.accept::(Ok(())); + } + let selected = crate::ocr::Error::InvalidRequest("provider".into()); + assert_eq!( + lifecycle.accept(Err(HostFailure::Error(selected.clone()))), + Some(selected) + ); + lifecycle.accept::(Ok(())); + for phase in [ + HostPhase::DeploymentFailure, + HostPhase::Failure, + HostPhase::AsyncFailure, + ] { + assert_eq!(lifecycle.phase(), phase); + assert_eq!( + lifecycle.accept(Err(HostFailure::Error(crate::ocr::Error::InvalidRequest( + "callback".into() + )))), + None + ); + } + assert_eq!(lifecycle.phase(), HostPhase::Complete); + } + + #[test] + fn cancellation_skips_terminal_dispatch() { + let mut lifecycle = HostLifecycle::new(true); + let error = crate::ocr::Error::InvalidRequest("cancelled".into()); + assert_eq!( + lifecycle.accept(Err(HostFailure::Cancelled(error.clone()))), + Some(error) + ); + assert_eq!(lifecycle.phase(), HostPhase::Complete); + } +} diff --git a/litellm-rust/crates/core/src/call_lifecycle/mod.rs b/litellm-rust/crates/core/src/call_lifecycle/mod.rs index dce240c3d2b..fcd9cbd2ab8 100644 --- a/litellm-rust/crates/core/src/call_lifecycle/mod.rs +++ b/litellm-rust/crates/core/src/call_lifecycle/mod.rs @@ -2,9 +2,6 @@ use std::future::Future; use std::time::{Instant, SystemTime, UNIX_EPOCH}; pub mod host; -#[cfg(test)] -#[path = "../../tests/host_lifecycle.rs"] -mod host_tests; pub mod types; pub use types::{ diff --git a/litellm-rust/crates/core/src/llms/azure_ai/ocr/cohere_parse_transformation.rs b/litellm-rust/crates/core/src/llms/azure_ai/ocr/cohere_parse_transformation.rs index 97c327b9e69..503c826da67 100644 --- a/litellm-rust/crates/core/src/llms/azure_ai/ocr/cohere_parse_transformation.rs +++ b/litellm-rust/crates/core/src/llms/azure_ai/ocr/cohere_parse_transformation.rs @@ -1,12 +1,11 @@ use crate::llms::base_llm::ocr::transformation::{BaseOcrConfig, OcrRequestContext}; use crate::llms::cohere::ocr::transformation::{CohereParseConfig, CohereRequest}; -use crate::llms::cohere::ocr::{CohereOptions, CohereResponse, validate_document}; +use crate::llms::cohere::ocr::{CohereOptions, validate_document}; use crate::ocr::OcrClient; use crate::ocr::document::{inline_remote_document, validate_inline_document}; use crate::ocr::prepare::{credential_env, transform_request_body}; use crate::ocr::types::{LiteLLMOcrRequest, LiteLLMOcrResponse, OcrDocument}; use crate::url_utils::ApiUrl; -use litellm_auth_azure::AzureAuthInputs; const AZURE_AI_API_BASE_ENV: &str = "AZURE_AI_API_BASE"; @@ -16,7 +15,58 @@ pub(crate) struct AzureAICohereParseConfig; impl BaseOcrConfig for AzureAICohereParseConfig { type OcrParams = CohereOptions; type ProviderRequest = CohereRequest; - type ProviderResponse = CohereResponse; + type Environment = Vec<(String, String)>; + + fn get_api_key_env_var(&self) -> Option<&'static str> { + super::transformation::AzureAIOCRConfig.get_api_key_env_var() + } + + fn get_health_check_document(&self) -> OcrDocument { + CohereParseConfig.get_health_check_document() + } + + async fn validate_environment( + &self, + request: &LiteLLMOcrRequest, + client: &OcrClient, + ) -> Result { + BaseOcrConfig::validate_environment( + &super::transformation::AzureAIOCRConfig, + request, + client, + ) + .await + } + + fn get_complete_url( + &self, + request: &LiteLLMOcrRequest, + _params: &Self::OcrParams, + _environment: &Self::Environment, + ) -> Result { + let base = request + .connection + .api_base + .clone() + .or_else(|| credential_env(AZURE_AI_API_BASE_ENV)) + .filter(|base| !base.trim().is_empty()) + .ok_or_else(|| { + crate::ocr::Error::Auth(litellm_auth::Error::ProviderAuthentication( + "Missing Azure AI API Base - Set AZURE_AI_API_BASE or pass api_base".into(), + )) + })?; + self.get_complete_url(&base) + } + + fn transform_ocr_request( + &self, + model: &str, + document: OcrDocument, + params: &CohereOptions, + headers: &[(String, String)], + ) -> Result { + CohereParseConfig.transform_ocr_request(model, document, params, headers) + } fn get_supported_ocr_params(&self, model: &str) -> &'static [&'static str] { CohereParseConfig.get_supported_ocr_params(model) @@ -37,15 +87,16 @@ impl BaseOcrConfig for AzureAICohereParseConfig { context.connection, ) .await?; - CohereParseConfig.transform_ocr_request(model, document, optional_params, headers) + self.transform_ocr_request(model, document, optional_params, headers) } - fn normalize_response( + fn transform_ocr_response( &self, model: &str, - response: CohereResponse, + raw_response: &[u8], + request_format: crate::ocr::types::OcrResponseFormat, ) -> Result { - CohereParseConfig.normalize_response(model, response) + CohereParseConfig.transform_ocr_response(model, raw_response, request_format) } } @@ -56,28 +107,8 @@ impl AzureAICohereParseConfig { client: &OcrClient, ) -> Result { let params = self.map_ocr_params(&request.optional_params, &request.model)?; - let config = AzureAuthInputs { - azure_ad_token_provider: request.azure_ad_token_provider.clone(), - ..AzureAuthInputs::from_sourced_optional_params( - &request.optional_params, - &request.input_sources, - ) - .map_err(crate::ocr::Error::from)? - }; - let base = request - .connection - .api_base - .clone() - .or_else(|| credential_env(AZURE_AI_API_BASE_ENV)) - .filter(|base| !base.trim().is_empty()) - .ok_or_else(|| { - crate::ocr::Error::Auth(litellm_auth::Error::ProviderAuthentication( - "Missing Azure AI API Base - Set AZURE_AI_API_BASE or pass api_base".into(), - )) - })?; - let headers = super::transformation::AzureAIOCRConfig - .validate_environment(&request.connection, &config, &credential_env) - .await?; + let url = BaseOcrConfig::get_complete_url(self, request, ¶ms, &Vec::new())?; + let headers = self.validate_environment(request, client).await?; let remote = request.document.source().starts_with("http://") || request.document.source().starts_with("https://"); let body = self @@ -92,19 +123,11 @@ impl AzureAICohereParseConfig { }, ) .await?; - transform_request_body( - client, - request, - &self.get_complete_url(&base)?, - &headers, - !remote, - body, - |body| { - let document = crate::ocr::prepare::body_document(body)?; - validate_document(&document)?; - validate_inline_document(&document) - }, - ) + transform_request_body(client, request, &url, &headers, !remote, body, |body| { + let document = crate::ocr::prepare::body_document(body)?; + validate_document(&document)?; + validate_inline_document(&document) + }) .await } } diff --git a/litellm-rust/crates/core/src/llms/azure_ai/ocr/document_intelligence/transformation.rs b/litellm-rust/crates/core/src/llms/azure_ai/ocr/document_intelligence/transformation.rs index afafd5bb194..5b2a8f7838d 100644 --- a/litellm-rust/crates/core/src/llms/azure_ai/ocr/document_intelligence/transformation.rs +++ b/litellm-rust/crates/core/src/llms/azure_ai/ocr/document_intelligence/transformation.rs @@ -470,7 +470,63 @@ pub(crate) struct AzureDocumentIntelligenceOCRConfig; impl BaseOcrConfig for AzureDocumentIntelligenceOCRConfig { type OcrParams = DocumentIntelligenceParams; type ProviderRequest = DocumentIntelligenceRequest; - type ProviderResponse = AzureDocumentIntelligenceOperation; + type Environment = Vec<(String, String)>; + + fn get_api_key_env_var(&self) -> Option<&'static str> { + Some(AZURE_DI_API_KEY_ENV) + } + + fn resolve_connection_params( + &self, + api_key: Option>, + api_base: Option>, + dynamic_api_key: Option>, + dynamic_api_base: Option>, + ) -> ( + Option>, + Option>, + ) { + ( + api_key.and_then(|key| { + dynamic_api_key + .filter(|value| !value.value().is_empty()) + .or(Some(key)) + }), + api_base.and_then(|base| { + dynamic_api_base + .filter(|value| !value.value().is_empty()) + .or(Some(base)) + }), + ) + } + + async fn validate_environment( + &self, + request: &LiteLLMOcrRequest, + _client: &OcrClient, + ) -> Result { + let config = AzureAuthInputs { + azure_ad_token_provider: request.azure_ad_token_provider.clone(), + ..AzureAuthInputs::from_sourced_optional_params( + &request.optional_params, + &request.input_sources, + )? + }; + self.validate_environment(&request.connection, &config, &credential_env) + .await + } + + fn get_complete_url( + &self, + request: &LiteLLMOcrRequest, + params: &Self::OcrParams, + _environment: &Self::Environment, + ) -> Result { + let endpoint = nonblank(request.connection.api_base.clone()) + .or_else(|| nonblank(credential_env(AZURE_DI_ENDPOINT_ENV))) + .ok_or_else(|| crate::ocr::Error::Auth(litellm_auth::Error::ProviderAuthentication("Missing Azure Document Intelligence API Base - Set AZURE_DOCUMENT_INTELLIGENCE_ENDPOINT or pass api_base".into())))?; + self.get_complete_url(&endpoint, &request.model, params) + } fn get_supported_ocr_params(&self, _model: &str) -> &'static [&'static str] { &["pages", "features", "req_format"] @@ -498,12 +554,18 @@ impl BaseOcrConfig for AzureDocumentIntelligenceOCRConfig { self.transform_ocr_request(model, document, optional_params, headers) } - fn normalize_response( + fn transform_ocr_response( &self, model: &str, - response: AzureDocumentIntelligenceOperation, + raw_response: &[u8], + request_format: OcrResponseFormat, ) -> Result { - transform_completed_response(model, response) + crate::llms::base_llm::ocr::transformation::decode_and_normalize_response( + model, + raw_response, + request_format, + transform_completed_response, + ) } async fn async_transform_ocr_response( @@ -527,9 +589,6 @@ impl BaseOcrConfig for AzureDocumentIntelligenceOCRConfig { ..transform_completed_response(model, decoded.data)? }) } -} - -impl AzureDocumentIntelligenceOCRConfig { fn transform_ocr_request( &self, _model: &str, @@ -539,28 +598,17 @@ impl AzureDocumentIntelligenceOCRConfig { ) -> Result { build_request(document) } +} +impl AzureDocumentIntelligenceOCRConfig { pub(crate) async fn prepare_request( &self, request: &LiteLLMOcrRequest, client: &OcrClient, ) -> Result { let params = self.map_ocr_params(&request.optional_params, &request.model)?; - let config = AzureAuthInputs { - azure_ad_token_provider: request.azure_ad_token_provider.clone(), - ..AzureAuthInputs::from_sourced_optional_params( - &request.optional_params, - &request.input_sources, - ) - .map_err(crate::ocr::Error::from)? - }; - let headers = self - .validate_environment(&request.connection, &config, &credential_env) - .await?; - let endpoint = nonblank(request.connection.api_base.clone()) - .or_else(|| nonblank(credential_env(AZURE_DI_ENDPOINT_ENV))) - .ok_or_else(|| crate::ocr::Error::Auth(litellm_auth::Error::ProviderAuthentication("Missing Azure Document Intelligence API Base - Set AZURE_DOCUMENT_INTELLIGENCE_ENDPOINT or pass api_base".into())))?; - let url = self.get_complete_url(&endpoint, &request.model, ¶ms)?; + let headers = BaseOcrConfig::validate_environment(self, request, client).await?; + let url = BaseOcrConfig::get_complete_url(self, request, ¶ms, &headers)?; let body = self .async_transform_ocr_request( &request.model, @@ -627,7 +675,7 @@ impl AzureDocumentIntelligenceOCRConfig { let key = nonblank(connection.api_key.clone()) .map(|value| Sourced::new(value, connection.api_key_source)) .or_else(|| { - nonblank(env_lookup(AZURE_DI_API_KEY_ENV)) + nonblank(self.get_api_key_env_var().and_then(env_lookup)) .map(|value| Sourced::new(value, InputSource::Environment)) }); if let Some(key) = key { @@ -706,7 +754,7 @@ mod tests { #[test] fn response_numbers_follow_python_validation_before_dimension_conversion() { - let response = AzureDocumentIntelligenceOCRConfig.decode_and_normalize_response( + let response = AzureDocumentIntelligenceOCRConfig.transform_ocr_response( "model", br#"{"status":"succeeded","analyzeResult":{"pages":[{"pageNumber":2.0,"width":" 8.5 ","height":true}]}}"#, OcrResponseFormat::Litellm, @@ -814,4 +862,464 @@ mod tests { (AZURE_DI_SUBSCRIPTION_HEADER.into(), "request-key".into()) ); } + + use std::sync::{Arc, Mutex}; + + use crate::ocr::test_support::{MockResponse, mock_server, perform_ocr, wire_request}; + use crate::ocr::wire::{OcrWireRequest, decode_request}; + + fn query_value(url: &str, key: &str) -> Option { + url::Url::parse(url) + .unwrap() + .query_pairs() + .find_map(|(name, value)| (name == key).then(|| value.into_owned())) + } + + #[tokio::test] + async fn facade_maps_pages_features_and_url_document() { + let (base, seen, server) = mock_server(vec![MockResponse::json(json!({ + "status":"succeeded", + "analyzeResult":{"pages":[]} + }))]) + .await; + let mut request = wire_request( + "azure_ai/doc-intelligence/prebuilt-read", + &base, + json!({"pages":[2,0,0,1],"features":["keyValuePairs","languages"], "future_option": {"nested":null}, "extra_body":{"provider_option":false}}), + ); + request.document = serde_json::from_value(json!({ + "type":"document_url", + "document_url":"https://example.com/document.pdf" + })) + .unwrap(); + + perform_ocr(request).await.unwrap(); + server.await.unwrap(); + let request = &seen.lock().unwrap()[0]; + let target = request.split_whitespace().nth(1).unwrap(); + let url = format!("{base}{target}"); + assert_eq!(query_value(&url, "pages").as_deref(), Some("1,2,3")); + assert_eq!( + query_value(&url, "features").as_deref(), + Some("keyValuePairs,languages") + ); + let body: Value = serde_json::from_str(request.split_once("\r\n\r\n").unwrap().1).unwrap(); + assert_eq!( + body, + json!({"urlSource":"https://example.com/document.pdf", "future_option":{"nested":null}, "provider_option":false}) + ); + } + + #[tokio::test] + async fn rejects_invalid_pages_features_and_format() { + for options in [ + json!({"pages":[true]}), + json!({"pages":[1,"2"]}), + json!({"pages":[-1]}), + json!({"pages":"1&&features=bad"}), + json!({"features":"languages&pages=1"}), + json!({"req_format":"azure"}), + ] { + let result = decode_request(OcrWireRequest { + model: "azure_ai/doc-intelligence/prebuilt-read".into(), + document: json!({"type":"document_url","document_url":"https://example.com/a.pdf"}), + api_key: Some("key".into()), + api_base: Some("http://127.0.0.1:1".into()), + custom_llm_provider: None, + extra_headers: None, + optional_params: options.as_object().unwrap().clone().into(), + input_sources: Default::default(), + timeout_seconds: None, + }); + let rejected = match result { + Ok(request) => perform_ocr(request).await.is_err(), + Err(_) => true, + }; + assert!(rejected, "accepted {options}"); + } + } + + #[tokio::test] + async fn inline_document_decodes_to_base64_source() { + let (base, seen, server) = mock_server(vec![MockResponse::json(json!({ + "status":"succeeded" + }))]) + .await; + let request = wire_request("azure_ai/doc-intelligence/prebuilt-read", &base, json!({})); + + perform_ocr(request).await.unwrap(); + server.await.unwrap(); + let request = &seen.lock().unwrap()[0]; + let body: Value = serde_json::from_str(request.split_once("\r\n\r\n").unwrap().1).unwrap(); + assert_eq!(body, json!({"base64Source":"YWJj"})); + } + + #[tokio::test] + async fn immediate_response_normalizes_pages_and_preserves_native() { + let operation = json!({ + "status":"succeeded", + "operationExtension":42, + "analyzeResult":{ + "content":"A\n\nB", + "tables":[{"cells":[]}], + "keyValuePairs":[{"key":{"content":"A"}}], + "pages":[{ + "pageNumber":"2", + "width":"8.5", + "height":11, + "unit":"inch", + "lines":[{"content":"A"},{"content":null},{"content":"B"}] + }] + } + }); + let (base, _, server) = mock_server(vec![MockResponse::json(operation.clone())]).await; + let result = perform_ocr(wire_request( + "azure_ai/doc-intelligence/prebuilt-read", + &base, + json!({"req_format":"native"}), + )) + .await + .unwrap(); + server.await.unwrap(); + + assert_eq!(result.pages[0].index, 1); + assert_eq!(result.pages[0].markdown, "A\n\nB"); + assert_eq!( + serde_json::to_value(&result.pages[0].dimensions).unwrap(), + json!({"width":816,"height":1056,"dpi":96}) + ); + assert_eq!(result.usage_info.as_ref().unwrap().pages_processed, Some(1)); + let serialized = result.clone().into_json(); + assert_eq!(serialized["content"], "A\n\nB"); + assert_eq!(serialized["tables"], json!([{"cells":[]}])); + assert_eq!( + serialized["keyValuePairs"], + json!([{"key":{"content":"A"}}]) + ); + assert!(serialized.get("key_value_pairs").is_none()); + assert_eq!( + result.provider_native_response.as_ref(), + operation.as_object() + ); + } + + #[tokio::test] + async fn accepted_response_polls_to_success_with_only_credentials() { + let operation = json!({"status":"succeeded","analyzeResult":{"pages":[]}}); + let (base, seen, server) = mock_server(vec![ + MockResponse { + status: 202, + headers: vec![("Operation-Location", "{base}/operation".into())], + body: json!({}), + }, + MockResponse { + status: 200, + headers: vec![("Retry-After", "0".into())], + body: json!({"status":"running"}), + }, + MockResponse::json(operation.clone()), + ]) + .await; + let mut request = wire_request( + "azure_ai/doc-intelligence/prebuilt-read", + &base, + json!({"req_format":"native"}), + ); + request + .connection + .extra_headers + .push(("X-Trace".into(), "initial-only".into())); + + let result = perform_ocr(request).await.unwrap(); + server.await.unwrap(); + assert_eq!( + result.provider_native_response.as_ref(), + operation.as_object() + ); + let requests = seen.lock().unwrap(); + assert_eq!(requests.len(), 3); + assert!(requests[0].to_ascii_lowercase().contains("x-trace:")); + for poll in &requests[1..] { + assert!(!poll.to_ascii_lowercase().contains("x-trace:")); + assert!( + poll.to_ascii_lowercase() + .contains("ocp-apim-subscription-key: test-key") + ); + } + } + + struct SubmissionBoundary { + request_count: Arc>>, + post_calls: Arc>>, + } + + impl crate::ocr::hooks::OcrHooks for SubmissionBoundary { + fn post_call( + &self, + request: crate::ocr::hooks::OcrPostCallRequest, + ) -> crate::ocr::hooks::OcrHookFuture<'_, crate::ocr::hooks::OcrPostCallRequest> { + Box::pin(async move { + assert_eq!(self.request_count.lock().unwrap().len(), 1); + self.post_calls + .lock() + .unwrap() + .push(request.original_response.clone()); + Ok(request) + }) + } + } + + #[tokio::test] + async fn accepted_response_runs_post_call_once_before_polling() { + let (base, seen, server) = mock_server(vec![ + MockResponse { + status: 202, + headers: vec![("Operation-Location", "{base}/operation".into())], + body: json!({"submitted": true}), + }, + MockResponse::json(json!({"status":"succeeded"})), + ]) + .await; + let post_calls = Arc::new(Mutex::new(Vec::new())); + let request = crate::ocr::LiteLLMOcrRequest { + hooks: Arc::new(SubmissionBoundary { + request_count: seen.clone(), + post_calls: post_calls.clone(), + }), + ..wire_request("azure_ai/doc-intelligence/prebuilt-read", &base, json!({})) + }; + + perform_ocr(request).await.unwrap(); + server.await.unwrap(); + assert_eq!(seen.lock().unwrap().len(), 2); + assert_eq!( + *post_calls.lock().unwrap(), + [json!(r#"{"submitted":true}"#)] + ); + } + + #[tokio::test] + async fn polling_forwards_bearer_credentials() { + let (base, seen, server) = mock_server(vec![ + MockResponse { + status: 202, + headers: vec![("Operation-Location", "{base}/operation".into())], + body: json!({}), + }, + MockResponse::json(json!({"status":"succeeded"})), + ]) + .await; + let mut request = wire_request("azure_ai/doc-intelligence/prebuilt-read", &base, json!({})); + request.connection.api_key = None; + request.connection.extra_headers = vec![("Authorization".into(), "Bearer token".into())]; + + perform_ocr(request).await.unwrap(); + server.await.unwrap(); + let requests = seen.lock().unwrap(); + assert!( + requests[1] + .to_ascii_lowercase() + .contains("authorization: bearer token") + ); + } + + #[tokio::test] + async fn polling_does_not_follow_redirects() { + let (base, seen, server) = mock_server(vec![ + MockResponse { + status: 202, + headers: vec![("Operation-Location", "{base}/operation".into())], + body: json!({}), + }, + MockResponse { + status: 302, + headers: vec![("Location", "{base}/redirected".into())], + body: json!({}), + }, + MockResponse::json(json!({"status":"succeeded"})), + ]) + .await; + + let error = perform_ocr(wire_request( + "azure_ai/doc-intelligence/prebuilt-read", + &base, + json!({}), + )) + .await + .unwrap_err(); + + assert!(error.to_string().contains("status 302"), "{error}"); + assert_eq!(seen.lock().unwrap().len(), 2); + server.abort(); + } + + #[tokio::test] + async fn polling_rejects_terminal_failure() { + let (base, _, server) = mock_server(vec![ + MockResponse { + status: 202, + headers: vec![("Operation-Location", "{base}/operation".into())], + body: json!({}), + }, + MockResponse::json(json!({"status":"failed"})), + ]) + .await; + + let error = perform_ocr(wire_request( + "azure_ai/doc-intelligence/prebuilt-read", + &base, + json!({}), + )) + .await + .unwrap_err(); + server.await.unwrap(); + assert!(error.to_string().contains("status failed")); + } + + #[tokio::test] + async fn malformed_provider_pages_report_response_paths() { + for (analysis, path) in [ + (json!({"pages":null}), "pages"), + (json!({"pages":[null]}), "pages[0]"), + (json!({"pages":[{"lines":null}]}), "lines"), + (json!({"pages":[{"width":"bad"}]}), "width"), + ] { + let (base, _, server) = mock_server(vec![MockResponse::json(json!({ + "status":"succeeded", + "analyzeResult":analysis + }))]) + .await; + let error = perform_ocr(wire_request( + "azure_ai/doc-intelligence/prebuilt-read", + &base, + json!({}), + )) + .await + .unwrap_err(); + server.await.unwrap(); + assert!(error.to_string().contains(path), "{error}"); + } + } + + #[tokio::test] + async fn rejects_missing_invalid_and_cross_origin_operation_locations() { + for headers in [ + Vec::new(), + vec![("Operation-Location", "/relative".into())], + vec![("Operation-Location", "http://example.com/operation".into())], + vec![( + "Operation-Location", + "http://user:password@127.0.0.1/operation".into(), + )], + ] { + let (base, _, server) = mock_server(vec![MockResponse { + status: 202, + headers, + body: json!({}), + }]) + .await; + let error = perform_ocr(wire_request( + "azure_ai/doc-intelligence/prebuilt-read", + &base, + json!({}), + )) + .await + .unwrap_err(); + server.await.unwrap(); + assert!(error.to_string().contains("operation-location")); + } + } + + #[tokio::test] + async fn polling_deadline_bounds_retry_delay() { + let (base, _, server) = mock_server(vec![ + MockResponse { + status: 202, + headers: vec![("Operation-Location", "{base}/operation".into())], + body: json!({}), + }, + MockResponse { + status: 200, + headers: vec![("Retry-After", "9999".into())], + body: json!({"status":"notStarted"}), + }, + ]) + .await; + let mut request = wire_request("azure_ai/doc-intelligence/prebuilt-read", &base, json!({})); + request.connection.poll_timeout = std::time::Duration::from_millis(100); + + let error = tokio::time::timeout(std::time::Duration::from_secs(1), perform_ocr(request)) + .await + .unwrap() + .unwrap_err(); + server.await.unwrap(); + assert!(error.to_string().contains("timed out")); + } + + #[tokio::test] + async fn model_id_is_encoded_and_dot_segments_are_rejected() { + let (base, seen, server) = mock_server(vec![MockResponse::json(json!({ + "status":"succeeded" + }))]) + .await; + perform_ocr(wire_request( + "azure_ai/doc-intelligence/a ?#é", + &base, + json!({}), + )) + .await + .unwrap(); + server.await.unwrap(); + assert!(seen.lock().unwrap()[0].contains("a%20%3F%23%C3%A9:analyze")); + + for model in [ + "azure_ai/doc-intelligence/.", + "azure_ai/doc-intelligence/..", + ] { + let error = perform_ocr(wire_request(model, "http://127.0.0.1:1", json!({}))) + .await + .unwrap_err(); + assert!(error.to_string().contains("dot segment")); + } + } + + #[tokio::test] + async fn pre_call_guardrail_receives_caller_pages_before_mapping() { + use crate::ocr::hooks::{OcrHookFuture, OcrHooks, OcrPreCallRequest}; + use std::sync::Arc; + + struct RewritePages; + impl OcrHooks for RewritePages { + fn intercepts_requests(&self) -> bool { + true + } + + fn pre_call(&self, request: OcrPreCallRequest) -> OcrHookFuture<'_, OcrPreCallRequest> { + Box::pin(async move { + assert_eq!(request.optional_params["pages"], json!([0, 2])); + Ok(OcrPreCallRequest { + optional_params: json!({"pages": [1]}), + ..request + }) + }) + } + } + let (base, seen, server) = + mock_server(vec![MockResponse::json(json!({"status": "succeeded"}))]).await; + let request = wire_request( + "azure_ai/doc-intelligence/prebuilt-read", + &base, + json!({"pages": [0, 2]}), + ) + .with_host_hooks(Arc::new(RewritePages), None); + perform_ocr(request).await.unwrap(); + server.await.unwrap(); + let requests = seen.lock().unwrap(); + let target = requests[0].split_whitespace().nth(1).unwrap(); + assert_eq!( + query_value(&format!("{base}{target}"), "pages").as_deref(), + Some("2") + ); + assert_eq!(requests.len(), 1); + } } diff --git a/litellm-rust/crates/core/src/llms/azure_ai/ocr/transformation.rs b/litellm-rust/crates/core/src/llms/azure_ai/ocr/transformation.rs index 237b77598c3..ea1aca0b66b 100644 --- a/litellm-rust/crates/core/src/llms/azure_ai/ocr/transformation.rs +++ b/litellm-rust/crates/core/src/llms/azure_ai/ocr/transformation.rs @@ -1,6 +1,5 @@ use crate::constants::AZURE_AI_OCR_PATH; use crate::llms::base_llm::ocr::transformation::{BaseOcrConfig, OcrRequestContext}; -use crate::llms::mistral::ocr::MistralOcrResponse; use crate::llms::mistral::ocr::transformation::{MistralOCRConfig, MistralOcrRequest}; use crate::ocr::OcrClient; use crate::ocr::document::{inline_remote_document, validate_inline_document}; @@ -20,7 +19,46 @@ pub(crate) struct AzureAIOCRConfig; impl BaseOcrConfig for AzureAIOCRConfig { type OcrParams = OpaqueParams; type ProviderRequest = MistralOcrRequest; - type ProviderResponse = MistralOcrResponse; + type Environment = Vec<(String, String)>; + + fn get_api_key_env_var(&self) -> Option<&'static str> { + Some(AZURE_AI_API_KEY_ENV) + } + + async fn validate_environment( + &self, + request: &LiteLLMOcrRequest, + _client: &OcrClient, + ) -> Result { + let config = AzureAuthInputs { + azure_ad_token_provider: request.azure_ad_token_provider.clone(), + ..AzureAuthInputs::from_sourced_optional_params( + &request.optional_params, + &request.input_sources, + )? + }; + self.validate_environment(&request.connection, &config, &credential_env) + .await + } + + fn get_complete_url( + &self, + request: &LiteLLMOcrRequest, + _params: &Self::OcrParams, + _environment: &Self::Environment, + ) -> Result { + self.get_complete_url(request.connection.api_base.as_deref(), &credential_env) + } + + fn transform_ocr_request( + &self, + model: &str, + document: OcrDocument, + params: &OpaqueParams, + headers: &[(String, String)], + ) -> Result { + MistralOCRConfig.transform_ocr_request(model, document, params, headers) + } fn get_supported_ocr_params(&self, model: &str) -> &'static [&'static str] { MistralOCRConfig.get_supported_ocr_params(model) @@ -40,15 +78,16 @@ impl BaseOcrConfig for AzureAIOCRConfig { context.connection, ) .await?; - MistralOCRConfig.transform_ocr_request(model, document, optional_params, headers) + self.transform_ocr_request(model, document, optional_params, headers) } - fn normalize_response( + fn transform_ocr_response( &self, model: &str, - response: MistralOcrResponse, + raw_response: &[u8], + request_format: crate::ocr::types::OcrResponseFormat, ) -> Result { - MistralOCRConfig.normalize_response(model, response) + MistralOCRConfig.transform_ocr_response(model, raw_response, request_format) } } @@ -59,18 +98,8 @@ impl AzureAIOCRConfig { client: &OcrClient, ) -> Result { let params = self.map_ocr_params(&request.optional_params, &request.model)?; - let config = AzureAuthInputs { - azure_ad_token_provider: request.azure_ad_token_provider.clone(), - ..AzureAuthInputs::from_sourced_optional_params( - &request.optional_params, - &request.input_sources, - ) - .map_err(crate::ocr::Error::from)? - }; - let url = self.get_complete_url(request.connection.api_base.as_deref(), &credential_env)?; - let headers = self - .validate_environment(&request.connection, &config, &credential_env) - .await?; + let url = BaseOcrConfig::get_complete_url(self, request, ¶ms, &Vec::new())?; + let headers = BaseOcrConfig::validate_environment(self, request, client).await?; let retains_document = !request.document.source().starts_with("http://") && !request.document.source().starts_with("https://"); let body = self @@ -132,7 +161,7 @@ impl AzureAIOCRConfig { let key = nonblank(connection.api_key.clone()) .map(|value| Sourced::new(value, connection.api_key_source)) .or_else(|| { - nonblank(env_lookup(AZURE_AI_API_KEY_ENV)) + nonblank(self.get_api_key_env_var().and_then(env_lookup)) .map(|value| Sourced::new(value, InputSource::Environment)) }); if let Some(key) = key { @@ -259,4 +288,103 @@ mod tests { ("Authorization".into(), "Bearer request-key".into()) ); } + + use std::sync::Arc; + + use serde_json::{Value, json}; + + use crate::ocr::hooks::{OcrDuringCallRequest, OcrHookFuture, OcrHooks}; + use crate::ocr::test_support::{MockResponse, mock_server, perform_ocr, wire_request}; + + #[tokio::test] + async fn facade_executes_azure_mistral_with_prepared_auth() { + let (base, seen, server) = mock_server(vec![MockResponse::json(json!({ + "pages":[{"index":0,"markdown":"hello"}], + "usage_info":{"pages_processed":1} + }))]) + .await; + let mut request = wire_request( + "azure_ai/model", + &base, + json!({"include_image_base64":true}), + ); + request.connection.api_key = None; + request.connection.extra_headers = vec![( + "Authorization".into(), + "Bearer python-prepared-token".into(), + )]; + + let result = perform_ocr(request).await.unwrap(); + server.await.unwrap(); + assert_eq!(result.pages[0].markdown, "hello"); + let requests = seen.lock().unwrap(); + assert_eq!(requests.len(), 1); + assert!(requests[0].starts_with("POST /providers/mistral/azure/ocr ")); + assert!( + requests[0] + .to_ascii_lowercase() + .contains("authorization: bearer python-prepared-token\r\n") + ); + let body: Value = + serde_json::from_str(requests[0].split_once("\r\n\r\n").unwrap().1).unwrap(); + assert_eq!( + body, + json!({ + "model":"model", + "document":{"type":"document_url","document_url":"data:application/pdf;base64,YWJj"}, + "include_image_base64":true + }) + ); + } + + #[tokio::test] + async fn facade_acquires_supplied_entra_token_for_final_request() { + let (base, seen, server) = mock_server(vec![MockResponse::json(json!({"pages":[]}))]).await; + let mut request = wire_request( + "azure_ai/model", + &base, + json!({"azure_ad_token":"rust-owned-token"}), + ); + request.connection.api_key = None; + + perform_ocr(request).await.unwrap(); + server.await.unwrap(); + + let requests = seen.lock().unwrap(); + assert_eq!(requests.len(), 1); + assert!( + requests[0] + .to_ascii_lowercase() + .contains("authorization: bearer rust-owned-token\r\n") + ); + } + + struct ReplaceBodyDocument; + + impl OcrHooks for ReplaceBodyDocument { + fn intercepts_requests(&self) -> bool { + true + } + + fn during_call( + &self, + mut request: OcrDuringCallRequest, + ) -> OcrHookFuture<'_, OcrDuringCallRequest> { + Box::pin(async move { + request.body["document"] = json!({ + "type":"document_url", + "document_url":"https://example.com/not-inline.pdf" + }); + Ok(request) + }) + } + } + + #[tokio::test] + async fn rejects_non_inline_body_after_guardrails() { + let mut request = wire_request("azure_ai/model", "http://127.0.0.1:1", json!({})); + request.hooks = Arc::new(ReplaceBodyDocument); + let error = perform_ocr(request).await.unwrap_err(); + assert!(error.to_string().contains("data URI")); + } } diff --git a/litellm-rust/crates/core/src/llms/base_llm/ocr/transformation.rs b/litellm-rust/crates/core/src/llms/base_llm/ocr/transformation.rs index e7cc8eaa481..e8ccdda75ab 100644 --- a/litellm-rust/crates/core/src/llms/base_llm/ocr/transformation.rs +++ b/litellm-rust/crates/core/src/llms/base_llm/ocr/transformation.rs @@ -1,18 +1,64 @@ use std::future::Future; use std::sync::Arc; +use litellm_auth::Sourced; use serde::Serialize; use serde::de::DeserializeOwned; use crate::call_arguments::{CallArguments, parse_options}; use crate::ocr::OcrClient; use crate::ocr::hooks::OcrHooks; -use crate::ocr::types::{LiteLLMOcrResponse, OcrConnection, OcrDocument, OcrResponseFormat}; +use crate::ocr::types::{ + LiteLLMOcrRequest, LiteLLMOcrResponse, OcrConnection, OcrDocument, OcrResponseFormat, +}; + +const HEALTH_CHECK_PDF_DATA_URI: &str = "data:application/pdf;base64,JVBERi0xLjQKJeLjz9MKMyAwIG9iago8PC9UeXBlIC9QYWdlCi9QYXJlbnQgMSAwIFIKL01lZGlhQm94IFswIDAgNjEyIDc5Ml0KL0NvbnRlbnRzIDQgMCBSCi9SZXNvdXJjZXMgPDwvRm9udCA8PC9GMSAyIDAgUj4+Pj4+PgplbmRvYmoKNCAwIG9iago8PC9MZW5ndGggNDQ+PgpzdHJlYW0KQlQKL0YxIDI0IFRmCjEwMCA3MDAgVGQKKHRlc3QpIFRqCkVUCmVuZHN0cmVhbQplbmRvYmoKMiAwIG9iago8PC9UeXBlIC9Gb250Ci9TdWJ0eXBlIC9UeXBlMQovQmFzZUZvbnQgL0hlbHZldGljYT4+CmVuZG9iagoxIDAgb2JqCjw8L1R5cGUgL1BhZ2VzCi9LaWRzIFszIDAgUl0KL0NvdW50IDE+PgplbmRvYmoKNSAwIG9iago8PC9UeXBlIC9DYXRhbG9nCi9QYWdlcyAxIDAgUj4+CmVuZG9iagp0cmFpbGVyCjw8L1NpemUgNgovUm9vdCA1IDAgUj4+CnN0YXJ0eHJlZgozMjQKJSVFT0Y="; pub(crate) trait BaseOcrConfig: Send + Sync + Sized + 'static { type OcrParams: DeserializeOwned + Send + Sync; type ProviderRequest: Serialize + Send; - type ProviderResponse: DeserializeOwned + Send; + type Environment: Send + Sync; + + fn get_api_key_env_var(&self) -> Option<&'static str> { + None + } + + fn resolve_connection_params( + &self, + api_key: Option>, + api_base: Option>, + dynamic_api_key: Option>, + dynamic_api_base: Option>, + ) -> (Option>, Option>) { + ( + dynamic_api_key + .filter(|value| !value.value().is_empty()) + .or(api_key), + dynamic_api_base + .filter(|value| !value.value().is_empty()) + .or(api_base), + ) + } + + fn get_health_check_document(&self) -> OcrDocument { + OcrDocument::DocumentUrl { + document_url: HEALTH_CHECK_PDF_DATA_URI.into(), + extra_fields: Default::default(), + } + } + + fn validate_environment( + &self, + request: &LiteLLMOcrRequest, + client: &OcrClient, + ) -> impl Future> + Send; + + fn get_complete_url( + &self, + request: &LiteLLMOcrRequest, + optional_params: &Self::OcrParams, + environment: &Self::Environment, + ) -> Result; fn get_supported_ocr_params(&self, _model: &str) -> &'static [&'static str] { &[] @@ -30,36 +76,31 @@ pub(crate) trait BaseOcrConfig: Send + Sync + Sized + 'static { )?) } + fn transform_ocr_request( + &self, + model: &str, + document: OcrDocument, + optional_params: &Self::OcrParams, + headers: &[(String, String)], + ) -> Result; + fn async_transform_ocr_request( &self, model: &str, document: OcrDocument, optional_params: &Self::OcrParams, headers: &[(String, String)], - context: OcrRequestContext<'_>, - ) -> impl Future> + Send; + _context: OcrRequestContext<'_>, + ) -> impl Future> + Send { + async move { self.transform_ocr_request(model, document, optional_params, headers) } + } - fn normalize_response( - &self, - model: &str, - response: Self::ProviderResponse, - ) -> Result; - - fn decode_and_normalize_response( + fn transform_ocr_response( &self, model: &str, raw_response: &[u8], request_format: OcrResponseFormat, - ) -> Result { - let decoded = crate::ocr::wire::decode_response::( - raw_response, - request_format == OcrResponseFormat::Native, - )?; - Ok(LiteLLMOcrResponse { - provider_native_response: decoded.native, - ..self.normalize_response(model, decoded.data)? - }) - } + ) -> Result; fn async_transform_ocr_response( &self, @@ -74,9 +115,38 @@ pub(crate) trait BaseOcrConfig: Send + Sync + Sized + 'static { ) .await?; crate::ocr::handler::post_call(context.hooks, &bytes).await?; - self.decode_and_normalize_response(model, &bytes, context.request_format) + self.transform_ocr_response(model, &bytes, context.request_format) } } + + fn get_error_class( + &self, + error_message: String, + status_code: u16, + headers: Vec<(String, String)>, + ) -> crate::ocr::Error { + crate::ocr::Error::Provider { + status: status_code, + body: error_message, + headers, + } + } +} + +pub(crate) fn decode_and_normalize_response( + model: &str, + raw_response: &[u8], + request_format: OcrResponseFormat, + normalize: impl FnOnce(&str, T) -> Result, +) -> Result { + let decoded = crate::ocr::wire::decode_response( + raw_response, + request_format == OcrResponseFormat::Native, + )?; + Ok(LiteLLMOcrResponse { + provider_native_response: decoded.native, + ..normalize(model, decoded.data)? + }) } #[derive(Clone, Copy)] diff --git a/litellm-rust/crates/core/src/llms/cohere/ocr/mod.rs b/litellm-rust/crates/core/src/llms/cohere/ocr/mod.rs index 1c04ff00a67..9cbe4df56e5 100644 --- a/litellm-rust/crates/core/src/llms/cohere/ocr/mod.rs +++ b/litellm-rust/crates/core/src/llms/cohere/ocr/mod.rs @@ -1,3 +1,3 @@ pub(crate) mod transformation; -pub(crate) use transformation::{CohereOptions, CohereResponse, validate_document}; +pub(crate) use transformation::{CohereOptions, validate_document}; diff --git a/litellm-rust/crates/core/src/llms/cohere/ocr/transformation.rs b/litellm-rust/crates/core/src/llms/cohere/ocr/transformation.rs index 27b24b8b6db..d0b92127f81 100644 --- a/litellm-rust/crates/core/src/llms/cohere/ocr/transformation.rs +++ b/litellm-rust/crates/core/src/llms/cohere/ocr/transformation.rs @@ -1,8 +1,8 @@ -use crate::serde_compat::LaxI64; use serde::{Deserialize, Serialize}; use serde_json::{Map, Value}; use serde_with::serde_as; +use crate::serde_compat::LaxI64; use crate::constants::{COHERE_API_KEY_ENV, COHERE_PARSE_API_BASE}; use crate::llms::base_llm::ocr::transformation::{BaseOcrConfig, OcrRequestContext}; use crate::ocr::OcrClient; @@ -14,6 +14,8 @@ use crate::ocr::types::{ }; use crate::url_utils::ApiUrl; +const COHERE_PARSE_HEALTH_CHECK_IMAGE_DATA_URI: &str = "data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAIAAACQd1PeAAAADElEQVR4nGP4//8/AAX+Av4N70a4AAAAAElFTkSuQmCC"; + #[derive(Clone, Copy, Debug, Default, Deserialize, Serialize)] #[serde(rename_all = "lowercase")] pub(crate) enum OutputFormat { @@ -80,8 +82,46 @@ struct CohereBilledUnits { #[derive(Default)] pub(crate) struct CohereParseConfig; -impl CohereParseConfig { - pub(crate) fn transform_ocr_request( +impl BaseOcrConfig for CohereParseConfig { + type OcrParams = CohereOptions; + type ProviderRequest = CohereRequest; + type Environment = Vec<(String, String)>; + + fn get_api_key_env_var(&self) -> Option<&'static str> { + Some(COHERE_API_KEY_ENV) + } + + fn get_health_check_document(&self) -> OcrDocument { + OcrDocument::ImageUrl { + image_url: COHERE_PARSE_HEALTH_CHECK_IMAGE_DATA_URI.into(), + extra_fields: Default::default(), + } + } + + async fn validate_environment( + &self, + request: &LiteLLMOcrRequest, + _client: &OcrClient, + ) -> Result { + self.validate_environment(&request.connection, &credential_env) + } + + fn get_complete_url( + &self, + request: &LiteLLMOcrRequest, + _params: &Self::OcrParams, + _environment: &Self::Environment, + ) -> Result { + self.get_complete_url( + request + .connection + .api_base + .as_deref() + .unwrap_or(COHERE_PARSE_API_BASE), + ) + } + + fn transform_ocr_request( &self, model: &str, document: OcrDocument, @@ -91,12 +131,6 @@ impl CohereParseConfig { let image_url = image_url(document)?; Ok(build_request(model, image_url, optional_params)) } -} - -impl BaseOcrConfig for CohereParseConfig { - type OcrParams = CohereOptions; - type ProviderRequest = CohereRequest; - type ProviderResponse = CohereResponse; fn get_supported_ocr_params(&self, _model: &str) -> &'static [&'static str] { &["output_format", "req_format"] @@ -113,12 +147,18 @@ impl BaseOcrConfig for CohereParseConfig { self.transform_ocr_request(model, document, optional_params, headers) } - fn normalize_response( + fn transform_ocr_response( &self, model: &str, - response: CohereResponse, + raw_response: &[u8], + request_format: crate::ocr::types::OcrResponseFormat, ) -> Result { - normalize_response(model, response) + crate::llms::base_llm::ocr::transformation::decode_and_normalize_response( + model, + raw_response, + request_format, + normalize_response, + ) } } @@ -129,14 +169,8 @@ impl CohereParseConfig { client: &OcrClient, ) -> Result { let params = self.map_ocr_params(&request.optional_params, &request.model)?; - let headers = self.validate_environment(&request.connection, &credential_env)?; - let url = self.get_complete_url( - request - .connection - .api_base - .as_deref() - .unwrap_or(COHERE_PARSE_API_BASE), - )?; + let headers = BaseOcrConfig::validate_environment(self, request, client).await?; + let url = BaseOcrConfig::get_complete_url(self, request, ¶ms, &headers)?; let body = self .async_transform_ocr_request( &request.model, @@ -299,7 +333,11 @@ impl CohereParseConfig { .map(str::trim) .filter(|key| !key.is_empty()) .map(str::to_string) - .or_else(|| env_lookup(COHERE_API_KEY_ENV).filter(|key| !key.trim().is_empty())) + .or_else(|| { + self.get_api_key_env_var() + .and_then(env_lookup) + .filter(|key| !key.trim().is_empty()) + }) .ok_or_else(|| { crate::ocr::Error::Auth(litellm_auth::Error::ProviderAuthentication( "Missing COHERE_API_KEY - set it in the environment or pass api_key".into(), diff --git a/litellm-rust/crates/core/src/llms/mistral/ocr/mod.rs b/litellm-rust/crates/core/src/llms/mistral/ocr/mod.rs index 086afd5737f..080f0a1183f 100644 --- a/litellm-rust/crates/core/src/llms/mistral/ocr/mod.rs +++ b/litellm-rust/crates/core/src/llms/mistral/ocr/mod.rs @@ -1,3 +1 @@ pub(crate) mod transformation; - -pub(crate) use transformation::MistralOcrResponse; diff --git a/litellm-rust/crates/core/src/llms/mistral/ocr/transformation.rs b/litellm-rust/crates/core/src/llms/mistral/ocr/transformation.rs index 02640af480c..5401c819bc0 100644 --- a/litellm-rust/crates/core/src/llms/mistral/ocr/transformation.rs +++ b/litellm-rust/crates/core/src/llms/mistral/ocr/transformation.rs @@ -37,9 +37,34 @@ pub(crate) struct MistralOcrResponse { #[derive(Clone, Debug, Default)] pub(crate) struct MistralOCRConfig; -impl MistralOCRConfig { +impl BaseOcrConfig for MistralOCRConfig { + type OcrParams = OpaqueParams; + type ProviderRequest = MistralOcrRequest; + type Environment = Vec<(String, String)>; + + fn get_api_key_env_var(&self) -> Option<&'static str> { + Some(MISTRAL_API_KEY_ENV) + } + + async fn validate_environment( + &self, + request: &LiteLLMOcrRequest, + _client: &OcrClient, + ) -> Result { + self.validate_environment(&request.connection, &credential_env) + } + + fn get_complete_url( + &self, + request: &LiteLLMOcrRequest, + _params: &Self::OcrParams, + _environment: &Self::Environment, + ) -> Result { + self.get_complete_url(request.connection.api_base.as_deref()) + } + #[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)] - pub(crate) fn transform_ocr_request( + fn transform_ocr_request( &self, model: &str, document: OcrDocument, @@ -52,12 +77,6 @@ impl MistralOCRConfig { params: optional_params.clone(), }) } -} - -impl BaseOcrConfig for MistralOCRConfig { - type OcrParams = OpaqueParams; - type ProviderRequest = MistralOcrRequest; - type ProviderResponse = MistralOcrResponse; fn get_supported_ocr_params(&self, _model: &str) -> &'static [&'static str] { &[ @@ -88,12 +107,18 @@ impl BaseOcrConfig for MistralOCRConfig { self.transform_ocr_request(model, document, optional_params, headers) } - fn normalize_response( + fn transform_ocr_response( &self, model: &str, - response: MistralOcrResponse, + raw_response: &[u8], + request_format: crate::ocr::types::OcrResponseFormat, ) -> Result { - normalize_response(model, response) + crate::llms::base_llm::ocr::transformation::decode_and_normalize_response( + model, + raw_response, + request_format, + normalize_response, + ) } } @@ -104,8 +129,8 @@ impl MistralOCRConfig { client: &OcrClient, ) -> Result { let params = self.map_ocr_params(&request.optional_params, &request.model)?; - let headers = self.validate_environment(&request.connection, &credential_env)?; - let url = self.get_complete_url(request.connection.api_base.as_deref())?; + let headers = BaseOcrConfig::validate_environment(self, request, client).await?; + let url = BaseOcrConfig::get_complete_url(self, request, ¶ms, &headers)?; let body = self .async_transform_ocr_request( &request.model, @@ -170,7 +195,11 @@ impl MistralOCRConfig { .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())) + .or_else(|| { + self.get_api_key_env_var() + .and_then(env_lookup) + .filter(|key| !key.trim().is_empty()) + }) .ok_or(litellm_auth::Error::MissingApiKey { provider: "Mistral", environment_variable: MISTRAL_API_KEY_ENV, @@ -292,11 +321,7 @@ mod tests { fn raw_response_transform_keeps_native_payload_separate_from_typed_normalization() { let raw = br#"{"pages":[{"index":"2","markdown":"text"}],"provider_extension":false}"#; let response = MistralOCRConfig - .decode_and_normalize_response( - "model", - raw, - crate::ocr::types::OcrResponseFormat::Native, - ) + .transform_ocr_response("model", raw, crate::ocr::types::OcrResponseFormat::Native) .unwrap(); assert_eq!(response.pages[0].index, 2); let native = response.provider_native_response.unwrap(); @@ -305,7 +330,7 @@ mod tests { assert!(response.extra_fields.is_empty()); assert!( MistralOCRConfig - .decode_and_normalize_response( + .transform_ocr_response( "model", br#"{"pages":[{"index":0}]}"#, crate::ocr::types::OcrResponseFormat::Litellm diff --git a/litellm-rust/crates/core/src/llms/reducto/ocr/transformation.rs b/litellm-rust/crates/core/src/llms/reducto/ocr/transformation.rs index 92c50d06916..aa6ce333757 100644 --- a/litellm-rust/crates/core/src/llms/reducto/ocr/transformation.rs +++ b/litellm-rust/crates/core/src/llms/reducto/ocr/transformation.rs @@ -81,7 +81,37 @@ pub(crate) struct ReductoParseV3Config; impl BaseOcrConfig for ReductoParseV3Config { type OcrParams = ReductoV3Params; type ProviderRequest = ReductoV3Request; - type ProviderResponse = ReductoResponse; + type Environment = Vec<(String, String)>; + + async fn validate_environment( + &self, + request: &LiteLLMOcrRequest, + _client: &OcrClient, + ) -> Result { + validate_environment(&request.connection, &credential_env) + } + + fn get_complete_url( + &self, + request: &LiteLLMOcrRequest, + _params: &Self::OcrParams, + _environment: &Self::Environment, + ) -> Result { + get_complete_url(request.connection.api_base.as_deref()) + } + + fn transform_ocr_request( + &self, + _model: &str, + document: OcrDocument, + params: &Self::OcrParams, + _headers: &[(String, String)], + ) -> Result { + Ok(ReductoV3Request { + input: uploaded_file_id(document)?, + params: params.clone(), + }) + } fn get_supported_ocr_params(&self, _model: &str) -> &'static [&'static str] { &["formatting", "retrieval", "settings"] @@ -108,12 +138,18 @@ impl BaseOcrConfig for ReductoParseV3Config { }) } - fn normalize_response( + fn transform_ocr_response( &self, model: &str, - response: ReductoResponse, + raw_response: &[u8], + request_format: crate::ocr::types::OcrResponseFormat, ) -> Result { - normalize_response(model, response) + crate::llms::base_llm::ocr::transformation::decode_and_normalize_response( + model, + raw_response, + request_format, + normalize_response, + ) } } @@ -124,8 +160,8 @@ impl ReductoParseV3Config { client: &OcrClient, ) -> Result { let params = self.map_ocr_params(&request.optional_params, &request.model)?; - let headers = validate_environment(&request.connection, &credential_env)?; - let url = get_complete_url(request.connection.api_base.as_deref())?; + let headers = self.validate_environment(request, client).await?; + let url = self.get_complete_url(request, ¶ms, &headers)?; let (document, headers) = guardrail_document(request, &url, &headers).await?; let body = self .async_transform_ocr_request( @@ -154,7 +190,36 @@ pub(crate) struct ReductoParseLegacyConfig; impl BaseOcrConfig for ReductoParseLegacyConfig { type OcrParams = ReductoLegacyParams; type ProviderRequest = ReductoLegacyRequest; - type ProviderResponse = ReductoResponse; + type Environment = Vec<(String, String)>; + + async fn validate_environment( + &self, + request: &LiteLLMOcrRequest, + client: &OcrClient, + ) -> Result { + ReductoParseV3Config + .validate_environment(request, client) + .await + } + + fn get_complete_url( + &self, + request: &LiteLLMOcrRequest, + params: &Self::OcrParams, + environment: &Self::Environment, + ) -> Result { + ReductoParseV3Config.get_complete_url(request, params, environment) + } + + fn transform_ocr_request( + &self, + _model: &str, + document: OcrDocument, + params: &Self::OcrParams, + _headers: &[(String, String)], + ) -> Result { + Ok(build_legacy_body(uploaded_file_id(document)?, params)) + } fn get_supported_ocr_params(&self, _model: &str) -> &'static [&'static str] { &["enhance"] @@ -178,12 +243,13 @@ impl BaseOcrConfig for ReductoParseLegacyConfig { Ok(build_legacy_body(file_id, optional_params)) } - fn normalize_response( + fn transform_ocr_response( &self, model: &str, - response: ReductoResponse, + raw_response: &[u8], + request_format: crate::ocr::types::OcrResponseFormat, ) -> Result { - normalize_response(model, response) + ReductoParseV3Config.transform_ocr_response(model, raw_response, request_format) } } @@ -194,8 +260,8 @@ impl ReductoParseLegacyConfig { client: &OcrClient, ) -> Result { let params = self.map_ocr_params(&request.optional_params, &request.model)?; - let headers = validate_environment(&request.connection, &credential_env)?; - let url = get_complete_url(request.connection.api_base.as_deref())?; + let headers = self.validate_environment(request, client).await?; + let url = self.get_complete_url(request, ¶ms, &headers)?; let (document, headers) = guardrail_document(request, &url, &headers).await?; let body = self .async_transform_ocr_request( @@ -218,6 +284,21 @@ impl ReductoParseLegacyConfig { } } +fn uploaded_file_id(document: OcrDocument) -> Result { + if !document.source().starts_with(REDUCTO_ID_PREFIX) { + return Err(crate::ocr::Error::ReductoSource); + } + if document.source()[REDUCTO_ID_PREFIX.len()..] + .trim() + .is_empty() + { + return Err(crate::ocr::Error::RequestField { + path: "document file id".into(), + }); + } + Ok(ReductoFileId(document.source().into())) +} + fn present_nullable<'de, D: Deserializer<'de>, T: Deserialize<'de>>( deserializer: D, ) -> Result>, D::Error> { @@ -605,4 +686,335 @@ mod tests { connection.extra_headers ); } + + use std::sync::Arc; + + use rstest::rstest; + + use crate::ocr::hooks::{OcrDuringCallRequest, OcrHookFuture, OcrHooks, OcrPostCallRequest}; + use crate::ocr::test_support::{MockResponse, mock_server, perform_ocr, wire_request}; + + fn request_body(request: &str) -> Value { + serde_json::from_str(request.split_once("\r\n\r\n").unwrap().1).unwrap() + } + + #[rstest] + #[case( + "reducto/parse-v3", + json!({ + "formatting":{"table_output_format":"html"}, + "retrieval":{"chunk_mode":"section"}, + "settings":{"ocr_system":"standard"}, + "future_ocr_option":true, + "extra_body":{"provider_option":"value"} + }), + "reducto://already.pdf", + json!({ + "input":"reducto://already.pdf", + "formatting":{"table_output_format":"html"}, + "retrieval":{"chunk_mode":"section"}, + "settings":{"ocr_system":"standard"}, + "future_ocr_option":true, + "provider_option":"value" + }) + )] + #[case( + "reducto/parse-legacy", + json!({ + "enhance":{"agentic":[{"type":"table"}]}, + "future_ocr_option":true, + "extra_body":{"provider_option":"value"} + }), + "reducto://legacy.pdf", + json!({ + "document_url":"reducto://legacy.pdf", + "options":{"enhance":{"agentic":[{"type":"table"}]}}, + "future_ocr_option":true, + "provider_option":"value" + }) + )] + #[tokio::test] + async fn request_mapping_matches_python( + #[case] model: &str, + #[case] options: Value, + #[case] source: &str, + #[case] expected: Value, + ) { + let (base, seen, server) = mock_server(vec![MockResponse::json(json!({ + "result":{"chunks":[]} + }))]) + .await; + let mut request = wire_request(model, &base, options); + request.document = request.document.with_source(source.into()); + + perform_ocr(request).await.unwrap(); + server.await.unwrap(); + let requests = seen.lock().unwrap(); + assert_eq!(requests.len(), 1); + assert!(requests[0].starts_with("POST /parse ")); + assert_eq!(request_body(&requests[0]), expected); + } + + #[rstest] + #[case("parse-v3")] + #[case("parse-legacy")] + #[tokio::test] + async fn data_uri_upload_preserves_multipart_headers(#[case] model: &str) { + let (base, seen, server) = mock_server(vec![ + MockResponse::json(json!({"file_id":"reducto://uploaded.pdf"})), + MockResponse::json(json!({"result":{"chunks":[{"content":"hello"}]}})), + ]) + .await; + let mut request = wire_request(&format!("reducto/{model}"), &base, json!({})); + request.connection.extra_headers = vec![ + ("Content-Type".into(), "application/json".into()), + ("X-Trace".into(), "upload-test".into()), + ]; + + let response = perform_ocr(request).await.unwrap(); + server.await.unwrap(); + assert_eq!(response.pages[0].markdown, "hello"); + let requests = seen.lock().unwrap(); + assert_eq!(requests.len(), 2); + assert!(requests[0].starts_with("POST /upload ")); + assert!( + requests[0] + .to_ascii_lowercase() + .contains("content-type: multipart/form-data; boundary=") + ); + assert!(requests[0].contains("x-trace: upload-test")); + assert!(requests[0].contains("application/pdf")); + assert!(requests[0].contains("abc")); + assert!(requests[1].starts_with("POST /parse ")); + } + + struct ParseBoundary { + request_count: Arc>>, + } + + impl OcrHooks for ParseBoundary { + fn post_call(&self, request: OcrPostCallRequest) -> OcrHookFuture<'_, OcrPostCallRequest> { + Box::pin(async move { + assert_eq!(self.request_count.lock().unwrap().len(), 2); + assert_eq!( + request.original_response, + json!(r#"{"result":{"chunks":[]}}"#) + ); + Ok(request) + }) + } + } + + #[tokio::test] + async fn post_call_stays_after_reducto_upload_and_parse() { + let (base, seen, server) = mock_server(vec![ + MockResponse::json(json!({"file_id":"reducto://uploaded.pdf"})), + MockResponse::json(json!({"result":{"chunks":[]}})), + ]) + .await; + let request = crate::ocr::LiteLLMOcrRequest { + hooks: Arc::new(ParseBoundary { + request_count: seen.clone(), + }), + ..wire_request("reducto/parse-v3", &base, json!({})) + }; + + perform_ocr(request).await.unwrap(); + server.await.unwrap(); + assert_eq!(seen.lock().unwrap().len(), 2); + } + + #[rstest] + #[case(json!({"file_id":""}))] + #[case(json!({}))] + #[case(json!({"file_id":null}))] + #[tokio::test] + async fn invalid_upload_ids_stop_before_parse(#[case] response: Value) { + let (base, seen, server) = mock_server(vec![MockResponse::json(response)]).await; + let error = perform_ocr(wire_request("reducto/parse-v3", &base, json!({}))) + .await + .unwrap_err(); + server.await.unwrap(); + assert!(error.to_string().contains("file_id")); + assert_eq!(seen.lock().unwrap().len(), 1); + } + + #[tokio::test] + async fn upload_failure_stops_before_parse() { + let (base, seen, server) = mock_server(vec![MockResponse { + status: 503, + headers: vec![], + body: json!({"error":"unavailable"}), + }]) + .await; + assert!( + perform_ocr(wire_request("reducto/parse-v3", &base, json!({}))) + .await + .is_err() + ); + server.await.unwrap(); + assert_eq!(seen.lock().unwrap().len(), 1); + } + + #[rstest] + #[case("https://example.com/a.pdf")] + #[case("reducto://")] + #[case("data:application/pdf;base64")] + #[case("data:application/pdf;base64,INVALID!")] + #[tokio::test] + async fn rejects_invalid_document_sources_before_network(#[case] source: &str) { + let mut request = wire_request("reducto/parse-v3", "http://127.0.0.1:1", json!({})); + request.document = request.document.with_source(source.into()); + assert!(perform_ocr(request).await.is_err()); + } + + #[test] + fn response_normalization_groups_blocks_and_distinguishes_null_result() { + use crate::llms::reducto::ocr::transformation::{ReductoResponse, normalize_response}; + + let raw = json!({"usage":{"num_pages":"2","credits":"3"},"result":{"type":"full","chunks":[ + {"blocks":[{ + "type":"Table", + "content":"B", + "bbox":{"left":0.1,"top":0.2,"width":0.8,"height":0.3,"page":2,"original_page":4}, + "confidence":"high", + "granular_confidence":{"parse_confidence":0.95,"extract_confidence":null}, + "image_url":null + }]}, + {"blocks":[{"content":"A","bbox":{"page":1},"type":"Text"},{"content":"C","bbox":{"page":1}}]} + ]}}); + let response: ReductoResponse = serde_json::from_value(raw).unwrap(); + let normalized = normalize_response("parse-v3", response) + .unwrap() + .into_json(); + assert_eq!(normalized["pages"][0]["markdown"], "A\n\nC"); + assert_eq!(normalized["pages"][1]["markdown"], "B"); + assert_eq!(normalized["pages"][1]["blocks"][0]["type"], "Table"); + assert_eq!( + normalized["pages"][1]["blocks"][0]["bbox"], + json!({"left":0.1,"top":0.2,"width":0.8,"height":0.3,"page":2,"original_page":4}) + ); + assert_eq!(normalized["pages"][1]["blocks"][0]["confidence"], "high"); + assert_eq!( + normalized["pages"][1]["blocks"][0]["granular_confidence"]["parse_confidence"], + 0.95 + ); + assert!(normalized["pages"][1]["blocks"][0]["image_url"].is_null()); + assert_eq!(normalized["usage_info"]["pages_processed"], 2); + assert_eq!(normalized["usage_info"]["credits"], 3.0); + + let missing: ReductoResponse = + serde_json::from_value(json!({"chunks":[{"content":"text"}]})).unwrap(); + let missing = normalize_response("parse-v3", missing).unwrap(); + assert_eq!(missing.pages[0].markdown, "text"); + let null: ReductoResponse = serde_json::from_value( + json!({"result":null,"chunks":[{"content":"ignored"}],"usage":null}), + ) + .unwrap(); + let null = normalize_response("parse-v3", null).unwrap(); + assert!(null.pages.is_empty()); + } + + #[tokio::test] + async fn facade_omits_native_response_by_default_and_preserves_auth_priority() { + let raw = json!({"job_id":"job-1","result":{"chunks":[]}}); + let (base, seen, server) = mock_server(vec![MockResponse::json(raw)]).await; + let mut request = wire_request("reducto/parse-v3", &base, json!({})); + request.document = request.document.with_source("reducto://ready.pdf".into()); + request.connection.extra_headers = vec![("authorization".into(), "Bearer existing".into())]; + + let response = perform_ocr(request).await.unwrap(); + server.await.unwrap(); + assert_eq!(response.provider_native_response, None); + assert!( + seen.lock().unwrap()[0] + .to_ascii_lowercase() + .contains("authorization: bearer existing") + ); + } + + struct RewriteDocument; + + struct RewriteHeaders; + + impl OcrHooks for RewriteHeaders { + fn intercepts_requests(&self) -> bool { + true + } + + fn during_call( + &self, + request: OcrDuringCallRequest, + ) -> OcrHookFuture<'_, OcrDuringCallRequest> { + Box::pin(async move { + Ok(OcrDuringCallRequest { + headers: vec![("authorization".into(), "Bearer guarded".into())], + ..request + }) + }) + } + } + + #[rstest] + #[case("reducto/parse-v3")] + #[case("reducto/parse-legacy")] + #[tokio::test] + async fn guardrail_headers_reach_upload_and_parse(#[case] model: &str) { + let (base, seen, server) = mock_server(vec![ + MockResponse::json(json!({"file_id":"reducto://uploaded.pdf"})), + MockResponse::json(json!({"result":{"chunks":[]}})), + ]) + .await; + let mut request = wire_request(model, &base, json!({})); + request.connection.extra_headers = vec![("authorization".into(), "Bearer original".into())]; + request.hooks = Arc::new(RewriteHeaders); + + perform_ocr(request).await.unwrap(); + server.await.unwrap(); + let requests = seen.lock().unwrap(); + assert_eq!(requests.len(), 2); + assert!(requests[0].starts_with("POST /upload ")); + assert!(requests[1].starts_with("POST /parse ")); + for request in requests.iter() { + assert!(request.contains("authorization: Bearer guarded")); + assert!(!request.contains("Bearer original")); + } + } + + impl OcrHooks for RewriteDocument { + fn intercepts_requests(&self) -> bool { + true + } + + fn during_call( + &self, + request: OcrDuringCallRequest, + ) -> OcrHookFuture<'_, OcrDuringCallRequest> { + Box::pin(async move { + assert_eq!( + request.body["document_url"], + "data:application/pdf;base64,YWJj" + ); + Ok(OcrDuringCallRequest { + body: json!({"type":"document_url","document_url":"reducto://guarded.pdf"}), + ..request + }) + }) + } + } + + #[tokio::test] + async fn guardrail_rewrites_document_before_upload() { + let (base, seen, server) = + mock_server(vec![MockResponse::json(json!({"result":{"chunks":[]}}))]).await; + let mut request = wire_request("reducto/parse-v3", &base, json!({})); + request.hooks = Arc::new(RewriteDocument); + + perform_ocr(request).await.unwrap(); + server.await.unwrap(); + let requests = seen.lock().unwrap(); + assert_eq!(requests.len(), 1); + assert!(requests[0].starts_with("POST /parse ")); + assert!(requests[0].contains("reducto://guarded.pdf")); + } } diff --git a/litellm-rust/crates/core/src/llms/vertex_ai/ocr/deepseek_transformation.rs b/litellm-rust/crates/core/src/llms/vertex_ai/ocr/deepseek_transformation.rs index 21e69e89210..4c12c02cb02 100644 --- a/litellm-rust/crates/core/src/llms/vertex_ai/ocr/deepseek_transformation.rs +++ b/litellm-rust/crates/core/src/llms/vertex_ai/ocr/deepseek_transformation.rs @@ -100,7 +100,38 @@ pub(crate) struct VertexAIDeepSeekOCRConfig; impl BaseOcrConfig for VertexAIDeepSeekOCRConfig { type OcrParams = DeepSeekOcrParams; type ProviderRequest = DeepSeekOcrRequest; - type ProviderResponse = DeepSeekOcrResponse; + type Environment = vertex::VertexEnvironment; + + fn get_api_key_env_var(&self) -> Option<&'static str> { + VertexAIOCRConfig.get_api_key_env_var() + } + + async fn validate_environment( + &self, + request: &LiteLLMOcrRequest, + client: &OcrClient, + ) -> Result { + BaseOcrConfig::validate_environment(&VertexAIOCRConfig, request, client).await + } + + fn get_complete_url( + &self, + request: &LiteLLMOcrRequest, + _params: &Self::OcrParams, + environment: &Self::Environment, + ) -> Result { + let config = VertexConfig::from_sourced_optional_params( + &request.optional_params, + &request.input_sources, + )?; + let location = vertex::get_vertex_ai_location(&config, &credential_env) + .unwrap_or_else(|| DEFAULT_LOCATION.to_string()); + self.get_complete_url( + request.connection.api_base.as_deref(), + &environment.project_id, + &location, + ) + } async fn async_transform_ocr_request( &self, @@ -113,17 +144,21 @@ impl BaseOcrConfig for VertexAIDeepSeekOCRConfig { self.transform_ocr_request(model, document, optional_params, headers) } - fn normalize_response( + fn transform_ocr_response( &self, model: &str, - response: DeepSeekOcrResponse, + raw_response: &[u8], + request_format: crate::ocr::types::OcrResponseFormat, ) -> Result { - normalize_response(model, response) + crate::llms::base_llm::ocr::transformation::decode_and_normalize_response( + model, + raw_response, + request_format, + normalize_response, + ) } -} -impl VertexAIDeepSeekOCRConfig { - pub(crate) fn transform_ocr_request( + fn transform_ocr_request( &self, model: &str, document: OcrDocument, @@ -148,28 +183,18 @@ impl VertexAIDeepSeekOCRConfig { .collect(), }) } +} +impl VertexAIDeepSeekOCRConfig { pub(crate) async fn prepare_request( &self, request: &LiteLLMOcrRequest, client: &OcrClient, ) -> Result { let params = self.map_ocr_params(&request.optional_params, &request.model)?; - let config = VertexConfig::from_sourced_optional_params( - &request.optional_params, - &request.input_sources, - ) - .map_err(crate::ocr::Error::from)?; - let authentication = VertexAIOCRConfig - .validate_environment(&request.connection, &config, client) - .await?; - let location = vertex::get_vertex_ai_location(&config, &credential_env) - .unwrap_or_else(|| DEFAULT_LOCATION.to_string()); - let url = self.get_complete_url( - request.connection.api_base.as_deref(), - &authentication.project_id, - &location, - )?; + let authentication = self.validate_environment(request, client).await?; + let url = BaseOcrConfig::get_complete_url(self, request, ¶ms, &authentication)?; + let body = self .async_transform_ocr_request( &request.model, @@ -393,7 +418,11 @@ impl VertexAIDeepSeekOCRConfig { #[cfg(test)] mod tests { - use super::{VertexAIDeepSeekOCRConfig, provider_model}; + use super::{ + DeepSeekOcrParams, DeepSeekOcrResponse, VertexAIDeepSeekOCRConfig, normalize_response, + provider_model, + }; + use serde_json::{Value, json}; #[test] fn unconsumed_options_remain_available_for_body_composition() { @@ -437,4 +466,247 @@ mod tests { "https://aiplatform.googleapis.com/v1/projects/proj-1/locations/europe-west4/endpoints/openapi/chat/completions" ); } + + use rstest::rstest; + + use crate::llms::base_llm::ocr::transformation::BaseOcrConfig; + 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"]))] + #[case("temperature", json!(null))] + fn request_mapping_matches_python(#[case] name: &str, #[case] value: Value) { + let params: DeepSeekOcrParams = + serde_json::from_value(json!({name: value.clone(), "ignored": true})).unwrap(); + let result = serde_json::to_value( + VertexAIDeepSeekOCRConfig + .transform_ocr_request("deepseek-ai/deepseek-ocr-maas", document(), ¶ms, &[]) + .unwrap(), + ) + .unwrap(); + assert_eq!(result["model"], "deepseek-ai/deepseek-ocr-maas"); + assert_eq!( + result["messages"][0]["content"][0], + json!({"type":"image_url","image_url":"gs://bucket/a.png"}) + ); + assert_eq!(result[name], value); + assert!(result.get("ignored").is_none()); + } + + #[rstest] + #[case(json!({"type":"image_url","image_url":"data:image/png;base64,AA=="}))] + #[case(json!({"type":"document_url","document_url":"data:application/pdf;base64,AA=="}))] + fn request_maps_both_document_types_to_image_content(#[case] document: Value) { + let source = document + .get("image_url") + .or_else(|| document.get("document_url")) + .unwrap() + .clone(); + let request = VertexAIDeepSeekOCRConfig + .transform_ocr_request( + "deepseek-ai/deepseek-ocr-maas", + serde_json::from_value(document).unwrap(), + &DeepSeekOcrParams::default(), + &[], + ) + .unwrap(); + let result = serde_json::to_value(request).unwrap(); + assert_eq!( + result["messages"][0]["content"][0], + json!({"type":"image_url","image_url":source}) + ); + } + + #[rstest] + #[case(json!("# hello"), "# hello")] + #[case(json!("{broken"), "{broken")] + #[case(json!(" {\"pages\":[]} "), " {\"pages\":[]} ")] + #[case(json!({"pages":[]}), "")] + #[case(json!("[]"), "[]")] + #[case(json!("{\"pages\":[{\"markdown\":\"json text\"}]}"), "json text")] + #[case(json!({"pages":[{"markdown":"object"}]}), "object")] + fn response_transform_handles_text_json_and_objects( + #[case] content: Value, + #[case] expected: &str, + ) { + let has_pages = content + .as_object() + .is_some_and(|data| data.contains_key("pages")) + || content + .as_str() + .is_some_and(|text| text.contains("\"pages\"")); + let response: DeepSeekOcrResponse = serde_json::from_value( + json!({"choices":[{"message":{"content":content}}],"usage":{"prompt_tokens":1}}), + ) + .unwrap(); + let result = normalize_response("model", response).unwrap().into_json(); + assert_eq!(result["pages"][0]["markdown"], expected); + assert_eq!(result["pages"][0]["index"], 0); + if has_pages { + assert!(result["usage_info"].is_null()); + } else { + assert_eq!(result["usage_info"]["prompt_tokens"], 1); + } + } + + #[test] + fn structured_result_maps_pages_usage_model_and_annotation() { + let response: DeepSeekOcrResponse = serde_json::from_value(json!({ + "choices":[{"message":{"content":{ + "pages":[{"index":2,"markdown":"page","images":[{"id":"one"}],"dimensions":{"width":10}}], + "model":"provider-model", + "usage_info":{"pages_processed":1}, + "document_annotation":{"language":"en"}, + "future":"kept" + }}}] + })) + .unwrap(); + let result = normalize_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!(result.get("future").is_none()); + } + + #[test] + fn response_transform_rejects_missing_empty_and_malformed_content() { + for value in [ + json!({"choices":[]}), + json!({"choices":[{"message":{"content":{}}}]}), + 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| normalize_response("model", response).map_err(|_| ())); + assert!(result.is_err()); + } + } + + #[test] + fn structured_content_preserves_usage_presence_and_shared_page_defaults() { + for (usage, expected) in [(json!(null), None), (json!({"pages_processed":2}), Some(2))] { + let response = serde_json::from_value(json!({ + "choices":[{"message":{"content":{ + "pages":[42, {"index":"2", "images":[{"id":"kept"}], "ignored":true}], + "usage_info":usage + }}}], + "usage":{"pages_processed":99} + })) + .unwrap(); + let normalized = normalize_response("model", response).unwrap(); + assert_eq!(normalized.pages.len(), 1); + assert_eq!(normalized.pages[0].index, 2); + assert_eq!(normalized.pages[0].markdown, ""); + assert!(normalized.pages[0].extra_fields.is_empty()); + assert_eq!( + normalized + .usage_info + .and_then(|usage| usage.pages_processed), + expected + ); + } + } + + use crate::ocr::test_support::{MockResponse, mock_server, perform_ocr, wire_request}; + use litellm_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().extra_fields["prompt_tokens"], + 1 + ); + let requests = seen.lock().unwrap(); + assert!(requests[0].starts_with( + "POST /v1/projects/project-1/locations/europe-west4/endpoints/openapi/chat/completions " + )); + assert!( + requests[0] + .to_ascii_lowercase() + .contains("authorization: bearer test-key") + ); + let body = request_body(&requests[0]); + assert_eq!(body["model"], "deepseek-ai/deepseek-ocr-maas"); + assert_eq!(body["temperature"], 0.1); + assert_eq!(body["future_ocr_option"], true); + assert_eq!(body["provider_option"], "value"); + assert!(body.get("vertex_project").is_none()); + assert!(body.get("extra_body").is_none()); + assert_eq!( + body["messages"][0]["content"][0], + json!({"type":"image_url","image_url":"gs://bucket/document.pdf"}) + ); + } + + #[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/src/llms/vertex_ai/ocr/transformation.rs b/litellm-rust/crates/core/src/llms/vertex_ai/ocr/transformation.rs index a32e0eb55da..d81ea31bc2b 100644 --- a/litellm-rust/crates/core/src/llms/vertex_ai/ocr/transformation.rs +++ b/litellm-rust/crates/core/src/llms/vertex_ai/ocr/transformation.rs @@ -1,6 +1,7 @@ +use litellm_auth_gcp::{self as vertex, VertexConfig}; + use super::common_utils::validate_destination; use crate::llms::base_llm::ocr::transformation::{BaseOcrConfig, OcrRequestContext}; -use crate::llms::mistral::ocr::MistralOcrResponse; use crate::llms::mistral::ocr::transformation::{MistralOCRConfig, MistralOcrRequest}; use crate::ocr::OcrClient; use crate::ocr::document::{inline_remote_document, validate_inline_document}; @@ -8,7 +9,7 @@ use crate::ocr::prepare::{credential_env, transform_request_body}; use crate::ocr::types::{LiteLLMOcrRequest, LiteLLMOcrResponse, OcrConnection, OcrDocument}; use crate::params::OpaqueParams; use crate::url_utils::ApiUrl; -use litellm_auth_gcp::{self as vertex, VertexConfig}; + const DEFAULT_LOCATION: &str = "us-central1"; #[derive(Clone, Debug, Default)] @@ -17,7 +18,54 @@ pub(crate) struct VertexAIOCRConfig; impl BaseOcrConfig for VertexAIOCRConfig { type OcrParams = OpaqueParams; type ProviderRequest = MistralOcrRequest; - type ProviderResponse = MistralOcrResponse; + type Environment = vertex::VertexEnvironment; + + fn get_api_key_env_var(&self) -> Option<&'static str> { + Some("VERTEX_AI_API_KEY") + } + + async fn validate_environment( + &self, + request: &LiteLLMOcrRequest, + client: &OcrClient, + ) -> Result { + let config = VertexConfig::from_sourced_optional_params( + &request.optional_params, + &request.input_sources, + )?; + self.validate_environment(&request.connection, &config, client) + .await + } + + fn get_complete_url( + &self, + request: &LiteLLMOcrRequest, + _params: &Self::OcrParams, + environment: &Self::Environment, + ) -> Result { + let config = VertexConfig::from_sourced_optional_params( + &request.optional_params, + &request.input_sources, + )?; + let location = vertex::get_vertex_ai_location(&config, &credential_env) + .unwrap_or_else(|| DEFAULT_LOCATION.to_string()); + self.get_complete_url( + request.connection.api_base.as_deref(), + &environment.project_id, + &location, + &request.model, + ) + } + + fn transform_ocr_request( + &self, + model: &str, + document: OcrDocument, + params: &OpaqueParams, + headers: &[(String, String)], + ) -> Result { + MistralOCRConfig.transform_ocr_request(model, document, params, headers) + } fn get_supported_ocr_params(&self, model: &str) -> &'static [&'static str] { MistralOCRConfig.get_supported_ocr_params(model) @@ -37,15 +85,16 @@ impl BaseOcrConfig for VertexAIOCRConfig { context.connection, ) .await?; - MistralOCRConfig.transform_ocr_request(model, document, optional_params, headers) + self.transform_ocr_request(model, document, optional_params, headers) } - fn normalize_response( + fn transform_ocr_response( &self, model: &str, - response: MistralOcrResponse, + raw_response: &[u8], + request_format: crate::ocr::types::OcrResponseFormat, ) -> Result { - MistralOCRConfig.normalize_response(model, response) + MistralOCRConfig.transform_ocr_response(model, raw_response, request_format) } } @@ -56,22 +105,8 @@ impl VertexAIOCRConfig { client: &OcrClient, ) -> Result { let params = self.map_ocr_params(&request.optional_params, &request.model)?; - let config = VertexConfig::from_sourced_optional_params( - &request.optional_params, - &request.input_sources, - ) - .map_err(crate::ocr::Error::from)?; - let authentication = self - .validate_environment(&request.connection, &config, client) - .await?; - let location = vertex::get_vertex_ai_location(&config, &credential_env) - .unwrap_or_else(|| DEFAULT_LOCATION.to_string()); - let url = self.get_complete_url( - request.connection.api_base.as_deref(), - &authentication.project_id, - &location, - &request.model, - )?; + let authentication = BaseOcrConfig::validate_environment(self, request, client).await?; + let url = BaseOcrConfig::get_complete_url(self, request, ¶ms, &authentication)?; let retains_document = !request.document.source().starts_with("http://") && !request.document.source().starts_with("https://"); let body = self @@ -193,4 +228,173 @@ mod tests { .is_err() ); } + + use serde_json::{Value, json}; + + use crate::ocr::test_support::{MockResponse, mock_server, perform_ocr, wire_request}; + use litellm_auth::InputSource; + + fn request_body(request: &str) -> Value { + serde_json::from_str(request.split_once("\r\n\r\n").unwrap().1).unwrap() + } + + #[tokio::test] + async fn facade_executes_vertex_mistral_with_resolved_project_and_location() { + let (base, seen, server) = mock_server(vec![MockResponse::json(json!({ + "pages":[{"index":0,"markdown":"hello"}], + "usage_info":{"pages_processed":1} + }))]) + .await; + let request = wire_request( + "vertex_ai/mistral-ocr-maas", + &base, + json!({ + "vertex_project":"project-1", + "vertex_location":"europe-west4", + "extract_footer":true + }), + ); + + let response = perform_ocr(request).await.unwrap(); + server.await.unwrap(); + assert_eq!(response.pages[0].markdown, "hello"); + let requests = seen.lock().unwrap(); + assert_eq!(requests.len(), 1); + assert!(requests[0].starts_with( + "POST /v1/projects/project-1/locations/europe-west4/publishers/mistralai/models/mistral-ocr-maas:rawPredict " + )); + assert!( + requests[0] + .to_ascii_lowercase() + .contains("authorization: bearer test-key") + ); + assert_eq!( + request_body(&requests[0]), + json!({ + "model":"mistral-ocr-maas", + "document":{"type":"document_url","document_url":"data:application/pdf;base64,YWJj"}, + "extract_footer":true + }) + ); + } + + #[tokio::test] + async fn supplied_authorization_is_forwarded_without_a_static_token() { + let (base, seen, server) = mock_server(vec![MockResponse::json(json!({"pages":[]}))]).await; + let mut request = wire_request( + "vertex_ai/model", + &base, + json!({"vertex_project":"project-1"}), + ); + request.connection.api_key = None; + request.connection.extra_headers = vec![("authorization".into(), "Bearer supplied".into())]; + + perform_ocr(request).await.unwrap(); + server.await.unwrap(); + assert!( + seen.lock().unwrap()[0] + .to_ascii_lowercase() + .contains("authorization: bearer supplied") + ); + } + + #[tokio::test] + async fn invalid_credentials_fail_before_provider_http() { + let request = wire_request( + "vertex_ai/model", + "http://127.0.0.1:1", + json!({"vertex_credentials": true}), + ); + let error = perform_ocr(request).await.unwrap_err(); + assert!(error.to_string().contains("vertex_credentials")); + } + + #[tokio::test] + async fn request_controlled_api_base_is_rejected_before_vertex_auth() { + let mut request = wire_request( + "vertex_ai/mistral-ocr-maas", + "https://caller.example", + json!({"vertex_project":"project-1"}), + ); + request.connection.api_base_source = InputSource::Request; + + let error = perform_ocr(request).await.unwrap_err(); + assert!( + error + .to_string() + .contains("request-controlled Vertex AI endpoint") + ); + } + + #[tokio::test] + async fn configs_build_complete_requests_and_share_mistral_normalization() { + use std::time::Duration; + + use crate::llms::base_llm::ocr::transformation::BaseOcrConfig; + use crate::llms::mistral::ocr::transformation::MistralOCRConfig; + use crate::llms::vertex_ai::ocr::transformation::VertexAIOCRConfig; + use crate::ocr::test_support::ocr_client; + + let client = ocr_client(); + let options = json!({ + "pages": [0, 2], + "include_image_base64": true, + "vertex_project": "project-1", + "vertex_location": "us-central1", + "unknown": "preserved" + }); + let direct = wire_request( + "mistral/mistral-ocr-maas", + "https://mistral.test", + options.clone(), + ); + let vertex = wire_request("vertex_ai/mistral-ocr-maas", "https://vertex.test", options); + let direct_http = MistralOCRConfig + .prepare_request(&direct, &client) + .await + .unwrap(); + let vertex_http = VertexAIOCRConfig + .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, + "unknown": "preserved" + }) + ); + } + let payload = serde_json::to_vec( + &json!({"pages": [{"index": 0, "markdown": "hello"}], "extra": "preserved"}), + ) + .unwrap(); + let direct_response = MistralOCRConfig + .transform_ocr_response(&direct.model, &payload, Default::default()) + .unwrap() + .into_json(); + let vertex_response = VertexAIOCRConfig + .transform_ocr_response(&vertex.model, &payload, Default::default()) + .unwrap() + .into_json(); + assert_eq!(direct_response, vertex_response); + assert_eq!(direct_response["model"], "mistral-ocr-maas"); + assert_eq!(direct_response["object"], "ocr"); + assert!(direct_response.get("extra").is_none()); + } } diff --git a/litellm-rust/crates/core/src/ocr/error.rs b/litellm-rust/crates/core/src/ocr/error.rs index 19d916804ed..9f4c2f34ae1 100644 --- a/litellm-rust/crates/core/src/ocr/error.rs +++ b/litellm-rust/crates/core/src/ocr/error.rs @@ -1,5 +1,11 @@ #[derive(Clone, Debug, PartialEq, Eq, thiserror::Error)] pub enum Error { + #[error("upstream OCR error ({status}): {body}")] + Provider { + status: u16, + body: String, + headers: Vec<(String, String)>, + }, #[error("File is empty or could not be read")] EmptyFile, #[error("Invalid MIME type: {0}")] diff --git a/litellm-rust/crates/core/src/ocr/handler.rs b/litellm-rust/crates/core/src/ocr/handler.rs index 26976c044d8..73ebd38bf2e 100644 --- a/litellm-rust/crates/core/src/ocr/handler.rs +++ b/litellm-rust/crates/core/src/ocr/handler.rs @@ -23,7 +23,7 @@ pub(crate) async fn perform_ocr_request( let context = CallLifecycleContext::new( "ocr", request.model.clone(), - request.config.provider().as_str(), + request.provider_name(), request .litellm_call_id .clone() @@ -54,6 +54,7 @@ impl PreparedOcrCall { client: OcrClient, request: LiteLLMOcrRequest, ) -> Result { + let request = super::prepare::resolve_connection_params(request); let http = match request.config { OcrConfigKind::Cohere => CohereParseConfig.prepare_request(&request, &client).await?, OcrConfigKind::Mistral => MistralOCRConfig.prepare_request(&request, &client).await?, @@ -99,6 +100,30 @@ impl PreparedOcrCall { crate::http_utils::execute_http_request(self.client.provider_http(), self.http) .await .map_err(super::client::transport_error)?; + if !response.status().is_success() { + let headers = response + .headers() + .iter() + .filter_map(|(name, value)| { + value + .to_str() + .ok() + .map(|value| (name.to_string(), value.to_string())) + }) + .collect(); + return match super::client::read_response_bytes( + response, + self.request.connection.max_response_bytes, + ) + .await + { + Err(super::Error::Transport(crate::transport::Error::Http { status, body })) => { + Err(self.request.config.get_error_class(body, status, headers)) + } + Err(error) => Err(error), + Ok(_) => unreachable!("non-success response produces an HTTP error"), + }; + } let model = &self.request.model; let context = OcrResponseContext { client: &self.client, diff --git a/litellm-rust/crates/core/src/ocr/lifecycle.rs b/litellm-rust/crates/core/src/ocr/lifecycle.rs index b11f5688c59..f72094e2291 100644 --- a/litellm-rust/crates/core/src/ocr/lifecycle.rs +++ b/litellm-rust/crates/core/src/ocr/lifecycle.rs @@ -652,3 +652,966 @@ impl OcrHost for OcrHookHost { }) } } + +#[cfg(test)] +mod tests { + use std::sync::{Arc, Mutex}; + + use serde_json::{Value, json}; + + use crate::call_lifecycle::{CallLifecycleContext, CallLifecycleTiming}; + use crate::ocr::OcrClient; + use crate::ocr::hooks::{ + OcrDuringCallRequest, OcrHookFuture, OcrHooks, OcrLogFuture, OcrPostCallRequest, + OcrPreCallRequest, + }; + use crate::ocr::test_support::{MockResponse, mock_server, perform_ocr, wire_request}; + use crate::ocr::wire::{OcrWireRequest, decode_request}; + use crate::ocr::{ + NativeOutcome, NoopOcrHost, OcrAdmission, OcrCall, OcrCallStep, OcrDecline, OcrHost, + OcrHostOperation, OcrHostResult, + }; + + #[test] + fn request_boundary_selects_mistral_and_rejects_unknown_providers() { + let request = OcrWireRequest { + model: "mistral/model".into(), + document: json!({"type":"document_url","document_url":"https://example.com/doc.pdf"}), + api_key: Some("key".into()), + api_base: None, + custom_llm_provider: None, + extra_headers: None, + optional_params: json!({"extract_header":true,"unknown":42}) + .as_object() + .unwrap() + .clone() + .into(), + input_sources: Default::default(), + timeout_seconds: None, + }; + assert!(decode_request(request).is_ok()); + assert!( + decode_request(OcrWireRequest { + model: "model".into(), + document: json!({"type":"document_url","document_url":"https://example.com/doc.pdf"}), + api_key: Some("key".into()), + api_base: None, + custom_llm_provider: Some("unknown".into()), + extra_headers: None, + optional_params: Default::default(), + input_sources: Default::default(), + timeout_seconds: None, + }) + .is_err() + ); + } + + #[tokio::test] + async fn facade_executes_direct_mistral_once() { + let (base, seen, server) = mock_server(vec![MockResponse::json(json!({ + "pages":[{"index":0,"markdown":"hello","custom":"preserved"}], + "usage_info":{"pages_processed":1} + }))]) + .await; + let result = perform_ocr(wire_request( + "mistral/model", + &base, + json!({"pages":"0,2-4","extract_header":true,"unknown":{"nested":[null,false,0]}}), + )) + .await + .unwrap(); + server.await.unwrap(); + assert_eq!(result.pages[0].markdown, "hello"); + assert_eq!(result.pages[0].extra_fields["custom"], "preserved"); + let requests = seen.lock().unwrap(); + assert_eq!(requests.len(), 1); + assert!(requests[0].starts_with("POST /v1/ocr ")); + assert!( + requests[0] + .to_ascii_lowercase() + .contains("authorization: bearer test-key\r\n") + ); + let body: Value = + serde_json::from_str(requests[0].split_once("\r\n\r\n").unwrap().1).unwrap(); + assert_eq!( + body, + json!({ + "model":"model", + "document":{"type":"document_url","document_url":"data:application/pdf;base64,YWJj"}, + "pages":"0,2-4", + "extract_header":true, + "unknown":{"nested":[null,false,0]} + }) + ); + } + + #[tokio::test] + async fn facade_resolves_dynamic_connection_before_auth_and_url_preparation() { + use litellm_auth::{InputSource, Sourced}; + let (base, seen, server) = + mock_server(vec![MockResponse::json(json!({"pages": []}))]).await; + let request = wire_request("mistral/model", "https://unused.invalid", json!({})); + let request = crate::ocr::LiteLLMOcrRequest { + connection: crate::ocr::OcrConnection { + dynamic_api_key: Some(Sourced::new("dynamic-key".into(), InputSource::Deployment)), + dynamic_api_base: Some(Sourced::new(base, InputSource::Deployment)), + ..request.connection + }, + ..request + }; + perform_ocr(request).await.unwrap(); + server.await.unwrap(); + let requests = seen.lock().unwrap(); + assert_eq!(requests.len(), 1); + assert!(requests[0].starts_with("POST /v1/ocr ")); + assert!( + requests[0] + .to_ascii_lowercase() + .contains("authorization: bearer dynamic-key\r\n") + ); + } + + #[tokio::test] + async fn provider_error_factory_preserves_status_body_and_response_headers() { + let (base, seen, server) = mock_server(vec![MockResponse { + status: 429, + headers: vec![ + ("retry-after", "17".into()), + ("x-request-id", "ocr-request".into()), + ], + body: json!({"message": "rate limited"}), + }]) + .await; + let error = perform_ocr(wire_request("mistral/model", &base, json!({}))) + .await + .unwrap_err(); + server.await.unwrap(); + let crate::ocr::Error::Provider { + status, + body, + headers, + } = error + else { + panic!("expected provider error") + }; + assert_eq!(status, 429); + assert_eq!( + serde_json::from_str::(&body).unwrap(), + json!({"message": "rate limited"}) + ); + assert!(headers.contains(&("retry-after".into(), "17".into()))); + assert!(headers.contains(&("x-request-id".into(), "ocr-request".into()))); + assert_eq!(seen.lock().unwrap().len(), 1); + } + + #[tokio::test] + async fn facade_retains_native_response_when_requested() { + let provider_response = json!({ + "pages":[{"index":0,"markdown":"hello"}], + "usage_info":{"pages_processed":1}, + "provider_only":"preserved" + }); + let (base, _, server) = + mock_server(vec![MockResponse::json(provider_response.clone())]).await; + let response = perform_ocr(wire_request( + "mistral/model", + &base, + json!({"req_format":"native"}), + )) + .await + .unwrap(); + + server.await.unwrap(); + assert_eq!( + response.provider_native_response.as_ref(), + provider_response.as_object() + ); + } + + #[tokio::test] + async fn facade_uses_the_injected_http_client() { + let (base, seen, server) = mock_server(vec![MockResponse::json(json!({"pages":[]}))]).await; + let mut default_headers = reqwest::header::HeaderMap::new(); + default_headers.insert( + "x-transport-owner", + reqwest::header::HeaderValue::from_static("host"), + ); + let provider_http = reqwest::Client::builder() + .default_headers(default_headers) + .build() + .unwrap(); + OcrClient::new(provider_http) + .unwrap() + .perform(wire_request("mistral/model", &base, json!({}))) + .await + .unwrap(); + server.await.unwrap(); + assert!(seen.lock().unwrap()[0].contains("x-transport-owner: host")); + } + + struct RecordingHooks { + events: Arc>>, + block: bool, + } + + struct ExtensionHooks; + + impl OcrHooks for ExtensionHooks { + fn intercepts_requests(&self) -> bool { + true + } + + fn during_call( + &self, + mut request: OcrDuringCallRequest, + ) -> OcrHookFuture<'_, OcrDuringCallRequest> { + Box::pin(async move { + assert_eq!(request.body["pages"], json!([2])); + assert_eq!(request.body.get("future"), Some(&Value::Null)); + assert!( + !request + .retained_fields + .iter() + .any(|field| field == "pages" || field == "document") + ); + request.body.as_object_mut().unwrap().remove("future"); + request.body["hook_option"] = json!({"nested":[null,false,0]}); + Ok(request) + }) + } + } + + #[tokio::test] + async fn composed_extensions_reach_hooks_and_removed_fields_stay_removed() { + let (base, seen, server) = mock_server(vec![MockResponse::json(json!({"pages":[]}))]).await; + let request = crate::ocr::LiteLLMOcrRequest { + hooks: Arc::new(ExtensionHooks), + ..wire_request( + "mistral/model", + &base, + json!({ + "pages":[0], "future":null, "extra_body":{"pages":[2], + "document":{"type":"document_url","document_url":"data:application/pdf;base64,eHl6"}} + }), + ) + }; + perform_ocr(request).await.unwrap(); + server.await.unwrap(); + let requests = seen.lock().unwrap(); + let body: Value = + serde_json::from_str(requests[0].split_once("\r\n\r\n").unwrap().1).unwrap(); + assert_eq!(body["pages"], json!([2])); + assert_eq!( + body["document"]["document_url"], + "data:application/pdf;base64,eHl6" + ); + assert_eq!(body["hook_option"], json!({"nested":[null,false,0]})); + assert!(body.get("future").is_none()); + } + + impl OcrHooks for RecordingHooks { + fn intercepts_requests(&self) -> bool { + true + } + + fn pre_call(&self, request: OcrPreCallRequest) -> OcrHookFuture<'_, OcrPreCallRequest> { + Box::pin(async move { + self.events.lock().unwrap().push("pre"); + if self.block { + return Err(crate::ocr::Error::InvalidRequest("blocked".into())); + } + Ok(request) + }) + } + + fn during_call( + &self, + request: crate::ocr::hooks::OcrDuringCallRequest, + ) -> OcrHookFuture<'_, crate::ocr::hooks::OcrDuringCallRequest> { + Box::pin(async move { + self.events.lock().unwrap().push("during"); + Ok(request) + }) + } + + fn post_call(&self, request: OcrPostCallRequest) -> OcrHookFuture<'_, OcrPostCallRequest> { + Box::pin(async move { + self.events.lock().unwrap().push("post"); + Ok(request) + }) + } + + fn success<'a>( + &'a self, + _context: &'a CallLifecycleContext, + _response: &'a crate::ocr::LiteLLMOcrResponse, + _timing: &'a CallLifecycleTiming, + ) -> OcrLogFuture<'a> { + Box::pin(async move { + self.events.lock().unwrap().push("success"); + }) + } + + fn failure<'a>( + &'a self, + _context: &'a CallLifecycleContext, + _error: &'a crate::ocr::Error, + _timing: &'a CallLifecycleTiming, + ) -> OcrLogFuture<'a> { + Box::pin(async move { + self.events.lock().unwrap().push("failure"); + }) + } + } + + struct HeaderEditHooks; + + impl OcrHooks for HeaderEditHooks { + fn intercepts_requests(&self) -> bool { + true + } + + fn during_call( + &self, + mut request: OcrDuringCallRequest, + ) -> OcrHookFuture<'_, OcrDuringCallRequest> { + request + .headers + .push(("x-core-callback".into(), "edited".into())); + Box::pin(async move { Ok(request) }) + } + } + + #[tokio::test] + async fn lifecycle_sends_headers_returned_by_the_typed_during_call_operation() { + let (base, seen, server) = mock_server(vec![MockResponse::json(json!({"pages":[]}))]).await; + let request = crate::ocr::LiteLLMOcrRequest { + hooks: Arc::new(HeaderEditHooks), + ..wire_request("mistral/model", &base, json!({})) + }; + + perform_ocr(request).await.unwrap(); + server.await.unwrap(); + + assert!(seen.lock().unwrap()[0].contains("x-core-callback: edited")); + } + + #[tokio::test] + async fn lifecycle_orders_hooks_and_emits_one_success() { + let (base, seen, server) = mock_server(vec![MockResponse::json(json!({"pages":[]}))]).await; + let events = Arc::new(Mutex::new(Vec::new())); + let request = wire_request("mistral/model", &base, json!({})); + let request = crate::ocr::LiteLLMOcrRequest { + hooks: Arc::new(RecordingHooks { + events: events.clone(), + block: false, + }), + ..request + }; + perform_ocr(request).await.unwrap(); + server.await.unwrap(); + assert_eq!( + *events.lock().unwrap(), + ["pre", "during", "post", "success"] + ); + assert_eq!(seen.lock().unwrap().len(), 1); + } + + #[tokio::test] + async fn lifecycle_blocking_prevents_execution_and_emits_one_failure() { + let events = Arc::new(Mutex::new(Vec::new())); + let request = wire_request("mistral/model", "http://127.0.0.1:1", json!({})); + let request = crate::ocr::LiteLLMOcrRequest { + hooks: Arc::new(RecordingHooks { + events: events.clone(), + block: true, + }), + ..request + }; + let error = perform_ocr(request).await.unwrap_err(); + assert!(matches!(error, crate::ocr::Error::InvalidRequest(_))); + assert_eq!(*events.lock().unwrap(), ["pre", "failure"]); + } + + #[tokio::test] + async fn upstream_failure_emits_one_terminal_failure() { + let (base, seen, server) = mock_server(vec![MockResponse { + status: 500, + headers: vec![], + body: json!({"error":"failed"}), + }]) + .await; + let events = Arc::new(Mutex::new(Vec::new())); + let request = wire_request("mistral/model", &base, json!({})); + let request = crate::ocr::LiteLLMOcrRequest { + hooks: Arc::new(RecordingHooks { + events: events.clone(), + block: false, + }), + ..request + }; + assert!(perform_ocr(request).await.is_err()); + server.await.unwrap(); + assert_eq!(*events.lock().unwrap(), ["pre", "during", "failure"]); + assert_eq!(seen.lock().unwrap().len(), 1); + } + + struct AdmissionSpy { + effects: Arc>, + } + + impl OcrHooks for AdmissionSpy { + fn intercepts_requests(&self) -> bool { + *self.effects.lock().unwrap() += 1; + true + } + + fn pre_call(&self, request: OcrPreCallRequest) -> OcrHookFuture<'_, OcrPreCallRequest> { + *self.effects.lock().unwrap() += 1; + Box::pin(async move { Ok(request) }) + } + } + + #[test] + fn admission_declines_without_invoking_hooks_or_transport() { + for (admission, expected) in [ + ( + OcrAdmission { + provider_workflow: false, + host_operations: true, + asynchronous: false, + }, + OcrDecline::ProviderWorkflow, + ), + ( + OcrAdmission { + provider_workflow: true, + host_operations: false, + asynchronous: false, + }, + OcrDecline::HostOperations, + ), + ] { + let outcome = OcrCall::admit(crate::ocr::test_support::ocr_client(), admission); + assert!(matches!(outcome, NativeOutcome::Declined(reason) if reason == expected)); + } + } + + #[tokio::test] + async fn fallible_host_phases_do_not_replay_or_reach_transport() { + for failure_phase in ["pre", "during"] { + let request = crate::ocr::LiteLLMOcrRequest { + hooks: Arc::new(AdmissionSpy { + effects: Arc::new(Mutex::new(0)), + }), + ..wire_request("mistral/model", "http://127.0.0.1:1", json!({})) + }; + let NativeOutcome::Completed(mut call) = + OcrCall::admit(crate::ocr::test_support::ocr_client(), OcrAdmission::all()) + else { + panic!("supported call declined") + }; + let mut request = Some(request); + let mut result = None; + let mut phases = Vec::new(); + let error = loop { + match call.resume(result.take()).await { + Ok(OcrCallStep::Host(operation)) => match operation { + OcrHostOperation::Lifecycle(_) + | OcrHostOperation::ConstructResponse(_) + | OcrHostOperation::MapFailure(_) + | OcrHostOperation::Success { .. } + | OcrHostOperation::Failure { .. } => { + result = Some(OcrHostResult::Lifecycle(Ok(()))) + } + OcrHostOperation::ProjectRequest => { + result = Some(OcrHostResult::Request(Ok(( + Box::new(request.take().unwrap()), + false, + )))) + } + OcrHostOperation::AcquireAzureAdToken => { + panic!("test request has no token provider") + } + OcrHostOperation::PreCall(request) => { + phases.push("pre"); + result = Some(OcrHostResult::PreCall(if failure_phase == "pre" { + Err(crate::ocr::Error::InvalidRequest("pre failed".into())) + } else { + Ok(request) + })); + } + OcrHostOperation::DuringCall(request) => { + phases.push("during"); + result = + Some(OcrHostResult::DuringCall(if failure_phase == "during" { + Err(crate::ocr::Error::InvalidRequest("during failed".into())) + } else { + Ok(request) + })); + } + OcrHostOperation::PostCall(_) => panic!("transport should not be reached"), + }, + Err(error) => break error, + Ok(OcrCallStep::Complete(_)) => panic!("failed call completed"), + } + }; + assert!(matches!(error, crate::ocr::Error::InvalidRequest(_))); + assert_eq!( + phases + .iter() + .filter(|phase| **phase == failure_phase) + .count(), + 1 + ); + } + } + + #[tokio::test] + async fn invalid_provider_response_runs_post_call_before_normalization_failure() { + let (base, seen, server) = + mock_server(vec![MockResponse::json(json!({"pages":"invalid"}))]).await; + let mut request = Some(wire_request("mistral/model", &base, json!({}))); + let NativeOutcome::Completed(mut call) = + OcrCall::admit(crate::ocr::test_support::ocr_client(), OcrAdmission::all()) + else { + panic!("supported call declined") + }; + let host = NoopOcrHost; + let mut result = None; + let mut post_calls = Vec::new(); + let error = loop { + match call.resume(result.take()).await { + Ok(OcrCallStep::Host(OcrHostOperation::ProjectRequest)) => { + result = Some(OcrHostResult::Request(Ok(( + Box::new(request.take().unwrap()), + false, + )))); + } + Ok(OcrCallStep::Host(operation)) => { + if let OcrHostOperation::PostCall(request) = &operation { + post_calls.push(request.original_response.clone()); + } + result = Some(host.invoke(operation).await); + } + Err(error) => break error, + Ok(OcrCallStep::Complete(_)) => panic!("invalid provider response completed"), + } + }; + server.await.unwrap(); + assert!(matches!( + error, + crate::ocr::Error::ResponseField { ref path } if path == "pages" + )); + assert_eq!(seen.lock().unwrap().len(), 1); + assert_eq!(post_calls, [json!(r#"{"pages":"invalid"}"#)]); + } + + #[tokio::test] + async fn direct_native_host_drives_the_same_state_machine() { + let (base, seen, server) = mock_server(vec![MockResponse::json(json!({ + "pages":[{"index":0,"markdown":"native"}] + }))]) + .await; + let request = crate::ocr::LiteLLMOcrRequest { + hooks: Arc::new(AdmissionSpy { + effects: Arc::new(Mutex::new(0)), + }), + ..wire_request("mistral/model", &base, json!({})) + }; + let NativeOutcome::Completed(mut call) = OcrCall::admit( + crate::ocr::test_support::ocr_client(), + OcrAdmission { + asynchronous: true, + ..OcrAdmission::all() + }, + ) else { + panic!("supported call declined") + }; + let mut request = Some(request); + let host = NoopOcrHost; + let mut result = None; + let mut operations = Vec::new(); + let response = loop { + match call.resume(result.take()).await.unwrap() { + OcrCallStep::Host(operation) => { + operations.push(match &operation { + OcrHostOperation::ProjectRequest => "ProjectRequest".into(), + OcrHostOperation::Lifecycle(phase) => format!("{phase:?}"), + OcrHostOperation::PreCall(_) => "PreCall".into(), + OcrHostOperation::DuringCall(_) => "DuringCall".into(), + OcrHostOperation::PostCall(_) => "PostCall".into(), + OcrHostOperation::ConstructResponse(_) => "ConstructResponse".into(), + OcrHostOperation::Success { response, .. } => { + assert_eq!(response.pages[0].markdown, "native"); + "Success".into() + } + _ => panic!("unexpected OCR operation"), + }); + result = Some(match operation { + OcrHostOperation::ProjectRequest => { + OcrHostResult::Request(Ok((Box::new(request.take().unwrap()), false))) + } + operation => host.invoke(operation).await, + }); + } + OcrCallStep::Complete(response) => break response, + } + }; + server.await.unwrap(); + assert_eq!(response.pages[0].markdown, "native"); + assert_eq!(seen.lock().unwrap().len(), 1); + assert_eq!( + operations, + [ + "Setup", + "DeploymentPreCall", + "Prepare", + "ProjectRequest", + "PreCall", + "DuringCall", + "PostCall", + "ConstructResponse", + "DeploymentPostCall", + "Finalize", + "Success", + ] + ); + assert!(matches!( + call.resume(None).await, + Err(crate::ocr::Error::InvalidRequest(_)) + )); + } + + #[tokio::test] + async fn public_finalization_failure_never_dispatches_success_or_replays_provider() { + use crate::call_lifecycle::host::{HostFailure, HostPhase}; + + let (base, seen, server) = mock_server(vec![MockResponse::json(json!({"pages":[]}))]).await; + let mut request = Some(wire_request("mistral/model", &base, json!({}))); + let NativeOutcome::Completed(mut call) = OcrCall::admit( + crate::ocr::test_support::ocr_client(), + OcrAdmission { + asynchronous: true, + ..OcrAdmission::all() + }, + ) else { + panic!("supported call declined") + }; + let selected = crate::ocr::Error::InvalidRequest("public metadata failed".into()); + let host = NoopOcrHost; + let mut result = None; + let mut failures = Vec::new(); + let error = loop { + match call.resume(result.take()).await { + Ok(OcrCallStep::Host(operation)) => { + result = Some(match operation { + OcrHostOperation::Lifecycle(HostPhase::Finalize) => { + OcrHostResult::Lifecycle(Err(HostFailure::Error(selected.clone()))) + } + OcrHostOperation::Failure { error, .. } => { + assert_eq!(error, selected); + failures.push("sync"); + OcrHostResult::Lifecycle(Err(HostFailure::Error( + crate::ocr::Error::InvalidRequest("failure callback failed".into()), + ))) + } + OcrHostOperation::Lifecycle(HostPhase::AsyncFailure) => { + failures.push("async"); + OcrHostResult::Lifecycle(Ok(())) + } + OcrHostOperation::Success { .. } + | OcrHostOperation::MapFailure(_) + | OcrHostOperation::Lifecycle(HostPhase::DeploymentFailure) => { + panic!("finalization failure used provider/success dispatch") + } + OcrHostOperation::ProjectRequest => { + OcrHostResult::Request(Ok((Box::new(request.take().unwrap()), false))) + } + operation => host.invoke(operation).await, + }); + } + Ok(OcrCallStep::Complete(_)) => panic!("failed call completed successfully"), + Err(error) => break error, + } + }; + server.await.unwrap(); + assert_eq!(error, selected); + assert_eq!(failures, ["sync", "async"]); + assert_eq!(seen.lock().unwrap().len(), 1); + } + + #[tokio::test] + async fn cancellation_at_provider_hook_prevents_execution_and_further_resumption() { + use crate::call_lifecycle::host::HostFailure; + + let request = crate::ocr::LiteLLMOcrRequest { + hooks: Arc::new(AdmissionSpy { + effects: Arc::new(Mutex::new(0)), + }), + ..wire_request("mistral/model", "http://127.0.0.1:1", json!({})) + }; + let NativeOutcome::Completed(mut call) = + OcrCall::admit(crate::ocr::test_support::ocr_client(), OcrAdmission::all()) + else { + panic!("supported call declined") + }; + let mut request = Some(request); + let host = NoopOcrHost; + let mut result = None; + loop { + match call.resume(result.take()).await.unwrap() { + OcrCallStep::Host(OcrHostOperation::PreCall(_)) => break, + OcrCallStep::Host(OcrHostOperation::ProjectRequest) => { + result = Some(OcrHostResult::Request(Ok(( + Box::new(request.take().unwrap()), + false, + )))) + } + OcrCallStep::Host(operation) => result = Some(host.invoke(operation).await), + OcrCallStep::Complete(_) => panic!("provider executed before pre-call result"), + } + } + let selected = crate::ocr::Error::InvalidRequest("cancelled".into()); + assert!(matches!( + call.interrupt(HostFailure::Cancelled(selected.clone())).await, + Err(error) if error == selected + )); + assert!( + call.resume(Some(OcrHostResult::Lifecycle(Ok(())))) + .await + .is_err() + ); + } + + #[tokio::test] + async fn missing_host_result_preserves_pending_operation() { + use crate::call_lifecycle::host::HostPhase; + + let NativeOutcome::Completed(mut call) = + OcrCall::admit(crate::ocr::test_support::ocr_client(), OcrAdmission::all()) + else { + panic!("supported call declined") + }; + assert!(matches!( + call.resume(None).await.unwrap(), + OcrCallStep::Host(OcrHostOperation::Lifecycle(HostPhase::Setup)) + )); + assert!(call.resume(None).await.is_err()); + assert!(matches!( + call.resume(Some(OcrHostResult::Lifecycle(Ok(())))) + .await + .unwrap(), + OcrCallStep::Host(OcrHostOperation::Lifecycle(HostPhase::Prepare)) + )); + } + + async fn read_bounded_response( + response: Vec, + limit: usize, + ) -> Result { + use tokio::io::{AsyncReadExt, AsyncWriteExt}; + + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let address = listener.local_addr().unwrap(); + let server = tokio::spawn(async move { + let (mut socket, _) = listener.accept().await.unwrap(); + let mut request = [0; 4096]; + assert!(socket.read(&mut request).await.unwrap() > 0); + socket.write_all(&response).await.unwrap(); + std::future::pending::<()>().await; + }); + let response = reqwest::Client::new() + .get(format!("http://{address}")) + .send() + .await + .unwrap(); + let result = tokio::time::timeout( + std::time::Duration::from_secs(2), + crate::ocr::client::read_response_bytes(response, limit), + ) + .await; + server.abort(); + let _ = server.await; + result.expect("bounded reads must finish without waiting for the rest of an oversized body") + } + + #[tokio::test] + async fn response_limit_accepts_exact_size_and_rejects_declared_and_chunked_overflow() { + for response in [ + "HTTP/1.1 200 OK\r\nContent-Length: 8\r\n\r\nabcdefgh", + "HTTP/1.1 200 OK\r\nTransfer-Encoding: chunked\r\n\r\n4\r\nabcd\r\n4\r\nefgh\r\n0\r\n\r\n", + ] { + assert_eq!( + read_bounded_response(response.as_bytes().to_vec(), 8) + .await + .unwrap(), + "abcdefgh" + ); + } + for response in [ + "HTTP/1.1 200 OK\r\nContent-Length: 9\r\n\r\n", + "HTTP/1.1 200 OK\r\nTransfer-Encoding: chunked\r\n\r\n4\r\nabcd\r\n5\r\nefghi\r\n", + ] { + assert!(matches!( + read_bounded_response(response.as_bytes().to_vec(), 8).await, + Err(crate::ocr::Error::TooLarge { limit: 8 }) + )); + } + } + + #[tokio::test] + async fn oversized_error_retains_http_status_and_bounded_diagnostics_without_draining() { + let prefix = "x".repeat(4 * (crate::constants::UPSTREAM_ERROR_BODY_MAX_CHARS + 1)); + for headers in ["Content-Length: 1000000", "Transfer-Encoding: chunked"] { + let body = if headers.starts_with("Transfer") { + format!("{:x}\r\n{prefix}\r\n", prefix.len()) + } else { + prefix.clone() + }; + let response = format!("HTTP/1.1 429 Too Many Requests\r\n{headers}\r\n\r\n{body}"); + let error = read_bounded_response(response.into_bytes(), 4096) + .await + .unwrap_err(); + match error { + crate::ocr::Error::Transport(crate::transport::Error::Http { status, body }) => { + assert_eq!(status, 429); + assert_eq!( + body, + format!( + "{}... (truncated)", + "x".repeat(crate::constants::UPSTREAM_ERROR_BODY_MAX_CHARS) + ) + ); + } + error => panic!("unexpected error: {error}"), + } + } + } + + #[test] + fn response_limit_is_validated_and_not_forwarded_to_the_provider() { + let request = wire_request( + "mistral/model", + "http://localhost", + json!({"max_response_bytes": 123}), + ); + assert_eq!(request.connection.max_response_bytes, 123); + assert!(!request.optional_params.contains_key("max_response_bytes")); + for value in [ + json!(0), + json!(-1), + json!(true), + json!("123"), + json!(1.5), + json!(crate::constants::OCR_RESPONSE_MAX_BYTES + 1), + Value::Null, + ] { + let wire = serde_json::from_value(json!({ + "model": "mistral/model", "document": {"type": "document_url", "document_url": "data:application/pdf;base64,YWJj"}, + "optional_params": {"max_response_bytes": value} + })).unwrap(); + let Err(error) = decode_request(wire) else { + panic!("invalid response limit accepted") + }; + assert!(error.to_string().contains("max_response_bytes")); + } + } + + #[derive(Debug)] + struct PendingToken { + entered: Arc, + dropped: Arc, + } + + struct TokenFutureDrop(Arc); + + impl Drop for TokenFutureDrop { + fn drop(&mut self) { + self.0.store(true, std::sync::atomic::Ordering::SeqCst); + } + } + + impl litellm_auth::TokenProvider for PendingToken { + fn acquire(&self) -> litellm_auth::TokenFuture<'_> { + Box::pin(async move { + let _guard = TokenFutureDrop(self.dropped.clone()); + self.entered.notify_one(); + std::future::pending().await + }) + } + } + + #[tokio::test] + async fn cancellation_waits_for_provider_capture_drop_even_when_acknowledgement_is_cancelled() { + use crate::call_lifecycle::host::HostFailure; + use std::future::Future; + use std::sync::atomic::{AtomicBool, Ordering}; + use std::task::Poll; + + for interrupt_acknowledgement in [false, true] { + let entered = Arc::new(tokio::sync::Notify::new()); + let dropped = Arc::new(AtomicBool::new(false)); + let request = + wire_request("azure_ai/mistral-ocr", "https://example.invalid", json!({})); + let request = crate::ocr::LiteLLMOcrRequest { + connection: crate::ocr::OcrConnection { + extra_headers: vec![("authorization".into(), "Bearer test-key".into())], + ..request.connection + }, + azure_ad_token_provider: Some(litellm_auth::TokenProviderHandle::new(Arc::new( + PendingToken { + entered: entered.clone(), + dropped: dropped.clone(), + }, + ))), + ..request + }; + let NativeOutcome::Completed(mut call) = + OcrCall::admit(crate::ocr::test_support::ocr_client(), OcrAdmission::all()) + else { + panic!("supported call declined") + }; + let mut request = Some(request); + let mut result = None; + tokio::time::timeout(std::time::Duration::from_secs(2), async { + loop { + tokio::select! { + _ = entered.notified() => break, + step = call.resume(result.take()) => { + result = Some(match step.unwrap() { + OcrCallStep::Host(OcrHostOperation::ProjectRequest) => OcrHostResult::Request(Ok((Box::new(request.take().unwrap()), false))), + OcrCallStep::Host(operation) => NoopOcrHost.invoke(operation).await, + OcrCallStep::Complete(_) => panic!("pending provider completed"), + }); + } + } + } + }).await.unwrap(); + assert!(!dropped.load(Ordering::SeqCst)); + let selected = crate::ocr::Error::InvalidRequest("cancelled".into()); + if interrupt_acknowledgement { + let mut acknowledgement = + Box::pin(call.interrupt(HostFailure::Cancelled(selected.clone()))); + std::future::poll_fn(|cx| { + assert!(acknowledgement.as_mut().poll(cx).is_pending()); + Poll::Ready(()) + }) + .await; + drop(acknowledgement); + assert!(!dropped.load(Ordering::SeqCst)); + } + let result = tokio::time::timeout( + std::time::Duration::from_secs(2), + call.interrupt(HostFailure::Cancelled(selected.clone())), + ) + .await + .unwrap(); + assert!(matches!(result, Err(error) if error == selected)); + assert!( + dropped.load(Ordering::SeqCst), + "cancellation returned while provider captures were still alive" + ); + } + } +} diff --git a/litellm-rust/crates/core/src/ocr/mod.rs b/litellm-rust/crates/core/src/ocr/mod.rs index 89cb2165b60..e5d31a5fe59 100644 --- a/litellm-rust/crates/core/src/ocr/mod.rs +++ b/litellm-rust/crates/core/src/ocr/mod.rs @@ -16,32 +16,11 @@ pub use lifecycle::{ NativeOutcome, NativeResult, NoopOcrHost, OcrAdmission, OcrCall, OcrCallStep, OcrDecline, OcrHookHost, OcrHost, OcrHostOperation, OcrHostResult, }; +pub use provider_config::{get_api_key_env_var, get_health_check_document}; pub use types::{ LiteLLMOcrRequest, LiteLLMOcrResponse, OcrConnection, OcrDocument, OcrPage, OcrPageDimensions, OcrPageImage, OcrUsageInfo, }; #[cfg(test)] -#[path = "../../tests/azure_ai_ocr.rs"] -mod azure_ai_tests; -#[cfg(test)] -#[path = "../../tests/azure_document_intelligence_ocr.rs"] -mod azure_document_intelligence_tests; -#[cfg(test)] -#[path = "../../tests/deepseek_ocr.rs"] -mod deepseek_tests; -#[cfg(test)] -#[path = "../../tests/reducto_ocr.rs"] -mod reducto_tests; -#[cfg(test)] -#[path = "../../tests/ocr/support.rs"] pub(crate) mod test_support; -#[cfg(test)] -#[path = "../../tests/ocr.rs"] -pub(crate) mod tests; -#[cfg(test)] -#[path = "../../tests/vertex_ai_deepseek_ocr.rs"] -mod vertex_ai_deepseek_tests; -#[cfg(test)] -#[path = "../../tests/vertex_ai_ocr.rs"] -mod vertex_ai_tests; diff --git a/litellm-rust/crates/core/src/ocr/prepare.rs b/litellm-rust/crates/core/src/ocr/prepare.rs index e7625a2dc92..b071a8fb4b8 100644 --- a/litellm-rust/crates/core/src/ocr/prepare.rs +++ b/litellm-rust/crates/core/src/ocr/prepare.rs @@ -43,7 +43,7 @@ where .hooks .during_call(OcrDuringCallRequest { model: request.model.clone(), - custom_llm_provider: request.config.provider().as_str().into(), + custom_llm_provider: request.provider_name().into(), url: url.into(), headers: headers.to_vec(), body, @@ -93,7 +93,7 @@ pub(crate) async fn guardrail_document( .hooks .during_call(OcrDuringCallRequest { model: request.model.clone(), - custom_llm_provider: request.config.provider().as_str().into(), + custom_llm_provider: request.provider_name().into(), url: url.into(), headers: headers.to_vec(), body: serde_json::to_value(&request.document).map_err(|_| { @@ -126,3 +126,50 @@ pub(crate) fn body_document(body: &Value) -> Result { pub(crate) fn credential_env(name: &str) -> Option { std::env::var(name).ok() } + +pub(crate) fn resolve_connection_params(request: LiteLLMOcrRequest) -> LiteLLMOcrRequest { + use litellm_auth::{InputSource, Sourced}; + + let connection = request.connection; + let api_base_env = match request.config.provider() { + super::provider_config::OcrProvider::Mistral => Some("MISTRAL_API_BASE"), + super::provider_config::OcrProvider::AzureAi => Some("AZURE_AI_API_BASE"), + super::provider_config::OcrProvider::Cohere + | super::provider_config::OcrProvider::Reducto + | super::provider_config::OcrProvider::VertexAi => None, + }; + let dynamic_api_key = connection.dynamic_api_key.or_else(|| { + connection + .api_key + .clone() + .map(|value| Sourced::new(value, connection.api_key_source)) + .or_else(|| { + request + .config + .get_api_key_env_var() + .and_then(credential_env) + .map(|value| Sourced::new(value, InputSource::Environment)) + }) + }); + let dynamic_api_base = connection.dynamic_api_base.or_else(|| { + connection + .api_base + .clone() + .map(|value| Sourced::new(value, connection.api_base_source)) + .or_else(|| { + api_base_env + .and_then(credential_env) + .map(|value| Sourced::new(value, InputSource::Environment)) + }) + }); + LiteLLMOcrRequest { + connection: request + .config + .resolve_connection_params(super::OcrConnection { + dynamic_api_key, + dynamic_api_base, + ..connection + }), + ..request + } +} diff --git a/litellm-rust/crates/core/src/ocr/provider_config.rs b/litellm-rust/crates/core/src/ocr/provider_config.rs index 531a5e9f346..08d3b810083 100644 --- a/litellm-rust/crates/core/src/ocr/provider_config.rs +++ b/litellm-rust/crates/core/src/ocr/provider_config.rs @@ -1,3 +1,4 @@ +use super::types::{OcrConnection, OcrDocument}; use crate::llms::azure_ai::ocr::cohere_parse_transformation::AzureAICohereParseConfig; use crate::llms::azure_ai::ocr::document_intelligence::transformation::AzureDocumentIntelligenceOCRConfig; use crate::llms::azure_ai::ocr::transformation::AzureAIOCRConfig; @@ -8,6 +9,24 @@ use crate::llms::reducto::ocr::transformation::{ReductoParseLegacyConfig, Reduct use crate::llms::vertex_ai::ocr::deepseek_transformation::VertexAIDeepSeekOCRConfig; use crate::llms::vertex_ai::ocr::transformation::VertexAIOCRConfig; use crate::routing_utils::provider::{CustomLlmProvider, get_custom_llm_provider}; +use litellm_auth::Sourced; +use strum::{EnumString, IntoStaticStr}; + +macro_rules! dispatch_config { + ($config:expr, $method:ident($($argument:expr),* $(,)?)) => { + match $config { + OcrConfigKind::Cohere => CohereParseConfig.$method($($argument),*), + OcrConfigKind::Mistral => MistralOCRConfig.$method($($argument),*), + OcrConfigKind::AzureAi => AzureAIOCRConfig.$method($($argument),*), + OcrConfigKind::AzureCohere => AzureAICohereParseConfig.$method($($argument),*), + OcrConfigKind::AzureDocumentIntelligence => AzureDocumentIntelligenceOCRConfig.$method($($argument),*), + OcrConfigKind::ReductoLegacy => ReductoParseLegacyConfig.$method($($argument),*), + OcrConfigKind::ReductoV3 => ReductoParseV3Config.$method($($argument),*), + OcrConfigKind::VertexAi => VertexAIOCRConfig.$method($($argument),*), + OcrConfigKind::VertexDeepSeek => VertexAIDeepSeekOCRConfig.$method($($argument),*), + } + }; +} #[derive(Clone, Copy, Debug, PartialEq, Eq)] pub(crate) enum OcrConfigKind { @@ -36,23 +55,80 @@ impl OcrConfigKind { } pub(crate) fn get_supported_ocr_params(self, model: &str) -> &'static [&'static str] { - match self { - Self::Cohere => CohereParseConfig.get_supported_ocr_params(model), - Self::Mistral => MistralOCRConfig.get_supported_ocr_params(model), - Self::AzureAi => AzureAIOCRConfig.get_supported_ocr_params(model), - Self::AzureCohere => AzureAICohereParseConfig.get_supported_ocr_params(model), - Self::AzureDocumentIntelligence => { - AzureDocumentIntelligenceOCRConfig.get_supported_ocr_params(model) - } - Self::ReductoLegacy => ReductoParseLegacyConfig.get_supported_ocr_params(model), - Self::ReductoV3 => ReductoParseV3Config.get_supported_ocr_params(model), - Self::VertexAi => VertexAIOCRConfig.get_supported_ocr_params(model), - Self::VertexDeepSeek => VertexAIDeepSeekOCRConfig.get_supported_ocr_params(model), + dispatch_config!(self, get_supported_ocr_params(model)) + } + + pub(crate) fn get_api_key_env_var(self) -> Option<&'static str> { + dispatch_config!(self, get_api_key_env_var()) + } + + pub(crate) fn get_health_check_document(self) -> OcrDocument { + dispatch_config!(self, get_health_check_document()) + } + + pub(crate) fn resolve_connection_params(self, connection: OcrConnection) -> OcrConnection { + let api_key = connection + .api_key + .map(|value| Sourced::new(value, connection.api_key_source)); + let api_base = connection + .api_base + .map(|value| Sourced::new(value, connection.api_base_source)); + let (api_key, api_base) = dispatch_config!( + self, + resolve_connection_params( + api_key, + api_base, + connection.dynamic_api_key, + connection.dynamic_api_base, + ) + ); + OcrConnection { + api_key_source: api_key + .as_ref() + .map(Sourced::source) + .unwrap_or(connection.api_key_source), + api_base_source: api_base + .as_ref() + .map(Sourced::source) + .unwrap_or(connection.api_base_source), + api_key: api_key.map(Sourced::into_value), + api_base: api_base.map(Sourced::into_value), + dynamic_api_key: None, + dynamic_api_base: None, + ..connection } } + + pub(crate) fn get_error_class( + self, + message: String, + status: u16, + headers: Vec<(String, String)>, + ) -> super::Error { + dispatch_config!(self, get_error_class(message, status, headers)) + } } -#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub fn get_api_key_env_var( + model: &str, + custom_llm_provider: Option<&str>, +) -> Result, super::Error> { + Ok(resolve_provider_config(model, custom_llm_provider)? + .1 + .get_api_key_env_var()) +} + +pub fn get_health_check_document( + model: &str, + custom_llm_provider: Option<&str>, +) -> Result { + Ok(resolve_provider_config(model, custom_llm_provider)? + .1 + .get_health_check_document()) +} + +#[derive(Clone, Copy, Debug, EnumString, IntoStaticStr, PartialEq, Eq)] +#[strum(serialize_all = "snake_case")] pub(crate) enum OcrProvider { Cohere, Mistral, @@ -61,18 +137,6 @@ pub(crate) enum OcrProvider { VertexAi, } -impl OcrProvider { - pub(crate) const fn as_str(self) -> &'static str { - match self { - Self::Cohere => "cohere", - Self::Mistral => "mistral", - Self::AzureAi => "azure_ai", - Self::Reducto => "reducto", - Self::VertexAi => "vertex_ai", - } - } -} - pub(crate) fn resolve_provider_config( model: &str, custom_llm_provider: Option<&str>, @@ -80,30 +144,33 @@ pub(crate) fn resolve_provider_config( let provider = get_custom_llm_provider(model, custom_llm_provider).unwrap_or(CustomLlmProvider { model, - custom_llm_provider: OcrProvider::Mistral.as_str(), + custom_llm_provider: OcrProvider::Mistral.into(), }); - let config = match provider.custom_llm_provider { - "cohere" => OcrConfigKind::Cohere, - "mistral" => OcrConfigKind::Mistral, - "azure_ai" if is_document_intelligence_model(provider.model) => { + let ocr_provider = provider + .custom_llm_provider + .parse::() + .map_err(|_| super::Error::InvalidProvider(provider.custom_llm_provider.to_string()))?; + let config = match ocr_provider { + OcrProvider::Cohere => OcrConfigKind::Cohere, + OcrProvider::Mistral => OcrConfigKind::Mistral, + OcrProvider::AzureAi if is_document_intelligence_model(provider.model) => { OcrConfigKind::AzureDocumentIntelligence } - "azure_ai" + OcrProvider::AzureAi if provider.model.to_ascii_lowercase().contains("cohere") && provider.model.to_ascii_lowercase().contains("parse") => { OcrConfigKind::AzureCohere } - "azure_ai" => OcrConfigKind::AzureAi, - "reducto" if provider.model.eq_ignore_ascii_case("parse-legacy") => { + OcrProvider::AzureAi => OcrConfigKind::AzureAi, + OcrProvider::Reducto if provider.model.eq_ignore_ascii_case("parse-legacy") => { OcrConfigKind::ReductoLegacy } - "reducto" => OcrConfigKind::ReductoV3, - "vertex_ai" if provider.model.to_ascii_lowercase().contains("deepseek") => { + OcrProvider::Reducto => OcrConfigKind::ReductoV3, + OcrProvider::VertexAi if provider.model.to_ascii_lowercase().contains("deepseek") => { OcrConfigKind::VertexDeepSeek } - "vertex_ai" => OcrConfigKind::VertexAi, - value => return Err(super::Error::InvalidProvider(value.to_string())), + OcrProvider::VertexAi => OcrConfigKind::VertexAi, }; Ok((provider.model.to_string(), config)) } @@ -116,6 +183,149 @@ fn is_document_intelligence_model(model: &str) -> bool { #[cfg(test)] mod tests { use super::*; + use litellm_auth::InputSource; + + #[test] + fn provider_names_round_trip_exactly() { + for provider in ["cohere", "mistral", "azure_ai", "reducto", "vertex_ai"] { + let (_, config) = resolve_provider_config("model", Some(provider)).unwrap(); + let resolved: &'static str = config.provider().into(); + assert_eq!(resolved, provider); + } + + for provider in ["Mistral", "unknown"] { + assert_eq!( + resolve_provider_config("model", Some(provider)), + Err(crate::ocr::Error::InvalidProvider(provider.into())) + ); + } + } + + #[test] + fn health_check_documents_are_valid_for_each_provider() { + for model in [ + "mistral/ocr", + "azure_ai/ocr", + "azure_ai/doc-intelligence/prebuilt-layout", + "reducto/parse-v3", + "vertex_ai/mistral-ocr", + "vertex_ai/deepseek-ocr", + ] { + let document = get_health_check_document(model, None).unwrap(); + assert!(matches!(document, OcrDocument::DocumentUrl { .. })); + let inline = crate::ocr::document::InlineDocument::parse(document.source()) + .unwrap() + .unwrap(); + assert_eq!(inline.mime_type().to_string(), "application/pdf"); + assert!(inline.decode(4096).unwrap().starts_with(b"%PDF-")); + } + for model in ["cohere/parse", "azure_ai/cohere-parse"] { + let document = get_health_check_document(model, None).unwrap(); + crate::llms::cohere::ocr::validate_document(&document).unwrap(); + let inline = crate::ocr::document::InlineDocument::parse(document.source()) + .unwrap() + .unwrap(); + assert_eq!(inline.mime_type().to_string(), "image/png"); + assert!( + inline + .decode(4096) + .unwrap() + .starts_with(b"\x89PNG\r\n\x1a\n") + ); + } + } + + #[test] + fn api_key_metadata_follows_provider_overrides_and_python_defaults() { + for (model, expected) in [ + ("mistral/ocr", Some("MISTRAL_API_KEY")), + ("cohere/parse", Some("COHERE_API_KEY")), + ("azure_ai/ocr", Some("AZURE_AI_API_KEY")), + ("azure_ai/cohere-parse", Some("AZURE_AI_API_KEY")), + ( + "azure_ai/doc-intelligence/prebuilt-layout", + Some("AZURE_DOCUMENT_INTELLIGENCE_API_KEY"), + ), + ("vertex_ai/mistral-ocr", Some("VERTEX_AI_API_KEY")), + ("vertex_ai/deepseek-ocr", Some("VERTEX_AI_API_KEY")), + ("reducto/parse-v3", None), + ("reducto/parse-legacy", None), + ] { + assert_eq!( + get_api_key_env_var(model, None).unwrap(), + expected, + "{model}" + ); + } + } + + #[test] + fn connection_resolution_preserves_dynamic_precedence_and_input_sources() { + let connection = OcrConfigKind::Mistral.resolve_connection_params(OcrConnection { + api_key: Some("explicit-key".into()), + api_base: Some("https://explicit.test".into()), + dynamic_api_key: Some(Sourced::new("dynamic-key".into(), InputSource::Environment)), + dynamic_api_base: Some(Sourced::new( + "https://dynamic.test".into(), + InputSource::Request, + )), + ..Default::default() + }); + assert_eq!(connection.api_key.as_deref(), Some("dynamic-key")); + assert_eq!(connection.api_base.as_deref(), Some("https://dynamic.test")); + assert_eq!(connection.api_key_source, InputSource::Environment); + assert_eq!(connection.api_base_source, InputSource::Request); + for dynamic in [ + None, + Some(Sourced::new(String::new(), InputSource::Environment)), + ] { + let connection = OcrConfigKind::Mistral.resolve_connection_params(OcrConnection { + api_key: Some("explicit-key".into()), + api_base: Some("https://explicit.test".into()), + dynamic_api_key: dynamic.clone(), + dynamic_api_base: dynamic, + ..Default::default() + }); + assert_eq!(connection.api_key.as_deref(), Some("explicit-key")); + assert_eq!( + connection.api_base.as_deref(), + Some("https://explicit.test") + ); + } + } + + #[test] + fn document_intelligence_only_accepts_dynamic_values_for_explicit_fields() { + for (explicit_key, explicit_base) in [ + (None, None), + (Some("key"), None), + (None, Some("base")), + (Some("key"), Some("base")), + ] { + let connection = + OcrConfigKind::AzureDocumentIntelligence.resolve_connection_params(OcrConnection { + api_key: explicit_key.map(str::to_string), + api_base: explicit_base.map(str::to_string), + dynamic_api_key: Some(Sourced::new( + "dynamic-key".into(), + InputSource::Environment, + )), + dynamic_api_base: Some(Sourced::new( + "https://dynamic.test".into(), + InputSource::Deployment, + )), + ..Default::default() + }); + assert_eq!( + connection.api_key.as_deref(), + explicit_key.map(|_| "dynamic-key") + ); + assert_eq!( + connection.api_base.as_deref(), + explicit_base.map(|_| "https://dynamic.test") + ); + } + } #[test] fn provider_models_are_preserved_without_a_local_allowlist() { diff --git a/litellm-rust/crates/core/tests/ocr/support.rs b/litellm-rust/crates/core/src/ocr/test_support.rs similarity index 96% rename from litellm-rust/crates/core/tests/ocr/support.rs rename to litellm-rust/crates/core/src/ocr/test_support.rs index d6323c2124f..59f4acb8a5b 100644 --- a/litellm-rust/crates/core/tests/ocr/support.rs +++ b/litellm-rust/crates/core/src/ocr/test_support.rs @@ -37,13 +37,13 @@ pub(crate) fn wire_request(model: &str, base: &str, options: Value) -> LiteLLMOc } pub(crate) struct MockResponse { - pub status: u16, - pub headers: Vec<(&'static str, String)>, - pub body: Value, + pub(crate) status: u16, + pub(crate) headers: Vec<(&'static str, String)>, + pub(crate) body: Value, } impl MockResponse { - pub fn json(body: Value) -> Self { + pub(crate) fn json(body: Value) -> Self { Self { status: 200, headers: vec![], diff --git a/litellm-rust/crates/core/src/ocr/types.rs b/litellm-rust/crates/core/src/ocr/types.rs index f617bb7095a..9e2f99a0cf9 100644 --- a/litellm-rust/crates/core/src/ocr/types.rs +++ b/litellm-rust/crates/core/src/ocr/types.rs @@ -2,16 +2,17 @@ use std::collections::BTreeMap; use std::sync::Arc; use std::time::Duration; -use crate::serde_compat::{FiniteF64, LaxI64}; use serde::{Deserialize, Serialize}; use serde_json::{Map, Value}; use serde_with::serde_as; +use litellm_auth::{InputSource, Sourced, TokenProviderHandle}; + use super::hooks::{NoopOcrHooks, OcrHooks}; use super::provider_config::{OcrConfigKind, resolve_provider_config}; use crate::call_arguments::CallArguments; use crate::constants::OCR_HTTP_TIMEOUT_SECS; -use litellm_auth::{InputSource, TokenProviderHandle}; +use crate::serde_compat::{FiniteF64, LaxI64}; #[derive(Clone, Debug, PartialEq, Serialize, Deserialize)] #[serde(tag = "type")] @@ -63,8 +64,10 @@ pub enum OcrResponseFormat { #[derive(Clone)] pub struct OcrConnection { pub api_key: Option, + pub dynamic_api_key: Option>, pub api_key_source: InputSource, pub api_base: Option, + pub dynamic_api_base: Option>, pub api_base_source: InputSource, pub extra_headers: Vec<(String, String)>, pub extra_headers_source: InputSource, @@ -78,8 +81,10 @@ impl Default for OcrConnection { fn default() -> Self { Self { api_key: None, + dynamic_api_key: None, api_key_source: InputSource::Deployment, api_base: None, + dynamic_api_base: None, api_base_source: InputSource::Deployment, extra_headers: Vec::new(), extra_headers_source: InputSource::Deployment, @@ -137,7 +142,7 @@ impl LiteLLMOcrRequest { } pub fn provider_name(&self) -> &'static str { - self.config.provider().as_str() + self.config.provider().into() } pub fn with_host_hooks( diff --git a/litellm-rust/crates/core/src/ocr/wire.rs b/litellm-rust/crates/core/src/ocr/wire.rs index 996e5d3ecf4..ddf0c9e7d67 100644 --- a/litellm-rust/crates/core/src/ocr/wire.rs +++ b/litellm-rust/crates/core/src/ocr/wire.rs @@ -157,8 +157,10 @@ pub fn decode_request(wire: OcrWireRequest) -> Result bool { - true - } - - fn during_call( - &self, - mut request: OcrDuringCallRequest, - ) -> OcrHookFuture<'_, OcrDuringCallRequest> { - Box::pin(async move { - request.body["document"] = json!({ - "type":"document_url", - "document_url":"https://example.com/not-inline.pdf" - }); - Ok(request) - }) - } -} - -#[tokio::test] -async fn rejects_non_inline_body_after_guardrails() { - let mut request = wire_request("azure_ai/model", "http://127.0.0.1:1", json!({})); - request.hooks = Arc::new(ReplaceBodyDocument); - let error = perform_ocr(request).await.unwrap_err(); - assert!(error.to_string().contains("data URI")); -} diff --git a/litellm-rust/crates/core/tests/azure_document_intelligence_ocr.rs b/litellm-rust/crates/core/tests/azure_document_intelligence_ocr.rs deleted file mode 100644 index 880a8d625a4..00000000000 --- a/litellm-rust/crates/core/tests/azure_document_intelligence_ocr.rs +++ /dev/null @@ -1,460 +0,0 @@ -use serde_json::{Value, json}; -use std::sync::{Arc, Mutex}; - -use super::test_support::{MockResponse, mock_server, perform_ocr, wire_request}; -use super::wire::{OcrWireRequest, decode_request}; - -fn query_value(url: &str, key: &str) -> Option { - url::Url::parse(url) - .unwrap() - .query_pairs() - .find_map(|(name, value)| (name == key).then(|| value.into_owned())) -} - -#[tokio::test] -async fn facade_maps_pages_features_and_url_document() { - let (base, seen, server) = mock_server(vec![MockResponse::json(json!({ - "status":"succeeded", - "analyzeResult":{"pages":[]} - }))]) - .await; - let mut request = wire_request( - "azure_ai/doc-intelligence/prebuilt-read", - &base, - json!({"pages":[2,0,0,1],"features":["keyValuePairs","languages"], "future_option": {"nested":null}, "extra_body":{"provider_option":false}}), - ); - request.document = serde_json::from_value(json!({ - "type":"document_url", - "document_url":"https://example.com/document.pdf" - })) - .unwrap(); - - perform_ocr(request).await.unwrap(); - server.await.unwrap(); - let request = &seen.lock().unwrap()[0]; - let target = request.split_whitespace().nth(1).unwrap(); - let url = format!("{base}{target}"); - assert_eq!(query_value(&url, "pages").as_deref(), Some("1,2,3")); - assert_eq!( - query_value(&url, "features").as_deref(), - Some("keyValuePairs,languages") - ); - let body: Value = serde_json::from_str(request.split_once("\r\n\r\n").unwrap().1).unwrap(); - assert_eq!( - body, - json!({"urlSource":"https://example.com/document.pdf", "future_option":{"nested":null}, "provider_option":false}) - ); -} - -#[tokio::test] -async fn rejects_invalid_pages_features_and_format() { - for options in [ - json!({"pages":[true]}), - json!({"pages":[1,"2"]}), - json!({"pages":[-1]}), - json!({"pages":"1&&features=bad"}), - json!({"features":"languages&pages=1"}), - json!({"req_format":"azure"}), - ] { - let result = decode_request(OcrWireRequest { - model: "azure_ai/doc-intelligence/prebuilt-read".into(), - document: json!({"type":"document_url","document_url":"https://example.com/a.pdf"}), - api_key: Some("key".into()), - api_base: Some("http://127.0.0.1:1".into()), - custom_llm_provider: None, - extra_headers: None, - optional_params: options.as_object().unwrap().clone().into(), - input_sources: Default::default(), - timeout_seconds: None, - }); - let rejected = match result { - Ok(request) => perform_ocr(request).await.is_err(), - Err(_) => true, - }; - assert!(rejected, "accepted {options}"); - } -} - -#[tokio::test] -async fn inline_document_decodes_to_base64_source() { - let (base, seen, server) = mock_server(vec![MockResponse::json(json!({ - "status":"succeeded" - }))]) - .await; - let request = wire_request("azure_ai/doc-intelligence/prebuilt-read", &base, json!({})); - - perform_ocr(request).await.unwrap(); - server.await.unwrap(); - let request = &seen.lock().unwrap()[0]; - let body: Value = serde_json::from_str(request.split_once("\r\n\r\n").unwrap().1).unwrap(); - assert_eq!(body, json!({"base64Source":"YWJj"})); -} - -#[tokio::test] -async fn immediate_response_normalizes_pages_and_preserves_native() { - let operation = json!({ - "status":"succeeded", - "operationExtension":42, - "analyzeResult":{ - "content":"A\n\nB", - "tables":[{"cells":[]}], - "keyValuePairs":[{"key":{"content":"A"}}], - "pages":[{ - "pageNumber":"2", - "width":"8.5", - "height":11, - "unit":"inch", - "lines":[{"content":"A"},{"content":null},{"content":"B"}] - }] - } - }); - let (base, _, server) = mock_server(vec![MockResponse::json(operation.clone())]).await; - let result = perform_ocr(wire_request( - "azure_ai/doc-intelligence/prebuilt-read", - &base, - json!({"req_format":"native"}), - )) - .await - .unwrap(); - server.await.unwrap(); - - assert_eq!(result.pages[0].index, 1); - assert_eq!(result.pages[0].markdown, "A\n\nB"); - assert_eq!( - serde_json::to_value(&result.pages[0].dimensions).unwrap(), - json!({"width":816,"height":1056,"dpi":96}) - ); - assert_eq!(result.usage_info.as_ref().unwrap().pages_processed, Some(1)); - let serialized = result.clone().into_json(); - assert_eq!(serialized["content"], "A\n\nB"); - assert_eq!(serialized["tables"], json!([{"cells":[]}])); - assert_eq!( - serialized["keyValuePairs"], - json!([{"key":{"content":"A"}}]) - ); - assert!(serialized.get("key_value_pairs").is_none()); - assert_eq!( - result.provider_native_response.as_ref(), - operation.as_object() - ); -} - -#[tokio::test] -async fn accepted_response_polls_to_success_with_only_credentials() { - let operation = json!({"status":"succeeded","analyzeResult":{"pages":[]}}); - let (base, seen, server) = mock_server(vec![ - MockResponse { - status: 202, - headers: vec![("Operation-Location", "{base}/operation".into())], - body: json!({}), - }, - MockResponse { - status: 200, - headers: vec![("Retry-After", "0".into())], - body: json!({"status":"running"}), - }, - MockResponse::json(operation.clone()), - ]) - .await; - let mut request = wire_request( - "azure_ai/doc-intelligence/prebuilt-read", - &base, - json!({"req_format":"native"}), - ); - request - .connection - .extra_headers - .push(("X-Trace".into(), "initial-only".into())); - - let result = perform_ocr(request).await.unwrap(); - server.await.unwrap(); - assert_eq!( - result.provider_native_response.as_ref(), - operation.as_object() - ); - let requests = seen.lock().unwrap(); - assert_eq!(requests.len(), 3); - assert!(requests[0].to_ascii_lowercase().contains("x-trace:")); - for poll in &requests[1..] { - assert!(!poll.to_ascii_lowercase().contains("x-trace:")); - assert!( - poll.to_ascii_lowercase() - .contains("ocp-apim-subscription-key: test-key") - ); - } -} - -struct SubmissionBoundary { - request_count: Arc>>, - post_calls: Arc>>, -} - -impl super::hooks::OcrHooks for SubmissionBoundary { - fn post_call( - &self, - request: super::hooks::OcrPostCallRequest, - ) -> super::hooks::OcrHookFuture<'_, super::hooks::OcrPostCallRequest> { - Box::pin(async move { - assert_eq!(self.request_count.lock().unwrap().len(), 1); - self.post_calls - .lock() - .unwrap() - .push(request.original_response.clone()); - Ok(request) - }) - } -} - -#[tokio::test] -async fn accepted_response_runs_post_call_once_before_polling() { - let (base, seen, server) = mock_server(vec![ - MockResponse { - status: 202, - headers: vec![("Operation-Location", "{base}/operation".into())], - body: json!({"submitted": true}), - }, - MockResponse::json(json!({"status":"succeeded"})), - ]) - .await; - let post_calls = Arc::new(Mutex::new(Vec::new())); - let request = super::LiteLLMOcrRequest { - hooks: Arc::new(SubmissionBoundary { - request_count: seen.clone(), - post_calls: post_calls.clone(), - }), - ..wire_request("azure_ai/doc-intelligence/prebuilt-read", &base, json!({})) - }; - - perform_ocr(request).await.unwrap(); - server.await.unwrap(); - assert_eq!(seen.lock().unwrap().len(), 2); - assert_eq!( - *post_calls.lock().unwrap(), - [json!(r#"{"submitted":true}"#)] - ); -} - -#[tokio::test] -async fn polling_forwards_bearer_credentials() { - let (base, seen, server) = mock_server(vec![ - MockResponse { - status: 202, - headers: vec![("Operation-Location", "{base}/operation".into())], - body: json!({}), - }, - MockResponse::json(json!({"status":"succeeded"})), - ]) - .await; - let mut request = wire_request("azure_ai/doc-intelligence/prebuilt-read", &base, json!({})); - request.connection.api_key = None; - request.connection.extra_headers = vec![("Authorization".into(), "Bearer token".into())]; - - perform_ocr(request).await.unwrap(); - server.await.unwrap(); - let requests = seen.lock().unwrap(); - assert!( - requests[1] - .to_ascii_lowercase() - .contains("authorization: bearer token") - ); -} - -#[tokio::test] -async fn polling_does_not_follow_redirects() { - let (base, seen, server) = mock_server(vec![ - MockResponse { - status: 202, - headers: vec![("Operation-Location", "{base}/operation".into())], - body: json!({}), - }, - MockResponse { - status: 302, - headers: vec![("Location", "{base}/redirected".into())], - body: json!({}), - }, - MockResponse::json(json!({"status":"succeeded"})), - ]) - .await; - - let error = perform_ocr(wire_request( - "azure_ai/doc-intelligence/prebuilt-read", - &base, - json!({}), - )) - .await - .unwrap_err(); - - assert!(error.to_string().contains("status 302"), "{error}"); - assert_eq!(seen.lock().unwrap().len(), 2); - server.abort(); -} - -#[tokio::test] -async fn polling_rejects_terminal_failure() { - let (base, _, server) = mock_server(vec![ - MockResponse { - status: 202, - headers: vec![("Operation-Location", "{base}/operation".into())], - body: json!({}), - }, - MockResponse::json(json!({"status":"failed"})), - ]) - .await; - - let error = perform_ocr(wire_request( - "azure_ai/doc-intelligence/prebuilt-read", - &base, - json!({}), - )) - .await - .unwrap_err(); - server.await.unwrap(); - assert!(error.to_string().contains("status failed")); -} - -#[tokio::test] -async fn malformed_provider_pages_report_response_paths() { - for (analysis, path) in [ - (json!({"pages":null}), "pages"), - (json!({"pages":[null]}), "pages[0]"), - (json!({"pages":[{"lines":null}]}), "lines"), - (json!({"pages":[{"width":"bad"}]}), "width"), - ] { - let (base, _, server) = mock_server(vec![MockResponse::json(json!({ - "status":"succeeded", - "analyzeResult":analysis - }))]) - .await; - let error = perform_ocr(wire_request( - "azure_ai/doc-intelligence/prebuilt-read", - &base, - json!({}), - )) - .await - .unwrap_err(); - server.await.unwrap(); - assert!(error.to_string().contains(path), "{error}"); - } -} - -#[tokio::test] -async fn rejects_missing_invalid_and_cross_origin_operation_locations() { - for headers in [ - Vec::new(), - vec![("Operation-Location", "/relative".into())], - vec![("Operation-Location", "http://example.com/operation".into())], - vec![( - "Operation-Location", - "http://user:password@127.0.0.1/operation".into(), - )], - ] { - let (base, _, server) = mock_server(vec![MockResponse { - status: 202, - headers, - body: json!({}), - }]) - .await; - let error = perform_ocr(wire_request( - "azure_ai/doc-intelligence/prebuilt-read", - &base, - json!({}), - )) - .await - .unwrap_err(); - server.await.unwrap(); - assert!(error.to_string().contains("operation-location")); - } -} - -#[tokio::test] -async fn polling_deadline_bounds_retry_delay() { - let (base, _, server) = mock_server(vec![ - MockResponse { - status: 202, - headers: vec![("Operation-Location", "{base}/operation".into())], - body: json!({}), - }, - MockResponse { - status: 200, - headers: vec![("Retry-After", "9999".into())], - body: json!({"status":"notStarted"}), - }, - ]) - .await; - let mut request = wire_request("azure_ai/doc-intelligence/prebuilt-read", &base, json!({})); - request.connection.poll_timeout = std::time::Duration::from_millis(100); - - let error = tokio::time::timeout(std::time::Duration::from_secs(1), perform_ocr(request)) - .await - .unwrap() - .unwrap_err(); - server.await.unwrap(); - assert!(error.to_string().contains("timed out")); -} - -#[tokio::test] -async fn model_id_is_encoded_and_dot_segments_are_rejected() { - let (base, seen, server) = mock_server(vec![MockResponse::json(json!({ - "status":"succeeded" - }))]) - .await; - perform_ocr(wire_request( - "azure_ai/doc-intelligence/a ?#é", - &base, - json!({}), - )) - .await - .unwrap(); - server.await.unwrap(); - assert!(seen.lock().unwrap()[0].contains("a%20%3F%23%C3%A9:analyze")); - - for model in [ - "azure_ai/doc-intelligence/.", - "azure_ai/doc-intelligence/..", - ] { - let error = perform_ocr(wire_request(model, "http://127.0.0.1:1", json!({}))) - .await - .unwrap_err(); - assert!(error.to_string().contains("dot segment")); - } -} - -#[tokio::test] -async fn pre_call_guardrail_receives_caller_pages_before_mapping() { - use crate::ocr::hooks::{OcrHookFuture, OcrHooks, OcrPreCallRequest}; - use std::sync::Arc; - - struct RewritePages; - impl OcrHooks for RewritePages { - fn intercepts_requests(&self) -> bool { - true - } - - fn pre_call(&self, request: OcrPreCallRequest) -> OcrHookFuture<'_, OcrPreCallRequest> { - Box::pin(async move { - assert_eq!(request.optional_params["pages"], json!([0, 2])); - Ok(OcrPreCallRequest { - optional_params: json!({"pages": [1]}), - ..request - }) - }) - } - } - let (base, seen, server) = - mock_server(vec![MockResponse::json(json!({"status": "succeeded"}))]).await; - let request = wire_request( - "azure_ai/doc-intelligence/prebuilt-read", - &base, - json!({"pages": [0, 2]}), - ) - .with_host_hooks(Arc::new(RewritePages), None); - perform_ocr(request).await.unwrap(); - server.await.unwrap(); - let requests = seen.lock().unwrap(); - let target = requests[0].split_whitespace().nth(1).unwrap(); - assert_eq!( - query_value(&format!("{base}{target}"), "pages").as_deref(), - Some("2") - ); - assert_eq!(requests.len(), 1); -} diff --git a/litellm-rust/crates/core/tests/deepseek_ocr.rs b/litellm-rust/crates/core/tests/deepseek_ocr.rs deleted file mode 100644 index 37f89f43137..00000000000 --- a/litellm-rust/crates/core/tests/deepseek_ocr.rs +++ /dev/null @@ -1,158 +0,0 @@ -use rstest::rstest; -use serde_json::{Value, json}; - -use crate::llms::vertex_ai::ocr::deepseek_transformation::{ - DeepSeekOcrParams, DeepSeekOcrResponse, VertexAIDeepSeekOCRConfig, normalize_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"]))] -#[case("temperature", json!(null))] -fn request_mapping_matches_python(#[case] name: &str, #[case] value: Value) { - let params: DeepSeekOcrParams = - serde_json::from_value(json!({name: value.clone(), "ignored": true})).unwrap(); - let result = serde_json::to_value( - VertexAIDeepSeekOCRConfig - .transform_ocr_request("deepseek-ai/deepseek-ocr-maas", document(), ¶ms, &[]) - .unwrap(), - ) - .unwrap(); - assert_eq!(result["model"], "deepseek-ai/deepseek-ocr-maas"); - assert_eq!( - result["messages"][0]["content"][0], - json!({"type":"image_url","image_url":"gs://bucket/a.png"}) - ); - assert_eq!(result[name], value); - assert!(result.get("ignored").is_none()); -} - -#[rstest] -#[case(json!({"type":"image_url","image_url":"data:image/png;base64,AA=="}))] -#[case(json!({"type":"document_url","document_url":"data:application/pdf;base64,AA=="}))] -fn request_maps_both_document_types_to_image_content(#[case] document: Value) { - let source = document - .get("image_url") - .or_else(|| document.get("document_url")) - .unwrap() - .clone(); - let request = VertexAIDeepSeekOCRConfig - .transform_ocr_request( - "deepseek-ai/deepseek-ocr-maas", - serde_json::from_value(document).unwrap(), - &DeepSeekOcrParams::default(), - &[], - ) - .unwrap(); - let result = serde_json::to_value(request).unwrap(); - assert_eq!( - result["messages"][0]["content"][0], - json!({"type":"image_url","image_url":source}) - ); -} - -#[rstest] -#[case(json!("# hello"), "# hello")] -#[case(json!("{broken"), "{broken")] -#[case(json!(" {\"pages\":[]} "), " {\"pages\":[]} ")] -#[case(json!({"pages":[]}), "")] -#[case(json!("[]"), "[]")] -#[case(json!("{\"pages\":[{\"markdown\":\"json text\"}]}"), "json text")] -#[case(json!({"pages":[{"markdown":"object"}]}), "object")] -fn response_transform_handles_text_json_and_objects( - #[case] content: Value, - #[case] expected: &str, -) { - let has_pages = content - .as_object() - .is_some_and(|data| data.contains_key("pages")) - || content - .as_str() - .is_some_and(|text| text.contains("\"pages\"")); - let response: DeepSeekOcrResponse = serde_json::from_value( - json!({"choices":[{"message":{"content":content}}],"usage":{"prompt_tokens":1}}), - ) - .unwrap(); - let result = normalize_response("model", response).unwrap().into_json(); - assert_eq!(result["pages"][0]["markdown"], expected); - assert_eq!(result["pages"][0]["index"], 0); - if has_pages { - assert!(result["usage_info"].is_null()); - } else { - assert_eq!(result["usage_info"]["prompt_tokens"], 1); - } -} - -#[test] -fn structured_result_maps_pages_usage_model_and_annotation() { - let response: DeepSeekOcrResponse = serde_json::from_value(json!({ - "choices":[{"message":{"content":{ - "pages":[{"index":2,"markdown":"page","images":[{"id":"one"}],"dimensions":{"width":10}}], - "model":"provider-model", - "usage_info":{"pages_processed":1}, - "document_annotation":{"language":"en"}, - "future":"kept" - }}}] - })) - .unwrap(); - let result = normalize_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!(result.get("future").is_none()); -} - -#[test] -fn response_transform_rejects_missing_empty_and_malformed_content() { - for value in [ - json!({"choices":[]}), - json!({"choices":[{"message":{"content":{}}}]}), - 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| normalize_response("model", response).map_err(|_| ())); - assert!(result.is_err()); - } -} - -#[test] -fn structured_content_preserves_usage_presence_and_shared_page_defaults() { - for (usage, expected) in [(json!(null), None), (json!({"pages_processed":2}), Some(2))] { - let response = serde_json::from_value(json!({ - "choices":[{"message":{"content":{ - "pages":[42, {"index":"2", "images":[{"id":"kept"}], "ignored":true}], - "usage_info":usage - }}}], - "usage":{"pages_processed":99} - })) - .unwrap(); - let normalized = normalize_response("model", response).unwrap(); - assert_eq!(normalized.pages.len(), 1); - assert_eq!(normalized.pages[0].index, 2); - assert_eq!(normalized.pages[0].markdown, ""); - assert!(normalized.pages[0].extra_fields.is_empty()); - assert_eq!( - normalized - .usage_info - .and_then(|usage| usage.pages_processed), - expected - ); - } -} diff --git a/litellm-rust/crates/core/tests/host_lifecycle.rs b/litellm-rust/crates/core/tests/host_lifecycle.rs deleted file mode 100644 index a07979222a2..00000000000 --- a/litellm-rust/crates/core/tests/host_lifecycle.rs +++ /dev/null @@ -1,116 +0,0 @@ -use crate::call_lifecycle::host::{HostFailure, HostLifecycle, HostPhase}; - -fn run(fail_at: Option, asynchronous: bool) -> (Vec, Vec) { - let mut lifecycle = HostLifecycle::new(asynchronous); - let mut events = Vec::new(); - let mut failures = Vec::new(); - - while lifecycle.phase() != HostPhase::Complete { - let phase = lifecycle.phase(); - events.push(phase); - let result = if Some(phase) == fail_at { - Err(HostFailure::Error(crate::ocr::Error::InvalidRequest( - "selected failure".into(), - ))) - } else { - Ok(()) - }; - if let Some(error) = lifecycle.accept(result) { - failures.push(error); - } - } - (events, failures) -} - -#[test] -fn public_outcome_is_finalized_before_a_single_terminal_dispatch() { - for asynchronous in [false, true] { - let (events, failures) = run(None, asynchronous); - assert!(failures.is_empty()); - assert_eq!( - &events[events.len() - 2..], - &[HostPhase::Finalize, HostPhase::Success] - ); - assert_eq!( - events - .iter() - .filter(|phase| **phase == HostPhase::Execute) - .count(), - 1 - ); - assert_eq!( - events.contains(&HostPhase::DeploymentPostCall), - asynchronous - ); - } -} - -#[test] -fn only_provider_and_response_construction_failures_use_provider_mapping() { - for phase in [ - HostPhase::Setup, - HostPhase::DeploymentPreCall, - HostPhase::Prepare, - HostPhase::Execute, - HostPhase::ConstructResponse, - HostPhase::DeploymentPostCall, - HostPhase::Finalize, - ] { - let (events, failures) = run(Some(phase), true); - assert_eq!(failures.len(), 1); - assert!(!events.contains(&HostPhase::Success)); - let mapped = matches!(phase, HostPhase::Execute | HostPhase::ConstructResponse); - assert_eq!(events.contains(&HostPhase::MapFailure), mapped); - assert_eq!(events.contains(&HostPhase::DeploymentFailure), mapped); - assert_eq!( - &events[events.len() - 2..], - &[HostPhase::Failure, HostPhase::AsyncFailure] - ); - assert!( - events - .iter() - .filter(|phase| **phase == HostPhase::Execute) - .count() - <= 1 - ); - } -} - -#[test] -fn failure_handler_errors_do_not_replace_selected_failure_or_suppress_async_dispatch() { - let mut lifecycle = HostLifecycle::new(true); - while lifecycle.phase() != HostPhase::Execute { - lifecycle.accept::(Ok(())); - } - let selected = crate::ocr::Error::InvalidRequest("provider".into()); - assert_eq!( - lifecycle.accept(Err(HostFailure::Error(selected.clone()))), - Some(selected) - ); - lifecycle.accept::(Ok(())); - for phase in [ - HostPhase::DeploymentFailure, - HostPhase::Failure, - HostPhase::AsyncFailure, - ] { - assert_eq!(lifecycle.phase(), phase); - assert_eq!( - lifecycle.accept(Err(HostFailure::Error(crate::ocr::Error::InvalidRequest( - "callback".into() - )))), - None - ); - } - assert_eq!(lifecycle.phase(), HostPhase::Complete); -} - -#[test] -fn cancellation_skips_terminal_dispatch() { - let mut lifecycle = HostLifecycle::new(true); - let error = crate::ocr::Error::InvalidRequest("cancelled".into()); - assert_eq!( - lifecycle.accept(Err(HostFailure::Cancelled(error.clone()))), - Some(error) - ); - assert_eq!(lifecycle.phase(), HostPhase::Complete); -} diff --git a/litellm-rust/crates/core/tests/ocr.rs b/litellm-rust/crates/core/tests/ocr.rs deleted file mode 100644 index 69d62908a06..00000000000 --- a/litellm-rust/crates/core/tests/ocr.rs +++ /dev/null @@ -1,895 +0,0 @@ -use std::sync::{Arc, Mutex}; - -use serde_json::{Value, json}; - -use super::OcrClient; -use super::hooks::{ - OcrDuringCallRequest, OcrHookFuture, OcrHooks, OcrLogFuture, OcrPostCallRequest, - OcrPreCallRequest, -}; -use super::test_support::{MockResponse, mock_server, perform_ocr, wire_request}; -use super::wire::{OcrWireRequest, decode_request}; -use super::{ - NativeOutcome, NoopOcrHost, OcrAdmission, OcrCall, OcrCallStep, OcrDecline, OcrHost, - OcrHostOperation, OcrHostResult, -}; -use crate::call_lifecycle::{CallLifecycleContext, CallLifecycleTiming}; - -#[test] -fn request_boundary_selects_mistral_and_rejects_unknown_providers() { - let request = OcrWireRequest { - model: "mistral/model".into(), - document: json!({"type":"document_url","document_url":"https://example.com/doc.pdf"}), - api_key: Some("key".into()), - api_base: None, - custom_llm_provider: None, - extra_headers: None, - optional_params: json!({"extract_header":true,"unknown":42}) - .as_object() - .unwrap() - .clone() - .into(), - input_sources: Default::default(), - timeout_seconds: None, - }; - assert!(decode_request(request).is_ok()); - assert!( - decode_request(OcrWireRequest { - model: "model".into(), - document: json!({"type":"document_url","document_url":"https://example.com/doc.pdf"}), - api_key: Some("key".into()), - api_base: None, - custom_llm_provider: Some("unknown".into()), - extra_headers: None, - optional_params: Default::default(), - input_sources: Default::default(), - timeout_seconds: None, - }) - .is_err() - ); -} - -#[tokio::test] -async fn facade_executes_direct_mistral_once() { - let (base, seen, server) = mock_server(vec![MockResponse::json(json!({ - "pages":[{"index":0,"markdown":"hello","custom":"preserved"}], - "usage_info":{"pages_processed":1} - }))]) - .await; - let result = perform_ocr(wire_request( - "mistral/model", - &base, - json!({"pages":"0,2-4","extract_header":true,"unknown":{"nested":[null,false,0]}}), - )) - .await - .unwrap(); - server.await.unwrap(); - assert_eq!(result.pages[0].markdown, "hello"); - assert_eq!(result.pages[0].extra_fields["custom"], "preserved"); - let requests = seen.lock().unwrap(); - assert_eq!(requests.len(), 1); - assert!(requests[0].starts_with("POST /v1/ocr ")); - assert!( - requests[0] - .to_ascii_lowercase() - .contains("authorization: bearer test-key\r\n") - ); - let body: Value = serde_json::from_str(requests[0].split_once("\r\n\r\n").unwrap().1).unwrap(); - assert_eq!( - body, - json!({ - "model":"model", - "document":{"type":"document_url","document_url":"data:application/pdf;base64,YWJj"}, - "pages":"0,2-4", - "extract_header":true, - "unknown":{"nested":[null,false,0]} - }) - ); -} - -#[tokio::test] -async fn facade_retains_native_response_when_requested() { - let provider_response = json!({ - "pages":[{"index":0,"markdown":"hello"}], - "usage_info":{"pages_processed":1}, - "provider_only":"preserved" - }); - let (base, _, server) = mock_server(vec![MockResponse::json(provider_response.clone())]).await; - let response = perform_ocr(wire_request( - "mistral/model", - &base, - json!({"req_format":"native"}), - )) - .await - .unwrap(); - - server.await.unwrap(); - assert_eq!( - response.provider_native_response.as_ref(), - provider_response.as_object() - ); -} - -#[tokio::test] -async fn facade_uses_the_injected_http_client() { - let (base, seen, server) = mock_server(vec![MockResponse::json(json!({"pages":[]}))]).await; - let mut default_headers = reqwest::header::HeaderMap::new(); - default_headers.insert( - "x-transport-owner", - reqwest::header::HeaderValue::from_static("host"), - ); - let provider_http = reqwest::Client::builder() - .default_headers(default_headers) - .build() - .unwrap(); - OcrClient::new(provider_http) - .unwrap() - .perform(wire_request("mistral/model", &base, json!({}))) - .await - .unwrap(); - server.await.unwrap(); - assert!(seen.lock().unwrap()[0].contains("x-transport-owner: host")); -} - -struct RecordingHooks { - events: Arc>>, - block: bool, -} - -struct ExtensionHooks; - -impl OcrHooks for ExtensionHooks { - fn intercepts_requests(&self) -> bool { - true - } - - fn during_call( - &self, - mut request: OcrDuringCallRequest, - ) -> OcrHookFuture<'_, OcrDuringCallRequest> { - Box::pin(async move { - assert_eq!(request.body["pages"], json!([2])); - assert_eq!(request.body.get("future"), Some(&Value::Null)); - assert!( - !request - .retained_fields - .iter() - .any(|field| field == "pages" || field == "document") - ); - request.body.as_object_mut().unwrap().remove("future"); - request.body["hook_option"] = json!({"nested":[null,false,0]}); - Ok(request) - }) - } -} - -#[tokio::test] -async fn composed_extensions_reach_hooks_and_removed_fields_stay_removed() { - let (base, seen, server) = mock_server(vec![MockResponse::json(json!({"pages":[]}))]).await; - let request = super::LiteLLMOcrRequest { - hooks: Arc::new(ExtensionHooks), - ..wire_request( - "mistral/model", - &base, - json!({ - "pages":[0], "future":null, "extra_body":{"pages":[2], - "document":{"type":"document_url","document_url":"data:application/pdf;base64,eHl6"}} - }), - ) - }; - perform_ocr(request).await.unwrap(); - server.await.unwrap(); - let requests = seen.lock().unwrap(); - let body: Value = serde_json::from_str(requests[0].split_once("\r\n\r\n").unwrap().1).unwrap(); - assert_eq!(body["pages"], json!([2])); - assert_eq!( - body["document"]["document_url"], - "data:application/pdf;base64,eHl6" - ); - assert_eq!(body["hook_option"], json!({"nested":[null,false,0]})); - assert!(body.get("future").is_none()); -} - -impl OcrHooks for RecordingHooks { - fn intercepts_requests(&self) -> bool { - true - } - - fn pre_call(&self, request: OcrPreCallRequest) -> OcrHookFuture<'_, OcrPreCallRequest> { - Box::pin(async move { - self.events.lock().unwrap().push("pre"); - if self.block { - return Err(crate::ocr::Error::InvalidRequest("blocked".into())); - } - Ok(request) - }) - } - - fn during_call( - &self, - request: super::hooks::OcrDuringCallRequest, - ) -> OcrHookFuture<'_, super::hooks::OcrDuringCallRequest> { - Box::pin(async move { - self.events.lock().unwrap().push("during"); - Ok(request) - }) - } - - fn post_call(&self, request: OcrPostCallRequest) -> OcrHookFuture<'_, OcrPostCallRequest> { - Box::pin(async move { - self.events.lock().unwrap().push("post"); - Ok(request) - }) - } - - fn success<'a>( - &'a self, - _context: &'a CallLifecycleContext, - _response: &'a super::LiteLLMOcrResponse, - _timing: &'a CallLifecycleTiming, - ) -> OcrLogFuture<'a> { - Box::pin(async move { - self.events.lock().unwrap().push("success"); - }) - } - - fn failure<'a>( - &'a self, - _context: &'a CallLifecycleContext, - _error: &'a crate::ocr::Error, - _timing: &'a CallLifecycleTiming, - ) -> OcrLogFuture<'a> { - Box::pin(async move { - self.events.lock().unwrap().push("failure"); - }) - } -} - -struct HeaderEditHooks; - -impl OcrHooks for HeaderEditHooks { - fn intercepts_requests(&self) -> bool { - true - } - - fn during_call( - &self, - mut request: OcrDuringCallRequest, - ) -> OcrHookFuture<'_, OcrDuringCallRequest> { - request - .headers - .push(("x-core-callback".into(), "edited".into())); - Box::pin(async move { Ok(request) }) - } -} - -#[tokio::test] -async fn lifecycle_sends_headers_returned_by_the_typed_during_call_operation() { - let (base, seen, server) = mock_server(vec![MockResponse::json(json!({"pages":[]}))]).await; - let request = super::LiteLLMOcrRequest { - hooks: Arc::new(HeaderEditHooks), - ..wire_request("mistral/model", &base, json!({})) - }; - - perform_ocr(request).await.unwrap(); - server.await.unwrap(); - - assert!(seen.lock().unwrap()[0].contains("x-core-callback: edited")); -} - -#[tokio::test] -async fn lifecycle_orders_hooks_and_emits_one_success() { - let (base, seen, server) = mock_server(vec![MockResponse::json(json!({"pages":[]}))]).await; - let events = Arc::new(Mutex::new(Vec::new())); - let request = wire_request("mistral/model", &base, json!({})); - let request = super::LiteLLMOcrRequest { - hooks: Arc::new(RecordingHooks { - events: events.clone(), - block: false, - }), - ..request - }; - perform_ocr(request).await.unwrap(); - server.await.unwrap(); - assert_eq!( - *events.lock().unwrap(), - ["pre", "during", "post", "success"] - ); - assert_eq!(seen.lock().unwrap().len(), 1); -} - -#[tokio::test] -async fn lifecycle_blocking_prevents_execution_and_emits_one_failure() { - let events = Arc::new(Mutex::new(Vec::new())); - let request = wire_request("mistral/model", "http://127.0.0.1:1", json!({})); - let request = super::LiteLLMOcrRequest { - hooks: Arc::new(RecordingHooks { - events: events.clone(), - block: true, - }), - ..request - }; - let error = perform_ocr(request).await.unwrap_err(); - assert!(matches!(error, crate::ocr::Error::InvalidRequest(_))); - assert_eq!(*events.lock().unwrap(), ["pre", "failure"]); -} - -#[tokio::test] -async fn upstream_failure_emits_one_terminal_failure() { - let (base, seen, server) = mock_server(vec![MockResponse { - status: 500, - headers: vec![], - body: json!({"error":"failed"}), - }]) - .await; - let events = Arc::new(Mutex::new(Vec::new())); - let request = wire_request("mistral/model", &base, json!({})); - let request = super::LiteLLMOcrRequest { - hooks: Arc::new(RecordingHooks { - events: events.clone(), - block: false, - }), - ..request - }; - assert!(perform_ocr(request).await.is_err()); - server.await.unwrap(); - assert_eq!(*events.lock().unwrap(), ["pre", "during", "failure"]); - assert_eq!(seen.lock().unwrap().len(), 1); -} - -struct AdmissionSpy { - effects: Arc>, -} - -impl OcrHooks for AdmissionSpy { - fn intercepts_requests(&self) -> bool { - *self.effects.lock().unwrap() += 1; - true - } - - fn pre_call(&self, request: OcrPreCallRequest) -> OcrHookFuture<'_, OcrPreCallRequest> { - *self.effects.lock().unwrap() += 1; - Box::pin(async move { Ok(request) }) - } -} - -#[test] -fn admission_declines_without_invoking_hooks_or_transport() { - for (admission, expected) in [ - ( - OcrAdmission { - provider_workflow: false, - host_operations: true, - asynchronous: false, - }, - OcrDecline::ProviderWorkflow, - ), - ( - OcrAdmission { - provider_workflow: true, - host_operations: false, - asynchronous: false, - }, - OcrDecline::HostOperations, - ), - ] { - let outcome = OcrCall::admit(super::test_support::ocr_client(), admission); - assert!(matches!(outcome, NativeOutcome::Declined(reason) if reason == expected)); - } -} - -#[tokio::test] -async fn fallible_host_phases_do_not_replay_or_reach_transport() { - for failure_phase in ["pre", "during"] { - let request = super::LiteLLMOcrRequest { - hooks: Arc::new(AdmissionSpy { - effects: Arc::new(Mutex::new(0)), - }), - ..wire_request("mistral/model", "http://127.0.0.1:1", json!({})) - }; - let NativeOutcome::Completed(mut call) = - OcrCall::admit(super::test_support::ocr_client(), OcrAdmission::all()) - else { - panic!("supported call declined") - }; - let mut request = Some(request); - let mut result = None; - let mut phases = Vec::new(); - let error = loop { - match call.resume(result.take()).await { - Ok(OcrCallStep::Host(operation)) => match operation { - OcrHostOperation::Lifecycle(_) - | OcrHostOperation::ConstructResponse(_) - | OcrHostOperation::MapFailure(_) - | OcrHostOperation::Success { .. } - | OcrHostOperation::Failure { .. } => { - result = Some(OcrHostResult::Lifecycle(Ok(()))) - } - OcrHostOperation::ProjectRequest => { - result = Some(OcrHostResult::Request(Ok(( - Box::new(request.take().unwrap()), - false, - )))) - } - OcrHostOperation::AcquireAzureAdToken => { - panic!("test request has no token provider") - } - OcrHostOperation::PreCall(request) => { - phases.push("pre"); - result = Some(OcrHostResult::PreCall(if failure_phase == "pre" { - Err(crate::ocr::Error::InvalidRequest("pre failed".into())) - } else { - Ok(request) - })); - } - OcrHostOperation::DuringCall(request) => { - phases.push("during"); - result = Some(OcrHostResult::DuringCall(if failure_phase == "during" { - Err(crate::ocr::Error::InvalidRequest("during failed".into())) - } else { - Ok(request) - })); - } - OcrHostOperation::PostCall(_) => panic!("transport should not be reached"), - }, - Err(error) => break error, - Ok(OcrCallStep::Complete(_)) => panic!("failed call completed"), - } - }; - assert!(matches!(error, crate::ocr::Error::InvalidRequest(_))); - assert_eq!( - phases - .iter() - .filter(|phase| **phase == failure_phase) - .count(), - 1 - ); - } -} - -#[tokio::test] -async fn invalid_provider_response_runs_post_call_before_normalization_failure() { - let (base, seen, server) = - mock_server(vec![MockResponse::json(json!({"pages":"invalid"}))]).await; - let mut request = Some(wire_request("mistral/model", &base, json!({}))); - let NativeOutcome::Completed(mut call) = - OcrCall::admit(super::test_support::ocr_client(), OcrAdmission::all()) - else { - panic!("supported call declined") - }; - let host = NoopOcrHost; - let mut result = None; - let mut post_calls = Vec::new(); - let error = loop { - match call.resume(result.take()).await { - Ok(OcrCallStep::Host(OcrHostOperation::ProjectRequest)) => { - result = Some(OcrHostResult::Request(Ok(( - Box::new(request.take().unwrap()), - false, - )))); - } - Ok(OcrCallStep::Host(operation)) => { - if let OcrHostOperation::PostCall(request) = &operation { - post_calls.push(request.original_response.clone()); - } - result = Some(host.invoke(operation).await); - } - Err(error) => break error, - Ok(OcrCallStep::Complete(_)) => panic!("invalid provider response completed"), - } - }; - server.await.unwrap(); - assert!(matches!( - error, - crate::ocr::Error::ResponseField { ref path } if path == "pages" - )); - assert_eq!(seen.lock().unwrap().len(), 1); - assert_eq!(post_calls, [json!(r#"{"pages":"invalid"}"#)]); -} - -#[tokio::test] -async fn direct_native_host_drives_the_same_state_machine() { - let (base, seen, server) = mock_server(vec![MockResponse::json(json!({ - "pages":[{"index":0,"markdown":"native"}] - }))]) - .await; - let request = super::LiteLLMOcrRequest { - hooks: Arc::new(AdmissionSpy { - effects: Arc::new(Mutex::new(0)), - }), - ..wire_request("mistral/model", &base, json!({})) - }; - let NativeOutcome::Completed(mut call) = OcrCall::admit( - super::test_support::ocr_client(), - OcrAdmission { - asynchronous: true, - ..OcrAdmission::all() - }, - ) else { - panic!("supported call declined") - }; - let mut request = Some(request); - let host = NoopOcrHost; - let mut result = None; - let mut operations = Vec::new(); - let response = loop { - match call.resume(result.take()).await.unwrap() { - OcrCallStep::Host(operation) => { - operations.push(match &operation { - OcrHostOperation::ProjectRequest => "ProjectRequest".into(), - OcrHostOperation::Lifecycle(phase) => format!("{phase:?}"), - OcrHostOperation::PreCall(_) => "PreCall".into(), - OcrHostOperation::DuringCall(_) => "DuringCall".into(), - OcrHostOperation::PostCall(_) => "PostCall".into(), - OcrHostOperation::ConstructResponse(_) => "ConstructResponse".into(), - OcrHostOperation::Success { response, .. } => { - assert_eq!(response.pages[0].markdown, "native"); - "Success".into() - } - _ => panic!("unexpected OCR operation"), - }); - result = Some(match operation { - OcrHostOperation::ProjectRequest => { - OcrHostResult::Request(Ok((Box::new(request.take().unwrap()), false))) - } - operation => host.invoke(operation).await, - }); - } - OcrCallStep::Complete(response) => break response, - } - }; - server.await.unwrap(); - assert_eq!(response.pages[0].markdown, "native"); - assert_eq!(seen.lock().unwrap().len(), 1); - assert_eq!( - operations, - [ - "Setup", - "DeploymentPreCall", - "Prepare", - "ProjectRequest", - "PreCall", - "DuringCall", - "PostCall", - "ConstructResponse", - "DeploymentPostCall", - "Finalize", - "Success", - ] - ); - assert!(matches!( - call.resume(None).await, - Err(crate::ocr::Error::InvalidRequest(_)) - )); -} - -#[tokio::test] -async fn public_finalization_failure_never_dispatches_success_or_replays_provider() { - use crate::call_lifecycle::host::{HostFailure, HostPhase}; - - let (base, seen, server) = mock_server(vec![MockResponse::json(json!({"pages":[]}))]).await; - let mut request = Some(wire_request("mistral/model", &base, json!({}))); - let NativeOutcome::Completed(mut call) = OcrCall::admit( - super::test_support::ocr_client(), - OcrAdmission { - asynchronous: true, - ..OcrAdmission::all() - }, - ) else { - panic!("supported call declined") - }; - let selected = crate::ocr::Error::InvalidRequest("public metadata failed".into()); - let host = NoopOcrHost; - let mut result = None; - let mut failures = Vec::new(); - let error = loop { - match call.resume(result.take()).await { - Ok(OcrCallStep::Host(operation)) => { - result = Some(match operation { - OcrHostOperation::Lifecycle(HostPhase::Finalize) => { - OcrHostResult::Lifecycle(Err(HostFailure::Error(selected.clone()))) - } - OcrHostOperation::Failure { error, .. } => { - assert_eq!(error, selected); - failures.push("sync"); - OcrHostResult::Lifecycle(Err(HostFailure::Error( - crate::ocr::Error::InvalidRequest("failure callback failed".into()), - ))) - } - OcrHostOperation::Lifecycle(HostPhase::AsyncFailure) => { - failures.push("async"); - OcrHostResult::Lifecycle(Ok(())) - } - OcrHostOperation::Success { .. } - | OcrHostOperation::MapFailure(_) - | OcrHostOperation::Lifecycle(HostPhase::DeploymentFailure) => { - panic!("finalization failure used provider/success dispatch") - } - OcrHostOperation::ProjectRequest => { - OcrHostResult::Request(Ok((Box::new(request.take().unwrap()), false))) - } - operation => host.invoke(operation).await, - }); - } - Ok(OcrCallStep::Complete(_)) => panic!("failed call completed successfully"), - Err(error) => break error, - } - }; - server.await.unwrap(); - assert_eq!(error, selected); - assert_eq!(failures, ["sync", "async"]); - assert_eq!(seen.lock().unwrap().len(), 1); -} - -#[tokio::test] -async fn cancellation_at_provider_hook_prevents_execution_and_further_resumption() { - use crate::call_lifecycle::host::HostFailure; - - let request = super::LiteLLMOcrRequest { - hooks: Arc::new(AdmissionSpy { - effects: Arc::new(Mutex::new(0)), - }), - ..wire_request("mistral/model", "http://127.0.0.1:1", json!({})) - }; - let NativeOutcome::Completed(mut call) = - OcrCall::admit(super::test_support::ocr_client(), OcrAdmission::all()) - else { - panic!("supported call declined") - }; - let mut request = Some(request); - let host = NoopOcrHost; - let mut result = None; - loop { - match call.resume(result.take()).await.unwrap() { - OcrCallStep::Host(OcrHostOperation::PreCall(_)) => break, - OcrCallStep::Host(OcrHostOperation::ProjectRequest) => { - result = Some(OcrHostResult::Request(Ok(( - Box::new(request.take().unwrap()), - false, - )))) - } - OcrCallStep::Host(operation) => result = Some(host.invoke(operation).await), - OcrCallStep::Complete(_) => panic!("provider executed before pre-call result"), - } - } - let selected = crate::ocr::Error::InvalidRequest("cancelled".into()); - assert!(matches!( - call.interrupt(HostFailure::Cancelled(selected.clone())).await, - Err(error) if error == selected - )); - assert!( - call.resume(Some(OcrHostResult::Lifecycle(Ok(())))) - .await - .is_err() - ); -} - -#[tokio::test] -async fn missing_host_result_preserves_pending_operation() { - use crate::call_lifecycle::host::HostPhase; - - let NativeOutcome::Completed(mut call) = - OcrCall::admit(super::test_support::ocr_client(), OcrAdmission::all()) - else { - panic!("supported call declined") - }; - assert!(matches!( - call.resume(None).await.unwrap(), - OcrCallStep::Host(OcrHostOperation::Lifecycle(HostPhase::Setup)) - )); - assert!(call.resume(None).await.is_err()); - assert!(matches!( - call.resume(Some(OcrHostResult::Lifecycle(Ok(())))) - .await - .unwrap(), - OcrCallStep::Host(OcrHostOperation::Lifecycle(HostPhase::Prepare)) - )); -} - -async fn read_bounded_response( - response: Vec, - limit: usize, -) -> Result { - use tokio::io::{AsyncReadExt, AsyncWriteExt}; - - let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); - let address = listener.local_addr().unwrap(); - let server = tokio::spawn(async move { - let (mut socket, _) = listener.accept().await.unwrap(); - let mut request = [0; 4096]; - assert!(socket.read(&mut request).await.unwrap() > 0); - socket.write_all(&response).await.unwrap(); - std::future::pending::<()>().await; - }); - let response = reqwest::Client::new() - .get(format!("http://{address}")) - .send() - .await - .unwrap(); - let result = tokio::time::timeout( - std::time::Duration::from_secs(2), - super::client::read_response_bytes(response, limit), - ) - .await; - server.abort(); - let _ = server.await; - result.expect("bounded reads must finish without waiting for the rest of an oversized body") -} - -#[tokio::test] -async fn response_limit_accepts_exact_size_and_rejects_declared_and_chunked_overflow() { - for response in [ - "HTTP/1.1 200 OK\r\nContent-Length: 8\r\n\r\nabcdefgh", - "HTTP/1.1 200 OK\r\nTransfer-Encoding: chunked\r\n\r\n4\r\nabcd\r\n4\r\nefgh\r\n0\r\n\r\n", - ] { - assert_eq!( - read_bounded_response(response.as_bytes().to_vec(), 8) - .await - .unwrap(), - "abcdefgh" - ); - } - for response in [ - "HTTP/1.1 200 OK\r\nContent-Length: 9\r\n\r\n", - "HTTP/1.1 200 OK\r\nTransfer-Encoding: chunked\r\n\r\n4\r\nabcd\r\n5\r\nefghi\r\n", - ] { - assert!(matches!( - read_bounded_response(response.as_bytes().to_vec(), 8).await, - Err(crate::ocr::Error::TooLarge { limit: 8 }) - )); - } -} - -#[tokio::test] -async fn oversized_error_retains_http_status_and_bounded_diagnostics_without_draining() { - let prefix = "x".repeat(4 * (crate::constants::UPSTREAM_ERROR_BODY_MAX_CHARS + 1)); - for headers in ["Content-Length: 1000000", "Transfer-Encoding: chunked"] { - let body = if headers.starts_with("Transfer") { - format!("{:x}\r\n{prefix}\r\n", prefix.len()) - } else { - prefix.clone() - }; - let response = format!("HTTP/1.1 429 Too Many Requests\r\n{headers}\r\n\r\n{body}"); - let error = read_bounded_response(response.into_bytes(), 4096) - .await - .unwrap_err(); - match error { - crate::ocr::Error::Transport(crate::transport::Error::Http { status, body }) => { - assert_eq!(status, 429); - assert_eq!( - body, - format!( - "{}... (truncated)", - "x".repeat(crate::constants::UPSTREAM_ERROR_BODY_MAX_CHARS) - ) - ); - } - error => panic!("unexpected error: {error}"), - } - } -} - -#[test] -fn response_limit_is_validated_and_not_forwarded_to_the_provider() { - let request = wire_request( - "mistral/model", - "http://localhost", - json!({"max_response_bytes": 123}), - ); - assert_eq!(request.connection.max_response_bytes, 123); - assert!(!request.optional_params.contains_key("max_response_bytes")); - for value in [ - json!(0), - json!(-1), - json!(true), - json!("123"), - json!(1.5), - json!(crate::constants::OCR_RESPONSE_MAX_BYTES + 1), - Value::Null, - ] { - let wire = serde_json::from_value(json!({ - "model": "mistral/model", "document": {"type": "document_url", "document_url": "data:application/pdf;base64,YWJj"}, - "optional_params": {"max_response_bytes": value} - })).unwrap(); - let Err(error) = decode_request(wire) else { - panic!("invalid response limit accepted") - }; - assert!(error.to_string().contains("max_response_bytes")); - } -} - -#[derive(Debug)] -struct PendingToken { - entered: Arc, - dropped: Arc, -} - -struct TokenFutureDrop(Arc); - -impl Drop for TokenFutureDrop { - fn drop(&mut self) { - self.0.store(true, std::sync::atomic::Ordering::SeqCst); - } -} - -impl litellm_auth::TokenProvider for PendingToken { - fn acquire(&self) -> litellm_auth::TokenFuture<'_> { - Box::pin(async move { - let _guard = TokenFutureDrop(self.dropped.clone()); - self.entered.notify_one(); - std::future::pending().await - }) - } -} - -#[tokio::test] -async fn cancellation_waits_for_provider_capture_drop_even_when_acknowledgement_is_cancelled() { - use crate::call_lifecycle::host::HostFailure; - use std::future::Future; - use std::sync::atomic::{AtomicBool, Ordering}; - use std::task::Poll; - - for interrupt_acknowledgement in [false, true] { - let entered = Arc::new(tokio::sync::Notify::new()); - let dropped = Arc::new(AtomicBool::new(false)); - let request = wire_request("azure_ai/mistral-ocr", "https://example.invalid", json!({})); - let request = super::LiteLLMOcrRequest { - connection: super::OcrConnection { - extra_headers: vec![("authorization".into(), "Bearer test-key".into())], - ..request.connection - }, - azure_ad_token_provider: Some(litellm_auth::TokenProviderHandle::new(Arc::new( - PendingToken { - entered: entered.clone(), - dropped: dropped.clone(), - }, - ))), - ..request - }; - let NativeOutcome::Completed(mut call) = - OcrCall::admit(super::test_support::ocr_client(), OcrAdmission::all()) - else { - panic!("supported call declined") - }; - let mut request = Some(request); - let mut result = None; - tokio::time::timeout(std::time::Duration::from_secs(2), async { - loop { - tokio::select! { - _ = entered.notified() => break, - step = call.resume(result.take()) => { - result = Some(match step.unwrap() { - OcrCallStep::Host(OcrHostOperation::ProjectRequest) => OcrHostResult::Request(Ok((Box::new(request.take().unwrap()), false))), - OcrCallStep::Host(operation) => NoopOcrHost.invoke(operation).await, - OcrCallStep::Complete(_) => panic!("pending provider completed"), - }); - } - } - } - }).await.unwrap(); - assert!(!dropped.load(Ordering::SeqCst)); - let selected = crate::ocr::Error::InvalidRequest("cancelled".into()); - if interrupt_acknowledgement { - let mut acknowledgement = - Box::pin(call.interrupt(HostFailure::Cancelled(selected.clone()))); - std::future::poll_fn(|cx| { - assert!(acknowledgement.as_mut().poll(cx).is_pending()); - Poll::Ready(()) - }) - .await; - drop(acknowledgement); - assert!(!dropped.load(Ordering::SeqCst)); - } - let result = tokio::time::timeout( - std::time::Duration::from_secs(2), - call.interrupt(HostFailure::Cancelled(selected.clone())), - ) - .await - .unwrap(); - assert!(matches!(result, Err(error) if error == selected)); - assert!( - dropped.load(Ordering::SeqCst), - "cancellation returned while provider captures were still alive" - ); - } -} diff --git a/litellm-rust/crates/core/tests/reducto_ocr.rs b/litellm-rust/crates/core/tests/reducto_ocr.rs deleted file mode 100644 index d7bcf0b8d12..00000000000 --- a/litellm-rust/crates/core/tests/reducto_ocr.rs +++ /dev/null @@ -1,331 +0,0 @@ -use std::sync::Arc; - -use rstest::rstest; -use serde_json::{Value, json}; - -use super::hooks::{OcrDuringCallRequest, OcrHookFuture, OcrHooks, OcrPostCallRequest}; -use super::test_support::{MockResponse, mock_server, perform_ocr, wire_request}; - -fn request_body(request: &str) -> Value { - serde_json::from_str(request.split_once("\r\n\r\n").unwrap().1).unwrap() -} - -#[rstest] -#[case( - "reducto/parse-v3", - json!({ - "formatting":{"table_output_format":"html"}, - "retrieval":{"chunk_mode":"section"}, - "settings":{"ocr_system":"standard"}, - "future_ocr_option":true, - "extra_body":{"provider_option":"value"} - }), - "reducto://already.pdf", - json!({ - "input":"reducto://already.pdf", - "formatting":{"table_output_format":"html"}, - "retrieval":{"chunk_mode":"section"}, - "settings":{"ocr_system":"standard"}, - "future_ocr_option":true, - "provider_option":"value" - }) -)] -#[case( - "reducto/parse-legacy", - json!({ - "enhance":{"agentic":[{"type":"table"}]}, - "future_ocr_option":true, - "extra_body":{"provider_option":"value"} - }), - "reducto://legacy.pdf", - json!({ - "document_url":"reducto://legacy.pdf", - "options":{"enhance":{"agentic":[{"type":"table"}]}}, - "future_ocr_option":true, - "provider_option":"value" - }) -)] -#[tokio::test] -async fn request_mapping_matches_python( - #[case] model: &str, - #[case] options: Value, - #[case] source: &str, - #[case] expected: Value, -) { - let (base, seen, server) = mock_server(vec![MockResponse::json(json!({ - "result":{"chunks":[]} - }))]) - .await; - let mut request = wire_request(model, &base, options); - request.document = request.document.with_source(source.into()); - - perform_ocr(request).await.unwrap(); - server.await.unwrap(); - let requests = seen.lock().unwrap(); - assert_eq!(requests.len(), 1); - assert!(requests[0].starts_with("POST /parse ")); - assert_eq!(request_body(&requests[0]), expected); -} - -#[rstest] -#[case("parse-v3")] -#[case("parse-legacy")] -#[tokio::test] -async fn data_uri_upload_preserves_multipart_headers(#[case] model: &str) { - let (base, seen, server) = mock_server(vec![ - MockResponse::json(json!({"file_id":"reducto://uploaded.pdf"})), - MockResponse::json(json!({"result":{"chunks":[{"content":"hello"}]}})), - ]) - .await; - let mut request = wire_request(&format!("reducto/{model}"), &base, json!({})); - request.connection.extra_headers = vec![ - ("Content-Type".into(), "application/json".into()), - ("X-Trace".into(), "upload-test".into()), - ]; - - let response = perform_ocr(request).await.unwrap(); - server.await.unwrap(); - assert_eq!(response.pages[0].markdown, "hello"); - let requests = seen.lock().unwrap(); - assert_eq!(requests.len(), 2); - assert!(requests[0].starts_with("POST /upload ")); - assert!( - requests[0] - .to_ascii_lowercase() - .contains("content-type: multipart/form-data; boundary=") - ); - assert!(requests[0].contains("x-trace: upload-test")); - assert!(requests[0].contains("application/pdf")); - assert!(requests[0].contains("abc")); - assert!(requests[1].starts_with("POST /parse ")); -} - -struct ParseBoundary { - request_count: Arc>>, -} - -impl OcrHooks for ParseBoundary { - fn post_call(&self, request: OcrPostCallRequest) -> OcrHookFuture<'_, OcrPostCallRequest> { - Box::pin(async move { - assert_eq!(self.request_count.lock().unwrap().len(), 2); - assert_eq!( - request.original_response, - json!(r#"{"result":{"chunks":[]}}"#) - ); - Ok(request) - }) - } -} - -#[tokio::test] -async fn post_call_stays_after_reducto_upload_and_parse() { - let (base, seen, server) = mock_server(vec![ - MockResponse::json(json!({"file_id":"reducto://uploaded.pdf"})), - MockResponse::json(json!({"result":{"chunks":[]}})), - ]) - .await; - let request = super::LiteLLMOcrRequest { - hooks: Arc::new(ParseBoundary { - request_count: seen.clone(), - }), - ..wire_request("reducto/parse-v3", &base, json!({})) - }; - - perform_ocr(request).await.unwrap(); - server.await.unwrap(); - assert_eq!(seen.lock().unwrap().len(), 2); -} - -#[rstest] -#[case(json!({"file_id":""}))] -#[case(json!({}))] -#[case(json!({"file_id":null}))] -#[tokio::test] -async fn invalid_upload_ids_stop_before_parse(#[case] response: Value) { - let (base, seen, server) = mock_server(vec![MockResponse::json(response)]).await; - let error = perform_ocr(wire_request("reducto/parse-v3", &base, json!({}))) - .await - .unwrap_err(); - server.await.unwrap(); - assert!(error.to_string().contains("file_id")); - assert_eq!(seen.lock().unwrap().len(), 1); -} - -#[tokio::test] -async fn upload_failure_stops_before_parse() { - let (base, seen, server) = mock_server(vec![MockResponse { - status: 503, - headers: vec![], - body: json!({"error":"unavailable"}), - }]) - .await; - assert!( - perform_ocr(wire_request("reducto/parse-v3", &base, json!({}))) - .await - .is_err() - ); - server.await.unwrap(); - assert_eq!(seen.lock().unwrap().len(), 1); -} - -#[rstest] -#[case("https://example.com/a.pdf")] -#[case("reducto://")] -#[case("data:application/pdf;base64")] -#[case("data:application/pdf;base64,INVALID!")] -#[tokio::test] -async fn rejects_invalid_document_sources_before_network(#[case] source: &str) { - let mut request = wire_request("reducto/parse-v3", "http://127.0.0.1:1", json!({})); - request.document = request.document.with_source(source.into()); - assert!(perform_ocr(request).await.is_err()); -} - -#[test] -fn response_normalization_groups_blocks_and_distinguishes_null_result() { - use crate::llms::reducto::ocr::transformation::{ReductoResponse, normalize_response}; - - let raw = json!({"usage":{"num_pages":"2","credits":"3"},"result":{"type":"full","chunks":[ - {"blocks":[{ - "type":"Table", - "content":"B", - "bbox":{"left":0.1,"top":0.2,"width":0.8,"height":0.3,"page":2,"original_page":4}, - "confidence":"high", - "granular_confidence":{"parse_confidence":0.95,"extract_confidence":null}, - "image_url":null - }]}, - {"blocks":[{"content":"A","bbox":{"page":1},"type":"Text"},{"content":"C","bbox":{"page":1}}]} - ]}}); - let response: ReductoResponse = serde_json::from_value(raw).unwrap(); - let normalized = normalize_response("parse-v3", response) - .unwrap() - .into_json(); - assert_eq!(normalized["pages"][0]["markdown"], "A\n\nC"); - assert_eq!(normalized["pages"][1]["markdown"], "B"); - assert_eq!(normalized["pages"][1]["blocks"][0]["type"], "Table"); - assert_eq!( - normalized["pages"][1]["blocks"][0]["bbox"], - json!({"left":0.1,"top":0.2,"width":0.8,"height":0.3,"page":2,"original_page":4}) - ); - assert_eq!(normalized["pages"][1]["blocks"][0]["confidence"], "high"); - assert_eq!( - normalized["pages"][1]["blocks"][0]["granular_confidence"]["parse_confidence"], - 0.95 - ); - assert!(normalized["pages"][1]["blocks"][0]["image_url"].is_null()); - assert_eq!(normalized["usage_info"]["pages_processed"], 2); - assert_eq!(normalized["usage_info"]["credits"], 3.0); - - let missing: ReductoResponse = - serde_json::from_value(json!({"chunks":[{"content":"text"}]})).unwrap(); - let missing = normalize_response("parse-v3", missing).unwrap(); - assert_eq!(missing.pages[0].markdown, "text"); - let null: ReductoResponse = serde_json::from_value( - json!({"result":null,"chunks":[{"content":"ignored"}],"usage":null}), - ) - .unwrap(); - let null = normalize_response("parse-v3", null).unwrap(); - assert!(null.pages.is_empty()); -} - -#[tokio::test] -async fn facade_omits_native_response_by_default_and_preserves_auth_priority() { - let raw = json!({"job_id":"job-1","result":{"chunks":[]}}); - let (base, seen, server) = mock_server(vec![MockResponse::json(raw)]).await; - let mut request = wire_request("reducto/parse-v3", &base, json!({})); - request.document = request.document.with_source("reducto://ready.pdf".into()); - request.connection.extra_headers = vec![("authorization".into(), "Bearer existing".into())]; - - let response = perform_ocr(request).await.unwrap(); - server.await.unwrap(); - assert_eq!(response.provider_native_response, None); - assert!( - seen.lock().unwrap()[0] - .to_ascii_lowercase() - .contains("authorization: bearer existing") - ); -} - -struct RewriteDocument; - -struct RewriteHeaders; - -impl OcrHooks for RewriteHeaders { - fn intercepts_requests(&self) -> bool { - true - } - - fn during_call( - &self, - request: OcrDuringCallRequest, - ) -> OcrHookFuture<'_, OcrDuringCallRequest> { - Box::pin(async move { - Ok(OcrDuringCallRequest { - headers: vec![("authorization".into(), "Bearer guarded".into())], - ..request - }) - }) - } -} - -#[rstest] -#[case("reducto/parse-v3")] -#[case("reducto/parse-legacy")] -#[tokio::test] -async fn guardrail_headers_reach_upload_and_parse(#[case] model: &str) { - let (base, seen, server) = mock_server(vec![ - MockResponse::json(json!({"file_id":"reducto://uploaded.pdf"})), - MockResponse::json(json!({"result":{"chunks":[]}})), - ]) - .await; - let mut request = wire_request(model, &base, json!({})); - request.connection.extra_headers = vec![("authorization".into(), "Bearer original".into())]; - request.hooks = Arc::new(RewriteHeaders); - - perform_ocr(request).await.unwrap(); - server.await.unwrap(); - let requests = seen.lock().unwrap(); - assert_eq!(requests.len(), 2); - assert!(requests[0].starts_with("POST /upload ")); - assert!(requests[1].starts_with("POST /parse ")); - for request in requests.iter() { - assert!(request.contains("authorization: Bearer guarded")); - assert!(!request.contains("Bearer original")); - } -} - -impl OcrHooks for RewriteDocument { - fn intercepts_requests(&self) -> bool { - true - } - - fn during_call( - &self, - request: OcrDuringCallRequest, - ) -> OcrHookFuture<'_, OcrDuringCallRequest> { - Box::pin(async move { - assert_eq!( - request.body["document_url"], - "data:application/pdf;base64,YWJj" - ); - Ok(OcrDuringCallRequest { - body: json!({"type":"document_url","document_url":"reducto://guarded.pdf"}), - ..request - }) - }) - } -} - -#[tokio::test] -async fn guardrail_rewrites_document_before_upload() { - let (base, seen, server) = - mock_server(vec![MockResponse::json(json!({"result":{"chunks":[]}}))]).await; - let mut request = wire_request("reducto/parse-v3", &base, json!({})); - request.hooks = Arc::new(RewriteDocument); - - perform_ocr(request).await.unwrap(); - server.await.unwrap(); - let requests = seen.lock().unwrap(); - assert_eq!(requests.len(), 1); - assert!(requests[0].starts_with("POST /parse ")); - assert!(requests[0].contains("reducto://guarded.pdf")); -} diff --git a/litellm-rust/crates/core/tests/vertex_ai_deepseek_ocr.rs b/litellm-rust/crates/core/tests/vertex_ai_deepseek_ocr.rs deleted file mode 100644 index efda6339e5d..00000000000 --- a/litellm-rust/crates/core/tests/vertex_ai_deepseek_ocr.rs +++ /dev/null @@ -1,88 +0,0 @@ -use serde_json::{Value, json}; - -use super::test_support::{MockResponse, mock_server, perform_ocr, wire_request}; -use litellm_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().extra_fields["prompt_tokens"], - 1 - ); - let requests = seen.lock().unwrap(); - assert!(requests[0].starts_with( - "POST /v1/projects/project-1/locations/europe-west4/endpoints/openapi/chat/completions " - )); - assert!( - requests[0] - .to_ascii_lowercase() - .contains("authorization: bearer test-key") - ); - let body = request_body(&requests[0]); - assert_eq!(body["model"], "deepseek-ai/deepseek-ocr-maas"); - assert_eq!(body["temperature"], 0.1); - assert_eq!(body["future_ocr_option"], true); - assert_eq!(body["provider_option"], "value"); - assert!(body.get("vertex_project").is_none()); - assert!(body.get("extra_body").is_none()); - assert_eq!( - body["messages"][0]["content"][0], - json!({"type":"image_url","image_url":"gs://bucket/document.pdf"}) - ); -} - -#[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 deleted file mode 100644 index 12b0c6b9b33..00000000000 --- a/litellm-rust/crates/core/tests/vertex_ai_ocr.rs +++ /dev/null @@ -1,167 +0,0 @@ -use serde_json::{Value, json}; - -use super::test_support::{MockResponse, mock_server, perform_ocr, wire_request}; -use litellm_auth::InputSource; - -fn request_body(request: &str) -> Value { - serde_json::from_str(request.split_once("\r\n\r\n").unwrap().1).unwrap() -} - -#[tokio::test] -async fn facade_executes_vertex_mistral_with_resolved_project_and_location() { - let (base, seen, server) = mock_server(vec![MockResponse::json(json!({ - "pages":[{"index":0,"markdown":"hello"}], - "usage_info":{"pages_processed":1} - }))]) - .await; - let request = wire_request( - "vertex_ai/mistral-ocr-maas", - &base, - json!({ - "vertex_project":"project-1", - "vertex_location":"europe-west4", - "extract_footer":true - }), - ); - - let response = perform_ocr(request).await.unwrap(); - server.await.unwrap(); - assert_eq!(response.pages[0].markdown, "hello"); - let requests = seen.lock().unwrap(); - assert_eq!(requests.len(), 1); - assert!(requests[0].starts_with( - "POST /v1/projects/project-1/locations/europe-west4/publishers/mistralai/models/mistral-ocr-maas:rawPredict " - )); - assert!( - requests[0] - .to_ascii_lowercase() - .contains("authorization: bearer test-key") - ); - assert_eq!( - request_body(&requests[0]), - json!({ - "model":"mistral-ocr-maas", - "document":{"type":"document_url","document_url":"data:application/pdf;base64,YWJj"}, - "extract_footer":true - }) - ); -} - -#[tokio::test] -async fn supplied_authorization_is_forwarded_without_a_static_token() { - let (base, seen, server) = mock_server(vec![MockResponse::json(json!({"pages":[]}))]).await; - let mut request = wire_request( - "vertex_ai/model", - &base, - json!({"vertex_project":"project-1"}), - ); - request.connection.api_key = None; - request.connection.extra_headers = vec![("authorization".into(), "Bearer supplied".into())]; - - perform_ocr(request).await.unwrap(); - server.await.unwrap(); - assert!( - seen.lock().unwrap()[0] - .to_ascii_lowercase() - .contains("authorization: bearer supplied") - ); -} - -#[tokio::test] -async fn invalid_credentials_fail_before_provider_http() { - let request = wire_request( - "vertex_ai/model", - "http://127.0.0.1:1", - json!({"vertex_credentials": true}), - ); - let error = perform_ocr(request).await.unwrap_err(); - assert!(error.to_string().contains("vertex_credentials")); -} - -#[tokio::test] -async fn request_controlled_api_base_is_rejected_before_vertex_auth() { - let mut request = wire_request( - "vertex_ai/mistral-ocr-maas", - "https://caller.example", - json!({"vertex_project":"project-1"}), - ); - request.connection.api_base_source = InputSource::Request; - - let error = perform_ocr(request).await.unwrap_err(); - assert!( - error - .to_string() - .contains("request-controlled Vertex AI endpoint") - ); -} - -#[tokio::test] -async fn configs_build_complete_requests_and_share_mistral_normalization() { - use std::time::Duration; - - use crate::llms::base_llm::ocr::transformation::BaseOcrConfig; - use crate::llms::mistral::ocr::transformation::MistralOCRConfig; - use crate::llms::vertex_ai::ocr::transformation::VertexAIOCRConfig; - use crate::ocr::test_support::ocr_client; - - let client = ocr_client(); - let options = json!({ - "pages": [0, 2], - "include_image_base64": true, - "vertex_project": "project-1", - "vertex_location": "us-central1", - "unknown": "preserved" - }); - let direct = wire_request( - "mistral/mistral-ocr-maas", - "https://mistral.test", - options.clone(), - ); - let vertex = wire_request("vertex_ai/mistral-ocr-maas", "https://vertex.test", options); - let direct_http = MistralOCRConfig - .prepare_request(&direct, &client) - .await - .unwrap(); - let vertex_http = VertexAIOCRConfig - .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, - "unknown": "preserved" - }) - ); - } - let payload = json!({"pages": [{"index": 0, "markdown": "hello"}], "extra": "preserved"}); - let direct_response = MistralOCRConfig - .normalize_response( - &direct.model, - serde_json::from_value(payload.clone()).unwrap(), - ) - .unwrap() - .into_json(); - let vertex_response = VertexAIOCRConfig - .normalize_response(&vertex.model, 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!(direct_response.get("extra").is_none()); -} diff --git a/litellm-rust/crates/python-bridge/src/routes/ocr/errors.rs b/litellm-rust/crates/python-bridge/src/routes/ocr/errors.rs index 1035cf960c8..9f25de3dd14 100644 --- a/litellm-rust/crates/python-bridge/src/routes/ocr/errors.rs +++ b/litellm-rust/crates/python-bridge/src/routes/ocr/errors.rs @@ -6,6 +6,21 @@ use pyo3::prelude::*; use crate::errors::{RustUpstreamError, core_error_to_pyerr}; pub(super) fn to_pyerr(error: Error) -> PyErr { + if let Error::Provider { + status, + body, + headers, + } = error + { + let mapped = attach_status(RustUpstreamError::new_err((status, body)), Some(status)); + return Python::attach(|py| -> PyResult { + let headers = + pyo3::types::PyDict::from_sequence(&headers.into_pyobject(py)?.into_any())?; + mapped.value(py).setattr("headers", headers)?; + Ok(mapped) + }) + .unwrap_or_else(|error| error); + } let (mapped, status) = match error { Error::MissingDocumentUrl => ( PyValueError::new_err(Error::MissingDocumentUrl.to_string()), @@ -47,6 +62,39 @@ mod tests { use super::*; use pyo3::exceptions::PyValueError; + #[test] + fn provider_error_retains_headers_at_the_python_boundary() { + Python::initialize(); + Python::attach(|py| { + let error = to_pyerr(Error::Provider { + status: 429, + body: "rate limited".into(), + headers: vec![("retry-after".into(), "17".into())], + }); + assert!(error.is_instance_of::(py)); + assert_eq!( + error + .value(py) + .getattr("args") + .unwrap() + .extract::<(u16, String)>() + .unwrap(), + (429, "rate limited".into()) + ); + assert_eq!( + error + .value(py) + .getattr("headers") + .unwrap() + .get_item("retry-after") + .unwrap() + .extract::() + .unwrap(), + "17" + ); + }); + } + #[test] fn preserves_python_validation_and_provider_details() { Python::initialize();