mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-12 03:48:56 +00:00
1c6a31d009
Split /model/new from the data-plane wait so a propagation timeout still deletes the control-plane registration (greptile orphan-model concern)
122 lines
3.9 KiB
Python
122 lines
3.9 KiB
Python
"""Router suite's `client` fixture.
|
|
|
|
The shared lifecycle (resources/scoped_key), proxy liveness skip, and e2e marker
|
|
live in the parent tests/e2e/conftest.py. ComplexityRouterClient holds the shared
|
|
Gateway, so the `resources` fixture cleans up keys this suite creates.
|
|
|
|
Also registers `complexity-smart-router` via management /model/new when the
|
|
proxy does not already list it (compose has it in static config; stage does not).
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import time
|
|
from collections.abc import Iterator
|
|
|
|
import pytest
|
|
from requests import RequestException
|
|
|
|
from complexity_router_client import ComplexityRouterClient, build_client
|
|
from e2e_gateway import Gateway
|
|
from e2e_http import NoBody, Success, unwrap
|
|
from models import (
|
|
LiteLLMParamsBody,
|
|
ModelInfoBody,
|
|
ModelNewBody,
|
|
ModelNewResponse,
|
|
ModelsListResponse,
|
|
)
|
|
|
|
ROUTER_MODEL = "complexity-smart-router"
|
|
ROUTER_PARAMS = LiteLLMParamsBody(
|
|
model="auto_router/complexity_router",
|
|
complexity_router_config={
|
|
"classifier_type": "llm",
|
|
"classifier_llm_config": {"model": "gpt-5.5"},
|
|
"tiers": {
|
|
"SIMPLE": "gpt-5.5",
|
|
"MEDIUM": "claude-haiku-4-5",
|
|
"COMPLEX": "claude-haiku-4-5",
|
|
"REASONING": "claude-haiku-4-5",
|
|
},
|
|
},
|
|
)
|
|
|
|
|
|
@pytest.fixture(scope="session")
|
|
def client() -> ComplexityRouterClient:
|
|
return build_client()
|
|
|
|
|
|
def _model_is_servable(gateway: Gateway, model_name: str) -> bool:
|
|
result = gateway.transport.get(
|
|
"/v1/models",
|
|
headers=gateway.transport.master,
|
|
params=NoBody(),
|
|
response_type=ModelsListResponse,
|
|
)
|
|
return isinstance(result, Success) and any(entry.id == model_name for entry in result.data.data)
|
|
|
|
|
|
def _register_router_model(gateway: Gateway) -> str:
|
|
"""POST /model/new only; returns the proxy model_id before data-plane wait.
|
|
|
|
Split from create_model so a slow control→data propagation timeout still
|
|
leaves us a model_id for teardown (avoids orphaning complexity-smart-router).
|
|
"""
|
|
return unwrap(
|
|
gateway.transport.post(
|
|
"/model/new",
|
|
headers=gateway.transport.master,
|
|
json=ModelNewBody(
|
|
model_name=ROUTER_MODEL,
|
|
litellm_params=ROUTER_PARAMS,
|
|
model_info=ModelInfoBody(),
|
|
),
|
|
response_type=ModelNewResponse,
|
|
)
|
|
).model_id
|
|
|
|
|
|
def _await_router_model_servable(gateway: Gateway) -> None:
|
|
deadline = time.monotonic() + gateway.poll_timeout
|
|
while time.monotonic() < deadline:
|
|
if _model_is_servable(gateway, ROUTER_MODEL):
|
|
return
|
|
time.sleep(gateway.poll_interval)
|
|
raise AssertionError(
|
|
f"model {ROUTER_MODEL!r} was created but never became servable on the data "
|
|
f"plane within {gateway.poll_timeout}s of /model/new"
|
|
)
|
|
|
|
|
|
@pytest.fixture(scope="session", autouse=True)
|
|
def _ensure_complexity_smart_router( # pyright: ignore[reportUnusedFunction] # pytest autouse session fixture, wired by name
|
|
client: ComplexityRouterClient,
|
|
) -> Iterator[None]:
|
|
"""Ensure the complexity router virtual model exists for this session.
|
|
|
|
Compose already declares it in docker-compose.yml; stage does not. Register
|
|
via /model/new when missing and tear down only what we created.
|
|
"""
|
|
gateway = client.gateway
|
|
if _model_is_servable(gateway, ROUTER_MODEL):
|
|
yield
|
|
return
|
|
|
|
try:
|
|
model_id = _register_router_model(gateway)
|
|
except (AssertionError, RequestException) as exc:
|
|
if _model_is_servable(gateway, ROUTER_MODEL):
|
|
yield
|
|
return
|
|
raise AssertionError(
|
|
f"failed to register {ROUTER_MODEL!r} for the complexity router e2e "
|
|
f"(not listed on /v1/models and /model/new failed): {exc}"
|
|
) from exc
|
|
|
|
try:
|
|
_await_router_model_servable(gateway)
|
|
yield
|
|
finally:
|
|
gateway.delete_model(model_id)
|