From 64f2a3d098b697bf2370771b3c329cf155f144d8 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Wed, 16 Sep 2026 20:34:51 +0000 Subject: [PATCH] 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// 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> --- .../python-bridge/src/routes/ocr/callbacks.rs | 4 +- litellm/__init__.py | 2 +- litellm/llms/anthropic/chat/handler.py | 4 +- .../bedrock/audio_transcription/__init__.py | 2 +- litellm/llms/bedrock/chat/converse_handler.py | 4 +- litellm/llms/custom_httpx/llm_http_handler.py | 4 +- litellm/ocr/__init__.py | 2 +- litellm/ocr/input.py | 14 +- litellm/ocr/legacy.py | 413 ---------------- litellm/ocr/main.py | 454 +++++++++++++++--- litellm/ocr/rust.py | 83 ++++ litellm/rust_bridge/_native.pyi | 2 +- .../rust_bridge/chat_completions/__init__.py | 0 .../native.py} | 0 litellm/rust_bridge/messages/__init__.py | 0 .../{messages.py => messages/native.py} | 0 litellm/rust_bridge/ocr/__init__.py | 0 .../{ocr_lifecycle.py => ocr/lifecycle.py} | 2 +- litellm/rust_bridge/{ocr.py => ocr/native.py} | 0 litellm/rust_bridge/responses/__init__.py | 0 .../websocket.py} | 0 litellm/rust_bridge/transcription/__init__.py | 0 .../native.py} | 0 .../test_rust_bridge_messages.py | 2 +- .../chat/test_anthropic_chat_handler.py | 4 +- .../chat/test_bedrock_converse_handler.py | 4 +- .../ocr/{test_legacy.py => test_main.py} | 4 +- .../ocr/test_ocr_native_format.py | 2 +- .../responses/test_rust_bridge_websocket.py | 3 +- tests/test_litellm/rust_bridge/__init__.py | 0 .../rust_bridge/chat_completions/__init__.py | 0 .../test_native.py} | 2 +- .../test_litellm/rust_bridge/ocr/__init__.py | 0 .../test_lifecycle.py} | 14 +- .../test_audio_transcription_rust_bridge.py | 2 +- tests/test_litellm_rust/ocr/test_lifecycle.py | 2 +- tests/test_litellm_rust/test_ocr.py | 2 +- 37 files changed, 516 insertions(+), 515 deletions(-) delete mode 100644 litellm/ocr/legacy.py create mode 100644 litellm/ocr/rust.py create mode 100644 litellm/rust_bridge/chat_completions/__init__.py rename litellm/rust_bridge/{chat_completions.py => chat_completions/native.py} (100%) create mode 100644 litellm/rust_bridge/messages/__init__.py rename litellm/rust_bridge/{messages.py => messages/native.py} (100%) create mode 100644 litellm/rust_bridge/ocr/__init__.py rename litellm/rust_bridge/{ocr_lifecycle.py => ocr/lifecycle.py} (97%) rename litellm/rust_bridge/{ocr.py => ocr/native.py} (100%) create mode 100644 litellm/rust_bridge/responses/__init__.py rename litellm/rust_bridge/{responses_websocket.py => responses/websocket.py} (100%) create mode 100644 litellm/rust_bridge/transcription/__init__.py rename litellm/rust_bridge/{transcription.py => transcription/native.py} (100%) rename tests/test_litellm/ocr/{test_legacy.py => test_main.py} (98%) create mode 100644 tests/test_litellm/rust_bridge/__init__.py create mode 100644 tests/test_litellm/rust_bridge/chat_completions/__init__.py rename tests/test_litellm/rust_bridge/{test_chat_completions.py => chat_completions/test_native.py} (99%) create mode 100644 tests/test_litellm/rust_bridge/ocr/__init__.py rename tests/test_litellm/rust_bridge/{test_ocr_lifecycle.py => ocr/test_lifecycle.py} (94%) diff --git a/litellm-rust/crates/python-bridge/src/routes/ocr/callbacks.rs b/litellm-rust/crates/python-bridge/src/routes/ocr/callbacks.rs index c7e5f123c19..0febedc01c3 100644 --- a/litellm-rust/crates/python-bridge/src/routes/ocr/callbacks.rs +++ b/litellm-rust/crates/python-bridge/src/routes/ocr/callbacks.rs @@ -159,7 +159,7 @@ fn redact( } pub(super) fn response(py: Python<'_>, response: &LiteLLMOcrResponse) -> PyResult> { - 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> { Ok(py - .import("litellm.rust_bridge.ocr_lifecycle")? + .import("litellm.rust_bridge.ocr.lifecycle")? .getattr("map_failure")? .call1((error, request, provider))? .extract()?) diff --git a/litellm/__init__.py b/litellm/__init__.py index dde94d68d5f..a56d988e801 100644 --- a/litellm/__init__.py +++ b/litellm/__init__.py @@ -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 * diff --git a/litellm/llms/anthropic/chat/handler.py b/litellm/llms/anthropic/chat/handler.py index 4dd0deeb62b..dff3a0be3fc 100644 --- a/litellm/llms/anthropic/chat/handler.py +++ b/litellm/llms/anthropic/chat/handler.py @@ -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, diff --git a/litellm/llms/bedrock/audio_transcription/__init__.py b/litellm/llms/bedrock/audio_transcription/__init__.py index 948d4280a4c..8f35b8eac7a 100644 --- a/litellm/llms/bedrock/audio_transcription/__init__.py +++ b/litellm/llms/bedrock/audio_transcription/__init__.py @@ -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, diff --git a/litellm/llms/bedrock/chat/converse_handler.py b/litellm/llms/bedrock/chat/converse_handler.py index d397420cb17..df8c4133450 100644 --- a/litellm/llms/bedrock/chat/converse_handler.py +++ b/litellm/llms/bedrock/chat/converse_handler.py @@ -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 diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index e049c62d28f..98e74ddce81 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -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, diff --git a/litellm/ocr/__init__.py b/litellm/ocr/__init__.py index a39141c0b5a..a171009564f 100644 --- a/litellm/ocr/__init__.py +++ b/litellm/ocr/__init__.py @@ -1,5 +1,5 @@ """OCR module for LiteLLM.""" -from .main import aocr, ocr +from .rust import aocr, ocr __all__ = ["aocr", "ocr"] diff --git a/litellm/ocr/input.py b/litellm/ocr/input.py index a58c7246128..eff91ca232a 100644 --- a/litellm/ocr/input.py +++ b/litellm/ocr/input.py @@ -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) diff --git a/litellm/ocr/legacy.py b/litellm/ocr/legacy.py deleted file mode 100644 index a742be274b3..00000000000 --- a/litellm/ocr/legacy.py +++ /dev/null @@ -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, - ) diff --git a/litellm/ocr/main.py b/litellm/ocr/main.py index faec3092d2b..a742be274b3 100644 --- a/litellm/ocr/main.py +++ b/litellm/ocr/main.py @@ -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, + ) diff --git a/litellm/ocr/rust.py b/litellm/ocr/rust.py new file mode 100644 index 00000000000..5f290e58d14 --- /dev/null +++ b/litellm/ocr/rust.py @@ -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) diff --git a/litellm/rust_bridge/_native.pyi b/litellm/rust_bridge/_native.pyi index e62c85f4599..a20bc1c0811 100644 --- a/litellm/rust_bridge/_native.pyi +++ b/litellm/rust_bridge/_native.pyi @@ -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"] diff --git a/litellm/rust_bridge/chat_completions/__init__.py b/litellm/rust_bridge/chat_completions/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/litellm/rust_bridge/chat_completions.py b/litellm/rust_bridge/chat_completions/native.py similarity index 100% rename from litellm/rust_bridge/chat_completions.py rename to litellm/rust_bridge/chat_completions/native.py diff --git a/litellm/rust_bridge/messages/__init__.py b/litellm/rust_bridge/messages/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/litellm/rust_bridge/messages.py b/litellm/rust_bridge/messages/native.py similarity index 100% rename from litellm/rust_bridge/messages.py rename to litellm/rust_bridge/messages/native.py diff --git a/litellm/rust_bridge/ocr/__init__.py b/litellm/rust_bridge/ocr/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/litellm/rust_bridge/ocr_lifecycle.py b/litellm/rust_bridge/ocr/lifecycle.py similarity index 97% rename from litellm/rust_bridge/ocr_lifecycle.py rename to litellm/rust_bridge/ocr/lifecycle.py index 4161007cce4..b3a022e46b3 100644 --- a/litellm/rust_bridge/ocr_lifecycle.py +++ b/litellm/rust_bridge/ocr/lifecycle.py @@ -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): diff --git a/litellm/rust_bridge/ocr.py b/litellm/rust_bridge/ocr/native.py similarity index 100% rename from litellm/rust_bridge/ocr.py rename to litellm/rust_bridge/ocr/native.py diff --git a/litellm/rust_bridge/responses/__init__.py b/litellm/rust_bridge/responses/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/litellm/rust_bridge/responses_websocket.py b/litellm/rust_bridge/responses/websocket.py similarity index 100% rename from litellm/rust_bridge/responses_websocket.py rename to litellm/rust_bridge/responses/websocket.py diff --git a/litellm/rust_bridge/transcription/__init__.py b/litellm/rust_bridge/transcription/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/litellm/rust_bridge/transcription.py b/litellm/rust_bridge/transcription/native.py similarity index 100% rename from litellm/rust_bridge/transcription.py rename to litellm/rust_bridge/transcription/native.py diff --git a/tests/test_litellm/anthropic_interface/test_rust_bridge_messages.py b/tests/test_litellm/anthropic_interface/test_rust_bridge_messages.py index 9e26d56d4d0..9f7f1bc86c7 100644 --- a/tests/test_litellm/anthropic_interface/test_rust_bridge_messages.py +++ b/tests/test_litellm/anthropic_interface/test_rust_bridge_messages.py @@ -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] = { diff --git a/tests/test_litellm/llms/anthropic/chat/test_anthropic_chat_handler.py b/tests/test_litellm/llms/anthropic/chat/test_anthropic_chat_handler.py index f854c2a0b71..e45d655ff7f 100644 --- a/tests/test_litellm/llms/anthropic/chat/test_anthropic_chat_handler.py +++ b/tests/test_litellm/llms/anthropic/chat/test_anthropic_chat_handler.py @@ -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": []} diff --git a/tests/test_litellm/llms/bedrock/chat/test_bedrock_converse_handler.py b/tests/test_litellm/llms/bedrock/chat/test_bedrock_converse_handler.py index 2fe92aead8f..49a73e857fd 100644 --- a/tests/test_litellm/llms/bedrock/chat/test_bedrock_converse_handler.py +++ b/tests/test_litellm/llms/bedrock/chat/test_bedrock_converse_handler.py @@ -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 diff --git a/tests/test_litellm/ocr/test_legacy.py b/tests/test_litellm/ocr/test_main.py similarity index 98% rename from tests/test_litellm/ocr/test_legacy.py rename to tests/test_litellm/ocr/test_main.py index 8b87690aedb..e0d2b5cfeb0 100644 --- a/tests/test_litellm/ocr/test_legacy.py +++ b/tests/test_litellm/ocr/test_main.py @@ -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 diff --git a/tests/test_litellm/ocr/test_ocr_native_format.py b/tests/test_litellm/ocr/test_ocr_native_format.py index 4ad556f6941..87d3faaf0fc 100644 --- a/tests/test_litellm/ocr/test_ocr_native_format.py +++ b/tests/test_litellm/ocr/test_ocr_native_format.py @@ -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(): diff --git a/tests/test_litellm/responses/test_rust_bridge_websocket.py b/tests/test_litellm/responses/test_rust_bridge_websocket.py index fcb5c5680ec..00ae5eb970f 100644 --- a/tests/test_litellm/responses/test_rust_bridge_websocket.py +++ b/tests/test_litellm/responses/test_rust_bridge_websocket.py @@ -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: diff --git a/tests/test_litellm/rust_bridge/__init__.py b/tests/test_litellm/rust_bridge/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/test_litellm/rust_bridge/chat_completions/__init__.py b/tests/test_litellm/rust_bridge/chat_completions/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/test_litellm/rust_bridge/test_chat_completions.py b/tests/test_litellm/rust_bridge/chat_completions/test_native.py similarity index 99% rename from tests/test_litellm/rust_bridge/test_chat_completions.py rename to tests/test_litellm/rust_bridge/chat_completions/test_native.py index b66cf1bfc63..14f8113924d 100644 --- a/tests/test_litellm/rust_bridge/test_chat_completions.py +++ b/tests/test_litellm/rust_bridge/chat_completions/test_native.py @@ -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 = { diff --git a/tests/test_litellm/rust_bridge/ocr/__init__.py b/tests/test_litellm/rust_bridge/ocr/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/test_litellm/rust_bridge/test_ocr_lifecycle.py b/tests/test_litellm/rust_bridge/ocr/test_lifecycle.py similarity index 94% rename from tests/test_litellm/rust_bridge/test_ocr_lifecycle.py rename to tests/test_litellm/rust_bridge/ocr/test_lifecycle.py index c61c5d79855..fd0a1591305 100644 --- a/tests/test_litellm/rust_bridge/test_ocr_lifecycle.py +++ b/tests/test_litellm/rust_bridge/ocr/test_lifecycle.py @@ -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: diff --git a/tests/test_litellm/test_audio_transcription_rust_bridge.py b/tests/test_litellm/test_audio_transcription_rust_bridge.py index c8c6627a898..48832528cc8 100644 --- a/tests/test_litellm/test_audio_transcription_rust_bridge.py +++ b/tests/test_litellm/test_audio_transcription_rust_bridge.py @@ -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") diff --git a/tests/test_litellm_rust/ocr/test_lifecycle.py b/tests/test_litellm_rust/ocr/test_lifecycle.py index dfcd63d3019..8eeee1941e9 100644 --- a/tests/test_litellm_rust/ocr/test_lifecycle.py +++ b/tests/test_litellm_rust/ocr/test_lifecycle.py @@ -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 diff --git a/tests/test_litellm_rust/test_ocr.py b/tests/test_litellm_rust/test_ocr.py index e0e06d685b8..7657eee2872 100644 --- a/tests/test_litellm_rust/test_ocr.py +++ b/tests/test_litellm_rust/test_ocr.py @@ -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