mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-19 00:01:29 +00:00
319 lines
14 KiB
Python
319 lines
14 KiB
Python
from __future__ import annotations
|
|
|
|
import importlib
|
|
import os
|
|
from pathlib import Path
|
|
from types import SimpleNamespace
|
|
from typing import Final, cast
|
|
|
|
import pytest
|
|
|
|
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, mapping
|
|
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 SdkCall, collect_trace, execute_trace
|
|
|
|
|
|
def _fixture(_engine: Engine, _base_url: str) -> RouteFixture:
|
|
return RouteFixture(kwargs={}, provider_responses=())
|
|
|
|
|
|
def _case(*, surface: Surface = "sdk", function: SdkFunction = "ocr") -> HarnessCase:
|
|
return HarnessCase(
|
|
strategy_id="trace_parity",
|
|
strategy_label="Trace parity",
|
|
sdk_function=function,
|
|
spec=ModuleCaseSpec(coverage=Coverage.PARTIAL, module="example"),
|
|
surface=surface,
|
|
)
|
|
|
|
|
|
def test_scenario_filtering_and_occurrence_node_ids() -> None:
|
|
suite: Final = TraceSuite(
|
|
route=RouteSpec("ocr", ("ocr", "aocr"), ("ocr", "aocr"), _fixture),
|
|
scenarios=(
|
|
TraceScenario("sync-one", _fixture, (), asynchronous=False),
|
|
TraceScenario("async-one", _fixture, (), asynchronous=True),
|
|
TraceScenario("async-two", _fixture, (), asynchronous=True),
|
|
),
|
|
)
|
|
case: Final = _case()
|
|
|
|
nodes: Final = scenario_nodeids(suite, case, frozenset({"async-two"}))
|
|
|
|
assert tuple(nodeid for _, nodeid in nodes) == ("trace:sdk:ocr:async-two",)
|
|
|
|
|
|
def test_python_engine_is_separate_from_scenario_selection() -> None:
|
|
assert runner_selection(("mistral", "--engine=python")) == (frozenset({"mistral"}), "python")
|
|
|
|
|
|
def test_python_engine_skips_native_bridge(monkeypatch: pytest.MonkeyPatch, tmp_path: Path) -> None:
|
|
runner: Final = importlib.import_module("tests.rust-python-harness.strategies.trace_parity.runner")
|
|
case: Final = _case()
|
|
selected: list[tuple[frozenset[str], str]] = []
|
|
|
|
def reject_bridge(_repo_root: Path) -> str | None:
|
|
raise AssertionError("Python-only tracing must not inspect or build the native bridge")
|
|
|
|
def capture_case(
|
|
_run: HarnessRun,
|
|
_case: HarnessCase,
|
|
scenarios: frozenset[str],
|
|
_on_update: object,
|
|
engine: str,
|
|
) -> None:
|
|
selected.append((scenarios, engine))
|
|
|
|
monkeypatch.setattr(runner, "ensure_trace_bridge", reject_bridge)
|
|
monkeypatch.setattr(runner, "_run_case", capture_case)
|
|
|
|
exit_code, _ = run_trace_cases((case,), tmp_path, lambda _: None, ("mistral", "--engine=python"))
|
|
|
|
assert exit_code == 0
|
|
assert selected == [(frozenset({"mistral"}), "python")]
|
|
|
|
|
|
def test_python_trace_preserves_native_ocr_dispatch_setting(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.setattr(execution, "_collect", collect)
|
|
|
|
monkeypatch.setenv("LITELLM_RUST", "0")
|
|
collect_trace(route, "python", asynchronous=False)
|
|
monkeypatch.setenv("LITELLM_RUST", "1")
|
|
collect_trace(route, "python", asynchronous=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:
|
|
loaded: Final = importlib.import_module("tests.rust-python-harness.strategies.trace_parity.sdk.responses.case")
|
|
suite: Final = cast(TraceSuite, loaded.TRACE_SUITE)
|
|
scenario: Final = next(item for item in suite.scenarios if item.name == "async-openai-provider-error")
|
|
monkeypatch.setattr(litellm, "suppress_debug_info", False)
|
|
assert isinstance(suite.route, RouteSpec)
|
|
|
|
result: Final = execute_trace(suite.route, scenario, "sdk", engine="python")
|
|
|
|
assert result.python_error is None
|
|
assert "Give Feedback / Get Help" not in capsys.readouterr().out
|
|
assert litellm.suppress_debug_info is False
|
|
|
|
|
|
@pytest.mark.parametrize("asynchronous", (False, True))
|
|
def test_vertex_trace_keeps_unmapped_helpers_and_parents(asynchronous: bool) -> None:
|
|
loaded: Final = importlib.import_module("tests.rust-python-harness.strategies.trace_parity.sdk.ocr.case")
|
|
suite: Final = cast(TraceSuite, loaded.TRACE_SUITE)
|
|
name: Final = f"{'async' if asynchronous else 'sync'}-vertex-deepseek"
|
|
scenario: Final = next(item for item in suite.scenarios if item.name == name)
|
|
assert isinstance(suite.route, RouteSpec)
|
|
|
|
trace: Final = execute_trace(suite.route, scenario, "sdk", engine="python")
|
|
|
|
assert trace.python_error is None
|
|
url: Final = next(
|
|
event for event in trace.python if event.raw.endswith(" VertexAIDeepSeekOCRConfig.get_complete_url")
|
|
)
|
|
project: Final = next(
|
|
event for event in trace.python if event.raw.endswith(" VertexBase.safe_get_vertex_ai_project")
|
|
)
|
|
location: Final = next(
|
|
event for event in trace.python if event.raw.endswith(" VertexBase.safe_get_vertex_ai_location")
|
|
)
|
|
assert project.parent_id == location.parent_id == url.id
|
|
assert not any(event.raw.endswith(" VertexBase.get_access_token") for event in trace.python)
|
|
|
|
|
|
@pytest.mark.parametrize("asynchronous", (False, True))
|
|
def test_vertex_credentials_trace_runs_real_auth_helpers(asynchronous: bool, monkeypatch: pytest.MonkeyPatch) -> None:
|
|
loaded: Final = importlib.import_module("tests.rust-python-harness.strategies.trace_parity.sdk.ocr.case")
|
|
suite: Final = cast(TraceSuite, loaded.TRACE_SUITE)
|
|
name: Final = f"{'async' if asynchronous else 'sync'}-vertex-deepseek-credentials"
|
|
scenario: Final = next(item for item in suite.scenarios if item.name == name)
|
|
monkeypatch.setenv("VERTEXAI_CREDENTIALS", "original-credentials")
|
|
monkeypatch.setenv("VERTEX_AI_API_KEY", "original-api-key")
|
|
assert isinstance(suite.route, RouteSpec)
|
|
|
|
trace: Final = execute_trace(suite.route, scenario, "sdk", engine="python")
|
|
|
|
assert trace.python_error is None
|
|
validate: Final = next(
|
|
event for event in trace.python if event.raw.endswith(" VertexAIDeepSeekOCRConfig.validate_environment")
|
|
)
|
|
helpers: Final = (
|
|
"VertexBase.safe_get_vertex_ai_project",
|
|
"VertexBase.safe_get_vertex_ai_credentials",
|
|
"VertexBase.get_access_token",
|
|
)
|
|
assert tuple(event.raw.split(" ", 1)[1] for event in trace.python if event.parent_id == validate.id) == helpers
|
|
token: Final = next(event for event in trace.python if event.raw.endswith(" VertexBase.get_access_token"))
|
|
load: Final = next(event for event in trace.python if event.raw.endswith(" VertexBase.load_auth"))
|
|
refresh: Final = next(event for event in trace.python if event.raw.endswith(" VertexBase.refresh_auth"))
|
|
assert load.parent_id == token.id
|
|
assert refresh.parent_id == load.id
|
|
assert os.environ["VERTEXAI_CREDENTIALS"] == "original-credentials"
|
|
assert os.environ["VERTEX_AI_API_KEY"] == "original-api-key"
|
|
|
|
|
|
def test_gateway_trace_keeps_calls_outside_scenario_mappings(monkeypatch: pytest.MonkeyPatch) -> None:
|
|
execution: Final = importlib.import_module("tests.rust-python-harness.strategies.trace_parity.gateway.execution")
|
|
events: Final = (
|
|
FunctionTraceEvent(0, None, "route.py:1 entry"),
|
|
FunctionTraceEvent(1, 0, "auth.py:2 authenticate"),
|
|
FunctionTraceEvent(2, 1, "auth.py:3 credentials"),
|
|
)
|
|
scenario: Final = TraceScenario(
|
|
"async-gateway",
|
|
_fixture,
|
|
(mapping(rust_span="entry", python_frame=r" entry$"),),
|
|
asynchronous=True,
|
|
)
|
|
monkeypatch.setattr(execution, "_collect", lambda *_args: events)
|
|
|
|
trace: Final = execution.execute_gateway_trace(GatewayRouteSpec("messages"), scenario, engine="python")
|
|
|
|
assert trace.python_error is None
|
|
assert tuple((event.id, event.parent_id, event.raw) for event in trace.python) == tuple(
|
|
(event.id, event.parent_id, event.raw) for event in events
|
|
)
|
|
|
|
|
|
def test_default_trace_skips_unavailable_rust_sdk_entrypoint(monkeypatch: pytest.MonkeyPatch) -> None:
|
|
execution: Final = importlib.import_module("tests.rust-python-harness.strategies.trace_parity.sdk.execution")
|
|
route: Final = RouteSpec("responses", ("responses", "aresponses"), None, _fixture)
|
|
scenario: Final = TraceScenario("sync-openai", _fixture, (), asynchronous=False)
|
|
engines: list[Engine] = []
|
|
|
|
def collect(_route: RouteSpec, engine: Engine, *, asynchronous: bool) -> tuple[FunctionTraceEvent, ...]:
|
|
engines.append(engine)
|
|
return (FunctionTraceEvent(0, None, "responses"),)
|
|
|
|
monkeypatch.setattr(execution, "collect_trace", collect)
|
|
|
|
trace: Final = execution.execute_trace(route, scenario, "sdk")
|
|
|
|
assert engines == ["python"]
|
|
assert trace.engine == "python"
|
|
assert trace.rust_error is None
|
|
|
|
|
|
def test_default_trace_skips_unavailable_rust_gateway_route(monkeypatch: pytest.MonkeyPatch) -> None:
|
|
execution: Final = importlib.import_module("tests.rust-python-harness.strategies.trace_parity.gateway.execution")
|
|
route: Final = GatewayRouteSpec("responses", rust_supported=False)
|
|
scenario: Final = TraceScenario("async-openai", _fixture, (), asynchronous=True)
|
|
engines: list[Engine] = []
|
|
|
|
def collect(_route: GatewayRouteSpec, _scenario: TraceScenario, engine: Engine) -> tuple[FunctionTraceEvent, ...]:
|
|
engines.append(engine)
|
|
return (FunctionTraceEvent(0, None, "responses"),)
|
|
|
|
monkeypatch.setattr(execution, "_collect", collect)
|
|
|
|
trace: Final = execution.execute_gateway_trace(route, scenario)
|
|
|
|
assert engines == ["python"]
|
|
assert trace.engine == "python"
|
|
assert trace.rust_error is None
|
|
|
|
|
|
def test_scenario_validation_rejects_duplicate_and_unsafe_names() -> None:
|
|
route: Final = RouteSpec("ocr", ("ocr", "aocr"), ("ocr", "aocr"), _fixture)
|
|
duplicate: Final = TraceSuite(
|
|
route=route,
|
|
scenarios=(
|
|
TraceScenario("sync-same", _fixture, (), asynchronous=False),
|
|
TraceScenario("sync-same", _fixture, (), asynchronous=False),
|
|
),
|
|
)
|
|
unsafe: Final = TraceSuite(
|
|
route=route, scenarios=(TraceScenario("sync-bad:name", _fixture, (), asynchronous=False),)
|
|
)
|
|
case: Final = _case()
|
|
|
|
assert validate_trace_suite(duplicate, case) is not None
|
|
assert validate_trace_suite(unsafe, case) is not None
|
|
|
|
|
|
def test_scenario_validation_rejects_invalid_names_and_route_registration() -> None:
|
|
invalid_name: Final = TraceSuite(
|
|
route=RouteSpec("ocr", ("ocr", "aocr"), ("ocr", "aocr"), _fixture),
|
|
scenarios=(TraceScenario("bedrock", _fixture, (), asynchronous=True),),
|
|
)
|
|
wrong_function: Final = TraceSuite(
|
|
route=RouteSpec("messages", ("create", "acreate"), ("messages", "amessages"), _fixture),
|
|
scenarios=(TraceScenario("sync-one", _fixture, (), asynchronous=False),),
|
|
)
|
|
wrong_surface: Final = TraceSuite(
|
|
route=GatewayRouteSpec("ocr"),
|
|
scenarios=(TraceScenario("sync-one", _fixture, (), asynchronous=False),),
|
|
)
|
|
case: Final = _case()
|
|
|
|
assert "start with sync- or async-" in (validate_trace_suite(invalid_name, case) or "")
|
|
assert "does not match case function" in (validate_trace_suite(wrong_function, case) or "")
|
|
assert "must use RouteSpec" in (validate_trace_suite(wrong_surface, case) or "")
|
|
|
|
|
|
def test_invalid_route_dispatch_records_harness_error() -> None:
|
|
case: Final = _case()
|
|
run: Final = HarnessRun.from_cases((case,))
|
|
result: Final = run.results[case.key]
|
|
suite: Final = TraceSuite(
|
|
route=GatewayRouteSpec("ocr"),
|
|
scenarios=(TraceScenario("sync-one", _fixture, (), asynchronous=False),),
|
|
)
|
|
nodeid: Final = "trace:sdk:ocr:sync-one"
|
|
|
|
run_trace_scenario(run, result, suite, suite.scenarios[0], "sdk", nodeid, lambda _: None)
|
|
|
|
assert result.outcomes[nodeid] is RunStatus.ERROR
|
|
assert run.failures == [(nodeid, "gateway route cannot run on the sdk surface")]
|
|
|
|
|
|
def test_different_python_and_rust_traces_pass(monkeypatch: pytest.MonkeyPatch) -> None:
|
|
runner: Final = importlib.import_module("tests.rust-python-harness.strategies.trace_parity.runner")
|
|
case: Final = _case()
|
|
run: Final = HarnessRun.from_cases((case,))
|
|
result: Final = run.results[case.key]
|
|
suite: Final = TraceSuite(
|
|
route=RouteSpec("ocr", ("ocr", "aocr"), ("ocr", "aocr"), _fixture),
|
|
scenarios=(TraceScenario("sync-one", _fixture, (), asynchronous=False),),
|
|
)
|
|
trace: Final = TraceArtifact.from_traces(
|
|
surface="sdk",
|
|
sdk_function="ocr",
|
|
scenario="sync-one",
|
|
python=(PipelineStep(0, None, "python_step", "python.py:1 python_step"),),
|
|
rust=(PipelineStep(0, None, "rust_step", "rust_step"),),
|
|
)
|
|
monkeypatch.setattr(runner, "_execute_scenario", lambda *_args: trace)
|
|
|
|
run_trace_scenario(run, result, suite, suite.scenarios[0], "sdk", "trace:sdk:ocr:sync-one", lambda _: None)
|
|
|
|
assert result.outcomes["trace:sdk:ocr:sync-one"] is RunStatus.PASSED
|
|
assert run.failures == []
|