From a545c493d74e470ac7e1a3bec8227b7a1d4a4012 Mon Sep 17 00:00:00 2001 From: yucheng-berri Date: Thu, 25 Jun 2026 11:59:35 -0700 Subject: [PATCH 1/8] fix(otel): hashable scope for _emit_once when guardrail_mode is list (#31262) * fix(otel): hashable scope for _emit_once when guardrail_mode is list `_emit_once` keys `spans_logged` by `(class, id, *scope)`. When a guardrail entry's `guardrail_mode` arrives as a `List[GuardrailEventHooks]` (the shape Presidio expands to with `output_parse_pii: true`, and the shape `event_hook` carries for any `mode: [...]` in config), the tuple contains a list and `spans_logged.get(dedupe_key)` raises `TypeError: unhashable type: 'list'`. On the post-call path this fires inside the logging callback and is swallowed; the request returns 200 but the OTEL `guardrail` span is silently dropped. On the blocking path the same error surfaces as HTTP 500. Adds `_freeze_for_dedupe`, a small recursive normalizer that turns lists and tuples into tuples, sets into frozensets, dicts into frozensets of `(key, value)` pairs, and falls back to `repr` for arbitrary unhashables. Applied inside `_emit_once` before the dict lookup, so all three callsites are protected without touching the guardrail-specific callsite. Helper assumes acyclic input; `guardrail_mode` values are built fresh from config (str enums, lists of str enums, TypedDict of str/list-of-str), so no cycle can arise in practice. Regression tests in `TestOpenTelemetrySpanDedupe` cover the list crash, distinct-list-scope collision, dict and set scope parts, and an end-to-end `_create_guardrail_span` exercise that confirms exactly one `guardrail` span is emitted across repeated lifecycle entrypoints. Each new test fails on a reverted helper (4/4 mutation kill) * fix(otel): cap _freeze_for_dedupe recursion depth and ignore in recursive detector CI's recursive_detector blocks new recursive functions in litellm/ unless they are in the allowlist with a documented bound. Cap the helper at 16 levels and return repr(value) past the cap; this is well past the realistic depth of guardrail_mode (1-3 levels) and means a future caller passing a cyclic container can no longer push the proxy logging path into a RecursionError. Add a regression test that exercises the cycle path. * refactor(otel): annotate _freeze_for_dedupe return as a HashableScope union Per review feedback from @mateo-berri: replace the loose `-> object` annotation with a recursive `HashableScope` union (str | int | float | bool | bytes | None | Tuple[HashableScope, ...] | FrozenSet[HashableScope]) so the helper's contract is visible at the signature. Replace the `try/except hash(value); return value` passthrough with an explicit isinstance check over the hashable-scalar types so the type checker can narrow without requiring `cast(Hashable, value)` on the return. Symmetric: dict keys also flow through the freezer (a TypedDict key is already a string in practice, so behaviorally identical). All 16 regression tests still pass; mutation kill behavior preserved * fix: avoid explicit casting --------- Co-authored-by: Mateo Wang <277851410+mateo-berri@users.noreply.github.com> --- litellm/integrations/opentelemetry.py | 47 +++++++- .../code_coverage_tests/recursive_detector.py | 1 + .../integrations/test_opentelemetry.py | 103 ++++++++++++++++++ 3 files changed, 146 insertions(+), 5 deletions(-) diff --git a/litellm/integrations/opentelemetry.py b/litellm/integrations/opentelemetry.py index 6b50ef49b49..c652141a2b6 100644 --- a/litellm/integrations/opentelemetry.py +++ b/litellm/integrations/opentelemetry.py @@ -189,6 +189,37 @@ def _normalize_team_metadata_keys(value: Any) -> List[str]: return [str(item).strip() for item in value if str(item).strip()] +_FREEZE_MAX_DEPTH = 16 + +HashableScope = Union[ + str, + int, + float, + bool, + bytes, + None, + tuple["HashableScope", ...], + frozenset["HashableScope"], +] + + +def _freeze_for_dedupe(value: object, _depth: int = 0) -> HashableScope: + if _depth >= _FREEZE_MAX_DEPTH: + return repr(value) + if isinstance(value, (list, tuple)): + return tuple(_freeze_for_dedupe(item, _depth + 1) for item in value) + if isinstance(value, set): + return frozenset(_freeze_for_dedupe(item, _depth + 1) for item in value) + if isinstance(value, dict): + return frozenset( + (_freeze_for_dedupe(key, _depth + 1), _freeze_for_dedupe(item, _depth + 1)) + for key, item in value.items() + ) + if isinstance(value, (str, int, float, bytes)) or value is None: + return value + return repr(value) + + @dataclass class OpenTelemetryConfig: exporter: Union[str, SpanExporter] = "console" @@ -1073,10 +1104,12 @@ class OpenTelemetry(OTELGenAISemconvMixin, CustomLogger): can be re-read with mutated entries between calls, so dedupe must be at entry granularity. Scope: the entry's stable identity. - ``scope`` parts can be any hashable identity. The marker is stored - in ``kwargs["litellm_params"]["metadata"]["_otel_internal"]`` so it - is request-local (kwargs is shared across the sync/async callbacks - and lifecycle hooks for one request). + ``scope`` parts may include unhashable containers (list, dict, set); + they are normalized into a hashable shape via ``_freeze_for_dedupe`` + before keying the marker dict. The marker is stored in + ``kwargs["litellm_params"]["metadata"]["_otel_internal"]`` so it is + request-local (kwargs is shared across the sync/async callbacks and + lifecycle hooks for one request). """ litellm_params = kwargs.get("litellm_params") if not isinstance(litellm_params, dict): @@ -1098,7 +1131,11 @@ class OpenTelemetry(OTELGenAISemconvMixin, CustomLogger): spans_logged = {} _otel_internal["spans_logged"] = spans_logged - dedupe_key = (self.__class__.__name__, id(self), *scope) + dedupe_key = ( + self.__class__.__name__, + id(self), + *(_freeze_for_dedupe(part) for part in scope), + ) if spans_logged.get(dedupe_key) is True: return False diff --git a/tests/code_coverage_tests/recursive_detector.py b/tests/code_coverage_tests/recursive_detector.py index 1d11d676207..6d0e12314f0 100644 --- a/tests/code_coverage_tests/recursive_detector.py +++ b/tests/code_coverage_tests/recursive_detector.py @@ -51,6 +51,7 @@ IGNORE_FUNCTIONS = [ "_resolve", # OCI: $ref resolver bounded by `resolving_stack` cycle guard. "resolve_oci_schema_anyof", # OCI: bounded by JSON-schema tree depth (no cycles possible in well-formed input). "sanitize_oci_schema", # OCI: bounded by JSON-schema tree depth. + "_freeze_for_dedupe", # OTEL: max depth set (default 16, _FREEZE_MAX_DEPTH); fails closed by returning repr(value) at the cap. ] diff --git a/tests/test_litellm/integrations/test_opentelemetry.py b/tests/test_litellm/integrations/test_opentelemetry.py index e47e437a131..7ffd09b931f 100644 --- a/tests/test_litellm/integrations/test_opentelemetry.py +++ b/tests/test_litellm/integrations/test_opentelemetry.py @@ -4690,6 +4690,109 @@ class TestOpenTelemetrySpanDedupe(unittest.TestCase): self.assertTrue(otel._emit_once(kwargs, "success")) self.assertFalse(otel._emit_once(kwargs, "success")) + def test_emit_once_accepts_list_valued_scope_part(self): + """Regression for LIT-3428 / LIT-3764: a list-valued ``guardrail_mode`` + (the shape Presidio expands to with ``output_parse_pii: true``) must + not raise ``TypeError: unhashable type: 'list'`` when building the + dedupe key. Pre-fix, this call crashed inside ``dict.get``.""" + otel = OpenTelemetry() + kwargs = self._build_kwargs() + self.assertTrue( + otel._emit_once(kwargs, "guardrail", "pii", 1.0, ["pre_call", "post_call"]) + ) + self.assertFalse( + otel._emit_once(kwargs, "guardrail", "pii", 1.0, ["pre_call", "post_call"]), + "Same list scope must dedupe to False on the second call", + ) + + def test_emit_once_distinct_list_scopes_dont_collide(self): + """Two different list-valued scopes on the same handler/kwargs must + each emit exactly once. Catches a regression where every list collapses + to the same key (e.g. ``str(list)`` collisions on near-identical input).""" + otel = OpenTelemetry() + kwargs = self._build_kwargs() + self.assertTrue(otel._emit_once(kwargs, "guardrail", "pii", 1.0, ["pre_call"])) + self.assertTrue( + otel._emit_once(kwargs, "guardrail", "pii", 1.0, ["pre_call", "post_call"]), + "Distinct list scopes must produce distinct dedupe keys", + ) + self.assertFalse(otel._emit_once(kwargs, "guardrail", "pii", 1.0, ["pre_call"])) + self.assertFalse( + otel._emit_once(kwargs, "guardrail", "pii", 1.0, ["pre_call", "post_call"]) + ) + + def test_emit_once_accepts_dict_and_set_scope_parts(self): + """``guardrail_mode`` can also arrive as a ``GuardrailMode`` TypedDict + (i.e. a plain dict at runtime). Sets are not produced today but flow + through the same normalization. Both must hash without raising.""" + otel = OpenTelemetry() + kwargs = self._build_kwargs() + self.assertTrue( + otel._emit_once(kwargs, "guardrail", "pii", 1.0, {"tags": ["pre", "post"]}) + ) + self.assertFalse( + otel._emit_once(kwargs, "guardrail", "pii", 1.0, {"tags": ["pre", "post"]}) + ) + self.assertTrue(otel._emit_once(kwargs, "guardrail", "pii", 1.0, {"a", "b"})) + + def test_emit_once_handles_self_referential_scope_without_recursion_error(self): + """``_freeze_for_dedupe`` caps recursion at ``_FREEZE_MAX_DEPTH`` and + falls back to ``repr`` past the cap, so a self-referential container + in scope must not crash ``_emit_once``. ``guardrail_mode`` cannot + construct such input today, but the cap is the bound that justifies + recursion on the logging hot path.""" + otel = OpenTelemetry() + kwargs = self._build_kwargs() + cyclic: list = [] + cyclic.append(cyclic) + self.assertTrue(otel._emit_once(kwargs, "guardrail", "pii", 1.0, cyclic)) + self.assertFalse(otel._emit_once(kwargs, "guardrail", "pii", 1.0, cyclic)) + + def test_create_guardrail_span_does_not_raise_on_list_mode(self): + """End-to-end regression for LIT-3428: ``_create_guardrail_span`` + must produce exactly one span (not raise ``TypeError``) when the + guardrail entry's ``guardrail_mode`` is a list.""" + span_exporter = InMemorySpanExporter() + tracer_provider = TracerProvider() + tracer_provider.add_span_processor(SimpleSpanProcessor(span_exporter)) + + otel = OpenTelemetry(tracer_provider=tracer_provider) + otel.tracer = tracer_provider.get_tracer(__name__) + + kwargs = { + "model": "gpt-4", + "messages": [{"role": "user", "content": "Hello"}], + "litellm_params": {"custom_llm_provider": "openai", "metadata": {}}, + "standard_logging_object": { + "id": "test-id", + "call_type": "completion", + "metadata": {}, + "hidden_params": {}, + "guardrail_information": [ + { + "guardrail_name": "presidio-pii", + "guardrail_mode": ["pre_call", "post_call"], + "guardrail_response": "ok", + "start_time": 1.0, + "end_time": 2.0, + } + ], + }, + } + + otel._create_guardrail_span(kwargs=kwargs, context=None) + otel._create_guardrail_span(kwargs=kwargs, context=None) + + guardrail_spans = [ + s for s in span_exporter.get_finished_spans() if s.name == "guardrail" + ] + self.assertEqual( + len(guardrail_spans), + 1, + "List-valued guardrail_mode must emit exactly one guardrail span " + "across repeated lifecycle entrypoints", + ) + def test_handle_success_emits_single_litellm_request_span_on_double_call(self): """Sync + async callback paths firing for the same kwargs must result in exactly one litellm_request span.""" From d8ef1da49d9fbd13a48bd86837bc6897a0b11994 Mon Sep 17 00:00:00 2001 From: ishaan-berri <155045088+ishaan-berri@users.noreply.github.com> Date: Thu, 25 Jun 2026 12:32:55 -0700 Subject: [PATCH 2/8] feat: package Rust OCR bridge in LiteLLM wheel (#31267) * feat: package rust ocr bridge in litellm wheel * Install Rust in Windows CircleCI job * Address Rust wheel review feedback * Pin Windows rustup installer hash --- .circleci/config.yml | 23 +++++++ .dockerignore | 2 + Dockerfile | 1 + docker/Dockerfile.non_root | 1 + litellm-rust/crates/python-bridge/Cargo.toml | 2 +- litellm-rust/crates/python-bridge/build.rs | 6 ++ litellm-rust/crates/python-bridge/src/lib.rs | 2 +- litellm/ocr/rust_bridge.py | 17 ++--- litellm/rust_bridge/__init__.py | 8 +++ litellm/rust_bridge/loader.py | 28 ++++++++ pyproject.toml | 33 ++++----- tests/test_litellm/ocr/test_rust_bridge.py | 70 ++++++++++++++++++-- uv.lock | 6 +- 13 files changed, 164 insertions(+), 35 deletions(-) create mode 100644 litellm-rust/crates/python-bridge/build.rs create mode 100644 litellm/rust_bridge/__init__.py create mode 100644 litellm/rust_bridge/loader.py diff --git a/.circleci/config.yml b/.circleci/config.yml index abcdbf45187..d532d3106e5 100644 --- a/.circleci/config.yml +++ b/.circleci/config.yml @@ -205,6 +205,24 @@ jobs: environment: UV_HTTP_TIMEOUT: "300" command: | + $rustupInit = Join-Path $env:TEMP "rustup-init.exe" + $rustupVersion = "1.28.2" + $rustupUrl = "https://static.rust-lang.org/rustup/archive/$rustupVersion/x86_64-pc-windows-msvc/rustup-init.exe" + Invoke-WebRequest -Uri $rustupUrl -OutFile $rustupInit + $rustupExpected = "88d8258dcf6ae4f7a80c7d1088e1f36fa7025a1cfd1343731b4ee6f385121fc0" + $rustupActual = (Get-FileHash -Path $rustupInit -Algorithm SHA256).Hash.ToLower() + if ($rustupActual -ne $rustupExpected) { + throw "rustup installer hash mismatch: expected $rustupExpected got $rustupActual" + } + & $rustupInit -y --profile minimal --default-toolchain stable + if ($LASTEXITCODE -ne 0) { + exit $LASTEXITCODE + } + Remove-Item $rustupInit + $cargoBin = Join-Path $HOME ".cargo\bin" + $env:Path = "$cargoBin;$env:Path" + rustc --version + cargo --version $installer = Join-Path $env:TEMP "uv-install.ps1" Invoke-WebRequest -Uri https://astral.sh/uv/0.10.9/install.ps1 -OutFile $installer $expected = "d43ffff8d28e7d1e7d1831a212465f12b24a43c7f87f386e95e2a5915aee5d7d" @@ -222,6 +240,9 @@ jobs: if (-not (Select-String -Path $PROFILE -SimpleMatch $uvBin -Quiet)) { Add-Content -Path $PROFILE -Value "`$env:Path = `"$uvBin;`$env:Path`"" } + if (-not (Select-String -Path $PROFILE -SimpleMatch $cargoBin -Quiet)) { + Add-Content -Path $PROFILE -Value "`$env:Path = `"$cargoBin;`$env:Path`"" + } uv sync --frozen --group dev --python 3.11 - run: name: Run Windows-specific test @@ -232,6 +253,8 @@ jobs: environment: UV_HTTP_TIMEOUT: "300" command: | + $env:Path = "$HOME\.cargo\bin;$HOME\.local\bin;$env:Path" + cargo --version uv build --wheel --out-dir dist uv run --no-sync python tests/windows_tests/check_windows_wheel_install.py diff --git a/.dockerignore b/.dockerignore index a487d2a859a..6b80caeaf9f 100644 --- a/.dockerignore +++ b/.dockerignore @@ -49,6 +49,8 @@ build/ *.egg-info/ .DS_Store **/node_modules +litellm-rust/target/ +litellm/rust_bridge/_native*.so *.log .env .env.local diff --git a/Dockerfile b/Dockerfile index af49dc8d8cf..681681f28cc 100644 --- a/Dockerfile +++ b/Dockerfile @@ -21,6 +21,7 @@ RUN apk add --no-cache \ gcc \ python3 \ python3-dev \ + rust \ openssl \ openssl-dev \ nodejs \ diff --git a/docker/Dockerfile.non_root b/docker/Dockerfile.non_root index ab02b43d0f9..6bb925aa723 100644 --- a/docker/Dockerfile.non_root +++ b/docker/Dockerfile.non_root @@ -19,6 +19,7 @@ RUN for i in 1 2 3; do \ python3 \ python3-dev \ gcc \ + rust \ bash \ coreutils \ curl \ diff --git a/litellm-rust/crates/python-bridge/Cargo.toml b/litellm-rust/crates/python-bridge/Cargo.toml index f5b29f49cfd..4db32818604 100644 --- a/litellm-rust/crates/python-bridge/Cargo.toml +++ b/litellm-rust/crates/python-bridge/Cargo.toml @@ -6,7 +6,7 @@ license.workspace = true repository.workspace = true [lib] -name = "litellm_python_bridge" +name = "_native" crate-type = ["cdylib"] [dependencies] diff --git a/litellm-rust/crates/python-bridge/build.rs b/litellm-rust/crates/python-bridge/build.rs new file mode 100644 index 00000000000..0f7293007b2 --- /dev/null +++ b/litellm-rust/crates/python-bridge/build.rs @@ -0,0 +1,6 @@ +fn main() { + if std::env::var("CARGO_CFG_TARGET_OS").as_deref() == Ok("macos") { + println!("cargo:rustc-cdylib-link-arg=-undefined"); + println!("cargo:rustc-cdylib-link-arg=dynamic_lookup"); + } +} diff --git a/litellm-rust/crates/python-bridge/src/lib.rs b/litellm-rust/crates/python-bridge/src/lib.rs index 50ec7fceeac..8c8416b6bd3 100644 --- a/litellm-rust/crates/python-bridge/src/lib.rs +++ b/litellm-rust/crates/python-bridge/src/lib.rs @@ -93,7 +93,7 @@ fn gil_stats(py: Python<'_>) -> PyResult> { } #[pymodule] -fn litellm_python_bridge(module: &Bound<'_, PyModule>) -> PyResult<()> { +fn _native(module: &Bound<'_, PyModule>) -> PyResult<()> { module.add_function(wrap_pyfunction!(ocr, module)?)?; module.add_function(wrap_pyfunction!(gil_stats, module)?)?; Ok(()) diff --git a/litellm/ocr/rust_bridge.py b/litellm/ocr/rust_bridge.py index 61f9e8ca69a..4e688c42e1f 100644 --- a/litellm/ocr/rust_bridge.py +++ b/litellm/ocr/rust_bridge.py @@ -2,7 +2,7 @@ Optional Rust-backed OCR path. Enable with ``litellm.use_litellm_rust()``; the sync ``litellm.ocr()`` entrypoint -then routes supported Mistral calls through the compiled ``litellm_python_bridge`` +then routes supported Mistral calls through the compiled ``litellm.rust_bridge._native`` extension, which performs the whole OCR call (URL, headers, HTTP, parse) in Rust. No module-level ``litellm`` imports keep this a leaf so ``litellm/ocr/main.py`` @@ -15,7 +15,7 @@ from typing import Final, Protocol, cast class RustOcr(Protocol): - """Signature of the compiled ``litellm_python_bridge.ocr`` entrypoint.""" + """Signature of the compiled Rust OCR entrypoint.""" def __call__( self, @@ -41,7 +41,7 @@ _rust_ocr_impl: RustOcr | None = None def use_litellm_rust( enabled: bool = True, *, ocr: RustOcr | None | _Unset = _UNSET ) -> None: - """Route supported OCR calls through the Rust ``litellm_python_bridge`` extension. + """Route supported OCR calls through the packaged Rust extension. ``ocr`` injects the bridge callable; when omitted the compiled extension is loaded on demand and any previously injected bridge is preserved. Pass @@ -62,13 +62,14 @@ def load_rust_ocr() -> RustOcr | None: """Return the Rust OCR callable, or ``None`` when no bridge is available. Prefers an injected implementation, otherwise loads the compiled - ``litellm_python_bridge`` extension; a missing extension yields ``None`` so + ``litellm.rust_bridge._native`` extension; a missing extension yields ``None`` so the caller can fall back to the Python path instead of hard-failing. """ if _rust_ocr_impl is not None: return _rust_ocr_impl - try: - import litellm_python_bridge - except ImportError: + from litellm.rust_bridge import get_native_bridge + + native_bridge = get_native_bridge() + if native_bridge is None: return None - return cast(RustOcr, litellm_python_bridge.ocr) + return cast(RustOcr, native_bridge.ocr) diff --git a/litellm/rust_bridge/__init__.py b/litellm/rust_bridge/__init__.py new file mode 100644 index 00000000000..ec89e3b65b4 --- /dev/null +++ b/litellm/rust_bridge/__init__.py @@ -0,0 +1,8 @@ +"""LiteLLM Rust bridge package.""" + +from litellm.rust_bridge.loader import ( + get_native_bridge, + native_bridge_available, +) + +__all__ = ["get_native_bridge", "native_bridge_available"] diff --git a/litellm/rust_bridge/loader.py b/litellm/rust_bridge/loader.py new file mode 100644 index 00000000000..3ae0bf9307c --- /dev/null +++ b/litellm/rust_bridge/loader.py @@ -0,0 +1,28 @@ +"""Loader for the packaged LiteLLM Rust extension.""" + +from __future__ import annotations + +from types import ModuleType + +_BRIDGE_SENTINEL = object() +_cached_bridge: ModuleType | None | object = _BRIDGE_SENTINEL + + +def get_native_bridge() -> ModuleType | None: + """Return the packaged Rust extension, or ``None`` when unavailable.""" + global _cached_bridge + if _cached_bridge is not _BRIDGE_SENTINEL: + return _cached_bridge if isinstance(_cached_bridge, ModuleType) else None + + try: + from litellm.rust_bridge import _native + except ImportError: + _cached_bridge = None + return None + _cached_bridge = _native + return _native + + +def native_bridge_available() -> bool: + """Whether the packaged Rust extension is importable.""" + return get_native_bridge() is not None diff --git a/pyproject.toml b/pyproject.toml index ae612ae2184..30b9a84e09a 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -234,8 +234,24 @@ healthcheck = [ ] [build-system] -requires = ["uv_build==0.11.8"] -build-backend = "uv_build" +requires = ["maturin>=1.9.4,<2"] +build-backend = "maturin" + +[tool.maturin] +manifest-path = "litellm-rust/crates/python-bridge/Cargo.toml" +module-name = "litellm.rust_bridge._native" +python-source = "." +bindings = "pyo3" +exclude = [ + "litellm/proxy/enterprise", + "litellm/proxy/enterprise/**", + "**/__pycache__", + "**/__pycache__/**", + "**/.pytest_cache", + "**/.pytest_cache/**", + "**/.ruff_cache", + "**/.ruff_cache/**", +] [tool.uv] constraint-dependencies = [ @@ -253,18 +269,6 @@ litellm-enterprise = { workspace = true } [tool.uv.workspace] members = ["enterprise", "litellm-proxy-extras"] -[tool.uv.build-backend] -module-root = "" -source-exclude = [ - "litellm/proxy/enterprise", - "**/__pycache__", - "**/__pycache__/**", - "**/.pytest_cache", - "**/.pytest_cache/**", - "**/.ruff_cache", - "**/.ruff_cache/**", -] - [tool.isort] profile = "black" @@ -328,4 +332,3 @@ pytest_add_cli_args = [ [tool.coverage.run] source = ["litellm"] relative_files = true - diff --git a/tests/test_litellm/ocr/test_rust_bridge.py b/tests/test_litellm/ocr/test_rust_bridge.py index 7e028064e4c..11448aed828 100644 --- a/tests/test_litellm/ocr/test_rust_bridge.py +++ b/tests/test_litellm/ocr/test_rust_bridge.py @@ -1,7 +1,7 @@ """Tests for the optional Rust-backed OCR path (``litellm/ocr/rust_bridge.py``).""" import importlib -import sys +import builtins import types import httpx @@ -15,6 +15,7 @@ from litellm.llms.base_llm.ocr.transformation import OCRResponse # explicitly via importlib rather than attribute traversal. ocr_main = importlib.import_module("litellm.ocr.main") rust_bridge = importlib.import_module("litellm.ocr.rust_bridge") +rust_bridge_loader = importlib.import_module("litellm.rust_bridge.loader") MODEL = "mistral/mistral-ocr-latest" DOCUMENT = {"type": "document_url", "document_url": "https://example.com/doc.pdf"} @@ -80,8 +81,10 @@ class FakeOCRConfig: def _reset_rust_flag(): """Keep the global toggle isolated between tests.""" rust_bridge.use_litellm_rust(False, ocr=None) + rust_bridge_loader._cached_bridge = rust_bridge_loader._BRIDGE_SENTINEL yield rust_bridge.use_litellm_rust(False, ocr=None) + rust_bridge_loader._cached_bridge = rust_bridge_loader._BRIDGE_SENTINEL @pytest.fixture @@ -106,6 +109,44 @@ def test_load_rust_ocr_returns_injected_impl(): assert rust_bridge.load_rust_ocr() is bridge +def test_native_bridge_loader_returns_none_when_extension_absent(monkeypatch): + real_import = builtins.__import__ + + def fake_import(name, globals=None, locals=None, fromlist=(), level=0): + if name == "litellm.rust_bridge" and "_native" in fromlist: + raise ImportError + return real_import(name, globals, locals, fromlist, level) + + monkeypatch.setattr(builtins, "__import__", fake_import) + + assert rust_bridge_loader.get_native_bridge() is None + + +def test_native_bridge_loader_caches_absent_extension(monkeypatch): + real_import = builtins.__import__ + attempts = 0 + + def fake_import(name, globals=None, locals=None, fromlist=(), level=0): + nonlocal attempts + if name == "litellm.rust_bridge" and "_native" in fromlist: + attempts += 1 + raise ImportError + return real_import(name, globals, locals, fromlist, level) + + monkeypatch.setattr(builtins, "__import__", fake_import) + + assert rust_bridge_loader.get_native_bridge() is None + assert rust_bridge_loader.get_native_bridge() is None + assert attempts == 1 + + +def test_native_bridge_available_reflects_loader(monkeypatch): + fake_module = types.ModuleType("litellm.rust_bridge._native") + monkeypatch.setattr(rust_bridge_loader, "get_native_bridge", lambda: fake_module) + + assert rust_bridge_loader.native_bridge_available() is True + + def test_toggle_without_ocr_arg_preserves_injected_impl(): """Regression: routine enable/disable calls must not clobber a prior injection. @@ -122,7 +163,12 @@ def test_toggle_without_ocr_arg_preserves_injected_impl(): assert rust_bridge.load_rust_ocr() is bridge -def test_explicit_ocr_none_clears_injected_impl(): +def test_explicit_ocr_none_clears_injected_impl(monkeypatch): + monkeypatch.setattr( + importlib.import_module("litellm.rust_bridge"), + "get_native_bridge", + lambda: None, + ) bridge = RecordingBridge() litellm.use_litellm_rust(True, ocr=bridge) @@ -130,20 +176,29 @@ def test_explicit_ocr_none_clears_injected_impl(): assert rust_bridge.load_rust_ocr() is None -def test_load_rust_ocr_none_when_extension_absent(): +def test_load_rust_ocr_none_when_extension_absent(monkeypatch): """With no injected impl and no compiled wheel, the loader returns None so the caller degrades to the Python path instead of raising ImportError.""" + monkeypatch.setattr( + importlib.import_module("litellm.rust_bridge"), + "get_native_bridge", + lambda: None, + ) litellm.use_litellm_rust(True) # no impl injected; extension isn't built in CI assert rust_bridge.load_rust_ocr() is None def test_load_rust_ocr_uses_compiled_extension(monkeypatch): - """With no injected impl but a compiled ``litellm_python_bridge`` importable, + """With no injected impl but a packaged ``litellm.rust_bridge._native`` importable, the loader returns the extension's ``ocr`` callable. The native wheel isn't - built in CI, so stand in a fake module via ``sys.modules``.""" - fake_module = types.ModuleType("litellm_python_bridge") + built in CI, so stand in a fake module via the bridge loader.""" + fake_module = types.ModuleType("litellm.rust_bridge._native") fake_module.ocr = lambda **kwargs: dict(FAKE_OCR_RESPONSE) # type: ignore[attr-defined] - monkeypatch.setitem(sys.modules, "litellm_python_bridge", fake_module) + monkeypatch.setattr( + importlib.import_module("litellm.rust_bridge"), + "get_native_bridge", + lambda: fake_module, + ) litellm.use_litellm_rust(True) # enabled, no impl injected -> import the extension assert rust_bridge.load_rust_ocr() is fake_module.ocr @@ -317,6 +372,7 @@ def test_ocr_does_not_route_to_rust_when_disabled(): def test_ocr_falls_back_to_python_when_bridge_unavailable(monkeypatch): """Rust enabled but no bridge available (no injected impl, no compiled wheel): ocr() must degrade to the Python HTTP handler instead of raising.""" + monkeypatch.setattr(ocr_main, "load_rust_ocr", lambda: None) litellm.use_litellm_rust(True) # enabled, but load_rust_ocr() returns None in CI captured = {} diff --git a/uv.lock b/uv.lock index 8193bf14f63..d81d6e5d8a0 100644 --- a/uv.lock +++ b/uv.lock @@ -3121,15 +3121,15 @@ wheels = [ [[package]] name = "langgraph-checkpoint" -version = "4.1.0" +version = "4.1.1" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "langchain-core" }, { name = "ormsgpack" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/02/b4/6005c5dd88ad484fe6235d4c43a0d2cee7e91b08ad85a180985c2662df87/langgraph_checkpoint-4.1.0.tar.gz", hash = "sha256:e5bb304e30fc1363ac8fcb5f7dee5ca2185d77fe475b0d01de2c5f91324c2c21", size = 181942, upload-time = "2026-05-12T03:33:49.888Z" } +sdist = { url = "https://files.pythonhosted.org/packages/83/47/886af6f886f0bff2273164a45f008694e48a96ff3cd25ff0228f2aa9480e/langgraph_checkpoint-4.1.1.tar.gz", hash = "sha256:6c2bdb530c91f91d7d9c1bd100925d0fc4f498d418c17f3587d1526279482a25", size = 184020, upload-time = "2026-05-22T16:57:38.503Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/93/74/d3be2b41955e20ccd624dba5f6fe9d38dcee385ba470a6e13ed86732fc86/langgraph_checkpoint-4.1.0-py3-none-any.whl", hash = "sha256:8bc2a0466a20c38b865ce6671b42093fd5c041133f32351cae4222e0eeaf7fb5", size = 56047, upload-time = "2026-05-12T03:33:48.548Z" }, + { url = "https://files.pythonhosted.org/packages/bd/b4/71425e3e38be92611300b9cc5e46a5bf98ab23f5ea8a75b73d02a2f1413c/langgraph_checkpoint-4.1.1-py3-none-any.whl", hash = "sha256:25d29144b082827218e7bc3f1e9b0566a4bb007895cd6cc26f66a8428739f56e", size = 56212, upload-time = "2026-05-22T16:57:37.203Z" }, ] [[package]] From f98e9355040cbc8e188a73698c613a99b5529231 Mon Sep 17 00:00:00 2001 From: Mateo Wang <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 25 Jun 2026 14:28:49 -0700 Subject: [PATCH 3/8] chore: gitignore rust bridge build artifacts (#31349) Ignore the compiled, platform-specific Rust extension output (litellm/rust_bridge/_native*.so/.pyd) and the litellm-rust/target/ build dir so local maturin/cargo builds don't show up as untracked files. Also drop the two stale self-referential .gitignore entries; .gitignore is tracked, so ignoring it did nothing except add confusion. Co-authored-by: Cursor Agent Co-authored-by: Mateo Wang --- .gitignore | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/.gitignore b/.gitignore index 3563d7c8c2d..59fa5803abe 100644 --- a/.gitignore +++ b/.gitignore @@ -9,12 +9,17 @@ litellm/proxy/myenv/* litellm_uuid.txt __pycache__/ *.pyc + +# Rust bridge build artifacts (compiled, platform-specific; regenerated by maturin/cargo) +litellm/rust_bridge/_native*.so +litellm/rust_bridge/_native*.pyd +litellm-rust/target/ + bun.lockb **/.DS_Store .aider* litellm_results.jsonl secrets.toml -.gitignore litellm/proxy/litellm_secrets.toml litellm/proxy/api_log.json .idea/ @@ -36,7 +41,6 @@ litellm/tests/dynamo*.log .vscode/settings.json litellm/proxy/log.txt proxy_server_config_@.yaml -.gitignore proxy_server_config_2.yaml litellm/proxy/secret_managers/credentials.json hosted_config.yaml From a2d04ccdbb274e9a124068ef58de0ec76b7ca08e Mon Sep 17 00:00:00 2001 From: ishaan-berri <155045088+ishaan-berri@users.noreply.github.com> Date: Thu, 25 Jun 2026 14:31:05 -0700 Subject: [PATCH 4/8] ci: harden cargo fetches during maturin builds (#31348) --- {litellm-rust/.cargo => .cargo}/config.toml | 9 ++++++ .circleci/config.yml | 15 ++++++++- .github/scripts/uv_sync_with_retries.sh | 31 +++++++++++++++++++ .github/workflows/_test-unit-base.yml | 2 +- .github/workflows/check-ui-api-types.yml | 2 +- .github/workflows/mutation-test.yml | 2 +- .github/workflows/test-mcp.yml | 2 +- .github/workflows/test-unit-documentation.yml | 2 +- .github/workflows/test-unit-proxy-legacy.yml | 2 +- pyproject.toml | 2 +- 10 files changed, 61 insertions(+), 8 deletions(-) rename {litellm-rust/.cargo => .cargo}/config.toml (66%) create mode 100755 .github/scripts/uv_sync_with_retries.sh diff --git a/litellm-rust/.cargo/config.toml b/.cargo/config.toml similarity index 66% rename from litellm-rust/.cargo/config.toml rename to .cargo/config.toml index c7fb2592542..e6afd7ff530 100644 --- a/litellm-rust/.cargo/config.toml +++ b/.cargo/config.toml @@ -1,3 +1,12 @@ +[http] +# CI has seen transient crates.io failures from libcurl's HTTP/2 multiplexing +# during `maturin` metadata resolution. Disable multiplexing and retry more +# aggressively so editable `uv sync` builds are not failed by one flaky frame. +multiplexing = false + +[net] +retry = 5 + # PyO3 cdylib (`litellm-python-bridge`) links against the host interpreter's # symbols, which are not present at link time when building an extension module. # On macOS, tell the linker to resolve undefined `_Py*` symbols dynamically at diff --git a/.circleci/config.yml b/.circleci/config.yml index d532d3106e5..337f4b5b3f3 100644 --- a/.circleci/config.yml +++ b/.circleci/config.yml @@ -190,6 +190,8 @@ jobs: working_directory: ~/project environment: UV_PYTHON: "3.11" + CARGO_HTTP_MULTIPLEXING: "false" + CARGO_NET_RETRY: "5" steps: - checkout - run: @@ -243,7 +245,17 @@ jobs: if (-not (Select-String -Path $PROFILE -SimpleMatch $cargoBin -Quiet)) { Add-Content -Path $PROFILE -Value "`$env:Path = `"$cargoBin;`$env:Path`"" } - uv sync --frozen --group dev --python 3.11 + for ($attempt = 1; $attempt -le 5; $attempt++) { + Write-Host "uv sync attempt $attempt/5" + uv sync --frozen --group dev --python 3.11 + if ($LASTEXITCODE -eq 0) { + break + } + if ($attempt -eq 5) { + exit $LASTEXITCODE + } + Start-Sleep -Seconds 15 + } - run: name: Run Windows-specific test command: | @@ -255,6 +267,7 @@ jobs: command: | $env:Path = "$HOME\.cargo\bin;$HOME\.local\bin;$env:Path" cargo --version + Get-ChildItem -Path "litellm\rust_bridge" -Filter "_native*" -File -ErrorAction SilentlyContinue | Remove-Item -Force uv build --wheel --out-dir dist uv run --no-sync python tests/windows_tests/check_windows_wheel_install.py diff --git a/.github/scripts/uv_sync_with_retries.sh b/.github/scripts/uv_sync_with_retries.sh new file mode 100755 index 00000000000..85ed75af566 --- /dev/null +++ b/.github/scripts/uv_sync_with_retries.sh @@ -0,0 +1,31 @@ +#!/usr/bin/env bash +set -euo pipefail + +max_attempts="${UV_SYNC_MAX_ATTEMPTS:-5}" +delay_seconds="${UV_SYNC_RETRY_DELAY_SECONDS:-15}" + +export CARGO_HTTP_MULTIPLEXING="${CARGO_HTTP_MULTIPLEXING:-false}" +export CARGO_NET_RETRY="${CARGO_NET_RETRY:-5}" + +if [[ "$#" -eq 0 ]]; then + echo "usage: $0 " >&2 + exit 2 +fi + +for attempt in $(seq 1 "${max_attempts}"); do + echo "uv sync attempt ${attempt}/${max_attempts}" + status=0 + if uv sync "$@"; then + exit 0 + else + status=$? + fi + + if [[ "${attempt}" -eq "${max_attempts}" ]]; then + echo "uv sync failed after ${max_attempts} attempts" >&2 + exit "${status}" + fi + + echo "uv sync failed; retrying in ${delay_seconds}s..." + sleep "${delay_seconds}" +done diff --git a/.github/workflows/_test-unit-base.yml b/.github/workflows/_test-unit-base.yml index a42b2f8f9df..25c6d4a7019 100644 --- a/.github/workflows/_test-unit-base.yml +++ b/.github/workflows/_test-unit-base.yml @@ -73,7 +73,7 @@ jobs: - name: Install dependencies run: | - uv sync --frozen --group ci --group proxy-dev --extra google --extra proxy --extra semantic-router + .github/scripts/uv_sync_with_retries.sh --frozen --group ci --group proxy-dev --extra google --extra proxy --extra semantic-router - name: Generate Prisma client env: diff --git a/.github/workflows/check-ui-api-types.yml b/.github/workflows/check-ui-api-types.yml index d8053c15683..439126aa1ee 100644 --- a/.github/workflows/check-ui-api-types.yml +++ b/.github/workflows/check-ui-api-types.yml @@ -46,7 +46,7 @@ jobs: ${{ runner.os }}-uv- - name: Install backend dependencies - run: uv sync --frozen --group ci --group proxy-dev --extra google --extra proxy --extra semantic-router + run: .github/scripts/uv_sync_with_retries.sh --frozen --group ci --group proxy-dev --extra google --extra proxy --extra semantic-router - name: Generate Prisma client env: diff --git a/.github/workflows/mutation-test.yml b/.github/workflows/mutation-test.yml index 8094ca57467..183f12f969c 100644 --- a/.github/workflows/mutation-test.yml +++ b/.github/workflows/mutation-test.yml @@ -55,7 +55,7 @@ jobs: - name: Install dependencies run: | - uv sync --frozen --group ci --group proxy-dev --extra google --extra proxy --extra semantic-router + .github/scripts/uv_sync_with_retries.sh --frozen --group ci --group proxy-dev --extra google --extra proxy --extra semantic-router - name: Generate Prisma client env: diff --git a/.github/workflows/test-mcp.yml b/.github/workflows/test-mcp.yml index cefb77980ac..5b5290880c1 100644 --- a/.github/workflows/test-mcp.yml +++ b/.github/workflows/test-mcp.yml @@ -39,7 +39,7 @@ jobs: - name: Install dependencies run: | uv lock --check - uv sync --frozen --group proxy-dev --extra proxy --extra semantic-router + .github/scripts/uv_sync_with_retries.sh --frozen --group proxy-dev --extra proxy --extra semantic-router - name: Run MCP tests run: | diff --git a/.github/workflows/test-unit-documentation.yml b/.github/workflows/test-unit-documentation.yml index 8ad9fb6a73b..4cef791a9b3 100644 --- a/.github/workflows/test-unit-documentation.yml +++ b/.github/workflows/test-unit-documentation.yml @@ -54,7 +54,7 @@ jobs: - name: Install dependencies run: | - uv sync --frozen --group ci --group proxy-dev --extra google --extra proxy --extra semantic-router + .github/scripts/uv_sync_with_retries.sh --frozen --group ci --group proxy-dev --extra google --extra proxy --extra semantic-router - name: Generate Prisma client env: diff --git a/.github/workflows/test-unit-proxy-legacy.yml b/.github/workflows/test-unit-proxy-legacy.yml index 922e3c2eddf..8db218cd1fc 100644 --- a/.github/workflows/test-unit-proxy-legacy.yml +++ b/.github/workflows/test-unit-proxy-legacy.yml @@ -71,7 +71,7 @@ jobs: - name: Install dependencies run: | - uv sync --frozen --group ci --group proxy-dev --extra google --extra proxy --extra semantic-router + .github/scripts/uv_sync_with_retries.sh --frozen --group ci --group proxy-dev --extra google --extra proxy --extra semantic-router - name: Generate Prisma client env: diff --git a/pyproject.toml b/pyproject.toml index 30b9a84e09a..6a6a47f540e 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -234,7 +234,7 @@ healthcheck = [ ] [build-system] -requires = ["maturin>=1.9.4,<2"] +requires = ["maturin==1.9.4"] build-backend = "maturin" [tool.maturin] From 92d0788da21c78d5ceec57072b9fc2621cb39634 Mon Sep 17 00:00:00 2001 From: Mateo Wang <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 25 Jun 2026 14:43:45 -0700 Subject: [PATCH 5/8] chore(lint): widen ANN slack to 10% of baseline and drop PLR0913 from the strict gate (#31335) * chore(lint): widen ruff budget slack to 10% of baseline for high-volume ANN rules and PLR0913 Co-authored-by: Mateo Wang * chore(lint): drop PLR0913 from strict gate to roll out rules gradually Co-authored-by: Mateo Wang * fix(lint): ratchet-guard rising baselines even when slack is cut to mask them Co-authored-by: Mateo Wang --------- Co-authored-by: Cursor Agent Co-authored-by: Mateo Wang --- ruff-strict-budget.json | 16 ++- ruff-strict.toml | 2 +- scripts/budget_ratchet_check.py | 97 ++++++++++++++----- .../test_litellm/test_budget_ratchet_check.py | 33 ++++++- 4 files changed, 109 insertions(+), 39 deletions(-) diff --git a/ruff-strict-budget.json b/ruff-strict-budget.json index ae46f020de1..10c820324ea 100644 --- a/ruff-strict-budget.json +++ b/ruff-strict-budget.json @@ -1,7 +1,7 @@ { "ANN001": { "baseline": 2865, - "slack": 50 + "slack": 287 }, "ANN002": { "baseline": 64, @@ -9,19 +9,19 @@ }, "ANN003": { "baseline": 759, - "slack": 30 + "slack": 76 }, "ANN201": { "baseline": 1944, - "slack": 50 + "slack": 194 }, "ANN202": { "baseline": 858, - "slack": 30 + "slack": 86 }, "ANN204": { "baseline": 658, - "slack": 20 + "slack": 66 }, "ANN205": { "baseline": 117, @@ -33,7 +33,7 @@ }, "ANN401": { "baseline": 1886, - "slack": 50 + "slack": 189 }, "ASYNC230": { "baseline": 11, @@ -231,10 +231,6 @@ "baseline": 6, "slack": 3 }, - "PLR0913": { - "baseline": 1813, - "slack": 50 - }, "PLR1704": { "baseline": 3, "slack": 3 diff --git a/ruff-strict.toml b/ruff-strict.toml index 1caa3567872..d58885fe848 100644 --- a/ruff-strict.toml +++ b/ruff-strict.toml @@ -2,7 +2,7 @@ extend = "ruff.toml" [lint] preview = true -select = ["ANN", "ASYNC230", "B004", "B006", "B008", "B009", "B010", "B018", "B019", "B021", "B026", "B033", "BLE", "C401", "C404", "C405", "C408", "C414", "C419", "C901", "D419", "DTZ001", "DTZ003", "DTZ005", "DTZ006", "DTZ007", "DTZ011", "EXE001", "EXE002", "F401", "FURB136", "FURB168", "FURB188", "I001", "LOG015", "N999", "PERF102", "PERF401", "PERF402", "PERF403", "PIE790", "PIE800", "PIE804", "PIE810", "PLC0206", "PLC0208", "PLC0414", "PLR0124", "PLR0206", "PLR0402", "PLR0913", "PLR1704", "PLR1711", "PLR1714", "PLR1730", "PLR2044", "PLW0127", "PLW0133", "PLW0602", "PLW0603", "PLW1508", "PLW1510", "PYI030", "PYI036", "PYI041", "PYI064", "RET501", "RET504", "RUF010", "RUF012", "RUF015", "RUF019", "RUF022", "RUF023", "RUF046", "RUF051", "RUF059", "RUF100", "S110", "S112", "SIM101", "SIM102", "SIM103", "SIM113", "SIM114", "SIM115", "SIM117", "SIM118", "SIM201", "SIM210", "SIM211", "SIM222", "SIM401", "TC004", "TC005", "TID251", "TRY002", "TRY004", "TRY201", "TRY203", "TRY300", "UP006", "UP007", "UP008", "UP012", "UP018", "UP024", "UP028", "UP031", "UP032", "UP034", "UP035", "UP036", "UP037", "UP045"] +select = ["ANN", "ASYNC230", "B004", "B006", "B008", "B009", "B010", "B018", "B019", "B021", "B026", "B033", "BLE", "C401", "C404", "C405", "C408", "C414", "C419", "C901", "D419", "DTZ001", "DTZ003", "DTZ005", "DTZ006", "DTZ007", "DTZ011", "EXE001", "EXE002", "F401", "FURB136", "FURB168", "FURB188", "I001", "LOG015", "N999", "PERF102", "PERF401", "PERF402", "PERF403", "PIE790", "PIE800", "PIE804", "PIE810", "PLC0206", "PLC0208", "PLC0414", "PLR0124", "PLR0206", "PLR0402", "PLR1704", "PLR1711", "PLR1714", "PLR1730", "PLR2044", "PLW0127", "PLW0133", "PLW0602", "PLW0603", "PLW1508", "PLW1510", "PYI030", "PYI036", "PYI041", "PYI064", "RET501", "RET504", "RUF010", "RUF012", "RUF015", "RUF019", "RUF022", "RUF023", "RUF046", "RUF051", "RUF059", "RUF100", "S110", "S112", "SIM101", "SIM102", "SIM103", "SIM113", "SIM114", "SIM115", "SIM117", "SIM118", "SIM201", "SIM210", "SIM211", "SIM222", "SIM401", "TC004", "TC005", "TID251", "TRY002", "TRY004", "TRY201", "TRY203", "TRY300", "UP006", "UP007", "UP008", "UP012", "UP018", "UP024", "UP028", "UP031", "UP032", "UP034", "UP035", "UP036", "UP037", "UP045"] extend-select = [] [lint.mccabe] diff --git a/scripts/budget_ratchet_check.py b/scripts/budget_ratchet_check.py index 861d65489e8..df9815d6557 100644 --- a/scripts/budget_ratchet_check.py +++ b/scripts/budget_ratchet_check.py @@ -1,16 +1,19 @@ #!/usr/bin/env python3 -"""Non-gating ratchet guard: budget ceilings may only fall, never rise. +"""Non-gating ratchet guard: budget baselines and ceilings may only fall, never rise. Every `*-budget.json` file (ruff-strict, type-discipline, basedpyright-code) is a -one-way ratchet: each rule's ceiling is `baseline + slack`, and the whole point is -to drive that number DOWN over time. This check compares every budget file against -its own content at the merge-base with the target branch and fails (exits 1, red) if: +one-way ratchet: each rule's ceiling is `baseline + slack`, and both the recorded +`baseline` (the live violation count) and that ceiling are meant to be driven DOWN +over time. This check compares every budget file against its own content at the +merge-base with the target branch and fails (exits 1, red) if: - * a rule's ceiling went up, + * a rule's ceiling (`baseline + slack`) went up, + * a rule's `baseline` went up, even if `slack` was lowered to keep the ceiling + flat (a higher baseline bakes in more accepted debt and must be acknowledged), * a rule was dropped from a budget (its ceiling effectively became infinite), or * an entire budget file was deleted. -New rules and lowered/equal ceilings are fine. +New rules and lowered/equal baselines and ceilings are fine. This is deliberately NOT a gating check. It should turn the run red so that a loosening is impossible to miss in review, but it must stay OUT of the @@ -66,7 +69,12 @@ def _load_head(rel: str) -> dict | None: def _ref_is_commit(ref: str) -> bool: - return _run(["git", "rev-parse", "--verify", "--quiet", f"{ref}^{{commit}}"]).returncode == 0 + return ( + _run( + ["git", "rev-parse", "--verify", "--quiet", f"{ref}^{{commit}}"] + ).returncode + == 0 + ) def _load_base(rel: str, ref: str) -> dict | None: @@ -81,13 +89,55 @@ def _load_base(rel: str, ref: str) -> dict | None: return json.loads(proc.stdout) +def _baselines(budget: dict) -> dict[str, int]: + """Map each rule to its recorded baseline; skip malformed specs.""" + return { + rule: int(spec.get("baseline", 0)) + for rule, spec in budget.items() + if isinstance(spec, dict) + } + + def _caps(budget: dict) -> dict[str, int]: """Map each rule to its ceiling (baseline + slack); skip malformed specs.""" - caps: dict[str, int] = {} - for rule, spec in budget.items(): - if isinstance(spec, dict): - caps[rule] = int(spec.get("baseline", 0)) + int(spec.get("slack", 0)) - return caps + return { + rule: int(spec.get("baseline", 0)) + int(spec.get("slack", 0)) + for rule, spec in budget.items() + if isinstance(spec, dict) + } + + +def _regression_detail( + rule: str, + base_caps: dict[str, int], + head_caps: dict[str, int], + base_baselines: dict[str, int], + head_baselines: dict[str, int], +) -> str | None: + """Why `rule` regressed vs base, or None when it held flat or fell. + + A dropped rule is terminal; otherwise a raised ceiling and a raised baseline are + independent loosenings (the latter catches a baseline bump masked by a slack cut), + so both reasons are reported when both apply. + """ + base_cap = base_caps[rule] + if rule not in head_caps: + return f"rule dropped (ceiling {base_cap} -> removed)" + reasons = tuple( + message + for raised, message in ( + ( + head_caps[rule] > base_cap, + f"ceiling raised {base_cap} -> {head_caps[rule]}", + ), + ( + head_baselines[rule] > base_baselines[rule], + f"baseline raised {base_baselines[rule]} -> {head_baselines[rule]}", + ), + ) + if raised + ) + return "; ".join(reasons) or None def regressions_for(rel: str, base: dict | None, head: dict | None) -> list[Regression]: @@ -96,18 +146,17 @@ def regressions_for(rel: str, base: dict | None, head: dict | None) -> list[Regr if head is None: return [Regression(rel, "*", "budget file was deleted (every ceiling removed)")] - base_caps = _caps(base) - head_caps = _caps(head) + base_caps, head_caps = _caps(base), _caps(head) + base_baselines, head_baselines = _baselines(base), _baselines(head) return [ - Regression( - rel, - rule, - f"rule dropped (ceiling {base_cap} -> removed)" - if rule not in head_caps - else f"ceiling raised {base_cap} -> {head_caps[rule]}", + Regression(rel, rule, detail) + for rule in sorted(base_caps) + if ( + detail := _regression_detail( + rule, base_caps, head_caps, base_baselines, head_baselines + ) ) - for rule, base_cap in sorted(base_caps.items()) - if rule not in head_caps or head_caps[rule] > base_cap + is not None ] @@ -141,7 +190,9 @@ def main() -> int: regressions.extend(regressions_for(rel, base, head)) if regressions: - print(f"FAIL: budget ceiling(s) loosened vs base {args.base} (merge-base {ref[:12]}):") + print( + f"FAIL: budget baseline(s)/ceiling(s) loosened vs base {args.base} (merge-base {ref[:12]}):" + ) for reg in regressions: print(f" {reg.budget} {reg.rule}: {reg.detail}") print( diff --git a/tests/test_litellm/test_budget_ratchet_check.py b/tests/test_litellm/test_budget_ratchet_check.py index 9f19944fdba..77cee8a485c 100644 --- a/tests/test_litellm/test_budget_ratchet_check.py +++ b/tests/test_litellm/test_budget_ratchet_check.py @@ -1,7 +1,8 @@ """Tests for scripts/budget_ratchet_check.py. -The guard's whole contract is "ceilings may only fall": a raised ceiling, a dropped -rule, or a deleted file is a regression, while a lowered/equal ceiling, a brand-new +The guard's contract is "baselines and ceilings may only fall": a raised ceiling, a +raised baseline (even when slack is cut to keep the ceiling flat), a dropped rule, or +a deleted file is a regression, while a lowered/equal baseline and ceiling, a brand-new rule, or a brand-new budget file is fine. Each branch is pinned here. """ @@ -10,7 +11,9 @@ import subprocess import sys from pathlib import Path -_MODULE_PATH = Path(__file__).resolve().parents[2] / "scripts" / "budget_ratchet_check.py" +_MODULE_PATH = ( + Path(__file__).resolve().parents[2] / "scripts" / "budget_ratchet_check.py" +) _spec = importlib.util.spec_from_file_location("budget_ratchet_check", _MODULE_PATH) ratchet = importlib.util.module_from_spec(_spec) _spec.loader.exec_module(ratchet) @@ -35,10 +38,30 @@ def test_raised_ceiling_is_a_regression(): def test_lowered_or_equal_ceiling_is_clean(): base = {"LIT006": _spec_of(1013, 10)} + # baseline drops, slack flat -> ceiling falls assert ratchet.regressions_for("b.json", base, {"LIT006": _spec_of(1000, 10)}) == [] + # nothing changes assert ratchet.regressions_for("b.json", base, {"LIT006": _spec_of(1013, 10)}) == [] - # slack traded for baseline at the same ceiling is fine - assert ratchet.regressions_for("b.json", base, {"LIT006": _spec_of(1023, 0)}) == [] + # slack cut while baseline holds -> ceiling falls, baseline flat + assert ratchet.regressions_for("b.json", base, {"LIT006": _spec_of(1013, 0)}) == [] + + +def test_raised_baseline_is_a_regression_even_when_ceiling_held_flat(): + # baseline 1013 -> 1023 with slack cut 10 -> 0 keeps the ceiling at 1023, but a + # higher baseline bakes in more accepted debt and must still surface as a regression + base = {"LIT006": _spec_of(1013, 10)} + regs = ratchet.regressions_for("b.json", base, {"LIT006": _spec_of(1023, 0)}) + assert [r.rule for r in regs] == ["LIT006"] + assert "baseline raised 1013 -> 1023" in regs[0].detail + assert "ceiling raised" not in regs[0].detail + + +def test_raised_baseline_and_ceiling_report_both_reasons(): + base = {"LIT006": _spec_of(1013, 10)} + regs = ratchet.regressions_for("b.json", base, {"LIT006": _spec_of(1100, 10)}) + assert [r.rule for r in regs] == ["LIT006"] + assert "ceiling raised 1023 -> 1110" in regs[0].detail + assert "baseline raised 1013 -> 1100" in regs[0].detail def test_dropped_rule_is_a_regression(): From 62f93a33434fdec2758e7427984e73608a5546a1 Mon Sep 17 00:00:00 2001 From: ishaan-berri <155045088+ishaan-berri@users.noreply.github.com> Date: Thu, 25 Jun 2026 15:12:30 -0700 Subject: [PATCH 6/8] feat: add Rust OCR providers (#31272) * feat: port OCR providers to Rust gateway * chore(deps): update langgraph checkpoint lock * ci: scope ruff format check to changed files * ci: fix OCR lint and patch coverage * fix(ocr): block mapped IPv6 fetch targets * test(ocr): include rust bridge coverage in OCR shard * ci: rerun responses shard --- .circleci/config.yml | 4 +- .github/workflows/test-linting.yml | 11 +- codecov.yaml | 3 + litellm-rust/Cargo.lock | 162 +++- litellm-rust/Cargo.toml | 6 +- litellm-rust/crates/ai-gateway/Cargo.toml | 1 + litellm-rust/crates/ai-gateway/src/io/ocr.rs | 359 +++++++- .../ai-gateway/src/io/ocr/common_utils.rs | 448 ++++++++++ litellm-rust/crates/core/src/error.rs | 4 + .../crates/core/src/ocr/transformation.rs | 49 +- .../crates/core/src/providers/azure_ai/mod.rs | 1 + .../core/src/providers/azure_ai/ocr/mod.rs | 1 + .../providers/azure_ai/ocr/transformation.rs | 520 +++++++++++ .../providers/mistral/ocr/transformation.rs | 18 + litellm-rust/crates/core/src/providers/mod.rs | 2 + .../core/src/providers/vertex_ai/mod.rs | 1 + .../core/src/providers/vertex_ai/ocr/mod.rs | 1 + .../providers/vertex_ai/ocr/transformation.rs | 435 ++++++++++ litellm-rust/crates/python-bridge/Cargo.toml | 2 + litellm-rust/crates/python-bridge/src/lib.rs | 149 +++- .../document_intelligence/transformation.py | 19 +- litellm/llms/azure_ai/ocr/transformation.py | 19 +- litellm/llms/base_llm/ocr/transformation.py | 46 +- litellm/llms/mistral/ocr/transformation.py | 19 +- .../vertex_ai/ocr/deepseek_transformation.py | 72 +- litellm/llms/vertex_ai/ocr/transformation.py | 24 +- litellm/ocr/main.py | 806 +++++++++++------- litellm/ocr/rust_bridge.py | 65 +- tests/documentation_tests/test_env_keys.py | 7 + tests/e2e/gateway/litellm-config.yml | 30 +- tests/e2e/gateway/test_ocr_rust_e2e.py | 148 ++++ tests/test_litellm/ocr/test_rust_bridge.py | 480 +++++++++-- uv.lock | 2 +- 33 files changed, 3361 insertions(+), 553 deletions(-) create mode 100644 litellm-rust/crates/ai-gateway/src/io/ocr/common_utils.rs create mode 100644 litellm-rust/crates/core/src/providers/azure_ai/mod.rs create mode 100644 litellm-rust/crates/core/src/providers/azure_ai/ocr/mod.rs create mode 100644 litellm-rust/crates/core/src/providers/azure_ai/ocr/transformation.rs create mode 100644 litellm-rust/crates/core/src/providers/vertex_ai/mod.rs create mode 100644 litellm-rust/crates/core/src/providers/vertex_ai/ocr/mod.rs create mode 100644 litellm-rust/crates/core/src/providers/vertex_ai/ocr/transformation.rs create mode 100644 tests/e2e/gateway/test_ocr_rust_e2e.py diff --git a/.circleci/config.yml b/.circleci/config.yml index 337f4b5b3f3..f13e9bf66f1 100644 --- a/.circleci/config.yml +++ b/.circleci/config.yml @@ -1056,7 +1056,9 @@ jobs: name: Run tests command: | mkdir -p test-results - TEST_FILES=$(circleci tests glob "tests/ocr_tests/**/test_*.py") + TEST_FILES=$(printf "%s\n%s\n" \ + "$(circleci tests glob "tests/ocr_tests/**/test_*.py")" \ + "tests/test_litellm/ocr/test_rust_bridge.py") echo "$TEST_FILES" | circleci tests run \ --verbose \ --command="tr ' ' '\\n' | awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest \ diff --git a/.github/workflows/test-linting.yml b/.github/workflows/test-linting.yml index f7c53ff443a..ff6c40ac9ae 100644 --- a/.github/workflows/test-linting.yml +++ b/.github/workflows/test-linting.yml @@ -51,10 +51,15 @@ jobs: uv sync --frozen - name: Check ruff format + env: + BASE_SHA: ${{ github.event.pull_request.base.sha }} run: | - cd litellm - uv run --no-sync ruff format --check --line-length 88 --exclude '/enterprise/' . - cd .. + git diff --name-only "$BASE_SHA"...HEAD -- 'litellm/**/*.py' | grep -v '^litellm/enterprise/' > "$RUNNER_TEMP/ruff_format_files.txt" || true + if [ ! -s "$RUNNER_TEMP/ruff_format_files.txt" ]; then + echo "No changed litellm Python files to check with ruff format." + exit 0 + fi + xargs uv run --no-sync ruff format --check --line-length 88 --exclude '/enterprise/' < "$RUNNER_TEMP/ruff_format_files.txt" - name: Debug - Check file state run: | diff --git a/codecov.yaml b/codecov.yaml index 3baea13e2d3..f5acdd39136 100644 --- a/codecov.yaml +++ b/codecov.yaml @@ -3,6 +3,9 @@ codecov: notify: wait_for_ci: false # post as soon as expected uploads arrive, don't wait on CI +ignore: + - "litellm-rust/**" + # Uploads are flagged per workflow/shard (GHA) or "circleci". carryforward makes # a re-upload of a flag replace its prior session instead of accumulating a # conflicting one, and lets a commit reuse a flag from its parent when that flag diff --git a/litellm-rust/Cargo.lock b/litellm-rust/Cargo.lock index 6fe84f1cfbc..ce86a0ee6ac 100644 --- a/litellm-rust/Cargo.lock +++ b/litellm-rust/Cargo.lock @@ -206,12 +206,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" @@ -221,6 +233,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" @@ -237,12 +264,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" @@ -261,8 +310,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", @@ -307,6 +358,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" @@ -368,6 +444,7 @@ dependencies = [ "bytes", "futures-channel", "futures-core", + "h2", "http", "http-body", "httparse", @@ -521,6 +598,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" @@ -544,9 +631,9 @@ checksum = "8f42a60cbdf9a97f5d2305f08a87dc4e09308d1276d28c869c684d7777685682" [[package]] name = "js-sys" -version = "0.3.102" +version = "0.3.103" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "03d04c30968dffe80775bd4d7fb676131cd04a1fb46d2686dbffbaec2d9dfd31" +checksum = "53b44bfcdb3f8d5837a46dae1ca9660a837176eee74a28b229bc626816589102" dependencies = [ "cfg-if", "futures-util", @@ -564,6 +651,7 @@ name = "litellm-ai-gateway" version = "0.1.0" dependencies = [ "axum", + "base64", "futures-channel", "futures-util", "litellm-core", @@ -594,7 +682,9 @@ dependencies = [ "litellm-ai-gateway", "litellm-core", "pyo3", + "pyo3-async-runtimes", "serde_json", + "tokio", ] [[package]] @@ -728,6 +818,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" @@ -913,6 +1016,7 @@ dependencies = [ "futures-channel", "futures-core", "futures-util", + "h2", "http", "http-body", "http-body-util", @@ -932,12 +1036,14 @@ dependencies = [ "sync_wrapper", "tokio", "tokio-rustls", + "tokio-util", "tower", "tower-http", "tower-service", "url", "wasm-bindgen", "wasm-bindgen-futures", + "wasm-streams", "web-sys", "webpki-roots", ] @@ -1335,6 +1441,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" @@ -1507,9 +1626,9 @@ dependencies = [ [[package]] name = "wasm-bindgen" -version = "0.2.125" +version = "0.2.126" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "8ddb3f79143bced6de84270411622a2699cee572fc0875aeaf1e7867cf9fca1a" +checksum = "4b067c0c11094aef6b7a801c1e34a26affafdf3d051dba08456b868789aaf9a4" dependencies = [ "cfg-if", "once_cell", @@ -1520,9 +1639,9 @@ dependencies = [ [[package]] name = "wasm-bindgen-futures" -version = "0.4.75" +version = "0.4.76" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "503b14d284f2c8dac03b819967e155ea753f573586193b2b2c95990cb5d69280" +checksum = "c62df1340f32221cb9c54d6a27b030e3dba64361d4a95bed55f9aacb44da291d" dependencies = [ "js-sys", "wasm-bindgen", @@ -1530,9 +1649,9 @@ dependencies = [ [[package]] name = "wasm-bindgen-macro" -version = "0.2.125" +version = "0.2.126" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "4e21a184b13fb19e157296e2c46056aec9092264fab83e4ba59e68c61b323c3d" +checksum = "167ce5e579f6bcf889c4f7175a8a5a585de84e8ff93976ce393efa5f2837aab1" dependencies = [ "quote", "wasm-bindgen-macro-support", @@ -1540,9 +1659,9 @@ dependencies = [ [[package]] name = "wasm-bindgen-macro-support" -version = "0.2.125" +version = "0.2.126" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "fecefd9c35bd935a20fc3fc344b5f29138961e4f47fb03297d88f2587afb5ebd" +checksum = "f3997c7839262f4ef12cf90b818d6340c18e80f263f1a94bf157d0ec4420380e" dependencies = [ "bumpalo", "proc-macro2", @@ -1553,18 +1672,31 @@ dependencies = [ [[package]] name = "wasm-bindgen-shared" -version = "0.2.125" +version = "0.2.126" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "23939e44bb9a5d7576fa2b563dc2e136628f1224e88a8deed09e04858b77871f" +checksum = "dc1b4cb0cc549fcf58d7dfc081778139b3d283a081644e833e84682ad71cea24" dependencies = [ "unicode-ident", ] [[package]] -name = "web-sys" -version = "0.3.102" +name = "wasm-streams" +version = "0.4.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a6430a72df5eb332242960fe84b3002a241163998241eb596d4f739b9757061d" +checksum = "15053d8d85c7eccdbefef60f06769760a563c7f0a9d6902a13d35c7800b0ad65" +dependencies = [ + "futures-util", + "js-sys", + "wasm-bindgen", + "wasm-bindgen-futures", + "web-sys", +] + +[[package]] +name = "web-sys" +version = "0.3.103" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8622dcb61c0bcc9fffa6938bed81210af2da9a7e4a1a834b2e37a59b6dfb6141" dependencies = [ "js-sys", "wasm-bindgen", diff --git a/litellm-rust/Cargo.toml b/litellm-rust/Cargo.toml index 25ee2213040..5842ed5ba9b 100644 --- a/litellm-rust/Cargo.toml +++ b/litellm-rust/Cargo.toml @@ -16,13 +16,15 @@ litellm-core = { path = "crates/core" } litellm-ai-gateway = { path = "crates/ai-gateway", default-features = false } axum = "0.7" pyo3 = "0.23.5" +pyo3-async-runtimes = { version = "0.23.0", features = ["tokio-runtime"] } rand = "0.8" -reqwest = { version = "0.12", default-features = false, features = ["blocking", "json", "rustls-tls"] } +reqwest = { version = "0.12", default-features = false, features = ["blocking", "json", "rustls-tls", "http2", "stream"] } serde = { version = "1.0", features = ["derive"] } serde_json = "1.0" sha2 = "0.10" subtle = "2" thiserror = "2.0" -tokio = { version = "1", features = ["rt-multi-thread", "macros", "time"] } +tokio = { version = "1", features = ["rt-multi-thread", "macros", "time", "net"] } tokio-tungstenite = { version = "0.24", default-features = false, features = ["connect", "rustls-tls-native-roots"] } futures-util = { version = "0.3", default-features = false, features = ["sink", "std"] } +base64 = "0.22" diff --git a/litellm-rust/crates/ai-gateway/Cargo.toml b/litellm-rust/crates/ai-gateway/Cargo.toml index b08fb89d5e8..2f414159158 100644 --- a/litellm-rust/crates/ai-gateway/Cargo.toml +++ b/litellm-rust/crates/ai-gateway/Cargo.toml @@ -23,6 +23,7 @@ tokio = { workspace = true, features = ["rt-multi-thread", "macros", "net", "tim tokio-tungstenite.workspace = true futures-util.workspace = true serde_json.workspace = true +base64.workspace = true axum = { workspace = true, features = ["ws"], optional = true } serde = { workspace = true, optional = true } subtle = { workspace = true, optional = true } diff --git a/litellm-rust/crates/ai-gateway/src/io/ocr.rs b/litellm-rust/crates/ai-gateway/src/io/ocr.rs index 5c32157bc6f..35e511fa982 100644 --- a/litellm-rust/crates/ai-gateway/src/io/ocr.rs +++ b/litellm-rust/crates/ai-gateway/src/io/ocr.rs @@ -1,6 +1,6 @@ //! End-to-end OCR orchestration. //! -//! Owns the whole Mistral OCR call so the Python side stays a thin bridge: +//! Owns supported OCR provider calls so the Python side stays a thin bridge: //! 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. @@ -8,75 +8,135 @@ use std::sync::OnceLock; use std::time::Duration; use litellm_core::error::CoreError; -use litellm_core::ocr::transformation::OcrProviderConfig; +use litellm_core::ocr::transformation::{OcrAuthStrategy, OcrResponseHandling}; use litellm_core::CoreResult; use serde_json::{Map, Value}; -use litellm_core::providers::mistral::ocr::transformation as mistral; -use litellm_core::providers::mistral::ocr::transformation::MISTRAL_OCR_CONFIG; +mod common_utils; + +use common_utils::{ + convert_document_url_to_data_uri, has_header, ocr_provider_config, poll_document_intelligence, + string_headers, truncate_error_body, +}; /// OCR over large documents can take a while; bound it generously rather than /// hanging forever on an unresponsive upstream. The client-level limit is the /// outer ceiling; callers can tighten it per request via ``run_ocr``'s ``timeout``. const OCR_TIMEOUT_SECS: u64 = 600; -/// Maximum upstream body characters retained in error messages. OCR responses -/// can echo document contents and prompts; keep enough for debugging without -/// 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). +/// +/// The Python fallback path uses LiteLLM's standard `BaseLLMHTTPHandler`. This +/// Rust path is opt-in and owns end-to-end OCR I/O, so it cannot call the +/// Python handler directly; keep this route-scoped until litellm-rust has a +/// shared HTTP abstraction. +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") }) } -fn truncate_error_body(body: &str) -> String { - if body.chars().count() <= ERROR_BODY_MAX_CHARS { - return body.to_string(); - } - let truncated: String = body.chars().take(ERROR_BODY_MAX_CHARS).collect(); - format!("{truncated}... (truncated)") +fn upstream_headers( + headers: &[(String, String)], + auth_strategy: OcrAuthStrategy, + api_key: Option<&str>, +) -> Vec<(String, String)> { + let auth_header = api_key.map(|api_key| match auth_strategy { + OcrAuthStrategy::Bearer => ("Authorization".to_string(), format!("Bearer {api_key}")), + OcrAuthStrategy::Header(header_name) => (header_name.to_string(), api_key.to_string()), + }); + auth_header + .into_iter() + .chain(headers.iter().cloned()) + .collect() } -/// Perform a Mistral OCR call end to end and return the normalized response as +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 an 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 config = ocr_provider_config(request.custom_llm_provider, model) + .ok_or_else(|| CoreError::InvalidProvider(request.custom_llm_provider.to_string()))?; + let env_lookup = |key: &str| std::env::var(key).ok(); - 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); + let headers = string_headers(request.extra_headers)?; + let auth_strategy = config.auth_strategy(); + let api_key = (!has_header(&headers, auth_strategy.header_name())) + .then(|| config.resolve_api_key(request.api_key, &env_lookup)) + .transpose()?; + let url = config.complete_url( + request.api_base, + model, + &request.optional_params, + &env_lookup, + )?; + let filtered_params = config.map_ocr_params(&request.optional_params); + let document = if config.requires_data_uri_document() { + convert_document_url_to_data_uri(request.document).await? + } else { + request.document + }; let body = config .transform_ocr_request(model, document, filtered_params)? .data; + let upstream_headers = upstream_headers(&headers, auth_strategy, api_key.as_deref()); - 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).json(&body); + for (key, value) in &upstream_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(); + if config.response_handling() == OcrResponseHandling::AzureDocumentIntelligencePoll + && status.as_u16() == 202 + { + let operation_url = response + .headers() + .get("operation-location") + .and_then(|value| value.to_str().ok()) + .map(str::to_string) + .ok_or_else(|| { + CoreError::InvalidResponse( + "Azure Document Intelligence returned 202 but no Operation-Location header found" + .to_string(), + ) + })?; + let response_json = + poll_document_intelligence(&operation_url, &url, &upstream_headers, request.timeout) + .await?; + return Ok(config + .transform_ocr_response(model, response_json)? + .into_json()); + } + let text = response .text() + .await .map_err(|err| CoreError::Network(err.to_string()))?; if !status.is_success() { @@ -97,6 +157,25 @@ pub fn run_ocr( #[cfg(test)] mod tests { use super::*; + use serde_json::json; + use tokio::io::{AsyncReadExt, AsyncWriteExt}; + use tokio::net::{TcpListener, TcpStream}; + + async fn read_http_headers(socket: &mut TcpStream) -> String { + let mut request = Vec::new(); + let mut buffer = [0_u8; 1024]; + loop { + let n = socket.read(&mut buffer).await.expect("reads request"); + if n == 0 { + break; + } + request.extend_from_slice(&buffer[..n]); + if request.windows(4).any(|window| window == b"\r\n\r\n") { + break; + } + } + String::from_utf8(request).expect("request is utf8") + } #[test] fn truncate_error_body_passes_short_strings_through() { @@ -106,7 +185,7 @@ mod tests { #[test] fn truncate_error_body_caps_long_payloads() { - let body = "x".repeat(ERROR_BODY_MAX_CHARS + 50); + let body = "x".repeat(306); let truncated = truncate_error_body(&body); assert!(truncated.ends_with("... (truncated)")); @@ -115,13 +194,213 @@ mod tests { .expect("truncated marker present") .chars() .count(); - assert_eq!(prefix_chars, ERROR_BODY_MAX_CHARS); + assert_eq!(prefix_chars, 256); } #[test] fn truncate_error_body_does_not_split_multibyte_chars() { - let body = "é".repeat(ERROR_BODY_MAX_CHARS + 10); + let body = "é".repeat(266); let truncated = truncate_error_body(&body); assert!(truncated.is_char_boundary(truncated.len())); } + + #[test] + fn ocr_dispatch_supports_migrated_providers() { + assert!(ocr_provider_config("mistral", "mistral-ocr-latest").is_some()); + assert!(ocr_provider_config("azure_ai", "pixtral-12b-2409") + .expect("azure ai config resolves") + .requires_data_uri_document()); + assert_eq!( + ocr_provider_config("azure_ai", "doc-intelligence/prebuilt-read") + .expect("document intelligence config resolves") + .response_handling(), + OcrResponseHandling::AzureDocumentIntelligencePoll + ); + assert!(ocr_provider_config("vertex_ai", "deepseek-ocr-maas") + .expect("vertex deepseek config resolves") + .supported_ocr_params() + .contains(&"temperature")); + assert!(ocr_provider_config("openai", "gpt-4o").is_none()); + } + + #[test] + fn string_headers_accepts_string_values() { + let headers = json!({ + "x-trace-id": "trace-1" + }) + .as_object() + .unwrap() + .clone(); + + assert_eq!( + string_headers(Some(headers)).expect("string headers accepted"), + vec![("x-trace-id".to_string(), "trace-1".to_string())] + ); + } + + #[test] + fn auth_header_detection_is_case_insensitive() { + let headers = vec![ + ("x-trace-id".to_string(), "trace-1".to_string()), + ("authorization".to_string(), "Bearer sk-test".to_string()), + ]; + + assert!(has_header(&headers, "authorization")); + + let headers = vec![("Authorization".to_string(), "Bearer sk-test".to_string())]; + assert!(has_header(&headers, "authorization")); + + let headers = vec![("x-trace-id".to_string(), "trace-1".to_string())]; + assert!(!has_header(&headers, "authorization")); + } + + #[tokio::test] + async fn ocr_does_not_duplicate_authorization_header_when_header_is_supplied() { + let listener = TcpListener::bind("127.0.0.1:0") + .await + .expect("test listener binds"); + let addr = listener.local_addr().expect("listener has local addr"); + + let server = tokio::spawn(async move { + let (mut socket, _) = listener.accept().await.expect("accepts one request"); + let request = read_http_headers(&mut socket).await; + + let response_body = r#"{"pages":[{"index":0,"markdown":"ok"}],"model":"mistral-ocr-latest","usage_info":{"pages_processed":1}}"#; + let response = format!( + "HTTP/1.1 200 OK\r\ncontent-type: application/json\r\ncontent-length: {}\r\nconnection: close\r\n\r\n{}", + response_body.len(), + response_body + ); + socket + .write_all(response.as_bytes()) + .await + .expect("writes response"); + request + }); + + let mut headers = Map::new(); + headers.insert( + "Authorization".to_string(), + Value::String("Bearer sk-from-python".to_string()), + ); + headers.insert( + "x-trace-id".to_string(), + Value::String("trace-1".to_string()), + ); + + let response = ocr(OcrRequest { + model: "mistral-ocr-latest", + document: json!({ + "type": "document_url", + "document_url": "https://example.com/doc.pdf" + }), + api_key: Some("sk-for-rust-fallback"), + api_base: Some(&format!("http://{addr}")), + custom_llm_provider: "mistral", + extra_headers: Some(headers), + optional_params: Map::new(), + timeout: Some(Duration::from_secs(5)), + }) + .await + .expect("ocr request succeeds"); + + assert_eq!(response["pages"][0]["markdown"], "ok"); + + let request = server.await.expect("server task completes"); + let authorization_count = request + .lines() + .filter(|line| line.to_ascii_lowercase().starts_with("authorization:")) + .count(); + assert_eq!(authorization_count, 1, "{request}"); + assert!( + request.contains("authorization: Bearer sk-from-python") + || request.contains("Authorization: Bearer sk-from-python"), + "{request}" + ); + } + + #[tokio::test] + async fn document_intelligence_poll_uses_resolved_subscription_key() { + let listener = TcpListener::bind("127.0.0.1:0") + .await + .expect("test listener binds"); + let addr = listener.local_addr().expect("listener has local addr"); + let operation_url = format!("http://{addr}/operations/1"); + + let server = tokio::spawn(async move { + let (mut post_socket, _) = listener.accept().await.expect("accepts post request"); + let post_request = read_http_headers(&mut post_socket).await; + let post_response = format!( + "HTTP/1.1 202 Accepted\r\noperation-location: {operation_url}\r\ncontent-length: 0\r\nconnection: close\r\n\r\n" + ); + post_socket + .write_all(post_response.as_bytes()) + .await + .expect("writes post response"); + + let (mut poll_socket, _) = listener.accept().await.expect("accepts poll request"); + let poll_request = read_http_headers(&mut poll_socket).await; + let response_body = r#"{"status":"succeeded","analyzeResult":{"pages":[{"pageNumber":1,"lines":[{"content":"ok"}]}]}}"#; + let poll_response = format!( + "HTTP/1.1 200 OK\r\ncontent-type: application/json\r\ncontent-length: {}\r\nconnection: close\r\n\r\n{}", + response_body.len(), + response_body + ); + poll_socket + .write_all(poll_response.as_bytes()) + .await + .expect("writes poll response"); + (post_request, poll_request) + }); + + let response = ocr(OcrRequest { + model: "prebuilt-read", + document: json!({ + "type": "document_url", + "document_url": "https://example.com/doc.pdf" + }), + api_key: Some("di-key"), + api_base: Some(&format!("http://{addr}")), + custom_llm_provider: "azure_ai/doc-intelligence", + extra_headers: None, + optional_params: Map::new(), + timeout: Some(Duration::from_secs(5)), + }) + .await + .expect("document intelligence request succeeds"); + + assert_eq!(response["pages"][0]["markdown"], "ok"); + + let (post_request, poll_request) = server.await.expect("server task completes"); + assert!( + post_request + .to_ascii_lowercase() + .contains("ocp-apim-subscription-key: di-key"), + "{post_request}" + ); + assert!( + poll_request + .to_ascii_lowercase() + .contains("ocp-apim-subscription-key: di-key"), + "{poll_request}" + ); + } + + #[test] + fn string_headers_rejects_non_string_values() { + let headers = json!({ + "x-retry-count": 3 + }) + .as_object() + .unwrap() + .clone(); + + let err = string_headers(Some(headers)).expect_err("non-string header rejected"); + assert_eq!( + err, + CoreError::InvalidRequest( + "OCR extra_headers.x-retry-count must be a string, got number".to_string() + ) + ); + } } diff --git a/litellm-rust/crates/ai-gateway/src/io/ocr/common_utils.rs b/litellm-rust/crates/ai-gateway/src/io/ocr/common_utils.rs new file mode 100644 index 00000000000..ee0d86c3000 --- /dev/null +++ b/litellm-rust/crates/ai-gateway/src/io/ocr/common_utils.rs @@ -0,0 +1,448 @@ +use std::net::IpAddr; +use std::time::{Duration, Instant}; + +use base64::engine::general_purpose::STANDARD as BASE64_STANDARD; +use base64::Engine; +use litellm_core::error::CoreError; +use litellm_core::ocr::transformation::OcrProviderConfig; +use litellm_core::CoreResult; +use reqwest::Url; +use serde_json::{Map, Value}; + +use litellm_core::providers::azure_ai::ocr::transformation::{ + AZURE_AI_OCR_CONFIG, AZURE_DOCUMENT_INTELLIGENCE_OCR_CONFIG, +}; +use litellm_core::providers::mistral::ocr::transformation::MISTRAL_OCR_CONFIG; +use litellm_core::providers::vertex_ai::ocr::transformation as vertex_ai; +use litellm_core::providers::vertex_ai::ocr::transformation::{ + VERTEX_AI_DEEPSEEK_OCR_CONFIG, VERTEX_AI_OCR_CONFIG, +}; + +use super::http_client; + +const ERROR_BODY_MAX_CHARS: usize = 256; +const AZURE_DOCUMENT_INTELLIGENCE_POLL_TIMEOUT_SECS: u64 = 120; +const DEFAULT_MAX_IMAGE_URL_DOWNLOAD_SIZE_MB: f64 = 50.0; +const MAX_SAFE_FETCH_REDIRECTS: usize = 10; + +pub(super) fn truncate_error_body(body: &str) -> String { + if body.chars().count() <= ERROR_BODY_MAX_CHARS { + return body.to_string(); + } + let truncated: String = body.chars().take(ERROR_BODY_MAX_CHARS).collect(); + format!("{truncated}... (truncated)") +} + +pub(super) fn ocr_provider_config( + provider: &str, + model: &str, +) -> Option<&'static dyn OcrProviderConfig> { + match provider { + "mistral" => Some(&MISTRAL_OCR_CONFIG), + "azure_ai" if is_azure_document_intelligence_model(model) => { + Some(&AZURE_DOCUMENT_INTELLIGENCE_OCR_CONFIG) + } + "azure_ai/doc-intelligence" => Some(&AZURE_DOCUMENT_INTELLIGENCE_OCR_CONFIG), + "azure_ai" => Some(&AZURE_AI_OCR_CONFIG), + "vertex_ai" if vertex_ai::is_deepseek_model(model) => Some(&VERTEX_AI_DEEPSEEK_OCR_CONFIG), + "vertex_ai" => Some(&VERTEX_AI_OCR_CONFIG), + _ => None, + } +} + +fn is_azure_document_intelligence_model(model: &str) -> bool { + let model = model.to_ascii_lowercase(); + model.contains("doc-intelligence") || model.contains("documentintelligence") +} + +pub(super) fn string_headers( + extra_headers: Option>, +) -> CoreResult> { + extra_headers + .unwrap_or_default() + .into_iter() + .map(|(key, value)| { + value + .as_str() + .map(|value| (key.clone(), value.to_string())) + .ok_or_else(|| { + CoreError::InvalidRequest(format!( + "OCR extra_headers.{key} must be a string, got {}", + litellm_core::error::json_type_name(&value) + )) + }) + }) + .collect() +} + +pub(super) fn has_header(headers: &[(String, String)], name: &str) -> bool { + headers + .iter() + .any(|(key, _)| key.eq_ignore_ascii_case(name)) +} + +fn document_url_field(document: &Value) -> CoreResult> { + let Some(object) = document.as_object() else { + return Ok(None); + }; + let Some(doc_type) = object.get("type").and_then(Value::as_str) else { + return Ok(None); + }; + let field = match doc_type { + "document_url" => "document_url", + "image_url" => "image_url", + _ => return Ok(None), + }; + let Some(url) = object.get(field).and_then(Value::as_str) else { + return Ok(None); + }; + Ok(Some((field, url))) +} + +fn is_url_requiring_fetch(url: &str) -> bool { + !url.starts_with("data:") && (url.starts_with("http://") || url.starts_with("https://")) +} + +fn max_document_download_bytes() -> u64 { + let max_size_mb = std::env::var("MAX_IMAGE_URL_DOWNLOAD_SIZE_MB") + .ok() + .and_then(|value| value.parse::().ok()) + .unwrap_or(DEFAULT_MAX_IMAGE_URL_DOWNLOAD_SIZE_MB); + (max_size_mb.max(0.0) * 1024.0 * 1024.0) as u64 +} + +fn is_blocked_ip(ip: IpAddr) -> bool { + match ip { + IpAddr::V4(ip) => { + ip.is_private() + || ip.is_loopback() + || ip.is_link_local() + || ip.is_broadcast() + || ip.is_multicast() + || ip.is_unspecified() + } + IpAddr::V6(ip) => { + let first_segment = ip.segments()[0]; + let is_unique_local = (first_segment & 0xfe00) == 0xfc00; + let is_link_local = (first_segment & 0xffc0) == 0xfe80; + ip.is_loopback() + || ip.is_unspecified() + || ip.is_multicast() + || is_unique_local + || is_link_local + || ip + .to_ipv4_mapped() + .or_else(|| ip.to_ipv4()) + .map(|v4| is_blocked_ip(IpAddr::V4(v4))) + .unwrap_or(false) + } + } +} + +fn blocked_url_error(url: &Url) -> CoreError { + CoreError::InvalidRequest(format!( + "OCR document URL rejected by SSRF protection: {url}" + )) +} + +async fn validate_safe_fetch_url(url: &Url) -> CoreResult<()> { + if !matches!(url.scheme(), "http" | "https") { + return Err(blocked_url_error(url)); + } + + let host = url.host_str().ok_or_else(|| blocked_url_error(url))?; + if let Ok(ip) = host.parse::() { + if is_blocked_ip(ip) { + return Err(blocked_url_error(url)); + } + return Ok(()); + } + + let port = url + .port_or_known_default() + .ok_or_else(|| blocked_url_error(url))?; + let addresses = tokio::net::lookup_host((host, port)) + .await + .map_err(|err| CoreError::Network(err.to_string()))?; + let mut saw_address = false; + for address in addresses { + saw_address = true; + if is_blocked_ip(address.ip()) { + return Err(blocked_url_error(url)); + } + } + if !saw_address { + return Err(blocked_url_error(url)); + } + Ok(()) +} + +fn redirect_location(response: &reqwest::Response, url: &Url) -> CoreResult { + let location = response + .headers() + .get(reqwest::header::LOCATION) + .and_then(|value| value.to_str().ok()) + .ok_or_else(|| { + CoreError::InvalidResponse("OCR document redirect missing Location header".to_string()) + })?; + url.join(location) + .map_err(|err| CoreError::InvalidResponse(format!("invalid OCR document redirect: {err}"))) +} + +async fn safe_get_document_url(url: &str) -> CoreResult<(Url, reqwest::Response)> { + let client = reqwest::Client::builder() + .redirect(reqwest::redirect::Policy::none()) + .build() + .map_err(|err| CoreError::Network(err.to_string()))?; + let mut current_url = Url::parse(url) + .map_err(|err| CoreError::InvalidRequest(format!("invalid OCR document URL: {err}")))?; + + for _ in 0..MAX_SAFE_FETCH_REDIRECTS { + validate_safe_fetch_url(¤t_url).await?; + let response = client + .get(current_url.clone()) + .send() + .await + .map_err(|err| CoreError::Network(err.to_string()))?; + if !response.status().is_redirection() { + return Ok((current_url, response)); + } + current_url = redirect_location(&response, ¤t_url)?; + } + + Err(CoreError::InvalidRequest( + "Too many redirects while fetching OCR document URL".to_string(), + )) +} + +fn enforce_download_size(content_length: u64, max_bytes: u64, url: &Url) -> CoreResult<()> { + if max_bytes == 0 { + return Err(CoreError::InvalidRequest(format!( + "OCR document URL download is disabled (MAX_IMAGE_URL_DOWNLOAD_SIZE_MB=0). url={url}" + ))); + } + if content_length > max_bytes { + let size_mb = content_length as f64 / (1024.0 * 1024.0); + let max_size_mb = max_bytes as f64 / (1024.0 * 1024.0); + return Err(CoreError::InvalidRequest(format!( + "OCR document size ({size_mb:.2}MB) exceeds maximum allowed size ({max_size_mb:.2}MB). url={url}" + ))); + } + Ok(()) +} + +async fn read_response_with_limit( + mut response: reqwest::Response, + url: &Url, +) -> CoreResult> { + let max_bytes = max_document_download_bytes(); + if let Some(content_length) = response.content_length() { + enforce_download_size(content_length, max_bytes, url)?; + } else { + enforce_download_size(0, max_bytes, url)?; + } + + let mut bytes = Vec::new(); + let mut bytes_downloaded: u64 = 0; + while let Some(chunk) = response + .chunk() + .await + .map_err(|err| CoreError::Network(err.to_string()))? + { + bytes_downloaded += chunk.len() as u64; + enforce_download_size(bytes_downloaded, max_bytes, url)?; + bytes.extend_from_slice(&chunk); + } + Ok(bytes) +} + +pub(super) async fn convert_document_url_to_data_uri(document: Value) -> CoreResult { + let Some((field, url)) = document_url_field(&document)? else { + return Ok(document); + }; + if !is_url_requiring_fetch(url) { + return Ok(document); + } + + let (final_url, response) = safe_get_document_url(url).await?; + let status = response.status(); + if !status.is_success() { + let body = response.text().await.unwrap_or_default(); + return Err(CoreError::Http { + status: status.as_u16(), + body: truncate_error_body(&body), + }); + } + let content_type = response + .headers() + .get(reqwest::header::CONTENT_TYPE) + .and_then(|value| value.to_str().ok()) + .and_then(|value| value.split(';').next()) + .map(str::trim) + .filter(|value| !value.is_empty()) + .unwrap_or("application/octet-stream") + .to_string(); + let bytes = read_response_with_limit(response, &final_url).await?; + let data_uri = format!( + "data:{content_type};base64,{}", + BASE64_STANDARD.encode(bytes) + ); + + let mut transformed = document + .as_object() + .cloned() + .ok_or_else(|| CoreError::InvalidRequest("OCR document must be an object".to_string()))?; + transformed.insert(field.to_string(), Value::String(data_uri)); + Ok(Value::Object(transformed)) +} + +fn same_origin(left: &str, right: &str) -> bool { + let Ok(left) = reqwest::Url::parse(left) else { + return false; + }; + let Ok(right) = reqwest::Url::parse(right) else { + return false; + }; + left.scheme() == right.scheme() + && left.host_str() == right.host_str() + && left.port_or_known_default() == right.port_or_known_default() +} + +fn retry_after_secs(response: &reqwest::Response) -> u64 { + response + .headers() + .get(reqwest::header::RETRY_AFTER) + .and_then(|value| value.to_str().ok()) + .and_then(|value| value.parse::().ok()) + .unwrap_or(2) +} + +fn operation_status(response_json: &Value) -> CoreResult<&str> { + let status = response_json + .get("status") + .and_then(Value::as_str) + .ok_or(CoreError::MissingField("status"))?; + match status { + "succeeded" => Ok("succeeded"), + "running" | "notStarted" => Ok("running"), + "failed" => { + let message = response_json + .get("error") + .and_then(|error| error.get("message")) + .and_then(Value::as_str) + .unwrap_or("Unknown error"); + Err(CoreError::InvalidResponse(format!( + "Azure Document Intelligence analysis failed: {message}" + ))) + } + other => Err(CoreError::InvalidResponse(format!( + "Unknown operation status: {other}" + ))), + } +} + +pub(super) async fn poll_document_intelligence( + operation_url: &str, + original_url: &str, + headers: &[(String, String)], + timeout: Option, +) -> CoreResult { + if !same_origin(operation_url, original_url) { + return Err(CoreError::InvalidResponse( + "Azure Document Intelligence: rejected cross-origin polling URL".to_string(), + )); + } + + let start = Instant::now(); + let timeout = timeout.unwrap_or(Duration::from_secs( + AZURE_DOCUMENT_INTELLIGENCE_POLL_TIMEOUT_SECS, + )); + loop { + if start.elapsed() > timeout { + return Err(CoreError::Network(format!( + "Azure Document Intelligence operation polling timed out after {} seconds", + timeout.as_secs() + ))); + } + + let mut request_builder = http_client().get(operation_url); + for (key, value) in headers { + if key.eq_ignore_ascii_case("ocp-apim-subscription-key") { + request_builder = request_builder.header(key, value); + } + } + let response = request_builder + .send() + .await + .map_err(|err| CoreError::Network(err.to_string()))?; + let retry_after = retry_after_secs(&response); + let status = response.status(); + let text = response + .text() + .await + .map_err(|err| CoreError::Network(err.to_string()))?; + if !status.is_success() { + return Err(CoreError::Http { + status: status.as_u16(), + body: truncate_error_body(&text), + }); + } + let response_json: Value = serde_json::from_str(&text).map_err(|err| { + CoreError::InvalidResponse(format!("invalid Azure DI poll response JSON: {err}")) + })?; + if operation_status(&response_json)? == "succeeded" { + return Ok(response_json); + } + tokio::time::sleep(Duration::from_secs(retry_after)).await; + } +} + +#[cfg(test)] +mod tests { + use super::*; + use serde_json::json; + + #[test] + fn blocks_private_and_metadata_ips() { + assert!(is_blocked_ip("127.0.0.1".parse().unwrap())); + assert!(is_blocked_ip("10.0.0.1".parse().unwrap())); + assert!(is_blocked_ip("169.254.169.254".parse().unwrap())); + assert!(is_blocked_ip("::1".parse().unwrap())); + assert!(is_blocked_ip("fd00::1".parse().unwrap())); + assert!(is_blocked_ip("fe80::1".parse().unwrap())); + assert!(is_blocked_ip("::ffff:169.254.169.254".parse().unwrap())); + assert!(is_blocked_ip("::ffff:10.0.0.1".parse().unwrap())); + assert!(!is_blocked_ip("8.8.8.8".parse().unwrap())); + assert!(!is_blocked_ip("::ffff:8.8.8.8".parse().unwrap())); + } + + #[tokio::test] + async fn convert_document_url_rejects_loopback_fetch() { + let error = convert_document_url_to_data_uri(json!({ + "type": "image_url", + "image_url": "http://127.0.0.1/image.png" + })) + .await + .unwrap_err(); + + assert!(matches!( + error, + CoreError::InvalidRequest(message) + if message.contains("SSRF protection") + )); + } + + #[tokio::test] + async fn convert_document_url_leaves_data_uri_untouched() { + let document = json!({ + "type": "image_url", + "image_url": "data:image/png;base64,abcd" + }); + + let transformed = convert_document_url_to_data_uri(document.clone()) + .await + .unwrap(); + + assert_eq!(transformed, document); + } +} diff --git a/litellm-rust/crates/core/src/error.rs b/litellm-rust/crates/core/src/error.rs index 9b29260cca4..b3e0519b772 100644 --- a/litellm-rust/crates/core/src/error.rs +++ b/litellm-rust/crates/core/src/error.rs @@ -13,6 +13,10 @@ pub enum CoreError { MissingField(&'static str), #[error("invalid response: {0}")] InvalidResponse(String), + #[error("invalid provider: {0}")] + InvalidProvider(String), + #[error("invalid request: {0}")] + InvalidRequest(String), #[error("{0}")] Auth(String), #[error("OCR request failed with status {status}: {body}")] diff --git a/litellm-rust/crates/core/src/ocr/transformation.rs b/litellm-rust/crates/core/src/ocr/transformation.rs index 7353d9d22c4..cb3e735e533 100644 --- a/litellm-rust/crates/core/src/ocr/transformation.rs +++ b/litellm-rust/crates/core/src/ocr/transformation.rs @@ -4,7 +4,28 @@ use crate::CoreResult; use super::types::{OcrRequestData, OcrResponseData}; -pub trait OcrProviderConfig { +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum OcrAuthStrategy { + Bearer, + Header(&'static str), +} + +impl OcrAuthStrategy { + pub fn header_name(self) -> &'static str { + match self { + Self::Bearer => "authorization", + Self::Header(header_name) => header_name, + } + } +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum OcrResponseHandling { + Json, + AzureDocumentIntelligencePoll, +} + +pub trait OcrProviderConfig: Sync { fn supported_ocr_params(&self) -> &'static [&'static str]; fn map_ocr_params(&self, non_default_params: &Map) -> Map { @@ -29,4 +50,30 @@ pub trait OcrProviderConfig { model: &str, response_json: Value, ) -> CoreResult; + + fn complete_url( + &self, + api_base: Option<&str>, + model: &str, + optional_params: &Map, + env_lookup: &dyn Fn(&str) -> Option, + ) -> CoreResult; + + fn resolve_api_key( + &self, + api_key: Option<&str>, + env_lookup: &dyn Fn(&str) -> Option, + ) -> CoreResult; + + fn auth_strategy(&self) -> OcrAuthStrategy { + OcrAuthStrategy::Bearer + } + + fn requires_data_uri_document(&self) -> bool { + false + } + + fn response_handling(&self) -> OcrResponseHandling { + OcrResponseHandling::Json + } } diff --git a/litellm-rust/crates/core/src/providers/azure_ai/mod.rs b/litellm-rust/crates/core/src/providers/azure_ai/mod.rs new file mode 100644 index 00000000000..3621ff6a2fd --- /dev/null +++ b/litellm-rust/crates/core/src/providers/azure_ai/mod.rs @@ -0,0 +1 @@ +pub mod ocr; diff --git a/litellm-rust/crates/core/src/providers/azure_ai/ocr/mod.rs b/litellm-rust/crates/core/src/providers/azure_ai/ocr/mod.rs new file mode 100644 index 00000000000..f239b6921fa --- /dev/null +++ b/litellm-rust/crates/core/src/providers/azure_ai/ocr/mod.rs @@ -0,0 +1 @@ +pub mod transformation; diff --git a/litellm-rust/crates/core/src/providers/azure_ai/ocr/transformation.rs b/litellm-rust/crates/core/src/providers/azure_ai/ocr/transformation.rs new file mode 100644 index 00000000000..060073acd47 --- /dev/null +++ b/litellm-rust/crates/core/src/providers/azure_ai/ocr/transformation.rs @@ -0,0 +1,520 @@ +use std::collections::BTreeSet; + +use crate::error::{json_type_name, CoreError, CoreResult}; +use crate::ocr::transformation::{OcrAuthStrategy, OcrProviderConfig, OcrResponseHandling}; +use crate::ocr::types::{OcrRequestData, OcrResponseData}; +use serde_json::{json, Map, Value}; + +use crate::providers::mistral::ocr::transformation::MISTRAL_OCR_CONFIG; + +const AZURE_AI_API_KEY_ENV: &str = "AZURE_AI_API_KEY"; +const AZURE_AI_API_BASE_ENV: &str = "AZURE_AI_API_BASE"; +const AZURE_DOCUMENT_INTELLIGENCE_API_KEY_ENV: &str = "AZURE_DOCUMENT_INTELLIGENCE_API_KEY"; +const AZURE_DOCUMENT_INTELLIGENCE_ENDPOINT_ENV: &str = "AZURE_DOCUMENT_INTELLIGENCE_ENDPOINT"; +const AZURE_DOCUMENT_INTELLIGENCE_API_VERSION: &str = "2024-11-30"; +const AZURE_DOCUMENT_INTELLIGENCE_DEFAULT_DPI: i64 = 96; + +const AZURE_DOCUMENT_INTELLIGENCE_SUPPORTED_OCR_PARAMS: &[&str] = &["pages"]; + +pub struct AzureAiOcrConfig; +pub struct AzureDocumentIntelligenceOcrConfig; + +pub const AZURE_AI_OCR_CONFIG: AzureAiOcrConfig = AzureAiOcrConfig; +pub const AZURE_DOCUMENT_INTELLIGENCE_OCR_CONFIG: AzureDocumentIntelligenceOcrConfig = + AzureDocumentIntelligenceOcrConfig; + +fn non_empty(value: Option<&str>) -> Option<&str> { + value.map(str::trim).filter(|value| !value.is_empty()) +} + +fn resolve_value( + explicit: Option<&str>, + env_name: &str, + env_lookup: &dyn Fn(&str) -> Option, + missing_message: &str, +) -> CoreResult { + non_empty(explicit) + .map(str::to_string) + .or_else(|| env_lookup(env_name).filter(|value| !value.trim().is_empty())) + .ok_or_else(|| CoreError::Auth(missing_message.to_string())) +} + +pub fn resolve_azure_ai_api_key( + api_key: Option<&str>, + env_lookup: &dyn Fn(&str) -> Option, +) -> CoreResult { + resolve_value( + api_key, + AZURE_AI_API_KEY_ENV, + env_lookup, + "Missing Azure AI API Key - A call is being made to Azure AI but no key is set either in the environment variables or via params", + ) +} + +pub fn resolve_azure_ai_api_base( + api_base: Option<&str>, + env_lookup: &dyn Fn(&str) -> Option, +) -> CoreResult { + resolve_value( + api_base, + AZURE_AI_API_BASE_ENV, + env_lookup, + "Missing Azure AI API Base - Set AZURE_AI_API_BASE environment variable or pass api_base parameter", + ) +} + +pub fn complete_azure_ai_url( + api_base: Option<&str>, + env_lookup: &dyn Fn(&str) -> Option, +) -> CoreResult { + let base = resolve_azure_ai_api_base(api_base, env_lookup)?; + Ok(format!( + "{}/providers/mistral/azure/ocr", + base.trim_end_matches('/') + )) +} + +pub fn resolve_document_intelligence_api_key( + api_key: Option<&str>, + env_lookup: &dyn Fn(&str) -> Option, +) -> CoreResult { + resolve_value( + api_key, + AZURE_DOCUMENT_INTELLIGENCE_API_KEY_ENV, + env_lookup, + "Missing Azure Document Intelligence API Key - Set AZURE_DOCUMENT_INTELLIGENCE_API_KEY environment variable or pass api_key parameter", + ) +} + +pub fn resolve_document_intelligence_endpoint( + api_base: Option<&str>, + env_lookup: &dyn Fn(&str) -> Option, +) -> CoreResult { + resolve_value( + api_base, + AZURE_DOCUMENT_INTELLIGENCE_ENDPOINT_ENV, + env_lookup, + "Missing Azure Document Intelligence Endpoint - Set AZURE_DOCUMENT_INTELLIGENCE_ENDPOINT environment variable or pass api_base parameter", + ) +} + +fn encode_model_id(model: &str) -> String { + let model_id = model.rsplit('/').next().unwrap_or(model); + model_id + .bytes() + .flat_map(|byte| match byte { + b'A'..=b'Z' | b'a'..=b'z' | b'0'..=b'9' | b'-' | b'_' | b'.' | b'~' => { + vec![byte as char] + } + _ => format!("%{byte:02X}").chars().collect(), + }) + .collect() +} + +fn pages_token_is_valid(token: &str) -> bool { + let mut parts = token.split('-'); + let Some(start) = parts.next() else { + return false; + }; + if start.is_empty() || !start.chars().all(|ch| ch.is_ascii_digit()) { + return false; + } + match parts.next() { + None => true, + Some(end) => { + !end.is_empty() && end.chars().all(|ch| ch.is_ascii_digit()) && parts.next().is_none() + } + } +} + +fn normalize_pages_param(pages: &Value) -> CoreResult> { + match pages { + Value::String(value) => { + let normalized = value + .split(',') + .map(str::trim) + .collect::>() + .join(","); + if normalized.split(',').all(pages_token_is_valid) { + Ok(Some(normalized)) + } else { + Err(CoreError::InvalidRequest(format!( + "Invalid `pages` string for Azure Document Intelligence: {value:?}. Expected format like '1-3,5,7-9'." + ))) + } + } + Value::Array(values) => { + if values.is_empty() { + return Ok(None); + } + if values.iter().all(Value::is_i64) { + let mut pages = BTreeSet::new(); + for value in values { + let page = value.as_i64().expect("checked is_i64"); + if page < 0 { + return Err(CoreError::InvalidRequest( + "`pages` integers must be >= 0 (Mistral 0-based indices)".to_string(), + )); + } + pages.insert(page + 1); + } + return Ok(Some( + pages + .into_iter() + .map(|page| page.to_string()) + .collect::>() + .join(","), + )); + } + if values.iter().all(Value::is_string) { + let normalized = values + .iter() + .filter_map(Value::as_str) + .map(str::trim) + .collect::>() + .join(","); + if normalized.split(',').all(pages_token_is_valid) { + return Ok(Some(normalized)); + } + return Err(CoreError::InvalidRequest(format!( + "Invalid `pages` list for Azure Document Intelligence: {values:?}. Expected tokens like '1' or '3-5'." + ))); + } + Err(CoreError::InvalidRequest( + "`pages` must be a list[int] (0-based, Mistral-style) or a string like '1-3,5,7-9'." + .to_string(), + )) + } + _ => Err(CoreError::InvalidRequest( + "`pages` must be a list[int] (0-based, Mistral-style) or a string like '1-3,5,7-9'." + .to_string(), + )), + } +} + +pub fn complete_document_intelligence_url( + api_base: Option<&str>, + model: &str, + optional_params: &Map, + env_lookup: &dyn Fn(&str) -> Option, +) -> CoreResult { + let endpoint = resolve_document_intelligence_endpoint(api_base, env_lookup)?; + let mut url = format!( + "{}/documentintelligence/documentModels/{}:analyze?api-version={}", + endpoint.trim_end_matches('/'), + encode_model_id(model), + AZURE_DOCUMENT_INTELLIGENCE_API_VERSION + ); + + if let Some(pages) = optional_params.get("pages") { + if let Some(normalized) = normalize_pages_param(pages)? { + url.push_str("&pages="); + url.push_str(&normalized); + } + } + + Ok(url) +} + +fn document_url_from_mistral_document(document: &Value) -> CoreResult<&str> { + let object = document.as_object().ok_or_else(|| CoreError::InvalidType { + expected: "object", + actual: json_type_name(document), + })?; + let doc_type = object + .get("type") + .and_then(Value::as_str) + .ok_or(CoreError::MissingField("document.type"))?; + let field_name = match doc_type { + "document_url" => "document_url", + "image_url" => "image_url", + other => { + return Err(CoreError::InvalidRequest(format!( + "Invalid document type: {other}. Must be 'document_url' or 'image_url'" + ))) + } + }; + object + .get(field_name) + .and_then(Value::as_str) + .filter(|value| !value.is_empty()) + .ok_or(CoreError::MissingField(field_name)) +} + +fn extract_base64_from_data_uri(data_uri: &str) -> &str { + data_uri + .split_once(',') + .map(|(_, data)| data) + .unwrap_or(data_uri) +} + +fn page_markdown(page: &Map) -> String { + page.get("lines") + .and_then(Value::as_array) + .map(|lines| { + lines + .iter() + .filter_map(|line| line.get("content").and_then(Value::as_str)) + .collect::>() + .join("\n") + }) + .unwrap_or_default() +} + +fn page_dimensions(page: &Map) -> Value { + let width = page.get("width").and_then(Value::as_f64).unwrap_or(8.5); + let height = page.get("height").and_then(Value::as_f64).unwrap_or(11.0); + let unit = page.get("unit").and_then(Value::as_str).unwrap_or("inch"); + let (width, height) = if unit == "inch" { + ( + (width * AZURE_DOCUMENT_INTELLIGENCE_DEFAULT_DPI as f64) as i64, + (height * AZURE_DOCUMENT_INTELLIGENCE_DEFAULT_DPI as f64) as i64, + ) + } else { + (width as i64, height as i64) + }; + json!({ + "width": width, + "height": height, + "dpi": AZURE_DOCUMENT_INTELLIGENCE_DEFAULT_DPI, + }) +} + +impl OcrProviderConfig for AzureAiOcrConfig { + fn supported_ocr_params(&self) -> &'static [&'static str] { + MISTRAL_OCR_CONFIG.supported_ocr_params() + } + + fn transform_ocr_request( + &self, + model: &str, + document: Value, + optional_params: Map, + ) -> CoreResult { + MISTRAL_OCR_CONFIG.transform_ocr_request(model, document, optional_params) + } + + fn transform_ocr_response( + &self, + model: &str, + response_json: Value, + ) -> CoreResult { + MISTRAL_OCR_CONFIG.transform_ocr_response(model, response_json) + } + + fn complete_url( + &self, + api_base: Option<&str>, + _model: &str, + _optional_params: &Map, + env_lookup: &dyn Fn(&str) -> Option, + ) -> CoreResult { + complete_azure_ai_url(api_base, env_lookup) + } + + fn resolve_api_key( + &self, + api_key: Option<&str>, + env_lookup: &dyn Fn(&str) -> Option, + ) -> CoreResult { + resolve_azure_ai_api_key(api_key, env_lookup) + } + + fn requires_data_uri_document(&self) -> bool { + true + } +} + +impl OcrProviderConfig for AzureDocumentIntelligenceOcrConfig { + fn supported_ocr_params(&self) -> &'static [&'static str] { + AZURE_DOCUMENT_INTELLIGENCE_SUPPORTED_OCR_PARAMS + } + + fn transform_ocr_request( + &self, + _model: &str, + document: Value, + _optional_params: Map, + ) -> CoreResult { + let document_url = document_url_from_mistral_document(&document)?; + let mut data = Map::new(); + if document_url.starts_with("data:") { + data.insert( + "base64Source".to_string(), + Value::String(extract_base64_from_data_uri(document_url).to_string()), + ); + } else { + data.insert( + "urlSource".to_string(), + Value::String(document_url.to_string()), + ); + } + Ok(OcrRequestData { + data: Value::Object(data), + files: None, + }) + } + + fn transform_ocr_response( + &self, + model: &str, + response_json: Value, + ) -> CoreResult { + let response = response_json + .as_object() + .ok_or_else(|| CoreError::InvalidType { + expected: "object", + actual: json_type_name(&response_json), + })?; + let status = response + .get("status") + .and_then(Value::as_str) + .ok_or(CoreError::MissingField("status"))?; + if status != "succeeded" { + return Err(CoreError::InvalidResponse(format!( + "Azure Document Intelligence analysis failed with status: {status}" + ))); + } + + let azure_pages = response + .get("analyzeResult") + .and_then(|result| result.get("pages")) + .and_then(Value::as_array) + .cloned() + .unwrap_or_default(); + + let pages = azure_pages + .iter() + .filter_map(Value::as_object) + .map(|page| { + let page_number = page.get("pageNumber").and_then(Value::as_i64).unwrap_or(1); + json!({ + "index": page_number - 1, + "markdown": page_markdown(page), + "dimensions": page_dimensions(page), + }) + }) + .collect::>(); + + Ok(OcrResponseData { + usage_info: Some(json!({ + "pages_processed": pages.len(), + "doc_size_bytes": null, + })), + pages, + model: model.to_string(), + document_annotation: None, + object: "ocr".to_string(), + }) + } + + fn complete_url( + &self, + api_base: Option<&str>, + model: &str, + optional_params: &Map, + env_lookup: &dyn Fn(&str) -> Option, + ) -> CoreResult { + complete_document_intelligence_url(api_base, model, optional_params, env_lookup) + } + + fn resolve_api_key( + &self, + api_key: Option<&str>, + env_lookup: &dyn Fn(&str) -> Option, + ) -> CoreResult { + resolve_document_intelligence_api_key(api_key, env_lookup) + } + + fn auth_strategy(&self) -> OcrAuthStrategy { + OcrAuthStrategy::Header("Ocp-Apim-Subscription-Key") + } + + fn response_handling(&self) -> OcrResponseHandling { + OcrResponseHandling::AzureDocumentIntelligencePoll + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn azure_ai_reuses_mistral_body_transform() { + let body = AZURE_AI_OCR_CONFIG + .transform_ocr_request( + "pixtral-12b-2409", + json!({"type": "document_url", "document_url": "data:application/pdf;base64,abc"}), + serde_json::Map::from_iter([("include_image_base64".to_string(), json!(true))]), + ) + .expect("request transforms") + .data; + + assert_eq!(body["model"], "pixtral-12b-2409"); + assert_eq!(body["include_image_base64"], true); + assert_eq!( + body["document"]["document_url"], + "data:application/pdf;base64,abc" + ); + } + + #[test] + fn document_intelligence_url_normalizes_zero_based_pages() { + let params = serde_json::Map::from_iter([("pages".to_string(), json!([2, 0, 2]))]); + let url = complete_document_intelligence_url( + Some("https://example.cognitiveservices.azure.com/"), + "azure_ai/doc-intelligence/prebuilt-layout", + ¶ms, + &|_| None, + ) + .expect("url builds"); + + assert_eq!( + url, + "https://example.cognitiveservices.azure.com/documentintelligence/documentModels/prebuilt-layout:analyze?api-version=2024-11-30&pages=1,3" + ); + } + + #[test] + fn document_intelligence_request_uses_base64_source_for_data_uri() { + let body = AZURE_DOCUMENT_INTELLIGENCE_OCR_CONFIG + .transform_ocr_request( + "prebuilt-read", + json!({"type": "document_url", "document_url": "data:application/pdf;base64,abc123"}), + Map::new(), + ) + .expect("request transforms") + .data; + + assert_eq!(body, json!({"base64Source": "abc123"})); + } + + #[test] + fn document_intelligence_response_normalizes_pages() { + let response = AZURE_DOCUMENT_INTELLIGENCE_OCR_CONFIG + .transform_ocr_response( + "prebuilt-layout", + json!({ + "status": "succeeded", + "analyzeResult": { + "pages": [{ + "pageNumber": 2, + "width": 8.5, + "height": 11, + "unit": "inch", + "lines": [{"content": "hello"}, {"content": "world"}] + }] + } + }), + ) + .expect("response transforms"); + + assert_eq!(response.pages[0]["index"], 1); + assert_eq!(response.pages[0]["markdown"], "hello\nworld"); + assert_eq!(response.pages[0]["dimensions"]["width"], 816); + assert_eq!( + response.usage_info, + Some(json!({"pages_processed": 1, "doc_size_bytes": null})) + ); + } +} diff --git a/litellm-rust/crates/core/src/providers/mistral/ocr/transformation.rs b/litellm-rust/crates/core/src/providers/mistral/ocr/transformation.rs index d5155991448..386457f8b84 100644 --- a/litellm-rust/crates/core/src/providers/mistral/ocr/transformation.rs +++ b/litellm-rust/crates/core/src/providers/mistral/ocr/transformation.rs @@ -132,6 +132,24 @@ impl OcrProviderConfig for MistralOcrConfig { object: "ocr".to_string(), }) } + + fn complete_url( + &self, + api_base: Option<&str>, + _model: &str, + _optional_params: &Map, + _env_lookup: &dyn Fn(&str) -> Option, + ) -> CoreResult { + Ok(complete_url(api_base)) + } + + fn resolve_api_key( + &self, + api_key: Option<&str>, + env_lookup: &dyn Fn(&str) -> Option, + ) -> CoreResult { + resolve_api_key(api_key, env_lookup) + } } pub fn supported_ocr_params() -> &'static [&'static str] { diff --git a/litellm-rust/crates/core/src/providers/mod.rs b/litellm-rust/crates/core/src/providers/mod.rs index 42207f0de0a..d75e750a0ba 100644 --- a/litellm-rust/crates/core/src/providers/mod.rs +++ b/litellm-rust/crates/core/src/providers/mod.rs @@ -1,2 +1,4 @@ +pub mod azure_ai; pub mod mistral; pub mod openai; +pub mod vertex_ai; diff --git a/litellm-rust/crates/core/src/providers/vertex_ai/mod.rs b/litellm-rust/crates/core/src/providers/vertex_ai/mod.rs new file mode 100644 index 00000000000..3621ff6a2fd --- /dev/null +++ b/litellm-rust/crates/core/src/providers/vertex_ai/mod.rs @@ -0,0 +1 @@ +pub mod ocr; diff --git a/litellm-rust/crates/core/src/providers/vertex_ai/ocr/mod.rs b/litellm-rust/crates/core/src/providers/vertex_ai/ocr/mod.rs new file mode 100644 index 00000000000..f239b6921fa --- /dev/null +++ b/litellm-rust/crates/core/src/providers/vertex_ai/ocr/mod.rs @@ -0,0 +1 @@ +pub mod transformation; diff --git a/litellm-rust/crates/core/src/providers/vertex_ai/ocr/transformation.rs b/litellm-rust/crates/core/src/providers/vertex_ai/ocr/transformation.rs new file mode 100644 index 00000000000..8639926c435 --- /dev/null +++ b/litellm-rust/crates/core/src/providers/vertex_ai/ocr/transformation.rs @@ -0,0 +1,435 @@ +use crate::error::{json_type_name, CoreError, CoreResult}; +use crate::ocr::transformation::OcrProviderConfig; +use crate::ocr::types::{OcrRequestData, OcrResponseData}; +use serde_json::{json, Map, Value}; + +use crate::providers::mistral::ocr::transformation::MISTRAL_OCR_CONFIG; + +const VERTEX_DEFAULT_LOCATION: &str = "us-central1"; +const VERTEX_DEFAULT_DEEPSEEK_API_BASE: &str = "https://aiplatform.googleapis.com"; +const VERTEX_AI_API_KEY_ENV: &str = "VERTEX_AI_API_KEY"; +const VERTEXAI_API_KEY_ENV: &str = "VERTEXAI_API_KEY"; +const VERTEXAI_PROJECT_ENV: &str = "VERTEXAI_PROJECT"; +const VERTEXAI_LOCATION_ENV: &str = "VERTEXAI_LOCATION"; +const VERTEX_LOCATION_ENV: &str = "VERTEX_LOCATION"; + +#[rustfmt::skip] +const DEEPSEEK_SUPPORTED_OCR_PARAMS: &[&str] = &[ + "stream", + "temperature", + "max_tokens", + "top_p", + "n", + "stop", +]; + +pub struct VertexAiOcrConfig; +pub struct VertexAiDeepSeekOcrConfig; + +pub const VERTEX_AI_OCR_CONFIG: VertexAiOcrConfig = VertexAiOcrConfig; +pub const VERTEX_AI_DEEPSEEK_OCR_CONFIG: VertexAiDeepSeekOcrConfig = VertexAiDeepSeekOcrConfig; + +fn string_param<'a>(params: &'a Map, keys: &[&str]) -> Option<&'a str> { + keys.iter() + .find_map(|key| params.get(*key).and_then(Value::as_str)) + .map(str::trim) + .filter(|value| !value.is_empty()) +} + +pub fn is_deepseek_model(model: &str) -> bool { + model.to_ascii_lowercase().contains("deepseek") +} + +pub fn resolve_vertex_api_key( + api_key: Option<&str>, + env_lookup: &dyn Fn(&str) -> Option, +) -> CoreResult { + api_key + .map(str::trim) + .filter(|key| !key.is_empty()) + .map(str::to_string) + .or_else(|| env_lookup(VERTEX_AI_API_KEY_ENV).filter(|key| !key.trim().is_empty())) + .or_else(|| env_lookup(VERTEXAI_API_KEY_ENV).filter(|key| !key.trim().is_empty())) + .ok_or_else(|| { + CoreError::Auth( + "Missing Vertex AI access token - pass api_key or provide Authorization via extra_headers" + .to_string(), + ) + }) +} + +fn vertex_project( + params: &Map, + env_lookup: &dyn Fn(&str) -> Option, +) -> CoreResult { + string_param(params, &["vertex_project", "vertex_ai_project"]) + .map(str::to_string) + .or_else(|| env_lookup(VERTEXAI_PROJECT_ENV).filter(|value| !value.trim().is_empty())) + .ok_or_else(|| { + CoreError::InvalidRequest( + "Missing vertex_project - Set VERTEXAI_PROJECT environment variable or pass vertex_project parameter" + .to_string(), + ) + }) +} + +fn vertex_location( + params: &Map, + env_lookup: &dyn Fn(&str) -> Option, +) -> String { + string_param(params, &["vertex_location", "vertex_ai_location"]) + .map(str::to_string) + .or_else(|| env_lookup(VERTEXAI_LOCATION_ENV).filter(|value| !value.trim().is_empty())) + .or_else(|| env_lookup(VERTEX_LOCATION_ENV).filter(|value| !value.trim().is_empty())) + .unwrap_or_else(|| VERTEX_DEFAULT_LOCATION.to_string()) +} + +fn vertex_mistral_api_base(api_base: Option<&str>, location: &str) -> String { + api_base + .map(str::trim) + .filter(|value| !value.is_empty()) + .map(str::to_string) + .unwrap_or_else(|| format!("https://{location}-aiplatform.googleapis.com")) + .trim_end_matches('/') + .to_string() +} + +pub fn complete_vertex_mistral_url( + api_base: Option<&str>, + model: &str, + optional_params: &Map, + env_lookup: &dyn Fn(&str) -> Option, +) -> CoreResult { + let project = vertex_project(optional_params, env_lookup)?; + let location = vertex_location(optional_params, env_lookup); + let base = vertex_mistral_api_base(api_base, &location); + Ok(format!( + "{base}/v1/projects/{project}/locations/{location}/publishers/mistralai/models/{model}:rawPredict" + )) +} + +pub fn complete_vertex_deepseek_url( + api_base: Option<&str>, + optional_params: &Map, + env_lookup: &dyn Fn(&str) -> Option, +) -> CoreResult { + let project = vertex_project(optional_params, env_lookup)?; + let location = vertex_location(optional_params, env_lookup); + let base = api_base + .map(str::trim) + .filter(|value| !value.is_empty()) + .unwrap_or(VERTEX_DEFAULT_DEEPSEEK_API_BASE) + .trim_end_matches('/'); + Ok(format!( + "{base}/v1/projects/{project}/locations/{location}/endpoints/openapi/chat/completions" + )) +} + +fn document_content_item(document: &Value) -> CoreResult { + let object = document.as_object().ok_or_else(|| CoreError::InvalidType { + expected: "object", + actual: json_type_name(document), + })?; + let doc_type = object + .get("type") + .and_then(Value::as_str) + .ok_or(CoreError::MissingField("document.type"))?; + let url_field = match doc_type { + "image_url" => "image_url", + "document_url" => "document_url", + other => { + return Err(CoreError::InvalidRequest(format!( + "Unsupported document type: {other}. Expected 'image_url' or 'document_url'" + ))) + } + }; + let url = object + .get(url_field) + .and_then(Value::as_str) + .filter(|value| !value.is_empty()) + .ok_or(CoreError::MissingField(url_field))?; + + Ok(json!({ + "type": "image_url", + "image_url": url, + })) +} + +fn deepseek_model_name(model: &str) -> String { + if model.starts_with("deepseek-ai/") { + model.to_string() + } else { + format!("deepseek-ai/{model}") + } +} + +fn first_choice_content(response: &Value) -> CoreResult { + response + .get("choices") + .and_then(Value::as_array) + .and_then(|choices| choices.first()) + .and_then(|choice| choice.get("message")) + .and_then(|message| message.get("content")) + .cloned() + .filter(|content| match content { + Value::String(value) => !value.is_empty(), + Value::Object(_) => true, + _ => false, + }) + .ok_or_else(|| { + CoreError::InvalidResponse("No content in DeepSeek OCR response".to_string()) + }) +} + +fn ocr_data_from_content(content: Value, usage: Option, model: &str) -> Value { + match content { + Value::String(content) => { + if content.trim_start().starts_with('{') { + serde_json::from_str(&content).unwrap_or_else(|_| { + json!({ + "pages": [{"index": 0, "markdown": content}], + "model": model, + "usage_info": usage.unwrap_or_else(|| json!({})), + }) + }) + } else { + json!({ + "pages": [{"index": 0, "markdown": content}], + "model": model, + "usage_info": usage.unwrap_or_else(|| json!({})), + }) + } + } + Value::Object(_) => content, + other => json!({ + "pages": [{"index": 0, "markdown": other.to_string()}], + "model": model, + "usage_info": usage.unwrap_or_else(|| json!({})), + }), + } +} + +impl OcrProviderConfig for VertexAiOcrConfig { + fn supported_ocr_params(&self) -> &'static [&'static str] { + MISTRAL_OCR_CONFIG.supported_ocr_params() + } + + fn transform_ocr_request( + &self, + model: &str, + document: Value, + optional_params: Map, + ) -> CoreResult { + MISTRAL_OCR_CONFIG.transform_ocr_request(model, document, optional_params) + } + + fn transform_ocr_response( + &self, + model: &str, + response_json: Value, + ) -> CoreResult { + MISTRAL_OCR_CONFIG.transform_ocr_response(model, response_json) + } + + fn complete_url( + &self, + api_base: Option<&str>, + model: &str, + optional_params: &Map, + env_lookup: &dyn Fn(&str) -> Option, + ) -> CoreResult { + complete_vertex_mistral_url(api_base, model, optional_params, env_lookup) + } + + fn resolve_api_key( + &self, + api_key: Option<&str>, + env_lookup: &dyn Fn(&str) -> Option, + ) -> CoreResult { + resolve_vertex_api_key(api_key, env_lookup) + } + + fn requires_data_uri_document(&self) -> bool { + true + } +} + +impl OcrProviderConfig for VertexAiDeepSeekOcrConfig { + fn supported_ocr_params(&self) -> &'static [&'static str] { + DEEPSEEK_SUPPORTED_OCR_PARAMS + } + + fn transform_ocr_request( + &self, + model: &str, + document: Value, + optional_params: Map, + ) -> CoreResult { + let mut data = Map::new(); + data.insert( + "model".to_string(), + Value::String(deepseek_model_name(model)), + ); + data.insert( + "messages".to_string(), + json!([{"role": "user", "content": [document_content_item(&document)?]}]), + ); + for (key, value) in optional_params { + if DEEPSEEK_SUPPORTED_OCR_PARAMS.contains(&key.as_str()) { + data.insert(key, value); + } + } + Ok(OcrRequestData { + data: Value::Object(data), + files: None, + }) + } + + fn transform_ocr_response( + &self, + model: &str, + response_json: Value, + ) -> CoreResult { + let response = response_json + .as_object() + .ok_or_else(|| CoreError::InvalidType { + expected: "object", + actual: json_type_name(&response_json), + })?; + let usage = response.get("usage").cloned(); + let content = first_choice_content(&response_json)?; + let mut ocr_data = ocr_data_from_content(content.clone(), usage.clone(), model); + + if !ocr_data.get("pages").is_some_and(Value::is_array) { + ocr_data = json!({ + "pages": [{ + "index": 0, + "markdown": match content { + Value::String(value) => value, + other => other.to_string(), + } + }], + "model": ocr_data.get("model").and_then(Value::as_str).unwrap_or(model), + "usage_info": ocr_data.get("usage_info").cloned().or(usage).unwrap_or_else(|| json!({})), + }); + } + + let object = ocr_data.as_object().ok_or_else(|| CoreError::InvalidType { + expected: "object", + actual: json_type_name(&ocr_data), + })?; + let pages = object + .get("pages") + .and_then(Value::as_array) + .cloned() + .unwrap_or_default(); + let usage_info = object + .get("usage_info") + .cloned() + .or_else(|| response.get("usage").cloned()); + Ok(OcrResponseData { + pages, + model: object + .get("model") + .and_then(Value::as_str) + .unwrap_or(model) + .to_string(), + document_annotation: object.get("document_annotation").cloned(), + usage_info, + object: "ocr".to_string(), + }) + } + + fn complete_url( + &self, + api_base: Option<&str>, + _model: &str, + optional_params: &Map, + env_lookup: &dyn Fn(&str) -> Option, + ) -> CoreResult { + complete_vertex_deepseek_url(api_base, optional_params, env_lookup) + } + + fn resolve_api_key( + &self, + api_key: Option<&str>, + env_lookup: &dyn Fn(&str) -> Option, + ) -> CoreResult { + resolve_vertex_api_key(api_key, env_lookup) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn vertex_mistral_url_uses_project_location_and_model() { + let params = Map::from_iter([ + ("vertex_project".to_string(), json!("proj-1")), + ("vertex_location".to_string(), json!("europe-west4")), + ]); + + let url = complete_vertex_mistral_url(None, "mistral-ocr-maas", ¶ms, &|_| None) + .expect("url builds"); + + assert_eq!( + url, + "https://europe-west4-aiplatform.googleapis.com/v1/projects/proj-1/locations/europe-west4/publishers/mistralai/models/mistral-ocr-maas:rawPredict" + ); + } + + #[test] + fn vertex_mistral_reuses_mistral_body_transform() { + let body = VERTEX_AI_OCR_CONFIG + .transform_ocr_request( + "mistral-ocr-maas", + json!({"type": "image_url", "image_url": "data:image/png;base64,abc"}), + Map::new(), + ) + .expect("request transforms") + .data; + + assert_eq!(body["model"], "mistral-ocr-maas"); + assert_eq!(body["document"]["image_url"], "data:image/png;base64,abc"); + } + + #[test] + fn vertex_deepseek_request_uses_ocr_endpoint_shape() { + let body = VERTEX_AI_DEEPSEEK_OCR_CONFIG + .transform_ocr_request( + "deepseek-ocr-maas", + json!({"type": "document_url", "document_url": "gs://bucket/doc.pdf"}), + Map::from_iter([("temperature".to_string(), json!(0.1))]), + ) + .expect("request transforms") + .data; + + assert_eq!(body["model"], "deepseek-ai/deepseek-ocr-maas"); + assert_eq!(body["temperature"], 0.1); + assert_eq!( + body["messages"][0]["content"][0], + json!({"type": "image_url", "image_url": "gs://bucket/doc.pdf"}) + ); + } + + #[test] + fn vertex_deepseek_response_wraps_markdown_content() { + let response = VERTEX_AI_DEEPSEEK_OCR_CONFIG + .transform_ocr_response( + "deepseek-ocr-maas", + json!({ + "choices": [{"message": {"content": "# OCR text"}}], + "usage": {"prompt_tokens": 1} + }), + ) + .expect("response transforms"); + + assert_eq!( + response.pages, + vec![json!({"index": 0, "markdown": "# OCR text"})] + ); + assert_eq!(response.model, "deepseek-ocr-maas"); + assert_eq!(response.usage_info, Some(json!({"prompt_tokens": 1}))); + } +} diff --git a/litellm-rust/crates/python-bridge/Cargo.toml b/litellm-rust/crates/python-bridge/Cargo.toml index 4db32818604..83e163c38f1 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-ai-gateway = { workspace = true, default-features = false } 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 8c8416b6bd3..82024e1bf47 100644 --- a/litellm-rust/crates/python-bridge/src/lib.rs +++ b/litellm-rust/crates/python-bridge/src/lib.rs @@ -1,6 +1,6 @@ use std::time::Duration; -use litellm_ai_gateway::io::ocr::run_ocr; +use litellm_ai_gateway::io::ocr::{ocr as run_ocr, OcrRequest}; use litellm_core::error::CoreError; use pyo3::exceptions::{PyRuntimeError, PyValueError}; use pyo3::prelude::*; @@ -9,6 +9,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()?; @@ -22,59 +29,93 @@ fn json_to_py(py: Python<'_>, value: Value) -> PyResult> { Ok(json.call_method1("loads", (encoded,))?.unbind()) } -/// Map a core error to the closest Python exception. Caller-input problems -/// (auth, bad types, missing fields) -> `ValueError`; everything else -/// (network, upstream status, parse failures) -> `RuntimeError`. 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::InvalidRequest(_) + | CoreError::InvalidType { .. } + | CoreError::MissingField(_) => PyValueError::new_err(err.to_string()), other => PyRuntimeError::new_err(other.to_string()), } } -/// Perform a Mistral OCR call end to end and return the response as a dict. +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)) +} + #[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). let result = gil::release_gil(py, || { - run_ocr( - &model, + pyo3_async_runtimes::tokio::get_runtime().block_on(run_ocr(OcrRequest { + model: &model, document, - api_key.as_deref(), - api_base.as_deref(), + 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 +124,47 @@ fn ocr( } } -/// Bridge GIL accounting, e.g. `{"releases": 12}`. Lets the Python side observe -/// how often the bridge has dropped the GIL for blocking work. +#[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 = run_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)) + }) +} + #[pyfunction] fn gil_stats(py: Python<'_>) -> PyResult> { let stats = PyDict::new(py); @@ -95,6 +175,7 @@ fn gil_stats(py: Python<'_>) -> PyResult> { #[pymodule] fn _native(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/llms/azure_ai/ocr/document_intelligence/transformation.py b/litellm/llms/azure_ai/ocr/document_intelligence/transformation.py index d4144a75718..cc65ad706ab 100644 --- a/litellm/llms/azure_ai/ocr/document_intelligence/transformation.py +++ b/litellm/llms/azure_ai/ocr/document_intelligence/transformation.py @@ -11,7 +11,7 @@ The operation location must be polled until the analysis completes. import asyncio import re import time -from typing import Any, Dict, Optional +from typing import Any, Dict from urllib.parse import quote import httpx @@ -35,6 +35,8 @@ from litellm.llms.base_llm.ocr.transformation import ( ) from litellm.secret_managers.main import get_secret_str +AZURE_DOCUMENT_INTELLIGENCE_API_KEY_ENV_VAR = "AZURE_DOCUMENT_INTELLIGENCE_API_KEY" + class AzureDocumentIntelligenceOCRConfig(BaseOCRConfig): """ @@ -54,6 +56,9 @@ class AzureDocumentIntelligenceOCRConfig(BaseOCRConfig): def __init__(self) -> None: super().__init__() + def get_api_key_env_var(self) -> str | None: + return AZURE_DOCUMENT_INTELLIGENCE_API_KEY_ENV_VAR + def get_supported_ocr_params(self, model: str) -> list: """ Get supported OCR parameters for Azure Document Intelligence. @@ -144,9 +149,9 @@ class AzureDocumentIntelligenceOCRConfig(BaseOCRConfig): self, headers: Dict, model: str, - api_key: Optional[str] = None, - api_base: Optional[str] = None, - litellm_params: Optional[dict] = None, + api_key: str | None = None, + api_base: str | None = None, + litellm_params: dict | None = None, **kwargs, ) -> Dict: """ @@ -156,7 +161,7 @@ class AzureDocumentIntelligenceOCRConfig(BaseOCRConfig): """ # Get API key from environment if not provided if api_key is None: - api_key = get_secret_str("AZURE_DOCUMENT_INTELLIGENCE_API_KEY") + api_key = get_secret_str(AZURE_DOCUMENT_INTELLIGENCE_API_KEY_ENV_VAR) if api_key is None: raise ValueError( @@ -182,10 +187,10 @@ class AzureDocumentIntelligenceOCRConfig(BaseOCRConfig): def get_complete_url( self, - api_base: Optional[str], + api_base: str | None, model: str, optional_params: dict, - litellm_params: Optional[dict] = None, + litellm_params: dict | None = None, **kwargs, ) -> str: """ diff --git a/litellm/llms/azure_ai/ocr/transformation.py b/litellm/llms/azure_ai/ocr/transformation.py index f661ddb9ebc..ee35fc28994 100644 --- a/litellm/llms/azure_ai/ocr/transformation.py +++ b/litellm/llms/azure_ai/ocr/transformation.py @@ -2,7 +2,7 @@ Azure AI OCR transformation implementation. """ -from typing import Dict, Optional +from typing import Dict from litellm._logging import verbose_logger from litellm.litellm_core_utils.prompt_templates.image_handling import ( @@ -13,6 +13,8 @@ from litellm.llms.base_llm.ocr.transformation import DocumentType, OCRRequestDat from litellm.llms.mistral.ocr.transformation import MistralOCRConfig from litellm.secret_managers.main import get_secret_str +AZURE_AI_OCR_API_KEY_ENV_VAR = "AZURE_AI_API_KEY" + class AzureAIOCRConfig(MistralOCRConfig): """ @@ -30,13 +32,16 @@ class AzureAIOCRConfig(MistralOCRConfig): def __init__(self) -> None: super().__init__() + def get_api_key_env_var(self) -> str | None: + return AZURE_AI_OCR_API_KEY_ENV_VAR + def validate_environment( self, headers: Dict, model: str, - api_key: Optional[str] = None, - api_base: Optional[str] = None, - litellm_params: Optional[dict] = None, + api_key: str | None = None, + api_base: str | None = None, + litellm_params: dict | None = None, **kwargs, ) -> Dict: """ @@ -46,7 +51,7 @@ class AzureAIOCRConfig(MistralOCRConfig): """ # Get API key from environment if not provided if api_key is None: - api_key = get_secret_str("AZURE_AI_API_KEY") + api_key = get_secret_str(AZURE_AI_OCR_API_KEY_ENV_VAR) if api_key is None: raise ValueError( @@ -72,10 +77,10 @@ class AzureAIOCRConfig(MistralOCRConfig): def get_complete_url( self, - api_base: Optional[str], + api_base: str | None, model: str, optional_params: dict, - litellm_params: Optional[dict] = None, + litellm_params: dict | None = None, **kwargs, ) -> str: """ diff --git a/litellm/llms/base_llm/ocr/transformation.py b/litellm/llms/base_llm/ocr/transformation.py index 263e0c094ce..a2946c62506 100644 --- a/litellm/llms/base_llm/ocr/transformation.py +++ b/litellm/llms/base_llm/ocr/transformation.py @@ -2,7 +2,7 @@ Base OCR transformation configuration. """ -from typing import TYPE_CHECKING, Any, Dict, List, Optional, Union +from typing import TYPE_CHECKING, Any, Dict, List, Union import httpx from pydantic import PrivateAttr @@ -25,16 +25,16 @@ DocumentType = Dict[str, str] class OCRPageDimensions(LiteLLMPydanticObjectBase): """Page dimensions from OCR response.""" - dpi: Optional[int] = None - height: Optional[int] = None - width: Optional[int] = None + dpi: int | None = None + height: int | None = None + width: int | None = None class OCRPageImage(LiteLLMPydanticObjectBase): """Image extracted from OCR page.""" - image_base64: Optional[str] = None - bbox: Optional[Dict[str, Any]] = None + image_base64: str | None = None + bbox: Dict[str, Any] | None = None model_config = {"extra": "allow"} @@ -44,8 +44,8 @@ class OCRPage(LiteLLMPydanticObjectBase): index: int markdown: str - images: Optional[List[OCRPageImage]] = None - dimensions: Optional[OCRPageDimensions] = None + images: List[OCRPageImage] | None = None + dimensions: OCRPageDimensions | None = None model_config = {"extra": "allow"} @@ -53,9 +53,9 @@ class OCRPage(LiteLLMPydanticObjectBase): class OCRUsageInfo(LiteLLMPydanticObjectBase): """Usage information from OCR response.""" - pages_processed: Optional[int] = None - credits: Optional[float] = None - doc_size_bytes: Optional[int] = None + pages_processed: int | None = None + credits: float | None = None + doc_size_bytes: int | None = None model_config = {"extra": "allow"} @@ -68,8 +68,8 @@ class OCRResponse(LiteLLMPydanticObjectBase): pages: List[OCRPage] model: str - document_annotation: Optional[Any] = None - usage_info: Optional[OCRUsageInfo] = None + document_annotation: Any | None = None + usage_info: OCRUsageInfo | None = None object: str = "ocr" model_config = {"extra": "allow"} @@ -81,8 +81,8 @@ class OCRResponse(LiteLLMPydanticObjectBase): class OCRRequestData(LiteLLMPydanticObjectBase): """OCR request data structure.""" - data: Optional[Union[Dict, bytes]] = None - files: Optional[Dict[str, Any]] = None + data: Union[Dict, bytes] | None = None + files: Dict[str, Any] | None = None class BaseOCRConfig: @@ -101,6 +101,12 @@ class BaseOCRConfig: """ return [] + def get_api_key_env_var(self) -> str | None: + """ + Return the provider-specific API key environment variable name, if any. + """ + return None + def map_ocr_params( self, non_default_params: dict, @@ -114,9 +120,9 @@ class BaseOCRConfig: self, headers: Dict, model: str, - api_key: Optional[str] = None, - api_base: Optional[str] = None, - litellm_params: Optional[dict] = None, + api_key: str | None = None, + api_base: str | None = None, + litellm_params: dict | None = None, **kwargs, ) -> Dict: """ @@ -127,10 +133,10 @@ class BaseOCRConfig: def get_complete_url( self, - api_base: Optional[str], + api_base: str | None, model: str, optional_params: dict, - litellm_params: Optional[dict] = None, + litellm_params: dict | None = None, **kwargs, ) -> str: """ diff --git a/litellm/llms/mistral/ocr/transformation.py b/litellm/llms/mistral/ocr/transformation.py index 21e0e27a314..3c0460cd51e 100644 --- a/litellm/llms/mistral/ocr/transformation.py +++ b/litellm/llms/mistral/ocr/transformation.py @@ -2,7 +2,7 @@ Mistral OCR transformation implementation. """ -from typing import Any, Dict, Optional +from typing import Any, Dict import httpx @@ -15,6 +15,8 @@ from litellm.llms.base_llm.ocr.transformation import ( ) from litellm.secret_managers.main import get_secret_str +MISTRAL_OCR_API_KEY_ENV_VAR = "MISTRAL_API_KEY" + class MistralOCRConfig(BaseOCRConfig): """ @@ -59,6 +61,9 @@ class MistralOCRConfig(BaseOCRConfig): "id", ] + def get_api_key_env_var(self) -> str | None: + return MISTRAL_OCR_API_KEY_ENV_VAR + def map_ocr_params( self, non_default_params: dict, @@ -85,9 +90,9 @@ class MistralOCRConfig(BaseOCRConfig): self, headers: Dict, model: str, - api_key: Optional[str] = None, - api_base: Optional[str] = None, - litellm_params: Optional[dict] = None, + api_key: str | None = None, + api_base: str | None = None, + litellm_params: dict | None = None, **kwargs, ) -> Dict: """ @@ -95,7 +100,7 @@ class MistralOCRConfig(BaseOCRConfig): """ # Get API key from environment if not provided if api_key is None: - api_key = get_secret_str("MISTRAL_API_KEY") + api_key = get_secret_str(MISTRAL_OCR_API_KEY_ENV_VAR) if api_key is None: raise ValueError( @@ -113,10 +118,10 @@ class MistralOCRConfig(BaseOCRConfig): def get_complete_url( self, - api_base: Optional[str], + api_base: str | None, model: str, optional_params: dict, - litellm_params: Optional[dict] = None, + litellm_params: dict | None = None, **kwargs, ) -> str: """ diff --git a/litellm/llms/vertex_ai/ocr/deepseek_transformation.py b/litellm/llms/vertex_ai/ocr/deepseek_transformation.py index 516ee03ba55..a98311d04eb 100644 --- a/litellm/llms/vertex_ai/ocr/deepseek_transformation.py +++ b/litellm/llms/vertex_ai/ocr/deepseek_transformation.py @@ -3,7 +3,7 @@ Vertex AI DeepSeek OCR transformation implementation. """ import json -from typing import TYPE_CHECKING, Any, Dict, Optional +from typing import TYPE_CHECKING, Any, Dict import httpx @@ -18,6 +18,8 @@ from litellm.llms.base_llm.ocr.transformation import ( ) from litellm.llms.vertex_ai.vertex_llm_base import VertexBase +VERTEX_AI_DEEPSEEK_OCR_API_KEY_ENV_VAR = "VERTEX_AI_API_KEY" + if TYPE_CHECKING: from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj else: @@ -28,21 +30,24 @@ class VertexAIDeepSeekOCRConfig(BaseOCRConfig): """ Vertex AI DeepSeek OCR transformation configuration. - Vertex AI DeepSeek OCR uses the chat completion API format through the openapi endpoint. - This transformation converts OCR requests to chat completion format and vice versa. + This transformation converts standard LiteLLM OCR requests to the + Vertex AI DeepSeek OCR OpenAPI endpoint shape and normalizes the response. """ def __init__(self) -> None: super().__init__() self.vertex_base = VertexBase() + def get_api_key_env_var(self) -> str | None: + return VERTEX_AI_DEEPSEEK_OCR_API_KEY_ENV_VAR + def validate_environment( self, headers: Dict, model: str, - api_key: Optional[str] = None, - api_base: Optional[str] = None, - litellm_params: Optional[dict] = None, + api_key: str | None = None, + api_base: str | None = None, + litellm_params: dict | None = None, **kwargs, ) -> Dict: """ @@ -50,6 +55,13 @@ class VertexAIDeepSeekOCRConfig(BaseOCRConfig): Vertex AI uses Bearer token authentication with access token from credentials. """ + if api_key is not None: + return { + "Authorization": f"Bearer {api_key}", + "Content-Type": "application/json", + **headers, + } + # Extract Vertex AI parameters using safe helpers from VertexBase # Use safe_get_* methods that don't mutate litellm_params dict litellm_params = litellm_params or {} @@ -77,18 +89,15 @@ class VertexAIDeepSeekOCRConfig(BaseOCRConfig): def get_complete_url( self, - api_base: Optional[str], + api_base: str | None, model: str, optional_params: dict, - litellm_params: Optional[dict] = None, + litellm_params: dict | None = None, **kwargs, ) -> str: """ Get complete URL for Vertex AI DeepSeek OCR endpoint. - Vertex AI endpoint format: - https://{location}-aiplatform.googleapis.com/v1/projects/{project}/locations/{location}/endpoints/openapi/chat/completions - Args: api_base: Vertex AI API base URL (optional) model: Model name (e.g., "deepseek-ai/deepseek-ocr-maas") @@ -123,8 +132,6 @@ class VertexAIDeepSeekOCRConfig(BaseOCRConfig): # Ensure no trailing slash api_base = api_base.rstrip("/") - # Vertex AI DeepSeek OCR endpoint format - # Format: https://{region}-aiplatform.googleapis.com/v1/projects/{project}/locations/{region}/endpoints/openapi/chat/completions return f"{api_base}/v1/projects/{vertex_project}/locations/{vertex_location}/endpoints/openapi/chat/completions" def transform_ocr_request( @@ -136,9 +143,9 @@ class VertexAIDeepSeekOCRConfig(BaseOCRConfig): **kwargs, ) -> OCRRequestData: """ - Transform OCR request to chat completion format for Vertex AI DeepSeek OCR. + Transform OCR request for Vertex AI DeepSeek OCR. - Converts OCR document format to chat completion messages format: + Converts OCR document format to the Vertex AI DeepSeek OCR payload: - Input: {"type": "image_url", "image_url": "gs://..."} - Output: {"model": "deepseek-ai/deepseek-ocr-maas", "messages": [{"role": "user", "content": [{"type": "image_url", "image_url": "gs://..."}]}]} @@ -150,7 +157,7 @@ class VertexAIDeepSeekOCRConfig(BaseOCRConfig): **kwargs: Additional arguments Returns: - OCRRequestData with JSON data in chat completion format + OCRRequestData with JSON data for the DeepSeek OCR endpoint """ verbose_logger.debug( "Vertex AI DeepSeek OCR transform_ocr_request (sync) called" @@ -173,7 +180,7 @@ class VertexAIDeepSeekOCRConfig(BaseOCRConfig): f"Unsupported document type: {doc_type}. Expected 'image_url' or 'document_url'" ) - # Build chat completion message content + # Build DeepSeek OCR message content content_item = {} if image_url: content_item = {"type": "image_url", "image_url": image_url} @@ -181,25 +188,21 @@ class VertexAIDeepSeekOCRConfig(BaseOCRConfig): # For document URLs, we use image_url type as well (Vertex AI supports both) content_item = {"type": "image_url", "image_url": document_url} - # Build chat completion request + # Build DeepSeek OCR request data = { "model": "deepseek-ai/" + model, "messages": [{"role": "user", "content": [content_item]}], } # Add optional parameters (stream, temperature, etc.) - # Filter out OCR-specific params that don't apply to chat completion - chat_completion_params = {} + deepseek_ocr_params = {} for key, value in optional_params.items(): - # Include common chat completion params if key in ["stream", "temperature", "max_tokens", "top_p", "n", "stop"]: - chat_completion_params[key] = value + deepseek_ocr_params[key] = value - data.update(chat_completion_params) + data.update(deepseek_ocr_params) - verbose_logger.debug( - "Vertex AI DeepSeek OCR: Transformed request to chat completion format" - ) + verbose_logger.debug("Vertex AI DeepSeek OCR: Transformed request") return OCRRequestData(data=data, files=None) @@ -212,7 +215,7 @@ class VertexAIDeepSeekOCRConfig(BaseOCRConfig): **kwargs, ) -> OCRRequestData: """ - Transform OCR request to chat completion format for Vertex AI DeepSeek OCR (async). + Transform OCR request for Vertex AI DeepSeek OCR (async). Same as sync version - no async-specific logic needed. @@ -224,7 +227,7 @@ class VertexAIDeepSeekOCRConfig(BaseOCRConfig): **kwargs: Additional arguments Returns: - OCRRequestData with JSON data in chat completion format + OCRRequestData with JSON data for the DeepSeek OCR endpoint """ return self.transform_ocr_request( model=model, @@ -242,12 +245,11 @@ class VertexAIDeepSeekOCRConfig(BaseOCRConfig): **kwargs, ) -> OCRResponse: """ - Transform chat completion response to OCR format. + Transform Vertex AI DeepSeek OCR response to OCR format. - Vertex AI DeepSeek OCR returns chat completion format: + Vertex AI DeepSeek OCR returns an OpenAPI response: { "id": "...", - "object": "chat.completion", "choices": [{ "message": { "role": "assistant", @@ -274,16 +276,16 @@ class VertexAIDeepSeekOCRConfig(BaseOCRConfig): try: response_json = raw_response.json() - # Extract content from chat completion response + # Extract OCR content from provider response choices = response_json.get("choices", []) if not choices: - raise ValueError("No choices in chat completion response") + raise ValueError("No choices in DeepSeek OCR response") message = choices[0].get("message", {}) content = message.get("content", "") if not content: - raise ValueError("No content in chat completion response") + raise ValueError("No content in DeepSeek OCR response") # Try to parse content as JSON (OCR result might be JSON string) ocr_data = None @@ -376,7 +378,7 @@ class VertexAIDeepSeekOCRConfig(BaseOCRConfig): **kwargs, ) -> OCRResponse: """ - Async transform chat completion response to OCR format. + Async transform Vertex AI DeepSeek OCR response to OCR format. Same as sync version - no async-specific logic needed. diff --git a/litellm/llms/vertex_ai/ocr/transformation.py b/litellm/llms/vertex_ai/ocr/transformation.py index cbf15803132..a725762b3c5 100644 --- a/litellm/llms/vertex_ai/ocr/transformation.py +++ b/litellm/llms/vertex_ai/ocr/transformation.py @@ -2,7 +2,7 @@ Vertex AI Mistral OCR transformation implementation. """ -from typing import Dict, Optional +from typing import Dict from litellm._logging import verbose_logger from litellm.litellm_core_utils.prompt_templates.image_handling import ( @@ -14,6 +14,8 @@ from litellm.llms.mistral.ocr.transformation import MistralOCRConfig from litellm.llms.vertex_ai.common_utils import get_vertex_base_url from litellm.llms.vertex_ai.vertex_llm_base import VertexBase +VERTEX_AI_OCR_API_KEY_ENV_VAR = "VERTEX_AI_API_KEY" + class VertexAIOCRConfig(MistralOCRConfig): """ @@ -32,13 +34,16 @@ class VertexAIOCRConfig(MistralOCRConfig): super().__init__() self.vertex_base = VertexBase() + def get_api_key_env_var(self) -> str | None: + return VERTEX_AI_OCR_API_KEY_ENV_VAR + def validate_environment( self, headers: Dict, model: str, - api_key: Optional[str] = None, - api_base: Optional[str] = None, - litellm_params: Optional[dict] = None, + api_key: str | None = None, + api_base: str | None = None, + litellm_params: dict | None = None, **kwargs, ) -> Dict: """ @@ -46,6 +51,13 @@ class VertexAIOCRConfig(MistralOCRConfig): Vertex AI uses Bearer token authentication with access token from credentials. """ + if api_key is not None: + return { + "Authorization": f"Bearer {api_key}", + "Content-Type": "application/json", + **headers, + } + # Extract Vertex AI parameters using safe helpers from VertexBase # Use safe_get_* methods that don't mutate litellm_params dict litellm_params = litellm_params or {} @@ -73,10 +85,10 @@ class VertexAIOCRConfig(MistralOCRConfig): def get_complete_url( self, - api_base: Optional[str], + api_base: str | None, model: str, optional_params: dict, - litellm_params: Optional[dict] = None, + litellm_params: dict | None = None, **kwargs, ) -> str: """ diff --git a/litellm/ocr/main.py b/litellm/ocr/main.py index 3a9ef8db804..6a196d41768 100644 --- a/litellm/ocr/main.py +++ b/litellm/ocr/main.py @@ -4,13 +4,12 @@ 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 +from typing import Any, Callable, Coroutine, Union, cast import httpx @@ -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,9 +34,40 @@ base_llm_http_handler = BaseLLMHTTPHandler() ################################################# +@dataclass +class _PreparedOCRRequest: + model: str + document: dict[str, Any] + api_key: str | None + api_base: str | None + custom_llm_provider: str + extra_headers: dict[str, object] | None + provider_config: BaseOCRConfig + optional_params: dict[str, object] + litellm_params: dict[str, object] + effective_timeout: Union[float, httpx.Timeout] + litellm_logging_obj: LiteLLMLoggingObj + + +@dataclass +class _PreparedRustOCRCall: + api_key: str | None + api_base: str | None + headers: dict[str, object] + optional_params: dict[str, object] + + +_RUST_OCR_PROVIDERS = { + "mistral", + "azure_ai", + "azure_ai/doc-intelligence", + "vertex_ai", +} + + def _timeout_to_seconds( - timeout: Optional[Union[float, httpx.Timeout]], -) -> Optional[float]: + timeout: Union[float, httpx.Timeout] | None, +) -> float | None: """Convert the Python OCR timeout to a single seconds value for the Rust bridge. The Rust HTTP client takes one duration; ``httpx.Timeout`` carries separate @@ -45,18 +81,206 @@ def _timeout_to_seconds( return float(timeout) +def _prepare_ocr_request( + model: str, + document: dict[str, Any], + api_key: str | None, + api_base: str | None, + timeout: Union[float, httpx.Timeout] | None, + custom_llm_provider: str | None, + extra_headers: dict[str, Any] | None, + kwargs: dict[str, Any], +) -> _PreparedOCRRequest: + litellm_logging_obj = cast(LiteLLMLoggingObj, kwargs.pop("litellm_logging_obj")) + litellm_call_id = cast(str | None, 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=document, + api_key=api_key, + api_base=api_base, + custom_llm_provider=custom_llm_provider, + extra_headers=cast(dict[str, object] | None, 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 _rust_ocr_supported(prepared_request: _PreparedOCRRequest) -> bool: + return prepared_request.custom_llm_provider in _RUST_OCR_PROVIDERS + + +def _rust_bridge_optional_params( + prepared_request: _PreparedOCRRequest, + resolve_secret: Callable[[str], str | None], +) -> dict[str, object]: + optional_params = dict(prepared_request.optional_params) + if prepared_request.custom_llm_provider == "vertex_ai": + vertex_project = ( + prepared_request.litellm_params.get("vertex_project") + or prepared_request.litellm_params.get("vertex_ai_project") + or litellm.vertex_project + or resolve_secret("VERTEXAI_PROJECT") + ) + vertex_location = ( + prepared_request.litellm_params.get("vertex_location") + or prepared_request.litellm_params.get("vertex_ai_location") + or litellm.vertex_location + or resolve_secret("VERTEXAI_LOCATION") + or resolve_secret("VERTEX_LOCATION") + ) + if vertex_project is not None: + optional_params["vertex_project"] = vertex_project + if vertex_location is not None: + optional_params["vertex_location"] = vertex_location + return optional_params + + +def _rust_bridge_api_base( + prepared_request: _PreparedOCRRequest, + resolve_secret: Callable[[str], str | None], +) -> str | None: + if prepared_request.api_base is not None: + return prepared_request.api_base + if prepared_request.custom_llm_provider == "azure_ai/doc-intelligence": + return resolve_secret("AZURE_DOCUMENT_INTELLIGENCE_ENDPOINT") + if prepared_request.custom_llm_provider == "azure_ai": + if ( + "doc-intelligence" in prepared_request.model + or "documentintelligence" in prepared_request.model + ): + return resolve_secret("AZURE_DOCUMENT_INTELLIGENCE_ENDPOINT") + return resolve_secret("AZURE_AI_API_BASE") + return None + + +def _prepare_rust_ocr_call( + prepared_request: _PreparedOCRRequest, + resolve_api_key: Callable[[str], str | None], +) -> _PreparedRustOCRCall: + provider_config = prepared_request.provider_config + api_key_env_var = provider_config.get_api_key_env_var() + resolved_api_key = prepared_request.api_key or ( + resolve_api_key(api_key_env_var) if api_key_env_var is not None else None + ) + resolved_headers = provider_config.validate_environment( + headers=prepared_request.extra_headers or {}, + model=prepared_request.model, + api_key=resolved_api_key, + api_base=prepared_request.api_base, + litellm_params=prepared_request.litellm_params, + ) + resolved_complete_url = provider_config.get_complete_url( + api_base=prepared_request.api_base, + model=prepared_request.model, + optional_params=prepared_request.optional_params, + litellm_params=prepared_request.litellm_params, + ) + rust_api_base = _rust_bridge_api_base(prepared_request, resolve_api_key) + rust_optional_params = _rust_bridge_optional_params( + prepared_request, resolve_api_key + ) + prepared_request.litellm_logging_obj.pre_call( + input="OCR document processing", + api_key=resolved_api_key, + additional_args={ + "complete_input_dict": { + "model": prepared_request.model, + "document": prepared_request.document, + **rust_optional_params, + }, + "api_base": resolved_complete_url, + "headers": resolved_headers, + }, + ) + return _PreparedRustOCRCall( + api_key=resolved_api_key, + api_base=rust_api_base, + headers=cast(dict[str, object], resolved_headers), + optional_params=rust_optional_params, + ) + + def _run_rust_ocr( rust_ocr: RustOcr, - 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], - optional_params: dict[str, object], - litellm_params: dict[str, object], - timeout_seconds: Optional[float], + prepared_request: _PreparedOCRRequest, + resolve_api_key: Callable[[str], str | None], ) -> OCRResponse: """Run the Mistral OCR call through the Rust bridge and wrap the result. @@ -66,41 +290,43 @@ def _run_rust_ocr( headers) is mirrored into pre_call so logs match the wire. Dependencies are injected so this stays unit-testable without patching module globals. """ - resolved_api_key = api_key or resolve_api_key("MISTRAL_API_KEY") - resolved_headers = provider_config.validate_environment( - headers={}, - 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, - }, + prepared = _prepare_rust_ocr_call( + prepared_request=prepared_request, + resolve_api_key=resolve_api_key, ) return OCRResponse.model_validate( rust_ocr( - model=model, - document=document, - api_key=resolved_api_key, - api_base=api_base, - optional_params=optional_params, - timeout_seconds=timeout_seconds, + model=prepared_request.model, + document=cast(dict[str, object], prepared_request.document), + api_key=prepared.api_key, + api_base=prepared.api_base, + custom_llm_provider=prepared_request.custom_llm_provider, + extra_headers=prepared.headers, + optional_params=prepared.optional_params, + timeout_seconds=_timeout_to_seconds(prepared_request.effective_timeout), + ) + ) + + +async def _run_rust_aocr( + rust_aocr: RustAocr, + prepared_request: _PreparedOCRRequest, + resolve_api_key: Callable[[str], str | None], +) -> OCRResponse: + prepared = _prepare_rust_ocr_call( + prepared_request=prepared_request, + resolve_api_key=resolve_api_key, + ) + return OCRResponse.model_validate( + await rust_aocr( + model=prepared_request.model, + document=cast(dict[str, object], prepared_request.document), + api_key=prepared.api_key, + api_base=prepared.api_base, + custom_llm_provider=prepared_request.custom_llm_provider, + extra_headers=prepared.headers, + optional_params=prepared.optional_params, + timeout_seconds=_timeout_to_seconds(prepared_request.effective_timeout), ) ) @@ -108,12 +334,12 @@ def _run_rust_ocr( @client async def aocr( model: str, - document: Dict[str, Any], - api_key: Optional[str] = None, - api_base: Optional[str] = None, - timeout: Optional[Union[float, httpx.Timeout]] = None, - custom_llm_provider: Optional[str] = None, - extra_headers: Optional[Dict[str, Any]] = None, + document: dict[str, Any], + api_key: str | None = None, + api_base: str | None = None, + timeout: Union[float, httpx.Timeout] | None = None, + custom_llm_provider: str | None = None, + extra_headers: dict[str, Any] | None = None, **kwargs, ) -> OCRResponse: """ @@ -174,19 +400,18 @@ async def aocr( ) ``` """ - local_vars = locals() + completion_kwargs: dict[str, object] = { + "model": model, + "document": document, + "api_key": api_key, + "api_base": api_base, + "timeout": timeout, + "custom_llm_provider": custom_llm_provider, + "extra_headers": extra_headers, + "kwargs": kwargs, + } 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 +419,47 @@ async def aocr( timeout=timeout, custom_llm_provider=custom_llm_provider, extra_headers=extra_headers, - **kwargs, + kwargs=kwargs, + ) + model = prepared.model + custom_llm_provider = prepared.custom_llm_provider + completion_kwargs.update( + {"model": model, "custom_llm_provider": custom_llm_provider} ) - ctx = contextvars.copy_context() - func_with_context = partial(ctx.run, func) - init_response = await loop.run_in_executor(None, func_with_context) + if _rust_ocr_supported(prepared) 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, + prepared_request=prepared, + resolve_api_key=get_secret_str, + ) + 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( @@ -217,20 +472,145 @@ async def aocr( model=model, custom_llm_provider=custom_llm_provider, original_exception=e, - completion_kwargs=local_vars, + completion_kwargs=completion_kwargs, extra_kwargs=kwargs, ) +################################################# +# Public utilities — used by the SDK and the proxy +################################################# + +_MIME_PATTERN = re.compile(r"^[\w.+-]+/[\w.+-]+$") + +_MIME_TYPE_MAP = { + ".pdf": "application/pdf", + ".png": "image/png", + ".jpg": "image/jpeg", + ".jpeg": "image/jpeg", + ".gif": "image/gif", + ".webp": "image/webp", + ".tiff": "image/tiff", + ".tif": "image/tiff", + ".bmp": "image/bmp", +} + + +def get_mime_type(file_path: str) -> str: + """ + Determine MIME type from file path extension. + + Falls back to mimetypes.guess_type, then to 'application/octet-stream'. + """ + ext = os.path.splitext(file_path)[1].lower() + mime = _MIME_TYPE_MAP.get(ext) + if mime: + return mime + guessed, _ = mimetypes.guess_type(file_path) + return guessed or "application/octet-stream" + + +def convert_file_document_to_url_document(document: dict[str, Any]) -> dict[str, str]: + """ + Convert a file-type document dict to a document_url-type document dict + with an inline base64 data URI. + + Accepts document dicts like: + {"type": "file", "file": Path("/path/to/doc.pdf")} # pathlib.Path + {"type": "file", "file": } # file-like object (BinaryIO) + {"type": "file", "file": b"raw bytes"} # raw bytes + + Bare ``str`` paths are not accepted — pass a ``pathlib.Path`` or + ``open(path, "rb")`` instead. See the str check below for the rationale. + + Returns: + {"type": "document_url", "document_url": "data:;base64,"} + or {"type": "image_url", "image_url": "data:;base64,"} + """ + file_input = document.get("file") + if file_input is None: + raise ValueError( + "document with type='file' must include a 'file' field containing " + "a pathlib.Path, file-like object, or bytes" + ) + + file_bytes: bytes + mime_type: str = "application/octet-stream" + file_name: str | None = None + + if isinstance(file_input, str): + # Bare strings are rejected here. The OCR ``document`` accepts a + # ``{"type": "file", "file": }`` shape, and when this helper + # runs in a proxy request handler ```` is attacker-controlled. + # Opening it as a path is an arbitrary local file read on the proxy + # host, which is then base64-encoded and forwarded to the OCR + # provider — an exfiltration primitive. + raise ValueError( + "OCR file input does not accept bare str values. Pass bytes, " + "a pathlib.Path, or a file-like object. To OCR a local file " + "from a path, call open(path, 'rb') yourself." + ) + if isinstance(file_input, os.PathLike): + # os.PathLike (pathlib.Path and custom __fspath__ classes) is a + # Python-level type that HTTP form values can't fabricate. + file_path = str(file_input) + if not os.path.isfile(file_path): + raise FileNotFoundError(f"File not found: {file_path}") + mime_type = get_mime_type(file_path) + file_name = os.path.basename(file_path) + with open(file_path, "rb") as f: + file_bytes = f.read() + elif isinstance(file_input, bytes): + file_bytes = file_input + elif isinstance(file_input, IOBase) or hasattr(file_input, "read"): + if hasattr(file_input, "name"): + file_name = getattr(file_input, "name", None) + if file_name: + mime_type = get_mime_type(file_name) + file_bytes = file_input.read() + if isinstance(file_bytes, str): + file_bytes = file_bytes.encode("utf-8") + else: + raise ValueError( + f"Unsupported file input type: {type(file_input)}. " + "Expected pathlib.Path, bytes, or a file-like object." + ) + + if not file_bytes: + raise ValueError("File is empty or could not be read") + + if "mime_type" in document: + mime_type = document["mime_type"] + + if not _MIME_PATTERN.match(mime_type): + raise ValueError(f"Invalid MIME type: {mime_type}") + + base64_data = base64.b64encode(file_bytes).decode("utf-8") + data_uri = f"data:{mime_type};base64,{base64_data}" + + if mime_type.startswith("image/"): + verbose_logger.debug( + f"OCR file input: Converted file to image_url data URI " + f"(mime={mime_type}, size={len(file_bytes)} bytes, name={file_name})" + ) + return {"type": "image_url", "image_url": data_uri} + + verbose_logger.debug( + f"OCR file input: Converted file to document_url data URI " + f"(mime={mime_type}, size={len(file_bytes)} bytes, name={file_name})" + ) + return {"type": "document_url", "document_url": data_uri} + + @client def ocr( model: str, - document: Dict[str, Any], - api_key: Optional[str] = None, - api_base: Optional[str] = None, - timeout: Optional[Union[float, httpx.Timeout]] = None, - custom_llm_provider: Optional[str] = None, - extra_headers: Optional[Dict[str, Any]] = None, + document: dict[str, Any], + api_key: str | None = None, + api_base: str | None = None, + timeout: Union[float, httpx.Timeout] | None = None, + custom_llm_provider: str | None = None, + extra_headers: dict[str, Any] | None = None, **kwargs, ) -> Union[OCRResponse, Coroutine[Any, Any, OCRResponse]]: """ @@ -295,96 +675,37 @@ def ocr( print(f"Page {page.index}: {page.markdown}") ``` """ - local_vars = locals() + completion_kwargs: dict[str, object] = { + "model": model, + "document": document, + "api_key": api_key, + "api_base": api_base, + "timeout": timeout, + "custom_llm_provider": custom_llm_provider, + "extra_headers": extra_headers, + "kwargs": kwargs, + } 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( + completion_kwargs["aocr"] = _is_async + 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, + ) + model = prepared.model + custom_llm_provider = prepared.custom_llm_provider + completion_kwargs.update( + {"model": model, "custom_llm_provider": custom_llm_provider} ) - # Optional Rust path: hand the whole Mistral OCR call to the Rust bridge. - if custom_llm_provider == "mistral" and rust_ocr_enabled(): + # Optional Rust path: hand supported OCR calls to the Rust bridge. + if _rust_ocr_supported(prepared) and rust_ocr_enabled(): rust_ocr = load_rust_ocr() if rust_ocr is None: verbose_logger.debug( @@ -395,31 +716,23 @@ def ocr( return _run_rust_ocr( rust_ocr=rust_ocr, - logging_obj=litellm_logging_obj, - provider_config=ocr_provider_config, + prepared_request=prepared, 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), ) 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 @@ -428,131 +741,6 @@ def ocr( model=model, custom_llm_provider=custom_llm_provider, original_exception=e, - completion_kwargs=local_vars, + completion_kwargs=completion_kwargs, extra_kwargs=kwargs, ) - - -################################################# -# Public utilities — used by the SDK and the proxy -################################################# - -_MIME_PATTERN = re.compile(r"^[\w.+-]+/[\w.+-]+$") - -_MIME_TYPE_MAP = { - ".pdf": "application/pdf", - ".png": "image/png", - ".jpg": "image/jpeg", - ".jpeg": "image/jpeg", - ".gif": "image/gif", - ".webp": "image/webp", - ".tiff": "image/tiff", - ".tif": "image/tiff", - ".bmp": "image/bmp", -} - - -def get_mime_type(file_path: str) -> str: - """ - Determine MIME type from file path extension. - - Falls back to mimetypes.guess_type, then to 'application/octet-stream'. - """ - ext = os.path.splitext(file_path)[1].lower() - mime = _MIME_TYPE_MAP.get(ext) - if mime: - return mime - guessed, _ = mimetypes.guess_type(file_path) - return guessed or "application/octet-stream" - - -def convert_file_document_to_url_document(document: Dict[str, Any]) -> Dict[str, str]: - """ - Convert a file-type document dict to a document_url-type document dict - with an inline base64 data URI. - - Accepts document dicts like: - {"type": "file", "file": Path("/path/to/doc.pdf")} # pathlib.Path - {"type": "file", "file": } # file-like object (BinaryIO) - {"type": "file", "file": b"raw bytes"} # raw bytes - - Bare ``str`` paths are not accepted — pass a ``pathlib.Path`` or - ``open(path, "rb")`` instead. See the str check below for the rationale. - - Returns: - {"type": "document_url", "document_url": "data:;base64,"} - or {"type": "image_url", "image_url": "data:;base64,"} - """ - file_input = document.get("file") - if file_input is None: - raise ValueError( - "document with type='file' must include a 'file' field containing " - "a pathlib.Path, file-like object, or bytes" - ) - - file_bytes: bytes - mime_type: str = "application/octet-stream" - file_name: Optional[str] = None - - if isinstance(file_input, str): - # Bare strings are rejected here. The OCR ``document`` accepts a - # ``{"type": "file", "file": }`` shape, and when this helper - # runs in a proxy request handler ```` is attacker-controlled. - # Opening it as a path is an arbitrary local file read on the proxy - # host, which is then base64-encoded and forwarded to the OCR - # provider — an exfiltration primitive. - raise ValueError( - "OCR file input does not accept bare str values. Pass bytes, " - "a pathlib.Path, or a file-like object. To OCR a local file " - "from a path, call open(path, 'rb') yourself." - ) - if isinstance(file_input, os.PathLike): - # os.PathLike (pathlib.Path and custom __fspath__ classes) is a - # Python-level type that HTTP form values can't fabricate. - file_path = str(file_input) - if not os.path.isfile(file_path): - raise FileNotFoundError(f"File not found: {file_path}") - mime_type = get_mime_type(file_path) - file_name = os.path.basename(file_path) - with open(file_path, "rb") as f: - file_bytes = f.read() - elif isinstance(file_input, bytes): - file_bytes = file_input - elif isinstance(file_input, IOBase) or hasattr(file_input, "read"): - if hasattr(file_input, "name"): - file_name = getattr(file_input, "name", None) - if file_name: - mime_type = get_mime_type(file_name) - file_bytes = file_input.read() - if isinstance(file_bytes, str): - file_bytes = file_bytes.encode("utf-8") - else: - raise ValueError( - f"Unsupported file input type: {type(file_input)}. " - "Expected pathlib.Path, bytes, or a file-like object." - ) - - if not file_bytes: - raise ValueError("File is empty or could not be read") - - if "mime_type" in document: - mime_type = document["mime_type"] - - if not _MIME_PATTERN.match(mime_type): - raise ValueError(f"Invalid MIME type: {mime_type}") - - base64_data = base64.b64encode(file_bytes).decode("utf-8") - data_uri = f"data:{mime_type};base64,{base64_data}" - - if mime_type.startswith("image/"): - verbose_logger.debug( - f"OCR file input: Converted file to image_url data URI " - f"(mime={mime_type}, size={len(file_bytes)} bytes, name={file_name})" - ) - return {"type": "image_url", "image_url": data_uri} - else: - verbose_logger.debug( - f"OCR file input: Converted file to document_url data URI " - f"(mime={mime_type}, size={len(file_bytes)} bytes, name={file_name})" - ) - return {"type": "document_url", "document_url": data_uri} diff --git a/litellm/ocr/rust_bridge.py b/litellm/ocr/rust_bridge.py index 4e688c42e1f..631fd4c63c5 100644 --- a/litellm/ocr/rust_bridge.py +++ b/litellm/ocr/rust_bridge.py @@ -11,7 +11,8 @@ can import it statically without forming an import cycle. from __future__ import annotations -from typing import Final, Protocol, cast +import os +from typing import Awaitable, Final, Protocol, cast class RustOcr(Protocol): @@ -23,9 +24,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]: ... + ) -> dict[str, object]: + raise NotImplementedError + + +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]]: + raise NotImplementedError class _Unset: @@ -34,23 +55,39 @@ class _Unset: _UNSET: Final[_Unset] = _Unset() -_rust_ocr_enabled = False + +def _env_enables_rust_ocr() -> bool: + return os.getenv("LITELLM_USE_RUST_OCR", "").strip().lower() in { + "1", + "true", + "yes", + "on", + } + + +_rust_ocr_enabled = _env_enables_rust_ocr() _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 packaged Rust 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: @@ -73,3 +110,15 @@ def load_rust_ocr() -> RustOcr | None: if native_bridge is None: return None return cast(RustOcr, native_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 + from litellm.rust_bridge import get_native_bridge + + native_bridge = get_native_bridge() + if native_bridge is None: + return None + return cast(RustAocr, getattr(native_bridge, "aocr", None)) diff --git a/tests/documentation_tests/test_env_keys.py b/tests/documentation_tests/test_env_keys.py index 681cd536259..60fbd505d67 100644 --- a/tests/documentation_tests/test_env_keys.py +++ b/tests/documentation_tests/test_env_keys.py @@ -24,6 +24,12 @@ EXCLUDED_GUARD_ONLY_VARS = { "MAVVRIK_FOCUS_FREQUENCY", } +# Temporary/internal rollout flags are intentionally not added to the public +# environment settings docs until the feature is ready for broad use. +EXCLUDED_ROLLOUT_FLAGS = { + "LITELLM_USE_RUST_OCR", +} + EXCLUDED_TERMINAL_VARS = { "TERM", "TERM_PROGRAM", @@ -71,6 +77,7 @@ for root, dirs, files in os.walk(repo_base): for match in getenv_matches if match not in EXCLUDED_TERMINAL_VARS and match not in EXCLUDED_GUARD_ONLY_VARS + and match not in EXCLUDED_ROLLOUT_FLAGS ) # Extract only the key part, excluding terminal vars # Find all keys using litellm.get_secret() diff --git a/tests/e2e/gateway/litellm-config.yml b/tests/e2e/gateway/litellm-config.yml index f4ca48cfee0..e059ac62429 100644 --- a/tests/e2e/gateway/litellm-config.yml +++ b/tests/e2e/gateway/litellm-config.yml @@ -140,6 +140,35 @@ model_list: model_info: mode: realtime + - model_name: rust-ocr-mistral + litellm_params: + model: mistral/mistral-ocr-latest + api_key: os.environ/MISTRAL_API_KEY + + - model_name: rust-ocr-azure-ai + litellm_params: + model: azure_ai/mistral-document-ai-2505 + api_base: os.environ/AZURE_API_BASE + api_key: os.environ/AZURE_API_KEY + + - model_name: rust-ocr-azure-document-intelligence + litellm_params: + model: azure_ai/doc-intelligence/prebuilt-layout + api_base: os.environ/AZURE_DOCUMENT_INTELLIGENCE_ENDPOINT + api_key: os.environ/AZURE_DOCUMENT_INTELLIGENCE_API_KEY + + - model_name: rust-ocr-vertex-mistral + litellm_params: + model: vertex_ai/mistral-ocr-2505 + vertex_project: os.environ/VERTEXAI_PROJECT + vertex_location: us-central1 + + - model_name: rust-ocr-vertex-deepseek + litellm_params: + model: vertex_ai/deepseek-ocr-maas + vertex_project: os.environ/VERTEXAI_PROJECT + vertex_location: us-central1 + mcp_servers: deepwiki_mcp: @@ -167,4 +196,3 @@ guardrails: US_SSN: BLOCK PHONE_NUMBER: BLOCK - diff --git a/tests/e2e/gateway/test_ocr_rust_e2e.py b/tests/e2e/gateway/test_ocr_rust_e2e.py new file mode 100644 index 00000000000..6ce59b2b5ac --- /dev/null +++ b/tests/e2e/gateway/test_ocr_rust_e2e.py @@ -0,0 +1,148 @@ +""" +Gateway E2E smoke for Rust-backed OCR. + +Start the proxy with: + +LITELLM_USE_RUST_OCR=1 litellm --config tests/e2e/gateway/litellm-config.yml --port 4000 +""" + +from __future__ import annotations + +import os +from dataclasses import dataclass +from pathlib import Path +from typing import Any + +import httpx +import pytest +import yaml + +TEST_PDF_URL = ( + "https://cdn.jsdelivr.net/gh/BerriAI/litellm" + "@d769e81c90d453240c61fc572cdb27fae06a89d0" + "/tests/llm_translation/fixtures/dummy.pdf" +) +TEST_IMAGE_URL = ( + "https://cdn.jsdelivr.net/gh/BerriAI/litellm" + "@d769e81c90d453240c61fc572cdb27fae06a89d0" + "/tests/image_gen_tests/test_image.png" +) + +RUST_OCR_GATEWAY_CASES = [ + pytest.param( + "rust-ocr-mistral", + {"type": "document_url", "document_url": TEST_PDF_URL}, + id="mistral", + ), + pytest.param( + "rust-ocr-azure-ai", + {"type": "document_url", "document_url": TEST_PDF_URL}, + id="azure_ai", + ), + pytest.param( + "rust-ocr-azure-document-intelligence", + {"type": "document_url", "document_url": TEST_PDF_URL}, + id="azure_document_intelligence", + ), + pytest.param( + "rust-ocr-vertex-mistral", + {"type": "document_url", "document_url": TEST_PDF_URL}, + id="vertex_mistral", + ), + pytest.param( + "rust-ocr-vertex-deepseek", + { + "type": "image_url", + "image_url": os.getenv("RUST_OCR_IMAGE_URL", TEST_IMAGE_URL), + }, + id="vertex_deepseek", + ), +] + +CONFIG_PATH = Path(__file__).with_name("litellm-config.yml") + + +@dataclass(frozen=True) +class OcrGateway: + base_url: str + master_key: str + + def model_names(self) -> set[str]: + with httpx.Client( + timeout=float(os.getenv("E2E_REQUEST_TIMEOUT", "120")) + ) as client: + response = client.get( + f"{self.base_url.rstrip('/')}/model/info", + headers={"Authorization": f"Bearer {self.master_key}"}, + ) + assert response.status_code == 200, response.text + return { + model["model_name"] + for model in response.json().get("data", []) + if "model_name" in model + } + + def ocr(self, model: str, document: dict[str, str]) -> httpx.Response: + with httpx.Client( + timeout=float(os.getenv("E2E_REQUEST_TIMEOUT", "120")) + ) as client: + return client.post( + f"{self.base_url.rstrip('/')}/v1/ocr", + headers={"Authorization": f"Bearer {self.master_key}"}, + json={"model": model, "document": document}, + ) + + +@dataclass(frozen=True) +class OcrResources: + gateway: OcrGateway + + +@pytest.fixture +def resources() -> OcrResources: + proxy_url = os.getenv("LITELLM_PROXY_URL") + if not proxy_url: + pytest.skip( + "Start a Rust OCR proxy and set LITELLM_PROXY_URL, e.g. http://localhost:4000" + ) + return OcrResources( + gateway=OcrGateway( + base_url=proxy_url, + master_key=os.getenv("LITELLM_MASTER_KEY", "sk-1234"), + ) + ) + + +def _assert_ocr_response_shape(response_json: dict[str, Any]) -> None: + assert response_json["object"] == "ocr" + assert response_json["model"] + assert isinstance(response_json["pages"], list) + assert len(response_json["pages"]) > 0 + assert "index" in response_json["pages"][0] + assert "markdown" in response_json["pages"][0] + + +class TestRustOcrGateway: + def test_rust_ocr_models_are_on_gateway_config(self) -> None: + config = yaml.safe_load(CONFIG_PATH.read_text()) + configured_models = { + model_config["model_name"] for model_config in config["model_list"] + } + + expected_models = {case.values[0] for case in RUST_OCR_GATEWAY_CASES} + assert expected_models.issubset(configured_models) + + def test_running_gateway_loaded_rust_ocr_models( + self, resources: OcrResources + ) -> None: + expected_models = {case.values[0] for case in RUST_OCR_GATEWAY_CASES} + assert expected_models.issubset(resources.gateway.model_names()) + + @pytest.mark.parametrize(("model", "document"), RUST_OCR_GATEWAY_CASES) + def test_rust_ocr_model_gateway_response( + self, resources: OcrResources, model: str, document: dict[str, str] + ) -> None: + response = resources.gateway.ocr(model, document) + + assert response.status_code == 200, response.text + _assert_ocr_response_shape(response.json()) diff --git a/tests/test_litellm/ocr/test_rust_bridge.py b/tests/test_litellm/ocr/test_rust_bridge.py index 11448aed828..7e23e441f50 100644 --- a/tests/test_litellm/ocr/test_rust_bridge.py +++ b/tests/test_litellm/ocr/test_rust_bridge.py @@ -3,6 +3,7 @@ import importlib import builtins import types +from typing import Any import httpx import pytest @@ -18,9 +19,12 @@ rust_bridge = importlib.import_module("litellm.ocr.rust_bridge") rust_bridge_loader = importlib.import_module("litellm.rust_bridge.loader") MODEL = "mistral/mistral-ocr-latest" -DOCUMENT = {"type": "document_url", "document_url": "https://example.com/doc.pdf"} +DOCUMENT: dict[str, object] = { + "type": "document_url", + "document_url": "https://example.com/doc.pdf", +} -FAKE_OCR_RESPONSE = { +FAKE_OCR_RESPONSE: dict[str, object] = { "pages": [{"index": 0, "markdown": "hello world"}], "model": "mistral-ocr-2505-completion", "document_annotation": None, @@ -29,21 +33,35 @@ FAKE_OCR_RESPONSE = { } +class CapturedException(Exception): + pass + + class RecordingBridge: """A fake ``RustOcr`` callable that records the args it was handed.""" - def __init__(self): - self.calls = [] + def __init__(self) -> None: + self.calls: list[dict[str, object]] = [] def __call__( - self, model, document, api_key, api_base, optional_params, timeout_seconds - ): + 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, + ) -> dict[str, object]: 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, } @@ -51,13 +69,81 @@ class RecordingBridge: return dict(FAKE_OCR_RESPONSE) +class RecordingAsyncBridge: + """A fake async ``RustAocr`` callable that records the args it was handed.""" + + def __init__(self) -> None: + self.calls: list[dict[str, object]] = [] + + async 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, + ) -> dict[str, object]: + 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, + } + ) + return dict(FAKE_OCR_RESPONSE) + + +class RaisingBridge: + 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, + ) -> dict[str, object]: + raise RuntimeError("bridge failed") + + +class RaisingAsyncBridge: + async 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, + ) -> dict[str, object]: + raise RuntimeError("bridge failed") + + class RecordingLogging: """A spy standing in for ``LiteLLMLoggingObj`` to capture ``pre_call``.""" - def __init__(self): - self.pre_call_kwargs = None + def __init__(self) -> None: + self.pre_call_kwargs: dict[str, object] | None = None - def pre_call(self, *, input, api_key, additional_args): + def pre_call( + self, + *, + input: str, + api_key: str | None, + additional_args: dict[str, object], + ) -> None: self.pre_call_kwargs = { "input": input, "api_key": api_key, @@ -68,22 +154,70 @@ class RecordingLogging: class FakeOCRConfig: """A stand-in ``BaseOCRConfig`` that echoes the request it would build.""" - def validate_environment( - self, *, headers, model, api_key, api_base, litellm_params - ): - return {"authorization": f"Bearer {api_key}"} + def __init__(self, api_key_env_var: str = "MISTRAL_API_KEY") -> None: + self.api_key_env_var = api_key_env_var - def get_complete_url(self, *, api_base, model, optional_params, litellm_params): + def get_api_key_env_var(self) -> str: + return self.api_key_env_var + + def validate_environment( + self, + *, + headers: dict[str, object], + model: str, + api_key: str | None, + api_base: str | None, + litellm_params: dict[str, object], + ) -> dict[str, object]: + return {"Authorization": f"Bearer {api_key}", **headers} + + def get_complete_url( + self, + *, + api_base: str | None, + model: str, + optional_params: dict[str, object], + litellm_params: dict[str, object], + ) -> str: return f"{api_base or 'https://api.mistral.ai/v1'}/ocr" +def build_prepared_request( + *, + logging_obj: RecordingLogging | None = None, + provider_config: FakeOCRConfig | None = None, + model: str = "mistral-ocr-latest", + document: dict[str, object] = DOCUMENT, + api_key: str | None = "sk-test", + api_base: str | None = None, + custom_llm_provider: str = "mistral", + extra_headers: dict[str, object] | None = None, + optional_params: dict[str, object] | None = None, + litellm_params: dict[str, object] | None = None, + timeout: float | httpx.Timeout | None = 12.5, +) -> Any: + return ocr_main._PreparedOCRRequest( + model=model, + document=document, + api_key=api_key, + api_base=api_base, + custom_llm_provider=custom_llm_provider, + extra_headers=extra_headers, + provider_config=provider_config or FakeOCRConfig(), + optional_params=optional_params or {}, + litellm_params=litellm_params or {}, + effective_timeout=timeout, + litellm_logging_obj=logging_obj or RecordingLogging(), + ) + + @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) rust_bridge_loader._cached_bridge = rust_bridge_loader._BRIDGE_SENTINEL yield - rust_bridge.use_litellm_rust(False, ocr=None) + rust_bridge.use_litellm_rust(False, ocr=None, aocr=None) rust_bridge_loader._cached_bridge = rust_bridge_loader._BRIDGE_SENTINEL @@ -95,6 +229,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() @@ -103,6 +245,11 @@ def test_use_litellm_rust_toggles_flag(): assert rust_bridge.rust_ocr_enabled() is False +def test_env_var_enables_rust_ocr(monkeypatch): + monkeypatch.setenv("LITELLM_USE_RUST_OCR", "1") + assert rust_bridge._env_enables_rust_ocr() is True + + def test_load_rust_ocr_returns_injected_impl(): bridge = RecordingBridge() litellm.use_litellm_rust(True, ocr=bridge) @@ -147,6 +294,12 @@ def test_native_bridge_available_reflects_loader(monkeypatch): assert rust_bridge_loader.native_bridge_available() is True +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. @@ -155,12 +308,15 @@ 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(monkeypatch): @@ -170,10 +326,12 @@ def test_explicit_ocr_none_clears_injected_impl(monkeypatch): lambda: None, ) 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(monkeypatch): @@ -186,6 +344,7 @@ def test_load_rust_ocr_none_when_extension_absent(monkeypatch): ) 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): @@ -194,6 +353,7 @@ def test_load_rust_ocr_uses_compiled_extension(monkeypatch): built in CI, so stand in a fake module via the bridge loader.""" fake_module = types.ModuleType("litellm.rust_bridge._native") 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.setattr( importlib.import_module("litellm.rust_bridge"), "get_native_bridge", @@ -202,6 +362,7 @@ def test_load_rust_ocr_uses_compiled_extension(monkeypatch): 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(): @@ -216,16 +377,14 @@ def test_run_rust_ocr_forwards_args_and_wraps_response(): response = ocr_main._run_rust_ocr( rust_ocr=bridge, - logging_obj=logging_obj, - provider_config=FakeOCRConfig(), + prepared_request=build_prepared_request( + logging_obj=logging_obj, + api_base="https://proxy.internal", + extra_headers={"x-trace-id": "trace-1"}, + optional_params={"include_image_base64": True}, + timeout=12.5, + ), resolve_api_key=lambda _name: None, - model="mistral-ocr-latest", - document=DOCUMENT, - api_key="sk-test", - api_base="https://proxy.internal", - optional_params={"include_image_base64": True}, - litellm_params={}, - timeout_seconds=12.5, ) assert isinstance(response, OCRResponse) @@ -236,6 +395,11 @@ 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": { + "Authorization": "Bearer sk-test", + "x-trace-id": "trace-1", + }, "optional_params": {"include_image_base64": True}, "timeout_seconds": 12.5, } @@ -248,23 +412,127 @@ def test_run_rust_ocr_resolves_key_via_secret_manager_when_missing(): ocr_main._run_rust_ocr( rust_ocr=bridge, - logging_obj=RecordingLogging(), - provider_config=FakeOCRConfig(), + prepared_request=build_prepared_request(api_key=None, timeout=None), resolve_api_key=lambda name: ( "sk-from-vault" if name == "MISTRAL_API_KEY" else None ), - model="mistral-ocr-latest", - document=DOCUMENT, - api_key=None, - api_base=None, - optional_params={}, - litellm_params={}, - timeout_seconds=None, ) assert bridge.calls[0]["api_key"] == "sk-from-vault" +def test_run_rust_ocr_uses_provider_api_key_env_var(): + bridge = RecordingBridge() + resolver_calls = [] + + def _resolver(name): + resolver_calls.append(name) + return "sk-provider-env" + + ocr_main._run_rust_ocr( + rust_ocr=bridge, + prepared_request=build_prepared_request( + provider_config=FakeOCRConfig(api_key_env_var="PROVIDER_OCR_API_KEY"), + model="provider-ocr-model", + api_key=None, + timeout=None, + ), + resolve_api_key=_resolver, + ) + + assert resolver_calls == ["PROVIDER_OCR_API_KEY"] + assert bridge.calls[0]["api_key"] == "sk-provider-env" + + +def test_prepare_rust_ocr_call_forwards_vertex_routing_metadata(): + bridge = RecordingBridge() + + ocr_main._run_rust_ocr( + rust_ocr=bridge, + prepared_request=build_prepared_request( + custom_llm_provider="vertex_ai", + model="mistral-ocr-maas", + litellm_params={ + "vertex_project": "project-1", + "vertex_location": "us-central1", + "vertex_credentials": "redacted", + }, + optional_params={"include_image_base64": True}, + timeout=None, + ), + resolve_api_key=lambda _name: None, + ) + + assert bridge.calls[0]["optional_params"] == { + "include_image_base64": True, + "vertex_project": "project-1", + "vertex_location": "us-central1", + } + + +def test_prepare_rust_ocr_call_resolves_vertex_routing_metadata_from_secret_manager(): + bridge = RecordingBridge() + + def _resolver(name: str) -> str | None: + return { + "VERTEXAI_PROJECT": "project-from-secret", + "VERTEXAI_LOCATION": "us-east5", + }.get(name) + + ocr_main._run_rust_ocr( + rust_ocr=bridge, + prepared_request=build_prepared_request( + custom_llm_provider="vertex_ai", + model="mistral-ocr-maas", + timeout=None, + ), + resolve_api_key=_resolver, + ) + + assert bridge.calls[0]["optional_params"]["vertex_project"] == "project-from-secret" + assert bridge.calls[0]["optional_params"]["vertex_location"] == "us-east5" + + +def test_prepare_rust_ocr_call_resolves_azure_ai_api_base_from_secret_manager(): + bridge = RecordingBridge() + + ocr_main._run_rust_ocr( + rust_ocr=bridge, + prepared_request=build_prepared_request( + custom_llm_provider="azure_ai", + model="pixtral-12b-2409", + api_base=None, + timeout=None, + ), + resolve_api_key=lambda name: ( + "https://azure.example.com" if name == "AZURE_AI_API_BASE" else None + ), + ) + + assert bridge.calls[0]["api_base"] == "https://azure.example.com" + + +def test_prepare_rust_ocr_call_resolves_document_intelligence_endpoint(): + bridge = RecordingBridge() + + ocr_main._run_rust_ocr( + rust_ocr=bridge, + prepared_request=build_prepared_request( + custom_llm_provider="azure_ai/doc-intelligence", + model="prebuilt-layout", + api_base=None, + timeout=None, + ), + resolve_api_key=lambda name: ( + "https://document-intelligence.example.com" + if name == "AZURE_DOCUMENT_INTELLIGENCE_ENDPOINT" + else None + ), + ) + + assert bridge.calls[0]["api_base"] == "https://document-intelligence.example.com" + + def test_run_rust_ocr_prefers_explicit_key_over_resolver(): bridge = RecordingBridge() resolver_calls = [] @@ -275,16 +543,8 @@ def test_run_rust_ocr_prefers_explicit_key_over_resolver(): ocr_main._run_rust_ocr( rust_ocr=bridge, - logging_obj=RecordingLogging(), - provider_config=FakeOCRConfig(), + prepared_request=build_prepared_request(api_key="sk-explicit", timeout=None), resolve_api_key=_resolver, - model="mistral-ocr-latest", - document=DOCUMENT, - api_key="sk-explicit", - api_base=None, - optional_params={}, - litellm_params={}, - timeout_seconds=None, ) assert bridge.calls[0]["api_key"] == "sk-explicit" @@ -297,16 +557,14 @@ def test_run_rust_ocr_runs_pre_call_logging(): ocr_main._run_rust_ocr( rust_ocr=RecordingBridge(), - logging_obj=logging_obj, - provider_config=FakeOCRConfig(), + prepared_request=build_prepared_request( + logging_obj=logging_obj, + api_base="https://api.mistral.ai/v1", + extra_headers={"x-trace-id": "trace-1"}, + optional_params={"include_image_base64": True}, + timeout=None, + ), resolve_api_key=lambda _name: None, - model="mistral-ocr-latest", - document=DOCUMENT, - api_key="sk-test", - api_base="https://api.mistral.ai/v1", - optional_params={"include_image_base64": True}, - litellm_params={}, - timeout_seconds=None, ) assert logging_obj.pre_call_kwargs is not None @@ -317,7 +575,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): @@ -325,6 +586,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, ) @@ -336,10 +598,93 @@ 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"] == { + "Authorization": "Bearer sk-test", + "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 +def test_ocr_routes_azure_ai_to_rust_when_enabled(fake_bridge): + response = litellm.ocr( + model="azure_ai/pixtral-12b-2409", + document=DOCUMENT, + api_key="sk-test", + api_base="https://example.services.ai.azure.com", + ) + + assert isinstance(response, OCRResponse) + assert len(fake_bridge.calls) == 1 + assert fake_bridge.calls[0]["model"] == "pixtral-12b-2409" + assert fake_bridge.calls[0]["custom_llm_provider"] == "azure_ai" + + +def test_ocr_exception_type_uses_resolved_provider_context( + monkeypatch: pytest.MonkeyPatch, +): + captured: dict[str, object] = {} + + def fake_exception_type(**kwargs: object) -> CapturedException: + captured.update(kwargs) + return CapturedException("wrapped") + + monkeypatch.setattr(ocr_main.litellm, "exception_type", fake_exception_type) + litellm.use_litellm_rust(True, ocr=RaisingBridge()) + + with pytest.raises(CapturedException): + litellm.ocr(model=MODEL, document=DOCUMENT, api_key="sk-test") + + assert captured["model"] == "mistral-ocr-latest" + assert captured["custom_llm_provider"] == "mistral" + + +@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"] == { + "Authorization": "Bearer sk-test", + "x-trace-id": "trace-1", + } + assert call["optional_params"].get("include_image_base64") is True + + +@pytest.mark.asyncio +async def test_aocr_exception_type_uses_resolved_provider_context( + monkeypatch: pytest.MonkeyPatch, +): + captured: dict[str, object] = {} + + def fake_exception_type(**kwargs: object) -> CapturedException: + captured.update(kwargs) + return CapturedException("wrapped") + + monkeypatch.setattr(ocr_main.litellm, "exception_type", fake_exception_type) + litellm.use_litellm_rust(True, aocr=RaisingAsyncBridge()) + + with pytest.raises(CapturedException): + await litellm.aocr(model=MODEL, document=DOCUMENT, api_key="sk-test") + + assert captured["model"] == "mistral-ocr-latest" + assert captured["custom_llm_provider"] == "mistral" + + 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.""" @@ -387,3 +732,26 @@ def test_ocr_falls_back_to_python_when_bridge_unavailable(monkeypatch): assert captured.get("called") is True # Python path was used assert isinstance(response, OCRResponse) + + +def test_ocr_provider_configs_expose_api_key_env_vars(): + from litellm.llms.azure_ai.ocr.document_intelligence.transformation import ( + AzureDocumentIntelligenceOCRConfig, + ) + from litellm.llms.azure_ai.ocr.transformation import AzureAIOCRConfig + from litellm.llms.base_llm.ocr.transformation import BaseOCRConfig + from litellm.llms.mistral.ocr.transformation import MistralOCRConfig + from litellm.llms.vertex_ai.ocr.deepseek_transformation import ( + VertexAIDeepSeekOCRConfig, + ) + from litellm.llms.vertex_ai.ocr.transformation import VertexAIOCRConfig + + assert BaseOCRConfig().get_api_key_env_var() is None + assert MistralOCRConfig().get_api_key_env_var() == "MISTRAL_API_KEY" + assert AzureAIOCRConfig().get_api_key_env_var() == "AZURE_AI_API_KEY" + assert ( + AzureDocumentIntelligenceOCRConfig().get_api_key_env_var() + == "AZURE_DOCUMENT_INTELLIGENCE_API_KEY" + ) + assert VertexAIOCRConfig().get_api_key_env_var() == "VERTEX_AI_API_KEY" + assert VertexAIDeepSeekOCRConfig().get_api_key_env_var() == "VERTEX_AI_API_KEY" diff --git a/uv.lock b/uv.lock index d81d6e5d8a0..917dff39e38 100644 --- a/uv.lock +++ b/uv.lock @@ -9,7 +9,7 @@ resolution-markers = [ ] [options] -exclude-newer = "2026-06-22T14:51:56.5801Z" +exclude-newer = "2026-06-22T19:05:55.080417Z" exclude-newer-span = "P3D" [manifest] From 01035499da489318786a04ce6f53bc0b86f7f4fe Mon Sep 17 00:00:00 2001 From: Yassin Kortam Date: Fri, 26 Jun 2026 01:39:07 +0300 Subject: [PATCH 7/8] fix(cache): apply Redis namespace to all key operations (#31288) The namespace configured under cache_params was only applied to get/set/ increment paths. Operations that take keys through other code paths (the Lua scripts registered via async_register_script, delete, scan_iter, rpush, lpop, get_ttl, and the sync increment_cache) hit raw keys. With a namespace set, the rate limiter ({key}:tokens/requests/window), pod-lock release, and budget limiters wrote keys outside the configured prefix, breaking multi-tenant key isolation and leaving those operations reading keys the namespaced writes never created. check_and_fix_namespace is now applied uniformly across every key-taking RedisCache operation. It is a no-op when no namespace is configured, so deployments without a namespace are unaffected. The prefix is prepended ahead of any {hash-tag}, so Redis Cluster slotting is preserved. Resolves LIT-3374 --- litellm/caching/redis_cache.py | 33 +++- .../test_litellm/caching/test_redis_cache.py | 174 ++++++++++++++++++ 2 files changed, 203 insertions(+), 4 deletions(-) diff --git a/litellm/caching/redis_cache.py b/litellm/caching/redis_cache.py index ba07511448a..5c9f74e4c18 100644 --- a/litellm/caching/redis_cache.py +++ b/litellm/caching/redis_cache.py @@ -435,6 +435,7 @@ class RedisCache(BaseCache): _redis_client = self.redis_client start_time = time.time() set_ttl = self.get_ttl(ttl=ttl) + key = self.check_and_fix_namespace(key=key) try: start_time = time.time() result: int = _redis_client.incr(name=key, amount=value) # type: ignore @@ -498,6 +499,7 @@ class RedisCache(BaseCache): ) return [] + pattern = self.check_and_fix_namespace(key=pattern) async for key in _redis_client.scan_iter(match=pattern + "*", count=count): # type: ignore keys.append(key) if len(keys) >= count: @@ -538,6 +540,11 @@ class RedisCache(BaseCache): Register a Lua script with Redis asynchronously. Works with both standalone Redis and Redis Cluster. + The returned callable namespaces every key it is invoked with, so Lua + scripts hit the same prefixed keys as get/set/increment. Without this, + scripts would operate on raw keys while the rest of the cache uses the + namespace, leaving rate-limit and lock keys outside the configured prefix. + Args: script (str): The Lua script to register @@ -548,7 +555,15 @@ class RedisCache(BaseCache): _redis_client = self.init_async_client() # For standalone Redis if hasattr(_redis_client, "register_script"): - return _redis_client.register_script(script) # type: ignore + registered_script = _redis_client.register_script(script) # type: ignore + + async def namespaced_script( + keys: list[str], args: list[Any], client: Any = None + ) -> Any: + keys = [self.check_and_fix_namespace(key=key) for key in keys] + return await registered_script(keys=keys, args=args, client=client) + + return namespaced_script # For Redis Cluster elif hasattr(_redis_client, "script_load"): # Load the script and get its SHA @@ -556,6 +571,7 @@ class RedisCache(BaseCache): # Return a callable that uses evalsha async def script_callable(keys: List[str], args: List[Any]) -> Any: + keys = [self.check_and_fix_namespace(key=key) for key in keys] return _redis_client.evalsha(script_sha, len(keys), *keys, *args) # type: ignore return script_callable @@ -1257,6 +1273,7 @@ class RedisCache(BaseCache): async def delete_cache_keys(self, keys): # typed as Any, redis python lib has incomplete type stubs for RedisCluster and does not include `delete` _redis_client: Any = self.init_async_client() + keys = [self.check_and_fix_namespace(key=key) for key in keys] # keys is a list, unpack it so it gets passed as individual elements to delete await _redis_client.delete(*keys) @@ -1322,10 +1339,12 @@ class RedisCache(BaseCache): async def async_delete_cache(self, key: str): # typed as Any, redis python lib has incomplete type stubs for RedisCluster and does not include `delete` _redis_client: Any = self.init_async_client() + key = self.check_and_fix_namespace(key=key) # keys is str return await _redis_client.delete(key) def delete_cache(self, key): + key = self.check_and_fix_namespace(key=key) self.redis_client.delete(key) async def _pipeline_increment_helper( @@ -1432,6 +1451,7 @@ class RedisCache(BaseCache): try: # typed as Any, redis python lib has incomplete type stubs for RedisCluster and does not include `ttl` _redis_client: Any = self.init_async_client() + key = self.check_and_fix_namespace(key=key) ttl = await _redis_client.ttl(key) if ttl <= -1: # -1 means the key does not exist, -2 key does not exist return None @@ -1460,6 +1480,7 @@ class RedisCache(BaseCache): int: The length of the list after the push operation """ _redis_client: Any = self.init_async_client() + key = self.check_and_fix_namespace(key=key) start_time = time.time() try: response = await _redis_client.rpush(key, *values) @@ -1499,7 +1520,8 @@ class RedisCache(BaseCache): ) -> List[int]: """Helper function for pipeline rpush operations""" for rpush_op in rpush_list: - pipe.rpush(rpush_op["key"], *rpush_op["values"]) + key = self.check_and_fix_namespace(key=rpush_op["key"]) + pipe.rpush(key, *rpush_op["values"]) results = await pipe.execute() # Preserve positional correspondence — raise on per-command errors for r in results: @@ -1586,6 +1608,7 @@ class RedisCache(BaseCache): **kwargs, ) -> Union[Any, List[Any]]: _redis_client: Any = self.init_async_client() + key = self.check_and_fix_namespace(key=key) start_time = time.time() print_verbose(f"LPOP from Redis list: key: {key}, count: {count}") try: @@ -1658,17 +1681,19 @@ class RedisCache(BaseCache): if major_version >= 7: for lpop_op in lpop_list: - pipe.lpop(lpop_op["key"], lpop_op["count"]) + key = self.check_and_fix_namespace(key=lpop_op["key"]) + pipe.lpop(key, lpop_op["count"]) raw_results = await pipe.execute() else: # For Redis < 7, LPOP doesn't support count param. # Issue `count` individual LPOP commands per key, all in one pipeline. counts: List[int] = [] for lpop_op in lpop_list: + key = self.check_and_fix_namespace(key=lpop_op["key"]) count = lpop_op["count"] or 1 counts.append(count) for _ in range(count): - pipe.lpop(lpop_op["key"]) + pipe.lpop(key) flat_results = await pipe.execute() # Re-group the flat results back into per-key lists diff --git a/tests/test_litellm/caching/test_redis_cache.py b/tests/test_litellm/caching/test_redis_cache.py index 78192400fb0..6aaadce93ce 100644 --- a/tests/test_litellm/caching/test_redis_cache.py +++ b/tests/test_litellm/caching/test_redis_cache.py @@ -517,3 +517,177 @@ async def test_async_lpop_with_float_redis_version( # Verify the method completed without error assert result is not None + + +# LIT-3374: the namespace must be applied uniformly across every key-taking +# Redis operation, not just get/set/increment. Before the fix these paths wrote +# or read raw keys, so with a namespace configured the prefixed keys other +# operations created were silently missed. + + +@pytest.mark.parametrize( + "namespace, raw_keys, expected_keys", + [ + (None, ["{k:v}:tokens", "{k:v}:requests"], ["{k:v}:tokens", "{k:v}:requests"]), + ( + "litellm_sandbox", + ["{k:v}:tokens", "{k:v}:requests"], + ["litellm_sandbox:{k:v}:tokens", "litellm_sandbox:{k:v}:requests"], + ), + ], +) +@pytest.mark.asyncio +async def test_async_register_script_namespaces_keys( + namespace, raw_keys, expected_keys, monkeypatch, redis_no_ping +): + """The callable returned by async_register_script (used by the rate limiter + Lua scripts, pod-lock release, and budget limiters) must namespace every key + it is invoked with. The hash tag is preserved so cluster slotting is intact.""" + monkeypatch.setenv("REDIS_HOST", "https://my-test-host") + redis_cache = RedisCache(namespace=namespace) + + registered_script = AsyncMock(return_value="ok") + mock_redis_instance = MagicMock() + mock_redis_instance.register_script = MagicMock(return_value=registered_script) + + with patch.object( + redis_cache, "init_async_client", return_value=mock_redis_instance + ): + script = redis_cache.async_register_script("return 1") + result = await script(keys=raw_keys, args=[60]) + + assert result == "ok" + registered_script.assert_awaited_once_with( + keys=expected_keys, args=[60], client=None + ) + + +@pytest.mark.parametrize("namespace, expected", [(None, "k"), ("ns", "ns:k")]) +@pytest.mark.asyncio +async def test_async_delete_cache_namespaces_key( + namespace, expected, monkeypatch, redis_no_ping +): + monkeypatch.setenv("REDIS_HOST", "https://my-test-host") + redis_cache = RedisCache(namespace=namespace) + mock_redis_instance = AsyncMock() + with patch.object( + redis_cache, "init_async_client", return_value=mock_redis_instance + ): + await redis_cache.async_delete_cache("k") + mock_redis_instance.delete.assert_awaited_once_with(expected) + + +@pytest.mark.parametrize("namespace, expected", [(None, "k"), ("ns", "ns:k")]) +@pytest.mark.asyncio +async def test_delete_cache_keys_namespaces_keys( + namespace, expected, monkeypatch, redis_no_ping +): + monkeypatch.setenv("REDIS_HOST", "https://my-test-host") + redis_cache = RedisCache(namespace=namespace) + mock_redis_instance = AsyncMock() + with patch.object( + redis_cache, "init_async_client", return_value=mock_redis_instance + ): + await redis_cache.delete_cache_keys(["k"]) + mock_redis_instance.delete.assert_awaited_once_with(expected) + + +@pytest.mark.parametrize("namespace, expected", [(None, "k"), ("ns", "ns:k")]) +@pytest.mark.asyncio +async def test_async_get_ttl_namespaces_key( + namespace, expected, monkeypatch, redis_no_ping +): + monkeypatch.setenv("REDIS_HOST", "https://my-test-host") + redis_cache = RedisCache(namespace=namespace) + mock_redis_instance = AsyncMock() + mock_redis_instance.ttl = AsyncMock(return_value=42) + with patch.object( + redis_cache, "init_async_client", return_value=mock_redis_instance + ): + ttl = await redis_cache.async_get_ttl("k") + assert ttl == 42 + mock_redis_instance.ttl.assert_awaited_once_with(expected) + + +@pytest.mark.parametrize("namespace, expected", [(None, "k"), ("ns", "ns:k")]) +@pytest.mark.asyncio +async def test_async_lpop_namespaces_key( + namespace, expected, monkeypatch, redis_no_ping +): + monkeypatch.setenv("REDIS_HOST", "https://my-test-host") + redis_cache = RedisCache(namespace=namespace) + mock_redis_instance = AsyncMock() + mock_redis_instance.lpop = AsyncMock(return_value=b"value") + with patch.object( + redis_cache, "init_async_client", return_value=mock_redis_instance + ): + await redis_cache.async_lpop(key="k") + mock_redis_instance.lpop.assert_awaited_once_with(expected, None) + + +@pytest.mark.parametrize("namespace, expected", [(None, "k"), ("ns", "ns:k")]) +@pytest.mark.asyncio +async def test_async_rpush_namespaces_key( + namespace, expected, monkeypatch, redis_no_ping +): + monkeypatch.setenv("REDIS_HOST", "https://my-test-host") + redis_cache = RedisCache(namespace=namespace) + mock_redis_instance = AsyncMock() + mock_redis_instance.rpush = AsyncMock(return_value=1) + with patch.object( + redis_cache, "init_async_client", return_value=mock_redis_instance + ): + await redis_cache.async_rpush("k", ["v"]) + mock_redis_instance.rpush.assert_awaited_once_with(expected, "v") + + +@pytest.mark.parametrize("namespace, expected_match", [(None, "k*"), ("ns", "ns:k*")]) +@pytest.mark.asyncio +async def test_async_scan_iter_namespaces_pattern( + namespace, expected_match, monkeypatch, redis_no_ping +): + monkeypatch.setenv("REDIS_HOST", "https://my-test-host") + redis_cache = RedisCache(namespace=namespace) + + captured = {} + + def scan_iter(match, count): + captured["match"] = match + + async def gen(): + for _ in (): + yield _ + + return gen() + + mock_redis_instance = MagicMock() + mock_redis_instance.scan_iter = scan_iter + with patch.object( + redis_cache, "init_async_client", return_value=mock_redis_instance + ): + await redis_cache.async_scan_iter(pattern="k") + assert captured["match"] == expected_match + + +@pytest.mark.parametrize("namespace, expected", [(None, "k"), ("ns", "ns:k")]) +def test_increment_cache_namespaces_key( + namespace, expected, monkeypatch, redis_no_ping +): + monkeypatch.setenv("REDIS_HOST", "https://my-test-host") + redis_cache = RedisCache(namespace=namespace) + mock_client = MagicMock() + mock_client.incr.return_value = 5 + mock_client.ttl.return_value = 100 + redis_cache.redis_client = mock_client + redis_cache.increment_cache(key="k", value=1) + mock_client.incr.assert_called_once_with(name=expected, amount=1) + + +@pytest.mark.parametrize("namespace, expected", [(None, "k"), ("ns", "ns:k")]) +def test_delete_cache_namespaces_key(namespace, expected, monkeypatch, redis_no_ping): + monkeypatch.setenv("REDIS_HOST", "https://my-test-host") + redis_cache = RedisCache(namespace=namespace) + mock_client = MagicMock() + redis_cache.redis_client = mock_client + redis_cache.delete_cache(key="k") + mock_client.delete.assert_called_once_with(expected) From 7ffce15766e6d12145aca5e1500645f60126dc6d Mon Sep 17 00:00:00 2001 From: milan-berri Date: Fri, 26 Jun 2026 00:40:54 +0200 Subject: [PATCH 8/8] Add GA pricing for gemini-3-pro-image and gemini-3.1-flash-image. (#30022) Fixes #29794. Adds bare, gemini/, and vertex_ai/ entries copied from preview models so proxy cost tracking works for GA model names. Co-authored-by: Cursor --- model_prices_and_context_window.json | 192 +++++++++++++++++++++++++++ 1 file changed, 192 insertions(+) diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 5f3f2294147..4ec4d8bb884 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -16153,6 +16153,47 @@ "supports_service_tier": true, "supports_image_size": false }, + "gemini-3-pro-image": { + "input_cost_per_image": 0.0011, + "input_cost_per_token": 2e-06, + "input_cost_per_token_batches": 1e-06, + "litellm_provider": "vertex_ai-language-models", + "max_input_tokens": 65536, + "max_output_tokens": 32768, + "max_tokens": 32768, + "mode": "image_generation", + "output_cost_per_image": 0.134, + "output_cost_per_image_token": 0.00012, + "output_cost_per_token": 1.2e-05, + "output_cost_per_token_batches": 6e-06, + "source": "https://ai.google.dev/gemini-api/docs/pricing#gemini-3-pro-image", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/completions", + "/v1/batch" + ], + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text", + "image" + ], + "supports_function_calling": false, + "supports_prompt_caching": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_vision": true, + "supports_web_search": true, + "search_context_cost_per_query": { + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014, + "search_context_size_high": 0.014 + }, + "web_search_billing_unit": "per_query", + "supports_service_tier": true + }, "gemini-3-pro-image-preview": { "input_cost_per_image": 0.0011, "input_cost_per_token": 2e-06, @@ -16194,6 +16235,44 @@ "web_search_billing_unit": "per_query", "supports_service_tier": true }, + "gemini-3.1-flash-image": { + "input_cost_per_image": 0.00056, + "input_cost_per_token": 5e-07, + "litellm_provider": "vertex_ai-language-models", + "max_input_tokens": 65536, + "max_output_tokens": 32768, + "max_tokens": 32768, + "mode": "image_generation", + "output_cost_per_image": 0.0672, + "output_cost_per_image_token": 6e-05, + "output_cost_per_token": 3e-06, + "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing#gemini-models", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/completions", + "/v1/batch" + ], + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text", + "image" + ], + "supports_function_calling": false, + "supports_prompt_caching": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_vision": true, + "supports_web_search": true, + "search_context_cost_per_query": { + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014, + "search_context_size_high": 0.014 + }, + "web_search_billing_unit": "per_query" + }, "gemini-3.1-flash-image-preview": { "input_cost_per_image": 0.00056, "input_cost_per_token": 5e-07, @@ -17765,6 +17844,49 @@ "supports_service_tier": true, "supports_image_size": false }, + "gemini/gemini-3-pro-image": { + "input_cost_per_image": 0.0011, + "input_cost_per_token": 2e-06, + "input_cost_per_token_batches": 1e-06, + "litellm_provider": "gemini", + "max_input_tokens": 65536, + "max_output_tokens": 32768, + "max_tokens": 32768, + "mode": "image_generation", + "output_cost_per_image": 0.134, + "output_cost_per_image_token": 0.00012, + "output_cost_per_token": 1.2e-05, + "rpm": 1000, + "tpm": 4000000, + "output_cost_per_token_batches": 6e-06, + "source": "https://ai.google.dev/gemini-api/docs/pricing#gemini-3-pro-image", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/completions", + "/v1/batch" + ], + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text", + "image" + ], + "supports_function_calling": false, + "supports_prompt_caching": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_vision": true, + "supports_web_search": true, + "search_context_cost_per_query": { + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014, + "search_context_size_high": 0.014 + }, + "web_search_billing_unit": "per_query", + "supports_service_tier": true + }, "gemini/gemini-3-pro-image-preview": { "input_cost_per_image": 0.0011, "input_cost_per_token": 2e-06, @@ -17808,6 +17930,48 @@ "web_search_billing_unit": "per_query", "supports_service_tier": true }, + "gemini/gemini-3.1-flash-image": { + "input_cost_per_token": 2.5e-07, + "input_cost_per_token_batches": 1.25e-07, + "litellm_provider": "gemini", + "max_input_tokens": 65536, + "max_output_tokens": 32768, + "max_tokens": 32768, + "mode": "image_generation", + "output_cost_per_image": 0.045, + "output_cost_per_image_token": 6e-05, + "output_cost_per_image_token_batches": 3e-05, + "output_cost_per_token": 1.5e-06, + "output_cost_per_token_batches": 7.5e-07, + "rpm": 1000, + "tpm": 4000000, + "source": "https://ai.google.dev/gemini-api/docs/pricing#gemini-3.1-flash-image", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/completions", + "/v1/batch" + ], + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text", + "image" + ], + "supports_function_calling": false, + "supports_prompt_caching": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_vision": true, + "supports_web_search": true, + "search_context_cost_per_query": { + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014, + "search_context_size_high": 0.014 + }, + "web_search_billing_unit": "per_query" + }, "gemini/gemini-3.1-flash-image-preview": { "input_cost_per_token": 2.5e-07, "input_cost_per_token_batches": 1.25e-07, @@ -35573,6 +35737,21 @@ "tpm": 8000000, "supports_image_size": false }, + "vertex_ai/gemini-3-pro-image": { + "input_cost_per_image": 0.0011, + "input_cost_per_token": 2e-06, + "input_cost_per_token_batches": 1e-06, + "litellm_provider": "vertex_ai-language-models", + "max_input_tokens": 65536, + "max_output_tokens": 32768, + "max_tokens": 32768, + "mode": "image_generation", + "output_cost_per_image": 0.134, + "output_cost_per_image_token": 0.00012, + "output_cost_per_token": 1.2e-05, + "output_cost_per_token_batches": 6e-06, + "source": "https://docs.cloud.google.com/vertex-ai/generative-ai/docs/models/gemini/3-pro-image" + }, "vertex_ai/gemini-3-pro-image-preview": { "input_cost_per_image": 0.0011, "input_cost_per_token": 2e-06, @@ -35588,6 +35767,19 @@ "output_cost_per_token_batches": 6e-06, "source": "https://docs.cloud.google.com/vertex-ai/generative-ai/docs/models/gemini/3-pro-image" }, + "vertex_ai/gemini-3.1-flash-image": { + "input_cost_per_image": 0.00056, + "input_cost_per_token": 5e-07, + "litellm_provider": "vertex_ai-language-models", + "max_input_tokens": 65536, + "max_output_tokens": 32768, + "max_tokens": 32768, + "mode": "image_generation", + "output_cost_per_image": 0.0672, + "output_cost_per_image_token": 6e-05, + "output_cost_per_token": 3e-06, + "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing#gemini-models" + }, "vertex_ai/gemini-3.1-flash-image-preview": { "input_cost_per_image": 0.00056, "input_cost_per_token": 5e-07,