mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-17 04:34:25 +00:00
18831b818c
* 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
716 lines
32 KiB
Python
716 lines
32 KiB
Python
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
|