mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-11 22:51:28 +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.
32 lines
1,022 B
Python
32 lines
1,022 B
Python
"""Client for the quota-management suite: the shared ProxyClient plus raw chat
|
|
calls judged by HTTP status, body, and headers (a rate-limit block is a 429
|
|
whose body and retry-after header carry the contract, not a typed success
|
|
model)."""
|
|
|
|
from __future__ import annotations
|
|
|
|
from dataclasses import dataclass
|
|
|
|
from proxy_client import ProxyClient
|
|
from e2e_http import StreamingResponse
|
|
from models import ChatBody, ChatMessage
|
|
|
|
|
|
@dataclass(frozen=True, slots=True)
|
|
class QuotaClient:
|
|
proxy: ProxyClient
|
|
|
|
def chat(self, key: str, model: str, content: str, *, max_tokens: int = 16) -> StreamingResponse:
|
|
return self.proxy.transport.send(
|
|
"/chat/completions",
|
|
headers=self.proxy.transport.bearer(key),
|
|
json=ChatBody(
|
|
model=model,
|
|
messages=[ChatMessage(role="user", content=content)],
|
|
max_tokens=max_tokens,
|
|
),
|
|
)
|
|
|
|
|
|
def build_client(proxy: ProxyClient) -> QuotaClient:
|
|
return QuotaClient(proxy=proxy)
|