mirror of
https://github.com/tiennm99/litellm.git
synced 2026-08-06 16:24:46 +00:00
[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>
This commit is contained in:
Ishaan Jaff
co-authored by
Copilot Autofix powered by AI <62310815+github-advanced-security[bot]@users.noreply.github.com>
Copilot Autofix powered by AI <62310815+github-advanced-security[bot]@users.noreply.github.com>
parent
74bfbeb439
commit
bc26845ec4
+52
-1
@@ -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
|
||||
|
||||
@@ -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) |
|
||||
|
||||
@@ -347,6 +347,7 @@ const sidebars = {
|
||||
]
|
||||
},
|
||||
"moderation",
|
||||
"ocr",
|
||||
{
|
||||
type: "category",
|
||||
label: "Pass-through Endpoints (Anthropic SDK, etc.)",
|
||||
|
||||
@@ -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 *
|
||||
|
||||
@@ -0,0 +1,5 @@
|
||||
"""Azure AI OCR module."""
|
||||
from .transformation import AzureAIOCRConfig
|
||||
|
||||
__all__ = ["AzureAIOCRConfig"]
|
||||
|
||||
@@ -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://<api_base>/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,
|
||||
)
|
||||
|
||||
@@ -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",
|
||||
]
|
||||
@@ -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,
|
||||
)
|
||||
|
||||
@@ -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",
|
||||
],
|
||||
):
|
||||
|
||||
@@ -0,0 +1,2 @@
|
||||
"""Mistral OCR transformation module."""
|
||||
|
||||
@@ -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": "<https-url or data-uri>"
|
||||
},
|
||||
"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
|
||||
|
||||
@@ -0,0 +1,5 @@
|
||||
"""OCR module for LiteLLM."""
|
||||
from .main import aocr, ocr
|
||||
|
||||
__all__ = ["ocr", "aocr"]
|
||||
|
||||
@@ -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,
|
||||
)
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -0,0 +1,2 @@
|
||||
# OCR Endpoints
|
||||
|
||||
@@ -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,
|
||||
)
|
||||
|
||||
@@ -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
|
||||
@@ -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)
|
||||
|
||||
@@ -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",
|
||||
],
|
||||
):
|
||||
"""
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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}")
|
||||
|
||||
@@ -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"),
|
||||
}
|
||||
@@ -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)}")
|
||||
|
||||
Reference in New Issue
Block a user