litellm/tests/integration/_support/client.py

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, *, name: str | None = None, model_info: Mapping[str, JsonValue] | None = None, **parameters: JsonValue
    ) -> str:
        model_name: Final = name if name is not None else f"integration-{uuid.uuid4().hex}"
        created: Final = self.gateway.post(
            "/model/new",
            {
                "model_name": model_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(model_name, identity)
        return model_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 ())