mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-16 04:29:31 +00:00
fix(decisions): read a deployment's own supports_box_questions flag before the cost map (#46696)
* fix(decisions): read a deployment's own supports_box_questions flag before the cost map The box-question check only read the shared cost-map row and required a decisions mode there, so an OpenAI-format deployment whose model_info sets supports_box_questions was refused with a 400, which kept the wire test from #46633 red on main * fix(decisions): keep the box flag on its deployment and out of request bodies A model_info that is not a plain dict no longer fails the call, the proxy decisions endpoint drops a caller-sent model_info, and the router stops copying supports_box_questions onto the shared cost-map key. * fix(decisions): keep the box flag shared with same-model deployments like main Scoping supports_box_questions to one deployment turned box questions on an unflagged sibling of the same backend model from reaching the provider into a 400. Siblings inherit the flag through the shared cost-map row again, as they do on main. Also formats two new tests. * test(decisions): pin box flag inheritance by an unflagged same-model deployment * test(decisions): audit box-question deployment flag across providers, routes, and chaos cases
This commit is contained in:
Mateo Wang
GitHub
parent
fac0135d82
commit
1938f537bf
8 changed files with 1118 additions and 13 deletions
|
|
@ -5,7 +5,7 @@ from types import MappingProxyType
|
|||
from typing import Final, Literal, TypeAlias
|
||||
|
||||
import httpx
|
||||
from pydantic import ConfigDict, TypeAdapter, ValidationError
|
||||
from pydantic import BaseModel, ConfigDict, StrictBool, TypeAdapter, ValidationError
|
||||
from typing_extensions import assert_never
|
||||
|
||||
import litellm
|
||||
|
|
@ -290,6 +290,7 @@ def _unsupported_params(
|
|||
model: str,
|
||||
custom_llm_provider: str,
|
||||
headers: Mapping[str, str],
|
||||
supports_box_questions: bool,
|
||||
) -> tuple[str, ...]:
|
||||
refuses_safety_identifier: Final = (
|
||||
ir_request.safety_identifier is not None and not provider_config.supports_safety_identifier
|
||||
|
|
@ -298,9 +299,7 @@ def _unsupported_params(
|
|||
None if provider_config.supports_images_in(ir_request.input, headers) else image_param(ir_request.input)
|
||||
)
|
||||
refuses_assets: Final = bool(ir_request.assets) and not provider_config.supports_assets(headers)
|
||||
refuses_box_questions: Final = _has_box_question(ir_request.questions) and not (
|
||||
provider_config.supports_box_questions(model, custom_llm_provider, headers)
|
||||
)
|
||||
refuses_box_questions: Final = _has_box_question(ir_request.questions) and not supports_box_questions
|
||||
refuses_audio: Final = bool(ir_request.audio) and not provider_config.supports_audio_input(
|
||||
model, custom_llm_provider
|
||||
)
|
||||
|
|
@ -328,9 +327,8 @@ def _is_not_box_question(question: DecisionsIRQuestion) -> bool:
|
|||
def _without_unsupported_params(
|
||||
ir_request: DecisionsIRRequest,
|
||||
provider_config: BaseDecisionsConfig,
|
||||
model: str,
|
||||
custom_llm_provider: str,
|
||||
headers: Mapping[str, str],
|
||||
supports_box_questions: bool,
|
||||
) -> DecisionsIRRequest:
|
||||
return replace(
|
||||
ir_request,
|
||||
|
|
@ -343,7 +341,7 @@ def _without_unsupported_params(
|
|||
assets=ir_request.assets if provider_config.supports_assets(headers) else MappingProxyType({}),
|
||||
questions=(
|
||||
ir_request.questions
|
||||
if provider_config.supports_box_questions(model, custom_llm_provider, headers)
|
||||
if supports_box_questions
|
||||
else tuple(filter(_is_not_box_question, ir_request.questions))
|
||||
),
|
||||
)
|
||||
|
|
@ -365,6 +363,19 @@ def _without_malformed_images(
|
|||
return well_formed
|
||||
|
||||
|
||||
class _DeploymentBoxQuestions(BaseModel):
|
||||
model_config = ConfigDict(frozen=True, from_attributes=True)
|
||||
|
||||
supports_box_questions: StrictBool | None = None
|
||||
|
||||
|
||||
def _deployment_supports_box_questions(kwargs: Mapping[str, object]) -> bool | None:
|
||||
try:
|
||||
return _DeploymentBoxQuestions.model_validate(kwargs.get("model_info") or {}).supports_box_questions
|
||||
except ValidationError:
|
||||
return None
|
||||
|
||||
|
||||
def _provider_ir_request(
|
||||
ir_request: DecisionsIRRequest,
|
||||
*,
|
||||
|
|
@ -374,12 +385,17 @@ def _provider_ir_request(
|
|||
headers: Mapping[str, str],
|
||||
kwargs: Mapping[str, object],
|
||||
) -> DecisionsIRRequest | _UnsupportedParams | _NothingLeftAfterDrop:
|
||||
unsupported: Final = _unsupported_params(ir_request, provider_config, model, custom_llm_provider, headers)
|
||||
supports_box_questions: Final = provider_config.supports_box_questions(
|
||||
model, custom_llm_provider, headers, _deployment_supports_box_questions(kwargs)
|
||||
)
|
||||
unsupported: Final = _unsupported_params(
|
||||
ir_request, provider_config, model, custom_llm_provider, headers, supports_box_questions
|
||||
)
|
||||
if not unsupported:
|
||||
return ir_request
|
||||
if not _drops_params(kwargs):
|
||||
return _UnsupportedParams(names=unsupported)
|
||||
supported: Final = _without_unsupported_params(ir_request, provider_config, model, custom_llm_provider, headers)
|
||||
supported: Final = _without_unsupported_params(ir_request, provider_config, headers, supports_box_questions)
|
||||
if _leaves_no_input(supported.input):
|
||||
return _NothingLeftAfterDrop(names=unsupported, missing="input to decide on")
|
||||
if not supported.questions:
|
||||
|
|
|
|||
|
|
@ -537,7 +537,11 @@ class BaseDecisionsConfig(ABC):
|
|||
def supports_assets(self, headers: Mapping[str, str]) -> bool:
|
||||
return False
|
||||
|
||||
def supports_box_questions(self, model: str, custom_llm_provider: str, headers: Mapping[str, str]) -> bool:
|
||||
def supports_box_questions(
|
||||
self, model: str, custom_llm_provider: str, headers: Mapping[str, str], deployment_flag: bool | None
|
||||
) -> bool:
|
||||
if deployment_flag is not None:
|
||||
return deployment_flag
|
||||
return _is_decisions_model(model, custom_llm_provider) and litellm.supports_box_questions(
|
||||
model=model, custom_llm_provider=custom_llm_provider
|
||||
)
|
||||
|
|
|
|||
|
|
@ -94,8 +94,12 @@ class TypeSafeDecisionsConfig(BaseDecisionsConfig):
|
|||
name.lower() == _IMAGES_HEADER and value.strip().lower() == "true" for name, value in headers.items()
|
||||
)
|
||||
|
||||
def supports_box_questions(self, model: str, custom_llm_provider: str, headers: Mapping[str, str]) -> bool:
|
||||
return self.supports_assets(headers) and super().supports_box_questions(model, custom_llm_provider, headers)
|
||||
def supports_box_questions(
|
||||
self, model: str, custom_llm_provider: str, headers: Mapping[str, str], deployment_flag: bool | None
|
||||
) -> bool:
|
||||
return self.supports_assets(headers) and super().supports_box_questions(
|
||||
model, custom_llm_provider, headers, deployment_flag
|
||||
)
|
||||
|
||||
def transform_decisions_request(
|
||||
self,
|
||||
|
|
|
|||
|
|
@ -99,12 +99,16 @@ def _fields_checked_before_routing(data: Mapping[str, object]) -> Mapping[str, o
|
|||
return MappingProxyType({key: value for key, value in data.items() if key not in _FIELDS_CHECKED_AFTER_ROUTING})
|
||||
|
||||
|
||||
_DEPLOYMENT_ONLY_FIELDS: Final = frozenset({"model_info"})
|
||||
|
||||
|
||||
async def _request_data(request: Request, user_api_key_dict: UserAPIKeyAuth) -> dict[str, object]:
|
||||
body: Final = await request.body()
|
||||
try:
|
||||
return _REQUEST_DATA_ADAPTER.validate_json(body)
|
||||
data: Final = _REQUEST_DATA_ADAPTER.validate_json(body)
|
||||
except ValidationError as error:
|
||||
raise await _invalid_request(raw_data={}, error=error, user_api_key_dict=user_api_key_dict)
|
||||
return {key: value for key, value in data.items() if key not in _DEPLOYMENT_ONLY_FIELDS}
|
||||
|
||||
|
||||
async def _process_decisions(
|
||||
|
|
|
|||
|
|
@ -0,0 +1,302 @@
|
|||
import asyncio
|
||||
import json
|
||||
import threading
|
||||
import uuid
|
||||
from collections.abc import Mapping
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
from queue import SimpleQueue
|
||||
from types import MappingProxyType
|
||||
from typing import Final
|
||||
from urllib.parse import urlsplit
|
||||
|
||||
import httpx
|
||||
import psutil
|
||||
import pytest
|
||||
import yaml
|
||||
from integration._support.client import Gateway, eventually
|
||||
from integration._support.database import read_rows
|
||||
from integration._support.process import graceful_stop_seconds, owned_proxy_process
|
||||
from integration._support.wire import Reply, Request, Wire, wire_server
|
||||
from pydantic import JsonValue, TypeAdapter
|
||||
|
||||
_API_KEY: Final = "synthetic-decisions-key"
|
||||
_JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue])
|
||||
_BOX: Final[dict[str, JsonValue]] = {"top": 0.1, "bottom": 0.4, "left": 0.2, "right": 0.6}
|
||||
_BOX_QUESTION: Final[dict[str, JsonValue]] = {"type": "box", "name": "crack", "instructions": "Where is the crack?"}
|
||||
_SYSTEM_ONE_BOX_QUESTION: Final[dict[str, JsonValue]] = {"type": "box", "instructions": "Where is the crack?"}
|
||||
_SYSTEM_ONE_BOX_ANSWER: Final[dict[str, JsonValue]] = {"type": "box", "box": _BOX, "confidence": 0.7}
|
||||
_BOX_ANSWER: Final[dict[str, JsonValue]] = {"type": "box", "name": "crack", "box": _BOX, "confidence": 0.7}
|
||||
_REFUSED_BOX: Final = "does not support parameters: ['box questions']"
|
||||
_OPENAI_FORMAT_ROUTES: Final = ("/v1/decisions", "/decisions")
|
||||
_ROUTES: Final = (*_OPENAI_FORMAT_ROUTES, "/v1/systemone", "/systemone")
|
||||
_FLAGS: Final = ("true", "false", "unset")
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class _Call:
|
||||
route: str
|
||||
model: str
|
||||
marker: str
|
||||
answers: bool
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class _Served:
|
||||
call: _Call
|
||||
status: int
|
||||
call_id: str
|
||||
text: str
|
||||
|
||||
|
||||
def _marker_of(request: Request) -> str:
|
||||
state: Final = _JSON_OBJECT.validate_json(request.body)["state"]
|
||||
assert isinstance(state, str), request.body
|
||||
return state
|
||||
|
||||
|
||||
def _reply(request: Request) -> Reply:
|
||||
marker: Final = _marker_of(request)
|
||||
answer: Final = {
|
||||
"model": f"model-{marker}",
|
||||
"answers": {"crack": _SYSTEM_ONE_BOX_ANSWER},
|
||||
"usage": {"input_tokens": 12, "output_tokens": 1},
|
||||
}
|
||||
return Reply(body=json.dumps(answer).encode())
|
||||
|
||||
|
||||
def _body(call: _Call) -> dict[str, JsonValue]:
|
||||
if call.route in _OPENAI_FORMAT_ROUTES:
|
||||
return {"model": call.model, "input": call.marker, "questions": [_BOX_QUESTION], "num_retries": 0}
|
||||
return {
|
||||
"model": call.model,
|
||||
"state": call.marker,
|
||||
"questions": {"crack": _SYSTEM_ONE_BOX_QUESTION},
|
||||
"num_retries": 0,
|
||||
}
|
||||
|
||||
|
||||
async def _send(client: httpx.AsyncClient, key: str, call: _Call) -> _Served:
|
||||
response: Final = await client.post(call.route, json=_body(call), headers={"Authorization": f"Bearer {key}"})
|
||||
return _Served(call, response.status_code, response.headers.get("x-litellm-call-id", ""), response.text)
|
||||
|
||||
|
||||
async def _burst(
|
||||
base_url: str, key: str, calls: tuple[_Call, ...], *, tolerate_transport_errors: bool = False
|
||||
) -> tuple[_Served, ...]:
|
||||
async with httpx.AsyncClient(base_url=base_url, timeout=60, trust_env=False) as client:
|
||||
results: Final = await asyncio.gather(
|
||||
*(_send(client, key, call) for call in calls), return_exceptions=tolerate_transport_errors
|
||||
)
|
||||
assert all(not isinstance(result, BaseException) or isinstance(result, httpx.TransportError) for result in results)
|
||||
return tuple(result for result in results if isinstance(result, _Served))
|
||||
|
||||
|
||||
def _assert_served_its_own(served: _Served) -> None:
|
||||
if not served.call.answers:
|
||||
assert served.status == 400, (served.status, served.text)
|
||||
assert _REFUSED_BOX in served.text, served.text
|
||||
return
|
||||
assert served.status == 200, (served.status, served.text)
|
||||
body: Final = _JSON_OBJECT.validate_json(served.text)
|
||||
assert body["model"] == f"model-{served.call.marker}", served.text
|
||||
expected: Final[JsonValue] = (
|
||||
[_BOX_ANSWER] if served.call.route in _OPENAI_FORMAT_ROUTES else {"crack": _SYSTEM_ONE_BOX_ANSWER}
|
||||
)
|
||||
assert body["answers"] == expected, served.text
|
||||
|
||||
|
||||
def _statuses_by_call_id(call_ids: tuple[str, ...]) -> dict[str, JsonValue]:
|
||||
placeholders: Final = ", ".join("%s" for _ in call_ids)
|
||||
rows: Final = eventually(
|
||||
lambda: read_rows(
|
||||
f'SELECT request_id, status FROM "LiteLLM_SpendLogs" WHERE request_id IN ({placeholders})', call_ids
|
||||
),
|
||||
lambda found: len(found) >= len(call_ids),
|
||||
seconds=70,
|
||||
)
|
||||
assert len(rows) == len(call_ids), rows
|
||||
return {str(row["request_id"]): row["status"] for row in rows}
|
||||
|
||||
|
||||
def _answered(served: tuple[_Served, ...]) -> tuple[_Served, ...]:
|
||||
return tuple(item for item in served if item.call.answers)
|
||||
|
||||
|
||||
def _assert_each_served_its_own_and_logged_once(served: tuple[_Served, ...]) -> None:
|
||||
for item in served:
|
||||
_assert_served_its_own(item)
|
||||
answered: Final = _answered(served)
|
||||
assert len({item.call_id for item in answered}) == len(answered), answered
|
||||
assert set(_statuses_by_call_id(tuple(item.call_id for item in answered)).values()) == {"success"}
|
||||
|
||||
|
||||
def _backend() -> str:
|
||||
return f"perplexity/box-{uuid.uuid4().hex[:12]}"
|
||||
|
||||
|
||||
def _calls(models: Mapping[str, str], count: int) -> tuple[_Call, ...]:
|
||||
return tuple(
|
||||
_Call(
|
||||
route=_ROUTES[index % len(_ROUTES)],
|
||||
model=models[_FLAGS[(index // len(_ROUTES)) % len(_FLAGS)]],
|
||||
marker=f"{_FLAGS[(index // len(_ROUTES)) % len(_FLAGS)]}-{uuid.uuid4().hex}",
|
||||
answers=_FLAGS[(index // len(_ROUTES)) % len(_FLAGS)] == "true",
|
||||
)
|
||||
for index in range(count)
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.timeout(180)
|
||||
async def test_a_burst_over_every_route_and_flag_sends_boxes_only_for_flagged_deployments_and_logs_each_once(
|
||||
gateway: Gateway,
|
||||
) -> None:
|
||||
with wire_server(_reply) as wire, gateway.scenario() as scenario:
|
||||
flags: Final[Mapping[str, Mapping[str, JsonValue] | None]] = MappingProxyType(
|
||||
{"true": {"supports_box_questions": True}, "false": {"supports_box_questions": False}, "unset": None}
|
||||
)
|
||||
models: Final = MappingProxyType(
|
||||
{
|
||||
flag: scenario.model(model=_backend(), api_base=wire.url, api_key=_API_KEY, model_info=model_info)
|
||||
for flag, model_info in flags.items()
|
||||
}
|
||||
)
|
||||
calls: Final = _calls(models, 36)
|
||||
served: Final = await _burst(str(gateway.client.base_url), gateway.key, calls)
|
||||
assert len(served) == 36
|
||||
for item in served:
|
||||
_assert_served_its_own(item)
|
||||
answered: Final = _answered(served)
|
||||
assert len(answered) == 12
|
||||
assert len({item.call_id for item in answered}) == 12
|
||||
assert set(_statuses_by_call_id(tuple(item.call_id for item in answered)).values()) == {"success"}
|
||||
received: Final = wire.drain()
|
||||
assert sorted(_marker_of(request) for request in received) == sorted(item.call.marker for item in answered)
|
||||
assert {request.target for request in received} == {"/v1/decisions"}, received
|
||||
|
||||
|
||||
@pytest.mark.timeout(180)
|
||||
async def test_a_flagged_deployment_recovers_box_answers_after_an_upstream_outage_on_the_same_port(
|
||||
gateway: Gateway,
|
||||
) -> None:
|
||||
base_url: Final = str(gateway.client.base_url)
|
||||
with gateway.scenario() as scenario:
|
||||
with wire_server(_reply) as wire:
|
||||
port: Final = urlsplit(wire.url).port
|
||||
assert port is not None
|
||||
model: Final = scenario.model(
|
||||
model=_backend(), api_base=wire.url, api_key=_API_KEY, model_info={"supports_box_questions": True}
|
||||
)
|
||||
before: Final = await _burst(base_url, gateway.key, _calls(MappingProxyType({"true": model}), 4))
|
||||
during: Final = await _burst(base_url, gateway.key, _calls(MappingProxyType({"true": model}), 4))
|
||||
assert [item.status for item in during] == [500] * 4, [item.text for item in during]
|
||||
with wire_server(_reply, port=port) as recovered:
|
||||
after: Final = await _burst(base_url, gateway.key, _calls(MappingProxyType({"true": model}), 4))
|
||||
assert sorted(map(_marker_of, recovered.drain())) == sorted(item.call.marker for item in after)
|
||||
for item in (*before, *after):
|
||||
_assert_served_its_own(item)
|
||||
statuses: Final = _statuses_by_call_id(tuple(item.call_id for item in (*before, *during, *after)))
|
||||
assert statuses == {
|
||||
**{item.call_id: "success" for item in (*before, *after)},
|
||||
**{item.call_id: "failure" for item in during},
|
||||
}
|
||||
|
||||
|
||||
_YAML_FLAGS: Final[tuple[tuple[str, str, Mapping[str, JsonValue] | None, bool], ...]] = (
|
||||
("box-yaml-true", "own", {"supports_box_questions": True}, True),
|
||||
("box-yaml-false-on-true-row", "shared", {"supports_box_questions": False}, False),
|
||||
("box-yaml-true-row", "shared", {"supports_box_questions": True, "mode": "evaluation"}, True),
|
||||
("box-yaml-unset-on-true-row", "shared", None, True),
|
||||
("box-yaml-string", "string", {"supports_box_questions": "true"}, False),
|
||||
)
|
||||
|
||||
|
||||
def _yaml_config(wire: Wire, tmp_path: Path, backends: Mapping[str, str]) -> Path:
|
||||
base: Final = _JSON_OBJECT.validate_python(yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text()))
|
||||
model_list: Final = [
|
||||
{
|
||||
"model_name": name,
|
||||
"litellm_params": {"model": backends[backend], "api_base": wire.url, "api_key": _API_KEY},
|
||||
**({} if model_info is None else {"model_info": dict(model_info)}),
|
||||
}
|
||||
for name, backend, model_info, _ in _YAML_FLAGS
|
||||
]
|
||||
path: Final = tmp_path / "decisions-box-flags.yaml"
|
||||
path.write_text(yaml.safe_dump({**base, "model_list": model_list}))
|
||||
return path
|
||||
|
||||
|
||||
def _yaml_calls(per_model: int) -> tuple[_Call, ...]:
|
||||
return tuple(
|
||||
_Call(
|
||||
route=_ROUTES[index % len(_ROUTES)],
|
||||
model=_YAML_FLAGS[index % len(_YAML_FLAGS)][0],
|
||||
marker=f"{_YAML_FLAGS[index % len(_YAML_FLAGS)][0]}-{uuid.uuid4().hex}",
|
||||
answers=_YAML_FLAGS[index % len(_YAML_FLAGS)][3],
|
||||
)
|
||||
for index in range(per_model * len(_YAML_FLAGS))
|
||||
)
|
||||
|
||||
|
||||
def _open_upstream_connections(pid: int, upstream: str) -> int:
|
||||
port: Final = urlsplit(upstream).port
|
||||
return sum(
|
||||
1
|
||||
for connection in psutil.Process(pid).net_connections(kind="tcp")
|
||||
if connection.status == psutil.CONN_ESTABLISHED and connection.raddr and connection.raddr.port == port
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.timeout(2 * graceful_stop_seconds() + 120)
|
||||
async def test_yaml_box_flags_survive_a_worker_sigkill_mid_burst_and_hold_on_respawned_workers(
|
||||
gateway: Gateway, tmp_path: Path
|
||||
) -> None:
|
||||
release: Final = threading.Event()
|
||||
held_markers: Final[SimpleQueue[str]] = SimpleQueue()
|
||||
backends: Final = MappingProxyType({"own": _backend(), "shared": _backend(), "string": _backend()})
|
||||
|
||||
def held(request: Request) -> Reply:
|
||||
marker: Final = _marker_of(request)
|
||||
if marker.startswith("held-"):
|
||||
held_markers.put(marker)
|
||||
assert release.wait(timeout=60), "The burst was never released"
|
||||
return _reply(request)
|
||||
|
||||
with wire_server(held) as wire:
|
||||
config: Final = _yaml_config(wire, tmp_path, backends)
|
||||
with owned_proxy_process(gateway, tmp_path, {}, config=config, workers=2) as owned:
|
||||
base_url: Final = str(owned.gateway.client.base_url)
|
||||
key: Final = owned.gateway.key
|
||||
first: Final = await _burst(base_url, key, _yaml_calls(2))
|
||||
_assert_each_served_its_own_and_logged_once(first)
|
||||
workers: Final = owned.worker_pids()
|
||||
assert len(workers) == 2, workers
|
||||
calls: Final = tuple(
|
||||
_Call(
|
||||
route=_ROUTES[index % len(_ROUTES)],
|
||||
model="box-yaml-true",
|
||||
marker=f"held-{uuid.uuid4().hex}",
|
||||
answers=True,
|
||||
)
|
||||
for index in range(20)
|
||||
)
|
||||
burst: Final = asyncio.create_task(_burst(base_url, key, calls, tolerate_transport_errors=True))
|
||||
await asyncio.to_thread(eventually, held_markers.qsize, lambda size: size == 20, 60)
|
||||
held_by: Final = MappingProxyType({pid: _open_upstream_connections(pid, wire.url) for pid in workers})
|
||||
assert sum(held_by.values()) == 20, held_by
|
||||
victim_pid, survivor_pid = sorted(workers, key=held_by.__getitem__)
|
||||
with owned.replacing_worker(victim_pid) as kill:
|
||||
kill.freeze()
|
||||
await asyncio.to_thread(kill)
|
||||
release.set()
|
||||
survived: Final = await burst
|
||||
assert len(survived) == held_by[survivor_pid], (held_by, len(survived))
|
||||
await asyncio.to_thread(_assert_each_served_its_own_and_logged_once, survived)
|
||||
await asyncio.to_thread(owned.kill_worker, survivor_pid)
|
||||
assert not workers & owned.worker_pids(), (workers, owned.worker_pids())
|
||||
respawned: Final = await _burst(base_url, key, _yaml_calls(4))
|
||||
await asyncio.to_thread(_assert_each_served_its_own_and_logged_once, respawned)
|
||||
received: Final = tuple(map(_marker_of, wire.drain()))
|
||||
assert sorted(received) == sorted(
|
||||
(*(call.marker for call in calls), *(item.call.marker for item in _answered((*first, *respawned))))
|
||||
)
|
||||
|
|
@ -0,0 +1,607 @@
|
|||
import json
|
||||
import uuid
|
||||
from collections.abc import Mapping, Sequence
|
||||
from dataclasses import dataclass, field
|
||||
from pathlib import Path
|
||||
from queue import SimpleQueue
|
||||
from types import MappingProxyType
|
||||
from typing import Final, Literal
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
from integration._support.client import Gateway, Scenario, eventually, object_value, string_value
|
||||
from integration._support.database import read_rows
|
||||
from integration._support.process import configless_proxy_process, graceful_stop_seconds
|
||||
from integration._support.wire import Reply, Request, Wire, wire_server
|
||||
from pydantic import JsonValue, TypeAdapter
|
||||
|
||||
import litellm
|
||||
from litellm.constants import PROXY_CONFIG_RELOAD_INTERVAL_SECONDS
|
||||
|
||||
_API_KEY: Final = "synthetic-decisions-key"
|
||||
_IMAGES_HEADER: Final = "TypeSafe-Danger-Zone-Images"
|
||||
_JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue])
|
||||
_JSON_LIST: Final = TypeAdapter(list[JsonValue])
|
||||
_INPUT: Final = "Photo description: a phone screen with a diagonal crack from the top-left corner to the middle."
|
||||
_BOX: Final[dict[str, JsonValue]] = {"top": 0.1, "bottom": 0.4, "left": 0.2, "right": 0.6}
|
||||
_PREDICATE: Final[dict[str, JsonValue]] = {
|
||||
"type": "predicate",
|
||||
"name": "cracked",
|
||||
"instructions": "Is the screen cracked?",
|
||||
}
|
||||
_BOX_QUESTION: Final[dict[str, JsonValue]] = {"type": "box", "name": "crack", "instructions": "Where is the crack?"}
|
||||
_QUESTIONS: Final[tuple[dict[str, JsonValue], ...]] = (_PREDICATE, _BOX_QUESTION)
|
||||
_PREDICATE_ANSWER: Final[dict[str, JsonValue]] = {"type": "predicate", "name": "cracked", "probability": 0.9}
|
||||
_BOX_ANSWER: Final[dict[str, JsonValue]] = {"type": "box", "name": "crack", "box": _BOX, "confidence": 0.7}
|
||||
_BOX_REFUSAL: Final[dict[str, JsonValue]] = {"type": "refusal", "name": "crack"}
|
||||
_SYSTEM_ONE_QUESTIONS: Final[dict[str, JsonValue]] = {
|
||||
"cracked": {"type": "noul", "instructions": "Is the screen cracked?"},
|
||||
"crack": {"type": "box", "instructions": "Where is the crack?"},
|
||||
}
|
||||
_SYSTEM_ONE_ANSWERS: Final[dict[str, JsonValue]] = {
|
||||
"cracked": {"type": "noul", "noul": 0.9},
|
||||
"crack": {"type": "box", "box": _BOX, "confidence": 0.7},
|
||||
}
|
||||
_USAGE: Final[dict[str, JsonValue]] = {"input_tokens": 41, "output_tokens": 2}
|
||||
_OPENAI_USAGE: Final[dict[str, JsonValue]] = {
|
||||
"input_tokens": 41,
|
||||
"input_tokens_details": {"cached_tokens": 0, "cache_write_tokens": 0},
|
||||
"output_tokens": 2,
|
||||
"output_tokens_details": {"reasoning_tokens": 0},
|
||||
"total_tokens": 43,
|
||||
}
|
||||
_REFUSED_BOX: Final = "does not support parameters: ['box questions']"
|
||||
_SPEND_QUERY: Final = (
|
||||
'SELECT status, call_type, custom_llm_provider, model_group, model_id FROM "LiteLLM_SpendLogs" '
|
||||
"WHERE request_id = %s"
|
||||
)
|
||||
_OPENAI_FORMAT_ROUTES: Final = ("/v1/decisions", "/decisions")
|
||||
_SYSTEM_ONE_ROUTES: Final = ("/v1/systemone", "/systemone")
|
||||
|
||||
Envelope = Literal["plain", "result", "ai_decide"]
|
||||
|
||||
|
||||
def _systemone_question(question: Mapping[str, JsonValue]) -> dict[str, JsonValue]:
|
||||
kind: Final = "box" if question["type"] == "box" else "noul"
|
||||
return {"type": kind, "instructions": question["instructions"]}
|
||||
|
||||
|
||||
def _systemone_answer(question: JsonValue, *, probability_key: str) -> JsonValue:
|
||||
if object_value(question)["type"] == "box":
|
||||
return {"type": "box", "box": _BOX, "confidence": 0.7}
|
||||
return {"type": "noul", probability_key: 0.9}
|
||||
|
||||
|
||||
def _openai_answer(question: JsonValue) -> JsonValue:
|
||||
fields: Final = object_value(question)
|
||||
template: Final = _BOX_ANSWER if fields["type"] == "box" else _PREDICATE_ANSWER
|
||||
return {**template, "name": fields["name"]}
|
||||
|
||||
|
||||
def _expected_answers(questions: Sequence[Mapping[str, JsonValue]]) -> list[JsonValue]:
|
||||
return [_openai_answer(dict(question)) for question in questions]
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class _Member:
|
||||
provider: str
|
||||
model: str
|
||||
path: str
|
||||
body_model: str | None
|
||||
api_key: str | None = _API_KEY
|
||||
speaks_openai: bool = False
|
||||
envelope: Envelope = "plain"
|
||||
headers: Mapping[str, str] = field(default_factory=lambda: MappingProxyType({}))
|
||||
|
||||
def deployment(
|
||||
self, scenario: Scenario, api_base: str, model_info: Mapping[str, JsonValue] | None, **extra: JsonValue
|
||||
) -> str:
|
||||
header_params: Final[dict[str, JsonValue]] = {"extra_headers": dict(self.headers)} if self.headers else {}
|
||||
return scenario.model(
|
||||
model=self.model,
|
||||
api_base=api_base,
|
||||
api_key=self.api_key,
|
||||
model_info=model_info,
|
||||
**header_params,
|
||||
**extra,
|
||||
)
|
||||
|
||||
def upstream_body(self, questions: Sequence[Mapping[str, JsonValue]]) -> dict[str, JsonValue]:
|
||||
if self.speaks_openai:
|
||||
return {"model": self.body_model, "input": _INPUT, "questions": [dict(question) for question in questions]}
|
||||
keyed: Final[dict[str, JsonValue]] = {
|
||||
string_value(question["name"]): _systemone_question(question) for question in questions
|
||||
}
|
||||
model_field: Final[dict[str, JsonValue]] = {} if self.body_model is None else {"model": self.body_model}
|
||||
return {**model_field, "state": _INPUT, "questions": keyed}
|
||||
|
||||
def reply(self, request: Request) -> Reply:
|
||||
body: Final = _JSON_OBJECT.validate_json(request.body)
|
||||
if self.speaks_openai:
|
||||
questions: Final = _JSON_LIST.validate_python(body["questions"])
|
||||
answer: Final = {
|
||||
"model": self.body_model,
|
||||
"answers": list(map(_openai_answer, questions)),
|
||||
"usage": _OPENAI_USAGE,
|
||||
}
|
||||
return Reply(body=json.dumps(answer).encode())
|
||||
probability_key: Final = "probability" if self.envelope == "ai_decide" else "noul"
|
||||
answers: Final = {
|
||||
key: _systemone_answer(question, probability_key=probability_key)
|
||||
for key, question in object_value(body["questions"]).items()
|
||||
}
|
||||
return Reply(body=json.dumps(self._enveloped(answers)).encode())
|
||||
|
||||
def _enveloped(self, answers: Mapping[str, JsonValue]) -> dict[str, JsonValue]:
|
||||
plain: Final[dict[str, JsonValue]] = {"model": self.body_model, "answers": dict(answers), "usage": _USAGE}
|
||||
if self.envelope == "result":
|
||||
return {"result": plain, "success": True}
|
||||
if self.envelope == "ai_decide":
|
||||
return {"response": {"answers": dict(answers)}, "metadata": {"version": "1.0"}}
|
||||
return plain
|
||||
|
||||
|
||||
def _members(tag: str) -> Mapping[str, _Member]:
|
||||
name: Final = f"box-{tag}"
|
||||
return MappingProxyType(
|
||||
{
|
||||
"perplexity": _Member("perplexity", f"perplexity/{name}", "/v1/decisions", name),
|
||||
"typesafe": _Member(
|
||||
"typesafe",
|
||||
f"typesafe/{name}",
|
||||
"/v1/systemone",
|
||||
name,
|
||||
headers=MappingProxyType({_IMAGES_HEADER: "true"}),
|
||||
),
|
||||
"openrouter": _Member("openrouter", f"openrouter/typesafe/{name}", "/alpha/decisions", f"typesafe/{name}"),
|
||||
"fireworks_ai": _Member(
|
||||
"fireworks_ai",
|
||||
f"fireworks_ai/accounts/fireworks/routers/{name}",
|
||||
"/v1/systemone",
|
||||
f"accounts/fireworks/routers/{name}",
|
||||
),
|
||||
"cloudflare": _Member(
|
||||
"cloudflare", f"cloudflare/{name}", f"/ai/run/@cf/cloudflare/{name}", name, envelope="result"
|
||||
),
|
||||
"strands_decider": _Member("strands_decider", f"strands_decider/{name}", "/v1/systemone", name, None),
|
||||
"databricks": _Member("databricks", f"databricks/{name}", f"/{name}/invocations", name),
|
||||
"databricks_ai_decide": _Member(
|
||||
"databricks", "databricks/ai_decide", "/api/2.0/ai-functions/ai-decide", None, envelope="ai_decide"
|
||||
),
|
||||
"hosted_vllm": _Member("hosted_vllm", f"hosted_vllm/{name}", "/v1/systemone", name, None),
|
||||
"laya": _Member("laya", "laya/typed-decisions", "/v1/systemone", "typed-decisions", None),
|
||||
"bespoke": _Member(
|
||||
"bespoke",
|
||||
"bespoke/bespokelabs/Bespoke-Nimble-9B",
|
||||
"/v1/systemone",
|
||||
"bespokelabs/Bespoke-Nimble-9B",
|
||||
None,
|
||||
),
|
||||
"azure_ai": _Member("azure_ai", f"azure_ai/{name}", "/providers/microsoft/v1/systemone", name),
|
||||
"openai": _Member("openai", f"openai/{name}", "/v1/decisions", name, speaks_openai=True),
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
_MEMBER_LABELS: Final = tuple(_members("labels"))
|
||||
|
||||
|
||||
def _tag() -> str:
|
||||
return uuid.uuid4().hex[:12]
|
||||
|
||||
|
||||
def _perplexity() -> _Member:
|
||||
return _members(_tag())["perplexity"]
|
||||
|
||||
|
||||
def _decide(
|
||||
gateway: Gateway,
|
||||
model: str,
|
||||
questions: Sequence[Mapping[str, JsonValue]] = _QUESTIONS,
|
||||
*,
|
||||
route: str = "/v1/decisions",
|
||||
**extra: JsonValue,
|
||||
) -> httpx.Response:
|
||||
body: Final[dict[str, JsonValue]] = {
|
||||
"model": model,
|
||||
"input": _INPUT,
|
||||
"questions": [dict(question) for question in questions],
|
||||
"num_retries": 0,
|
||||
**extra,
|
||||
}
|
||||
return gateway.request("POST", route, body)
|
||||
|
||||
|
||||
def _decide_in_system_one_format(gateway: Gateway, model: str, route: str, **extra: JsonValue) -> httpx.Response:
|
||||
body: Final[dict[str, JsonValue]] = {
|
||||
"model": model,
|
||||
"state": _INPUT,
|
||||
"questions": _SYSTEM_ONE_QUESTIONS,
|
||||
"num_retries": 0,
|
||||
**extra,
|
||||
}
|
||||
return gateway.request("POST", route, body)
|
||||
|
||||
|
||||
def _spend_row(call_id: str) -> dict[str, JsonValue]:
|
||||
rows: Final = eventually(lambda: read_rows(_SPEND_QUERY, (call_id,)), lambda found: len(found) == 1, seconds=70)
|
||||
return rows[0]
|
||||
|
||||
|
||||
def _assert_refused_box(response: httpx.Response) -> None:
|
||||
assert response.status_code == 400, response.text
|
||||
assert _REFUSED_BOX in response.text, response.text
|
||||
|
||||
|
||||
def _authorization(member: _Member) -> str:
|
||||
return f"Bearer {member.api_key}" if member.api_key else ""
|
||||
|
||||
|
||||
@pytest.mark.parametrize("label", _MEMBER_LABELS)
|
||||
def test_a_flagged_deployment_of_every_decisions_provider_sends_box_questions_and_returns_the_box(
|
||||
gateway: Gateway, label: str
|
||||
) -> None:
|
||||
member: Final = _members(_tag())[label]
|
||||
with wire_server(member.reply) as wire, gateway.scenario() as scenario:
|
||||
model: Final = member.deployment(scenario, wire.url, {"supports_box_questions": True})
|
||||
response: Final = _decide(gateway, model)
|
||||
assert response.status_code == 200, response.text
|
||||
assert response.json()["answers"] == _expected_answers(_QUESTIONS), response.text
|
||||
(call,) = wire.drain()
|
||||
assert (call.method, call.target) == ("POST", member.path)
|
||||
assert call.headers.get("authorization", "") == _authorization(member)
|
||||
assert _JSON_OBJECT.validate_json(call.body) == member.upstream_body(_QUESTIONS)
|
||||
assert all(call.headers.get(name.lower()) == value for name, value in member.headers.items()), call.headers
|
||||
row: Final = _spend_row(response.headers["x-litellm-call-id"])
|
||||
assert (row["status"], row["call_type"], row["custom_llm_provider"], row["model_group"]) == (
|
||||
"success",
|
||||
"adecisions",
|
||||
member.provider,
|
||||
model,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("route", (*_OPENAI_FORMAT_ROUTES, *_SYSTEM_ONE_ROUTES))
|
||||
def test_every_route_answers_box_questions_for_a_flagged_deployment_and_refuses_them_for_an_unflagged_one(
|
||||
gateway: Gateway, route: str
|
||||
) -> None:
|
||||
flagged_member: Final = _perplexity()
|
||||
plain_member: Final = _perplexity()
|
||||
send: Final = (
|
||||
(lambda model: _decide(gateway, model, route=route))
|
||||
if route in _OPENAI_FORMAT_ROUTES
|
||||
else (lambda model: _decide_in_system_one_format(gateway, model, route))
|
||||
)
|
||||
with wire_server(flagged_member.reply) as wire, gateway.scenario() as scenario:
|
||||
flagged: Final = flagged_member.deployment(scenario, wire.url, {"supports_box_questions": True})
|
||||
plain: Final = plain_member.deployment(scenario, wire.url, None)
|
||||
refused: Final = send(plain)
|
||||
_assert_refused_box(refused)
|
||||
assert wire.drain() == ()
|
||||
answered: Final = send(flagged)
|
||||
assert answered.status_code == 200, answered.text
|
||||
expected: Final[JsonValue] = (
|
||||
_expected_answers(_QUESTIONS) if route in _OPENAI_FORMAT_ROUTES else _SYSTEM_ONE_ANSWERS
|
||||
)
|
||||
assert answered.json()["answers"] == expected, answered.text
|
||||
(call,) = wire.drain()
|
||||
assert call.target == flagged_member.path
|
||||
assert _JSON_OBJECT.validate_json(call.body) == flagged_member.upstream_body(_QUESTIONS)
|
||||
|
||||
|
||||
def test_repeated_identical_box_requests_each_reach_the_upstream_and_each_log_once(gateway: Gateway) -> None:
|
||||
member: Final = _perplexity()
|
||||
with wire_server(member.reply) as wire, gateway.scenario() as scenario:
|
||||
model: Final = member.deployment(scenario, wire.url, {"supports_box_questions": True})
|
||||
responses: Final = tuple(_decide(gateway, model, (_BOX_QUESTION,)) for _ in range(2))
|
||||
assert [response.status_code for response in responses] == [200, 200], [r.text for r in responses]
|
||||
call_ids: Final = tuple(response.headers["x-litellm-call-id"] for response in responses)
|
||||
assert len(set(call_ids)) == 2, call_ids
|
||||
assert [_JSON_OBJECT.validate_json(call.body) for call in wire.drain()] == [
|
||||
member.upstream_body((_BOX_QUESTION,))
|
||||
] * 2
|
||||
assert [_spend_row(call_id)["status"] for call_id in call_ids] == ["success", "success"]
|
||||
|
||||
|
||||
async def test_the_sdk_honors_a_model_info_flag_passed_by_its_caller_on_sync_and_async_calls() -> None:
|
||||
member: Final = _perplexity()
|
||||
with wire_server(member.reply) as wire:
|
||||
common: Final[dict[str, object]] = {
|
||||
"model": member.model,
|
||||
"input": _INPUT,
|
||||
"questions": [dict(question) for question in _QUESTIONS],
|
||||
"api_base": wire.url,
|
||||
"api_key": _API_KEY,
|
||||
}
|
||||
synchronous: Final = litellm.decisions(**common, model_info={"supports_box_questions": True})
|
||||
asynchronous: Final = await litellm.adecisions(**common, model_info={"supports_box_questions": True})
|
||||
for response in (synchronous, asynchronous):
|
||||
assert response.model_dump(mode="json")["answers"] == _expected_answers(_QUESTIONS)
|
||||
with pytest.raises(litellm.UnsupportedParamsError, match=r"\['box questions'\]"):
|
||||
await litellm.adecisions(**common)
|
||||
with pytest.raises(litellm.UnsupportedParamsError, match=r"\['box questions'\]"):
|
||||
litellm.decisions(**common, model_info={"supports_box_questions": "true"})
|
||||
assert [_JSON_OBJECT.validate_json(call.body) for call in wire.drain()] == [
|
||||
member.upstream_body(_QUESTIONS)
|
||||
] * 2
|
||||
|
||||
|
||||
async def test_a_router_deployment_flagged_false_refuses_boxes_even_when_its_backend_row_says_true() -> None:
|
||||
member: Final = _perplexity()
|
||||
with wire_server(member.reply) as wire:
|
||||
params: Final[dict[str, object]] = {"model": member.model, "api_base": wire.url, "api_key": _API_KEY}
|
||||
router: Final = litellm.Router(
|
||||
model_list=[
|
||||
{"model_name": "box-off", "litellm_params": params, "model_info": {"supports_box_questions": False}},
|
||||
{
|
||||
"model_name": "box-on",
|
||||
"litellm_params": params,
|
||||
"model_info": {"supports_box_questions": True, "mode": "evaluation"},
|
||||
},
|
||||
{"model_name": "box-unset", "litellm_params": params},
|
||||
]
|
||||
)
|
||||
common: Final[dict[str, object]] = {"input": _INPUT, "questions": [dict(question) for question in _QUESTIONS]}
|
||||
with pytest.raises(litellm.UnsupportedParamsError, match=r"\['box questions'\]"):
|
||||
await router.adecisions(model="box-off", **common)
|
||||
assert wire.drain() == ()
|
||||
answered: Final = [await router.adecisions(model=model, **common) for model in ("box-on", "box-unset")]
|
||||
assert [response.model_dump(mode="json")["answers"] for response in answered] == [
|
||||
_expected_answers(_QUESTIONS)
|
||||
] * 2
|
||||
assert [_JSON_OBJECT.validate_json(call.body) for call in wire.drain()] == [
|
||||
member.upstream_body(_QUESTIONS)
|
||||
] * 2
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"value", ("true", 1, [True], "", "x" * 5000, None), ids=("string", "int", "list", "empty", "5kb", "null")
|
||||
)
|
||||
def test_a_non_boolean_deployment_flag_counts_as_no_flag(gateway: Gateway, value: JsonValue) -> None:
|
||||
member: Final = _perplexity()
|
||||
with wire_server(member.reply) as wire, gateway.scenario() as scenario:
|
||||
model: Final = member.deployment(scenario, wire.url, {"supports_box_questions": value})
|
||||
_assert_refused_box(_decide(gateway, model))
|
||||
assert wire.drain() == ()
|
||||
predicate_only: Final = _decide(gateway, model, (_PREDICATE,))
|
||||
assert predicate_only.status_code == 200, predicate_only.text
|
||||
assert predicate_only.json()["answers"] == [_PREDICATE_ANSWER]
|
||||
(call,) = wire.drain()
|
||||
assert _JSON_OBJECT.validate_json(call.body) == member.upstream_body((_PREDICATE,))
|
||||
|
||||
|
||||
_BODY_MODEL_INFO: Final[tuple[tuple[str, JsonValue], ...]] = (
|
||||
("flag", {"supports_box_questions": True}),
|
||||
("int", 5),
|
||||
("list", [{"supports_box_questions": True}]),
|
||||
("empty", ""),
|
||||
("5kb", "x" * 5000),
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("route", ("/v1/decisions", "/v1/systemone"))
|
||||
@pytest.mark.parametrize(
|
||||
"value", [value for _, value in _BODY_MODEL_INFO], ids=[label for label, _ in _BODY_MODEL_INFO]
|
||||
)
|
||||
def test_a_body_model_info_never_turns_box_questions_on_and_never_reaches_the_provider(
|
||||
gateway: Gateway, route: str, value: JsonValue
|
||||
) -> None:
|
||||
member: Final = _perplexity()
|
||||
with wire_server(member.reply) as wire, gateway.scenario() as scenario:
|
||||
model: Final = member.deployment(scenario, wire.url, None)
|
||||
box: Final = (
|
||||
_decide(gateway, model, route=route, model_info=value)
|
||||
if route == "/v1/decisions"
|
||||
else _decide_in_system_one_format(gateway, model, route, model_info=value)
|
||||
)
|
||||
_assert_refused_box(box)
|
||||
assert wire.drain() == ()
|
||||
predicate_only: Final = _decide(gateway, model, (_PREDICATE,), model_info=value)
|
||||
assert predicate_only.status_code == 200, predicate_only.text
|
||||
(call,) = wire.drain()
|
||||
assert _JSON_OBJECT.validate_json(call.body) == member.upstream_body((_PREDICATE,))
|
||||
|
||||
|
||||
def test_a_duplicated_or_metadata_model_info_and_a_forged_id_change_nothing_behind_the_router(
|
||||
gateway: Gateway,
|
||||
) -> None:
|
||||
member: Final = _perplexity()
|
||||
forged: Final = f"forged-{uuid.uuid4().hex}"
|
||||
with wire_server(member.reply) as wire, gateway.scenario() as scenario:
|
||||
model: Final = member.deployment(scenario, wire.url, None)
|
||||
first: Final = json.dumps(
|
||||
{
|
||||
"model": model,
|
||||
"input": _INPUT,
|
||||
"questions": list(_QUESTIONS),
|
||||
"num_retries": 0,
|
||||
"model_info": {"supports_box_questions": False},
|
||||
}
|
||||
)
|
||||
duplicated: Final = (
|
||||
f'{first[:-1]}, "model_info": {json.dumps({"supports_box_questions": True, "id": forged})}}}'
|
||||
)
|
||||
raw: Final = gateway.client.post(
|
||||
"/v1/decisions",
|
||||
content=duplicated.encode(),
|
||||
headers={"Authorization": f"Bearer {gateway.key}", "Content-Type": "application/json"},
|
||||
)
|
||||
_assert_refused_box(raw)
|
||||
_assert_refused_box(_decide(gateway, model, metadata={"model_info": {"supports_box_questions": True}}))
|
||||
assert wire.drain() == ()
|
||||
answered: Final = _decide(gateway, model, (_PREDICATE,), model_info={"id": forged})
|
||||
assert answered.status_code == 200, answered.text
|
||||
assert answered.headers.get("x-litellm-model-id", "") not in ("", forged), answered.headers
|
||||
assert _spend_row(answered.headers["x-litellm-call-id"])["model_id"] == answered.headers["x-litellm-model-id"]
|
||||
(call,) = wire.drain()
|
||||
assert _JSON_OBJECT.validate_json(call.body) == member.upstream_body((_PREDICATE,))
|
||||
|
||||
|
||||
@pytest.mark.parametrize("status", (401, 404, 413))
|
||||
def test_a_flagged_box_request_brings_the_provider_error_back_after_one_attempt(gateway: Gateway, status: int) -> None:
|
||||
member: Final = _perplexity()
|
||||
marker: Final = f"provider says {status} {uuid.uuid4().hex}"
|
||||
|
||||
def failing(_request: Request) -> Reply:
|
||||
return Reply(status=status, body=json.dumps({"error": {"message": marker}}).encode())
|
||||
|
||||
with wire_server(failing) as wire, gateway.scenario() as scenario:
|
||||
model: Final = member.deployment(scenario, wire.url, {"supports_box_questions": True})
|
||||
response: Final = _decide(gateway, model)
|
||||
assert response.status_code == status, response.text
|
||||
assert marker in response.text, response.text
|
||||
(call,) = wire.drain()
|
||||
assert _JSON_OBJECT.validate_json(call.body) == member.upstream_body(_QUESTIONS)
|
||||
|
||||
|
||||
def test_a_box_question_answered_with_the_wrong_answer_type_reaches_the_caller_as_a_refusal(gateway: Gateway) -> None:
|
||||
member: Final = _perplexity()
|
||||
|
||||
def mismatched(_request: Request) -> Reply:
|
||||
answers: Final = {"cracked": {"type": "noul", "noul": 0.9}, "crack": {"type": "noul", "noul": 0.4}}
|
||||
return Reply(body=json.dumps({"model": member.body_model, "answers": answers, "usage": _USAGE}).encode())
|
||||
|
||||
with wire_server(mismatched) as wire, gateway.scenario() as scenario:
|
||||
model: Final = member.deployment(scenario, wire.url, {"supports_box_questions": True})
|
||||
response: Final = _decide(gateway, model)
|
||||
assert response.status_code == 200, response.text
|
||||
assert response.json()["answers"] == [_PREDICATE_ANSWER, _BOX_REFUSAL], response.text
|
||||
assert len(wire.drain()) == 1
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("header", "answers"), ((None, False), ("false", False), ("true", True)), ids=("absent", "false", "true")
|
||||
)
|
||||
def test_typesafe_needs_its_images_header_for_boxes_even_on_a_flagged_deployment(
|
||||
gateway: Gateway, header: str | None, answers: bool
|
||||
) -> None:
|
||||
template: Final = _members(_tag())["typesafe"]
|
||||
member: Final = _Member(
|
||||
template.provider,
|
||||
template.model,
|
||||
template.path,
|
||||
template.body_model,
|
||||
headers=MappingProxyType({} if header is None else {_IMAGES_HEADER: header}),
|
||||
)
|
||||
with wire_server(member.reply) as wire, gateway.scenario() as scenario:
|
||||
model: Final = member.deployment(scenario, wire.url, {"supports_box_questions": True})
|
||||
response: Final = _decide(gateway, model)
|
||||
if not answers:
|
||||
_assert_refused_box(response)
|
||||
assert wire.drain() == ()
|
||||
return
|
||||
assert response.status_code == 200, response.text
|
||||
assert response.json()["answers"] == _expected_answers(_QUESTIONS)
|
||||
(call,) = wire.drain()
|
||||
assert call.headers.get(_IMAGES_HEADER.lower()) == "true"
|
||||
assert _JSON_OBJECT.validate_json(call.body) == member.upstream_body(_QUESTIONS)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("flag", "keeps_box"), ((True, True), (False, False), (None, False)), ids=("true", "false", "unset")
|
||||
)
|
||||
def test_drop_params_drops_box_questions_unless_the_deployment_flag_is_true(
|
||||
gateway: Gateway, flag: bool | None, keeps_box: bool
|
||||
) -> None:
|
||||
member: Final = _perplexity()
|
||||
with wire_server(member.reply) as wire, gateway.scenario() as scenario:
|
||||
model_info: Final[dict[str, JsonValue] | None] = None if flag is None else {"supports_box_questions": flag}
|
||||
model: Final = member.deployment(scenario, wire.url, model_info, drop_params=True)
|
||||
response: Final = _decide(gateway, model)
|
||||
assert response.status_code == 200, response.text
|
||||
sent: Final = _QUESTIONS if keeps_box else (_PREDICATE,)
|
||||
answers: Final = _expected_answers(_QUESTIONS) if keeps_box else [_PREDICATE_ANSWER, _BOX_REFUSAL]
|
||||
assert response.json()["answers"] == answers, response.text
|
||||
(call,) = wire.drain()
|
||||
assert _JSON_OBJECT.validate_json(call.body) == member.upstream_body(sent)
|
||||
if keeps_box:
|
||||
return
|
||||
nothing_left: Final = _decide(gateway, model, (_BOX_QUESTION,))
|
||||
assert nothing_left.status_code == 400, nothing_left.text
|
||||
assert "leaves no questions to answer" in nothing_left.text, nothing_left.text
|
||||
assert wire.drain() == ()
|
||||
|
||||
|
||||
@pytest.mark.timeout(2 * graceful_stop_seconds() + 120)
|
||||
def test_a_proxy_without_a_router_ignores_a_body_model_info_for_box_questions_and_model_ids(
|
||||
gateway: Gateway, tmp_path: Path
|
||||
) -> None:
|
||||
member: Final = _perplexity()
|
||||
forged: Final = f"forged-{uuid.uuid4().hex}"
|
||||
body_model_info: Final[dict[str, JsonValue]] = {"supports_box_questions": True, "id": forged}
|
||||
with wire_server(member.reply) as wire:
|
||||
with configless_proxy_process(
|
||||
gateway, tmp_path, {"PERPLEXITYAI_API_KEY": _API_KEY}, model=member.model, api_base=wire.url
|
||||
) as owned:
|
||||
box: Final = _decide(owned.gateway, member.model, model_info=body_model_info)
|
||||
predicate_only: Final = _decide(owned.gateway, member.model, (_PREDICATE,), model_info=body_model_info)
|
||||
system_one: Final = owned.gateway.request(
|
||||
"POST",
|
||||
"/v1/systemone",
|
||||
{
|
||||
"model": member.model,
|
||||
"state": _INPUT,
|
||||
"questions": {"cracked": _systemone_question(_PREDICATE)},
|
||||
"model_info": body_model_info,
|
||||
},
|
||||
)
|
||||
_assert_refused_box(box)
|
||||
for answered in (predicate_only, system_one):
|
||||
assert answered.status_code == 200, answered.text
|
||||
assert answered.headers.get("x-litellm-model-id", "") != forged, answered.headers
|
||||
assert _spend_row(answered.headers["x-litellm-call-id"])["model_id"] != forged
|
||||
assert [_JSON_OBJECT.validate_json(call.body) for call in wire.drain()] == [
|
||||
member.upstream_body((_PREDICATE,))
|
||||
] * 2
|
||||
|
||||
|
||||
def _new_deployment(gateway: Gateway, scenario: Scenario, member: _Member, wire: Wire) -> tuple[str, str]:
|
||||
name: Final = f"integration-{uuid.uuid4().hex}"
|
||||
created: Final = gateway.post(
|
||||
"/model/new",
|
||||
{
|
||||
"model_name": name,
|
||||
"litellm_params": {"model": member.model, "api_base": wire.url, "api_key": _API_KEY},
|
||||
"model_info": {"supports_box_questions": True},
|
||||
},
|
||||
)
|
||||
identity: Final = string_value(object_value(created["model_info"])["id"])
|
||||
scenario.cleanups.callback(scenario.delete_model, identity)
|
||||
return name, identity
|
||||
|
||||
|
||||
def _box_outcome(gateway: Gateway, model: str, outcomes: SimpleQueue[int]) -> int:
|
||||
response: Final = _decide(gateway, model, (_BOX_QUESTION,))
|
||||
outcomes.put(response.status_code)
|
||||
assert response.status_code in (200, 400), response.text
|
||||
if response.status_code == 400:
|
||||
assert _REFUSED_BOX in response.text, response.text
|
||||
return 400
|
||||
assert response.json()["answers"] == [_BOX_ANSWER], response.text
|
||||
return 200
|
||||
|
||||
|
||||
@pytest.mark.timeout(240)
|
||||
def test_turning_the_flag_off_under_traffic_converges_without_sending_boxes_to_a_refusing_deployment(
|
||||
gateway: Gateway,
|
||||
) -> None:
|
||||
member: Final = _perplexity()
|
||||
outcomes: Final[SimpleQueue[int]] = SimpleQueue()
|
||||
with wire_server(member.reply) as wire, gateway.scenario() as scenario:
|
||||
model, identity = _new_deployment(gateway, scenario, member, wire)
|
||||
before: Final = tuple(_box_outcome(gateway, model, outcomes) for _ in range(6))
|
||||
assert before == (200,) * 6, before
|
||||
patched: Final = gateway.request(
|
||||
"PATCH", f"/model/{identity}/update", {"model_info": {"id": identity, "supports_box_questions": False}}
|
||||
)
|
||||
assert patched.status_code == 200, patched.text
|
||||
during: Final = eventually(
|
||||
lambda: tuple(_box_outcome(gateway, model, outcomes) for _ in range(8)),
|
||||
lambda outcomes: outcomes == (400,) * 8,
|
||||
seconds=3 * PROXY_CONFIG_RELOAD_INTERVAL_SECONDS,
|
||||
)
|
||||
assert during == (400,) * 8
|
||||
received: Final = wire.drain()
|
||||
statuses: Final = tuple(outcomes.get_nowait() for _ in range(outcomes.qsize()))
|
||||
assert all(_JSON_OBJECT.validate_json(call.body) == member.upstream_body((_BOX_QUESTION,)) for call in received)
|
||||
assert len(received) == statuses.count(200), statuses
|
||||
|
|
@ -2,6 +2,7 @@ from __future__ import annotations
|
|||
|
||||
import asyncio
|
||||
import base64
|
||||
import copy
|
||||
import json
|
||||
import re
|
||||
import threading
|
||||
|
|
@ -33,6 +34,7 @@ from litellm.types.decisions import (
|
|||
OpenAIDecisionResponse,
|
||||
ScoreAnswer,
|
||||
)
|
||||
from litellm.types.router import ModelInfo
|
||||
from litellm.types.utils import LlmProviders
|
||||
from litellm.utils import ProviderConfigManager
|
||||
|
||||
|
|
@ -3188,6 +3190,147 @@ async def test_box_questions_are_dropped_under_drop_params_and_answered_as_refus
|
|||
assert response.model_dump(mode="json")["answers"] == _DROPPED_BOX_ANSWERS[api]
|
||||
|
||||
|
||||
_CRACK_BOX_QUESTION: Final[Mapping[str, object]] = MappingProxyType(
|
||||
{"type": "box", "name": "crack", "instructions": "the crack on the screen"}
|
||||
)
|
||||
_CRACK_BOX_ANSWER: Final[Mapping[str, object]] = MappingProxyType(
|
||||
{"type": "box", "name": "crack", "box": {"top": 0.1, "bottom": 0.4, "left": 0.2, "right": 0.6}, "confidence": 0.7}
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_deployment_that_advertises_box_questions_gets_box_answers_from_an_openai_format_provider(
|
||||
respx_mock: respx.MockRouter,
|
||||
) -> None:
|
||||
route: Final = respx_mock.post("https://api.openai.com/v1/decisions").respond(
|
||||
json={**_OPENAI_RESPONSE, "answers": [dict(_CRACK_BOX_ANSWER)]}
|
||||
)
|
||||
|
||||
response: Final = await litellm.adecisions(
|
||||
model="openai/gpt-6-luna",
|
||||
input="review",
|
||||
questions=[_CRACK_BOX_QUESTION],
|
||||
api_key="deployment-key",
|
||||
model_info={"supports_box_questions": True},
|
||||
)
|
||||
|
||||
assert json.loads(route.calls.last.request.content)["questions"] == [dict(_CRACK_BOX_QUESTION)]
|
||||
assert [answer.model_dump(mode="json") for answer in response.answers] == [dict(_CRACK_BOX_ANSWER)]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_deployment_that_turns_box_questions_off_is_refused_even_when_the_cost_map_allows_them(
|
||||
respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch
|
||||
) -> None:
|
||||
monkeypatch.setattr(litellm, "drop_params", False)
|
||||
route: Final = respx_mock.post("https://api.typesafe.ai/v1/systemone").respond(json=_RESPONSE)
|
||||
|
||||
with pytest.raises(litellm.UnsupportedParamsError, match=r"\['box questions'\]"):
|
||||
await litellm.adecisions(
|
||||
model="typesafe/jev-1.14.0",
|
||||
input="review",
|
||||
questions=[_CRACK_BOX_QUESTION],
|
||||
api_key="deployment-key",
|
||||
extra_headers={"TypeSafe-Danger-Zone-Images": "true"},
|
||||
model_info={"supports_box_questions": False},
|
||||
)
|
||||
|
||||
assert not route.called
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_typesafe_deployment_that_advertises_box_questions_still_needs_the_images_header(
|
||||
respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch
|
||||
) -> None:
|
||||
monkeypatch.setattr(litellm, "drop_params", False)
|
||||
route: Final = respx_mock.post("https://api.typesafe.ai/v1/systemone").respond(json=_RESPONSE)
|
||||
|
||||
with pytest.raises(litellm.UnsupportedParamsError, match=r"\['box questions'\]"):
|
||||
await litellm.adecisions(
|
||||
model="typesafe/jev-1.14.0",
|
||||
input="review",
|
||||
questions=[_CRACK_BOX_QUESTION],
|
||||
api_key="deployment-key",
|
||||
model_info={"supports_box_questions": True},
|
||||
)
|
||||
|
||||
assert not route.called
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_deployment_flag_that_is_not_a_boolean_leaves_the_cost_map_in_charge(
|
||||
respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch
|
||||
) -> None:
|
||||
monkeypatch.setattr(litellm, "drop_params", False)
|
||||
route: Final = respx_mock.post("https://api.openai.com/v1/decisions").respond(json=_OPENAI_RESPONSE)
|
||||
|
||||
with pytest.raises(litellm.UnsupportedParamsError, match=r"\['box questions'\]"):
|
||||
await litellm.adecisions(
|
||||
model="openai/gpt-6-luna",
|
||||
input="review",
|
||||
questions=[_CRACK_BOX_QUESTION],
|
||||
api_key="deployment-key",
|
||||
model_info={"supports_box_questions": "true"},
|
||||
)
|
||||
|
||||
assert not route.called
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_an_unflagged_deployment_keeps_box_questions_a_same_model_deployment_advertises(
|
||||
respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch
|
||||
) -> None:
|
||||
monkeypatch.setattr(litellm, "drop_params", False)
|
||||
backend: Final = "perplexity/pplx-decider-v1-27b"
|
||||
monkeypatch.setitem(litellm.model_cost, backend, copy.deepcopy(litellm.model_cost[backend]))
|
||||
for deployment_id in ("box-flagged", "box-sibling"):
|
||||
monkeypatch.setitem(litellm.model_cost, deployment_id, {})
|
||||
route: Final = respx_mock.post(_PERPLEXITY_DECISIONS_URL).respond(
|
||||
json={
|
||||
"model": "pplx-decider-v1-27b",
|
||||
"answers": {"crack": {k: v for k, v in _CRACK_BOX_ANSWER.items() if k != "name"}},
|
||||
"usage": {"input_tokens": _INPUT_TOKENS, "output_tokens": _OUTPUT_TOKENS},
|
||||
}
|
||||
)
|
||||
router: Final = litellm.Router(
|
||||
model_list=[
|
||||
{
|
||||
"model_name": "flagged",
|
||||
"litellm_params": {"model": backend, "api_key": "k"},
|
||||
"model_info": {"id": "box-flagged", "supports_box_questions": True},
|
||||
},
|
||||
{
|
||||
"model_name": "sibling",
|
||||
"litellm_params": {"model": backend, "api_key": "k"},
|
||||
"model_info": {"id": "box-sibling"},
|
||||
},
|
||||
]
|
||||
)
|
||||
|
||||
response: Final = await router.adecisions(model="sibling", input="review", questions=[_CRACK_BOX_QUESTION])
|
||||
|
||||
assert "crack" in json.loads(route.calls.last.request.content)["questions"]
|
||||
assert [answer.model_dump(mode="json") for answer in response.answers] == [dict(_CRACK_BOX_ANSWER)]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_model_info_object_instead_of_a_dict_still_reaches_the_provider(
|
||||
respx_mock: respx.MockRouter,
|
||||
) -> None:
|
||||
route: Final = respx_mock.post("https://api.openai.com/v1/decisions").respond(json=_OPENAI_RESPONSE)
|
||||
|
||||
response: Final = await litellm.adecisions(
|
||||
model="openai/gpt-6-luna",
|
||||
input="review",
|
||||
questions=[{"type": "predicate", "name": "is_defect", "instructions": "Is this a defect?"}],
|
||||
api_key="deployment-key",
|
||||
model_info=ModelInfo(id="deployment-1"),
|
||||
)
|
||||
|
||||
assert route.called
|
||||
assert response.answers[0].model_dump(mode="json")["name"] == "is_defect"
|
||||
|
||||
|
||||
def test_dropping_every_question_answers_400_before_http(respx_mock: respx.MockRouter) -> None:
|
||||
route: Final = respx_mock.post(_PERPLEXITY_DECISIONS_URL).respond(json=_RESPONSE)
|
||||
|
||||
|
|
|
|||
|
|
@ -422,6 +422,31 @@ def test_proxy_decisions_without_model_uses_the_proxy_default_model(
|
|||
assert json.loads(upstream.calls[0].request.content)["model"] == "pplx-decider-v1-27b"
|
||||
|
||||
|
||||
def test_a_request_body_cannot_turn_on_box_questions_for_a_proxy_without_a_router(
|
||||
client: TestClient,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
respx_mock: respx.MockRouter,
|
||||
) -> None:
|
||||
monkeypatch.setattr(litellm, "drop_params", False)
|
||||
monkeypatch.setattr(litellm.proxy.proxy_server, "llm_router", None)
|
||||
monkeypatch.setattr(litellm.proxy.proxy_server, "user_model", "openai/gpt-6-luna")
|
||||
upstream: Final = respx_mock.post("https://fake-openai.example/v1/decisions").respond(json={})
|
||||
|
||||
response: Final = client.post(
|
||||
"/v1/decisions",
|
||||
json={
|
||||
"model": "openai/gpt-6-luna",
|
||||
"input": "A phone screen with a diagonal crack.",
|
||||
"questions": [{"type": "box", "name": "crack", "instructions": "Where is the crack?"}],
|
||||
"model_info": {"supports_box_questions": True},
|
||||
},
|
||||
)
|
||||
|
||||
assert response.status_code == 400, response.text
|
||||
assert "['box questions']" in response.json()["error"]["message"]
|
||||
assert not upstream.called
|
||||
|
||||
|
||||
_OPENAI_FORMAT_REQUEST: Final[Mapping[str, object]] = {
|
||||
"model": "decider",
|
||||
"input": [
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue