mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-16 04:29:31 +00:00
fix: keep model capability flags out of provider requests (#46647)
* fix: keep model capability flags out of provider requests A supports_* flag or another model map field set on a deployment's litellm_params was forwarded to the provider as a request field, so Bedrock answered 400 Extra inputs are not permitted on chat traffic and on the background health check alike. The shared owned-kwarg predicate now treats every ProviderSpecificModelInfo field as litellm-owned. Builds on the health-check-only approach in #40952 by leilei3167. * fix: treat every supports_ model map flag as a capability field and warn when one sits in litellm_params * fix(router): skip declared litellm params in the capability flag warning and warn for runtime-added deployments * fix(router): strip model capability flags from the router's deployments instead of every request * fix(router): keep capability flags out of health probes and read them from the loaded catalog * perf(router): read the catalog's supports_ field set once per cost map load Cache the loaded catalog's field names when the cost map snapshot is taken so the capability-flag check stops walking every catalog entry per deployment, and add integration cells that drive the flagged and control deployments through chat, streaming, messages, responses, embeddings, health and model management against a vendor that rejects the flags
This commit is contained in:
Mateo Wang
GitHub
parent
56dad9077d
commit
18831b818c
8 changed files with 934 additions and 22 deletions
|
|
@ -21,6 +21,7 @@ from collections.abc import Awaitable, Callable, Mapping
|
|||
from dataclasses import dataclass, replace
|
||||
from datetime import datetime, timezone
|
||||
from importlib.resources import files
|
||||
from itertools import chain
|
||||
from pathlib import Path
|
||||
from types import MappingProxyType
|
||||
from typing import Final, Protocol
|
||||
|
|
@ -93,16 +94,23 @@ class GetModelCostMap:
|
|||
return GetModelCostMap.load_local_model_cost_map_with_revision().model_cost_map
|
||||
|
||||
_loaded_catalog: Mapping[str, Mapping[str, object]] = MappingProxyType({})
|
||||
_loaded_catalog_fields: frozenset[str] = frozenset()
|
||||
|
||||
@classmethod
|
||||
def loaded_model_cost_map(cls) -> Mapping[str, Mapping[str, object]]:
|
||||
"""The catalog as last loaded (bundled or remote), untouched by ``register_model`` or router registrations."""
|
||||
return cls._loaded_catalog
|
||||
|
||||
@classmethod
|
||||
def loaded_model_cost_map_fields(cls) -> frozenset[str]:
|
||||
"""Every field name that at least one entry of the loaded catalog carries."""
|
||||
return cls._loaded_catalog_fields
|
||||
|
||||
@classmethod
|
||||
def _snapshot_loaded_catalog(cls, model_cost: Mapping[str, object]) -> None:
|
||||
raw: Final = _CATALOG_ADAPTER.validate_python(model_cost)
|
||||
cls._loaded_catalog = MappingProxyType({key: MappingProxyType(entry) for key, entry in raw.items()})
|
||||
cls._loaded_catalog_fields = frozenset(chain.from_iterable(entry.keys() for entry in raw.values()))
|
||||
|
||||
@classmethod
|
||||
def get_backup_model_count(cls) -> int:
|
||||
|
|
|
|||
|
|
@ -32,6 +32,7 @@ from litellm.router_utils.auto_router_model_naming import (
|
|||
classify_strategy_router_model,
|
||||
strategy_router_dependencies,
|
||||
)
|
||||
from litellm.router_utils.common_utils import is_model_capability_field
|
||||
from litellm.types.utils import secret_bearing_wif_litellm_params, server_owned_wif_litellm_params
|
||||
|
||||
# Provider routing and server-owned federation or OAuth token-exchange fields. Allowed for proxy
|
||||
|
|
@ -777,6 +778,9 @@ def _resolve_health_check_max_tokens(model_info: dict, litellm_params: dict) ->
|
|||
return None
|
||||
|
||||
|
||||
_PROBE_PARAMS_ADAPTER: Final = TypeAdapter(dict[str, object])
|
||||
|
||||
|
||||
def update_litellm_params_for_health_check(model_info: dict, litellm_params: dict) -> dict:
|
||||
"""
|
||||
Update the litellm params for health check.
|
||||
|
|
@ -867,7 +871,8 @@ def update_litellm_params_for_health_check(model_info: dict, litellm_params: dic
|
|||
"bedrock"
|
||||
)
|
||||
|
||||
return litellm_params
|
||||
probe_params: Final = _PROBE_PARAMS_ADAPTER.validate_python(litellm_params)
|
||||
return {key: value for key, value in probe_params.items() if not is_model_capability_field(key)}
|
||||
|
||||
|
||||
_update_litellm_params_for_health_check: Final = update_litellm_params_for_health_check
|
||||
|
|
|
|||
|
|
@ -201,7 +201,9 @@ from litellm.router_utils.common_utils import (
|
|||
provider_for_generic_call,
|
||||
resolve_model_group_alias,
|
||||
truncate_fallback_error_detail,
|
||||
warn_on_ignored_model_capability_flags,
|
||||
warn_on_provider_credential_mismatch,
|
||||
without_model_capability_flags,
|
||||
without_router_only_kwargs,
|
||||
)
|
||||
from litellm.router_utils.cooldown_cache import CooldownCache
|
||||
|
|
@ -503,6 +505,7 @@ def _stream_chunks_have_generated_content(chunks: Sequence[ModelResponseStream])
|
|||
_NO_SESSION_KWARGS: Final[Mapping[str, Mapping[str, object]]] = MappingProxyType({})
|
||||
_SESSION_ADAPTER: Final = TypeAdapter(Mapping[str, object])
|
||||
_MODEL_INFO_ADAPTER: Final = TypeAdapter(Mapping[str, object])
|
||||
_LITELLM_PARAMS_ADAPTER: Final = TypeAdapter(Mapping[str, object])
|
||||
_SILENT_MODEL_ADAPTER: Final = TypeAdapter(str | list[str])
|
||||
_RESOLVED_RETRY_POLICY_ADAPTER: Final = TypeAdapter(RetryPolicy | None)
|
||||
_ROUTING_KWARGS_ADAPTER: Final[TypeAdapter[Mapping[str, object] | None]] = TypeAdapter(Mapping[str, object] | None)
|
||||
|
|
@ -9491,12 +9494,16 @@ class Router:
|
|||
_litellm_params if zeroed_pricing is None else MappingProxyType({**_litellm_params, **zeroed_pricing})
|
||||
)
|
||||
litellm_params: Final[LiteLLM_Params] = LiteLLM_Params.model_validate(dict(**merged_params))
|
||||
warn_on_provider_credential_mismatch(model_name=_model_name, litellm_params=_litellm_params)
|
||||
deployment = Deployment(
|
||||
**deployment_info,
|
||||
model_name=_model_name,
|
||||
litellm_params=litellm_params,
|
||||
model_info=_model_info,
|
||||
declared_params: Final = _LITELLM_PARAMS_ADAPTER.validate_python(_litellm_params)
|
||||
warn_on_provider_credential_mismatch(model_name=_model_name, litellm_params=declared_params)
|
||||
warn_on_ignored_model_capability_flags(model_name=_model_name, litellm_params=declared_params)
|
||||
deployment = without_model_capability_flags(
|
||||
Deployment(
|
||||
**deployment_info,
|
||||
model_name=_model_name,
|
||||
litellm_params=litellm_params,
|
||||
model_info=_model_info,
|
||||
)
|
||||
)
|
||||
for field in CustomPricingLiteLLMParams.model_fields:
|
||||
if deployment.litellm_params.get(field) is not None:
|
||||
|
|
@ -10317,23 +10324,23 @@ class Router:
|
|||
if _deployment_model_id and self.has_model_id(_deployment_model_id):
|
||||
return None
|
||||
|
||||
warn_on_provider_credential_mismatch(
|
||||
model_name=deployment.model_name,
|
||||
litellm_params=deployment.litellm_params.model_dump(exclude_none=True),
|
||||
)
|
||||
declared_params: Final = deployment.litellm_params.model_dump(exclude_none=True)
|
||||
warn_on_provider_credential_mismatch(model_name=deployment.model_name, litellm_params=declared_params)
|
||||
warn_on_ignored_model_capability_flags(model_name=deployment.model_name, litellm_params=declared_params)
|
||||
served: Final = without_model_capability_flags(deployment)
|
||||
|
||||
# add to model list
|
||||
_deployment: Final = deployment.to_json(exclude_none=True)
|
||||
_deployment: Final = served.to_json(exclude_none=True)
|
||||
# initialize client
|
||||
self._add_deployment(deployment=deployment)
|
||||
self._add_deployment(deployment=served)
|
||||
|
||||
Router._register_deployment_pricing(deployment=deployment)
|
||||
Router._register_deployment_pricing(deployment=served)
|
||||
|
||||
# add to model names
|
||||
self._add_model_to_list_and_index_map(model=_deployment, model_id=deployment.model_info.id)
|
||||
self.model_names.add(deployment.model_name)
|
||||
self._sync_deployment_budget_config(deployment=deployment)
|
||||
return deployment
|
||||
self._add_model_to_list_and_index_map(model=_deployment, model_id=served.model_info.id)
|
||||
self.model_names.add(served.model_name)
|
||||
self._sync_deployment_budget_config(deployment=served)
|
||||
return served
|
||||
|
||||
def _update_deployment_indices_after_removal(self, model_id: str, removal_idx: int) -> None:
|
||||
"""
|
||||
|
|
@ -10463,12 +10470,13 @@ class Router:
|
|||
"""
|
||||
_deployment_model_id: Final = deployment.model_info.id or ""
|
||||
_deployment_on_router: Final[Deployment | None] = self.get_deployment(model_id=_deployment_model_id)
|
||||
served: Final = without_model_capability_flags(deployment)
|
||||
try:
|
||||
if _deployment_on_router is not None:
|
||||
# deployment with this model_id exists on the router
|
||||
if (
|
||||
deployment.model_name == _deployment_on_router.model_name
|
||||
and deployment.litellm_params == _deployment_on_router.litellm_params
|
||||
and served.litellm_params == _deployment_on_router.litellm_params
|
||||
and deployment.model_info == _deployment_on_router.model_info
|
||||
):
|
||||
# No need to update
|
||||
|
|
|
|||
|
|
@ -14,9 +14,16 @@ from litellm._logging import verbose_logger, verbose_router_logger
|
|||
from litellm.constants import ROUTER_FALLBACK_ERROR_DETAIL_MAX_CHARS
|
||||
from litellm.exceptions import BadRequestError
|
||||
from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider
|
||||
from litellm.litellm_core_utils.get_model_cost_map import GetModelCostMap
|
||||
from litellm.litellm_core_utils.sensitive_data_masker import mask_sensitive_structure
|
||||
from litellm.types.router import CredentialLiteLLMParams, FallbackSkipReason, SkippedFallbackTarget
|
||||
from litellm.types.utils import LlmProviders
|
||||
from litellm.types.router import (
|
||||
CredentialLiteLLMParams,
|
||||
Deployment,
|
||||
FallbackSkipReason,
|
||||
LiteLLM_Params,
|
||||
SkippedFallbackTarget,
|
||||
)
|
||||
from litellm.types.utils import LlmProviders, ProviderSpecificModelInfo
|
||||
|
||||
_V = TypeVar("_V")
|
||||
|
||||
|
|
@ -398,6 +405,43 @@ def provider_for_generic_call(litellm_params: Mapping[str, object]) -> str | Non
|
|||
return inferred
|
||||
|
||||
|
||||
PROVIDER_SPECIFIC_CAPABILITY_FIELDS: Final = frozenset(ProviderSpecificModelInfo.__annotations__)
|
||||
|
||||
|
||||
_LITELLM_PARAMS_ADAPTER: Final = TypeAdapter(Mapping[str, object])
|
||||
|
||||
|
||||
def is_model_capability_field(name: str) -> bool:
|
||||
return name in PROVIDER_SPECIFIC_CAPABILITY_FIELDS or (
|
||||
name.startswith("supports_") and name in GetModelCostMap.loaded_model_cost_map_fields()
|
||||
)
|
||||
|
||||
|
||||
def without_model_capability_flags(deployment: Deployment) -> Deployment:
|
||||
flags: Final = frozenset(
|
||||
key for key in (deployment.litellm_params.model_extra or {}) if is_model_capability_field(key)
|
||||
)
|
||||
if not flags:
|
||||
return deployment
|
||||
declared_params: Final = _LITELLM_PARAMS_ADAPTER.validate_python(deployment.litellm_params.model_dump())
|
||||
served_params: Final = {key: value for key, value in declared_params.items() if key not in flags}
|
||||
return deployment.model_copy(update={"litellm_params": LiteLLM_Params.model_validate(served_params)})
|
||||
|
||||
|
||||
def warn_on_ignored_model_capability_flags(model_name: str, litellm_params: Mapping[str, object]) -> str | None:
|
||||
flags: Final = sorted(
|
||||
key for key, value in litellm_params.items() if value is not None and is_model_capability_field(key)
|
||||
)
|
||||
if not flags:
|
||||
return None
|
||||
warning: Final = (
|
||||
f"Deployment '{model_name}' sets {flags} under litellm_params, where litellm ignores them and never "
|
||||
"sends them to the provider. Move them under model_info for litellm to use them."
|
||||
)
|
||||
verbose_router_logger.warning(warning)
|
||||
return warning
|
||||
|
||||
|
||||
def warn_on_provider_credential_mismatch(model_name: str, litellm_params: Mapping[str, object]) -> str | None:
|
||||
"""
|
||||
Warn when a deployment carries one provider's credentials but resolves to another.
|
||||
|
|
|
|||
716
tests/integration/routing/test_model_capability_flags.py
Normal file
716
tests/integration/routing/test_model_capability_flags.py
Normal file
|
|
@ -0,0 +1,716 @@
|
|||
import asyncio
|
||||
import json
|
||||
import re
|
||||
import threading
|
||||
import uuid
|
||||
from collections.abc import Callable, Iterator, Mapping
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
from dataclasses import dataclass
|
||||
from itertools import chain
|
||||
from pathlib import Path
|
||||
from types import MappingProxyType
|
||||
from typing import Final
|
||||
from urllib.parse import unquote
|
||||
|
||||
import anthropic
|
||||
import httpx
|
||||
import openai
|
||||
import pytest
|
||||
import yaml
|
||||
from integration._support.anthropic_sse import error_body, message_json
|
||||
from integration._support.client import Gateway, Scenario, eventually, list_value, object_value, string_value
|
||||
from integration._support.database import read_rows
|
||||
from integration._support.openai_wire import answering_model_discovery, chat_reply, responses_reply
|
||||
from integration._support.process import WorkerHealthcheck, graceful_stop_seconds, owned_proxy_process
|
||||
from integration._support.provider import MODEL_DISCOVERY
|
||||
from integration._support.upstream import aws_event_stream_frame
|
||||
from integration._support.wire import Reply, Request, Wire, wire_server
|
||||
from pydantic import JsonValue, TypeAdapter
|
||||
|
||||
_MARKER: Final = re.compile(r"capflag-[0-9a-f]{32}")
|
||||
_FLAG: Final = re.compile(r"^(supports_|bedrock_output_config_effort_ceiling$)")
|
||||
_OPENAI_MODEL: Final = "openai/gpt-4o-mini"
|
||||
_EMBEDDING_MODEL: Final = "openai/text-embedding-3-small"
|
||||
_CLAUDE_MODEL: Final = "anthropic/claude-sonnet-4-5-20250929"
|
||||
_BEDROCK_ID: Final = "anthropic.claude-3-haiku-20240307-v1:0"
|
||||
_PROVIDER_KEY: Final = "integration-provider-key"
|
||||
_REVOKED_KEY: Final = "sk-revoked-capflag"
|
||||
_AWS: Final[Mapping[str, JsonValue]] = MappingProxyType(
|
||||
{
|
||||
"aws_access_key_id": "AKIASCRIPTEDPROVIDER",
|
||||
"aws_secret_access_key": "scripted-secret",
|
||||
"aws_region_name": "us-east-1",
|
||||
}
|
||||
)
|
||||
_VISION: Final[Mapping[str, JsonValue]] = MappingProxyType({"supports_vision": True})
|
||||
_EMBEDDING: Final = [0.25, 0.5]
|
||||
_JSON: Final = TypeAdapter(dict[str, JsonValue])
|
||||
_PROXY_CONFIG: Final = Path(__file__).resolve().parents[1] / "proxy_config.yaml"
|
||||
_FRESH: Final = httpx.Limits(max_keepalive_connections=0)
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class _Vendor:
|
||||
wire: Wire
|
||||
outage: threading.Event
|
||||
|
||||
def outbound(self) -> tuple[Request, ...]:
|
||||
return tuple(request for request in self.wire.drain() if (request.method, request.target) != MODEL_DISCOVERY)
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class _Exchange:
|
||||
marker: str
|
||||
text: str
|
||||
request: Request
|
||||
|
||||
|
||||
Send = Callable[[str, str], str]
|
||||
|
||||
|
||||
def _marker() -> str:
|
||||
return f"capflag-{uuid.uuid4().hex}"
|
||||
|
||||
|
||||
def _echo(marker: str) -> str:
|
||||
return f"echo {marker}"
|
||||
|
||||
|
||||
def _body(request: Request) -> dict[str, JsonValue]:
|
||||
return _JSON.validate_json(request.body) if request.body else {}
|
||||
|
||||
|
||||
def _flags(value: JsonValue) -> tuple[str, ...]:
|
||||
if isinstance(value, dict):
|
||||
own: Final = tuple(key for key in value if _FLAG.match(key))
|
||||
return own + tuple(chain.from_iterable(_flags(item) for item in value.values()))
|
||||
if isinstance(value, list):
|
||||
return tuple(chain.from_iterable(_flags(item) for item in value))
|
||||
return ()
|
||||
|
||||
|
||||
def _openai_error(status: int, message: str, code: str, param: str | None) -> Reply:
|
||||
error: Final = {"message": message, "type": "invalid_request_error", "param": param, "code": code}
|
||||
return Reply(status=status, body=json.dumps({"error": error}).encode())
|
||||
|
||||
|
||||
def _rejection(target: str, flag: str) -> Reply:
|
||||
if target.startswith("/model/"):
|
||||
return Reply(
|
||||
status=400,
|
||||
body=json.dumps(
|
||||
{"message": f"The model returned the following errors: {flag}: Extra inputs are not permitted"}
|
||||
).encode(),
|
||||
headers=MappingProxyType({"x-amzn-ErrorType": "ValidationException"}),
|
||||
)
|
||||
if target == "/v1/messages":
|
||||
return Reply(status=400, body=error_body(400, f"{flag}: Extra inputs are not permitted"))
|
||||
return _openai_error(400, f"Unknown parameter: '{flag}'.", "unknown_parameter", flag)
|
||||
|
||||
|
||||
def _converse_reply(text: str) -> Reply:
|
||||
return Reply(
|
||||
body=json.dumps(
|
||||
{
|
||||
"output": {"message": {"role": "assistant", "content": [{"text": text}]}},
|
||||
"stopReason": "end_turn",
|
||||
"usage": {"inputTokens": 11, "outputTokens": 4, "totalTokens": 15},
|
||||
"metrics": {"latencyMs": 1},
|
||||
}
|
||||
).encode()
|
||||
)
|
||||
|
||||
|
||||
def _converse_stream_reply(text: str) -> Reply:
|
||||
events: Final[tuple[tuple[str, dict[str, JsonValue]], ...]] = (
|
||||
("messageStart", {"role": "assistant"}),
|
||||
("contentBlockDelta", {"delta": {"text": text}, "contentBlockIndex": 0}),
|
||||
("messageStop", {"stopReason": "end_turn"}),
|
||||
("metadata", {"usage": {"inputTokens": 11, "outputTokens": 4, "totalTokens": 15}}),
|
||||
)
|
||||
return Reply(
|
||||
content_type="application/vnd.amazon.eventstream",
|
||||
chunks=tuple(
|
||||
aws_event_stream_frame(
|
||||
{":event-type": kind, ":content-type": "application/json", ":message-type": "event"},
|
||||
json.dumps(payload).encode(),
|
||||
)
|
||||
for kind, payload in events
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
def _embedding_reply(model: str) -> Reply:
|
||||
return Reply(
|
||||
body=json.dumps(
|
||||
{
|
||||
"object": "list",
|
||||
"data": [{"object": "embedding", "index": 0, "embedding": _EMBEDDING}],
|
||||
"model": model,
|
||||
"usage": {"prompt_tokens": 2, "total_tokens": 2},
|
||||
}
|
||||
).encode()
|
||||
)
|
||||
|
||||
|
||||
def _answer(target: str, body: Mapping[str, JsonValue], marker: str) -> Reply:
|
||||
text: Final = _echo(marker)
|
||||
model: Final = str(body.get("model", "scripted"))
|
||||
streamed: Final = body.get("stream") is True
|
||||
match target:
|
||||
case "/v1/chat/completions":
|
||||
return chat_reply(f"chatcmpl-{uuid.uuid4().hex}", model, text, stream=streamed)
|
||||
case "/v1/responses":
|
||||
return responses_reply(f"resp_{uuid.uuid4().hex}", model, text, stream=streamed)
|
||||
case "/v1/embeddings":
|
||||
return _embedding_reply(model)
|
||||
case "/v1/messages":
|
||||
return Reply(body=message_json(f"msg_{uuid.uuid4().hex}", model, text))
|
||||
case _ if target.endswith("/converse-stream"):
|
||||
return _converse_stream_reply(text)
|
||||
case _ if target.endswith("/converse"):
|
||||
return _converse_reply(text)
|
||||
case _ if target.endswith("/invoke"):
|
||||
return Reply(body=message_json(f"msg_{uuid.uuid4().hex}", "claude-3-haiku", text))
|
||||
case _:
|
||||
return Reply(status=404, body=b'{"error": "unscripted target"}')
|
||||
|
||||
|
||||
def _vendor(outage: threading.Event) -> Callable[[Request], Reply]:
|
||||
def respond(request: Request) -> Reply:
|
||||
if outage.is_set():
|
||||
return Reply(drop_connection=True)
|
||||
if request.headers.get("authorization", "") == f"Bearer {_REVOKED_KEY}":
|
||||
return _openai_error(401, "Incorrect API key provided: sk-revo****flag.", "invalid_api_key", None)
|
||||
target: Final = unquote(request.target.split("?")[0])
|
||||
body: Final = _body(request)
|
||||
flags: Final = _flags(body)
|
||||
if flags:
|
||||
return _rejection(target, flags[0])
|
||||
found: Final = _MARKER.search(request.body.decode())
|
||||
return _answer(target, body, found.group(0) if found else "probe")
|
||||
|
||||
return answering_model_discovery(respond)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def vendor() -> Iterator[_Vendor]:
|
||||
outage: Final = threading.Event()
|
||||
with wire_server(_vendor(outage)) as wire:
|
||||
yield _Vendor(wire, outage)
|
||||
|
||||
|
||||
def _proxy(gateway: Gateway) -> str:
|
||||
return str(gateway.client.base_url).rstrip("/")
|
||||
|
||||
|
||||
def _openai_client(gateway: Gateway) -> openai.OpenAI:
|
||||
return openai.OpenAI(base_url=f"{_proxy(gateway)}/v1", api_key=gateway.key, max_retries=0)
|
||||
|
||||
|
||||
def _chat(gateway: Gateway, extra_body: Mapping[str, JsonValue] | None = None) -> Send:
|
||||
def send(model: str, marker: str) -> str:
|
||||
with _openai_client(gateway) as client:
|
||||
completion: Final = client.chat.completions.create(
|
||||
model=model, messages=[{"role": "user", "content": marker}], extra_body=dict(extra_body or {})
|
||||
)
|
||||
return completion.choices[0].message.content or ""
|
||||
|
||||
return send
|
||||
|
||||
|
||||
def _chat_stream(gateway: Gateway) -> Send:
|
||||
async def streamed(model: str, marker: str) -> str:
|
||||
async with openai.AsyncOpenAI(base_url=f"{_proxy(gateway)}/v1", api_key=gateway.key, max_retries=0) as client:
|
||||
stream: Final = await client.chat.completions.create(
|
||||
model=model, messages=[{"role": "user", "content": marker}], stream=True
|
||||
)
|
||||
return "".join([chunk.choices[0].delta.content or "" async for chunk in stream if chunk.choices])
|
||||
|
||||
def send(model: str, marker: str) -> str:
|
||||
return asyncio.run(streamed(model, marker))
|
||||
|
||||
return send
|
||||
|
||||
|
||||
def _responses(gateway: Gateway) -> Send:
|
||||
def send(model: str, marker: str) -> str:
|
||||
with _openai_client(gateway) as client:
|
||||
return client.responses.create(model=model, input=marker).output_text
|
||||
|
||||
return send
|
||||
|
||||
|
||||
def _embeddings(gateway: Gateway) -> Send:
|
||||
def send(model: str, marker: str) -> str:
|
||||
with _openai_client(gateway) as client:
|
||||
return json.dumps(client.embeddings.create(model=model, input=marker).data[0].embedding)
|
||||
|
||||
return send
|
||||
|
||||
|
||||
def _messages(gateway: Gateway) -> Send:
|
||||
def send(model: str, marker: str) -> str:
|
||||
with anthropic.Anthropic(base_url=_proxy(gateway), api_key=gateway.key, max_retries=0) as client:
|
||||
message: Final = client.messages.create(
|
||||
model=model, max_tokens=16, messages=[{"role": "user", "content": marker}]
|
||||
)
|
||||
return "".join(block.text for block in message.content if block.type == "text")
|
||||
|
||||
return send
|
||||
|
||||
|
||||
def _family(name: str, wire: Wire) -> dict[str, JsonValue]:
|
||||
match name:
|
||||
case "openai":
|
||||
return {"model": _OPENAI_MODEL, "api_base": f"{wire.url}/v1"}
|
||||
case "anthropic":
|
||||
return {"model": _CLAUDE_MODEL, "api_base": wire.url, "api_key": "sk-ant-scripted"}
|
||||
case "converse":
|
||||
return {"model": f"bedrock/converse/{_BEDROCK_ID}", "api_base": wire.url, "api_key": None, **_AWS}
|
||||
case "invoke":
|
||||
return {
|
||||
"model": f"bedrock/invoke/{_BEDROCK_ID}",
|
||||
"api_base": wire.url,
|
||||
"aws_bedrock_runtime_endpoint": wire.url,
|
||||
"api_key": None,
|
||||
**_AWS,
|
||||
}
|
||||
case _:
|
||||
raise AssertionError(f"unknown family {name}")
|
||||
|
||||
|
||||
def _deployment(
|
||||
scenario: Scenario,
|
||||
params: Mapping[str, JsonValue],
|
||||
model_info: Mapping[str, JsonValue] | None = None,
|
||||
name: str | None = None,
|
||||
) -> str:
|
||||
model_name: Final = name or f"capability-flag-{uuid.uuid4().hex}"
|
||||
body: Final[dict[str, JsonValue]] = {
|
||||
"model_name": model_name,
|
||||
"litellm_params": {"api_key": _PROVIDER_KEY, **params},
|
||||
"model_info": dict(model_info or {}),
|
||||
}
|
||||
created: Final = scenario.gateway.post("/model/new", body)
|
||||
scenario.cleanups.callback(scenario.delete_model, string_value(object_value(created["model_info"])["id"]))
|
||||
return model_name
|
||||
|
||||
|
||||
def _idle(gateway: Gateway) -> dict[str, JsonValue]:
|
||||
return {"model": _OPENAI_MODEL, "api_base": f"{gateway.upstream_url}/v1"}
|
||||
|
||||
|
||||
def _pair(scenario: Scenario, params: Mapping[str, JsonValue], flag: Mapping[str, JsonValue]) -> tuple[str, str]:
|
||||
return _deployment(scenario, {**params, **flag}), _deployment(scenario, params)
|
||||
|
||||
|
||||
def _exchange(vendor: _Vendor, send: Send, model: str) -> _Exchange:
|
||||
marker: Final = _marker()
|
||||
text: Final = send(model, marker)
|
||||
(request,) = vendor.outbound()
|
||||
return _Exchange(marker, text, request)
|
||||
|
||||
|
||||
def _shape(request: Request) -> tuple[str, JsonValue]:
|
||||
return unquote(request.target), json.loads(_MARKER.sub("capflag-marker", request.body.decode()))
|
||||
|
||||
|
||||
def _probe_shape(request: Request) -> tuple[str, dict[str, JsonValue]]:
|
||||
return unquote(request.target), {key: value for key, value in _body(request).items() if key != "messages"}
|
||||
|
||||
|
||||
def _assert_sends_the_control_request(
|
||||
vendor: _Vendor, send: Send, flagged: str, control: str, expected: Callable[[str], str] = _echo
|
||||
) -> None:
|
||||
sent: Final = _exchange(vendor, send, flagged)
|
||||
baseline: Final = _exchange(vendor, send, control)
|
||||
assert sent.text == expected(sent.marker), sent.text
|
||||
assert baseline.text == expected(baseline.marker), baseline.text
|
||||
assert _flags(_body(sent.request)) == (), sent.request.body
|
||||
assert _shape(sent.request) == _shape(baseline.request)
|
||||
|
||||
|
||||
def _entry(gateway: Gateway, path: str, name: str, params: Mapping[str, str] | None = None) -> dict[str, JsonValue]:
|
||||
entries: Final = list_value(gateway.get(path, params)["data"])
|
||||
(found,) = tuple(object_value(entry) for entry in entries if object_value(entry)["model_name"] == name)
|
||||
return found
|
||||
|
||||
|
||||
def _model_id(gateway: Gateway, name: str) -> str:
|
||||
return string_value(object_value(_entry(gateway, "/model/info", name)["model_info"])["id"])
|
||||
|
||||
|
||||
def _health(gateway: Gateway, name: str) -> tuple[int, JsonValue, JsonValue]:
|
||||
response: Final = gateway.request("GET", "/health", params={"model": name})
|
||||
health: Final = _JSON.validate_json(response.content)
|
||||
return response.status_code, health["healthy_count"], health["unhealthy_count"]
|
||||
|
||||
|
||||
def _fresh_request(gateway: Gateway, method: str, path: str, body: Mapping[str, JsonValue] | None) -> httpx.Response:
|
||||
with httpx.Client(base_url=_proxy(gateway), timeout=60, trust_env=False, limits=_FRESH) as client:
|
||||
return client.request(method, path, json=body, headers={"Authorization": f"Bearer {gateway.key}"})
|
||||
|
||||
|
||||
def _fresh_chat(gateway: Gateway, model: str, marker: str) -> str:
|
||||
response: Final = _fresh_request(
|
||||
gateway, "POST", "/v1/chat/completions", {"model": model, "messages": [{"role": "user", "content": marker}]}
|
||||
)
|
||||
assert response.status_code == 200, response.text
|
||||
choice: Final = object_value(list_value(_JSON.validate_json(response.content)["choices"])[0])
|
||||
return string_value(object_value(choice["message"])["content"])
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("family", "flag"),
|
||||
[
|
||||
("openai", {"supports_vision": True}),
|
||||
("openai", {"supports_speed": True}),
|
||||
("openai", {"bedrock_output_config_effort_ceiling": "high"}),
|
||||
("anthropic", {"supports_vision": True}),
|
||||
("converse", {"supports_max_reasoning_effort": False}),
|
||||
("invoke", {"supports_max_reasoning_effort": False}),
|
||||
],
|
||||
ids=["openai-static", "openai-catalog_only", "openai-ceiling", "anthropic", "bedrock_converse", "bedrock_invoke"],
|
||||
)
|
||||
def test_chat_completions_send_a_flagged_deployment_the_control_request(
|
||||
gateway: Gateway, vendor: _Vendor, family: str, flag: Mapping[str, JsonValue]
|
||||
) -> None:
|
||||
with gateway.scenario() as scenario:
|
||||
flagged, control = _pair(scenario, _family(family, vendor.wire), flag)
|
||||
_assert_sends_the_control_request(vendor, _chat(gateway), flagged, control)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("family", ["openai", "converse"])
|
||||
def test_streamed_chat_completions_send_a_flagged_deployment_the_control_request(
|
||||
gateway: Gateway, vendor: _Vendor, family: str
|
||||
) -> None:
|
||||
flag: Final = {"supports_max_reasoning_effort": False} if family == "converse" else _VISION
|
||||
with gateway.scenario() as scenario:
|
||||
flagged, control = _pair(scenario, _family(family, vendor.wire), flag)
|
||||
_assert_sends_the_control_request(vendor, _chat_stream(gateway), flagged, control)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("family", ["openai", "anthropic", "converse"])
|
||||
def test_messages_send_a_flagged_deployment_the_control_request(gateway: Gateway, vendor: _Vendor, family: str) -> None:
|
||||
with gateway.scenario() as scenario:
|
||||
flagged, control = _pair(scenario, _family(family, vendor.wire), _VISION)
|
||||
_assert_sends_the_control_request(vendor, _messages(gateway), flagged, control)
|
||||
|
||||
|
||||
def test_responses_send_a_flagged_deployment_the_control_request(gateway: Gateway, vendor: _Vendor) -> None:
|
||||
with gateway.scenario() as scenario:
|
||||
flagged, control = _pair(scenario, _family("openai", vendor.wire), _VISION)
|
||||
_assert_sends_the_control_request(vendor, _responses(gateway), flagged, control)
|
||||
|
||||
|
||||
def test_embeddings_send_a_flagged_deployment_the_control_request(gateway: Gateway, vendor: _Vendor) -> None:
|
||||
params: Final = {"model": _EMBEDDING_MODEL, "api_base": f"{vendor.wire.url}/v1"}
|
||||
with gateway.scenario() as scenario:
|
||||
flagged, control = _pair(scenario, params, {"supports_embedding_image_input": True})
|
||||
_assert_sends_the_control_request(
|
||||
vendor, _embeddings(gateway), flagged, control, expected=lambda _marker: json.dumps(_EMBEDDING)
|
||||
)
|
||||
|
||||
|
||||
def _wildcard(scenario: Scenario, wire: Wire, flag: Mapping[str, JsonValue]) -> str:
|
||||
prefix: Final = f"capflag{uuid.uuid4().hex}"
|
||||
_deployment(scenario, {"model": "openai/*", "api_base": f"{wire.url}/v1", **flag}, name=f"{prefix}/*")
|
||||
return f"{prefix}/gpt-4o-mini"
|
||||
|
||||
|
||||
def test_a_flagged_wildcard_deployment_sends_the_control_request(gateway: Gateway, vendor: _Vendor) -> None:
|
||||
with gateway.scenario() as scenario:
|
||||
flagged: Final = _wildcard(scenario, vendor.wire, _VISION)
|
||||
control: Final = _wildcard(scenario, vendor.wire, {})
|
||||
_assert_sends_the_control_request(vendor, _chat(gateway), flagged, control)
|
||||
|
||||
|
||||
def test_health_probes_a_flagged_deployment_with_the_control_probe(gateway: Gateway, vendor: _Vendor) -> None:
|
||||
with gateway.scenario() as scenario:
|
||||
flagged, control = _pair(scenario, _family("openai", vendor.wire), _VISION)
|
||||
flagged_health: Final = _health(gateway, flagged)
|
||||
(flagged_probe,) = vendor.outbound()
|
||||
control_health: Final = _health(gateway, control)
|
||||
(control_probe,) = vendor.outbound()
|
||||
assert flagged_health == control_health == (200, 1, 0)
|
||||
assert _flags(_body(flagged_probe)) == (), flagged_probe.body
|
||||
assert _probe_shape(flagged_probe) == _probe_shape(control_probe)
|
||||
|
||||
|
||||
def test_test_connection_probes_flagged_params_with_the_unflagged_probe(gateway: Gateway, vendor: _Vendor) -> None:
|
||||
params: Final[dict[str, JsonValue]] = {
|
||||
"model": _OPENAI_MODEL,
|
||||
"api_base": f"{vendor.wire.url}/v1",
|
||||
"api_key": _PROVIDER_KEY,
|
||||
}
|
||||
flagged: Final = gateway.post("/health/test_connection", {"mode": "chat", "litellm_params": {**params, **_VISION}})
|
||||
(flagged_probe,) = vendor.outbound()
|
||||
unflagged: Final = gateway.post("/health/test_connection", {"mode": "chat", "litellm_params": params})
|
||||
(unflagged_probe,) = vendor.outbound()
|
||||
assert (flagged["status"], unflagged["status"]) == ("success", "success"), (flagged, unflagged)
|
||||
assert _flags(_body(flagged_probe)) == (), flagged_probe.body
|
||||
assert _probe_shape(flagged_probe) == _probe_shape(unflagged_probe)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"value",
|
||||
[7, ["vision"], "", "v" * 5120, None, False],
|
||||
ids=["int", "list", "empty_string", "five_kb", "null", "false"],
|
||||
)
|
||||
def test_any_flag_value_sends_the_control_request(gateway: Gateway, vendor: _Vendor, value: JsonValue) -> None:
|
||||
with gateway.scenario() as scenario:
|
||||
flagged, control = _pair(scenario, _family("openai", vendor.wire), {"supports_vision": value})
|
||||
_assert_sends_the_control_request(vendor, _chat(gateway), flagged, control)
|
||||
|
||||
|
||||
def test_a_flag_in_litellm_params_and_model_info_keeps_the_model_info_value(gateway: Gateway, vendor: _Vendor) -> None:
|
||||
with gateway.scenario() as scenario:
|
||||
flagged: Final = _deployment(
|
||||
scenario, {**_family("openai", vendor.wire), **_VISION}, {"supports_vision": False}
|
||||
)
|
||||
control: Final = _deployment(scenario, _family("openai", vendor.wire))
|
||||
_assert_sends_the_control_request(vendor, _chat(gateway), flagged, control)
|
||||
model_info: Final = object_value(_entry(gateway, "/model/info", flagged)["model_info"])
|
||||
assert model_info["supports_vision"] is False, model_info
|
||||
|
||||
|
||||
@pytest.mark.parametrize("key", ["supports_litellm_qa_unknown", "supports_system_prompt"])
|
||||
def test_an_undeclared_or_provider_supports_key_still_reaches_the_vendor(
|
||||
gateway: Gateway, vendor: _Vendor, key: str
|
||||
) -> None:
|
||||
with gateway.scenario() as scenario:
|
||||
model: Final = _deployment(scenario, {**_family("openai", vendor.wire), key: True})
|
||||
with pytest.raises(openai.BadRequestError) as raised:
|
||||
_chat(gateway)(model, _marker())
|
||||
(request,) = vendor.outbound()
|
||||
assert f"Unknown parameter: '{key}'." in str(raised.value), raised.value
|
||||
assert _flags(_body(request)) == (key,), request.body
|
||||
|
||||
|
||||
def test_a_flagged_deployment_with_a_revoked_key_returns_the_vendor_401(gateway: Gateway, vendor: _Vendor) -> None:
|
||||
with gateway.scenario() as scenario:
|
||||
flagged: Final = _deployment(scenario, {**_family("openai", vendor.wire), "api_key": _REVOKED_KEY, **_VISION})
|
||||
control: Final = _deployment(scenario, _family("openai", vendor.wire))
|
||||
with pytest.raises(openai.AuthenticationError) as raised:
|
||||
_chat(gateway)(flagged, _marker())
|
||||
assert len(vendor.outbound()) == 1
|
||||
marker: Final = _marker()
|
||||
answered: Final = _chat(gateway)(control, marker)
|
||||
assert "Incorrect API key provided" in str(raised.value), raised.value
|
||||
assert answered == _echo(marker)
|
||||
|
||||
|
||||
def test_a_fallback_to_a_flagged_deployment_sends_the_control_request(gateway: Gateway, vendor: _Vendor) -> None:
|
||||
with gateway.scenario() as scenario:
|
||||
broken: Final = _deployment(scenario, {**_family("openai", vendor.wire), "api_key": _REVOKED_KEY})
|
||||
flagged, control = _pair(scenario, _family("openai", vendor.wire), _VISION)
|
||||
marker: Final = _marker()
|
||||
answered: Final = _chat(gateway, {"fallbacks": [flagged]})(broken, marker)
|
||||
rejected, fallback = vendor.outbound()
|
||||
baseline: Final = _exchange(vendor, _chat(gateway), control)
|
||||
assert answered == _echo(marker)
|
||||
assert rejected.headers["authorization"] == f"Bearer {_REVOKED_KEY}"
|
||||
assert _flags(_body(fallback)) == (), fallback.body
|
||||
assert _shape(fallback) == _shape(baseline.request)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("path", "by_name"), [("/model/info", False), ("/v2/model/info", True)], ids=["model_info", "v2_model_info"]
|
||||
)
|
||||
def test_model_info_lists_the_served_params_without_the_flag(gateway: Gateway, path: str, by_name: bool) -> None:
|
||||
with gateway.scenario() as scenario:
|
||||
flagged: Final = _deployment(scenario, {**_idle(gateway), **_VISION})
|
||||
entry: Final = _entry(gateway, path, flagged, {"search": flagged} if by_name else None)
|
||||
litellm_params: Final = object_value(entry["litellm_params"])
|
||||
assert litellm_params["model"] == _OPENAI_MODEL, litellm_params
|
||||
assert "supports_vision" not in litellm_params, litellm_params
|
||||
|
||||
|
||||
def test_model_group_info_reports_a_flagged_group_like_a_plain_one(gateway: Gateway) -> None:
|
||||
with gateway.scenario() as scenario:
|
||||
flagged: Final = _deployment(scenario, {**_idle(gateway), **_VISION})
|
||||
plain: Final = _deployment(scenario, _idle(gateway))
|
||||
groups: Final = tuple(
|
||||
object_value(list_value(gateway.get("/model_group/info", {"model_group": name})["data"])[0])
|
||||
for name in (flagged, plain)
|
||||
)
|
||||
flagged_group, plain_group = (
|
||||
{key: value for key, value in group.items() if key != "model_group"} for group in groups
|
||||
)
|
||||
assert flagged_group == plain_group
|
||||
|
||||
|
||||
def test_supports_system_message_false_still_folds_the_system_prompt(gateway: Gateway, vendor: _Vendor) -> None:
|
||||
instruction: Final = f"answer briefly {uuid.uuid4().hex}"
|
||||
marker: Final = _marker()
|
||||
with gateway.scenario() as scenario:
|
||||
model: Final = _deployment(scenario, {**_family("openai", vendor.wire), "supports_system_message": False})
|
||||
with _openai_client(gateway) as client:
|
||||
completion: Final = client.chat.completions.create(
|
||||
model=model,
|
||||
messages=[{"role": "system", "content": instruction}, {"role": "user", "content": marker}],
|
||||
)
|
||||
(request,) = vendor.outbound()
|
||||
messages: Final = list_value(_body(request)["messages"])
|
||||
assert completion.choices[0].message.content == _echo(marker)
|
||||
assert [object_value(message)["role"] for message in messages] == ["user"], messages
|
||||
assert instruction in json.dumps(messages) and marker in json.dumps(messages), messages
|
||||
assert _flags(_body(request)) == (), request.body
|
||||
|
||||
|
||||
def test_three_identical_requests_to_a_flagged_deployment_each_log_once(gateway: Gateway, vendor: _Vendor) -> None:
|
||||
marker: Final = _marker()
|
||||
with gateway.scenario() as scenario:
|
||||
flagged: Final = _deployment(scenario, {**_family("openai", vendor.wire), **_VISION})
|
||||
with _openai_client(gateway) as client:
|
||||
identities: Final = tuple(
|
||||
client.chat.completions.create(
|
||||
model=flagged,
|
||||
messages=[{"role": "user", "content": marker}],
|
||||
extra_body={"cache": {"no-cache": True}},
|
||||
).id
|
||||
for _ in range(3)
|
||||
)
|
||||
received: Final = vendor.outbound()
|
||||
rows: Final = eventually(
|
||||
lambda: read_rows(
|
||||
'SELECT request_id FROM "LiteLLM_SpendLogs" WHERE request_id IN (%s, %s, %s)', identities
|
||||
),
|
||||
lambda found: sorted(str(row["request_id"]) for row in found) == sorted(identities),
|
||||
seconds=70,
|
||||
)
|
||||
assert len(set(identities)) == 3, identities
|
||||
assert len(rows) == 3, rows
|
||||
assert len(received) == 3 and len({request.body for request in received}) == 1, received
|
||||
assert _flags(_body(received[0])) == (), received[0].body
|
||||
|
||||
|
||||
def _rpm_on_fresh_connections(gateway: Gateway, name: str) -> tuple[JsonValue, ...]:
|
||||
def read() -> JsonValue:
|
||||
response: Final = _fresh_request(gateway, "GET", "/model/info", None)
|
||||
entries: Final = list_value(_JSON.validate_json(response.content)["data"])
|
||||
(entry,) = tuple(object_value(item) for item in entries if object_value(item)["model_name"] == name)
|
||||
return object_value(entry["litellm_params"]).get("rpm")
|
||||
|
||||
return tuple(read() for _ in range(8))
|
||||
|
||||
|
||||
def test_a_flag_patched_onto_a_db_deployment_never_reaches_the_vendor(gateway: Gateway, vendor: _Vendor) -> None:
|
||||
with gateway.scenario() as scenario:
|
||||
patched: Final = _deployment(scenario, _family("openai", vendor.wire))
|
||||
control: Final = _deployment(scenario, _family("openai", vendor.wire))
|
||||
response: Final = gateway.request(
|
||||
"PATCH",
|
||||
f"/model/{_model_id(gateway, patched)}/update",
|
||||
{"litellm_params": {"supports_vision": True, "rpm": 7}},
|
||||
)
|
||||
assert response.status_code == 200, response.text
|
||||
eventually(lambda: _rpm_on_fresh_connections(gateway, patched), lambda seen: seen == (7,) * 8, seconds=60)
|
||||
for _ in range(3):
|
||||
_assert_sends_the_control_request(vendor, _chat(gateway), patched, control)
|
||||
|
||||
|
||||
def test_updates_during_a_burst_never_send_the_flag(gateway: Gateway, vendor: _Vendor) -> None:
|
||||
markers: Final = tuple(_marker() for _ in range(20))
|
||||
with gateway.scenario() as scenario:
|
||||
flagged: Final = _deployment(scenario, {**_family("openai", vendor.wire), **_VISION})
|
||||
model_id: Final = _model_id(gateway, flagged)
|
||||
send: Final = _chat(gateway)
|
||||
with ThreadPoolExecutor(max_workers=8) as pool:
|
||||
answers: Final = tuple(pool.submit(send, flagged, marker) for marker in markers)
|
||||
updates: Final = tuple(
|
||||
gateway.request(
|
||||
"PATCH", f"/model/{model_id}/update", {"litellm_params": {"rpm": rpm, "supports_vision": True}}
|
||||
).status_code
|
||||
for rpm in (11, 12, 13)
|
||||
)
|
||||
texts: Final = tuple(answer.result() for answer in answers)
|
||||
received: Final = vendor.outbound()
|
||||
assert updates == (200, 200, 200), updates
|
||||
assert texts == tuple(_echo(marker) for marker in markers), texts
|
||||
assert len(received) == 20, received
|
||||
assert [_flags(_body(request)) for request in received] == [()] * 20, received
|
||||
|
||||
|
||||
def _config_with_flagged_model(directory: Path, name: str, wire: Wire) -> Path:
|
||||
base: Final = _JSON.validate_python(yaml.safe_load(_PROXY_CONFIG.read_text()))
|
||||
flagged: Final[dict[str, JsonValue]] = {
|
||||
"model_name": name,
|
||||
"litellm_params": {"model": _OPENAI_MODEL, "api_key": _PROVIDER_KEY, "api_base": f"{wire.url}/v1", **_VISION},
|
||||
}
|
||||
path: Final = directory / "capability_flags_config.yaml"
|
||||
path.write_text(yaml.safe_dump({**base, "model_list": [*list_value(base["model_list"]), flagged]}))
|
||||
return path
|
||||
|
||||
|
||||
@pytest.mark.timeout(int(2 * graceful_stop_seconds() + 120))
|
||||
def test_each_worker_and_a_replacement_send_config_and_db_flagged_deployments_without_the_flag(
|
||||
gateway: Gateway, vendor: _Vendor, tmp_path: Path
|
||||
) -> None:
|
||||
configured: Final = f"capflag-config-{uuid.uuid4().hex}"
|
||||
config: Final = _config_with_flagged_model(tmp_path, configured, vendor.wire)
|
||||
with owned_proxy_process(
|
||||
gateway, tmp_path, {}, config=config, workers=2, healthcheck=WorkerHealthcheck.LENIENT
|
||||
) as owned:
|
||||
with owned.gateway.scenario() as scenario:
|
||||
stored: Final = _deployment(scenario, {**_family("openai", vendor.wire), **_VISION})
|
||||
stored_id: Final = _model_id(owned.gateway, stored)
|
||||
eventually(lambda: owned.serves_only_on_every_worker(stored, stored_id), bool, seconds=120)
|
||||
first, second = sorted(eventually(owned.worker_pids, lambda pids: len(pids) == 2))
|
||||
with owned.suspended(second):
|
||||
on_first: Final = tuple((model, _marker()) for model in (configured, stored))
|
||||
first_texts: Final = tuple(_fresh_chat(owned.gateway, model, marker) for model, marker in on_first)
|
||||
with owned.suspended(first):
|
||||
on_second: Final = tuple((model, _marker()) for model in (configured, stored))
|
||||
second_texts: Final = tuple(_fresh_chat(owned.gateway, model, marker) for model, marker in on_second)
|
||||
owned.kill_worker(first)
|
||||
eventually(
|
||||
lambda: owned.serves_only_on_every_worker(stored, stored_id),
|
||||
bool,
|
||||
seconds=graceful_stop_seconds() + 60,
|
||||
)
|
||||
replacement: Final = next(iter(owned.worker_pids() - {second}))
|
||||
with owned.suspended(second):
|
||||
on_replacement: Final = tuple((model, _marker()) for model in (configured, stored))
|
||||
replacement_texts: Final = tuple(
|
||||
_fresh_chat(owned.gateway, model, marker) for model, marker in on_replacement
|
||||
)
|
||||
assert replacement not in (first, second)
|
||||
received: Final = vendor.outbound()
|
||||
sent: Final = (*on_first, *on_second, *on_replacement)
|
||||
assert (*first_texts, *second_texts, *replacement_texts) == tuple(_echo(marker) for _, marker in sent)
|
||||
assert len(received) == 6, received
|
||||
assert [_flags(_body(request)) for request in received] == [()] * 6, received
|
||||
|
||||
|
||||
def _burst(send: Send, model: str, size: int) -> tuple[str, ...]:
|
||||
def attempt(marker: str) -> str:
|
||||
try:
|
||||
return "ok" if send(model, marker) == _echo(marker) else "wrong text"
|
||||
except openai.APIError as error:
|
||||
return type(error).__name__
|
||||
|
||||
with ThreadPoolExecutor(max_workers=size) as pool:
|
||||
return tuple(pool.map(attempt, tuple(_marker() for _ in range(size))))
|
||||
|
||||
|
||||
def test_a_vendor_outage_between_bursts_reports_unhealthy_then_healthy_and_never_sends_the_flag(
|
||||
gateway: Gateway, vendor: _Vendor
|
||||
) -> None:
|
||||
with gateway.scenario() as scenario:
|
||||
flagged: Final = _deployment(scenario, {**_family("openai", vendor.wire), **_VISION})
|
||||
send: Final = _chat(gateway)
|
||||
before: Final = _burst(send, flagged, 10)
|
||||
healthy_before: Final = _health(gateway, flagged)
|
||||
vendor.outage.set()
|
||||
during: Final = _burst(send, flagged, 10)
|
||||
down: Final = _health(gateway, flagged)
|
||||
vendor.outage.clear()
|
||||
up: Final = _health(gateway, flagged)
|
||||
after: Final = _burst(send, flagged, 10)
|
||||
received: Final = vendor.outbound()
|
||||
assert before == after == ("ok",) * 10, (before, after)
|
||||
assert "ok" not in during, during
|
||||
assert (healthy_before, down, up) == ((200, 1, 0), (503, 0, 1), (200, 1, 0))
|
||||
assert [_flags(_body(request)) for request in received] == [()] * len(received), received
|
||||
|
|
@ -8,6 +8,7 @@ import asyncio, importlib, importlib.resources as importlib_get_model, json, lit
|
|||
import os
|
||||
import sys
|
||||
import threading
|
||||
from typing import Final
|
||||
|
||||
import pytest
|
||||
|
||||
|
|
@ -440,6 +441,22 @@ async def test_loaded_catalog_snapshot_follows_the_fetched_map_and_ignores_later
|
|||
assert GetModelCostMap.loaded_model_cost_map()["gpt-5.4-mini"]["max_input_tokens"] == 777
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_loaded_catalog_fields_follow_each_fetched_map():
|
||||
added = json.loads(_real_map_bytes())
|
||||
added["gpt-5.4-mini"]["supports_litellm_qa_capability"] = True
|
||||
with_field, _ = _mock_client([httpx.Response(200, content=json.dumps(added).encode())])
|
||||
without_field, _ = _mock_client([httpx.Response(200, content=_real_map_bytes())])
|
||||
|
||||
await refetch_model_cost_map(url=_URL, sleep=_SleepRecorder(), rng=random.Random(0), client=with_field)
|
||||
fields_after_adding: Final = GetModelCostMap.loaded_model_cost_map_fields()
|
||||
await refetch_model_cost_map(url=_URL, sleep=_SleepRecorder(), rng=random.Random(0), client=without_field)
|
||||
|
||||
assert "supports_litellm_qa_capability" in fields_after_adding
|
||||
assert "supports_litellm_qa_capability" not in GetModelCostMap.loaded_model_cost_map_fields()
|
||||
assert "max_input_tokens" in GetModelCostMap.loaded_model_cost_map_fields()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_refetch_revision_follows_the_bytes_not_the_url():
|
||||
edited = json.loads(_real_map_bytes())
|
||||
|
|
|
|||
|
|
@ -34,6 +34,22 @@ async def test_update_litellm_params_max_tokens_default(monkeypatch):
|
|||
assert updated_params["max_tokens"] == 16
|
||||
|
||||
|
||||
def test_update_litellm_params_for_health_check_drops_model_capability_flags():
|
||||
probe_params: Final = update_litellm_params_for_health_check(
|
||||
model_info={},
|
||||
litellm_params={
|
||||
"model": "bedrock/invoke/us.anthropic.claude-sonnet-4-5-20250929-v1:0",
|
||||
"aws_region_name": "us-west-2",
|
||||
"supports_max_reasoning_effort": False,
|
||||
"supports_system_message": False,
|
||||
},
|
||||
)
|
||||
|
||||
assert "supports_max_reasoning_effort" not in probe_params
|
||||
assert probe_params["aws_region_name"] == "us-west-2"
|
||||
assert probe_params["supports_system_message"] is False
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_update_litellm_params_max_tokens_custom():
|
||||
"""
|
||||
|
|
|
|||
|
|
@ -1,11 +1,14 @@
|
|||
import logging
|
||||
from itertools import chain
|
||||
from typing import Dict, Final, List, Optional, Union
|
||||
from unittest.mock import Mock
|
||||
|
||||
import pytest
|
||||
|
||||
import litellm
|
||||
from litellm import Router
|
||||
from litellm.constants import ROUTER_FALLBACK_ERROR_DETAIL_MAX_CHARS
|
||||
from litellm.litellm_core_utils.get_model_cost_map import GetModelCostMap
|
||||
from litellm.proxy._types import UserAPIKeyAuth
|
||||
from litellm.router_utils.common_utils import (
|
||||
_deployment_supports_web_search,
|
||||
|
|
@ -20,7 +23,8 @@ from litellm.router_utils.common_utils import (
|
|||
PROVIDER_SCOPED_CREDENTIAL_PARAMS,
|
||||
warn_on_provider_credential_mismatch,
|
||||
)
|
||||
from litellm.types.router import FallbackSkipReason, SkippedFallbackTarget
|
||||
from litellm.types.router import Deployment, FallbackSkipReason, LiteLLM_Params, ModelInfo, SkippedFallbackTarget
|
||||
from litellm.types.utils import ProviderSpecificModelInfo
|
||||
|
||||
|
||||
class TestFilterTeamBasedModels:
|
||||
|
|
@ -846,3 +850,97 @@ class TestWarnOnProviderCredentialMismatch:
|
|||
def test_provider_for_generic_call(litellm_params, expected, monkeypatch):
|
||||
monkeypatch.setenv("AZURE_AI_API_BASE", "https://unrelated.openai.azure.com")
|
||||
assert provider_for_generic_call(litellm_params) == expected
|
||||
|
||||
|
||||
def test_router_warns_only_for_a_capability_flag_set_under_litellm_params(caplog):
|
||||
bedrock_model: Final = "bedrock/invoke/us.anthropic.claude-sonnet-4-5-20250929-v1:0"
|
||||
with caplog.at_level(logging.WARNING, logger="LiteLLM Router"):
|
||||
Router(
|
||||
model_list=[
|
||||
{
|
||||
"model_name": "flag-in-litellm-params",
|
||||
"litellm_params": {"model": bedrock_model, "supports_max_reasoning_effort": False},
|
||||
},
|
||||
{
|
||||
"model_name": "flag-in-model-info",
|
||||
"litellm_params": {"model": bedrock_model},
|
||||
"model_info": {"supports_max_reasoning_effort": False},
|
||||
},
|
||||
{
|
||||
"model_name": "declared-litellm-param",
|
||||
"litellm_params": {"model": bedrock_model, "supports_system_message": False},
|
||||
},
|
||||
]
|
||||
)
|
||||
|
||||
capability_warnings: Final = [r.getMessage() for r in caplog.records if "model_info" in r.getMessage()]
|
||||
assert len(capability_warnings) == 1, capability_warnings
|
||||
assert "flag-in-litellm-params" in capability_warnings[0]
|
||||
assert "supports_max_reasoning_effort" in capability_warnings[0]
|
||||
|
||||
|
||||
def test_router_serves_deployments_without_model_capability_flags():
|
||||
capability_flags: Final = frozenset(ProviderSpecificModelInfo.__annotations__) | frozenset(
|
||||
key for key in chain.from_iterable(GetModelCostMap.loaded_model_cost_map().values()) if key.startswith("supports_")
|
||||
)
|
||||
declared_params: Final = {"supports_system_message": False, "supports_system_prompt": True}
|
||||
router: Final = Router(
|
||||
model_list=[
|
||||
{
|
||||
"model_name": "every-flag",
|
||||
"litellm_params": {
|
||||
"model": "replicate/meta/meta-llama-3-8b-instruct",
|
||||
**dict.fromkeys(capability_flags - declared_params.keys(), True),
|
||||
**declared_params,
|
||||
},
|
||||
}
|
||||
]
|
||||
)
|
||||
|
||||
served_params: Final = router.get_model_list()[0]["litellm_params"]
|
||||
assert sorted(capability_flags & served_params.keys() - declared_params.keys()) == []
|
||||
assert {key: served_params.get(key) for key in declared_params} == declared_params
|
||||
|
||||
|
||||
def test_router_keeps_a_provider_param_that_another_deployments_model_info_names():
|
||||
Router(
|
||||
model_list=[
|
||||
{
|
||||
"model_name": "registers-the-name",
|
||||
"litellm_params": {"model": "openai/litellm-qa-registered-capability"},
|
||||
"model_info": {"supports_system_prompt": True},
|
||||
}
|
||||
]
|
||||
)
|
||||
router: Final = Router(
|
||||
model_list=[
|
||||
{
|
||||
"model_name": "replicate-system-prompt",
|
||||
"litellm_params": {"model": "replicate/meta/meta-llama-3-8b-instruct", "supports_system_prompt": True},
|
||||
}
|
||||
]
|
||||
)
|
||||
|
||||
assert router.get_model_list()[0]["litellm_params"].get("supports_system_prompt") is True
|
||||
|
||||
|
||||
def test_router_warns_once_and_serves_no_flag_for_a_deployment_synced_from_the_db(caplog):
|
||||
router: Final = Router(model_list=[])
|
||||
db_deployment: Final = Deployment(
|
||||
model_name="db-model",
|
||||
litellm_params=LiteLLM_Params(
|
||||
model="bedrock/invoke/us.anthropic.claude-sonnet-4-5-20250929-v1:0",
|
||||
supports_max_reasoning_effort=False,
|
||||
),
|
||||
model_info=ModelInfo(id="db-model-id"),
|
||||
)
|
||||
with caplog.at_level(logging.WARNING, logger="LiteLLM Router"):
|
||||
first_sync: Final = router.upsert_deployment(db_deployment)
|
||||
second_sync: Final = router.upsert_deployment(db_deployment)
|
||||
|
||||
capability_warnings: Final = [r.getMessage() for r in caplog.records if "model_info" in r.getMessage()]
|
||||
assert first_sync is not None
|
||||
assert second_sync is None
|
||||
assert len(capability_warnings) == 1, capability_warnings
|
||||
assert "db-model" in capability_warnings[0]
|
||||
assert "supports_max_reasoning_effort" not in router.get_model_list()[0]["litellm_params"]
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue