mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-05 08:07:05 +00:00
08fa25042c
The shared proxy wrapper in tests/e2e/e2e_gateway.py was misnamed: Gateway is not a gateway server, it is the client every suite uses to talk to the proxy (keys, models, chat/embed/ocr, spend read-backs, poll helpers). Rename the module to proxy_client.py and the class to ProxyClient, with build_gateway becoming build_proxy_client and the GatewayProvider protocol becoming ProxyClientProvider. The .gateway attribute suites held is now .proxy. Only identifiers changed; prose and string literals that use the word gateway for the proxy-server concept were left alone. Each suite previously built its own instance through a per-suite build_client() that called build_gateway() inside, duplicating the proxy wiring across suites. There is now one session-scoped proxy fixture in tests/e2e/conftest.py; every suite's client fixture depends on it and injects it, so the wiring lives in one place. claude_code keeps building its own client directly since it has its own harness and does not use the shared fixtures. Behavior is unchanged: shared transport, data-plane/control-plane split routing, poll budget, typed request/response models, and resource cleanup all go through the same object.
56 lines
1.6 KiB
Python
56 lines
1.6 KiB
Python
"""Client for the access-control e2e suite."""
|
|
|
|
from __future__ import annotations
|
|
|
|
from dataclasses import dataclass
|
|
|
|
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"
|
|
|
|
|
|
@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)
|