diff --git a/tests/route_parity/fixture_generator.py b/tests/route_parity/fixture_generator.py index 0b4eae55c99..de51f09b8b7 100644 --- a/tests/route_parity/fixture_generator.py +++ b/tests/route_parity/fixture_generator.py @@ -2,10 +2,10 @@ from __future__ import annotations import argparse import logging -from collections.abc import Callable, Sequence +from collections.abc import Callable, Mapping, Sequence from dataclasses import dataclass from pathlib import Path -from typing import Final, Generic, TypeVar, cast +from typing import Final, Generic, Protocol, TypeVar, cast from hypothesis.strategies import SearchStrategy from pydantic import BaseModel @@ -34,6 +34,26 @@ class FixtureTarget(Generic[InputT]): required_inputs: tuple[InputT, ...] = () +class FixtureSdkCall(Protocol): + def __call__(self, **kwargs: object) -> object: ... + + +class FixtureProvider(Protocol[InputT]): + def targets( + self, + environ: Mapping[str, str], + sdk_call: FixtureSdkCall, + ) -> tuple[FixtureTarget[InputT], ...]: ... + + +def discover_fixture_targets( + providers: tuple[FixtureProvider[InputT], ...], + environ: Mapping[str, str], + sdk_call: FixtureSdkCall, +) -> tuple[FixtureTarget[InputT], ...]: + return tuple(target for provider in providers for target in provider.targets(environ, sdk_call)) + + def generate_target_fixtures( target: FixtureTarget[InputT], root: Path, diff --git a/tests/route_parity/test_fixture_generator.py b/tests/route_parity/test_fixture_generator.py new file mode 100644 index 00000000000..5b2030d8888 --- /dev/null +++ b/tests/route_parity/test_fixture_generator.py @@ -0,0 +1,72 @@ +from __future__ import annotations + +import queue +from collections.abc import Mapping +from dataclasses import dataclass +from typing import Final + +from hypothesis import strategies as st + +from tests.route_parity.fixture_generator import FixtureSdkCall, FixtureTarget, discover_fixture_targets +from tests.route_parity.fixture_models import SdkInputBase +from tests.route_parity.fixture_recorder import ProviderSpec + + +class ExampleSdkInput(SdkInputBase): + model: str + + +@dataclass(frozen=True, slots=True) +class ExampleProvider: + name: str + key_name: str + + def targets( + self, + environ: Mapping[str, str], + sdk_call: FixtureSdkCall, + ) -> tuple[FixtureTarget[ExampleSdkInput], ...]: + api_key: Final = environ.get(self.key_name) + if not api_key: + return () + + def invoke(api_base: str, case_input: ExampleSdkInput) -> object: + return sdk_call(api_base=api_base, api_key=api_key, **case_input.as_sdk_kwargs()) + + case_input: Final = ExampleSdkInput(model=f"{self.name}/model") + return ( + FixtureTarget( + name=self.name, + provider_spec=ProviderSpec(upstream_base=f"https://{self.name}.example"), + strategy=st.just(case_input), + invoke=invoke, + required_inputs=(case_input,), + ), + ) + + +def test_discover_fixture_targets_flattens_configured_providers_and_injects_sdk_call() -> None: + calls: Final[queue.SimpleQueue[dict[str, object]]] = queue.SimpleQueue() + + def sdk_call(**kwargs: object) -> object: + calls.put(kwargs) + return "response" + + providers: Final = ( + ExampleProvider(name="first", key_name="FIRST_KEY"), + ExampleProvider(name="skipped", key_name="SKIPPED_KEY"), + ExampleProvider(name="second", key_name="SECOND_KEY"), + ) + targets: Final = discover_fixture_targets( + providers, + {"FIRST_KEY": "first-secret", "SECOND_KEY": "second-secret"}, + sdk_call, + ) + + assert tuple(target.name for target in targets) == ("first", "second") + assert targets[1].invoke("http://127.0.0.1:1234", targets[1].required_inputs[0]) == "response" + assert calls.get_nowait() == { + "api_base": "http://127.0.0.1:1234", + "api_key": "second-secret", + "model": "second/model", + } diff --git a/tests/test_litellm/ocr/conftest.py b/tests/test_litellm/ocr/conftest.py index e724bdb4e31..6018c5e1dda 100644 --- a/tests/test_litellm/ocr/conftest.py +++ b/tests/test_litellm/ocr/conftest.py @@ -6,7 +6,7 @@ from typing import Final import pytest from tests.route_parity.fixture_recorder import fixture_id, parametrize_recorded_fixtures -from tests.test_litellm.ocr.fixture_models import OcrParityCase +from tests.test_litellm.ocr.fixtures.models import OcrParityCase FIXTURE_DIR_ENV: Final = "LITELLM_OCR_FIXTURE_DIR" @@ -19,7 +19,7 @@ def _fixture_id(fixture: OcrParityCase) -> str: def pytest_generate_tests(metafunc: pytest.Metafunc) -> None: - default_directory: Final = Path(__file__).with_name("fixtures") + default_directory: Final = Path(__file__).with_name("fixtures") / "data" parametrize_recorded_fixtures( metafunc, fixture_name="ocr_fixture", @@ -27,7 +27,7 @@ def pytest_generate_tests(metafunc: pytest.Metafunc) -> None: env_var=FIXTURE_DIR_ENV, default_directory=default_directory, regeneration_command=( - f"uv run python -m tests.test_litellm.ocr.generate_fixtures --fixture-dir {default_directory}" + f"uv run python -m tests.test_litellm.ocr.fixtures.generate --fixture-dir {default_directory}" ), id_builder=_fixture_id, ) diff --git a/tests/test_litellm/ocr/fixtures/__init__.py b/tests/test_litellm/ocr/fixtures/__init__.py new file mode 100644 index 00000000000..8b137891791 --- /dev/null +++ b/tests/test_litellm/ocr/fixtures/__init__.py @@ -0,0 +1 @@ + diff --git a/tests/test_litellm/ocr/fixtures/azure.py b/tests/test_litellm/ocr/fixtures/azure.py new file mode 100644 index 00000000000..c4b23e0d669 --- /dev/null +++ b/tests/test_litellm/ocr/fixtures/azure.py @@ -0,0 +1,115 @@ +from __future__ import annotations + +from collections.abc import Mapping +from typing import Final, cast + +from hypothesis import strategies as st +from hypothesis.strategies import DrawFn, SearchStrategy + +from tests.route_parity.fixture_generator import FixtureSdkCall +from tests.route_parity.fixture_recorder import ProviderSpec +from tests.test_litellm.ocr.fixtures.common import ( + OcrFixtureTarget, + invoke_with_api_key, + pdf_document, + public_document_strategy, +) +from tests.test_litellm.ocr.fixtures.mistral import ( + MISTRAL_MODEL, + mistral_input_strategy, + required_mistral_inputs, +) +from tests.test_litellm.ocr.fixtures.models import ( + AzureDocumentIntelligenceOcrSdkInput, + AzureMistralOcrSdkInput, + MistralOcrSdkInput, + OcrSdkInputBase, +) + + +def _as_azure_mistral(case_input: MistralOcrSdkInput, model: str) -> AzureMistralOcrSdkInput: + values: Final = case_input.model_dump(mode="python", exclude={"boundary", "model", "custom_llm_provider"}) + return AzureMistralOcrSdkInput.model_validate({**values, "model": model}) + + +def _required_document_intelligence_inputs() -> tuple[AzureDocumentIntelligenceOcrSdkInput, ...]: + document: Final = pdf_document() + model: Final = "azure_ai/doc-intelligence/prebuilt-layout" + return ( + AzureDocumentIntelligenceOcrSdkInput(model=model, document=document), + AzureDocumentIntelligenceOcrSdkInput(model=model, document=document, pages=[0, 1]), + AzureDocumentIntelligenceOcrSdkInput(model=model, document=document, features=["languages"]), + AzureDocumentIntelligenceOcrSdkInput(model=model, document=document, req_format="litellm"), + ) + + +@st.composite +def azure_document_intelligence_input_strategy(draw: DrawFn) -> AzureDocumentIntelligenceOcrSdkInput: + optional_params: Final = draw( + st.fixed_dictionaries( + {}, + optional={ + "pages": st.sampled_from(([0], [0, 1], "1-2")), + "features": st.sampled_from((["languages"], ["keyValuePairs"], "languages,keyValuePairs")), + }, + ) + ) + return AzureDocumentIntelligenceOcrSdkInput.model_validate( + { + "model": draw( + st.sampled_from( + ( + "azure_ai/doc-intelligence/prebuilt-read", + "azure_ai/doc-intelligence/prebuilt-layout", + "azure_ai/doc-intelligence/prebuilt-document", + ) + ) + ), + "document": draw(public_document_strategy()), + **optional_params, + } + ) + + +class AzureMistralFixtureProvider: + def targets(self, environ: Mapping[str, str], sdk_call: FixtureSdkCall) -> tuple[OcrFixtureTarget, ...]: + api_key: Final = environ.get("AZURE_AI_API_KEY") + upstream_base: Final = environ.get("AZURE_AI_API_BASE") + configured_model: Final = environ.get("AZURE_AI_OCR_MODEL") + if not api_key or not upstream_base or not configured_model: + return () + model: Final = configured_model if configured_model.startswith("azure_ai/") else f"azure_ai/{configured_model}" + return ( + OcrFixtureTarget( + name="azure-mistral", + provider_spec=ProviderSpec(upstream_base=upstream_base.rstrip("/")), + strategy=cast( + SearchStrategy[OcrSdkInputBase], + mistral_input_strategy(MISTRAL_MODEL).map(lambda case_input: _as_azure_mistral(case_input, model)), + ), + invoke=invoke_with_api_key(sdk_call, api_key), + required_inputs=cast( + tuple[OcrSdkInputBase, ...], + tuple( + _as_azure_mistral(case_input, model) for case_input in required_mistral_inputs(MISTRAL_MODEL) + ), + ), + ), + ) + + +class AzureDocumentIntelligenceFixtureProvider: + def targets(self, environ: Mapping[str, str], sdk_call: FixtureSdkCall) -> tuple[OcrFixtureTarget, ...]: + api_key: Final = environ.get("AZURE_DOCUMENT_INTELLIGENCE_API_KEY") + upstream_base: Final = environ.get("AZURE_DOCUMENT_INTELLIGENCE_ENDPOINT") + if not api_key or not upstream_base: + return () + return ( + OcrFixtureTarget( + name="azure-document-intelligence", + provider_spec=ProviderSpec(upstream_base=upstream_base.rstrip("/")), + strategy=cast(SearchStrategy[OcrSdkInputBase], azure_document_intelligence_input_strategy()), + invoke=invoke_with_api_key(sdk_call, api_key), + required_inputs=cast(tuple[OcrSdkInputBase, ...], _required_document_intelligence_inputs()), + ), + ) diff --git a/tests/test_litellm/ocr/fixtures/common.py b/tests/test_litellm/ocr/fixtures/common.py new file mode 100644 index 00000000000..2d21f15dc15 --- /dev/null +++ b/tests/test_litellm/ocr/fixtures/common.py @@ -0,0 +1,64 @@ +from __future__ import annotations + +import base64 +from collections.abc import Callable +from pathlib import Path +from typing import Final +from urllib.parse import quote + +from hypothesis import strategies as st +from hypothesis.strategies import SearchStrategy + +from tests.route_parity.fixture_generator import FixtureSdkCall, FixtureTarget +from tests.test_litellm.ocr.fixtures.models import ( + JsonSchemaDefinition, + JsonSchemaResponseFormat, + MistralDocumentUrlDocument, + MistralImageUrlDocument, + OcrSdkInputBase, +) + +OcrFixtureTarget = FixtureTarget[OcrSdkInputBase] + + +def image_document(text: str, font_size: int) -> MistralImageUrlDocument: + url: Final = f"https://dummyjson.com/image/800x300/ffffff/000000?text={quote(text)}&fontSize={font_size}" + return MistralImageUrlDocument(type="image_url", image_url=url) + + +def fixture_pdf_data_uri() -> str: + fixture: Final = Path(__file__).resolve().parents[3] / "llm_translation" / "fixtures" / "dummy.pdf" + encoded: Final = base64.b64encode(fixture.read_bytes()).decode("ascii") + return f"data:application/pdf;base64,{encoded}" + + +def pdf_document() -> MistralDocumentUrlDocument: + return MistralDocumentUrlDocument(type="document_url", document_url=fixture_pdf_data_uri()) + + +def public_document_strategy() -> SearchStrategy[MistralImageUrlDocument | MistralDocumentUrlDocument]: + return st.sampled_from((image_document("invoice 123", 24), pdf_document())) + + +def annotation_format(name: str) -> JsonSchemaResponseFormat: + return JsonSchemaResponseFormat( + type="json_schema", + json_schema=JsonSchemaDefinition( + name=name, + description="Extract the visible document fields", + schema={ + "type": "object", + "properties": {"title": {"type": "string"}}, + "required": ["title"], + "additionalProperties": False, + }, + strict=True, + ), + ) + + +def invoke_with_api_key(sdk_call: FixtureSdkCall, api_key: str) -> Callable[[str, OcrSdkInputBase], object]: + def invoke(api_base: str, case_input: OcrSdkInputBase) -> object: + return sdk_call(api_base=api_base, api_key=api_key, **case_input.as_sdk_kwargs()) + + return invoke diff --git a/tests/test_litellm/ocr/fixtures/6aa18aa0bb41abea634cac9c77fcb73438844c5eb7b4f5eae1a86ecce7b60742.json b/tests/test_litellm/ocr/fixtures/data/6aa18aa0bb41abea634cac9c77fcb73438844c5eb7b4f5eae1a86ecce7b60742.json similarity index 100% rename from tests/test_litellm/ocr/fixtures/6aa18aa0bb41abea634cac9c77fcb73438844c5eb7b4f5eae1a86ecce7b60742.json rename to tests/test_litellm/ocr/fixtures/data/6aa18aa0bb41abea634cac9c77fcb73438844c5eb7b4f5eae1a86ecce7b60742.json diff --git a/tests/test_litellm/ocr/fixtures/ecfb251cbe6cb13e52dd11a4706f788cf6f2b3d5f4fdae5d7b9b620c398f47dc.json b/tests/test_litellm/ocr/fixtures/data/ecfb251cbe6cb13e52dd11a4706f788cf6f2b3d5f4fdae5d7b9b620c398f47dc.json similarity index 100% rename from tests/test_litellm/ocr/fixtures/ecfb251cbe6cb13e52dd11a4706f788cf6f2b3d5f4fdae5d7b9b620c398f47dc.json rename to tests/test_litellm/ocr/fixtures/data/ecfb251cbe6cb13e52dd11a4706f788cf6f2b3d5f4fdae5d7b9b620c398f47dc.json diff --git a/tests/test_litellm/ocr/fixtures/mistral-ocr/0ab748e6b7315f0c7a36195cc4564eff8b03523dfa908ea59ff1bb7bd185ede8.json b/tests/test_litellm/ocr/fixtures/data/mistral-ocr/0ab748e6b7315f0c7a36195cc4564eff8b03523dfa908ea59ff1bb7bd185ede8.json similarity index 100% rename from tests/test_litellm/ocr/fixtures/mistral-ocr/0ab748e6b7315f0c7a36195cc4564eff8b03523dfa908ea59ff1bb7bd185ede8.json rename to tests/test_litellm/ocr/fixtures/data/mistral-ocr/0ab748e6b7315f0c7a36195cc4564eff8b03523dfa908ea59ff1bb7bd185ede8.json diff --git a/tests/test_litellm/ocr/fixtures/mistral-ocr/0c9204f5ed668c628835b9d60806a6cea47e1f67c0f48c9804cb09d1c4695f42.json b/tests/test_litellm/ocr/fixtures/data/mistral-ocr/0c9204f5ed668c628835b9d60806a6cea47e1f67c0f48c9804cb09d1c4695f42.json similarity index 100% rename from tests/test_litellm/ocr/fixtures/mistral-ocr/0c9204f5ed668c628835b9d60806a6cea47e1f67c0f48c9804cb09d1c4695f42.json rename to tests/test_litellm/ocr/fixtures/data/mistral-ocr/0c9204f5ed668c628835b9d60806a6cea47e1f67c0f48c9804cb09d1c4695f42.json diff --git a/tests/test_litellm/ocr/fixtures/mistral-ocr/216aacc1aa47a6060445f2ea49a2e0a8f4a8bde179faaaa28f39b049e7886f07.json b/tests/test_litellm/ocr/fixtures/data/mistral-ocr/216aacc1aa47a6060445f2ea49a2e0a8f4a8bde179faaaa28f39b049e7886f07.json similarity index 100% rename from tests/test_litellm/ocr/fixtures/mistral-ocr/216aacc1aa47a6060445f2ea49a2e0a8f4a8bde179faaaa28f39b049e7886f07.json rename to tests/test_litellm/ocr/fixtures/data/mistral-ocr/216aacc1aa47a6060445f2ea49a2e0a8f4a8bde179faaaa28f39b049e7886f07.json diff --git a/tests/test_litellm/ocr/fixtures/mistral-ocr/2c893c555fe38cdfa0f03b43b5e678b8acd5c2dc460cd06f5302a0bf05f3089c.json b/tests/test_litellm/ocr/fixtures/data/mistral-ocr/2c893c555fe38cdfa0f03b43b5e678b8acd5c2dc460cd06f5302a0bf05f3089c.json similarity index 100% rename from tests/test_litellm/ocr/fixtures/mistral-ocr/2c893c555fe38cdfa0f03b43b5e678b8acd5c2dc460cd06f5302a0bf05f3089c.json rename to tests/test_litellm/ocr/fixtures/data/mistral-ocr/2c893c555fe38cdfa0f03b43b5e678b8acd5c2dc460cd06f5302a0bf05f3089c.json diff --git a/tests/test_litellm/ocr/fixtures/mistral-ocr/3e4f8a006942b85eb954833d42cc19740883cf42bb87221eafd16df3e597ceda.json b/tests/test_litellm/ocr/fixtures/data/mistral-ocr/3e4f8a006942b85eb954833d42cc19740883cf42bb87221eafd16df3e597ceda.json similarity index 100% rename from tests/test_litellm/ocr/fixtures/mistral-ocr/3e4f8a006942b85eb954833d42cc19740883cf42bb87221eafd16df3e597ceda.json rename to tests/test_litellm/ocr/fixtures/data/mistral-ocr/3e4f8a006942b85eb954833d42cc19740883cf42bb87221eafd16df3e597ceda.json diff --git a/tests/test_litellm/ocr/fixtures/mistral-ocr/49328aed7d9a4833427d80c7ac4f54fc03ab0db9076004caf5262ce3fed4f06b.json b/tests/test_litellm/ocr/fixtures/data/mistral-ocr/49328aed7d9a4833427d80c7ac4f54fc03ab0db9076004caf5262ce3fed4f06b.json similarity index 100% rename from tests/test_litellm/ocr/fixtures/mistral-ocr/49328aed7d9a4833427d80c7ac4f54fc03ab0db9076004caf5262ce3fed4f06b.json rename to tests/test_litellm/ocr/fixtures/data/mistral-ocr/49328aed7d9a4833427d80c7ac4f54fc03ab0db9076004caf5262ce3fed4f06b.json diff --git a/tests/test_litellm/ocr/fixtures/mistral-ocr/67a074c43483176b6a8da74f21869eb5c13a7327ed36f77a40c0b192c3a1bf4f.json b/tests/test_litellm/ocr/fixtures/data/mistral-ocr/67a074c43483176b6a8da74f21869eb5c13a7327ed36f77a40c0b192c3a1bf4f.json similarity index 100% rename from tests/test_litellm/ocr/fixtures/mistral-ocr/67a074c43483176b6a8da74f21869eb5c13a7327ed36f77a40c0b192c3a1bf4f.json rename to tests/test_litellm/ocr/fixtures/data/mistral-ocr/67a074c43483176b6a8da74f21869eb5c13a7327ed36f77a40c0b192c3a1bf4f.json diff --git a/tests/test_litellm/ocr/fixtures/mistral-ocr/9f220bfe8a2042fa74e1dada33facf8c80e20574f6ef9782997eaef9adc10e0b.json b/tests/test_litellm/ocr/fixtures/data/mistral-ocr/9f220bfe8a2042fa74e1dada33facf8c80e20574f6ef9782997eaef9adc10e0b.json similarity index 100% rename from tests/test_litellm/ocr/fixtures/mistral-ocr/9f220bfe8a2042fa74e1dada33facf8c80e20574f6ef9782997eaef9adc10e0b.json rename to tests/test_litellm/ocr/fixtures/data/mistral-ocr/9f220bfe8a2042fa74e1dada33facf8c80e20574f6ef9782997eaef9adc10e0b.json diff --git a/tests/test_litellm/ocr/fixtures/mistral-ocr/afc42ea89caa932d7ea33666d91e4f1a486bbe3165f5be777e67376075fd89a8.json b/tests/test_litellm/ocr/fixtures/data/mistral-ocr/afc42ea89caa932d7ea33666d91e4f1a486bbe3165f5be777e67376075fd89a8.json similarity index 100% rename from tests/test_litellm/ocr/fixtures/mistral-ocr/afc42ea89caa932d7ea33666d91e4f1a486bbe3165f5be777e67376075fd89a8.json rename to tests/test_litellm/ocr/fixtures/data/mistral-ocr/afc42ea89caa932d7ea33666d91e4f1a486bbe3165f5be777e67376075fd89a8.json diff --git a/tests/test_litellm/ocr/fixtures/mistral-ocr/c74f9b5ad64bd97d8b4c48772c9ee17598a3a1265a630f9a6e42a2df352906bf.json b/tests/test_litellm/ocr/fixtures/data/mistral-ocr/c74f9b5ad64bd97d8b4c48772c9ee17598a3a1265a630f9a6e42a2df352906bf.json similarity index 100% rename from tests/test_litellm/ocr/fixtures/mistral-ocr/c74f9b5ad64bd97d8b4c48772c9ee17598a3a1265a630f9a6e42a2df352906bf.json rename to tests/test_litellm/ocr/fixtures/data/mistral-ocr/c74f9b5ad64bd97d8b4c48772c9ee17598a3a1265a630f9a6e42a2df352906bf.json diff --git a/tests/test_litellm/ocr/fixtures/mistral-ocr/cd2461bbef1fc222e80b3397231c79a7f6895b1d6ef5c8f544bae77854bfd951.json b/tests/test_litellm/ocr/fixtures/data/mistral-ocr/cd2461bbef1fc222e80b3397231c79a7f6895b1d6ef5c8f544bae77854bfd951.json similarity index 100% rename from tests/test_litellm/ocr/fixtures/mistral-ocr/cd2461bbef1fc222e80b3397231c79a7f6895b1d6ef5c8f544bae77854bfd951.json rename to tests/test_litellm/ocr/fixtures/data/mistral-ocr/cd2461bbef1fc222e80b3397231c79a7f6895b1d6ef5c8f544bae77854bfd951.json diff --git a/tests/test_litellm/ocr/fixtures/mistral-ocr/d1aa25a0353a8ff15c49088247a6d73ea777503f904a6bff9960a25046085d0d.json b/tests/test_litellm/ocr/fixtures/data/mistral-ocr/d1aa25a0353a8ff15c49088247a6d73ea777503f904a6bff9960a25046085d0d.json similarity index 100% rename from tests/test_litellm/ocr/fixtures/mistral-ocr/d1aa25a0353a8ff15c49088247a6d73ea777503f904a6bff9960a25046085d0d.json rename to tests/test_litellm/ocr/fixtures/data/mistral-ocr/d1aa25a0353a8ff15c49088247a6d73ea777503f904a6bff9960a25046085d0d.json diff --git a/tests/test_litellm/ocr/fixtures/mistral-ocr/e362d4eb7aa5f50706b54d0bbb777eebedf342e1844f9114d8db8b9de8a8b08b.json b/tests/test_litellm/ocr/fixtures/data/mistral-ocr/e362d4eb7aa5f50706b54d0bbb777eebedf342e1844f9114d8db8b9de8a8b08b.json similarity index 100% rename from tests/test_litellm/ocr/fixtures/mistral-ocr/e362d4eb7aa5f50706b54d0bbb777eebedf342e1844f9114d8db8b9de8a8b08b.json rename to tests/test_litellm/ocr/fixtures/data/mistral-ocr/e362d4eb7aa5f50706b54d0bbb777eebedf342e1844f9114d8db8b9de8a8b08b.json diff --git a/tests/test_litellm/ocr/fixtures/mistral-ocr/e9217f2a27a061f2ed26a567677d176db3731e07107f45318fb681740ee49a74.json b/tests/test_litellm/ocr/fixtures/data/mistral-ocr/e9217f2a27a061f2ed26a567677d176db3731e07107f45318fb681740ee49a74.json similarity index 100% rename from tests/test_litellm/ocr/fixtures/mistral-ocr/e9217f2a27a061f2ed26a567677d176db3731e07107f45318fb681740ee49a74.json rename to tests/test_litellm/ocr/fixtures/data/mistral-ocr/e9217f2a27a061f2ed26a567677d176db3731e07107f45318fb681740ee49a74.json diff --git a/tests/test_litellm/ocr/fixtures/mistral-ocr/ea3c537a691b95003afc21b351549d5a62e67fa9e6c51ba30de4eaf4bba5a73f.json b/tests/test_litellm/ocr/fixtures/data/mistral-ocr/ea3c537a691b95003afc21b351549d5a62e67fa9e6c51ba30de4eaf4bba5a73f.json similarity index 100% rename from tests/test_litellm/ocr/fixtures/mistral-ocr/ea3c537a691b95003afc21b351549d5a62e67fa9e6c51ba30de4eaf4bba5a73f.json rename to tests/test_litellm/ocr/fixtures/data/mistral-ocr/ea3c537a691b95003afc21b351549d5a62e67fa9e6c51ba30de4eaf4bba5a73f.json diff --git a/tests/test_litellm/ocr/fixtures/mistral-ocr/effe3c112538d2aa6ac98e8f47033086cb79e87d2bdb99bdbb3b7e11751c43d7.json b/tests/test_litellm/ocr/fixtures/data/mistral-ocr/effe3c112538d2aa6ac98e8f47033086cb79e87d2bdb99bdbb3b7e11751c43d7.json similarity index 100% rename from tests/test_litellm/ocr/fixtures/mistral-ocr/effe3c112538d2aa6ac98e8f47033086cb79e87d2bdb99bdbb3b7e11751c43d7.json rename to tests/test_litellm/ocr/fixtures/data/mistral-ocr/effe3c112538d2aa6ac98e8f47033086cb79e87d2bdb99bdbb3b7e11751c43d7.json diff --git a/tests/test_litellm/ocr/fixtures/mistral-ocr/f430ede1f4069abab794731b9c233573014c6353f8d2486cc34753ceab6c7c1a.json b/tests/test_litellm/ocr/fixtures/data/mistral-ocr/f430ede1f4069abab794731b9c233573014c6353f8d2486cc34753ceab6c7c1a.json similarity index 100% rename from tests/test_litellm/ocr/fixtures/mistral-ocr/f430ede1f4069abab794731b9c233573014c6353f8d2486cc34753ceab6c7c1a.json rename to tests/test_litellm/ocr/fixtures/data/mistral-ocr/f430ede1f4069abab794731b9c233573014c6353f8d2486cc34753ceab6c7c1a.json diff --git a/tests/test_litellm/ocr/fixtures/reducto-legacy/5671d085483943d56d3e16d95a0a8cf08e7dc5bbe1f7286d7c04a0454baf5c66.json b/tests/test_litellm/ocr/fixtures/data/reducto-legacy/5671d085483943d56d3e16d95a0a8cf08e7dc5bbe1f7286d7c04a0454baf5c66.json similarity index 100% rename from tests/test_litellm/ocr/fixtures/reducto-legacy/5671d085483943d56d3e16d95a0a8cf08e7dc5bbe1f7286d7c04a0454baf5c66.json rename to tests/test_litellm/ocr/fixtures/data/reducto-legacy/5671d085483943d56d3e16d95a0a8cf08e7dc5bbe1f7286d7c04a0454baf5c66.json diff --git a/tests/test_litellm/ocr/fixtures/reducto-legacy/ea3262579a0882bcb61972e2a2aa4b0c7ae2e28ce957493e996a006ef2fcc6be.json b/tests/test_litellm/ocr/fixtures/data/reducto-legacy/ea3262579a0882bcb61972e2a2aa4b0c7ae2e28ce957493e996a006ef2fcc6be.json similarity index 100% rename from tests/test_litellm/ocr/fixtures/reducto-legacy/ea3262579a0882bcb61972e2a2aa4b0c7ae2e28ce957493e996a006ef2fcc6be.json rename to tests/test_litellm/ocr/fixtures/data/reducto-legacy/ea3262579a0882bcb61972e2a2aa4b0c7ae2e28ce957493e996a006ef2fcc6be.json diff --git a/tests/test_litellm/ocr/fixtures/reducto-legacy/fbf21ae169e9ddbca638f445bad921e3e782e4e9f10be9a4ba000a8caf1e9ab5.json b/tests/test_litellm/ocr/fixtures/data/reducto-legacy/fbf21ae169e9ddbca638f445bad921e3e782e4e9f10be9a4ba000a8caf1e9ab5.json similarity index 100% rename from tests/test_litellm/ocr/fixtures/reducto-legacy/fbf21ae169e9ddbca638f445bad921e3e782e4e9f10be9a4ba000a8caf1e9ab5.json rename to tests/test_litellm/ocr/fixtures/data/reducto-legacy/fbf21ae169e9ddbca638f445bad921e3e782e4e9f10be9a4ba000a8caf1e9ab5.json diff --git a/tests/test_litellm/ocr/fixtures/reducto-v3/00312c58e0dccd4699086d6c37b45a56b19b4775393c22bf9fa6ead2c3cc323c.json b/tests/test_litellm/ocr/fixtures/data/reducto-v3/00312c58e0dccd4699086d6c37b45a56b19b4775393c22bf9fa6ead2c3cc323c.json similarity index 100% rename from tests/test_litellm/ocr/fixtures/reducto-v3/00312c58e0dccd4699086d6c37b45a56b19b4775393c22bf9fa6ead2c3cc323c.json rename to tests/test_litellm/ocr/fixtures/data/reducto-v3/00312c58e0dccd4699086d6c37b45a56b19b4775393c22bf9fa6ead2c3cc323c.json diff --git a/tests/test_litellm/ocr/fixtures/reducto-v3/0b4ca483494e37fce75cc861ba62f06ce970d117de9828d72ed76b71b6f5adc1.json b/tests/test_litellm/ocr/fixtures/data/reducto-v3/0b4ca483494e37fce75cc861ba62f06ce970d117de9828d72ed76b71b6f5adc1.json similarity index 100% rename from tests/test_litellm/ocr/fixtures/reducto-v3/0b4ca483494e37fce75cc861ba62f06ce970d117de9828d72ed76b71b6f5adc1.json rename to tests/test_litellm/ocr/fixtures/data/reducto-v3/0b4ca483494e37fce75cc861ba62f06ce970d117de9828d72ed76b71b6f5adc1.json diff --git a/tests/test_litellm/ocr/fixtures/reducto-v3/3a512e502dc6310e671e8967a7652520370e0c873939813c2167799cffacba06.json b/tests/test_litellm/ocr/fixtures/data/reducto-v3/3a512e502dc6310e671e8967a7652520370e0c873939813c2167799cffacba06.json similarity index 100% rename from tests/test_litellm/ocr/fixtures/reducto-v3/3a512e502dc6310e671e8967a7652520370e0c873939813c2167799cffacba06.json rename to tests/test_litellm/ocr/fixtures/data/reducto-v3/3a512e502dc6310e671e8967a7652520370e0c873939813c2167799cffacba06.json diff --git a/tests/test_litellm/ocr/fixtures/reducto-v3/5227441615ac16de68a18458e9704be5c772291d89dd5a67480c21e4e2ac3834.json b/tests/test_litellm/ocr/fixtures/data/reducto-v3/5227441615ac16de68a18458e9704be5c772291d89dd5a67480c21e4e2ac3834.json similarity index 100% rename from tests/test_litellm/ocr/fixtures/reducto-v3/5227441615ac16de68a18458e9704be5c772291d89dd5a67480c21e4e2ac3834.json rename to tests/test_litellm/ocr/fixtures/data/reducto-v3/5227441615ac16de68a18458e9704be5c772291d89dd5a67480c21e4e2ac3834.json diff --git a/tests/test_litellm/ocr/fixtures/reducto-v3/75876320a008dc29db48c33e4d4ccb9e95c3aae1d774f03869cf7cacfaa9aa07.json b/tests/test_litellm/ocr/fixtures/data/reducto-v3/75876320a008dc29db48c33e4d4ccb9e95c3aae1d774f03869cf7cacfaa9aa07.json similarity index 100% rename from tests/test_litellm/ocr/fixtures/reducto-v3/75876320a008dc29db48c33e4d4ccb9e95c3aae1d774f03869cf7cacfaa9aa07.json rename to tests/test_litellm/ocr/fixtures/data/reducto-v3/75876320a008dc29db48c33e4d4ccb9e95c3aae1d774f03869cf7cacfaa9aa07.json diff --git a/tests/test_litellm/ocr/fixtures/reducto-v3/aee30a1d7a51a3fd8837b56d6baa3a43364ea319d905b5030014f4d85e7c94e9.json b/tests/test_litellm/ocr/fixtures/data/reducto-v3/aee30a1d7a51a3fd8837b56d6baa3a43364ea319d905b5030014f4d85e7c94e9.json similarity index 100% rename from tests/test_litellm/ocr/fixtures/reducto-v3/aee30a1d7a51a3fd8837b56d6baa3a43364ea319d905b5030014f4d85e7c94e9.json rename to tests/test_litellm/ocr/fixtures/data/reducto-v3/aee30a1d7a51a3fd8837b56d6baa3a43364ea319d905b5030014f4d85e7c94e9.json diff --git a/tests/test_litellm/ocr/fixtures/reducto-v3/cfb4093138a106055f36402de37409b1ca17c23b9e491b4f8b6ac705acd05637.json b/tests/test_litellm/ocr/fixtures/data/reducto-v3/cfb4093138a106055f36402de37409b1ca17c23b9e491b4f8b6ac705acd05637.json similarity index 100% rename from tests/test_litellm/ocr/fixtures/reducto-v3/cfb4093138a106055f36402de37409b1ca17c23b9e491b4f8b6ac705acd05637.json rename to tests/test_litellm/ocr/fixtures/data/reducto-v3/cfb4093138a106055f36402de37409b1ca17c23b9e491b4f8b6ac705acd05637.json diff --git a/tests/test_litellm/ocr/fixtures/reducto-v3/d5a49e5897e920719ca28697413ee0396ea754023bd5ab34065efcc852914f15.json b/tests/test_litellm/ocr/fixtures/data/reducto-v3/d5a49e5897e920719ca28697413ee0396ea754023bd5ab34065efcc852914f15.json similarity index 100% rename from tests/test_litellm/ocr/fixtures/reducto-v3/d5a49e5897e920719ca28697413ee0396ea754023bd5ab34065efcc852914f15.json rename to tests/test_litellm/ocr/fixtures/data/reducto-v3/d5a49e5897e920719ca28697413ee0396ea754023bd5ab34065efcc852914f15.json diff --git a/tests/test_litellm/ocr/fixtures/generate.py b/tests/test_litellm/ocr/fixtures/generate.py new file mode 100644 index 00000000000..60bbe851b5d --- /dev/null +++ b/tests/test_litellm/ocr/fixtures/generate.py @@ -0,0 +1,92 @@ +from __future__ import annotations + +import logging +import os +from collections.abc import Mapping +from pathlib import Path +from typing import Final, cast + +from dotenv import load_dotenv + +import litellm +from litellm.rust_bridge.ocr import use_litellm_rust +from tests.route_parity.fixture_generator import ( + FixtureProvider, + FixtureSdkCall, + discover_fixture_targets, + generate_target_fixtures, + parse_generator_args, +) +from tests.route_parity.fixture_generator import require_targets as require_fixture_targets +from tests.route_parity.fixture_recorder import fixture_directory +from tests.test_litellm.ocr.fixtures.azure import ( + AzureDocumentIntelligenceFixtureProvider, + AzureMistralFixtureProvider, + azure_document_intelligence_input_strategy, +) +from tests.test_litellm.ocr.fixtures.common import OcrFixtureTarget +from tests.test_litellm.ocr.fixtures.mistral import ( + MistralFixtureProvider, + mistral_input_strategy, +) +from tests.test_litellm.ocr.fixtures.models import OcrParityCase, OcrSdkInputBase +from tests.test_litellm.ocr.fixtures.reducto import ( + ReductoFixtureProvider, + reducto_legacy_input_strategy, + reducto_v3_input_strategy, +) +from tests.test_litellm.ocr.fixtures.vertex import ( + VertexFixtureProvider, + vertex_deepseek_input_strategy, +) + +__all__ = ( + "azure_document_intelligence_input_strategy", + "mistral_input_strategy", + "reducto_legacy_input_strategy", + "reducto_v3_input_strategy", + "vertex_deepseek_input_strategy", +) + +FIXTURE_DIR_ENV: Final = "LITELLM_OCR_FIXTURE_DIR" +OCR_FIXTURE_PROVIDERS: Final[tuple[FixtureProvider[OcrSdkInputBase], ...]] = ( + MistralFixtureProvider(), + AzureMistralFixtureProvider(), + AzureDocumentIntelligenceFixtureProvider(), + VertexFixtureProvider(), + ReductoFixtureProvider(), +) + + +def discover_targets( + environ: Mapping[str, str], + sdk_call: FixtureSdkCall, +) -> tuple[OcrFixtureTarget, ...]: + return discover_fixture_targets(OCR_FIXTURE_PROVIDERS, environ, sdk_call) + + +def require_targets(targets: tuple[OcrFixtureTarget, ...]) -> tuple[OcrFixtureTarget, ...]: + return require_fixture_targets( + targets, + "No OCR fixture providers are configured. Set a supported provider API key and endpoint", + ) + + +def main() -> None: + logging.basicConfig(level=logging.INFO, format="%(message)s") + load_dotenv() + args: Final = parse_generator_args() + sdk_call: Final = cast(FixtureSdkCall, litellm.ocr) + targets: Final = require_targets(discover_targets(os.environ, sdk_call)) + root: Final = fixture_directory( + args.fixture_dir, + os.environ.get(FIXTURE_DIR_ENV), + Path(__file__).with_name("data"), + ) + use_litellm_rust(False, ocr=None, aocr=None) + for target in targets: + generate_target_fixtures(target, root, args.examples, args.concurrency, OcrParityCase) + + +if __name__ == "__main__": + main() diff --git a/tests/test_litellm/ocr/fixtures/mistral.py b/tests/test_litellm/ocr/fixtures/mistral.py new file mode 100644 index 00000000000..fc585b3ec2d --- /dev/null +++ b/tests/test_litellm/ocr/fixtures/mistral.py @@ -0,0 +1,98 @@ +from __future__ import annotations + +from collections.abc import Mapping +from typing import Final, cast + +from hypothesis import strategies as st +from hypothesis.strategies import DrawFn, SearchStrategy + +from tests.route_parity.fixture_generator import FixtureSdkCall +from tests.route_parity.fixture_recorder import ProviderSpec +from tests.test_litellm.ocr.fixtures.common import ( + OcrFixtureTarget, + annotation_format, + image_document, + invoke_with_api_key, + public_document_strategy, +) +from tests.test_litellm.ocr.fixtures.models import MistralOcrSdkInput, OcrSdkInputBase + +MISTRAL_MODEL: Final = "mistral/mistral-ocr-latest" +_VALUE_TEXT: Final = st.just("case-1") + + +@st.composite +def mistral_input_strategy(draw: DrawFn, model: str) -> MistralOcrSdkInput: + document: Final = draw(public_document_strategy()) + annotation: Final = draw( + st.sampled_from( + ( + {}, + {"document_annotation_format": annotation_format("document_title")}, + { + "document_annotation_format": annotation_format("prompted_document_title"), + "document_annotation_prompt": "Extract the visible title", + }, + ) + ) + ) + optional_params: Final = draw( + st.fixed_dictionaries( + {}, + optional={ + "pages": st.just([0]), + "include_image_base64": st.booleans(), + "image_limit": st.just(1), + "image_min_size": st.just(300), + "bbox_annotation_format": st.just(annotation_format("bounding_boxes")), + "extract_header": st.booleans(), + "extract_footer": st.booleans(), + "table_format": st.just("markdown"), + "confidence_scores_granularity": st.just("page"), + "include_blocks": st.booleans(), + "id": _VALUE_TEXT, + }, + ) + ) + return MistralOcrSdkInput.model_validate({"model": model, "document": document, **annotation, **optional_params}) + + +def required_mistral_inputs(model: str) -> tuple[MistralOcrSdkInput, ...]: + document: Final = image_document("invoice 123", 24) + annotation: Final = annotation_format("document_title") + bbox_annotation: Final = annotation_format("bounding_boxes") + cases: Final[tuple[dict[str, object], ...]] = ( + {}, + {"pages": [0]}, + {"include_image_base64": True}, + {"image_limit": 1}, + {"image_min_size": 300}, + {"bbox_annotation_format": bbox_annotation}, + {"document_annotation_format": annotation}, + {"document_annotation_format": annotation, "document_annotation_prompt": "Extract the visible title"}, + {"extract_header": True}, + {"extract_footer": True}, + {"table_format": "markdown"}, + {"confidence_scores_granularity": "page"}, + {"include_blocks": False}, + {"id": "case-1"}, + ) + return tuple(MistralOcrSdkInput.model_validate({"model": model, "document": document, **case}) for case in cases) + + +class MistralFixtureProvider: + def targets(self, environ: Mapping[str, str], sdk_call: FixtureSdkCall) -> tuple[OcrFixtureTarget, ...]: + api_key: Final = environ.get("MISTRAL_API_KEY") + if not api_key: + return () + configured: Final = environ.get("MISTRAL_API_BASE", "https://api.mistral.ai").rstrip("/") + upstream_base: Final = configured.removesuffix("/v1") + return ( + OcrFixtureTarget( + name="mistral-ocr", + provider_spec=ProviderSpec(upstream_base=upstream_base), + strategy=cast(SearchStrategy[OcrSdkInputBase], mistral_input_strategy(MISTRAL_MODEL)), + invoke=invoke_with_api_key(sdk_call, api_key), + required_inputs=cast(tuple[OcrSdkInputBase, ...], required_mistral_inputs(MISTRAL_MODEL)), + ), + ) diff --git a/tests/test_litellm/ocr/fixture_models.py b/tests/test_litellm/ocr/fixtures/models.py similarity index 100% rename from tests/test_litellm/ocr/fixture_models.py rename to tests/test_litellm/ocr/fixtures/models.py diff --git a/tests/test_litellm/ocr/fixtures/reducto.py b/tests/test_litellm/ocr/fixtures/reducto.py new file mode 100644 index 00000000000..bc28c483f70 --- /dev/null +++ b/tests/test_litellm/ocr/fixtures/reducto.py @@ -0,0 +1,187 @@ +from __future__ import annotations + +from collections.abc import Mapping +from typing import Final, cast + +from hypothesis import strategies as st +from hypothesis.strategies import DrawFn, SearchStrategy + +from tests.route_parity.fixture_generator import FixtureSdkCall +from tests.route_parity.fixture_recorder import ProviderSpec +from tests.test_litellm.ocr.fixtures.common import ( + OcrFixtureTarget, + fixture_pdf_data_uri, + invoke_with_api_key, +) +from tests.test_litellm.ocr.fixtures.models import ( + OcrSdkInputBase, + ReductoChunking, + ReductoDocumentUrlDocument, + ReductoFormatting, + ReductoParseLegacySdkInput, + ReductoParseV3SdkInput, + ReductoRetrieval, + ReductoSettings, +) + +_REDUCTO_API_BASE: Final = "https://platform.reducto.ai" + + +def _formatting_strategy() -> SearchStrategy[ReductoFormatting]: + return st.builds( + ReductoFormatting, + add_page_markers=st.booleans(), + table_output_format=st.sampled_from(("html", "json", "md", "jsonbbox", "dynamic", "csv")), + merge_tables=st.booleans(), + include=st.sampled_from( + ( + [], + ["hyperlinks"], + ["change_tracking", "highlight", "comments"], + ["signatures", "ignore_watermarks"], + ) + ), + ) + + +def _chunking_strategy() -> SearchStrategy[ReductoChunking]: + return st.one_of( + st.builds( + ReductoChunking, + chunk_mode=st.sampled_from(("section", "page", "disabled", "block", "page_sections")), + chunk_size=st.just(None), + chunk_overlap=st.just(0), + ), + st.builds( + ReductoChunking, + chunk_mode=st.just("variable"), + chunk_size=st.sampled_from((250, 1000, 1500)), + chunk_overlap=st.sampled_from((0, 32, 128)), + ), + ) + + +def _retrieval_strategy() -> SearchStrategy[ReductoRetrieval]: + return st.builds( + ReductoRetrieval, + chunking=_chunking_strategy(), + filter_blocks=st.sampled_from( + ( + [], + ["Header"], + ["Header", "Footer", "Page Number"], + ["Figure", "Table", "Key Value"], + ) + ), + embedding_optimized=st.booleans(), + ) + + +def _settings_strategy() -> SearchStrategy[ReductoSettings]: + return st.builds( + ReductoSettings, + ocr_system=st.sampled_from(("standard", "legacy")), + extraction_mode=st.sampled_from(("ocr", "hybrid")), + force_url_result=st.just(False), + return_ocr_data=st.booleans(), + return_images=st.sampled_from(([], ["figure"], ["table"], ["page"], ["figure", "table", "page"])), + embed_pdf_metadata=st.booleans(), + embed_pdf_metadata_dpi=st.sampled_from((50, 100, 250)), + persist_results=st.just(False), + timeout=st.sampled_from((None, 300.0, 900.0)), + page_range=st.sampled_from((None, [1], [1, 2], ["Sheet1"])), + ) + + +@st.composite +def reducto_v3_input_strategy( + draw: DrawFn, document: ReductoDocumentUrlDocument | None = None +) -> ReductoParseV3SdkInput: + model, custom_llm_provider = draw(st.sampled_from((("reducto/parse-v3", None), ("parse-v3", "reducto")))) + options: Final = draw( + st.fixed_dictionaries( + {}, + optional={ + "formatting": _formatting_strategy(), + "retrieval": _retrieval_strategy(), + "settings": _settings_strategy(), + }, + ) + ) + return ReductoParseV3SdkInput.model_validate( + { + "model": model, + "custom_llm_provider": custom_llm_provider, + "document": document + or ReductoDocumentUrlDocument(type="document_url", document_url="reducto://fixture-document.pdf"), + **options, + } + ) + + +def reducto_legacy_input_strategy( + document: ReductoDocumentUrlDocument | None = None, +) -> SearchStrategy[ReductoParseLegacySdkInput]: + selected_document: Final = document or ReductoDocumentUrlDocument( + type="document_url", document_url="reducto://fixture-document.pdf" + ) + return st.sampled_from( + ( + ReductoParseLegacySdkInput(model="reducto/parse-legacy", document=selected_document), + ReductoParseLegacySdkInput(model="parse-legacy", custom_llm_provider="reducto", document=selected_document), + ReductoParseLegacySdkInput(model="reducto/parse-legacy", document=selected_document, enhance={}), + ) + ) + + +def _required_v3_inputs(document: ReductoDocumentUrlDocument) -> tuple[ReductoParseV3SdkInput, ...]: + return ( + ReductoParseV3SdkInput(model="reducto/parse-v3", document=document), + ReductoParseV3SdkInput( + model="reducto/parse-v3", + document=document, + formatting=ReductoFormatting(table_output_format="md"), + ), + ReductoParseV3SdkInput( + model="reducto/parse-v3", + document=document, + retrieval=ReductoRetrieval(chunking=ReductoChunking(chunk_mode="page")), + ), + ReductoParseV3SdkInput( + model="reducto/parse-v3", + document=document, + settings=ReductoSettings(return_ocr_data=True), + ), + ) + + +class ReductoFixtureProvider: + def targets(self, environ: Mapping[str, str], sdk_call: FixtureSdkCall) -> tuple[OcrFixtureTarget, ...]: + api_key: Final = environ.get("REDUCTO_API_KEY") + if not api_key: + return () + upstream_base: Final = environ.get("REDUCTO_API_BASE", _REDUCTO_API_BASE).rstrip("/") + document: Final = ReductoDocumentUrlDocument(type="document_url", document_url=fixture_pdf_data_uri()) + invoke: Final = invoke_with_api_key(sdk_call, api_key) + return ( + OcrFixtureTarget( + name="reducto-v3", + provider_spec=ProviderSpec(upstream_base=upstream_base), + strategy=cast(SearchStrategy[OcrSdkInputBase], reducto_v3_input_strategy(document)), + invoke=invoke, + required_inputs=cast(tuple[OcrSdkInputBase, ...], _required_v3_inputs(document)), + ), + OcrFixtureTarget( + name="reducto-legacy", + provider_spec=ProviderSpec(upstream_base=upstream_base), + strategy=cast(SearchStrategy[OcrSdkInputBase], reducto_legacy_input_strategy(document)), + invoke=invoke, + required_inputs=cast( + tuple[OcrSdkInputBase, ...], + ( + ReductoParseLegacySdkInput(model="reducto/parse-legacy", document=document), + ReductoParseLegacySdkInput(model="reducto/parse-legacy", document=document, enhance={}), + ), + ), + ), + ) diff --git a/tests/test_litellm/ocr/fixtures/vertex.py b/tests/test_litellm/ocr/fixtures/vertex.py new file mode 100644 index 00000000000..aa79b0f9b48 --- /dev/null +++ b/tests/test_litellm/ocr/fixtures/vertex.py @@ -0,0 +1,113 @@ +from __future__ import annotations + +from collections.abc import Mapping +from typing import Final, cast + +from hypothesis import strategies as st +from hypothesis.strategies import DrawFn, SearchStrategy + +from tests.route_parity.fixture_generator import FixtureSdkCall +from tests.route_parity.fixture_recorder import ProviderSpec +from tests.test_litellm.ocr.fixtures.common import ( + OcrFixtureTarget, + image_document, + invoke_with_api_key, + public_document_strategy, +) +from tests.test_litellm.ocr.fixtures.mistral import ( + MISTRAL_MODEL, + mistral_input_strategy, + required_mistral_inputs, +) +from tests.test_litellm.ocr.fixtures.models import ( + MistralOcrSdkInput, + OcrSdkInputBase, + VertexDeepSeekOcrSdkInput, + VertexMistralOcrSdkInput, +) + + +def _as_vertex_mistral(case_input: MistralOcrSdkInput, project: str, location: str) -> VertexMistralOcrSdkInput: + values: Final = case_input.model_dump(mode="python", exclude={"boundary", "model", "custom_llm_provider"}) + return VertexMistralOcrSdkInput.model_validate({**values, "vertex_project": project, "vertex_location": location}) + + +def _required_deepseek_inputs(project: str, location: str) -> tuple[VertexDeepSeekOcrSdkInput, ...]: + document: Final = image_document("invoice 123", 24) + common: Final = {"document": document, "vertex_project": project, "vertex_location": location} + cases: Final[tuple[dict[str, object], ...]] = ( + {}, + {"stream": False}, + {"temperature": 0.5}, + {"max_tokens": 256}, + {"top_p": 0.9}, + {"n": 1}, + {"stop": ["END", "STOP"]}, + ) + return tuple(VertexDeepSeekOcrSdkInput.model_validate({**common, **case}) for case in cases) + + +@st.composite +def vertex_deepseek_input_strategy(draw: DrawFn, project: str, location: str) -> VertexDeepSeekOcrSdkInput: + optional_params: Final = draw( + st.fixed_dictionaries( + {}, + optional={ + "stream": st.just(False), + "temperature": st.sampled_from((0.0, 0.5, 1.0)), + "max_tokens": st.sampled_from((1, 256, 1024)), + "top_p": st.sampled_from((0.1, 0.9, 1.0)), + "n": st.just(1), + "stop": st.sampled_from(("END", ["END", "STOP"])), + }, + ) + ) + return VertexDeepSeekOcrSdkInput.model_validate( + { + "document": draw(public_document_strategy()), + "vertex_project": project, + "vertex_location": location, + **optional_params, + } + ) + + +class VertexFixtureProvider: + def targets(self, environ: Mapping[str, str], sdk_call: FixtureSdkCall) -> tuple[OcrFixtureTarget, ...]: + api_key: Final = environ.get("VERTEX_AI_API_KEY") + project: Final = environ.get("VERTEXAI_PROJECT") or environ.get("VERTEX_PROJECT") + location: Final = environ.get("VERTEXAI_LOCATION") or environ.get("VERTEX_LOCATION") or "us-central1" + if not api_key or not project: + return () + upstream_base: Final = environ.get("VERTEX_AI_API_BASE") or f"https://{location}-aiplatform.googleapis.com" + invoke: Final = invoke_with_api_key(sdk_call, api_key) + return ( + OcrFixtureTarget( + name="vertex-mistral", + provider_spec=ProviderSpec(upstream_base=upstream_base.rstrip("/")), + strategy=cast( + SearchStrategy[OcrSdkInputBase], + st.builds( + _as_vertex_mistral, + case_input=mistral_input_strategy(MISTRAL_MODEL), + project=st.just(project), + location=st.just(location), + ), + ), + invoke=invoke, + required_inputs=cast( + tuple[OcrSdkInputBase, ...], + tuple( + _as_vertex_mistral(case_input, project, location) + for case_input in required_mistral_inputs(MISTRAL_MODEL) + ), + ), + ), + OcrFixtureTarget( + name="vertex-deepseek", + provider_spec=ProviderSpec(upstream_base=upstream_base.rstrip("/")), + strategy=cast(SearchStrategy[OcrSdkInputBase], vertex_deepseek_input_strategy(project, location)), + invoke=invoke, + required_inputs=cast(tuple[OcrSdkInputBase, ...], _required_deepseek_inputs(project, location)), + ), + ) diff --git a/tests/test_litellm/ocr/generate_fixtures.py b/tests/test_litellm/ocr/generate_fixtures.py deleted file mode 100644 index 82533aecf20..00000000000 --- a/tests/test_litellm/ocr/generate_fixtures.py +++ /dev/null @@ -1,605 +0,0 @@ -from __future__ import annotations - -import base64 -import logging -import os -from collections.abc import Callable, Mapping -from pathlib import Path -from typing import Final, cast -from urllib.parse import quote - -from dotenv import load_dotenv -from hypothesis import strategies as st -from hypothesis.strategies import DrawFn, SearchStrategy - -import litellm -from litellm.rust_bridge.ocr import use_litellm_rust -from tests.route_parity.fixture_generator import ( - FixtureTarget, - generate_target_fixtures, - parse_generator_args, -) -from tests.route_parity.fixture_generator import ( - require_targets as require_fixture_targets, -) -from tests.route_parity.fixture_recorder import ( - ProviderSpec, - fixture_directory, -) -from tests.test_litellm.ocr.fixture_models import ( - AzureDocumentIntelligenceOcrSdkInput, - AzureMistralOcrSdkInput, - JsonSchemaDefinition, - JsonSchemaResponseFormat, - MistralDocumentUrlDocument, - MistralImageUrlDocument, - MistralOcrSdkInput, - OcrParityCase, - OcrSdkInputBase, - ReductoChunking, - ReductoDocumentUrlDocument, - ReductoFormatting, - ReductoParseLegacySdkInput, - ReductoParseV3SdkInput, - ReductoRetrieval, - ReductoSettings, - VertexDeepSeekOcrSdkInput, - VertexMistralOcrSdkInput, -) - -FIXTURE_DIR_ENV: Final = "LITELLM_OCR_FIXTURE_DIR" -_VALUE_TEXT: Final = st.just("case-1") -_MISTRAL_MODEL: Final = "mistral/mistral-ocr-latest" -_REDUCTO_API_BASE: Final = "https://platform.reducto.ai" - - -OcrFixtureTarget = FixtureTarget[OcrSdkInputBase] - - -def _image_document(text: str, font_size: int) -> MistralImageUrlDocument: - url: Final = f"https://dummyjson.com/image/800x300/ffffff/000000?text={quote(text)}&fontSize={font_size}" - return MistralImageUrlDocument(type="image_url", image_url=url) - - -def _fixture_pdf_data_uri() -> str: - fixture: Final = Path(__file__).resolve().parents[2] / "llm_translation" / "fixtures" / "dummy.pdf" - encoded: Final = base64.b64encode(fixture.read_bytes()).decode("ascii") - return f"data:application/pdf;base64,{encoded}" - - -def _pdf_document() -> MistralDocumentUrlDocument: - return MistralDocumentUrlDocument(type="document_url", document_url=_fixture_pdf_data_uri()) - - -def _public_document_strategy() -> SearchStrategy[MistralImageUrlDocument | MistralDocumentUrlDocument]: - return st.sampled_from((_image_document("invoice 123", 24), _pdf_document())) - - -def _annotation_format(name: str) -> JsonSchemaResponseFormat: - return JsonSchemaResponseFormat( - type="json_schema", - json_schema=JsonSchemaDefinition( - name=name, - description="Extract the visible document fields", - schema={ - "type": "object", - "properties": {"title": {"type": "string"}}, - "required": ["title"], - "additionalProperties": False, - }, - strict=True, - ), - ) - - -def _mistral_confidence_strategy() -> SearchStrategy[str]: - return st.just("page") - - -@st.composite -def mistral_input_strategy(draw: DrawFn, model: str) -> MistralOcrSdkInput: - document: Final = draw(_public_document_strategy()) - annotation: Final = draw( - st.sampled_from( - ( - {}, - {"document_annotation_format": _annotation_format("document_title")}, - { - "document_annotation_format": _annotation_format("prompted_document_title"), - "document_annotation_prompt": "Extract the visible title", - }, - ) - ) - ) - optional_params: Final = draw( - st.fixed_dictionaries( - {}, - optional={ - "pages": st.just([0]), - "include_image_base64": st.booleans(), - "image_limit": st.just(1), - "image_min_size": st.just(300), - "bbox_annotation_format": st.just(_annotation_format("bounding_boxes")), - "extract_header": st.booleans(), - "extract_footer": st.booleans(), - "table_format": st.just("markdown"), - "confidence_scores_granularity": _mistral_confidence_strategy(), - "include_blocks": st.booleans(), - "id": _VALUE_TEXT, - }, - ) - ) - return MistralOcrSdkInput.model_validate( - { - "model": model, - "document": document, - **annotation, - **optional_params, - } - ) - - -def _as_azure_mistral(case_input: MistralOcrSdkInput, model: str) -> AzureMistralOcrSdkInput: - values: Final = case_input.model_dump(mode="python", exclude={"boundary", "model", "custom_llm_provider"}) - return AzureMistralOcrSdkInput.model_validate({**values, "model": model}) - - -def _as_vertex_mistral(case_input: MistralOcrSdkInput, project: str, location: str) -> VertexMistralOcrSdkInput: - values: Final = case_input.model_dump(mode="python", exclude={"boundary", "model", "custom_llm_provider"}) - return VertexMistralOcrSdkInput.model_validate( - {**values, "vertex_project": project, "vertex_location": location} - ) - - -def _required_mistral_inputs(model: str) -> tuple[MistralOcrSdkInput, ...]: - document: Final = _image_document("invoice 123", 24) - annotation: Final = _annotation_format("document_title") - bbox_annotation: Final = _annotation_format("bounding_boxes") - cases: Final[tuple[dict[str, object], ...]] = ( - {}, - {"pages": [0]}, - {"include_image_base64": True}, - {"image_limit": 1}, - {"image_min_size": 300}, - {"bbox_annotation_format": bbox_annotation}, - {"document_annotation_format": annotation}, - {"document_annotation_format": annotation, "document_annotation_prompt": "Extract the visible title"}, - {"extract_header": True}, - {"extract_footer": True}, - {"table_format": "markdown"}, - {"confidence_scores_granularity": "page"}, - {"include_blocks": False}, - {"id": "case-1"}, - ) - return tuple( - MistralOcrSdkInput.model_validate({"model": model, "document": document, **case}) for case in cases - ) - - -def _required_azure_document_intelligence_inputs() -> tuple[AzureDocumentIntelligenceOcrSdkInput, ...]: - document: Final = _pdf_document() - model: Final = "azure_ai/doc-intelligence/prebuilt-layout" - return ( - AzureDocumentIntelligenceOcrSdkInput(model=model, document=document), - AzureDocumentIntelligenceOcrSdkInput(model=model, document=document, pages=[0, 1]), - AzureDocumentIntelligenceOcrSdkInput(model=model, document=document, features=["languages"]), - AzureDocumentIntelligenceOcrSdkInput(model=model, document=document, req_format="litellm"), - ) - - -def _required_vertex_deepseek_inputs(project: str, location: str) -> tuple[VertexDeepSeekOcrSdkInput, ...]: - document: Final = _image_document("invoice 123", 24) - common: Final = {"document": document, "vertex_project": project, "vertex_location": location} - cases: Final[tuple[dict[str, object], ...]] = ( - {}, - {"stream": False}, - {"temperature": 0.5}, - {"max_tokens": 256}, - {"top_p": 0.9}, - {"n": 1}, - {"stop": ["END", "STOP"]}, - ) - return tuple(VertexDeepSeekOcrSdkInput.model_validate({**common, **case}) for case in cases) - - -def _required_reducto_v3_inputs( - document: ReductoDocumentUrlDocument, -) -> tuple[ReductoParseV3SdkInput, ...]: - return ( - ReductoParseV3SdkInput(model="reducto/parse-v3", document=document), - ReductoParseV3SdkInput( - model="reducto/parse-v3", - document=document, - formatting=ReductoFormatting(table_output_format="md"), - ), - ReductoParseV3SdkInput( - model="reducto/parse-v3", - document=document, - retrieval=ReductoRetrieval(chunking=ReductoChunking(chunk_mode="page")), - ), - ReductoParseV3SdkInput( - model="reducto/parse-v3", - document=document, - settings=ReductoSettings(return_ocr_data=True), - ), - ) - - -@st.composite -def azure_document_intelligence_input_strategy(draw: DrawFn) -> AzureDocumentIntelligenceOcrSdkInput: - optional_params: Final = draw( - st.fixed_dictionaries( - {}, - optional={ - "pages": st.sampled_from(([0], [0, 1], "1-2")), - "features": st.sampled_from((["languages"], ["keyValuePairs"], "languages,keyValuePairs")), - }, - ) - ) - return AzureDocumentIntelligenceOcrSdkInput.model_validate( - { - "model": draw( - st.sampled_from( - ( - "azure_ai/doc-intelligence/prebuilt-read", - "azure_ai/doc-intelligence/prebuilt-layout", - "azure_ai/doc-intelligence/prebuilt-document", - ) - ) - ), - "document": draw(_public_document_strategy()), - **optional_params, - } - ) - - -@st.composite -def vertex_deepseek_input_strategy(draw: DrawFn, project: str, location: str) -> VertexDeepSeekOcrSdkInput: - optional_params: Final = draw( - st.fixed_dictionaries( - {}, - optional={ - "stream": st.just(False), - "temperature": st.sampled_from((0.0, 0.5, 1.0)), - "max_tokens": st.sampled_from((1, 256, 1024)), - "top_p": st.sampled_from((0.1, 0.9, 1.0)), - "n": st.just(1), - "stop": st.sampled_from(("END", ["END", "STOP"])), - }, - ) - ) - return VertexDeepSeekOcrSdkInput.model_validate( - { - "document": draw(_public_document_strategy()), - "vertex_project": project, - "vertex_location": location, - **optional_params, - } - ) - - -def _reducto_formatting_strategy() -> SearchStrategy[ReductoFormatting]: - return st.builds( - ReductoFormatting, - add_page_markers=st.booleans(), - table_output_format=st.sampled_from(("html", "json", "md", "jsonbbox", "dynamic", "csv")), - merge_tables=st.booleans(), - include=st.sampled_from( - ( - [], - ["hyperlinks"], - ["change_tracking", "highlight", "comments"], - ["signatures", "ignore_watermarks"], - ) - ), - ) - - -def _reducto_chunking_strategy() -> SearchStrategy[ReductoChunking]: - return st.one_of( - st.builds( - ReductoChunking, - chunk_mode=st.sampled_from(("section", "page", "disabled", "block", "page_sections")), - chunk_size=st.just(None), - chunk_overlap=st.just(0), - ), - st.builds( - ReductoChunking, - chunk_mode=st.just("variable"), - chunk_size=st.sampled_from((250, 1000, 1500)), - chunk_overlap=st.sampled_from((0, 32, 128)), - ), - ) - - -def _reducto_retrieval_strategy() -> SearchStrategy[ReductoRetrieval]: - return st.builds( - ReductoRetrieval, - chunking=_reducto_chunking_strategy(), - filter_blocks=st.sampled_from( - ( - [], - ["Header"], - ["Header", "Footer", "Page Number"], - ["Figure", "Table", "Key Value"], - ) - ), - embedding_optimized=st.booleans(), - ) - - -def _reducto_settings_strategy() -> SearchStrategy[ReductoSettings]: - return st.builds( - ReductoSettings, - ocr_system=st.sampled_from(("standard", "legacy")), - extraction_mode=st.sampled_from(("ocr", "hybrid")), - force_url_result=st.just(False), - return_ocr_data=st.booleans(), - return_images=st.sampled_from(([], ["figure"], ["table"], ["page"], ["figure", "table", "page"])), - embed_pdf_metadata=st.booleans(), - embed_pdf_metadata_dpi=st.sampled_from((50, 100, 250)), - persist_results=st.just(False), - timeout=st.sampled_from((None, 300.0, 900.0)), - page_range=st.sampled_from((None, [1], [1, 2], ["Sheet1"])), - ) - - -@st.composite -def reducto_v3_input_strategy( - draw: DrawFn, document: ReductoDocumentUrlDocument | None = None -) -> ReductoParseV3SdkInput: - model, custom_llm_provider = draw(st.sampled_from((("reducto/parse-v3", None), ("parse-v3", "reducto")))) - options: Final = draw( - st.fixed_dictionaries( - {}, - optional={ - "formatting": _reducto_formatting_strategy(), - "retrieval": _reducto_retrieval_strategy(), - "settings": _reducto_settings_strategy(), - }, - ) - ) - return ReductoParseV3SdkInput.model_validate( - { - "model": model, - "custom_llm_provider": custom_llm_provider, - "document": document - or ReductoDocumentUrlDocument(type="document_url", document_url="reducto://fixture-document.pdf"), - **options, - } - ) - - -def reducto_legacy_input_strategy( - document: ReductoDocumentUrlDocument | None = None, -) -> SearchStrategy[ReductoParseLegacySdkInput]: - selected_document: Final = document or ReductoDocumentUrlDocument( - type="document_url", document_url="reducto://fixture-document.pdf" - ) - return st.sampled_from( - ( - ReductoParseLegacySdkInput( - model="reducto/parse-legacy", - document=selected_document, - ), - ReductoParseLegacySdkInput( - model="parse-legacy", - custom_llm_provider="reducto", - document=selected_document, - ), - ReductoParseLegacySdkInput( - model="reducto/parse-legacy", - document=selected_document, - enhance={}, - ), - ) - ) - - -def _generate_examples( - target: OcrFixtureTarget, - root: Path, - examples: int, - concurrency: int, -) -> None: - generate_target_fixtures(target, root, examples, concurrency, OcrParityCase) - - -def _mistral_upstream_base(environ: Mapping[str, str]) -> str: - configured: Final = environ.get("MISTRAL_API_BASE", "https://api.mistral.ai").rstrip("/") - return configured.removesuffix("/v1") - - -def _mistral_target( - environ: Mapping[str, str], - sdk_call: Callable[..., object], -) -> OcrFixtureTarget | None: - api_key: Final = environ.get("MISTRAL_API_KEY") - if not api_key: - return None - - def invoke(api_base: str, case_input: OcrSdkInputBase) -> object: - return sdk_call(api_base=api_base, api_key=api_key, **case_input.as_sdk_kwargs()) - - return OcrFixtureTarget( - name="mistral-ocr", - provider_spec=ProviderSpec(upstream_base=_mistral_upstream_base(environ)), - strategy=cast(SearchStrategy[OcrSdkInputBase], mistral_input_strategy(_MISTRAL_MODEL)), - invoke=invoke, - required_inputs=cast(tuple[OcrSdkInputBase, ...], _required_mistral_inputs(_MISTRAL_MODEL)), - ) - - -def _azure_mistral_target( - environ: Mapping[str, str], sdk_call: Callable[..., object] -) -> OcrFixtureTarget | None: - api_key: Final = environ.get("AZURE_AI_API_KEY") - upstream_base: Final = environ.get("AZURE_AI_API_BASE") - configured_model: Final = environ.get("AZURE_AI_OCR_MODEL") - if not api_key or not upstream_base or not configured_model: - return None - model: Final = configured_model if configured_model.startswith("azure_ai/") else f"azure_ai/{configured_model}" - - def invoke(api_base: str, case_input: OcrSdkInputBase) -> object: - return sdk_call(api_base=api_base, api_key=api_key, **case_input.as_sdk_kwargs()) - - return OcrFixtureTarget( - name="azure-mistral", - provider_spec=ProviderSpec(upstream_base=upstream_base.rstrip("/")), - strategy=cast( - SearchStrategy[OcrSdkInputBase], - mistral_input_strategy(_MISTRAL_MODEL).map(lambda case_input: _as_azure_mistral(case_input, model)), - ), - invoke=invoke, - required_inputs=cast( - tuple[OcrSdkInputBase, ...], - tuple(_as_azure_mistral(case_input, model) for case_input in _required_mistral_inputs(_MISTRAL_MODEL)), - ), - ) - - -def _azure_document_intelligence_target( - environ: Mapping[str, str], sdk_call: Callable[..., object] -) -> OcrFixtureTarget | None: - api_key: Final = environ.get("AZURE_DOCUMENT_INTELLIGENCE_API_KEY") - upstream_base: Final = environ.get("AZURE_DOCUMENT_INTELLIGENCE_ENDPOINT") - if not api_key or not upstream_base: - return None - - def invoke(api_base: str, case_input: OcrSdkInputBase) -> object: - return sdk_call(api_base=api_base, api_key=api_key, **case_input.as_sdk_kwargs()) - - return OcrFixtureTarget( - name="azure-document-intelligence", - provider_spec=ProviderSpec(upstream_base=upstream_base.rstrip("/")), - strategy=cast(SearchStrategy[OcrSdkInputBase], azure_document_intelligence_input_strategy()), - invoke=invoke, - required_inputs=cast( - tuple[OcrSdkInputBase, ...], _required_azure_document_intelligence_inputs() - ), - ) - - -def _vertex_targets( - environ: Mapping[str, str], sdk_call: Callable[..., object] -) -> tuple[OcrFixtureTarget, ...]: - api_key: Final = environ.get("VERTEX_AI_API_KEY") - project: Final = environ.get("VERTEXAI_PROJECT") or environ.get("VERTEX_PROJECT") - location: Final = environ.get("VERTEXAI_LOCATION") or environ.get("VERTEX_LOCATION") or "us-central1" - if not api_key or not project: - return () - upstream_base: Final = environ.get("VERTEX_AI_API_BASE") or f"https://{location}-aiplatform.googleapis.com" - - def invoke(api_base: str, case_input: OcrSdkInputBase) -> object: - return sdk_call(api_base=api_base, api_key=api_key, **case_input.as_sdk_kwargs()) - - return ( - OcrFixtureTarget( - name="vertex-mistral", - provider_spec=ProviderSpec(upstream_base=upstream_base.rstrip("/")), - strategy=cast( - SearchStrategy[OcrSdkInputBase], - st.builds( - _as_vertex_mistral, - case_input=mistral_input_strategy(_MISTRAL_MODEL), - project=st.just(project), - location=st.just(location), - ), - ), - invoke=invoke, - required_inputs=cast( - tuple[OcrSdkInputBase, ...], - tuple( - _as_vertex_mistral(case_input, project, location) - for case_input in _required_mistral_inputs(_MISTRAL_MODEL) - ), - ), - ), - OcrFixtureTarget( - name="vertex-deepseek", - provider_spec=ProviderSpec(upstream_base=upstream_base.rstrip("/")), - strategy=cast( - SearchStrategy[OcrSdkInputBase], vertex_deepseek_input_strategy(project, location) - ), - invoke=invoke, - required_inputs=cast( - tuple[OcrSdkInputBase, ...], _required_vertex_deepseek_inputs(project, location) - ), - ), - ) - - -def _reducto_targets( - environ: Mapping[str, str], sdk_call: Callable[..., object] -) -> tuple[OcrFixtureTarget, ...]: - api_key: Final = environ.get("REDUCTO_API_KEY") - if not api_key: - return () - upstream_base: Final = environ.get("REDUCTO_API_BASE", _REDUCTO_API_BASE).rstrip("/") - document: Final = ReductoDocumentUrlDocument(type="document_url", document_url=_fixture_pdf_data_uri()) - - def invoke(api_base: str, case_input: OcrSdkInputBase) -> object: - return sdk_call(api_base=api_base, api_key=api_key, **case_input.as_sdk_kwargs()) - - return ( - OcrFixtureTarget( - name="reducto-v3", - provider_spec=ProviderSpec(upstream_base=upstream_base), - strategy=cast(SearchStrategy[OcrSdkInputBase], reducto_v3_input_strategy(document)), - invoke=invoke, - required_inputs=cast(tuple[OcrSdkInputBase, ...], _required_reducto_v3_inputs(document)), - ), - OcrFixtureTarget( - name="reducto-legacy", - provider_spec=ProviderSpec(upstream_base=upstream_base), - strategy=cast(SearchStrategy[OcrSdkInputBase], reducto_legacy_input_strategy(document)), - invoke=invoke, - required_inputs=cast( - tuple[OcrSdkInputBase, ...], - ( - ReductoParseLegacySdkInput(model="reducto/parse-legacy", document=document), - ReductoParseLegacySdkInput(model="reducto/parse-legacy", document=document, enhance={}), - ), - ), - ), - ) - - -def discover_targets( - environ: Mapping[str, str], - sdk_call: Callable[..., object], -) -> tuple[OcrFixtureTarget, ...]: - optional_targets: Final = ( - _mistral_target(environ, sdk_call), - _azure_mistral_target(environ, sdk_call), - _azure_document_intelligence_target(environ, sdk_call), - ) - direct_targets: Final = tuple(target for target in optional_targets if target is not None) - return (*direct_targets, *_vertex_targets(environ, sdk_call), *_reducto_targets(environ, sdk_call)) - - -def require_targets(targets: tuple[OcrFixtureTarget, ...]) -> tuple[OcrFixtureTarget, ...]: - return require_fixture_targets( - targets, - "No OCR fixture providers are configured. Set a supported provider API key and endpoint", - ) - - -def main() -> None: - logging.basicConfig(level=logging.INFO, format="%(message)s") - load_dotenv() - args: Final = parse_generator_args() - sdk_call: Final = cast(Callable[..., object], litellm.ocr) - targets: Final = require_targets(discover_targets(os.environ, sdk_call)) - root: Final = fixture_directory( - args.fixture_dir, - os.environ.get(FIXTURE_DIR_ENV), - Path(__file__).with_name("fixtures"), - ) - use_litellm_rust(False, ocr=None, aocr=None) - for target in targets: - _generate_examples(target, root, args.examples, args.concurrency) - - -if __name__ == "__main__": - main() diff --git a/tests/test_litellm/ocr/test_fixture_models.py b/tests/test_litellm/ocr/test_fixture_models.py index e44dd5c573c..fc563ed657d 100644 --- a/tests/test_litellm/ocr/test_fixture_models.py +++ b/tests/test_litellm/ocr/test_fixture_models.py @@ -11,7 +11,14 @@ from litellm.llms.base_llm.ocr.transformation import BaseOCRConfig from litellm.llms.mistral.ocr.transformation import MistralOCRConfig from litellm.llms.reducto.ocr.transformation import ReductoParseLegacyConfig, ReductoParseV3Config from litellm.llms.vertex_ai.ocr.deepseek_transformation import VertexAIDeepSeekOCRConfig -from tests.test_litellm.ocr.fixture_models import ( +from tests.test_litellm.ocr.fixtures.generate import ( + azure_document_intelligence_input_strategy, + mistral_input_strategy, + reducto_legacy_input_strategy, + reducto_v3_input_strategy, + vertex_deepseek_input_strategy, +) +from tests.test_litellm.ocr.fixtures.models import ( AzureDocumentIntelligenceOcrSdkInput, AzureMistralOcrSdkInput, JsonSchemaDefinition, @@ -32,13 +39,6 @@ from tests.test_litellm.ocr.fixture_models import ( VertexDeepSeekOcrSdkInput, VertexMistralOcrSdkInput, ) -from tests.test_litellm.ocr.generate_fixtures import ( - azure_document_intelligence_input_strategy, - mistral_input_strategy, - reducto_legacy_input_strategy, - reducto_v3_input_strategy, - vertex_deepseek_input_strategy, -) COMMON_FIELDS: Final = frozenset( {"boundary", "model", "document", "custom_llm_provider", "vertex_project", "vertex_location"} diff --git a/tests/test_litellm/ocr/test_generate_fixtures.py b/tests/test_litellm/ocr/test_generate_fixtures.py index 83f867e86ca..e2e2360b480 100644 --- a/tests/test_litellm/ocr/test_generate_fixtures.py +++ b/tests/test_litellm/ocr/test_generate_fixtures.py @@ -7,7 +7,7 @@ from typing import Final import pytest from tests.route_parity.fixture_recorder import generate_case_inputs -from tests.test_litellm.ocr.generate_fixtures import ( +from tests.test_litellm.ocr.fixtures.generate import ( discover_targets, parse_generator_args, require_targets, @@ -80,9 +80,9 @@ def test_azure_mistral_discovery_requires_and_normalizes_deployment_model() -> N } assert discover_targets(incomplete, _unused_sdk_call) == () - target: Final = discover_targets( - {**incomplete, "AZURE_AI_OCR_MODEL": "mistral-ocr-deployment"}, _unused_sdk_call - )[0] + target: Final = discover_targets({**incomplete, "AZURE_AI_OCR_MODEL": "mistral-ocr-deployment"}, _unused_sdk_call)[ + 0 + ] assert target.required_inputs[0].model == "azure_ai/mistral-ocr-deployment" diff --git a/tests/test_litellm/ocr/test_sdk_parity.py b/tests/test_litellm/ocr/test_sdk_parity.py index 2fb8ca1ad43..1c6539da40a 100644 --- a/tests/test_litellm/ocr/test_sdk_parity.py +++ b/tests/test_litellm/ocr/test_sdk_parity.py @@ -38,7 +38,7 @@ from tests.route_parity.runner import ( parity_worker_main, run_execution, ) -from tests.test_litellm.ocr.fixture_models import OcrParityCase, OcrSdkInput +from tests.test_litellm.ocr.fixtures.models import OcrParityCase, OcrSdkInput API_KEY: Final = "test-key" PYTHON_HTTP_SENTINEL: Final = "python-ocr-parity-fallback" @@ -318,7 +318,7 @@ def sdk_workers() -> Generator[tuple[PythonScriptWorker, PythonScriptWorker]]: @pytest.fixture(scope="module") def startup_ocr_fixture() -> OcrParityCase: - default_directory: Final = Path(__file__).with_name("fixtures") + default_directory: Final = Path(__file__).with_name("fixtures") / "data" configured: Final = os.environ.get(FIXTURE_DIR_ENV) directory: Final = Path(configured).expanduser() if configured is not None else default_directory fixtures: Final = recorded_fixtures(directory, OcrParityCase)