mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-19 00:01:29 +00:00
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:
parent
9484595fa2
commit
64f2a3d098
37 changed files with 516 additions and 515 deletions
|
|
@ -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()?)
|
||||
|
|
|
|||
|
|
@ -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 *
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -1,5 +1,5 @@
|
|||
"""OCR module for LiteLLM."""
|
||||
|
||||
from .main import aocr, ocr
|
||||
from .rust import aocr, ocr
|
||||
|
||||
__all__ = ["aocr", "ocr"]
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
)
|
||||
|
|
@ -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
83
litellm/ocr/rust.py
Normal 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)
|
||||
|
|
@ -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"]
|
||||
|
||||
|
|
|
|||
0
litellm/rust_bridge/chat_completions/__init__.py
Normal file
0
litellm/rust_bridge/chat_completions/__init__.py
Normal file
0
litellm/rust_bridge/messages/__init__.py
Normal file
0
litellm/rust_bridge/messages/__init__.py
Normal file
0
litellm/rust_bridge/ocr/__init__.py
Normal file
0
litellm/rust_bridge/ocr/__init__.py
Normal 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):
|
||||
0
litellm/rust_bridge/responses/__init__.py
Normal file
0
litellm/rust_bridge/responses/__init__.py
Normal file
0
litellm/rust_bridge/transcription/__init__.py
Normal file
0
litellm/rust_bridge/transcription/__init__.py
Normal 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] = {
|
||||
|
|
|
|||
|
|
@ -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": []}
|
||||
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
@ -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():
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
0
tests/test_litellm/rust_bridge/__init__.py
Normal file
0
tests/test_litellm/rust_bridge/__init__.py
Normal 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 = {
|
||||
0
tests/test_litellm/rust_bridge/ocr/__init__.py
Normal file
0
tests/test_litellm/rust_bridge/ocr/__init__.py
Normal 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:
|
||||
|
|
@ -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")
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue