diff --git a/litellm/ocr/main.py b/litellm/ocr/main.py index 8bd87ae21d3..80fde8ff9f3 100644 --- a/litellm/ocr/main.py +++ b/litellm/ocr/main.py @@ -336,7 +336,16 @@ async def aocr( ) ``` """ - local_vars = locals() + completion_kwargs: dict[str, object] = { + "model": model, + "document": document, + "api_key": api_key, + "api_base": api_base, + "timeout": timeout, + "custom_llm_provider": custom_llm_provider, + "extra_headers": extra_headers, + "kwargs": kwargs, + } try: prepared = _prepare_ocr_request( model=model, @@ -350,7 +359,9 @@ async def aocr( ) model = prepared.model custom_llm_provider = prepared.custom_llm_provider - local_vars.update({"model": model, "custom_llm_provider": custom_llm_provider}) + completion_kwargs.update( + {"model": model, "custom_llm_provider": custom_llm_provider} + ) if prepared.custom_llm_provider == "mistral" and rust_ocr_enabled(): rust_aocr = load_rust_aocr() @@ -397,139 +408,7 @@ async def aocr( model=model, custom_llm_provider=custom_llm_provider, original_exception=e, - completion_kwargs=local_vars, - extra_kwargs=kwargs, - ) - - -@client -def ocr( - model: str, - document: dict[str, Any], - api_key: Optional[str] = None, - api_base: Optional[str] = None, - timeout: Optional[Union[float, httpx.Timeout]] = None, - custom_llm_provider: Optional[str] = None, - extra_headers: Optional[dict[str, Any]] = None, - **kwargs, -) -> Union[OCRResponse, Coroutine[Any, Any, OCRResponse]]: - """ - Synchronous OCR function. - - Args: - model: Model name (e.g., "mistral/mistral-ocr-latest") - document: Document to process in Mistral format: - {"type": "document_url", "document_url": "https://..."} for PDFs/docs, - {"type": "image_url", "image_url": "https://..."} for images, or - {"type": "file", "file": } for local files - api_key: Optional API key - api_base: Optional API base URL - timeout: Optional timeout - custom_llm_provider: Optional custom LLM provider - extra_headers: Optional extra headers - **kwargs: Additional parameters (e.g., include_image_base64, pages, image_limit) - - Returns: - OCRResponse in Mistral OCR format with pages, model, usage_info, etc. - - Example: - ```python - import litellm - - # OCR with PDF - response = litellm.ocr( - model="mistral/mistral-ocr-latest", - document={ - "type": "document_url", - "document_url": "https://arxiv.org/pdf/2201.04234" - }, - include_image_base64=True - ) - - # OCR with image - response = litellm.ocr( - model="mistral/mistral-ocr-latest", - document={ - "type": "image_url", - "image_url": "https://example.com/image.png" - } - ) - - # OCR with base64 encoded PDF - response = litellm.ocr( - model="mistral/mistral-ocr-latest", - document={ - "type": "document_url", - "document_url": f"data:application/pdf;base64,{base64_pdf}" - } - ) - - # OCR with local file - response = litellm.ocr( - model="mistral/mistral-ocr-latest", - document={"type": "file", "file": "/path/to/document.pdf"} - ) - - # Access pages - for page in response.pages: - print(f"Page {page.index}: {page.markdown}") - ``` - """ - local_vars = locals() - try: - _is_async = kwargs.pop("aocr", False) is True - prepared = _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 - local_vars.update({"model": model, "custom_llm_provider": custom_llm_provider}) - - # Optional Rust path: hand the whole Mistral OCR call to the Rust bridge. - if prepared.custom_llm_provider == "mistral" and rust_ocr_enabled(): - rust_ocr = load_rust_ocr() - if rust_ocr is None: - verbose_logger.debug( - "Rust OCR bridge unavailable; falling back to Python path" - ) - else: - from litellm.secret_managers.main import get_secret_str - - return _run_rust_ocr( - rust_ocr=rust_ocr, - prepared_request=prepared, - resolve_api_key=get_secret_str, - ) - - response = base_llm_http_handler.ocr( - model=prepared.model, - document=prepared.document, - optional_params=prepared.optional_params, - timeout=prepared.effective_timeout, - logging_obj=prepared.litellm_logging_obj, - api_key=prepared.api_key, - api_base=prepared.api_base, - custom_llm_provider=prepared.custom_llm_provider, - aocr=_is_async, - headers=prepared.extra_headers, - provider_config=prepared.provider_config, - litellm_params=prepared.litellm_params, - ) - - return response - except Exception as e: - raise litellm.exception_type( - model=model, - custom_llm_provider=custom_llm_provider, - original_exception=e, - completion_kwargs=local_vars, + completion_kwargs=completion_kwargs, extra_kwargs=kwargs, ) @@ -651,9 +530,153 @@ def convert_file_document_to_url_document(document: dict[str, Any]) -> dict[str, f"(mime={mime_type}, size={len(file_bytes)} bytes, name={file_name})" ) return {"type": "image_url", "image_url": data_uri} - else: - verbose_logger.debug( - f"OCR file input: Converted file to document_url data URI " - f"(mime={mime_type}, size={len(file_bytes)} bytes, name={file_name})" + + verbose_logger.debug( + f"OCR file input: Converted file to document_url data URI " + f"(mime={mime_type}, size={len(file_bytes)} bytes, name={file_name})" + ) + return {"type": "document_url", "document_url": data_uri} + + +@client +def ocr( + model: str, + document: dict[str, Any], + api_key: Optional[str] = None, + api_base: Optional[str] = None, + timeout: Optional[Union[float, httpx.Timeout]] = None, + custom_llm_provider: Optional[str] = None, + extra_headers: Optional[dict[str, Any]] = None, + **kwargs, +) -> Union[OCRResponse, Coroutine[Any, Any, OCRResponse]]: + """ + Synchronous OCR function. + + Args: + model: Model name (e.g., "mistral/mistral-ocr-latest") + document: Document to process in Mistral format: + {"type": "document_url", "document_url": "https://..."} for PDFs/docs, + {"type": "image_url", "image_url": "https://..."} for images, or + {"type": "file", "file": } for local files + api_key: Optional API key + api_base: Optional API base URL + timeout: Optional timeout + custom_llm_provider: Optional custom LLM provider + extra_headers: Optional extra headers + **kwargs: Additional parameters (e.g., include_image_base64, pages, image_limit) + + Returns: + OCRResponse in Mistral OCR format with pages, model, usage_info, etc. + + Example: + ```python + import litellm + + # OCR with PDF + response = litellm.ocr( + model="mistral/mistral-ocr-latest", + document={ + "type": "document_url", + "document_url": "https://arxiv.org/pdf/2201.04234" + }, + include_image_base64=True + ) + + # OCR with image + response = litellm.ocr( + model="mistral/mistral-ocr-latest", + document={ + "type": "image_url", + "image_url": "https://example.com/image.png" + } + ) + + # OCR with base64 encoded PDF + response = litellm.ocr( + model="mistral/mistral-ocr-latest", + document={ + "type": "document_url", + "document_url": f"data:application/pdf;base64,{base64_pdf}" + } + ) + + # OCR with local file + response = litellm.ocr( + model="mistral/mistral-ocr-latest", + document={"type": "file", "file": "/path/to/document.pdf"} + ) + + # Access pages + for page in response.pages: + print(f"Page {page.index}: {page.markdown}") + ``` + """ + completion_kwargs: dict[str, object] = { + "model": model, + "document": document, + "api_key": api_key, + "api_base": api_base, + "timeout": timeout, + "custom_llm_provider": custom_llm_provider, + "extra_headers": extra_headers, + "kwargs": kwargs, + } + try: + _is_async = kwargs.pop("aocr", False) is True + completion_kwargs["aocr"] = _is_async + prepared = _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} + ) + + # Optional Rust path: hand the whole Mistral OCR call to the Rust bridge. + if prepared.custom_llm_provider == "mistral" and rust_ocr_enabled(): + rust_ocr = load_rust_ocr() + if rust_ocr is None: + verbose_logger.debug( + "Rust OCR bridge unavailable; falling back to Python path" + ) + else: + from litellm.secret_managers.main import get_secret_str + + return _run_rust_ocr( + rust_ocr=rust_ocr, + prepared_request=prepared, + resolve_api_key=get_secret_str, + ) + + response = base_llm_http_handler.ocr( + model=prepared.model, + document=prepared.document, + optional_params=prepared.optional_params, + timeout=prepared.effective_timeout, + logging_obj=prepared.litellm_logging_obj, + api_key=prepared.api_key, + api_base=prepared.api_base, + custom_llm_provider=prepared.custom_llm_provider, + aocr=_is_async, + headers=prepared.extra_headers, + provider_config=prepared.provider_config, + litellm_params=prepared.litellm_params, + ) + + return response + except Exception as e: + raise litellm.exception_type( + model=model, + custom_llm_provider=custom_llm_provider, + original_exception=e, + completion_kwargs=completion_kwargs, + extra_kwargs=kwargs, ) - return {"type": "document_url", "document_url": data_uri} diff --git a/litellm/proxy/ocr_endpoints/endpoints.py b/litellm/proxy/ocr_endpoints/endpoints.py index e32fee6afc5..823b4a5e86c 100644 --- a/litellm/proxy/ocr_endpoints/endpoints.py +++ b/litellm/proxy/ocr_endpoints/endpoints.py @@ -8,7 +8,6 @@ from fastapi import APIRouter, Depends, Request, Response, UploadFile from fastapi.responses import ORJSONResponse from litellm._logging import verbose_proxy_logger -from litellm.ocr.main import convert_file_document_to_url_document, get_mime_type from litellm.proxy._types import * from litellm.proxy.auth.user_api_key_auth import UserAPIKeyAuth, user_api_key_auth from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing @@ -27,6 +26,8 @@ def _build_document_from_upload( Delegates to convert_file_document_to_url_document after resolving MIME type from the upload's content_type header or filename. """ + from litellm.ocr.main import convert_file_document_to_url_document, get_mime_type + mime_type = content_type.split(";")[0].strip() if content_type else None if not mime_type or mime_type == "application/octet-stream": if filename: