mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-16 23:41:43 +00:00
fix(harness): trace OCR Python and native dispatch paths
This commit is contained in:
Yujong Lee
parent
d4869ba710
commit
209fc7afb0
5 changed files with 132 additions and 58 deletions
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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 ()
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
),
|
||||
),
|
||||
)
|
||||
|
|
|
|||
|
|
@ -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),
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue