diff --git a/litellm-rust/Cargo.lock b/litellm-rust/Cargo.lock index 2ec86fc00e7..5f9c3a0b270 100644 --- a/litellm-rust/Cargo.lock +++ b/litellm-rust/Cargo.lock @@ -137,12 +137,24 @@ dependencies = [ "syn", ] +[[package]] +name = "equivalent" +version = "1.0.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "877a4ace8713b0bcf2a4e7eec82529c029f1d0619886d18145fea96c3ffe5c0f" + [[package]] name = "find-msvc-tools" version = "0.1.9" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "5baebc0774151f905a1a2cc41989300b1e6fbb29aff0ceffa1064fdd3088d582" +[[package]] +name = "fnv" +version = "1.0.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3f9eec918d3f24069decb9af1554cad7c880e2da24a9afd88aca000531ab82c1" + [[package]] name = "form_urlencoded" version = "1.2.2" @@ -152,6 +164,21 @@ dependencies = [ "percent-encoding", ] +[[package]] +name = "futures" +version = "0.3.32" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8b147ee9d1f6d097cef9ce628cd2ee62288d963e16fb287bd9286455b241382d" +dependencies = [ + "futures-channel", + "futures-core", + "futures-executor", + "futures-io", + "futures-sink", + "futures-task", + "futures-util", +] + [[package]] name = "futures-channel" version = "0.3.32" @@ -168,12 +195,34 @@ version = "0.3.32" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "7e3450815272ef58cec6d564423f6e755e25379b217b0bc688e295ba24df6b1d" +[[package]] +name = "futures-executor" +version = "0.3.32" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "baf29c38818342a3b26b5b923639e7b1f4a61fc5e76102d4b1981c6dc7a7579d" +dependencies = [ + "futures-core", + "futures-task", + "futures-util", +] + [[package]] name = "futures-io" version = "0.3.32" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "cecba35d7ad927e23624b22ad55235f2239cfa44fd10428eecbeba6d6a717718" +[[package]] +name = "futures-macro" +version = "0.3.32" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e835b70203e41293343137df5c0664546da5745f82ec9b84d40be8336958447b" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "futures-sink" version = "0.3.32" @@ -192,8 +241,10 @@ version = "0.3.32" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "389ca41296e6190b48053de0321d02a77f32f8a5d2461dd38762c0593805c6d6" dependencies = [ + "futures-channel", "futures-core", "futures-io", + "futures-macro", "futures-sink", "futures-task", "memchr", @@ -238,6 +289,31 @@ dependencies = [ "wasm-bindgen", ] +[[package]] +name = "h2" +version = "0.4.15" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6cb093c84e8bd9b188d4c4a8cb6579fc016968d14c99882163cd3ff402a4f155" +dependencies = [ + "atomic-waker", + "bytes", + "fnv", + "futures-core", + "futures-sink", + "http", + "indexmap", + "slab", + "tokio", + "tokio-util", + "tracing", +] + +[[package]] +name = "hashbrown" +version = "0.17.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ed5909b6e89a2db4456e54cd5f673791d7eca6732202bbf2a9cc504fe2f9b84a" + [[package]] name = "heck" version = "0.5.0" @@ -293,6 +369,7 @@ dependencies = [ "bytes", "futures-channel", "futures-core", + "h2", "http", "http-body", "httparse", @@ -445,6 +522,16 @@ dependencies = [ "icu_properties", ] +[[package]] +name = "indexmap" +version = "2.14.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d466e9454f08e4a911e14806c24e16fba1b4c121d1ea474396f396069cf949d9" +dependencies = [ + "equivalent", + "hashbrown", +] + [[package]] name = "indoc" version = "2.0.7" @@ -511,7 +598,9 @@ dependencies = [ "litellm-core", "litellm-providers", "pyo3", + "pyo3-async-runtimes", "serde_json", + "tokio", ] [[package]] @@ -633,6 +722,19 @@ dependencies = [ "unindent", ] +[[package]] +name = "pyo3-async-runtimes" +version = "0.23.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "977dc837525cfd22919ba6a831413854beb7c99a256c03bf8624ad707e45810e" +dependencies = [ + "futures", + "once_cell", + "pin-project-lite", + "pyo3", + "tokio", +] + [[package]] name = "pyo3-build-config" version = "0.23.5" @@ -815,9 +917,8 @@ checksum = "eddd3ca559203180a307f12d114c268abf583f59b03cb906fd0b3ff8646c1147" dependencies = [ "base64", "bytes", - "futures-channel", "futures-core", - "futures-util", + "h2", "http", "http-body", "http-body-util", @@ -1218,6 +1319,19 @@ dependencies = [ "tungstenite", ] +[[package]] +name = "tokio-util" +version = "0.7.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9ae9cec805b01e8fc3fd2fe289f89149a9b66dd16786abd8b19cfa7b48cb0098" +dependencies = [ + "bytes", + "futures-core", + "futures-sink", + "pin-project-lite", + "tokio", +] + [[package]] name = "tower" version = "0.5.3" diff --git a/litellm-rust/Cargo.toml b/litellm-rust/Cargo.toml index 0b1528672a7..63c7d6e8b9d 100644 --- a/litellm-rust/Cargo.toml +++ b/litellm-rust/Cargo.toml @@ -15,7 +15,8 @@ repository = "https://github.com/BerriAI/litellm" litellm-core = { path = "crates/core" } litellm-providers = { path = "crates/providers" } pyo3 = "0.23.5" -reqwest = { version = "0.12", default-features = false, features = ["blocking", "json", "rustls-tls"] } +pyo3-async-runtimes = { version = "0.23.0", features = ["tokio-runtime"] } +reqwest = { version = "0.12", default-features = false, features = ["json", "rustls-tls", "http2"] } serde = { version = "1.0", features = ["derive"] } serde_json = "1.0" thiserror = "2.0" diff --git a/litellm-rust/crates/core/src/error.rs b/litellm-rust/crates/core/src/error.rs index 645e261f76d..b57fb2f5de8 100644 --- a/litellm-rust/crates/core/src/error.rs +++ b/litellm-rust/crates/core/src/error.rs @@ -13,6 +13,8 @@ pub enum CoreError { MissingField(&'static str), #[error("invalid response: {0}")] InvalidResponse(String), + #[error("invalid provider: {0}")] + InvalidProvider(String), #[error("{0}")] Auth(String), #[error("OCR request failed with status {status}: {body}")] diff --git a/litellm-rust/crates/core/src/lib.rs b/litellm-rust/crates/core/src/lib.rs index 5d4d5bfd142..6142cece16f 100644 --- a/litellm-rust/crates/core/src/lib.rs +++ b/litellm-rust/crates/core/src/lib.rs @@ -1,5 +1,7 @@ pub mod error; pub mod ocr; +pub mod providers; pub mod realtime; pub use error::{CoreError, CoreResult}; +pub use providers::LlmProvider; diff --git a/litellm-rust/crates/core/src/ocr/transformation.rs b/litellm-rust/crates/core/src/ocr/transformation.rs index 7353d9d22c4..c24f3519354 100644 --- a/litellm-rust/crates/core/src/ocr/transformation.rs +++ b/litellm-rust/crates/core/src/ocr/transformation.rs @@ -4,7 +4,7 @@ use crate::CoreResult; use super::types::{OcrRequestData, OcrResponseData}; -pub trait OcrProviderConfig { +pub trait OcrProviderConfig: Sync { fn supported_ocr_params(&self) -> &'static [&'static str]; fn map_ocr_params(&self, non_default_params: &Map) -> Map { diff --git a/litellm-rust/crates/core/src/providers/mod.rs b/litellm-rust/crates/core/src/providers/mod.rs new file mode 100644 index 00000000000..3d6f90356a1 --- /dev/null +++ b/litellm-rust/crates/core/src/providers/mod.rs @@ -0,0 +1,3 @@ +pub mod provider; + +pub use provider::LlmProvider; diff --git a/litellm-rust/crates/core/src/providers/provider.rs b/litellm-rust/crates/core/src/providers/provider.rs new file mode 100644 index 00000000000..f14152aae18 --- /dev/null +++ b/litellm-rust/crates/core/src/providers/provider.rs @@ -0,0 +1,659 @@ +use std::fmt; +use std::str::FromStr; + +use crate::error::CoreError; + +#[allow(clippy::enum_variant_names)] +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum LlmProvider { + Openai, + Chatgpt, + OpenaiLike, + JinaAi, + Xai, + Zai, + CustomOpenai, + TextCompletionOpenai, + Cohere, + CohereChat, + Clarifai, + Anthropic, + AnthropicText, + Bytez, + Replicate, + Reducto, + Runwayml, + AwsPolly, + Huggingface, + TogetherAi, + Openrouter, + Datarobot, + VertexAi, + VertexAiBeta, + Gemini, + Ai21, + Baseten, + BlackForestLabs, + Azure, + AzureText, + AzureAi, + Sagemaker, + SagemakerChat, + SagemakerNova, + Bedrock, + Vllm, + NlpCloud, + Petals, + Oobabooga, + Ollama, + OllamaChat, + Deepinfra, + Perplexity, + Mistral, + Milvus, + Groq, + A2a, + Gigachat, + NvidiaNim, + NvidiaRiva, + Soniox, + Cerebras, + Ai21Chat, + Volcengine, + Codestral, + TextCompletionCodestral, + Dashscope, + Modelscope, + Moonshot, + Publicai, + V0, + Morph, + LambdaAi, + Inception, + TextCompletionInception, + Deepseek, + Sambanova, + Maritalk, + Voyage, + Cloudflare, + Xinference, + FireworksAi, + Friendliai, + FeatherlessAi, + Watsonx, + WatsonxText, + Triton, + Predibase, + Databricks, + Empower, + Github, + Ragflow, + Compactifai, + DockerModelRunner, + Custom, + LitellmProxy, + HostedVllm, + Llamafile, + LmStudio, + Galadriel, + Nebius, + Infinity, + Deepgram, + Elevenlabs, + Novita, + AiohttpOpenai, + Langfuse, + Humanloop, + Topaz, + Sap, + Assemblyai, + CharityEngine, + GithubCopilot, + Snowflake, + GradientAi, + MetaLlama, + Nscale, + PgVector, + S3Vectors, + Helicone, + Hyperbolic, + Recraft, + FalAi, + Stability, + Heroku, + Aiml, + Cometapi, + Oci, + AutoRouter, + VercelAiGateway, + Dotprompt, + Manus, + Wandb, + Ovhcloud, + Scaleway, + Lemonade, + AmazonNova, + A2aAgent, + Langgraph, + Langflow, + Minimax, + Synthetic, + Apertis, + NanoGpt, + Poe, + Chutes, + Neosantara, + Parasail, + XiaomiMimo, + Tensormesh, + Libertai, + Pinstripes, + Darkbloom, + LitellmAgent, + Cursor, + BedrockMantle, +} + +impl LlmProvider { + #[allow(clippy::too_many_lines)] + pub fn as_str(self) -> &'static str { + match self { + LlmProvider::Openai => "openai", + LlmProvider::Chatgpt => "chatgpt", + LlmProvider::OpenaiLike => "openai_like", + LlmProvider::JinaAi => "jina_ai", + LlmProvider::Xai => "xai", + LlmProvider::Zai => "zai", + LlmProvider::CustomOpenai => "custom_openai", + LlmProvider::TextCompletionOpenai => "text-completion-openai", + LlmProvider::Cohere => "cohere", + LlmProvider::CohereChat => "cohere_chat", + LlmProvider::Clarifai => "clarifai", + LlmProvider::Anthropic => "anthropic", + LlmProvider::AnthropicText => "anthropic_text", + LlmProvider::Bytez => "bytez", + LlmProvider::Replicate => "replicate", + LlmProvider::Reducto => "reducto", + LlmProvider::Runwayml => "runwayml", + LlmProvider::AwsPolly => "aws_polly", + LlmProvider::Huggingface => "huggingface", + LlmProvider::TogetherAi => "together_ai", + LlmProvider::Openrouter => "openrouter", + LlmProvider::Datarobot => "datarobot", + LlmProvider::VertexAi => "vertex_ai", + LlmProvider::VertexAiBeta => "vertex_ai_beta", + LlmProvider::Gemini => "gemini", + LlmProvider::Ai21 => "ai21", + LlmProvider::Baseten => "baseten", + LlmProvider::BlackForestLabs => "black_forest_labs", + LlmProvider::Azure => "azure", + LlmProvider::AzureText => "azure_text", + LlmProvider::AzureAi => "azure_ai", + LlmProvider::Sagemaker => "sagemaker", + LlmProvider::SagemakerChat => "sagemaker_chat", + LlmProvider::SagemakerNova => "sagemaker_nova", + LlmProvider::Bedrock => "bedrock", + LlmProvider::Vllm => "vllm", + LlmProvider::NlpCloud => "nlp_cloud", + LlmProvider::Petals => "petals", + LlmProvider::Oobabooga => "oobabooga", + LlmProvider::Ollama => "ollama", + LlmProvider::OllamaChat => "ollama_chat", + LlmProvider::Deepinfra => "deepinfra", + LlmProvider::Perplexity => "perplexity", + LlmProvider::Mistral => "mistral", + LlmProvider::Milvus => "milvus", + LlmProvider::Groq => "groq", + LlmProvider::A2a => "a2a", + LlmProvider::Gigachat => "gigachat", + LlmProvider::NvidiaNim => "nvidia_nim", + LlmProvider::NvidiaRiva => "nvidia_riva", + LlmProvider::Soniox => "soniox", + LlmProvider::Cerebras => "cerebras", + LlmProvider::Ai21Chat => "ai21_chat", + LlmProvider::Volcengine => "volcengine", + LlmProvider::Codestral => "codestral", + LlmProvider::TextCompletionCodestral => "text-completion-codestral", + LlmProvider::Dashscope => "dashscope", + LlmProvider::Modelscope => "modelscope", + LlmProvider::Moonshot => "moonshot", + LlmProvider::Publicai => "publicai", + LlmProvider::V0 => "v0", + LlmProvider::Morph => "morph", + LlmProvider::LambdaAi => "lambda_ai", + LlmProvider::Inception => "inception", + LlmProvider::TextCompletionInception => "text-completion-inception", + LlmProvider::Deepseek => "deepseek", + LlmProvider::Sambanova => "sambanova", + LlmProvider::Maritalk => "maritalk", + LlmProvider::Voyage => "voyage", + LlmProvider::Cloudflare => "cloudflare", + LlmProvider::Xinference => "xinference", + LlmProvider::FireworksAi => "fireworks_ai", + LlmProvider::Friendliai => "friendliai", + LlmProvider::FeatherlessAi => "featherless_ai", + LlmProvider::Watsonx => "watsonx", + LlmProvider::WatsonxText => "watsonx_text", + LlmProvider::Triton => "triton", + LlmProvider::Predibase => "predibase", + LlmProvider::Databricks => "databricks", + LlmProvider::Empower => "empower", + LlmProvider::Github => "github", + LlmProvider::Ragflow => "ragflow", + LlmProvider::Compactifai => "compactifai", + LlmProvider::DockerModelRunner => "docker_model_runner", + LlmProvider::Custom => "custom", + LlmProvider::LitellmProxy => "litellm_proxy", + LlmProvider::HostedVllm => "hosted_vllm", + LlmProvider::Llamafile => "llamafile", + LlmProvider::LmStudio => "lm_studio", + LlmProvider::Galadriel => "galadriel", + LlmProvider::Nebius => "nebius", + LlmProvider::Infinity => "infinity", + LlmProvider::Deepgram => "deepgram", + LlmProvider::Elevenlabs => "elevenlabs", + LlmProvider::Novita => "novita", + LlmProvider::AiohttpOpenai => "aiohttp_openai", + LlmProvider::Langfuse => "langfuse", + LlmProvider::Humanloop => "humanloop", + LlmProvider::Topaz => "topaz", + LlmProvider::Sap => "sap", + LlmProvider::Assemblyai => "assemblyai", + LlmProvider::CharityEngine => "charity_engine", + LlmProvider::GithubCopilot => "github_copilot", + LlmProvider::Snowflake => "snowflake", + LlmProvider::GradientAi => "gradient_ai", + LlmProvider::MetaLlama => "meta_llama", + LlmProvider::Nscale => "nscale", + LlmProvider::PgVector => "pg_vector", + LlmProvider::S3Vectors => "s3_vectors", + LlmProvider::Helicone => "helicone", + LlmProvider::Hyperbolic => "hyperbolic", + LlmProvider::Recraft => "recraft", + LlmProvider::FalAi => "fal_ai", + LlmProvider::Stability => "stability", + LlmProvider::Heroku => "heroku", + LlmProvider::Aiml => "aiml", + LlmProvider::Cometapi => "cometapi", + LlmProvider::Oci => "oci", + LlmProvider::AutoRouter => "auto_router", + LlmProvider::VercelAiGateway => "vercel_ai_gateway", + LlmProvider::Dotprompt => "dotprompt", + LlmProvider::Manus => "manus", + LlmProvider::Wandb => "wandb", + LlmProvider::Ovhcloud => "ovhcloud", + LlmProvider::Scaleway => "scaleway", + LlmProvider::Lemonade => "lemonade", + LlmProvider::AmazonNova => "amazon_nova", + LlmProvider::A2aAgent => "a2a_agent", + LlmProvider::Langgraph => "langgraph", + LlmProvider::Langflow => "langflow", + LlmProvider::Minimax => "minimax", + LlmProvider::Synthetic => "synthetic", + LlmProvider::Apertis => "apertis", + LlmProvider::NanoGpt => "nano-gpt", + LlmProvider::Poe => "poe", + LlmProvider::Chutes => "chutes", + LlmProvider::Neosantara => "neosantara", + LlmProvider::Parasail => "parasail", + LlmProvider::XiaomiMimo => "xiaomi_mimo", + LlmProvider::Tensormesh => "tensormesh", + LlmProvider::Libertai => "libertai", + LlmProvider::Pinstripes => "pinstripes", + LlmProvider::Darkbloom => "darkbloom", + LlmProvider::LitellmAgent => "litellm_agent", + LlmProvider::Cursor => "cursor", + LlmProvider::BedrockMantle => "bedrock_mantle", + } + } + + #[allow(clippy::too_many_lines)] + pub const ALL: [LlmProvider; 146] = [ + LlmProvider::Openai, + LlmProvider::Chatgpt, + LlmProvider::OpenaiLike, + LlmProvider::JinaAi, + LlmProvider::Xai, + LlmProvider::Zai, + LlmProvider::CustomOpenai, + LlmProvider::TextCompletionOpenai, + LlmProvider::Cohere, + LlmProvider::CohereChat, + LlmProvider::Clarifai, + LlmProvider::Anthropic, + LlmProvider::AnthropicText, + LlmProvider::Bytez, + LlmProvider::Replicate, + LlmProvider::Reducto, + LlmProvider::Runwayml, + LlmProvider::AwsPolly, + LlmProvider::Huggingface, + LlmProvider::TogetherAi, + LlmProvider::Openrouter, + LlmProvider::Datarobot, + LlmProvider::VertexAi, + LlmProvider::VertexAiBeta, + LlmProvider::Gemini, + LlmProvider::Ai21, + LlmProvider::Baseten, + LlmProvider::BlackForestLabs, + LlmProvider::Azure, + LlmProvider::AzureText, + LlmProvider::AzureAi, + LlmProvider::Sagemaker, + LlmProvider::SagemakerChat, + LlmProvider::SagemakerNova, + LlmProvider::Bedrock, + LlmProvider::Vllm, + LlmProvider::NlpCloud, + LlmProvider::Petals, + LlmProvider::Oobabooga, + LlmProvider::Ollama, + LlmProvider::OllamaChat, + LlmProvider::Deepinfra, + LlmProvider::Perplexity, + LlmProvider::Mistral, + LlmProvider::Milvus, + LlmProvider::Groq, + LlmProvider::A2a, + LlmProvider::Gigachat, + LlmProvider::NvidiaNim, + LlmProvider::NvidiaRiva, + LlmProvider::Soniox, + LlmProvider::Cerebras, + LlmProvider::Ai21Chat, + LlmProvider::Volcengine, + LlmProvider::Codestral, + LlmProvider::TextCompletionCodestral, + LlmProvider::Dashscope, + LlmProvider::Modelscope, + LlmProvider::Moonshot, + LlmProvider::Publicai, + LlmProvider::V0, + LlmProvider::Morph, + LlmProvider::LambdaAi, + LlmProvider::Inception, + LlmProvider::TextCompletionInception, + LlmProvider::Deepseek, + LlmProvider::Sambanova, + LlmProvider::Maritalk, + LlmProvider::Voyage, + LlmProvider::Cloudflare, + LlmProvider::Xinference, + LlmProvider::FireworksAi, + LlmProvider::Friendliai, + LlmProvider::FeatherlessAi, + LlmProvider::Watsonx, + LlmProvider::WatsonxText, + LlmProvider::Triton, + LlmProvider::Predibase, + LlmProvider::Databricks, + LlmProvider::Empower, + LlmProvider::Github, + LlmProvider::Ragflow, + LlmProvider::Compactifai, + LlmProvider::DockerModelRunner, + LlmProvider::Custom, + LlmProvider::LitellmProxy, + LlmProvider::HostedVllm, + LlmProvider::Llamafile, + LlmProvider::LmStudio, + LlmProvider::Galadriel, + LlmProvider::Nebius, + LlmProvider::Infinity, + LlmProvider::Deepgram, + LlmProvider::Elevenlabs, + LlmProvider::Novita, + LlmProvider::AiohttpOpenai, + LlmProvider::Langfuse, + LlmProvider::Humanloop, + LlmProvider::Topaz, + LlmProvider::Sap, + LlmProvider::Assemblyai, + LlmProvider::CharityEngine, + LlmProvider::GithubCopilot, + LlmProvider::Snowflake, + LlmProvider::GradientAi, + LlmProvider::MetaLlama, + LlmProvider::Nscale, + LlmProvider::PgVector, + LlmProvider::S3Vectors, + LlmProvider::Helicone, + LlmProvider::Hyperbolic, + LlmProvider::Recraft, + LlmProvider::FalAi, + LlmProvider::Stability, + LlmProvider::Heroku, + LlmProvider::Aiml, + LlmProvider::Cometapi, + LlmProvider::Oci, + LlmProvider::AutoRouter, + LlmProvider::VercelAiGateway, + LlmProvider::Dotprompt, + LlmProvider::Manus, + LlmProvider::Wandb, + LlmProvider::Ovhcloud, + LlmProvider::Scaleway, + LlmProvider::Lemonade, + LlmProvider::AmazonNova, + LlmProvider::A2aAgent, + LlmProvider::Langgraph, + LlmProvider::Langflow, + LlmProvider::Minimax, + LlmProvider::Synthetic, + LlmProvider::Apertis, + LlmProvider::NanoGpt, + LlmProvider::Poe, + LlmProvider::Chutes, + LlmProvider::Neosantara, + LlmProvider::Parasail, + LlmProvider::XiaomiMimo, + LlmProvider::Tensormesh, + LlmProvider::Libertai, + LlmProvider::Pinstripes, + LlmProvider::Darkbloom, + LlmProvider::LitellmAgent, + LlmProvider::Cursor, + LlmProvider::BedrockMantle, + ]; +} + +impl fmt::Display for LlmProvider { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.write_str(self.as_str()) + } +} + +impl FromStr for LlmProvider { + type Err = CoreError; + + fn from_str(value: &str) -> Result { + LlmProvider::ALL + .iter() + .copied() + .find(|provider| provider.as_str() == value) + .ok_or_else(|| CoreError::InvalidProvider(value.to_string())) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + const PYTHON_LLM_PROVIDER_VALUES: &[&str] = &[ + "openai", + "chatgpt", + "openai_like", + "jina_ai", + "xai", + "zai", + "custom_openai", + "text-completion-openai", + "cohere", + "cohere_chat", + "clarifai", + "anthropic", + "anthropic_text", + "bytez", + "replicate", + "reducto", + "runwayml", + "aws_polly", + "huggingface", + "together_ai", + "openrouter", + "datarobot", + "vertex_ai", + "vertex_ai_beta", + "gemini", + "ai21", + "baseten", + "black_forest_labs", + "azure", + "azure_text", + "azure_ai", + "sagemaker", + "sagemaker_chat", + "sagemaker_nova", + "bedrock", + "vllm", + "nlp_cloud", + "petals", + "oobabooga", + "ollama", + "ollama_chat", + "deepinfra", + "perplexity", + "mistral", + "milvus", + "groq", + "a2a", + "gigachat", + "nvidia_nim", + "nvidia_riva", + "soniox", + "cerebras", + "ai21_chat", + "volcengine", + "codestral", + "text-completion-codestral", + "dashscope", + "modelscope", + "moonshot", + "publicai", + "v0", + "morph", + "lambda_ai", + "inception", + "text-completion-inception", + "deepseek", + "sambanova", + "maritalk", + "voyage", + "cloudflare", + "xinference", + "fireworks_ai", + "friendliai", + "featherless_ai", + "watsonx", + "watsonx_text", + "triton", + "predibase", + "databricks", + "empower", + "github", + "ragflow", + "compactifai", + "docker_model_runner", + "custom", + "litellm_proxy", + "hosted_vllm", + "llamafile", + "lm_studio", + "galadriel", + "nebius", + "infinity", + "deepgram", + "elevenlabs", + "novita", + "aiohttp_openai", + "langfuse", + "humanloop", + "topaz", + "sap", + "assemblyai", + "charity_engine", + "github_copilot", + "snowflake", + "gradient_ai", + "meta_llama", + "nscale", + "pg_vector", + "s3_vectors", + "helicone", + "hyperbolic", + "recraft", + "fal_ai", + "stability", + "heroku", + "aiml", + "cometapi", + "oci", + "auto_router", + "vercel_ai_gateway", + "dotprompt", + "manus", + "wandb", + "ovhcloud", + "scaleway", + "lemonade", + "amazon_nova", + "a2a_agent", + "langgraph", + "langflow", + "minimax", + "synthetic", + "apertis", + "nano-gpt", + "poe", + "chutes", + "neosantara", + "parasail", + "xiaomi_mimo", + "tensormesh", + "libertai", + "pinstripes", + "darkbloom", + "litellm_agent", + "cursor", + "bedrock_mantle", + ]; + + #[test] + fn provider_values_match_python_enum_order() { + assert_eq!(LlmProvider::ALL.len(), 146); + assert_eq!( + LlmProvider::ALL + .iter() + .map(|provider| provider.as_str()) + .collect::>(), + PYTHON_LLM_PROVIDER_VALUES + ); + } + + #[test] + fn from_str_round_trips_all_providers() { + for provider in LlmProvider::ALL { + assert_eq!(LlmProvider::from_str(provider.as_str()), Ok(provider)); + assert_eq!(provider.to_string(), provider.as_str()); + } + } + + #[test] + fn from_str_rejects_unknown_provider() { + assert_eq!( + LlmProvider::from_str("not-a-provider"), + Err(CoreError::InvalidProvider("not-a-provider".to_string())) + ); + } +} diff --git a/litellm-rust/crates/providers/src/ocr.rs b/litellm-rust/crates/providers/src/ocr.rs index dcd56a5f0b4..a52d5532afc 100644 --- a/litellm-rust/crates/providers/src/ocr.rs +++ b/litellm-rust/crates/providers/src/ocr.rs @@ -4,12 +4,14 @@ //! resolve the API key, build the URL + body via the pure transforms, POST it, //! and normalize the response. The HTTP client is built once and reused. +use std::str::FromStr; use std::sync::OnceLock; use std::time::Duration; use litellm_core::error::CoreError; use litellm_core::ocr::transformation::OcrProviderConfig; use litellm_core::CoreResult; +use litellm_core::LlmProvider; use serde_json::{Map, Value}; use crate::mistral::ocr::transformation as mistral; @@ -25,11 +27,11 @@ const OCR_TIMEOUT_SECS: u64 = 600; /// forwarding sensitive payloads across the host boundary. const ERROR_BODY_MAX_CHARS: usize = 256; -/// Process-wide blocking HTTP client (connection pool + TLS reused across calls). -fn http_client() -> &'static reqwest::blocking::Client { - static CLIENT: OnceLock = OnceLock::new(); +/// Process-wide async HTTP client (connection pool + TLS reused across calls). +fn http_client() -> &'static reqwest::Client { + static CLIENT: OnceLock = OnceLock::new(); CLIENT.get_or_init(|| { - reqwest::blocking::Client::builder() + reqwest::Client::builder() .timeout(Duration::from_secs(OCR_TIMEOUT_SECS)) .build() .expect("failed to build reqwest client") @@ -44,39 +46,69 @@ fn truncate_error_body(body: &str) -> String { format!("{truncated}... (truncated)") } +fn ocr_config_for(provider: LlmProvider) -> Option<&'static dyn OcrProviderConfig> { + match provider { + LlmProvider::Mistral => Some(&MISTRAL_OCR_CONFIG), + _ => None, + } +} + +fn string_headers(extra_headers: Option>) -> Vec<(String, String)> { + extra_headers + .unwrap_or_default() + .into_iter() + .filter_map(|(key, value)| value.as_str().map(|value| (key, value.to_string()))) + .collect() +} + +pub struct OcrRequest<'a> { + pub model: &'a str, + pub document: Value, + pub api_key: Option<&'a str>, + pub api_base: Option<&'a str>, + pub custom_llm_provider: &'a str, + pub extra_headers: Option>, + pub optional_params: Map, + pub timeout: Option, +} + /// Perform a Mistral OCR call end to end and return the normalized response as /// JSON (the shape the Python `OCRResponse` model expects). /// -/// Blocking: intended to be called with the GIL released from the Python bridge. -pub fn run_ocr( - model: &str, - document: Value, - api_key: Option<&str>, - api_base: Option<&str>, - optional_params: Map, - timeout: Option, -) -> CoreResult { - let config = &MISTRAL_OCR_CONFIG; +/// Async: intended to be awaited directly by the Python bridge's async entrypoint. +pub async fn ocr(request: OcrRequest<'_>) -> CoreResult { + let model = request.model; + let provider = LlmProvider::from_str(request.custom_llm_provider)?; + let config = + ocr_config_for(provider).ok_or_else(|| CoreError::InvalidProvider(provider.to_string()))?; - let api_key = mistral::resolve_api_key(api_key, &|key| std::env::var(key).ok())?; - let url = mistral::complete_url(api_base); - let filtered_params = config.map_ocr_params(&optional_params); + // TODO: key and URL resolution are still Mistral-specific while Mistral is + // the only Rust OCR provider. Move these onto the trait when another OCR + // provider is added here. + let api_key = mistral::resolve_api_key(request.api_key, &|key| std::env::var(key).ok())?; + let url = mistral::complete_url(request.api_base); + let filtered_params = config.map_ocr_params(&request.optional_params); let body = config - .transform_ocr_request(model, document, filtered_params)? + .transform_ocr_request(model, request.document, filtered_params)? .data; - let mut request = http_client().post(&url).bearer_auth(&api_key).json(&body); - if let Some(duration) = timeout { - request = request.timeout(duration); + let mut request_builder = http_client().post(&url).bearer_auth(&api_key).json(&body); + for (key, value) in string_headers(request.extra_headers) { + request_builder = request_builder.header(&key, value); + } + if let Some(duration) = request.timeout { + request_builder = request_builder.timeout(duration); } - let response = request + let response = request_builder .send() + .await .map_err(|err| CoreError::Network(err.to_string()))?; let status = response.status(); let text = response .text() + .await .map_err(|err| CoreError::Network(err.to_string()))?; if !status.is_success() { @@ -97,6 +129,7 @@ pub fn run_ocr( #[cfg(test)] mod tests { use super::*; + use serde_json::json; #[test] fn truncate_error_body_passes_short_strings_through() { @@ -124,4 +157,27 @@ mod tests { let truncated = truncate_error_body(&body); assert!(truncated.is_char_boundary(truncated.len())); } + + #[test] + fn ocr_registry_supports_only_mistral() { + assert!(ocr_config_for(LlmProvider::Mistral).is_some()); + assert!(ocr_config_for(LlmProvider::Openai).is_none()); + } + + #[test] + fn string_headers_keeps_only_string_values() { + let headers = json!({ + "x-trace-id": "trace-1", + "x-number": 42, + "x-bool": true + }) + .as_object() + .unwrap() + .clone(); + + assert_eq!( + string_headers(Some(headers)), + vec![("x-trace-id".to_string(), "trace-1".to_string())] + ); + } } diff --git a/litellm-rust/crates/python-bridge/Cargo.toml b/litellm-rust/crates/python-bridge/Cargo.toml index 80b6478daac..d0a7cd56cc6 100644 --- a/litellm-rust/crates/python-bridge/Cargo.toml +++ b/litellm-rust/crates/python-bridge/Cargo.toml @@ -13,4 +13,6 @@ crate-type = ["cdylib"] litellm-core.workspace = true litellm-providers.workspace = true pyo3 = { workspace = true, features = ["extension-module"] } +pyo3-async-runtimes.workspace = true serde_json.workspace = true +tokio.workspace = true diff --git a/litellm-rust/crates/python-bridge/src/lib.rs b/litellm-rust/crates/python-bridge/src/lib.rs index 15e93f7b00c..46ff700f456 100644 --- a/litellm-rust/crates/python-bridge/src/lib.rs +++ b/litellm-rust/crates/python-bridge/src/lib.rs @@ -1,7 +1,6 @@ use std::time::Duration; use litellm_core::error::CoreError; -use litellm_providers::ocr::run_ocr; use pyo3::exceptions::{PyRuntimeError, PyValueError}; use pyo3::prelude::*; use pyo3::types::{PyAny, PyDict}; @@ -9,6 +8,13 @@ use serde_json::{Map, Value}; mod gil; +type MarshaledOcrInputs = ( + Value, + Option>, + Map, + Option, +); + fn py_to_json(py: Python<'_>, value: &Bound<'_, PyAny>) -> PyResult { let json = py.import("json")?; let encoded: String = json.call_method1("dumps", (value,))?.extract()?; @@ -28,53 +34,93 @@ fn json_to_py(py: Python<'_>, value: Value) -> PyResult> { fn core_error_to_pyerr(err: CoreError) -> PyErr { match err { CoreError::Auth(message) => PyValueError::new_err(message), - CoreError::InvalidType { .. } | CoreError::MissingField(_) => { - PyValueError::new_err(err.to_string()) - } + CoreError::InvalidProvider(_) + | CoreError::InvalidType { .. } + | CoreError::MissingField(_) => PyValueError::new_err(err.to_string()), other => PyRuntimeError::new_err(other.to_string()), } } +fn optional_object_to_map( + py: Python<'_>, + name: &'static str, + value: Option>, +) -> PyResult> { + match value { + Some(value) => match py_to_json(py, value.bind(py))? { + Value::Object(map) => Ok(map), + _ => Err(PyValueError::new_err(format!("{name} must be a dict"))), + }, + None => Ok(Map::new()), + } +} + +fn optional_timeout(timeout_seconds: Option) -> Option { + timeout_seconds.and_then(|secs| { + if secs.is_finite() && secs > 0.0 { + Some(Duration::from_secs_f64(secs)) + } else { + None + } + }) +} + +fn marshal_inputs( + py: Python<'_>, + document: Py, + extra_headers: Option>, + optional_params: Option>, + timeout_seconds: Option, +) -> PyResult { + let document = py_to_json(py, document.bind(py))?; + let extra_headers = match extra_headers { + Some(headers) => Some(optional_object_to_map(py, "extra_headers", Some(headers))?), + None => None, + }; + let optional_params = optional_object_to_map(py, "optional_params", optional_params)?; + let timeout = optional_timeout(timeout_seconds); + + Ok((document, extra_headers, optional_params, timeout)) +} + /// Perform a Mistral OCR call end to end and return the response as a dict. #[pyfunction] -#[pyo3(signature = (model, document, api_key=None, api_base=None, optional_params=None, timeout_seconds=None))] +#[pyo3(signature = (model, document, api_key=None, api_base=None, custom_llm_provider=None, extra_headers=None, optional_params=None, timeout_seconds=None))] +#[allow(clippy::too_many_arguments)] fn ocr( py: Python<'_>, model: String, document: Py, api_key: Option, api_base: Option, + custom_llm_provider: Option, + extra_headers: Option>, optional_params: Option>, timeout_seconds: Option, ) -> PyResult> { - let document = py_to_json(py, document.bind(py))?; + let custom_llm_provider = custom_llm_provider.unwrap_or_else(|| "mistral".to_string()); + let (document, extra_headers, optional_params, timeout) = marshal_inputs( + py, + document, + extra_headers, + optional_params, + timeout_seconds, + )?; - let optional_params = match optional_params { - Some(params) => match py_to_json(py, params.bind(py))? { - Value::Object(map) => map, - _ => return Err(PyValueError::new_err("optional_params must be a dict")), - }, - None => Map::new(), - }; - - let timeout = timeout_seconds.and_then(|secs| { - if secs.is_finite() && secs > 0.0 { - Some(Duration::from_secs_f64(secs)) - } else { - None - } - }); - - // Release the GIL during the blocking HTTP call (counted for observability). + // Release the GIL while the sync API waits on async Rust work. let result = gil::release_gil(py, || { - run_ocr( - &model, - document, - api_key.as_deref(), - api_base.as_deref(), - optional_params, - timeout, - ) + pyo3_async_runtimes::tokio::get_runtime().block_on(litellm_providers::ocr::ocr( + litellm_providers::ocr::OcrRequest { + model: &model, + document, + api_key: api_key.as_deref(), + api_base: api_base.as_deref(), + custom_llm_provider: &custom_llm_provider, + extra_headers, + optional_params, + timeout, + }, + )) }); match result { @@ -83,8 +129,50 @@ fn ocr( } } +/// Perform an OCR call end to end and return an asyncio awaitable. +#[pyfunction] +#[pyo3(signature = (model, document, api_key=None, api_base=None, custom_llm_provider=None, extra_headers=None, optional_params=None, timeout_seconds=None))] +#[allow(clippy::too_many_arguments)] +fn aocr( + py: Python<'_>, + model: String, + document: Py, + api_key: Option, + api_base: Option, + custom_llm_provider: Option, + extra_headers: Option>, + optional_params: Option>, + timeout_seconds: Option, +) -> PyResult> { + let custom_llm_provider = custom_llm_provider.unwrap_or_else(|| "mistral".to_string()); + let (document, extra_headers, optional_params, timeout) = marshal_inputs( + py, + document, + extra_headers, + optional_params, + timeout_seconds, + )?; + + pyo3_async_runtimes::tokio::future_into_py(py, async move { + let value = litellm_providers::ocr::ocr(litellm_providers::ocr::OcrRequest { + model: &model, + document, + api_key: api_key.as_deref(), + api_base: api_base.as_deref(), + custom_llm_provider: &custom_llm_provider, + extra_headers, + optional_params, + timeout, + }) + .await + .map_err(core_error_to_pyerr)?; + + Python::with_gil(|py| json_to_py(py, value)) + }) +} + /// Bridge GIL accounting, e.g. `{"releases": 12}`. Lets the Python side observe -/// how often the bridge has dropped the GIL for blocking work. +/// how often the sync bridge has dropped the GIL while awaiting Rust work. #[pyfunction] fn gil_stats(py: Python<'_>) -> PyResult> { let stats = PyDict::new(py); @@ -95,6 +183,7 @@ fn gil_stats(py: Python<'_>) -> PyResult> { #[pymodule] fn litellm_python_bridge(module: &Bound<'_, PyModule>) -> PyResult<()> { module.add_function(wrap_pyfunction!(ocr, module)?)?; + module.add_function(wrap_pyfunction!(aocr, module)?)?; module.add_function(wrap_pyfunction!(gil_stats, module)?)?; Ok(()) } diff --git a/litellm/ocr/main.py b/litellm/ocr/main.py index 3a9ef8db804..98f6c6af8f2 100644 --- a/litellm/ocr/main.py +++ b/litellm/ocr/main.py @@ -4,11 +4,10 @@ Main OCR function for LiteLLM. import asyncio import base64 -import contextvars import mimetypes import os import re -from functools import partial +from dataclasses import dataclass from io import IOBase from typing import Any, Callable, Coroutine, Dict, Optional, Union, cast @@ -20,7 +19,13 @@ from litellm.constants import request_timeout from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.llms.base_llm.ocr.transformation import BaseOCRConfig, OCRResponse from litellm.llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler -from litellm.ocr.rust_bridge import RustOcr, load_rust_ocr, rust_ocr_enabled +from litellm.ocr.rust_bridge import ( + RustAocr, + RustOcr, + load_rust_aocr, + load_rust_ocr, + rust_ocr_enabled, +) from litellm.types.router import GenericLiteLLMParams from litellm.utils import ProviderConfigManager, client @@ -29,6 +34,21 @@ base_llm_http_handler = BaseLLMHTTPHandler() ################################################# +@dataclass +class _PreparedOCRRequest: + model: str + document: dict[str, object] + api_key: Optional[str] + api_base: Optional[str] + custom_llm_provider: str + extra_headers: Optional[dict[str, object]] + provider_config: BaseOCRConfig + optional_params: dict[str, object] + litellm_params: dict[str, object] + effective_timeout: Union[float, httpx.Timeout] + litellm_logging_obj: LiteLLMLoggingObj + + def _timeout_to_seconds( timeout: Optional[Union[float, httpx.Timeout]], ) -> Optional[float]: @@ -45,6 +65,107 @@ def _timeout_to_seconds( return float(timeout) +def _prepare_ocr_request( + model: str, + document: Dict[str, Any], + api_key: Optional[str], + api_base: Optional[str], + timeout: Optional[Union[float, httpx.Timeout]], + custom_llm_provider: Optional[str], + extra_headers: Optional[Dict[str, Any]], + kwargs: dict[str, object], +) -> _PreparedOCRRequest: + litellm_logging_obj = cast(LiteLLMLoggingObj, kwargs.pop("litellm_logging_obj")) + litellm_call_id = cast(Optional[str], kwargs.get("litellm_call_id", None)) + + if not isinstance(document, dict): + raise ValueError( + f"document must be a dict with 'type' and URL/file field, got {type(document)}" + ) + + doc_type = 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'" + ) + + ( + model, + custom_llm_provider, + dynamic_api_key, + dynamic_api_base, + ) = litellm.get_llm_provider( + model=model, + custom_llm_provider=custom_llm_provider, + api_base=api_base, + api_key=api_key, + ) + + if dynamic_api_key: + api_key = dynamic_api_key + if dynamic_api_base: + api_base = dynamic_api_base + + ocr_provider_config = ProviderConfigManager.get_provider_ocr_config( + model=model, + provider=litellm.LlmProviders(custom_llm_provider), + ) + + if ocr_provider_config is None: + raise ValueError(f"OCR is not supported for provider: {custom_llm_provider}") + + verbose_logger.debug(f"OCR call - model: {model}, provider: {custom_llm_provider}") + + litellm_params = GenericLiteLLMParams(**kwargs) + + supported_params = ocr_provider_config.get_supported_ocr_params(model=model) + non_default_params = {} + for param in supported_params: + if param in kwargs: + non_default_params[param] = kwargs.pop(param) + + optional_params = ocr_provider_config.map_ocr_params( + non_default_params=non_default_params, + optional_params={}, + model=model, + ) + + verbose_logger.debug(f"OCR optional_params after mapping: {optional_params}") + + effective_timeout = timeout or request_timeout + + litellm_logging_obj.update_from_kwargs( + kwargs=kwargs, + model=model, + optional_params=optional_params, + litellm_params={ + "litellm_call_id": litellm_call_id, + "api_base": api_base, + }, + custom_llm_provider=custom_llm_provider, + ) + + return _PreparedOCRRequest( + model=model, + document=cast(dict[str, object], document), + api_key=api_key, + api_base=api_base, + custom_llm_provider=custom_llm_provider, + extra_headers=cast(Optional[dict[str, object]], extra_headers), + provider_config=ocr_provider_config, + optional_params=cast(dict[str, object], optional_params), + litellm_params=dict(litellm_params), + effective_timeout=effective_timeout, + litellm_logging_obj=litellm_logging_obj, + ) + + def _run_rust_ocr( rust_ocr: RustOcr, logging_obj: LiteLLMLoggingObj, @@ -54,6 +175,8 @@ def _run_rust_ocr( document: dict[str, object], api_key: Optional[str], api_base: Optional[str], + custom_llm_provider: str, + extra_headers: Optional[dict[str, object]], optional_params: dict[str, object], litellm_params: dict[str, object], timeout_seconds: Optional[float], @@ -68,7 +191,7 @@ def _run_rust_ocr( """ resolved_api_key = api_key or resolve_api_key("MISTRAL_API_KEY") resolved_headers = provider_config.validate_environment( - headers={}, + headers=extra_headers or {}, model=model, api_key=resolved_api_key, api_base=api_base, @@ -99,6 +222,64 @@ def _run_rust_ocr( document=document, api_key=resolved_api_key, api_base=api_base, + custom_llm_provider=custom_llm_provider, + extra_headers=extra_headers, + optional_params=optional_params, + timeout_seconds=timeout_seconds, + ) + ) + + +async def _run_rust_aocr( + rust_aocr: RustAocr, + logging_obj: LiteLLMLoggingObj, + provider_config: BaseOCRConfig, + resolve_api_key: Callable[[str], Optional[str]], + model: str, + document: dict[str, object], + api_key: Optional[str], + api_base: Optional[str], + custom_llm_provider: str, + extra_headers: Optional[dict[str, object]], + optional_params: dict[str, object], + litellm_params: dict[str, object], + timeout_seconds: Optional[float], +) -> OCRResponse: + resolved_api_key = api_key or resolve_api_key("MISTRAL_API_KEY") + resolved_headers = provider_config.validate_environment( + headers=extra_headers or {}, + model=model, + api_key=resolved_api_key, + api_base=api_base, + litellm_params=litellm_params, + ) + resolved_complete_url = provider_config.get_complete_url( + api_base=api_base, + model=model, + optional_params=optional_params, + litellm_params=litellm_params, + ) + logging_obj.pre_call( + input="OCR document processing", + api_key=resolved_api_key, + additional_args={ + "complete_input_dict": { + "model": model, + "document": document, + **optional_params, + }, + "api_base": resolved_complete_url, + "headers": resolved_headers, + }, + ) + return OCRResponse.model_validate( + await rust_aocr( + model=model, + document=document, + api_key=resolved_api_key, + api_base=api_base, + custom_llm_provider=custom_llm_provider, + extra_headers=extra_headers, optional_params=optional_params, timeout_seconds=timeout_seconds, ) @@ -176,17 +357,7 @@ async def aocr( """ local_vars = locals() try: - loop = asyncio.get_event_loop() - kwargs["aocr"] = True - - # Get custom llm provider - if custom_llm_provider is None: - _, custom_llm_provider, _, _ = litellm.get_llm_provider( - model=model, api_base=api_base - ) - - func = partial( - ocr, + prepared = _prepare_ocr_request( model=model, document=document, api_key=api_key, @@ -194,17 +365,56 @@ async def aocr( timeout=timeout, custom_llm_provider=custom_llm_provider, extra_headers=extra_headers, - **kwargs, + kwargs=kwargs, ) - ctx = contextvars.copy_context() - func_with_context = partial(ctx.run, func) - init_response = await loop.run_in_executor(None, func_with_context) + if prepared.custom_llm_provider == "mistral" and rust_ocr_enabled(): + rust_aocr = load_rust_aocr() + if rust_aocr is None: + verbose_logger.debug( + "Async Rust OCR bridge unavailable; falling back to Python path" + ) + else: + from litellm.secret_managers.main import get_secret_str - if asyncio.iscoroutine(init_response): - response = await init_response - else: - response = init_response + response = await _run_rust_aocr( + rust_aocr=rust_aocr, + logging_obj=prepared.litellm_logging_obj, + provider_config=prepared.provider_config, + resolve_api_key=get_secret_str, + model=prepared.model, + document=prepared.document, + api_key=prepared.api_key, + api_base=prepared.api_base, + custom_llm_provider=prepared.custom_llm_provider, + extra_headers=prepared.extra_headers, + optional_params=prepared.optional_params, + litellm_params=prepared.litellm_params, + timeout_seconds=_timeout_to_seconds(prepared.effective_timeout), + ) + if response is None: + raise ValueError( + f"Got an unexpected None response from the OCR API: {response}" + ) + return response + + response = base_llm_http_handler.ocr( + model=prepared.model, + document=prepared.document, + optional_params=prepared.optional_params, + timeout=prepared.effective_timeout, + logging_obj=prepared.litellm_logging_obj, + api_key=prepared.api_key, + api_base=prepared.api_base, + custom_llm_provider=prepared.custom_llm_provider, + aocr=True, + headers=prepared.extra_headers, + provider_config=prepared.provider_config, + litellm_params=prepared.litellm_params, + ) + + if asyncio.iscoroutine(response): + response = await response if response is None: raise ValueError( @@ -297,94 +507,20 @@ def ocr( """ local_vars = locals() try: - litellm_logging_obj = cast(LiteLLMLoggingObj, kwargs.pop("litellm_logging_obj")) - litellm_call_id: Optional[str] = kwargs.get("litellm_call_id", None) _is_async = kwargs.pop("aocr", False) is True - - # Validate document parameter format - if not isinstance(document, dict): - raise ValueError( - f"document must be a dict with 'type' and URL/file field, got {type(document)}" - ) - - doc_type = document.get("type") - - # Handle file type: convert to document_url/image_url with base64 data URI - 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'" - ) - - ( - model, - custom_llm_provider, - dynamic_api_key, - dynamic_api_base, - ) = litellm.get_llm_provider( + prepared = _prepare_ocr_request( model=model, - custom_llm_provider=custom_llm_provider, - api_base=api_base, + document=document, api_key=api_key, - ) - - # Update with dynamic values if available - if dynamic_api_key: - api_key = dynamic_api_key - if dynamic_api_base: - api_base = dynamic_api_base - - ocr_provider_config: Optional[BaseOCRConfig] = ( - ProviderConfigManager.get_provider_ocr_config( - model=model, - provider=litellm.LlmProviders(custom_llm_provider), - ) - ) - - if ocr_provider_config is None: - raise ValueError( - f"OCR is not supported for provider: {custom_llm_provider}" - ) - - verbose_logger.debug( - f"OCR call - model: {model}, provider: {custom_llm_provider}" - ) - - litellm_params = GenericLiteLLMParams(**kwargs) - - supported_params = ocr_provider_config.get_supported_ocr_params(model=model) - non_default_params = {} - for param in supported_params: - if param in kwargs: - non_default_params[param] = kwargs.pop(param) - - optional_params = ocr_provider_config.map_ocr_params( - non_default_params=non_default_params, - optional_params={}, - model=model, - ) - - verbose_logger.debug(f"OCR optional_params after mapping: {optional_params}") - - effective_timeout = timeout or request_timeout - - litellm_logging_obj.update_from_kwargs( + api_base=api_base, kwargs=kwargs, - model=model, - optional_params=optional_params, - litellm_params={ - "litellm_call_id": litellm_call_id, - "api_base": api_base, - }, custom_llm_provider=custom_llm_provider, + extra_headers=extra_headers, + timeout=timeout, ) # Optional Rust path: hand the whole Mistral OCR call to the Rust bridge. - if custom_llm_provider == "mistral" and rust_ocr_enabled(): + if prepared.custom_llm_provider == "mistral" and rust_ocr_enabled(): rust_ocr = load_rust_ocr() if rust_ocr is None: verbose_logger.debug( @@ -395,31 +531,33 @@ def ocr( return _run_rust_ocr( rust_ocr=rust_ocr, - logging_obj=litellm_logging_obj, - provider_config=ocr_provider_config, + logging_obj=prepared.litellm_logging_obj, + provider_config=prepared.provider_config, resolve_api_key=get_secret_str, - model=model, - document=document, - api_key=api_key, - api_base=api_base, - optional_params=optional_params, - litellm_params=dict(litellm_params), - timeout_seconds=_timeout_to_seconds(effective_timeout), + model=prepared.model, + document=prepared.document, + api_key=prepared.api_key, + api_base=prepared.api_base, + custom_llm_provider=prepared.custom_llm_provider, + extra_headers=prepared.extra_headers, + optional_params=prepared.optional_params, + litellm_params=prepared.litellm_params, + timeout_seconds=_timeout_to_seconds(prepared.effective_timeout), ) response = base_llm_http_handler.ocr( - model=model, - document=document, - optional_params=optional_params, - timeout=effective_timeout, - logging_obj=litellm_logging_obj, - api_key=api_key, - api_base=api_base, - custom_llm_provider=custom_llm_provider, + model=prepared.model, + document=prepared.document, + optional_params=prepared.optional_params, + timeout=prepared.effective_timeout, + logging_obj=prepared.litellm_logging_obj, + api_key=prepared.api_key, + api_base=prepared.api_base, + custom_llm_provider=prepared.custom_llm_provider, aocr=_is_async, - headers=extra_headers, - provider_config=ocr_provider_config, - litellm_params=dict(litellm_params), + headers=prepared.extra_headers, + provider_config=prepared.provider_config, + litellm_params=prepared.litellm_params, ) return response diff --git a/litellm/ocr/rust_bridge.py b/litellm/ocr/rust_bridge.py index 61f9e8ca69a..0dec57d9168 100644 --- a/litellm/ocr/rust_bridge.py +++ b/litellm/ocr/rust_bridge.py @@ -11,7 +11,7 @@ can import it statically without forming an import cycle. from __future__ import annotations -from typing import Final, Protocol, cast +from typing import Awaitable, Final, Protocol, cast class RustOcr(Protocol): @@ -23,11 +23,29 @@ class RustOcr(Protocol): document: dict[str, object], api_key: str | None, api_base: str | None, + custom_llm_provider: str, + extra_headers: dict[str, object] | None, optional_params: dict[str, object], timeout_seconds: float | None, ) -> dict[str, object]: ... +class RustAocr(Protocol): + """Signature of the compiled ``litellm_python_bridge.aocr`` entrypoint.""" + + def __call__( + self, + model: str, + document: dict[str, object], + api_key: str | None, + api_base: str | None, + custom_llm_provider: str, + extra_headers: dict[str, object] | None, + optional_params: dict[str, object], + timeout_seconds: float | None, + ) -> Awaitable[dict[str, object]]: ... + + class _Unset: """Sentinel type so ``ocr=None`` can clear a prior injection while omission preserves it.""" @@ -36,21 +54,27 @@ _UNSET: Final[_Unset] = _Unset() _rust_ocr_enabled = False _rust_ocr_impl: RustOcr | None = None +_rust_aocr_impl: RustAocr | None = None def use_litellm_rust( - enabled: bool = True, *, ocr: RustOcr | None | _Unset = _UNSET + enabled: bool = True, + *, + ocr: RustOcr | None | _Unset = _UNSET, + aocr: RustAocr | None | _Unset = _UNSET, ) -> None: """Route supported OCR calls through the Rust ``litellm_python_bridge`` extension. - ``ocr`` injects the bridge callable; when omitted the compiled extension is - loaded on demand and any previously injected bridge is preserved. Pass - ``ocr=None`` explicitly to clear a prior injection. + ``ocr`` and ``aocr`` inject bridge callables; when omitted the compiled + extension is loaded on demand and any previously injected bridge is + preserved. Pass ``None`` explicitly to clear a prior injection. """ - global _rust_ocr_enabled, _rust_ocr_impl + global _rust_ocr_enabled, _rust_ocr_impl, _rust_aocr_impl _rust_ocr_enabled = enabled if not isinstance(ocr, _Unset): _rust_ocr_impl = ocr + if not isinstance(aocr, _Unset): + _rust_aocr_impl = aocr def rust_ocr_enabled() -> bool: @@ -72,3 +96,14 @@ def load_rust_ocr() -> RustOcr | None: except ImportError: return None return cast(RustOcr, litellm_python_bridge.ocr) + + +def load_rust_aocr() -> RustAocr | None: + """Return the async Rust OCR callable, or ``None`` when unavailable.""" + if _rust_aocr_impl is not None: + return _rust_aocr_impl + try: + import litellm_python_bridge + except ImportError: + return None + return cast(RustAocr, getattr(litellm_python_bridge, "aocr", None)) diff --git a/tests/test_litellm/ocr/test_rust_bridge.py b/tests/test_litellm/ocr/test_rust_bridge.py index 7e028064e4c..aee8506b84b 100644 --- a/tests/test_litellm/ocr/test_rust_bridge.py +++ b/tests/test_litellm/ocr/test_rust_bridge.py @@ -35,7 +35,15 @@ class RecordingBridge: self.calls = [] def __call__( - self, model, document, api_key, api_base, optional_params, timeout_seconds + self, + model, + document, + api_key, + api_base, + custom_llm_provider, + extra_headers, + optional_params, + timeout_seconds, ): self.calls.append( { @@ -43,6 +51,40 @@ class RecordingBridge: "document": document, "api_key": api_key, "api_base": api_base, + "custom_llm_provider": custom_llm_provider, + "extra_headers": extra_headers, + "optional_params": optional_params, + "timeout_seconds": timeout_seconds, + } + ) + return dict(FAKE_OCR_RESPONSE) + + +class RecordingAsyncBridge: + """A fake async ``RustAocr`` callable that records the args it was handed.""" + + def __init__(self): + self.calls = [] + + async def __call__( + self, + model, + document, + api_key, + api_base, + custom_llm_provider, + extra_headers, + optional_params, + timeout_seconds, + ): + self.calls.append( + { + "model": model, + "document": document, + "api_key": api_key, + "api_base": api_base, + "custom_llm_provider": custom_llm_provider, + "extra_headers": extra_headers, "optional_params": optional_params, "timeout_seconds": timeout_seconds, } @@ -70,7 +112,7 @@ class FakeOCRConfig: def validate_environment( self, *, headers, model, api_key, api_base, litellm_params ): - return {"authorization": f"Bearer {api_key}"} + return {"Authorization": f"Bearer {api_key}", **headers} def get_complete_url(self, *, api_base, model, optional_params, litellm_params): return f"{api_base or 'https://api.mistral.ai/v1'}/ocr" @@ -79,9 +121,9 @@ class FakeOCRConfig: @pytest.fixture(autouse=True) def _reset_rust_flag(): """Keep the global toggle isolated between tests.""" - rust_bridge.use_litellm_rust(False, ocr=None) + rust_bridge.use_litellm_rust(False, ocr=None, aocr=None) yield - rust_bridge.use_litellm_rust(False, ocr=None) + rust_bridge.use_litellm_rust(False, ocr=None, aocr=None) @pytest.fixture @@ -92,6 +134,14 @@ def fake_bridge(): return bridge +@pytest.fixture +def fake_async_bridge(): + """Enable the async Rust path with an injected recording bridge.""" + bridge = RecordingAsyncBridge() + litellm.use_litellm_rust(True, aocr=bridge) + return bridge + + def test_use_litellm_rust_toggles_flag(): assert rust_bridge.rust_ocr_enabled() is False litellm.use_litellm_rust() @@ -106,6 +156,12 @@ def test_load_rust_ocr_returns_injected_impl(): assert rust_bridge.load_rust_ocr() is bridge +def test_load_rust_aocr_returns_injected_impl(): + bridge = RecordingAsyncBridge() + litellm.use_litellm_rust(True, aocr=bridge) + assert rust_bridge.load_rust_aocr() is bridge + + def test_toggle_without_ocr_arg_preserves_injected_impl(): """Regression: routine enable/disable calls must not clobber a prior injection. @@ -114,20 +170,25 @@ def test_toggle_without_ocr_arg_preserves_injected_impl(): a caller toggled the flag without re-passing ``ocr=``. """ bridge = RecordingBridge() - litellm.use_litellm_rust(True, ocr=bridge) + async_bridge = RecordingAsyncBridge() + litellm.use_litellm_rust(True, ocr=bridge, aocr=async_bridge) litellm.use_litellm_rust(False) assert rust_bridge.load_rust_ocr() is bridge + assert rust_bridge.load_rust_aocr() is async_bridge litellm.use_litellm_rust(True) assert rust_bridge.load_rust_ocr() is bridge + assert rust_bridge.load_rust_aocr() is async_bridge def test_explicit_ocr_none_clears_injected_impl(): bridge = RecordingBridge() - litellm.use_litellm_rust(True, ocr=bridge) + async_bridge = RecordingAsyncBridge() + litellm.use_litellm_rust(True, ocr=bridge, aocr=async_bridge) - litellm.use_litellm_rust(True, ocr=None) + litellm.use_litellm_rust(True, ocr=None, aocr=None) assert rust_bridge.load_rust_ocr() is None + assert rust_bridge.load_rust_aocr() is None def test_load_rust_ocr_none_when_extension_absent(): @@ -135,6 +196,7 @@ def test_load_rust_ocr_none_when_extension_absent(): caller degrades to the Python path instead of raising ImportError.""" litellm.use_litellm_rust(True) # no impl injected; extension isn't built in CI assert rust_bridge.load_rust_ocr() is None + assert rust_bridge.load_rust_aocr() is None def test_load_rust_ocr_uses_compiled_extension(monkeypatch): @@ -143,10 +205,12 @@ def test_load_rust_ocr_uses_compiled_extension(monkeypatch): built in CI, so stand in a fake module via ``sys.modules``.""" fake_module = types.ModuleType("litellm_python_bridge") fake_module.ocr = lambda **kwargs: dict(FAKE_OCR_RESPONSE) # type: ignore[attr-defined] + fake_module.aocr = lambda **kwargs: dict(FAKE_OCR_RESPONSE) # type: ignore[attr-defined] monkeypatch.setitem(sys.modules, "litellm_python_bridge", fake_module) litellm.use_litellm_rust(True) # enabled, no impl injected -> import the extension assert rust_bridge.load_rust_ocr() is fake_module.ocr + assert rust_bridge.load_rust_aocr() is fake_module.aocr def test_timeout_to_seconds_handles_float_timeout_and_none(): @@ -168,6 +232,8 @@ def test_run_rust_ocr_forwards_args_and_wraps_response(): document=DOCUMENT, api_key="sk-test", api_base="https://proxy.internal", + custom_llm_provider="mistral", + extra_headers={"x-trace-id": "trace-1"}, optional_params={"include_image_base64": True}, litellm_params={}, timeout_seconds=12.5, @@ -181,6 +247,8 @@ def test_run_rust_ocr_forwards_args_and_wraps_response(): "document": DOCUMENT, "api_key": "sk-test", "api_base": "https://proxy.internal", + "custom_llm_provider": "mistral", + "extra_headers": {"x-trace-id": "trace-1"}, "optional_params": {"include_image_base64": True}, "timeout_seconds": 12.5, } @@ -202,6 +270,8 @@ def test_run_rust_ocr_resolves_key_via_secret_manager_when_missing(): document=DOCUMENT, api_key=None, api_base=None, + custom_llm_provider="mistral", + extra_headers=None, optional_params={}, litellm_params={}, timeout_seconds=None, @@ -227,6 +297,8 @@ def test_run_rust_ocr_prefers_explicit_key_over_resolver(): document=DOCUMENT, api_key="sk-explicit", api_base=None, + custom_llm_provider="mistral", + extra_headers=None, optional_params={}, litellm_params={}, timeout_seconds=None, @@ -249,6 +321,8 @@ def test_run_rust_ocr_runs_pre_call_logging(): document=DOCUMENT, api_key="sk-test", api_base="https://api.mistral.ai/v1", + custom_llm_provider="mistral", + extra_headers={"x-trace-id": "trace-1"}, optional_params={"include_image_base64": True}, litellm_params={}, timeout_seconds=None, @@ -262,7 +336,10 @@ def test_run_rust_ocr_runs_pre_call_logging(): assert complete_input["include_image_base64"] is True # The logged request mirrors what Rust sends: resolved URL + headers. assert additional_args["api_base"] == "https://api.mistral.ai/v1/ocr" - assert additional_args["headers"] == {"authorization": "Bearer sk-test"} + assert additional_args["headers"] == { + "Authorization": "Bearer sk-test", + "x-trace-id": "trace-1", + } def test_ocr_routes_to_rust_when_enabled(fake_bridge): @@ -270,6 +347,7 @@ def test_ocr_routes_to_rust_when_enabled(fake_bridge): model=MODEL, document=DOCUMENT, api_key="sk-test", + extra_headers={"x-trace-id": "trace-1"}, include_image_base64=True, ) @@ -281,10 +359,34 @@ def test_ocr_routes_to_rust_when_enabled(fake_bridge): assert call["model"] == "mistral-ocr-latest" assert call["document"] == DOCUMENT assert call["api_key"] == "sk-test" + assert call["custom_llm_provider"] == "mistral" + assert call["extra_headers"] == {"x-trace-id": "trace-1"} # Raw OCR params ride along in optional_params; Rust filters to supported keys. assert call["optional_params"].get("include_image_base64") is True +@pytest.mark.asyncio +async def test_aocr_routes_to_async_rust_when_enabled(fake_async_bridge): + response = await litellm.aocr( + model=MODEL, + document=DOCUMENT, + api_key="sk-test", + extra_headers={"x-trace-id": "trace-1"}, + include_image_base64=True, + ) + + assert isinstance(response, OCRResponse) + assert response.pages[0].markdown == "hello world" + assert len(fake_async_bridge.calls) == 1 + call = fake_async_bridge.calls[0] + assert call["model"] == "mistral-ocr-latest" + assert call["document"] == DOCUMENT + assert call["api_key"] == "sk-test" + assert call["custom_llm_provider"] == "mistral" + assert call["extra_headers"] == {"x-trace-id": "trace-1"} + assert call["optional_params"].get("include_image_base64") is True + + def test_ocr_forwards_timeout_to_rust(fake_bridge): """Caller-supplied timeout must flow into the Rust bridge so the fixed 600s client ceiling doesn't silently override shorter deadlines."""