litellm/tests/test_litellm_rust/support/requests.py
yujonglee 347b642bdd
refactor(ocr): complete native lifecycle and preserve Azure auth (#40734)
* refactor(ocr): extract call completion boundary

* fix(ocr): release completion state after dispatch

* test(ocr): prove wrapper completion handoff

* test(ocr): narrow mapped failure assertion

* fix(ocr): preserve wrapper invocation kwargs

* fix(ocr): retain completion through finalization

* fix(ocr): make completion ownership explicit

* refactor(ocr): resolve logging executor explicitly

* fix(callbacks): preserve completion lifecycle behavior

* refactor(ocr): move public OCR into native lifecycle

* refactor(ocr): remove unused rust bridge capability

* wip

* wip

* refactor

* wip

* fix(ocr): preserve reducto native compatibility

* wip

* fix(ocr): document native callable casts

* perf(ocr): bound responses and reduce native scheduling overhead

* refactor(python-bridge): organize placeholder routes

* refactor test

* fix(ocr): normalize DeepSeek document content

* perf(ocr): skip unused callback work and benchmark callback overhead

* fix(ocr): align conversion contracts

* test(ocr): cover official provider response shapes

* fix(ocr): restore Python fallback and honor Rust opt-out

* fixes and refactor

* fix(ocr): preserve Azure Document Intelligence authentication

* fix(rust): enforce OCR response limits and lint contracts

* test(rust): align native OCR contract coverage

* test(ocr): isolate Azure auth precedence coverage
2026-09-12 11:56:49 -07:00

58 lines
2 KiB
Python

from typing import Final
import litellm
from litellm.llms.base_llm.ocr.transformation import OCRResponse
from tests.test_litellm_rust.support.recording_server import RecordingServer
OCR_DOCUMENT: Final = {"type": "document_url", "document_url": "data:application/pdf;base64,YWJj"}
OCR_MODEL: Final = "mistral/mistral-ocr-latest"
OCR_RESPONSE: Final = {
"pages": [{"index": 0, "markdown": "native OCR response", "images": [], "dimensions": None}],
"model": "mistral-ocr-latest",
"usage_info": {"pages_processed": 1, "doc_size_bytes": 3},
}
def ocr_arguments(server: RecordingServer, **kwargs: object) -> dict[str, object]:
return {
"model": OCR_MODEL,
"document": dict(OCR_DOCUMENT),
"api_key": "test-key",
"api_base": server.base_url,
**kwargs,
}
def call_ocr(server: RecordingServer, **kwargs: object) -> OCRResponse:
response: Final = litellm.ocr(**ocr_arguments(server, **kwargs))
if not isinstance(response, OCRResponse):
raise TypeError(f"Expected OCRResponse, got {type(response).__name__}")
return response
async def call_aocr(server: RecordingServer, **kwargs: object) -> OCRResponse:
return await litellm.aocr(**ocr_arguments(server, **kwargs))
def call_native_ocr(server: RecordingServer, **kwargs: object) -> OCRResponse:
return call_ocr(server, **kwargs)
async def call_native_aocr(server: RecordingServer, **kwargs: object) -> OCRResponse:
return await call_aocr(server, **kwargs)
def request_body(kwargs: dict[str, object]) -> dict[str, object]:
additional_args = kwargs["additional_args"]
assert isinstance(additional_args, dict)
body = additional_args["complete_input_dict"]
assert isinstance(body, dict)
return body
def request_headers(kwargs: dict[str, object]) -> dict[str, object]:
additional_args = kwargs["additional_args"]
assert isinstance(additional_args, dict)
headers = additional_args["headers"]
assert isinstance(headers, dict)
return headers