litellm/tests/e2e/coverage_registry/schema.py
Yassin Kortam d5b91b94d3
test(e2e): replay a real tool-search assistant turn back to Bedrock Invoke (#36856)
The tool_search x bedrock_invoke cell only ever probed the first turn, so
nothing in the suite has sent a server_tool_use block back to a provider.
Every turn of a real Claude Code session after the first carries the
server_tool_use and tool_search_tool_result blocks the previous turn
produced, and that path was uncovered.

Adds probe_tool_search_multiturn, which takes the real assistant turn
back, answers any client-side tool_use with the id the model actually
emitted, and replays the whole thing as history with the tools still
declared. The assertion refuses to go green unless both server-tool
blocks made it into the replayed history, so a first turn truncated at
max_tokens reads as a failure instead of a vacuous pass.

The replay assertion's red paths never run in a green cell, so they get
markerless harness tests of their own alongside the existing
_builder_unit_tests tree.

No production code.
2026-08-17 11:59:26 -07:00

215 lines
4.6 KiB
Python

"""Registry row schema: the contract every denominator cell validates against.
A cell is one customer-noticeable behavior a single e2e test can assert pass/fail
on. `module` is the id's segment-1 prefix (eight of them); dashboard rollups can
split or merge those prefixes. The union is discriminated on `module`, so an LLM
row cannot carry a guardrail field and vice versa.
"""
from __future__ import annotations
from enum import Enum
from typing import Annotated, Literal
from pydantic import BaseModel, ConfigDict, Field, TypeAdapter
class Tier(str, Enum):
P0 = "P0"
P1 = "P1"
P2 = "P2"
class FailBeforeFix(str, Enum):
proven = "proven"
unproven = "unproven"
LlmEndpoint = Literal[
"chat_completions",
"completions",
"messages",
"responses",
"embeddings",
"batches",
"files",
"rerank",
"images_generations",
"images_edits",
"audio_speech",
"audio_transcriptions",
"moderations",
"realtime",
"google_native",
"vector_stores",
"ocr",
"bedrock_native",
]
LlmRoute = Literal[
"anthropic",
"azure_foundry",
"azure_openai",
"bedrock_converse",
"bedrock_invoke",
"cohere",
"gemini",
"hosted_vllm",
"openai",
"together_ai",
"vertex",
]
LlmCapability = Literal[
"assume_role",
"basic",
"count_tokens",
"input_validation",
"long_context_1m",
"mid_conversation_system",
"multi_turn",
"pdf_input",
"prompt_cache_1h",
"prompt_cache_5m",
"service_tier",
"structured_output",
"thinking",
"thinking_with_tool_use",
"tool_search",
"tool_search_history",
"tool_use",
"vision",
"web_search",
"web_search_server_tool",
]
class _Base(BaseModel):
model_config = ConfigDict(frozen=True, extra="forbid")
id: str
tier: Tier
assertions: tuple[str, ...]
source: str
rationale: str = ""
fail_before_fix: FailBeforeFix = FailBeforeFix.unproven
supported: bool = True
class LlmCell(_Base):
module: Literal["llm"]
subject_endpoint: LlmEndpoint
route: LlmRoute
capability: LlmCapability
streaming: Literal["stream", "nonstream", "na"]
class MgmtCell(_Base):
module: Literal["mgmt"]
surface: Literal["api", "ui"]
class McpCell(_Base):
module: Literal["mcp"]
operation: str
auth_family: Literal["none", "api_key", "bearer", "oauth"]
class ReliabilityCell(_Base):
module: Literal["reliability"]
behavior: str
variant: str
exercised_on: tuple[str, ...]
class QuotaCell(_Base):
module: Literal["quota_management"]
behavior: Literal["ratelimit", "budget", "spend_tracking"]
variant: str
exercised_on: tuple[str, ...]
class LoggingCell(_Base):
module: Literal["logging"]
event: str
exercised_on: tuple[str, ...]
class GuardrailCell(_Base):
module: Literal["guardrail"]
hook_point: str
exercised_on: tuple[str, ...]
class OtherCell(_Base):
module: Literal["other"]
area: str
Cell = Annotated[
LlmCell
| MgmtCell
| McpCell
| ReliabilityCell
| QuotaCell
| LoggingCell
| GuardrailCell
| OtherCell,
Field(discriminator="module"),
]
CELL_ADAPTER: TypeAdapter[Cell] = TypeAdapter(Cell)
CORE_LLM_ENDPOINTS: frozenset[str] = frozenset(
{
"chat_completions",
"messages",
"responses",
}
)
PREFIX_ROLLUP: dict[str, str] = {
"mcp": "MCPs",
"mgmt": "Management/UI",
"reliability": "Reliability & Performance",
"quota_management": "Quota Management",
"logging": "Logging & Guardrails",
"guardrail": "Logging & Guardrails",
"other": "Other",
}
MODULE_ORDER: tuple[str, ...] = (
"Core LLMs",
"Non-Core LLMs",
"MCPs",
"Management/UI",
"Reliability & Performance",
"Quota Management",
"Logging & Guardrails",
"Other",
)
LOKI_MODULE_LABELS: dict[str, str] = {
"Core LLMs": "core_llms",
"Non-Core LLMs": "non_core_llms",
"MCPs": "mcp",
"Management/UI": "management_ui",
"Reliability & Performance": "reliability_performance",
"Quota Management": "quota_management",
"Logging & Guardrails": "logging_guardrails",
"Other": "other",
}
def dashboard_module(cell: Cell) -> str:
"""Return the Grafana/reporting module for a registry cell."""
if isinstance(cell, LlmCell):
if cell.subject_endpoint in CORE_LLM_ENDPOINTS:
return "Core LLMs"
return "Non-Core LLMs"
return PREFIX_ROLLUP[cell.module]
def loki_module_label(module: str) -> str:
"""Return the log-safe Loki label for a dashboard module."""
return LOKI_MODULE_LABELS[module]