litellm/tests/rust-python-harness/strategies/trace_parity/sdk/ocr/case.py
2026-09-14 12:17:15 -07:00

509 lines
20 KiB
Python

from __future__ import annotations
import json
from typing import Final
from .....shared.parity.recorded_http import HttpHeader, RecordedHttpResponse
from .....shared.tracing.steps import Engine, mapping
from ...models import RouteFixture, RouteSpec, TraceScenario, TraceSuite
COMMON_MAPPINGS: Final = (
mapping(rust_span="ocr", python_frame=r"ocr/main\.py:\d+ a?ocr$"),
mapping(rust_span="prepare_ocr_call", python_frame=r"ocr/main\.py:\d+ _prepare_ocr_request$"),
mapping(rust_span="ocr_provider_config", python_frame=r"ProviderConfigManager\.get_provider_ocr_config$"),
mapping(rust_span="supported_ocr_params", python_frame=r"get_supported_ocr_params$"),
mapping(rust_span="map_ocr_params", python_frame=r"(?<!async_)map_ocr_params$"),
mapping(rust_span="validate_environment", python_frame=r"(?<!_)validate_environment$"),
mapping(rust_span="complete_url", python_frame=r"get_complete_url$"),
mapping(rust_span="transform_ocr_request", python_frame=r"(?<!async_)transform_ocr_request$"),
mapping(rust_span="http_request", python_frame=r"AsyncHTTPHandler\.post$|HTTPHandler\.post$"),
)
SUCCESS_CALLBACK_SYNC_MAPPING: Final = mapping(
rust_span="success_callback",
python_frame=r"BoundedLoggingThreadPoolExecutor\.submit$",
)
SUCCESS_CALLBACK_ASYNC_MAPPING: Final = mapping(
rust_span="success_callback",
python_frame=r"Logging\.async_success_handler$",
)
FAILURE_CALLBACK_MAPPING: Final = mapping(
rust_span="failure_callback",
python_frame=r"Logging\.(?:async_)?failure_handler$",
)
IGNORED_SUCCESS_CALLBACK_MAPPING: Final = mapping(rust_span="success_callback")
SYNC_MAPPINGS: Final = (
*COMMON_MAPPINGS,
mapping(rust_span="execute_ocr_provider_call", python_frame=r"BaseLLMHTTPHandler\.ocr$"),
mapping(span="python_transform_ocr_response_wrapper", python_frame=r"BaseLLMHTTPHandler\._transform_ocr_response$"),
mapping(
rust_span="transform_ocr_response",
python_frame=r"MistralOCRConfig\.transform_ocr_response$",
),
)
ASYNC_MAPPINGS: Final = (
*COMMON_MAPPINGS,
mapping(span="python_ocr_wrapper", python_frame=r"BaseLLMHTTPHandler\.ocr$"),
mapping(rust_span="execute_ocr_provider_call", python_frame=r"BaseLLMHTTPHandler\.async_ocr$"),
mapping(
rust_span="transform_ocr_response",
python_frame=r"MistralOCRConfig\.transform_ocr_response$",
),
)
PUBLIC_RUST_DISPATCH_MAPPINGS: Final = (
mapping(span="public_sdk_entrypoint", python_frame=r"ocr/main\.py:\d+ a?ocr$"),
mapping(span="public_request", python_frame=r"ocr/main\.py:\d+ _public_request$"),
mapping(span="bind_request", python_frame=r"ocr/main\.py:\d+ _bind_request$"),
mapping(span="rust_ocr_enabled", python_frame=r"rust_bridge/configuration\.py:\d+ rust_ocr_enabled$"),
mapping(span="select_native_ocr", python_frame=r"rust_bridge/ocr_lifecycle\.py:\d+ select$"),
mapping(span="load_native_bridge", python_frame=r"rust_bridge/bindings\.py:\d+ NativeBinding\.load$"),
mapping(span="native_call_setup", python_frame=r"rust_bridge/lifecycle\.py:\d+ setup$"),
mapping(span="native_response", python_frame=r"rust_bridge/ocr\.py:\d+ _response$"),
mapping(span="native_call_finalize", python_frame=r"rust_bridge/lifecycle\.py:\d+ finalize$"),
mapping(
span="native_success_bookkeeping",
python_frame=r"rust_bridge/lifecycle\.py:\d+ success_bookkeeping$",
),
*(mapping(rust_span=item.rust) for item in SYNC_MAPPINGS if item.rust is not None),
)
CALLBACK_SUCCESS_SYNC_MAPPINGS: Final = (*SYNC_MAPPINGS, SUCCESS_CALLBACK_SYNC_MAPPING)
CALLBACK_SUCCESS_ASYNC_MAPPINGS: Final = (*ASYNC_MAPPINGS, SUCCESS_CALLBACK_ASYNC_MAPPING)
CALLBACK_FAILURE_SYNC_MAPPINGS: Final = (
*COMMON_MAPPINGS,
mapping(rust_span="execute_ocr_provider_call", python_frame=r"BaseLLMHTTPHandler\.ocr$"),
FAILURE_CALLBACK_MAPPING,
)
CALLBACK_FAILURE_ASYNC_MAPPINGS: Final = (
*COMMON_MAPPINGS,
mapping(span="python_ocr_wrapper", python_frame=r"BaseLLMHTTPHandler\.ocr$"),
mapping(rust_span="execute_ocr_provider_call", python_frame=r"BaseLLMHTTPHandler\.async_ocr$"),
FAILURE_CALLBACK_MAPPING,
)
AZURE_COMMON_MAPPINGS: Final = (
*COMMON_MAPPINGS[:7],
mapping(
rust_span="transform_ocr_request",
python_frame=(
r"AzureAIOCRConfig\.(?:async_)?transform_ocr_request$"
r"|MistralOCRConfig\.transform_ocr_request$"
),
),
COMMON_MAPPINGS[-1],
)
AZURE_SYNC_MAPPINGS: Final = (
*AZURE_COMMON_MAPPINGS,
mapping(rust_span="execute_ocr_provider_call", python_frame=r"BaseLLMHTTPHandler\.ocr$"),
mapping(
span="python_transform_ocr_response_wrapper",
python_frame=r"BaseLLMHTTPHandler\._transform_ocr_response$",
),
mapping(
rust_span="transform_ocr_response",
python_frame=r"MistralOCRConfig\.transform_ocr_response$",
),
)
AZURE_ASYNC_MAPPINGS: Final = (
*AZURE_COMMON_MAPPINGS,
mapping(span="python_ocr_wrapper", python_frame=r"BaseLLMHTTPHandler\.ocr$"),
mapping(rust_span="execute_ocr_provider_call", python_frame=r"BaseLLMHTTPHandler\.async_ocr$"),
mapping(
rust_span="transform_ocr_response",
python_frame=r"MistralOCRConfig\.transform_ocr_response$",
),
)
def _fixture(engine: Engine, model: str, document: dict[str, str] | None = None) -> RouteFixture:
response: Final = json.dumps(
{
"pages": [{"index": 0, "markdown": "hello"}],
"model": "mistral-ocr-latest",
"usage_info": {"pages_processed": 1},
}
).encode()
return RouteFixture(
kwargs={
"model": model,
"document": document or {"type": "document_url", "document_url": "https://example.com/document.pdf"},
**({"optional_params": {"pages": [0]}} if engine == "rust" else {"pages": [0]}),
},
provider_responses=(
RecordedHttpResponse.from_bytes(
200, (HttpHeader(name="content-type", value="application/json"),), response
),
),
)
def _mistral_fixture(engine: Engine, _base_url: str) -> RouteFixture:
return _fixture(engine, "mistral/mistral-ocr-latest")
def _callback_fixture(engine: Engine, *, failure: bool) -> RouteFixture:
fixture: Final = _fixture(engine, "mistral/mistral-ocr-latest")
provider_responses: Final = (
(
RecordedHttpResponse.from_bytes(
400,
(HttpHeader(name="content-type", value="application/json"),),
b'{"message":"trace callback provider failure"}',
),
)
if failure
else fixture.provider_responses
)
return RouteFixture(
kwargs=fixture.kwargs,
provider_responses=provider_responses,
expected_failure=failure,
)
def _mistral_callback_success_fixture(engine: Engine, _base_url: str) -> RouteFixture:
return _callback_fixture(engine, failure=False)
def _mistral_callback_failure_fixture(engine: Engine, _base_url: str) -> RouteFixture:
return _callback_fixture(engine, failure=True)
def _azure_fixture(engine: Engine, _base_url: str) -> RouteFixture:
return _fixture(
engine,
"azure_ai/pixtral-12b-2409",
{"type": "image_url", "image_url": "data:image/png;base64,aGVsbG8="},
)
def _vertex_deepseek_fixture(engine: Engine, _base_url: str) -> RouteFixture:
vertex: Final = {"vertex_project": "trace-project", "vertex_location": "us-central1"}
return RouteFixture(
kwargs={
"model": "vertex_ai/deepseek-ocr-maas",
"document": {"type": "image_url", "image_url": "data:image/png;base64,aGVsbG8="},
**({"optional_params": vertex} if engine == "rust" else vertex),
},
provider_responses=(
RecordedHttpResponse.from_bytes(
200,
(HttpHeader(name="content-type", value="application/json"),),
json.dumps(
{
"choices": [{"message": {"role": "assistant", "content": "hello"}}],
"usage": {"prompt_tokens": 1, "completion_tokens": 1},
}
).encode(),
),
),
)
def _vertex_deepseek_credentials_fixture(engine: Engine, base_url: str) -> RouteFixture:
from cryptography.hazmat.primitives import serialization
from cryptography.hazmat.primitives.asymmetric import rsa
fixture: Final = _vertex_deepseek_fixture(engine, base_url)
private_key: Final = rsa.generate_private_key(public_exponent=65537, key_size=2048)
credentials: Final = json.dumps(
{
"type": "service_account",
"project_id": "trace-project",
"private_key_id": "trace-key",
"private_key": private_key.private_bytes(
serialization.Encoding.PEM,
serialization.PrivateFormat.PKCS8,
serialization.NoEncryption(),
).decode(),
"client_email": "trace@trace-project.iam.gserviceaccount.com",
"token_uri": f"{base_url}/token",
}
)
return RouteFixture(
kwargs={**fixture.kwargs, "api_key": None},
environment=(("VERTEXAI_CREDENTIALS", credentials), ("VERTEX_AI_API_KEY", "")),
provider_responses=(
RecordedHttpResponse.from_bytes(
200,
(HttpHeader(name="content-type", value="application/json"),),
b'{"access_token":"trace-token","token_type":"Bearer","expires_in":3600}',
),
*fixture.provider_responses,
),
)
def _cohere_fixture(engine: Engine, _base_url: str) -> RouteFixture:
return RouteFixture(
kwargs={
"model": "cohere/parse-v5.0",
"document": {"type": "image_url", "image_url": "data:image/png;base64,aGVsbG8="},
**({"optional_params": {"output_format": "blocks"}} if engine == "rust" else {"output_format": "blocks"}),
},
provider_responses=(
RecordedHttpResponse.from_bytes(
200,
(HttpHeader(name="content-type", value="application/json"),),
json.dumps(
{
"pages": [{"index": 0, "blocks": [{"type": "text", "text": {"content": "hello"}}]}],
"meta": {"billed_units": {"pages": 1}},
}
).encode(),
),
),
)
def _azure_document_intelligence_fixture(engine: Engine, base_url: str) -> RouteFixture:
completed: Final = json.dumps(
{
"status": "succeeded",
"analyzeResult": {
"content": "hello",
"pages": [
{
"pageNumber": 1,
"width": 8.5,
"height": 11,
"unit": "inch",
"lines": [{"content": "hello"}],
}
],
},
}
).encode()
return RouteFixture(
kwargs={
"model": "azure_ai/doc-intelligence/prebuilt-read",
"document": {
"type": "document_url",
"document_url": "data:application/pdf;base64,aGVsbG8=",
},
**({"optional_params": {"pages": [0]}} if engine == "rust" else {"pages": [0]}),
},
provider_responses=(
RecordedHttpResponse.from_bytes(
202,
(
HttpHeader(name="content-type", value="application/json"),
HttpHeader(name="operation-location", value=f"{base_url}/operations/trace"),
),
b"{}",
),
RecordedHttpResponse.from_bytes(
200,
(HttpHeader(name="content-type", value="application/json"),),
completed,
),
),
)
DEEPSEEK_COMMON_MAPPINGS: Final = (
mapping(rust_span="ocr", python_frame=r"ocr/main\.py:\d+ a?ocr$"),
mapping(rust_span="prepare_ocr_call", python_frame=r"ocr/main\.py:\d+ _prepare_ocr_request$"),
mapping(rust_span="ocr_provider_config", python_frame=r"ProviderConfigManager\.get_provider_ocr_config$"),
mapping(rust_span="supported_ocr_params", python_frame=r"get_supported_ocr_params$"),
mapping(rust_span="map_ocr_params", python_frame=r"(?<!async_)map_ocr_params$"),
mapping(rust_span="validate_environment", python_frame=r"(?<!_)validate_environment$"),
mapping(rust_span="complete_url", python_frame=r"get_complete_url$"),
mapping(
rust_span="transform_ocr_request",
python_frame=r"VertexAIDeepSeekOCRConfig\.transform_ocr_request$",
),
mapping(rust_span="http_request", python_frame=r"AsyncHTTPHandler\.post$|HTTPHandler\.post$"),
mapping(
rust_span="transform_ocr_response",
python_frame=r"VertexAIDeepSeekOCRConfig\.transform_ocr_response$",
),
)
DEEPSEEK_SYNC_MAPPINGS: Final = (
*DEEPSEEK_COMMON_MAPPINGS,
mapping(rust_span="execute_ocr_provider_call", python_frame=r"BaseLLMHTTPHandler\.ocr$"),
mapping(span="python_transform_ocr_response_wrapper", python_frame=r"BaseLLMHTTPHandler\._transform_ocr_response$"),
)
DEEPSEEK_ASYNC_MAPPINGS: Final = (
*DEEPSEEK_COMMON_MAPPINGS,
mapping(span="python_ocr_wrapper", python_frame=r"BaseLLMHTTPHandler\.ocr$"),
mapping(
span="python_async_transform_ocr_request",
python_frame=r"VertexAIDeepSeekOCRConfig\.async_transform_ocr_request$",
),
mapping(rust_span="execute_ocr_provider_call", python_frame=r"BaseLLMHTTPHandler\.async_ocr$"),
)
DOCUMENT_INTELLIGENCE_COMMON_MAPPINGS: Final = (
mapping(rust_span="ocr", python_frame=r"ocr/main\.py:\d+ a?ocr$"),
mapping(rust_span="prepare_ocr_call", python_frame=r"ocr/main\.py:\d+ _prepare_ocr_request$"),
mapping(rust_span="ocr_provider_config", python_frame=r"ProviderConfigManager\.get_provider_ocr_config$"),
mapping(
rust_span="supported_ocr_params", python_frame=r"AzureDocumentIntelligenceOCRConfig\.get_supported_ocr_params$"
),
mapping(rust_span="map_ocr_params", python_frame=r"AzureDocumentIntelligenceOCRConfig\.map_ocr_params$"),
mapping(
rust_span="validate_environment", python_frame=r"AzureDocumentIntelligenceOCRConfig\.validate_environment$"
),
mapping(rust_span="complete_url", python_frame=r"AzureDocumentIntelligenceOCRConfig\.get_complete_url$"),
mapping(
rust_span="transform_ocr_request", python_frame=r"AzureDocumentIntelligenceOCRConfig\.transform_ocr_request$"
),
mapping(rust_span="http_request", python_frame=r"AsyncHTTPHandler\.post$|HTTPHandler\.post$"),
mapping(
rust_span="poll_document_intelligence",
python_frame=r"AzureDocumentIntelligenceOCRConfig\._poll_operation_(?:sync|async)$",
),
mapping(
rust_span="transform_ocr_response",
python_frame=r"AzureDocumentIntelligenceOCRConfig\._transform_completed_response$",
),
)
DOCUMENT_INTELLIGENCE_SYNC_MAPPINGS: Final = (
*DOCUMENT_INTELLIGENCE_COMMON_MAPPINGS,
mapping(rust_span="execute_ocr_provider_call", python_frame=r"BaseLLMHTTPHandler\.ocr$"),
mapping(span="python_transform_ocr_response_wrapper", python_frame=r"BaseLLMHTTPHandler\._transform_ocr_response$"),
mapping(
span="python_provider_transform_response",
python_frame=r"AzureDocumentIntelligenceOCRConfig\.transform_ocr_response$",
),
mapping(span="python_poll_http_request", python_frame=r"HTTPHandler\.get$"),
)
DOCUMENT_INTELLIGENCE_ASYNC_MAPPINGS: Final = (
*DOCUMENT_INTELLIGENCE_COMMON_MAPPINGS,
mapping(span="python_ocr_wrapper", python_frame=r"BaseLLMHTTPHandler\.ocr$"),
mapping(rust_span="execute_ocr_provider_call", python_frame=r"BaseLLMHTTPHandler\.async_ocr$"),
mapping(
span="python_provider_transform_response",
python_frame=r"AzureDocumentIntelligenceOCRConfig\.async_transform_ocr_response$",
),
mapping(span="python_poll_http_request", python_frame=r"AsyncHTTPHandler\.get$"),
)
COHERE_COMMON_MAPPINGS: Final = (
*COMMON_MAPPINGS[:7],
mapping(
rust_span="transform_ocr_request",
python_frame=r"CohereParseConfig\.(?:async_)?transform_ocr_request$",
),
COMMON_MAPPINGS[-1],
mapping(rust_span="transform_ocr_response", python_frame=r"CohereParseConfig\.transform_ocr_response$"),
)
COHERE_ASYNC_MAPPINGS: Final = (
*COHERE_COMMON_MAPPINGS,
mapping(span="python_ocr_wrapper", python_frame=r"BaseLLMHTTPHandler\.ocr$"),
mapping(rust_span="execute_ocr_provider_call", python_frame=r"BaseLLMHTTPHandler\.async_ocr$"),
)
SPEC: Final = RouteSpec("ocr", ("ocr", "aocr"), ("ocr", "aocr"), _mistral_fixture)
TRACE_SUITE: Final = TraceSuite(
route=SPEC,
scenarios=(
TraceScenario(
name="sync-mistral",
fixture=_mistral_fixture,
mappings=(*SYNC_MAPPINGS, IGNORED_SUCCESS_CALLBACK_MAPPING),
asynchronous=False,
),
TraceScenario(
name="async-mistral",
fixture=_mistral_fixture,
mappings=(*ASYNC_MAPPINGS, IGNORED_SUCCESS_CALLBACK_MAPPING),
asynchronous=True,
),
TraceScenario(
name="sync-mistral-callback-success",
fixture=_mistral_callback_success_fixture,
mappings=CALLBACK_SUCCESS_SYNC_MAPPINGS,
asynchronous=False,
),
TraceScenario(
name="async-mistral-callback-success",
fixture=_mistral_callback_success_fixture,
mappings=CALLBACK_SUCCESS_ASYNC_MAPPINGS,
asynchronous=True,
),
TraceScenario(
name="sync-mistral-callback-failure",
fixture=_mistral_callback_failure_fixture,
mappings=CALLBACK_FAILURE_SYNC_MAPPINGS,
asynchronous=False,
),
TraceScenario(
name="async-mistral-callback-failure",
fixture=_mistral_callback_failure_fixture,
mappings=CALLBACK_FAILURE_ASYNC_MAPPINGS,
asynchronous=True,
),
TraceScenario(
name="sync-azure-ai",
fixture=_azure_fixture,
mappings=(*AZURE_SYNC_MAPPINGS, IGNORED_SUCCESS_CALLBACK_MAPPING),
asynchronous=False,
),
TraceScenario(
name="async-azure-ai",
fixture=_azure_fixture,
mappings=(*AZURE_ASYNC_MAPPINGS, IGNORED_SUCCESS_CALLBACK_MAPPING),
asynchronous=True,
),
TraceScenario(
name="sync-azure-document-intelligence",
fixture=_azure_document_intelligence_fixture,
mappings=(*DOCUMENT_INTELLIGENCE_SYNC_MAPPINGS, IGNORED_SUCCESS_CALLBACK_MAPPING),
asynchronous=False,
),
TraceScenario(
name="async-azure-document-intelligence",
fixture=_azure_document_intelligence_fixture,
mappings=(*DOCUMENT_INTELLIGENCE_ASYNC_MAPPINGS, IGNORED_SUCCESS_CALLBACK_MAPPING),
asynchronous=True,
),
TraceScenario(
name="sync-vertex-deepseek",
fixture=_vertex_deepseek_fixture,
mappings=(*DEEPSEEK_SYNC_MAPPINGS, IGNORED_SUCCESS_CALLBACK_MAPPING),
asynchronous=False,
),
TraceScenario(
name="async-vertex-deepseek",
fixture=_vertex_deepseek_fixture,
mappings=(*DEEPSEEK_ASYNC_MAPPINGS, IGNORED_SUCCESS_CALLBACK_MAPPING),
asynchronous=True,
),
TraceScenario(
name="sync-vertex-deepseek-credentials",
fixture=_vertex_deepseek_credentials_fixture,
mappings=DEEPSEEK_SYNC_MAPPINGS,
asynchronous=False,
),
TraceScenario(
name="async-vertex-deepseek-credentials",
fixture=_vertex_deepseek_credentials_fixture,
mappings=DEEPSEEK_ASYNC_MAPPINGS,
asynchronous=True,
),
TraceScenario(
name="async-cohere",
fixture=_cohere_fixture,
mappings=(*COHERE_ASYNC_MAPPINGS, IGNORED_SUCCESS_CALLBACK_MAPPING),
asynchronous=True,
),
TraceScenario(
name="sync-public-rust-dispatch",
fixture=_mistral_fixture,
mappings=PUBLIC_RUST_DISPATCH_MAPPINGS,
asynchronous=False,
),
TraceScenario(
name="async-public-rust-dispatch",
fixture=_mistral_fixture,
mappings=PUBLIC_RUST_DISPATCH_MAPPINGS,
asynchronous=True,
),
),
)