ai-workspace-services/litellm
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Enable local file support for OCR (#22133)

Browse Source 
* [Docs] Enable local file support

Implemented internal handling for converting file-type documents to the required format for OCR processing, ensuring seamless integration with various providers.

* Refactor OCR file handling and improve security checks

Removed deprecated MIME type mapping and file conversion functions, replacing them with updated implementations. Enhanced security by rejecting 'file' document types in JSON requests, ensuring file uploads are handled via multipart/form-data. Updated tests to reflect these changes and ensure proper functionality.

* Enhance MIME type validation in OCR processing

Added a regular expression check to validate MIME types in the convert_file_document_to_url_document function, raising a ValueError for invalid types. Updated tests to ensure proper error handling for unsupported MIME types.

* Enhance type safety in OCR file handling

Added type casting for the uploaded file in the _parse_multipart_form function to ensure proper handling of UploadFile instances. This change improves type safety and reduces potential runtime errors during file processing.

* Refactor MIME type handling in document uploads

Updated the MIME type extraction logic to strip parameters from the Content-Type header, ensuring only the base type is used. Added tests to verify that MIME parameters are correctly handled and stripped in various scenarios.

* Update OCR documentation for MIME type recommendations and remove unnecessary tips

Clarified the recommended usage of MIME types for raw bytes in document uploads. Simplified the documentation by removing the tip about multipart file uploads from tools like Postman, ensuring a more concise and focused guide.

* Enhance multipart form handling in OCR endpoints

Updated the _parse_multipart_form function to ignore both 'file' and 'document' fields during form parsing, ensuring that the document built from the uploaded file is not overridden. Added a new test to verify that injected document fields do not affect the constructed document, improving security and robustness of the file upload process.
This commit is contained in:
Noah Nistler 2026-02-27 12:50:02 -06:00 committed by GitHub
parent adb9d94833
commit d13508c1c5
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
6 changed files with 939 additions and 65 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

98
docs/my-website/docs/ocr.md
View File

@ -61,6 +61,52 @@ async def test_async_ocr():
asyncio.run(test_async_ocr())
```
### Using Local Files
LiteLLM can read local files directly — no manual base64 encoding needed:
```python
from litellm import ocr
# OCR with a local PDF file path
response = ocr(
model="mistral/mistral-ocr-latest",
document={
"type": "file",
"file": "/path/to/document.pdf"
}
)
# OCR with a file object
response = ocr(
model="mistral/mistral-ocr-latest",
document={
"type": "file",
"file": open("document.pdf", "rb")
}
)
# OCR with raw bytes
with open("document.pdf", "rb") as f:
pdf_bytes = f.read()
response = ocr(
model="mistral/mistral-ocr-latest",
document={
"type": "file",
"file": pdf_bytes,
"mime_type": "application/pdf" # recommended for raw bytes (auto-detected from extension for file paths)
}
)
```
The `file` field accepts:
- **File path** (`str` or `pathlib.Path`) — LiteLLM reads the file and detects the MIME type from the extension
- **File object** (binary file-like object) — e.g. `open("doc.pdf", "rb")`
- **Raw bytes** (`bytes`) — use `mime_type` to specify the content type
LiteLLM automatically converts file inputs to base64 data URIs internally, so all providers work seamlessly.
### Using Base64 Encoded Documents
```python
@ -121,7 +167,7 @@ litellm --config /path/to/config.yaml
# RUNNING on http://0.0.0.0:4000
```
Test request
**Test request — JSON body**
```bash
curl http://0.0.0.0:4000/v1/ocr \
@ -136,6 +182,27 @@ curl http://0.0.0.0:4000/v1/ocr \
}'
```
**Test request — multipart file upload**
Upload a file directly using multipart form data. No need to base64-encode the file yourself.
```bash
curl http://0.0.0.0:4000/v1/ocr \
-H "Authorization: Bearer sk-1234" \
-F "model=mistral-ocr" \
-F "file=@/path/to/document.pdf"
```
You can also pass optional parameters as additional form fields:
```bash
curl http://0.0.0.0:4000/v1/ocr \
-H "Authorization: Bearer sk-1234" \
-F "model=mistral-ocr" \
-F "file=@screenshot.png" \
-F 'pages=[0,1,2]' \
-F "include_image_base64=true"
```
## **Request/Response Format**
@ -168,10 +235,12 @@ See the [official Mistral OCR documentation](https://docs.mistral.ai/capabilitie
| 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"`) |
| `document` | object | Yes | Document to process. Must contain `type` and the corresponding field |
| `document.type` | string | Yes | `"document_url"` for PDFs/docs, `"image_url"` for images, or `"file"` for local files |
| `document.document_url` | string | Conditional | URL or data URI to the document (required if `type` is `"document_url"`) |
| `document.image_url` | string | Conditional | URL or data URI to the image (required if `type` is `"image_url"`) |
| `document.file` | string/bytes/file | Conditional | File path, bytes, or file-like object (required if `type` is `"file"`) |
| `document.mime_type` | string | No | Explicit MIME type for file inputs (auto-detected from extension if not provided) |
| `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 |
@ -179,7 +248,7 @@ See the [official Mistral OCR documentation](https://docs.mistral.ai/capabilitie
#### Document Format Examples
**For PDFs and documents:**
**For PDFs and documents (URL):**
```json
{
"type": "document_url",
@ -187,7 +256,7 @@ See the [official Mistral OCR documentation](https://docs.mistral.ai/capabilitie
}
```
**For images:**
**For images (URL):**
```json
{
"type": "image_url",
@ -203,6 +272,21 @@ See the [official Mistral OCR documentation](https://docs.mistral.ai/capabilitie
}
```
**For local files (SDK):**
```python
{"type": "file", "file": "/path/to/document.pdf"}
{"type": "file", "file": open("image.png", "rb")}
{"type": "file", "file": pdf_bytes, "mime_type": "application/pdf"}
```
**For file uploads (Proxy — multipart form):**
```bash
curl http://0.0.0.0:4000/v1/ocr \
-H "Authorization: Bearer sk-1234" \
-F "model=mistral-ocr" \
-F "file=@document.pdf"
```
### Response Format
The response follows Mistral's OCR format with the following structure:

10
litellm/llms/base_llm/ocr/transformation.py
View File

@ -15,7 +15,9 @@ else:
LiteLLMLoggingObj = Any
# DocumentType for OCR - Mistral format document dict
# DocumentType for OCR - providers always receive a dict with
# type="document_url" or type="image_url" (str values only).
# File-type inputs are preprocessed to this format in litellm/ocr/main.py.
DocumentType = Dict[str, str]
@ -141,9 +143,13 @@ class BaseOCRConfig:
Transform OCR request to provider-specific format.
Override in provider-specific implementations.
Note: By the time this method is called, any file-type documents have already
been converted to document_url/image_url format with base64 data URIs by
the preprocessing in litellm/ocr/main.py.
Args:
model: Model name
document: Document to process (Mistral format dict, or file path, bytes, etc.)
document: Document to process - always a dict with type="document_url" or type="image_url"
optional_params: Optional parameters for the request
headers: Request headers

226
litellm/ocr/main.py
View File

@ -2,8 +2,14 @@
Main OCR function for LiteLLM.
"""
import asyncio
import base64
import contextvars
import mimetypes
import os
import re
from functools import partial
from io import IOBase
from pathlib import Path
from typing import Any, Coroutine, Dict, Optional, Union
import httpx
@ -25,7 +31,7 @@ base_llm_http_handler = BaseLLMHTTPHandler()
@client
async def aocr(
model: str,
document: Dict[str, str],
document: Dict[str, Any],
api_key: Optional[str] = None,
api_base: Optional[str] = None,
timeout: Optional[Union[float, httpx.Timeout]] = None,
@ -35,26 +41,27 @@ async def aocr(
) -> 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
{"type": "document_url", "document_url": "https://..."} for PDFs/docs,
{"type": "image_url", "image_url": "https://..."} for images, or
{"type": "file", "file": <path/bytes/file-obj>} for local files
api_key: Optional API key
api_base: Optional API base URL
timeout: Optional timeout
custom_llm_provider: Optional custom LLM provider
extra_headers: Optional extra headers
**kwargs: Additional parameters (e.g., include_image_base64, pages, image_limit)
Returns:
OCRResponse in Mistral OCR format with pages, model, usage_info, etc.
Example:
```python
import litellm
# OCR with PDF
response = await litellm.aocr(
model="mistral/mistral-ocr-latest",
@ -64,7 +71,7 @@ async def aocr(
},
include_image_base64=True
)
# OCR with image
response = await litellm.aocr(
model="mistral/mistral-ocr-latest",
@ -73,7 +80,7 @@ async def aocr(
"image_url": "https://example.com/image.png"
}
)
# OCR with base64 encoded PDF
response = await litellm.aocr(
model="mistral/mistral-ocr-latest",
@ -82,6 +89,12 @@ async def aocr(
"document_url": f"data:application/pdf;base64,{base64_pdf}"
}
)
# OCR with local file
response = await litellm.aocr(
model="mistral/mistral-ocr-latest",
document={"type": "file", "file": "/path/to/document.pdf"}
)
```
"""
local_vars = locals()
@ -135,7 +148,7 @@ async def aocr(
@client
def ocr(
model: str,
document: Dict[str, str],
document: Dict[str, Any],
api_key: Optional[str] = None,
api_base: Optional[str] = None,
timeout: Optional[Union[float, httpx.Timeout]] = None,
@ -145,26 +158,27 @@ def ocr(
) -> 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
{"type": "document_url", "document_url": "https://..."} for PDFs/docs,
{"type": "image_url", "image_url": "https://..."} for images, or
{"type": "file", "file": <path/bytes/file-obj>} for local files
api_key: Optional API key
api_base: Optional API base URL
timeout: Optional timeout
custom_llm_provider: Optional custom LLM provider
extra_headers: Optional extra headers
**kwargs: Additional parameters (e.g., include_image_base64, pages, image_limit)
Returns:
OCRResponse in Mistral OCR format with pages, model, usage_info, etc.
Example:
```python
import litellm
# OCR with PDF
response = litellm.ocr(
model="mistral/mistral-ocr-latest",
@ -174,7 +188,7 @@ def ocr(
},
include_image_base64=True
)
# OCR with image
response = litellm.ocr(
model="mistral/mistral-ocr-latest",
@ -183,7 +197,7 @@ def ocr(
"image_url": "https://example.com/image.png"
}
)
# OCR with base64 encoded PDF
response = litellm.ocr(
model="mistral/mistral-ocr-latest",
@ -192,7 +206,13 @@ def ocr(
"document_url": f"data:application/pdf;base64,{base64_pdf}"
}
)
# OCR with local file
response = litellm.ocr(
model="mistral/mistral-ocr-latest",
document={"type": "file", "file": "/path/to/document.pdf"}
)
# Access pages
for page in response.pages:
print(f"Page {page.index}: {page.markdown}")
@ -203,24 +223,38 @@ def ocr(
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,
# Validate document parameter format
if not isinstance(document, dict):
raise ValueError(
f"document must be a dict with 'type' and URL/file field, got {type(document)}"
)
doc_type = document.get("type")
# Handle file type: convert to document_url/image_url with base64 data URI
if doc_type == "file":
document = convert_file_document_to_url_document(document)
doc_type = document.get("type")
if doc_type not in ["document_url", "image_url"]:
raise ValueError(
f"Invalid document type: {doc_type}. "
"Must be 'document_url', 'image_url', or 'file'"
)
(
model,
custom_llm_provider,
dynamic_api_key,
dynamic_api_base,
) = litellm.get_llm_provider(
model=model,
custom_llm_provider=custom_llm_provider,
api_base=api_base,
api_key=api_key,
)
# Update with dynamic values if available
if dynamic_api_key:
api_key = dynamic_api_key
@ -228,11 +262,11 @@ def ocr(
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),
)
ocr_provider_config: Optional[
BaseOCRConfig
] = ProviderConfigManager.get_provider_ocr_config(
model=model,
provider=litellm.LlmProviders(custom_llm_provider),
)
if ocr_provider_config is None:
@ -246,21 +280,21 @@ def ocr(
# Get litellm params using GenericLiteLLMParams (same as responses API)
litellm_params = GenericLiteLLMParams(**kwargs)
# 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
@ -300,3 +334,111 @@ def ocr(
extra_kwargs=kwargs,
)
#################################################
# Public utilities — used by the SDK and the proxy
#################################################
_MIME_PATTERN = re.compile(r"^[\w.+-]+/[\w.+-]+$")
_MIME_TYPE_MAP = {
".pdf": "application/pdf",
".png": "image/png",
".jpg": "image/jpeg",
".jpeg": "image/jpeg",
".gif": "image/gif",
".webp": "image/webp",
".tiff": "image/tiff",
".tif": "image/tiff",
".bmp": "image/bmp",
}
def get_mime_type(file_path: str) -> str:
"""
Determine MIME type from file path extension.
Falls back to mimetypes.guess_type, then to 'application/octet-stream'.
"""
ext = os.path.splitext(file_path)[1].lower()
mime = _MIME_TYPE_MAP.get(ext)
if mime:
return mime
guessed, _ = mimetypes.guess_type(file_path)
return guessed or "application/octet-stream"
def convert_file_document_to_url_document(document: Dict[str, Any]) -> Dict[str, str]:
"""
Convert a file-type document dict to a document_url-type document dict
with an inline base64 data URI.
Accepts document dicts like:
{"type": "file", "file": "/path/to/document.pdf"} # file path string
{"type": "file", "file": Path("/path/to/doc.pdf")} # pathlib.Path
{"type": "file", "file": <binary file-like object>} # file-like object (BinaryIO)
{"type": "file", "file": b"raw bytes"} # raw bytes
Returns:
{"type": "document_url", "document_url": "data:<mime>;base64,<data>"}
or {"type": "image_url", "image_url": "data:<mime>;base64,<data>"}
"""
file_input = document.get("file")
if file_input is None:
raise ValueError(
"document with type='file' must include a 'file' field containing "
"a file path (str), pathlib.Path, file-like object, or bytes"
)
file_bytes: bytes
mime_type: str = "application/octet-stream"
file_name: Optional[str] = None
if isinstance(file_input, (str, Path)):
file_path = str(file_input)
if not os.path.isfile(file_path):
raise FileNotFoundError(f"File not found: {file_path}")
mime_type = get_mime_type(file_path)
file_name = os.path.basename(file_path)
with open(file_path, "rb") as f:
file_bytes = f.read()
elif isinstance(file_input, bytes):
file_bytes = file_input
elif isinstance(file_input, IOBase) or hasattr(file_input, "read"):
if hasattr(file_input, "name"):
file_name = getattr(file_input, "name", None)
if file_name:
mime_type = get_mime_type(file_name)
file_bytes = file_input.read()
if isinstance(file_bytes, str):
file_bytes = file_bytes.encode("utf-8")
else:
raise ValueError(
f"Unsupported file input type: {type(file_input)}. "
"Expected str (file path), pathlib.Path, bytes, or a file-like object."
)
if not file_bytes:
raise ValueError("File is empty or could not be read")
if "mime_type" in document:
mime_type = document["mime_type"]
if not _MIME_PATTERN.match(mime_type):
raise ValueError(f"Invalid MIME type: {mime_type}")
base64_data = base64.b64encode(file_bytes).decode("utf-8")
data_uri = f"data:{mime_type};base64,{base64_data}"
if mime_type.startswith("image/"):
verbose_logger.debug(
f"OCR file input: Converted file to image_url data URI "
f"(mime={mime_type}, size={len(file_bytes)} bytes, name={file_name})"
)
return {"type": "image_url", "image_url": data_uri}
else:
verbose_logger.debug(
f"OCR file input: Converted file to document_url data URI "
f"(mime={mime_type}, size={len(file_bytes)} bytes, name={file_name})"
)
return {"type": "document_url", "document_url": data_uri}

206
litellm/proxy/ocr_endpoints/endpoints.py
View File

@ -1,9 +1,14 @@
#### OCR Endpoints #####
import json
from typing import Any, Dict, Optional, cast
import orjson
from fastapi import APIRouter, Depends, Request, Response
from fastapi import APIRouter, Depends, Request, Response, UploadFile
from fastapi.responses import ORJSONResponse
from litellm._logging import verbose_proxy_logger
from litellm.ocr.main import convert_file_document_to_url_document, get_mime_type
from litellm.proxy._types import *
from litellm.proxy.auth.user_api_key_auth import UserAPIKeyAuth, user_api_key_auth
from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing
@ -11,6 +16,171 @@ from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessin
router = APIRouter()
def _build_document_from_upload(
file_content: bytes,
filename: Optional[str],
content_type: Optional[str],
) -> Dict[str, str]:
"""
Convert uploaded file bytes into a Mistral-format document dict with base64 data URI.
Delegates to convert_file_document_to_url_document after resolving MIME type
from the upload's content_type header or filename.
"""
mime_type = content_type.split(";")[0].strip() if content_type else None
if not mime_type or mime_type == "application/octet-stream":
if filename:
mime_type = get_mime_type(filename)
return convert_file_document_to_url_document(
{
"type": "file",
"file": file_content,
"mime_type": mime_type or "application/octet-stream",
}
)
async def _parse_multipart_form(request: Request) -> Dict[str, Any]:
"""
Extract OCR data from a multipart form request.
Uses the cached form if already parsed by auth middleware,
otherwise parses the form from the request.
Returns:
A dict with 'document', 'model', and any other OCR params.
"""
try:
form = await request.form()
except Exception as e:
raise ValueError(
f"Failed to parse multipart form data: {str(e)}. "
"When using curl with --form/-F, do NOT set the Content-Type header "
"manually — curl will set it automatically with the required boundary."
)
uploaded_file = form.get("file")
# request.form() may return either a FastAPI or Starlette UploadFile
# depending on middleware; check both via isinstance (FastAPI's UploadFile
# is a subclass of Starlette's) and fall back to duck-type check.
if uploaded_file is None or (
not isinstance(uploaded_file, UploadFile) and not hasattr(uploaded_file, "read")
):
raise ValueError(
"Multipart OCR request must include a 'file' field with the document to process"
)
uploaded_file = cast(UploadFile, uploaded_file)
# Seek to start in case the file was already partially read by middleware
await uploaded_file.seek(0)
file_content = await uploaded_file.read()
if not file_content:
raise ValueError("Uploaded file is empty")
document = _build_document_from_upload(
file_content=file_content,
filename=uploaded_file.filename,
content_type=uploaded_file.content_type,
)
data: Dict[str, Any] = {"document": document}
for field_name, field_value in form.items():
if field_name in ("file", "document"):
continue
# Try to parse JSON values (e.g. pages=[0,1,2])
if isinstance(field_value, str):
try:
data[field_name] = json.loads(field_value)
except (json.JSONDecodeError, ValueError):
data[field_name] = field_value
else:
data[field_name] = field_value
verbose_proxy_logger.debug(
f"OCR multipart form request parsed - model: {data.get('model')}, "
f"document_type: {document['type']}, "
f"filename: {uploaded_file.filename}"
)
return data
async def _parse_ocr_request(request: Request) -> Dict[str, Any]:
"""
Parse an OCR request, supporting both JSON and multipart form data.
JSON body (existing behavior):
{
"model": "mistral/mistral-ocr-latest",
"document": {"type": "document_url", "document_url": "https://..."}
}
Multipart form data (new):
- file: the uploaded file
- model: model name (form field)
- Any other OCR params as form fields (pages, include_image_base64, etc.)
Returns:
A dict suitable for passing to the OCR processing pipeline.
"""
content_type = request.headers.get("content-type", "")
if "multipart/form-data" in content_type.lower():
return await _parse_multipart_form(request)
# --- JSON body (existing behavior) ---
try:
body = await request.body()
except RuntimeError:
# Body stream was consumed by auth middleware (e.g., form parsing).
body = b""
if not body:
# The body may be empty because the auth middleware already parsed
# it as form data (e.g., _read_request_body called request.form()).
# Check if form data is available.
if getattr(request, "_form", None) is not None:
verbose_proxy_logger.debug(
"OCR request body is empty but form data is available from middleware — "
"processing as multipart form."
)
return await _parse_multipart_form(request)
raise ValueError(
"Empty request body. For file uploads, use multipart/form-data content type "
"with a file field. When using curl with --form/-F, do NOT set the Content-Type "
"header manually."
)
try:
data = orjson.loads(body)
except orjson.JSONDecodeError as e:
raise ValueError(
f"Invalid JSON in request body: {e}. "
"Ensure the request body is valid JSON with Content-Type: application/json, "
"or use multipart/form-data for file uploads."
)
# Security: reject type="file" documents received via JSON.
# The "file" document type is designed for local SDK usage where the
# caller and the process share a filesystem. In the proxy context the
# caller is remote, so allowing a file-path string would let an
# authenticated user read arbitrary files from the server's filesystem.
# File uploads must go through multipart/form-data instead.
doc = data.get("document") if isinstance(data, dict) else None
if isinstance(doc, dict) and doc.get("type") == "file":
raise ValueError(
"document type 'file' is not supported through the JSON API. "
"To upload a local file, use multipart/form-data with a 'file' field. "
"For JSON requests, use 'document_url' or 'image_url' document types."
)
return data
@router.post(
"/v1/ocr",
dependencies=[Depends(user_api_key_auth)],
@ -30,23 +200,30 @@ async def ocr(
):
"""
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:
Supports two input modes:
**1. JSON body** (Mistral OCR API compatible):
```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",
"model": "mistral-ocr",
"document": {
"type": "document_url",
"document_url": "https://arxiv.org/pdf/2201.04234"
}
}'
```
**2. Multipart form file upload**:
```bash
curl -X POST "http://localhost:4000/v1/ocr" \
-H "Authorization: Bearer sk-1234" \
-F "model=mistral-ocr" \
-F "file=@document.pdf"
```
"""
from litellm.proxy.proxy_server import (
general_settings,
@ -62,13 +239,14 @@ async def ocr(
version,
)
# Read request body
body = await request.body()
data = orjson.loads(body)
# Process request using ProxyBaseLLMRequestProcessing
processor = ProxyBaseLLMRequestProcessing(data=data)
data: dict = {}
try:
# Parse request body (JSON or multipart form)
data = await _parse_ocr_request(request)
# Process request using ProxyBaseLLMRequestProcessing
processor = ProxyBaseLLMRequestProcessing(data=data)
return await processor.base_process_llm_request(
request=request,
fastapi_response=fastapi_response,
@ -88,10 +266,10 @@ async def ocr(
version=version,
)
except Exception as e:
processor = ProxyBaseLLMRequestProcessing(data=data)
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,
)

0
tests/test_litellm/ocr/__init__.py Normal file
View File

464
tests/test_litellm/ocr/test_ocr_file_input.py Normal file
View File

@ -0,0 +1,464 @@
"""
Tests for OCR file input support.
Tests that:
1. The SDK document parameter with type="file" correctly converts file paths,
file objects, and raw bytes to base64 data URIs before sending to providers.
2. The proxy _build_document_from_upload helper correctly handles uploaded file bytes.
3. The proxy rejects type="file" documents received via JSON (security guard).
4. The proxy returns user-friendly errors for invalid JSON bodies.
"""
import base64
import os
import tempfile
from io import BytesIO
from pathlib import Path
from unittest.mock import AsyncMock, MagicMock
import orjson
import pytest
from litellm.ocr.main import convert_file_document_to_url_document, get_mime_type
class TestGetMimeType:
def test_should_detect_pdf_mime_type(self):
assert get_mime_type("document.pdf") == "application/pdf"
def test_should_detect_png_mime_type(self):
assert get_mime_type("image.png") == "image/png"
def test_should_detect_jpg_mime_type(self):
assert get_mime_type("photo.jpg") == "image/jpeg"
def test_should_detect_jpeg_mime_type(self):
assert get_mime_type("photo.jpeg") == "image/jpeg"
def test_should_detect_gif_mime_type(self):
assert get_mime_type("animation.gif") == "image/gif"
def test_should_detect_webp_mime_type(self):
assert get_mime_type("image.webp") == "image/webp"
def test_should_detect_tiff_mime_type(self):
assert get_mime_type("scan.tiff") == "image/tiff"
def test_should_detect_tif_mime_type(self):
assert get_mime_type("scan.tif") == "image/tiff"
def test_should_detect_bmp_mime_type(self):
assert get_mime_type("bitmap.bmp") == "image/bmp"
def test_should_be_case_insensitive(self):
assert get_mime_type("DOCUMENT.PDF") == "application/pdf"
assert get_mime_type("IMAGE.PNG") == "image/png"
def test_should_fallback_for_unknown_extension(self):
result = get_mime_type("file.xyz123")
assert isinstance(result, str)
class TestConvertFileDocumentToUrlDocument:
def test_should_convert_pdf_file_path_to_document_url(self):
"""File path to a PDF should produce type=document_url with base64 data URI."""
pdf_content = b"%PDF-1.4 test content"
with tempfile.NamedTemporaryFile(suffix=".pdf", delete=False) as f:
f.write(pdf_content)
f.flush()
tmp_path = f.name
try:
result = convert_file_document_to_url_document(
{"type": "file", "file": tmp_path}
)
assert result["type"] == "document_url"
assert result["document_url"].startswith("data:application/pdf;base64,")
b64_data = result["document_url"].split(";base64,")[1]
assert base64.b64decode(b64_data) == pdf_content
finally:
os.unlink(tmp_path)
def test_should_convert_image_file_path_to_image_url(self):
"""File path to a PNG image should produce type=image_url with base64 data URI."""
png_content = b"\x89PNG\r\n\x1a\n fake png content"
with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as f:
f.write(png_content)
f.flush()
tmp_path = f.name
try:
result = convert_file_document_to_url_document(
{"type": "file", "file": tmp_path}
)
assert result["type"] == "image_url"
assert result["image_url"].startswith("data:image/png;base64,")
b64_data = result["image_url"].split(";base64,")[1]
assert base64.b64decode(b64_data) == png_content
finally:
os.unlink(tmp_path)
def test_should_convert_pathlib_path(self):
"""pathlib.Path objects should work the same as string paths."""
content = b"test pdf content"
with tempfile.NamedTemporaryFile(suffix=".pdf", delete=False) as f:
f.write(content)
f.flush()
tmp_path = Path(f.name)
try:
result = convert_file_document_to_url_document(
{"type": "file", "file": tmp_path}
)
assert result["type"] == "document_url"
assert result["document_url"].startswith("data:application/pdf;base64,")
finally:
os.unlink(str(tmp_path))
def test_should_convert_raw_bytes(self):
"""Raw bytes should be converted using a fallback MIME type."""
content = b"raw bytes content"
result = convert_file_document_to_url_document(
{"type": "file", "file": content}
)
assert result["type"] == "document_url"
assert "base64," in result["document_url"]
b64_data = result["document_url"].split(";base64,")[1]
assert base64.b64decode(b64_data) == content
def test_should_convert_raw_bytes_with_explicit_mime_type(self):
"""Raw bytes with explicit mime_type should use the specified MIME type."""
content = b"raw pdf content"
result = convert_file_document_to_url_document(
{"type": "file", "file": content, "mime_type": "application/pdf"}
)
assert result["type"] == "document_url"
assert result["document_url"].startswith("data:application/pdf;base64,")
def test_should_convert_raw_bytes_with_image_mime_type(self):
"""Raw bytes with an image MIME type should produce type=image_url."""
content = b"raw image content"
result = convert_file_document_to_url_document(
{"type": "file", "file": content, "mime_type": "image/jpeg"}
)
assert result["type"] == "image_url"
assert result["image_url"].startswith("data:image/jpeg;base64,")
def test_should_convert_file_like_object(self):
"""BytesIO and other file-like objects should be supported."""
content = b"file-like content"
file_obj = BytesIO(content)
result = convert_file_document_to_url_document(
{"type": "file", "file": file_obj}
)
assert result["type"] == "document_url"
assert "base64," in result["document_url"]
def test_should_convert_file_like_object_with_name(self):
"""File-like objects with a .name attribute should detect MIME from the name."""
content = b"file-like png content"
file_obj = BytesIO(content)
file_obj.name = "test_image.png"
result = convert_file_document_to_url_document(
{"type": "file", "file": file_obj}
)
assert result["type"] == "image_url"
assert result["image_url"].startswith("data:image/png;base64,")
def test_should_raise_error_for_missing_file_field(self):
"""Missing 'file' field should raise ValueError."""
with pytest.raises(ValueError, match="must include a 'file' field"):
convert_file_document_to_url_document({"type": "file"})
def test_should_raise_error_for_nonexistent_file_path(self):
"""Non-existent file path should raise FileNotFoundError."""
with pytest.raises(FileNotFoundError, match="File not found"):
convert_file_document_to_url_document(
{"type": "file", "file": "/nonexistent/path/to/file.pdf"}
)
def test_should_raise_error_for_empty_file(self):
"""Empty file should raise ValueError."""
with tempfile.NamedTemporaryFile(suffix=".pdf", delete=False) as f:
tmp_path = f.name
try:
with pytest.raises(ValueError, match="File is empty"):
convert_file_document_to_url_document(
{"type": "file", "file": tmp_path}
)
finally:
os.unlink(tmp_path)
def test_should_raise_error_for_unsupported_type(self):
"""Unsupported file input types should raise ValueError."""
with pytest.raises(ValueError, match="Unsupported file input type"):
convert_file_document_to_url_document({"type": "file", "file": 12345})
def test_should_raise_error_for_invalid_mime_type(self):
"""MIME types with special characters should be rejected."""
content = b"some content"
with pytest.raises(ValueError, match="Invalid MIME type"):
convert_file_document_to_url_document(
{"type": "file", "file": content, "mime_type": "text/html; charset=utf-8\nX-Injected: true"}
)
def test_should_override_mime_type_for_file_path(self):
"""Explicit mime_type should override auto-detection from extension."""
content = b"some content"
with tempfile.NamedTemporaryFile(suffix=".pdf", delete=False) as f:
f.write(content)
f.flush()
tmp_path = f.name
try:
result = convert_file_document_to_url_document(
{"type": "file", "file": tmp_path, "mime_type": "image/png"}
)
assert result["type"] == "image_url"
assert result["image_url"].startswith("data:image/png;base64,")
finally:
os.unlink(tmp_path)
class TestBuildDocumentFromUpload:
"""Test the proxy endpoint's file upload to document conversion helper."""
@pytest.fixture(autouse=True)
def _import_helper(self):
"""Import the proxy helper, skip if proxy deps aren't installed."""
try:
from litellm.proxy.ocr_endpoints.endpoints import (
_build_document_from_upload,
)
self._build = _build_document_from_upload
except ImportError:
pytest.skip("Proxy dependencies (fastapi/orjson) not installed")
def test_should_build_document_url_for_pdf(self):
content = b"%PDF-1.4 test content"
result = self._build(
file_content=content,
filename="document.pdf",
content_type="application/pdf",
)
assert result["type"] == "document_url"
assert result["document_url"].startswith("data:application/pdf;base64,")
b64_data = result["document_url"].split(";base64,")[1]
assert base64.b64decode(b64_data) == content
def test_should_build_image_url_for_png(self):
content = b"\x89PNG fake png"
result = self._build(
file_content=content,
filename="screenshot.png",
content_type="image/png",
)
assert result["type"] == "image_url"
assert result["image_url"].startswith("data:image/png;base64,")
def test_should_build_image_url_for_jpeg(self):
content = b"\xff\xd8\xff fake jpeg"
result = self._build(
file_content=content,
filename="photo.jpg",
content_type="image/jpeg",
)
assert result["type"] == "image_url"
assert result["image_url"].startswith("data:image/jpeg;base64,")
def test_should_detect_mime_from_filename_when_content_type_is_octet_stream(self):
content = b"pdf content"
result = self._build(
file_content=content,
filename="report.pdf",
content_type="application/octet-stream",
)
assert result["type"] == "document_url"
assert result["document_url"].startswith("data:application/pdf;base64,")
def test_should_detect_mime_from_filename_when_content_type_is_none(self):
content = b"png content"
result = self._build(
file_content=content,
filename="image.png",
content_type=None,
)
assert result["type"] == "image_url"
assert result["image_url"].startswith("data:image/png;base64,")
def test_should_fallback_to_octet_stream_for_unknown(self):
content = b"unknown content"
result = self._build(
file_content=content,
filename=None,
content_type=None,
)
assert result["type"] == "document_url"
assert "application/octet-stream" in result["document_url"]
def test_should_preserve_base64_content_correctly(self):
content = b"Hello, World! \x00\x01\x02\xff"
result = self._build(
file_content=content,
filename="test.pdf",
content_type="application/pdf",
)
b64_data = result["document_url"].split(";base64,")[1]
assert base64.b64decode(b64_data) == content
def test_should_strip_mime_parameters_from_content_type(self):
"""Content-Type with parameters (e.g. charset) should be stripped to the base MIME type."""
content = b"%PDF-1.4 test"
result = self._build(
file_content=content,
filename="doc.pdf",
content_type="application/pdf; charset=utf-8",
)
assert result["type"] == "document_url"
assert result["document_url"].startswith("data:application/pdf;base64,")
def test_should_strip_mime_parameters_with_multiple_params(self):
"""Content-Type with multiple parameters should still be stripped correctly."""
content = b"image data"
result = self._build(
file_content=content,
filename="img.png",
content_type="image/png; charset=utf-8; boundary=something",
)
assert result["type"] == "image_url"
assert result["image_url"].startswith("data:image/png;base64,")
class TestProxySecurityGuard:
"""Test that the proxy rejects type='file' documents in JSON requests
and that multipart form fields cannot override the constructed document."""
@pytest.fixture(autouse=True)
def _import_helpers(self):
"""Import the proxy helpers, skip if proxy deps aren't installed."""
try:
from litellm.proxy.ocr_endpoints.endpoints import (
_parse_multipart_form,
_parse_ocr_request,
)
self._parse = _parse_ocr_request
self._parse_multipart = _parse_multipart_form
except ImportError:
pytest.skip("Proxy dependencies (fastapi/orjson) not installed")
@pytest.mark.asyncio
async def test_should_reject_file_type_document_in_json_body(self):
"""type='file' in a JSON body must be rejected to prevent server-side file reads."""
body = orjson.dumps(
{
"model": "mistral/mistral-ocr-latest",
"document": {"type": "file", "file": "/etc/passwd"},
}
)
mock_request = MagicMock()
mock_request.headers = {"content-type": "application/json"}
mock_request.body = AsyncMock(return_value=body)
mock_request._form = None
with pytest.raises(ValueError, match="not supported through the JSON API"):
await self._parse(mock_request)
@pytest.mark.asyncio
async def test_should_accept_document_url_type_in_json_body(self):
"""type='document_url' in a JSON body should pass through normally."""
expected = {
"model": "mistral/mistral-ocr-latest",
"document": {
"type": "document_url",
"document_url": "https://example.com/doc.pdf",
},
}
body = orjson.dumps(expected)
mock_request = MagicMock()
mock_request.headers = {"content-type": "application/json"}
mock_request.body = AsyncMock(return_value=body)
mock_request._form = None
result = await self._parse(mock_request)
assert result["document"]["type"] == "document_url"
@pytest.mark.asyncio
async def test_should_raise_on_invalid_json_body(self):
"""Invalid JSON should produce a user-friendly ValueError."""
mock_request = MagicMock()
mock_request.headers = {"content-type": "application/json"}
mock_request.body = AsyncMock(return_value=b"not valid json{{{")
mock_request._form = None
with pytest.raises(ValueError, match="Invalid JSON in request body"):
await self._parse(mock_request)
@pytest.mark.asyncio
async def test_should_ignore_document_form_field_injection(self):
"""A 'document' form field must not override the document built from the uploaded file."""
from starlette.datastructures import UploadFile
file_content = b"%PDF-1.4 legit content"
upload = UploadFile(filename="legit.pdf", file=BytesIO(file_content))
injected = '{"type": "file", "file": "/etc/passwd"}'
mock_form = {
"file": upload,
"model": "mistral/mistral-ocr-latest",
"document": injected,
}
mock_request = MagicMock()
mock_request.headers = {"content-type": "multipart/form-data; boundary=---"}
mock_request.form = AsyncMock(return_value=mock_form)
result = await self._parse_multipart(mock_request)
assert result["document"]["type"] == "document_url"
assert result["document"]["document_url"].startswith("data:application/pdf;base64,")
assert result["model"] == "mistral/mistral-ocr-latest"