scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-08-28 05:25:59 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 10
litellm/tests/e2e/management/management_client.py
mubashir1osmani a780d4e4e3
test(musty_leopard): cover customer chat/messages cost + streaming paths (#34164) 
* test(e2e): cover customer chat/messages cost + streaming paths

Fills five uncovered P0 registry cells matching the customer's confirmed stack
(OpenAI SDK, Bedrock, /v1/messages) and their per-request cost dependency:
- /v1/messages logs cost that matches the x-litellm-response-cost header (LIT-4076)
- OpenAI /chat/completions streams real content, and a non-streamed call is costed
- Bedrock Converse /chat/completions returns real content non-streamed and streamed

The streaming checks aggregate delta content and parse every chunk as JSON, so a
clean-but-empty stream or a truncated chunk fails instead of passing on a bare 200.

* test(e2e): add tool-use coverage for openai, bedrock converse, anthropic responses

Function-calling regression guards on the paths the customer's agentic SDK usage
exercises: OpenAI and Bedrock Converse /chat/completions, and Anthropic
/v1/responses. The model is forced to call a weather tool and the test asserts the
returned tool call names the function and carries JSON-parseable arguments with the
expected field, so a dropped tool_call or malformed argument JSON fails instead of
passing on a bare 200. Adds a minimal tool_calls field to the response OutMessage.

* test(e2e): cover bedrock converse responses + thinking

Adds llm.responses.bedrock_converse.basic/tool_use and
llm.chat_completions.bedrock_converse.thinking. The thinking test enables extended
thinking and requires reasoning_content plus a real answer, so a path that drops
the reasoning block fails rather than passing.

* test(e2e): cover bedrock embeddings + openai structured output and reasoning

Bedrock Titan embeddings return a real vector; OpenAI structured output must yield
schema-conforming JSON with the correct extracted values (age==42, not just valid
JSON); an OpenAI reasoning call must report reasoning tokens, so a non-reasoning
fallback fails. Adds response_format to ChatBody and reasoning-token details to Usage.

* test(e2e): cover vision + streaming tool calls on openai and bedrock converse

Vision on both providers must describe the image (not just 200); the streamed
OpenAI tool call is reassembled from its fragments and its argument JSON parsed, so
a stream that never completes the call or splits its JSON fails. Extends ChatMessage
content to a typed text/image union.

* test(e2e): cover openai prompt caching hit on repeated large prefix

A repeated large-prefix prompt must report cached prompt tokens on the second call,
so a cache regression that stops reusing the prefix (and silently re-bills full
input) fails here.

* test(e2e): cover openai audio speech + bedrock rerank and image generation

Marks the OpenAI TTS cell and adds Bedrock Titan rerank (top_n honored, scored) and
Bedrock Titan image generation (returns b64/url), the customer's non-chat AWS
surfaces.

* test(e2e): cover end-user (customer) create persistence

mgmt.end_user.new.happy_path: create an end-user via /customer/new and confirm
/customer/info reports it, the end-user-identity surface the customer relies on for
per-customer controls. Adds customer models + management-client methods.

* test(e2e): enforce key model allow-list on the passthrough route

other.auth.passthrough.model_allowlist_enforced: a key scoped to gemini must be
denied a claude call through the anthropic passthrough route (403), so custom-auth
scoping is not bypassable by going through passthrough instead of /chat/completions.

* test(e2e): address Greptile - assert stream data events, correlate messages spend by key

- streaming: assert len(stream_events) > 1 instead of chunks > 1, since chunks
  counts the terminal data: [DONE] marker and would pass a single content event
- messages cost: correlate the spend row by the unique scoped key rather than the
  Anthropic response id, which need not equal the proxy spend-log request_id
2026-07-21 18:57:11 -07:00

455 lines
15 KiB
Python
Raw Permalink Blame History

"""Client for the management-routes e2e suite: the shared ProxyClient plus the
key/team/user/organization writes, the info/list read-backs the tests assert,
and the raw-status calls judged by HTTP outcome (chat under a scoped key, an
llm-only key hitting a management route).
"""
from __future__ import annotations
import time
from dataclasses import dataclass
from proxy_client import ProxyClient
from e2e_http import NoBody, ProbeResult, Result, StreamingResponse, Success, UnknownApiError, unwrap
from models import (
ChatBody,
ChatMessage,
CustomerDeleteBody,
CustomerInfoParams,
CustomerNewBody,
CustomerResponse,
KeyBlockBody,
KeyDeleteBody,
KeyGenerateBody,
KeyGenerateResponse,
KeyListParams,
KeyListResponse,
KeyRegenerateBody,
KeyUpdateBody,
ModelDeleteBody,
OrgDeleteBody,
OrgInfoParams,
OrgInfoResponse,
OrgNewBody,
OrgNewResponse,
OrgUpdateBody,
TagDeleteBody,
TagListEntry,
TagListResponse,
TagNewBody,
TeamData,
TeamDeleteBody,
TeamInfoParams,
TeamInfoResponse,
TeamListResponse,
TeamMemberAddBody,
TeamMemberDeleteBody,
TeamMemberEntry,
TeamNewBody,
TeamNewResponse,
TeamUpdateBody,
UserDeleteBody,
UserDeleteResponse,
UserInfoParams,
UserInfoResponse,
UserListParams,
UserListResponse,
UserNewBody,
UserNewResponse,
UserUpdateBody,
)
MODEL_ACCESS_DENIED_MARKER = "key_model_access_denied"
ROUTE_NOT_ALLOWED_MARKER = "not allowed to call this route"
_TEAM_READY_ATTEMPTS = 15
_TEAM_READY_SLEEP_SECONDS = 0.4
@dataclass(frozen=True, slots=True)
class ManagementClient:
proxy: ProxyClient
def llm_only_key(self) -> str:
return self.proxy.generate_key(KeyGenerateBody(models=[], allowed_routes=["llm_api_routes"]))
def update_key_models(self, key: str, models: list[str]) -> None:
last: Result[NoBody] | None = None
for attempt in range(5):
last = self.proxy.transport.post(
"/key/update",
headers=self.proxy.transport.master,
json=KeyUpdateBody(key=key, models=models),
response_type=NoBody,
)
match last:
case Success():
return
case UnknownApiError(body=body) if (
"connecting to redis" in body.lower() or "name resolution" in body.lower()
):
time.sleep(0.5 * (attempt + 1))
continue
case _:
break
assert last is not None
raise AssertionError(last)
def delete_key_strict(self, key: str) -> None:
"""Strict delete for the act phase of a test: a failed delete is a hard
failure, unlike the warn-only ProxyClient.delete_key used at teardown."""
_ = unwrap(
self.proxy.transport.post(
"/key/delete",
headers=self.proxy.transport.master,
json=KeyDeleteBody(keys=[key]),
response_type=NoBody,
)
)
def delete_model_strict(self, model_id: str) -> None:
"""Strict delete for the act phase of a test: a failed delete is a hard
failure, unlike the warn-only ProxyClient.delete_model used at teardown."""
_ = unwrap(
self.proxy.transport.post(
"/model/delete",
headers=self.proxy.transport.master,
json=ModelDeleteBody(id=model_id),
response_type=NoBody,
)
)
def block_key(self, key: str) -> None:
_ = unwrap(
self.proxy.transport.post(
"/key/block",
headers=self.proxy.transport.master,
json=KeyBlockBody(key=key),
response_type=NoBody,
)
)
def regenerate_key(self, key: str) -> str:
return unwrap(
self.proxy.transport.post(
"/key/regenerate",
headers=self.proxy.transport.master,
json=KeyRegenerateBody(key=key),
response_type=KeyGenerateResponse,
)
).key
def key_alias_count(self, key_alias: str) -> int:
return unwrap(
self.proxy.transport.get(
"/key/list",
headers=self.proxy.transport.master,
params=KeyListParams(key_alias=key_alias),
response_type=KeyListResponse,
)
).total_count
def create_team(self, body: TeamNewBody) -> str:
team_id = unwrap(
self.proxy.transport.post(
"/team/new",
headers=self.proxy.transport.master,
json=body,
response_type=TeamNewResponse,
)
).team_id
self._wait_for_team(team_id)
return team_id
def update_team(self, body: TeamUpdateBody) -> None:
last: Result[NoBody] | None = None
for attempt in range(5):
last = self.proxy.transport.post(
"/team/update",
headers=self.proxy.transport.master,
json=body,
response_type=NoBody,
)
match last:
case Success():
return
case UnknownApiError(body=body_text) if (
"connecting to redis" in body_text.lower() or "name resolution" in body_text.lower()
):
time.sleep(0.5 * (attempt + 1))
continue
case _:
break
assert last is not None
raise AssertionError(last)
def delete_team(self, team_id: str) -> None:
_ = self.proxy.transport.post(
"/team/delete",
headers=self.proxy.transport.master,
json=TeamDeleteBody(team_ids=[team_id]),
response_type=NoBody,
)
def team_info(self, team_id: str) -> TeamData:
return unwrap(
self.proxy.transport.get(
"/team/info",
headers=self.proxy.transport.master,
params=TeamInfoParams(team_id=team_id),
response_type=TeamInfoResponse,
)
).team_info
def team_list_ids(self) -> tuple[str, ...]:
return tuple(
entry.team_id
for entry in unwrap(
self.proxy.transport.get(
"/team/list",
headers=self.proxy.transport.master,
params=NoBody(),
response_type=TeamListResponse,
)
).root
)
def team_info_status(self, team_id: str) -> ProbeResult:
return self.proxy.transport.probe("/team/info", params=TeamInfoParams(team_id=team_id))
def _wait_for_team(self, team_id: str) -> None:
last: Result[TeamInfoResponse] | None = None
for _ in range(_TEAM_READY_ATTEMPTS):
last = self.proxy.transport.get(
"/team/info",
headers=self.proxy.transport.master,
params=TeamInfoParams(team_id=team_id),
response_type=TeamInfoResponse,
)
match last:
case Success():
return
case _:
time.sleep(_TEAM_READY_SLEEP_SECONDS)
assert last is not None
raise AssertionError(last)
def add_team_member(self, team_id: str, user_id: str) -> None:
last: Result[NoBody] | None = None
for attempt in range(_TEAM_READY_ATTEMPTS):
last = self.proxy.transport.post(
"/team/member_add",
headers=self.proxy.transport.master,
json=TeamMemberAddBody(team_id=team_id, member=TeamMemberEntry(role="user", user_id=user_id)),
response_type=NoBody,
)
match last:
case Success():
return
case UnknownApiError(body=body) if (
"doesn't exist" in body and attempt + 1 < _TEAM_READY_ATTEMPTS
):
time.sleep(_TEAM_READY_SLEEP_SECONDS)
continue
case _:
break
assert last is not None
raise AssertionError(last)
def delete_team_member(self, team_id: str, user_id: str) -> None:
_ = unwrap(
self.proxy.transport.post(
"/team/member_delete",
headers=self.proxy.transport.master,
json=TeamMemberDeleteBody(team_id=team_id, user_id=user_id),
response_type=NoBody,
)
)
def create_user(self, body: UserNewBody) -> str:
return unwrap(
self.proxy.transport.post(
"/user/new",
headers=self.proxy.transport.master,
json=body,
response_type=UserNewResponse,
)
).user_id
def create_customer(self, user_id: str) -> str:
_ = unwrap(
self.proxy.transport.post(
"/customer/new",
headers=self.proxy.transport.master,
json=CustomerNewBody(user_id=user_id),
response_type=CustomerResponse,
)
)
return user_id
def customer_info(self, end_user_id: str) -> CustomerResponse:
return unwrap(
self.proxy.transport.get(
"/customer/info",
headers=self.proxy.transport.master,
params=CustomerInfoParams(end_user_id=end_user_id),
response_type=CustomerResponse,
)
)
def delete_customer(self, user_id: str) -> None:
_ = self.proxy.transport.post(
"/customer/delete",
headers=self.proxy.transport.master,
json=CustomerDeleteBody(user_ids=[user_id]),
response_type=NoBody,
)
def update_user(self, body: UserUpdateBody) -> None:
_ = unwrap(
self.proxy.transport.post(
"/user/update",
headers=self.proxy.transport.master,
json=body,
response_type=NoBody,
)
)
def delete_user(self, user_id: str) -> None:
_ = self.proxy.transport.post(
"/user/delete",
headers=self.proxy.transport.master,
json=UserDeleteBody(user_ids=[user_id]),
response_type=NoBody,
)
def delete_user_strict(self, user_id: str) -> None:
"""Strict delete for the act phase of a test: a failed delete is a hard
failure, unlike the warn-only delete_user used at teardown."""
_ = unwrap(
self.proxy.transport.post(
"/user/delete",
headers=self.proxy.transport.master,
json=UserDeleteBody(user_ids=[user_id]),
response_type=UserDeleteResponse,
)
)
def user_info(self, user_id: str) -> UserInfoResponse:
return unwrap(
self.proxy.transport.get(
"/user/info",
headers=self.proxy.transport.master,
params=UserInfoParams(user_id=user_id),
response_type=UserInfoResponse,
)
)
def user_count(self, user_id: str) -> int:
return unwrap(
self.proxy.transport.get(
"/user/list",
headers=self.proxy.transport.master,
params=UserListParams(user_ids=user_id),
response_type=UserListResponse,
)
).total
def user_list_ids(self, user_id: str) -> tuple[str, ...]:
listing = unwrap(
self.proxy.transport.get(
"/user/list",
headers=self.proxy.transport.master,
params=UserListParams(user_ids=user_id),
response_type=UserListResponse,
)
)
return tuple(row.user_id for row in listing.users)
def create_org(self, body: OrgNewBody) -> str:
return unwrap(
self.proxy.transport.post(
"/organization/new",
headers=self.proxy.transport.master,
json=body,
response_type=OrgNewResponse,
)
).organization_id
def update_org(self, body: OrgUpdateBody) -> None:
_ = unwrap(
self.proxy.transport.patch(
"/organization/update",
headers=self.proxy.transport.master,
json=body,
response_type=NoBody,
)
)
def delete_org(self, organization_id: str) -> None:
_ = self.proxy.transport.delete(
"/organization/delete",
headers=self.proxy.transport.master,
json=OrgDeleteBody(organization_ids=[organization_id]),
response_type=NoBody,
)
def org_info(self, organization_id: str) -> OrgInfoResponse:
return unwrap(
self.proxy.transport.get(
"/organization/info",
headers=self.proxy.transport.master,
params=OrgInfoParams(organization_id=organization_id),
response_type=OrgInfoResponse,
)
)
def org_info_status(self, organization_id: str) -> ProbeResult:
return self.proxy.transport.probe("/organization/info", params=OrgInfoParams(organization_id=organization_id))
def create_tag(self, body: TagNewBody) -> None:
_ = unwrap(
self.proxy.transport.post(
"/tag/new",
headers=self.proxy.transport.master,
json=body,
response_type=NoBody,
)
)
def delete_tag(self, name: str) -> None:
_ = self.proxy.transport.post(
"/tag/delete",
headers=self.proxy.transport.master,
json=TagDeleteBody(name=name),
response_type=NoBody,
)
def tag_list(self) -> tuple[TagListEntry, ...]:
return tuple(
unwrap(
self.proxy.transport.get(
"/tag/list",
headers=self.proxy.transport.master,
params=NoBody(),
response_type=TagListResponse,
)
).root
)
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)], max_tokens=16),
)
def key_generate_status(self, key: str, body: KeyGenerateBody) -> StreamingResponse:
return self.proxy.transport.send("/key/generate", headers=self.proxy.transport.bearer(key), json=body)
def team_new_status(self, key: str, body: TeamNewBody) -> StreamingResponse:
return self.proxy.transport.send("/team/new", headers=self.proxy.transport.bearer(key), json=body)
def user_new_status(self, key: str, body: UserNewBody) -> StreamingResponse:
return self.proxy.transport.send("/user/new", headers=self.proxy.transport.bearer(key), json=body)
def build_client(proxy: ProxyClient) -> ManagementClient:
return ManagementClient(proxy=proxy)
Reference in a new issue View git blame Copy permalink