diff --git a/tests/test_litellm/_fixture_recorder.py b/tests/test_litellm/_fixture_recorder.py index 849efd15040..94fecac85a1 100644 --- a/tests/test_litellm/_fixture_recorder.py +++ b/tests/test_litellm/_fixture_recorder.py @@ -5,6 +5,7 @@ import hashlib import queue import threading from collections.abc import Callable, Generator +from concurrent.futures import ThreadPoolExecutor from contextlib import contextmanager from dataclasses import dataclass from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer @@ -51,6 +52,7 @@ class RecorderResult: @dataclass(frozen=True, slots=True) class GeneratorArgs: + concurrency: int examples: int fixture_dir: Path | None model: str @@ -66,7 +68,7 @@ def _excluded_headers(headers: tuple[tuple[str, str], ...]) -> frozenset[str]: def _end_to_end_headers(headers: httpx.Headers) -> tuple[HttpHeader, ...]: decoded: Final = tuple((name.decode("ascii"), value.decode("latin-1")) for name, value in headers.raw) - excluded: Final = _excluded_headers(decoded) | {"content-length"} + excluded: Final = _excluded_headers(decoded) | {"content-encoding", "content-length"} return tuple(HttpHeader(name=name, value=value) for name, value in decoded if name.lower() not in excluded) @@ -110,7 +112,7 @@ class _RecordingHandler(BaseHTTPRequestHandler): content=request_body, timeout=120, ) as upstream: - response_body: Final = b"".join(upstream.iter_raw()) + response_body: Final = b"".join(upstream.iter_bytes()) recorded_response: Final = RecordedHttpResponse.from_bytes( status_code=upstream.status_code, headers=_end_to_end_headers(upstream.headers), @@ -172,18 +174,39 @@ def record_case( sdk_call(api_base=recorder.url, api_key=spec.api_key, **case_input.as_sdk_kwargs()) upstream_response: Final = recorder.take_response() - case: Final = OcrParityCase(input=case_input, upstream_response=upstream_response) + case: Final = OcrParityCase(litellm_input=case_input, provider_response=upstream_response) cache.put(cache_key, cast(dict[str, object], case.model_dump(mode="json", exclude_unset=True))) return RecorderResult(case=case, cache_hit=False) +def record_cases( + spec: ProviderSpec, + root: Path, + case_inputs: tuple[MistralOcrParityInput, ...], + sdk_call: Callable[..., object], + max_concurrency: int, +) -> tuple[RecorderResult, ...]: + if max_concurrency < 1: + raise ValueError("max_concurrency must be at least 1") + unique_inputs: Final = tuple( + {canonical_json(fixture_cache_key(case_input)): case_input for case_input in case_inputs}.values() + ) + with ThreadPoolExecutor(max_workers=max_concurrency) as executor: + futures: Final = tuple( + executor.submit(record_case, spec, root, case_input, sdk_call) for case_input in unique_inputs + ) + return tuple(future.result() for future in futures) + + def parse_generator_args() -> GeneratorArgs: parser: Final = argparse.ArgumentParser() + parser.add_argument("--concurrency", type=int, default=4) parser.add_argument("--examples", type=int, default=4) parser.add_argument("--fixture-dir", type=Path) parser.add_argument("--model", default="mistral/mistral-ocr-latest") namespace: Final = parser.parse_args() return GeneratorArgs( + concurrency=cast(int, namespace.concurrency), examples=cast(int, namespace.examples), fixture_dir=cast(Path | None, namespace.fixture_dir), model=cast(str, namespace.model), @@ -202,13 +225,13 @@ def recorded_fixtures(directory: Path) -> tuple[OcrParityCase, ...]: fixtures.append(OcrParityCase.model_validate(raw_fixture)) except ValidationError as error: raise ValueError( - f"invalid OCR parity fixture {path}: expected exactly `input` and `upstream_response` " + f"invalid OCR parity fixture {path}: expected exactly `litellm_input` and `provider_response` " f"({len(error.errors())} validation errors)" ) from error return tuple(fixtures) def fixture_id(fixture: OcrParityCase) -> str: - input_json: Final = canonical_json(fixture.input.canonical_input()) + input_json: Final = canonical_json(fixture.litellm_input.canonical_input()) digest: Final = hashlib.sha256(input_json.encode("utf-8")).hexdigest()[:8] - return f"mistral-{fixture.input.model.rsplit('/', 1)[-1]}-{digest}" + return f"mistral-{fixture.litellm_input.model.rsplit('/', 1)[-1]}-{digest}" diff --git a/tests/test_litellm/ocr/conftest.py b/tests/test_litellm/ocr/conftest.py index b5410325905..6a3947e2fdb 100644 --- a/tests/test_litellm/ocr/conftest.py +++ b/tests/test_litellm/ocr/conftest.py @@ -30,7 +30,7 @@ def pytest_generate_tests(metafunc: pytest.Metafunc) -> None: except (ValidationError, ValueError) as error: raise pytest.UsageError( f"Invalid OCR parity fixture bundle at {directory}. " - "Each fixture must contain exactly `input` and `upstream_response`. " + "Each fixture must contain exactly `litellm_input` and `provider_response`. " "Record fresh fixtures in an empty directory with: " f"`uv run python tests/test_litellm/ocr/generate_fixtures.py --fixture-dir {directory}`. " f"Validation details: {error}" diff --git a/tests/test_litellm/ocr/fixture_models.py b/tests/test_litellm/ocr/fixture_models.py index 2b06bf3392e..ab4432cc349 100644 --- a/tests/test_litellm/ocr/fixture_models.py +++ b/tests/test_litellm/ocr/fixture_models.py @@ -87,5 +87,5 @@ class RecordedHttpResponse(_FixtureModel): class OcrParityCase(_FixtureModel): - input: MistralOcrParityInput - upstream_response: RecordedHttpResponse + litellm_input: MistralOcrParityInput + provider_response: RecordedHttpResponse diff --git a/tests/test_litellm/ocr/generate_fixtures.py b/tests/test_litellm/ocr/generate_fixtures.py index 99f2ca05648..0cdd9237680 100644 --- a/tests/test_litellm/ocr/generate_fixtures.py +++ b/tests/test_litellm/ocr/generate_fixtures.py @@ -2,6 +2,7 @@ from __future__ import annotations import logging import os +import queue from collections.abc import Callable from pathlib import Path from typing import Final, cast @@ -18,13 +19,10 @@ from tests.test_litellm._fixture_recorder import ( ProviderSpec, fixture_directory, parse_generator_args, - record_case, + record_cases, ) from tests.test_litellm.ocr.fixture_models import ( - AnnotationFormat, - DocumentUrlDocument, ImageUrlDocument, - JsonSchemaDefinition, MistralOcrParityInput, ) @@ -32,16 +30,6 @@ FIXTURE_DIR_ENV: Final = "LITELLM_OCR_FIXTURE_DIR" LOGGER: Final = logging.getLogger(__name__) _TEXT: Final = st.from_regex(r"[A-Za-z0-9 ]{1,24}", fullmatch=True) _VALUE_TEXT: Final = st.text(alphabet="abcdefghijklmnopqrstuvwxyz0123456789 -_", min_size=1, max_size=32) -_ANNOTATION_FORMAT: Final[SearchStrategy[AnnotationFormat]] = st.builds( - AnnotationFormat, - type=st.just("json_schema"), - json_schema=st.builds( - JsonSchemaDefinition, - name=_VALUE_TEXT, - schema_value=st.just({"type": "object", "properties": {}, "additionalProperties": False}), - strict=st.one_of(st.none(), st.booleans()), - ), -) def _image_document(text: str, font_size: int) -> ImageUrlDocument: @@ -50,36 +38,18 @@ def _image_document(text: str, font_size: int) -> ImageUrlDocument: def _input_strategy(model: str) -> SearchStrategy[MistralOcrParityInput]: - document_strategy: Final = st.one_of( - st.builds(_image_document, _TEXT, st.integers(min_value=12, max_value=36)), - st.just( - DocumentUrlDocument( - type="document_url", - document_url="https://arxiv.org/pdf/2201.04234", - ) - ), - ) + document_strategy: Final = st.builds(_image_document, _TEXT, st.integers(min_value=12, max_value=36)) input_values: Final = st.fixed_dictionaries( {"model": st.just(model), "document": document_strategy}, optional={ - "pages": st.one_of( - st.none(), - st.lists(st.integers(min_value=0, max_value=20), min_size=1, max_size=5, unique=True), - ), - "include_image_base64": st.one_of(st.none(), st.booleans()), - "image_limit": st.one_of(st.none(), st.integers(min_value=1, max_value=100)), - "image_min_size": st.one_of(st.none(), st.integers(min_value=0, max_value=10_000)), - "bbox_annotation_format": st.one_of(st.none(), _ANNOTATION_FORMAT), - "document_annotation_format": st.one_of(st.none(), _ANNOTATION_FORMAT), - "document_annotation_prompt": st.one_of(st.none(), _VALUE_TEXT), - "extract_header": st.one_of(st.none(), st.booleans()), - "extract_footer": st.one_of(st.none(), st.booleans()), - "table_format": st.one_of(st.none(), st.sampled_from(("markdown", "html"))), - "confidence_scores_granularity": st.one_of( - st.none(), st.sampled_from(("word", "page", "block")) - ), - "include_blocks": st.one_of(st.none(), st.booleans()), - "id": st.one_of(st.none(), _VALUE_TEXT), + "include_image_base64": st.booleans(), + "image_limit": st.integers(min_value=1, max_value=100), + "image_min_size": st.integers(min_value=0, max_value=10_000), + "extract_header": st.booleans(), + "extract_footer": st.booleans(), + "table_format": st.sampled_from(("markdown", "html")), + "include_blocks": st.booleans(), + "id": _VALUE_TEXT, }, ) return input_values.map(MistralOcrParityInput.model_validate) @@ -89,15 +59,22 @@ def _generate_examples( spec: ProviderSpec, root: Path, examples: int, + concurrency: int, sdk_call: Callable[..., object], ) -> None: + generated: Final[queue.SimpleQueue[MistralOcrParityInput | None]] = queue.SimpleQueue() + @settings(max_examples=examples, deadline=None, derandomize=True) @given(case_input=_input_strategy(spec.model)) def generate_case(case_input: MistralOcrParityInput) -> None: - result: Final = record_case(spec, root, case_input, sdk_call) - LOGGER.info("%s %s", "cached" if result.cache_hit else "recorded", result.case.input.model) + generated.put(case_input) generate_case() + generated.put(None) + case_inputs: Final = tuple(iter(generated.get, None)) + results: Final = record_cases(spec, root, case_inputs, sdk_call, concurrency) + for result in results: + LOGGER.info("%s %s", "cached" if result.cache_hit else "recorded", result.case.litellm_input.model) def _mistral_upstream_base() -> str: @@ -119,7 +96,7 @@ def main() -> None: ) spec: Final = ProviderSpec(model=args.model, upstream_base=_mistral_upstream_base(), api_key=api_key) use_litellm_rust(False, ocr=None, aocr=None) - _generate_examples(spec, root, args.examples, cast(Callable[..., object], litellm.ocr)) + _generate_examples(spec, root, args.examples, args.concurrency, cast(Callable[..., object], litellm.ocr)) if __name__ == "__main__": diff --git a/tests/test_litellm/ocr/test_sdk_parity.py b/tests/test_litellm/ocr/test_sdk_parity.py index aa68cfd0b0c..0a463232afb 100644 --- a/tests/test_litellm/ocr/test_sdk_parity.py +++ b/tests/test_litellm/ocr/test_sdk_parity.py @@ -51,7 +51,7 @@ def _execute_sdk_case(sdk_input: MistralOcrParityInput, route: SDKRoute, mock_ur def test_recorded_ocr_sdk_parity(ocr_fixture: OcrParityCase, route: SDKRoute, tmp_path: Path) -> None: case_file: Final = tmp_path / f"{route.value}-ocr-parity-case.json" case_file.write_text(ocr_fixture.model_dump_json(indent=2, exclude_unset=True), encoding="utf-8") - response: Final = ocr_fixture.upstream_response + response: Final = ocr_fixture.provider_response response_body: Final = response.body_bytes() response_headers: Final = tuple((header.name, header.value) for header in response.headers) runner: Final = PythonScriptRunner( @@ -90,7 +90,7 @@ def _child_main() -> None: mock_url: Final = sys.argv[3] report_file: Final = Path(sys.argv[4]) case: Final = OcrParityCase.model_validate_json(case_file.read_text(encoding="utf-8")) - report: Final = _execute_sdk_case(case.input, route, mock_url) + report: Final = _execute_sdk_case(case.litellm_input, route, mock_url) report_file.write_text(report.model_dump_json(indent=2), encoding="utf-8") diff --git a/tests/test_litellm/test__fixture_recorder.py b/tests/test_litellm/test__fixture_recorder.py new file mode 100644 index 00000000000..f457f97ee31 --- /dev/null +++ b/tests/test_litellm/test__fixture_recorder.py @@ -0,0 +1,102 @@ +from __future__ import annotations + +import threading +from collections.abc import Generator +from contextlib import contextmanager +from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer +from pathlib import Path +from typing import Final, cast + +import httpx + +from tests.test_litellm._fixture_recorder import ProviderSpec, record_cases +from tests.test_litellm.ocr.fixture_models import ImageUrlDocument, MistralOcrParityInput + + +class _ControlledUpstream(ThreadingHTTPServer): + daemon_threads = True + + def __init__(self) -> None: + super().__init__(("127.0.0.1", 0), _ControlledUpstreamHandler) + self.lock: Final = threading.Lock() + self.two_requests_started: Final = threading.Event() + self.active_requests: int = 0 + self.max_active_requests: int = 0 + self.request_count: int = 0 + + @property + def url(self) -> str: + return f"http://127.0.0.1:{self.server_address[1]}" + + def start_request(self) -> None: + with self.lock: + self.active_requests += 1 + self.request_count += 1 + self.max_active_requests = max(self.max_active_requests, self.active_requests) + if self.active_requests == 2: + self.two_requests_started.set() + self.two_requests_started.wait(timeout=2) + + def end_tracked_request(self) -> None: + with self.lock: + self.active_requests -= 1 + + +class _ControlledUpstreamHandler(BaseHTTPRequestHandler): + protocol_version = "HTTP/1.1" + + def do_POST(self) -> None: + upstream: Final = self.server + assert isinstance(upstream, _ControlledUpstream) + length: Final = int(self.headers.get("content-length") or "0") + self.rfile.read(length) + upstream.start_request() + try: + body: Final = b"{}" + self.send_response(200) + self.send_header("content-type", "application/json") + self.send_header("content-length", str(len(body))) + self.end_headers() + self.wfile.write(body) + finally: + upstream.end_tracked_request() + + def log_message(self, format: str, *args: object) -> None: + return + + +@contextmanager +def _controlled_upstream() -> Generator[_ControlledUpstream]: + server: Final = _ControlledUpstream() + 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 _case(identifier: str) -> MistralOcrParityInput: + return MistralOcrParityInput( + model="mistral/mistral-ocr-latest", + document=ImageUrlDocument(type="image_url", image_url="https://example.com/image.png"), + id=identifier, + ) + + +def _sdk_call(**kwargs: object) -> object: + api_base: Final = cast(str, kwargs["api_base"]) + return httpx.post(f"{api_base}/v1/ocr", content=b"{}", timeout=5) + + +def test_record_cases_deduplicates_and_limits_concurrency(tmp_path: Path) -> None: + case_inputs: Final = (_case("one"), _case("two"), _case("one"), _case("three")) + with _controlled_upstream() as upstream: + spec: Final = ProviderSpec(model="mistral/mistral-ocr-latest", upstream_base=upstream.url, api_key="test-key") + results: Final = record_cases(spec, tmp_path, case_inputs, _sdk_call, max_concurrency=2) + + assert len(results) == 3 + assert upstream.request_count == 3 + assert upstream.max_active_requests == 2