From 78a195b147fb73628ba5a6737e236095b304c97e Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Sat, 29 Aug 2026 12:02:05 -0700 Subject: [PATCH] wip --- litellm/llms/reducto/ocr/transformation.py | 18 +- litellm/ocr/main.py | 22 +- tests/test_litellm/_fixture_recorder.py | 321 ++++++++ .../llms/reducto/test_parse_legacy.py | 17 +- .../llms/reducto/test_parse_v3.py | 31 +- tests/test_litellm/ocr/conftest.py | 21 +- tests/test_litellm/ocr/fixture_models.py | 34 - tests/test_litellm/ocr/generate_fixtures.py | 692 +++++++++++------- tests/test_litellm/ocr/test_rust_bridge.py | 80 +- tests/test_litellm/ocr/test_sdk_parity.py | 28 +- tests/test_litellm/parity/compare.py | 20 +- tests/test_litellm/parity/runner.py | 5 + tests/test_litellm/test_fixture_recorder.py | 106 +++ 13 files changed, 1029 insertions(+), 366 deletions(-) create mode 100644 tests/test_litellm/_fixture_recorder.py delete mode 100644 tests/test_litellm/ocr/fixture_models.py create mode 100644 tests/test_litellm/test_fixture_recorder.py diff --git a/litellm/llms/reducto/ocr/transformation.py b/litellm/llms/reducto/ocr/transformation.py index a7216e4ec40..109875c8999 100644 --- a/litellm/llms/reducto/ocr/transformation.py +++ b/litellm/llms/reducto/ocr/transformation.py @@ -1,4 +1,4 @@ -from typing import TYPE_CHECKING, Any, Final +from typing import TYPE_CHECKING, Final import httpx @@ -150,8 +150,8 @@ class _BaseReductoOCRConfig(BaseOCRConfig): class ReductoParseV3Config(_BaseReductoOCRConfig): - def get_supported_ocr_params(self, model: str) -> list: - return ["formatting", "retrieval", "settings"] + def get_supported_ocr_params(self, model: str) -> list[str]: + return ["enhance", "retrieval", "formatting", "spreadsheet", "settings"] def transform_ocr_request( self, @@ -187,15 +187,11 @@ class ReductoParseV3Config(_BaseReductoOCRConfig): class ReductoParseLegacyConfig(_BaseReductoOCRConfig): - def get_supported_ocr_params(self, model: str) -> list: - return ["enhance"] + def get_supported_ocr_params(self, model: str) -> list[str]: + return ["options", "advanced_options", "experimental_options", "priority"] - def _build_legacy_body(self, file_id: str, optional_params: dict) -> dict[str, Any]: - body: Final[dict[str, Any]] = {"document_url": file_id} - enhance: Final = optional_params.get("enhance") - if enhance is not None: - body["options"] = {"enhance": enhance} - return body + def _build_legacy_body(self, file_id: str, optional_params: dict[str, object]) -> dict[str, object]: + return {"document_url": file_id, **optional_params} def transform_ocr_request( self, diff --git a/litellm/ocr/main.py b/litellm/ocr/main.py index b918f013700..96836937817 100644 --- a/litellm/ocr/main.py +++ b/litellm/ocr/main.py @@ -62,8 +62,7 @@ class _PreparedRustOCRCall: _RUST_OCR_PROVIDERS: Final = { "mistral", - "azure_ai", - "vertex_ai", + "reducto", } @@ -191,6 +190,11 @@ def _prepare_ocr_request( def _rust_ocr_supported(prepared_request: _PreparedOCRRequest) -> bool: if prepared_request.optional_params.get(OCR_REQUEST_FORMAT_PARAM) == "native": return False + if prepared_request.custom_llm_provider == "reducto": + source: Final = prepared_request.document.get( + "document_url" + ) or prepared_request.document.get("image_url") + return prepared_request.model == "parse-v3" and isinstance(source, str) and source.startswith("reducto://") return prepared_request.custom_llm_provider in _RUST_OCR_PROVIDERS @@ -300,7 +304,17 @@ def _run_rust_ocr( ) if rust_response is None: return None - return OCRResponse.model_validate(rust_response) + return _ocr_response_from_rust(rust_response) + + +def _ocr_response_from_rust(rust_response: dict[str, object]) -> OCRResponse: + reducto_raw: Final = rust_response.get("reducto_raw") + response: Final = OCRResponse.model_validate( + {key: value for key, value in rust_response.items() if key != "reducto_raw"} + ) + if isinstance(reducto_raw, dict): + response._hidden_params["reducto_raw"] = reducto_raw + return response async def _run_rust_aocr( @@ -325,7 +339,7 @@ async def _run_rust_aocr( ) if rust_response is None: return None - return OCRResponse.model_validate(rust_response) + return _ocr_response_from_rust(rust_response) @client diff --git a/tests/test_litellm/_fixture_recorder.py b/tests/test_litellm/_fixture_recorder.py new file mode 100644 index 00000000000..fa9d73d47dd --- /dev/null +++ b/tests/test_litellm/_fixture_recorder.py @@ -0,0 +1,321 @@ +from __future__ import annotations + +import argparse +import hashlib +import json +import queue +import threading +from collections.abc import Callable, Generator, Mapping +from contextlib import contextmanager +from dataclasses import dataclass +from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer +from pathlib import Path +from typing import Final, cast + +import httpx +from pydantic import BaseModel, ConfigDict, TypeAdapter + +from tests.test_litellm._json_fs_cache import JsonFileCache, canonical_json + +JSON_OBJECT: Final = TypeAdapter(dict[str, object]) + + +class ProviderWireRequest(BaseModel): + model_config = ConfigDict(frozen=True) + + method: str + path: str + body: dict[str, object] + + +class FixtureRequest(BaseModel): + model_config = ConfigDict(frozen=True) + + provider: str + sdk_kwargs: dict[str, object] + provider_request: ProviderWireRequest + + +class FixtureResponse(BaseModel): + model_config = ConfigDict(frozen=True) + + status_code: int + headers: dict[str, str] + body: dict[str, object] + + +class Fixture(BaseModel): + model_config = ConfigDict(frozen=True) + + request: FixtureRequest + response: FixtureResponse + + +@dataclass(frozen=True, slots=True) +class ProviderSpec: + name: str + model: str + upstream_base: str + api_key: str | None + upstream_model: str | None = None + + +@dataclass(frozen=True, slots=True) +class RecorderResult: + request: FixtureRequest + response: FixtureResponse | None + cache_hit: bool + + +@dataclass(frozen=True, slots=True) +class GeneratorArgs: + providers: tuple[str, ...] + examples: int + fixture_dir: Path | None + requests_only: bool + responses_only: bool + + +class _RecordingProvider(ThreadingHTTPServer): + daemon_threads = True + + def __init__( + self, + spec: ProviderSpec, + sdk_kwargs: dict[str, object], + cache: JsonFileCache, + requests_only: bool, + ) -> None: + super().__init__(("127.0.0.1", 0), _RecordingHandler) + self.spec: Final = spec + self.sdk_kwargs: Final = sdk_kwargs + self.cache: Final = cache + self.requests_only: Final = requests_only + self.results: queue.Queue[RecorderResult] = queue.Queue() + + @property + def url(self) -> str: + return f"http://127.0.0.1:{self.server_address[1]}" + + def take_result(self) -> RecorderResult: + try: + return self.results.get(timeout=5) + except queue.Empty as error: + raise RuntimeError("successful SDK call did not produce a recorder result") from error + + +class _RecordingHandler(BaseHTTPRequestHandler): + protocol_version = "HTTP/1.1" + + def do_POST(self) -> None: + provider: Final = self.server + assert isinstance(provider, _RecordingProvider) + length: Final = int(self.headers.get("content-length") or "0") + body: Final = JSON_OBJECT.validate_json(self.rfile.read(length)) + fixture_request: Final = FixtureRequest( + provider=provider.spec.name, + sdk_kwargs=provider.sdk_kwargs, + provider_request=ProviderWireRequest(method=self.command, path=self.path, body=body), + ) + cache_key: Final = fixture_cache_key( + provider.spec.name, + fixture_request.sdk_kwargs, + fixture_request.provider_request, + ) + cached_value: Final = provider.cache.get(cache_key) + if cached_value is not None: + cached_request: Final = FixtureRequest.model_validate(cached_value["request"]) + raw_cached_response: Final = cached_value.get("response") + if raw_cached_response is not None: + cached_response: Final = FixtureResponse.model_validate(raw_cached_response) + provider.results.put(RecorderResult(request=cached_request, response=cached_response, cache_hit=True)) + self._send_fixture_response(cached_response) + return + if provider.requests_only: + provider.results.put(RecorderResult(request=cached_request, response=None, cache_hit=True)) + self._send_response(200, {"content-type": "application/json"}, b"{}") + return + + if provider.requests_only: + provider.results.put(RecorderResult(request=fixture_request, response=None, cache_hit=False)) + self._send_response(200, {"content-type": "application/json"}, b"{}") + return + + upstream_url: Final = f"{provider.spec.upstream_base.rstrip('/')}{self.path}" + forwarded_headers: Final = { + name: value + for name, value in self.headers.items() + if name.lower() not in {"host", "content-length", "accept-encoding", "x-parity-case"} + } + upstream_body: Final = ( + {**body, "model": provider.spec.upstream_model} if provider.spec.upstream_model is not None else body + ) + try: + upstream_response: Final = httpx.post( + upstream_url, + headers=forwarded_headers, + content=json.dumps(upstream_body, separators=(",", ":")), + timeout=120, + ) + except httpx.HTTPError as error: + error_body: Final = json.dumps({"error": str(error)}).encode() + self._send_response(502, {"content-type": "application/json"}, error_body) + return + + raw_content_type: Final = cast(object, upstream_response.headers.get("content-type", "application/json")) + content_type: Final = raw_content_type if isinstance(raw_content_type, str) else "application/json" + response_headers: Final = {"content-type": content_type.split(";", 1)[0]} + if not upstream_response.is_success: + self._send_response(upstream_response.status_code, response_headers, upstream_response.content) + return + + upstream_response_body: Final = JSON_OBJECT.validate_json(upstream_response.content) + fixture_response: Final = FixtureResponse( + status_code=upstream_response.status_code, + headers=response_headers, + body=upstream_response_body, + ) + provider.results.put(RecorderResult(request=fixture_request, response=fixture_response, cache_hit=False)) + self._send_fixture_response(fixture_response) + + def _send_fixture_response(self, response: FixtureResponse) -> None: + response_body: Final = json.dumps(response.body, separators=(",", ":")).encode() + self._send_response(response.status_code, response.headers, response_body) + + def _send_response(self, status_code: int, headers: Mapping[str, str], body: bytes) -> None: + self.send_response(status_code) + for name, value in headers.items(): + self.send_header(name, value) + self.send_header("content-length", str(len(body))) + self.end_headers() + self.wfile.write(body) + + def log_message(self, format: str, *args: object) -> None: + return + + +@contextmanager +def _recording_provider( + spec: ProviderSpec, + sdk_kwargs: dict[str, object], + cache: JsonFileCache, + requests_only: bool, +) -> Generator[_RecordingProvider]: + server: Final = _RecordingProvider( + spec=spec, + sdk_kwargs=sdk_kwargs, + cache=cache, + requests_only=requests_only, + ) + thread: Final = threading.Thread(target=server.serve_forever, daemon=True) + thread.start() + try: + yield server + finally: + server.shutdown() + server.server_close() + thread.join(timeout=5) + + +def fixture_cache_key( + provider: str, + sdk_kwargs: dict[str, object], + request: ProviderWireRequest, +) -> dict[str, object]: + return { + "provider": provider, + "sdk_kwargs": sdk_kwargs, + "request": request.model_dump(mode="json"), + } + + +def record_case( + spec: ProviderSpec, + root: Path, + sdk_kwargs: dict[str, object], + requests_only: bool, + sdk_call: Callable[..., object], +) -> RecorderResult: + cache: Final = JsonFileCache(root / spec.name) + with _recording_provider( + spec=spec, + sdk_kwargs=sdk_kwargs, + cache=cache, + requests_only=requests_only, + ) as recorder: + try: + sdk_call(api_base=recorder.url, api_key=spec.api_key, **sdk_kwargs) + except Exception: + if not requests_only: + raise + result: Final = recorder.take_result() + + if not result.cache_hit: + value: Final = ( + Fixture(request=result.request, response=result.response).model_dump(mode="json") + if result.response is not None + else {"request": result.request.model_dump(mode="json")} + ) + cache.put( + fixture_cache_key(spec.name, result.request.sdk_kwargs, result.request.provider_request), + value, + ) + return result + + +def pending_requests(cache: JsonFileCache) -> tuple[FixtureRequest, ...]: + return tuple( + FixtureRequest.model_validate(value["request"]) for value in cache.values() if value.get("response") is None + ) + + +def fill_missing_responses( + specs: tuple[ProviderSpec, ...], + root: Path, + sdk_call: Callable[..., object], +) -> tuple[RecorderResult, ...]: + return tuple( + record_case(spec, root, request.sdk_kwargs, requests_only=False, sdk_call=sdk_call) + for spec in specs + for request in pending_requests(JsonFileCache(root / spec.name)) + if request.sdk_kwargs.get("model") == spec.model + ) + + +def parse_generator_args(provider_names: tuple[str, ...]) -> GeneratorArgs: + parser: Final = argparse.ArgumentParser() + parser.add_argument("--provider", action="append", choices=provider_names) + parser.add_argument("--examples", type=int, default=4) + parser.add_argument("--fixture-dir", type=Path) + mode: Final = parser.add_mutually_exclusive_group() + mode.add_argument("--requests-only", action="store_true", help="record deterministic requests without API calls") + mode.add_argument("--responses-only", action="store_true", help="fill responses for saved pending requests") + namespace: Final = parser.parse_args() + providers: Final = cast(list[str] | None, namespace.provider) + return GeneratorArgs( + providers=tuple(providers) if providers else provider_names, + examples=cast(int, namespace.examples), + fixture_dir=cast(Path | None, namespace.fixture_dir), + requests_only=cast(bool, namespace.requests_only), + responses_only=cast(bool, namespace.responses_only), + ) + + +def fixture_directory(configured: Path | None, env_value: str | None, default: Path) -> Path: + return (configured or Path(env_value or default)).expanduser() + + +def recorded_fixtures(directory: Path) -> tuple[Fixture, ...]: + return tuple( + Fixture.model_validate(raw_fixture) + for raw_fixture in JsonFileCache(directory).values() + if raw_fixture.get("response") is not None + ) + + +def fixture_id(fixture: Fixture) -> str: + raw_model: Final = fixture.request.sdk_kwargs.get("model") + model: Final = raw_model if isinstance(raw_model, str) else "unknown-model" + request_json: Final = canonical_json(fixture.request.provider_request.model_dump(mode="json")) + digest: Final = hashlib.sha256(request_json.encode("utf-8")).hexdigest()[:8] + return f"{fixture.request.provider}-{model.rsplit('/', 1)[-1]}-{digest}" diff --git a/tests/test_litellm/llms/reducto/test_parse_legacy.py b/tests/test_litellm/llms/reducto/test_parse_legacy.py index db19460baa3..823b7dff1f7 100644 --- a/tests/test_litellm/llms/reducto/test_parse_legacy.py +++ b/tests/test_litellm/llms/reducto/test_parse_legacy.py @@ -1,8 +1,9 @@ import json -import litellm import pytest +import litellm + @pytest.fixture() def disable_aiohttp_transport(): @@ -17,9 +18,7 @@ def disable_aiohttp_transport(): @pytest.mark.asyncio -async def test_parse_legacy_wraps_enhance_under_options( - disable_aiohttp_transport, respx_mock -): +async def test_parse_legacy_forwards_legacy_option_groups(disable_aiohttp_transport, respx_mock): upload_route = respx_mock.post("https://platform.reducto.ai/upload").respond( json={"file_id": "reducto://legacy.pdf"} ) @@ -46,7 +45,10 @@ async def test_parse_legacy_wraps_enhance_under_options( }, api_key="legacy-key", api_base="https://platform.reducto.ai", - enhance={"agentic": [{"type": "table"}]}, + options={"ocr_mode": "agentic", "chunking": {"chunk_mode": "section"}}, + advanced_options={"table_output_format": "html", "ocr_system": "highres"}, + experimental_options={"enable_checkboxes": True}, + priority=False, ) assert upload_route.called @@ -54,6 +56,9 @@ async def test_parse_legacy_wraps_enhance_under_options( request_body = json.loads(parse_route.calls[0].request.read()) assert request_body == { "document_url": "reducto://legacy.pdf", - "options": {"enhance": {"agentic": [{"type": "table"}]}}, + "options": {"ocr_mode": "agentic", "chunking": {"chunk_mode": "section"}}, + "advanced_options": {"table_output_format": "html", "ocr_system": "highres"}, + "experimental_options": {"enable_checkboxes": True}, + "priority": False, } assert response.pages[0].markdown == "Legacy parse" diff --git a/tests/test_litellm/llms/reducto/test_parse_v3.py b/tests/test_litellm/llms/reducto/test_parse_v3.py index 140b9737dc0..07b8a732d7c 100644 --- a/tests/test_litellm/llms/reducto/test_parse_v3.py +++ b/tests/test_litellm/llms/reducto/test_parse_v3.py @@ -1,8 +1,9 @@ import json -import litellm import pytest +import litellm + def _reducto_parse_response() -> dict: return { @@ -68,15 +69,11 @@ def disable_aiohttp_transport(): @pytest.mark.asyncio -async def test_parse_v3_file_upload_and_response_mapping( - disable_aiohttp_transport, respx_mock -): +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() - ) + parse_route = respx_mock.post("https://platform.reducto.ai/parse").respond(json=_reducto_parse_response()) response = await litellm.aocr( model="reducto/parse-v3", @@ -87,8 +84,10 @@ async def test_parse_v3_file_upload_and_response_mapping( }, api_key="test-key", api_base="https://platform.reducto.ai", + enhance={"agentic": [{"scope": "table", "mode": "max"}]}, formatting={"table_output_format": "html"}, - retrieval={"chunk_mode": "section"}, + retrieval={"chunking": {"chunk_mode": "section"}}, + spreadsheet={"clustering": "fast"}, settings={"ocr_system": "standard"}, ) @@ -106,8 +105,10 @@ async def test_parse_v3_file_upload_and_response_mapping( parse_request_body = json.loads(parse_route.calls[0].request.read()) assert parse_request_body["input"] == "reducto://uploaded.pdf" + assert parse_request_body["enhance"] == {"agentic": [{"scope": "table", "mode": "max"}]} assert parse_request_body["formatting"] == {"table_output_format": "html"} - assert parse_request_body["retrieval"] == {"chunk_mode": "section"} + assert parse_request_body["retrieval"] == {"chunking": {"chunk_mode": "section"}} + assert parse_request_body["spreadsheet"] == {"clustering": "fast"} assert parse_request_body["settings"] == {"ocr_system": "standard"} assert response.usage_info is not None @@ -123,15 +124,11 @@ async def test_parse_v3_file_upload_and_response_mapping( @pytest.mark.asyncio -async def test_parse_v3_reducto_id_passthrough_skips_upload( - disable_aiohttp_transport, respx_mock -): +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() - ) + parse_route = respx_mock.post("https://platform.reducto.ai/parse").respond(json=_reducto_parse_response()) response = await litellm.aocr( model="reducto/parse-v3", @@ -141,12 +138,12 @@ async def test_parse_v3_reducto_id_passthrough_skips_upload( }, api_key="test-key", api_base="https://platform.reducto.ai", - retrieval={"chunk_mode": "section"}, + retrieval={"chunking": {"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 parse_request_body["retrieval"]["chunking"]["chunk_mode"] == "section" assert response.pages[0].markdown.startswith("Page 1 block A") diff --git a/tests/test_litellm/ocr/conftest.py b/tests/test_litellm/ocr/conftest.py index 6e0b033e28d..6da67be1984 100644 --- a/tests/test_litellm/ocr/conftest.py +++ b/tests/test_litellm/ocr/conftest.py @@ -1,14 +1,12 @@ from __future__ import annotations -import hashlib import os from pathlib import Path from typing import Final import pytest -from tests.test_litellm._json_fs_cache import JsonFileCache, canonical_json -from tests.test_litellm.ocr.fixture_models import OcrFixture +from tests.test_litellm._fixture_recorder import fixture_id, recorded_fixtures FIXTURE_DIR_ENV: Final = "LITELLM_OCR_FIXTURE_DIR" pytest_plugins: Final = ("tests.test_litellm.parity.pytest_plugin",) @@ -19,23 +17,10 @@ def _fixture_directory() -> Path: return Path(configured).expanduser() if configured else Path(__file__).with_name(".fixtures") -def _recorded_fixtures() -> tuple[OcrFixture, ...]: - raw_fixtures: Final = JsonFileCache(_fixture_directory()).values() - return tuple(OcrFixture.model_validate(raw_fixture) for raw_fixture in raw_fixtures) - - -def _fixture_id(fixture: OcrFixture) -> str: - raw_model: Final = fixture.request.sdk_kwargs.get("model") - model: Final = raw_model if isinstance(raw_model, str) else "unknown-model" - request_json: Final = canonical_json(fixture.request.provider_request.model_dump(mode="json")) - digest: Final = hashlib.sha256(request_json.encode("utf-8")).hexdigest()[:8] - return f"{fixture.request.provider}-{model.rsplit('/', 1)[-1]}-{digest}" - - def pytest_generate_tests(metafunc: pytest.Metafunc) -> None: if "ocr_fixture" not in metafunc.fixturenames: return - fixtures: Final = _recorded_fixtures() + fixtures: Final = recorded_fixtures(_fixture_directory()) if not fixtures: metafunc.parametrize( "ocr_fixture", @@ -48,4 +33,4 @@ def pytest_generate_tests(metafunc: pytest.Metafunc) -> None: ), ) return - metafunc.parametrize("ocr_fixture", fixtures, ids=tuple(_fixture_id(fixture) for fixture in fixtures)) + metafunc.parametrize("ocr_fixture", fixtures, ids=tuple(fixture_id(fixture) for fixture in fixtures)) diff --git a/tests/test_litellm/ocr/fixture_models.py b/tests/test_litellm/ocr/fixture_models.py deleted file mode 100644 index bd2c73ffc8e..00000000000 --- a/tests/test_litellm/ocr/fixture_models.py +++ /dev/null @@ -1,34 +0,0 @@ -from __future__ import annotations - -from pydantic import BaseModel, ConfigDict - - -class ProviderWireRequest(BaseModel): - model_config = ConfigDict(frozen=True) - - method: str - path: str - body: dict[str, object] - - -class OcrFixtureRequest(BaseModel): - model_config = ConfigDict(frozen=True) - - provider: str - sdk_kwargs: dict[str, object] - provider_request: ProviderWireRequest - - -class OcrFixtureResponse(BaseModel): - model_config = ConfigDict(frozen=True) - - status_code: int - headers: dict[str, str] - body: dict[str, object] - - -class OcrFixture(BaseModel): - model_config = ConfigDict(frozen=True) - - request: OcrFixtureRequest - response: OcrFixtureResponse diff --git a/tests/test_litellm/ocr/generate_fixtures.py b/tests/test_litellm/ocr/generate_fixtures.py index 8783cf74ac6..afcac1c4bfe 100644 --- a/tests/test_litellm/ocr/generate_fixtures.py +++ b/tests/test_litellm/ocr/generate_fixtures.py @@ -1,16 +1,9 @@ from __future__ import annotations -import argparse import base64 -import json import logging import os -import queue -import threading -from collections.abc import Callable, Generator -from contextlib import contextmanager -from dataclasses import dataclass -from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer +from collections.abc import Callable from pathlib import Path from typing import Final, cast from urllib.parse import quote @@ -19,178 +12,341 @@ import httpx from dotenv import load_dotenv from hypothesis import given, settings from hypothesis import strategies as st +from hypothesis.strategies import SearchStrategy from pydantic import TypeAdapter import litellm -from tests.test_litellm._json_fs_cache import JsonFileCache -from tests.test_litellm.ocr.fixture_models import ( - OcrFixture, - OcrFixtureRequest, - OcrFixtureResponse, - ProviderWireRequest, +from litellm.llms.reducto.common import ( + REDUCTO_API_BASE, + extract_file_id_or_bytes, + upload_bytes_sync, +) +from tests.test_litellm._fixture_recorder import ( + ProviderSpec, + RecorderResult, + fill_missing_responses, + fixture_directory, + parse_generator_args, + record_case, ) FIXTURE_DIR_ENV: Final = "LITELLM_OCR_FIXTURE_DIR" JSON_OBJECT: Final = TypeAdapter(dict[str, object]) -PROVIDER_NAMES: Final = ("mistral", "azure_ai", "vertex_ai") +PROVIDER_NAMES: Final = ("mistral", "reducto") LOGGER: Final = logging.getLogger(__name__) _TEXT: Final = st.from_regex(r"[A-Za-z0-9 ]{1,24}", fullmatch=True) -_OPTIONS: Final = st.fixed_dictionaries( - { - "pages": st.just([0]), - "include_image_base64": st.booleans(), - "image_limit": st.integers(min_value=1, max_value=4), - "image_min_size": st.integers(min_value=0, max_value=64), - "extract_header": st.booleans(), - "extract_footer": st.booleans(), - "table_format": st.sampled_from(("markdown", "html")), - "confidence_scores_granularity": st.sampled_from(("word", "page")), - "include_blocks": st.booleans(), - "id": _TEXT, - } +_VALUE_TEXT: Final = st.text(alphabet="abcdefghijklmnopqrstuvwxyz0123456789 -_", min_size=1, max_size=32) +_NULLABLE_TEXT: Final = st.one_of(st.none(), _VALUE_TEXT) +_POSITIVE_INTEGER: Final = st.integers(min_value=1, max_value=10_000) +_NON_NEGATIVE_INTEGER: Final = st.integers(min_value=0, max_value=10_000) +_SMALL_NUMBER: Final = st.floats(min_value=0.01, max_value=30.0, allow_nan=False, allow_infinity=False) +_REQUEST_ONLY_IMAGE: Final = ( + "data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAQAAAC1HAwCAAAAC0lEQVR42mNk+A8AAQUBAScY42YAAAAASUVORK5CYII=" +) +_BLOCK_TYPES: Final = ( + "Header", + "Footer", + "Title", + "Section Header", + "Page Number", + "List Item", + "Figure", + "Table", + "Key Value", + "Text", + "Comment", + "Signature", ) -@dataclass(frozen=True, slots=True) -class ProviderSpec: - name: str - model: str - upstream_base: str - api_key: str | None - upstream_model: str | None = None - vertex_project: str | None = None - vertex_location: str | None = None +def _optional_object(optional: dict[str, SearchStrategy[object]]) -> SearchStrategy[dict[str, object]]: + return st.fixed_dictionaries({}, optional=optional) -@dataclass(frozen=True, slots=True) -class RecorderResult: - fixture: OcrFixture - cache_hit: bool +def _merge_objects(left: dict[str, object], right: dict[str, object]) -> dict[str, object]: + return {**left, **right} -@dataclass(frozen=True, slots=True) -class GeneratorArgs: - providers: tuple[str, ...] - examples: int - fixture_dir: Path | None +def _page_range(start: int, length: int) -> dict[str, object]: + return {"start": start, "end": start + length} -class _RecordingProvider(ThreadingHTTPServer): - daemon_threads = True - - def __init__( - self, - spec: ProviderSpec, - sdk_kwargs: dict[str, object], - cache: JsonFileCache, - ) -> None: - super().__init__(("127.0.0.1", 0), _RecordingHandler) - self.spec: Final = spec - self.sdk_kwargs: Final = sdk_kwargs - self.cache: Final = cache - self.results: queue.Queue[RecorderResult] = queue.Queue() - - @property - def url(self) -> str: - return f"http://127.0.0.1:{self.server_address[1]}" - - def take_result(self) -> RecorderResult: - try: - return self.results.get(timeout=5) - except queue.Empty as error: - raise RuntimeError("successful OCR call did not produce a recorder result") from error +def _annotation_format(name: str, strict: bool) -> dict[str, object]: + return { + "type": "json_schema", + "json_schema": { + "name": name, + "schema": {"type": "object", "properties": {}, "additionalProperties": False}, + "strict": strict, + }, + } -class _RecordingHandler(BaseHTTPRequestHandler): - protocol_version = "HTTP/1.1" +def _annotation_format_only(value: dict[str, object]) -> dict[str, object]: + return {"document_annotation_format": value} - def do_POST(self) -> None: - provider: Final = self.server - assert isinstance(provider, _RecordingProvider) - length: Final = int(self.headers.get("content-length") or "0") - body: Final = JSON_OBJECT.validate_json(self.rfile.read(length)) - fixture_request: Final = OcrFixtureRequest( - provider=provider.spec.name, - sdk_kwargs=provider.sdk_kwargs, - provider_request=ProviderWireRequest(method=self.command, path=self.path, body=body), - ) - cache_key: Final = _fixture_cache_key(provider.spec.name, fixture_request.provider_request) - cached_value: Final = provider.cache.get(cache_key) - if cached_value is not None: - cached_fixture: Final = OcrFixture.model_validate(cached_value) - provider.results.put(RecorderResult(fixture=cached_fixture, cache_hit=True)) - self._send_fixture_response(cached_fixture.response) - return - upstream_url: Final = f"{provider.spec.upstream_base.rstrip('/')}{self.path}" - forwarded_headers: Final = { - name: value - for name, value in self.headers.items() - if name.lower() not in {"host", "content-length", "accept-encoding", "x-parity-case"} +def _null_annotation_format() -> dict[str, object]: + return {"document_annotation_format": None} + + +def _annotation_format_and_prompt(value: dict[str, object], prompt: str | None) -> dict[str, object]: + return {"document_annotation_format": value, "document_annotation_prompt": prompt} + + +def _table_agentic(prompt: str | None, mode: str) -> dict[str, object]: + return {"scope": "table", "prompt": prompt, "mode": mode} + + +def _figure_agentic(prompt: str | None, advanced: bool, overlays: bool) -> dict[str, object]: + return { + "scope": "figure", + "prompt": prompt, + "advanced_chart_agent": advanced, + "return_overlays": overlays, + } + + +def _text_agentic(prompt: str | None) -> dict[str, object]: + return {"scope": "text", "prompt": prompt} + + +def _agentic_scope(value: dict[str, object]) -> object: + return value["scope"] + + +_PAGE_RANGE: Final = st.builds( + _page_range, st.integers(min_value=1, max_value=1000), st.integers(min_value=0, max_value=50) +) +_PAGE_RANGE_VALUE: Final = st.one_of( + _PAGE_RANGE, + st.lists(_PAGE_RANGE, min_size=1, max_size=3), + st.lists(st.integers(min_value=1, max_value=1000), min_size=1, max_size=5, unique=True), + st.lists(_VALUE_TEXT, min_size=1, max_size=3, unique=True), +) +_ANNOTATION_FORMAT: Final[SearchStrategy[dict[str, object]]] = st.builds(_annotation_format, _VALUE_TEXT, st.booleans()) +_DOCUMENT_ANNOTATION: Final[SearchStrategy[dict[str, object]]] = st.one_of( + st.just(dict[str, object]()), + st.just(_null_annotation_format()), + st.builds(_annotation_format_only, _ANNOTATION_FORMAT), + st.builds(_annotation_format_and_prompt, _ANNOTATION_FORMAT, _NULLABLE_TEXT), +) + + +def _mistral_options(model: str) -> SearchStrategy[dict[str, object]]: + model_name: Final = model.rsplit("/", 1)[-1] + confidence_values: Final = ("word", "page") if model_name == "mistral-ocr-4-0" else ("word", "page", "block") + confidence_option: Final[dict[str, SearchStrategy[object]]] = ( + {} + if model_name == "mistral-ocr-2512" + else {"confidence_scores_granularity": st.one_of(st.none(), st.sampled_from(confidence_values))} + ) + independent: Final = _optional_object( + { + "pages": st.one_of( + st.none(), + st.lists(_NON_NEGATIVE_INTEGER, min_size=1, max_size=5, unique=True), + st.sampled_from(("0", "0,1,2", "0-5", "0,2-4")), + ), + "include_image_base64": st.one_of(st.none(), st.booleans()), + "image_limit": st.one_of(st.none(), _POSITIVE_INTEGER), + "image_min_size": st.one_of(st.none(), _NON_NEGATIVE_INTEGER), + "bbox_annotation_format": st.one_of(st.none(), _ANNOTATION_FORMAT), + "extract_header": st.booleans(), + "extract_footer": st.booleans(), + "table_format": st.one_of(st.none(), st.sampled_from(("markdown", "html"))), + "include_blocks": st.booleans(), + **confidence_option, } - upstream_body: Final = ( - {**body, "model": provider.spec.upstream_model} if provider.spec.upstream_model is not None else body - ) - try: - upstream_response: Final = httpx.post( - upstream_url, - headers=forwarded_headers, - content=json.dumps(upstream_body, separators=(",", ":")), - timeout=120, - ) - except httpx.HTTPError as error: - error_body: Final = json.dumps({"error": str(error)}).encode() - self._send_response(502, {"content-type": "application/json"}, error_body) - return - - raw_content_type: Final = cast(object, upstream_response.headers.get("content-type", "application/json")) - content_type: Final = raw_content_type if isinstance(raw_content_type, str) else "application/json" - response_headers: Final = {"content-type": content_type.split(";", 1)[0]} - if not upstream_response.is_success: - self._send_response(upstream_response.status_code, response_headers, upstream_response.content) - return - - upstream_response_body: Final = JSON_OBJECT.validate_json(upstream_response.content) - fixture_response: Final = OcrFixtureResponse( - status_code=upstream_response.status_code, - headers=response_headers, - body=upstream_response_body, - ) - recorded_fixture: Final = OcrFixture(request=fixture_request, response=fixture_response) - provider.results.put(RecorderResult(fixture=recorded_fixture, cache_hit=False)) - self._send_fixture_response(fixture_response) - - def _send_fixture_response(self, response: OcrFixtureResponse) -> None: - response_body: Final = json.dumps(response.body, separators=(",", ":")).encode() - self._send_response(response.status_code, response.headers, response_body) - - def _send_response(self, status_code: int, headers: dict[str, str], body: bytes) -> None: - self.send_response(status_code) - for name, value in headers.items(): - self.send_header(name, value) - self.send_header("content-length", str(len(body))) - self.end_headers() - self.wfile.write(body) - - def log_message(self, format: str, *args: object) -> None: - return + ) + return st.builds(_merge_objects, independent, _DOCUMENT_ANNOTATION) -@contextmanager -def _recording_provider( - spec: ProviderSpec, - sdk_kwargs: dict[str, object], - cache: JsonFileCache, -) -> Generator[_RecordingProvider]: - server: Final = _RecordingProvider(spec=spec, sdk_kwargs=sdk_kwargs, cache=cache) - thread: Final = threading.Thread(target=server.serve_forever, daemon=True) - thread.start() - try: - yield server - finally: - server.shutdown() - server.server_close() - thread.join(timeout=5) +_AGENTIC_TABLE: Final[SearchStrategy[dict[str, object]]] = st.builds( + _table_agentic, + _NULLABLE_TEXT, + st.sampled_from(("default", "auto", "max")), +) +_AGENTIC_FIGURE: Final[SearchStrategy[dict[str, object]]] = st.builds( + _figure_agentic, + _NULLABLE_TEXT, + st.booleans(), + st.booleans(), +) +_AGENTIC_TEXT: Final[SearchStrategy[dict[str, object]]] = st.builds(_text_agentic, _NULLABLE_TEXT) +_REDUCTO_ENHANCE: Final = _optional_object( + { + "agentic": st.lists( + st.one_of(_AGENTIC_TABLE, _AGENTIC_FIGURE, _AGENTIC_TEXT), + max_size=3, + unique_by=_agentic_scope, + ), + "summarize_figures": st.booleans(), + "intelligent_ordering": st.booleans(), + } +) +_CHUNKING: Final = _optional_object( + { + "chunk_mode": st.sampled_from(("variable", "section", "page", "disabled", "block", "page_sections")), + "chunk_size": st.one_of(st.none(), _POSITIVE_INTEGER), + "chunk_overlap": _NON_NEGATIVE_INTEGER, + } +) +_LEGACY_CHUNKING: Final = _optional_object( + { + "chunk_mode": st.sampled_from(("variable", "section", "page", "disabled", "block", "page_sections")), + "chunk_size": _POSITIVE_INTEGER, + "chunk_overlap": _NON_NEGATIVE_INTEGER, + } +) +_REDUCTO_RETRIEVAL: Final = _optional_object( + { + "chunking": _CHUNKING, + "filter_blocks": st.lists(st.sampled_from(_BLOCK_TYPES), max_size=len(_BLOCK_TYPES), unique=True), + "embedding_optimized": st.booleans(), + } +) +_REDUCTO_FORMATTING: Final = _optional_object( + { + "add_page_markers": st.booleans(), + "table_output_format": st.sampled_from(("html", "json", "md", "jsonbbox", "dynamic", "csv")), + "merge_tables": st.booleans(), + "include": st.lists( + st.sampled_from( + ("change_tracking", "highlight", "comments", "hyperlinks", "signatures", "ignore_watermarks") + ), + max_size=6, + unique=True, + ), + } +) +_SPLIT_TABLE_SIZE: Final = st.one_of( + _POSITIVE_INTEGER, + _optional_object( + {"row": st.one_of(st.none(), _POSITIVE_INTEGER), "column": st.one_of(st.none(), _POSITIVE_INTEGER)} + ), +) +_REDUCTO_SPREADSHEET: Final = _optional_object( + { + "split_large_tables": _optional_object({"enabled": st.booleans(), "size": _SPLIT_TABLE_SIZE}), + "include": st.lists(st.sampled_from(("cell_colors", "formula", "dropdowns")), max_size=3, unique=True), + "clustering": st.sampled_from(("accurate", "fast", "disabled")), + "exclude": st.lists( + st.sampled_from(("hidden_sheets", "hidden_rows", "hidden_cols", "styling", "spreadsheet_images")), + max_size=5, + unique=True, + ), + "max_cell_count": st.one_of(st.none(), _POSITIVE_INTEGER), + } +) +_TENANT_THROTTLING: Final = st.fixed_dictionaries( + {"tenant_id": st.text(alphabet="abcdefghijklmnopqrstuvwxyz0123456789-_", min_size=1, max_size=256)}, + optional={"max_share": st.floats(min_value=0.01, max_value=1.0, allow_nan=False, allow_infinity=False)}, +) +_REDUCTO_SETTINGS: Final = _optional_object( + { + "ocr_system": st.sampled_from(("standard", "legacy")), + "extraction_mode": st.sampled_from(("ocr", "hybrid")), + "force_url_result": st.booleans(), + "force_file_extension": _NULLABLE_TEXT, + "return_ocr_data": st.booleans(), + "return_images": st.lists(st.sampled_from(("figure", "table", "page")), max_size=3, unique=True), + "embed_pdf_metadata": st.booleans(), + "embed_pdf_metadata_dpi": st.integers(min_value=50, max_value=250), + "persist_results": st.booleans(), + "tenant_throttling": st.one_of(st.none(), _TENANT_THROTTLING), + "timeout": st.one_of(st.none(), _SMALL_NUMBER), + "page_range": st.one_of(st.none(), _PAGE_RANGE_VALUE), + "document_password": _NULLABLE_TEXT, + "hybrid_vpc": _optional_object({"environment": _NULLABLE_TEXT}), + } +) +_REDUCTO_V3_OPTIONS: Final = _optional_object( + { + "enhance": _REDUCTO_ENHANCE, + "retrieval": _REDUCTO_RETRIEVAL, + "formatting": _REDUCTO_FORMATTING, + "spreadsheet": _REDUCTO_SPREADSHEET, + "settings": _REDUCTO_SETTINGS, + } +) +_LEGACY_SUMMARY: Final = _optional_object( + {"enabled": st.booleans(), "prompt": _VALUE_TEXT, "override": st.booleans(), "advanced_chart_agent": st.booleans()} +) +_REDUCTO_LEGACY_OPTIONS: Final = _optional_object( + { + "ocr_mode": st.sampled_from(("standard", "agentic")), + "extraction_mode": st.sampled_from(("ocr", "metadata", "hybrid")), + "chunking": _LEGACY_CHUNKING, + "table_summary": _optional_object({"enabled": st.booleans(), "prompt": _VALUE_TEXT}), + "figure_summary": _LEGACY_SUMMARY, + "filter_blocks": st.lists(st.sampled_from(_BLOCK_TYPES), max_size=len(_BLOCK_TYPES), unique=True), + "force_url_result": st.booleans(), + } +) +_REDUCTO_LEGACY_ADVANCED: Final = _optional_object( + { + "ocr_system": st.sampled_from(("highres", "multilingual", "combined", "reducto", "legacy")), + "table_output_format": st.sampled_from(("html", "json", "md", "jsonbbox", "dynamic", "ai_json", "csv")), + "merge_tables": st.booleans(), + "include_formula_information": st.booleans(), + "include_color_information": st.booleans(), + "include_dropdown_information": st.booleans(), + "continue_hierarchy": st.booleans(), + "keep_line_breaks": st.booleans(), + "page_range": _PAGE_RANGE_VALUE, + "force_file_extension": _VALUE_TEXT, + "large_table_chunking": _optional_object({"enabled": st.booleans(), "size": _POSITIVE_INTEGER}), + "spreadsheet_table_clustering": st.sampled_from(("default", "disabled", "intelligent")), + "max_cell_count": st.one_of(st.none(), _POSITIVE_INTEGER), + "add_page_markers": st.booleans(), + "remove_text_formatting": st.booleans(), + "return_ocr_data": st.booleans(), + "document_password": _VALUE_TEXT, + "filter_line_numbers": st.booleans(), + "read_comments": st.booleans(), + "persist_results": st.booleans(), + "exclude_hidden_sheets": st.booleans(), + "exclude_hidden_rows_cols": st.booleans(), + "enable_change_tracking": st.booleans(), + "enable_highlight_detection": st.booleans(), + "ignore_watermarks": st.booleans(), + } +) +_REDUCTO_LEGACY_EXPERIMENTAL: Final = _optional_object( + { + "enrich": _optional_object( + { + "enabled": st.booleans(), + "mode": st.sampled_from(("standard", "page", "table", "table_auto")), + "prompt": _VALUE_TEXT, + } + ), + "layout_enrichment": st.booleans(), + "enable_checkboxes": st.booleans(), + "enable_equations": st.booleans(), + "rotate_pages": st.booleans(), + "rotate_figures": st.booleans(), + "enable_scripts": st.booleans(), + "return_figure_images": st.booleans(), + "return_table_images": st.booleans(), + "return_page_images": st.booleans(), + "layout_model": st.sampled_from(("default", "beta")), + "embed_text_metadata_pdf": st.booleans(), + "embed_pdf_metadata_dpi": st.integers(min_value=50, max_value=250), + "detect_signatures": st.booleans(), + "danger_filter_wide_boxes": st.booleans(), + "user_specified_timeout_seconds": st.one_of(st.none(), _SMALL_NUMBER), + } +) +_REDUCTO_LEGACY_ROOT: Final = _optional_object( + { + "options": _REDUCTO_LEGACY_OPTIONS, + "advanced_options": _REDUCTO_LEGACY_ADVANCED, + "experimental_options": _REDUCTO_LEGACY_EXPERIMENTAL, + "priority": st.booleans(), + } +) def _mistral_upstream_base() -> str: @@ -198,56 +354,60 @@ def _mistral_upstream_base() -> str: return configured.removesuffix("/v1") -def _fixture_cache_key(provider: str, request: ProviderWireRequest) -> dict[str, object]: - return {"provider": provider, "request": request.model_dump(mode="json")} - - -def _provider_specs(selected: tuple[str, ...]) -> tuple[ProviderSpec, ...]: - specs: Final[dict[str, ProviderSpec | None]] = { - "mistral": ( +def _provider_specs( + selected: tuple[str, ...], + requests_only: bool = False, + all_models: bool = False, +) -> tuple[ProviderSpec, ...]: + mistral_key: Final = os.environ.get("MISTRAL_API_KEY") or os.environ.get("LITELLM_API_KEY") + reducto_key: Final = os.environ.get("REDUCTO_API_KEY") + mistral_model: Final = os.environ.get("MISTRAL_OCR_MODEL") + reducto_model: Final = os.environ.get("REDUCTO_OCR_MODEL") + mistral_models: Final = ( + (mistral_model,) + if mistral_model is not None + else ( + ("mistral/mistral-ocr-2512", "mistral/mistral-ocr-4-0", "mistral/mistral-ocr-4-1") + if requests_only or all_models + else ("mistral/mistral-ocr-latest",) + ) + ) + reducto_models: Final = ( + (reducto_model,) if reducto_model is not None else ("reducto/parse-v3", "reducto/parse-legacy") + ) + mistral_specs: Final = ( + tuple( ProviderSpec( name="mistral", - model=os.environ.get("MISTRAL_OCR_MODEL", "mistral/mistral-ocr-latest"), + model=model, upstream_base=_mistral_upstream_base(), - api_key=os.environ.get("MISTRAL_API_KEY") or os.environ.get("LITELLM_API_KEY"), + api_key=mistral_key or "request-only-key", upstream_model=os.environ.get("MISTRAL_OCR_UPSTREAM_MODEL"), ) - if os.environ.get("MISTRAL_API_KEY") or os.environ.get("LITELLM_API_KEY") - else None - ), - "azure_ai": ( + for model in mistral_models + ) + if "mistral" in selected and (mistral_key is not None or requests_only) + else () + ) + reducto_specs: Final = ( + tuple( ProviderSpec( - name="azure_ai", - model=os.environ.get("AZURE_AI_OCR_MODEL", "azure_ai/mistral-document-ai-2512"), - upstream_base=os.environ["AZURE_AI_API_BASE"], - api_key=os.environ["AZURE_AI_API_KEY"], + name="reducto", + model=model, + upstream_base=os.environ.get("REDUCTO_API_BASE", REDUCTO_API_BASE), + api_key=reducto_key or "request-only-key", ) - if os.environ.get("AZURE_AI_API_BASE") and os.environ.get("AZURE_AI_API_KEY") - else None - ), - "vertex_ai": _vertex_spec(), - } - missing: Final = tuple(name for name in selected if specs[name] is None) + for model in reducto_models + ) + if "reducto" in selected and (reducto_key is not None or requests_only) + else () + ) + specs: Final = (*mistral_specs, *reducto_specs) + present: Final = frozenset(spec.name for spec in specs) + missing: Final = tuple(name for name in selected if name not in present) if missing: LOGGER.warning("Skipping providers without credentials: %s", ", ".join(missing)) - return tuple(spec for name in selected if (spec := specs[name]) is not None) - - -def _vertex_spec() -> ProviderSpec | None: - project: Final = os.environ.get("VERTEXAI_PROJECT") - credentials_available: Final = bool(os.environ.get("VERTEX_AI_API_KEY") or os.environ.get("VERTEXAI_CREDENTIALS")) - if project is None or not credentials_available: - return None - location: Final = os.environ.get("VERTEXAI_LOCATION", os.environ.get("VERTEX_LOCATION", "us-central1")) - upstream_base: Final = os.environ.get("VERTEX_AI_API_BASE", f"https://{location}-aiplatform.googleapis.com") - return ProviderSpec( - name="vertex_ai", - model=os.environ.get("VERTEX_AI_OCR_MODEL", "vertex_ai/mistral-ocr-2505"), - upstream_base=upstream_base, - api_key=os.environ.get("VERTEX_AI_API_KEY"), - vertex_project=project, - vertex_location=location, - ) + return specs def _image_data_uri(text: str, font_size: int) -> str: @@ -266,43 +426,36 @@ def _sdk_kwargs( image_data_uri: str, options: dict[str, object], ) -> dict[str, object]: - provider_kwargs: Final = ( - {"vertex_project": spec.vertex_project, "vertex_location": spec.vertex_location} - if spec.name == "vertex_ai" - else {} - ) return { "model": spec.model, "document": {"type": "image_url", "image_url": image_data_uri}, **options, - **provider_kwargs, } -def _record_case(spec: ProviderSpec, root: Path, sdk_kwargs: dict[str, object]) -> RecorderResult: - cache: Final = JsonFileCache(root / spec.name) - with _recording_provider(spec=spec, sdk_kwargs=sdk_kwargs, cache=cache) as recorder: - ocr_call: Final = cast(Callable[..., object], litellm.ocr) - ocr_call(api_base=recorder.url, api_key=spec.api_key, **sdk_kwargs) - result: Final = recorder.take_result() - - if not result.cache_hit: - cache.put( - _fixture_cache_key(spec.name, result.fixture.request.provider_request), - result.fixture.model_dump(mode="json"), - ) - return result - - -def _generate(specs: tuple[ProviderSpec, ...], root: Path, examples: int) -> None: - @settings(max_examples=examples, deadline=None, derandomize=True) - @given(text=_TEXT, font_size=st.integers(min_value=12, max_value=36), options=_OPTIONS) - def generate_case(text: str, font_size: int, options: dict[str, object]) -> None: - image_data_uri: Final = _image_data_uri(text, font_size) - for spec in specs: - _generate_provider_case(spec, root, image_data_uri, options) - - generate_case() +def _upload_reducto_document( + spec: ProviderSpec, + sdk_kwargs: dict[str, object], + requests_only: bool, +) -> dict[str, object]: + if spec.name != "reducto": + return sdk_kwargs + if requests_only: + return {**sdk_kwargs, "document": {"type": "image_url", "image_url": "reducto://fixture"}} + if spec.api_key is None: + raise ValueError("Reducto response fixture generation requires REDUCTO_API_KEY") + document: Final = JSON_OBJECT.validate_python(sdk_kwargs["document"]) + image_data_uri: Final = document.get("image_url") + if not isinstance(image_data_uri, str): + raise ValueError("Reducto fixture generation requires an image_url data URI") + _, raw_bytes, mime = extract_file_id_or_bytes(image_data_uri, model=spec.model) + file_id: Final = upload_bytes_sync( + raw_bytes=raw_bytes or b"", + mime=mime, + api_key=spec.api_key, + api_base=spec.upstream_base, + ) + return {**sdk_kwargs, "document": {"type": "image_url", "image_url": file_id}} def _generate_provider_case( @@ -310,38 +463,85 @@ def _generate_provider_case( root: Path, image_data_uri: str, options: dict[str, object], + requests_only: bool, + sdk_call: Callable[..., object], ) -> None: - sdk_kwargs: Final = _sdk_kwargs(spec, image_data_uri, options) - result: Final = _record_case(spec, root, sdk_kwargs) - state: Final = "cached" if result.cache_hit else "recorded" - LOGGER.info("%s %s %s", state, spec.name, result.fixture.request.provider_request.path) + sdk_kwargs: Final = _upload_reducto_document(spec, _sdk_kwargs(spec, image_data_uri, options), requests_only) + result: Final = record_case(spec, root, sdk_kwargs, requests_only, sdk_call) + state: Final = "cached" if result.cache_hit else "recorded request" if requests_only else "recorded" + LOGGER.info("%s %s %s", state, spec.name, result.request.provider_request.path) -def _parse_args() -> GeneratorArgs: - parser: Final = argparse.ArgumentParser() - parser.add_argument("--provider", action="append", choices=PROVIDER_NAMES) - parser.add_argument("--examples", type=int, default=4) - parser.add_argument("--fixture-dir", type=Path) - namespace: Final = parser.parse_args() - providers: Final = cast(list[str] | None, namespace.provider) - return GeneratorArgs( - providers=tuple(providers) if providers else PROVIDER_NAMES, - examples=cast(int, namespace.examples), - fixture_dir=cast(Path | None, namespace.fixture_dir), +def _generate_provider_examples( + spec: ProviderSpec, + root: Path, + examples: int, + requests_only: bool, + sdk_call: Callable[..., object], +) -> None: + options_strategy: Final = _options_strategy(spec) + image_strategy: Final = ( + st.just(_REQUEST_ONLY_IMAGE) + if requests_only + else st.builds(_image_data_uri, _TEXT, st.integers(min_value=12, max_value=36)) ) + @settings(max_examples=examples, deadline=None, derandomize=True) + @given(image_data_uri=image_strategy, options=options_strategy) + def generate_case(image_data_uri: str, options: dict[str, object]) -> None: + _generate_provider_case(spec, root, image_data_uri, options, requests_only, sdk_call) + + generate_case() + + +def _options_strategy(spec: ProviderSpec) -> SearchStrategy[dict[str, object]]: + model: Final = spec.model.rsplit("/", 1)[-1] + if spec.name == "mistral": + return _mistral_options(spec.model) + if model == "parse-v3": + return _REDUCTO_V3_OPTIONS + if model == "parse-legacy": + return _REDUCTO_LEGACY_ROOT + raise ValueError(f"Unsupported Reducto OCR fixture model: {spec.model}") + + +def _generate( + specs: tuple[ProviderSpec, ...], + root: Path, + examples: int, + requests_only: bool, + sdk_call: Callable[..., object], +) -> None: + for spec in specs: + _generate_provider_examples(spec, root, examples, requests_only, sdk_call) + + +def _log_filled_responses(results: tuple[RecorderResult, ...]) -> None: + for result in results: + LOGGER.info("filled response %s %s", result.request.provider, result.request.provider_request.path) + def main() -> None: logging.basicConfig(level=logging.INFO, format="%(message)s") load_dotenv() - args: Final = _parse_args() - root: Final = ( - args.fixture_dir or Path(os.environ.get(FIXTURE_DIR_ENV, Path(__file__).with_name(".fixtures"))).expanduser() + args: Final = parse_generator_args(PROVIDER_NAMES) + root: Final = fixture_directory( + args.fixture_dir, + os.environ.get(FIXTURE_DIR_ENV), + Path(__file__).with_name(".fixtures"), + ) + specs: Final = _provider_specs( + args.providers, + requests_only=args.requests_only, + all_models=args.responses_only, ) - specs: Final = _provider_specs(args.providers) if not specs: raise SystemExit("No selected provider has the required credentials") - _generate(specs, root, args.examples) + ocr_call: Final = cast(Callable[..., object], litellm.ocr) + if args.responses_only: + _log_filled_responses(fill_missing_responses(specs, root, ocr_call)) + return + _generate(specs, root, args.examples, args.requests_only, ocr_call) if __name__ == "__main__": diff --git a/tests/test_litellm/ocr/test_rust_bridge.py b/tests/test_litellm/ocr/test_rust_bridge.py index acad249a2bb..5cd7fee0df7 100644 --- a/tests/test_litellm/ocr/test_rust_bridge.py +++ b/tests/test_litellm/ocr/test_rust_bridge.py @@ -1,7 +1,7 @@ """Tests for the optional Rust-backed OCR path.""" -import importlib import builtins +import importlib import types from typing import Any @@ -471,9 +471,7 @@ def test_run_rust_ocr_resolves_key_via_secret_manager_when_missing(): ocr_main._run_rust_ocr( prepared_request=build_prepared_request(api_key=None, timeout=None), - resolve_api_key=lambda name: ( - "sk-from-vault" if name == "MISTRAL_API_KEY" else None - ), + resolve_api_key=lambda name: "sk-from-vault" if name == "MISTRAL_API_KEY" else None, ) assert bridge.calls[0]["api_key"] == "sk-from-vault" @@ -580,9 +578,7 @@ def test_prepare_rust_ocr_call_resolves_azure_ai_api_base_from_secret_manager(): api_base=None, timeout=None, ), - resolve_api_key=lambda name: ( - "https://azure.example.com" if name == "AZURE_AI_API_BASE" else None - ), + resolve_api_key=lambda name: "https://azure.example.com" if name == "AZURE_AI_API_BASE" else None, ) assert bridge.calls[0]["api_base"] == "https://azure.example.com" @@ -600,9 +596,7 @@ def test_prepare_rust_ocr_call_resolves_document_intelligence_endpoint(): timeout=None, ), resolve_api_key=lambda name: ( - "https://document-intelligence.example.com" - if name == "AZURE_DOCUMENT_INTELLIGENCE_ENDPOINT" - else None + "https://document-intelligence.example.com" if name == "AZURE_DOCUMENT_INTELLIGENCE_ENDPOINT" else None ), ) @@ -662,18 +656,67 @@ def test_ocr_routes_to_rust_when_enabled(fake_bridge): assert call["optional_params"].get("include_image_base64") is True -def test_ocr_routes_azure_ai_to_rust_when_enabled(fake_bridge): +def test_ocr_routes_reducto_parse_v3_file_id_to_rust(fake_bridge: RecordingBridge) -> None: + document: dict[str, object] = { + "type": "document_url", + "document_url": "reducto://fixture", + } + response = litellm.ocr( - model="azure_ai/pixtral-12b-2409", - document=DOCUMENT, + model="reducto/parse-v3", + document=document, api_key="sk-test", - api_base="https://example.services.ai.azure.com", ) assert isinstance(response, OCRResponse) assert len(fake_bridge.calls) == 1 - assert fake_bridge.calls[0]["model"] == "pixtral-12b-2409" - assert fake_bridge.calls[0]["custom_llm_provider"] == "azure_ai" + call = fake_bridge.calls[0] + assert call["model"] == "parse-v3" + assert call["document"] == document + assert call["custom_llm_provider"] == "reducto" + + +@pytest.mark.parametrize( + ("model", "document"), + ( + ("azure_ai/pixtral-12b-2409", DOCUMENT), + ("vertex_ai/mistral-ocr-2505", DOCUMENT), + ( + "reducto/parse-v3", + { + "type": "document_url", + "document_url": "data:application/pdf;base64,AA==", + }, + ), + ( + "reducto/parse-legacy", + {"type": "document_url", "document_url": "reducto://fixture"}, + ), + ), +) +def test_ocr_falls_back_to_python_for_unsupported_rust_case( + monkeypatch: pytest.MonkeyPatch, + fake_bridge: RecordingBridge, + model: str, + document: dict[str, object], +) -> None: + captured: dict[str, object] = {} + + def fake_handler_ocr(**kwargs: object) -> OCRResponse: + captured.update(kwargs) + return OCRResponse(pages=[], model=model, object="ocr") + + monkeypatch.setattr(ocr_main.base_llm_http_handler, "ocr", fake_handler_ocr) + response = litellm.ocr( + model=model, + document=document, + api_key="sk-test", + api_base="https://example.com", + ) + + assert isinstance(response, OCRResponse) + assert fake_bridge.calls == [] + assert captured["model"] in {"pixtral-12b-2409", "mistral-ocr-2505"} def test_ocr_rust_path_converts_file_document_before_bridge(fake_bridge): @@ -815,9 +858,6 @@ def test_ocr_provider_configs_expose_api_key_env_vars(): assert BaseOCRConfig().get_api_key_env_var() is None assert MistralOCRConfig().get_api_key_env_var() == "MISTRAL_API_KEY" assert AzureAIOCRConfig().get_api_key_env_var() == "AZURE_AI_API_KEY" - assert ( - AzureDocumentIntelligenceOCRConfig().get_api_key_env_var() - == "AZURE_DOCUMENT_INTELLIGENCE_API_KEY" - ) + assert AzureDocumentIntelligenceOCRConfig().get_api_key_env_var() == "AZURE_DOCUMENT_INTELLIGENCE_API_KEY" assert VertexAIOCRConfig().get_api_key_env_var() == "VERTEX_AI_API_KEY" assert VertexAIDeepSeekOCRConfig().get_api_key_env_var() == "VERTEX_AI_API_KEY" diff --git a/tests/test_litellm/ocr/test_sdk_parity.py b/tests/test_litellm/ocr/test_sdk_parity.py index d587ce99281..a8b154c8de8 100644 --- a/tests/test_litellm/ocr/test_sdk_parity.py +++ b/tests/test_litellm/ocr/test_sdk_parity.py @@ -10,7 +10,7 @@ from typing import Final, Protocol, cast import pytest from pydantic import BaseModel, ConfigDict, TypeAdapter -from tests.test_litellm.ocr.fixture_models import OcrFixture, OcrFixtureResponse, ProviderWireRequest +from tests.test_litellm._fixture_recorder import Fixture, FixtureResponse, ProviderWireRequest from tests.test_litellm.parity.compare import assert_parity from tests.test_litellm.parity.models import ( CapturedRequest, @@ -27,6 +27,7 @@ from tests.test_litellm.parity.runner import PythonScriptRunner, run_execution API_KEY: Final = "test-key" PYTHON_HTTP_SENTINEL: Final = "python-ocr-parity-fallback" GIL_STATS: Final = TypeAdapter(dict[str, int]) +SUPPORTED_BY_RUST_PROVIDERS: Final = frozenset({"mistral", "reducto"}) class SDKRoute(str, Enum): @@ -39,6 +40,7 @@ class SDKInput(BaseModel): route: SDKRoute kwargs: dict[str, object] + native_expected: bool class _NativeBridge(Protocol): @@ -124,6 +126,9 @@ def _execute_sdk_case(sdk_input: SDKInput, mock_url: str) -> SDKReport: raise RuntimeError("LITELLM_USE_RUST_OCR=1 but the native bridge is unavailable") native_bridge: Final = cast(_NativeBridge, raw_native_bridge) + if not sdk_input.native_expected: + return _native_sdk_report(_capture_sdk_call(sdk_input, mock_url), native_handled_case=False) + native_callable: Final = load_rust_ocr() if sdk_input.route is SDKRoute.OCR else load_rust_aocr() expected_native_callable: Final = native_bridge.ocr if sdk_input.route is SDKRoute.OCR else native_bridge.aocr if native_callable is None or native_callable is not expected_native_callable: @@ -139,7 +144,7 @@ def _execute_sdk_case(sdk_input: SDKInput, mock_url: str) -> SDKReport: return _native_sdk_report(trace, native_handled_case=after_gil_releases == before_gil_releases + 1) -def _replay_response(response: OcrFixtureResponse) -> ReplayResponse: +def _replay_response(response: FixtureResponse) -> ReplayResponse: return ReplayResponse(status_code=response.status_code, headers=response.headers, body=response.body) @@ -147,9 +152,22 @@ def _provider_wire_request(request: CapturedRequest) -> ProviderWireRequest: return ProviderWireRequest(method=request.method, path=request.path, body=request.body) +def _native_expected(ocr_fixture: Fixture) -> bool: + provider: Final = ocr_fixture.request.provider + if provider != "reducto": + return provider in SUPPORTED_BY_RUST_PROVIDERS + model: Final = ocr_fixture.request.sdk_kwargs.get("model") + document: Final = ocr_fixture.request.sdk_kwargs.get("document") + if model != "reducto/parse-v3" or not isinstance(document, dict): + return False + source: Final = document.get("document_url") or document.get("image_url") + return isinstance(source, str) and source.startswith("reducto://") + + @pytest.mark.parametrize("route", tuple(SDKRoute), ids=tuple(route.value for route in SDKRoute)) -def test_recorded_ocr_sdk_parity(ocr_fixture: OcrFixture, route: SDKRoute, tmp_path: Path) -> None: - sdk_input: Final = SDKInput(route=route, kwargs=ocr_fixture.request.sdk_kwargs) +def test_recorded_ocr_sdk_parity(ocr_fixture: Fixture, route: SDKRoute, tmp_path: Path) -> None: + native_expected: Final = _native_expected(ocr_fixture) + sdk_input: Final = SDKInput(route=route, kwargs=ocr_fixture.request.sdk_kwargs, native_expected=native_expected) case_file: Final = tmp_path / f"{ocr_fixture.request.provider}-{route.value}-sdk-input.json" case_file.write_text(sdk_input.model_dump_json(indent=2), encoding="utf-8") expected_request: Final = ocr_fixture.request.provider_request @@ -178,7 +196,7 @@ def test_recorded_ocr_sdk_parity(ocr_fixture: OcrFixture, route: SDKRoute, tmp_p rust_enabled=True, ) - assert_parity(python, rust, PYTHON_HTTP_SENTINEL) + assert_parity(python, rust, PYTHON_HTTP_SENTINEL, native_expected) assert tuple(_provider_wire_request(request) for request in python.requests) == (expected_request,) assert tuple(_provider_wire_request(request) for request in rust.requests) == (expected_request,) diff --git a/tests/test_litellm/parity/compare.py b/tests/test_litellm/parity/compare.py index 323524c1164..4ffd4fef822 100644 --- a/tests/test_litellm/parity/compare.py +++ b/tests/test_litellm/parity/compare.py @@ -12,17 +12,27 @@ def _stable_request(request: CapturedRequest) -> dict[str, object]: return request.model_dump(mode="json", exclude={"user_agent"}) -def assert_parity(python: Execution, rust: Execution, python_user_agent: str) -> None: +def assert_parity( + python: Execution, + rust: Execution, + python_user_agent: str, + native_expected: bool, +) -> None: assert python.report.native.rust_enabled is False assert rust.report.native.rust_enabled is True - assert rust.report.native.native_callable_loaded is True - assert rust.report.native.native_handled_case is True assert tuple(request.user_agent for request in python.requests) == (python_user_agent,) * len(python.requests) - assert all(request.user_agent != python_user_agent for request in rust.requests) + if native_expected: + assert rust.report.native.native_callable_loaded is True + assert rust.report.native.native_handled_case is True + assert all(request.user_agent != python_user_agent for request in rust.requests) + else: + assert rust.report.native.native_handled_case is False + assert tuple(request.user_agent for request in rust.requests) == (python_user_agent,) * len(rust.requests) assert tuple(_stable_request(request) for request in rust.requests) == tuple( _stable_request(request) for request in python.requests ) - assert rust.report.trace == python.report.trace + trace_differences: Final = _diff(rust.report.trace, python.report.trace) + assert not trace_differences, "\n".join(trace_differences) assert python.report.trace.exception is None assert python.report.trace.outputs diff --git a/tests/test_litellm/parity/runner.py b/tests/test_litellm/parity/runner.py index a136d9901ee..b771408e707 100644 --- a/tests/test_litellm/parity/runner.py +++ b/tests/test_litellm/parity/runner.py @@ -34,10 +34,15 @@ def run_execution( provider: JsonReplayServer, rust_enabled: bool, ) -> Execution: + project_root: Final = str(runner.entrypoint.resolve().parents[3]) + existing_pythonpath: Final = os.environ.get("PYTHONPATH") env: Final = { **os.environ, runner.rust_env_var: "1" if rust_enabled else "0", "LITELLM_USER_AGENT": runner.python_user_agent, + "PYTHONPATH": os.pathsep.join( + path for path in (project_root, existing_pythonpath) if path + ), } completed: Final = subprocess.run( runner.command(case_file, provider.url, report_file), diff --git a/tests/test_litellm/test_fixture_recorder.py b/tests/test_litellm/test_fixture_recorder.py new file mode 100644 index 00000000000..70f12809b35 --- /dev/null +++ b/tests/test_litellm/test_fixture_recorder.py @@ -0,0 +1,106 @@ +from __future__ import annotations + +import json +from collections.abc import Callable +from pathlib import Path +from typing import Final, cast + +import httpx + +from tests.test_litellm._fixture_recorder import ( + ProviderSpec, + pending_requests, + record_case, +) +from tests.test_litellm._json_fs_cache import JsonFileCache + + +def _request_value(model: str = "test/model") -> dict[str, object]: + return { + "request": { + "provider": "test-provider", + "sdk_kwargs": {"model": model}, + "provider_request": {"method": "POST", "path": "/v1/test", "body": {"model": model}}, + } + } + + +def _sdk_call(**kwargs: object) -> None: + api_base: Final = kwargs["api_base"] + model: Final = kwargs["model"] + assert isinstance(api_base, str) + assert isinstance(model, str) + response: Final = httpx.post( + f"{api_base}/v1/test", + content=json.dumps({"model": model}).encode(), + headers={"content-type": "application/json"}, + ) + response.raise_for_status() + + +def _same_wire_sdk_call(**kwargs: object) -> None: + api_base: Final = kwargs["api_base"] + assert isinstance(api_base, str) + response: Final = httpx.post( + f"{api_base}/v1/test", + content=b'{"constant":true}', + headers={"content-type": "application/json"}, + ) + response.raise_for_status() + + +def test_request_only_recording_persists_pending_fixture_without_calling_upstream(tmp_path: Path) -> None: + spec: Final = ProviderSpec( + name="test-provider", + model="test/model", + upstream_base="http://127.0.0.1:1", + api_key="test-key", + ) + sdk_kwargs: Final[dict[str, object]] = {"model": spec.model} + + result: Final = record_case( + spec, + tmp_path, + sdk_kwargs, + requests_only=True, + sdk_call=cast(Callable[..., object], _sdk_call), + ) + values: Final = JsonFileCache(tmp_path / spec.name).values() + + assert result.response is None + assert len(values) == 1 + assert values[0] == {"request": result.request.model_dump(mode="json")} + + +def test_pending_requests_excludes_completed_fixtures(tmp_path: Path) -> None: + cache: Final = JsonFileCache(tmp_path) + pending: Final = _request_value() + response: Final[dict[str, object]] = {"status_code": 200, "headers": {}, "body": {}} + completed: Final[dict[str, object]] = { + **_request_value("test/completed"), + "response": response, + } + cache.put({"case": "pending"}, pending) + cache.put({"case": "completed"}, completed) + + requests: Final = pending_requests(cache) + + assert len(requests) == 1 + assert requests[0].sdk_kwargs["model"] == "test/model" + + +def test_request_cache_distinguishes_sdk_inputs_with_identical_wire_requests(tmp_path: Path) -> None: + spec: Final = ProviderSpec( + name="test-provider", + model="test/model-a", + upstream_base="http://127.0.0.1:1", + api_key="test-key", + ) + sdk_call: Final = cast(Callable[..., object], _same_wire_sdk_call) + + first: Final = record_case(spec, tmp_path, {"model": "test/model-a"}, True, sdk_call) + second: Final = record_case(spec, tmp_path, {"model": "test/model-b"}, True, sdk_call) + + assert not first.cache_hit + assert not second.cache_hit + assert len(JsonFileCache(tmp_path / spec.name).values()) == 2