mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-25 05:54:53 +00:00
352 lines
14 KiB
Python
352 lines
14 KiB
Python
from __future__ import annotations
|
|
|
|
import os
|
|
import time
|
|
import uuid
|
|
from collections.abc import Callable, Iterator, Mapping, Sequence
|
|
from concurrent.futures import ThreadPoolExecutor
|
|
from contextlib import ExitStack, contextmanager
|
|
from dataclasses import dataclass
|
|
from hashlib import sha256
|
|
from typing import Final, TypeVar
|
|
|
|
import httpx
|
|
from pydantic import JsonValue, TypeAdapter
|
|
|
|
from tests.integration._support.database import read_rows
|
|
|
|
JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue])
|
|
DEPLOYMENT_IDS: Final = TypeAdapter(dict[str, list[str]])
|
|
GATEWAY_LIMITS: Final = httpx.Limits(keepalive_expiry=2)
|
|
T = TypeVar("T")
|
|
|
|
|
|
def object_value(value: JsonValue) -> dict[str, JsonValue]:
|
|
return JSON_OBJECT.validate_python(value)
|
|
|
|
|
|
def string_value(value: JsonValue) -> str:
|
|
assert isinstance(value, str), f"Expected a string, received {type(value).__name__}"
|
|
return value
|
|
|
|
|
|
def list_value(value: JsonValue) -> Sequence[JsonValue]:
|
|
assert isinstance(value, list), f"Expected a list, received {type(value).__name__}"
|
|
return value
|
|
|
|
|
|
def delete_key_if_present(candidate: Gateway, key: str) -> None:
|
|
digest: Final = sha256(key.encode()).hexdigest()
|
|
if candidate.rows('SELECT token FROM "LiteLLM_VerificationToken" WHERE token=%s', (digest,)):
|
|
candidate.post("/key/delete", {"keys": [key]})
|
|
assert candidate.rows('SELECT token FROM "LiteLLM_VerificationToken" WHERE token=%s', (digest,)) == []
|
|
|
|
|
|
def eventually(
|
|
read: Callable[[], T],
|
|
satisfied: Callable[[T], bool],
|
|
seconds: float = 10,
|
|
return_last_on_timeout: bool = False,
|
|
) -> T:
|
|
deadline: Final = time.monotonic() + seconds
|
|
while True:
|
|
observed: Final = read()
|
|
if satisfied(observed):
|
|
return observed
|
|
if return_last_on_timeout and time.monotonic() >= deadline:
|
|
return observed
|
|
assert time.monotonic() < deadline, f"State did not converge: {observed!r}"
|
|
time.sleep(0.1)
|
|
|
|
|
|
def until_asserted(check: Callable[[], T], *, seconds: float = 10) -> T:
|
|
deadline: Final = time.monotonic() + seconds
|
|
while True:
|
|
try:
|
|
return check()
|
|
except AssertionError:
|
|
if time.monotonic() >= deadline:
|
|
raise
|
|
time.sleep(0.1)
|
|
|
|
|
|
def held(read: Callable[[], T], satisfied: Callable[[T], bool], *, holding: float, seconds: float = 10) -> T:
|
|
"""`eventually`, and then `satisfied` must hold on every read for `holding` more seconds; a miss restarts the hold."""
|
|
deadline: Final = time.monotonic() + seconds
|
|
while True:
|
|
converged: Final = eventually(read, satisfied, seconds=max(deadline - time.monotonic(), 0.0))
|
|
misses: Final = _misses_within(read, satisfied, holding)
|
|
if not misses:
|
|
return converged
|
|
assert time.monotonic() < deadline, f"State did not hold: {misses[0]!r}"
|
|
time.sleep(0.1)
|
|
|
|
|
|
def _misses_within(read: Callable[[], T], satisfied: Callable[[T], bool], holding: float) -> tuple[T, ...]:
|
|
until: Final = time.monotonic() + holding
|
|
while time.monotonic() < until:
|
|
time.sleep(0.1)
|
|
observed: Final = read()
|
|
if not satisfied(observed):
|
|
return (observed,)
|
|
return ()
|
|
|
|
|
|
@dataclass(frozen=True, slots=True)
|
|
class Gateway:
|
|
client: httpx.Client
|
|
key: str
|
|
upstream_url: str
|
|
database_url: str | None = None
|
|
workers: int = 1
|
|
|
|
def rows(self, query: str, parameters: tuple[str, ...]) -> list[dict[str, JsonValue]]:
|
|
return read_rows(query, parameters, database_url=self.database_url)
|
|
|
|
def request(
|
|
self,
|
|
method: str,
|
|
path: str,
|
|
body: Mapping[str, JsonValue] | None = None,
|
|
*,
|
|
key: str | None = None,
|
|
params: Mapping[str, str] | None = None,
|
|
headers: Mapping[str, str] | None = None,
|
|
) -> httpx.Response:
|
|
request_headers: Final = {
|
|
"Authorization": f"Bearer {self.key if key is None else key}",
|
|
**(headers or {}),
|
|
}
|
|
return self.client.request(
|
|
method,
|
|
path,
|
|
json=body,
|
|
params=params,
|
|
headers=request_headers,
|
|
)
|
|
|
|
def request_multipart(
|
|
self,
|
|
path: str,
|
|
fields: Mapping[str, str],
|
|
files: Mapping[str, tuple[str, bytes, str]],
|
|
*,
|
|
key: str | None = None,
|
|
) -> httpx.Response:
|
|
return self.client.post(
|
|
path,
|
|
data=fields,
|
|
files=files,
|
|
headers={"Authorization": f"Bearer {self.key if key is None else key}"},
|
|
)
|
|
|
|
def post(self, path: str, body: Mapping[str, JsonValue], *, key: str | None = None) -> dict[str, JsonValue]:
|
|
response: Final = self.request("POST", path, body, key=key)
|
|
assert response.status_code == 200, f"POST {path}: {response.status_code} {response.text}"
|
|
return JSON_OBJECT.validate_json(response.content)
|
|
|
|
def get(self, path: str, params: Mapping[str, str] | None = None) -> dict[str, JsonValue]:
|
|
response: Final = self.request("GET", path, params=params)
|
|
assert response.status_code == 200, f"GET {path}: {response.status_code} {response.text}"
|
|
return JSON_OBJECT.validate_json(response.content)
|
|
|
|
def chat(self, model: str, *, key: str | None = None, text: str = "integration control") -> dict[str, JsonValue]:
|
|
return self.post(
|
|
"/v1/chat/completions",
|
|
{"model": model, "messages": [{"role": "user", "content": text}]},
|
|
key=key,
|
|
)
|
|
|
|
@contextmanager
|
|
def scenario(self) -> Iterator[Scenario]:
|
|
with ExitStack() as cleanups:
|
|
yield Scenario(self, cleanups)
|
|
|
|
|
|
@dataclass(frozen=True, slots=True)
|
|
class Scenario:
|
|
gateway: Gateway
|
|
cleanups: ExitStack
|
|
|
|
def key(self, **fields: JsonValue) -> str:
|
|
created: Final = self.gateway.post("/key/generate", fields)
|
|
token: Final = string_value(created["key"])
|
|
self.cleanups.callback(self.delete_key, token)
|
|
return token
|
|
|
|
def team(self, **fields: JsonValue) -> str:
|
|
created: Final = self.gateway.post("/team/new", {"team_alias": f"integration-{uuid.uuid4().hex}", **fields})
|
|
identity: Final = string_value(created["team_id"])
|
|
self.cleanups.callback(self.delete_team, identity)
|
|
return identity
|
|
|
|
def delete_team(self, identity: str) -> None:
|
|
self.gateway.post("/team/delete", {"team_ids": [identity]})
|
|
assert self.gateway.rows('SELECT team_id FROM "LiteLLM_TeamTable" WHERE team_id = %s', (identity,)) == []
|
|
|
|
def project(self, team_id: str, **fields: JsonValue) -> str:
|
|
created: Final = self.gateway.post(
|
|
"/project/new", {"team_id": team_id, "project_alias": f"integration-{uuid.uuid4().hex}", **fields}
|
|
)
|
|
identity: Final = string_value(created["project_id"])
|
|
self.cleanups.callback(self.delete_project, identity)
|
|
return identity
|
|
|
|
def delete_project(self, identity: str) -> None:
|
|
response: Final = self.gateway.request("DELETE", "/project/delete", {"project_ids": [identity]})
|
|
assert response.status_code == 200, response.text
|
|
assert (
|
|
self.gateway.rows('SELECT project_id FROM "LiteLLM_ProjectTable" WHERE project_id = %s', (identity,)) == []
|
|
)
|
|
|
|
def budget(self, **fields: JsonValue) -> str:
|
|
created: Final = self.gateway.post("/budget/new", fields)
|
|
identity: Final = string_value(created["budget_id"])
|
|
self.cleanups.callback(self.delete_budget, identity)
|
|
return identity
|
|
|
|
def delete_budget(self, identity: str) -> None:
|
|
self.gateway.post("/budget/delete", {"id": identity})
|
|
assert self.gateway.rows('SELECT budget_id FROM "LiteLLM_BudgetTable" WHERE budget_id = %s', (identity,)) == []
|
|
|
|
def organization(self, **fields: JsonValue) -> str:
|
|
created: Final = self.gateway.post(
|
|
"/organization/new", {"organization_alias": f"integration-{uuid.uuid4().hex}", **fields}
|
|
)
|
|
identity: Final = string_value(created["organization_id"])
|
|
self.cleanups.callback(self.delete_organization, identity, string_value(created["budget_id"]))
|
|
return identity
|
|
|
|
def delete_organization(self, identity: str, budget_id: str) -> None:
|
|
response: Final = self.gateway.request("DELETE", "/organization/delete", {"organization_ids": [identity]})
|
|
assert response.status_code == 200, response.text
|
|
assert (
|
|
self.gateway.rows(
|
|
'SELECT organization_id FROM "LiteLLM_OrganizationTable" WHERE organization_id = %s', (identity,)
|
|
)
|
|
== []
|
|
)
|
|
self.delete_budget(budget_id)
|
|
|
|
def org_member(self, organization_id: str, role: str) -> str:
|
|
user_id: Final = self.user(user_role="internal_user")
|
|
self.gateway.post(
|
|
"/organization/member_add",
|
|
{"organization_id": organization_id, "member": {"role": role, "user_id": user_id}},
|
|
)
|
|
return user_id
|
|
|
|
def user(self, **fields: JsonValue) -> str:
|
|
created: Final = self.gateway.post(
|
|
"/user/new", {"user_id": f"integration-{uuid.uuid4().hex}", "auto_create_key": False, **fields}
|
|
)
|
|
identity: Final = string_value(created["user_id"])
|
|
self.cleanups.callback(self.delete_user, identity)
|
|
return identity
|
|
|
|
def delete_user(self, identity: str) -> None:
|
|
response: Final = self.gateway.request("POST", "/user/delete", {"user_ids": [identity]})
|
|
assert response.status_code == 200 and response.json() == 1, response.text
|
|
assert self.gateway.rows('SELECT user_id FROM "LiteLLM_UserTable" WHERE user_id = %s', (identity,)) == []
|
|
|
|
def member(self, team_id: str, role: str = "user") -> str:
|
|
"""Create an internal user and add them to ``team_id``; deleting the user later removes the membership."""
|
|
user_id: Final = self.user(user_role="internal_user")
|
|
self.gateway.post("/team/member_add", {"team_id": team_id, "member": {"role": role, "user_id": user_id}})
|
|
return user_id
|
|
|
|
def delete_key(self, token: str) -> None:
|
|
self.gateway.post("/key/delete", {"keys": [token]})
|
|
hashed: Final = sha256(token.encode()).hexdigest()
|
|
assert self.gateway.rows('SELECT token FROM "LiteLLM_VerificationToken" WHERE token = %s', (hashed,)) == []
|
|
info: Final = object_value(self.gateway.get("/key/info", {"key": hashed})["info"])
|
|
assert info["status"] == "deleted", f"Deleted key still served as live: {info['status']}"
|
|
|
|
def served_model_ids(self) -> frozenset[JsonValue]:
|
|
entries: Final = self.gateway.get("/model/info")["data"]
|
|
assert isinstance(entries, list)
|
|
return frozenset(object_value(object_value(entry)["model_info"])["id"] for entry in entries)
|
|
|
|
def delete_model(self, identity: str) -> None:
|
|
from litellm.constants import PROXY_CONFIG_RELOAD_INTERVAL_SECONDS
|
|
|
|
self.gateway.post("/model/delete", {"id": identity})
|
|
eventually(
|
|
self.served_model_ids,
|
|
lambda served: identity not in served,
|
|
seconds=2 * PROXY_CONFIG_RELOAD_INTERVAL_SECONDS,
|
|
)
|
|
assert (
|
|
self.gateway.rows('SELECT model_id FROM "LiteLLM_ProxyModelTable" WHERE model_id = %s', (identity,)) == []
|
|
)
|
|
|
|
def model(self, *, model_info: Mapping[str, JsonValue] | None = None, **parameters: JsonValue) -> str:
|
|
name: Final = f"integration-{uuid.uuid4().hex}"
|
|
created: Final = self.gateway.post(
|
|
"/model/new",
|
|
{
|
|
"model_name": name,
|
|
"litellm_params": {
|
|
"model": "openai/gpt-4o-mini",
|
|
"api_key": "integration-provider-key",
|
|
"api_base": f"{self.gateway.upstream_url}/v1",
|
|
**parameters,
|
|
},
|
|
"model_info": dict(model_info) if model_info is not None else {},
|
|
},
|
|
)
|
|
identity: Final = string_value(object_value(created["model_info"])["id"])
|
|
self.cleanups.callback(self.delete_model, identity)
|
|
if self.gateway.workers > 1:
|
|
self._await_every_worker(name, identity)
|
|
return name
|
|
|
|
def _await_every_worker(self, name: str, identity: str) -> None:
|
|
from litellm.constants import PROXY_CONFIG_RELOAD_INTERVAL_SECONDS
|
|
|
|
eventually(
|
|
lambda: self._deployment_ids_by_worker(name),
|
|
lambda answers: len({pid for pid, _ in answers}) >= self.gateway.workers
|
|
and all(identity in model_ids for _, model_ids in answers),
|
|
seconds=2 * PROXY_CONFIG_RELOAD_INTERVAL_SECONDS,
|
|
)
|
|
|
|
def _deployment_ids_by_worker(self, name: str) -> tuple[tuple[str, frozenset[str]], ...]:
|
|
with ThreadPoolExecutor(max_workers=8) as pool:
|
|
return tuple(pool.map(lambda _: self._deployment_ids_on_a_worker(name), range(16)))
|
|
|
|
def _deployment_ids_on_a_worker(self, name: str) -> tuple[str, frozenset[str]]:
|
|
response: Final = self.gateway.request(
|
|
"GET",
|
|
"/health/integration-deployment-ids",
|
|
params={"model_name": name},
|
|
headers={"connection": "close"},
|
|
)
|
|
assert response.status_code == 200, f"Deployment ids on a worker: {response.status_code} {response.text}"
|
|
return response.headers["x-integration-worker-pid"], frozenset(
|
|
DEPLOYMENT_IDS.validate_json(response.content)["model_ids"]
|
|
)
|
|
|
|
|
|
@contextmanager
|
|
def gateway_from_environment() -> Iterator[Gateway]:
|
|
url: Final = os.environ["INTEGRATION_PROXY_URL"]
|
|
upstream: Final = os.environ["INTEGRATION_UPSTREAM_URL"]
|
|
with httpx.Client(base_url=url, timeout=15, trust_env=False, limits=GATEWAY_LIMITS) as client:
|
|
yield Gateway(client, os.environ["INTEGRATION_MASTER_KEY"], upstream)
|
|
|
|
|
|
def _set_team_admin_permissions(gateway: Gateway, fields: Sequence[str]) -> None:
|
|
response: Final = gateway.request("PATCH", "/update/ui_settings", {"team_admin_editable_team_fields": list(fields)})
|
|
assert response.status_code == 200, response.text
|
|
|
|
|
|
@contextmanager
|
|
def team_admin_permissions(gateway: Gateway, fields: Sequence[str]) -> Iterator[None]:
|
|
"""Grant team admins ``fields`` proxy-wide for the block, then restore the prior grant."""
|
|
original: Final = object_value(gateway.get("/get/ui_settings")["values"]).get("team_admin_editable_team_fields")
|
|
_set_team_admin_permissions(gateway, fields)
|
|
try:
|
|
yield
|
|
finally:
|
|
_set_team_admin_permissions(gateway, [str(field) for field in original] if isinstance(original, list) else ())
|