From edfa01da81d2456fa9182beeff6e12278c04468b Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Wed, 16 Sep 2026 19:48:26 -0700 Subject: [PATCH 01/15] refactor(ocr): mirror Python provider layout and preserve tests --- litellm-rust/Cargo.lock | 147 +- litellm-rust/Cargo.toml | 1 + litellm-rust/crates/core/Cargo.toml | 3 +- .../crates/core/src/call_arguments.rs | 467 ++++++ litellm-rust/crates/core/src/lib.rs | 4 + .../crates/core/src/llms/azure_ai/mod.rs | 1 + .../ocr/cohere_parse_transformation.rs | 165 +++ .../azure_ai/ocr/common_utils.rs} | 24 +- .../azure_ai/ocr/document_intelligence/mod.rs | 1 + .../document_intelligence/transformation.rs | 1260 +++++++++++++++++ .../crates/core/src/llms/azure_ai/ocr/mod.rs | 4 + .../src/llms/azure_ai/ocr/transformation.rs | 399 ++++++ .../crates/core/src/llms/base_llm/mod.rs | 1 + .../crates/core/src/llms/base_llm/ocr/mod.rs | 1 + .../src/llms/base_llm/ocr/transformation.rs | 211 +++ .../crates/core/src/llms/cohere/mod.rs | 1 + .../crates/core/src/llms/cohere/ocr/mod.rs | 3 + .../src/llms/cohere/ocr/transformation.rs | 740 ++++++++++ .../crates/core/src/llms/mistral/mod.rs | 1 + .../crates/core/src/llms/mistral/ocr/mod.rs | 1 + .../src/llms/mistral/ocr/transformation.rs | 626 ++++++++ litellm-rust/crates/core/src/llms/mod.rs | 6 + .../crates/core/src/llms/reducto/mod.rs | 1 + .../crates/core/src/llms/reducto/ocr/mod.rs | 1 + .../src/llms/reducto/ocr/transformation.rs | 1018 +++++++++++++ .../crates/core/src/llms/vertex_ai/mod.rs | 1 + .../src/llms/vertex_ai/ocr/common_utils.rs | 9 + .../vertex_ai/ocr/deepseek_transformation.rs | 705 +++++++++ .../crates/core/src/llms/vertex_ai/ocr/mod.rs | 3 + .../src/llms/vertex_ai/ocr/transformation.rs | 395 ++++++ .../core/src/ocr/adapters/azure/cohere.rs | 131 -- .../azure/document_intelligence/mod.rs | 214 --- .../azure/document_intelligence/polling.rs | 119 -- .../core/src/ocr/adapters/azure/mistral.rs | 229 --- .../crates/core/src/ocr/adapters/cohere.rs | 123 -- .../crates/core/src/ocr/adapters/mistral.rs | 147 -- .../crates/core/src/ocr/adapters/mod.rs | 91 -- .../core/src/ocr/adapters/reducto/legacy.rs | 45 - .../core/src/ocr/adapters/reducto/mod.rs | 148 -- .../core/src/ocr/adapters/reducto/v3.rs | 45 - .../core/src/ocr/adapters/vertex/deepseek.rs | 140 -- .../core/src/ocr/adapters/vertex/mistral.rs | 157 -- .../core/src/ocr/adapters/vertex/mod.rs | 18 - litellm-rust/crates/core/src/ocr/arguments.rs | 101 ++ litellm-rust/crates/core/src/ocr/client.rs | 62 +- .../crates/core/src/ocr/codecs/cohere.rs | 254 ---- .../core/src/ocr/codecs/deepseek/mod.rs | 5 - .../src/ocr/codecs/deepseek/transformation.rs | 101 -- .../core/src/ocr/codecs/deepseek/types.rs | 95 -- .../ocr/codecs/document_intelligence/mod.rs | 9 - .../codecs/document_intelligence/params.rs | 219 --- .../document_intelligence/transformation.rs | 107 -- .../ocr/codecs/document_intelligence/types.rs | 138 -- .../crates/core/src/ocr/codecs/mistral/mod.rs | 5 - .../src/ocr/codecs/mistral/transformation.rs | 250 ---- .../core/src/ocr/codecs/mistral/types.rs | 60 - .../crates/core/src/ocr/codecs/mod.rs | 5 - .../crates/core/src/ocr/codecs/reducto/mod.rs | 9 - .../src/ocr/codecs/reducto/transformation.rs | 103 -- .../core/src/ocr/codecs/reducto/types.rs | 128 -- litellm-rust/crates/core/src/ocr/document.rs | 49 +- litellm-rust/crates/core/src/ocr/error.rs | 241 ++-- litellm-rust/crates/core/src/ocr/handler.rs | 128 +- litellm-rust/crates/core/src/ocr/hooks.rs | 22 +- litellm-rust/crates/core/src/ocr/json.rs | 62 + litellm-rust/crates/core/src/ocr/lifecycle.rs | 4 +- litellm-rust/crates/core/src/ocr/mod.rs | 21 +- litellm-rust/crates/core/src/ocr/prepare.rs | 248 ++-- .../crates/core/src/ocr/provider_config.rs | 411 ++++++ litellm-rust/crates/core/src/ocr/registry.rs | 132 -- litellm-rust/crates/core/src/ocr/types.rs | 626 +++++++- litellm-rust/crates/core/src/ocr/wire.rs | 309 +--- litellm-rust/crates/core/src/params.rs | 231 +++ litellm-rust/crates/core/src/providers/mod.rs | 1 + .../crates/core/src/providers/model.rs | 219 +++ litellm-rust/crates/core/src/serde_compat.rs | 151 ++ .../crates/core/tests/azure_ai_ocr.rs | 8 +- .../tests/azure_document_intelligence_ocr.rs | 31 +- .../crates/core/tests/deepseek_ocr.rs | 40 +- .../crates/core/tests/host_lifecycle.rs | 23 +- litellm-rust/crates/core/tests/ocr.rs | 80 +- litellm-rust/crates/core/tests/ocr/support.rs | 14 + litellm-rust/crates/core/tests/reducto_ocr.rs | 27 +- .../core/tests/vertex_ai_deepseek_ocr.rs | 20 +- .../crates/core/tests/vertex_ai_ocr.rs | 43 +- .../crates/python-bridge/src/errors.rs | 20 +- .../python-bridge/src/routes/ocr/errors.rs | 18 +- .../python-bridge/src/routes/ocr/project.rs | 2 +- 88 files changed, 8518 insertions(+), 4121 deletions(-) create mode 100644 litellm-rust/crates/core/src/call_arguments.rs create mode 100644 litellm-rust/crates/core/src/llms/azure_ai/mod.rs create mode 100644 litellm-rust/crates/core/src/llms/azure_ai/ocr/cohere_parse_transformation.rs rename litellm-rust/crates/core/src/{ocr/adapters/azure/mod.rs => llms/azure_ai/ocr/common_utils.rs} (65%) create mode 100644 litellm-rust/crates/core/src/llms/azure_ai/ocr/document_intelligence/mod.rs create mode 100644 litellm-rust/crates/core/src/llms/azure_ai/ocr/document_intelligence/transformation.rs create mode 100644 litellm-rust/crates/core/src/llms/azure_ai/ocr/mod.rs create mode 100644 litellm-rust/crates/core/src/llms/azure_ai/ocr/transformation.rs create mode 100644 litellm-rust/crates/core/src/llms/base_llm/mod.rs create mode 100644 litellm-rust/crates/core/src/llms/base_llm/ocr/mod.rs create mode 100644 litellm-rust/crates/core/src/llms/base_llm/ocr/transformation.rs create mode 100644 litellm-rust/crates/core/src/llms/cohere/mod.rs create mode 100644 litellm-rust/crates/core/src/llms/cohere/ocr/mod.rs create mode 100644 litellm-rust/crates/core/src/llms/cohere/ocr/transformation.rs create mode 100644 litellm-rust/crates/core/src/llms/mistral/mod.rs create mode 100644 litellm-rust/crates/core/src/llms/mistral/ocr/mod.rs create mode 100644 litellm-rust/crates/core/src/llms/mistral/ocr/transformation.rs create mode 100644 litellm-rust/crates/core/src/llms/mod.rs create mode 100644 litellm-rust/crates/core/src/llms/reducto/mod.rs create mode 100644 litellm-rust/crates/core/src/llms/reducto/ocr/mod.rs create mode 100644 litellm-rust/crates/core/src/llms/reducto/ocr/transformation.rs create mode 100644 litellm-rust/crates/core/src/llms/vertex_ai/mod.rs create mode 100644 litellm-rust/crates/core/src/llms/vertex_ai/ocr/common_utils.rs create mode 100644 litellm-rust/crates/core/src/llms/vertex_ai/ocr/deepseek_transformation.rs create mode 100644 litellm-rust/crates/core/src/llms/vertex_ai/ocr/mod.rs create mode 100644 litellm-rust/crates/core/src/llms/vertex_ai/ocr/transformation.rs delete mode 100644 litellm-rust/crates/core/src/ocr/adapters/azure/cohere.rs delete mode 100644 litellm-rust/crates/core/src/ocr/adapters/azure/document_intelligence/mod.rs delete mode 100644 litellm-rust/crates/core/src/ocr/adapters/azure/document_intelligence/polling.rs delete mode 100644 litellm-rust/crates/core/src/ocr/adapters/azure/mistral.rs delete mode 100644 litellm-rust/crates/core/src/ocr/adapters/cohere.rs delete mode 100644 litellm-rust/crates/core/src/ocr/adapters/mistral.rs delete mode 100644 litellm-rust/crates/core/src/ocr/adapters/mod.rs delete mode 100644 litellm-rust/crates/core/src/ocr/adapters/reducto/legacy.rs delete mode 100644 litellm-rust/crates/core/src/ocr/adapters/reducto/mod.rs delete mode 100644 litellm-rust/crates/core/src/ocr/adapters/reducto/v3.rs delete mode 100644 litellm-rust/crates/core/src/ocr/adapters/vertex/deepseek.rs delete mode 100644 litellm-rust/crates/core/src/ocr/adapters/vertex/mistral.rs delete mode 100644 litellm-rust/crates/core/src/ocr/adapters/vertex/mod.rs create mode 100644 litellm-rust/crates/core/src/ocr/arguments.rs delete mode 100644 litellm-rust/crates/core/src/ocr/codecs/cohere.rs delete mode 100644 litellm-rust/crates/core/src/ocr/codecs/deepseek/mod.rs delete mode 100644 litellm-rust/crates/core/src/ocr/codecs/deepseek/transformation.rs delete mode 100644 litellm-rust/crates/core/src/ocr/codecs/deepseek/types.rs delete mode 100644 litellm-rust/crates/core/src/ocr/codecs/document_intelligence/mod.rs delete mode 100644 litellm-rust/crates/core/src/ocr/codecs/document_intelligence/params.rs delete mode 100644 litellm-rust/crates/core/src/ocr/codecs/document_intelligence/transformation.rs delete mode 100644 litellm-rust/crates/core/src/ocr/codecs/document_intelligence/types.rs delete mode 100644 litellm-rust/crates/core/src/ocr/codecs/mistral/mod.rs delete mode 100644 litellm-rust/crates/core/src/ocr/codecs/mistral/transformation.rs delete mode 100644 litellm-rust/crates/core/src/ocr/codecs/mistral/types.rs delete mode 100644 litellm-rust/crates/core/src/ocr/codecs/mod.rs delete mode 100644 litellm-rust/crates/core/src/ocr/codecs/reducto/mod.rs delete mode 100644 litellm-rust/crates/core/src/ocr/codecs/reducto/transformation.rs delete mode 100644 litellm-rust/crates/core/src/ocr/codecs/reducto/types.rs create mode 100644 litellm-rust/crates/core/src/ocr/json.rs create mode 100644 litellm-rust/crates/core/src/ocr/provider_config.rs delete mode 100644 litellm-rust/crates/core/src/ocr/registry.rs create mode 100644 litellm-rust/crates/core/src/params.rs create mode 100644 litellm-rust/crates/core/src/providers/model.rs create mode 100644 litellm-rust/crates/core/src/serde_compat.rs diff --git a/litellm-rust/Cargo.lock b/litellm-rust/Cargo.lock index 1cc200a7bec..afa1eecc13f 100644 --- a/litellm-rust/Cargo.lock +++ b/litellm-rust/Cargo.lock @@ -948,8 +948,18 @@ version = "0.20.11" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "fc7f46116c46ff9ab3eb1597a45688b6715c6e628b5c133e288e709a29bcb4ee" dependencies = [ - "darling_core", - "darling_macro", + "darling_core 0.20.11", + "darling_macro 0.20.11", +] + +[[package]] +name = "darling" +version = "0.21.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9cdf337090841a411e2a7f3deb9187445851f91b309c0c0a29e05f74a00a48c0" +dependencies = [ + "darling_core 0.21.3", + "darling_macro 0.21.3", ] [[package]] @@ -966,13 +976,38 @@ dependencies = [ "syn 2.0.119", ] +[[package]] +name = "darling_core" +version = "0.21.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1247195ecd7e3c85f83c8d2a366e4210d588e802133e1e355180a9870b517ea4" +dependencies = [ + "fnv", + "ident_case", + "proc-macro2", + "quote", + "strsim", + "syn 2.0.119", +] + [[package]] name = "darling_macro" version = "0.20.11" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "fc34b93ccb385b40dc71c6fceac4b2ad23662c7eeb248cf10d529b7e055b6ead" dependencies = [ - "darling_core", + "darling_core 0.20.11", + "quote", + "syn 2.0.119", +] + +[[package]] +name = "darling_macro" +version = "0.21.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d38308df82d1080de0afee5d069fa14b0326a88c14f15c5ccda35b4a6c414c81" +dependencies = [ + "darling_core 0.21.3", "quote", "syn 2.0.119", ] @@ -1022,7 +1057,7 @@ version = "0.20.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "2d5bcf7b024d6835cfb3d473887cd966994907effbe9227e8c8219824d06c4e8" dependencies = [ - "darling", + "darling 0.20.11", "proc-macro2", "quote", "syn 2.0.119", @@ -1363,7 +1398,7 @@ dependencies = [ "futures-sink", "futures-util", "http 0.2.12", - "indexmap", + "indexmap 2.14.0", "slab", "tokio", "tokio-util", @@ -1382,7 +1417,7 @@ dependencies = [ "futures-core", "futures-sink", "http 1.4.2", - "indexmap", + "indexmap 2.14.0", "slab", "tokio", "tokio-util", @@ -1400,6 +1435,12 @@ dependencies = [ "zerocopy", ] +[[package]] +name = "hashbrown" +version = "0.12.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8a9ee70c43aaf417c914396645a0fa852624801b24ebb7ae78fe8272889ac888" + [[package]] name = "hashbrown" version = "0.17.1" @@ -1736,6 +1777,17 @@ dependencies = [ "icu_properties", ] +[[package]] +name = "indexmap" +version = "1.9.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bd070e393353796e801d209ad339e89596eb4c8d430d18ede6a1cced8fafbd99" +dependencies = [ + "autocfg", + "hashbrown 0.12.3", + "serde", +] + [[package]] name = "indexmap" version = "2.14.0" @@ -1743,7 +1795,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "d466e9454f08e4a911e14806c24e16fba1b4c121d1ea474396f396069cf949d9" dependencies = [ "equivalent", - "hashbrown", + "hashbrown 0.17.1", "serde", "serde_core", ] @@ -1971,6 +2023,7 @@ dependencies = [ "serde", "serde_json", "serde_path_to_error", + "serde_with", "sha2 0.10.9", "strum", "subtle", @@ -2032,7 +2085,7 @@ version = "0.1.0" dependencies = [ "base64 0.22.1", "criterion", - "indexmap", + "indexmap 2.14.0", "itoa", "rand 0.8.7", "rstest", @@ -2753,6 +2806,26 @@ dependencies = [ "bitflags", ] +[[package]] +name = "ref-cast" +version = "1.0.27" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7e440fb4e4b4147295338efb76001ab9e4efc0e5839df2c47fc5ac2381d365c3" +dependencies = [ + "ref-cast-impl", +] + +[[package]] +name = "ref-cast-impl" +version = "1.0.27" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "92ecd8964f8453721699a1ed72037b0db49ce2f5a5138486ee89bed6f67cdf3a" +dependencies = [ + "proc-macro2", + "quote", + "syn 3.0.0", +] + [[package]] name = "regex" version = "1.13.1" @@ -3075,6 +3148,30 @@ dependencies = [ "windows-sys 0.61.2", ] +[[package]] +name = "schemars" +version = "0.9.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4cd191f9397d57d581cddd31014772520aa448f65ef991055d7f61582c65165f" +dependencies = [ + "dyn-clone", + "ref-cast", + "serde", + "serde_json", +] + +[[package]] +name = "schemars" +version = "1.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "687274d293b6cdc6e73e0fee520bf2049650090d7164f87672d212a3c530cf4a" +dependencies = [ + "dyn-clone", + "ref-cast", + "serde", + "serde_json", +] + [[package]] name = "scopeguard" version = "1.2.0" @@ -3156,6 +3253,7 @@ version = "1.0.150" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "e8014e44b4736ed0538adeecded0fce2a272f22dc9578a7eb6b2d9993c74cfb9" dependencies = [ + "indexmap 2.14.0", "itoa", "memchr", "serde", @@ -3186,6 +3284,37 @@ dependencies = [ "serde", ] +[[package]] +name = "serde_with" +version = "3.16.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4fa237f2807440d238e0364a218270b98f767a00d3dada77b1c53ae88940e2e7" +dependencies = [ + "base64 0.22.1", + "chrono", + "hex", + "indexmap 1.9.3", + "indexmap 2.14.0", + "schemars 0.9.0", + "schemars 1.2.2", + "serde_core", + "serde_json", + "serde_with_macros", + "time", +] + +[[package]] +name = "serde_with_macros" +version = "3.16.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "52a8e3ca0ca629121f70ab50f95249e5a6f925cc0f6ffe8256c45b728875706c" +dependencies = [ + "darling 0.21.3", + "proc-macro2", + "quote", + "syn 2.0.119", +] + [[package]] name = "sha1" version = "0.10.7" @@ -3661,7 +3790,7 @@ version = "0.25.13+spec-1.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6975367e4d2ef766d86af01ffad14b622fecc8d4357a998fbc4deb6e9bacaf9b" dependencies = [ - "indexmap", + "indexmap 2.14.0", "toml_datetime", "toml_parser", "winnow", diff --git a/litellm-rust/Cargo.toml b/litellm-rust/Cargo.toml index 879090870d8..8f5b19f096c 100644 --- a/litellm-rust/Cargo.toml +++ b/litellm-rust/Cargo.toml @@ -29,6 +29,7 @@ rustls = { version = "0.23", default-features = false, features = ["ring", "std" rustls-native-certs = "0.8" serde = { version = "1.0", features = ["derive"] } serde_json = { version = "1.0", features = ["float_roundtrip"] } +serde_with = { version = "=3.16.1", default-features = false, features = ["std", "macros"] } sha2 = "0.10" subtle = "2" thiserror = "2.0" diff --git a/litellm-rust/crates/core/Cargo.toml b/litellm-rust/crates/core/Cargo.toml index ededfeef8af..ccca7be4971 100644 --- a/litellm-rust/crates/core/Cargo.toml +++ b/litellm-rust/crates/core/Cargo.toml @@ -22,7 +22,8 @@ reqwest.workspace = true rustls.workspace = true rustls-native-certs.workspace = true serde.workspace = true -serde_json.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 diff --git a/litellm-rust/crates/core/src/call_arguments.rs b/litellm-rust/crates/core/src/call_arguments.rs new file mode 100644 index 00000000000..67852cef27d --- /dev/null +++ b/litellm-rust/crates/core/src/call_arguments.rs @@ -0,0 +1,467 @@ +use std::ops::Deref; + +use serde::{Deserialize, Serialize, de::DeserializeOwned}; +use serde_json::{Map, Value}; + +#[derive(Clone, Debug, Default, PartialEq, Serialize, Deserialize)] +#[serde(transparent)] +pub struct CallArguments(Map); + +impl CallArguments { + pub(crate) fn select(&self, names: &[&str]) -> Map { + self.iter() + .filter(|(name, _)| names.contains(&name.as_str())) + .map(|(name, value)| (name.clone(), value.clone())) + .collect() + } +} + +#[derive(Clone, Debug, PartialEq, Eq, thiserror::Error)] +#[error("invalid argument: {path}")] +pub struct ArgumentError { + pub path: String, +} + +pub fn parse_options(arguments: &CallArguments) -> Result { + let deserializer = serde::de::value::MapDeserializer::new( + arguments.iter().map(|(name, value)| (name.as_str(), value)), + ); + serde_path_to_error::deserialize(deserializer).map_err(|error| ArgumentError { + path: error.path().to_string(), + }) +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub struct ArgumentSpec { + pub name: &'static str, + pub secret: bool, +} + +pub fn should_project(name: &str, consumed: &[ArgumentSpec], bound_fields: &[&str]) -> bool { + consumed.iter().any(|field| field.name == name) + || (!bound_fields.contains(&name) && !is_control(name)) +} + +pub fn is_control(name: &str) -> bool { + crate::params::is_control_param(name) || HOST_CONTROLS.contains(&name) +} + +const HOST_CONTROLS: &[&str] = &[ + "_agentic_loop_api_surface", + "_agentic_loop_depth", + "_agentic_loop_fingerprints", + "_code_interpreter_interception_active", + "_code_interpreter_interception_converted_stream", + "_code_interpreter_interception_sandbox_key", + "_code_interpreter_interception_session_scoped", + "_headroom_interception_converted_stream", + "_litellm_strip_stream_usage", + "_router_weights", + "_websearch_interception_converted_stream", + "_websearch_interception_emit_native_blocks", + "acompletion", + "adaptive_router_config", + "adaptive_router_default_model", + "aembedding", + "aimg_generation", + "allm_passthrough_route", + "allow_client_keepalive_override", + "allowed_model_region", + "allowed_openai_params", + "annotation_cost_per_page", + "api_version", + "arize_api_key", + "arize_space_id", + "arize_space_key", + "assistant_continue_message", + "async_call", + "atext_completion", + "attempted_targets", + "auto_router_config", + "auto_router_config_path", + "auto_router_default_model", + "auto_router_embedding_model", + "auto_router_max_input_chars", + "auto_router_model_compression", + "auto_router_routing_compression", + "aws_batch_role_arn", + "azure", + "azure_password", + "azure_username", + "base_model", + "bedrock_tags", + "bos_token", + "budget_duration", + "cache", + "cache_creation_input_audio_token_cost", + "cache_creation_input_token_cost", + "cache_creation_input_token_cost_above_1hr", + "cache_creation_input_token_cost_above_200k_tokens", + "cache_creation_input_token_cost_above_272k_tokens", + "cache_creation_input_token_cost_above_272k_tokens_flex", + "cache_creation_input_token_cost_above_272k_tokens_priority", + "cache_creation_input_token_cost_flex", + "cache_creation_input_token_cost_priority", + "cache_creation_input_token_cost_ultrafast", + "cache_key", + "cache_read_input_audio_token_cost", + "cache_read_input_token_cost", + "cache_read_input_token_cost_above_200k_tokens", + "cache_read_input_token_cost_above_200k_tokens_priority", + "cache_read_input_token_cost_above_272k_tokens", + "cache_read_input_token_cost_above_272k_tokens_flex", + "cache_read_input_token_cost_above_272k_tokens_priority", + "cache_read_input_token_cost_above_512k_tokens", + "cache_read_input_token_cost_flex", + "cache_read_input_token_cost_priority", + "cache_read_input_token_cost_ultrafast", + "caching", + "caching_groups", + "citation_cost_per_token", + "client", + "client_side_timeout", + "complete_response", + "completion_call_id", + "complexity_router_config", + "complexity_router_default_model", + "configurable_clientside_auth_params", + "context_window_fallback_dict", + "cooldown_time", + "cost_per_query", + "custom_prompt_dict", + "data_residency", + "dd_agent_host", + "dd_agent_port", + "dd_api_key", + "dd_site", + "default_api_key_rpm_limit", + "default_api_key_tpm_limit", + "disable_add_transform_inline_image_block", + "enable_json_schema_validation", + "enable_prompt_caching", + "enable_tag_filtering", + "ensure_alternating_roles", + "eos_token", + "fallback_depth", + "fallbacks", + "fastest_response", + "final_prompt_value", + "force_timeout", + "gcs_bucket_name", + "gcs_path_service_account", + "google_maps_grounding_cost_per_query", + "headers", + "hf_model_name", + "humanloop_api_key", + "id", + "input_cost_per_audio_per_second", + "input_cost_per_audio_per_second_above_128k_tokens", + "input_cost_per_audio_token", + "input_cost_per_audio_token_batches", + "input_cost_per_character", + "input_cost_per_character_above_128k_tokens", + "input_cost_per_image", + "input_cost_per_image_above_128k_tokens", + "input_cost_per_image_token", + "input_cost_per_image_token_batches", + "input_cost_per_pixel", + "input_cost_per_query", + "input_cost_per_second", + "input_cost_per_token", + "input_cost_per_token_above_128k_tokens", + "input_cost_per_token_above_200k_tokens", + "input_cost_per_token_above_200k_tokens_priority", + "input_cost_per_token_above_272k_tokens", + "input_cost_per_token_above_272k_tokens_flex", + "input_cost_per_token_above_272k_tokens_priority", + "input_cost_per_token_above_512k_tokens", + "input_cost_per_token_batches", + "input_cost_per_token_cache_hit", + "input_cost_per_token_flex", + "input_cost_per_token_priority", + "input_cost_per_token_ultrafast", + "input_cost_per_video_per_second", + "input_cost_per_video_per_second_above_128k_tokens", + "input_cost_per_video_per_second_above_15s_interval", + "input_cost_per_video_per_second_above_8s_interval", + "input_cost_per_video_token", + "input_cost_per_video_token_batches", + "itpm", + "keepalive_seconds", + "langfuse_environment", + "langfuse_host", + "langfuse_prompt_version", + "langfuse_public_key", + "langfuse_secret", + "langfuse_secret_key", + "langsmith_api_key", + "langsmith_base_url", + "langsmith_project", + "langsmith_sampling_rate", + "langsmith_tenant_id", + "litellm_credential_name", + "litellm_disabled_callbacks", + "litellm_request_debug", + "litellm_session_id", + "litellm_system_prompt", + "litellm_trace_id", + "litellm_trusted_callback_vars", + "logger_fn", + "max_agentic_loops", + "max_budget", + "max_fallbacks", + "max_parallel_requests", + "merge_reasoning_content_in_choices", + "metadata", + "mock_response", + "mock_timeout", + "model_alias_map", + "model_config", + "model_file_id_mapping", + "model_info", + "model_list", + "newrelic_api_key", + "newrelic_region", + "no-log", + "num_retries", + "ocr_cost_per_credit", + "ocr_cost_per_page", + "order", + "otpm", + "output_cost_per_audio_per_second", + "output_cost_per_audio_token", + "output_cost_per_character", + "output_cost_per_character_above_128k_tokens", + "output_cost_per_image", + "output_cost_per_image_token", + "output_cost_per_pixel", + "output_cost_per_reasoning_token", + "output_cost_per_reasoning_token_flex", + "output_cost_per_reasoning_token_priority", + "output_cost_per_second", + "output_cost_per_second_1080p", + "output_cost_per_second_480p", + "output_cost_per_second_4k", + "output_cost_per_second_720p", + "output_cost_per_token", + "output_cost_per_token_above_128k_tokens", + "output_cost_per_token_above_200k_tokens", + "output_cost_per_token_above_200k_tokens_priority", + "output_cost_per_token_above_272k_tokens", + "output_cost_per_token_above_272k_tokens_flex", + "output_cost_per_token_above_272k_tokens_priority", + "output_cost_per_token_above_512k_tokens", + "output_cost_per_token_batches", + "output_cost_per_token_flex", + "output_cost_per_token_priority", + "output_cost_per_token_ultrafast", + "output_cost_per_video_per_second", + "output_cost_per_video_token", + "output_vector_size", + "posthog_api_key", + "posthog_api_url", + "preset_cache_key", + "prompt_environment", + "prompt_id", + "prompt_label", + "prompt_variables", + "prompt_version", + "provider_specific_header", + "quality_router_config", + "quality_router_default_model", + "region_name", + "regional_endpoint_uplift_multiplier", + "regional_processing_uplift_multiplier_eu", + "regional_processing_uplift_multiplier_us", + "retry_policy", + "retry_strategy", + "roles", + "routing_strategy", + "rpm", + "rust", + "s3_bucket_name", + "s3_output_bucket_name", + "s3_region_name", + "search_context_cost_per_query", + "search_tool_name", + "secret_fields", + "self", + "shared_session", + "ssl_verify", + "stream_response", + "stream_timeout", + "supports_system_message", + "tags", + "text_completion", + "tiered_pricing", + "tpm", + "ttl", + "turn_off_message_logging", + "use_chat_completions_api", + "use_client", + "use_in_pass_through", + "use_litellm_proxy", + "use_xai_oauth", + "user_continue_message", + "verbose", + "wandb_api_key", + "weave_project_id", + "weight", +]; + +pub fn compose_body( + arguments: &CallArguments, + body: &B, + consumed: &[&str], +) -> Result { + let Value::Object(fields) = + serde_json::to_value(body).map_err(|_| crate::params::Error::Body)? + else { + return Err(crate::params::Error::Body); + }; + let overrides = match arguments.get("extra_body") { + None | Some(Value::Null) => None, + Some(Value::Object(fields)) => Some(fields), + Some(_) => return Err(crate::params::Error::ExtraBody), + }; + let extensions = arguments.iter().filter(|(name, _)| { + !consumed.contains(&name.as_str()) && name.as_str() != "extra_body" && !is_control(name) + }); + Ok(Value::Object( + fields + .into_iter() + .chain( + extensions + .chain(overrides.into_iter().flatten()) + .filter(|(name, _)| { + name.as_str() != "model" + && name.as_str() != "extra_body" + && !crate::params::is_control_param(name) + }) + .map(|(name, value)| (name.clone(), value.clone())), + ) + .collect(), + )) +} + +impl Deref for CallArguments { + type Target = Map; + + fn deref(&self) -> &Self::Target { + &self.0 + } +} + +impl From> for CallArguments { + fn from(values: Map) -> Self { + Self(values) + } +} + +impl From for Map { + fn from(arguments: CallArguments) -> Self { + arguments.0 + } +} + +impl FromIterator<(String, Value)> for CallArguments { + fn from_iter>(iter: T) -> Self { + Self(iter.into_iter().collect()) + } +} + +impl IntoIterator for CallArguments { + type Item = (String, Value); + type IntoIter = serde_json::map::IntoIter; + + fn into_iter(self) -> Self::IntoIter { + self.0.into_iter() + } +} + +#[cfg(test)] +mod tests { + use super::*; + use serde_json::json; + + #[test] + fn composition_preserves_extensions_and_applies_shallow_explicit_overrides() { + let original = json!({ + "known": false, "future": {"old": 1}, "null": null, "zero": 0, + "metadata": {"host": true}, "shared_session": "host", "api_key": "secret", + "extra_body": { + "known": null, "future": {"new": [false, 0, null]}, + "metadata": {"provider": true}, "model": "ignored", "api_key": "ignored" + } + }); + let arguments = serde_json::from_value(original.clone()).unwrap(); + let body = compose_body( + &arguments, + &json!({"model":"resolved", "known":false}), + &["known"], + ) + .unwrap(); + assert_eq!( + body, + json!({ + "model":"resolved", "known":null, "future":{"new":[false,0,null]}, + "null":null, "zero":0, "metadata":{"provider":true} + }) + ); + assert_eq!(serde_json::to_value(arguments).unwrap(), original); + } + + #[test] + fn projection_prioritizes_consumed_fields_and_keeps_unknown_names() { + let fields = [ArgumentSpec { + name: "id", + secret: false, + }]; + assert!(should_project("id", &fields, &[])); + assert!(!should_project("id", &[], &[])); + assert!(should_project("future_option", &[], &[])); + assert!(!should_project("document", &fields, &["document"])); + assert!(!should_project("metadata", &fields, &[])); + assert!(!should_project("callbacks", &fields, &[])); + assert!(!should_project("ocr_cost_per_page", &fields, &[])); + } + + #[test] + fn invalid_extra_body_is_rejected_without_coercing_it_to_empty() { + for value in [json!(false), json!(0), json!([]), json!("")] { + let arguments = serde_json::from_value(json!({"extra_body":value})).unwrap(); + assert_eq!( + compose_body(&arguments, &json!({}), &[]), + Err(crate::params::Error::ExtraBody) + ); + } + let arguments = serde_json::from_value(json!({"extra_body":null})).unwrap(); + assert_eq!( + compose_body(&arguments, &json!({}), &[]).unwrap(), + json!({}) + ); + } + + #[test] + fn typed_views_preserve_missing_and_explicit_null_in_the_source() { + #[derive(Deserialize)] + struct Options { + enabled: Option, + } + let arguments: CallArguments = + serde_json::from_value(json!({"enabled":null,"future":0})).unwrap(); + assert!( + parse_options::(&arguments) + .unwrap() + .enabled + .is_none() + ); + assert_eq!(arguments.get("enabled"), Some(&Value::Null)); + assert_eq!(arguments.get("missing"), None); + let invalid = serde_json::from_value(json!({"enabled":0})).unwrap(); + assert_eq!( + parse_options::(&invalid).err().unwrap().path, + "enabled" + ); + } +} diff --git a/litellm-rust/crates/core/src/lib.rs b/litellm-rust/crates/core/src/lib.rs index b028b7bc9b1..288bde52ce4 100644 --- a/litellm-rust/crates/core/src/lib.rs +++ b/litellm-rust/crates/core/src/lib.rs @@ -1,14 +1,18 @@ pub mod audio_transcription; +pub mod call_arguments; pub mod call_lifecycle; pub mod chat_completions; pub mod constants; pub mod error; pub mod http_utils; +pub(crate) mod llms; mod media; pub mod messages; pub mod ocr; +pub mod params; pub mod providers; pub mod responses; +mod serde_compat; pub mod transport; mod url_utils; diff --git a/litellm-rust/crates/core/src/llms/azure_ai/mod.rs b/litellm-rust/crates/core/src/llms/azure_ai/mod.rs new file mode 100644 index 00000000000..079e0c41eae --- /dev/null +++ b/litellm-rust/crates/core/src/llms/azure_ai/mod.rs @@ -0,0 +1 @@ +pub(crate) mod ocr; 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 new file mode 100644 index 00000000000..add70c2596d --- /dev/null +++ b/litellm-rust/crates/core/src/llms/azure_ai/ocr/cohere_parse_transformation.rs @@ -0,0 +1,165 @@ +use crate::call_arguments::CallArguments; +use crate::llms::base_llm::ocr::transformation::{BaseOcrConfig, OcrRequestContext}; +use crate::llms::cohere::ocr::transformation::{CohereParseConfig, CohereRequest}; +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::types::{LiteLLMOcrResponse, OcrDocument, PreparedOcrRequest}; +use crate::url_utils::ApiUrl; +use serde_json::Value; + +#[derive(Default)] +pub(crate) struct AzureAICohereParseConfig; + +impl BaseOcrConfig for AzureAICohereParseConfig { + type OcrParams = CohereOptions; + type ProviderRequest = CohereRequest; + 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: &PreparedOcrRequest, + client: &OcrClient, + ) -> Result { + BaseOcrConfig::validate_environment( + &super::transformation::AzureAIOCRConfig, + request, + client, + ) + .await + } + + fn get_complete_url( + &self, + request: &PreparedOcrRequest, + _params: &Self::OcrParams, + _environment: &Self::Environment, + ) -> Result { + let base = super::transformation::AzureAIOCRConfig::resolve_api_base( + request.connection.api_base.as_deref(), + &crate::ocr::prepare::credential_env, + )?; + 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) + } + + fn map_ocr_params( + &self, + arguments: &CallArguments, + model: &str, + ) -> Result { + CohereParseConfig.map_ocr_params(arguments, model) + } + + async fn async_transform_ocr_request( + &self, + model: &str, + document: OcrDocument, + optional_params: &CohereOptions, + headers: &[(String, String)], + context: OcrRequestContext<'_>, + ) -> Result { + validate_document(&document)?; + let document = inline_remote_document( + context.client.document_fetcher(), + document, + context.connection, + ) + .await?; + self.transform_ocr_request(model, document, optional_params, headers) + } + + fn transform_ocr_response( + &self, + model: &str, + raw_response: &[u8], + request_format: crate::ocr::types::OcrResponseFormat, + ) -> Result { + CohereParseConfig.transform_ocr_response(model, raw_response, request_format) + } + + fn validate_request_body(&self, body: &Value) -> Result<(), crate::ocr::Error> { + let document = crate::ocr::prepare::body_document(body)?; + validate_document(&document)?; + validate_inline_document(&document) + } +} + +impl AzureAICohereParseConfig { + fn get_complete_url(&self, base: &str) -> Result { + let mut url = reqwest::Url::parse(base).map_err(|_| invalid_api_base())?; + if !matches!(url.scheme(), "http" | "https") { + return Err(invalid_api_base()); + } + let path = url.path().trim_end_matches('/').to_string(); + if path.ends_with("/v2/parse") { + url.set_path(&path); + return Ok(url.into()); + } + url.set_path(path.strip_suffix("/models").unwrap_or(&path)); + ApiUrl::parse(url.as_str()) + .and_then(|url| url.complete_path(&["providers", "cohere", "v2", "parse"])) + .map(|url| url.into_string()) + .map_err(|_| invalid_api_base()) + } +} + +fn invalid_api_base() -> crate::ocr::Error { + crate::ocr::Error::RequestField { + path: "api_base".into(), + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn completes_foundry_urls_without_duplicate_paths_and_preserves_queries() { + for suffix in [ + "", + "/models", + "/providers/cohere/v2", + "/providers/cohere/v2/parse", + ] { + assert_eq!( + AzureAICohereParseConfig + .get_complete_url(&format!("https://example.com{suffix}?tenant=a")) + .unwrap(), + "https://example.com/providers/cohere/v2/parse?tenant=a" + ); + } + assert_eq!( + AzureAICohereParseConfig + .get_complete_url("https://example.com/v2/parse?tenant=a") + .unwrap(), + "https://example.com/v2/parse?tenant=a" + ); + assert!( + AzureAICohereParseConfig + .get_complete_url("relative/path") + .is_err() + ); + } +} diff --git a/litellm-rust/crates/core/src/ocr/adapters/azure/mod.rs b/litellm-rust/crates/core/src/llms/azure_ai/ocr/common_utils.rs similarity index 65% rename from litellm-rust/crates/core/src/ocr/adapters/azure/mod.rs rename to litellm-rust/crates/core/src/llms/azure_ai/ocr/common_utils.rs index 0b2fcb0f4cb..c381e39eaae 100644 --- a/litellm-rust/crates/core/src/ocr/adapters/azure/mod.rs +++ b/litellm-rust/crates/core/src/llms/azure_ai/ocr/common_utils.rs @@ -1,25 +1,13 @@ -mod cohere; -mod document_intelligence; -mod mistral; - use std::sync::OnceLock; -use crate::ocr::Error; - -use crate::ocr::error::OcrError; use crate::ocr::types::OcrConnection; use litellm_auth::{InputSource, Sourced}; use litellm_auth_azure::{AzureAuthInputs, AzureAuthService}; -pub(crate) use cohere::AzureCohereAdapter; -pub(crate) use document_intelligence::AzureDocumentIntelligenceAdapter; -pub(crate) use mistral::AzureMistralAdapter; -pub(super) use mistral::validate_environment as validate_ai_environment; - -async fn resolve_entra( +pub(super) async fn resolve_entra( config: &AzureAuthInputs, env_lookup: &(dyn Fn(&str) -> Option + Sync), -) -> Result>, Error> { +) -> Result>, crate::ocr::Error> { static SERVICE: OnceLock = OnceLock::new(); SERVICE .get_or_init(AzureAuthService::default) @@ -36,18 +24,18 @@ async fn resolve_entra( Sourced::new(value, source) }) }) - .map_err(Error::from) + .map_err(crate::ocr::Error::from) } -fn validate_destination( +pub(super) fn validate_destination( connection: &OcrConnection, credential_source: InputSource, -) -> Result<(), OcrError> { +) -> Result<(), crate::ocr::Error> { if connection.api_base.is_some() && connection.api_base_source == InputSource::Request && credential_source != InputSource::Request { - return Err(Error::from(litellm_auth::Error::RequestAzureCredentialDestination).into()); + return Err(litellm_auth::Error::RequestAzureCredentialDestination.into()); } Ok(()) } diff --git a/litellm-rust/crates/core/src/llms/azure_ai/ocr/document_intelligence/mod.rs b/litellm-rust/crates/core/src/llms/azure_ai/ocr/document_intelligence/mod.rs new file mode 100644 index 00000000000..080f0a1183f --- /dev/null +++ b/litellm-rust/crates/core/src/llms/azure_ai/ocr/document_intelligence/mod.rs @@ -0,0 +1 @@ +pub(crate) mod transformation; 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 new file mode 100644 index 00000000000..ae13944c06b --- /dev/null +++ b/litellm-rust/crates/core/src/llms/azure_ai/ocr/document_intelligence/transformation.rs @@ -0,0 +1,1260 @@ +use std::collections::BTreeSet; +use std::sync::Arc; +use std::time::Duration; + +use base64::{Engine, engine::general_purpose::STANDARD}; +use reqwest::Url; +use serde::{Deserialize, Deserializer, Serialize}; +use serde_json::{Map, Value}; +use serde_with::serde_as; +use tokio::time::Instant; + +use litellm_auth::{InputSource, Sourced}; +use litellm_auth_azure::AzureAuthInputs; + +use crate::call_arguments::CallArguments; +use crate::constants::{ + AZURE_DI_API_VERSION, AZURE_DI_DEFAULT_DPI, AZURE_DI_DEFAULT_HEIGHT, AZURE_DI_DEFAULT_WIDTH, + AZURE_DI_SUBSCRIPTION_HEADER, OCR_POLL_RETRY_SECS, +}; +use crate::llms::base_llm::ocr::transformation::{ + BaseOcrConfig, OcrRequestContext, OcrResponseContext, +}; +use crate::ocr::OcrClient; +use crate::ocr::client::read_json_response; +use crate::ocr::document::InlineDocument; +use crate::ocr::hooks::OcrHooks; +use crate::ocr::json::DecodedOcrResponse; +use crate::ocr::prepare::credential_env; +use crate::ocr::types::{ + LiteLLMOcrResponse, OcrConnection, OcrCredentialInputs, OcrDocument, OcrPage, + OcrPageDimensions, OcrResponseFormat, OcrUsageInfo, PreparedOcrRequest, ResolvedOcrCredentials, +}; +use crate::serde_compat::{FiniteF64, LaxI64}; +use crate::url_utils::ApiUrl; + +#[derive(Clone, Debug, PartialEq, Serialize)] +pub(crate) struct DocumentIntelligenceParams { + #[serde(skip_serializing_if = "Option::is_none")] + pub pages: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub features: Option, +} + +#[derive(Clone, Debug, Serialize, Deserialize)] +#[serde(untagged)] +pub(crate) enum DocumentIntelligenceRequest { + UrlSource { + #[serde(rename = "urlSource")] + url_source: String, + }, + Base64Source { + #[serde(rename = "base64Source")] + base64_source: String, + }, +} + +#[derive(Clone, Debug, PartialEq)] +enum OperationStatus { + Succeeded, + Running, + NotStarted, + Failed, + Unknown(String), +} + +impl<'de> Deserialize<'de> for OperationStatus { + fn deserialize>(deserializer: D) -> Result { + Ok(match String::deserialize(deserializer)?.as_str() { + "succeeded" => Self::Succeeded, + "running" => Self::Running, + "notStarted" => Self::NotStarted, + "failed" => Self::Failed, + value => Self::Unknown(value.to_string()), + }) + } +} + +impl std::fmt::Display for OperationStatus { + fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + formatter.write_str(match self { + Self::Succeeded => "succeeded", + Self::Running => "running", + Self::NotStarted => "notStarted", + Self::Failed => "failed", + Self::Unknown(value) => value, + }) + } +} + +#[derive(Clone, Debug, Deserialize)] +pub(crate) struct AzureDocumentIntelligenceOperation { + status: Option, + #[serde(rename = "analyzeResult")] + analyze_result: Option, +} + +#[derive(Clone, Debug, Default, Deserialize)] +struct AzureDocumentIntelligenceAnalyzeResult { + pub content: Option, + #[serde(default)] + pub pages: Vec, + pub tables: Option>>, + #[serde(rename = "keyValuePairs")] + pub key_value_pairs: Option>>, +} + +#[serde_as] +#[derive(Clone, Debug, Deserialize)] +struct AzureDocumentIntelligencePage { + #[serde(rename = "pageNumber")] + #[serde_as(deserialize_as = "Option")] + pub page_number: Option, + #[serde_as(deserialize_as = "Option")] + pub width: Option, + #[serde_as(deserialize_as = "Option")] + pub height: Option, + pub unit: Option, + #[serde(default)] + pub lines: Vec, +} + +#[derive(Clone, Debug, Deserialize)] +struct AzureDocumentIntelligenceLine { + pub content: Option, +} + +fn normalize_pages(pages: Option<&Value>) -> Result, crate::ocr::Error> { + let normalized = match pages { + None | Some(Value::Null) => return Ok(None), + Some(Value::Array(pages)) if pages.is_empty() => return Ok(None), + Some(Value::Array(pages)) if pages.iter().all(Value::is_number) => pages + .iter() + .map(|page| { + let page = page + .as_i64() + .ok_or_else(|| crate::ocr::Error::Pages("page index is out of range".into()))?; + if page < 0 { + return Err(crate::ocr::Error::Pages("negative page index".into())); + } + page.checked_add(1) + .ok_or_else(|| crate::ocr::Error::Pages("page index is out of range".into())) + }) + .collect::, _>>()? + .into_iter() + .map(|page| page.to_string()) + .collect::>() + .join(","), + Some(Value::Array(tokens)) => tokens + .iter() + .map(|token| { + token.as_str().map(str::trim).ok_or_else(|| { + crate::ocr::Error::Pages("expected only integers or only strings".into()) + }) + }) + .collect::, _>>()? + .join(","), + Some(Value::String(range)) => range + .split(',') + .map(str::trim) + .collect::>() + .join(","), + Some(_) => { + return Err(crate::ocr::Error::Pages( + "expected an array of integers or strings, or a native page range".into(), + )); + } + }; + if !normalized.split(',').all(valid_page_token) { + return Err(crate::ocr::Error::Pages("invalid native page range".into())); + } + Ok(Some(normalized)) +} + +fn valid_page_token(token: &str) -> bool { + let mut parts = token.split('-'); + let start = parts.next().unwrap_or_default(); + if start.is_empty() || !start.chars().all(|character| character.is_ascii_digit()) { + return false; + } + match parts.next() { + None => true, + Some(end) => { + !end.is_empty() + && end.chars().all(|character| character.is_ascii_digit()) + && parts.next().is_none() + } + } +} + +fn normalize_features(features: Option<&Value>) -> Result, crate::ocr::Error> { + let tokens = match features { + None | Some(Value::Null) => return Ok(None), + Some(Value::Array(names)) => names + .iter() + .map(|name| name.as_str().ok_or(crate::ocr::Error::Features)) + .collect::, _>>()?, + Some(Value::String(names)) => names.split(',').collect(), + Some(_) => return Err(crate::ocr::Error::Features), + }; + if tokens.is_empty() { + return Ok(None); + } + let normalized = tokens.iter().map(|token| token.trim()).collect::>(); + if !normalized.iter().all(|token| { + let Some((first, rest)) = token.as_bytes().split_first() else { + return false; + }; + first.is_ascii_alphabetic() && rest.iter().all(u8::is_ascii_alphanumeric) + }) { + return Err(crate::ocr::Error::Features); + } + Ok(Some(normalized.join(","))) +} + +fn build_request(document: OcrDocument) -> Result { + let source = document.source(); + if source.is_empty() { + return Err(crate::ocr::Error::MissingDocumentUrl); + } + Ok(if let Some(document) = InlineDocument::parse(source)? { + DocumentIntelligenceRequest::Base64Source { + base64_source: STANDARD + .encode(document.decode(crate::constants::OCR_INLINE_MAX_BYTES)?), + } + } else { + DocumentIntelligenceRequest::UrlSource { + url_source: source.to_string(), + } + }) +} + +fn transform_completed_response( + model: &str, + response: AzureDocumentIntelligenceOperation, +) -> Result { + if response.status != Some(OperationStatus::Succeeded) { + return Err(crate::ocr::Error::OperationStatus( + response + .status + .map(|status| status.to_string()) + .unwrap_or_else(|| "None".into()), + )); + } + let result = response.analyze_result.unwrap_or_default(); + let pages = result + .pages + .into_iter() + .map(transform_azure_page) + .collect::, _>>()?; + let pages_processed = + i64::try_from(pages.len()).map_err(|_| crate::ocr::Error::NumericRange("pages"))?; + Ok(LiteLLMOcrResponse { + content: result.content, + tables: result.tables, + key_value_pairs: result.key_value_pairs, + usage_info: Some(OcrUsageInfo { + pages_processed: Some(pages_processed), + ..Default::default() + }), + ..LiteLLMOcrResponse::new(model, pages) + }) +} + +fn transform_azure_page(page: AzureDocumentIntelligencePage) -> Result { + let index = page + .page_number + .unwrap_or(1) + .checked_sub(1) + .ok_or(crate::ocr::Error::NumericRange("page.pageNumber"))?; + let dimensions = convert_dimensions( + page.width.unwrap_or(AZURE_DI_DEFAULT_WIDTH), + page.height.unwrap_or(AZURE_DI_DEFAULT_HEIGHT), + page.unit.as_deref().unwrap_or("inch"), + )?; + let markdown = page + .lines + .iter() + .map(|line| line.content.as_deref().unwrap_or_default()) + .collect::>() + .join("\n"); + Ok(OcrPage { + index, + markdown, + dimensions: Some(dimensions), + ..Default::default() + }) +} + +fn convert_dimensions( + width: f64, + height: f64, + unit: &str, +) -> Result { + let scale = if unit == "inch" { + AZURE_DI_DEFAULT_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), + }) +} + +fn pixel_dimension(value: f64, scale: f64, field: &'static str) -> Result { + let value = value * scale; + if !value.is_finite() || value < i64::MIN as f64 || value >= -(i64::MIN as f64) { + return Err(crate::ocr::Error::NumericRange(field)); + } + Ok(value.trunc() as i64) +} + +async fn read_operation_response( + http_client: &reqwest::Client, + response: reqwest::Response, + original_url: &str, + headers: &[(String, String)], + connection: &OcrConnection, + native: bool, + hooks: &Arc, +) -> Result, crate::ocr::Error> { + if response.status() != reqwest::StatusCode::ACCEPTED { + let bytes = + crate::ocr::client::read_response_bytes(response, connection.max_response_bytes) + .await?; + crate::ocr::handler::post_call(hooks, &bytes).await?; + return crate::ocr::json::decode_response(&bytes, native); + } + let location = response + .headers() + .get("operation-location") + .and_then(|value| value.to_str().ok()) + .ok_or(crate::ocr::Error::PollLocation)? + .to_string(); + let original = Url::parse(original_url).map_err(|_| crate::ocr::Error::PollOrigin)?; + let operation = Url::parse(&location).map_err(|_| crate::ocr::Error::PollOrigin)?; + if original.origin() != operation.origin() + || !operation.username().is_empty() + || operation.password().is_some() + { + return Err(crate::ocr::Error::PollOrigin); + } + let bytes = + crate::ocr::client::read_response_bytes(response, connection.max_response_bytes).await?; + crate::ocr::handler::post_call(hooks, &bytes).await?; + poll_operation(http_client, operation, headers, connection, native).await +} + +async fn poll_operation( + http_client: &reqwest::Client, + url: Url, + headers: &[(String, String)], + connection: &OcrConnection, + native: bool, +) -> Result, crate::ocr::Error> { + let deadline = Instant::now() + .checked_add(connection.poll_timeout) + .ok_or(crate::ocr::Error::PollTimeout)?; + + loop { + let remaining = deadline + .checked_duration_since(Instant::now()) + .filter(|remaining| !remaining.is_zero()) + .ok_or(crate::ocr::Error::PollTimeout)?; + let builder = http_client + .get(url.clone()) + .timeout(remaining.min(connection.timeout)); + let builder = crate::http_utils::with_headers( + builder, + headers, + crate::http_utils::HeaderPolicy::Only(&[AZURE_DI_SUBSCRIPTION_HEADER, "authorization"]), + ); + let response = tokio::time::timeout_at(deadline, crate::http_utils::http_request(builder)) + .await + .map_err(|_| crate::ocr::Error::PollTimeout)? + .map_err(crate::transport::Error::from)?; + let retry = response + .headers() + .get(reqwest::header::RETRY_AFTER) + .and_then(|value| value.to_str().ok()) + .and_then(|value| value.parse::().ok()) + .unwrap_or(OCR_POLL_RETRY_SECS) + .max(1); + let decoded = tokio::time::timeout_at( + deadline, + read_json_response::( + response, + native, + connection.max_response_bytes, + ), + ) + .await + .map_err(|_| crate::ocr::Error::PollTimeout)??; + match &decoded.data.status { + Some(OperationStatus::Succeeded) => return Ok(decoded), + Some(OperationStatus::Running | OperationStatus::NotStarted) => { + tokio::time::timeout_at(deadline, tokio::time::sleep(Duration::from_secs(retry))) + .await + .map_err(|_| crate::ocr::Error::PollTimeout)?; + } + status => { + return Err(crate::ocr::Error::OperationStatus( + status + .as_ref() + .map(ToString::to_string) + .unwrap_or_else(|| "None".into()), + )); + } + } + } +} + +const AZURE_DI_API_KEY_ENV: &str = "AZURE_DOCUMENT_INTELLIGENCE_API_KEY"; +const AZURE_DI_ENDPOINT_ENV: &str = "AZURE_DOCUMENT_INTELLIGENCE_ENDPOINT"; + +#[derive(Clone, Debug)] +pub(crate) struct AzureDocumentIntelligenceOCRConfig; + +impl BaseOcrConfig for AzureDocumentIntelligenceOCRConfig { + type OcrParams = DocumentIntelligenceParams; + type ProviderRequest = DocumentIntelligenceRequest; + 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, inputs: OcrCredentialInputs) -> ResolvedOcrCredentials { + ResolvedOcrCredentials { + api_key: inputs.api_key.and_then(|key| { + inputs + .dynamic_api_key + .filter(|value| !value.value().is_empty()) + .or(Some(key)) + }), + api_base: inputs.api_base.and_then(|base| { + inputs + .dynamic_api_base + .filter(|value| !value.value().is_empty()) + .or(Some(base)) + }), + } + } + + async fn validate_environment( + &self, + request: &PreparedOcrRequest, + _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: &PreparedOcrRequest, + 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"] + } + + fn map_ocr_params( + &self, + arguments: &CallArguments, + _model: &str, + ) -> Result { + Ok(DocumentIntelligenceParams { + pages: normalize_pages(arguments.get("pages"))?, + features: normalize_features(arguments.get("features"))?, + }) + } + + async fn async_transform_ocr_request( + &self, + model: &str, + document: OcrDocument, + optional_params: &DocumentIntelligenceParams, + headers: &[(String, String)], + _context: OcrRequestContext<'_>, + ) -> Result { + self.transform_ocr_request(model, document, optional_params, headers) + } + + fn transform_ocr_response( + &self, + model: &str, + raw_response: &[u8], + request_format: OcrResponseFormat, + ) -> Result { + crate::llms::base_llm::ocr::transformation::decode_and_normalize_response( + model, + raw_response, + request_format, + transform_completed_response, + ) + } + + async fn async_transform_ocr_response( + &self, + model: &str, + raw_response: reqwest::Response, + context: OcrResponseContext<'_>, + ) -> Result { + let decoded = read_operation_response( + context.client.polling_http(), + raw_response, + context.url, + context.headers, + context.connection, + context.request_format == OcrResponseFormat::Native, + context.hooks, + ) + .await?; + Ok(LiteLLMOcrResponse { + provider_native_response: decoded.native, + ..transform_completed_response(model, decoded.data)? + }) + } + fn transform_ocr_request( + &self, + _model: &str, + document: OcrDocument, + _optional_params: &DocumentIntelligenceParams, + _headers: &[(String, String)], + ) -> Result { + build_request(document) + } +} + +impl AzureDocumentIntelligenceOCRConfig { + fn get_complete_url( + &self, + endpoint: &str, + model: &str, + params: &DocumentIntelligenceParams, + ) -> Result { + let model = format!("{}:analyze", model_id(model)?); + ApiUrl::parse(endpoint) + .and_then(|url| url.complete_path(&["documentintelligence", "documentModels", &model])) + .map(|url| { + url.append_query_pairs( + [("api-version", AZURE_DI_API_VERSION)] + .into_iter() + .chain(params.pages.iter().map(|pages| ("pages", pages.as_str()))) + .chain( + params + .features + .iter() + .map(|features| ("features", features.as_str())), + ), + ) + .into_string() + }) + .map_err(|_| crate::ocr::Error::RequestField { + path: "api_base".into(), + }) + } + + async fn validate_environment( + &self, + connection: &OcrConnection, + config: &AzureAuthInputs, + env_lookup: &(dyn Fn(&str) -> Option + Sync), + ) -> Result, crate::ocr::Error> { + if crate::http_utils::has_header(&connection.extra_headers, "authorization") + || crate::http_utils::has_header( + &connection.extra_headers, + AZURE_DI_SUBSCRIPTION_HEADER, + ) + { + super::super::common_utils::validate_destination( + connection, + connection.extra_headers_source, + )?; + return Ok(connection.extra_headers.clone()); + } + let key = nonblank(connection.api_key.clone()) + .map(|value| Sourced::new(value, connection.api_key_source)) + .or_else(|| { + nonblank(self.get_api_key_env_var().and_then(env_lookup)) + .map(|value| Sourced::new(value, InputSource::Environment)) + }); + if let Some(key) = key { + super::super::common_utils::validate_destination(connection, key.source())?; + return Ok( + std::iter::once((AZURE_DI_SUBSCRIPTION_HEADER.into(), key.into_value())) + .chain(connection.extra_headers.clone()) + .collect(), + ); + } + let token = super::super::common_utils::resolve_entra(config, env_lookup) + .await? + .ok_or(crate::ocr::Error::MissingAzureDocumentIntelligenceCredentials)?; + super::super::common_utils::validate_destination(connection, token.source())?; + Ok( + std::iter::once(("Authorization".into(), format!("Bearer {}", token.value()))) + .chain(connection.extra_headers.clone()) + .collect(), + ) + } +} + +fn model_id(model: &str) -> Result<&str, crate::ocr::Error> { + let model = model.rsplit('/').next().unwrap_or(model); + if matches!(model, "." | "..") { + return Err(crate::ocr::Error::DotModel); + } + Ok(model) +} + +fn nonblank(value: Option) -> Option { + value + .map(|value| value.trim().to_string()) + .filter(|value| !value.is_empty()) +} + +#[cfg(test)] +mod tests { + use super::*; + use rstest::rstest; + use serde_json::{Value, json}; + + fn map(value: Value) -> Result { + let arguments = serde_json::from_value(value).unwrap(); + AzureDocumentIntelligenceOCRConfig.map_ocr_params(&arguments, "model") + } + + #[test] + fn empty_options_do_not_create_query_fields() { + let overrides = + serde_json::from_value(json!({"pages":[], "features":null, "req_format":"native"})) + .unwrap(); + let mapped = AzureDocumentIntelligenceOCRConfig + .map_ocr_params(&overrides, "model") + .unwrap(); + assert_eq!(serde_json::to_value(mapped).unwrap(), json!({})); + } + + #[test] + fn input_params_retain_unknown_fields() { + let arguments = serde_json::from_value(json!({ + "pages": [0], + "future_ocr_option": true, + "extra_body": {"provider_option": "value"} + })) + .unwrap(); + let mapped = AzureDocumentIntelligenceOCRConfig + .map_ocr_params(&arguments, "model") + .unwrap(); + assert_eq!(mapped.pages.as_deref(), Some("1")); + assert_eq!(mapped.features, None); + assert_eq!(arguments["pages"], json!([0])); + assert_eq!(arguments["future_ocr_option"], true); + assert_eq!(arguments["extra_body"], json!({"provider_option": "value"})); + } + + #[test] + fn options_normalize_query_fields_without_consuming_extensions() { + let arguments = serde_json::from_value(json!({ + "pages":"4", "features":"languages", "extension":true + })) + .unwrap(); + let mapped = AzureDocumentIntelligenceOCRConfig + .map_ocr_params(&arguments, "model") + .unwrap(); + assert_eq!( + serde_json::to_value(mapped).unwrap(), + json!({ + "pages":"4", "features":"languages" + }) + ); + assert_eq!(arguments["extension"], true); + } + + #[test] + fn response_numbers_follow_python_validation_before_dimension_conversion() { + let response = AzureDocumentIntelligenceOCRConfig.transform_ocr_response( + "model", + br#"{"status":"succeeded","analyzeResult":{"pages":[{"pageNumber":2.0,"width":" 8.5 ","height":true}]}}"#, + OcrResponseFormat::Litellm, + ).unwrap(); + assert_eq!(response.pages[0].index, 1); + let dimensions = response.pages[0].dimensions.as_ref().unwrap(); + assert_eq!(dimensions.width, Some(816)); + assert_eq!(dimensions.height, Some(96)); + assert!(pixel_dimension(9_223_372_036_854_775_808.0, 1.0, "width").is_err()); + } + + #[rstest] + #[case(json!([0, 1, 2]), Some("1,2,3"))] + #[case(json!([2, 0, 0, 1]), Some("1,2,3"))] + #[case(json!([]), None)] + #[case(Value::Null, None)] + #[case(json!([i64::MAX - 1]), Some("9223372036854775807"))] + #[case(json!("3-9"), Some("3-9"))] + #[case(json!("1-3, 5"), Some("1-3,5"))] + #[case(json!(["1", "3-5"]), Some("1,3-5"))] + fn page_mapping_matches_python(#[case] input: Value, #[case] expected: Option<&str>) { + assert_eq!( + map(json!({"pages": input})).unwrap().pages.as_deref(), + expected + ); + } + + #[rstest] + #[case(json!("a,b"))] + #[case(json!([-1]))] + #[case(json!([true, false]))] + #[case(json!([1, "2"]))] + #[case(json!(["1", 2]))] + #[case(json!([1.0]))] + #[case(json!([i64::MAX]))] + #[case(json!([u64::MAX]))] + #[case(json!([null]))] + #[case(json!([[1]]))] + #[case(json!(5))] + fn page_mapping_rejects_invalid_shapes_and_overflow(#[case] input: Value) { + assert!(map(json!({"pages": input})).is_err()); + } + + #[rstest] + #[case(json!(["keyValuePairs"]), "keyValuePairs")] + #[case(json!(["keyValuePairs", "languages"]), "keyValuePairs,languages")] + #[case(json!("keyValuePairs"), "keyValuePairs")] + #[case(json!("keyValuePairs,languages"), "keyValuePairs,languages")] + #[case(json!("keyValuePairs, languages"), "keyValuePairs,languages")] + fn feature_mapping_matches_python(#[case] input: Value, #[case] expected: &str) { + assert_eq!( + map(json!({"features": input})).unwrap().features.as_deref(), + Some(expected) + ); + } + + #[rstest] + #[case(json!("keyValuePairs&pages=9"))] + #[case(json!("key value pairs"))] + #[case(json!(""))] + #[case(json!([1, 2]))] + #[case(json!([["keyValuePairs"]]))] + #[case(json!({"feature":"keyValuePairs"}))] + #[case(json!(5))] + fn invalid_feature_mapping_matches_python(#[case] input: Value) { + assert!(map(json!({"features": input})).is_err()); + } + + #[test] + fn empty_feature_list_is_omitted() { + assert_eq!(map(json!({"features": []})).unwrap().features, None); + } + + #[tokio::test] + async fn request_endpoint_cannot_receive_environment_key() { + let connection = OcrConnection { + api_base: Some("https://request.example".into()), + api_base_source: InputSource::Request, + ..Default::default() + }; + + let error = AzureDocumentIntelligenceOCRConfig + .validate_environment(&connection, &Default::default(), &|name| { + (name == AZURE_DI_API_KEY_ENV).then(|| "environment-key".into()) + }) + .await + .unwrap_err(); + + assert!( + error + .to_string() + .contains("request-controlled Azure endpoint") + ); + } + + #[tokio::test] + async fn request_endpoint_accepts_request_owned_key() { + let connection = OcrConnection { + api_key: Some("request-key".into()), + api_key_source: InputSource::Request, + api_base: Some("https://request.example".into()), + api_base_source: InputSource::Request, + ..Default::default() + }; + + let headers = AzureDocumentIntelligenceOCRConfig + .validate_environment(&connection, &Default::default(), &|_| None) + .await + .unwrap(); + + assert_eq!( + headers[0], + (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}; + + 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() + .into(); + + 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 request = wire_request( + "azure_ai/doc-intelligence/prebuilt-read", + "http://127.0.0.1:1", + options.clone(), + ); + let rejected = perform_ocr(request).await.is_err(); + 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 + .transport + .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.credentials.api_key = None; + request.transport.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.transport.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/mod.rs b/litellm-rust/crates/core/src/llms/azure_ai/ocr/mod.rs new file mode 100644 index 00000000000..e106f50b0a7 --- /dev/null +++ b/litellm-rust/crates/core/src/llms/azure_ai/ocr/mod.rs @@ -0,0 +1,4 @@ +pub(crate) mod cohere_parse_transformation; +pub(crate) mod common_utils; +pub(crate) mod document_intelligence; +pub(crate) mod transformation; 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 new file mode 100644 index 00000000000..dffe0aa9b05 --- /dev/null +++ b/litellm-rust/crates/core/src/llms/azure_ai/ocr/transformation.rs @@ -0,0 +1,399 @@ +use crate::call_arguments::CallArguments; +use crate::constants::AZURE_AI_OCR_PATH; +use crate::llms::base_llm::ocr::transformation::{BaseOcrConfig, OcrRequestContext}; +use crate::llms::mistral::ocr::transformation::{MistralOCRConfig, MistralOcrRequest}; +use crate::ocr::OcrClient; +use crate::ocr::document::{inline_remote_document, validate_inline_document}; +use crate::ocr::prepare::credential_env; +use crate::ocr::types::{LiteLLMOcrResponse, OcrConnection, OcrDocument, PreparedOcrRequest}; +use crate::params::OpaqueParams; +use crate::url_utils::ApiUrl; +use litellm_auth::{InputSource, Sourced}; +use litellm_auth_azure::AzureAuthInputs; +use serde_json::Value; + +const AZURE_AI_API_KEY_ENV: &str = "AZURE_AI_API_KEY"; +const AZURE_AI_API_BASE_ENV: &str = "AZURE_AI_API_BASE"; + +#[derive(Clone, Debug, Default)] +pub(crate) struct AzureAIOCRConfig; + +impl BaseOcrConfig for AzureAIOCRConfig { + type OcrParams = OpaqueParams; + type ProviderRequest = MistralOcrRequest; + 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: &PreparedOcrRequest, + _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: &PreparedOcrRequest, + _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) + } + + fn map_ocr_params( + &self, + arguments: &CallArguments, + model: &str, + ) -> Result { + MistralOCRConfig.map_ocr_params(arguments, model) + } + + async fn async_transform_ocr_request( + &self, + model: &str, + document: OcrDocument, + optional_params: &OpaqueParams, + headers: &[(String, String)], + context: OcrRequestContext<'_>, + ) -> Result { + let document = inline_remote_document( + context.client.document_fetcher(), + document, + context.connection, + ) + .await?; + self.transform_ocr_request(model, document, optional_params, headers) + } + + fn transform_ocr_response( + &self, + model: &str, + raw_response: &[u8], + request_format: crate::ocr::types::OcrResponseFormat, + ) -> Result { + MistralOCRConfig.transform_ocr_response(model, raw_response, request_format) + } + + fn validate_request_body(&self, body: &Value) -> Result<(), crate::ocr::Error> { + validate_inline_document(&crate::ocr::prepare::body_document(body)?) + } +} + +impl AzureAIOCRConfig { + /// Python `AzureAIOCRConfig.validate_environment` requires the endpoint + /// before it resolves credentials; keep that order so a missing base is + /// reported without invoking any token provider. + pub(super) fn resolve_api_base( + api_base: Option<&str>, + env_lookup: &dyn Fn(&str) -> Option, + ) -> Result { + nonblank(api_base.map(str::to_string)) + .or_else(|| nonblank(env_lookup(AZURE_AI_API_BASE_ENV))) + .ok_or(crate::ocr::Error::Auth( + litellm_auth::Error::MissingApiBase { + provider: "Azure AI", + environment_variable: AZURE_AI_API_BASE_ENV, + }, + )) + } + + fn get_complete_url( + &self, + api_base: Option<&str>, + env_lookup: &dyn Fn(&str) -> Option, + ) -> Result { + let base = Self::resolve_api_base(api_base, env_lookup)?; + let path: Vec<&str> = AZURE_AI_OCR_PATH.trim_matches('/').split('/').collect(); + ApiUrl::parse(&base) + .and_then(|url| url.complete_path(&path)) + .map(|url| url.into_string()) + .map_err(|_| crate::ocr::Error::RequestField { + path: "api_base".into(), + }) + } + + pub(super) async fn validate_environment( + &self, + connection: &OcrConnection, + config: &AzureAuthInputs, + env_lookup: &(dyn Fn(&str) -> Option + Sync), + ) -> Result, crate::ocr::Error> { + Self::resolve_api_base(connection.api_base.as_deref(), env_lookup)?; + if crate::http_utils::has_header(&connection.extra_headers, "authorization") { + if config.azure_ad_token_provider.is_some() { + super::common_utils::resolve_entra(config, env_lookup).await?; + } + super::common_utils::validate_destination(connection, connection.extra_headers_source)?; + return Ok(connection.extra_headers.clone()); + } + let key = nonblank(connection.api_key.clone()) + .map(|value| Sourced::new(value, connection.api_key_source)) + .or_else(|| { + nonblank(self.get_api_key_env_var().and_then(env_lookup)) + .map(|value| Sourced::new(value, InputSource::Environment)) + }); + if let Some(key) = key { + super::common_utils::validate_destination(connection, key.source())?; + return Ok(bearer_headers(connection, key.value())); + } + let key = super::common_utils::resolve_entra(config, env_lookup) + .await? + .ok_or(crate::ocr::Error::MissingAzureAiCredentials)?; + super::common_utils::validate_destination(connection, key.source())?; + Ok(bearer_headers(connection, key.value())) + } +} + +fn bearer_headers(connection: &OcrConnection, key: &str) -> Vec<(String, String)> { + std::iter::once(("Authorization".into(), format!("Bearer {key}"))) + .chain(connection.extra_headers.clone()) + .collect() +} + +fn nonblank(value: Option) -> Option { + value + .map(|value| value.trim().to_string()) + .filter(|value| !value.is_empty()) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn completes_azure_path_and_preserves_query() { + assert_eq!( + AzureAIOCRConfig + .get_complete_url(Some("https://example.com/?tenant=a"), &|_| None) + .unwrap(), + "https://example.com/providers/mistral/azure/ocr?tenant=a" + ); + assert_eq!( + AzureAIOCRConfig + .get_complete_url( + Some("https://example.com/providers/mistral/azure/ocr"), + &|_| None + ) + .unwrap(), + "https://example.com/providers/mistral/azure/ocr" + ); + } + + #[test] + fn missing_api_base_is_structured() { + assert!(matches!( + AzureAIOCRConfig::resolve_api_base(None, &|_| None), + Err(crate::ocr::Error::Auth( + litellm_auth::Error::MissingApiBase { + provider: "Azure AI", + environment_variable: AZURE_AI_API_BASE_ENV, + } + )) + )); + } + + #[tokio::test] + async fn supplied_authorization_precedes_keys() { + let connection = OcrConnection { + api_key: Some("request-key".into()), + api_base: Some("https://example.com".into()), + extra_headers: vec![("authorization".into(), "Bearer prepared".into())], + ..Default::default() + }; + assert_eq!( + AzureAIOCRConfig + .validate_environment(&connection, &Default::default(), &|_| { + Some("environment-key".into()) + }) + .await + .unwrap(), + connection.extra_headers + ); + } + + #[tokio::test] + async fn request_key_precedes_environment_key() { + let connection = OcrConnection { + api_key: Some("request-key".into()), + api_base: Some("https://example.com".into()), + ..Default::default() + }; + assert_eq!( + AzureAIOCRConfig + .validate_environment(&connection, &Default::default(), &|_| { + Some("environment-key".into()) + }) + .await + .unwrap()[0], + ("Authorization".into(), "Bearer request-key".into()) + ); + } + + #[tokio::test] + async fn request_endpoint_cannot_receive_environment_key() { + let connection = OcrConnection { + api_base: Some("https://request.example".into()), + api_base_source: InputSource::Request, + ..Default::default() + }; + + let error = AzureAIOCRConfig + .validate_environment(&connection, &Default::default(), &|name| { + (name == AZURE_AI_API_KEY_ENV).then(|| "environment-key".into()) + }) + .await + .unwrap_err(); + + assert!( + error + .to_string() + .contains("request-controlled Azure endpoint") + ); + } + + #[tokio::test] + async fn request_endpoint_accepts_request_owned_key() { + let connection = OcrConnection { + api_key: Some("request-key".into()), + api_key_source: InputSource::Request, + api_base: Some("https://request.example".into()), + api_base_source: InputSource::Request, + ..Default::default() + }; + + let headers = AzureAIOCRConfig + .validate_environment(&connection, &Default::default(), &|_| None) + .await + .unwrap(); + + assert_eq!( + headers[0], + ("Authorization".into(), "Bearer request-key".into()) + ); + } + + use std::sync::Arc; + + use serde_json::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.credentials.api_key = None; + request.transport.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.credentials.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/mod.rs b/litellm-rust/crates/core/src/llms/base_llm/mod.rs new file mode 100644 index 00000000000..079e0c41eae --- /dev/null +++ b/litellm-rust/crates/core/src/llms/base_llm/mod.rs @@ -0,0 +1 @@ +pub(crate) mod ocr; diff --git a/litellm-rust/crates/core/src/llms/base_llm/ocr/mod.rs b/litellm-rust/crates/core/src/llms/base_llm/ocr/mod.rs new file mode 100644 index 00000000000..080f0a1183f --- /dev/null +++ b/litellm-rust/crates/core/src/llms/base_llm/ocr/mod.rs @@ -0,0 +1 @@ +pub(crate) mod transformation; 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 new file mode 100644 index 00000000000..8af304b7d8d --- /dev/null +++ b/litellm-rust/crates/core/src/llms/base_llm/ocr/transformation.rs @@ -0,0 +1,211 @@ +use std::future::Future; +use std::sync::Arc; + +use serde::Serialize; +use serde::de::DeserializeOwned; +use serde_json::Value; + +use crate::call_arguments::CallArguments; +use crate::ocr::OcrClient; +use crate::ocr::hooks::OcrHooks; +use crate::ocr::types::{ + LiteLLMOcrResponse, OcrConnection, OcrCredentialInputs, OcrDocument, OcrResponseFormat, + PreparedOcrRequest, ResolvedOcrCredentials, +}; + +/// Output of `validate_environment`: whatever a provider resolves up front +/// (headers at minimum; Vertex also carries the project id). +pub(crate) trait OcrEnvironment: Send + Sync { + fn headers(&self) -> &[(String, String)]; +} + +impl OcrEnvironment for Vec<(String, String)> { + fn headers(&self) -> &[(String, String)] { + self + } +} + +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: Send + Sync; + type ProviderRequest: Serialize + Send; + type Environment: OcrEnvironment; + + fn get_api_key_env_var(&self) -> Option<&'static str> { + None + } + + fn resolve_connection_params(&self, inputs: OcrCredentialInputs) -> ResolvedOcrCredentials { + ResolvedOcrCredentials { + api_key: inputs + .dynamic_api_key + .filter(|value| !value.value().is_empty()) + .or(inputs.api_key), + api_base: inputs + .dynamic_api_base + .filter(|value| !value.value().is_empty()) + .or(inputs.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: &PreparedOcrRequest, + client: &OcrClient, + ) -> impl Future> + Send; + + fn get_complete_url( + &self, + request: &PreparedOcrRequest, + optional_params: &Self::OcrParams, + environment: &Self::Environment, + ) -> Result; + + fn get_supported_ocr_params(&self, _model: &str) -> &'static [&'static str] { + &[] + } + + fn map_ocr_params( + &self, + arguments: &CallArguments, + model: &str, + ) -> Result; + + 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 { + async move { self.transform_ocr_request(model, document, optional_params, headers) } + } + + fn transform_ocr_response( + &self, + model: &str, + raw_response: &[u8], + request_format: OcrResponseFormat, + ) -> Result; + + fn async_transform_ocr_response( + &self, + model: &str, + raw_response: reqwest::Response, + context: OcrResponseContext<'_>, + ) -> impl Future> + Send { + async move { + let bytes = crate::ocr::client::read_response_bytes( + raw_response, + context.connection.max_response_bytes, + ) + .await?; + crate::ocr::handler::post_call(context.hooks, &bytes).await?; + 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, + } + } + + /// Provider-specific check applied to the composed body, both before and + /// after guardrail hooks. Defaults to accepting any body. + fn validate_request_body(&self, _body: &Value) -> Result<(), crate::ocr::Error> { + Ok(()) + } + + /// Rust counterpart of `BaseLLMHTTPHandler._async_prepare_ocr_request`: + /// map params, validate environment, build URL, transform, compose body. + fn prepare_request( + &self, + request: &PreparedOcrRequest, + client: &OcrClient, + ) -> impl Future> + Send { + async move { + let params = self.map_ocr_params(&request.optional_params, &request.model)?; + let environment = self.validate_environment(request, client).await?; + let url = self.get_complete_url(request, ¶ms, &environment)?; + let headers = environment.headers(); + let body = self + .async_transform_ocr_request( + &request.model, + request.document.clone(), + ¶ms, + headers, + OcrRequestContext { + client, + connection: &request.connection, + }, + ) + .await?; + crate::ocr::prepare::transform_request_body( + client, + request, + &url, + headers, + body, + |body| self.validate_request_body(body), + ) + .await + } + } +} + +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::json::decode_response( + raw_response, + request_format == OcrResponseFormat::Native, + )?; + Ok(LiteLLMOcrResponse { + provider_native_response: decoded.native, + ..normalize(model, decoded.data)? + }) +} + +#[derive(Clone, Copy)] +pub(crate) struct OcrRequestContext<'a> { + pub client: &'a OcrClient, + pub connection: &'a OcrConnection, +} + +#[derive(Clone, Copy)] +pub(crate) struct OcrResponseContext<'a> { + pub client: &'a OcrClient, + pub connection: &'a OcrConnection, + pub hooks: &'a Arc, + pub request_format: OcrResponseFormat, + pub url: &'a str, + pub headers: &'a [(String, String)], +} diff --git a/litellm-rust/crates/core/src/llms/cohere/mod.rs b/litellm-rust/crates/core/src/llms/cohere/mod.rs new file mode 100644 index 00000000000..079e0c41eae --- /dev/null +++ b/litellm-rust/crates/core/src/llms/cohere/mod.rs @@ -0,0 +1 @@ +pub(crate) mod ocr; diff --git a/litellm-rust/crates/core/src/llms/cohere/ocr/mod.rs b/litellm-rust/crates/core/src/llms/cohere/ocr/mod.rs new file mode 100644 index 00000000000..9cbe4df56e5 --- /dev/null +++ b/litellm-rust/crates/core/src/llms/cohere/ocr/mod.rs @@ -0,0 +1,3 @@ +pub(crate) mod transformation; + +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 new file mode 100644 index 00000000000..fc11f62833c --- /dev/null +++ b/litellm-rust/crates/core/src/llms/cohere/ocr/transformation.rs @@ -0,0 +1,740 @@ +use serde::{Deserialize, Serialize}; +use serde_json::{Map, Value}; +use serde_with::serde_as; + +use crate::call_arguments::{CallArguments, parse_options}; +use crate::constants::{COHERE_API_KEY_ENV, COHERE_PARSE_API_BASE}; +use crate::llms::base_llm::ocr::transformation::{BaseOcrConfig, OcrRequestContext}; +use crate::ocr::OcrClient; +use crate::ocr::document::InlineDocument; +use crate::ocr::prepare::credential_env; +use crate::ocr::types::{ + LiteLLMOcrResponse, OcrConnection, OcrDocument, OcrPage, OcrPageImage, OcrUsageInfo, + PreparedOcrRequest, +}; +use crate::serde_compat::LaxI64; +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 { + #[default] + Markdown, + Blocks, +} + +#[derive(Default, Deserialize, Serialize)] +pub(crate) struct CohereOptions { + #[serde(skip_serializing_if = "Option::is_none")] + pub output_format: Option, +} + +#[derive(Deserialize, Serialize)] +pub(crate) struct CohereRequest { + pub model: String, + pub document: CohereParseDocument, + pub output_format: String, +} + +#[derive(Deserialize, Serialize)] +#[serde(tag = "type")] +pub(crate) enum CohereParseDocument { + #[serde(rename = "image_url")] + ImageUrl { image_url: String }, +} + +#[derive(Deserialize)] +pub(crate) struct CohereResponse { + #[serde(default)] + pages: Vec, + meta: Option, +} + +#[serde_as] +#[derive(Deserialize)] +struct CoherePage { + #[serde_as(deserialize_as = "Option")] + index: Option, + markdown: Option, + blocks: Option>>, +} + +#[derive(Deserialize, Serialize)] +struct CohereMarkdown { + #[serde(default)] + content: String, + images: Option>>, +} + +#[derive(Deserialize)] +struct CohereMeta { + billed_units: Option, +} + +#[serde_as] +#[derive(Deserialize)] +struct CohereBilledUnits { + #[serde_as(deserialize_as = "Option")] + pages: Option, +} + +#[derive(Default)] +pub(crate) struct CohereParseConfig; + +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: &PreparedOcrRequest, + _client: &OcrClient, + ) -> Result { + self.validate_environment(&request.connection, &credential_env) + } + + fn get_complete_url( + &self, + request: &PreparedOcrRequest, + _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, + optional_params: &CohereOptions, + _headers: &[(String, String)], + ) -> Result { + let image_url = image_url(document)?; + Ok(build_request(model, image_url, optional_params)) + } + + fn get_supported_ocr_params(&self, _model: &str) -> &'static [&'static str] { + &["output_format", "req_format"] + } + + fn map_ocr_params( + &self, + arguments: &CallArguments, + _model: &str, + ) -> Result { + Ok(parse_options(arguments)?) + } + + async fn async_transform_ocr_request( + &self, + model: &str, + document: OcrDocument, + optional_params: &CohereOptions, + headers: &[(String, String)], + _context: OcrRequestContext<'_>, + ) -> Result { + self.transform_ocr_request(model, document, optional_params, headers) + } + + fn transform_ocr_response( + &self, + model: &str, + raw_response: &[u8], + request_format: crate::ocr::types::OcrResponseFormat, + ) -> Result { + crate::llms::base_llm::ocr::transformation::decode_and_normalize_response( + model, + raw_response, + request_format, + normalize_response, + ) + } + + fn validate_request_body(&self, body: &Value) -> Result<(), crate::ocr::Error> { + validate_document(&crate::ocr::prepare::body_document(body)?) + } +} + +pub(crate) fn validate_document(document: &OcrDocument) -> Result<(), crate::ocr::Error> { + let OcrDocument::ImageUrl { image_url, .. } = document else { + return Err(crate::ocr::Error::CohereImageOnly); + }; + if image_url.is_empty() { + return Err(crate::ocr::Error::CohereImageOnly); + } + if let Some(inline) = InlineDocument::parse(image_url)? { + if !inline.mime_type().type_.eq_ignore_ascii_case("image") { + return Err(crate::ocr::Error::CohereImageOnly); + } + inline.decode(crate::constants::OCR_INLINE_MAX_BYTES)?; + } + Ok(()) +} + +pub(crate) fn normalize_response( + model: &str, + response: CohereResponse, +) -> Result { + let pages_processed = billed_pages(&response).map(Ok).unwrap_or_else(|| { + i64::try_from(response.pages.len()).map_err(|_| crate::ocr::Error::NumericRange("pages")) + })?; + let pages = response + .pages + .into_iter() + .enumerate() + .map(|(position, page)| normalize_page(page, position)) + .collect::, crate::ocr::Error>>()?; + Ok(LiteLLMOcrResponse { + usage_info: Some(OcrUsageInfo { + pages_processed: Some(pages_processed), + ..Default::default() + }), + ..LiteLLMOcrResponse::new(model, pages) + }) +} + +fn image_url(document: OcrDocument) -> Result { + validate_document(&document)?; + let OcrDocument::ImageUrl { image_url, .. } = document else { + return Err(crate::ocr::Error::CohereImageOnly); + }; + Ok(image_url) +} + +fn build_request(model: &str, image_url: String, params: &CohereOptions) -> CohereRequest { + CohereRequest { + model: model.into(), + document: CohereParseDocument::ImageUrl { image_url }, + output_format: match params.output_format.unwrap_or_default() { + OutputFormat::Markdown => "markdown", + OutputFormat::Blocks => "blocks", + } + .into(), + } +} + +fn page_image( + mut image: Map, + path: &str, +) -> Result { + if let Some(Value::Object(bbox)) = image.get("bounding_box") { + image.insert("bbox".into(), Value::Object(bbox.clone())); + } + crate::ocr::json::decode_response_value(Value::Object(image), path) +} + +fn normalize_page(page: CoherePage, position: usize) -> Result { + let index = page.index.map(Ok).unwrap_or_else(|| { + i64::try_from(position).map_err(|_| crate::ocr::Error::NumericRange("page index")) + })?; + let (markdown, images) = match page.markdown { + Some(markdown) => { + let images = markdown + .images + .filter(|images| !images.is_empty()) + .map(|images| { + images + .into_iter() + .enumerate() + .map(|(image_index, image)| { + page_image( + image, + &format!("pages[{position}].markdown.images[{image_index}]"), + ) + }) + .collect::, _>>() + }) + .transpose()?; + (markdown.content, images) + } + None => (String::new(), None), + }; + let extra_fields = page + .blocks + .map(|blocks| { + ( + "blocks".into(), + Value::Array(blocks.into_iter().map(Value::Object).collect()), + ) + }) + .into_iter() + .collect(); + Ok(OcrPage { + index, + markdown, + images, + extra_fields, + ..Default::default() + }) +} + +fn billed_pages(response: &CohereResponse) -> Option { + response.meta.as_ref()?.billed_units.as_ref()?.pages +} + +impl CohereParseConfig { + fn get_complete_url(&self, base: &str) -> Result { + let parsed = reqwest::Url::parse(base).map_err(|_| invalid_api_base())?; + if !matches!(parsed.scheme(), "http" | "https") { + return Err(invalid_api_base()); + } + ApiUrl::parse(base) + .and_then(|url| url.complete_path(&["v2", "parse"])) + .map(|url| url.into_string()) + .map_err(|_| invalid_api_base()) + } + + fn validate_environment( + &self, + connection: &OcrConnection, + env_lookup: &(dyn Fn(&str) -> Option + Sync), + ) -> Result, crate::ocr::Error> { + if crate::http_utils::has_header(&connection.extra_headers, "authorization") { + return Ok(connection.extra_headers.clone()); + } + let key = connection + .api_key + .as_deref() + .map(str::trim) + .filter(|key| !key.is_empty()) + .map(str::to_string) + .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(), + )) + })?; + Ok( + std::iter::once(("Authorization".into(), format!("Bearer {key}"))) + .chain(connection.extra_headers.clone()) + .collect(), + ) + } +} + +fn invalid_api_base() -> crate::ocr::Error { + crate::ocr::Error::RequestField { + path: "api_base".into(), + } +} + +#[cfg(test)] +mod tests { + use super::*; + use serde_json::json; + + #[tokio::test] + async fn composed_body_preserves_native_document_fields_and_untyped_overrides() { + let request = crate::ocr::test_support::wire_request( + "cohere/parse", + "https://example.com", + json!({ + "output_format":"markdown", "metadata":{"host":true}, + "extra_body":{ + "output_format": {"future":true}, + "document":{"type":"image_url","image_url":"https://example.com/a.png", + "provider_options":{"nested":[false,0,null]}} + } + }), + ); + let request = request.with_document( + serde_json::from_value(json!({ + "type":"image_url","image_url":"https://example.com/original.png" + })) + .unwrap(), + ); + let request = crate::ocr::prepare::prepare_request(request); + let http = CohereParseConfig + .prepare_request(&request, &crate::ocr::test_support::ocr_client()) + .await + .unwrap(); + let body: Value = serde_json::from_slice(http.body().unwrap().as_bytes().unwrap()).unwrap(); + assert_eq!( + body, + json!({ + "model":"parse", "output_format":{"future":true}, + "document":{"type":"image_url","image_url":"https://example.com/a.png", + "provider_options":{"nested":[false,0,null]}} + }) + ); + } + + #[test] + fn options_read_known_fields_without_changing_arguments() { + let arguments = serde_json::from_value(json!({ + "output_format":"blocks", "req_format":"native", "extension":false + })) + .unwrap(); + for config in [false, true] { + let mapped = if config { + crate::llms::azure_ai::ocr::cohere_parse_transformation::AzureAICohereParseConfig + .map_ocr_params(&arguments, "parse") + } else { + CohereParseConfig.map_ocr_params(&arguments, "parse") + } + .unwrap(); + assert_eq!( + serde_json::to_value(mapped).unwrap(), + json!({"output_format":"blocks"}) + ); + } + assert_eq!(arguments["req_format"], "native"); + assert_eq!(arguments["extension"], false); + let invalid = serde_json::from_value(json!({"output_format":"html"})).unwrap(); + assert!(matches!( + CohereParseConfig.map_ocr_params(&invalid, "parse"), + Err(crate::ocr::Error::RequestField { path }) + if path == "optional_params.output_format" + )); + } + + #[test] + fn billed_pages_accept_integral_doubles_and_reject_fractional_counts() { + let response = serde_json::from_str::( + r#"{"pages":[],"meta":{"billed_units":{"pages":1.0}}}"#, + ) + .unwrap(); + let normalized = normalize_response("parse", response).unwrap(); + assert_eq!(normalized.usage_info.unwrap().pages_processed, Some(1)); + assert!( + serde_json::from_str::( + r#"{"pages":[],"meta":{"billed_units":{"pages":1.5}}}"#, + ) + .is_err() + ); + } + + #[test] + fn response_preserves_python_mapping_shapes_and_extensions() { + let blocks = json!([ + {"type":"text", "text":"Total Due: $4.00"}, + {"type":"future", "payload":{"nested":[null,false,0]}} + ]); + let response = serde_json::from_value(json!({ + "pages":[{ + "index":"2", + "markdown":{"content":"receipt", "images":[ + {"bounding_box":{"x":1}, "bbox":"replaced", "category":"future", "extension":null}, + {"image_base64":"encoded"} + ]}, + "blocks":blocks + }], + "meta":{"billed_units":{"pages":0}} + })).unwrap(); + let response = normalize_response("parse", response).unwrap(); + assert_eq!(response.pages[0].index, 2); + assert_eq!(response.usage_info.unwrap().pages_processed, Some(0)); + assert_eq!(response.pages[0].extra_fields["blocks"], blocks); + let images = response.pages[0].images.as_ref().unwrap(); + assert_eq!(images[0].bbox.as_ref().unwrap()["x"], 1); + assert_eq!(images[0].extra_fields["category"], "future"); + assert_eq!(images[0].extra_fields.get("extension"), Some(&Value::Null)); + assert_eq!(images[1].image_base64.as_deref(), Some("encoded")); + assert!(images[1].bbox.is_none()); + } + + #[test] + fn malformed_normalized_image_fields_report_the_original_path() { + let response = serde_json::from_value(json!({ + "pages":[{"markdown":{"images":[{"image_base64":42}]}}] + })) + .unwrap(); + assert!(matches!( + normalize_response("parse", response).unwrap_err(), + crate::ocr::Error::ResponseField { path } + if path == "pages[0].markdown.images[0].image_base64" + )); + } + + #[test] + fn provider_options_exclude_response_controls_and_extensions() { + let arguments = serde_json::from_value( + json!({"output_format":"blocks","req_format":"native","unknown":true}), + ) + .unwrap(); + let params = CohereParseConfig + .map_ocr_params(&arguments, "parse") + .unwrap(); + assert_eq!( + serde_json::to_value(¶ms).unwrap(), + json!({"output_format":"blocks"}) + ); + let document = serde_json::from_value( + json!({"type":"image_url","image_url":"https://example.com/a.png","ignored":"field"}), + ) + .unwrap(); + let body = CohereParseConfig + .transform_ocr_request("parse", document, ¶ms, &[]) + .unwrap(); + assert_eq!( + serde_json::to_value(body).unwrap(), + json!({ + "model":"parse", "document":{"type":"image_url","image_url":"https://example.com/a.png"}, "output_format":"blocks" + }) + ); + } + + #[tokio::test] + async fn explicit_null_options_use_defaults_before_http() { + let request = crate::ocr::test_support::wire_request( + "cohere/parse", + "https://example.com", + json!({"output_format":null,"req_format":null}), + ); + let request = request.with_document( + serde_json::from_value( + json!({"type":"image_url","image_url":"https://example.com/a.png"}), + ) + .unwrap(), + ); + assert_eq!( + request.response_format().unwrap(), + crate::ocr::types::OcrResponseFormat::Litellm + ); + let request = crate::ocr::prepare::prepare_request(request); + let http = CohereParseConfig + .prepare_request(&request, &crate::ocr::test_support::ocr_client()) + .await + .unwrap(); + let body: Value = serde_json::from_slice(http.body().unwrap().as_bytes().unwrap()).unwrap(); + assert_eq!(body["output_format"], "markdown"); + assert!(body.get("req_format").is_none()); + } + + #[test] + fn response_normalizes_markdown_images_blocks_and_billed_pages() { + let response = serde_json::from_value(json!({ + "pages": [ + { + "type":"markdown", + "index":4, + "markdown":{ + "content":"receipt", + "images":[{ + "id":"image", + "bounding_box":{ + "top_left_x":1, + "top_left_y":2, + "bottom_right_x":48, + "bottom_right_y":49 + }, + "bounding_box_normalized":{ + "top_left_x":0.04, + "top_left_y":0.05, + "bottom_right_x":0.15, + "bottom_right_y":0.16 + }, + "description":"scan", + "category":"logo", + "provider_extension":"preserved" + }] + } + }, + {"type":"blocks","blocks":[{"type":"text","text":{"content":"total"}}]} + ], + "meta":{"api_version":{"version":"2"},"billed_units":{"pages":3}} + })) + .unwrap(); + let normalized = normalize_response("parse-v5.0", response).unwrap(); + assert_eq!(normalized.pages[0].index, 4); + assert_eq!(normalized.pages[0].markdown, "receipt"); + let image = &normalized.pages[0].images.as_ref().unwrap()[0]; + assert_eq!(image.bbox.as_ref().unwrap()["top_left_x"], 1); + assert_eq!( + image.extra_fields["bounding_box_normalized"]["bottom_right_x"], + 0.15 + ); + assert_eq!(image.extra_fields["description"], "scan"); + assert_eq!(image.extra_fields["category"], "logo"); + assert_eq!(image.extra_fields["provider_extension"], "preserved"); + assert_eq!(normalized.pages[1].index, 1); + assert_eq!(normalized.pages[1].markdown, ""); + assert_eq!( + normalized.pages[1].extra_fields["blocks"][0]["text"]["content"], + "total" + ); + assert_eq!(normalized.usage_info.unwrap().pages_processed, Some(3)); + } + + #[test] + fn response_defaults_and_invalid_fields() { + for value in [ + json!({}), + json!({"meta":null}), + json!({"pages":[],"meta":{"billed_units":null}}), + ] { + let normalized = + normalize_response("parse", serde_json::from_value(value).unwrap()).unwrap(); + assert!(normalized.pages.is_empty()); + assert_eq!(normalized.usage_info.unwrap().pages_processed, Some(0)); + } + for value in [ + json!({"pages":null}), + json!({"pages":[{"markdown":"text"}]}), + json!({"pages":[{"index":"bad"}]}), + ] { + assert!(serde_json::from_value::(value).is_err()); + } + let normalized = normalize_response( + "parse", + serde_json::from_value(json!({"pages":[{"markdown":null}]})).unwrap(), + ) + .unwrap(); + assert_eq!(normalized.usage_info.unwrap().pages_processed, Some(1)); + assert!(normalized.pages[0].images.is_none()); + } + + #[test] + fn response_types_documented_block_variants() { + let response = serde_json::from_value(json!({ + "pages": [{ + "type": "blocks", + "index": 0, + "blocks": [ + {"type": "text", "text": {"content": "hello"}}, + { + "type": "image", + "image": { + "id": "img-0", + "description": "logo", + "category": "logo", + "bounding_box": { + "top_left_x": 1, + "top_left_y": 2, + "bottom_right_x": 3, + "bottom_right_y": 4 + }, + "bounding_box_normalized": { + "top_left_x": 0.1, + "top_left_y": 0.2, + "bottom_right_x": 0.3, + "bottom_right_y": 0.4 + } + } + }, + { + "type": "table", + "table": { + "type": "html", + "html": "
", + "bounding_box": { + "top_left_x": 5, + "top_left_y": 6, + "bottom_right_x": 7, + "bottom_right_y": 8 + }, + "bounding_box_normalized": { + "top_left_x": 0.5, + "top_left_y": 0.6, + "bottom_right_x": 0.7, + "bottom_right_y": 0.8 + }, + "title": "Totals" + } + } + ] + }] + })) + .unwrap(); + let normalized = normalize_response("parse-v5.0", response).unwrap(); + let blocks = normalized.pages[0].extra_fields["blocks"] + .as_array() + .unwrap(); + assert_eq!(blocks[0]["text"]["content"], "hello"); + assert_eq!(blocks[1]["image"]["category"], "logo"); + assert_eq!(blocks[2]["table"]["type"], "html"); + assert_eq!(blocks[2]["table"]["title"], "Totals"); + } + + #[test] + fn request_requires_image_and_supported_output_format() { + for value in [ + json!({"type":"document_url","document_url":"https://example.com/a.pdf"}), + json!({"type":"image_url","image_url":""}), + json!({"type":"image_url","image_url":"data:application/pdf;base64,YQ=="}), + ] { + assert!(matches!( + validate_document(&serde_json::from_value(value).unwrap()), + Err(crate::ocr::Error::CohereImageOnly) + )); + } + assert!(serde_json::from_value::(json!({"output_format":"html"})).is_err()); + for format in ["markdown", "blocks"] { + assert!( + serde_json::from_value::(json!({"output_format":format})).is_ok() + ); + } + let request = CohereParseConfig + .transform_ocr_request( + "parse-v5.0", + serde_json::from_value(json!({ + "type":"image_url", + "image_url":"https://example.com/image.png" + })) + .unwrap(), + &serde_json::from_value(json!({})).unwrap(), + &[], + ) + .unwrap(); + assert_eq!( + serde_json::to_value(request).unwrap()["output_format"], + "markdown" + ); + } + + #[test] + fn completes_provider_urls_without_duplicate_paths_and_preserves_queries() { + for suffix in ["", "/v2", "/v2/parse"] { + assert_eq!( + CohereParseConfig + .get_complete_url(&format!("https://example.com{suffix}?tenant=a")) + .unwrap(), + "https://example.com/v2/parse?tenant=a" + ); + } + } + + #[test] + fn rejects_invalid_urls_and_blank_keys() { + assert!(CohereParseConfig.get_complete_url("relative/path").is_err()); + assert!( + CohereParseConfig + .get_complete_url("ftp://example.com") + .is_err() + ); + assert!(matches!( + CohereParseConfig.validate_environment( + &OcrConnection { + api_key: Some(" ".into()), + ..Default::default() + }, + &|_| None, + ), + Err(crate::ocr::Error::Auth(_)) + )); + } +} diff --git a/litellm-rust/crates/core/src/llms/mistral/mod.rs b/litellm-rust/crates/core/src/llms/mistral/mod.rs new file mode 100644 index 00000000000..079e0c41eae --- /dev/null +++ b/litellm-rust/crates/core/src/llms/mistral/mod.rs @@ -0,0 +1 @@ +pub(crate) mod ocr; diff --git a/litellm-rust/crates/core/src/llms/mistral/ocr/mod.rs b/litellm-rust/crates/core/src/llms/mistral/ocr/mod.rs new file mode 100644 index 00000000000..080f0a1183f --- /dev/null +++ b/litellm-rust/crates/core/src/llms/mistral/ocr/mod.rs @@ -0,0 +1 @@ +pub(crate) mod transformation; diff --git a/litellm-rust/crates/core/src/llms/mistral/ocr/transformation.rs b/litellm-rust/crates/core/src/llms/mistral/ocr/transformation.rs new file mode 100644 index 00000000000..d90bfeff2a7 --- /dev/null +++ b/litellm-rust/crates/core/src/llms/mistral/ocr/transformation.rs @@ -0,0 +1,626 @@ +use serde::{Deserialize, Serialize}; +use serde_json::Value; + +use crate::call_arguments::CallArguments; +use crate::constants::MISTRAL_OCR_API_BASE; +use crate::llms::base_llm::ocr::transformation::{BaseOcrConfig, OcrRequestContext}; +use crate::ocr::OcrClient; +use crate::ocr::prepare::credential_env; +use crate::ocr::types::{ + LiteLLMOcrResponse, OcrConnection, OcrDocument, OcrPage, OcrUsageInfo, PreparedOcrRequest, +}; +use crate::params::OpaqueParams; +use crate::url_utils::ApiUrl; + +const MISTRAL_API_KEY_ENV: &str = "MISTRAL_API_KEY"; + +#[derive(Clone, Debug, Serialize, Deserialize)] +pub(crate) struct MistralOcrRequest { + pub model: String, + pub document: OcrDocument, + #[serde(flatten)] + pub params: OpaqueParams, +} + +#[derive(Clone, Debug, Default, Deserialize)] +pub(crate) struct MistralOcrResponse { + #[serde(flatten)] + pub extra_fields: serde_json::Map, + #[serde(default)] + pub pages: Vec, + #[serde( + default, + deserialize_with = "serde_with::rust::double_option::deserialize" + )] + pub model: Option>, + pub document_annotation: Option, + pub usage_info: Option, +} + +#[derive(Clone, Debug, Default)] +pub(crate) struct 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: &PreparedOcrRequest, + _client: &OcrClient, + ) -> Result { + self.validate_environment(&request.connection, &credential_env) + } + + fn get_complete_url( + &self, + request: &PreparedOcrRequest, + _params: &Self::OcrParams, + _environment: &Self::Environment, + ) -> Result { + self.get_complete_url(request.connection.api_base.as_deref()) + } + + fn transform_ocr_request( + &self, + model: &str, + document: OcrDocument, + optional_params: &OpaqueParams, + _headers: &[(String, String)], + ) -> Result { + Ok(MistralOcrRequest { + model: model.to_string(), + document, + params: optional_params.clone(), + }) + } + + fn get_supported_ocr_params(&self, _model: &str) -> &'static [&'static str] { + &[ + "pages", + "include_image_base64", + "image_limit", + "image_min_size", + "bbox_annotation_format", + "document_annotation_format", + "document_annotation_prompt", + "extract_header", + "extract_footer", + "table_format", + "confidence_scores_granularity", + "include_blocks", + "id", + ] + } + + fn map_ocr_params( + &self, + arguments: &CallArguments, + model: &str, + ) -> Result { + Ok(arguments + .select(self.get_supported_ocr_params(model)) + .into()) + } + + async fn async_transform_ocr_request( + &self, + model: &str, + document: OcrDocument, + optional_params: &OpaqueParams, + headers: &[(String, String)], + _context: OcrRequestContext<'_>, + ) -> Result { + self.transform_ocr_request(model, document, optional_params, headers) + } + + fn transform_ocr_response( + &self, + model: &str, + raw_response: &[u8], + request_format: crate::ocr::types::OcrResponseFormat, + ) -> Result { + crate::llms::base_llm::ocr::transformation::decode_and_normalize_response( + model, + raw_response, + request_format, + normalize_response, + ) + } +} + +pub(crate) fn normalize_response( + model: &str, + response: MistralOcrResponse, +) -> Result { + let model = match response.model { + Some(Some(model)) => model, + Some(None) => { + return Err(crate::ocr::Error::ResponseField { + path: "model".into(), + }); + } + None => model.to_string(), + }; + Ok(LiteLLMOcrResponse { + extra_fields: response.extra_fields, + document_annotation: response.document_annotation, + usage_info: response.usage_info, + ..LiteLLMOcrResponse::new(model, response.pages) + }) +} + +impl MistralOCRConfig { + fn get_complete_url(&self, api_base: Option<&str>) -> Result { + let base = api_base + .map(str::trim) + .filter(|base| !base.is_empty()) + .unwrap_or(MISTRAL_OCR_API_BASE); + ApiUrl::parse(base) + .and_then(|url| url.complete_path(&["v1", "ocr"])) + .map(|url| url.into_string()) + .map_err(|_| crate::ocr::Error::RequestField { + path: "api_base".into(), + }) + } + + fn validate_environment( + &self, + connection: &OcrConnection, + env_lookup: &(dyn Fn(&str) -> Option + Sync), + ) -> Result, crate::ocr::Error> { + if crate::http_utils::has_header(&connection.extra_headers, "authorization") { + return Ok(connection.extra_headers.clone()); + } + let api_key = connection + .api_key + .as_deref() + .map(str::trim) + .filter(|key| !key.is_empty()) + .map(str::to_string) + .or_else(|| { + 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, + })?; + Ok( + std::iter::once(("Authorization".into(), format!("Bearer {api_key}"))) + .chain(connection.extra_headers.clone()) + .collect(), + ) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use rstest::rstest; + use serde_json::{Value, json}; + + #[test] + fn explicit_null_model_does_not_use_the_missing_model_default() { + let response = serde_json::from_value(json!({"model":null})).unwrap(); + assert!(matches!( + normalize_response("fallback", response).unwrap_err(), + crate::ocr::Error::ResponseField { path } if path == "model" + )); + } + + #[test] + fn response_validates_normalized_shapes_at_the_provider_boundary() { + for (payload, path) in [ + (json!({"pages":[42]}), "pages[0]"), + (json!({"pages":[{"index":0}]}), "pages[0]"), + ( + json!({"pages":[{"index":0,"markdown":42}]}), + "pages[0].markdown", + ), + ( + json!({"pages":[{"index":0,"markdown":"","images":[42]}]}), + "pages[0].images[0]", + ), + ( + json!({"pages":[{"index":0,"markdown":"","dimensions":{"width":1.5}}]}), + "pages[0].dimensions.width", + ), + ( + json!({"usage_info":{"pages_processed":"bad"}}), + "usage_info.pages_processed", + ), + ] { + let error = crate::ocr::json::decode_response::( + &serde_json::to_vec(&payload).unwrap(), + false, + ) + .unwrap_err(); + assert!(matches!( + error, + crate::ocr::Error::ResponseField { path: actual } if actual == path + )); + } + } + + #[test] + fn response_normalizes_python_numeric_inputs_and_shared_defaults() { + let response = serde_json::from_value(json!({ + "pages":[{"index":"2","markdown":"text","dimensions":{"width":1.0},"extension":false}], + "usage_info":{"pages_processed":true,"credits":"1.5","custom":0}, + "extra":"ignored" + })) + .unwrap(); + let response = normalize_response("model", response).unwrap(); + assert_eq!(response.pages[0].index, 2); + assert_eq!( + response.pages[0].dimensions.as_ref().unwrap().width, + Some(1) + ); + assert_eq!( + response.usage_info.as_ref().unwrap().pages_processed, + Some(1) + ); + assert_eq!(response.usage_info.as_ref().unwrap().credits, Some(1.5)); + let serialized = response.into_json(); + assert_eq!(serialized["pages"][0]["extension"], false); + assert!(serialized["pages"][0]["images"].is_null()); + assert!(serialized["usage_info"]["doc_size_bytes"].is_null()); + assert_eq!(serialized["usage_info"]["custom"], 0); + assert!(serialized["content"].is_null()); + assert_eq!(serialized["extra"], "ignored"); + } + + #[test] + fn map_ocr_params_selects_known_fields_without_changing_arguments() { + let input = + serde_json::from_value(json!({"pages":null,"extract_header":false,"unknown":true})) + .unwrap(); + let params = MistralOCRConfig.map_ocr_params(&input, "model").unwrap(); + assert_eq!( + serde_json::to_value(params).unwrap(), + json!({"pages":null,"extract_header":false}) + ); + assert_eq!(input["unknown"], true); + assert_eq!(input.get("pages"), Some(&Value::Null)); + } + + #[test] + fn request_transform_uses_already_mapped_params_without_filtering_again() { + let params = serde_json::from_value(json!({"extension":{"nested":null}})).unwrap(); + let body = MistralOCRConfig + .transform_ocr_request("model", document(), ¶ms, &[]) + .unwrap(); + assert_eq!( + serde_json::to_value(body).unwrap()["extension"], + json!({"nested":null}) + ); + } + + #[test] + 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 + .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(); + assert_eq!(native["pages"][0]["index"], "2"); + assert_eq!(native["provider_extension"], false); + assert_eq!(response.extra_fields["provider_extension"], false); + assert!( + MistralOCRConfig + .transform_ocr_response( + "model", + br#"{"pages":[{"index":0}]}"#, + crate::ocr::types::OcrResponseFormat::Litellm + ) + .is_err() + ); + } + + fn mapped_params(value: Value) -> Value { + let params = serde_json::from_value(value).unwrap(); + serde_json::to_value(MistralOCRConfig.map_ocr_params(¶ms, "model").unwrap()).unwrap() + } + + fn document() -> OcrDocument { + serde_json::from_value( + json!({"type":"document_url","document_url":"https://example.com/a.pdf"}), + ) + .unwrap() + } + + #[rstest] + fn extract_header_is_a_supported_ocr_param() { + assert_eq!( + mapped_params(json!({"extract_header":true}))["extract_header"], + true + ); + } + + #[rstest] + fn extract_footer_is_a_supported_ocr_param() { + assert_eq!( + mapped_params(json!({"extract_footer":false}))["extract_footer"], + false + ); + } + + #[rstest] + fn existing_ocr_params_remain_supported() { + let mapped = mapped_params(json!({ + "pages":[0,2], + "include_image_base64":true, + "image_limit":2, + "image_min_size":100, + "bbox_annotation_format":{"type":"json_schema"}, + "document_annotation_format":{"type":"json_schema"} + })); + assert_eq!(mapped["pages"], json!([0, 2])); + assert_eq!(mapped["include_image_base64"], true); + assert_eq!(mapped["image_limit"], 2); + assert_eq!(mapped["image_min_size"], 100); + assert_eq!(mapped["bbox_annotation_format"]["type"], "json_schema"); + assert_eq!(mapped["document_annotation_format"]["type"], "json_schema"); + } + + #[rstest] + fn map_ocr_params_forwards_extract_header() { + assert_eq!( + mapped_params(json!({"extract_header":true}))["extract_header"], + true + ); + } + + #[rstest] + fn map_ocr_params_forwards_extract_footer() { + assert_eq!( + mapped_params(json!({"extract_footer":true}))["extract_footer"], + true + ); + } + + #[rstest] + fn map_ocr_params_forwards_extract_header_and_footer() { + let mapped = mapped_params(json!({"extract_header":true,"extract_footer":false})); + assert_eq!(mapped["extract_header"], true); + assert_eq!(mapped["extract_footer"], false); + } + + #[rstest] + fn map_ocr_params_excludes_extensions_from_the_provider_options() { + let mapped = mapped_params(json!({"extract_header":true,"unsupported_param":"value"})); + assert_eq!(mapped["extract_header"], true); + assert!(mapped.get("unsupported_param").is_none()); + } + + #[rstest] + fn map_ocr_params_preserves_unvalidated_values_and_explicit_null() { + let mapped = mapped_params(json!({ + "pages":{"future":"shape"}, + "include_image_base64":null + })); + assert_eq!(mapped["pages"], json!({"future":"shape"})); + assert!(mapped.get("include_image_base64").unwrap().is_null()); + } + + #[rstest] + #[case("table_format", json!("html"))] + #[case("confidence_scores_granularity", json!("word"))] + #[case("confidence_scores_granularity", json!("block"))] + #[case("document_annotation_prompt", json!("extract"))] + #[case("include_blocks", json!(true))] + #[case("id", json!("req-123"))] + fn new_ocr_params_are_supported(#[case] name: &str, #[case] value: Value) { + assert_eq!(mapped_params(json!({name:value.clone()}))[name], value); + } + + #[rstest] + #[case("table_format", json!("html"))] + #[case("confidence_scores_granularity", json!("word"))] + #[case("document_annotation_prompt", json!("extract"))] + #[case("include_blocks", json!(true))] + #[case("id", json!("req-123"))] + fn map_ocr_params_forwards_new_ocr_params(#[case] name: &str, #[case] value: Value) { + assert_eq!(mapped_params(json!({name:value.clone()}))[name], value); + } + + #[rstest] + #[case("pages", json!([0, 2]))] + #[case("pages", json!("0,2-4"))] + #[case("include_image_base64", json!(true))] + #[case("image_limit", json!(2))] + #[case("image_min_size", json!(100))] + #[case("bbox_annotation_format", json!({"type":"json_schema"}))] + #[case("document_annotation_format", json!({"type":"json_schema"}))] + #[case("document_annotation_prompt", json!("extract"))] + #[case("extract_header", json!(true))] + #[case("extract_footer", json!(false))] + #[case("table_format", json!("html"))] + #[case("confidence_scores_granularity", json!("word"))] + #[case("include_blocks", json!(true))] + #[case("id", json!("req-123"))] + fn request_mapping_matches_python(#[case] name: &str, #[case] value: Value) { + let params: OpaqueParams = serde_json::from_value(json!({name: value.clone()})).unwrap(); + let result = serde_json::to_value( + MistralOCRConfig + .transform_ocr_request("model", document(), ¶ms, &[]) + .unwrap(), + ) + .unwrap(); + assert_eq!(result["model"], "model"); + assert_eq!(result[name], value); + } + + #[rstest] + #[case("table_format", json!("html"))] + #[case("confidence_scores_granularity", json!("word"))] + #[case("document_annotation_prompt", json!("extract"))] + #[case("id", json!("req-123"))] + #[case("extract_header", json!(true))] + #[case("include_blocks", json!(true))] + #[case("pages", json!([0,1]))] + fn transform_ocr_request_includes_each_optional_param( + #[case] name: &str, + #[case] value: Value, + ) { + let params: OpaqueParams = serde_json::from_value(json!({name:value.clone()})).unwrap(); + let result = serde_json::to_value( + MistralOCRConfig + .transform_ocr_request("mistral-ocr-latest", document(), ¶ms, &[]) + .unwrap(), + ) + .unwrap(); + assert_eq!(result[name], value); + assert_eq!(result["model"], "mistral-ocr-latest"); + } + + #[rstest] + fn transform_ocr_request_includes_multiple_new_params() { + let params: OpaqueParams = serde_json::from_value(json!({ + "table_format":"html", + "confidence_scores_granularity":"page", + "extract_header":true + })) + .unwrap(); + let result = serde_json::to_value( + MistralOCRConfig + .transform_ocr_request("mistral-ocr-latest", document(), ¶ms, &[]) + .unwrap(), + ) + .unwrap(); + assert_eq!(result["table_format"], "html"); + assert_eq!(result["confidence_scores_granularity"], "page"); + assert_eq!(result["extract_header"], true); + } + + #[rstest] + fn transform_ocr_response_preserves_blocks_and_confidence_scores() { + let response: MistralOcrResponse = serde_json::from_value(json!({ + "pages":[{ + "index":0, + "markdown":"hello", + "images":[{"id":"img-0","image_base64":"data:image/png;base64,AA=="}], + "dimensions":{"width":612,"height":792,"dpi":72}, + "blocks":[{"type":"title","bbox":{"x":1},"confidence_scores":{"mean":0.98}}], + "confidence_scores":{"average_page_confidence_score":0.99,"minimum_page_confidence_score":0.97} + }], + "model":"returned-model", + "document_annotation":"{\"language\":\"en\"}", + "usage_info":{"pages_processed":1} + })) + .unwrap(); + let result = normalize_response("model", response).unwrap().into_json(); + assert_eq!(result["pages"][0]["blocks"][0]["type"], "title"); + assert_eq!(result["pages"][0]["blocks"][0]["bbox"]["x"], 1); + assert_eq!( + result["pages"][0]["blocks"][0]["confidence_scores"]["mean"], + 0.98 + ); + assert_eq!( + result["pages"][0]["confidence_scores"]["average_page_confidence_score"], + 0.99 + ); + assert_eq!(result["pages"][0]["images"][0]["id"], "img-0"); + assert_eq!(result["pages"][0]["dimensions"]["dpi"], 72); + assert_eq!(result["model"], "returned-model"); + assert_eq!(result["document_annotation"], "{\"language\":\"en\"}"); + assert_eq!(result["usage_info"]["pages_processed"], 1); + } + + #[rstest] + fn transform_ocr_response_preserves_ocr4_page_fields() { + let page = json!({ + "index":0, + "markdown":"table page", + "tables":[{"rows":2,"cols":3}], + "hyperlinks":["https://example.com"], + "header":"header", + "footer":"footer" + }); + let response: MistralOcrResponse = + serde_json::from_value(json!({"pages":[page.clone()]})).unwrap(); + let result = normalize_response("model", response).unwrap().into_json(); + assert_eq!(result["pages"][0]["tables"], page["tables"]); + assert_eq!(result["pages"][0]["hyperlinks"], page["hyperlinks"]); + assert_eq!(result["pages"][0]["header"], page["header"]); + assert_eq!(result["pages"][0]["footer"], page["footer"]); + assert!(result["pages"][0]["images"].is_null()); + assert!(result["pages"][0]["dimensions"].is_null()); + } + + #[test] + fn complete_url_defaults_and_dedupes_v1() { + assert_eq!( + MistralOCRConfig.get_complete_url(None).unwrap(), + "https://api.mistral.ai/v1/ocr" + ); + assert_eq!( + MistralOCRConfig + .get_complete_url(Some("https://example.com/v1?tenant=a")) + .unwrap(), + "https://example.com/v1/ocr?tenant=a" + ); + assert_eq!( + MistralOCRConfig + .get_complete_url(Some("https://example.com/v1/ocr?tenant=a")) + .unwrap(), + "https://example.com/v1/ocr?tenant=a" + ); + } + + #[test] + fn environment_prefers_explicit_key_then_environment() { + let explicit = OcrConnection { + api_key: Some("explicit".into()), + ..OcrConnection::default() + }; + assert_eq!( + MistralOCRConfig + .validate_environment(&explicit, &|_| Some("environment".into())) + .unwrap()[0], + ("Authorization".into(), "Bearer explicit".into()) + ); + + assert_eq!( + MistralOCRConfig + .validate_environment(&OcrConnection::default(), &|_| Some("environment".into())) + .unwrap()[0], + ("Authorization".into(), "Bearer environment".into()) + ); + } + + #[test] + fn environment_preserves_forwarded_authorization() { + let connection = OcrConnection { + extra_headers: vec![("authorization".into(), "Bearer forwarded".into())], + ..OcrConnection::default() + }; + assert_eq!( + MistralOCRConfig + .validate_environment(&connection, &|_| None) + .unwrap(), + connection.extra_headers + ); + } + + #[test] + fn environment_rejects_missing_key() { + assert!(matches!( + MistralOCRConfig.validate_environment(&OcrConnection::default(), &|_| None), + Err(crate::ocr::Error::Auth( + litellm_auth::Error::MissingApiKey { + provider: "Mistral", + environment_variable: MISTRAL_API_KEY_ENV, + } + )) + )); + } +} diff --git a/litellm-rust/crates/core/src/llms/mod.rs b/litellm-rust/crates/core/src/llms/mod.rs new file mode 100644 index 00000000000..3dad380f833 --- /dev/null +++ b/litellm-rust/crates/core/src/llms/mod.rs @@ -0,0 +1,6 @@ +pub(crate) mod azure_ai; +pub(crate) mod base_llm; +pub(crate) mod cohere; +pub(crate) mod mistral; +pub(crate) mod reducto; +pub(crate) mod vertex_ai; diff --git a/litellm-rust/crates/core/src/llms/reducto/mod.rs b/litellm-rust/crates/core/src/llms/reducto/mod.rs new file mode 100644 index 00000000000..079e0c41eae --- /dev/null +++ b/litellm-rust/crates/core/src/llms/reducto/mod.rs @@ -0,0 +1 @@ +pub(crate) mod ocr; diff --git a/litellm-rust/crates/core/src/llms/reducto/ocr/mod.rs b/litellm-rust/crates/core/src/llms/reducto/ocr/mod.rs new file mode 100644 index 00000000000..080f0a1183f --- /dev/null +++ b/litellm-rust/crates/core/src/llms/reducto/ocr/mod.rs @@ -0,0 +1 @@ +pub(crate) mod transformation; diff --git a/litellm-rust/crates/core/src/llms/reducto/ocr/transformation.rs b/litellm-rust/crates/core/src/llms/reducto/ocr/transformation.rs new file mode 100644 index 00000000000..f4ed5946fac --- /dev/null +++ b/litellm-rust/crates/core/src/llms/reducto/ocr/transformation.rs @@ -0,0 +1,1018 @@ +use std::collections::BTreeMap; + +use serde::{Deserialize, Deserializer, Serialize}; +use serde_json::{Map, Value, json}; + +use crate::call_arguments::{CallArguments, compose_body}; +use crate::constants::{REDUCTO_API_BASE, REDUCTO_API_KEY_ENV, REDUCTO_ID_PREFIX}; +use crate::llms::base_llm::ocr::transformation::{BaseOcrConfig, OcrRequestContext}; +use crate::ocr::OcrClient; +use crate::ocr::document::InlineDocument; +use crate::ocr::prepare::{build_http_request, credential_env, guardrail_document}; +use crate::ocr::types::{ + LiteLLMOcrResponse, OcrConnection, OcrDocument, OcrPage, OcrUsageInfo, PreparedOcrRequest, +}; +use crate::params::OpaqueParams; +use crate::url_utils::ApiUrl; + +#[derive(Clone, Debug, Serialize, Deserialize)] +#[serde(transparent)] +pub(crate) struct ReductoFileId(String); + +pub(crate) type ReductoV3Params = OpaqueParams; +pub(crate) type ReductoLegacyParams = OpaqueParams; + +#[derive(Clone, Debug, Serialize, Deserialize)] +pub(crate) struct ReductoV3Request { + pub input: ReductoFileId, + #[serde(flatten)] + pub params: ReductoV3Params, +} + +#[derive(Clone, Debug, Serialize, Deserialize)] +pub(crate) struct ReductoLegacyRequest { + pub document_url: ReductoFileId, + #[serde(skip_serializing_if = "Option::is_none")] + pub options: Option, +} + +#[derive(Clone, Debug, Serialize, Deserialize)] +pub(crate) struct ReductoLegacyOptions { + pub enhance: Value, +} + +#[derive(Deserialize)] +struct ReductoUploadResponse { + pub file_id: Option, +} + +#[derive(Clone, Debug, Deserialize)] +pub(crate) struct ReductoResponse { + #[serde(default, deserialize_with = "present_nullable")] + result: Option>, + usage: Option, + #[serde(default)] + chunks: Option>, +} + +#[derive(Clone, Debug, Default, Deserialize)] +struct ReductoResult { + pub chunks: Option>, +} + +#[serde_with::serde_as] +#[derive(Clone, Debug, Default, Deserialize)] +struct ReductoUsage { + #[serde_as(deserialize_as = "Option")] + pub num_pages: Option, + #[serde_as(deserialize_as = "Option")] + pub credits: Option, +} + +#[derive(Clone, Debug, Deserialize)] +struct ReductoChunk { + pub content: Option, + pub blocks: Option>>, +} + +#[derive(Clone, Debug)] +pub(crate) struct ReductoParseV3Config; + +impl BaseOcrConfig for ReductoParseV3Config { + type OcrParams = ReductoV3Params; + type ProviderRequest = ReductoV3Request; + type Environment = Vec<(String, String)>; + + async fn validate_environment( + &self, + request: &PreparedOcrRequest, + _client: &OcrClient, + ) -> Result { + validate_environment(&request.connection, &credential_env) + } + + fn get_complete_url( + &self, + request: &PreparedOcrRequest, + _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"] + } + + fn map_ocr_params( + &self, + arguments: &CallArguments, + model: &str, + ) -> Result { + Ok(arguments + .select(self.get_supported_ocr_params(model)) + .into()) + } + + async fn async_transform_ocr_request( + &self, + _model: &str, + document: OcrDocument, + optional_params: &ReductoV3Params, + headers: &[(String, String)], + context: OcrRequestContext<'_>, + ) -> Result { + let file_id = ensure_file_id_async(document, headers, context).await?; + Ok(ReductoV3Request { + input: file_id, + params: optional_params.clone(), + }) + } + + fn transform_ocr_response( + &self, + model: &str, + raw_response: &[u8], + request_format: crate::ocr::types::OcrResponseFormat, + ) -> Result { + crate::llms::base_llm::ocr::transformation::decode_and_normalize_response( + model, + raw_response, + request_format, + normalize_response, + ) + } + + async fn prepare_request( + &self, + request: &PreparedOcrRequest, + client: &OcrClient, + ) -> Result { + prepare_upload_request(self, request, client).await + } +} + +#[derive(Clone, Debug)] +pub(crate) struct ReductoParseLegacyConfig; + +impl BaseOcrConfig for ReductoParseLegacyConfig { + type OcrParams = ReductoLegacyParams; + type ProviderRequest = ReductoLegacyRequest; + type Environment = Vec<(String, String)>; + + async fn validate_environment( + &self, + request: &PreparedOcrRequest, + client: &OcrClient, + ) -> Result { + ReductoParseV3Config + .validate_environment(request, client) + .await + } + + fn get_complete_url( + &self, + request: &PreparedOcrRequest, + 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"] + } + + fn map_ocr_params( + &self, + arguments: &CallArguments, + model: &str, + ) -> Result { + Ok(arguments + .select(self.get_supported_ocr_params(model)) + .into()) + } + + async fn async_transform_ocr_request( + &self, + _model: &str, + document: OcrDocument, + optional_params: &ReductoLegacyParams, + headers: &[(String, String)], + context: OcrRequestContext<'_>, + ) -> Result { + let file_id = ensure_file_id_async(document, headers, context).await?; + Ok(build_legacy_body(file_id, optional_params)) + } + + fn transform_ocr_response( + &self, + model: &str, + raw_response: &[u8], + request_format: crate::ocr::types::OcrResponseFormat, + ) -> Result { + ReductoParseV3Config.transform_ocr_response(model, raw_response, request_format) + } + + async fn prepare_request( + &self, + request: &PreparedOcrRequest, + client: &OcrClient, + ) -> Result { + prepare_upload_request(self, request, client).await + } +} + +/// Reducto differs from the shared `BaseOcrConfig::prepare_request` flow: +/// guardrails see the *source* document before it is uploaded, because the +/// final body only carries the opaque Reducto file id. +async fn prepare_upload_request>>( + config: &C, + request: &PreparedOcrRequest, + client: &OcrClient, +) -> Result { + let params = config.map_ocr_params(&request.optional_params, &request.model)?; + let headers = config.validate_environment(request, client).await?; + let url = config.get_complete_url(request, ¶ms, &headers)?; + let (document, headers) = guardrail_document(request, &url, &headers).await?; + let body = config + .async_transform_ocr_request( + &request.model, + document, + ¶ms, + &headers, + OcrRequestContext { + client, + connection: &request.connection, + }, + ) + .await?; + let body = compose_body( + &request.optional_params, + &body, + config.get_supported_ocr_params(&request.model), + )?; + build_http_request(client, request, &url, &headers, &body) +} + +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> { + Option::::deserialize(deserializer).map(Some) +} + +fn block_page_number(value: &Value) -> Option { + match value { + Value::Number(number) => number + .as_i64() + .or_else(|| number.as_f64().and_then(checked_truncated_i64)), + Value::String(value) => value.trim().parse::().ok(), + Value::Bool(value) => Some(i64::from(*value)), + _ => None, + } +} + +fn checked_truncated_i64(value: f64) -> Option { + (value.is_finite() && value >= i64::MIN as f64 && value <= i64::MAX as f64) + .then(|| value.trunc() as i64) +} + +pub(crate) fn normalize_response( + model: &str, + response: ReductoResponse, +) -> Result { + let result = match response.result { + Some(result) => result.unwrap_or_default(), + None => ReductoResult { + chunks: response.chunks, + }, + }; + let usage = response.usage.unwrap_or_default(); + Ok(LiteLLMOcrResponse { + usage_info: Some(OcrUsageInfo { + pages_processed: usage.num_pages, + credits: usage.credits, + ..Default::default() + }), + ..LiteLLMOcrResponse::new( + model, + build_pages_from_reducto(result.chunks.unwrap_or_default())?, + ) + }) +} + +fn build_pages_from_reducto(chunks: Vec) -> Result, crate::ocr::Error> { + let blocks_by_page = chunks + .iter() + .flat_map(|chunk| chunk.blocks.iter().flatten()) + .filter_map(|block| { + block_page_number(block.get("bbox")?.get("page")?).map(|page| (page, block)) + }) + .fold( + BTreeMap::>>::new(), + |mut pages, (page, block)| { + pages.entry(page).or_default().push(block); + pages + }, + ); + if blocks_by_page.is_empty() { + let markdown = join_content(chunks.iter().map(|chunk| chunk.content.as_deref())); + return Ok(if markdown.is_empty() { + Vec::new() + } else { + vec![page(0, markdown, None)] + }); + } + blocks_by_page + .into_iter() + .map(|(index, blocks)| { + let content = blocks + .iter() + .map(|block| match block.get("content") { + None | Some(Value::Null) => Ok(None), + Some(Value::String(content)) => Ok(Some(content.as_str())), + Some(_) => Err(crate::ocr::Error::ResponseField { + path: "result.chunks.blocks.content".into(), + }), + }) + .collect::, _>>()?; + let markdown = join_content(content.into_iter()); + Ok(page( + index.saturating_sub(1).max(0), + markdown, + Some(json!(blocks)), + )) + }) + .collect() +} + +fn join_content<'a>(content: impl Iterator>) -> String { + content + .flatten() + .filter(|text| !text.is_empty()) + .collect::>() + .join("\n\n") +} + +fn page(index: i64, markdown: String, blocks: Option) -> OcrPage { + OcrPage { + index, + markdown, + extra_fields: blocks + .map(|blocks| ("blocks".into(), blocks)) + .into_iter() + .collect(), + ..Default::default() + } +} +fn get_complete_url(api_base: Option<&str>) -> Result { + complete_endpoint_url(api_base, "parse") +} + +fn complete_endpoint_url(api_base: Option<&str>, path: &str) -> Result { + let base = api_base + .map(str::trim) + .filter(|base| !base.is_empty()) + .unwrap_or(REDUCTO_API_BASE); + ApiUrl::parse(base) + .and_then(|url| url.complete_path(&[path])) + .map(|url| url.into_string()) + .map_err(|_| crate::ocr::Error::RequestField { + path: "api_base".into(), + }) +} + +fn validate_environment( + connection: &OcrConnection, + env_lookup: &(dyn Fn(&str) -> Option + Sync), +) -> Result, crate::ocr::Error> { + if crate::http_utils::has_header(&connection.extra_headers, "authorization") { + return Ok(connection.extra_headers.clone()); + } + let api_key = connection + .api_key + .as_deref() + .map(str::trim) + .filter(|key| !key.is_empty()) + .map(str::to_string) + .or_else(|| { + env_lookup(REDUCTO_API_KEY_ENV) + .map(|key| key.trim().to_string()) + .filter(|key| !key.is_empty()) + }) + .ok_or(crate::ocr::Error::MissingReductoApiKey)?; + Ok( + std::iter::once(("Authorization".into(), format!("Bearer {api_key}"))) + .chain(connection.extra_headers.clone()) + .collect(), + ) +} + +fn build_legacy_body( + file_id: ReductoFileId, + optional_params: &ReductoLegacyParams, +) -> ReductoLegacyRequest { + ReductoLegacyRequest { + document_url: file_id, + options: optional_params + .get("enhance") + .filter(|value| !value.is_null()) + .map(|enhance| ReductoLegacyOptions { + enhance: enhance.clone(), + }), + } +} + +async fn ensure_file_id_async( + document: OcrDocument, + headers: &[(String, String)], + context: OcrRequestContext<'_>, +) -> Result { + if document.source().starts_with(REDUCTO_ID_PREFIX) { + if document.source()[REDUCTO_ID_PREFIX.len()..] + .trim() + .is_empty() + { + return Err(crate::ocr::Error::RequestField { + path: "document file id".into(), + }); + } + return Ok(ReductoFileId(document.source().to_string())); + } + let inline = + InlineDocument::parse(document.source())?.ok_or(crate::ocr::Error::ReductoSource)?; + let mime = inline.mime_type().to_string(); + let bytes = inline.decode(crate::constants::OCR_INLINE_MAX_BYTES)?; + upload_bytes_async(bytes, &mime, headers, context).await +} + +async fn upload_bytes_async( + bytes: Vec, + mime: &str, + headers: &[(String, String)], + context: OcrRequestContext<'_>, +) -> Result { + let OcrRequestContext { client, connection } = context; + let part = reqwest::multipart::Part::bytes(bytes) + .file_name("document") + .mime_str(mime) + .map_err(|_| crate::ocr::Error::InvalidDataUri)?; + let builder = client + .provider_http() + .post(complete_endpoint_url( + connection.api_base.as_deref(), + "upload", + )?) + .multipart(reqwest::multipart::Form::new().part("file", part)) + .timeout(connection.timeout); + let builder = crate::http_utils::with_headers( + builder, + headers, + crate::http_utils::HeaderPolicy::Except(&["content-type", "content-length"]), + ); + let response = crate::http_utils::http_request(builder) + .await + .map_err(crate::transport::Error::from)?; + let uploaded = crate::ocr::client::read_json_response::( + response, + false, + connection.max_response_bytes, + ) + .await? + .data; + let file_id = uploaded + .file_id + .as_deref() + .map(str::trim) + .filter(|id| !id.is_empty()); + let Some(file_id) = file_id else { + return Err(crate::ocr::Error::ResponseField { + path: "file_id".into(), + }); + }; + Ok(ReductoFileId(file_id.to_string())) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn options_preserve_null_and_select_the_provider_fields() { + let overrides = serde_json::from_value(json!({ + "formatting":null, "enhance":null, "ignored":true + })) + .unwrap(); + let v3 = ReductoParseV3Config + .map_ocr_params(&overrides, "parse-v3") + .unwrap(); + assert_eq!( + serde_json::to_value(v3).unwrap(), + json!({ + "formatting":null + }) + ); + let legacy = ReductoParseLegacyConfig + .map_ocr_params(&overrides, "parse-legacy") + .unwrap(); + assert_eq!( + serde_json::to_value(legacy).unwrap(), + json!({ + "enhance":null + }) + ); + } + + #[test] + fn usage_uses_shared_validation_while_block_page_numbers_are_best_effort() { + for usage in [ + json!({"num_pages":1.5}), + json!({"num_pages":[]}), + json!({"credits":{}}), + ] { + assert!(serde_json::from_value::(json!({"usage":usage})).is_err()); + } + let response = serde_json::from_value(json!({"result":{"chunks":[{"blocks":[ + {"content":"ignored", "bbox":{"page":"invalid"}}, + {"content":"kept", "bbox":{"page":2.5}, "extra":null} + ]}]}, "usage":{"num_pages":2.0, "credits":true}})) + .unwrap(); + let normalized = normalize_response("model", response).unwrap(); + assert_eq!(normalized.pages[0].index, 1); + assert_eq!(normalized.pages[0].markdown, "kept"); + assert_eq!( + normalized.pages[0].extra_fields["blocks"][0]["bbox"]["page"], + 2.5 + ); + assert_eq!(normalized.usage_info.unwrap().credits, Some(1.0)); + } + + #[tokio::test] + async fn v3_options_preserve_explicit_null() { + let overrides = + serde_json::from_value(json!({"formatting":null,"settings":{},"unknown":true})) + .unwrap(); + let params = ReductoParseV3Config + .map_ocr_params(&overrides, "parse-v3") + .unwrap(); + let client = crate::ocr::test_support::ocr_client(); + let connection = OcrConnection::default(); + let document = serde_json::from_value( + json!({"type":"document_url","document_url":"reducto://ready.pdf"}), + ) + .unwrap(); + let body = ReductoParseV3Config + .async_transform_ocr_request( + "parse-v3", + document, + ¶ms, + &[], + OcrRequestContext { + client: &client, + connection: &connection, + }, + ) + .await + .unwrap(); + assert_eq!( + serde_json::to_value(body).unwrap(), + json!({ + "input":"reducto://ready.pdf", "formatting":null, "settings":{} + }) + ); + let absent = ReductoParseV3Config + .map_ocr_params(&crate::call_arguments::CallArguments::default(), "parse-v3") + .unwrap(); + assert_eq!(serde_json::to_value(absent).unwrap(), json!({})); + } + + #[test] + fn legacy_body_omits_null_enhance_and_wraps_mapped_options() { + for (value, expected) in [ + (json!(null), json!({"document_url":"reducto://ready.pdf"})), + ( + json!({}), + json!({"document_url":"reducto://ready.pdf","options":{"enhance":{}}}), + ), + ] { + let overrides = + serde_json::from_value(json!({"enhance":value,"unknown":true})).unwrap(); + let params = ReductoParseLegacyConfig + .map_ocr_params(&overrides, "parse-legacy") + .unwrap(); + assert_eq!( + serde_json::to_value(build_legacy_body( + ReductoFileId("reducto://ready.pdf".into()), + ¶ms + )) + .unwrap(), + expected + ); + } + } + + #[test] + fn explicit_key_precedes_environment_key() { + let connection = OcrConnection { + api_key: Some("passed-key".into()), + ..Default::default() + }; + let headers = validate_environment(&connection, &|_| Some("env-key".into())).unwrap(); + assert_eq!(headers[0].1, "Bearer passed-key"); + } + + #[test] + fn blank_explicit_key_uses_environment_key() { + let connection = OcrConnection { + api_key: Some(" ".into()), + ..Default::default() + }; + let headers = validate_environment(&connection, &|_| Some(" env-key ".into())).unwrap(); + assert_eq!(headers[0].1, "Bearer env-key"); + } + + #[test] + fn existing_authorization_skips_key_lookup() { + let connection = OcrConnection { + extra_headers: vec![("authorization".into(), "Bearer existing".into())], + ..Default::default() + }; + assert_eq!( + validate_environment(&connection, &|_| None).unwrap(), + connection.extra_headers + ); + } + + 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 request = + crate::ocr::test_support::with_source(wire_request(model, &base, options), source); + + 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.transport.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 request = crate::ocr::test_support::with_source( + wire_request("reducto/parse-v3", "http://127.0.0.1:1", json!({})), + source, + ); + 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 = crate::ocr::test_support::with_source( + wire_request("reducto/parse-v3", &base, json!({})), + "reducto://ready.pdf", + ); + request.transport.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.transport.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/mod.rs b/litellm-rust/crates/core/src/llms/vertex_ai/mod.rs new file mode 100644 index 00000000000..079e0c41eae --- /dev/null +++ b/litellm-rust/crates/core/src/llms/vertex_ai/mod.rs @@ -0,0 +1 @@ +pub(crate) mod ocr; diff --git a/litellm-rust/crates/core/src/llms/vertex_ai/ocr/common_utils.rs b/litellm-rust/crates/core/src/llms/vertex_ai/ocr/common_utils.rs new file mode 100644 index 00000000000..6340084ad7f --- /dev/null +++ b/litellm-rust/crates/core/src/llms/vertex_ai/ocr/common_utils.rs @@ -0,0 +1,9 @@ +use crate::ocr::types::OcrConnection; +use litellm_auth::InputSource; + +pub(super) fn validate_destination(connection: &OcrConnection) -> Result<(), crate::ocr::Error> { + if connection.api_base.is_some() && connection.api_base_source == InputSource::Request { + return Err(litellm_auth::Error::RequestVertexCredentialDestination.into()); + } + Ok(()) +} 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 new file mode 100644 index 00000000000..7caa4656678 --- /dev/null +++ b/litellm-rust/crates/core/src/llms/vertex_ai/ocr/deepseek_transformation.rs @@ -0,0 +1,705 @@ +use serde::{Deserialize, Serialize}; +use serde_json::{Map, Value}; + +use litellm_auth_gcp::{self as vertex, VertexConfig}; + +use super::transformation::VertexAIOCRConfig; +use crate::call_arguments::CallArguments; +use crate::llms::base_llm::ocr::transformation::{BaseOcrConfig, OcrRequestContext}; +use crate::ocr::OcrClient; +use crate::ocr::prepare::credential_env; +use crate::ocr::types::{ + LiteLLMOcrResponse, OcrDocument, OcrPage, OcrPageDimensions, OcrPageImage, OcrUsageInfo, + PreparedOcrRequest, +}; +use crate::params::OpaqueParams; +use crate::providers::model::{ModelNamespace, ProviderModel, RoutedModel}; +use crate::url_utils::ApiUrl; + +const DEFAULT_API_BASE: &str = "https://aiplatform.googleapis.com"; +const MODEL_NAMESPACE: &str = "deepseek-ai"; +const DEFAULT_LOCATION: &str = "us-central1"; +const DEEPSEEK_OCR_PARAMS: &[&str] = &["stream", "temperature", "max_tokens", "top_p", "n", "stop"]; + +pub(crate) type DeepSeekOcrParams = OpaqueParams; + +#[derive(Clone, Debug, Serialize, Deserialize)] +pub(crate) struct DeepSeekOcrRequest { + pub model: ProviderModel, + pub messages: Vec, + #[serde(flatten)] + pub params: OpaqueParams, +} + +#[derive(Clone, Debug, Serialize, Deserialize)] +pub(crate) struct DeepSeekOcrMessage { + pub role: UserRole, + pub content: Vec, +} + +#[derive(Clone, Debug, Serialize, Deserialize)] +#[serde(tag = "type")] +pub(crate) enum DeepSeekDocument { + #[serde(rename = "image_url")] + ImageUrl { image_url: String }, +} + +#[derive(Clone, Debug, Serialize, Deserialize)] +#[serde(rename_all = "lowercase")] +pub(crate) enum UserRole { + User, +} + +#[derive(Clone, Debug, Deserialize)] +pub(crate) struct DeepSeekOcrResponse { + #[serde(default)] + choices: Vec, + #[serde(default = "empty_object")] + usage: Value, +} + +#[derive(Clone, Debug, Deserialize)] +struct DeepSeekChoice { + #[serde(default)] + message: DeepSeekResponseMessage, +} + +#[derive(Clone, Debug, Default, Deserialize)] +struct DeepSeekResponseMessage { + content: Option, +} + +#[derive(Clone, Debug, Deserialize)] +#[serde(untagged)] +enum DeepSeekContent { + Text(String), + Object(Map), +} + +#[serde_with::serde_as] +#[derive(Deserialize)] +struct DeepSeekPage { + #[serde(default)] + #[serde_as(deserialize_as = "crate::serde_compat::LaxI64")] + index: i64, + #[serde(default)] + markdown: String, + images: Option>, + dimensions: Option, +} + +#[derive(Clone, Debug)] +pub(crate) struct DeepSeekAi; + +impl ModelNamespace for DeepSeekAi { + const NAME: &'static str = MODEL_NAMESPACE; +} + +#[derive(Clone, Debug)] +pub(crate) struct VertexAIDeepSeekOCRConfig; + +impl BaseOcrConfig for VertexAIDeepSeekOCRConfig { + type OcrParams = DeepSeekOcrParams; + type ProviderRequest = DeepSeekOcrRequest; + type Environment = vertex::VertexEnvironment; + + fn get_api_key_env_var(&self) -> Option<&'static str> { + VertexAIOCRConfig.get_api_key_env_var() + } + + fn map_ocr_params( + &self, + _arguments: &CallArguments, + _model: &str, + ) -> Result { + Ok(DeepSeekOcrParams::default()) + } + + async fn validate_environment( + &self, + request: &PreparedOcrRequest, + client: &OcrClient, + ) -> Result { + BaseOcrConfig::validate_environment(&VertexAIOCRConfig, request, client).await + } + + fn get_complete_url( + &self, + request: &PreparedOcrRequest, + _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, + model: &str, + document: OcrDocument, + optional_params: &DeepSeekOcrParams, + headers: &[(String, String)], + _context: OcrRequestContext<'_>, + ) -> Result { + self.transform_ocr_request(model, document, optional_params, headers) + } + + fn transform_ocr_response( + &self, + model: &str, + raw_response: &[u8], + request_format: crate::ocr::types::OcrResponseFormat, + ) -> Result { + crate::llms::base_llm::ocr::transformation::decode_and_normalize_response( + model, + raw_response, + request_format, + normalize_response, + ) + } + + fn transform_ocr_request( + &self, + model: &str, + document: OcrDocument, + optional_params: &DeepSeekOcrParams, + _headers: &[(String, String)], + ) -> Result { + if document.source().is_empty() { + return Err(crate::ocr::Error::MissingDocumentUrl); + } + Ok(DeepSeekOcrRequest { + model: provider_model(model)?, + messages: vec![DeepSeekOcrMessage { + role: UserRole::User, + content: vec![DeepSeekDocument::ImageUrl { + image_url: document.source().to_string(), + }], + }], + params: optional_params + .iter() + .filter(|(name, _)| DEEPSEEK_OCR_PARAMS.contains(&name.as_str())) + .map(|(name, value)| (name.clone(), value.clone())) + .collect(), + }) + } +} + +pub(crate) fn normalize_response( + model: &str, + response: DeepSeekOcrResponse, +) -> Result { + let content = response + .choices + .into_iter() + .next() + .and_then(|choice| choice.message.content) + .ok_or(crate::ocr::Error::EmptyContent)?; + let (ocr_data, fallback_markdown) = match content { + DeepSeekContent::Text(text) if text.is_empty() => { + return Err(crate::ocr::Error::EmptyContent); + } + DeepSeekContent::Text(text) => { + let parsed = text + .trim_start() + .starts_with('{') + .then(|| serde_json::from_str::>(&text).ok()) + .flatten(); + (parsed.unwrap_or_default(), text) + } + DeepSeekContent::Object(data) if data.is_empty() => { + return Err(crate::ocr::Error::EmptyContent); + } + DeepSeekContent::Object(data) => { + let fallback = if data.contains_key("pages") { + String::new() + } else { + let mut output = Vec::new(); + data.serialize(&mut serde_json::Serializer::with_formatter( + &mut output, + PythonJsonFormatter, + )) + .map_err(|_| response_field("content"))?; + String::from_utf8(output).map_err(|_| response_field("content"))? + }; + (data, fallback) + } + }; + let has_pages = ocr_data.contains_key("pages"); + let pages = match ocr_data.get("pages") { + Some(Value::Array(pages)) => pages + .iter() + .enumerate() + .filter(|(_, page)| page.is_object()) + .map(|(position, page)| { + let page: DeepSeekPage = crate::ocr::json::decode_response_value( + page.clone(), + &format!("choices[0].message.content.pages[{position}]"), + )?; + Ok(OcrPage { + index: page.index, + markdown: page.markdown, + images: page.images, + dimensions: page.dimensions, + ..Default::default() + }) + }) + .collect::, crate::ocr::Error>>()?, + Some(_) => return Err(response_field("pages")), + None => Vec::new(), + }; + let usage = ocr_data + .get("usage_info") + .or_else(|| (!has_pages).then_some(&response.usage)); + let usage_info: Option = usage + .filter(|usage| usage.is_object()) + .map(|usage| crate::ocr::json::decode_response_value(usage.clone(), "usage_info")) + .transpose()?; + let model = match ocr_data.get("model") { + Some(Value::String(model)) => model.clone(), + Some(_) => return Err(response_field("model")), + None => model.to_string(), + }; + Ok(LiteLLMOcrResponse { + extra_fields: ocr_data + .iter() + .filter(|(name, _)| { + !matches!( + name.as_str(), + "pages" + | "model" + | "document_annotation" + | "usage_info" + | "object" + | "content" + | "tables" + | "keyValuePairs" + | "provider_native_response" + ) + }) + .map(|(name, value)| (name.clone(), value.clone())) + .collect(), + document_annotation: has_pages + .then(|| ocr_data.get("document_annotation").cloned()) + .flatten(), + usage_info, + ..LiteLLMOcrResponse::new( + model, + if pages.is_empty() { + vec![OcrPage { + markdown: fallback_markdown, + ..Default::default() + }] + } else { + pages + }, + ) + }) +} + +fn empty_object() -> Value { + Value::Object(Map::new()) +} + +struct PythonJsonFormatter; + +impl serde_json::ser::Formatter for PythonJsonFormatter { + fn begin_array_value( + &mut self, + writer: &mut W, + first: bool, + ) -> std::io::Result<()> { + if first { + Ok(()) + } else { + writer.write_all(b", ") + } + } + + fn begin_object_key( + &mut self, + writer: &mut W, + first: bool, + ) -> std::io::Result<()> { + if first { + Ok(()) + } else { + writer.write_all(b", ") + } + } + + fn begin_object_value( + &mut self, + writer: &mut W, + ) -> std::io::Result<()> { + writer.write_all(b": ") + } + + fn write_string_fragment( + &mut self, + writer: &mut W, + fragment: &str, + ) -> std::io::Result<()> { + for character in fragment.chars() { + if character.is_ascii() && character != '\u{7f}' { + writer.write_all(&[character as u8])?; + } else { + for unit in character.encode_utf16(&mut [0; 2]) { + write!(writer, "\\u{unit:04x}")?; + } + } + } + Ok(()) + } +} + +fn response_field(field: &str) -> crate::ocr::Error { + crate::ocr::Error::ResponseField { + path: format!("choices[0].message.content.{field}"), + } +} + +pub(crate) fn provider_model(model: &str) -> Result, crate::ocr::Error> { + RoutedModel::new(model) + .and_then(RoutedModel::into_provider::) + .map_err(|_| crate::ocr::Error::RequestField { + path: "model".into(), + }) +} + +impl VertexAIDeepSeekOCRConfig { + fn get_complete_url( + &self, + api_base: Option<&str>, + project: &str, + location: &str, + ) -> Result { + let base = api_base + .map(str::trim) + .filter(|base| !base.is_empty()) + .unwrap_or(DEFAULT_API_BASE); + ApiUrl::parse(base) + .and_then(|url| { + url.complete_path(&[ + "v1", + "projects", + project, + "locations", + location, + "endpoints", + "openapi", + "chat", + "completions", + ]) + }) + .map(|url| url.into_string()) + .map_err(|_| crate::ocr::Error::RequestField { + path: "api_base".into(), + }) + } +} + +#[cfg(test)] +mod tests { + use super::{ + DeepSeekOcrParams, DeepSeekOcrResponse, VertexAIDeepSeekOCRConfig, normalize_response, + provider_model, + }; + use serde_json::{Value, json}; + + #[test] + fn unconsumed_options_remain_available_for_body_composition() { + use crate::llms::base_llm::ocr::transformation::BaseOcrConfig; + use serde_json::json; + + let arguments = + serde_json::from_value(json!({"temperature":0.5,"extension":null})).unwrap(); + assert_eq!( + serde_json::to_value( + VertexAIDeepSeekOCRConfig + .map_ocr_params(&arguments, "deepseek-ocr") + .unwrap() + ) + .unwrap(), + json!({}) + ); + assert_eq!( + crate::call_arguments::compose_body(&arguments, &json!({"model":"deepseek-ocr"}), &[]) + .unwrap(), + json!({"model":"deepseek-ocr","temperature":0.5,"extension":null}) + ); + } + + #[test] + fn config_owns_model_namespace_and_endpoint() { + assert_eq!( + provider_model("deepseek-ocr-maas").unwrap().as_str(), + "deepseek-ai/deepseek-ocr-maas" + ); + assert_eq!( + provider_model("deepseek-ai/deepseek-ocr-maas") + .unwrap() + .as_str(), + "deepseek-ai/deepseek-ocr-maas" + ); + assert_eq!( + VertexAIDeepSeekOCRConfig + .get_complete_url(None, "proj-1", "europe-west4") + .unwrap(), + "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_eq!(result["future"], "kept"); + } + + #[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 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"} + }), + ); + let request = crate::ocr::test_support::with_source(request, "gs://bucket/document.pdf"); + + 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::is_supported_request( + "deepseek-ocr-maas", + Some("vertex_ai") + )); + assert!(crate::ocr::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.credentials.api_base = Some(litellm_auth::Sourced::new( + "https://caller.example".into(), + 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/mod.rs b/litellm-rust/crates/core/src/llms/vertex_ai/ocr/mod.rs new file mode 100644 index 00000000000..f894ec145f8 --- /dev/null +++ b/litellm-rust/crates/core/src/llms/vertex_ai/ocr/mod.rs @@ -0,0 +1,3 @@ +pub(crate) mod common_utils; +pub(crate) mod deepseek_transformation; +pub(crate) mod transformation; 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 new file mode 100644 index 00000000000..f71a295e7dd --- /dev/null +++ b/litellm-rust/crates/core/src/llms/vertex_ai/ocr/transformation.rs @@ -0,0 +1,395 @@ +use litellm_auth_gcp::{self as vertex, VertexConfig}; +use serde_json::Value; + +use super::common_utils::validate_destination; +use crate::call_arguments::CallArguments; +use crate::llms::base_llm::ocr::transformation::{ + BaseOcrConfig, OcrEnvironment, OcrRequestContext, +}; +use crate::llms::mistral::ocr::transformation::{MistralOCRConfig, MistralOcrRequest}; +use crate::ocr::OcrClient; +use crate::ocr::document::{inline_remote_document, validate_inline_document}; +use crate::ocr::prepare::credential_env; +use crate::ocr::types::{LiteLLMOcrResponse, OcrConnection, OcrDocument, PreparedOcrRequest}; +use crate::params::OpaqueParams; +use crate::url_utils::ApiUrl; + +const DEFAULT_LOCATION: &str = "us-central1"; + +#[derive(Clone, Debug, Default)] +pub(crate) struct VertexAIOCRConfig; + +impl BaseOcrConfig for VertexAIOCRConfig { + type OcrParams = OpaqueParams; + type ProviderRequest = MistralOcrRequest; + 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: &PreparedOcrRequest, + 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: &PreparedOcrRequest, + _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) + } + + fn map_ocr_params( + &self, + arguments: &CallArguments, + model: &str, + ) -> Result { + MistralOCRConfig.map_ocr_params(arguments, model) + } + + async fn async_transform_ocr_request( + &self, + model: &str, + document: OcrDocument, + optional_params: &OpaqueParams, + headers: &[(String, String)], + context: OcrRequestContext<'_>, + ) -> Result { + let document = inline_remote_document( + context.client.document_fetcher(), + document, + context.connection, + ) + .await?; + self.transform_ocr_request(model, document, optional_params, headers) + } + + fn transform_ocr_response( + &self, + model: &str, + raw_response: &[u8], + request_format: crate::ocr::types::OcrResponseFormat, + ) -> Result { + MistralOCRConfig.transform_ocr_response(model, raw_response, request_format) + } + + fn validate_request_body(&self, body: &Value) -> Result<(), crate::ocr::Error> { + validate_inline_document(&crate::ocr::prepare::body_document(body)?) + } +} + +impl OcrEnvironment for vertex::VertexEnvironment { + fn headers(&self) -> &[(String, String)] { + &self.headers + } +} + +impl VertexAIOCRConfig { + pub(super) async fn validate_environment( + &self, + connection: &OcrConnection, + config: &VertexConfig, + client: &OcrClient, + ) -> Result { + validate_destination(connection)?; + client + .vertex_auth() + .validate_environment( + connection.extra_headers.clone(), + connection.api_key.as_deref(), + config, + &credential_env, + ) + .await + .map_err(crate::ocr::Error::from) + } + + fn get_complete_url( + &self, + api_base: Option<&str>, + project: &str, + location: &str, + model: &str, + ) -> Result { + validate_location(location)?; + let default_base = format!("https://{location}-aiplatform.googleapis.com"); + let base = api_base + .map(str::trim) + .filter(|base| !base.is_empty()) + .unwrap_or(&default_base); + let prediction = format!("{model}:rawPredict"); + ApiUrl::parse(base) + .and_then(|url| { + url.complete_path(&[ + "v1", + "projects", + project, + "locations", + location, + "publishers", + "mistralai", + "models", + &prediction, + ]) + }) + .map(|url| url.into_string()) + .map_err(|_| crate::ocr::Error::RequestField { + path: "api_base".into(), + }) + } +} + +fn validate_location(location: &str) -> Result<(), crate::ocr::Error> { + let valid = !location.is_empty() + && location + .bytes() + .all(|value| value.is_ascii_lowercase() || value.is_ascii_digit() || value == b'-') + && location + .as_bytes() + .first() + .is_some_and(u8::is_ascii_alphanumeric) + && location + .as_bytes() + .last() + .is_some_and(u8::is_ascii_alphanumeric); + if valid { + return Ok(()); + } + Err(crate::ocr::Error::RequestField { + path: "vertex_location".into(), + }) +} + +#[cfg(test)] +mod tests { + use super::VertexAIOCRConfig; + + #[test] + fn endpoint_uses_location_project_and_model() { + assert_eq!( + VertexAIOCRConfig + .get_complete_url(None, "proj-1", "europe-west4", "mistral-ocr-maas") + .unwrap(), + "https://europe-west4-aiplatform.googleapis.com/v1/projects/proj-1/locations/europe-west4/publishers/mistralai/models/mistral-ocr-maas:rawPredict" + ); + assert!( + VertexAIOCRConfig + .get_complete_url(None, "proj-1", "attacker.example/path", "model") + .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.credentials.api_key = None; + request.transport.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.credentials.api_base = Some(litellm_auth::Sourced::new( + "https://caller.example".into(), + 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 = crate::ocr::prepare::prepare_request( + crate::ocr::test_support::resolved_request(direct), + ); + let vertex = crate::ocr::prepare::prepare_request( + crate::ocr::test_support::resolved_request(vertex), + ); + 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_eq!(direct_response["extra"], "preserved"); + } +} diff --git a/litellm-rust/crates/core/src/ocr/adapters/azure/cohere.rs b/litellm-rust/crates/core/src/ocr/adapters/azure/cohere.rs deleted file mode 100644 index 3691e9e1809..00000000000 --- a/litellm-rust/crates/core/src/ocr/adapters/azure/cohere.rs +++ /dev/null @@ -1,131 +0,0 @@ -use super::super::OcrAdapter; -use crate::ocr::Error; -use crate::ocr::OcrClient; -use crate::ocr::codecs::cohere::{ - CohereParams, CohereResponse, transform_request, transform_response, validate_document, -}; -use crate::ocr::document::{inline_remote_document, validate_inline_document}; -use crate::ocr::error::{OcrError, OcrRequestError, OcrResponseError}; -use crate::ocr::prepare::{credential_env, transform_request_body}; -use crate::ocr::registry::OcrProvider; -use crate::ocr::types::{LiteLLMOcrRequest, LiteLLMOcrResponse}; -use crate::url_utils::ApiUrl; -use litellm_auth_azure::AzureAuthInputs; - -const AZURE_AI_API_BASE_ENV: &str = "AZURE_AI_API_BASE"; - -pub(crate) struct AzureCohereAdapter; - -impl OcrAdapter for AzureCohereAdapter { - type ProviderResponse = CohereResponse; - const PROVIDER: OcrProvider = OcrProvider::AzureAi; - - async fn prepare_request( - &self, - request: &LiteLLMOcrRequest, - client: &OcrClient, - ) -> Result { - let params = super::super::super::wire::decode_request_value::( - serde_json::Value::Object(request.optional_params.clone()), - "optional_params", - )?; - let mut config = AzureAuthInputs::from_sourced_optional_params( - &request.optional_params, - &request.input_sources, - ) - .map_err(Error::from)?; - config.azure_ad_token_provider = request.azure_ad_token_provider.clone(); - 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(|| { - Error::Auth( - "Missing Azure AI API Base - Set AZURE_AI_API_BASE or pass api_base".into(), - ) - })?; - let headers = - super::validate_ai_environment(&request.connection, &config, &credential_env).await?; - validate_document(&request.document)?; - let remote = request.document.source().starts_with("http://") - || request.document.source().starts_with("https://"); - let document = inline_remote_document( - client.document_fetcher(), - request.document.clone(), - &request.connection, - ) - .await?; - let body = transform_request(&request.model, document, params)?; - transform_request_body( - client, - request, - &complete_url(&base)?, - &headers, - !remote, - body, - |body| { - validate_document(&body.document)?; - validate_inline_document(&body.document) - }, - ) - .await - } - - fn transform_ocr_response( - &self, - request: &LiteLLMOcrRequest, - response: Self::ProviderResponse, - ) -> Result { - transform_response(&request.model, response) - } -} - -fn complete_url(base: &str) -> Result { - let mut url = reqwest::Url::parse(base).map_err(|_| invalid_api_base())?; - if !matches!(url.scheme(), "http" | "https") { - return Err(invalid_api_base().into()); - } - let path = url.path().trim_end_matches('/').to_string(); - if path.ends_with("/v2/parse") { - url.set_path(&path); - return Ok(url.into()); - } - url.set_path(path.strip_suffix("/models").unwrap_or(&path)); - ApiUrl::parse(url.as_str()) - .and_then(|url| url.complete_path(&["providers", "cohere", "v2", "parse"])) - .map(|url| url.into_string()) - .map_err(|_| invalid_api_base().into()) -} - -fn invalid_api_base() -> OcrRequestError { - OcrRequestError::RequestField { - path: "api_base".into(), - } -} - -#[cfg(test)] -mod tests { - use super::*; - - #[test] - fn completes_foundry_urls_without_duplicate_paths_and_preserves_queries() { - for suffix in [ - "", - "/models", - "/providers/cohere/v2", - "/providers/cohere/v2/parse", - ] { - assert_eq!( - complete_url(&format!("https://example.com{suffix}?tenant=a")).unwrap(), - "https://example.com/providers/cohere/v2/parse?tenant=a" - ); - } - assert_eq!( - complete_url("https://example.com/v2/parse?tenant=a").unwrap(), - "https://example.com/v2/parse?tenant=a" - ); - assert!(complete_url("relative/path").is_err()); - } -} diff --git a/litellm-rust/crates/core/src/ocr/adapters/azure/document_intelligence/mod.rs b/litellm-rust/crates/core/src/ocr/adapters/azure/document_intelligence/mod.rs deleted file mode 100644 index eba300908f1..00000000000 --- a/litellm-rust/crates/core/src/ocr/adapters/azure/document_intelligence/mod.rs +++ /dev/null @@ -1,214 +0,0 @@ -use super::super::OcrAdapter; -use crate::constants::{AZURE_DI_API_VERSION, AZURE_DI_SUBSCRIPTION_HEADER}; -use crate::ocr::Error; -use crate::ocr::OcrClient; -use crate::ocr::codecs::document_intelligence::{ - self, AzureDocumentIntelligenceOperation, DocumentIntelligenceParams, -}; -use crate::ocr::error::{OcrError, OcrRequestError, OcrResponseError}; -use crate::ocr::prepare::{credential_env, transform_request_body}; -use crate::ocr::registry::OcrProvider; -use crate::ocr::types::{LiteLLMOcrRequest, LiteLLMOcrResponse, OcrConnection, OcrResponseFormat}; -use crate::url_utils::ApiUrl; -use litellm_auth::{InputSource, Sourced}; -use litellm_auth_azure::AzureAuthInputs; - -mod polling; - -const AZURE_DI_API_KEY_ENV: &str = "AZURE_DOCUMENT_INTELLIGENCE_API_KEY"; -const AZURE_DI_ENDPOINT_ENV: &str = "AZURE_DOCUMENT_INTELLIGENCE_ENDPOINT"; - -#[derive(Clone, Debug)] -pub(crate) struct AzureDocumentIntelligenceAdapter; - -impl OcrAdapter for AzureDocumentIntelligenceAdapter { - type ProviderResponse = AzureDocumentIntelligenceOperation; - const PROVIDER: OcrProvider = OcrProvider::AzureAi; - - async fn prepare_request( - &self, - request: &LiteLLMOcrRequest, - client: &OcrClient, - ) -> Result { - let params = map_ocr_params(request)?; - let mut config = AzureAuthInputs::from_sourced_optional_params( - &request.optional_params, - &request.input_sources, - ) - .map_err(Error::from)?; - config.azure_ad_token_provider = request.azure_ad_token_provider.clone(); - let headers = validate_environment(&request.connection, &config, &credential_env).await?; - let endpoint = nonblank(request.connection.api_base.clone()) - .or_else(|| nonblank(credential_env(AZURE_DI_ENDPOINT_ENV))) - .ok_or_else(|| Error::Auth("Missing Azure Document Intelligence API Base - Set AZURE_DOCUMENT_INTELLIGENCE_ENDPOINT or pass api_base".into()))?; - let url = get_complete_url(&endpoint, &request.model, ¶ms)?; - let body = document_intelligence::transform_ocr_request(request.document.clone())?; - transform_request_body(client, request, &url, &headers, false, body, |_| Ok(())).await - } - - fn transform_ocr_response( - &self, - request: &LiteLLMOcrRequest, - response: Self::ProviderResponse, - ) -> Result { - document_intelligence::transform_ocr_response(&request.model, response) - } - - async fn read_response( - &self, - client: &OcrClient, - response: reqwest::Response, - url: &str, - headers: &[(String, String)], - request: &LiteLLMOcrRequest, - ) -> Result, OcrError> { - polling::read_operation_response( - client.polling_http(), - response, - url, - headers, - &request.connection, - request.response_format()? == OcrResponseFormat::Native, - &request.hooks, - ) - .await - } -} - -fn map_ocr_params( - request: &LiteLLMOcrRequest, -) -> Result { - let params = document_intelligence::decode_input_params( - request.optional_params.clone(), - "optional_params", - )?; - let crate::ocr::prepare::ParsedProviderParams { - known: params, - extra_params: _extra_params, - } = params; - document_intelligence::map_ocr_params(params) -} - -fn get_complete_url( - endpoint: &str, - model: &str, - params: &DocumentIntelligenceParams, -) -> Result { - let model = format!("{}:analyze", model_id(model)?); - ApiUrl::parse(endpoint) - .and_then(|url| url.complete_path(&["documentintelligence", "documentModels", &model])) - .map(|url| { - url.append_query_pairs( - [("api-version", AZURE_DI_API_VERSION)] - .into_iter() - .chain(params.pages.iter().map(|pages| ("pages", pages.as_str()))) - .chain( - params - .features - .iter() - .map(|features| ("features", features.as_str())), - ), - ) - .into_string() - }) - .map_err(|_| OcrRequestError::RequestField { - path: "api_base".into(), - }) - .map_err(OcrError::from) -} - -async fn validate_environment( - connection: &OcrConnection, - config: &AzureAuthInputs, - env_lookup: &(dyn Fn(&str) -> Option + Sync), -) -> Result, OcrError> { - if crate::http_utils::has_header(&connection.extra_headers, "authorization") - || crate::http_utils::has_header(&connection.extra_headers, AZURE_DI_SUBSCRIPTION_HEADER) - { - super::validate_destination(connection, connection.extra_headers_source)?; - return Ok(connection.extra_headers.clone()); - } - let key = nonblank(connection.api_key.clone()) - .map(|value| Sourced::new(value, connection.api_key_source)) - .or_else(|| { - nonblank(env_lookup(AZURE_DI_API_KEY_ENV)) - .map(|value| Sourced::new(value, InputSource::Environment)) - }); - if let Some(key) = key { - super::validate_destination(connection, key.source())?; - return Ok( - std::iter::once((AZURE_DI_SUBSCRIPTION_HEADER.into(), key.into_value())) - .chain(connection.extra_headers.clone()) - .collect(), - ); - } - let token = super::resolve_entra(config, env_lookup) - .await? - .ok_or(Error::MissingAzureDocumentIntelligenceCredentials)?; - super::validate_destination(connection, token.source())?; - Ok( - std::iter::once(("Authorization".into(), format!("Bearer {}", token.value()))) - .chain(connection.extra_headers.clone()) - .collect(), - ) -} - -fn model_id(model: &str) -> Result<&str, OcrRequestError> { - let model = model.rsplit('/').next().unwrap_or(model); - if matches!(model, "." | "..") { - return Err(OcrRequestError::DotModel); - } - Ok(model) -} - -fn nonblank(value: Option) -> Option { - value - .map(|value| value.trim().to_string()) - .filter(|value| !value.is_empty()) -} - -#[cfg(test)] -mod tests { - use super::*; - - #[tokio::test] - async fn request_endpoint_cannot_receive_environment_key() { - let connection = OcrConnection { - api_base: Some("https://request.example".into()), - api_base_source: InputSource::Request, - ..Default::default() - }; - - let error = validate_environment(&connection, &Default::default(), &|name| { - (name == AZURE_DI_API_KEY_ENV).then(|| "environment-key".into()) - }) - .await - .unwrap_err(); - - assert!( - error - .to_string() - .contains("request-controlled Azure endpoint") - ); - } - - #[tokio::test] - async fn request_endpoint_accepts_request_owned_key() { - let connection = OcrConnection { - api_key: Some("request-key".into()), - api_key_source: InputSource::Request, - api_base: Some("https://request.example".into()), - api_base_source: InputSource::Request, - ..Default::default() - }; - - let headers = validate_environment(&connection, &Default::default(), &|_| None) - .await - .unwrap(); - - assert_eq!( - headers[0], - (AZURE_DI_SUBSCRIPTION_HEADER.into(), "request-key".into()) - ); - } -} diff --git a/litellm-rust/crates/core/src/ocr/adapters/azure/document_intelligence/polling.rs b/litellm-rust/crates/core/src/ocr/adapters/azure/document_intelligence/polling.rs deleted file mode 100644 index 87378dccdb7..00000000000 --- a/litellm-rust/crates/core/src/ocr/adapters/azure/document_intelligence/polling.rs +++ /dev/null @@ -1,119 +0,0 @@ -use std::sync::Arc; -use std::time::Duration; - -use reqwest::Url; -use tokio::time::Instant; - -use crate::constants::{AZURE_DI_SUBSCRIPTION_HEADER, OCR_POLL_RETRY_SECS}; -use crate::ocr::client::read_json_response; -use crate::ocr::codecs::document_intelligence::{ - AzureDocumentIntelligenceOperation, OperationStatus, -}; -use crate::ocr::error::{OcrError, OcrPollingError, OcrResponseError}; -use crate::ocr::hooks::OcrHooks; -use crate::ocr::types::OcrConnection; -use crate::ocr::wire::DecodedOcrResponse; - -pub(super) async fn read_operation_response( - http_client: &reqwest::Client, - response: reqwest::Response, - original_url: &str, - headers: &[(String, String)], - connection: &OcrConnection, - native: bool, - hooks: &Arc, -) -> Result, OcrError> { - if response.status() != reqwest::StatusCode::ACCEPTED { - let bytes = - crate::ocr::client::read_response_bytes(response, connection.max_response_bytes) - .await?; - crate::ocr::handler::post_call(hooks, &bytes).await?; - return Ok(crate::ocr::wire::decode_response(&bytes, native)?); - } - let location = response - .headers() - .get("operation-location") - .and_then(|value| value.to_str().ok()) - .ok_or(OcrPollingError::PollLocation)? - .to_string(); - let original = Url::parse(original_url).map_err(|_| OcrPollingError::PollOrigin)?; - let operation = Url::parse(&location).map_err(|_| OcrPollingError::PollOrigin)?; - if original.origin() != operation.origin() - || !operation.username().is_empty() - || operation.password().is_some() - { - return Err(OcrPollingError::PollOrigin.into()); - } - let bytes = - crate::ocr::client::read_response_bytes(response, connection.max_response_bytes).await?; - crate::ocr::handler::post_call(hooks, &bytes).await?; - poll_operation(http_client, operation, headers, connection, native, hooks).await -} - -async fn poll_operation( - http_client: &reqwest::Client, - url: Url, - headers: &[(String, String)], - connection: &OcrConnection, - native: bool, - hooks: &Arc, -) -> Result, OcrError> { - let deadline = Instant::now() - .checked_add(connection.poll_timeout) - .ok_or(OcrPollingError::PollTimeout)?; - loop { - let remaining = deadline - .checked_duration_since(Instant::now()) - .filter(|remaining| !remaining.is_zero()) - .ok_or(OcrPollingError::PollTimeout)?; - let builder = http_client - .get(url.clone()) - .timeout(remaining.min(connection.timeout)); - let builder = crate::http_utils::with_headers( - builder, - headers, - crate::http_utils::HeaderPolicy::Only(&[AZURE_DI_SUBSCRIPTION_HEADER, "authorization"]), - ); - let response = tokio::time::timeout_at(deadline, crate::http_utils::http_request(builder)) - .await - .map_err(|_| OcrPollingError::PollTimeout)? - .map_err(crate::transport::Error::from)?; - let retry = response - .headers() - .get(reqwest::header::RETRY_AFTER) - .and_then(|value| value.to_str().ok()) - .and_then(|value| value.parse::().ok()) - .unwrap_or(OCR_POLL_RETRY_SECS) - .max(1); - let decoded = tokio::time::timeout_at( - deadline, - read_json_response::( - response, - native, - connection.max_response_bytes, - ), - ) - .await - .map_err(|_| OcrPollingError::PollTimeout)??; - match &decoded.data.status { - Some(OperationStatus::Succeeded) => { - crate::ocr::handler::post_call(hooks, decoded.text.as_bytes()).await?; - return Ok(decoded); - } - Some(OperationStatus::Running | OperationStatus::NotStarted) => { - tokio::time::timeout_at(deadline, tokio::time::sleep(Duration::from_secs(retry))) - .await - .map_err(|_| OcrPollingError::PollTimeout)?; - } - status => { - return Err(OcrResponseError::OperationStatus( - status - .as_ref() - .map(ToString::to_string) - .unwrap_or_else(|| "None".into()), - ) - .into()); - } - } - } -} diff --git a/litellm-rust/crates/core/src/ocr/adapters/azure/mistral.rs b/litellm-rust/crates/core/src/ocr/adapters/azure/mistral.rs deleted file mode 100644 index 28e09cdc80f..00000000000 --- a/litellm-rust/crates/core/src/ocr/adapters/azure/mistral.rs +++ /dev/null @@ -1,229 +0,0 @@ -use super::super::OcrAdapter; -use crate::constants::AZURE_AI_OCR_PATH; -use crate::ocr::Error; -use crate::ocr::OcrClient; -use crate::ocr::codecs::mistral::{self, MistralOcrParams, MistralOcrResponse}; -use crate::ocr::document::{inline_remote_document, validate_inline_document}; -use crate::ocr::error::{OcrError, OcrRequestError, OcrResponseError}; -use crate::ocr::prepare::{ - _prepare_ocr_request, ParsedProviderParams, credential_env, transform_request_body, -}; -use crate::ocr::registry::OcrProvider; -use crate::ocr::types::{LiteLLMOcrRequest, LiteLLMOcrResponse, OcrConnection}; -use crate::url_utils::ApiUrl; -use litellm_auth::{InputSource, Sourced}; -use litellm_auth_azure::AzureAuthInputs; - -const AZURE_AI_API_KEY_ENV: &str = "AZURE_AI_API_KEY"; -const AZURE_AI_API_BASE_ENV: &str = "AZURE_AI_API_BASE"; - -#[derive(Clone, Debug)] -pub(crate) struct AzureMistralAdapter; - -impl OcrAdapter for AzureMistralAdapter { - type ProviderResponse = MistralOcrResponse; - const PROVIDER: OcrProvider = OcrProvider::AzureAi; - - async fn prepare_request( - &self, - request: &LiteLLMOcrRequest, - client: &OcrClient, - ) -> Result { - let ParsedProviderParams { - known: params, - extra_params: _extra_params, - } = _prepare_ocr_request::(request)?; - let mut config = AzureAuthInputs::from_sourced_optional_params( - &request.optional_params, - &request.input_sources, - ) - .map_err(Error::from)?; - config.azure_ad_token_provider = request.azure_ad_token_provider.clone(); - let url = get_complete_url(request.connection.api_base.as_deref(), &credential_env)?; - let headers = validate_environment(&request.connection, &config, &credential_env).await?; - let retains_document = !request.document.source().starts_with("http://") - && !request.document.source().starts_with("https://"); - let document = inline_remote_document( - client.document_fetcher(), - request.document.clone(), - &request.connection, - ) - .await?; - let body = mistral::transform_ocr_request(&request.model, document, ¶ms)?; - transform_request_body( - client, - request, - &url, - &headers, - retains_document, - body, - |body| validate_inline_document(&body.document), - ) - .await - } - - fn transform_ocr_response( - &self, - request: &LiteLLMOcrRequest, - response: Self::ProviderResponse, - ) -> Result { - mistral::transform_ocr_response(&request.model, response) - } -} - -fn get_complete_url( - api_base: Option<&str>, - env_lookup: &dyn Fn(&str) -> Option, -) -> Result { - let base = nonblank(api_base.map(str::to_string)) - .or_else(|| nonblank(env_lookup(AZURE_AI_API_BASE_ENV))) - .ok_or_else(|| Error::Auth( - "Missing Azure AI API Base - Set AZURE_AI_API_BASE environment variable or pass api_base parameter".into(), - ))?; - let path: Vec<&str> = AZURE_AI_OCR_PATH.trim_matches('/').split('/').collect(); - ApiUrl::parse(&base) - .and_then(|url| url.complete_path(&path)) - .map(|url| url.into_string()) - .map_err(|_| { - OcrRequestError::RequestField { - path: "api_base".into(), - } - .into() - }) -} - -pub(in crate::ocr::adapters) async fn validate_environment( - connection: &OcrConnection, - config: &AzureAuthInputs, - env_lookup: &(dyn Fn(&str) -> Option + Sync), -) -> Result, OcrError> { - if crate::http_utils::has_header(&connection.extra_headers, "authorization") { - if config.azure_ad_token_provider.is_some() { - super::resolve_entra(config, env_lookup).await?; - } - super::validate_destination(connection, connection.extra_headers_source)?; - return Ok(connection.extra_headers.clone()); - } - let key = nonblank(connection.api_key.clone()) - .map(|value| Sourced::new(value, connection.api_key_source)) - .or_else(|| { - nonblank(env_lookup(AZURE_AI_API_KEY_ENV)) - .map(|value| Sourced::new(value, InputSource::Environment)) - }); - if let Some(key) = key { - super::validate_destination(connection, key.source())?; - return Ok(bearer_headers(connection, key.value())); - } - let key = super::resolve_entra(config, env_lookup) - .await? - .ok_or(Error::MissingAzureAiCredentials)?; - super::validate_destination(connection, key.source())?; - Ok(bearer_headers(connection, key.value())) -} - -fn bearer_headers(connection: &OcrConnection, key: &str) -> Vec<(String, String)> { - std::iter::once(("Authorization".into(), format!("Bearer {key}"))) - .chain(connection.extra_headers.clone()) - .collect() -} - -fn nonblank(value: Option) -> Option { - value - .map(|value| value.trim().to_string()) - .filter(|value| !value.is_empty()) -} - -#[cfg(test)] -mod tests { - use super::*; - - #[test] - fn completes_azure_path_and_preserves_query() { - assert_eq!( - get_complete_url(Some("https://example.com/?tenant=a"), &|_| None).unwrap(), - "https://example.com/providers/mistral/azure/ocr?tenant=a" - ); - assert_eq!( - get_complete_url( - Some("https://example.com/providers/mistral/azure/ocr"), - &|_| None - ) - .unwrap(), - "https://example.com/providers/mistral/azure/ocr" - ); - } - - #[tokio::test] - async fn supplied_authorization_precedes_keys() { - let connection = OcrConnection { - api_key: Some("request-key".into()), - extra_headers: vec![("authorization".into(), "Bearer prepared".into())], - ..Default::default() - }; - assert_eq!( - validate_environment(&connection, &Default::default(), &|_| { - Some("environment-key".into()) - }) - .await - .unwrap(), - connection.extra_headers - ); - } - - #[tokio::test] - async fn request_key_precedes_environment_key() { - let connection = OcrConnection { - api_key: Some("request-key".into()), - ..Default::default() - }; - assert_eq!( - validate_environment(&connection, &Default::default(), &|_| { - Some("environment-key".into()) - }) - .await - .unwrap()[0], - ("Authorization".into(), "Bearer request-key".into()) - ); - } - - #[tokio::test] - async fn request_endpoint_cannot_receive_environment_key() { - let connection = OcrConnection { - api_base: Some("https://request.example".into()), - api_base_source: InputSource::Request, - ..Default::default() - }; - - let error = validate_environment(&connection, &Default::default(), &|name| { - (name == AZURE_AI_API_KEY_ENV).then(|| "environment-key".into()) - }) - .await - .unwrap_err(); - - assert!( - error - .to_string() - .contains("request-controlled Azure endpoint") - ); - } - - #[tokio::test] - async fn request_endpoint_accepts_request_owned_key() { - let connection = OcrConnection { - api_key: Some("request-key".into()), - api_key_source: InputSource::Request, - api_base: Some("https://request.example".into()), - api_base_source: InputSource::Request, - ..Default::default() - }; - - let headers = validate_environment(&connection, &Default::default(), &|_| None) - .await - .unwrap(); - - assert_eq!( - headers[0], - ("Authorization".into(), "Bearer request-key".into()) - ); - } -} diff --git a/litellm-rust/crates/core/src/ocr/adapters/cohere.rs b/litellm-rust/crates/core/src/ocr/adapters/cohere.rs deleted file mode 100644 index d1faeeb7b1d..00000000000 --- a/litellm-rust/crates/core/src/ocr/adapters/cohere.rs +++ /dev/null @@ -1,123 +0,0 @@ -use super::OcrAdapter; -use crate::constants::{COHERE_API_KEY_ENV, COHERE_PARSE_API_BASE}; -use crate::ocr::Error; -use crate::ocr::OcrClient; -use crate::ocr::codecs::cohere::{ - CohereParams, CohereResponse, transform_request, transform_response, validate_document, -}; -use crate::ocr::error::{OcrError, OcrRequestError, OcrResponseError}; -use crate::ocr::prepare::{credential_env, transform_request_body}; -use crate::ocr::registry::OcrProvider; -use crate::ocr::types::{LiteLLMOcrRequest, LiteLLMOcrResponse, OcrConnection}; -use crate::url_utils::ApiUrl; - -pub(crate) struct CohereAdapter; - -impl OcrAdapter for CohereAdapter { - type ProviderResponse = CohereResponse; - const PROVIDER: OcrProvider = OcrProvider::Cohere; - - async fn prepare_request( - &self, - request: &LiteLLMOcrRequest, - client: &OcrClient, - ) -> Result { - let params = super::super::wire::decode_request_value::( - serde_json::Value::Object(request.optional_params.clone()), - "optional_params", - )?; - let headers = validate_environment(&request.connection, &credential_env)?; - let url = complete_url( - request - .connection - .api_base - .as_deref() - .unwrap_or(COHERE_PARSE_API_BASE), - )?; - let body = transform_request(&request.model, request.document.clone(), params)?; - transform_request_body(client, request, &url, &headers, true, body, |body| { - validate_document(&body.document) - }) - .await - } - - fn transform_ocr_response( - &self, - request: &LiteLLMOcrRequest, - response: Self::ProviderResponse, - ) -> Result { - transform_response(&request.model, response) - } -} - -fn complete_url(base: &str) -> Result { - let parsed = reqwest::Url::parse(base).map_err(|_| invalid_api_base())?; - if !matches!(parsed.scheme(), "http" | "https") { - return Err(invalid_api_base().into()); - } - ApiUrl::parse(base) - .and_then(|url| url.complete_path(&["v2", "parse"])) - .map(|url| url.into_string()) - .map_err(|_| invalid_api_base().into()) -} - -fn invalid_api_base() -> OcrRequestError { - OcrRequestError::RequestField { - path: "api_base".into(), - } -} - -fn validate_environment( - connection: &OcrConnection, - env_lookup: &(dyn Fn(&str) -> Option + Sync), -) -> Result, OcrError> { - if crate::http_utils::has_header(&connection.extra_headers, "authorization") { - return Ok(connection.extra_headers.clone()); - } - let key = connection - .api_key - .as_deref() - .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())) - .ok_or_else(|| { - Error::Auth("Missing COHERE_API_KEY - set it in the environment or pass api_key".into()) - })?; - Ok( - std::iter::once(("Authorization".into(), format!("Bearer {key}"))) - .chain(connection.extra_headers.clone()) - .collect(), - ) -} - -#[cfg(test)] -mod tests { - use super::*; - - #[test] - fn completes_provider_urls_without_duplicate_paths_and_preserves_queries() { - for suffix in ["", "/v2", "/v2/parse"] { - assert_eq!( - complete_url(&format!("https://example.com{suffix}?tenant=a")).unwrap(), - "https://example.com/v2/parse?tenant=a" - ); - } - } - - #[test] - fn rejects_invalid_urls_and_blank_keys() { - assert!(complete_url("relative/path").is_err()); - assert!(complete_url("ftp://example.com").is_err()); - assert!(matches!( - validate_environment( - &OcrConnection { - api_key: Some(" ".into()), - ..Default::default() - }, - &|_| None, - ), - Err(OcrError::Public(Error::Auth(_))) - )); - } -} diff --git a/litellm-rust/crates/core/src/ocr/adapters/mistral.rs b/litellm-rust/crates/core/src/ocr/adapters/mistral.rs deleted file mode 100644 index c379462c089..00000000000 --- a/litellm-rust/crates/core/src/ocr/adapters/mistral.rs +++ /dev/null @@ -1,147 +0,0 @@ -use super::OcrAdapter; -use crate::constants::MISTRAL_OCR_API_BASE; -use crate::ocr::Error; -use crate::ocr::OcrClient; -use crate::ocr::codecs::mistral::{self, MistralOcrParams, MistralOcrResponse}; -use crate::ocr::error::{OcrError, OcrRequestError, OcrResponseError}; -use crate::ocr::prepare::{ - _prepare_ocr_request, ParsedProviderParams, credential_env, transform_request_body, -}; -use crate::ocr::registry::OcrProvider; -use crate::ocr::types::{LiteLLMOcrRequest, LiteLLMOcrResponse, OcrConnection}; -use crate::url_utils::ApiUrl; - -const MISTRAL_API_KEY_ENV: &str = "MISTRAL_API_KEY"; - -#[derive(Clone, Debug)] -pub(crate) struct MistralAdapter; - -impl OcrAdapter for MistralAdapter { - type ProviderResponse = MistralOcrResponse; - const PROVIDER: OcrProvider = OcrProvider::Mistral; - - async fn prepare_request( - &self, - request: &LiteLLMOcrRequest, - client: &OcrClient, - ) -> Result { - let ParsedProviderParams { - known: params, - extra_params: _extra_params, - } = _prepare_ocr_request::(request)?; - let headers = validate_environment(&request.connection, &credential_env)?; - let url = get_complete_url(request.connection.api_base.as_deref())?; - let body = - mistral::transform_ocr_request(&request.model, request.document.clone(), ¶ms)?; - transform_request_body(client, request, &url, &headers, true, body, |_| Ok(())).await - } - - fn transform_ocr_response( - &self, - request: &LiteLLMOcrRequest, - response: Self::ProviderResponse, - ) -> Result { - mistral::transform_ocr_response(&request.model, response) - } -} - -pub(crate) fn get_complete_url(api_base: Option<&str>) -> Result { - let base = api_base - .map(str::trim) - .filter(|base| !base.is_empty()) - .unwrap_or(MISTRAL_OCR_API_BASE); - ApiUrl::parse(base) - .and_then(|url| url.complete_path(&["v1", "ocr"])) - .map(|url| url.into_string()) - .map_err(|_| { - OcrRequestError::RequestField { - path: "api_base".into(), - } - .into() - }) -} - -fn validate_environment( - connection: &OcrConnection, - env_lookup: &(dyn Fn(&str) -> Option + Sync), -) -> Result, OcrError> { - if crate::http_utils::has_header(&connection.extra_headers, "authorization") { - return Ok(connection.extra_headers.clone()); - } - let api_key = connection - .api_key - .as_deref() - .map(str::trim) - .filter(|key| !key.is_empty()) - .map(str::to_string) - .or_else(|| env_lookup(MISTRAL_API_KEY_ENV).filter(|key| !key.trim().is_empty())) - .ok_or(Error::MissingApiKey { - provider: "Mistral", - })?; - Ok( - std::iter::once(("Authorization".into(), format!("Bearer {api_key}"))) - .chain(connection.extra_headers.clone()) - .collect(), - ) -} - -#[cfg(test)] -mod tests { - use super::*; - - #[test] - fn complete_url_defaults_and_dedupes_v1() { - assert_eq!( - get_complete_url(None).unwrap(), - "https://api.mistral.ai/v1/ocr" - ); - assert_eq!( - get_complete_url(Some("https://example.com/v1?tenant=a")).unwrap(), - "https://example.com/v1/ocr?tenant=a" - ); - assert_eq!( - get_complete_url(Some("https://example.com/v1/ocr?tenant=a")).unwrap(), - "https://example.com/v1/ocr?tenant=a" - ); - } - - #[test] - fn environment_prefers_explicit_key_then_environment() { - let explicit = OcrConnection { - api_key: Some("explicit".into()), - ..OcrConnection::default() - }; - assert_eq!( - validate_environment(&explicit, &|_| Some("environment".into())).unwrap()[0], - ("Authorization".into(), "Bearer explicit".into()) - ); - - assert_eq!( - validate_environment(&OcrConnection::default(), &|_| Some("environment".into())) - .unwrap()[0], - ("Authorization".into(), "Bearer environment".into()) - ); - } - - #[test] - fn environment_preserves_forwarded_authorization() { - let connection = OcrConnection { - extra_headers: vec![("authorization".into(), "Bearer forwarded".into())], - ..OcrConnection::default() - }; - assert_eq!( - validate_environment(&connection, &|_| None).unwrap(), - connection.extra_headers - ); - } - - #[test] - fn environment_rejects_missing_key() { - assert!(matches!( - validate_environment(&OcrConnection::default(), &|_| None), - Err(OcrError::Public(Error::MissingApiKey { - provider: "Mistral" - })) - )); - } -} diff --git a/litellm-rust/crates/core/src/ocr/adapters/mod.rs b/litellm-rust/crates/core/src/ocr/adapters/mod.rs deleted file mode 100644 index d473fcad280..00000000000 --- a/litellm-rust/crates/core/src/ocr/adapters/mod.rs +++ /dev/null @@ -1,91 +0,0 @@ -use std::future::Future; - -use serde::de::DeserializeOwned; - -use super::OcrClient; -use super::error::{OcrError, OcrResponseError}; -use super::registry::OcrProvider; -use super::types::{LiteLLMOcrRequest, LiteLLMOcrResponse}; - -mod azure; -mod cohere; -mod mistral; -mod reducto; -mod vertex; - -pub(crate) use azure::{AzureCohereAdapter, AzureDocumentIntelligenceAdapter, AzureMistralAdapter}; -pub(crate) use cohere::CohereAdapter; -pub(crate) use mistral::MistralAdapter; -pub(crate) use reducto::{ReductoLegacyAdapter, ReductoV3Adapter}; -pub(crate) use vertex::{VertexDeepSeekAdapter, VertexMistralAdapter}; - -/// Converts a complete LiteLLM OCR call to provider HTTP and normalizes its response. -pub(crate) trait OcrAdapter: Send + Sync + Sized + 'static { - /// Provider JSON schema; direct and Vertex Mistral share `MistralOcrResponse`. - type ProviderResponse: DeserializeOwned + Send; - - const PROVIDER: OcrProvider; - - /// Prepares the complete provider HTTP request. - /// `request` contains the model, document, connection, and unmapped caller options. - /// `client` supplies reusable provider and document HTTP clients. - /// Returns the complete HTTP request, whereas Python returns body data. - fn prepare_request( - &self, - request: &LiteLLMOcrRequest, - client: &OcrClient, - ) -> impl Future> + Send; - - /// Python: `transform_ocr_response`. - /// `request` supplies caller context, including the fallback model. - /// `response` is the decoded provider payload; the output is the shared LiteLLM schema. - fn transform_ocr_response( - &self, - request: &LiteLLMOcrRequest, - response: Self::ProviderResponse, - ) -> Result; - - /// Decodes provider HTTP; adapters may override this to poll asynchronous operations. - /// Python performs that polling inside `async_transform_ocr_response`. - /// `client` is reused for polling; `response` is the initial HTTP response. - /// `url` and `headers` describe the submitted call; `request` supplies limits and format. - fn read_response( - &self, - _client: &OcrClient, - response: reqwest::Response, - _url: &str, - _headers: &[(String, String)], - request: &LiteLLMOcrRequest, - ) -> impl Future< - Output = Result, OcrError>, - > + Send { - async move { - let bytes = - super::client::read_response_bytes(response, request.connection.max_response_bytes) - .await?; - super::handler::post_call(&request.hooks, &bytes).await?; - Ok(super::wire::decode_response( - &bytes, - request.response_format()? == super::types::OcrResponseFormat::Native, - )?) - } - } -} - -macro_rules! for_each_ocr_adapter { - ($callback:ident) => { - $callback! { - Cohere, $crate::ocr::adapters::CohereAdapter, $crate::ocr::adapters::CohereAdapter, Cohere; - AzureCohere, $crate::ocr::adapters::AzureCohereAdapter, $crate::ocr::adapters::AzureCohereAdapter, AzureAi; - Mistral, $crate::ocr::adapters::MistralAdapter, $crate::ocr::adapters::MistralAdapter, Mistral; - AzureMistral, $crate::ocr::adapters::AzureMistralAdapter, $crate::ocr::adapters::AzureMistralAdapter, AzureAi; - AzureDocumentIntelligence, $crate::ocr::adapters::AzureDocumentIntelligenceAdapter, $crate::ocr::adapters::AzureDocumentIntelligenceAdapter, AzureAi; - ReductoLegacy, $crate::ocr::adapters::ReductoLegacyAdapter, $crate::ocr::adapters::ReductoLegacyAdapter, Reducto; - ReductoV3, $crate::ocr::adapters::ReductoV3Adapter, $crate::ocr::adapters::ReductoV3Adapter, Reducto; - VertexMistral, $crate::ocr::adapters::VertexMistralAdapter, $crate::ocr::adapters::VertexMistralAdapter, VertexAi; - VertexDeepSeek, $crate::ocr::adapters::VertexDeepSeekAdapter, $crate::ocr::adapters::VertexDeepSeekAdapter, VertexAi; - } - }; -} - -pub(crate) use for_each_ocr_adapter; diff --git a/litellm-rust/crates/core/src/ocr/adapters/reducto/legacy.rs b/litellm-rust/crates/core/src/ocr/adapters/reducto/legacy.rs deleted file mode 100644 index 8889bcd1b45..00000000000 --- a/litellm-rust/crates/core/src/ocr/adapters/reducto/legacy.rs +++ /dev/null @@ -1,45 +0,0 @@ -use super::super::OcrAdapter; -use crate::ocr::OcrClient; -use crate::ocr::codecs::reducto::{self, ReductoLegacyParams, ReductoResponse}; -use crate::ocr::error::{OcrError, OcrResponseError}; -use crate::ocr::prepare::{ - _prepare_ocr_request, ParsedProviderParams, build_http_request, credential_env, - guardrail_document, merge_extra_params, -}; -use crate::ocr::registry::OcrProvider; -use crate::ocr::types::{LiteLLMOcrRequest, LiteLLMOcrResponse}; - -#[derive(Clone, Debug)] -pub(crate) struct ReductoLegacyAdapter; - -impl OcrAdapter for ReductoLegacyAdapter { - type ProviderResponse = ReductoResponse; - const PROVIDER: OcrProvider = OcrProvider::Reducto; - - async fn prepare_request( - &self, - request: &LiteLLMOcrRequest, - client: &OcrClient, - ) -> Result { - let ParsedProviderParams { - known: params, - extra_params, - } = _prepare_ocr_request::(request)?; - let headers = super::validate_environment(&request.connection, &credential_env)?; - let url = super::get_complete_url(request.connection.api_base.as_deref(), "parse")?; - let (document, headers) = guardrail_document(request, &url, &headers).await?; - let document = - super::prepare_document(client, document, &request.connection, &headers).await?; - let body = reducto::transform_legacy_ocr_request(&request.model, document, ¶ms)?; - let body = merge_extra_params(&body, extra_params)?; - build_http_request(client, request, &url, &headers, &body) - } - - fn transform_ocr_response( - &self, - request: &LiteLLMOcrRequest, - response: Self::ProviderResponse, - ) -> Result { - reducto::transform_ocr_response(&request.model, response) - } -} diff --git a/litellm-rust/crates/core/src/ocr/adapters/reducto/mod.rs b/litellm-rust/crates/core/src/ocr/adapters/reducto/mod.rs deleted file mode 100644 index 40cefa05373..00000000000 --- a/litellm-rust/crates/core/src/ocr/adapters/reducto/mod.rs +++ /dev/null @@ -1,148 +0,0 @@ -mod legacy; -mod v3; - -use crate::constants::{REDUCTO_API_BASE, REDUCTO_API_KEY_ENV, REDUCTO_ID_PREFIX}; -use crate::ocr::Error; -use crate::ocr::document::InlineDocument; -use crate::ocr::error::{OcrError, OcrRequestError, OcrResponseError}; -use crate::ocr::types::{OcrConnection, OcrDocument}; -use crate::url_utils::ApiUrl; - -pub(crate) use legacy::ReductoLegacyAdapter; -pub(crate) use v3::ReductoV3Adapter; - -pub(super) fn get_complete_url(api_base: Option<&str>, path: &str) -> Result { - let base = api_base - .map(str::trim) - .filter(|base| !base.is_empty()) - .unwrap_or(REDUCTO_API_BASE); - ApiUrl::parse(base) - .and_then(|url| url.complete_path(&[path])) - .map(|url| url.into_string()) - .map_err(|_| { - OcrRequestError::RequestField { - path: "api_base".into(), - } - .into() - }) -} - -pub(super) fn validate_environment( - connection: &OcrConnection, - env_lookup: &(dyn Fn(&str) -> Option + Sync), -) -> Result, OcrError> { - if crate::http_utils::has_header(&connection.extra_headers, "authorization") { - return Ok(connection.extra_headers.clone()); - } - let api_key = connection - .api_key - .as_deref() - .map(str::trim) - .filter(|key| !key.is_empty()) - .map(str::to_string) - .or_else(|| { - env_lookup(REDUCTO_API_KEY_ENV) - .map(|key| key.trim().to_string()) - .filter(|key| !key.is_empty()) - }) - .ok_or(Error::MissingReductoApiKey)?; - Ok( - std::iter::once(("Authorization".into(), format!("Bearer {api_key}"))) - .chain(connection.extra_headers.clone()) - .collect(), - ) -} - -pub(super) async fn prepare_document( - client: &crate::ocr::OcrClient, - document: OcrDocument, - connection: &OcrConnection, - headers: &[(String, String)], -) -> Result { - if document.source().starts_with(REDUCTO_ID_PREFIX) { - if document.source()[REDUCTO_ID_PREFIX.len()..] - .trim() - .is_empty() - { - return Err(OcrRequestError::RequestField { - path: "document file id".into(), - } - .into()); - } - return Ok(document); - } - let inline = InlineDocument::parse(document.source())?.ok_or(OcrRequestError::ReductoSource)?; - let mime = inline.mime_type().to_string(); - let bytes = inline.decode(crate::constants::OCR_INLINE_MAX_BYTES)?; - let part = reqwest::multipart::Part::bytes(bytes) - .file_name("document") - .mime_str(&mime) - .map_err(|_| OcrRequestError::InvalidDataUri)?; - let builder = client - .provider_http() - .post(get_complete_url(connection.api_base.as_deref(), "upload")?) - .multipart(reqwest::multipart::Form::new().part("file", part)) - .timeout(connection.timeout); - let builder = crate::http_utils::with_headers( - builder, - headers, - crate::http_utils::HeaderPolicy::Except(&["content-type", "content-length"]), - ); - let response = crate::http_utils::http_request(builder) - .await - .map_err(crate::transport::Error::from)?; - let uploaded = crate::ocr::client::read_json_response::< - crate::ocr::codecs::reducto::ReductoUploadResponse, - >(response, false, connection.max_response_bytes) - .await? - .data; - let file_id = uploaded - .file_id - .as_deref() - .map(str::trim) - .filter(|id| !id.is_empty()); - let Some(file_id) = file_id else { - return Err(OcrResponseError::ResponseField { - path: "file_id".into(), - } - .into()); - }; - Ok(document.with_source(file_id.to_string())) -} - -#[cfg(test)] -mod tests { - use super::*; - - #[test] - fn explicit_key_precedes_environment_key() { - let connection = OcrConnection { - api_key: Some("passed-key".into()), - ..Default::default() - }; - let headers = validate_environment(&connection, &|_| Some("env-key".into())).unwrap(); - assert_eq!(headers[0].1, "Bearer passed-key"); - } - - #[test] - fn blank_explicit_key_uses_environment_key() { - let connection = OcrConnection { - api_key: Some(" ".into()), - ..Default::default() - }; - let headers = validate_environment(&connection, &|_| Some(" env-key ".into())).unwrap(); - assert_eq!(headers[0].1, "Bearer env-key"); - } - - #[test] - fn existing_authorization_skips_key_lookup() { - let connection = OcrConnection { - extra_headers: vec![("authorization".into(), "Bearer existing".into())], - ..Default::default() - }; - assert_eq!( - validate_environment(&connection, &|_| None).unwrap(), - connection.extra_headers - ); - } -} diff --git a/litellm-rust/crates/core/src/ocr/adapters/reducto/v3.rs b/litellm-rust/crates/core/src/ocr/adapters/reducto/v3.rs deleted file mode 100644 index c272d31b67e..00000000000 --- a/litellm-rust/crates/core/src/ocr/adapters/reducto/v3.rs +++ /dev/null @@ -1,45 +0,0 @@ -use super::super::OcrAdapter; -use crate::ocr::OcrClient; -use crate::ocr::codecs::reducto::{self, ReductoResponse, ReductoV3Params}; -use crate::ocr::error::{OcrError, OcrResponseError}; -use crate::ocr::prepare::{ - _prepare_ocr_request, ParsedProviderParams, build_http_request, credential_env, - guardrail_document, merge_extra_params, -}; -use crate::ocr::registry::OcrProvider; -use crate::ocr::types::{LiteLLMOcrRequest, LiteLLMOcrResponse}; - -#[derive(Clone, Debug)] -pub(crate) struct ReductoV3Adapter; - -impl OcrAdapter for ReductoV3Adapter { - type ProviderResponse = ReductoResponse; - const PROVIDER: OcrProvider = OcrProvider::Reducto; - - async fn prepare_request( - &self, - request: &LiteLLMOcrRequest, - client: &OcrClient, - ) -> Result { - let ParsedProviderParams { - known: params, - extra_params, - } = _prepare_ocr_request::(request)?; - let headers = super::validate_environment(&request.connection, &credential_env)?; - let url = super::get_complete_url(request.connection.api_base.as_deref(), "parse")?; - let (document, headers) = guardrail_document(request, &url, &headers).await?; - let document = - super::prepare_document(client, document, &request.connection, &headers).await?; - let body = reducto::transform_v3_ocr_request(&request.model, document, ¶ms)?; - let body = merge_extra_params(&body, extra_params)?; - build_http_request(client, request, &url, &headers, &body) - } - - fn transform_ocr_response( - &self, - request: &LiteLLMOcrRequest, - response: Self::ProviderResponse, - ) -> Result { - reducto::transform_ocr_response(&request.model, response) - } -} diff --git a/litellm-rust/crates/core/src/ocr/adapters/vertex/deepseek.rs b/litellm-rust/crates/core/src/ocr/adapters/vertex/deepseek.rs deleted file mode 100644 index fc24dbe489c..00000000000 --- a/litellm-rust/crates/core/src/ocr/adapters/vertex/deepseek.rs +++ /dev/null @@ -1,140 +0,0 @@ -use super::super::OcrAdapter; -use super::validate_destination; -use crate::ocr::Error; -use crate::ocr::OcrClient; -use crate::ocr::codecs::deepseek::{self, DeepSeekOcrParams, DeepSeekOcrResponse}; -use crate::ocr::error::{OcrError, OcrRequestError, OcrResponseError}; -use crate::ocr::prepare::{ - _prepare_ocr_request, ParsedProviderParams, credential_env, transform_request_body, -}; -use crate::ocr::registry::OcrProvider; -use crate::ocr::types::{LiteLLMOcrRequest, LiteLLMOcrResponse}; -use crate::url_utils::ApiUrl; -use litellm_auth_gcp::{self as vertex, VertexConfig}; -const DEFAULT_API_BASE: &str = "https://aiplatform.googleapis.com"; -const MODEL_NAMESPACE: &str = "deepseek-ai"; -const DEFAULT_LOCATION: &str = "us-central1"; - -#[derive(Clone, Debug)] -pub(crate) struct VertexDeepSeekAdapter; - -impl OcrAdapter for VertexDeepSeekAdapter { - type ProviderResponse = DeepSeekOcrResponse; - const PROVIDER: OcrProvider = OcrProvider::VertexAi; - - async fn prepare_request( - &self, - request: &LiteLLMOcrRequest, - client: &OcrClient, - ) -> Result { - validate_destination(&request.connection)?; - let ParsedProviderParams { - known: params, - extra_params: _extra_params, - } = _prepare_ocr_request::(request)?; - let config = VertexConfig::from_sourced_optional_params( - &request.optional_params, - &request.input_sources, - ) - .map_err(Error::from)?; - let authentication = client - .vertex_auth() - .validate_environment( - request.connection.extra_headers.clone(), - request.connection.api_key.as_deref(), - &config, - &credential_env, - ) - .await - .map_err(Error::from)?; - let location = vertex::get_vertex_ai_location(&config, &credential_env) - .unwrap_or_else(|| DEFAULT_LOCATION.to_string()); - let url = get_complete_url( - request.connection.api_base.as_deref(), - &authentication.project_id, - &location, - )?; - let document = request.document.clone(); - let body = - deepseek::transform_ocr_request(&provider_model(&request.model), document, ¶ms)?; - transform_request_body( - client, - request, - &url, - &authentication.headers, - false, - body, - |_| Ok(()), - ) - .await - } - - fn transform_ocr_response( - &self, - request: &LiteLLMOcrRequest, - response: Self::ProviderResponse, - ) -> Result { - deepseek::transform_ocr_response(&request.model, response) - } -} - -fn provider_model(model: &str) -> String { - if model.starts_with(&format!("{MODEL_NAMESPACE}/")) { - model.to_string() - } else { - format!("{MODEL_NAMESPACE}/{model}") - } -} - -fn get_complete_url( - api_base: Option<&str>, - project: &str, - location: &str, -) -> Result { - let base = api_base - .map(str::trim) - .filter(|base| !base.is_empty()) - .unwrap_or(DEFAULT_API_BASE); - ApiUrl::parse(base) - .and_then(|url| { - url.complete_path(&[ - "v1", - "projects", - project, - "locations", - location, - "endpoints", - "openapi", - "chat", - "completions", - ]) - }) - .map(|url| url.into_string()) - .map_err(|_| { - OcrRequestError::RequestField { - path: "api_base".into(), - } - .into() - }) -} - -#[cfg(test)] -mod tests { - use super::{get_complete_url, provider_model}; - - #[test] - fn adapter_owns_model_namespace_and_endpoint() { - assert_eq!( - provider_model("deepseek-ocr-maas"), - "deepseek-ai/deepseek-ocr-maas" - ); - assert_eq!( - provider_model("deepseek-ai/deepseek-ocr-maas"), - "deepseek-ai/deepseek-ocr-maas" - ); - assert_eq!( - get_complete_url(None, "proj-1", "europe-west4").unwrap(), - "https://aiplatform.googleapis.com/v1/projects/proj-1/locations/europe-west4/endpoints/openapi/chat/completions" - ); - } -} diff --git a/litellm-rust/crates/core/src/ocr/adapters/vertex/mistral.rs b/litellm-rust/crates/core/src/ocr/adapters/vertex/mistral.rs deleted file mode 100644 index 3a1abf47ddf..00000000000 --- a/litellm-rust/crates/core/src/ocr/adapters/vertex/mistral.rs +++ /dev/null @@ -1,157 +0,0 @@ -use super::super::OcrAdapter; -use super::validate_destination; -use crate::ocr::Error; -use crate::ocr::OcrClient; -use crate::ocr::codecs::mistral::{self, MistralOcrParams, MistralOcrResponse}; -use crate::ocr::document::{inline_remote_document, validate_inline_document}; -use crate::ocr::error::{OcrError, OcrRequestError, OcrResponseError}; -use crate::ocr::prepare::{ - _prepare_ocr_request, ParsedProviderParams, credential_env, transform_request_body, -}; -use crate::ocr::registry::OcrProvider; -use crate::ocr::types::{LiteLLMOcrRequest, LiteLLMOcrResponse}; -use crate::url_utils::ApiUrl; -use litellm_auth_gcp::{self as vertex, VertexConfig}; -const DEFAULT_LOCATION: &str = "us-central1"; - -#[derive(Clone, Debug)] -pub(crate) struct VertexMistralAdapter; - -impl OcrAdapter for VertexMistralAdapter { - type ProviderResponse = MistralOcrResponse; - const PROVIDER: OcrProvider = OcrProvider::VertexAi; - - async fn prepare_request( - &self, - request: &LiteLLMOcrRequest, - client: &OcrClient, - ) -> Result { - validate_destination(&request.connection)?; - let ParsedProviderParams { - known: params, - extra_params: _extra_params, - } = _prepare_ocr_request::(request)?; - let config = VertexConfig::from_sourced_optional_params( - &request.optional_params, - &request.input_sources, - ) - .map_err(Error::from)?; - let authentication = client - .vertex_auth() - .validate_environment( - request.connection.extra_headers.clone(), - request.connection.api_key.as_deref(), - &config, - &credential_env, - ) - .await - .map_err(Error::from)?; - let location = vertex::get_vertex_ai_location(&config, &credential_env) - .unwrap_or_else(|| DEFAULT_LOCATION.to_string()); - let url = get_complete_url( - request.connection.api_base.as_deref(), - &authentication.project_id, - &location, - &request.model, - )?; - let retains_document = !request.document.source().starts_with("http://") - && !request.document.source().starts_with("https://"); - let document = inline_remote_document( - client.document_fetcher(), - request.document.clone(), - &request.connection, - ) - .await?; - let body = mistral::transform_ocr_request(&request.model, document, ¶ms)?; - transform_request_body( - client, - request, - &url, - &authentication.headers, - retains_document, - body, - |body| validate_inline_document(&body.document), - ) - .await - } - - fn transform_ocr_response( - &self, - request: &LiteLLMOcrRequest, - response: Self::ProviderResponse, - ) -> Result { - mistral::transform_ocr_response(&request.model, response) - } -} - -fn get_complete_url( - api_base: Option<&str>, - project: &str, - location: &str, - model: &str, -) -> Result { - validate_location(location)?; - let default_base = format!("https://{location}-aiplatform.googleapis.com"); - let base = api_base - .map(str::trim) - .filter(|base| !base.is_empty()) - .unwrap_or(&default_base); - let prediction = format!("{model}:rawPredict"); - ApiUrl::parse(base) - .and_then(|url| { - url.complete_path(&[ - "v1", - "projects", - project, - "locations", - location, - "publishers", - "mistralai", - "models", - &prediction, - ]) - }) - .map(|url| url.into_string()) - .map_err(|_| { - OcrRequestError::RequestField { - path: "api_base".into(), - } - .into() - }) -} - -fn validate_location(location: &str) -> Result<(), OcrError> { - let valid = !location.is_empty() - && location - .bytes() - .all(|value| value.is_ascii_lowercase() || value.is_ascii_digit() || value == b'-') - && location - .as_bytes() - .first() - .is_some_and(u8::is_ascii_alphanumeric) - && location - .as_bytes() - .last() - .is_some_and(u8::is_ascii_alphanumeric); - if valid { - return Ok(()); - } - Err(OcrRequestError::RequestField { - path: "vertex_location".into(), - } - .into()) -} - -#[cfg(test)] -mod tests { - use super::get_complete_url; - - #[test] - fn endpoint_uses_location_project_and_model() { - assert_eq!( - get_complete_url(None, "proj-1", "europe-west4", "mistral-ocr-maas").unwrap(), - "https://europe-west4-aiplatform.googleapis.com/v1/projects/proj-1/locations/europe-west4/publishers/mistralai/models/mistral-ocr-maas:rawPredict" - ); - assert!(get_complete_url(None, "proj-1", "attacker.example/path", "model").is_err()); - } -} diff --git a/litellm-rust/crates/core/src/ocr/adapters/vertex/mod.rs b/litellm-rust/crates/core/src/ocr/adapters/vertex/mod.rs deleted file mode 100644 index 798510e7405..00000000000 --- a/litellm-rust/crates/core/src/ocr/adapters/vertex/mod.rs +++ /dev/null @@ -1,18 +0,0 @@ -mod deepseek; -mod mistral; - -use crate::ocr::Error; -use litellm_auth::InputSource; - -use crate::ocr::error::OcrError; -use crate::ocr::types::OcrConnection; - -pub(crate) use deepseek::VertexDeepSeekAdapter; -pub(crate) use mistral::VertexMistralAdapter; - -fn validate_destination(connection: &OcrConnection) -> Result<(), OcrError> { - if connection.api_base.is_some() && connection.api_base_source == InputSource::Request { - return Err(Error::from(litellm_auth::Error::RequestVertexCredentialDestination).into()); - } - Ok(()) -} diff --git a/litellm-rust/crates/core/src/ocr/arguments.rs b/litellm-rust/crates/core/src/ocr/arguments.rs new file mode 100644 index 00000000000..293931e8bbb --- /dev/null +++ b/litellm-rust/crates/core/src/ocr/arguments.rs @@ -0,0 +1,101 @@ +use crate::call_arguments::ArgumentSpec; + +use super::provider_config::{OcrConfigKind, resolve_provider_config}; + +const COMMON_OPTION_FIELDS: &[&str] = &["req_format", "extra_body", "max_response_bytes"]; +const AZURE_AUTH_OPTION_FIELDS: &[&str] = &[ + "azure_ad_token", + "tenant_id", + "client_id", + "client_secret", + "azure_scope", + "azure_authority_host", + "azure_credential", + "azure_federated_token_file", + "enable_azure_ad_token_refresh", +]; +const VERTEX_AUTH_OPTION_FIELDS: &[&str] = &[ + "vertex_credentials", + "vertex_ai_credentials", + "vertex_project", + "vertex_ai_project", + "vertex_location", + "vertex_ai_location", +]; + +pub fn is_supported_request(model: &str, custom_llm_provider: Option<&str>) -> bool { + resolve_provider_config(model, custom_llm_provider).is_ok() +} + +pub fn consumed_optional_param_names( + model: &str, + custom_llm_provider: Option<&str>, +) -> Result, super::Error> { + let (model, config) = resolve_provider_config(model, custom_llm_provider)?; + let provider_fields = config.get_supported_ocr_params(&model); + let auth_fields: &[&str] = match config { + OcrConfigKind::AzureAi + | OcrConfigKind::AzureDocumentIntelligence + | OcrConfigKind::AzureCohere => AZURE_AUTH_OPTION_FIELDS, + OcrConfigKind::VertexAi | OcrConfigKind::VertexDeepSeek => VERTEX_AUTH_OPTION_FIELDS, + _ => &[], + }; + Ok(COMMON_OPTION_FIELDS + .iter() + .chain(provider_fields) + .chain(auth_fields) + .copied() + .collect()) +} + +pub fn consumed_optional_params( + model: &str, + custom_llm_provider: Option<&str>, +) -> Result, super::Error> { + consumed_optional_param_names(model, custom_llm_provider).map(|names| { + names + .into_iter() + .map(|name| ArgumentSpec { + name, + secret: matches!( + name, + "azure_ad_token" + | "client_secret" + | "azure_federated_token_file" + | "vertex_credentials" + | "vertex_ai_credentials" + ), + }) + .collect() + }) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn consumed_params_include_provider_options_and_mark_credentials() { + let mistral = consumed_optional_param_names("mistral/model", None).unwrap(); + assert!(mistral.contains(&"pages")); + assert!(mistral.contains(&"req_format")); + assert!(!mistral.contains(&"vertex_project")); + + let vertex = consumed_optional_param_names("vertex_ai/deepseek-ocr", None).unwrap(); + assert!(!vertex.contains(&"temperature")); + assert!(vertex.contains(&"vertex_credentials")); + assert!(!vertex.contains(&"pages")); + + let azure = consumed_optional_params("model", Some("azure_ai")).unwrap(); + assert!( + azure + .iter() + .any(|spec| spec.name == "client_secret" && spec.secret) + ); + assert!( + azure + .iter() + .any(|spec| spec.name == "tenant_id" && !spec.secret) + ); + } +} diff --git a/litellm-rust/crates/core/src/ocr/client.rs b/litellm-rust/crates/core/src/ocr/client.rs index 9a30b2f8e04..5881519855c 100644 --- a/litellm-rust/crates/core/src/ocr/client.rs +++ b/litellm-rust/crates/core/src/ocr/client.rs @@ -4,12 +4,10 @@ use std::time::Duration; use bytes::{Bytes, BytesMut}; use serde::de::DeserializeOwned; -use super::error::{Error, OcrError, OcrResponseError}; +use super::json::{DecodedOcrResponse, decode_response}; use super::types::{LiteLLMOcrRequest, LiteLLMOcrResponse}; -use super::wire::{DecodedOcrResponse, decode_response}; use crate::constants::OCR_CONNECT_TIMEOUT_SECS; use crate::media::MediaFetcher; -use crate::transport::Error as TransportError; use litellm_auth_gcp::VertexAuth; #[derive(Clone)] @@ -21,8 +19,8 @@ pub struct OcrClient { } impl OcrClient { - pub fn new(provider_http: reqwest::Client) -> Result { - let document_fetcher = MediaFetcher::new().map_err(TransportError::from)?; + pub fn new(provider_http: reqwest::Client) -> Result { + let document_fetcher = MediaFetcher::new().map_err(crate::transport::Error::from)?; Ok(Self { provider_http, polling_http: no_redirect_http()?, @@ -31,11 +29,14 @@ impl OcrClient { }) } - pub fn shared() -> Result { + pub fn shared() -> Result { shared_client() } - pub async fn perform(&self, request: LiteLLMOcrRequest) -> Result { + pub async fn perform( + &self, + request: LiteLLMOcrRequest, + ) -> Result { use super::{ NativeOutcome, OcrAdmission, OcrCall, OcrCallStep, OcrHookHost, OcrHost, OcrHostOperation, OcrHostResult, @@ -45,7 +46,7 @@ impl OcrClient { let mut request = Some(request); let NativeOutcome::Completed(mut call) = OcrCall::admit(self.clone(), OcrAdmission::all()) else { - return Err(Error::InvalidRequest( + return Err(crate::ocr::Error::InvalidRequest( "native OCR host admission declined".into(), )); }; @@ -54,16 +55,11 @@ impl OcrClient { match call.resume(result.take()).await? { OcrCallStep::Host(OcrHostOperation::ProjectRequest) => { result = Some(OcrHostResult::Request(Ok(( - Box::new( - request - .take() - .ok_or_else(|| { - Error::InvalidRequest( - "OCR request was already projected".into(), - ) - })? - .into(), - ), + Box::new(request.take().ok_or_else(|| { + crate::ocr::Error::InvalidRequest( + "OCR request was already projected".into(), + ) + })?), false, )))) } @@ -100,29 +96,29 @@ impl OcrClient { } } -fn no_redirect_http() -> Result { +fn no_redirect_http() -> Result { reqwest::Client::builder() .connect_timeout(Duration::from_secs(OCR_CONNECT_TIMEOUT_SECS)) .redirect(reqwest::redirect::Policy::none()) .build() - .map_err(TransportError::from) + .map_err(crate::transport::Error::from) } -pub(crate) fn shared_client() -> Result { - static CLIENT: OnceLock> = OnceLock::new(); +pub(crate) fn shared_client() -> Result { + static CLIENT: OnceLock> = OnceLock::new(); let client = CLIENT .get_or_init(|| { reqwest::Client::builder() .connect_timeout(Duration::from_secs(OCR_CONNECT_TIMEOUT_SECS)) .build() - .map_err(TransportError::from) + .map_err(crate::transport::Error::from) .and_then(OcrClient::new) }) .clone()?; Ok(client) } -pub async fn ocr(request: LiteLLMOcrRequest) -> Result { +pub async fn ocr(request: LiteLLMOcrRequest) -> Result { shared_client()?.perform(request).await } @@ -130,15 +126,15 @@ pub async fn read_json_response( response: reqwest::Response, native: bool, max_response_bytes: usize, -) -> Result, OcrError> { +) -> Result, crate::ocr::Error> { let bytes = read_response_bytes(response, max_response_bytes).await?; - Ok(decode_response(&bytes, native)?) + decode_response(&bytes, native) } pub(crate) async fn read_response_bytes( mut response: reqwest::Response, max_response_bytes: usize, -) -> Result { +) -> Result { let status = response.status(); let limit = if status.is_success() { max_response_bytes @@ -150,13 +146,13 @@ pub(crate) async fn read_response_bytes( .content_length() .is_some_and(|length| length > limit as u64) { - return Err(OcrResponseError::TooLarge { limit }.into()); + return Err(crate::ocr::Error::TooLarge { limit }); } let mut bytes = BytesMut::new(); while let Some(chunk) = response.chunk().await.map_err(transport_error)? { let remaining = limit.saturating_sub(bytes.len()); if status.is_success() && chunk.len() > remaining { - return Err(OcrResponseError::TooLarge { limit }.into()); + return Err(crate::ocr::Error::TooLarge { limit }); } bytes.extend_from_slice(&chunk[..chunk.len().min(remaining)]); if !status.is_success() && bytes.len() == limit { @@ -173,12 +169,12 @@ pub(crate) async fn read_response_bytes( Ok(bytes.freeze()) } -pub(crate) fn transport_error(error: reqwest::Error) -> Error { +pub(crate) fn transport_error(error: reqwest::Error) -> crate::ocr::Error { if error.is_timeout() { - return Error::Http { + return crate::ocr::Error::Transport(crate::transport::Error::Http { status: 408, body: "OCR request timed out".into(), - }; + }); } crate::transport::Error::from(error).into() } @@ -203,7 +199,7 @@ mod tests { .unwrap_err(); assert!(matches!( transport_error(error), - Error::Http { status: 408, .. } + crate::ocr::Error::Transport(crate::transport::Error::Http { status: 408, .. }) )); server.abort(); } diff --git a/litellm-rust/crates/core/src/ocr/codecs/cohere.rs b/litellm-rust/crates/core/src/ocr/codecs/cohere.rs deleted file mode 100644 index 649432f39d3..00000000000 --- a/litellm-rust/crates/core/src/ocr/codecs/cohere.rs +++ /dev/null @@ -1,254 +0,0 @@ -use serde::{Deserialize, Serialize}; -use serde_json::{Map, Value, json}; - -use crate::ocr::document::InlineDocument; -use crate::ocr::error::{OcrRequestError, OcrResponseError}; -use crate::ocr::types::{LiteLLMOcrResponse, OcrDocument}; - -#[derive(Clone, Copy, Debug, Default, Deserialize, Serialize)] -#[serde(rename_all = "lowercase")] -pub(crate) enum OutputFormat { - #[default] - Markdown, - Blocks, -} - -#[derive(Deserialize)] -pub(crate) struct CohereParams { - #[serde(default)] - pub output_format: OutputFormat, -} - -#[derive(Deserialize, Serialize)] -pub(crate) struct CohereRequest { - pub model: String, - pub document: OcrDocument, - pub output_format: OutputFormat, -} - -pub(crate) fn validate_document(document: &OcrDocument) -> Result<(), OcrRequestError> { - let OcrDocument::ImageUrl { image_url, .. } = document else { - return Err(OcrRequestError::CohereImageOnly); - }; - if image_url.is_empty() { - return Err(OcrRequestError::CohereImageOnly); - } - if let Some(inline) = InlineDocument::parse(image_url)? { - if !inline.mime_type().type_.eq_ignore_ascii_case("image") { - return Err(OcrRequestError::CohereImageOnly); - } - inline.decode(crate::constants::OCR_INLINE_MAX_BYTES)?; - } - Ok(()) -} - -#[derive(Deserialize)] -pub(crate) struct CohereResponse { - #[serde(default)] - pages: Vec, - meta: Option, -} - -#[derive(Deserialize)] -struct CoherePage { - index: Option, - markdown: Option, - blocks: Option>>, -} - -#[derive(Deserialize)] -struct CohereMarkdown { - #[serde(default)] - content: String, - images: Option>>, -} - -#[derive(Deserialize)] -struct CohereMeta { - billed_units: Option, -} - -#[derive(Deserialize)] -struct CohereBilledUnits { - pages: Option, -} - -pub(crate) fn transform_response( - model: &str, - response: CohereResponse, -) -> Result { - let pages_processed = response - .meta - .and_then(|meta| meta.billed_units) - .and_then(|units| units.pages) - .map(Ok) - .unwrap_or_else(|| { - i64::try_from(response.pages.len()).map_err(|_| OcrResponseError::NumericRange("pages")) - })?; - let pages = response - .pages - .into_iter() - .enumerate() - .map(|(position, page)| { - let index = page.index.map(Ok).unwrap_or_else(|| { - i64::try_from(position).map_err(|_| OcrResponseError::NumericRange("page index")) - })?; - let (content, images) = page - .markdown - .map(|markdown| { - let images = - markdown - .images - .filter(|images| !images.is_empty()) - .map(|images| { - images - .into_iter() - .map(|mut image| { - if let Some(Value::Object(bbox)) = - image.get("bounding_box").cloned() - { - image.insert("bbox".into(), Value::Object(bbox)); - } - Value::Object(image) - }) - .collect::>() - }); - (markdown.content, images) - }) - .unwrap_or_default(); - let mut normalized = json!({"index": index, "markdown": content, "images": images}); - if let Some(blocks) = page.blocks { - normalized["blocks"] = json!(blocks); - } - Ok(normalized) - }) - .collect::, OcrResponseError>>()?; - Ok(LiteLLMOcrResponse { - pages, - model: model.into(), - document_annotation: None, - usage_info: Some(json!({"pages_processed": pages_processed})), - object: "ocr".into(), - extra_fields: Map::new(), - provider_native_response: None, - }) -} - -pub(crate) fn transform_request( - model: &str, - document: OcrDocument, - params: CohereParams, -) -> Result { - validate_document(&document)?; - Ok(CohereRequest { - model: model.into(), - document, - output_format: params.output_format, - }) -} - -#[cfg(test)] -mod tests { - use super::*; - - #[test] - fn response_normalizes_markdown_images_blocks_and_billed_pages() { - let response = serde_json::from_value(json!({ - "pages": [ - { - "type":"markdown", - "index":4, - "markdown":{ - "content":"receipt", - "images":[{ - "id":"image", - "bounding_box":{"top_left_x":1,"bottom_right_x":48}, - "bounding_box_normalized":{"top_left_x":0.04,"bottom_right_x":0.15}, - "description":"scan", - "category":"logo" - }] - } - }, - {"type":"blocks","blocks":[{"type":"text","text":{"content":"total"}}]} - ], - "meta":{"api_version":{"version":"2"},"billed_units":{"pages":3}} - })) - .unwrap(); - let normalized = transform_response("parse-v5.0", response).unwrap(); - assert_eq!(normalized.pages[0]["index"], 4); - assert_eq!(normalized.pages[0]["markdown"], "receipt"); - assert_eq!(normalized.pages[0]["images"][0]["bbox"]["top_left_x"], 1); - assert_eq!( - normalized.pages[0]["images"][0]["bounding_box_normalized"]["bottom_right_x"], - 0.15 - ); - assert_eq!(normalized.pages[0]["images"][0]["description"], "scan"); - assert_eq!(normalized.pages[0]["images"][0]["category"], "logo"); - assert_eq!(normalized.pages[1]["index"], 1); - assert_eq!(normalized.pages[1]["markdown"], ""); - assert_eq!(normalized.pages[1]["blocks"][0]["text"]["content"], "total"); - assert_eq!(normalized.usage_info.unwrap()["pages_processed"], 3); - } - - #[test] - fn response_defaults_and_invalid_fields() { - for value in [ - json!({}), - json!({"meta":null}), - json!({"pages":[],"meta":{"billed_units":null}}), - ] { - let normalized = - transform_response("parse", serde_json::from_value(value).unwrap()).unwrap(); - assert!(normalized.pages.is_empty()); - assert_eq!(normalized.usage_info.unwrap()["pages_processed"], 0); - } - for value in [ - json!({"pages":null}), - json!({"pages":[{"markdown":"text"}]}), - json!({"pages":[{"index":"bad"}]}), - ] { - assert!(serde_json::from_value::(value).is_err()); - } - let normalized = transform_response( - "parse", - serde_json::from_value(json!({"pages":[{"markdown":null}]})).unwrap(), - ) - .unwrap(); - assert_eq!(normalized.usage_info.unwrap()["pages_processed"], 1); - assert!(normalized.pages[0]["images"].is_null()); - } - - #[test] - fn request_requires_image_and_supported_output_format() { - for value in [ - json!({"type":"document_url","document_url":"https://example.com/a.pdf"}), - json!({"type":"image_url","image_url":""}), - json!({"type":"image_url","image_url":"data:application/pdf;base64,YQ=="}), - ] { - assert_eq!( - validate_document(&serde_json::from_value(value).unwrap()), - Err(OcrRequestError::CohereImageOnly) - ); - } - assert!(serde_json::from_value::(json!({"output_format":"html"})).is_err()); - for format in ["markdown", "blocks"] { - assert!( - serde_json::from_value::(json!({"output_format":format})).is_ok() - ); - } - let request = transform_request( - "parse-v5.0", - serde_json::from_value(json!({ - "type":"image_url", - "image_url":"https://example.com/image.png" - })) - .unwrap(), - serde_json::from_value(json!({})).unwrap(), - ) - .unwrap(); - assert_eq!( - serde_json::to_value(request).unwrap()["output_format"], - "markdown" - ); - } -} diff --git a/litellm-rust/crates/core/src/ocr/codecs/deepseek/mod.rs b/litellm-rust/crates/core/src/ocr/codecs/deepseek/mod.rs deleted file mode 100644 index 682b3addde7..00000000000 --- a/litellm-rust/crates/core/src/ocr/codecs/deepseek/mod.rs +++ /dev/null @@ -1,5 +0,0 @@ -mod transformation; -mod types; - -pub(crate) use transformation::{transform_ocr_request, transform_ocr_response}; -pub(crate) use types::{DeepSeekOcrParams, DeepSeekOcrResponse}; diff --git a/litellm-rust/crates/core/src/ocr/codecs/deepseek/transformation.rs b/litellm-rust/crates/core/src/ocr/codecs/deepseek/transformation.rs deleted file mode 100644 index 999ac6cf032..00000000000 --- a/litellm-rust/crates/core/src/ocr/codecs/deepseek/transformation.rs +++ /dev/null @@ -1,101 +0,0 @@ -use serde::de::IntoDeserializer; -use serde_json::{Value, json}; - -use super::types::*; -use crate::ocr::error::{OcrRequestError, OcrResponseError}; -use crate::ocr::types::{LiteLLMOcrResponse, OcrDocument}; - -pub(crate) fn transform_ocr_request( - provider_model: &str, - document: OcrDocument, - params: &DeepSeekOcrParams, -) -> Result { - if document.source().is_empty() { - return Err(OcrRequestError::MissingDocumentUrl); - } - let content = OcrDocument::ImageUrl { - image_url: document.source().to_string(), - extra_fields: serde_json::Map::new(), - }; - Ok(DeepSeekOcrRequest { - model: provider_model.to_string(), - messages: vec![DeepSeekOcrMessage { - role: UserRole::User, - content: vec![content], - }], - params: params.clone(), - }) -} - -pub(crate) fn transform_ocr_response( - model: &str, - response: DeepSeekOcrResponse, -) -> Result { - let content = response - .choices - .into_iter() - .next() - .and_then(|choice| choice.message.content) - .ok_or(OcrResponseError::EmptyContent)?; - let decoded = decode_content(content)?; - let pages = match decoded.result.pages { - Some(pages) if !pages.is_empty() => pages - .into_iter() - .map(|page| serde_json::to_value(page).expect("DeepSeek page serializes")) - .collect(), - _ => vec![json!({ - "index":0, - "markdown":decoded.fallback_markdown, - "images":null - })], - }; - Ok(LiteLLMOcrResponse { - pages, - model: decoded.result.model.unwrap_or_else(|| model.to_string()), - document_annotation: decoded.result.document_annotation, - usage_info: decoded.result.usage_info.or(response.usage), - object: "ocr".into(), - extra_fields: decoded.result.extra_fields, - provider_native_response: None, - }) -} - -struct DecodedContent { - result: DeepSeekOcrResult, - fallback_markdown: String, -} - -fn decode_content(content: DeepSeekContent) -> Result { - let (result, fallback_markdown) = match content { - DeepSeekContent::Text(text) if text.is_empty() => { - return Err(OcrResponseError::EmptyContent); - } - DeepSeekContent::Text(text) => (decode_json_content(&text)?, text), - DeepSeekContent::Object(object) => { - let fallback = - serde_json::to_string(&object).map_err(|_| OcrResponseError::ResponseField { - path: "choices[0].message.content".into(), - })?; - (Some(object), fallback) - } - }; - Ok(DecodedContent { - result: result.unwrap_or_default(), - fallback_markdown, - }) -} - -fn decode_json_content(text: &str) -> Result, OcrResponseError> { - if !text.trim_start().starts_with('{') { - return Ok(None); - } - let value = match serde_json::from_str::(text) { - Ok(value) => value, - Err(_) => return Ok(None), - }; - serde_path_to_error::deserialize(value.into_deserializer()) - .map(Some) - .map_err(|error| OcrResponseError::ResponseField { - path: format!("choices[0].message.content.{}", error.path()), - }) -} diff --git a/litellm-rust/crates/core/src/ocr/codecs/deepseek/types.rs b/litellm-rust/crates/core/src/ocr/codecs/deepseek/types.rs deleted file mode 100644 index 0ce2d9913f7..00000000000 --- a/litellm-rust/crates/core/src/ocr/codecs/deepseek/types.rs +++ /dev/null @@ -1,95 +0,0 @@ -use serde::{Deserialize, Serialize}; -use serde_json::{Map, Value}; - -#[derive(Clone, Debug, Default, Serialize, Deserialize)] -pub(crate) struct DeepSeekOcrParams { - #[serde(skip_serializing_if = "Option::is_none")] - pub stream: Option, - #[serde(skip_serializing_if = "Option::is_none")] - pub temperature: Option, - #[serde(skip_serializing_if = "Option::is_none")] - pub max_tokens: Option, - #[serde(skip_serializing_if = "Option::is_none")] - pub top_p: Option, - #[serde(skip_serializing_if = "Option::is_none")] - pub n: Option, - #[serde(skip_serializing_if = "Option::is_none")] - pub stop: Option, -} - -#[derive(Clone, Debug, Serialize, Deserialize)] -#[serde(untagged)] -pub(crate) enum StopSequences { - One(String), - Many(Vec), -} - -#[derive(Clone, Debug, Serialize, Deserialize)] -pub(crate) struct DeepSeekOcrRequest { - pub model: String, - pub messages: Vec, - #[serde(flatten)] - pub params: DeepSeekOcrParams, -} - -#[derive(Clone, Debug, Serialize, Deserialize)] -pub(crate) struct DeepSeekOcrMessage { - pub role: UserRole, - pub content: Vec, -} - -#[derive(Clone, Debug, Serialize, Deserialize)] -#[serde(rename_all = "lowercase")] -pub(crate) enum UserRole { - User, -} - -#[derive(Clone, Debug, Deserialize)] -pub(crate) struct DeepSeekOcrResponse { - #[serde(default)] - pub choices: Vec, - pub usage: Option, -} - -#[derive(Clone, Debug, Deserialize)] -pub(crate) struct DeepSeekChoice { - pub message: DeepSeekResponseMessage, -} - -#[derive(Clone, Debug, Deserialize)] -pub(crate) struct DeepSeekResponseMessage { - pub content: Option, -} - -#[derive(Clone, Debug, Deserialize)] -#[serde(untagged)] -pub(crate) enum DeepSeekContent { - Text(String), - Object(DeepSeekOcrResult), -} - -#[derive(Clone, Debug, Default, Serialize, Deserialize)] -pub(crate) struct DeepSeekOcrResult { - #[serde(skip_serializing_if = "Option::is_none")] - pub pages: Option>, - #[serde(skip_serializing_if = "Option::is_none")] - pub model: Option, - #[serde(skip_serializing_if = "Option::is_none")] - pub usage_info: Option, - #[serde(skip_serializing_if = "Option::is_none")] - pub document_annotation: Option, - #[serde(flatten)] - pub extra_fields: Map, -} - -#[derive(Clone, Debug, Serialize, Deserialize)] -pub(crate) struct DeepSeekPage { - #[serde(default)] - pub index: i64, - #[serde(default)] - pub markdown: String, - pub images: Option, - pub dimensions: Option, - #[serde(flatten)] - pub extra_fields: Map, -} diff --git a/litellm-rust/crates/core/src/ocr/codecs/document_intelligence/mod.rs b/litellm-rust/crates/core/src/ocr/codecs/document_intelligence/mod.rs deleted file mode 100644 index 8031f2124a3..00000000000 --- a/litellm-rust/crates/core/src/ocr/codecs/document_intelligence/mod.rs +++ /dev/null @@ -1,9 +0,0 @@ -mod params; -mod transformation; -mod types; - -pub(crate) use params::{decode_input_params, map_ocr_params}; -pub(crate) use transformation::{transform_ocr_request, transform_ocr_response}; -pub(crate) use types::{ - AzureDocumentIntelligenceOperation, DocumentIntelligenceParams, OperationStatus, -}; diff --git a/litellm-rust/crates/core/src/ocr/codecs/document_intelligence/params.rs b/litellm-rust/crates/core/src/ocr/codecs/document_intelligence/params.rs deleted file mode 100644 index 9389f93b8e3..00000000000 --- a/litellm-rust/crates/core/src/ocr/codecs/document_intelligence/params.rs +++ /dev/null @@ -1,219 +0,0 @@ -use std::collections::BTreeSet; - -use serde_json::{Map, Value}; - -use super::types::{ - DocumentIntelligenceInputParams, DocumentIntelligenceParams, FeaturesInput, PagesInput, -}; -use crate::ocr::error::OcrRequestError; -use crate::ocr::prepare::ParsedProviderParams; - -pub(crate) fn decode_input_params( - params: Map, - prefix: &str, -) -> Result, OcrRequestError> { - if let Some(Value::Array(pages)) = params.get("pages") { - if pages.iter().any(Value::is_boolean) { - return Err(OcrRequestError::Pages("boolean page index".into())); - } - if pages - .iter() - .any(|page| page.is_number() && page.as_i64().is_none()) - { - return Err(OcrRequestError::Pages("page index is out of range".into())); - } - if !pages.iter().all(Value::is_i64) && !pages.iter().all(Value::is_string) { - return Err(OcrRequestError::Pages("mixed page element types".into())); - } - } - crate::ocr::wire::decode_request_value(Value::Object(params), prefix) -} - -pub(crate) fn map_ocr_params( - params: DocumentIntelligenceInputParams, -) -> Result { - Ok(DocumentIntelligenceParams { - pages: params.pages.map(normalize_pages).transpose()?.flatten(), - features: params - .features - .map(normalize_features) - .transpose()? - .flatten(), - }) -} - -fn normalize_pages(pages: PagesInput) -> Result, OcrRequestError> { - let normalized = match pages { - PagesInput::ZeroBasedIndices(indices) => { - if indices.is_empty() { - return Ok(None); - } - indices - .into_iter() - .map(|page| { - if page < 0 { - return Err(OcrRequestError::Pages("negative page index".into())); - } - page.checked_add(1) - .ok_or_else(|| OcrRequestError::Pages("page index is out of range".into())) - }) - .collect::, _>>()? - .into_iter() - .map(|page| page.to_string()) - .collect::>() - .join(",") - } - PagesInput::NativeTokens(tokens) => { - if tokens.is_empty() { - return Ok(None); - } - tokens - .iter() - .map(|token| token.trim()) - .collect::>() - .join(",") - } - PagesInput::NativeRange(range) => range - .split(',') - .map(str::trim) - .collect::>() - .join(","), - }; - if !normalized.split(',').all(valid_page_token) { - return Err(OcrRequestError::Pages("invalid native page range".into())); - } - Ok(Some(normalized)) -} - -fn valid_page_token(token: &str) -> bool { - let mut parts = token.split('-'); - let start = parts.next().unwrap_or_default(); - if start.is_empty() || !start.chars().all(|character| character.is_ascii_digit()) { - return false; - } - match parts.next() { - None => true, - Some(end) => { - !end.is_empty() - && end.chars().all(|character| character.is_ascii_digit()) - && parts.next().is_none() - } - } -} - -fn normalize_features(features: FeaturesInput) -> Result, OcrRequestError> { - let tokens = match features { - FeaturesInput::Names(names) => names, - FeaturesInput::CommaSeparated(names) => names.split(',').map(str::to_string).collect(), - }; - if tokens.is_empty() { - return Ok(None); - } - let normalized = tokens.iter().map(|token| token.trim()).collect::>(); - if !normalized.iter().all(|token| { - let Some((first, rest)) = token.as_bytes().split_first() else { - return false; - }; - first.is_ascii_alphabetic() && rest.iter().all(u8::is_ascii_alphanumeric) - }) { - return Err(OcrRequestError::Features); - } - Ok(Some(normalized.join(","))) -} - -#[cfg(test)] -mod tests { - use rstest::rstest; - use serde_json::{Value, json}; - - use super::*; - - fn map(value: Value) -> Result { - let fields = value.as_object().unwrap().clone(); - map_ocr_params(decode_input_params(fields, "optional_params")?.known) - } - - #[test] - fn input_params_retain_unknown_fields() { - let parsed = decode_input_params( - json!({ - "pages": [0], - "future_ocr_option": true, - "extra_body": {"provider_option": "value"} - }) - .as_object() - .unwrap() - .clone(), - "optional_params", - ) - .unwrap(); - - assert_eq!( - parsed.known.pages, - Some(PagesInput::ZeroBasedIndices(vec![0])) - ); - assert_eq!(parsed.extra_params["future_ocr_option"], true); - assert_eq!( - parsed.extra_params["extra_body"], - json!({"provider_option": "value"}) - ); - assert_eq!( - serde_json::to_value(map_ocr_params(parsed.known).unwrap()).unwrap(), - json!({"pages": "1", "features": null}) - ); - } - - #[rstest] - #[case(json!([0, 1, 2]), Some("1,2,3"))] - #[case(json!([2, 0, 0, 1]), Some("1,2,3"))] - #[case(json!([]), None)] - #[case(json!("3-9"), Some("3-9"))] - #[case(json!("1-3, 5"), Some("1-3,5"))] - #[case(json!(["1", "3-5"]), Some("1,3-5"))] - fn page_mapping_matches_python(#[case] input: Value, #[case] expected: Option<&str>) { - assert_eq!( - map(json!({"pages": input})).unwrap().pages.as_deref(), - expected - ); - } - - #[rstest] - #[case(json!("a,b"))] - #[case(json!([-1]))] - #[case(json!([true, false]))] - #[case(json!([1, "2"]))] - #[case(json!(5))] - fn invalid_page_mapping_matches_python(#[case] input: Value) { - assert!(map(json!({"pages": input})).is_err()); - } - - #[rstest] - #[case(json!(["keyValuePairs"]), "keyValuePairs")] - #[case(json!(["keyValuePairs", "languages"]), "keyValuePairs,languages")] - #[case(json!("keyValuePairs"), "keyValuePairs")] - #[case(json!("keyValuePairs,languages"), "keyValuePairs,languages")] - #[case(json!("keyValuePairs, languages"), "keyValuePairs,languages")] - fn feature_mapping_matches_python(#[case] input: Value, #[case] expected: &str) { - assert_eq!( - map(json!({"features": input})).unwrap().features.as_deref(), - Some(expected) - ); - } - - #[rstest] - #[case(json!("keyValuePairs&pages=9"))] - #[case(json!("key value pairs"))] - #[case(json!(""))] - #[case(json!([1, 2]))] - #[case(json!([["keyValuePairs"]]))] - #[case(json!({"feature":"keyValuePairs"}))] - #[case(json!(5))] - fn invalid_feature_mapping_matches_python(#[case] input: Value) { - assert!(map(json!({"features": input})).is_err()); - } - - #[test] - fn empty_feature_list_is_omitted() { - assert_eq!(map(json!({"features": []})).unwrap().features, None); - } -} diff --git a/litellm-rust/crates/core/src/ocr/codecs/document_intelligence/transformation.rs b/litellm-rust/crates/core/src/ocr/codecs/document_intelligence/transformation.rs deleted file mode 100644 index 018d7eb9c65..00000000000 --- a/litellm-rust/crates/core/src/ocr/codecs/document_intelligence/transformation.rs +++ /dev/null @@ -1,107 +0,0 @@ -use base64::{Engine, engine::general_purpose::STANDARD}; -use serde_json::{Map, Value, json}; - -use super::types::*; -use crate::constants::{AZURE_DI_DEFAULT_DPI, AZURE_DI_DEFAULT_HEIGHT, AZURE_DI_DEFAULT_WIDTH}; -use crate::ocr::document::InlineDocument; -use crate::ocr::error::{OcrRequestError, OcrResponseError}; -use crate::ocr::types::{LiteLLMOcrResponse, OcrDocument}; - -pub(crate) fn transform_ocr_request( - document: OcrDocument, -) -> Result { - let source = document.source(); - if source.is_empty() { - return Err(OcrRequestError::MissingDocumentUrl); - } - Ok(if let Some(document) = InlineDocument::parse(source)? { - DocumentIntelligenceRequest::Base64Source( - STANDARD.encode(document.decode(crate::constants::OCR_INLINE_MAX_BYTES)?), - ) - } else { - DocumentIntelligenceRequest::UrlSource(source.to_string()) - }) -} - -pub(crate) fn transform_ocr_response( - model: &str, - response: AzureDocumentIntelligenceOperation, -) -> Result { - if response.status != Some(OperationStatus::Succeeded) { - return Err(OcrResponseError::OperationStatus( - response - .status - .map(|status| status.to_string()) - .unwrap_or_else(|| "None".into()), - )); - } - let result = response.analyze_result.unwrap_or_default(); - let pages = result - .pages - .into_iter() - .map(normalize_page) - .collect::, _>>()?; - let pages_processed = pages.len(); - let mut extra_fields = Map::new(); - extra_fields.insert("content".into(), option_value(result.content)); - extra_fields.insert("tables".into(), option_value(result.tables)); - extra_fields.insert("keyValuePairs".into(), option_value(result.key_value_pairs)); - Ok(LiteLLMOcrResponse { - pages, - model: model.into(), - document_annotation: None, - usage_info: Some(json!({"pages_processed":pages_processed})), - object: "ocr".into(), - extra_fields, - provider_native_response: None, - }) -} - -fn normalize_page(page: AzureDocumentIntelligencePage) -> Result { - let index = page - .page_number - .unwrap_or(1) - .checked_sub(1) - .ok_or(OcrResponseError::NumericRange("page.pageNumber"))?; - let scale = if page.unit.as_deref().unwrap_or("inch") == "inch" { - AZURE_DI_DEFAULT_DPI as f64 - } else { - 1.0 - }; - let width = pixel_dimension( - page.width.unwrap_or(AZURE_DI_DEFAULT_WIDTH), - scale, - "page.width", - )?; - let height = pixel_dimension( - page.height.unwrap_or(AZURE_DI_DEFAULT_HEIGHT), - scale, - "page.height", - )?; - let markdown = page - .lines - .iter() - .map(|line| line.content.as_deref().unwrap_or_default()) - .collect::>() - .join("\n"); - Ok(json!({ - "index":index, - "markdown":markdown, - "images":null, - "dimensions":{"width":width,"height":height,"dpi":AZURE_DI_DEFAULT_DPI} - })) -} - -fn pixel_dimension(value: f64, scale: f64, field: &'static str) -> Result { - let value = value * scale; - if !value.is_finite() || value < i64::MIN as f64 || value > i64::MAX as f64 { - return Err(OcrResponseError::NumericRange(field)); - } - Ok(value.trunc() as i64) -} - -fn option_value(value: Option) -> Value { - value - .and_then(|value| serde_json::to_value(value).ok()) - .unwrap_or(Value::Null) -} diff --git a/litellm-rust/crates/core/src/ocr/codecs/document_intelligence/types.rs b/litellm-rust/crates/core/src/ocr/codecs/document_intelligence/types.rs deleted file mode 100644 index 793f4547e99..00000000000 --- a/litellm-rust/crates/core/src/ocr/codecs/document_intelligence/types.rs +++ /dev/null @@ -1,138 +0,0 @@ -use serde::{Deserialize, Deserializer, Serialize}; -use serde_json::{Map, Value}; - -#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)] -#[serde(untagged)] -pub(crate) enum PagesInput { - ZeroBasedIndices(Vec), - NativeTokens(Vec), - NativeRange(String), -} - -#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)] -#[serde(untagged)] -pub(crate) enum FeaturesInput { - Names(Vec), - CommaSeparated(String), -} - -#[derive(Clone, Debug, Default, PartialEq, Serialize, Deserialize)] -pub(crate) struct DocumentIntelligenceInputParams { - pub pages: Option, - pub features: Option, -} - -#[derive(Clone, Debug, PartialEq, Serialize)] -pub(crate) struct DocumentIntelligenceParams { - pub pages: Option, - pub features: Option, -} - -#[derive(Clone, Debug, Serialize, Deserialize)] -pub(crate) enum DocumentIntelligenceRequest { - #[serde(rename = "urlSource")] - UrlSource(String), - #[serde(rename = "base64Source")] - Base64Source(String), -} - -#[derive(Clone, Debug, PartialEq)] -pub(crate) enum OperationStatus { - Succeeded, - Running, - NotStarted, - Failed, - Unknown(String), -} - -impl<'de> Deserialize<'de> for OperationStatus { - fn deserialize>(deserializer: D) -> Result { - Ok(match String::deserialize(deserializer)?.as_str() { - "succeeded" => Self::Succeeded, - "running" => Self::Running, - "notStarted" => Self::NotStarted, - "failed" => Self::Failed, - value => Self::Unknown(value.to_string()), - }) - } -} - -impl std::fmt::Display for OperationStatus { - fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - formatter.write_str(match self { - Self::Succeeded => "succeeded", - Self::Running => "running", - Self::NotStarted => "notStarted", - Self::Failed => "failed", - Self::Unknown(value) => value, - }) - } -} - -#[derive(Clone, Debug, Deserialize)] -pub(crate) struct AzureDocumentIntelligenceOperation { - pub status: Option, - #[serde(rename = "analyzeResult")] - pub analyze_result: Option, -} - -#[derive(Clone, Debug, Default, Deserialize)] -pub(crate) struct AzureDocumentIntelligenceAnalyzeResult { - pub content: Option, - #[serde(default)] - pub pages: Vec, - pub tables: Option>>, - #[serde(rename = "keyValuePairs")] - pub key_value_pairs: Option>>, -} - -#[derive(Clone, Debug, Deserialize)] -pub(crate) struct AzureDocumentIntelligencePage { - #[serde(rename = "pageNumber", default, deserialize_with = "optional_i64")] - pub page_number: Option, - #[serde(default, deserialize_with = "optional_f64")] - pub width: Option, - #[serde(default, deserialize_with = "optional_f64")] - pub height: Option, - pub unit: Option, - #[serde(default)] - pub lines: Vec, -} - -#[derive(Clone, Debug, Deserialize)] -pub(crate) struct AzureDocumentIntelligenceLine { - pub content: Option, -} - -fn optional_i64<'de, D: Deserializer<'de>>(deserializer: D) -> Result, D::Error> { - match Option::::deserialize(deserializer)? { - None | Some(Value::Null) => Ok(None), - Some(Value::Number(number)) => number - .as_i64() - .map(Some) - .ok_or_else(|| serde::de::Error::custom("expected an integer")), - Some(Value::String(value)) => value - .parse::() - .map(Some) - .map_err(|_| serde::de::Error::custom("expected an integer")), - Some(_) => Err(serde::de::Error::custom("expected an integer")), - } -} - -fn optional_f64<'de, D: Deserializer<'de>>(deserializer: D) -> Result, D::Error> { - match Option::::deserialize(deserializer)? { - None | Some(Value::Null) => Ok(None), - Some(Value::Number(number)) => number - .as_f64() - .filter(|value| value.is_finite()) - .map(Some) - .ok_or_else(|| serde::de::Error::custom("expected a finite number")), - Some(Value::String(value)) => value - .parse::() - .ok() - .filter(|value| value.is_finite()) - .map(Some) - .ok_or_else(|| serde::de::Error::custom("expected a finite number")), - Some(_) => Err(serde::de::Error::custom("expected a number")), - } -} diff --git a/litellm-rust/crates/core/src/ocr/codecs/mistral/mod.rs b/litellm-rust/crates/core/src/ocr/codecs/mistral/mod.rs deleted file mode 100644 index eea4254779e..00000000000 --- a/litellm-rust/crates/core/src/ocr/codecs/mistral/mod.rs +++ /dev/null @@ -1,5 +0,0 @@ -mod transformation; -mod types; - -pub(crate) use transformation::{transform_ocr_request, transform_ocr_response}; -pub(crate) use types::{MistralOcrParams, MistralOcrRequest, MistralOcrResponse}; diff --git a/litellm-rust/crates/core/src/ocr/codecs/mistral/transformation.rs b/litellm-rust/crates/core/src/ocr/codecs/mistral/transformation.rs deleted file mode 100644 index e8073905548..00000000000 --- a/litellm-rust/crates/core/src/ocr/codecs/mistral/transformation.rs +++ /dev/null @@ -1,250 +0,0 @@ -use super::{MistralOcrParams, MistralOcrRequest, MistralOcrResponse}; -use crate::ocr::error::{OcrRequestError, OcrResponseError}; -use crate::ocr::types::{LiteLLMOcrResponse, OcrDocument}; - -pub(crate) fn transform_ocr_request( - model: &str, - document: OcrDocument, - params: &MistralOcrParams, -) -> Result { - Ok(MistralOcrRequest { - model: model.to_string(), - document, - params: params.clone(), - }) -} - -pub(crate) fn transform_ocr_response( - model: &str, - response: MistralOcrResponse, -) -> Result { - Ok(LiteLLMOcrResponse { - pages: response.pages, - model: response.model.unwrap_or_else(|| model.to_string()), - document_annotation: response.document_annotation, - usage_info: response.usage_info, - object: "ocr".to_string(), - extra_fields: response.extra_fields, - provider_native_response: None, - }) -} - -#[cfg(test)] -mod tests { - use super::*; - use rstest::rstest; - use serde_json::{Value, json}; - - fn mapped_params(value: Value) -> Value { - serde_json::to_value(serde_json::from_value::(value).unwrap()).unwrap() - } - - fn document() -> OcrDocument { - serde_json::from_value( - json!({"type":"document_url","document_url":"https://example.com/a.pdf"}), - ) - .unwrap() - } - - #[rstest] - fn extract_header_is_a_supported_ocr_param() { - assert_eq!( - mapped_params(json!({"extract_header":true}))["extract_header"], - true - ); - } - - #[rstest] - fn extract_footer_is_a_supported_ocr_param() { - assert_eq!( - mapped_params(json!({"extract_footer":false}))["extract_footer"], - false - ); - } - - #[rstest] - fn existing_ocr_params_remain_supported() { - let mapped = mapped_params(json!({ - "pages":[0,2], - "include_image_base64":true, - "image_limit":2, - "image_min_size":100, - "bbox_annotation_format":{"type":"json_schema"}, - "document_annotation_format":{"type":"json_schema"} - })); - assert_eq!(mapped["pages"], json!([0, 2])); - assert_eq!(mapped["include_image_base64"], true); - assert_eq!(mapped["image_limit"], 2); - assert_eq!(mapped["image_min_size"], 100); - assert_eq!(mapped["bbox_annotation_format"]["type"], "json_schema"); - assert_eq!(mapped["document_annotation_format"]["type"], "json_schema"); - } - - #[rstest] - fn map_ocr_params_forwards_extract_header() { - assert_eq!( - mapped_params(json!({"extract_header":true}))["extract_header"], - true - ); - } - - #[rstest] - fn map_ocr_params_forwards_extract_footer() { - assert_eq!( - mapped_params(json!({"extract_footer":true}))["extract_footer"], - true - ); - } - - #[rstest] - fn map_ocr_params_forwards_extract_header_and_footer() { - let mapped = mapped_params(json!({"extract_header":true,"extract_footer":false})); - assert_eq!(mapped["extract_header"], true); - assert_eq!(mapped["extract_footer"], false); - } - - #[rstest] - fn map_ocr_params_drops_unknown_params() { - let mapped = mapped_params(json!({"extract_header":true,"unsupported_param":"value"})); - assert_eq!(mapped["extract_header"], true); - assert!(mapped.get("unsupported_param").is_none()); - } - - #[rstest] - #[case("table_format", json!("html"))] - #[case("confidence_scores_granularity", json!("word"))] - #[case("confidence_scores_granularity", json!("block"))] - #[case("document_annotation_prompt", json!("extract"))] - #[case("include_blocks", json!(true))] - #[case("id", json!("req-123"))] - fn new_ocr_params_are_supported(#[case] name: &str, #[case] value: Value) { - assert_eq!(mapped_params(json!({name:value.clone()}))[name], value); - } - - #[rstest] - #[case("table_format", json!("html"))] - #[case("confidence_scores_granularity", json!("word"))] - #[case("document_annotation_prompt", json!("extract"))] - #[case("include_blocks", json!(true))] - #[case("id", json!("req-123"))] - fn map_ocr_params_forwards_new_ocr_params(#[case] name: &str, #[case] value: Value) { - assert_eq!(mapped_params(json!({name:value.clone()}))[name], value); - } - - #[rstest] - #[case("pages", json!([0, 2]))] - #[case("pages", json!("0,2-4"))] - #[case("include_image_base64", json!(true))] - #[case("image_limit", json!(2))] - #[case("image_min_size", json!(100))] - #[case("bbox_annotation_format", json!({"type":"json_schema"}))] - #[case("document_annotation_format", json!({"type":"json_schema"}))] - #[case("document_annotation_prompt", json!("extract"))] - #[case("extract_header", json!(true))] - #[case("extract_footer", json!(false))] - #[case("table_format", json!("html"))] - #[case("confidence_scores_granularity", json!("word"))] - #[case("include_blocks", json!(true))] - #[case("id", json!("req-123"))] - fn request_mapping_matches_python(#[case] name: &str, #[case] value: Value) { - let params: MistralOcrParams = - serde_json::from_value(json!({name: value.clone()})).unwrap(); - let result = - serde_json::to_value(transform_ocr_request("model", document(), ¶ms).unwrap()) - .unwrap(); - assert_eq!(result["model"], "model"); - assert_eq!(result[name], value); - } - - #[rstest] - #[case("table_format", json!("html"))] - #[case("confidence_scores_granularity", json!("word"))] - #[case("document_annotation_prompt", json!("extract"))] - #[case("id", json!("req-123"))] - #[case("extract_header", json!(true))] - #[case("include_blocks", json!(true))] - #[case("pages", json!([0,1]))] - fn transform_ocr_request_includes_each_optional_param( - #[case] name: &str, - #[case] value: Value, - ) { - let params: MistralOcrParams = serde_json::from_value(json!({name:value.clone()})).unwrap(); - let result = serde_json::to_value( - transform_ocr_request("mistral-ocr-latest", document(), ¶ms).unwrap(), - ) - .unwrap(); - assert_eq!(result[name], value); - assert_eq!(result["model"], "mistral-ocr-latest"); - } - - #[rstest] - fn transform_ocr_request_includes_multiple_new_params() { - let params: MistralOcrParams = serde_json::from_value(json!({ - "table_format":"html", - "confidence_scores_granularity":"page", - "extract_header":true - })) - .unwrap(); - let result = serde_json::to_value( - transform_ocr_request("mistral-ocr-latest", document(), ¶ms).unwrap(), - ) - .unwrap(); - assert_eq!(result["table_format"], "html"); - assert_eq!(result["confidence_scores_granularity"], "page"); - assert_eq!(result["extract_header"], true); - } - - #[rstest] - fn transform_ocr_response_preserves_blocks_and_confidence_scores() { - let response: MistralOcrResponse = serde_json::from_value(json!({ - "pages":[{ - "index":0, - "markdown":"hello", - "images":[{"id":"img-0","image_base64":"data:image/png;base64,AA=="}], - "dimensions":{"width":612,"height":792,"dpi":72}, - "blocks":[{"type":"title","bbox":{"x":1},"confidence_scores":{"mean":0.98}}], - "confidence_scores":{"average_page_confidence_score":0.99,"minimum_page_confidence_score":0.97} - }], - "model":"returned-model", - "document_annotation":"{\"language\":\"en\"}", - "usage_info":{"pages_processed":1} - })) - .unwrap(); - let result = transform_ocr_response("model", response) - .unwrap() - .into_json(); - assert_eq!(result["pages"][0]["blocks"][0]["type"], "title"); - assert_eq!(result["pages"][0]["blocks"][0]["bbox"]["x"], 1); - assert_eq!( - result["pages"][0]["blocks"][0]["confidence_scores"]["mean"], - 0.98 - ); - assert_eq!( - result["pages"][0]["confidence_scores"]["average_page_confidence_score"], - 0.99 - ); - assert_eq!(result["pages"][0]["images"][0]["id"], "img-0"); - assert_eq!(result["pages"][0]["dimensions"]["dpi"], 72); - assert_eq!(result["model"], "returned-model"); - assert_eq!(result["document_annotation"], "{\"language\":\"en\"}"); - assert_eq!(result["usage_info"]["pages_processed"], 1); - } - - #[rstest] - fn transform_ocr_response_preserves_ocr4_page_fields() { - let page = json!({ - "index":0, - "markdown":"table page", - "tables":[{"rows":2,"cols":3}], - "hyperlinks":["https://example.com"], - "header":"header", - "footer":"footer" - }); - let response: MistralOcrResponse = - serde_json::from_value(json!({"pages":[page.clone()]})).unwrap(); - let result = transform_ocr_response("model", response) - .unwrap() - .into_json(); - assert_eq!(result["pages"][0], page); - } -} diff --git a/litellm-rust/crates/core/src/ocr/codecs/mistral/types.rs b/litellm-rust/crates/core/src/ocr/codecs/mistral/types.rs deleted file mode 100644 index e0bc8a267d2..00000000000 --- a/litellm-rust/crates/core/src/ocr/codecs/mistral/types.rs +++ /dev/null @@ -1,60 +0,0 @@ -use serde::{Deserialize, Serialize}; -use serde_json::{Map, Value}; - -use crate::ocr::types::OcrDocument; - -#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)] -#[serde(untagged)] -pub(crate) enum MistralOcrPages { - Range(String), - Indices(Vec), -} - -#[derive(Clone, Debug, Default, PartialEq, Serialize, Deserialize)] -pub(crate) struct MistralOcrParams { - #[serde(skip_serializing_if = "Option::is_none")] - pub pages: Option, - #[serde(skip_serializing_if = "Option::is_none")] - pub include_image_base64: Option, - #[serde(skip_serializing_if = "Option::is_none")] - pub image_limit: Option, - #[serde(skip_serializing_if = "Option::is_none")] - pub image_min_size: Option, - #[serde(skip_serializing_if = "Option::is_none")] - pub bbox_annotation_format: Option, - #[serde(skip_serializing_if = "Option::is_none")] - pub document_annotation_format: Option, - #[serde(skip_serializing_if = "Option::is_none")] - pub document_annotation_prompt: Option, - #[serde(skip_serializing_if = "Option::is_none")] - pub extract_header: Option, - #[serde(skip_serializing_if = "Option::is_none")] - pub extract_footer: Option, - #[serde(skip_serializing_if = "Option::is_none")] - pub table_format: Option, - #[serde(skip_serializing_if = "Option::is_none")] - pub confidence_scores_granularity: Option, - #[serde(skip_serializing_if = "Option::is_none")] - pub include_blocks: Option, - #[serde(skip_serializing_if = "Option::is_none")] - pub id: Option, -} - -#[derive(Clone, Debug, Serialize, Deserialize)] -pub(crate) struct MistralOcrRequest { - pub model: String, - pub document: OcrDocument, - #[serde(flatten)] - pub params: MistralOcrParams, -} - -#[derive(Clone, Debug, Default, Deserialize)] -pub(crate) struct MistralOcrResponse { - #[serde(default)] - pub pages: Vec, - pub model: Option, - pub document_annotation: Option, - pub usage_info: Option, - #[serde(flatten)] - pub extra_fields: Map, -} diff --git a/litellm-rust/crates/core/src/ocr/codecs/mod.rs b/litellm-rust/crates/core/src/ocr/codecs/mod.rs deleted file mode 100644 index 639b985b9ae..00000000000 --- a/litellm-rust/crates/core/src/ocr/codecs/mod.rs +++ /dev/null @@ -1,5 +0,0 @@ -pub(crate) mod cohere; -pub(crate) mod deepseek; -pub(crate) mod document_intelligence; -pub(crate) mod mistral; -pub(crate) mod reducto; diff --git a/litellm-rust/crates/core/src/ocr/codecs/reducto/mod.rs b/litellm-rust/crates/core/src/ocr/codecs/reducto/mod.rs deleted file mode 100644 index 3fff40451c6..00000000000 --- a/litellm-rust/crates/core/src/ocr/codecs/reducto/mod.rs +++ /dev/null @@ -1,9 +0,0 @@ -mod transformation; -mod types; - -pub(crate) use transformation::{ - transform_legacy_ocr_request, transform_ocr_response, transform_v3_ocr_request, -}; -pub(crate) use types::{ - ReductoLegacyParams, ReductoResponse, ReductoUploadResponse, ReductoV3Params, -}; diff --git a/litellm-rust/crates/core/src/ocr/codecs/reducto/transformation.rs b/litellm-rust/crates/core/src/ocr/codecs/reducto/transformation.rs deleted file mode 100644 index f4c8338c134..00000000000 --- a/litellm-rust/crates/core/src/ocr/codecs/reducto/transformation.rs +++ /dev/null @@ -1,103 +0,0 @@ -use std::collections::BTreeMap; - -use serde_json::{Value, json}; - -use super::types::*; -use crate::ocr::error::{OcrRequestError, OcrResponseError}; -use crate::ocr::types::{LiteLLMOcrResponse, OcrDocument}; - -pub(crate) fn transform_v3_ocr_request( - _model: &str, - document: OcrDocument, - params: &ReductoV3Params, -) -> Result { - Ok(ReductoV3Request { - input: document.source().to_string(), - params: params.clone(), - }) -} - -pub(crate) fn transform_legacy_ocr_request( - _model: &str, - document: OcrDocument, - params: &ReductoLegacyParams, -) -> Result { - Ok(ReductoLegacyRequest { - document_url: document.source().to_string(), - options: params.enhance.as_ref().map(|_| params.clone()), - }) -} - -pub(crate) fn transform_ocr_response( - model: &str, - response: ReductoResponse, -) -> Result { - let result = match response.result { - Some(result) => result.unwrap_or_default(), - None => ReductoResult { - chunks: response.chunks, - }, - }; - let usage = response.usage.unwrap_or_default(); - Ok(LiteLLMOcrResponse { - pages: build_pages(result.chunks.unwrap_or_default()), - model: model.to_string(), - document_annotation: None, - usage_info: Some(json!({ - "pages_processed": usage.num_pages, - "credits": usage.credits, - })), - object: "ocr".to_string(), - extra_fields: serde_json::Map::new(), - provider_native_response: None, - }) -} - -fn build_pages(chunks: Vec) -> Vec { - let blocks_by_page = chunks - .iter() - .flat_map(|chunk| chunk.blocks.iter().flatten()) - .filter_map(|block| block.bbox.as_ref()?.page.map(|page| (page, block))) - .fold( - BTreeMap::>::new(), - |mut pages, (page, block)| { - pages.entry(page).or_default().push(block); - pages - }, - ); - if blocks_by_page.is_empty() { - let markdown = join_content(chunks.iter().map(|chunk| chunk.content.as_deref())); - return if markdown.is_empty() { - Vec::new() - } else { - vec![page(0, markdown, None)] - }; - } - blocks_by_page - .into_iter() - .map(|(index, blocks)| { - let markdown = join_content(blocks.iter().map(|block| block.content.as_deref())); - page( - index.saturating_sub(1).max(0), - markdown, - Some(json!(blocks)), - ) - }) - .collect() -} - -fn join_content<'a>(content: impl Iterator>) -> String { - content - .flatten() - .filter(|text| !text.is_empty()) - .collect::>() - .join("\n\n") -} - -fn page(index: i64, markdown: String, blocks: Option) -> Value { - let mut result = json!({"index":index,"markdown":markdown,"images":null}); - if let (Value::Object(fields), Some(blocks)) = (&mut result, blocks) { - fields.insert("blocks".into(), blocks); - } - result -} diff --git a/litellm-rust/crates/core/src/ocr/codecs/reducto/types.rs b/litellm-rust/crates/core/src/ocr/codecs/reducto/types.rs deleted file mode 100644 index c03720cc8ae..00000000000 --- a/litellm-rust/crates/core/src/ocr/codecs/reducto/types.rs +++ /dev/null @@ -1,128 +0,0 @@ -use serde::{Deserialize, Deserializer, Serialize}; -use serde_json::{Map, Value}; - -#[derive(Clone, Debug, Default, Serialize, Deserialize)] -pub(crate) struct ReductoV3Params { - #[serde(skip_serializing_if = "Option::is_none")] - pub formatting: Option>, - #[serde(skip_serializing_if = "Option::is_none")] - pub retrieval: Option>, - #[serde(skip_serializing_if = "Option::is_none")] - pub settings: Option>, -} - -#[derive(Clone, Debug, Default, Serialize, Deserialize)] -pub(crate) struct ReductoLegacyParams { - #[serde(skip_serializing_if = "Option::is_none")] - pub enhance: Option>, -} - -#[derive(Clone, Debug, Serialize, Deserialize)] -pub(crate) struct ReductoV3Request { - pub input: String, - #[serde(flatten)] - pub params: ReductoV3Params, -} - -#[derive(Clone, Debug, Serialize, Deserialize)] -pub(crate) struct ReductoLegacyRequest { - pub document_url: String, - #[serde(skip_serializing_if = "Option::is_none")] - pub options: Option, -} - -#[derive(Deserialize)] -pub(crate) struct ReductoUploadResponse { - pub file_id: Option, -} - -#[derive(Clone, Debug, Deserialize)] -pub(crate) struct ReductoResponse { - #[serde(default, deserialize_with = "present_nullable")] - pub result: Option>, - pub usage: Option, - #[serde(default)] - pub chunks: Option>, -} - -fn present_nullable<'de, D: Deserializer<'de>, T: Deserialize<'de>>( - deserializer: D, -) -> Result>, D::Error> { - Option::::deserialize(deserializer).map(Some) -} - -#[derive(Clone, Debug, Default, Deserialize)] -pub(crate) struct ReductoResult { - pub chunks: Option>, -} - -#[derive(Clone, Debug, Default, Deserialize)] -pub(crate) struct ReductoUsage { - #[serde(default, deserialize_with = "optional_i64")] - pub num_pages: Option, - #[serde(default, deserialize_with = "optional_f64")] - pub credits: Option, -} - -#[derive(Clone, Debug, Deserialize)] -pub(crate) struct ReductoChunk { - pub content: Option, - pub blocks: Option>, -} - -#[derive(Clone, Debug, Serialize, Deserialize)] -pub(crate) struct ReductoBlock { - #[serde(skip_serializing_if = "Option::is_none")] - pub content: Option, - #[serde(skip_serializing_if = "Option::is_none")] - pub bbox: Option, - #[serde(flatten)] - pub extra_fields: Map, -} - -#[derive(Clone, Debug, Serialize, Deserialize)] -pub(crate) struct ReductoBoundingBox { - #[serde(default, deserialize_with = "optional_i64")] - pub page: Option, - #[serde(flatten)] - pub extra_fields: Map, -} - -fn optional_i64<'de, D: Deserializer<'de>>(deserializer: D) -> Result, D::Error> { - match Option::::deserialize(deserializer)? { - None | Some(Value::Null) => Ok(None), - Some(Value::Number(number)) => number - .as_i64() - .or_else(|| number.as_f64().and_then(checked_truncated_i64)) - .map(Some) - .ok_or_else(|| serde::de::Error::custom("expected an integer")), - Some(Value::String(value)) => value - .trim() - .parse::() - .map(Some) - .map_err(|_| serde::de::Error::custom("expected an integer")), - Some(Value::Bool(value)) => Ok(Some(i64::from(value))), - Some(_) => Ok(None), - } -} - -fn optional_f64<'de, D: Deserializer<'de>>(deserializer: D) -> Result, D::Error> { - match Option::::deserialize(deserializer)? { - None | Some(Value::Null) => Ok(None), - Some(Value::Number(number)) => number - .as_f64() - .map(Some) - .ok_or_else(|| serde::de::Error::custom("expected a number")), - Some(Value::String(value)) => value - .trim() - .parse::() - .map(Some) - .map_err(|_| serde::de::Error::custom("expected a number")), - Some(_) => Ok(None), - } -} - -fn checked_truncated_i64(value: f64) -> Option { - (value.is_finite() && value >= i64::MIN as f64 && value <= i64::MAX as f64) - .then(|| value.trunc() as i64) -} diff --git a/litellm-rust/crates/core/src/ocr/document.rs b/litellm-rust/crates/core/src/ocr/document.rs index 1b3d2dada44..fbb54f0bbd1 100644 --- a/litellm-rust/crates/core/src/ocr/document.rs +++ b/litellm-rust/crates/core/src/ocr/document.rs @@ -5,9 +5,11 @@ use base64::{Engine, engine::general_purpose::STANDARD}; use data_url::mime::Mime; use data_url::{DataUrl, DataUrlError, forgiving_base64::DecodeError}; use reqwest::Url; -use serde_json::Map; +use std::collections::BTreeMap as Map; -use super::error::{OcrError, OcrRequestError, OcrResponseError}; +use super::Error as OcrError; +use super::Error as OcrRequestError; +use super::Error as OcrResponseError; use super::types::{OcrConnection, OcrDocument, OcrDocumentInput}; use crate::constants::{OCR_INLINE_MAX_BYTES, OCR_MAX_FETCH_REDIRECTS}; use crate::media::Error as MediaError; @@ -47,11 +49,10 @@ pub fn read_path_document( }) .map_err(|source| super::Error::FileRead { path: path.to_owned(), - kind: source.kind(), - message: source.to_string(), + source: std::sync::Arc::new(source), })?; let name = path.file_name().map(|name| name.to_string_lossy()); - Ok(encode_file_document(&bytes, name.as_deref(), mime_type)?) + encode_file_document(&bytes, name.as_deref(), mime_type) } pub fn encode_file_document( @@ -164,7 +165,7 @@ pub(crate) async fn inline_remote_document( connection: &OcrConnection, ) -> Result { let source = document.source(); - if !source.starts_with("http://") && !source.starts_with("https://") { + if !document.is_remote() { validate_inline_document(&document)?; return Ok(document); } @@ -193,12 +194,12 @@ pub(crate) async fn inline_remote_document( fn map_media_error(error: MediaError) -> OcrError { match error { - MediaError::BlockedUrl => OcrRequestError::BlockedDocumentUrl.into(), - MediaError::DownloadDisabled => OcrRequestError::DownloadDisabled.into(), - MediaError::DownloadTooLarge => OcrRequestError::DownloadTooLarge.into(), - MediaError::TooManyRedirects => OcrRequestError::TooManyRedirects.into(), - MediaError::MissingRedirectLocation => OcrResponseError::MissingRedirectLocation.into(), - MediaError::InvalidRedirect => OcrResponseError::InvalidRedirect.into(), + MediaError::BlockedUrl => OcrRequestError::BlockedDocumentUrl, + MediaError::DownloadDisabled => OcrRequestError::DownloadDisabled, + MediaError::DownloadTooLarge => OcrRequestError::DownloadTooLarge, + MediaError::TooManyRedirects => OcrRequestError::TooManyRedirects, + MediaError::MissingRedirectLocation => OcrResponseError::MissingRedirectLocation, + MediaError::InvalidRedirect => OcrResponseError::InvalidRedirect, MediaError::Http(status) => TransportError::Http { status, body: "OCR document download failed".into(), @@ -216,7 +217,7 @@ fn map_media_error(error: MediaError) -> OcrError { #[cfg(test)] mod tests { use super::*; - use serde_json::Map; + use std::collections::BTreeMap as Map; fn document(source: &str) -> OcrDocument { OcrDocument::DocumentUrl { @@ -286,17 +287,17 @@ mod tests { document("data:application/pdf;base64,YWJj") ); std::fs::write(&path, vec![b'a'; OCR_INLINE_MAX_BYTES + 1]).unwrap(); - assert_eq!( + assert!(matches!( prepare_document(OcrDocumentInput::Path { path: path.clone(), mime_type: None, }), - Err(OcrRequestError::InlineDocumentTooLarge.into()) - ); + Err(OcrRequestError::InlineDocumentTooLarge) + )); std::fs::remove_dir_all(&dir).unwrap(); let missing = dir.join("missing.pdf"); - let Err(super::super::Error::FileRead { path, kind, .. }) = + let Err(super::super::Error::FileRead { path, source, .. }) = prepare_document(OcrDocumentInput::Path { path: missing.clone(), mime_type: None, @@ -305,7 +306,7 @@ mod tests { panic!("missing paths must surface a file read error"); }; assert_eq!(path, missing); - assert_eq!(kind, std::io::ErrorKind::NotFound); + assert_eq!(source.kind(), std::io::ErrorKind::NotFound); } #[test] @@ -325,10 +326,10 @@ mod tests { #[test] fn file_encoding_enforces_decoded_size_limit() { let bytes = vec![b'a'; OCR_INLINE_MAX_BYTES + 1]; - assert_eq!( + assert!(matches!( encode_file_document(&bytes, None, None), Err(OcrRequestError::InlineDocumentTooLarge) - ); + )); let document = encode_file_document(&bytes[..OCR_INLINE_MAX_BYTES], None, None).unwrap(); let inline = InlineDocument::parse(document.source()).unwrap().unwrap(); assert_eq!( @@ -359,10 +360,10 @@ mod tests { ] { let inline = InlineDocument::parse(source).unwrap().unwrap(); assert_eq!(inline.decode(expected.len()).unwrap(), expected); - assert_eq!( + assert!(matches!( inline.decode(expected.len() - 1), Err(OcrRequestError::InlineDocumentTooLarge) - ); + )); } } @@ -427,7 +428,7 @@ mod tests { client.document_fetcher(), OcrDocument::ImageUrl { image_url: format!("http://{address}/image"), - extra_fields: Map::from_iter([("detail".into(), serde_json::json!("high"))]), + extra_fields: Map::from_iter([("detail".into(), "high".into())]), }, &OcrConnection::default(), ) @@ -439,7 +440,7 @@ mod tests { converted, OcrDocument::ImageUrl { image_url: "data:image/png;base64,YWJj".into(), - extra_fields: Map::from_iter([("detail".into(), serde_json::json!("high"))]), + extra_fields: Map::from_iter([("detail".into(), "high".into())]), } ); assert!(!request.to_ascii_lowercase().contains("authorization")); diff --git a/litellm-rust/crates/core/src/ocr/error.rs b/litellm-rust/crates/core/src/ocr/error.rs index 0c92b511a38..7685875709e 100644 --- a/litellm-rust/crates/core/src/ocr/error.rs +++ b/litellm-rust/crates/core/src/ocr/error.rs @@ -1,117 +1,21 @@ -use thiserror::Error; - -use crate::transport::Error as TransportError; - -#[derive(Clone, Debug, Error, PartialEq, Eq)] +#[derive(Clone, Debug, thiserror::Error)] pub enum Error { - #[error("expected {expected}, got {actual}")] - InvalidType { - expected: &'static str, - actual: &'static str, + #[error("upstream OCR error ({status}): {body}")] + Provider { + status: u16, + body: String, + headers: Vec<(String, String)>, }, - #[error("missing required field: {0}")] - MissingField(&'static str), - #[error("Document URL is required")] - MissingDocumentUrl, - #[error("invalid response: {0}")] - InvalidResponse(String), - #[error("invalid provider: {0}")] - InvalidProvider(String), - #[error("invalid request: {0}")] - InvalidRequest(String), - #[error("{0}")] - Auth(String), - #[error( - "Missing {provider} API Key - A call is being made to {provider} but no key is set either in the environment variables or via params" - )] - MissingApiKey { provider: &'static str }, - #[error( - "invalid authentication configuration: Missing Azure AI credentials - set AZURE_AI_API_KEY or configure Entra ID" - )] - MissingAzureAiCredentials, - #[error( - "invalid authentication configuration: Missing Azure Document Intelligence credentials - set AZURE_DOCUMENT_INTELLIGENCE_API_KEY or configure Entra ID" - )] - MissingAzureDocumentIntelligenceCredentials, - #[error( - "Missing REDUCTO_API_KEY - set it in the environment or pass api_key to litellm.ocr()/litellm.aocr()" - )] - MissingReductoApiKey, - #[error("upstream request failed with status {status}: {body}")] - Http { status: u16, body: String }, - #[error("upstream network error: {0}")] - Network(String), - /// The provider was never reached: DNS, TCP, TLS or proxy setup failed - /// before any byte of the request went out. Nothing was billed, so a host - /// that keeps a reference implementation can serve the request itself. - /// A timeout is deliberately not this, since the provider may have received - /// and answered the request already. - #[error("could not reach the provider: {0}")] - Connect(String), - #[error("routing error: {0}")] - Routing(String), - #[error("Failed to read OCR file {}: {message}", path.display())] - FileRead { - path: std::path::PathBuf, - kind: std::io::ErrorKind, - message: String, - }, - /// The request is outside the surface this route covers in Rust. Hosts that - /// keep a reference implementation treat this as "fall back", not "fail". - #[error("unsupported by the rust path: {0}")] - Unsupported(&'static str), -} - -impl Error { - pub const fn http_status_code(&self) -> Option { - match self { - Self::InvalidRequest(_) => Some(400), - Self::MissingDocumentUrl => Some(500), - Self::Http { status, .. } => Some(*status), - _ => None, - } - } -} - -impl From for Error { - fn from(error: OcrRequestError) -> Self { - match error { - OcrRequestError::MissingField(field) => Self::MissingField(field), - OcrRequestError::MissingDocumentUrl => Self::MissingDocumentUrl, - error => Self::InvalidRequest(error.to_string()), - } - } -} - -impl From for Error { - fn from(error: OcrResponseError) -> Self { - Self::InvalidResponse(error.to_string()) - } -} - -impl From for Error { - fn from(error: TransportError) -> Self { - match error { - TransportError::Http { status, body } => Self::Http { status, body }, - TransportError::Network(message) => Self::Network(message), - TransportError::Connect(message) => Self::Connect(message), - } - } -} - -impl From for Error { - fn from(error: litellm_auth::Error) -> Self { - match error { - litellm_auth::Error::MissingApiKey { provider, .. } => Self::MissingApiKey { provider }, - error => Self::Auth(error.to_string()), - } - } -} - -#[derive(Debug, Clone, PartialEq, Eq, Error)] -pub enum OcrRequestError { #[error("File is empty or could not be read")] EmptyFile, + #[error("Failed to read OCR file {}: {source}", path.display())] + FileRead { + path: std::path::PathBuf, + #[source] + source: std::sync::Arc, + }, + #[error("OCR document preparation task failed: {0}")] + DocumentTask(#[source] std::sync::Arc), #[error("Invalid MIME type: {0}")] InvalidMimeType(String), #[error( @@ -148,10 +52,6 @@ pub enum OcrRequestError { Features, #[error("OCR model cannot be a dot segment")] DotModel, -} - -#[derive(Debug, Clone, PartialEq, Eq, Error)] -pub enum OcrResponseError { #[error("OCR response exceeds the size limit of {limit} bytes")] TooLarge { limit: usize }, #[error("invalid OCR response field: {path}")] @@ -166,40 +66,101 @@ pub enum OcrResponseError { OperationStatus(String), #[error("OCR response numeric value is out of range: {0}")] NumericRange(&'static str), -} - -#[derive(Debug, Clone, PartialEq, Eq, Error)] -pub enum OcrPollingError { #[error("OCR accepted response is missing a valid operation-location")] PollLocation, #[error("OCR operation-location must use the submission origin without credentials")] PollOrigin, #[error("OCR polling timed out")] PollTimeout, + #[error("unsupported by the rust path: {0}")] + Unsupported(&'static str), + #[error("invalid provider: {0}")] + InvalidProvider(String), + #[error("invalid request: {0}")] + InvalidRequest(String), + #[error("invalid response: {0}")] + InvalidResponse(String), + #[error( + "invalid authentication configuration: Missing Azure AI credentials - set AZURE_AI_API_KEY or configure Entra ID" + )] + MissingAzureAiCredentials, + #[error( + "invalid authentication configuration: Missing Azure Document Intelligence credentials - set AZURE_DOCUMENT_INTELLIGENCE_API_KEY or configure Entra ID" + )] + MissingAzureDocumentIntelligenceCredentials, + #[error( + "Missing REDUCTO_API_KEY - set it in the environment or pass api_key to litellm.ocr()/litellm.aocr()" + )] + MissingReductoApiKey, + #[error(transparent)] + Auth(#[from] litellm_auth::Error), + #[error(transparent)] + Transport(#[from] crate::transport::Error), + #[error(transparent)] + Params(#[from] crate::params::Error), + #[error(transparent)] + Headers(#[from] crate::http_utils::HeaderError), } -#[derive(Debug, Error)] -pub enum OcrError { - #[error("{0}")] - Request(#[from] OcrRequestError), - #[error("{0}")] - Response(#[from] OcrResponseError), - #[error("{0}")] - Transport(#[from] TransportError), - #[error("{0}")] - Polling(#[from] OcrPollingError), - #[error("{0}")] - Public(#[from] Error), -} - -impl From for Error { - fn from(error: OcrError) -> Self { - match error { - OcrError::Request(error) => error.into(), - OcrError::Response(error) => error.into(), - OcrError::Transport(error) => error.into(), - OcrError::Polling(error) => Error::InvalidResponse(error.to_string()), - OcrError::Public(error) => error, +impl From for Error { + fn from(error: crate::call_arguments::ArgumentError) -> Self { + Self::RequestField { + path: format!("optional_params.{}", error.path), } } } + +impl Error { + pub fn http_status_code(&self) -> Option { + match self { + Self::MissingDocumentUrl => Some(500), + Self::Provider { status, .. } + | Self::Transport(crate::transport::Error::Http { status, .. }) => Some(*status), + error if error.is_request() => Some(400), + _ => None, + } + } + + pub fn is_request(&self) -> bool { + matches!( + self, + Self::EmptyFile + | Self::InvalidMimeType(_) + | Self::CohereImageOnly + | Self::RequestFormat + | Self::RequestField { .. } + | Self::MissingField(_) + | Self::MissingDocumentUrl + | Self::InvalidDataUri + | Self::ReductoSource + | Self::InlineDocumentTooLarge + | Self::BlockedDocumentUrl + | Self::DownloadDisabled + | Self::DownloadTooLarge + | Self::TooManyRedirects + | Self::Pages(_) + | Self::Features + | Self::DotModel + | Self::InvalidRequest(_) + | Self::Params(_) + | Self::Headers(_) + ) + } + + pub fn is_response(&self) -> bool { + matches!( + self, + Self::TooLarge { .. } + | Self::ResponseField { .. } + | Self::EmptyContent + | Self::MissingRedirectLocation + | Self::InvalidRedirect + | Self::OperationStatus(_) + | Self::NumericRange(_) + | Self::PollLocation + | Self::PollOrigin + | Self::PollTimeout + | Self::InvalidResponse(_) + ) + } +} diff --git a/litellm-rust/crates/core/src/ocr/handler.rs b/litellm-rust/crates/core/src/ocr/handler.rs index 1ec02f3b622..7e42111da0a 100644 --- a/litellm-rust/crates/core/src/ocr/handler.rs +++ b/litellm-rust/crates/core/src/ocr/handler.rs @@ -1,21 +1,20 @@ -use super::OcrClient; -use super::adapters::OcrAdapter; -use super::hooks::{OcrHooks, OcrLifecycleHooks, OcrPostCallRequest}; -use super::registry::OcrAdapterKind; -use super::types::{LiteLLMOcrRequest, LiteLLMOcrResponse}; -use crate::call_lifecycle::{CallLifecycle, CallLifecycleContext}; -use crate::ocr::Error; use std::sync::Arc; +use super::OcrClient; +use super::hooks::{OcrHooks, OcrLifecycleHooks, OcrPostCallRequest}; +use super::types::{LiteLLMOcrResponse, PreparedOcrRequest, ResolvedOcrRequest}; +use crate::call_lifecycle::{CallLifecycle, CallLifecycleContext}; +use crate::llms::base_llm::ocr::transformation::OcrResponseContext; + pub(crate) async fn perform_ocr_request( client: &OcrClient, - request: LiteLLMOcrRequest, -) -> Result { + request: ResolvedOcrRequest, +) -> Result { request.response_format()?; let context = CallLifecycleContext::new( "ocr", request.model.clone(), - request.adapter.provider().as_str(), + request.provider_name(), request .litellm_call_id .clone() @@ -30,31 +29,24 @@ pub(crate) async fn perform_ocr_request( PreparedOcrCall::prepare(client.clone(), request) .await? .execute() - .await? - .normalize() + .await }) .await } pub(crate) struct PreparedOcrCall { client: OcrClient, - request: LiteLLMOcrRequest, + request: PreparedOcrRequest, http: reqwest::Request, } impl PreparedOcrCall { pub(crate) async fn prepare( client: OcrClient, - request: LiteLLMOcrRequest, - ) -> Result { - macro_rules! prepare_adapter { - ($( $variant:ident, $adapter:ty, $instance:expr, $provider:ident; )+) => { - match request.adapter { - $( OcrAdapterKind::$variant => $instance.prepare_request(&request, &client).await?, )+ - } - }; - } - let http = super::adapters::for_each_ocr_adapter!(prepare_adapter); + request: ResolvedOcrRequest, + ) -> Result { + let request = super::prepare::prepare_request(request); + let http = request.config.prepare_request(&request, &client).await?; Ok(Self { client, request, @@ -62,33 +54,54 @@ impl PreparedOcrCall { }) } - pub(crate) async fn execute(self) -> Result { + pub(crate) async fn execute(self) -> Result { let url = self.http.url().to_string(); let headers = request_headers(&self.http)?; - let response = crate::http_utils::http_request(reqwest::RequestBuilder::from_parts( - self.client.provider_http().clone(), - self.http, - )) - .await - .map_err(super::client::transport_error)?; - macro_rules! read_adapter { - ($( $variant:ident, $adapter:ty, $instance:expr, $provider:ident; )+) => { - match self.request.adapter { - $( OcrAdapterKind::$variant => { - let decoded = $instance.read_response(&self.client, response, &url, &headers, &self.request).await?; - Ok(OcrProviderResponse { - request: self.request, - data: OcrProviderData::$variant(decoded), - }) - }, )+ + let response = + 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"), }; } - super::adapters::for_each_ocr_adapter!(read_adapter) + let model = &self.request.model; + let context = OcrResponseContext { + client: &self.client, + connection: &self.request.connection, + hooks: &self.request.hooks, + request_format: self.request.response_format()?, + url: &url, + headers: &headers, + }; + self.request + .config + .async_transform_ocr_response(model, response, context) + .await } } -fn request_headers(request: &reqwest::Request) -> Result, Error> { +fn request_headers(request: &reqwest::Request) -> Result, super::Error> { request .headers() .iter() @@ -96,44 +109,17 @@ fn request_headers(request: &reqwest::Request) -> Result, value .to_str() .map(|value| (name.to_string(), value.to_string())) - .map_err(|_| super::error::OcrRequestError::RequestField { + .map_err(|_| super::Error::RequestField { path: "headers".into(), }) - .map_err(Error::from) }) .collect() } -macro_rules! provider_data { - ($( $variant:ident, $adapter:ty, $instance:expr, $provider:ident; )+) => { - enum OcrProviderData { - $( $variant(super::wire::DecodedOcrResponse<<$adapter as OcrAdapter>::ProviderResponse>), )+ - } - - impl OcrProviderResponse { - pub(crate) fn normalize(self) -> Result { - match self.data { - $( OcrProviderData::$variant(decoded) => { - let response = $instance.transform_ocr_response(&self.request, decoded.data)?; - Ok(LiteLLMOcrResponse { provider_native_response: decoded.native, ..response }) - }, )+ - } - } - } - }; -} - -pub(crate) struct OcrProviderResponse { - request: LiteLLMOcrRequest, - data: OcrProviderData, -} - -pub(crate) async fn post_call(hooks: &Arc, bytes: &[u8]) -> Result<(), Error> { +pub(crate) async fn post_call(hooks: &Arc, bytes: &[u8]) -> Result<(), super::Error> { let original_response = serde_json::Value::String(String::from_utf8_lossy(bytes).into_owned()); hooks .post_call(OcrPostCallRequest { original_response }) .await?; Ok(()) } - -super::adapters::for_each_ocr_adapter!(provider_data); diff --git a/litellm-rust/crates/core/src/ocr/hooks.rs b/litellm-rust/crates/core/src/ocr/hooks.rs index 1d8c5953fa7..8a14afb7c50 100644 --- a/litellm-rust/crates/core/src/ocr/hooks.rs +++ b/litellm-rust/crates/core/src/ocr/hooks.rs @@ -2,7 +2,7 @@ use std::future::Future; use std::pin::Pin; use std::sync::Arc; -use super::types::{LiteLLMOcrRequest, LiteLLMOcrResponse, OcrDocument}; +use super::types::{LiteLLMOcrRequest, LiteLLMOcrResponse, OcrDocument, ResolvedOcrRequest}; use crate::call_lifecycle::{CallLifecycleContext, CallLifecycleHooks, CallLifecycleTiming}; use crate::ocr::Error; use serde::Serialize; @@ -23,6 +23,7 @@ pub struct OcrPreCallRequest { pub struct OcrDuringCallRequest { pub model: String, pub custom_llm_provider: String, + pub api_key: Option, pub url: String, pub headers: Vec<(String, String)>, pub body: Value, @@ -77,19 +78,19 @@ pub(crate) struct OcrLifecycleHooks { pub provider_name: String, } -impl CallLifecycleHooks +impl CallLifecycleHooks for OcrLifecycleHooks { type Error = crate::ocr::Error; - type PreCallFuture<'a> = OcrHookFuture<'a, LiteLLMOcrRequest>; - type DuringCallFuture<'a> = OcrHookFuture<'a, LiteLLMOcrRequest>; + type PreCallFuture<'a> = OcrHookFuture<'a, ResolvedOcrRequest>; + type DuringCallFuture<'a> = OcrHookFuture<'a, ResolvedOcrRequest>; type SuccessFuture<'a> = OcrLogFuture<'a>; type FailureFuture<'a> = OcrLogFuture<'a>; fn async_pre_call_hook<'a>( &'a self, _context: &'a CallLifecycleContext, - request: LiteLLMOcrRequest, + request: ResolvedOcrRequest, ) -> Self::PreCallFuture<'a> { Box::pin(async move { if !self.hooks.intercepts_requests() { @@ -101,18 +102,17 @@ impl CallLifecycleHooks( &'a self, _context: &'a CallLifecycleContext, - request: LiteLLMOcrRequest, + request: ResolvedOcrRequest, ) -> Self::DuringCallFuture<'a> { Box::pin(async move { Ok(request) }) } diff --git a/litellm-rust/crates/core/src/ocr/json.rs b/litellm-rust/crates/core/src/ocr/json.rs new file mode 100644 index 00000000000..d4651838a2d --- /dev/null +++ b/litellm-rust/crates/core/src/ocr/json.rs @@ -0,0 +1,62 @@ +use serde::de::{DeserializeOwned, IntoDeserializer}; +use serde_json::{Map, Value}; + +#[derive(Debug)] +pub struct DecodedOcrResponse { + pub data: T, + pub native: Option>, + pub text: String, +} + +pub(crate) fn decode_request_value( + value: Value, + prefix: &str, +) -> Result { + serde_path_to_error::deserialize(value.into_deserializer()).map_err(|error| { + crate::ocr::Error::RequestField { + path: format!("{prefix}.{}", error.path()), + } + }) +} + +pub(crate) fn decode_response_value( + value: Value, + prefix: &str, +) -> Result { + serde_path_to_error::deserialize(value.into_deserializer()).map_err(|error| { + crate::ocr::Error::ResponseField { + path: format!("{prefix}.{}", error.path()), + } + }) +} + +pub(crate) fn decode_response( + bytes: &[u8], + native: bool, +) -> Result, crate::ocr::Error> { + let mut deserializer = serde_json::Deserializer::from_slice(bytes); + let data = serde_path_to_error::deserialize(&mut deserializer).map_err(|error| { + crate::ocr::Error::ResponseField { + path: error.path().to_string(), + } + })?; + deserializer + .end() + .map_err(|_| crate::ocr::Error::ResponseField { + path: "response".into(), + })?; + let native = if native { + Some( + serde_json::from_slice(bytes).map_err(|_| crate::ocr::Error::ResponseField { + path: "response".into(), + })?, + ) + } else { + None + }; + Ok(DecodedOcrResponse { + data, + native, + text: String::from_utf8_lossy(bytes).into_owned(), + }) +} diff --git a/litellm-rust/crates/core/src/ocr/lifecycle.rs b/litellm-rust/crates/core/src/ocr/lifecycle.rs index 994a9698459..dee34526001 100644 --- a/litellm-rust/crates/core/src/ocr/lifecycle.rs +++ b/litellm-rust/crates/core/src/ocr/lifecycle.rs @@ -380,7 +380,7 @@ impl OcrExecution { self.execution = None; self.completed = true; result - .map_err(|error| Error::Network(format!("OCR execution task failed: {error}")))? + .map_err(|error| Error::Transport(crate::transport::Error::Network(format!("OCR execution task failed: {error}"))))? .map(OcrCallStep::Complete) } } @@ -431,7 +431,7 @@ impl OcrExecution { async fn prepare_request_document( request: LiteLLMOcrRequest, hooks: &ProtocolHooks, -) -> Result { +) -> Result { let request = match &request.document { OcrDocumentInput::HostReader { mime_type } => { let mime_type = mime_type.clone(); diff --git a/litellm-rust/crates/core/src/ocr/mod.rs b/litellm-rust/crates/core/src/ocr/mod.rs index f2e7aa4f46d..943d99c74e3 100644 --- a/litellm-rust/crates/core/src/ocr/mod.rs +++ b/litellm-rust/crates/core/src/ocr/mod.rs @@ -1,26 +1,31 @@ -mod adapters; +mod arguments; pub mod client; -mod codecs; -mod document; +pub(crate) mod document; pub mod error; pub use error::Error; -mod handler; +pub(crate) mod handler; pub mod hooks; +pub(crate) mod json; mod lifecycle; -mod prepare; -mod registry; +pub(crate) mod prepare; +mod provider_config; pub mod types; pub mod wire; +pub use arguments::{ + consumed_optional_param_names, consumed_optional_params, is_supported_request, +}; pub use client::{OcrClient, ocr}; pub use document::{encode_file_document, mime_type_for_name, read_path_document}; 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, OcrDocumentInput, - OcrFileContent, + LiteLLMOcrRequest, LiteLLMOcrResponse, OcrConnection, OcrConnectionInputs, OcrCredentialInputs, + OcrDocument, OcrDocumentInput, OcrFileContent, OcrPage, OcrPageDimensions, OcrPageImage, + OcrTransportConfig, OcrUsageInfo, }; #[cfg(test)] diff --git a/litellm-rust/crates/core/src/ocr/prepare.rs b/litellm-rust/crates/core/src/ocr/prepare.rs index 5a48206d53c..aa4ca94bf0c 100644 --- a/litellm-rust/crates/core/src/ocr/prepare.rs +++ b/litellm-rust/crates/core/src/ocr/prepare.rs @@ -1,117 +1,72 @@ -use serde::{Deserialize, Serialize, de::DeserializeOwned}; -use serde_json::{Map, Value}; +use serde::Serialize; +use serde_json::Value; use super::OcrClient; -use super::error::{OcrError, OcrRequestError}; use super::hooks::OcrDuringCallRequest; -use super::types::{LiteLLMOcrRequest, OcrDocument}; - -#[derive(Debug, Deserialize)] -pub(crate) struct ParsedProviderParams { - #[serde(flatten)] - pub known: T, - #[serde(default, flatten)] - pub extra_params: Map, -} - -pub(crate) fn _prepare_ocr_request( - request: &LiteLLMOcrRequest, -) -> Result, OcrRequestError> { - super::wire::decode_request_value( - Value::Object(request.optional_params.clone()), - "optional_params", - ) -} - -pub(crate) fn merge_extra_params( - body: &B, - extra_params: Map, -) -> Result { - let Value::Object(fields) = - serde_json::to_value(body).map_err(|_| OcrRequestError::RequestField { - path: "body".into(), - })? - else { - return Err(OcrRequestError::RequestField { - path: "body".into(), - }); - }; - let extra_body = extra_params - .get("extra_body") - .and_then(Value::as_object) - .cloned() - .unwrap_or_default() - .into_iter() - .collect::>(); - Ok(Value::Object( - fields - .into_iter() - .chain( - extra_params - .into_iter() - .filter(|(name, _)| name != "extra_body"), - ) - .chain(extra_body) - .collect(), - )) -} +use super::types::{OcrConnection, OcrDocument, PreparedOcrRequest, ResolvedOcrRequest}; pub(crate) async fn transform_request_body( client: &OcrClient, - request: &LiteLLMOcrRequest, + request: &PreparedOcrRequest, url: &str, headers: &[(String, String)], - retains_document: bool, body: B, - validate: impl FnOnce(&B) -> Result<(), OcrRequestError>, -) -> Result + validate: impl Fn(&Value) -> Result<(), super::Error>, +) -> Result where - B: Serialize + DeserializeOwned, + B: Serialize, { + let composed = crate::call_arguments::compose_body( + &request.optional_params, + &body, + request.config.get_supported_ocr_params(&request.model), + )?; + validate(&composed)?; + let retained_fields = request + .optional_params + .keys() + .filter(|name| composed.get(*name).is_some()) + .cloned() + .chain( + composed + .get("document") + .is_some() + .then(|| "document".to_string()), + ) + .collect(); let (body, headers) = if request.hooks.intercepts_requests() { - let body = serde_json::to_value(body).map_err(|_| OcrRequestError::RequestField { - path: "body".into(), - })?; - let retained_fields = request - .optional_params - .keys() - .filter(|name| body.get(*name).is_some()) - .cloned() - .chain(retains_document.then(|| "document".to_string())) - .collect(); let changed = request .hooks .during_call(OcrDuringCallRequest { model: request.model.clone(), - custom_llm_provider: request.adapter.provider().as_str().into(), + custom_llm_provider: request.provider_name().into(), + api_key: request.connection.api_key.clone(), url: url.into(), headers: headers.to_vec(), - body, + body: composed, retained_fields, }) .await?; - let body = OcrWireBody::::decode(changed.body)?; - validate(&body.body)?; - (body, changed.headers) + if !changed.body.is_object() { + return Err(super::Error::RequestField { + path: "guardrail.body".into(), + }); + } + validate(&changed.body)?; + (changed.body, changed.headers) } else { - ( - OcrWireBody { - body, - extra: Map::new(), - }, - headers.to_vec(), - ) + (composed, headers.to_vec()) }; build_http_request(client, request, url, &headers, &body) } pub(crate) fn build_http_request( client: &OcrClient, - request: &LiteLLMOcrRequest, + request: &PreparedOcrRequest, url: &str, headers: &[(String, String)], body: &B, -) -> Result { +) -> Result { let builder = client .provider_http() .post(url) @@ -120,14 +75,14 @@ pub(crate) fn build_http_request( crate::http_utils::with_headers(builder, headers, crate::http_utils::HeaderPolicy::All) .build() .map_err(crate::transport::Error::from) - .map_err(OcrError::from) + .map_err(super::Error::from) } pub(crate) async fn guardrail_document( - request: &LiteLLMOcrRequest, + request: &PreparedOcrRequest, url: &str, headers: &[(String, String)], -) -> Result<(OcrDocument, Vec<(String, String)>), OcrError> { +) -> Result<(OcrDocument, Vec<(String, String)>), super::Error> { if !request.hooks.intercepts_requests() { return Ok((request.document.clone(), headers.to_vec())); } @@ -135,80 +90,113 @@ pub(crate) async fn guardrail_document( .hooks .during_call(OcrDuringCallRequest { model: request.model.clone(), - custom_llm_provider: request.adapter.provider().as_str().into(), + custom_llm_provider: request.provider_name().into(), + api_key: request.connection.api_key.clone(), url: url.into(), headers: headers.to_vec(), body: serde_json::to_value(&request.document).map_err(|_| { - OcrRequestError::RequestField { + super::Error::RequestField { path: "document".into(), } })?, retained_fields: Vec::new(), }) .await?; - let document = super::wire::decode_request_value(changed.body, "guardrail.document")?; + let document = super::json::decode_request_value(changed.body, "guardrail.document")?; Ok((document, changed.headers)) } -#[derive(Serialize)] -struct OcrWireBody { - #[serde(flatten)] - body: B, - #[serde(flatten)] - extra: Map, -} - -impl OcrWireBody { - fn decode(value: Value) -> Result { - let body: B = super::wire::decode_request_value(value.clone(), "guardrail.body")?; - let Value::Object(fields) = value else { - return Err(OcrRequestError::RequestField { - path: "guardrail.body".into(), - }); - }; - let known = serde_json::to_value(&body).map_err(|_| OcrRequestError::RequestField { - path: "guardrail.body".into(), +pub(crate) fn body_document(body: &Value) -> Result { + let document = body + .get("document") + .and_then(Value::as_object) + .ok_or_else(|| super::Error::RequestField { + path: "body.document".into(), })?; - let extra = fields - .into_iter() - .filter(|(key, _)| known.get(key).is_none()) - .collect(); - Ok(Self { body, extra }) - } + let source = document + .iter() + .filter(|(name, _)| matches!(name.as_str(), "type" | "image_url" | "document_url")) + .map(|(name, value)| (name.clone(), value.clone())) + .collect(); + super::json::decode_request_value(Value::Object(source), "body.document") } pub(crate) fn credential_env(name: &str) -> Option { std::env::var(name).ok() } + +pub(crate) fn prepare_request(request: ResolvedOcrRequest) -> PreparedOcrRequest { + use litellm_auth::{InputSource, Sourced}; + + let credentials = request.credentials.clone(); + 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 = credentials.dynamic_api_key.or_else(|| { + credentials.api_key.clone().or_else(|| { + request + .config + .get_api_key_env_var() + .and_then(credential_env) + .map(|value| Sourced::new(value, InputSource::Environment)) + }) + }); + let dynamic_api_base = credentials.dynamic_api_base.or_else(|| { + credentials.api_base.clone().or_else(|| { + api_base_env + .and_then(credential_env) + .map(|value| Sourced::new(value, InputSource::Environment)) + }) + }); + let resolved = request + .config + .resolve_connection_params(super::types::OcrCredentialInputs { + dynamic_api_key, + dynamic_api_base, + ..credentials + }); + let transport = request.transport.clone(); + PreparedOcrRequest::new(request, OcrConnection::new(resolved, transport)) +} + #[cfg(test)] mod tests { + use crate::call_arguments::{CallArguments, compose_body, parse_options}; use serde_json::json; - use super::*; - - #[derive(Debug, Deserialize, PartialEq)] + #[derive(serde::Deserialize)] struct KnownParams { pages: Option>, } #[test] fn parsed_provider_params_separates_known_and_extra_params() { - let parsed: ParsedProviderParams = super::super::wire::decode_request_value( - json!({ - "pages": [0, 2], - "future_ocr_option": true, - "extra_body": {"provider_option": "value"} - }), - "optional_params", - ) + let arguments: CallArguments = serde_json::from_value(json!({ + "pages": [0, 2], + "future_ocr_option": true, + "extra_body": {"provider_option": "value"} + })) .unwrap(); - - assert_eq!(parsed.known.pages, Some(vec![0, 2])); - assert_eq!(parsed.extra_params["future_ocr_option"], true); + let known: KnownParams = parse_options(&arguments).unwrap(); + assert_eq!(known.pages, Some(vec![0, 2])); + assert_eq!(arguments["future_ocr_option"], true); + assert_eq!(arguments["extra_body"], json!({"provider_option": "value"})); assert_eq!( - parsed.extra_params["extra_body"], - json!({"provider_option": "value"}) + arguments + .iter() + .filter(|(name, _)| name.as_str() != "pages") + .count(), + 2 + ); + assert_eq!( + compose_body(&arguments, &json!({"pages": known.pages}), &["pages"]).unwrap(), + json!({ + "pages": [0, 2], "future_ocr_option": true, "provider_option": "value" + }) ); - assert_eq!(parsed.extra_params.len(), 2); } } diff --git a/litellm-rust/crates/core/src/ocr/provider_config.rs b/litellm-rust/crates/core/src/ocr/provider_config.rs new file mode 100644 index 00000000000..9fb89812664 --- /dev/null +++ b/litellm-rust/crates/core/src/ocr/provider_config.rs @@ -0,0 +1,411 @@ +use super::OcrClient; +use super::types::{ + LiteLLMOcrResponse, OcrCredentialInputs, OcrDocument, PreparedOcrRequest, + ResolvedOcrCredentials, +}; +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; +use crate::llms::base_llm::ocr::transformation::{BaseOcrConfig, OcrResponseContext}; +use crate::llms::cohere::ocr::transformation::CohereParseConfig; +use crate::llms::mistral::ocr::transformation::MistralOCRConfig; +use crate::llms::reducto::ocr::transformation::{ReductoParseLegacyConfig, ReductoParseV3Config}; +use crate::llms::vertex_ai::ocr::deepseek_transformation::VertexAIDeepSeekOCRConfig; +use crate::llms::vertex_ai::ocr::transformation::VertexAIOCRConfig; +use crate::providers::custom_llm_provider::{CustomLlmProvider, get_custom_llm_provider}; +use strum::{EnumString, IntoStaticStr}; + +macro_rules! dispatch_config { + ($config:expr, $method:ident($($argument:expr),* $(,)?)) => { + dispatch_config!(@arms $config, $method($($argument),*), ) + }; + ($config:expr, $method:ident($($argument:expr),* $(,)?).await) => { + dispatch_config!(@arms $config, $method($($argument),*), .await) + }; + (@arms $config:expr, $method:ident($($argument:expr),*), $($suffix:tt)*) => { + match $config { + OcrConfigKind::Cohere => CohereParseConfig.$method($($argument),*)$($suffix)*, + OcrConfigKind::Mistral => MistralOCRConfig.$method($($argument),*)$($suffix)*, + OcrConfigKind::AzureAi => AzureAIOCRConfig.$method($($argument),*)$($suffix)*, + OcrConfigKind::AzureCohere => AzureAICohereParseConfig.$method($($argument),*)$($suffix)*, + OcrConfigKind::AzureDocumentIntelligence => AzureDocumentIntelligenceOCRConfig.$method($($argument),*)$($suffix)*, + OcrConfigKind::ReductoLegacy => ReductoParseLegacyConfig.$method($($argument),*)$($suffix)*, + OcrConfigKind::ReductoV3 => ReductoParseV3Config.$method($($argument),*)$($suffix)*, + OcrConfigKind::VertexAi => VertexAIOCRConfig.$method($($argument),*)$($suffix)*, + OcrConfigKind::VertexDeepSeek => VertexAIDeepSeekOCRConfig.$method($($argument),*)$($suffix)*, + } + }; +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub(crate) enum OcrConfigKind { + Cohere, + Mistral, + AzureAi, + AzureCohere, + AzureDocumentIntelligence, + ReductoLegacy, + ReductoV3, + VertexAi, + VertexDeepSeek, +} + +impl OcrConfigKind { + pub(crate) const fn provider(self) -> OcrProvider { + match self { + Self::Cohere => OcrProvider::Cohere, + Self::Mistral => OcrProvider::Mistral, + Self::AzureAi | Self::AzureCohere | Self::AzureDocumentIntelligence => { + OcrProvider::AzureAi + } + Self::ReductoLegacy | Self::ReductoV3 => OcrProvider::Reducto, + Self::VertexAi | Self::VertexDeepSeek => OcrProvider::VertexAi, + } + } + + pub(crate) fn get_supported_ocr_params(self, model: &str) -> &'static [&'static str] { + 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, + inputs: OcrCredentialInputs, + ) -> ResolvedOcrCredentials { + dispatch_config!(self, resolve_connection_params(inputs)) + } + + 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)) + } + + pub(crate) async fn prepare_request( + self, + request: &PreparedOcrRequest, + client: &OcrClient, + ) -> Result { + dispatch_config!(self, prepare_request(request, client).await) + } + + pub(crate) async fn async_transform_ocr_response( + self, + model: &str, + raw_response: reqwest::Response, + context: OcrResponseContext<'_>, + ) -> Result { + dispatch_config!( + self, + async_transform_ocr_response(model, raw_response, context).await + ) + } +} + +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, + AzureAi, + Reducto, + VertexAi, +} + +pub(crate) fn resolve_provider_config( + model: &str, + custom_llm_provider: Option<&str>, +) -> Result<(String, OcrConfigKind), super::Error> { + let provider = + get_custom_llm_provider(model, custom_llm_provider).unwrap_or(CustomLlmProvider { + model, + custom_llm_provider: OcrProvider::Mistral.into(), + }); + 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 + } + OcrProvider::AzureAi + if provider.model.to_ascii_lowercase().contains("cohere") + && provider.model.to_ascii_lowercase().contains("parse") => + { + OcrConfigKind::AzureCohere + } + OcrProvider::AzureAi => OcrConfigKind::AzureAi, + OcrProvider::Reducto if provider.model.eq_ignore_ascii_case("parse-legacy") => { + OcrConfigKind::ReductoLegacy + } + OcrProvider::Reducto => OcrConfigKind::ReductoV3, + OcrProvider::VertexAi if provider.model.to_ascii_lowercase().contains("deepseek") => { + OcrConfigKind::VertexDeepSeek + } + OcrProvider::VertexAi => OcrConfigKind::VertexAi, + }; + Ok((provider.model.to_string(), config)) +} + +fn is_document_intelligence_model(model: &str) -> bool { + let model = model.to_ascii_lowercase(); + model.contains("doc-intelligence") || model.contains("documentintelligence") +} + +#[cfg(test)] +mod tests { + use super::*; + use litellm_auth::{InputSource, Sourced}; + use rstest::rstest; + + #[rstest] + #[case("cohere")] + #[case("mistral")] + #[case("azure_ai")] + #[case("reducto")] + #[case("vertex_ai")] + fn provider_names_round_trip_exactly(#[case] provider: &str) { + let (_, config) = resolve_provider_config("model", Some(provider)).unwrap(); + let resolved: &'static str = config.provider().into(); + assert_eq!(resolved, provider); + } + + #[rstest] + #[case("Mistral")] + #[case("unknown")] + fn invalid_provider_names_are_rejected(#[case] provider: &str) { + assert!(matches!( + resolve_provider_config("model", Some(provider)), + Err(crate::ocr::Error::InvalidProvider(value)) if value == provider + )); + } + + #[rstest] + #[case("mistral/ocr")] + #[case("azure_ai/ocr")] + #[case("azure_ai/doc-intelligence/prebuilt-layout")] + #[case("reducto/parse-v3")] + #[case("vertex_ai/mistral-ocr")] + #[case("vertex_ai/deepseek-ocr")] + fn pdf_health_check_documents_are_valid(#[case] model: &str) { + 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-")); + } + + #[rstest] + #[case("cohere/parse")] + #[case("azure_ai/cohere-parse")] + fn png_health_check_documents_are_valid(#[case] model: &str) { + 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") + ); + } + + #[rstest] + #[case("mistral/ocr", Some("MISTRAL_API_KEY"))] + #[case("cohere/parse", Some("COHERE_API_KEY"))] + #[case("azure_ai/ocr", Some("AZURE_AI_API_KEY"))] + #[case("azure_ai/cohere-parse", Some("AZURE_AI_API_KEY"))] + #[case( + "azure_ai/doc-intelligence/prebuilt-layout", + Some("AZURE_DOCUMENT_INTELLIGENCE_API_KEY") + )] + #[case("vertex_ai/mistral-ocr", Some("VERTEX_AI_API_KEY"))] + #[case("vertex_ai/deepseek-ocr", Some("VERTEX_AI_API_KEY"))] + #[case("reducto/parse-v3", None)] + #[case("reducto/parse-legacy", None)] + fn api_key_metadata_follows_provider_overrides_and_python_defaults( + #[case] model: &str, + #[case] expected: Option<&str>, + ) { + assert_eq!(get_api_key_env_var(model, None).unwrap(), expected); + } + + #[test] + fn connection_resolution_preserves_dynamic_precedence_and_input_sources() { + let connection = OcrConfigKind::Mistral.resolve_connection_params(OcrCredentialInputs { + api_key: Some(Sourced::new("explicit-key".into(), InputSource::Deployment)), + api_base: Some(Sourced::new( + "https://explicit.test".into(), + InputSource::Deployment, + )), + dynamic_api_key: Some(Sourced::new("dynamic-key".into(), InputSource::Environment)), + dynamic_api_base: Some(Sourced::new( + "https://dynamic.test".into(), + InputSource::Request, + )), + }); + assert_eq!( + connection + .api_key + .as_ref() + .map(|value| value.value().as_str()), + Some("dynamic-key") + ); + assert_eq!( + connection + .api_base + .as_ref() + .map(|value| value.value().as_str()), + Some("https://dynamic.test") + ); + assert_eq!( + connection.api_key.as_ref().map(Sourced::source), + Some(InputSource::Environment) + ); + assert_eq!( + connection.api_base.as_ref().map(Sourced::source), + Some(InputSource::Request) + ); + } + + #[rstest] + #[case(None)] + #[case(Some(""))] + fn empty_or_missing_dynamic_credentials_preserve_explicit_values( + #[case] dynamic_value: Option<&str>, + ) { + let dynamic = + dynamic_value.map(|value| Sourced::new(value.into(), InputSource::Environment)); + let connection = OcrConfigKind::Mistral.resolve_connection_params(OcrCredentialInputs { + api_key: Some(Sourced::new("explicit-key".into(), InputSource::Deployment)), + api_base: Some(Sourced::new( + "https://explicit.test".into(), + InputSource::Deployment, + )), + dynamic_api_key: dynamic.clone(), + dynamic_api_base: dynamic, + }); + assert_eq!( + connection + .api_key + .as_ref() + .map(|value| value.value().as_str()), + Some("explicit-key") + ); + assert_eq!( + connection + .api_base + .as_ref() + .map(|value| value.value().as_str()), + Some("https://explicit.test") + ); + } + + #[rstest] + #[case(None, None)] + #[case(Some("key"), None)] + #[case(None, Some("base"))] + #[case(Some("key"), Some("base"))] + fn document_intelligence_only_accepts_dynamic_values_for_explicit_fields( + #[case] explicit_key: Option<&str>, + #[case] explicit_base: Option<&str>, + ) { + let connection = OcrConfigKind::AzureDocumentIntelligence.resolve_connection_params( + OcrCredentialInputs { + api_key: explicit_key + .map(|value| Sourced::new(value.into(), InputSource::Deployment)), + api_base: explicit_base + .map(|value| Sourced::new(value.into(), InputSource::Deployment)), + dynamic_api_key: Some(Sourced::new("dynamic-key".into(), InputSource::Environment)), + dynamic_api_base: Some(Sourced::new( + "https://dynamic.test".into(), + InputSource::Deployment, + )), + }, + ); + assert_eq!( + connection + .api_key + .as_ref() + .map(|value| value.value().as_str()), + explicit_key.map(|_| "dynamic-key") + ); + assert_eq!( + connection + .api_base + .as_ref() + .map(|value| value.value().as_str()), + explicit_base.map(|_| "https://dynamic.test") + ); + } + + #[rstest] + #[case("mistral/future-ocr-model", OcrConfigKind::Mistral)] + #[case("azure_ai/future-ocr-model", OcrConfigKind::AzureAi)] + fn provider_models_are_preserved_without_a_local_allowlist( + #[case] qualified_model: &str, + #[case] expected_config: OcrConfigKind, + ) { + let expected_model = qualified_model.split_once('/').unwrap().1; + let (model, config) = resolve_provider_config(qualified_model, None).unwrap(); + assert_eq!(model, expected_model); + assert_eq!(config, expected_config); + } + + #[rstest] + #[case("reducto/parse-legacy", OcrConfigKind::ReductoLegacy)] + #[case("reducto/future-parse-model", OcrConfigKind::ReductoV3)] + #[case( + "azure_ai/doc-intelligence/prebuilt-layout", + OcrConfigKind::AzureDocumentIntelligence + )] + fn provider_specific_models_select_their_config( + #[case] model: &str, + #[case] expected_config: OcrConfigKind, + ) { + assert_eq!( + resolve_provider_config(model, None).unwrap().1, + expected_config + ); + assert_eq!( + resolve_provider_config(model, None).unwrap().0, + model.split_once('/').unwrap().1 + ); + } +} diff --git a/litellm-rust/crates/core/src/ocr/registry.rs b/litellm-rust/crates/core/src/ocr/registry.rs deleted file mode 100644 index 17185a02020..00000000000 --- a/litellm-rust/crates/core/src/ocr/registry.rs +++ /dev/null @@ -1,132 +0,0 @@ -use super::adapters::OcrAdapter; -use crate::ocr::Error; -use crate::providers::custom_llm_provider::{CustomLlmProvider, get_custom_llm_provider}; - -macro_rules! define_adapter_types { - ($( $variant:ident, $adapter:ty, $instance:expr, $provider:ident; )+) => { - #[derive(Clone, Copy, Debug, PartialEq, Eq)] - pub(crate) enum OcrAdapterKind { - $( $variant, )+ - } - - impl OcrAdapterKind { - pub(crate) const fn provider(self) -> OcrProvider { - match self { - $( Self::$variant => <$adapter>::PROVIDER, )+ - } - } - } - }; -} - -super::adapters::for_each_ocr_adapter!(define_adapter_types); - -#[derive(Clone, Copy, Debug, PartialEq, Eq)] -pub(crate) enum OcrProvider { - Cohere, - Mistral, - AzureAi, - Reducto, - 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_wire_adapter( - model: &str, - custom_llm_provider: Option<&str>, -) -> Result<(String, OcrAdapterKind), Error> { - let provider = - get_custom_llm_provider(model, custom_llm_provider).unwrap_or(CustomLlmProvider { - model, - custom_llm_provider: OcrProvider::Mistral.as_str(), - }); - let typed_provider = match provider.custom_llm_provider { - "cohere" => OcrProvider::Cohere, - "mistral" => OcrProvider::Mistral, - "azure_ai" => OcrProvider::AzureAi, - "reducto" => OcrProvider::Reducto, - "vertex_ai" => OcrProvider::VertexAi, - value => return Err(Error::InvalidProvider(value.to_string())), - }; - let adapter = match typed_provider { - OcrProvider::Cohere => OcrAdapterKind::Cohere, - OcrProvider::Mistral => OcrAdapterKind::Mistral, - OcrProvider::AzureAi if is_document_intelligence_model(provider.model) => { - OcrAdapterKind::AzureDocumentIntelligence - } - OcrProvider::AzureAi - if provider.model.to_ascii_lowercase().contains("cohere") - && provider.model.to_ascii_lowercase().contains("parse") => - { - OcrAdapterKind::AzureCohere - } - OcrProvider::AzureAi => OcrAdapterKind::AzureMistral, - OcrProvider::Reducto if provider.model.eq_ignore_ascii_case("parse-legacy") => { - OcrAdapterKind::ReductoLegacy - } - OcrProvider::Reducto if provider.model.eq_ignore_ascii_case("parse-v3") => { - OcrAdapterKind::ReductoV3 - } - OcrProvider::Reducto => OcrAdapterKind::ReductoV3, - OcrProvider::VertexAi if provider.model.to_ascii_lowercase().contains("deepseek") => { - OcrAdapterKind::VertexDeepSeek - } - OcrProvider::VertexAi => OcrAdapterKind::VertexMistral, - }; - Ok((provider.model.to_string(), adapter)) -} - -fn is_document_intelligence_model(model: &str) -> bool { - let model = model.to_ascii_lowercase(); - model.contains("doc-intelligence") || model.contains("documentintelligence") -} - -#[cfg(test)] -mod tests { - use super::*; - - #[test] - fn provider_models_are_preserved_without_a_local_allowlist() { - let cases = [ - ("mistral/future-ocr-model", OcrAdapterKind::Mistral), - ("azure_ai/future-ocr-model", OcrAdapterKind::AzureMistral), - ]; - - for (qualified_model, expected_adapter) in cases { - let expected_model = qualified_model.split_once('/').unwrap().1; - let (model, adapter) = resolve_wire_adapter(qualified_model, None).unwrap(); - assert_eq!(model, expected_model); - assert_eq!(adapter, expected_adapter); - } - } - - #[test] - fn unknown_reducto_models_use_the_current_protocol() { - let (model, adapter) = resolve_wire_adapter("reducto/future-parse-model", None).unwrap(); - assert_eq!(model, "future-parse-model"); - assert_eq!(adapter, OcrAdapterKind::ReductoV3); - } - - #[test] - fn known_protocol_models_still_select_specialized_adapters() { - let (model, adapter) = resolve_wire_adapter("reducto/parse-legacy", None).unwrap(); - assert_eq!(model, "parse-legacy"); - assert_eq!(adapter, OcrAdapterKind::ReductoLegacy); - - let (model, adapter) = - resolve_wire_adapter("azure_ai/doc-intelligence/prebuilt-layout", None).unwrap(); - assert_eq!(model, "doc-intelligence/prebuilt-layout"); - assert_eq!(adapter, OcrAdapterKind::AzureDocumentIntelligence); - } -} diff --git a/litellm-rust/crates/core/src/ocr/types.rs b/litellm-rust/crates/core/src/ocr/types.rs index bb212674b33..449ba34b593 100644 --- a/litellm-rust/crates/core/src/ocr/types.rs +++ b/litellm-rust/crates/core/src/ocr/types.rs @@ -1,5 +1,4 @@ use std::collections::BTreeMap; -use std::convert::Infallible; use std::path::PathBuf; use std::sync::Arc; use std::time::Duration; @@ -7,12 +6,15 @@ use std::time::Duration; use bytes::Bytes; 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::registry::{OcrAdapterKind, resolve_wire_adapter}; +use super::provider_config::{OcrConfigKind, resolve_provider_config}; +use crate::call_arguments::CallArguments; use crate::constants::OCR_HTTP_TIMEOUT_SECS; -use crate::ocr::Error; -use litellm_auth::{InputSource, TokenProviderHandle}; +use crate::serde_compat::{FiniteF64, LaxI64}; #[derive(Clone, Debug, PartialEq, Serialize, Deserialize)] #[serde(tag = "type")] @@ -21,13 +23,13 @@ pub enum OcrDocument { DocumentUrl { document_url: String, #[serde(flatten)] - extra_fields: Map, + extra_fields: BTreeMap, }, #[serde(rename = "image_url")] ImageUrl { image_url: String, #[serde(flatten)] - extra_fields: Map, + extra_fields: BTreeMap, }, } @@ -39,6 +41,11 @@ impl OcrDocument { } } + pub(crate) fn is_remote(&self) -> bool { + let source = self.source(); + source.starts_with("http://") || source.starts_with("https://") + } + pub(crate) fn with_source(self, source: String) -> Self { match self { Self::DocumentUrl { extra_fields, .. } => Self::DocumentUrl { @@ -53,6 +60,14 @@ impl OcrDocument { } } +impl TryFrom for OcrDocument { + type Error = super::Error; + + fn try_from(value: Value) -> Result { + super::json::decode_request_value(value, "document") + } +} + #[derive(Clone, Debug, PartialEq)] pub enum OcrDocumentInput { Document(OcrDocument), @@ -76,6 +91,15 @@ impl From for OcrDocumentInput { } } +impl From for OcrDocumentInput { + fn from(path: PathBuf) -> Self { + Self::Path { + path, + mime_type: None, + } + } +} + #[derive(Clone, Debug, PartialEq, Eq)] pub struct OcrFileContent { pub bytes: Bytes, @@ -90,6 +114,107 @@ pub enum OcrResponseFormat { Native, } +#[derive(Clone, Default)] +pub struct OcrCredentialInputs { + pub api_key: Option>, + pub dynamic_api_key: Option>, + pub api_base: Option>, + pub dynamic_api_base: Option>, +} + +impl OcrCredentialInputs { + pub fn new( + api_key: Option, + api_key_source: InputSource, + api_base: Option, + api_base_source: InputSource, + ) -> Self { + Self { + api_key: nonblank(api_key).map(|value| Sourced::new(value, api_key_source)), + dynamic_api_key: None, + api_base: nonblank(api_base).map(|value| Sourced::new(value, api_base_source)), + dynamic_api_base: None, + } + } +} + +#[derive(Clone)] +pub struct OcrTransportConfig { + 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, +} + +impl Default for OcrTransportConfig { + fn default() -> Self { + Self { + extra_headers: Vec::new(), + extra_headers_source: InputSource::Deployment, + timeout: Duration::from_secs(OCR_HTTP_TIMEOUT_SECS), + max_download_bytes: crate::constants::OCR_DOWNLOAD_MAX_BYTES, + max_response_bytes: crate::constants::OCR_RESPONSE_MAX_BYTES, + poll_timeout: Duration::from_secs(crate::constants::OCR_POLL_TIMEOUT_SECS), + } + } +} + +impl OcrTransportConfig { + pub fn with_overrides( + self, + extra_headers: Vec<(String, String)>, + extra_headers_source: InputSource, + timeout: Option, + ) -> Self { + Self { + extra_headers, + extra_headers_source, + timeout: timeout.unwrap_or(self.timeout), + ..self + } + } +} + +fn nonblank(value: Option) -> Option { + value + .map(|value| value.trim().to_string()) + .filter(|value| !value.is_empty()) +} + +/// Caller-supplied connection overrides for a [`LiteLLMOcrRequest`], in the +/// shape hosts receive them: JSON-ish headers, optional timeout, optional +/// credentials, and per-field provenance in `input_sources`. +#[derive(Clone, Debug, Default)] +pub struct OcrConnectionInputs { + pub api_key: Option, + pub api_base: Option, + pub extra_headers: Map, + pub timeout: Option, + pub input_sources: BTreeMap, +} + +impl OcrConnectionInputs { + fn source(&self, name: &str) -> InputSource { + self.input_sources.get(name).copied().unwrap_or_default() + } + + fn header_pairs(&self) -> Result, super::Error> { + self.extra_headers + .iter() + .map(|(name, value)| { + value + .as_str() + .map(|value| (name.clone(), value.to_string())) + .ok_or_else(|| super::Error::RequestField { + path: format!("extra_headers.{name}"), + }) + }) + .collect() + } +} + #[derive(Clone)] pub struct OcrConnection { pub api_key: Option, @@ -104,72 +229,154 @@ pub struct OcrConnection { pub poll_timeout: Duration, } -impl Default for OcrConnection { - fn default() -> Self { +impl OcrConnection { + pub(crate) fn new(credentials: ResolvedOcrCredentials, transport: OcrTransportConfig) -> Self { + let api_key_source = credentials + .api_key + .as_ref() + .map(Sourced::source) + .unwrap_or(InputSource::Deployment); + let api_base_source = credentials + .api_base + .as_ref() + .map(Sourced::source) + .unwrap_or(InputSource::Deployment); Self { - api_key: None, - api_key_source: InputSource::Deployment, - api_base: None, - api_base_source: InputSource::Deployment, - extra_headers: Vec::new(), - extra_headers_source: InputSource::Deployment, - timeout: Duration::from_secs(OCR_HTTP_TIMEOUT_SECS), - max_download_bytes: crate::constants::OCR_DOWNLOAD_MAX_BYTES, - max_response_bytes: crate::constants::OCR_RESPONSE_MAX_BYTES, - poll_timeout: Duration::from_secs(crate::constants::OCR_POLL_TIMEOUT_SECS), + api_key: credentials.api_key.map(Sourced::into_value), + api_key_source, + api_base: credentials.api_base.map(Sourced::into_value), + 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, + max_response_bytes: transport.max_response_bytes, + poll_timeout: transport.poll_timeout, } } } -pub struct LiteLLMOcrRequest { - pub model: String, - pub document: D, - pub connection: OcrConnection, - pub hooks: Arc, - pub litellm_call_id: Option, - pub optional_params: Map, - pub input_sources: BTreeMap, - pub azure_ad_token_provider: Option, - pub(crate) adapter: OcrAdapterKind, +impl Default for OcrConnection { + fn default() -> Self { + Self::new( + ResolvedOcrCredentials::default(), + OcrTransportConfig::default(), + ) + } } -impl LiteLLMOcrRequest { +#[derive(Clone, Default)] +pub(crate) struct ResolvedOcrCredentials { + pub api_key: Option>, + pub api_base: Option>, +} + +pub struct LiteLLMOcrRequest { + pub model: String, + pub document: D, + pub credentials: OcrCredentialInputs, + pub transport: OcrTransportConfig, + pub hooks: Arc, + pub litellm_call_id: Option, + pub optional_params: CallArguments, + pub input_sources: BTreeMap, + pub azure_ad_token_provider: Option, + pub(crate) config: OcrConfigKind, +} + +impl LiteLLMOcrRequest { pub fn new( model: String, - document: D, + document: impl Into, custom_llm_provider: Option<&str>, - optional_params: Map, - ) -> Result { - let (model, adapter_kind) = resolve_wire_adapter(&model, custom_llm_provider)?; + optional_params: CallArguments, + ) -> Result { + let (model, config) = resolve_provider_config(&model, custom_llm_provider)?; + let default_transport = OcrTransportConfig::default(); + let max_response_bytes = optional_params + .get("max_response_bytes") + .map(|value| { + value + .as_u64() + .and_then(|value| usize::try_from(value).ok()) + .filter(|value| *value > 0 && *value <= default_transport.max_response_bytes) + .ok_or_else(|| super::Error::RequestField { + path: "max_response_bytes".into(), + }) + }) + .transpose()? + .unwrap_or(default_transport.max_response_bytes); + let transport = OcrTransportConfig { + max_response_bytes, + ..default_transport + }; + let optional_params = optional_params + .into_iter() + .filter(|(name, _)| name != "max_response_bytes") + .collect(); Ok(Self { model, - document, - connection: OcrConnection::default(), + document: document.into(), + credentials: OcrCredentialInputs::default(), + transport, hooks: Arc::new(NoopOcrHooks), litellm_call_id: None, optional_params, input_sources: BTreeMap::new(), azure_ad_token_provider: None, - adapter: adapter_kind, + config, + }) + } +} + +impl LiteLLMOcrRequest { + pub fn map_document( + self, + map: impl FnOnce(D) -> Result, + ) -> Result, E> { + Ok(LiteLLMOcrRequest { + model: self.model, + document: map(self.document)?, + credentials: self.credentials, + transport: self.transport, + hooks: self.hooks, + litellm_call_id: self.litellm_call_id, + optional_params: self.optional_params, + input_sources: self.input_sources, + azure_ad_token_provider: self.azure_ad_token_provider, + config: self.config, }) } - pub(crate) fn response_format( - &self, - ) -> Result { + pub fn with_document(self, document: T) -> LiteLLMOcrRequest { + LiteLLMOcrRequest { + model: self.model, + document, + credentials: self.credentials, + transport: self.transport, + hooks: self.hooks, + litellm_call_id: self.litellm_call_id, + optional_params: self.optional_params, + input_sources: self.input_sources, + azure_ad_token_provider: self.azure_ad_token_provider, + config: self.config, + } + } + + pub(crate) fn response_format(&self) -> Result { self.optional_params .get("req_format") + .filter(|value| !value.is_null()) .map(|value| { - serde_json::from_value(value.clone()) - .map_err(|_| super::error::OcrRequestError::RequestFormat) + serde_json::from_value(value.clone()).map_err(|_| super::Error::RequestFormat) }) .transpose() .map(|format| format.unwrap_or_default()) } pub fn provider_name(&self) -> &'static str { - self.adapter.provider().as_str() + self.config.provider().into() } pub fn with_host_hooks( @@ -184,61 +391,335 @@ impl LiteLLMOcrRequest { } } - pub fn map_document( + pub fn with_connection_inputs( self, - map: impl FnOnce(D) -> Result, - ) -> Result, E> { - Ok(LiteLLMOcrRequest { - model: self.model, - document: map(self.document)?, - connection: self.connection, - hooks: self.hooks, - litellm_call_id: self.litellm_call_id, - optional_params: self.optional_params, - input_sources: self.input_sources, - azure_ad_token_provider: self.azure_ad_token_provider, - adapter: self.adapter, - }) - } - - pub fn with_document(self, document: T) -> LiteLLMOcrRequest { - let Ok(request) = self.map_document(|_| Ok::(document)); - request + credentials: OcrCredentialInputs, + transport: OcrTransportConfig, + input_sources: BTreeMap, + ) -> Self { + Self { + credentials, + transport, + input_sources, + ..self + } } } -impl From for LiteLLMOcrRequest { - fn from(request: LiteLLMOcrRequest) -> Self { - let Ok(request) = request - .map_document(|document| Ok::<_, Infallible>(OcrDocumentInput::Document(document))); - request +impl LiteLLMOcrRequest { + /// Builds a request from host-shaped inputs in one step: provider + /// resolution, optional-param validation, header/timeout overrides and + /// sourced credentials. Hosts should prefer this over sequencing + /// [`Self::new`], [`OcrTransportConfig::with_overrides`] and + /// [`Self::with_connection_inputs`] by hand. + pub fn from_inputs( + model: String, + document: impl Into, + custom_llm_provider: Option<&str>, + optional_params: CallArguments, + connection: OcrConnectionInputs, + ) -> Result { + let request = Self::new(model, document, custom_llm_provider, optional_params)?; + let transport = request.transport.clone().with_overrides( + connection.header_pairs()?, + connection.source("extra_headers"), + connection.timeout, + ); + let (api_key_source, api_base_source) = + (connection.source("api_key"), connection.source("api_base")); + let credentials = OcrCredentialInputs::new( + connection.api_key, + api_key_source, + connection.api_base, + api_base_source, + ); + Ok(request.with_connection_inputs(credentials, transport, connection.input_sources)) } } +pub(crate) type ResolvedOcrRequest = LiteLLMOcrRequest; + +pub(crate) struct PreparedOcrRequest { + pub model: String, + pub document: OcrDocument, + pub connection: OcrConnection, + pub hooks: Arc, + pub optional_params: CallArguments, + pub input_sources: BTreeMap, + pub azure_ad_token_provider: Option, + pub(crate) config: OcrConfigKind, +} + +impl PreparedOcrRequest { + pub(crate) fn new(request: ResolvedOcrRequest, connection: OcrConnection) -> Self { + let LiteLLMOcrRequest { + model, + document, + credentials: _, + transport: _, + hooks, + litellm_call_id: _, + optional_params, + input_sources, + azure_ad_token_provider, + config, + } = request; + Self { + model, + document, + connection, + hooks, + optional_params, + input_sources, + azure_ad_token_provider, + config, + } + } + + pub(crate) fn response_format(&self) -> Result { + self.optional_params + .get("req_format") + .filter(|value| !value.is_null()) + .map(|value| { + serde_json::from_value(value.clone()).map_err(|_| super::Error::RequestFormat) + }) + .transpose() + .map(|format| format.unwrap_or_default()) + } + + pub(crate) fn provider_name(&self) -> &'static str { + self.config.provider().into() + } +} + +#[serde_as] +#[derive(Clone, Debug, Default, PartialEq, Serialize, Deserialize)] +pub struct OcrPageDimensions { + #[serde_as(deserialize_as = "Option")] + pub dpi: Option, + #[serde_as(deserialize_as = "Option")] + pub height: Option, + #[serde_as(deserialize_as = "Option")] + pub width: Option, +} + +#[derive(Clone, Debug, Default, PartialEq, Serialize, Deserialize)] +pub struct OcrPageImage { + pub image_base64: Option, + pub bbox: Option>, + #[serde(flatten)] + pub extra_fields: Map, +} + +#[serde_as] +#[derive(Clone, Debug, Default, PartialEq, Serialize, Deserialize)] +pub struct OcrPage { + #[serde_as(deserialize_as = "LaxI64")] + pub index: i64, + pub markdown: String, + pub images: Option>, + pub dimensions: Option, + #[serde(flatten)] + pub extra_fields: Map, +} + +#[serde_as] +#[derive(Clone, Debug, Default, PartialEq, Serialize, Deserialize)] +pub struct OcrUsageInfo { + #[serde_as(deserialize_as = "Option")] + pub pages_processed: Option, + #[serde_as(deserialize_as = "Option")] + pub pages_processed_annotation: Option, + #[serde_as(deserialize_as = "Option")] + pub credits: Option, + #[serde_as(deserialize_as = "Option")] + pub doc_size_bytes: Option, + #[serde(flatten)] + pub extra_fields: Map, +} + #[derive(Clone, Debug, PartialEq, Serialize, Deserialize)] pub struct LiteLLMOcrResponse { - pub pages: Vec, + pub pages: Vec, pub model: String, pub document_annotation: Option, - pub usage_info: Option, + pub usage_info: Option, + pub content: Option, + pub tables: Option>>, + #[serde(rename = "keyValuePairs")] + pub key_value_pairs: Option>>, + #[serde(default = "ocr_object")] pub object: String, #[serde(flatten)] pub extra_fields: Map, #[serde(skip_serializing_if = "Option::is_none")] - pub provider_native_response: Option, + pub provider_native_response: Option>, } impl LiteLLMOcrResponse { + pub fn new(model: impl Into, pages: Vec) -> Self { + Self { + pages, + model: model.into(), + document_annotation: None, + usage_info: None, + content: None, + tables: None, + key_value_pairs: None, + object: ocr_object(), + extra_fields: Map::new(), + provider_native_response: None, + } + } + pub fn into_json(self) -> Value { serde_json::to_value(self).expect("OCR response fields are JSON-compatible") } } +fn ocr_object() -> String { + "ocr".into() +} + #[cfg(test)] mod tests { use super::*; use serde_json::json; + fn document() -> OcrDocument { + OcrDocument::try_from( + json!({"type":"document_url","document_url":"data:application/pdf;base64,YWJj"}), + ) + .unwrap() + } + + #[test] + fn from_inputs_applies_connection_overrides_with_field_sources() { + let request = LiteLLMOcrRequest::from_inputs( + "mistral/model".into(), + document(), + None, + Default::default(), + OcrConnectionInputs { + api_key: Some(" key ".into()), + api_base: Some("".into()), + extra_headers: json!({"x-a": "1"}).as_object().unwrap().clone(), + timeout: Some(Duration::from_secs(7)), + input_sources: [ + ("api_key".to_string(), InputSource::Request), + ("extra_headers".to_string(), InputSource::Request), + ] + .into(), + }, + ) + .unwrap(); + + let api_key = request.credentials.api_key.as_ref().unwrap(); + assert_eq!(api_key.clone().into_value(), "key"); + assert_eq!(api_key.source(), InputSource::Request); + assert!(request.credentials.api_base.is_none()); + assert_eq!( + request.transport.extra_headers, + 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.input_sources.len(), 2); + + let defaulted = LiteLLMOcrRequest::from_inputs( + "mistral/model".into(), + document(), + None, + Default::default(), + OcrConnectionInputs::default(), + ) + .unwrap(); + assert_eq!( + defaulted.transport.timeout, + OcrTransportConfig::default().timeout + ); + assert_eq!( + defaulted.transport.extra_headers_source, + InputSource::Deployment + ); + } + + #[test] + fn from_inputs_rejects_non_string_header_values_by_path() { + let Err(error) = LiteLLMOcrRequest::from_inputs( + "mistral/model".into(), + document(), + None, + Default::default(), + OcrConnectionInputs { + extra_headers: json!({"x-a": 1}).as_object().unwrap().clone(), + ..Default::default() + }, + ) else { + panic!("non-string header value accepted"); + }; + assert!(matches!( + error, + super::super::Error::RequestField { ref path } if path == "extra_headers.x-a" + )); + } + + #[test] + fn normalized_response_rejects_invalid_shared_fields() { + for fields in [ + json!({"pages":[{}]}), + json!({"pages":[{"index":0,"markdown":false}]}), + json!({"pages":[{"index":0,"markdown":"","images":[{"bbox":[]}]}]}), + json!({"usage_info":{"pages_processed":1.5}}), + json!({"tables":[false]}), + json!({"keyValuePairs":[[]]}), + json!({"provider_native_response":[]}), + ] { + let payload: Map = json!({"model":"model", "pages":[]}) + .as_object() + .unwrap() + .iter() + .chain(fields.as_object().unwrap()) + .map(|(key, value)| (key.clone(), value.clone())) + .collect(); + assert!(serde_json::from_value::(Value::Object(payload)).is_err()); + } + assert!( + serde_json::from_value::(json!({ + "type":"image_url", "image_url":"https://example.com/image", "detail":42 + })) + .is_err() + ); + } + + #[test] + fn numeric_coercion_preserves_integer_precision_and_rejects_fractional_values() { + for (value, expected) in [ + (json!("9007199254740993.0"), 9_007_199_254_740_993), + (json!("+2.000"), 2), + (json!("1_000"), 1000), + (json!(true), 1), + (json!(2.0), 2), + ] { + let page: OcrPage = + serde_json::from_value(json!({"index":value,"markdown":""})).unwrap(); + assert_eq!(page.index, expected); + } + for value in [ + json!("1e2"), + json!(".0"), + json!("2."), + json!("_2"), + json!("2__0"), + json!(2.5), + json!(null), + ] { + assert!( + serde_json::from_value::(json!({"index":value,"markdown":""})).is_err() + ); + } + } + #[test] fn document_variants_preserve_provider_fields_when_rewriting_sources() { for (value, original, replacement, expected) in [ @@ -283,16 +764,11 @@ mod tests { #[test] fn response_serialization_flattens_extra_fields_and_omits_absent_native_response() { let response = LiteLLMOcrResponse { - pages: vec![], - model: "model".into(), - document_annotation: None, - usage_info: None, - object: "ocr".into(), extra_fields: json!({"provider_field":"kept"}) .as_object() .unwrap() .clone(), - provider_native_response: None, + ..LiteLLMOcrResponse::new("model", vec![]) }; let serialized = response.into_json(); assert_eq!(serialized["provider_field"], "kept"); diff --git a/litellm-rust/crates/core/src/ocr/wire.rs b/litellm-rust/crates/core/src/ocr/wire.rs index f0cad2b4e93..b05f388a277 100644 --- a/litellm-rust/crates/core/src/ocr/wire.rs +++ b/litellm-rust/crates/core/src/ocr/wire.rs @@ -1,69 +1,40 @@ -use crate::ocr::error::OcrRequestError; -use crate::ocr::error::OcrResponseError; use std::collections::BTreeMap; use std::time::Duration; -use super::types::{LiteLLMOcrRequest, OcrConnection, OcrDocument}; -use crate::ocr::Error; use litellm_auth::InputSource; -use serde::{ - Deserialize, - de::{DeserializeOwned, IntoDeserializer}, -}; +use serde::Deserialize; use serde_json::{Map, Value}; -const COMMON_OPTION_FIELDS: &[&str] = &["req_format", "extra_body", "max_response_bytes"]; -const MISTRAL_OPTION_FIELDS: &[&str] = &[ - "pages", - "include_image_base64", - "image_limit", - "image_min_size", - "bbox_annotation_format", - "document_annotation_format", - "document_annotation_prompt", - "extract_header", - "extract_footer", - "table_format", - "confidence_scores_granularity", - "include_blocks", - "id", -]; -const DEEPSEEK_OPTION_FIELDS: &[&str] = - &["stream", "temperature", "max_tokens", "top_p", "n", "stop"]; -const DOCUMENT_INTELLIGENCE_OPTION_FIELDS: &[&str] = &["pages", "features"]; -const REDUCTO_V3_OPTION_FIELDS: &[&str] = &["formatting", "retrieval", "settings"]; -const REDUCTO_LEGACY_OPTION_FIELDS: &[&str] = &["enhance"]; -const AZURE_AUTH_OPTION_FIELDS: &[&str] = &[ - "azure_ad_token", - "tenant_id", - "client_id", - "client_secret", - "azure_scope", - "azure_authority_host", - "azure_credential", - "azure_federated_token_file", - "enable_azure_ad_token_refresh", -]; -const VERTEX_AUTH_OPTION_FIELDS: &[&str] = &[ - "vertex_credentials", - "vertex_ai_credentials", - "vertex_project", - "vertex_ai_project", - "vertex_location", - "vertex_ai_location", -]; +pub use super::is_supported_request; +use super::{Error, LiteLLMOcrRequest, OcrConnectionInputs, OcrDocument, OcrDocumentInput}; -#[derive(Clone, Copy, Debug, PartialEq, Eq)] -pub struct OptionalParamSpec { - pub name: &'static str, - pub secret: bool, +pub fn consumed_optional_params( + model: &str, + provider: Option<&str>, +) -> Result, Error> { + let specs = super::consumed_optional_params(model, provider)?; + Ok(consumed_optional_param_names(model, provider)? + .into_iter() + .map(|name| crate::call_arguments::ArgumentSpec { + name, + secret: specs.iter().any(|spec| spec.name == name && spec.secret), + }) + .collect()) } -#[derive(Debug)] -pub struct DecodedOcrResponse { - pub data: T, - pub native: Option, - pub text: String, +pub fn consumed_optional_param_names( + model: &str, + provider: Option<&str>, +) -> Result, Error> { + let names = super::consumed_optional_param_names(model, provider)?; + let (_, config) = super::provider_config::resolve_provider_config(model, provider)?; + if config == super::provider_config::OcrConfigKind::VertexDeepSeek { + return Ok(names + .into_iter() + .chain(["stream", "temperature", "max_tokens", "top_p", "n", "stop"]) + .collect()); + } + Ok(names) } #[derive(Deserialize)] @@ -82,216 +53,54 @@ pub struct OcrWireRequest { pub timeout_seconds: Option, } -pub fn is_supported_request(model: &str, custom_llm_provider: Option<&str>) -> bool { - super::registry::resolve_wire_adapter(model, custom_llm_provider).is_ok() -} - -pub fn consumed_optional_param_names( - model: &str, - custom_llm_provider: Option<&str>, -) -> Result, Error> { - use super::registry::OcrAdapterKind; - - let (_, adapter) = super::registry::resolve_wire_adapter(model, custom_llm_provider)?; - let provider_fields: &[&str] = match adapter { - OcrAdapterKind::Cohere | OcrAdapterKind::AzureCohere => &["output_format"], - OcrAdapterKind::Mistral | OcrAdapterKind::AzureMistral | OcrAdapterKind::VertexMistral => { - MISTRAL_OPTION_FIELDS - } - OcrAdapterKind::AzureDocumentIntelligence => DOCUMENT_INTELLIGENCE_OPTION_FIELDS, - OcrAdapterKind::ReductoV3 => REDUCTO_V3_OPTION_FIELDS, - OcrAdapterKind::ReductoLegacy => REDUCTO_LEGACY_OPTION_FIELDS, - OcrAdapterKind::VertexDeepSeek => DEEPSEEK_OPTION_FIELDS, - }; - let auth_fields: &[&str] = match adapter { - OcrAdapterKind::AzureMistral - | OcrAdapterKind::AzureDocumentIntelligence - | OcrAdapterKind::AzureCohere => AZURE_AUTH_OPTION_FIELDS, - OcrAdapterKind::VertexMistral | OcrAdapterKind::VertexDeepSeek => VERTEX_AUTH_OPTION_FIELDS, - _ => &[], - }; - Ok(COMMON_OPTION_FIELDS - .iter() - .chain(provider_fields) - .chain(auth_fields) - .copied() - .collect()) -} - -pub fn consumed_optional_params( - model: &str, - custom_llm_provider: Option<&str>, -) -> Result, Error> { - consumed_optional_param_names(model, custom_llm_provider).map(|names| { - names - .into_iter() - .map(|name| OptionalParamSpec { - name, - secret: matches!( - name, - "azure_ad_token" - | "client_secret" - | "azure_federated_token_file" - | "vertex_credentials" - | "vertex_ai_credentials" - ), - }) - .collect() - }) -} - pub fn decode_request(wire: OcrWireRequest) -> Result { - let OcrWireRequest { - model, - document, - api_key, - api_base, - custom_llm_provider, - extra_headers, - optional_params, - input_sources, - timeout_seconds, - } = wire; decode_request_input(OcrWireRequest { - model, - document: decode_document(document)?, - api_key, - api_base, - custom_llm_provider, - extra_headers, - optional_params, - input_sources, - timeout_seconds, + model: wire.model, + document: decode_document(wire.document)?, + api_key: wire.api_key, + api_base: wire.api_base, + custom_llm_provider: wire.custom_llm_provider, + extra_headers: wire.extra_headers, + optional_params: wire.optional_params, + input_sources: wire.input_sources, + timeout_seconds: wire.timeout_seconds, }) } -pub fn decode_request_input(wire: OcrWireRequest) -> Result, Error> { - let api_key_source = source_for(&wire.input_sources, "api_key"); - let api_base_source = source_for(&wire.input_sources, "api_base"); - let extra_headers_source = source_for(&wire.input_sources, "extra_headers"); - let headers = wire - .extra_headers - .unwrap_or_default() - .into_iter() - .map(|(name, value)| { - let value = value - .as_str() - .ok_or_else(|| OcrRequestError::RequestField { - path: format!("extra_headers.{name}"), - })?; - Ok((name, value.to_string())) - }) - .collect::, OcrRequestError>>()?; +pub fn decode_request_input>( + wire: OcrWireRequest, +) -> Result { let timeout = wire .timeout_seconds .map(|seconds| { - Duration::try_from_secs_f64(seconds).map_err(|_| OcrRequestError::RequestField { + Duration::try_from_secs_f64(seconds).map_err(|_| Error::RequestField { path: "timeout_seconds".into(), }) }) .transpose()?; - let defaults = OcrConnection::default(); - let max_response_bytes = wire - .optional_params - .get("max_response_bytes") - .map(|value| { - value - .as_u64() - .and_then(|value| usize::try_from(value).ok()) - .filter(|value| *value > 0 && *value <= defaults.max_response_bytes) - .ok_or_else(|| OcrRequestError::RequestField { - path: "max_response_bytes".into(), - }) - }) - .transpose()? - .unwrap_or(defaults.max_response_bytes); - let request = LiteLLMOcrRequest::new( + LiteLLMOcrRequest::from_inputs( wire.model, wire.document, wire.custom_llm_provider.as_deref(), - wire.optional_params - .into_iter() - .filter(|(name, _)| name != "max_response_bytes") - .collect(), - )?; - let connection = OcrConnection { - api_key: nonblank(wire.api_key), - api_key_source, - api_base: nonblank(wire.api_base), - api_base_source, - extra_headers: headers, - extra_headers_source, - timeout: timeout.unwrap_or(defaults.timeout), - max_download_bytes: defaults.max_download_bytes, - max_response_bytes, - poll_timeout: defaults.poll_timeout, - }; - Ok(LiteLLMOcrRequest { - connection, - input_sources: wire.input_sources, - ..request - }) + wire.optional_params.into(), + OcrConnectionInputs { + api_key: wire.api_key, + api_base: wire.api_base, + extra_headers: wire.extra_headers.unwrap_or_default(), + timeout, + input_sources: wire.input_sources, + }, + ) } pub fn decode_document(value: Value) -> Result { let kind = value.get("type").and_then(Value::as_str); - let missing_url = matches!(kind, Some("document_url")) && value.get("document_url").is_none() - || matches!(kind, Some("image_url")) && value.get("image_url").is_none(); - if missing_url { - return Err(OcrRequestError::MissingDocumentUrl.into()); + if matches!(kind, Some("document_url")) && value.get("document_url").is_none() + || matches!(kind, Some("image_url")) && value.get("image_url").is_none() + { + return Err(Error::MissingDocumentUrl); } - Ok(decode_request_value(value, "document")?) -} - -fn source_for(sources: &BTreeMap, name: &str) -> InputSource { - sources.get(name).copied().unwrap_or_default() -} - -fn nonblank(value: Option) -> Option { - value - .map(|s| s.trim().to_string()) - .filter(|s| !s.is_empty()) -} -pub fn decode_request_value( - value: Value, - prefix: &str, -) -> Result { - serde_path_to_error::deserialize(value.into_deserializer()).map_err(|error| { - OcrRequestError::RequestField { - path: format!("{prefix}.{}", error.path()), - } - }) -} - -pub fn decode_response( - bytes: &[u8], - native: bool, -) -> Result, OcrResponseError> { - let mut deserializer = serde_json::Deserializer::from_slice(bytes); - let data = serde_path_to_error::deserialize(&mut deserializer).map_err(|error| { - OcrResponseError::ResponseField { - path: error.path().to_string(), - } - })?; - deserializer - .end() - .map_err(|_| OcrResponseError::ResponseField { - path: "response".into(), - })?; - let native = if native { - Some( - serde_json::from_slice(bytes).map_err(|_| OcrResponseError::ResponseField { - path: "response".into(), - })?, - ) - } else { - None - }; - Ok(DecodedOcrResponse { - data, - native, - text: String::from_utf8_lossy(bytes).into_owned(), - }) + super::json::decode_request_value(value, "document") } #[cfg(test)] @@ -305,7 +114,6 @@ mod tests { assert!(mistral.contains(&"req_format")); assert!(!mistral.contains(&"vertex_project")); assert!(!mistral.contains(&"opaque_extension")); - let vertex = consumed_optional_param_names("vertex_ai/deepseek-ocr", None).unwrap(); assert!(vertex.contains(&"temperature")); assert!(vertex.contains(&"vertex_credentials")); @@ -358,7 +166,10 @@ mod tests { serde_json::json!({"type": "document_url"}), serde_json::json!({"type": "image_url"}), ] { - assert_eq!(decode_document(document), Err(Error::MissingDocumentUrl)); + assert!(matches!( + decode_document(document), + Err(Error::MissingDocumentUrl) + )); } } } diff --git a/litellm-rust/crates/core/src/params.rs b/litellm-rust/crates/core/src/params.rs new file mode 100644 index 00000000000..cea410db816 --- /dev/null +++ b/litellm-rust/crates/core/src/params.rs @@ -0,0 +1,231 @@ +#[derive(Clone, Debug, PartialEq, Eq, thiserror::Error)] +pub enum Error { + #[error("invalid request: extra_body must be an object")] + ExtraBody, + #[error("invalid request: body must be a JSON object")] + Body, +} + +use std::ops::{Deref, DerefMut}; + +use serde::{Deserialize, Serialize}; +use serde_json::{Map, Value}; + +#[derive(Clone, Debug, Default, PartialEq, Serialize, Deserialize)] +#[serde(transparent)] +pub struct OpaqueParams(Map); + +pub fn is_control_param(name: &str) -> bool { + matches!( + name, + "api_key" + | "api_base" + | "custom_llm_provider" + | "extra_headers" + | "timeout" + | "timeout_seconds" + | "request_timeout" + | "max_retries" + | "req_format" + | "max_response_bytes" + | "litellm_call_id" + | "litellm_logging_obj" + | "litellm_metadata" + | "proxy_server_request" + | "callbacks" + | "success_callback" + | "failure_callback" + | "guardrails" + | "azure_ad_token" + | "azure_ad_token_provider" + | "tenant_id" + | "client_id" + | "client_secret" + | "azure_scope" + | "azure_authority_host" + | "azure_credential" + | "azure_federated_token_file" + | "enable_azure_ad_token_refresh" + | "vertex_credentials" + | "vertex_ai_credentials" + | "vertex_project" + | "vertex_ai_project" + | "vertex_location" + | "vertex_ai_location" + | "aws_access_key_id" + | "aws_secret_access_key" + | "aws_session_token" + | "aws_region_name" + | "aws_session_name" + | "aws_profile_name" + | "aws_role_name" + | "aws_web_identity_token" + | "aws_sts_endpoint" + | "aws_external_id" + | "aws_bedrock_runtime_endpoint" + ) +} + +impl OpaqueParams { + pub fn into_inner(self) -> Map { + self.0 + } + + pub fn without(&self, names: &[&str]) -> Self { + self.iter() + .filter(|(name, _)| !names.contains(&name.as_str())) + .map(|(name, value)| (name.clone(), value.clone())) + .collect() + } + + pub fn provider_params(&self) -> Self { + self.iter() + .filter(|(name, _)| !is_control_param(name)) + .map(|(name, value)| (name.clone(), value.clone())) + .collect() + } + + pub fn into_provider_body(self) -> Result, Error> { + let mut fields = self.0; + let overrides = match fields.remove("extra_body") { + None | Some(Value::Null) => Map::new(), + Some(Value::Object(fields)) => fields, + Some(_) => { + return Err(Error::ExtraBody); + } + }; + Ok(fields + .into_iter() + .chain(overrides) + .filter(|(name, _)| name != "extra_body" && !is_control_param(name)) + .collect()) + } +} + +#[cfg(test)] +fn merge_extra_params(body: &B, extra_params: OpaqueParams) -> Result { + let Value::Object(fields) = serde_json::to_value(body).map_err(|_| Error::Body)? else { + return Err(Error::Body); + }; + Ok(Value::Object( + fields + .into_iter() + .chain( + extra_params + .into_provider_body()? + .into_iter() + .filter(|(name, _)| name != "model"), + ) + .collect(), + )) +} + +impl Deref for OpaqueParams { + type Target = Map; + + fn deref(&self) -> &Self::Target { + &self.0 + } +} + +impl DerefMut for OpaqueParams { + fn deref_mut(&mut self) -> &mut Self::Target { + &mut self.0 + } +} + +impl From> for OpaqueParams { + fn from(value: Map) -> Self { + Self(value) + } +} + +impl From for Map { + fn from(value: OpaqueParams) -> Self { + value.0 + } +} + +impl FromIterator<(String, Value)> for OpaqueParams { + fn from_iter>(iter: T) -> Self { + Self(iter.into_iter().collect()) + } +} + +impl IntoIterator for OpaqueParams { + type Item = (String, Value); + type IntoIter = serde_json::map::IntoIter; + + fn into_iter(self) -> Self::IntoIter { + self.0.into_iter() + } +} + +#[cfg(test)] +mod tests { + use serde_json::json; + + use super::*; + + #[test] + fn extras_merge_shallowly_and_preserve_values_without_leaking_controls() { + let extras: OpaqueParams = serde_json::from_value(json!({ + "future": {"nested": [false, 0, null]}, + "explicit_null": null, + "azure_ad_token": "secret", + "req_format": "native", + "extra_body": { + "future": {"replacement": true}, + "temperature": 0.5, + "model": "override", + "aws_secret_access_key": "secret" + } + })) + .unwrap(); + let body = + merge_extra_params(&json!({"model":"resolved", "temperature":0.1}), extras).unwrap(); + assert_eq!( + body, + json!({ + "model":"resolved", "temperature":0.5, + "future":{"replacement":true}, "explicit_null":null + }) + ); + } + + #[test] + fn invalid_extra_body_is_rejected_and_null_is_empty() { + for value in [json!(false), json!([]), json!("value"), json!(1)] { + let params: OpaqueParams = serde_json::from_value(json!({"extra_body":value})).unwrap(); + assert!(params.into_provider_body().is_err()); + } + let params: OpaqueParams = + serde_json::from_value(json!({"extra_body":null,"future":null})).unwrap(); + assert_eq!( + Value::Object(params.into_provider_body().unwrap()), + json!({"future":null}) + ); + } + + #[test] + fn provider_params_preserve_opaque_values() { + let params: OpaqueParams = serde_json::from_value(json!({ + "object": {"future": [1, null]}, + "null": null, + "azure_ad_token": "secret" + })) + .unwrap(); + + let retained = params.provider_params(); + + assert_eq!( + serde_json::to_value(retained).unwrap(), + json!({"object": {"future": [1, null]}, "null": null}) + ); + } + + #[test] + fn outer_value_must_be_an_object() { + assert!(serde_json::from_value::(json!(["value"])).is_err()); + } +} diff --git a/litellm-rust/crates/core/src/providers/mod.rs b/litellm-rust/crates/core/src/providers/mod.rs index 70ca4386fff..79eb3404ece 100644 --- a/litellm-rust/crates/core/src/providers/mod.rs +++ b/litellm-rust/crates/core/src/providers/mod.rs @@ -2,4 +2,5 @@ pub mod anthropic; pub mod azure_ai; pub mod bedrock; pub mod custom_llm_provider; +pub(crate) mod model; pub mod openai; diff --git a/litellm-rust/crates/core/src/providers/model.rs b/litellm-rust/crates/core/src/providers/model.rs new file mode 100644 index 00000000000..fcedc4b023a --- /dev/null +++ b/litellm-rust/crates/core/src/providers/model.rs @@ -0,0 +1,219 @@ +use std::marker::PhantomData; + +use serde::{Deserialize, Deserializer, Serialize, Serializer}; + +#[derive(Clone, Copy, Debug, Eq, PartialEq, thiserror::Error)] +pub(crate) enum ModelNameError { + #[error("model name cannot be empty")] + EmptyModel, + #[error("model namespace must be one non-empty path segment: {0}")] + InvalidNamespace(&'static str), +} + +pub(crate) trait ModelNamespace { + const NAME: &'static str; +} + +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub(crate) struct RoutedModel<'a>(&'a str); + +impl<'a> RoutedModel<'a> { + pub(crate) fn new(value: &'a str) -> Result { + if value.is_empty() { + return Err(ModelNameError::EmptyModel); + } + Ok(Self(value)) + } + + pub(crate) fn into_provider( + self, + ) -> Result, ModelNameError> { + let namespace = N::NAME; + if namespace.is_empty() || namespace.contains('/') { + return Err(ModelNameError::InvalidNamespace(namespace)); + } + let prefix = format!("{namespace}/"); + let local_model = self.0.trim_start_matches(prefix.as_str()); + if local_model.is_empty() { + return Err(ModelNameError::EmptyModel); + } + Ok(ProviderModel { + value: format!("{prefix}{local_model}"), + namespace: PhantomData, + }) + } +} + +#[derive(Clone, Debug, Eq, PartialEq)] +pub(crate) struct ProviderModel { + value: String, + namespace: PhantomData, +} + +impl ProviderModel { + #[cfg(test)] + pub(crate) fn as_str(&self) -> &str { + &self.value + } +} + +impl Serialize for ProviderModel { + fn serialize(&self, serializer: S) -> Result + where + S: Serializer, + { + self.value.serialize(serializer) + } +} + +impl<'de, N: ModelNamespace> Deserialize<'de> for ProviderModel { + fn deserialize(deserializer: D) -> Result + where + D: Deserializer<'de>, + { + let value = String::deserialize(deserializer)?; + RoutedModel::new(&value) + .and_then(RoutedModel::into_provider::) + .map_err(::custom) + } +} + +#[cfg(test)] +mod tests { + use serde_json::json; + + use super::*; + + #[derive(Clone, Debug, Eq, PartialEq)] + struct DeepSeekAi; + + impl ModelNamespace for DeepSeekAi { + const NAME: &'static str = "deepseek-ai"; + } + + #[derive(Clone, Debug, Eq, PartialEq)] + struct FalAi; + + impl ModelNamespace for FalAi { + const NAME: &'static str = "fal-ai"; + } + + #[test] + fn qualifies_a_bare_model() { + let model = RoutedModel::new("deepseek-ocr-maas") + .and_then(RoutedModel::into_provider::) + .unwrap(); + + assert_eq!(model.as_str(), "deepseek-ai/deepseek-ocr-maas"); + } + + #[test] + fn preserves_an_already_qualified_model() { + let model = RoutedModel::new("deepseek-ai/deepseek-ocr-maas") + .and_then(RoutedModel::into_provider::) + .unwrap(); + + assert_eq!(model.as_str(), "deepseek-ai/deepseek-ocr-maas"); + } + + #[test] + fn collapses_repeated_owned_namespaces() { + let model = RoutedModel::new("deepseek-ai/deepseek-ai/deepseek-ai/deepseek-ocr-maas") + .and_then(RoutedModel::into_provider::) + .unwrap(); + + assert_eq!(model.as_str(), "deepseek-ai/deepseek-ocr-maas"); + } + + #[test] + fn matches_the_namespace_as_a_complete_segment() { + let model = RoutedModel::new("deepseek-ai-v2/model") + .and_then(RoutedModel::into_provider::) + .unwrap(); + + assert_eq!(model.as_str(), "deepseek-ai/deepseek-ai-v2/model"); + } + + #[test] + fn preserves_nested_provider_model_paths() { + let model = RoutedModel::new("publishers/vendor/models/model-v1") + .and_then(RoutedModel::into_provider::) + .unwrap(); + + assert_eq!( + model.as_str(), + "deepseek-ai/publishers/vendor/models/model-v1" + ); + } + + #[test] + fn namespace_markers_select_different_wire_names() { + let routed = RoutedModel::new("model-v1").unwrap(); + let deepseek = routed.into_provider::().unwrap(); + let fal = routed.into_provider::().unwrap(); + + assert_eq!(deepseek.as_str(), "deepseek-ai/model-v1"); + assert_eq!(fal.as_str(), "fal-ai/model-v1"); + } + + #[test] + fn rejects_empty_routed_models() { + assert_eq!(RoutedModel::new(""), Err(ModelNameError::EmptyModel)); + } + + #[test] + fn rejects_a_namespace_without_a_model() { + let result = + RoutedModel::new("deepseek-ai/").and_then(RoutedModel::into_provider::); + + assert_eq!(result, Err(ModelNameError::EmptyModel)); + } + + #[test] + fn rejects_invalid_namespace_markers() { + struct Empty; + impl ModelNamespace for Empty { + const NAME: &'static str = ""; + } + struct MultipleSegments; + impl ModelNamespace for MultipleSegments { + const NAME: &'static str = "one/two"; + } + + assert!(matches!( + RoutedModel::new("model").and_then(RoutedModel::into_provider::), + Err(ModelNameError::InvalidNamespace("")) + )); + assert!(matches!( + RoutedModel::new("model").and_then(RoutedModel::into_provider::), + Err(ModelNameError::InvalidNamespace("one/two")) + )); + } + + #[test] + fn provider_models_serialize_as_plain_strings() { + let model = RoutedModel::new("deepseek-ocr-maas") + .and_then(RoutedModel::into_provider::) + .unwrap(); + + assert_eq!( + serde_json::to_value(model).unwrap(), + json!("deepseek-ai/deepseek-ocr-maas") + ); + } + + #[test] + fn deserialization_reestablishes_the_namespace_invariant() { + let model: ProviderModel = + serde_json::from_value(json!("deepseek-ai/deepseek-ai/model-v1")).unwrap(); + + assert_eq!(model.as_str(), "deepseek-ai/model-v1"); + } + + #[test] + fn deserialization_rejects_missing_model_names() { + let result = serde_json::from_value::>(json!("deepseek-ai/")); + + assert!(result.is_err()); + } +} diff --git a/litellm-rust/crates/core/src/serde_compat.rs b/litellm-rust/crates/core/src/serde_compat.rs new file mode 100644 index 00000000000..5a2d0688c33 --- /dev/null +++ b/litellm-rust/crates/core/src/serde_compat.rs @@ -0,0 +1,151 @@ +use serde::{Deserialize, Deserializer, de::Error}; +use serde_json::Value; +use serde_with::DeserializeAs; + +pub(crate) struct LaxI64; +pub(crate) struct FiniteF64; + +impl<'de> DeserializeAs<'de, i64> for LaxI64 { + fn deserialize_as>(deserializer: D) -> Result { + match Value::deserialize(deserializer)? { + Value::Number(number) if number.is_f64() => number.as_f64().and_then(integral_float), + Value::Number(number) => number.as_i64(), + Value::String(value) => integer_string(value.trim()), + Value::Bool(value) => Some(i64::from(value)), + _ => None, + } + .ok_or_else(|| D::Error::custom("expected an integer in the i64 range")) + } +} + +impl<'de> DeserializeAs<'de, f64> for FiniteF64 { + fn deserialize_as>(deserializer: D) -> Result { + match Value::deserialize(deserializer)? { + Value::Number(number) => number.as_f64(), + Value::String(value) => value.trim().parse::().ok(), + Value::Bool(value) => Some(f64::from(value)), + _ => None, + } + .filter(|value| value.is_finite()) + .ok_or_else(|| D::Error::custom("expected a finite number")) + } +} + +fn integer_string(value: &str) -> Option { + let integer = match value.split_once('.') { + Some((integer, fraction)) => { + if fraction.is_empty() || !fraction.bytes().all(|byte| byte == b'0') { + return None; + } + integer + } + None => value, + }; + if integer.starts_with('_') || integer.ends_with('_') || integer.contains("__") { + return None; + } + let digits = integer.strip_prefix(['+', '-']).unwrap_or(integer); + if digits.is_empty() + || digits.starts_with('_') + || !digits + .bytes() + .all(|byte| byte.is_ascii_digit() || byte == b'_') + { + return None; + } + integer.replace('_', "").parse().ok() +} + +fn integral_float(value: f64) -> Option { + (value.is_finite() + && value.fract() == 0.0 + && value >= i64::MIN as f64 + && value < -(i64::MIN as f64)) + .then_some(value as i64) +} + +#[cfg(test)] +mod tests { + use super::*; + use serde::Serialize; + use serde_json::json; + use serde_with::serde_as; + + #[serde_as] + #[derive(Debug, Deserialize, Serialize, PartialEq)] + struct Numbers { + #[serde_as(deserialize_as = "Option>")] + integers: Option>, + #[serde_as(deserialize_as = "Option")] + float: Option, + } + + #[test] + fn adapters_compose_and_serialize_as_numbers() { + let numbers: Numbers = serde_json::from_value(json!({ + "integers": ["9007199254740993.0", "1_000", " +2.000 ", 3.0, true], + "float": " 1.5 " + })) + .unwrap(); + assert_eq!( + serde_json::to_value(numbers).unwrap(), + json!({ + "integers": [9_007_199_254_740_993_i64, 1000, 2, 3, 1], "float": 1.5 + }) + ); + for input in [json!({}), json!({"integers": null, "float": null})] { + assert_eq!( + serde_json::from_value::(input).unwrap(), + Numbers { + integers: None, + float: None, + } + ); + } + } + + #[test] + fn integer_bounds_and_invalid_values_are_checked() { + for input in [ + json!(i64::MIN), + json!(i64::MAX), + json!(i64::MAX.to_string()), + ] { + assert!(serde_json::from_value::(json!({"integers": [input]})).is_ok()); + } + for input in [ + json!(u64::MAX), + json!(9_223_372_036_854_775_808_u64), + json!(9_223_372_036_854_775_808.0), + json!("-9223372036854775809"), + json!("1.0000000000000001"), + json!("1e3"), + json!("2."), + json!(".0"), + json!("_2"), + json!("2__0"), + json!(2.5), + json!(null), + json!({}), + ] { + assert!(serde_json::from_value::(json!({"integers": [input]})).is_err()); + } + } + + #[test] + fn floats_reject_nonfinite_and_invalid_values() { + for input in [ + json!("NaN"), + json!("inf"), + json!("-inf"), + json!("1e999"), + json!([]), + ] { + assert!(serde_json::from_value::(json!({"float": input})).is_err()); + } + for (input, expected) in [(json!(2), 2.0), (json!(2.5), 2.5), (json!(true), 1.0)] { + let numbers: Numbers = serde_json::from_value(json!({"float": input})).unwrap(); + assert_eq!(numbers.float, Some(expected)); + } + } +} diff --git a/litellm-rust/crates/core/tests/azure_ai_ocr.rs b/litellm-rust/crates/core/tests/azure_ai_ocr.rs index b6dc8d90b93..253d2582acc 100644 --- a/litellm-rust/crates/core/tests/azure_ai_ocr.rs +++ b/litellm-rust/crates/core/tests/azure_ai_ocr.rs @@ -17,15 +17,15 @@ async fn facade_executes_azure_mistral_with_prepared_auth() { &base, json!({"include_image_base64":true}), ); - request.connection.api_key = None; - request.connection.extra_headers = vec![( + request.credentials.api_key = None; + request.transport.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"); + 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 ")); @@ -53,7 +53,7 @@ async fn facade_acquires_supplied_entra_token_for_final_request() { &base, json!({"azure_ad_token":"rust-owned-token"}), ); - request.connection.api_key = None; + request.credentials.api_key = None; perform_ocr(request).await.unwrap(); server.await.unwrap(); diff --git a/litellm-rust/crates/core/tests/azure_document_intelligence_ocr.rs b/litellm-rust/crates/core/tests/azure_document_intelligence_ocr.rs index 3fca59033cc..5682e8ad5be 100644 --- a/litellm-rust/crates/core/tests/azure_document_intelligence_ocr.rs +++ b/litellm-rust/crates/core/tests/azure_document_intelligence_ocr.rs @@ -23,11 +23,12 @@ async fn facade_maps_pages_features_and_url_document() { &base, json!({"pages":[2,0,0,1],"features":["keyValuePairs","languages"]}), ); - request.document = serde_json::from_value(json!({ + request.document = serde_json::from_value::(json!({ "type":"document_url", "document_url":"https://example.com/document.pdf" })) - .unwrap(); + .unwrap() + .into(); perform_ocr(request).await.unwrap(); server.await.unwrap(); @@ -118,13 +119,13 @@ async fn immediate_response_normalizes_pages_and_preserves_native() { .unwrap(); server.await.unwrap(); - assert_eq!(result.pages[0]["index"], 1); - assert_eq!(result.pages[0]["markdown"], "A\n\nB"); + assert_eq!(result.pages[0].index, 1); + assert_eq!(result.pages[0].markdown, "A\n\nB"); assert_eq!( - result.pages[0]["dimensions"], + serde_json::to_value(&result.pages[0].dimensions).unwrap(), json!({"width":816,"height":1056,"dpi":96}) ); - assert_eq!(result.usage_info, Some(json!({"pages_processed":1}))); + 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":[]}])); @@ -133,7 +134,10 @@ async fn immediate_response_normalizes_pages_and_preserves_native() { json!([{"key":{"content":"A"}}]) ); assert!(serialized.get("key_value_pairs").is_none()); - assert_eq!(result.provider_native_response, Some(operation)); + assert_eq!( + result.provider_native_response.map(Value::Object), + Some(operation) + ); } #[tokio::test] @@ -159,13 +163,16 @@ async fn accepted_response_polls_to_success_with_only_credentials() { json!({"req_format":"native"}), ); request - .connection + .transport .extra_headers .push(("X-Trace".into(), "initial-only".into())); let result = perform_ocr(request).await.unwrap(); server.await.unwrap(); - assert_eq!(result.provider_native_response, Some(operation)); + assert_eq!( + result.provider_native_response.map(Value::Object), + Some(operation) + ); let requests = seen.lock().unwrap(); assert_eq!(requests.len(), 3); assert!(requests[0].to_ascii_lowercase().contains("x-trace:")); @@ -239,8 +246,8 @@ async fn polling_forwards_bearer_credentials() { ]) .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())]; + request.credentials.api_key = None; + request.transport.extra_headers = vec![("Authorization".into(), "Bearer token".into())]; perform_ocr(request).await.unwrap(); server.await.unwrap(); @@ -375,7 +382,7 @@ async fn polling_deadline_bounds_retry_delay() { ]) .await; let mut request = wire_request("azure_ai/doc-intelligence/prebuilt-read", &base, json!({})); - request.connection.poll_timeout = std::time::Duration::from_millis(100); + request.transport.poll_timeout = std::time::Duration::from_millis(100); let error = tokio::time::timeout(std::time::Duration::from_secs(1), perform_ocr(request)) .await diff --git a/litellm-rust/crates/core/tests/deepseek_ocr.rs b/litellm-rust/crates/core/tests/deepseek_ocr.rs index 4ba39561dcd..3129f1e60a9 100644 --- a/litellm-rust/crates/core/tests/deepseek_ocr.rs +++ b/litellm-rust/crates/core/tests/deepseek_ocr.rs @@ -1,8 +1,10 @@ use rstest::rstest; use serde_json::{Value, json}; -use crate::ocr::codecs::deepseek::{ - DeepSeekOcrParams, DeepSeekOcrResponse, transform_ocr_request, transform_ocr_response, +use crate::llms::base_llm::ocr::transformation::BaseOcrConfig; +use crate::llms::vertex_ai::ocr::deepseek_transformation::{ + DeepSeekOcrParams, DeepSeekOcrResponse, VertexAIDeepSeekOCRConfig, + normalize_response as transform_ocr_response, }; use crate::ocr::types::OcrDocument; @@ -22,7 +24,9 @@ fn request_mapping_matches_python(#[case] name: &str, #[case] value: Value) { let params: DeepSeekOcrParams = serde_json::from_value(json!({name: value.clone(), "ignored": true})).unwrap(); let result = serde_json::to_value( - transform_ocr_request("deepseek-ai/deepseek-ocr-maas", document(), ¶ms).unwrap(), + VertexAIDeepSeekOCRConfig + .transform_ocr_request("deepseek-ai/deepseek-ocr-maas", document(), ¶ms, &[]) + .unwrap(), ) .unwrap(); assert_eq!(result["model"], "deepseek-ai/deepseek-ocr-maas"); @@ -43,12 +47,14 @@ fn request_maps_both_document_types_to_image_content(#[case] document: Value) { .or_else(|| document.get("document_url")) .unwrap() .clone(); - let request = transform_ocr_request( - "deepseek-ai/deepseek-ocr-maas", - serde_json::from_value(document).unwrap(), - &DeepSeekOcrParams::default(), - ) - .unwrap(); + 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], @@ -60,12 +66,17 @@ fn request_maps_both_document_types_to_image_content(#[case] document: Value) { #[case(json!("# hello"), "# hello")] #[case(json!("{broken"), "{broken")] #[case(json!(" {\"pages\":[]} "), " {\"pages\":[]} ")] -#[case(json!({"pages":[]}), "{\"pages\":[]}")] -#[case(json!({}), "{}")] +#[case(json!({"pages":[]}), "")] #[case(json!("[]"), "[]")] #[case(json!("{\"pages\":[{\"markdown\":\"json text\"}]}"), "json text")] #[case(json!({"pages":[{"markdown":"object"}]}), "object")] fn response_codec_handles_text_json_and_objects(#[case] content: Value, #[case] expected: &str) { + let structured = content + .as_object() + .is_some_and(|object| object.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}}), ) @@ -75,7 +86,11 @@ fn response_codec_handles_text_json_and_objects(#[case] content: Value, #[case] .into_json(); assert_eq!(result["pages"][0]["markdown"], expected); assert_eq!(result["pages"][0]["index"], 0); - assert_eq!(result["usage_info"]["prompt_tokens"], 1); + if structured { + assert!(result["usage_info"].is_null()); + } else { + assert_eq!(result["usage_info"]["prompt_tokens"], 1); + } } #[test] @@ -104,6 +119,7 @@ fn structured_result_maps_pages_usage_model_and_annotation() { #[test] fn response_codec_rejects_missing_empty_and_malformed_content() { for value in [ + json!({"choices":[{"message":{"content":{}}}]}), json!({"choices":[]}), json!({"choices":[{"message":{"content":""}}]}), json!({"choices":[{"message":{"content":"{\"pages\":[{\"markdown\":42}]}"}}]}), diff --git a/litellm-rust/crates/core/tests/host_lifecycle.rs b/litellm-rust/crates/core/tests/host_lifecycle.rs index 0e58462af1a..cdf9a7a2c8a 100644 --- a/litellm-rust/crates/core/tests/host_lifecycle.rs +++ b/litellm-rust/crates/core/tests/host_lifecycle.rs @@ -83,10 +83,10 @@ fn failure_handler_errors_do_not_replace_selected_failure_or_suppress_async_disp lifecycle.accept::(Ok(())); } let selected = Error::InvalidRequest("provider".into()); - assert_eq!( + assert!(matches!( lifecycle.accept(Err(HostFailure::Error(selected.clone()))), - Some(selected) - ); + Some(Error::InvalidRequest(message)) if message == "provider" + )); lifecycle.accept::(Ok(())); for phase in [ HostPhase::DeploymentFailure, @@ -94,11 +94,12 @@ fn failure_handler_errors_do_not_replace_selected_failure_or_suppress_async_disp HostPhase::AsyncFailure, ] { assert_eq!(lifecycle.phase(), phase); - assert_eq!( - lifecycle.accept(Err(HostFailure::Error(Error::InvalidRequest( - "callback".into() - )))), - None + assert!( + lifecycle + .accept(Err(HostFailure::Error(Error::InvalidRequest( + "callback".into() + )))) + .is_none() ); } assert_eq!(lifecycle.phase(), HostPhase::Complete); @@ -108,9 +109,9 @@ fn failure_handler_errors_do_not_replace_selected_failure_or_suppress_async_disp fn cancellation_skips_terminal_dispatch() { let mut lifecycle = HostLifecycle::new(true); let error = Error::InvalidRequest("cancelled".into()); - assert_eq!( + assert!(matches!( lifecycle.accept(Err(HostFailure::Cancelled(error.clone()))), - Some(error) - ); + Some(Error::InvalidRequest(message)) if message == "cancelled" + )); assert_eq!(lifecycle.phase(), HostPhase::Complete); } diff --git a/litellm-rust/crates/core/tests/ocr.rs b/litellm-rust/crates/core/tests/ocr.rs index a24d960422d..302ed91701e 100644 --- a/litellm-rust/crates/core/tests/ocr.rs +++ b/litellm-rust/crates/core/tests/ocr.rs @@ -63,8 +63,8 @@ async fn facade_executes_direct_mistral_once() { .await .unwrap(); server.await.unwrap(); - assert_eq!(result.pages[0]["markdown"], "hello"); - assert_eq!(result.pages[0]["custom"], "preserved"); + 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 ")); @@ -80,7 +80,8 @@ async fn facade_executes_direct_mistral_once() { "model":"model", "document":{"type":"document_url","document_url":"data:application/pdf;base64,YWJj"}, "pages":"0,2-4", - "extract_header":true + "extract_header":true, + "unknown":"ignored" }) ); } @@ -102,7 +103,10 @@ async fn facade_retains_native_response_when_requested() { .unwrap(); server.await.unwrap(); - assert_eq!(response.provider_native_response, Some(provider_response)); + assert_eq!( + response.provider_native_response.map(Value::Object), + Some(provider_response) + ); } #[tokio::test] @@ -348,7 +352,7 @@ async fn fallible_host_phases_do_not_replay_or_reach_transport() { } OcrHostOperation::ProjectRequest => { result = Some(OcrHostResult::Request(Ok(( - Box::new(request.take().unwrap().into()), + Box::new(request.take().unwrap()), false, )))) } @@ -406,7 +410,7 @@ async fn invalid_provider_response_runs_post_call_before_normalization_failure() match call.resume(result.take()).await { Ok(OcrCallStep::Host(OcrHostOperation::ProjectRequest)) => { result = Some(OcrHostResult::Request(Ok(( - Box::new(request.take().unwrap().into()), + Box::new(request.take().unwrap()), false, )))); } @@ -421,7 +425,7 @@ async fn invalid_provider_response_runs_post_call_before_normalization_failure() } }; server.await.unwrap(); - assert!(matches!(error, crate::ocr::Error::InvalidResponse(_))); + assert!(matches!(error, crate::ocr::Error::ResponseField { .. })); assert_eq!(seen.lock().unwrap().len(), 1); assert_eq!(post_calls, [json!(r#"{"pages":"invalid"}"#)]); } @@ -462,16 +466,15 @@ async fn direct_native_host_drives_the_same_state_machine() { OcrHostOperation::PostCall(_) => "PostCall".into(), OcrHostOperation::ConstructResponse(_) => "ConstructResponse".into(), OcrHostOperation::Success { response, .. } => { - assert_eq!(response.pages[0]["markdown"], "native"); + 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().into()), - false, - ))), + OcrHostOperation::ProjectRequest => { + OcrHostResult::Request(Ok((Box::new(request.take().unwrap()), false))) + } operation => host.invoke(operation).await, }); } @@ -479,7 +482,7 @@ async fn direct_native_host_drives_the_same_state_machine() { } }; server.await.unwrap(); - assert_eq!(response.pages[0]["markdown"], "native"); + assert_eq!(response.pages[0].markdown, "native"); assert_eq!(seen.lock().unwrap().len(), 1); assert_eq!( operations, @@ -556,7 +559,7 @@ async fn host_reader_documents_are_read_once_at_the_core_selected_point_and_enco ) .await; server.await.unwrap(); - assert_eq!(response.unwrap().pages[0]["markdown"], "file"); + assert_eq!(response.unwrap().pages[0].markdown, "file"); assert_eq!(reads, 1); assert!(seen.lock().unwrap()[0].contains("data:application/pdf;base64,YWJj")); } @@ -571,7 +574,9 @@ async fn host_reader_failures_and_empty_files_fail_before_the_provider_is_called Err(failure.clone()), ) .await; - assert_eq!(response.unwrap_err(), failure); + assert!( + matches!(response.unwrap_err(), crate::ocr::Error::InvalidRequest(message) if message == "reader exploded") + ); assert_eq!(reads, 1); let request = wire_request("mistral/model", &base, json!({})); @@ -585,7 +590,7 @@ async fn host_reader_failures_and_empty_files_fail_before_the_provider_is_called .await; assert!(matches!( response.unwrap_err(), - crate::ocr::Error::InvalidRequest(_) + crate::ocr::Error::EmptyFile )); assert!(seen.lock().unwrap().is_empty()); } @@ -613,7 +618,7 @@ async fn path_documents_are_read_by_core_without_a_host_operation() { .await; server.await.unwrap(); std::fs::remove_dir_all(&dir).unwrap(); - assert_eq!(response.unwrap().pages[0]["markdown"], "path"); + assert_eq!(response.unwrap().pages[0].markdown, "path"); assert_eq!(reads, 0); assert!(seen.lock().unwrap()[0].contains("data:image/png;base64,YWJj")); @@ -629,7 +634,7 @@ async fn path_documents_are_read_by_core_without_a_host_operation() { .await; assert!(matches!( response.unwrap_err(), - crate::ocr::Error::FileRead { path: failed, kind: std::io::ErrorKind::NotFound, .. } if failed == path + crate::ocr::Error::FileRead { path: failed, source } if failed == path && source.kind() == std::io::ErrorKind::NotFound )); assert!(seen.lock().unwrap().is_empty()); } @@ -661,7 +666,9 @@ async fn public_finalization_failure_never_dispatches_success_or_replays_provide OcrHostResult::Lifecycle(Err(HostFailure::Error(selected.clone()))) } OcrHostOperation::Failure { error, .. } => { - assert_eq!(error, selected); + assert!( + matches!(error, crate::ocr::Error::InvalidRequest(message) if message == "public metadata failed") + ); failures.push("sync"); OcrHostResult::Lifecycle(Err(HostFailure::Error( crate::ocr::Error::InvalidRequest("failure callback failed".into()), @@ -676,10 +683,9 @@ async fn public_finalization_failure_never_dispatches_success_or_replays_provide | OcrHostOperation::Lifecycle(HostPhase::DeploymentFailure) => { panic!("finalization failure used provider/success dispatch") } - OcrHostOperation::ProjectRequest => OcrHostResult::Request(Ok(( - Box::new(request.take().unwrap().into()), - false, - ))), + OcrHostOperation::ProjectRequest => { + OcrHostResult::Request(Ok((Box::new(request.take().unwrap()), false))) + } operation => host.invoke(operation).await, }); } @@ -688,7 +694,9 @@ async fn public_finalization_failure_never_dispatches_success_or_replays_provide } }; server.await.unwrap(); - assert_eq!(error, selected); + assert!( + matches!(error, crate::ocr::Error::InvalidRequest(message) if message == "public metadata failed") + ); assert_eq!(failures, ["sync", "async"]); assert_eq!(seen.lock().unwrap().len(), 1); } @@ -716,7 +724,7 @@ async fn cancellation_at_provider_hook_prevents_execution_and_further_resumption OcrCallStep::Host(OcrHostOperation::PreCall(_)) => break, OcrCallStep::Host(OcrHostOperation::ProjectRequest) => { result = Some(OcrHostResult::Request(Ok(( - Box::new(request.take().unwrap().into()), + Box::new(request.take().unwrap()), false, )))) } @@ -727,7 +735,7 @@ async fn cancellation_at_provider_hook_prevents_execution_and_further_resumption let selected = crate::ocr::Error::InvalidRequest("cancelled".into()); assert!(matches!( call.interrupt(HostFailure::Cancelled(selected.clone())).await, - Err(error) if error == selected + Err(crate::ocr::Error::InvalidRequest(message)) if message == "cancelled" )); assert!( call.resume(Some(OcrHostResult::Lifecycle(Ok(())))) @@ -761,7 +769,7 @@ async fn missing_host_result_preserves_pending_operation() { async fn read_bounded_response( response: Vec, limit: usize, -) -> Result { +) -> Result { use tokio::io::{AsyncReadExt, AsyncWriteExt}; let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); @@ -790,7 +798,7 @@ async fn read_bounded_response( #[tokio::test] async fn response_limit_accepts_exact_size_and_rejects_declared_and_chunked_overflow() { - use super::error::{OcrError, OcrResponseError}; + use super::Error; for response in [ "HTTP/1.1 200 OK\r\nContent-Length: 8\r\n\r\nabcdefgh", @@ -809,7 +817,7 @@ async fn response_limit_accepts_exact_size_and_rejects_declared_and_chunked_over ] { assert!(matches!( read_bounded_response(response.as_bytes().to_vec(), 8).await, - Err(OcrError::Response(OcrResponseError::TooLarge { limit: 8 })) + Err(Error::TooLarge { limit: 8 }) )); } } @@ -828,7 +836,7 @@ async fn oversized_error_retains_http_status_and_bounded_diagnostics_without_dra .await .unwrap_err(); match error { - super::error::OcrError::Transport(crate::transport::Error::Http { status, body }) => { + super::Error::Transport(crate::transport::Error::Http { status, body }) => { assert_eq!(status, 429); assert_eq!( body, @@ -850,7 +858,7 @@ fn response_limit_is_validated_and_not_forwarded_to_the_provider() { "http://localhost", json!({"max_response_bytes": 123}), ); - assert_eq!(request.connection.max_response_bytes, 123); + assert_eq!(request.transport.max_response_bytes, 123); assert!(!request.optional_params.contains_key("max_response_bytes")); for value in [ json!(0), @@ -908,9 +916,9 @@ async fn cancellation_waits_for_provider_capture_drop_even_when_acknowledgement_ 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 { + transport: super::OcrTransportConfig { extra_headers: vec![("authorization".into(), "Bearer test-key".into())], - ..request.connection + ..request.transport }, azure_ad_token_provider: Some(litellm_auth::TokenProviderHandle::new(Arc::new( PendingToken { @@ -933,7 +941,7 @@ async fn cancellation_waits_for_provider_capture_drop_even_when_acknowledgement_ _ = 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().into()), false))), + 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"), }); @@ -960,7 +968,9 @@ async fn cancellation_waits_for_provider_capture_drop_even_when_acknowledgement_ ) .await .unwrap(); - assert!(matches!(result, Err(error) if error == selected)); + assert!( + matches!(result, Err(crate::ocr::Error::InvalidRequest(message)) if message == "cancelled") + ); assert!( dropped.load(Ordering::SeqCst), "cancellation returned while provider captures were still alive" diff --git a/litellm-rust/crates/core/tests/ocr/support.rs b/litellm-rust/crates/core/tests/ocr/support.rs index c7b64e300f0..44fd0462bbf 100644 --- a/litellm-rust/crates/core/tests/ocr/support.rs +++ b/litellm-rust/crates/core/tests/ocr/support.rs @@ -36,6 +36,20 @@ pub(crate) fn wire_request(model: &str, base: &str, options: Value) -> LiteLLMOc .unwrap() } +pub(crate) fn resolved_request( + request: LiteLLMOcrRequest, +) -> crate::ocr::types::ResolvedOcrRequest { + request + .map_document(crate::ocr::document::prepare_document) + .unwrap() +} + +pub(crate) fn with_source(request: LiteLLMOcrRequest, source: &str) -> LiteLLMOcrRequest { + let request = resolved_request(request); + let document = request.document.clone().with_source(source.into()); + request.with_document(document.into()) +} + pub(crate) struct MockResponse { pub status: u16, pub headers: Vec<(&'static str, String)>, diff --git a/litellm-rust/crates/core/tests/reducto_ocr.rs b/litellm-rust/crates/core/tests/reducto_ocr.rs index a15e9cae5b5..0a7053b7429 100644 --- a/litellm-rust/crates/core/tests/reducto_ocr.rs +++ b/litellm-rust/crates/core/tests/reducto_ocr.rs @@ -56,8 +56,7 @@ async fn request_mapping_matches_python( "result":{"chunks":[]} }))]) .await; - let mut request = wire_request(model, &base, options); - request.document = request.document.with_source(source.into()); + let request = super::test_support::with_source(wire_request(model, &base, options), source); perform_ocr(request).await.unwrap(); server.await.unwrap(); @@ -78,14 +77,14 @@ async fn data_uri_upload_preserves_multipart_headers(#[case] model: &str) { ]) .await; let mut request = wire_request(&format!("reducto/{model}"), &base, json!({})); - request.connection.extra_headers = vec![ + request.transport.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"); + assert_eq!(response.pages[0].markdown, "hello"); let requests = seen.lock().unwrap(); assert_eq!(requests.len(), 2); assert!(requests[0].starts_with("POST /upload ")); @@ -175,14 +174,18 @@ async fn upload_failure_stops_before_parse() { #[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()); + let request = super::test_support::with_source( + wire_request("reducto/parse-v3", "http://127.0.0.1:1", json!({})), + source, + ); assert!(perform_ocr(request).await.is_err()); } #[test] fn response_normalization_groups_blocks_and_distinguishes_null_result() { - use crate::ocr::codecs::reducto::{ReductoResponse, transform_ocr_response}; + use crate::llms::reducto::ocr::transformation::{ + ReductoResponse, normalize_response as transform_ocr_response, + }; let raw = json!({"usage":{"num_pages":"2","credits":"3"},"result":{"type":"full","chunks":[ {"blocks":[{ @@ -218,7 +221,7 @@ fn response_normalization_groups_blocks_and_distinguishes_null_result() { let missing: ReductoResponse = serde_json::from_value(json!({"chunks":[{"content":"text"}]})).unwrap(); let missing = transform_ocr_response("parse-v3", missing).unwrap(); - assert_eq!(missing.pages[0]["markdown"], "text"); + assert_eq!(missing.pages[0].markdown, "text"); let null: ReductoResponse = serde_json::from_value( json!({"result":null,"chunks":[{"content":"ignored"}],"usage":null}), ) @@ -231,9 +234,11 @@ fn response_normalization_groups_blocks_and_distinguishes_null_result() { 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 mut request = super::test_support::with_source( + wire_request("reducto/parse-v3", &base, json!({})), + "reducto://ready.pdf", + ); + request.transport.extra_headers = vec![("authorization".into(), "Bearer existing".into())]; let response = perform_ocr(request).await.unwrap(); server.await.unwrap(); diff --git a/litellm-rust/crates/core/tests/vertex_ai_deepseek_ocr.rs b/litellm-rust/crates/core/tests/vertex_ai_deepseek_ocr.rs index a73c1e7710a..be0898e1135 100644 --- a/litellm-rust/crates/core/tests/vertex_ai_deepseek_ocr.rs +++ b/litellm-rust/crates/core/tests/vertex_ai_deepseek_ocr.rs @@ -14,7 +14,7 @@ async fn facade_executes_vertex_deepseek_at_the_openai_endpoint() { "usage":{"prompt_tokens":1} }))]) .await; - let mut request = wire_request( + let request = wire_request( "vertex_ai/deepseek-ocr-maas", &base, json!({ @@ -25,14 +25,15 @@ async fn facade_executes_vertex_deepseek_at_the_openai_endpoint() { "extra_body":{"provider_option":"value"} }), ); - request.document = request - .document - .with_source("gs://bucket/document.pdf".into()); + let request = super::test_support::with_source(request, "gs://bucket/document.pdf"); let response = perform_ocr(request).await.unwrap(); server.await.unwrap(); - assert_eq!(response.pages[0]["markdown"], "recognized"); - assert_eq!(response.usage_info.unwrap()["prompt_tokens"], 1); + 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 " @@ -45,7 +46,7 @@ async fn facade_executes_vertex_deepseek_at_the_openai_endpoint() { let body = request_body(&requests[0]); assert_eq!(body["model"], "deepseek-ai/deepseek-ocr-maas"); assert_eq!(body["temperature"], 0.1); - assert!(body.get("future_ocr_option").is_none()); + assert_eq!(body["future_ocr_option"], true); assert!(body.get("extra_body").is_none()); assert_eq!( body["messages"][0]["content"][0], @@ -72,7 +73,10 @@ async fn request_controlled_api_base_is_rejected_before_vertex_auth() { "https://caller.example", json!({"vertex_project":"project-1"}), ); - request.connection.api_base_source = InputSource::Request; + request.credentials.api_base = Some(litellm_auth::Sourced::new( + "https://caller.example".into(), + InputSource::Request, + )); let error = perform_ocr(request).await.unwrap_err(); assert!( diff --git a/litellm-rust/crates/core/tests/vertex_ai_ocr.rs b/litellm-rust/crates/core/tests/vertex_ai_ocr.rs index 93e9efca849..27e4802b00d 100644 --- a/litellm-rust/crates/core/tests/vertex_ai_ocr.rs +++ b/litellm-rust/crates/core/tests/vertex_ai_ocr.rs @@ -26,7 +26,7 @@ async fn facade_executes_vertex_mistral_with_resolved_project_and_location() { let response = perform_ocr(request).await.unwrap(); server.await.unwrap(); - assert_eq!(response.pages[0]["markdown"], "hello"); + assert_eq!(response.pages[0].markdown, "hello"); let requests = seen.lock().unwrap(); assert_eq!(requests.len(), 1); assert!(requests[0].starts_with( @@ -55,8 +55,8 @@ async fn supplied_authorization_is_forwarded_without_a_static_token() { &base, json!({"vertex_project":"project-1"}), ); - request.connection.api_key = None; - request.connection.extra_headers = vec![("authorization".into(), "Bearer supplied".into())]; + request.credentials.api_key = None; + request.transport.extra_headers = vec![("authorization".into(), "Bearer supplied".into())]; perform_ocr(request).await.unwrap(); server.await.unwrap(); @@ -85,7 +85,10 @@ async fn request_controlled_api_base_is_rejected_before_vertex_auth() { "https://caller.example", json!({"vertex_project":"project-1"}), ); - request.connection.api_base_source = InputSource::Request; + request.credentials.api_base = Some(litellm_auth::Sourced::new( + "https://caller.example".into(), + InputSource::Request, + )); let error = perform_ocr(request).await.unwrap_err(); assert!( @@ -99,7 +102,9 @@ async fn request_controlled_api_base_is_rejected_before_vertex_auth() { async fn adapters_build_complete_requests_and_share_mistral_normalization() { use std::time::Duration; - use crate::ocr::adapters::{MistralAdapter, OcrAdapter, VertexMistralAdapter}; + 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(); @@ -116,11 +121,15 @@ async fn adapters_build_complete_requests_and_share_mistral_normalization() { options.clone(), ); let vertex = wire_request("vertex_ai/mistral-ocr-maas", "https://vertex.test", options); - let direct_http = MistralAdapter + let direct = + crate::ocr::prepare::prepare_request(super::test_support::resolved_request(direct)); + let vertex = + crate::ocr::prepare::prepare_request(super::test_support::resolved_request(vertex)); + let direct_http = MistralOCRConfig .prepare_request(&direct, &client) .await .unwrap(); - let vertex_http = VertexMistralAdapter + let vertex_http = VertexAIOCRConfig .prepare_request(&vertex, &client) .await .unwrap(); @@ -141,17 +150,27 @@ async fn adapters_build_complete_requests_and_share_mistral_normalization() { "model": "mistral-ocr-maas", "document": {"type": "document_url", "document_url": "data:application/pdf;base64,YWJj"}, "pages": [0, 2], - "include_image_base64": true + "include_image_base64": true, + "unknown": "ignored" }) ); } let payload = json!({"pages": [{"index": 0, "markdown": "hello"}], "extra": "preserved"}); - let direct_response = MistralAdapter - .transform_ocr_response(&direct, serde_json::from_value(payload.clone()).unwrap()) + let raw = serde_json::to_vec(&payload).unwrap(); + let direct_response = MistralOCRConfig + .transform_ocr_response( + &direct.model, + &raw, + crate::ocr::types::OcrResponseFormat::Litellm, + ) .unwrap() .into_json(); - let vertex_response = VertexMistralAdapter - .transform_ocr_response(&vertex, serde_json::from_value(payload).unwrap()) + let vertex_response = VertexAIOCRConfig + .transform_ocr_response( + &vertex.model, + &raw, + crate::ocr::types::OcrResponseFormat::Litellm, + ) .unwrap() .into_json(); assert_eq!(direct_response, vertex_response); diff --git a/litellm-rust/crates/python-bridge/src/errors.rs b/litellm-rust/crates/python-bridge/src/errors.rs index 7ca86b3ccfa..d54b2755e89 100644 --- a/litellm-rust/crates/python-bridge/src/errors.rs +++ b/litellm-rust/crates/python-bridge/src/errors.rs @@ -35,15 +35,17 @@ pub(crate) fn responses_error_to_pyerr(error: responses::Error) -> PyErr { pub(crate) fn core_error_to_pyerr(error: Error) -> PyErr { let value_error = match &error { - Error::Ocr(error) => matches!( - error, - ocr::Error::Auth(_) - | ocr::Error::InvalidProvider(_) - | ocr::Error::InvalidRequest(_) - | ocr::Error::InvalidType { .. } - | ocr::Error::MissingField(_) - | ocr::Error::MissingDocumentUrl - ), + Error::Ocr(error) => { + error.is_request() + || matches!( + error, + ocr::Error::Auth(_) + | ocr::Error::InvalidProvider(_) + | ocr::Error::InvalidRequest(_) + | ocr::Error::MissingField(_) + | ocr::Error::MissingDocumentUrl + ) + } Error::Messages(error) => match error { messages::Error::Auth(source) => auth_is_value_error(source), messages::Error::InvalidProvider(_) 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 7dbc35289ff..d943a053a61 100644 --- a/litellm-rust/crates/python-bridge/src/routes/ocr/errors.rs +++ b/litellm-rust/crates/python-bridge/src/routes/ocr/errors.rs @@ -7,13 +7,14 @@ use crate::errors::{RustUpstreamError, core_error_to_pyerr}; pub(super) fn to_pyerr(error: Error) -> PyErr { let status = error.http_status_code(); let mapped = match error { - Error::Http { status, body } => RustUpstreamError::new_err((status, body)), - Error::FileRead { - path, - kind: std::io::ErrorKind::NotFound, - .. - } => PyFileNotFoundError::new_err(format!("File not found: {}", path.display())), - Error::FileRead { message, .. } => PyOSError::new_err(message), + Error::Provider { status, body, .. } + | Error::Transport(litellm_core::transport::Error::Http { status, body }) => { + RustUpstreamError::new_err((status, body)) + } + Error::FileRead { path, source } if source.kind() == std::io::ErrorKind::NotFound => { + PyFileNotFoundError::new_err(format!("File not found: {}", path.display())) + } + Error::FileRead { source, .. } => PyOSError::new_err(source.to_string()), other => core_error_to_pyerr(other.into()), }; attach_status(mapped, status) @@ -51,9 +52,10 @@ mod tests { .unwrap(), 500 ); - let mapped = to_pyerr(Error::Http { + let mapped = to_pyerr(Error::Provider { status: 429, body: r#"{"message":"rate limited"}"#.to_string(), + headers: Vec::new(), }); assert!(mapped.is_instance_of::(py)); let args: (u16, String) = mapped diff --git a/litellm-rust/crates/python-bridge/src/routes/ocr/project.rs b/litellm-rust/crates/python-bridge/src/routes/ocr/project.rs index ad223645c62..3076895c1c4 100644 --- a/litellm-rust/crates/python-bridge/src/routes/ocr/project.rs +++ b/litellm-rust/crates/python-bridge/src/routes/ocr/project.rs @@ -215,7 +215,7 @@ mod tests { fn url_document(url: &str) -> OcrDocumentInput { litellm_core::ocr::OcrDocument::DocumentUrl { document_url: url.into(), - extra_fields: Map::new(), + extra_fields: Default::default(), } .into() } From e0ce9980912b9f6f77e1123d019f82628bc6c9ab Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Wed, 16 Sep 2026 20:21:51 -0700 Subject: [PATCH 02/15] fmt --- .../crates/core/src/audio_transcription/handler.rs | 3 +-- .../crates/core/src/audio_transcription/mod.rs | 3 +-- .../crates/core/src/audio_transcription/prepare.rs | 5 ++--- .../core/src/audio_transcription/transformation.rs | 2 +- litellm-rust/crates/core/src/call_arguments.rs | 3 ++- litellm-rust/crates/core/src/call_lifecycle/mod.rs | 3 ++- .../crates/core/src/chat_completions/common_utils.rs | 6 +++--- .../crates/core/src/chat_completions/conversation.rs | 6 +++--- .../crates/core/src/chat_completions/handler.rs | 3 +-- litellm-rust/crates/core/src/chat_completions/mod.rs | 3 +-- .../crates/core/src/chat_completions/prepare.rs | 5 ++--- .../crates/core/src/chat_completions/tests.rs | 3 +-- .../core/src/chat_completions/transformation.rs | 2 +- litellm-rust/crates/core/src/http_utils.rs | 3 ++- .../llms/azure_ai/ocr/cohere_parse_transformation.rs | 3 ++- .../core/src/llms/azure_ai/ocr/common_utils.rs | 3 ++- .../ocr/document_intelligence/transformation.rs | 11 ++++++----- .../core/src/llms/azure_ai/ocr/transformation.rs | 7 ++++--- .../core/src/llms/cohere/ocr/transformation.rs | 3 ++- .../core/src/llms/mistral/ocr/transformation.rs | 3 ++- .../core/src/llms/vertex_ai/ocr/common_utils.rs | 3 ++- .../llms/vertex_ai/ocr/deepseek_transformation.rs | 12 +++++++----- .../core/src/llms/vertex_ai/ocr/transformation.rs | 2 +- litellm-rust/crates/core/src/media.rs | 4 +++- .../crates/core/src/messages/common_utils.rs | 9 ++++----- litellm-rust/crates/core/src/messages/handler.rs | 5 ++--- litellm-rust/crates/core/src/messages/prepare.rs | 6 +++--- litellm-rust/crates/core/src/messages/tests.rs | 1 - litellm-rust/crates/core/src/ocr/arguments.rs | 3 +-- litellm-rust/crates/core/src/ocr/client.rs | 2 +- litellm-rust/crates/core/src/ocr/document.rs | 5 +++-- litellm-rust/crates/core/src/ocr/hooks.rs | 5 +++-- litellm-rust/crates/core/src/ocr/lifecycle.rs | 4 ++-- litellm-rust/crates/core/src/ocr/prepare.rs | 3 ++- litellm-rust/crates/core/src/ocr/provider_config.rs | 6 ++++-- litellm-rust/crates/core/src/ocr/types.rs | 6 +++--- .../providers/anthropic/chat_completions/tests.rs | 3 ++- .../anthropic/chat_completions/transformation.rs | 3 +-- .../providers/azure_ai/messages/transformation.rs | 6 ++++-- .../src/providers/bedrock/audio_transcription.rs | 5 ++--- .../src/providers/bedrock/chat_completions/tests.rs | 3 ++- .../bedrock/chat_completions/transformation.rs | 5 ++--- litellm-rust/crates/core/src/serde_compat.rs | 3 ++- .../core/tests/azure_document_intelligence_ocr.rs | 6 ++++-- litellm-rust/crates/core/tests/ocr.rs | 3 ++- .../crates/core/tests/vertex_ai_deepseek_ocr.rs | 2 +- litellm-rust/crates/core/tests/vertex_ai_ocr.rs | 2 +- 47 files changed, 105 insertions(+), 92 deletions(-) diff --git a/litellm-rust/crates/core/src/audio_transcription/handler.rs b/litellm-rust/crates/core/src/audio_transcription/handler.rs index bd1740a8b93..2a7afccf9ea 100644 --- a/litellm-rust/crates/core/src/audio_transcription/handler.rs +++ b/litellm-rust/crates/core/src/audio_transcription/handler.rs @@ -1,10 +1,9 @@ use serde_json::Value; use super::Error; -use crate::http_utils::{http_request, truncate_error_body}; - use super::client::http_client; use super::types::ProviderAudioTranscriptionRequest; +use crate::http_utils::{http_request, truncate_error_body}; pub async fn execute_audio_transcription_provider_call( request: ProviderAudioTranscriptionRequest, diff --git a/litellm-rust/crates/core/src/audio_transcription/mod.rs b/litellm-rust/crates/core/src/audio_transcription/mod.rs index 87f6c41d80f..47b1e8bb151 100644 --- a/litellm-rust/crates/core/src/audio_transcription/mod.rs +++ b/litellm-rust/crates/core/src/audio_transcription/mod.rs @@ -6,10 +6,9 @@ mod prepare; pub mod transformation; pub mod types; -use serde_json::Value; - pub use handler::execute_audio_transcription_provider_call; pub use prepare::prepare_audio_transcription_provider_call; +use serde_json::Value; pub use types::{AudioTranscriptionRequest, ProviderAudioTranscriptionRequest}; pub async fn audio_transcription(request: AudioTranscriptionRequest<'_>) -> Result { diff --git a/litellm-rust/crates/core/src/audio_transcription/prepare.rs b/litellm-rust/crates/core/src/audio_transcription/prepare.rs index 82f85ba85ce..416ada2491e 100644 --- a/litellm-rust/crates/core/src/audio_transcription/prepare.rs +++ b/litellm-rust/crates/core/src/audio_transcription/prepare.rs @@ -1,11 +1,10 @@ use super::Error; +use super::transformation::{AudioTranscriptionAuth, AudioTranscriptionProviderConfig}; +use super::types::{AudioTranscriptionRequest, ProviderAudioTranscriptionRequest}; use crate::http_utils::{has_header, string_headers}; use crate::providers::bedrock::audio_transcription::BEDROCK_AUDIO_TRANSCRIPTION_CONFIG; use crate::providers::custom_llm_provider::{CustomLlmProvider, get_custom_llm_provider}; -use super::transformation::{AudioTranscriptionAuth, AudioTranscriptionProviderConfig}; -use super::types::{AudioTranscriptionRequest, ProviderAudioTranscriptionRequest}; - fn provider_config(provider: &str) -> Option<&'static dyn AudioTranscriptionProviderConfig> { if provider == "bedrock" { return Some(&BEDROCK_AUDIO_TRANSCRIPTION_CONFIG); diff --git a/litellm-rust/crates/core/src/audio_transcription/transformation.rs b/litellm-rust/crates/core/src/audio_transcription/transformation.rs index a849f052e12..f8082991241 100644 --- a/litellm-rust/crates/core/src/audio_transcription/transformation.rs +++ b/litellm-rust/crates/core/src/audio_transcription/transformation.rs @@ -1,6 +1,6 @@ -use super::Error; use serde_json::{Map, Value}; +use super::Error; use super::types::{AudioTranscriptionRequestData, AudioTranscriptionResponseData}; #[derive(Clone, Debug, PartialEq, Eq)] diff --git a/litellm-rust/crates/core/src/call_arguments.rs b/litellm-rust/crates/core/src/call_arguments.rs index 67852cef27d..3b9183c739a 100644 --- a/litellm-rust/crates/core/src/call_arguments.rs +++ b/litellm-rust/crates/core/src/call_arguments.rs @@ -381,9 +381,10 @@ impl IntoIterator for CallArguments { #[cfg(test)] mod tests { - use super::*; use serde_json::json; + use super::*; + #[test] fn composition_preserves_extensions_and_applies_shallow_explicit_overrides() { let original = json!({ diff --git a/litellm-rust/crates/core/src/call_lifecycle/mod.rs b/litellm-rust/crates/core/src/call_lifecycle/mod.rs index dce240c3d2b..e012961e005 100644 --- a/litellm-rust/crates/core/src/call_lifecycle/mod.rs +++ b/litellm-rust/crates/core/src/call_lifecycle/mod.rs @@ -228,10 +228,11 @@ fn epoch_seconds() -> f64 { #[cfg(test)] mod tests { - use super::*; use std::pin::Pin; use std::sync::Mutex; + use super::*; + type BoxFuture<'a, T> = Pin + Send + 'a>>; #[derive(Default)] diff --git a/litellm-rust/crates/core/src/chat_completions/common_utils.rs b/litellm-rust/crates/core/src/chat_completions/common_utils.rs index 9ebc5ae0efa..c89450aeb77 100644 --- a/litellm-rust/crates/core/src/chat_completions/common_utils.rs +++ b/litellm-rust/crates/core/src/chat_completions/common_utils.rs @@ -1,9 +1,9 @@ -use super::Error; -use crate::http_utils::string_headers as shared_string_headers; -use crate::providers::anthropic::chat_completions::transformation::ANTHROPIC_CHAT_COMPLETIONS_CONFIG; use serde_json::{Map, Value}; +use super::Error; use super::transformation::ChatCompletionsProviderConfig; +use crate::http_utils::string_headers as shared_string_headers; +use crate::providers::anthropic::chat_completions::transformation::ANTHROPIC_CHAT_COMPLETIONS_CONFIG; const HEADER_CONTEXT: &str = "chat completions"; diff --git a/litellm-rust/crates/core/src/chat_completions/conversation.rs b/litellm-rust/crates/core/src/chat_completions/conversation.rs index f7bdc60af37..1f1984ed8be 100644 --- a/litellm-rust/crates/core/src/chat_completions/conversation.rs +++ b/litellm-rust/crates/core/src/chat_completions/conversation.rs @@ -10,9 +10,8 @@ //! `_bedrock_converse_messages_pt` for the text-only surface this route //! accepts; anything richer is declined upstream by the capability gate. -use crate::constants::EMPTY_TEXT_PLACEHOLDER; - use super::types::{ChatMessage, ChatMessageContent}; +use crate::constants::EMPTY_TEXT_PLACEHOLDER; #[derive(Clone, Copy, Debug, PartialEq, Eq)] pub enum TurnRole { @@ -132,9 +131,10 @@ pub fn build_conversation(messages: &[ChatMessage]) -> Conversation { #[cfg(test)] mod tests { - use super::*; use serde_json::json; + use super::*; + fn messages(value: serde_json::Value) -> Vec { serde_json::from_value(value).expect("valid messages") } diff --git a/litellm-rust/crates/core/src/chat_completions/handler.rs b/litellm-rust/crates/core/src/chat_completions/handler.rs index d4527e99a10..2d192e971b0 100644 --- a/litellm-rust/crates/core/src/chat_completions/handler.rs +++ b/litellm-rust/crates/core/src/chat_completions/handler.rs @@ -1,8 +1,6 @@ use serde_json::Value; use super::Error; -use crate::http_utils::{http_request, truncate_error_body}; - use super::client::http_client; use super::prepare::prepare_provider_request; use super::transformation::ChatCompletionsAuth; @@ -10,6 +8,7 @@ use super::types::{ ChatCompletionsResponse, ProviderChatCompletionsRequest, ProviderChatResponseData, ResolvedChatCompletionsRequest, }; +use crate::http_utils::{http_request, truncate_error_body}; pub(super) async fn execute_chat_completions_provider_call( request: ResolvedChatCompletionsRequest<'_>, diff --git a/litellm-rust/crates/core/src/chat_completions/mod.rs b/litellm-rust/crates/core/src/chat_completions/mod.rs index 401eef609f2..b31ceaffb5c 100644 --- a/litellm-rust/crates/core/src/chat_completions/mod.rs +++ b/litellm-rust/crates/core/src/chat_completions/mod.rs @@ -17,10 +17,9 @@ pub mod response_utils; pub mod transformation; pub mod types; -use serde_json::{Map, Value}; - use handler::execute_chat_completions_provider_call; use prepare::{parse_messages, resolve_provider_config, resolve_request}; +use serde_json::{Map, Value}; use types::{ChatCompletionsRequest, ChatCompletionsResponse}; pub async fn chat_completions( diff --git a/litellm-rust/crates/core/src/chat_completions/prepare.rs b/litellm-rust/crates/core/src/chat_completions/prepare.rs index e8d8d70f271..b2360021ef7 100644 --- a/litellm-rust/crates/core/src/chat_completions/prepare.rs +++ b/litellm-rust/crates/core/src/chat_completions/prepare.rs @@ -1,15 +1,14 @@ use serde_json::Value; use super::Error; -use crate::http_utils::has_header; -use crate::providers::custom_llm_provider::{CustomLlmProvider, get_custom_llm_provider}; - use super::common_utils::{chat_completions_provider_config, string_headers}; use super::transformation::{ChatCompletionsAuth, ChatCompletionsProviderConfig}; use super::types::{ ChatCompletionsRequest, ChatMessage, ProviderChatCompletionsRequest, ResolvedChatCompletionsRequest, }; +use crate::http_utils::has_header; +use crate::providers::custom_llm_provider::{CustomLlmProvider, get_custom_llm_provider}; pub(super) fn resolve_provider_config<'a>( model: &'a str, diff --git a/litellm-rust/crates/core/src/chat_completions/tests.rs b/litellm-rust/crates/core/src/chat_completions/tests.rs index 39fabe27f44..b860b5f7206 100644 --- a/litellm-rust/crates/core/src/chat_completions/tests.rs +++ b/litellm-rust/crates/core/src/chat_completions/tests.rs @@ -1,7 +1,6 @@ use serde_json::{Map, Value, json}; use super::Error; - use super::prepare::{prepare_provider_request, resolve_request}; use super::transformation::ChatCompletionsAuth; use super::types::{ChatCompletionsRequest, ProviderChatCompletionsRequest}; @@ -588,10 +587,10 @@ fn the_gate_agrees_with_prepare_on_every_case_it_accepts() { } mod round_trip { - use super::*; use tokio::io::{AsyncReadExt, AsyncWriteExt}; use tokio::net::{TcpListener, TcpStream}; + use super::*; use crate::chat_completions::chat_completions; async fn read_http_request(socket: &mut TcpStream) -> String { diff --git a/litellm-rust/crates/core/src/chat_completions/transformation.rs b/litellm-rust/crates/core/src/chat_completions/transformation.rs index 1000dbaa673..2325e22e019 100644 --- a/litellm-rust/crates/core/src/chat_completions/transformation.rs +++ b/litellm-rust/crates/core/src/chat_completions/transformation.rs @@ -1,6 +1,6 @@ -use super::Error; use serde_json::{Map, Value}; +use super::Error; use super::types::{ ChatCompletionsResponse, ChatMessage, ChatMessageContent, ProviderChatRequestData, ProviderChatResponseData, diff --git a/litellm-rust/crates/core/src/http_utils.rs b/litellm-rust/crates/core/src/http_utils.rs index 53d2f961bd5..060559322ea 100644 --- a/litellm-rust/crates/core/src/http_utils.rs +++ b/litellm-rust/crates/core/src/http_utils.rs @@ -131,9 +131,10 @@ pub fn json_type_name(value: &serde_json::Value) -> &'static str { #[cfg(test)] mod tests { - use super::*; use serde_json::json; + use super::*; + #[rstest::rstest] #[case(HeaderPolicy::All, true, true)] #[case(HeaderPolicy::Only(&["authorization"]), true, false)] 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 add70c2596d..bdd18cbf4df 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,3 +1,5 @@ +use serde_json::Value; + use crate::call_arguments::CallArguments; use crate::llms::base_llm::ocr::transformation::{BaseOcrConfig, OcrRequestContext}; use crate::llms::cohere::ocr::transformation::{CohereParseConfig, CohereRequest}; @@ -6,7 +8,6 @@ use crate::ocr::OcrClient; use crate::ocr::document::{inline_remote_document, validate_inline_document}; use crate::ocr::types::{LiteLLMOcrResponse, OcrDocument, PreparedOcrRequest}; use crate::url_utils::ApiUrl; -use serde_json::Value; #[derive(Default)] pub(crate) struct AzureAICohereParseConfig; diff --git a/litellm-rust/crates/core/src/llms/azure_ai/ocr/common_utils.rs b/litellm-rust/crates/core/src/llms/azure_ai/ocr/common_utils.rs index c381e39eaae..4e7be1620ae 100644 --- a/litellm-rust/crates/core/src/llms/azure_ai/ocr/common_utils.rs +++ b/litellm-rust/crates/core/src/llms/azure_ai/ocr/common_utils.rs @@ -1,9 +1,10 @@ use std::sync::OnceLock; -use crate::ocr::types::OcrConnection; use litellm_auth::{InputSource, Sourced}; use litellm_auth_azure::{AzureAuthInputs, AzureAuthService}; +use crate::ocr::types::OcrConnection; + pub(super) async fn resolve_entra( config: &AzureAuthInputs, env_lookup: &(dyn Fn(&str) -> Option + Sync), 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 ae13944c06b..e20ec29132d 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 @@ -3,15 +3,14 @@ use std::sync::Arc; use std::time::Duration; use base64::{Engine, engine::general_purpose::STANDARD}; +use litellm_auth::{InputSource, Sourced}; +use litellm_auth_azure::AzureAuthInputs; use reqwest::Url; use serde::{Deserialize, Deserializer, Serialize}; use serde_json::{Map, Value}; use serde_with::serde_as; use tokio::time::Instant; -use litellm_auth::{InputSource, Sourced}; -use litellm_auth_azure::AzureAuthInputs; - use crate::call_arguments::CallArguments; use crate::constants::{ AZURE_DI_API_VERSION, AZURE_DI_DEFAULT_DPI, AZURE_DI_DEFAULT_HEIGHT, AZURE_DI_DEFAULT_WIDTH, @@ -632,10 +631,11 @@ fn nonblank(value: Option) -> Option { #[cfg(test)] mod tests { - use super::*; use rstest::rstest; use serde_json::{Value, json}; + use super::*; + fn map(value: Value) -> Result { let arguments = serde_json::from_value(value).unwrap(); AzureDocumentIntelligenceOCRConfig.map_ocr_params(&arguments, "model") @@ -1220,9 +1220,10 @@ mod tests { #[tokio::test] async fn pre_call_guardrail_receives_caller_pages_before_mapping() { - use crate::ocr::hooks::{OcrHookFuture, OcrHooks, OcrPreCallRequest}; use std::sync::Arc; + use crate::ocr::hooks::{OcrHookFuture, OcrHooks, OcrPreCallRequest}; + struct RewritePages; impl OcrHooks for RewritePages { fn intercepts_requests(&self) -> bool { 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 dffe0aa9b05..1a909abc2d6 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,3 +1,7 @@ +use litellm_auth::{InputSource, Sourced}; +use litellm_auth_azure::AzureAuthInputs; +use serde_json::Value; + use crate::call_arguments::CallArguments; use crate::constants::AZURE_AI_OCR_PATH; use crate::llms::base_llm::ocr::transformation::{BaseOcrConfig, OcrRequestContext}; @@ -8,9 +12,6 @@ use crate::ocr::prepare::credential_env; use crate::ocr::types::{LiteLLMOcrResponse, OcrConnection, OcrDocument, PreparedOcrRequest}; use crate::params::OpaqueParams; use crate::url_utils::ApiUrl; -use litellm_auth::{InputSource, Sourced}; -use litellm_auth_azure::AzureAuthInputs; -use serde_json::Value; const AZURE_AI_API_KEY_ENV: &str = "AZURE_AI_API_KEY"; const AZURE_AI_API_BASE_ENV: &str = "AZURE_AI_API_BASE"; 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 fc11f62833c..09dd8d49757 100644 --- a/litellm-rust/crates/core/src/llms/cohere/ocr/transformation.rs +++ b/litellm-rust/crates/core/src/llms/cohere/ocr/transformation.rs @@ -344,9 +344,10 @@ fn invalid_api_base() -> crate::ocr::Error { #[cfg(test)] mod tests { - use super::*; use serde_json::json; + use super::*; + #[tokio::test] async fn composed_body_preserves_native_document_fields_and_untyped_overrides() { let request = crate::ocr::test_support::wire_request( 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 d90bfeff2a7..ffabce84d05 100644 --- a/litellm-rust/crates/core/src/llms/mistral/ocr/transformation.rs +++ b/litellm-rust/crates/core/src/llms/mistral/ocr/transformation.rs @@ -202,10 +202,11 @@ impl MistralOCRConfig { #[cfg(test)] mod tests { - use super::*; use rstest::rstest; use serde_json::{Value, json}; + use super::*; + #[test] fn explicit_null_model_does_not_use_the_missing_model_default() { let response = serde_json::from_value(json!({"model":null})).unwrap(); diff --git a/litellm-rust/crates/core/src/llms/vertex_ai/ocr/common_utils.rs b/litellm-rust/crates/core/src/llms/vertex_ai/ocr/common_utils.rs index 6340084ad7f..08ffbc43cd5 100644 --- a/litellm-rust/crates/core/src/llms/vertex_ai/ocr/common_utils.rs +++ b/litellm-rust/crates/core/src/llms/vertex_ai/ocr/common_utils.rs @@ -1,6 +1,7 @@ -use crate::ocr::types::OcrConnection; use litellm_auth::InputSource; +use crate::ocr::types::OcrConnection; + pub(super) fn validate_destination(connection: &OcrConnection) -> Result<(), crate::ocr::Error> { if connection.api_base.is_some() && connection.api_base_source == InputSource::Request { return Err(litellm_auth::Error::RequestVertexCredentialDestination.into()); 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 7caa4656678..335d6e49dd3 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 @@ -1,8 +1,7 @@ +use litellm_auth_gcp::{self as vertex, VertexConfig}; use serde::{Deserialize, Serialize}; use serde_json::{Map, Value}; -use litellm_auth_gcp::{self as vertex, VertexConfig}; - use super::transformation::VertexAIOCRConfig; use crate::call_arguments::CallArguments; use crate::llms::base_llm::ocr::transformation::{BaseOcrConfig, OcrRequestContext}; @@ -410,17 +409,19 @@ impl VertexAIDeepSeekOCRConfig { #[cfg(test)] mod tests { + use serde_json::{Value, json}; + use super::{ DeepSeekOcrParams, DeepSeekOcrResponse, VertexAIDeepSeekOCRConfig, normalize_response, provider_model, }; - use serde_json::{Value, json}; #[test] fn unconsumed_options_remain_available_for_body_composition() { - use crate::llms::base_llm::ocr::transformation::BaseOcrConfig; use serde_json::json; + use crate::llms::base_llm::ocr::transformation::BaseOcrConfig; + let arguments = serde_json::from_value(json!({"temperature":0.5,"extension":null})).unwrap(); assert_eq!( @@ -615,9 +616,10 @@ mod tests { } } - use crate::ocr::test_support::{MockResponse, mock_server, perform_ocr, wire_request}; use litellm_auth::InputSource; + 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() } 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 f71a295e7dd..337fa76cfe2 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 @@ -215,10 +215,10 @@ mod tests { ); } + use litellm_auth::InputSource; 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() diff --git a/litellm-rust/crates/core/src/media.rs b/litellm-rust/crates/core/src/media.rs index ba26f431e57..0b5bc7f575d 100644 --- a/litellm-rust/crates/core/src/media.rs +++ b/litellm-rust/crates/core/src/media.rs @@ -279,11 +279,13 @@ impl Resolve for PublicDnsResolver { #[cfg(test)] mod tests { - use super::*; use std::collections::HashSet; + use tokio::io::{AsyncReadExt, AsyncWriteExt}; use tokio::net::TcpListener; + use super::*; + async fn serve(response: &'static [u8]) -> (Url, tokio::task::JoinHandle<()>) { let listener = TcpListener::bind("127.0.0.1:0") .await diff --git a/litellm-rust/crates/core/src/messages/common_utils.rs b/litellm-rust/crates/core/src/messages/common_utils.rs index cbaf92b4986..73e9a964749 100644 --- a/litellm-rust/crates/core/src/messages/common_utils.rs +++ b/litellm-rust/crates/core/src/messages/common_utils.rs @@ -1,12 +1,11 @@ -use super::Error; -use crate::http_utils::string_headers as shared_string_headers; -use crate::providers::anthropic::messages::transformation::ANTHROPIC_MESSAGES_CONFIG; -use crate::providers::azure_ai::messages::transformation::AZURE_ANTHROPIC_MESSAGES_CONFIG; use serde_json::{Map, Value}; +use super::Error; use super::transformation::AnthropicMessagesProviderConfig; - +use crate::http_utils::string_headers as shared_string_headers; pub(super) use crate::http_utils::{has_bearer_auth, has_header, truncate_error_body}; +use crate::providers::anthropic::messages::transformation::ANTHROPIC_MESSAGES_CONFIG; +use crate::providers::azure_ai::messages::transformation::AZURE_ANTHROPIC_MESSAGES_CONFIG; const HEADER_CONTEXT: &str = "messages"; diff --git a/litellm-rust/crates/core/src/messages/handler.rs b/litellm-rust/crates/core/src/messages/handler.rs index d7d593f2d57..8d1d4432627 100644 --- a/litellm-rust/crates/core/src/messages/handler.rs +++ b/litellm-rust/crates/core/src/messages/handler.rs @@ -1,11 +1,10 @@ use super::Error; -use crate::constants::ANTHROPIC_MESSAGES_PROVIDER; -use crate::http_utils::http_request; - use super::client::http_client; use super::common_utils::truncate_error_body; use super::prepare::prepare_provider_request; use super::types::{AnthropicMessagesResponse, MessagesRequest}; +use crate::constants::ANTHROPIC_MESSAGES_PROVIDER; +use crate::http_utils::http_request; pub(super) async fn execute_messages_provider_call( request: MessagesRequest<'_>, diff --git a/litellm-rust/crates/core/src/messages/prepare.rs b/litellm-rust/crates/core/src/messages/prepare.rs index b10e03ea9c0..0deb42a34ae 100644 --- a/litellm-rust/crates/core/src/messages/prepare.rs +++ b/litellm-rust/crates/core/src/messages/prepare.rs @@ -1,10 +1,10 @@ -use super::Error; -use crate::providers::custom_llm_provider::{CustomLlmProvider, get_custom_llm_provider}; +use serde_json::{Map, Value}; +use super::Error; use super::common_utils::{has_bearer_auth, has_header, messages_provider_config, string_headers}; use super::transformation::{AnthropicMessagesProviderConfig, MessagesAuthStrategy}; use super::types::{MessagesRequest, ProviderMessagesRequest}; -use serde_json::{Map, Value}; +use crate::providers::custom_llm_provider::{CustomLlmProvider, get_custom_llm_provider}; pub(super) fn prepare_provider_request( request: MessagesRequest<'_>, diff --git a/litellm-rust/crates/core/src/messages/tests.rs b/litellm-rust/crates/core/src/messages/tests.rs index f454effd7b5..212096fbd53 100644 --- a/litellm-rust/crates/core/src/messages/tests.rs +++ b/litellm-rust/crates/core/src/messages/tests.rs @@ -5,7 +5,6 @@ use tokio::io::{AsyncReadExt, AsyncWriteExt}; use tokio::net::{TcpListener, TcpStream}; use super::Error; - use super::common_utils::{ has_bearer_auth, has_header, messages_provider_config, string_headers, truncate_error_body, }; diff --git a/litellm-rust/crates/core/src/ocr/arguments.rs b/litellm-rust/crates/core/src/ocr/arguments.rs index 293931e8bbb..a657ef0dc8a 100644 --- a/litellm-rust/crates/core/src/ocr/arguments.rs +++ b/litellm-rust/crates/core/src/ocr/arguments.rs @@ -1,6 +1,5 @@ -use crate::call_arguments::ArgumentSpec; - use super::provider_config::{OcrConfigKind, resolve_provider_config}; +use crate::call_arguments::ArgumentSpec; const COMMON_OPTION_FIELDS: &[&str] = &["req_format", "extra_body", "max_response_bytes"]; const AZURE_AUTH_OPTION_FIELDS: &[&str] = &[ diff --git a/litellm-rust/crates/core/src/ocr/client.rs b/litellm-rust/crates/core/src/ocr/client.rs index 5881519855c..8dba37bb00b 100644 --- a/litellm-rust/crates/core/src/ocr/client.rs +++ b/litellm-rust/crates/core/src/ocr/client.rs @@ -2,13 +2,13 @@ use std::sync::OnceLock; use std::time::Duration; use bytes::{Bytes, BytesMut}; +use litellm_auth_gcp::VertexAuth; use serde::de::DeserializeOwned; use super::json::{DecodedOcrResponse, decode_response}; use super::types::{LiteLLMOcrRequest, LiteLLMOcrResponse}; use crate::constants::OCR_CONNECT_TIMEOUT_SECS; use crate::media::MediaFetcher; -use litellm_auth_gcp::VertexAuth; #[derive(Clone)] pub struct OcrClient { diff --git a/litellm-rust/crates/core/src/ocr/document.rs b/litellm-rust/crates/core/src/ocr/document.rs index fbb54f0bbd1..c3ffac701b3 100644 --- a/litellm-rust/crates/core/src/ocr/document.rs +++ b/litellm-rust/crates/core/src/ocr/document.rs @@ -1,3 +1,4 @@ +use std::collections::BTreeMap as Map; use std::io::Read; use std::path::Path; @@ -5,7 +6,6 @@ use base64::{Engine, engine::general_purpose::STANDARD}; use data_url::mime::Mime; use data_url::{DataUrl, DataUrlError, forgiving_base64::DecodeError}; use reqwest::Url; -use std::collections::BTreeMap as Map; use super::Error as OcrError; use super::Error as OcrRequestError; @@ -216,9 +216,10 @@ fn map_media_error(error: MediaError) -> OcrError { #[cfg(test)] mod tests { - use super::*; use std::collections::BTreeMap as Map; + use super::*; + fn document(source: &str) -> OcrDocument { OcrDocument::DocumentUrl { document_url: source.into(), diff --git a/litellm-rust/crates/core/src/ocr/hooks.rs b/litellm-rust/crates/core/src/ocr/hooks.rs index 8a14afb7c50..fdcf4fa05ba 100644 --- a/litellm-rust/crates/core/src/ocr/hooks.rs +++ b/litellm-rust/crates/core/src/ocr/hooks.rs @@ -2,11 +2,12 @@ use std::future::Future; use std::pin::Pin; use std::sync::Arc; +use serde::Serialize; +use serde_json::Value; + use super::types::{LiteLLMOcrRequest, LiteLLMOcrResponse, OcrDocument, ResolvedOcrRequest}; use crate::call_lifecycle::{CallLifecycleContext, CallLifecycleHooks, CallLifecycleTiming}; use crate::ocr::Error; -use serde::Serialize; -use serde_json::Value; pub type OcrHookFuture<'a, T> = Pin> + Send + 'a>>; pub type OcrLogFuture<'a> = Pin + Send + 'a>>; diff --git a/litellm-rust/crates/core/src/ocr/lifecycle.rs b/litellm-rust/crates/core/src/ocr/lifecycle.rs index dee34526001..b8b81a6b672 100644 --- a/litellm-rust/crates/core/src/ocr/lifecycle.rs +++ b/litellm-rust/crates/core/src/ocr/lifecycle.rs @@ -2,6 +2,8 @@ use std::future::Future; use std::pin::Pin; use std::sync::Arc; +use litellm_auth::Error as AuthError; +use litellm_auth::{ResolvedCredential, TokenFuture, TokenProvider, TokenProviderHandle}; use tokio::sync::{mpsc, oneshot}; use super::handler::perform_ocr_request; @@ -16,8 +18,6 @@ use crate::call_lifecycle::host::{ }; use crate::call_lifecycle::{CallLifecycleContext, CallLifecycleTiming}; use crate::ocr::Error; -use litellm_auth::Error as AuthError; -use litellm_auth::{ResolvedCredential, TokenFuture, TokenProvider, TokenProviderHandle}; pub type NativeResult = Result, Error>; diff --git a/litellm-rust/crates/core/src/ocr/prepare.rs b/litellm-rust/crates/core/src/ocr/prepare.rs index aa4ca94bf0c..91da5a9613d 100644 --- a/litellm-rust/crates/core/src/ocr/prepare.rs +++ b/litellm-rust/crates/core/src/ocr/prepare.rs @@ -165,9 +165,10 @@ pub(crate) fn prepare_request(request: ResolvedOcrRequest) -> PreparedOcrRequest #[cfg(test)] mod tests { - use crate::call_arguments::{CallArguments, compose_body, parse_options}; use serde_json::json; + use crate::call_arguments::{CallArguments, compose_body, parse_options}; + #[derive(serde::Deserialize)] struct KnownParams { pages: Option>, diff --git a/litellm-rust/crates/core/src/ocr/provider_config.rs b/litellm-rust/crates/core/src/ocr/provider_config.rs index 9fb89812664..ef9de23c913 100644 --- a/litellm-rust/crates/core/src/ocr/provider_config.rs +++ b/litellm-rust/crates/core/src/ocr/provider_config.rs @@ -1,3 +1,5 @@ +use strum::{EnumString, IntoStaticStr}; + use super::OcrClient; use super::types::{ LiteLLMOcrResponse, OcrCredentialInputs, OcrDocument, PreparedOcrRequest, @@ -13,7 +15,6 @@ 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::providers::custom_llm_provider::{CustomLlmProvider, get_custom_llm_provider}; -use strum::{EnumString, IntoStaticStr}; macro_rules! dispatch_config { ($config:expr, $method:ident($($argument:expr),* $(,)?)) => { @@ -185,10 +186,11 @@ fn is_document_intelligence_model(model: &str) -> bool { #[cfg(test)] mod tests { - use super::*; use litellm_auth::{InputSource, Sourced}; use rstest::rstest; + use super::*; + #[rstest] #[case("cohere")] #[case("mistral")] diff --git a/litellm-rust/crates/core/src/ocr/types.rs b/litellm-rust/crates/core/src/ocr/types.rs index 449ba34b593..facfd04fe8e 100644 --- a/litellm-rust/crates/core/src/ocr/types.rs +++ b/litellm-rust/crates/core/src/ocr/types.rs @@ -4,12 +4,11 @@ use std::sync::Arc; use std::time::Duration; use bytes::Bytes; +use litellm_auth::{InputSource, Sourced, TokenProviderHandle}; 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; @@ -583,9 +582,10 @@ fn ocr_object() -> String { #[cfg(test)] mod tests { - use super::*; use serde_json::json; + use super::*; + fn document() -> OcrDocument { OcrDocument::try_from( json!({"type":"document_url","document_url":"data:application/pdf;base64,YWJj"}), diff --git a/litellm-rust/crates/core/src/providers/anthropic/chat_completions/tests.rs b/litellm-rust/crates/core/src/providers/anthropic/chat_completions/tests.rs index 2cc94751fb4..81bc8f02a66 100644 --- a/litellm-rust/crates/core/src/providers/anthropic/chat_completions/tests.rs +++ b/litellm-rust/crates/core/src/providers/anthropic/chat_completions/tests.rs @@ -1,6 +1,7 @@ +use serde_json::json; + use super::*; use crate::chat_completions::Error; -use serde_json::json; fn messages(value: Value) -> Vec { serde_json::from_value(value).expect("valid messages") diff --git a/litellm-rust/crates/core/src/providers/anthropic/chat_completions/transformation.rs b/litellm-rust/crates/core/src/providers/anthropic/chat_completions/transformation.rs index ba1a1e1d350..dd0830edab7 100644 --- a/litellm-rust/crates/core/src/providers/anthropic/chat_completions/transformation.rs +++ b/litellm-rust/crates/core/src/providers/anthropic/chat_completions/transformation.rs @@ -2,6 +2,7 @@ use serde_json::{Map, Value, json}; use crate::chat_completions::Error; use crate::chat_completions::conversation::{Conversation, build_conversation}; +use crate::chat_completions::response_utils::{finish_reason_for, unix_now, usage_from_parts}; use crate::chat_completions::transformation::{ ChatCompletionsAuth, ChatCompletionsProviderConfig, Unsupported, unsupported_message, unsupported_param, @@ -15,8 +16,6 @@ use crate::providers::anthropic::messages::transformation::{ complete_anthropic_url, resolve_anthropic_api_key, }; -use crate::chat_completions::response_utils::{finish_reason_for, unix_now, usage_from_parts}; - /// Anthropic parameter names, post `map_openai_params`, that the Rust path can /// place verbatim in the Messages body. /// diff --git a/litellm-rust/crates/core/src/providers/azure_ai/messages/transformation.rs b/litellm-rust/crates/core/src/providers/azure_ai/messages/transformation.rs index 182aea84ab2..1929f86a1d6 100644 --- a/litellm-rust/crates/core/src/providers/azure_ai/messages/transformation.rs +++ b/litellm-rust/crates/core/src/providers/azure_ai/messages/transformation.rs @@ -1,3 +1,5 @@ +use serde_json::{Map, Value}; + use crate::messages::Error; use crate::messages::transformation::{AnthropicMessagesProviderConfig, MessagesAuthStrategy}; use crate::messages::types::{ @@ -7,7 +9,6 @@ use crate::messages::types::{ use crate::providers::anthropic::messages::transformation::{ ANTHROPIC_MESSAGES_CONFIG, AnthropicMessagesConfig, non_empty, }; -use serde_json::{Map, Value}; const AZURE_API_KEY_ENV: &str = "AZURE_API_KEY"; const AZURE_API_BASE_ENV: &str = "AZURE_API_BASE"; @@ -191,9 +192,10 @@ impl AnthropicMessagesProviderConfig for AzureAnthropicMessagesConfig { #[cfg(test)] mod tests { - use super::*; use serde_json::json; + use super::*; + fn request_from(value: serde_json::Value) -> AnthropicMessagesRequest { serde_json::from_value(value).expect("valid request") } diff --git a/litellm-rust/crates/core/src/providers/bedrock/audio_transcription.rs b/litellm-rust/crates/core/src/providers/bedrock/audio_transcription.rs index a418e860b92..12ea91672e8 100644 --- a/litellm-rust/crates/core/src/providers/bedrock/audio_transcription.rs +++ b/litellm-rust/crates/core/src/providers/bedrock/audio_transcription.rs @@ -1,5 +1,7 @@ use serde_json::{Map, Value, json}; +pub use super::aws_base::{aws_auth_config, bedrock_model_id_and_region, resolve_bedrock_region}; +use super::constants::{BEDROCK_RUNTIME_ENDPOINT_TEMPLATE, BEDROCK_SERVICE}; use crate::audio_transcription::Error; use crate::audio_transcription::transformation::{ AudioTranscriptionAuth, AudioTranscriptionProviderConfig, @@ -9,9 +11,6 @@ use crate::audio_transcription::types::{ }; use crate::http_utils::json_type_name; -pub use super::aws_base::{aws_auth_config, bedrock_model_id_and_region, resolve_bedrock_region}; -use super::constants::{BEDROCK_RUNTIME_ENDPOINT_TEMPLATE, BEDROCK_SERVICE}; - const SUPPORTED_PARAMS: &[&str] = &["language", "prompt", "temperature", "response_format"]; pub static BEDROCK_AUDIO_TRANSCRIPTION_CONFIG: BedrockAudioTranscriptionConfig = diff --git a/litellm-rust/crates/core/src/providers/bedrock/chat_completions/tests.rs b/litellm-rust/crates/core/src/providers/bedrock/chat_completions/tests.rs index 74716a2200b..08ebac9dea1 100644 --- a/litellm-rust/crates/core/src/providers/bedrock/chat_completions/tests.rs +++ b/litellm-rust/crates/core/src/providers/bedrock/chat_completions/tests.rs @@ -1,6 +1,7 @@ +use serde_json::json; + use super::*; use crate::chat_completions::Error; -use serde_json::json; fn messages(value: Value) -> Vec { serde_json::from_value(value).expect("valid messages") diff --git a/litellm-rust/crates/core/src/providers/bedrock/chat_completions/transformation.rs b/litellm-rust/crates/core/src/providers/bedrock/chat_completions/transformation.rs index 19efaf833bd..53d3842955c 100644 --- a/litellm-rust/crates/core/src/providers/bedrock/chat_completions/transformation.rs +++ b/litellm-rust/crates/core/src/providers/bedrock/chat_completions/transformation.rs @@ -1,5 +1,7 @@ use serde_json::{Map, Value, json}; +use super::super::aws_base::{bedrock_model_id_and_region, resolve_bedrock_region}; +use super::super::constants::{AWS_BEARER_TOKEN_BEDROCK, BEDROCK_RUNTIME_ENDPOINT_TEMPLATE}; use crate::chat_completions::Error; use crate::chat_completions::conversation::{Conversation, TurnRole, build_conversation}; use crate::chat_completions::response_utils::{finish_reason_for, unix_now, usage_from_parts}; @@ -13,9 +15,6 @@ use crate::chat_completions::types::{ ProviderChatResponseData, }; -use super::super::aws_base::{bedrock_model_id_and_region, resolve_bedrock_region}; -use super::super::constants::{AWS_BEARER_TOKEN_BEDROCK, BEDROCK_RUNTIME_ENDPOINT_TEMPLATE}; - /// Converse parameter names, post `map_openai_params`, that the Rust path can /// place verbatim in `inferenceConfig`. /// diff --git a/litellm-rust/crates/core/src/serde_compat.rs b/litellm-rust/crates/core/src/serde_compat.rs index 5a2d0688c33..3ec869b40e2 100644 --- a/litellm-rust/crates/core/src/serde_compat.rs +++ b/litellm-rust/crates/core/src/serde_compat.rs @@ -66,11 +66,12 @@ fn integral_float(value: f64) -> Option { #[cfg(test)] mod tests { - use super::*; use serde::Serialize; use serde_json::json; use serde_with::serde_as; + use super::*; + #[serde_as] #[derive(Debug, Deserialize, Serialize, PartialEq)] struct Numbers { diff --git a/litellm-rust/crates/core/tests/azure_document_intelligence_ocr.rs b/litellm-rust/crates/core/tests/azure_document_intelligence_ocr.rs index 5682e8ad5be..1da340b57d4 100644 --- a/litellm-rust/crates/core/tests/azure_document_intelligence_ocr.rs +++ b/litellm-rust/crates/core/tests/azure_document_intelligence_ocr.rs @@ -1,6 +1,7 @@ -use serde_json::{Value, json}; use std::sync::{Arc, Mutex}; +use serde_json::{Value, json}; + use super::test_support::{MockResponse, mock_server, perform_ocr, wire_request}; use super::wire::{OcrWireRequest, decode_request}; @@ -421,9 +422,10 @@ async fn model_id_is_encoded_and_dot_segments_are_rejected() { #[tokio::test] async fn pre_call_guardrail_receives_caller_pages_before_mapping() { - use crate::ocr::hooks::{OcrHookFuture, OcrHooks, OcrPreCallRequest}; use std::sync::Arc; + use crate::ocr::hooks::{OcrHookFuture, OcrHooks, OcrPreCallRequest}; + struct RewritePages; impl OcrHooks for RewritePages { fn intercepts_requests(&self) -> bool { diff --git a/litellm-rust/crates/core/tests/ocr.rs b/litellm-rust/crates/core/tests/ocr.rs index 302ed91701e..78dfd5a2c9f 100644 --- a/litellm-rust/crates/core/tests/ocr.rs +++ b/litellm-rust/crates/core/tests/ocr.rs @@ -906,11 +906,12 @@ impl litellm_auth::TokenProvider for PendingToken { #[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; + use crate::call_lifecycle::host::HostFailure; + for interrupt_acknowledgement in [false, true] { let entered = Arc::new(tokio::sync::Notify::new()); let dropped = Arc::new(AtomicBool::new(false)); diff --git a/litellm-rust/crates/core/tests/vertex_ai_deepseek_ocr.rs b/litellm-rust/crates/core/tests/vertex_ai_deepseek_ocr.rs index be0898e1135..6be30f784c4 100644 --- a/litellm-rust/crates/core/tests/vertex_ai_deepseek_ocr.rs +++ b/litellm-rust/crates/core/tests/vertex_ai_deepseek_ocr.rs @@ -1,7 +1,7 @@ +use litellm_auth::InputSource; 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() diff --git a/litellm-rust/crates/core/tests/vertex_ai_ocr.rs b/litellm-rust/crates/core/tests/vertex_ai_ocr.rs index 27e4802b00d..ebee4046e23 100644 --- a/litellm-rust/crates/core/tests/vertex_ai_ocr.rs +++ b/litellm-rust/crates/core/tests/vertex_ai_ocr.rs @@ -1,7 +1,7 @@ +use litellm_auth::InputSource; 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() From 85e70ea3746c7981bb379f1344dc2eed4286f7b8 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Wed, 16 Sep 2026 20:49:31 -0700 Subject: [PATCH 03/15] fix(ocr): await blocking preparation on cancellation --- litellm-rust/crates/core/src/ocr/lifecycle.rs | 56 +++++++++++-- litellm-rust/crates/core/tests/ocr.rs | 80 +++++++++++++++++++ 2 files changed, 129 insertions(+), 7 deletions(-) diff --git a/litellm-rust/crates/core/src/ocr/lifecycle.rs b/litellm-rust/crates/core/src/ocr/lifecycle.rs index b8b81a6b672..f2e5479b361 100644 --- a/litellm-rust/crates/core/src/ocr/lifecycle.rs +++ b/litellm-rust/crates/core/src/ocr/lifecycle.rs @@ -1,10 +1,11 @@ use std::future::Future; use std::pin::Pin; use std::sync::Arc; +use std::sync::atomic::{AtomicBool, Ordering}; use litellm_auth::Error as AuthError; use litellm_auth::{ResolvedCredential, TokenFuture, TokenProvider, TokenProviderHandle}; -use tokio::sync::{mpsc, oneshot}; +use tokio::sync::{Notify, mpsc, oneshot}; use super::handler::perform_ocr_request; use super::hooks::{ @@ -321,6 +322,7 @@ struct OcrExecution { operations_rx: mpsc::UnboundedReceiver, pending_result: Option>, execution: Option>>, + blocking_preparation: Arc, completed: bool, azure_ad_token_provider: bool, terminal: Arc>>, @@ -336,6 +338,7 @@ impl OcrExecution { operations_rx, pending_result: None, execution: None, + blocking_preparation: Arc::new(BlockingPreparation::default()), completed: false, azure_ad_token_provider: false, terminal: Arc::default(), @@ -406,8 +409,9 @@ impl OcrExecution { terminal: self.terminal.clone(), }); request.hooks = hooks.clone(); + let blocking_preparation = self.blocking_preparation.clone(); self.execution = Some(tokio::spawn(async move { - let request = prepare_request_document(request, &hooks).await?; + let request = prepare_request_document(request, &hooks, blocking_preparation).await?; perform_ocr_request(&client, request).await })); } @@ -424,13 +428,47 @@ impl OcrExecution { if let Some(execution) = self.execution.as_mut() { let _ = execution.await; } + self.blocking_preparation.wait().await; self.execution = None; } } +#[derive(Default)] +struct BlockingPreparation { + running: AtomicBool, + finished: Notify, +} + +impl BlockingPreparation { + fn start(self: &Arc) -> BlockingPreparationGuard { + self.running.store(true, Ordering::Release); + BlockingPreparationGuard(self.clone()) + } + + async fn wait(&self) { + loop { + let finished = self.finished.notified(); + if !self.running.load(Ordering::Acquire) { + return; + } + finished.await; + } + } +} + +struct BlockingPreparationGuard(Arc); + +impl Drop for BlockingPreparationGuard { + fn drop(&mut self) { + self.0.running.store(false, Ordering::Release); + self.0.finished.notify_waiters(); + } +} + async fn prepare_request_document( request: LiteLLMOcrRequest, hooks: &ProtocolHooks, + blocking_preparation: Arc, ) -> Result { let request = match &request.document { OcrDocumentInput::HostReader { mime_type } => { @@ -454,11 +492,15 @@ async fn prepare_request_document( if let OcrDocumentInput::Document(_) = &request.document { return request.map_document(super::document::prepare_document); } - tokio::task::spawn_blocking(move || request.map_document(super::document::prepare_document)) - .await - .map_err(|error| { - Error::InvalidRequest(format!("OCR document preparation task failed: {error}")) - })? + let guard = blocking_preparation.start(); + tokio::task::spawn_blocking(move || { + let _guard = guard; + request.map_document(super::document::prepare_document) + }) + .await + .map_err(|error| { + Error::InvalidRequest(format!("OCR document preparation task failed: {error}")) + })? } impl Drop for OcrExecution { diff --git a/litellm-rust/crates/core/tests/ocr.rs b/litellm-rust/crates/core/tests/ocr.rs index 78dfd5a2c9f..480774d1ad1 100644 --- a/litellm-rust/crates/core/tests/ocr.rs +++ b/litellm-rust/crates/core/tests/ocr.rs @@ -744,6 +744,86 @@ async fn cancellation_at_provider_hook_prevents_execution_and_further_resumption ); } +#[cfg(unix)] +#[tokio::test] +async fn cancellation_acknowledges_blocking_preparation_completion() { + use std::future::Future; + use std::io::Write; + use std::task::Poll; + + use crate::call_lifecycle::host::HostFailure; + + let path = std::env::temp_dir().join(format!("litellm-ocr-{}.fifo", rand::random::())); + assert!( + std::process::Command::new("mkfifo") + .arg(&path) + .status() + .unwrap() + .success() + ); + let request = wire_request("mistral/model", "http://127.0.0.1:1", json!({})).with_document( + super::OcrDocumentInput::Path { + path: path.clone(), + mime_type: Some("application/pdf".into()), + }, + ); + 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; + loop { + match call.resume(result.take()).await.unwrap() { + OcrCallStep::Host(OcrHostOperation::ProjectRequest) => break, + OcrCallStep::Host(operation) => result = Some(NoopOcrHost.invoke(operation).await), + OcrCallStep::Complete(_) => panic!("provider executed before request projection"), + } + } + let mut preparation = Box::pin(call.resume(Some(OcrHostResult::Request(Ok(( + Box::new(request.take().unwrap()), + false, + )))))); + std::future::poll_fn(|cx| { + assert!(preparation.as_mut().poll(cx).is_pending()); + Poll::Ready(()) + }) + .await; + drop(preparation); + + let (entered_tx, entered_rx) = tokio::sync::oneshot::channel(); + let (release_tx, release_rx) = std::sync::mpsc::channel(); + let writer_path = path.clone(); + let writer = tokio::task::spawn_blocking(move || { + let mut fifo = std::fs::File::options() + .write(true) + .open(writer_path) + .unwrap(); + entered_tx.send(()).unwrap(); + release_rx.recv().unwrap(); + fifo.write_all(b"document").unwrap(); + }); + tokio::time::timeout(std::time::Duration::from_secs(2), entered_rx) + .await + .unwrap() + .unwrap(); + + let selected = crate::ocr::Error::InvalidRequest("cancelled".into()); + 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; + release_tx.send(()).unwrap(); + assert!( + matches!(acknowledgement.await, Err(crate::ocr::Error::InvalidRequest(message)) if message == "cancelled") + ); + writer.await.unwrap(); + std::fs::remove_file(path).unwrap(); +} + #[tokio::test] async fn missing_host_result_preserves_pending_operation() { use crate::call_lifecycle::host::HostPhase; From 0f636c5db1b4a6c55bf2c092d34dde12528950a3 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Thu, 17 Sep 2026 06:51:52 -0700 Subject: [PATCH 04/15] refactor(core): use string for DeepSeek model --- .../vertex_ai/ocr/deepseek_transformation.rs | 30 +-- litellm-rust/crates/core/src/providers/mod.rs | 1 - .../crates/core/src/providers/model.rs | 219 ------------------ 3 files changed, 11 insertions(+), 239 deletions(-) delete mode 100644 litellm-rust/crates/core/src/providers/model.rs 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 335d6e49dd3..43bee24b860 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 @@ -12,11 +12,10 @@ use crate::ocr::types::{ PreparedOcrRequest, }; use crate::params::OpaqueParams; -use crate::providers::model::{ModelNamespace, ProviderModel, RoutedModel}; use crate::url_utils::ApiUrl; const DEFAULT_API_BASE: &str = "https://aiplatform.googleapis.com"; -const MODEL_NAMESPACE: &str = "deepseek-ai"; +const MODEL_PREFIX: &str = "deepseek-ai/"; const DEFAULT_LOCATION: &str = "us-central1"; const DEEPSEEK_OCR_PARAMS: &[&str] = &["stream", "temperature", "max_tokens", "top_p", "n", "stop"]; @@ -24,7 +23,7 @@ pub(crate) type DeepSeekOcrParams = OpaqueParams; #[derive(Clone, Debug, Serialize, Deserialize)] pub(crate) struct DeepSeekOcrRequest { - pub model: ProviderModel, + pub model: String, pub messages: Vec, #[serde(flatten)] pub params: OpaqueParams, @@ -87,13 +86,6 @@ struct DeepSeekPage { dimensions: Option, } -#[derive(Clone, Debug)] -pub(crate) struct DeepSeekAi; - -impl ModelNamespace for DeepSeekAi { - const NAME: &'static str = MODEL_NAMESPACE; -} - #[derive(Clone, Debug)] pub(crate) struct VertexAIDeepSeekOCRConfig; @@ -367,12 +359,14 @@ fn response_field(field: &str) -> crate::ocr::Error { } } -pub(crate) fn provider_model(model: &str) -> Result, crate::ocr::Error> { - RoutedModel::new(model) - .and_then(RoutedModel::into_provider::) - .map_err(|_| crate::ocr::Error::RequestField { +pub(crate) fn provider_model(model: &str) -> Result { + let local_model = model.trim_start_matches(MODEL_PREFIX); + if local_model.is_empty() { + return Err(crate::ocr::Error::RequestField { path: "model".into(), - }) + }); + } + Ok(format!("{MODEL_PREFIX}{local_model}")) } impl VertexAIDeepSeekOCRConfig { @@ -443,13 +437,11 @@ mod tests { #[test] fn config_owns_model_namespace_and_endpoint() { assert_eq!( - provider_model("deepseek-ocr-maas").unwrap().as_str(), + provider_model("deepseek-ocr-maas").unwrap(), "deepseek-ai/deepseek-ocr-maas" ); assert_eq!( - provider_model("deepseek-ai/deepseek-ocr-maas") - .unwrap() - .as_str(), + provider_model("deepseek-ai/deepseek-ocr-maas").unwrap(), "deepseek-ai/deepseek-ocr-maas" ); assert_eq!( diff --git a/litellm-rust/crates/core/src/providers/mod.rs b/litellm-rust/crates/core/src/providers/mod.rs index 79eb3404ece..70ca4386fff 100644 --- a/litellm-rust/crates/core/src/providers/mod.rs +++ b/litellm-rust/crates/core/src/providers/mod.rs @@ -2,5 +2,4 @@ pub mod anthropic; pub mod azure_ai; pub mod bedrock; pub mod custom_llm_provider; -pub(crate) mod model; pub mod openai; diff --git a/litellm-rust/crates/core/src/providers/model.rs b/litellm-rust/crates/core/src/providers/model.rs deleted file mode 100644 index fcedc4b023a..00000000000 --- a/litellm-rust/crates/core/src/providers/model.rs +++ /dev/null @@ -1,219 +0,0 @@ -use std::marker::PhantomData; - -use serde::{Deserialize, Deserializer, Serialize, Serializer}; - -#[derive(Clone, Copy, Debug, Eq, PartialEq, thiserror::Error)] -pub(crate) enum ModelNameError { - #[error("model name cannot be empty")] - EmptyModel, - #[error("model namespace must be one non-empty path segment: {0}")] - InvalidNamespace(&'static str), -} - -pub(crate) trait ModelNamespace { - const NAME: &'static str; -} - -#[derive(Clone, Copy, Debug, Eq, PartialEq)] -pub(crate) struct RoutedModel<'a>(&'a str); - -impl<'a> RoutedModel<'a> { - pub(crate) fn new(value: &'a str) -> Result { - if value.is_empty() { - return Err(ModelNameError::EmptyModel); - } - Ok(Self(value)) - } - - pub(crate) fn into_provider( - self, - ) -> Result, ModelNameError> { - let namespace = N::NAME; - if namespace.is_empty() || namespace.contains('/') { - return Err(ModelNameError::InvalidNamespace(namespace)); - } - let prefix = format!("{namespace}/"); - let local_model = self.0.trim_start_matches(prefix.as_str()); - if local_model.is_empty() { - return Err(ModelNameError::EmptyModel); - } - Ok(ProviderModel { - value: format!("{prefix}{local_model}"), - namespace: PhantomData, - }) - } -} - -#[derive(Clone, Debug, Eq, PartialEq)] -pub(crate) struct ProviderModel { - value: String, - namespace: PhantomData, -} - -impl ProviderModel { - #[cfg(test)] - pub(crate) fn as_str(&self) -> &str { - &self.value - } -} - -impl Serialize for ProviderModel { - fn serialize(&self, serializer: S) -> Result - where - S: Serializer, - { - self.value.serialize(serializer) - } -} - -impl<'de, N: ModelNamespace> Deserialize<'de> for ProviderModel { - fn deserialize(deserializer: D) -> Result - where - D: Deserializer<'de>, - { - let value = String::deserialize(deserializer)?; - RoutedModel::new(&value) - .and_then(RoutedModel::into_provider::) - .map_err(::custom) - } -} - -#[cfg(test)] -mod tests { - use serde_json::json; - - use super::*; - - #[derive(Clone, Debug, Eq, PartialEq)] - struct DeepSeekAi; - - impl ModelNamespace for DeepSeekAi { - const NAME: &'static str = "deepseek-ai"; - } - - #[derive(Clone, Debug, Eq, PartialEq)] - struct FalAi; - - impl ModelNamespace for FalAi { - const NAME: &'static str = "fal-ai"; - } - - #[test] - fn qualifies_a_bare_model() { - let model = RoutedModel::new("deepseek-ocr-maas") - .and_then(RoutedModel::into_provider::) - .unwrap(); - - assert_eq!(model.as_str(), "deepseek-ai/deepseek-ocr-maas"); - } - - #[test] - fn preserves_an_already_qualified_model() { - let model = RoutedModel::new("deepseek-ai/deepseek-ocr-maas") - .and_then(RoutedModel::into_provider::) - .unwrap(); - - assert_eq!(model.as_str(), "deepseek-ai/deepseek-ocr-maas"); - } - - #[test] - fn collapses_repeated_owned_namespaces() { - let model = RoutedModel::new("deepseek-ai/deepseek-ai/deepseek-ai/deepseek-ocr-maas") - .and_then(RoutedModel::into_provider::) - .unwrap(); - - assert_eq!(model.as_str(), "deepseek-ai/deepseek-ocr-maas"); - } - - #[test] - fn matches_the_namespace_as_a_complete_segment() { - let model = RoutedModel::new("deepseek-ai-v2/model") - .and_then(RoutedModel::into_provider::) - .unwrap(); - - assert_eq!(model.as_str(), "deepseek-ai/deepseek-ai-v2/model"); - } - - #[test] - fn preserves_nested_provider_model_paths() { - let model = RoutedModel::new("publishers/vendor/models/model-v1") - .and_then(RoutedModel::into_provider::) - .unwrap(); - - assert_eq!( - model.as_str(), - "deepseek-ai/publishers/vendor/models/model-v1" - ); - } - - #[test] - fn namespace_markers_select_different_wire_names() { - let routed = RoutedModel::new("model-v1").unwrap(); - let deepseek = routed.into_provider::().unwrap(); - let fal = routed.into_provider::().unwrap(); - - assert_eq!(deepseek.as_str(), "deepseek-ai/model-v1"); - assert_eq!(fal.as_str(), "fal-ai/model-v1"); - } - - #[test] - fn rejects_empty_routed_models() { - assert_eq!(RoutedModel::new(""), Err(ModelNameError::EmptyModel)); - } - - #[test] - fn rejects_a_namespace_without_a_model() { - let result = - RoutedModel::new("deepseek-ai/").and_then(RoutedModel::into_provider::); - - assert_eq!(result, Err(ModelNameError::EmptyModel)); - } - - #[test] - fn rejects_invalid_namespace_markers() { - struct Empty; - impl ModelNamespace for Empty { - const NAME: &'static str = ""; - } - struct MultipleSegments; - impl ModelNamespace for MultipleSegments { - const NAME: &'static str = "one/two"; - } - - assert!(matches!( - RoutedModel::new("model").and_then(RoutedModel::into_provider::), - Err(ModelNameError::InvalidNamespace("")) - )); - assert!(matches!( - RoutedModel::new("model").and_then(RoutedModel::into_provider::), - Err(ModelNameError::InvalidNamespace("one/two")) - )); - } - - #[test] - fn provider_models_serialize_as_plain_strings() { - let model = RoutedModel::new("deepseek-ocr-maas") - .and_then(RoutedModel::into_provider::) - .unwrap(); - - assert_eq!( - serde_json::to_value(model).unwrap(), - json!("deepseek-ai/deepseek-ocr-maas") - ); - } - - #[test] - fn deserialization_reestablishes_the_namespace_invariant() { - let model: ProviderModel = - serde_json::from_value(json!("deepseek-ai/deepseek-ai/model-v1")).unwrap(); - - assert_eq!(model.as_str(), "deepseek-ai/model-v1"); - } - - #[test] - fn deserialization_rejects_missing_model_names() { - let result = serde_json::from_value::>(json!("deepseek-ai/")); - - assert!(result.is_err()); - } -} From 3ad91fc27e1769e1a69c40f01e0d07d0d3a590e0 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Thu, 17 Sep 2026 06:53:19 -0700 Subject: [PATCH 05/15] fix(ocr): run hooks on completed Azure poll --- .../document_intelligence/transformation.rs | 26 ++++++++++++------- 1 file changed, 16 insertions(+), 10 deletions(-) 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 e20ec29132d..1016ed02783 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 @@ -343,7 +343,7 @@ async fn read_operation_response( let bytes = crate::ocr::client::read_response_bytes(response, connection.max_response_bytes).await?; crate::ocr::handler::post_call(hooks, &bytes).await?; - poll_operation(http_client, operation, headers, connection, native).await + poll_operation(http_client, operation, headers, connection, native, hooks).await } async fn poll_operation( @@ -352,6 +352,7 @@ async fn poll_operation( headers: &[(String, String)], connection: &OcrConnection, native: bool, + hooks: &Arc, ) -> Result, crate::ocr::Error> { let deadline = Instant::now() .checked_add(connection.poll_timeout) @@ -392,7 +393,10 @@ async fn poll_operation( .await .map_err(|_| crate::ocr::Error::PollTimeout)??; match &decoded.data.status { - Some(OperationStatus::Succeeded) => return Ok(decoded), + Some(OperationStatus::Succeeded) => { + crate::ocr::handler::post_call(hooks, decoded.text.as_bytes()).await?; + return Ok(decoded); + } Some(OperationStatus::Running | OperationStatus::NotStarted) => { tokio::time::timeout_at(deadline, tokio::time::sleep(Duration::from_secs(retry))) .await @@ -985,7 +989,7 @@ mod tests { struct SubmissionBoundary { request_count: Arc>>, - post_calls: Arc>>, + post_calls: Arc>>, } impl crate::ocr::hooks::OcrHooks for SubmissionBoundary { @@ -994,18 +998,17 @@ mod tests { 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()); + self.post_calls.lock().unwrap().push(( + self.request_count.lock().unwrap().len(), + request.original_response.clone(), + )); Ok(request) }) } } #[tokio::test] - async fn accepted_response_runs_post_call_once_before_polling() { + async fn accepted_response_runs_post_call_for_submission_and_completed_poll() { let (base, seen, server) = mock_server(vec![ MockResponse { status: 202, @@ -1029,7 +1032,10 @@ mod tests { assert_eq!(seen.lock().unwrap().len(), 2); assert_eq!( *post_calls.lock().unwrap(), - [json!(r#"{"submitted":true}"#)] + [ + (1, json!(r#"{"submitted":true}"#)), + (2, json!(r#"{"status":"succeeded"}"#)), + ] ); } From 0e5f41bc93f55c9b9a7dafca0184a7b5a154ceca Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Thu, 17 Sep 2026 06:54:18 -0700 Subject: [PATCH 06/15] refactor(rust): drop unused OpaqueParams body-composition helpers --- litellm-rust/crates/core/src/params.rs | 113 +------------------------ 1 file changed, 1 insertion(+), 112 deletions(-) diff --git a/litellm-rust/crates/core/src/params.rs b/litellm-rust/crates/core/src/params.rs index cea410db816..bdeb178c940 100644 --- a/litellm-rust/crates/core/src/params.rs +++ b/litellm-rust/crates/core/src/params.rs @@ -66,60 +66,6 @@ pub fn is_control_param(name: &str) -> bool { ) } -impl OpaqueParams { - pub fn into_inner(self) -> Map { - self.0 - } - - pub fn without(&self, names: &[&str]) -> Self { - self.iter() - .filter(|(name, _)| !names.contains(&name.as_str())) - .map(|(name, value)| (name.clone(), value.clone())) - .collect() - } - - pub fn provider_params(&self) -> Self { - self.iter() - .filter(|(name, _)| !is_control_param(name)) - .map(|(name, value)| (name.clone(), value.clone())) - .collect() - } - - pub fn into_provider_body(self) -> Result, Error> { - let mut fields = self.0; - let overrides = match fields.remove("extra_body") { - None | Some(Value::Null) => Map::new(), - Some(Value::Object(fields)) => fields, - Some(_) => { - return Err(Error::ExtraBody); - } - }; - Ok(fields - .into_iter() - .chain(overrides) - .filter(|(name, _)| name != "extra_body" && !is_control_param(name)) - .collect()) - } -} - -#[cfg(test)] -fn merge_extra_params(body: &B, extra_params: OpaqueParams) -> Result { - let Value::Object(fields) = serde_json::to_value(body).map_err(|_| Error::Body)? else { - return Err(Error::Body); - }; - Ok(Value::Object( - fields - .into_iter() - .chain( - extra_params - .into_provider_body()? - .into_iter() - .filter(|(name, _)| name != "model"), - ) - .collect(), - )) -} - impl Deref for OpaqueParams { type Target = Map; @@ -165,64 +111,7 @@ impl IntoIterator for OpaqueParams { mod tests { use serde_json::json; - use super::*; - - #[test] - fn extras_merge_shallowly_and_preserve_values_without_leaking_controls() { - let extras: OpaqueParams = serde_json::from_value(json!({ - "future": {"nested": [false, 0, null]}, - "explicit_null": null, - "azure_ad_token": "secret", - "req_format": "native", - "extra_body": { - "future": {"replacement": true}, - "temperature": 0.5, - "model": "override", - "aws_secret_access_key": "secret" - } - })) - .unwrap(); - let body = - merge_extra_params(&json!({"model":"resolved", "temperature":0.1}), extras).unwrap(); - assert_eq!( - body, - json!({ - "model":"resolved", "temperature":0.5, - "future":{"replacement":true}, "explicit_null":null - }) - ); - } - - #[test] - fn invalid_extra_body_is_rejected_and_null_is_empty() { - for value in [json!(false), json!([]), json!("value"), json!(1)] { - let params: OpaqueParams = serde_json::from_value(json!({"extra_body":value})).unwrap(); - assert!(params.into_provider_body().is_err()); - } - let params: OpaqueParams = - serde_json::from_value(json!({"extra_body":null,"future":null})).unwrap(); - assert_eq!( - Value::Object(params.into_provider_body().unwrap()), - json!({"future":null}) - ); - } - - #[test] - fn provider_params_preserve_opaque_values() { - let params: OpaqueParams = serde_json::from_value(json!({ - "object": {"future": [1, null]}, - "null": null, - "azure_ad_token": "secret" - })) - .unwrap(); - - let retained = params.provider_params(); - - assert_eq!( - serde_json::to_value(retained).unwrap(), - json!({"object": {"future": [1, null]}, "null": null}) - ); - } + use super::OpaqueParams; #[test] fn outer_value_must_be_an_object() { From ab1f966a17939299324cbdb38178188f562c8880 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Thu, 17 Sep 2026 07:10:49 -0700 Subject: [PATCH 07/15] test coverage --- .../src/llms/cohere/ocr/transformation.rs | 91 ++++++++++----- .../src/llms/mistral/ocr/transformation.rs | 60 +++++----- .../crates/core/src/ocr/provider_config.rs | 7 ++ .../tests/azure_document_intelligence_ocr.rs | 105 +++++++++++++----- litellm-rust/crates/core/tests/reducto_ocr.rs | 72 ++++++++++-- 5 files changed, 247 insertions(+), 88 deletions(-) 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 09dd8d49757..996d9e462ab 100644 --- a/litellm-rust/crates/core/src/llms/cohere/ocr/transformation.rs +++ b/litellm-rust/crates/core/src/llms/cohere/ocr/transformation.rs @@ -344,6 +344,7 @@ fn invalid_api_base() -> crate::ocr::Error { #[cfg(test)] mod tests { + use rstest::rstest; use serde_json::json; use super::*; @@ -471,10 +472,13 @@ mod tests { )); } - #[test] - fn provider_options_exclude_response_controls_and_extensions() { + #[rstest] + fn provider_options_exclude_response_controls_and_extensions( + #[values("markdown", "blocks")] output_format: &str, + #[values("https://example.com/a.png", "data:image/png;base64,YWJj")] source: &str, + ) { let arguments = serde_json::from_value( - json!({"output_format":"blocks","req_format":"native","unknown":true}), + json!({"output_format":output_format,"req_format":"native","unknown":true}), ) .unwrap(); let params = CohereParseConfig @@ -482,10 +486,10 @@ mod tests { .unwrap(); assert_eq!( serde_json::to_value(¶ms).unwrap(), - json!({"output_format":"blocks"}) + json!({"output_format":output_format}) ); let document = serde_json::from_value( - json!({"type":"image_url","image_url":"https://example.com/a.png","ignored":"field"}), + json!({"type":"image_url","image_url":source,"ignored":"field"}), ) .unwrap(); let body = CohereParseConfig @@ -494,7 +498,7 @@ mod tests { assert_eq!( serde_json::to_value(body).unwrap(), json!({ - "model":"parse", "document":{"type":"image_url","image_url":"https://example.com/a.png"}, "output_format":"blocks" + "model":"parse", "document":{"type":"image_url","image_url":source}, "output_format":output_format }) ); } @@ -526,9 +530,9 @@ mod tests { assert!(body.get("req_format").is_none()); } - #[test] + #[rstest] fn response_normalizes_markdown_images_blocks_and_billed_pages() { - let response = serde_json::from_value(json!({ + let payload = json!({ "pages": [ { "type":"markdown", @@ -558,17 +562,22 @@ mod tests { {"type":"blocks","blocks":[{"type":"text","text":{"content":"total"}}]} ], "meta":{"api_version":{"version":"2"},"billed_units":{"pages":3}} - })) - .unwrap(); + }); + let response = serde_json::from_value(payload.clone()).unwrap(); let normalized = normalize_response("parse-v5.0", response).unwrap(); assert_eq!(normalized.pages[0].index, 4); assert_eq!(normalized.pages[0].markdown, "receipt"); let image = &normalized.pages[0].images.as_ref().unwrap()[0]; - assert_eq!(image.bbox.as_ref().unwrap()["top_left_x"], 1); + let original_image = &payload["pages"][0]["markdown"]["images"][0]; assert_eq!( - image.extra_fields["bounding_box_normalized"]["bottom_right_x"], - 0.15 + serde_json::to_value(&image.bbox).unwrap(), + original_image["bounding_box"] ); + assert_eq!( + image.extra_fields["bounding_box_normalized"], + original_image["bounding_box_normalized"] + ); + assert_eq!(image.extra_fields["id"], original_image["id"]); assert_eq!(image.extra_fields["description"], "scan"); assert_eq!(image.extra_fields["category"], "logo"); assert_eq!(image.extra_fields["provider_extension"], "preserved"); @@ -609,9 +618,15 @@ mod tests { assert!(normalized.pages[0].images.is_none()); } - #[test] - fn response_types_documented_block_variants() { - let response = serde_json::from_value(json!({ + #[rstest] + fn response_types_documented_block_variants( + #[values( + crate::ocr::types::OcrResponseFormat::Litellm, + crate::ocr::types::OcrResponseFormat::Native + )] + response_format: crate::ocr::types::OcrResponseFormat, + ) { + let payload = json!({ "pages": [{ "type": "blocks", "index": 0, @@ -654,21 +669,45 @@ mod tests { "bottom_right_x": 0.7, "bottom_right_y": 0.8 }, - "title": "Totals" + "title": "Totals", + "description": "Invoice totals" } } ] }] - })) - .unwrap(); - let normalized = normalize_response("parse-v5.0", response).unwrap(); - let blocks = normalized.pages[0].extra_fields["blocks"] - .as_array() + }); + let normalized = CohereParseConfig + .transform_ocr_response( + "parse-v5.0", + &serde_json::to_vec(&payload).unwrap(), + response_format, + ) .unwrap(); - assert_eq!(blocks[0]["text"]["content"], "hello"); - assert_eq!(blocks[1]["image"]["category"], "logo"); - assert_eq!(blocks[2]["table"]["type"], "html"); - assert_eq!(blocks[2]["table"]["title"], "Totals"); + assert_eq!( + normalized.pages[0].extra_fields["blocks"], + payload["pages"][0]["blocks"] + ); + assert_eq!(normalized.pages[0].markdown, ""); + assert_eq!(normalized.pages[0].index, 0); + assert_eq!( + normalized.usage_info.as_ref().unwrap().pages_processed, + Some(1) + ); + match response_format { + crate::ocr::types::OcrResponseFormat::Litellm => { + assert!(normalized.provider_native_response.is_none()); + } + crate::ocr::types::OcrResponseFormat::Native => { + assert_eq!( + normalized.provider_native_response.as_ref(), + payload.as_object() + ); + } + } + assert_eq!( + normalized.into_json()["pages"][0]["blocks"], + payload["pages"][0]["blocks"] + ); } #[test] 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 ffabce84d05..0f982e5e88a 100644 --- a/litellm-rust/crates/core/src/llms/mistral/ocr/transformation.rs +++ b/litellm-rust/crates/core/src/llms/mistral/ocr/transformation.rs @@ -425,7 +425,9 @@ mod tests { #[rstest] #[case("table_format", json!("html"))] + #[case("table_format", json!("markdown"))] #[case("confidence_scores_granularity", json!("word"))] + #[case("confidence_scores_granularity", json!("page"))] #[case("document_annotation_prompt", json!("extract"))] #[case("include_blocks", json!(true))] #[case("id", json!("req-123"))] @@ -436,7 +438,9 @@ mod tests { #[rstest] #[case("pages", json!([0, 2]))] #[case("pages", json!("0,2-4"))] + #[case("pages", Value::Null)] #[case("include_image_base64", json!(true))] + #[case("include_image_base64", json!(false))] #[case("image_limit", json!(2))] #[case("image_min_size", json!(100))] #[case("bbox_annotation_format", json!({"type":"json_schema"}))] @@ -445,19 +449,28 @@ mod tests { #[case("extract_header", json!(true))] #[case("extract_footer", json!(false))] #[case("table_format", json!("html"))] + #[case("table_format", json!("markdown"))] #[case("confidence_scores_granularity", json!("word"))] + #[case("confidence_scores_granularity", json!("page"))] + #[case("confidence_scores_granularity", json!("block"))] #[case("include_blocks", json!(true))] + #[case("include_blocks", json!(false))] #[case("id", json!("req-123"))] - fn request_mapping_matches_python(#[case] name: &str, #[case] value: Value) { - let params: OpaqueParams = serde_json::from_value(json!({name: value.clone()})).unwrap(); + fn request_mapping_preserves_supplied_options(#[case] name: &str, #[case] value: Value) { + let arguments = serde_json::from_value(json!({name: value.clone()})).unwrap(); + let params = MistralOCRConfig + .map_ocr_params(&arguments, "model") + .unwrap(); let result = serde_json::to_value( MistralOCRConfig .transform_ocr_request("model", document(), ¶ms, &[]) .unwrap(), ) .unwrap(); - assert_eq!(result["model"], "model"); - assert_eq!(result[name], value); + assert_eq!( + result, + json!({"model":"model", "document":document(), name:value}) + ); } #[rstest] @@ -504,30 +517,25 @@ mod tests { #[rstest] fn transform_ocr_response_preserves_blocks_and_confidence_scores() { - let response: MistralOcrResponse = serde_json::from_value(json!({ - "pages":[{ - "index":0, - "markdown":"hello", - "images":[{"id":"img-0","image_base64":"data:image/png;base64,AA=="}], - "dimensions":{"width":612,"height":792,"dpi":72}, - "blocks":[{"type":"title","bbox":{"x":1},"confidence_scores":{"mean":0.98}}], - "confidence_scores":{"average_page_confidence_score":0.99,"minimum_page_confidence_score":0.97} - }], - "model":"returned-model", - "document_annotation":"{\"language\":\"en\"}", - "usage_info":{"pages_processed":1} - })) - .unwrap(); + let payload = json!({ + "pages":[{ + "index":0, + "markdown":"hello", + "images":[{"id":"img-0","image_base64":"data:image/png;base64,AA=="}], + "dimensions":{"width":612,"height":792,"dpi":72}, + "blocks":[{"type":"title","bbox":{"x":1},"confidence_scores":{"mean":0.98}}], + "confidence_scores":{"average_page_confidence_score":0.99,"minimum_page_confidence_score":0.97} + }], + "model":"returned-model", + "document_annotation":"{\"language\":\"en\"}", + "usage_info":{"pages_processed":1} + }); + let response: MistralOcrResponse = serde_json::from_value(payload.clone()).unwrap(); let result = normalize_response("model", response).unwrap().into_json(); - assert_eq!(result["pages"][0]["blocks"][0]["type"], "title"); - assert_eq!(result["pages"][0]["blocks"][0]["bbox"]["x"], 1); + assert_eq!(result["pages"][0]["blocks"], payload["pages"][0]["blocks"]); assert_eq!( - result["pages"][0]["blocks"][0]["confidence_scores"]["mean"], - 0.98 - ); - assert_eq!( - result["pages"][0]["confidence_scores"]["average_page_confidence_score"], - 0.99 + result["pages"][0]["confidence_scores"], + payload["pages"][0]["confidence_scores"] ); assert_eq!(result["pages"][0]["images"][0]["id"], "img-0"); assert_eq!(result["pages"][0]["dimensions"]["dpi"], 72); diff --git a/litellm-rust/crates/core/src/ocr/provider_config.rs b/litellm-rust/crates/core/src/ocr/provider_config.rs index ef9de23c913..37cc924fcc0 100644 --- a/litellm-rust/crates/core/src/ocr/provider_config.rs +++ b/litellm-rust/crates/core/src/ocr/provider_config.rs @@ -393,6 +393,13 @@ mod tests { #[rstest] #[case("reducto/parse-legacy", OcrConfigKind::ReductoLegacy)] #[case("reducto/future-parse-model", OcrConfigKind::ReductoV3)] + #[case("azure_ai/Cohere-parse-v5", OcrConfigKind::AzureCohere)] + #[case("azure_ai/cohere-parse-v5", OcrConfigKind::AzureCohere)] + #[case("azure_ai/cohere/parse-v5", OcrConfigKind::AzureCohere)] + #[case("azure_ai/invoice-parser", OcrConfigKind::AzureAi)] + #[case("azure_ai/parse-v5", OcrConfigKind::AzureAi)] + #[case("azure_ai/mistral-ocr-4-0", OcrConfigKind::AzureAi)] + #[case("azure_ai/mistral-document-ai-2512", OcrConfigKind::AzureAi)] #[case( "azure_ai/doc-intelligence/prebuilt-layout", OcrConfigKind::AzureDocumentIntelligence diff --git a/litellm-rust/crates/core/tests/azure_document_intelligence_ocr.rs b/litellm-rust/crates/core/tests/azure_document_intelligence_ocr.rs index 1da340b57d4..41fe0c734cf 100644 --- a/litellm-rust/crates/core/tests/azure_document_intelligence_ocr.rs +++ b/litellm-rust/crates/core/tests/azure_document_intelligence_ocr.rs @@ -1,5 +1,6 @@ use std::sync::{Arc, Mutex}; +use rstest::rstest; use serde_json::{Value, json}; use super::test_support::{MockResponse, mock_server, perform_ocr, wire_request}; @@ -48,33 +49,87 @@ async fn facade_maps_pages_features_and_url_document() { ); } +#[rstest] +#[case(json!({"pages":[true]}), crate::ocr::Error::Pages("expected only integers or only strings".into()))] +#[case(json!({"pages":[1,"2"]}), crate::ocr::Error::Pages("expected only integers or only strings".into()))] +#[case(json!({"pages":[-1]}), crate::ocr::Error::Pages("negative page index".into()))] +#[case(json!({"pages":"1&&features=bad"}), crate::ocr::Error::Pages("invalid native page range".into()))] +#[case(json!({"features":"languages&pages=1"}), crate::ocr::Error::Features)] +#[case(json!({"req_format":"azure"}), crate::ocr::Error::RequestFormat)] #[tokio::test] -async fn rejects_invalid_pages_features_and_format() { - for options in [ - json!({"pages":[true]}), - json!({"pages":[1,"2"]}), - json!({"pages":[-1]}), - json!({"pages":"1&&features=bad"}), - json!({"features":"languages&pages=1"}), - json!({"req_format":"azure"}), - ] { - let result = decode_request(OcrWireRequest { - model: "azure_ai/doc-intelligence/prebuilt-read".into(), - document: json!({"type":"document_url","document_url":"https://example.com/a.pdf"}), - api_key: Some("key".into()), - api_base: Some("http://127.0.0.1:1".into()), - custom_llm_provider: None, - extra_headers: None, - optional_params: options.as_object().unwrap().clone(), - input_sources: Default::default(), - timeout_seconds: None, - }); - let rejected = match result { - Ok(request) => perform_ocr(request).await.is_err(), - Err(_) => true, - }; - assert!(rejected, "accepted {options}"); +async fn rejects_invalid_pages_features_and_format( + #[case] options: Value, + #[case] expected: super::Error, +) { + let (base, seen, server) = mock_server(vec![MockResponse::json(json!({}))]).await; + 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(base), + custom_llm_provider: None, + extra_headers: None, + optional_params: options.as_object().unwrap().clone(), + input_sources: Default::default(), + timeout_seconds: Some(2.0), + }); + let result = match result { + Ok(request) => perform_ocr(request).await, + Err(error) => Err(error), + }; + server.abort(); + let _ = server.await; + assert!( + seen.lock().unwrap().is_empty(), + "sent invalid options: {options}" + ); + let error = result.unwrap_err(); + assert_eq!( + std::mem::discriminant(&error), + std::mem::discriminant(&expected) + ); + assert_eq!(error.http_status_code(), Some(400)); + assert_eq!(error.to_string(), expected.to_string()); +} + +#[rstest] +#[case(json!({}))] +#[case(json!({"req_format":"litellm"}))] +#[tokio::test] +async fn missing_native_fields_keep_page_text_without_retaining_raw_response( + #[case] options: Value, +) { + let operation = json!({ + "status":"succeeded", + "analyzeResult":{"pages":[{"pageNumber":1,"lines":[{"content":"hello"}]}]} + }); + let (base, seen, server) = mock_server(vec![MockResponse::json(operation)]).await; + let response = perform_ocr(wire_request( + "azure_ai/doc-intelligence/prebuilt-read", + &base, + options, + )) + .await + .unwrap(); + server.await.unwrap(); + + assert_eq!(response.pages.len(), 1); + assert_eq!(response.pages[0].index, 0); + assert_eq!(response.pages[0].markdown, "hello"); + assert_eq!(response.provider_native_response, None); + let serialized = response.into_json(); + assert_eq!(serialized.get("content"), Some(&Value::Null)); + assert_eq!(serialized.get("tables"), Some(&Value::Null)); + assert_eq!(serialized.get("keyValuePairs"), Some(&Value::Null)); + let requests = seen.lock().unwrap(); + assert_eq!(requests.len(), 1); + let target = requests[0].split_whitespace().nth(1).unwrap(); + let url = format!("{base}{target}"); + for field in ["pages", "features", "req_format"] { + assert_eq!(query_value(&url, field), None); } + let body: Value = serde_json::from_str(requests[0].split_once("\r\n\r\n").unwrap().1).unwrap(); + assert_eq!(body, json!({"base64Source":"YWJj"})); } #[tokio::test] diff --git a/litellm-rust/crates/core/tests/reducto_ocr.rs b/litellm-rust/crates/core/tests/reducto_ocr.rs index 0a7053b7429..0c25fd7a051 100644 --- a/litellm-rust/crates/core/tests/reducto_ocr.rs +++ b/litellm-rust/crates/core/tests/reducto_ocr.rs @@ -70,13 +70,26 @@ async fn request_mapping_matches_python( #[case("parse-v3")] #[case("parse-legacy")] #[tokio::test] -async fn data_uri_upload_preserves_multipart_headers(#[case] model: &str) { +async fn data_uri_upload_preserves_multipart_headers( + #[case] model: &str, + #[values("application/pdf", "image/png")] mime_type: &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!({})); + let document = if mime_type.starts_with("image/") { + json!({"type":"image_url","image_url":format!("data:{mime_type};base64,YWJj")}) + } else { + json!({"type":"document_url","document_url":format!("data:{mime_type};base64,YWJj")}) + }; + let mut request = super::LiteLLMOcrRequest { + document: serde_json::from_value::(document) + .unwrap() + .into(), + ..wire_request(&format!("reducto/{model}"), &base, json!({})) + }; request.transport.extra_headers = vec![ ("Content-Type".into(), "application/json".into()), ("X-Trace".into(), "upload-test".into()), @@ -94,9 +107,26 @@ async fn data_uri_upload_preserves_multipart_headers(#[case] model: &str) { .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")); + let multipart = requests[0].split_once("\r\n\r\n").unwrap().1; + assert!(multipart.contains(&format!("Content-Type: {mime_type}\r\n"))); + assert!(multipart.contains("\r\n\r\nabc\r\n--")); assert!(requests[1].starts_with("POST /parse ")); + let source_field = if model == "parse-legacy" { + "document_url" + } else { + "input" + }; + assert_eq!( + request_body(&requests[1]), + json!({source_field:"reducto://uploaded.pdf"}) + ); + for request in requests.iter() { + assert!( + request + .to_ascii_lowercase() + .contains("authorization: bearer test-key\r\n") + ); + } } struct ParseBoundary { @@ -168,17 +198,37 @@ async fn upload_failure_stops_before_parse() { } #[rstest] -#[case("https://example.com/a.pdf")] -#[case("reducto://")] -#[case("data:application/pdf;base64")] -#[case("data:application/pdf;base64,INVALID!")] +#[case("https://example.com/a.pdf", crate::ocr::Error::ReductoSource)] +#[case("reducto://", crate::ocr::Error::RequestField { path: "document file id".into() })] +#[case("data:application/pdf;base64", crate::ocr::Error::InvalidDataUri)] +#[case( + "data:application/pdf;base64,INVALID!", + crate::ocr::Error::InvalidDataUri +)] #[tokio::test] -async fn rejects_invalid_document_sources_before_network(#[case] source: &str) { +async fn rejects_invalid_document_sources_before_network( + #[case] source: &str, + #[case] expected: super::Error, +) { + let (base, seen, server) = mock_server(vec![MockResponse::json(json!({}))]).await; let request = super::test_support::with_source( - wire_request("reducto/parse-v3", "http://127.0.0.1:1", json!({})), + wire_request("reducto/parse-v3", &base, json!({})), source, ); - assert!(perform_ocr(request).await.is_err()); + let result = perform_ocr(request).await; + server.abort(); + let _ = server.await; + assert!( + seen.lock().unwrap().is_empty(), + "sent invalid source: {source}" + ); + let error = result.unwrap_err(); + assert_eq!( + std::mem::discriminant(&error), + std::mem::discriminant(&expected) + ); + assert_eq!(error.http_status_code(), Some(400)); + assert_eq!(error.to_string(), expected.to_string()); } #[test] From 27ccf7326bf02389b426755e1684a213b0536b75 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Thu, 17 Sep 2026 07:33:06 -0700 Subject: [PATCH 08/15] mistral alignment --- litellm-rust/crates/core/AGENTS.md | 18 +- .../src/llms/azure_ai/ocr/transformation.rs | 10 +- .../src/llms/base_llm/ocr/transformation.rs | 54 +-- .../src/llms/mistral/ocr/transformation.rs | 420 +++++++++--------- .../src/llms/vertex_ai/ocr/transformation.rs | 16 +- .../crates/core/src/ocr/provider_config.rs | 4 +- .../crates/core/tests/vertex_ai_ocr.rs | 6 +- 7 files changed, 271 insertions(+), 257 deletions(-) diff --git a/litellm-rust/crates/core/AGENTS.md b/litellm-rust/crates/core/AGENTS.md index 9ba7bfb5323..d591d241512 100644 --- a/litellm-rust/crates/core/AGENTS.md +++ b/litellm-rust/crates/core/AGENTS.md @@ -1,7 +1,23 @@ litellm-core is the LiteLLM SDK in Rust — it makes the LLM call. Each top-level call is a module under `src//` exposing a public entrypoint named after the route (`messages::messages()`, the Rust equivalent of `litellm.messages()`): you call it and get a typed non-streaming response back. -A route module owns everything the call needs: types, the provider template trait, provider transforms (under `providers/`), provider/auth/URL resolution, and the handler that performs the HTTP call. Handlers belong here, never in a host crate. +A route module owns the call entrypoint, runtime types, provider/auth/URL resolution, and the handler that performs the HTTP call. OCR provider transforms and the base provider trait live under `src/llms/`, mirroring their Python source paths. Other routes still use `src/providers/` and route-local traits. Handlers belong in core, never in a host crate Not here: serving HTTP (axum routes, extractors), config file reading, rollout state, databases, or host-specific callback execution. Core owns lifecycle sequencing and callback payload construction; hosts execute the selected integrations. Env reads are limited to credential fallback in a route's `prepare.rs`. Routes (messages, ocr, realtime) and providers (anthropic, mistral, openai) are modules, not crates. + +## Python/Rust transformation pairs + +Use the base OCR and Mistral OCR pairs as the reference when aligning transformations. Derive `src/.rs` from `litellm/.py`, preserving meaningful basenames such as `messages_transformation` + +Keep corresponding operation names and parameter names when their responsibilities match. Rust types retain the Python semantic name with Rust acronym casing (`BaseOCRConfig` / `BaseOcrConfig`, `MistralOCRConfig` / `MistralOcrConfig`). Private Python helpers can drop their leading underscore. Give Rust adapter helpers distinct responsibility names rather than duplicating trait method names + +Order OCR config methods as supported parameters, credential metadata and connection resolution, health-check input, parameter mapping, environment validation, URL construction, request transformation, async request transformation, response transformation, async response transformation, and error conversion. Put constants and data types before the config, private helpers after it in operation order, and tests last. Rust-only trait hooks follow the corresponding Python methods + +Use trait defaults for unchanged inherited behavior and explicit delegation for shared provider behavior. Keep typed inputs, ownership, `Result`, and async I/O idiomatic. A matching path or symbol identifies the counterpart, not a claim of full behavioral parity + +Use named `#[rstest]` cases for independent input/output scenarios instead of loops or repeated calls in one test. Inject reusable setup with `#[fixture]` arguments and use `#[with(...)]` for fixture overrides. Keep assertions about the same result together + +For base OCR, Python response models correspond to `src/ocr/types.rs`; Rust context/environment types support the runtime. `BaseOcrConfig::prepare_request` corresponds to Python's HTTP-handler preparation rather than a `BaseOCRConfig` method, and `validate_request_body` is a Rust-only hook + +For Mistral, `async_transform_ocr_request` uses the base default in both languages. `resolve_headers` and `build_ocr_url` implement the respective environment and URL operations, and `normalize_response` implements the typed part of response transformation. Existing auth key/header handling and top-level response-extra preservation differ between languages; layout refactors must preserve those behaviors and verify them with the existing tests 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 1a909abc2d6..122f1dbce53 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 @@ -5,7 +5,7 @@ use serde_json::Value; use crate::call_arguments::CallArguments; use crate::constants::AZURE_AI_OCR_PATH; use crate::llms::base_llm::ocr::transformation::{BaseOcrConfig, OcrRequestContext}; -use crate::llms::mistral::ocr::transformation::{MistralOCRConfig, MistralOcrRequest}; +use crate::llms::mistral::ocr::transformation::{MistralOcrConfig, MistralOcrRequest}; use crate::ocr::OcrClient; use crate::ocr::document::{inline_remote_document, validate_inline_document}; use crate::ocr::prepare::credential_env; @@ -60,11 +60,11 @@ impl BaseOcrConfig for AzureAIOCRConfig { params: &OpaqueParams, headers: &[(String, String)], ) -> Result { - MistralOCRConfig.transform_ocr_request(model, document, params, headers) + 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) + MistralOcrConfig.get_supported_ocr_params(model) } fn map_ocr_params( @@ -72,7 +72,7 @@ impl BaseOcrConfig for AzureAIOCRConfig { arguments: &CallArguments, model: &str, ) -> Result { - MistralOCRConfig.map_ocr_params(arguments, model) + MistralOcrConfig.map_ocr_params(arguments, model) } async fn async_transform_ocr_request( @@ -98,7 +98,7 @@ impl BaseOcrConfig for AzureAIOCRConfig { raw_response: &[u8], request_format: crate::ocr::types::OcrResponseFormat, ) -> Result { - MistralOCRConfig.transform_ocr_response(model, raw_response, request_format) + MistralOcrConfig.transform_ocr_response(model, raw_response, request_format) } fn validate_request_body(&self, body: &Value) -> Result<(), crate::ocr::Error> { 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 8af304b7d8d..4c4b7a066ef 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 @@ -13,6 +13,8 @@ use crate::ocr::types::{ PreparedOcrRequest, ResolvedOcrCredentials, }; +const HEALTH_CHECK_PDF_DATA_URI: &str = "data:application/pdf;base64,JVBERi0xLjQKJeLjz9MKMyAwIG9iago8PC9UeXBlIC9QYWdlCi9QYXJlbnQgMSAwIFIKL01lZGlhQm94IFswIDAgNjEyIDc5Ml0KL0NvbnRlbnRzIDQgMCBSCi9SZXNvdXJjZXMgPDwvRm9udCA8PC9GMSAyIDAgUj4+Pj4+PgplbmRvYmoKNCAwIG9iago8PC9MZW5ndGggNDQ+PgpzdHJlYW0KQlQKL0YxIDI0IFRmCjEwMCA3MDAgVGQKKHRlc3QpIFRqCkVUCmVuZHN0cmVhbQplbmRvYmoKMiAwIG9iago8PC9UeXBlIC9Gb250Ci9TdWJ0eXBlIC9UeXBlMQovQmFzZUZvbnQgL0hlbHZldGljYT4+CmVuZG9iagoxIDAgb2JqCjw8L1R5cGUgL1BhZ2VzCi9LaWRzIFszIDAgUl0KL0NvdW50IDE+PgplbmRvYmoKNSAwIG9iago8PC9UeXBlIC9DYXRhbG9nCi9QYWdlcyAxIDAgUj4+CmVuZG9iagp0cmFpbGVyCjw8L1NpemUgNgovUm9vdCA1IDAgUj4+CnN0YXJ0eHJlZgozMjQKJSVFT0Y="; + /// Output of `validate_environment`: whatever a provider resolves up front /// (headers at minimum; Vertex also carries the project id). pub(crate) trait OcrEnvironment: Send + Sync { @@ -25,13 +27,31 @@ impl OcrEnvironment for Vec<(String, String)> { } } -const HEALTH_CHECK_PDF_DATA_URI: &str = "data:application/pdf;base64,JVBERi0xLjQKJeLjz9MKMyAwIG9iago8PC9UeXBlIC9QYWdlCi9QYXJlbnQgMSAwIFIKL01lZGlhQm94IFswIDAgNjEyIDc5Ml0KL0NvbnRlbnRzIDQgMCBSCi9SZXNvdXJjZXMgPDwvRm9udCA8PC9GMSAyIDAgUj4+Pj4+PgplbmRvYmoKNCAwIG9iago8PC9MZW5ndGggNDQ+PgpzdHJlYW0KQlQKL0YxIDI0IFRmCjEwMCA3MDAgVGQKKHRlc3QpIFRqCkVUCmVuZHN0cmVhbQplbmRvYmoKMiAwIG9iago8PC9UeXBlIC9Gb250Ci9TdWJ0eXBlIC9UeXBlMQovQmFzZUZvbnQgL0hlbHZldGljYT4+CmVuZG9iagoxIDAgb2JqCjw8L1R5cGUgL1BhZ2VzCi9LaWRzIFszIDAgUl0KL0NvdW50IDE+PgplbmRvYmoKNSAwIG9iago8PC9UeXBlIC9DYXRhbG9nCi9QYWdlcyAxIDAgUj4+CmVuZG9iagp0cmFpbGVyCjw8L1NpemUgNgovUm9vdCA1IDAgUj4+CnN0YXJ0eHJlZgozMjQKJSVFT0Y="; +#[derive(Clone, Copy)] +pub(crate) struct OcrRequestContext<'a> { + pub client: &'a OcrClient, + pub connection: &'a OcrConnection, +} + +#[derive(Clone, Copy)] +pub(crate) struct OcrResponseContext<'a> { + pub client: &'a OcrClient, + pub connection: &'a OcrConnection, + pub hooks: &'a Arc, + pub request_format: OcrResponseFormat, + pub url: &'a str, + pub headers: &'a [(String, String)], +} pub(crate) trait BaseOcrConfig: Send + Sync + Sized + 'static { type OcrParams: Send + Sync; type ProviderRequest: Serialize + Send; type Environment: OcrEnvironment; + fn get_supported_ocr_params(&self, _model: &str) -> &'static [&'static str] { + &[] + } + fn get_api_key_env_var(&self) -> Option<&'static str> { None } @@ -56,6 +76,12 @@ pub(crate) trait BaseOcrConfig: Send + Sync + Sized + 'static { } } + fn map_ocr_params( + &self, + non_default_params: &CallArguments, + model: &str, + ) -> Result; + fn validate_environment( &self, request: &PreparedOcrRequest, @@ -69,16 +95,6 @@ pub(crate) trait BaseOcrConfig: Send + Sync + Sized + 'static { environment: &Self::Environment, ) -> Result; - fn get_supported_ocr_params(&self, _model: &str) -> &'static [&'static str] { - &[] - } - - fn map_ocr_params( - &self, - arguments: &CallArguments, - model: &str, - ) -> Result; - fn transform_ocr_request( &self, model: &str, @@ -193,19 +209,3 @@ pub(crate) fn decode_and_normalize_response( ..normalize(model, decoded.data)? }) } - -#[derive(Clone, Copy)] -pub(crate) struct OcrRequestContext<'a> { - pub client: &'a OcrClient, - pub connection: &'a OcrConnection, -} - -#[derive(Clone, Copy)] -pub(crate) struct OcrResponseContext<'a> { - pub client: &'a OcrClient, - pub connection: &'a OcrConnection, - pub hooks: &'a Arc, - pub request_format: OcrResponseFormat, - pub url: &'a str, - pub headers: &'a [(String, String)], -} 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 0f982e5e88a..71dcf88cd0f 100644 --- a/litellm-rust/crates/core/src/llms/mistral/ocr/transformation.rs +++ b/litellm-rust/crates/core/src/llms/mistral/ocr/transformation.rs @@ -3,16 +3,17 @@ use serde_json::Value; use crate::call_arguments::CallArguments; use crate::constants::MISTRAL_OCR_API_BASE; -use crate::llms::base_llm::ocr::transformation::{BaseOcrConfig, OcrRequestContext}; +use crate::llms::base_llm::ocr::transformation::{BaseOcrConfig, decode_and_normalize_response}; use crate::ocr::OcrClient; use crate::ocr::prepare::credential_env; use crate::ocr::types::{ - LiteLLMOcrResponse, OcrConnection, OcrDocument, OcrPage, OcrUsageInfo, PreparedOcrRequest, + LiteLLMOcrResponse, OcrConnection, OcrDocument, OcrPage, OcrResponseFormat, OcrUsageInfo, + PreparedOcrRequest, }; use crate::params::OpaqueParams; use crate::url_utils::ApiUrl; -const MISTRAL_API_KEY_ENV: &str = "MISTRAL_API_KEY"; +const MISTRAL_OCR_API_KEY_ENV_VAR: &str = "MISTRAL_API_KEY"; #[derive(Clone, Debug, Serialize, Deserialize)] pub(crate) struct MistralOcrRequest { @@ -24,8 +25,6 @@ pub(crate) struct MistralOcrRequest { #[derive(Clone, Debug, Default, Deserialize)] pub(crate) struct MistralOcrResponse { - #[serde(flatten)] - pub extra_fields: serde_json::Map, #[serde(default)] pub pages: Vec, #[serde( @@ -35,51 +34,19 @@ pub(crate) struct MistralOcrResponse { pub model: Option>, pub document_annotation: Option, pub usage_info: Option, + + #[serde(flatten)] + pub extra_fields: serde_json::Map, } #[derive(Clone, Debug, Default)] -pub(crate) struct MistralOCRConfig; +pub(crate) struct MistralOcrConfig; -impl BaseOcrConfig for 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: &PreparedOcrRequest, - _client: &OcrClient, - ) -> Result { - self.validate_environment(&request.connection, &credential_env) - } - - fn get_complete_url( - &self, - request: &PreparedOcrRequest, - _params: &Self::OcrParams, - _environment: &Self::Environment, - ) -> Result { - self.get_complete_url(request.connection.api_base.as_deref()) - } - - fn transform_ocr_request( - &self, - model: &str, - document: OcrDocument, - optional_params: &OpaqueParams, - _headers: &[(String, String)], - ) -> Result { - Ok(MistralOcrRequest { - model: model.to_string(), - document, - params: optional_params.clone(), - }) - } - fn get_supported_ocr_params(&self, _model: &str) -> &'static [&'static str] { &[ "pages", @@ -98,40 +65,104 @@ impl BaseOcrConfig for MistralOCRConfig { ] } + fn get_api_key_env_var(&self) -> Option<&'static str> { + Some(MISTRAL_OCR_API_KEY_ENV_VAR) + } + fn map_ocr_params( &self, - arguments: &CallArguments, + non_default_params: &CallArguments, model: &str, ) -> Result { - Ok(arguments + Ok(non_default_params .select(self.get_supported_ocr_params(model)) .into()) } - async fn async_transform_ocr_request( + async fn validate_environment( + &self, + request: &PreparedOcrRequest, + _client: &OcrClient, + ) -> Result { + self.resolve_headers(&request.connection, &credential_env) + } + + fn get_complete_url( + &self, + request: &PreparedOcrRequest, + _optional_params: &Self::OcrParams, + _environment: &Self::Environment, + ) -> Result { + self.build_ocr_url(request.connection.api_base.as_deref()) + } + + fn transform_ocr_request( &self, model: &str, document: OcrDocument, optional_params: &OpaqueParams, - headers: &[(String, String)], - _context: OcrRequestContext<'_>, + _headers: &[(String, String)], ) -> Result { - self.transform_ocr_request(model, document, optional_params, headers) + Ok(MistralOcrRequest { + model: model.to_string(), + document, + params: optional_params.clone(), + }) } fn transform_ocr_response( &self, model: &str, raw_response: &[u8], - request_format: crate::ocr::types::OcrResponseFormat, + request_format: OcrResponseFormat, ) -> Result { - crate::llms::base_llm::ocr::transformation::decode_and_normalize_response( - model, - raw_response, - request_format, - normalize_response, + decode_and_normalize_response(model, raw_response, request_format, normalize_response) + } +} + +impl MistralOcrConfig { + fn resolve_headers( + &self, + connection: &OcrConnection, + env_lookup: &(dyn Fn(&str) -> Option + Sync), + ) -> Result, crate::ocr::Error> { + if crate::http_utils::has_header(&connection.extra_headers, "authorization") { + return Ok(connection.extra_headers.clone()); + } + let api_key = connection + .api_key + .as_deref() + .map(str::trim) + .filter(|key| !key.is_empty()) + .map(str::to_string) + .or_else(|| { + 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_OCR_API_KEY_ENV_VAR, + })?; + Ok( + std::iter::once(("Authorization".into(), format!("Bearer {api_key}"))) + .chain(connection.extra_headers.clone()) + .collect(), ) } + + fn build_ocr_url(&self, api_base: Option<&str>) -> Result { + let base = api_base + .map(str::trim) + .filter(|base| !base.is_empty()) + .unwrap_or(MISTRAL_OCR_API_BASE); + ApiUrl::parse(base) + .and_then(|url| url.complete_path(&["v1", "ocr"])) + .map(|url| url.into_string()) + .map_err(|_| crate::ocr::Error::RequestField { + path: "api_base".into(), + }) + } } pub(crate) fn normalize_response( @@ -155,58 +186,33 @@ pub(crate) fn normalize_response( }) } -impl MistralOCRConfig { - fn get_complete_url(&self, api_base: Option<&str>) -> Result { - let base = api_base - .map(str::trim) - .filter(|base| !base.is_empty()) - .unwrap_or(MISTRAL_OCR_API_BASE); - ApiUrl::parse(base) - .and_then(|url| url.complete_path(&["v1", "ocr"])) - .map(|url| url.into_string()) - .map_err(|_| crate::ocr::Error::RequestField { - path: "api_base".into(), - }) - } - - fn validate_environment( - &self, - connection: &OcrConnection, - env_lookup: &(dyn Fn(&str) -> Option + Sync), - ) -> Result, crate::ocr::Error> { - if crate::http_utils::has_header(&connection.extra_headers, "authorization") { - return Ok(connection.extra_headers.clone()); - } - let api_key = connection - .api_key - .as_deref() - .map(str::trim) - .filter(|key| !key.is_empty()) - .map(str::to_string) - .or_else(|| { - 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, - })?; - Ok( - std::iter::once(("Authorization".into(), format!("Bearer {api_key}"))) - .chain(connection.extra_headers.clone()) - .collect(), - ) - } -} - #[cfg(test)] mod tests { - use rstest::rstest; + use rstest::{fixture, rstest}; use serde_json::{Value, json}; use super::*; + #[fixture] + fn document() -> OcrDocument { + serde_json::from_value( + json!({"type":"document_url","document_url":"https://example.com/a.pdf"}), + ) + .unwrap() + } + + #[fixture] + fn connection( + #[default(None)] api_key: Option<&str>, + #[default(vec![])] extra_headers: Vec<(String, String)>, + ) -> OcrConnection { + OcrConnection { + api_key: api_key.map(str::to_string), + extra_headers, + ..OcrConnection::default() + } + } + #[test] fn explicit_null_model_does_not_use_the_missing_model_default() { let response = serde_json::from_value(json!({"model":null})).unwrap(); @@ -216,38 +222,38 @@ mod tests { )); } - #[test] - fn response_validates_normalized_shapes_at_the_provider_boundary() { - for (payload, path) in [ - (json!({"pages":[42]}), "pages[0]"), - (json!({"pages":[{"index":0}]}), "pages[0]"), - ( - json!({"pages":[{"index":0,"markdown":42}]}), - "pages[0].markdown", - ), - ( - json!({"pages":[{"index":0,"markdown":"","images":[42]}]}), - "pages[0].images[0]", - ), - ( - json!({"pages":[{"index":0,"markdown":"","dimensions":{"width":1.5}}]}), - "pages[0].dimensions.width", - ), - ( - json!({"usage_info":{"pages_processed":"bad"}}), - "usage_info.pages_processed", - ), - ] { - let error = crate::ocr::json::decode_response::( - &serde_json::to_vec(&payload).unwrap(), - false, - ) - .unwrap_err(); - assert!(matches!( - error, - crate::ocr::Error::ResponseField { path: actual } if actual == path - )); - } + #[rstest] + #[case::non_object_page(json!({"pages":[42]}), "pages[0]")] + #[case::missing_markdown(json!({"pages":[{"index":0}]}), "pages[0]")] + #[case::non_string_markdown( + json!({"pages":[{"index":0,"markdown":42}]}), + "pages[0].markdown" + )] + #[case::non_object_image( + json!({"pages":[{"index":0,"markdown":"","images":[42]}]}), + "pages[0].images[0]" + )] + #[case::fractional_width( + json!({"pages":[{"index":0,"markdown":"","dimensions":{"width":1.5}}]}), + "pages[0].dimensions.width" + )] + #[case::invalid_page_count( + json!({"usage_info":{"pages_processed":"bad"}}), + "usage_info.pages_processed" + )] + fn response_validates_normalized_shapes_at_the_provider_boundary( + #[case] payload: Value, + #[case] path: &str, + ) { + let error = crate::ocr::json::decode_response::( + &serde_json::to_vec(&payload).unwrap(), + false, + ) + .unwrap_err(); + assert!(matches!( + error, + crate::ocr::Error::ResponseField { path: actual } if actual == path + )); } #[test] @@ -283,7 +289,7 @@ mod tests { let input = serde_json::from_value(json!({"pages":null,"extract_header":false,"unknown":true})) .unwrap(); - let params = MistralOCRConfig.map_ocr_params(&input, "model").unwrap(); + let params = MistralOcrConfig.map_ocr_params(&input, "model").unwrap(); assert_eq!( serde_json::to_value(params).unwrap(), json!({"pages":null,"extract_header":false}) @@ -292,11 +298,11 @@ mod tests { assert_eq!(input.get("pages"), Some(&Value::Null)); } - #[test] - fn request_transform_uses_already_mapped_params_without_filtering_again() { + #[rstest] + fn request_transform_uses_already_mapped_params_without_filtering_again(document: OcrDocument) { let params = serde_json::from_value(json!({"extension":{"nested":null}})).unwrap(); - let body = MistralOCRConfig - .transform_ocr_request("model", document(), ¶ms, &[]) + let body = MistralOcrConfig + .transform_ocr_request("model", document, ¶ms, &[]) .unwrap(); assert_eq!( serde_json::to_value(body).unwrap()["extension"], @@ -307,7 +313,7 @@ mod tests { #[test] 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 + let response = MistralOcrConfig .transform_ocr_response("model", raw, crate::ocr::types::OcrResponseFormat::Native) .unwrap(); assert_eq!(response.pages[0].index, 2); @@ -315,27 +321,23 @@ mod tests { assert_eq!(native["pages"][0]["index"], "2"); assert_eq!(native["provider_extension"], false); assert_eq!(response.extra_fields["provider_extension"], false); + } + + #[rstest] + fn raw_response_transform_rejects_invalid_page( + #[values(OcrResponseFormat::Litellm, OcrResponseFormat::Native)] + request_format: OcrResponseFormat, + ) { assert!( - MistralOCRConfig - .transform_ocr_response( - "model", - br#"{"pages":[{"index":0}]}"#, - crate::ocr::types::OcrResponseFormat::Litellm - ) + MistralOcrConfig + .transform_ocr_response("model", br#"{"pages":[{"index":0}]}"#, request_format) .is_err() ); } fn mapped_params(value: Value) -> Value { let params = serde_json::from_value(value).unwrap(); - serde_json::to_value(MistralOCRConfig.map_ocr_params(¶ms, "model").unwrap()).unwrap() - } - - fn document() -> OcrDocument { - serde_json::from_value( - json!({"type":"document_url","document_url":"https://example.com/a.pdf"}), - ) - .unwrap() + serde_json::to_value(MistralOcrConfig.map_ocr_params(¶ms, "model").unwrap()).unwrap() } #[rstest] @@ -456,20 +458,24 @@ mod tests { #[case("include_blocks", json!(true))] #[case("include_blocks", json!(false))] #[case("id", json!("req-123"))] - fn request_mapping_preserves_supplied_options(#[case] name: &str, #[case] value: Value) { + fn request_mapping_preserves_supplied_options( + document: OcrDocument, + #[case] name: &str, + #[case] value: Value, + ) { let arguments = serde_json::from_value(json!({name: value.clone()})).unwrap(); - let params = MistralOCRConfig + let params = MistralOcrConfig .map_ocr_params(&arguments, "model") .unwrap(); let result = serde_json::to_value( - MistralOCRConfig - .transform_ocr_request("model", document(), ¶ms, &[]) + MistralOcrConfig + .transform_ocr_request("model", document.clone(), ¶ms, &[]) .unwrap(), ) .unwrap(); assert_eq!( result, - json!({"model":"model", "document":document(), name:value}) + json!({"model":"model", "document":document, name:value}) ); } @@ -482,13 +488,14 @@ mod tests { #[case("include_blocks", json!(true))] #[case("pages", json!([0,1]))] fn transform_ocr_request_includes_each_optional_param( + document: OcrDocument, #[case] name: &str, #[case] value: Value, ) { let params: OpaqueParams = serde_json::from_value(json!({name:value.clone()})).unwrap(); let result = serde_json::to_value( - MistralOCRConfig - .transform_ocr_request("mistral-ocr-latest", document(), ¶ms, &[]) + MistralOcrConfig + .transform_ocr_request("mistral-ocr-latest", document, ¶ms, &[]) .unwrap(), ) .unwrap(); @@ -497,7 +504,7 @@ mod tests { } #[rstest] - fn transform_ocr_request_includes_multiple_new_params() { + fn transform_ocr_request_includes_multiple_new_params(document: OcrDocument) { let params: OpaqueParams = serde_json::from_value(json!({ "table_format":"html", "confidence_scores_granularity":"page", @@ -505,8 +512,8 @@ mod tests { })) .unwrap(); let result = serde_json::to_value( - MistralOCRConfig - .transform_ocr_request("mistral-ocr-latest", document(), ¶ms, &[]) + MistralOcrConfig + .transform_ocr_request("mistral-ocr-latest", document, ¶ms, &[]) .unwrap(), ) .unwrap(); @@ -565,69 +572,60 @@ mod tests { assert!(result["pages"][0]["dimensions"].is_null()); } - #[test] - fn complete_url_defaults_and_dedupes_v1() { + #[rstest] + #[case::default_base(None, "https://api.mistral.ai/v1/ocr")] + #[case::versioned_base( + Some("https://example.com/v1?tenant=a"), + "https://example.com/v1/ocr?tenant=a" + )] + #[case::complete_endpoint( + Some("https://example.com/v1/ocr?tenant=a"), + "https://example.com/v1/ocr?tenant=a" + )] + fn complete_url_defaults_and_dedupes_v1( + #[case] api_base: Option<&str>, + #[case] expected: &str, + ) { + assert_eq!(MistralOcrConfig.build_ocr_url(api_base).unwrap(), expected); + } + + #[rstest] + #[case::explicit_key(Some("explicit"), "Bearer explicit")] + #[case::environment_fallback(None, "Bearer environment")] + fn environment_prefers_explicit_key_then_environment( + #[case] _api_key: Option<&str>, + #[case] expected: &str, + #[with(_api_key)] connection: OcrConnection, + ) { assert_eq!( - MistralOCRConfig.get_complete_url(None).unwrap(), - "https://api.mistral.ai/v1/ocr" - ); - assert_eq!( - MistralOCRConfig - .get_complete_url(Some("https://example.com/v1?tenant=a")) - .unwrap(), - "https://example.com/v1/ocr?tenant=a" - ); - assert_eq!( - MistralOCRConfig - .get_complete_url(Some("https://example.com/v1/ocr?tenant=a")) - .unwrap(), - "https://example.com/v1/ocr?tenant=a" + MistralOcrConfig + .resolve_headers(&connection, &|_| Some("environment".into())) + .unwrap()[0], + ("Authorization".into(), expected.into()) ); } - #[test] - fn environment_prefers_explicit_key_then_environment() { - let explicit = OcrConnection { - api_key: Some("explicit".into()), - ..OcrConnection::default() - }; + #[rstest] + fn environment_preserves_forwarded_authorization( + #[with(None, vec![("authorization".into(), "Bearer forwarded".into())])] + connection: OcrConnection, + ) { assert_eq!( - MistralOCRConfig - .validate_environment(&explicit, &|_| Some("environment".into())) - .unwrap()[0], - ("Authorization".into(), "Bearer explicit".into()) - ); - - assert_eq!( - MistralOCRConfig - .validate_environment(&OcrConnection::default(), &|_| Some("environment".into())) - .unwrap()[0], - ("Authorization".into(), "Bearer environment".into()) - ); - } - - #[test] - fn environment_preserves_forwarded_authorization() { - let connection = OcrConnection { - extra_headers: vec![("authorization".into(), "Bearer forwarded".into())], - ..OcrConnection::default() - }; - assert_eq!( - MistralOCRConfig - .validate_environment(&connection, &|_| None) + MistralOcrConfig + .resolve_headers(&connection, &|_| None) .unwrap(), connection.extra_headers ); } - #[test] - fn environment_rejects_missing_key() { + #[rstest] + fn environment_rejects_missing_key(connection: OcrConnection) { assert!(matches!( - MistralOCRConfig.validate_environment(&OcrConnection::default(), &|_| None), + MistralOcrConfig.resolve_headers(&connection, &|_| None), Err(crate::ocr::Error::Auth( litellm_auth::Error::MissingApiKey { provider: "Mistral", - environment_variable: MISTRAL_API_KEY_ENV, + environment_variable: MISTRAL_OCR_API_KEY_ENV_VAR, } )) )); 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 337fa76cfe2..1183043e9ee 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 @@ -6,7 +6,7 @@ use crate::call_arguments::CallArguments; use crate::llms::base_llm::ocr::transformation::{ BaseOcrConfig, OcrEnvironment, OcrRequestContext, }; -use crate::llms::mistral::ocr::transformation::{MistralOCRConfig, MistralOcrRequest}; +use crate::llms::mistral::ocr::transformation::{MistralOcrConfig, MistralOcrRequest}; use crate::ocr::OcrClient; use crate::ocr::document::{inline_remote_document, validate_inline_document}; use crate::ocr::prepare::credential_env; @@ -68,11 +68,11 @@ impl BaseOcrConfig for VertexAIOCRConfig { params: &OpaqueParams, headers: &[(String, String)], ) -> Result { - MistralOCRConfig.transform_ocr_request(model, document, params, headers) + 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) + MistralOcrConfig.get_supported_ocr_params(model) } fn map_ocr_params( @@ -80,7 +80,7 @@ impl BaseOcrConfig for VertexAIOCRConfig { arguments: &CallArguments, model: &str, ) -> Result { - MistralOCRConfig.map_ocr_params(arguments, model) + MistralOcrConfig.map_ocr_params(arguments, model) } async fn async_transform_ocr_request( @@ -106,7 +106,7 @@ impl BaseOcrConfig for VertexAIOCRConfig { raw_response: &[u8], request_format: crate::ocr::types::OcrResponseFormat, ) -> Result { - MistralOCRConfig.transform_ocr_response(model, raw_response, request_format) + MistralOcrConfig.transform_ocr_response(model, raw_response, request_format) } fn validate_request_body(&self, body: &Value) -> Result<(), crate::ocr::Error> { @@ -320,7 +320,7 @@ mod tests { use std::time::Duration; use crate::llms::base_llm::ocr::transformation::BaseOcrConfig; - use crate::llms::mistral::ocr::transformation::MistralOCRConfig; + use crate::llms::mistral::ocr::transformation::MistralOcrConfig; use crate::llms::vertex_ai::ocr::transformation::VertexAIOCRConfig; use crate::ocr::test_support::ocr_client; @@ -344,7 +344,7 @@ mod tests { let vertex = crate::ocr::prepare::prepare_request( crate::ocr::test_support::resolved_request(vertex), ); - let direct_http = MistralOCRConfig + let direct_http = MistralOcrConfig .prepare_request(&direct, &client) .await .unwrap(); @@ -379,7 +379,7 @@ mod tests { &json!({"pages": [{"index": 0, "markdown": "hello"}], "extra": "preserved"}), ) .unwrap(); - let direct_response = MistralOCRConfig + let direct_response = MistralOcrConfig .transform_ocr_response(&direct.model, &payload, Default::default()) .unwrap() .into_json(); diff --git a/litellm-rust/crates/core/src/ocr/provider_config.rs b/litellm-rust/crates/core/src/ocr/provider_config.rs index 37cc924fcc0..fcbea54779f 100644 --- a/litellm-rust/crates/core/src/ocr/provider_config.rs +++ b/litellm-rust/crates/core/src/ocr/provider_config.rs @@ -10,7 +10,7 @@ use crate::llms::azure_ai::ocr::document_intelligence::transformation::AzureDocu use crate::llms::azure_ai::ocr::transformation::AzureAIOCRConfig; use crate::llms::base_llm::ocr::transformation::{BaseOcrConfig, OcrResponseContext}; use crate::llms::cohere::ocr::transformation::CohereParseConfig; -use crate::llms::mistral::ocr::transformation::MistralOCRConfig; +use crate::llms::mistral::ocr::transformation::MistralOcrConfig; use crate::llms::reducto::ocr::transformation::{ReductoParseLegacyConfig, ReductoParseV3Config}; use crate::llms::vertex_ai::ocr::deepseek_transformation::VertexAIDeepSeekOCRConfig; use crate::llms::vertex_ai::ocr::transformation::VertexAIOCRConfig; @@ -26,7 +26,7 @@ macro_rules! dispatch_config { (@arms $config:expr, $method:ident($($argument:expr),*), $($suffix:tt)*) => { match $config { OcrConfigKind::Cohere => CohereParseConfig.$method($($argument),*)$($suffix)*, - OcrConfigKind::Mistral => MistralOCRConfig.$method($($argument),*)$($suffix)*, + OcrConfigKind::Mistral => MistralOcrConfig.$method($($argument),*)$($suffix)*, OcrConfigKind::AzureAi => AzureAIOCRConfig.$method($($argument),*)$($suffix)*, OcrConfigKind::AzureCohere => AzureAICohereParseConfig.$method($($argument),*)$($suffix)*, OcrConfigKind::AzureDocumentIntelligence => AzureDocumentIntelligenceOCRConfig.$method($($argument),*)$($suffix)*, diff --git a/litellm-rust/crates/core/tests/vertex_ai_ocr.rs b/litellm-rust/crates/core/tests/vertex_ai_ocr.rs index ebee4046e23..1908c7aa347 100644 --- a/litellm-rust/crates/core/tests/vertex_ai_ocr.rs +++ b/litellm-rust/crates/core/tests/vertex_ai_ocr.rs @@ -103,7 +103,7 @@ async fn adapters_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::mistral::ocr::transformation::MistralOcrConfig; use crate::llms::vertex_ai::ocr::transformation::VertexAIOCRConfig; use crate::ocr::test_support::ocr_client; @@ -125,7 +125,7 @@ async fn adapters_build_complete_requests_and_share_mistral_normalization() { crate::ocr::prepare::prepare_request(super::test_support::resolved_request(direct)); let vertex = crate::ocr::prepare::prepare_request(super::test_support::resolved_request(vertex)); - let direct_http = MistralOCRConfig + let direct_http = MistralOcrConfig .prepare_request(&direct, &client) .await .unwrap(); @@ -157,7 +157,7 @@ async fn adapters_build_complete_requests_and_share_mistral_normalization() { } let payload = json!({"pages": [{"index": 0, "markdown": "hello"}], "extra": "preserved"}); let raw = serde_json::to_vec(&payload).unwrap(); - let direct_response = MistralOCRConfig + let direct_response = MistralOcrConfig .transform_ocr_response( &direct.model, &raw, From b063ffe88399417f4578034c8112c91de6fa8767 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Thu, 17 Sep 2026 07:52:14 -0700 Subject: [PATCH 09/15] providers folder is gone --- litellm-rust/crates/core/AGENTS.md | 6 +- .../core/src/audio_transcription/handler.rs | 7 +- .../core/src/audio_transcription/mod.rs | 1 - .../core/src/audio_transcription/prepare.rs | 16 +- .../core/src/audio_transcription/types.rs | 6 +- .../core/src/chat_completions/common_utils.rs | 10 +- .../core/src/chat_completions/handler.rs | 4 +- .../crates/core/src/chat_completions/mod.rs | 1 - .../core/src/chat_completions/prepare.rs | 12 +- .../crates/core/src/chat_completions/tests.rs | 2 +- .../crates/core/src/chat_completions/types.rs | 6 +- litellm-rust/crates/core/src/lib.rs | 4 +- .../get_llm_provider_logic.rs} | 0 .../crates/core/src/litellm_core_utils/mod.rs | 1 + .../anthropic/chat}/mod.rs | 0 .../anthropic/chat}/tests.rs | 2 +- .../anthropic/chat}/transformation.rs | 194 ++++++------ .../messages/mod.rs | 0 .../messages/transformation.rs | 44 ++- .../experimental_pass_through}/mod.rs | 0 .../crates/core/src/llms/anthropic/mod.rs | 2 + .../anthropic/messages_transformation.rs} | 135 ++++---- .../core/src/llms/azure_ai/anthropic/mod.rs | 1 + .../crates/core/src/llms/azure_ai/mod.rs | 1 + .../ocr/cohere_parse_transformation.rs | 6 +- .../document_intelligence/transformation.rs | 288 +++++++++-------- .../src/llms/azure_ai/ocr/transformation.rs | 143 +++++---- .../base_llm/anthropic_messages}/mod.rs | 0 .../anthropic_messages}/transformation.rs | 38 +-- .../base_llm/audio_transcription}/mod.rs | 0 .../audio_transcription/transformation.rs | 50 +-- .../responses => llms/base_llm/chat}/mod.rs | 0 .../base_llm/chat}/transformation.rs | 52 +-- .../crates/core/src/llms/base_llm/mod.rs | 3 + .../bedrock/audio_transcription/mod.rs} | 28 +- .../bedrock/chat/converse_transformation.rs} | 266 ++++++++-------- .../crates/core/src/llms/bedrock/chat/mod.rs | 1 + .../bedrock/chat}/tests.rs | 12 +- .../crates/core/src/llms/bedrock/mod.rs | 2 + .../src/llms/cohere/ocr/transformation.rs | 298 +++++++++--------- litellm-rust/crates/core/src/llms/mod.rs | 7 +- .../src/{providers => llms}/openai/mod.rs | 0 .../core/src/llms/openai/responses/mod.rs | 1 + .../openai/responses/transformation.rs | 6 +- .../src/llms/reducto/ocr/transformation.rs | 131 ++++---- .../vertex_ai/ocr/deepseek_transformation.rs | 8 +- .../src/llms/vertex_ai/ocr/transformation.rs | 110 ++++--- .../crates/core/src/messages/common_utils.rs | 8 +- .../crates/core/src/messages/handler.rs | 4 +- litellm-rust/crates/core/src/messages/mod.rs | 1 - .../crates/core/src/messages/prepare.rs | 14 +- .../crates/core/src/messages/types.rs | 4 +- .../crates/core/src/ocr/provider_config.rs | 16 +- .../core/src/providers/anthropic/mod.rs | 2 - .../core/src/providers/bedrock/aws_base.rs | 1 - .../core/src/providers/bedrock/constants.rs | 1 - .../crates/core/src/providers/bedrock/mod.rs | 8 - litellm-rust/crates/core/src/providers/mod.rs | 5 - .../crates/core/tests/vertex_ai_ocr.rs | 6 +- 59 files changed, 1002 insertions(+), 973 deletions(-) rename litellm-rust/crates/core/src/{providers/custom_llm_provider.rs => litellm_core_utils/get_llm_provider_logic.rs} (100%) create mode 100644 litellm-rust/crates/core/src/litellm_core_utils/mod.rs rename litellm-rust/crates/core/src/{providers/anthropic/chat_completions => llms/anthropic/chat}/mod.rs (100%) rename litellm-rust/crates/core/src/{providers/anthropic/chat_completions => llms/anthropic/chat}/tests.rs (99%) rename litellm-rust/crates/core/src/{providers/anthropic/chat_completions => llms/anthropic/chat}/transformation.rs (90%) rename litellm-rust/crates/core/src/{providers/anthropic => llms/anthropic/experimental_pass_through}/messages/mod.rs (100%) rename litellm-rust/crates/core/src/{providers/anthropic => llms/anthropic/experimental_pass_through}/messages/transformation.rs (93%) rename litellm-rust/crates/core/src/{providers/azure_ai => llms/anthropic/experimental_pass_through}/mod.rs (100%) create mode 100644 litellm-rust/crates/core/src/llms/anthropic/mod.rs rename litellm-rust/crates/core/src/{providers/azure_ai/messages/transformation.rs => llms/azure_ai/anthropic/messages_transformation.rs} (93%) create mode 100644 litellm-rust/crates/core/src/llms/azure_ai/anthropic/mod.rs rename litellm-rust/crates/core/src/{providers/azure_ai/messages => llms/base_llm/anthropic_messages}/mod.rs (100%) rename litellm-rust/crates/core/src/{messages => llms/base_llm/anthropic_messages}/transformation.rs (82%) rename litellm-rust/crates/core/src/{providers/bedrock/chat_completions => llms/base_llm/audio_transcription}/mod.rs (100%) rename litellm-rust/crates/core/src/{ => llms/base_llm}/audio_transcription/transformation.rs (61%) rename litellm-rust/crates/core/src/{providers/openai/responses => llms/base_llm/chat}/mod.rs (100%) rename litellm-rust/crates/core/src/{chat_completions => llms/base_llm/chat}/transformation.rs (95%) rename litellm-rust/crates/core/src/{providers/bedrock/audio_transcription.rs => llms/bedrock/audio_transcription/mod.rs} (91%) rename litellm-rust/crates/core/src/{providers/bedrock/chat_completions/transformation.rs => llms/bedrock/chat/converse_transformation.rs} (93%) create mode 100644 litellm-rust/crates/core/src/llms/bedrock/chat/mod.rs rename litellm-rust/crates/core/src/{providers/bedrock/chat_completions => llms/bedrock/chat}/tests.rs (98%) create mode 100644 litellm-rust/crates/core/src/llms/bedrock/mod.rs rename litellm-rust/crates/core/src/{providers => llms}/openai/mod.rs (100%) create mode 100644 litellm-rust/crates/core/src/llms/openai/responses/mod.rs rename litellm-rust/crates/core/src/{providers => llms}/openai/responses/transformation.rs (86%) delete mode 100644 litellm-rust/crates/core/src/providers/anthropic/mod.rs delete mode 100644 litellm-rust/crates/core/src/providers/bedrock/aws_base.rs delete mode 100644 litellm-rust/crates/core/src/providers/bedrock/constants.rs delete mode 100644 litellm-rust/crates/core/src/providers/bedrock/mod.rs delete mode 100644 litellm-rust/crates/core/src/providers/mod.rs diff --git a/litellm-rust/crates/core/AGENTS.md b/litellm-rust/crates/core/AGENTS.md index d591d241512..541b3b7e3d5 100644 --- a/litellm-rust/crates/core/AGENTS.md +++ b/litellm-rust/crates/core/AGENTS.md @@ -1,6 +1,6 @@ litellm-core is the LiteLLM SDK in Rust — it makes the LLM call. Each top-level call is a module under `src//` exposing a public entrypoint named after the route (`messages::messages()`, the Rust equivalent of `litellm.messages()`): you call it and get a typed non-streaming response back. -A route module owns the call entrypoint, runtime types, provider/auth/URL resolution, and the handler that performs the HTTP call. OCR provider transforms and the base provider trait live under `src/llms/`, mirroring their Python source paths. Other routes still use `src/providers/` and route-local traits. Handlers belong in core, never in a host crate +A route module owns the call entrypoint, runtime types, provider/auth/URL resolution, and the handler that performs the HTTP call. Provider code and base config traits live under `src/llms/`, mirroring their Python source paths. This applies to every API surface: shared orchestration stays in its route module (`ocr/`, `chat_completions/`, `messages/`, `audio_transcription/`, or `responses/`), while provider transformations live under the corresponding Python-mirrored `llms//` path. Import implementations directly from their canonical paths; do not add a `src/providers/` layer or compatibility re-exports. Shared provider resolution lives under `src/litellm_core_utils/get_llm_provider_logic.rs`. Handlers belong in core, never in a host crate Not here: serving HTTP (axum routes, extractors), config file reading, rollout state, databases, or host-specific callback execution. Core owns lifecycle sequencing and callback payload construction; hosts execute the selected integrations. Env reads are limited to credential fallback in a route's `prepare.rs`. @@ -21,3 +21,7 @@ Use named `#[rstest]` cases for independent input/output scenarios instead of lo For base OCR, Python response models correspond to `src/ocr/types.rs`; Rust context/environment types support the runtime. `BaseOcrConfig::prepare_request` corresponds to Python's HTTP-handler preparation rather than a `BaseOCRConfig` method, and `validate_request_body` is a Rust-only hook For Mistral, `async_transform_ocr_request` uses the base default in both languages. `resolve_headers` and `build_ocr_url` implement the respective environment and URL operations, and `normalize_response` implements the typed part of response transformation. Existing auth key/header handling and top-level response-extra preservation differ between languages; layout refactors must preserve those behaviors and verify them with the existing tests + +For non-OCR pairs, order corresponding methods as parameter support/mapping, environment validation, URL construction, request transformation, and response transformation, followed by Rust-only runtime hooks. Auth resolution remains split between configs and route preparation. Chat `supported_openai_param_mappings` describes accepted OpenAI/provider name pairs, unlike Python's `get_supported_openai_params` name list. Audio `map_transcription_params` remains a Rust filtering helper + +Azure Messages maps to `llms/azure_ai/anthropic/messages_transformation.py`; Bedrock Converse maps to `llms/bedrock/chat/converse_transformation.py`. `AnthropicConfig`, `AmazonConverseConfig`, and the non-OCR base traits are partial ports. `OpenAiResponsesApiConfig` currently implements only the WebSocket surface. Preserve their acceptance gates, passthrough behavior, and host fallback contracts when aligning layout diff --git a/litellm-rust/crates/core/src/audio_transcription/handler.rs b/litellm-rust/crates/core/src/audio_transcription/handler.rs index 2a7afccf9ea..4c48b6b5ede 100644 --- a/litellm-rust/crates/core/src/audio_transcription/handler.rs +++ b/litellm-rust/crates/core/src/audio_transcription/handler.rs @@ -36,7 +36,7 @@ pub async fn execute_audio_transcription_provider_call( .map_err(|error| Error::InvalidResponse(format!("invalid audio response JSON: {error}")))?; Ok(request .config - .transform_transcription_response(&request.model, response_json)? + .transform_audio_transcription_response(&request.model, response_json)? .into_json()) } @@ -47,9 +47,8 @@ async fn signed_headers( use std::collections::BTreeMap; use std::time::SystemTime; - use crate::audio_transcription::transformation::AudioTranscriptionAuth; - use crate::providers::bedrock::audio_transcription::aws_auth_config; - use crate::providers::bedrock::aws_base::{resolve_credentials, sign_bedrock_post}; + use crate::llms::base_llm::audio_transcription::transformation::AudioTranscriptionAuth; + use litellm_auth_aws::{aws_auth_config, resolve_credentials, sign_bedrock_post}; let AudioTranscriptionAuth::AwsSigV4 { region, .. } = &request.auth else { return Ok(request.upstream_headers.clone()); diff --git a/litellm-rust/crates/core/src/audio_transcription/mod.rs b/litellm-rust/crates/core/src/audio_transcription/mod.rs index 47b1e8bb151..fafc29a2d2a 100644 --- a/litellm-rust/crates/core/src/audio_transcription/mod.rs +++ b/litellm-rust/crates/core/src/audio_transcription/mod.rs @@ -3,7 +3,6 @@ pub use error::Error; mod client; mod handler; mod prepare; -pub mod transformation; pub mod types; pub use handler::execute_audio_transcription_provider_call; diff --git a/litellm-rust/crates/core/src/audio_transcription/prepare.rs b/litellm-rust/crates/core/src/audio_transcription/prepare.rs index 416ada2491e..26e705408e0 100644 --- a/litellm-rust/crates/core/src/audio_transcription/prepare.rs +++ b/litellm-rust/crates/core/src/audio_transcription/prepare.rs @@ -1,11 +1,15 @@ use super::Error; -use super::transformation::{AudioTranscriptionAuth, AudioTranscriptionProviderConfig}; use super::types::{AudioTranscriptionRequest, ProviderAudioTranscriptionRequest}; use crate::http_utils::{has_header, string_headers}; -use crate::providers::bedrock::audio_transcription::BEDROCK_AUDIO_TRANSCRIPTION_CONFIG; -use crate::providers::custom_llm_provider::{CustomLlmProvider, get_custom_llm_provider}; +use crate::litellm_core_utils::get_llm_provider_logic::{ + CustomLlmProvider, get_custom_llm_provider, +}; +use crate::llms::base_llm::audio_transcription::transformation::{ + AudioTranscriptionAuth, BaseAudioTranscriptionConfig, +}; +use crate::llms::bedrock::audio_transcription::BEDROCK_AUDIO_TRANSCRIPTION_CONFIG; -fn provider_config(provider: &str) -> Option<&'static dyn AudioTranscriptionProviderConfig> { +fn provider_config(provider: &str) -> Option<&'static dyn BaseAudioTranscriptionConfig> { if provider == "bedrock" { return Some(&BEDROCK_AUDIO_TRANSCRIPTION_CONFIG); } @@ -45,7 +49,7 @@ pub fn prepare_audio_transcription_provider_call( if !has_header(&headers, "content-type") { headers.push(("Content-Type".to_string(), "application/json".to_string())); } - let url = config.complete_url( + let url = config.get_complete_url( request.api_base, &model, &request.optional_params, @@ -53,7 +57,7 @@ pub fn prepare_audio_transcription_provider_call( )?; let filtered_params = config.map_transcription_params(&request.optional_params); let transformed = - config.transform_transcription_request(&model, request.audio, filtered_params)?; + config.transform_audio_transcription_request(&model, request.audio, filtered_params)?; Ok(ProviderAudioTranscriptionRequest { model, custom_llm_provider: provider_info.custom_llm_provider.to_string(), diff --git a/litellm-rust/crates/core/src/audio_transcription/types.rs b/litellm-rust/crates/core/src/audio_transcription/types.rs index 1f90f61c0da..1ec1f224f6b 100644 --- a/litellm-rust/crates/core/src/audio_transcription/types.rs +++ b/litellm-rust/crates/core/src/audio_transcription/types.rs @@ -3,7 +3,9 @@ use std::time::Duration; use serde::{Deserialize, Serialize}; use serde_json::{Map, Value}; -use super::transformation::{AudioTranscriptionAuth, AudioTranscriptionProviderConfig}; +use crate::llms::base_llm::audio_transcription::transformation::{ + AudioTranscriptionAuth, BaseAudioTranscriptionConfig, +}; pub struct AudioTranscriptionRequest<'a> { pub model: &'a str, @@ -20,7 +22,7 @@ pub struct AudioTranscriptionRequest<'a> { pub struct ProviderAudioTranscriptionRequest { pub(super) model: String, pub(super) custom_llm_provider: String, - pub(super) config: &'static dyn AudioTranscriptionProviderConfig, + pub(super) config: &'static dyn BaseAudioTranscriptionConfig, pub(super) url: String, pub(super) body: Value, pub(super) upstream_headers: Vec<(String, String)>, diff --git a/litellm-rust/crates/core/src/chat_completions/common_utils.rs b/litellm-rust/crates/core/src/chat_completions/common_utils.rs index c89450aeb77..8b966c7a173 100644 --- a/litellm-rust/crates/core/src/chat_completions/common_utils.rs +++ b/litellm-rust/crates/core/src/chat_completions/common_utils.rs @@ -1,19 +1,17 @@ use serde_json::{Map, Value}; use super::Error; -use super::transformation::ChatCompletionsProviderConfig; use crate::http_utils::string_headers as shared_string_headers; -use crate::providers::anthropic::chat_completions::transformation::ANTHROPIC_CHAT_COMPLETIONS_CONFIG; +use crate::llms::anthropic::chat::transformation::ANTHROPIC_CHAT_COMPLETIONS_CONFIG; +use crate::llms::base_llm::chat::transformation::BaseConfig; const HEADER_CONTEXT: &str = "chat completions"; -pub(super) fn chat_completions_provider_config( - provider: &str, -) -> Option<&'static dyn ChatCompletionsProviderConfig> { +pub(super) fn chat_completions_provider_config(provider: &str) -> Option<&'static dyn BaseConfig> { match provider { "anthropic" => Some(&ANTHROPIC_CHAT_COMPLETIONS_CONFIG), "bedrock" => Some( - &crate::providers::bedrock::chat_completions::transformation::BEDROCK_CHAT_COMPLETIONS_CONFIG, + &crate::llms::bedrock::chat::converse_transformation::BEDROCK_CHAT_COMPLETIONS_CONFIG, ), _ => None, } diff --git a/litellm-rust/crates/core/src/chat_completions/handler.rs b/litellm-rust/crates/core/src/chat_completions/handler.rs index 2d192e971b0..5090d481f6f 100644 --- a/litellm-rust/crates/core/src/chat_completions/handler.rs +++ b/litellm-rust/crates/core/src/chat_completions/handler.rs @@ -3,12 +3,12 @@ use serde_json::Value; use super::Error; use super::client::http_client; use super::prepare::prepare_provider_request; -use super::transformation::ChatCompletionsAuth; use super::types::{ ChatCompletionsResponse, ProviderChatCompletionsRequest, ProviderChatResponseData, ResolvedChatCompletionsRequest, }; use crate::http_utils::{http_request, truncate_error_body}; +use crate::llms::base_llm::chat::transformation::ChatCompletionsAuth; pub(super) async fn execute_chat_completions_provider_call( request: ResolvedChatCompletionsRequest<'_>, @@ -86,7 +86,7 @@ pub(super) async fn signed_headers( use std::collections::BTreeMap; use std::time::SystemTime; - use crate::providers::bedrock::aws_base::{ + use litellm_auth_aws::{ aws_auth_config, aws_signature_headers, host_supplied_credentials, is_sigv4_computed_header, resolve_credentials, sign_bedrock_post, }; diff --git a/litellm-rust/crates/core/src/chat_completions/mod.rs b/litellm-rust/crates/core/src/chat_completions/mod.rs index b31ceaffb5c..b5c231eb42d 100644 --- a/litellm-rust/crates/core/src/chat_completions/mod.rs +++ b/litellm-rust/crates/core/src/chat_completions/mod.rs @@ -14,7 +14,6 @@ pub mod conversation; pub(crate) mod handler; mod prepare; pub mod response_utils; -pub mod transformation; pub mod types; use handler::execute_chat_completions_provider_call; diff --git a/litellm-rust/crates/core/src/chat_completions/prepare.rs b/litellm-rust/crates/core/src/chat_completions/prepare.rs index b2360021ef7..983fbdf4f1d 100644 --- a/litellm-rust/crates/core/src/chat_completions/prepare.rs +++ b/litellm-rust/crates/core/src/chat_completions/prepare.rs @@ -2,18 +2,20 @@ use serde_json::Value; use super::Error; use super::common_utils::{chat_completions_provider_config, string_headers}; -use super::transformation::{ChatCompletionsAuth, ChatCompletionsProviderConfig}; use super::types::{ ChatCompletionsRequest, ChatMessage, ProviderChatCompletionsRequest, ResolvedChatCompletionsRequest, }; use crate::http_utils::has_header; -use crate::providers::custom_llm_provider::{CustomLlmProvider, get_custom_llm_provider}; +use crate::litellm_core_utils::get_llm_provider_logic::{ + CustomLlmProvider, get_custom_llm_provider, +}; +use crate::llms::base_llm::chat::transformation::{BaseConfig, ChatCompletionsAuth}; pub(super) fn resolve_provider_config<'a>( model: &'a str, custom_llm_provider: Option<&'a str>, -) -> Result<(String, &'static dyn ChatCompletionsProviderConfig), Error> { +) -> Result<(String, &'static dyn BaseConfig), Error> { let provider_info = get_custom_llm_provider(model, custom_llm_provider) .or_else(|| { custom_llm_provider.map(|provider| CustomLlmProvider { @@ -64,7 +66,7 @@ pub(super) fn resolve_request( fn validate_environment( request: &ResolvedChatCompletionsRequest<'_>, model: &str, - config: &dyn ChatCompletionsProviderConfig, + config: &dyn BaseConfig, ) -> Result<(Vec<(String, String)>, ChatCompletionsAuth), Error> { let env_lookup = |key: &str| std::env::var(key).ok(); let mut headers = string_headers(request.extra_headers.clone())?; @@ -121,7 +123,7 @@ pub(super) fn prepare_provider_request( let model = request.model; let config = request.config; let env_lookup = |key: &str| std::env::var(key).ok(); - let url = config.complete_url( + let url = config.get_complete_url( request.api_base, &model, &request.optional_params, diff --git a/litellm-rust/crates/core/src/chat_completions/tests.rs b/litellm-rust/crates/core/src/chat_completions/tests.rs index b860b5f7206..86ac6c6ca35 100644 --- a/litellm-rust/crates/core/src/chat_completions/tests.rs +++ b/litellm-rust/crates/core/src/chat_completions/tests.rs @@ -2,8 +2,8 @@ use serde_json::{Map, Value, json}; use super::Error; use super::prepare::{prepare_provider_request, resolve_request}; -use super::transformation::ChatCompletionsAuth; use super::types::{ChatCompletionsRequest, ProviderChatCompletionsRequest}; +use crate::llms::base_llm::chat::transformation::ChatCompletionsAuth; fn prepare_chat_completions_call( request: ChatCompletionsRequest<'_>, diff --git a/litellm-rust/crates/core/src/chat_completions/types.rs b/litellm-rust/crates/core/src/chat_completions/types.rs index 7178d594870..6e6b3d7063d 100644 --- a/litellm-rust/crates/core/src/chat_completions/types.rs +++ b/litellm-rust/crates/core/src/chat_completions/types.rs @@ -3,7 +3,7 @@ use std::time::Duration; use serde::{Deserialize, Serialize}; use serde_json::{Map, Value}; -use super::transformation::{ChatCompletionsAuth, ChatCompletionsProviderConfig}; +use crate::llms::base_llm::chat::transformation::{BaseConfig, ChatCompletionsAuth}; /// A `/chat/completions` call as it crosses into the core. /// @@ -24,7 +24,7 @@ pub struct ChatCompletionsRequest<'a> { pub(super) struct ResolvedChatCompletionsRequest<'a> { pub(super) model: String, - pub(super) config: &'static dyn ChatCompletionsProviderConfig, + pub(super) config: &'static dyn BaseConfig, pub(super) messages: Vec, pub(super) optional_params: Map, pub(super) api_key: Option<&'a str>, @@ -35,7 +35,7 @@ pub(super) struct ResolvedChatCompletionsRequest<'a> { pub(super) struct ProviderChatCompletionsRequest { pub(super) model: String, - pub(super) config: &'static dyn ChatCompletionsProviderConfig, + pub(super) config: &'static dyn BaseConfig, pub(super) url: String, pub(super) body: Value, pub(super) upstream_headers: Vec<(String, String)>, diff --git a/litellm-rust/crates/core/src/lib.rs b/litellm-rust/crates/core/src/lib.rs index 288bde52ce4..6d540ceaa6f 100644 --- a/litellm-rust/crates/core/src/lib.rs +++ b/litellm-rust/crates/core/src/lib.rs @@ -5,12 +5,12 @@ pub mod chat_completions; pub mod constants; pub mod error; pub mod http_utils; -pub(crate) mod llms; +pub mod litellm_core_utils; +pub mod llms; mod media; pub mod messages; pub mod ocr; pub mod params; -pub mod providers; pub mod responses; mod serde_compat; pub mod transport; diff --git a/litellm-rust/crates/core/src/providers/custom_llm_provider.rs b/litellm-rust/crates/core/src/litellm_core_utils/get_llm_provider_logic.rs similarity index 100% rename from litellm-rust/crates/core/src/providers/custom_llm_provider.rs rename to litellm-rust/crates/core/src/litellm_core_utils/get_llm_provider_logic.rs diff --git a/litellm-rust/crates/core/src/litellm_core_utils/mod.rs b/litellm-rust/crates/core/src/litellm_core_utils/mod.rs new file mode 100644 index 00000000000..7e3b3e96dda --- /dev/null +++ b/litellm-rust/crates/core/src/litellm_core_utils/mod.rs @@ -0,0 +1 @@ +pub mod get_llm_provider_logic; diff --git a/litellm-rust/crates/core/src/providers/anthropic/chat_completions/mod.rs b/litellm-rust/crates/core/src/llms/anthropic/chat/mod.rs similarity index 100% rename from litellm-rust/crates/core/src/providers/anthropic/chat_completions/mod.rs rename to litellm-rust/crates/core/src/llms/anthropic/chat/mod.rs diff --git a/litellm-rust/crates/core/src/providers/anthropic/chat_completions/tests.rs b/litellm-rust/crates/core/src/llms/anthropic/chat/tests.rs similarity index 99% rename from litellm-rust/crates/core/src/providers/anthropic/chat_completions/tests.rs rename to litellm-rust/crates/core/src/llms/anthropic/chat/tests.rs index 81bc8f02a66..25c2f5e49f4 100644 --- a/litellm-rust/crates/core/src/providers/anthropic/chat_completions/tests.rs +++ b/litellm-rust/crates/core/src/llms/anthropic/chat/tests.rs @@ -420,7 +420,7 @@ fn resolves_the_messages_url_and_x_api_key_auth() { let config = &ANTHROPIC_CHAT_COMPLETIONS_CONFIG; assert_eq!( config - .complete_url(None, "claude-sonnet-4-5", &Map::new(), &|_| None) + .get_complete_url(None, "claude-sonnet-4-5", &Map::new(), &|_| None) .expect("url builds"), "https://api.anthropic.com/v1/messages" ); diff --git a/litellm-rust/crates/core/src/providers/anthropic/chat_completions/transformation.rs b/litellm-rust/crates/core/src/llms/anthropic/chat/transformation.rs similarity index 90% rename from litellm-rust/crates/core/src/providers/anthropic/chat_completions/transformation.rs rename to litellm-rust/crates/core/src/llms/anthropic/chat/transformation.rs index dd0830edab7..fc48ef6d74f 100644 --- a/litellm-rust/crates/core/src/providers/anthropic/chat_completions/transformation.rs +++ b/litellm-rust/crates/core/src/llms/anthropic/chat/transformation.rs @@ -3,18 +3,17 @@ use serde_json::{Map, Value, json}; use crate::chat_completions::Error; use crate::chat_completions::conversation::{Conversation, build_conversation}; use crate::chat_completions::response_utils::{finish_reason_for, unix_now, usage_from_parts}; -use crate::chat_completions::transformation::{ - ChatCompletionsAuth, ChatCompletionsProviderConfig, Unsupported, unsupported_message, - unsupported_param, -}; use crate::chat_completions::types::{ ChatCompletionsChoice, ChatCompletionsChoiceMessage, ChatCompletionsResponse, ChatMessage, ProviderChatRequestData, ProviderChatResponseData, }; use crate::constants::ANTHROPIC_OAUTH_TOKEN_PREFIX; -use crate::providers::anthropic::messages::transformation::{ +use crate::llms::anthropic::experimental_pass_through::messages::transformation::{ complete_anthropic_url, resolve_anthropic_api_key, }; +use crate::llms::base_llm::chat::transformation::{ + BaseConfig, ChatCompletionsAuth, Unsupported, unsupported_message, unsupported_param, +}; /// Anthropic parameter names, post `map_openai_params`, that the Rust path can /// place verbatim in the Messages body. @@ -33,46 +32,16 @@ const SUPPORTED_PARAMS: &[(&str, &str)] = &[ ("stop", "stop_sequences"), ]; -pub struct AnthropicChatCompletionsConfig; +pub struct AnthropicConfig; -pub const ANTHROPIC_CHAT_COMPLETIONS_CONFIG: AnthropicChatCompletionsConfig = - AnthropicChatCompletionsConfig; +pub const ANTHROPIC_CHAT_COMPLETIONS_CONFIG: AnthropicConfig = AnthropicConfig; -fn text_block(text: &str) -> Value { - json!({"type": "text", "text": text}) -} +impl BaseConfig for AnthropicConfig { + fn supported_openai_param_mappings(&self) -> &'static [(&'static str, &'static str)] { + SUPPORTED_PARAMS + } -fn anthropic_body(model: &str, conversation: &Conversation, params: Map) -> Value { - let messages: Vec = conversation - .turns - .iter() - .map(|turn| { - json!({ - "role": turn.role.as_str(), - "content": turn.texts.iter().map(|text| text_block(text)).collect::>(), - }) - }) - .collect(); - - let system: Vec = conversation.system.iter().map(|s| text_block(s)).collect(); - - let body = Map::from_iter( - [ - ("model".to_string(), json!(model)), - ("messages".to_string(), json!(messages)), - ] - .into_iter() - // Python builds `{"model", "messages", **optional_params}` with - // `system` already folded into optional_params, so a caller-supplied - // key of the same name wins here too. - .chain((!system.is_empty()).then(|| ("system".to_string(), json!(system)))) - .chain(params), - ); - Value::Object(body) -} - -impl ChatCompletionsProviderConfig for AnthropicChatCompletionsConfig { - fn complete_url( + fn get_complete_url( &self, api_base: Option<&str>, _model: &str, @@ -82,60 +51,6 @@ impl ChatCompletionsProviderConfig for AnthropicChatCompletionsConfig { Ok(complete_anthropic_url(api_base, env_lookup)) } - fn auth( - &self, - api_key: Option<&str>, - _model: &str, - _optional_params: &Map, - env_lookup: &dyn Fn(&str) -> Option, - ) -> Result { - Ok(ChatCompletionsAuth::Header { - name: "x-api-key", - value: resolve_anthropic_api_key(api_key, env_lookup)?, - }) - } - - fn default_headers(&self) -> &'static [(&'static str, &'static str)] { - &[ - ("anthropic-version", "2023-06-01"), - ("content-type", "application/json"), - ] - } - - /// An OAuth bearer is the whole credential: Python's `validate_environment` - /// authenticates with it and drops `x-api-key` rather than resolving one, so - /// the resolved key must not be applied over the top. Any other forwarded - /// `authorization` is unrelated to this header and does not defer, which is - /// also what Python does: it sends the deployment's `x-api-key` alongside. - fn defers_to_forwarded_auth(&self, headers: &[(String, String)]) -> bool { - headers.iter().any(|(name, value)| { - name.eq_ignore_ascii_case("authorization") - && value - .strip_prefix("Bearer ") - .is_some_and(|token| token.starts_with(ANTHROPIC_OAUTH_TOKEN_PREFIX)) - }) - } - - fn supported_openai_params(&self) -> &'static [(&'static str, &'static str)] { - SUPPORTED_PARAMS - } - - fn unsupported_reason( - &self, - messages: &[ChatMessage], - optional_params: &Map, - ) -> Option { - unsupported_param(self.supported_openai_params(), &[], optional_params) - .or_else(|| messages.iter().find_map(unsupported_message)) - // Anthropic rejects a request whose first turn is not a user turn. - // Python only repairs that under `litellm.modify_params`, which the - // core cannot observe, so decline instead of guessing. - .or_else(|| { - (!build_conversation(messages).opens_on_user_turn()) - .then_some(Unsupported("conversation does not open on a user turn")) - }) - } - fn transform_request( &self, model: &str, @@ -209,6 +124,93 @@ impl ChatCompletionsProviderConfig for AnthropicChatCompletionsConfig { ), }) } + + fn auth( + &self, + api_key: Option<&str>, + _model: &str, + _optional_params: &Map, + env_lookup: &dyn Fn(&str) -> Option, + ) -> Result { + Ok(ChatCompletionsAuth::Header { + name: "x-api-key", + value: resolve_anthropic_api_key(api_key, env_lookup)?, + }) + } + + fn default_headers(&self) -> &'static [(&'static str, &'static str)] { + &[ + ("anthropic-version", "2023-06-01"), + ("content-type", "application/json"), + ] + } + + /// An OAuth bearer is the whole credential: Python's `validate_environment` + /// authenticates with it and drops `x-api-key` rather than resolving one, so + /// the resolved key must not be applied over the top. Any other forwarded + /// `authorization` is unrelated to this header and does not defer, which is + /// also what Python does: it sends the deployment's `x-api-key` alongside. + fn defers_to_forwarded_auth(&self, headers: &[(String, String)]) -> bool { + headers.iter().any(|(name, value)| { + name.eq_ignore_ascii_case("authorization") + && value + .strip_prefix("Bearer ") + .is_some_and(|token| token.starts_with(ANTHROPIC_OAUTH_TOKEN_PREFIX)) + }) + } + + fn unsupported_reason( + &self, + messages: &[ChatMessage], + optional_params: &Map, + ) -> Option { + unsupported_param(self.supported_openai_param_mappings(), &[], optional_params) + .or_else(|| messages.iter().find_map(unsupported_message)) + // Anthropic rejects a request whose first turn is not a user turn. + // Python only repairs that under `litellm.modify_params`, which the + // core cannot observe, so decline instead of guessing. + .or_else(|| { + (!build_conversation(messages).opens_on_user_turn()) + .then_some(Unsupported("conversation does not open on a user turn")) + }) + } +} + +fn text_block(text: &str) -> Value { + json!({"type": "text", "text": text}) +} + +fn anthropic_body( + model: &str, + conversation: &Conversation, + optional_params: Map, +) -> Value { + let messages: Vec = conversation + .turns + .iter() + .map(|turn| { + json!({ + "role": turn.role.as_str(), + "content": turn.texts.iter().map(|text| text_block(text)).collect::>(), + }) + }) + .collect(); + + let system: Vec = conversation.system.iter().map(|s| text_block(s)).collect(); + + let body = Map::from_iter( + [ + ("model".to_string(), json!(model)), + ("messages".to_string(), json!(messages)), + ] + .into_iter() + // Python builds `{"model", "messages", **optional_params}` with + // `system` already folded into optional_params, so a caller-supplied + // key of the same name wins here too. + .chain((!system.is_empty()).then(|| ("system".to_string(), json!(system)))) + .chain(optional_params), + ); + Value::Object(body) } #[cfg(test)] diff --git a/litellm-rust/crates/core/src/providers/anthropic/messages/mod.rs b/litellm-rust/crates/core/src/llms/anthropic/experimental_pass_through/messages/mod.rs similarity index 100% rename from litellm-rust/crates/core/src/providers/anthropic/messages/mod.rs rename to litellm-rust/crates/core/src/llms/anthropic/experimental_pass_through/messages/mod.rs diff --git a/litellm-rust/crates/core/src/providers/anthropic/messages/transformation.rs b/litellm-rust/crates/core/src/llms/anthropic/experimental_pass_through/messages/transformation.rs similarity index 93% rename from litellm-rust/crates/core/src/providers/anthropic/messages/transformation.rs rename to litellm-rust/crates/core/src/llms/anthropic/experimental_pass_through/messages/transformation.rs index 080f11c8cac..a4dc7d2aaa3 100644 --- a/litellm-rust/crates/core/src/providers/anthropic/messages/transformation.rs +++ b/litellm-rust/crates/core/src/llms/anthropic/experimental_pass_through/messages/transformation.rs @@ -1,5 +1,5 @@ +use crate::llms::base_llm::anthropic_messages::transformation::BaseAnthropicMessagesConfig; use crate::messages::Error; -use crate::messages::transformation::{AnthropicMessagesProviderConfig, MessagesAuthStrategy}; const ANTHROPIC_API_KEY_ENV: &str = "ANTHROPIC_API_KEY"; const ANTHROPIC_API_BASE_ENV: &str = "ANTHROPIC_API_BASE"; @@ -10,6 +10,25 @@ pub struct AnthropicMessagesConfig; pub const ANTHROPIC_MESSAGES_CONFIG: AnthropicMessagesConfig = AnthropicMessagesConfig; +impl BaseAnthropicMessagesConfig for AnthropicMessagesConfig { + fn get_complete_url( + &self, + api_base: Option<&str>, + _model: &str, + env_lookup: &dyn Fn(&str) -> Option, + ) -> Result { + Ok(complete_anthropic_url(api_base, env_lookup)) + } + + fn resolve_api_key( + &self, + api_key: Option<&str>, + env_lookup: &dyn Fn(&str) -> Option, + ) -> Result { + resolve_anthropic_api_key(api_key, env_lookup).map_err(Error::from) + } +} + pub fn non_empty(value: Option<&str>) -> Option<&str> { value.map(str::trim).filter(|value| !value.is_empty()) } @@ -43,29 +62,6 @@ pub fn complete_anthropic_url( format!("{api_base}{MESSAGES_PATH_SUFFIX}") } -impl AnthropicMessagesProviderConfig for AnthropicMessagesConfig { - fn complete_url( - &self, - api_base: Option<&str>, - _model: &str, - env_lookup: &dyn Fn(&str) -> Option, - ) -> Result { - Ok(complete_anthropic_url(api_base, env_lookup)) - } - - fn resolve_api_key( - &self, - api_key: Option<&str>, - env_lookup: &dyn Fn(&str) -> Option, - ) -> Result { - resolve_anthropic_api_key(api_key, env_lookup).map_err(Error::from) - } - - fn auth_strategy(&self) -> MessagesAuthStrategy { - MessagesAuthStrategy::Header("x-api-key") - } -} - #[cfg(test)] mod tests { use super::*; diff --git a/litellm-rust/crates/core/src/providers/azure_ai/mod.rs b/litellm-rust/crates/core/src/llms/anthropic/experimental_pass_through/mod.rs similarity index 100% rename from litellm-rust/crates/core/src/providers/azure_ai/mod.rs rename to litellm-rust/crates/core/src/llms/anthropic/experimental_pass_through/mod.rs diff --git a/litellm-rust/crates/core/src/llms/anthropic/mod.rs b/litellm-rust/crates/core/src/llms/anthropic/mod.rs new file mode 100644 index 00000000000..4943d80a45c --- /dev/null +++ b/litellm-rust/crates/core/src/llms/anthropic/mod.rs @@ -0,0 +1,2 @@ +pub mod chat; +pub mod experimental_pass_through; diff --git a/litellm-rust/crates/core/src/providers/azure_ai/messages/transformation.rs b/litellm-rust/crates/core/src/llms/azure_ai/anthropic/messages_transformation.rs similarity index 93% rename from litellm-rust/crates/core/src/providers/azure_ai/messages/transformation.rs rename to litellm-rust/crates/core/src/llms/azure_ai/anthropic/messages_transformation.rs index 1929f86a1d6..feaee0375c4 100644 --- a/litellm-rust/crates/core/src/providers/azure_ai/messages/transformation.rs +++ b/litellm-rust/crates/core/src/llms/azure_ai/anthropic/messages_transformation.rs @@ -1,14 +1,16 @@ use serde_json::{Map, Value}; +use crate::llms::anthropic::experimental_pass_through::messages::transformation::{ + ANTHROPIC_MESSAGES_CONFIG, AnthropicMessagesConfig, non_empty, +}; +use crate::llms::base_llm::anthropic_messages::transformation::{ + BaseAnthropicMessagesConfig, MessagesAuthStrategy, +}; use crate::messages::Error; -use crate::messages::transformation::{AnthropicMessagesProviderConfig, MessagesAuthStrategy}; use crate::messages::types::{ AnthropicMessage, AnthropicMessagesRequest, AnthropicMessagesResponse, ContentBlock, MessageContent, SystemPrompt, }; -use crate::providers::anthropic::messages::transformation::{ - ANTHROPIC_MESSAGES_CONFIG, AnthropicMessagesConfig, non_empty, -}; const AZURE_API_KEY_ENV: &str = "AZURE_API_KEY"; const AZURE_API_BASE_ENV: &str = "AZURE_API_BASE"; @@ -26,6 +28,61 @@ pub const AZURE_ANTHROPIC_MESSAGES_CONFIG: AzureAnthropicMessagesConfig = anthropic: ANTHROPIC_MESSAGES_CONFIG, }; +impl BaseAnthropicMessagesConfig for AzureAnthropicMessagesConfig { + fn get_complete_url( + &self, + api_base: Option<&str>, + _model: &str, + env_lookup: &dyn Fn(&str) -> Option, + ) -> Result { + complete_azure_anthropic_url(api_base, env_lookup) + } + + fn transform_anthropic_messages_request( + &self, + request: AnthropicMessagesRequest, + ) -> Result { + let mut request = fold_system_role_messages(request); + if let Some(system) = request.system.as_mut() { + strip_scope_from_system(system); + } + request + .messages + .iter_mut() + .for_each(strip_scope_from_message); + self.anthropic.transform_anthropic_messages_request(request) + } + + fn transform_anthropic_messages_response( + &self, + model: &str, + response: AnthropicMessagesResponse, + ) -> Result { + self.anthropic + .transform_anthropic_messages_response(model, response) + } + + fn resolve_api_key( + &self, + api_key: Option<&str>, + env_lookup: &dyn Fn(&str) -> Option, + ) -> Result { + resolve_azure_api_key(api_key, env_lookup) + } + + fn auth_strategy(&self) -> MessagesAuthStrategy { + self.anthropic.auth_strategy() + } + + fn accepts_bearer_auth(&self) -> bool { + true + } + + fn default_headers(&self) -> &'static [(&'static str, &'static str)] { + self.anthropic.default_headers() + } +} + pub fn resolve_azure_api_key( api_key: Option<&str>, env_lookup: &dyn Fn(&str) -> Option, @@ -136,60 +193,6 @@ fn fold_system_role_messages(request: AnthropicMessagesRequest) -> AnthropicMess } } -impl AnthropicMessagesProviderConfig for AzureAnthropicMessagesConfig { - fn complete_url( - &self, - api_base: Option<&str>, - _model: &str, - env_lookup: &dyn Fn(&str) -> Option, - ) -> Result { - complete_azure_anthropic_url(api_base, env_lookup) - } - - fn resolve_api_key( - &self, - api_key: Option<&str>, - env_lookup: &dyn Fn(&str) -> Option, - ) -> Result { - resolve_azure_api_key(api_key, env_lookup) - } - - fn auth_strategy(&self) -> MessagesAuthStrategy { - self.anthropic.auth_strategy() - } - - fn accepts_bearer_auth(&self) -> bool { - true - } - - fn default_headers(&self) -> &'static [(&'static str, &'static str)] { - self.anthropic.default_headers() - } - - fn transform_request( - &self, - request: AnthropicMessagesRequest, - ) -> Result { - let mut request = fold_system_role_messages(request); - if let Some(system) = request.system.as_mut() { - strip_scope_from_system(system); - } - request - .messages - .iter_mut() - .for_each(strip_scope_from_message); - self.anthropic.transform_request(request) - } - - fn transform_response( - &self, - model: &str, - response: AnthropicMessagesResponse, - ) -> Result { - self.anthropic.transform_response(model, response) - } -} - #[cfg(test)] mod tests { use serde_json::json; @@ -339,7 +342,7 @@ mod tests { let transformed = to_value( AZURE_ANTHROPIC_MESSAGES_CONFIG - .transform_request(request) + .transform_anthropic_messages_request(request) .expect("request transforms"), ); @@ -366,10 +369,10 @@ mod tests { "messages": [{"role": "user", "content": "hi"}] })); let once = AZURE_ANTHROPIC_MESSAGES_CONFIG - .transform_request(request) + .transform_anthropic_messages_request(request) .expect("request transforms"); let twice = AZURE_ANTHROPIC_MESSAGES_CONFIG - .transform_request(once.clone()) + .transform_anthropic_messages_request(once.clone()) .expect("request transforms"); assert_eq!(once, twice); assert_eq!(to_value(once)["system"], json!("plain string system")); @@ -403,7 +406,7 @@ mod tests { }); let transformed = to_value( AZURE_ANTHROPIC_MESSAGES_CONFIG - .transform_request(request_from(body.clone())) + .transform_anthropic_messages_request(request_from(body.clone())) .expect("request transforms"), ); assert_eq!(transformed, body); @@ -423,7 +426,7 @@ mod tests { let transformed = to_value( AZURE_ANTHROPIC_MESSAGES_CONFIG - .transform_request(request) + .transform_anthropic_messages_request(request) .expect("request transforms"), ); @@ -453,7 +456,7 @@ mod tests { let transformed = to_value( AZURE_ANTHROPIC_MESSAGES_CONFIG - .transform_request(request) + .transform_anthropic_messages_request(request) .expect("request transforms"), ); @@ -480,7 +483,7 @@ mod tests { }); let transformed = to_value( AZURE_ANTHROPIC_MESSAGES_CONFIG - .transform_request(request_from(body.clone())) + .transform_anthropic_messages_request(request_from(body.clone())) .expect("request transforms"), ); assert_eq!(transformed, body); @@ -507,7 +510,7 @@ mod tests { })) .expect("valid response"); let transformed = AZURE_ANTHROPIC_MESSAGES_CONFIG - .transform_response("claude-sonnet-4-5", response) + .transform_anthropic_messages_response("claude-sonnet-4-5", response) .expect("response transforms"); let value = serde_json::to_value(transformed).expect("serializable"); assert_eq!(value["stop_reason"], json!("end_turn")); diff --git a/litellm-rust/crates/core/src/llms/azure_ai/anthropic/mod.rs b/litellm-rust/crates/core/src/llms/azure_ai/anthropic/mod.rs new file mode 100644 index 00000000000..eb8d16a4616 --- /dev/null +++ b/litellm-rust/crates/core/src/llms/azure_ai/anthropic/mod.rs @@ -0,0 +1 @@ +pub mod messages_transformation; diff --git a/litellm-rust/crates/core/src/llms/azure_ai/mod.rs b/litellm-rust/crates/core/src/llms/azure_ai/mod.rs index 079e0c41eae..8a52bda45be 100644 --- a/litellm-rust/crates/core/src/llms/azure_ai/mod.rs +++ b/litellm-rust/crates/core/src/llms/azure_ai/mod.rs @@ -1 +1,2 @@ +pub mod anthropic; pub(crate) mod ocr; 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 bdd18cbf4df..0b60c793c9d 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 @@ -18,7 +18,7 @@ impl BaseOcrConfig for AzureAICohereParseConfig { type Environment = Vec<(String, String)>; fn get_api_key_env_var(&self) -> Option<&'static str> { - super::transformation::AzureAIOCRConfig.get_api_key_env_var() + super::transformation::AzureAiOcrConfig.get_api_key_env_var() } fn get_health_check_document(&self) -> OcrDocument { @@ -31,7 +31,7 @@ impl BaseOcrConfig for AzureAICohereParseConfig { client: &OcrClient, ) -> Result { BaseOcrConfig::validate_environment( - &super::transformation::AzureAIOCRConfig, + &super::transformation::AzureAiOcrConfig, request, client, ) @@ -44,7 +44,7 @@ impl BaseOcrConfig for AzureAICohereParseConfig { _params: &Self::OcrParams, _environment: &Self::Environment, ) -> Result { - let base = super::transformation::AzureAIOCRConfig::resolve_api_base( + let base = super::transformation::AzureAiOcrConfig::resolve_api_base( request.connection.api_base.as_deref(), &crate::ocr::prepare::credential_env, )?; 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 1016ed02783..78841274f39 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 @@ -17,7 +17,7 @@ use crate::constants::{ AZURE_DI_SUBSCRIPTION_HEADER, OCR_POLL_RETRY_SECS, }; use crate::llms::base_llm::ocr::transformation::{ - BaseOcrConfig, OcrRequestContext, OcrResponseContext, + BaseOcrConfig, OcrResponseContext, decode_and_normalize_response, }; use crate::ocr::OcrClient; use crate::ocr::client::read_json_response; @@ -32,6 +32,9 @@ use crate::ocr::types::{ use crate::serde_compat::{FiniteF64, LaxI64}; use crate::url_utils::ApiUrl; +const AZURE_DI_API_KEY_ENV: &str = "AZURE_DOCUMENT_INTELLIGENCE_API_KEY"; +const AZURE_DI_ENDPOINT_ENV: &str = "AZURE_DOCUMENT_INTELLIGENCE_ENDPOINT"; + #[derive(Clone, Debug, PartialEq, Serialize)] pub(crate) struct DocumentIntelligenceParams { #[serde(skip_serializing_if = "Option::is_none")] @@ -123,7 +126,126 @@ struct AzureDocumentIntelligenceLine { pub content: Option, } -fn normalize_pages(pages: Option<&Value>) -> Result, crate::ocr::Error> { +#[derive(Clone, Debug)] +pub(crate) struct AzureDocumentIntelligenceOcrConfig; + +impl BaseOcrConfig for AzureDocumentIntelligenceOcrConfig { + type OcrParams = DocumentIntelligenceParams; + type ProviderRequest = DocumentIntelligenceRequest; + type Environment = Vec<(String, String)>; + + fn get_supported_ocr_params(&self, _model: &str) -> &'static [&'static str] { + &["pages", "features", "req_format"] + } + + fn get_api_key_env_var(&self) -> Option<&'static str> { + Some(AZURE_DI_API_KEY_ENV) + } + + fn resolve_connection_params(&self, inputs: OcrCredentialInputs) -> ResolvedOcrCredentials { + ResolvedOcrCredentials { + api_key: inputs.api_key.and_then(|key| { + inputs + .dynamic_api_key + .filter(|value| !value.value().is_empty()) + .or(Some(key)) + }), + api_base: inputs.api_base.and_then(|base| { + inputs + .dynamic_api_base + .filter(|value| !value.value().is_empty()) + .or(Some(base)) + }), + } + } + + fn map_ocr_params( + &self, + non_default_params: &CallArguments, + _model: &str, + ) -> Result { + Ok(DocumentIntelligenceParams { + pages: normalize_pages_param(non_default_params.get("pages"))?, + features: normalize_features_param(non_default_params.get("features"))?, + }) + } + + async fn validate_environment( + &self, + request: &PreparedOcrRequest, + _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.resolve_headers(&request.connection, &config, &credential_env) + .await + } + + fn get_complete_url( + &self, + request: &PreparedOcrRequest, + optional_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.build_ocr_url(&endpoint, &request.model, optional_params) + } + + fn transform_ocr_request( + &self, + _model: &str, + document: OcrDocument, + _optional_params: &DocumentIntelligenceParams, + _headers: &[(String, String)], + ) -> Result { + build_request(document) + } + + fn transform_ocr_response( + &self, + model: &str, + raw_response: &[u8], + request_format: OcrResponseFormat, + ) -> Result { + decode_and_normalize_response( + model, + raw_response, + request_format, + transform_completed_response, + ) + } + + async fn async_transform_ocr_response( + &self, + model: &str, + raw_response: reqwest::Response, + context: OcrResponseContext<'_>, + ) -> Result { + let decoded = read_operation_response( + context.client.polling_http(), + raw_response, + context.url, + context.headers, + context.connection, + context.request_format == OcrResponseFormat::Native, + context.hooks, + ) + .await?; + Ok(LiteLLMOcrResponse { + provider_native_response: decoded.native, + ..transform_completed_response(model, decoded.data)? + }) + } +} + +fn normalize_pages_param(pages: Option<&Value>) -> Result, crate::ocr::Error> { let normalized = match pages { None | Some(Value::Null) => return Ok(None), Some(Value::Array(pages)) if pages.is_empty() => return Ok(None), @@ -186,7 +308,7 @@ fn valid_page_token(token: &str) -> bool { } } -fn normalize_features(features: Option<&Value>) -> Result, crate::ocr::Error> { +fn normalize_features_param(features: Option<&Value>) -> Result, crate::ocr::Error> { let tokens = match features { None | Some(Value::Null) => return Ok(None), Some(Value::Array(names)) => names @@ -414,140 +536,8 @@ async fn poll_operation( } } -const AZURE_DI_API_KEY_ENV: &str = "AZURE_DOCUMENT_INTELLIGENCE_API_KEY"; -const AZURE_DI_ENDPOINT_ENV: &str = "AZURE_DOCUMENT_INTELLIGENCE_ENDPOINT"; - -#[derive(Clone, Debug)] -pub(crate) struct AzureDocumentIntelligenceOCRConfig; - -impl BaseOcrConfig for AzureDocumentIntelligenceOCRConfig { - type OcrParams = DocumentIntelligenceParams; - type ProviderRequest = DocumentIntelligenceRequest; - 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, inputs: OcrCredentialInputs) -> ResolvedOcrCredentials { - ResolvedOcrCredentials { - api_key: inputs.api_key.and_then(|key| { - inputs - .dynamic_api_key - .filter(|value| !value.value().is_empty()) - .or(Some(key)) - }), - api_base: inputs.api_base.and_then(|base| { - inputs - .dynamic_api_base - .filter(|value| !value.value().is_empty()) - .or(Some(base)) - }), - } - } - - async fn validate_environment( - &self, - request: &PreparedOcrRequest, - _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: &PreparedOcrRequest, - 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"] - } - - fn map_ocr_params( - &self, - arguments: &CallArguments, - _model: &str, - ) -> Result { - Ok(DocumentIntelligenceParams { - pages: normalize_pages(arguments.get("pages"))?, - features: normalize_features(arguments.get("features"))?, - }) - } - - async fn async_transform_ocr_request( - &self, - model: &str, - document: OcrDocument, - optional_params: &DocumentIntelligenceParams, - headers: &[(String, String)], - _context: OcrRequestContext<'_>, - ) -> Result { - self.transform_ocr_request(model, document, optional_params, headers) - } - - fn transform_ocr_response( - &self, - model: &str, - raw_response: &[u8], - request_format: OcrResponseFormat, - ) -> Result { - crate::llms::base_llm::ocr::transformation::decode_and_normalize_response( - model, - raw_response, - request_format, - transform_completed_response, - ) - } - - async fn async_transform_ocr_response( - &self, - model: &str, - raw_response: reqwest::Response, - context: OcrResponseContext<'_>, - ) -> Result { - let decoded = read_operation_response( - context.client.polling_http(), - raw_response, - context.url, - context.headers, - context.connection, - context.request_format == OcrResponseFormat::Native, - context.hooks, - ) - .await?; - Ok(LiteLLMOcrResponse { - provider_native_response: decoded.native, - ..transform_completed_response(model, decoded.data)? - }) - } - fn transform_ocr_request( - &self, - _model: &str, - document: OcrDocument, - _optional_params: &DocumentIntelligenceParams, - _headers: &[(String, String)], - ) -> Result { - build_request(document) - } -} - -impl AzureDocumentIntelligenceOCRConfig { - fn get_complete_url( +impl AzureDocumentIntelligenceOcrConfig { + fn build_ocr_url( &self, endpoint: &str, model: &str, @@ -575,7 +565,7 @@ impl AzureDocumentIntelligenceOCRConfig { }) } - async fn validate_environment( + async fn resolve_headers( &self, connection: &OcrConnection, config: &AzureAuthInputs, @@ -642,7 +632,7 @@ mod tests { fn map(value: Value) -> Result { let arguments = serde_json::from_value(value).unwrap(); - AzureDocumentIntelligenceOCRConfig.map_ocr_params(&arguments, "model") + AzureDocumentIntelligenceOcrConfig.map_ocr_params(&arguments, "model") } #[test] @@ -650,7 +640,7 @@ mod tests { let overrides = serde_json::from_value(json!({"pages":[], "features":null, "req_format":"native"})) .unwrap(); - let mapped = AzureDocumentIntelligenceOCRConfig + let mapped = AzureDocumentIntelligenceOcrConfig .map_ocr_params(&overrides, "model") .unwrap(); assert_eq!(serde_json::to_value(mapped).unwrap(), json!({})); @@ -664,7 +654,7 @@ mod tests { "extra_body": {"provider_option": "value"} })) .unwrap(); - let mapped = AzureDocumentIntelligenceOCRConfig + let mapped = AzureDocumentIntelligenceOcrConfig .map_ocr_params(&arguments, "model") .unwrap(); assert_eq!(mapped.pages.as_deref(), Some("1")); @@ -680,7 +670,7 @@ mod tests { "pages":"4", "features":"languages", "extension":true })) .unwrap(); - let mapped = AzureDocumentIntelligenceOCRConfig + let mapped = AzureDocumentIntelligenceOcrConfig .map_ocr_params(&arguments, "model") .unwrap(); assert_eq!( @@ -694,7 +684,7 @@ mod tests { #[test] fn response_numbers_follow_python_validation_before_dimension_conversion() { - let response = AzureDocumentIntelligenceOCRConfig.transform_ocr_response( + let response = AzureDocumentIntelligenceOcrConfig.transform_ocr_response( "model", br#"{"status":"succeeded","analyzeResult":{"pages":[{"pageNumber":2.0,"width":" 8.5 ","height":true}]}}"#, OcrResponseFormat::Litellm, @@ -703,6 +693,10 @@ mod tests { let dimensions = response.pages[0].dimensions.as_ref().unwrap(); assert_eq!(dimensions.width, Some(816)); assert_eq!(dimensions.height, Some(96)); + } + + #[test] + fn pixel_dimension_rejects_out_of_range_value() { assert!(pixel_dimension(9_223_372_036_854_775_808.0, 1.0, "width").is_err()); } @@ -776,8 +770,8 @@ mod tests { ..Default::default() }; - let error = AzureDocumentIntelligenceOCRConfig - .validate_environment(&connection, &Default::default(), &|name| { + let error = AzureDocumentIntelligenceOcrConfig + .resolve_headers(&connection, &Default::default(), &|name| { (name == AZURE_DI_API_KEY_ENV).then(|| "environment-key".into()) }) .await @@ -800,8 +794,8 @@ mod tests { ..Default::default() }; - let headers = AzureDocumentIntelligenceOCRConfig - .validate_environment(&connection, &Default::default(), &|_| None) + let headers = AzureDocumentIntelligenceOcrConfig + .resolve_headers(&connection, &Default::default(), &|_| None) .await .unwrap(); 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 122f1dbce53..36a07fca8a9 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 @@ -17,17 +17,29 @@ const AZURE_AI_API_KEY_ENV: &str = "AZURE_AI_API_KEY"; const AZURE_AI_API_BASE_ENV: &str = "AZURE_AI_API_BASE"; #[derive(Clone, Debug, Default)] -pub(crate) struct AzureAIOCRConfig; +pub(crate) struct AzureAiOcrConfig; -impl BaseOcrConfig for AzureAIOCRConfig { +impl BaseOcrConfig for AzureAiOcrConfig { type OcrParams = OpaqueParams; type ProviderRequest = MistralOcrRequest; type Environment = Vec<(String, String)>; + fn get_supported_ocr_params(&self, model: &str) -> &'static [&'static str] { + MistralOcrConfig.get_supported_ocr_params(model) + } + fn get_api_key_env_var(&self) -> Option<&'static str> { Some(AZURE_AI_API_KEY_ENV) } + fn map_ocr_params( + &self, + non_default_params: &CallArguments, + model: &str, + ) -> Result { + MistralOcrConfig.map_ocr_params(non_default_params, model) + } + async fn validate_environment( &self, request: &PreparedOcrRequest, @@ -40,39 +52,27 @@ impl BaseOcrConfig for AzureAIOCRConfig { &request.input_sources, )? }; - self.validate_environment(&request.connection, &config, &credential_env) + self.resolve_headers(&request.connection, &config, &credential_env) .await } fn get_complete_url( &self, request: &PreparedOcrRequest, - _params: &Self::OcrParams, + _optional_params: &Self::OcrParams, _environment: &Self::Environment, ) -> Result { - self.get_complete_url(request.connection.api_base.as_deref(), &credential_env) + self.build_ocr_url(request.connection.api_base.as_deref(), &credential_env) } fn transform_ocr_request( &self, model: &str, document: OcrDocument, - params: &OpaqueParams, + optional_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) - } - - fn map_ocr_params( - &self, - arguments: &CallArguments, - model: &str, - ) -> Result { - MistralOcrConfig.map_ocr_params(arguments, model) + MistralOcrConfig.transform_ocr_request(model, document, optional_params, headers) } async fn async_transform_ocr_request( @@ -106,7 +106,7 @@ impl BaseOcrConfig for AzureAIOCRConfig { } } -impl AzureAIOCRConfig { +impl AzureAiOcrConfig { /// Python `AzureAIOCRConfig.validate_environment` requires the endpoint /// before it resolves credentials; keep that order so a missing base is /// reported without invoking any token provider. @@ -124,22 +124,7 @@ impl AzureAIOCRConfig { )) } - fn get_complete_url( - &self, - api_base: Option<&str>, - env_lookup: &dyn Fn(&str) -> Option, - ) -> Result { - let base = Self::resolve_api_base(api_base, env_lookup)?; - let path: Vec<&str> = AZURE_AI_OCR_PATH.trim_matches('/').split('/').collect(); - ApiUrl::parse(&base) - .and_then(|url| url.complete_path(&path)) - .map(|url| url.into_string()) - .map_err(|_| crate::ocr::Error::RequestField { - path: "api_base".into(), - }) - } - - pub(super) async fn validate_environment( + async fn resolve_headers( &self, connection: &OcrConnection, config: &AzureAuthInputs, @@ -169,6 +154,21 @@ impl AzureAIOCRConfig { super::common_utils::validate_destination(connection, key.source())?; Ok(bearer_headers(connection, key.value())) } + + fn build_ocr_url( + &self, + api_base: Option<&str>, + env_lookup: &dyn Fn(&str) -> Option, + ) -> Result { + let base = Self::resolve_api_base(api_base, env_lookup)?; + let path: Vec<&str> = AZURE_AI_OCR_PATH.trim_matches('/').split('/').collect(); + ApiUrl::parse(&base) + .and_then(|url| url.complete_path(&path)) + .map(|url| url.into_string()) + .map_err(|_| crate::ocr::Error::RequestField { + path: "api_base".into(), + }) + } } fn bearer_headers(connection: &OcrConnection, key: &str) -> Vec<(String, String)> { @@ -185,31 +185,41 @@ fn nonblank(value: Option) -> Option { #[cfg(test)] mod tests { + use rstest::{fixture, rstest}; + use super::*; - #[test] - fn completes_azure_path_and_preserves_query() { + #[fixture] + fn connection() -> OcrConnection { + OcrConnection { + api_key: Some("request-key".into()), + api_base: Some("https://example.com".into()), + ..Default::default() + } + } + + #[rstest] + #[case::base_with_query( + "https://example.com/?tenant=a", + "https://example.com/providers/mistral/azure/ocr?tenant=a" + )] + #[case::complete_endpoint( + "https://example.com/providers/mistral/azure/ocr", + "https://example.com/providers/mistral/azure/ocr" + )] + fn completes_azure_path_and_preserves_query(#[case] api_base: &str, #[case] expected: &str) { assert_eq!( - AzureAIOCRConfig - .get_complete_url(Some("https://example.com/?tenant=a"), &|_| None) + AzureAiOcrConfig + .build_ocr_url(Some(api_base), &|_| None) .unwrap(), - "https://example.com/providers/mistral/azure/ocr?tenant=a" - ); - assert_eq!( - AzureAIOCRConfig - .get_complete_url( - Some("https://example.com/providers/mistral/azure/ocr"), - &|_| None - ) - .unwrap(), - "https://example.com/providers/mistral/azure/ocr" + expected ); } #[test] fn missing_api_base_is_structured() { assert!(matches!( - AzureAIOCRConfig::resolve_api_base(None, &|_| None), + AzureAiOcrConfig::resolve_api_base(None, &|_| None), Err(crate::ocr::Error::Auth( litellm_auth::Error::MissingApiBase { provider: "Azure AI", @@ -219,17 +229,16 @@ mod tests { )); } + #[rstest] #[tokio::test] - async fn supplied_authorization_precedes_keys() { + async fn supplied_authorization_precedes_keys(connection: OcrConnection) { let connection = OcrConnection { - api_key: Some("request-key".into()), - api_base: Some("https://example.com".into()), extra_headers: vec![("authorization".into(), "Bearer prepared".into())], - ..Default::default() + ..connection }; assert_eq!( - AzureAIOCRConfig - .validate_environment(&connection, &Default::default(), &|_| { + AzureAiOcrConfig + .resolve_headers(&connection, &Default::default(), &|_| { Some("environment-key".into()) }) .await @@ -238,16 +247,12 @@ mod tests { ); } + #[rstest] #[tokio::test] - async fn request_key_precedes_environment_key() { - let connection = OcrConnection { - api_key: Some("request-key".into()), - api_base: Some("https://example.com".into()), - ..Default::default() - }; + async fn request_key_precedes_environment_key(connection: OcrConnection) { assert_eq!( - AzureAIOCRConfig - .validate_environment(&connection, &Default::default(), &|_| { + AzureAiOcrConfig + .resolve_headers(&connection, &Default::default(), &|_| { Some("environment-key".into()) }) .await @@ -264,8 +269,8 @@ mod tests { ..Default::default() }; - let error = AzureAIOCRConfig - .validate_environment(&connection, &Default::default(), &|name| { + let error = AzureAiOcrConfig + .resolve_headers(&connection, &Default::default(), &|name| { (name == AZURE_AI_API_KEY_ENV).then(|| "environment-key".into()) }) .await @@ -288,8 +293,8 @@ mod tests { ..Default::default() }; - let headers = AzureAIOCRConfig - .validate_environment(&connection, &Default::default(), &|_| None) + let headers = AzureAiOcrConfig + .resolve_headers(&connection, &Default::default(), &|_| None) .await .unwrap(); diff --git a/litellm-rust/crates/core/src/providers/azure_ai/messages/mod.rs b/litellm-rust/crates/core/src/llms/base_llm/anthropic_messages/mod.rs similarity index 100% rename from litellm-rust/crates/core/src/providers/azure_ai/messages/mod.rs rename to litellm-rust/crates/core/src/llms/base_llm/anthropic_messages/mod.rs diff --git a/litellm-rust/crates/core/src/messages/transformation.rs b/litellm-rust/crates/core/src/llms/base_llm/anthropic_messages/transformation.rs similarity index 82% rename from litellm-rust/crates/core/src/messages/transformation.rs rename to litellm-rust/crates/core/src/llms/base_llm/anthropic_messages/transformation.rs index 2719e62d280..37bf8884ec0 100644 --- a/litellm-rust/crates/core/src/messages/transformation.rs +++ b/litellm-rust/crates/core/src/llms/base_llm/anthropic_messages/transformation.rs @@ -1,5 +1,5 @@ -use super::Error; -use super::types::{AnthropicMessagesRequest, AnthropicMessagesResponse}; +use crate::messages::Error; +use crate::messages::types::{AnthropicMessagesRequest, AnthropicMessagesResponse}; #[derive(Clone, Copy, Debug, PartialEq, Eq)] pub enum MessagesAuthStrategy { @@ -16,14 +16,29 @@ impl MessagesAuthStrategy { } } -pub trait AnthropicMessagesProviderConfig: Sync { - fn complete_url( +pub trait BaseAnthropicMessagesConfig: Sync { + fn get_complete_url( &self, api_base: Option<&str>, model: &str, env_lookup: &dyn Fn(&str) -> Option, ) -> Result; + fn transform_anthropic_messages_request( + &self, + request: AnthropicMessagesRequest, + ) -> Result { + Ok(request) + } + + fn transform_anthropic_messages_response( + &self, + _model: &str, + response: AnthropicMessagesResponse, + ) -> Result { + Ok(response) + } + fn resolve_api_key( &self, api_key: Option<&str>, @@ -44,19 +59,4 @@ pub trait AnthropicMessagesProviderConfig: Sync { ("content-type", "application/json"), ] } - - fn transform_request( - &self, - request: AnthropicMessagesRequest, - ) -> Result { - Ok(request) - } - - fn transform_response( - &self, - _model: &str, - response: AnthropicMessagesResponse, - ) -> Result { - Ok(response) - } } diff --git a/litellm-rust/crates/core/src/providers/bedrock/chat_completions/mod.rs b/litellm-rust/crates/core/src/llms/base_llm/audio_transcription/mod.rs similarity index 100% rename from litellm-rust/crates/core/src/providers/bedrock/chat_completions/mod.rs rename to litellm-rust/crates/core/src/llms/base_llm/audio_transcription/mod.rs diff --git a/litellm-rust/crates/core/src/audio_transcription/transformation.rs b/litellm-rust/crates/core/src/llms/base_llm/audio_transcription/transformation.rs similarity index 61% rename from litellm-rust/crates/core/src/audio_transcription/transformation.rs rename to litellm-rust/crates/core/src/llms/base_llm/audio_transcription/transformation.rs index f8082991241..b478bd4caab 100644 --- a/litellm-rust/crates/core/src/audio_transcription/transformation.rs +++ b/litellm-rust/crates/core/src/llms/base_llm/audio_transcription/transformation.rs @@ -1,7 +1,9 @@ use serde_json::{Map, Value}; -use super::Error; -use super::types::{AudioTranscriptionRequestData, AudioTranscriptionResponseData}; +use crate::audio_transcription::Error; +use crate::audio_transcription::types::{ + AudioTranscriptionRequestData, AudioTranscriptionResponseData, +}; #[derive(Clone, Debug, PartialEq, Eq)] pub enum AudioTranscriptionAuth { @@ -12,34 +14,21 @@ pub enum AudioTranscriptionAuth { }, } -pub trait AudioTranscriptionProviderConfig: Sync { - fn supported_transcription_params(&self) -> &'static [&'static str]; +pub trait BaseAudioTranscriptionConfig: Sync { + fn get_supported_openai_params(&self) -> &'static [&'static str]; - fn map_transcription_params(&self, params: &Map) -> Map { - params + fn map_transcription_params( + &self, + non_default_params: &Map, + ) -> Map { + non_default_params .iter() - .filter(|(key, _)| { - self.supported_transcription_params() - .contains(&key.as_str()) - }) + .filter(|(key, _)| self.get_supported_openai_params().contains(&key.as_str())) .map(|(key, value)| (key.clone(), value.clone())) .collect() } - fn transform_transcription_request( - &self, - model: &str, - audio: Value, - optional_params: Map, - ) -> Result; - - fn transform_transcription_response( - &self, - model: &str, - response_json: Value, - ) -> Result; - - fn complete_url( + fn get_complete_url( &self, api_base: Option<&str>, model: &str, @@ -47,6 +36,19 @@ pub trait AudioTranscriptionProviderConfig: Sync { env_lookup: &dyn Fn(&str) -> Option, ) -> Result; + fn transform_audio_transcription_request( + &self, + model: &str, + audio: Value, + optional_params: Map, + ) -> Result; + + fn transform_audio_transcription_response( + &self, + model: &str, + response_json: Value, + ) -> Result; + fn auth_strategy( &self, model: &str, diff --git a/litellm-rust/crates/core/src/providers/openai/responses/mod.rs b/litellm-rust/crates/core/src/llms/base_llm/chat/mod.rs similarity index 100% rename from litellm-rust/crates/core/src/providers/openai/responses/mod.rs rename to litellm-rust/crates/core/src/llms/base_llm/chat/mod.rs diff --git a/litellm-rust/crates/core/src/chat_completions/transformation.rs b/litellm-rust/crates/core/src/llms/base_llm/chat/transformation.rs similarity index 95% rename from litellm-rust/crates/core/src/chat_completions/transformation.rs rename to litellm-rust/crates/core/src/llms/base_llm/chat/transformation.rs index 2325e22e019..cb340db7326 100644 --- a/litellm-rust/crates/core/src/chat_completions/transformation.rs +++ b/litellm-rust/crates/core/src/llms/base_llm/chat/transformation.rs @@ -1,11 +1,17 @@ use serde_json::{Map, Value}; -use super::Error; -use super::types::{ +use crate::chat_completions::Error; +use crate::chat_completions::types::{ ChatCompletionsResponse, ChatMessage, ChatMessageContent, ProviderChatRequestData, ProviderChatResponseData, }; +pub const STREAM_PARAM: &str = "stream"; + +/// Message fields that carry no meaning for the upstream body, so their +/// presence does not make a request untranslatable. +const IGNORABLE_MESSAGE_FIELDS: &[&str] = &["name"]; + /// How the upstream call is authenticated. API-key strategies are resolved in /// `prepare`; SigV4 needs the serialized body, so the handler signs it. #[derive(Clone, Debug, PartialEq, Eq)] @@ -25,14 +31,11 @@ pub enum ChatCompletionsAuth { #[derive(Clone, Copy, Debug, PartialEq, Eq)] pub struct Unsupported(pub &'static str); -pub const STREAM_PARAM: &str = "stream"; +pub trait BaseConfig: Sync { + /// Supported OpenAI parameter names paired with their provider names. + fn supported_openai_param_mappings(&self) -> &'static [(&'static str, &'static str)]; -/// Message fields that carry no meaning for the upstream body, so their -/// presence does not make a request untranslatable. -const IGNORABLE_MESSAGE_FIELDS: &[&str] = &["name"]; - -pub trait ChatCompletionsProviderConfig: Sync { - fn complete_url( + fn get_complete_url( &self, api_base: Option<&str>, model: &str, @@ -40,6 +43,19 @@ pub trait ChatCompletionsProviderConfig: Sync { env_lookup: &dyn Fn(&str) -> Option, ) -> Result; + fn transform_request( + &self, + model: &str, + messages: Vec, + optional_params: Map, + ) -> Result; + + fn transform_response( + &self, + model: &str, + response: ProviderChatResponseData, + ) -> Result; + fn auth( &self, api_key: Option<&str>, @@ -62,9 +78,6 @@ pub trait ChatCompletionsProviderConfig: Sync { false } - /// Supported OpenAI parameter names paired with their provider names. - fn supported_openai_params(&self) -> &'static [(&'static str, &'static str)]; - /// Parameters consumed as call configuration (credentials, endpoints) /// rather than placed in the body. Accepted, never serialized. fn config_params(&self) -> &'static [&'static str] { @@ -77,25 +90,12 @@ pub trait ChatCompletionsProviderConfig: Sync { optional_params: &Map, ) -> Option { unsupported_param( - self.supported_openai_params(), + self.supported_openai_param_mappings(), self.config_params(), optional_params, ) .or_else(|| messages.iter().find_map(unsupported_message)) } - - fn transform_request( - &self, - model: &str, - messages: Vec, - optional_params: Map, - ) -> Result; - - fn transform_response( - &self, - model: &str, - response: ProviderChatResponseData, - ) -> Result; } pub fn unsupported_param( diff --git a/litellm-rust/crates/core/src/llms/base_llm/mod.rs b/litellm-rust/crates/core/src/llms/base_llm/mod.rs index 079e0c41eae..5cd48a21fb6 100644 --- a/litellm-rust/crates/core/src/llms/base_llm/mod.rs +++ b/litellm-rust/crates/core/src/llms/base_llm/mod.rs @@ -1 +1,4 @@ +pub mod anthropic_messages; +pub mod audio_transcription; +pub mod chat; pub(crate) mod ocr; diff --git a/litellm-rust/crates/core/src/providers/bedrock/audio_transcription.rs b/litellm-rust/crates/core/src/llms/bedrock/audio_transcription/mod.rs similarity index 91% rename from litellm-rust/crates/core/src/providers/bedrock/audio_transcription.rs rename to litellm-rust/crates/core/src/llms/bedrock/audio_transcription/mod.rs index 12ea91672e8..49397e00901 100644 --- a/litellm-rust/crates/core/src/providers/bedrock/audio_transcription.rs +++ b/litellm-rust/crates/core/src/llms/bedrock/audio_transcription/mod.rs @@ -1,15 +1,15 @@ use serde_json::{Map, Value, json}; -pub use super::aws_base::{aws_auth_config, bedrock_model_id_and_region, resolve_bedrock_region}; -use super::constants::{BEDROCK_RUNTIME_ENDPOINT_TEMPLATE, BEDROCK_SERVICE}; use crate::audio_transcription::Error; -use crate::audio_transcription::transformation::{ - AudioTranscriptionAuth, AudioTranscriptionProviderConfig, -}; use crate::audio_transcription::types::{ AudioTranscriptionRequestData, AudioTranscriptionResponseData, }; use crate::http_utils::json_type_name; +use crate::llms::base_llm::audio_transcription::transformation::{ + AudioTranscriptionAuth, BaseAudioTranscriptionConfig, +}; +use litellm_auth_aws::constants::{BEDROCK_RUNTIME_ENDPOINT_TEMPLATE, BEDROCK_SERVICE}; +use litellm_auth_aws::{bedrock_model_id_and_region, resolve_bedrock_region}; const SUPPORTED_PARAMS: &[&str] = &["language", "prompt", "temperature", "response_format"]; @@ -45,12 +45,12 @@ fn optional_string<'a>(params: &'a Map, key: &str) -> Option<&'a .filter(|value| !value.is_empty()) } -impl AudioTranscriptionProviderConfig for BedrockAudioTranscriptionConfig { - fn supported_transcription_params(&self) -> &'static [&'static str] { +impl BaseAudioTranscriptionConfig for BedrockAudioTranscriptionConfig { + fn get_supported_openai_params(&self) -> &'static [&'static str] { SUPPORTED_PARAMS } - fn transform_transcription_request( + fn transform_audio_transcription_request( &self, _model: &str, audio: Value, @@ -83,7 +83,7 @@ impl AudioTranscriptionProviderConfig for BedrockAudioTranscriptionConfig { }) } - fn transform_transcription_response( + fn transform_audio_transcription_response( &self, _model: &str, response_json: Value, @@ -105,7 +105,7 @@ impl AudioTranscriptionProviderConfig for BedrockAudioTranscriptionConfig { Ok(AudioTranscriptionResponseData { text }) } - fn complete_url( + fn get_complete_url( &self, api_base: Option<&str>, model: &str, @@ -160,7 +160,7 @@ mod tests { ]); let params = BEDROCK_AUDIO_TRANSCRIPTION_CONFIG.map_transcription_params(¶ms); let result = BEDROCK_AUDIO_TRANSCRIPTION_CONFIG - .transform_transcription_request( + .transform_audio_transcription_request( "mistral.voxtral-mini-3b-2507", json!({"data": "AQI=", "format": "wav", "filename": "sample.wav"}), params, @@ -185,7 +185,7 @@ mod tests { #[test] fn response_concatenates_content_blocks() { let result = BEDROCK_AUDIO_TRANSCRIPTION_CONFIG - .transform_transcription_response( + .transform_audio_transcription_response( "model", json!({"output": {"message": {"content": [{"text": "hello "}, {"text": "world"}]}}}), ) @@ -196,7 +196,7 @@ mod tests { #[test] fn invalid_audio_is_rejected() { - let result = BEDROCK_AUDIO_TRANSCRIPTION_CONFIG.transform_transcription_request( + let result = BEDROCK_AUDIO_TRANSCRIPTION_CONFIG.transform_audio_transcription_request( "model", json!({"data": "AQI="}), Map::new(), @@ -208,7 +208,7 @@ mod tests { fn region_and_url_precedence_match_python() { let params = Map::from_iter([("aws_region_name".to_string(), json!("eu-west-1"))]); let url = BEDROCK_AUDIO_TRANSCRIPTION_CONFIG - .complete_url( + .get_complete_url( None, "bedrock/us-east-1/mistral.voxtral-mini-3b-2507", ¶ms, diff --git a/litellm-rust/crates/core/src/providers/bedrock/chat_completions/transformation.rs b/litellm-rust/crates/core/src/llms/bedrock/chat/converse_transformation.rs similarity index 93% rename from litellm-rust/crates/core/src/providers/bedrock/chat_completions/transformation.rs rename to litellm-rust/crates/core/src/llms/bedrock/chat/converse_transformation.rs index 53d3842955c..525bb6d7abc 100644 --- a/litellm-rust/crates/core/src/providers/bedrock/chat_completions/transformation.rs +++ b/litellm-rust/crates/core/src/llms/bedrock/chat/converse_transformation.rs @@ -1,19 +1,18 @@ use serde_json::{Map, Value, json}; -use super::super::aws_base::{bedrock_model_id_and_region, resolve_bedrock_region}; -use super::super::constants::{AWS_BEARER_TOKEN_BEDROCK, BEDROCK_RUNTIME_ENDPOINT_TEMPLATE}; use crate::chat_completions::Error; use crate::chat_completions::conversation::{Conversation, TurnRole, build_conversation}; use crate::chat_completions::response_utils::{finish_reason_for, unix_now, usage_from_parts}; -use crate::chat_completions::transformation::{ - ChatCompletionsAuth, ChatCompletionsProviderConfig, Unsupported, unsupported_message, - unsupported_param, -}; use crate::chat_completions::types::{ ChatCompletionsChoice, ChatCompletionsChoiceMessage, ChatCompletionsResponse, ChatCompletionsUsage, ChatMessage, ChatMessageContent, ProviderChatRequestData, ProviderChatResponseData, }; +use crate::llms::base_llm::chat::transformation::{ + BaseConfig, ChatCompletionsAuth, Unsupported, unsupported_message, unsupported_param, +}; +use litellm_auth_aws::constants::{AWS_BEARER_TOKEN_BEDROCK, BEDROCK_RUNTIME_ENDPOINT_TEMPLATE}; +use litellm_auth_aws::{bedrock_model_id_and_region, resolve_bedrock_region}; /// Converse parameter names, post `map_openai_params`, that the Rust path can /// place verbatim in `inferenceConfig`. @@ -49,62 +48,16 @@ const CONFIG_PARAMS: &[&str] = &[ const CONVERSE_PATH_SUFFIX: &str = "/converse"; -pub struct BedrockChatCompletionsConfig; +pub struct AmazonConverseConfig; -pub const BEDROCK_CHAT_COMPLETIONS_CONFIG: BedrockChatCompletionsConfig = - BedrockChatCompletionsConfig; +pub const BEDROCK_CHAT_COMPLETIONS_CONFIG: AmazonConverseConfig = AmazonConverseConfig; -fn converse_body(conversation: &Conversation, params: &Map) -> Value { - let messages: Vec = conversation - .turns - .iter() - .map(|turn| { - json!({ - "role": turn.role.as_str(), - "content": turn.texts.iter().map(|text| json!({"text": text})).collect::>(), - }) - }) - .collect(); - - let inference_config = Map::from_iter(SUPPORTED_PARAMS.iter().filter_map(|(_, name)| { - params - .get(*name) - .map(|value| ((*name).to_string(), value.clone())) - })); - - let system: Vec = conversation - .system - .iter() - .map(|text| json!({"text": text})) - .collect(); - - Value::Object(Map::from_iter( - [ - ( - "inferenceConfig".to_string(), - Value::Object(inference_config), - ), - ("messages".to_string(), json!(messages)), - ] - .into_iter() - .chain((!system.is_empty()).then(|| ("system".to_string(), json!(system)))), - )) -} - -fn has_blank_text(message: &ChatMessage) -> bool { - match &message.content { - None => false, - Some(ChatMessageContent::Text(text)) => text.trim().is_empty(), - Some(ChatMessageContent::Parts(parts)) => parts.iter().any(|part| { - part.get("text") - .and_then(Value::as_str) - .is_none_or(|text| text.trim().is_empty()) - }), +impl BaseConfig for AmazonConverseConfig { + fn supported_openai_param_mappings(&self) -> &'static [(&'static str, &'static str)] { + SUPPORTED_PARAMS } -} -impl ChatCompletionsProviderConfig for BedrockChatCompletionsConfig { - fn complete_url( + fn get_complete_url( &self, api_base: Option<&str>, model: &str, @@ -131,82 +84,6 @@ impl ChatCompletionsProviderConfig for BedrockChatCompletionsConfig { Ok(format!("{endpoint}/model/{model_id}{CONVERSE_PATH_SUFFIX}")) } - fn auth( - &self, - api_key: Option<&str>, - model: &str, - optional_params: &Map, - env_lookup: &dyn Fn(&str) -> Option, - ) -> Result { - // Python reads `api_key` as the Bedrock bearer token and consults the - // env only when the caller passed none, so a caller-supplied empty key - // falls through to SigV4 without reaching for the environment. An - // all-whitespace token stays a bearer token here because Python sends - // it too: treating it as absent would sign as the host principal - // instead, which is the identity swap this branch exists to prevent. - let bearer = match api_key { - Some(key) => Some(key.to_string()), - None => env_lookup(AWS_BEARER_TOKEN_BEDROCK), - } - .filter(|token| !token.is_empty()); - if let Some(token) = bearer { - return Ok(ChatCompletionsAuth::Bearer { token }); - } - let (_, model_region) = bedrock_model_id_and_region(model); - Ok(ChatCompletionsAuth::AwsSigV4 { - region: resolve_bedrock_region(model_region.as_deref(), optional_params, env_lookup), - }) - } - - fn default_headers(&self) -> &'static [(&'static str, &'static str)] { - &[("Content-Type", "application/json")] - } - - fn supported_openai_params(&self) -> &'static [(&'static str, &'static str)] { - SUPPORTED_PARAMS - } - - fn config_params(&self) -> &'static [&'static str] { - CONFIG_PARAMS - } - - fn unsupported_reason( - &self, - messages: &[ChatMessage], - optional_params: &Map, - ) -> Option { - unsupported_param( - self.supported_openai_params(), - CONFIG_PARAMS, - optional_params, - ) - .or_else(|| messages.iter().find_map(unsupported_message)) - // Python's Converse translation drops blank text blocks instead of - // substituting the placeholder the shared conversation builder - // applies, so decline blank text rather than diverge. - .or_else(|| { - messages - .iter() - .any(has_blank_text) - .then_some(Unsupported("blank message text")) - }) - // Converse has no assistant prefill: Python inserts a continue turn - // when a conversation opens or closes on an assistant message, and - // only under `litellm.modify_params`, which the core cannot see. - // Declining both ends also keeps the shared builder's final - // assistant right-strip (an Anthropic rule) unreachable here. - .or_else(|| { - let conversation = build_conversation(messages); - let ends_on_assistant = conversation - .turns - .last() - .is_some_and(|turn| turn.role == TurnRole::Assistant); - (!conversation.opens_on_user_turn() || ends_on_assistant).then_some(Unsupported( - "conversation does not run user turn to user turn", - )) - }) - } - fn transform_request( &self, _model: &str, @@ -295,6 +172,127 @@ impl ChatCompletionsProviderConfig for BedrockChatCompletionsConfig { usage, }) } + + fn auth( + &self, + api_key: Option<&str>, + model: &str, + optional_params: &Map, + env_lookup: &dyn Fn(&str) -> Option, + ) -> Result { + // Python reads `api_key` as the Bedrock bearer token and consults the + // env only when the caller passed none, so a caller-supplied empty key + // falls through to SigV4 without reaching for the environment. An + // all-whitespace token stays a bearer token here because Python sends + // it too: treating it as absent would sign as the host principal + // instead, which is the identity swap this branch exists to prevent. + let bearer = match api_key { + Some(key) => Some(key.to_string()), + None => env_lookup(AWS_BEARER_TOKEN_BEDROCK), + } + .filter(|token| !token.is_empty()); + if let Some(token) = bearer { + return Ok(ChatCompletionsAuth::Bearer { token }); + } + let (_, model_region) = bedrock_model_id_and_region(model); + Ok(ChatCompletionsAuth::AwsSigV4 { + region: resolve_bedrock_region(model_region.as_deref(), optional_params, env_lookup), + }) + } + + fn default_headers(&self) -> &'static [(&'static str, &'static str)] { + &[("Content-Type", "application/json")] + } + + fn config_params(&self) -> &'static [&'static str] { + CONFIG_PARAMS + } + + fn unsupported_reason( + &self, + messages: &[ChatMessage], + optional_params: &Map, + ) -> Option { + unsupported_param( + self.supported_openai_param_mappings(), + CONFIG_PARAMS, + optional_params, + ) + .or_else(|| messages.iter().find_map(unsupported_message)) + // Python's Converse translation drops blank text blocks instead of + // substituting the placeholder the shared conversation builder + // applies, so decline blank text rather than diverge. + .or_else(|| { + messages + .iter() + .any(has_blank_text) + .then_some(Unsupported("blank message text")) + }) + // Converse has no assistant prefill: Python inserts a continue turn + // when a conversation opens or closes on an assistant message, and + // only under `litellm.modify_params`, which the core cannot see. + // Declining both ends also keeps the shared builder's final + // assistant right-strip (an Anthropic rule) unreachable here. + .or_else(|| { + let conversation = build_conversation(messages); + let ends_on_assistant = conversation + .turns + .last() + .is_some_and(|turn| turn.role == TurnRole::Assistant); + (!conversation.opens_on_user_turn() || ends_on_assistant).then_some(Unsupported( + "conversation does not run user turn to user turn", + )) + }) + } +} + +fn converse_body(conversation: &Conversation, optional_params: &Map) -> Value { + let messages: Vec = conversation + .turns + .iter() + .map(|turn| { + json!({ + "role": turn.role.as_str(), + "content": turn.texts.iter().map(|text| json!({"text": text})).collect::>(), + }) + }) + .collect(); + + let inference_config = Map::from_iter(SUPPORTED_PARAMS.iter().filter_map(|(_, name)| { + optional_params + .get(*name) + .map(|value| ((*name).to_string(), value.clone())) + })); + + let system: Vec = conversation + .system + .iter() + .map(|text| json!({"text": text})) + .collect(); + + Value::Object(Map::from_iter( + [ + ( + "inferenceConfig".to_string(), + Value::Object(inference_config), + ), + ("messages".to_string(), json!(messages)), + ] + .into_iter() + .chain((!system.is_empty()).then(|| ("system".to_string(), json!(system)))), + )) +} + +fn has_blank_text(message: &ChatMessage) -> bool { + match &message.content { + None => false, + Some(ChatMessageContent::Text(text)) => text.trim().is_empty(), + Some(ChatMessageContent::Parts(parts)) => parts.iter().any(|part| { + part.get("text") + .and_then(Value::as_str) + .is_none_or(|text| text.trim().is_empty()) + }), + } } #[cfg(test)] diff --git a/litellm-rust/crates/core/src/llms/bedrock/chat/mod.rs b/litellm-rust/crates/core/src/llms/bedrock/chat/mod.rs new file mode 100644 index 00000000000..a41ad86ef49 --- /dev/null +++ b/litellm-rust/crates/core/src/llms/bedrock/chat/mod.rs @@ -0,0 +1 @@ +pub mod converse_transformation; diff --git a/litellm-rust/crates/core/src/providers/bedrock/chat_completions/tests.rs b/litellm-rust/crates/core/src/llms/bedrock/chat/tests.rs similarity index 98% rename from litellm-rust/crates/core/src/providers/bedrock/chat_completions/tests.rs rename to litellm-rust/crates/core/src/llms/bedrock/chat/tests.rs index 08ebac9dea1..ed34a46c431 100644 --- a/litellm-rust/crates/core/src/providers/bedrock/chat_completions/tests.rs +++ b/litellm-rust/crates/core/src/llms/bedrock/chat/tests.rs @@ -226,7 +226,7 @@ fn builds_the_converse_url_from_the_region_in_the_model_id() { let config = &BEDROCK_CHAT_COMPLETIONS_CONFIG; assert_eq!( config - .complete_url(None, "us-east-1/anthropic.claude-v2", &Map::new(), &|_| { + .get_complete_url(None, "us-east-1/anthropic.claude-v2", &Map::new(), &|_| { None }) .expect("url builds"), @@ -240,13 +240,13 @@ fn falls_back_to_the_region_env_then_the_default_region() { let with_env = |key: &str| (key == "AWS_REGION_NAME").then(|| "eu-west-1".to_string()); assert_eq!( config - .complete_url(None, "anthropic.claude-v2", &Map::new(), &with_env) + .get_complete_url(None, "anthropic.claude-v2", &Map::new(), &with_env) .expect("url builds"), "https://bedrock-runtime.eu-west-1.amazonaws.com/model/anthropic.claude-v2/converse" ); assert_eq!( config - .complete_url(None, "anthropic.claude-v2", &Map::new(), &|_| None) + .get_complete_url(None, "anthropic.claude-v2", &Map::new(), &|_| None) .expect("url builds"), "https://bedrock-runtime.us-west-2.amazonaws.com/model/anthropic.claude-v2/converse" ); @@ -258,7 +258,7 @@ fn prefers_an_explicit_runtime_endpoint_over_the_api_base() { let overrides = params(json!({"aws_bedrock_runtime_endpoint": "https://vpce.internal/"})); assert_eq!( config - .complete_url( + .get_complete_url( Some("https://ignored.example"), "anthropic.claude-v2", &overrides, @@ -540,7 +540,7 @@ fn leaves_a_complete_converse_url_untouched() { "https://bedrock-runtime.us-east-1.amazonaws.com/model/us.anthropic.claude-v2%3A0/converse"; assert_eq!( config - .complete_url( + .get_complete_url( Some(already_built), "anthropic.claude-v2", &Map::new(), @@ -554,7 +554,7 @@ fn leaves_a_complete_converse_url_untouched() { #[test] fn host_supplied_credentials_outrank_ambient_profile_and_role_state() { - use crate::providers::bedrock::aws_base::host_supplied_credentials; + use litellm_auth_aws::host_supplied_credentials; let supplied = params(json!({ "aws_access_key_id": "AKIAHOST", diff --git a/litellm-rust/crates/core/src/llms/bedrock/mod.rs b/litellm-rust/crates/core/src/llms/bedrock/mod.rs new file mode 100644 index 00000000000..695aeb8af5e --- /dev/null +++ b/litellm-rust/crates/core/src/llms/bedrock/mod.rs @@ -0,0 +1,2 @@ +pub mod audio_transcription; +pub mod chat; 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 996d9e462ab..925e20c8947 100644 --- a/litellm-rust/crates/core/src/llms/cohere/ocr/transformation.rs +++ b/litellm-rust/crates/core/src/llms/cohere/ocr/transformation.rs @@ -4,13 +4,13 @@ use serde_with::serde_as; use crate::call_arguments::{CallArguments, parse_options}; use crate::constants::{COHERE_API_KEY_ENV, COHERE_PARSE_API_BASE}; -use crate::llms::base_llm::ocr::transformation::{BaseOcrConfig, OcrRequestContext}; +use crate::llms::base_llm::ocr::transformation::{BaseOcrConfig, decode_and_normalize_response}; use crate::ocr::OcrClient; use crate::ocr::document::InlineDocument; use crate::ocr::prepare::credential_env; use crate::ocr::types::{ - LiteLLMOcrResponse, OcrConnection, OcrDocument, OcrPage, OcrPageImage, OcrUsageInfo, - PreparedOcrRequest, + LiteLLMOcrResponse, OcrConnection, OcrDocument, OcrPage, OcrPageImage, OcrResponseFormat, + OcrUsageInfo, PreparedOcrRequest, }; use crate::serde_compat::LaxI64; use crate::url_utils::ApiUrl; @@ -88,6 +88,10 @@ impl BaseOcrConfig for CohereParseConfig { type ProviderRequest = CohereRequest; type Environment = Vec<(String, String)>; + fn get_supported_ocr_params(&self, _model: &str) -> &'static [&'static str] { + &["output_format", "req_format"] + } + fn get_api_key_env_var(&self) -> Option<&'static str> { Some(COHERE_API_KEY_ENV) } @@ -99,21 +103,29 @@ impl BaseOcrConfig for CohereParseConfig { } } + fn map_ocr_params( + &self, + non_default_params: &CallArguments, + _model: &str, + ) -> Result { + Ok(parse_options(non_default_params)?) + } + async fn validate_environment( &self, request: &PreparedOcrRequest, _client: &OcrClient, ) -> Result { - self.validate_environment(&request.connection, &credential_env) + self.resolve_headers(&request.connection, &credential_env) } fn get_complete_url( &self, request: &PreparedOcrRequest, - _params: &Self::OcrParams, + _optional_params: &Self::OcrParams, _environment: &Self::Environment, ) -> Result { - self.get_complete_url( + self.build_ocr_url( request .connection .api_base @@ -133,41 +145,13 @@ impl BaseOcrConfig for CohereParseConfig { Ok(build_request(model, image_url, optional_params)) } - fn get_supported_ocr_params(&self, _model: &str) -> &'static [&'static str] { - &["output_format", "req_format"] - } - - fn map_ocr_params( - &self, - arguments: &CallArguments, - _model: &str, - ) -> Result { - Ok(parse_options(arguments)?) - } - - async fn async_transform_ocr_request( - &self, - model: &str, - document: OcrDocument, - optional_params: &CohereOptions, - headers: &[(String, String)], - _context: OcrRequestContext<'_>, - ) -> Result { - self.transform_ocr_request(model, document, optional_params, headers) - } - fn transform_ocr_response( &self, model: &str, raw_response: &[u8], - request_format: crate::ocr::types::OcrResponseFormat, + request_format: OcrResponseFormat, ) -> Result { - crate::llms::base_llm::ocr::transformation::decode_and_normalize_response( - model, - raw_response, - request_format, - normalize_response, - ) + decode_and_normalize_response(model, raw_response, request_format, normalize_response) } fn validate_request_body(&self, body: &Value) -> Result<(), crate::ocr::Error> { @@ -175,6 +159,50 @@ impl BaseOcrConfig for CohereParseConfig { } } +impl CohereParseConfig { + fn resolve_headers( + &self, + connection: &OcrConnection, + env_lookup: &(dyn Fn(&str) -> Option + Sync), + ) -> Result, crate::ocr::Error> { + if crate::http_utils::has_header(&connection.extra_headers, "authorization") { + return Ok(connection.extra_headers.clone()); + } + let key = connection + .api_key + .as_deref() + .map(str::trim) + .filter(|key| !key.is_empty()) + .map(str::to_string) + .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(), + )) + })?; + Ok( + std::iter::once(("Authorization".into(), format!("Bearer {key}"))) + .chain(connection.extra_headers.clone()) + .collect(), + ) + } + + fn build_ocr_url(&self, api_base: &str) -> Result { + let parsed = reqwest::Url::parse(api_base).map_err(|_| invalid_api_base())?; + if !matches!(parsed.scheme(), "http" | "https") { + return Err(invalid_api_base()); + } + ApiUrl::parse(api_base) + .and_then(|url| url.complete_path(&["v2", "parse"])) + .map(|url| url.into_string()) + .map_err(|_| invalid_api_base()) + } +} + pub(crate) fn validate_document(document: &OcrDocument) -> Result<(), crate::ocr::Error> { let OcrDocument::ImageUrl { image_url, .. } = document else { return Err(crate::ocr::Error::CohereImageOnly); @@ -292,50 +320,6 @@ fn billed_pages(response: &CohereResponse) -> Option { response.meta.as_ref()?.billed_units.as_ref()?.pages } -impl CohereParseConfig { - fn get_complete_url(&self, base: &str) -> Result { - let parsed = reqwest::Url::parse(base).map_err(|_| invalid_api_base())?; - if !matches!(parsed.scheme(), "http" | "https") { - return Err(invalid_api_base()); - } - ApiUrl::parse(base) - .and_then(|url| url.complete_path(&["v2", "parse"])) - .map(|url| url.into_string()) - .map_err(|_| invalid_api_base()) - } - - fn validate_environment( - &self, - connection: &OcrConnection, - env_lookup: &(dyn Fn(&str) -> Option + Sync), - ) -> Result, crate::ocr::Error> { - if crate::http_utils::has_header(&connection.extra_headers, "authorization") { - return Ok(connection.extra_headers.clone()); - } - let key = connection - .api_key - .as_deref() - .map(str::trim) - .filter(|key| !key.is_empty()) - .map(str::to_string) - .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(), - )) - })?; - Ok( - std::iter::once(("Authorization".into(), format!("Bearer {key}"))) - .chain(connection.extra_headers.clone()) - .collect(), - ) - } -} - fn invalid_api_base() -> crate::ocr::Error { crate::ocr::Error::RequestField { path: "api_base".into(), @@ -385,27 +369,31 @@ mod tests { ); } - #[test] - fn options_read_known_fields_without_changing_arguments() { + #[rstest] + #[case::cohere(false)] + #[case::azure(true)] + fn options_read_known_fields_without_changing_arguments(#[case] azure: bool) { let arguments = serde_json::from_value(json!({ "output_format":"blocks", "req_format":"native", "extension":false })) .unwrap(); - for config in [false, true] { - let mapped = if config { - crate::llms::azure_ai::ocr::cohere_parse_transformation::AzureAICohereParseConfig - .map_ocr_params(&arguments, "parse") - } else { - CohereParseConfig.map_ocr_params(&arguments, "parse") - } - .unwrap(); - assert_eq!( - serde_json::to_value(mapped).unwrap(), - json!({"output_format":"blocks"}) - ); + let mapped = if azure { + crate::llms::azure_ai::ocr::cohere_parse_transformation::AzureAICohereParseConfig + .map_ocr_params(&arguments, "parse") + } else { + CohereParseConfig.map_ocr_params(&arguments, "parse") } + .unwrap(); + assert_eq!( + serde_json::to_value(mapped).unwrap(), + json!({"output_format":"blocks"}) + ); assert_eq!(arguments["req_format"], "native"); assert_eq!(arguments["extension"], false); + } + + #[test] + fn options_reject_invalid_output_format() { let invalid = serde_json::from_value(json!({"output_format":"html"})).unwrap(); assert!(matches!( CohereParseConfig.map_ocr_params(&invalid, "parse"), @@ -415,13 +403,17 @@ mod tests { } #[test] - fn billed_pages_accept_integral_doubles_and_reject_fractional_counts() { + fn billed_pages_accept_integral_doubles() { let response = serde_json::from_str::( r#"{"pages":[],"meta":{"billed_units":{"pages":1.0}}}"#, ) .unwrap(); let normalized = normalize_response("parse", response).unwrap(); assert_eq!(normalized.usage_info.unwrap().pages_processed, Some(1)); + } + + #[test] + fn billed_pages_reject_fractional_counts() { assert!( serde_json::from_str::( r#"{"pages":[],"meta":{"billed_units":{"pages":1.5}}}"#, @@ -590,25 +582,27 @@ mod tests { assert_eq!(normalized.usage_info.unwrap().pages_processed, Some(3)); } + #[rstest] + #[case::empty(json!({}))] + #[case::null_meta(json!({"meta":null}))] + #[case::null_billed_units(json!({"pages":[],"meta":{"billed_units":null}}))] + fn response_defaults(#[case] value: Value) { + let normalized = + normalize_response("parse", serde_json::from_value(value).unwrap()).unwrap(); + assert!(normalized.pages.is_empty()); + assert_eq!(normalized.usage_info.unwrap().pages_processed, Some(0)); + } + + #[rstest] + #[case::null_pages(json!({"pages":null}))] + #[case::invalid_markdown(json!({"pages":[{"markdown":"text"}]}))] + #[case::invalid_index(json!({"pages":[{"index":"bad"}]}))] + fn response_rejects_invalid_fields(#[case] value: Value) { + assert!(serde_json::from_value::(value).is_err()); + } + #[test] - fn response_defaults_and_invalid_fields() { - for value in [ - json!({}), - json!({"meta":null}), - json!({"pages":[],"meta":{"billed_units":null}}), - ] { - let normalized = - normalize_response("parse", serde_json::from_value(value).unwrap()).unwrap(); - assert!(normalized.pages.is_empty()); - assert_eq!(normalized.usage_info.unwrap().pages_processed, Some(0)); - } - for value in [ - json!({"pages":null}), - json!({"pages":[{"markdown":"text"}]}), - json!({"pages":[{"index":"bad"}]}), - ] { - assert!(serde_json::from_value::(value).is_err()); - } + fn null_markdown_uses_page_defaults() { let normalized = normalize_response( "parse", serde_json::from_value(json!({"pages":[{"markdown":null}]})).unwrap(), @@ -710,24 +704,30 @@ mod tests { ); } + #[rstest] + #[case::document_url(json!({"type":"document_url","document_url":"https://example.com/a.pdf"}))] + #[case::empty_image_url(json!({"type":"image_url","image_url":""}))] + #[case::pdf_data_uri(json!({"type":"image_url","image_url":"data:application/pdf;base64,YQ=="}))] + fn request_requires_image(#[case] value: Value) { + assert!(matches!( + validate_document(&serde_json::from_value(value).unwrap()), + Err(crate::ocr::Error::CohereImageOnly) + )); + } + + #[rstest] + #[case::markdown("markdown", true)] + #[case::blocks("blocks", true)] + #[case::unsupported("html", false)] + fn request_requires_supported_output_format(#[case] format: &str, #[case] valid: bool) { + assert_eq!( + serde_json::from_value::(json!({"output_format":format})).is_ok(), + valid + ); + } + #[test] - fn request_requires_image_and_supported_output_format() { - for value in [ - json!({"type":"document_url","document_url":"https://example.com/a.pdf"}), - json!({"type":"image_url","image_url":""}), - json!({"type":"image_url","image_url":"data:application/pdf;base64,YQ=="}), - ] { - assert!(matches!( - validate_document(&serde_json::from_value(value).unwrap()), - Err(crate::ocr::Error::CohereImageOnly) - )); - } - assert!(serde_json::from_value::(json!({"output_format":"html"})).is_err()); - for format in ["markdown", "blocks"] { - assert!( - serde_json::from_value::(json!({"output_format":format})).is_ok() - ); - } + fn request_defaults_to_markdown() { let request = CohereParseConfig .transform_ocr_request( "parse-v5.0", @@ -746,28 +746,30 @@ mod tests { ); } - #[test] - fn completes_provider_urls_without_duplicate_paths_and_preserves_queries() { - for suffix in ["", "/v2", "/v2/parse"] { - assert_eq!( - CohereParseConfig - .get_complete_url(&format!("https://example.com{suffix}?tenant=a")) - .unwrap(), - "https://example.com/v2/parse?tenant=a" - ); - } + #[rstest] + #[case::base("")] + #[case::version("/v2")] + #[case::complete("/v2/parse")] + fn completes_provider_urls_without_duplicate_paths_and_preserves_queries(#[case] suffix: &str) { + assert_eq!( + CohereParseConfig + .build_ocr_url(&format!("https://example.com{suffix}?tenant=a")) + .unwrap(), + "https://example.com/v2/parse?tenant=a" + ); + } + + #[rstest] + #[case::relative("relative/path")] + #[case::unsupported_scheme("ftp://example.com")] + fn rejects_invalid_urls(#[case] api_base: &str) { + assert!(CohereParseConfig.build_ocr_url(api_base).is_err()); } #[test] - fn rejects_invalid_urls_and_blank_keys() { - assert!(CohereParseConfig.get_complete_url("relative/path").is_err()); - assert!( - CohereParseConfig - .get_complete_url("ftp://example.com") - .is_err() - ); + fn rejects_blank_keys() { assert!(matches!( - CohereParseConfig.validate_environment( + CohereParseConfig.resolve_headers( &OcrConnection { api_key: Some(" ".into()), ..Default::default() diff --git a/litellm-rust/crates/core/src/llms/mod.rs b/litellm-rust/crates/core/src/llms/mod.rs index 3dad380f833..635d381561c 100644 --- a/litellm-rust/crates/core/src/llms/mod.rs +++ b/litellm-rust/crates/core/src/llms/mod.rs @@ -1,6 +1,9 @@ -pub(crate) mod azure_ai; -pub(crate) mod base_llm; +pub mod anthropic; +pub mod azure_ai; +pub mod base_llm; +pub mod bedrock; pub(crate) mod cohere; pub(crate) mod mistral; +pub mod openai; pub(crate) mod reducto; pub(crate) mod vertex_ai; diff --git a/litellm-rust/crates/core/src/providers/openai/mod.rs b/litellm-rust/crates/core/src/llms/openai/mod.rs similarity index 100% rename from litellm-rust/crates/core/src/providers/openai/mod.rs rename to litellm-rust/crates/core/src/llms/openai/mod.rs diff --git a/litellm-rust/crates/core/src/llms/openai/responses/mod.rs b/litellm-rust/crates/core/src/llms/openai/responses/mod.rs new file mode 100644 index 00000000000..f239b6921fa --- /dev/null +++ b/litellm-rust/crates/core/src/llms/openai/responses/mod.rs @@ -0,0 +1 @@ +pub mod transformation; diff --git a/litellm-rust/crates/core/src/providers/openai/responses/transformation.rs b/litellm-rust/crates/core/src/llms/openai/responses/transformation.rs similarity index 86% rename from litellm-rust/crates/core/src/providers/openai/responses/transformation.rs rename to litellm-rust/crates/core/src/llms/openai/responses/transformation.rs index 6203b195d5e..220933d3db0 100644 --- a/litellm-rust/crates/core/src/providers/openai/responses/transformation.rs +++ b/litellm-rust/crates/core/src/llms/openai/responses/transformation.rs @@ -2,11 +2,11 @@ use crate::responses::Error; use crate::responses::types::{ResponsesWsEvent, ResponsesWsTransformResult}; use crate::responses::websocket::{ResponsesWebSocketProviderConfig, enforce_model}; -pub struct OpenAIResponsesWsConfig; +pub struct OpenAiResponsesApiConfig; -pub const OPENAI_RESPONSES_WS_CONFIG: OpenAIResponsesWsConfig = OpenAIResponsesWsConfig; +pub const OPENAI_RESPONSES_WS_CONFIG: OpenAiResponsesApiConfig = OpenAiResponsesApiConfig; -impl ResponsesWebSocketProviderConfig for OpenAIResponsesWsConfig { +impl ResponsesWebSocketProviderConfig for OpenAiResponsesApiConfig { fn supports_native_websocket(&self) -> bool { true } 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 f4ed5946fac..98f981a239d 100644 --- a/litellm-rust/crates/core/src/llms/reducto/ocr/transformation.rs +++ b/litellm-rust/crates/core/src/llms/reducto/ocr/transformation.rs @@ -5,12 +5,15 @@ use serde_json::{Map, Value, json}; use crate::call_arguments::{CallArguments, compose_body}; use crate::constants::{REDUCTO_API_BASE, REDUCTO_API_KEY_ENV, REDUCTO_ID_PREFIX}; -use crate::llms::base_llm::ocr::transformation::{BaseOcrConfig, OcrRequestContext}; +use crate::llms::base_llm::ocr::transformation::{ + BaseOcrConfig, OcrRequestContext, decode_and_normalize_response, +}; use crate::ocr::OcrClient; use crate::ocr::document::InlineDocument; use crate::ocr::prepare::{build_http_request, credential_env, guardrail_document}; use crate::ocr::types::{ - LiteLLMOcrResponse, OcrConnection, OcrDocument, OcrPage, OcrUsageInfo, PreparedOcrRequest, + LiteLLMOcrResponse, OcrConnection, OcrDocument, OcrPage, OcrResponseFormat, OcrUsageInfo, + PreparedOcrRequest, }; use crate::params::OpaqueParams; use crate::url_utils::ApiUrl; @@ -83,50 +86,50 @@ impl BaseOcrConfig for ReductoParseV3Config { type ProviderRequest = ReductoV3Request; type Environment = Vec<(String, String)>; - async fn validate_environment( - &self, - request: &PreparedOcrRequest, - _client: &OcrClient, - ) -> Result { - validate_environment(&request.connection, &credential_env) - } - - fn get_complete_url( - &self, - request: &PreparedOcrRequest, - _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"] } fn map_ocr_params( &self, - arguments: &CallArguments, + non_default_params: &CallArguments, model: &str, ) -> Result { - Ok(arguments + Ok(non_default_params .select(self.get_supported_ocr_params(model)) .into()) } + async fn validate_environment( + &self, + request: &PreparedOcrRequest, + _client: &OcrClient, + ) -> Result { + resolve_headers(&request.connection, &credential_env) + } + + fn get_complete_url( + &self, + request: &PreparedOcrRequest, + _optional_params: &Self::OcrParams, + _environment: &Self::Environment, + ) -> Result { + build_ocr_url(request.connection.api_base.as_deref()) + } + + fn transform_ocr_request( + &self, + _model: &str, + document: OcrDocument, + optional_params: &Self::OcrParams, + _headers: &[(String, String)], + ) -> Result { + Ok(ReductoV3Request { + input: uploaded_file_id(document)?, + params: optional_params.clone(), + }) + } + async fn async_transform_ocr_request( &self, _model: &str, @@ -146,14 +149,9 @@ impl BaseOcrConfig for ReductoParseV3Config { &self, model: &str, raw_response: &[u8], - request_format: crate::ocr::types::OcrResponseFormat, + request_format: OcrResponseFormat, ) -> Result { - crate::llms::base_llm::ocr::transformation::decode_and_normalize_response( - model, - raw_response, - request_format, - normalize_response, - ) + decode_and_normalize_response(model, raw_response, request_format, normalize_response) } async fn prepare_request( @@ -173,6 +171,20 @@ impl BaseOcrConfig for ReductoParseLegacyConfig { type ProviderRequest = ReductoLegacyRequest; type Environment = Vec<(String, String)>; + fn get_supported_ocr_params(&self, _model: &str) -> &'static [&'static str] { + &["enhance"] + } + + fn map_ocr_params( + &self, + non_default_params: &CallArguments, + model: &str, + ) -> Result { + Ok(non_default_params + .select(self.get_supported_ocr_params(model)) + .into()) + } + async fn validate_environment( &self, request: &PreparedOcrRequest, @@ -186,34 +198,23 @@ impl BaseOcrConfig for ReductoParseLegacyConfig { fn get_complete_url( &self, request: &PreparedOcrRequest, - params: &Self::OcrParams, + optional_params: &Self::OcrParams, environment: &Self::Environment, ) -> Result { - ReductoParseV3Config.get_complete_url(request, params, environment) + ReductoParseV3Config.get_complete_url(request, optional_params, environment) } fn transform_ocr_request( &self, _model: &str, document: OcrDocument, - params: &Self::OcrParams, + optional_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"] - } - - fn map_ocr_params( - &self, - arguments: &CallArguments, - model: &str, - ) -> Result { - Ok(arguments - .select(self.get_supported_ocr_params(model)) - .into()) + Ok(build_legacy_body( + uploaded_file_id(document)?, + optional_params, + )) } async fn async_transform_ocr_request( @@ -232,7 +233,7 @@ impl BaseOcrConfig for ReductoParseLegacyConfig { &self, model: &str, raw_response: &[u8], - request_format: crate::ocr::types::OcrResponseFormat, + request_format: OcrResponseFormat, ) -> Result { ReductoParseV3Config.transform_ocr_response(model, raw_response, request_format) } @@ -403,7 +404,7 @@ fn page(index: i64, markdown: String, blocks: Option) -> OcrPage { ..Default::default() } } -fn get_complete_url(api_base: Option<&str>) -> Result { +fn build_ocr_url(api_base: Option<&str>) -> Result { complete_endpoint_url(api_base, "parse") } @@ -420,7 +421,7 @@ fn complete_endpoint_url(api_base: Option<&str>, path: &str) -> Result Option + Sync), ) -> Result, crate::ocr::Error> { @@ -655,7 +656,7 @@ mod tests { api_key: Some("passed-key".into()), ..Default::default() }; - let headers = validate_environment(&connection, &|_| Some("env-key".into())).unwrap(); + let headers = resolve_headers(&connection, &|_| Some("env-key".into())).unwrap(); assert_eq!(headers[0].1, "Bearer passed-key"); } @@ -665,7 +666,7 @@ mod tests { api_key: Some(" ".into()), ..Default::default() }; - let headers = validate_environment(&connection, &|_| Some(" env-key ".into())).unwrap(); + let headers = resolve_headers(&connection, &|_| Some(" env-key ".into())).unwrap(); assert_eq!(headers[0].1, "Bearer env-key"); } @@ -676,7 +677,7 @@ mod tests { ..Default::default() }; assert_eq!( - validate_environment(&connection, &|_| None).unwrap(), + resolve_headers(&connection, &|_| None).unwrap(), connection.extra_headers ); } 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 43bee24b860..ffa0fd28202 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 @@ -2,7 +2,7 @@ use litellm_auth_gcp::{self as vertex, VertexConfig}; use serde::{Deserialize, Serialize}; use serde_json::{Map, Value}; -use super::transformation::VertexAIOCRConfig; +use super::transformation::VertexAiOcrConfig; use crate::call_arguments::CallArguments; use crate::llms::base_llm::ocr::transformation::{BaseOcrConfig, OcrRequestContext}; use crate::ocr::OcrClient; @@ -95,7 +95,7 @@ impl BaseOcrConfig for VertexAIDeepSeekOCRConfig { type Environment = vertex::VertexEnvironment; fn get_api_key_env_var(&self) -> Option<&'static str> { - VertexAIOCRConfig.get_api_key_env_var() + VertexAiOcrConfig.get_api_key_env_var() } fn map_ocr_params( @@ -111,7 +111,9 @@ impl BaseOcrConfig for VertexAIDeepSeekOCRConfig { request: &PreparedOcrRequest, client: &OcrClient, ) -> Result { - BaseOcrConfig::validate_environment(&VertexAIOCRConfig, request, client).await + VertexAiOcrConfig + .validate_environment(request, client) + .await } fn get_complete_url( 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 1183043e9ee..28c2b8a09da 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 @@ -17,17 +17,29 @@ use crate::url_utils::ApiUrl; const DEFAULT_LOCATION: &str = "us-central1"; #[derive(Clone, Debug, Default)] -pub(crate) struct VertexAIOCRConfig; +pub(crate) struct VertexAiOcrConfig; -impl BaseOcrConfig for VertexAIOCRConfig { +impl BaseOcrConfig for VertexAiOcrConfig { type OcrParams = OpaqueParams; type ProviderRequest = MistralOcrRequest; type Environment = vertex::VertexEnvironment; + fn get_supported_ocr_params(&self, model: &str) -> &'static [&'static str] { + MistralOcrConfig.get_supported_ocr_params(model) + } + fn get_api_key_env_var(&self) -> Option<&'static str> { Some("VERTEX_AI_API_KEY") } + fn map_ocr_params( + &self, + non_default_params: &CallArguments, + model: &str, + ) -> Result { + MistralOcrConfig.map_ocr_params(non_default_params, model) + } + async fn validate_environment( &self, request: &PreparedOcrRequest, @@ -37,14 +49,14 @@ impl BaseOcrConfig for VertexAIOCRConfig { &request.optional_params, &request.input_sources, )?; - self.validate_environment(&request.connection, &config, client) + self.resolve_environment(&request.connection, &config, client) .await } fn get_complete_url( &self, request: &PreparedOcrRequest, - _params: &Self::OcrParams, + _optional_params: &Self::OcrParams, environment: &Self::Environment, ) -> Result { let config = VertexConfig::from_sourced_optional_params( @@ -53,7 +65,7 @@ impl BaseOcrConfig for VertexAIOCRConfig { )?; let location = vertex::get_vertex_ai_location(&config, &credential_env) .unwrap_or_else(|| DEFAULT_LOCATION.to_string()); - self.get_complete_url( + self.build_ocr_url( request.connection.api_base.as_deref(), &environment.project_id, &location, @@ -65,22 +77,10 @@ impl BaseOcrConfig for VertexAIOCRConfig { &self, model: &str, document: OcrDocument, - params: &OpaqueParams, + optional_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) - } - - fn map_ocr_params( - &self, - arguments: &CallArguments, - model: &str, - ) -> Result { - MistralOcrConfig.map_ocr_params(arguments, model) + MistralOcrConfig.transform_ocr_request(model, document, optional_params, headers) } async fn async_transform_ocr_request( @@ -120,8 +120,8 @@ impl OcrEnvironment for vertex::VertexEnvironment { } } -impl VertexAIOCRConfig { - pub(super) async fn validate_environment( +impl VertexAiOcrConfig { + async fn resolve_environment( &self, connection: &OcrConnection, config: &VertexConfig, @@ -140,7 +140,7 @@ impl VertexAIOCRConfig { .map_err(crate::ocr::Error::from) } - fn get_complete_url( + fn build_ocr_url( &self, api_base: Option<&str>, project: &str, @@ -198,19 +198,24 @@ fn validate_location(location: &str) -> Result<(), crate::ocr::Error> { #[cfg(test)] mod tests { - use super::VertexAIOCRConfig; + use super::VertexAiOcrConfig; + use rstest::rstest; #[test] fn endpoint_uses_location_project_and_model() { assert_eq!( - VertexAIOCRConfig - .get_complete_url(None, "proj-1", "europe-west4", "mistral-ocr-maas") + VertexAiOcrConfig + .build_ocr_url(None, "proj-1", "europe-west4", "mistral-ocr-maas") .unwrap(), "https://europe-west4-aiplatform.googleapis.com/v1/projects/proj-1/locations/europe-west4/publishers/mistralai/models/mistral-ocr-maas:rawPredict" ); + } + + #[test] + fn endpoint_rejects_invalid_location() { assert!( - VertexAIOCRConfig - .get_complete_url(None, "proj-1", "attacker.example/path", "model") + VertexAiOcrConfig + .build_ocr_url(None, "proj-1", "attacker.example/path", "model") .is_err() ); } @@ -315,13 +320,18 @@ mod tests { ); } + #[rstest] + #[case::mistral(false)] + #[case::vertex(true)] #[tokio::test] - async fn configs_build_complete_requests_and_share_mistral_normalization() { + async fn configs_build_complete_requests_and_share_mistral_normalization( + #[case] use_vertex: bool, + ) { 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::llms::vertex_ai::ocr::transformation::VertexAiOcrConfig; use crate::ocr::test_support::ocr_client; let client = ocr_client(); @@ -348,7 +358,7 @@ mod tests { .prepare_request(&direct, &client) .await .unwrap(); - let vertex_http = VertexAIOCRConfig + let vertex_http = VertexAiOcrConfig .prepare_request(&vertex, &client) .await .unwrap(); @@ -357,24 +367,26 @@ mod tests { 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 http = if use_vertex { + &vertex_http + } else { + &direct_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"}), ) @@ -383,7 +395,7 @@ mod tests { .transform_ocr_response(&direct.model, &payload, Default::default()) .unwrap() .into_json(); - let vertex_response = VertexAIOCRConfig + let vertex_response = VertexAiOcrConfig .transform_ocr_response(&vertex.model, &payload, Default::default()) .unwrap() .into_json(); diff --git a/litellm-rust/crates/core/src/messages/common_utils.rs b/litellm-rust/crates/core/src/messages/common_utils.rs index 73e9a964749..a0a120c34a9 100644 --- a/litellm-rust/crates/core/src/messages/common_utils.rs +++ b/litellm-rust/crates/core/src/messages/common_utils.rs @@ -1,17 +1,17 @@ use serde_json::{Map, Value}; use super::Error; -use super::transformation::AnthropicMessagesProviderConfig; use crate::http_utils::string_headers as shared_string_headers; pub(super) use crate::http_utils::{has_bearer_auth, has_header, truncate_error_body}; -use crate::providers::anthropic::messages::transformation::ANTHROPIC_MESSAGES_CONFIG; -use crate::providers::azure_ai::messages::transformation::AZURE_ANTHROPIC_MESSAGES_CONFIG; +use crate::llms::anthropic::experimental_pass_through::messages::transformation::ANTHROPIC_MESSAGES_CONFIG; +use crate::llms::azure_ai::anthropic::messages_transformation::AZURE_ANTHROPIC_MESSAGES_CONFIG; +use crate::llms::base_llm::anthropic_messages::transformation::BaseAnthropicMessagesConfig; const HEADER_CONTEXT: &str = "messages"; pub(super) fn messages_provider_config( provider: &str, -) -> Option<&'static dyn AnthropicMessagesProviderConfig> { +) -> Option<&'static dyn BaseAnthropicMessagesConfig> { match provider { "anthropic" => Some(&ANTHROPIC_MESSAGES_CONFIG), "azure_ai" => Some(&AZURE_ANTHROPIC_MESSAGES_CONFIG), diff --git a/litellm-rust/crates/core/src/messages/handler.rs b/litellm-rust/crates/core/src/messages/handler.rs index 8d1d4432627..a7393e33a92 100644 --- a/litellm-rust/crates/core/src/messages/handler.rs +++ b/litellm-rust/crates/core/src/messages/handler.rs @@ -37,7 +37,9 @@ pub(super) async fn execute_messages_provider_call( let response = serde_json::from_str(&text) .map_err(|err| Error::InvalidResponse(format!("invalid messages response JSON: {err}")))?; - request.config.transform_response(&request.model, response) + request + .config + .transform_anthropic_messages_response(&request.model, response) } pub(super) async fn execute_messages_provider_stream( diff --git a/litellm-rust/crates/core/src/messages/mod.rs b/litellm-rust/crates/core/src/messages/mod.rs index 156f42056f1..8f6fffcaf7f 100644 --- a/litellm-rust/crates/core/src/messages/mod.rs +++ b/litellm-rust/crates/core/src/messages/mod.rs @@ -13,7 +13,6 @@ mod client; mod common_utils; mod handler; mod prepare; -pub mod transformation; pub mod types; use handler::{execute_messages_provider_call, execute_messages_provider_stream}; diff --git a/litellm-rust/crates/core/src/messages/prepare.rs b/litellm-rust/crates/core/src/messages/prepare.rs index 0deb42a34ae..a3c93746d3e 100644 --- a/litellm-rust/crates/core/src/messages/prepare.rs +++ b/litellm-rust/crates/core/src/messages/prepare.rs @@ -2,9 +2,13 @@ use serde_json::{Map, Value}; use super::Error; use super::common_utils::{has_bearer_auth, has_header, messages_provider_config, string_headers}; -use super::transformation::{AnthropicMessagesProviderConfig, MessagesAuthStrategy}; use super::types::{MessagesRequest, ProviderMessagesRequest}; -use crate::providers::custom_llm_provider::{CustomLlmProvider, get_custom_llm_provider}; +use crate::litellm_core_utils::get_llm_provider_logic::{ + CustomLlmProvider, get_custom_llm_provider, +}; +use crate::llms::base_llm::anthropic_messages::transformation::{ + BaseAnthropicMessagesConfig, MessagesAuthStrategy, +}; pub(super) fn prepare_provider_request( request: MessagesRequest<'_>, @@ -36,14 +40,14 @@ pub(super) fn prepare_provider_request( let typed_request = serde_json::from_value(request.body).map_err(|err| { Error::InvalidRequest(format!("invalid Anthropic messages request: {err}")) })?; - let transformed = config.transform_request(typed_request)?; + let transformed = config.transform_anthropic_messages_request(typed_request)?; let body = serde_json::to_value(transformed).map_err(|err| { Error::InvalidRequest(format!( "failed to serialize Anthropic messages request: {err}" )) })?; - let url = config.complete_url(request.api_base, &model, &env_lookup)?; + let url = config.get_complete_url(request.api_base, &model, &env_lookup)?; Ok(ProviderMessagesRequest { provider: provider.to_string(), @@ -57,7 +61,7 @@ pub(super) fn prepare_provider_request( } fn validate_environment( - config: &dyn AnthropicMessagesProviderConfig, + config: &dyn BaseAnthropicMessagesConfig, extra_headers: Option>, api_key: Option<&str>, env_lookup: &dyn Fn(&str) -> Option, diff --git a/litellm-rust/crates/core/src/messages/types.rs b/litellm-rust/crates/core/src/messages/types.rs index b9f807c29fd..32cf4b29faf 100644 --- a/litellm-rust/crates/core/src/messages/types.rs +++ b/litellm-rust/crates/core/src/messages/types.rs @@ -3,7 +3,7 @@ use std::time::Duration; use serde::{Deserialize, Serialize}; use serde_json::{Map, Value}; -use super::transformation::AnthropicMessagesProviderConfig; +use crate::llms::base_llm::anthropic_messages::transformation::BaseAnthropicMessagesConfig; pub struct MessagesRequest<'a> { pub model: &'a str, @@ -18,7 +18,7 @@ pub struct MessagesRequest<'a> { pub(super) struct ProviderMessagesRequest { pub(super) provider: String, pub(super) model: String, - pub(super) config: &'static dyn AnthropicMessagesProviderConfig, + pub(super) config: &'static dyn BaseAnthropicMessagesConfig, pub(super) url: String, pub(super) body: Value, pub(super) upstream_headers: Vec<(String, String)>, diff --git a/litellm-rust/crates/core/src/ocr/provider_config.rs b/litellm-rust/crates/core/src/ocr/provider_config.rs index fcbea54779f..dcce6258a12 100644 --- a/litellm-rust/crates/core/src/ocr/provider_config.rs +++ b/litellm-rust/crates/core/src/ocr/provider_config.rs @@ -5,16 +5,18 @@ use super::types::{ LiteLLMOcrResponse, OcrCredentialInputs, OcrDocument, PreparedOcrRequest, ResolvedOcrCredentials, }; +use crate::litellm_core_utils::get_llm_provider_logic::{ + CustomLlmProvider, get_custom_llm_provider, +}; 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; +use crate::llms::azure_ai::ocr::document_intelligence::transformation::AzureDocumentIntelligenceOcrConfig; +use crate::llms::azure_ai::ocr::transformation::AzureAiOcrConfig; use crate::llms::base_llm::ocr::transformation::{BaseOcrConfig, OcrResponseContext}; use crate::llms::cohere::ocr::transformation::CohereParseConfig; use crate::llms::mistral::ocr::transformation::MistralOcrConfig; use crate::llms::reducto::ocr::transformation::{ReductoParseLegacyConfig, ReductoParseV3Config}; use crate::llms::vertex_ai::ocr::deepseek_transformation::VertexAIDeepSeekOCRConfig; -use crate::llms::vertex_ai::ocr::transformation::VertexAIOCRConfig; -use crate::providers::custom_llm_provider::{CustomLlmProvider, get_custom_llm_provider}; +use crate::llms::vertex_ai::ocr::transformation::VertexAiOcrConfig; macro_rules! dispatch_config { ($config:expr, $method:ident($($argument:expr),* $(,)?)) => { @@ -27,12 +29,12 @@ macro_rules! dispatch_config { match $config { OcrConfigKind::Cohere => CohereParseConfig.$method($($argument),*)$($suffix)*, OcrConfigKind::Mistral => MistralOcrConfig.$method($($argument),*)$($suffix)*, - OcrConfigKind::AzureAi => AzureAIOCRConfig.$method($($argument),*)$($suffix)*, + OcrConfigKind::AzureAi => AzureAiOcrConfig.$method($($argument),*)$($suffix)*, OcrConfigKind::AzureCohere => AzureAICohereParseConfig.$method($($argument),*)$($suffix)*, - OcrConfigKind::AzureDocumentIntelligence => AzureDocumentIntelligenceOCRConfig.$method($($argument),*)$($suffix)*, + OcrConfigKind::AzureDocumentIntelligence => AzureDocumentIntelligenceOcrConfig.$method($($argument),*)$($suffix)*, OcrConfigKind::ReductoLegacy => ReductoParseLegacyConfig.$method($($argument),*)$($suffix)*, OcrConfigKind::ReductoV3 => ReductoParseV3Config.$method($($argument),*)$($suffix)*, - OcrConfigKind::VertexAi => VertexAIOCRConfig.$method($($argument),*)$($suffix)*, + OcrConfigKind::VertexAi => VertexAiOcrConfig.$method($($argument),*)$($suffix)*, OcrConfigKind::VertexDeepSeek => VertexAIDeepSeekOCRConfig.$method($($argument),*)$($suffix)*, } }; diff --git a/litellm-rust/crates/core/src/providers/anthropic/mod.rs b/litellm-rust/crates/core/src/providers/anthropic/mod.rs deleted file mode 100644 index 0bb20991ff7..00000000000 --- a/litellm-rust/crates/core/src/providers/anthropic/mod.rs +++ /dev/null @@ -1,2 +0,0 @@ -pub mod chat_completions; -pub mod messages; diff --git a/litellm-rust/crates/core/src/providers/bedrock/aws_base.rs b/litellm-rust/crates/core/src/providers/bedrock/aws_base.rs deleted file mode 100644 index b51cef7545c..00000000000 --- a/litellm-rust/crates/core/src/providers/bedrock/aws_base.rs +++ /dev/null @@ -1 +0,0 @@ -pub use litellm_auth_aws::*; diff --git a/litellm-rust/crates/core/src/providers/bedrock/constants.rs b/litellm-rust/crates/core/src/providers/bedrock/constants.rs deleted file mode 100644 index 663f887c1fd..00000000000 --- a/litellm-rust/crates/core/src/providers/bedrock/constants.rs +++ /dev/null @@ -1 +0,0 @@ -pub use litellm_auth_aws::constants::*; diff --git a/litellm-rust/crates/core/src/providers/bedrock/mod.rs b/litellm-rust/crates/core/src/providers/bedrock/mod.rs deleted file mode 100644 index 5c849064989..00000000000 --- a/litellm-rust/crates/core/src/providers/bedrock/mod.rs +++ /dev/null @@ -1,8 +0,0 @@ -//! User-directed exception: this base provider owns AWS auth I/O for parity -//! with Python's `BaseAWSLLM`; the broader core purity guidance is reconciled -//! separately. - -pub mod audio_transcription; -pub mod aws_base; -pub mod chat_completions; -mod constants; diff --git a/litellm-rust/crates/core/src/providers/mod.rs b/litellm-rust/crates/core/src/providers/mod.rs deleted file mode 100644 index 70ca4386fff..00000000000 --- a/litellm-rust/crates/core/src/providers/mod.rs +++ /dev/null @@ -1,5 +0,0 @@ -pub mod anthropic; -pub mod azure_ai; -pub mod bedrock; -pub mod custom_llm_provider; -pub mod openai; diff --git a/litellm-rust/crates/core/tests/vertex_ai_ocr.rs b/litellm-rust/crates/core/tests/vertex_ai_ocr.rs index 1908c7aa347..858fee1ba3e 100644 --- a/litellm-rust/crates/core/tests/vertex_ai_ocr.rs +++ b/litellm-rust/crates/core/tests/vertex_ai_ocr.rs @@ -104,7 +104,7 @@ async fn adapters_build_complete_requests_and_share_mistral_normalization() { 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::llms::vertex_ai::ocr::transformation::VertexAiOcrConfig; use crate::ocr::test_support::ocr_client; let client = ocr_client(); @@ -129,7 +129,7 @@ async fn adapters_build_complete_requests_and_share_mistral_normalization() { .prepare_request(&direct, &client) .await .unwrap(); - let vertex_http = VertexAIOCRConfig + let vertex_http = VertexAiOcrConfig .prepare_request(&vertex, &client) .await .unwrap(); @@ -165,7 +165,7 @@ async fn adapters_build_complete_requests_and_share_mistral_normalization() { ) .unwrap() .into_json(); - let vertex_response = VertexAIOCRConfig + let vertex_response = VertexAiOcrConfig .transform_ocr_response( &vertex.model, &raw, From 370cdaabf9f75a270dc822efd73278ed8ce8742c Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Thu, 17 Sep 2026 08:04:52 -0700 Subject: [PATCH 10/15] encode failing tests --- .../crates/core/src/ocr/provider_config.rs | 14 ++ litellm-rust/crates/core/src/ocr/wire.rs | 29 +++ litellm-rust/crates/core/tests/ocr.rs | 54 +++++ tests/test_litellm_rust/ocr/test_requests.py | 194 +++++++++++++++++- 4 files changed, 290 insertions(+), 1 deletion(-) diff --git a/litellm-rust/crates/core/src/ocr/provider_config.rs b/litellm-rust/crates/core/src/ocr/provider_config.rs index dcce6258a12..b798fd95841 100644 --- a/litellm-rust/crates/core/src/ocr/provider_config.rs +++ b/litellm-rust/crates/core/src/ocr/provider_config.rs @@ -419,4 +419,18 @@ mod tests { model.split_once('/').unwrap().1 ); } + + #[rstest] + #[case::prefix("not_a_provider/model", None)] + #[case::explicit("model", Some("not_a_provider"))] + fn ocr_contract_unknown_provider_is_bad_request( + #[case] model: &str, + #[case] provider: Option<&str>, + ) { + let error = resolve_provider_config(model, provider).unwrap_err(); + assert!( + matches!(&error, crate::ocr::Error::InvalidProvider(provider) if provider == "not_a_provider") + ); + assert_eq!(error.http_status_code(), Some(400)); + } } diff --git a/litellm-rust/crates/core/src/ocr/wire.rs b/litellm-rust/crates/core/src/ocr/wire.rs index b05f388a277..b2f07caa754 100644 --- a/litellm-rust/crates/core/src/ocr/wire.rs +++ b/litellm-rust/crates/core/src/ocr/wire.rs @@ -106,6 +106,35 @@ pub fn decode_document(value: Value) -> Result { #[cfg(test)] mod tests { use super::*; + use rstest::rstest; + use serde_json::json; + + #[rstest] + #[case::omitted(json!({"type":"document_url", "document_url":"https://example.com/a.pdf"}))] + #[case::null(json!({"type":"document_url", "document_url":"https://example.com/a.pdf", "document_name":null}))] + fn ocr_contract_optional_document_name(#[case] document: Value) { + let decoded = decode_document(document).unwrap(); + assert_eq!(decoded.source(), "https://example.com/a.pdf"); + } + + #[rstest] + #[case::non_object(json!([]), "document")] + #[case::missing_type(json!({"document_url":"https://example.com/a.pdf"}), "document")] + #[case::unsupported_type(json!({"type":"text"}), "document")] + #[case::missing_document_url(json!({"type":"document_url"}), "Document URL")] + #[case::missing_image_url(json!({"type":"image_url"}), "Document URL")] + fn ocr_contract_malformed_document_is_bad_request( + #[case] document: Value, + #[case] field: &str, + ) { + let error = decode_document(document).unwrap_err(); + assert!(matches!( + error, + Error::RequestField { .. } | Error::MissingDocumentUrl + )); + assert_eq!(error.http_status_code(), Some(400)); + assert!(error.to_string().contains(field)); + } #[test] fn option_projection_is_provider_specific_and_excludes_opaque_fields() { diff --git a/litellm-rust/crates/core/tests/ocr.rs b/litellm-rust/crates/core/tests/ocr.rs index 480774d1ad1..c094000ee06 100644 --- a/litellm-rust/crates/core/tests/ocr.rs +++ b/litellm-rust/crates/core/tests/ocr.rs @@ -1,5 +1,6 @@ use std::sync::{Arc, Mutex}; +use rstest::rstest; use serde_json::{Value, json}; use super::OcrClient; @@ -15,6 +16,59 @@ use super::{ }; use crate::call_lifecycle::{CallLifecycleContext, CallLifecycleTiming}; +#[rstest] +#[case::mistral("mistral/model", json!({}))] +#[case::vertex("vertex_ai/mistral-ocr-latest", json!({"vertex_project":"test-project", "vertex_location":"us-central1"}))] +#[tokio::test] +async fn ocr_contract_upstream_error_preserves_status_body_and_headers( + #[case] model: &str, + #[case] options: Value, +) { + let payload = json!({"message": format!("{} END-OF-PROVIDER-BODY", "x".repeat(4096))}); + let expected_body = serde_json::to_string(&payload).unwrap(); + let (base, seen, server) = mock_server(vec![MockResponse { + status: 422, + headers: vec![ + ("Retry-After", "17".into()), + ("X-Request-ID", "request-123".into()), + ("X-Future-Header", "retained".into()), + ], + body: payload, + }]) + .await; + let error = perform_ocr(wire_request(model, &base, options)) + .await + .unwrap_err(); + server.await.unwrap(); + assert_eq!(seen.lock().unwrap().len(), 1); + let super::Error::Provider { + status, + body, + headers, + } = error + else { + panic!("expected provider error, got {error:?}"); + }; + assert_eq!(status, 422); + for (name, value) in [ + ("retry-after", "17"), + ("x-request-id", "request-123"), + ("x-future-header", "retained"), + ] { + assert!( + headers + .iter() + .any(|(key, actual)| key.eq_ignore_ascii_case(name) && actual == value) + ); + } + assert_eq!( + body.len(), + expected_body.len(), + "provider error body was truncated" + ); + assert_eq!(body, expected_body); +} + #[test] fn request_boundary_selects_mistral_and_rejects_unknown_providers() { let request = OcrWireRequest { diff --git a/tests/test_litellm_rust/ocr/test_requests.py b/tests/test_litellm_rust/ocr/test_requests.py index 58bb6a77537..3e95258fb36 100644 --- a/tests/test_litellm_rust/ocr/test_requests.py +++ b/tests/test_litellm_rust/ocr/test_requests.py @@ -1,7 +1,10 @@ +import json from pathlib import Path from typing import Final +import httpx import pytest +from pydantic import JsonValue import litellm from litellm.llms.base_llm.ocr.transformation import OCRResponse @@ -17,6 +20,193 @@ from tests.test_litellm_rust.support.requests import ( pytestmark = pytest.mark.requires_rust_extension +@pytest.fixture(params=[False, True], ids=["python", "rust"]) +def ocr_backend(request: pytest.FixtureRequest, monkeypatch: pytest.MonkeyPatch) -> bool: + enabled: Final = bool(request.param) + monkeypatch.setenv("LITELLM_RUST", "1" if enabled else "0") + return enabled + + +@pytest.mark.asyncio +@pytest.mark.parametrize("asynchronous", [False, True], ids=["sync", "async"]) +async def test_ocr_contract_upstream_status( + ocr_server: RecordingServer, + ocr_backend: bool, + asynchronous: bool, +) -> None: + upstream: Final = ResponseSpec(body={"detail": "invalid provider option"}, status=422) + ocr_server.enqueue(upstream) + arguments: Final = { + "model": "vertex_ai/mistral-ocr-latest", + "vertex_project": "test-project", + "vertex_location": "us-central1", + "num_retries": 0, + } + with pytest.raises(litellm.BadRequestError) as caught: + if asynchronous: + await call_native_aocr(ocr_server, **arguments) + else: + call_native_ocr(ocr_server, **arguments) + assert caught.value.status_code == upstream.status + assert caught.value.response.status_code == upstream.status + + +@pytest.mark.asyncio +@pytest.mark.parametrize("asynchronous", [False, True], ids=["sync", "async"]) +@pytest.mark.parametrize("preserved", ["body", "headers"]) +async def test_ocr_contract_provider_error_details( + ocr_server: RecordingServer, + ocr_backend: bool, + asynchronous: bool, + preserved: str, +) -> None: + payload: Final = {"message": "rate limited"} + headers: Final = {"Retry-After": "17", "X-Request-ID": "ocr-request-123", "X-Future-Header": "retained"} + ocr_server.enqueue(ResponseSpec(body=payload, status=429, headers=headers)) + with pytest.raises(litellm.RateLimitError) as caught: + if asynchronous: + await call_native_aocr(ocr_server, num_retries=0) + else: + call_native_ocr(ocr_server, num_retries=0) + response: Final = caught.value.response + assert isinstance(response, httpx.Response) + if preserved == "body": + assert response.content == json.dumps(payload).encode() + else: + for name, value in headers.items(): + assert response.headers.get(name.lower()) == value + assert response.headers.get(name.upper()) == value + + +@pytest.mark.asyncio +@pytest.mark.parametrize("asynchronous", [False, True], ids=["sync", "async"]) +async def test_ocr_contract_invalid_response_format( + ocr_server: RecordingServer, + ocr_backend: bool, + asynchronous: bool, +) -> None: + ocr_server.expected_requests = 0 + with pytest.raises(litellm.UnsupportedParamsError) as caught: + if asynchronous: + await call_native_aocr(ocr_server, req_format="bogus", num_retries=0) + else: + call_native_ocr(ocr_server, req_format="bogus", num_retries=0) + assert caught.value.status_code == 400 + for value in ("req_format", "bogus", "native", "litellm"): + assert value in str(caught.value) + assert ocr_server.requests == [] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("asynchronous", [False, True], ids=["sync", "async"]) +@pytest.mark.parametrize( + "document,field", + [ + ([], "document"), + ({"document_url": "https://example.com/a.pdf"}, "type"), + ({"type": "text"}, "type"), + ], +) +async def test_ocr_contract_malformed_document_is_actionable( + ocr_server: RecordingServer, + ocr_backend: bool, + asynchronous: bool, + document: JsonValue, + field: str, +) -> None: + ocr_server.expected_requests = None + with pytest.raises(litellm.BadRequestError) as caught: + if asynchronous: + await call_native_aocr(ocr_server, document=document, num_retries=0) + else: + call_native_ocr(ocr_server, document=document, num_retries=0) + assert caught.value.status_code == 400 + assert field.lower() in str(caught.value).lower() + assert "NoneType: None" not in str(caught.value) + assert "indices must be" not in str(caught.value) + assert ocr_server.requests == [] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("asynchronous", [False, True], ids=["sync", "async"]) +@pytest.mark.parametrize("option,value,field", [("pages", [-1], "pages"), ("features", [1], "features")]) +async def test_ocr_contract_azure_invalid_options_are_bad_requests( + ocr_server: RecordingServer, + ocr_backend: bool, + asynchronous: bool, + option: str, + value: JsonValue, + field: str, +) -> None: + ocr_server.expected_requests = 0 + arguments: Final = {"model": "azure_ai/doc-intelligence/prebuilt-read", option: value, "num_retries": 0} + with pytest.raises(litellm.BadRequestError) as caught: + if asynchronous: + await call_native_aocr(ocr_server, **arguments) + else: + call_native_ocr(ocr_server, **arguments) + assert caught.value.status_code == 400 + assert field in str(caught.value) + assert ocr_server.requests == [] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("asynchronous", [False, True], ids=["sync", "async"]) +@pytest.mark.parametrize("model", ["mistral/mistral-ocr-latest", "azure_ai/mistral-ocr-latest", "reducto/parse-v3"]) +async def test_ocr_contract_native_format_supported( + ocr_server: RecordingServer, + ocr_backend: bool, + asynchronous: bool, + model: str, +) -> None: + ocr_server.expected_requests = None + payload: Final = ( + {"result": {"chunks": [{"content": "native OCR response"}]}, "usage": {"num_pages": 1}} + if model.startswith("reducto/") + else OCR_RESPONSE + ) + ocr_server.default_response = ResponseSpec(body=payload) + arguments: Final = { + "model": model, + "req_format": "native", + "num_retries": 0, + "document": {"type": "document_url", "document_url": "reducto://ready.pdf"} + if model.startswith("reducto/") + else OCR_DOCUMENT, + } + response: Final = ( + await call_native_aocr(ocr_server, **arguments) if asynchronous else call_native_ocr(ocr_server, **arguments) + ) + assert response.pages[0].markdown == "native OCR response" + assert response.get_provider_native_response() == payload + assert len(ocr_server.requests) == 1 + if ocr_backend: + assert_native_request(ocr_server) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("asynchronous", [False, True], ids=["sync", "async"]) +async def test_ocr_contract_unknown_reducto_model_reaches_provider( + ocr_server: RecordingServer, + ocr_backend: bool, + asynchronous: bool, +) -> None: + ocr_server.default_response = ResponseSpec(body={"result": {"chunks": [{"content": "future model response"}]}}) + arguments: Final = { + "model": "reducto/future-parse-model", + "document": {"type": "document_url", "document_url": "reducto://ready.pdf"}, + "num_retries": 0, + } + response: Final = ( + await call_native_aocr(ocr_server, **arguments) if asynchronous else call_native_ocr(ocr_server, **arguments) + ) + assert response.model == "future-parse-model" + assert response.pages[0].markdown == "future model response" + assert len(ocr_server.requests) == 1 + assert ocr_server.requests[0].path == "/parse" + assert ocr_server.requests[0].body == {"input": "reducto://ready.pdf"} + + @pytest.mark.asyncio @pytest.mark.parametrize("asynchronous", [False, True], ids=["sync", "async"]) async def test_native_azure_ocr_uses_token_provider_result_as_bearer_token( @@ -595,7 +785,9 @@ def test_native_file_preparation_rejects_unsupported_reader_results(ocr_server: @pytest.mark.parametrize("kind", ["bytes", "path", "reader"]) -def test_native_file_preparation_rejects_oversized_input(ocr_server: RecordingServer, kind: str, tmp_path: Path) -> None: +def test_native_file_preparation_rejects_oversized_input( + ocr_server: RecordingServer, kind: str, tmp_path: Path +) -> None: ocr_server.expected_requests = 0 limit: Final = 50 * 1024 * 1024 path: Final = tmp_path / "large.pdf" From f1ea94fee70dbaa85ffdbb7bc52010814e9003ce Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Thu, 17 Sep 2026 08:55:59 -0700 Subject: [PATCH 11/15] make test pass --- litellm-rust/crates/core/src/ocr/client.rs | 9 +-- litellm-rust/crates/core/src/ocr/document.rs | 4 +- litellm-rust/crates/core/src/ocr/error.rs | 2 +- litellm-rust/crates/core/src/ocr/types.rs | 61 ++++++------------ litellm-rust/crates/core/tests/ocr.rs | 45 +++++++------- .../python-bridge/src/routes/ocr/errors.rs | 58 +++++++++++++---- .../python-bridge/src/routes/ocr/project.rs | 24 +++++-- litellm/exceptions.py | 1 + litellm/llms/custom_httpx/llm_http_handler.py | 27 ++++++-- litellm/ocr/legacy.py | 62 ++++++++++--------- litellm/rust_bridge/ocr_lifecycle.py | 15 ++++- 11 files changed, 182 insertions(+), 126 deletions(-) diff --git a/litellm-rust/crates/core/src/ocr/client.rs b/litellm-rust/crates/core/src/ocr/client.rs index 8dba37bb00b..18d0f3b7498 100644 --- a/litellm-rust/crates/core/src/ocr/client.rs +++ b/litellm-rust/crates/core/src/ocr/client.rs @@ -133,14 +133,9 @@ pub async fn read_json_response( pub(crate) async fn read_response_bytes( mut response: reqwest::Response, - max_response_bytes: usize, + limit: usize, ) -> Result { let status = response.status(); - let limit = if status.is_success() { - max_response_bytes - } else { - max_response_bytes.min(4 * (crate::constants::UPSTREAM_ERROR_BODY_MAX_CHARS + 1)) - }; if status.is_success() && response .content_length() @@ -162,7 +157,7 @@ pub(crate) async fn read_response_bytes( if !status.is_success() { return Err(crate::transport::Error::Http { status: status.as_u16(), - body: crate::http_utils::truncate_error_body(&String::from_utf8_lossy(&bytes)), + body: String::from_utf8_lossy(&bytes).into_owned(), } .into()); } diff --git a/litellm-rust/crates/core/src/ocr/document.rs b/litellm-rust/crates/core/src/ocr/document.rs index c3ffac701b3..5d1f0dd9ab4 100644 --- a/litellm-rust/crates/core/src/ocr/document.rs +++ b/litellm-rust/crates/core/src/ocr/document.rs @@ -429,7 +429,7 @@ mod tests { client.document_fetcher(), OcrDocument::ImageUrl { image_url: format!("http://{address}/image"), - extra_fields: Map::from_iter([("detail".into(), "high".into())]), + extra_fields: Map::from_iter([("detail".into(), Some("high".into()))]), }, &OcrConnection::default(), ) @@ -441,7 +441,7 @@ mod tests { converted, OcrDocument::ImageUrl { image_url: "data:image/png;base64,YWJj".into(), - extra_fields: Map::from_iter([("detail".into(), "high".into())]), + extra_fields: Map::from_iter([("detail".into(), Some("high".into()))]), } ); assert!(!request.to_ascii_lowercase().contains("authorization")); diff --git a/litellm-rust/crates/core/src/ocr/error.rs b/litellm-rust/crates/core/src/ocr/error.rs index 7685875709e..4906b5515b9 100644 --- a/litellm-rust/crates/core/src/ocr/error.rs +++ b/litellm-rust/crates/core/src/ocr/error.rs @@ -113,7 +113,6 @@ impl From for Error { impl Error { pub fn http_status_code(&self) -> Option { match self { - Self::MissingDocumentUrl => Some(500), Self::Provider { status, .. } | Self::Transport(crate::transport::Error::Http { status, .. }) => Some(*status), error if error.is_request() => Some(400), @@ -142,6 +141,7 @@ impl Error { | Self::Features | Self::DotModel | Self::InvalidRequest(_) + | Self::InvalidProvider(_) | Self::Params(_) | Self::Headers(_) ) diff --git a/litellm-rust/crates/core/src/ocr/types.rs b/litellm-rust/crates/core/src/ocr/types.rs index facfd04fe8e..fe7e41a6128 100644 --- a/litellm-rust/crates/core/src/ocr/types.rs +++ b/litellm-rust/crates/core/src/ocr/types.rs @@ -22,13 +22,13 @@ pub enum OcrDocument { DocumentUrl { document_url: String, #[serde(flatten)] - extra_fields: BTreeMap, + extra_fields: BTreeMap>, }, #[serde(rename = "image_url")] ImageUrl { image_url: String, #[serde(flatten)] - extra_fields: BTreeMap, + extra_fields: BTreeMap>, }, } @@ -720,45 +720,24 @@ mod tests { } } - #[test] - fn document_variants_preserve_provider_fields_when_rewriting_sources() { - for (value, original, replacement, expected) in [ - ( - json!({ - "type":"document_url", - "document_url":"https://example.com/input.pdf", - "document_name":"input.pdf" - }), - "https://example.com/input.pdf", - "data:application/pdf;base64,AA==", - json!({ - "type":"document_url", - "document_url":"data:application/pdf;base64,AA==", - "document_name":"input.pdf" - }), - ), - ( - json!({ - "type":"image_url", - "image_url":"https://example.com/input.png", - "detail":"high" - }), - "https://example.com/input.png", - "data:image/png;base64,AA==", - json!({ - "type":"image_url", - "image_url":"data:image/png;base64,AA==", - "detail":"high" - }), - ), - ] { - let document: OcrDocument = serde_json::from_value(value).unwrap(); - assert_eq!(document.source(), original); - assert_eq!( - serde_json::to_value(document.with_source(replacement.into())).unwrap(), - expected - ); - } + #[rstest::rstest] + #[case::document_url("document_url", "document_name", "application/pdf")] + #[case::image_url("image_url", "detail", "image/png")] + fn document_variants_preserve_provider_fields_when_rewriting_sources( + #[case] kind: &str, + #[case] field: &str, + #[case] mime_type: &str, + #[values(json!("kept"), Value::Null)] extra: Value, + ) { + let original = "https://example.com/input"; + let replacement = format!("data:{mime_type};base64,AA=="); + let document: OcrDocument = + serde_json::from_value(json!({"type": kind, kind: original, field: extra})).unwrap(); + assert_eq!(document.source(), original); + assert_eq!( + serde_json::to_value(document.with_source(replacement.clone())).unwrap(), + json!({"type": kind, kind: replacement, field: extra}) + ); } #[test] diff --git a/litellm-rust/crates/core/tests/ocr.rs b/litellm-rust/crates/core/tests/ocr.rs index c094000ee06..58762fb4d93 100644 --- a/litellm-rust/crates/core/tests/ocr.rs +++ b/litellm-rust/crates/core/tests/ocr.rs @@ -956,32 +956,29 @@ async fn response_limit_accepts_exact_size_and_rejects_declared_and_chunked_over } } +#[rstest] +#[case::declared("Content-Length: 1000000")] +#[case::chunked("Transfer-Encoding: chunked")] #[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 { - super::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}"), +async fn oversized_error_retains_http_status_and_bounded_diagnostics_without_draining( + #[case] headers: &str, +) { + let prefix = "x".repeat(4096); + 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(), prefix.len()) + .await + .unwrap_err(); + match error { + super::Error::Transport(crate::transport::Error::Http { status, body }) => { + assert_eq!(status, 429); + assert_eq!(body, prefix); } + error => panic!("unexpected error: {error}"), } } 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 d943a053a61..02d2ccbdeea 100644 --- a/litellm-rust/crates/python-bridge/src/routes/ocr/errors.rs +++ b/litellm-rust/crates/python-bridge/src/routes/ocr/errors.rs @@ -1,25 +1,59 @@ use litellm_core::ocr::Error; use pyo3::exceptions::{PyFileNotFoundError, PyOSError}; use pyo3::prelude::*; +use pyo3::types::PyDict; use crate::errors::{RustUpstreamError, core_error_to_pyerr}; pub(super) fn to_pyerr(error: Error) -> PyErr { let status = error.http_status_code(); - let mapped = match error { - Error::Provider { status, body, .. } - | Error::Transport(litellm_core::transport::Error::Http { status, body }) => { - RustUpstreamError::new_err((status, body)) - } - Error::FileRead { path, source } if source.kind() == std::io::ErrorKind::NotFound => { - PyFileNotFoundError::new_err(format!("File not found: {}", path.display())) - } - Error::FileRead { source, .. } => PyOSError::new_err(source.to_string()), - other => core_error_to_pyerr(other.into()), - }; + let mapped = Python::attach(|py| -> PyResult { + Ok(match error { + Error::Provider { + status, + body, + headers, + } => upstream_error(py, status, body, headers)?, + Error::Transport(litellm_core::transport::Error::Http { status, body }) => { + upstream_error(py, status, body, Vec::new())? + } + Error::RequestFormat => { + let error = core_error_to_pyerr(Error::RequestFormat.into()); + error + .value(py) + .setattr("ocr_request_format_error", true) + .ok(); + error + } + Error::FileRead { path, source } if source.kind() == std::io::ErrorKind::NotFound => { + PyFileNotFoundError::new_err(format!("File not found: {}", path.display())) + } + Error::FileRead { source, .. } => PyOSError::new_err(source.to_string()), + other => core_error_to_pyerr(other.into()), + }) + }) + .unwrap_or_else(|error| error); attach_status(mapped, status) } +fn upstream_error( + py: Python<'_>, + status: u16, + body: String, + headers: Vec<(String, String)>, +) -> PyResult { + let kwargs = PyDict::new(py); + kwargs.set_item("content", &body)?; + kwargs.set_item("headers", headers)?; + let response = py + .import("httpx")? + .getattr("Response")? + .call((status,), Some(&kwargs))?; + let error = RustUpstreamError::new_err((status, body)); + error.value(py).setattr("response", response)?; + Ok(error) +} + fn attach_status(error: PyErr, status: Option) -> PyErr { if let Some(status) = status { Python::attach(|py| { @@ -50,7 +84,7 @@ mod tests { .unwrap() .extract::() .unwrap(), - 500 + 400 ); let mapped = to_pyerr(Error::Provider { status: 429, diff --git a/litellm-rust/crates/python-bridge/src/routes/ocr/project.rs b/litellm-rust/crates/python-bridge/src/routes/ocr/project.rs index 3076895c1c4..e2fe7ae4109 100644 --- a/litellm-rust/crates/python-bridge/src/routes/ocr/project.rs +++ b/litellm-rust/crates/python-bridge/src/routes/ocr/project.rs @@ -88,7 +88,21 @@ enum ProjectedDocument { impl ProjectedDocument { fn project(document: &Bound<'_, PyAny>) -> PyResult { - let kind: String = document.get_item("type")?.extract()?; + let kind: String = document + .get_item("type") + .and_then(|value| value.extract()) + .map_err(|error| { + let py = document.py(); + if error.is_instance_of::(py) + || error.is_instance_of::(py) + { + ocr_error_to_pyerr(litellm_core::ocr::Error::RequestField { + path: "document.type".into(), + }) + } else { + error + } + })?; if kind != "file" { return Ok(Self::Other { wire: from_py(document)?, @@ -185,7 +199,7 @@ pub(super) fn admitted_call(outcome: NativeOutcome) -> PyResult(py) + .is_instance_of::(py) ); let non_string = py.eval(c"{'type': 1}", None, None).unwrap(); assert!( project_document(&non_string) .unwrap_err() - .is_instance_of::(py) + .is_instance_of::(py) ); let locals = eval( diff --git a/litellm/exceptions.py b/litellm/exceptions.py index 3f22a4b2dcd..23f9c1f2a12 100644 --- a/litellm/exceptions.py +++ b/litellm/exceptions.py @@ -500,6 +500,7 @@ class RateLimitError(openai.RateLimitError): self.response = httpx.Response( status_code=429, headers=_response_headers, + content=response.content if response is not None else None, request=httpx.Request( method="POST", url=" https://cloud.google.com/vertex-ai/", diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index 2fe4130a310..fd941c0d8bc 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -59,7 +59,7 @@ from litellm.llms.base_llm.image_edit.transformation import BaseImageEditConfig from litellm.llms.base_llm.image_generation.transformation import ( BaseImageGenerationConfig, ) -from litellm.llms.base_llm.ocr.transformation import BaseOCRConfig, OCRResponse +from litellm.llms.base_llm.ocr.transformation import OCR_REQUEST_FORMAT_PARAM, BaseOCRConfig, OCRResponse from litellm.llms.base_llm.realtime.http_transformation import BaseRealtimeHTTPConfig from litellm.llms.base_llm.realtime.transformation import BaseRealtimeConfig from litellm.llms.base_llm.rerank.transformation import BaseRerankConfig @@ -1568,7 +1568,7 @@ class BaseLLMHTTPHandler: transformed_result: Final = provider_config.transform_ocr_request( model=model, document=document, - optional_params=optional_params, + optional_params={key: value for key, value in optional_params.items() if key != OCR_REQUEST_FORMAT_PARAM}, headers=headers, api_key=api_key, api_base=api_base, @@ -1634,7 +1634,7 @@ class BaseLLMHTTPHandler: transformed_result: Final = await provider_config.async_transform_ocr_request( model=model, document=document, - optional_params=optional_params, + optional_params={key: value for key, value in optional_params.items() if key != OCR_REQUEST_FORMAT_PARAM}, headers=headers, api_key=api_key, api_base=api_base, @@ -1672,12 +1672,26 @@ class BaseLLMHTTPHandler: optional_params: Mapping[str, object], ) -> OCRResponse: """Shared logic for transforming OCR responses.""" - return provider_config.transform_ocr_response( + normalized: Final = provider_config.transform_ocr_response( model=model, raw_response=response, logging_obj=logging_obj, optional_params=optional_params, ) + return self._finalize_ocr_response(normalized, response, optional_params) + + @staticmethod + def _finalize_ocr_response( + normalized: OCRResponse, + response: httpx.Response, + optional_params: Mapping[str, object], + ) -> OCRResponse: + if ( + optional_params.get(OCR_REQUEST_FORMAT_PARAM) == "native" + and normalized.get_provider_native_response() is None + ): + normalized.set_provider_native_response(response.json()) + return normalized def ocr( self, @@ -1823,12 +1837,13 @@ class BaseLLMHTTPHandler: ) # Use async response transform for async operations - return await provider_config.async_transform_ocr_response( + normalized: Final = await provider_config.async_transform_ocr_response( model=model, raw_response=response, logging_obj=logging_obj, optional_params=optional_params, ) + return self._finalize_ocr_response(normalized, response, optional_params) def search( self, @@ -6157,6 +6172,8 @@ class BaseLLMHTTPHandler: status_code=status_code, headers=error_headers, ) + if isinstance(provider_config, BaseOCRConfig) and isinstance(error_response, httpx.Response): + provider_error.response = error_response if not isinstance(received_status_code, int): provider_error.status_code_is_synthesized = True raise provider_error diff --git a/litellm/ocr/legacy.py b/litellm/ocr/legacy.py index f0cf6cc82cc..1c9e1c2c28f 100644 --- a/litellm/ocr/legacy.py +++ b/litellm/ocr/legacy.py @@ -70,16 +70,27 @@ def _prepare_ocr_request( ) if not isinstance(document, dict): - raise ValueError(f"document must be a dict with 'type' and URL/file field, got {type(document)}") + raise litellm.BadRequestError( + message="document must be a dict with 'type' and URL/file field", + model=model, + llm_provider=custom_llm_provider or "", + ) - doc_type = document.get("type") + normalized_document: Final = ( + convert_file_document_to_url_document(document) if document.get("type") == "file" else document + ) + doc_type: Final = normalized_document.get("type") - if doc_type == "file": - document = convert_file_document_to_url_document(document) - doc_type = document.get("type") - - if doc_type not in ["document_url", "image_url"]: - raise ValueError(f"Invalid document type: {doc_type}. Must be 'document_url', 'image_url', or 'file'") + if doc_type not in ("document_url", "image_url"): + raise litellm.BadRequestError( + message=f"Invalid document type: {doc_type}. Must be 'document_url', 'image_url', or 'file'", + model=model, + llm_provider=custom_llm_provider or "", + ) + if not normalized_document.get(doc_type): + raise litellm.BadRequestError( + message="Document URL is required", model=model, llm_provider=custom_llm_provider or "" + ) ( model, @@ -116,31 +127,26 @@ def _prepare_ocr_request( requested_format: Final = kwargs.get(OCR_REQUEST_FORMAT_PARAM) if requested_format is not None: try: - parsed_format: Final = parse_ocr_request_format(requested_format) + parse_ocr_request_format(requested_format) except ValueError as e: raise litellm.exceptions.UnsupportedParamsError( message=f"{e}", model=model, llm_provider=custom_llm_provider ) from e - if OCR_REQUEST_FORMAT_PARAM not in supported_params and parsed_format == "native": - raise litellm.exceptions.UnsupportedParamsError( - message=( - f"`{OCR_REQUEST_FORMAT_PARAM}='native'` is not supported for provider: {custom_llm_provider}, " - f"model: {model}" - ), - model=model, - llm_provider=custom_llm_provider, - ) - non_default_params: Final = {} - for param in supported_params: - if param in kwargs: - non_default_params[param] = kwargs.pop(param) + non_default_params: Final = {param: kwargs.pop(param) for param in supported_params if param in kwargs} - optional_params: Final = ocr_provider_config.map_ocr_params( - non_default_params=non_default_params, - optional_params={}, - model=model, - ) + try: + mapped_params: Final = ocr_provider_config.map_ocr_params( + non_default_params=non_default_params, + optional_params={}, + model=model, + ) + except ValueError as error: + raise litellm.BadRequestError(message=str(error), model=model, llm_provider=custom_llm_provider) from error + optional_params: Final = { + **mapped_params, + **({OCR_REQUEST_FORMAT_PARAM: requested_format} if requested_format is not None else {}), + } verbose_logger.debug("OCR optional_params after mapping: %s", optional_params) @@ -160,7 +166,7 @@ def _prepare_ocr_request( return _PreparedOCRRequest( model=model, - document=document, + document=normalized_document, api_key=resolved_api_key, api_base=resolved_api_base, custom_llm_provider=custom_llm_provider, diff --git a/litellm/rust_bridge/ocr_lifecycle.py b/litellm/rust_bridge/ocr_lifecycle.py index 5ca584e1c11..1958fdf8cf3 100644 --- a/litellm/rust_bridge/ocr_lifecycle.py +++ b/litellm/rust_bridge/ocr_lifecycle.py @@ -3,6 +3,8 @@ from __future__ import annotations from collections.abc import Awaitable, Mapping, Sequence from typing import Final, Protocol, cast # noqa: TID251 # validates dynamically loaded native callables +import httpx + import litellm from litellm.llms.base_llm.ocr.transformation import OCRResponse from litellm.rust_bridge.bindings import NativeBinding @@ -51,17 +53,28 @@ def arguments(request: LiteLLMOcrRequest) -> Mapping[str, object]: def map_failure(error: Exception, request: LiteLLMOcrRequest, request_provider: str) -> Exception: + model: Final = request.model.removeprefix(f"{request_provider}/") + if getattr(error, "ocr_request_format_error", False): + return litellm.UnsupportedParamsError( + message=f"Invalid `req_format`: {request.kwargs.get('req_format')!r}. Expected 'native' or 'litellm'.", + model=model, + llm_provider=request_provider, + ) mapper: Final = cast( # cast-ok: bounded adapter for the legacy public exception mapper ExceptionMapper, litellm.exception_type ) try: return mapper( - model=request.model.removeprefix(f"{request_provider}/"), + model=model, custom_llm_provider=request_provider, original_exception=error, completion_kwargs=dict(arguments(request)), # mutable-ok: exception mapper requires owned kwargs extra_kwargs=dict(request.kwargs), # mutable-ok: exception mapper requires owned kwargs ) except Exception as public_error: + response: Final = getattr(error, "response", None) + if isinstance(response, httpx.Response): + public_error.response = response + public_error.status_code = response.status_code public_error.__context__ = error return public_error From f91d1f7ea170ed90c35a3f909891030a6c879904 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Thu, 17 Sep 2026 10:30:32 -0700 Subject: [PATCH 12/15] fix wrong assertion --- .../test_model_access_group_e2e.py | 31 ++++++------------- 1 file changed, 10 insertions(+), 21 deletions(-) diff --git a/tests/e2e/access_control/test_model_access_group_e2e.py b/tests/e2e/access_control/test_model_access_group_e2e.py index 5cc062ea096..6dab51805fb 100644 --- a/tests/e2e/access_control/test_model_access_group_e2e.py +++ b/tests/e2e/access_control/test_model_access_group_e2e.py @@ -23,7 +23,7 @@ from access_control_client import ( MODEL_ACCESS_DENIED_MARKER, TEAM_MODEL_ACCESS_DENIED_MARKER, ) -from e2e_config import unique_marker +from e2e_config import settle_propagation, unique_marker from lifecycle import ResourceManager from models import ( ChatResponse, @@ -31,6 +31,7 @@ from models import ( LiteLLMParamsBody, ModelInfoBody, ModelNewBody, + TeamInfoResponse, ) pytestmark = pytest.mark.e2e @@ -111,24 +112,6 @@ def _await_group_members(client: AccessControlClient, access_group: str, expecte ) -def _await_team_allowlist(client: AccessControlClient, grant_key: str, access_group: str) -> None: - """Registering a team-scoped deployment appends its public name to the team's - allow-list, and a wildcard sitting there directly would grant the model under test - on its own. Poll a denial until the message enumerates the allow-list the test - means to exercise: the group, and nothing else.""" - allowlist: Final = f"models=['{access_group}']" - deadline = time.monotonic() + client.proxy.poll_timeout - body = "" - while time.monotonic() < deadline: - body = client.chat_status( - grant_key, UNCOVERED_OPENAI_MODEL, f"{PROMPT} {unique_marker()}", MAX_COMPLETION_TOKENS - ).body - if allowlist in body: - return - time.sleep(client.proxy.poll_interval) - pytest.fail(f"the team's allow-list never settled to {allowlist}; last denial read {body[:300]}") - - @pytest.fixture(scope="module") def grouped(client: AccessControlClient) -> Iterator[GroupedDeployments]: marker: Final = unique_marker() @@ -172,9 +155,15 @@ def team_grant(client: AccessControlClient) -> Iterator[TeamGrant]: ), listed_for=key, ) - client.set_team_models(team_id, team_alias, [access_group]) try: - _await_team_allowlist(client, key, access_group) + client.set_team_models(team_id, team_alias, [access_group]) + written_at: Final = time.monotonic() + _ = client.proxy.read_body_back_everywhere( + f"/team/info?team_id={team_id}", + TeamInfoResponse, + settled=lambda response: response.team_id == team_id and response.team_info.models == [access_group], + ) + settle_propagation(written_at) yield TeamGrant(access_group=access_group, team_id=team_id, key=key) finally: client.proxy.delete_model(model_id) From 4ecc55ec704db85a90d95fad1477382144414e8f Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Thu, 17 Sep 2026 17:34:18 +0000 Subject: [PATCH 13/15] fix(ocr): build upstream httpx response in Python and satisfy PT012 The Rust bridge imported httpx to construct the provider error response, which fails in the isolated wheel check where httpx is absent. Rust now raises RustUpstreamError with a headers attribute and the Python lifecycle wraps it in a typed UpstreamFailure carrying the httpx.Response before legacy mapping. Test helpers gained call_native so pytest.raises blocks hold a single call Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../python-bridge/src/routes/ocr/errors.rs | 18 +++++------ litellm/rust_bridge/ocr_lifecycle.py | 32 ++++++++++++++++--- tests/test_litellm_rust/ocr/test_requests.py | 26 ++++----------- tests/test_litellm_rust/support/requests.py | 4 +++ 4 files changed, 45 insertions(+), 35 deletions(-) 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 02d2ccbdeea..9bd29ce601f 100644 --- a/litellm-rust/crates/python-bridge/src/routes/ocr/errors.rs +++ b/litellm-rust/crates/python-bridge/src/routes/ocr/errors.rs @@ -1,7 +1,6 @@ use litellm_core::ocr::Error; use pyo3::exceptions::{PyFileNotFoundError, PyOSError}; use pyo3::prelude::*; -use pyo3::types::PyDict; use crate::errors::{RustUpstreamError, core_error_to_pyerr}; @@ -42,15 +41,8 @@ fn upstream_error( body: String, headers: Vec<(String, String)>, ) -> PyResult { - let kwargs = PyDict::new(py); - kwargs.set_item("content", &body)?; - kwargs.set_item("headers", headers)?; - let response = py - .import("httpx")? - .getattr("Response")? - .call((status,), Some(&kwargs))?; let error = RustUpstreamError::new_err((status, body)); - error.value(py).setattr("response", response)?; + error.value(py).setattr("headers", headers)?; Ok(error) } @@ -89,9 +81,15 @@ mod tests { let mapped = to_pyerr(Error::Provider { status: 429, body: r#"{"message":"rate limited"}"#.to_string(), - headers: Vec::new(), + headers: vec![("Retry-After".to_string(), "17".to_string())], }); assert!(mapped.is_instance_of::(py)); + let headers: Vec<(String, String)> = mapped + .value(py) + .getattr("headers") + .and_then(|headers| headers.extract()) + .expect("OCR failures retain provider headers"); + assert_eq!(headers, vec![("Retry-After".to_string(), "17".to_string())]); let args: (u16, String) = mapped .value(py) .getattr("args") diff --git a/litellm/rust_bridge/ocr_lifecycle.py b/litellm/rust_bridge/ocr_lifecycle.py index 1958fdf8cf3..e22722d22c4 100644 --- a/litellm/rust_bridge/ocr_lifecycle.py +++ b/litellm/rust_bridge/ocr_lifecycle.py @@ -4,6 +4,7 @@ from collections.abc import Awaitable, Mapping, Sequence from typing import Final, Protocol, cast # noqa: TID251 # validates dynamically loaded native callables import httpx +from pydantic import TypeAdapter, ValidationError import litellm from litellm.llms.base_llm.ocr.transformation import OCRResponse @@ -41,6 +42,27 @@ def _binding(value: object) -> NativeOcrLifecycle | None: NATIVE_OCR_LIFECYCLE: Final = NativeBinding("_ocr_lifecycle", validate=_binding) +_UPSTREAM_ARGS: Final = TypeAdapter(tuple[int, str]) +_UPSTREAM_HEADERS: Final = TypeAdapter(list[tuple[str, str]]) + + +class UpstreamFailure(Exception): + def __init__(self, response: httpx.Response, cause: Exception) -> None: + super().__init__(str(cause)) + self.message: Final = str(cause) + self.response: Final = response + self.status_code: Final = response.status_code + self.__cause__ = cause + + +def _upstream_failure(error: Exception) -> Exception: + try: + status, body = _UPSTREAM_ARGS.validate_python(error.args) + headers: Final = _UPSTREAM_HEADERS.validate_python(getattr(error, "headers", None)) + except ValidationError: + return error + return UpstreamFailure(httpx.Response(status, content=body.encode(), headers=headers), error) + def select(request: LiteLLMOcrRequest) -> NativeOcrLifecycle | None: if request.kwargs.get("aocr"): @@ -63,18 +85,18 @@ def map_failure(error: Exception, request: LiteLLMOcrRequest, request_provider: mapper: Final = cast( # cast-ok: bounded adapter for the legacy public exception mapper ExceptionMapper, litellm.exception_type ) + original: Final = _upstream_failure(error) try: return mapper( model=model, custom_llm_provider=request_provider, - original_exception=error, + original_exception=original, completion_kwargs=dict(arguments(request)), # mutable-ok: exception mapper requires owned kwargs extra_kwargs=dict(request.kwargs), # mutable-ok: exception mapper requires owned kwargs ) except Exception as public_error: - response: Final = getattr(error, "response", None) - if isinstance(response, httpx.Response): - public_error.response = response - public_error.status_code = response.status_code + if isinstance(original, UpstreamFailure): + public_error.response = original.response + public_error.status_code = original.status_code public_error.__context__ = error return public_error diff --git a/tests/test_litellm_rust/ocr/test_requests.py b/tests/test_litellm_rust/ocr/test_requests.py index 3e95258fb36..e360401a435 100644 --- a/tests/test_litellm_rust/ocr/test_requests.py +++ b/tests/test_litellm_rust/ocr/test_requests.py @@ -13,6 +13,7 @@ from tests.test_litellm_rust.support.recording_server import RecordingServer, Re from tests.test_litellm_rust.support.requests import ( OCR_DOCUMENT, OCR_RESPONSE, + call_native, call_native_aocr, call_native_ocr, ) @@ -43,10 +44,7 @@ async def test_ocr_contract_upstream_status( "num_retries": 0, } with pytest.raises(litellm.BadRequestError) as caught: - if asynchronous: - await call_native_aocr(ocr_server, **arguments) - else: - call_native_ocr(ocr_server, **arguments) + await call_native(ocr_server, asynchronous, **arguments) assert caught.value.status_code == upstream.status assert caught.value.response.status_code == upstream.status @@ -64,10 +62,7 @@ async def test_ocr_contract_provider_error_details( headers: Final = {"Retry-After": "17", "X-Request-ID": "ocr-request-123", "X-Future-Header": "retained"} ocr_server.enqueue(ResponseSpec(body=payload, status=429, headers=headers)) with pytest.raises(litellm.RateLimitError) as caught: - if asynchronous: - await call_native_aocr(ocr_server, num_retries=0) - else: - call_native_ocr(ocr_server, num_retries=0) + await call_native(ocr_server, asynchronous, num_retries=0) response: Final = caught.value.response assert isinstance(response, httpx.Response) if preserved == "body": @@ -87,10 +82,7 @@ async def test_ocr_contract_invalid_response_format( ) -> None: ocr_server.expected_requests = 0 with pytest.raises(litellm.UnsupportedParamsError) as caught: - if asynchronous: - await call_native_aocr(ocr_server, req_format="bogus", num_retries=0) - else: - call_native_ocr(ocr_server, req_format="bogus", num_retries=0) + await call_native(ocr_server, asynchronous, req_format="bogus", num_retries=0) assert caught.value.status_code == 400 for value in ("req_format", "bogus", "native", "litellm"): assert value in str(caught.value) @@ -116,10 +108,7 @@ async def test_ocr_contract_malformed_document_is_actionable( ) -> None: ocr_server.expected_requests = None with pytest.raises(litellm.BadRequestError) as caught: - if asynchronous: - await call_native_aocr(ocr_server, document=document, num_retries=0) - else: - call_native_ocr(ocr_server, document=document, num_retries=0) + await call_native(ocr_server, asynchronous, document=document, num_retries=0) assert caught.value.status_code == 400 assert field.lower() in str(caught.value).lower() assert "NoneType: None" not in str(caught.value) @@ -141,10 +130,7 @@ async def test_ocr_contract_azure_invalid_options_are_bad_requests( ocr_server.expected_requests = 0 arguments: Final = {"model": "azure_ai/doc-intelligence/prebuilt-read", option: value, "num_retries": 0} with pytest.raises(litellm.BadRequestError) as caught: - if asynchronous: - await call_native_aocr(ocr_server, **arguments) - else: - call_native_ocr(ocr_server, **arguments) + await call_native(ocr_server, asynchronous, **arguments) assert caught.value.status_code == 400 assert field in str(caught.value) assert ocr_server.requests == [] diff --git a/tests/test_litellm_rust/support/requests.py b/tests/test_litellm_rust/support/requests.py index 7114e42a59e..b60cf5eac02 100644 --- a/tests/test_litellm_rust/support/requests.py +++ b/tests/test_litellm_rust/support/requests.py @@ -42,6 +42,10 @@ async def call_native_aocr(server: RecordingServer, **kwargs: object) -> OCRResp return await call_aocr(server, **kwargs) +async def call_native(server: RecordingServer, asynchronous: bool, **kwargs: object) -> OCRResponse: + return await call_native_aocr(server, **kwargs) if asynchronous else call_native_ocr(server, **kwargs) + + def request_body(kwargs: dict[str, object]) -> dict[str, object]: additional_args = kwargs["additional_args"] assert isinstance(additional_args, dict) From 4dcbef0558bf2ef9c76a10012389f7ec71a79243 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Thu, 17 Sep 2026 17:42:44 +0000 Subject: [PATCH 14/15] refactor(ocr): drop mutable collection builds flagged by LIT002 gate Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/ocr/legacy.py | 13 ++++++------- 1 file changed, 6 insertions(+), 7 deletions(-) diff --git a/litellm/ocr/legacy.py b/litellm/ocr/legacy.py index 1c9e1c2c28f..27e72195b60 100644 --- a/litellm/ocr/legacy.py +++ b/litellm/ocr/legacy.py @@ -143,10 +143,9 @@ def _prepare_ocr_request( ) except ValueError as error: raise litellm.BadRequestError(message=str(error), model=model, llm_provider=custom_llm_provider) from error - optional_params: Final = { - **mapped_params, - **({OCR_REQUEST_FORMAT_PARAM: requested_format} if requested_format is not None else {}), - } + optional_params: Final = ( + mapped_params if requested_format is None else {**mapped_params, OCR_REQUEST_FORMAT_PARAM: requested_format} + ) verbose_logger.debug("OCR optional_params after mapping: %s", optional_params) @@ -185,7 +184,7 @@ def _error_provider(model: str, custom_llm_provider: str | None) -> str | None: if custom_llm_provider is not None: return custom_llm_provider prefix: Final = model.partition("/")[0] - if prefix in {"mistral", "azure_ai", "vertex_ai"}: + if prefix in ("mistral", "azure_ai", "vertex_ai"): return prefix return "mistral" if model.startswith("mistral-ocr") else None @@ -224,7 +223,7 @@ async def aocr( ) model = prepared.model custom_llm_provider = prepared.custom_llm_provider - completion_kwargs.update({"model": model, "custom_llm_provider": custom_llm_provider}) + completion_kwargs.update(model=model, custom_llm_provider=custom_llm_provider) response = base_llm_http_handler.ocr( model=prepared.model, @@ -390,7 +389,7 @@ def ocr( ) model = prepared.model custom_llm_provider = prepared.custom_llm_provider - completion_kwargs.update({"model": model, "custom_llm_provider": custom_llm_provider}) + completion_kwargs.update(model=model, custom_llm_provider=custom_llm_provider) response: Final = base_llm_http_handler.ocr( model=prepared.model, From cd4d78a26a39ffc2b6005dfb1e8a307f75f070b0 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Thu, 17 Sep 2026 19:06:28 +0000 Subject: [PATCH 15/15] fix(ocr): narrow public error attribute writes and cover callback failure mapping Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/llms/custom_httpx/llm_http_handler.py | 10 ++-- litellm/rust_bridge/ocr/callbacks.py | 6 +- tests/test_litellm/ocr/test_main.py | 16 ++++++ .../rust_bridge/ocr/test_callbacks.py | 56 +++++++++++++++++++ 4 files changed, 82 insertions(+), 6 deletions(-) diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index 7fe0d92b8cc..857adf5b9f1 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -44,7 +44,7 @@ from litellm.llms.base_llm.base_model_iterator import ( MockResponseIterator, ) from litellm.llms.base_llm.batches.transformation import BaseBatchesConfig -from litellm.llms.base_llm.chat.transformation import BaseConfig +from litellm.llms.base_llm.chat.transformation import BaseConfig, BaseLLMException from litellm.llms.base_llm.containers.transformation import BaseContainerConfig from litellm.llms.base_llm.embedding.transformation import BaseEmbeddingConfig from litellm.llms.base_llm.evals.transformation import BaseEvalsAPIConfig @@ -6060,8 +6060,6 @@ class BaseLLMHTTPHandler: error_headers = {} if provider_config is None: - from litellm.llms.base_llm.chat.transformation import BaseLLMException - raise BaseLLMException( status_code=status_code, message=error_text, @@ -6074,7 +6072,11 @@ class BaseLLMHTTPHandler: status_code=status_code, headers=error_headers, ) - if isinstance(provider_config, BaseOCRConfig) and isinstance(error_response, httpx.Response): + if ( + isinstance(provider_config, BaseOCRConfig) + and isinstance(provider_error, BaseLLMException) + and isinstance(error_response, httpx.Response) + ): provider_error.response = error_response if not isinstance(received_status_code, int): provider_error.status_code_is_synthesized = True diff --git a/litellm/rust_bridge/ocr/callbacks.py b/litellm/rust_bridge/ocr/callbacks.py index 4e6a2d054af..0bc7b383eea 100644 --- a/litellm/rust_bridge/ocr/callbacks.py +++ b/litellm/rust_bridge/ocr/callbacks.py @@ -5,6 +5,7 @@ from types import MappingProxyType from typing import Final import httpx +import openai from pydantic import TypeAdapter, ValidationError import litellm @@ -59,7 +60,8 @@ def map_failure(error: Exception, request: LiteLLMOcrRequest, request_provider: original: Final = _upstream_failure(error) public_error: Final = failures.map_failure(original, request.model, request_provider, arguments(request)) if isinstance(original, UpstreamFailure) and public_error.__context__ is original: - public_error.response = original.response - public_error.status_code = original.status_code public_error.__context__ = error + if isinstance(public_error, openai.APIStatusError): + public_error.response = original.response + public_error.status_code = original.status_code return public_error diff --git a/tests/test_litellm/ocr/test_main.py b/tests/test_litellm/ocr/test_main.py index 712a3438ddd..5531a2639c0 100644 --- a/tests/test_litellm/ocr/test_main.py +++ b/tests/test_litellm/ocr/test_main.py @@ -17,6 +17,7 @@ from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler from litellm.ocr.main import _prepare_ocr_request from litellm.rust_bridge import bindings, configuration, runtime from litellm.rust_bridge.ocr.entrypoints import NATIVE_AOCR, NATIVE_OCR +from litellm.utils import ProviderConfigManager @pytest.fixture @@ -277,6 +278,7 @@ def _prepare(model: str, document: object, **kwargs: object) -> object: ( ("https://example.com/file.pdf", "document must be a dict"), ({"type": "video_url", "video_url": "https://example.com/clip.mp4"}, "Invalid document type: video_url"), + ({"type": "document_url", "document_url": ""}, "Document URL is required"), ), ) def test_prepare_ocr_request_rejects_malformed_documents(document: object, match: str) -> None: @@ -284,6 +286,20 @@ def test_prepare_ocr_request_rejects_malformed_documents(document: object, match _prepare("mistral/mistral-ocr-latest", document) +def test_prepare_ocr_request_maps_param_mapping_errors_to_bad_request(monkeypatch: pytest.MonkeyPatch) -> None: + config: Final = Mock() + config.resolve_connection_params.return_value = ("test-key", None) + config.get_supported_ocr_params.return_value = ["pages"] + config.map_ocr_params.side_effect = ValueError("pages must be a list") + monkeypatch.setattr(ProviderConfigManager, "get_provider_ocr_config", Mock(return_value=config)) + + with pytest.raises(litellm.BadRequestError, match="pages must be a list") as error: + _prepare("mistral/mistral-ocr-latest", dict(PRICING_DOCUMENT), pages="1") + + assert error.value.llm_provider == "mistral" + assert isinstance(error.value.__cause__, ValueError) + + def test_prepare_ocr_request_rejects_provider_without_ocr_support() -> None: with pytest.raises(ValueError, match="OCR is not supported for provider: openai"): _prepare("openai/gpt-4o", dict(PRICING_DOCUMENT)) diff --git a/tests/test_litellm/rust_bridge/ocr/test_callbacks.py b/tests/test_litellm/rust_bridge/ocr/test_callbacks.py index c5e9d60ff86..a85940aa049 100644 --- a/tests/test_litellm/rust_bridge/ocr/test_callbacks.py +++ b/tests/test_litellm/rust_bridge/ocr/test_callbacks.py @@ -1,4 +1,32 @@ +from typing import Final + +import pytest + +import litellm +from litellm.rust_bridge.ocr.callbacks import UpstreamFailure, map_failure from litellm.rust_bridge.ocr.callbacks import response as build_ocr_response +from litellm.rust_bridge.ocr.entrypoints import LiteLLMOcrRequest + +REQUEST: Final = LiteLLMOcrRequest( + model="mistral/mistral-ocr-latest", + document={"type": "document_url", "document_url": "https://example.com/file.pdf"}, + api_key="test-key", + api_base=None, + timeout=None, + custom_llm_provider=None, + extra_headers=None, + kwargs={"req_format": "markdown"}, +) + + +class RustUpstreamError(Exception): + def __init__(self, status: int, body: str, headers: tuple[tuple[str, str], ...]) -> None: + super().__init__(status, body) + self.headers: Final = list(headers) + + +class RustFormatError(Exception): + ocr_request_format_error: Final = True def test_rust_ocr_response_retains_provider_native_response(): @@ -16,3 +44,31 @@ def test_rust_ocr_response_retains_provider_native_response(): assert response.get_provider_native_response() == provider_response assert response.model_dump().get("provider_native_response") is None + + +def test_map_failure_builds_public_error_from_upstream_status_and_headers() -> None: + error: Final = RustUpstreamError(429, '{"message": "slow down"}', (("retry-after", "7"),)) + + public_error: Final = map_failure(error, REQUEST, "mistral") + + assert isinstance(public_error, litellm.RateLimitError) + assert public_error.status_code == 429 + assert public_error.response.headers["retry-after"] == "7" + assert public_error.response.text == '{"message": "slow down"}' + assert public_error.__context__ is error + assert public_error.llm_provider == "mistral" + + +def test_map_failure_leaves_non_upstream_errors_unwrapped() -> None: + error: Final = RuntimeError("bridge exploded") + + public_error: Final = map_failure(error, REQUEST, "mistral") + + assert not isinstance(public_error, UpstreamFailure) + assert isinstance(public_error, litellm.APIConnectionError) + assert "bridge exploded" in str(public_error) + + +def test_map_failure_reports_invalid_request_format_as_unsupported_params() -> None: + with pytest.raises(litellm.UnsupportedParamsError, match="Invalid `req_format`: 'markdown'"): + raise map_failure(RustFormatError(), REQUEST, "mistral")