diff --git a/tests/rust-python-harness/strategies/trace_parity/models.py b/tests/rust-python-harness/strategies/trace_parity/models.py index 078af2b316d..be7375dc604 100644 --- a/tests/rust-python-harness/strategies/trace_parity/models.py +++ b/tests/rust-python-harness/strategies/trace_parity/models.py @@ -64,6 +64,7 @@ class TraceScenario: fixture: Callable[[Engine, str], RouteFixture] mappings: tuple[TraceMapping, ...] asynchronous: bool + python_rust_enabled: bool = False @dataclass(frozen=True, slots=True) diff --git a/tests/rust-python-harness/strategies/trace_parity/sdk/execution.py b/tests/rust-python-harness/strategies/trace_parity/sdk/execution.py index c249d09fa73..f8a674451b8 100644 --- a/tests/rust-python-harness/strategies/trace_parity/sdk/execution.py +++ b/tests/rust-python-harness/strategies/trace_parity/sdk/execution.py @@ -1,10 +1,13 @@ from __future__ import annotations import asyncio +import os from collections.abc import AsyncIterable, Awaitable, Iterable +from contextlib import nullcontext from dataclasses import dataclass from pathlib import Path from typing import Final, Protocol, cast +from unittest.mock import patch from ....shared.parity.replay import replay_server from ....shared.reporting.models import Surface @@ -105,7 +108,7 @@ def _collect( def collect_trace( - spec: RouteSpec, engine: Engine, *, asynchronous: bool + spec: RouteSpec, engine: Engine, *, asynchronous: bool, python_rust_enabled: bool = False ) -> tuple[FunctionTraceEvent, ...] | TraceExecutionFailure: function: Final = _entrypoint(spec, engine, asynchronous=asynchronous) if isinstance(function, TraceExecutionFailure): @@ -126,7 +129,13 @@ def collect_trace( expected_failure=base_fixture.expected_failure, consume_stream=base_fixture.consume_stream, ) - collected: Final = _collect(function, fixture, engine, asynchronous=asynchronous) + environment: Final = ( + patch.dict(os.environ, {"LITELLM_RUST": "1" if python_rust_enabled else "0"}) + if engine == "python" + else nullcontext() + ) + with environment: + collected: Final = _collect(function, fixture, engine, asynchronous=asynchronous) provider.take_requests(len(fixture.provider_responses)) except Exception as error: return TraceExecutionFailure(engine, f"{type(error).__name__}: {error}") @@ -159,7 +168,14 @@ def execute_trace( fixture=scenario.fixture, ) python_trace: Final = ( - collect_trace(scenario_route, "python", asynchronous=scenario.asynchronous) if engine != "rust" else () + collect_trace( + scenario_route, + "python", + asynchronous=scenario.asynchronous, + python_rust_enabled=scenario.python_rust_enabled, + ) + if engine != "rust" + else () ) rust_trace: Final = ( collect_trace(scenario_route, "rust", asynchronous=scenario.asynchronous) if engine != "python" else () diff --git a/tests/rust-python-harness/strategies/trace_parity/sdk/ocr/case.py b/tests/rust-python-harness/strategies/trace_parity/sdk/ocr/case.py index effd1a0b4f6..b2c0e4c7acc 100644 --- a/tests/rust-python-harness/strategies/trace_parity/sdk/ocr/case.py +++ b/tests/rust-python-harness/strategies/trace_parity/sdk/ocr/case.py @@ -1,7 +1,7 @@ from __future__ import annotations import json -from typing import Final, cast +from typing import Final from .....shared.parity.recorded_http import HttpHeader, RecordedHttpResponse from .....shared.tracing.steps import Engine, mapping @@ -53,6 +53,23 @@ ASYNC_MAPPINGS: Final = ( ), ) +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 = ( @@ -164,23 +181,6 @@ def _azure_fixture(engine: Engine, _base_url: str) -> RouteFixture: ) -def _vertex_fixture(engine: Engine, _base_url: str) -> RouteFixture: - fixture: Final = _fixture( - engine, - "vertex_ai/mistral-ocr-maas", - {"type": "image_url", "image_url": "data:image/png;base64,aGVsbG8="}, - ) - vertex: Final = {"vertex_project": "trace-project", "vertex_location": "us-central1"} - optional_params: Final = cast(dict[str, object], fixture.kwargs.get("optional_params", {})) - return RouteFixture( - kwargs={ - **fixture.kwargs, - **({"optional_params": {**optional_params, **vertex}} if engine == "rust" else vertex), - }, - provider_responses=fixture.provider_responses, - ) - - def _vertex_deepseek_fixture(engine: Engine, _base_url: str) -> RouteFixture: vertex: Final = {"vertex_project": "trace-project", "vertex_location": "us-central1"} return RouteFixture( @@ -204,6 +204,28 @@ def _vertex_deepseek_fixture(engine: Engine, _base_url: str) -> RouteFixture: ) +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( { @@ -249,30 +271,6 @@ def _azure_document_intelligence_fixture(engine: Engine, base_url: str) -> Route ) -VERTEX_COMMON_MAPPINGS: Final = ( - *COMMON_MAPPINGS[:7], - mapping( - rust_span="transform_ocr_request", - python_frame=( - r"VertexAIOCRConfig\.(?:async_)?transform_ocr_request$" - r"|MistralOCRConfig\.transform_ocr_request$" - ), - ), - COMMON_MAPPINGS[-1], -) -VERTEX_SYNC_MAPPINGS: Final = ( - *VERTEX_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$"), -) -VERTEX_ASYNC_MAPPINGS: Final = ( - *VERTEX_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$"), -) - 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$"), @@ -352,6 +350,20 @@ DOCUMENT_INTELLIGENCE_ASYNC_MAPPINGS: Final = ( 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( @@ -417,18 +429,6 @@ TRACE_SUITE: Final = TraceSuite( mappings=(*DOCUMENT_INTELLIGENCE_ASYNC_MAPPINGS, IGNORED_SUCCESS_CALLBACK_MAPPING), asynchronous=True, ), - TraceScenario( - name="sync-vertex-ai", - fixture=_vertex_fixture, - mappings=(*VERTEX_SYNC_MAPPINGS, IGNORED_SUCCESS_CALLBACK_MAPPING), - asynchronous=False, - ), - TraceScenario( - name="async-vertex-ai", - fixture=_vertex_fixture, - mappings=(*VERTEX_ASYNC_MAPPINGS, IGNORED_SUCCESS_CALLBACK_MAPPING), - asynchronous=True, - ), TraceScenario( name="sync-vertex-deepseek", fixture=_vertex_deepseek_fixture, @@ -441,5 +441,25 @@ TRACE_SUITE: Final = TraceSuite( mappings=(*DEEPSEEK_ASYNC_MAPPINGS, IGNORED_SUCCESS_CALLBACK_MAPPING), 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, + python_rust_enabled=True, + ), + TraceScenario( + name="async-public-rust-dispatch", + fixture=_mistral_fixture, + mappings=PUBLIC_RUST_DISPATCH_MAPPINGS, + asynchronous=True, + python_rust_enabled=True, + ), ), ) diff --git a/tests/rust-python-harness/strategies/trace_parity/sdk/test_core_scenario_matrix.py b/tests/rust-python-harness/strategies/trace_parity/sdk/test_core_scenario_matrix.py index d0921398795..d0dbd281a97 100644 --- a/tests/rust-python-harness/strategies/trace_parity/sdk/test_core_scenario_matrix.py +++ b/tests/rust-python-harness/strategies/trace_parity/sdk/test_core_scenario_matrix.py @@ -16,6 +16,7 @@ def _suite(module: str) -> TraceSuite: def test_core_sdk_scenario_matrix_keeps_distinct_migration_paths() -> None: chat: Final = _suite("tests.rust-python-harness.strategies.trace_parity.sdk.chat_completions.case") messages: Final = _suite("tests.rust-python-harness.strategies.trace_parity.sdk.messages.case") + ocr: Final = _suite("tests.rust-python-harness.strategies.trace_parity.sdk.ocr.case") responses: Final = _suite("tests.rust-python-harness.strategies.trace_parity.sdk.responses.case") assert {(scenario.name, scenario.asynchronous) for scenario in chat.scenarios} >= { @@ -38,6 +39,11 @@ def test_core_sdk_scenario_matrix_keeps_distinct_migration_paths() -> None: ("async-bedrock-invalid-thinking-retry", True), ("sync-unsupported", False), } + assert {(scenario.name, scenario.asynchronous) for scenario in ocr.scenarios} >= { + ("async-cohere", True), + ("sync-public-rust-dispatch", False), + ("async-public-rust-dispatch", True), + } assert {(scenario.name, scenario.asynchronous) for scenario in responses.scenarios} >= { ("sync-openai", False), ("async-openai", True), diff --git a/tests/rust-python-harness/strategies/trace_parity/test_runner.py b/tests/rust-python-harness/strategies/trace_parity/test_runner.py index 4992854e21d..9dfd0d51c89 100644 --- a/tests/rust-python-harness/strategies/trace_parity/test_runner.py +++ b/tests/rust-python-harness/strategies/trace_parity/test_runner.py @@ -1,7 +1,9 @@ from __future__ import annotations import importlib +import os from pathlib import Path +from types import SimpleNamespace from typing import Final, cast import pytest @@ -10,11 +12,12 @@ import litellm from ...shared.reporting.models import Coverage, HarnessCase, HarnessRun, RunStatus, SdkFunction, Surface from ...shared.reporting.strategy import ModuleCaseSpec +from ...shared.tracing.profiler import FunctionTraceEvent from ...shared.tracing.steps import Engine, PipelineStep from .models import GatewayRouteSpec, RouteFixture, RouteSpec, TraceScenario, TraceSuite from .reporting import TraceArtifact from .runner import run_trace_cases, run_trace_scenario, runner_selection, scenario_nodeids, validate_trace_suite -from .sdk.execution import execute_trace +from .sdk.execution import SdkCall, collect_trace, execute_trace def _fixture(_engine: Engine, _base_url: str) -> RouteFixture: @@ -77,6 +80,34 @@ def test_python_engine_skips_native_bridge(monkeypatch: pytest.MonkeyPatch, tmp_ assert selected == [(frozenset({"mistral"}), "python")] +def test_python_trace_controls_native_ocr_dispatch(monkeypatch: pytest.MonkeyPatch) -> None: + execution: Final = importlib.import_module("tests.rust-python-harness.strategies.trace_parity.sdk.execution") + route: Final = RouteSpec("ocr", ("ocr", "aocr"), ("ocr", "aocr"), _fixture) + observed: list[str | None] = [] + + def collect( + _function: SdkCall, + _fixture: RouteFixture, + _engine: Engine, + *, + asynchronous: bool, + ) -> SimpleNamespace: + observed.append(os.environ.get("LITELLM_RUST")) + return SimpleNamespace( + events=(FunctionTraceEvent(0, None, "aocr" if asynchronous else "ocr"),), + error=None, + ) + + monkeypatch.setenv("LITELLM_RUST", "1") + monkeypatch.setattr(execution, "_collect", collect) + + collect_trace(route, "python", asynchronous=False) + collect_trace(route, "python", asynchronous=True, python_rust_enabled=True) + + assert observed == ["0", "1"] + assert os.environ["LITELLM_RUST"] == "1" + + def test_expected_provider_failure_omits_feedback_banner( capsys: pytest.CaptureFixture[str], monkeypatch: pytest.MonkeyPatch ) -> None: