litellm/tests/e2e/perf/setup/performance_test_case.py

"""The schema of a perf test: a `PerformanceTestCase` names what the gateway is deployed with
(model groups and guardrails, both written out in the gateway's config) and lists what the test
itself creates and sends: teams, Keycloak users, virtual keys, the request shapes, the traffic plan
and the `Limit`s that decide pass or fail. Add a case by adding a module to `cases/`.

The same shape as `CostTrackingTestCase` in the cost suite: frozen dataclasses with plain defaults,
so a case is a literal a reviewer can read top to bottom.
"""

from __future__ import annotations

from collections.abc import Mapping
from dataclasses import dataclass, field
from typing import Final, Literal

Endpoint = Literal["chat_completions", "responses", "messages"]
MockProvider = Literal["openai", "anthropic", "bedrock"]
GuardrailMode = Literal["pre_call", "post_call"]
Metric = Literal[
    "unexpected_status_count",
    "gateway_added_seconds",
    "time_to_first_token_seconds",
    "memory_growth_megabytes",
    "postgres_calls_per_request",
    "cpu_seconds_per_request",
    "redis_calls_per_request",
]


@dataclass(frozen=True, slots=True)
class Team:
    """One team (`/team/new`). `models` names the model groups the team may call; empty means all."""

    name: str
    models: tuple[str, ...] = ()
    max_budget: float | None = None
    rpm_limit: int | None = None
    tpm_limit: int | None = None


@dataclass(frozen=True, slots=True)
class User:
    """One Keycloak user who calls the gateway with a JWT. `team` puts them in that team's group."""

    name: str
    team: str | None = None
    max_budget: float | None = None


@dataclass(frozen=True, slots=True)
class VirtualKey:
    """One virtual key (`/key/generate`), scoped to a team and/or a user when named."""

    name: str
    team: str | None = None
    user: str | None = None
    models: tuple[str, ...] = ()
    rpm_limit: int | None = None
    max_budget: float | None = None


@dataclass(frozen=True, slots=True)
class Guardrail:
    """A guardrail the gateway's config defines (a `generic_guardrail_api` pointed at the mock
    guardrail, `default_on`); the case names it so the plan knows which hooks every request sees."""

    name: str
    mode: GuardrailMode


@dataclass(frozen=True, slots=True)
class RequestShape:
    """What one request looks like: a coding-agent session (`setup/corpus/session.py`, the system prompt, the
    Read/Edit/Write/Bash/Grep/Glob tools when `tools` is set, then file-reading turns carrying real source files)
    whose size in input tokens is drawn per request from the `input_tokens` range, and the output cap. The mocks
    answer with a reply generated for the request, sized from that cap, a tool call when `tools` is set."""

    name: str
    input_tokens: tuple[int, int]
    tools: bool = False
    output_tokens: int = 1_024


SHORT: Final = RequestShape("short", input_tokens=(5_000, 100_000), tools=True)
LONG: Final = RequestShape("long", input_tokens=(300_000, 1_000_000), tools=True)


@dataclass(frozen=True, slots=True)
class Traffic:
    """The plan: `total_requests` drawn from `seed`, sent at `rps` by `concurrency` simulated
    users. Shares are fractions of the requests: `endpoints` and `shapes` must each sum to 1,
    `intervene_share` picks non-streaming requests the guardrail rewrites, `block_share` picks
    requests it blocks."""

    seed: int
    rps: float
    total_requests: int
    endpoints: Mapping[Endpoint, float] = field(default_factory=lambda: {"chat_completions": 1.0})
    shapes: Mapping[RequestShape, float] = field(default_factory=lambda: {SHORT: 1.0})
    stream_share: float = 0.0
    intervene_share: float = 0.0
    block_share: float = 0.0
    concurrency: int = 20

    @property
    def expected_seconds(self) -> float:
        return self.total_requests / self.rps


