scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-16 04:29:31 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 139
litellm/tests/integration/providers/test_databricks_decisions_wire.py
devin-ai-integration[bot] 7efce7f8e0
test(integration): give chaos and timing tests a harness API and gate raw signals (#46031) 
* test(integration): give chaos and timing tests a harness API and gate raw signals

Chaos tests killed and suspended proxy workers with raw os.kill and psutil
calls and bounded elapsed time with fixed literals, so they raced the uvicorn
worker healthcheck and failed on slow CI boxes

OwnedProxy now offers kill_worker and replacing_worker, which return only once
a replacement answers readiness with its own pid, and suspended, which is
refused unless the proxy was launched with a lenient worker healthcheck.
timing.py bounds durations against a measured control or the injected delay.
run.py measures the proxy import time once and derives readiness deadlines and
the per-test timeout from it, writing boot-calibration.json

check_test_quality.py adds TQ010 (raw process signals) and TQ011 (elapsed time
asserted under a numeric literal) for tests/integration outside _support, and
every call site is migrated so both report zero with no suppressions

* test(integration): steer the assistants wire burst onto both workers

The mixed-status burst assumed the OS would spread 48 connections across both workers and failed when one worker took them all. Each half is now sent while the other worker is suspended, and the test asserts an exact 24 and 24 split

* test(integration): restore product-enforced timing bounds lost in the migration

Some TQ011 migrations dropped or loosened a bound whose time is the behaviour
under test. The websocket session cap only kept its floor, so a cap closing at
300 s would have passed, and several timeouts were bounded by the injected stall
instead of the configured product value

assert_enforced_within bounds an elapsed time by the configured value plus
explicit slack, with an optional floor. The websocket cap, the MCP validation
timeout, the guardrail litellm_params timeout, the prisma command timeouts, the
baseline projection lock timeout and the key metadata statement timeout use it.
The device code refusal keeps its 15 s margin, and the langfuse retried prompt
fetch is now bounded by a cold fetch of a steady prompt as well as by the
ignored retry-after

* fix(tests): wait for the supervisor to reap killed workers and fix real-kill migrations

The lenient worker healthcheck was the whole case budget, so uvicorn's supervisor sat
in a ping on the killed worker for up to that long before respawning it. It now uses
the graceful stop budget, and a worker kill returns only after the victim is reaped,
so pooled connections no longer race the dead worker's socket.

Several tests on main picked the multiprocessing resource tracker as their victim and
never killed a worker. Now that they kill a real one, presidio counts streams cut by
the death, delete_model waits for the config reload that removes a deleted model from
a replacement worker, and the wedged coordination Redis test kills before the wedge
and waits for every parked publish to fail before resuming Redis.

The redis cluster test imports the harness through the tests package so it collects
in the unit job, and the per-model fallback rig exposes the owned proxy it now needs.

* test(observability): only require S3 deliveries the killed worker could not have buffered

A real SIGKILL drops the victim's unflushed batch, so requests it already answered can't be expected to land. The test now requires every request served after the replacement to land, some of the surviving worker's burst share to land, and nothing unsent to land

---------

Co-authored-by: Yuneng Jiang <yuneng@berri.ai>
2026-10-12 05:27:33 -07:00

657 lines
31 KiB
Python
Raw Blame History

import threading
import time
import uuid
from collections.abc import Mapping, Sequence
from pathlib import Path
from typing import Final
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 graceful_stop_seconds, owned_proxy_process
from integration._support.timing import assert_within_factor_of_control
from integration._support.upstream import ScenarioHandle, delete_scenario, register_scenario
from integration.cost_calculation.cost_tracking_case import JsonResponse
from pydantic import JsonValue
_MODEL: Final = "ai_decide"
_ROUTE: Final = "/api/2.0/ai-functions/ai-decide"
_ENDPOINT: Final = "databricks-openjev-qwen35-4b"
_API_KEY: Final = "synthetic-databricks-key"
_ENV_KEY: Final = "synthetic-databricks-env-key"
_ENV_TOKEN: Final = "synthetic-databricks-env-token"
_STATE: Final[dict[str, JsonValue]] = {"ticket": "The export job hangs at 99%", "component": "billing"}
_TIER_CRITERIA: Final[dict[str, JsonValue]] = {"SIMPLE": "a lookup", "MEDIUM": "some work", "COMPLEX": "deep work"}
_QUESTIONS: Final[dict[str, JsonValue]] = {"tier": {"type": "choice", "criteria": _TIER_CRITERIA}}
_OPENAI_INPUT: Final = "The export job hangs at 99% in billing"
_OPENAI_TIER_QUESTION: Final[dict[str, JsonValue]] = {
"type": "choice",
"name": "tier",
"instructions": "Which tier?",
"choices": [{"value": name, "description": description} for name, description in _TIER_CRITERIA.items()],
}
_ANSWERS: Final[dict[str, JsonValue]] = {
"tier": {
"type": "choice",
"choice": "COMPLEX",
"confidence": 0.91,
"probabilities": {"SIMPLE": 0.03, "MEDIUM": 0.06, "COMPLEX": 0.91},
}
}
_USAGE: Final[dict[str, JsonValue]] = {"input_tokens": 367, "output_tokens": 3}
_ANSWER_BODY: Final[dict[str, JsonValue]] = {"model": _ENDPOINT, "answers": _ANSWERS, "usage": _USAGE}
_METADATA: Final[dict[str, JsonValue]] = {"version": "1.0"}
_AI_DECIDE_BODY: Final[dict[str, JsonValue]] = {"response": {"answers": _ANSWERS}, "metadata": _METADATA}
_SYSTEM_ONE_BODY: Final[dict[str, JsonValue]] = {
"model": None,
"answers": _ANSWERS,
"usage": {"input_tokens": 0, "output_tokens": 0},
"metadata": _METADATA,
}
_MODEL_RULE: Final = (
"Databricks decisions model must be 'ai_decide' (the ai_decide AI Function) or a bare serving endpoint name "
"(letters, digits, '-', '_' and '.', with no '/', '?', '#', spaces, or a leading '.'), e.g. "
"databricks-openjev-qwen35-4b"
)
_MISSING_KEY: Final = "Missing API key for Decisions provider 'databricks'"
_NON_BARE_NAMES: Final[tuple[tuple[str, str], ...]] = (
("slash", "foo/bar"),
("space", "a b"),
("leading-dot", ".hidden"),
("question-mark", "a?b"),
("hash", "a#b"),
("route-path", "serving-endpoints/ai_decide"),
)
_DATABRICKS_ERRORS: Final[tuple[tuple[str, str, int, dict[str, JsonValue]], ...]] = (
(_MODEL, _ROUTE, 400, {"error_code": "BAD_REQUEST", "message": "Invalid input: questions must not be empty"}),
(_MODEL, _ROUTE, 403, {"error_code": "PERMISSION_DENIED", "message": "ai_decide is not enabled for this workspace"}),
(
_ENDPOINT,
f"/{_ENDPOINT}/invocations",
404,
{"error_code": "RESOURCE_DOES_NOT_EXIST", "message": f"Endpoint with name '{_ENDPOINT}' does not exist."},
),
)
_SPEND_QUERY: Final = (
'SELECT spend, status, call_type, model_group, custom_llm_provider, api_base FROM "LiteLLM_SpendLogs" '
"WHERE request_id = %s"
)
_OWNED_PROXY_BUDGET: Final = int(2 * graceful_stop_seconds() + 120)
_MEGABYTE_NAME: Final = "a" * 1_000_000 + "/"
_TENTH_MEGABYTE_NAME: Final = "a" * 100_000 + "/"
def _number(value: JsonValue) -> float:
assert isinstance(value, (int, float)) and not isinstance(value, bool), value
return float(value)
def _register(scenario: Scenario, body: dict[str, JsonValue], *, status: int = 200) -> ScenarioHandle:
handle: Final = register_scenario(
f"databricks-{uuid.uuid4().hex[:12]}", JsonResponse(content_type="application/json", body=body, status=status)
)
scenario.cleanups.callback(delete_scenario, handle)
return handle
def _decide(gateway: Gateway, model: str, *, key: str | None = None) -> httpx.Response:
return gateway.request("POST", "/v1/systemone", {"model": model, "state": _STATE, "questions": _QUESTIONS}, key=key)
def _observed(gateway: Gateway) -> tuple[dict[str, JsonValue], ...]:
with httpx.Client(base_url=gateway.upstream_url, timeout=5, trust_env=False) as upstream:
return tuple(map(object_value, upstream.get("/__observations").json()["requests"]))
def _calls_to(requests: Sequence[dict[str, JsonValue]], handle: ScenarioHandle) -> list[dict[str, JsonValue]]:
return [request for request in requests if string_value(request["path"]).startswith(f"/{handle.scenario_id}/")]
def _upstream_calls(gateway: Gateway, handle: ScenarioHandle) -> list[dict[str, JsonValue]]:
return _calls_to(_observed(gateway), handle)
def _timed_decide_under_liveliness_probes(
gateway: Gateway, model: str
) -> tuple[httpx.Response, float, tuple[tuple[int, float], ...]]:
liveliness: Final[list[tuple[int, float]]] = []
stop: Final = threading.Event()
def poll() -> None:
while not stop.is_set():
started: Final = time.perf_counter()
probe: Final = gateway.client.get("/health/liveliness")
liveliness.append((probe.status_code, time.perf_counter() - started))
poller: Final = threading.Thread(target=poll)
poller.start()
try:
started: Final = time.perf_counter()
response: Final = _decide(gateway, model)
elapsed: Final = time.perf_counter() - started
finally:
stop.set()
poller.join()
return response, elapsed, tuple(liveliness)
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 _wildcard_deployment(scenario: Scenario, **parameters: JsonValue) -> str:
prefix: Final = f"databricks-{uuid.uuid4().hex[:8]}"
created: Final = scenario.gateway.post(
"/model/new",
{"model_name": f"{prefix}/*", "litellm_params": {"model": "databricks/*", **parameters}, "model_info": {}},
)
scenario.cleanups.callback(scenario.delete_model, string_value(object_value(created["model_info"])["id"]))
return prefix
def _router_config(classifier: Mapping[str, JsonValue], simple: str, complex_: str) -> dict[str, JsonValue]:
return {
"classifier_type": "oss_classifier",
"opensource_classifier_config": dict(classifier),
"tiers": {"SIMPLE": simple, "MEDIUM": simple, "COMPLEX": complex_, "REASONING": complex_},
}
def _environment_classifier() -> dict[str, JsonValue]:
return {"provider": "databricks", "model": _MODEL, "timeout_ms": 20000, "circuit_breaker_enabled": False}
def _router(scenario: Scenario, classifier: Mapping[str, JsonValue]) -> str:
simple: Final = scenario.model(model="openai/simple-tier")
complex_: Final = scenario.model(model="openai/complex-tier")
name: Final = f"router-{uuid.uuid4().hex[:10]}"
created: Final = scenario.gateway.post(
"/model/new",
{
"model_name": name,
"litellm_params": {
"model": "auto_router/complexity_router",
"complexity_router_config": _router_config(classifier, simple, complex_),
},
"model_info": {},
},
)
scenario.cleanups.callback(scenario.delete_model, string_value(object_value(created["model_info"])["id"]))
return name
def _chat(gateway: Gateway, model: str, *, key: str | None = None) -> httpx.Response:
return gateway.request(
"POST",
"/v1/chat/completions",
{"model": model, "messages": [{"role": "user", "content": f"design a distributed cache {uuid.uuid4().hex}"}]},
key=key,
)
def _assert_classified_by_the_complex_tier(response: httpx.Response, router: str) -> None:
assert response.status_code == 200, response.text
assert response.json()["choices"][0]["message"]["content"], response.text
assert (
response.headers["x-litellm-model-group"],
response.headers["x-litellm-model-name"],
response.headers["x-litellm-complexity-router-tier"],
response.headers["x-litellm-complexity-router-cause"],
) == (router, "openai/complex-tier", "COMPLEX", "jev_classifier"), dict(response.headers)
assert "x-litellm-classifier-cost" not in response.headers, dict(response.headers)
def _assert_classifier_call(
call: Mapping[str, JsonValue], *, path: str, authorization: str, prompt_marker: str
) -> None:
assert call["path"] == path, call
assert call["authorization"] == authorization, call
body: Final = object_value(call["body"])
assert "model" not in body, body
assert prompt_marker in string_value(body["state"]), body
assert sorted(object_value(object_value(body["questions"])["tier"])) == ["criteria", "instructions", "type"], body
def _assert_refused_before_any_upstream_call(gateway: Gateway, model: str, message: str) -> None:
with gateway.scenario() as scenario:
handle: Final = _register(scenario, _AI_DECIDE_BODY)
deployment: Final = scenario.model(model=model, api_base=handle.api_base(), api_key=_API_KEY)
response: Final = _decide(gateway, deployment)
assert response.status_code == 400, response.text
assert message in response.text, response.text
assert _upstream_calls(gateway, handle) == []
row: Final = _spend_row(response.headers["x-litellm-call-id"])
assert (row["status"], row["call_type"], _number(row["spend"])) == ("failure", "asystemone", 0.0), row
@pytest.mark.parametrize("name", [name for _, name in _NON_BARE_NAMES], ids=[label for label, _ in _NON_BARE_NAMES])
def test_a_deployment_naming_a_non_bare_endpoint_is_refused_before_any_upstream_call(
gateway: Gateway, name: str
) -> None:
_assert_refused_before_any_upstream_call(gateway, f"databricks/{name}", _MODEL_RULE)
def test_a_deployment_naming_no_model_is_refused_before_any_upstream_call(gateway: Gateway) -> None:
_assert_refused_before_any_upstream_call(gateway, "databricks/", "A model name is required for the Decisions API")
def test_an_openai_format_request_at_v1_decisions_reaches_the_serving_endpoint_as_system_one(gateway: Gateway) -> None:
with gateway.scenario() as scenario:
handle: Final = _register(scenario, _ANSWER_BODY)
deployment: Final = scenario.model(
model=f"databricks/{_ENDPOINT}", api_base=handle.api_base(), api_key=_API_KEY
)
response: Final = gateway.request(
"POST",
"/v1/decisions",
{
"model": deployment,
"input": _OPENAI_INPUT,
"questions": [_OPENAI_TIER_QUESTION],
},
)
assert response.status_code == 200, response.text
assert response.json() == {
"model": _ENDPOINT,
"answers": [
{
"type": "choice",
"name": "tier",
"choice": "COMPLEX",
"probabilities": [
{"value": "SIMPLE", "probability": 0.03},
{"value": "MEDIUM", "probability": 0.06},
{"value": "COMPLEX", "probability": 0.91},
],
"confidence": 0.91,
}
],
"usage": {
"input_tokens": 367,
"input_tokens_details": {"cached_tokens": 0, "cache_write_tokens": 0},
"output_tokens": 3,
"output_tokens_details": {"reasoning_tokens": 0},
"total_tokens": 370,
},
}, response.text
assert response.headers["x-litellm-model-name"] == f"databricks/{_ENDPOINT}", dict(response.headers)
(call,) = _upstream_calls(gateway, handle)
assert call["path"] == f"/{handle.scenario_id}/{_ENDPOINT}/invocations", call
assert call["authorization"] == f"Bearer {_API_KEY}", call
assert call["body"] == {
"model": _ENDPOINT,
"state": _OPENAI_INPUT,
"questions": {"tier": {"type": "choice", "instructions": "Which tier?", "criteria": _TIER_CRITERIA}},
}, call
row: Final = _spend_row(response.headers["x-litellm-call-id"])
assert (
row["status"],
row["call_type"],
row["custom_llm_provider"],
row["model_group"],
row["api_base"],
_number(row["spend"]),
) == ("success", "adecisions", "databricks", deployment, f"{handle.api_base()}/{_ENDPOINT}/invocations", 0.0), (
row
)
def test_an_openai_format_request_at_v1_decisions_reaches_the_ai_decide_route_without_a_model_field(
gateway: Gateway,
) -> None:
with gateway.scenario() as scenario:
handle: Final = _register(scenario, _AI_DECIDE_BODY)
deployment: Final = scenario.model(model=f"databricks/{_MODEL}", api_base=handle.api_base(), api_key=_API_KEY)
response: Final = gateway.request(
"POST",
"/v1/decisions",
{
"model": deployment,
"input": _OPENAI_INPUT,
"questions": [_OPENAI_TIER_QUESTION],
},
)
assert response.status_code == 200, response.text
assert response.json() == {
"model": f"databricks/{_MODEL}",
"answers": [
{
"type": "choice",
"name": "tier",
"choice": "COMPLEX",
"probabilities": [
{"value": "SIMPLE", "probability": 0.03},
{"value": "MEDIUM", "probability": 0.06},
{"value": "COMPLEX", "probability": 0.91},
],
"confidence": 0.91,
}
],
"usage": {
"input_tokens": 0,
"input_tokens_details": {"cached_tokens": 0, "cache_write_tokens": 0},
"output_tokens": 0,
"output_tokens_details": {"reasoning_tokens": 0},
"total_tokens": 0,
},
}, response.text
assert response.headers["x-litellm-model-name"] == f"databricks/{_MODEL}", dict(response.headers)
(call,) = _upstream_calls(gateway, handle)
assert call["path"] == f"/{handle.scenario_id}{_ROUTE}", call
assert call["authorization"] == f"Bearer {_API_KEY}", call
assert call["body"] == {
"state": _OPENAI_INPUT,
"questions": {"tier": {"type": "choice", "instructions": "Which tier?", "criteria": _TIER_CRITERIA}},
}, call
row: Final = _spend_row(response.headers["x-litellm-call-id"])
assert (
row["status"],
row["call_type"],
row["custom_llm_provider"],
row["model_group"],
row["api_base"],
_number(row["spend"]),
) == ("success", "adecisions", "databricks", deployment, f"{handle.api_base()}{_ROUTE}", 0.0), row
def test_a_noul_probability_reaches_the_caller_as_the_system_one_noul_answer(gateway: Gateway) -> None:
with gateway.scenario() as scenario:
handle: Final = _register(
scenario, {"response": {"answers": {"defect": {"type": "noul", "probability": 0.8}}}, "metadata": _METADATA}
)
deployment: Final = scenario.model(model=f"databricks/{_MODEL}", api_base=handle.api_base(), api_key=_API_KEY)
response: Final = gateway.request(
"POST",
"/v1/systemone",
{"model": deployment, "state": _STATE, "questions": {"defect": {"type": "noul", "instructions": "Bug?"}}},
)
assert response.status_code == 200, response.text
assert response.json()["answers"] == {"defect": {"type": "noul", "noul": 0.8}}, response.text
(call,) = _upstream_calls(gateway, handle)
assert call["body"] == {"state": _STATE, "questions": {"defect": {"type": "noul", "instructions": "Bug?"}}}
def test_a_wildcard_deployment_sends_each_request_to_the_named_endpoint_and_bills_each_once(gateway: Gateway) -> None:
with gateway.scenario() as scenario:
handle: Final = _register(scenario, _ANSWER_BODY)
prefix: Final = _wildcard_deployment(scenario, api_base=handle.api_base(), api_key=_API_KEY)
model: Final = f"{prefix}/{_ENDPOINT}"
responses: Final = tuple(_decide(gateway, model) for _ in range(2))
for response in responses:
assert response.status_code == 200, response.text
assert response.json() == _ANSWER_BODY, response.text
assert response.headers["x-litellm-model-group"] == model, dict(response.headers)
assert response.headers["x-litellm-model-name"] == f"databricks/{_ENDPOINT}", dict(response.headers)
call_ids: Final = tuple(response.headers["x-litellm-call-id"] for response in responses)
assert len(set(call_ids)) == 2, call_ids
calls: Final = _upstream_calls(gateway, handle)
assert len(calls) == 2, calls
for call in calls:
assert call["path"] == f"/{handle.scenario_id}/{_ENDPOINT}/invocations", call
assert call["authorization"] == f"Bearer {_API_KEY}", call
assert call["body"] == {"model": _ENDPOINT, "state": _STATE, "questions": _QUESTIONS}, call
for call_id in call_ids:
row: Final = _spend_row(call_id)
assert (
row["status"],
row["call_type"],
row["custom_llm_provider"],
row["model_group"],
row["api_base"],
_number(row["spend"]),
) == ("success", "asystemone", "databricks", model, f"{handle.api_base()}/{_ENDPOINT}/invocations", 0.0), (
row
)
def test_a_wildcard_deployment_sends_each_ai_decide_request_to_the_route_and_logs_each_once(gateway: Gateway) -> None:
with gateway.scenario() as scenario:
handle: Final = _register(scenario, _AI_DECIDE_BODY)
prefix: Final = _wildcard_deployment(scenario, api_base=handle.api_base(), api_key=_API_KEY)
model: Final = f"{prefix}/{_MODEL}"
responses: Final = tuple(_decide(gateway, model) for _ in range(2))
for response in responses:
assert response.status_code == 200, response.text
assert response.json() == _SYSTEM_ONE_BODY, response.text
assert response.headers["x-litellm-model-group"] == model, dict(response.headers)
assert response.headers["x-litellm-model-name"] == f"databricks/{_MODEL}", dict(response.headers)
call_ids: Final = tuple(response.headers["x-litellm-call-id"] for response in responses)
assert len(set(call_ids)) == 2, call_ids
calls: Final = _upstream_calls(gateway, handle)
assert len(calls) == 2, calls
for call in calls:
assert call["path"] == f"/{handle.scenario_id}{_ROUTE}", call
assert call["authorization"] == f"Bearer {_API_KEY}", call
assert call["body"] == {"state": _STATE, "questions": _QUESTIONS}, call
for call_id in call_ids:
row: Final = _spend_row(call_id)
assert (
row["status"],
row["call_type"],
row["custom_llm_provider"],
row["model_group"],
row["api_base"],
_number(row["spend"]),
) == ("success", "asystemone", "databricks", model, f"{handle.api_base()}{_ROUTE}", 0.0), row
@pytest.mark.parametrize("model", (5, ["a"]), ids=("int", "list"))
def test_a_non_string_model_is_refused_at_the_gateway_without_an_upstream_call(
gateway: Gateway, model: JsonValue
) -> None:
with gateway.scenario() as scenario:
handle: Final = _register(scenario, _AI_DECIDE_BODY)
_wildcard_deployment(scenario, api_base=handle.api_base(), api_key=_API_KEY)
response: Final = gateway.request(
"POST", "/v1/systemone", {"model": model, "state": _STATE, "questions": _QUESTIONS}
)
assert 400 <= response.status_code < 500, response.text
assert _upstream_calls(gateway, handle) == []
def test_a_five_kilobyte_endpoint_name_reaches_the_upstream_and_its_not_found_answer_is_kept(gateway: Gateway) -> None:
name: Final = "n" * 5000
with gateway.scenario() as scenario:
handle: Final = _register(
scenario,
{"error_code": "RESOURCE_DOES_NOT_EXIST", "message": f"Endpoint with name '{name}' does not exist."},
status=404,
)
prefix: Final = _wildcard_deployment(scenario, api_base=handle.api_base(), api_key=_API_KEY)
response: Final = _decide(gateway, f"{prefix}/{name}")
assert response.status_code == 404, response.text[:500]
assert "RESOURCE_DOES_NOT_EXIST" in response.text, response.text[:500]
(call,) = _upstream_calls(gateway, handle)
assert call["path"] == f"/{handle.scenario_id}/{name}/invocations", call["path"][:200]
assert call["authorization"] == f"Bearer {_API_KEY}", call["authorization"]
assert call["body"] == {"model": name, "state": _STATE, "questions": _QUESTIONS}
row: Final = _spend_row(response.headers["x-litellm-call-id"])
assert (row["status"], row["call_type"], row["model_group"], _number(row["spend"])) == (
"failure",
"asystemone",
f"{prefix}/{name}",
0.0,
), row
@pytest.mark.parametrize(
"name,path,status,body", _DATABRICKS_ERRORS, ids=("ai_decide-400", "ai_decide-403", "serving_endpoint-404")
)
def test_a_databricks_error_keeps_its_status_and_message_and_bills_nothing(
gateway: Gateway, name: str, path: str, status: int, body: dict[str, JsonValue]
) -> None:
with gateway.scenario() as scenario:
handle: Final = _register(scenario, body, status=status)
model: Final = scenario.model(model=f"databricks/{name}", api_base=handle.api_base(), api_key=_API_KEY)
response: Final = _decide(gateway, model)
assert response.status_code == status, response.text
assert string_value(body["message"]) in response.text, response.text
(call,) = _upstream_calls(gateway, handle)
assert call["path"] == f"/{handle.scenario_id}{path}", call
row: Final = _spend_row(response.headers["x-litellm-call-id"])
assert (row["status"], row["call_type"], row["model_group"], _number(row["spend"])) == (
"failure",
"asystemone",
model,
0.0,
), row
def test_a_deployment_without_a_key_and_no_environment_key_is_refused_before_any_upstream_call(
gateway: Gateway,
) -> None:
with gateway.scenario() as scenario:
handle: Final = _register(scenario, _AI_DECIDE_BODY)
model: Final = scenario.model(model=f"databricks/{_MODEL}", api_base=handle.api_base(), api_key=None)
response: Final = _decide(gateway, model)
assert response.status_code == 401, response.text
assert _MISSING_KEY in response.text, response.text
assert _upstream_calls(gateway, handle) == []
@pytest.mark.timeout(_OWNED_PROXY_BUDGET)
def test_the_environment_key_and_base_serve_a_deployment_a_classifier_and_a_team_member_that_leave_them_unset(
gateway: Gateway, tmp_path: Path
) -> None:
with gateway.scenario() as rig_scenario:
handle: Final = _register(rig_scenario, _AI_DECIDE_BODY)
environment_base: Final = f"{handle.api_base()}/environment-workspace/serving-endpoints"
environment_path: Final = f"/{handle.scenario_id}/environment-workspace{_ROUTE}"
environment: Final = {
"DATABRICKS_API_KEY": _ENV_KEY,
"DATABRICKS_TOKEN": _ENV_TOKEN,
"DATABRICKS_API_BASE": environment_base,
}
with owned_proxy_process(gateway, tmp_path, environment) as owned, owned.gateway.scenario() as scenario:
candidate: Final = owned.gateway
with_base: Final = scenario.model(model=f"databricks/{_MODEL}", api_base=handle.api_base(), api_key=None)
bare: Final = scenario.model(model=f"databricks/{_MODEL}", api_base=None, api_key=None)
decided: Final = tuple(_decide(candidate, model) for model in (with_base, bare))
for response in decided:
assert response.status_code == 200, response.text
assert response.json() == _SYSTEM_ONE_BODY, response.text
assert [(call["path"], call["authorization"]) for call in _upstream_calls(candidate, handle)] == [
(f"/{handle.scenario_id}{_ROUTE}", f"Bearer {_ENV_KEY}"),
(environment_path, f"Bearer {_ENV_KEY}"),
]
for response, model in zip(decided, (with_base, bare)):
row: Final = _spend_row(response.headers["x-litellm-call-id"])
assert (row["status"], row["model_group"], _number(row["spend"])) == ("success", model, 0.0), row
assert (
_spend_row(decided[1].headers["x-litellm-call-id"])["api_base"]
== f"{handle.api_base()}/environment-workspace{_ROUTE}"
)
router: Final = _router(scenario, _environment_classifier())
routed: Final = _chat(candidate, router)
_assert_classified_by_the_complex_tier(routed, router)
(judge_call,) = _upstream_calls(candidate, handle)
_assert_classifier_call(
judge_call,
path=environment_path,
authorization=f"Bearer {_ENV_KEY}",
prompt_marker="design a distributed cache",
)
team: Final = scenario.team(team_member_permissions=["/auto_router/manage"])
member: Final = scenario.member(team)
member_key: Final = scenario.key(user_id=member, team_id=team)
simple: Final = scenario.model(model="openai/simple-tier")
complex_: Final = scenario.model(model="openai/complex-tier")
team_router: Final = f"team-router-{uuid.uuid4().hex[:8]}"
created: Final = candidate.request(
"POST",
"/model/new",
{
"model_name": team_router,
"litellm_params": {
"model": "auto_router/complexity_router",
"complexity_router_config": _router_config(_environment_classifier(), simple, complex_),
},
"model_info": {"team_id": team},
},
key=member_key,
)
assert created.status_code == 200, created.text
team_router_id: Final = string_value(object_value(object_value(created.json())["model_info"])["id"])
scenario.cleanups.callback(scenario.delete_model, team_router_id)
stored: Final = candidate.request("GET", "/model/info", params={"litellm_model_id": team_router_id})
assert stored.status_code == 200, stored.text
(entry,) = stored.json()["data"]
stored_classifier: Final = object_value(
object_value(object_value(object_value(entry)["litellm_params"])["complexity_router_config"])[
"opensource_classifier_config"
]
)
assert (stored_classifier["provider"], stored_classifier["model"]) == ("databricks", _MODEL), entry
assert "api_key" not in stored_classifier and "api_base" not in stored_classifier, entry
member_chat: Final = _chat(candidate, team_router, key=member_key)
_assert_classified_by_the_complex_tier(member_chat, team_router)
(member_judge_call,) = _upstream_calls(candidate, handle)
_assert_classifier_call(
member_judge_call,
path=environment_path,
authorization=f"Bearer {_ENV_KEY}",
prompt_marker="design a distributed cache",
)
@pytest.mark.timeout(_OWNED_PROXY_BUDGET)
def test_the_environment_token_serves_a_deployment_and_a_classifier_when_no_environment_key_is_set(
gateway: Gateway, tmp_path: Path
) -> None:
with gateway.scenario() as rig_scenario:
handle: Final = _register(rig_scenario, _AI_DECIDE_BODY)
environment_path: Final = f"/{handle.scenario_id}/environment-workspace{_ROUTE}"
environment: Final = {
"DATABRICKS_TOKEN": _ENV_TOKEN,
"DATABRICKS_API_BASE": f"{handle.api_base()}/environment-workspace/serving-endpoints",
}
with (
owned_proxy_process(gateway, tmp_path, environment, remove_environment=("DATABRICKS_API_KEY",)) as owned,
owned.gateway.scenario() as scenario,
):
candidate: Final = owned.gateway
bare: Final = scenario.model(model=f"databricks/{_MODEL}", api_base=None, api_key=None)
decided: Final = _decide(candidate, bare)
assert decided.status_code == 200, decided.text
assert decided.json() == _SYSTEM_ONE_BODY, decided.text
(decision_call,) = _upstream_calls(candidate, handle)
assert (decision_call["path"], decision_call["authorization"]) == (environment_path, f"Bearer {_ENV_TOKEN}")
assert _spend_row(decided.headers["x-litellm-call-id"])["status"] == "success"
router: Final = _router(scenario, _environment_classifier())
routed: Final = _chat(candidate, router)
_assert_classified_by_the_complex_tier(routed, router)
(judge_call,) = _upstream_calls(candidate, handle)
_assert_classifier_call(
judge_call,
path=environment_path,
authorization=f"Bearer {_ENV_TOKEN}",
prompt_marker="design a distributed cache",
)
def test_a_megabyte_model_name_is_refused_in_linear_time_while_liveliness_stays_up(gateway: Gateway) -> None:
with gateway.scenario() as scenario:
handle: Final = _register(scenario, _AI_DECIDE_BODY)
prefix: Final = _wildcard_deployment(scenario, api_base=handle.api_base(), api_key=_API_KEY)
model: Final = f"{prefix}/{_MEGABYTE_NAME}"
anonymous: Final = gateway.client.post(
"/v1/systemone", json={"model": model, "state": _STATE, "questions": _QUESTIONS}
)
assert anonymous.status_code == 401, anonymous.text[:300]
tenth_response, tenth_elapsed, _ = _timed_decide_under_liveliness_probes(
gateway, f"{prefix}/{_TENTH_MEGABYTE_NAME}"
)
response, elapsed, liveliness = _timed_decide_under_liveliness_probes(gateway, model)
assert tenth_response.status_code == 400, tenth_response.text[:300]
assert_within_factor_of_control(elapsed, control=tenth_elapsed, factor=25)
assert liveliness and all(status == 200 for status, _ in liveliness), liveliness
assert max(seconds for _, seconds in liveliness) < max(1.0, elapsed / 2), (elapsed, liveliness)
assert response.status_code == 400, response.text[:300]
assert _MODEL_RULE in response.text, response.text[:400]
assert _upstream_calls(gateway, handle) == []
row: Final = _spend_row(response.headers["x-litellm-call-id"])
assert (row["status"], row["call_type"], _number(row["spend"])) == ("failure", "asystemone", 0.0), row
Reference in a new issue View git blame Copy permalink