From bc26845ec4f000404e0c00b4abe0fe975e59aef2 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Wed, 15 Oct 2025 17:20:01 -0700 Subject: [PATCH] [Feat] Native /ocr endpoint support (#15573) * [Feat] Add native litellm.ocr() functions (#15567) * fix get_supported_ocr_params * add get_provider_ocr_config * init OCR * init ocr functions * add OCRResponse Base Model * add ocr to llm http handlers * add main.py for OCR * fix linting for OCR * TestMistralOCR * update to use DocumentType for Mistral * fix _prepare_ocr_request * fix transform * add main.py for OCR * add spec to init * fix OCR * TestMistralOCR * ruff fix * Potential fix for code scanning alert no. 3521: Clear-text logging of sensitive information Co-authored-by: Copilot Autofix powered by AI <62310815+github-advanced-security[bot]@users.noreply.github.com> --------- Co-authored-by: Copilot Autofix powered by AI <62310815+github-advanced-security[bot]@users.noreply.github.com> * [Feat] Add /ocr route on LiteLLM AI Gateway - Adds support for native mistral ocr calling (#15571) * fix get_supported_ocr_params * add get_provider_ocr_config * init OCR * init ocr functions * add OCRResponse Base Model * add ocr to llm http handlers * add main.py for OCR * fix linting for OCR * TestMistralOCR * update to use DocumentType for Mistral * fix _prepare_ocr_request * fix transform * add main.py for OCR * add spec to init * fix OCR * TestMistralOCR * ruff fix * add router.ocr() methods * add OCR routes * feat add ocr routes * add OCR routes * feat: add OCR routes in proxy server * working /ocr routes * test_router_aocr_with_mistral * docs Mistral OCR * docs OCR * [Feat] Add Azure AI Mistral OCR Integration (#15572) * fix get_supported_ocr_params * add get_provider_ocr_config * init OCR * init ocr functions * add OCRResponse Base Model * add ocr to llm http handlers * add main.py for OCR * fix linting for OCR * TestMistralOCR * update to use DocumentType for Mistral * fix _prepare_ocr_request * fix transform * add main.py for OCR * add spec to init * fix OCR * TestMistralOCR * ruff fix * add router.ocr() methods * add OCR routes * feat add ocr routes * add OCR routes * feat: add OCR routes in proxy server * working /ocr routes * test_router_aocr_with_mistral * docs Mistral OCR * docs OCR * add azure ai to get_provider_ocr_config * add AzureAIOCRConfig * TestAzureAIOCR * TestAzureAIOCR * test fixes for azure ai ocr * fix async OCR transform for Azure * fix transform_ocr_request --------- Co-authored-by: Copilot Autofix powered by AI <62310815+github-advanced-security[bot]@users.noreply.github.com> --- .circleci/config.yml | 53 ++- docs/my-website/docs/ocr.md | 257 +++++++++++++++ docs/my-website/sidebars.js | 1 + litellm/__init__.py | 1 + litellm/llms/azure_ai/ocr/__init__.py | 5 + litellm/llms/azure_ai/ocr/transformation.py | 268 ++++++++++++++++ litellm/llms/base_llm/ocr/__init__.py | 22 ++ litellm/llms/base_llm/ocr/transformation.py | 207 ++++++++++++ litellm/llms/custom_httpx/llm_http_handler.py | 285 +++++++++++++++++ litellm/llms/mistral/ocr/__init__.py | 2 + litellm/llms/mistral/ocr/transformation.py | 223 +++++++++++++ litellm/ocr/__init__.py | 5 + litellm/ocr/main.py | 301 ++++++++++++++++++ litellm/proxy/common_request_processing.py | 2 + litellm/proxy/ocr_endpoints/__init__.py | 2 + litellm/proxy/ocr_endpoints/endpoints.py | 97 ++++++ litellm/proxy/proxy_config.yaml | 14 +- litellm/proxy/proxy_server.py | 10 +- litellm/proxy/route_llm_request.py | 2 + litellm/router.py | 14 + litellm/utils.py | 21 ++ tests/ocr_tests/base_ocr_unit_tests.py | 140 ++++++++ tests/ocr_tests/test_ocr_azure_ai.py | 27 ++ tests/ocr_tests/test_ocr_mistral.py | 88 +++++ 24 files changed, 2032 insertions(+), 15 deletions(-) create mode 100644 docs/my-website/docs/ocr.md create mode 100644 litellm/llms/azure_ai/ocr/__init__.py create mode 100644 litellm/llms/azure_ai/ocr/transformation.py create mode 100644 litellm/llms/base_llm/ocr/__init__.py create mode 100644 litellm/llms/base_llm/ocr/transformation.py create mode 100644 litellm/llms/mistral/ocr/__init__.py create mode 100644 litellm/llms/mistral/ocr/transformation.py create mode 100644 litellm/ocr/__init__.py create mode 100644 litellm/ocr/main.py create mode 100644 litellm/proxy/ocr_endpoints/__init__.py create mode 100644 litellm/proxy/ocr_endpoints/endpoints.py create mode 100644 tests/ocr_tests/base_ocr_unit_tests.py create mode 100644 tests/ocr_tests/test_ocr_azure_ai.py create mode 100644 tests/ocr_tests/test_ocr_mistral.py diff --git a/.circleci/config.yml b/.circleci/config.yml index 0bfbbcb440..8ae399c5c5 100644 --- a/.circleci/config.yml +++ b/.circleci/config.yml @@ -1041,6 +1041,49 @@ jobs: paths: - llm_responses_api_coverage.xml - llm_responses_api_coverage + ocr_testing: + docker: + - image: cimg/python:3.11 + auth: + username: ${DOCKERHUB_USERNAME} + password: ${DOCKERHUB_PASSWORD} + working_directory: ~/project + + steps: + - checkout + - setup_google_dns + - run: + name: Install Dependencies + command: | + python -m pip install --upgrade pip + python -m pip install -r requirements.txt + pip install "pytest==7.3.1" + pip install "pytest-retry==1.6.3" + pip install "pytest-cov==5.0.0" + pip install "pytest-asyncio==0.21.1" + pip install "respx==0.22.0" + # Run pytest and generate JUnit XML report + - run: + name: Run tests + command: | + pwd + ls + python -m pytest -vv tests/ocr_tests --cov=litellm --cov-report=xml -x -s -v --junitxml=test-results/junit.xml --durations=5 + no_output_timeout: 120m + - run: + name: Rename the coverage files + command: | + mv coverage.xml ocr_coverage.xml + mv .coverage ocr_coverage + + # Store test results + - store_test_results: + path: test-results + - persist_to_workspace: + root: . + paths: + - ocr_coverage.xml + - ocr_coverage litellm_mapped_tests: docker: - image: cimg/python:3.11 @@ -2741,7 +2784,7 @@ jobs: python -m venv venv . venv/bin/activate pip install coverage - coverage combine llm_translation_coverage llm_responses_api_coverage mcp_coverage logging_coverage litellm_router_coverage local_testing_coverage litellm_assistants_api_coverage auth_ui_unit_tests_coverage langfuse_coverage caching_coverage litellm_proxy_unit_tests_coverage image_gen_coverage pass_through_unit_tests_coverage batches_coverage litellm_security_tests_coverage guardrails_coverage + coverage combine llm_translation_coverage llm_responses_api_coverage ocr_coverage mcp_coverage logging_coverage litellm_router_coverage local_testing_coverage litellm_assistants_api_coverage auth_ui_unit_tests_coverage langfuse_coverage caching_coverage litellm_proxy_unit_tests_coverage image_gen_coverage pass_through_unit_tests_coverage batches_coverage litellm_security_tests_coverage guardrails_coverage coverage xml - codecov/upload: file: ./coverage.xml @@ -3289,6 +3332,12 @@ workflows: only: - main - /litellm_.*/ + - ocr_testing: + filters: + branches: + only: + - main + - /litellm_.*/ - litellm_mapped_enterprise_tests: filters: branches: @@ -3338,6 +3387,7 @@ workflows: - google_generate_content_endpoint_testing - guardrails_testing - llm_responses_api_testing + - ocr_testing - litellm_mapped_tests - litellm_mapped_enterprise_tests - batches_testing @@ -3400,6 +3450,7 @@ workflows: - mcp_testing - google_generate_content_endpoint_testing - llm_responses_api_testing + - ocr_testing - litellm_mapped_tests - litellm_mapped_enterprise_tests - batches_testing diff --git a/docs/my-website/docs/ocr.md b/docs/my-website/docs/ocr.md new file mode 100644 index 0000000000..d966097b32 --- /dev/null +++ b/docs/my-website/docs/ocr.md @@ -0,0 +1,257 @@ +# /ocr + +:::tip + +LiteLLM follows the [Mistral API request/response for the OCR API](https://docs.mistral.ai/capabilities/vision/#optical-character-recognition-ocr) + +::: + +## **LiteLLM Python SDK Usage** +### Quick Start + +```python +from litellm import ocr +import os + +os.environ["MISTRAL_API_KEY"] = "sk-.." + +response = ocr( + model="mistral/mistral-ocr-latest", + document={ + "type": "document_url", + "document_url": "https://arxiv.org/pdf/2201.04234" + } +) + +# Access extracted text +for page in response.pages: + print(f"Page {page.index}:") + print(page.markdown) +``` + +### Async Usage + +```python +from litellm import aocr +import os, asyncio + +os.environ["MISTRAL_API_KEY"] = "sk-.." + +async def test_async_ocr(): + response = await aocr( + model="mistral/mistral-ocr-latest", + document={ + "type": "document_url", + "document_url": "https://arxiv.org/pdf/2201.04234" + } + ) + + # Access extracted text + for page in response.pages: + print(f"Page {page.index}:") + print(page.markdown) + +asyncio.run(test_async_ocr()) +``` + +### Using Base64 Encoded Documents + +```python +import base64 +from litellm import ocr + +# Encode PDF to base64 +with open("document.pdf", "rb") as f: + base64_pdf = base64.b64encode(f.read()).decode('utf-8') + +response = ocr( + model="mistral/mistral-ocr-latest", + document={ + "type": "document_url", + "document_url": f"data:application/pdf;base64,{base64_pdf}" + } +) +``` + +### Optional Parameters + +```python +response = ocr( + model="mistral/mistral-ocr-latest", + document={ + "type": "document_url", + "document_url": "https://example.com/doc.pdf" + }, + # Optional Mistral parameters + pages=[0, 1, 2], # Only process specific pages + include_image_base64=True, # Include extracted images + image_limit=10, # Max images to return + image_min_size=100 # Min image size to include +) +``` + +## **LiteLLM Proxy Usage** + +LiteLLM provides a Mistral API compatible `/ocr` endpoint for OCR calls. + +**Setup** + +Add this to your litellm proxy config.yaml + +```yaml +model_list: + - model_name: mistral-ocr + litellm_params: + model: mistral/mistral-ocr-latest + api_key: os.environ/MISTRAL_API_KEY +``` + +Start litellm + +```bash +litellm --config /path/to/config.yaml + +# RUNNING on http://0.0.0.0:4000 +``` + +Test request + +```bash +curl http://0.0.0.0:4000/v1/ocr \ + -H "Authorization: Bearer sk-1234" \ + -H "Content-Type: application/json" \ + -d '{ + "model": "mistral-ocr", + "document": { + "type": "document_url", + "document_url": "https://arxiv.org/pdf/2201.04234" + } + }' +``` + + +## **Request/Response Format** + +:::info + +LiteLLM follows the **Mistral OCR API specification**. + +See the [official Mistral OCR documentation](https://docs.mistral.ai/capabilities/vision/#optical-character-recognition-ocr) for complete details. + +::: + +### Example Request + +```python +{ + "model": "mistral/mistral-ocr-latest", + "document": { + "type": "document_url", + "document_url": "https://arxiv.org/pdf/2201.04234" + }, + "pages": [0, 1, 2], # Optional: specific pages to process + "include_image_base64": True, # Optional: include extracted images + "image_limit": 10, # Optional: max images to return + "image_min_size": 100 # Optional: min image size in pixels +} +``` + +### Request Parameters + +| Parameter | Type | Required | Description | +|-----------|------|----------|-------------| +| `model` | string | Yes | The OCR model to use (e.g., `"mistral/mistral-ocr-latest"`) | +| `document` | object | Yes | Document to process. Must contain `type` and URL field | +| `document.type` | string | Yes | Either `"document_url"` for PDFs/docs or `"image_url"` for images | +| `document.document_url` | string | Conditional | URL to the document (required if `type` is `"document_url"`) | +| `document.image_url` | string | Conditional | URL to the image (required if `type` is `"image_url"`) | +| `pages` | array | No | List of specific page indices to process (0-indexed) | +| `include_image_base64` | boolean | No | Whether to include extracted images as base64 strings | +| `image_limit` | integer | No | Maximum number of images to return | +| `image_min_size` | integer | No | Minimum size (in pixels) for images to include | + +#### Document Format Examples + +**For PDFs and documents:** +```json +{ + "type": "document_url", + "document_url": "https://example.com/document.pdf" +} +``` + +**For images:** +```json +{ + "type": "image_url", + "image_url": "https://example.com/image.png" +} +``` + +**For base64-encoded content:** +```json +{ + "type": "document_url", + "document_url": "data:application/pdf;base64,JVBERi0xLjQKJ..." +} +``` + +### Response Format + +The response follows Mistral's OCR format with the following structure: + +```json +{ + "pages": [ + { + "index": 0, + "markdown": "# Document Title\n\nExtracted text content...", + "dimensions": { + "dpi": 200, + "height": 2200, + "width": 1700 + }, + "images": [ + { + "image_base64": "base64string...", + "bbox": { + "x": 100, + "y": 200, + "width": 300, + "height": 400 + } + } + ] + } + ], + "model": "mistral-ocr-2505-completion", + "usage_info": { + "pages_processed": 29, + "doc_size_bytes": 3002783 + }, + "document_annotation": null, + "object": "ocr" +} +``` + +#### Response Fields + +| Field | Type | Description | +|-------|------|-------------| +| `pages` | array | List of processed pages with extracted content | +| `pages[].index` | integer | Page number (0-indexed) | +| `pages[].markdown` | string | Extracted text in Markdown format | +| `pages[].dimensions` | object | Page dimensions (dpi, height, width in pixels) | +| `pages[].images` | array | Extracted images from the page (if `include_image_base64=true`) | +| `model` | string | The model used for OCR processing | +| `usage_info` | object | Processing statistics (pages processed, document size) | +| `document_annotation` | object | Optional document-level annotations | +| `object` | string | Always `"ocr"` for OCR responses | + + +## **Supported Providers** + +| Provider | Link to Usage | +|-------------|--------------------| +| Mistral AI | [Usage](#quick-start) | + diff --git a/docs/my-website/sidebars.js b/docs/my-website/sidebars.js index 53577029e8..73a17f6c55 100644 --- a/docs/my-website/sidebars.js +++ b/docs/my-website/sidebars.js @@ -347,6 +347,7 @@ const sidebars = { ] }, "moderation", + "ocr", { type: "category", label: "Pass-through Endpoints (Anthropic SDK, etc.)", diff --git a/litellm/__init__.py b/litellm/__init__.py index e461c88efd..fc20eecbb0 100644 --- a/litellm/__init__.py +++ b/litellm/__init__.py @@ -1324,6 +1324,7 @@ from .batch_completion.main import * # type: ignore from .rerank_api.main import * from .llms.anthropic.experimental_pass_through.messages.handler import * from .responses.main import * +from .ocr.main import * from .realtime_api.main import _arealtime from .fine_tuning.main import * from .files.main import * diff --git a/litellm/llms/azure_ai/ocr/__init__.py b/litellm/llms/azure_ai/ocr/__init__.py new file mode 100644 index 0000000000..86f7e53d60 --- /dev/null +++ b/litellm/llms/azure_ai/ocr/__init__.py @@ -0,0 +1,5 @@ +"""Azure AI OCR module.""" +from .transformation import AzureAIOCRConfig + +__all__ = ["AzureAIOCRConfig"] + diff --git a/litellm/llms/azure_ai/ocr/transformation.py b/litellm/llms/azure_ai/ocr/transformation.py new file mode 100644 index 0000000000..eade2dd765 --- /dev/null +++ b/litellm/llms/azure_ai/ocr/transformation.py @@ -0,0 +1,268 @@ +""" +Azure AI OCR transformation implementation. +""" +from typing import Dict, Optional + +from litellm._logging import verbose_logger +from litellm.litellm_core_utils.prompt_templates.image_handling import ( + async_convert_url_to_base64, + convert_url_to_base64, +) +from litellm.llms.base_llm.ocr.transformation import DocumentType, OCRRequestData +from litellm.llms.mistral.ocr.transformation import MistralOCRConfig +from litellm.secret_managers.main import get_secret_str + + +class AzureAIOCRConfig(MistralOCRConfig): + """ + Azure AI OCR transformation configuration. + + Azure AI uses Mistral's OCR API but with a different endpoint format. + Inherits transformation logic from MistralOCRConfig since they use the same format. + + Reference: Azure AI Foundry OCR documentation + + Important: Azure AI only supports base64 data URIs (data:image/..., data:application/pdf;base64,...). + Regular URLs are not supported. + """ + + def __init__(self) -> None: + super().__init__() + + def validate_environment( + self, + headers: Dict, + model: str, + api_key: Optional[str] = None, + api_base: Optional[str] = None, + **kwargs, + ) -> Dict: + """ + Validate environment and return headers for Azure AI OCR. + + Azure AI uses Bearer token authentication with AZURE_AI_API_KEY. + """ + # Get API key from environment if not provided + if api_key is None: + api_key = get_secret_str("AZURE_AI_API_KEY") + + if api_key is None: + raise ValueError( + "Missing Azure AI API Key - A call is being made to Azure AI but no key is set either in the environment variables or via params" + ) + + # Validate API base is provided + if api_base is None: + api_base = get_secret_str("AZURE_AI_API_BASE") + + if api_base is None: + raise ValueError( + "Missing Azure AI API Base - Set AZURE_AI_API_BASE environment variable or pass api_base parameter" + ) + + headers = { + "Authorization": f"Bearer {api_key}", + "Content-Type": "application/json", + **headers, + } + + return headers + + def get_complete_url( + self, + api_base: Optional[str], + model: str, + optional_params: dict, + **kwargs, + ) -> str: + """ + Get complete URL for Azure AI OCR endpoint. + + Azure AI endpoint format: https:///providers/mistral/azure/ocr + + Args: + api_base: Azure AI API base URL + model: Model name (not used in URL construction) + optional_params: Optional parameters + + Returns: Complete URL for Azure AI OCR endpoint + """ + if api_base is None: + raise ValueError( + "Missing Azure AI API Base - Set AZURE_AI_API_BASE environment variable or pass api_base parameter" + ) + + # Ensure no trailing slash + api_base = api_base.rstrip("/") + + # Azure AI OCR endpoint format + return f"{api_base}/providers/mistral/azure/ocr" + + def _convert_url_to_data_uri_sync(self, url: str) -> str: + """ + Synchronously convert a URL to a base64 data URI. + + Azure AI OCR doesn't have internet access, so we need to fetch URLs + and convert them to base64 data URIs. + + Args: + url: The URL to convert + + Returns: + Base64 data URI string + """ + verbose_logger.debug(f"Azure AI OCR: Converting URL to base64 data URI (sync): {url}") + + # Fetch and convert to base64 data URI + # convert_url_to_base64 already returns a full data URI like "data:image/jpeg;base64,..." + data_uri = convert_url_to_base64(url=url) + + verbose_logger.debug(f"Azure AI OCR: Converted URL to data URI (length: {len(data_uri)})") + + return data_uri + + async def _convert_url_to_data_uri_async(self, url: str) -> str: + """ + Asynchronously convert a URL to a base64 data URI. + + Azure AI OCR doesn't have internet access, so we need to fetch URLs + and convert them to base64 data URIs. + + Args: + url: The URL to convert + + Returns: + Base64 data URI string + """ + verbose_logger.debug(f"Azure AI OCR: Converting URL to base64 data URI (async): {url}") + + # Fetch and convert to base64 data URI asynchronously + # async_convert_url_to_base64 already returns a full data URI like "data:image/jpeg;base64,..." + data_uri = await async_convert_url_to_base64(url=url) + + verbose_logger.debug(f"Azure AI OCR: Converted URL to data URI (length: {len(data_uri)})") + + return data_uri + + def transform_ocr_request( + self, + model: str, + document: DocumentType, + optional_params: dict, + headers: dict, + **kwargs, + ) -> OCRRequestData: + """ + Transform OCR request for Azure AI, converting URLs to base64 data URIs (sync). + + Azure AI OCR doesn't have internet access, so we automatically fetch + any URLs and convert them to base64 data URIs synchronously. + + Args: + model: Model name + document: Document dict from user + optional_params: Already mapped optional parameters + headers: Request headers + **kwargs: Additional arguments + + Returns: + OCRRequestData with JSON data + """ + verbose_logger.debug(f"Azure AI OCR transform_ocr_request (sync) - model: {model}") + + if not isinstance(document, dict): + raise ValueError(f"Expected document dict, got {type(document)}") + + # Check if we need to convert URL to base64 + doc_type = document.get("type") + transformed_document = document.copy() + + if doc_type == "document_url": + document_url = document.get("document_url", "") + # If it's not already a data URI, convert it + if document_url and not document_url.startswith("data:"): + verbose_logger.debug( + "Azure AI OCR: Converting document URL to base64 data URI (sync)" + ) + data_uri = self._convert_url_to_data_uri_sync(url=document_url) + transformed_document["document_url"] = data_uri + elif doc_type == "image_url": + image_url = document.get("image_url", "") + # If it's not already a data URI, convert it + if image_url and not image_url.startswith("data:"): + verbose_logger.debug( + "Azure AI OCR: Converting image URL to base64 data URI (sync)" + ) + data_uri = self._convert_url_to_data_uri_sync(url=image_url) + transformed_document["image_url"] = data_uri + + # Call parent's transform to build the request + return super().transform_ocr_request( + model=model, + document=transformed_document, + optional_params=optional_params, + headers=headers, + **kwargs, + ) + + async def async_transform_ocr_request( + self, + model: str, + document: DocumentType, + optional_params: dict, + headers: dict, + **kwargs, + ) -> OCRRequestData: + """ + Transform OCR request for Azure AI, converting URLs to base64 data URIs (async). + + Azure AI OCR doesn't have internet access, so we automatically fetch + any URLs and convert them to base64 data URIs asynchronously. + + Args: + model: Model name + document: Document dict from user + optional_params: Already mapped optional parameters + headers: Request headers + **kwargs: Additional arguments + + Returns: + OCRRequestData with JSON data + """ + verbose_logger.debug(f"Azure AI OCR async_transform_ocr_request - model: {model}") + + if not isinstance(document, dict): + raise ValueError(f"Expected document dict, got {type(document)}") + + # Check if we need to convert URL to base64 + doc_type = document.get("type") + transformed_document = document.copy() + + if doc_type == "document_url": + document_url = document.get("document_url", "") + # If it's not already a data URI, convert it + if document_url and not document_url.startswith("data:"): + verbose_logger.debug( + "Azure AI OCR: Converting document URL to base64 data URI (async)" + ) + data_uri = await self._convert_url_to_data_uri_async(url=document_url) + transformed_document["document_url"] = data_uri + elif doc_type == "image_url": + image_url = document.get("image_url", "") + # If it's not already a data URI, convert it + if image_url and not image_url.startswith("data:"): + verbose_logger.debug( + "Azure AI OCR: Converting image URL to base64 data URI (async)" + ) + data_uri = await self._convert_url_to_data_uri_async(url=image_url) + transformed_document["image_url"] = data_uri + + # Call parent's transform to build the request + return super().transform_ocr_request( + model=model, + document=transformed_document, + optional_params=optional_params, + headers=headers, + **kwargs, + ) + diff --git a/litellm/llms/base_llm/ocr/__init__.py b/litellm/llms/base_llm/ocr/__init__.py new file mode 100644 index 0000000000..5965af5f2b --- /dev/null +++ b/litellm/llms/base_llm/ocr/__init__.py @@ -0,0 +1,22 @@ +"""Base OCR transformation module.""" +from .transformation import ( + BaseOCRConfig, + DocumentType, + OCRPage, + OCRPageDimensions, + OCRPageImage, + OCRRequestData, + OCRResponse, + OCRUsageInfo, +) + +__all__ = [ + "BaseOCRConfig", + "DocumentType", + "OCRResponse", + "OCRPage", + "OCRPageDimensions", + "OCRPageImage", + "OCRUsageInfo", + "OCRRequestData", +] diff --git a/litellm/llms/base_llm/ocr/transformation.py b/litellm/llms/base_llm/ocr/transformation.py new file mode 100644 index 0000000000..41d7d31e6b --- /dev/null +++ b/litellm/llms/base_llm/ocr/transformation.py @@ -0,0 +1,207 @@ +""" +Base OCR transformation configuration. +""" +from typing import TYPE_CHECKING, Any, Dict, List, Optional, Union + +import httpx +from pydantic import BaseModel + +from litellm.llms.base_llm.chat.transformation import BaseLLMException + +if TYPE_CHECKING: + from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj +else: + LiteLLMLoggingObj = Any + + +# DocumentType for OCR - Mistral format document dict +DocumentType = Dict[str, str] + + +class OCRPageDimensions(BaseModel): + """Page dimensions from OCR response.""" + dpi: Optional[int] = None + height: Optional[int] = None + width: Optional[int] = None + + +class OCRPageImage(BaseModel): + """Image extracted from OCR page.""" + image_base64: Optional[str] = None + bbox: Optional[Dict[str, Any]] = None + + model_config = {"extra": "allow"} + + +class OCRPage(BaseModel): + """Single page from OCR response.""" + index: int + markdown: str + images: Optional[List[OCRPageImage]] = None + dimensions: Optional[OCRPageDimensions] = None + + model_config = {"extra": "allow"} + + +class OCRUsageInfo(BaseModel): + """Usage information from OCR response.""" + pages_processed: Optional[int] = None + doc_size_bytes: Optional[int] = None + + model_config = {"extra": "allow"} + + +class OCRResponse(BaseModel): + """ + Standard OCR response format. + Standardized to Mistral OCR format - other providers should transform to this format. + """ + pages: List[OCRPage] + model: str + document_annotation: Optional[Any] = None + usage_info: Optional[OCRUsageInfo] = None + object: str = "ocr" + + model_config = {"extra": "allow"} + + +class OCRRequestData(BaseModel): + """OCR request data structure.""" + data: Optional[Union[Dict, bytes]] = None + files: Optional[Dict[str, Any]] = None + + +class BaseOCRConfig: + """ + Base configuration for OCR transformations. + Handles provider-agnostic OCR operations. + """ + + def __init__(self) -> None: + pass + + def get_supported_ocr_params(self, model: str) -> list: + """ + Get supported OCR parameters for this provider. + Override this method in provider-specific implementations. + """ + return [] + + def map_ocr_params( + self, + non_default_params: dict, + optional_params: dict, + model: str, + ) -> dict: + """Map OCR parameters to provider-specific parameters.""" + return optional_params + + def validate_environment( + self, + headers: Dict, + model: str, + api_key: Optional[str] = None, + api_base: Optional[str] = None, + **kwargs, + ) -> Dict: + """ + Validate environment and return headers. + Override in provider-specific implementations. + """ + return headers + + def get_complete_url( + self, + api_base: Optional[str], + model: str, + optional_params: dict, + **kwargs, + ) -> str: + """ + Get complete URL for OCR endpoint. + Override in provider-specific implementations. + """ + raise NotImplementedError("get_complete_url must be implemented by provider") + + def transform_ocr_request( + self, + model: str, + document: DocumentType, + optional_params: dict, + headers: dict, + **kwargs, + ) -> OCRRequestData: + """ + Transform OCR request to provider-specific format. + Override in provider-specific implementations. + + Args: + model: Model name + document: Document to process (Mistral format dict, or file path, bytes, etc.) + optional_params: Optional parameters for the request + headers: Request headers + + Returns: + OCRRequestData with data and files fields + """ + raise NotImplementedError("transform_ocr_request must be implemented by provider") + + async def async_transform_ocr_request( + self, + model: str, + document: DocumentType, + optional_params: dict, + headers: dict, + **kwargs, + ) -> OCRRequestData: + """ + Async transform OCR request to provider-specific format. + Optional method - providers can override if they need async transformations + (e.g., Azure AI for URL-to-base64 conversion). + + Default implementation falls back to sync transform_ocr_request. + + Args: + model: Model name + document: Document to process (Mistral format dict, or file path, bytes, etc.) + optional_params: Optional parameters for the request + headers: Request headers + + Returns: + OCRRequestData with data and files fields + """ + # Default implementation: call sync version + return self.transform_ocr_request( + model=model, + document=document, + optional_params=optional_params, + headers=headers, + **kwargs, + ) + + def transform_ocr_response( + self, + model: str, + raw_response: httpx.Response, + logging_obj: LiteLLMLoggingObj, + **kwargs, + ) -> OCRResponse: + """ + Transform provider-specific OCR response to standard format. + Override in provider-specific implementations. + """ + raise NotImplementedError("transform_ocr_response must be implemented by provider") + + def get_error_class( + self, + error_message: str, + status_code: int, + headers: dict, + ) -> Exception: + """Get appropriate error class for the provider.""" + return BaseLLMException( + status_code=status_code, + message=error_message, + headers=headers, + ) + diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index 5037f4d8d4..d7b7987b67 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -39,6 +39,7 @@ from litellm.llms.base_llm.image_edit.transformation import BaseImageEditConfig from litellm.llms.base_llm.image_generation.transformation import ( BaseImageGenerationConfig, ) +from litellm.llms.base_llm.ocr.transformation import BaseOCRConfig, OCRResponse from litellm.llms.base_llm.realtime.transformation import BaseRealtimeConfig from litellm.llms.base_llm.rerank.transformation import BaseRerankConfig from litellm.llms.base_llm.responses.transformation import BaseResponsesAPIConfig @@ -1256,6 +1257,289 @@ class BaseLLMHTTPHandler: api_key=api_key, ) + def _prepare_ocr_request( + self, + model: str, + document: Dict[str, str], + optional_params: dict, + logging_obj: LiteLLMLoggingObj, + api_key: Optional[str], + api_base: Optional[str], + headers: Optional[Dict[str, Any]], + provider_config: BaseOCRConfig, + litellm_params: dict, + ) -> Tuple[Dict[str, Any], str, Dict[str, Any], None]: + """ + Shared logic for preparing OCR requests. + Returns: (headers, complete_url, data, files) + """ + from litellm.llms.base_llm.ocr.transformation import OCRRequestData + + headers = provider_config.validate_environment( + api_key=api_key, + api_base=api_base, + headers=headers or {}, + model=model, + ) + + complete_url = provider_config.get_complete_url( + api_base=api_base, + model=model, + optional_params=optional_params, + ) + + # Transform the request to get data and files + transformed_result = provider_config.transform_ocr_request( + model=model, + document=document, + optional_params=optional_params, + headers=headers, + ) + + # All providers return OCRRequestData + if not isinstance(transformed_result, OCRRequestData): + raise ValueError( + f"Provider {provider_config.__class__.__name__} must return OCRRequestData" + ) + + # Data is always a dict for Mistral OCR format + if not isinstance(transformed_result.data, dict): + raise ValueError(f"Expected dict data for OCR request, got {type(transformed_result.data)}") + + data = transformed_result.data + + ## LOGGING + logging_obj.pre_call( + input="OCR document processing", + api_key=api_key, + additional_args={ + "complete_input_dict": data, + "api_base": complete_url, + "headers": headers, + }, + ) + + return headers, complete_url, data, None + + async def _async_prepare_ocr_request( + self, + model: str, + document: Dict[str, str], + optional_params: dict, + logging_obj: LiteLLMLoggingObj, + api_key: Optional[str], + api_base: Optional[str], + headers: Optional[Dict[str, Any]], + provider_config: BaseOCRConfig, + litellm_params: dict, + ) -> Tuple[Dict[str, Any], str, Dict[str, Any], None]: + """ + Async version of _prepare_ocr_request for providers that need async transforms. + Returns: (headers, complete_url, data, files) + """ + from litellm.llms.base_llm.ocr.transformation import OCRRequestData + + headers = provider_config.validate_environment( + api_key=api_key, + api_base=api_base, + headers=headers or {}, + model=model, + ) + + complete_url = provider_config.get_complete_url( + api_base=api_base, + model=model, + optional_params=optional_params, + ) + + # Use async transform (providers can override this method if they need async operations) + transformed_result = await provider_config.async_transform_ocr_request( + model=model, + document=document, + optional_params=optional_params, + headers=headers, + ) + + # All providers return OCRRequestData + if not isinstance(transformed_result, OCRRequestData): + raise ValueError( + f"Provider {provider_config.__class__.__name__} must return OCRRequestData" + ) + + # Data is always a dict for Mistral OCR format + if not isinstance(transformed_result.data, dict): + raise ValueError(f"Expected dict data for OCR request, got {type(transformed_result.data)}") + + data = transformed_result.data + + ## LOGGING + logging_obj.pre_call( + input="OCR document processing", + api_key=api_key, + additional_args={ + "complete_input_dict": data, + "api_base": complete_url, + "headers": headers, + }, + ) + + return headers, complete_url, data, None + + def _transform_ocr_response( + self, + provider_config: BaseOCRConfig, + model: str, + response: httpx.Response, + logging_obj: LiteLLMLoggingObj, + ) -> OCRResponse: + """Shared logic for transforming OCR responses.""" + return provider_config.transform_ocr_response( + model=model, + raw_response=response, + logging_obj=logging_obj, + ) + + def ocr( + self, + model: str, + document: Dict[str, str], + optional_params: dict, + timeout: Union[float, httpx.Timeout], + logging_obj: LiteLLMLoggingObj, + api_key: Optional[str], + api_base: Optional[str], + custom_llm_provider: str, + client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, + aocr: bool = False, + headers: Optional[Dict[str, Any]] = None, + provider_config: Optional[BaseOCRConfig] = None, + litellm_params: Optional[dict] = None, + ) -> Union[OCRResponse, Coroutine[Any, Any, OCRResponse]]: + """ + Sync OCR handler. + """ + if provider_config is None: + raise ValueError( + f"No provider config found for model: {model} and provider: {custom_llm_provider}" + ) + + if litellm_params is None: + litellm_params = {} + + if aocr is True: + return self.async_ocr( + model=model, + document=document, + optional_params=optional_params, + timeout=timeout, + logging_obj=logging_obj, + api_key=api_key, + api_base=api_base, + custom_llm_provider=custom_llm_provider, + client=client, + headers=headers, + provider_config=provider_config, + litellm_params=litellm_params, + ) + + # Prepare the request + headers, complete_url, data, files = self._prepare_ocr_request( + model=model, + document=document, + optional_params=optional_params, + logging_obj=logging_obj, + api_key=api_key, + api_base=api_base, + headers=headers, + provider_config=provider_config, + litellm_params=litellm_params, + ) + + if client is None or not isinstance(client, HTTPHandler): + client = _get_httpx_client() + + try: + # Make the POST request with JSON data (Mistral format) + response = client.post( + url=complete_url, + headers=headers, + json=data, + timeout=timeout, + ) + except Exception as e: + raise self._handle_error(e=e, provider_config=provider_config) + + return self._transform_ocr_response( + provider_config=provider_config, + model=model, + response=response, + logging_obj=logging_obj, + ) + + async def async_ocr( + self, + model: str, + document: Dict[str, str], + optional_params: dict, + timeout: Union[float, httpx.Timeout], + logging_obj: LiteLLMLoggingObj, + api_key: Optional[str], + api_base: Optional[str], + custom_llm_provider: str, + client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, + headers: Optional[Dict[str, Any]] = None, + provider_config: Optional[BaseOCRConfig] = None, + litellm_params: Optional[dict] = None, + ) -> OCRResponse: + """ + Async OCR handler. + """ + if provider_config is None: + raise ValueError( + f"No provider config found for model: {model} and provider: {custom_llm_provider}" + ) + + if litellm_params is None: + litellm_params = {} + + # Prepare the request using async prepare method + headers, complete_url, data, files = await self._async_prepare_ocr_request( + model=model, + document=document, + optional_params=optional_params, + logging_obj=logging_obj, + api_key=api_key, + api_base=api_base, + headers=headers, + provider_config=provider_config, + litellm_params=litellm_params, + ) + + if client is None or not isinstance(client, AsyncHTTPHandler): + async_httpx_client = get_async_httpx_client( + llm_provider=litellm.LlmProviders(custom_llm_provider), + ) + else: + async_httpx_client = client + + try: + # Make the async POST request with JSON data (Mistral format) + response = await async_httpx_client.post( + url=complete_url, + headers=headers, + json=data, + timeout=timeout, + ) + except Exception as e: + raise self._handle_error(e=e, provider_config=provider_config) + + return self._transform_ocr_response( + provider_config=provider_config, + model=model, + response=response, + logging_obj=logging_obj, + ) + async def async_anthropic_messages_handler( self, model: str, @@ -2995,6 +3279,7 @@ class BaseLLMHTTPHandler: BaseGoogleGenAIGenerateContentConfig, BaseAnthropicMessagesConfig, BaseBatchesConfig, + BaseOCRConfig, "BasePassthroughConfig", ], ): diff --git a/litellm/llms/mistral/ocr/__init__.py b/litellm/llms/mistral/ocr/__init__.py new file mode 100644 index 0000000000..40cc62696b --- /dev/null +++ b/litellm/llms/mistral/ocr/__init__.py @@ -0,0 +1,2 @@ +"""Mistral OCR transformation module.""" + diff --git a/litellm/llms/mistral/ocr/transformation.py b/litellm/llms/mistral/ocr/transformation.py new file mode 100644 index 0000000000..f17c872f53 --- /dev/null +++ b/litellm/llms/mistral/ocr/transformation.py @@ -0,0 +1,223 @@ +""" +Mistral OCR transformation implementation. +""" +from typing import Any, Dict, Optional + +import httpx + +from litellm._logging import verbose_logger +from litellm.llms.base_llm.ocr.transformation import ( + BaseOCRConfig, + DocumentType, + OCRRequestData, + OCRResponse, +) +from litellm.secret_managers.main import get_secret_str + + +class MistralOCRConfig(BaseOCRConfig): + """ + Mistral OCR transformation configuration. + + Reference: https://docs.mistral.ai/api/#tag/ocr + """ + + def __init__(self) -> None: + super().__init__() + + def get_supported_ocr_params(self, model: str) -> list: + """ + Get supported OCR parameters for Mistral OCR. + + Mistral OCR supports: + - pages: List of page numbers to process + - include_image_base64: Whether to include base64 encoded images + - image_limit: Maximum number of images to return + - image_min_size: Minimum size of images to include + - bbox_annotation_format: Format for bounding box annotations + - document_annotation_format: Format for document annotations + """ + return [ + "pages", + "include_image_base64", + "image_limit", + "image_min_size", + "bbox_annotation_format", + "document_annotation_format", + ] + + def map_ocr_params( + self, + non_default_params: dict, + optional_params: dict, + model: str, + ) -> dict: + """ + Map OCR parameters to Mistral-specific format. + + Mistral accepts these parameters directly, so no transformation needed. + Just filter out unsupported params. + """ + supported_params = self.get_supported_ocr_params(model=model) + + # Only include params that are in the supported list + mapped_params = {} + for param, value in non_default_params.items(): + if param in supported_params: + mapped_params[param] = value + + return mapped_params + + def validate_environment( + self, + headers: Dict, + model: str, + api_key: Optional[str] = None, + api_base: Optional[str] = None, + **kwargs, + ) -> Dict: + """ + Validate environment and return headers for Mistral OCR. + """ + # Get API key from environment if not provided + if api_key is None: + api_key = ( + get_secret_str("MISTRAL_API_KEY") + ) + + if api_key is None: + raise ValueError( + "Missing Mistral API Key - A call is being made to Mistral but no key is set either in the environment variables or via params" + ) + + headers = { + "Authorization": f"Bearer {api_key}", + **headers, + } + + # Don't set Content-Type for multipart/form-data - httpx will handle it + + return headers + + def get_complete_url( + self, + api_base: Optional[str], + model: str, + optional_params: dict, + **kwargs, + ) -> str: + """ + Get complete URL for Mistral OCR endpoint. + + Returns: https://api.mistral.ai/v1/ocr + """ + if api_base is None: + api_base = "https://api.mistral.ai/v1" + + # Ensure no trailing slash + api_base = api_base.rstrip("/") + + # Remove /v1 if it's already in the base to avoid duplication + if api_base.endswith("/v1"): + return f"{api_base}/ocr" + + return f"{api_base}/v1/ocr" + + + def transform_ocr_request( + self, + model: str, + document: DocumentType, + optional_params: dict, + headers: dict, + **kwargs, + ) -> OCRRequestData: + """ + Transform OCR request to Mistral-specific format. + + Mistral OCR API accepts: + { + "model": "mistral-ocr-latest", + "document": { + "type": "document_url", + "document_url": "" + }, + "pages": [0], # optional + "include_image_base64": false, # optional + ... + } + + Args: + model: Model name (e.g., "mistral-ocr-latest") + document: Document dict from user (Mistral format) - already validated in main.py + optional_params: Already mapped optional parameters + headers: Request headers + + Returns: + OCRRequestData with JSON data + """ + verbose_logger.debug(f"Mistral OCR transform_ocr_request - model: {model}") + + # Document parameter is the Mistral-format dict from the user + # Just pass it through as-is to the Mistral API + if not isinstance(document, dict): + raise ValueError(f"Expected document dict, got {type(document)}") + + # Build request data - use document dict directly + data = { + "model": model, + "document": document, # Pass through the Mistral-format document dict + } + + # Add all optional parameters from the already-mapped optional_params + data.update(optional_params) + + # No multipart files - using JSON + return OCRRequestData(data=data, files=None) + + def transform_ocr_response( + self, + model: str, + raw_response: httpx.Response, + logging_obj: Any, + **kwargs, + ) -> OCRResponse: + """ + Return Mistral OCR response in native format. + + Mistral OCR is the standard format for LiteLLM OCR responses. + No transformation needed - return native response. + + Mistral OCR returns: + { + "pages": [ + { + "index": 0, + "markdown": "extracted text content", + "images": [...], + "dimensions": {...} + }, + ... + ], + "model": "mistral-ocr-2505-completion", + "document_annotation": null, + "usage_info": {...} + } + """ + try: + response_json = raw_response.json() + + verbose_logger.debug(f"Mistral OCR response keys: {response_json.keys()}") + + # Return native Mistral format - no transformation + return OCRResponse( + pages=response_json.get("pages", []), + model=response_json.get("model", model), + document_annotation=response_json.get("document_annotation"), + usage_info=response_json.get("usage_info"), + object="ocr", + ) + except Exception as e: + verbose_logger.error(f"Error parsing Mistral OCR response: {e}") + raise e + diff --git a/litellm/ocr/__init__.py b/litellm/ocr/__init__.py new file mode 100644 index 0000000000..53f455619d --- /dev/null +++ b/litellm/ocr/__init__.py @@ -0,0 +1,5 @@ +"""OCR module for LiteLLM.""" +from .main import aocr, ocr + +__all__ = ["ocr", "aocr"] + diff --git a/litellm/ocr/main.py b/litellm/ocr/main.py new file mode 100644 index 0000000000..62172b0fba --- /dev/null +++ b/litellm/ocr/main.py @@ -0,0 +1,301 @@ +""" +Main OCR function for LiteLLM. +""" +import asyncio +import contextvars +from functools import partial +from typing import Any, Coroutine, Dict, Optional, Union + +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 BaseOCRConfig, OCRResponse +from litellm.llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler +from litellm.utils import ProviderConfigManager, client + +####### ENVIRONMENT VARIABLES ################### +base_llm_http_handler = BaseLLMHTTPHandler() +################################################# + + +@client +async def aocr( + model: str, + document: Dict[str, str], + 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, +) -> OCRResponse: + """ + Async 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 or + {"type": "image_url", "image_url": "https://..."} for images + 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 = await litellm.aocr( + 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 = await litellm.aocr( + model="mistral/mistral-ocr-latest", + document={ + "type": "image_url", + "image_url": "https://example.com/image.png" + } + ) + + # OCR with base64 encoded PDF + response = await litellm.aocr( + model="mistral/mistral-ocr-latest", + document={ + "type": "document_url", + "document_url": f"data:application/pdf;base64,{base64_pdf}" + } + ) + ``` + """ + local_vars = locals() + try: + loop = asyncio.get_event_loop() + kwargs["aocr"] = True + + # Get custom llm provider + if custom_llm_provider is None: + _, custom_llm_provider, _, _ = litellm.get_llm_provider( + model=model, api_base=api_base + ) + + func = partial( + ocr, + 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, + ) + + ctx = contextvars.copy_context() + func_with_context = partial(ctx.run, func) + init_response = await loop.run_in_executor(None, func_with_context) + + if asyncio.iscoroutine(init_response): + response = await init_response + else: + response = init_response + + if response is None: + raise ValueError( + f"Got an unexpected None response from the OCR API: {response}" + ) + + 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, + extra_kwargs=kwargs, + ) + + +@client +def ocr( + model: str, + document: Dict[str, str], + 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 or + {"type": "image_url", "image_url": "https://..."} for images + 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}" + } + ) + + # Access pages + for page in response.pages: + print(f"Page {page.index}: {page.markdown}") + ``` + """ + local_vars = locals() + try: + litellm_logging_obj: LiteLLMLoggingObj = kwargs.pop("litellm_logging_obj") # type: ignore + litellm_call_id: Optional[str] = kwargs.get("litellm_call_id", None) + _is_async = kwargs.pop("aocr", False) is True + + # Validate document parameter format (Mistral spec) + if not isinstance(document, dict): + raise ValueError(f"document must be a dict with 'type' and URL field, got {type(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' or 'image_url'") + + 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, + ) + ) + + # Update with dynamic values if available + if dynamic_api_key: + api_key = dynamic_api_key + if dynamic_api_base: + api_base = dynamic_api_base + + # Get provider config + ocr_provider_config: Optional[BaseOCRConfig] = ( + ProviderConfigManager.get_provider_ocr_config( + model=model, + provider=litellm.LlmProviders(custom_llm_provider), + ) + ) + + if ocr_provider_config is None: + raise ValueError( + f"OCR is not supported for provider: {custom_llm_provider}" + ) + + verbose_logger.debug( + f"OCR call - model: {model}, provider: {custom_llm_provider}" + ) + + # Extract OCR-specific parameters from kwargs + supported_params = ocr_provider_config.get_supported_ocr_params(model=model) + non_default_params = {} + for param in supported_params: + if param in kwargs: + non_default_params[param] = kwargs.pop(param) + + # Map parameters to provider-specific format + optional_params = ocr_provider_config.map_ocr_params( + non_default_params=non_default_params, + optional_params={}, + model=model, + ) + + verbose_logger.debug(f"OCR optional_params after mapping: {optional_params}") + + # Pre Call logging + litellm_logging_obj.update_environment_variables( + model=model, + optional_params=optional_params, + litellm_params={ + "litellm_call_id": litellm_call_id, + "api_base": api_base, + }, + custom_llm_provider=custom_llm_provider, + ) + + # Call the handler - pass document dict directly + response = base_llm_http_handler.ocr( + model=model, + document=document, # Pass the entire document dict + optional_params=optional_params, + timeout=timeout or request_timeout, + logging_obj=litellm_logging_obj, + api_key=api_key, + api_base=api_base, + custom_llm_provider=custom_llm_provider, + aocr=_is_async, + headers=extra_headers, + provider_config=ocr_provider_config, + litellm_params={ + "api_base": api_base, + "api_key": api_key, + }, + ) + + 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, + extra_kwargs=kwargs, + ) + diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index 9263142dc9..9b19abd9c3 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -277,6 +277,7 @@ class ProxyBaseLLMRequestProcessing: "allm_passthrough_route", "avector_store_search", "avector_store_create", + "aocr", ], version: Optional[str] = None, user_model: Optional[str] = None, @@ -367,6 +368,7 @@ class ProxyBaseLLMRequestProcessing: "allm_passthrough_route", "avector_store_search", "avector_store_create", + "aocr", ], proxy_logging_obj: ProxyLogging, general_settings: dict, diff --git a/litellm/proxy/ocr_endpoints/__init__.py b/litellm/proxy/ocr_endpoints/__init__.py new file mode 100644 index 0000000000..3488912f66 --- /dev/null +++ b/litellm/proxy/ocr_endpoints/__init__.py @@ -0,0 +1,2 @@ +# OCR Endpoints + diff --git a/litellm/proxy/ocr_endpoints/endpoints.py b/litellm/proxy/ocr_endpoints/endpoints.py new file mode 100644 index 0000000000..c1092a06b4 --- /dev/null +++ b/litellm/proxy/ocr_endpoints/endpoints.py @@ -0,0 +1,97 @@ +#### OCR Endpoints ##### + +import orjson +from fastapi import APIRouter, Depends, Request, Response +from fastapi.responses import ORJSONResponse + +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 + +router = APIRouter() + + +@router.post( + "/v1/ocr", + dependencies=[Depends(user_api_key_auth)], + response_class=ORJSONResponse, + tags=["ocr"], +) +@router.post( + "/ocr", + dependencies=[Depends(user_api_key_auth)], + response_class=ORJSONResponse, + tags=["ocr"], +) +async def ocr( + request: Request, + fastapi_response: Response, + user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), +): + """ + OCR endpoint for extracting text from documents and images. + + Follows the Mistral OCR API spec: + https://docs.mistral.ai/capabilities/vision/#optical-character-recognition-ocr + + Example: + ```bash + curl -X POST "http://localhost:4000/v1/ocr" \ + -H "Authorization: Bearer sk-1234" \ + -H "Content-Type: application/json" \ + -d '{ + "model": "mistral/mistral-ocr-latest", + "document": { + "type": "document_url", + "document_url": "https://arxiv.org/pdf/2201.04234" + } + }' + ``` + """ + from litellm.proxy.proxy_server import ( + general_settings, + llm_router, + proxy_config, + proxy_logging_obj, + select_data_generator, + user_api_base, + user_max_tokens, + user_model, + user_request_timeout, + user_temperature, + version, + ) + + # Read request body + body = await request.body() + data = orjson.loads(body) + + # Process request using ProxyBaseLLMRequestProcessing + processor = ProxyBaseLLMRequestProcessing(data=data) + try: + return await processor.base_process_llm_request( + request=request, + fastapi_response=fastapi_response, + user_api_key_dict=user_api_key_dict, + route_type="aocr", + proxy_logging_obj=proxy_logging_obj, + llm_router=llm_router, + general_settings=general_settings, + proxy_config=proxy_config, + select_data_generator=select_data_generator, + model=None, + user_model=user_model, + user_temperature=user_temperature, + user_request_timeout=user_request_timeout, + user_max_tokens=user_max_tokens, + user_api_base=user_api_base, + version=version, + ) + except Exception as e: + raise await processor._handle_llm_api_exception( + e=e, + user_api_key_dict=user_api_key_dict, + proxy_logging_obj=proxy_logging_obj, + version=version, + ) + diff --git a/litellm/proxy/proxy_config.yaml b/litellm/proxy/proxy_config.yaml index 338fa98118..0cf9556d90 100644 --- a/litellm/proxy/proxy_config.yaml +++ b/litellm/proxy/proxy_config.yaml @@ -1,16 +1,6 @@ model_list: - - model_name: db-openai-endpoint + - model_name: mistral/* litellm_params: - model: openai/gm - api_key: hi - api_base: https://exampleopenaiendpoint-production.up.railway.app/ + model: mistral/* -litellm_settings: - callbacks: ["dynamic_rate_limiter_v3"] - priority_reservation: - "prod": 0.9 # 90% reserved for production (9 RPM) - "dev": 0.1 # 10% reserved for development (1 RPM) - priority_reservation_settings: - default_priority: 0.2 # Weight (0%) assigned to keys without explicit priority metadata - saturation_threshold: 0.50 # A model is saturated if it has hit 50% of its RPM limit \ No newline at end of file diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 8f5a046be0..540d6d7844 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -259,7 +259,9 @@ from litellm.proxy.management_endpoints.customer_endpoints import ( from litellm.proxy.management_endpoints.internal_user_endpoints import ( router as internal_user_router, ) -from litellm.proxy.management_endpoints.internal_user_endpoints import user_update +from litellm.proxy.management_endpoints.internal_user_endpoints import ( + user_update, +) from litellm.proxy.management_endpoints.key_management_endpoints import ( delete_verification_tokens, duration_in_seconds, @@ -303,10 +305,13 @@ from litellm.proxy.management_endpoints.user_agent_analytics_endpoints import ( ) from litellm.proxy.management_helpers.audit_logs import create_audit_log_for_update from litellm.proxy.middleware.prometheus_auth_middleware import PrometheusAuthMiddleware +from litellm.proxy.ocr_endpoints.endpoints import router as ocr_router from litellm.proxy.openai_files_endpoints.files_endpoints import ( router as openai_files_router, ) -from litellm.proxy.openai_files_endpoints.files_endpoints import set_files_config +from litellm.proxy.openai_files_endpoints.files_endpoints import ( + set_files_config, +) from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import ( passthrough_endpoint_router, ) @@ -9767,6 +9772,7 @@ app.include_router(response_router) app.include_router(batches_router) app.include_router(public_endpoints_router) app.include_router(rerank_router) +app.include_router(ocr_router) app.include_router(image_router) app.include_router(fine_tuning_router) app.include_router(vector_store_router) diff --git a/litellm/proxy/route_llm_request.py b/litellm/proxy/route_llm_request.py index 1ae87637be..f7b4cf0cbe 100644 --- a/litellm/proxy/route_llm_request.py +++ b/litellm/proxy/route_llm_request.py @@ -25,6 +25,7 @@ ROUTE_ENDPOINT_MAPPING = { "alist_input_items": "/responses/{response_id}/input_items", "aimage_edit": "/images/edits", "acancel_responses": "/responses/{response_id}/cancel", + "aocr": "/ocr", } @@ -98,6 +99,7 @@ async def route_request( "allm_passthrough_route", "avector_store_search", "avector_store_create", + "aocr", ], ): """ diff --git a/litellm/router.py b/litellm/router.py index 5972b06f01..b136dcbae8 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -872,6 +872,14 @@ class Router: generate_content_stream, call_type="generate_content_stream" ) + ######################################################### + # OCR routes + ######################################################### + from litellm.ocr import aocr, ocr + + self.aocr = self.factory_function(aocr, call_type="aocr") + self.ocr = self.factory_function(ocr, call_type="ocr") + def validate_fallbacks(self, fallback_param: Optional[List]): """ Validate the fallbacks parameter. @@ -3537,6 +3545,9 @@ class Router: "avector_store_create", "vector_store_search", "vector_store_create", + "aocr", + "ocr", + "aadapter_generate_content" ] = "assistants", ): """ @@ -3553,6 +3564,7 @@ class Router: "generate_content_stream", "vector_store_search", "vector_store_create", + "ocr", ): def sync_wrapper( @@ -3595,6 +3607,8 @@ class Router: "aimage_edit", "agenerate_content", "agenerate_content_stream", + "aocr", + "ocr", ): return await self._ageneric_api_call_with_fallbacks( original_function=original_function, diff --git a/litellm/utils.py b/litellm/utils.py index 5861d703a3..80ba37307d 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -143,8 +143,10 @@ from litellm.litellm_core_utils.token_counter import get_modified_max_tokens from litellm.llms.base_llm.google_genai.transformation import ( BaseGoogleGenAIGenerateContentConfig, ) +from litellm.llms.base_llm.ocr.transformation import BaseOCRConfig from litellm.llms.bedrock.common_utils import BedrockModelInfo from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler +from litellm.llms.mistral.ocr.transformation import MistralOCRConfig from litellm.router_utils.get_retry_from_policy import ( get_num_retries_from_retry_policy, reset_retry_policy, @@ -7578,6 +7580,25 @@ class ProviderConfigManager: return LiteLLMProxyImageEditConfig() return None + @staticmethod + def get_provider_ocr_config( + model: str, + provider: LlmProviders, + ) -> Optional["BaseOCRConfig"]: + """ + Get OCR configuration for a given provider. + """ + from litellm.llms.azure_ai.ocr.transformation import AzureAIOCRConfig + + PROVIDER_TO_CONFIG_MAP = { + litellm.LlmProviders.MISTRAL: MistralOCRConfig, + litellm.LlmProviders.AZURE_AI: AzureAIOCRConfig, + } + config_class = PROVIDER_TO_CONFIG_MAP.get(provider, None) + if config_class is None: + return None + return config_class() + @staticmethod def get_provider_google_genai_generate_content_config( model: str, diff --git a/tests/ocr_tests/base_ocr_unit_tests.py b/tests/ocr_tests/base_ocr_unit_tests.py new file mode 100644 index 0000000000..1226484479 --- /dev/null +++ b/tests/ocr_tests/base_ocr_unit_tests.py @@ -0,0 +1,140 @@ +""" +Base test class for OCR functionality across different providers. + +This follows the same pattern as BaseLLMChatTest in tests/llm_translation/base_llm_unit_tests.py +""" +import pytest +import litellm +from abc import ABC, abstractmethod + + +# Test resources +TEST_IMAGE_PATH = "test_image_edit.png" +TEST_PDF_URL = "https://arxiv.org/pdf/2201.04234" + + +class BaseOCRTest(ABC): + """ + Abstract base test class that enforces common OCR tests across all providers. + + Each provider-specific test class should inherit from this and implement + get_base_ocr_call_args() to return provider-specific configuration. + """ + + @abstractmethod + def get_base_ocr_call_args(self) -> dict: + """Must return the base OCR call args for the specific provider""" + pass + + @pytest.fixture(autouse=True) + def _handle_rate_limits(self): + """Fixture to handle rate limit errors for all test methods""" + try: + yield + except litellm.RateLimitError: + pytest.skip("Rate limit exceeded") + except litellm.InternalServerError: + pytest.skip("Model is overloaded") + + @pytest.mark.parametrize("sync_mode", [True, False]) + @pytest.mark.asyncio + async def test_basic_ocr_with_url(self, sync_mode): + """ + Test basic OCR with a public URL. + """ + litellm._turn_on_debug() + base_ocr_call_args = self.get_base_ocr_call_args() + print("BASE OCR Call args=", base_ocr_call_args) + + try: + if sync_mode: + response = litellm.ocr( + document={ + "type": "document_url", + "document_url": TEST_PDF_URL + }, + **base_ocr_call_args, + ) + else: + response = await litellm.aocr( + document={ + "type": "document_url", + "document_url": TEST_PDF_URL + }, + **base_ocr_call_args, + ) + + print(f"\n{'='*80}") + print(f"Sync Mode: {sync_mode}") + print(f"Response type: {type(response)}") + print(f"Response object: {response.object if hasattr(response, 'object') else 'N/A'}") + + # Check if response has expected OCR format + assert hasattr(response, "pages"), "Response should have 'pages' attribute" + assert hasattr(response, "model"), "Response should have 'model' attribute" + assert hasattr(response, "object"), "Response should have 'object' attribute" + assert response.object == "ocr", f"Expected object='ocr', got '{response.object}'" + + # Validate pages structure + assert isinstance(response.pages, list), "pages should be a list" + assert len(response.pages) > 0, "Should have at least one page" + + # Check first page structure + first_page = response.pages[0] + assert hasattr(first_page, "index"), "Page should have 'index' attribute" + assert hasattr(first_page, "markdown"), "Page should have 'markdown' attribute" + + # Extract text from all pages for validation + total_text = "\n\n".join(page.markdown for page in response.pages if page.markdown) + print(f"Total pages: {len(response.pages)}") + print(f"Total extracted text length: {len(total_text)} characters") + print(f"First 200 chars: {total_text[:200]}") + print(f"Model: {response.model}") + if response.usage_info: + print(f"Pages processed: {response.usage_info.pages_processed}") + print(f"{'='*80}\n") + + assert len(total_text) > 0, "Should extract some text from the document" + + except Exception as e: + pytest.fail(f"OCR call failed: {str(e)}") + + def test_ocr_response_structure(self): + """ + Test that the OCR response has the correct structure. + """ + litellm.set_verbose = True + base_ocr_call_args = self.get_base_ocr_call_args() + + response = litellm.ocr( + document={ + "type": "document_url", + "document_url": TEST_PDF_URL + }, + **base_ocr_call_args, + ) + + # Validate response structure + assert hasattr(response, "pages"), "Response should have 'pages' attribute" + assert hasattr(response, "model"), "Response should have 'model' attribute" + assert hasattr(response, "object"), "Response should have 'object' attribute" + assert hasattr(response, "usage_info"), "Response should have 'usage_info' attribute" + + assert isinstance(response.pages, list), "pages should be a list" + assert len(response.pages) > 0, "Should have at least one page" + assert response.object == "ocr", "object should be 'ocr'" + + # Validate first page structure + first_page = response.pages[0] + assert hasattr(first_page, "index"), "Page should have 'index' attribute" + assert hasattr(first_page, "markdown"), "Page should have 'markdown' attribute" + assert isinstance(first_page.markdown, str), "markdown should be a string" + + print(f"\nResponse structure validated:") + print(f" - object: {response.object}") + print(f" - model: {response.model}") + print(f" - pages: {len(response.pages)}") + if response.usage_info: + print(f" - pages_processed: {response.usage_info.pages_processed}") + print(f" - doc_size_bytes: {response.usage_info.doc_size_bytes}") + diff --git a/tests/ocr_tests/test_ocr_azure_ai.py b/tests/ocr_tests/test_ocr_azure_ai.py new file mode 100644 index 0000000000..34fd9b37ba --- /dev/null +++ b/tests/ocr_tests/test_ocr_azure_ai.py @@ -0,0 +1,27 @@ +""" +Test OCR functionality with Azure AI API. + +Note: Azure AI OCR automatically converts URLs to base64 data URIs since +the Azure AI endpoint doesn't have internet access. +""" +import os +from base_ocr_unit_tests import BaseOCRTest + +class TestAzureAIOCR(BaseOCRTest): + """ + Test class for Azure AI OCR functionality. + Inherits from BaseOCRTest and provides Azure AI-specific configuration. + + Note: For Azure AI, LiteLLM will automatically convert URLs to base64 data URIs before + sending to the API, since Azure AI OCR endpoint doesn't have internet access. + """ + + def get_base_ocr_call_args(self) -> dict: + """ + Return the base OCR call args for Azure AI. + """ + return { + "model": "azure_ai/mistral-document-ai-2505", + "api_key": os.getenv("AZURE_AI_API_KEY_MISTRAL"), + "api_base": os.getenv("AZURE_AI_API_BASE_MISTRAL"), + } diff --git a/tests/ocr_tests/test_ocr_mistral.py b/tests/ocr_tests/test_ocr_mistral.py new file mode 100644 index 0000000000..ea64745927 --- /dev/null +++ b/tests/ocr_tests/test_ocr_mistral.py @@ -0,0 +1,88 @@ +""" +Test OCR functionality with Mistral API. +""" +import os +import sys +import pytest +import litellm +from litellm import Router +from base_ocr_unit_tests import BaseOCRTest, TEST_PDF_URL + + +class TestMistralOCR(BaseOCRTest): + """ + Test class for Mistral OCR functionality. + """ + + def get_base_ocr_call_args(self) -> dict: + """Return the base OCR call args for Mistral""" + return { + "model": "mistral/mistral-ocr-latest", + "api_key": os.getenv("MISTRAL_API_KEY"), + } + +@pytest.mark.asyncio +async def test_router_aocr_with_mistral(): + """ + Test OCR with Router using Mistral OCR deployment. + """ + litellm.set_verbose = True + + # Create router with Mistral OCR deployment + router = Router( + model_list=[ + { + "model_name": "mistral-ocr", + "litellm_params": { + "model": "mistral/mistral-ocr-latest", + "api_key": os.getenv("MISTRAL_API_KEY"), + }, + } + ] + ) + + try: + # Call OCR through router + response = await router.aocr( + model="mistral-ocr", + document={ + "type": "document_url", + "document_url": TEST_PDF_URL + }, + ) + + print(f"\n{'='*80}") + print("Router OCR Test") + print(f"Response type: {type(response)}") + print(f"Response object: {response.object if hasattr(response, 'object') else 'N/A'}") + + # Check if response has expected Mistral OCR format + assert hasattr(response, "pages"), "Response should have 'pages' attribute" + assert hasattr(response, "model"), "Response should have 'model' attribute" + assert hasattr(response, "object"), "Response should have 'object' attribute" + assert response.object == "ocr", f"Expected object='ocr', got '{response.object}'" + + # Validate pages structure + assert isinstance(response.pages, list), "pages should be a list" + assert len(response.pages) > 0, "Should have at least one page" + + # Check first page structure + first_page = response.pages[0] + assert hasattr(first_page, "index"), "Page should have 'index' attribute" + assert hasattr(first_page, "markdown"), "Page should have 'markdown' attribute" + + # Extract text from all pages for validation + total_text = "\n\n".join(page.markdown for page in response.pages if page.markdown) + print(f"Total pages: {len(response.pages)}") + print(f"Total extracted text length: {len(total_text)} characters") + print(f"First 200 chars: {total_text[:200]}") + print(f"Model: {response.model}") + if response.usage_info: + print(f"Pages processed: {response.usage_info.pages_processed}") + print(f"{'='*80}\n") + + assert len(total_text) > 0, "Should extract some text from the document" + + except Exception as e: + pytest.fail(f"Router OCR call failed: {str(e)}") +