@dataclass(frozen=True, slots=True)
class Limit:
    """One pass/fail rule: `metric` (at `percentile`, or the whole-run value) must not exceed `max`.
    Metrics without a Limit are still reported."""

    metric: Metric
    max: float
    percentile: int | None = None

    @property
    def label(self) -> str:
        return f"{self.metric} percentile {self.percentile}" if self.percentile is not None else self.metric


@dataclass(frozen=True, slots=True)
class MockLatency:
    """Seconds a mock provider waits before its first byte, plus or minus up to `jitter_seconds` per request
    (drawn from the request marker, so a rerun waits the same). Zero means the mocks answer at once."""

    seconds: float = 0.0
    jitter_seconds: float = 0.0


@dataclass(frozen=True, slots=True)
class PerformanceTestCase:
    """`model_groups` and `guardrails` name what the gateway config deploys (setup checks they are
    there); everything else the test creates or sends. `mock_latency` is how long every mock provider waits
    before answering, set on the mock pods at setup, so the provider side of each request is a known
    constant and the rest of the client's time is the gateway's."""

    name: str
    teams: tuple[Team, ...]
    users: tuple[User, ...]
    virtual_keys: tuple[VirtualKey, ...]
    model_groups: tuple[str, ...]
    guardrails: tuple[Guardrail, ...]
    traffic: Traffic
    limits: tuple[Limit, ...]
    mock_latency: MockLatency = field(default_factory=lambda: MockLatency())

    def team(self, name: str) -> Team:
        return next(team for team in self.teams if team.name == name)

    @property
    def has_post_guardrail(self) -> bool:
        return any(guardrail.mode == "post_call" for guardrail in self.guardrails)

    def validate(self) -> None:
        """Fail fast on a case that cannot run: dangling names, shares that do not add up, a model
        group no caller may use."""
        team_names: Final = {team.name for team in self.teams}
        user_names: Final = {user.name for user in self.users}
        group_names: Final = set(self.model_groups)
        for names, label in (
            ([team.name for team in self.teams], "team"),
            ([user.name for user in self.users], "user"),
            ([key.name for key in self.virtual_keys], "key"),
            (list(self.model_groups), "model group"),
            ([guardrail.name for guardrail in self.guardrails], "guardrail"),
        ):
            assert len(names) == len(set(names)), f"{self.name}: duplicate {label} names in {names}"
        for team in self.teams:
            assert set(team.models) <= group_names, f"{self.name}: team {team.name} names unknown models {team.models}"
        for user in self.users:
            assert user.team is None or user.team in team_names, (
                f"{self.name}: user {user.name} in unknown team {user.team}"
            )
        for key in self.virtual_keys:
            assert key.team is None or key.team in team_names, f"{self.name}: key {key.name} in unknown team {key.team}"
            assert key.user is None or key.user in user_names, (
                f"{self.name}: key {key.name} for unknown user {key.user}"
            )
            assert set(key.models) <= group_names, f"{self.name}: key {key.name} names unknown models {key.models}"
        assert abs(sum(self.traffic.endpoints.values()) - 1.0) < 1e-9, f"{self.name}: endpoint shares must sum to 1"
        assert abs(sum(self.traffic.shapes.values()) - 1.0) < 1e-9, f"{self.name}: shape shares must sum to 1"
        assert 0 <= self.traffic.stream_share <= 1, f"{self.name}: stream_share out of range"
        assert self.traffic.intervene_share + self.traffic.block_share <= 1, f"{self.name}: guardrail shares exceed 1"
        assert self.traffic.total_requests > 0 and self.traffic.rps > 0, f"{self.name}: empty traffic"
        assert self.model_groups, f"{self.name}: no model groups"
        assert 0 <= self.mock_latency.jitter_seconds <= self.mock_latency.seconds or self.mock_latency.seconds == 0, (
            f"{self.name}: mock latency jitter {self.mock_latency.jitter_seconds} s exceeds {self.mock_latency.seconds} s"
        )