mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-12 23:01:41 +00:00
ed01f7316b
Every case in TestAccessControl asserted that something was refused. A gateway
that denied the allow-listed model too would have passed all of them, so the
suite could not tell "denied correctly" from "broken outright".
Adds the positive half: a key allow-listed for gemini-2.5-flash can call it and
gets back a real completion rather than a 200-wrapped error.
Also tightens the unknown-model case. It accepted any valid JSON, so a bare
"{}" or even "null" satisfied it. It now requires the OpenAI-shaped error
envelope with a message a client can actually surface, parsed through a typed
model instead of json.loads.
76 lines
2.1 KiB
Python
76 lines
2.1 KiB
Python
"""Client for the access-control e2e suite."""
|
|
|
|
from __future__ import annotations
|
|
|
|
from dataclasses import dataclass
|
|
|
|
from pydantic import BaseModel, ValidationError
|
|
|
|
from proxy_client import ProxyClient
|
|
from e2e_http import StreamingResponse
|
|
from models import (
|
|
ChatBody,
|
|
ChatMessage,
|
|
KeyGenerateBody,
|
|
LiteLLMParamsBody,
|
|
ModelInfoBody,
|
|
ModelNewBody,
|
|
)
|
|
|
|
MODEL_ACCESS_DENIED_MARKER = "key_model_access_denied"
|
|
ROUTE_NOT_ALLOWED_MARKER = "not allowed to call this route"
|
|
|
|
|
|
class ApiErrorDetail(BaseModel):
|
|
message: str | None = None
|
|
type: str | None = None
|
|
code: str | int | None = None
|
|
|
|
|
|
class ApiErrorEnvelope(BaseModel):
|
|
error: ApiErrorDetail
|
|
|
|
|
|
def error_envelope(body: str) -> ApiErrorEnvelope | None:
|
|
"""The OpenAI-shaped `{"error": {...}}` a client parses, or None if absent."""
|
|
try:
|
|
return ApiErrorEnvelope.model_validate_json(body)
|
|
except ValidationError:
|
|
return None
|
|
|
|
|
|
@dataclass(frozen=True, slots=True)
|
|
class AccessControlClient:
|
|
proxy: ProxyClient
|
|
|
|
def llm_only_key(self) -> str:
|
|
return self.proxy.generate_key(
|
|
KeyGenerateBody(models=[], allowed_routes=["llm_api_routes"])
|
|
)
|
|
|
|
def delete_key(self, key: str) -> None:
|
|
self.proxy.delete_key(key)
|
|
|
|
def chat_status(self, key: str, model: str, content: str) -> StreamingResponse:
|
|
return self.proxy.transport.send(
|
|
"/chat/completions",
|
|
headers=self.proxy.transport.bearer(key),
|
|
json=ChatBody(
|
|
model=model, messages=[ChatMessage(role="user", content=content)]
|
|
),
|
|
)
|
|
|
|
def create_model_status(self, key: str, model_name: str) -> StreamingResponse:
|
|
return self.proxy.transport.send(
|
|
"/model/new",
|
|
headers=self.proxy.transport.bearer(key),
|
|
json=ModelNewBody(
|
|
model_name=model_name,
|
|
litellm_params=LiteLLMParamsBody(model="openai/gpt-4o-mini"),
|
|
model_info=ModelInfoBody(id=model_name),
|
|
),
|
|
)
|
|
|
|
|
|
def build_client(proxy: ProxyClient) -> AccessControlClient:
|
|
return AccessControlClient(proxy=proxy)
|