refactor(rust_bridge): group route modules into packages and split ocr into main and rust

Move each route's bridge module under litellm/rust_bridge/<route>/ so a folder
means a Rust implementation exists while the catalog row says whether it is
used. OCR now keeps the Python implementation in litellm/ocr/main.py and the
Rust selection in litellm/ocr/rust.py, removing litellm/ocr/legacy.py

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
Yujong Lee 2026-09-16 20:34:51 +00:00
parent 9484595fa2
commit 64f2a3d098
37 changed files with 516 additions and 515 deletions

View file

@ -159,7 +159,7 @@ fn redact(
}
pub(super) fn response(py: Python<'_>, response: &LiteLLMOcrResponse) -> PyResult<Py<PyAny>> {
py.import("litellm.rust_bridge.ocr")?
py.import("litellm.rust_bridge.ocr.native")?
.getattr("_response")?
.call1((to_py(py, response)?,))
.map(Bound::unbind)
@ -172,7 +172,7 @@ pub(super) fn map_failure(
provider: &str,
) -> PyResult<Py<PyBaseException>> {
Ok(py
.import("litellm.rust_bridge.ocr_lifecycle")?
.import("litellm.rust_bridge.ocr.lifecycle")?
.getattr("map_failure")?
.call1((error, request, provider))?
.extract()?)

View file

@ -1434,7 +1434,7 @@ from .skills.main import (
adelete_skill,
)
from .containers.main import *
from .ocr.main import *
from .ocr.rust import *
from .rust_bridge import rust
from .rag.main import *
from .sandbox.main import *

View file

@ -25,8 +25,8 @@ from litellm.llms.custom_httpx.http_handler import (
_get_httpx_client,
get_async_httpx_client,
)
from litellm.rust_bridge import chat_completions as rust_chat_completions_bridge
from litellm.rust_bridge.chat_completions import rust_chat_completions_accepts
from litellm.rust_bridge.chat_completions import native as rust_chat_completions_bridge
from litellm.rust_bridge.chat_completions.native import rust_chat_completions_accepts
from litellm.types.llms.anthropic import (
ContentBlockDelta,
ContentBlockStart,

View file

@ -7,7 +7,7 @@ from litellm.litellm_core_utils.audio_utils.utils import process_audio_file
from litellm.rust_bridge import runtime
from litellm.rust_bridge.catalog import Context, Route
from litellm.rust_bridge.timeouts import timeout_to_seconds
from litellm.rust_bridge.transcription import (
from litellm.rust_bridge.transcription.native import (
NATIVE_ATRANSCRIPTION,
NATIVE_TRANSCRIPTION,
RustAtranscription,

View file

@ -16,8 +16,8 @@ from litellm.llms.custom_httpx.http_handler import (
_get_httpx_client,
get_async_httpx_client,
)
from litellm.rust_bridge import chat_completions as rust_chat_completions_bridge
from litellm.rust_bridge.chat_completions import rust_chat_completions_accepts
from litellm.rust_bridge.chat_completions import native as rust_chat_completions_bridge
from litellm.rust_bridge.chat_completions.native import rust_chat_completions_accepts
from litellm.types.utils import ModelResponse
from litellm.utils import CustomStreamWrapper

View file

@ -2464,7 +2464,7 @@ class BaseLLMHTTPHandler:
if has_agentic_hook:
return None
from litellm.rust_bridge import messages as rust_messages_bridge
from litellm.rust_bridge.messages import native as rust_messages_bridge
upstream_body: Final = {key: value for key, value in request_body.items() if key != "stream"}
try:
@ -6659,7 +6659,7 @@ class BaseLLMHTTPHandler:
@asynccontextmanager
async def _backend_connection():
if _rust_responses_websocket_enabled(custom_llm_provider):
from litellm.rust_bridge import responses_websocket as rust_responses_websocket
from litellm.rust_bridge.responses import websocket as rust_responses_websocket
rust_backend: Final = await rust_responses_websocket.connect(
url=ws_url,

View file

@ -1,5 +1,5 @@
"""OCR module for LiteLLM."""
from .main import aocr, ocr
from .rust import aocr, ocr
__all__ = ["aocr", "ocr"]

View file

@ -75,9 +75,9 @@ def _native_helpers_selected() -> bool:
def get_mime_type(file_path: str) -> str:
native: Final = _MIME_TYPE.load() if _native_helpers_selected() else None
if native is None:
from litellm.ocr import legacy
from litellm.ocr import main
return legacy.get_mime_type(file_path)
return main.get_mime_type(file_path)
return native(file_path)
@ -91,9 +91,9 @@ def get_max_file_bytes() -> int:
def convert_file_document_to_url_document(document: FileDocument) -> dict[str, str]:
native: Final = _FILE_DOCUMENT.load() if _native_helpers_selected() else None
if native is None:
from litellm.ocr import legacy
from litellm.ocr import main
return legacy.convert_file_document_to_url_document(document)
return main.convert_file_document_to_url_document(document)
return native(document)
@ -102,17 +102,17 @@ def convert_upload_to_url_document(
) -> dict[str, str]:
native: Final = _UPLOAD_DOCUMENT.load() if _native_helpers_selected() else None
if native is None:
from litellm.ocr import legacy
from litellm.ocr import main
if len(file_content) > _PYTHON_MAX_FILE_BYTES:
raise ValueError("OCR file exceeds the size limit")
content_mime: Final = content_type.split(";")[0].strip() if content_type else None
mime_type: Final = (
legacy.get_mime_type(filename)
main.get_mime_type(filename)
if filename and (not content_mime or content_mime == "application/octet-stream")
else content_mime or "application/octet-stream"
)
return legacy.convert_file_document_to_url_document(
return main.convert_file_document_to_url_document(
{"type": "file", "file": file_content, "mime_type": mime_type}
)
return native(file_content, filename, content_type)

View file

@ -1,413 +0,0 @@
"""
Main OCR function for LiteLLM.
"""
import asyncio
import base64
import mimetypes
import os
import re
from collections.abc import Coroutine, Mapping
from dataclasses import dataclass
from io import IOBase
from types import MappingProxyType
from typing import Final, cast # noqa: TID251 # adapters preserve the legacy untyped contracts
import httpx
import litellm
from litellm._logging import verbose_logger
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 (
OCR_REQUEST_FORMAT_PARAM,
BaseOCRConfig,
OCRResponse,
parse_ocr_request_format,
)
from litellm.llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler
from litellm.ocr.input import FileReader
from litellm.types.router import GenericLiteLLMParams
from litellm.types.utils import CustomPricingLiteLLMParams
from litellm.utils import ProviderConfigManager, client
base_llm_http_handler: Final = BaseLLMHTTPHandler()
@dataclass(frozen=True, slots=True)
class _PreparedOCRRequest:
model: str
document: Mapping[str, object]
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: float | httpx.Timeout
litellm_logging_obj: LiteLLMLoggingObj
def _prepare_ocr_request(
model: str,
document: Mapping[str, object],
api_key: str | None,
api_base: str | None,
timeout: float | httpx.Timeout | None,
custom_llm_provider: str | None,
extra_headers: dict[str, object] | None,
kwargs: dict[str, object],
) -> _PreparedOCRRequest:
litellm_logging_obj: Final = cast( # cast-ok: @client supplies the logging object; preserve legacy failure behavior
LiteLLMLoggingObj, kwargs.pop("litellm_logging_obj")
)
litellm_call_id: Final = cast( # cast-ok: @client supplies the call id without coercion
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,
)
ocr_provider_config: Final = 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}")
resolved_api_key, resolved_api_base = ocr_provider_config.resolve_connection_params(
api_key=api_key,
api_base=api_base,
dynamic_api_key=dynamic_api_key,
dynamic_api_base=dynamic_api_base,
)
verbose_logger.debug("OCR call - model: %s, provider: %s", model, custom_llm_provider)
litellm_params: Final = GenericLiteLLMParams.model_validate(kwargs)
supported_params: Final = ocr_provider_config.get_supported_ocr_params(model=model)
requested_format: Final = kwargs.get(OCR_REQUEST_FORMAT_PARAM)
if requested_format is not None:
try:
parsed_format: Final = parse_ocr_request_format(requested_format)
except ValueError as e:
raise litellm.exceptions.UnsupportedParamsError(
message=f"{e}", model=model, llm_provider=custom_llm_provider
) from e
if OCR_REQUEST_FORMAT_PARAM not in supported_params and parsed_format == "native":
raise litellm.exceptions.UnsupportedParamsError(
message=(
f"`{OCR_REQUEST_FORMAT_PARAM}='native'` is not supported for provider: {custom_llm_provider}, "
f"model: {model}"
),
model=model,
llm_provider=custom_llm_provider,
)
non_default_params: Final = {}
for param in supported_params:
if param in kwargs:
non_default_params[param] = kwargs.pop(param)
optional_params: Final = ocr_provider_config.map_ocr_params(
non_default_params=non_default_params,
optional_params={},
model=model,
)
verbose_logger.debug("OCR optional_params after mapping: %s", optional_params)
effective_timeout: Final = 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": resolved_api_base,
**litellm_params.model_dump(include=frozenset(CustomPricingLiteLLMParams.model_fields), exclude_none=True),
},
custom_llm_provider=custom_llm_provider,
)
return _PreparedOCRRequest(
model=model,
document=document,
api_key=resolved_api_key,
api_base=resolved_api_base,
custom_llm_provider=custom_llm_provider,
extra_headers=extra_headers,
provider_config=ocr_provider_config,
optional_params=cast(
dict[str, object], optional_params
), # cast-ok: provider configs return heterogeneous OCR options
litellm_params=dict(litellm_params),
effective_timeout=effective_timeout,
litellm_logging_obj=litellm_logging_obj,
)
def _error_provider(model: str, custom_llm_provider: str | None) -> str | None:
if custom_llm_provider is not None:
return custom_llm_provider
prefix: Final = model.partition("/")[0]
if prefix in {"mistral", "azure_ai", "vertex_ai"}:
return prefix
return "mistral" if model.startswith("mistral-ocr") else None
@client
async def aocr(
model: str,
document: Mapping[str, object],
api_key: str | None = None,
api_base: str | None = None,
timeout: float | httpx.Timeout | None = None,
custom_llm_provider: str | None = None,
extra_headers: dict[str, object] | None = None,
**kwargs: object, # kwargs-ok: public OCR accepts provider-specific options
) -> OCRResponse:
completion_kwargs: Final[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:
prepared: Final = _prepare_ocr_request(
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,
)
model = prepared.model
custom_llm_provider = prepared.custom_llm_provider
completion_kwargs.update({"model": model, "custom_llm_provider": custom_llm_provider})
response = base_llm_http_handler.ocr(
model=prepared.model,
document=cast( # cast-ok: preserve legacy document fields for provider validation
dict[str, str], 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(f"Got an unexpected None response from the OCR API: {response}")
return response
except Exception as e:
error_provider: Final = _error_provider(model, custom_llm_provider)
error_model: Final = model.removeprefix(f"{error_provider}/") if error_provider else model
raise litellm.exception_type(
model=error_model,
custom_llm_provider=error_provider,
original_exception=e,
completion_kwargs=completion_kwargs,
extra_kwargs=kwargs,
)
_MIME_PATTERN: Final = re.compile(r"^[\w.+-]+/[\w.+-]+$")
_MIME_TYPE_MAP: Final = MappingProxyType(
{
".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:
ext: Final = os.path.splitext(file_path)[1].lower()
mime: Final = _MIME_TYPE_MAP.get(ext)
if mime:
return mime
guessed, _ = mimetypes.guess_type(file_path)
return guessed or "application/octet-stream"
def _read_file(file_input: object) -> tuple[bytes, str, str | None]:
if isinstance(file_input, str):
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):
file_path: Final = str(cast(object, file_input)) # cast-ok: preserve staging's str(PathLike) conversion
if not os.path.isfile(file_path):
raise FileNotFoundError(f"File not found: {file_path}")
mime_type: Final = get_mime_type(file_path)
with open(file_path, "rb") as stream:
return stream.read(), mime_type, os.path.basename(file_path)
if isinstance(file_input, bytes):
return file_input, "application/octet-stream", None
if isinstance(file_input, IOBase) or hasattr(file_input, "read"):
file_name: Final = cast( # cast-ok: retain legacy validation and errors for file-like metadata
str | None, getattr(file_input, "name", None)
)
inferred_mime: Final = get_mime_type(file_name) if file_name else "application/octet-stream"
reader: Final = cast(FileReader, file_input) # cast-ok: legacy accepts duck-typed file readers
content: Final = reader.read()
return content.encode("utf-8") if isinstance(content, str) else content, inferred_mime, file_name
raise ValueError(
f"Unsupported file input type: {type(file_input)}. Expected pathlib.Path, bytes, or a file-like object."
)
def convert_file_document_to_url_document(document: Mapping[str, object]) -> dict[str, str]:
file_input: Final = 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, inferred_mime, file_name = _read_file(file_input)
if not file_bytes:
raise ValueError("File is empty or could not be read")
mime_type: Final = cast( # cast-ok: keep staging's MIME validation errors
str, document.get("mime_type", inferred_mime)
)
if not _MIME_PATTERN.match(mime_type):
raise ValueError(f"Invalid MIME type: {mime_type}")
base64_data: Final = base64.b64encode(file_bytes).decode("utf-8")
data_uri: Final = f"data:{mime_type};base64,{base64_data}"
if mime_type.startswith("image/"):
verbose_logger.debug(
"OCR file input: Converted file to image_url data URI (mime=%s, size=%s bytes, name=%s)",
mime_type,
len(file_bytes),
file_name,
)
return {"type": "image_url", "image_url": data_uri}
verbose_logger.debug(
"OCR file input: Converted file to document_url data URI (mime=%s, size=%s bytes, name=%s)",
mime_type,
len(file_bytes),
file_name,
)
return {"type": "document_url", "document_url": data_uri}
@client
def ocr(
model: str,
document: Mapping[str, object],
api_key: str | None = None,
api_base: str | None = None,
timeout: float | httpx.Timeout | None = None,
custom_llm_provider: str | None = None,
extra_headers: dict[str, object] | None = None,
**kwargs: object, # kwargs-ok: public OCR accepts provider-specific options
) -> OCRResponse | Coroutine[object, object, OCRResponse]:
completion_kwargs: Final[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:
_is_async: Final = kwargs.pop("aocr", False) is True
completion_kwargs["aocr"] = _is_async
prepared: Final = _prepare_ocr_request(
model=model,
document=document,
api_key=api_key,
api_base=api_base,
kwargs=kwargs,
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})
response: Final = base_llm_http_handler.ocr(
model=prepared.model,
document=cast( # cast-ok: preserve legacy document fields for provider validation
dict[str, str], 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=prepared.extra_headers,
provider_config=prepared.provider_config,
litellm_params=prepared.litellm_params,
)
return response
except Exception as e:
error_provider: Final = _error_provider(model, custom_llm_provider)
error_model: Final = model.removeprefix(f"{error_provider}/") if error_provider else model
raise litellm.exception_type(
model=error_model,
custom_llm_provider=error_provider,
original_exception=e,
completion_kwargs=completion_kwargs,
extra_kwargs=kwargs,
)

View file

@ -1,20 +1,188 @@
from collections.abc import Awaitable, Callable, Coroutine, Mapping
from typing import Final, cast # noqa: TID251 # native binding selects a sync result or an async awaitable
"""
Main OCR function for LiteLLM.
"""
import asyncio
import base64
import mimetypes
import os
import re
from collections.abc import Coroutine, Mapping
from dataclasses import dataclass
from io import IOBase
from types import MappingProxyType
from typing import Final, cast # noqa: TID251 # adapters preserve the legacy untyped contracts
import httpx
from litellm.llms.base_llm.ocr.transformation import OCRResponse
from litellm.ocr import legacy
from litellm.ocr.input import convert_file_document_to_url_document, get_mime_type
from litellm.rust_bridge.catalog import Context, Route
from litellm.rust_bridge.ocr import LiteLLMOcrRequest
from litellm.rust_bridge.ocr_lifecycle import NATIVE_OCR_LIFECYCLE, NativeOcrLifecycle
from litellm.rust_bridge.runtime import arun, run
import litellm
from litellm._logging import verbose_logger
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 (
OCR_REQUEST_FORMAT_PARAM,
BaseOCRConfig,
OCRResponse,
parse_ocr_request_format,
)
from litellm.llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler
from litellm.ocr.input import FileReader
from litellm.types.router import GenericLiteLLMParams
from litellm.types.utils import CustomPricingLiteLLMParams
from litellm.utils import ProviderConfigManager, client
__all__ = ("aocr", "convert_file_document_to_url_document", "get_mime_type", "ocr")
base_llm_http_handler: Final = BaseLLMHTTPHandler()
def _bind_request(
@dataclass(frozen=True, slots=True)
class _PreparedOCRRequest:
model: str
document: Mapping[str, object]
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: float | httpx.Timeout
litellm_logging_obj: LiteLLMLoggingObj
def _prepare_ocr_request(
model: str,
document: Mapping[str, object],
api_key: str | None,
api_base: str | None,
timeout: float | httpx.Timeout | None,
custom_llm_provider: str | None,
extra_headers: dict[str, object] | None,
kwargs: dict[str, object],
) -> _PreparedOCRRequest:
litellm_logging_obj: Final = cast( # cast-ok: @client supplies the logging object; preserve legacy failure behavior
LiteLLMLoggingObj, kwargs.pop("litellm_logging_obj")
)
litellm_call_id: Final = cast( # cast-ok: @client supplies the call id without coercion
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,
)
ocr_provider_config: Final = 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}")
resolved_api_key, resolved_api_base = ocr_provider_config.resolve_connection_params(
api_key=api_key,
api_base=api_base,
dynamic_api_key=dynamic_api_key,
dynamic_api_base=dynamic_api_base,
)
verbose_logger.debug("OCR call - model: %s, provider: %s", model, custom_llm_provider)
litellm_params: Final = GenericLiteLLMParams.model_validate(kwargs)
supported_params: Final = ocr_provider_config.get_supported_ocr_params(model=model)
requested_format: Final = kwargs.get(OCR_REQUEST_FORMAT_PARAM)
if requested_format is not None:
try:
parsed_format: Final = parse_ocr_request_format(requested_format)
except ValueError as e:
raise litellm.exceptions.UnsupportedParamsError(
message=f"{e}", model=model, llm_provider=custom_llm_provider
) from e
if OCR_REQUEST_FORMAT_PARAM not in supported_params and parsed_format == "native":
raise litellm.exceptions.UnsupportedParamsError(
message=(
f"`{OCR_REQUEST_FORMAT_PARAM}='native'` is not supported for provider: {custom_llm_provider}, "
f"model: {model}"
),
model=model,
llm_provider=custom_llm_provider,
)
non_default_params: Final = {}
for param in supported_params:
if param in kwargs:
non_default_params[param] = kwargs.pop(param)
optional_params: Final = ocr_provider_config.map_ocr_params(
non_default_params=non_default_params,
optional_params={},
model=model,
)
verbose_logger.debug("OCR optional_params after mapping: %s", optional_params)
effective_timeout: Final = 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": resolved_api_base,
**litellm_params.model_dump(include=frozenset(CustomPricingLiteLLMParams.model_fields), exclude_none=True),
},
custom_llm_provider=custom_llm_provider,
)
return _PreparedOCRRequest(
model=model,
document=document,
api_key=resolved_api_key,
api_base=resolved_api_base,
custom_llm_provider=custom_llm_provider,
extra_headers=extra_headers,
provider_config=ocr_provider_config,
optional_params=cast(
dict[str, object], optional_params
), # cast-ok: provider configs return heterogeneous OCR options
litellm_params=dict(litellm_params),
effective_timeout=effective_timeout,
litellm_logging_obj=litellm_logging_obj,
)
def _error_provider(model: str, custom_llm_provider: str | None) -> str | None:
if custom_llm_provider is not None:
return custom_llm_provider
prefix: Final = model.partition("/")[0]
if prefix in {"mistral", "azure_ai", "vertex_ai"}:
return prefix
return "mistral" if model.startswith("mistral-ocr") else None
@client
async def aocr(
model: str,
document: Mapping[str, object],
api_key: str | None = None,
@ -23,61 +191,223 @@ def _bind_request(
custom_llm_provider: str | None = None,
extra_headers: dict[str, object] | None = None,
**kwargs: object, # kwargs-ok: public OCR accepts provider-specific options
) -> LiteLLMOcrRequest:
return LiteLLMOcrRequest(
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,
)
def _public_request(name: str, args: tuple[object, ...], kwargs: dict[str, object]) -> LiteLLMOcrRequest:
) -> OCRResponse:
completion_kwargs: Final[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:
return _bind_request(*args, **kwargs) # pyright: ignore[reportArgumentType] # Python binds the public arguments before native validation
except TypeError as error:
raise TypeError(str(error).replace("_bind_request()", f"{name}()")) from None
prepared: Final = _prepare_ocr_request(
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,
)
model = prepared.model
custom_llm_provider = prepared.custom_llm_provider
completion_kwargs.update({"model": model, "custom_llm_provider": custom_llm_provider})
def ocr(
*args: object,
**kwargs: object, # kwargs-ok: preserve the public OCR call shape
) -> OCRResponse | Coroutine[object, object, OCRResponse]:
request: Final = _public_request("ocr", args, kwargs)
fallback: Final = cast( # cast-ok: forward the original call shape through the legacy @client decorator
Callable[..., OCRResponse | Coroutine[object, object, OCRResponse]], legacy.ocr
)
if request.kwargs.get("aocr"):
return fallback(*args, **kwargs)
return run(
_context(request),
binding=NATIVE_OCR_LIFECYCLE,
native=lambda hook: cast( # cast-ok: False selects the synchronous result
OCRResponse, hook(request, args, kwargs, False)
),
python=lambda: fallback(*args, **kwargs),
)
async def aocr(*args: object, **kwargs: object) -> OCRResponse: # kwargs-ok: preserve the public OCR call shape
request: Final = _public_request("aocr", args, kwargs)
fallback: Final = cast( # cast-ok: forward the original call shape through the legacy @client decorator
Callable[..., Awaitable[OCRResponse]], legacy.aocr
)
async def native(hook: NativeOcrLifecycle) -> OCRResponse:
return await cast( # cast-ok: True selects the asynchronous result
Awaitable[OCRResponse], hook(request, args, kwargs, True)
response = base_llm_http_handler.ocr(
model=prepared.model,
document=cast( # cast-ok: preserve legacy document fields for provider validation
dict[str, str], 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,
)
return await arun(
_context(request), binding=NATIVE_OCR_LIFECYCLE, native=native, python=lambda: fallback(*args, **kwargs)
if asyncio.iscoroutine(response):
response = await response
if response is None:
raise ValueError(f"Got an unexpected None response from the OCR API: {response}")
return response
except Exception as e:
error_provider: Final = _error_provider(model, custom_llm_provider)
error_model: Final = model.removeprefix(f"{error_provider}/") if error_provider else model
raise litellm.exception_type(
model=error_model,
custom_llm_provider=error_provider,
original_exception=e,
completion_kwargs=completion_kwargs,
extra_kwargs=kwargs,
)
_MIME_PATTERN: Final = re.compile(r"^[\w.+-]+/[\w.+-]+$")
_MIME_TYPE_MAP: Final = MappingProxyType(
{
".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:
ext: Final = os.path.splitext(file_path)[1].lower()
mime: Final = _MIME_TYPE_MAP.get(ext)
if mime:
return mime
guessed, _ = mimetypes.guess_type(file_path)
return guessed or "application/octet-stream"
def _read_file(file_input: object) -> tuple[bytes, str, str | None]:
if isinstance(file_input, str):
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):
file_path: Final = str(cast(object, file_input)) # cast-ok: preserve staging's str(PathLike) conversion
if not os.path.isfile(file_path):
raise FileNotFoundError(f"File not found: {file_path}")
mime_type: Final = get_mime_type(file_path)
with open(file_path, "rb") as stream:
return stream.read(), mime_type, os.path.basename(file_path)
if isinstance(file_input, bytes):
return file_input, "application/octet-stream", None
if isinstance(file_input, IOBase) or hasattr(file_input, "read"):
file_name: Final = cast( # cast-ok: retain legacy validation and errors for file-like metadata
str | None, getattr(file_input, "name", None)
)
inferred_mime: Final = get_mime_type(file_name) if file_name else "application/octet-stream"
reader: Final = cast(FileReader, file_input) # cast-ok: legacy accepts duck-typed file readers
content: Final = reader.read()
return content.encode("utf-8") if isinstance(content, str) else content, inferred_mime, file_name
raise ValueError(
f"Unsupported file input type: {type(file_input)}. Expected pathlib.Path, bytes, or a file-like object."
)
def _context(request: LiteLLMOcrRequest) -> Context:
return Context(Route.OCR, provider=request.custom_llm_provider, model=request.model)
def convert_file_document_to_url_document(document: Mapping[str, object]) -> dict[str, str]:
file_input: Final = 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, inferred_mime, file_name = _read_file(file_input)
if not file_bytes:
raise ValueError("File is empty or could not be read")
mime_type: Final = cast( # cast-ok: keep staging's MIME validation errors
str, document.get("mime_type", inferred_mime)
)
if not _MIME_PATTERN.match(mime_type):
raise ValueError(f"Invalid MIME type: {mime_type}")
base64_data: Final = base64.b64encode(file_bytes).decode("utf-8")
data_uri: Final = f"data:{mime_type};base64,{base64_data}"
if mime_type.startswith("image/"):
verbose_logger.debug(
"OCR file input: Converted file to image_url data URI (mime=%s, size=%s bytes, name=%s)",
mime_type,
len(file_bytes),
file_name,
)
return {"type": "image_url", "image_url": data_uri}
verbose_logger.debug(
"OCR file input: Converted file to document_url data URI (mime=%s, size=%s bytes, name=%s)",
mime_type,
len(file_bytes),
file_name,
)
return {"type": "document_url", "document_url": data_uri}
@client
def ocr(
model: str,
document: Mapping[str, object],
api_key: str | None = None,
api_base: str | None = None,
timeout: float | httpx.Timeout | None = None,
custom_llm_provider: str | None = None,
extra_headers: dict[str, object] | None = None,
**kwargs: object, # kwargs-ok: public OCR accepts provider-specific options
) -> OCRResponse | Coroutine[object, object, OCRResponse]:
completion_kwargs: Final[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:
_is_async: Final = kwargs.pop("aocr", False) is True
completion_kwargs["aocr"] = _is_async
prepared: Final = _prepare_ocr_request(
model=model,
document=document,
api_key=api_key,
api_base=api_base,
kwargs=kwargs,
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})
response: Final = base_llm_http_handler.ocr(
model=prepared.model,
document=cast( # cast-ok: preserve legacy document fields for provider validation
dict[str, str], 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=prepared.extra_headers,
provider_config=prepared.provider_config,
litellm_params=prepared.litellm_params,
)
return response
except Exception as e:
error_provider: Final = _error_provider(model, custom_llm_provider)
error_model: Final = model.removeprefix(f"{error_provider}/") if error_provider else model
raise litellm.exception_type(
model=error_model,
custom_llm_provider=error_provider,
original_exception=e,
completion_kwargs=completion_kwargs,
extra_kwargs=kwargs,
)

83
litellm/ocr/rust.py Normal file
View file

@ -0,0 +1,83 @@
from collections.abc import Awaitable, Callable, Coroutine, Mapping
from typing import Final, cast # noqa: TID251 # native binding selects a sync result or an async awaitable
import httpx
from litellm.llms.base_llm.ocr.transformation import OCRResponse
from litellm.ocr import main
from litellm.ocr.input import convert_file_document_to_url_document, get_mime_type
from litellm.rust_bridge.catalog import Context, Route
from litellm.rust_bridge.ocr.lifecycle import NATIVE_OCR_LIFECYCLE, NativeOcrLifecycle
from litellm.rust_bridge.ocr.native import LiteLLMOcrRequest
from litellm.rust_bridge.runtime import arun, run
__all__ = ("aocr", "convert_file_document_to_url_document", "get_mime_type", "ocr")
def _bind_request(
model: str,
document: Mapping[str, object],
api_key: str | None = None,
api_base: str | None = None,
timeout: float | httpx.Timeout | None = None,
custom_llm_provider: str | None = None,
extra_headers: dict[str, object] | None = None,
**kwargs: object, # kwargs-ok: public OCR accepts provider-specific options
) -> LiteLLMOcrRequest:
return LiteLLMOcrRequest(
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,
)
def _public_request(name: str, args: tuple[object, ...], kwargs: dict[str, object]) -> LiteLLMOcrRequest:
try:
return _bind_request(*args, **kwargs) # pyright: ignore[reportArgumentType] # Python binds the public arguments before native validation
except TypeError as error:
raise TypeError(str(error).replace("_bind_request()", f"{name}()")) from None
def ocr(
*args: object,
**kwargs: object, # kwargs-ok: preserve the public OCR call shape
) -> OCRResponse | Coroutine[object, object, OCRResponse]:
request: Final = _public_request("ocr", args, kwargs)
fallback: Final = cast( # cast-ok: forward the original call shape through the Python @client decorator
Callable[..., OCRResponse | Coroutine[object, object, OCRResponse]], main.ocr
)
if request.kwargs.get("aocr"):
return fallback(*args, **kwargs)
return run(
_context(request),
binding=NATIVE_OCR_LIFECYCLE,
native=lambda hook: cast( # cast-ok: False selects the synchronous result
OCRResponse, hook(request, args, kwargs, False)
),
python=lambda: fallback(*args, **kwargs),
)
async def aocr(*args: object, **kwargs: object) -> OCRResponse: # kwargs-ok: preserve the public OCR call shape
request: Final = _public_request("aocr", args, kwargs)
fallback: Final = cast( # cast-ok: forward the original call shape through the Python @client decorator
Callable[..., Awaitable[OCRResponse]], main.aocr
)
async def native(hook: NativeOcrLifecycle) -> OCRResponse:
return await cast( # cast-ok: True selects the asynchronous result
Awaitable[OCRResponse], hook(request, args, kwargs, True)
)
return await arun(
_context(request), binding=NATIVE_OCR_LIFECYCLE, native=native, python=lambda: fallback(*args, **kwargs)
)
def _context(request: LiteLLMOcrRequest) -> Context:
return Context(Route.OCR, provider=request.custom_llm_provider, model=request.model)

View file

@ -3,7 +3,7 @@ from collections.abc import Coroutine, Mapping, Sequence
from typing import Literal, Never, TypeAlias, final
from litellm.llms.base_llm.ocr.transformation import OCRResponse
from litellm.rust_bridge.ocr import LiteLLMOcrRequest
from litellm.rust_bridge.ocr.native import LiteLLMOcrRequest
_InputSource: TypeAlias = Literal["request", "deployment", "environment"]

View file

View file

View file

@ -6,7 +6,7 @@ from typing import Final, Protocol, cast # noqa: TID251 # validates dynamicall
import litellm
from litellm.llms.base_llm.ocr.transformation import OCRResponse
from litellm.rust_bridge.bindings import NativeBinding
from litellm.rust_bridge.ocr import LiteLLMOcrRequest
from litellm.rust_bridge.ocr.native import LiteLLMOcrRequest
class NativeOcrLifecycle(Protocol):

View file

@ -14,7 +14,7 @@ from litellm.types.llms.anthropic_messages.anthropic_response import (
)
from litellm.types.router import GenericLiteLLMParams
rust_messages = importlib.import_module("litellm.rust_bridge.messages")
rust_messages = importlib.import_module("litellm.rust_bridge.messages.native")
rust_bridge_loader = importlib.import_module("litellm.rust_bridge.loader")
FAKE_MESSAGES_RESPONSE: dict[str, object] = {

View file

@ -2339,7 +2339,7 @@ class TestRustChatCompletionsHook:
@pytest.fixture(autouse=True)
def _reset_bridge(self, monkeypatch):
from litellm.rust_bridge import chat_completions as bridge
from litellm.rust_bridge.chat_completions import native as bridge
from litellm.rust_bridge import configuration
monkeypatch.setenv("LITELLM_RUST", "1")
@ -2376,7 +2376,7 @@ class TestRustChatCompletionsHook:
@staticmethod
def _inject():
from litellm.rust_bridge import chat_completions as bridge
from litellm.rust_bridge.chat_completions import native as bridge
seen = {"gate": [], "call": []}

View file

@ -14,13 +14,13 @@ from unittest.mock import MagicMock, patch
import boto3
import httpx
import pytest
from botocore.credentials import Credentials
from botocore.exceptions import ClientError
from litellm.llms.bedrock.chat.converse_handler import BedrockConverseLLM
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler
from litellm.rust_bridge import chat_completions as bridge
from litellm.rust_bridge import configuration
from litellm.rust_bridge.chat_completions import native as bridge
from litellm.types.utils import ModelResponse
from tests.test_litellm.llms.bedrock.event_loop_probe import EventLoopProbe

View file

@ -14,9 +14,9 @@ from litellm.litellm_core_utils.litellm_logging import Logging, use_custom_prici
from litellm.llms.base_llm.ocr.transformation import OCRPage, OCRResponse, OCRUsageInfo
from litellm.llms.custom_httpx import llm_http_handler
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler
from litellm.ocr.legacy import _prepare_ocr_request
from litellm.ocr.main import _prepare_ocr_request
from litellm.rust_bridge import bindings, configuration, runtime
from litellm.rust_bridge.ocr_lifecycle import NATIVE_OCR_LIFECYCLE
from litellm.rust_bridge.ocr.lifecycle import NATIVE_OCR_LIFECYCLE
@pytest.fixture

View file

@ -2,7 +2,7 @@
Tests for the OCR `req_format` option in the SDK request path.
"""
from litellm.rust_bridge import ocr as rust_ocr_bridge
from litellm.rust_bridge.ocr import native as rust_ocr_bridge
def test_rust_ocr_response_retains_provider_native_response():

View file

@ -2,7 +2,8 @@ from __future__ import annotations
import pytest
from litellm.rust_bridge import configuration, responses_websocket
from litellm.rust_bridge import configuration
from litellm.rust_bridge.responses import websocket as responses_websocket
class _FakeNativeConnection:

View file

@ -10,7 +10,7 @@ from __future__ import annotations
import pytest
from litellm.rust_bridge import configuration
from litellm.rust_bridge import chat_completions as bridge
from litellm.rust_bridge.chat_completions import native as bridge
from litellm.types.utils import ModelResponse
RUST_RESPONSE = {

View file

@ -6,10 +6,10 @@ import pytest
import litellm
from litellm.llms.base_llm.ocr.transformation import OCRResponse
from litellm.ocr import legacy
from litellm.ocr import main as python_ocr
from litellm.rust_bridge import bindings, configuration, runtime
from litellm.rust_bridge.ocr import LiteLLMOcrRequest
from litellm.rust_bridge.ocr_lifecycle import NATIVE_OCR_LIFECYCLE
from litellm.rust_bridge.ocr.lifecycle import NATIVE_OCR_LIFECYCLE
from litellm.rust_bridge.ocr.native import LiteLLMOcrRequest
@pytest.fixture(autouse=True)
@ -26,7 +26,7 @@ def isolated_ocr_configuration(monkeypatch: pytest.MonkeyPatch) -> Generator[Non
async def test_unavailable_native_uses_legacy(monkeypatch: pytest.MonkeyPatch, asynchronous: bool) -> None:
response: Final = OCRResponse(pages=[], model="mistral-ocr-latest")
fallback: Final = AsyncMock(return_value=response) if asynchronous else Mock(return_value=response)
monkeypatch.setattr(legacy, "aocr" if asynchronous else "ocr", fallback)
monkeypatch.setattr(python_ocr, "aocr" if asynchronous else "ocr", fallback)
NATIVE_OCR_LIFECYCLE.override(None)
document: Final = {"type": "document_url", "document_url": "https://example.com"}
@ -150,7 +150,7 @@ async def test_environment_opt_out_never_loads_native(
monkeypatch.setenv("LITELLM_RUST", "0")
response: Final = OCRResponse(pages=[], model="mistral-ocr-latest")
fallback: Final = AsyncMock(return_value=response) if asynchronous else Mock(return_value=response)
monkeypatch.setattr(legacy, "aocr" if asynchronous else "ocr", fallback)
monkeypatch.setattr(python_ocr, "aocr" if asynchronous else "ocr", fallback)
load: Final = Mock(side_effect=AssertionError("native must not be loaded"))
monkeypatch.setattr(bindings, "get_native_bridge", load)
litellm.rust(enabled)
@ -179,7 +179,7 @@ async def test_native_is_enabled_by_default(
native: Final = AsyncMock(return_value=response) if asynchronous else Mock(return_value=response)
NATIVE_OCR_LIFECYCLE.override(native)
fallback: Final = Mock(side_effect=AssertionError("legacy must not run"))
monkeypatch.setattr(legacy, "aocr" if asynchronous else "ocr", fallback)
monkeypatch.setattr(python_ocr, "aocr" if asynchronous else "ocr", fallback)
result: Final = (
await litellm.aocr("mistral/mistral-ocr-latest", {})
@ -212,7 +212,7 @@ async def test_only_native_declines_replay_on_legacy(
monkeypatch.setattr(runtime, "native_exception_types", lambda: (Declined, Upstream))
response: Final = OCRResponse(pages=[], model="mistral-ocr-latest")
fallback: Final = AsyncMock(return_value=response) if asynchronous else Mock(return_value=response)
monkeypatch.setattr(legacy, "aocr" if asynchronous else "ocr", fallback)
monkeypatch.setattr(python_ocr, "aocr" if asynchronous else "ocr", fallback)
document: Final = {"type": "file", "file": b"pdf"}
async def call() -> object:

View file

@ -9,7 +9,7 @@ import pytest
import litellm
from litellm.llms.bedrock.audio_transcription import BedrockAudioTranscriptionRustDispatch
from litellm.rust_bridge import bindings, configuration
from litellm.rust_bridge.transcription import NATIVE_ATRANSCRIPTION, NATIVE_TRANSCRIPTION
from litellm.rust_bridge.transcription.native import NATIVE_ATRANSCRIPTION, NATIVE_TRANSCRIPTION
MODEL: Final = "bedrock/mistral.voxtral-mini-3b-2507"
AUDIO_FILE: Final = ("audio.wav", b"audio", "audio/wav")

View file

@ -580,7 +580,7 @@ async def test_retained_argument_aliases_and_body_roots_survive_envelope_replace
def test_unstarted_native_coroutine_releases_input_without_reading_file(ocr_server: RecordingServer) -> None:
from litellm.ocr.main import _public_request
from litellm.ocr.rust import _public_request
from litellm.rust_bridge import _native
ocr_server.expected_requests = 0

View file

@ -8,7 +8,7 @@ from typing import Final
import pytest
import litellm
from litellm.rust_bridge import ocr as rust_ocr_bridge
from litellm.rust_bridge.ocr import native as rust_ocr_bridge
pytestmark = pytest.mark.requires_rust_extension