mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-15 09:12:18 +00:00
127 lines
4.2 KiB
Python
127 lines
4.2 KiB
Python
"""The credentials are `os.environ/` references the gateway resolves, the same names the Claude Code matrix's
|
|
`test_config.yaml` uses, and Bedrock rides the gateway's default AWS identity."""
|
|
|
|
from __future__ import annotations
|
|
|
|
from collections.abc import Mapping
|
|
from dataclasses import dataclass, replace
|
|
from typing import Final, Literal
|
|
|
|
import pytest
|
|
from _pytest.mark.structures import ParameterSet
|
|
from cli_harness.rate_limiter import RateBucket
|
|
from e2e_metadata import Domain, Mode, Provider, Route, Subject, meta
|
|
from models import LiteLLMParamsBody
|
|
|
|
GPT_ON_BEDROCK_MANTLE: Final = "bedrock_mantle/openai.gpt-6.1-sol"
|
|
GPT_ON_BEDROCK_RUNTIME: Final = "bedrock/us.openai.gpt-6.1-sol"
|
|
GPT_ON_AZURE_OPENAI: Final = "azure/gpt-6.1-sol"
|
|
GPT_ON_OPENAI: Final = "openai/gpt-6.1-sol"
|
|
CLAUDE_ON_BEDROCK: Final = "bedrock/us.anthropic.claude-sonnet-5-5"
|
|
BEDROCK_REGION: Final = "us-east-1"
|
|
|
|
type SetupId = Literal["bedrock_mantle", "bedrock_runtime", "azure_openai", "openai", "claude_bedrock"]
|
|
|
|
|
|
@dataclass(frozen=True, slots=True)
|
|
class Deployment:
|
|
litellm_params: LiteLLMParamsBody
|
|
rate_bucket: RateBucket
|
|
|
|
|
|
@dataclass(frozen=True, slots=True)
|
|
class ModelSetup:
|
|
model: str
|
|
rate_bucket: RateBucket
|
|
|
|
|
|
def _drives(provider: Provider, model: str) -> pytest.MarkDecorator:
|
|
return meta(
|
|
Subject(
|
|
domain=Domain.LLM_TRANSLATION,
|
|
route=Route.RESPONSES,
|
|
providers=(provider,),
|
|
models=(model,),
|
|
mode=Mode.STREAM,
|
|
)
|
|
)
|
|
|
|
|
|
SETUPS: Final = (
|
|
pytest.param(
|
|
Deployment(
|
|
LiteLLMParamsBody(model=GPT_ON_BEDROCK_MANTLE, aws_region_name=BEDROCK_REGION), RateBucket.BEDROCK_MANTLE
|
|
),
|
|
id="bedrock_mantle",
|
|
marks=_drives(Provider.BEDROCK_MANTLE, GPT_ON_BEDROCK_MANTLE),
|
|
),
|
|
pytest.param(
|
|
Deployment(
|
|
LiteLLMParamsBody(model=GPT_ON_BEDROCK_RUNTIME, aws_region_name=BEDROCK_REGION),
|
|
RateBucket.BEDROCK_CONVERSE,
|
|
),
|
|
id="bedrock_runtime",
|
|
marks=_drives(Provider.BEDROCK, GPT_ON_BEDROCK_RUNTIME),
|
|
),
|
|
pytest.param(
|
|
Deployment(
|
|
LiteLLMParamsBody(
|
|
model=GPT_ON_AZURE_OPENAI, api_base="os.environ/AZURE_API_BASE", api_key="os.environ/AZURE_API_KEY"
|
|
),
|
|
RateBucket.AZURE_OPENAI,
|
|
),
|
|
id="azure_openai",
|
|
marks=_drives(Provider.AZURE, GPT_ON_AZURE_OPENAI),
|
|
),
|
|
pytest.param(
|
|
Deployment(LiteLLMParamsBody(model=GPT_ON_OPENAI, api_key="os.environ/OPENAI_API_KEY"), RateBucket.OPENAI),
|
|
id="openai",
|
|
marks=_drives(Provider.OPENAI, GPT_ON_OPENAI),
|
|
),
|
|
pytest.param(
|
|
Deployment(
|
|
LiteLLMParamsBody(model=CLAUDE_ON_BEDROCK, aws_region_name=BEDROCK_REGION), RateBucket.BEDROCK_CONVERSE
|
|
),
|
|
id="claude_bedrock",
|
|
marks=_drives(Provider.BEDROCK, CLAUDE_ON_BEDROCK),
|
|
),
|
|
)
|
|
|
|
|
|
def with_expected_failure(setup_id: SetupId, *, reason: str) -> tuple[ParameterSet, ...]:
|
|
return tuple(
|
|
pytest.param(*setup.values, id=setup.id, marks=(*setup.marks, pytest.mark.xfail(strict=True, reason=reason)))
|
|
if setup.id == setup_id
|
|
else setup
|
|
for setup in SETUPS
|
|
)
|
|
|
|
|
|
@dataclass(frozen=True, slots=True)
|
|
class ExtraModel:
|
|
provider: Provider
|
|
litellm_params: LiteLLMParamsBody
|
|
|
|
|
|
def also_driving(setups: tuple[ParameterSet, ...], extra: Mapping[str, ExtraModel]) -> tuple[ParameterSet, ...]:
|
|
return tuple(_also_driving(setup, extra) for setup in setups)
|
|
|
|
|
|
def _also_driving(setup: ParameterSet, extra: Mapping[str, ExtraModel]) -> ParameterSet:
|
|
deployment: Final = setup.values[0]
|
|
assert isinstance(deployment, Deployment)
|
|
added: Final = extra[deployment.litellm_params.model]
|
|
return pytest.param(*setup.values, id=setup.id, marks=tuple(_widened(mark, added) for mark in setup.marks))
|
|
|
|
|
|
def _widened(mark: pytest.MarkDecorator | pytest.Mark, added: ExtraModel) -> pytest.MarkDecorator | pytest.Mark:
|
|
subject: Final = mark.args[0] if mark.name == "meta" and mark.args else None
|
|
if not isinstance(subject, Subject):
|
|
return mark
|
|
return meta(
|
|
replace(
|
|
subject,
|
|
providers=(*subject.providers, added.provider),
|
|
models=(*subject.models, added.litellm_params.model),
|
|
)
|
|
)
|