mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-24 00:52:24 +00:00
fix(rust): read OCR env-backed constants instead of hardcoding their defaults
Native OCR hardcoded the default of five Python constants that come from env vars, so an operator setting them saw no effect: REQUEST_TIMEOUT (Rust used 600s, Python 6000s), MAX_IMAGE_URL_DOWNLOAD_SIZE_MB (0 disables document downloads), AZURE_OPERATION_POLLING_TIMEOUT, AZURE_DOCUMENT_INTELLIGENCE_API_VERSION and AZURE_DOCUMENT_INTELLIGENCE_DEFAULT_DPI. OcrSettings reads them through Lookup with Python's parsing, the bridge builds it per call and OcrClient carries it into the connection. A zero per-call timeout now falls back to REQUEST_TIMEOUT, matching `timeout or request_timeout`. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
This commit is contained in:
parent
d77c144c6c
commit
c0705f31b4
13 changed files with 297 additions and 54 deletions
|
|
@ -22,7 +22,7 @@ pub(crate) async fn perform_ocr_request(
|
|||
) -> Result<LiteLLMOcrResponse, Error> {
|
||||
request.response_format()?;
|
||||
let config = request.config;
|
||||
let request = prepare_request(request, caller_document);
|
||||
let request = prepare_request(request, caller_document, client.settings());
|
||||
let hooks = OcrCallHooks::new(host.clone(), &request, config);
|
||||
config.ocr(client, &request, &hooks).await
|
||||
}
|
||||
|
|
|
|||
|
|
@ -1,6 +1,7 @@
|
|||
use litellm_auth::{InputSource, SecretValue, Sourced};
|
||||
use litellm_llms::base_llm::ocr::transformation::{
|
||||
OcrConnection, OcrCredentialInputs, PreparedOcrRequest, credential_env,
|
||||
use litellm_llms::base_llm::ocr::{
|
||||
settings::OcrSettings,
|
||||
transformation::{OcrConnection, OcrCredentialInputs, PreparedOcrRequest, credential_env},
|
||||
};
|
||||
|
||||
use super::provider_config::OcrProvider;
|
||||
|
|
@ -9,6 +10,7 @@ use crate::ocr::types::{LiteLLMOcrRequest, ResolvedOcrRequest};
|
|||
pub(crate) fn prepare_request(
|
||||
request: ResolvedOcrRequest,
|
||||
caller_document: bool,
|
||||
settings: &OcrSettings,
|
||||
) -> PreparedOcrRequest {
|
||||
let credentials = request.credentials.clone();
|
||||
let api_base_env = match request.config.provider() {
|
||||
|
|
@ -51,7 +53,7 @@ pub(crate) fn prepare_request(
|
|||
PreparedOcrRequest {
|
||||
model,
|
||||
document,
|
||||
connection: OcrConnection::new(resolved, transport),
|
||||
connection: OcrConnection::new(resolved, transport, settings.clone()),
|
||||
caller_document,
|
||||
optional_params,
|
||||
input_sources,
|
||||
|
|
@ -61,7 +63,7 @@ pub(crate) fn prepare_request(
|
|||
|
||||
#[cfg(test)]
|
||||
pub(crate) fn prepare_request_for_test(request: ResolvedOcrRequest) -> PreparedOcrRequest {
|
||||
prepare_request(request, true)
|
||||
prepare_request(request, true, &OcrSettings::default())
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
|
|
|
|||
|
|
@ -277,7 +277,7 @@ mod tests {
|
|||
vec![("x-a".to_string(), "1".to_string())]
|
||||
);
|
||||
assert_eq!(request.transport.extra_headers_source, InputSource::Request);
|
||||
assert_eq!(request.transport.timeout, Duration::from_secs(7));
|
||||
assert_eq!(request.transport.timeout, Some(Duration::from_secs(7)));
|
||||
assert_eq!(request.input_sources.len(), 2);
|
||||
|
||||
let defaulted = LiteLLMOcrRequest::from_inputs(
|
||||
|
|
|
|||
|
|
@ -1,10 +1,12 @@
|
|||
use litellm_host::event::{CallEvent, MachineEvent};
|
||||
use litellm_llms::base_llm::ocr::error::Error;
|
||||
use litellm_llms::base_llm::ocr::{error::Error, settings::OcrSettings};
|
||||
use rstest::rstest;
|
||||
use serde_json::{Value, json};
|
||||
|
||||
use super::{
|
||||
test_support::{MockResponse, mock_server, perform_ocr, perform_ocr_with, wire_request},
|
||||
test_support::{
|
||||
MockResponse, mock_server, ocr_client, perform_ocr, perform_ocr_with, wire_request,
|
||||
},
|
||||
wire::{OcrWireRequest, decode_request},
|
||||
};
|
||||
use crate::ocr::route::LocalOcrHost;
|
||||
|
|
@ -200,6 +202,42 @@ async fn immediate_response_normalizes_pages_and_preserves_native() {
|
|||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn client_settings_choose_the_api_version_and_the_inch_to_pixel_dpi() {
|
||||
let (base, seen, server) = mock_server(vec![MockResponse::json(json!({
|
||||
"status":"succeeded",
|
||||
"analyzeResult":{"pages":[{"pageNumber":1,"width":8.5,"height":11,"unit":"inch"}]}
|
||||
}))])
|
||||
.await;
|
||||
let client = ocr_client().with_settings(OcrSettings {
|
||||
document_intelligence_api_version: "2099-01-01".into(),
|
||||
document_intelligence_dpi: 72,
|
||||
..OcrSettings::default()
|
||||
});
|
||||
|
||||
let result = crate::ocr::client::perform(
|
||||
&client,
|
||||
wire_request("azure_ai/doc-intelligence/prebuilt-read", &base, json!({})),
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
server.await.unwrap();
|
||||
|
||||
let target = seen.lock().unwrap()[0]
|
||||
.split_whitespace()
|
||||
.nth(1)
|
||||
.unwrap()
|
||||
.to_string();
|
||||
assert_eq!(
|
||||
query_value(&format!("{base}{target}"), "api-version").as_deref(),
|
||||
Some("2099-01-01")
|
||||
);
|
||||
assert_eq!(
|
||||
serde_json::to_value(&result.pages[0].dimensions).unwrap(),
|
||||
json!({"width":612,"height":792,"dpi":72})
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn accepted_response_polls_to_success_with_only_credentials() {
|
||||
let operation = json!({"status":"succeeded","analyzeResult":{"pages":[]}});
|
||||
|
|
@ -425,13 +463,19 @@ async fn polling_deadline_bounds_retry_delay() {
|
|||
},
|
||||
])
|
||||
.await;
|
||||
let mut request = wire_request("azure_ai/doc-intelligence/prebuilt-read", &base, json!({}));
|
||||
request.transport.poll_timeout = std::time::Duration::from_millis(100);
|
||||
let request = wire_request("azure_ai/doc-intelligence/prebuilt-read", &base, json!({}));
|
||||
let client = ocr_client().with_settings(OcrSettings {
|
||||
poll_timeout: std::time::Duration::from_millis(100),
|
||||
..OcrSettings::default()
|
||||
});
|
||||
|
||||
let error = tokio::time::timeout(std::time::Duration::from_secs(1), perform_ocr(request))
|
||||
.await
|
||||
.unwrap()
|
||||
.unwrap_err();
|
||||
let error = tokio::time::timeout(
|
||||
std::time::Duration::from_secs(1),
|
||||
crate::ocr::client::perform(&client, request),
|
||||
)
|
||||
.await
|
||||
.unwrap()
|
||||
.unwrap_err();
|
||||
server.await.unwrap();
|
||||
assert!(error.to_string().contains("timed out"));
|
||||
}
|
||||
|
|
|
|||
|
|
@ -13,6 +13,7 @@ use litellm_http::{
|
|||
use litellm_llms::base_llm::ocr::{
|
||||
error::Error as OcrError,
|
||||
handler::OcrClient,
|
||||
settings::OcrSettings,
|
||||
transformation::{LiteLLMOcrResponse, OCR_RESPONSE_MAX_BYTES, OcrTransportConfig},
|
||||
};
|
||||
use rstest::rstest;
|
||||
|
|
@ -185,6 +186,7 @@ async fn ocr_client_uses_the_injected_http_pool_configuration() {
|
|||
&Resolution::from(&settings).config,
|
||||
UrlPolicy::default(),
|
||||
VertexAuth::default(),
|
||||
OcrSettings::default(),
|
||||
)
|
||||
.unwrap();
|
||||
crate::ocr::client::perform(&client, wire_request("mistral/model", &base, json!({})))
|
||||
|
|
|
|||
|
|
@ -18,6 +18,7 @@ use crate::base_llm::ocr::{
|
|||
document::InlineDocument,
|
||||
error::Error,
|
||||
handler::{CallHooks, OcrClient, read_json_response},
|
||||
settings::OcrSettings,
|
||||
transformation::{
|
||||
BaseOcrConfig, DecodedOcrResponse, LiteLLMOcrResponse, OCR_INLINE_MAX_BYTES,
|
||||
OCR_POLL_RETRY_SECS, OcrConnection, OcrCredentialInputs, OcrDocument, OcrPage,
|
||||
|
|
@ -26,9 +27,7 @@ use crate::base_llm::ocr::{
|
|||
},
|
||||
};
|
||||
|
||||
const AZURE_DI_API_VERSION: &str = "2024-11-30";
|
||||
const AZURE_DI_SUBSCRIPTION_HEADER: &str = "Ocp-Apim-Subscription-Key";
|
||||
const AZURE_DI_DEFAULT_DPI: i64 = 96;
|
||||
const AZURE_DI_DEFAULT_WIDTH: f64 = 8.5;
|
||||
const AZURE_DI_DEFAULT_HEIGHT: f64 = 11.0;
|
||||
|
||||
|
|
@ -195,7 +194,15 @@ impl BaseOcrConfig for AzureDocumentIntelligenceOcrConfig {
|
|||
let endpoint = nonblank(request.connection.api_base.clone())
|
||||
.or_else(|| nonblank(credential_env(AZURE_DI_ENDPOINT_ENV)))
|
||||
.ok_or_else(|| Error::Auth(litellm_auth::Error::ProviderAuthentication("Missing Azure Document Intelligence API Base - Set AZURE_DOCUMENT_INTELLIGENCE_ENDPOINT or pass api_base".into())))?;
|
||||
self.build_ocr_url(&endpoint, &request.model, optional_params)
|
||||
self.build_ocr_url(
|
||||
&endpoint,
|
||||
&request.model,
|
||||
optional_params,
|
||||
&request
|
||||
.connection
|
||||
.settings
|
||||
.document_intelligence_api_version,
|
||||
)
|
||||
}
|
||||
|
||||
fn transform_ocr_request(
|
||||
|
|
@ -214,12 +221,13 @@ impl BaseOcrConfig for AzureDocumentIntelligenceOcrConfig {
|
|||
raw_response: &[u8],
|
||||
request_format: OcrResponseFormat,
|
||||
) -> Result<LiteLLMOcrResponse, Error> {
|
||||
decode_and_normalize_response(
|
||||
model,
|
||||
raw_response,
|
||||
request_format,
|
||||
transform_completed_response,
|
||||
)
|
||||
decode_and_normalize_response(model, raw_response, request_format, |model, response| {
|
||||
transform_completed_response(
|
||||
model,
|
||||
response,
|
||||
OcrSettings::default().document_intelligence_dpi,
|
||||
)
|
||||
})
|
||||
}
|
||||
|
||||
async fn async_transform_ocr_response(
|
||||
|
|
@ -240,7 +248,11 @@ impl BaseOcrConfig for AzureDocumentIntelligenceOcrConfig {
|
|||
.await?;
|
||||
Ok(LiteLLMOcrResponse {
|
||||
provider_native_response: decoded.native,
|
||||
..transform_completed_response(model, decoded.data)?
|
||||
..transform_completed_response(
|
||||
model,
|
||||
decoded.data,
|
||||
context.connection.settings.document_intelligence_dpi,
|
||||
)?
|
||||
})
|
||||
}
|
||||
}
|
||||
|
|
@ -353,6 +365,7 @@ fn build_request(document: OcrDocument) -> Result<DocumentIntelligenceRequest, E
|
|||
fn transform_completed_response(
|
||||
model: &str,
|
||||
response: AzureDocumentIntelligenceOperation,
|
||||
dpi: i64,
|
||||
) -> Result<LiteLLMOcrResponse, Error> {
|
||||
if response.status != Some(OperationStatus::Succeeded) {
|
||||
return Err(Error::OperationStatus(
|
||||
|
|
@ -366,7 +379,7 @@ fn transform_completed_response(
|
|||
let pages = result
|
||||
.pages
|
||||
.into_iter()
|
||||
.map(transform_azure_page)
|
||||
.map(|page| transform_azure_page(page, dpi))
|
||||
.collect::<Result<Vec<_>, _>>()?;
|
||||
let pages_processed = i64::try_from(pages.len()).map_err(|_| Error::NumericRange("pages"))?;
|
||||
Ok(LiteLLMOcrResponse {
|
||||
|
|
@ -381,7 +394,7 @@ fn transform_completed_response(
|
|||
})
|
||||
}
|
||||
|
||||
fn transform_azure_page(page: AzureDocumentIntelligencePage) -> Result<OcrPage, Error> {
|
||||
fn transform_azure_page(page: AzureDocumentIntelligencePage, dpi: i64) -> Result<OcrPage, Error> {
|
||||
let index = page
|
||||
.page_number
|
||||
.unwrap_or(1)
|
||||
|
|
@ -391,6 +404,7 @@ fn transform_azure_page(page: AzureDocumentIntelligencePage) -> Result<OcrPage,
|
|||
page.width.unwrap_or(AZURE_DI_DEFAULT_WIDTH),
|
||||
page.height.unwrap_or(AZURE_DI_DEFAULT_HEIGHT),
|
||||
page.unit.as_deref().unwrap_or("inch"),
|
||||
dpi,
|
||||
)?;
|
||||
let markdown = page
|
||||
.lines
|
||||
|
|
@ -406,16 +420,17 @@ fn transform_azure_page(page: AzureDocumentIntelligencePage) -> Result<OcrPage,
|
|||
})
|
||||
}
|
||||
|
||||
fn convert_dimensions(width: f64, height: f64, unit: &str) -> Result<OcrPageDimensions, Error> {
|
||||
let scale = if unit == "inch" {
|
||||
AZURE_DI_DEFAULT_DPI as f64
|
||||
} else {
|
||||
1.0
|
||||
};
|
||||
fn convert_dimensions(
|
||||
width: f64,
|
||||
height: f64,
|
||||
unit: &str,
|
||||
dpi: i64,
|
||||
) -> Result<OcrPageDimensions, Error> {
|
||||
let scale = if unit == "inch" { dpi as f64 } else { 1.0 };
|
||||
Ok(OcrPageDimensions {
|
||||
width: Some(pixel_dimension(width, scale, "page.width")?),
|
||||
height: Some(pixel_dimension(height, scale, "page.height")?),
|
||||
dpi: Some(AZURE_DI_DEFAULT_DPI),
|
||||
dpi: Some(dpi),
|
||||
})
|
||||
}
|
||||
|
||||
|
|
@ -475,7 +490,7 @@ async fn poll_operation(
|
|||
hooks: &dyn CallHooks<Error>,
|
||||
) -> Result<DecodedOcrResponse<AzureDocumentIntelligenceOperation>, Error> {
|
||||
let deadline = Instant::now()
|
||||
.checked_add(connection.poll_timeout)
|
||||
.checked_add(connection.settings.poll_timeout)
|
||||
.ok_or(Error::PollTimeout)?;
|
||||
|
||||
loop {
|
||||
|
|
@ -544,13 +559,14 @@ impl AzureDocumentIntelligenceOcrConfig {
|
|||
endpoint: &str,
|
||||
model: &str,
|
||||
params: &DocumentIntelligenceParams,
|
||||
api_version: &str,
|
||||
) -> Result<String, Error> {
|
||||
let model = format!("{}:analyze", model_id(model)?);
|
||||
ApiUrl::parse(endpoint)
|
||||
.and_then(|url| url.complete_path(&["documentintelligence", "documentModels", &model]))
|
||||
.map(|url| {
|
||||
url.append_query_pairs(
|
||||
[("api-version", AZURE_DI_API_VERSION)]
|
||||
[("api-version", api_version)]
|
||||
.into_iter()
|
||||
.chain(params.pages.iter().map(|pages| ("pages", pages.as_str())))
|
||||
.chain(
|
||||
|
|
|
|||
|
|
@ -68,7 +68,7 @@ pub async fn inline_remote_document(
|
|||
url,
|
||||
DownloadPolicy {
|
||||
timeout: connection.timeout,
|
||||
max_bytes: connection.max_download_bytes,
|
||||
max_bytes: connection.settings.max_download_bytes,
|
||||
max_redirects: OCR_MAX_FETCH_REDIRECTS,
|
||||
},
|
||||
)
|
||||
|
|
|
|||
|
|
@ -13,6 +13,7 @@ use serde_json::Value;
|
|||
|
||||
use crate::base_llm::ocr::{
|
||||
error::Error,
|
||||
settings::OcrSettings,
|
||||
transformation::{
|
||||
BaseOcrConfig, DecodedOcrResponse, LiteLLMOcrResponse, OcrDocument, OcrResponseContext,
|
||||
PreparedOcrRequest, decode_request_value, decode_response,
|
||||
|
|
@ -33,6 +34,7 @@ pub struct OcrClient {
|
|||
polling_http: reqwest::Client,
|
||||
document_fetcher: MediaFetcher,
|
||||
vertex_auth: VertexAuth,
|
||||
settings: OcrSettings,
|
||||
}
|
||||
|
||||
impl OcrClient {
|
||||
|
|
@ -41,12 +43,14 @@ impl OcrClient {
|
|||
config: &HttpClientConfig,
|
||||
url_policy: UrlPolicy,
|
||||
vertex_auth: VertexAuth,
|
||||
settings: OcrSettings,
|
||||
) -> Result<Self, litellm_http::Error> {
|
||||
Ok(Self {
|
||||
provider_http: pool.client(config, ClientVariant::Provider)?,
|
||||
polling_http: pool.client(config, ClientVariant::NoRedirect)?,
|
||||
document_fetcher: MediaFetcher::new(pool, config, url_policy)?,
|
||||
vertex_auth,
|
||||
settings,
|
||||
})
|
||||
}
|
||||
|
||||
|
|
@ -66,6 +70,10 @@ impl OcrClient {
|
|||
&self.vertex_auth
|
||||
}
|
||||
|
||||
pub fn settings(&self) -> &OcrSettings {
|
||||
&self.settings
|
||||
}
|
||||
|
||||
#[cfg(any(test, feature = "test-support"))]
|
||||
pub fn for_test(provider_http: reqwest::Client, document_http: reqwest::Client) -> Self {
|
||||
Self {
|
||||
|
|
@ -76,8 +84,14 @@ impl OcrClient {
|
|||
.expect("test polling client builds"),
|
||||
document_fetcher: MediaFetcher::for_test(document_http),
|
||||
vertex_auth: VertexAuth::default(),
|
||||
settings: OcrSettings::default(),
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(any(test, feature = "test-support"))]
|
||||
pub fn with_settings(self, settings: OcrSettings) -> Self {
|
||||
Self { settings, ..self }
|
||||
}
|
||||
}
|
||||
|
||||
/// Rust counterpart of `BaseLLMHTTPHandler.async_ocr`: prepare the provider request,
|
||||
|
|
|
|||
|
|
@ -1,4 +1,5 @@
|
|||
pub mod document;
|
||||
pub mod error;
|
||||
pub mod handler;
|
||||
pub mod settings;
|
||||
pub mod transformation;
|
||||
|
|
|
|||
137
litellm-rust/crates/llms/src/base_llm/ocr/settings.rs
Normal file
137
litellm-rust/crates/llms/src/base_llm/ocr/settings.rs
Normal file
|
|
@ -0,0 +1,137 @@
|
|||
use std::time::Duration;
|
||||
|
||||
use litellm_core_utils::settings::Lookup;
|
||||
|
||||
#[derive(Clone, Debug, PartialEq)]
|
||||
pub struct OcrSettings {
|
||||
pub request_timeout: Duration,
|
||||
pub max_download_bytes: u64,
|
||||
pub poll_timeout: Duration,
|
||||
pub document_intelligence_api_version: String,
|
||||
pub document_intelligence_dpi: i64,
|
||||
}
|
||||
|
||||
impl Default for OcrSettings {
|
||||
fn default() -> Self {
|
||||
Self {
|
||||
request_timeout: Duration::from_secs(6000),
|
||||
max_download_bytes: megabytes(50.0),
|
||||
poll_timeout: Duration::from_secs(120),
|
||||
document_intelligence_api_version: "2024-11-30".into(),
|
||||
document_intelligence_dpi: 96,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl OcrSettings {
|
||||
pub fn from_environment(env: &impl Lookup) -> Self {
|
||||
let defaults = Self::default();
|
||||
Self {
|
||||
request_timeout: env
|
||||
.parsed::<f64>("REQUEST_TIMEOUT")
|
||||
.and_then(|seconds| Duration::try_from_secs_f64(seconds).ok())
|
||||
.unwrap_or(defaults.request_timeout),
|
||||
max_download_bytes: env
|
||||
.parsed::<f64>("MAX_IMAGE_URL_DOWNLOAD_SIZE_MB")
|
||||
.filter(|size| size.is_finite())
|
||||
.map_or(defaults.max_download_bytes, megabytes),
|
||||
poll_timeout: env
|
||||
.parsed::<i64>("AZURE_OPERATION_POLLING_TIMEOUT")
|
||||
.map_or(defaults.poll_timeout, |seconds| {
|
||||
Duration::from_secs(seconds.max(0).unsigned_abs())
|
||||
}),
|
||||
document_intelligence_api_version: env
|
||||
.get("AZURE_DOCUMENT_INTELLIGENCE_API_VERSION")
|
||||
.unwrap_or(defaults.document_intelligence_api_version),
|
||||
document_intelligence_dpi: env
|
||||
.parsed("AZURE_DOCUMENT_INTELLIGENCE_DEFAULT_DPI")
|
||||
.unwrap_or(defaults.document_intelligence_dpi),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn megabytes(size: f64) -> u64 {
|
||||
(size * 1024.0 * 1024.0) as u64
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use rstest::rstest;
|
||||
|
||||
use super::*;
|
||||
|
||||
fn env_of(values: &'static [(&'static str, &'static str)]) -> impl Fn(&str) -> Option<String> {
|
||||
move |name| {
|
||||
values
|
||||
.iter()
|
||||
.find(|(key, _)| *key == name)
|
||||
.map(|(_, value)| value.to_string())
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn an_empty_environment_keeps_the_python_defaults() {
|
||||
assert_eq!(
|
||||
OcrSettings::from_environment(&env_of(&[])),
|
||||
OcrSettings::default()
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn every_setting_follows_its_environment_variable() {
|
||||
let settings = OcrSettings::from_environment(&env_of(&[
|
||||
("REQUEST_TIMEOUT", "30.5"),
|
||||
("MAX_IMAGE_URL_DOWNLOAD_SIZE_MB", "0.5"),
|
||||
("AZURE_OPERATION_POLLING_TIMEOUT", " 600 "),
|
||||
("AZURE_DOCUMENT_INTELLIGENCE_API_VERSION", "2025-01-01"),
|
||||
("AZURE_DOCUMENT_INTELLIGENCE_DEFAULT_DPI", "72"),
|
||||
]));
|
||||
assert_eq!(
|
||||
settings,
|
||||
OcrSettings {
|
||||
request_timeout: Duration::from_millis(30_500),
|
||||
max_download_bytes: 512 * 1024,
|
||||
poll_timeout: Duration::from_secs(600),
|
||||
document_intelligence_api_version: "2025-01-01".into(),
|
||||
document_intelligence_dpi: 72,
|
||||
}
|
||||
);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::zero_disables_downloads("0", 0)]
|
||||
#[case::negative_rejects_every_download("-1", 0)]
|
||||
#[case::fraction_truncates_like_int("0.0000001", 0)]
|
||||
#[case::unparsable_keeps_the_default("big", 50 * 1024 * 1024)]
|
||||
fn download_size_converts_megabytes_like_python(
|
||||
#[case] value: &'static str,
|
||||
#[case] bytes: u64,
|
||||
) {
|
||||
let env =
|
||||
move |name: &str| (name == "MAX_IMAGE_URL_DOWNLOAD_SIZE_MB").then(|| value.to_string());
|
||||
assert_eq!(
|
||||
OcrSettings::from_environment(&env).max_download_bytes,
|
||||
bytes
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn a_negative_polling_timeout_expires_immediately() {
|
||||
let env =
|
||||
|name: &str| (name == "AZURE_OPERATION_POLLING_TIMEOUT").then(|| "-5".to_string());
|
||||
assert_eq!(
|
||||
OcrSettings::from_environment(&env).poll_timeout,
|
||||
Duration::ZERO
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn an_empty_api_version_is_sent_as_is_like_python_str_of_getenv() {
|
||||
let env =
|
||||
|name: &str| (name == "AZURE_DOCUMENT_INTELLIGENCE_API_VERSION").then(String::new);
|
||||
assert_eq!(
|
||||
OcrSettings::from_environment(&env).document_intelligence_api_version,
|
||||
""
|
||||
);
|
||||
}
|
||||
}
|
||||
|
|
@ -15,14 +15,12 @@ use serde_with::serde_as;
|
|||
use crate::base_llm::ocr::{
|
||||
error::Error,
|
||||
handler::{CallHooks, OcrClient, read_response_bytes, transform_request_body},
|
||||
settings::OcrSettings,
|
||||
};
|
||||
|
||||
pub const OCR_RESPONSE_MAX_BYTES: usize = 64 * 1024 * 1024;
|
||||
pub const OCR_HTTP_TIMEOUT_SECS: u64 = 600;
|
||||
pub const OCR_INLINE_MAX_BYTES: usize = 50 * 1024 * 1024;
|
||||
pub const OCR_DOWNLOAD_MAX_BYTES: u64 = 50 * 1024 * 1024;
|
||||
pub const OCR_MAX_FETCH_REDIRECTS: usize = 10;
|
||||
pub const OCR_POLL_TIMEOUT_SECS: u64 = 120;
|
||||
pub const OCR_POLL_RETRY_SECS: u64 = 2;
|
||||
|
||||
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
|
||||
|
|
@ -114,10 +112,8 @@ impl OcrCredentialInputs {
|
|||
pub struct OcrTransportConfig {
|
||||
pub extra_headers: Vec<(String, String)>,
|
||||
pub extra_headers_source: InputSource,
|
||||
pub timeout: Duration,
|
||||
pub max_download_bytes: u64,
|
||||
pub timeout: Option<Duration>,
|
||||
pub max_response_bytes: usize,
|
||||
pub poll_timeout: Duration,
|
||||
}
|
||||
|
||||
impl Default for OcrTransportConfig {
|
||||
|
|
@ -125,10 +121,8 @@ impl Default for OcrTransportConfig {
|
|||
Self {
|
||||
extra_headers: Vec::new(),
|
||||
extra_headers_source: InputSource::Deployment,
|
||||
timeout: Duration::from_secs(OCR_HTTP_TIMEOUT_SECS),
|
||||
max_download_bytes: OCR_DOWNLOAD_MAX_BYTES,
|
||||
timeout: None,
|
||||
max_response_bytes: OCR_RESPONSE_MAX_BYTES,
|
||||
poll_timeout: Duration::from_secs(OCR_POLL_TIMEOUT_SECS),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
|
@ -143,7 +137,7 @@ impl OcrTransportConfig {
|
|||
Self {
|
||||
extra_headers,
|
||||
extra_headers_source,
|
||||
timeout: timeout.unwrap_or(self.timeout),
|
||||
timeout: timeout.or(self.timeout),
|
||||
..self
|
||||
}
|
||||
}
|
||||
|
|
@ -164,13 +158,16 @@ pub struct OcrConnection {
|
|||
pub extra_headers: Vec<(String, String)>,
|
||||
pub extra_headers_source: InputSource,
|
||||
pub timeout: Duration,
|
||||
pub max_download_bytes: u64,
|
||||
pub max_response_bytes: usize,
|
||||
pub poll_timeout: Duration,
|
||||
pub settings: OcrSettings,
|
||||
}
|
||||
|
||||
impl OcrConnection {
|
||||
pub fn new(credentials: ResolvedOcrCredentials, transport: OcrTransportConfig) -> Self {
|
||||
pub fn new(
|
||||
credentials: ResolvedOcrCredentials,
|
||||
transport: OcrTransportConfig,
|
||||
settings: OcrSettings,
|
||||
) -> Self {
|
||||
let api_key_source = credentials
|
||||
.api_key
|
||||
.as_ref()
|
||||
|
|
@ -188,10 +185,12 @@ impl OcrConnection {
|
|||
api_base_source,
|
||||
extra_headers: transport.extra_headers,
|
||||
extra_headers_source: transport.extra_headers_source,
|
||||
timeout: transport.timeout,
|
||||
max_download_bytes: transport.max_download_bytes,
|
||||
timeout: transport
|
||||
.timeout
|
||||
.filter(|timeout| !timeout.is_zero())
|
||||
.unwrap_or(settings.request_timeout),
|
||||
max_response_bytes: transport.max_response_bytes,
|
||||
poll_timeout: transport.poll_timeout,
|
||||
settings,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
|
@ -201,6 +200,7 @@ impl Default for OcrConnection {
|
|||
Self::new(
|
||||
ResolvedOcrCredentials::default(),
|
||||
OcrTransportConfig::default(),
|
||||
OcrSettings::default(),
|
||||
)
|
||||
}
|
||||
}
|
||||
|
|
@ -573,6 +573,31 @@ mod tests {
|
|||
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn connection_timeout_falls_back_to_the_request_timeout_setting_like_a_python_or() {
|
||||
let settings = OcrSettings {
|
||||
request_timeout: Duration::from_secs(42),
|
||||
..OcrSettings::default()
|
||||
};
|
||||
let timeout = |call: Option<Duration>| {
|
||||
OcrConnection::new(
|
||||
ResolvedOcrCredentials::default(),
|
||||
OcrTransportConfig {
|
||||
timeout: call,
|
||||
..OcrTransportConfig::default()
|
||||
},
|
||||
settings.clone(),
|
||||
)
|
||||
.timeout
|
||||
};
|
||||
assert_eq!(timeout(None), Duration::from_secs(42));
|
||||
assert_eq!(timeout(Some(Duration::ZERO)), Duration::from_secs(42));
|
||||
assert_eq!(
|
||||
timeout(Some(Duration::from_secs(5))),
|
||||
Duration::from_secs(5)
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn normalized_response_rejects_invalid_shared_fields() {
|
||||
for fields in [
|
||||
|
|
|
|||
|
|
@ -9,7 +9,8 @@ use host::OcrRouteHost;
|
|||
use litellm_auth_gcp::VertexAuth;
|
||||
use litellm_callbacks_legacy::{LegacySurface, PublicCall, run_legacy_call};
|
||||
use litellm_core::ocr::route::ocr_machine;
|
||||
use litellm_llms::base_llm::ocr::handler::OcrClient;
|
||||
use litellm_core_utils::settings::ProcessEnvironment;
|
||||
use litellm_llms::base_llm::ocr::{handler::OcrClient, settings::OcrSettings};
|
||||
use pyo3::{
|
||||
prelude::*,
|
||||
types::{PyDict, PyTuple},
|
||||
|
|
@ -43,6 +44,7 @@ fn run_ocr(
|
|||
&config,
|
||||
http::url_policy(py)?,
|
||||
VERTEX_AUTH.clone(),
|
||||
OcrSettings::from_environment(&ProcessEnvironment),
|
||||
)
|
||||
.map_err(|error| RustBridgeDeclined::new_err(error.to_string()))?;
|
||||
run_legacy_call(
|
||||
|
|
|
|||
|
|
@ -592,7 +592,7 @@ kwargs = {
|
|||
);
|
||||
assert_eq!(
|
||||
projected.transport.timeout,
|
||||
std::time::Duration::from_secs(5)
|
||||
Some(std::time::Duration::from_secs(5))
|
||||
);
|
||||
});
|
||||
}
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue