mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-15 23:31:29 +00:00
feat(ocr): add Reducto parse OCR support (#26068)
* feat(ocr): add Reducto parse OCR support
* fix(reducto): address OCR review feedback
* chore: refresh uv lockfile
* Revert "chore: refresh uv lockfile"
This reverts commit 47200c0e60.
This commit is contained in:
parent
b8c0f8dbec
commit
91c355dfd2
19 changed files with 1119 additions and 4 deletions
103
docs/my-website/docs/providers/reducto.md
Normal file
103
docs/my-website/docs/providers/reducto.md
Normal file
|
|
@ -0,0 +1,103 @@
|
|||
# Reducto
|
||||
|
||||
## Overview
|
||||
|
||||
| Property | Details |
|
||||
|-------|-------|
|
||||
| Description | Reducto parse support over LiteLLM's existing OCR API |
|
||||
| Provider Route on LiteLLM | `reducto/` |
|
||||
| Supported Operations | `/ocr` |
|
||||
| Supported Models | `reducto/parse-v3`, `reducto/parse-legacy` |
|
||||
| Link to Provider Doc | [Reducto ↗](https://platform.reducto.ai/) |
|
||||
|
||||
Reducto is exposed through LiteLLM's OCR surface, so this provider uses `litellm.ocr()` and `litellm.aocr()`.
|
||||
|
||||
## Quick Start
|
||||
|
||||
### LiteLLM SDK
|
||||
|
||||
```python showLineNumbers title="SDK Usage"
|
||||
import litellm
|
||||
import os
|
||||
|
||||
os.environ["REDUCTO_API_KEY"] = "your-api-key"
|
||||
|
||||
response = litellm.ocr(
|
||||
model="reducto/parse-v3",
|
||||
document={"type": "file", "file": "document.pdf"},
|
||||
)
|
||||
|
||||
for page in response.pages:
|
||||
print(page.markdown)
|
||||
```
|
||||
|
||||
You can also override credentials per call with `api_key=` and `api_base=`.
|
||||
|
||||
### LiteLLM Proxy
|
||||
|
||||
```yaml showLineNumbers title="proxy_config.yaml"
|
||||
model_list:
|
||||
- model_name: reducto-parse
|
||||
litellm_params:
|
||||
model: reducto/parse-v3
|
||||
api_key: os.environ/REDUCTO_API_KEY
|
||||
model_info:
|
||||
mode: ocr
|
||||
```
|
||||
|
||||
## Parse V3
|
||||
|
||||
`reducto/parse-v3` maps to Reducto's current parse API and accepts:
|
||||
|
||||
- `formatting`
|
||||
- `retrieval`
|
||||
- `settings`
|
||||
|
||||
```python showLineNumbers title="Parse V3"
|
||||
response = await litellm.aocr(
|
||||
model="reducto/parse-v3",
|
||||
document={"type": "file", "file": "document.pdf"},
|
||||
formatting={"table_output_format": "html"},
|
||||
retrieval={"chunk_mode": "section"},
|
||||
settings={"ocr_system": "standard"},
|
||||
)
|
||||
```
|
||||
|
||||
## Parse Legacy
|
||||
|
||||
`reducto/parse-legacy` keeps the legacy request shape and accepts `enhance`.
|
||||
|
||||
```python showLineNumbers title="Parse Legacy"
|
||||
response = litellm.ocr(
|
||||
model="reducto/parse-legacy",
|
||||
document={"type": "file", "file": "document.pdf"},
|
||||
enhance={"agentic": [{"type": "table"}]},
|
||||
)
|
||||
```
|
||||
|
||||
## Upload Behavior
|
||||
|
||||
- `document={"type":"file","file":...}` is auto-converted by LiteLLM into a data URI, then uploaded to Reducto's `/upload` endpoint before `/parse`.
|
||||
- `document_url="reducto://..."` is passed through directly and skips upload.
|
||||
- Plain `http(s)` document URLs are rejected for Reducto. Upload the file first or pass a local file to LiteLLM.
|
||||
- Image files also work through `type="file"`; LiteLLM normalizes them to `image_url` data URIs before the Reducto upload step.
|
||||
|
||||
## Cost Tracking
|
||||
|
||||
Reducto returns OCR usage in credits. LiteLLM supports credit-priced OCR models via `ocr_cost_per_credit`.
|
||||
|
||||
If you want spend tracking, register pricing for your deployment:
|
||||
|
||||
```python showLineNumbers title="Register OCR Credit Pricing"
|
||||
import litellm
|
||||
|
||||
litellm.register_model(
|
||||
{
|
||||
"reducto/parse-v3": {
|
||||
"litellm_provider": "reducto",
|
||||
"mode": "ocr",
|
||||
"ocr_cost_per_credit": 0.003,
|
||||
}
|
||||
}
|
||||
)
|
||||
```
|
||||
|
|
@ -996,6 +996,7 @@ const sidebars = {
|
|||
"providers/predibase",
|
||||
"providers/pydantic_ai_agent",
|
||||
"providers/ragflow",
|
||||
"providers/reducto",
|
||||
"providers/recraft",
|
||||
"providers/replicate",
|
||||
{
|
||||
|
|
|
|||
|
|
@ -1801,10 +1801,6 @@ def ocr_cost(
|
|||
if response.usage_info is None:
|
||||
raise ValueError("OCR response usage_info is None")
|
||||
|
||||
pages_processed = response.usage_info.pages_processed
|
||||
if pages_processed is None:
|
||||
raise ValueError("OCR response pages_processed is None")
|
||||
|
||||
try:
|
||||
model_info: Optional[ModelInfo] = litellm.get_model_info(
|
||||
model=model, custom_llm_provider=custom_llm_provider
|
||||
|
|
@ -1812,6 +1808,17 @@ def ocr_cost(
|
|||
except Exception:
|
||||
model_info = None
|
||||
|
||||
credits = getattr(response.usage_info, "credits", None)
|
||||
cost_per_credit = None
|
||||
if model_info is not None:
|
||||
cost_per_credit = model_info.get("ocr_cost_per_credit")
|
||||
if credits is not None and cost_per_credit:
|
||||
return cost_per_credit * credits, 0.0
|
||||
|
||||
pages_processed = response.usage_info.pages_processed
|
||||
if pages_processed is None:
|
||||
return 0.0, 0.0
|
||||
|
||||
ocr_cost_per_page: float = 0.0
|
||||
if model_info is not None:
|
||||
ocr_cost_per_page = model_info.get("ocr_cost_per_page") or 0.0
|
||||
|
|
|
|||
|
|
@ -54,6 +54,7 @@ class OCRUsageInfo(LiteLLMPydanticObjectBase):
|
|||
"""Usage information from OCR response."""
|
||||
|
||||
pages_processed: Optional[int] = None
|
||||
credits: Optional[float] = None
|
||||
doc_size_bytes: Optional[int] = None
|
||||
|
||||
model_config = {"extra": "allow"}
|
||||
|
|
|
|||
1
litellm/llms/reducto/__init__.py
Normal file
1
litellm/llms/reducto/__init__.py
Normal file
|
|
@ -0,0 +1 @@
|
|||
|
||||
141
litellm/llms/reducto/common.py
Normal file
141
litellm/llms/reducto/common.py
Normal file
|
|
@ -0,0 +1,141 @@
|
|||
import base64
|
||||
import binascii
|
||||
from collections import defaultdict
|
||||
from typing import TYPE_CHECKING, Any, Dict, List, NoReturn, Optional, Tuple
|
||||
|
||||
from litellm.constants import request_timeout
|
||||
|
||||
REDUCTO_API_BASE = "https://platform.reducto.ai"
|
||||
REDUCTO_ID_PREFIX = "reducto://"
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from litellm.llms.base_llm.ocr.transformation import OCRPage
|
||||
|
||||
|
||||
def _normalize_api_base(api_base: Optional[str]) -> str:
|
||||
return (api_base or REDUCTO_API_BASE).rstrip("/")
|
||||
|
||||
|
||||
def _raise_bad_request(message: str, model: str) -> NoReturn:
|
||||
import litellm
|
||||
|
||||
raise litellm.BadRequestError(
|
||||
message=message,
|
||||
model=model,
|
||||
llm_provider="reducto",
|
||||
)
|
||||
|
||||
|
||||
def extract_file_id_or_bytes(
|
||||
source_url: str,
|
||||
model: str,
|
||||
) -> Tuple[Optional[str], Optional[bytes], Optional[str]]:
|
||||
if source_url.startswith(REDUCTO_ID_PREFIX):
|
||||
return source_url, None, None
|
||||
|
||||
if source_url.startswith("http://") or source_url.startswith("https://"):
|
||||
_raise_bad_request(
|
||||
"Reducto requires type='file' (auto-uploaded) or a reducto:// id. Plain http(s) URLs are not supported; upload the file first.",
|
||||
model=model,
|
||||
)
|
||||
|
||||
if not source_url.startswith("data:"):
|
||||
_raise_bad_request(
|
||||
"Reducto requires a reducto:// id or a base64 data URI after OCR preprocessing.",
|
||||
model=model,
|
||||
)
|
||||
|
||||
try:
|
||||
header, encoded = source_url.split(",", 1)
|
||||
except ValueError:
|
||||
_raise_bad_request("Invalid Reducto data URI provided.", model=model)
|
||||
|
||||
if ";base64" not in header:
|
||||
_raise_bad_request(
|
||||
"Reducto only supports base64-encoded data URIs.", model=model
|
||||
)
|
||||
|
||||
mime = header.removeprefix("data:").split(";")[0] or "application/octet-stream"
|
||||
try:
|
||||
raw_bytes = base64.b64decode(encoded, validate=True)
|
||||
except (binascii.Error, ValueError):
|
||||
_raise_bad_request("Invalid Reducto base64 payload provided.", model=model)
|
||||
|
||||
return None, raw_bytes, mime
|
||||
|
||||
|
||||
def upload_bytes_sync(
|
||||
raw_bytes: bytes,
|
||||
mime: Optional[str],
|
||||
api_key: str,
|
||||
api_base: Optional[str],
|
||||
) -> str:
|
||||
import litellm
|
||||
|
||||
response = litellm.module_level_client.post(
|
||||
url="{}{}".format(_normalize_api_base(api_base), "/upload"),
|
||||
headers={"Authorization": f"Bearer {api_key}"},
|
||||
files={"file": ("document", raw_bytes, mime or "application/octet-stream")},
|
||||
timeout=request_timeout,
|
||||
)
|
||||
response.raise_for_status()
|
||||
return response.json()["file_id"]
|
||||
|
||||
|
||||
async def upload_bytes_async(
|
||||
raw_bytes: bytes,
|
||||
mime: Optional[str],
|
||||
api_key: str,
|
||||
api_base: Optional[str],
|
||||
) -> str:
|
||||
import litellm
|
||||
|
||||
response = await litellm.module_level_aclient.post(
|
||||
url="{}{}".format(_normalize_api_base(api_base), "/upload"),
|
||||
headers={"Authorization": f"Bearer {api_key}"},
|
||||
files={"file": ("document", raw_bytes, mime or "application/octet-stream")},
|
||||
timeout=request_timeout,
|
||||
)
|
||||
response.raise_for_status()
|
||||
return response.json()["file_id"]
|
||||
|
||||
|
||||
def build_pages_from_reducto(result: Dict[str, Any]) -> List["OCRPage"]:
|
||||
from litellm.llms.base_llm.ocr.transformation import OCRPage
|
||||
|
||||
chunks = result.get("chunks", []) or []
|
||||
blocks_by_page: Dict[int, List[Dict[str, Any]]] = defaultdict(list)
|
||||
|
||||
for chunk in chunks:
|
||||
for block in chunk.get("blocks", []) or []:
|
||||
page_no = (block.get("bbox") or {}).get("page")
|
||||
if page_no is None:
|
||||
continue
|
||||
try:
|
||||
normalized_page = int(page_no)
|
||||
except (TypeError, ValueError):
|
||||
continue
|
||||
blocks_by_page[normalized_page].append(block)
|
||||
|
||||
if not blocks_by_page:
|
||||
fallback_markdown = "\n\n".join(
|
||||
chunk.get("content", "") for chunk in chunks if chunk.get("content")
|
||||
)
|
||||
if fallback_markdown == "":
|
||||
return []
|
||||
return [OCRPage(index=0, markdown=fallback_markdown)]
|
||||
|
||||
pages: List["OCRPage"] = []
|
||||
for page_no, blocks in sorted(blocks_by_page.items()):
|
||||
markdown = "\n\n".join(
|
||||
block.get("content", "") for block in blocks if block.get("content")
|
||||
)
|
||||
page_index = max(page_no - 1, 0)
|
||||
pages.append(
|
||||
OCRPage(
|
||||
index=page_index,
|
||||
markdown=markdown,
|
||||
blocks=blocks,
|
||||
)
|
||||
)
|
||||
return pages
|
||||
1
litellm/llms/reducto/ocr/__init__.py
Normal file
1
litellm/llms/reducto/ocr/__init__.py
Normal file
|
|
@ -0,0 +1 @@
|
|||
|
||||
205
litellm/llms/reducto/ocr/transformation.py
Normal file
205
litellm/llms/reducto/ocr/transformation.py
Normal file
|
|
@ -0,0 +1,205 @@
|
|||
from typing import Any, Dict, Optional
|
||||
|
||||
import httpx
|
||||
|
||||
from litellm.llms.base_llm.ocr.transformation import (
|
||||
BaseOCRConfig,
|
||||
DocumentType,
|
||||
OCRRequestData,
|
||||
OCRResponse,
|
||||
OCRUsageInfo,
|
||||
)
|
||||
from litellm.llms.reducto.common import (
|
||||
REDUCTO_API_BASE,
|
||||
build_pages_from_reducto,
|
||||
extract_file_id_or_bytes,
|
||||
upload_bytes_async,
|
||||
upload_bytes_sync,
|
||||
)
|
||||
|
||||
|
||||
class _BaseReductoOCRConfig(BaseOCRConfig):
|
||||
def __init__(self) -> None:
|
||||
super().__init__()
|
||||
self._api_key: Optional[str] = None
|
||||
self._api_base: Optional[str] = None
|
||||
|
||||
def map_ocr_params(
|
||||
self,
|
||||
non_default_params: dict,
|
||||
optional_params: dict,
|
||||
model: str,
|
||||
) -> dict:
|
||||
mapped_params = dict(optional_params)
|
||||
supported_params = self.get_supported_ocr_params(model=model)
|
||||
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,
|
||||
litellm_params: Optional[dict] = None,
|
||||
**kwargs,
|
||||
) -> Dict:
|
||||
from litellm.secret_managers.main import get_secret_str
|
||||
|
||||
resolved_key = api_key or get_secret_str("REDUCTO_API_KEY")
|
||||
if resolved_key is None:
|
||||
raise ValueError(
|
||||
"Missing REDUCTO_API_KEY - set it in the environment or pass api_key to litellm.ocr()/litellm.aocr()"
|
||||
)
|
||||
|
||||
self._api_key = resolved_key
|
||||
self._api_base = (api_base or REDUCTO_API_BASE).rstrip("/")
|
||||
|
||||
return {
|
||||
"Authorization": f"Bearer {resolved_key}",
|
||||
"Content-Type": "application/json",
|
||||
**headers,
|
||||
}
|
||||
|
||||
def get_complete_url(
|
||||
self,
|
||||
api_base: Optional[str],
|
||||
model: str,
|
||||
optional_params: dict,
|
||||
litellm_params: Optional[dict] = None,
|
||||
**kwargs,
|
||||
) -> str:
|
||||
return "{}/parse".format((api_base or REDUCTO_API_BASE).rstrip("/"))
|
||||
|
||||
def _get_source_url(self, document: DocumentType, model: str) -> str:
|
||||
source_url = document.get("document_url") or document.get("image_url")
|
||||
if source_url is None:
|
||||
raise ValueError(
|
||||
"Reducto expected OCR preprocessing to produce document_url or image_url for model={}".format(
|
||||
model
|
||||
)
|
||||
)
|
||||
return source_url
|
||||
|
||||
def _ensure_file_id_sync(self, model: str, document: DocumentType) -> str:
|
||||
source_url = self._get_source_url(document=document, model=model)
|
||||
file_id, raw_bytes, mime = extract_file_id_or_bytes(source_url, model=model)
|
||||
if file_id is not None:
|
||||
return file_id
|
||||
if self._api_key is None:
|
||||
raise ValueError("Reducto API key was not initialized before OCR upload.")
|
||||
return upload_bytes_sync(
|
||||
raw_bytes=raw_bytes or b"",
|
||||
mime=mime,
|
||||
api_key=self._api_key,
|
||||
api_base=self._api_base,
|
||||
)
|
||||
|
||||
async def _ensure_file_id_async(self, model: str, document: DocumentType) -> str:
|
||||
source_url = self._get_source_url(document=document, model=model)
|
||||
file_id, raw_bytes, mime = extract_file_id_or_bytes(source_url, model=model)
|
||||
if file_id is not None:
|
||||
return file_id
|
||||
if self._api_key is None:
|
||||
raise ValueError("Reducto API key was not initialized before OCR upload.")
|
||||
return await upload_bytes_async(
|
||||
raw_bytes=raw_bytes or b"",
|
||||
mime=mime,
|
||||
api_key=self._api_key,
|
||||
api_base=self._api_base,
|
||||
)
|
||||
|
||||
def transform_ocr_response(
|
||||
self,
|
||||
model: str,
|
||||
raw_response: httpx.Response,
|
||||
logging_obj: Any,
|
||||
**kwargs,
|
||||
) -> OCRResponse:
|
||||
response_json = raw_response.json()
|
||||
result = response_json.get("result", response_json) or {}
|
||||
usage = response_json.get("usage", {}) or {}
|
||||
response = OCRResponse(
|
||||
pages=build_pages_from_reducto(result),
|
||||
model=model,
|
||||
usage_info=OCRUsageInfo(
|
||||
pages_processed=usage.get("num_pages"),
|
||||
credits=usage.get("credits"),
|
||||
),
|
||||
object="ocr",
|
||||
)
|
||||
response._hidden_params["reducto_raw"] = response_json
|
||||
return response
|
||||
|
||||
|
||||
class ReductoParseV3Config(_BaseReductoOCRConfig):
|
||||
def get_supported_ocr_params(self, model: str) -> list:
|
||||
return ["formatting", "retrieval", "settings"]
|
||||
|
||||
def transform_ocr_request(
|
||||
self,
|
||||
model: str,
|
||||
document: DocumentType,
|
||||
optional_params: dict,
|
||||
headers: dict,
|
||||
**kwargs,
|
||||
) -> OCRRequestData:
|
||||
file_id = self._ensure_file_id_sync(model=model, document=document)
|
||||
return OCRRequestData(data={"input": file_id, **optional_params}, files=None)
|
||||
|
||||
async def async_transform_ocr_request(
|
||||
self,
|
||||
model: str,
|
||||
document: DocumentType,
|
||||
optional_params: dict,
|
||||
headers: dict,
|
||||
**kwargs,
|
||||
) -> OCRRequestData:
|
||||
file_id = await self._ensure_file_id_async(model=model, document=document)
|
||||
return OCRRequestData(data={"input": file_id, **optional_params}, files=None)
|
||||
|
||||
|
||||
class ReductoParseLegacyConfig(_BaseReductoOCRConfig):
|
||||
def get_supported_ocr_params(self, model: str) -> list:
|
||||
return ["enhance"]
|
||||
|
||||
def _build_legacy_body(self, file_id: str, optional_params: dict) -> Dict[str, Any]:
|
||||
body: Dict[str, Any] = {"document_url": file_id}
|
||||
enhance = optional_params.get("enhance")
|
||||
if enhance is not None:
|
||||
body["options"] = {"enhance": enhance}
|
||||
return body
|
||||
|
||||
def transform_ocr_request(
|
||||
self,
|
||||
model: str,
|
||||
document: DocumentType,
|
||||
optional_params: dict,
|
||||
headers: dict,
|
||||
**kwargs,
|
||||
) -> OCRRequestData:
|
||||
file_id = self._ensure_file_id_sync(model=model, document=document)
|
||||
return OCRRequestData(
|
||||
data=self._build_legacy_body(
|
||||
file_id=file_id, optional_params=optional_params
|
||||
),
|
||||
files=None,
|
||||
)
|
||||
|
||||
async def async_transform_ocr_request(
|
||||
self,
|
||||
model: str,
|
||||
document: DocumentType,
|
||||
optional_params: dict,
|
||||
headers: dict,
|
||||
**kwargs,
|
||||
) -> OCRRequestData:
|
||||
file_id = await self._ensure_file_id_async(model=model, document=document)
|
||||
return OCRRequestData(
|
||||
data=self._build_legacy_body(
|
||||
file_id=file_id, optional_params=optional_params
|
||||
),
|
||||
files=None,
|
||||
)
|
||||
|
|
@ -27379,6 +27379,20 @@
|
|||
"supports_tool_choice": true,
|
||||
"source": "https://aws.amazon.com/bedrock/pricing/"
|
||||
},
|
||||
"reducto/parse-legacy": {
|
||||
"litellm_provider": "reducto",
|
||||
"mode": "ocr",
|
||||
"supported_endpoints": [
|
||||
"/v1/ocr"
|
||||
]
|
||||
},
|
||||
"reducto/parse-v3": {
|
||||
"litellm_provider": "reducto",
|
||||
"mode": "ocr",
|
||||
"supported_endpoints": [
|
||||
"/v1/ocr"
|
||||
]
|
||||
},
|
||||
"recraft/recraftv2": {
|
||||
"litellm_provider": "recraft",
|
||||
"mode": "image_generation",
|
||||
|
|
|
|||
|
|
@ -239,6 +239,7 @@ class ModelInfoBase(ProviderSpecificModelInfo, total=False):
|
|||
float
|
||||
] # video_generation tier: key output_cost_per_second_<resolution> (e.g. 1080p, 720p)
|
||||
ocr_cost_per_page: Optional[float] # for OCR models
|
||||
ocr_cost_per_credit: Optional[float] # for OCR models priced by credit
|
||||
annotation_cost_per_page: Optional[float] # for OCR models
|
||||
search_context_cost_per_query: Optional[
|
||||
SearchContextCostPerQuery
|
||||
|
|
@ -256,6 +257,7 @@ class ModelInfoBase(ProviderSpecificModelInfo, total=False):
|
|||
"chat",
|
||||
"audio_transcription",
|
||||
"responses",
|
||||
"ocr",
|
||||
]
|
||||
]
|
||||
tpm: Optional[int]
|
||||
|
|
@ -3191,6 +3193,7 @@ class LlmProviders(str, Enum):
|
|||
ANTHROPIC_TEXT = "anthropic_text"
|
||||
BYTEZ = "bytez"
|
||||
REPLICATE = "replicate"
|
||||
REDUCTO = "reducto"
|
||||
RUNWAYML = "runwayml"
|
||||
AWS_POLLY = "aws_polly"
|
||||
HUGGINGFACE = "huggingface"
|
||||
|
|
|
|||
|
|
@ -5909,6 +5909,7 @@ def _get_model_info_helper( # noqa: PLR0915
|
|||
tpm=_model_info.get("tpm", None),
|
||||
rpm=_model_info.get("rpm", None),
|
||||
ocr_cost_per_page=_model_info.get("ocr_cost_per_page", None),
|
||||
ocr_cost_per_credit=_model_info.get("ocr_cost_per_credit", None),
|
||||
annotation_cost_per_page=_model_info.get(
|
||||
"annotation_cost_per_page", None
|
||||
),
|
||||
|
|
@ -9128,6 +9129,18 @@ class ProviderConfigManager:
|
|||
|
||||
return get_vertex_ai_ocr_config(model=model)
|
||||
|
||||
if provider == litellm.LlmProviders.REDUCTO:
|
||||
from litellm.llms.reducto.ocr.transformation import (
|
||||
ReductoParseLegacyConfig,
|
||||
ReductoParseV3Config,
|
||||
)
|
||||
|
||||
if model == "parse-v3":
|
||||
return ReductoParseV3Config()
|
||||
if model == "parse-legacy":
|
||||
return ReductoParseLegacyConfig()
|
||||
return None
|
||||
|
||||
MistralOCRConfig = getattr(sys.modules[__name__], "MistralOCRConfig")
|
||||
PROVIDER_TO_CONFIG_MAP = {
|
||||
litellm.LlmProviders.MISTRAL: MistralOCRConfig,
|
||||
|
|
|
|||
82
tests/proxy_unit_tests/test_reducto_ocr_route.py
Normal file
82
tests/proxy_unit_tests/test_reducto_ocr_route.py
Normal file
|
|
@ -0,0 +1,82 @@
|
|||
import asyncio
|
||||
import os
|
||||
from unittest.mock import AsyncMock, patch
|
||||
|
||||
import litellm
|
||||
import pytest
|
||||
from fastapi.testclient import TestClient
|
||||
|
||||
from litellm.llms.base_llm.ocr.transformation import OCRPage, OCRResponse, OCRUsageInfo
|
||||
from litellm.proxy.proxy_server import app, initialize
|
||||
|
||||
|
||||
@pytest.fixture(scope="function")
|
||||
def fake_env_vars(monkeypatch):
|
||||
monkeypatch.setenv("OPENAI_API_KEY", "fake_openai_api_key")
|
||||
monkeypatch.setenv("OPENAI_API_BASE", "http://fake-openai-api-base")
|
||||
monkeypatch.setenv("AZURE_AI_API_BASE", "http://fake-azure-api-base")
|
||||
monkeypatch.setenv("AZURE_AI_API_KEY", "fake_azure_api_key")
|
||||
monkeypatch.setenv("AZURE_OPENAI_API_KEY", "fake_azure_openai_api_key")
|
||||
monkeypatch.setenv("AZURE_SWEDEN_API_BASE", "http://fake-azure-sweden-api-base")
|
||||
monkeypatch.setenv("AZURE_SWEDEN_API_KEY", "fake_azure_sweden_api_key")
|
||||
monkeypatch.setenv("REDIS_HOST", "localhost")
|
||||
|
||||
|
||||
@pytest.fixture(scope="function")
|
||||
def client_no_auth(fake_env_vars):
|
||||
from litellm.proxy.proxy_server import cleanup_router_config_variables
|
||||
|
||||
original_disable_aiohttp = litellm.disable_aiohttp_transport
|
||||
litellm.disable_aiohttp_transport = True
|
||||
litellm.in_memory_llm_clients_cache.flush_cache()
|
||||
cleanup_router_config_variables()
|
||||
|
||||
filepath = os.path.dirname(os.path.abspath(__file__))
|
||||
config_fp = os.path.join(filepath, "test_configs", "test_config_no_auth.yaml")
|
||||
asyncio.run(initialize(config=config_fp, debug=True))
|
||||
|
||||
try:
|
||||
yield TestClient(app)
|
||||
finally:
|
||||
litellm.disable_aiohttp_transport = original_disable_aiohttp
|
||||
litellm.in_memory_llm_clients_cache.flush_cache()
|
||||
|
||||
|
||||
def test_proxy_reducto_ocr_json_passthrough(client_no_auth):
|
||||
mocked_response = OCRResponse(
|
||||
pages=[OCRPage(index=0, markdown="Proxy OCR")],
|
||||
model="parse-v3",
|
||||
usage_info=OCRUsageInfo(pages_processed=1, credits=1),
|
||||
)
|
||||
|
||||
with patch(
|
||||
"litellm.proxy.proxy_server.llm_router.aocr",
|
||||
new=AsyncMock(return_value=mocked_response),
|
||||
) as mock_aocr:
|
||||
response = client_no_auth.post(
|
||||
"/v1/ocr",
|
||||
json={
|
||||
"model": "reducto/parse-v3",
|
||||
"document": {
|
||||
"type": "document_url",
|
||||
"document_url": "reducto://proxy.pdf",
|
||||
},
|
||||
"api_key": "proxy-key",
|
||||
"api_base": "https://platform.reducto.ai",
|
||||
},
|
||||
)
|
||||
|
||||
assert response.status_code == 200
|
||||
assert mock_aocr.await_count == 1
|
||||
assert mock_aocr.await_args.kwargs["model"] == "reducto/parse-v3"
|
||||
assert mock_aocr.await_args.kwargs["document"] == {
|
||||
"type": "document_url",
|
||||
"document_url": "reducto://proxy.pdf",
|
||||
}
|
||||
assert mock_aocr.await_args.kwargs["api_key"] == "proxy-key"
|
||||
assert mock_aocr.await_args.kwargs["api_base"] == "https://platform.reducto.ai"
|
||||
|
||||
response_body = response.json()
|
||||
assert response_body["object"] == "ocr"
|
||||
assert response_body["usage_info"]["credits"] == 1
|
||||
assert response_body["pages"][0]["markdown"] == "Proxy OCR"
|
||||
1
tests/test_litellm/llms/reducto/__init__.py
Normal file
1
tests/test_litellm/llms/reducto/__init__.py
Normal file
|
|
@ -0,0 +1 @@
|
|||
|
||||
73
tests/test_litellm/llms/reducto/test_cost.py
Normal file
73
tests/test_litellm/llms/reducto/test_cost.py
Normal file
|
|
@ -0,0 +1,73 @@
|
|||
import litellm
|
||||
|
||||
from litellm.cost_calculator import completion_cost
|
||||
from litellm.llms.base_llm.ocr.transformation import OCRPage, OCRResponse, OCRUsageInfo
|
||||
|
||||
|
||||
def test_ocr_cost_prefers_credit_pricing_when_pages_processed_is_none(monkeypatch):
|
||||
monkeypatch.setattr(
|
||||
litellm,
|
||||
"get_model_info",
|
||||
lambda model, custom_llm_provider=None: {"ocr_cost_per_credit": 0.003},
|
||||
)
|
||||
|
||||
response = OCRResponse(
|
||||
pages=[OCRPage(index=0, markdown="credit priced")],
|
||||
model="parse-v3",
|
||||
usage_info=OCRUsageInfo(pages_processed=None, credits=10),
|
||||
)
|
||||
|
||||
cost = completion_cost(
|
||||
completion_response=response,
|
||||
model="reducto/parse-v3",
|
||||
custom_llm_provider="reducto",
|
||||
call_type="ocr",
|
||||
)
|
||||
|
||||
assert cost == 0.03
|
||||
|
||||
|
||||
def test_ocr_cost_falls_back_to_page_pricing(monkeypatch):
|
||||
monkeypatch.setattr(
|
||||
litellm,
|
||||
"get_model_info",
|
||||
lambda model, custom_llm_provider=None: {"ocr_cost_per_page": 0.5},
|
||||
)
|
||||
|
||||
response = OCRResponse(
|
||||
pages=[OCRPage(index=0, markdown="page priced")],
|
||||
model="mistral-ocr-latest",
|
||||
usage_info=OCRUsageInfo(pages_processed=2),
|
||||
)
|
||||
|
||||
cost = completion_cost(
|
||||
completion_response=response,
|
||||
model="mistral/mistral-ocr-latest",
|
||||
custom_llm_provider="mistral",
|
||||
call_type="ocr",
|
||||
)
|
||||
|
||||
assert cost == 1.0
|
||||
|
||||
|
||||
def test_ocr_cost_returns_zero_when_no_pricing_and_no_pages(monkeypatch):
|
||||
monkeypatch.setattr(
|
||||
litellm,
|
||||
"get_model_info",
|
||||
lambda model, custom_llm_provider=None: {},
|
||||
)
|
||||
|
||||
response = OCRResponse(
|
||||
pages=[OCRPage(index=0, markdown="unpriced")],
|
||||
model="parse-v3",
|
||||
usage_info=OCRUsageInfo(pages_processed=None, credits=5),
|
||||
)
|
||||
|
||||
cost = completion_cost(
|
||||
completion_response=response,
|
||||
model="reducto/parse-v3",
|
||||
custom_llm_provider="reducto",
|
||||
call_type="ocr",
|
||||
)
|
||||
|
||||
assert cost == 0.0
|
||||
44
tests/test_litellm/llms/reducto/test_model_info.py
Normal file
44
tests/test_litellm/llms/reducto/test_model_info.py
Normal file
|
|
@ -0,0 +1,44 @@
|
|||
import uuid
|
||||
|
||||
import litellm
|
||||
|
||||
from litellm.utils import _invalidate_model_cost_lowercase_map
|
||||
|
||||
|
||||
def test_reducto_provider_registration():
|
||||
model, custom_llm_provider, _, _ = litellm.get_llm_provider(
|
||||
model="reducto/parse-v3"
|
||||
)
|
||||
|
||||
assert model == "parse-v3"
|
||||
assert custom_llm_provider == "reducto"
|
||||
|
||||
|
||||
def test_get_model_info_preserves_ocr_cost_per_credit():
|
||||
test_model_name = f"reducto/test-cost-propagation-{uuid.uuid4().hex[:12]}"
|
||||
previous_model_entry = litellm.model_cost.get(test_model_name)
|
||||
_invalidate_model_cost_lowercase_map()
|
||||
|
||||
try:
|
||||
litellm.register_model(
|
||||
{
|
||||
test_model_name: {
|
||||
"litellm_provider": "reducto",
|
||||
"mode": "ocr",
|
||||
"ocr_cost_per_credit": 0.003,
|
||||
}
|
||||
}
|
||||
)
|
||||
|
||||
model_info = litellm.get_model_info(
|
||||
model=test_model_name,
|
||||
custom_llm_provider="reducto",
|
||||
)
|
||||
|
||||
assert model_info.get("ocr_cost_per_credit") == 0.003
|
||||
finally:
|
||||
if previous_model_entry is None:
|
||||
litellm.model_cost.pop(test_model_name, None)
|
||||
else:
|
||||
litellm.model_cost[test_model_name] = previous_model_entry
|
||||
_invalidate_model_cost_lowercase_map()
|
||||
59
tests/test_litellm/llms/reducto/test_parse_legacy.py
Normal file
59
tests/test_litellm/llms/reducto/test_parse_legacy.py
Normal file
|
|
@ -0,0 +1,59 @@
|
|||
import json
|
||||
|
||||
import litellm
|
||||
import pytest
|
||||
|
||||
|
||||
@pytest.fixture()
|
||||
def disable_aiohttp_transport():
|
||||
original_disable_aiohttp = litellm.disable_aiohttp_transport
|
||||
litellm.disable_aiohttp_transport = True
|
||||
litellm.in_memory_llm_clients_cache.flush_cache()
|
||||
try:
|
||||
yield
|
||||
finally:
|
||||
litellm.disable_aiohttp_transport = original_disable_aiohttp
|
||||
litellm.in_memory_llm_clients_cache.flush_cache()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_parse_legacy_wraps_enhance_under_options(
|
||||
disable_aiohttp_transport, respx_mock
|
||||
):
|
||||
upload_route = respx_mock.post("https://platform.reducto.ai/upload").respond(
|
||||
json={"file_id": "reducto://legacy.pdf"}
|
||||
)
|
||||
parse_route = respx_mock.post("https://platform.reducto.ai/parse").respond(
|
||||
json={
|
||||
"usage": {"num_pages": 1, "credits": 1},
|
||||
"result": {
|
||||
"chunks": [
|
||||
{
|
||||
"content": "Legacy parse",
|
||||
"blocks": [{"content": "Legacy parse", "bbox": {"page": 1}}],
|
||||
}
|
||||
]
|
||||
},
|
||||
}
|
||||
)
|
||||
|
||||
response = await litellm.aocr(
|
||||
model="reducto/parse-legacy",
|
||||
document={
|
||||
"type": "file",
|
||||
"file": b"%PDF-1.4 legacy",
|
||||
"mime_type": "application/pdf",
|
||||
},
|
||||
api_key="legacy-key",
|
||||
api_base="https://platform.reducto.ai",
|
||||
enhance={"agentic": [{"type": "table"}]},
|
||||
)
|
||||
|
||||
assert upload_route.called
|
||||
assert parse_route.called
|
||||
request_body = json.loads(parse_route.calls[0].request.read())
|
||||
assert request_body == {
|
||||
"document_url": "reducto://legacy.pdf",
|
||||
"options": {"enhance": {"agentic": [{"type": "table"}]}},
|
||||
}
|
||||
assert response.pages[0].markdown == "Legacy parse"
|
||||
152
tests/test_litellm/llms/reducto/test_parse_v3.py
Normal file
152
tests/test_litellm/llms/reducto/test_parse_v3.py
Normal file
|
|
@ -0,0 +1,152 @@
|
|||
import json
|
||||
|
||||
import litellm
|
||||
import pytest
|
||||
|
||||
|
||||
def _reducto_parse_response() -> dict:
|
||||
return {
|
||||
"job_id": "job_123",
|
||||
"usage": {"num_pages": 3, "credits": 3},
|
||||
"result": {
|
||||
"chunks": [
|
||||
{
|
||||
"content": "Page 1 block A",
|
||||
"blocks": [
|
||||
{
|
||||
"content": "Page 1 block A",
|
||||
"bbox": {"page": 1},
|
||||
"kind": "text",
|
||||
}
|
||||
],
|
||||
},
|
||||
{
|
||||
"content": "Page 2 block A",
|
||||
"blocks": [
|
||||
{
|
||||
"content": "Page 2 block A",
|
||||
"bbox": {"page": 2},
|
||||
"kind": "table",
|
||||
}
|
||||
],
|
||||
},
|
||||
{
|
||||
"content": "Page 1 block B",
|
||||
"blocks": [
|
||||
{
|
||||
"content": "Page 1 block B",
|
||||
"bbox": {"page": 1},
|
||||
"kind": "text",
|
||||
}
|
||||
],
|
||||
},
|
||||
{
|
||||
"content": "Page 3 block A",
|
||||
"blocks": [
|
||||
{
|
||||
"content": "Page 3 block A",
|
||||
"bbox": {"page": 3},
|
||||
"kind": "figure",
|
||||
}
|
||||
],
|
||||
},
|
||||
]
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
@pytest.fixture()
|
||||
def disable_aiohttp_transport():
|
||||
original_disable_aiohttp = litellm.disable_aiohttp_transport
|
||||
litellm.disable_aiohttp_transport = True
|
||||
litellm.in_memory_llm_clients_cache.flush_cache()
|
||||
try:
|
||||
yield
|
||||
finally:
|
||||
litellm.disable_aiohttp_transport = original_disable_aiohttp
|
||||
litellm.in_memory_llm_clients_cache.flush_cache()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_parse_v3_file_upload_and_response_mapping(
|
||||
disable_aiohttp_transport, respx_mock
|
||||
):
|
||||
upload_route = respx_mock.post("https://platform.reducto.ai/upload").respond(
|
||||
json={"file_id": "reducto://uploaded.pdf"}
|
||||
)
|
||||
parse_route = respx_mock.post("https://platform.reducto.ai/parse").respond(
|
||||
json=_reducto_parse_response()
|
||||
)
|
||||
|
||||
response = await litellm.aocr(
|
||||
model="reducto/parse-v3",
|
||||
document={
|
||||
"type": "file",
|
||||
"file": b"%PDF-1.4 reducto",
|
||||
"mime_type": "application/pdf",
|
||||
},
|
||||
api_key="test-key",
|
||||
api_base="https://platform.reducto.ai",
|
||||
formatting={"table_output_format": "html"},
|
||||
retrieval={"chunk_mode": "section"},
|
||||
settings={"ocr_system": "standard"},
|
||||
)
|
||||
|
||||
assert upload_route.called
|
||||
assert parse_route.called
|
||||
assert len(upload_route.calls) == 1
|
||||
assert len(parse_route.calls) == 1
|
||||
|
||||
upload_request = upload_route.calls[0].request
|
||||
assert upload_request.headers["authorization"] == "Bearer test-key"
|
||||
assert "application/json" not in upload_request.headers["content-type"]
|
||||
upload_body = upload_request.read()
|
||||
assert b'filename="document"' in upload_body
|
||||
assert b"application/pdf" in upload_body
|
||||
|
||||
parse_request_body = json.loads(parse_route.calls[0].request.read())
|
||||
assert parse_request_body["input"] == "reducto://uploaded.pdf"
|
||||
assert parse_request_body["formatting"] == {"table_output_format": "html"}
|
||||
assert parse_request_body["retrieval"] == {"chunk_mode": "section"}
|
||||
assert parse_request_body["settings"] == {"ocr_system": "standard"}
|
||||
|
||||
assert response.usage_info is not None
|
||||
assert response.usage_info.credits == 3
|
||||
assert response.usage_info.pages_processed == 3
|
||||
assert len(response.pages) == 3
|
||||
assert response.pages[0].index == 0
|
||||
assert response.pages[0].markdown == "Page 1 block A\n\nPage 1 block B"
|
||||
assert getattr(response.pages[0], "blocks")[0]["bbox"]["page"] == 1
|
||||
assert response.pages[1].markdown == "Page 2 block A"
|
||||
assert response.pages[2].markdown == "Page 3 block A"
|
||||
assert response._hidden_params["reducto_raw"]["usage"]["credits"] == 3
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_parse_v3_reducto_id_passthrough_skips_upload(
|
||||
disable_aiohttp_transport, respx_mock
|
||||
):
|
||||
upload_route = respx_mock.post("https://platform.reducto.ai/upload").respond(
|
||||
json={"file_id": "reducto://should-not-upload.pdf"}
|
||||
)
|
||||
parse_route = respx_mock.post("https://platform.reducto.ai/parse").respond(
|
||||
json=_reducto_parse_response()
|
||||
)
|
||||
|
||||
response = await litellm.aocr(
|
||||
model="reducto/parse-v3",
|
||||
document={
|
||||
"type": "document_url",
|
||||
"document_url": "reducto://already-uploaded.pdf",
|
||||
},
|
||||
api_key="test-key",
|
||||
api_base="https://platform.reducto.ai",
|
||||
retrieval={"chunk_mode": "section"},
|
||||
)
|
||||
|
||||
assert not upload_route.called
|
||||
assert parse_route.called
|
||||
parse_request_body = json.loads(parse_route.calls[0].request.read())
|
||||
assert parse_request_body["input"] == "reducto://already-uploaded.pdf"
|
||||
assert parse_request_body["retrieval"]["chunk_mode"] == "section"
|
||||
assert response.pages[0].markdown.startswith("Page 1 block A")
|
||||
213
tests/test_litellm/llms/reducto/test_upload.py
Normal file
213
tests/test_litellm/llms/reducto/test_upload.py
Normal file
|
|
@ -0,0 +1,213 @@
|
|||
import json
|
||||
import os
|
||||
from unittest.mock import AsyncMock, Mock
|
||||
|
||||
import httpx
|
||||
import litellm
|
||||
import pytest
|
||||
|
||||
from litellm.llms.reducto.common import (
|
||||
extract_file_id_or_bytes,
|
||||
upload_bytes_async,
|
||||
upload_bytes_sync,
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture()
|
||||
def disable_aiohttp_transport(monkeypatch):
|
||||
original_disable_aiohttp = litellm.disable_aiohttp_transport
|
||||
litellm.disable_aiohttp_transport = True
|
||||
litellm.in_memory_llm_clients_cache.flush_cache()
|
||||
monkeypatch.setenv("REDUCTO_API_KEY", "env-reducto-key")
|
||||
try:
|
||||
yield
|
||||
finally:
|
||||
litellm.disable_aiohttp_transport = original_disable_aiohttp
|
||||
litellm.in_memory_llm_clients_cache.flush_cache()
|
||||
os.environ.pop("REDUCTO_API_KEY", None)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_parse_v3_rejects_plain_http_urls(disable_aiohttp_transport):
|
||||
with pytest.raises(litellm.BadRequestError, match="upload the file first"):
|
||||
await litellm.aocr(
|
||||
model="reducto/parse-v3",
|
||||
document={
|
||||
"type": "document_url",
|
||||
"document_url": "https://example.com/document.pdf",
|
||||
},
|
||||
api_key="test-key",
|
||||
api_base="https://platform.reducto.ai",
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_parse_v3_image_data_uri_upload_uses_image_mime(
|
||||
disable_aiohttp_transport, respx_mock
|
||||
):
|
||||
upload_route = respx_mock.post("https://custom.reducto.test/upload").respond(
|
||||
json={"file_id": "reducto://uploaded-image.png"}
|
||||
)
|
||||
parse_route = respx_mock.post("https://custom.reducto.test/parse").respond(
|
||||
json={
|
||||
"usage": {"num_pages": 1, "credits": 1},
|
||||
"result": {
|
||||
"chunks": [
|
||||
{
|
||||
"content": "Image OCR",
|
||||
"blocks": [{"content": "Image OCR", "bbox": {"page": 1}}],
|
||||
}
|
||||
]
|
||||
},
|
||||
}
|
||||
)
|
||||
|
||||
response = await litellm.aocr(
|
||||
model="reducto/parse-v3",
|
||||
document={
|
||||
"type": "file",
|
||||
"file": b"\x89PNG\r\n\x1a\npng",
|
||||
"mime_type": "image/png",
|
||||
},
|
||||
api_key="programmatic-key",
|
||||
api_base="https://custom.reducto.test/",
|
||||
)
|
||||
|
||||
assert upload_route.called
|
||||
assert parse_route.called
|
||||
upload_request = upload_route.calls[0].request
|
||||
assert upload_request.headers["authorization"] == "Bearer programmatic-key"
|
||||
assert b"image/png" in upload_request.read()
|
||||
|
||||
parse_request_body = json.loads(parse_route.calls[0].request.read())
|
||||
assert parse_request_body["input"] == "reducto://uploaded-image.png"
|
||||
assert response.pages[0].markdown == "Image OCR"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_parse_v3_uses_programmatic_api_key_over_env(
|
||||
disable_aiohttp_transport, respx_mock
|
||||
):
|
||||
upload_route = respx_mock.post("https://platform.reducto.ai/upload").respond(
|
||||
json={"file_id": "reducto://uploaded.pdf"}
|
||||
)
|
||||
parse_route = respx_mock.post("https://platform.reducto.ai/parse").respond(
|
||||
json={
|
||||
"usage": {"num_pages": 1, "credits": 1},
|
||||
"result": {
|
||||
"chunks": [
|
||||
{
|
||||
"content": "Programmatic auth",
|
||||
"blocks": [
|
||||
{"content": "Programmatic auth", "bbox": {"page": 1}}
|
||||
],
|
||||
}
|
||||
]
|
||||
},
|
||||
}
|
||||
)
|
||||
|
||||
await litellm.aocr(
|
||||
model="reducto/parse-v3",
|
||||
document={
|
||||
"type": "file",
|
||||
"file": b"%PDF-1.4 auth",
|
||||
"mime_type": "application/pdf",
|
||||
},
|
||||
api_key="passed-key",
|
||||
api_base="https://platform.reducto.ai",
|
||||
)
|
||||
|
||||
assert upload_route.calls[0].request.headers["authorization"] == "Bearer passed-key"
|
||||
assert parse_route.calls[0].request.headers["authorization"] == "Bearer passed-key"
|
||||
|
||||
|
||||
def test_upload_bytes_sync_uses_shared_client(monkeypatch):
|
||||
captured = {}
|
||||
|
||||
def fake_post(*, url, headers, files, timeout):
|
||||
captured["url"] = url
|
||||
captured["headers"] = headers
|
||||
captured["files"] = files
|
||||
captured["timeout"] = timeout
|
||||
return httpx.Response(
|
||||
200,
|
||||
json={"file_id": "reducto://sync-upload"},
|
||||
request=httpx.Request("POST", url),
|
||||
)
|
||||
|
||||
sync_post = Mock(side_effect=fake_post)
|
||||
monkeypatch.setattr(litellm.module_level_client, "post", sync_post)
|
||||
|
||||
class ForbiddenSyncClient:
|
||||
def __init__(self, *args, **kwargs):
|
||||
raise AssertionError("should not construct")
|
||||
|
||||
monkeypatch.setattr(httpx, "Client", ForbiddenSyncClient)
|
||||
|
||||
file_id = upload_bytes_sync(
|
||||
raw_bytes=b"%PDF-1.4 sync",
|
||||
mime="application/pdf",
|
||||
api_key="sync-key",
|
||||
api_base="https://sync.reducto.test/",
|
||||
)
|
||||
|
||||
assert file_id == "reducto://sync-upload"
|
||||
sync_post.assert_called_once()
|
||||
assert captured["url"] == "https://sync.reducto.test/upload"
|
||||
assert captured["headers"] == {"Authorization": "Bearer sync-key"}
|
||||
assert captured["files"]["file"] == (
|
||||
"document",
|
||||
b"%PDF-1.4 sync",
|
||||
"application/pdf",
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_upload_bytes_async_uses_shared_aclient(monkeypatch):
|
||||
captured = {}
|
||||
|
||||
async def fake_post(*, url, headers, files, timeout):
|
||||
captured["url"] = url
|
||||
captured["headers"] = headers
|
||||
captured["files"] = files
|
||||
captured["timeout"] = timeout
|
||||
return httpx.Response(
|
||||
200,
|
||||
json={"file_id": "reducto://async-upload"},
|
||||
request=httpx.Request("POST", url),
|
||||
)
|
||||
|
||||
async_post = AsyncMock(side_effect=fake_post)
|
||||
monkeypatch.setattr(litellm.module_level_aclient, "post", async_post)
|
||||
|
||||
class ForbiddenAsyncClient:
|
||||
def __init__(self, *args, **kwargs):
|
||||
raise AssertionError("should not construct")
|
||||
|
||||
monkeypatch.setattr(httpx, "AsyncClient", ForbiddenAsyncClient)
|
||||
|
||||
file_id = await upload_bytes_async(
|
||||
raw_bytes=b"%PDF-1.4 async",
|
||||
mime="application/pdf",
|
||||
api_key="async-key",
|
||||
api_base="https://async.reducto.test/",
|
||||
)
|
||||
|
||||
assert file_id == "reducto://async-upload"
|
||||
async_post.assert_awaited_once()
|
||||
assert captured["url"] == "https://async.reducto.test/upload"
|
||||
assert captured["headers"] == {"Authorization": "Bearer async-key"}
|
||||
assert captured["files"]["file"] == (
|
||||
"document",
|
||||
b"%PDF-1.4 async",
|
||||
"application/pdf",
|
||||
)
|
||||
|
||||
|
||||
def test_extract_file_id_or_bytes_raises_on_malformed_data_uri():
|
||||
with pytest.raises(litellm.BadRequestError, match="Invalid Reducto data URI"):
|
||||
extract_file_id_or_bytes("data:application/pdf", model="reducto/parse-v3")
|
||||
|
||||
with pytest.raises(litellm.BadRequestError, match="Invalid Reducto base64 payload"):
|
||||
extract_file_id_or_bytes("data:;base64,!!!not-base64", model="reducto/parse-v3")
|
||||
|
|
@ -675,6 +675,7 @@ def test_aaamodel_prices_and_context_window_json_is_valid():
|
|||
"input_dbu_cost_per_token": {"type": "number"},
|
||||
"annotation_cost_per_page": {"type": "number"},
|
||||
"ocr_cost_per_page": {"type": "number"},
|
||||
"ocr_cost_per_credit": {"type": "number"},
|
||||
"code_interpreter_cost_per_session": {"type": "number"},
|
||||
"inference_geo": {"type": "string"},
|
||||
"litellm_provider": {"type": "string"},
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue