Coverage for .venv/lib/python3.13/site-packages/litellm/proxy/ocr_endpoints/endpoints.py: 30%
112 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-10-10 12:01 +0000
1#### OCR Endpoints #####
3import io
4import json
5from collections.abc import Mapping
6from typing import Final, cast
8import orjson
9from fastapi import APIRouter, Depends, HTTPException, Request, Response, UploadFile
10from fastapi.responses import ORJSONResponse
12from litellm._logging import verbose_proxy_logger
13from litellm.llms.base_llm.ocr.transformation import (
14 OCR_REQUEST_FORMAT_HEADER,
15 OCR_REQUEST_FORMAT_PARAM,
16 OCRResponse,
17 parse_ocr_request_format,
18)
19from litellm.proxy._types import *
20from litellm.proxy.auth.user_api_key_auth import UserAPIKeyAuth, user_api_key_auth
21from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing
23router: Final = APIRouter()
24_MAX_FILE_BYTES: Final = 50 * 1024 * 1024
27class _NamedUpload(io.BytesIO):
28 name: str | None
30 def __init__(self, content: bytes, name: str | None) -> None:
31 super().__init__(content)
32 self.name = name
35def _build_document_from_upload(
36 file_content: bytes,
37 filename: str | None,
38 content_type: str | None,
39) -> dict[str, object]:
40 supplied_mime: Final = content_type.split(";")[0].strip() if content_type else None
41 upload: Final = _NamedUpload(file_content, filename)
42 if supplied_mime and supplied_mime != "application/octet-stream":
43 return {"type": "file", "file": upload, "mime_type": supplied_mime}
44 return {"type": "file", "file": upload}
47def _with_request_format(data: Mapping[str, object], request: Request) -> Mapping[str, object]:
48 """
49 Resolve the requested response format from the body or the `x-req-format` header.
51 An explicit `req_format` in the body wins over the header.
52 """
53 body_value: Final = data.get(OCR_REQUEST_FORMAT_PARAM)
54 header_value: Final = request.headers.get(OCR_REQUEST_FORMAT_HEADER)
55 raw_value: Final = body_value if body_value is not None else header_value
56 if raw_value is None:
57 return data
58 try:
59 request_format: Final = parse_ocr_request_format(
60 raw_value.strip().lower() if isinstance(raw_value, str) else raw_value
61 )
62 except ValueError as e:
63 raise HTTPException(status_code=400, detail={"error": f"{e}"})
64 return {**data, OCR_REQUEST_FORMAT_PARAM: request_format}
67def _native_response(response: object, fastapi_response: Response) -> Response | None:
68 """
69 Return the provider's native payload when the caller asked for
70 `req_format=native` and the provider config captured it, carrying over the
71 LiteLLM response headers (cost, call id, etc.) built for the normalized response.
72 """
73 if not isinstance(response, OCRResponse):
74 return None
75 native_payload: Final = response.get_provider_native_response()
76 if native_payload is None:
77 return None
78 return Response(
79 content=orjson.dumps(native_payload),
80 media_type="application/json",
81 headers={
82 key: value
83 for key, value in fastapi_response.headers.items()
84 if key.lower() not in ("content-length", "content-type")
85 },
86 )
89async def _parse_multipart_form(request: Request) -> dict[str, object]:
90 """
91 Extract OCR data from a multipart form request.
93 Uses the cached form if already parsed by auth middleware,
94 otherwise parses the form from the request.
96 Returns:
97 A dict with 'document', 'model', and any other OCR params.
98 """
99 try:
100 form: Final = await request.form()
101 except Exception as e:
102 raise ValueError(
103 f"Failed to parse multipart form data: {e}. "
104 "When using curl with --form/-F, do NOT set the Content-Type header "
105 "manually — curl will set it automatically with the required boundary."
106 )
108 uploaded_file = form.get("file")
109 # request.form() may return either a FastAPI or Starlette UploadFile
110 # depending on middleware; check both via isinstance (FastAPI's UploadFile
111 # is a subclass of Starlette's) and fall back to duck-type check.
112 if uploaded_file is None or (not isinstance(uploaded_file, UploadFile) and not hasattr(uploaded_file, "read")):
113 raise ValueError("Multipart OCR request must include a 'file' field with the document to process")
115 uploaded_file = cast(UploadFile, uploaded_file)
117 # Seek to start in case the file was already partially read by middleware
118 await uploaded_file.seek(0)
119 file_content: Final = await uploaded_file.read(_MAX_FILE_BYTES + 1)
120 if not file_content:
121 raise ValueError("Uploaded file is empty")
122 if len(file_content) > _MAX_FILE_BYTES:
123 raise ValueError("OCR file exceeds the size limit")
125 document: Final = _build_document_from_upload(
126 file_content=file_content,
127 filename=uploaded_file.filename,
128 content_type=uploaded_file.content_type,
129 )
131 data: Final[dict[str, object]] = {"document": document}
133 for field_name, field_value in form.items():
134 if field_name in ("file", "document"):
135 continue
136 # Try to parse JSON values (e.g. pages=[0,1,2])
137 if isinstance(field_value, str):
138 try:
139 data[field_name] = json.loads(field_value)
140 except (json.JSONDecodeError, ValueError):
141 data[field_name] = field_value
142 else:
143 data[field_name] = field_value
145 verbose_proxy_logger.debug(
146 "OCR multipart form request parsed - model: %s, document_type: %s, filename: %s",
147 data.get("model"),
148 document["type"],
149 uploaded_file.filename,
150 )
152 return data
155async def _parse_ocr_request(request: Request) -> Mapping[str, object]:
156 """Parse an OCR request and apply the `x-req-format` header, if any."""
157 return _with_request_format(await _parse_ocr_request_body(request), request)
160async def _parse_ocr_request_body(request: Request) -> dict[str, object]:
161 """
162 Parse an OCR request, supporting both JSON and multipart form data.
164 JSON body (existing behavior):
165 {
166 "model": "mistral/mistral-ocr-latest",
167 "document": {"type": "document_url", "document_url": "https://..."}
168 }
170 Multipart form data (new):
171 - file: the uploaded file
172 - model: model name (form field)
173 - Any other OCR params as form fields (pages, include_image_base64, etc.)
175 Returns:
176 A dict suitable for passing to the OCR processing pipeline.
177 """
178 content_type: Final = request.headers.get("content-type", "")
180 if "multipart/form-data" in content_type.lower(): 180 ↛ 181line 180 didn't jump to line 181 because the condition on line 180 was never true
181 return await _parse_multipart_form(request)
183 # --- JSON body (existing behavior) ---
184 try:
185 body = await request.body()
186 except RuntimeError:
187 # Body stream was consumed by auth middleware (e.g., form parsing).
188 body = b""
190 if not body: 190 ↛ 206line 190 didn't jump to line 206 because the condition on line 190 was always true
191 # The body may be empty because the auth middleware already parsed
192 # it as form data (e.g., _read_request_body called request.form()).
193 # Check if form data is available.
194 if getattr(request, "_form", None) is not None: 194 ↛ 195line 194 didn't jump to line 195 because the condition on line 194 was never true
195 verbose_proxy_logger.debug(
196 "OCR request body is empty but form data is available from middleware — processing as multipart form."
197 )
198 return await _parse_multipart_form(request)
200 raise ValueError(
201 "Empty request body. For file uploads, use multipart/form-data content type "
202 "with a file field. When using curl with --form/-F, do NOT set the Content-Type "
203 "header manually."
204 )
206 try:
207 data: Final = orjson.loads(body)
208 except orjson.JSONDecodeError as e:
209 raise ValueError(
210 f"Invalid JSON in request body: {e}. "
211 "Ensure the request body is valid JSON with Content-Type: application/json, "
212 "or use multipart/form-data for file uploads."
213 )
215 # Security: reject type="file" documents received via JSON.
216 # The "file" document type is designed for local SDK usage where the
217 # caller and the process share a filesystem. In the proxy context the
218 # caller is remote, so allowing a file-path string would let an
219 # authenticated user read arbitrary files from the server's filesystem.
220 # File uploads must go through multipart/form-data instead.
221 doc: Final = data.get("document") if isinstance(data, dict) else None
222 if isinstance(doc, dict) and doc.get("type") == "file":
223 raise ValueError(
224 "document type 'file' is not supported through the JSON API. "
225 "To upload a local file, use multipart/form-data with a 'file' field. "
226 "For JSON requests, use 'document_url' or 'image_url' document types."
227 )
229 # Security: reject provider-native file IDs (e.g. reducto://) received via
230 # JSON. These IDs are not scoped to the LiteLLM proxy user/key, so an
231 # authenticated user who obtains another user's file ID could submit it
232 # here and receive the OCR result using the proxy's shared provider
233 # credentials. Force callers to upload fresh content per request via
234 # multipart/form-data or an inline base64 data URI, both of which produce
235 # a server-mediated upload bound to the current request.
236 if isinstance(doc, dict):
237 for url_field in ("document_url", "image_url"):
238 url_value = doc.get(url_field)
239 if isinstance(url_value, str) and url_value.startswith("reducto://"):
240 raise ValueError(
241 "reducto:// file IDs are not accepted through the proxy "
242 "OCR API; upload the file in the same request via "
243 "multipart/form-data with a 'file' field, or pass an "
244 "inline base64 data URI as the document URL."
245 )
247 return data
250@router.post(
251 "/v1/ocr",
252 dependencies=[Depends(user_api_key_auth)],
253 response_class=ORJSONResponse,
254 tags=["ocr"],
255)
256@router.post(
257 "/ocr",
258 dependencies=[Depends(user_api_key_auth)],
259 response_class=ORJSONResponse,
260 tags=["ocr"],
261)
262async def ocr(
263 request: Request,
264 fastapi_response: Response,
265 user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
266):
267 """
268 OCR endpoint for extracting text from documents and images.
270 Supports two input modes:
272 **1. JSON body** (Mistral OCR API compatible):
273 ```bash
274 curl -X POST "http://localhost:4000/v1/ocr" \
275 -H "Authorization: Bearer sk-1234" \
276 -H "Content-Type: application/json" \
277 -d '{
278 "model": "mistral-ocr",
279 "document": {
280 "type": "document_url",
281 "document_url": "https://arxiv.org/pdf/2201.04234"
282 }
283 }'
284 ```
286 **2. Multipart form file upload**:
287 ```bash
288 curl -X POST "http://localhost:4000/v1/ocr" \
289 -H "Authorization: Bearer sk-1234" \
290 -F "model=mistral-ocr" \
291 -F "file=@document.pdf"
292 ```
294 Response format is normalized to the LiteLLM OCR schema by default. Providers
295 that support it (Azure Document Intelligence) can return their own payload
296 instead, with cost tracking unchanged, via `x-req-format: native` (or
297 `"req_format": "native"` in the body).
298 """
299 from litellm.proxy.proxy_server import (
300 general_settings,
301 llm_router,
302 proxy_config,
303 proxy_logging_obj,
304 select_data_generator,
305 user_api_base,
306 user_max_tokens,
307 user_model,
308 user_request_timeout,
309 user_temperature,
310 version,
311 )
313 data: dict = {}
314 try:
315 # Parse request body (JSON or multipart form)
316 data = dict(await _parse_ocr_request(request))
318 # Process request using ProxyBaseLLMRequestProcessing
319 processor = ProxyBaseLLMRequestProcessing(data=data)
321 response: Final[object] = await processor.base_process_llm_request(
322 request=request,
323 fastapi_response=fastapi_response,
324 user_api_key_dict=user_api_key_dict,
325 route_type="aocr",
326 proxy_logging_obj=proxy_logging_obj,
327 llm_router=llm_router,
328 general_settings=general_settings,
329 proxy_config=proxy_config,
330 select_data_generator=select_data_generator,
331 model=None,
332 user_model=user_model,
333 user_temperature=user_temperature,
334 user_request_timeout=user_request_timeout,
335 user_max_tokens=user_max_tokens,
336 user_api_base=user_api_base,
337 version=version,
338 )
340 return _native_response(response, fastapi_response) or response
341 except Exception as e:
342 processor = ProxyBaseLLMRequestProcessing(data=data)
343 raise await processor._handle_llm_api_exception(
344 e=e,
345 user_api_key_dict=user_api_key_dict,
346 proxy_logging_obj=proxy_logging_obj,
347 version=version,
348 )